manager.py

import argparse
import pathlib
from datetime import datetime
import subprocess
import json
import time

from agent_common import ensure_dir, write_text


BANNED_SUBSTRINGS = [
    "```",
    "Test plan",
    "Acceptance criteria",
    "Here is the code",
    "```diff",
]


ALLOWED_PATCH_FILES = {"scripts/evcore_loop.py"}


def sh(cmd, cwd=None):
    r = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
    return r.returncode, r.stdout, r.stderr


def log_event(run_dir: pathlib.Path, msg: str) -> None:
    ts = datetime.now().isoformat(timespec="seconds")
    (run_dir / "events.log").open("a", encoding="utf-8").write(f"[{ts}] {msg}\n")


def write_status(run_dir: pathlib.Path, **fields) -> None:
    p = run_dir / "status.json"
    data = {}
    if p.exists():
        try:
            data = json.loads(p.read_text(encoding="utf-8"))
        except Exception:
            data = {}
    data.update(fields)
    data["updated_ts"] = time.time()
    p.write_text(json.dumps(data, indent=2, sort_keys=True), encoding="utf-8")


def run_worker(role, goal, context_path, out_path, model=None):
    cmd = [
        "python",
        "worker.py",
        "--role",
        role,
        "--goal",
        goal,
        "--context",
        str(context_path),
        "--out",
        str(out_path),
    ]
    if model:
        cmd += ["--model", model]

    rc, out, err = sh(cmd, cwd=pathlib.Path(__file__).parent)
    if rc != 0:
        raise RuntimeError(f"worker failed role={role}\nSTDOUT:\n{out}\nSTDERR:\n{err}")


def strip_markdown_fences_inplace(patch_path: pathlib.Path) -> bool:
    """
    Removes common markdown code fences from patch output:
      ```diff
      ...
      ```
    Also removes any line starting with ``` anywhere.
    Returns True if it changed the file.
    """
    txt = patch_path.read_text(encoding="utf-8", errors="replace")
    lines = txt.splitlines()

    new_lines = []
    changed = False

    for line in lines:
        s = line.strip()
        if s.startswith("```"):
            changed = True
            continue
        new_lines.append(line)

    # normalise final newline
    new_txt = "\n".join(new_lines).lstrip("\n")
    if new_txt and not new_txt.endswith("\n"):
        new_txt += "\n"

    if new_txt != txt:
        patch_path.write_text(new_txt, encoding="utf-8")
        return True
    return changed


def normalise_patch_inplace(patch_path: pathlib.Path) -> bool:
    """
    Make patch git-apply friendly:
    - remove markdown fences
    - if diff header exists: drop anything before it and keep only that one file section
    - else if a hunk exists (@@): drop anything before first @@ and synthesize header for scripts/evcore_loop.py
    - ensure trailing newline
    """
    original = patch_path.read_text(encoding="utf-8", errors="replace")
    lines = original.splitlines(True)

    # 1) Strip common markdown fences
    cleaned = []
    for ln in lines:
        s = ln.strip()
        if s in ("```", "```diff", "```patch"):
            continue
        cleaned.append(ln)
    lines = cleaned

    header = "diff --git a/scripts/evcore_loop.py b/scripts/evcore_loop.py"

    # 2) If we have a proper diff header, keep only that section
    start = None
    for i, ln in enumerate(lines):
        if ln.startswith(header):
            start = i
            break

    if start is not None:
        lines = lines[start:]

        # Drop any subsequent diff sections (only one file allowed)
        out = []
        seen_first = False
        for ln in lines:
            if ln.startswith("diff --git "):
                if not seen_first:
                    seen_first = True
                    out.append(ln)
                else:
                    break
            else:
                out.append(ln)

        txt = "".join(out)
        if txt and not txt.endswith("\n"):
            txt += "\n"

        changed = (txt != original)
        patch_path.write_text(txt, encoding="utf-8")
        return changed

    # 3) No diff header: if there is a hunk, synthesise header and keep from first @@
    hunk_i = None
    for i, ln in enumerate(lines):
        if ln.startswith("@@"):
            hunk_i = i
            break

    if hunk_i is None:
        # Nothing we can do; leave as-is
        patch_path.write_text("".join(lines), encoding="utf-8")
        return False

    out = []
    out.append(header + "\n")
    out.append("--- a/scripts/evcore_loop.py\n")
    out.append("+++ b/scripts/evcore_loop.py\n")
    out.extend(lines[hunk_i:])

    txt = "".join(out)
    if txt and not txt.endswith("\n"):
        txt += "\n"

    changed = (txt != original)
    patch_path.write_text(txt, encoding="utf-8")
    return changed


