From fecfc5f6dbf832e230e914505ce9656c4a5782b6 Mon Sep 17 00:00:00 2001 From: dobby-d-elf Date: Sun, 10 May 2026 10:31:14 -0600 Subject: [PATCH] fix: reanchor live context usage updates --- api/streaming.py | 8 ++++++- static/messages.js | 8 ------- tests/test_issue1617_tps_message_header.py | 26 ++++++++++++++++++++++ 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/api/streaming.py b/api/streaming.py index 1fe28521..2653f6d3 100644 --- a/api/streaming.py +++ b/api/streaming.py @@ -1939,6 +1939,7 @@ def _run_agent_streaming( agent = None _live_prompt_estimate_tokens = [0] + _live_prompt_exact_tokens = [0] _live_prompt_estimate_seen_ids = set() def _seed_live_prompt_estimate() -> int: @@ -1961,6 +1962,7 @@ def _run_agent_streaming( except Exception: _base = 0 _live_prompt_estimate_tokens[0] = int(_base or 0) + _live_prompt_exact_tokens[0] = _live_prompt_estimate_tokens[0] return _live_prompt_estimate_tokens[0] def _bump_live_prompt_estimate(messages) -> int: @@ -2023,7 +2025,11 @@ def _run_agent_streaming( except Exception: pass - if _live_prompt_estimate_tokens[0] > (_usage.get('last_prompt_tokens') or 0): + _real_prompt_tokens = int(_usage.get('last_prompt_tokens') or 0) + if _real_prompt_tokens and _real_prompt_tokens != _live_prompt_exact_tokens[0]: + _live_prompt_exact_tokens[0] = _real_prompt_tokens + _live_prompt_estimate_tokens[0] = _real_prompt_tokens + elif _live_prompt_estimate_tokens[0] > _real_prompt_tokens: _usage['last_prompt_tokens'] = _live_prompt_estimate_tokens[0] return _usage diff --git a/static/messages.js b/static/messages.js index fd5cd036..6ee504fd 100644 --- a/static/messages.js +++ b/static/messages.js @@ -1161,14 +1161,6 @@ function attachLiveStream(activeSid, streamId, uploaded=[], options={}){ if((d.session_id||activeSid)!==activeSid) return; if(d.usage&&typeof _syncCtxIndicator==='function'){ S.lastUsage={...(S.lastUsage||{}),...d.usage}; - if(S.session&&S.session.session_id===activeSid){ - S.session.input_tokens=d.usage.input_tokens??S.session.input_tokens; - S.session.output_tokens=d.usage.output_tokens??S.session.output_tokens; - S.session.estimated_cost=d.usage.estimated_cost??S.session.estimated_cost; - S.session.context_length=d.usage.context_length??S.session.context_length; - S.session.threshold_tokens=d.usage.threshold_tokens??S.session.threshold_tokens; - S.session.last_prompt_tokens=d.usage.last_prompt_tokens??S.session.last_prompt_tokens; - } _syncCtxIndicator(S.lastUsage); } if(d.estimated===true||d.tps_available!==true||typeof d.tps!=='number'||d.tps<=0){ diff --git a/tests/test_issue1617_tps_message_header.py b/tests/test_issue1617_tps_message_header.py index b3c49db2..814d2ea7 100644 --- a/tests/test_issue1617_tps_message_header.py +++ b/tests/test_issue1617_tps_message_header.py @@ -46,6 +46,32 @@ def test_live_metering_updates_only_real_tps_and_never_placeholders(): ) +def test_live_metering_usage_is_provisional_until_done(): + listener_start = MESSAGES_JS.find("source.addEventListener('metering'") + assert listener_start != -1, "messages.js should listen for metering SSE events" + listener_end = MESSAGES_JS.find("source.addEventListener('apperror'", listener_start) + assert listener_end != -1, "apperror listener should follow metering listener" + listener = MESSAGES_JS[listener_start:listener_end] + + assert "S.lastUsage={...(S.lastUsage||{}),...d.usage}" in listener, ( + "live usage should update the transient usage cache for the indicator" + ) + assert "_syncCtxIndicator(S.lastUsage)" in listener, ( + "live usage should refresh the context indicator" + ) + assert "S.session.input_tokens=d.usage.input_tokens" not in listener + assert "S.session.last_prompt_tokens=d.usage.last_prompt_tokens" not in listener + + +def test_live_prompt_estimate_reanchors_to_fresh_exact_prompt_tokens(): + assert "_live_prompt_exact_tokens = [0]" in STREAMING_PY, ( + "live prompt estimates need a separate exact-token anchor" + ) + assert "_real_prompt_tokens = int(_usage.get('last_prompt_tokens') or 0)" in STREAMING_PY + assert "_real_prompt_tokens != _live_prompt_exact_tokens[0]" in STREAMING_PY + assert "_live_prompt_estimate_tokens[0] = _real_prompt_tokens" in STREAMING_PY + + def test_done_payload_persists_final_tps_when_exact_usage_available(): assert "usage['tps']" in STREAMING_PY, "done usage payload should include final exact TPS when available" assert "output_tokens" in STREAMING_PY and "duration_seconds" in STREAMING_PY, (