fix(pilot): preserve partial results and stop on provider billing errors

experiments/P0-pilot/run_pilot.py

@@ -50,37 +50,64 @@ def judge(client, passage: str, lean: str) -> dict:
             "reason": str(reply.get("reason", ""))[:300]}
 
 
-def run_item(client, item: dict) -> dict:
-    if not item.get("source_text", "").strip():
-        raise ValueError("Missing source_text; use a private source-inclusive export")
-    code = item.get("lean_code") or ""
-    row: dict = {"id": item["id"], "compiled": bool(item.get("verified")), "errors": {}}
-    proj = project(code)
-    row["projection"] = proj.as_dict()
-    row["lint"] = lint(proj)
-    try:
-        src = srcsimplex.extract_source_simplex(client, item.get("source_text", ""),
-                                                item.get("label", ""))
-        al = align(src, proj)
-        row["source"] = [r.__dict__ for r in src.roles]
-        row["role"] = al.as_dict()
-    except CallBudgetExceeded:
-        raise
-    except Exception as error:  # a failed extraction is data, not a crash
-        row["errors"]["role"] = str(error)[:300]
-    try:
-        row["judge"] = judge(client, item.get("source_text", ""), code)
-    except CallBudgetExceeded:
-        raise
-    except Exception as error:
-        row["errors"]["judge"] = str(error)[:300]
-    return row
+def provider_blocked(error: Exception) -> bool:
+    text = str(error).lower()
+    return any(code in text for code in ("arrearage", "overdue-payment", "invalid_api_key", "http 401", "http 403"))
+
+
+def complete_row(row: dict) -> bool:
+    return not row.get("errors") and "role" in row and "judge" in row
+
+
+class ProviderBlocked(RuntimeError):
+    """A billing or authentication failure requires external recovery before resuming."""
+
+    def __init__(self, row):
+        super().__init__("Provider billing/authentication blocked; partial results were preserved")
+        self.row = row
 
 
 class CallBudgetExceeded(RuntimeError):
     """Stop before issuing a request that would exceed the cumulative budget."""
 
 
+def run_item(client, item: dict, previous: dict | None = None) -> dict:
+    if not item.get("source_text", "").strip():
+        raise ValueError("Missing source_text; use a private source-inclusive export")
+    code = item.get("lean_code") or ""
+    row: dict = {**(previous or {}), "id": item["id"], "compiled": bool(item.get("verified"))}
+    row["errors"] = dict(row.get("errors", {}))
+    proj = project(code)
+    row["projection"] = proj.as_dict()
+    row["lint"] = lint(proj)
+    if "role" not in row:
+        try:
+            src = srcsimplex.extract_source_simplex(client, item["source_text"], item.get("label", ""))
+            al = align(src, proj)
+            row["source"] = [r.__dict__ for r in src.roles]
+            row["role"] = al.as_dict()
+            row["errors"].pop("role", None)
+        except CallBudgetExceeded as error:
+            error.row = row
+            raise
+        except Exception as error:
+            row["errors"]["role"] = str(error)[:300]
+            if provider_blocked(error):
+                raise ProviderBlocked(row) from error
+    if "judge" not in row:
+        try:
+            row["judge"] = judge(client, item["source_text"], code)
+            row["errors"].pop("judge", None)
+        except CallBudgetExceeded as error:
+            error.row = row
+            raise
+        except Exception as error:
+            row["errors"]["judge"] = str(error)[:300]
+            if provider_blocked(error):
+                raise ProviderBlocked(row) from error
+    return row
+
+
 def digest(path: Path) -> str:
     return hashlib.sha256(path.read_bytes()).hexdigest()
 
@@ -106,8 +133,9 @@ def execute(args, items, settings) -> int:
     if args.out.exists():
         rows = [json.loads(x) for x in args.out.read_text(encoding="utf-8").splitlines() if x.strip()]
     ids = {item["id"] for item in items}
