Skip to content

Commit efd15d8

Browse files
findleyrgopherbot
findleyr
authored and
gopherbot
committed
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 <[email protected]> Reviewed-by: Jonathan Amsterdam <[email protected]> Auto-Submit: Robert Findley <[email protected]>
1 parent 80e0fd8 commit efd15d8

14 files changed

+282
-104
lines changed

‎internal/mcp/client.go

+2-13
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ package mcp
77
import (
88
"context"
99
"encoding/json"
10-
"errors"
1110
"fmt"
1211
"iter"
1312
"slices"
@@ -194,7 +193,7 @@ func (sc *ServerConnection) ListTools(ctx context.Context) ([]protocol.Tool, err
194193
// TODO(jba): make the following true:
195194
// If the provided arguments do not conform to the schema for the given tool,
196195
// the call fails.
197-
func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[string]any) (_ []Content, err error) {
196+
func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[string]any) (_ *protocol.CallToolResult, err error) {
198197
defer func() {
199198
if err != nil {
200199
err = fmt.Errorf("calling tool %q: %w", name, err)
@@ -218,15 +217,5 @@ func (sc *ServerConnection) CallTool(ctx context.Context, name string, args map[
218217
if err := call(ctx, sc.conn, "tools/call", params, &result); err != nil {
219218
return nil, err
220219
}
221-
content, err := unmarshalContent(result.Content)
222-
if err != nil {
223-
return nil, fmt.Errorf("unmarshaling tool content: %v", err)
224-
}
225-
if result.IsError {
226-
if len(content) != 1 || !is[TextContent](content[0]) {
227-
return nil, errors.New("malformed error content")
228-
}
229-
return nil, errors.New(content[0].(TextContent).Text)
230-
}
231-
return content, nil
220+
return &result, nil
232221
}

‎internal/mcp/cmd_test.go

+4-1
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313

1414
"github.com/google/go-cmp/cmp"
1515
"golang.org/x/tools/internal/mcp"
16+
"golang.org/x/tools/internal/mcp/internal/protocol"
1617
)
1718

1819
const runAsServer = "_MCP_RUN_AS_SERVER"
@@ -57,7 +58,9 @@ func TestCmdTransport(t *testing.T) {
5758
if err != nil {
5859
log.Fatal(err)
5960
}
60-
want := []mcp.Content{mcp.TextContent{Text: "Hi user"}}
61+
want := &protocol.CallToolResult{
62+
Content: []protocol.Content{{Type: "text", Text: "Hi user"}},
63+
}
6164
if diff := cmp.Diff(want, got); diff != "" {
6265
t.Errorf("greet returned unexpected content (-want +got):\n%s", diff)
6366
}

‎internal/mcp/content.go

+100-40
Original file line numberDiff line numberDiff line change
@@ -5,60 +5,120 @@
55
package mcp
66

77
import (
8-
"encoding/json"
98
"fmt"
109

1110
"golang.org/x/tools/internal/mcp/internal/protocol"
1211
)
1312

14-
// Content is the abstract result of a Tool call.
13+
// Content is the union of supported content types: [TextContent],
14+
// [ImageContent], [AudioContent], and [ResourceContent].
1515
//
16-
// TODO: support all content types.
16+
// ToWire converts content to its jsonrpc2 wire format.
1717
type Content interface {
18-
toProtocol() any
18+
ToWire() protocol.Content
1919
}
2020

21-
func marshalContent(content []Content) []json.RawMessage {
22-
var msgs []json.RawMessage
23-
for _, c := range content {
24-
msg, err := json.Marshal(c.toProtocol())
25-
if err != nil {
26-
panic(fmt.Sprintf("marshaling content: %v", err))
27-
}
28-
msgs = append(msgs, msg)
29-
}
30-
return msgs
21+
// TextContent is a textual content.
22+
type TextContent struct {
23+
Text string
3124
}
3225

33-
func unmarshalContent(msgs []json.RawMessage) ([]Content, error) {
34-
var content []Content
35-
for _, msg := range msgs {
36-
var allContent struct {
37-
Type string `json:"type"`
38-
Text json.RawMessage
39-
}
40-
if err := json.Unmarshal(msg, &allContent); err != nil {
41-
return nil, fmt.Errorf("content missing \"type\"")
42-
}
43-
switch allContent.Type {
44-
case "text":
45-
var text string
46-
if err := json.Unmarshal(allContent.Text, &text); err != nil {
47-
return nil, fmt.Errorf("unmarshalling text content: %v", err)
48-
}
49-
content = append(content, TextContent{Text: text})
50-
default:
51-
return nil, fmt.Errorf("unsupported content type %q", allContent.Type)
52-
}
26+
func (c TextContent) ToWire() protocol.Content {
27+
return protocol.Content{Type: "text", Text: c.Text}
28+
}
29+
30+
// ImageContent contains base64-encoded image data.
31+
type ImageContent struct {
32+
Data string
33+
MimeType string
34+
}
35+
36+
func (c ImageContent) ToWire() protocol.Content {
37+
return protocol.Content{Type: "image", MIMEType: c.MimeType, Data: c.Data}
38+
}
39+
40+
// AudioContent contains base64-encoded audio data.
41+
type AudioContent struct {
42+
Data string
43+
MimeType string
44+
}
45+
46+
func (c AudioContent) ToWire() protocol.Content {
47+
return protocol.Content{Type: "audio", MIMEType: c.MimeType, Data: c.Data}
48+
}
49+
50+
// ResourceContent contains embedded resources.
51+
type ResourceContent struct {
52+
Resource Resource
53+
}
54+
55+
func (r ResourceContent) ToWire() protocol.Content {
56+
res := r.Resource.ToWire()
57+
return protocol.Content{Type: "resource", Resource: &res}
58+
}
59+
60+
type Resource interface {
61+
ToWire() protocol.Resource
62+
}
63+
64+
type TextResource struct {
65+
URI string
66+
MimeType string
67+
Text string
68+
}
69+
70+
func (r TextResource) ToWire() protocol.Resource {
71+
return protocol.Resource{
72+
URI: r.URI,
73+
MIMEType: r.MimeType,
74+
Text: r.Text,
5375
}
54-
return content, nil
5576
}
5677

57-
// TextContent is a textual content.
58-
type TextContent struct {
59-
Text string
78+
type BlobResource struct {
79+
URI string
80+
MimeType string
81+
Blob string
6082
}
6183

62-
func (c TextContent) toProtocol() any {
63-
return protocol.TextContent{Type: "text", Text: c.Text}
84+
func (r BlobResource) ToWire() protocol.Resource {
85+
blob := r.Blob
86+
return protocol.Resource{
87+
URI: r.URI,
88+
MIMEType: r.MimeType,
89+
Blob: &blob,
90+
}
91+
}
92+
93+
// ContentFromWireContent converts content from the jsonrpc2 wire format to a
94+
// typed Content value.
95+
func ContentFromWireContent(c protocol.Content) Content {
96+
switch c.Type {
97+
case "text":
98+
return TextContent{Text: c.Text}
99+
case "image":
100+
return ImageContent{Data: c.Data, MimeType: c.MIMEType}
101+
case "audio":
102+
return AudioContent{Data: c.Data, MimeType: c.MIMEType}
103+
case "resource":
104+
r := ResourceContent{}
105+
if c.Resource != nil {
106+
if c.Resource.Blob != nil {
107+
r.Resource = BlobResource{
108+
URI: c.Resource.URI,
109+
MimeType: c.Resource.MIMEType,
110+
Blob: *c.Resource.Blob,
111+
}
112+
} else {
113+
r.Resource = TextResource{
114+
URI: c.Resource.URI,
115+
MimeType: c.Resource.MIMEType,
116+
Text: c.Resource.Text,
117+
}
118+
}
119+
}
120+
return r
121+
default:
122+
panic(fmt.Sprintf("unrecognized wire content type %q", c.Type))
123+
}
64124
}

‎internal/mcp/content_test.go

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
// Copyright 2025 The Go Authors. All rights reserved.
2+
// Use of this source code is governed by a BSD-style
3+
// license that can be found in the LICENSE file.
4+
5+
package mcp_test
6+
7+
import (
8+
"testing"
9+
10+
"github.com/google/go-cmp/cmp"
11+
"golang.org/x/tools/internal/mcp"
12+
"golang.org/x/tools/internal/mcp/internal/protocol"
13+
)
14+
15+
func TestContent(t *testing.T) {
16+
tests := []struct {
17+
in mcp.Content
18+
want protocol.Content
19+
}{
20+
{mcp.TextContent{Text: "hello"}, protocol.Content{Type: "text", Text: "hello"}},
21+
{
22+
mcp.ImageContent{Data: "a1b2c3", MimeType: "image/png"},
23+
protocol.Content{Type: "image", Data: "a1b2c3", MIMEType: "image/png"},
24+
},
25+
{
26+
mcp.AudioContent{Data: "a1b2c3", MimeType: "audio/wav"},
27+
protocol.Content{Type: "audio", Data: "a1b2c3", MIMEType: "audio/wav"},
28+
},
29+
{
30+
mcp.ResourceContent{
31+
Resource: mcp.TextResource{
32+
URI: "file://foo",
33+
MimeType: "text",
34+
Text: "abc",
35+
},
36+
},
37+
protocol.Content{
38+
Type: "resource",
39+
Resource: &protocol.Resource{
40+
URI: "file://foo",
41+
MIMEType: "text",
42+
Text: "abc",
43+
},
44+
},
45+
},
46+
{
47+
mcp.ResourceContent{
48+
Resource: mcp.BlobResource{
49+
URI: "file://foo",
50+
MimeType: "text",
51+
Blob: "a1b2c3",
52+
},
53+
},
54+
protocol.Content{
55+
Type: "resource",
56+
Resource: &protocol.Resource{
57+
URI: "file://foo",
58+
MIMEType: "text",
59+
Blob: ptr("a1b2c3"),
60+
},
61+
},
62+
},
63+
}
64+
65+
for _, test := range tests {
66+
got := test.in.ToWire()
67+
if diff := cmp.Diff(test.want, got); diff != "" {
68+
t.Errorf("ToWire mismatch (-want +got):\n%s", diff)
69+
}
70+
}
71+
}
72+
73+
func ptr[T any](t T) *T {
74+
return &t
75+
}

‎internal/mcp/examples/hello/main.go

+1-11
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ package main
66

77
import (
88
"context"
9-
"encoding/json"
109
"flag"
1110
"fmt"
1211
"net/http"
@@ -29,19 +28,10 @@ func SayHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) ([]m
2928
}
3029

3130
func PromptHi(ctx context.Context, cc *mcp.ClientConnection, params *HiParams) (*protocol.GetPromptResult, error) {
32-
// (see related TODOs about cleaning up content construction)
33-
content, err := json.Marshal(protocol.TextContent{
34-
Type: "text",
35-
Text: "Say hi to " + params.Name,
36-
})
37-
if err != nil {
38-
return nil, err
39-
}
4031
return &protocol.GetPromptResult{
4132
Description: "Code review prompt",
4233
Messages: []protocol.PromptMessage{
43-
// TODO: move 'Content' to the protocol package.
44-
{Role: "user", Content: json.RawMessage(content)},
34+
{Role: "user", Content: mcp.TextContent{Text: "Say hi to " + params.Name}.ToWire()},
4535
},
4636
}, nil
4737
}

‎internal/mcp/internal/protocol/content.go

+42
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
// Copyright 2025 The Go Authors. All rights reserved.
2+
// Use of this source code is governed by a BSD-style
3+
// license that can be found in the LICENSE file.
4+
5+
package protocol
6+
7+
import (
8+
"encoding/json"
9+
"fmt"
10+
)
11+
12+
// Content is the wire format for content, including all fields.
13+
type Content struct {
14+
Type string `json:"type"`
15+
Text string `json:"text,omitempty"`
16+
MIMEType string `json:"mimeType,omitempty"`
17+
Data string `json:"data,omitempty"`
18+
Resource *Resource `json:"resource,omitempty"`
19+
}
20+
21+
// Resource is the wire format for embedded resources, including all fields.
22+
type Resource struct {
23+
URI string `json:"uri,"`
24+
MIMEType string `json:"mimeType,omitempty"`
25+
Text string `json:"text"`
26+
Blob *string `json:"blob"` // blob is a pointer to distinguish empty from missing data
27+
}
28+
29+
func (c *Content) UnmarshalJSON(data []byte) error {
30+
type wireContent Content // for naive unmarshaling
31+
var c2 wireContent
32+
if err := json.Unmarshal(data, &c2); err != nil {
33+
return err
34+
}
35+
switch c2.Type {
36+
case "text", "image", "audio", "resource":
37+
default:
38+
return fmt.Errorf("unrecognized content type %s", c.Type)
39+
}
40+
*c = Content(c2)
41+
return nil
42+
}

‎internal/mcp/internal/protocol/generate.go

+9-3
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ var declarations = config{
9191
"Tools": {Name: "ToolCapabilities"},
9292
},
9393
},
94-
"TextContent": {Name: "TextContent"},
9594
"Tool": {
9695
Name: "Tool",
9796
Fields: config{"InputSchema": {Substitute: "*jsonschema.Schema"}},
@@ -249,8 +248,15 @@ func writeType(w io.Writer, config *typeConfig, def *jsonschema.Schema, named ma
249248
}
250249

251250
if def.Type == "" {
252-
// E.g. union types.
253-
fmt.Fprintf(w, "json.RawMessage")
251+
// special case: recognize Content
252+
if slices.ContainsFunc(def.AnyOf, func(s *jsonschema.Schema) bool {
253+
return s.Ref == "#/definitions/TextContent"
254+
}) {
255+
fmt.Fprintf(w, "Content")
256+
} else {
257+
// E.g. union types.
258+
fmt.Fprintf(w, "json.RawMessage")
259+
}
254260
} else {
255261
switch def.Type {
256262
case "array":

0 commit comments

Comments
 (0)