From ce5b6cd9269bd7709a8b14c2128c23902f52d221 Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Sat, 17 May 2025 12:59:45 -0400 Subject: [PATCH] internal/mcp: add notifications Add notification logic for intialization and list changes. Change-Id: I1d392bb49d5995f2046c400da955a033ab2715ce Reviewed-on: https://go-review.googlesource.com/c/tools/+/673775 Reviewed-by: Sam Thanawalla LUCI-TryBot-Result: Go LUCI Reviewed-by: Robert Findley --- mcp/client.go | 59 ++++++++++++++++++---- mcp/generate.go | 24 +++++++-- mcp/mcp_test.go | 49 +++++++++++++++++- mcp/protocol.go | 68 +++++++++++++++++-------- mcp/server.go | 130 +++++++++++++++++++++++++++++++----------------- mcp/shared.go | 27 ++++++++++ 6 files changed, 276 insertions(+), 81 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 0367758..f793b99 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -49,6 +49,10 @@ type ClientOptions struct { // Handler for sampling. // Called when a server calls CreateMessage. CreateMessageHandler func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) + // Handlers for notifications from the server. + ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams) + PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams) + ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams) } // bind implements the binder[*ClientSession] interface, so that Clients can @@ -86,10 +90,13 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e if err != nil { return nil, err } + caps := &ClientCapabilities{} + caps.Roots.ListChanged = true if c.opts.CreateMessageHandler != nil { caps.Sampling = &SamplingCapabilities{} } + params := &InitializeParams{ ClientInfo: &implementation{Name: c.name, Version: c.version}, Capabilities: caps, @@ -133,11 +140,13 @@ func (c *ClientSession) Wait() error { // AddRoots adds the given roots to the client, // replacing any with the same URIs, // and notifies any connected servers. -// TODO: notification func (c *Client) AddRoots(roots ...*Root) { - c.mu.Lock() - defer c.mu.Unlock() - c.roots.add(roots...) + // Only notify if something could change. + if len(roots) == 0 { + return + } + c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{}, + func() bool { c.roots.add(roots...); return true }) } // RemoveRoots removes the roots with the given URIs, @@ -145,9 +154,22 @@ func (c *Client) AddRoots(roots ...*Root) { // It is not an error to remove a nonexistent root. // TODO: notification func (c *Client) RemoveRoots(uris ...string) { + c.changeAndNotify(notificationRootsListChanged, &RootsListChangedParams{}, + func() bool { return c.roots.remove(uris...) }) +} + +// changeAndNotify is called when a feature is added or removed. +// It calls change, which should do the work and report whether a change actually occurred. +// If there was a change, it notifies a snapshot of the sessions. +func (c *Client) changeAndNotify(notification string, params any, change func() bool) { + var sessions []*ClientSession + // Lock for the change, but not for the notification. c.mu.Lock() - defer c.mu.Unlock() - c.roots.remove(uris...) + if change() { + sessions = slices.Clone(c.sessions) + } + c.mu.Unlock() + notifySessions(sessions, notification, params) } func (c *Client) listRoots(_ context.Context, _ *ClientSession, _ *ListRootsParams) (*ListRootsResult, error) { @@ -180,10 +202,12 @@ func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) { // clientMethodInfos maps from the RPC method name to serverMethodInfos. var clientMethodInfos = map[string]methodInfo[ClientSession]{ - methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)), - methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)), - methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)), - // TODO: notifications + methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)), + methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)), + methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)), + notificationToolListChanged: newMethodInfo(clientMethod((*Client).callToolChangedHandler)), + notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler)), + notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler)), } var _ session[ClientSession] = (*ClientSession)(nil) @@ -202,6 +226,9 @@ func (cs *ClientSession) methodHandler() MethodHandler[ClientSession] { return cs.client.methodHandler_ } +// getConn implements [session.getConn]. +func (cs *ClientSession) getConn() *jsonrpc2.Connection { return cs.conn } + func (c *ClientSession) ping(ct context.Context, params *PingParams) (struct{}, error) { return struct{}{}, nil } @@ -263,6 +290,18 @@ func (c *ClientSession) ReadResource(ctx context.Context, params *ReadResourcePa return standardCall[ReadResourceResult](ctx, c.conn, methodReadResource, params) } +func (c *Client) callToolChangedHandler(ctx context.Context, s *ClientSession, params *ToolListChangedParams) (any, error) { + return callNotificationHandler(ctx, c.opts.ToolListChangedHandler, s, params) +} + +func (c *Client) callPromptChangedHandler(ctx context.Context, s *ClientSession, params *PromptListChangedParams) (any, error) { + return callNotificationHandler(ctx, c.opts.PromptListChangedHandler, s, params) +} + +func (c *Client) callResourceChangedHandler(ctx context.Context, s *ClientSession, params *ResourceListChangedParams) (any, error) { + return callNotificationHandler(ctx, c.opts.ResourceListChangedHandler, s, params) +} + func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) { var result TRes if err := call(ctx, conn, method, params, &result); err != nil { diff --git a/mcp/generate.go b/mcp/generate.go index 12998d1..a2de04d 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -120,8 +120,12 @@ var declarations = config{ "Prompt": {}, "PromptMessage": {}, "PromptArgument": {}, - "ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string - "RequestId": {Name: "-", Substitute: "any"}, // null|number|string + "PromptListChangedNotification": { + Name: "-", + Fields: config{"Params": {Name: "PromptListChangedParams"}}, + }, + "ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string + "RequestId": {Name: "-", Substitute: "any"}, // null|number|string "ReadResourceRequest": { Name: "-", Fields: config{"Params": {Name: "ReadResourceParams"}}, @@ -130,8 +134,16 @@ var declarations = config{ Fields: config{"Contents": {Substitute: "*ResourceContents"}}, }, "Resource": {}, - "Role": {}, - "Root": {}, + "ResourceListChangedNotification": { + Name: "-", + Fields: config{"Params": {Name: "ResourceListChangedParams"}}, + }, + "Role": {}, + "Root": {}, + "RootsListChangedNotification": { + Name: "-", + Fields: config{"Params": {Name: "RootsListChangedParams"}}, + }, "SamplingCapabilities": {Substitute: "struct{}"}, "SamplingMessage": {}, @@ -147,6 +159,10 @@ var declarations = config{ Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}}, }, "ToolAnnotations": {}, + "ToolListChangedNotification": { + Name: "-", + Fields: config{"Params": {Name: "ToolListChangedParams"}}, + }, } func main() { diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 0165b98..3dc74c5 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -39,7 +39,26 @@ func TestEndToEnd(t *testing.T) { ctx := context.Background() ct, st := NewInMemoryTransports() - s := NewServer("testServer", "v1.0.0", nil) + // Channels to check if notification callbacks happened. + notificationChans := map[string]chan int{} + for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources"} { + notificationChans[name] = make(chan int, 1) + } + waitForNotification := func(t *testing.T, name string) { + t.Helper() + select { + case <-notificationChans[name]: + case <-time.After(time.Second): + t.Fatalf("%s handler never called", name) + } + } + + sopts := &ServerOptions{ + InitializedHandler: func(context.Context, *ServerSession, *InitializedParams) { notificationChans["initialized"] <- 0 }, + RootsListChangedHandler: func(context.Context, *ServerSession, *RootsListChangedParams) { notificationChans["roots"] <- 0 }, + } + + s := NewServer("testServer", "v1.0.0", sopts) // The 'greet' tool says hi. s.AddTools(NewTool("greet", "say hi", sayHi)) @@ -89,6 +108,9 @@ func TestEndToEnd(t *testing.T) { CreateMessageHandler: func(context.Context, *ClientSession, *CreateMessageParams) (*CreateMessageResult, error) { return &CreateMessageResult{Model: "aModel"}, nil }, + ToolListChangedHandler: func(context.Context, *ClientSession, *ToolListChangedParams) { notificationChans["tools"] <- 0 }, + PromptListChangedHandler: func(context.Context, *ClientSession, *PromptListChangedParams) { notificationChans["prompts"] <- 0 }, + ResourceListChangedHandler: func(context.Context, *ClientSession, *ResourceListChangedParams) { notificationChans["resources"] <- 0 }, } c := NewClient("testClient", "v1.0.0", opts) rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files")) @@ -103,6 +125,7 @@ func TestEndToEnd(t *testing.T) { t.Fatal(err) } + waitForNotification(t, "initialized") if err := cs.Ping(ctx, nil); err != nil { t.Fatalf("ping failed: %v", err) } @@ -141,6 +164,11 @@ func TestEndToEnd(t *testing.T) { if _, err := cs.GetPrompt(ctx, &GetPromptParams{Name: "fail"}); err == nil || !strings.Contains(err.Error(), failure.Error()) { t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) } + + s.AddPrompts(&ServerPrompt{Prompt: &Prompt{Name: "T"}}) + waitForNotification(t, "prompts") + s.RemovePrompts("T") + waitForNotification(t, "prompts") }) t.Run("tools", func(t *testing.T) { @@ -198,6 +226,11 @@ func TestEndToEnd(t *testing.T) { if diff := cmp.Diff(wantFail, gotFail); diff != "" { t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) } + + s.AddTools(&ServerTool{Tool: &Tool{Name: "T"}}) + waitForNotification(t, "tools") + s.RemoveTools("T") + waitForNotification(t, "tools") }) t.Run("resources", func(t *testing.T) { @@ -254,6 +287,11 @@ func TestEndToEnd(t *testing.T) { } } } + + s.AddResources(&ServerResource{Resource: &Resource{URI: "http://U"}}) + waitForNotification(t, "resources") + s.RemoveResources("http://U") + waitForNotification(t, "resources") }) t.Run("roots", func(t *testing.T) { rootRes, err := ss.ListRoots(ctx, &ListRootsParams{}) @@ -265,6 +303,11 @@ func TestEndToEnd(t *testing.T) { if diff := cmp.Diff(wantRoots, gotRoots); diff != "" { t.Errorf("roots/list mismatch (-want +got):\n%s", diff) } + + c.AddRoots(&Root{URI: "U"}) + waitForNotification(t, "roots") + c.RemoveRoots("U") + waitForNotification(t, "roots") }) t.Run("sampling", func(t *testing.T) { // TODO: test that a client that doesn't have the handler returns CodeUnsupportedMethod. @@ -462,6 +505,10 @@ func TestMiddleware(t *testing.T) { 2 >initialize 2 notifications/initialized +2 >notifications/initialized +2 tools/list 2 >tools/list 2