From 74340f088dc2f1ec0de98a719b9b5d160341342a Mon Sep 17 00:00:00 2001 From: Jonathan Amsterdam Date: Sat, 17 May 2025 13:02:49 -0400 Subject: [PATCH] internal/mcp: generate method names Augment the generator to create constants for all method names. Change-Id: I858e516b1afec0aa3cf3b9852cbdfc16aec4c553 Reviewed-on: https://go-review.googlesource.com/c/tools/+/673177 LUCI-TryBot-Result: Go LUCI Reviewed-by: Robert Findley Reviewed-by: Sam Thanawalla --- mcp/client.go | 22 +++++++++++----------- mcp/generate.go | 18 ++++++++++++++++++ mcp/protocol.go | 27 +++++++++++++++++++++++++++ mcp/server.go | 20 ++++++++++---------- 4 files changed, 66 insertions(+), 21 deletions(-) diff --git a/mcp/client.go b/mcp/client.go index 8de5050..432c9c7 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -99,7 +99,7 @@ func (c *Client) Connect(ctx context.Context, t Transport) (cs *ClientSession, e _ = cs.Close() return nil, err } - if err := cs.conn.Notify(ctx, "notifications/initialized", &InitializedParams{}); err != nil { + if err := cs.conn.Notify(ctx, notificationInitialized, &InitializedParams{}); err != nil { _ = cs.Close() return nil, err } @@ -181,9 +181,9 @@ func (c *Client) AddMiddleware(middleware ...Middleware[ClientSession]) { // clientMethodInfos maps from the RPC method name to serverMethodInfos. var clientMethodInfos = map[string]methodInfo[ClientSession]{ - "ping": newMethodInfo(sessionMethod((*ClientSession).ping)), - "roots/list": newMethodInfo(clientMethod((*Client).listRoots)), - "sampling/createMessage": newMethodInfo(clientMethod((*Client).createMessage)), + methodPing: newMethodInfo(sessionMethod((*ClientSession).ping)), + methodListRoots: newMethodInfo(clientMethod((*Client).listRoots)), + methodCreateMessage: newMethodInfo(clientMethod((*Client).createMessage)), // TODO: notifications } @@ -209,22 +209,22 @@ func (c *ClientSession) ping(ct context.Context, params *PingParams) (struct{}, // Ping makes an MCP "ping" request to the server. func (c *ClientSession) Ping(ctx context.Context, params *PingParams) error { - return call(ctx, c.conn, "ping", params, nil) + return call(ctx, c.conn, methodPing, params, nil) } // ListPrompts lists prompts that are currently available on the server. func (c *ClientSession) ListPrompts(ctx context.Context, params *ListPromptsParams) (*ListPromptsResult, error) { - return standardCall[ListPromptsResult](ctx, c.conn, "prompts/list", params) + return standardCall[ListPromptsResult](ctx, c.conn, methodListPrompts, params) } // GetPrompt gets a prompt from the server. func (c *ClientSession) GetPrompt(ctx context.Context, params *GetPromptParams) (*GetPromptResult, error) { - return standardCall[GetPromptResult](ctx, c.conn, "prompts/get", params) + return standardCall[GetPromptResult](ctx, c.conn, methodGetPrompt, params) } // ListTools lists tools that are currently available on the server. func (c *ClientSession) ListTools(ctx context.Context, params *ListToolsParams) (*ListToolsResult, error) { - return standardCall[ListToolsResult](ctx, c.conn, "tools/list", params) + return standardCall[ListToolsResult](ctx, c.conn, methodListTools, params) } // CallTool calls the tool with the given name and arguments. @@ -244,7 +244,7 @@ func (c *ClientSession) CallTool(ctx context.Context, name string, args map[stri Name: name, Arguments: json.RawMessage(data), } - return standardCall[CallToolResult](ctx, c.conn, "tools/call", params) + return standardCall[CallToolResult](ctx, c.conn, methodCallTool, params) } // NOTE: the following struct should consist of all fields of callToolParams except name and arguments. @@ -256,12 +256,12 @@ type CallToolOptions struct { // ListResources lists the resources that are currently available on the server. func (c *ClientSession) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { - return standardCall[ListResourcesResult](ctx, c.conn, "resources/list", params) + return standardCall[ListResourcesResult](ctx, c.conn, methodListResources, params) } // ReadResource ask the server to read a resource and return its contents. func (c *ClientSession) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { - return standardCall[ReadResourceResult](ctx, c.conn, "resources/read", params) + return standardCall[ReadResourceResult](ctx, c.conn, methodReadResource, params) } func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) { diff --git a/mcp/generate.go b/mcp/generate.go index e291f47..12998d1 100644 --- a/mcp/generate.go +++ b/mcp/generate.go @@ -202,6 +202,24 @@ import ( fmt.Fprintln(buf) fmt.Fprint(buf, b.String()) } + // Write out method names. + fmt.Fprintln(buf, `const (`) + for name, s := range schema.Definitions { + prefix := "method" + method, found := strings.CutSuffix(name, "Request") + if !found { + prefix = "notification" + method, found = strings.CutSuffix(name, "Notification") + } + if found { + if ms, ok := s.Properties["method"]; ok { + if c := ms.Const; c != nil { + fmt.Fprintf(buf, "%s%s = %q\n", prefix, method, *c) + } + } + } + } + fmt.Fprintln(buf, `)`) formatted, err := format.Source(buf.Bytes()) if err != nil { diff --git a/mcp/protocol.go b/mcp/protocol.go index 12bfc35..75eaa5b 100644 --- a/mcp/protocol.go +++ b/mcp/protocol.go @@ -494,3 +494,30 @@ type toolCapabilities struct { // Whether this server supports notifications for changes to the tool list. ListChanged bool `json:"listChanged,omitempty"` } + +const ( + notificationCancelled = "notifications/cancelled" + methodInitialize = "initialize" + notificationProgress = "notifications/progress" + methodSetLevel = "logging/setLevel" + methodCreateMessage = "sampling/createMessage" + notificationResourceListChanged = "notifications/resources/list_changed" + notificationInitialized = "notifications/initialized" + methodUnsubscribe = "resources/unsubscribe" + notificationLoggingMessage = "notifications/message" + methodSubscribe = "resources/subscribe" + methodComplete = "completion/complete" + methodCallTool = "tools/call" + notificationPromptListChanged = "notifications/prompts/list_changed" + methodReadResource = "resources/read" + methodListResourceTemplates = "resources/templates/list" + methodListRoots = "roots/list" + notificationToolListChanged = "notifications/tools/list_changed" + methodGetPrompt = "prompts/get" + methodListPrompts = "prompts/list" + methodPing = "ping" + notificationRootsListChanged = "notifications/roots/list_changed" + methodListTools = "tools/list" + methodListResources = "resources/list" + notificationResourceUpdated = "notifications/resources/updated" +) diff --git a/mcp/server.go b/mcp/server.go index 642b16d..d3a258b 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -325,12 +325,12 @@ func (ss *ServerSession) Ping(ctx context.Context, _ *PingParams) error { // ListRoots lists the client roots. func (ss *ServerSession) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { - return standardCall[ListRootsResult](ctx, ss.conn, "roots/list", params) + return standardCall[ListRootsResult](ctx, ss.conn, methodListRoots, params) } // CreateMessage sends a sampling request to the client. func (ss *ServerSession) CreateMessage(ctx context.Context, params *CreateMessageParams) (*CreateMessageResult, error) { - return standardCall[CreateMessageResult](ctx, ss.conn, "sampling/createMessage", params) + return standardCall[CreateMessageResult](ctx, ss.conn, methodCreateMessage, params) } // AddMiddleware wraps the server's current method handler using the provided @@ -347,14 +347,14 @@ func (s *Server) AddMiddleware(middleware ...Middleware[ServerSession]) { // serverMethodInfos maps from the RPC method name to serverMethodInfos. var serverMethodInfos = map[string]methodInfo[ServerSession]{ - "initialize": newMethodInfo(sessionMethod((*ServerSession).initialize)), - "ping": newMethodInfo(sessionMethod((*ServerSession).ping)), - "prompts/list": newMethodInfo(serverMethod((*Server).listPrompts)), - "prompts/get": newMethodInfo(serverMethod((*Server).getPrompt)), - "tools/list": newMethodInfo(serverMethod((*Server).listTools)), - "tools/call": newMethodInfo(serverMethod((*Server).callTool)), - "resources/list": newMethodInfo(serverMethod((*Server).listResources)), - "resources/read": newMethodInfo(serverMethod((*Server).readResource)), + methodInitialize: newMethodInfo(sessionMethod((*ServerSession).initialize)), + methodPing: newMethodInfo(sessionMethod((*ServerSession).ping)), + methodListPrompts: newMethodInfo(serverMethod((*Server).listPrompts)), + methodGetPrompt: newMethodInfo(serverMethod((*Server).getPrompt)), + methodListTools: newMethodInfo(serverMethod((*Server).listTools)), + methodCallTool: newMethodInfo(serverMethod((*Server).callTool)), + methodListResources: newMethodInfo(serverMethod((*Server).listResources)), + methodReadResource: newMethodInfo(serverMethod((*Server).readResource)), // TODO: notifications }