Add candidate-patch rollout hook to reward.py (#2a)

README.md

@@ -419,6 +419,14 @@ For v1, the reward is binary:
 
 `error` is `null` on success. On a plain test failure it carries `"Test command exited with code <n>"`; on an infrastructure failure (missing artifacts, Docker build failure, or timeout) it carries the corresponding message.
 
+To score an agent's candidate fix, pass it as a unified diff:
+
+```bash
+forge reward taskpacks/click-pr-001 --patch candidate.diff
+```
+
+The patch is applied with `git apply` onto a throwaway copy of the packaged `repo/` — the base snapshot is never mutated — then the image is rebuilt from the patched copy and the tests run inside the container. If the candidate patch does not apply, the reward is `0.0` with an explanatory `error`. Without `--patch`, the base snapshot is scored as-is, which is `0.0` by construction for a real bug-fix task. Set `FORGE_DOCKER_BIN` to use a different container engine (for example `podman`).
+
 ### `forge report <task_id>`
 
 Prints a compact quality report:

forge/cli.py

@@ -166,10 +166,20 @@ def package_command(task_id: str = typer.Argument(..., help="Task id to package.
 
 
 @app.command("reward")
-def reward_command(taskpack_path: Path = typer.Argument(..., exists=True, file_okay=False, readable=True, help="Path to a taskpack directory.")) -> None:
+def reward_command(
+    taskpack_path: Path = typer.Argument(..., exists=True, file_okay=False, readable=True, help="Path to a taskpack directory."),
+    patch: Path | None = typer.Option(
+        None,
+        "--patch",
+        exists=True,
+        dir_okay=False,
+        readable=True,
+        help="Candidate fix (unified diff) to apply onto a copy of repo/ before scoring.",
+    ),
+) -> None:
     """Run a taskpack reward script and print JSON."""
 
-    result = run_reward_script(taskpack_path)
+    result = run_reward_script(taskpack_path, patch_path=patch)
     print(json.dumps(result.model_dump(), indent=2))
 
 

forge/reward_runner.py

@@ -43,18 +43,22 @@ def _reward_timeout_seconds(taskpack_path: Path) -> int:
     return max(120, configured) * 3 + 120
 
 
-def run_reward_script(taskpack_path: Path, *, timeout_seconds: float | None = None) -> RewardResult:
+def run_reward_script(taskpack_path: Path, *, patch_path: Path | None = None, timeout_seconds: float | None = None) -> RewardResult:
     """Run a taskpack's standalone reward.py and parse its JSON output."""
 
     taskpack_path = taskpack_path.resolve()
     reward_script = taskpack_path / "reward.py"
     if not reward_script.exists():
         return RewardResult(score=0.0, tests_passed=False, error=f"Missing reward script: {reward_script}")
 
+    argv = [sys.executable, str(reward_script)]
+    if patch_path is not None:
+        argv += ["--patch", str(Path(patch_path).resolve())]
+
     timeout = timeout_seconds if timeout_seconds is not None else _reward_timeout_seconds(taskpack_path)
     try:
         completed = subprocess.run(
-            [sys.executable, str(reward_script)],
+            argv,
             cwd=taskpack_path,
             text=True,
             encoding="utf-8",

forge/task_builder.py

@@ -270,12 +270,26 @@ def _prompt(metadata: TaskMetadata) -> str:
 
 def _reward_script() -> str:
     return r'''#!/usr/bin/env python3
+"""Standalone, self-contained reward for this task.
+
+Usage:
+  python reward.py                     score the base repo snapshot as-is
+  python reward.py --patch fix.diff    apply a candidate fix, then score
+
+The candidate patch is applied (git apply) onto a throwaway copy of repo/, so the
+base snapshot is never mutated and reruns stay independent. Emits one JSON line:
+{"score": 0.0|1.0, "tests_passed": bool, "error": str|null}. Set FORGE_DOCKER_BIN
+to use a different container engine (e.g. podman).
+"""
 from __future__ import annotations
 
+import argparse
 import json
+import os
 import re
 import shutil
 import subprocess
+import tempfile
 import time
 import uuid
 from pathlib import Path
@@ -291,7 +305,36 @@ def slug(value: str) -> str:
     return normalized or "task"
 
 
+def stage_repo_with_patch(repo_dir: Path, patch_file: Path) -> tuple[Path | None, str | None]:
+    """Copy repo_dir into a temp dir and apply the candidate patch. Returns (work_dir, error)."""
+
+    if not patch_file.exists():
+        return None, f"Candidate patch not found: {patch_file}"
+    if shutil.which("git") is None:
+        return None, "git executable not found on PATH (required to apply --patch)"
+    stage_root = Path(tempfile.mkdtemp(prefix="forge-reward-stage-"))
+    work = stage_root / "repo"
+    shutil.copytree(repo_dir, work)
+    applied = subprocess.run(
+        ["git", "apply", "--whitespace=nowarn", str(patch_file)],
+        cwd=work,
+        text=True,
+        encoding="utf-8",
+        errors="replace",
+        capture_output=True,
+    )
+    if applied.returncode != 0:
+        shutil.rmtree(stage_root, ignore_errors=True)
+        detail = (applied.stderr or applied.stdout).strip()[-500:]
+        return None, f"Candidate patch did not apply: {detail}"
+    return work, None
+
+
 def main() -> int:
+    parser = argparse.ArgumentParser(description="Score a candidate fix against this task's tests.")
+    parser.add_argument("--patch", type=Path, default=None, help="Unified diff with the candidate fix to apply before testing.")
+    args = parser.parse_args()
+
     package_dir = Path(__file__).resolve().parent
     task_path = package_dir / "task.json"
     repo_dir = package_dir / "repo"
@@ -303,55 +346,70 @@ def main() -> int:
         return emit(0.0, False, f"Missing repository directory: {repo_dir}")
     if not dockerfile.exists():
         return emit(0.0, False, f"Missing Dockerfile: {dockerfile}")
-    if shutil.which("docker") is None:
-        return emit(0.0, False, "Docker executable not found on PATH")
 
-    task = json.loads(task_path.read_text(encoding="utf-8"))
-    test_command = task["test_command"]
-    timeout_seconds = int(task.get("timeout_seconds", 300))
-    image_tag = f"forge-reward-{slug(task['id'])}-{uuid.uuid4().hex[:12]}"
-    container_name = f"{image_tag}-run"
+    build_context = repo_dir
+    stage_root: Path | None = None
+    if args.patch is not None:
+        staged, error = stage_repo_with_patch(repo_dir, args.patch.resolve())
+        if error is not None:
+            return emit(0.0, False, error)
+        build_context = staged
+        stage_root = staged.parent
 
-    started_at = time.monotonic()
     try:
-        build = subprocess.run(
-            ["docker", "build", "--file", str(dockerfile), "--tag", image_tag, str(repo_dir)],
-            text=True,
-            encoding="utf-8",
-            errors="replace",
-            capture_output=True,
-            timeout=max(120, timeout_seconds * 2),
-        )
-    except subprocess.TimeoutExpired:
-        return emit(0.0, False, "Docker build timed out")
-
-    if build.returncode != 0:
-        detail = (build.stderr or build.stdout).strip()[-1000:]
-        return emit(0.0, False, f"Docker build failed with code {build.returncode}: {detail}")
+        docker_bin = os.environ.get("FORGE_DOCKER_BIN", "docker")
+        if shutil.which(docker_bin) is None:
+            return emit(0.0, False, f"Docker executable not found on PATH: {docker_bin}")
 
-    try:
-        run = subprocess.run(
-            [
-                "docker", "run", "--rm", "--name", container_name,
-                "--network", "none", "--memory", "4g", "--cpus", "2",
-                "--pids-limit", "512", "--security-opt", "no-new-privileges",
-                image_tag, "sh", "-lc", test_command,
-            ],
-            text=True,
-            encoding="utf-8",
-            errors="replace",
-            capture_output=True,
-            timeout=timeout_seconds,
-        )
-        if run.returncode == 0:
-            return emit(1.0, True, None)
-        return emit(0.0, False, f"Test command exited with code {run.returncode}")
-    except subprocess.TimeoutExpired:
-        subprocess.run(["docker", "rm", "-f", container_name], text=True, encoding="utf-8", errors="replace", capture_output=True)
-        elapsed = round(time.monotonic() - started_at, 2)
-        return emit(0.0, False, f"Test command timed out after {timeout_seconds} seconds ({elapsed}s elapsed)")
+        task = json.loads(task_path.read_text(encoding="utf-8"))
+        test_command = task["test_command"]
+        timeout_seconds = int(task.get("timeout_seconds", 300))
+        image_tag = f"forge-reward-{slug(task['id'])}-{uuid.uuid4().hex[:12]}"
+        container_name = f"{image_tag}-run"
+
+        started_at = time.monotonic()
+        try:
+            build = subprocess.run(
+                [docker_bin, "build", "--file", str(dockerfile), "--tag", image_tag, str(build_context)],
+                text=True,
+                encoding="utf-8",
+                errors="replace",
+                capture_output=True,
+                timeout=max(120, timeout_seconds * 2),
+            )
+        except subprocess.TimeoutExpired:
+            return emit(0.0, False, "Docker build timed out")
+
+        if build.returncode != 0:
+            detail = (build.stderr or build.stdout).strip()[-1000:]
+            return emit(0.0, False, f"Docker build failed with code {build.returncode}: {detail}")
+
+        try:
+            run = subprocess.run(
+                [
+                    docker_bin, "run", "--rm", "--name", container_name,
+                    "--network", "none", "--memory", "4g", "--cpus", "2",
+                    "--pids-limit", "512", "--security-opt", "no-new-privileges",
+                    image_tag, "sh", "-lc", test_command,
+                ],
+                text=True,
+                encoding="utf-8",
+                errors="replace",
+                capture_output=True,
+                timeout=timeout_seconds,
+            )
+            if run.returncode == 0:
+                return emit(1.0, True, None)
+            return emit(0.0, False, f"Test command exited with code {run.returncode}")
+        except subprocess.TimeoutExpired:
+            subprocess.run([docker_bin, "rm", "-f", container_name], text=True, encoding="utf-8", errors="replace", capture_output=True)
+            elapsed = round(time.monotonic() - started_at, 2)
+            return emit(0.0, False, f"Test command timed out after {timeout_seconds} seconds ({elapsed}s elapsed)")
+        finally:
+            subprocess.run([docker_bin, "image", "rm", "-f", image_tag], text=True, encoding="utf-8", errors="replace", capture_output=True)
     finally:
-        subprocess.run(["docker", "image", "rm", "-f", image_tag], text=True, encoding="utf-8", errors="replace", capture_output=True)
+        if stage_root is not None:
+            shutil.rmtree(stage_root, ignore_errors=True)
 
 
 if __name__ == "__main__":

tests/test_reward_rollout.py

@@ -0,0 +1,154 @@
+"""Coverage for the candidate-patch rollout hook in the generated reward.py.
+
+The reward script is a self-contained template (no `forge` import), so these
+tests run the *generated* script as a subprocess. Docker is forced absent via a
+bogus FORGE_DOCKER_BIN, which lets us prove patch-staging behavior without a
+container engine: a good patch applies and the flow reaches the docker check,
+while a bad patch is rejected before it.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+
+import pytest
+
+from forge import task_builder
+from forge.reward_runner import RewardResult, run_reward_script
+from forge.task_builder import (
+    _reward_script,
+    gold_patch_path,
+    metadata_path,
+    package_task,
+    verification_path,
+)
+from forge.task_schema import TaskMetadata
+
+GOOD_PATCH = (
+    b"diff --git a/m.py b/m.py\n"
+    b"--- a/m.py\n+++ b/m.py\n@@ -1,2 +1,2 @@\n def f():\n-    return 1\n+    return 2\n"
+)
+BAD_PATCH = (
+    b"diff --git a/m.py b/m.py\n"
+    b"--- a/m.py\n+++ b/m.py\n@@ -1,2 +1,2 @@\n def g():\n-    return 1\n+    return 2\n"
+)
+
+requires_git = pytest.mark.skipif(shutil.which("git") is None, reason="git required to apply candidate patches")
+
+
+def _metadata() -> TaskMetadata:
+    return TaskMetadata(
+        id="rollout-001",
+        repo_url="https://github.com/example/project.git",
+        pr_number=3,
+        repo_name="example/project",
+        pr_title="Fix f",
+        pr_body="body",
+        base_commit="a" * 40,
+        head_commit="b" * 40,
+        test_command="pytest",
+        language="python",
+        timeout_seconds=60,
+    )
+
+
+def _build_taskpack(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
+    metadata = _metadata()
+    metadata.write_json(metadata_path(tmp_path, metadata.id))
+    gold_patch_path(tmp_path, metadata.id).write_bytes(b"diff --git a/m.py b/m.py\n")
+    verification_path(tmp_path, metadata.id).write_text("{}", encoding="utf-8")
+
+    repo_dir = tmp_path / ".forge" / "repos" / metadata.id
+    repo_dir.mkdir(parents=True)
+    (repo_dir / "m.py").write_text("def f():\n    return 1\n", encoding="utf-8")
+
+    monkeypatch.setattr(task_builder, "checkout_clean", lambda *a, **k: None)
+    return package_task(metadata.id, root=tmp_path)
+
+
+def _run_reward(package_dir: Path, *args: str) -> dict:
+    env = os.environ.copy()
+    env["FORGE_DOCKER_BIN"] = "forge-no-such-docker-bin"  # force the docker-absent branch
+    completed = subprocess.run(
+        [sys.executable, str(package_dir / "reward.py"), *args],
+        cwd=package_dir,
+        env=env,
+        capture_output=True,
+        text=True,
+    )
+    return json.loads(completed.stdout.strip().splitlines()[-1])
+
+
+def test_generated_reward_script_is_valid_and_supports_patch() -> None:
+    src = _reward_script()
+    compile(src, "reward.py", "exec")  # syntactically valid
+    assert "--patch" in src
+    assert "stage_repo_with_patch" in src
+
+
+@requires_git
+def test_good_candidate_patch_applies_then_reaches_docker(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    package_dir = _build_taskpack(tmp_path, monkeypatch)
+    patch = tmp_path / "good.patch"
+    patch.write_bytes(GOOD_PATCH)
+
+    payload = _run_reward(package_dir, "--patch", str(patch))
+
+    # Patch applied cleanly, so the failure is the (forced) missing docker, not the patch.
+    assert payload["tests_passed"] is False
+    assert "Docker executable not found" in payload["error"]
+    # The base snapshot must remain unmodified.
+    assert (package_dir / "repo" / "m.py").read_text(encoding="utf-8") == "def f():\n    return 1\n"
+
+
+@requires_git
+def test_bad_candidate_patch_is_rejected_before_docker(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    package_dir = _build_taskpack(tmp_path, monkeypatch)
+    patch = tmp_path / "bad.patch"
+    patch.write_bytes(BAD_PATCH)
+
+    payload = _run_reward(package_dir, "--patch", str(patch))
+
+    assert payload["score"] == 0.0
+    assert payload["tests_passed"] is False
+    assert "Candidate patch did not apply" in payload["error"]
+
+
+def test_missing_candidate_patch_is_reported(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    package_dir = _build_taskpack(tmp_path, monkeypatch)
+
+    payload = _run_reward(package_dir, "--patch", str(tmp_path / "nope.patch"))
+
+    assert payload["score"] == 0.0
+    assert "Candidate patch not found" in payload["error"]
+
+
+def test_no_patch_preserves_base_scoring_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    package_dir = _build_taskpack(tmp_path, monkeypatch)
+
+    payload = _run_reward(package_dir)
+
+    # Without --patch the script scores the base repo directly; here it only reaches
+    # the forced-absent docker check, proving the base path is intact.
+    assert payload["tests_passed"] is False
+    assert "Docker executable not found" in payload["error"]
+
+
+def test_run_reward_script_passes_patch_through(tmp_path: Path) -> None:
+    taskpack = tmp_path / "taskpack"
+    taskpack.mkdir()
+    (taskpack / "reward.py").write_text(
+        "import sys, json\nprint(json.dumps({'score': 0.0, 'tests_passed': False, 'error': ' '.join(sys.argv[1:])}))\n",
+        encoding="utf-8",
+    )
+
+    result = run_reward_script(taskpack, patch_path=Path("candidate.patch"))
+
+    assert isinstance(result, RewardResult)
+    assert "--patch" in (result.error or "")
+    assert (result.error or "").endswith("candidate.patch")