def ensure_patch_is_safe(patch_path: pathlib.Path) -> None:
    txt = patch_path.read_text(encoding="utf-8", errors="replace")

    # Ensure header is truly the first line (after hygiene)
    if txt.startswith("\ufeff"):
        txt = txt.lstrip("\ufeff")

    # Reject patches that try to create/delete the file or use /dev/null headers
    banned_structural = (
        "new file mode",
        "deleted file mode",
        "--- /dev/null",
        "+++ /dev/null",
    )
    for s in banned_structural:
        if s in txt:
            raise RuntimeError(f"Patch rejected: not allowed structural change: {s}")

    # Catch common empty-file / rewrite headers
    if "index 0000000..0000000" in txt:
        raise RuntimeError("Patch rejected: suspicious empty-file header (index 0000000..0000000)")

    # Reject patches with no hunks (no real changes)
    if "\n@@ " not in txt:
        raise RuntimeError("Patch rejected: no @@ hunks (empty/no-op patch)")

    for bad in BANNED_SUBSTRINGS:
        if bad in txt:
            raise RuntimeError(f"Patch rejected: contains banned text: {bad}")

    if not txt.startswith("diff --git a/scripts/evcore_loop.py b/scripts/evcore_loop.py"):
        raise RuntimeError("Patch rejected: must start with diff --git for scripts/evcore_loop.py")

    if "index e69de29" in txt:
        raise RuntimeError("Patch rejected: looks like an empty-file/full-rewrite patch (index e69de29)")

    # Parse touched files from diff headers
    touched = set()
    for line in txt.splitlines():
        if line.startswith("diff --git "):
            parts = line.strip().split()
            if len(parts) >= 4:
                a_path = parts[2].removeprefix("a/")
                b_path = parts[3].removeprefix("b/")
                touched.add(a_path)
                touched.add(b_path)

    touched = {t for t in touched if t}  # remove empties
    if touched != ALLOWED_PATCH_FILES:
        raise RuntimeError(f"Patch rejected: touched files {sorted(touched)} not allowed (allowed: {sorted(ALLOWED_PATCH_FILES)})")


def ensure_patch_matches_goal(goal: str, patch_path: pathlib.Path) -> None:
    """
    Extra safety layer:
    Ensure the patch context relates to the goal, without rejecting valid tiny hunks.
    """
    txt = patch_path.read_text(encoding="utf-8", errors="replace")
    goal_l = goal.lower()

    # Specific /status goal
    if "/status" in goal_l:
        if not any(
            marker in txt
            for marker in [
                '"/status"',
                "/status",
                "Mode=",
                "memory_turns",
                "Status:",
            ]
        ):
            raise RuntimeError(
                "Patch rejected: goal references /status but patch does not contain matching /status context."
            )
        return

    # Specific /reset goal
    if "/reset" in goal_l:
        # Must clearly relate to reset logic
        if not any(
            marker in txt
            for marker in [
                '"/reset"',
                "/reset",
                "Memory cleared.",
                "Memory and context cleared.",
                "Reset:",
                "clear_memory()",
                "clear_pending()",
            ]
        ):
            raise RuntimeError(
                "Patch rejected: goal references /reset but patch does not contain matching /reset context."
            )

        if "_is_exit" in txt or "Goodbye!" in txt or "Bye!" in txt:
            raise RuntimeError(
                "Patch rejected: goal references /reset but patch appears to modify exit logic instead."
            )

        # Must NOT be changing the exit/bye path
        if "Bye!" in txt or "Reset: Bye!" in txt:
            raise RuntimeError(
                "Patch rejected: goal references /reset but patch appears to modify exit/bye logic instead."
            )
        return

    # Generic CLI / interactive goals only when there is no specific command anchor
    if "cli" in goal_l or "interactive" in goal_l:
        if not any(
            marker in txt
            for marker in [
                "user_text.lower()",
                "while True",
                "input(",
                "Mode=",
                "memory_turns",
                "Memory cleared.",
                "Memory and context cleared.",
                "Reset:",
                "Status:",
            ]
        ):
            raise RuntimeError(
                "Patch rejected: goal references CLI handler but patch does not touch recognizable CLI logic."
            )


