"""
DialogueRunner executes dialogue nodes with user interaction.

This module implements the core agentic loop that:
1. Streams LLM completions (text and tool calls)
2. Executes tools and handles their responses
3. Manages agent handoffs with auto-continue for seamless transitions
4. Classifies tools as "display" or "save" to determine loop behavior

## Tool Classification (Key Concept)

Tools are classified into two types based on their content_type and actions:

- **Display Tools**: Require user interaction before continuing
  - INTERACTIVE content_type (always)
  - RICH content_type WITH actions array (user must click something)

- **Save Tools**: Allow the loop to continue autonomously
  - TEXT content_type
  - RICH content_type WITHOUT actions (e.g., Plan display)

This classification replaced the previous `force_text_response` flag which caused
bugs when agents were handed off after RICH tool responses.

## Auto-Continue Flow

When a save tool triggers an agent handoff:
1. Tool executes and returns agent reference
2. Loop continues with new agent
3. New agent generates response (may be text-only initially)
4. If no tool calls AND just_switched_agent=True:
   - Inject synthetic "continue" user message
   - Loop continues, giving new agent another turn with tools
5. This allows workflow chains like: Agent A -> save_tool -> Agent B -> plan_tool

See docs/architecture/AUTO_CONTINUE.md for detailed flow diagrams.
"""

import copy
import json
import logging
from collections import defaultdict
from dataclasses import dataclass, field
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Dict,
    List,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
)
from uuid import uuid4

from openai.types.chat.chat_completion_message_tool_call import (
    ChatCompletionMessageToolCall,
    Function,
)

from swarm.agent import Agent
from swarm.chunk import (
    AssistantMessageChunk,
    ChunkType,
    DelimiterChunk,
    ExecutionMetadata,
    TokenUsage,
    ToolCallChunk,
    ToolResponseChunk,
    ToolResponseContent,
)
from swarm.config.protocols import (
    MessageStoreProtocol,
    MetricsCollectorProtocol,
    StreamManagerProtocol,
)
from swarm.context import ContextProtocol, RunContext
from swarm.guardrails import (
    GuardrailResult,
    OutputGuardrailContext,
    _extract_last_assistant_text,
    run_input_guardrails,
    run_output_guardrails,
)
from swarm.hooks import fire_hook_on_llm_end, fire_hook_on_llm_start
from swarm.message import SwarmAssistantMessage, SwarmMessage, SwarmToolMessage
from swarm.providers.types import CompletionRequest, ToolCallDelta
from swarm.tool import (
    ToolExecutionResponse,
    execute_tools,
    prepare_tools_for_completion,
)
from swarm.types import ContentType
from swarm.util import extract_interactive_display_text, merge_chunk

logger = logging.getLogger(__name__)


@dataclass
class _LoopState:
    """Encapsulates all mutable state for the run_streamed loop.

    Using a dataclass instead of separate variables:
    - Makes state dependencies explicit
    - Simplifies passing state to helper methods
    - Prevents accidental variable shadowing
    """

    # Conversation history (mix of dicts from DB and SwarmMessage objects)
    message_history: List[Union[Dict[str, Any], SwarmMessage]]

    # UUID for the current assistant message being generated
    current_message_uuid: str

    # Length of message_history at loop start (for max_turns calculation)
    initial_len: int

    # The currently active agent (may change during handoffs)
    active_agent: Agent

    # System prompt for the active agent
    system_prompt: str

    # Context variables passed between tools and agents
    context_variables: Dict[str, Any]

    # Flag: True after agent handoff, triggers auto-continue if LLM returns text-only
    # This enables seamless agent transitions without user intervention
    just_switched_agent: bool = False

    # Token usage from the current streaming completion (reset before each call)
    current_call_usage: Optional[TokenUsage] = None

    # Stop-stream: track last assistant tool_calls and completed tool_call_ids
    # Used to compute dangling tool calls on stop for history finalization
    last_tool_calls: List[ChatCompletionMessageToolCall] = field(default_factory=list)
    completed_tool_call_ids: Set[str] = field(default_factory=set)

    # Observability metrics for the execution
    execution_metadata: Dict[str, Any] = field(
        default_factory=lambda: {
            "agent_transitions": 0,  # Number of agent handoffs
            "tool_executions": 0,  # Total tool calls executed
            "thinking_turns": 0,  # LLM responses without tool calls
            "agent_completed": False,  # Agent signaled completion
            "agent_paused": False,  # Agent paused for user input
            "completion_reason": None,  # Why the loop ended
        }
    )


