motosan_agent_loop/core/

engine.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::core::decision::FlowDecision;
use crate::core::event::{AgentEvent, CoreEvent, ExtensionEvent};
use crate::error::AgentError;
use crate::llm::{LlmClient, TokenUsage, ToolCallItem};
use crate::message::{Message, ToolCallRef};
use crate::Result;

/// The current operating mode of the agent loop.
///
/// When plan mode is enabled, the agent can enter `Planning` mode to explore
/// and brainstorm before committing to an implementation approach.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum AgentMode {
    /// Normal execution mode — the agent implements directly.
    #[default]
    Normal,
    /// Planning mode — the agent explores, asks questions, and proposes a plan.
    Planning,
}

/// Schedule controlling when `next_step_prompt` is injected.
///
/// Used by `StuckDetectionExtension` (in `src/extensions/stuck_detection/`)
/// to control the cadence at which the think-before-act prompt is prepended
/// to the conversation.
///
/// The default is `Every(1)` (inject every iteration), but injection only
/// happens when a `next_step_prompt` has been configured on the extension.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NextStepSchedule {
    /// Inject only before the first iteration.
    FirstOnly,
    /// Inject every `n` iterations (1 = every iteration, 2 = every other, …).
    ///
    /// A value of `0` is treated as `1`.
    Every(usize),
}

impl Default for NextStepSchedule {
    fn default() -> Self {
        Self::Every(1)
    }
}

impl NextStepSchedule {
    /// Returns `true` if the prompt should be injected for the given
    /// 1-based `iteration`.
    // TODO(phase-4): remove once StuckDetectionExtension is the only caller.
    #[allow(dead_code)]
    pub(crate) fn should_inject(&self, iteration: usize) -> bool {
        match self {
            Self::FirstOnly => iteration == 1,
            Self::Every(n) => {
                let n = if *n == 0 { 1 } else { *n };
                iteration == 1 || (iteration - 1).is_multiple_of(n)
            }
        }
    }
}

/// 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 [`CoreEvent::OpDropped`]
    /// instead of [`CoreEvent::OpRejected`], allowing callers to distinguish intent.
    ///
    /// [`CoreEvent::OpDropped`]: crate::core::CoreEvent::OpDropped
    /// [`CoreEvent::OpRejected`]: crate::core::CoreEvent::OpRejected
    DropOldest,
    /// Reject the new item and emit a telemetry event.
    Reject,
}

/// Configuration for bounded channel capacities used by [`AgentSession`](crate::AgentSession).
///
/// All capacities have sensible defaults (64) and can be overridden via the
/// builder pattern on [`EngineBuilder`].
#[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,
    /// Capacity of the stream channel used by [`Engine::run`].
    ///
    /// `None` means use the default of `256`. The channel is bounded; if it
    /// fills up, intermediate `Event` items are dropped (best-effort) but the
    /// terminal item is always delivered (sent via the awaiting `send`).
    pub stream_capacity: Option<usize>,
}

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

// ---------------------------------------------------------------------------
// 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
    }
}

/// Outcome of processing an LLM text response via
/// `Engine::handle_text_response`.
enum TextResponseOutcome {
    /// The turn should finalize with this text as the answer. The caller
    /// still needs to push `Message::assistant(text)` onto the state
    /// and build the `AgentResult` (the helper doesn't do this because
    /// the caller returns different types — `AgentResult` vs the
    /// `(Result<AgentResult>, Vec<Message>)` tuple).
    Finalize(String),
    /// An extension injected a continuation message. The helper already
    /// pushed `Message::assistant(text)` and `state.messages.push(msg)`,
    /// and updated `state.continuation_text`. The caller should loop back
    /// to another LLM call immediately.
    Continue,
    /// An extension halted the turn. The caller should return
    /// `Err(AgentError::Internal(...))` with the halt reason's message.
    Halt(String),
}

/// Accumulator for per-run mutable state shared across turn iterations.
struct RunTerminalMeta {
    turn_result: crate::core::extension::TurnResult,
    iteration: usize,
    token_usage: TokenUsage,
}

struct TurnState {
    messages: Vec<Message>,
    all_tool_calls: Vec<(String, serde_json::Value)>,
    total_usage: TokenUsage,
    /// Accumulated text across all token-budget continuation segments.
    continuation_text: String,
}

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

    /// 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.
    ///
    /// If there were continuations, the accumulated continuation text is
    /// prepended to the final segment so the full answer is returned.
    fn into_result(self, answer: String, iteration: usize) -> AgentResult {
        let full_answer = if self.continuation_text.is_empty() {
            answer
        } else {
            let mut buf = self.continuation_text;
            buf.push_str(&answer);
            buf
        };
        AgentResult {
            answer: full_answer,
            tool_calls: self.all_tool_calls,
            iterations: iteration,
            usage: self.total_usage,
            messages: self.messages,
        }
    }
}

impl RunTerminalMeta {
    fn success(agent_result: &AgentResult) -> Self {
        Self {
            turn_result: crate::core::extension::TurnResult::Success,
            iteration: agent_result.iterations,
            token_usage: agent_result.usage,
        }
    }

    fn interrupted(iteration: usize, token_usage: TokenUsage) -> Self {
        Self {
            turn_result: crate::core::extension::TurnResult::Interrupted,
            iteration,
            token_usage,
        }
    }

    fn failure(iteration: usize, token_usage: TokenUsage, err: &AgentError) -> Self {
        Self {
            turn_result: crate::core::extension::TurnResult::Failure(err.to_string()),
            iteration,
            token_usage,
        }
    }
}

fn run_return_success(
    res: AgentResult,
    snapshot: Vec<Message>,
) -> (Result<AgentResult>, Vec<Message>, RunTerminalMeta) {
    let meta = RunTerminalMeta::success(&res);
    (Ok(res), snapshot, meta)
}

fn run_return_interrupted(
    res: AgentResult,
    snapshot: Vec<Message>,
) -> (Result<AgentResult>, Vec<Message>, RunTerminalMeta) {
    let meta = RunTerminalMeta::interrupted(res.iterations, res.usage);
    (Ok(res), snapshot, meta)
}

fn run_return_failure(
    err: AgentError,
    messages: Vec<Message>,
    iteration: usize,
    token_usage: TokenUsage,
) -> (Result<AgentResult>, Vec<Message>, RunTerminalMeta) {
    let meta = RunTerminalMeta::failure(iteration, token_usage, &err);
    (Err(err), messages, meta)
}

/// 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::Core(CoreEvent::ToolStarted {
            name: tc.name.clone(),
        }));
    }
}

/// Merge streamed tool results with batch-executed results for any tools that
/// were not streamed, reassembling everything in the canonical order from the
/// `Done` payload.
///
/// Returns `(final_items, final_results)` in the same order as `canonical_items`.
async fn merge_streamed_tool_results(
    canonical_items: &[ToolCallItem],
    streamed_items: Vec<ToolCallItem>,
    streamed_results: Vec<ToolResult>,
    tool_map: &HashMap<String, Arc<dyn Tool>>,
    tool_timeout: Option<std::time::Duration>,
    tool_context: &ToolContext,
    on_event: &(impl Fn(AgentEvent) + Send + Sync),
) -> (Vec<ToolCallItem>, Vec<ToolResult>) {
    let streamed_ids: std::collections::HashSet<&str> =
        streamed_items.iter().map(|i| i.id.as_str()).collect();

    // Find tools in the canonical payload that were NOT streamed.
    let remaining: Vec<ToolCallItem> = canonical_items
        .iter()
        .filter(|item| !streamed_ids.contains(item.id.as_str()))
        .cloned()
        .collect();

    // Build lookup of streamed results by ID.
    let streamed_map: HashMap<&str, ToolResult> = streamed_items
        .iter()
        .zip(streamed_results)
        .map(|(item, result)| (item.id.as_str(), result))
        .collect();

    // Execute remaining tools (if any) and build their lookup.
    let remaining_map: HashMap<String, ToolResult> = if !remaining.is_empty() {
        emit_tool_started(&remaining, on_event);
        let remaining_results =
            Engine::execute_tools_parallel(tool_map, &remaining, tool_timeout, tool_context).await;
        remaining
            .iter()
            .zip(remaining_results)
            .map(|(item, result)| (item.id.clone(), result))
            .collect()
    } else {
        HashMap::new()
    };

    // Reassemble in canonical order, with a defensive fallback for missing IDs.
    let mut final_items = Vec::with_capacity(canonical_items.len());
    let mut final_results = Vec::with_capacity(canonical_items.len());

    for item in canonical_items {
        let result = if let Some(r) = streamed_map.get(item.id.as_str()).cloned() {
            r
        } else if let Some(r) = remaining_map.get(item.id.as_str()).cloned() {
            r
        } else {
            // Defensive: this should never happen if the backend follows the
            // contract, but we must not silently drop a tool call — the LLM
            // expects a result for every call it made.
            ToolResult::error(format!(
                "internal error: no result for tool call '{}'",
                item.id
            ))
        };
        final_items.push(item.clone());
        final_results.push(result);
    }

    (final_items, final_results)
}

/// A single question within an `ask_user` interaction.
#[derive(Debug, Clone)]
pub struct AskUserQuestion {
    /// The question text to display to the user.
    pub question: String,
    /// An optional short header (max 12 chars) for grouping or labeling.
    pub header: Option<String>,
    /// The available options for the user to choose from.
    pub options: Vec<AskUserOption>,
    /// Whether the user can select multiple options.
    pub multi_select: bool,
}

/// A single option within an [`AskUserQuestion`].
#[derive(Debug, Clone)]
pub struct AskUserOption {
    /// The display label for this option.
    pub label: String,
    /// An optional longer description of this option.
    pub description: Option<String>,
    /// An optional preview string (e.g. a code snippet or summary).
    pub preview: Option<String>,
}

/// Commands that can be sent into a running [`Engine`].
#[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,
    },
    /// Approve or reject a pending plan from `exit_plan_mode`.
    ApprovePlan {
        /// Whether the plan is approved.
        approved: bool,
        /// Optional feedback (used when rejecting).
        feedback: Option<String>,
    },
}

/// The final outcome produced by a completed agent run — the return
/// type of [`RunBuilder::result`] and [`RunBuilder::callback`], and the
/// payload carried by the terminal item of [`RunBuilder::stream`].
#[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 [`Engine`] with validated configuration.
pub struct EngineBuilder {
    tools: Vec<Arc<dyn Tool>>,
    context_providers: Vec<Box<dyn ContextProvider>>,
    max_iterations: usize,
    tool_timeout: Option<std::time::Duration>,
    tool_context: Option<ToolContext>,
    channel_config: ChannelConfig,
    #[cfg(feature = "mcp-client")]
    mcp_servers: Vec<Arc<dyn crate::mcp::McpServer>>,
    /// Extensions pending registration; drained into `ExtensionSet` in `build()`.
    pending_extensions: Vec<(Box<dyn crate::core::Extension>, crate::core::ErrorPolicy)>,
}

impl EngineBuilder {
    /// 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 [`Engine::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 [`Engine`]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
    }

    /// Set the capacity of the stream channel used by
    /// [`Engine::run`].
    ///
    /// Default: `256`. The channel is bounded; intermediate `Event` items may
    /// be dropped (best-effort) if the consumer is slower than the producer,
    /// but the terminal item is always delivered.
    pub fn stream_capacity(mut self, capacity: usize) -> Self {
        assert!(capacity > 0, "stream_capacity must be >= 1");
        self.channel_config.stream_capacity = Some(capacity);
        self
    }

    /// Register an extension. See [`crate::core::Extension`] and the
    /// extension architecture spec for usage. Extensions are dispatched
    /// in registration order.
    pub fn extension(mut self, ext: Box<dyn crate::core::Extension>) -> Self {
        self.pending_extensions
            .push((ext, crate::core::ErrorPolicy::Fallback));
        self
    }

    /// Register an extension with an explicit [`crate::core::ErrorPolicy`].
    ///
    /// `ErrorPolicy::Fallback` (the default for [`Self::extension`]) logs hook
    /// errors and treats them as the hook's default value. `ErrorPolicy::Abort`
    /// propagates hook errors as turn failures via [`AgentError::Internal`].
    pub fn extension_with_policy(
        mut self,
        ext: Box<dyn crate::core::Extension>,
        policy: crate::core::ErrorPolicy,
    ) -> Self {
        self.pending_extensions.push((ext, policy));
        self
    }

    /// Register the built-in extensions from a declarative
    /// [`crate::extensions::config::ExtensionsConfig`] plus runtime
    /// [`crate::extensions::config::ExtensionWiring`].
    ///
    /// This is a one-call convenience over [`Self::extension`]. The
    /// builder reuses its own [`Self::tool_context`] so callers don't
    /// pass it twice — set `.tool_context(...)` **before** calling
    /// `with_extensions_config` if you need a non-default context.
    /// The context is snapshotted at call time; a `.tool_context()`
    /// call placed *after* this method will not reach the extensions
    /// registered by this call (though it will still apply to the
    /// resulting `Engine`).
    ///
    /// Returns `Err` if the config/wiring pair is invalid (see
    /// [`crate::extensions::config::ExtensionBuildError`]).
    ///
    /// This method is fully compatible with the direct
    /// [`Self::extension`] path — you can mix config-driven
    /// registration with manual `.extension(...)` calls in either
    /// order.
    pub fn with_extensions_config(
        mut self,
        config: crate::extensions::config::ExtensionsConfig,
        wiring: crate::extensions::config::ExtensionWiring,
    ) -> std::result::Result<Self, crate::extensions::config::ExtensionBuildError> {
        let tool_context = self.tool_context.clone().unwrap_or_default();
        let set = crate::extensions::config::build_extension_set(&config, wiring, tool_context)?;
        self.pending_extensions.extend(set.into_inner());
        Ok(self)
    }

    /// Consume the builder and produce an [`Engine`].
    ///
    /// # Panics
    ///
    /// Panics if `input_capacity` or `ops_capacity` is zero, since
    /// `tokio::sync::mpsc::channel(0)` panics at runtime.
    pub fn build(self) -> Engine {
        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();
        // Aggregate tool defs from registered extensions (tool_defs() trait method).
        // This lets extensions expose their tools (ask_user, enter/exit_plan_mode,
        // planning, etc.) without any builder-level configuration.
        for (ext, _) in &self.pending_extensions {
            tool_defs.extend(ext.tool_defs());
        }
        let mut extensions = crate::core::ExtensionSet::new();
        for (ext, policy) in self.pending_extensions {
            extensions
                .add_with_policy(ext, policy)
                .unwrap_or_else(|e| panic!("duplicate extension name: {}", e.message()));
        }

        Engine {
            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(),
            channel_config: self.channel_config,
            extensions: tokio::sync::Mutex::new(extensions),
            deferred_calls: tokio::sync::Mutex::new(std::collections::HashMap::new()),
            #[cfg(feature = "mcp-client")]
            mcp_servers: self.mcp_servers,
            #[cfg(test)]
            test_parked_notify: std::sync::Mutex::new(None),
        }
    }
}

