From bbd253ee4cbd80e0069a422348c283829bb231b4 Mon Sep 17 00:00:00 2001
From: Jonathan Amsterdam <jba@google.com>
Date: Wed, 14 May 2025 07:24:51 -0400
Subject: [PATCH] internal/mcp: use RawMessage for tool args

Change the type of CallToolParams.Arguments to json.RawMessage.

Change-Id: Ief16ef261b9f41566a1a16efc60a9382a8d414eb
Reviewed-on: https://go-review.googlesource.com/c/tools/+/672615
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>
---
 mcp/client.go   | 14 +++++---------
 mcp/generate.go | 11 +++++++++--
 mcp/protocol.go |  4 ++--
 mcp/tool.go     |  8 +-------
 4 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/mcp/client.go b/mcp/client.go
index 0a24930..d1ff374 100644
--- a/mcp/client.go
+++ b/mcp/client.go
@@ -181,18 +181,14 @@ func (c *Client) CallTool(ctx context.Context, name string, args map[string]any,
 			err = fmt.Errorf("calling tool %q: %w", name, err)
 		}
 	}()
-	argsJSON := make(map[string]json.RawMessage)
-	for name, arg := range args {
-		argJSON, err := json.Marshal(arg)
-		if err != nil {
-			return nil, fmt.Errorf("marshaling argument %s: %v", name, err)
-		}
-		argsJSON[name] = argJSON
-	}
 
+	data, err := json.Marshal(args)
+	if err != nil {
+		return nil, fmt.Errorf("marshaling arguments: %w", err)
+	}
 	params := &CallToolParams{
 		Name:      name,
-		Arguments: argsJSON,
+		Arguments: json.RawMessage(data),
 	}
 	return standardCall[CallToolResult](ctx, c.conn, "tools/call", params)
 }
diff --git a/mcp/generate.go b/mcp/generate.go
index 37089ab..6ac0a65 100644
--- a/mcp/generate.go
+++ b/mcp/generate.go
@@ -53,8 +53,15 @@ type config map[string]*typeConfig
 var declarations = config{
 	"Annotations": {},
 	"CallToolRequest": {
-		Name:   "-",
-		Fields: config{"Params": {Name: "CallToolParams"}},
+		Name: "-",
+		Fields: config{
+			"Params": {
+				Name: "CallToolParams",
+				Fields: config{
+					"Arguments": {Substitute: "json.RawMessage"},
+				},
+			},
+		},
 	},
 	"CallToolResult": {},
 	"CancelledNotification": {
diff --git a/mcp/protocol.go b/mcp/protocol.go
index 5b78a9e..2664ce3 100644
--- a/mcp/protocol.go
+++ b/mcp/protocol.go
@@ -29,8 +29,8 @@ type Annotations struct {
 }
 
 type CallToolParams struct {
-	Arguments map[string]json.RawMessage `json:"arguments,omitempty"`
-	Name      string                     `json:"name"`
+	Arguments json.RawMessage `json:"arguments,omitempty"`
+	Name      string          `json:"name"`
 }
 
 // The server's response to a tool call.
diff --git a/mcp/tool.go b/mcp/tool.go
index 2cc9a50..8dc5c5c 100644
--- a/mcp/tool.go
+++ b/mcp/tool.go
@@ -40,14 +40,8 @@ func NewTool[TReq any](name, description string, handler func(context.Context, *
 		panic(err)
 	}
 	wrapped := func(ctx context.Context, cc *ServerConnection, params *CallToolParams) (*CallToolResult, error) {
-		// For simplicity, just marshal and unmarshal the arguments.
-		// This could be avoided in the future.
-		rawArgs, err := json.Marshal(params.Arguments)
-		if err != nil {
-			return nil, err
-		}
 		var v TReq
-		if err := unmarshalSchema(rawArgs, schema, &v); err != nil {
+		if err := unmarshalSchema(params.Arguments, schema, &v); err != nil {
 			return nil, err
 		}
 		content, err := handler(ctx, cc, v)