Skip to content

Commit 957edf5

Browse files
authored
fix(codex): exclude compaction estimates from model call count (#252)
* fix(codex): exclude compaction estimates from model call count * fix(codex): harden compaction usage accounting regressions
1 parent 4779486 commit 957edf5

4 files changed

Lines changed: 185 additions & 14 deletions

File tree

‎.claude/notes/agents.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,21 @@ intentionally brief and out of scope; trimming for DISPLAY belongs in the render
8080
table's to state, not this file's. Full table + rationale:
8181
docs/agents/HARNESS_PARITY.md.
8282

83+
## Codex compaction usage metadata
84+
85+
Codex emits usage for the compaction model response, then emits another usage update
86+
after estimating the replacement context. The estimate has zero input, output, cache
87+
and reasoning deltas. Within a compaction lifecycle, that later bookkeeping update
88+
still updates usage and transcript capture but does not begin another model call.
89+
The first compaction usage update counts even when its reported deltas are zero;
90+
completion clears the lifecycle before ordinary model and tool calls resume.
91+
92+
The source contract is visible in the pinned CLI's `openai/codex` tag `rust-v0.156.1`:
93+
`compact.rs` calls `update_token_usage_info` and then `recompute_token_usage`, each of
94+
which emits a `TokenCount` event. The app-server forwards both as
95+
`thread/tokenUsage/updated`. An ordinary tool call's next-call reservation remains
96+
unchanged. An estimate outside compaction keeps the existing accounting behavior.
97+
8398
## Shared turn lifecycle
8499

85100
Every adapter drives the same skeleton, on the base class: `_begin_turn()` resets the

‎docs/agents/HARNESS_PARITY.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -553,8 +553,11 @@ Each harness finds the call boundary in its own stream (the table above):
553553
`error_max_turns` stop. The CLI does not apply it on every route, so the harness
554554
also counts distinct main-thread `message_id`s and ends the turn itself when call
555555
N+1 begins.
556-
- **Codex** sends `thread/tokenUsage/updated` once per call, after that call's
557-
tools finish. An item starts as its tool runs, so waiting for the next call's
556+
- **Codex** sends `thread/tokenUsage/updated` after a call's tools finish. Context
557+
compaction also sends a zero-delta context-size estimate; that bookkeeping
558+
update does not count as another call. The compaction model call itself counts,
559+
unlike Claude Code's compaction, which emits no main-thread assistant message.
560+
An item starts as its tool runs, so waiting for the next call's
558561
first item would let one tool of call N+1 act. A call that ran tools therefore
559562
opens the next call at its `tokenUsage`, since the results always go back to the
560563
model, and the cap fires before call N+1 can run anything.

‎src/coder_eval/agents/codex_agent.py‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -557,6 +557,8 @@ def __init__(
557557
self.api_calls = 0
558558
self.in_api_call = False
559559
self.call_ran_tools = False
560+
self.in_context_compaction = False
561+
self.compaction_usage_seen = False
560562

561563
# Finalize inputs, COMMITTED by communicate after a clean pump return.
562564
# Defaults are the crash values (no terminal usage; format from messages).
@@ -755,6 +757,9 @@ def on_item_started(self, notification: Any) -> None:
755757
if item_id is not None and started_at_ms is not None:
756758
self.start_ms_by_id[item_id] = started_at_ms
757759
root_type = getattr(root, "type", None)
760+
if root_type == "contextCompaction":
761+
self.in_context_compaction = True
762+
self.compaction_usage_seen = False
758763
if not self.in_api_call and root_type not in _NON_MODEL_ITEM_TYPES:
759764
self._open_api_call()
760765
# Any item that isn't transcript content is a tool call (generic capture).
@@ -790,6 +795,8 @@ def on_item_completed(self, notification: Any) -> None:
790795
return
791796
completed_ms = getattr(notification.payload, "completed_at_ms", None)
792797
root_type = getattr(root, "type", None)
798+
if root_type == "contextCompaction":
799+
self.in_context_compaction = False
793800
if root_type is not None and root_type not in _CONTENT_ITEM_TYPES:
794801
tool_id = getattr(root, "id", None) or f"{root_type}_{self.next_sequence}"
795802
seq = self.seq_by_id.get(tool_id, self.next_sequence)
@@ -875,18 +882,39 @@ def on_agent_message_delta(self, notification: Any) -> None:
875882
self.emit.on_event(TextChunkEvent(task_id=self.task_id, turn_id=self.turn_id, text=delta))
876883

877884
def on_token_usage_updated(self, notification: Any) -> None:
878-
"""One per generation → cut a message. Carries `total` (cumulative over the
879-
whole THREAD, i.e. every turn so far) and `last` (this generation's delta)."""
880-
if not self.in_api_call:
881-
self.api_calls += 1
882-
self.in_api_call = False
883-
# Tool results always go back to the model, so the next call begins here,
884-
# before Codex can run the next call's tools (an item starts as its tool runs).
885-
if self.call_ran_tools:
886-
self._open_api_call()
885+
"""Capture generation usage and compaction's context-size estimates.
886+
887+
Rationale: .claude/notes/agents.md § Codex compaction usage metadata
888+
"""
889+
usage = getattr(notification.payload, "token_usage", None) if notification.payload else None
890+
last = getattr(usage, "last", None)
891+
estimate = (
892+
self.in_context_compaction
893+
and self.compaction_usage_seen
894+
and last is not None
895+
and all(
896+
(getattr(last, key, 0) or 0) == 0
897+
for key in (
898+
"input_tokens",
899+
"output_tokens",
900+
"cached_input_tokens",
901+
"cache_write_input_tokens",
902+
"reasoning_output_tokens",
903+
)
904+
)
905+
)
906+
if self.in_context_compaction:
907+
self.compaction_usage_seen = True
908+
if not estimate:
909+
if not self.in_api_call:
910+
self.api_calls += 1
911+
self.in_api_call = False
912+
# Tool results reserve the next model call before its tools can run.
913+
if self.call_ran_tools:
914+
self._open_api_call()
887915
if notification.payload:
888-
self.latest_token_usage = getattr(notification.payload, "token_usage", None)
889-
self._flush_message(getattr(self.latest_token_usage, "last", None))
916+
self.latest_token_usage = usage
917+
self._flush_message(last)
890918

891919
def on_turn_completed(self, notification: Any) -> bool:
892920
"""Capture the final Turn. Returns True (break the pump) iff the payload is

‎tests/test_codex_agent.py‎

Lines changed: 126 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
_CODEX_APPROVAL_MODE,
2020
CodexAgent,
2121
)
22-
from coder_eval.models import AgentConfig, AgentKind, parse_agent_config
22+
from coder_eval.models import AgentConfig, AgentKind, AssistantMessage, ReconciliationMessage, parse_agent_config
2323

2424

2525
class TestCodexAgentInitialization:
@@ -504,6 +504,7 @@ def test_get_state_returns_current_state():
504504

505505
from openai_codex.generated.v2_all import ( # noqa: E402
506506
SubAgentActivityThreadItem,
507+
TokenUsageBreakdown,
507508
Turn,
508509
TurnCompletedNotification,
509510
)
@@ -2229,6 +2230,130 @@ def _cmd_notifications(count: int) -> list:
22292230
notifications.append(_turn_completed())
22302231
return notifications
22312232

2233+
@staticmethod
2234+
def _compaction_notifications(
2235+
*,
2236+
estimate: bool = True,
2237+
empty_response: bool = False,
2238+
item_id: str = "compact",
2239+
cache_write: int | None = 0,
2240+
) -> list:
2241+
item = SimpleNamespace(type="contextCompaction", id=item_id)
2242+
response = _token_usage(inp=10, out=5, cached=0)
2243+
response.payload.token_usage.last = TokenUsageBreakdown(
2244+
input_tokens=0 if empty_response else 10,
2245+
output_tokens=0 if empty_response else 5,
2246+
cached_input_tokens=0,
2247+
cache_write_input_tokens=0,
2248+
reasoning_output_tokens=0,
2249+
total_tokens=0 if empty_response else 15,
2250+
)
2251+
notifications = [_item_notification("item/started", item), response]
2252+
if estimate:
2253+
recomputed = _token_usage(inp=20, out=5, cached=0)
2254+
recomputed.payload.token_usage.last = TokenUsageBreakdown(
2255+
input_tokens=0,
2256+
output_tokens=0,
2257+
cached_input_tokens=0,
2258+
cache_write_input_tokens=cache_write,
2259+
reasoning_output_tokens=0,
2260+
total_tokens=20,
2261+
)
2262+
notifications.append(recomputed)
2263+
notifications.append(_item_notification("item/completed", item))
2264+
return notifications
2265+
2266+
@pytest.mark.parametrize(
2267+
"estimate,empty_response,cache_write",
2268+
[(True, False, 0), (False, False, 0), (True, True, 0), (True, False, None)],
2269+
)
2270+
async def test_compaction_estimate_does_not_spend_model_call_budget(self, estimate, empty_response, cache_write):
2271+
notifications = self._cmd_notifications(1)
2272+
notifications[5:5] = self._compaction_notifications(
2273+
estimate=estimate, empty_response=empty_response, cache_write=cache_write
2274+
)
2275+
agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), notifications)
2276+
2277+
record = await agent.communicate("go", max_turns=3)
2278+
2279+
assert record.num_turns == 3
2280+
assert record.max_turns_exhausted is False
2281+
assert len(record.commands) == 1
2282+
2283+
async def test_repeated_compaction_resets_before_normal_tool_calls(self):
2284+
notifications = self._cmd_notifications(2)
2285+
notifications[10:10] = self._compaction_notifications(item_id="compact-2", empty_response=True)
2286+
notifications[5:5] = self._compaction_notifications(item_id="compact-1")
2287+
agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), notifications)
2288+
2289+
record = await agent.communicate("go", max_turns=None)
2290+
2291+
assert record.num_turns == 5
2292+
assert record.max_turns_exhausted is False
2293+
assert len(record.commands) == 2
2294+
2295+
async def test_zero_usage_outside_compaction_still_counts(self):
2296+
update = _token_usage(inp=10, out=5, cached=0)
2297+
update.payload.token_usage.last = TokenUsageBreakdown(
2298+
input_tokens=0,
2299+
output_tokens=0,
2300+
cached_input_tokens=0,
2301+
cache_write_input_tokens=0,
2302+
reasoning_output_tokens=0,
2303+
total_tokens=20,
2304+
)
2305+
notifications = [_token_usage(inp=10, out=5, cached=0), update, _turn_completed()]
2306+
agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), notifications)
2307+
2308+
record = await agent.communicate("go", max_turns=None)
2309+
2310+
assert record.num_turns == 2
2311+
assert record.max_turns_exhausted is False
2312+
2313+
async def test_final_compaction_estimate_updates_usage_and_message_buckets(self):
2314+
reply = SimpleNamespace(type="agentMessage", id="preface", text="Preparing a shorter context.")
2315+
notifications = [
2316+
_item_notification("item/started", reply),
2317+
_item_notification("item/completed", reply),
2318+
*self._compaction_notifications(),
2319+
_turn_completed(),
2320+
]
2321+
agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), notifications)
2322+
2323+
record = await agent.communicate("go", max_turns=None)
2324+
2325+
assert record.token_usage is not None
2326+
assert record.token_usage.uncached_input_tokens == 20
2327+
assert record.token_usage.output_tokens == 5
2328+
buckets = [
2329+
message for message in record.messages if isinstance(message, (AssistantMessage, ReconciliationMessage))
2330+
]
2331+
assert sum(message.input_tokens for message in buckets) == record.token_usage.uncached_input_tokens
2332+
assert sum(message.output_tokens for message in buckets) == record.token_usage.output_tokens
2333+
assert sum(message.cache_read_tokens for message in buckets) == record.token_usage.cache_read_input_tokens
2334+
assert (
2335+
sum(message.cache_creation_tokens for message in buckets) == record.token_usage.cache_creation_input_tokens
2336+
)
2337+
2338+
async def test_zero_usage_after_compaction_completion_counts_next_model_call(self):
2339+
update = _token_usage(inp=0, out=0, cached=0)
2340+
update.payload.token_usage.last = SimpleNamespace(
2341+
input_tokens=0,
2342+
output_tokens=0,
2343+
cached_input_tokens=0,
2344+
cache_write_input_tokens=0,
2345+
reasoning_output_tokens=0,
2346+
total_tokens=0,
2347+
)
2348+
notifications = self._compaction_notifications()
2349+
notifications.extend([update, _turn_completed()])
2350+
agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), notifications)
2351+
2352+
record = await agent.communicate("go", max_turns=2)
2353+
2354+
assert record.num_turns == 2
2355+
assert record.max_turns_exhausted is False
2356+
22322357
async def test_cap_stops_before_the_next_call_runs_a_tool(self):
22332358
agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), self._cmd_notifications(5))
22342359

0 commit comments

Comments
 (0)