Bound opt-in native receiver setup recovery by the existing outer budget

benchmark/README.md

@@ -1543,3 +1543,24 @@ no account response, login values or token is emitted. The ordinary startup
 probe and automatic retry path keep their previous behavior. This observation
 is not a reservation of future provider capacity; a later failure is retained
 and requires another explicit repair.
+
+For a **new** native Goal or heartbeat receiver, the startup-only account probe
+keeps its default single 25-second total window. An operator may explicitly pass
+`--native-readiness-total-seconds 75 --native-readiness-attempt-seconds 25` to
+`benchmark.edgebench.run` when bounded setup recovery is authorized. These caps
+are at most 75 seconds total and 25 seconds per attempt, with at most three
+attempts. The total is additionally clipped to the native runner's original
+remaining execution snapshot, charging entry delay, preparation, reads, cleanup
+and backoff; terminal cleanup can finish after the deadline but grants no more
+requests. A missing, invalid or expired snapshot refuses the observer.
+
+This opt-in retries only locally typed `account/read` or
+`account/rateLimits/read` transport timeouts. Its bounded `account_timeout`
+diagnostic preserves that distinction without exposing remote payloads.
+Initialization failure, identity or
+protocol rejection, disconnect, unknown failure, Stop and exhausted budget do
+not authorize another attempt. It performs no inference, token refresh, account
+switch, model retry or Goal mutation, and starts no solver phase until a probe
+succeeds. It is unavailable for Official/single profiles, custom API providers
+and warm checkpoint continuations. Omit both options to retain existing startup
+behavior; neither option changes submission counting, cooldown or solver budget.

benchmark/edgebench/run.py

@@ -135,6 +135,12 @@ def main(argv=None):
     parser.add_argument("--native-goal-continue", action="store_true",
                         help="Opt-in: when a native Goal declares completion with budget left, start "
                              "continuation Goals until the budget edge (native-goal worker only)")
+    parser.add_argument("--native-readiness-total-seconds", type=float,
+                        help="Opt-in fresh native receiver setup budget, at most 75s and clipped to "
+                             "the existing outer allowance; only read-only account timeouts retry")
+    parser.add_argument("--native-readiness-attempt-seconds", type=float,
+                        help="Per-attempt cap for explicit native readiness setup, at most 25s; "
+                             "omitted keeps 25s (not a model retry)")
     # Preserve the existing unique abbreviation after adding --managed-run-json.
     parser.add_argument("--model", "--m", required=True)
     parser.add_argument("--effort", choices=REASONING_EFFORTS, required=True)
@@ -214,6 +220,16 @@ def main(argv=None):
     api_key = os.environ.get("OPENAI_API_KEY") if api_base_url else None
     if api_base_url and not api_key:
         parser.error("OPENAI_BASE_URL requires OPENAI_API_KEY; no login fallback")
+    from benchmark.runtime.codex_readiness import native_readiness_budget
+    try:
+        readiness_budget = native_readiness_budget(
+            args.native_readiness_total_seconds, args.native_readiness_attempt_seconds)
+    except ValueError as error:
+        parser.error(str(error))
+    if readiness_budget is not None and (
+        args.worker in {"official", "single"} or api_base_url or args.restore_checkpoint is not None
+    ):
+        parser.error("Native readiness budget requires a fresh native Goal or heartbeat worker")
     args.timeout = _task_default(args.task, "timeout_seconds", args.timeout, DEFAULT_TIMEOUT_SECONDS)
     args.eval_interval = (0 if args.feedback == "native" and args.eval_interval is None else
                          _task_default(args.task, "eval_interval_seconds", args.eval_interval, 300))
@@ -418,6 +434,9 @@ def main(argv=None):
                             if args.execution_observation is not None else {}),
                          **({"provider_request_admission": True} if args.provider_request_admission else {}),
                          native_goal_continue=args.native_goal_continue,
+                         **({"native_readiness_total_seconds": readiness_budget[0],
+                             "native_readiness_attempt_seconds": readiness_budget[1]}
+                            if readiness_budget is not None else {}),
                          replan_after_turns=args.replan_after_turns,
                          replan_after_todos=args.replan_after_todos,
                          planning_horizon_goal=args.planning_horizon_goal,

benchmark/runtime/codex_readiness.py

