|
19 | 19 | _CODEX_APPROVAL_MODE, |
20 | 20 | CodexAgent, |
21 | 21 | ) |
22 | | -from coder_eval.models import AgentConfig, AgentKind, parse_agent_config |
| 22 | +from coder_eval.models import AgentConfig, AgentKind, AssistantMessage, ReconciliationMessage, parse_agent_config |
23 | 23 |
|
24 | 24 |
|
25 | 25 | class TestCodexAgentInitialization: |
@@ -504,6 +504,7 @@ def test_get_state_returns_current_state(): |
504 | 504 |
|
505 | 505 | from openai_codex.generated.v2_all import ( # noqa: E402 |
506 | 506 | SubAgentActivityThreadItem, |
| 507 | + TokenUsageBreakdown, |
507 | 508 | Turn, |
508 | 509 | TurnCompletedNotification, |
509 | 510 | ) |
@@ -2229,6 +2230,130 @@ def _cmd_notifications(count: int) -> list: |
2229 | 2230 | notifications.append(_turn_completed()) |
2230 | 2231 | return notifications |
2231 | 2232 |
|
| 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 | + |
2232 | 2357 | async def test_cap_stops_before_the_next_call_runs_a_tool(self): |
2233 | 2358 | agent = _started_agent(parse_agent_config(type=AgentKind.CODEX), self._cmd_notifications(5)) |
2234 | 2359 |
|
|
0 commit comments