686f6c61
c0aaa238f6
The usage anchor (real usage.prompt_tokens + delta estimate of what was appended since) identified the priced transcript by id() of the last message, so it was None on EVERY gateway turn (history is re-read from the DB each turn) and in every fresh process (--resume, desktop per-turn serve). Those are exactly the surfaces where the bytes/4 estimate then fired local compression against payloads the provider priced far under threshold (#99421, #104462). - agent/usage_anchor.py owns the anchor: content fingerprint instead of id(), persisted on the session row (model_config._usage_anchor) via set_usage_anchor(), restored on the first resumed turn while the durable transcript still matches, cleared with the row on compaction / codex-native rewrite / session reset. - Callers repointed from model_metadata (the compat table follows). Design and persistence slot from #99585 by @686f6c61; re-authored against the Sep 2026 layout (the branch predates the model_metadata / agent_init split).
199 lines
8.6 KiB
Python
199 lines
8.6 KiB
Python
"""Turn-base display anchor: the context meter shows durable-transcript cost.
|
|
|
|
On reasoning models a long tool loop replays the current turn's thinking +
|
|
scaffolding on every request, so the LAST request's ``prompt_tokens`` can
|
|
exceed the durable transcript by hundreds of K — all of which evaporates at
|
|
the turn boundary. Display surfaces (CLI status bar, /context breakdown)
|
|
therefore anchor on the turn's FIRST response (``_turn_base_usage_anchor``)
|
|
plus a stale-thinking-free delta estimate, instead of the raw last-request
|
|
figure. Compression trigger math is unchanged (real last-request usage).
|
|
|
|
Covers:
|
|
* anchored_context_tokens(charge_stale_thinking=False) excludes stale
|
|
reasoning text in the delta while keeping the newest assistant turn;
|
|
* the CLI status snapshot prefers the turn-base anchored figure over
|
|
compressor.last_prompt_tokens and falls back cleanly without an anchor;
|
|
* compute_session_context_breakdown prefers the turn-base anchor over the
|
|
last-response anchor;
|
|
* invalidation sites clear _turn_base_usage_anchor alongside _usage_anchor.
|
|
"""
|
|
|
|
from types import SimpleNamespace
|
|
|
|
from agent.model_metadata import estimate_messages_tokens_rough
|
|
from agent.usage_anchor import anchored_context_tokens, capture_usage_anchor
|
|
|
|
|
|
def _msg(role, content, **extra):
|
|
m = {"role": role, "content": content}
|
|
m.update(extra)
|
|
return m
|
|
|
|
|
|
class TestChargeStaleThinkingKwarg:
|
|
def test_delta_excludes_stale_reasoning(self):
|
|
messages = [_msg("user", "start"), _msg("assistant", "base reply")]
|
|
anchor = capture_usage_anchor(10_000, 100, messages)
|
|
assert anchor is not None
|
|
|
|
# Simulate a tool loop appending reasoning-heavy assistant turns.
|
|
big_thinking = "deliberation " * 5_000 # ~65K chars ≈ 16K tokens
|
|
messages.append(_msg("assistant", "the anchored reply itself"))
|
|
messages.append(
|
|
_msg("assistant", "step one", reasoning_content=big_thinking)
|
|
)
|
|
messages.append(_msg("tool", "tool output", tool_call_id="c1"))
|
|
messages.append(
|
|
_msg("assistant", "step two", reasoning_content=big_thinking)
|
|
)
|
|
|
|
charged = anchored_context_tokens(messages, anchor)
|
|
uncharged = anchored_context_tokens(
|
|
messages, anchor, charge_stale_thinking=False
|
|
)
|
|
assert charged is not None and uncharged is not None
|
|
# Stale thinking on the non-newest assistant message is excluded;
|
|
# the newest assistant message keeps its reasoning charge.
|
|
one_thinking_tokens = estimate_messages_tokens_rough(
|
|
[_msg("assistant", "", reasoning_content=big_thinking)]
|
|
)
|
|
assert charged - uncharged >= one_thinking_tokens * 0.9
|
|
assert uncharged >= 10_000 + 100 # anchor base still counted exactly
|
|
|
|
def test_default_remains_full_charge(self):
|
|
messages = [_msg("user", "s"), _msg("assistant", "r")]
|
|
anchor = capture_usage_anchor(1_000, 10, messages)
|
|
messages.append(_msg("assistant", "reply"))
|
|
assert anchored_context_tokens(messages, anchor) == anchored_context_tokens(
|
|
messages, anchor, charge_stale_thinking=True
|
|
)
|
|
|
|
|
|
class TestCliStatusSnapshotPrefersTurnBaseAnchor:
|
|
def _agent_with(self, last_prompt_tokens, messages, anchor):
|
|
compressor = SimpleNamespace(
|
|
last_prompt_tokens=last_prompt_tokens,
|
|
context_length=1_000_000,
|
|
compression_count=0,
|
|
)
|
|
return SimpleNamespace(
|
|
context_compressor=compressor,
|
|
_session_messages=messages,
|
|
_turn_base_usage_anchor=anchor,
|
|
)
|
|
|
|
def _snapshot_context_tokens(self, agent):
|
|
"""Mirror the cli.py snapshot block's context_tokens resolution."""
|
|
compressor = agent.context_compressor
|
|
context_tokens = getattr(compressor, "last_prompt_tokens", 0) or 0
|
|
if context_tokens < 0:
|
|
context_tokens = 0
|
|
msgs = getattr(agent, "_session_messages", None)
|
|
anchored = anchored_context_tokens(
|
|
msgs if isinstance(msgs, list) else [],
|
|
getattr(agent, "_turn_base_usage_anchor", None),
|
|
charge_stale_thinking=False,
|
|
)
|
|
if anchored is not None and anchored > 0:
|
|
context_tokens = anchored
|
|
return context_tokens
|
|
|
|
def test_turn_base_anchor_wins_over_inflated_last_request(self):
|
|
messages = [_msg("user", "start"), _msg("assistant", "reply")]
|
|
anchor = capture_usage_anchor(600_000, 500, messages)
|
|
messages.append(_msg("assistant", "anchored reply"))
|
|
agent = self._agent_with(850_000, messages, anchor)
|
|
# Bar shows the durable figure, not the inflated last request.
|
|
tokens = self._snapshot_context_tokens(agent)
|
|
assert 600_000 <= tokens < 650_000
|
|
|
|
def test_fallback_without_anchor(self):
|
|
agent = self._agent_with(123_456, [_msg("user", "x")], None)
|
|
assert self._snapshot_context_tokens(agent) == 123_456
|
|
|
|
def test_stale_anchor_falls_back(self):
|
|
messages = [_msg("user", "start"), _msg("assistant", "reply")]
|
|
anchor = capture_usage_anchor(50_000, 10, messages)
|
|
agent = self._agent_with(77_000, [_msg("user", "rebuilt")], anchor)
|
|
# Compaction rebuilt the list: structural check fails, raw fallback.
|
|
assert self._snapshot_context_tokens(agent) == 77_000
|
|
|
|
def test_negative_sentinel_still_clamped(self):
|
|
agent = self._agent_with(-1, [], None)
|
|
assert self._snapshot_context_tokens(agent) == 0
|
|
|
|
|
|
class TestContextBreakdownPrefersTurnBaseAnchor:
|
|
def test_breakdown_uses_turn_base_over_last_response(self, monkeypatch):
|
|
from agent import context_breakdown as cb
|
|
|
|
messages = [_msg("user", "start"), _msg("assistant", "reply")]
|
|
turn_base = capture_usage_anchor(400_000, 200, messages)
|
|
messages.append(_msg("assistant", "anchored reply"))
|
|
last_anchor = capture_usage_anchor(900_000, 50, messages)
|
|
|
|
agent = SimpleNamespace(
|
|
_usage_anchor=last_anchor,
|
|
_turn_base_usage_anchor=turn_base,
|
|
_memory_store=None,
|
|
tools=[],
|
|
model="test/model",
|
|
context_compressor=SimpleNamespace(
|
|
context_length=1_000_000, last_prompt_tokens=900_000
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
"agent.system_prompt.build_system_prompt_parts",
|
|
lambda a: {"stable": "sys", "context": "", "volatile": ""},
|
|
)
|
|
payload = cb.compute_session_context_breakdown(agent, messages)
|
|
assert 400_000 <= payload["context_used"] < 450_000
|
|
|
|
def test_breakdown_falls_back_to_last_response_anchor(self, monkeypatch):
|
|
from agent import context_breakdown as cb
|
|
|
|
messages = [_msg("user", "start"), _msg("assistant", "reply")]
|
|
last_anchor = capture_usage_anchor(300_000, 50, messages)
|
|
|
|
agent = SimpleNamespace(
|
|
_usage_anchor=last_anchor,
|
|
_turn_base_usage_anchor=None,
|
|
_memory_store=None,
|
|
tools=[],
|
|
model="test/model",
|
|
context_compressor=SimpleNamespace(
|
|
context_length=1_000_000, last_prompt_tokens=1
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
"agent.system_prompt.build_system_prompt_parts",
|
|
lambda a: {"stable": "sys", "context": "", "volatile": ""},
|
|
)
|
|
payload = cb.compute_session_context_breakdown(agent, messages)
|
|
assert payload["context_used"] >= 300_000
|
|
|
|
|
|
class TestInvalidationSitesClearTurnBaseAnchor:
|
|
def test_clearing_the_anchor_clears_the_turn_base_too(self):
|
|
"""Compaction and the codex-native rewrite clear via set_usage_anchor(None): both the
|
|
last-response anchor and the display turn-base anchor go; a later same-turn capture
|
|
(turn_base=False) leaves the turn-base untouched."""
|
|
from agent.usage_anchor import set_usage_anchor
|
|
|
|
messages = [_msg("user", "a"), _msg("assistant", "b")]
|
|
agent = SimpleNamespace(_usage_anchor=None, _turn_base_usage_anchor=None, _session_db=None, session_id=None)
|
|
first = capture_usage_anchor(1_000, 10, messages)
|
|
set_usage_anchor(agent, first, turn_base=True)
|
|
set_usage_anchor(agent, capture_usage_anchor(2_000, 10, messages))
|
|
assert agent._turn_base_usage_anchor is first
|
|
set_usage_anchor(agent, None)
|
|
assert agent._usage_anchor is None and agent._turn_base_usage_anchor is None
|
|
|
|
def test_agent_init_defines_turn_base_anchor(self):
|
|
import inspect
|
|
from agent import agent_init
|
|
|
|
src = inspect.getsource(agent_init)
|
|
# Init defaults are declared in the ``_USAGE_STATE`` table.
|
|
assert '"_turn_base_usage_anchor": None' in src
|