bkinnightskytw
/
agent-orchestrator
mirror of https://github.com/ComposioHQ/agent-orchestrator.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: tighten zellij runtime attach and launch

This commit is contained in:
harshitsinghbhandari 2026-05-28 16:18:27 +05:30
parent 586931b3ff
commit db6975eb71
3 changed files with 244 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

116
backend/internal/adapters/runtime/zellij/commands.go
View File

@ -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:]
}

47
backend/internal/adapters/runtime/zellij/zellij.go
View File

@ -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 == "" {

94
backend/internal/adapters/runtime/zellij/zellij_test.go
View File

@ -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