WIP: mcp: support custom notifications

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.

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

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

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 {