From 6863d59cb9ed1df1135e78a48bb70753530247a4 Mon Sep 17 00:00:00 2001 From: Delaney Gillilan Date: Wed, 5 Aug 2026 07:14:23 -0700 Subject: [PATCH 1/4] mcp: support custom notifications Protocol extensions can define custom JSON-RPC notifications, but the SDK only exposes helpers for standard notifications. Add SendNotification to client and server sessions. Route custom notifications through sending middleware and preserve arbitrary parameters. Fixes modelcontextprotocol/go-sdk#745. --- mcp/client.go | 10 ++++++++++ mcp/mcp_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++ mcp/server.go | 10 ++++++++++ mcp/shared.go | 24 ++++++++++++++++++++++++ 4 files changed, 87 insertions(+) diff --git a/mcp/client.go b/mcp/client.go index 2ad0ea41..3235d048 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -1543,6 +1543,16 @@ func (cs *ClientSession) NotifyProgress(ctx context.Context, params *ProgressNot return handleNotify(ctx, notificationProgress, newClientRequest(cs, orZero[Params](params))) } +// SendNotification sends a custom notification to the server associated with +// this session. It supports protocol extensions such as notifications/foobar/stats. +func (cs *ClientSession) SendNotification(ctx context.Context, method string, params any) error { + return handleNotify( + ctx, + "x-notifications/"+method, + newClientRequest(cs, Params(&customNotificationParams{payload: params})), + ) +} + // Tools provides an iterator for all tools available on the server, // automatically fetching pages and managing cursors. // The params argument can set the initial cursor. diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index d9d9b3af..a3c280e4 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -839,6 +839,49 @@ func (b *safeBuffer) Bytes() []byte { return b.buf.Bytes() } +func TestSendNotification(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + ctx := context.Background() + ct, st := NewInMemoryTransports() + var clientLog, serverLog safeBuffer + + server := NewServer(testImpl, nil) + ss, err := server.Connect(ctx, &LoggingTransport{Transport: st, Writer: &serverLog}, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ss.Close() }) + + client := NewClient(testImpl, nil) + cs, err := client.Connect(ctx, &LoggingTransport{Transport: ct, Writer: &clientLog}, nil) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = cs.Close() }) + + if err := cs.SendNotification(ctx, "notifications/foobar/stats", map[string]any{"status": "ok"}); err != nil { + t.Fatal(err) + } + if err := ss.SendNotification(ctx, "notifications/foobar/stats", nil); err != nil { + t.Fatal(err) + } + synctest.Wait() + + for _, test := range []struct { + name string + log *safeBuffer + want string + }{ + {"client", &clientLog, `"method":"notifications/foobar/stats","params":{"status":"ok"}`}, + {"server", &serverLog, `"method":"notifications/foobar/stats","params":{}`}, + } { + if !bytes.Contains(test.log.Bytes(), []byte(test.want)) { + t.Errorf("%s log does not contain %q:\n%s", test.name, test.want, test.log.Bytes()) + } + } + }) +} + func TestNoJSONNull(t *testing.T) { ctx := context.Background() var ct, st Transport = NewInMemoryTransports() diff --git a/mcp/server.go b/mcp/server.go index c189a8ee..dcf56347 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1460,6 +1460,16 @@ func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNot return handleNotify(ctx, notificationProgress, newServerRequest(ss, orZero[Params](params))) } +// SendNotification sends a custom notification to the client associated with +// this session. It supports protocol extensions such as notifications/foobar/stats. +func (ss *ServerSession) SendNotification(ctx context.Context, method string, params any) error { + return handleNotify( + ctx, + "x-notifications/"+method, + newServerRequest(ss, Params(&customNotificationParams{payload: params})), + ) +} + // notifySubscriptionAcked sends a "notifications/subscriptions/acknowledged" // notification on the listen stream represented by this session, indicating // the subscription filter the server accepted (SEP-2575). diff --git a/mcp/shared.go b/mcp/shared.go index 5069a470..12c2ec52 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -117,6 +117,14 @@ func addMiddleware(handlerp *MethodHandler, middleware []Middleware) { } func defaultSendingMethodHandler(ctx context.Context, method string, req Request) (Result, error) { + if strings.HasPrefix(method, "x-notifications/") { + return nil, req.GetSession().getConn().Notify( + ctx, + strings.TrimPrefix(method, "x-notifications/"), + req.GetParams(), + ) + } + info, ok := req.GetSession().sendingMethodInfos()[method] if !ok { // This can be called from user code, with an arbitrary value for method. @@ -275,6 +283,22 @@ const ( missingParamsOK // params may be missing or null ) +type customNotificationParams struct { + payload any +} + +func (*customNotificationParams) GetMeta() map[string]any { return nil } +func (*customNotificationParams) SetMeta(map[string]any) {} +func (*customNotificationParams) isParams() {} +func (p *customNotificationParams) isNil() bool { return p == nil } + +func (p customNotificationParams) MarshalJSON() ([]byte, error) { + if p.payload == nil { + return []byte("{}"), nil + } + return json.Marshal(p.payload) +} + func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHandler[P, R], flags methodFlags) methodInfo { mi := newMethodInfo[P, R](flags) mi.newRequest = func(s Session, p Params, _ *RequestExtra) Request { From a5ff027e28009d85fa54e7f9b32e1da790fa9181 Mon Sep 17 00:00:00 2001 From: Delaney Gillilan Date: Thu, 6 Aug 2026 08:58:10 -0700 Subject: [PATCH 2/4] mcp: receive typed custom notifications --- mcp/client.go | 46 +++++++++++++++++++++++++++++++++++++++++++++- mcp/mcp_test.go | 23 +++++++++++++++++++++-- mcp/server.go | 17 +++++++++++++++++ mcp/shared.go | 33 +++++++++++++++++++++++++++------ 4 files changed, 110 insertions(+), 9 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 3235d048..33b8a483 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -39,6 +39,11 @@ type Client struct { // serverMethodInfos) plus any custom methods registered via // [AddSendingCustomMethod]. sendMethods map[string]methodInfo + // receiveMethods is the list of methods this client may receive from a + // server: it always contains the standard client methods (from + // clientMethodInfos) plus any custom notifications registered via + // [AddReceivingCustomNotification]. + receiveMethods map[string]methodInfo } // NewClient creates a new [Client]. @@ -67,6 +72,8 @@ func NewClient(impl *Implementation, options *ClientOptions) *Client { sendMethods := make(map[string]methodInfo, len(serverMethodInfos)) maps.Copy(sendMethods, serverMethodInfos) + receiveMethods := make(map[string]methodInfo, len(clientMethodInfos)) + maps.Copy(receiveMethods, clientMethodInfos) c := &Client{ impl: impl, @@ -75,6 +82,7 @@ func NewClient(impl *Implementation, options *ClientOptions) *Client { sendingMethodHandler_: defaultSendingMethodHandler, receivingMethodHandler_: defaultReceivingMethodHandler[*ClientSession], sendMethods: sendMethods, + receiveMethods: receiveMethods, } if opts.MultiRoundTrip == nil || !opts.MultiRoundTrip.Disabled { c.AddSendingMiddleware(clientMultiRoundTripMiddleware()) @@ -1177,7 +1185,9 @@ func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo { } func (cs *ClientSession) receivingMethodInfos() map[string]methodInfo { - return clientMethodInfos + cs.client.mu.Lock() + defer cs.client.mu.Unlock() + return cs.client.receiveMethods } func (cs *ClientSession) handle(ctx context.Context, req *jsonrpc.Request) (any, error) { @@ -1628,6 +1638,40 @@ func paginate[P listParams, R listResult[T], T any](ctx context.Context, params } } +// AddReceivingCustomNotification registers a handler for a custom JSON-RPC +// notification from a server. +// +// The method must start with "notifications/". Params are unmarshaled into P +// before handler is called. P must embed [ParamsBase]. +// +// Registration must occur before [Client.Connect]. Registering the same custom +// notification twice replaces the previous handler. +func AddReceivingCustomNotification[P paramsPtr[T], T any]( + c *Client, + method string, + handler func(context.Context, *ClientSession, P), +) error { + if !strings.HasPrefix(method, "notifications/") { + return fmt.Errorf("mcp: AddReceivingCustomNotification: %q is not a notification method", method) + } + if _, ok := clientMethodInfos[method]; ok { + return fmt.Errorf("mcp: AddReceivingCustomNotification: %q shadows a standard MCP notification", method) + } + if handler == nil { + return errors.New("mcp: AddReceivingCustomNotification: nil handler") + } + + typed := typedClientMethodHandler[P, *emptyResult](func(ctx context.Context, req *ClientRequest[P]) (*emptyResult, error) { + handler(ctx, req.Session, req.Params) + return nil, nil + }) + + c.mu.Lock() + defer c.mu.Unlock() + c.receiveMethods[method] = newClientMethodInfo(typed, notification) + return nil +} + // AddSendingCustomMethod registers a custom JSON-RPC method // that the client may send to the server. // diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index a3c280e4..10f2d67c 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -839,6 +839,11 @@ func (b *safeBuffer) Bytes() []byte { return b.buf.Bytes() } +type statsNotificationParams struct { + ParamsBase + Status string `json:"status"` +} + func TestSendNotification(t *testing.T) { synctest.Test(t, func(t *testing.T) { ctx := context.Background() @@ -853,6 +858,12 @@ func TestSendNotification(t *testing.T) { t.Cleanup(func() { _ = ss.Close() }) client := NewClient(testImpl, nil) + received := make(chan string, 1) + if err := AddReceivingCustomNotification(client, "notifications/foobar/stats", func(_ context.Context, _ *ClientSession, params *statsNotificationParams) { + received <- params.Status + }); err != nil { + t.Fatal(err) + } cs, err := client.Connect(ctx, &LoggingTransport{Transport: ct, Writer: &clientLog}, nil) if err != nil { t.Fatal(err) @@ -862,10 +873,18 @@ func TestSendNotification(t *testing.T) { if err := cs.SendNotification(ctx, "notifications/foobar/stats", map[string]any{"status": "ok"}); err != nil { t.Fatal(err) } - if err := ss.SendNotification(ctx, "notifications/foobar/stats", nil); err != nil { + if err := ss.SendNotification(ctx, "notifications/foobar/stats", map[string]any{"status": "ok"}); err != nil { t.Fatal(err) } synctest.Wait() + select { + case status := <-received: + if status != "ok" { + t.Errorf("received status %q, want ok", status) + } + default: + t.Error("custom notification handler was not called") + } for _, test := range []struct { name string @@ -873,7 +892,7 @@ func TestSendNotification(t *testing.T) { want string }{ {"client", &clientLog, `"method":"notifications/foobar/stats","params":{"status":"ok"}`}, - {"server", &serverLog, `"method":"notifications/foobar/stats","params":{}`}, + {"server", &serverLog, `"method":"notifications/foobar/stats","params":{"status":"ok"}`}, } { if !bytes.Contains(test.log.Bytes(), []byte(test.want)) { t.Errorf("%s log does not contain %q:\n%s", test.name, test.want, test.log.Bytes()) diff --git a/mcp/server.go b/mcp/server.go index dcf56347..0c6fa2d1 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1470,6 +1470,23 @@ func (ss *ServerSession) SendNotification(ctx context.Context, method string, pa ) } +// SendSubscriptionNotification sends a custom notification on the +// subscriptions/listen stream represented by ctx. It adds the subscription ID +// to the notification metadata. +func (ss *ServerSession) SendSubscriptionNotification(ctx context.Context, method string, params any) error { + requestID, ok := ctx.Value(idContextKey{}).(jsonrpc.ID) + if !ok || !requestID.IsValid() { + return fmt.Errorf("mcp: SendSubscriptionNotification: context has no subscription ID") + } + customParams := &customNotificationParams{payload: params} + injectMetaSubscriptionID(customParams, requestID) + return handleNotify( + ctx, + "x-notifications/"+method, + newServerRequest(ss, Params(customParams)), + ) +} + // notifySubscriptionAcked sends a "notifications/subscriptions/acknowledged" // notification on the listen stream represented by this session, indicating // the subscription filter the server accepted (SEP-2575). diff --git a/mcp/shared.go b/mcp/shared.go index 12c2ec52..cad09cd1 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -284,19 +284,40 @@ const ( ) type customNotificationParams struct { + meta map[string]any payload any } -func (*customNotificationParams) GetMeta() map[string]any { return nil } -func (*customNotificationParams) SetMeta(map[string]any) {} -func (*customNotificationParams) isParams() {} -func (p *customNotificationParams) isNil() bool { return p == nil } +func (p *customNotificationParams) GetMeta() map[string]any { return p.meta } +func (p *customNotificationParams) SetMeta(meta map[string]any) { + p.meta = meta +} +func (*customNotificationParams) isParams() {} +func (p *customNotificationParams) isNil() bool { return p == nil } func (p customNotificationParams) MarshalJSON() ([]byte, error) { - if p.payload == nil { + if p.payload == nil && p.meta == nil { return []byte("{}"), nil } - return json.Marshal(p.payload) + encoded, err := json.Marshal(p.payload) + if err != nil { + return nil, err + } + var object map[string]json.RawMessage + if err := json.Unmarshal(encoded, &object); err != nil { + return nil, fmt.Errorf("custom notification params must be an object: %w", err) + } + if object == nil { + object = map[string]json.RawMessage{} + } + if p.meta != nil { + encodedMeta, err := json.Marshal(p.meta) + if err != nil { + return nil, err + } + object["_meta"] = encodedMeta + } + return json.Marshal(object) } func newClientMethodInfo[P paramsPtr[T], R Result, T any](d typedClientMethodHandler[P, R], flags methodFlags) methodInfo { From 23b47e606d79dd531d1f40aa962d31368b6b27d3 Mon Sep 17 00:00:00 2001 From: Delaney Gillilan Date: Thu, 6 Aug 2026 08:59:07 -0700 Subject: [PATCH 3/4] mcp: test custom subscription notifications --- mcp/mcp_test.go | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 10f2d67c..93338619 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -858,9 +858,9 @@ func TestSendNotification(t *testing.T) { t.Cleanup(func() { _ = ss.Close() }) client := NewClient(testImpl, nil) - received := make(chan string, 1) + received := make(chan *statsNotificationParams, 1) if err := AddReceivingCustomNotification(client, "notifications/foobar/stats", func(_ context.Context, _ *ClientSession, params *statsNotificationParams) { - received <- params.Status + received <- params }); err != nil { t.Fatal(err) } @@ -873,14 +873,22 @@ func TestSendNotification(t *testing.T) { if err := cs.SendNotification(ctx, "notifications/foobar/stats", map[string]any{"status": "ok"}); err != nil { t.Fatal(err) } - if err := ss.SendNotification(ctx, "notifications/foobar/stats", map[string]any{"status": "ok"}); err != nil { + subscriptionID, err := jsonrpc.MakeID("subscription-1") + if err != nil { + t.Fatal(err) + } + subscriptionCtx := context.WithValue(ctx, idContextKey{}, subscriptionID) + if err := ss.SendSubscriptionNotification(subscriptionCtx, "notifications/foobar/stats", map[string]any{"status": "ok"}); err != nil { t.Fatal(err) } synctest.Wait() select { - case status := <-received: - if status != "ok" { - t.Errorf("received status %q, want ok", status) + case params := <-received: + if params.Status != "ok" { + t.Errorf("received status %q, want ok", params.Status) + } + if got := params.Meta[MetaKeySubscriptionID]; got != "subscription-1" { + t.Errorf("received subscription ID %v, want subscription-1", got) } default: t.Error("custom notification handler was not called") @@ -892,7 +900,7 @@ func TestSendNotification(t *testing.T) { want string }{ {"client", &clientLog, `"method":"notifications/foobar/stats","params":{"status":"ok"}`}, - {"server", &serverLog, `"method":"notifications/foobar/stats","params":{"status":"ok"}`}, + {"server", &serverLog, `"method":"notifications/foobar/stats","params":{"_meta":{"io.modelcontextprotocol/subscriptionId":"subscription-1"},"status":"ok"}`}, } { if !bytes.Contains(test.log.Bytes(), []byte(test.want)) { t.Errorf("%s log does not contain %q:\n%s", test.name, test.want, test.log.Bytes()) From dfe9d6b0360f872f7801a8db6aa3b0774f387c60 Mon Sep 17 00:00:00 2001 From: Delaney Gillilan Date: Thu, 6 Aug 2026 09:11:14 -0700 Subject: [PATCH 4/4] mcp: cancel HTTP listen streams with their carrier --- mcp/client.go | 3 +++ mcp/mcp_test.go | 42 +++++++++++++++++++----------------------- mcp/server.go | 5 ++++- mcp/shared.go | 3 ++- mcp/streamable.go | 2 ++ mcp/transport.go | 22 ++++++++++++++++++---- 6 files changed, 48 insertions(+), 29 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 33b8a483..d8cdd25b 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -1212,6 +1212,9 @@ func (cs *ClientSession) receivingMethodHandler() MethodHandler { // getConn implements [Session.getConn]. func (cs *ClientSession) getConn() *jsonrpc2.Connection { return cs.conn } +// getMCPConn implements [Session.getMCPConn]. +func (cs *ClientSession) getMCPConn() Connection { return cs.mcpConn } + func (*ClientSession) ping(context.Context, *PingParams) (*emptyResult, error) { return &emptyResult{}, nil } diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 93338619..d0b930c6 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -2810,22 +2810,8 @@ type resourceSubEvent struct { id string // _meta subscription ID, stringified } -// TestResourceSubscriptionsSEP2575_Streamable verifies the Subscribe -> -// ResourceUpdated path on a stateless Streamable HTTP server. -// -// Caveat: per-subscription Unsubscribe is intentionally NOT verified here. -// In stateless Streamable HTTP mode the subscriptions/listen handler blocks -// on its request context, and neither the HTTP POST disconnect nor the -// separate notifications/cancelled POST currently propagates to that -// handler's context. The handler only unwinds when the server next attempts -// a write to the (now-dead) SSE stream and the writeErr branch in the -// jsonrpc2 layer cancels the in-flight request. To keep the test -// hermetic we therefore trigger a write at the end by adding a resource, -// which fires notifications/resources/list_changed on the auto-listen path -// (if any) and on the per-URI listen, causing the listen handler to unwind. -// The spec-correct fix is to plumb the POST's request context down to the -// subscriptionsListen handler so HTTP disconnect is observed directly; this -// is tracked separately. +// TestResourceSubscriptions_Streamable verifies resource subscription +// delivery and cancellation on a stateless Streamable HTTP server. func TestResourceSubscriptions_Streamable(t *testing.T) { subCh := make(chan string, 8) @@ -2882,13 +2868,23 @@ func TestResourceSubscriptions_Streamable(t *testing.T) { t.Fatal("timed out waiting for resource update") } - // See test header comment for the explanation of this teardown ritual: - // close the client, then drop server-side TCP, then drive a write that - // will fail (any extra ResourceUpdated for our URI), to unblock the - // in-flight listen handler so httpServer.Close can return. - _ = cs.Close() - httpServer.CloseClientConnections() - server.ResourceUpdated(ctx, &ResourceUpdatedNotificationParams{URI: "file:///r1"}) + if err := cs.Unsubscribe(ctx, &UnsubscribeParams{URI: "file:///r1"}); err != nil { + t.Fatalf("unsubscribe r1: %v", err) + } + select { + case got := <-unsubCh: + if got != "file:///r1" { + t.Fatalf("got URI %q, want %q", got, "file:///r1") + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for UnsubscribeHandler") + } + if _, err := cs.ListTools(ctx, nil); err != nil { + t.Fatalf("list tools after unsubscribe: %v", err) + } + if err := cs.Close(); err != nil { + t.Fatalf("close client: %v", err) + } httpServer.Close() } diff --git a/mcp/server.go b/mcp/server.go index 0c6fa2d1..6f77d90a 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1880,9 +1880,12 @@ func (ss *ServerSession) receivingMethodHandler() MethodHandler { return s.receivingMethodHandler_ } -// getConn implements [session.getConn]. +// getConn implements [Session.getConn]. func (ss *ServerSession) getConn() *jsonrpc2.Connection { return ss.conn } +// getMCPConn implements [Session.getMCPConn]. +func (ss *ServerSession) getMCPConn() Connection { return ss.mcpConn } + // handle invokes the method described by the given JSON RPC request. func (ss *ServerSession) handle(ctx context.Context, req *jsonrpc.Request) (any, error) { ss.mu.Lock() diff --git a/mcp/shared.go b/mcp/shared.go index cad09cd1..8fcb9f39 100644 --- a/mcp/shared.go +++ b/mcp/shared.go @@ -104,6 +104,7 @@ type Session interface { sendingMethodHandler() MethodHandler receivingMethodHandler() MethodHandler getConn() *jsonrpc2.Connection + getMCPConn() Connection } // Middleware is a function from [MethodHandler] to [MethodHandler]. @@ -146,7 +147,7 @@ func defaultSendingMethodHandler(ctx context.Context, method string, req Request // The concrete type of the result is the return type of the receiving function. res := info.newResult() if method == methodSubscriptionsListen { - callSubscriptionsListen(ctx, req.GetSession().getConn(), method, params) + callSubscriptionsListen(ctx, req.GetSession().getConn(), req.GetSession().getMCPConn(), method, params) } else { if err := call(ctx, req.GetSession().getConn(), method, params, res); err != nil { return nil, err diff --git a/mcp/streamable.go b/mcp/streamable.go index db81a37b..4a89b14d 100644 --- a/mcp/streamable.go +++ b/mcp/streamable.go @@ -2125,6 +2125,8 @@ type streamableClientConn struct { var _ clientConnection = (*streamableClientConn)(nil) +func (*streamableClientConn) cancelsListenWithContext() bool { return true } + func (c *streamableClientConn) sessionUpdated(state clientSessionState) { c.mu.Lock() c.initializedResult = state.InitializeResult diff --git a/mcp/transport.go b/mcp/transport.go index 55c73d74..9d02a012 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -254,18 +254,27 @@ func (c *canceller) Preempt(ctx context.Context, req *jsonrpc.Request) (result a // response, if ever delivered, only marks subscription teardown — so the // caller has nothing useful to block on. // -// Cancellation is driven by ctx: when it is cancelled, a background goroutine -// sends a "notifications/cancelled" notification referencing the listen's -// request ID and retires the call from the connection's outgoing-calls map. -func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params) { +// Cancellation is driven by ctx. A carrier-bound transport closes its request. +// Other transports send "notifications/cancelled" with the listen request ID. +func callSubscriptionsListen(ctx context.Context, conn *jsonrpc2.Connection, mcpConn Connection, method string, params Params) { call := conn.Call(ctx, method, params) go func() { <-ctx.Done() + if carrier, ok := mcpConn.(listenContextCanceller); ok && carrier.cancelsListenWithContext() { + conn.Retire(call, ctx.Err()) + return + } _ = cancelCall(ctx, conn, call) }() } +// listenContextCanceller identifies transports that cancel a listen request by +// closing its carrier when the request context ends. +type listenContextCanceller interface { + cancelsListenWithContext() bool +} + // call executes and awaits a jsonrpc2 call on the given connection, // translating errors into the mcp domain. func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params Params, result Result) error { @@ -334,6 +343,11 @@ type loggingConn struct { func (c *loggingConn) SessionID() string { return c.delegate.SessionID() } +func (c *loggingConn) cancelsListenWithContext() bool { + carrier, ok := c.delegate.(listenContextCanceller) + return ok && carrier.cancelsListenWithContext() +} + // Read is a stream middleware that logs incoming messages. func (s *loggingConn) Read(ctx context.Context) (jsonrpc.Message, error) { msg, err := s.delegate.Read(ctx)