def repair_hunk_prefixes_inplace(patch_path: pathlib.Path) -> bool:
    """
    Some models output hunk context lines without the required leading ' '.
    This rewrites any such line inside @@ hunks to be prefixed with a single space.
    Returns True if it changed the file.
    """
    txt = patch_path.read_text(encoding="utf-8", errors="replace").splitlines(True)
    out = []
    changed = False
    in_hunk = False

    for line in txt:
        if line.startswith("diff --git ") or line.startswith("index ") or line.startswith("--- ") or line.startswith("+++ "):
            in_hunk = False
            out.append(line)
            continue

        if line.startswith("@@"):
            in_hunk = True
            out.append(line)
            continue

        if in_hunk:
            if line.startswith((" ", "+", "-", "\\")):
                out.append(line)
            else:
                # Illegal inside hunk -> treat as context line
                out.append(" " + line)
                changed = True
        else:
            out.append(line)

    if changed:
        patch_path.write_text("".join(out), encoding="utf-8")
    return changed


def is_noop_patch(patch_path: pathlib.Path) -> bool:
    txt = patch_path.read_text(encoding="utf-8", errors="replace")
    # A real patch must contain at least one hunk header
    return "\n@@ " not in txt


def extract_anchor_window(text: str, anchor: str, radius: int = 80) -> str | None:
    """
    Return a small excerpt around the first occurrence of `anchor`.
    radius is measured in lines above/below the anchor line.
    """
    lines = text.splitlines()
    for i, ln in enumerate(lines):
        if anchor in ln:
            lo = max(0, i - radius)
            hi = min(len(lines), i + radius)
            return "\n".join(lines[lo:hi]) + "\n"
    return None


def extract_line_window(text: str, line_no_1based: int, radius: int = 80) -> str | None:
    lines = text.splitlines()
    if line_no_1based < 1 or line_no_1based > len(lines):
        return None
    start = max(0, line_no_1based - 1 - radius)
    end = min(len(lines), line_no_1based - 1 + radius + 1)
    return "\n".join(lines[start:end]) + "\n"


def parse_apply_failure_line(stderr: str) -> int | None:
    import re
    m = re.search(r"scripts/evcore_loop\.py:(\d+)", stderr or "")
    return int(m.group(1)) if m else None


