diff --git a/api/streaming.py b/api/streaming.py index 35233482..5081d70e 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -2849,6 +2849,50 @@ def _assistant_reply_added_after_current_turn(result_messages, previous_context, _TOOL_RESULT_SNIPPET_MAX = 4000 +_LIVE_TOOL_PROMPT_DELTA_MAX = 12_000 + + +def _bounded_live_tool_prompt_delta(messages, *, cap: int = _LIVE_TOOL_PROMPT_DELTA_MAX) -> int: + """Return a bounded rough token delta for live tool metering. + + Tool-result callbacks can fire before the agent's next exact prompt accounting + is available. The live usage ring should show a conservative in-flight hint, + not replay a full large tool payload into `last_prompt_tokens`. + """ + if not messages: + return 0 + try: + from agent.model_metadata import estimate_messages_tokens_rough + delta = int(estimate_messages_tokens_rough(messages) or 0) + except Exception: + delta = 0 + if delta <= 0: + return 0 + return min(delta, int(cap or 0)) + + +def live_usage_prompt_estimate_after_tool_delta( + *, + base_prompt_tokens: int, + exact_prompt_tokens: int = 0, + messages=None, + cap: int = _LIVE_TOOL_PROMPT_DELTA_MAX, +) -> dict: + """Compute the live `last_prompt_tokens` estimate after a tool update. + + Exact compressor/provider prompt accounting wins. When no newer exact prompt + is available, add only a bounded live tool delta to the persisted base. + """ + base = int(base_prompt_tokens or 0) + exact = int(exact_prompt_tokens or 0) + if exact and exact != base: + return {'last_prompt_tokens': exact, 'estimated': False} + return { + 'last_prompt_tokens': base + _bounded_live_tool_prompt_delta(messages, cap=cap), + 'estimated': True, + } + + def _tool_result_snippet(raw, limit: int = _TOOL_RESULT_SNIPPET_MAX) -> str: """Extract a bounded result preview from a stored tool message payload.""" if limit <= 0: @@ -3380,11 +3424,7 @@ def _run_agent_streaming( """Increment a rough next-prompt estimate from live tool activity.""" if not messages: return _live_prompt_estimate_tokens[0] - try: - from agent.model_metadata import estimate_messages_tokens_rough - _delta = int(estimate_messages_tokens_rough(messages) or 0) - except Exception: - _delta = 0 + _delta = _bounded_live_tool_prompt_delta(messages) if _delta > 0: _seed_live_prompt_estimate() _live_prompt_estimate_tokens[0] += _delta diff --git a/pr-artifacts/context-replay-failure-cases.md b/pr-artifacts/context-replay-failure-cases.md index 3901a459..46a9a72c 100644 --- a/pr-artifacts/context-replay-failure-cases.md +++ b/pr-artifacts/context-replay-failure-cases.md @@ -141,3 +141,20 @@ Key invariant for upstream: If context_messages exists, it is the authoritative model-facing prefix. State/db/display histories may be fuller/noisier and should only contribute messages that are demonstrably newer than that prefix. ``` + + +## F. Live metering over-counts large in-flight tool results after cancel/retry + +Observed after deploying the reconciliation fix and retrying `continue` in session `20260521_060755_294aed`: + +```text +persisted sidecar context: 69 messages, ~49,993 rough content tokens +state.db messages: 90 messages +state delta after context: 21 messages, ~21,226 rough content tokens +turn-start context after fix: 90 messages, ~71,219 rough content tokens +last_prompt_tokens persisted: 86,723 (~33.9%) +``` + +The previous full-transcript turn-start replay was gone, but the ring still jumped during the run. The new jump came from live metering, not persisted context reconciliation. The run executed several large `read_file` tool calls (5k / 17k / 13k chars). `_record_live_tool_complete()` fed each bounded preview into `_bump_live_prompt_estimate()`, which added the full rough tool-result tokens to `last_prompt_tokens` before any exact next-prompt accounting was available. Repeated cancel/retry makes this look like context replay even when final sidecar context remains clean. + +Regression target: live tool metering should be a conservative UI hint and must not inflate `last_prompt_tokens` by the full content of large in-flight tool results. Exact provider/compressor prompt accounting should still win when available. diff --git a/tests/test_streaming_live_usage_estimate.py b/tests/test_streaming_live_usage_estimate.py new file mode 100644 index 00000000..57d87c08 --- /dev/null +++ b/tests/test_streaming_live_usage_estimate.py @@ -0,0 +1,22 @@ +from api.streaming import live_usage_prompt_estimate_after_tool_delta + + +def test_live_usage_estimate_caps_tool_delta_against_previous_prompt(): + usage = live_usage_prompt_estimate_after_tool_delta( + base_prompt_tokens=86_723, + exact_prompt_tokens=86_723, + messages=[{"role": "tool", "content": "x" * 80_000}], + ) + + assert usage["last_prompt_tokens"] <= 86_723 + 12_000 + assert usage["last_prompt_tokens"] < 120_000 + + +def test_live_usage_estimate_preserves_real_prompt_when_exact_prompt_advances(): + usage = live_usage_prompt_estimate_after_tool_delta( + base_prompt_tokens=86_723, + exact_prompt_tokens=136_000, + messages=[{"role": "tool", "content": "x" * 80_000}], + ) + + assert usage["last_prompt_tokens"] == 136_000