2 files changed, 132 insertions, 9 deletions
diff --git a/internal/executor/container.go b/internal/executor/container.go
index d90a273..61ac29c 100644
--- a/internal/executor/container.go
+++ b/internal/executor/container.go
@@ -58,6 +58,48 @@ func (r *ContainerRunner) ExecLogDir(execID string) string {
 	return filepath.Join(r.LogDir, execID)
 }
 
+// ensureStoryBranch checks whether branchName exists in remoteURL and creates
+// it from main if not. Uses localPath as a reference clone for speed if set.
+func (r *ContainerRunner) ensureStoryBranch(ctx context.Context, remoteURL, branchName, localPath string) error {
+	// Check if branch already exists.
+	out, err := r.command(ctx, "git", "ls-remote", "--heads", remoteURL, branchName).CombinedOutput()
+	if err == nil && len(strings.TrimSpace(string(out))) > 0 {
+		return nil // already exists
+	}
+
+	r.Logger.Info("story branch missing, creating from main", "branch", branchName, "remote", remoteURL)
+
+	// Clone into a temp dir so we can create the branch.
+	tmp, err := os.MkdirTemp("", "claudomator-branchsetup-*")
+	if err != nil {
+		return fmt.Errorf("mktemp for branch setup: %w", err)
+	}
+	defer os.RemoveAll(tmp)
+
+	// Remove the dir git clone expects to create.
+	if err := os.Remove(tmp); err != nil {
+		return fmt.Errorf("removing tmp dir before clone: %w", err)
+	}
+
+	var cloneArgs []string
+	if localPath != "" {
+		cloneArgs = []string{"clone", "--reference", localPath, remoteURL, tmp}
+	} else {
+		cloneArgs = []string{"clone", remoteURL, tmp}
+	}
+	if out, err := r.command(ctx, "git", cloneArgs...).CombinedOutput(); err != nil {
+		return fmt.Errorf("git clone for branch setup: %w\n%s", err, string(out))
+	}
+	if out, err := r.command(ctx, "git", "-C", tmp, "checkout", "-b", branchName).CombinedOutput(); err != nil {
+		return fmt.Errorf("git checkout -b %q: %w\n%s", branchName, err, string(out))
+	}
+	if out, err := r.command(ctx, "git", "-C", tmp, "push", "origin", branchName).CombinedOutput(); err != nil {
+		return fmt.Errorf("git push %q: %w\n%s", branchName, err, string(out))
+	}
+	r.Logger.Info("story branch created and pushed", "branch", branchName)
+	return nil
+}
+
 func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Execution) error {
 	var err error
 	repoURL := t.RepositoryURL
@@ -114,7 +156,16 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 		storyBranch = t.BranchName
 	}
 
