diff --git a/agent/context_engine.py b/agent/context_engine.py index 17480f578260..8cc9d512788e 100644 --- a/agent/context_engine.py +++ b/agent/context_engine.py @@ -289,6 +289,16 @@ class ContextEngine(ABC): ``task_id``, ``api_call_count``, ``interrupted``, ``failed``, and ``turn_exit_reason``. + ``usage`` carries the completed turn's canonical token usage (the same + dict shape passed to ``update_from_response`` — ``prompt_tokens`` / + ``completion_tokens`` / ``total_tokens`` plus the canonical + ``input_tokens`` / ``output_tokens`` / ``cache_read_tokens`` / + ``cache_write_tokens`` / ``reasoning_tokens`` buckets) so an engine can + weigh how large/expensive the selected context actually was when + deciding the next ``select_context()``. It is ``None`` on turns that + never reached a provider response (early failure / interrupt); engines + must treat it as optional. + Default is a no-op. """ return None diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index e164b0aa5934..425d1ff1eecd 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -956,6 +956,13 @@ def run_conversation( # over instead of spinning. Reset here so each turn starts fresh. See #26080. agent._auth_pool_refresh_counts = {} + # Reset the per-turn usage holder forwarded to the context engine's + # on_turn_complete() observation hook. Set after each successful provider + # response (see below); left as None on turns that never reach a response + # (early failure / interrupt) so the hook receives None rather than a + # stale prior turn's usage. + agent._last_turn_usage = None + # Optional opt-in runtime: if api_mode == codex_app_server, hand the # turn to the codex app-server subprocess (terminal/file ops/patching # all run inside Codex). Default Hermes path is bypassed entirely. @@ -2741,6 +2748,14 @@ def run_conversation( "reasoning_tokens": canonical_usage.reasoning_tokens, } agent.context_compressor.update_from_response(usage_dict) + + # Stash this response's canonical usage so the post-turn + # on_turn_complete() observation hook can forward it (the + # same dict shape passed to update_from_response). A turn + # may make several API calls; the engine's per-turn signal + # of interest is the cost/size of the latest assembled + # request, so we keep the most recent call's usage. + agent._last_turn_usage = dict(usage_dict) elif getattr( agent.context_compressor, "awaiting_real_usage_after_compression", diff --git a/agent/turn_finalizer.py b/agent/turn_finalizer.py index f2ddb947ba48..4e2d318b2e2c 100644 --- a/agent/turn_finalizer.py +++ b/agent/turn_finalizer.py @@ -531,10 +531,17 @@ def finalize_turn( # observation after the turn). No-op default, fail-open. try: from agent.conversation_loop import _notify_context_engine_turn_complete + # Forward the turn's canonical usage when the host has it. The loop + # stashes the most recent API response's usage dict (the same + # canonical buckets fed to ``update_from_response``) on the agent as + # ``_last_turn_usage``. It is ``None`` on turns that never reached a + # provider response (early failure / interrupt), which is exactly the + # contract: real usage when available, ``None`` otherwise. + _turn_usage = getattr(agent, "_last_turn_usage", None) _notify_context_engine_turn_complete( agent, messages, - usage=None, + usage=_turn_usage, logger=logger, turn_id=turn_id, task_id=effective_task_id, diff --git a/tests/agent/test_context_engine_on_turn_complete_usage.py b/tests/agent/test_context_engine_on_turn_complete_usage.py new file mode 100644 index 000000000000..6a80b001bba1 --- /dev/null +++ b/tests/agent/test_context_engine_on_turn_complete_usage.py @@ -0,0 +1,108 @@ +"""Integration test: ``finalize_turn`` forwards the turn's real usage to the +``ContextEngine.on_turn_complete()`` observation hook. + +The hook is the engine's post-turn observation point, so it must receive the +completed turn's canonical token usage (prompt/completion/total + the canonical +``input_tokens`` / ``output_tokens`` / ``cache_read_tokens`` / +``cache_write_tokens`` / ``reasoning_tokens`` buckets) when the host has it — +not a hardcoded ``None`` — so the engine can weigh how large/expensive the +selected context was before the next ``select_context()``. + +The conversation loop stashes the most recent provider response's usage on the +agent as ``_last_turn_usage`` (the same dict shape fed to +``update_from_response``); ``finalize_turn`` forwards it. On turns that never +reach a provider response (early failure / interrupt) it stays ``None`` and the +hook receives ``None``. These tests pin both ends of that contract through the +real ``finalize_turn`` call site (the path that previously passed +``usage=None``). +""" + +from __future__ import annotations + +from typing import Any, Dict, List + +from agent.context_engine import ContextEngine + +# Reuse the minimal agent harness that exercises the real finalize_turn path. +from tests.agent.test_turn_finalizer_cleanup_guard import _StubAgent, _run + + +class _CapturingEngine(ContextEngine): + """Engine that records what on_turn_complete() receives.""" + + last_prompt_tokens = 0 + + def __init__(self) -> None: + self.captured: Dict[str, Any] = {} + + @property + def name(self) -> str: + return "capturing" + + def update_from_response(self, usage: Dict[str, Any]) -> None: + pass + + def should_compress(self, prompt_tokens: int = None) -> bool: + return False + + def compress( + self, + messages: List[Dict[str, Any]], + current_tokens: int = None, + focus_topic: str = None, + ) -> List[Dict[str, Any]]: + return messages + + def on_turn_complete(self, messages, usage=None, **kwargs): + self.captured["seen"] = True + self.captured["usage"] = usage + self.captured["kwargs"] = kwargs + + +CANONICAL_USAGE = { + "prompt_tokens": 1200, + "completion_tokens": 80, + "total_tokens": 1280, + "input_tokens": 1200, + "output_tokens": 80, + "cache_read_tokens": 1024, + "cache_write_tokens": 0, + "reasoning_tokens": 16, +} + + +def _agent_with_engine() -> _StubAgent: + agent = _StubAgent(raise_in=()) + agent.context_compressor = _CapturingEngine() + return agent + + +def test_finalize_turn_forwards_canonical_usage_when_available(): + """A completed turn forwards the stashed canonical usage dict intact.""" + agent = _agent_with_engine() + agent._last_turn_usage = dict(CANONICAL_USAGE) + + _run(agent, final_response="done") + + captured = agent.context_compressor.captured + assert captured.get("seen") is True + # The full canonical bucket set is forwarded unchanged — the engine relies + # on cache_read/write + reasoning, not just the legacy aggregate keys. + assert captured["usage"] == CANONICAL_USAGE + # Turn metadata still rides alongside usage. + assert captured["kwargs"]["turn_id"] == "turn-1" + + +def test_finalize_turn_forwards_none_when_no_response_usage(): + """An early-failure/interrupt turn (no stashed usage) forwards None.""" + agent = _agent_with_engine() + # _last_turn_usage left unset, mirroring a turn that never reached a + # provider response. + if hasattr(agent, "_last_turn_usage"): + delattr(agent, "_last_turn_usage") + + _run(agent, final_response="done") + + captured = agent.context_compressor.captured + assert captured.get("seen") is True + assert captured["usage"] is None