From aa258b983d3acbadababfb5db7e90738df8fae67 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 24 May 2026 17:11:42 +0000
Subject: [PATCH] Studio: account for every Codex turn in fan-out usage chunk
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 .../backend/core/inference/codex_provider.py  | 16 +++++++-
 studio/backend/tests/test_codex_provider.py   | 40 +++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/studio/backend/core/inference/codex_provider.py b/studio/backend/core/inference/codex_provider.py
index db790f623e..ca3dad0e37 100644
--- a/studio/backend/core/inference/codex_provider.py
+++ b/studio/backend/core/inference/codex_provider.py
@@ -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)
 
diff --git a/studio/backend/tests/test_codex_provider.py b/studio/backend/tests/test_codex_provider.py
index 2b52bc5f94..27f40c7e77 100644
--- a/studio/backend/tests/test_codex_provider.py
+++ b/studio/backend/tests/test_codex_provider.py
@@ -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.