internal/mcp/design: further clean-up

Do a pass through the design doc, adding a bit more clean-up and
commentary.

Additionally, update ServerTool and ServerPrompt to be consistent with
the design.

Change-Id: I1b75c4690d86fe8946d512cb36a1fc2938b3c883
Reviewed-on: https://go-review.googlesource.com/c/tools/+/672576
Reviewed-by: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Sam Thanawalla <samthanawalla@google.com>
Auto-Submit: Robert Findley <rfindley@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

mcp/client.go

@@ -208,7 +208,7 @@ func (c *ClientSession) CallTool(ctx context.Context, name string, args map[stri
 
 // NOTE: the following struct should consist of all fields of callToolParams except name and arguments.
 
-// CallToolOptions contains options to [Client.CallTools].
+// CallToolOptions contains options to [ClientSession.CallTool].
 type CallToolOptions struct {
 	ProgressToken any // string or int
 }

mcp/mcp.go

@@ -10,13 +10,13 @@
 // To get started, create either a [Client] or [Server], and connect it to a
 // peer using a [Transport]. The diagram below illustrates how this works:
 //
-//	Client                                   Server
-//	  ⇅          (jsonrpc2)                    ⇅
-//	Client Transport ⇄ Server Transport ⇄ ServerSession
+//	Client                                                   Server
+//	 ⇅                          (jsonrpc2)                     ⇅
+//	ClientSession ⇄ Client Transport ⇄ Server Transport ⇄ ServerSession
 //
 // A [Client] is an MCP client, which can be configured with various client
 // capabilities. Clients may be connected to a [Server] instance
-// using the [Client.Start] method.
+// using the [Client.Connect] method.
 //
 // Similarly, a [Server] is an MCP server, which can be configured with various
 // server capabilities. Servers may be connected to one or more [Client]

mcp/prompt.go

@@ -20,8 +20,8 @@ type PromptHandler func(context.Context, *ServerSession, map[string]string) (*Ge
 
 // A Prompt is a prompt definition bound to a prompt handler.
 type ServerPrompt struct {
-	Definition *Prompt
-	Handler    PromptHandler
+	Prompt  *Prompt
+	Handler PromptHandler
 }
 
 // NewPrompt is a helper to use reflection to create a prompt for the given
@@ -41,7 +41,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
 		panic(fmt.Sprintf("handler request type must be a struct"))
 	}
 	prompt := &ServerPrompt{
-		Definition: &Prompt{
+		Prompt: &Prompt{
 			Name:        name,
 			Description: description,
 		},
@@ -54,7 +54,7 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context,
 		if prop.Type != "string" {
 			panic(fmt.Sprintf("handler type must consist only of string fields"))
 		}
-		prompt.Definition.Arguments = append(prompt.Definition.Arguments, &PromptArgument{
+		prompt.Prompt.Arguments = append(prompt.Prompt.Arguments, &PromptArgument{
 			Name:        name,
 			Description: prop.Description,
 			Required:    required[name],
@@ -95,16 +95,16 @@ func (s promptSetter) set(p *ServerPrompt) { s(p) }
 // Required and Description, and panics when encountering any other option.
 func Argument(name string, opts ...SchemaOption) PromptOption {
 	return promptSetter(func(p *ServerPrompt) {
-		i := slices.IndexFunc(p.Definition.Arguments, func(arg *PromptArgument) bool {
+		i := slices.IndexFunc(p.Prompt.Arguments, func(arg *PromptArgument) bool {
 			return arg.Name == name
 		})
 		var arg *PromptArgument
 		if i < 0 {
-			i = len(p.Definition.Arguments)
+			i = len(p.Prompt.Arguments)
 			arg = &PromptArgument{Name: name}
-			p.Definition.Arguments = append(p.Definition.Arguments, arg)
+			p.Prompt.Arguments = append(p.Prompt.Arguments, arg)
 		} else {
-			arg = p.Definition.Arguments[i]
+			arg = p.Prompt.Arguments[i]
 		}
 		for _, opt := range opts {
 			switch v := opt.(type) {
@@ -116,6 +116,6 @@ func Argument(name string, opts ...SchemaOption) PromptOption {
 				panic(fmt.Sprintf("unsupported prompt argument schema option %T", opt))
 			}
 		}
-		p.Definition.Arguments[i] = arg
+		p.Prompt.Arguments[i] = arg
 	})
 }

mcp/prompt_test.go

@@ -46,8 +46,8 @@ func TestNewPrompt(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		if diff := cmp.Diff(test.want, test.prompt.Definition.Arguments); diff != "" {
-			t.Errorf("NewPrompt(%v) mismatch (-want +got):\n%s", test.prompt.Definition.Name, diff)
+		if diff := cmp.Diff(test.want, test.prompt.Prompt.Arguments); diff != "" {
+			t.Errorf("NewPrompt(%v) mismatch (-want +got):\n%s", test.prompt.Prompt.Name, diff)
 		}
 	}
 }

mcp/server.go

@@ -53,8 +53,8 @@ func NewServer(name, version string, opts *ServerOptions) *Server {
 		name:      name,
 		version:   version,
 		opts:      *opts,
-		prompts:   newFeatureSet(func(p *ServerPrompt) string { return p.Definition.Name }),
-		tools:     newFeatureSet(func(t *ServerTool) string { return t.Definition.Name }),
+		prompts:   newFeatureSet(func(p *ServerPrompt) string { return p.Prompt.Name }),
+		tools:     newFeatureSet(func(t *ServerTool) string { return t.Tool.Name }),
 		resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }),
 	}
 }
@@ -170,7 +170,7 @@ func (s *Server) listPrompts(_ context.Context, _ *ServerSession, params *ListPr
 	defer s.mu.Unlock()
 	res := new(ListPromptsResult)
 	for p := range s.prompts.all() {
-		res.Prompts = append(res.Prompts, p.Definition)
+		res.Prompts = append(res.Prompts, p.Prompt)
 	}
 	return res, nil
 }
@@ -191,7 +191,7 @@ func (s *Server) listTools(_ context.Context, _ *ServerSession, params *ListTool
 	defer s.mu.Unlock()
 	res := new(ListToolsResult)
 	for t := range s.tools.all() {
-		res.Tools = append(res.Tools, t.Definition)
+		res.Tools = append(res.Tools, t.Tool)
 	}
 	return res, nil
 }

mcp/tool.go

@@ -17,8 +17,8 @@ type ToolHandler func(context.Context, *ServerSession, *CallToolParams) (*CallTo
 
 // A Tool is a tool definition that is bound to a tool handler.
 type ServerTool struct {
-	Definition *Tool
-	Handler    ToolHandler
+	Tool    *Tool
+	Handler ToolHandler
 }
 
 // NewTool is a helper to make a tool using reflection on the given handler.
@@ -59,7 +59,7 @@ func NewTool[TReq any](name, description string, handler func(context.Context, *
 		return res, nil
 	}
 	t := &ServerTool{
-		Definition: &Tool{
+		Tool: &Tool{
 			Name:        name,
 			Description: description,
 			InputSchema: schema,
@@ -94,7 +94,7 @@ func (s toolSetter) set(t *ServerTool) { s(t) }
 func Input(opts ...SchemaOption) ToolOption {
 	return toolSetter(func(t *ServerTool) {
 		for _, opt := range opts {
-			opt.set(t.Definition.InputSchema)
+			opt.set(t.Tool.InputSchema)
 		}
 	})
 }

mcp/tool_test.go

@@ -83,8 +83,8 @@ func TestNewTool(t *testing.T) {
 		},
 	}
 	for _, test := range tests {
-		if diff := cmp.Diff(test.want, test.tool.Definition.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
-			t.Errorf("NewTool(%v) mismatch (-want +got):\n%s", test.tool.Definition.Name, diff)
+		if diff := cmp.Diff(test.want, test.tool.Tool.InputSchema, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
+			t.Errorf("NewTool(%v) mismatch (-want +got):\n%s", test.tool.Tool.Name, diff)
 		}
 	}
 }