diff --git a/docs/mcpgodebug.md b/docs/mcpgodebug.md index 39a76a8b..d25b9b49 100644 --- a/docs/mcpgodebug.md +++ b/docs/mcpgodebug.md @@ -34,6 +34,16 @@ Options listed below were added and will be removed in the 1.9.0 version of the existing renegotiation logic recover and prevents the failure from tearing down the underlying connection. +- `blockingcancelnotify` added. If set to `1`, a cancelled call waits + synchronously for the best-effort `notifications/cancelled` message to be + 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. + ### 1.7.0 Options listed below were added and will be removed in the 1.9.0 version of the SDK. diff --git a/internal/docs/mcpgodebug.src.md b/internal/docs/mcpgodebug.src.md index efb02da2..32f8af26 100644 --- a/internal/docs/mcpgodebug.src.md +++ b/internal/docs/mcpgodebug.src.md @@ -33,6 +33,16 @@ Options listed below were added and will be removed in the 1.9.0 version of the existing renegotiation logic recover and prevents the failure from tearing down the underlying connection. +- `blockingcancelnotify` added. If set to `1`, a cancelled call waits + synchronously for the best-effort `notifications/cancelled` message to be + 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. + ### 1.7.0 Options listed below were added and will be removed in the 1.9.0 version of the SDK. diff --git a/mcp/streamable_test.go b/mcp/streamable_test.go index 18d6b5fd..504cda66 100644 --- a/mcp/streamable_test.go +++ b/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 diff --git a/mcp/transport.go b/mcp/transport.go index 55c73d74..d72f9d1b 100644 --- a/mcp/transport.go +++ b/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) }