mcp: allow asynchronous cancel notification (#1151)

Fixes #1150

mcp/streamable_test.go

@@ -3599,6 +3599,79 @@ func TestStreamableStateful_RejectsNewProtocol_LegacyPlainText(t *testing.T) {
 	}
 }
 
+// TestCallCancellation_FastReturn verifies that a cancelled tool call
+// returns as soon as its context is cancelled, even when the peer is slow
+// to accept the follow-up notifications/cancelled POST.
+func TestCallCancellation_FastReturn(t *testing.T) {
+	// stallDuration is the time the middleware holds a
+	// notifications/cancelled POST. It must be larger than
+	// notifyCancellationTimeout so the blocking branch clearly overshoots
+	// even after cushioning for scheduler jitter.
+	const stallDuration = notifyCancellationTimeout + 2*time.Second
+
+	// callerDeadline is the deadline the client passes into CallTool. It
+	// must be short enough that any waiting on the cancel notification is
+	// obvious in the measured elapsed time.
+	const callerDeadline = 100 * time.Millisecond
+
+	// slack accounts for scheduling jitter around the deadline.
+	const slack = 500 * time.Millisecond
+
+	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)
+				}
+			}
+			h.ServeHTTP(w, r)
+		})
+	}
+
+	server := NewServer(testImpl, nil)
+	AddTool(server, &Tool{Name: "slow"},
+		func(ctx context.Context, req *CallToolRequest, args struct{}) (*CallToolResult, any, error) {
+			// Bound the handler above the caller deadline so the test
+			// always exercises the cancellation path, but not so long
+			// that a stuck cancellation makes the test slow.
+			select {
+			case <-ctx.Done():
+			case <-time.After(2 * stallDuration):
+			}
+			return &CallToolResult{Content: []Content{&TextContent{Text: "ok"}}}, nil, nil
+		})
+
+	handler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil)
+	httpServer := httptest.NewServer(stallCancelNotify(handler))
+	t.Cleanup(httpServer.Close)
+
+	client := NewClient(testImpl, nil)
+	cs, err := client.Connect(context.Background(),
+		&StreamableClientTransport{Endpoint: httpServer.URL}, nil)
+	if err != nil {
+		t.Fatalf("client connect: %v", err)
+	}
+	t.Cleanup(func() { cs.Close() })
+
+	callCtx, cancel := context.WithTimeout(context.Background(), callerDeadline)
+	defer cancel()
+
+	start := time.Now()
+	_, err = cs.CallTool(callCtx, &CallToolParams{Name: "slow"})
+	elapsed := time.Since(start)
+
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("CallTool error = %v, want context.DeadlineExceeded", err)
+	}
+	// The caller should return within slack of its own deadline.
+	if elapsed > callerDeadline+slack {
+		t.Errorf("CallTool returned after %v; want <= %v (caller deadline %v + slack %v)",
+			elapsed, callerDeadline+slack, callerDeadline, slack)
+	}
+}
+
 // TestStreamableStateless_AcceptsNewProtocol is the positive control:
 // confirms that a stateless server still accepts new-protocol requests
 // (the rejection in TestStreamableStateful_RejectsNewProtocol must not

mcp/transport.go

@@ -18,6 +18,7 @@ import (
 
 	internaljson "github.com/modelcontextprotocol/go-sdk/internal/json"
 	"github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2"
+	"github.com/modelcontextprotocol/go-sdk/internal/mcpgodebug"
 	"github.com/modelcontextprotocol/go-sdk/jsonrpc"
 )
 
@@ -28,6 +29,13 @@ import (
 // re-trigger expensive recovery on its behalf. See issue #882.
 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.
+var blockingcancelnotify = mcpgodebug.Value("blockingcancelnotify")
+
 // ErrConnectionClosed is returned when sending a message to a connection that
 // is closed or in the process of closing.
 var ErrConnectionClosed = errors.New("connection closed")
@@ -276,8 +284,28 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params
 	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:
-		err := cancelCall(ctx, conn, call)
-		return errors.Join(ctx.Err(), err)
+		// 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.
+		//
+		// Setting MCPGODEBUG=blockingcancelnotify=1 restores the previous
+		// behavior of waiting synchronously for delivery inside cancelCall.
+		if blockingcancelnotify == "1" {
+			err := cancelCall(ctx, conn, call)
+			return errors.Join(ctx.Err(), err)
+		}
+		conn.Retire(call, ctx.Err())
+		go func() {
+			notifyCtx, stop := context.WithTimeout(context.WithoutCancel(ctx), notifyCancellationTimeout)
+			defer stop()
+			_ = conn.Notify(notifyCtx, notificationCancelled, &CancelledParams{
+				Reason:    ctx.Err().Error(),
+				RequestID: call.ID().Raw(),
+			})
+		}()
+		return ctx.Err()
 	case err != nil:
 		return fmt.Errorf("calling %q: %w", method, err)
 	}