openswarm/backend/apps/agents/manager/run/TurnRunner.py

"""The streaming turn + capacity-retry, lifted out of run_agent_loop so agent_manager stays under
the file ceiling. Faithful relocation as a mixin method (self.sessions / self.live_partial resolve
across the MRO unchanged); turn/thinking are created by the caller and passed in so the loop's
except-handlers can still read them after a mid-stream failure."""

import asyncio
import logging
import time
from typing import Dict, List, Union
from typeguard import typechecked

from backend.apps.agents.core.models import AgentSession
from backend.apps.agents.core.ws_manager import ws_manager
from backend.apps.agents.core.error_classify import CAPACITY_BACKOFFS, capacity_retry_wait
from backend.apps.agents.manager.streaming.state import ThinkingState, TurnState
from backend.apps.agents.manager.streaming.handle_stream_event import handle_stream_event
from backend.apps.agents.manager.streaming.handle_assistant_message import handle_assistant_message
from backend.apps.agents.manager.streaming.handle_result_message import handle_result_message
from backend.apps.agents.manager.streaming import thinking as thinking_mod
from backend.apps.settings.models import AppSettings

logger = logging.getLogger(__name__)


from backend.apps.agents.manager.AgentManagerProtocol import AgentManagerProtocol


class TurnRunner(AgentManagerProtocol):
    # `options` is the SDK ClaudeAgentOptions, lazy-imported below (so mock-mode can import the manager without the SDK present), so it's left unannotated; everything else is typed.
    @typechecked
    async def run_turn_with_retry(self, session: AgentSession, session_id: str,
                                    prompt_content: Union[str, List], options,
                                    options_kwargs: Dict, turn: TurnState, thinking: ThinkingState,
                                    p_stderr_buffer: List[str], resolved_model: str, api_type: str,
                                    global_settings: AppSettings) -> None:
        from claude_agent_sdk import query, ClaudeAgentOptions, AssistantMessage, ResultMessage
        from claude_agent_sdk.types import StreamEvent, SystemMessage

        async def prompt_stream():
            yield {
                "type": "user",
                "message": {"role": "user", "content": prompt_content},
            }

        async def p_run_streaming_turn():
            # Per-turn thinking aggregation trackers (added for the "Thought for Ns · M tokens" persisted label). Without nonlocal, the int reassignments at AssistantMessage emission below shadow them as locals and the dict access at content_block_start crashes with UnboundLocalError.
            async for message in query(
                prompt=prompt_stream(),
                options=options,
            ):
                if isinstance(message, ResultMessage):
                    turn.current_turn_emitted = False
                else:
                    turn.current_turn_emitted = True
                    # Stamp the turn's wall-clock start at the FIRST non-Result message we see, this is when the user actually started waiting. We use the same timestamp as the basis for "Thought for Ns" so the duration covers thinking + tool exec + assistant text generation.
                    if turn.started_ts is None:
                        turn.started_ts = time.time()
                        # 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_fresh", 0) or 0)
                                turn.baseline_session_out = int(session.tokens.get("output", 0) or 0)
                            p_ch_in = 0
                            p_ch_out = 0
                            for p_child in self.sessions.values():
                                if getattr(p_child, "parent_session_id", None) != session.id:
                                    continue
                                p_ct = getattr(p_child, "tokens", None)
                                if not isinstance(p_ct, dict):
                                    continue
                                p_ch_in += int(p_ct.get("input_fresh", 0) or 0)
                                p_ch_out += int(p_ct.get("output", 0) or 0)
                            turn.baseline_children_in = p_ch_in
                            turn.baseline_children_out = p_ch_out
                            turn.baseline_captured = True
                        except Exception:
                            pass
                        # Pre-emit thinking pill for routes whose translator strips reasoning content (cx/, gc/, ag/, gemini/). Without this, the pill emits at turn end and lands BELOW the assistant text in session.messages, visually wrong. Pre-emitting here gives the pill the same ordering as Anthropic's natural streaming path. Updates in place at turn end via the stable thinking.msg_id dedupe.
                        try:
                            p_route_strips_reasoning_pre = (
                                isinstance(resolved_model, str)
                                and resolved_model.startswith(("cx/", "gc/", "ag/", "gemini/"))
                            )
                            if p_route_strips_reasoning_pre:
                                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")

                if turn.first_event:
                    logger.info(f"[MCP-DEBUG] First event received: {type(message).__name__}")
                    turn.first_event = False

                # Log system messages (MCP server status, errors, etc.)
                if isinstance(message, SystemMessage):
                    raw = message.__dict__ if hasattr(message, '__dict__') else str(message)
                    logger.info(f"[MCP-DEBUG] SystemMessage: {raw}")

                if isinstance(message, StreamEvent):
                    await handle_stream_event(
                        message, session, session_id, turn, thinking, self.live_partial
                    )

                elif isinstance(message, AssistantMessage):
                    await handle_assistant_message(
                        message, session, session_id, turn, thinking, self.live_partial, self.sessions
                    )
                elif isinstance(message, ResultMessage):
                    await handle_result_message(
                        message, session, session_id, turn, thinking, self.sessions,
                        resolved_model, api_type, global_settings,
                    )

        capacity_retry_attempt = 0
        while True:
            try:
                await p_run_streaming_turn()
                break
            except Exception as e:
                # Make sure the consolidated-thinking ticker doesn't outlive the turn on error/retry. Without this, an exception mid-stream leaves a dangling task that keeps re-emitting against a stale msg id.
                if thinking.ticker_task is not None and not thinking.ticker_task.done():
                    thinking.ticker_task.cancel()
                    try:
                        await thinking.ticker_task
                    except (asyncio.CancelledError, Exception):
                        pass
                thinking.ticker_task = None
                stderr_snapshot = "\n".join(p_stderr_buffer[-50:])
                wait = capacity_retry_wait(e, capacity_retry_attempt, extra_text=stderr_snapshot)
                if wait is not None:
                    capacity_retry_attempt += 1
                    mid_stream = turn.current_turn_emitted
                    logger.warning(
                        f"Transient upstream error on session {session_id} "
                        f"(attempt {capacity_retry_attempt}/{len(CAPACITY_BACKOFFS)}, "
                        f"mid_stream={mid_stream}); sleeping {wait}s before retry. "
                        f"exc={e!r} stderr_tail={stderr_snapshot[-400:]!r}"
                    )
                    # Finalize any in-flight stream messages so the UI doesn't leave them pinned as "still streaming" while we wait and restart. On resume the CLI re-runs the last turn from scratch (Anthropic doesn't persist in-progress responses), so the partial assistant text / tool call we emitted is now orphaned, cap it with stream_end and start the fresh turn under a new message id.
                    if turn.stream_text_msg_id:
                        await ws_manager.send_to_session(session_id, "agent:stream_end", {
                            "session_id": session_id,
                            "message_id": turn.stream_text_msg_id,
                        })
                        turn.stream_text_msg_id = None
                    turn.stream_text_accum = ""
                    self.live_partial.pop(session_id, None)
                    for p_tool_msg_id in turn.stream_tool_msg_ids_ordered:
                        await ws_manager.send_to_session(session_id, "agent:stream_end", {
                            "session_id": session_id,
                            "message_id": p_tool_msg_id,
                        })
                    turn.stream_tool_msg_ids_ordered = []
                    turn.stream_block_index_map = {}
                    turn.current_turn_emitted = False
                    await asyncio.sleep(wait)
                    p_stderr_buffer.clear()
                    if session.sdk_session_id:
                        options_kwargs["resume"] = session.sdk_session_id
                        options = ClaudeAgentOptions(**options_kwargs)
                    continue
                raise