@@ -6,11 +6,14 @@
 """
 from __future__ import annotations
 
+import argparse
 import json
+import math
 import os
 import shutil
 import tempfile
 import time
+from dataclasses import dataclass
 from pathlib import Path
 
 from loopx.control_plane.turn_driver.codex_stdio import (
@@ -29,6 +32,26 @@
     "app_server_stream_closed", "app_server_stdin_closed",
     "app_server_write_failed", "app_server_exited_before_response",
 })
+_ACCOUNT_TIMEOUTS = _TRANSPORT_TIMEOUTS - {"app_server_response_timeout:initialize"}
+
+
+@dataclass(frozen=True)
+class _Observation:
+    receipt: dict
+    account_timeout: bool = False
+
+
+def native_readiness_budget(total_seconds=None, attempt_seconds=None) -> tuple[float, float] | None:
+    """Validate an explicit startup-only budget; omission preserves the old path."""
+    if total_seconds is None and attempt_seconds is None:
+        return None
+    attempt_seconds = 25 if attempt_seconds is None else attempt_seconds
+    if any(type(value) not in (int, float) or not math.isfinite(value)
+           for value in (total_seconds, attempt_seconds)):
+        raise ValueError("Native readiness budgets must be finite numbers")
+    if not 0 < total_seconds <= 75 or not 0 < attempt_seconds <= 25:
+        raise ValueError("Native readiness requires total seconds in (0, 75] and attempt seconds in (0, 25]")
+    return float(total_seconds), float(attempt_seconds)
 
 
 class _ObservationExpired(TimeoutError):
@@ -64,7 +87,7 @@ def _remaining(deadline: float) -> float:
     return remaining
 
 
-def _observe_once(env: dict[str, str], deadline: float, *, require_allowance: bool = False) -> dict:
+def _observe_once(env: dict[str, str], deadline: float, *, require_allowance: bool = False) -> _Observation:
     started = time.monotonic()
     stage = "inspection_home"
     allowance = None
@@ -119,30 +142,67 @@ def request(method, params):
                     allowance = {"availability": "available", "observed_at": observed_at}
     except Exception as exc:
         error_kind, retryable = _failure(exc)
-        return {"ok": False, "stage": stage, "error_kind": error_kind,
+        receipt = {"ok": False, "stage": stage, "error_kind": error_kind,
                 "retryable": retryable, "elapsed_seconds": round(time.monotonic() - started, 3)}
-    return {"ok": True, "stage": stage, **({"allowance": allowance} if allowance else {}),
-            "elapsed_seconds": round(time.monotonic() - started, 3)}
+        # Retry permission comes from the local transport exception, never a
+        # remote message or the sanitized category printed in the receipt.
+        return _Observation(receipt, isinstance(exc, NativeGoalProtocolError)
+                            and str(exc) in _ACCOUNT_TIMEOUTS)
+    return _Observation({"ok": True, "stage": stage, **({"allowance": allowance} if allowance else {}),
+                         "elapsed_seconds": round(time.monotonic() - started, 3)})
+
+
+def _outer_remaining(env: dict[str, str]) -> float:
+    # The native runner snapshots the existing execution allowance immediately
+    # before exec. Charge entry/setup delay and convert it only once to a local
+    # monotonic deadline. Missing or invalid provenance cannot grant more time.
+    try:
+        remaining = float(env["SFORGE_AGENT_REMAINING_SECONDS"])
+        sampled_at = float(env["SFORGE_AGENT_SEGMENT_STARTED_AT"])
+    except (KeyError, TypeError, ValueError):
+        raise _InvalidObservation("unavailable_outer_budget") from None
+    age = time.time() - sampled_at
+    if not all(math.isfinite(value) for value in (remaining, sampled_at, age)) or age < 0:
+        raise _InvalidObservation("unavailable_outer_budget")
+    if remaining - age <= 0:
+        raise _ObservationExpired
+    return remaining - age
 
 
-def check_native_receiver(env: dict[str, str]) -> dict:
-    # Retain the existing total observation budget, including home preparation,
-    # initialization, reads and retry backoff. Process cleanup may finish later;
+def check_native_receiver(env: dict[str, str], *, total_seconds=None, attempt_seconds=None) -> dict:
+    # Retain the unconfigured total budget; an explicit startup cap also charges
+    # preparation, initialization, reads and retry backoff. Cleanup may finish later;
     # it never grants another request beyond this deadline or starts a model.
     started = time.monotonic()
     deadline = started + 25
     attempts = []
+    budget = native_readiness_budget(total_seconds, attempt_seconds)
+    if budget is not None:
+        try:
+            # Reading the snapshot is part of setup cost too.
+            deadline = min(started + budget[0], time.monotonic() + _outer_remaining(env))
+        except (_InvalidObservation, _ObservationExpired) as exc:
+            kind, _ = _failure(exc)
+            attempts.append({"ok": False, "stage": "outer_budget", "error_kind": kind,
+                             "retryable": False, "account_timeout": False,
+                             "elapsed_seconds": round(time.monotonic() - started, 3)})
     for index in range(3):
+        if attempts and not attempts[-1]["retryable"]:
+            break
         if index:
             remaining = deadline - time.monotonic()
             if remaining <= 0:
                 break
             time.sleep(min(0.2, remaining))
             if time.monotonic() >= deadline:
                 break
-        observation = _observe_once(env, deadline)
-        attempts.append(observation)
-        if observation["ok"] or not observation["retryable"]:
+        attempt_deadline = min(deadline, time.monotonic() + budget[1]) if budget else deadline
+        observation = _observe_once(env, attempt_deadline)
+        attempts.append({**observation.receipt,
+                         **({"account_timeout": observation.account_timeout} if budget is not None else {})})
+        if observation.receipt["ok"] or not observation.receipt["retryable"]:
+            break
+        if budget is not None and not observation.account_timeout:
             break
     ok = attempts[-1]["ok"]
     result = {"ok": ok, "inference_invoked": False, "attempts": attempts,
@@ -159,17 +219,26 @@ def require_native_receiver_allowance(env: dict[str, str]) -> dict:
     refuses unknown/exhausted allowance and missing optional provider support;
     it never refreshes login, switches accounts, or starts inference.
     """
