motosan_agent_loop/

loop_.rs

use std::collections::HashMap;
use std::sync::Arc;

use futures::future::{join_all, try_join_all};
use motosan_agent_tool::{Tool, ToolContext, ToolDef, ToolResult};
use tokio::sync::mpsc;

use crate::context::ContextProvider;
use crate::error::AgentError;
use crate::llm::{LlmClient, TokenUsage, ToolCallItem};
use crate::message::{Message, ToolCallRef};
use crate::Result;

/// Policy for handling channel backpressure when a bounded queue is full.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum BackpressurePolicy {
    /// Block the sender until capacity is available (default).
    #[default]
    Block,
    /// Drop the **incoming** (newest) op when the queue is full.
    ///
    /// Note: true drop-oldest semantics would require a ring-buffer channel.
    /// This variant currently behaves like [`Reject`] but emits [`AgentEvent::OpDropped`]
    /// instead of [`AgentEvent::OpRejected`], allowing callers to distinguish intent.
    DropOldest,
    /// Reject the new item and emit a telemetry event.
    Reject,
}

/// Configuration for bounded channel capacities used by [`AgentSession`].
///
/// All capacities have sensible defaults (64) and can be overridden via the
/// builder pattern on [`AgentLoopBuilder`].
#[derive(Debug, Clone)]
pub struct ChannelConfig {
    /// Capacity of the user-input channel (messages sent via `AgentSession::send`).
    pub input_capacity: usize,
    /// Capacity of the per-turn operations channel (ops like Interrupt, InjectHint).
    pub ops_capacity: usize,
    /// Backpressure policy applied when the ops channel is full.
    pub ops_backpressure: BackpressurePolicy,
}

impl Default for ChannelConfig {
    fn default() -> Self {
        Self {
            input_capacity: 64,
            ops_capacity: 64,
            ops_backpressure: BackpressurePolicy::Block,
        }
    }
}

// ---------------------------------------------------------------------------
// Pipeline stage helpers
// ---------------------------------------------------------------------------

/// Resolved tool registry for a single run, combining static and dynamic tools.
///
/// This always clones the base maps so it can own the data. When `extra_tools`
/// is empty the clone is the only cost; when extra tools are present they are
/// inserted into the cloned maps.
struct MergedTools {
    map: HashMap<String, Arc<dyn Tool>>,
    defs: Vec<ToolDef>,
}

impl MergedTools {
    /// Clone base maps and merge any extra tools into them.
    ///
    /// Note: this always clones `base_map` and `base_defs` because `MergedTools`
    /// owns its data. If `extra_tools` is empty the clone is the only cost.
    fn new(
        base_map: &HashMap<String, Arc<dyn Tool>>,
        base_defs: &[ToolDef],
        extra_tools: &[Arc<dyn Tool>],
    ) -> Self {
        if extra_tools.is_empty() {
            return Self {
                map: base_map.clone(),
                defs: base_defs.to_vec(),
            };
        }
        let mut map = base_map.clone();
        let mut defs = base_defs.to_vec();
        for t in extra_tools {
            map.insert(t.def().name.clone(), Arc::clone(t));
            defs.push(t.def());
        }
        Self { map, defs }
    }

    fn tool_map(&self) -> &HashMap<String, Arc<dyn Tool>> {
        &self.map
    }

    fn tool_defs(&self) -> &[ToolDef] {
        &self.defs
    }
}

/// Accumulator for per-run mutable state shared across turn iterations.
struct TurnState {
    messages: Vec<Message>,
    all_tool_calls: Vec<(String, serde_json::Value)>,
    total_usage: TokenUsage,
}

impl TurnState {
    fn new(messages: Vec<Message>) -> Self {
        Self {
            messages,
            all_tool_calls: Vec::new(),
            total_usage: TokenUsage::default(),
        }
    }

    /// Record token usage from an LLM call.
    fn accumulate_usage(&mut self, usage: Option<TokenUsage>) {
        if let Some(u) = usage {
            self.total_usage.accumulate(u);
        }
    }

    /// Build the final result after the LLM produces a text answer.
    fn into_result(self, answer: String, iteration: usize) -> AgentResult {
        AgentResult {
            answer,
            tool_calls: self.all_tool_calls,
            iterations: iteration,
            usage: self.total_usage,
            messages: self.messages,
        }
    }
}

/// Stage: execute tool calls, emit events, and append messages.
///
/// This is the common "tool-call planning/execution + result merge" stage
/// extracted from every turn loop variant.
fn execute_and_record_tool_calls(
    items: &[ToolCallItem],
    results: Vec<ToolResult>,
    state: &mut TurnState,
    on_event: &(impl Fn(AgentEvent) + Send + Sync),
) {
    // Emit ToolCompleted events and record tool calls.
    for (tc, result) in items.iter().zip(results.iter()) {
        on_event(AgentEvent::ToolCompleted {
            name: tc.name.clone(),
            result: result.clone(),
        });
        state
            .all_tool_calls
            .push((tc.name.clone(), tc.args.clone()));
    }

    // Build a single assistant message with all tool call refs.
    let tool_call_refs: Vec<ToolCallRef> = items
        .iter()
        .map(|tc| ToolCallRef {
            id: tc.id.clone(),
            name: tc.name.clone(),
            args: tc.args.clone(),
        })
        .collect();
    state
        .messages
        .push(Message::assistant_with_tool_calls("", tool_call_refs));

    // Append individual tool result messages (order matches tool call order).
    for (tc, result) in items.iter().zip(results.iter()) {
        state
            .messages
            .push(Message::tool_result(&tc.id, &tool_result_to_string(result)));
    }
}

/// Stage: emit ToolStarted events for all items in a batch.
fn emit_tool_started(items: &[ToolCallItem], on_event: &(impl Fn(AgentEvent) + Send + Sync)) {
    for tc in items {
        on_event(AgentEvent::ToolStarted {
            name: tc.name.clone(),
        });
    }
}

/// Events emitted by the agent loop for observability.
#[derive(Debug, Clone)]
pub enum AgentEvent {
    /// A tool is about to be executed.
    ToolStarted { name: String },
    /// A tool execution completed.
    ToolCompleted { name: String, result: ToolResult },
    /// A text chunk from the LLM's streaming response.
    TextChunk(String),
    /// The full accumulated text after streaming completes.
    TextDone(String),
    /// Emitted at the start of each reasoning iteration.
    IterationStarted(usize),
    /// Loop was interrupted via `AgentOp::Interrupt`.
    Interrupted,
    /// Agent requested user input via the built-in `ask_user` tool.
    AskUser {
        call_id: String,
        question: String,
        options: Vec<String>,
    },
    /// A pending `ask_user` request timed out.
    AskUserTimeout { call_id: String, question: String },
    /// The ops channel reached its capacity limit.
    ///
    /// Emitted once when the queue becomes full, not on every blocked/rejected send.
    OpsSaturated {
        /// Current queue capacity.
        capacity: usize,
    },
    /// An operation was dropped due to backpressure policy.
    OpDropped {
        /// Description of the dropped operation.
        reason: String,
    },
    /// An operation was rejected due to backpressure policy.
    OpRejected {
        /// Description of the rejected operation.
        reason: String,
    },
}

/// Commands that can be sent into a running [`AgentLoop`].
#[derive(Debug, Clone)]
pub enum AgentOp {
    /// Stop the current turn at the next safe checkpoint.
    Interrupt,
    /// Append a user message before the next LLM iteration.
    InjectUserMessage(String),
    /// Append a user-visible note before the next LLM iteration.
    InjectHint(String),
    /// Answer a pending `ask_user` request.
    AskUserAnswer {
        call_id: Option<String>,
        answer: String,
    },
}

/// The final outcome produced by [`AgentLoop::run`].
#[derive(Debug, Clone)]
pub struct AgentResult {
    /// The assistant's final textual answer.
    pub answer: String,
    /// History of tool calls made: (tool_name, arguments).
    pub tool_calls: Vec<(String, serde_json::Value)>,
    /// Number of LLM round-trips performed.
    pub iterations: usize,
    /// Accumulated token usage across all LLM calls.
    pub usage: TokenUsage,
    /// Full conversation history including tool call/result pairs.
    ///
    /// Callers can pass this to a subsequent `run()` call to continue a
    /// multi-turn conversation.
    pub messages: Vec<Message>,
}

/// Builder for constructing an [`AgentLoop`] with validated configuration.
pub struct AgentLoopBuilder {
    tools: Vec<Arc<dyn Tool>>,
    context_providers: Vec<Box<dyn ContextProvider>>,
    max_iterations: usize,
    tool_timeout: Option<std::time::Duration>,
    tool_context: Option<ToolContext>,
    ask_user_enabled: bool,
    ask_user_timeout: Option<std::time::Duration>,
    channel_config: ChannelConfig,
    #[cfg(feature = "mcp-client")]
    mcp_servers: Vec<Arc<dyn crate::mcp::McpServer>>,
}

impl AgentLoopBuilder {
    /// Set the maximum number of LLM round-trips before aborting.
    pub fn max_iterations(mut self, n: usize) -> Self {
        self.max_iterations = n;
        self
    }

    /// Register a tool with the agent loop.
    pub fn tool(mut self, tool: Arc<dyn Tool>) -> Self {
        self.tools.push(tool);
        self
    }

    /// Register multiple tools at once.
    pub fn tools(mut self, tools: impl IntoIterator<Item = Arc<dyn Tool>>) -> Self {
        self.tools.extend(tools);
        self
    }

    /// Set a static system prompt that is injected as a `System` message
    /// at the beginning of every conversation.
    ///
    /// This is a convenience shortcut for registering a [`ContextProvider`]
    /// that always returns the given string.
    pub fn system_prompt(self, prompt: impl Into<String>) -> Self {
        self.context(crate::context::StringContextProvider(prompt.into()))
    }

    /// Register a context provider that injects dynamic context into the
    /// conversation before each [`AgentLoop::run`] call.
    ///
    /// Multiple providers can be registered; they are invoked in registration
    /// order and each non-empty result becomes a `System` message.
    pub fn context(mut self, provider: impl ContextProvider + 'static) -> Self {
        self.context_providers.push(Box::new(provider));
        self
    }

    /// Register multiple context providers at once.
    ///
    /// This is the batch equivalent of [`context()`](Self::context), mirroring
    /// the [`tools()`](Self::tools) / [`tool()`](Self::tool) pattern.
    pub fn contexts(
        mut self,
        providers: impl IntoIterator<Item = Box<dyn ContextProvider>>,
    ) -> Self {
        self.context_providers.extend(providers);
        self
    }

    /// Register an MCP server whose tools will be connected lazily on `run()`.
    ///
    /// Requires the `mcp-client` feature.
    #[cfg(feature = "mcp-client")]
    pub fn mcp_server(mut self, server: impl crate::mcp::McpServer + 'static) -> Self {
        self.mcp_servers.push(std::sync::Arc::new(server));
        self
    }

    /// Register a pre-shared `Arc<dyn McpServer>` whose tools will be connected
    /// lazily on `run()`.
    ///
    /// This is useful when the same MCP server instance must be shared between
    /// multiple [`AgentLoop`]s or retained by the caller.
    ///
    /// Requires the `mcp-client` feature.
    #[cfg(feature = "mcp-client")]
    pub fn mcp_server_arc(mut self, server: std::sync::Arc<dyn crate::mcp::McpServer>) -> Self {
        self.mcp_servers.push(server);
        self
    }

