fix(cli): completion with direct connections (--connect, --context, --uncloud-config) (fixes #377)

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

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