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