    /// Set a per-tool-call timeout.  When a tool takes longer than `duration`
    /// it returns a `ToolResult::error("timed out")` instead of blocking.
    ///
    /// By default there is no timeout.
    pub fn tool_timeout(mut self, duration: std::time::Duration) -> Self {
        self.tool_timeout = Some(duration);
        self
    }

    /// Set a custom [`ToolContext`] that will be passed to every tool invocation.
    ///
    /// By default, [`ToolContext::default()`] is used. Use this when tools need
    /// access to session-level state such as a working directory or environment
    /// variables.
    pub fn tool_context(mut self, ctx: ToolContext) -> Self {
        self.tool_context = Some(ctx);
        self
    }

    /// Set the full channel configuration for bounded queues.
    ///
    /// See [`ChannelConfig`] for details on each field.
    pub fn channel_config(mut self, config: ChannelConfig) -> Self {
        self.channel_config = config;
        self
    }

    /// Set the capacity of the user-input channel.
    ///
    /// Default: 64.
    pub fn input_channel_capacity(mut self, capacity: usize) -> Self {
        self.channel_config.input_capacity = capacity;
        self
    }

    /// Set the capacity of the per-turn operations channel.
    ///
    /// Default: 64.
    pub fn ops_channel_capacity(mut self, capacity: usize) -> Self {
        self.channel_config.ops_capacity = capacity;
        self
    }

    /// Set the backpressure policy for the operations channel.
    ///
    /// Default: [`BackpressurePolicy::Block`].
    pub fn ops_backpressure(mut self, policy: BackpressurePolicy) -> Self {
        self.channel_config.ops_backpressure = policy;
        self
    }

    /// Register the built-in `ask_user` tool with the default timeout (30s).
    pub fn with_ask_user(mut self) -> Self {
        self.ask_user_enabled = true;
        if self.ask_user_timeout.is_none() {
            self.ask_user_timeout = Some(std::time::Duration::from_secs(30));
        }
        self
    }

    /// Register the built-in `ask_user` tool with a custom timeout.
    pub fn with_ask_user_timeout(mut self, timeout: std::time::Duration) -> Self {
        self.ask_user_enabled = true;
        self.ask_user_timeout = Some(timeout);
        self
    }

    /// Consume the builder and produce an [`AgentLoop`].
    ///
    /// # Panics
    ///
    /// Panics if `input_capacity` or `ops_capacity` is zero, since
    /// `tokio::sync::mpsc::channel(0)` panics at runtime.
    pub fn build(self) -> AgentLoop {
        assert!(
            self.channel_config.input_capacity > 0,
            "input_capacity must be >= 1"
        );
        assert!(
            self.channel_config.ops_capacity > 0,
            "ops_capacity must be >= 1"
        );
        let tool_map: HashMap<String, Arc<dyn Tool>> = self
            .tools
            .iter()
            .map(|t| (t.def().name.clone(), Arc::clone(t)))
            .collect();
        let mut tool_defs: Vec<ToolDef> = self.tools.iter().map(|t| t.def()).collect();
        if self.ask_user_enabled {
            tool_defs.push(ask_user_tool_def());
        }
        AgentLoop {
            tool_map,
            tool_defs,
            context_providers: self.context_providers,
            max_iterations: self.max_iterations,
            tool_timeout: self.tool_timeout,
            tool_context: self.tool_context.unwrap_or_default(),
            ask_user_timeout: if self.ask_user_enabled {
                self.ask_user_timeout
            } else {
                None
            },
            channel_config: self.channel_config,
            #[cfg(feature = "mcp-client")]
            mcp_servers: self.mcp_servers,
        }
    }
}

/// The core ReAct agent loop.
///
/// Drives an LLM through iterative reasoning and tool execution until the
/// model produces a final text answer or the iteration limit is reached.
///
/// The loop does **not** own the LLM client; instead, `run()` takes
/// `&dyn LlmClient` so the same loop can be reused with different backends.
pub struct AgentLoop {
    /// Pre-built lookup map for static tools (excludes per-run MCP tools).
    tool_map: HashMap<String, Arc<dyn Tool>>,
    /// Pre-built tool definitions for static tools (excludes per-run MCP tools).
    tool_defs: Vec<ToolDef>,
    context_providers: Vec<Box<dyn ContextProvider>>,
    max_iterations: usize,
    /// Optional per-tool-call timeout.  When set, any tool that takes longer
    /// returns `ToolResult::error("timed out")` rather than blocking indefinitely.
    pub(crate) tool_timeout: Option<std::time::Duration>,
    /// Context passed to every tool invocation.
    tool_context: ToolContext,
    /// Optional timeout for the built-in `ask_user` tool.
    ask_user_timeout: Option<std::time::Duration>,
    /// Channel configuration for bounded queues and backpressure.
    channel_config: ChannelConfig,
    #[cfg(feature = "mcp-client")]
    mcp_servers: Vec<Arc<dyn crate::mcp::McpServer>>,
}

impl AgentLoop {
    /// Create a builder with default settings.
    pub fn builder() -> AgentLoopBuilder {
        AgentLoopBuilder {
            tools: Vec::new(),
            context_providers: Vec::new(),
            max_iterations: 10,
            tool_timeout: None,
            tool_context: None,
            ask_user_enabled: false,
            ask_user_timeout: None,
            channel_config: ChannelConfig::default(),
            #[cfg(feature = "mcp-client")]
            mcp_servers: Vec::new(),
        }
    }

    /// Returns the channel configuration for this agent loop.
    ///
    /// Used by [`AgentSession`](crate::AgentSession) to create bounded
    /// channels with the configured capacities and policies.
    pub fn channel_config(&self) -> &ChannelConfig {
        &self.channel_config
    }

