feat(mcp): build flagset from mcp input schema to parse args (#1215)

- parse input schema to build flagset for mcp command

Author: burmudar

cmd/src/mcp.go

@@ -0,0 +1,115 @@
+package main
+
+import (
+	"flag"
+	"fmt"
+
+	"github.com/sourcegraph/src-cli/internal/mcp"
+)
+
+func init() {
+	flagSet := flag.NewFlagSet("mcp", flag.ExitOnError)
+	commands = append(commands, &command{
+		flagSet: flagSet,
+		handler: mcpMain,
+	})
+}
+func mcpMain(args []string) error {
+	fmt.Println("NOTE: This command is still experimental")
+	tools, err := mcp.LoadToolDefinitions()
+	if err != nil {
+		return err
+	}
+
+	subcmd := args[0]
+	if subcmd == "list-tools" {
+		fmt.Println("The following tools are available:")
+		for name := range tools {
+			fmt.Printf(" • %s\n", name)
+		}
+		fmt.Println("\nUSAGE:")
+		fmt.Printf(" • Invoke a tool\n")
+		fmt.Printf("     src mcp <tool-name> <flags>\n")
+		fmt.Printf("\n • View the Input / Output Schema of a tool\n")
+		fmt.Printf("     src mcp <tool-name> schema\n")
+		fmt.Printf("\n • List the available flags of a tool\n")
+		fmt.Printf("     src mcp <tool-name> -h\n")
+		fmt.Printf("\n • View the Input / Output Schema of a tool\n")
+		fmt.Printf("     src mcp <tool-name> schema\n")
+		return nil
+	}
+
+	tool, ok := tools[subcmd]
+	if !ok {
+		return fmt.Errorf("tool definition for %q not found - run src mcp list-tools to see a list of available tools", subcmd)
+	}
+
+	flagArgs := args[1:] // skip subcommand name
+	if len(args) > 1 && args[1] == "schema" {
+		return printSchemas(tool)
+	}
+
+	flags, vars, err := mcp.BuildArgFlagSet(tool)
+	if err != nil {
+		return err
+	}
+	if err := flags.Parse(flagArgs); err != nil {
+		return err
+	}
+	mcp.DerefFlagValues(vars)
+
+	if err := validateToolArgs(tool.InputSchema, args, vars); err != nil {
+		return err
+	}
+
+	apiClient := cfg.apiClient(nil, flags.Output())
+	return handleMcpTool(context.Background(), apiClient, tool, vars)
+}
+
+func printSchemas(tool *mcp.ToolDef) error {
+	input, err := json.MarshalIndent(tool.InputSchema, "", " ")
+	if err != nil {
+		return err
+	}
+	output, err := json.MarshalIndent(tool.OutputSchema, "", " ")
+	if err != nil {
+		return err
+	}
+
+	fmt.Printf("Input:\n%v\nOutput:\n%v\n", string(input), string(output))
+	return nil
+}
+
+func validateToolArgs(inputSchema mcp.SchemaObject, args []string, vars map[string]any) error {
+	for _, reqName := range inputSchema.Required {
+		if vars[reqName] == nil {
+			return errors.Newf("no value provided for required flag --%s", reqName)
+		}
+	}
+
+	if len(args) < len(inputSchema.Required) {
+		return errors.Newf("not enough arguments provided - the following flags are required:\n%s", strings.Join(inputSchema.Required, "\n"))
+	}
+
+	return nil
+}
+
+func handleMcpTool(ctx context.Context, client api.Client, tool *mcp.ToolDef, vars map[string]any) error {
+	resp, err := mcp.DoToolRequest(ctx, client, tool, vars)
+	if err != nil {
+		return err
+	}
+
+	result, err := mcp.DecodeToolResponse(resp)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	output, err := json.MarshalIndent(result, "", "  ")
+	if err != nil {
+		return err
+	}
+	fmt.Println(string(output))
+	return nil
+}

internal/mcp/mcp_args.go

@@ -0,0 +1,94 @@
+package mcp
+
+import (
+	"flag"
+	"fmt"
+	"reflect"
+	"strings"
+
+	"github.com/sourcegraph/sourcegraph/lib/errors"
+)
+
+var _ flag.Value = (*strSliceFlag)(nil)
+
+type strSliceFlag struct {
+	vals []string
+}
+
+func (s *strSliceFlag) Set(v string) error {
+	s.vals = append(s.vals, v)
+	return nil
+}
+
+func (s *strSliceFlag) String() string {
+	return strings.Join(s.vals, ",")
+}
+
+func DerefFlagValues(vars map[string]any) {
+	for k, v := range vars {
+		rfl := reflect.ValueOf(v)
+		if rfl.Kind() == reflect.Pointer {
+			vv := rfl.Elem().Interface()
+			if slice, ok := vv.(strSliceFlag); ok {
+				vv = slice.vals
+			}
+			if isNil(vv) {
+				delete(vars, k)
+			} else {
+				vars[k] = vv
+			}
+		}
+	}
+}
+
+func isNil(v any) bool {
+	if v == nil {
+		return true
+	}
+	rv := reflect.ValueOf(v)
+	switch rv.Kind() {
+	case reflect.Slice, reflect.Map, reflect.Pointer, reflect.Interface:
+		return rv.IsNil()
+	default:
+		return false
+	}
+}
+
+func BuildArgFlagSet(tool *ToolDef) (*flag.FlagSet, map[string]any, error) {
+	if tool == nil {
+		return nil, nil, errors.New("cannot build flagset on nil Tool Definition")
+	}
+	fs := flag.NewFlagSet(tool.Name, flag.ContinueOnError)
+	flagVars := map[string]any{}
+
+	for name, pVal := range tool.InputSchema.Properties {
+		switch pv := pVal.(type) {
+		case *SchemaPrimitive:
+			switch pv.Type {
+			case "integer":
+				dst := fs.Int(name, 0, pv.Description)
+				flagVars[name] = dst
+
+			case "boolean":
+				dst := fs.Bool(name, false, pv.Description)
+				flagVars[name] = dst
+			case "string":
+				dst := fs.String(name, "", pv.Description)
+				flagVars[name] = dst
+			default:
+				return nil, nil, fmt.Errorf("unknown schema primitive kind %q", pv.Type)
+
+			}
+		case *SchemaArray:
+			strSlice := new(strSliceFlag)
+			fs.Var(strSlice, name, pv.Description)
+			flagVars[name] = strSlice
+		case *SchemaObject:
+			// TODO(burmudar): we can support SchemaObject as part of stdin echo '{ stuff }' | sg mcp commit-search
+			// not supported yet
+			// Also support sg mcp commit-search --json '{ stuff }'
+		}
+	}
+
+	return fs, flagVars, nil
+}

internal/mcp/mcp_args_test.go

@@ -0,0 +1,97 @@
+package mcp
+
+import (
+	"testing"
+)
+
+func TestFlagSetParse(t *testing.T) {
+	toolJSON := []byte(`{
+	  "tools": [
+		{
+		  "name": "sg_test_tool",
+		  "description": "test description",
+		  "inputSchema": {
+			"type": "object",
+			"$schema": "https://localhost/schema-draft/2025-07",
+			"required": ["values"],
+			"properties": {
+			  "repos": {
+				"type": "array",
+				"items": {
+				  "type": "string"
+				}
+			  },
+			  "tag": {
+				"type": "string",
+				"items": true
+			  },
+			  "count": {
+				"type": "integer"
+	          },
+	          "boolFlag": {
+				"type": "boolean"
+	          }
+			}
+		  },
+		  "outputSchema": {
+			"type": "object",
+			"$schema": "https://localhost/schema-draft/2025-07",
+			"properties": {
+			  "result": { "type": "string" }
+			}
+		  }
+		}
+	  ]
+	}`)
+
+	defs, err := loadToolDefinitions(toolJSON)
+	if err != nil {
+		t.Fatalf("failed to load tool json: %v", err)
+	}
+
+	flagSet, vars, err := BuildArgFlagSet(defs["test-tool"])
+	if err != nil {
+		t.Fatalf("failed to build flagset from mcp tool definition: %v", err)
+	}
+
+	if len(vars) == 0 {
+		t.Fatalf("vars from buildArgFlagSet should not be empty")
+	}
+
+	args := []string{"-repos=A", "-repos=B", "-count=10", "-boolFlag", "-tag=testTag"}
+
+	if err := flagSet.Parse(args); err != nil {
+		t.Fatalf("flagset parsing failed: %v", err)
+	}
+	DerefFlagValues(vars)
+
+	if v, ok := vars["repos"].([]string); ok {
+		if len(v) != 2 {
+			t.Fatalf("expected flag 'repos' values to have length %d but got %d", 2, len(v))
+		}
+	} else {
+		t.Fatalf("expected flag 'repos' to have type of []string but got %T", v)
+	}
+	if v, ok := vars["tag"].(string); ok {
+		if v != "testTag" {
+			t.Fatalf("expected flag 'tag' values to have value %q but got %q", "testTag", v)
+		}
+	} else {
+		t.Fatalf("expected flag 'tag' to have type of string but got %T", v)
+	}
+	if v, ok := vars["count"].(int); ok {
+		if v != 10 {
+			t.Fatalf("expected flag 'count' values to have value %d but got %d", 10, v)
+		}
+	} else {
+		t.Fatalf("expected flag 'count' to have type of int but got %T", v)
+	}
+	if v, ok := vars["boolFlag"].(bool); ok {
+		if v != true {
+			t.Fatalf("expected flag 'boolFlag' values to have value %v but got %v", true, v)
+		}
+	} else {
+		t.Fatalf("expected flag 'boolFlag' to have type of bool but got %T", v)
+	}
+
+}