#!/usr/bin/env python3

"""The full task loop, end to end — the product's signature motion.

oak:  mount -> edit -> commit -> push -> desc -> end, then REMOUNT for a

      follow-up task (the warm-cache story) and repeat edit/commit/push/end.

git:  worktree add -> branch edit -> add/commit -> push -> worktree remove,

      then a second worktree for the follow-up. The task description lives in

      the commit message git already paid for; oak pays a separate `desc`

      call. That asymmetry is the measurement, not a harness gap

      (docs/command-semantics.md): this is agent-default-track, whole-task

      evidence, never a step-by-step equivalence claim.

Transport honesty: oak rows always cross the network to a real oak server.

git runs against GitHub (--git-remote github, transport-comparable) or a

local bare repo (--git-remote local, the core-equivalent track). Rows carry

remote_transport/remote_server; latency comparisons across transports are

forbidden — cross-subject evidence at unlike transports is tool calls and

tokens only.

Setup:

  OAK_BENCH_SYNC_REPO=oak/bench-sync-tmp \\

  [email protected]:oakdotspace/bench-task-loop-tmp.git \\

  python3 scripts/task_loop.py --runs 3

Unconfigured remotes emit skip rows (returncode 77), never silent absence.

"""

from __future__ import annotations

import argparse

import platform

import re

import shutil

import subprocess

import tempfile

import time

from datetime import datetime, timezone

from pathlib import Path

from typing import Any

from oakbench import remotes as oakbench_remotes

from oakbench import tokens as oakbench_tokens

from oakbench.environment import ENV_ISOLATION_VERSION, base_env as oakbench_base_env

from oakbench.environment import command_display

from oakbench.execution import run_timed

from oakbench.results import ResultsStore

from oakbench.rows import row_returncode

from oakbench.runlock import measurement_lock

from oakbench.subjects import (

    DEFAULT_OAK_REPO,

    Subject,

    load_subjects,

    source_metadata,

    subject_details,

    subject_versions,

)

ROOT = Path(__file__).resolve().parents[1]

DEFAULT_WORKDIR = Path(tempfile.gettempdir()) / "oak-task-loop"

SKIP_RETURNCODE = 77

VIRTUAL_BRANCH_RE = re.compile(r"virtual branch ([^\s)]+)")

PHASES = ("task", "followup")

def base_env() -> dict[str, str]:

    return oakbench_base_env(

        author_name="Oak Task Loop",

        author_email="[email protected]",

        oak_author="oak-task-loop",

    )

def parse_args() -> argparse.Namespace:

    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("--subjects", help="Comma-separated subject names from config/subjects.toml")

    parser.add_argument("--config", type=Path, default=ROOT / "config" / "subjects.toml")

    parser.add_argument("--git-bin", type=Path)

    parser.add_argument("--oak-installed-bin", type=Path)

    parser.add_argument("--oak-local-bin", type=Path)

    parser.add_argument("--runs", type=int, default=3)

    parser.add_argument(

        "--git-remote",

        choices=("github", "local"),

        default="github",

        help=(

            "github: agent-default track, real push to GIT_BENCH_REMOTE "

            "(transport-comparable to oak). local: core-equivalent track, "

            "local bare remote (no network in git's loop)."

        ),

    )

    parser.add_argument("--workdir", type=Path, default=DEFAULT_WORKDIR)

    parser.add_argument("--results", type=Path, default=ROOT / "results" / "task-loop")

    parser.add_argument("--oak-repo", type=Path, default=DEFAULT_OAK_REPO)

    parser.add_argument("--admitted-output-chars", type=int, default=20_000)

    parser.add_argument("--keep-workdirs", action="store_true")

    return parser.parse_args()

def run_untimed(command: list[str], cwd: Path | None, timeout: float = 300.0) -> subprocess.CompletedProcess[str]:

    return subprocess.run(

        command,

        cwd=cwd,

        env=base_env(),

        text=True,

        stdout=subprocess.PIPE,

        stderr=subprocess.PIPE,

        timeout=timeout,

        check=False,

    )

