DRU-989 -- A run waiting for the usage window shows as running

backend/druks/agents.py

@@ -5,7 +5,7 @@
 from collections.abc import Awaitable, Callable
 from contextvars import ContextVar
 from dataclasses import dataclass, field
-from datetime import UTC, datetime
+from datetime import UTC, datetime, timedelta
 from typing import TYPE_CHECKING, Any
 
 from dbos import DBOS, StepOptions
@@ -284,6 +284,10 @@ async def _quota_retry_wait() -> float:
                     if wait_seconds > _QUOTA_MAX_WAIT_SECONDS:
                         raise
                     wait_seconds += random.uniform(60, 120)
+                    await set_run_phase(
+                        "waiting_for_quota",
+                        ends_at=datetime.now(UTC) + timedelta(seconds=wait_seconds),
+                    )
                 else:
                     # Permanent or unclassified — never retried.
                     raise

backend/druks/durable/activity.py

@@ -1,13 +1,21 @@
+from datetime import datetime
+
 from dbos import DBOS
 
 
-async def set_run_phase(phase: str) -> None:
+async def set_run_phase(phase: str, *, ends_at: datetime | None = None) -> None:
     # The live infra/agent phase, pushed as a DBOS event (transient by nature —
     # overwritten each phase, gone when the run ends). NOT a durable column.
     # A chat or browser box provisions outside a run: nothing shows its phase.
     if DBOS.workflow_id:
         await DBOS.set_event_async("phase", phase)
+        if ends_at:
+            await DBOS.set_event_async("phase_ends_at", ends_at)
 
 
 async def get_run_phase(run_id: str) -> str | None:
     return await DBOS.get_event_async(run_id, "phase", timeout_seconds=0)
+
+
+async def get_run_phase_ends_at(run_id: str) -> datetime | None:
+    return await DBOS.get_event_async(run_id, "phase_ends_at", timeout_seconds=0)

backend/druks/durable/reads.py

@@ -10,7 +10,7 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from druks.database import session_scope
-from druks.durable.activity import get_run_phase
+from druks.durable.activity import get_run_phase, get_run_phase_ends_at
 from druks.durable.exceptions import AgentCallNotFound
 from druks.durable.live import keepalive_comment, serialize_model_event
 
