diff --git a/examples/hello/main.go b/examples/hello/main.go index 82e3d6c..fcdd70e 100644 --- a/examples/hello/main.go +++ b/examples/hello/main.go @@ -6,7 +6,6 @@ package main import ( "context" - "encoding/json" "flag" "fmt" "net/http" @@ -29,19 +28,10 @@ func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) ([]m } func PromptHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) (*protocol.GetPromptResult, error) { - // (see related TODOs about cleaning up content construction) - content, err := json.Marshal(protocol.TextContent{ - Type: "text", - Text: "Say hi to " + params.Name, - }) - if err != nil { - return nil, err - } return &protocol.GetPromptResult{ Description: "Code review prompt", Messages: []protocol.PromptMessage{ - // TODO: move 'Content' to the protocol package. - {Role: "user", Content: json.RawMessage(content)}, + {Role: "user", Content: mcp.TextContent{Text: "Say hi to " + params.Name}.ToWire()}, }, }, nil } diff --git a/internal/protocol/content.go b/internal/protocol/content.go new file mode 100644 index 0000000..f615db8 --- /dev/null +++ b/internal/protocol/content.go @@ -0,0 +1,42 @@ +// 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" +) + +// Content is the wire format for content, including all fields. +type Content struct { + Type string `json:"type"` + Text string `json:"text,omitempty"` + MIMEType string `json:"mimeType,omitempty"` + Data string `json:"data,omitempty"` + Resource *Resource `json:"resource,omitempty"` +} + +// Resource is the wire format for embedded resources, including all fields. +type Resource struct { + URI string `json:"uri,"` + MIMEType string `json:"mimeType,omitempty"` + Text string `json:"text"` + Blob *string `json:"blob"` // blob is a pointer to distinguish empty from missing data +} + +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/internal/protocol/generate.go b/internal/protocol/generate.go index 5172315..60b0f0a 100644 --- a/internal/protocol/generate.go +++ b/internal/protocol/generate.go @@ -91,7 +91,6 @@ var declarations = config{ "Tools": {Name: "ToolCapabilities"}, }, }, - "TextContent": {Name: "TextContent"}, "Tool": { Name: "Tool", Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}}, @@ -249,8 +248,15 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma } if def.Type == "" { - // E.g. union types. - fmt.Fprintf(w, "json.RawMessage") + // special case: recognize Content + if slices.ContainsFunc(def.AnyOf, func(s *jsonschema.Schema) bool { + return s.Ref == "#/definitions/TextContent" + }) { + fmt.Fprintf(w, "Content") + } else { + // E.g. union types. + fmt.Fprintf(w, "json.RawMessage") + } } else { switch def.Type { case "array": diff --git a/internal/protocol/protocol.go b/internal/protocol/protocol.go index 1af7cff..126816c 100644 --- a/internal/protocol/protocol.go +++ b/internal/protocol/protocol.go @@ -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 []json.RawMessage `json:"content"` + Content []Content `json:"content"` // Whether the tool call ended in an error. // // If not set, this is assumed to be false (the call was successful). @@ -198,8 +198,8 @@ type PromptCapabilities struct { // This is similar to `SamplingMessage`, but also supports the embedding of // resources from the MCP server. type PromptMessage struct { - Content json.RawMessage `json:"content"` - Role Role `json:"role"` + Content Content `json:"content"` + Role Role `json:"role"` } // Present if the server offers any resources to read. @@ -231,15 +231,6 @@ type ServerCapabilities struct { Tools *ToolCapabilities `json:"tools,omitempty"` } -// Text provided to or from an LLM. -type TextContent struct { - // Optional annotations for the client. - Annotations *Annotations `json:"annotations,omitempty"` - // The text content of the message. - Text string `json:"text"` - Type string `json:"type"` -} - // Definition for a tool the client can call. type Tool struct { // Optional additional tool information. diff --git a/internal/util/util.go b/internal/util/util.go index 03a5dbb..0cd3eac 100644 --- a/internal/util/util.go +++ b/internal/util/util.go @@ -10,6 +10,15 @@ import ( "slices" ) +// Apply returns a new slice resulting from applying f to each element of x. +func Apply[S ~[]E, E, F any](x S, f func(E) F) []F { + y := make([]F, len(x)) + for i, e := range x { + y[i] = f(e) + } + return y +} + // Helpers below are copied from gopls' moremaps package. // sorted returns an iterator over the entries of m in key order. diff --git a/mcp/client.go b/mcp/client.go index 540715e..ab0b119 100644 --- a/mcp/client.go +++ b/mcp/client.go @@ -7,7 +7,6 @@ package mcp import ( "context" "encoding/json" - "errors" "fmt" "iter" "slices" @@ -194,7 +193,7 @@ func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, err // TODO(jba): make the following true: // If the provided arguments do not conform to the schema for the given tool, // the call fails. -func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[string]any) (_ []Content, err error) { +func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[string]any) (_ *protocol.CallToolResult, err error) { defer func() { if err != nil { err = fmt.Errorf("calling tool %q: %w", name, err) @@ -218,15 +217,5 @@ func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[ if err := call(ctx, sc.conn, "tools/call", params, &result); err != nil { return nil, err } - content, err := unmarshalContent(result.Content) - if err != nil { - return nil, fmt.Errorf("unmarshaling tool content: %v", err) - } - if result.IsError { - if len(content) != 1 || !is[TextContent](content[0]) { - return nil, errors.New("malformed error content") - } - return nil, errors.New(content[0].(TextContent).Text) - } - return content, nil + return &result, nil } diff --git a/mcp/cmd_test.go b/mcp/cmd_test.go index 7d6215d..4d64a03 100644 --- a/mcp/cmd_test.go +++ b/mcp/cmd_test.go @@ -13,6 +13,7 @@ import ( "github.com/google/go-cmp/cmp" "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/internal/protocol" ) const runAsServer = "_MCP_RUN_AS_SERVER" @@ -57,7 +58,9 @@ func TestCmdTransport(t *testing.T) { if err != nil { log.Fatal(err) } - want := []mcp.Content{mcp.TextContent{Text: "Hi user"}} + want := &protocol.CallToolResult{ + Content: []protocol.Content{{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 987ee0e..093b939 100644 --- a/mcp/content.go +++ b/mcp/content.go @@ -5,53 +5,17 @@ package mcp import ( - "encoding/json" "fmt" "github.com/modelcontextprotocol/go-sdk/internal/protocol" ) -// Content is the abstract result of a Tool call. +// Content is the union of supported content types: [TextContent], +// [ImageContent], [AudioContent], and [ResourceContent]. // -// TODO: support all content types. +// ToWire converts content to its jsonrpc2 wire format. type Content interface { - toProtocol() any -} - -func marshalContent(content []Content) []json.RawMessage { - var msgs []json.RawMessage - for _, c := range content { - msg, err := json.Marshal(c.toProtocol()) - if err != nil { - panic(fmt.Sprintf("marshaling content: %v", err)) - } - msgs = append(msgs, msg) - } - return msgs -} - -func unmarshalContent(msgs []json.RawMessage) ([]Content, error) { - var content []Content - for _, msg := range msgs { - var allContent struct { - Type string `json:"type"` - Text json.RawMessage - } - if err := json.Unmarshal(msg, &allContent); err != nil { - return nil, fmt.Errorf("content missing \"type\"") - } - switch allContent.Type { - case "text": - var text string - if err := json.Unmarshal(allContent.Text, &text); err != nil { - return nil, fmt.Errorf("unmarshalling text content: %v", err) - } - content = append(content, TextContent{Text: text}) - default: - return nil, fmt.Errorf("unsupported content type %q", allContent.Type) - } - } - return content, nil + ToWire() protocol.Content } // TextContent is a textual content. @@ -59,6 +23,102 @@ type TextContent struct { Text string } -func (c TextContent) toProtocol() any { - return protocol.TextContent{Type: "text", Text: c.Text} +func (c TextContent) ToWire() protocol.Content { + return protocol.Content{Type: "text", Text: c.Text} +} + +// ImageContent contains base64-encoded image data. +type ImageContent struct { + Data string + MimeType string +} + +func (c ImageContent) ToWire() protocol.Content { + return protocol.Content{Type: "image", MIMEType: c.MimeType, Data: c.Data} +} + +// AudioContent contains base64-encoded audio data. +type AudioContent struct { + Data string + MimeType string +} + +func (c AudioContent) ToWire() protocol.Content { + return protocol.Content{Type: "audio", MIMEType: c.MimeType, Data: c.Data} +} + +// ResourceContent contains embedded resources. +type ResourceContent struct { + Resource Resource +} + +func (r ResourceContent) ToWire() protocol.Content { + res := r.Resource.ToWire() + return protocol.Content{Type: "resource", Resource: &res} +} + +type Resource interface { + ToWire() protocol.Resource +} + +type TextResource struct { + URI string + MimeType string + Text string +} + +func (r TextResource) ToWire() protocol.Resource { + return protocol.Resource{ + URI: r.URI, + MIMEType: r.MimeType, + Text: r.Text, + } +} + +type BlobResource struct { + URI string + MimeType string + Blob string +} + +func (r BlobResource) ToWire() protocol.Resource { + blob := r.Blob + return protocol.Resource{ + URI: r.URI, + MIMEType: r.MimeType, + Blob: &blob, + } +} + +// ContentFromWireContent converts content from the jsonrpc2 wire format to a +// typed Content value. +func ContentFromWireContent(c protocol.Content) Content { + switch c.Type { + case "text": + return TextContent{Text: c.Text} + case "image": + return ImageContent{Data: c.Data, MimeType: c.MIMEType} + case "audio": + return AudioContent{Data: c.Data, MimeType: c.MIMEType} + case "resource": + r := ResourceContent{} + if c.Resource != nil { + if c.Resource.Blob != nil { + r.Resource = BlobResource{ + URI: c.Resource.URI, + MimeType: c.Resource.MIMEType, + Blob: *c.Resource.Blob, + } + } else { + r.Resource = TextResource{ + URI: c.Resource.URI, + MimeType: c.Resource.MIMEType, + Text: c.Resource.Text, + } + } + } + return r + default: + panic(fmt.Sprintf("unrecognized wire content type %q", c.Type)) + } } diff --git a/mcp/content_test.go b/mcp/content_test.go new file mode 100644 index 0000000..d9774ff --- /dev/null +++ b/mcp/content_test.go @@ -0,0 +1,75 @@ +// 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 mcp_test + +import ( + "testing" + + "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/internal/protocol" +) + +func TestContent(t *testing.T) { + tests := []struct { + in mcp.Content + want protocol.Content + }{ + {mcp.TextContent{Text: "hello"}, protocol.Content{Type: "text", Text: "hello"}}, + { + mcp.ImageContent{Data: "a1b2c3", MimeType: "image/png"}, + protocol.Content{Type: "image", Data: "a1b2c3", MIMEType: "image/png"}, + }, + { + mcp.AudioContent{Data: "a1b2c3", MimeType: "audio/wav"}, + protocol.Content{Type: "audio", Data: "a1b2c3", MIMEType: "audio/wav"}, + }, + { + mcp.ResourceContent{ + Resource: mcp.TextResource{ + URI: "file://foo", + MimeType: "text", + Text: "abc", + }, + }, + protocol.Content{ + Type: "resource", + Resource: &protocol.Resource{ + URI: "file://foo", + MIMEType: "text", + Text: "abc", + }, + }, + }, + { + mcp.ResourceContent{ + Resource: mcp.BlobResource{ + URI: "file://foo", + MimeType: "text", + Blob: "a1b2c3", + }, + }, + protocol.Content{ + Type: "resource", + Resource: &protocol.Resource{ + URI: "file://foo", + MIMEType: "text", + Blob: ptr("a1b2c3"), + }, + }, + }, + } + + for _, test := range tests { + got := test.in.ToWire() + if diff := cmp.Diff(test.want, got); diff != "" { + t.Errorf("ToWire mismatch (-want +got):\n%s", diff) + } + } +} + +func ptr[T any](t T) *T { + return &t +} diff --git a/mcp/mcp_test.go b/mcp/mcp_test.go index 418ba09..8fbac1f 100644 --- a/mcp/mcp_test.go +++ b/mcp/mcp_test.go @@ -6,7 +6,6 @@ package mcp import ( "context" - "encoding/json" "errors" "fmt" "slices" @@ -50,19 +49,10 @@ func TestEndToEnd(t *testing.T) { s.AddPrompts( MakePrompt("code_review", "do a code review", func(_ context.Context, _ *ClientConnection, params struct{ Code string }) (*protocol.GetPromptResult, error) { - // TODO(rfindley): clean up this handling of content. - content, err := json.Marshal(protocol.TextContent{ - Type: "text", - Text: "Please review the following code: " + params.Code, - }) - if err != nil { - return nil, err - } return &protocol.GetPromptResult{ Description: "Code review prompt", Messages: []protocol.PromptMessage{ - // TODO: move 'Content' to the protocol package. - {Role: "user", Content: json.RawMessage(content)}, + {Role: "user", Content: TextContent{Text: "Please review the following code: " + params.Code}.ToWire()}, }, }, nil }), @@ -126,10 +116,17 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } - // TODO: assert on the full review, once content is easier to create. - if got, want := gotReview.Description, "Code review prompt"; got != want { - t.Errorf("prompts/get 'code_review': got description %q, want %q", got, want) + wantReview := &protocol.GetPromptResult{ + Description: "Code review prompt", + Messages: []protocol.PromptMessage{{ + Content: TextContent{Text: "Please review the following code: 1+1"}.ToWire(), + Role: "user", + }}, } + if diff := cmp.Diff(wantReview, gotReview); diff != "" { + t.Errorf("prompts/get 'code_review' mismatch (-want +got):\n%s", diff) + } + if _, err := sc.GetPrompt(ctx, "fail", map[string]string{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) } @@ -165,13 +162,25 @@ func TestEndToEnd(t *testing.T) { if err != nil { t.Fatal(err) } - wantHi := []Content{TextContent{Text: "hi user"}} + wantHi := &protocol.CallToolResult{ + Content: []protocol.Content{{Type: "text", Text: "hi user"}}, + } if diff := cmp.Diff(wantHi, gotHi); diff != "" { t.Errorf("tools/call 'greet' mismatch (-want +got):\n%s", diff) } - if _, err := sc.CallTool(ctx, "fail", map[string]any{}); err == nil || !strings.Contains(err.Error(), failure.Error()) { - t.Errorf("fail returned unexpected error: got %v, want containing %v", err, failure) + gotFail, err := sc.CallTool(ctx, "fail", map[string]any{}) + // Counter-intuitively, when a tool fails, we don't expect an RPC error for + // call tool: instead, the failure is embedded in the result. + if err != nil { + t.Fatal(err) + } + wantFail := &protocol.CallToolResult{ + IsError: true, + Content: []protocol.Content{{Type: "text", Text: failure.Error()}}, + } + if diff := cmp.Diff(wantFail, gotFail); diff != "" { + t.Errorf("tools/call 'fail' mismatch (-want +got):\n%s", diff) } // Disconnect. diff --git a/mcp/server_example_test.go b/mcp/server_example_test.go index 63826b5..eabc812 100644 --- a/mcp/server_example_test.go +++ b/mcp/server_example_test.go @@ -40,11 +40,11 @@ func ExampleServer() { log.Fatal(err) } - content, err := serverConnection.CallTool(ctx, "greet", map[string]any{"name": "user"}) + res, err := serverConnection.CallTool(ctx, "greet", map[string]any{"name": "user"}) if err != nil { log.Fatal(err) } - fmt.Println(content[0].(mcp.TextContent).Text) + fmt.Println(res.Content[0].Text) serverConnection.Close() clientConnection.Wait() diff --git a/mcp/sse_example_test.go b/mcp/sse_example_test.go index baf6dc8..2f53a6c 100644 --- a/mcp/sse_example_test.go +++ b/mcp/sse_example_test.go @@ -40,11 +40,11 @@ func ExampleSSEHandler() { } defer serverConn.Close() - content, err := serverConn.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2}) + res, err := serverConn.CallTool(ctx, "add", map[string]any{"x": 1, "y": 2}) if err != nil { log.Fatal(err) } - fmt.Println(content[0].(mcp.TextContent).Text) + fmt.Println(res.Content[0].Text) // Output: 3 } diff --git a/mcp/sse_test.go b/mcp/sse_test.go index 52135c8..d259f4e 100644 --- a/mcp/sse_test.go +++ b/mcp/sse_test.go @@ -12,6 +12,7 @@ import ( "testing" "github.com/google/go-cmp/cmp" + "github.com/modelcontextprotocol/go-sdk/internal/protocol" ) func TestSSEServer(t *testing.T) { @@ -48,7 +49,9 @@ func TestSSEServer(t *testing.T) { if err != nil { t.Fatal(err) } - wantHi := []Content{TextContent{Text: "hi user"}} + wantHi := &protocol.CallToolResult{ + Content: []protocol.Content{{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 3de32d5..02748bb 100644 --- a/mcp/tool.go +++ b/mcp/tool.go @@ -11,6 +11,7 @@ import ( "github.com/modelcontextprotocol/go-sdk/internal/jsonschema" "github.com/modelcontextprotocol/go-sdk/internal/protocol" + "github.com/modelcontextprotocol/go-sdk/internal/util" ) // A ToolHandler handles a call to tools/call. @@ -56,12 +57,12 @@ func MakeTool[TReq any](name, description string, handler func(context.Context, // rather than returned as jsonrpc2 server errors. if err != nil { return &protocol.CallToolResult{ - Content: marshalContent([]Content{TextContent{Text: err.Error()}}), + Content: []protocol.Content{TextContent{Text: err.Error()}.ToWire()}, IsError: true, }, nil } res := &protocol.CallToolResult{ - Content: marshalContent(content), + Content: util.Apply(content, Content.ToWire), } return res, nil }