-    observation = _observe_once(env, time.monotonic() + 25, require_allowance=True)
+    observation = _observe_once(env, time.monotonic() + 25, require_allowance=True).receipt
     if not observation["ok"]:
         raise ValueError("Explicit Host continuation requires fresh native allowance: "
                          + observation["stage"] + "/" + observation["error_kind"])
     return observation["allowance"]
 
 
-def main() -> int:
+def main(argv=None) -> int:
     # Emit only bounded local stage/category/timing diagnostics. Never include
     # raw exception text, native stderr, login values or account/RPC payloads.
-    result = check_native_receiver(dict(os.environ))
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--total-seconds", type=float)
+    parser.add_argument("--attempt-seconds", type=float)
+    args = parser.parse_args(argv)
+    try:
+        native_readiness_budget(args.total_seconds, args.attempt_seconds)
+    except ValueError as exc:
+        parser.error(str(exc))
+    result = check_native_receiver(dict(os.environ), total_seconds=args.total_seconds,
+                                   attempt_seconds=args.attempt_seconds)
     print(json.dumps(result), flush=True)
     return 0 if result["ok"] else 1
 

benchmark/runtime/sforge.py

@@ -22,6 +22,7 @@
 
 from .codex import DEFAULT_REPLAN_AFTER_TURNS, Execution, codex_settings, native_child_instructions, prepare_codex_home
 from .codex_offline import CodexOffline, MODEL_CATALOG_REMOTE, stage_portable_python
+from .codex_readiness import native_readiness_budget
 from .harbor import (
     BenchmarkCodex, _CONTROL, _CODEX_HOME, _PYTHON, _SCHEDULER_STATE, _SRC,
 )
