From 351698c280e1a2368bbb5dde3e3a4c8991b15bb9 Mon Sep 17 00:00:00 2001 From: Pasha Sviderski Date: Tue, 21 Jul 2026 17:01:21 +1000 Subject: [PATCH] fix(cli): completion with direct connections (--connect, --context, --uncloud-config) (fixes #377) --- cmd/uc/main.go | 32 ++++++++ cmd/uc/main_test.go | 116 +++++++++++++++++++++++++++++ internal/cli/completion/context.go | 6 +- internal/cli/completion/machine.go | 3 +- internal/cli/completion/service.go | 3 +- internal/cli/completion/volume.go | 3 +- 6 files changed, 159 insertions(+), 4 deletions(-) create mode 100644 cmd/uc/main_test.go diff --git a/cmd/uc/main.go b/cmd/uc/main.go index a81e1501..e15f6d81 100644 --- a/cmd/uc/main.go +++ b/cmd/uc/main.go @@ -24,6 +24,7 @@ import ( "github.com/psviderski/uncloud/internal/machine" "github.com/psviderski/uncloud/internal/version" "github.com/spf13/cobra" + "github.com/spf13/pflag" ) type globalOptions struct { @@ -42,6 +43,13 @@ func main() { SilenceUsage: true, SilenceErrors: true, PersistentPreRunE: func(cmd *cobra.Command, args []string) error { + // Shell completion runs through the hidden __complete command which has flag parsing disabled, + // so the global flags from the completed command line are never parsed. Apply them manually to make + // completion work with --connect, --context, and --uncloud-config. + if cmd.Name() == cobra.ShellCompRequestCmd { + applyGlobalFlagsFromCompletionArgs(cmd.Root().PersistentFlags(), os.Args[1:]) + } + cli.BindEnvToFlag(cmd, "connect", "UNCLOUD_CONNECT") cli.BindEnvToFlag(cmd, "context", "UNCLOUD_CONTEXT") cli.BindEnvToFlag(cmd, "uncloud-config", "UNCLOUD_CONFIG") @@ -159,3 +167,27 @@ func main() { cobra.CheckErr(err) } } + +// applyGlobalFlagsFromCompletionArgs parses the global flags from the raw arguments of a __complete command and applies +// the ones found to flags. The trailing word being completed, unknown flags, and positional arguments are ignored. +func applyGlobalFlagsFromCompletionArgs(flags *pflag.FlagSet, args []string) { + // The shell always passes the word being completed as the last argument, even if it's empty. + // Exclude it from parsing as its value may not be complete yet. + if len(args) == 0 { + return + } + args = args[:len(args)-1] + + fset := pflag.NewFlagSet("global", pflag.ContinueOnError) + fset.ParseErrorsAllowlist.UnknownFlags = true + fset.String("connect", "", "") + fset.StringP("context", "c", "", "") + fset.String("uncloud-config", "", "") + // Parsing an incomplete command line may fail, apply the flags parsed so far anyway. + _ = fset.Parse(args) + + fset.Visit(func(f *pflag.Flag) { + // Setting the flag marks it as changed so it takes precedence over environment variables. + _ = flags.Set(f.Name, f.Value.String()) + }) +} diff --git a/cmd/uc/main_test.go b/cmd/uc/main_test.go new file mode 100644 index 00000000..c2bec5eb --- /dev/null +++ b/cmd/uc/main_test.go @@ -0,0 +1,116 @@ +package main + +import ( + "slices" + "testing" + + "github.com/spf13/pflag" + "github.com/stretchr/testify/assert" +) + +func TestApplyGlobalFlagsFromCompletionArgs(t *testing.T) { + defaultConfigPath := "~/.config/uncloud/config.yaml" + + tests := []struct { + name string + args []string + wantConnect string + wantContext string + wantConfigPath string + // Flag names expected to be marked as changed on the target flag set. + wantChanged []string + }{ + { + name: "no flags", + args: []string{"__complete", "inspect", ""}, + }, + { + name: "connect with space", + args: []string{"__complete", "--connect", "ssh://user@host", "inspect", ""}, + wantConnect: "ssh://user@host", + wantChanged: []string{"connect"}, + }, + { + name: "connect with equals", + args: []string{"__complete", "--connect=tcp://127.0.0.1:51000", "inspect", ""}, + wantConnect: "tcp://127.0.0.1:51000", + wantChanged: []string{"connect"}, + }, + { + name: "context shorthand", + args: []string{"__complete", "-c", "prod", "inspect", ""}, + wantContext: "prod", + wantChanged: []string{"context"}, + }, + { + name: "all flags", + args: []string{"__complete", "--connect", "user@host", "-c", "prod", "--uncloud-config", "/tmp/uncloud.yaml", "inspect", ""}, + wantConnect: "user@host", + wantContext: "prod", + wantConfigPath: "/tmp/uncloud.yaml", + wantChanged: []string{"connect", "context", "uncloud-config"}, + }, + { + name: "unknown flags are ignored", + args: []string{"__complete", "--quiet", "-n", "5", "--connect", "user@host", "logs", ""}, + wantConnect: "user@host", + wantChanged: []string{"connect"}, + }, + { + name: "flags after double dash are ignored", + args: []string{"__complete", "exec", "svc", "--", "sh", "--connect", "user@host"}, + }, + { + name: "flags before double dash are applied", + args: []string{"__complete", "--connect", "user@host", "exec", "svc", "--", "sh", "-c", "env"}, + wantConnect: "user@host", + wantChanged: []string{"connect"}, + }, + { + name: "partial flag name being completed is excluded", + args: []string{"__complete", "--connect", "user@host", "inspect", "--context"}, + wantConnect: "user@host", + wantChanged: []string{"connect"}, + }, + { + name: "partial flag value being completed is excluded", + args: []string{"__complete", "--uncloud-config", "/tmp/"}, + }, + { + name: "partial connect value being completed is excluded", + args: []string{"__complete", "--connect", "tcp://127.0.0.1:5"}, + }, + { + name: "completed flag value with partial command word", + args: []string{"__complete", "--uncloud-config", "/tmp/uncloud.yaml", "insp"}, + wantConfigPath: "/tmp/uncloud.yaml", + wantChanged: []string{"uncloud-config"}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Mirror the global persistent flags defined on the root command. + var opts globalOptions + flags := pflag.NewFlagSet("test", pflag.ContinueOnError) + flags.StringVar(&opts.connect, "connect", "", "") + flags.StringVarP(&opts.context, "context", "c", "", "") + flags.StringVar(&opts.configPath, "uncloud-config", defaultConfigPath, "") + + applyGlobalFlagsFromCompletionArgs(flags, tt.args) + + assert.Equal(t, tt.wantConnect, opts.connect) + assert.Equal(t, tt.wantContext, opts.context) + wantConfigPath := tt.wantConfigPath + if wantConfigPath == "" { + wantConfigPath = defaultConfigPath + } + assert.Equal(t, wantConfigPath, opts.configPath) + + for _, name := range []string{"connect", "context", "uncloud-config"} { + assert.Equal(t, slices.Contains(tt.wantChanged, name), flags.Changed(name), + "changed status of flag '%s'", name) + } + }) + } +} diff --git a/internal/cli/completion/context.go b/internal/cli/completion/context.go index 77fa5b55..c6580338 100644 --- a/internal/cli/completion/context.go +++ b/internal/cli/completion/context.go @@ -11,6 +11,11 @@ import ( ) func Contexts(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) { + // There are no contexts to complete when the CLI uses a direct machine connection (--connect) without a config. + if uncli.Config == nil { + return nil, cobra.ShellCompDirectiveNoFileComp + } + contexts := slices.Sorted(maps.Keys(uncli.Config.Contexts)) names := []cobra.Completion{} @@ -21,7 +26,6 @@ func Contexts(ctx context.Context, uncli *cli.CLI, args []string, toComplete str if strings.HasPrefix(context, toComplete) { names = append(names, context) } - names = append(names, context) } return names, cobra.ShellCompDirectiveNoFileComp diff --git a/internal/cli/completion/machine.go b/internal/cli/completion/machine.go index 33b2a569..4ef81169 100644 --- a/internal/cli/completion/machine.go +++ b/internal/cli/completion/machine.go @@ -10,7 +10,8 @@ import ( ) func Machines(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) { - client, err := uncli.ConnectCluster(ctx) + // Disable the connection progress output to not interfere with the shell completion output. + client, err := uncli.ConnectClusterWithOptions(ctx, cli.ConnectOptions{}) if err != nil { return nil, cobra.ShellCompDirectiveError } diff --git a/internal/cli/completion/service.go b/internal/cli/completion/service.go index c91f0b1c..9e49a662 100644 --- a/internal/cli/completion/service.go +++ b/internal/cli/completion/service.go @@ -11,7 +11,8 @@ import ( ) func Services(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) { - client, err := uncli.ConnectCluster(ctx) + // Disable the connection progress output to not interfere with the shell completion output. + client, err := uncli.ConnectClusterWithOptions(ctx, cli.ConnectOptions{}) if err != nil { return nil, cobra.ShellCompDirectiveError } diff --git a/internal/cli/completion/volume.go b/internal/cli/completion/volume.go index 51e482c9..6d439c1b 100644 --- a/internal/cli/completion/volume.go +++ b/internal/cli/completion/volume.go @@ -11,7 +11,8 @@ import ( ) func Volumes(ctx context.Context, uncli *cli.CLI, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) { - client, err := uncli.ConnectCluster(ctx) + // Disable the connection progress output to not interfere with the shell completion output. + client, err := uncli.ConnectClusterWithOptions(ctx, cli.ConnectOptions{}) if err != nil { return nil, cobra.ShellCompDirectiveError }