Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions docs/mcpgodebug.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
10 changes: 10 additions & 0 deletions internal/docs/mcpgodebug.src.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
73 changes: 73 additions & 0 deletions mcp/streamable_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
32 changes: 30 additions & 2 deletions mcp/transport.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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")
Expand Down Expand Up @@ -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)
}
Expand Down
Loading