/// 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.
///
/// # Starting a turn
///
/// Every agent turn goes through [`Engine::run`], which returns a
/// [`RunBuilder`]. Chain axis setters (`.ops`, `.cancel`, `.chunked`)
/// in any order, then call one of the three terminators
/// (`.result`, `.callback`, `.stream`) to execute the turn:
///
/// ```ignore
/// use std::sync::Arc;
/// use motosan_agent_loop::{Engine, Message};
///
/// # async fn demo(llm: Arc<dyn motosan_agent_loop::LlmClient>) -> motosan_agent_loop::Result<()> {
/// let agent = Arc::new(Engine::builder().build());
///
/// // Simplest: just the final result.
/// let result = Arc::clone(&agent)
///     .run(llm.clone(), vec![Message::user("Hi!")])
///     .result()
///     .await?;
///
/// // Live token-by-token UI with AskUser support:
/// # let (_tx, ops_rx) = tokio::sync::mpsc::channel(8);
/// Arc::clone(&agent)
///     .run(llm, vec![Message::user("Ask me a question")])
///     .chunked()
///     .ops(ops_rx)
///     .callback(|event| println!("{event:?}"))
///     .await?;
/// # Ok(()) }
/// ```
///
/// See [`RunBuilder`] for the full axis/terminator matrix and
/// [`Engine::run`] for the entry point.
///
/// The loop does **not** own the LLM client; instead, [`RunBuilder`]
/// takes `Arc<dyn LlmClient>` so the same loop can be reused with
/// different backends.
pub struct Engine {
    /// 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,
    /// Channel configuration for bounded queues and backpressure.
    channel_config: ChannelConfig,
    /// Registered extensions dispatched on each iteration. Wrapped in a
    /// `Mutex` so hooks can be called through the `&self` methods of `Engine`.
    extensions: tokio::sync::Mutex<crate::core::ExtensionSet>,
    /// Tool calls suspended by extensions via `ToolDecision::Defer`,
    /// awaiting a matching `OpDecision::ResumeDeferred` from an
    /// extension's `on_op` hook. See spec §8.3.
    ///
    /// `tokio::sync::Mutex` because both `intercept_tool_call` (writes
    /// new entries) and `on_op` (removes resolved entries) can run
    /// from different async contexts within a single turn.
    #[allow(dead_code)] // Used in Phase 3B Tasks 4-5
    deferred_calls: tokio::sync::Mutex<std::collections::HashMap<String, DeferredCall>>,
    #[cfg(feature = "mcp-client")]
    mcp_servers: Vec<Arc<dyn crate::mcp::McpServer>>,

    /// Test-only hook fired by `resolve_deferred_slots` each time it
    /// is about to park on `ops_rx.recv()`. Lets tests synchronize on
    /// "engine is now waiting for an op" without wall-clock sleeps.
    /// Set via `Engine::set_test_parked_notify`. Never touched in
    /// release builds.
    ///
    /// Note: this hook fires on EVERY park, not just the first. If a
    /// deferred-slot resolution takes multiple iterations (e.g. a tool
    /// call that defers again after its first resolve), the notifier
    /// fires once per iteration. Tests that need to distinguish first-
    /// park from subsequent parks should use a dedicated Notify per
    /// expected park point, or set the field to `None` between assertions.
    #[cfg(test)]
    test_parked_notify: std::sync::Mutex<Option<Arc<tokio::sync::Notify>>>,
}

/// Result from consuming a streaming LLM response.
struct ConsumeStreamResult {
    accumulated_text: String,
    response: crate::LlmResponse,
    streaming_results: Option<(Vec<ToolCallItem>, Vec<ToolResult>)>,
    stop_reason: Option<crate::llm::StopReason>,
    /// Output tokens consumed by this stream call (delta from before/after).
    output_tokens: Option<u64>,
}

impl Engine {
    /// Create a builder with default settings.
    pub fn builder() -> EngineBuilder {
        EngineBuilder {
            tools: Vec::new(),
            context_providers: Vec::new(),
            max_iterations: 10,
            tool_timeout: None,
            tool_context: None,
            channel_config: ChannelConfig::default(),
            #[cfg(feature = "mcp-client")]
            mcp_servers: Vec::new(),
            pending_extensions: 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
    }

    /// Install a notifier fired each time `resolve_deferred_slots`
    /// is about to park on `ops_rx.recv()`. The same notifier is
    /// reused across parks within a single turn; see the field doc
    /// on `test_parked_notify` for multi-iteration caveats.
    /// Test-only — not part of the public API.
    #[cfg(test)]
    #[allow(dead_code)]
    pub(crate) fn set_test_parked_notify(&self, notify: Arc<tokio::sync::Notify>) {
        *self.test_parked_notify.lock().unwrap() = Some(notify);
    }

    /// 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)
    }

