From a2497338ebbf57b31424c5a58bcfdfad9bc05f2c Mon Sep 17 00:00:00 2001 From: Vishal Punjabi <10079537+vshalpnjabi@users.noreply.github.com> Date: Fri, 6 Nov 2026 08:29:19 +0000 Subject: [PATCH] mcp: deliver notifications/cancelled when closing after a cancelled call Since #1151, call() retires a cancelled request and sends notifications/cancelled from a goroutine. Closing the session immediately after the call returns marks the jsonrpc2 connection shutting down with no remaining outgoing calls, so Notify rejects the write. On a stateful streamable HTTP server the tool handler then keeps running and ClientSession.Close blocks for closeDeleteTimeout. Reserve the outgoing notification before returning from call() so Close waits for the write and Notify is allowed during shutdown. The caller still returns immediately, and the call is still retired immediately. For #1505 Co-authored-by: Vishal Punjabi --- docs/mcpgodebug.md | 9 +++-- internal/docs/mcpgodebug.src.md | 9 +++-- internal/jsonrpc2/conn.go | 36 ++++++++++++++++-- internal/jsonrpc2/conn_test.go | 59 ++++++++++++++++++++++++++++++ mcp/streamable_test.go | 65 ++++++++++++++++++++++++++++++++- mcp/transport.go | 20 ++++++---- 6 files changed, 177 insertions(+), 21 deletions(-) diff --git a/docs/mcpgodebug.md b/docs/mcpgodebug.md index 47563523..83567025 100644 --- a/docs/mcpgodebug.md +++ b/docs/mcpgodebug.md @@ -100,10 +100,11 @@ Options listed below were added and will be removed in the 1.9.0 version of the delivered (up to `notifyCancellationTimeout`, currently 5 s) before returning to the caller, restoring the previous behavior. The delivery error is joined into the caller's returned error. The default behavior was - changed so that the call is retired immediately and the notification is - sent asynchronously off the caller's return path: the caller returns as - soon as its context is cancelled and cannot be delayed by a slow or - unresponsive peer. See issue #1150. + changed so that the notification is sent asynchronously off the caller's + return path: the caller returns as soon as its context is cancelled and + cannot be delayed by a slow or unresponsive peer. The pending write is + reserved before return so a subsequent Close still delivers it. See + issues #1150 and #1505. Options below were removed, according to plan: diff --git a/internal/docs/mcpgodebug.src.md b/internal/docs/mcpgodebug.src.md index 66f36ff2..3d2404aa 100644 --- a/internal/docs/mcpgodebug.src.md +++ b/internal/docs/mcpgodebug.src.md @@ -99,10 +99,11 @@ Options listed below were added and will be removed in the 1.9.0 version of the delivered (up to `notifyCancellationTimeout`, currently 5 s) before returning to the caller, restoring the previous behavior. The delivery error is joined into the caller's returned error. The default behavior was - changed so that the call is retired immediately and the notification is - sent asynchronously off the caller's return path: the caller returns as - soon as its context is cancelled and cannot be delayed by a slow or - unresponsive peer. See issue #1150. + changed so that the notification is sent asynchronously off the caller's + return path: the caller returns as soon as its context is cancelled and + cannot be delayed by a slow or unresponsive peer. The pending write is + reserved before return so a subsequent Close still delivers it. See + issues #1150 and #1505. Options below were removed, according to plan: diff --git a/internal/jsonrpc2/conn.go b/internal/jsonrpc2/conn.go index 22e6faee..d72c02ba 100644 --- a/internal/jsonrpc2/conn.go +++ b/internal/jsonrpc2/conn.go @@ -293,10 +293,12 @@ func (c *Connection) Notify(ctx context.Context, method string, params any) (err c.updateInFlight(func(s *inFlightState) { // If the connection is shutting down, allow outgoing notifications only if - // there is at least one call still in flight. The number of calls in flight - // cannot increase once shutdown begins, and allowing outgoing notifications - // may permit notifications that will cancel in-flight calls. - if len(s.outgoingCalls) == 0 && len(s.incomingByID) == 0 { + // there is at least one call still in flight, or a notification already + // reserved with [Connection.BeginOutgoingNotification]. The number of + // calls in flight cannot increase once shutdown begins, and allowing + // outgoing notifications may permit notifications that will cancel + // in-flight calls. + if len(s.outgoingCalls) == 0 && len(s.incomingByID) == 0 && s.outgoingNotifications == 0 { err = s.shuttingDown(ErrConnClosing) if err != nil { return @@ -317,6 +319,32 @@ func (c *Connection) Notify(ctx context.Context, method string, params any) (err return c.write(ctx, notify) } +// BeginOutgoingNotification records that a notification will be sent +// asynchronously. Close waits until the returned function is called, and +// [Connection.Notify] is allowed to write during shutdown while the +// reservation is held. If the connection is already done, the reservation is +// a no-op. +func (c *Connection) BeginOutgoingNotification() (end func()) { + accepted := false + c.updateInFlight(func(s *inFlightState) { + select { + case <-c.done: + return + default: + } + s.outgoingNotifications++ + accepted = true + }) + if !accepted { + return func() {} + } + return sync.OnceFunc(func() { + c.updateInFlight(func(s *inFlightState) { + s.outgoingNotifications-- + }) + }) +} + // Call invokes the target method and returns an object that can be used to await the response. // The params will be marshaled to JSON before sending over the wire, and will // be handed to the method invoked. diff --git a/internal/jsonrpc2/conn_test.go b/internal/jsonrpc2/conn_test.go index 0893953c..c67a1710 100644 --- a/internal/jsonrpc2/conn_test.go +++ b/internal/jsonrpc2/conn_test.go @@ -14,6 +14,7 @@ import ( "strings" "sync" "testing" + "time" ) type fakeReader struct { @@ -259,3 +260,61 @@ func (w *dropWriter) drops() []ID { defer w.mu.Unlock() return slices.Clone(w.dropped) } + +// TestBeginOutgoingNotificationAllowsNotifyDuringClose verifies that a +// notification reserved with BeginOutgoingNotification is still written when +// Close races with the sender, which is how MCP delivers notifications/cancelled +// after a cancelled call (issue #1505). +func TestBeginOutgoingNotificationAllowsNotifyDuringClose(t *testing.T) { + in := &chanReader{msgs: make(chan Message)} + w := fakeWriter{out: make(chan Message, 2)} + conn := NewConnection(context.Background(), ConnectionConfig{ + Reader: in, + Writer: w, + Closer: in, + Bind: func(*Connection) Handler { + return HandlerFunc(func(context.Context, *Request) (any, error) { + return nil, ErrNotHandled + }) + }, + }) + + callCtx, cancel := context.WithCancel(context.Background()) + call := conn.Call(callCtx, "slow", nil) + got := <-w.out + if req, ok := got.(*Request); !ok || !req.IsCall() { + t.Fatalf("got %#v, want the outgoing call", got) + } + cancel() + if err := call.Await(callCtx, nil); !errors.Is(err, context.Canceled) { + t.Fatalf("Await: %v, want context.Canceled", err) + } + + end := conn.BeginOutgoingNotification() + conn.Retire(call, callCtx.Err()) + go func() { + defer end() + _ = conn.Notify(context.Background(), "notifications/cancelled", nil) + }() + + closed := make(chan error, 1) + go func() { closed <- conn.Close() }() + + select { + case msg := <-w.out: + req, ok := msg.(*Request) + if !ok || req.IsCall() || req.Method != "notifications/cancelled" { + t.Fatalf("got %#v, want notifications/cancelled", msg) + } + case <-time.After(2 * time.Second): + t.Fatal("notification was not written") + } + select { + case err := <-closed: + if err != nil { + t.Errorf("Close: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("Close did not return") + } +} diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 013c77d7..8b13f3f4 100644 --- a/mcp/streamable_test.go +++ b/mcp/streamable_test.go @@ -4032,13 +4032,17 @@ func TestCallCancellation_FastReturn(t *testing.T) { // slack accounts for scheduling jitter around the deadline. const slack = 500 * time.Millisecond + stopStall := make(chan struct{}) stallCancelNotify := func(h http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodPost && r.Body != nil { body, _ := io.ReadAll(r.Body) r.Body = io.NopCloser(bytes.NewReader(body)) if bytes.Contains(body, []byte(notificationCancelled)) { - time.Sleep(stallDuration) + select { + case <-time.After(stallDuration): + case <-stopStall: + } } } h.ServeHTTP(w, r) @@ -4068,7 +4072,10 @@ func TestCallCancellation_FastReturn(t *testing.T) { if err != nil { t.Fatalf("client connect: %v", err) } - t.Cleanup(func() { cs.Close() }) + t.Cleanup(func() { + close(stopStall) + cs.Close() + }) callCtx, cancel := context.WithTimeout(context.Background(), callerDeadline) defer cancel() @@ -4087,6 +4094,60 @@ func TestCallCancellation_FastReturn(t *testing.T) { } } +// TestCloseAfterCanceledCallCancelsTheServerHandler verifies that closing the +// session immediately after a cancelled tools/call still delivers +// notifications/cancelled. Otherwise a stateful streamable handler keeps +// running and ClientSession.Close blocks for closeDeleteTimeout. See #1505. +func TestCloseAfterCanceledCallCancelsTheServerHandler(t *testing.T) { + entered := make(chan struct{}, 1) + canceled := make(chan struct{}, 1) + release := make(chan struct{}) + + server := NewServer(testImpl, nil) + AddTool(server, &Tool{Name: "wait"}, + func(ctx context.Context, _ *CallToolRequest, _ struct{}) (*CallToolResult, any, error) { + entered <- struct{}{} + select { + case <-ctx.Done(): + canceled <- struct{}{} + return nil, nil, ctx.Err() + case <-release: + return &CallToolResult{}, nil, nil + } + }) + httpServer := httptest.NewServer(NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)) + t.Cleanup(httpServer.Close) + t.Cleanup(func() { close(release) }) + + client := NewClient(testImpl, nil) + session, err := client.Connect(context.Background(), &StreamableClientTransport{Endpoint: httpServer.URL}, nil) + if err != nil { + t.Fatal(err) + } + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + <-entered + cancel() + }() + if _, err := session.CallTool(ctx, &CallToolParams{Name: "wait"}); !errors.Is(err, context.Canceled) { + t.Fatalf("CallTool: got %v, want context.Canceled", err) + } + + start := time.Now() + _ = session.Close() + closed := time.Since(start) + + select { + case <-canceled: + if closed >= closeDeleteTimeout { + t.Errorf("Close took %v; want less than closeDeleteTimeout (%v)", closed, closeDeleteTimeout) + } + case <-time.After(2 * time.Second): + t.Fatalf("server handler never saw the cancellation; Close took %v", closed.Round(time.Millisecond)) + } +} + // TestStreamableStateless_AcceptsNewProtocol is the positive control: // confirms that a stateless server still accepts new-protocol requests // (the rejection in TestStreamableStateful_RejectsNewProtocol must not diff --git a/mcp/transport.go b/mcp/transport.go index 724f134b..9c95b1d1 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -32,8 +32,8 @@ const notifyCancellationTimeout = 5 * time.Second // blockingcancelnotify, when set to "1" via MCPGODEBUG, restores the previous // behavior of blocking the caller's return on delivery of the best-effort // notifications/cancelled message (up to notifyCancellationTimeout). By -// default, the call is retired immediately and the notification is sent -// asynchronously so the caller returns as soon as its context is cancelled. +// default, the notification is sent asynchronously so the caller returns as +// soon as its context is cancelled. var blockingcancelnotify = mcpgodebug.Value("blockingcancelnotify") // ErrConnectionClosed is returned when sending a message to a connection that @@ -319,11 +319,15 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params } return fmt.Errorf("%w: calling %q", ErrConnectionClosed, method) case ctx.Err() != nil: - // The notifications/cancelled message is best-effort. Retire the call - // immediately (so an unresponsive peer cannot delay the eager - // retirement that cancelCall's docstring promises) and send the - // notification off the caller's return path so a slow or unresponsive - // peer cannot delay the caller past its own deadline. See issue #1150. + // The notifications/cancelled message is best-effort. Send it off the + // caller's return path so a slow or unresponsive peer cannot delay + // the caller past its own deadline (issue #1150). + // + // Reserve the write before returning so Close, which follows a + // cancelled call on graceful shutdown, waits for it and does not + // reject Notify. Retire immediately so an unresponsive peer's missing + // response cannot delay Close beyond notifyCancellationTimeout. + // See issue #1505. // // Setting MCPGODEBUG=blockingcancelnotify=1 restores the previous // behavior of waiting synchronously for delivery inside cancelCall. @@ -331,8 +335,10 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params err := cancelCall(ctx, conn, call) return errors.Join(ctx.Err(), err) } + endNotify := conn.BeginOutgoingNotification() conn.Retire(call, ctx.Err()) go func() { + defer endNotify() notifyCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout) defer stop() _ = conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{