def step_row(

    subject: Subject,

    scenario: str,

    operation: str,

    command: list[str],

    cwd: Path,

    run_index: int,

    metadata: dict[str, Any],

    admitted_output_chars: int,

    *,

    kind: str = "vcs",

    expected_returncodes: tuple[int, ...] = (0,),

) -> dict[str, Any]:

    capture = run_timed(command, cwd, base_env(), admitted_output_chars)

    expected = capture.returncode in expected_returncodes

    is_vcs = 1 if kind == "vcs" else 0

    row = {

        **metadata,

        "subject": subject.name,

        "subject_kind": subject.kind,

        "subject_label": subject.label,

        "scenario": scenario,

        "run": run_index,

        "operation": operation,

        "step_kind": kind,

        "elapsed_ms": round(capture.elapsed_ms, 3),

        "returncode": 0 if expected else (capture.returncode or 1),

        "process_returncode": capture.returncode,

        "expected_returncodes": list(expected_returncodes),

        "command": command,

        "tool_call_count": 1,

        "terminal_tool_call_count": 1,

        "vcs_tool_call_count": is_vcs,

        **oakbench_tokens.interaction_token_fields(

            command_display(command),

            capture.stdout_text,

            capture.stderr_text,

            capture.stdout_bytes,

            capture.stderr_bytes,

            capture.stdout_truncated,

            capture.stderr_truncated,

        ),

    }

    if not expected:

        row["stderr"] = capture.stderr_text[-4000:]

    return row

def skip_row(

    subject: Subject,

    scenario: str,

    run_index: int,

    metadata: dict[str, Any],

    reason: str,

) -> dict[str, Any]:

    return {

        **metadata,

        "subject": subject.name,

        "subject_kind": subject.kind,

        "subject_label": subject.label,

        "scenario": scenario,

        "run": run_index,

        "operation": "task_loop.skipped",

        "step_kind": "summary",

        "elapsed_ms": 0.0,

        "returncode": SKIP_RETURNCODE,

        "command": [],

        "skipped": True,

        "skip_reason": reason,

        "tool_call_count": 0,

        "terminal_tool_call_count": 0,

        "vcs_tool_call_count": 0,

    }

def phase_total_row(

    rows: list[dict[str, Any]],

    subject: Subject,

    scenario: str,

    run_index: int,

    metadata: dict[str, Any],

    operation: str,

) -> dict[str, Any]:

    failures = [row for row in rows if row_returncode(row) != 0]

    return {

        **metadata,

        "subject": subject.name,

        "subject_kind": subject.kind,

        "subject_label": subject.label,

        "scenario": scenario,

        "run": run_index,

        "operation": operation,

        "step_kind": "summary",

        "elapsed_ms": round(sum(float(row.get("elapsed_ms", 0)) for row in rows), 3),

        "returncode": 1 if failures else 0,

        "command": [row["command"] for row in rows],

        "tool_call_count": sum(oakbench_tokens.int_or_zero(row.get("tool_call_count")) for row in rows),

        "terminal_tool_call_count": sum(

            oakbench_tokens.int_or_zero(row.get("terminal_tool_call_count")) for row in rows

        ),

        "vcs_tool_call_count": sum(oakbench_tokens.int_or_zero(row.get("vcs_tool_call_count")) for row in rows),

        **oakbench_tokens.summed_token_fields(

            rows, "sum_of_task_loop_steps_command_plus_admitted_output_chars_div_4"

        ),

        "summarized_operations": [row["operation"] for row in rows],

    }

def has_lane_failures(rows: list[dict[str, Any]]) -> bool:

    return any(row_returncode(row) not in (0, SKIP_RETURNCODE) for row in rows)

def enforce_oracle_result(total: dict[str, Any], oracle_ok: bool | None) -> dict[str, Any]:

    if oracle_ok is False:

        total["returncode"] = 1

        total["stderr"] = "oracle_payload_missing_from_remote_branch"

    return total

def edit_step(

    subject: Subject,

    scenario: str,

    operation: str,

    workspace: Path,

    payload: str,

    run_index: int,

    metadata: dict[str, Any],

    admitted: int,

) -> dict[str, Any]:

    command = ["/bin/zsh", "-c", f"echo {payload!r} >> task-notes.md"]

    return step_row(

        subject, scenario, operation, command, workspace, run_index, metadata, admitted, kind="edit"

    )

def oak_virtual_branch(subject: Subject, mount_dir: Path) -> str | None:

    proc = run_untimed([str(subject.bin), "status"], mount_dir)

    match = VIRTUAL_BRANCH_RE.search(proc.stdout + proc.stderr)

    return match.group(1) if match else None