class DialogueRunner:
    """
    Executes dialogue nodes with user interaction.
    """

    def __init__(
        self,
        metrics_collector: Optional[MetricsCollectorProtocol] = None,
    ) -> None:
        """
        Initialize the dialogue runner.

        Args:
            metrics_collector: Optional metrics collector for observability
        """
        self.metrics_collector = metrics_collector

    async def run_streamed(
        self,
        agent: Agent,
        run_context: RunContext[ContextProtocol],
        max_turns: int = 25,
        is_new_conversation: bool = False,
        check_stop_fn: Optional[Callable[[], Awaitable[bool]]] = None,
        heartbeat_fn: Optional[Callable[[], Awaitable[None]]] = None,
    ) -> AsyncIterator[ChunkType]:
        """Run the agentic loop, streaming chunks back to the client.

        This is the main entry point for dialogue execution. It implements an
        autonomous loop that:
        1. Generates LLM responses (streaming text and tool calls)
        2. Executes tools when the LLM calls them
        3. Classifies tool responses as "display" or "save"
        4. Handles agent handoffs with auto-continue for seamless transitions

        Args:
            agent: The Swarm agent to run (may change during handoffs)
            run_context: Runtime context with dependencies (message store, router, etc.)
            max_turns: Maximum messages to add before stopping (safety limit)
            is_new_conversation: Whether this is a fresh conversation
            check_stop_fn: Optional async callable returning True if stop was signalled
            heartbeat_fn: Optional async callable to refresh the running-key TTL

        Yields:
            ChunkType: Streaming chunks (DelimiterChunk, AssistantMessageChunk,
                      ToolCallChunk, ToolResponseChunk)

        Loop Exit Conditions:
            - Display tool executed (wait for user interaction)
            - No tool calls AND not just_switched_agent (conversation complete)
            - Agent signaled completion/pause
            - max_turns reached
            - Stop signal detected via check_stop_fn callback
        """
        # Extract dependencies from RunContext
        context = run_context.deps
        stream_manager = run_context.stream_manager
        message_store = run_context.message_store
        router = run_context.router

        # Initialize loop state with conversation history
        context_history = context.get_history_for_completion()
        message_history: List[Union[Dict[str, Any], SwarmMessage]] = list(
            context_history
        )

        state = _LoopState(
            message_history=message_history,
            current_message_uuid=str(uuid4()),
            initial_len=len(message_history),
            active_agent=agent,
            system_prompt=(
                agent.instructions(context)
                if callable(agent.instructions)
                else agent.instructions
            ),
            context_variables=copy.deepcopy(context.current_context_variables),
        )

        logger.info(
            "dialogue_runner.run_streamed",
            extra={
                "agent": state.active_agent.name,
                "messages": state.message_history,
                "location": "dialogue_runner.run_streamed",
                "context_variables": state.context_variables,
                "message_count": state.initial_len,
            },
        )

        await stream_manager.publish_status("Thinking...")

        # Run input guardrails before any LLM call
        await run_input_guardrails(
            agent, list(state.message_history), state.context_variables
        )

        # Track the agent that produced the last output — updated each
        # iteration, used for output guardrails after the loop exits.
        # Initialised here so it's defined even when max_turns=0.
        output_guardrail_agent = agent

        # ┌─────────────────────────────────────────────────────────────────────┐
        # │                         MAIN AGENTIC LOOP                          │
        # │                                                                     │
        # │  Each iteration:                                                   │
        # │  1. Generate LLM completion (stream text/tool calls)               │
        # │  2. If no tool calls: check auto-continue or break                 │
        # │  3. Execute tools and classify as display/save                     │
        # │  4. Display tools: break (wait for user)                           │
        # │  5. Save tools: continue loop (LLM responds to tool result)        │
        # └─────────────────────────────────────────────────────────────────────┘
        while len(state.message_history) - state.initial_len < max_turns:
            # ── Safe Point: check stop + heartbeat ────────────────────────
            if check_stop_fn is not None and await check_stop_fn():
                state.execution_metadata["completion_reason"] = "stopped"
                break
            if heartbeat_fn is not None:
                await heartbeat_fn()

            # Initialize accumulator for streaming response chunks
            message = self._init_message_accumulator(state)
            request = self._build_completion_request(state, state.system_prompt)

            # Reset per-call usage before each streamed completion
            state.current_call_usage = None

            # Fire LLM lifecycle hooks
            await fire_hook_on_llm_start(
                state.active_agent,
                list(state.message_history),
                state.context_variables,
            )

            # Stream LLM completion with try/finally for hook guarantees
            llm_response: Optional[SwarmAssistantMessage] = None
            llm_error: BaseException | None = None
            yield DelimiterChunk(delim="start")
            await stream_manager.publish_status("Generating response...")

            try:
                async for chunk in router.stream_completion(request):
                    for yielded in self._process_stream_chunk(chunk, message, state):
                        yield yielded
            except BaseException as e:
                llm_error = e
                raise
            finally:
                # Finalize assistant message inside finally so llm_response is
                # available for the hook even when the stream succeeds.
                if llm_error is None:
                    llm_response = self._finalize_assistant_message(message, state)
                yield DelimiterChunk(delim="end")
                await fire_hook_on_llm_end(
                    state.active_agent, llm_response, error=llm_error
                )

            # Update snapshot: this is the agent that produced this iteration's
            # output. Tool execution below may handoff to a different agent.
            output_guardrail_agent = state.active_agent

            # Persist finalized assistant message to database
            # llm_response is guaranteed non-None here — if the stream had
            # failed, the re-raise above would have prevented reaching this line.
            assert llm_response is not None
            await self._persist_messages(
                [llm_response],
                context,
                state.context_variables,
                state.active_agent,
                message_store,
            )

            # ── No Tool Calls Path ────────────────────────────────────────────
            # If LLM generated text only (no tool calls):
            # - If just_switched_agent: inject "continue" and loop again
            # - Otherwise: conversation turn is complete, break
            if not message["tool_calls"]:
                should_continue = await self._handle_no_tool_calls(
                    state, context, message_store
                )
                if should_continue:
                    continue
                break

            # ── Tool Execution Path ───────────────────────────────────────────
            await stream_manager.publish_status("Processing actions")
            tool_calls = self._build_tool_calls(message["tool_calls"])

            # Track tool calls for stop-stream finalization
            state.last_tool_calls = tool_calls
            state.completed_tool_call_ids = set()

            # ── Safe Point: check stop before tool execution ──────────────
            if check_stop_fn is not None and await check_stop_fn():
                state.execution_metadata["completion_reason"] = "stopped"
                break
            if heartbeat_fn is not None:
                await heartbeat_fn()

            tool_response = await execute_tools(
                tool_calls,
                state.active_agent.tools,
                run_context,
                current_agent_name=state.active_agent.name,
                hooks=state.active_agent.hooks,
                agent=state.active_agent,
                check_stop_fn=check_stop_fn,
            )

            # Track completed tool call IDs for dangling detection
            for msg in tool_response.messages:
                state.completed_tool_call_ids.add(msg.tool_call_id)

            # Add tool messages to history and persist immediately
            state.message_history.extend(tool_response.messages)
            self._update_execution_metadata(state, tool_response)

            if tool_response.messages:
                await self._persist_messages(
                    tool_response.messages,
                    context,
                    state.context_variables,
                    state.active_agent,
                    message_store,
                )

            # Sync context variables so subsequent tools see updated values
            state.context_variables = tool_response.context_variables.copy()
            run_context.deps.workflow_context.update(tool_response.context_variables)

            # ── Stop check after tool execution ───────────────────────────────
            if tool_response.stopped or (
                check_stop_fn is not None and await check_stop_fn()
            ):
                state.execution_metadata["completion_reason"] = "stopped"
                break

            # ── Tool Classification ───────────────────────────────────────────
            # Classify tool as "display" (needs user interaction) or "save" (can continue)
            tool_msg = self._get_first_tool_message(tool_response.messages)
            tool_type = self._classify_tool_type(tool_msg) if tool_msg else "save"

            logger.info(
                "dialogue_runner.run_streamed.tool_classified",
                extra={
                    "location": "dialogue_runner.run_streamed",
                    "agent": state.active_agent.name,
                    "message_id": state.current_message_uuid,
                    "tool_type": tool_type,
                    "content_type": tool_msg.content_type if tool_msg else None,
                },
            )

            # ── Display Tool: Break Loop ──────────────────────────────────────
            # Display tools (INTERACTIVE or RICH with actions) require user input.
            # Stream the display text and tool response, then exit the loop.
            if tool_type == "display" and tool_msg is not None:
                async for chunk in self._handle_display_tool(
                    tool_msg,
                    tool_response,
                    state,
                    context,
                    message_store,
                    stream_manager,
                ):
                    yield chunk
                break  # Wait for user interaction

            # ── Save Tool: Continue Loop ──────────────────────────────────────
            # Save tools (TEXT or RICH without actions) don't need user input.
            # Stream the tool response, handle any agent switch, then continue.
            tool_contents = self._build_tool_contents(tool_response.messages)
            if tool_contents:
                yield ToolResponseChunk(tool_response={"messages": tool_contents})

            # Handle agent handoff if tool returned a new agent
            if tool_response.agent:
                await self._handle_agent_switch(tool_response.agent, state, context)
            else:
                # No handoff - clear the flag to prevent spurious auto-continue
                state.just_switched_agent = False

            # Check if agent signaled completion or pause
            should_exit, reason = self._should_exit_loop(tool_response)
            if should_exit:
                if reason:
                    state.execution_metadata["completion_reason"] = reason
                logger.info(
                    "dialogue_runner.run_streamed.exiting_on_signal",
                    extra={
                        "location": "dialogue_runner.run_streamed",
                        "message_id": state.current_message_uuid,
                        "agent": state.active_agent.name,
                        "completion_reason": reason,
                    },
                )
                break

            # Generate new UUID for the next iteration's assistant message
            state.current_message_uuid = str(uuid4())

        # If completion_reason not set, loop exited due to max_turns
        if state.execution_metadata["completion_reason"] is None:
            state.execution_metadata["completion_reason"] = "max_turns"

        stopped = state.execution_metadata["completion_reason"] == "stopped"

        # Signal completion of the conversation turn
        logger.info(
            "swarm.core.run_and_stream.complete",
            extra={
                "agent_name": state.active_agent.name,
                "location": "swarm.core.run_and_stream",
                "message_id": state.current_message_uuid,
                "message_count": len(state.message_history) - state.initial_len,
                "stopped": stopped,
            },
        )

        # Run output guardrails (advisory-only — content already streamed)
        # Skip guardrails on stop — partial output may not be meaningful.
        output_guardrail_results: list[GuardrailResult] = []
        if not stopped and getattr(output_guardrail_agent, "output_guardrails", None):
            agent_output = _extract_last_assistant_text(state.message_history)
            output_ctx = OutputGuardrailContext(
                agent_name=output_guardrail_agent.name,
                agent_output=agent_output,
                context_variables=state.context_variables,
            )
            output_guardrail_results = await run_output_guardrails(
                output_guardrail_agent, output_ctx
            )
            for gr in output_guardrail_results:
                if gr.tripwire_triggered:
                    logger.warning(
                        "dialogue_runner.output_guardrail_tripped_advisory",
                        extra={
                            "guardrail_name": gr.guardrail_name,
                            "output_info": gr.output_info,
                            "location": "dialogue_runner.run_streamed",
                        },
                    )

        # Compute dangling tool calls for stop finalization
        dangling_tool_calls: list[ChatCompletionMessageToolCall] = []
        if stopped and state.last_tool_calls:
            dangling_tool_calls = [
                tc
                for tc in state.last_tool_calls
                if tc.id not in state.completed_tool_call_ids
            ]

        # Log execution summary for observability
        logger.info(
            "dialogue_runner.execution_summary",
            extra={
                "location": "dialogue_runner.run_streamed",
                "agent": state.active_agent.name,
                "execution_metadata": state.execution_metadata,
                "description": "Autonomous loop execution completed",
                "dangling_tool_calls": len(dangling_tool_calls),
            },
        )

        execution_metadata = ExecutionMetadata(
            **state.execution_metadata,
            output_guardrail_results=output_guardrail_results,
            dangling_tool_calls=dangling_tool_calls,
        )
        yield DelimiterChunk(
            delim="stopped" if stopped else "complete",
            execution_metadata=execution_metadata,
        )

    def _init_message_accumulator(self, state: _LoopState) -> Dict[str, Any]:
        """Initialize an accumulator dict for streaming LLM response chunks.

        The accumulator collects partial content and tool calls as they stream in.
        Using defaultdict for tool_calls allows incremental building of each tool call
        as chunks arrive (arguments are streamed in pieces).

        Args:
            state: Current loop state containing agent info and message UUID.

        Returns:
            Dict with keys: id, content, sender, role, function_call, tool_calls.
            The tool_calls value is a defaultdict for incremental building.
        """
        return {
            "id": state.current_message_uuid,
            "content": "",
            "sender": state.active_agent.id,
            "role": "assistant",
            "function_call": None,
            "tool_calls": defaultdict(
                lambda: {
                    "function": {"arguments": "", "name": ""},
                    "id": "",
                    "type": "",
                }
            ),
        }

    def _process_stream_chunk(
        self, chunk: Any, message: Dict[str, Any], state: _LoopState
    ) -> List[ChunkType]:
        """Process a single stream chunk and return chunks to yield.

        Handles both content chunks (text) and tool call chunks. Content chunks
        become AssistantMessageChunk, tool calls become ToolCallChunk. The chunk
        is also merged into the message accumulator.

        Args:
            chunk: Raw chunk from the LLM completion stream.
            message: Accumulator dict being built up across chunks.
            state: Current loop state for agent name and message UUID.

        Returns:
            List of ChunkType to yield (0-1 items typically).
        """
        self._capture_chunk_usage(chunk, state)

        if not chunk.choices:
            return []

        delta = self._parse_chunk_delta(chunk)
        if delta is None:
            return []

        delta["id"] = state.current_message_uuid
        delta["sender"] = state.active_agent.id
        chunks = self._build_stream_output_chunks(delta, state)
        self._merge_delta_into_message(message, delta)
        return chunks

    def _capture_chunk_usage(self, chunk: Any, state: _LoopState) -> None:
        """Capture token usage from streaming chunks when available."""
        if not (hasattr(chunk, "usage") and chunk.usage):
            return

        usage = TokenUsage(
            model=self._resolve_usage_model(chunk, state),
            input_tokens=getattr(chunk.usage, "prompt_tokens", 0) or 0,
            output_tokens=getattr(chunk.usage, "completion_tokens", 0) or 0,
        )
        state.current_call_usage = usage
        state.execution_metadata.setdefault("accumulated_usage", []).append(usage)

    @staticmethod
    def _resolve_usage_model(chunk: Any, state: _LoopState) -> str:
        """Resolve a single model string for usage records."""
        chunk_model = getattr(chunk, "model", None)
        if isinstance(chunk_model, str) and chunk_model:
            return chunk_model
        if isinstance(state.active_agent.model, list) and state.active_agent.model:
            return state.active_agent.model[0]
        if isinstance(state.active_agent.model, str):
            return state.active_agent.model
        return "unknown"

    def _parse_chunk_delta(self, chunk: Any) -> Optional[Dict[str, Any]]:
        """Parse JSON delta payload from a streaming chunk."""
        try:
            parsed = json.loads(chunk.choices[0].delta.json())
            if isinstance(parsed, dict):
                return parsed
            logger.warning(
                "dialogue_runner.invalid_chunk_delta",
                extra={"type": type(parsed).__name__},
            )
            return None
        except json.JSONDecodeError as e:
            logger.warning(
                "dialogue_runner.json_decode_error",
                extra={"error": str(e), "chunk": str(chunk)},
            )
            return None

    def _build_stream_output_chunks(
        self, delta: Dict[str, Any], state: _LoopState
    ) -> List[ChunkType]:
        """Build stream output chunks from parsed delta."""
        if delta.get("tool_calls"):
            return [self._create_tool_call_chunk(delta, state)]
        if delta.get("content"):
            delta["role"] = "assistant"
            return [AssistantMessageChunk(**delta)]
        return []

    @staticmethod
    def _merge_delta_into_message(
        message: Dict[str, Any], delta: Dict[str, Any]
    ) -> None:
        """Merge a parsed delta payload into the message accumulator."""
        clean_delta = {
            key: value
            for key, value in delta.items()
            if key not in {"role", "sender", "id"}
        }
        merge_chunk(message, clean_delta)

    def _create_tool_call_chunk(
        self, delta: Dict[str, Any], state: _LoopState
    ) -> ToolCallChunk:
        """Create a ToolCallChunk from delta.

        Args:
            delta: Parsed delta dict containing tool_calls array.
            state: Current loop state for message UUID and history.

        Returns:
            ToolCallChunk with tool call deltas and parent message ID.
        """
        parent_id = None
        if state.message_history:
            last = state.message_history[-1]
            parent_id = (
                last.get("id") if isinstance(last, dict) else getattr(last, "id", None)
            )

        return ToolCallChunk(
            id=state.current_message_uuid,
            tool_calls=[ToolCallDelta(**tc) for tc in delta["tool_calls"]],
            sender=state.active_agent.id,
            parent_id=parent_id,
        )

    def _finalize_assistant_message(
        self, message: Dict[str, Any], state: _LoopState
    ) -> SwarmAssistantMessage:
        """Finalize message accumulator and add to history.

        Converts the tool_calls defaultdict to a list, creates a SwarmAssistantMessage,
        and appends it to the message history.

        Args:
            message: Accumulated message dict from streaming.
            state: Current loop state containing message history.

        Returns:
            The finalized SwarmAssistantMessage that was added to history.
        """
        message["tool_calls"] = list(message.get("tool_calls", {}).values())
        if not message["tool_calls"]:
            message["tool_calls"] = None

        # Attach per-call usage to the message if present
        if state.current_call_usage is not None:
            message["usage"] = state.current_call_usage

        designer_message = SwarmAssistantMessage(**message)
        state.message_history.append(designer_message)
        return designer_message

    def _build_tool_calls(
        self, tool_calls_raw: List[Dict[str, Any]]
    ) -> List[ChatCompletionMessageToolCall]:
        """Convert raw tool calls to OpenAI ChatCompletionMessageToolCall format.

        Args:
            tool_calls_raw: List of tool call dicts from the message accumulator.

        Returns:
            List of ChatCompletionMessageToolCall objects for tool execution.
        """
        return [
            ChatCompletionMessageToolCall(
                id=tc["id"],
                function=Function(
                    arguments=tc["function"]["arguments"],
                    name=tc["function"]["name"],
                ),
                type=tc["type"],
            )
            for tc in tool_calls_raw
        ]

    def _update_execution_metadata(
        self, state: _LoopState, tool_response: ToolExecutionResponse
    ) -> None:
        """Update execution metadata from tool response.

        Increments tool_executions count and sets agent_completed/agent_paused
        flags based on tool response metadata.

        Args:
            state: Loop state containing execution_metadata dict.
            tool_response: Response from tool execution with metadata.
        """
        tool_count = len(
            [msg for msg in tool_response.messages if isinstance(msg, SwarmToolMessage)]
        )
        state.execution_metadata["tool_executions"] += tool_count

        if tool_response.metadata.agent_completed:
            state.execution_metadata["agent_completed"] = True
        if tool_response.metadata.agent_paused:
            state.execution_metadata["agent_paused"] = True

    def _get_first_tool_message(
        self, messages: Sequence[SwarmMessage]
    ) -> Optional[SwarmToolMessage]:
        """Get the first tool message from the list.

        Args:
            messages: Sequence of SwarmMessage (may include assistant messages).

        Returns:
            First SwarmToolMessage found, or None if no tool messages exist.
        """
        for msg in messages:
            if isinstance(msg, SwarmToolMessage):
                return msg
        return None

    async def _handle_display_tool(
        self,
        tool_msg: SwarmToolMessage,
        tool_response: ToolExecutionResponse,
        state: _LoopState,
        context: ContextProtocol,
        message_store: MessageStoreProtocol,
        stream_manager: StreamManagerProtocol,
    ) -> AsyncIterator[ChunkType]:
        """Handle display tool: extract text, stream word-by-word, then break loop.

        Display tools (INTERACTIVE or RICH with actions) require user interaction
        before the workflow can continue. This method:

        1. Extracts human-readable text from component's data.text field
        2. For RICH tools: streams text word-by-word as deltas (like LLM streaming)
           For INTERACTIVE tools: skips text streaming (component has title/text fields)
        3. Persists the text message for LLM context
        4. Publishes status and streams the tool response (for UI component rendering)
        5. Sets completion_reason to "display_tool"

        The loop breaks after this method returns, waiting for user input.

        Note: Even display tools can trigger agent handoffs. The new agent will
        be active when the user next interacts, but won't auto-continue since
        the loop is breaking here.

        Args:
            tool_msg: The tool message with component data.
            tool_response: Full tool execution response.
            state: Current loop state.
            context: Conversation context.
            message_store: Message persistence store.
            stream_manager: Stream manager for publishing status updates.

        Yields:
            ChunkType: DelimiterChunk, AssistantMessageChunk (word-by-word deltas),
                      and ToolResponseChunk.
        """
        is_interactive = tool_msg.content_type == ContentType.INTERACTIVE

        # Extract display text: use fallback chain for INTERACTIVE, simple extraction for RICH
        display_text: Optional[str]
        if is_interactive:
            content_json = (
                tool_msg.content[0]["text"]
                if tool_msg.content and len(tool_msg.content) > 0
                else ""
            )
            display_text = extract_interactive_display_text(
                content_json, tool_name=tool_msg.tool_name
            )
        else:
            display_text = self._extract_display_text(tool_msg)

        if is_interactive or display_text:
            text_msg_uuid = str(uuid4())

            if not is_interactive and display_text:
                # RICH display tools: stream text word-by-word (unchanged)
                yield DelimiterChunk(delim="start")

                words = display_text.split(" ")
                for i, word in enumerate(words):
                    content = f" {word}" if i > 0 else word
                    yield AssistantMessageChunk(
                        id=text_msg_uuid,
                        content=content,
                        sender=state.active_agent.id,
                        role="assistant",
                    )

                yield DelimiterChunk(delim="end")

            # Always persist the text message for LLM context
            persist_text = display_text or ""
            text_message = SwarmAssistantMessage(
                id=text_msg_uuid,
                content=[{"type": "text", "text": persist_text}],
                sender=state.active_agent.id,
                role="assistant",
            )
            state.message_history.append(text_message)
            await self._persist_messages(
                [text_message],
                context,
                state.context_variables,
                state.active_agent,
                message_store,
                is_client_message_override=False if is_interactive else None,
            )

        # Publish status before streaming tool response
        await stream_manager.publish_status("Designing interactive component...")

        # Stream tool response
        tool_contents = self._build_tool_contents(tool_response.messages)
        if tool_contents:
            yield ToolResponseChunk(tool_response={"messages": tool_contents})

        # Handle agent switch even for display tools
        if tool_response.agent:
            await self._handle_agent_switch(tool_response.agent, state, context)

        state.execution_metadata["completion_reason"] = "display_tool"

    def _classify_tool_type(self, msg: SwarmToolMessage) -> str:
        """Classify tool message as 'display' or 'save'.

        This classification determines whether the loop should:
        - break (display): Tool shows UI that needs user interaction
        - continue (save): Tool completed work, LLM should respond to result

        Classification Rules:
        - INTERACTIVE → display (always needs user input, e.g., form)
        - RICH with actions → display (has buttons/links user must click)
        - RICH without actions → save (informational display, e.g., Plan card)
        - TEXT → save (simple result, continue loop)

        This replaces the old `force_text_response` flag which incorrectly
        disabled tools for the next agent after any RICH response.

        Args:
            msg: The tool message to classify.

        Returns:
            Either "display" or "save" string.
        """
        if msg.content_type == ContentType.INTERACTIVE:
            return "display"
        elif msg.content_type == ContentType.RICH:
            return "display" if self._has_user_actions(msg) else "save"
        return "save"

    def _has_user_actions(self, msg: SwarmToolMessage) -> bool:
        """Check if tool response has user actions.

        Parses the tool message content as JSON and checks for a non-empty
        'actions' array in the response.

        Args:
            msg: The tool message to check.

        Returns:
            True if actions array exists and has items, False otherwise.
        """
        try:
            content = json.loads(msg.content[0]["text"]) if msg.content else {}
            return len(content.get("actions", [])) > 0
        except (json.JSONDecodeError, KeyError, IndexError):
            return False

    def _extract_display_text(self, msg: SwarmToolMessage) -> Optional[str]:
        """Extract text from component data for display streaming.

        Only extracts the 'text' field (not title) since title is already displayed
        as part of the UI component rendering. The text provides the conversational
        context that gets streamed to the chat.

        Args:
            msg: The tool message containing component data.

        Returns:
            The text field value if present and parseable, None otherwise.
        """
        try:
            content = json.loads(msg.content[0]["text"]) if msg.content else {}
            data = content.get("data", {})
            text = data.get("text")
            return str(text) if text is not None else None
        except (json.JSONDecodeError, KeyError, IndexError):
            return None

    def _build_completion_request(
        self, state: _LoopState, system_prompt: str
    ) -> CompletionRequest:
        """Build LLM completion request with messages and tools.

        Constructs the full request including system prompt, message history,
        and available tools from the active agent.

        Args:
            state: Loop state with active_agent, message_history, and model settings.
            system_prompt: The system prompt for the active agent.

        Returns:
            CompletionRequest ready to send to the model router.
        """
        models = (
            state.active_agent.model
            if isinstance(state.active_agent.model, list)
            else [state.active_agent.model]
        )

        # Always include tools (no force_text_response needed)
        include_tools = bool(state.active_agent.tools)

        return CompletionRequest(
            model=models,
            messages=[
                {"role": "system", "content": system_prompt},
                *[
                    msg if isinstance(msg, dict) else msg.to_openai_dict()
                    for msg in state.message_history
                ],
            ],
            tools=prepare_tools_for_completion(state.active_agent.tools)
            if include_tools
            else None,
            **(
                state.active_agent.model_settings.to_completion_params(
                    has_tools=include_tools
                )
                if state.active_agent.model_settings
                else {}
            ),
        )

    def _build_tool_contents(
        self, messages: Sequence[SwarmMessage]
    ) -> List[ToolResponseContent]:
        """Convert tool messages to ToolResponseContent for streaming.

        Filters for SwarmToolMessage instances and converts them to the
        ToolResponseContent format used by ToolResponseChunk.

        Args:
            messages: Sequence of SwarmMessage (filters for tool messages only).

        Returns:
            List of ToolResponseContent objects for streaming to frontend.
        """
        return [
            ToolResponseContent(
                id=msg.id,
                role=msg.role,
                tool_call_id=msg.tool_call_id,
                tool_name=msg.tool_name,
                content=(
                    msg.content[0]["text"]
                    if msg.content and len(msg.content) > 0
                    else ""
                ),
                component_type=msg.component_type,
                content_type=msg.content_type,
                is_error=msg.is_error,
                metadata=msg.metadata,
            )
            for msg in messages
            if isinstance(msg, SwarmToolMessage)
        ]

    async def _handle_agent_switch(
        self,
        new_agent: Agent,
        state: _LoopState,
        context: ContextProtocol,
    ) -> None:
        """Update state for agent transition.

        Sets the new active agent, regenerates system prompt, sets the
        just_switched_agent flag for auto-continue, and increments the
        agent_transitions counter.

        Args:
            new_agent: The agent to switch to.
            state: Loop state to update.
            context: Context for regenerating system prompt.
        """
        state.active_agent = new_agent
        state.just_switched_agent = True
        state.execution_metadata["agent_transitions"] += 1

        # Regenerate system_prompt for the new agent
        state.system_prompt = (
            new_agent.instructions(context)
            if callable(new_agent.instructions)
            else new_agent.instructions
        )

        prompt_preview = state.system_prompt.strip().split("\n")[0][:100]
        logger.info(
            "dialogue_runner.run_streamed.switched_to_agent",
            extra={
                "description": f"Switched to agent {new_agent.name}",
                "location": "dialogue_runner.run_streamed",
                "agent_name": new_agent.name,
                "system_prompt_preview": prompt_preview,
                "message_id": state.current_message_uuid,
                "context_variables": state.context_variables,
            },
        )

    async def _handle_no_tool_calls(
        self,
        state: _LoopState,
        context: ContextProtocol,
        message_store: MessageStoreProtocol,
    ) -> bool:
        """Handle case when LLM generates text without tool calls.

        This is the core of the auto-continue mechanism for seamless agent handoffs.

        Scenario:
        1. Agent A calls save_discovery_details which returns Plan + next agent
        2. Agent B is now active, generates greeting text (no tool calls)
        3. Without auto-continue, loop would break here
        4. With auto-continue: inject "continue" message, loop again
        5. Agent B now calls plan_tasks (with tool access!)

        Args:
            state: Loop state with just_switched_agent flag and message history.
            context: Context for conversation ID.
            message_store: Store for persisting the synthetic continue message.

        Returns:
            True if loop should continue (auto-continue triggered),
            False if loop should break (conversation turn complete).
        """
        from swarm.message import SwarmUserMessage

        state.execution_metadata["thinking_turns"] += 1

        if state.just_switched_agent:
            # Auto-continue: Agent just switched, give it another turn with tools
            # This synthetic message appears in conversation history as a "continue"
            # from the user, allowing the new agent to proceed with its workflow.
            user_msg = SwarmUserMessage(
                id=str(uuid4()),
                role="user",
                content=[{"type": "text", "text": "continue"}],
                metadata={"synthetic": True, "auto_continue": True},
            )
            state.message_history.append(user_msg)

            await self._persist_messages(
                [user_msg],
                context,
                state.context_variables,
                state.active_agent,
                message_store,
            )

            state.just_switched_agent = False
            state.current_message_uuid = str(uuid4())
            logger.info(
                "dialogue_runner.auto_continue_after_agent_switch",
                extra={
                    "location": "dialogue_runner.run_streamed",
                    "message_id": state.current_message_uuid,
                    "agent": state.active_agent.name,
                },
            )
            return True  # Continue loop - give new agent another turn

        # No agent switch pending - this is a natural conversation end
        state.execution_metadata["completion_reason"] = "no_tool_calls"
        logger.info(
            "dialogue_runner.run_streamed.ending_turn",
            extra={
                "location": "dialogue_runner.run_streamed",
                "message_id": state.current_message_uuid,
                "agent": state.active_agent.name,
            },
        )
        return False  # Break loop - conversation turn complete

    def _should_exit_loop(
        self, tool_response: ToolExecutionResponse
    ) -> Tuple[bool, Optional[str]]:
        """Check if loop should exit based on tool response signals.

        Checks agent_paused and agent_completed flags in tool response metadata.
        Exit if paused, or if completed without a next agent handoff.

        Args:
            tool_response: The tool execution response with metadata.

        Returns:
            Tuple of (should_exit, completion_reason). Reason is "agent_completed",
            "agent_paused", or None.
        """
        # Exit if paused OR (completed AND no next agent)
        should_exit = tool_response.metadata.agent_paused or (
            tool_response.metadata.agent_completed and not tool_response.agent
        )

        if not should_exit:
            return (False, None)

        if tool_response.metadata.agent_completed:
            return (True, "agent_completed")
        elif tool_response.metadata.agent_paused:
            return (True, "agent_paused")
        return (True, None)

    async def _persist_messages(
        self,
        messages: Sequence[SwarmMessage],
        context: ContextProtocol,
        context_variables: Dict[str, Any],
        agent: Agent,
        message_store: MessageStoreProtocol,
        is_client_message_override: Optional[bool] = None,
    ) -> None:
        """
        Persist messages to storage (resilient against stream failures).

        Args:
            messages: List of messages to persist
            context: The dialogue context (provides conversation_id and designer_fullname)
            context_variables: Current context variables
            agent: The active agent
            message_store: Storage protocol implementation
            is_client_message_override: When not None, forces is_client_message
                on assistant messages after transformer conversion.
        """
        try:
            # Safely extract designer model
            if isinstance(agent.model, str):
                designer_model = agent.model
            elif isinstance(agent.model, list) and len(agent.model) > 0:
                designer_model = str(agent.model[0])
            else:
                designer_model = "unknown"

            # Extract billing context from deps (if available)
            raw_user_id = getattr(context, "user_id", None)

            await message_store.persist_messages(
                messages=messages,
                conversation_id=str(context.conversation_id),
                context_variables=context_variables,
                designer_model=designer_model,
                designer_fullname=context.designer_fullname,
                user_id=str(raw_user_id) if raw_user_id is not None else None,
                audience=getattr(context, "audience", None),
                node_key=getattr(context, "node_key", None),
                is_client_message_override=is_client_message_override,
            )

            logger.info(
                "dialogue_runner.messages_persisted",
                extra={
                    "conversation_id": str(context.conversation_id),
                    "message_count": len(messages),
                    "message_types": [msg.role for msg in messages],
                    "location": "dialogue_runner._persist_messages",
                },
            )
        except Exception as e:
            logger.error(
                "dialogue_runner.persist_messages_failed",
                extra={
                    "error": str(e),
                    "conversation_id": str(context.conversation_id),
                    "message_count": len(messages),
                    "location": "dialogue_runner._persist_messages",
                },
                exc_info=True,
            )
            # Continue despite persistence failure - streaming is more important
