Studio: account for every Codex turn in fan-out usage chunk

The fan-out path used to report usage as if a single Codex call had
run: `prompt_tokens = max(1, len(prompt)//4)` and
`completion_tokens = len(synthesis)//4`. In reality it had spawned
N parallel worker turns (each carrying the same prompt) plus a
synthesis turn that re-sent the prompt and every tab's output. For
`parallel_calls=20` that meant the cost / context widget under-
reported the request by roughly 20x.

Now sums:

- `prompt_tokens ≈ (N * prompt + synthesis_prompt) / 4` where
  `synthesis_prompt = sum(tab_outputs) + prompt`.
- `completion_tokens ≈ (sum(tab_output_chars) + synthesis_chars) / 4`.

Tests: new regression
`test_parallel_usage_accounts_for_all_calls` runs a 4-way fan-out
against a fake SDK with deterministic chunk lengths and asserts
the reported usage scales with N, not the single-call shape.
This commit is contained in:
Daniel Han 2026-05-24 17:11:42 +00:00
commit aa258b983d
2 changed files with 54 additions and 2 deletions

View file

@ -1001,10 +1001,22 @@ async def _stream_codex_parallel(
if synthesis_text:
yield _chunk_text(completion_id, synthesis_text)
# Account for ALL Codex turns the fan-out spawned: N parallel
# workers each ran the same prompt (≈ N * prompt_tokens), and the
# synthesis turn re-sent the prompt plus every tab's output. Without
# this the cost / context display is off by the fan-out factor and
# users see a wildly inaccurate token count for the request.
total_tab_completion_chars = sum(len(t) for t in per_tab_texts)
synthesis_prompt_chars = sum(len(t) for t in per_tab_texts) + len(prompt)
yield _chunk_usage(
completion_id,
prompt_tokens = max(1, len(prompt) // 4),
completion_tokens = max(0, len(synthesis_text) // 4),
# n worker prompts (same prompt each) + synthesis prompt (which
# carries the prompt again plus every tab's output).
prompt_tokens = max(1, (n * len(prompt) + synthesis_prompt_chars) // 4),
# Sum of every worker's output plus the synthesis text.
completion_tokens = max(
0, (total_tab_completion_chars + len(synthesis_text)) // 4
),
)
yield _chunk_stop(completion_id)

View file

@ -1424,6 +1424,46 @@ class TestCodexHardenedRegressions:
assert url_events and url_events[0]["url"].endswith("/codex/device")
assert code_events and code_events[0]["code"] == "ABCD-EFGH"
def test_parallel_usage_accounts_for_all_calls(self, monkeypatch):
"""The fan-out path runs N worker calls + 1 synthesis call.
The reported usage must reflect that, not just one call's
worth, otherwise the cost / context display is off by the
fan-out factor.
"""
_install_fake_codex_sdk(
monkeypatch,
lambda: _FakeAsyncCodex(
chunks = ["AAAAAAAAAA"], # 10 chars per tab
final = "SYNTHESISED" * 10, # 110 chars synthesis
),
)
from core.inference.codex_provider import stream_codex
n = 4
long_prompt = "a" * 200 # 200 chars
lines = _collect_stream(
stream_codex(
messages = [{"role": "user", "content": long_prompt}],
model = "gpt-5.4",
parallel_calls = n,
)
)
chunks = _parse_sse_chunks(lines)
usage_chunks = [c for c in chunks if c.get("choices") == [] and c.get("usage")]
assert len(usage_chunks) == 1
usage = usage_chunks[0]["usage"]
# Single-call prompt would be ~200/4 = 50 tokens. For n=4 with
# synthesis, prompt should be much larger: n*200 + (n*10 + 200)
# = 800 + 240 = 1040 chars ~= 260 tokens.
assert (
usage["prompt_tokens"] >= 200
), f"prompt_tokens not scaled for fan-out: {usage['prompt_tokens']}"
# Completion = n*10 (tab outputs) + 110 (synthesis) = 150 chars
# ~= 37 tokens. Definitely > the synthesis-only count of 27.
assert (
usage["completion_tokens"] >= 30
), f"completion_tokens not scaled for fan-out: {usage['completion_tokens']}"
def test_buffered_result_none_final_does_not_emit_repr(self, monkeypatch):
"""A buffered TurnResult whose final_response is None must NOT
send a Python object repr (``TurnResult(...)``) to the user.