From e09050a89c6f4dac4bd7c4540ef1c32cae7df591 Mon Sep 17 00:00:00 2001 From: Anzal Husain Abidi Date: Wed, 23 Sep 2026 22:09:32 +0530 Subject: [PATCH] fix(cli): honour ax ssh --help and stop global flag parsing at -- ax ssh --help was taken as a task name and ended in a NotFound error after resolving the server. The ssh arguments are now parsed up front, so -h/--help prints ssh usage and a missing task name fails fast, before any server lookup or port-forward. The global flag parser also kept consuming -a, -n, --server and --context after "--", so ax ssh t -- grep -n foo file silently set the namespace to foo. Parsing now stops at "--" and hands the rest to the command unchanged. Fixes #373 Fixes #410 --- cmd/ax/main.go | 166 ++++++++++++++++++++++++++++++-------------- cmd/ax/main_test.go | 90 ++++++++++++++++++++++++ 2 files changed, 202 insertions(+), 54 deletions(-) diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 9771f464..c8470a8e 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -45,53 +45,16 @@ func main() { os.Exit(1) } + opts := parseGlobalArgs(os.Args[1:]) var ( - cmd string - cleanArgs []string - atespace = "default" - explicitServer = "" - kubeContext = "" - axNamespace = "ax-system" + cmd = opts.cmd + cleanArgs = opts.args + atespace = opts.atespace + explicitServer = opts.server + kubeContext = opts.kubeContext + axNamespace = opts.namespace ) - args := os.Args[1:] - for i := 0; i < len(args); i++ { - arg := args[i] - if arg == "-a" || arg == "--atespace" { - if i+1 < len(args) { - atespace = args[i+1] - i++ - } - } else if strings.HasPrefix(arg, "--atespace=") { - atespace = strings.TrimPrefix(arg, "--atespace=") - } else if arg == "--server" { - if i+1 < len(args) { - explicitServer = args[i+1] - i++ - } - } else if strings.HasPrefix(arg, "--server=") { - explicitServer = strings.TrimPrefix(arg, "--server=") - } else if arg == "--context" { - if i+1 < len(args) { - kubeContext = args[i+1] - i++ - } - } else if strings.HasPrefix(arg, "--context=") { - kubeContext = strings.TrimPrefix(arg, "--context=") - } else if arg == "-n" || arg == "--namespace" { - if i+1 < len(args) { - axNamespace = args[i+1] - i++ - } - } else if strings.HasPrefix(arg, "--namespace=") { - axNamespace = strings.TrimPrefix(arg, "--namespace=") - } else if cmd == "" && !strings.HasPrefix(arg, "-") { - cmd = arg - } else { - cleanArgs = append(cleanArgs, arg) - } - } - if cmd == "" { printUsage() os.Exit(1) @@ -117,6 +80,17 @@ func main() { os.Exit(1) } return + case "ssh": + // Answer help and usage errors before resolving the server, which may + // start a port-forward to the cluster. + if _, _, help, err := parseSSHArgs(cleanArgs); help || err != nil { + fmt.Println(sshUsage) + if err != nil { + fmt.Fprintf(os.Stderr, "Error: %v\n", err) + os.Exit(1) + } + return + } } // Resolve the AX server URL (auto-tunneling to active kube context if not explicitly set) @@ -159,6 +133,64 @@ func main() { } } +// globalArgs holds the command name, the flags shared by every command, and the +// remaining arguments for the command itself. +type globalArgs struct { + cmd string + args []string + atespace string + server string + kubeContext string + namespace string +} + +// parseGlobalArgs extracts the command name and the global flags from args. It +// stops interpreting flags at "--", so everything after it (for example the +// command run by `ax ssh -- ...`) is passed through to the command intact. +func parseGlobalArgs(args []string) globalArgs { + g := globalArgs{atespace: "default", namespace: "ax-system"} + for i := 0; i < len(args); i++ { + arg := args[i] + if arg == "--" { + g.args = append(g.args, args[i:]...) + break + } else if arg == "-a" || arg == "--atespace" { + if i+1 < len(args) { + g.atespace = args[i+1] + i++ + } + } else if strings.HasPrefix(arg, "--atespace=") { + g.atespace = strings.TrimPrefix(arg, "--atespace=") + } else if arg == "--server" { + if i+1 < len(args) { + g.server = args[i+1] + i++ + } + } else if strings.HasPrefix(arg, "--server=") { + g.server = strings.TrimPrefix(arg, "--server=") + } else if arg == "--context" { + if i+1 < len(args) { + g.kubeContext = args[i+1] + i++ + } + } else if strings.HasPrefix(arg, "--context=") { + g.kubeContext = strings.TrimPrefix(arg, "--context=") + } else if arg == "-n" || arg == "--namespace" { + if i+1 < len(args) { + g.namespace = args[i+1] + i++ + } + } else if strings.HasPrefix(arg, "--namespace=") { + g.namespace = strings.TrimPrefix(arg, "--namespace=") + } else if g.cmd == "" && !strings.HasPrefix(arg, "-") { + g.cmd = arg + } else { + g.args = append(g.args, arg) + } + } + return g +} + func printUsage() { fmt.Println(`AX CLI - Autonomous agent execution control @@ -1008,24 +1040,50 @@ func runTunnel(args []string) error { } } -func runSSH(serverURL, atespace, kubeContext string, args []string) error { - if len(args) == 0 { - return fmt.Errorf("usage: ax ssh [-- command...]") +const sshUsage = `Usage: + ax ssh [-- command...] + +Run a command inside a running task's container, or /bin/sh when no command is +given. The task must be Running and have spec.debug: true. + +Examples: + ax ssh task123 + ax ssh task123 -- ls -la /workspace + ax ssh -a my-atespace task123 -- python3 main.py` + +// parseSSHArgs splits the arguments of `ax ssh` into the task name and the command +// to run, defaulting the command to /bin/sh. help reports a -h or --help given in +// place of the task name. Arguments after the task name form the command, with or +// without a separating "--". +func parseSSHArgs(args []string) (taskName string, command []string, help bool, err error) { + if len(args) == 0 || args[0] == "--" { + return "", nil, false, errors.New("missing task name") + } + if args[0] == "-h" || args[0] == "--help" { + return "", nil, true, nil + } + if strings.HasPrefix(args[0], "-") { + return "", nil, false, fmt.Errorf("unknown flag %q", args[0]) } - taskName := args[0] - var cmdToRun []string + taskName = args[0] for i := 1; i < len(args); i++ { if args[i] == "--" { - cmdToRun = args[i+1:] + command = append(command, args[i+1:]...) break - } else { - cmdToRun = append(cmdToRun, args[i]) } + command = append(command, args[i]) + } + if len(command) == 0 { + command = []string{"/bin/sh"} } + return taskName, command, false, nil +} - if len(cmdToRun) == 0 { - cmdToRun = []string{"/bin/sh"} +func runSSH(serverURL, atespace, kubeContext string, args []string) error { + taskName, cmdToRun, _, err := parseSSHArgs(args) + if err != nil { + return err } client, conn, err := getAXClient(serverURL) diff --git a/cmd/ax/main_test.go b/cmd/ax/main_test.go index d1c9d031..e295a1fd 100644 --- a/cmd/ax/main_test.go +++ b/cmd/ax/main_test.go @@ -19,6 +19,7 @@ import ( "net" "os" "path/filepath" + "reflect" "strings" "testing" @@ -206,3 +207,92 @@ func TestRunGetResourceAliases(t *testing.T) { } } } + +func TestParseGlobalArgs(t *testing.T) { + tests := []struct { + name string + args []string + want globalArgs + }{ + { + name: "flags before and after the command", + args: []string{"-a", "team", "get", "tasks", "--namespace=ax"}, + want: globalArgs{cmd: "get", args: []string{"tasks"}, atespace: "team", namespace: "ax"}, + }, + { + name: "flags after -- belong to the remote command", + args: []string{"ssh", "task123", "--", "grep", "-n", "foo", "-a", "file"}, + want: globalArgs{ + cmd: "ssh", + args: []string{"task123", "--", "grep", "-n", "foo", "-a", "file"}, + atespace: "default", + namespace: "ax-system", + }, + }, + { + name: "global flags before -- still apply", + args: []string{"ssh", "-a", "team", "task123", "--", "ls", "--context=x"}, + want: globalArgs{ + cmd: "ssh", + args: []string{"task123", "--", "ls", "--context=x"}, + atespace: "team", + namespace: "ax-system", + }, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := parseGlobalArgs(tt.args); !reflect.DeepEqual(got, tt.want) { + t.Errorf("parseGlobalArgs(%q) = %+v, want %+v", tt.args, got, tt.want) + } + }) + } +} + +func TestParseSSHArgs(t *testing.T) { + tests := []struct { + name string + args []string + wantTask string + wantCommand []string + wantHelp bool + wantErr bool + }{ + {name: "long help", args: []string{"--help"}, wantHelp: true}, + {name: "short help", args: []string{"-h"}, wantHelp: true}, + {name: "no args", args: nil, wantErr: true}, + {name: "no task before --", args: []string{"--", "ls"}, wantErr: true}, + {name: "unknown flag", args: []string{"--verbose", "task123"}, wantErr: true}, + {name: "default shell", args: []string{"task123"}, wantTask: "task123", wantCommand: []string{"/bin/sh"}}, + { + name: "command after --", + args: []string{"task123", "--", "ls", "-la", "/workspace"}, + wantTask: "task123", + wantCommand: []string{"ls", "-la", "/workspace"}, + }, + { + name: "command without --", + args: []string{"task123", "python3", "main.py"}, + wantTask: "task123", + wantCommand: []string{"python3", "main.py"}, + }, + { + name: "help after the task name goes to the remote command", + args: []string{"task123", "--", "git", "--help"}, + wantTask: "task123", + wantCommand: []string{"git", "--help"}, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + task, command, help, err := parseSSHArgs(tt.args) + if (err != nil) != tt.wantErr { + t.Fatalf("parseSSHArgs(%q) error = %v, wantErr %v", tt.args, err, tt.wantErr) + } + if task != tt.wantTask || help != tt.wantHelp || !reflect.DeepEqual(command, tt.wantCommand) { + t.Errorf("parseSSHArgs(%q) = (%q, %q, %v), want (%q, %q, %v)", + tt.args, task, command, help, tt.wantTask, tt.wantCommand, tt.wantHelp) + } + }) + } +}