search-with-proof: cluster-unavailable detection, kept messages, run-level abort

Zero running workers behind "no free workers available" now raises
ClusterUnavailable instead of waiting out the capacity deadline; CapacityDeadline
carries its message; cleanup failures no longer mask the original exception; the
pilot driver aborts the run on infrastructure failure before any model call.

Co-Authored-By: Claude Fable 5.1 <[email protected]>

tests/test_ax_batch.py

@@ -246,6 +246,8 @@ def test_capacity_backpressure_retries_same_manifest_until_deadline(tmp_path, mo
     busy = [True]
 
     def run(argv, **kwargs):
+        if argv[0] == "kubectl":
+            return subprocess.CompletedProcess(argv, 0, json.dumps({"items": [{"status": {"phase": "Running"}}]}), "")
         if argv[1] == "apply":
             applies.append(kwargs["input"])
             return subprocess.CompletedProcess(argv, 1 if busy[0] else 0, "", "no free workers available secret-url")
@@ -326,3 +328,60 @@ def test_shared_capacity_across_workloads(tmp_path, monkeypatch):
         results = list(pool.map(lambda pair: pair[0](item(), pair[1], 0), zip(lanes, paths, strict=True)))
     assert all(r["status"] == "ok" for r in results)
     assert peak[0] == 2
+
+
+@pytest.mark.parametrize("phases", [[], ["Pending", "Pending"], ["Running", "Running"]])
+def test_capacity_health_uses_lane_runner_and_throttles(monkeypatch, phases):
+    import subprocess
+    from types import SimpleNamespace
+
+    from benchmarks.ax import batch
+
+    now, checks, applies = [0.0], [], []
+    monkeypatch.setenv("S3_ENDPOINT", "https://files.example")
+
+    def jitter(low, high):
+        assert (low, high) == ((2, 5) if now[0] < 60 else (15, 30))
+        return 5 if high == 5 else 20
+
+    monkeypatch.setattr(batch.random, "uniform", jitter)
+
+    def runner(argv, **kwargs):
+        if argv[0] == "kubectl":
+            checks.append(now[0])
+            assert argv == ["kubectl", "-n", "ax-workers", "get", "pods", "-o", "json"]
+            assert kwargs["env"]["KUBECONFIG"] == "/tmp/test-kubeconfig"
+            assert kwargs["timeout"] <= 10
+            return subprocess.CompletedProcess(
+                argv, 0, json.dumps({"items": [{"metadata": {"name": "private-name"}, "status": {"phase": phase}} for phase in phases]}), ""
+            )
+        if argv[1] == "apply":
+            applies.append(kwargs["input"])
+            return subprocess.CompletedProcess(argv, 1, "", "no free workers available")
+        return subprocess.CompletedProcess(argv, 0, "", "")
+
+    lane = batch.AXLane(
+        SimpleNamespace(ax_bin="ax", kubeconfig="/tmp/test-kubeconfig", atespace="default"),
+        workload(),
+        store=object(),
+        runner=runner,
+        clock=lambda: now[0],
+        pause=lambda seconds: now.__setitem__(0, now[0] + seconds),
+    )
+    expected = batch.CapacityDeadline if "Running" in phases else batch.ClusterUnavailable
+    with pytest.raises(expected) as caught:
+        batch.apply_when_available(lane, {"metadata": {"name": "same-task"}}, 190)
+    assert checks == ([60, 120, 180] if "Running" in phases else [60])
+    assert len(set(applies)) == 1
+    assert "private-name" not in str(caught.value)
+    assert str(caught.value)
+
+
+def test_executor_slot_deadline_has_reason(monkeypatch):
+    from types import SimpleNamespace
+
+    from benchmarks.ax import batch
+
+    monkeypatch.setattr(batch, "AX_SLOTS", SimpleNamespace(acquire=lambda **kw: False))
+    with pytest.raises(batch.CapacityDeadline, match="no free workers available before deadline"):
+        batch.SearchExecutorTask(object(), b"archive").__enter__()

tests/test_search_executor_smoke.py

@@ -261,3 +261,28 @@ def test_dispatch_rejects_duplicate_items_before_launch(tmp_path):
 
     with pytest.raises(ValueError, match="duplicate dispatch item"):
         batch.dispatch([item, item], batch.Workload("image", frozenset(), 1, None), tmp_path, {"ax": unexpected}, lanes=("ax",))
+
+
+@pytest.mark.parametrize("stage", ["enter", "body", "normal"])
+def test_cleanup_preserves_original_and_logs(tmp_path, monkeypatch, caplog, stage):
+    _, lane, _ = setup(tmp_path, monkeypatch)
+    original = ValueError("original apply failure")
+
+    def ax(*args, **kwargs):
+        if args[0] == "apply" and stage == "enter":
+            raise original
+        if args[0] == "delete":
+            raise ValueError("delete failed https://private.example/?token=secret")
+
+    lane.ax = ax
+    # Retained reservations must not leak into other tests.
+    monkeypatch.setattr(batch, "AX_SLOTS", __import__("threading").BoundedSemaphore(2))
+    with pytest.raises(RuntimeError if stage == "normal" else ValueError) as caught, batch.SearchExecutorTask(lane, b"archive"):
+        if stage == "body":
+            raise original
+    if stage != "normal":
+        assert caught.value is original
+        assert str(caught.value).startswith("original apply failure; executor cleanup failed:")
+    assert "delete failed" in str(caught.value)
+    assert "executor cleanup failed" in caplog.text
+    assert "private.example" not in caplog.text + str(caught.value)

tests/test_search_pilot_run.py

@@ -362,3 +362,99 @@ def test_explore_exception_keeps_sanitized_message(tmp_path, monkeypatch, throug
     assert "app.py" not in failure["reason"]
     assert "https://" not in failure["reason"]
     assert "private-token-1234" not in failure["reason"]
+
+
+@pytest.mark.parametrize(
+    "kind,calls,aborts", [("cluster", 0, True), ("cluster", 1, True), ("infra", 0, True), ("infra", 1, False), ("repo", 0, False)]
+)
+def test_pilot_classifies_run_abort(tmp_path, monkeypatch, kind, calls, aborts):
+    from benchmarks.ax.batch import ClusterUnavailable
+    from benchmarks.search import pilot_run
+
+    pilot = make_pilot(tmp_path)
+    error = {
+        "cluster": ClusterUnavailable("worker pods: total=2, running=0, pending=2"),
+        "infra": RuntimeError("port-forward unavailable"),
+        "repo": pilot_run.RepositoryFailure("checkout failed"),
+    }[kind]
+
+    def dispatch(task):
+        pilot.worker.metrics["model_calls"] = calls
+        raise error
+
+    monkeypatch.setattr(pilot, "dispatch", dispatch)
+    record = pilot.run()
+    assert bool(record.get("aborted_reason")) is aborts
+    assert record["end_reason"] == "instrument_failure"
+    if aborts:
+        assert record["aborted_reason"] == record["failures"][-1]["reason"]
+
+
+@pytest.mark.parametrize("setup_failure", [False, True])
+def test_live_driver_stops_remaining_pairs_on_infrastructure_failure(tmp_path, monkeypatch, setup_failure):
+    from benchmarks.ax.batch import ClusterUnavailable
+    from benchmarks.search import pilot_run
+
+    monkeypatch.chdir(tmp_path)
+    root = Path("benchmarks/search/private")
+    root.mkdir(parents=True)
+    row = dict(
+        repo="https://github.com/private/project",
+        vulnerable="a" * 40,
+        fixed="b" * 40,
+        family="path",
+        language="python",
+        file="app.py",
+        function="main",
+    )
+    manifest = root / "pilot.json"
+    manifest.write_text(json.dumps({"pairs": [row] * 7}))
+    monkeypatch.setattr(pilot_run, "load_dotenv", lambda: None)
+    monkeypatch.setenv(pilot_run.DEEPSEEK_KEY_ENV, "fake-test-key")
+    monkeypatch.setattr(pilot_run, "OpenRouterChatClient", lambda **kw: __import__("types").SimpleNamespace(max_attempts=1))
+
+    def lane(*args):
+        if setup_failure:
+            raise RuntimeError("store unavailable")
+        return object()
+
+    monkeypatch.setattr(pilot_run, "AXLane", lane)
+    attempts = []
+
+    def enter(task):
+        attempts.append(task.name)
+        raise ClusterUnavailable("worker pods: total=2, running=0, pending=2")
+
+    monkeypatch.setattr(pilot_run.SearchExecutorTask, "__enter__", enter)
+    output = tmp_path / "records.jsonl"
+    status = main(
+        [
+            "--side",
+            "vulnerable",
+            "--pairs",
+            "2,5,6",
+            "--ceiling-usd",
+            "1",
+            "--verify-lane",
+            "ax",
+            "--image",
+            "image@sha256:" + "a" * 64,
+            "--out",
+            str(output),
+        ]
+    )
+    records = [json.loads(line) for line in output.read_text().splitlines()]
+    assert status == 1
+    if setup_failure:
+        assert not attempts and len(records) == 1
+        assert records[0]["aborted_reason"] == "RuntimeError: store unavailable"
+        assert records[0]["skipped_pairs"] == [2, 5, 6]
+        assert records[0]["searches"] == 0
+        return
+    assert len(attempts) == 1
+    search, summary = records
+    assert search["pair_index"] == 2 and search["model_calls"] == 0
+    assert summary["searches"] == 1 and summary["skipped_pairs"] == [5, 6]
+    assert summary["end_reason"] == "instrument_failure"
+    assert summary["aborted_reason"] == search["aborted_reason"]
+    assert "running=0" in summary["aborted_reason"]

tests/test_search_task_storage.py

@@ -222,3 +222,21 @@ def test_scratch_failure_preserves_exit_and_stderr(workspace, monkeypatch):
     assert result["exit_code"] is not None
     assert result["stderr"] == "pip build detail"
     assert result["reason"].startswith("ScratchLimit:")
+
+
+def test_capacity_deadline_plain_diagnostic_survives():
+    from benchmarks.ax.batch import CapacityDeadline
+
+    message = "no free workers available before deadline"
+    assert task_storage.diagnostic(message) == message
+    assert task_storage.exception_reason(CapacityDeadline(message)) == "CapacityDeadline: " + message
+
+
+def test_cleanup_diagnostic_retains_both_causes():
+    error = ValueError("original apply failed")
+    error.executor_cleanup_reason = "executor cleanup failed: " + "details " * 100 + "delete failed"
+    error.args = (str(error) + "; " + error.executor_cleanup_reason,)
+    reason = task_storage.exception_reason(error)
+    assert reason.startswith("ValueError: original apply failed;")
+    assert reason.endswith("delete failed")
+    assert len(reason) <= 300