internal/mcp: reorganize the tool API, keeping generics

CallToolParams.Arguments is any.

Clients can send whatever they want.

Tool authors who create their own ToolHandler (bypassing NewTool)
will receive the arguments as a map[string]any which has been
validated against the input schema.

Tool authors who call NewTool can choose the type that the arguments
unmarshal into.

There is no double-unmarshaling and no downcasting.

There is no way for tool authors to avoid the unmarshal, and no way to
avoid the validation unless the omit an input schema, which is probably
bad for LLMs. If tool authors want this optimization, we can provide
it later.

Change-Id: Ieb201e56c7fb8a23bb1a93b27e946630ec9e79ad
Reviewed-on: https://go-review.googlesource.com/c/tools/+/679595
TryBot-Bypass: Jonathan Amsterdam <jba@google.com>
Auto-Submit: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Robert Findley <rfindley@google.com>

Author: jba

gopls/internal/mcp/context.go

@@ -32,7 +32,7 @@ type ContextParams struct {
 	Location protocol.Location `json:"location"`
 }
 
-func contextHandler(ctx context.Context, session *cache.Session, params *mcp.CallToolParams[ContextParams]) (*mcp.CallToolResult, error) {
+func contextHandler(ctx context.Context, session *cache.Session, params *mcp.CallToolParamsFor[ContextParams]) (*mcp.CallToolResultFor[struct{}], error) {
 	fh, snapshot, release, err := session.FileOf(ctx, params.Arguments.Location.URI)
 	if err != nil {
 		return nil, err
@@ -137,7 +137,7 @@ func contextHandler(ctx context.Context, session *cache.Session, params *mcp.Cal
 		}
 	}
 
-	return &mcp.CallToolResult{
+	return &mcp.CallToolResultFor[struct{}]{
 		Content: []*mcp.Content{
 			mcp.NewTextContent(result.String()),
 		},

gopls/internal/mcp/mcp.go

@@ -117,7 +117,7 @@ func newServer(session *cache.Session) *mcp.Server {
 		mcp.NewTool(
 			"context",
 			"Provide context for a region within a Go file",
-			func(ctx context.Context, _ *mcp.ServerSession, request *mcp.CallToolParams[ContextParams]) (*mcp.CallToolResult, error) {
+			func(ctx context.Context, _ *mcp.ServerSession, request *mcp.CallToolParamsFor[ContextParams]) (*mcp.CallToolResultFor[struct{}], error) {
 				return contextHandler(ctx, session, request)
 			},
 			mcp.Input(

gopls/internal/test/marker/marker_test.go

@@ -297,7 +297,7 @@ func Test(t *testing.T) {
 				t.Errorf("formatTest: %v", err)
 			} else if *update {
 				filename := filepath.Join(dir, test.name)
-				if err := os.WriteFile(filename, formatted, 0644); err != nil {
+				if err := os.WriteFile(filename, formatted, 0o644); err != nil {
 					t.Error(err)
 				}
 			} else if !t.Failed() {
@@ -1414,7 +1414,6 @@ func (sm stringMatcher) check(mark marker, got string) {
 		if !sm.pattern.MatchString(got) {
 			mark.errorf("got %q, does not match pattern %#q", got, sm.pattern)
 		}
-
 	} else if !strings.Contains(got, sm.substr) {
 		// Content must contain the expected substring.
 		mark.errorf("got %q, want substring %q", got, sm.substr)
@@ -1428,7 +1427,6 @@ func checkChangedFiles(mark marker, changed map[string][]byte, golden *Golden) {
 		if want, ok := golden.Get(mark.T(), filename, got); !ok {
 			mark.errorf("%s: unexpected change to file %s; got:\n%s",
 				mark.note.Name, filename, got)
-
 		} else if string(got) != string(want) {
 			mark.errorf("%s: wrong file content for %s: got:\n%s\nwant:\n%s\ndiff:\n%s",
 				mark.note.Name, filename, got, want,
@@ -1481,7 +1479,6 @@ func checkDiffs(mark marker, changed map[string][]byte, golden *Golden) {
 		if want, ok := golden.Get(mark.T(), filename, []byte(got)); !ok {
 			mark.errorf("%s: unexpected change to file %s; got diff:\n%s",
 				mark.note.Name, filename, got)
-
 		} else if got != string(want) {
 			mark.errorf("%s: wrong diff for %s:\n\ngot:\n%s\n\nwant:\n%s\n",
 				mark.note.Name, filename, got, want)
@@ -1680,7 +1677,6 @@ func acceptCompletionMarker(mark marker, src protocol.Location, label string, go
 	filename := mark.path()
 	mapper := mark.mapper()
 	patched, _, err := protocol.ApplyEdits(mapper, append([]protocol.TextEdit{edit}, selected.AdditionalTextEdits...))
-
 	if err != nil {
 		mark.errorf("ApplyProtocolEdits failed: %v", err)
 		return
@@ -1799,6 +1795,7 @@ func highlightLocationMarker(mark marker, loc protocol.Location, kindName expect
 		Kind:  kind,
 	}
 }
+
 func sortDocumentHighlights(s []protocol.DocumentHighlight) {
 	sort.Slice(s, func(i, j int) bool {
 		return protocol.CompareRange(s[i].Range, s[j].Range) < 0
@@ -2452,7 +2449,7 @@ func mcpToolMarker(mark marker, tool string, rawArgs string, loc protocol.Locati
 	// TODO(hxjiang): Make the "location" key configurable.
 	args["location"] = loc
 
-	res, err := mcp.CallTool(mark.ctx(), mark.run.env.MCPSession, &mcp.CallToolParams[map[string]any]{
+	res, err := mark.run.env.MCPSession.CallTool(mark.ctx(), &mcp.CallToolParams{
 		Name:      tool,
 		Arguments: args,
 	})

internal/mcp/README.md

@@ -50,11 +50,11 @@ func main() {
 	}
 	defer session.Close()
 	// Call a tool on the server.
-	params := &mcp.CallToolParams[map[string]any]{
+	params := &mcp.CallToolParams{
 		Name:      "greet",
 		Arguments: map[string]any{"name": "you"},
 	}
-	if res, err := mcp.CallTool(ctx, session, params); err != nil {
+	if res, err := session.CallTool(ctx, params); err != nil {
 		log.Printf("CallTool failed: %v", err)
 	} else {
 		if res.IsError {
@@ -82,8 +82,8 @@ type HiParams struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[HiParams]) (*mcp.CallToolResult, error) {
-	return &mcp.CallToolResult{
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[HiParams]) (*mcp.CallToolResultFor[string], error) {
+	return &mcp.CallToolResultFor[string]{
 		Content: []*mcp.Content{mcp.NewTextContent("Hi " + params.Name)},
 	}, nil
 }

internal/mcp/client.go

@@ -6,8 +6,6 @@ package mcp
 
 import (
 	"context"
-	"encoding/json"
-	"fmt"
 	"iter"
 	"slices"
 	"sync"
@@ -288,34 +286,13 @@ func (cs *ClientSession) ListTools(ctx context.Context, params *ListToolsParams)
 }
 
 // CallTool calls the tool with the given name and arguments.
-// Pass a [CallToolOptions] to provide additional request fields.
-func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams[json.RawMessage]) (*CallToolResult, error) {
-	return handleSend[*CallToolResult](ctx, cs, methodCallTool, params)
-}
-
-// CallTool is a helper to call a tool with any argument type. It returns an
-// error if params.Arguments fails to marshal to JSON.
-func CallTool[TArgs any](ctx context.Context, cs *ClientSession, params *CallToolParams[TArgs]) (*CallToolResult, error) {
-	wireParams, err := toWireParams(params)
-	if err != nil {
-		return nil, err
+// The arguments can be any value that marshals into a JSON object.
+func (cs *ClientSession) CallTool(ctx context.Context, params *CallToolParams) (*CallToolResult, error) {
+	if params.Arguments == nil {
+		// Avoid sending nil over the wire.
+		params.Arguments = map[string]any{}
 	}
-	return cs.CallTool(ctx, wireParams)
-}
-
-func toWireParams[TArgs any](params *CallToolParams[TArgs]) (*CallToolParams[json.RawMessage], error) {
-	data, err := json.Marshal(params.Arguments)
-	if err != nil {
-		return nil, fmt.Errorf("failed to marshal arguments: %v", err)
-	}
-	// The field mapping here must be kept up to date with the CallToolParams.
-	// This is partially enforced by TestToWireParams, which verifies that all
-	// comparable fields are mapped.
-	return &CallToolParams[json.RawMessage]{
-		Meta:      params.Meta,
-		Name:      params.Name,
-		Arguments: data,
-	}, nil
+	return handleSend[*CallToolResult](ctx, cs, methodCallTool, params)
 }
 
 func (cs *ClientSession) SetLevel(ctx context.Context, params *SetLevelParams) error {

internal/mcp/client_test.go

@@ -7,7 +7,6 @@ package mcp
 import (
 	"context"
 	"fmt"
-	"reflect"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -200,28 +199,3 @@ func TestClientPaginateVariousPageSizes(t *testing.T) {
 		})
 	}
 }
-
-func TestToWireParams(t *testing.T) {
-	// This test verifies that toWireParams maps all fields.
-	// The Meta and Arguments fields are not comparable, so can't be checked by
-	// this simple test. However, this test will fail if new fields are added,
-	// and not handled by toWireParams.
-	params := &CallToolParams[struct{}]{
-		Name: "tool",
-	}
-	wireParams, err := toWireParams(params)
-	if err != nil {
-		t.Fatal(err)
-	}
-	v := reflect.ValueOf(wireParams).Elem()
-	for i := range v.Type().NumField() {
-		f := v.Type().Field(i)
-		if f.Name == "Meta" || f.Name == "Arguments" {
-			continue // not comparable
-		}
-		fv := v.Field(i)
-		if fv.Interface() == reflect.Zero(f.Type).Interface() {
-			t.Fatalf("toWireParams: unmapped field %q", f.Name)
-		}
-	}
-}

internal/mcp/cmd_test.go

@@ -56,7 +56,7 @@ func TestCmdTransport(t *testing.T) {
 	if err != nil {
 		log.Fatal(err)
 	}
-	got, err := mcp.CallTool(ctx, session, &mcp.CallToolParams[map[string]any]{
+	got, err := session.CallTool(ctx, &mcp.CallToolParams{
 		Name:      "greet",
 		Arguments: map[string]any{"name": "user"},
 	})

internal/mcp/examples/hello/main.go

@@ -21,8 +21,8 @@ type HiArgs struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParams[HiArgs]) (*mcp.CallToolResult, error) {
-	return &mcp.CallToolResult{
+func SayHi(ctx context.Context, ss *mcp.ServerSession, params *mcp.CallToolParamsFor[HiArgs]) (*mcp.CallToolResultFor[any], error) {
+	return &mcp.CallToolResultFor[any]{
 		Content: []*mcp.Content{
 			mcp.NewTextContent("Hi " + params.Name),
 		},

internal/mcp/examples/sse/main.go

@@ -19,8 +19,8 @@ type SayHiParams struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParams[SayHiParams]) (*mcp.CallToolResult, error) {
-	return &mcp.CallToolResult{
+func SayHi(ctx context.Context, cc *mcp.ServerSession, params *mcp.CallToolParamsFor[SayHiParams]) (*mcp.CallToolResultFor[any], error) {
+	return &mcp.CallToolResultFor[any]{
 		Content: []*mcp.Content{
 			mcp.NewTextContent("Hi " + params.Name),
 		},

internal/mcp/features_test.go

@@ -18,8 +18,8 @@ type SayHiParams struct {
 	Name string `json:"name"`
 }
 
-func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParams[SayHiParams]) (*CallToolResult, error) {
-	return &CallToolResult{
+func SayHi(ctx context.Context, cc *ServerSession, params *CallToolParamsFor[SayHiParams]) (*CallToolResultFor[any], error) {
+	return &CallToolResultFor[any]{
 		Content: []*Content{
 			NewTextContent("Hi " + params.Name),
 		},
@@ -71,7 +71,6 @@ func TestFeatureSetAbove(t *testing.T) {
 		got := slices.Collect(fs.above(tc.above))
 		if diff := cmp.Diff(got, tc.want, cmpopts.IgnoreUnexported(jsonschema.Schema{})); diff != "" {
 			t.Errorf("expected %v, got %v, (-want +got):\n%s", tc.want, got, diff)
-
 		}
 	}
 }