    #[cfg(feature = "cancellation")]
    async fn run_streaming_with_cancel_inner(
        &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>, Vec<Message>, RunTerminalMeta) {
        use futures::StreamExt;

        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let prepared = match self.prepare_messages(messages).await {
            Ok(m) => m,
            Err(e) => return run_return_failure(e, Vec::new(), 0, TokenUsage::default()),
        };
        let mut state = TurnState::new(prepared);

        for iteration in 1..=self.max_iterations {
            // Pre-turn: cancellation check.
            if cancel.is_cancelled() {
                return run_return_failure(
                    AgentError::Cancelled,
                    state.messages,
                    iteration,
                    state.total_usage,
                );
            }
            on_event(AgentEvent::Core(CoreEvent::IterationStarted { iteration }));

            // before_iteration extension hook (after IterationStarted, before request build).
            if let Err(e) = self
                .dispatch_before_iteration(iteration, &mut state.messages, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            // Autocompact check before LLM call (after IterationStarted).
            if let Err(e) = self
                .maybe_compact_via_extensions(&mut state.messages, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            loop {
                // Streaming LLM step with cancellation race and eager tool execution.
                let (
                    accumulated_text,
                    response,
                    streaming_results,
                    stream_stop_reason,
                    output_tokens,
                ) = {
                    let mut stream = llm.chat_stream(&state.messages, tools.tool_defs());
                    let mut accumulated = String::new();
                    let mut final_response: Option<crate::LlmResponse> = None;
                    let mut executor = crate::streaming_executor::StreamingToolExecutor::new();
                    let mut submitted_ids: std::collections::HashSet<String> =
                        std::collections::HashSet::new();
                    let mut stop_reason: Option<crate::llm::StopReason> = None;

                    // Snapshot output tokens before this stream to compute the delta.
                    let output_tokens_before = state.total_usage.output_tokens;

                    loop {
                        tokio::select! {
                            chunk_opt = stream.next() => {
                                match chunk_opt {
                                    Some(chunk_result) => {
                                        let chunk = match chunk_result {
                                            Ok(c) => c,
                                            Err(e) => return run_return_failure(e, state.messages.clone(), iteration, state.total_usage),
                                        };
                                        match chunk {
                                            crate::llm::StreamChunk::TextDelta(delta) => {
                                                accumulated.push_str(&delta);
                                                on_event(AgentEvent::Core(CoreEvent::TextChunk(delta)));
                                            }
                                            crate::llm::StreamChunk::ToolUse { id, name, args } => {
                                                if !submitted_ids.contains(&id) {
                                                    let item = crate::llm::ToolCallItem {
                                                        id: id.clone(),
                                                        name: name.clone(),
                                                        args,
                                                    };
                                                    on_event(AgentEvent::Core(CoreEvent::ToolStarted {
                                                        name: name.clone(),
                                                    }));
                                                    executor.submit(
                                                        item,
                                                        tools.tool_map(),
                                                        self.tool_timeout,
                                                        &self.tool_context,
                                                    );
                                                    submitted_ids.insert(id);
                                                }
                                            }
                                            crate::llm::StreamChunk::Done(resp) => {
                                                final_response = Some(resp);
                                            }
                                            crate::llm::StreamChunk::Usage(usage) => {
                                                state.total_usage.accumulate(usage);
                                            }
                                            crate::llm::StreamChunk::StopReason(reason) => {
                                                stop_reason = Some(reason);
                                            }
                                        }
                                    }
                                    None => break,
                                }
                            }
                            _ = cancel.cancelled() => {
                                return run_return_failure(AgentError::Cancelled, state.messages.clone(), iteration, state.total_usage);
                            }
                        }
                    }

                    let resp = final_response
                        .unwrap_or_else(|| crate::LlmResponse::Message(accumulated.clone()));
                    let sr = if executor.has_pending() {
                        Some(executor.collect().await)
                    } else {
                        None
                    };
                    let delta = state.total_usage.output_tokens - output_tokens_before;
                    let ot = if delta > 0 { Some(delta) } else { None };
                    (accumulated, resp, sr, stop_reason, ot)
                };

                match response {
                    crate::LlmResponse::Message(text) => {
                        if accumulated_text.is_empty() && !text.is_empty() {
                            on_event(AgentEvent::Core(CoreEvent::TextChunk(text.clone())));
                        }
                        on_event(AgentEvent::Core(CoreEvent::TextDone(text.clone())));

                        let meta = crate::llm::LlmResponseMeta {
                            stop_reason: stream_stop_reason,
                            output_tokens,
                        };
                        let outcome = match self
                            .handle_text_response(text, meta, &mut state, &on_event)
                            .await
                        {
                            Ok(o) => o,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };

                        match outcome {
                            TextResponseOutcome::Continue => continue,
                            TextResponseOutcome::Halt(msg) => {
                                return run_return_failure(
                                    AgentError::Internal(format!(
                                        "Turn halted by extension: {}",
                                        msg
                                    )),
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                            TextResponseOutcome::Finalize(text) => {
                                state.messages.push(Message::assistant(&text));
                                let res = state.into_result(text, iteration);
                                let snapshot = res.messages.clone();
                                return run_return_success(res, snapshot);
                            }
                        }
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        if let Some((streamed_items, streamed_results)) = streaming_results {
                            let (final_items, final_results) = merge_streamed_tool_results(
                                &items,
                                streamed_items,
                                streamed_results,
                                tools.tool_map(),
                                self.tool_timeout,
                                &self.tool_context,
                                &on_event,
                            )
                            .await;
                            if let Err(e) = self
                                .finalize_tool_call_batch(
                                    &final_items,
                                    final_results,
                                    iteration,
                                    tools.tool_defs(),
                                    &mut state,
                                    &on_event,
                                )
                                .await
                            {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                        } else {
                            emit_tool_started(&items, &on_event);
                            let slots = self.dispatch_intercept_tool_calls(&items, &on_event).await;
                            // Resolve any Deferred slots. No ops channel in this path.
                            let mut tmp_ops_state = OpsState::default();
                            let slots = match self
                                .resolve_deferred_slots(
                                    slots,
                                    &mut tmp_ops_state,
                                    &mut None,
                                    &on_event,
                                )
                                .await
                            {
                                Ok(s) => s,
                                Err(e) => {
                                    return run_return_failure(
                                        e,
                                        state.messages,
                                        iteration,
                                        state.total_usage,
                                    )
                                }
                            };
                            let mut remaining_items: Vec<crate::llm::ToolCallItem> = Vec::new();
                            let mut remaining_indices: Vec<usize> = Vec::new();
                            let mut final_results: Vec<Option<motosan_agent_tool::ToolResult>> =
                                slots
                                    .into_iter()
                                    .map(|s| match s {
                                        InterceptedSlot::Resolved(r) => Some(r),
                                        InterceptedSlot::Pending => None,
                                        InterceptedSlot::Deferred(_) => {
                                            unreachable!(
                                                "InterceptedSlot::Deferred should have been \
                                                 resolved by resolve_deferred_slots"
                                            );
                                        }
                                    })
                                    .collect();
                            for (i, item) in items.iter().enumerate() {
                                if final_results[i].is_none() {
                                    remaining_items.push(item.clone());
                                    remaining_indices.push(i);
                                }
                            }
                            let execution_results = Self::execute_tools_parallel(
                                tools.tool_map(),
                                &remaining_items,
                                self.tool_timeout,
                                &self.tool_context,
                            )
                            .await;
                            for (exec_idx, original_idx) in remaining_indices.iter().enumerate() {
                                final_results[*original_idx] =
                                    Some(execution_results[exec_idx].clone());
                            }
                            let results: Vec<motosan_agent_tool::ToolResult> = final_results
                                .into_iter()
                                .map(|opt| opt.expect("all slots filled after intercept + execute"))
                                .collect();
                            if let Err(e) = self
                                .finalize_tool_call_batch(
                                    &items,
                                    results,
                                    iteration,
                                    tools.tool_defs(),
                                    &mut state,
                                    &on_event,
                                )
                                .await
                            {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                        }
                        break; // Back to outer iteration loop.
                    }
                }
            }
        }

        run_return_failure(
            AgentError::MaxIterations(self.max_iterations),
            state.messages,
            self.max_iterations,
            state.total_usage,
        )
    }

    /// Consume a streaming LLM response, forwarding text deltas as events,
    /// accumulating usage, and eagerly submitting tool calls via
    /// [`StreamingToolExecutor`](crate::StreamingToolExecutor).
    ///
    /// When a [`StreamChunk::ToolUse`] is received mid-stream the tool starts
    /// executing immediately. After the stream ends the executor's results
    /// are collected. This overlaps tool execution with LLM output generation.
    ///
    /// Returns the accumulated text, the final [`LlmResponse`](crate::LlmResponse),
    /// and optionally the collected streaming tool results (items + results)
    /// if any tools were submitted during the stream.
    #[allow(clippy::too_many_arguments)]
    async fn consume_stream(
        llm: &dyn LlmClient,
        messages: &[Message],
        tool_defs: &[ToolDef],
        tool_map: &HashMap<String, Arc<dyn Tool>>,
        tool_timeout: Option<std::time::Duration>,
        tool_context: &motosan_agent_tool::ToolContext,
        total_usage: &mut TokenUsage,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<ConsumeStreamResult> {
        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;
        let mut executor = crate::streaming_executor::StreamingToolExecutor::new();
        let mut submitted_ids: std::collections::HashSet<String> = std::collections::HashSet::new();
        let mut stop_reason: Option<crate::llm::StopReason> = None;

        // Snapshot output tokens before this stream to compute the delta.
        let output_tokens_before = total_usage.output_tokens;

        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::Core(CoreEvent::TextChunk(delta)));
                }
                crate::llm::StreamChunk::ToolUse { id, name, args } => {
                    if !submitted_ids.contains(&id) {
                        let item = ToolCallItem {
                            id: id.clone(),
                            name: name.clone(),
                            args,
                        };
                        on_event(AgentEvent::Core(CoreEvent::ToolStarted {
                            name: name.clone(),
                        }));
                        executor.submit(item, tool_map, tool_timeout, tool_context);
                        submitted_ids.insert(id);
                    }
                }
                crate::llm::StreamChunk::Done(response) => {
                    final_response = Some(response);
                }
                crate::llm::StreamChunk::Usage(usage) => {
                    total_usage.accumulate(usage);
                }
                crate::llm::StreamChunk::StopReason(reason) => {
                    stop_reason = Some(reason);
                }
            }
        }

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

        let streaming_results = if executor.has_pending() {
            Some(executor.collect().await)
        } else {
            None
        };

        // Compute output token delta for this call.
        let output_tokens_delta = total_usage.output_tokens - output_tokens_before;
        let output_tokens = if output_tokens_delta > 0 {
            Some(output_tokens_delta)
        } else {
            None
        };

        Ok(ConsumeStreamResult {
            accumulated_text: accumulated,
            response,
            streaming_results,
            stop_reason,
            output_tokens,
        })
    }

    /// 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.text())
            .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>, Vec<Message>, RunTerminalMeta) {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let prepared = match self.prepare_messages(messages).await {
            Ok(m) => m,
            Err(e) => return run_return_failure(e, Vec::new(), 0, TokenUsage::default()),
        };
        let mut state = TurnState::new(prepared);

        for iteration in 1..=self.max_iterations {
            // Pre-turn: cancellation check.
            if cancel.is_cancelled() {
                return run_return_failure(
                    AgentError::Cancelled,
                    state.messages,
                    iteration,
                    state.total_usage,
                );
            }
            on_event(AgentEvent::Core(CoreEvent::IterationStarted { iteration }));

            // before_iteration extension hook (after IterationStarted, before request build).
            if let Err(e) = self
                .dispatch_before_iteration(iteration, &mut state.messages, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            // Autocompact check before LLM call (after IterationStarted).
            if let Err(e) = self
                .maybe_compact_via_extensions(&mut state.messages, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            loop {
                // LLM step with cancellation race.
                let output = tokio::select! {
                    output = llm.chat(&state.messages, tools.tool_defs()) => match output {
                        Ok(o) => o,
                        Err(e) => return run_return_failure(e, state.messages, iteration, state.total_usage),
                    },
                    _ = cancel.cancelled() => return run_return_failure(AgentError::Cancelled, state.messages, iteration, state.total_usage),
                };

                state.accumulate_usage(output.usage);
                let stop_reason = output.stop_reason;
                let output_tokens = output.usage.map(|u| u.output_tokens);

                match output.response {
                    crate::LlmResponse::Message(text) => {
                        on_event(AgentEvent::Core(CoreEvent::TextChunk(text.clone())));

                        let meta = crate::llm::LlmResponseMeta {
                            stop_reason,
                            output_tokens,
                        };
                        let outcome = match self
                            .handle_text_response(text, meta, &mut state, on_event)
                            .await
                        {
                            Ok(o) => o,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };

                        match outcome {
                            TextResponseOutcome::Continue => continue,
                            TextResponseOutcome::Halt(msg) => {
                                return run_return_failure(
                                    AgentError::Internal(format!(
                                        "Turn halted by extension: {}",
                                        msg
                                    )),
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                            TextResponseOutcome::Finalize(text) => {
                                state.messages.push(Message::assistant(&text));
                                let res = state.into_result(text, iteration);
                                let snapshot = res.messages.clone();
                                return run_return_success(res, snapshot);
                            }
                        }
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        emit_tool_started(&items, on_event);
                        let slots = self.dispatch_intercept_tool_calls(&items, on_event).await;
                        // Resolve any Deferred slots. No ops channel in this path.
                        let mut tmp_ops_state = OpsState::default();
                        let slots = match self
                            .resolve_deferred_slots(slots, &mut tmp_ops_state, &mut None, on_event)
                            .await
                        {
                            Ok(s) => s,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };
                        let mut remaining_items: Vec<crate::llm::ToolCallItem> = Vec::new();
                        let mut remaining_indices: Vec<usize> = Vec::new();
                        let mut final_results: Vec<Option<motosan_agent_tool::ToolResult>> = slots
                            .into_iter()
                            .map(|s| match s {
                                InterceptedSlot::Resolved(r) => Some(r),
                                InterceptedSlot::Pending => None,
                                InterceptedSlot::Deferred(_) => {
                                    unreachable!(
                                        "InterceptedSlot::Deferred should have been \
                                         resolved by resolve_deferred_slots"
                                    );
                                }
                            })
                            .collect();
                        for (i, item) in items.iter().enumerate() {
                            if final_results[i].is_none() {
                                remaining_items.push(item.clone());
                                remaining_indices.push(i);
                            }
                        }
                        let execution_results = Self::execute_tools_parallel(
                            tools.tool_map(),
                            &remaining_items,
                            self.tool_timeout,
                            &self.tool_context,
                        )
                        .await;
                        for (exec_idx, original_idx) in remaining_indices.iter().enumerate() {
                            final_results[*original_idx] =
                                Some(execution_results[exec_idx].clone());
                        }
                        let results: Vec<motosan_agent_tool::ToolResult> = final_results
                            .into_iter()
                            .map(|opt| opt.expect("all slots filled after intercept + execute"))
                            .collect();
                        if let Err(e) = self
                            .finalize_tool_call_batch(
                                &items,
                                results,
                                iteration,
                                tools.tool_defs(),
                                &mut state,
                                on_event,
                            )
                            .await
                        {
                            return run_return_failure(
                                e,
                                state.messages,
                                iteration,
                                state.total_usage,
                            );
                        }
                        break;
                    }
                }
            }
        }

        run_return_failure(
            AgentError::MaxIterations(self.max_iterations),
            state.messages,
            self.max_iterations,
            state.total_usage,
        )
    }

    /// Internal turn driver.
    ///
    /// Returns a tuple of `(Result<AgentResult>, Vec<Message>)` so callers
    /// (notably `run`) can recover the final conversation history even
    /// when the loop terminated with an error. On success, the returned
    /// `Vec<Message>` is a clone of `AgentResult.messages` for ergonomic
    /// access; on failure it contains the partial history accumulated so far.
    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>, Vec<Message>, RunTerminalMeta) {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let prepared = match self.prepare_messages(messages).await {
            Ok(m) => m,
            Err(e) => return run_return_failure(e, Vec::new(), 0, TokenUsage::default()),
        };
        let mut state = TurnState::new(prepared);
        let mut ops_state = OpsState::default();

        for iteration in 1..=self.max_iterations {
            // Pre-turn: ingest pending ops.
            if let Err(e) = self
                .drain_ops(&mut state.messages, &mut ops_rx, &mut ops_state, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }
            if ops_state.interrupted {
                on_event(AgentEvent::Core(CoreEvent::Interrupted));
                let res =
                    state.into_result("(interrupted)".to_string(), iteration.saturating_sub(1));
                let snapshot = res.messages.clone();
                return run_return_interrupted(res, snapshot);
            }

            on_event(AgentEvent::Core(CoreEvent::IterationStarted { iteration }));

            // before_iteration extension hook (after IterationStarted, before request build).
            if let Err(e) = self
                .dispatch_before_iteration(iteration, &mut state.messages, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            // Autocompact check before LLM call (after IterationStarted).
            if let Err(e) = self
                .maybe_compact_via_extensions(&mut state.messages, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            loop {
                // LLM step.
                let output = match llm.chat(&state.messages, tools.tool_defs()).await {
                    Ok(o) => o,
                    Err(e) => {
                        return run_return_failure(e, state.messages, iteration, state.total_usage);
                    }
                };

                state.accumulate_usage(output.usage);
                let stop_reason = output.stop_reason;
                let output_tokens = output.usage.map(|u| u.output_tokens);

                match output.response {
                    crate::LlmResponse::Message(text) => {
                        on_event(AgentEvent::Core(CoreEvent::TextChunk(text.clone())));

                        let meta = crate::llm::LlmResponseMeta {
                            stop_reason,
                            output_tokens,
                        };
                        let outcome = match self
                            .handle_text_response(text, meta, &mut state, on_event)
                            .await
                        {
                            Ok(o) => o,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };

                        match outcome {
                            TextResponseOutcome::Continue => continue,
                            TextResponseOutcome::Halt(msg) => {
                                return run_return_failure(
                                    AgentError::Internal(format!(
                                        "Turn halted by extension: {}",
                                        msg
                                    )),
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                            TextResponseOutcome::Finalize(text) => {
                                state.messages.push(Message::assistant(&text));
                                let res = state.into_result(text, iteration);
                                let snapshot = res.messages.clone();
                                return run_return_success(res, snapshot);
                            }
                        }
                    }
                    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;
                        if let Err(e) = self
                            .finalize_tool_call_batch(
                                &items,
                                results,
                                iteration,
                                tools.tool_defs(),
                                &mut state,
                                on_event,
                            )
                            .await
                        {
                            return run_return_failure(
                                e,
                                state.messages,
                                iteration,
                                state.total_usage,
                            );
                        }
                        break; // Back to outer iteration loop.
                    }
                }
            }
        }

        run_return_failure(
            AgentError::MaxIterations(self.max_iterations),
            state.messages,
            self.max_iterations,
            state.total_usage,
        )
    }

    /// Inner function covering (batch LLM + ops_rx + cancel). Combines
    /// the ops drain-and-dispatch loop from `run_inner_with_ops` with
    /// the pre-iteration cancel check from `run_inner_cancel`. Returns
    /// the same `(Result<AgentResult>, Vec<Message>)` tuple shape as
    /// `run_inner_with_ops` for dispatch uniformity.
    ///
    /// # Cancel semantics
    ///
    /// The cancellation token is checked at two points per iteration:
    ///
    /// 1. **Before `drain_ops` at iteration top.** A pending cancel wins
    ///    over any queued ops — the function returns
    ///    `Err(AgentError::Cancelled)` with no `Interrupted` event and
    ///    any unread ops in the channel are dropped on the floor. This
    ///    matches the intuition that cancel is stronger than interrupt.
    ///
    /// 2. **Inside the `tokio::select!` around `llm.chat()`.** If the
    ///    token is tripped during the LLM call, the in-flight chat future
    ///    is dropped and the function returns `Err(AgentError::Cancelled)`.
    ///
    /// # Known latency gap: tool-call deferred-slot resolution
    ///
    /// During tool-call execution, this function calls
    /// `execute_tools_with_policy`, which in turn calls
    /// `resolve_deferred_slots` to wait for matching ops (e.g. the
    /// `AskUserAnswer` for an `ask_user` tool call). The cancel token is
    /// **not** raced inside `resolve_deferred_slots`, so a cancel that
    /// arrives while an `ask_user` slot is still awaiting its answer will
    /// not be observed until the slot resolves (or the 60 s defer
    /// timeout fires). This is a known limitation inherited from the
    /// current `execute_tools_with_policy` signature, which does not
    /// accept a cancel token. Fixing it requires threading cancel through
    /// to `resolve_deferred_slots`, tracked as a follow-up task.
    #[cfg(feature = "cancellation")]
    #[allow(dead_code)] // Called by RunBuilder dispatch in a later task
    async fn run_inner_with_ops_and_cancel(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        extra_tools: &[Arc<dyn Tool>],
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
        ops_rx: mpsc::Receiver<AgentOp>,
        cancel: &tokio_util::sync::CancellationToken,
    ) -> (Result<AgentResult>, Vec<Message>, RunTerminalMeta) {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let prepared = match self.prepare_messages(messages).await {
            Ok(m) => m,
            Err(e) => return run_return_failure(e, Vec::new(), 0, TokenUsage::default()),
        };
        let mut state = TurnState::new(prepared);
        let mut ops_state = OpsState::default();
        let mut ops_rx_opt = Some(ops_rx);

        for iteration in 1..=self.max_iterations {
            // Cancel check (from run_inner_cancel).
            if cancel.is_cancelled() {
                return run_return_failure(
                    AgentError::Cancelled,
                    state.messages,
                    iteration,
                    state.total_usage,
                );
            }

            // Drain pending ops (from run_inner_with_ops).
            if let Err(e) = self
                .drain_ops(
                    &mut state.messages,
                    &mut ops_rx_opt,
                    &mut ops_state,
                    on_event,
                )
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }
            if ops_state.interrupted {
                on_event(AgentEvent::Core(CoreEvent::Interrupted));
                let res =
                    state.into_result("(interrupted)".to_string(), iteration.saturating_sub(1));
                let snapshot = res.messages.clone();
                return run_return_interrupted(res, snapshot);
            }

            on_event(AgentEvent::Core(CoreEvent::IterationStarted { iteration }));

            if let Err(e) = self
                .dispatch_before_iteration(iteration, &mut state.messages, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            if let Err(e) = self
                .maybe_compact_via_extensions(&mut state.messages, on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            loop {
                // LLM step raced against cancel token.
                let output = tokio::select! {
                    output = llm.chat(&state.messages, tools.tool_defs()) => match output {
                        Ok(o) => o,
                        Err(e) => return run_return_failure(e, state.messages, iteration, state.total_usage),
                    },
                    _ = cancel.cancelled() => {
                        return run_return_failure(AgentError::Cancelled, state.messages, iteration, state.total_usage);
                    }
                };

                state.accumulate_usage(output.usage);
                let stop_reason = output.stop_reason;
                let output_tokens = output.usage.map(|u| u.output_tokens);

                match output.response {
                    crate::LlmResponse::Message(text) => {
                        on_event(AgentEvent::Core(CoreEvent::TextChunk(text.clone())));

                        let meta = crate::llm::LlmResponseMeta {
                            stop_reason,
                            output_tokens,
                        };
                        let outcome = match self
                            .handle_text_response(text, meta, &mut state, on_event)
                            .await
                        {
                            Ok(o) => o,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };

                        match outcome {
                            TextResponseOutcome::Continue => continue,
                            TextResponseOutcome::Halt(msg) => {
                                return run_return_failure(
                                    AgentError::Internal(format!(
                                        "Turn halted by extension: {}",
                                        msg
                                    )),
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                            TextResponseOutcome::Finalize(text) => {
                                state.messages.push(Message::assistant(&text));
                                let res = state.into_result(text, iteration);
                                let snapshot = res.messages.clone();
                                return run_return_success(res, snapshot);
                            }
                        }
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        emit_tool_started(&items, on_event);
                        let results = self
                            .execute_tools_with_policy(
                                tools.tool_map(),
                                &items,
                                &mut state.messages,
                                &mut ops_rx_opt,
                                &mut ops_state,
                                on_event,
                            )
                            .await;
                        if let Err(e) = self
                            .finalize_tool_call_batch(
                                &items,
                                results,
                                iteration,
                                tools.tool_defs(),
                                &mut state,
                                on_event,
                            )
                            .await
                        {
                            return run_return_failure(
                                e,
                                state.messages,
                                iteration,
                                state.total_usage,
                            );
                        }
                        break;
                    }
                }
            }
        }

        run_return_failure(
            AgentError::MaxIterations(self.max_iterations),
            state.messages,
            self.max_iterations,
            state.total_usage,
        )
    }

    /// Estimate the token count of a message list using a simple heuristic
    /// (4 characters per token). Used in tests.
    #[allow(dead_code)]
    fn estimate_tokens(messages: &[Message]) -> usize {
        messages.iter().map(|m| m.approx_visible_chars() / 4).sum()
    }

    // TODO(phase-4): delete maybe_compact_via_extensions entirely.
    // After the public API cutover, consumers observe CoreEvent +
    // ExtensionEvent directly and the AgentEvent::AutoCompacted
    // forwarding adapter below (plus the variant itself in the
    // AgentEvent enum) becomes unnecessary.

    /// Run the extension pipeline's `transform_request` hook and forward any
    /// `AutocompactEvent` as the legacy `AgentEvent::AutoCompacted` variant.
    ///
    /// Short-circuits immediately when no extensions are registered so the
    /// hot path pays no overhead.
    async fn maybe_compact_via_extensions(
        &self,
        messages: &mut Vec<Message>,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<()> {
        // Build an AgentState snapshot for the hook.
        let tools_snapshot: Vec<motosan_agent_tool::ToolDef> = self.tool_defs.to_vec();
        let agent_state = crate::core::AgentState {
            iteration: 0,
            messages: messages.as_slice(),
            tools: &tools_snapshot,
            turn_started_at: std::time::Instant::now(),
            token_usage: TokenUsage::default(),
        };

        // CaptureSink collects any events the extension emits.
        let mut sink = crate::core::CaptureSink::default();

        // Lock the ExtensionSet and call the extension pipeline.
        // Short-circuit on empty set (common case when no extensions are registered).
        // ErrorPolicy::Fallback is handled inside ExtensionSet::transform_request —
        // a failing extension is logged and skipped; messages are left unchanged.
        // Only ErrorPolicy::Abort extensions surface an error here.
        let transformed = {
            let mut set = self.extensions.lock().await;
            if set.is_empty() {
                return Ok(());
            }
            set.transform_request(messages.clone(), &agent_state, &mut sink)
                .await
                .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?
        };

        *messages = transformed;

        // Forward any extension events captured during transform_request.
        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        Ok(())
    }

    /// Dispatch `Extension::before_iteration` at the top of each iteration.
    ///
    /// Mirrors `maybe_compact_via_extensions` in shape: builds an
    /// `AgentState` snapshot, locks the `ExtensionSet`, and propagates
    /// per-extension `ErrorPolicy` through `ExtensionSet::before_iteration`.
    /// `ErrorPolicy::Fallback` is handled inside the set; only `Abort`
    /// surfaces an error here, which is wrapped as `AgentError::Internal`.
    ///
    /// Returned `FlowDecision`:
    /// - `Continue`: caller proceeds to the LLM step normally.
    /// - `Inject(msg)`: the message is appended to history and the caller
    ///   proceeds to the LLM step (same iteration) — semantically the
    ///   extension nudged the conversation before the model sees it.
    /// - `Halt(reason)`: the turn is terminated with
    ///   `AgentError::Internal(reason.message)`, mirroring how
    ///   `handle_text_response` treats `Halt`.
    ///
    /// Short-circuits immediately when no extensions are registered so the
    /// hot path pays no overhead.
    async fn dispatch_before_iteration(
        &self,
        iteration: usize,
        messages: &mut Vec<Message>,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<()> {
        if self.extensions.lock().await.is_empty() {
            return Ok(());
        }

        let tools_snapshot: Vec<motosan_agent_tool::ToolDef> = self.tool_defs.to_vec();
        let agent_state = crate::core::AgentState {
            iteration,
            messages: messages.as_slice(),
            tools: &tools_snapshot,
            turn_started_at: std::time::Instant::now(),
            token_usage: TokenUsage::default(),
        };

        let mut sink = crate::core::CaptureSink::default();
        let decision = self
            .extensions
            .lock()
            .await
            .before_iteration(&agent_state, &mut sink)
            .await
            .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?;

        // Forward any captured extension events.
        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        match decision {
            FlowDecision::Continue => Ok(()),
            FlowDecision::Inject(msg) => {
                messages.push(msg);
                Ok(())
            }
            FlowDecision::Halt(reason) => Err(crate::error::AgentError::Internal(reason.message)),
        }
    }

    // TODO(phase-4): delete dispatch_after_llm_response once the public API
    // cutover deletes AgentEvent::TokenBudgetContinuation.

    /// Dispatch `after_llm_response` via the extension system.
    ///
    /// Returns the `FlowDecision` from the first extension that returned
    /// something other than `Continue`, or `Continue` if all extensions
    /// passed. The caller is responsible for interpreting the decision:
    /// - `Continue`: proceed to tool dispatch as normal
    /// - `Inject(msg)`: append `msg` to history and loop back to another
    ///   LLM call (do NOT proceed to tool dispatch)
    /// - `Halt(reason)`: terminate the turn with the halt reason
    ///
    /// Also forwards `TokenBudgetEvent::Continuation` to the legacy
    /// `AgentEvent::TokenBudgetContinuation` event stream for backward
    /// compatibility. Phase 4 will delete the forwarding.
    async fn dispatch_after_llm_response(
        &self,
        resp: &crate::llm::LlmResponse,
        meta: &crate::llm::LlmResponseMeta,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<FlowDecision> {
        if self.extensions.lock().await.is_empty() {
            return Ok(FlowDecision::Continue);
        }

        let tools_snapshot = self.tool_defs.to_vec();
        // TODO(phase-4): pass the real messages through instead of &[].
        // TokenBudgetExtension doesn't read state.messages so this is OK now.
        let empty_messages: Vec<crate::message::Message> = vec![];
        let agent_state = crate::core::AgentState {
            iteration: 0,
            messages: &empty_messages,
            tools: &tools_snapshot,
            turn_started_at: std::time::Instant::now(),
            token_usage: crate::llm::TokenUsage::default(),
        };

        let mut sink = crate::core::CaptureSink::default();
        let decision = self
            .extensions
            .lock()
            .await
            .after_llm_response(resp, meta, &agent_state, &mut sink)
            .await
            .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?;

        // Forward captured extension events.
        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        Ok(decision)
    }

    fn agent_state_snapshot<'a>(
        &self,
        iteration: usize,
        messages: &'a [crate::message::Message],
        tools: &'a [motosan_agent_tool::ToolDef],
        token_usage: crate::llm::TokenUsage,
    ) -> crate::core::AgentState<'a> {
        crate::core::AgentState {
            iteration,
            messages,
            tools,
            turn_started_at: std::time::Instant::now(),
            token_usage,
        }
    }

    /// Dispatch `Extension::after_tool_result` for a single tool result.
    /// Runs the full pipeline; returns the first non-`Continue` decision
    /// encountered, or `Continue` if all extensions returned `Continue`.
    ///
    /// Captured extension events are forwarded via `on_event`.
    async fn dispatch_after_tool_result(
        &self,
        result: &motosan_agent_tool::ToolResult,
        agent_state: &crate::core::AgentState<'_>,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<FlowDecision> {
        if self.extensions.lock().await.is_empty() {
            return Ok(FlowDecision::Continue);
        }

        let mut sink = crate::core::CaptureSink::default();
        let decision = self
            .extensions
            .lock()
            .await
            .after_tool_result(result, agent_state, &mut sink)
            .await
            .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?;

        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        Ok(decision)
    }

    /// Dispatch `Extension::rewrite_tool_result` pipeline for a single
    /// tool call's result. Returns the final (possibly rewritten) result.
    /// Captured extension events are forwarded via `on_event`.
    async fn dispatch_rewrite_tool_result(
        &self,
        call: &crate::llm::ToolCallItem,
        initial: motosan_agent_tool::ToolResult,
        agent_state: &crate::core::AgentState<'_>,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<motosan_agent_tool::ToolResult> {
        if self.extensions.lock().await.is_empty() {
            return Ok(initial);
        }

        let mut sink = crate::core::CaptureSink::default();
        let final_result = self
            .extensions
            .lock()
            .await
            .rewrite_tool_result(call, initial, agent_state, &mut sink)
            .await
            .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?;

        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        Ok(final_result)
    }

    /// Dispatch `Extension::on_terminal` across all extensions.
    async fn dispatch_on_terminal(
        &self,
        result: &crate::core::extension::TurnResult,
        iteration: usize,
        messages: &[crate::message::Message],
        tools: &[motosan_agent_tool::ToolDef],
        token_usage: crate::llm::TokenUsage,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<()> {
        if self.extensions.lock().await.is_empty() {
            return Ok(());
        }

        let agent_state = self.agent_state_snapshot(iteration, messages, tools, token_usage);
        let mut sink = crate::core::CaptureSink::default();
        self.extensions
            .lock()
            .await
            .on_terminal(result, &agent_state, &mut sink)
            .await
            .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?;

        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        Ok(())
    }

    async fn dispatch_on_terminal_from_meta(
        &self,
        meta: &RunTerminalMeta,
        messages: &[crate::message::Message],
        tools: &[motosan_agent_tool::ToolDef],
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<()> {
        self.dispatch_on_terminal(
            &meta.turn_result,
            meta.iteration,
            messages,
            tools,
            meta.token_usage,
            on_event,
        )
        .await
    }

    /// Handle a text response from the LLM inside a run method's inner loop.
    ///
    /// Dispatches `after_llm_response` to the extension system, interprets
    /// the returned `FlowDecision`, and returns a `TextResponseOutcome` that
    /// tells the caller what to do next:
    ///
    /// - `Finalize`: the caller should push `Message::assistant(text)` and
    ///   build the `AgentResult`
    /// - `Continue`: the extension injected a continuation; the caller's
    ///   inner LLM loop should run another iteration
    /// - `Halt`: the extension halted the turn; the caller should return
    ///   `Err(AgentError::Internal(...))`
    ///
    /// The `TextChunk`/`TextDone` events are NOT emitted here — call sites
    /// emit them before calling this helper (streaming variants emit chunks
    /// during the stream itself; non-streaming variants emit `TextChunk`
    /// synchronously before dispatching).
    async fn handle_text_response(
        &self,
        text: String,
        meta: crate::llm::LlmResponseMeta,
        state: &mut TurnState,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<TextResponseOutcome> {
        let resp_for_ext = crate::LlmResponse::Message(text.clone());
        let decision = self
            .dispatch_after_llm_response(&resp_for_ext, &meta, on_event)
            .await?;

        match decision {
            FlowDecision::Continue => Ok(TextResponseOutcome::Finalize(text)),
            FlowDecision::Inject(msg) => {
                state.continuation_text.push_str(&text);
                state.messages.push(Message::assistant(&text));
                state.messages.push(msg);
                Ok(TextResponseOutcome::Continue)
            }
            FlowDecision::Halt(reason) => Ok(TextResponseOutcome::Halt(reason.message)),
        }
    }

    /// Shared body of the streaming iteration
    /// loop. `ops_rx` is threaded through so deferred tool slots and
    /// the pre-iteration drain stage both have access to the real
    /// channel when present.
    async fn run_streaming_inner(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        extra_tools: &[Arc<dyn Tool>],
        mut ops_rx: Option<mpsc::Receiver<AgentOp>>,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> (Result<AgentResult>, Vec<Message>, RunTerminalMeta) {
        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let prepared = match self.prepare_messages(messages).await {
            Ok(m) => m,
            Err(e) => return run_return_failure(e, Vec::new(), 0, TokenUsage::default()),
        };
        let mut state = TurnState::new(prepared);
        let mut ops_state = OpsState::default();

        for iteration in 1..=self.max_iterations {
            // Pre-iteration: ingest any pending ops routed through the
            // extension pipeline (dispatch_on_op). This matches the
            // ordering in `run_inner_with_ops` so callback and
            // stream-item paths stay dispatch-equivalent.
            if let Err(e) = self
                .drain_ops(&mut state.messages, &mut ops_rx, &mut ops_state, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }
            if ops_state.interrupted {
                on_event(AgentEvent::Core(CoreEvent::Interrupted));
                let res =
                    state.into_result("(interrupted)".to_string(), iteration.saturating_sub(1));
                let snapshot = res.messages.clone();
                return run_return_interrupted(res, snapshot);
            }

            on_event(AgentEvent::Core(CoreEvent::IterationStarted { iteration }));

            // before_iteration extension hook (after IterationStarted, before request build).
            if let Err(e) = self
                .dispatch_before_iteration(iteration, &mut state.messages, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            // Autocompact check before LLM call (after IterationStarted).
            if let Err(e) = self
                .maybe_compact_via_extensions(&mut state.messages, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            loop {
                // Streaming LLM step with eager tool execution.
                let stream_result = match Self::consume_stream(
                    llm,
                    &state.messages,
                    tools.tool_defs(),
                    tools.tool_map(),
                    self.tool_timeout,
                    &self.tool_context,
                    &mut state.total_usage,
                    &on_event,
                )
                .await
                {
                    Ok(r) => r,
                    Err(e) => {
                        return run_return_failure(e, state.messages, iteration, state.total_usage)
                    }
                };

                match stream_result.response {
                    crate::LlmResponse::Message(text) => {
                        if stream_result.accumulated_text.is_empty() && !text.is_empty() {
                            on_event(AgentEvent::Core(CoreEvent::TextChunk(text.clone())));
                        }
                        on_event(AgentEvent::Core(CoreEvent::TextDone(text.clone())));

                        let meta = crate::llm::LlmResponseMeta {
                            stop_reason: stream_result.stop_reason,
                            output_tokens: stream_result.output_tokens,
                        };
                        let outcome = match self
                            .handle_text_response(text, meta, &mut state, &on_event)
                            .await
                        {
                            Ok(o) => o,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };

                        match outcome {
                            TextResponseOutcome::Continue => continue,
                            TextResponseOutcome::Halt(msg) => {
                                return run_return_failure(
                                    AgentError::Internal(format!(
                                        "Turn halted by extension: {}",
                                        msg
                                    )),
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                            TextResponseOutcome::Finalize(text) => {
                                state.messages.push(Message::assistant(&text));
                                let res = state.into_result(text, iteration);
                                let snapshot = res.messages.clone();
                                return run_return_success(res, snapshot);
                            }
                        }
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        if let Some((streamed_items, streamed_results)) =
                            stream_result.streaming_results
                        {
                            let (final_items, final_results) = merge_streamed_tool_results(
                                &items,
                                streamed_items,
                                streamed_results,
                                tools.tool_map(),
                                self.tool_timeout,
                                &self.tool_context,
                                &on_event,
                            )
                            .await;
                            if let Err(e) = self
                                .finalize_tool_call_batch(
                                    &final_items,
                                    final_results,
                                    iteration,
                                    tools.tool_defs(),
                                    &mut state,
                                    &on_event,
                                )
                                .await
                            {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                        } else {
                            // No streaming tool execution — fall back to batch execution.
                            emit_tool_started(&items, &on_event);
                            let slots = self.dispatch_intercept_tool_calls(&items, &on_event).await;
                            // Resolve any Deferred slots via the real ops
                            // channel when one is available. Previously this
                            // path always passed `&mut None`, which meant
                            // ask_user / exit_plan_mode slots could not
                            // complete on the streaming callback path.
                            let slots = match self
                                .resolve_deferred_slots(
                                    slots,
                                    &mut ops_state,
                                    &mut ops_rx,
                                    &on_event,
                                )
                                .await
                            {
                                Ok(s) => s,
                                Err(e) => {
                                    return run_return_failure(
                                        e,
                                        state.messages,
                                        iteration,
                                        state.total_usage,
                                    )
                                }
                            };
                            let mut remaining_items: Vec<crate::llm::ToolCallItem> = Vec::new();
                            let mut remaining_indices: Vec<usize> = Vec::new();
                            let mut final_results: Vec<Option<motosan_agent_tool::ToolResult>> =
                                slots
                                    .into_iter()
                                    .map(|s| match s {
                                        InterceptedSlot::Resolved(r) => Some(r),
                                        InterceptedSlot::Pending => None,
                                        InterceptedSlot::Deferred(_) => {
                                            unreachable!(
                                                "InterceptedSlot::Deferred should have been \
                                                 resolved by resolve_deferred_slots"
                                            );
                                        }
                                    })
                                    .collect();
                            for (i, item) in items.iter().enumerate() {
                                if final_results[i].is_none() {
                                    remaining_items.push(item.clone());
                                    remaining_indices.push(i);
                                }
                            }
                            let execution_results = Self::execute_tools_parallel(
                                tools.tool_map(),
                                &remaining_items,
                                self.tool_timeout,
                                &self.tool_context,
                            )
                            .await;
                            for (exec_idx, original_idx) in remaining_indices.iter().enumerate() {
                                final_results[*original_idx] =
                                    Some(execution_results[exec_idx].clone());
                            }
                            let results: Vec<motosan_agent_tool::ToolResult> = final_results
                                .into_iter()
                                .map(|opt| opt.expect("all slots filled after intercept + execute"))
                                .collect();
                            if let Err(e) = self
                                .finalize_tool_call_batch(
                                    &items,
                                    results,
                                    iteration,
                                    tools.tool_defs(),
                                    &mut state,
                                    &on_event,
                                )
                                .await
                            {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                        }
                        break; // Back to outer iteration loop.
                    }
                }
            }
        }

        run_return_failure(
            AgentError::MaxIterations(self.max_iterations),
            state.messages,
            self.max_iterations,
            state.total_usage,
        )
    }

    /// Inner function covering (chunked LLM + ops_rx + cancel). Mirrors
    /// `run_streaming_inner`, but races the streaming LLM step against the
    /// cancellation token so terminal hooks observe the real in-progress
    /// messages / usage when cancellation wins.
    ///
    /// # Known latency gap
    ///
    /// Same as `run_inner_with_ops_and_cancel`: cancellation is not
    /// observed during `resolve_deferred_slots` inside the tool-call
    /// path. A cancel that arrives while waiting for a deferred slot
    /// (e.g. `ask_user` answer) will not be observed until the slot
    /// resolves or the 60s defer timeout fires.
    #[cfg(feature = "cancellation")]
    #[allow(dead_code)] // Called by RunBuilder dispatch in a later task
    async fn run_streaming_inner_with_cancel_and_ops(
        &self,
        llm: &dyn LlmClient,
        messages: Vec<Message>,
        extra_tools: &[Arc<dyn Tool>],
        ops_rx: mpsc::Receiver<AgentOp>,
        cancel: &tokio_util::sync::CancellationToken,
        on_event: impl Fn(AgentEvent) + Send + Sync,
    ) -> (Result<AgentResult>, Vec<Message>, RunTerminalMeta) {
        use futures::StreamExt;

        let tools = MergedTools::new(&self.tool_map, &self.tool_defs, extra_tools);
        let prepared = match self.prepare_messages(messages).await {
            Ok(m) => m,
            Err(e) => return run_return_failure(e, Vec::new(), 0, TokenUsage::default()),
        };
        let mut state = TurnState::new(prepared);
        let mut ops_state = OpsState::default();
        let mut ops_rx_opt = Some(ops_rx);

        for iteration in 1..=self.max_iterations {
            if cancel.is_cancelled() {
                return run_return_failure(
                    AgentError::Cancelled,
                    state.messages,
                    iteration,
                    state.total_usage,
                );
            }

            if let Err(e) = self
                .drain_ops(
                    &mut state.messages,
                    &mut ops_rx_opt,
                    &mut ops_state,
                    &on_event,
                )
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }
            if ops_state.interrupted {
                on_event(AgentEvent::Core(CoreEvent::Interrupted));
                let res =
                    state.into_result("(interrupted)".to_string(), iteration.saturating_sub(1));
                let snapshot = res.messages.clone();
                return run_return_interrupted(res, snapshot);
            }

            on_event(AgentEvent::Core(CoreEvent::IterationStarted { iteration }));

            if let Err(e) = self
                .dispatch_before_iteration(iteration, &mut state.messages, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            if let Err(e) = self
                .maybe_compact_via_extensions(&mut state.messages, &on_event)
                .await
            {
                return run_return_failure(e, state.messages, iteration, state.total_usage);
            }

            loop {
                let (
                    accumulated_text,
                    response,
                    streaming_results,
                    stream_stop_reason,
                    output_tokens,
                ) = {
                    let mut stream = llm.chat_stream(&state.messages, tools.tool_defs());
                    let mut accumulated = String::new();
                    let mut final_response: Option<crate::LlmResponse> = None;
                    let mut executor = crate::streaming_executor::StreamingToolExecutor::new();
                    let mut submitted_ids: std::collections::HashSet<String> =
                        std::collections::HashSet::new();
                    let mut stop_reason: Option<crate::llm::StopReason> = None;
                    let output_tokens_before = state.total_usage.output_tokens;

                    loop {
                        tokio::select! {
                            chunk_opt = stream.next() => {
                                match chunk_opt {
                                    Some(chunk_result) => {
                                        let chunk = match chunk_result {
                                            Ok(c) => c,
                                            Err(e) => return run_return_failure(e, state.messages.clone(), iteration, state.total_usage),
                                        };
                                        match chunk {
                                            crate::llm::StreamChunk::TextDelta(delta) => {
                                                accumulated.push_str(&delta);
                                                on_event(AgentEvent::Core(CoreEvent::TextChunk(delta)));
                                            }
                                            crate::llm::StreamChunk::ToolUse { id, name, args } => {
                                                if !submitted_ids.contains(&id) {
                                                    let item = crate::llm::ToolCallItem {
                                                        id: id.clone(),
                                                        name: name.clone(),
                                                        args,
                                                    };
                                                    on_event(AgentEvent::Core(CoreEvent::ToolStarted {
                                                        name: name.clone(),
                                                    }));
                                                    executor.submit(
                                                        item,
                                                        tools.tool_map(),
                                                        self.tool_timeout,
                                                        &self.tool_context,
                                                    );
                                                    submitted_ids.insert(id);
                                                }
                                            }
                                            crate::llm::StreamChunk::Done(resp) => {
                                                final_response = Some(resp);
                                            }
                                            crate::llm::StreamChunk::Usage(usage) => {
                                                state.total_usage.accumulate(usage);
                                            }
                                            crate::llm::StreamChunk::StopReason(reason) => {
                                                stop_reason = Some(reason);
                                            }
                                        }
                                    }
                                    None => break,
                                }
                            }
                            _ = cancel.cancelled() => {
                                return run_return_failure(AgentError::Cancelled, state.messages.clone(), iteration, state.total_usage);
                            }
                        }
                    }

                    let resp = final_response
                        .unwrap_or_else(|| crate::LlmResponse::Message(accumulated.clone()));
                    let sr = if executor.has_pending() {
                        Some(executor.collect().await)
                    } else {
                        None
                    };
                    let delta = state.total_usage.output_tokens - output_tokens_before;
                    let ot = if delta > 0 { Some(delta) } else { None };
                    (accumulated, resp, sr, stop_reason, ot)
                };

                match response {
                    crate::LlmResponse::Message(text) => {
                        if accumulated_text.is_empty() && !text.is_empty() {
                            on_event(AgentEvent::Core(CoreEvent::TextChunk(text.clone())));
                        }
                        on_event(AgentEvent::Core(CoreEvent::TextDone(text.clone())));

                        let meta = crate::llm::LlmResponseMeta {
                            stop_reason: stream_stop_reason,
                            output_tokens,
                        };
                        let outcome = match self
                            .handle_text_response(text, meta, &mut state, &on_event)
                            .await
                        {
                            Ok(o) => o,
                            Err(e) => {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                )
                            }
                        };

                        match outcome {
                            TextResponseOutcome::Continue => continue,
                            TextResponseOutcome::Halt(msg) => {
                                return run_return_failure(
                                    AgentError::Internal(format!(
                                        "Turn halted by extension: {}",
                                        msg
                                    )),
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                            TextResponseOutcome::Finalize(text) => {
                                state.messages.push(Message::assistant(&text));
                                let res = state.into_result(text, iteration);
                                let snapshot = res.messages.clone();
                                return run_return_success(res, snapshot);
                            }
                        }
                    }
                    crate::LlmResponse::ToolCalls(items) => {
                        if let Some((streamed_items, streamed_results)) = streaming_results {
                            let (final_items, final_results) = merge_streamed_tool_results(
                                &items,
                                streamed_items,
                                streamed_results,
                                tools.tool_map(),
                                self.tool_timeout,
                                &self.tool_context,
                                &on_event,
                            )
                            .await;
                            if let Err(e) = self
                                .finalize_tool_call_batch(
                                    &final_items,
                                    final_results,
                                    iteration,
                                    tools.tool_defs(),
                                    &mut state,
                                    &on_event,
                                )
                                .await
                            {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                        } else {
                            emit_tool_started(&items, &on_event);
                            let slots = self.dispatch_intercept_tool_calls(&items, &on_event).await;
                            let slots = match self
                                .resolve_deferred_slots(
                                    slots,
                                    &mut ops_state,
                                    &mut ops_rx_opt,
                                    &on_event,
                                )
                                .await
                            {
                                Ok(s) => s,
                                Err(e) => {
                                    return run_return_failure(
                                        e,
                                        state.messages,
                                        iteration,
                                        state.total_usage,
                                    )
                                }
                            };
                            let mut remaining_items: Vec<crate::llm::ToolCallItem> = Vec::new();
                            let mut remaining_indices: Vec<usize> = Vec::new();
                            let mut final_results: Vec<Option<motosan_agent_tool::ToolResult>> =
                                slots
                                    .into_iter()
                                    .map(|s| match s {
                                        InterceptedSlot::Resolved(r) => Some(r),
                                        InterceptedSlot::Pending => None,
                                        InterceptedSlot::Deferred(_) => {
                                            unreachable!(
                                                "InterceptedSlot::Deferred should have been \
                                                 resolved by resolve_deferred_slots"
                                            );
                                        }
                                    })
                                    .collect();
                            for (i, item) in items.iter().enumerate() {
                                if final_results[i].is_none() {
                                    remaining_items.push(item.clone());
                                    remaining_indices.push(i);
                                }
                            }
                            let execution_results = Self::execute_tools_parallel(
                                tools.tool_map(),
                                &remaining_items,
                                self.tool_timeout,
                                &self.tool_context,
                            )
                            .await;
                            for (exec_idx, original_idx) in remaining_indices.iter().enumerate() {
                                final_results[*original_idx] =
                                    Some(execution_results[exec_idx].clone());
                            }
                            let results: Vec<motosan_agent_tool::ToolResult> = final_results
                                .into_iter()
                                .map(|opt| opt.expect("all slots filled after intercept + execute"))
                                .collect();
                            if let Err(e) = self
                                .finalize_tool_call_batch(
                                    &items,
                                    results,
                                    iteration,
                                    tools.tool_defs(),
                                    &mut state,
                                    &on_event,
                                )
                                .await
                            {
                                return run_return_failure(
                                    e,
                                    state.messages,
                                    iteration,
                                    state.total_usage,
                                );
                            }
                        }
                        break;
                    }
                }
            }
        }

        run_return_failure(
            AgentError::MaxIterations(self.max_iterations),
            state.messages,
            self.max_iterations,
            state.total_usage,
        )
    }

    /// Drain all immediately available ops from `ops_rx` and route each one
    /// through `dispatch_on_op` first so that extensions (e.g.
    /// `AskUserExtension`, `PlanningExtension`) can buffer ops that arrive
    /// before the corresponding deferred tool call exists.  Only ops for
    /// which every extension returns `OpDecision::Pass` fall through to the
    /// legacy `apply_op` handler (Interrupt, InjectUserMessage, etc.).
    ///
    /// This fixes the pre-queue bug (#127): previously `drain_ops` called
    /// `apply_op` directly, which was a no-op for `AskUserAnswer` and
    /// `ApprovePlan`, silently dropping them.
    async fn drain_ops(
        &self,
        messages: &mut Vec<Message>,
        ops_rx: &mut Option<mpsc::Receiver<AgentOp>>,
        ops_state: &mut OpsState,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<()> {
        let Some(rx) = ops_rx.as_mut() else {
            return Ok(());
        };
        while let Ok(op) = rx.try_recv() {
            // Route to extensions first. If no extension claims the op, fall
            // through to the legacy inline handler.
            let decision = self.dispatch_on_op(&op, on_event).await?;
            match decision {
                crate::core::OpDecision::Pass => {
                    Self::apply_op(op, messages, ops_state);
                }
                crate::core::OpDecision::Handled | crate::core::OpDecision::Reject(_) => {
                    // Extension consumed it. Nothing more to do.
                }
                crate::core::OpDecision::ResumeDeferred { call_id, .. } => {
                    debug_assert!(
                        false,
                        "extension returned ResumeDeferred from drain_ops for call_id '{}', \
                         but no deferred calls can exist at drain time (on_terminal clears \
                         pending state, tool calls only happen inside iterations). \
                         This indicates an extension bug.",
                        call_id
                    );
                    // In release builds, ignore the return — best-effort resilience.
                }
            }
        }
        Ok(())
    }

    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 { .. } => {
                // Handled by AskUserExtension via the Defer/ResumeDeferred protocol.
                // apply_op is called for ops outside the resolve_deferred_slots loop;
                // AskUserAnswer arriving here is a no-op (no pending answers queue).
            }
            AgentOp::ApprovePlan { .. } => {
                // Handled by PlanningExtension via the Defer/ResumeDeferred protocol.
                // apply_op is called for ops outside the resolve_deferred_slots loop;
                // ApprovePlan arriving here is a no-op (no pending approval queue).
            }
        }
    }

    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 => {
                let slots = self.dispatch_intercept_tool_calls(items, on_event).await;

                // Separate Pending items (need normal execution) from Deferred/Resolved.
                let mut pending_items: Vec<crate::llm::ToolCallItem> = Vec::new();
                let mut pending_indices: Vec<usize> = Vec::new();
                for (i, (slot, item)) in slots.iter().zip(items.iter()).enumerate() {
                    if matches!(slot, InterceptedSlot::Pending) {
                        pending_items.push(item.clone());
                        pending_indices.push(i);
                    }
                }

                // Run Pending tools in parallel with Deferred slot resolution.
                // This preserves the old behavior where non-ask_user tools run
                // concurrently while waiting for an ask_user answer.
                let (resolved_slots, pending_results) = futures::join!(
                    self.resolve_deferred_slots(slots, ops_state, ops_rx, on_event),
                    Self::execute_tools_parallel(
                        tool_map,
                        &pending_items,
                        self.tool_timeout,
                        &self.tool_context,
                    )
                );
                let resolved_slots = resolved_slots.unwrap_or_else(|e| {
                    eprintln!("[motosan-agent-loop] resolve_deferred_slots error: {}", e);
                    Vec::new()
                });

                let mut final_results: Vec<Option<motosan_agent_tool::ToolResult>> = resolved_slots
                    .into_iter()
                    .map(|s| match s {
                        InterceptedSlot::Resolved(r) => Some(r),
                        InterceptedSlot::Pending => None,
                        InterceptedSlot::Deferred(_) => {
                            unreachable!(
                                "InterceptedSlot::Deferred should have been \
                                 resolved by resolve_deferred_slots"
                            );
                        }
                    })
                    .collect();
                for (exec_idx, original_idx) in pending_indices.iter().enumerate() {
                    final_results[*original_idx] = Some(pending_results[exec_idx].clone());
                }
                final_results
                    .into_iter()
                    .map(|opt| opt.expect("all slots filled after intercept + execute"))
                    .collect()
            }
        }
    }

    /// Run each tool call through the extension `intercept_tool_call` pipeline.
    ///
    /// Returns a slot-aligned `Vec<InterceptedSlot>` matching `items`.
    /// See [`InterceptedSlot`] for the three possible outcomes.
    ///
    /// Forwards any `DelegationEvent::Started/Completed` as the legacy
    /// `AgentEvent::DelegateStarted/DelegateCompleted` for backward compatibility.
    /// TODO(phase-4): remove the downcast block once the public API cutover removes
    /// `AgentEvent::DelegateStarted` and `AgentEvent::DelegateCompleted`.
    async fn dispatch_intercept_tool_calls(
        &self,
        items: &[crate::llm::ToolCallItem],
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Vec<InterceptedSlot> {
        // Fast path: no extensions registered, nothing to intercept.
        if self.extensions.lock().await.is_empty() {
            return items.iter().map(|_| InterceptedSlot::Pending).collect();
        }

        let tools_snapshot = self.tool_defs.to_vec();
        let empty_messages: Vec<crate::message::Message> = vec![];
        let agent_state = crate::core::AgentState {
            iteration: 0,
            messages: &empty_messages,
            tools: &tools_snapshot,
            turn_started_at: std::time::Instant::now(),
            token_usage: crate::llm::TokenUsage::default(),
        };

        let mut slots: Vec<InterceptedSlot> = Vec::with_capacity(items.len());

        for item in items.iter() {
            let mut sink = crate::core::CaptureSink::default();
            let decision = {
                let mut set = self.extensions.lock().await;
                set.intercept_tool_call(item.clone(), &agent_state, &mut sink)
                    .await
                    .unwrap_or_else(|e| {
                        eprintln!(
                            "[motosan-agent-loop] dispatch_intercept_tool_calls failed: {}. \
                             Treating as Proceed.",
                            e
                        );
                        crate::core::ToolDecision::Proceed(item.clone())
                    })
            };

            // Forward extension events → AgentEvent::Extension stream.
            for (_name, evt) in sink.captured() {
                forward_extension_event(evt.as_ref(), on_event);
            }

            let slot = match decision {
                crate::core::ToolDecision::ShortCircuit(result) => {
                    InterceptedSlot::Resolved(result)
                }
                crate::core::ToolDecision::Defer { call_id, reason } => {
                    // Record the deferred call. Task 5's resolve_deferred_slots
                    // will read this map when it processes ResumeDeferred ops.
                    let mut deferred = self.deferred_calls.lock().await;
                    if deferred.contains_key(&call_id) {
                        eprintln!(
                            "[motosan-agent-loop] WARNING: duplicate Defer call_id '{}' \
                             (reason: {}). Treating as Pending — the previously deferred \
                             call may be lost.",
                            call_id, reason
                        );
                        InterceptedSlot::Pending
                    } else {
                        deferred.insert(
                            call_id.clone(),
                            DeferredCall {
                                call: item.clone(),
                                by_extension: "unknown",
                                // TODO(phase-3b): plumb the real extension name
                                // through ExtensionSet::intercept_tool_call so the
                                // diagnostic can name which extension deferred the
                                // call. For now, "unknown" is acceptable.
                                at: std::time::Instant::now(),
                            },
                        );
                        InterceptedSlot::Deferred(call_id)
                    }
                }
                // Proceed (with possibly modified call) or Replace —
                // both fall through to normal dispatch.
                _ => InterceptedSlot::Pending,
            };

            slots.push(slot);
        }

        slots
    }

    /// Route an incoming `AgentOp` to the extension system via
    /// `ExtensionSet::on_op` and return the resulting `OpDecision`.
    ///
    /// Returns the `OpDecision` from the first extension that returned
    /// something other than `Pass`. `Pass` means "no extension handled
    /// this op" — the caller should fall back to legacy inline handling
    /// (until Tasks 11 and 15 delete those handlers).
    ///
    /// Forwards extension events to the legacy `AgentEvent` stream.
    /// Phase 3B Tasks 11 and 15 will populate the forwarding loop with
    /// AskUserEvent and PlanningEvent downcasts. For Task 3, the loop
    /// is intentionally a no-op skeleton.
    ///
    /// TODO(phase-4): delete this method once the public API cutover
    /// routes consumers to CoreEvent + ExtensionEvent directly.
    async fn dispatch_on_op(
        &self,
        op: &AgentOp,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<crate::core::OpDecision> {
        if self.extensions.lock().await.is_empty() {
            return Ok(crate::core::OpDecision::Pass);
        }

        let tools_snapshot = self.tool_defs.to_vec();
        let empty_messages: Vec<crate::message::Message> = vec![];
        let agent_state = crate::core::AgentState {
            iteration: 0,
            messages: &empty_messages,
            tools: &tools_snapshot,
            turn_started_at: std::time::Instant::now(),
            token_usage: crate::llm::TokenUsage::default(),
        };

        let mut sink = crate::core::CaptureSink::default();
        let decision = self
            .extensions
            .lock()
            .await
            .on_op(op, &agent_state, &mut sink)
            .await
            .map_err(|e| crate::error::AgentError::Internal(e.message().to_string()))?;

        // Forward captured events to the AgentEvent stream.
        for (_name, evt) in sink.captured() {
            forward_extension_event(evt.as_ref(), on_event);
        }

        Ok(decision)
    }

    /// Wait for all `Deferred` slots in the given vec to be resolved
    /// via incoming `AgentOp`s and `dispatch_on_op` returning
    /// `ResumeDeferred`. Returns the same vec with all `Deferred`
    /// entries replaced by `Resolved`.
    ///
    /// On timeout (default 60s; eventually extension config should drive
    /// this), unresolved deferred slots are filled with an error
    /// `ToolResult` and a diagnostic is logged via eprintln.
    ///
    /// During the wait, ops that no extension claims (returning
    /// `OpDecision::Pass`) fall through to legacy inline handling
    /// (mirroring apply_op). Tasks 11 and 15 will replace the legacy
    /// handlers with extension-based ones, but for now this method
    /// preserves backward compatibility with the existing AskUserAnswer /
    /// ApprovePlan inline op handlers.
    async fn resolve_deferred_slots(
        &self,
        mut slots: Vec<InterceptedSlot>,
        ops_state: &mut OpsState,
        ops_rx: &mut Option<tokio::sync::mpsc::Receiver<AgentOp>>,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<Vec<InterceptedSlot>> {
        // Fast path: nothing deferred.
        if !slots
            .iter()
            .any(|s| matches!(s, InterceptedSlot::Deferred(_)))
        {
            return Ok(slots);
        }

        // Use the minimum timeout from registered extensions (e.g. AskUserConfig::timeout),
        // falling back to a 60s ceiling if no extension constrains it.
        const DEFAULT_DEFER_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);
        let overall_timeout = self
            .extensions
            .lock()
            .await
            .min_deferred_call_timeout()
            .unwrap_or(DEFAULT_DEFER_TIMEOUT);
        let deadline = std::time::Instant::now() + overall_timeout;

        loop {
            // Check if all deferred slots are now resolved.
            if !slots
                .iter()
                .any(|s| matches!(s, InterceptedSlot::Deferred(_)))
            {
                break;
            }

            // Compute remaining time for this iteration's wait.
            let now = std::time::Instant::now();
            if now >= deadline {
                // Timeout: resolve every still-Deferred slot with an error.
                // Emit AskUserTimeout for any ask_user deferred calls so that
                // the legacy AgentEvent::AskUserTimeout contract is preserved.
                for slot in slots.iter_mut() {
                    if let InterceptedSlot::Deferred(call_id) = slot {
                        let cid = call_id.clone();
                        self.deferred_calls.lock().await.remove(&cid);
                        // Notify all extensions of the timeout so they can emit
                        // their own timeout events (e.g. AskUserEvent::Timeout).
                        // This replaces the old hardcoded ask_user name check.
                        {
                            let tools_snapshot = self.tool_defs.to_vec();
                            let empty_messages: Vec<crate::message::Message> = vec![];
                            let agent_state = crate::core::AgentState {
                                iteration: 0,
                                messages: &empty_messages,
                                tools: &tools_snapshot,
                                turn_started_at: std::time::Instant::now(),
                                token_usage: crate::llm::TokenUsage::default(),
                            };
                            let mut sink = crate::core::CaptureSink::default();
                            let _ = self
                                .extensions
                                .lock()
                                .await
                                .on_defer_timeout(&cid, &agent_state, &mut sink)
                                .await;
                            for (_name, evt) in sink.captured() {
                                forward_extension_event(evt.as_ref(), on_event);
                            }
                        }
                        eprintln!(
                            "[motosan-agent-loop] Defer call '{}' timed out after {:?}",
                            cid, overall_timeout
                        );
                        *slot = InterceptedSlot::Resolved(motosan_agent_tool::ToolResult::error(
                            format!(
                                "Deferred call '{}' timed out after {:?}",
                                cid, overall_timeout
                            ),
                        ));
                    }
                }
                break;
            }
            let remaining = deadline - now;

            // Wait for the next op or for the timeout.
            let op = match ops_rx {
                Some(rx) => {
                    // TEST-ONLY: fire parked notifier before awaiting ops_rx.
                    #[cfg(test)]
                    if let Some(notify) = self.test_parked_notify.lock().unwrap().as_ref() {
                        notify.notify_one();
                    }
                    tokio::select! {
                        received = rx.recv() => {
                            match received {
                                Some(op) => op,
                                None => {
                                    // Ops channel closed. Resolve remaining as errors.
                                    for slot in slots.iter_mut() {
                                        if let InterceptedSlot::Deferred(call_id) = slot {
                                            let cid = call_id.clone();
                                            self.deferred_calls.lock().await.remove(&cid);
                                            eprintln!(
                                                "[motosan-agent-loop] Defer call '{}' aborted: ops channel closed",
                                                cid
                                            );
                                            *slot = InterceptedSlot::Resolved(
                                                motosan_agent_tool::ToolResult::error(format!(
                                                    "Deferred call '{}' aborted: ops channel closed",
                                                    cid
                                                )),
                                            );
                                        }
                                    }
                                    break;
                                }
                            }
                        }
                        _ = tokio::time::sleep(remaining) => {
                            continue; // loop will detect timeout on next iteration
                        }
                    }
                }
                None => {
                    // No ops channel — cannot wait for resume. Resolve all
                    // deferred slots as errors immediately and notify extensions
                    // via on_defer_timeout so they can emit their own events.
                    for slot in slots.iter_mut() {
                        if let InterceptedSlot::Deferred(call_id) = slot {
                            let cid = call_id.clone();
                            self.deferred_calls.lock().await.remove(&cid);
                            // Notify all extensions of the timeout/abort.
                            // This replaces the old hardcoded ask_user name check.
                            {
                                let tools_snapshot = self.tool_defs.to_vec();
                                let empty_messages: Vec<crate::message::Message> = vec![];
                                let agent_state = crate::core::AgentState {
                                    iteration: 0,
                                    messages: &empty_messages,
                                    tools: &tools_snapshot,
                                    turn_started_at: std::time::Instant::now(),
                                    token_usage: crate::llm::TokenUsage::default(),
                                };
                                let mut sink = crate::core::CaptureSink::default();
                                let _ = self
                                    .extensions
                                    .lock()
                                    .await
                                    .on_defer_timeout(&cid, &agent_state, &mut sink)
                                    .await;
                                for (_name, evt) in sink.captured() {
                                    forward_extension_event(evt.as_ref(), on_event);
                                }
                            }
                            eprintln!(
                                "[motosan-agent-loop] Defer call '{}' aborted: no ops channel",
                                cid
                            );
                            *slot = InterceptedSlot::Resolved(
                                motosan_agent_tool::ToolResult::error(format!(
                                    "Deferred call '{}' aborted: no ops channel available",
                                    cid
                                )),
                            );
                        }
                    }
                    break;
                }
            };

            // Dispatch the op via the extension system first.
            let decision = self.dispatch_on_op(&op, on_event).await?;

            match decision {
                crate::core::OpDecision::ResumeDeferred { call_id, result } => {
                    // Find the slot with matching call_id and fill it.
                    let mut found = false;
                    for slot in slots.iter_mut() {
                        if let InterceptedSlot::Deferred(slot_call_id) = slot {
                            if slot_call_id == &call_id {
                                *slot = InterceptedSlot::Resolved(result.clone());
                                self.deferred_calls.lock().await.remove(&call_id);
                                found = true;
                                break;
                            }
                        }
                    }
                    if !found {
                        eprintln!(
                            "[motosan-agent-loop] WARNING: extension returned ResumeDeferred \
                             for unknown call_id '{}' — no matching deferred slot.",
                            call_id
                        );
                    }
                }
                crate::core::OpDecision::Handled | crate::core::OpDecision::Reject(_) => {
                    // Op was handled by an extension (without resuming a
                    // deferred call) — continue the wait.
                }
                crate::core::OpDecision::Pass => {
                    // No extension claimed the op. Inline handling for ops
                    // not managed by any extension.
                    // AskUserAnswer is fully handled by AskUserExtension.
                    // ApprovePlan is fully handled by PlanningExtension.
                    match &op {
                        AgentOp::Interrupt => {
                            ops_state.interrupted = true;
                            on_event(AgentEvent::Core(CoreEvent::Interrupted));
                            // Resolve all remaining deferred slots with an error
                            // so the caller gets a result, then exit.
                            for slot in slots.iter_mut() {
                                if let InterceptedSlot::Deferred(call_id) = slot {
                                    let cid = call_id.clone();
                                    self.deferred_calls.lock().await.remove(&cid);
                                    *slot = InterceptedSlot::Resolved(
                                        motosan_agent_tool::ToolResult::error(format!(
                                            "Deferred call '{}' aborted: loop interrupted",
                                            cid
                                        )),
                                    );
                                }
                            }
                            break;
                        }
                        AgentOp::AskUserAnswer { .. } | AgentOp::ApprovePlan { .. } => {
                            // Handled by AskUserExtension / PlanningExtension above via
                            // ResumeDeferred. If they fall through to Pass, there was no
                            // matching deferred call_id — ignore silently.
                        }
                        _ => {
                            // Other ops not handled during Defer waits.
                        }
                    }
                }
            }
        }

        Ok(slots)
    }

    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
        }
    }

    /// Finalize a batch of tool call results.
    async fn finalize_tool_call_batch(
        &self,
        items: &[ToolCallItem],
        raw_results: Vec<ToolResult>,
        iteration: usize,
        tools: &[ToolDef],
        state: &mut TurnState,
        on_event: &(impl Fn(AgentEvent) + Send + Sync),
    ) -> Result<()> {
        let mut results: Vec<ToolResult> = Vec::with_capacity(raw_results.len());
        for (tc, raw) in items.iter().zip(raw_results) {
            let agent_state =
                self.agent_state_snapshot(iteration, &state.messages, tools, state.total_usage);
            let rewritten = self
                .dispatch_rewrite_tool_result(tc, raw, &agent_state, on_event)
                .await?;
            results.push(rewritten);
        }

        for (tc, result) in items.iter().zip(results.iter()) {
            on_event(AgentEvent::Core(CoreEvent::ToolCompleted {
                name: tc.name.clone(),
                result: result.clone(),
            }));
            state
                .all_tool_calls
                .push((tc.name.clone(), tc.args.clone()));
        }

        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));

        for (tc, result) in items.iter().zip(results.iter()) {
            state
                .messages
                .push(Message::tool_result(&tc.id, &tool_result_to_string(result)));
        }

        for result in &results {
            let agent_state =
                self.agent_state_snapshot(iteration, &state.messages, tools, state.total_usage);
            match self
                .dispatch_after_tool_result(result, &agent_state, on_event)
                .await?
            {
                FlowDecision::Continue => {}
                FlowDecision::Inject(msg) => state.messages.push(msg),
                FlowDecision::Halt(reason) => {
                    return Err(crate::error::AgentError::Internal(reason.message));
                }
            }
        }

        Ok(())
    }
}

/// `ToolExecutionPolicy` controls how a batch of tool calls is executed.
///
/// `enter_plan_mode` / `exit_plan_mode` are now handled by `PlanningExtension`
/// via the Defer/ResumeDeferred protocol (just like `ask_user`), so all batches
/// now use the `ParallelOnly` path.
#[derive(Debug, Clone, Copy)]
enum ToolExecutionPolicy {
    ParallelOnly,
}

impl ToolExecutionPolicy {
    /// Determine the execution policy for a batch of tool calls.
    ///
    /// All interactive tools (`ask_user`, `enter_plan_mode`, `exit_plan_mode`)
    /// are now handled by extensions via the Defer protocol — every batch uses
    /// `ParallelOnly`.
    fn from_items(_items: &[crate::llm::ToolCallItem]) -> Self {
        Self::ParallelOnly
    }
}

#[derive(Default)]
struct OpsState {
    interrupted: bool,
}

/// Result of dispatching a single tool call through the extension
/// `intercept_tool_call` pipeline.
///
/// - `Pending`: no extension claimed the call; the core's normal
///   tool-map dispatch should run it.
/// - `Resolved`: an extension returned `ShortCircuit(result)`; the
///   call has its final result and skips normal dispatch.
/// - `Deferred`: an extension returned `Defer { call_id }`; the
///   call is suspended, awaiting a matching `OpDecision::ResumeDeferred`.
///   The `call_id` is the key into `Engine::deferred_calls`.
#[derive(Debug)]
enum InterceptedSlot {
    Pending,
    Resolved(motosan_agent_tool::ToolResult),
    Deferred(String),
}

/// A tool call that was suspended by an extension via
/// `ToolDecision::Defer`. The core records these in a map and
/// blocks the turn until a matching `OpDecision::ResumeDeferred`
/// arrives via the ops channel.
///
/// See spec §8.3 for the full protocol.
#[derive(Debug)]
#[allow(dead_code)] // Used in Phase 3B Tasks 4-5
struct DeferredCall {
    /// The original tool call that was deferred. Stored so the
    /// resume path knows which slot to fill in the tool dispatch
    /// batch.
    call: crate::llm::ToolCallItem,
    /// Name of the extension that returned `Defer`. Used for
    /// diagnostics — if a turn times out with this call still
    /// deferred, the error message includes which extension was
    /// responsible.
    by_extension: &'static str,
    /// Wall-clock instant when the call was deferred. Used for
    /// timeout enforcement.
    at: std::time::Instant,
}

/// Forward a type-erased extension event captured in a `CaptureSink` to the
/// `AgentEvent::Extension(...)` stream.
///
/// Attempts each known extension event type in turn; if none match, the
/// event is silently discarded (unknown extension). Called from all four
/// dispatch helpers: `maybe_compact_via_extensions`, `dispatch_after_llm_response`,
/// `dispatch_intercept_tool_calls`, `dispatch_on_op` (and `resolve_deferred_slots`).
fn forward_extension_event(
    evt: &dyn crate::core::hook_ctx::ExtEvent,
    on_event: &(impl Fn(AgentEvent) + Send + Sync),
) {
    #[cfg(feature = "redact")]
    use crate::extensions::redact::RedactEvent;
    use crate::extensions::{
        ask_user::AskUserEvent, autocompact::AutocompactEvent, delegation::DelegationEvent,
        follow_up::FollowUpEvent, planning::PlanningEvent, stuck_detection::StuckDetectionEvent,
        token_budget::TokenBudgetEvent,
    };
    let any = evt.as_any();
    if let Some(e) = any.downcast_ref::<AutocompactEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::Autocompact(
            e.clone(),
        )));
    } else if let Some(e) = any.downcast_ref::<TokenBudgetEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::TokenBudget(
            e.clone(),
        )));
    } else if let Some(e) = any.downcast_ref::<StuckDetectionEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::StuckDetection(
            e.clone(),
        )));
    } else if let Some(e) = any.downcast_ref::<DelegationEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::Delegation(e.clone())));
    } else if let Some(e) = any.downcast_ref::<AskUserEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::AskUser(e.clone())));
    } else if let Some(e) = any.downcast_ref::<PlanningEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::Planning(e.clone())));
    } else if let Some(e) = any.downcast_ref::<FollowUpEvent>() {
        on_event(AgentEvent::Extension(ExtensionEvent::FollowUp(e.clone())));
    } else {
        #[cfg(feature = "redact")]
        if let Some(e) = any.downcast_ref::<RedactEvent>() {
            on_event(AgentEvent::Extension(ExtensionEvent::Redact(e.clone())));
        }
    }
}

/// 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())
        }
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// RunBuilder — the 0.13.0 unified agent run API
// ─────────────────────────────────────────────────────────────────────────────