def run_oak_loop(

    subject: Subject,

    remote: oakbench_remotes.RemoteResolution,

    run_root: Path,

    run_index: int,

    metadata: dict[str, Any],

    admitted: int,

) -> list[dict[str, Any]]:

    scenario = "task_loop"

    meta = {**metadata, **remote.row_fields()}

    rows: list[dict[str, Any]] = []

    branches: dict[str, str | None] = {}

    payloads: dict[str, str] = {}

    mounts: list[Path] = []

    vcs = str(subject.bin)

    try:

        for phase in PHASES:

            workspace = run_root / f"tl-{subject.name}-r{run_index}-{phase}"

            payload = f"payload {phase} run={run_index} bench={metadata['bench_id']}"

            payloads[phase] = payload

            acquire = step_row(

                subject, scenario, f"{phase}.acquire", [vcs, "mount", remote.repo or "", str(workspace)],

                run_root, run_index, meta, admitted,

            )

            rows.append(acquire)

            if int(acquire["returncode"]) != 0:

                return rows

            mounts.append(workspace)

            branches[phase] = oak_virtual_branch(subject, workspace)

            rows.append(edit_step(subject, scenario, f"{phase}.edit", workspace, payload, run_index, meta, admitted))

            rows.append(step_row(subject, scenario, f"{phase}.snapshot", [vcs, "commit", "--no-verify"], workspace, run_index, meta, admitted))

            rows.append(step_row(subject, scenario, f"{phase}.publish", [vcs, "push"], workspace, run_index, meta, admitted))

            if phase == "task":

                rows.append(

                    step_row(

                        subject, scenario, "task.describe",

                        [vcs, "desc", f"bench task loop {metadata['bench_id']} r{run_index}"],

                        workspace, run_index, meta, admitted,

                    )

                )

            end = step_row(

                subject, scenario, f"{phase}.end", [vcs, "mount", "end", str(workspace)],

                run_root, run_index, meta, admitted,

            )

            rows.append(end)

            if int(end["returncode"]) == 0 and workspace in mounts:

                mounts.remove(workspace)

        # Untimed oracle: each phase's payload must be on its pushed branch.

        verifier = run_root / f"tl-verify-{subject.name}-r{run_index}"

        clone = run_untimed([vcs, "clone", remote.repo or "", str(verifier)], run_root)

        oracle_ok: bool | None = None

        if clone.returncode == 0:

            oracle_ok = True

            for phase in PHASES:

                branch = branches.get(phase)

                if not branch:

                    oracle_ok = None

                    break

                switch = run_untimed([vcs, "switch", branch], verifier)

                content = (verifier / "task-notes.md")

                seen = content.read_text() if content.exists() else ""

                if switch.returncode != 0 or payloads[phase] not in seen:

                    oracle_ok = False

        total = {

            **phase_total_row(rows, subject, scenario, run_index, meta, "task_loop.total"),

            "oracle_payload_on_remote_branch": oracle_ok,

            "workspace_branches": branches,

            "measurement_source": "verifier_clone_switch_grep; null oracle means branch name unparsed",

        }

        rows.append(enforce_oracle_result(total, oracle_ok))

        return rows

    finally:

        for mount_dir in mounts:

            run_untimed([vcs, "mount", "end", str(mount_dir), "-f"], run_root)

def ensure_git_remote_seeded(subject: Subject, remote: oakbench_remotes.RemoteResolution, scratch: Path) -> str | None:

    """Untimed: make sure the git remote has a main with a seed commit."""

    vcs = str(subject.bin)

    probe = run_untimed([vcs, "ls-remote", str(remote.repo), "refs/heads/main"], None)

    if probe.returncode == 0 and probe.stdout.strip():

        return None

    seed = scratch / "git-seed"

    if seed.exists():

        shutil.rmtree(seed)

    seed.mkdir(parents=True)

    (seed / "README.md").write_text("Task-loop benchmark repo (disposable).\n")

    (seed / "task-notes.md").write_text("# task notes\n")

    for command in (

        [vcs, "init", "-q", "-b", "main"],

        [vcs, "add", "."],

        [vcs, "commit", "-q", "-m", "seed"],

        [vcs, "push", "-q", str(remote.repo), "main:main"],

    ):

        proc = run_untimed(command, seed)

        if proc.returncode != 0:

            return f"git remote seed `{' '.join(command[1:3])}` failed: {proc.stderr.strip()[:200]}"

    return None

