fix(api_server): GET /v1/runs/{id} reports the served fallback runtime and cache-read tokens

GET /v1/runs/{run_id} and the run.completed SSE event only echoed the
requested model and three token counters. After a fallback_providers
switch the run record still named the requested model and usage had no
cache figures, so a supervisor polling /v1/runs for cost attribution
booked the whole run to the wrong provider at the wrong price, with
cache reads counted as full-price input.

agent.provider / agent.model still hold the fallback pair when
run_conversation() returns (the primary is restored only at the start of
the next turn), and agent.session_cache_read_tokens has the cache reads.
Surface them on the completed run status and the terminal event:

- usage gains cache_read_tokens / cache_write_tokens (_USAGE_FIELDS)
- _run_agent_sync returns the served {provider, model} as a third value
  and _finish stamps it as `runtime` on both the pollable status and the
  run.completed event (the persisted idempotent record inherits it).

Ported onto the refactored _run_agent_sync / _finish shape from #102161;
the served pair is read from the agent (the reporter's diff in #102101)
rather than the turn record, so a run whose result dict lacks the keys is
still attributed. Tests trimmed to two invariants.

Fixes #102101
This commit is contained in:
Tranquil-Flow
2026-09-03 14:20:28 +02:00
committed by Teknium
parent 62277ebe29
commit 1b02df86e3
3 changed files with 118 additions and 7 deletions

View File

@@ -37,10 +37,12 @@ _SUBAGENT_EVENT_KEYS = (
"output_tokens", "reasoning_tokens", "api_calls", "cost_usd", "files_read", "files_written",
"output_tail")
_SUBAGENT_TEXT_KEYS = ("goal", "summary", "output_tail")
# Terminal usage payload: (wire key, agent attribute), in wire order.
# Terminal usage payload: (wire key, agent attribute), in wire order. Cache reads ride along so a
# cost poller does not book them as full-price input (#102101).
_USAGE_FIELDS = (
("input_tokens", "session_prompt_tokens"), ("output_tokens", "session_completion_tokens"),
("total_tokens", "session_total_tokens"))
("total_tokens", "session_total_tokens"), ("cache_read_tokens", "session_cache_read_tokens"),
("cache_write_tokens", "session_cache_write_tokens"))
# Tool-progress event -> SSE payload fields (tool_name, preview, kwargs); key order is wire format.
_FIXED_EVENT_FIELDS = {
"tool.started": lambda tool, preview, kw: {"tool": tool, "preview": preview},
@@ -625,8 +627,30 @@ async def _handle_runs(self, request: "web.Request", *, _api_server) -> "web.Res
return _accepted_response(run_id, "started", gateway_session_key, replayed=False)
def _run_usage(agent) -> Dict[str, int]:
"""Terminal ``usage`` payload from the agent's session counters; a missing or non-numeric
counter (test doubles, agents without cache accounting) reads as ``0``."""
usage = {}
for key, attr in _USAGE_FIELDS:
value = getattr(agent, attr, 0)
usage[key] = int(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else 0
return usage
def _served_runtime(agent) -> Dict[str, str]:
"""The ``{provider, model}`` pair that actually served the turn. After a ``fallback_providers``
switch the agent keeps the fallback runtime until the NEXT turn restores the primary, so when
``run_conversation()`` returns these attributes name the served pair — the run record's
``model`` field only echoes the request (#102101). Non-string attributes read as ``""``."""
pair = {}
for key in ("provider", "model"):
value = getattr(agent, key, "")
pair[key] = value if isinstance(value, str) else ""
return pair
def _run_agent_sync(self, run: _RunLaunch, agent, approval_notify, *, _api_server):
"""Executor-thread body of one run; returns ``(result, usage)``."""
"""Executor-thread body of one run; returns ``(result, usage, served_runtime)``."""
from gateway.session_context import clear_session_vars
from gateway.hosted_room_execution_policy import (
RoomExecutionPolicy, bind_room_execution_policy, reset_room_execution_policy)
@@ -690,7 +714,7 @@ def _run_agent_sync(self, run: _RunLaunch, agent, approval_notify, *, _api_serve
for token, reset in resets:
with suppress(Exception):
reset(token)
return r, {key: getattr(agent, attr, 0) or 0 for key, attr in _USAGE_FIELDS}
return r, _run_usage(agent), _served_runtime(agent)
def _make_approval_notify(self, run: _RunLaunch, *, _api_server) -> Callable[[Dict[str, Any]], None]:
@@ -755,7 +779,7 @@ async def _execute_run(self, run: _RunLaunch, *, _api_server) -> None:
**run.agent_kwargs)
self._active_run_agents[run_id] = agent
approval_notify = _make_approval_notify(self, run, _api_server=_api_server)
result, usage = await loop.run_in_executor(
result, usage, served_runtime = await loop.run_in_executor(
None, lambda: _run_agent_sync(self, run, agent, approval_notify, _api_server=_api_server))
if not isinstance(result, dict):
result = {}
@@ -766,7 +790,8 @@ async def _execute_run(self, run: _RunLaunch, *, _api_server) -> None:
# Non-retryable client errors (401/400) return failed=True rather than raising.
_finish("failed", fields, error=_redact_api_error_text(result.get("error") or "agent run failed"))
else:
_finish(status, fields, output=result.get("final_response", ""), usage=usage)
# ``runtime`` rides on both the pollable status and the run.completed event via _finish.
_finish(status, fields, output=result.get("final_response", ""), usage=usage, runtime=served_runtime)
except asyncio.CancelledError:
_finish("cancelled")
raise

View File

@@ -408,6 +408,47 @@ class TestRunStatus:
assert mock_agent.run_conversation.call_args.kwargs["task_id"] == "space-session"
assert status["session_id"] == "space-session"
@pytest.mark.asyncio
async def test_status_completed_run_reports_served_runtime_and_cache_tokens(self, adapter):
"""After a fallback_providers switch the run record carries the runtime that actually
served the turn plus cache-read tokens, next to the requested ``model`` (#102101).
``agent.provider`` / ``agent.model`` still hold the fallback pair when
``run_conversation()`` returns: the primary is only restored at the start of the NEXT
turn, so they are the served pair, while the top-level ``model`` echoes the request.
"""
app = _create_runs_app(adapter)
async with TestClient(TestServer(app)) as cli:
with patch.object(adapter, "_create_agent") as mock_create:
mock_agent = MagicMock()
mock_agent.run_conversation.return_value = {"final_response": "done"}
mock_agent.provider = "openai-codex"
mock_agent.model = "gpt-5.6-luna"
mock_agent.session_prompt_tokens = 100
mock_agent.session_completion_tokens = 5
mock_agent.session_total_tokens = 105
mock_agent.session_cache_read_tokens = 84
mock_agent.session_cache_write_tokens = 11
mock_create.return_value = mock_agent
resp = await cli.post("/v1/runs", json={"input": "hello", "model": "deepseek-v4-pro"})
run_id = (await resp.json())["run_id"]
for _ in range(40):
status = await (await cli.get(f"/v1/runs/{run_id}")).json()
if status["status"] == "completed":
break
await asyncio.sleep(0.05)
assert status["status"] == "completed"
# Top-level model still echoes the request; the served pair is disclosed alongside.
assert status["model"] == "deepseek-v4-pro"
assert status["runtime"] == {"provider": "openai-codex", "model": "gpt-5.6-luna"}
assert status["usage"] == {
"input_tokens": 100, "output_tokens": 5, "total_tokens": 105,
"cache_read_tokens": 84, "cache_write_tokens": 11,
}
# ---------------------------------------------------------------------------
# GET /v1/runs/{run_id}/events — SSE event stream
@@ -466,6 +507,48 @@ class TestRunEvents:
assert "run.completed" in body
assert "Hello!" in body
@pytest.mark.asyncio
async def test_completed_event_carries_served_runtime_and_cache_tokens(self, adapter):
"""The run.completed SSE event discloses the same served runtime and cache tokens as the
pollable status, so streaming clients get identical cost-attribution data (#102101)."""
import json as _json
app = _create_runs_app(adapter)
async with TestClient(TestServer(app)) as cli:
with patch.object(adapter, "_create_agent") as mock_create:
mock_agent = MagicMock()
mock_agent.run_conversation.return_value = {"final_response": "served"}
mock_agent.provider = "openai-codex"
mock_agent.model = "gpt-5.6-luna"
mock_agent.session_prompt_tokens = 774050
mock_agent.session_completion_tokens = 6286
mock_agent.session_total_tokens = 780336
mock_agent.session_cache_read_tokens = 650000
mock_agent.session_cache_write_tokens = 42
mock_create.return_value = mock_agent
resp = await cli.post("/v1/runs", json={"input": "hello", "model": "deepseek-v4-pro"})
run_id = (await resp.json())["run_id"]
events_resp = await cli.get(f"/v1/runs/{run_id}/events")
assert events_resp.status == 200
body = await events_resp.text()
completed = None
for frame in body.split("\n"):
if frame.startswith("data: "):
try:
payload = _json.loads(frame[len("data: "):])
except ValueError:
continue
if payload.get("event") == "run.completed":
completed = payload
break
assert completed is not None, "run.completed event missing from stream"
assert completed["runtime"] == {"provider": "openai-codex", "model": "gpt-5.6-luna"}
assert completed["usage"]["cache_read_tokens"] == 650000
assert completed["usage"]["cache_write_tokens"] == 42
@pytest.mark.asyncio
async def test_approval_resolve_all_is_scoped_to_target_run(self, auth_adapter):

View File

@@ -473,10 +473,13 @@ Poll the current run state. This is useful for dashboards that need status witho
"session_id": "space-session",
"model": "hermes-agent",
"output": "Done.",
"usage": {"input_tokens": 50, "output_tokens": 200, "total_tokens": 250}
"usage": {"input_tokens": 50, "output_tokens": 200, "total_tokens": 250, "cache_read_tokens": 40, "cache_write_tokens": 0},
"runtime": {"provider": "openai", "model": "gpt-5"}
}
```
`model` echoes what the request asked for. On a completed run, `runtime` is the provider/model pair that actually served the turn — after a [fallback provider](fallback-providers.md) switch it names the fallback pair, so a cost-attribution poller books the run to the right provider. `usage.cache_read_tokens` / `usage.cache_write_tokens` are the session's prompt-cache reads and writes, so cached input is not priced as full-price input. The `run.completed` event on the events stream carries the same `usage` and `runtime` fields.
Statuses are retained briefly after terminal states (`completed`, `failed`, `cancelled`, or `interrupted`) for polling and UI reconciliation. When the gateway shuts down while a run is active, the run is persisted as `interrupted` (error `Gateway shutdown interrupted the run.`, terminal event `run.interrupted`) before the agent is asked to stop, so a durable run never survives a restart as `running`; a late result from the interrupted turn cannot overwrite it.
### GET /v1/runs/\{run_id\}/events