/// Configuration builder for a single agent turn, returned by
/// [`Engine::run`]. Chain axis setters (`.ops`, `.cancel`, `.chunked`)
/// in any order, then call one of the terminators (`.result`,
/// `.callback`, `.stream`) to execute the turn.
///
/// `RunBuilder` is `#[must_use]`. Dropping it without calling a
/// terminator emits a compiler warning. Under
/// `#![deny(unused_must_use)]` this becomes a hard error:
///
/// ```compile_fail
/// #![deny(unused_must_use)]
/// use std::sync::Arc;
/// use motosan_agent_loop::{Engine, LlmClient, Message};
///
/// fn drop_builder(agent: Arc<Engine>, llm: Arc<dyn LlmClient>) {
///     agent.run(llm, vec![Message::user("hi")]); // no terminator — denied
/// }
/// ```
///
/// Positive control — the same body with a terminator bound to `_`
/// compiles cleanly under the same deny, proving the `compile_fail`
/// above is failing on `unused_must_use` and not some incidental
/// rename or trait-resolution error:
///
/// ```no_run
/// #![deny(unused_must_use)]
/// use std::sync::Arc;
/// use motosan_agent_loop::{Engine, LlmClient, Message};
///
/// async fn use_builder(agent: Arc<Engine>, llm: Arc<dyn LlmClient>) {
///     let _ = agent.run(llm, vec![Message::user("hi")]).result().await;
/// }
/// ```
///
/// All fields are private. External code cannot construct a
/// `RunBuilder` directly — use [`Engine::run`].
#[must_use = "call .result(), .callback(cb), or .stream() to execute the run"]
pub struct RunBuilder {
    engine: Arc<Engine>,
    llm: Arc<dyn LlmClient>,
    messages: Vec<Message>,
    ops_rx: Option<mpsc::Receiver<AgentOp>>,
    #[cfg(feature = "cancellation")]
    cancel: Option<tokio_util::sync::CancellationToken>,
    chunked: bool,
}

