feat(mcpgen,mcpserver): @mcpToolArg for tool-only arguments

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 {

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}},

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)

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

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

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",

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