fix(codex): report per-turn tokens instead of the thread-cumulative total

src/coder_eval/agents/codex_agent.py

@@ -12,7 +12,7 @@
 from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
-from typing import Any, ClassVar
+from typing import Any, ClassVar, NamedTuple
 from urllib.parse import urlparse
 
 from coder_eval.agent import Agent, AgentState
@@ -195,6 +195,34 @@ def _fresh_input_tokens(raw_input: int, cached: int) -> int:
     return max(raw_input - cached, 0)
 
 
+class _ThreadTotals(NamedTuple):
+    """A snapshot of the Codex SDK's thread-cumulative ``ThreadTokenUsage.total``.
+
+    Held on the agent across turns (the thread outlives the turn) so each turn can
+    report its own slice instead of the running total. ``input`` is the full prompt
+    count, cached prefix included — the SDK's convention, not ours.
+    """
+
+    input: int = 0
+    output: int = 0
+    cached: int = 0
+
+    def since(self, baseline: "_ThreadTotals") -> "_ThreadTotals":
+        """This turn's tokens = the cumulative snapshot minus the previous one.
+
+        A total that moved BACKWARDS means the thread restarted under us (a fresh
+        thread counts from zero), so the snapshot is already turn-local: return it
+        whole rather than clamping every bucket to zero and losing the turn.
+        """
+        if self.input < baseline.input or self.output < baseline.output or self.cached < baseline.cached:
+            return self
+        return _ThreadTotals(
+            input=self.input - baseline.input,
+            output=self.output - baseline.output,
+            cached=self.cached - baseline.cached,
+        )
+
+
 def _message_uncached_input(m: AssistantMessage) -> int:
     """A captured generation's fresh (uncached) input.
 
@@ -569,8 +597,8 @@ def on_agent_message_delta(self, notification: Any) -> None:
                 self.emit.on_event(TextChunkEvent(task_id=self.task_id, turn_id=self.turn_id, text=delta))
 
     def on_token_usage_updated(self, notification: Any) -> None:
-        """One per generation → cut a message. Carries `total` (cumulative turn
-        figure) and `last` (this generation's delta)."""
+        """One per generation → cut a message. Carries `total` (cumulative over the
+        whole THREAD, i.e. every turn so far) and `last` (this generation's delta)."""
         if notification.payload:
             self.latest_token_usage = getattr(notification.payload, "token_usage", None)
             self._flush_message(getattr(self.latest_token_usage, "last", None))
@@ -608,13 +636,19 @@ def finalize(self, status: AgentEndStatus, *, crashed: bool = False, crash_reaso
             return
         self.finalized = True
 
-        # Prefer the SDK total; on crash/timeout it stays None, so fall back to the
-        # per-generation tokens already captured on the messages.
-        token_usage = self._agent._token_usage_from_sdk(self.sdk_token_usage) or self._agent._token_usage_from_messages(
-            self.messages
-        )
+        # Prefer the SDK total (deltas off the thread-cumulative figure); on
+        # crash/timeout it stays None, so fall back to the per-generation tokens
+        # already captured on the messages — those are per-turn to begin with, but
+        # the thread baseline still has to move past them or the NEXT turn's delta
+        # re-books this one.
+        token_usage = self._agent._token_usage_from_sdk(self.sdk_token_usage)
+        if token_usage is None:
+            token_usage = self._agent._token_usage_from_messages(self.messages)
+            self._agent._advance_usage_baseline(token_usage)
         # Codex bills sub-agents on separate threads, so fold the recovered child
         # generations into the turn total — matching Claude's bubbled-up totals.
+        # Folded AFTER the baseline advance: the SDK total covers the parent thread
+        # only, so child tokens must not shift the parent's baseline.
         token_usage = self._agent._fold_subagent_tokens(token_usage, self.messages)
 
         self.emit.on_event(
@@ -688,6 +722,10 @@ def __init__(
         self.route = route or DirectRoute()
         self.codex_client: Any = None
         self.thread: Any = None
+        # Thread-cumulative token snapshot as of the END of the last finalized turn.
+        # The thread outlives the turn, so this is what makes each turn's usage its
+        # own delta rather than the running total (see _token_usage_from_sdk).
+        self._thread_usage_baseline = _ThreadTotals()
         self.working_directory: Path | None = None
         self._env_path_prepend: list[str] = []
         self._login_shell_home: Path | None = None
@@ -845,6 +883,8 @@ async def communicate(
                 if self.working_directory:
                     thread_kwargs["cwd"] = str(self.working_directory)
                 self.thread = await self._run_async(self.codex_client.thread_start, **thread_kwargs)
+                # A fresh thread counts its cumulative total from zero.
+                self._thread_usage_baseline = _ThreadTotals()
 
             def _on_turn_timeout() -> None:
                 state.timeout_hit = True
@@ -944,6 +984,7 @@ async def stop(self) -> None:
         """
         self._close_client()
         self.thread = None
+        self._thread_usage_baseline = _ThreadTotals()
         self._active_turn_handle = None
         self._cleanup_login_shell_home()
         self._mark_stopped()
@@ -2205,13 +2246,23 @@ def _extract_file_change_telemetry(
             return None
 
     def _token_usage_from_sdk(self, sdk_token_usage: Any) -> TokenUsage | None:
-        """Convert the Codex SDK's ThreadTokenUsage to our TokenUsage.
+        """This turn's own slice of the Codex SDK's thread-cumulative total.
 
         Single conversion site for both the TurnEndEvent and the AgentEndEvent,
         so cached-input tokens can't be captured in one path but dropped in the
         other. The Codex SDK does not surface cost, so we derive it from the
         pricing table keyed on the effective model (None if the model is unpriced).
 
+        ``ThreadTokenUsage.total`` counts the whole THREAD, not the turn — that is
+        the SDK's contract, and ``last`` is the per-generation delta beside it. The
+        Codex thread is created once per task and reused for every turn (see
+        ``communicate``), so by turn N ``total`` still carries turns 1..N-1. The
+        orchestrator sums per-turn usages into the task total, so handing it the
+        cumulative figure books turn 1 again on turn 2, turns 1-2 again on turn 3,
+        and so on: the task total becomes a sum of prefix sums, inflating an
+        N-turn task by roughly (N+1)/2. Subtracting the baseline captured at the
+        end of the previous turn leaves just this turn.
+
         Cache-bucket convention (Codex/OpenAI): the SDK's ``input_tokens`` is the
         FULL prompt count, *inclusive* of the cached prefix. The fresh slice
         (``input_tokens - cached``) is the uncached input (OpenAI bills no separate
@@ -2231,24 +2282,48 @@ def _token_usage_from_sdk(self, sdk_token_usage: Any) -> TokenUsage | None:
         total = getattr(sdk_token_usage, "total", None)
         if not total:
             return None
-        input_tokens = getattr(total, "input_tokens", 0) or 0
-        output_tokens = getattr(total, "output_tokens", 0) or 0
-        cached_input = getattr(total, "cached_input_tokens", 0) or 0
+        cumulative = _ThreadTotals(
+            input=getattr(total, "input_tokens", 0) or 0,
+            output=getattr(total, "output_tokens", 0) or 0,
+            cached=getattr(total, "cached_input_tokens", 0) or 0,
+        )
+        turn = cumulative.since(self._thread_usage_baseline)
+        self._thread_usage_baseline = cumulative
         # Fresh (uncached) prompt slice = full prompt minus the cached prefix.
-        uncached = _fresh_input_tokens(input_tokens, cached_input)
+        uncached = _fresh_input_tokens(turn.input, turn.cached)
         cost = calculate_cost(
             self._effective_model() or "",
             uncached_input_tokens=uncached,
-            output_tokens=output_tokens,
-            cache_read_tokens=cached_input,
+            output_tokens=turn.output,
+            cache_read_tokens=turn.cached,
         )
         return TokenUsage(
             uncached_input_tokens=uncached,
-            output_tokens=output_tokens,
-            cache_read_input_tokens=cached_input,
+            output_tokens=turn.output,
+            cache_read_input_tokens=turn.cached,
             total_cost_usd=cost,
         )
 
+    def _advance_usage_baseline(self, usage: TokenUsage | None) -> None:
+        """Move the thread baseline past a turn whose SDK total never arrived.
+
+        The crash/timeout fallback (``_token_usage_from_messages``) reads
+        per-generation tokens straight off the flushed messages, so the crashed
+        turn itself is right — but the thread's cumulative total kept climbing on
+        the SDK side. Without advancing past it here, the next turn's delta would
+        re-book everything the crashed turn already reported.
+        """
+        if usage is None:
+            return
+        base = self._thread_usage_baseline
+        # SDK ``input_tokens`` is the full prompt, cached prefix included, so the
+        # input baseline advances by uncached + cache_read.
+        self._thread_usage_baseline = _ThreadTotals(
+            input=base.input + usage.uncached_input_tokens + usage.cache_read_input_tokens,
+            output=base.output + usage.output_tokens,
+            cached=base.cached + usage.cache_read_input_tokens,
+        )
+
     def _fold_subagent_tokens(self, parent: TokenUsage | None, messages: list[TranscriptMessage]) -> TokenUsage | None:
         """Add recovered sub-agent (child-thread) tokens to the parent turn total.
 

tests/test_codex_token_mapping.py

@@ -11,9 +11,11 @@
 
 from __future__ import annotations
 
+from types import SimpleNamespace
+
 import pytest
 
-from coder_eval.agents.codex_agent import _fresh_input_tokens
+from coder_eval.agents.codex_agent import _fresh_input_tokens, _ThreadTotals
 from coder_eval.models import TokenUsage
 
 
@@ -58,3 +60,128 @@ def test_fresh_is_uncached_and_cache_creation_is_zero(self):
         assert tu.input_tokens == 1000
         # Cost bills the uncached slice, never the derived total.
         assert tu.uncached_input_tokens < tu.input_tokens
+
+
+def _sdk_usage(input_tokens: int, output_tokens: int, cached: int):
+    """A stand-in for the SDK's ThreadTokenUsage — only ``.total`` is read."""
+    return SimpleNamespace(
+        total=SimpleNamespace(
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cached_input_tokens=cached,
+        )
+    )
+
+
+def _agent():
+    from coder_eval.agents.codex_agent import CodexAgent
+    from coder_eval.models import CodexAgentConfig
+
+    return CodexAgent(CodexAgentConfig(type="codex", model="gpt-5.6-terra"))
+
+
+class TestThreadTotals:
+    """``since`` turns the thread-cumulative snapshot into a per-turn delta."""
+
+    def test_delta_against_baseline(self):
+        assert _ThreadTotals(1000, 100, 400).since(_ThreadTotals(600, 40, 250)) == _ThreadTotals(400, 60, 150)
+
+    def test_first_turn_is_the_whole_snapshot(self):
+        assert _ThreadTotals(1000, 100, 400).since(_ThreadTotals()) == _ThreadTotals(1000, 100, 400)
+
+    def test_backwards_total_means_a_restarted_thread(self):
+        # A fresh thread counts from zero, so the snapshot is already turn-local.
+        # Returning it whole beats clamping every bucket to 0 and losing the turn.
+        restarted = _ThreadTotals(300, 20, 100)
+        assert restarted.since(_ThreadTotals(9000, 800, 5000)) == restarted
+
+
+class TestPerTurnUsageFromCumulativeSdkTotal:
+    """The Codex thread is created once per task and reused for every turn, so the
+    SDK's ``total`` keeps climbing. Each turn must report only its own slice —
+    ``orchestrator._aggregate_token_usage`` SUMS the per-turn usages, so a
+    cumulative figure makes the task total a sum of prefix sums."""
+
+    def test_each_turn_reports_its_own_delta(self):
+        agent = _agent()
+        # Thread-cumulative totals as the SDK reports them at the end of turns 1-3.
+        turn1 = agent._token_usage_from_sdk(_sdk_usage(10_000, 500, 6_000))
+        turn2 = agent._token_usage_from_sdk(_sdk_usage(26_000, 1_300, 18_000))
+        turn3 = agent._token_usage_from_sdk(_sdk_usage(40_000, 2_000, 30_000))
+
+        assert turn1 is not None and turn2 is not None and turn3 is not None
+        # Turn 1: 10k prompt, 6k of it cached.
+        assert (turn1.uncached_input_tokens, turn1.output_tokens, turn1.cache_read_input_tokens) == (4_000, 500, 6_000)
+        # Turn 2 spent 16k input / 800 output / 12k cached — NOT the 26k running total.
+        assert (turn2.uncached_input_tokens, turn2.output_tokens, turn2.cache_read_input_tokens) == (4_000, 800, 12_000)
+        assert (turn3.uncached_input_tokens, turn3.output_tokens, turn3.cache_read_input_tokens) == (2_000, 700, 12_000)
+
+    def test_summed_turns_equal_the_final_cumulative_total(self):
+        # The invariant the orchestrator depends on: Σ(per-turn usage) over a task
+        # == the thread's final cumulative total. Reporting `total` per turn broke
+        # this, inflating an N-turn task by roughly (N+1)/2.
+        agent = _agent()
+        cumulative = [(10_000, 500, 6_000), (26_000, 1_300, 18_000), (40_000, 2_000, 30_000)]
+        turns = [agent._token_usage_from_sdk(_sdk_usage(*c)) for c in cumulative]
+
+        final_input, final_output, final_cached = cumulative[-1]
+        assert sum(t.output_tokens for t in turns if t) == final_output
+        assert sum(t.cache_read_input_tokens for t in turns if t) == final_cached
+        # uncached + cache_read reconstitutes the full prompt count.
+        assert sum(t.uncached_input_tokens + t.cache_read_input_tokens for t in turns if t) == final_input
+
+    def test_cost_is_per_turn_not_cumulative(self):
+        agent = _agent()
+        agent._token_usage_from_sdk(_sdk_usage(100_000, 5_000, 80_000))
+        second = agent._token_usage_from_sdk(_sdk_usage(101_000, 5_100, 80_500))
+        assert second is not None and second.total_cost_usd is not None
+        # Turn 2 was a 1k-prompt sliver (500 fresh + 500 cached, 100 output).
+        # Pricing the 101k running total instead would be roughly 100x this.
+        assert second.total_cost_usd < 0.01
+
+    def test_baseline_is_per_agent_not_global(self):
+        first, second = _agent(), _agent()
+        first._token_usage_from_sdk(_sdk_usage(50_000, 900, 30_000))
+        fresh = second._token_usage_from_sdk(_sdk_usage(8_000, 200, 5_000))
+        assert fresh is not None
+        assert fresh.uncached_input_tokens == 3_000
+
+    @pytest.mark.parametrize("sdk_usage", [None, SimpleNamespace(total=None)])
+    def test_absent_total_yields_none_and_leaves_baseline_untouched(self, sdk_usage):
+        agent = _agent()
+        agent._token_usage_from_sdk(_sdk_usage(10_000, 500, 6_000))
+        before = agent._thread_usage_baseline
+        assert agent._token_usage_from_sdk(sdk_usage) is None
+        assert agent._thread_usage_baseline == before
+
+
+class TestCrashFallbackBaseline:
+    """A crashed turn never yields an SDK total, so ``_token_usage_from_messages``
+    reports it. The baseline still has to advance past it, or the next turn's
+    delta re-books the crashed turn's tokens."""
+
+    def test_advance_moves_baseline_past_a_message_derived_turn(self):
+        agent = _agent()
+        agent._advance_usage_baseline(
+            TokenUsage(uncached_input_tokens=4_000, output_tokens=500, cache_read_input_tokens=6_000)
+        )
+        # SDK input is the full prompt, cached prefix included → 4k + 6k.
+        assert agent._thread_usage_baseline == _ThreadTotals(input=10_000, output=500, cached=6_000)
+
+    def test_turn_after_a_crash_is_not_inflated(self):
+        agent = _agent()
+        # Turn 1 completes normally: 10k prompt / 6k cached.
+        agent._token_usage_from_sdk(_sdk_usage(10_000, 500, 6_000))
+        # Turn 2 crashes; its tokens come off the flushed messages instead.
+        agent._advance_usage_baseline(
+            TokenUsage(uncached_input_tokens=2_000, output_tokens=300, cache_read_input_tokens=8_000)
+        )
+        # Turn 3's cumulative total includes turns 1 and 2; only turn 3 is new.
+        turn3 = agent._token_usage_from_sdk(_sdk_usage(32_000, 1_100, 22_000))
+        assert turn3 is not None
+        assert (turn3.uncached_input_tokens, turn3.output_tokens, turn3.cache_read_input_tokens) == (4_000, 300, 8_000)
+
+    def test_advance_is_a_noop_for_none(self):
+        agent = _agent()
+        agent._advance_usage_baseline(None)
+        assert agent._thread_usage_baseline == _ThreadTotals()