fix: address final container execution issues and cleanup review docs

1 files changed, 76 insertions, 22 deletions

diff --git a/internal/executor/container.go b/internal/executor/container.go
index d21aea3..45758d2 100644
--- a/internal/executor/container.go
+++ b/internal/executor/container.go
@@ -17,12 +17,23 @@ import (
 
 // ContainerRunner executes an agent inside a container.
 type ContainerRunner struct {
-	Image      string // default image if not specified in task
-	Logger     *slog.Logger
-	LogDir     string
-	APIURL     string
-	DropsDir   string
-	SSHAuthSock string // optional path to host SSH agent
+	Image        string // default image if not specified in task
+	Logger       *slog.Logger
+	LogDir       string
+	APIURL       string
+	DropsDir     string
+	SSHAuthSock  string // optional path to host SSH agent
+	ClaudeBinary string // optional path to claude binary in container
+	GeminiBinary string // optional path to gemini binary in container
+	// Command allows mocking exec.CommandContext for tests.
+	Command func(ctx context.Context, name string, arg ...string) *exec.Cmd
+}
+
+func (r *ContainerRunner) command(ctx context.Context, name string, arg ...string) *exec.Cmd {
+	if r.Command != nil {
+		return r.Command(ctx, name, arg...)
+	}
+	return exec.CommandContext(ctx, name, arg...)
 }
 
 func (r *ContainerRunner) ExecLogDir(execID string) string {
@@ -88,7 +99,11 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 	// 2. Clone repo into workspace if not resuming
 	if !isResume {
 		r.Logger.Info("cloning repository", "url", repoURL, "workspace", workspace)
-		if out, err := exec.CommandContext(ctx, "git", "clone", repoURL, workspace).CombinedOutput(); err != nil {
+		if out, err := r.command(ctx, "git", "clone", repoURL, workspace).CombinedOutput(); err != nil {
+			// If it looks like a remote URL, fail fast.
+			if strings.HasPrefix(repoURL, "http") || strings.HasPrefix(repoURL, "git@") || strings.HasPrefix(repoURL, "ssh://") {
+				return fmt.Errorf("git clone failed for remote repository: %w\n%s", err, string(out))
+			}
 			r.Logger.Warn("git clone failed, attempting fallback init", "url", repoURL, "error", err)
 			if initErr := r.fallbackGitInit(repoURL, workspace); initErr != nil {
 				return fmt.Errorf("git clone and fallback init failed: %w\n%s", err, string(out))
@@ -143,7 +158,7 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 	fullArgs = append(fullArgs, innerCmd...)
 
 	r.Logger.Info("starting container", "image", image, "taskID", t.ID)
-	cmd := exec.CommandContext(ctx, "docker", fullArgs...)
+	cmd := r.command(ctx, "docker", fullArgs...)
 	cmd.Stderr = stderrFile
 	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 
@@ -162,6 +177,18 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 	}
 	stdoutW.Close()
 
+	// Watch for context cancellation to kill the process group (Issue 1)
+	done := make(chan struct{})
+	defer close(done)
+	go func() {
+		select {
+		case <-ctx.Done():
+			r.Logger.Info("killing container process group due to context cancellation", "taskID", t.ID)
+			syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+		case <-done:
+		}
+	}()
+
 	// Stream stdout to the log file and parse cost/errors.
 	var costUSD float64
 	var sessionID string
@@ -193,6 +220,9 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 		} else {
 			isBlocked = true
 			success = true // We consider BLOCKED as a "success" for workspace preservation
+			if e.SessionID == "" {
+				r.Logger.Warn("missing session ID; resume will start fresh", "taskID", e.TaskID)
+			}
 			return &BlockedError{
 				QuestionJSON: questionJSON,
 				SessionID:    e.SessionID,
@@ -210,14 +240,24 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 
 	// 5. Post-execution: push changes if successful
 	if waitErr == nil && streamErr == nil {
-		r.Logger.Info("pushing changes back to remote", "url", repoURL)
-		// We assume the sandbox has committed changes (the agent image should enforce this)
-		if out, err := exec.CommandContext(ctx, "git", "-C", workspace, "push", "origin", "HEAD").CombinedOutput(); err != nil {
-			r.Logger.Warn("git push failed or no changes", "error", err, "output", string(out))
-			// Only set success = true if we consider this "good enough".
-			// Review says: "If the remote is missing or the push fails, the task is marked FAILED and the host-side workspace is preserved"
-			// So we MUST return error here.
-			return fmt.Errorf("git push failed: %w\n%s", err, string(out))
+		// Check if there are any commits to push (Issue 10)
+		// We use rev-list to see if HEAD is ahead of origin/HEAD.
+		// If origin/HEAD doesn't exist (e.g. fresh init), we just attempt to push.
+		hasCommits := true
+		if out, err := r.command(ctx, "git", "-C", workspace, "rev-list", "origin/HEAD..HEAD").CombinedOutput(); err == nil {
+			if len(strings.TrimSpace(string(out))) == 0 {
+				hasCommits = false
+			}
+		}
+
+		if hasCommits {
+			r.Logger.Info("pushing changes back to remote", "url", repoURL)
+			if out, err := r.command(ctx, "git", "-C", workspace, "push", "origin", "HEAD").CombinedOutput(); err != nil {
+				r.Logger.Warn("git push failed", "error", err, "output", string(out))
+				return fmt.Errorf("git push failed: %w\n%s", err, string(out))
+			}
+		} else {
+			r.Logger.Info("no new commits to push", "taskID", t.ID)
 		}
 		success = true
 	}
@@ -235,7 +275,7 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 func (r *ContainerRunner) buildDockerArgs(workspace, taskID string) []string {
 	// --env-file takes a HOST path.
 	hostEnvFile := filepath.Join(workspace, ".claudomator-env")
-	return []string{
+	args := []string{
 		"run", "--rm",
 		"-v", workspace + ":/workspace",
 		"-w", "/workspace",
@@ -244,28 +284,42 @@ func (r *ContainerRunner) buildDockerArgs(workspace, taskID string) []string {
 		"-e", "CLAUDOMATOR_TASK_ID=" + taskID,
 		"-e", "CLAUDOMATOR_DROP_DIR=" + r.DropsDir,
 	}
+	if r.SSHAuthSock != "" {
+		args = append(args, "-v", r.SSHAuthSock+":/tmp/ssh-auth.sock", "-e", "SSH_AUTH_SOCK=/tmp/ssh-auth.sock")
+	}
+	return args
 }
 
 func (r *ContainerRunner) buildInnerCmd(t *task.Task, e *storage.Execution, isResume bool) []string {
 	// Claude CLI uses -p for prompt text. To pass a file, we use a shell to cat it.
 	// We use a shell variable to capture the expansion to avoid quoting issues with instructions contents.
 	// The outer single quotes around the sh -c argument prevent host-side expansion.
-	
+
+	claudeBin := r.ClaudeBinary
+	if claudeBin == "" {
+		claudeBin = "claude"
+	}
+	geminiBin := r.GeminiBinary
+	if geminiBin == "" {
+		geminiBin = "gemini"
+	}
+
 	if t.Agent.Type == "gemini" {
-		return []string{"sh", "-c", "INST=$(cat /workspace/.claudomator-instructions.txt); gemini -p \"$INST\""}
+		return []string{"sh", "-c", fmt.Sprintf("INST=$(cat /workspace/.claudomator-instructions.txt); %s -p \"$INST\"", geminiBin)}
 	}
 
 	// Claude
 	var claudeCmd strings.Builder
-	claudeCmd.WriteString("INST=$(cat /workspace/.claudomator-instructions.txt); claude -p \"$INST\"")
+	claudeCmd.WriteString(fmt.Sprintf("INST=$(cat /workspace/.claudomator-instructions.txt); %s -p \"$INST\"", claudeBin))
 	if isResume && e.ResumeSessionID != "" {
 		claudeCmd.WriteString(fmt.Sprintf(" --resume %s", e.ResumeSessionID))
 	}
 	claudeCmd.WriteString(" --output-format stream-json --verbose --permission-mode bypassPermissions")
-	
+
 	return []string{"sh", "-c", claudeCmd.String()}
 }
 
+
 func (r *ContainerRunner) fallbackGitInit(repoURL, workspace string) error {
 	// Ensure directory exists
 	if err := os.MkdirAll(workspace, 0755); err != nil {
@@ -281,7 +335,7 @@ func (r *ContainerRunner) fallbackGitInit(repoURL, workspace string) error {
 	// git clone handle local paths fine if they are repos.
 	// This fallback is only if it's NOT a repo.
 	for _, args := range cmds {
-		if out, err := exec.Command("git", args...).CombinedOutput(); err != nil {
+		if out, err := r.command(context.Background(), "git", args...).CombinedOutput(); err != nil {
 			return fmt.Errorf("git init failed: %w\n%s", err, out)
 		}
 	}