From 11b17f5858d393684f20a93909ef786f13147ae4 Mon Sep 17 00:00:00 2001 From: Sam Thanawalla Date: Thu, 12 Jun 2025 14:55:39 +0000 Subject: [PATCH] internal/mcp: add NotifyProgress This CL adds a NotifyProgress function and a ProgressNotificationHandler to enable progress sending and handling. Subsequent CL's to follow to handle progress updates for other methods such as list or readresource. Change-Id: Ifcc422e6f6231d7d773d5432c1cf64032d796709 Reviewed-on: https://go-review.googlesource.com/c/tools/+/681116 LUCI-TryBot-Result: Go LUCI Reviewed-by: Robert Findley Auto-Submit: Sam Thanawalla --- mcp/client.go | 22 ++++++++++++++++++---- mcp/generate.go | 4 ++++ mcp/mcp_test.go | 24 +++++++++++++++++++++++- mcp/protocol.go | 18 ++++++++++++++++++ mcp/server.go | 15 +++++++++++++++ 5 files changed, 78 insertions(+), 5 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 5ec7eeb..55d80f3 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -51,10 +51,11 @@ type ClientOptions struct { // 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) - LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams) + ToolListChangedHandler func(context.Context, *ClientSession, *ToolListChangedParams) + PromptListChangedHandler func(context.Context, *ClientSession, *PromptListChangedParams) + ResourceListChangedHandler func(context.Context, *ClientSession, *ResourceListChangedParams) + LoggingMessageHandler func(context.Context, *ClientSession, *LoggingMessageParams) + ProgressNotificationHandler func(context.Context, *ClientSession, *ProgressNotificationParams) } // bind implements the binder[*ClientSession] interface, so that Clients can @@ -231,6 +232,7 @@ var clientMethodInfos = map[string]methodInfo{ notificationPromptListChanged: newMethodInfo(clientMethod((*Client).callPromptChangedHandler)), notificationResourceListChanged: newMethodInfo(clientMethod((*Client).callResourceChangedHandler)), notificationLoggingMessage: newMethodInfo(clientMethod((*Client).callLoggingHandler)), + notificationProgress: newMethodInfo(sessionMethod((*ClientSession).callProgressNotificationHandler)), } func (cs *ClientSession) sendingMethodInfos() map[string]methodInfo { @@ -334,6 +336,18 @@ func (c *Client) callLoggingHandler(ctx context.Context, cs *ClientSession, para return nil, nil } +func (cs *ClientSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) { + return callNotificationHandler(ctx, cs.client.opts.ProgressNotificationHandler, cs, params) +} + +// NotifyProgress sends a progress notification from the client to the server +// associated with this session. +// This can be used if the client is performing a long-running task that was +// initiated by the server +func (cs *ClientSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error { + return handleNotify(ctx, cs, notificationProgress, 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/generate.go b/mcp/generate.go index 8ed7678..4ab7f1a 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -142,6 +142,10 @@ var declarations = config{ Name: "-", Fields: config{"Params": {Name: "PromptListChangedParams"}}, }, + "ProgressNotification": { + Name: "-", + Fields: config{"Params": {Name: "ProgressNotificationParams"}}, + }, "ProgressToken": {Name: "-", Substitute: "any"}, // null|number|string "RequestId": {Name: "-", Substitute: "any"}, // null|number|string "ReadResourceRequest": { diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index fa0ddb3..ca958a6 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -45,7 +45,7 @@ func TestEndToEnd(t *testing.T) { // Channels to check if notification callbacks happened. notificationChans := map[string]chan int{} - for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources"} { + for _, name := range []string{"initialized", "roots", "tools", "prompts", "resources", "progress_server", "progress_client"} { notificationChans[name] = make(chan int, 1) } waitForNotification := func(t *testing.T, name string) { @@ -60,6 +60,9 @@ func TestEndToEnd(t *testing.T) { sopts := &ServerOptions{ InitializedHandler: func(context.Context, *ServerSession, *InitializedParams) { notificationChans["initialized"] <- 0 }, RootsListChangedHandler: func(context.Context, *ServerSession, *RootsListChangedParams) { notificationChans["roots"] <- 0 }, + ProgressNotificationHandler: func(context.Context, *ServerSession, *ProgressNotificationParams) { + notificationChans["progress_server"] <- 0 + }, } s := NewServer("testServer", "v1.0.0", sopts) add(tools, s.AddTools, "greet", "fail") @@ -96,6 +99,9 @@ func TestEndToEnd(t *testing.T) { LoggingMessageHandler: func(_ context.Context, _ *ClientSession, lm *LoggingMessageParams) { loggingMessages <- lm }, + ProgressNotificationHandler: func(context.Context, *ClientSession, *ProgressNotificationParams) { + notificationChans["progress_client"] <- 0 + }, } c := NewClient("testClient", "v1.0.0", opts) rootAbs, err := filepath.Abs(filepath.FromSlash("testdata/files")) @@ -394,6 +400,22 @@ func TestEndToEnd(t *testing.T) { check(t) }) }) + t.Run("progress", func(t *testing.T) { + ss.NotifyProgress(ctx, &ProgressNotificationParams{ + ProgressToken: "token-xyz", + Message: "progress update", + Progress: 0.5, + Total: 2, + }) + waitForNotification(t, "progress_client") + cs.NotifyProgress(ctx, &ProgressNotificationParams{ + ProgressToken: "token-abc", + Message: "progress update", + Progress: 1, + Total: 2, + }) + waitForNotification(t, "progress_server") + }) // Disconnect. cs.Close() diff --git a/mcp/protocol.go b/mcp/protocol.go index 279c3e0..ac84a20 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -431,6 +431,24 @@ type PingParams struct { func (x *PingParams) GetMeta() *Meta { return &x.Meta } +type ProgressNotificationParams struct { + // This property is reserved by the protocol to allow clients and servers to + // attach additional metadata to their responses. + Meta Meta `json:"_meta,omitempty"` + // An optional message describing the current progress. + Message string `json:"message,omitempty"` + // The progress thus far. This should increase every time progress is made, even + // if the total is unknown. + Progress float64 `json:"progress"` + // The progress token which was given in the initial request, used to associate + // this notification with the request that is proceeding. + ProgressToken any `json:"progressToken"` + // Total number of items to process (or total progress required), if known. + Total float64 `json:"total,omitempty"` +} + +func (x *ProgressNotificationParams) GetMeta() *Meta { return &x.Meta } + // A prompt or prompt template that the server offers. type Prompt struct { // A list of arguments to use for templating the prompt. diff --git a/mcp/server.go b/mcp/server.go index b5e6d5a..13f1cba 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -55,6 +55,8 @@ type ServerOptions struct { PageSize int // If non-nil, called when "notifications/roots/list_changed" is received. RootsListChangedHandler func(context.Context, *ServerSession, *RootsListChangedParams) + // If non-nil, called when "notifications/progress" is received. + ProgressNotificationHandler func(context.Context, *ServerSession, *ProgressNotificationParams) } // NewServer creates a new MCP server. The resulting server has no features: @@ -465,6 +467,18 @@ func (s *Server) callRootsListChangedHandler(ctx context.Context, ss *ServerSess return callNotificationHandler(ctx, s.opts.RootsListChangedHandler, ss, params) } +func (ss *ServerSession) callProgressNotificationHandler(ctx context.Context, params *ProgressNotificationParams) (Result, error) { + return callNotificationHandler(ctx, ss.server.opts.ProgressNotificationHandler, ss, params) +} + +// NotifyProgress sends a progress notification from the server to the client +// associated with this session. +// This is typically used to report on the status of a long-running request +// that was initiated by the client. +func (ss *ServerSession) NotifyProgress(ctx context.Context, params *ProgressNotificationParams) error { + return handleNotify(ctx, ss, notificationProgress, params) +} + // A ServerSession is a logical connection from a single MCP client. Its // methods can be used to send requests or notifications to the client. Create // a session by calling [Server.Connect]. @@ -559,6 +573,7 @@ var serverMethodInfos = map[string]methodInfo{ methodSetLevel: newMethodInfo(sessionMethod((*ServerSession).setLevel)), notificationInitialized: newMethodInfo(serverMethod((*Server).callInitializedHandler)), notificationRootsListChanged: newMethodInfo(serverMethod((*Server).callRootsListChangedHandler)), + notificationProgress: newMethodInfo(sessionMethod((*ServerSession).callProgressNotificationHandler)), } func (ss *ServerSession) sendingMethodInfos() map[string]methodInfo { return clientMethodInfos }