From db6975eb7147e20e7b4d18ee3c8d9f69290f1636 Mon Sep 17 00:00:00 2001 From: harshitsinghbhandari <24b4506@iitb.ac.in> Date: Thu, 28 May 2026 16:18:27 +0530 Subject: [PATCH] fix: tighten zellij runtime attach and launch --- .../adapters/runtime/zellij/commands.go | 116 +++++++++++++++++- .../adapters/runtime/zellij/zellij.go | 47 +++++-- .../adapters/runtime/zellij/zellij_test.go | 94 +++++++++++++- 3 files changed, 244 insertions(+), 13 deletions(-) diff --git a/backend/internal/adapters/runtime/zellij/commands.go b/backend/internal/adapters/runtime/zellij/commands.go index 2a866a0d7..4cdf86542 100644 --- a/backend/internal/adapters/runtime/zellij/commands.go +++ b/backend/internal/adapters/runtime/zellij/commands.go @@ -68,15 +68,46 @@ func terminalPaneID(id int) string { } func buildLayout(cfg ports.RuntimeConfig, shellPath string) string { + spec := shellLaunchSpecFor(shellPath) + shellCommand := shellLaunchCommand(cfg, shellPath, spec) + return layoutString(cfg.WorkspacePath, shellPath, spec.args, shellCommand) +} + +type shellLaunchSpec struct { + args []string +} + +func shellLaunchSpecFor(shellPath string) shellLaunchSpec { + base := strings.ToLower(filepathBase(shellPath)) + if strings.Contains(base, "cmd") { + return shellLaunchSpec{args: []string{"/D", "/S", "/K"}} + } + if strings.Contains(base, "powershell") || strings.Contains(base, "pwsh") { + return shellLaunchSpec{args: []string{"-NoLogo", "-NoProfile", "-NoExit", "-Command"}} + } + return shellLaunchSpec{args: []string{"-lc"}} +} + +func layoutString(workspacePath, shellPath string, shellArgs []string, shellCommand string) string { return "layout {\n" + - " cwd " + kdlQuote(cfg.WorkspacePath) + "\n" + + " cwd " + kdlQuote(workspacePath) + "\n" + " pane command=" + kdlQuote(shellPath) + " name=" + kdlQuote(agentPaneName) + " {\n" + - " args " + kdlQuote("-lc") + " " + kdlQuote(wrapLaunchCommand(cfg, shellPath)) + "\n" + + " args " + kdlJoin(shellArgs) + " " + kdlQuote(shellCommand) + "\n" + " }\n" + "}\n" } -func wrapLaunchCommand(cfg ports.RuntimeConfig, shellPath string) string { +func shellLaunchCommand(cfg ports.RuntimeConfig, shellPath string, spec shellLaunchSpec) string { + if len(spec.args) > 0 && spec.args[0] == "-NoLogo" { + return wrapLaunchCommandPowerShell(cfg, shellPath) + } + if len(spec.args) > 0 && spec.args[0] == "/D" { + return wrapLaunchCommandCmd(cfg) + } + return wrapLaunchCommandUnix(cfg, shellPath) +} + +func wrapLaunchCommandUnix(cfg ports.RuntimeConfig, shellPath string) string { path := cfg.Env["PATH"] if path == "" { path = getenv("PATH") @@ -105,6 +136,58 @@ func wrapLaunchCommand(cfg ports.RuntimeConfig, shellPath string) string { return b.String() } +func wrapLaunchCommandPowerShell(cfg ports.RuntimeConfig, shellPath string) string { + path := cfg.Env["PATH"] + if path == "" { + path = getenv("PATH") + } + + var b strings.Builder + for _, key := range sortedKeys(cfg.Env) { + if key == "PATH" { + continue + } + b.WriteString("$env:") + b.WriteString(key) + b.WriteString(" = ") + b.WriteString(psQuote(cfg.Env[key])) + b.WriteString("; ") + } + if path != "" { + b.WriteString("$env:PATH = ") + b.WriteString(psQuote(path)) + b.WriteString("; ") + } + b.WriteString(cfg.LaunchCommand) + return b.String() +} + +func wrapLaunchCommandCmd(cfg ports.RuntimeConfig) string { + path := cfg.Env["PATH"] + if path == "" { + path = getenv("PATH") + } + + var b strings.Builder + for _, key := range sortedKeys(cfg.Env) { + if key == "PATH" { + continue + } + b.WriteString("set \"") + b.WriteString(key) + b.WriteString("=") + b.WriteString(cmdQuote(cfg.Env[key])) + b.WriteString("\" && ") + } + if path != "" { + b.WriteString("set \"PATH=") + b.WriteString(cmdQuote(path)) + b.WriteString("\" && ") + } + b.WriteString(cfg.LaunchCommand) + return b.String() +} + func sortedKeys(m map[string]string) []string { keys := make([]string, 0, len(m)) for k := range m { @@ -118,6 +201,33 @@ func shellQuote(s string) string { return "'" + strings.ReplaceAll(s, "'", "'\\''") + "'" } +func psQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", "''") + "'" +} + +func cmdQuote(s string) string { + return strings.ReplaceAll(s, "\"", "\"\"") +} + func kdlQuote(s string) string { return strconv.Quote(s) } + +func kdlJoin(args []string) string { + parts := make([]string, 0, len(args)) + for _, arg := range args { + parts = append(parts, kdlQuote(arg)) + } + return strings.Join(parts, " ") +} + +func filepathBase(path string) string { + if path == "" { + return "" + } + i := strings.LastIndexAny(path, `/\`) + if i < 0 { + return path + } + return path[i+1:] +} diff --git a/backend/internal/adapters/runtime/zellij/zellij.go b/backend/internal/adapters/runtime/zellij/zellij.go index 1d110748e..deeb38a3c 100644 --- a/backend/internal/adapters/runtime/zellij/zellij.go +++ b/backend/internal/adapters/runtime/zellij/zellij.go @@ -11,6 +11,7 @@ import ( "os" "os/exec" "regexp" + "runtime" "strconv" "strings" "time" @@ -81,7 +82,11 @@ func New(opts Options) *Runtime { shellPath = os.Getenv("SHELL") } if shellPath == "" { - shellPath = "/bin/sh" + if runtime.GOOS == "windows" { + shellPath = "powershell.exe" + } else { + shellPath = "/bin/sh" + } } chunkSize := opts.ChunkSize if chunkSize <= 0 { @@ -189,9 +194,12 @@ func (r *Runtime) AttachCommand(handle ports.RuntimeHandle) ([]string, error) { if err != nil { return nil, err } - args := append([]string{r.binary}, r.baseArgs()...) + args := append([]string{}, r.baseArgs()...) args = append(args, attachArgs(id)...) - return args, nil + if r.socketDir == "" { + return append([]string{r.binary}, args...), nil + } + return attachCommandWithEnv(r.binary, r.socketDir, args...), nil } func (r *Runtime) ensureSupportedVersion(ctx context.Context) error { @@ -232,14 +240,15 @@ func (r *Runtime) findAgentPane(ctx context.Context, id string) (string, error) var lastErr error for { out, err := r.run(ctx, listPanesArgs(id)...) - if err != nil { - return "", fmt.Errorf("zellij runtime: list panes %s: %w", id, err) - } - paneID, err := agentPaneID(out) if err == nil { - return paneID, nil + paneID, parseErr := agentPaneID(out) + if parseErr == nil { + return paneID, nil + } + lastErr = parseErr + } else { + lastErr = err } - lastErr = err if time.Now().After(deadline) { return "", fmt.Errorf("zellij runtime: list panes %s: %w", id, lastErr) } @@ -280,6 +289,26 @@ func (r *Runtime) env() []string { return []string{"ZELLIJ_SOCKET_DIR=" + r.socketDir} } +func attachCommandWithEnv(binary, socketDir string, args ...string) []string { + if socketDir == "" { + return append([]string{binary}, args...) + } + if runtime.GOOS == "windows" { + command := strings.Builder{} + command.WriteString("$env:ZELLIJ_SOCKET_DIR = ") + command.WriteString(psQuote(socketDir)) + command.WriteString("; & ") + command.WriteString(psQuote(binary)) + for _, arg := range args { + command.WriteByte(' ') + command.WriteString(psQuote(arg)) + } + return []string{"powershell.exe", "-NoLogo", "-NoProfile", "-Command", command.String()} + } + return append([]string{"env", "ZELLIJ_SOCKET_DIR=" + socketDir, binary}, args...) +} + + func zellijSessionName(id domain.SessionID) (string, error) { raw := string(id) if raw == "" { diff --git a/backend/internal/adapters/runtime/zellij/zellij_test.go b/backend/internal/adapters/runtime/zellij/zellij_test.go index 72ab19e29..a690af038 100644 --- a/backend/internal/adapters/runtime/zellij/zellij_test.go +++ b/backend/internal/adapters/runtime/zellij/zellij_test.go @@ -5,6 +5,7 @@ import ( "errors" "os/exec" "reflect" + "runtime" "strings" "testing" "time" @@ -15,7 +16,11 @@ import ( func TestNewDefaultsToPortableShell(t *testing.T) { t.Setenv("SHELL", "") r := New(Options{}) - if got, want := r.shell, "/bin/sh"; got != want { + want := "/bin/sh" + if runtime.GOOS == "windows" { + want = "powershell.exe" + } + if got := r.shell; got != want { t.Fatalf("default shell = %q, want %q", got, want) } } @@ -121,6 +126,56 @@ func TestBuildLayoutExportsEnvAndKeepsPaneAlive(t *testing.T) { } } +func TestBuildLayoutUsesPowerShellLaunchOnWindowsShells(t *testing.T) { + oldGetenv := getenv + getenv = func(key string) string { + if key == "PATH" { + return `C:\custom\bin` + } + return "" + } + defer func() { getenv = oldGetenv }() + + got := buildLayout(ports.RuntimeConfig{WorkspacePath: `C:\ws`, LaunchCommand: "Write-Host ready", Env: map[string]string{ + "AO_SESSION_ID": "sess-1", + }}, `C:\Program Files\PowerShell\7\pwsh.exe`) + + for _, want := range []string{ + `pane command="C:\\Program Files\\PowerShell\\7\\pwsh.exe" name="agent"`, + `args "-NoLogo" "-NoProfile" "-NoExit" "-Command"`, + "$env:AO_SESSION_ID = 'sess-1';", + "$env:PATH = ", + "Write-Host ready", + } { + if !strings.Contains(got, want) { + t.Fatalf("powershell layout missing %q in %q", want, got) + } + } +} + +func TestBuildLayoutUsesCmdLaunchOnCmdShells(t *testing.T) { + oldGetenv := getenv + getenv = func(key string) string { + return "" + } + defer func() { getenv = oldGetenv }() + + got := buildLayout(ports.RuntimeConfig{WorkspacePath: `C:\ws`, LaunchCommand: "echo ready", Env: map[string]string{ + "AO_SESSION_ID": "sess-1", + }}, `C:\Windows\System32\cmd.exe`) + + for _, want := range []string{ + `pane command="C:\\Windows\\System32\\cmd.exe" name="agent"`, + `args "/D" "/S" "/K"`, + `AO_SESSION_ID=sess-1`, + "echo ready", + } { + if !strings.Contains(got, want) { + t.Fatalf("cmd layout missing %q in %q", want, got) + } + } +} + func TestCreateStartsSessionAndDiscoversPane(t *testing.T) { fr := &fakeRunner{outputs: [][]byte{[]byte("zellij 0.44.3"), nil, []byte(`[{"id":0,"is_plugin":true,"title":"zellij:tab-bar"},{"id":3,"is_plugin":false,"title":"agent"}]`)}} r := New(Options{Binary: "zellij-test", Timeout: time.Second, Shell: "/bin/zsh", SocketDir: "/tmp/zj", ConfigDir: "/tmp/cfg"}) @@ -155,6 +210,43 @@ func TestCreateStartsSessionAndDiscoversPane(t *testing.T) { } } +func TestAttachCommandUsesSocketDir(t *testing.T) { + r := New(Options{SocketDir: "/tmp/zj"}) + args, err := r.AttachCommand(ports.RuntimeHandle{ID: "sess-1/terminal_0", RuntimeName: runtimeName}) + if err != nil { + t.Fatalf("AttachCommand: %v", err) + } + if runtime.GOOS == "windows" { + if len(args) < 4 || args[0] != "powershell.exe" { + t.Fatalf("attach command = %#v, want powershell wrapper", args) + } + if !strings.Contains(strings.Join(args, " "), "ZELLIJ_SOCKET_DIR") { + t.Fatalf("attach command missing socket dir env: %#v", args) + } + return + } + if got, want := args[:3], []string{"env", "ZELLIJ_SOCKET_DIR=/tmp/zj", r.binary}; !reflect.DeepEqual(got, want) { + t.Fatalf("attach prefix = %#v, want %#v", got, want) + } +} + +func TestFindAgentPaneRetriesTransientErrors(t *testing.T) { + fr := &fakeRunner{outputs: [][]byte{[]byte("boom"), []byte(`[{"id":0,"is_plugin":false,"title":"agent"}]`)}} + r := New(Options{Timeout: time.Second}) + r.runner = fr + + got, err := r.findAgentPane(context.Background(), "sess-1") + if err != nil { + t.Fatalf("findAgentPane: %v", err) + } + if got != "terminal_0" { + t.Fatalf("findAgentPane = %q, want terminal_0", got) + } + if len(fr.calls) != 2 { + t.Fatalf("calls = %d, want 2", len(fr.calls)) + } +} + func TestParseVersion(t *testing.T) { for _, tc := range []struct { in string