impl RunBuilder {
    // ─── Terminators ────────────────────────────────────────────

    /// Run the turn to completion and return the final `AgentResult`.
    /// Intermediate events are generated internally but not delivered
    /// anywhere. Use this when you only need the final answer.
    pub async fn result(self) -> Result<AgentResult> {
        self.dispatch_callback_internal(|_| {}).await
    }

    /// Run the turn and push each `AgentEvent` through the callback
    /// as it happens. Returns the final `AgentResult` when the turn
    /// terminates.
    pub async fn callback(
        self,
        on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    ) -> Result<AgentResult> {
        self.dispatch_callback_internal(on_event).await
    }

    // ─── Internal dispatch ──────────────────────────────────────

    /// Assemble the axis state and delegate to the appropriate inner
    /// function. Called by `.result()` (with a no-op callback) and
    /// `.callback(cb)` (with the user's callback, in Task 8).
    async fn dispatch_callback_internal(
        self,
        on_event: impl Fn(AgentEvent) + Send + Sync + 'static,
    ) -> Result<AgentResult> {
        // Wrap in Arc so all match arms can use it without move conflicts.
        let on_event = std::sync::Arc::new(on_event);

        let RunBuilder {
            engine,
            llm,
            messages,
            ops_rx,
            #[cfg(feature = "cancellation")]
            cancel,
            chunked,
        } = self;

        // MCP connect (external to the inner functions after Task 1).
        #[cfg(feature = "mcp-client")]
        let mcp_tools: Vec<Arc<dyn Tool>> = engine.connect_mcp_servers().await?;
        #[cfg(not(feature = "mcp-client"))]
        let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

        // Dispatch based on axis combination.
        let (result, _final_messages, terminal_meta) = {
            #[cfg(feature = "cancellation")]
            {
                match (chunked, cancel, ops_rx) {
                    // batch LLM, no cancel, ops or no ops
                    (false, None, ops) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_inner_with_ops(llm.as_ref(), messages, &mcp_tools, &*cb, ops)
                            .await
                    }
                    // batch LLM, cancel, no ops
                    (false, Some(token), None) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_inner_cancel(llm.as_ref(), messages, &mcp_tools, &*cb, &token)
                            .await
                    }
                    // batch LLM, cancel + ops
                    (false, Some(token), Some(rx)) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_inner_with_ops_and_cancel(
                                llm.as_ref(),
                                messages,
                                &mcp_tools,
                                &*cb,
                                rx,
                                &token,
                            )
                            .await
                    }
                    // chunked LLM, no cancel, ops or no ops
                    (true, None, ops) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_streaming_inner(
                                llm.as_ref(),
                                messages,
                                &mcp_tools,
                                ops,
                                move |e| cb(e),
                            )
                            .await
                    }
                    // chunked LLM, cancel, no ops
                    (true, Some(token), None) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_streaming_with_cancel_inner(
                                llm.as_ref(),
                                messages,
                                &mcp_tools,
                                move |e| cb(e),
                                &token,
                            )
                            .await
                    }
                    // chunked LLM, cancel + ops
                    (true, Some(token), Some(rx)) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_streaming_inner_with_cancel_and_ops(
                                llm.as_ref(),
                                messages,
                                &mcp_tools,
                                rx,
                                &token,
                                move |e| cb(e),
                            )
                            .await
                    }
                }
            }
            #[cfg(not(feature = "cancellation"))]
            {
                match (chunked, ops_rx) {
                    (false, ops) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_inner_with_ops(llm.as_ref(), messages, &mcp_tools, &*cb, ops)
                            .await
                    }
                    (true, ops) => {
                        let cb = std::sync::Arc::clone(&on_event);
                        engine
                            .run_streaming_inner(
                                llm.as_ref(),
                                messages,
                                &mcp_tools,
                                ops,
                                move |e| cb(e),
                            )
                            .await
                    }
                }
            }
        };

        let merged_tools = MergedTools::new(&engine.tool_map, &engine.tool_defs, &mcp_tools);
        let terminal_result = match engine
            .dispatch_on_terminal_from_meta(
                &terminal_meta,
                &_final_messages,
                merged_tools.tool_defs(),
                &*on_event,
            )
            .await
        {
            Ok(()) => result,
            Err(e) => Err(e),
        };

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

        terminal_result
    }

    /// Return a pulled stream of `AgentStreamItem`s. The consumer
    /// drives the turn by calling `.next().await`. The final
    /// `Terminal` item is guaranteed to be the last yielded before
    /// the underlying channel closes.
    ///
    /// Use `.callback(cb)` or `.result()` if you want a push-based
    /// delivery model or a final-result-only future.
    pub fn stream(
        self,
    ) -> impl futures::Stream<Item = crate::stream::AgentStreamItem> + Send + 'static {
        self.dispatch_stream()
    }

    /// Spawns a tokio task running the same 8-way dispatch as
    /// `dispatch_callback_internal`, with events piped into an mpsc
    /// channel. The receiver is returned wrapped in
    /// `ReceiverStream`, giving consumers a pull-based delivery
    /// model while internally reusing the same axis routing.
    fn dispatch_stream(
        self,
    ) -> impl futures::Stream<Item = crate::stream::AgentStreamItem> + Send + 'static {
        use crate::stream::{AgentStreamItem, AgentTerminal};

        let capacity = self.engine.channel_config.stream_capacity.unwrap_or(256);
        let (tx, rx) = mpsc::channel::<AgentStreamItem>(capacity);

        let RunBuilder {
            engine,
            llm,
            messages,
            ops_rx,
            #[cfg(feature = "cancellation")]
            cancel,
            chunked,
        } = self;
        let tx_for_events = tx.clone();

        tokio::spawn(async move {
            let on_event = move |e: AgentEvent| {
                let _ = tx_for_events.try_send(AgentStreamItem::Event(e));
            };
            let on_event = std::sync::Arc::new(on_event);

            // MCP connect.
            #[cfg(feature = "mcp-client")]
            let mcp_tools: Vec<Arc<dyn Tool>> = match engine.connect_mcp_servers().await {
                Ok(t) => t,
                Err(e) => {
                    let _ = tx
                        .send(AgentStreamItem::Terminal(AgentTerminal {
                            result: Err(e),
                            messages: Vec::new(),
                        }))
                        .await;
                    return;
                }
            };
            #[cfg(not(feature = "mcp-client"))]
            let mcp_tools: Vec<Arc<dyn Tool>> = vec![];

            // Same 8-way dispatch match as dispatch_callback_internal.
            let (result, final_messages, terminal_meta) = {
                #[cfg(feature = "cancellation")]
                {
                    match (chunked, cancel, ops_rx) {
                        (false, None, ops) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_inner_with_ops(llm.as_ref(), messages, &mcp_tools, &*cb, ops)
                                .await
                        }
                        (false, Some(token), None) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_inner_cancel(llm.as_ref(), messages, &mcp_tools, &*cb, &token)
                                .await
                        }
                        (false, Some(token), Some(rx_ops)) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_inner_with_ops_and_cancel(
                                    llm.as_ref(),
                                    messages,
                                    &mcp_tools,
                                    &*cb,
                                    rx_ops,
                                    &token,
                                )
                                .await
                        }
                        (true, None, ops) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_streaming_inner(
                                    llm.as_ref(),
                                    messages,
                                    &mcp_tools,
                                    ops,
                                    move |e| cb(e),
                                )
                                .await
                        }
                        (true, Some(token), None) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_streaming_with_cancel_inner(
                                    llm.as_ref(),
                                    messages,
                                    &mcp_tools,
                                    move |e| cb(e),
                                    &token,
                                )
                                .await
                        }
                        (true, Some(token), Some(rx_ops)) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_streaming_inner_with_cancel_and_ops(
                                    llm.as_ref(),
                                    messages,
                                    &mcp_tools,
                                    rx_ops,
                                    &token,
                                    move |e| cb(e),
                                )
                                .await
                        }
                    }
                }
                #[cfg(not(feature = "cancellation"))]
                {
                    match (chunked, ops_rx) {
                        (false, ops) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_inner_with_ops(llm.as_ref(), messages, &mcp_tools, &*cb, ops)
                                .await
                        }
                        (true, ops) => {
                            let cb = std::sync::Arc::clone(&on_event);
                            engine
                                .run_streaming_inner(
                                    llm.as_ref(),
                                    messages,
                                    &mcp_tools,
                                    ops,
                                    move |e| cb(e),
                                )
                                .await
                        }
                    }
                }
            };

            let merged_tools = MergedTools::new(&engine.tool_map, &engine.tool_defs, &mcp_tools);
            let terminal_result = match engine
                .dispatch_on_terminal_from_meta(
                    &terminal_meta,
                    &final_messages,
                    merged_tools.tool_defs(),
                    &*on_event,
                )
                .await
            {
                Ok(()) => result,
                Err(e) => Err(e),
            };

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

            // CRITICAL: send Terminal with awaiting send (not try_send) to
            // guarantee delivery even if the consumer is slow.
            let _ = tx
                .send(AgentStreamItem::Terminal(AgentTerminal {
                    result: terminal_result,
                    messages: final_messages,
                }))
                .await;
        });

        tokio_stream::wrappers::ReceiverStream::new(rx)
    }

    // ─── Axis setters ────────────────────────────────────────────

    /// Route `AgentOp` values from an external channel into the
    /// running turn. Calling twice overwrites the first receiver
    /// (last-write-wins).
    pub fn ops(mut self, rx: mpsc::Receiver<AgentOp>) -> Self {
        self.ops_rx = Some(rx);
        self
    }

    /// Attach a cancellation token. When tripped, the turn exits
    /// with `AgentError::Cancelled`. Calling twice overwrites.
    #[cfg(feature = "cancellation")]
    pub fn cancel(mut self, token: tokio_util::sync::CancellationToken) -> Self {
        self.cancel = Some(token);
        self
    }

    /// Use [`LlmClient::chat_stream`] instead of [`LlmClient::chat`],
    /// so `CoreEvent::TextChunk` events fire per LLM token delta.
    /// Observable via `.callback(cb)` or by draining the stream
    /// returned by `.stream()`. Without an event observer, the chunks
    /// are still generated but not delivered anywhere.
    ///
    /// Equivalent to `stream=True` on the OpenAI / Anthropic SDKs.
    pub fn chunked(mut self) -> Self {
        self.chunked = true;
        self
    }
}

