From a4795d68fc5381f1ff48d043fe7554355e5899fb Mon Sep 17 00:00:00 2001
From: Peter Stone <thepeterstone@gmail.com>
Date: Wed, 18 Mar 2026 07:54:27 +0000
Subject: fix: address final container execution issues and cleanup review docs

---
 internal/executor/container_test.go | 52 ++++++++++++++++++++++++++++++-------
 1 file changed, 43 insertions(+), 9 deletions(-)

(limited to 'internal/executor/container_test.go')

diff --git a/internal/executor/container_test.go b/internal/executor/container_test.go
index 0e36def..d4d591e 100644
--- a/internal/executor/container_test.go
+++ b/internal/executor/container_test.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"log/slog"
 	"os"
+	"os/exec"
 	"strings"
 	"testing"
 
@@ -15,14 +16,15 @@ import (
 
 func TestContainerRunner_BuildDockerArgs(t *testing.T) {
 	runner := &ContainerRunner{
-		APIURL:   "http://localhost:8484",
-		DropsDir: "/data/drops",
+		APIURL:      "http://localhost:8484",
+		DropsDir:    "/data/drops",
+		SSHAuthSock: "/tmp/ssh.sock",
 	}
 	workspace := "/tmp/ws"
 	taskID := "task-123"
 
 	args := runner.buildDockerArgs(workspace, taskID)
-	
+
 	expected := []string{
 		"run", "--rm",
 		"-v", "/tmp/ws:/workspace",
@@ -31,11 +33,12 @@ func TestContainerRunner_BuildDockerArgs(t *testing.T) {
 		"-e", "CLAUDOMATOR_API_URL=http://localhost:8484",
 		"-e", "CLAUDOMATOR_TASK_ID=task-123",
 		"-e", "CLAUDOMATOR_DROP_DIR=/data/drops",
+		"-v", "/tmp/ssh.sock:/tmp/ssh-auth.sock",
+		"-e", "SSH_AUTH_SOCK=/tmp/ssh-auth.sock",
 	}
 
-
 	if len(args) != len(expected) {
-		t.Fatalf("expected %d args, got %d", len(expected), len(args))
+		t.Fatalf("expected %d args, got %d. Got: %v", len(expected), len(args), args)
 	}
 	for i, v := range args {
 		if v != expected[i] {
@@ -76,12 +79,31 @@ func TestContainerRunner_BuildInnerCmd(t *testing.T) {
 		tk := &task.Task{Agent: task.AgentConfig{Type: "gemini"}}
 		exec := &storage.Execution{}
 		cmd := runner.buildInnerCmd(tk, exec, false)
-		
+
 		cmdStr := strings.Join(cmd, " ")
 		if !strings.Contains(cmdStr, "gemini -p \"$INST\"") {
 			t.Errorf("expected gemini command with safer quoting, got %q", cmdStr)
 		}
 	})
+
+	t.Run("custom-binaries", func(t *testing.T) {
+		runnerCustom := &ContainerRunner{
+			ClaudeBinary: "/usr/bin/claude-v2",
+			GeminiBinary: "/usr/local/bin/gemini-pro",
+		}
+		
+		tkClaude := &task.Task{Agent: task.AgentConfig{Type: "claude"}}
+		cmdClaude := runnerCustom.buildInnerCmd(tkClaude, &storage.Execution{}, false)
+		if !strings.Contains(strings.Join(cmdClaude, " "), "/usr/bin/claude-v2 -p") {
+			t.Errorf("expected custom claude binary, got %q", cmdClaude)
+		}
+
+		tkGemini := &task.Task{Agent: task.AgentConfig{Type: "gemini"}}
+		cmdGemini := runnerCustom.buildInnerCmd(tkGemini, &storage.Execution{}, false)
+		if !strings.Contains(strings.Join(cmdGemini, " "), "/usr/local/bin/gemini-pro -p") {
+			t.Errorf("expected custom gemini binary, got %q", cmdGemini)
+		}
+	})
 }
 
 func TestContainerRunner_Run_PreservesWorkspaceOnFailure(t *testing.T) {
@@ -89,19 +111,31 @@ func TestContainerRunner_Run_PreservesWorkspaceOnFailure(t *testing.T) {
 	runner := &ContainerRunner{
 		Logger: logger,
 		Image:  "busybox",
+		Command: func(ctx context.Context, name string, arg ...string) *exec.Cmd {
+			// Mock docker run to exit 1
+			if name == "docker" {
+				return exec.Command("sh", "-c", "exit 1")
+			}
+			// Mock git clone to succeed and create the directory
+			if name == "git" && len(arg) > 0 && arg[0] == "clone" {
+				dir := arg[len(arg)-1]
+				os.MkdirAll(dir, 0755)
+				return exec.Command("true")
+			}
+			return exec.Command("true")
+		},
 	}
 
-	// Use an invalid repo URL to trigger failure.
 	tk := &task.Task{
 		ID:            "test-task",
-		RepositoryURL: "/nonexistent/repo",
+		RepositoryURL: "https://github.com/example/repo.git",
 		Agent:         task.AgentConfig{Type: "claude"},
 	}
 	exec := &storage.Execution{ID: "test-exec", TaskID: "test-task"}
 
 	err := runner.Run(context.Background(), tk, exec)
 	if err == nil {
-		t.Fatal("expected error due to invalid repo")
+		t.Fatal("expected error due to mocked docker failure")
 	}
 
 	// Verify SandboxDir was set and directory exists.
-- 
cgit v1.2.3