-    done = {row["id"] for row in rows}
-    if len(done) != len(rows) or not done <= ids:
+    by_id = {row["id"]: row for row in rows}
+    done = {row["id"] for row in rows if complete_row(row)}
+    if len(by_id) != len(rows) or not set(by_id) <= ids:
         raise ValueError("Results have duplicate or unknown item IDs")
     hashes = {
         "source": hashlib.sha256((srcsimplex.SYSTEM_PROMPT + inspect.getsource(srcsimplex.passage_prompt)).encode()).hexdigest(),
@@ -129,11 +157,24 @@ def execute(args, items, settings) -> int:
     elif rows:
         raise ValueError("Cannot resume results without their audit metadata")
     todo = [item for item in items if item["id"] not in done]
+    missing = sum(2 - int("role" in by_id.get(i["id"], {}))
+                  - int("judge" in by_id.get(i["id"], {})) for i in todo)
     print(f"items={len(items)} done={len(done)} todo={len(todo)} model={settings.llm.model} "
-          f"nominal_calls={len(todo) * 2} request_attempts={meta['request_attempts']} cap={args.max_calls}")
+          f"missing_components={missing} request_attempts={meta['request_attempts']} cap={args.max_calls}")
     print("prompt_hashes:", json.dumps(hashes))
     if args.dry_run:
         return 0
+    if getattr(args, "audit_only", False):
+        meta["completed_items"] = len(done)
+        meta["processed_items"] = len(rows)
+        meta["status"] = ("provider_blocked" if any(provider_blocked(RuntimeError(message))
+                          for row in rows for message in row.get("errors", {}).values())
+                          else "detector_errors" if todo else "complete")
+        if args.out.exists():
+            meta["results_sha256"] = digest(args.out)
+        atomic_json(meta_path, meta)
+        print(json.dumps(meta))
+        return 0 if meta["status"] == "complete" else 2
     if settings.llm.model != "qwen-max" or settings.llm.temperature != 0:
         raise ValueError("P0 is preregistered for qwen-max at temperature 0")
     client = AuditedClient.from_settings(settings.llm)
@@ -142,6 +183,7 @@ def execute(args, items, settings) -> int:
     def checkpoint():
         meta["usage"] = client.usage.as_dict()
         meta["completed_items"] = len(done)
+        meta["processed_items"] = len(by_id)
         atomic_json(meta_path, meta)
 
     def reserve():
@@ -155,19 +197,30 @@ def execute(args, items, settings) -> int:
     meta["max_calls"] = args.max_calls
     meta["status"] = "running"
     checkpoint()
+    def save_row(row):
+        by_id[row["id"]] = row
+        if complete_row(row):
+            done.add(row["id"])
+        temporary = args.out.with_suffix(args.out.suffix + ".tmp")
+        with temporary.open("w", encoding="utf-8") as stream:
+            for item in items:
+                if item["id"] in by_id:
+                    stream.write(json.dumps(by_id[item["id"]], ensure_ascii=False) + "\n")
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, args.out)
+        checkpoint()
+
     try:
-        with args.out.open("a", encoding="utf-8") as stream:
-            for item in todo:
-                row = run_item(client, item)
-                stream.write(json.dumps(row, ensure_ascii=False) + "\n")
-                stream.flush()
-                os.fsync(stream.fileno())
-                done.add(item["id"])
-                checkpoint()
-                print(f"completed={len(done)}/{len(items)} requests={meta['request_attempts']}", flush=True)
-        meta["status"] = "complete"
-    except CallBudgetExceeded as error:
-        meta["status"] = "budget_exhausted"
+        for item in todo:
+            row = run_item(client, item, by_id.get(item["id"]))
+            save_row(row)
+            print(f"completed={len(done)}/{len(items)} requests={meta['request_attempts']}", flush=True)
+        meta["status"] = "complete" if len(done) == len(items) else "detector_errors"
+    except (CallBudgetExceeded, ProviderBlocked) as error:
+        if getattr(error, "row", None) is not None:
+            save_row(error.row)
+        meta["status"] = "provider_blocked" if isinstance(error, ProviderBlocked) else "budget_exhausted"
         print(str(error), file=sys.stderr)
     except BaseException:
         meta["status"] = "interrupted"
@@ -186,6 +239,7 @@ def main() -> int:
     ap.add_argument("--out", required=True, type=Path)
     ap.add_argument("--max-calls", type=int, default=400)
     ap.add_argument("--dry-run", action="store_true")
+    ap.add_argument("--audit-only", action="store_true", help="Reconcile saved state without making API calls")
     ap.add_argument("--labels", type=Path, help="Frozen labels; their hash is bound to the run")
     args = ap.parse_args()
     if not 1 <= args.max_calls <= 400: