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