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>

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