"""The consolidated-thinking pill: build the running 'Thought for Ns · N tokens · N tools' aggregate message and broadcast it, plus the 1s ticker that keeps the elapsed counter moving. Lifted out of the agent loop; operates on the passed TurnState/ThinkingState + session.""" import asyncio import time from typing import Dict, Optional from uuid import uuid4 from typeguard import typechecked from backend.apps.agents.core.models import AgentSession, Message from backend.apps.agents.core.ws_manager import ws_manager from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState import logging logger = logging.getLogger(__name__) @typechecked async def emit_consolidated_thinking(thinking: ThinkingState, turn: TurnState, session: AgentSession, session_id: str, sessions: Dict[str, AgentSession], force_provider_unavailable: bool = False) -> None: """Build the running aggregate Message and broadcast it. Safe to call multiple times, uses a stable per-turn id so the frontend dedupes by id and updates the bubble in place. Emission rule: emit when ANY of the following is true: 1. Reasoning text exists (Anthropic happy path). 2. Upstream provider reported reasoning tokens via 9Router (best-effort path for GPT/Gemini). 3. force_provider_unavailable=True, caller has determined this turn went through a translator that doesn't carry reasoning content (cx/ or gc/), and the user should see a "provider doesn't expose reasoning text" pill regardless of metric availability. This is what makes GPT/Gemini turns show a pill even when 9Router can't surface a token count. """ upstream_reasoning_tokens: Optional[int] = None # Probe 9Router for the upstream reasoning-token count # whenever (a) there's no in-process text, OR (b) the # caller flagged this as a force-emit for a route that # strips reasoning. Case (b) is what makes the FINAL # emit on GPT/Gemini show the real reasoning count # (e.g. 196) instead of the heuristic chars/3.6 of the # answer text (e.g. 13). if not thinking.text_parts or force_provider_unavailable: try: from backend.apps.nine_router import ( get_latest_reasoning_tokens, is_running as nine_router_running, ) if nine_router_running(): rt = await get_latest_reasoning_tokens(model_hint=session.model) if rt and rt > 0: upstream_reasoning_tokens = rt except Exception: pass if ( not thinking.text_parts and upstream_reasoning_tokens is None and not force_provider_unavailable ): # No text, no upstream signal, and caller didn't # ask for the unavailable-pill, nothing to show. return joined_text = "\n".join(thinking.text_parts) # Total turn output token estimate. Combines two sources: # - SDK usage.output_tokens summed across completed # AssistantMessages (authoritative for finished # blocks). # - chars/3.6 heuristic over the running streams of # thinking + assistant-text + tool-input JSON # (covers in-flight blocks the SDK hasn't billed # yet, i.e. the answer the user is currently # reading). # Take the max so the number doesn't visually shrink as # the SDK's authoritative count overtakes our running # heuristic. running_chars = ( len(joined_text) + turn.assistant_text_chars + turn.tool_input_chars ) heuristic_tokens = max(1, round(running_chars / 3.6)) if running_chars else 0 turn_tokens: Optional[int] = None # Priority order: # 1. Upstream reasoning-token count from 9Router (the # only honest signal for GPT/Gemini, captured above). # 2. SDK-reported usage.output_tokens (Anthropic). # 3. chars/3.6 heuristic over running streams (live UI). if upstream_reasoning_tokens and upstream_reasoning_tokens > 0: turn_tokens = upstream_reasoning_tokens elif turn.output_tokens > 0 or heuristic_tokens > 0: turn_tokens = max(turn.output_tokens, heuristic_tokens) else: try: from backend.apps.nine_router import ( get_latest_reasoning_tokens, is_running as nine_router_running, ) if nine_router_running(): rt = await get_latest_reasoning_tokens(model_hint=session.model) if rt and rt > 0: turn_tokens = rt except Exception: pass if turn.started_ts is not None: turn.total_ms = int((time.time() - turn.started_ts) * 1000) # Accumulate into session-level "agent active time" and # the per-model breakdown so a session that spans # multiple turns reports the total wall-clock time the # agent was running. Per-model bucket uses the model # active *now* (model can be switched mid-turn but the # current value is the right attribution for the work # just produced). try: session.agent_active_ms = int(getattr(session, "agent_active_ms", 0) or 0) + turn.total_ms m = session.model or "unknown" session.time_per_model[m] = int(session.time_per_model.get(m, 0)) + turn.total_ms except Exception: pass if thinking.msg_id is None: thinking.msg_id = uuid4().hex # Combined token total for the pill, input + output for # the parent turn PLUS any work delegated to subagents # (browser agents, invoke-agent forks) and tool MCP # servers that produced their own usage on this turn. # The user-visible answer to "how big is this turn" is # the all-in sum, not just the primary's output. We sum # every reachable source: # - parent's input (session.tokens["input"], # ResultMessage.usage at line ~2886) # - parent's output (session.tokens["output"], same # ResultMessage) # - every direct sub-session whose parent_session_id # points at this session (browser agents, sub-agent # forks, invoke-agent calls book their own usage at # subprocess return time, agent_manager.py:1365 + # browser_agent.py:1000-1001) # This mirrors how billing accumulates per-turn, caches, # tool MCP servers that talk to LLMs (e.g. summarizers), # and subagent reasoning all show up under the parent's # "session.tokens" once their result lands. # Read cumulative session totals + cumulative subagent # totals at this moment, then subtract the turn-start # 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_fresh", 0) or 0) cum_out = int(session.tokens.get("output", 0) or 0) cum_children_in = 0 cum_children_out = 0 try: for child in sessions.values(): if getattr(child, "parent_session_id", None) != session.id: continue ct = getattr(child, "tokens", None) if not isinstance(ct, dict): continue cum_children_in += int(ct.get("input_fresh", 0) or 0) cum_children_out += int(ct.get("output", 0) or 0) except Exception: pass # 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) else: parent_in = cum_in parent_out = cum_out 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: Optional[int] = ( parent_in + parent_out + children_in + children_out ) if not turn_total_tokens or turn_total_tokens <= 0: turn_total_tokens = None consolidated = Message( id=thinking.msg_id, role="thinking", content=joined_text, branch_id=session.active_branch_id, elapsed_ms=turn.total_ms or None, tokens=turn_tokens, input_tokens=turn_total_tokens, tool_count=turn.tool_count or None, ) existing_idx = next( (i for i, m in enumerate(session.messages) if m.id == thinking.msg_id), -1, ) if existing_idx >= 0: session.messages[existing_idx] = consolidated else: session.messages.append(consolidated) try: await ws_manager.send_to_session(session_id, "agent:message", { "session_id": session_id, "message": consolidated.model_dump(mode="json"), }) except Exception: logger.exception("Failed to emit consolidated thinking message") @typechecked async def ticker_loop(thinking: ThinkingState, turn: TurnState, session: AgentSession, session_id: str, sessions: Dict[str, AgentSession]) -> None: """Re-emit the consolidated thinking message every 1s so the elapsed-time counter keeps ticking through gaps where no SDK events fire (e.g. while a tool is running or while assistant text is being generated). Cancelled at turn boundaries from `ResultMessage`.""" try: while True: await asyncio.sleep(1.0) await emit_consolidated_thinking(thinking, turn, session, session_id, sessions) except asyncio.CancelledError: pass