diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index c90fc3fb..9f97a15e 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -1971,11 +1971,6 @@ class AgentManager: # that don't use signatures. # session.tokens accumulates SDK running totals across turns, # so subtract the turn-start baseline to get this turn's delta. - _turn_baseline_session_in: int = 0 - _turn_baseline_session_out: int = 0 - _turn_baseline_children_in: int = 0 - _turn_baseline_children_out: int = 0 - _turn_baseline_captured: bool = False # Background ticker handle. Re-emits the consolidated # thinking message every 1s so the elapsed counter keeps # ticking through gaps where no SDK events fire (tool @@ -2148,11 +2143,11 @@ class AgentManager: # Fall back to cumulative if the baseline wasn't captured # (degenerate empty turn, better than showing zero). - if _turn_baseline_captured: - _parent_in = max(0, _cum_in - _turn_baseline_session_in) - _parent_out = max(0, _cum_out - _turn_baseline_session_out) - _children_in = max(0, _cum_children_in - _turn_baseline_children_in) - _children_out = max(0, _cum_children_out - _turn_baseline_children_out) + if turn.baseline_captured: + _parent_in = max(0, _cum_in - turn.baseline_session_in) + _parent_out = max(0, _cum_out - turn.baseline_session_out) + _children_in = max(0, _cum_children_in - turn.baseline_children_in) + _children_out = max(0, _cum_children_out - turn.baseline_children_out) else: _parent_in = _cum_in _parent_out = _cum_out @@ -2237,8 +2232,8 @@ class AgentManager: # Baselines track the SAME fresh lane the pill reads, # so the per-turn delta is fresh-minus-fresh. if isinstance(session.tokens, dict): - _turn_baseline_session_in = int(session.tokens.get("input_fresh", 0) or 0) - _turn_baseline_session_out = int(session.tokens.get("output", 0) or 0) + turn.baseline_session_in = int(session.tokens.get("input_fresh", 0) or 0) + turn.baseline_session_out = int(session.tokens.get("output", 0) or 0) _ch_in = 0 _ch_out = 0 for _child in self.sessions.values(): @@ -2249,9 +2244,9 @@ class AgentManager: continue _ch_in += int(_ct.get("input_fresh", 0) or 0) _ch_out += int(_ct.get("output", 0) or 0) - _turn_baseline_children_in = _ch_in - _turn_baseline_children_out = _ch_out - _turn_baseline_captured = True + turn.baseline_children_in = _ch_in + turn.baseline_children_out = _ch_out + turn.baseline_captured = True except Exception: pass # Pre-emit thinking pill for routes whose @@ -2690,11 +2685,11 @@ class AgentManager: turn.assistant_text_chars = 0 turn.tool_input_chars = 0 thinking.thought_signature = None - _turn_baseline_session_in = 0 - _turn_baseline_session_out = 0 - _turn_baseline_children_in = 0 - _turn_baseline_children_out = 0 - _turn_baseline_captured = False + turn.baseline_session_in = 0 + turn.baseline_session_out = 0 + turn.baseline_children_in = 0 + turn.baseline_children_out = 0 + turn.baseline_captured = False thinking.total_ms = 0 thinking.total_chars = 0 thinking.block_starts = {} diff --git a/backend/apps/agents/manager/streaming/state.py b/backend/apps/agents/manager/streaming/state.py index 69f75675..937f73d5 100644 --- a/backend/apps/agents/manager/streaming/state.py +++ b/backend/apps/agents/manager/streaming/state.py @@ -49,3 +49,10 @@ class TurnState(BaseModel): output_tokens: int = 0 assistant_text_chars: int = 0 tool_input_chars: int = 0 + # Cumulative-token snapshot taken at turn start; subtracted at emit time so the thinking + # pill shows THIS turn's new tokens, not the whole session's running total. + baseline_session_in: int = 0 + baseline_session_out: int = 0 + baseline_children_in: int = 0 + baseline_children_out: int = 0 + baseline_captured: bool = False diff --git a/backend/tests/test_streaming_harness.py b/backend/tests/test_streaming_harness.py index 0b8522f3..df5a0b0b 100644 --- a/backend/tests/test_streaming_harness.py +++ b/backend/tests/test_streaming_harness.py @@ -141,3 +141,35 @@ def test_transient_capacity_error_is_retried_then_succeeds(monkeypatch): assert state["n"] == 2 # retried exactly once assert any(m.role == "assistant" and "Recovered" in str(m.content) for m in session.messages) assert session.status == "completed" + + +def test_thinking_pill_shows_per_turn_delta_not_cumulative(monkeypatch): + # The pill's token total must reflect THIS turn's new tokens, not the whole session's + # running cumulative (the baseline-delta fix: capture-at-turn-start, subtract-at-emit, + # unified through TurnState). Prior turns left 1500 tokens on the session; this turn adds + # 100 in + 50 out = 150. Before the fix the baseline writes leaked into a closure-local + # and the pill showed the cumulative 1650; now it shows 150. + pills = [] + + async def fake_send(sid, event, data): + msg = data.get("message") if isinstance(data, dict) else None + if isinstance(msg, dict) and msg.get("role") == "thinking": + pills.append(msg) + + async def q(*a, **k): + yield _assistant([ThinkingBlock(thinking="reasoning", signature="s"), TextBlock(text="answer")], + usage={"input_tokens": 100, "output_tokens": 50}) + yield _result(usage={"input_tokens": 1100, "output_tokens": 550}) + + monkeypatch.setattr(ws_mod.ws_manager, "send_to_session", fake_send, raising=True) + monkeypatch.setattr(claude_agent_sdk, "query", q, raising=True) + + mgr = AgentManager() + from backend.apps.agents.core.models import AgentSession + session = AgentSession(name="t", model="sonnet", dashboard_id="d") + session.tokens = {"input_fresh": 1000, "output": 500} # prior-turn accumulation + mgr.sessions[session.id] = session + asyncio.run(mgr._run_agent_loop(session.id, "hi")) + + assert pills, "expected a consolidated thinking pill" + assert pills[-1]["input_tokens"] == 150 # (1100-1000)+(550-500), not the cumulative 1650