def detect_goal_already_satisfied(goal: str, target_text: str) -> str | None:
    """
    Return a human-readable reason if the requested goal is already satisfied.
    Otherwise return None.
    """
    goal_l = goal.lower()

    if "/status" in goal_l and "status:" in goal_l:
        if 'speak(f"Status: Mode={mode}, memory_turns={len(history)//2}.")' in target_text:
            return "The /status CLI message already starts with 'Status:'."

    if "/reset" in goal_l and "reset:" in goal_l:
        if 'speak("Reset: Memory cleared.")' in target_text:
            return "The /reset CLI message already starts with 'Reset:'."

    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--goal", required=True, help="High-level goal for this run")
    ap.add_argument("--tag", default="run", help="Short tag for folder name")
    ap.add_argument("--model-fast", default=None, help="Model for fast roles")
    ap.add_argument("--model-build", default=None, help="Model for builder role")
    ap.add_argument(
        "--apply",
        action="store_true",
        help="If set, allow interactive approval + merge into current branch. Default: preview only.",
    )
    args = ap.parse_args()

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    run_dir = pathlib.Path("runs") / f"{ts}_{args.tag}"
    ensure_dir(run_dir)

    write_status(
        run_dir,
        state="created",
        tag=args.tag,
        goal=args.goal,
        started_ts=time.time(),
    )
    log_event(run_dir, "Run created")

    repo_root = pathlib.Path(__file__).parent.parent
    agents_dir = pathlib.Path(__file__).parent

    # Create a smarter builder context: prefer a small excerpt around an anchor
    target_file = repo_root / "scripts" / "evcore_loop.py"
    target_text = target_file.read_text(encoding="utf-8", errors="replace")

    already_ok = detect_goal_already_satisfied(args.goal, target_text)
    if already_ok:
        log_event(run_dir, f"no_changes: {already_ok}")
        write_status(run_dir, state="no_changes", finished_ts=time.time())
        print(f"\nNO CHANGES ✅ ({already_ok})\nRun folder: {run_dir}\n")
        return

    builder_context_path = run_dir / "builder_context.md"
    goal_l = args.goal.lower()

    # Heuristic anchors (add more over time as you discover failure modes)
    anchors: list[str] = []

    if "/status" in goal_l:
        # Prefer CLI /status first; include Flask route as fallback
        anchors += ['user_text.lower() == "/status"', '@app.route("/status"']

    if "/reset" in goal_l:
        anchors += ['user_text.lower() == "/reset"']

    # Generic CLI hints
    if "cli" in goal_l or "interactive" in goal_l:
        anchors += ["while True", "input(", "user_text ="]

    excerpt = None
    chosen = None
    for a in anchors:
        excerpt = extract_anchor_window(target_text, a, radius=80)
        if excerpt:
            chosen = a
            break

    if chosen:
        log_event(run_dir, f"builder_context: anchor={chosen}")
    else:
        log_event(run_dir, "builder_context: anchor=None (full file)")

    parts = [
        f"GOAL:\n{args.goal}\n",
        "TARGET FILE: scripts/evcore_loop.py",
    ]

    if excerpt:
        parts += [
            f"ANCHOR USED: {chosen}",
            "AUTHORITATIVE EXCERPT (patch must apply to this exact text):",
            excerpt,
            "",
            "RULE: Only change what is required inside this excerpt unless the goal explicitly requires changes elsewhere.",
        ]
    else:
        parts += [
            "AUTHORITATIVE CURRENT CONTENT (patch must apply to this exact text):",
            target_text,
        ]

    builder_context_path.write_text("\n".join(parts) + "\n", encoding="utf-8")

    # ---------------------------------------------------
    # Preflight: repository must be clean before running
    # ---------------------------------------------------
    rc, out, err = sh(["git", "status", "--porcelain"], cwd=repo_root)
    if rc != 0:
        raise RuntimeError(f"git status failed\nSTDOUT:\n{out}\nSTDERR:\n{err}")

    if out.strip():
        raise RuntimeError(
            "Working tree is dirty. Please commit/stash or run:\n"
            "  git restore --staged .\n"
            "  git restore .\n"
            f"Dirty files:\n{out}"
        )

    # Remember current branch early
    rc, cur_branch, _ = sh(["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=repo_root)
    current_branch = (cur_branch or "").strip() if rc == 0 else "main"

    # ------------------------------
    # Create context pack
    # ------------------------------
    context_path = run_dir / "context.md"
    write_status(run_dir, state="context_pack")
    log_event(run_dir, "context_pack: start")
    rc, out, err = sh(["python", "context_pack.py", "--out", str(context_path)], cwd=agents_dir)
    if rc != 0:
        write_status(
            run_dir,
            state="failed",
            error=f"context_pack failed\nSTDOUT:\n{out}\nSTDERR:\n{err}",
            finished_ts=time.time(),
        )
        raise RuntimeError(f"context_pack failed\nSTDOUT:\n{out}\nSTDERR:\n{err}")
    log_event(run_dir, "context_pack: done")

    # Paths
    research_path = run_dir / "research.md"
    plan_path = run_dir / "plan.md"
    patch_path = run_dir / "changes.patch"
    test_path = run_dir / "test_plan.md"

    try:
        # ------------------------------
        # Research
        # ------------------------------
        write_status(run_dir, state="researcher")
        log_event(run_dir, "researcher: start")
        run_worker("researcher", args.goal, context_path, research_path, model=args.model_fast)
        log_event(run_dir, "researcher: done")

        # ------------------------------
        # Plan
        # ------------------------------
        write_status(run_dir, state="planner")
        log_event(run_dir, "planner: start")
        run_worker("planner", args.goal, context_path, plan_path, model=args.model_fast)
        log_event(run_dir, "planner: done")

        # ------------------------------
        # Builder (PATCH MODE) with retry
        # ------------------------------
        max_build_attempts = 2
        last_apply_err = ""

        for attempt in range(1, max_build_attempts + 1):
            write_status(run_dir, state="builder")
            log_event(run_dir, f"builder: start (attempt {attempt}/{max_build_attempts})")

            run_worker("builder", args.goal, builder_context_path, patch_path, model=args.model_build)
            log_event(run_dir, "builder: done")

            if not patch_path.exists() or patch_path.stat().st_size == 0:
                raise RuntimeError("builder did not produce patch file")

            # Patch hygiene BEFORE safety checks
            if strip_markdown_fences_inplace(patch_path):
                log_event(run_dir, "builder: stripped markdown fences")

            if normalise_patch_inplace(patch_path):
                log_event(run_dir, "builder: normalised patch")

            if repair_hunk_prefixes_inplace(patch_path):
                log_event(run_dir, "builder: repaired hunk prefixes")

            # No-op handling
            if is_noop_patch(patch_path):
                log_event(run_dir, "no_changes: patch has no hunks; treating as no-op")
                write_status(run_dir, state="no_changes", finished_ts=time.time())
                print(f"\nNO CHANGES ✅ (patch contains no hunks)\nRun folder: {run_dir}\n")
                return

            # Safety checks
            ensure_patch_is_safe(patch_path)
            ensure_patch_matches_goal(args.goal, patch_path)

            rc, out_txt, err_txt = sh(["python", "validate_patch.py", str(patch_path)], cwd=agents_dir)
            if rc != 0:
                last_apply_err = f"validate_patch failed\nSTDOUT:\n{out_txt}\nSTDERR:\n{err_txt}"
                log_event(run_dir, f"builder: patch invalid (attempt {attempt})")

                if attempt == max_build_attempts:
                    raise RuntimeError(last_apply_err)
                continue

            rc2, out2, err2 = sh(["git", "apply", "--check", str(patch_path.resolve())], cwd=repo_root)
            if rc2 == 0:
                log_event(run_dir, "builder: patch applies cleanly")
                break

            last_apply_err = (err2 or out2 or "").strip()
            log_event(run_dir, f"builder: apply-check failed (attempt {attempt}): {last_apply_err}")

            if attempt < max_build_attempts:
                # For command-specific goals, keep using the original anchor excerpt.
                # Do NOT tighten around the failing line, because the failing line may come
                # from a hallucinated/old version of the file.
                if "/reset" in goal_l or "/status" in goal_l:
                    log_event(run_dir, "builder_context: keeping original anchor excerpt for retry")
                else:
                    fail_line = parse_apply_failure_line(last_apply_err)
                    if fail_line:
                        tighter = extract_line_window(target_text, fail_line, radius=120)
                        if tighter:
                            builder_context_path.write_text(
                                "\n".join(
                                    [
                                        f"GOAL:\n{args.goal}\n",
                                        "TARGET FILE: scripts/evcore_loop.py",
                                        "NOTE: Previous patch failed to apply. Error:",
                                        last_apply_err,
                                        "",
                                        "AUTHORITATIVE EXCERPT (patch must apply to this exact text):",
                                        tighter,
                                        "",
                                        "RULE: Only change what is required inside this excerpt.",
                                    ]
                                ) + "\n",
                                encoding="utf-8",
                            )
                            log_event(run_dir, f"builder_context: tightened around failing line {fail_line}")
                continue

            raise RuntimeError(f"git apply --check failed\nSTDOUT:\n{out2}\nSTDERR:\n{err2}")

        # ------------------------------
        # Patch validation (before touching git)
        # ------------------------------
        write_status(run_dir, state="patch_validate")
        log_event(run_dir, "patch_validate: start")

        ensure_patch_is_safe(patch_path)
        ensure_patch_matches_goal(args.goal, patch_path)

        rc, out_txt, err_txt = sh(["python", "validate_patch.py", str(patch_path)], cwd=agents_dir)
        if rc != 0:
            raise RuntimeError(f"validate_patch failed\nSTDOUT:\n{out_txt}\nSTDERR:\n{err_txt}")

        rc2, out2, err2 = sh(["git", "apply", "--check", str(patch_path.resolve())], cwd=repo_root)
        if rc2 != 0:
            raise RuntimeError(f"git apply --check failed\nSTDOUT:\n{out2}\nSTDERR:\n{err2}")

        log_event(run_dir, "patch_validate: done")

        # ------------------------------
        # Tester (does NOT modify code)
        # ------------------------------
        write_status(run_dir, state="tester")
        log_event(run_dir, "tester: start")
        run_worker("tester", args.goal, context_path, test_path, model=args.model_fast)
        log_event(run_dir, "tester: done")

        # ------------------------------
        # Create agent branch + apply patch there
        # ------------------------------
        write_status(run_dir, state="branch_apply")
        log_event(run_dir, "branch_apply: start")
        branch = f"agent/{ts}_{args.tag}"

        rc, _, err = sh(["git", "checkout", "-b", branch], cwd=repo_root)
        if rc != 0:
            raise RuntimeError(f"could not create branch:\n{err}")

        rc, outa, erra = sh(["git", "apply", str(patch_path.resolve())], cwd=repo_root)
        if rc != 0:
            raise RuntimeError(f"git apply failed\nSTDOUT:\n{outa}\nSTDERR:\n{erra}")

        # Compile gate on modified file
        rc, outc, errc = sh(["python", "-m", "py_compile", "scripts/evcore_loop.py"], cwd=repo_root)
        if rc != 0:
            raise RuntimeError(f"py_compile failed\nSTDOUT:\n{outc}\nSTDERR:\n{errc}")

        # --------------------------------------------------
        # NEW: Commit the patch so preview branches contain
        # a real commit (so git diff main..branch works)
        # --------------------------------------------------
        sh(["git", "add", "scripts/evcore_loop.py"], cwd=repo_root)

        rcq, _, _ = sh(["git", "diff", "--cached", "--quiet"], cwd=repo_root)
        if rcq == 0:
            log_event(run_dir, "no_changes: nothing staged after apply")
            write_status(run_dir, state="no_changes", finished_ts=time.time())
            print(f"\nNO CHANGES ✅ (nothing staged)\nRun folder: {run_dir}\n")
            return

        sh(["git", "commit", "-m", f"Agent preview: {args.tag}"], cwd=repo_root)

        log_event(run_dir, "branch_apply: done")

        # Compile gate on modified file
        rc, outc, errc = sh(["python", "-m", "py_compile", "scripts/evcore_loop.py"], cwd=repo_root)
        if rc != 0:
            raise RuntimeError(f"py_compile failed\nSTDOUT:\n{outc}\nSTDERR:\n{errc}")

        log_event(run_dir, "branch_apply: done")

        # ------------------------------
        # PREVIEW: show diff vs current branch
        # ------------------------------
        write_status(run_dir, state="review")
        log_event(run_dir, "review: git diff")
        rc, diff_out, diff_err = sh(
            ["git", "--no-pager", "diff", f"{current_branch}..{branch}"],
            cwd=repo_root,
        )
        print(diff_out if diff_out else diff_err)

        if not args.apply:
            write_status(run_dir, state="preview_only", finished_ts=time.time())
            log_event(run_dir, "preview_only: leaving agent branch for manual review")
            sh(["git", "checkout", current_branch], cwd=repo_root)
            print(f"\nPREVIEW ONLY ✅\nLeft branch for review: {branch}\nReturned to: {current_branch}\nRun folder: {run_dir}\n")
            return

        # ------------------------------
        # APPLY (interactive approval + merge)
        # ------------------------------
        print(f"\nPatch is on branch: {branch}")
        print(f"Current branch: {current_branch}")
        ans = input("Approve merge into current branch? Type YES to merge: ").strip()

        if ans != "YES":
            write_status(run_dir, state="awaiting_approval", finished_ts=time.time())
            log_event(run_dir, "approval denied: leaving agent branch for manual review")
            print(f"\nNot merged. Branch left in place for review: {branch}")
            print(f"To inspect: git checkout {branch} && git --no-pager diff {current_branch}..{branch}")
            return

        write_status(run_dir, state="merge")
        log_event(run_dir, "merge: committing + merging")

        sh(["git", "add", "scripts/evcore_loop.py"], cwd=repo_root)
        sh(["git", "commit", "-m", f"Agent: {args.tag}"], cwd=repo_root)

        sh(["git", "checkout", current_branch], cwd=repo_root)
        rc, outm, errm = sh(["git", "merge", "--ff-only", branch], cwd=repo_root)
        if rc != 0:
            raise RuntimeError(f"merge failed\nSTDOUT:\n{outm}\nSTDERR:\n{errm}")

        sh(["git", "branch", "-D", branch], cwd=repo_root)

        write_status(run_dir, state="done", finished_ts=time.time())
        log_event(run_dir, "Run complete")
        print(f"\nMerged into {current_branch}. Done.")
        print(str(run_dir))

    except Exception as e:
        # Roll back safely if we created/checked out an agent branch
        try:
            # If we're not on current_branch, try to get back
            rc, here, _ = sh(["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd=repo_root)
            here_branch = (here or "").strip()

            if here_branch and here_branch != current_branch:
                sh(["git", "reset", "--hard"], cwd=repo_root)
                sh(["git", "checkout", current_branch], cwd=repo_root)

                # If it looks like an agent branch, attempt to delete it
                if here_branch.startswith("agent/"):
                    sh(["git", "branch", "-D", here_branch], cwd=repo_root)
        except Exception:
            pass

        write_status(run_dir, state="failed", error=str(e), finished_ts=time.time())
        log_event(run_dir, f"FAILED: {e}")
        raise


if __name__ == "__main__":
    main()