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 <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
Auto-Submit: Sam Thanawalla <samthanawalla@google.com>

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.

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": {

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

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.

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 }