impl Engine {
    /// Build a run configuration for this agent.
    ///
    /// Chain axis setters (`.ops(rx)`, `.cancel(token)`, `.chunked()`)
    /// in any order, then call one of the terminators (`.result()`,
    /// `.callback(cb)`, `.stream()`) to execute the turn.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use std::sync::Arc;
    /// use motosan_agent_loop::{Engine, Message};
    ///
    /// # async fn demo(llm: Arc<dyn motosan_agent_loop::LlmClient>) -> motosan_agent_loop::Result<()> {
    /// let agent = Arc::new(Engine::builder().build());
    ///
    /// // Simplest: just the final result.
    /// let result = Arc::clone(&agent)
    ///     .run(llm.clone(), vec![Message::user("Hi!")])
    ///     .result()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    ///
    /// See the [`RunBuilder`] struct documentation for the full set
    /// of axis setters and terminators.
    pub fn run(self: Arc<Self>, llm: Arc<dyn LlmClient>, messages: Vec<Message>) -> RunBuilder {
        RunBuilder {
            engine: self,
            llm,
            messages,
            ops_rx: None,
            #[cfg(feature = "cancellation")]
            cancel: None,
            chunked: false,
        }
    }
}

#[cfg(test)]
#[path = "engine_tests.rs"]
mod tests;

