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, cast
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, is_router_unreachable_error
from backend.apps.agents.core import flight_recorder
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 TurnResultError, handle_result_message
from backend.apps.agents.manager.streaming.note_provider_retry import note_provider_retry, settle_provider_retries
from backend.apps.agents.manager.run.client_pool import (
    SdkClientLike,
    acquire_client,
    boot_fingerprint,
    dispose_client,
    persistent_client_enabled,
)
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, force_respawn: bool = False) -> 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(p_stream=None):
            # 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.
            # p_stream lets the persistent-client path feed receive_response() through this same consumption loop (one body, two transports).
            async for message in (p_stream if p_stream is not None else query(prompt=prompt_stream(), options=options)):
                # MCPActivate tells the model its new tools are not callable yet and to stop. Asking
                # was not enough: it kept going and guessed names like `send email`, which is what
                # made every MCP task look broken. The activated tools genuinely do not exist until
                # the transport is rebuilt, so end the turn here and let the auto-continuation fire
                # with the real names. Checked before the message is handled, so the model's next
                # move after activating never runs.
                if getattr(session, "pending_continuation", False):
                    break
                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__}")
                    flight_recorder.crumb(session_id, "first-event", kind=type(message).__name__)
                    turn.first_event = False

                # Active-time ledger: count the gap since the previous event, capped so a stall
                # (approval wait, paused workflow, provider backoff) can't book its wall-clock.
                p_now_ts = time.time()
                if turn.last_event_ts is not None:
                    turn.active_ms += int(min(p_now_ts - turn.last_event_ts, 30.0) * 1000)
                turn.last_event_ts = p_now_ts

                # 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}")
                    p_subtype = getattr(message, "subtype", "")
                    if p_subtype == "compact_boundary":
                        turn.compact_boundaries += 1
                    elif p_subtype == "api_retry":
                        note_provider_retry(session_id, raw, turn)

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

                elif isinstance(message, AssistantMessage):
                    flight_recorder.crumb(session_id, "assistant-msg")
                    await handle_assistant_message(
                        message, session, session_id, turn, thinking, self.live_partial, self.sessions
                    )
                elif isinstance(message, ResultMessage):
                    flight_recorder.crumb(session_id, "result-msg", subtype=str(getattr(message, "subtype", "")))
                    await handle_result_message(
                        message, session, session_id, turn, thinking, self.sessions,
                        resolved_model, api_type, global_settings,
                    )

        async def p_run_streaming_turn_persistent():
            from claude_agent_sdk import ClaudeSDKClient

            async def p_connect():
                p_client = ClaudeSDKClient(options=options)
                logger.info(f"[SPAWN-PHASE] cli-connect start session={session_id[:8]} t={time.monotonic():.3f}")
                flight_recorder.crumb(session_id, "cli-connect-start")
                await p_client.connect()
                logger.info(f"[SPAWN-PHASE] cli-connect done session={session_id[:8]} t={time.monotonic():.3f}")
                flight_recorder.crumb(session_id, "cli-connect-done")
                return p_client

            fp = boot_fingerprint(options_kwargs, session)
            logger.info(f"[SPAWN-PHASE] client-acquire start session={session_id[:8]} t={time.monotonic():.3f}")
            flight_recorder.crumb(session_id, "client-acquire")
            handle = await acquire_client(
                self.client_pool, session_id, fp, p_connect, force_respawn=force_respawn,
            )
            async with handle.lock:
                handle.turns_served += 1
                try:
                    sdk = cast(SdkClientLike, handle.client)
                    await sdk.query(prompt_stream())
                    await p_run_streaming_turn(p_stream=sdk.receive_response())
                    # LRU by turn-END so a session mid-long-turn isn't first cap-evicted the instant it finishes.
                    handle.last_used = time.monotonic()
                except BaseException:
                    # Fail-safe: an error or stop mid-turn poisons the live conversation; drop the client so the next attempt/turn reconnects fresh (== today's one-shot behavior, never worse). Pool pop is sync-first, so even a cancelled disconnect can't leave a reusable stale handle.
                    await dispose_client(self.client_pool, session_id)
                    raise

        async def p_finalize_interrupted_stream():
            # 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

        p_use_persistent = persistent_client_enabled()
        capacity_retry_attempt = 0
        p_router_retry_attempt = 0
        # Baseline crumb so even a first-call failure's envelope names the turn it died in.
        flight_recorder.crumb(session_id, "turn-start", model=resolved_model, api=api_type)
        while True:
            try:
                if p_use_persistent:
                    await p_run_streaming_turn_persistent()
                else:
                    await p_run_streaming_turn()
                # The near-miss ledger: a turn that needed retries and still finished is a net that
                # FIRED, and "how often do the nets fire" needs a denominator in analytics.
                if p_router_retry_attempt or capacity_retry_attempt:
                    flight_recorder.record_recovery(
                        session_id,
                        net="router-resume" if p_router_retry_attempt else "transient-backoff",
                        model=resolved_model,
                        attempts=p_router_retry_attempt + capacity_retry_attempt,
                        sessions=self.sessions,
                    )
                settle_provider_retries(session_id, turn, resolved_model, self.sessions)
                break
            except TurnResultError as p_result_err:
                # "Unable to connect" in a turn result is the CLI failing to reach our own localhost
                # router, which a dev reload kills and the watchdog revives within seconds. The CLI
                # transcript keeps the tools that already ran, so a resume continues the SAME
                # conversation without re-executing side effects: re-ensure the router, resume, go.
                if p_router_retry_attempt < 2 and is_router_unreachable_error(str(p_result_err)):
                    p_router_retry_attempt += 1
                    flight_recorder.crumb(session_id, "router-retry", attempt=p_router_retry_attempt, err=str(p_result_err)[:160])
                    logger.warning(
                        f"Router unreachable mid-turn on session {session_id} "
                        f"(attempt {p_router_retry_attempt}/2); re-ensuring router and resuming. "
                        f"err={p_result_err!s}"
                    )
                    try:
                        from backend.apps.nine_router.process import ensure_running
                        await ensure_running()
                    except Exception:
                        logger.exception("Router re-ensure failed; resuming anyway after the wait")
                    await p_finalize_interrupted_stream()
                    await asyncio.sleep(2.0 if p_router_retry_attempt == 1 else 5.0)
                    p_stderr_buffer.clear()
                    if session.sdk_session_id:
                        options_kwargs["resume"] = session.sdk_session_id
                        options = ClaudeAgentOptions(**options_kwargs)
                    continue
                # Any other error-shaped result: the CLI already ran the whole turn (tools executed) and then reported failure; a resume-retry would re-execute side effects, so this goes straight to the error card.
                raise
            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)
                # Persistent-client fail-safe: a dead/wedged CLI raises a connection-class error that the capacity classifier won't retry. The client is already disposed (see p_run_streaming_turn_persistent), so ONE immediate retry reconnects fresh == today's cold behavior; a second failure surfaces normally.
                if wait is None and p_use_persistent and capacity_retry_attempt == 0 and not turn.current_turn_emitted:
                    p_name = type(e).__name__
                    if "CLIConnection" in p_name or "ProcessError" in p_name or "Transport" in p_name:
                        logger.warning(f"[client-pool] {session_id}: dead client ({p_name}); one transparent respawn retry")
                        wait = 0.0
                if wait is not None:
                    capacity_retry_attempt += 1
                    flight_recorder.crumb(session_id, "transient-retry", attempt=capacity_retry_attempt, wait=wait, err=str(e)[:160])
                    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}"
                    )
                    await p_finalize_interrupted_stream()
                    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