fix(autocomplete): rejoin trailing colon command tokens

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"}} {

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

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

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