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:
parent
cb0680ebaa
commit
aa258b983d
2 changed files with 54 additions and 2 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue