diff --git a/backend/apps/agents/agent_manager.py b/backend/apps/agents/agent_manager.py index 9f97a15e..5424540d 100644 --- a/backend/apps/agents/agent_manager.py +++ b/backend/apps/agents/agent_manager.py @@ -57,6 +57,7 @@ from backend.apps.agents.manager.session import lifecycle from backend.apps.agents.manager.permissions import path_gate from backend.apps.agents.manager import context_budget from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState +from backend.apps.agents.manager.streaming import thinking as thinking_mod from backend.apps.agents.manager.session.workspace_git import _detect_git_identity, _ensure_cwd_git_repo from backend.apps.agents.manager.prompt.tool_catalog import ( FULL_TOOLS, @@ -1989,221 +1990,6 @@ class AgentManager: # (auth, plan limit, invalid args) fall through to the existing # error handler unchanged. - async def _emit_consolidated_thinking(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: int | None = 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 _9r_running, - ) - if _9r_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: int | None = 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 _9r_running, - ) - if _9r_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 self.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: int | None = ( - _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") - - async def _ticker_loop(): - """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() - except asyncio.CancelledError: - pass - async def _run_streaming_turn(): # Per-turn thinking aggregation trackers (added for the # "Thought for Ns · M tokens" persisted label). Without @@ -2264,7 +2050,7 @@ class AgentManager: and resolved_model.startswith(("cx/", "gc/", "ag/", "gemini/")) ) if _route_strips_reasoning_pre: - await _emit_consolidated_thinking(force_provider_unavailable=True) + await thinking_mod.emit_consolidated_thinking(thinking, turn, session, session_id, self.sessions, force_provider_unavailable=True) except Exception: logger.exception("pre-emit thinking pill failed; continuing") @@ -2506,12 +2292,12 @@ class AgentManager: # ticks even during tool execution / slow text # generation gaps. if thinking.text_parts: - await _emit_consolidated_thinking() + await thinking_mod.emit_consolidated_thinking(thinking, turn, session, session_id, self.sessions) # Start the 1Hz ticker once we have a # consolidated message in flight so the # bubble keeps updating between SDK events. if thinking.ticker_task is None or thinking.ticker_task.done(): - thinking.ticker_task = asyncio.create_task(_ticker_loop()) + thinking.ticker_task = asyncio.create_task(thinking_mod.ticker_loop(thinking, turn, session, session_id, self.sessions)) if content_parts: _asst_text = "\n".join(content_parts) @@ -2619,7 +2405,7 @@ class AgentManager: # Pre-populate session.tokens BEFORE emitting the # final consolidated thinking pill. Order matters: - # _emit_consolidated_thinking reads + # emit_consolidated_thinking reads # session.tokens["input"]/["output"] for the # combined-total stamp on the pill. If we emit # first, the pill freezes with input=0 because @@ -2664,7 +2450,8 @@ class AgentManager: ) if thinking.text_parts or _route_strips_reasoning: try: - await _emit_consolidated_thinking( + await thinking_mod.emit_consolidated_thinking( + thinking, turn, session, session_id, self.sessions, force_provider_unavailable=_route_strips_reasoning, ) except Exception: diff --git a/backend/apps/agents/manager/streaming/thinking.py b/backend/apps/agents/manager/streaming/thinking.py new file mode 100644 index 00000000..96753779 --- /dev/null +++ b/backend/apps/agents/manager/streaming/thinking.py @@ -0,0 +1,238 @@ +"""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 +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: int | None = 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 _9r_running, + ) + if _9r_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: int | None = 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 _9r_running, + ) + if _9r_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: int | None = ( + _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 +