    /// Run the agent loop to completion.
    ///
    /// Iteratively calls the LLM, executes any requested tools, and feeds
    /// tool results back into the conversation until the LLM produces a
    /// final text response or `max_iterations` is exceeded.
    ///
    /// The `on_event` callback is invoked for each notable event (tool
    /// started, tool completed, text chunk). Pass `|_| {}` for a no-op.
    pub async fn run(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> Result<AgentResult> {
        // Connect MCP servers and collect their tools (mcp-client feature).
        #[cfg(feature = "mcp-client")]
        let mcp_tools: Vec<Arc<dyn Tool>> = self.connect_mcp_servers().await?;
        #[cfg(not(feature = "mcp-client"))]
        let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

        let result = self.run_inner(llm, messages, &mcp_tools, &on_event).await;

        // Disconnect all MCP servers (best-effort) regardless of run outcome.
        #[cfg(feature = "mcp-client")]
        for server in &self.mcp_servers {
            let _ = server.disconnect().await;
        }

        result
    }

    /// Run with an optional interactive operations channel.
    ///
    /// Pass `ops_rx: None` for behavior equivalent to [`run`](Self::run).
    pub async fn run_with_ops(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        ops_rx: Option<mpsc::Receiver<AgentOp>>,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> Result<AgentResult> {
        #[cfg(feature = "mcp-client")]
        let mcp_tools: Vec<Arc<dyn Tool>> = self.connect_mcp_servers().await?;
        #[cfg(not(feature = "mcp-client"))]
        let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

        let result = self
            .run_inner_with_ops(llm, messages, &mcp_tools, &on_event, ops_rx)
            .await;

        #[cfg(feature = "mcp-client")]
        for server in &self.mcp_servers {
            let _ = server.disconnect().await;
        }

        result
    }

    /// Connect all MCP servers, collecting their tools.
    ///
    /// If any server fails to connect, all previously-connected servers are
    /// disconnected (best-effort) before the error is returned.
    #[cfg(feature = "mcp-client")]
    async fn connect_mcp_servers(&self) -> Result<Vec<Arc<dyn Tool>>> {
        use crate::mcp::adapter::McpToolAdapter;

        let mut connected: Vec<&Arc<dyn crate::mcp::McpServer>> = Vec::new();
        let mut tools: Vec<Arc<dyn Tool>> = Vec::new();

        for server in &self.mcp_servers {
            match server.connect().await {
                Ok(()) => {
                    connected.push(server);
                    match McpToolAdapter::from_server(Arc::clone(server)).await {
                        Ok(adapter_tools) => tools.extend(adapter_tools),
                        Err(e) => {
                            for s in &connected {
                                let _ = s.disconnect().await;
                            }
                            return Err(e);
                        }
                    }
                }
                Err(e) => {
                    for s in &connected {
                        let _ = s.disconnect().await;
                    }
                    return Err(e);
                }
            }
        }

        Ok(tools)
    }

    /// Run the agent loop to completion with cancellation support.
    ///
    /// Behaves identically to [`run`](Self::run) but accepts a
    /// [`CancellationToken`](tokio_util::sync::CancellationToken). When the
    /// token is cancelled the loop exits early with
    /// [`AgentError::Cancelled`].
    ///
    /// Cancellation is checked at the start of each iteration **and** races
    /// against the LLM call so an in-flight request is interrupted promptly.
    ///
    /// Requires the `cancellation` feature.
    #[cfg(feature = "cancellation")]
    pub async fn run_with_cancel(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        cancel: tokio_util::sync::CancellationToken,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> Result<AgentResult> {
        #[cfg(feature = "mcp-client")]
        let mcp_tools: Vec<Arc<dyn Tool>> = {
            use crate::mcp::adapter::McpToolAdapter;
            let mut tools = Vec::new();
            for server in &self.mcp_servers {
                server.connect().await?;
                tools.extend(McpToolAdapter::from_server(Arc::clone(server)).await?);
            }
            tools
        };
        #[cfg(not(feature = "mcp-client"))]
        let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

        let result = self
            .run_inner_cancel(llm, messages, &mcp_tools, &on_event, &cancel)
            .await;

        #[cfg(feature = "mcp-client")]
        for server in &self.mcp_servers {
            let _ = server.disconnect().await;
        }

        result
    }

    /// Streaming variant of [`run_with_cancel`](Self::run_with_cancel).
    ///
    /// Uses [`LlmClient::chat_stream`] and races each chunk against the
    /// cancellation token so the loop exits promptly when cancelled.
    ///
    /// Requires the `cancellation` feature.
    #[cfg(feature = "cancellation")]
    pub async fn run_streaming_with_cancel(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        cancel: tokio_util::sync::CancellationToken,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> Result<AgentResult> {
        use futures::StreamExt;

        #[cfg(feature = "mcp-client")]
        let mcp_tools: Vec<Arc<dyn Tool>> = {
            use crate::mcp::adapter::McpToolAdapter;
            let mut tools = Vec::new();
            for server in &self.mcp_servers {
                server.connect().await?;
                tools.extend(McpToolAdapter::from_server(Arc::clone(server)).await?);
            }
            tools
        };
        #[cfg(not(feature = "mcp-client"))]
        let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, &mcp_tools);
        let mut state = TurnState::new(self.prepare_messages(messages).await?);

        let result = async {
            for iteration in 1..=self.max_iterations {
                // Pre-turn: cancellation check.
                if cancel.is_cancelled() {
                    return Err(AgentError::Cancelled);
                }
                on_event(AgentEvent::IterationStarted(iteration));

                // Streaming LLM step with cancellation race.
                let (accumulated_text, response) = {
                    let mut stream = llm.chat_stream(&state.messages, tools.tool_defs());
                    let mut accumulated = String::new();
                    let mut final_response: Option<crate::LlmResponse> = None;

                    loop {
                        tokio::select! {
                            chunk_opt = stream.next() => {
                                match chunk_opt {
                                    Some(chunk_result) => {
                                        let chunk = chunk_result?;
                                        match chunk {
                                            crate::llm::StreamChunk::TextDelta(delta) => {
                                                accumulated.push_str(&delta);
                                                on_event(AgentEvent::TextChunk(delta));
                                            }
                                            crate::llm::StreamChunk::Done(resp) => {
                                                final_response = Some(resp);
                                            }
                                            crate::llm::StreamChunk::Usage(usage) => {
                                                state.total_usage.accumulate(usage);
                                            }
                                        }
                                    }
                                    None => break,
                                }
                            }
                            _ = cancel.cancelled() => {
                                return Err(AgentError::Cancelled);
                            }
                        }
                    }

                    let resp = final_response
                        .unwrap_or_else(|| crate::LlmResponse::Message(accumulated.clone()));
                    (accumulated, resp)
                };

                match response {
                    crate::LlmResponse::Message(text) => {
                        if accumulated_text.is_empty() && !text.is_empty() {
                            on_event(AgentEvent::TextChunk(text.clone()));
                        }
                        on_event(AgentEvent::TextDone(text.clone()));
                        return Ok(state.into_result(text, iteration));
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        emit_tool_started(&items, &on_event);
                        let results = Self::execute_tools_parallel(
                            tools.tool_map(),
                            &items,
                            self.tool_timeout,
                            &self.tool_context,
                        )
                        .await;
                        execute_and_record_tool_calls(&items, results, &mut state, &on_event);
                    }
                }
            }

            Err(AgentError::MaxIterations(self.max_iterations))
        }
        .await;

        #[cfg(feature = "mcp-client")]
        for server in &self.mcp_servers {
            let _ = server.disconnect().await;
        }

        result
    }

    /// Consume a streaming LLM response, forwarding text deltas as events
    /// and accumulating usage.  Returns the accumulated text and the final
    /// [`LlmResponse`](crate::LlmResponse).
    async fn consume_stream(
        llm: &dyn LlmClient,
        messages: &[Message],
        tool_defs: &[ToolDef],
        total_usage: &mut TokenUsage,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<(String, crate::LlmResponse)> {
        use futures::StreamExt;

        let mut stream = llm.chat_stream(messages, tool_defs);
        let mut accumulated = String::new();
        let mut final_response: Option<crate::LlmResponse> = None;

        while let Some(chunk_result) = stream.next().await {
            let chunk = chunk_result?;
            match chunk {
                crate::llm::StreamChunk::TextDelta(delta) => {
                    accumulated.push_str(&delta);
                    on_event(AgentEvent::TextChunk(delta));
                }
                crate::llm::StreamChunk::Done(response) => {
                    final_response = Some(response);
                }
                crate::llm::StreamChunk::Usage(usage) => {
                    total_usage.accumulate(usage);
                }
            }
        }

        let response =
            final_response.unwrap_or_else(|| crate::LlmResponse::Message(accumulated.clone()));
        Ok((accumulated, response))
    }

    /// Prepend context-provider messages to the conversation.
    ///
    /// Each non-empty result is inserted in registration order starting at
    /// index 0 (i.e. at the *beginning* of the conversation, as documented
    /// on [`ContextProvider`]).
    async fn prepare_messages(&self, mut messages: Vec<Message>) -> Result<Vec<Message>> {
        if self.context_providers.is_empty() {
            return Ok(messages);
        }
        let query: String = messages
            .iter()
            .rev()
            .find(|m| m.role == crate::message::Role::User)
            .map(|m| m.content.clone())
            .unwrap_or_default();

        let contexts = try_join_all(self.context_providers.iter().map(|p| p.build(&query))).await?;

        let mut insert_idx = 0;
        for ctx in contexts {
            if !ctx.is_empty() {
                messages.insert(insert_idx, Message::system(&ctx));
                insert_idx += 1;
            }
        }
        Ok(messages)
    }

    #[cfg(feature = "cancellation")]
    async fn run_inner_cancel(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        extra_tools: &[Arc<dyn Tool>],
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
        cancel: &tokio_util::sync::CancellationToken,
    ) -> Result<AgentResult> {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let mut state = TurnState::new(self.prepare_messages(messages).await?);

        for iteration in 1..=self.max_iterations {
            // Pre-turn: cancellation check.
            if cancel.is_cancelled() {
                return Err(AgentError::Cancelled);
            }
            on_event(AgentEvent::IterationStarted(iteration));

            // LLM step with cancellation race.
            let output = tokio::select! {
                output = llm.chat(&state.messages, tools.tool_defs()) => output?,
                _ = cancel.cancelled() => return Err(AgentError::Cancelled),
            };
            state.accumulate_usage(output.usage);

            match output.response {
                crate::LlmResponse::Message(text) => {
                    on_event(AgentEvent::TextChunk(text.clone()));
                    return Ok(state.into_result(text, iteration));
                }
                crate::LlmResponse::ToolCalls(items) => {
                    emit_tool_started(&items, on_event);
                    let results = Self::execute_tools_parallel(
                        tools.tool_map(),
                        &items,
                        self.tool_timeout,
                        &self.tool_context,
                    )
                    .await;
                    execute_and_record_tool_calls(&items, results, &mut state, on_event);
                }
            }
        }

        Err(AgentError::MaxIterations(self.max_iterations))
    }

    async fn run_inner_with_ops(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        extra_tools: &[Arc<dyn Tool>],
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
        mut ops_rx: Option<mpsc::Receiver<AgentOp>>,
    ) -> Result<AgentResult> {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let mut state = TurnState::new(self.prepare_messages(messages).await?);
        let mut ops_state = OpsState::default();

        for iteration in 1..=self.max_iterations {
            // Pre-turn: ingest pending ops.
            Self::drain_ops(&mut state.messages, &mut ops_rx, &mut ops_state);
            if ops_state.interrupted {
                on_event(AgentEvent::Interrupted);
                return Ok(
                    state.into_result("(interrupted)".to_string(), iteration.saturating_sub(1))
                );
            }

            on_event(AgentEvent::IterationStarted(iteration));

            // LLM step.
            let output = llm.chat(&state.messages, tools.tool_defs()).await?;
            state.accumulate_usage(output.usage);

            match output.response {
                crate::LlmResponse::Message(text) => {
                    on_event(AgentEvent::TextChunk(text.clone()));
                    state.messages.push(Message::assistant(&text));
                    return Ok(state.into_result(text, iteration));
                }
                crate::LlmResponse::ToolCalls(items) => {
                    // Tool-call execution stage (with ask_user policy).
                    emit_tool_started(&items, on_event);
                    let results = self
                        .execute_tools_with_policy(
                            tools.tool_map(),
                            &items,
                            &mut state.messages,
                            &mut ops_rx,
                            &mut ops_state,
                            on_event,
                        )
                        .await;
                    execute_and_record_tool_calls(&items, results, &mut state, on_event);
                }
            }
        }

        Err(AgentError::MaxIterations(self.max_iterations))
    }

    async fn wait_for_ask_user_answer(
        &self,
        call_id: &str,
        question: &str,
        messages: &mut Vec<Message>,
        ops_rx: &mut Option<mpsc::Receiver<AgentOp>>,
        ops_state: &mut OpsState,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> String {
        if let Some(answer) = pop_matching_answer(&mut ops_state.pending_answers, call_id) {
            return answer;
        }

        let timeout = self.ask_user_timeout;
        let started = tokio::time::Instant::now();

        loop {
            let next_op = if let Some(rx) = ops_rx.as_mut() {
                if let Some(limit) = timeout {
                    let elapsed = started.elapsed();
                    if elapsed >= limit {
                        on_event(AgentEvent::AskUserTimeout {
                            call_id: call_id.to_string(),
                            question: question.to_string(),
                        });
                        return String::new();
                    }
                    let remaining = limit - elapsed;
                    match tokio::time::timeout(remaining, rx.recv()).await {
                        Ok(op) => op,
                        Err(_) => {
                            on_event(AgentEvent::AskUserTimeout {
                                call_id: call_id.to_string(),
                                question: question.to_string(),
                            });
                            return String::new();
                        }
                    }
                } else {
                    rx.recv().await
                }
            } else {
                on_event(AgentEvent::AskUserTimeout {
                    call_id: call_id.to_string(),
                    question: question.to_string(),
                });
                return String::new();
            };

            let Some(op) = next_op else {
                on_event(AgentEvent::AskUserTimeout {
                    call_id: call_id.to_string(),
                    question: question.to_string(),
                });
                return String::new();
            };
            Self::apply_op(op, messages, ops_state);
            if ops_state.interrupted {
                return String::new();
            }

            if let Some(answer) = pop_matching_answer(&mut ops_state.pending_answers, call_id) {
                return answer;
            }
        }
    }

    async fn run_inner(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        extra_tools: &[Arc<dyn Tool>],
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<AgentResult> {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let mut state = TurnState::new(self.prepare_messages(messages).await?);

        for iteration in 1..=self.max_iterations {
            on_event(AgentEvent::IterationStarted(iteration));

            // LLM step.
            let output = llm.chat(&state.messages, tools.tool_defs()).await?;
            state.accumulate_usage(output.usage);

            match output.response {
                crate::LlmResponse::Message(text) => {
                    on_event(AgentEvent::TextChunk(text.clone()));
                    state.messages.push(Message::assistant(&text));
                    return Ok(state.into_result(text, iteration));
                }
                crate::LlmResponse::ToolCalls(items) => {
                    // Tool-call execution stage.
                    emit_tool_started(&items, on_event);
                    let results = Self::execute_tools_parallel(
                        tools.tool_map(),
                        &items,
                        self.tool_timeout,
                        &self.tool_context,
                    )
                    .await;
                    execute_and_record_tool_calls(&items, results, &mut state, on_event);
                }
            }
        }

        Err(AgentError::MaxIterations(self.max_iterations))
    }

    /// Streaming variant of [`run`](Self::run).
    ///
    /// Uses [`LlmClient::chat_stream`] so that `on_event` receives
    /// [`AgentEvent::TextChunk`] for each text delta as it arrives, plus
    /// [`AgentEvent::TextDone`] with the full accumulated text and
    /// [`AgentEvent::IterationStarted`] at each reasoning iteration.
    ///
    /// Tool execution is identical to the non-streaming path.
    pub async fn run_streaming(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> Result<AgentResult> {
        // Connect MCP servers and collect their tools (mcp-client feature).
        #[cfg(feature = "mcp-client")]
        let mcp_tools: Vec<Arc<dyn Tool>> = self.connect_mcp_servers().await?;
        #[cfg(not(feature = "mcp-client"))]
        let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, &mcp_tools);
        let mut state = TurnState::new(self.prepare_messages(messages).await?);

        let result = async {
            for iteration in 1..=self.max_iterations {
                on_event(AgentEvent::IterationStarted(iteration));

                // Streaming LLM step.
                let (accumulated_text, response) = Self::consume_stream(
                    llm,
                    &state.messages,
                    tools.tool_defs(),
                    &mut state.total_usage,
                    &on_event,
                )
                .await?;

                match response {
                    crate::LlmResponse::Message(text) => {
                        if accumulated_text.is_empty() && !text.is_empty() {
                            on_event(AgentEvent::TextChunk(text.clone()));
                        }
                        on_event(AgentEvent::TextDone(text.clone()));
                        state.messages.push(Message::assistant(&text));
                        return Ok(state.into_result(text, iteration));
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        emit_tool_started(&items, &on_event);
                        let results = Self::execute_tools_parallel(
                            tools.tool_map(),
                            &items,
                            self.tool_timeout,
                            &self.tool_context,
                        )
                        .await;
                        execute_and_record_tool_calls(&items, results, &mut state, &on_event);
                    }
                }
            }

            Err(AgentError::MaxIterations(self.max_iterations))
        }
        .await;

        // Disconnect all MCP servers (best-effort) regardless of run outcome.
        #[cfg(feature = "mcp-client")]
        for server in &self.mcp_servers {
            let _ = server.disconnect().await;
        }

        result
    }

    fn drain_ops(
        messages: &mut Vec<Message>,
        ops_rx: &mut Option<mpsc::Receiver<AgentOp>>,
        ops_state: &mut OpsState,
    ) {
        if let Some(rx) = ops_rx.as_mut() {
            while let Ok(op) = rx.try_recv() {
                Self::apply_op(op, messages, ops_state);
            }
        }
    }

    fn apply_op(op: AgentOp, messages: &mut Vec<Message>, ops_state: &mut OpsState) {
        match op {
            AgentOp::Interrupt => {
                ops_state.interrupted = true;
            }
            AgentOp::InjectUserMessage(text) => {
                messages.push(Message::user(&text));
            }
            AgentOp::InjectHint(hint) => {
                messages.push(Message::user(&format!("[Note: {hint}]")));
            }
            AgentOp::AskUserAnswer { call_id, answer } => {
                ops_state
                    .pending_answers
                    .push(PendingAskUserAnswer { call_id, answer });
            }
        }
    }

    async fn execute_tools_with_policy(
        &self,
        tool_map: &HashMap<String, Arc<dyn Tool>>,
        items: &[crate::llm::ToolCallItem],
        messages: &mut Vec<Message>,
        ops_rx: &mut Option<mpsc::Receiver<AgentOp>>,
        ops_state: &mut OpsState,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Vec<ToolResult> {
        let policy = ToolExecutionPolicy::from_items(items);
        match policy {
            ToolExecutionPolicy::ParallelOnly => {
                Self::execute_tools_parallel(tool_map, items, self.tool_timeout, &self.tool_context)
                    .await
            }
            ToolExecutionPolicy::InteractiveAskUser => {
                let non_ask_future = join_all(
                    items
                        .iter()
                        .enumerate()
                        .filter(|(_, tc)| tc.name != "ask_user")
                        .map(|(idx, tc)| async move {
                            (
                                idx,
                                Self::execute_tool(
                                    tool_map,
                                    &tc.name,
                                    tc.args.clone(),
                                    self.tool_timeout,
                                    &self.tool_context,
                                )
                                .await,
                            )
                        }),
                );

                let ask_future = async {
                    let mut ask_results = Vec::new();
                    for (idx, tc) in items
                        .iter()
                        .enumerate()
                        .filter(|(_, tc)| tc.name == "ask_user")
                    {
                        let (question, options) = parse_ask_user_args(&tc.args);
                        on_event(AgentEvent::AskUser {
                            call_id: tc.id.clone(),
                            question: question.clone(),
                            options,
                        });

                        let answer = self
                            .wait_for_ask_user_answer(
                                &tc.id, &question, messages, ops_rx, ops_state, on_event,
                            )
                            .await;
                        ask_results.push((idx, ToolResult::text(answer)));
                    }
                    ask_results
                };

                let (non_ask_results, ask_results) = futures::join!(non_ask_future, ask_future);
                let mut merged: Vec<Option<ToolResult>> = vec![None; items.len()];
                for (idx, result) in non_ask_results.into_iter().chain(ask_results) {
                    merged[idx] = Some(result);
                }
                merged
                    .into_iter()
                    .map(|result| result.expect("tool result must be present"))
                    .collect()
            }
        }
    }

    async fn execute_tools_parallel(
        tool_map: &HashMap<String, Arc<dyn Tool>>,
        items: &[crate::llm::ToolCallItem],
        timeout: Option<std::time::Duration>,
        ctx: &ToolContext,
    ) -> Vec<ToolResult> {
        join_all(
            items
                .iter()
                .map(|tc| Self::execute_tool(tool_map, &tc.name, tc.args.clone(), timeout, ctx)),
        )
        .await
    }

    /// Execute a single tool by name using a pre-built lookup map.
    /// Returns an error `ToolResult` if the tool is not found or times out.
    async fn execute_tool(
        tool_map: &HashMap<String, Arc<dyn Tool>>,
        name: &str,
        args: serde_json::Value,
        timeout: Option<std::time::Duration>,
        ctx: &ToolContext,
    ) -> ToolResult {
        let fut = async {
            if let Some(tool) = tool_map.get(name) {
                tool.call(args, &ctx).await
            } else {
                ToolResult::error(format!("unknown tool: {name}"))
            }
        };
        if let Some(dur) = timeout {
            match tokio::time::timeout(dur, fut).await {
                Ok(result) => result,
                Err(_) => ToolResult::error(format!("tool '{name}' timed out after {dur:?}")),
            }
        } else {
            fut.await
        }
    }
}

#[derive(Debug, Clone, Copy)]
enum ToolExecutionPolicy {
    ParallelOnly,
    InteractiveAskUser,
}

impl ToolExecutionPolicy {
    fn from_items(items: &[crate::llm::ToolCallItem]) -> Self {
        if items.iter().any(|tc| tc.name == "ask_user") {
            Self::InteractiveAskUser
        } else {
            Self::ParallelOnly
        }
    }
}

#[derive(Default)]
struct OpsState {
    interrupted: bool,
    pending_answers: Vec<PendingAskUserAnswer>,
}

struct PendingAskUserAnswer {
    call_id: Option<String>,
    answer: String,
}

fn pop_matching_answer(pending: &mut Vec<PendingAskUserAnswer>, call_id: &str) -> Option<String> {
    let pos = pending
        .iter()
        .position(|item| item.call_id.as_deref() == Some(call_id) || item.call_id.is_none())?;
    Some(pending.remove(pos).answer)
}

fn parse_ask_user_args(args: &serde_json::Value) -> (String, Vec<String>) {
    let question = args
        .get("question")
        .and_then(|value| value.as_str())
        .unwrap_or("")
        .to_string();
    let options = args
        .get("options")
        .and_then(|value| value.as_array())
        .map(|list| {
            list.iter()
                .filter_map(|entry| entry.as_str().map(ToString::to_string))
                .collect::<Vec<_>>()
        })
        .unwrap_or_default();
    (question, options)
}

fn ask_user_tool_def() -> ToolDef {
    ToolDef {
        name: "ask_user".to_string(),
        description: "Ask the user a question and wait for a reply.".to_string(),
        input_schema: serde_json::json!({
            "type": "object",
            "properties": {
                "question": { "type": "string" },
                "options": {
                    "type": "array",
                    "items": { "type": "string" }
                }
            },
            "required": ["question"]
        }),
    }
}

/// Convert a [`ToolResult`] into a string suitable for message content.
fn tool_result_to_string(result: &ToolResult) -> String {
    match result.as_text() {
        Some(text) => text.to_string(),
        None => {
            // Fall back to JSON serialization of the content.
            serde_json::to_string(&result.content).unwrap_or_else(|_| "<no content>".to_string())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::context::ContextProvider;
    use crate::llm::{ChatOutput, LlmResponse, TokenUsage, ToolCallItem};
    use async_trait::async_trait;
    use std::sync::{Arc, Mutex};

    // -----------------------------------------------------------------------
    // Mock LLM client
    // -----------------------------------------------------------------------

    /// A mock LLM client that returns pre-programmed responses in sequence.
    struct MockLlm {
        responses: Mutex<Vec<LlmResponse>>,
        usage_per_call: Option<TokenUsage>,
    }

    impl MockLlm {
        fn new(responses: Vec<LlmResponse>) -> Self {
            Self {
                responses: Mutex::new(responses),
                usage_per_call: None,
            }
        }

        fn with_usage(responses: Vec<LlmResponse>, usage: TokenUsage) -> Self {
            Self {
                responses: Mutex::new(responses),
                usage_per_call: Some(usage),
            }
        }
    }

    #[async_trait]
    impl LlmClient for MockLlm {
        async fn chat(
            &self,
            _messages: &[Message],
            _tools: &[ToolDef],
        ) -> crate::Result<ChatOutput> {
            let mut responses = self.responses.lock().unwrap();
            if responses.is_empty() {
                panic!("MockLlm: no more responses");
            }
            let response = responses.remove(0);
            Ok(ChatOutput {
                response,
                usage: self.usage_per_call,
            })
        }
    }

    // -----------------------------------------------------------------------
    // Mock tool
    // -----------------------------------------------------------------------

    struct MockTool {
        name: String,
        result_text: String,
    }

    impl MockTool {
        fn new(name: &str, result_text: &str) -> Self {
            Self {
                name: name.to_string(),
                result_text: result_text.to_string(),
            }
        }
    }

    impl Tool for MockTool {
        fn def(&self) -> ToolDef {
            ToolDef {
                name: self.name.clone(),
                description: format!("Mock tool: {}", self.name),
                input_schema: serde_json::json!({
                    "type": "object",
                    "properties": {
                        "input": { "type": "string" }
                    },
                    "required": []
                }),
            }
        }

        fn call(
            &self,
            _args: serde_json::Value,
            _ctx: &ToolContext,
        ) -> std::pin::Pin<Box<dyn std::future::Future<Output = ToolResult> + Send + '_>> {
            let text = self.result_text.clone();
            Box::pin(async move { ToolResult::text(text) })
        }
    }

    // -----------------------------------------------------------------------
    // Tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn direct_message_response() {
        let llm = MockLlm::new(vec![LlmResponse::Message("Hello!".to_string())]);

        let agent = AgentLoop::builder().build();
        let result = agent
            .run(&llm, vec![Message::user("Hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "Hello!");
        assert!(result.tool_calls.is_empty());
        assert_eq!(result.iterations, 1);
    }

    #[tokio::test]
    async fn single_tool_call_then_message() {
        let llm = MockLlm::new(vec![
            LlmResponse::single_tool_call(
                "call_1".to_string(),
                "search".to_string(),
                serde_json::json!({"input": "rust"}),
            ),
            LlmResponse::Message("Found results about Rust.".to_string()),
        ]);

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "Rust is a systems language"));

        let agent = AgentLoop::builder().tool(tool).build();
        let result = agent
            .run(&llm, vec![Message::user("Search for rust")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "Found results about Rust.");
        assert_eq!(result.tool_calls.len(), 1);
        assert_eq!(result.tool_calls[0].0, "search");
        assert_eq!(result.iterations, 2);
    }

    #[tokio::test]
    async fn parallel_tool_calls() {
        let llm = MockLlm::new(vec![
            LlmResponse::ToolCalls(vec![
                ToolCallItem {
                    id: "call_1".to_string(),
                    name: "search".to_string(),
                    args: serde_json::json!({"input": "a"}),
                },
                ToolCallItem {
                    id: "call_2".to_string(),
                    name: "fetch".to_string(),
                    args: serde_json::json!({"input": "b"}),
                },
            ]),
            LlmResponse::Message("Combined results.".to_string()),
        ]);

        let search: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result_a"));
        let fetch: Arc<dyn Tool> = Arc::new(MockTool::new("fetch", "result_b"));

        let agent = AgentLoop::builder().tool(search).tool(fetch).build();
        let result = agent
            .run(&llm, vec![Message::user("Do both")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "Combined results.");
        assert_eq!(result.tool_calls.len(), 2);
        assert_eq!(result.tool_calls[0].0, "search");
        assert_eq!(result.tool_calls[1].0, "fetch");
        assert_eq!(result.iterations, 2);
    }

    #[tokio::test]
    async fn max_iterations_exceeded() {
        // LLM keeps requesting tool calls indefinitely.
        let responses: Vec<LlmResponse> = (0..5)
            .map(|i| {
                LlmResponse::single_tool_call(
                    format!("call_{i}"),
                    "search".to_string(),
                    serde_json::json!({}),
                )
            })
            .collect();
        let llm = MockLlm::new(responses);

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result"));

        let agent = AgentLoop::builder().tool(tool).max_iterations(3).build();
        let err = agent
            .run(&llm, vec![Message::user("loop forever")], |_| {})
            .await
            .unwrap_err();

        assert!(matches!(err, AgentError::MaxIterations(3)));
    }

    #[tokio::test]
    async fn events_emitted_correctly() {
        let llm = MockLlm::new(vec![
            LlmResponse::single_tool_call(
                "call_1".to_string(),
                "search".to_string(),
                serde_json::json!({}),
            ),
            LlmResponse::Message("Done.".to_string()),
        ]);

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result"));
        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder().tool(tool).build();
        let _ = agent
            .run(&llm, vec![Message::user("test")], move |event| {
                let label = match &event {
                    AgentEvent::ToolStarted { name } => format!("started:{name}"),
                    AgentEvent::ToolCompleted { name, .. } => format!("completed:{name}"),
                    AgentEvent::TextChunk(t) => format!("text:{t}"),
                    AgentEvent::TextDone(t) => format!("done:{t}"),
                    AgentEvent::IterationStarted(n) => format!("iter:{n}"),
                    AgentEvent::Interrupted => "interrupted".to_string(),
                    AgentEvent::AskUser { question, .. } => format!("ask_user:{question}"),
                    AgentEvent::AskUserTimeout { call_id, .. } => {
                        format!("ask_user_timeout:{call_id}")
                    }
                    AgentEvent::OpsSaturated { capacity } => {
                        format!("ops_saturated:{capacity}")
                    }
                    AgentEvent::OpDropped { reason } => format!("op_dropped:{reason}"),
                    AgentEvent::OpRejected { reason } => format!("op_rejected:{reason}"),
                };
                events_clone.lock().unwrap().push(label);
            })
            .await
            .unwrap();

        let events = events.lock().unwrap();
        assert_eq!(
            *events,
            vec![
                "iter:1",
                "started:search",
                "completed:search",
                "iter:2",
                "text:Done.",
            ]
        );
    }

    #[tokio::test]
    async fn unknown_tool_produces_error_result() {
        let llm = MockLlm::new(vec![
            LlmResponse::single_tool_call(
                "call_1".to_string(),
                "nonexistent".to_string(),
                serde_json::json!({}),
            ),
            LlmResponse::Message("Handled error.".to_string()),
        ]);

        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder().build();
        let result = agent
            .run(&llm, vec![Message::user("call missing")], move |event| {
                if let AgentEvent::ToolCompleted { result, .. } = &event {
                    if result.is_error {
                        events_clone
                            .lock()
                            .unwrap()
                            .push("error_tool_result".to_string());
                    }
                }
            })
            .await
            .unwrap();

        assert_eq!(result.answer, "Handled error.");
        let events = events.lock().unwrap();
        assert!(events.contains(&"error_tool_result".to_string()));
    }

    #[tokio::test]
    async fn builder_tools_method() {
        let tools: Vec<Arc<dyn Tool>> = vec![
            Arc::new(MockTool::new("a", "ra")),
            Arc::new(MockTool::new("b", "rb")),
        ];

        let llm = MockLlm::new(vec![LlmResponse::Message("ok".to_string())]);

        let agent = AgentLoop::builder().tools(tools).build();
        let result = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "ok");
    }

    #[tokio::test]
    async fn noop_event_callback() {
        let llm = MockLlm::new(vec![LlmResponse::Message("ok".to_string())]);
        let agent = AgentLoop::builder().build();
        // Verify that |_| {} compiles and works as the noop callback.
        let result = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();
        assert_eq!(result.answer, "ok");
    }

    // -----------------------------------------------------------------------
    // Mock context provider
    // -----------------------------------------------------------------------

    struct MockContextProvider {
        context: String,
    }

    impl MockContextProvider {
        fn new(context: &str) -> Self {
            Self {
                context: context.to_string(),
            }
        }
    }

    #[async_trait]
    impl ContextProvider for MockContextProvider {
        async fn build(&self, _query: &str) -> crate::Result<String> {
            Ok(self.context.clone())
        }
    }

    /// A context provider that echoes the query back, useful for verifying
    /// that the correct query string is passed.
    struct EchoContextProvider;

    #[async_trait]
    impl ContextProvider for EchoContextProvider {
        async fn build(&self, query: &str) -> crate::Result<String> {
            Ok(format!("echo: {query}"))
        }
    }

    // -----------------------------------------------------------------------
    // Context provider tests
    // -----------------------------------------------------------------------

    /// A capturing LLM that records the messages it receives, then returns
    /// a canned response.
    struct CapturingLlm {
        captured: Mutex<Vec<Vec<Message>>>,
        responses: Mutex<Vec<LlmResponse>>,
    }

    impl CapturingLlm {
        fn new(responses: Vec<LlmResponse>) -> Self {
            Self {
                captured: Mutex::new(Vec::new()),
                responses: Mutex::new(responses),
            }
        }

        fn captured_messages(&self) -> Vec<Vec<Message>> {
            self.captured.lock().unwrap().clone()
        }
    }

    #[async_trait]
    impl LlmClient for CapturingLlm {
        async fn chat(
            &self,
            messages: &[Message],
            _tools: &[motosan_agent_tool::ToolDef],
        ) -> crate::Result<ChatOutput> {
            self.captured.lock().unwrap().push(messages.to_vec());
            let mut responses = self.responses.lock().unwrap();
            if responses.is_empty() {
                panic!("CapturingLlm: no more responses");
            }
            Ok(ChatOutput::new(responses.remove(0)))
        }
    }

    #[tokio::test]
    async fn context_provider_injects_system_message() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("answer".to_string())]);

        let agent = AgentLoop::builder()
            .context(MockContextProvider::new("You have access to RAG docs."))
            .build();

        let result = agent
            .run(&llm, vec![Message::user("tell me about rust")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "answer");

        // Verify the LLM received the context as a system message
        // *prepended* before the user message.
        let calls = llm.captured_messages();
        assert_eq!(calls.len(), 1);
        let msgs = &calls[0];
        // Should have: system context message + user message
        assert_eq!(msgs.len(), 2);
        assert_eq!(msgs[0].role, crate::message::Role::System);
        assert_eq!(msgs[0].content, "You have access to RAG docs.");
        assert_eq!(msgs[1].role, crate::message::Role::User);
        assert_eq!(msgs[1].content, "tell me about rust");
    }

    #[tokio::test]
    async fn empty_context_is_skipped() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("answer".to_string())]);

        let agent = AgentLoop::builder()
            .context(MockContextProvider::new("")) // empty — should be skipped
            .build();

        let result = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "answer");

        let calls = llm.captured_messages();
        assert_eq!(calls.len(), 1);
        let msgs = &calls[0];
        // Only the user message, no system context.
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].role, crate::message::Role::User);
    }

    #[tokio::test]
    async fn multiple_context_providers() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("done".to_string())]);

        let agent = AgentLoop::builder()
            .context(MockContextProvider::new("RAG context here"))
            .context(MockContextProvider::new("")) // empty — skipped
            .context(MockContextProvider::new("User profile: premium"))
            .build();

        let result = agent
            .run(&llm, vec![Message::user("query")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "done");

        let calls = llm.captured_messages();
        assert_eq!(calls.len(), 1);
        let msgs = &calls[0];
        // 2 non-empty context messages (prepended) + user
        assert_eq!(msgs.len(), 3);
        assert_eq!(msgs[0].role, crate::message::Role::System);
        assert_eq!(msgs[0].content, "RAG context here");
        assert_eq!(msgs[1].role, crate::message::Role::System);
        assert_eq!(msgs[1].content, "User profile: premium");
        assert_eq!(msgs[2].role, crate::message::Role::User);
    }

    #[tokio::test]
    async fn context_provider_receives_user_query() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("ok".to_string())]);

        let agent = AgentLoop::builder().context(EchoContextProvider).build();

        let result = agent
            .run(&llm, vec![Message::user("my question")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "ok");

        let calls = llm.captured_messages();
        let msgs = &calls[0];
        assert_eq!(msgs.len(), 2);
        assert_eq!(msgs[0].role, crate::message::Role::System);
        assert_eq!(msgs[0].content, "echo: my question");
    }

    #[tokio::test]
    async fn no_context_providers_leaves_messages_unchanged() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("ok".to_string())]);

        let agent = AgentLoop::builder().build(); // no context providers

        let _ = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();

        let calls = llm.captured_messages();
        let msgs = &calls[0];
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].content, "hi");
    }

    #[tokio::test]
    async fn builder_contexts_batch_method() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("ok".to_string())]);

        let providers: Vec<Box<dyn ContextProvider>> = vec![
            Box::new(MockContextProvider::new("ctx-a")),
            Box::new(MockContextProvider::new("ctx-b")),
        ];

        let agent = AgentLoop::builder().contexts(providers).build();
        let result = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "ok");

        let calls = llm.captured_messages();
        let msgs = &calls[0];
        // 2 context messages prepended + user message
        assert_eq!(msgs.len(), 3);
        assert_eq!(msgs[0].role, crate::message::Role::System);
        assert_eq!(msgs[0].content, "ctx-a");
        assert_eq!(msgs[1].role, crate::message::Role::System);
        assert_eq!(msgs[1].content, "ctx-b");
        assert_eq!(msgs[2].role, crate::message::Role::User);
        assert_eq!(msgs[2].content, "hi");
    }

    /// A context provider that sleeps for a given duration, used to verify
    /// parallel execution.
    struct DelayContextProvider {
        context: String,
        delay: std::time::Duration,
    }

    impl DelayContextProvider {
        fn new(context: &str, delay: std::time::Duration) -> Self {
            Self {
                context: context.to_string(),
                delay,
            }
        }
    }

    #[async_trait]
    impl ContextProvider for DelayContextProvider {
        async fn build(&self, _query: &str) -> crate::Result<String> {
            tokio::time::sleep(self.delay).await;
            Ok(self.context.clone())
        }
    }

    #[tokio::test]
    async fn context_providers_run_in_parallel() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("ok".to_string())]);
        let delay = std::time::Duration::from_millis(100);

        let agent = AgentLoop::builder()
            .context(DelayContextProvider::new("ctx-a", delay))
            .context(DelayContextProvider::new("ctx-b", delay))
            .context(DelayContextProvider::new("ctx-c", delay))
            .build();

        let start = std::time::Instant::now();
        let result = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();
        let elapsed = start.elapsed();

        // If sequential, ~300ms; if parallel, ~100ms.
        // Use 250ms as a generous threshold.
        assert!(
            elapsed < std::time::Duration::from_millis(250),
            "Expected parallel execution (<250ms), but took {elapsed:?}",
        );

        assert_eq!(result.answer, "ok");

        // Verify order is preserved.
        let calls = llm.captured_messages();
        let msgs = &calls[0];
        assert_eq!(msgs.len(), 4); // 3 context + 1 user
        assert_eq!(msgs[0].content, "ctx-a");
        assert_eq!(msgs[1].content, "ctx-b");
        assert_eq!(msgs[2].content, "ctx-c");
        assert_eq!(msgs[3].content, "hi");
    }

    // -----------------------------------------------------------------------
    // system_prompt shortcut tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn system_prompt_injects_system_message() {
        let llm = CapturingLlm::new(vec![LlmResponse::Message("ok".to_string())]);

        let agent = AgentLoop::builder()
            .system_prompt("You are a helpful assistant.")
            .build();

        let result = agent
            .run(&llm, vec![Message::user("hello")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.answer, "ok");

        let calls = llm.captured_messages();
        assert_eq!(calls.len(), 1);
        let msgs = &calls[0];
        assert_eq!(msgs.len(), 2);
        assert_eq!(msgs[0].role, crate::message::Role::System);
        assert_eq!(msgs[0].content, "You are a helpful assistant.");
        assert_eq!(msgs[1].role, crate::message::Role::User);
        assert_eq!(msgs[1].content, "hello");
    }

    // -----------------------------------------------------------------------
    // Streaming tests
    // -----------------------------------------------------------------------

    /// A mock LLM that emits text deltas via `chat_stream` before the final
    /// `Done` chunk, simulating true streaming behaviour.
    struct StreamingMockLlm {
        /// Each entry is a sequence of StreamChunks for one call.
        responses: Mutex<Vec<Vec<crate::llm::StreamChunk>>>,
    }

    impl StreamingMockLlm {
        fn new(responses: Vec<Vec<crate::llm::StreamChunk>>) -> Self {
            Self {
                responses: Mutex::new(responses),
            }
        }
    }

    #[async_trait]
    impl LlmClient for StreamingMockLlm {
        async fn chat(
            &self,
            _messages: &[Message],
            _tools: &[ToolDef],
        ) -> crate::Result<ChatOutput> {
            panic!("StreamingMockLlm: chat() should not be called");
        }

        fn chat_stream<'a>(
            &'a self,
            _messages: &'a [Message],
            _tools: &'a [ToolDef],
        ) -> std::pin::Pin<
            Box<dyn futures::Stream<Item = crate::Result<crate::llm::StreamChunk>> + Send + 'a>,
        > {
            let chunks = {
                let mut responses = self.responses.lock().unwrap();
                if responses.is_empty() {
                    panic!("StreamingMockLlm: no more responses");
                }
                responses.remove(0)
            };
            Box::pin(futures::stream::iter(chunks.into_iter().map(Ok)))
        }
    }

    #[tokio::test]
    async fn run_streaming_emits_text_chunks_and_done() {
        let llm = StreamingMockLlm::new(vec![vec![
            crate::llm::StreamChunk::TextDelta("Hel".into()),
            crate::llm::StreamChunk::TextDelta("lo!".into()),
            crate::llm::StreamChunk::Done(LlmResponse::Message("Hello!".into())),
        ]]);

        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder().build();
        let result = agent
            .run_streaming(&llm, vec![Message::user("Hi")], move |event| {
                let label = match &event {
                    AgentEvent::TextChunk(t) => format!("chunk:{t}"),
                    AgentEvent::TextDone(t) => format!("done:{t}"),
                    AgentEvent::IterationStarted(n) => format!("iter:{n}"),
                    AgentEvent::ToolStarted { name } => format!("started:{name}"),
                    AgentEvent::ToolCompleted { name, .. } => format!("completed:{name}"),
                    AgentEvent::Interrupted => "interrupted".to_string(),
                    AgentEvent::AskUser { question, .. } => format!("ask_user:{question}"),
                    AgentEvent::AskUserTimeout { call_id, .. } => {
                        format!("ask_user_timeout:{call_id}")
                    }
                    _ => format!("{event:?}"),
                };
                events_clone.lock().unwrap().push(label);
            })
            .await
            .unwrap();

        assert_eq!(result.answer, "Hello!");
        assert_eq!(result.iterations, 1);

        let events = events.lock().unwrap();
        assert_eq!(
            *events,
            vec!["iter:1", "chunk:Hel", "chunk:lo!", "done:Hello!"]
        );
    }

    #[tokio::test]
    async fn run_streaming_with_tool_calls() {
        let llm = StreamingMockLlm::new(vec![
            // First call: tool call (no text deltas)
            vec![crate::llm::StreamChunk::Done(
                LlmResponse::single_tool_call(
                    "call_1".into(),
                    "search".into(),
                    serde_json::json!({"input": "rust"}),
                ),
            )],
            // Second call: streamed text response
            vec![
                crate::llm::StreamChunk::TextDelta("Found ".into()),
                crate::llm::StreamChunk::TextDelta("it.".into()),
                crate::llm::StreamChunk::Done(LlmResponse::Message("Found it.".into())),
            ],
        ]);

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "Rust is great"));
        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder().tool(tool).build();
        let result = agent
            .run_streaming(&llm, vec![Message::user("search rust")], move |event| {
                let label = match &event {
                    AgentEvent::TextChunk(t) => format!("chunk:{t}"),
                    AgentEvent::TextDone(t) => format!("done:{t}"),
                    AgentEvent::IterationStarted(n) => format!("iter:{n}"),
                    AgentEvent::ToolStarted { name } => format!("started:{name}"),
                    AgentEvent::ToolCompleted { name, .. } => format!("completed:{name}"),
                    AgentEvent::Interrupted => "interrupted".to_string(),
                    AgentEvent::AskUser { question, .. } => format!("ask_user:{question}"),
                    AgentEvent::AskUserTimeout { call_id, .. } => {
                        format!("ask_user_timeout:{call_id}")
                    }
                    _ => format!("{event:?}"),
                };
                events_clone.lock().unwrap().push(label);
            })
            .await
            .unwrap();

        assert_eq!(result.answer, "Found it.");
        assert_eq!(result.tool_calls.len(), 1);
        assert_eq!(result.iterations, 2);

        let events = events.lock().unwrap();
        assert_eq!(
            *events,
            vec![
                "iter:1",
                "started:search",
                "completed:search",
                "iter:2",
                "chunk:Found ",
                "chunk:it.",
                "done:Found it.",
            ]
        );
    }

    #[tokio::test]
    async fn run_streaming_fallback_non_streaming_llm() {
        // MockLlm doesn't override chat_stream, so it uses the default
        // fallback that delegates to chat().
        let llm = MockLlm::new(vec![LlmResponse::Message("Hello!".into())]);

        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder().build();
        let result = agent
            .run_streaming(&llm, vec![Message::user("Hi")], move |event| {
                let label = match &event {
                    AgentEvent::TextChunk(t) => format!("chunk:{t}"),
                    AgentEvent::TextDone(t) => format!("done:{t}"),
                    AgentEvent::IterationStarted(n) => format!("iter:{n}"),
                    _ => "other".into(),
                };
                events_clone.lock().unwrap().push(label);
            })
            .await
            .unwrap();

        assert_eq!(result.answer, "Hello!");

        let events = events.lock().unwrap();
        // Default fallback: no text deltas, so run_streaming emits
        // the full text as a TextChunk, then TextDone.
        assert_eq!(*events, vec!["iter:1", "chunk:Hello!", "done:Hello!"]);
    }

    // -----------------------------------------------------------------------
    // Token usage tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn usage_is_nonzero_after_mocked_llm_call() {
        let usage = TokenUsage {
            input_tokens: 100,
            output_tokens: 50,
        };
        let llm = MockLlm::with_usage(vec![LlmResponse::Message("Hello!".to_string())], usage);

        let agent = AgentLoop::builder().build();
        let result = agent
            .run(&llm, vec![Message::user("Hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.usage.input_tokens, 100);
        assert_eq!(result.usage.output_tokens, 50);
    }

    #[tokio::test]
    async fn usage_accumulates_across_iterations() {
        let usage = TokenUsage {
            input_tokens: 10,
            output_tokens: 20,
        };
        let llm = MockLlm::with_usage(
            vec![
                LlmResponse::single_tool_call(
                    "call_1".to_string(),
                    "search".to_string(),
                    serde_json::json!({}),
                ),
                LlmResponse::Message("Done.".to_string()),
            ],
            usage,
        );

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result"));
        let agent = AgentLoop::builder().tool(tool).build();
        let result = agent
            .run(&llm, vec![Message::user("test")], |_| {})
            .await
            .unwrap();

        // Two LLM calls, each with usage (10, 20)
        assert_eq!(result.usage.input_tokens, 20);
        assert_eq!(result.usage.output_tokens, 40);
        assert_eq!(result.iterations, 2);
    }

    #[tokio::test]
    async fn usage_zero_when_llm_reports_no_usage() {
        let llm = MockLlm::new(vec![LlmResponse::Message("ok".to_string())]);

        let agent = AgentLoop::builder().build();
        let result = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.usage.input_tokens, 0);
        assert_eq!(result.usage.output_tokens, 0);
    }

    #[tokio::test]
    async fn streaming_usage_accumulates() {
        let llm = StreamingMockLlm::new(vec![vec![
            crate::llm::StreamChunk::TextDelta("Hi".into()),
            crate::llm::StreamChunk::Done(LlmResponse::Message("Hi".into())),
            crate::llm::StreamChunk::Usage(TokenUsage {
                input_tokens: 50,
                output_tokens: 25,
            }),
        ]]);

        let agent = AgentLoop::builder().build();
        let result = agent
            .run_streaming(&llm, vec![Message::user("Hi")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.usage.input_tokens, 50);
        assert_eq!(result.usage.output_tokens, 25);
    }

    #[tokio::test]
    async fn run_streaming_max_iterations() {
        let responses: Vec<Vec<crate::llm::StreamChunk>> = (0..5)
            .map(|i| {
                vec![crate::llm::StreamChunk::Done(
                    LlmResponse::single_tool_call(
                        format!("call_{i}"),
                        "search".into(),
                        serde_json::json!({}),
                    ),
                )]
            })
            .collect();
        let llm = StreamingMockLlm::new(responses);

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result"));
        let agent = AgentLoop::builder().tool(tool).max_iterations(3).build();
        let err = agent
            .run_streaming(&llm, vec![Message::user("loop")], |_| {})
            .await
            .unwrap_err();

        assert!(matches!(err, AgentError::MaxIterations(3)));
    }

    // -----------------------------------------------------------------------
    // messages field tests
    // -----------------------------------------------------------------------

    #[tokio::test]
    async fn result_messages_contains_full_history() {
        let llm = MockLlm::new(vec![LlmResponse::Message("Hello!".to_string())]);

        let agent = AgentLoop::builder().build();
        let result = agent
            .run(&llm, vec![Message::user("Hi")], |_| {})
            .await
            .unwrap();

        // Should contain the user message + assistant reply.
        assert_eq!(result.messages.len(), 2);
        assert_eq!(result.messages[0].role, crate::message::Role::User);
        assert_eq!(result.messages[0].content, "Hi");
        assert_eq!(result.messages[1].role, crate::message::Role::Assistant);
        assert_eq!(result.messages[1].content, "Hello!");
    }

    #[tokio::test]
    async fn result_messages_includes_tool_call_pairs() {
        let llm = MockLlm::new(vec![
            LlmResponse::single_tool_call(
                "call_1".to_string(),
                "search".to_string(),
                serde_json::json!({"input": "rust"}),
            ),
            LlmResponse::Message("Found it.".to_string()),
        ]);

        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result text"));

        let agent = AgentLoop::builder().tool(tool).build();
        let result = agent
            .run(&llm, vec![Message::user("Search")], |_| {})
            .await
            .unwrap();

        // user -> assistant(tool_call) -> tool_result -> assistant(final)
        assert_eq!(result.messages.len(), 4);
        assert_eq!(result.messages[0].role, crate::message::Role::User);
        assert_eq!(result.messages[1].role, crate::message::Role::Assistant);
        assert_eq!(result.messages[1].tool_calls.len(), 1);
        assert_eq!(result.messages[1].tool_calls[0].name, "search");
        assert_eq!(result.messages[2].role, crate::message::Role::Tool);
        assert_eq!(result.messages[3].role, crate::message::Role::Assistant);
        assert_eq!(result.messages[3].content, "Found it.");
    }

    #[tokio::test]
    async fn multi_turn_continuation_via_messages() {
        // First turn: simple Q&A
        let llm1 = MockLlm::new(vec![LlmResponse::Message("I'm fine!".to_string())]);
        let agent = AgentLoop::builder().build();
        let result1 = agent
            .run(&llm1, vec![Message::user("How are you?")], |_| {})
            .await
            .unwrap();

        // Second turn: continue using result.messages
        let llm2 = MockLlm::new(vec![LlmResponse::Message("Goodbye!".to_string())]);
        let mut next_messages = result1.messages;
        next_messages.push(Message::user("Bye!"));

        let result2 = agent.run(&llm2, next_messages, |_| {}).await.unwrap();

        assert_eq!(result2.answer, "Goodbye!");
        // History: user("How are you?") + assistant("I'm fine!") + user("Bye!") + assistant("Goodbye!")
        assert_eq!(result2.messages.len(), 4);
        assert_eq!(result2.messages[0].content, "How are you?");
        assert_eq!(result2.messages[1].content, "I'm fine!");
        assert_eq!(result2.messages[2].content, "Bye!");
        assert_eq!(result2.messages[3].content, "Goodbye!");
    }

    #[tokio::test]
    async fn streaming_result_messages_contains_full_history() {
        let llm = StreamingMockLlm::new(vec![vec![
            crate::llm::StreamChunk::TextDelta("Hi".to_string()),
            crate::llm::StreamChunk::TextDelta(" there".to_string()),
            crate::llm::StreamChunk::Done(LlmResponse::Message("Hi there".to_string())),
        ]]);

        let agent = AgentLoop::builder().build();
        let result = agent
            .run_streaming(&llm, vec![Message::user("Hello")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.messages.len(), 2);
        assert_eq!(result.messages[0].role, crate::message::Role::User);
        assert_eq!(result.messages[0].content, "Hello");
        assert_eq!(result.messages[1].role, crate::message::Role::Assistant);
        assert_eq!(result.messages[1].content, "Hi there");
    }

    #[test]
    fn tool_execution_policy_detects_ask_user_presence() {
        let parallel_items = vec![ToolCallItem {
            id: "call_1".to_string(),
            name: "search".to_string(),
            args: serde_json::json!({}),
        }];
        let interactive_items = vec![
            ToolCallItem {
                id: "call_1".to_string(),
                name: "search".to_string(),
                args: serde_json::json!({}),
            },
            ToolCallItem {
                id: "call_2".to_string(),
                name: "ask_user".to_string(),
                args: serde_json::json!({"question": "continue?"}),
            },
        ];

        assert!(matches!(
            ToolExecutionPolicy::from_items(&parallel_items),
            ToolExecutionPolicy::ParallelOnly
        ));
        assert!(matches!(
            ToolExecutionPolicy::from_items(&interactive_items),
            ToolExecutionPolicy::InteractiveAskUser
        ));
    }

    // -----------------------------------------------------------------------
    // Pipeline stage regression tests
    // -----------------------------------------------------------------------

    /// Verify that tool results are appended in the same order as tool calls.
    /// This is a regression test for the `execute_and_record_tool_calls` stage:
    /// tool result messages must correlate 1:1 with the tool call items.
    #[tokio::test]
    async fn tool_result_ordering_matches_call_ordering() {
        let llm = MockLlm::new(vec![
            LlmResponse::ToolCalls(vec![
                ToolCallItem {
                    id: "c1".to_string(),
                    name: "alpha".to_string(),
                    args: serde_json::json!({}),
                },
                ToolCallItem {
                    id: "c2".to_string(),
                    name: "beta".to_string(),
                    args: serde_json::json!({}),
                },
                ToolCallItem {
                    id: "c3".to_string(),
                    name: "gamma".to_string(),
                    args: serde_json::json!({}),
                },
            ]),
            LlmResponse::Message("done".to_string()),
        ]);

        let alpha: Arc<dyn Tool> = Arc::new(MockTool::new("alpha", "res_alpha"));
        let beta: Arc<dyn Tool> = Arc::new(MockTool::new("beta", "res_beta"));
        let gamma: Arc<dyn Tool> = Arc::new(MockTool::new("gamma", "res_gamma"));

        let agent = AgentLoop::builder()
            .tool(alpha)
            .tool(beta)
            .tool(gamma)
            .build();

        let result = agent
            .run(&llm, vec![Message::user("go")], |_| {})
            .await
            .unwrap();

        // Verify tool_calls record order.
        assert_eq!(result.tool_calls[0].0, "alpha");
        assert_eq!(result.tool_calls[1].0, "beta");
        assert_eq!(result.tool_calls[2].0, "gamma");

        // Verify message history: user -> assistant(tool_calls) -> 3x tool_result -> assistant(done)
        assert_eq!(result.messages.len(), 6);
        // The assistant message at index 1 should carry all 3 tool call refs.
        assert_eq!(result.messages[1].tool_calls.len(), 3);
        assert_eq!(result.messages[1].tool_calls[0].id, "c1");
        assert_eq!(result.messages[1].tool_calls[1].id, "c2");
        assert_eq!(result.messages[1].tool_calls[2].id, "c3");
        // Tool result messages at indices 2, 3, 4 must match call order.
        assert_eq!(result.messages[2].tool_call_id.as_deref(), Some("c1"));
        assert_eq!(result.messages[2].content, "res_alpha");
        assert_eq!(result.messages[3].tool_call_id.as_deref(), Some("c2"));
        assert_eq!(result.messages[3].content, "res_beta");
        assert_eq!(result.messages[4].tool_call_id.as_deref(), Some("c3"));
        assert_eq!(result.messages[4].content, "res_gamma");
    }

    /// Verify that the pipeline stages emit events in the correct order:
    /// IterationStarted -> ToolStarted (all) -> ToolCompleted (all) -> next iteration.
    /// This tests the boundary between the LLM step and tool-call execution stage.
    #[tokio::test]
    async fn stage_boundary_event_ordering() {
        let llm = MockLlm::new(vec![
            LlmResponse::ToolCalls(vec![
                ToolCallItem {
                    id: "c1".to_string(),
                    name: "x".to_string(),
                    args: serde_json::json!({}),
                },
                ToolCallItem {
                    id: "c2".to_string(),
                    name: "y".to_string(),
                    args: serde_json::json!({}),
                },
            ]),
            LlmResponse::Message("final".to_string()),
        ]);

        let x: Arc<dyn Tool> = Arc::new(MockTool::new("x", "rx"));
        let y: Arc<dyn Tool> = Arc::new(MockTool::new("y", "ry"));

        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder().tool(x).tool(y).build();
        let _ = agent
            .run(&llm, vec![Message::user("go")], move |event| {
                let label = match &event {
                    AgentEvent::IterationStarted(n) => format!("iter:{n}"),
                    AgentEvent::ToolStarted { name } => format!("started:{name}"),
                    AgentEvent::ToolCompleted { name, .. } => format!("completed:{name}"),
                    AgentEvent::TextChunk(t) => format!("text:{t}"),
                    AgentEvent::TextDone(_) => "done".to_string(),
                    _ => "other".to_string(),
                };
                events_clone.lock().unwrap().push(label);
            })
            .await
            .unwrap();

        let events = events.lock().unwrap();
        // Stage boundaries:
        // 1. iter:1 (iteration starts)
        // 2. started:x, started:y (tool-call planning emits all ToolStarted before execution)
        // 3. completed:x, completed:y (execution completes, results recorded)
        // 4. iter:2 (next iteration)
        // 5. text:final (LLM returns message)
        assert_eq!(events[0], "iter:1");
        assert_eq!(events[1], "started:x");
        assert_eq!(events[2], "started:y");
        assert_eq!(events[3], "completed:x");
        assert_eq!(events[4], "completed:y");
        assert_eq!(events[5], "iter:2");
        assert_eq!(events[6], "text:final");
    }

    /// Verify that `TurnState::into_result` correctly captures all accumulated
    /// state (tool calls, usage, messages).
    #[tokio::test]
    async fn turn_state_accumulation_across_iterations() {
        let usage = TokenUsage {
            input_tokens: 7,
            output_tokens: 3,
        };
        let llm = MockLlm::with_usage(
            vec![
                LlmResponse::single_tool_call(
                    "c1".to_string(),
                    "t1".to_string(),
                    serde_json::json!({}),
                ),
                LlmResponse::single_tool_call(
                    "c2".to_string(),
                    "t2".to_string(),
                    serde_json::json!({}),
                ),
                LlmResponse::Message("end".to_string()),
            ],
            usage,
        );

        let t1: Arc<dyn Tool> = Arc::new(MockTool::new("t1", "r1"));
        let t2: Arc<dyn Tool> = Arc::new(MockTool::new("t2", "r2"));

        let agent = AgentLoop::builder().tool(t1).tool(t2).build();
        let result = agent
            .run(&llm, vec![Message::user("go")], |_| {})
            .await
            .unwrap();

        assert_eq!(result.iterations, 3);
        assert_eq!(result.tool_calls.len(), 2);
        // 3 LLM calls * 7 input tokens each
        assert_eq!(result.usage.input_tokens, 21);
        assert_eq!(result.usage.output_tokens, 9);
        // Message history: user + (assistant+tool_result) * 2 + assistant(final)
        // user, assistant(tc), tool_result, assistant(tc), tool_result, assistant(final)
        assert_eq!(result.messages.len(), 6);
    }

    /// Verify that the streaming path produces identical message history
    /// to the non-streaming path for the same logical conversation.
    #[tokio::test]
    async fn streaming_and_non_streaming_produce_same_messages() {
        // Non-streaming
        let llm_sync = MockLlm::new(vec![
            LlmResponse::single_tool_call(
                "c1".to_string(),
                "search".to_string(),
                serde_json::json!({"q": "rust"}),
            ),
            LlmResponse::Message("Found it.".to_string()),
        ]);
        let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result text"));
        let agent = AgentLoop::builder().tool(tool).build();
        let result_sync = agent
            .run(&llm_sync, vec![Message::user("Search")], |_| {})
            .await
            .unwrap();

        // Streaming (using MockLlm which falls back to chat via default chat_stream)
        let llm_stream = MockLlm::new(vec![
            LlmResponse::single_tool_call(
                "c1".to_string(),
                "search".to_string(),
                serde_json::json!({"q": "rust"}),
            ),
            LlmResponse::Message("Found it.".to_string()),
        ]);
        let tool2: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result text"));
        let agent2 = AgentLoop::builder().tool(tool2).build();
        let result_stream = agent2
            .run_streaming(&llm_stream, vec![Message::user("Search")], |_| {})
            .await
            .unwrap();

        // Both paths should produce the same message history.
        assert_eq!(result_sync.messages.len(), result_stream.messages.len());
        for (a, b) in result_sync
            .messages
            .iter()
            .zip(result_stream.messages.iter())
        {
            assert_eq!(a.role, b.role);
            assert_eq!(a.content, b.content);
            assert_eq!(a.tool_call_id, b.tool_call_id);
            assert_eq!(a.tool_calls.len(), b.tool_calls.len());
        }
        assert_eq!(result_sync.iterations, result_stream.iterations);
        assert_eq!(result_sync.tool_calls.len(), result_stream.tool_calls.len());
    }

    // -----------------------------------------------------------------------
    // Cancellation tests (cancellation feature)
    // -----------------------------------------------------------------------

    #[cfg(feature = "cancellation")]
    mod cancellation_tests {
        use super::*;
        use tokio_util::sync::CancellationToken;

        #[tokio::test]
        async fn cancel_before_run_returns_cancelled() {
            let llm = MockLlm::new(vec![LlmResponse::Message("Hello!".into())]);
            let agent = AgentLoop::builder().build();

            let token = CancellationToken::new();
            token.cancel();

            let err = agent
                .run_with_cancel(&llm, vec![Message::user("Hi")], token, |_| {})
                .await
                .unwrap_err();

            assert!(matches!(err, AgentError::Cancelled));
        }

        #[tokio::test]
        async fn cancel_mid_run_returns_cancelled() {
            let responses: Vec<LlmResponse> = (0..5)
                .map(|i| {
                    LlmResponse::single_tool_call(
                        format!("call_{i}"),
                        "search".to_string(),
                        serde_json::json!({}),
                    )
                })
                .collect();
            let llm = MockLlm::new(responses);
            let tool: Arc<dyn Tool> = Arc::new(MockTool::new("search", "result"));

            let agent = AgentLoop::builder().tool(tool).max_iterations(5).build();

            let token = CancellationToken::new();
            let child = token.child_token();
            let iterations = Arc::new(Mutex::new(0usize));
            let iterations_clone = iterations.clone();
            let token_clone = token.clone();

            let err = agent
                .run_with_cancel(&llm, vec![Message::user("loop")], child, move |event| {
                    if let AgentEvent::IterationStarted(_) = &event {
                        let mut count = iterations_clone.lock().unwrap();
                        *count += 1;
                        if *count >= 2 {
                            token_clone.cancel();
                        }
                    }
                })
                .await
                .unwrap_err();

            assert!(matches!(err, AgentError::Cancelled));
            let count = *iterations.lock().unwrap();
            assert!(
                count >= 2 && count <= 3,
                "unexpected iteration count: {count}"
            );
        }

        #[tokio::test]
        async fn cancel_streaming_returns_cancelled() {
            let llm = MockLlm::new(vec![LlmResponse::Message("Hello!".into())]);
            let agent = AgentLoop::builder().build();

            let token = CancellationToken::new();
            token.cancel();

            let err = agent
                .run_streaming_with_cancel(&llm, vec![Message::user("Hi")], token, |_| {})
                .await
                .unwrap_err();

            assert!(matches!(err, AgentError::Cancelled));
        }

        #[tokio::test]
        async fn uncancelled_token_runs_normally() {
            let llm = MockLlm::new(vec![LlmResponse::Message("Hello!".into())]);
            let agent = AgentLoop::builder().build();

            let token = CancellationToken::new();

            let result = agent
                .run_with_cancel(&llm, vec![Message::user("Hi")], token, |_| {})
                .await
                .unwrap();

            assert_eq!(result.answer, "Hello!");
            assert_eq!(result.iterations, 1);
        }
    }
}

#[cfg(all(test, feature = "mcp-client"))]
mod mcp_integration_tests {
    use super::*;
    use crate::mcp::McpServer;
    use async_trait::async_trait;
    use motosan_agent_tool::ToolDef;
    use serde_json::json;

    struct EchoMcpServer {
        name: String,
    }

    #[async_trait]
    impl McpServer for EchoMcpServer {
        fn name(&self) -> &str {
            &self.name
        }
        async fn connect(&self) -> crate::Result<()> {
            Ok(())
        }
        async fn list_tools(&self) -> crate::Result<Vec<ToolDef>> {
            Ok(vec![ToolDef {
                name: "echo".to_string(),
                description: "Echo input".to_string(),
                input_schema: json!({"type": "object", "properties": {"msg": {"type": "string"}}}),
            }])
        }
        async fn call_tool(&self, _name: &str, args: serde_json::Value) -> crate::Result<String> {
            Ok(format!(
                "echo: {}",
                args.get("msg").and_then(|v| v.as_str()).unwrap_or("")
            ))
        }
        async fn disconnect(&self) -> crate::Result<()> {
            Ok(())
        }
    }

    #[test]
    fn builder_accepts_mcp_server() {
        let agent = AgentLoop::builder()
            .mcp_server(EchoMcpServer {
                name: "test_mcp".to_string(),
            })
            .max_iterations(5)
            .build();
        assert_eq!(agent.mcp_servers.len(), 1);
    }

    #[test]
    fn builder_accepts_shared_arc_mcp_server() {
        let shared: Arc<dyn McpServer> = Arc::new(EchoMcpServer {
            name: "shared_mcp".to_string(),
        });

        let agent_a = AgentLoop::builder()
            .mcp_server_arc(Arc::clone(&shared))
            .max_iterations(5)
            .build();

        let agent_b = AgentLoop::builder()
            .mcp_server_arc(Arc::clone(&shared))
            .max_iterations(5)
            .build();

        assert_eq!(agent_a.mcp_servers.len(), 1);
        assert_eq!(agent_b.mcp_servers.len(), 1);
        assert!(Arc::ptr_eq(
            &agent_a.mcp_servers[0],
            &agent_b.mcp_servers[0]
        ));
    }

    /// A simple mock LLM for MCP streaming tests.
    /// Uses the default `chat_stream` fallback (delegates to `chat`).
    struct McpTestLlm {
        responses: std::sync::Mutex<Vec<crate::llm::LlmResponse>>,
    }

    impl McpTestLlm {
        fn new(responses: Vec<crate::llm::LlmResponse>) -> Self {
            Self {
                responses: std::sync::Mutex::new(responses),
            }
        }
    }

    #[async_trait]
    impl crate::llm::LlmClient for McpTestLlm {
        async fn chat(
            &self,
            _messages: &[Message],
            _tools: &[ToolDef],
        ) -> crate::Result<crate::llm::ChatOutput> {
            let mut responses = self.responses.lock().unwrap();
            assert!(!responses.is_empty(), "McpTestLlm: no more responses");
            let response = responses.remove(0);
            Ok(crate::llm::ChatOutput {
                response,
                usage: None,
            })
        }
    }

    #[tokio::test]
    async fn run_streaming_with_mcp_server() {
        use crate::llm::{LlmResponse, ToolCallItem};
        use std::sync::Mutex;

        // LLM first requests the MCP tool (namespaced as server__tool),
        // then returns a final message.
        // Uses the default chat_stream fallback (delegates to chat).
        let llm = McpTestLlm::new(vec![
            LlmResponse::ToolCalls(vec![ToolCallItem {
                id: "call_mcp".to_string(),
                name: "test_mcp__echo".to_string(),
                args: json!({"msg": "hello"}),
            }]),
            LlmResponse::Message("MCP says: echo hello".to_string()),
        ]);

        let events: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let events_clone = events.clone();

        let agent = AgentLoop::builder()
            .mcp_server(EchoMcpServer {
                name: "test_mcp".to_string(),
            })
            .max_iterations(5)
            .build();

        let result = agent
            .run_streaming(&llm, vec![Message::user("call echo")], move |event| {
                let label = match &event {
                    AgentEvent::ToolStarted { name } => format!("started:{name}"),
                    AgentEvent::ToolCompleted { name, result } => {
                        format!("completed:{name}:{}", result.as_text().unwrap_or(""))
                    }
                    AgentEvent::TextChunk(t) => format!("chunk:{t}"),
                    AgentEvent::TextDone(t) => format!("done:{t}"),
                    AgentEvent::IterationStarted(n) => format!("iter:{n}"),
                    AgentEvent::Interrupted => "interrupted".to_string(),
                    AgentEvent::AskUser { question, .. } => format!("ask_user:{question}"),
                    AgentEvent::AskUserTimeout { call_id, .. } => {
                        format!("ask_user_timeout:{call_id}")
                    }
                };
                events_clone.lock().unwrap().push(label);
            })
            .await
            .unwrap();

        assert_eq!(result.answer, "MCP says: echo hello");
        assert_eq!(result.tool_calls.len(), 1);
        assert_eq!(result.tool_calls[0].0, "test_mcp__echo");

        let events = events.lock().unwrap();
        // The MCP tool should have been executed and returned "echo: hello".
        assert!(events
            .iter()
            .any(|e| e.starts_with("completed:test_mcp__echo:echo: hello")));
    }

    /// MCP server that tracks connect/disconnect calls via shared counters.
    struct TrackingMcpServer {
        name: String,
        fail_connect: bool,
        connected: Arc<std::sync::atomic::AtomicBool>,
        disconnect_count: Arc<std::sync::atomic::AtomicUsize>,
    }

    #[async_trait]
    impl McpServer for TrackingMcpServer {
        fn name(&self) -> &str {
            &self.name
        }
        async fn connect(&self) -> crate::Result<()> {
            if self.fail_connect {
                Err(crate::AgentError::Mcp("connect failed".into()))
            } else {
                self.connected
                    .store(true, std::sync::atomic::Ordering::SeqCst);
                Ok(())
            }
        }
        async fn list_tools(&self) -> crate::Result<Vec<ToolDef>> {
            Ok(vec![])
        }
        async fn call_tool(&self, _name: &str, _args: serde_json::Value) -> crate::Result<String> {
            Ok(String::new())
        }
        async fn disconnect(&self) -> crate::Result<()> {
            self.disconnect_count
                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            Ok(())
        }
    }

    #[tokio::test]
    async fn partial_connect_failure_disconnects_already_connected() {
        use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

        let server1_connected = Arc::new(AtomicBool::new(false));
        let server1_disconnects = Arc::new(AtomicUsize::new(0));
        let server2_disconnects = Arc::new(AtomicUsize::new(0));

        let server1: Arc<dyn McpServer> = Arc::new(TrackingMcpServer {
            name: "ok_server".to_string(),
            fail_connect: false,
            connected: server1_connected.clone(),
            disconnect_count: server1_disconnects.clone(),
        });
        let server2: Arc<dyn McpServer> = Arc::new(TrackingMcpServer {
            name: "fail_server".to_string(),
            fail_connect: true,
            connected: Arc::new(AtomicBool::new(false)),
            disconnect_count: server2_disconnects.clone(),
        });

        let agent = AgentLoop::builder()
            .mcp_server_arc(server1)
            .mcp_server_arc(server2)
            .max_iterations(1)
            .build();

        let llm = McpTestLlm::new(vec![]);
        let err = agent
            .run(&llm, vec![Message::user("hi")], |_| {})
            .await
            .unwrap_err();

        assert!(
            matches!(err, crate::AgentError::Mcp(_)),
            "expected connect error"
        );
        // Server 1 was successfully connected, so it must have been disconnected.
        assert!(server1_connected.load(Ordering::SeqCst));
        assert_eq!(server1_disconnects.load(Ordering::SeqCst), 1);
        // Server 2 never connected, so disconnect should not have been called.
        assert_eq!(server2_disconnects.load(Ordering::SeqCst), 0);
    }
}