diff --git a/cmd/openai/main_requestheaders_test.go b/cmd/openai/main_requestheaders_test.go index 215a132d..f2a08892 100644 --- a/cmd/openai/main_requestheaders_test.go +++ b/cmd/openai/main_requestheaders_test.go @@ -327,6 +327,11 @@ func TestMainRequestHeadersHelpAndCompletion(t *testing.T) { t.Errorf("header help = %q, want the environment input documented", got.stdout) } } + got := runMainDispatch(t, "zsh", "openai", "__complete", "--header", "chat:", "completions", "create", "--mo") + if got.code != 0 || !strings.Contains(got.stdout, "--model:") { + t.Errorf("zsh completion after colon-ending header = %+v, want --model suggestion", got) + } + for _, style := range []string{"bash", "zsh", "fish", "pwsh"} { t.Run(style, func(t *testing.T) { for _, scope := range [][]string{nil, {"models", "retrieve"}} { diff --git a/internal/autocomplete/autocomplete.go b/internal/autocomplete/autocomplete.go index 82718af6..0edae6e4 100644 --- a/internal/autocomplete/autocomplete.go +++ b/internal/autocomplete/autocomplete.go @@ -312,7 +312,7 @@ func getAllPossibleCompletions(completionStyle CompletionStyle, root *cli.Comman func ExecuteShellCompletion(ctx context.Context, cmd *cli.Command) error { root := cmd.Root() - args := rebuildColonSeparatedArgs(root.Args().Slice()[1:]) + args := rebuildColonSeparatedArgs(root, root.Args().Slice()[1:]) var completionStyle CompletionStyle if style, ok := os.LookupEnv("COMPLETION_STYLE"); ok { @@ -350,41 +350,111 @@ func ExecuteShellCompletion(ctx context.Context, cmd *cli.Command) error { return cli.Exit("", int(result.Behavior)) } -// When CLI arguments are passed in, they are separated on word barriers. -// Most commonly this is whitespace but in some cases that may also be colons. -// We wish to allow arguments with colons. To handle this, we append/prepend colons to their neighboring -// arguments. -// -// Example: `rebuildColonSeparatedArgs(["a", "b", ":", "c", "d"])` => `["a", "b:c", "d"]` -func rebuildColonSeparatedArgs(args []string) []string { +// When CLI arguments are passed in, shell word breaking can split a colon +// command into adjacent tokens. Rejoin explicit colon separators, and only +// rejoin a token that already ends in ':' when the combined value can still +// name a command. This avoids swallowing an ordinary following argument such +// as `--instructions Prefix: --model`. +func rebuildColonSeparatedArgs(root *cli.Command, args []string) []string { if len(args) == 0 { return args } result := []string{} + cmd := root + lineage := []*cli.Command{root} + flags := completionFlags(lineage) i := 0 for i < len(args) { current := args[i] - // Keep joining while the next element is ":" or the current element ends with ":" - for i+1 < len(args) && (args[i+1] == ":" || strings.HasSuffix(current, ":")) { - if args[i+1] == ":" { - current += ":" + // A value-taking flag owns the next shell word. Do not let a trailing colon + // in that value absorb the command that follows it. + if isFlag(current) { + result = append(result, current) + if flag := findFlag(flags, current); flag != nil { + if docFlag, ok := (*flag).(cli.DocGenerationFlag); ok && docFlag.TakesValue() && i+1 < len(args) { + value := args[i+1] + i += 2 + + // Bash includes ':' in COMP_WORDBREAKS, so a single flag value such as + // `X:completions` can arrive as `X`, `:`, `completions`. Rebuild the + // split value here, but stop at a colon when the following tokens form + // a valid command path. That preserves both `X:completions models ...` + // and a value ending in a colon, such as `chat: completions create ...`. + for i < len(args) && args[i] == ":" { + value += ":" + i++ + if i >= len(args) || commandTailStartsAt(cmd, args[i:]) { + break + } + value += args[i] + i++ + } + result = append(result, value) + continue + } + } + i++ + continue + } + + for i+1 < len(args) { + next := args[i+1] + if next == ":" { + current += next i++ - // Check if there's a following element after the ":" if i+1 < len(args) && args[i+1] != ":" { current += args[i+1] i++ } - } else { - break + continue + } + if strings.HasSuffix(current, ":") && hasCommandPrefix(root, current+next) { + current += next + i++ + continue } + break } result = append(result, current) + if child := findChild(cmd, current); child != nil { + cmd = child + lineage = append(lineage, child) + flags = completionFlags(lineage) + } i++ } return result } + +func commandTailStartsAt(cmd *cli.Command, args []string) bool { + matched := false + for _, arg := range args { + if isFlag(arg) { + return matched + } + child := findChild(cmd, arg) + if child == nil { + return false + } + matched = true + cmd = child + } + return matched +} + +func hasCommandPrefix(cmd *cli.Command, prefix string) bool { + if cmd == nil { + return false + } + for _, child := range cmd.Commands { + if strings.HasPrefix(child.Name, prefix) || hasCommandPrefix(child, prefix) { + return true + } + } + return false +} diff --git a/internal/autocomplete/colon_args_test.go b/internal/autocomplete/colon_args_test.go new file mode 100644 index 00000000..9a7da172 --- /dev/null +++ b/internal/autocomplete/colon_args_test.go @@ -0,0 +1,72 @@ +package autocomplete + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/urfave/cli/v3" +) + +func TestRebuildColonSeparatedArgs(t *testing.T) { + t.Parallel() + + root := &cli.Command{ + Flags: []cli.Flag{&cli.StringFlag{Name: "header"}}, + Commands: []*cli.Command{ + {Name: "config:get"}, + {Name: "config:set"}, + {Name: "completions", Commands: []*cli.Command{{Name: "create"}}}, + {Name: "models", Commands: []*cli.Command{{Name: "retrieve"}}}, + {Name: "chat:completions", Commands: []*cli.Command{{Name: "create"}}}, + }, + } + + tests := map[string]struct { + args []string + want []string + }{ + "standalone colon": { + args: []string{"a", "b", ":", "c", "d"}, + want: []string{"a", "b:c", "d"}, + }, + "trailing colon": { + args: []string{"config:", "get"}, + want: []string{"config:get"}, + }, + "repeated colons": { + args: []string{"a", ":", ":", "b"}, + want: []string{"a::b"}, + }, + "ordinary arguments": { + args: []string{"a", "b", "c"}, + want: []string{"a", "b", "c"}, + }, + "colon-ending value before flag": { + args: []string{"--instructions", "Prefix:", "--mo"}, + want: []string{"--instructions", "Prefix:", "--mo"}, + }, + "colon-ending ordinary value": { + args: []string{"Prefix:", "value"}, + want: []string{"Prefix:", "value"}, + }, + "flag value ending in colon before command": { + args: []string{"--header", "chat:", "completions", "create", "--mo"}, + want: []string{"--header", "chat:", "completions", "create", "--mo"}, + }, + "bash-split colon inside flag value": { + args: []string{"--header", "X", ":", "completions", "models", "retrieve", "--mo"}, + want: []string{"--header", "X:completions", "models", "retrieve", "--mo"}, + }, + "bash-split trailing colon before command": { + args: []string{"--header", "chat", ":", "completions", "create", "--mo"}, + want: []string{"--header", "chat:", "completions", "create", "--mo"}, + }, + } + + for name, test := range tests { + t.Run(name, func(t *testing.T) { + t.Parallel() + assert.Equal(t, test.want, rebuildColonSeparatedArgs(root, test.args)) + }) + } +} diff --git a/internal/autocomplete/protocol_test.go b/internal/autocomplete/protocol_test.go index 907f9351..24c98a2b 100644 --- a/internal/autocomplete/protocol_test.go +++ b/internal/autocomplete/protocol_test.go @@ -25,10 +25,12 @@ func TestShellCompletionProtocolHelper(t *testing.T) { Flags: []cli.Flag{ &cli.StringFlag{Name: "format"}, &cli.StringFlag{Name: "file", TakesFile: true}, + &cli.StringFlag{Name: "header"}, }, Commands: []*cli.Command{ {Name: "models", Commands: []*cli.Command{ {Name: "list", Flags: []cli.Flag{&cli.IntFlag{Name: "max-items"}}}, + {Name: "retrieve", Flags: []cli.Flag{&cli.StringFlag{Name: "model"}}}, }}, {Name: "__complete", Hidden: true, SkipFlagParsing: true, Action: ExecuteShellCompletion}, }, @@ -57,17 +59,26 @@ func TestShellCompletionProtocol(t *testing.T) { for _, test := range []struct { name string args []string + bashArgs []string code int output string candidates string }{ - {"root value", []string{"--format", "candidate-"}, 11, "", ""}, - {"nested local value", []string{"models", "list", "--max-items", "candidate-"}, 11, "", ""}, - {"file value", []string{"--file", "candidate-"}, 10, "", "candidate-fixture.txt\n"}, - {"spaced preceding value", []string{"--format", "two words", "--file", "candidate-"}, 10, "", "candidate-fixture.txt\n"}, - {"empty preceding value", []string{"--format", "", "--file", "candidate-"}, 10, "", "candidate-fixture.txt\n"}, - {"explicit file prefix", []string{"--format", "@candidate-"}, 11, "", "@candidate-fixture.txt\n"}, - {"command prefix", []string{"mo"}, 0, "models\n", "models\n"}, + {name: "root value", args: []string{"--format", "candidate-"}, code: 11}, + {name: "nested local value", args: []string{"models", "list", "--max-items", "candidate-"}, code: 11}, + {name: "file value", args: []string{"--file", "candidate-"}, code: 10, candidates: "candidate-fixture.txt\n"}, + {name: "spaced preceding value", args: []string{"--format", "two words", "--file", "candidate-"}, code: 10, candidates: "candidate-fixture.txt\n"}, + {name: "empty preceding value", args: []string{"--format", "", "--file", "candidate-"}, code: 10, candidates: "candidate-fixture.txt\n"}, + {name: "explicit file prefix", args: []string{"--format", "@candidate-"}, code: 11, candidates: "@candidate-fixture.txt\n"}, + {name: "command prefix", args: []string{"mo"}, code: 0, output: "models\n", candidates: "models\n"}, + { + name: "bash colon-split header value", + args: []string{"--header", "X:completions", "models", "retrieve", "--mo"}, + bashArgs: []string{"--header", "X", ":", "completions", "models", "retrieve", "--mo"}, + code: 0, + output: "--model\n", + candidates: "--model\n", + }, } { t.Run(test.name, func(t *testing.T) { t.Parallel() @@ -117,7 +128,11 @@ for candidate in "${COMPREPLY[@]}"; do printf '%s\n' "$candidate" done ` - args := append([]string{"-c", probe, "completion-probe", binary}, test.args...) + bashArgs := test.args + if test.bashArgs != nil { + bashArgs = test.bashArgs + } + args := append([]string{"-c", probe, "completion-probe", binary}, bashArgs...) command := exec.Command(bash, args...) command.Dir, command.Env = dir, env stdout.Reset() @@ -128,7 +143,7 @@ done require.Equal(t, test.candidates, stdout.String()) argv, err := os.ReadFile(filepath.Join(dir, "helper.argv")) require.NoError(t, err) - wantArgs := append([]string{"__complete", "--"}, test.args...) + wantArgs := append([]string{"__complete", "--"}, bashArgs...) require.Equal(t, strings.Join(wantArgs, "\x00")+"\x00", string(argv)) }) })