From 4f6510b9ae339e61d19d1d21a6495b55e08daa4a Mon Sep 17 00:00:00 2001 From: Rob Findley Date: Tue, 13 May 2025 20:58:47 +0000 Subject: [PATCH] internal/mcp: merge the protocol package into mcp In accordance with the updated design, merge the protocol package into the mcp package. To do so, rename protocol.Content to WireContent and protocol.ResourceContents to WireResource. Also, remove some stray logging in readResource. Change-Id: I76f61ce8d67dd3ac520ef2e1959ade0a0785a480 Reviewed-on: https://go-review.googlesource.com/c/tools/+/672375 LUCI-TryBot-Result: Go LUCI Reviewed-by: Jonathan Amsterdam --- examples/hello/main.go | 7 +- mcp/client.go | 51 +++++------ mcp/cmd_test.go | 5 +- mcp/content.go | 79 ++++++++++++---- mcp/content_test.go | 17 ++-- mcp/{protocol => }/generate.go | 21 +++-- mcp/mcp.go | 2 + mcp/mcp_test.go | 47 +++++----- mcp/prompt.go | 31 +++---- mcp/prompt_test.go | 11 +-- mcp/{protocol => }/protocol.go | 162 ++++++++++++++++----------------- mcp/protocol/content.go | 56 ------------ mcp/protocol/doc.go | 13 --- mcp/server.go | 56 ++++++------ mcp/sse_test.go | 5 +- mcp/tool.go | 29 +++--- mcp/tool_test.go | 2 +- mcp/transport.go | 5 +- 18 files changed, 282 insertions(+), 317 deletions(-) rename mcp/{protocol => }/generate.go (96%) rename mcp/{protocol => }/protocol.go (94%) delete mode 100644 mcp/protocol/content.go delete mode 100644 mcp/protocol/doc.go diff --git a/examples/hello/main.go b/examples/hello/main.go index c75960e..20087b4 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -12,7 +12,6 @@ import ( "os" "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) var httpAddr = flag.String("http", "", "if set, use SSE HTTP at this address, instead of stdin/stdout") @@ -27,10 +26,10 @@ func SayHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) ([]m }, nil } -func PromptHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) (*protocol.GetPromptResult, error) { - return &protocol.GetPromptResult{ +func PromptHi(ctx context.Context, cc *mcp.ServerConnection, params *HiParams) (*mcp.GetPromptResult, error) { + return &mcp.GetPromptResult{ Description: "Code review prompt", - Messages: []protocol.PromptMessage{ + Messages: []mcp.PromptMessage{ {Role: "user", Content: mcp.TextContent{Text: "Say hi to " + params.Name}.ToWire()}, }, }, nil diff --git a/mcp/client.go b/mcp/client.go index b6b1d84..b0f2756 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -13,7 +13,6 @@ import ( "sync" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) // A Client is an MCP client, which may be connected to an MCP server @@ -25,8 +24,8 @@ type Client struct { opts ClientOptions mu sync.Mutex conn *jsonrpc2.Connection - roots *featureSet[protocol.Root] - initializeResult *protocol.InitializeResult + roots *featureSet[Root] + initializeResult *initializeResult } // NewClient creates a new Client. @@ -39,7 +38,7 @@ func NewClient(name, version string, t Transport, opts *ClientOptions) *Client { name: name, version: version, transport: t, - roots: newFeatureSet(func(r protocol.Root) string { return r.URI }), + roots: newFeatureSet(func(r Root) string { return r.URI }), } if opts != nil { c.opts = *opts @@ -84,13 +83,13 @@ func (c *Client) Start(ctx context.Context) (err error) { if err != nil { return err } - params := &protocol.InitializeParams{ - ClientInfo: protocol.Implementation{Name: c.name, Version: c.version}, + params := &initializeParams{ + ClientInfo: implementation{Name: c.name, Version: c.version}, } if err := call(ctx, c.conn, "initialize", params, &c.initializeResult); err != nil { return err } - if err := c.conn.Notify(ctx, "notifications/initialized", &protocol.InitializedParams{}); err != nil { + if err := c.conn.Notify(ctx, "notifications/initialized", &initializedParams{}); err != nil { return err } return nil @@ -113,7 +112,7 @@ func (c *Client) Wait() error { // replacing any with the same URIs, // and notifies any connected servers. // TODO: notification -func (c *Client) AddRoots(roots ...protocol.Root) { +func (c *Client) AddRoots(roots ...Root) { c.mu.Lock() defer c.mu.Unlock() c.roots.add(roots...) @@ -129,10 +128,10 @@ func (c *Client) RemoveRoots(uris ...string) { c.roots.remove(uris...) } -func (c *Client) listRoots(_ context.Context, _ *protocol.ListRootsParams) (*protocol.ListRootsResult, error) { +func (c *Client) listRoots(_ context.Context, _ *ListRootsParams) (*ListRootsResult, error) { c.mu.Lock() defer c.mu.Unlock() - return &protocol.ListRootsResult{ + return &ListRootsResult{ Roots: slices.Collect(c.roots.all()), }, nil } @@ -160,10 +159,10 @@ func (c *Client) Ping(ctx context.Context) error { } // ListPrompts lists prompts that are currently available on the server. -func (c *Client) ListPrompts(ctx context.Context) ([]protocol.Prompt, error) { +func (c *Client) ListPrompts(ctx context.Context) ([]Prompt, error) { var ( - params = &protocol.ListPromptsParams{} - result protocol.ListPromptsResult + params = &ListPromptsParams{} + result ListPromptsResult ) if err := call(ctx, c.conn, "prompts/list", params, &result); err != nil { return nil, err @@ -172,13 +171,13 @@ func (c *Client) ListPrompts(ctx context.Context) ([]protocol.Prompt, error) { } // GetPrompt gets a prompt from the server. -func (c *Client) GetPrompt(ctx context.Context, name string, args map[string]string) (*protocol.GetPromptResult, error) { +func (c *Client) GetPrompt(ctx context.Context, name string, args map[string]string) (*GetPromptResult, error) { var ( - params = &protocol.GetPromptParams{ + params = &GetPromptParams{ Name: name, Arguments: args, } - result = &protocol.GetPromptResult{} + result = &GetPromptResult{} ) if err := call(ctx, c.conn, "prompts/get", params, result); err != nil { return nil, err @@ -187,10 +186,10 @@ func (c *Client) GetPrompt(ctx context.Context, name string, args map[string]str } // ListTools lists tools that are currently available on the server. -func (c *Client) ListTools(ctx context.Context) ([]protocol.Tool, error) { +func (c *Client) ListTools(ctx context.Context) ([]Tool, error) { var ( - params = &protocol.ListToolsParams{} - result protocol.ListToolsResult + params = &ListToolsParams{} + result ListToolsResult ) if err := call(ctx, c.conn, "tools/list", params, &result); err != nil { return nil, err @@ -199,7 +198,7 @@ func (c *Client) ListTools(ctx context.Context) ([]protocol.Tool, error) { } // CallTool calls the tool with the given name and arguments. -func (c *Client) CallTool(ctx context.Context, name string, args map[string]any) (_ *protocol.CallToolResult, err error) { +func (c *Client) CallTool(ctx context.Context, name string, args map[string]any) (_ *CallToolResult, err error) { defer func() { if err != nil { err = fmt.Errorf("calling tool %q: %w", name, err) @@ -214,21 +213,21 @@ func (c *Client) CallTool(ctx context.Context, name string, args map[string]any) argsJSON[name] = argJSON } - params := &protocol.CallToolParams{ + params := &CallToolParams{ Name: name, Arguments: argsJSON, } - return standardCall[protocol.CallToolResult](ctx, c.conn, "tools/call", params) + return standardCall[CallToolResult](ctx, c.conn, "tools/call", params) } // ListResources lists the resources that are currently available on the server. -func (c *Client) ListResources(ctx context.Context, params *protocol.ListResourcesParams) (*protocol.ListResourcesResult, error) { - return standardCall[protocol.ListResourcesResult](ctx, c.conn, "resources/list", params) +func (c *Client) ListResources(ctx context.Context, params *ListResourcesParams) (*ListResourcesResult, error) { + return standardCall[ListResourcesResult](ctx, c.conn, "resources/list", params) } // ReadResource ask the server to read a resource and return its contents. -func (c *Client) ReadResource(ctx context.Context, params *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) { - return standardCall[protocol.ReadResourceResult](ctx, c.conn, "resources/read", params) +func (c *Client) ReadResource(ctx context.Context, params *ReadResourceParams) (*ReadResourceResult, error) { + return standardCall[ReadResourceResult](ctx, c.conn, "resources/read", params) } func standardCall[TRes, TParams any](ctx context.Context, conn *jsonrpc2.Connection, method string, params TParams) (*TRes, error) { diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index 28a8b1c..9d17749 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -13,7 +13,6 @@ import ( "github.com/google/go-cmp/cmp" "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) const runAsServer = "_MCP_RUN_AS_SERVER" @@ -57,8 +56,8 @@ func TestCmdTransport(t *testing.T) { if err != nil { log.Fatal(err) } - want := &protocol.CallToolResult{ - Content: []protocol.Content{{Type: "text", Text: "Hi user"}}, + want := &mcp.CallToolResult{ + Content: []mcp.WireContent{{Type: "text", Text: "Hi user"}}, } if diff := cmp.Diff(want, got); diff != "" { t.Errorf("greet returned unexpected content (-want +got):\n%s", diff) diff --git a/mcp/content.go b/mcp/content.go index d77674b..7d783e6 100644 --- a/mcp/content.go +++ b/mcp/content.go @@ -5,18 +5,63 @@ package mcp import ( + "encoding/json" "fmt" - - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) +// The []byte fields below are marked omitzero, not omitempty: +// we want to marshal an empty byte slice. + +// WireContent is the wire format for content. +// It represents the protocol types TextContent, ImageContent, AudioContent +// and EmbeddedResource. +// The Type field distinguishes them. In the protocol, each type has a constant +// value for the field. +// At most one of Text, Data, and Resource is non-zero. +type WireContent struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + MIMEType string `json:"mimeType,omitempty"` + Data []byte `json:"data,omitzero"` + Resource *WireResource `json:"resource,omitempty"` + Annotations *Annotations `json:"annotations,omitempty"` +} + +// A WireResource is either a TextResourceContents or a BlobResourceContents. +// See https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2025-03-26/schema.ts#L524-L551 +// for the inheritance structure. +// If Blob is nil, this is a TextResourceContents; otherwise it's a BlobResourceContents. +// +// The URI field describes the resource location. +type WireResource struct { + URI string `json:"uri,"` + MIMEType string `json:"mimeType,omitempty"` + Text string `json:"text"` + Blob []byte `json:"blob,omitzero"` +} + +func (c *WireContent) UnmarshalJSON(data []byte) error { + type wireContent WireContent // for naive unmarshaling + var c2 wireContent + if err := json.Unmarshal(data, &c2); err != nil { + return err + } + switch c2.Type { + case "text", "image", "audio", "resource": + default: + return fmt.Errorf("unrecognized content type %s", c.Type) + } + *c = WireContent(c2) + return nil +} + // Content is the union of supported content types: [TextContent], // [ImageContent], [AudioContent], and [ResourceContent]. // // ToWire converts content to its jsonrpc2 wire format. type Content interface { // TODO: unexport this, and move the tests that use it to this package. - ToWire() protocol.Content + ToWire() WireContent } // TextContent is a textual content. @@ -24,8 +69,8 @@ type TextContent struct { Text string } -func (c TextContent) ToWire() protocol.Content { - return protocol.Content{Type: "text", Text: c.Text} +func (c TextContent) ToWire() WireContent { + return WireContent{Type: "text", Text: c.Text} } // ImageContent contains base64-encoded image data. @@ -34,8 +79,8 @@ type ImageContent struct { MIMEType string } -func (c ImageContent) ToWire() protocol.Content { - return protocol.Content{Type: "image", MIMEType: c.MIMEType, Data: c.Data} +func (c ImageContent) ToWire() WireContent { + return WireContent{Type: "image", MIMEType: c.MIMEType, Data: c.Data} } // AudioContent contains base64-encoded audio data. @@ -44,8 +89,8 @@ type AudioContent struct { MIMEType string } -func (c AudioContent) ToWire() protocol.Content { - return protocol.Content{Type: "audio", MIMEType: c.MIMEType, Data: c.Data} +func (c AudioContent) ToWire() WireContent { + return WireContent{Type: "audio", MIMEType: c.MIMEType, Data: c.Data} } // ResourceContent contains embedded resources. @@ -53,13 +98,13 @@ type ResourceContent struct { Resource EmbeddedResource } -func (r ResourceContent) ToWire() protocol.Content { +func (r ResourceContent) ToWire() WireContent { res := r.Resource.toWire() - return protocol.Content{Type: "resource", Resource: &res} + return WireContent{Type: "resource", Resource: &res} } type EmbeddedResource interface { - toWire() protocol.ResourceContents + toWire() WireResource } // The {Text,Blob}ResourceContents types match the protocol definitions, @@ -72,8 +117,8 @@ type TextResourceContents struct { Text string } -func (r TextResourceContents) toWire() protocol.ResourceContents { - return protocol.ResourceContents{ +func (r TextResourceContents) toWire() WireResource { + return WireResource{ URI: r.URI, MIMEType: r.MIMEType, Text: r.Text, @@ -88,8 +133,8 @@ type BlobResourceContents struct { Blob []byte } -func (r BlobResourceContents) toWire() protocol.ResourceContents { - return protocol.ResourceContents{ +func (r BlobResourceContents) toWire() WireResource { + return WireResource{ URI: r.URI, MIMEType: r.MIMEType, Blob: r.Blob, @@ -98,7 +143,7 @@ func (r BlobResourceContents) toWire() protocol.ResourceContents { // ContentFromWireContent converts content from the jsonrpc2 wire format to a // typed Content value. -func ContentFromWireContent(c protocol.Content) Content { +func ContentFromWireContent(c WireContent) Content { switch c.Type { case "text": return TextContent{Text: c.Text} diff --git a/mcp/content_test.go b/mcp/content_test.go index 4abf35e..152bc27 100644 --- a/mcp/content_test.go +++ b/mcp/content_test.go @@ -9,22 +9,21 @@ import ( "github.com/google/go-cmp/cmp" "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) func TestContent(t *testing.T) { tests := []struct { in mcp.Content - want protocol.Content + want mcp.WireContent }{ - {mcp.TextContent{Text: "hello"}, protocol.Content{Type: "text", Text: "hello"}}, + {mcp.TextContent{Text: "hello"}, mcp.WireContent{Type: "text", Text: "hello"}}, { mcp.ImageContent{Data: []byte("a1b2c3"), MIMEType: "image/png"}, - protocol.Content{Type: "image", Data: []byte("a1b2c3"), MIMEType: "image/png"}, + mcp.WireContent{Type: "image", Data: []byte("a1b2c3"), MIMEType: "image/png"}, }, { mcp.AudioContent{Data: []byte("a1b2c3"), MIMEType: "audio/wav"}, - protocol.Content{Type: "audio", Data: []byte("a1b2c3"), MIMEType: "audio/wav"}, + mcp.WireContent{Type: "audio", Data: []byte("a1b2c3"), MIMEType: "audio/wav"}, }, { mcp.ResourceContent{ @@ -34,9 +33,9 @@ func TestContent(t *testing.T) { Text: "abc", }, }, - protocol.Content{ + mcp.WireContent{ Type: "resource", - Resource: &protocol.ResourceContents{ + Resource: &mcp.WireResource{ URI: "file://foo", MIMEType: "text", Text: "abc", @@ -51,9 +50,9 @@ func TestContent(t *testing.T) { Blob: []byte("a1b2c3"), }, }, - protocol.Content{ + mcp.WireContent{ Type: "resource", - Resource: &protocol.ResourceContents{ + Resource: &mcp.WireResource{ URI: "file://foo", MIMEType: "text", Blob: []byte("a1b2c3"), diff --git a/mcp/protocol/generate.go b/mcp/generate.go similarity index 96% rename from mcp/protocol/generate.go rename to mcp/generate.go index 887245f..36b2bf6 100644 --- a/mcp/protocol/generate.go +++ b/mcp/generate.go @@ -67,15 +67,15 @@ var declarations = config{ Fields: config{"Params": {Name: "GetPromptParams"}}, }, "GetPromptResult": {}, - "Implementation": {}, + "Implementation": {Name: "implementation"}, "InitializeRequest": { Name: "-", - Fields: config{"Params": {Name: "InitializeParams"}}, + Fields: config{"Params": {Name: "initializeParams"}}, }, - "InitializeResult": {}, + "InitializeResult": {Name: "initializeResult"}, "InitializedNotification": { Name: "-", - Fields: config{"Params": {Name: "InitializedParams"}}, + Fields: config{"Params": {Name: "initializedParams"}}, }, "ListPromptsRequest": { Name: "-", @@ -107,17 +107,18 @@ var declarations = config{ Fields: config{"Params": {Name: "ReadResourceParams"}}, }, "ReadResourceResult": { - Fields: config{"Contents": {Substitute: "*ResourceContents"}}, + Fields: config{"Contents": {Substitute: "*WireResource"}}, }, "Resource": {}, "Role": {}, "Root": {}, "ServerCapabilities": { + Name: "serverCapabilities", Fields: config{ - "Prompts": {Name: "PromptCapabilities"}, - "Resources": {Name: "ResourceCapabilities"}, - "Tools": {Name: "ToolCapabilities"}, + "Prompts": {Name: "promptCapabilities"}, + "Resources": {Name: "resourceCapabilities"}, + "Tools": {Name: "toolCapabilities"}, }, }, "Tool": { @@ -161,7 +162,7 @@ func main() { // Code generated by generate.go. DO NOT EDIT. -package protocol +package mcp import ( "encoding/json" @@ -287,7 +288,7 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma if slices.ContainsFunc(def.AnyOf, func(s *jsonschema.Schema) bool { return s.Ref == "#/definitions/TextContent" }) { - fmt.Fprintf(w, "Content") + fmt.Fprintf(w, "WireContent") } else { // E.g. union types. fmt.Fprintf(w, "json.RawMessage") diff --git a/mcp/mcp.go b/mcp/mcp.go index 734b96c..eaf3e56 100644 --- a/mcp/mcp.go +++ b/mcp/mcp.go @@ -2,6 +2,8 @@ // Use of this source code is governed by an MIT-style // license that can be found in the LICENSE file. +//go:generate go run generate.go + // The mcp package provides an SDK for writing model context protocol clients // and servers. // diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index baa419f..05d1a89 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -18,7 +18,6 @@ import ( "github.com/google/go-cmp/cmp/cmpopts" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" "github.com/modelcontextprotocol/go-sdk/jsonschema" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) type hiParams struct { @@ -50,15 +49,15 @@ func TestEndToEnd(t *testing.T) { ) s.AddPrompts( - NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerConnection, params struct{ Code string }) (*protocol.GetPromptResult, error) { - return &protocol.GetPromptResult{ + NewPrompt("code_review", "do a code review", func(_ context.Context, _ *ServerConnection, params struct{ Code string }) (*GetPromptResult, error) { + return &GetPromptResult{ Description: "Code review prompt", - Messages: []protocol.PromptMessage{ + Messages: []PromptMessage{ {Role: "user", Content: TextContent{Text: "Please review the following code: " + params.Code}.ToWire()}, }, }, nil }), - NewPrompt("fail", "", func(_ context.Context, _ *ServerConnection, params struct{}) (*protocol.GetPromptResult, error) { + NewPrompt("fail", "", func(_ context.Context, _ *ServerConnection, params struct{}) (*GetPromptResult, error) { return nil, failure }), ) @@ -83,7 +82,7 @@ func TestEndToEnd(t *testing.T) { }() c := NewClient("testClient", "v1.0.0", ct, nil) - c.AddRoots(protocol.Root{URI: "file:///root"}) + c.AddRoots(Root{URI: "file:///root"}) // Connect the client. if err := c.Start(ctx); err != nil { @@ -98,11 +97,11 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Errorf("prompts/list failed: %v", err) } - wantPrompts := []protocol.Prompt{ + wantPrompts := []Prompt{ { Name: "code_review", Description: "do a code review", - Arguments: []protocol.PromptArgument{{Name: "Code", Required: true}}, + Arguments: []PromptArgument{{Name: "Code", Required: true}}, }, {Name: "fail"}, } @@ -114,9 +113,9 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } - wantReview := &protocol.GetPromptResult{ + wantReview := &GetPromptResult{ Description: "Code review prompt", - Messages: []protocol.PromptMessage{{ + Messages: []PromptMessage{{ Content: TextContent{Text: "Please review the following code: 1+1"}.ToWire(), Role: "user", }}, @@ -135,7 +134,7 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Errorf("tools/list failed: %v", err) } - wantTools := []protocol.Tool{ + wantTools := []Tool{ { Name: "fail", Description: "just fail", @@ -165,8 +164,8 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } - wantHi := &protocol.CallToolResult{ - Content: []protocol.Content{{Type: "text", Text: "hi user"}}, + wantHi := &CallToolResult{ + Content: []WireContent{{Type: "text", Text: "hi user"}}, } if diff := cmp.Diff(wantHi, gotHi); diff != "" { t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) @@ -178,9 +177,9 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } - wantFail := &protocol.CallToolResult{ + wantFail := &CallToolResult{ IsError: true, - Content: []protocol.Content{{Type: "text", Text: failure.Error()}}, + Content: []WireContent{{Type: "text", Text: failure.Error()}}, } if diff := cmp.Diff(wantFail, gotFail); diff != "" { t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) @@ -188,21 +187,21 @@ func TestEndToEnd(t *testing.T) { }) t.Run("resources", func(t *testing.T) { - resource1 := protocol.Resource{ + resource1 := Resource{ Name: "public", MIMEType: "text/plain", URI: "file:///file1.txt", } - resource2 := protocol.Resource{ + resource2 := Resource{ Name: "public", // names are not unique IDs MIMEType: "text/plain", URI: "file:///nonexistent.txt", } - readHandler := func(_ context.Context, r protocol.Resource, _ *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) { + readHandler := func(_ context.Context, r Resource, _ *ReadResourceParams) (*ReadResourceResult, error) { if r.URI == "file:///file1.txt" { - return &protocol.ReadResourceResult{ - Contents: &protocol.ResourceContents{ + return &ReadResourceResult{ + Contents: &WireResource{ Text: "file contents", }, }, nil @@ -217,7 +216,7 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } - if diff := cmp.Diff([]protocol.Resource{resource1, resource2}, lrres.Resources); diff != "" { + if diff := cmp.Diff([]Resource{resource1, resource2}, lrres.Resources); diff != "" { t.Errorf("resources/list mismatch (-want, +got):\n%s", diff) } @@ -229,7 +228,7 @@ func TestEndToEnd(t *testing.T) { {"file:///nonexistent.txt", ""}, // TODO(jba): add resource template cases when we implement them } { - rres, err := c.ReadResource(ctx, &protocol.ReadResourceParams{URI: tt.uri}) + rres, err := c.ReadResource(ctx, &ReadResourceParams{URI: tt.uri}) if err != nil { var werr *jsonrpc2.WireError if errors.As(err, &werr) && werr.Code == codeResourceNotFound { @@ -256,7 +255,7 @@ func TestEndToEnd(t *testing.T) { break } - rootRes, err := sc.ListRoots(ctx, &protocol.ListRootsParams{}) + rootRes, err := sc.ListRoots(ctx, &ListRootsParams{}) if err != nil { t.Fatal(err) } @@ -283,7 +282,7 @@ func TestEndToEnd(t *testing.T) { // // The caller should cancel either the client connection or server connection // when the connections are no longer needed. -func basicConnection(t *testing.T, tools ...*Tool) (*ServerConnection, *Client) { +func basicConnection(t *testing.T, tools ...*ServerTool) (*ServerConnection, *Client) { t.Helper() ctx := context.Background() diff --git a/mcp/prompt.go b/mcp/prompt.go index bd57190..2ec9228 100644 --- a/mcp/prompt.go +++ b/mcp/prompt.go @@ -13,15 +13,14 @@ import ( "github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/jsonschema" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) // A PromptHandler handles a call to prompts/get. -type PromptHandler func(context.Context, *ServerConnection, map[string]string) (*protocol.GetPromptResult, error) +type PromptHandler func(context.Context, *ServerConnection, map[string]string) (*GetPromptResult, error) // A Prompt is a prompt definition bound to a prompt handler. -type Prompt struct { - Definition protocol.Prompt +type ServerPrompt struct { + Definition Prompt Handler PromptHandler } @@ -33,7 +32,7 @@ type Prompt struct { // of type string or *string. The argument names for the resulting prompt // definition correspond to the JSON names of the request fields, and any // fields that are not marked "omitempty" are considered required. -func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) (*protocol.GetPromptResult, error), opts ...PromptOption) *Prompt { +func NewPrompt[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) (*GetPromptResult, error), opts ...PromptOption) *ServerPrompt { schema, err := jsonschema.For[TReq]() if err != nil { panic(err) @@ -41,8 +40,8 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context, if schema.Type != "object" || !reflect.DeepEqual(schema.AdditionalProperties, &jsonschema.Schema{Not: &jsonschema.Schema{}}) { panic(fmt.Sprintf("handler request type must be a struct")) } - prompt := &Prompt{ - Definition: protocol.Prompt{ + prompt := &ServerPrompt{ + Definition: Prompt{ Name: name, Description: description, }, @@ -55,13 +54,13 @@ 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, protocol.PromptArgument{ + prompt.Definition.Arguments = append(prompt.Definition.Arguments, PromptArgument{ Name: name, Description: prop.Description, Required: required[name], }) } - prompt.Handler = func(ctx context.Context, cc *ServerConnection, args map[string]string) (*protocol.GetPromptResult, error) { + prompt.Handler = func(ctx context.Context, cc *ServerConnection, args map[string]string) (*GetPromptResult, error) { // For simplicity, just marshal and unmarshal the arguments. // This could be avoided in the future. rawArgs, err := json.Marshal(args) @@ -82,12 +81,12 @@ func NewPrompt[TReq any](name, description string, handler func(context.Context, // A PromptOption configures the behavior of a Prompt. type PromptOption interface { - set(*Prompt) + set(*ServerPrompt) } -type promptSetter func(*Prompt) +type promptSetter func(*ServerPrompt) -func (s promptSetter) set(p *Prompt) { s(p) } +func (s promptSetter) set(p *ServerPrompt) { s(p) } // Argument configures the 'schema' of a prompt argument. // If the argument does not exist, it is added. @@ -95,14 +94,14 @@ func (s promptSetter) set(p *Prompt) { s(p) } // Since prompt arguments are not a full JSON schema, Argument only accepts // Required and Description, and panics when encountering any other option. func Argument(name string, opts ...SchemaOption) PromptOption { - return promptSetter(func(p *Prompt) { - i := slices.IndexFunc(p.Definition.Arguments, func(arg protocol.PromptArgument) bool { + return promptSetter(func(p *ServerPrompt) { + i := slices.IndexFunc(p.Definition.Arguments, func(arg PromptArgument) bool { return arg.Name == name }) - var arg protocol.PromptArgument + var arg PromptArgument if i < 0 { i = len(p.Definition.Arguments) - arg = protocol.PromptArgument{Name: name} + arg = PromptArgument{Name: name} p.Definition.Arguments = append(p.Definition.Arguments, arg) } else { arg = p.Definition.Arguments[i] diff --git a/mcp/prompt_test.go b/mcp/prompt_test.go index 35075a4..2da8f0f 100644 --- a/mcp/prompt_test.go +++ b/mcp/prompt_test.go @@ -10,18 +10,17 @@ import ( "github.com/google/go-cmp/cmp" "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) // testPromptHandler is used for type inference in TestNewPrompt. -func testPromptHandler[T any](context.Context, *mcp.ServerConnection, T) (*protocol.GetPromptResult, error) { +func testPromptHandler[T any](context.Context, *mcp.ServerConnection, T) (*mcp.GetPromptResult, error) { panic("not implemented") } func TestNewPrompt(t *testing.T) { tests := []struct { - prompt *mcp.Prompt - want []protocol.PromptArgument + prompt *mcp.ServerPrompt + want []mcp.PromptArgument }{ { mcp.NewPrompt("empty", "", testPromptHandler[struct{}]), @@ -29,7 +28,7 @@ func TestNewPrompt(t *testing.T) { }, { mcp.NewPrompt("add_arg", "", testPromptHandler[struct{}], mcp.Argument("x")), - []protocol.PromptArgument{{Name: "x"}}, + []mcp.PromptArgument{{Name: "x"}}, }, { mcp.NewPrompt("combo", "", testPromptHandler[struct { @@ -39,7 +38,7 @@ func TestNewPrompt(t *testing.T) { }], mcp.Argument("name", mcp.Description("the person's name")), mcp.Argument("State", mcp.Required(false))), - []protocol.PromptArgument{ + []mcp.PromptArgument{ {Name: "State"}, {Name: "country"}, {Name: "name", Required: true, Description: "the person's name"}, diff --git a/mcp/protocol/protocol.go b/mcp/protocol.go similarity index 94% rename from mcp/protocol/protocol.go rename to mcp/protocol.go index f43702f..b885d2b 100644 --- a/mcp/protocol/protocol.go +++ b/mcp/protocol.go @@ -4,7 +4,7 @@ // Code generated by generate.go. DO NOT EDIT. -package protocol +package mcp import ( "encoding/json" @@ -47,7 +47,7 @@ type CallToolResult struct { // This result property is reserved by the protocol to allow clients and servers // to attach additional metadata to their responses. Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Content []Content `json:"content"` + Content []WireContent `json:"content"` // Whether the tool call ended in an error. // // If not set, this is assumed to be false (the call was successful). @@ -99,46 +99,6 @@ type GetPromptResult struct { Messages []PromptMessage `json:"messages"` } -// Describes the name and version of an MCP implementation. -type Implementation struct { - Name string `json:"name"` - Version string `json:"version"` -} - -type InitializeParams struct { - Capabilities ClientCapabilities `json:"capabilities"` - ClientInfo Implementation `json:"clientInfo"` - // The latest version of the Model Context Protocol that the client supports. - // The client MAY decide to support older versions as well. - ProtocolVersion string `json:"protocolVersion"` -} - -// After receiving an initialize request from the client, the server sends this -// response. -type InitializeResult struct { - // This result property is reserved by the protocol to allow clients and servers - // to attach additional metadata to their responses. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Capabilities ServerCapabilities `json:"capabilities"` - // Instructions describing how to use the server and its features. - // - // This can be used by clients to improve the LLM's understanding of available - // tools, resources, etc. It can be thought of like a "hint" to the model. For - // example, this information MAY be added to the system prompt. - Instructions string `json:"instructions,omitempty"` - // The version of the Model Context Protocol that the server wants to use. This - // may not match the version that the client requested. If the client cannot - // support this version, it MUST disconnect. - ProtocolVersion string `json:"protocolVersion"` - ServerInfo Implementation `json:"serverInfo"` -} - -type InitializedParams struct { - // This parameter name is reserved by MCP to allow clients and servers to attach - // additional metadata to their notifications. - Meta map[string]json.RawMessage `json:"_meta,omitempty"` -} - type ListPromptsParams struct { // An opaque token representing the current pagination position. If provided, // the server should return results starting after this cursor. @@ -230,19 +190,13 @@ type PromptArgument struct { Required bool `json:"required,omitempty"` } -// Present if the server offers any prompt templates. -type PromptCapabilities struct { - // Whether this server supports notifications for changes to the prompt list. - ListChanged bool `json:"listChanged,omitempty"` -} - // Describes a message returned as part of a prompt. // // This is similar to `SamplingMessage`, but also supports the embedding of // resources from the MCP server. type PromptMessage struct { - Content Content `json:"content"` - Role Role `json:"role"` + Content WireContent `json:"content"` + Role Role `json:"role"` } type ReadResourceParams struct { @@ -256,7 +210,7 @@ type ReadResourceResult struct { // This result property is reserved by the protocol to allow clients and servers // to attach additional metadata to their responses. Meta map[string]json.RawMessage `json:"_meta,omitempty"` - Contents *ResourceContents `json:"contents"` + Contents *WireResource `json:"contents"` } // A known resource that the server is capable of reading. @@ -284,14 +238,6 @@ type Resource struct { URI string `json:"uri"` } -// Present if the server offers any resources to read. -type ResourceCapabilities struct { - // Whether this server supports notifications for changes to the resource list. - ListChanged bool `json:"listChanged,omitempty"` - // Whether this server supports subscribing to resource updates. - Subscribe bool `json:"subscribe,omitempty"` -} - // The sender or recipient of messages and data in a conversation. type Role string @@ -307,27 +253,6 @@ type Root struct { URI string `json:"uri"` } -// Capabilities that a server may support. Known capabilities are defined here, -// in this schema, but this is not a closed set: any server can define its own, -// additional capabilities. -type ServerCapabilities struct { - // Present if the server supports argument autocompletion suggestions. - Completions struct { - } `json:"completions,omitempty"` - // Experimental, non-standard capabilities that the server supports. - Experimental map[string]struct { - } `json:"experimental,omitempty"` - // Present if the server supports sending log messages to the client. - Logging struct { - } `json:"logging,omitempty"` - // Present if the server offers any prompt templates. - Prompts *PromptCapabilities `json:"prompts,omitempty"` - // Present if the server offers any resources to read. - Resources *ResourceCapabilities `json:"resources,omitempty"` - // Present if the server offers any tools to call. - Tools *ToolCapabilities `json:"tools,omitempty"` -} - // Definition for a tool the client can call. type Tool struct { // Optional additional tool information. @@ -380,8 +305,83 @@ type ToolAnnotations struct { Title string `json:"title,omitempty"` } +// Describes the name and version of an MCP implementation. +type implementation struct { + Name string `json:"name"` + Version string `json:"version"` +} + +type initializeParams struct { + Capabilities ClientCapabilities `json:"capabilities"` + ClientInfo implementation `json:"clientInfo"` + // The latest version of the Model Context Protocol that the client supports. + // The client MAY decide to support older versions as well. + ProtocolVersion string `json:"protocolVersion"` +} + +// After receiving an initialize request from the client, the server sends this +// response. +type initializeResult struct { + // This result property is reserved by the protocol to allow clients and servers + // to attach additional metadata to their responses. + Meta map[string]json.RawMessage `json:"_meta,omitempty"` + Capabilities serverCapabilities `json:"capabilities"` + // Instructions describing how to use the server and its features. + // + // This can be used by clients to improve the LLM's understanding of available + // tools, resources, etc. It can be thought of like a "hint" to the model. For + // example, this information MAY be added to the system prompt. + Instructions string `json:"instructions,omitempty"` + // The version of the Model Context Protocol that the server wants to use. This + // may not match the version that the client requested. If the client cannot + // support this version, it MUST disconnect. + ProtocolVersion string `json:"protocolVersion"` + ServerInfo implementation `json:"serverInfo"` +} + +type initializedParams struct { + // This parameter name is reserved by MCP to allow clients and servers to attach + // additional metadata to their notifications. + Meta map[string]json.RawMessage `json:"_meta,omitempty"` +} + +// Present if the server offers any prompt templates. +type promptCapabilities struct { + // Whether this server supports notifications for changes to the prompt list. + ListChanged bool `json:"listChanged,omitempty"` +} + +// Present if the server offers any resources to read. +type resourceCapabilities struct { + // Whether this server supports notifications for changes to the resource list. + ListChanged bool `json:"listChanged,omitempty"` + // Whether this server supports subscribing to resource updates. + Subscribe bool `json:"subscribe,omitempty"` +} + +// Capabilities that a server may support. Known capabilities are defined here, +// in this schema, but this is not a closed set: any server can define its own, +// additional capabilities. +type serverCapabilities struct { + // Present if the server supports argument autocompletion suggestions. + Completions struct { + } `json:"completions,omitempty"` + // Experimental, non-standard capabilities that the server supports. + Experimental map[string]struct { + } `json:"experimental,omitempty"` + // Present if the server supports sending log messages to the client. + Logging struct { + } `json:"logging,omitempty"` + // Present if the server offers any prompt templates. + Prompts *promptCapabilities `json:"prompts,omitempty"` + // Present if the server offers any resources to read. + Resources *resourceCapabilities `json:"resources,omitempty"` + // Present if the server offers any tools to call. + Tools *toolCapabilities `json:"tools,omitempty"` +} + // Present if the server offers any tools to call. -type ToolCapabilities struct { +type toolCapabilities struct { // Whether this server supports notifications for changes to the tool list. ListChanged bool `json:"listChanged,omitempty"` } diff --git a/mcp/protocol/content.go b/mcp/protocol/content.go deleted file mode 100644 index 1202704..0000000 --- a/mcp/protocol/content.go +++ /dev/null @@ -1,56 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -package protocol - -import ( - "encoding/json" - "fmt" -) - -// The []byte fields below are marked omitzero, not omitempty: -// we want to marshal an empty byte slice. - -// Content is the wire format for content. -// It represents the protocol types TextContent, ImageContent, AudioContent -// and EmbeddedResource. -// The Type field distinguishes them. In the protocol, each type has a constant -// value for the field. -// At most one of Text, Data, and Resource is non-zero. -type Content struct { - Type string `json:"type"` - Text string `json:"text,omitempty"` - MIMEType string `json:"mimeType,omitempty"` - Data []byte `json:"data,omitzero"` - Resource *ResourceContents `json:"resource,omitempty"` - Annotations *Annotations `json:"annotations,omitempty"` -} - -// A ResourceContents is either a TextResourceContents or a BlobResourceContents. -// See https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/schema/2025-03-26/schema.ts#L524-L551 -// for the inheritance structure. -// If Blob is nil, this is a TextResourceContents; otherwise it's a BlobResourceContents. -// -// The URI field describes the resource location. -type ResourceContents struct { - URI string `json:"uri,"` - MIMEType string `json:"mimeType,omitempty"` - Text string `json:"text"` - Blob []byte `json:"blob,omitzero"` -} - -func (c *Content) UnmarshalJSON(data []byte) error { - type wireContent Content // for naive unmarshaling - var c2 wireContent - if err := json.Unmarshal(data, &c2); err != nil { - return err - } - switch c2.Type { - case "text", "image", "audio", "resource": - default: - return fmt.Errorf("unrecognized content type %s", c.Type) - } - *c = Content(c2) - return nil -} diff --git a/mcp/protocol/doc.go b/mcp/protocol/doc.go deleted file mode 100644 index 4109d43..0000000 --- a/mcp/protocol/doc.go +++ /dev/null @@ -1,13 +0,0 @@ -// Copyright 2025 The Go MCP SDK Authors. All rights reserved. -// Use of this source code is governed by an MIT-style -// license that can be found in the LICENSE file. - -//go:generate go run generate.go - -// The protocol package contains types that define the MCP protocol. -// -// It is auto-generated from the MCP spec. Run go generate to update it. -// The generated set of types is intended to be minimal, in the sense that we -// only generate types that are actually used by the SDK. See generate.go for -// instructions on how to generate more (or different) types. -package protocol diff --git a/mcp/server.go b/mcp/server.go index a0d0444..a404c74 100644 --- a/mcp/server.go +++ b/mcp/server.go @@ -9,13 +9,11 @@ import ( "encoding/json" "fmt" "iter" - "log" "net/url" "slices" "sync" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) // A Server is an instance of an MCP server. @@ -29,8 +27,8 @@ type Server struct { opts ServerOptions mu sync.Mutex - prompts *featureSet[*Prompt] - tools *featureSet[*Tool] + prompts *featureSet[*ServerPrompt] + tools *featureSet[*ServerTool] resources *featureSet[*ServerResource] conns []*ServerConnection } @@ -55,15 +53,15 @@ func NewServer(name, version string, opts *ServerOptions) *Server { name: name, version: version, opts: *opts, - prompts: newFeatureSet(func(p *Prompt) string { return p.Definition.Name }), - tools: newFeatureSet(func(t *Tool) string { return t.Definition.Name }), + prompts: newFeatureSet(func(p *ServerPrompt) string { return p.Definition.Name }), + tools: newFeatureSet(func(t *ServerTool) string { return t.Definition.Name }), resources: newFeatureSet(func(r *ServerResource) string { return r.Resource.URI }), } } // AddPrompts adds the given prompts to the server, // replacing any with the same names. -func (s *Server) AddPrompts(prompts ...*Prompt) { +func (s *Server) AddPrompts(prompts ...*ServerPrompt) { s.mu.Lock() defer s.mu.Unlock() s.prompts.add(prompts...) @@ -84,7 +82,7 @@ func (s *Server) RemovePrompts(names ...string) { // AddTools adds the given tools to the server, // replacing any with the same names. -func (s *Server) AddTools(tools ...*Tool) { +func (s *Server) AddTools(tools ...*ServerTool) { s.mu.Lock() defer s.mu.Unlock() s.tools.add(tools...) @@ -123,11 +121,11 @@ const codeResourceNotFound = -31002 // A ReadResourceHandler is a function that reads a resource. // If it cannot find the resource, it should return the result of calling [ResourceNotFoundError]. -type ReadResourceHandler func(context.Context, protocol.Resource, *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) +type ReadResourceHandler func(context.Context, Resource, *ReadResourceParams) (*ReadResourceResult, error) // A ServerResource associates a Resource with its handler. type ServerResource struct { - Resource protocol.Resource + Resource Resource Handler ReadResourceHandler } @@ -168,17 +166,17 @@ func (s *Server) Clients() iter.Seq[*ServerConnection] { return slices.Values(clients) } -func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *protocol.ListPromptsParams) (*protocol.ListPromptsResult, error) { +func (s *Server) listPrompts(_ context.Context, _ *ServerConnection, params *ListPromptsParams) (*ListPromptsResult, error) { s.mu.Lock() defer s.mu.Unlock() - res := new(protocol.ListPromptsResult) + res := new(ListPromptsResult) for p := range s.prompts.all() { res.Prompts = append(res.Prompts, p.Definition) } return res, nil } -func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *protocol.GetPromptParams) (*protocol.GetPromptResult, error) { +func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *GetPromptParams) (*GetPromptResult, error) { s.mu.Lock() prompt, ok := s.prompts.get(params.Name) s.mu.Unlock() @@ -189,17 +187,17 @@ func (s *Server) getPrompt(ctx context.Context, cc *ServerConnection, params *pr return prompt.Handler(ctx, cc, params.Arguments) } -func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *protocol.ListToolsParams) (*protocol.ListToolsResult, error) { +func (s *Server) listTools(_ context.Context, _ *ServerConnection, params *ListToolsParams) (*ListToolsResult, error) { s.mu.Lock() defer s.mu.Unlock() - res := new(protocol.ListToolsResult) + res := new(ListToolsResult) for t := range s.tools.all() { res.Tools = append(res.Tools, t.Definition) } return res, nil } -func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *protocol.CallToolParams) (*protocol.CallToolResult, error) { +func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) { s.mu.Lock() tool, ok := s.tools.get(params.Name) s.mu.Unlock() @@ -209,19 +207,17 @@ func (s *Server) callTool(ctx context.Context, cc *ServerConnection, params *pro return tool.Handler(ctx, cc, params.Arguments) } -func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *protocol.ListResourcesParams) (*protocol.ListResourcesResult, error) { +func (s *Server) listResources(_ context.Context, _ *ServerConnection, params *ListResourcesParams) (*ListResourcesResult, error) { s.mu.Lock() defer s.mu.Unlock() - res := new(protocol.ListResourcesResult) + res := new(ListResourcesResult) for r := range s.resources.all() { res.Resources = append(res.Resources, r.Resource) } return res, nil } -func (s *Server) readResource(ctx context.Context, _ *ServerConnection, params *protocol.ReadResourceParams) (*protocol.ReadResourceResult, error) { - log.Printf("readResource") - defer log.Printf("done") +func (s *Server) readResource(ctx context.Context, _ *ServerConnection, params *ReadResourceParams) (*ReadResourceResult, error) { uri := params.URI // Look up the resource URI in the list we have. // This is a security check as well as an information lookup. @@ -302,7 +298,7 @@ type ServerConnection struct { conn *jsonrpc2.Connection mu sync.Mutex - initializeParams *protocol.InitializeParams + initializeParams *initializeParams initialized bool } @@ -311,8 +307,8 @@ func (cc *ServerConnection) Ping(ctx context.Context) error { return call(ctx, cc.conn, "ping", nil, nil) } -func (cc *ServerConnection) ListRoots(ctx context.Context, params *protocol.ListRootsParams) (*protocol.ListRootsResult, error) { - return standardCall[protocol.ListRootsResult](ctx, cc.conn, "roots/list", params) +func (cc *ServerConnection) ListRoots(ctx context.Context, params *ListRootsParams) (*ListRootsResult, error) { + return standardCall[ListRootsResult](ctx, cc.conn, "roots/list", params) } func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) (any, error) { @@ -366,7 +362,7 @@ func (cc *ServerConnection) handle(ctx context.Context, req *jsonrpc2.Request) ( return nil, jsonrpc2.ErrNotHandled } -func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *protocol.InitializeParams) (*protocol.InitializeResult, error) { +func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, params *initializeParams) (*initializeResult, error) { cc.mu.Lock() cc.initializeParams = params cc.mu.Unlock() @@ -382,19 +378,19 @@ func (cc *ServerConnection) initialize(ctx context.Context, _ *ServerConnection, cc.mu.Unlock() }() - return &protocol.InitializeResult{ + return &initializeResult{ // TODO(rfindley): support multiple protocol versions. ProtocolVersion: "2024-11-05", - Capabilities: protocol.ServerCapabilities{ - Prompts: &protocol.PromptCapabilities{ + Capabilities: serverCapabilities{ + Prompts: &promptCapabilities{ ListChanged: false, // not yet supported }, - Tools: &protocol.ToolCapabilities{ + Tools: &toolCapabilities{ ListChanged: false, // not yet supported }, }, Instructions: cc.server.opts.Instructions, - ServerInfo: protocol.Implementation{ + ServerInfo: implementation{ Name: cc.server.name, Version: cc.server.version, }, diff --git a/mcp/sse_test.go b/mcp/sse_test.go index b2e3d0e..e4533e6 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -12,7 +12,6 @@ import ( "testing" "github.com/google/go-cmp/cmp" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) func TestSSEServer(t *testing.T) { @@ -48,8 +47,8 @@ func TestSSEServer(t *testing.T) { if err != nil { t.Fatal(err) } - wantHi := &protocol.CallToolResult{ - Content: []protocol.Content{{Type: "text", Text: "hi user"}}, + wantHi := &CallToolResult{ + Content: []WireContent{{Type: "text", Text: "hi user"}}, } if diff := cmp.Diff(wantHi, gotHi); diff != "" { t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) diff --git a/mcp/tool.go b/mcp/tool.go index 3c8bacc..78a0df4 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -11,15 +11,14 @@ import ( "github.com/modelcontextprotocol/go-sdk/internal/util" "github.com/modelcontextprotocol/go-sdk/jsonschema" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" ) // A ToolHandler handles a call to tools/call. -type ToolHandler func(context.Context, *ServerConnection, map[string]json.RawMessage) (*protocol.CallToolResult, error) +type ToolHandler func(context.Context, *ServerConnection, map[string]json.RawMessage) (*CallToolResult, error) // A Tool is a tool definition that is bound to a tool handler. -type Tool struct { - Definition protocol.Tool +type ServerTool struct { + Definition Tool Handler ToolHandler } @@ -36,12 +35,12 @@ type Tool struct { // // TODO: just have the handler return a CallToolResult: returning []Content is // going to be inconsistent with other server features. -func NewTool[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) ([]Content, error), opts ...ToolOption) *Tool { +func NewTool[TReq any](name, description string, handler func(context.Context, *ServerConnection, TReq) ([]Content, error), opts ...ToolOption) *ServerTool { schema, err := jsonschema.For[TReq]() if err != nil { panic(err) } - wrapped := func(ctx context.Context, cc *ServerConnection, args map[string]json.RawMessage) (*protocol.CallToolResult, error) { + wrapped := func(ctx context.Context, cc *ServerConnection, args map[string]json.RawMessage) (*CallToolResult, error) { // For simplicity, just marshal and unmarshal the arguments. // This could be avoided in the future. rawArgs, err := json.Marshal(args) @@ -56,18 +55,18 @@ func NewTool[TReq any](name, description string, handler func(context.Context, * // TODO: investigate why server errors are embedded in this strange way, // rather than returned as jsonrpc2 server errors. if err != nil { - return &protocol.CallToolResult{ - Content: []protocol.Content{TextContent{Text: err.Error()}.ToWire()}, + return &CallToolResult{ + Content: []WireContent{TextContent{Text: err.Error()}.ToWire()}, IsError: true, }, nil } - res := &protocol.CallToolResult{ + res := &CallToolResult{ Content: util.Apply(content, Content.ToWire), } return res, nil } - t := &Tool{ - Definition: protocol.Tool{ + t := &ServerTool{ + Definition: Tool{ Name: name, Description: description, InputSchema: schema, @@ -90,17 +89,17 @@ func unmarshalSchema(data json.RawMessage, _ *jsonschema.Schema, v any) error { // A ToolOption configures the behavior of a Tool. type ToolOption interface { - set(*Tool) + set(*ServerTool) } -type toolSetter func(*Tool) +type toolSetter func(*ServerTool) -func (s toolSetter) set(t *Tool) { s(t) } +func (s toolSetter) set(t *ServerTool) { s(t) } // Input applies the provided [SchemaOption] configuration to the tool's input // schema. func Input(opts ...SchemaOption) ToolOption { - return toolSetter(func(t *Tool) { + return toolSetter(func(t *ServerTool) { for _, opt := range opts { opt.set(t.Definition.InputSchema) } diff --git a/mcp/tool_test.go b/mcp/tool_test.go index f6230b2..d23c610 100644 --- a/mcp/tool_test.go +++ b/mcp/tool_test.go @@ -21,7 +21,7 @@ func testToolHandler[T any](context.Context, *mcp.ServerConnection, T) ([]mcp.Co func TestNewTool(t *testing.T) { tests := []struct { - tool *mcp.Tool + tool *mcp.ServerTool want *jsonschema.Schema }{ { diff --git a/mcp/transport.go b/mcp/transport.go index bf474b0..29a29c3 100644 --- a/mcp/transport.go +++ b/mcp/transport.go @@ -15,7 +15,6 @@ import ( "sync" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" - "github.com/modelcontextprotocol/go-sdk/mcp/protocol" "github.com/modelcontextprotocol/go-sdk/internal/xcontext" ) @@ -131,7 +130,7 @@ type canceller struct { // Preempt implements jsonrpc2.Preempter. func (c *canceller) Preempt(ctx context.Context, req *jsonrpc2.Request) (result any, err error) { if req.Method == "notifications/cancelled" { - var params protocol.CancelledParams + var params CancelledParams if err := json.Unmarshal(req.Params, ¶ms); err != nil { return nil, err } @@ -156,7 +155,7 @@ func call(ctx context.Context, conn *jsonrpc2.Connection, method string, params, return fmt.Errorf("calling %q: %w", method, ErrConnectionClosed) case ctx.Err() != nil: // Notify the peer of cancellation. - err := conn.Notify(xcontext.Detach(ctx), "notifications/cancelled", &protocol.CancelledParams{ + err := conn.Notify(xcontext.Detach(ctx), "notifications/cancelled", &CancelledParams{ Reason: ctx.Err().Error(), RequestID: call.ID().Raw(), })