diff --git a/benchmarks/search/pilot_run.py b/benchmarks/search/pilot_run.py index 477dc35..e2b414e 100644 --- a/benchmarks/search/pilot_run.py +++ b/benchmarks/search/pilot_run.py @@ -88,6 +88,9 @@ class Pilot: fixed_effect_observed=None, outcome="no_evidence", end_reason="exhausted", + end_detail=None, + task_metrics=[], + executor_task_metrics=[], tasks_submitted=0, executor_tasks=0, verify_tasks=0, @@ -139,9 +142,35 @@ class Pilot: @contextmanager def executor(self, revision, deadline=900): + before = time.monotonic() + metrics = dict(self.worker.metrics) + self.executor_end = "ok" + submitted = self.record["executor_tasks"] + try: + with self._executor(revision, deadline) as task: + yield task + except BudgetExhausted as exc: + self.executor_end = exc.end_detail + raise + except Exception: + self.executor_end = "execution_failure" + raise + finally: + if self.record["executor_tasks"] > submitted: + self.record["executor_task_metrics"].append( + { + "executor_task": submitted + 1, + **{key: value - metrics[key] for key, value in self.worker.metrics.items()}, + "wall_seconds": time.monotonic() - before, + "end": self.executor_end, + } + ) + + @contextmanager + def _executor(self, revision, deadline=900): deadline = min(deadline, self.deadline - time.monotonic()) if deadline < 1: - raise BudgetExhausted("task deadline", usd=0, calls=0) + raise BudgetExhausted("task deadline", usd=0, calls=0, end_detail="task_wall") acquired_at = time.monotonic() empty = io.BytesIO() with tarfile.open(fileobj=empty, mode="w"): @@ -151,7 +180,7 @@ class Pilot: with SearchExecutorTask(self.lane, empty.getvalue(), deadline=deadline) as task: remaining = deadline - (time.monotonic() - acquired_at) if remaining < 1: - raise BudgetExhausted("acquisition deadline", usd=0, calls=0) + raise BudgetExhausted("acquisition deadline", usd=0, calls=0, end_detail="task_wall") timeout = min(180, remaining) result = task.client.submit( "run", @@ -163,7 +192,7 @@ class Pilot: raise RepositoryFailure("checkout failed") remaining = deadline - (time.monotonic() - acquired_at) if remaining < 1: - raise BudgetExhausted("acquisition deadline", usd=0, calls=0) + raise BudgetExhausted("acquisition deadline", usd=0, calls=0, end_detail="task_wall") result = task.client.submit("read_file", {"path": self.row["file"]}, timeout_seconds=min(30, remaining)) if not isinstance(result.get("input_bytes"), int) or result["input_bytes"] <= 0: self.failure("acquisition", "candidate_input_unreadable") @@ -173,7 +202,7 @@ class Pilot: def admit_task(self): if self.record["tasks_submitted"] >= 32: - raise BudgetExhausted("executor task cap", usd=0, calls=0) + raise BudgetExhausted("executor task cap", usd=0, calls=0, end_detail="tasks") self.record["tasks_submitted"] += 1 def prepare(self, revision, spec): @@ -247,6 +276,9 @@ class Pilot: self.phase = task.step before = time.monotonic() self.deadline = before + task.limits.get("task_wall_seconds", 900) + metrics_before = dict(self.worker.metrics) + result = {} + end = "execution_failure" try: if task.step == "verify": result = self.verification() @@ -267,6 +299,7 @@ class Pilot: self.failure("verify/" + side.outcome, "side_verification_failed") if result.outcome != "demonstrated": self.failure("verify", "no_consistent_differential" if result.outcome == "inconclusive" else result.outcome) + end = result.outcome return dict(status="ok", outcome=result.outcome, evidence_refs=["verify:owned"], cost_usd=0) if task.step == "explore": if self.args.dry_run: @@ -277,20 +310,32 @@ class Pilot: self.worker.executor = remote.client remaining = task.limits["task_wall_seconds"] - (time.monotonic() - before) if remaining < 1: - raise BudgetExhausted("explore deadline", usd=0, calls=0) + raise BudgetExhausted("explore deadline", usd=0, calls=0, end_detail="task_wall") result = self.worker(replace(task, limits={**task.limits, "task_wall_seconds": remaining})) + self.executor_end = result.get("end_detail") or result.get("status", "execution_failure") else: result = self.worker(task) - if self.worker.budget_exhausted: - raise BudgetExhausted("run ceiling", usd=0, calls=0) + end = result.get("end_detail") or result.get("status", "execution_failure") if result.get("status") != "ok": self.failure(task.step, result.get("failure", "worker_execution_failed")) - if self.worker.metrics["model_calls"] == 0: + if self.worker.metrics["model_calls"] == 0 and not result.get("budget_exhaustion"): raise RuntimeError(self.record["failures"][-1]["reason"]) if not self.verifications: self.record["outcome"] = "inconclusive" return result + except BudgetExhausted as exc: + end = exc.end_detail + raise finally: + self.record["task_metrics"].append( + { + "task_id": task.id, + "step": task.step, + **{key: value - metrics_before[key] for key, value in self.worker.metrics.items()}, + "wall_seconds": time.monotonic() - before, + "end": end, + } + ) wall = self.record["wall_seconds_by_phase"] wall[task.step] = wall.get(task.step, 0) + time.monotonic() - before @@ -300,6 +345,7 @@ class Pilot: try: state = coordinator.run() self.record["end_reason"] = state["end_reason"] + self.record["end_detail"] = state.get("end_detail") if not self.verifications and any(f.get("step") == "explore" for f in state["facts"]): self.record["outcome"] = "plausible_unproven" except Exception as exc: diff --git a/src/openultrasast/plane/budget.py b/src/openultrasast/plane/budget.py index df76415..aa6b1fe 100644 --- a/src/openultrasast/plane/budget.py +++ b/src/openultrasast/plane/budget.py @@ -23,8 +23,9 @@ _USAGE_FIELDS = ("prompt_tokens", "prompt_cache_hit_tokens", "completion_tokens" class BudgetExhausted(RuntimeError): """Raised before a call would start once the usd or calls ceiling is reached.""" - def __init__(self, message: str, *, usd: float | None, calls: int) -> None: - super().__init__(message) + def __init__(self, message: str, *, usd: float | None, calls: int, end_detail: str = "spend") -> None: + super().__init__(f"{end_detail}: {message}") + self.end_detail = end_detail self.usd = usd self.calls = calls diff --git a/src/openultrasast/search/board.py b/src/openultrasast/search/board.py index 94340a6..3f67f62 100644 --- a/src/openultrasast/search/board.py +++ b/src/openultrasast/search/board.py @@ -27,7 +27,7 @@ class Board: self._lock = threading.RLock() self.head: str | None = None self._state: dict[str, Any] = dict( - search_id=search_id, version=0, previous=None, facts=[], intents=[], hints=[], end_reason=None, checkpoint={} + search_id=search_id, version=0, previous=None, facts=[], intents=[], hints=[], end_reason=None, end_detail=None, checkpoint={} ) @property @@ -54,7 +54,7 @@ class Board: raise PermissionError("only the coordinator may write the board") if self._state["end_reason"] is not None: raise ValueError("search already ended") - if changes.keys() - {"facts", "intents", "hints", "end_reason", "checkpoint"}: + if changes.keys() - {"facts", "intents", "hints", "end_reason", "end_detail", "checkpoint"}: raise ValueError("unknown board fields") state = {**self.state, **copy.deepcopy(changes), "version": self._state["version"] + 1, "previous": self.head} self._validate(state) diff --git a/src/openultrasast/search/coordinator.py b/src/openultrasast/search/coordinator.py index e9c6c49..02606da 100644 --- a/src/openultrasast/search/coordinator.py +++ b/src/openultrasast/search/coordinator.py @@ -67,6 +67,7 @@ class Coordinator: self.board, self.executor = board, executor self.budget, self.clock, self.push_budget = budget or SearchBudget(), clock, push_budget self.state = board.state + self.state.setdefault("end_detail", None) self.progress = self.state["checkpoint"] or dict( started_at=clock(), spent_usd=0.0, @@ -84,10 +85,25 @@ class Coordinator: self.progress["wall_seconds"] = max(0, self.clock() - self.progress["started_at"]) self.state["checkpoint"] = self.progress self.board.commit( - writer=self.board.coordinator, **{k: self.state[k] for k in ("facts", "intents", "hints", "end_reason", "checkpoint")} + writer=self.board.coordinator, + **{k: self.state[k] for k in ("facts", "intents", "hints", "end_reason", "end_detail", "checkpoint")}, ) - def _end(self, reason: str) -> dict[str, Any]: + def _end(self, reason: str, end_detail: str | None = None) -> dict[str, Any]: + if reason == "budget_spent" and end_detail not in { + "spend", + "task_spend", + "reason_rounds", + "demonstrations", + "tasks", + "task_wall", + "search_wall", + "worker_steps", + "worker_wall", + "run_ceiling", + }: + raise ValueError("budget_spent requires an exhausted limit") + self.state["end_detail"] = end_detail self.state["end_reason"] = reason self._save() return self.board.state @@ -110,16 +126,21 @@ class Coordinator: while True: remaining = self._remaining() if remaining <= self.budget.cleanup_seconds or self.progress["tasks"] >= self.budget.tasks: - raise BudgetExhausted("task or wall budget exhausted", usd=self.meter.spent, calls=self.progress["tasks"]) + raise BudgetExhausted( + "task or wall budget exhausted", + usd=self.meter.spent, + calls=self.progress["tasks"], + end_detail="search_wall" if remaining <= self.budget.cleanup_seconds else "tasks", + ) maximum = self.budget.max_call_usd if step != "verify" else 0 reservation = self.meter.reserve(maximum) push_reservation = None try: if self.push_budget is not None: push_reservation = self.push_budget.reserve(maximum) - except BudgetExhausted: + except BudgetExhausted as exc: reservation.release() - raise + raise BudgetExhausted("push ceiling", usd=exc.usd, calls=exc.calls, end_detail="run_ceiling") from exc task_id = "task-" + str(self.progress["tasks"] + 1) limits = asdict(self.budget) limits["task_wall_seconds"] = min(self.budget.task_wall_seconds, remaining - self.budget.cleanup_seconds) @@ -164,7 +185,16 @@ class Coordinator: or self._remaining() <= 0 or self.clock() - dispatched_at > limits["task_wall_seconds"] ): - raise BudgetExhausted("executor exceeded cost or wall contract", usd=self.meter.spent, calls=self.progress["tasks"]) + raise BudgetExhausted( + "executor exceeded cost or wall contract", + usd=self.meter.spent, + calls=self.progress["tasks"], + end_detail="task_spend" + if result.get("status") == "invalid_cost" + else "search_wall" + if self._remaining() <= 0 + else "task_wall", + ) result = {**result, "task_id": task_id} # Persist results together with their effects in run(), not a "done" # checkpoint which could lose the result if the host then crashes. @@ -209,13 +239,13 @@ class Coordinator: self._save() while True: if self._remaining() <= self.budget.cleanup_seconds: - return self._end("budget_spent") + return self._end("budget_spent", "search_wall") try: intent = next((i for i in self.state["intents"] if i["status"] == "open"), None) demo = next((f for f in self.state["facts"] if f.get("demo") and f["id"] not in self.progress["verified"]), None) if demo is not None: if self.progress["demonstrations"] >= self.budget.demonstrations: - return self._end("budget_spent") + return self._end("budget_spent", "demonstrations") self.progress["demonstrations"] += 1 result = self._call("verify", {"fact_id": demo["id"], "demo": demo["demo"]}) self.progress["verified"].append(demo["id"]) @@ -259,7 +289,7 @@ class Coordinator: intent["failure"] = "invalid_result" if result.get("status") == "ok" else result.get("status", "execution_failure") else: if self.progress["rounds"] >= self.budget.reason_rounds: - return self._end("budget_spent") + return self._end("budget_spent", "reason_rounds") self.progress["rounds"] += 1 result = self._call("reason", {}) if result.get("status") == "ok" and not result.get("intents") and "complete" not in result: @@ -275,6 +305,8 @@ class Coordinator: if result.get("status") == "ok" and not result.get("intents") and "complete" not in result: self.progress["failures"].append({"task_id": result["task_id"], "reason": "reason_refused"}) return self._end("reason_refused") + if result.get("budget_exhaustion"): + return self._end("budget_spent", result["budget_exhaustion"]["end_detail"]) proposals = result.get("intents", []) if result.get("status") != "ok" or not isinstance(proposals, list) or len(proposals) > self.budget.intents_per_round: return self._end("exhausted") @@ -304,9 +336,11 @@ class Coordinator: ) if not added: return self._end("exhausted") + if result.get("budget_exhaustion"): + return self._end("budget_spent", result["budget_exhaustion"]["end_detail"]) self._save() - except BudgetExhausted: - return self._end("budget_spent") + except BudgetExhausted as exc: + return self._end("budget_spent", exc.end_detail) @staticmethod def _refs(refs: Any) -> bool: diff --git a/src/openultrasast/search/executor.py b/src/openultrasast/search/executor.py index 42cf849..16265f7 100644 --- a/src/openultrasast/search/executor.py +++ b/src/openultrasast/search/executor.py @@ -314,7 +314,7 @@ class InProcessExecutor: self.phase = "run" timed_out = False disk_exceeded = False - failure = None + failure: Exception | None = None # Monitor the full disk-backed workspace, including package caches. # An unnamespaced child can still hardcode RAM-backed paths: this # guard is not a replacement for the operator's memory limit. diff --git a/src/openultrasast/search/worker.py b/src/openultrasast/search/worker.py index e71e22a..9dba4a7 100644 --- a/src/openultrasast/search/worker.py +++ b/src/openultrasast/search/worker.py @@ -173,6 +173,8 @@ class Worker: self.run_budget = run_budget self.metrics = dict(reserved=0.0, settled=0.0, model_calls=0, tokens_in=0, tokens_out=0) self.budget_exhausted = False + self.task_calls_start = 0 + self.observations: list[tuple[str, str]] = [] selected_prices = prices if prices is not None else getattr(client, "prices", None) if not isinstance(selected_prices, Prices) or any(not math.isfinite(v) or v < 0 for v in vars(selected_prices).values()): raise ValueError("nonnegative finite model prices required before admission") @@ -192,26 +194,36 @@ class Worker: def _remaining(self) -> int: remaining = min(60, self.deadline - time.monotonic()) if remaining < 1: - raise BudgetExhausted("worker wall budget spent", usd=0, calls=0) + raise BudgetExhausted("worker wall budget spent", usd=0, calls=0, end_detail="worker_wall") return int(remaining) def _call(self, messages: list[dict[str, Any]], tools: list[dict[str, Any]], meter: SpendBudget) -> dict[str, Any]: timeout = self._remaining() - # UTF-8 bytes upper-bound content tokens; include conservative per-message - # chat framing. A provider with extra hidden billed tokens needs new pricing. - input_bound = len(json.dumps([messages, tools], ensure_ascii=False).encode()) + 4096 * (len(messages) + 1) + # Pilot measured ~27k input tokens against the old ~100k reservation (3-4x). + # Use bytes/2 plus 1024 framing tokens per message; retain max_tokens for output. + # This calibrated bound is checked at settlement, not a tokenizer guarantee. + input_bound = math.ceil(len(json.dumps([messages, tools], ensure_ascii=False).encode()) / 2) + 1024 * (len(messages) + 1) maximum = ( input_bound * max(self.prices.input_per_m, self.prices.cache_hit_per_m) + self.max_output_tokens * self.prices.output_per_m ) / 1_000_000 - reservation = meter.reserve(maximum) + try: + reservation = meter.reserve(maximum) + except BudgetExhausted as exc: + raise BudgetExhausted( + "next call reservation exceeds task allowance", + usd=exc.usd, + calls=int(self.metrics["model_calls"]) - self.task_calls_start, + end_detail="task_spend", + ) from exc run_reservation = None try: if self.run_budget is not None: run_reservation = self.run_budget.reserve(maximum) - except BudgetExhausted: + except BudgetExhausted as exc: reservation.release() - self.budget_exhausted = True - raise + raise BudgetExhausted( + "run ceiling", usd=exc.usd, calls=int(self.metrics["model_calls"]) - self.task_calls_start, end_detail="run_ceiling" + ) from exc self.metrics["reserved"] += maximum self.metrics["model_calls"] += 1 actual = maximum @@ -326,6 +338,7 @@ class Worker: {"role": "user", "content": task.snapshot + "\nIntent: " + json.dumps(task.payload)}, ] self.evidence = set() + self.observations = [] for index in range(self.max_steps): messages[0]["content"] = ( EXPLORE_PROMPT @@ -369,13 +382,20 @@ class Worker: ref = "tool:" + str(len(self.evidence) + 1) self.evidence.add(ref) result["evidence_ref"] = ref + if not result.get("error"): + text = json.dumps(result, ensure_ascii=False) + if len(text) > 2000: + text = text[:2000] + " [truncated]" + self.observations.append((ref, name + ": " + text)) except (ValueError, OSError, KeyError, TypeError) as exc: result = {"error": type(exc).__name__, "reason": exception_reason(exc)} messages.append({"role": "tool", "tool_call_id": call["id"], "content": json.dumps(result)}) - return {"status": "execution_failure", "failure": "tool step limit exceeded"} + raise BudgetExhausted("tool step limit exceeded", usd=meter.spent, calls=self.max_steps, end_detail="worker_steps") def __call__(self, task: SearchTask) -> dict[str, Any]: self.budget_exhausted = False + self.task_calls_start = int(self.metrics["model_calls"]) + self.observations = [] self.limits = task.limits self.deadline = time.monotonic() + task.limits.get("task_wall_seconds", 900) meter = SpendBudget(task.limits["max_call_usd"]) @@ -392,6 +412,27 @@ class Worker: "cost_usd": meter.spent, "isolation_mode": getattr(self.executor, "isolation_mode", "brain-only"), } + except BudgetExhausted as exc: + self.budget_exhausted = True + result = {"status": "execution_failure", "failure": exception_reason(exc)} + if task.step == "explore" and self.observations: + # Preserve literal executor observations, without inventing a security conclusion + # or making another paid call to summarize them. Bound board growth explicitly. + observations = self.observations[-6:] + result = { + "status": "ok", + "fact": { + "text": "Partial tool observations (search incomplete):\n" + "\n".join(text for _, text in observations), + "evidence_refs": [ref for ref, _ in observations], + }, + } + return { + **result, + "cost_usd": meter.spent, + "end_detail": exc.end_detail, + "budget_exhaustion": {"end_detail": exc.end_detail, "usd": exc.usd, "calls": exc.calls}, + "isolation_mode": getattr(self.executor, "isolation_mode", "brain-only"), + } except _sandbox.IsolationUnavailable as exc: return { "status": "instrument_failure", diff --git a/tests/test_search_coordinator.py b/tests/test_search_coordinator.py index 35edd61..7d9a629 100644 --- a/tests/test_search_coordinator.py +++ b/tests/test_search_coordinator.py @@ -164,3 +164,56 @@ def test_invalid_identifiers_refused_before_paths(tmp_path): Board(FileStore(tmp_path), bad, coordinator="host") with pytest.raises(ValueError): SearchTask(bad, "search-1", "reason", "", {}, asdict(SearchBudget())) + + +@pytest.mark.parametrize("detail", ["task_spend", "worker_steps", "worker_wall", "run_ceiling"]) +def test_worker_limit_preserves_fact_and_survives_resume(tmp_path, detail): + def execute(task): + if task.step == "reason": + return {"status": "ok", "cost_usd": 0, "intents": [{"description": "inspect"}]} + return { + "status": "ok", + "cost_usd": 0.01, + "budget_exhaustion": {"end_detail": detail, "usd": 0.01, "calls": 1}, + "fact": {"text": "Confirmed partial observation", "evidence_refs": ["tool:1"]}, + } + + b = board(tmp_path) + state = Coordinator(b, execute).run() + assert state["end_reason"] == "budget_spent" and state["end_detail"] == detail + assert state["facts"][0]["text"] == "Confirmed partial observation" + assert state["intents"][0]["status"] == "concluded" + assert state["checkpoint"]["spent_usd"] == 0.01 + assert Board.resume(b.store, "search-1", b.head, coordinator="host").state == state + + +@pytest.mark.parametrize("detail", ["spend", "tasks", "reason_rounds", "demonstrations", "task_wall", "search_wall", "run_ceiling"]) +def test_coordinator_identifies_each_own_limit(tmp_path, detail): + from openultrasast.plane.budget import SpendBudget + + now = [0] + + def execute(task): + if detail in {"task_wall", "search_wall"}: + now[0] = 1000 if detail == "task_wall" else 20000 + if task.step == "reason": + return {"status": "ok", "cost_usd": 0, "intents": [{"description": str(task.id)}]} + if task.step == "explore": + return {"status": "ok", "cost_usd": 0, "fact": {"text": "candidate", "evidence_refs": ["tool:1"], "demo": "demo/"}} + return {"status": "ok", "cost_usd": 0, "outcome": "inconclusive"} + + limits = { + "spend": {"spend_usd": 0}, + "tasks": {"tasks": 1}, + "reason_rounds": {"reason_rounds": 1}, + "demonstrations": {"demonstrations": 1}, + } + state = Coordinator( + board(tmp_path), + execute, + budget=SearchBudget(**limits.get(detail, {})), + clock=lambda: now[0], + push_budget=SpendBudget(0) if detail == "run_ceiling" else None, + ).run() + assert state["end_reason"] == "budget_spent" + assert state["end_detail"] == detail diff --git a/tests/test_search_pilot_run.py b/tests/test_search_pilot_run.py index 63f62a3..7dfa0ad 100644 --- a/tests/test_search_pilot_run.py +++ b/tests/test_search_pilot_run.py @@ -63,6 +63,8 @@ def test_ceiling_refuses_before_stub_call(tmp_path): record = json.loads(output.read_text().splitlines()[0]) assert record["model_calls"] == 0 assert record["end_reason"] == "budget_spent" + assert record["end_detail"] == "run_ceiling" + assert record["task_metrics"][0]["end"] == "run_ceiling" def make_pilot(tmp_path, *, side="fixed", dry_run=False): @@ -221,6 +223,13 @@ def test_live_wiring_with_fake_mailbox_no_network(tmp_path, monkeypatch, verify_ assert record["end_reason"] == "goal_met" assert record["model_calls"] == 3 assert record["executor_tasks"] == 3 + assert len(record["executor_task_metrics"]) == 3 + assert sum(t["model_calls"] for t in record["executor_task_metrics"]) == 3 + for row in record["executor_task_metrics"]: + assert row["reserved"] >= row["settled"] + assert row["wall_seconds"] > 0 + assert row["end"] == "ok" + assert {"tokens_in", "tokens_out"} <= row.keys() assert record["verify_tasks"] == (6 if verify_lane == "ax" else 0) assert record["tasks_submitted"] == (9 if verify_lane == "ax" else 3) assert record["input_bytes"] > 0 diff --git a/tests/test_search_worker.py b/tests/test_search_worker.py index 9c7c08e..7527750 100644 --- a/tests/test_search_worker.py +++ b/tests/test_search_worker.py @@ -307,3 +307,82 @@ def test_prompt_contracts_and_demo_example(tmp_path, step): assert len(prompt) / 4 < 1000 # rough offline token estimate, not a tokenizer else: assert "file:line" in prompt and "final call for finish" in prompt + + +@pytest.mark.parametrize("with_evidence", [False, True]) +def test_task_spend_finishes_with_confirmed_observations(tmp_path, with_evidence): + from openultrasast.model.endpoint import Prices + + (tmp_path / "input.txt").write_text("confirmed input\n" * 400) + + class UsageStub(StubModel): + def complete_chat_raw(self, **kwargs): + raw = super().complete_chat_raw(**kwargs) + raw["usage"] = {"prompt_tokens": 100, "completion_tokens": 100} + return raw + + model = UsageStub([tool("read_file", path="input.txt")]) + worker = Worker(model, repo=tmp_path, demo=tmp_path / "demo", prices=Prices(1, 1, 1), executor=InProcessExecutor(tmp_path)) + job = task("explore") + # First reservation fits; the second includes the real file contents and cannot fit. + job.limits["max_call_usd"] = 0.01 if with_evidence else 0.000001 + result = worker(job) + assert result["end_detail"] == result["budget_exhaustion"]["end_detail"] == "task_spend" + assert len(model.calls) == int(with_evidence) + assert worker.metrics["reserved"] >= worker.metrics["settled"] + if with_evidence: + assert result["status"] == "ok" + assert "confirmed input" in result["fact"]["text"] + assert result["fact"]["evidence_refs"] == ["tool:1"] + assert result["cost_usd"] == pytest.approx(0.0002) + else: + assert result["status"] == "execution_failure" + assert "task_spend" in result["failure"] + assert "fact" not in result + + +def test_calibrated_reservation_covers_reported_usage(tmp_path): + import math + + from openultrasast.model.endpoint import Prices + from openultrasast.plane.budget import SpendBudget + + class UsageStub(StubModel): + def complete_chat_raw(self, **kwargs): + raw = super().complete_chat_raw(**kwargs) + raw["usage"] = {"prompt_tokens": 27000, "completion_tokens": 1024} + return raw + + model = UsageStub([reply({})]) + worker = Worker(model, repo=tmp_path, demo=tmp_path / "demo", prices=Prices(1, 1, 2)) + worker.reason("facts: []\nnote: " + "x" * 85000, SpendBudget(1)) + call = model.calls[0] + expected_input = math.ceil(len(json.dumps([call["messages"], call["tools"]], ensure_ascii=False).encode()) / 2) + 1024 * ( + len(call["messages"]) + 1 + ) + assert worker.metrics["reserved"] == pytest.approx((expected_input + 1024 * 2) / 1_000_000) + assert worker.metrics["reserved"] >= worker.metrics["settled"] == pytest.approx(0.029048) + + +@pytest.mark.parametrize("detail", ["worker_steps", "worker_wall", "run_ceiling"]) +def test_worker_exhaustion_is_structured(tmp_path, monkeypatch, detail): + from openultrasast.model.endpoint import Prices + from openultrasast.plane.budget import SpendBudget + + worker = Worker( + StubModel([tool("list_files")]), + repo=tmp_path, + demo=tmp_path / "demo", + executor=InProcessExecutor(tmp_path), + max_steps=1, + prices=Prices(1, 1, 1), + run_budget=SpendBudget(0 if detail == "run_ceiling" else 1), + ) + if detail == "worker_wall": + from openultrasast.search import worker as module + + times = iter([0, 901]) + monkeypatch.setattr(module.time, "monotonic", lambda: next(times)) + result = worker(task("explore")) + assert result["budget_exhaustion"]["end_detail"] == detail + assert result["end_detail"] == detail