mcp: carry a cancelled notification's reason as the request's cancel cause (#1255)
The canceller read `CancelledParams.Reason` off the wire and dropped it: `Connection.Cancel` cancelled the request's context with a nil cause, so a handler that looked at `context.Cause` saw a bare `context.Canceled` and nothing of what the peer said. The specification asks implementations to log cancellation reasons, and the handler is the one place that can act on one; #1100 had already given the request context a cause, so the reason has a natural place to travel. `Connection.CancelCause` is `Cancel` with a cause. The canceller now cancels with an error that carries the peer's reason and unwraps to `context.Canceled`, so `errors.Is` keeps classifying the cancellation as one while `context.Cause` reads `request cancelled by the peer: <reason>`. It also logs the id and the reason at debug level through the logger the connection already carries, which is the SDK's own answer to the SHOULD. No exported API changes: the cause type stays unexported and is read through `context.Cause` and `Error()`. `TestCancellation` now asserts the cause the SDK's own client produces (its context's error, which the client sends as the reason), and `TestCancellationReason` writes the call and the cancelled notification to the connection directly with a reason the client would never send, on a session pinned to 2025-11-25 so the hand-written request needs no `_meta`. Both fail without the change. Fixes #1254 Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
This commit is contained in:
parent
e40f35d137
commit
3210c2b34c
3 changed files with 102 additions and 7 deletions
|
|
@ -453,12 +453,19 @@ func (ac *AsyncCall) Await(ctx context.Context, result any) error {
|
|||
// will not cause any messages that have not arrived yet with that ID to be
|
||||
// cancelled.
|
||||
func (c *Connection) Cancel(id ID) {
|
||||
c.CancelCause(id, nil)
|
||||
}
|
||||
|
||||
// CancelCause is like [Connection.Cancel], but records cause as the reason
|
||||
// the Context was cancelled, so that the Handle call can read it back through
|
||||
// [context.Cause]. A nil cause reads as [context.Canceled].
|
||||
func (c *Connection) CancelCause(id ID, cause error) {
|
||||
var req *incomingRequest
|
||||
c.updateInFlight(func(s *inFlightState) {
|
||||
req = s.incomingByID[id]
|
||||
})
|
||||
if req != nil {
|
||||
req.cancel(nil)
|
||||
req.cancel(cause)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -710,13 +710,13 @@ func TestCancellation(t *testing.T) {
|
|||
synctest.Test(t, func(t *testing.T) {
|
||||
var (
|
||||
start = make(chan struct{})
|
||||
cancelled = make(chan struct{}, 1) // don't block the request
|
||||
cancelled = make(chan error, 1) // don't block the request
|
||||
)
|
||||
slowTool := func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) {
|
||||
start <- struct{}{}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
cancelled <- struct{}{}
|
||||
cancelled <- context.Cause(ctx)
|
||||
case <-time.After(5 * time.Second):
|
||||
return nil, nil, nil
|
||||
}
|
||||
|
|
@ -732,7 +732,72 @@ func TestCancellation(t *testing.T) {
|
|||
<-start
|
||||
cancel()
|
||||
|
||||
<-cancelled
|
||||
// The client sends its context's error as the reason of its cancelled
|
||||
// notification, and the handler reads it back as the cause.
|
||||
cause := <-cancelled
|
||||
if !errors.Is(cause, context.Canceled) {
|
||||
t.Errorf("context.Cause = %v, want it to wrap context.Canceled", cause)
|
||||
}
|
||||
if want := "request cancelled by the peer: " + context.Canceled.Error(); cause == nil || cause.Error() != want {
|
||||
t.Errorf("context.Cause = %v, want %q", cause, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestCancellationReason verifies that the reason a peer gives in its
|
||||
// cancelled notification reaches the handler as the cause of its context,
|
||||
// which is what lets a server log it as the specification asks.
|
||||
//
|
||||
// The call and the notification are written to the connection directly, so
|
||||
// the reason is one the SDK's own client would never send; the session is
|
||||
// pinned to 2025-11-25 so that a request written that way needs no _meta.
|
||||
func TestCancellationReason(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
var (
|
||||
start = make(chan struct{})
|
||||
cancelled = make(chan error, 1) // don't block the request
|
||||
)
|
||||
slowTool := func(ctx context.Context, req *CallToolRequest, args any) (*CallToolResult, any, error) {
|
||||
start <- struct{}{}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
cancelled <- context.Cause(ctx)
|
||||
case <-time.After(5 * time.Second):
|
||||
cancelled <- nil
|
||||
}
|
||||
return nil, nil, nil
|
||||
}
|
||||
ctx := context.Background()
|
||||
ct, st := NewInMemoryTransports()
|
||||
s := NewServer(testImpl, nil)
|
||||
AddTool(s, &Tool{Name: "slow", InputSchema: &jsonschema.Schema{Type: "object"}}, slowTool)
|
||||
ss, err := s.Connect(ctx, st, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = ss.Close() })
|
||||
cs, err := NewClient(testImpl, nil).Connect(ctx, ct, &ClientSessionOptions{ProtocolVersion: protocolVersion20251125})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = cs.Close() })
|
||||
|
||||
call := cs.conn.Call(ctx, methodCallTool, &CallToolParams{Name: "slow"})
|
||||
<-start
|
||||
if err := cs.conn.Notify(ctx, notificationCancelled, &CancelledParams{RequestID: call.ID().Raw(), Reason: "user asked"}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
cause := <-cancelled
|
||||
if cause == nil {
|
||||
t.Fatal("the tool ran to completion, want it cancelled")
|
||||
}
|
||||
if !errors.Is(cause, context.Canceled) {
|
||||
t.Errorf("context.Cause = %v, want it to wrap context.Canceled", cause)
|
||||
}
|
||||
if got, want := cause.Error(), "request cancelled by the peer: user asked"; got != want {
|
||||
t.Errorf("context.Cause = %q, want %q", got, want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -207,7 +207,7 @@ func connect[H handler, State any](ctx context.Context, t Transport, b binder[H,
|
|||
reader, writer := jsonrpc2.Reader(mcpConn), jsonrpc2.Writer(mcpConn)
|
||||
var (
|
||||
h H
|
||||
preempter canceller
|
||||
preempter = canceller{logger: logger}
|
||||
)
|
||||
bind := func(conn *jsonrpc2.Connection) jsonrpc2.Handler {
|
||||
h = b.bind(mcpConn, conn, s, onClose)
|
||||
|
|
@ -252,7 +252,8 @@ type cancellationPropagator interface {
|
|||
// A canceller is a jsonrpc2.Preempter that cancels in-flight requests on MCP
|
||||
// cancelled notifications.
|
||||
type canceller struct {
|
||||
conn *jsonrpc2.Connection
|
||||
conn *jsonrpc2.Connection
|
||||
logger *slog.Logger
|
||||
}
|
||||
|
||||
// Preempt implements [jsonrpc2.Preempter].
|
||||
|
|
@ -266,11 +267,33 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result a
|
|||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
go c.conn.Cancel(id)
|
||||
// The spec says implementations should log cancellation reasons, and
|
||||
// the handler is the one place that can act on one, so the reason
|
||||
// travels as the cause of the request's context rather than being
|
||||
// dropped here.
|
||||
c.logger.Debug("request cancelled by the peer", "id", id.Raw(), "reason", params.Reason)
|
||||
go c.conn.CancelCause(id, &peerCancelledError{reason: params.Reason})
|
||||
}
|
||||
return nil, jsonrpc2.ErrNotHandled
|
||||
}
|
||||
|
||||
// A peerCancelledError is the cause a request's context carries once the peer
|
||||
// has sent a cancelled notification for it: [context.Cause] returns it to the
|
||||
// handler with the reason the peer gave, and it unwraps to [context.Canceled]
|
||||
// so that [errors.Is] keeps classifying the cancellation as one.
|
||||
type peerCancelledError struct {
|
||||
reason string
|
||||
}
|
||||
|
||||
func (e *peerCancelledError) Error() string {
|
||||
if e.reason == "" {
|
||||
return "request cancelled by the peer"
|
||||
}
|
||||
return "request cancelled by the peer: " + e.reason
|
||||
}
|
||||
|
||||
func (e *peerCancelledError) Unwrap() error { return context.Canceled }
|
||||
|
||||
// callSubscriptionsListen issues a "subscriptions/listen" call (SEP-2575)
|
||||
// without awaiting its JSON-RPC response. The call's logical lifetime is the
|
||||
// stream of notifications that follow on the same channel — the empty
|
||||
|
|
|
|||
Loading…
Reference in a new issue