diff --git a/cmd/api/api.go b/cmd/api/api.go index b6383707f..1c8697650 100644 --- a/cmd/api/api.go +++ b/cmd/api/api.go @@ -100,7 +100,7 @@ func NewCmdApiWithContext(ctx context.Context, f *cmdutil.Factory, runF func(*AP } return nil, cobra.ShellCompDirectiveNoFileComp } - _ = cmd.RegisterFlagCompletionFunc("format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + cmdutil.RegisterFlagCompletion(cmd, "format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { return []string{"json", "ndjson", "table", "csv"}, cobra.ShellCompDirectiveNoFileComp }) diff --git a/cmd/auth/login.go b/cmd/auth/login.go index 4b91dddfc..add8762bb 100644 --- a/cmd/auth/login.go +++ b/cmd/auth/login.go @@ -72,7 +72,7 @@ browser. Run it in the background and retrieve the verification URL from its out cmd.Flags().BoolVar(&opts.NoWait, "no-wait", false, "initiate device authorization and return immediately; use --device-code to complete") cmd.Flags().StringVar(&opts.DeviceCode, "device-code", "", "poll and complete authorization with a device code from a previous --no-wait call") - _ = cmd.RegisterFlagCompletionFunc("domain", func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) { + cmdutil.RegisterFlagCompletion(cmd, "domain", func(_ *cobra.Command, _ []string, toComplete string) ([]string, cobra.ShellCompDirective) { return completeDomain(toComplete), cobra.ShellCompDirectiveNoFileComp }) diff --git a/cmd/root.go b/cmd/root.go index 57c91d8b1..1f630bac2 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -85,6 +85,8 @@ func Execute() int { fmt.Fprintln(os.Stderr, "Error:", err) return 1 } + configureFlagCompletions(os.Args) + f, rootCmd := buildInternal( context.Background(), inv, WithIO(os.Stdin, os.Stdout, os.Stderr), @@ -153,6 +155,12 @@ func isCompletionCommand(args []string) bool { return false } +// configureFlagCompletions enables cmdutil.RegisterFlagCompletion only when +// the invocation will actually serve a __complete request. +func configureFlagCompletions(args []string) { + cmdutil.SetFlagCompletionsDisabled(!isCompletionCommand(args)) +} + // handleRootError dispatches a command error to the appropriate handler // and returns the process exit code. func handleRootError(f *cmdutil.Factory, err error) int { diff --git a/cmd/root_test.go b/cmd/root_test.go index 09af9587a..fe718841f 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -196,3 +196,28 @@ func TestRootLong_AgentSkillsLinkTargetsReadmeSection(t *testing.T) { t.Fatalf("root help should not reference the removed install-ai-agent-skills anchor, got:\n%s", rootLong) } } + +func TestConfigureFlagCompletions(t *testing.T) { + t.Cleanup(func() { cmdutil.SetFlagCompletionsDisabled(false) }) + + tests := []struct { + name string + args []string + wantDisabled bool + }{ + {"plain command", []string{"im", "+send"}, true}, + {"help flag", []string{"im", "--help"}, true}, + {"no args", []string{}, true}, + {"__complete request", []string{"__complete", "im", "+send", ""}, false}, + {"completion subcommand", []string{"completion", "bash"}, false}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + cmdutil.SetFlagCompletionsDisabled(!tc.wantDisabled) + configureFlagCompletions(tc.args) + if got := cmdutil.FlagCompletionsDisabled(); got != tc.wantDisabled { + t.Fatalf("FlagCompletionsDisabled() = %v, want %v", got, tc.wantDisabled) + } + }) + } +} diff --git a/cmd/schema/schema.go b/cmd/schema/schema.go index 152f37d24..38ecaa322 100644 --- a/cmd/schema/schema.go +++ b/cmd/schema/schema.go @@ -377,7 +377,7 @@ func NewCmdSchema(f *cmdutil.Factory, runF func(*SchemaOptions) error) *cobra.Co cmd.ValidArgsFunction = completeSchemaPath(f) cmd.Flags().StringVar(&opts.Format, "format", "json", "output format: json (default) | pretty") - _ = cmd.RegisterFlagCompletionFunc("format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + cmdutil.RegisterFlagCompletion(cmd, "format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { return []string{"json", "pretty"}, cobra.ShellCompDirectiveNoFileComp }) diff --git a/cmd/service/service.go b/cmd/service/service.go index 808c80077..cd1f7c364 100644 --- a/cmd/service/service.go +++ b/cmd/service/service.go @@ -189,7 +189,7 @@ func NewCmdServiceMethodWithContext(ctx context.Context, f *cmdutil.Factory, spe cmd.Flags().StringVar(&opts.File, "file", "", "file to upload ([field=]path, supports - for stdin)") } } - _ = cmd.RegisterFlagCompletionFunc("format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + cmdutil.RegisterFlagCompletion(cmd, "format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { return []string{"json", "ndjson", "table", "csv"}, cobra.ShellCompDirectiveNoFileComp }) diff --git a/internal/cmdutil/completion.go b/internal/cmdutil/completion.go new file mode 100644 index 000000000..c16160955 --- /dev/null +++ b/internal/cmdutil/completion.go @@ -0,0 +1,37 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package cmdutil + +import ( + "sync/atomic" + + "github.com/spf13/cobra" +) + +// Cobra keeps completion callbacks in a package-global map keyed by +// *pflag.Flag with no removal path, so registrations made for a *cobra.Command +// outlive the command itself. Skip registration when the current invocation +// will not serve a completion request. +var flagCompletionsDisabled atomic.Bool + +// SetFlagCompletionsDisabled switches RegisterFlagCompletion between +// registering and no-op. Typically set once at process start. +func SetFlagCompletionsDisabled(disabled bool) { + flagCompletionsDisabled.Store(disabled) +} + +// FlagCompletionsDisabled reports the current switch state. +func FlagCompletionsDisabled() bool { + return flagCompletionsDisabled.Load() +} + +// RegisterFlagCompletion wraps (*cobra.Command).RegisterFlagCompletionFunc +// and honors the package switch. The underlying error is swallowed to match +// the `_ = cmd.RegisterFlagCompletionFunc(...)` style already used here. +func RegisterFlagCompletion(cmd *cobra.Command, flagName string, fn cobra.CompletionFunc) { + if flagCompletionsDisabled.Load() { + return + } + _ = cmd.RegisterFlagCompletionFunc(flagName, fn) +} diff --git a/internal/cmdutil/completion_test.go b/internal/cmdutil/completion_test.go new file mode 100644 index 000000000..a20d290ba --- /dev/null +++ b/internal/cmdutil/completion_test.go @@ -0,0 +1,78 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package cmdutil + +import ( + "runtime" + "sync/atomic" + "testing" + "time" + + "github.com/spf13/cobra" +) + +func TestSetFlagCompletionsDisabled_RoundTrip(t *testing.T) { + t.Cleanup(func() { SetFlagCompletionsDisabled(false) }) + + if FlagCompletionsDisabled() { + t.Fatal("expected default false") + } + SetFlagCompletionsDisabled(true) + if !FlagCompletionsDisabled() { + t.Fatal("expected true after Set(true)") + } + SetFlagCompletionsDisabled(false) + if FlagCompletionsDisabled() { + t.Fatal("expected false after Set(false)") + } +} + +// When disabled, a *cobra.Command must be collectable after the caller drops +// its reference — i.e. the wrapper did not touch cobra's global map. +func TestRegisterFlagCompletion_Disabled_DoesNotRetainCommand(t *testing.T) { + SetFlagCompletionsDisabled(true) + t.Cleanup(func() { SetFlagCompletionsDisabled(false) }) + + const N = 5 + var collected atomic.Int32 + func() { + for range N { + cmd := &cobra.Command{Use: "x"} + cmd.Flags().String("foo", "", "") + RegisterFlagCompletion(cmd, "foo", func(_ *cobra.Command, _ []string, _ string) ([]cobra.Completion, cobra.ShellCompDirective) { + return nil, cobra.ShellCompDirectiveNoFileComp + }) + runtime.SetFinalizer(cmd, func(_ *cobra.Command) { collected.Add(1) }) + } + }() + // Finalizers run on a dedicated goroutine after GC; loop to give it time. + for range 30 { + runtime.GC() + time.Sleep(20 * time.Millisecond) + } + if got := collected.Load(); int(got) != N { + t.Fatalf("expected %d *cobra.Command finalizers to fire when completions disabled, got %d", N, got) + } +} + +// When enabled, the registered completion must be reachable via cobra. +func TestRegisterFlagCompletion_Enabled_DoesRegister(t *testing.T) { + SetFlagCompletionsDisabled(false) + + cmd := &cobra.Command{Use: "x"} + cmd.Flags().String("foo", "", "") + want := []cobra.Completion{"a", "b"} + RegisterFlagCompletion(cmd, "foo", func(_ *cobra.Command, _ []string, _ string) ([]cobra.Completion, cobra.ShellCompDirective) { + return want, cobra.ShellCompDirectiveNoFileComp + }) + + fn, ok := cmd.GetFlagCompletionFunc("foo") + if !ok { + t.Fatal("expected completion func to be registered") + } + got, _ := fn(cmd, nil, "") + if len(got) != 2 || got[0] != "a" || got[1] != "b" { + t.Fatalf("unexpected completion result: %v", got) + } +} diff --git a/internal/cmdutil/identity_flag.go b/internal/cmdutil/identity_flag.go index c99d5c628..1589120ea 100644 --- a/internal/cmdutil/identity_flag.go +++ b/internal/cmdutil/identity_flag.go @@ -54,7 +54,7 @@ func addIdentityFlag(ctx context.Context, cmd *cobra.Command, f *Factory, target } registerIdentityFlag(cmd, target, cfg.defaultValue, cfg.usage) - _ = cmd.RegisterFlagCompletionFunc("as", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + RegisterFlagCompletion(cmd, "as", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { return cfg.completionValues, cobra.ShellCompDirectiveNoFileComp }) } diff --git a/shortcuts/common/runner.go b/shortcuts/common/runner.go index 5b663ee1b..e81db01c7 100644 --- a/shortcuts/common/runner.go +++ b/shortcuts/common/runner.go @@ -876,7 +876,7 @@ func registerShortcutFlagsWithContext(ctx context.Context, cmd *cobra.Command, f } if len(fl.Enum) > 0 { vals := fl.Enum - _ = cmd.RegisterFlagCompletionFunc(fl.Name, func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + cmdutil.RegisterFlagCompletion(cmd, fl.Name, func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { return vals, cobra.ShellCompDirectiveNoFileComp }) } @@ -892,7 +892,7 @@ func registerShortcutFlagsWithContext(ctx context.Context, cmd *cobra.Command, f cmd.Flags().StringP("jq", "q", "", "jq expression to filter JSON output") cmdutil.AddShortcutIdentityFlag(ctx, cmd, f, s.AuthTypes) if s.HasFormat { - _ = cmd.RegisterFlagCompletionFunc("format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { + cmdutil.RegisterFlagCompletion(cmd, "format", func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) { return []string{"json", "pretty", "table", "ndjson", "csv"}, cobra.ShellCompDirectiveNoFileComp }) } diff --git a/shortcuts/common/runner_flag_completion_test.go b/shortcuts/common/runner_flag_completion_test.go new file mode 100644 index 000000000..35ce532fe --- /dev/null +++ b/shortcuts/common/runner_flag_completion_test.go @@ -0,0 +1,98 @@ +// Copyright (c) 2026 Lark Technologies Pte. Ltd. +// SPDX-License-Identifier: MIT + +package common + +import ( + "context" + "testing" + + "github.com/larksuite/cli/internal/cmdutil" + "github.com/spf13/cobra" +) + +// TestShortcutMount_FlagCompletionsRegistered exercises the two +// cmdutil.RegisterFlagCompletion call sites in registerShortcutFlagsWithContext: +// the per-flag enum completion (runner.go:879) and the auto-injected --format +// completion (runner.go:895). +func TestShortcutMount_FlagCompletionsRegistered(t *testing.T) { + t.Cleanup(func() { cmdutil.SetFlagCompletionsDisabled(false) }) + cmdutil.SetFlagCompletionsDisabled(false) + + f, _, _, _ := cmdutil.TestFactory(t, nil) + parent := &cobra.Command{Use: "root"} + shortcut := Shortcut{ + Service: "docs", + Command: "+fetch", + Description: "fetch doc", + HasFormat: true, + Flags: []Flag{ + {Name: "sort-by", Desc: "sort", Enum: []string{"asc", "desc"}}, + }, + Execute: func(context.Context, *RuntimeContext) error { return nil }, + } + shortcut.Mount(parent, f) + + cmd, _, err := parent.Find([]string{"+fetch"}) + if err != nil { + t.Fatalf("Find() error = %v", err) + } + + // Enum flag completion. + fn, ok := cmd.GetFlagCompletionFunc("sort-by") + if !ok { + t.Fatal("expected completion func for --sort-by") + } + got, _ := fn(cmd, nil, "") + if len(got) != 2 || got[0] != "asc" || got[1] != "desc" { + t.Fatalf("sort-by completion = %v, want [asc desc]", got) + } + + // HasFormat-injected --format completion. + fn, ok = cmd.GetFlagCompletionFunc("format") + if !ok { + t.Fatal("expected completion func for --format") + } + got, _ = fn(cmd, nil, "") + want := []string{"json", "pretty", "table", "ndjson", "csv"} + if len(got) != len(want) { + t.Fatalf("format completion = %v, want %v", got, want) + } + for i, v := range want { + if got[i] != v { + t.Fatalf("format completion[%d] = %q, want %q", i, got[i], v) + } + } +} + +// TestShortcutMount_FlagCompletionsDisabled verifies the switch actually +// prevents the two registrations from landing in cobra's global map. +func TestShortcutMount_FlagCompletionsDisabled(t *testing.T) { + t.Cleanup(func() { cmdutil.SetFlagCompletionsDisabled(false) }) + cmdutil.SetFlagCompletionsDisabled(true) + + f, _, _, _ := cmdutil.TestFactory(t, nil) + parent := &cobra.Command{Use: "root"} + shortcut := Shortcut{ + Service: "docs", + Command: "+fetch", + Description: "fetch doc", + HasFormat: true, + Flags: []Flag{ + {Name: "sort-by", Desc: "sort", Enum: []string{"asc", "desc"}}, + }, + Execute: func(context.Context, *RuntimeContext) error { return nil }, + } + shortcut.Mount(parent, f) + + cmd, _, err := parent.Find([]string{"+fetch"}) + if err != nil { + t.Fatalf("Find() error = %v", err) + } + if _, ok := cmd.GetFlagCompletionFunc("sort-by"); ok { + t.Fatal("did not expect completion func for --sort-by when disabled") + } + if _, ok := cmd.GetFlagCompletionFunc("format"); ok { + t.Fatal("did not expect completion func for --format when disabled") + } +}