mcp: support custom notifications

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.
 //

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()
 }
 

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()

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 {

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