internal/mcp: clean up handling of content

Add support for all content types, and formalize the conversion of
content to and from the wire format (as a discriminated union).

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

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
 }

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
+}

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":

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.

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.

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
 }

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

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))
+	}
 }

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
+}

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.

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

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
 }

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

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
 	}