#[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 = Arc::new(
            Engine::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 = Engine::builder()
            .mcp_server_arc(Arc::clone(&shared))
            .max_iterations(5)
            .build();

        let agent_b = Engine::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,
                stop_reason: 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 = Arc::new(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_cb = events.clone();

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

        let result = Arc::clone(&agent)
            .run(
                llm.clone() as Arc<dyn LlmClient>,
                vec![Message::user("call echo")],
            )
            .chunked()
            .callback(move |event| {
                use crate::extensions::ask_user::AskUserEvent;
                let label = match &event {
                    AgentEvent::Core(CoreEvent::ToolStarted { name }) => {
                        format!("started:{name}")
                    }
                    AgentEvent::Core(CoreEvent::ToolCompleted { name, result }) => {
                        format!("completed:{name}:{}", result.as_text().unwrap_or(""))
                    }
                    AgentEvent::Core(CoreEvent::TextChunk(t)) => format!("chunk:{t}"),
                    AgentEvent::Core(CoreEvent::TextDone(t)) => format!("done:{t}"),
                    AgentEvent::Core(CoreEvent::IterationStarted { iteration: n }) => {
                        format!("iter:{n}")
                    }
                    AgentEvent::Core(CoreEvent::Interrupted) => "interrupted".to_string(),
                    AgentEvent::Extension(ExtensionEvent::AskUser(AskUserEvent::Question {
                        questions,
                        ..
                    })) => format!(
                        "ask_user:{}",
                        questions.first().map(|q| q.question.as_str()).unwrap_or("")
                    ),
                    AgentEvent::Extension(ExtensionEvent::AskUser(AskUserEvent::Timeout {
                        call_id,
                        ..
                    })) => {
                        format!("ask_user_timeout:{call_id}")
                    }
                    _ => format!("{event:?}"),
                };
                events_cb.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 = Arc::new(
            Engine::builder()
                .mcp_server_arc(server1)
                .mcp_server_arc(server2)
                .max_iterations(1)
                .build(),
        );

        let llm = Arc::new(McpTestLlm::new(vec![]));
        let err = Arc::clone(&agent)
            .run(llm.clone() as Arc<dyn LlmClient>, vec![Message::user("hi")])
            .result()
            .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);
    }
}