def run_git_loop(

    subject: Subject,

    remote: oakbench_remotes.RemoteResolution,

    run_root: Path,

    run_index: int,

    metadata: dict[str, Any],

    admitted: int,

) -> list[dict[str, Any]]:

    scenario = "task_loop"

    meta = {**metadata, **remote.row_fields()}

    vcs = str(subject.bin)

    rows: list[dict[str, Any]] = []

    seed_error = ensure_git_remote_seeded(subject, remote, run_root)

    if seed_error:

        return [skip_row(subject, scenario, run_index, meta, seed_error)]

    # Untimed: the base clone is the long-lived checkout worktrees hang off;

    # an agent fleet pays it once, not per task.

    base = run_root / f"tl-git-base-r{run_index}"

    clone = run_untimed([vcs, "clone", "-q", str(remote.repo), str(base)], run_root)

    if clone.returncode != 0:

        return [skip_row(subject, scenario, run_index, meta, f"base clone failed: {clone.stderr.strip()[:200]}")]

    branches: dict[str, str] = {}

    payloads: dict[str, str] = {}

    for phase in PHASES:

        workspace = run_root / f"tl-{subject.name}-r{run_index}-{phase}"

        branch = oakbench_remotes.disposable_branch(

            str(metadata["bench_id"]), f"taskloop-{phase}", subject.name, run_index

        )

        branches[phase] = branch

        payload = f"payload {phase} run={run_index} bench={metadata['bench_id']}"

        payloads[phase] = payload

        acquire = step_row(

            subject, scenario, f"{phase}.acquire",

            [vcs, "worktree", "add", str(workspace), "-b", branch], base, run_index, meta, admitted,

        )

        rows.append(acquire)

        if int(acquire["returncode"]) != 0:

            return rows

        rows.append(edit_step(subject, scenario, f"{phase}.edit", workspace, payload, run_index, meta, admitted))

        # Agent-default snapshot: git pays stage+commit as two calls.

        rows.append(step_row(subject, scenario, f"{phase}.snapshot.add", [vcs, "add", "."], workspace, run_index, meta, admitted))

        rows.append(

            step_row(

                subject, scenario, f"{phase}.snapshot.commit",

                [vcs, "commit", "-q", "-m", f"bench task loop {metadata['bench_id']} r{run_index} {phase}"],

                workspace, run_index, meta, admitted,

            )

        )

        rows.append(

            step_row(

                subject, scenario, f"{phase}.publish",

                [vcs, "push", "-q", "-u", "origin", branch], workspace, run_index, meta, admitted,

            )

        )

        # No task.describe: git's description is the commit message already

        # paid in {phase}.snapshot.commit (asymmetry documented in the lane doc).

        rows.append(

            step_row(

                subject, scenario, f"{phase}.end",

                [vcs, "worktree", "remove", str(workspace)], base, run_index, meta, admitted,

            )

        )

    # Untimed oracle: payloads visible on the remote branches.

    oracle_ok: bool | None = True

    for phase in PHASES:

        fetch = run_untimed([vcs, "fetch", "-q", "origin", branches[phase]], base)

        show = run_untimed([vcs, "show", f"FETCH_HEAD:task-notes.md"], base)

        if fetch.returncode != 0 or payloads[phase] not in show.stdout:

            oracle_ok = False

    rows.append(

        {

            **phase_total_row(rows, subject, scenario, run_index, meta, "task_loop.total"),

            "oracle_payload_on_remote_branch": oracle_ok,

            "workspace_branches": branches,

            "measurement_source": "fetch_plus_show_grep",

        }

    )

    return rows

def loop_summary_fields(rows: list[dict[str, Any]]) -> dict[str, Any]:

    """followup-vs-task warm deltas computed from this run's step rows."""

    def op_ms(operation: str) -> float | None:

        for row in rows:

            if row.get("operation") == operation:

                return float(row.get("elapsed_ms", 0.0))

        return None

    def phase_ms(prefix: str) -> float:

        return sum(

            float(row.get("elapsed_ms", 0.0))

            for row in rows

            if str(row.get("operation", "")).startswith(prefix + ".")

        )

    task_acquire = op_ms("task.acquire")

    followup_acquire = op_ms("followup.acquire")

    fields: dict[str, Any] = {

        "task_phase_ms": round(phase_ms("task"), 3),

        "followup_phase_ms": round(phase_ms("followup"), 3),

    }

    if task_acquire is not None and followup_acquire is not None:

        fields["followup_warm_delta_ms"] = round(followup_acquire - task_acquire, 3)

        fields["acquire_ms_task"] = round(task_acquire, 3)

        fields["acquire_ms_followup"] = round(followup_acquire, 3)

    else:

        fields["followup_warm_delta_ms"] = None

    return fields

