From ba70c871a8ae65e1812229bbec6b3379b1849ebf Mon Sep 17 00:00:00 2001 From: Vitalii Cherepanov Date: Sat, 10 Oct 2026 11:11:46 +0200 Subject: [PATCH 1/2] mcp: keep a delivered response when the caller's context is cancelled call checked ctx.Err() before it checked whether Await had succeeded, so a context cancelled while the result was being decoded discarded the result and sent notifications/cancelled for a completed request. Return the result when Await succeeds; a request still in flight is retired and notified as before. Fixes #1368 --- mcp/transport.go | 6 +++ mcp/transport_test.go | 121 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+) diff --git a/mcp/transport.go b/mcp/transport.go index 35c3e70b6..6a2ddec67 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -301,6 +301,12 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params call := conn.Call(ctx, method, params) err := call.Await(ctx, result) switch { + case err == nil: + // The response arrived and decoded. It is returned even if ctx has + // been cancelled in the meantime: the request is complete, so there + // is nothing to cancel, and the spec permits notifications/cancelled + // only for requests believed to still be in progress. + return nil case errors.Is(err, jsonrpc2.ErrClientClosing), errors.Is(err, jsonrpc2.ErrServerClosing): return fmt.Errorf("%w: calling %q: %v", ErrConnectionClosed, method, err) case ctx.Err() != nil: diff --git a/mcp/transport_test.go b/mcp/transport_test.go index dadf196b5..72b9b6bbf 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -10,7 +10,9 @@ import ( "errors" "fmt" "io" + "slices" "strings" + "sync" "testing" "time" @@ -364,3 +366,122 @@ func TestIOConnFrameCap(t *testing.T) { }) } } + +// cancelAfterResponseTransport wraps a Transport so that a test can cancel a +// context at a point the client cannot observe from the outside: the first +// Read after the response to the given method has been delivered. The +// jsonrpc2 read loop dispatches a message before it reads the next one, so +// by then the call has been retired with its result while the caller may +// still be decoding it. The transport also records the methods written to +// it, so a test can check that no cancelled notification followed. +type cancelAfterResponseTransport struct { + Transport + method string + cancel context.CancelFunc + + mu sync.Mutex + requests map[int64]string // id -> method, for requests written by the client + written []string +} + +func (t *cancelAfterResponseTransport) Connect(ctx context.Context) (Connection, error) { + conn, err := t.Transport.Connect(ctx) + if err != nil { + return nil, err + } + t.requests = make(map[int64]string) + return &cancelAfterResponseConn{Connection: conn, t: t}, nil +} + +type cancelAfterResponseConn struct { + Connection + t *cancelAfterResponseTransport + + delivered bool // the response to t.method has been returned from Read + cancelled bool +} + +func (c *cancelAfterResponseConn) Read(ctx context.Context) (jsonrpc.Message, error) { + if c.delivered && !c.cancelled { + c.cancelled = true + c.t.cancel() + } + msg, err := c.Connection.Read(ctx) + if resp, ok := msg.(*jsonrpc.Response); ok { + if id, isInt := resp.ID.Raw().(int64); isInt { + c.t.mu.Lock() + c.delivered = c.t.requests[id] == c.t.method + c.t.mu.Unlock() + } + } + return msg, err +} + +func (c *cancelAfterResponseConn) Write(ctx context.Context, msg jsonrpc.Message) error { + if req, ok := msg.(*jsonrpc.Request); ok { + c.t.mu.Lock() + if id, isInt := req.ID.Raw().(int64); isInt { + c.t.requests[id] = req.Method + } + c.t.written = append(c.t.written, req.Method) + c.t.mu.Unlock() + } + return c.Connection.Write(ctx, msg) +} + +func (t *cancelAfterResponseTransport) writtenMethods() []string { + t.mu.Lock() + defer t.mu.Unlock() + return slices.Clone(t.written) +} + +// TestCallKeepsDeliveredResponseOnCancel verifies that a response that has +// already arrived is returned to the caller even if the caller's context is +// cancelled while the result is being decoded, and that no cancelled +// notification is sent for a request the client has the response to. The +// specification allows cancellation notifications only for requests that +// are believed to still be in progress. +func TestCallKeepsDeliveredResponseOnCancel(t *testing.T) { + // A large tool list keeps the client decoding the result for a while + // after the response has been retired, which is the window the test + // targets. + const numTools = 2000 + + ctx := context.Background() + server := NewServer(testImpl, &ServerOptions{PageSize: numTools}) + for i := range numTools { + server.AddTool(&Tool{Name: fmt.Sprintf("tool-%04d", i), InputSchema: map[string]any{"type": "object"}}, + func(context.Context, *CallToolRequest) (*CallToolResult, error) { + return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil + }) + } + ct, st := NewInMemoryTransports() + ss, err := server.Connect(ctx, st, nil) + if err != nil { + t.Fatalf("server connect: %v", err) + } + t.Cleanup(func() { _ = ss.Close() }) + + callCtx, cancel := context.WithCancel(ctx) + defer cancel() + transport := &cancelAfterResponseTransport{Transport: ct, method: methodListTools, cancel: cancel} + cs, err := NewClient(testImpl, nil).Connect(ctx, transport, nil) + if err != nil { + t.Fatalf("client connect: %v", err) + } + t.Cleanup(func() { _ = cs.Close() }) + + res, err := cs.ListTools(callCtx, nil) + if err != nil { + t.Fatalf("ListTools() = %v, want the delivered result", err) + } + if got := len(res.Tools); got != numTools { + t.Errorf("ListTools() returned %d tools, want %d", got, numTools) + } + if callCtx.Err() == nil { + t.Fatal("the context was not cancelled after the response was delivered; the test did not exercise the race") + } + if slices.Contains(transport.writtenMethods(), notificationCancelled) { + t.Errorf("client sent %s for a request whose response it had already received", notificationCancelled) + } +} From f8ad6b67fa286fe4300026d4adfbef9886d07dbd Mon Sep 17 00:00:00 2001 From: Vitalii Cherepanov Date: Mon, 12 Oct 2026 18:24:24 +0200 Subject: [PATCH 2/2] mcp: drop the regression test for the delivered response --- mcp/transport_test.go | 121 ------------------------------------------ 1 file changed, 121 deletions(-) diff --git a/mcp/transport_test.go b/mcp/transport_test.go index 72b9b6bbf..dadf196b5 100644 --- a/mcp/transport_test.go +++ b/mcp/transport_test.go @@ -10,9 +10,7 @@ import ( "errors" "fmt" "io" - "slices" "strings" - "sync" "testing" "time" @@ -366,122 +364,3 @@ func TestIOConnFrameCap(t *testing.T) { }) } } - -// cancelAfterResponseTransport wraps a Transport so that a test can cancel a -// context at a point the client cannot observe from the outside: the first -// Read after the response to the given method has been delivered. The -// jsonrpc2 read loop dispatches a message before it reads the next one, so -// by then the call has been retired with its result while the caller may -// still be decoding it. The transport also records the methods written to -// it, so a test can check that no cancelled notification followed. -type cancelAfterResponseTransport struct { - Transport - method string - cancel context.CancelFunc - - mu sync.Mutex - requests map[int64]string // id -> method, for requests written by the client - written []string -} - -func (t *cancelAfterResponseTransport) Connect(ctx context.Context) (Connection, error) { - conn, err := t.Transport.Connect(ctx) - if err != nil { - return nil, err - } - t.requests = make(map[int64]string) - return &cancelAfterResponseConn{Connection: conn, t: t}, nil -} - -type cancelAfterResponseConn struct { - Connection - t *cancelAfterResponseTransport - - delivered bool // the response to t.method has been returned from Read - cancelled bool -} - -func (c *cancelAfterResponseConn) Read(ctx context.Context) (jsonrpc.Message, error) { - if c.delivered && !c.cancelled { - c.cancelled = true - c.t.cancel() - } - msg, err := c.Connection.Read(ctx) - if resp, ok := msg.(*jsonrpc.Response); ok { - if id, isInt := resp.ID.Raw().(int64); isInt { - c.t.mu.Lock() - c.delivered = c.t.requests[id] == c.t.method - c.t.mu.Unlock() - } - } - return msg, err -} - -func (c *cancelAfterResponseConn) Write(ctx context.Context, msg jsonrpc.Message) error { - if req, ok := msg.(*jsonrpc.Request); ok { - c.t.mu.Lock() - if id, isInt := req.ID.Raw().(int64); isInt { - c.t.requests[id] = req.Method - } - c.t.written = append(c.t.written, req.Method) - c.t.mu.Unlock() - } - return c.Connection.Write(ctx, msg) -} - -func (t *cancelAfterResponseTransport) writtenMethods() []string { - t.mu.Lock() - defer t.mu.Unlock() - return slices.Clone(t.written) -} - -// TestCallKeepsDeliveredResponseOnCancel verifies that a response that has -// already arrived is returned to the caller even if the caller's context is -// cancelled while the result is being decoded, and that no cancelled -// notification is sent for a request the client has the response to. The -// specification allows cancellation notifications only for requests that -// are believed to still be in progress. -func TestCallKeepsDeliveredResponseOnCancel(t *testing.T) { - // A large tool list keeps the client decoding the result for a while - // after the response has been retired, which is the window the test - // targets. - const numTools = 2000 - - ctx := context.Background() - server := NewServer(testImpl, &ServerOptions{PageSize: numTools}) - for i := range numTools { - server.AddTool(&Tool{Name: fmt.Sprintf("tool-%04d", i), InputSchema: map[string]any{"type": "object"}}, - func(context.Context, *CallToolRequest) (*CallToolResult, error) { - return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil - }) - } - ct, st := NewInMemoryTransports() - ss, err := server.Connect(ctx, st, nil) - if err != nil { - t.Fatalf("server connect: %v", err) - } - t.Cleanup(func() { _ = ss.Close() }) - - callCtx, cancel := context.WithCancel(ctx) - defer cancel() - transport := &cancelAfterResponseTransport{Transport: ct, method: methodListTools, cancel: cancel} - cs, err := NewClient(testImpl, nil).Connect(ctx, transport, nil) - if err != nil { - t.Fatalf("client connect: %v", err) - } - t.Cleanup(func() { _ = cs.Close() }) - - res, err := cs.ListTools(callCtx, nil) - if err != nil { - t.Fatalf("ListTools() = %v, want the delivered result", err) - } - if got := len(res.Tools); got != numTools { - t.Errorf("ListTools() returned %d tools, want %d", got, numTools) - } - if callCtx.Err() == nil { - t.Fatal("the context was not cancelled after the response was delivered; the test did not exercise the race") - } - if slices.Contains(transport.writtenMethods(), notificationCancelled) { - t.Errorf("client sent %s for a request whose response it had already received", notificationCancelled) - } -}