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>

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.