diff --git a/mcp/client.go b/mcp/client.go index 2ad0ea41..d8cdd25b 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) { @@ -1202,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 } @@ -1543,6 +1556,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. @@ -1618,6 +1641,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 d9d9b3af..d0b930c6 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -839,6 +839,76 @@ 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() + 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) + received := make(chan *statsNotificationParams, 1) + if err := AddReceivingCustomNotification(client, "notifications/foobar/stats", func(_ context.Context, _ *ClientSession, params *statsNotificationParams) { + received <- params + }); err != nil { + t.Fatal(err) + } + 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) + } + 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 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") + } + + 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":{"_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()) + } + } + }) +} + func TestNoJSONNull(t *testing.T) { ctx := context.Background() var ct, st Transport = NewInMemoryTransports() @@ -2740,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) @@ -2812,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 c189a8ee..6f77d90a 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -1460,6 +1460,33 @@ 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})), + ) +} + +// 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). @@ -1853,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 5069a470..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]. @@ -117,6 +118,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. @@ -138,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 @@ -275,6 +284,43 @@ const ( missingParamsOK // params may be missing or null ) +type customNotificationParams struct { + meta map[string]any + payload any +} + +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 && p.meta == nil { + return []byte("{}"), nil + } + 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 { mi := newMethodInfo[P, R](flags) mi.newRequest = func(s Session, p Params, _ *RequestExtra) Request { 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 d72f9d1b..c8cccd8c 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -262,18 +262,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 { @@ -362,6 +371,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)