@@ -126,11 +126,14 @@ async def get_subject_response(
     # eager-loaded here.
     runs = await Run.list_for_subject(session, subject_type, subject_id, include_calls=True)
     latest = runs[0] if runs else None
+    phase = await get_run_phase(latest.id) if latest and latest.is_running else None
+    ends_at = await get_run_phase_ends_at(latest.id) if phase else None
     return SubjectResponse(
         summary=summary,
         status=await _status(latest),
         timeline=await _timeline(runs),
-        phase=await get_run_phase(latest.id) if latest and latest.is_running else None,
+        phase=phase,
+        phase_ends_at=ends_at if ends_at and ends_at > datetime.now(UTC) else None,
     )
 
 

backend/druks/durable/schemas.py

@@ -207,9 +207,10 @@ class SubjectResponse(Schema):
     status: SubjectStatus
     # The subject's runs, oldest first, each with its agent calls — the timeline.
     timeline: list[RunResponse] = Field(default_factory=list)
-    # The driving run's sandbox phase while it starts ("provisioning_vm"). The
-    # shell supplies the words.
+    # The driving run's live phase while it starts or waits. The shell supplies
+    # the words. A timed phase carries the time when it ends separately.
     phase: str | None = None
+    phase_ends_at: datetime | None = None
 
 
 class TranscriptChunk(Schema):

backend/druks/harnesses/base.py

@@ -78,7 +78,8 @@ async def build_invocation(self, manifest: Manifest, **kwargs: object) -> AgentI
     @abstractmethod
     def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> object:
         """Turn a finished run into the structured payload (and write the
-        cost/output sidecars under ``artifact_dir / run_id``)."""
+        cost/output sidecars under ``artifact_dir / run_id``). Write any
+        available cost sidecar before a run failure raises."""
 
     @classmethod
     def check_returncode(cls, result: HarnessRunResult) -> None:

backend/druks/harnesses/claude.py

@@ -101,6 +101,7 @@ def get_acp_session(
     # {…overloaded_error…}"), so "api error: 5" covers 529 and every 5xx; 429
     # and 401 are their own families.
     failure_markers = {
+        "monthly spend limit": exceptions.HarnessSpendLimitError,
         "session limit": exceptions.HarnessUsageLimitError,
         "usage limit": exceptions.HarnessUsageLimitError,
         "api error: 429": exceptions.HarnessRateLimitError,
@@ -213,18 +214,20 @@ async def build_invocation(
         )
 
     def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> Any:
-        self.check_returncode(result)
-
         try:
             envelope = collapse_claude_stream(result.stdout)
         except exceptions.StreamJsonError as error:
+            # Keep a provider failure when its stream is incomplete. Only a
+            # successful malformed response is an output failure.
+            self.check_returncode(result)
             transcript = artifact_dir / run_id / "stdout.jsonl"
             raise exceptions.HarnessInvalidOutputError(
                 f"claude wrote invalid stream-json. See {transcript}",
             ) from error
 
         cost_usd, cost_metadata = extract_claude_cost_from_envelope(envelope)
         write_cost(artifact_dir / run_id, cost_usd=cost_usd, metadata=cost_metadata)
+        self.check_returncode(result)
 
         structured: Any = envelope.get("structured_output")
         if structured is None:

backend/druks/harnesses/codex.py

@@ -520,21 +520,21 @@ async def build_invocation(
         )
 
     def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> Any:
+        call_dir = artifact_dir / run_id
+        cost_usd, cost_metadata = read_codex_cost_from_jsonl(
+            call_dir / "session.jsonl",
+            model=self.model_id,
+        )
+        write_cost(call_dir, cost_usd=cost_usd, metadata=cost_metadata)
         self.check_returncode(result)
 
-        call_dir = artifact_dir / run_id
         payload = read_result_json(
             call_dir / "output.json",
             name=self.name,
         )
         (call_dir / "output.json").write_text(
             json.dumps(payload, indent=2, sort_keys=True),
         )
-        cost_usd, cost_metadata = read_codex_cost_from_jsonl(
-            call_dir / "session.jsonl",
-            model=self.model_id,
-        )
-        write_cost(call_dir, cost_usd=cost_usd, metadata=cost_metadata)
         return payload
 
     def _mcp_flags(self, servers: tuple[McpServer, ...]) -> tuple[str, ...]:

backend/druks/harnesses/opencode.py

@@ -148,21 +148,9 @@ async def build_invocation(
         )
 
     def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> Any:
-        if result.returncode == 124:
-            raise exceptions.HarnessTimeoutError(
-                "opencode hit the agent deadline; the session was aborted."
-            )
         try:
             responses = [json.loads(line)["info"] for line in result.stdout.splitlines()]
             info = responses[-1]
-            if failure := info.get("error"):
-                data = failure.get("data") or {}
-                message = data.get("message")
-                detail = f"{failure['name']}: {message}" if message else failure["name"]
-                if failure["name"] == "APIError" and data.get("isRetryable"):
-                    raise exceptions.HarnessOverloadedError(detail)
-                raise _ERROR_TYPES.get(failure["name"], exceptions.HarnessError)(detail)
-            structured = info["structured"]
             tokens = [response["tokens"] for response in responses]
             metadata = {
                 "provider": self.name,
@@ -174,7 +162,26 @@ def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) ->
                 "cache_creation_tokens": sum(entry["cache"]["write"] for entry in tokens),
             }
             cost_usd = sum(float(response["cost"]) for response in responses)
+            output_dir = call_dir(artifact_dir, run_id)
+            write_cost(output_dir, cost_usd=cost_usd, metadata=metadata)
+            if result.returncode == 124:
+                raise exceptions.HarnessTimeoutError(
+                    "opencode hit the agent deadline; the session was aborted."
+                )
+            if failure := info.get("error"):
+                data = failure.get("data") or {}
+                message = data.get("message")
+                detail = f"{failure['name']}: {message}" if message else failure["name"]
+                if failure["name"] == "APIError" and data.get("isRetryable"):
+                    raise exceptions.HarnessOverloadedError(detail)
+                raise _ERROR_TYPES.get(failure["name"], exceptions.HarnessError)(detail)
+            structured = info["structured"]
+            self.check_returncode(result)
         except (ValueError, KeyError, TypeError, AttributeError, IndexError) as error:
+            if result.returncode == 124:
+                raise exceptions.HarnessTimeoutError(
+                    "opencode hit the agent deadline; the session was aborted."
+                ) from error
             try:
                 self.check_returncode(result)
             except exceptions.HarnessError as process_error:
@@ -183,8 +190,5 @@ def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) ->
                 "opencode wrote no usable response.",
             ) from error
 
-        self.check_returncode(result)
-        output_dir = call_dir(artifact_dir, run_id)
         (output_dir / "output.json").write_text(json.dumps(structured, indent=2, sort_keys=True))
-        write_cost(output_dir, cost_usd=cost_usd, metadata=metadata)
         return structured

backend/druks/harnesses/pi.py

@@ -126,7 +126,6 @@ async def build_invocation(
         )
 
     def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> Any:
-        self.check_returncode(result)
         try:
             events = [json.loads(line) for line in result.stdout.splitlines()]
             # pi repeats an assistant message on turn_end and agent_end;
@@ -136,6 +135,21 @@ def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) ->
                 for event in events
                 if event["type"] == "message_end" and event["message"]["role"] == "assistant"
             ]
+            usage = [message["usage"] for message in messages]
+            cost_usd = sum(entry["cost"]["total"] for entry in usage)
+            tokens = {
+                "input_tokens": sum(entry["input"] for entry in usage),
+                "output_tokens": sum(entry["output"] for entry in usage),
+                "cached_input_tokens": sum(entry["cacheRead"] for entry in usage),
+                "cache_creation_tokens": sum(entry["cacheWrite"] for entry in usage),
+            }
+            call_dir = artifact_dir / run_id
+            write_cost(
+                call_dir,
+                cost_usd=cost_usd,
+                metadata={"provider": self.name, "model": self.model, **tokens},
+            )
+            self.check_returncode(result)
             for message in messages:
                 # A runtime provider failure exits zero, so the HTTP status in
                 # ``errorMessage`` is what classifies it.
@@ -146,24 +160,11 @@ def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) ->
                     if code >= 500:
                         raise exceptions.HarnessOverloadedError(detail)
                     raise _STATUS_ERRORS.get(code, exceptions.HarnessError)(detail)
-            usage = [message["usage"] for message in messages]
-            cost_usd = sum(entry["cost"]["total"] for entry in usage)
-            tokens = {
-                "input_tokens": sum(entry["input"] for entry in usage),
-                "output_tokens": sum(entry["output"] for entry in usage),
-                "cached_input_tokens": sum(entry["cacheRead"] for entry in usage),
-                "cache_creation_tokens": sum(entry["cacheWrite"] for entry in usage),
-            }
         except (ValueError, KeyError, TypeError) as error:
+            self.check_returncode(result)
             raise exceptions.HarnessInvalidOutputError("pi wrote no usable stream.") from error
 
-        call_dir = artifact_dir / run_id
         output_path = call_dir / "output.json"
         payload = read_result_json(output_path, name=self.name)
         output_path.write_text(json.dumps(payload, indent=2, sort_keys=True))
-        write_cost(
-            call_dir,
-            cost_usd=cost_usd,
-            metadata={"provider": self.name, "model": self.model, **tokens},
-        )
         return payload

backend/tests/conftest.py

@@ -158,6 +158,7 @@ async def _dbos_cancel(workflow_id: str) -> None:
         mock.patch.object(Workflow, "start", classmethod(_noop)),
         mock.patch("druks.agents.set_run_phase", _phase_noop),
         mock.patch("druks.durable.reads.get_run_phase", _phase_noop),
+        mock.patch("druks.durable.reads.get_run_phase_ends_at", _phase_noop),
         mock.patch("dbos.DBOS.cancel_workflow_async", _dbos_cancel),
         mock.patch("dbos.DBOS.start_workflow_async", _noop),
     ):

backend/tests/software_factory/test_api_work_items.py

@@ -248,11 +248,33 @@ async def test_detail_carries_the_running_phase(client: TestClient, druks_db, mo
     await seed_build_run(druks_db, work_item_id=item.id, state="running")
 
     async def phase(_run_id):
-        return "provisioning_vm"
+        return "waiting_for_quota"
+
+    async def phase_ends_at(_run_id):
+        return datetime(2099, 8, 1, 14, 19, tzinfo=UTC)
+
+    monkeypatch.setattr("druks.durable.reads.get_run_phase", phase)
+    monkeypatch.setattr("druks.durable.reads.get_run_phase_ends_at", phase_ends_at)
+    detail = (await client.get(f"/api/software_factory/work_item/{item.id}")).json()
+    assert detail["phase"] == "waiting_for_quota"
+    assert detail["phaseEndsAt"] == "2099-08-01T14:19:00Z"
+
+
+async def test_detail_omits_a_phase_end_that_has_passed(client: TestClient, druks_db, monkeypatch):
+    item = await make_test_work_item(repo="ClawHaven/acme-app", title="x")
+    await seed_build_run(druks_db, work_item_id=item.id, state="running")
+
+    async def phase(_run_id):
+        return "waiting_for_quota"
+
+    async def phase_ends_at(_run_id):
+        return datetime(2000, 1, 1, tzinfo=UTC)
 
     monkeypatch.setattr("druks.durable.reads.get_run_phase", phase)
+    monkeypatch.setattr("druks.durable.reads.get_run_phase_ends_at", phase_ends_at)
     detail = (await client.get(f"/api/software_factory/work_item/{item.id}")).json()
-    assert detail["phase"] == "provisioning_vm"
+
+    assert not detail["phaseEndsAt"]
 
 
 async def test_detail_has_no_phase_when_not_running(client: TestClient, druks_db):
@@ -261,3 +283,4 @@ async def test_detail_has_no_phase_when_not_running(client: TestClient, druks_db
 
     detail = (await client.get(f"/api/software_factory/work_item/{item.id}")).json()
     assert not detail["phase"]
+    assert not detail["phaseEndsAt"]

backend/tests/test_agents.py

@@ -867,6 +867,8 @@ async def test_body_level_quota_waits_for_the_reset_once(
     )
     jitter = MagicMock(return_value=90.0)
     monkeypatch.setattr(agents.random, "uniform", jitter)
+    set_run_phase = AsyncMock()
+    monkeypatch.setattr(agents, "set_run_phase", set_run_phase)
     current_run._reap_run = AsyncMock()
     checkpoints, sleep = _inline_agent_steps
 
@@ -876,6 +878,10 @@ async def test_body_level_quota_waits_for_the_reset_once(
     assert excinfo.value is failures[-1]
     sleep.assert_awaited_once_with(60 * 60 + 90.0)
     jitter.assert_called_once_with(60.0, 120.0)
+    set_run_phase.assert_any_await(
+        "waiting_for_quota",
+        ends_at=now + timedelta(hours=1, seconds=90),
+    )
     current_run._reap_run.assert_awaited_once()
     assert [options["name"] for options in checkpoints] == [
         "test.agent.dummy",