diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 5725045a..52f07344 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -1106,6 +1106,9 @@ class AgentManager: usage = raw_response.get("usage", {}) if isinstance(usage, dict): sub_tokens["input"] = usage.get("input_tokens", 0) + usage.get("cache_creation_input_tokens", 0) + usage.get("cache_read_input_tokens", 0) + # Pill-only lane: NEW (uncached) input, excludes the cached + # static prefix so the bubble shows what this turn added. + sub_tokens["input_fresh"] = usage.get("input_tokens", 0) sub_tokens["output"] = usage.get("output_tokens", 0) if raw_response.get("total_cost_usd"): sub_cost = raw_response["total_cost_usd"] @@ -2394,10 +2397,13 @@ class AgentManager: # baseline to get THIS TURN'S delta. Without subtracting, # the second turn's pill would show turn-1 work added # to turn-2 work, the third would show all three, etc. + # Pill uses the FRESH lane (uncached input only). session.tokens + # ["input"] stays full for the context-fullness bar + cost; the + # bubble shows the NEW tokens this turn, not the cached re-reads. _cum_in = 0 _cum_out = 0 if isinstance(session.tokens, dict): - _cum_in = int(session.tokens.get("input", 0) or 0) + _cum_in = int(session.tokens.get("input_fresh", 0) or 0) _cum_out = int(session.tokens.get("output", 0) or 0) _cum_children_in = 0 _cum_children_out = 0 @@ -2408,7 +2414,7 @@ class AgentManager: _ct = getattr(_child, "tokens", None) if not isinstance(_ct, dict): continue - _cum_children_in += int(_ct.get("input", 0) or 0) + _cum_children_in += int(_ct.get("input_fresh", 0) or 0) _cum_children_out += int(_ct.get("output", 0) or 0) except Exception: pass @@ -2426,18 +2432,14 @@ class AgentManager: _children_in = _cum_children_in _children_out = _cum_children_out + # Fresh input + output = the NEW tokens this turn. The old + # framework-overhead subtraction is gone on purpose: it was an + # estimate to strip the cached static prefix out of the full + # input number, and the fresh lane already excludes that prefix + # exactly, so subtracting it again would double-discount to ~0. _turn_total_tokens: int | None = ( _parent_in + _parent_out + _children_in + _children_out ) - # Strip framework overhead so bubble shows what the user - # actually controls. Floor at output so over-estimates can't - # render absurdly small. - if _turn_total_tokens and session.framework_overhead_tokens > 0: - _adjusted = _turn_total_tokens - session.framework_overhead_tokens - _floor = _parent_out + _children_out - if _adjusted < _floor: - _adjusted = _floor - _turn_total_tokens = _adjusted if not _turn_total_tokens or _turn_total_tokens <= 0: _turn_total_tokens = None consolidated = Message( @@ -2514,8 +2516,10 @@ class AgentManager: # Snapshot cumulative tokens at turn start; # subtracted at emit time for per-turn deltas. try: + # 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", 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 @@ -2525,7 +2529,7 @@ class AgentManager: _ct = getattr(_child, "tokens", None) if not isinstance(_ct, dict): continue - _ch_in += int(_ct.get("input", 0) or 0) + _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 @@ -2919,6 +2923,9 @@ class AgentManager: _pre_out = int(_pre_usage.get("output_tokens", 0) or 0) if _pre_total_in > 0: session.tokens["input"] = _pre_total_in + # Pill reads the fresh lane: uncached input only, + # so re-read/cached context doesn't inflate it. + session.tokens["input_fresh"] = _pre_in if _pre_out > 0: session.tokens["output"] = _pre_out except Exception: @@ -2986,6 +2993,7 @@ class AgentManager: cache_read = usage.get("cache_read_input_tokens", 0) or 0 total_input = inp + cache_create + cache_read session.tokens["input"] = total_input + session.tokens["input_fresh"] = inp session.tokens["output"] = out cost = getattr(message, "total_cost_usd", None) diff --git a/backend/apps/agents/browser/browser_agent.py b/backend/apps/agents/browser/browser_agent.py index 12419b27..779ea6e0 100644 --- a/backend/apps/agents/browser/browser_agent.py +++ b/backend/apps/agents/browser/browser_agent.py @@ -1014,6 +1014,9 @@ async def run_browser_agent( _in = response.usage.input_tokens or 0 out_tokens_total += _out session.tokens["input"] = session.tokens.get("input", 0) + _in + # Already-uncached here (cache tracked separately below), so the + # fresh lane that feeds the parent's pill mirrors it 1:1. + session.tokens["input_fresh"] = session.tokens.get("input_fresh", 0) + _in session.tokens["output"] = session.tokens.get("output", 0) + _out _cr = getattr(response.usage, "cache_read_input_tokens", 0) or 0 _cw = getattr(response.usage, "cache_creation_input_tokens", 0) or 0