def summary_text(rows: list[dict[str, Any]]) -> str:

    totals = [row for row in rows if row.get("operation") == "task_loop.total"]

    lines = [

        "# Task Loop Summary",

        "",

        "Whole-task cost of the signature motion: acquire workspace -> edit -> snapshot ->",

        "publish -> describe -> end, then a follow-up task in a fresh workspace.",

        "Agent-default track: command sequences are NOT step-equivalent across subjects",

        "(oak pays desc; git's description is its commit message). Latency across unlike",

        "transports is never comparable; tool calls and tokens are.",

        "",

        "| Subject | Run | Transport | Wall ms | Tool calls | Cost-wt tokens (env) | Warm acquire delta ms | Oracle |",

        "| --- | ---: | --- | ---: | ---: | ---: | ---: | --- |",

    ]

    for row in totals:

        oracle = row.get("oracle_payload_on_remote_branch")

        lines.append(

            "| `{}` | {} | {} | {} | {} | {} | {} | {} |".format(

                row["subject"],

                row["run"],

                row.get("remote_transport", "unmeasured"),

                fmt_ms(row.get("elapsed_ms")),

                row.get("tool_call_count", "unmeasured"),

                row.get("estimated_cost_weighted_tokens_with_envelope", "unmeasured"),

                row.get("followup_warm_delta_ms", "unmeasured"),

                "pass" if oracle else ("FAIL" if oracle is False else "unmeasured"),

            )

        )

    skips = [row for row in rows if row.get("skipped")]

    for row in skips:

        lines.append(f"| `{row['subject']}` | {row['run']} | skipped | | | | | {row['skip_reason'][:70]} |")

    return "\n".join(lines) + "\n"

def fmt_ms(value: Any) -> str:

    try:

        return f"{float(value):.0f}"

    except (TypeError, ValueError):

        return "unmeasured"

def main() -> int:

    args = parse_args()

    subjects = load_subjects(args)

    timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")

    args.results.mkdir(parents=True, exist_ok=True)

    args.workdir.mkdir(parents=True, exist_ok=True)

    run_root = args.workdir / "runs" / timestamp

    run_root.mkdir(parents=True, exist_ok=True)

    metadata = {

        "bench_id": timestamp,

        "profile": "task-loop",

        "benchmark_track": "agent-default",

        "timestamp_utc": timestamp,

        "host": platform.node(),

        "platform": platform.platform(),

        "machine": platform.machine(),

        "python": platform.python_version(),

        "env_isolation_version": ENV_ISOLATION_VERSION,

        "subject_versions": subject_versions(subjects),

        "subject_details": subject_details(subjects),

        "source": source_metadata(args.oak_repo),

        "git_remote_mode": args.git_remote,

    }

    rows: list[dict[str, Any]] = []

    with measurement_lock("task_loop") as lock_info:

        metadata["measurement_lock_wait_ms"] = lock_info.wait_ms

        metadata["measurement_lock"] = "held" if lock_info.enabled else "disabled"

        for run_index in range(args.runs):

            for subject in subjects:

                print(f"[run] task_loop run={run_index} subject={subject.name}", flush=True)

                if subject.kind == "git":

                    if args.git_remote == "github":

                        remote = oakbench_remotes.resolve_git_github_remote()

                    else:

                        remote = oakbench_remotes.make_git_bare_remote(

                            subject.bin, run_root, f"task-loop-remote-r{run_index}.git"

                        )

                    if not remote.resolved:

                        rows.append(skip_row(subject, "task_loop", run_index, metadata, remote.skip_reason or ""))

                        continue

                    subject_rows = run_git_loop(subject, remote, run_root, run_index, metadata, args.admitted_output_chars)

                else:

                    remote = oakbench_remotes.resolve_oak_remote("sync")

                    if not remote.resolved:

                        rows.append(skip_row(subject, "task_loop", run_index, metadata, remote.skip_reason or ""))

                        continue

                    subject_rows = run_oak_loop(subject, remote, run_root, run_index, metadata, args.admitted_output_chars)

                total = next((row for row in subject_rows if row.get("operation") == "task_loop.total"), None)

                if total is not None:

                    total.update(loop_summary_fields(subject_rows))

                rows.extend(subject_rows)

    store = ResultsStore(args.results, lane="task-loop")

    raw_path, summary_path = store.write(timestamp, rows, summary_text(rows))

    if not args.keep_workdirs:

        shutil.rmtree(run_root, ignore_errors=True)

    print(f"[result] {raw_path}")

    print(f"[summary] {summary_path}")

    return 1 if has_lane_failures(rows) else 0

if __name__ == "__main__":

    raise SystemExit(main())