-	// 2. Clone repo into workspace if not resuming.
+	// 2. Ensure story branch exists in the remote before cloning.
+	// If the branch is missing (e.g. story approved before fix, or branch push failed),
+	// create it from main using the project local path as a reference repo.
+	if storyBranch != "" && !isResume {
+		if err := r.ensureStoryBranch(ctx, repoURL, storyBranch, storyLocalPath); err != nil {
+			r.Logger.Warn("ensureStoryBranch failed (will attempt checkout anyway)", "branch", storyBranch, "error", err)
+		}
+	}
+
+	// 3. Clone repo into workspace if not resuming.
 	// git clone requires the target directory to not exist; remove the MkdirTemp-created dir first.
 	if !isResume {
 		if err := os.Remove(workspace); err != nil {
@@ -133,14 +184,7 @@ func (r *ContainerRunner) Run(ctx context.Context, t *task.Task, e *storage.Exec
 		if storyBranch != "" {
 			r.Logger.Info("checking out story branch", "branch", storyBranch)
 			if out, err := r.command(ctx, "git", "-C", workspace, "checkout", storyBranch).CombinedOutput(); err != nil {
-				// Branch doesn't exist in the remote yet — create it from HEAD and push.
-				r.Logger.Warn("story branch not found, creating from HEAD", "branch", storyBranch)
-				if out2, err2 := r.command(ctx, "git", "-C", workspace, "checkout", "-b", storyBranch).CombinedOutput(); err2 != nil {
-					return fmt.Errorf("git checkout story branch %q failed: %w\n%s\ncreate attempt: %s", storyBranch, err, string(out), string(out2))
-				}
-				if out2, err2 := r.command(ctx, "git", "-C", workspace, "push", "origin", storyBranch).CombinedOutput(); err2 != nil {
-					r.Logger.Warn("push of auto-created story branch failed", "branch", storyBranch, "error", err2, "output", string(out2))
-				}
+				return fmt.Errorf("git checkout story branch %q failed: %w\n%s", storyBranch, err, string(out))
 			}
 		}
 		if err = os.Chmod(workspace, 0755); err != nil {
diff --git a/internal/executor/container_test.go b/internal/executor/container_test.go
index 15c147f..f840f85 100644
--- a/internal/executor/container_test.go
+++ b/internal/executor/container_test.go
@@ -606,3 +606,82 @@ func TestContainerRunner_ClonesDefaultBranchWhenNoBranchName(t *testing.T) {
 		}
 	}
 }
+
+func TestEnsureStoryBranch_CreatesMissingBranch(t *testing.T) {
+	// Set up a bare repo and a local clone to test branch creation.
+	dir := t.TempDir()
+	bare := filepath.Join(dir, "bare.git")
+	local := filepath.Join(dir, "local")
+
+	// Create bare repo with an initial commit.
+	if out, err := exec.Command("git", "init", "--bare", bare).CombinedOutput(); err != nil {
+		t.Fatalf("git init bare: %v\n%s", err, out)
+	}
+	if out, err := exec.Command("git", "clone", bare, local).CombinedOutput(); err != nil {
+		t.Fatalf("git clone: %v\n%s", err, out)
+	}
+	if out, err := exec.Command("git", "-C", local, "commit", "--allow-empty", "-m", "init").CombinedOutput(); err != nil {
+		t.Fatalf("git commit: %v\n%s", err, out)
+	}
+	if out, err := exec.Command("git", "-C", local, "push", "origin", "main").CombinedOutput(); err != nil {
+		// try master
+		if out2, err2 := exec.Command("git", "-C", local, "push", "origin", "HEAD:main").CombinedOutput(); err2 != nil {
+			t.Fatalf("git push main: %v\n%s\n%s", err, out, out2)
+		}
+	}
+
+	runner := &ContainerRunner{Logger: slog.Default()}
+
+	branch := "story/test-branch"
+
+	// Branch should not exist yet.
+	out, _ := exec.Command("git", "ls-remote", "--heads", bare, branch).CombinedOutput()
+	if len(strings.TrimSpace(string(out))) > 0 {
+		t.Fatal("branch should not exist before ensureStoryBranch")
+	}
+
+	if err := runner.ensureStoryBranch(context.Background(), bare, branch, ""); err != nil {
+		t.Fatalf("ensureStoryBranch: %v", err)
+	}
+
+	// Branch should now exist in the bare repo.
+	out, err := exec.Command("git", "ls-remote", "--heads", bare, branch).CombinedOutput()
+	if err != nil || len(strings.TrimSpace(string(out))) == 0 {
+		t.Errorf("branch %q not found in bare repo after ensureStoryBranch: %s", branch, out)
+	}
+}
+
+func TestEnsureStoryBranch_IdempotentIfExists(t *testing.T) {
+	dir := t.TempDir()
+	bare := filepath.Join(dir, "bare.git")
+	local := filepath.Join(dir, "local")
+
+	if out, err := exec.Command("git", "init", "--bare", bare).CombinedOutput(); err != nil {
+		t.Fatalf("git init bare: %v\n%s", err, out)
+	}
+	if out, err := exec.Command("git", "clone", bare, local).CombinedOutput(); err != nil {
+		t.Fatalf("git clone: %v\n%s", err, out)
+	}
+	if out, err := exec.Command("git", "-C", local, "commit", "--allow-empty", "-m", "init").CombinedOutput(); err != nil {
+		t.Fatalf("git commit: %v\n%s", err, out)
+	}
+	if _, err := exec.Command("git", "-C", local, "push", "origin", "HEAD:main").CombinedOutput(); err != nil {
+		t.Fatalf("push main: %v", err)
+	}
+
+	branch := "story/existing-branch"
+	// Pre-create the branch.
+	if out, err := exec.Command("git", "-C", local, "checkout", "-b", branch).CombinedOutput(); err != nil {
+		t.Fatalf("checkout -b: %v\n%s", err, out)
+	}
+	if out, err := exec.Command("git", "-C", local, "push", "origin", branch).CombinedOutput(); err != nil {
+		t.Fatalf("push branch: %v\n%s", err, out)
+	}
+
+	runner := &ContainerRunner{Logger: slog.Default()}
+
+	// Should be a no-op, not an error.
+	if err := runner.ensureStoryBranch(context.Background(), bare, branch, ""); err != nil {
+		t.Fatalf("ensureStoryBranch on existing branch: %v", err)
+	}
+}