|
8 | 8 | "fmt"
|
9 | 9 | "io"
|
10 | 10 | "path"
|
| 11 | + "slices" |
11 | 12 | "strings"
|
12 | 13 |
|
13 | 14 | "golang.org/x/xerrors"
|
@@ -473,20 +474,26 @@ func (s *MCPServer) handleCallTool(req JSONRPC2Request) {
|
473 | 474 | // Convert the arguments map to command-line args
|
474 | 475 | var cmdArgs []string
|
475 | 476 |
|
476 |
| - // Check for positional arguments (using "_" as the key) |
477 |
| - if posArgs, ok := args["_"]; ok { |
478 |
| - switch val := posArgs.(type) { |
479 |
| - case string: |
480 |
| - cmdArgs = append(cmdArgs, val) |
481 |
| - case []any: |
482 |
| - for _, item := range val { |
483 |
| - cmdArgs = append(cmdArgs, fmt.Sprintf("%v", item)) |
| 477 | + // Check for positional arguments prefix with `argN__<name>` |
| 478 | + deleteKeys := make([]string, 0) |
| 479 | + for k, v := range args { |
| 480 | + if strings.HasPrefix(k, "arg") && len(k) > 4 && k[3] >= '0' && k[3] <= '9' { |
| 481 | + deleteKeys = append(deleteKeys, k) |
| 482 | + switch val := v.(type) { |
| 483 | + case string: |
| 484 | + cmdArgs = append(cmdArgs, val) |
| 485 | + case []any: |
| 486 | + for _, item := range val { |
| 487 | + cmdArgs = append(cmdArgs, fmt.Sprintf("%v", item)) |
| 488 | + } |
| 489 | + default: |
| 490 | + cmdArgs = append(cmdArgs, fmt.Sprintf("%v", val)) |
484 | 491 | }
|
485 |
| - default: |
486 |
| - cmdArgs = append(cmdArgs, fmt.Sprintf("%v", val)) |
487 | 492 | }
|
488 |
| - // Remove the "_" key so it's not processed as a flag |
489 |
| - delete(args, "_") |
| 493 | + } |
| 494 | + // Delete any of the positional argument keys so they don't get processed below. |
| 495 | + for _, dk := range deleteKeys { |
| 496 | + delete(args, dk) |
490 | 497 | }
|
491 | 498 |
|
492 | 499 | // Process remaining arguments as flags
|
@@ -639,6 +646,15 @@ func (s *MCPServer) generateJSONSchema(cmd *Command) (json.RawMessage, error) {
|
639 | 646 | properties := schema["properties"].(map[string]any)
|
640 | 647 | requiredList := schema["required"].([]string)
|
641 | 648 |
|
| 649 | + // Add positional arguments based on the cmd usage. |
| 650 | + if posArgs, err := PosArgsFromCmdUsage(cmd.Use); err != nil { |
| 651 | + return nil, xerrors.Errorf("unable to process positional argument for command %q: %w", cmd.Name(), err) |
| 652 | + } else { |
| 653 | + for k, v := range posArgs { |
| 654 | + properties[k] = v |
| 655 | + } |
| 656 | + } |
| 657 | + |
642 | 658 | // Process each option in the command
|
643 | 659 | for _, opt := range cmd.Options {
|
644 | 660 | // Skip options that aren't exposed as flags
|
@@ -925,3 +941,108 @@ Commands with neither Tool nor Resource set will not be accessible via MCP.`,
|
925 | 941 | },
|
926 | 942 | }
|
927 | 943 | }
|
| 944 | + |
| 945 | +// PosArgsFromCmdUsage attempts to process a 'usage' string into a set of |
| 946 | +// arguments for display as tool parameters. |
| 947 | +// Example: the usage string `foo [flags] <bar> [baz] [razzle|dazzle]` |
| 948 | +// defines three arguments for the `foo` command: |
| 949 | +// - bar (required) |
| 950 | +// - baz (optional) |
| 951 | +// - the string `razzle` XOR `dazzle` (optional) |
| 952 | +// |
| 953 | +// The expected output of the above is as follows: |
| 954 | +// |
| 955 | +// { |
| 956 | +// "arg1:bar": { |
| 957 | +// "type": "string", |
| 958 | +// "description": "required argument", |
| 959 | +// }, |
| 960 | +// "arg2:baz": { |
| 961 | +// "type": "string", |
| 962 | +// "description": "optional argument", |
| 963 | +// }, |
| 964 | +// "arg3:razzle_dazzle": { |
| 965 | +// "type": "string", |
| 966 | +// "enum": ["razzle", "dazzle"] |
| 967 | +// }, |
| 968 | +// } |
| 969 | +// |
| 970 | +// The usage string is processed given the following assumptions: |
| 971 | +// 1. The first non-whitespace string of usage is the name of the command |
| 972 | +// and will be skipped. |
| 973 | +// 2. The pseudo-argument specifier [flags] will also be skipped, if present. |
| 974 | +// 3. Argument specifiers enclosed by [square brackets] are considered optional. |
| 975 | +// 4. All other argument specifiers are considered required. |
| 976 | +// 5. Invidiual argument specifiers are separated by a single whitespace character. |
| 977 | +// Argument specifiers that contain a space are considered invalid (e.g. `[foo bar]`) |
| 978 | +// |
| 979 | +// Variadic arguments [arg...] are treated as a single argument. |
| 980 | +func PosArgsFromCmdUsage(usage string) (map[string]any, error) { |
| 981 | + if len(usage) == 0 { |
| 982 | + return nil, xerrors.Errorf("usage may not be empty") |
| 983 | + } |
| 984 | + |
| 985 | + // Step 1: preprocessing. Skip the first token. |
| 986 | + parts := strings.Fields(usage) |
| 987 | + if len(parts) < 2 { |
| 988 | + return map[string]any{}, nil |
| 989 | + } |
| 990 | + parts = parts[1:] |
| 991 | + // Skip [flags], if present. |
| 992 | + parts = slices.DeleteFunc(parts, func(s string) bool { |
| 993 | + return s == "[flags]" |
| 994 | + }) |
| 995 | + |
| 996 | + result := make(map[string]any, len(parts)) |
| 997 | + |
| 998 | + // Process each argument token |
| 999 | + for i, part := range parts { |
| 1000 | + argIndex := i + 1 |
| 1001 | + argKey := fmt.Sprintf("arg%d__", argIndex) |
| 1002 | + |
| 1003 | + // Check for unbalanced brackets in the part. |
| 1004 | + // This catches cases like "command [flags] [a" or "command [flags] a b [c | d]" |
| 1005 | + // which would be split into multiple tokens by strings.Fields() |
| 1006 | + openSquare := strings.Count(part, "[") |
| 1007 | + closeSquare := strings.Count(part, "]") |
| 1008 | + openAngle := strings.Count(part, "<") |
| 1009 | + closeAngle := strings.Count(part, ">") |
| 1010 | + openBrace := strings.Count(part, "{") |
| 1011 | + closeBrace := strings.Count(part, "}") |
| 1012 | + |
| 1013 | + if openSquare != closeSquare { |
| 1014 | + return nil, xerrors.Errorf("malformed usage: unbalanced square bracket at %q", part) |
| 1015 | + } else if openAngle != closeAngle { |
| 1016 | + return nil, xerrors.Errorf("malformed usage: unbalanced angle bracket at %q", part) |
| 1017 | + } else if openBrace != closeBrace { |
| 1018 | + return nil, xerrors.Errorf("malformed usage: unbalanced brace at %q", part) |
| 1019 | + } |
| 1020 | + |
| 1021 | + // Determine if the argument is optional (enclosed in square brackets) |
| 1022 | + isOptional := openSquare > 0 |
| 1023 | + cleanName := strings.Trim(part, "[]{}<>.") |
| 1024 | + description := "required argument" |
| 1025 | + if isOptional { |
| 1026 | + description = "optional argument" |
| 1027 | + } |
| 1028 | + |
| 1029 | + argVal := map[string]any{ |
| 1030 | + "type": "string", |
| 1031 | + "description": description, |
| 1032 | + // "required": !isOptional, |
| 1033 | + } |
| 1034 | + |
| 1035 | + keyName := cleanName |
| 1036 | + // If an argument specifier contains a pipe, treat it as an enum. |
| 1037 | + if strings.Contains(cleanName, "|") { |
| 1038 | + choices := strings.Split(cleanName, "|") |
| 1039 | + // Create a name by joining alternatives with underscores |
| 1040 | + keyName = strings.Join(choices, "_") |
| 1041 | + argVal["enum"] = choices |
| 1042 | + } |
| 1043 | + argKey += keyName |
| 1044 | + result[argKey] = argVal |
| 1045 | + } |
| 1046 | + |
| 1047 | + return result, nil |
| 1048 | +} |
0 commit comments