From d654fab59450b0cd280efdb99f21aa587485b3aa Mon Sep 17 00:00:00 2001 From: zer0stars <74260741+zer0stars@users.noreply.github.com> Date: Thu, 23 Apr 2026 15:57:01 -0400 Subject: [PATCH] feat(mcpgen,mcpserver): @mcpToolArg for tool-only arguments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shortcut tools that need an argument-driven selection (e.g. signals time-series where each {name, agg} pair becomes a sub-field) currently require that argument to exist on the GraphQL field itself — which forces a schema change on queries that have been stable for years. Add a repeatable @mcpToolArg directive that declares arguments belonging only to the MCP tool. They feed SelectionTemplate via the same args map, and mcpserver strips them from the variables map before the GraphQL executor is called, so the underlying field never sees them. Usage: signals(tokenId: Int!, interval: String!, from: Time!, to: Time!, filter: SignalFilter): [SignalAggregations!] @mcpTool( name: "get_signals_time_series" description: "..." selection: "timestamp{{range .signalRequests}} {{.name}}(agg: {{.agg}}){{end}}" ) @mcpToolArg( name: "signalRequests" type: "[SignalAggregationRequest!]!" description: "..." ) The GraphQL field keeps its existing arguments and return type. The MCP tool gets `signalRequests` as a required input. At call time the template renders it into the selection, and `signalRequests` is dropped from the variables map before executing. ArgDefinition gets a `ToolOnly bool` field; mcpgen emits it alongside the other ArgDefinition fields and reuses the existing JSON-schema mapping via a small ad-hoc type-string parser so the `type:` argument of the directive can name any loaded GraphQL type (e.g. "[SignalAggregationRequest!]!", "[String!]!", "Int"). --- cmd/mcpgen/condensed.go | 39 ++++++++++++++ cmd/mcpgen/generate.go | 54 +++++++++++++++++++- cmd/mcpgen/generate_test.go | 59 ++++++++++++++++++++++ cmd/mcpgen/testdata/tool_only_arg.graphqls | 25 +++++++++ pkg/mcpserver/mcpserver.go | 5 ++ pkg/mcpserver/mcpserver_test.go | 52 +++++++++++++++++++ pkg/mcpserver/tools.go | 20 +++++++- 7 files changed, 252 insertions(+), 2 deletions(-) create mode 100644 cmd/mcpgen/testdata/tool_only_arg.graphqls diff --git a/cmd/mcpgen/condensed.go b/cmd/mcpgen/condensed.go index 8317b49..a6e9aa0 100644 --- a/cmd/mcpgen/condensed.go +++ b/cmd/mcpgen/condensed.go @@ -256,6 +256,45 @@ func extractTools(schema *ast.Schema, prefix string) ([]mcpserver.ToolDefinition }) } + // @mcpToolArg adds tool-only arguments that exist solely to feed the + // SelectionTemplate. They are stripped from GraphQL variables at call + // time and do not have to appear on the underlying field. + for _, d := range field.Directives { + if d.Name != "mcpToolArg" { + continue + } + argNameV := d.Arguments.ForName("name") + typeV := d.Arguments.ForName("type") + if argNameV == nil || typeV == nil { + return nil, fmt.Errorf("field %s: @mcpToolArg directive missing required argument(s)", field.Name) + } + parsedType, err := parseTypeString(typeV.Value.Raw) + if err != nil { + return nil, fmt.Errorf("field %s: @mcpToolArg(name: %q) invalid type: %w", field.Name, argNameV.Value.Raw, err) + } + argDesc := "" + if descV := d.Arguments.ForName("description"); descV != nil { + argDesc = descV.Value.Raw + } + if argDesc == "" { + if parsedType.NonNull { + argDesc = fmt.Sprintf("%s (%s, required)", argNameV.Value.Raw, parsedType.String()) + } else { + argDesc = fmt.Sprintf("%s (%s, optional)", argNameV.Value.Raw, parsedType.String()) + } + } + jsonType, itemsType := mapGraphQLType(parsedType, schema) + args = append(args, mcpserver.ArgDefinition{ + Name: argNameV.Value.Raw, + Type: jsonType, + Description: argDesc, + Required: parsedType.NonNull, + ItemsType: itemsType, + EnumValues: enumValues(parsedType, schema), + ToolOnly: true, + }) + } + selectionTemplate := "" querySelection := selection if isTemplated { diff --git a/cmd/mcpgen/generate.go b/cmd/mcpgen/generate.go index b0ed959..42b33e0 100644 --- a/cmd/mcpgen/generate.go +++ b/cmd/mcpgen/generate.go @@ -65,6 +65,58 @@ func isTemplatedSelection(selection string) bool { return strings.Contains(selection, "{{") } +// parseTypeString parses a GraphQL type reference (e.g. "Int", "String!", +// "[Foo!]!") into an *ast.Type. Used to resolve the `type:` argument of the +// @mcpToolArg directive so mcpgen can build the same JSON schema it builds +// for real GraphQL field arguments. +func parseTypeString(s string) (*ast.Type, error) { + t, rest, err := parseTypeInner(s) + if err != nil { + return nil, err + } + if strings.TrimSpace(rest) != "" { + return nil, fmt.Errorf("unexpected trailing text %q", strings.TrimSpace(rest)) + } + return t, nil +} + +func parseTypeInner(s string) (*ast.Type, string, error) { + s = strings.TrimLeft(s, " \t") + if s == "" { + return nil, "", fmt.Errorf("empty type") + } + var t *ast.Type + if s[0] == '[' { + inner, rest, err := parseTypeInner(s[1:]) + if err != nil { + return nil, "", err + } + rest = strings.TrimLeft(rest, " \t") + if rest == "" || rest[0] != ']' { + return nil, "", fmt.Errorf("missing closing ']' in type") + } + rest = rest[1:] + t = &ast.Type{Elem: inner} + s = rest + } else { + i := 0 + for i < len(s) && isIdentChar(rune(s[i])) { + i++ + } + if i == 0 { + return nil, "", fmt.Errorf("expected identifier at %q", s) + } + t = &ast.Type{NamedType: s[:i]} + s = s[i:] + } + s = strings.TrimLeft(s, " \t") + if strings.HasPrefix(s, "!") { + t.NonNull = true + s = s[1:] + } + return t, s, nil +} + // validateSelection checks that top-level field names in the selection exist on the type. func validateSelection(selection string, typeDef *ast.Definition) error { fields := extractTopLevelFields(selection) @@ -239,7 +291,7 @@ var MCPTools = []mcpserver.ToolDefinition{ Description: {{printf "%q" .Description}}, Args: []mcpserver.ArgDefinition{ {{- range .Args}} - {Name: {{printf "%q" .Name}}, Type: {{printf "%q" .Type}}, Description: {{printf "%q" .Description}}, Required: {{.Required}}, ItemsType: {{printf "%q" .ItemsType}}{{if .EnumValues}}, EnumValues: []string{ {{- range $i, $v := .EnumValues}}{{if $i}}, {{end}}{{printf "%q" $v}}{{end -}} }{{end}}}, + {Name: {{printf "%q" .Name}}, Type: {{printf "%q" .Type}}, Description: {{printf "%q" .Description}}, Required: {{.Required}}, ItemsType: {{printf "%q" .ItemsType}}{{if .EnumValues}}, EnumValues: []string{ {{- range $i, $v := .EnumValues}}{{if $i}}, {{end}}{{printf "%q" $v}}{{end -}} }{{end}}{{if .ToolOnly}}, ToolOnly: true{{end}}}, {{- end}} }, Query: {{printf "%q" .Query}}, diff --git a/cmd/mcpgen/generate_test.go b/cmd/mcpgen/generate_test.go index f572740..70e3295 100644 --- a/cmd/mcpgen/generate_test.go +++ b/cmd/mcpgen/generate_test.go @@ -236,6 +236,65 @@ func TestTemplatedSelectionGeneratedOutput(t *testing.T) { assert.Contains(t, output, "__MCPGEN_SELECTION__") } +func TestToolOnlyArg(t *testing.T) { + tools := loadTools(t, []string{"testdata/tool_only_arg.graphqls"}, "") + require.Len(t, tools, 1) + + tool := tools[0] + require.Len(t, tool.Args, 2, "tokenId (field arg) + signalRequests (tool-only)") + + assert.Equal(t, "tokenId", tool.Args[0].Name) + assert.False(t, tool.Args[0].ToolOnly, "tokenId is a real field arg") + + signalReqs := tool.Args[1] + assert.Equal(t, "signalRequests", signalReqs.Name) + assert.True(t, signalReqs.ToolOnly, "signalRequests is tool-only") + assert.Equal(t, "array", signalReqs.Type) + assert.Equal(t, "object", signalReqs.ItemsType) + assert.True(t, signalReqs.Required) + assert.Equal(t, "List of signal/aggregation pairs to render into the selection", signalReqs.Description) + + assert.NotContains(t, tool.Query, "$signalRequests", "tool-only arg must not become a GraphQL variable") + assert.Contains(t, tool.Query, "__MCPGEN_SELECTION__") + assert.Contains(t, tool.SelectionTemplate, ".signalRequests") +} + +func TestToolOnlyArgGeneratedOutput(t *testing.T) { + tools := loadTools(t, []string{"testdata/tool_only_arg.graphqls"}, "test") + output, err := generateGoFile("graph", tools, "") + require.NoError(t, err) + assert.Contains(t, output, "ToolOnly: true") +} + +func TestParseTypeString(t *testing.T) { + cases := []struct { + in string + named string + nonNull bool + listElem string + }{ + {"Int", "Int", false, ""}, + {"Int!", "Int", true, ""}, + {"[Foo!]", "", false, "Foo"}, + {"[Foo!]!", "", true, "Foo"}, + } + for _, c := range cases { + t.Run(c.in, func(t *testing.T) { + got, err := parseTypeString(c.in) + require.NoError(t, err) + if c.listElem != "" { + require.NotNil(t, got.Elem) + assert.Equal(t, c.listElem, got.Elem.NamedType) + } else { + assert.Equal(t, c.named, got.NamedType) + } + assert.Equal(t, c.nonNull, got.NonNull) + }) + } + _, err := parseTypeString("[Foo") + require.Error(t, err) +} + func TestMapGraphQLTypeAllScalars(t *testing.T) { tools := loadTools(t, []string{"testdata/all_types.graphqls"}, "") require.Len(t, tools, 1) diff --git a/cmd/mcpgen/testdata/tool_only_arg.graphqls b/cmd/mcpgen/testdata/tool_only_arg.graphqls new file mode 100644 index 0000000..1c9038d --- /dev/null +++ b/cmd/mcpgen/testdata/tool_only_arg.graphqls @@ -0,0 +1,25 @@ +directive @mcpTool(name: String!, description: String!, selection: String!, readOnly: Boolean = true) on FIELD_DEFINITION +directive @mcpToolArg(name: String!, type: String!, description: String) repeatable on FIELD_DEFINITION + +type Query { + signals(tokenId: Int!): [Bucket!] + @mcpTool( + name: "get_signals" + description: "Aggregated signals with dynamic selection" + selection: "timestamp{{range .signalRequests}} {{.name}}(agg: {{.agg}}){{end}}" + ) + @mcpToolArg( + name: "signalRequests" + type: "[SignalRequest!]!" + description: "List of signal/aggregation pairs to render into the selection" + ) +} + +type Bucket { + timestamp: String! +} + +input SignalRequest { + name: String! + agg: String! +} diff --git a/pkg/mcpserver/mcpserver.go b/pkg/mcpserver/mcpserver.go index e3b76cd..544aa3a 100644 --- a/pkg/mcpserver/mcpserver.go +++ b/pkg/mcpserver/mcpserver.go @@ -23,6 +23,11 @@ type ArgDefinition struct { Required bool ItemsType string // JSON Schema type for array elements EnumValues []string // Allowed values for enum types + // ToolOnly marks an argument that exists on the MCP tool for the sake of + // SelectionTemplate rendering but is not a real argument on the underlying + // GraphQL field. These are stripped from the variables map before the + // GraphQL executor is called. Emitted by mcpgen from @mcpToolArg directives. + ToolOnly bool } // SelectionPlaceholder is the marker mcpgen inserts into Query where a diff --git a/pkg/mcpserver/mcpserver_test.go b/pkg/mcpserver/mcpserver_test.go index 6af4efb..a00fa17 100644 --- a/pkg/mcpserver/mcpserver_test.go +++ b/pkg/mcpserver/mcpserver_test.go @@ -739,6 +739,58 @@ func TestSelectionTemplateRendering(t *testing.T) { assert.NotContains(t, capturedQuery, "{{", "template markers should not leak into executed query") } +func TestToolOnlyArgStrippedFromVariables(t *testing.T) { + var captured struct { + query string + vars map[string]any + } + exec := &mockExecutor{ + fn: func(ctx context.Context, query string, variables map[string]any) ([]byte, error) { + captured.query = query + captured.vars = variables + return []byte(`{"data":{}}`), nil + }, + } + + tool := ToolDefinition{ + Name: "get_signals", + Description: "Aggregated signals with dynamic selection", + Args: []ArgDefinition{ + {Name: "tokenId", Type: "integer", Required: true}, + {Name: "signalRequests", Type: "array", ItemsType: "object", Required: true, ToolOnly: true}, + }, + Query: `query($tokenId: Int!) { signals(tokenId: $tokenId) { __MCPGEN_SELECTION__ } }`, + SelectionTemplate: "timestamp{{range .signalRequests}} {{.name}}(agg: {{.agg}}){{end}}", + } + + mcpServer := mcp.NewServer(&mcp.Implementation{Name: "test", Version: "0.1.0"}, nil) + require.NoError(t, registerShortcutTools(mcpServer, exec, []ToolDefinition{tool}, nil)) + + serverTransport, clientTransport := mcp.NewInMemoryTransports() + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + go func() { _ = mcpServer.Run(ctx, serverTransport) }() + + client := mcp.NewClient(&mcp.Implementation{Name: "test-client", Version: "0.1.0"}, nil) + session, err := client.Connect(ctx, clientTransport, nil) + require.NoError(t, err) + + _, err = session.CallTool(ctx, &mcp.CallToolParams{ + Name: "get_signals", + Arguments: map[string]any{ + "tokenId": 42, + "signalRequests": []any{ + map[string]any{"name": "speed", "agg": "AVG"}, + }, + }, + }) + require.NoError(t, err) + + assert.Contains(t, captured.query, "timestamp speed(agg: AVG)", "template should render with the tool-only arg") + assert.Contains(t, captured.vars, "tokenId", "real field args must still go to the executor") + assert.NotContains(t, captured.vars, "signalRequests", "tool-only args must not be passed as GraphQL variables") +} + func TestSelectionTemplateInvalidRejectedAtRegistration(t *testing.T) { tool := ToolDefinition{ Name: "bad_tool", diff --git a/pkg/mcpserver/tools.go b/pkg/mcpserver/tools.go index ae57a6d..a2a3e8a 100644 --- a/pkg/mcpserver/tools.go +++ b/pkg/mcpserver/tools.go @@ -169,6 +169,14 @@ func registerShortcutTools(server *mcp.Server, exec GraphQLExecutor, tools []Too selTmpl = tmpl } + hasToolOnly := false + for _, a := range tool.Args { + if a.ToolOnly { + hasToolOnly = true + break + } + } + mcp.AddTool(server, &mcp.Tool{ Name: tool.Name, Description: tool.Description, @@ -187,7 +195,17 @@ func registerShortcutTools(server *mcp.Server, exec GraphQLExecutor, tools []Too } query = strings.Replace(tool.Query, SelectionPlaceholder, buf.String(), 1) } - return executeTool(ctx, tool.Name, exec, query, args, logger) + gqlArgs := args + if hasToolOnly { + gqlArgs = make(map[string]any, len(args)) + for k, v := range args { + if def, ok := argDefs[k]; ok && def.ToolOnly { + continue + } + gqlArgs[k] = v + } + } + return executeTool(ctx, tool.Name, exec, query, gqlArgs, logger) }) } return nil