@@ -98,6 +99,8 @@ def __init__(self, config, *, profile: str, cwd: str,
                  native_child_limit: int | None = None,
                  provider_request_admission: bool = False,
                  native_goal_continue: bool = False,
+                 native_readiness_total_seconds: float | None = None,
+                 native_readiness_attempt_seconds: float | None = None,
                  replan_after_turns: int | None = None,
                  replan_after_todos: int | None = None,
                  stride_strategy: str | None = None,
@@ -108,6 +111,12 @@ def __init__(self, config, *, profile: str, cwd: str,
         super().__init__(config)
         if profile not in PROFILES:
             raise ValueError("Unknown benchmark worker profile")
+        self.native_readiness_budget = native_readiness_budget(
+            native_readiness_total_seconds, native_readiness_attempt_seconds)
+        if self.native_readiness_budget is not None and (
+            profile in {"official", "single"} or config.agent_api_base_url or warm_start is not None
+        ):
+            raise ValueError("Native readiness budget requires a fresh native Goal or heartbeat worker")
         if not isinstance(turn_envelope, bool) or (
             turn_envelope and profile not in {"heartbeat-resume", "heartbeat-explore"}
         ):
@@ -452,6 +461,9 @@ def _install_worker(self, backend, handle, log_dir, logger):
             **(self.runtime._turn_envelope_receipt() if self.runtime is not None else {}),
             **({"provider_request_admission": "on"} if self.provider_request_admission else {}),
             **({"native_goal_continue": True} if self.native_goal_continue else {}),
+            **({"native_readiness_budget": dict(zip(("total_seconds", "attempt_seconds"),
+                                                    self.native_readiness_budget))}
+               if self.native_readiness_budget is not None else {}),
             "outer_resume": self.resume_cmd is not None,
             "explore_graph": self.profile == "heartbeat-explore",
             "explore_harness": self.profile == "heartbeat-explore",
@@ -561,8 +573,11 @@ def format_run_cmd(self, prompt_path, *, model=None, cwd="", internet=True, resu
         # their own admission path; no observer runs on every heartbeat wake.
         readiness = ""
         if not self._config.agent_api_base_url:
+            readiness_options = (f" --total-seconds {self.native_readiness_budget[0]}"
+                                 f" --attempt-seconds {self.native_readiness_budget[1]}"
+                                 if self.native_readiness_budget is not None else "")
             readiness = (f"env {exports} {shlex.quote(f'{_PYTHON}/bin/python3')} "
-                         "-m benchmark.runtime.codex_readiness; ")
+                         f"-m benchmark.runtime.codex_readiness{readiness_options}; ")
         return self._execution_command(
             f"set -eu; {readiness}{clock}"
             f"export LOOPX_PHASE_DEADLINE_EPOCH=$(cat {deadline}); "

benchmark/tests/test_codex_readiness.py

@@ -3,6 +3,7 @@
 import os
 import subprocess
 import sys
+import time
 from pathlib import Path
 
 import pytest
@@ -19,7 +20,7 @@ def native_fixture(tmp_path, mode):
     trace = tmp_path / "requests.jsonl"
     binary = tmp_path / "codex"
     binary.write_text(f"#!{sys.executable}\n" + '''
-import json, os, sys
+import json, os, sys, time
 from pathlib import Path
 counter = Path(os.environ["PROBE_ATTEMPTS"])
 attempt = int(counter.read_text()) + 1 if counter.exists() else 1
@@ -35,6 +36,10 @@ def native_fixture(tmp_path, mode):
     if "id" not in packet:
         continue
     method = packet["method"]
+    blocked = {"account_timeout": "account/read", "account_timeout_once": "account/read",
+               "usage_timeout_once": "account/rateLimits/read", "initialize_timeout": "initialize"}
+    if method == blocked.get(mode) and (not mode.endswith("_once") or attempt == 1):
+        time.sleep(60)  # The real stdio owner must time out and terminate this observer.
     if mode in ("disconnect", "disconnect_once") and method == "initialize":
         if mode == "disconnect" or attempt == 1:
             sys.exit(0)
@@ -298,24 +303,29 @@ def observer(*args, **kwargs):
 
 
 @pytest.mark.parametrize("profile", ["native-goal", "heartbeat-resume", "heartbeat-explore"])
-def test_sforge_rejects_bad_receiver_before_clock_or_scheduler(tmp_path, monkeypatch, profile):
+@pytest.mark.parametrize("budget", [False, True])
+def test_sforge_rejects_bad_receiver_before_clock_or_scheduler(tmp_path, monkeypatch, profile, budget):
     pytest.importorskip("sforge")
     pytest.importorskip("harbor")
     from sforge.harness.config import SForgeConfig
     from benchmark.runtime.sforge import SForgeWorker, BenchmarkCodex, _SRC, _PYTHON, _CODEX_HOME
     from pathlib import Path
 
     env, trace, home, pending = native_fixture(tmp_path, "account/read")
+    env.update(SFORGE_AGENT_REMAINING_SECONDS="100", SFORGE_AGENT_SEGMENT_STARTED_AT=str(time.time()))
     monkeypatch.setenv("CODEX_AUTH_JSON_PATH", str(home / "auth.json"))
 
     async def installed(self, environment):
         pass
 
     monkeypatch.setattr(BenchmarkCodex, "install", installed)
     worker = SForgeWorker(SForgeConfig(agent_model="fixture", agent_effort="xhigh"),
-                          profile=profile, cwd=str(tmp_path))
+                          profile=profile, cwd=str(tmp_path),
+                          **({"native_readiness_total_seconds": 75} if budget else {}))
     worker.install_stop_hook(None, None, tmp_path, None)
     worker.prepared = True
+    selected = json.loads((tmp_path / "worker-profile.json").read_text()).get("native_readiness_budget")
+    assert selected == ({"total_seconds": 75, "attempt_seconds": 25} if budget else None)
     command = worker.format_run_cmd("/task.md")
     checkout = str(Path(__file__).resolve().parents[2])
     runtime = tmp_path / "runtime"