pi/core/agent_session/

extension.rs

//! Extension binding, reload, and replaced-context seam.
//!
//! Ports `bindExtensions`, `extendResourcesFromExtensions`,
//! `createReplacedSessionContext`, `hasExtensionHandlers`, and `reload` from
//! `coding-agent/src/core/agent-session.ts`.
//!
//! All extension interaction goes through the [`ExtensionRunner`] trait seam
//! (defined in `extension_runner.rs`). This module owns:
//! - binding UI/mode/command/error listeners (recorded locally; pushed to the
//!   runner once the trait gains `set_ui_context` / `bind_command_context`)
//! - emitting the stored `session_start` event exactly once per session
//!   instance on the first `bind_extensions` call (under `bind_lock`)
//! - extension-driven resource discovery (skills/prompts/themes)
//! - reload (emits `session_shutdown{reload}` on the old host, preserves flag
//!   values, restarts the host, re-emits `session_start{reload}` on the new
//!   host, then re-discovers resources)
//! - the replaced-session context handed to `withSession` after runtime swap
//! - extension error isolation (host errors never abort the session)
//!
//! Divergence from upstream: reload emission is not gated on recorded
//! bindings (`hasBindings` in `agent-session.ts`). All Rust modes bind, and
//! `emit` self-gates on handler presence, so the gate would only suppress
//! correct emissions.

use crate::core::resources::{
    ResourceLoader, SlashCommandInfo, SlashCommandSource, SyntheticSourceInfoOptions,
    create_synthetic_source_info,
};
use std::sync::Arc;
#[cfg(test)]
use std::sync::Mutex;

use super::AgentSession;
use super::events::{
    AgentSessionEvent, SessionShutdownReason, SessionStartEvent, SessionStartReason,
};
use super::prompt::{CustomMessageInput, DeliverAs, PromptError};
use pi_ai::ImageContent;

/// Test-only restart factory for [`AgentSession::reload`].
///
/// When set, reload injects the returned host instead of spawning a real
/// process, so session cutover can be proven against in-memory transports.
#[cfg(test)]
pub(super) type ReloadRestartFactory = Arc<
    dyn Fn(
            Vec<String>,
            String,
            bool,
        ) -> futures::future::BoxFuture<
            'static,
            Result<
                Arc<crate::core::extension_host::HostExtensionRunner>,
                crate::core::extension_host::HostStartError,
            >,
        > + Send
        + Sync,
>;

/// Mode the session is bound to (mirrors `AppMode` minus `Interactive`).
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ExtensionMode {
    /// Interactive TUI.
    Tui,
    /// Print mode (text).
    Print,
    /// JSON print mode.
    Json,
    /// RPC server.
    Rpc,
}

impl ExtensionMode {
    /// Wire discriminant matching TS.
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Tui => "tui",
            Self::Print => "print",
            Self::Json => "json",
            Self::Rpc => "rpc",
        }
    }
}

/// UI context marker for extensions (opaque until pi-tui integration).
///
/// The interactive mode supplies a concrete `ExtensionUiContext`; other modes
/// pass `None`. Stored as an opaque tag so the host can detect presence.
#[derive(Clone, Debug, Default)]
pub struct ExtensionUiContext {
    /// Opaque caller-supplied tag (component handle, mode marker, etc.).
    pub tag: Option<String>,
}

/// Shared callback signature for structured extension error notifications.
pub type ExtensionErrorListener = Arc<dyn Fn(&str, &str, &str) + Send + Sync>;

/// Inputs accepted by [`AgentSession::bind_extensions`].
#[derive(Clone, Default)]
pub struct ExtensionBindings {
    /// Optional UI context (interactive only).
    pub ui_context: Option<ExtensionUiContext>,
    /// Mode the session is bound to.
    pub mode: Option<ExtensionMode>,
    /// Opaque command-context actions map (interactive/rpc).
    pub command_context_actions: Option<serde_json::Value>,
    /// Optional shutdown handler.
    pub shutdown_handler: Option<Arc<dyn Fn() + Send + Sync>>,
    /// Optional error listener invoked when the host reports an extension error.
    pub on_error: Option<ExtensionErrorListener>,
}

/// Context handed to `withSession` after a runtime replacement.
///
/// Every send targets the replacement session captured when this context is
/// created, even if the outer runtime later changes its active session again.
#[derive(Clone)]
pub struct ReplacedSessionContext {
    /// Session id of the new session.
    pub session_id: String,
    session: Arc<AgentSession>,
}

impl ReplacedSessionContext {
    /// Send a custom message through the replacement session.
    ///
    /// # Errors
    ///
    /// Returns the replacement session's prompt error unchanged.
    pub async fn send_custom_message(
        &self,
        message: CustomMessageInput,
        trigger_turn: bool,
        deliver_as: Option<DeliverAs>,
    ) -> Result<(), PromptError> {
        self.session
            .send_custom_message(message, trigger_turn, deliver_as)
            .await
    }

    /// Send a user message through the replacement session.
    ///
    /// # Errors
    ///
    /// Returns the replacement session's prompt error unchanged.
    pub async fn send_user_message(
        &self,
        text: &str,
        images: Vec<ImageContent>,
        deliver_as: Option<DeliverAs>,
    ) -> Result<(), PromptError> {
        self.session
            .send_user_message(text, images, deliver_as)
            .await
    }
}

/// Errors raised by extension binding / reload.
#[derive(Debug, thiserror::Error)]
pub enum ExtensionBindError {
    /// `resources_discover` failed.
    #[error("extension resource discovery failed: {0}")]
    ResourceDiscover(super::extension_runner::ExtensionRunnerError),
    /// Resource loader reload failed.
    #[error("resource reload failed: {0}")]
    ResourceReload(String),
    /// Host restart after reload failed.
    #[error("extension host restart failed: {0}")]
    HostRestart(String),
}

impl AgentSession {
    /// Returns true when at least one extension handler is registered for
    /// `event_type`. Cheap delegation to the runner; safe to call from any
    /// thread without locking session state.
    #[must_use]
    pub fn has_extension_handlers(&self, event_type: &str) -> bool {
        self.hooks.runner().has_handlers(event_type)
    }

    /// Bind extension UI/mode/error/shutdown listeners, emit the stored
    /// `session_start` event (first bind only), and drive resource discovery.
    ///
    /// The whole lifecycle runs under the session `bind_lock`, so concurrent
    /// binds are serialized: the losing bind waits for the winner's full
    /// session_start-then-discovery sequence. The stored event is consumed
    /// with `Option::take`, so repeated binds on the same session instance
    /// never re-emit.
    ///
    /// # Errors
    ///
    /// Returns [`ExtensionBindError::ResourceDiscover`] when the runner fails
    /// to discover resources.
    pub async fn bind_extensions(
        &self,
        bindings: ExtensionBindings,
    ) -> Result<(), ExtensionBindError> {
        let _bind_guard = self.bind_lock.lock().await;
        // Persist bindings on the inner state.
        {
            let mut inner = self.lock_inner();
            inner.extension_mode = bindings.mode;
            inner.extension_ui_tag = bindings.ui_context.as_ref().and_then(|c| c.tag.clone());
            inner
                .extension_shutdown_handler
                .clone_from(&bindings.shutdown_handler);
            inner
                .extension_error_listener
                .clone_from(&bindings.on_error);
            inner
                .extension_command_context
                .clone_from(&bindings.command_context_actions);
        }

        let pending = self.lock_inner().pending_session_start.take();
        if let Some(event) = &pending {
            // emit self-gates on has_handlers("session_start"); host errors
            // are isolated (reported via the host error listener).
            let _ = self
                .hooks
                .runner()
                .emit(AgentSessionEvent::SessionStart {
                    reason: event.reason,
                    previous_session_file: event.previous_session_file.clone(),
                })
                .await;
        }
        // Non-reload start reasons map to "startup" for resources_discover
        // (its wire contract only allows startup|reload).
        let discover_reason = match pending {
            Some(SessionStartEvent {
                reason: SessionStartReason::Reload,
                ..
            }) => "reload",
            _ => "startup",
        };
        self.extend_resources_from_extensions(discover_reason)
            .await?;
        Ok(())
    }

    /// Discover skills/prompts/themes from extensions and merge into the
    /// resource loader.
    ///
    /// No-op when no `resources_discover` handlers are registered.
    ///
    /// # Errors
    ///
    /// Returns [`ExtensionBindError::ResourceDiscover`] when the runner fails.
    pub async fn extend_resources_from_extensions(
        &self,
        reason: &str,
    ) -> Result<(), ExtensionBindError> {
        if reason == SessionStartReason::Startup.as_str()
            && self.lock_inner().initial_resources_discovered
        {
            return Ok(());
        }

        if reason == SessionStartReason::Reload.as_str()
            && let Some(loader) = &self.resource_loader
        {
            let mut loader = loader.lock().await;
            loader
                .reload()
                .await
                .map_err(|error| ExtensionBindError::ResourceReload(error.to_string()))?;
            self.apply_resource_snapshot(&loader);
        }

        let runner = self.hooks.runner();
        if !runner.has_handlers("resources_discover") {
            return Ok(());
        }
        let paths = runner
            .emit_resources_discover(&self.cwd, reason)
            .await
            .map_err(ExtensionBindError::ResourceDiscover)?;
        if let Some(loader) = &self.resource_loader {
            let mut loader = loader.lock().await;
            loader.extend_resources(paths);
            self.apply_resource_snapshot(&loader);
        }
        if reason == SessionStartReason::Startup.as_str() {
            self.lock_inner().initial_resources_discovered = true;
        }
        Ok(())
    }

    fn apply_resource_snapshot(&self, loader: &crate::core::resources::DefaultResourceLoader) {
        let skills = loader.get_skills().0.to_vec();
        let prompt_templates = loader.get_prompts().0.to_vec();
        let append = (!loader.get_append_system_prompt().is_empty())
            .then(|| loader.get_append_system_prompt().join("\n\n"));
        let selected_tools = self.lock_inner().active_tool_names.clone();
        let system_prompt = crate::core::system_prompt::build_system_prompt(
            &crate::core::system_prompt::BuildSystemPromptOptions {
                custom_prompt: loader.get_system_prompt().map(str::to_owned),
                selected_tools: Some(selected_tools),
                append,
                cwd: self.cwd.clone(),
                context_files: Some(loader.get_agents_files().to_vec()),
                skills: Some(skills.clone()),
                ..crate::core::system_prompt::BuildSystemPromptOptions::default()
            },
        );
        *self
            .skills
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = skills;
        *self
            .prompt_templates
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner) = prompt_templates;
        self.lock_inner()
            .base_system_prompt
            .clone_from(&system_prompt);
        self.hooks.set_base_system_prompt(system_prompt.clone());
        self.hooks.set_system_prompt_override(None);
        self.agent.set_system_prompt(system_prompt);
    }

    /// Source label for an extension path (`extension:<basename-without-ext>`).
    ///
    /// Angle-bracketed names (in-memory extensions) are emitted verbatim
    /// minus the brackets.
    #[must_use]
    pub fn get_extension_source_label(extension_path: &str) -> String {
        if extension_path.starts_with('<') {
            let trimmed = extension_path.trim_start_matches('<').trim_end_matches('>');
            return format!("extension:{trimmed}");
        }
        let base = std::path::Path::new(extension_path)
            .file_stem()
            .map_or_else(
                || extension_path.to_owned(),
                |s| s.to_string_lossy().into_owned(),
            );
        format!("extension:{base}")
    }

    /// Reload extensions.
    ///
    /// Mirrors TS `reload`:
    /// 1. Capture previous flag values (preserved across the swap).
    /// 2. Emit `session_shutdown{reload}` on the old runner (self-gated on
    ///    handler presence; host errors isolated).
    /// 3. When a concrete host is present: prepare and flag-sync its
    ///    replacement while the old host remains usable, then cut over
    ///    providers, host runner, and tools before reaping the old transport.
    /// 4. Emit `session_start{reload}` on the post-swap runner.
    /// 5. Reload base resources and re-discover extension resources.
    ///
    /// # Errors
    ///
    /// Returns [`ExtensionBindError`] on host restart or resource-discovery
    /// failure.
    pub async fn reload(&self) -> Result<(), ExtensionBindError> {
        let runner = self.hooks.runner();
        let previous_flag_values = runner.get_flag_values();

        // Lifecycle event on the old host. Emit self-gates on handler
        // presence; host transport reaping is handled below regardless.
        let _ = runner
            .emit(AgentSessionEvent::SessionShutdown {
                reason: SessionShutdownReason::Reload,
                target_session_file: None,
            })
            .await;

        if let Some(host) = self.host_extension_runner() {
            let Some(runtime) = self.model_runtime() else {
                // No runtime to re-register providers against: cut the
                // session surfaces over to the null runner while the old
                // host is still live, then drain and reap it.
                self.hooks
                    .set_runner(Arc::new(super::extension_runner::NullExtensionRunner));
                self.set_host_extension_runner(None);
                self.refresh_tool_registry(&super::tools::RefreshToolRegistryOptions {
                    active_tool_names: None,
                    include_all_extension_tools: true,
                });
                host.retire_after_cutover().await;
                self.emit_session_start_reload().await;
                self.extend_resources_from_extensions(SessionStartReason::Reload.as_str())
                    .await?;
                return Ok(());
            };
            let new_host = match self
                .restart_host_for_reload(&host, &runtime, previous_flag_values)
                .await
            {
                Ok(host) => host,
                Err(error) => {
                    // Preparation failures leave the old host, providers, and
                    // tool adapters intact. Re-open its lifecycle after the
                    // shutdown notification emitted above.
                    self.emit_session_start_reload().await;
                    return Err(ExtensionBindError::HostRestart(error.to_string()));
                }
            };
            // Cut the session surfaces over while the old host is still
            // live: trait runner + concrete host handle (no downcast), then
            // the tool registry so newly registered extension tools replace
            // the old set. Concurrent hooks/tools admitted on the old runner
            // keep a live transport until the drain below completes.
            self.hooks.set_runner(
                Arc::clone(&new_host) as Arc<dyn super::extension_runner::ExtensionRunner>
            );
            self.set_host_extension_runner(Some(new_host));
            self.refresh_tool_registry(&super::tools::RefreshToolRegistryOptions {
                active_tool_names: None,
                include_all_extension_tools: true,
            });
            // Only now reap the old transport: stop new admissions, drain
            // in-flight old-host traffic, then shut it down exactly once.
            host.retire_after_cutover().await;
            self.emit_session_start_reload().await;
            self.extend_resources_from_extensions(SessionStartReason::Reload.as_str())
                .await?;
            return Ok(());
        }

        // Trait-only / test path (no concrete host).
        self.emit_session_start_reload().await;
        self.extend_resources_from_extensions(SessionStartReason::Reload.as_str())
            .await?;
        Ok(())
    }

    /// Restart the host for reload, optionally via the test-only factory.
    async fn restart_host_for_reload(
        &self,
        host: &Arc<crate::core::extension_host::HostExtensionRunner>,
        runtime: &crate::core::model_runtime::ModelRuntime,
        previous_flag_values: std::collections::HashMap<String, serde_json::Value>,
    ) -> Result<
        Arc<crate::core::extension_host::HostExtensionRunner>,
        crate::core::extension_host::HostStartError,
    > {
        #[cfg(test)]
        {
            let factory = self
                .lock_inner()
                .reload_restart_factory
                .as_ref()
                .map(Arc::clone);
            if let Some(factory) = factory {
                return host
                    .restart_and_rewire_with(
                        runtime,
                        previous_flag_values,
                        move |paths, cwd, trusted| factory(paths, cwd, trusted),
                    )
                    .await;
            }
        }
        host.restart_and_rewire(runtime, previous_flag_values).await
    }

    /// Emit `session_start{reload}` on the current (post-swap) runner.
    async fn emit_session_start_reload(&self) {
        let _ = self
            .hooks
            .runner()
            .emit(AgentSessionEvent::SessionStart {
                reason: SessionStartReason::Reload,
                previous_session_file: None,
            })
            .await;
    }

    /// Build the [`ReplacedSessionContext`] for `withSession` callbacks.
    ///
    /// The context retains the replacement session so its send methods cannot
    /// accidentally target a later runtime session.
    pub async fn create_replaced_session_context(self: &Arc<Self>) -> ReplacedSessionContext {
        let session_id = self.session_id().await;
        ReplacedSessionContext {
            session_id,
            session: Arc::clone(self),
        }
    }

    /// Build the current extension/prompt/skill slash-command catalog.
    #[must_use]
    pub fn slash_commands(&self) -> Vec<SlashCommandInfo> {
        let mut commands = Vec::new();
        let mut extension_names = std::collections::HashSet::new();

        if let Some(host) = self.host_extension_runner() {
            for command in host.registry().commands() {
                extension_names.insert(command.name.clone());
                let path = command
                    .source
                    .clone()
                    .unwrap_or_else(|| "<extension>".to_owned());
                commands.push(SlashCommandInfo {
                    name: command.name.clone(),
                    description: command.description.clone(),
                    source: SlashCommandSource::Extension,
                    source_info: create_synthetic_source_info(
                        path,
                        SyntheticSourceInfoOptions {
                            source: "extension".to_owned(),
                            scope: None,
                            origin: None,
                            base_dir: None,
                        },
                    ),
                });
            }
        }

        for name in self.hooks.runner().get_registered_commands() {
            if extension_names.insert(name.clone()) {
                commands.push(SlashCommandInfo {
                    name,
                    description: None,
                    source: SlashCommandSource::Extension,
                    source_info: create_synthetic_source_info(
                        "<extension>",
                        SyntheticSourceInfoOptions {
                            source: "extension".to_owned(),
                            scope: None,
                            origin: None,
                            base_dir: None,
                        },
                    ),
                });
            }
        }

        commands.extend(
            self.prompt_templates
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .iter()
                .map(|template| SlashCommandInfo {
                    name: template.name.clone(),
                    description: Some(template.description.clone()),
                    source: SlashCommandSource::Prompt,
                    source_info: template.source_info.clone(),
                }),
        );
        commands.extend(
            self.skills
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .iter()
                .map(|skill| SlashCommandInfo {
                    name: format!("skill:{}", skill.name),
                    description: Some(skill.description.clone()),
                    source: SlashCommandSource::Skill,
                    source_info: skill.source_info.clone(),
                }),
        );
        commands
    }

    /// Report a structured extension error to the registered listener.
    pub fn report_extension_error(&self, extension_path: &str, event: &str, error: &str) {
        let listener = self.lock_inner().extension_error_listener.clone();
        if let Some(listener) = listener {
            let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
                listener(extension_path, event, error);
            }));
        }
    }

    /// Invoke the shutdown handler, if one is bound.
    pub fn invoke_extension_shutdown_handler(&self) {
        let handler = self.lock_inner().extension_shutdown_handler.clone();
        if let Some(handler) = handler {
            handler();
        }
    }

    /// Snapshot of the currently bound extension mode.
    #[must_use]
    pub fn extension_mode(&self) -> Option<ExtensionMode> {
        self.lock_inner().extension_mode
    }

    /// Concrete host runner handle (no trait downcast).
    #[must_use]
    pub fn host_extension_runner(
        &self,
    ) -> Option<Arc<crate::core::extension_host::HostExtensionRunner>> {
        self.host_extension_runner
            .read()
            .ok()
            .and_then(|guard| guard.clone())
    }

    /// Replace the concrete host runner handle (reload path).
    pub fn set_host_extension_runner(
        &self,
        runner: Option<Arc<crate::core::extension_host::HostExtensionRunner>>,
    ) {
        if let Ok(mut guard) = self.host_extension_runner.write() {
            *guard = runner;
        }
    }

    /// Install a test-only reload restart factory (in-memory host injection).
    #[cfg(test)]
    pub(super) fn set_reload_restart_factory(&self, factory: Option<ReloadRestartFactory>) {
        self.lock_inner().reload_restart_factory = factory;
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::agent_session::extension_runner::ExtensionRunner;
    use crate::core::agent_session::{AgentSession, AgentSessionConfig};
    use futures::future::BoxFuture;
    use futures::stream::{self, BoxStream, StreamExt};
    use pi_ai::{
        AssistantMessageEvent, Context, Model, ModelCost, ModelInput, Provider, ProviderError,
        StreamOptions,
    };
    use std::collections::HashMap;
    use std::collections::HashSet;
    use std::error::Error;
    use std::io;
    use std::sync::Mutex as StdMutex;
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::time::Duration;

    use pi_ext::client::HostClient;
    use pi_ext::protocol::{Frame, FrameKind, HelloAck, decode_frame_str, encode_frame};
    use serde_json::{Map, Value, json};
    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
    use tokio::sync::mpsc;

    use crate::core::extension_host::{HostExtensionRunner, HostStartError};
    use crate::core::model_runtime::ModelRuntime;

    type TestResult<T = ()> = Result<T, Box<dyn Error>>;

    fn test_model() -> Model {
        Model {
            id: "m".to_owned(),
            name: "m".to_owned(),
            api: "test-api".to_owned(),
            provider: "test-provider".to_owned(),
            base_url: String::new(),
            reasoning: false,
            thinking_level_map: None,
            input: vec![ModelInput::Text],
            cost: ModelCost::default(),
            context_window: 8_192,
            max_tokens: 1_024,
            headers: None,
            compat: None,
            extra: std::collections::BTreeMap::new(),
        }
    }

    #[derive(Clone)]
    struct StubProvider;

    impl Provider for StubProvider {
        fn stream(
            &self,
            _model: &Model,
            _context: Context,
            _options: StreamOptions,
        ) -> BoxStream<'static, Result<AssistantMessageEvent, ProviderError>> {
            stream::empty().boxed()
        }
    }

    fn make_session() -> Result<Arc<AgentSession>, crate::core::agent_session::AgentSessionError> {
        let config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        AgentSession::new(config)
    }

    fn locked_clone<T: Clone>(value: &Mutex<T>, label: &str) -> TestResult<T> {
        value
            .lock()
            .map(|guard| guard.clone())
            .map_err(|_| io::Error::other(format!("{label} mutex poisoned")).into())
    }

    /// Runner that records an ordered lifecycle log and supports toggling
    /// handler presence plus an optional emit delay (concurrency tests).
    struct TestRunner {
        has_start: AtomicBool,
        has_shutdown: AtomicBool,
        has_resources: AtomicBool,
        /// Unified ordered call log: `session_start:{reason}:{prev|-}`,
        /// `session_shutdown:{reason}:{target|-}`,
        /// `resources_discover:{reason}`.
        calls: Arc<Mutex<Vec<String>>>,
        emit_delay: Mutex<Option<std::time::Duration>>,
        flag_values: Arc<Mutex<HashMap<String, serde_json::Value>>>,
        resource_paths: Arc<Mutex<crate::core::resources::ResourceExtensionPaths>>,
    }

    impl TestRunner {
        fn new() -> Self {
            Self {
                has_start: AtomicBool::new(false),
                has_shutdown: AtomicBool::new(false),
                has_resources: AtomicBool::new(false),
                calls: Arc::new(Mutex::new(Vec::new())),
                emit_delay: Mutex::new(None),
                flag_values: Arc::new(Mutex::new(HashMap::new())),
                resource_paths: Arc::new(Mutex::new(
                    crate::core::resources::ResourceExtensionPaths::default(),
                )),
            }
        }

        fn record(&self, entry: String) {
            if let Ok(mut g) = self.calls.lock() {
                g.push(entry);
            }
        }

        fn lifecycle_label(event: &AgentSessionEvent) -> String {
            match event {
                AgentSessionEvent::SessionStart {
                    reason,
                    previous_session_file,
                } => format!(
                    "session_start:{}:{}",
                    reason.as_str(),
                    previous_session_file.as_deref().unwrap_or("-")
                ),
                AgentSessionEvent::SessionShutdown {
                    reason,
                    target_session_file,
                } => format!(
                    "session_shutdown:{}:{}",
                    reason.as_str(),
                    target_session_file.as_deref().unwrap_or("-")
                ),
                other => other.type_name().to_owned(),
            }
        }
    }

    impl ExtensionRunner for TestRunner {
        fn has_handlers(&self, event: &str) -> bool {
            match event {
                "session_start" => self.has_start.load(Ordering::SeqCst),
                "session_shutdown" => self.has_shutdown.load(Ordering::SeqCst),
                "resources_discover" => self.has_resources.load(Ordering::SeqCst),
                _ => false,
            }
        }

        fn emit(
            &self,
            event: AgentSessionEvent,
        ) -> BoxFuture<
            '_,
            Result<
                Option<super::super::extension_runner::CancelResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            let delay = self
                .emit_delay
                .lock()
                .map(|guard| *guard)
                .unwrap_or_default();
            let label = Self::lifecycle_label(&event);
            Box::pin(async move {
                if let Some(delay) = delay {
                    tokio::time::sleep(delay).await;
                }
                self.record(label);
                Ok(None)
            })
        }

        fn emit_message_end(
            &self,
            message: pi_agent::AgentMessage,
        ) -> BoxFuture<
            '_,
            Result<
                Option<pi_agent::AgentMessage>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async move { Ok(Some(message)) })
        }

        fn emit_tool_call(
            &self,
            _tool_name: &str,
            _tool_call_id: &str,
            _input: serde_json::Map<String, serde_json::Value>,
        ) -> BoxFuture<
            '_,
            Result<
                Option<pi_agent::BeforeToolCallResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(None) })
        }

        fn emit_tool_result(
            &self,
            _tool_name: &str,
            _tool_call_id: &str,
            _input: serde_json::Map<String, serde_json::Value>,
            _content: Vec<pi_ai::ToolResultContent>,
            _details: serde_json::Value,
            _is_error: bool,
        ) -> BoxFuture<
            '_,
            Result<
                Option<pi_agent::AfterToolCallResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(None) })
        }

        fn emit_input(
            &self,
            _text: &str,
            _images: Option<serde_json::Value>,
            _source: &str,
            _streaming_behavior: Option<&str>,
        ) -> BoxFuture<
            '_,
            Result<
                super::super::extension_runner::InputTransformResult,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(super::super::extension_runner::InputTransformResult::default()) })
        }

        fn emit_before_agent_start(
            &self,
            _prompt: &str,
            _images: Option<serde_json::Value>,
        ) -> BoxFuture<
            '_,
            Result<
                Option<super::super::extension_runner::BeforeAgentStartResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(None) })
        }

        fn emit_resources_discover(
            &self,
            cwd: &str,
            reason: &str,
        ) -> BoxFuture<
            '_,
            Result<
                crate::core::resources::ResourceExtensionPaths,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            self.record(format!("resources_discover:{reason}"));
            let _ = cwd;
            let paths = self
                .resource_paths
                .lock()
                .map(|paths| paths.clone())
                .unwrap_or_default();
            Box::pin(async move { Ok(paths) })
        }

        fn get_registered_commands(&self) -> Vec<String> {
            Vec::new()
        }

        fn get_all_registered_tools(&self) -> HashMap<String, Arc<dyn pi_agent::AgentTool>> {
            HashMap::new()
        }

        fn get_flag_values(&self) -> HashMap<String, serde_json::Value> {
            self.flag_values
                .lock()
                .map(|g| g.clone())
                .unwrap_or_default()
        }
        fn execute_command(
            &self,
            _name: &str,
            _args: &str,
        ) -> BoxFuture<'_, Result<bool, super::super::extension_runner::ExtensionRunnerError>>
        {
            Box::pin(async { Ok(false) })
        }

        fn invalidate(&self) {}

        fn emit_error(&self, _message: String) {}
    }

    #[tokio::test]
    async fn has_extension_handlers_delegates_to_runner() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;
        assert!(!session.has_extension_handlers("session_start"));
        runner.has_start.store(true, Ordering::SeqCst);
        assert!(session.has_extension_handlers("session_start"));
        Ok(())
    }

    #[tokio::test]
    async fn bind_extensions_records_bindings_and_discovers_resources() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        runner.has_resources.store(true, Ordering::SeqCst);
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;

        let error_hit = Arc::new(AtomicBool::new(false));
        let error_hit_clone = Arc::clone(&error_hit);
        let bindings = ExtensionBindings {
            mode: Some(ExtensionMode::Rpc),
            on_error: Some(Arc::new(move |_path: &str, _event: &str, _error: &str| {
                error_hit_clone.store(true, Ordering::SeqCst);
            })),
            ..Default::default()
        };
        session.bind_extensions(bindings).await?;

        // Resource discovery invoked with "startup".
        let calls = locked_clone(&runner.calls, "calls")?;
        assert!(calls.iter().any(|c| c == "resources_discover:startup"));

        // Bindings recorded.
        assert_eq!(session.extension_mode(), Some(ExtensionMode::Rpc));

        // Error listener is routed.
        session.report_extension_error("extension.ts", "agent_start", "boom");
        assert!(error_hit.load(Ordering::SeqCst));
        Ok(())
    }

    #[tokio::test]
    async fn bind_emits_stored_session_start_before_discovery() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        runner.has_start.store(true, Ordering::SeqCst);
        runner.has_resources.store(true, Ordering::SeqCst);
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;

        session
            .bind_extensions(ExtensionBindings::default())
            .await?;
        let calls = locked_clone(&runner.calls, "calls")?;
        assert_eq!(
            calls,
            vec![
                "session_start:startup:-".to_owned(),
                "resources_discover:startup".to_owned(),
            ],
            "first bind must emit session_start exactly once, before discovery"
        );

        // Second bind: no re-emission (take-guard), no second discovery.
        session
            .bind_extensions(ExtensionBindings::default())
            .await?;
        let calls = locked_clone(&runner.calls, "calls")?;
        assert_eq!(
            calls,
            vec![
                "session_start:startup:-".to_owned(),
                "resources_discover:startup".to_owned(),
            ],
            "second bind must not re-emit or rediscover"
        );
        Ok(())
    }

    #[tokio::test]
    async fn concurrent_binds_emit_start_once_before_any_discovery() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        runner.has_start.store(true, Ordering::SeqCst);
        runner.has_resources.store(true, Ordering::SeqCst);
        *runner
            .emit_delay
            .lock()
            .map_err(|_| io::Error::other("emit delay mutex poisoned"))? =
            Some(std::time::Duration::from_millis(25));
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;

        let (first, second) = tokio::join!(
            session.bind_extensions(ExtensionBindings::default()),
            session.bind_extensions(ExtensionBindings::default()),
        );
        first?;
        second?;

        let calls = locked_clone(&runner.calls, "calls")?;
        assert_eq!(
            calls,
            vec![
                "session_start:startup:-".to_owned(),
                "resources_discover:startup".to_owned(),
            ],
            "concurrent binds must serialize: one start strictly before one discovery"
        );
        Ok(())
    }

    #[tokio::test]
    async fn bind_emits_replacement_reason_and_previous_file() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        runner.has_start.store(true, Ordering::SeqCst);
        runner.has_resources.store(true, Ordering::SeqCst);
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        config.session_start_event = Some(SessionStartEvent {
            reason: SessionStartReason::New,
            previous_session_file: Some("prev.jsonl".into()),
        });
        let session = AgentSession::new(config)?;

        session
            .bind_extensions(ExtensionBindings::default())
            .await?;
        let calls = locked_clone(&runner.calls, "calls")?;
        assert_eq!(
            calls,
            vec![
                "session_start:new:prev.jsonl".to_owned(),
                // Non-reload replacement reasons map to "startup" for
                // resources_discover (wire contract: startup|reload only).
                "resources_discover:startup".to_owned(),
            ]
        );
        Ok(())
    }

    #[tokio::test]
    async fn extension_resources_refresh_session_skills_and_system_prompt() -> TestResult {
        let temp = tempfile::tempdir()?;
        let cwd = temp.path().join("project");
        let agent_dir = temp.path().join("agent");
        let extension_dir = temp.path().join("extension");
        let skill_dir = extension_dir.join("skills");
        std::fs::create_dir_all(&cwd)?;
        std::fs::create_dir_all(&agent_dir)?;
        std::fs::create_dir_all(&skill_dir)?;
        std::fs::write(
            skill_dir.join("SKILL.md"),
            "---\nname: extension-skill\ndescription: extension\n---\nbody\n",
        )?;
        let extension_path = extension_dir.join("plugin.ts");
        std::fs::write(&extension_path, "")?;
        let extension_path = extension_path.to_string_lossy().into_owned();

        let runner = Arc::new(TestRunner::new());
        runner.has_resources.store(true, Ordering::SeqCst);
        *runner
            .resource_paths
            .lock()
            .map_err(|_| io::Error::other("resource paths mutex poisoned"))? =
            crate::core::resources::ResourceExtensionPaths {
                skill_paths: vec![crate::core::resources::ExtensionResourcePath::discovered(
                    "skills".to_owned(),
                    &extension_path,
                )],
                ..crate::core::resources::ResourceExtensionPaths::default()
            };

        let settings = crate::core::settings::SettingsManager::create(
            &cwd,
            Some(&agent_dir),
            crate::core::settings::SettingsManagerCreateOptions::new().project_trusted(true),
        );
        let mut loader = crate::core::resources::DefaultResourceLoader::new(
            crate::core::resources::DefaultResourceLoaderOptions {
                cwd: cwd.clone(),
                agent_dir,
                settings_manager: Some(settings),
                ..crate::core::resources::DefaultResourceLoaderOptions::default()
            },
        );
        loader.reload().await?;
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.cwd = cwd.to_string_lossy().into_owned();
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        config.initial_active_tool_names = Some(vec!["read".to_owned()]);
        config.tools = vec![Arc::new(crate::core::tools::read::ReadTool::new(&cwd))];
        config.resource_loader = Some(loader);
        config.system_prompt = "stale".to_owned();
        let session = AgentSession::new(config)?;

        session
            .bind_extensions(ExtensionBindings::default())
            .await?;
        assert!(
            session
                .agent
                .state()
                .system_prompt
                .contains("<name>extension-skill</name>")
        );
        assert!(
            session
                .skills
                .lock()
                .map_err(|_| io::Error::other("skills mutex poisoned"))?
                .iter()
                .any(|skill| skill.name == "extension-skill")
        );

        *runner
            .resource_paths
            .lock()
            .map_err(|_| io::Error::other("resource paths mutex poisoned"))? =
            crate::core::resources::ResourceExtensionPaths::default();
        session.reload().await?;
        assert!(
            !session
                .agent
                .state()
                .system_prompt
                .contains("extension-skill")
        );
        assert!(
            session
                .skills
                .lock()
                .map_err(|_| io::Error::other("skills mutex poisoned"))?
                .is_empty()
        );
        Ok(())
    }

    #[tokio::test]
    async fn reload_emits_shutdown_start_discovery_in_order() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        runner.has_shutdown.store(true, Ordering::SeqCst);
        runner.has_start.store(true, Ordering::SeqCst);
        runner.has_resources.store(true, Ordering::SeqCst);
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;

        session
            .bind_extensions(ExtensionBindings {
                mode: Some(ExtensionMode::Rpc),
                ..Default::default()
            })
            .await?;
        runner
            .calls
            .lock()
            .map_err(|_| io::Error::other("calls mutex poisoned"))?
            .clear();

        session.reload().await?;

        let calls = locked_clone(&runner.calls, "calls")?;
        assert_eq!(
            calls,
            vec![
                "session_shutdown:reload:-".to_owned(),
                "session_start:reload:-".to_owned(),
                "resources_discover:reload".to_owned(),
            ],
            "reload must emit shutdown, then start, then rediscover"
        );
        Ok(())
    }

    #[tokio::test]
    async fn reload_rediscovers_resources_without_bindings() -> TestResult {
        let runner = Arc::new(TestRunner::new());
        runner.has_shutdown.store(true, Ordering::SeqCst);
        runner.has_resources.store(true, Ordering::SeqCst);
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(runner.clone() as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;
        session.reload().await?;
        let calls = locked_clone(&runner.calls, "calls")?;
        assert!(calls.iter().any(|c| c == "session_shutdown:reload:-"));
        assert!(
            calls.iter().any(|c| c == "resources_discover:reload"),
            "reload must refresh resources even without bindings, got {calls:?}"
        );
        Ok(())
    }

    /// Runner whose lifecycle `emit` fails (dead host simulation).
    struct FailingRunner;

    impl ExtensionRunner for FailingRunner {
        fn has_handlers(&self, _event: &str) -> bool {
            true
        }

        // Lifecycle emit fails; all other methods default to null-runner
        // behavior.
        fn emit(
            &self,
            _event: AgentSessionEvent,
        ) -> BoxFuture<
            '_,
            Result<
                Option<super::super::extension_runner::CancelResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async {
                Err(
                    super::super::extension_runner::ExtensionRunnerError::Failed(
                        "host gone".into(),
                    ),
                )
            })
        }

        fn emit_message_end(
            &self,
            message: pi_agent::AgentMessage,
        ) -> BoxFuture<
            '_,
            Result<
                Option<pi_agent::AgentMessage>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async move { Ok(Some(message)) })
        }

        fn emit_tool_call(
            &self,
            _tool_name: &str,
            _tool_call_id: &str,
            _input: serde_json::Map<String, serde_json::Value>,
        ) -> BoxFuture<
            '_,
            Result<
                Option<pi_agent::BeforeToolCallResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(None) })
        }

        fn emit_tool_result(
            &self,
            _tool_name: &str,
            _tool_call_id: &str,
            _input: serde_json::Map<String, serde_json::Value>,
            _content: Vec<pi_ai::ToolResultContent>,
            _details: serde_json::Value,
            _is_error: bool,
        ) -> BoxFuture<
            '_,
            Result<
                Option<pi_agent::AfterToolCallResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(None) })
        }

        fn emit_input(
            &self,
            _text: &str,
            _images: Option<serde_json::Value>,
            _source: &str,
            _streaming_behavior: Option<&str>,
        ) -> BoxFuture<
            '_,
            Result<
                super::super::extension_runner::InputTransformResult,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(super::super::extension_runner::InputTransformResult::default()) })
        }

        fn emit_before_agent_start(
            &self,
            _prompt: &str,
            _images: Option<serde_json::Value>,
        ) -> BoxFuture<
            '_,
            Result<
                Option<super::super::extension_runner::BeforeAgentStartResult>,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(None) })
        }

        fn emit_resources_discover(
            &self,
            _cwd: &str,
            _reason: &str,
        ) -> BoxFuture<
            '_,
            Result<
                crate::core::resources::ResourceExtensionPaths,
                super::super::extension_runner::ExtensionRunnerError,
            >,
        > {
            Box::pin(async { Ok(crate::core::resources::ResourceExtensionPaths::default()) })
        }

        fn get_registered_commands(&self) -> Vec<String> {
            Vec::new()
        }

        fn get_all_registered_tools(&self) -> HashMap<String, Arc<dyn pi_agent::AgentTool>> {
            HashMap::new()
        }

        fn get_flag_values(&self) -> HashMap<String, serde_json::Value> {
            HashMap::new()
        }

        fn execute_command(
            &self,
            _name: &str,
            _args: &str,
        ) -> BoxFuture<'_, Result<bool, super::super::extension_runner::ExtensionRunnerError>>
        {
            Box::pin(async { Ok(false) })
        }

        fn invalidate(&self) {}
        fn emit_error(&self, _message: String) {}
    }

    #[tokio::test]
    async fn reload_survives_lifecycle_emit_error() -> TestResult {
        // Lifecycle emit failures are isolated (host error reporting), never
        // fatal: reload must still complete resource rediscovery.
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(Arc::new(FailingRunner) as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;
        session.reload().await?;
        Ok(())
    }

    #[tokio::test]
    async fn bind_survives_lifecycle_emit_error() -> TestResult {
        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(Arc::new(FailingRunner) as Arc<dyn ExtensionRunner>);
        let session = AgentSession::new(config)?;
        session
            .bind_extensions(ExtensionBindings::default())
            .await?;
        Ok(())
    }

    #[tokio::test]
    async fn create_replaced_session_context_forwards_sends_to_session() -> TestResult {
        let session = make_session()?;
        let ctx = session.create_replaced_session_context().await;
        assert_eq!(ctx.session_id, session.session_id().await);

        ctx.send_custom_message(
            crate::core::agent_session::prompt::CustomMessageInput {
                custom_type: "replacement-note".to_owned(),
                content: crate::core::messages::CustomMessageContent::Text("hi".to_owned()),
                display: true,
                details: None,
            },
            false,
            None,
        )
        .await?;
        ctx.send_user_message("yo", Vec::new(), None).await?;

        let messages = session.agent.state().messages;
        assert!(messages.iter().any(|message| {
            matches!(
                message,
                pi_agent::AgentMessage::Custom(custom)
                    if custom.payload.get("customType").and_then(serde_json::Value::as_str)
                        == Some("replacement-note")
            )
        }));
        assert!(messages.iter().any(|message| {
            matches!(
                message,
                pi_agent::AgentMessage::Llm(message)
                    if matches!(
                        message.as_ref(),
                        pi_ai::Message::User(user)
                            if user.content == pi_ai::UserMessageContent::Text("yo".to_owned())
                    )
            )
        }));
        Ok(())
    }

    #[tokio::test]
    async fn get_extension_source_label_strips_extension() -> TestResult {
        assert_eq!(
            AgentSession::get_extension_source_label("/foo/bar/myext.ts"),
            "extension:myext"
        );
        assert_eq!(
            AgentSession::get_extension_source_label("<inline>"),
            "extension:inline"
        );
        Ok(())
    }

    #[tokio::test]
    async fn extension_error_listener_receives_structured_fields_and_is_isolated() -> TestResult {
        let session = make_session()?;
        let received = Arc::new(Mutex::new(Vec::new()));
        let received_clone = Arc::clone(&received);
        session
            .bind_extensions(ExtensionBindings {
                on_error: Some(Arc::new(move |path, event, error| {
                    received_clone
                        .lock()
                        .unwrap_or_else(std::sync::PoisonError::into_inner)
                        .push((path.to_owned(), event.to_owned(), error.to_owned()));
                })),
                ..Default::default()
            })
            .await?;
        session.report_extension_error("/ext/tool.ts", "tool_call", "host crashed");
        assert_eq!(
            received
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .as_slice(),
            &[(
                "/ext/tool.ts".to_owned(),
                "tool_call".to_owned(),
                "host crashed".to_owned(),
            )]
        );

        session
            .bind_extensions(ExtensionBindings {
                on_error: Some(Arc::new(|_, _, _| {
                    // Deliberate unwind: proves listener panics are isolated.
                    std::panic::resume_unwind(Box::new("listener panic"));
                })),
                ..Default::default()
            })
            .await?;
        session.report_extension_error("<runtime>", "reload", "boom");
        assert!(!session.session_id().await.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn bind_extensions_with_null_runner_is_noop() -> TestResult {
        let session = make_session()?;
        session
            .bind_extensions(ExtensionBindings {
                mode: Some(ExtensionMode::Print),
                ..Default::default()
            })
            .await?;
        assert!(!session.has_extension_handlers("session_start"));
        Ok(())
    }

    #[tokio::test]
    async fn invoke_extension_shutdown_handler_calls_bound_closure() -> TestResult {
        let session = make_session()?;
        let called = Arc::new(AtomicBool::new(false));
        let called_clone = Arc::clone(&called);
        session
            .bind_extensions(ExtensionBindings {
                shutdown_handler: Some(Arc::new(move || {
                    called_clone.store(true, Ordering::SeqCst);
                })),
                ..Default::default()
            })
            .await?;
        session.invoke_extension_shutdown_handler();
        assert!(called.load(Ordering::SeqCst));
        Ok(())
    }

    // -----------------------------------------------------------------------
    // Host-backed reload cutover (in-memory transports)
    // -----------------------------------------------------------------------

    enum FakeCmd {
        Emit(Frame),
    }

    #[derive(Clone)]
    struct FakeHost {
        cmd_tx: mpsc::Sender<FakeCmd>,
        drop_methods: Arc<StdMutex<HashSet<String>>>,
        requests: Arc<StdMutex<Vec<Frame>>>,
    }

    impl FakeHost {
        fn drop_method(&self, method: &str) {
            if let Ok(mut set) = self.drop_methods.lock() {
                set.insert(method.to_owned());
            }
        }

        async fn emit(&self, frame: Frame) {
            let _ = self.cmd_tx.send(FakeCmd::Emit(frame)).await;
        }

        async fn wait_for_request(&self, method: &str) -> TestResult {
            tokio::time::timeout(Duration::from_secs(1), async {
                loop {
                    if self.requests.lock().is_ok_and(|requests| {
                        requests.iter().any(|request| request.method == method)
                    }) {
                        return;
                    }
                    tokio::task::yield_now().await;
                }
            })
            .await
            .map_err(|_| io::Error::other(format!("fake host did not receive {method}")))?;
            Ok(())
        }
    }

    fn last_request_id(host: &FakeHost, method: &str) -> TestResult<u64> {
        host.requests
            .lock()
            .map_err(|_| io::Error::other("request lock poisoned"))?
            .iter()
            .rev()
            .find(|request| request.method == method)
            .map(|request| request.id)
            .ok_or_else(|| io::Error::other(format!("no recorded {method} request")).into())
    }

    fn recorded_request_count(host: &FakeHost) -> TestResult<usize> {
        Ok(host
            .requests
            .lock()
            .map_err(|_| io::Error::other("request lock poisoned"))?
            .len())
    }

    fn cutover_snapshot(tool_name: &str, provider_name: &str) -> Value {
        json!({
            "tools": [
                {
                    "name": tool_name,
                    "label": tool_name,
                    "description": "cutover tool",
                    "parameters": {"type": "object"}
                }
            ],
            "commands": [],
            "shortcuts": [],
            "flags": [{"name": "extFlag", "type": "string", "default": "x"}],
            "renderers": [],
            "providers": [{"name": provider_name}],
            "handlers": [
                "session_start",
                "session_shutdown",
                "resources_discover",
                "tool_call"
            ],
        })
    }

    fn dispatch(
        req: &Frame,
        snapshot: &Value,
        responses: &StdMutex<HashMap<String, Value>>,
        drop_methods: &StdMutex<HashSet<String>>,
    ) -> Option<Frame> {
        if drop_methods
            .lock()
            .is_ok_and(|set| set.contains(&req.method))
        {
            return None;
        }
        let payload = if req.method == "hello" {
            serde_json::to_value(HelloAck::local()).unwrap_or(Value::Null)
        } else if req.method == "extensions.load" {
            snapshot.clone()
        } else if let Some(payload) = responses
            .lock()
            .ok()
            .and_then(|map| map.get(&req.method).cloned())
        {
            payload
        } else if req.method == pi_ext::protocol::FLAGS_SET_METHOD {
            json!({"ok": true})
        } else {
            Value::Object(Map::new())
        };
        Some(Frame {
            id: req.id,
            kind: FrameKind::Res,
            method: req.method.clone(),
            payload,
        })
    }

    async fn fake_host_task(
        read: tokio::io::DuplexStream,
        mut write: tokio::io::DuplexStream,
        snapshot: Value,
        responses: Arc<StdMutex<HashMap<String, Value>>>,
        drop_methods: Arc<StdMutex<HashSet<String>>>,
        requests: Arc<StdMutex<Vec<Frame>>>,
        mut cmd_rx: mpsc::Receiver<FakeCmd>,
    ) {
        let mut reader = BufReader::new(read);
        let mut buf = String::new();
        loop {
            tokio::select! {
                biased;
                cmd = cmd_rx.recv() => {
                    match cmd {
                        Some(FakeCmd::Emit(frame)) => {
                            let bytes = encode_frame(&frame).unwrap_or_default();
                            if !bytes.is_empty() {
                                let _ = write.write_all(&bytes).await;
                                let _ = write.flush().await;
                            }
                        }
                        None => return,
                    }
                }
                n = reader.read_line(&mut buf) => {
                    match n {
                        Ok(0) | Err(_) => return,
                        Ok(_) => {
                            if let Ok(req) = decode_frame_str(&buf) {
                                if let Ok(mut recorded) = requests.lock() {
                                    recorded.push(req.clone());
                                }
                                if let Some(resp) =
                                    dispatch(&req, &snapshot, &responses, &drop_methods)
                                {
                                    let bytes = encode_frame(&resp).unwrap_or_default();
                                    let _ = write.write_all(&bytes).await;
                                    let _ = write.flush().await;
                                }
                            }
                            buf.clear();
                        }
                    }
                }
            }
        }
    }

    async fn make_host_runner(
        snapshot: Value,
        hook_timeout: Duration,
    ) -> TestResult<(Arc<HostExtensionRunner>, FakeHost)> {
        let (client_to_host, host_read) = tokio::io::duplex(64 * 1024);
        let (host_write, client_read) = tokio::io::duplex(64 * 1024);
        let (err_write, _err_read) = tokio::io::duplex(4096);
        let client = Arc::new(HostClient::connect_boxed(
            Box::new(client_to_host),
            Box::new(client_read),
            Box::new(err_write),
            None,
        ));
        let responses = Arc::new(StdMutex::new(HashMap::new()));
        let drop_methods = Arc::new(StdMutex::new(HashSet::new()));
        let requests = Arc::new(StdMutex::new(Vec::new()));
        let (cmd_tx, cmd_rx) = mpsc::channel(64);
        tokio::spawn(fake_host_task(
            host_read,
            host_write,
            snapshot,
            Arc::clone(&responses),
            Arc::clone(&drop_methods),
            Arc::clone(&requests),
            cmd_rx,
        ));
        let runner = HostExtensionRunner::connect_with_cwd_and_trust(
            client,
            vec![],
            "/workspace",
            false,
            hook_timeout,
        )
        .await?;
        Ok((
            runner,
            FakeHost {
                cmd_tx,
                drop_methods,
                requests,
            },
        ))
    }

    type ReplacementHostSlot = Arc<StdMutex<Option<(Arc<HostExtensionRunner>, FakeHost)>>>;

    struct ReloadCutoverFixture {
        runtime: Arc<ModelRuntime>,
        old_host: Arc<HostExtensionRunner>,
        old_fake: FakeHost,
        session: Arc<AgentSession>,
        replacement: ReplacementHostSlot,
    }

    async fn reload_cutover_fixture() -> TestResult<ReloadCutoverFixture> {
        let runtime = Arc::new(ModelRuntime::create_in_memory().await?);
        let (old_host, old_fake) = make_host_runner(
            cutover_snapshot("oldTool", "oldProv"),
            Duration::from_secs(5),
        )
        .await?;
        let registration = old_host.register_providers_on(runtime.as_ref());
        assert!(registration.iter().all(|(_, result)| result.is_ok()));
        assert!(
            runtime.get_registered_provider_config("oldProv").is_some(),
            "old provider must be registered before reload"
        );

        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(Arc::clone(&old_host) as Arc<dyn ExtensionRunner>);
        config.host_extension_runner = Some(Arc::clone(&old_host));
        config.model_runtime = Some(Arc::clone(&runtime));
        config.initial_active_tool_names = Some(vec!["oldTool".to_owned()]);
        let session = AgentSession::new(config)?;
        session.refresh_tool_registry(&super::super::tools::RefreshToolRegistryOptions {
            active_tool_names: Some(vec!["oldTool".to_owned()]),
            include_all_extension_tools: true,
        });
        assert!(
            session.get_tool("oldTool").is_some(),
            "session registry must start with the old host tool"
        );

        let replacement = Arc::new(StdMutex::new(None));
        let replacement_for_factory = Arc::clone(&replacement);
        session.set_reload_restart_factory(Some(Arc::new(move |_paths, _cwd, _trusted| {
            let replacement = Arc::clone(&replacement_for_factory);
            Box::pin(async move {
                let (runner, host) = make_host_runner(
                    cutover_snapshot("newTool", "newProv"),
                    Duration::from_secs(5),
                )
                .await
                .map_err(|error| HostStartError::Load(error.to_string()))?;
                *replacement
                    .lock()
                    .map_err(|_| HostStartError::Load("replacement lock poisoned".to_owned()))? =
                    Some((Arc::clone(&runner), host));
                Ok(runner)
            })
        })));

        Ok(ReloadCutoverFixture {
            runtime,
            old_host,
            old_fake,
            session,
            replacement,
        })
    }

    async fn wait_for_session_cutover(fixture: &ReloadCutoverFixture) -> TestResult {
        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                let host = fixture.session.host_extension_runner();
                let cut_over = host
                    .as_ref()
                    .is_some_and(|host| !Arc::ptr_eq(host, &fixture.old_host))
                    && fixture.session.get_tool("newTool").is_some()
                    && fixture.session.get_tool("oldTool").is_none()
                    && fixture
                        .runtime
                        .get_registered_provider_config("newProv")
                        .is_some()
                    && fixture
                        .runtime
                        .get_registered_provider_config("oldProv")
                        .is_none()
                    && fixture.old_host.is_running();
                if cut_over {
                    break;
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .map_err(|_| {
            io::Error::other("session surfaces never cut over while old host stayed live")
        })?;
        Ok(())
    }

    async fn verify_replacement_lifecycle(fixture: &ReloadCutoverFixture) -> TestResult {
        assert!(
            !fixture.old_host.is_running(),
            "old host must be reaped only after cutover completes"
        );
        assert!(
            fixture
                .session
                .host_extension_runner()
                .is_some_and(|host| host.is_running()),
            "replacement host remains live after reload"
        );
        assert!(
            fixture.session.get_tool("newTool").is_some(),
            "tool registry must retain the replacement tool"
        );
        assert!(
            fixture
                .runtime
                .get_registered_provider_config("newProv")
                .is_some(),
            "runtime must retain the replacement provider"
        );

        let old_methods: Vec<String> = fixture
            .old_fake
            .requests
            .lock()
            .map_err(|_| io::Error::other("old request lock poisoned"))?
            .iter()
            .map(|frame| frame.method.clone())
            .collect();
        assert!(
            old_methods
                .iter()
                .any(|method| method == "session_shutdown"),
            "old host must observe session_shutdown before cutover: {old_methods:?}"
        );
        let (replacement_runner, replacement_fake) = fixture
            .replacement
            .lock()
            .map_err(|_| io::Error::other("replacement lock poisoned"))?
            .clone()
            .ok_or("replacement runner missing")?;
        assert!(
            fixture
                .session
                .host_extension_runner()
                .is_some_and(|host| Arc::ptr_eq(&host, &replacement_runner)),
            "session host handle must be the factory-produced replacement"
        );
        let new_methods: Vec<String> = replacement_fake
            .requests
            .lock()
            .map_err(|_| io::Error::other("replacement request lock poisoned"))?
            .iter()
            .map(|frame| frame.method.clone())
            .collect();
        let start_idx = new_methods
            .iter()
            .position(|method| method == "session_start")
            .ok_or("replacement must receive session_start")?;
        let discover_idx = new_methods
            .iter()
            .position(|method| method == "resources_discover")
            .ok_or("replacement must receive resources_discover")?;
        assert!(
            start_idx < discover_idx,
            "lifecycle must be start then discover on the replacement: {new_methods:?}"
        );
        replacement_runner.shutdown_once().await;
        Ok(())
    }

    #[tokio::test]
    async fn reload_cuts_over_session_surfaces_before_reaping_old_host() -> TestResult {
        let fixture = reload_cutover_fixture().await?;

        fixture.old_fake.drop_method("tool_call");
        let blocked = {
            let old_host = Arc::clone(&fixture.old_host);
            tokio::spawn(async move {
                ExtensionRunner::emit_tool_call(
                    old_host.as_ref(),
                    "read",
                    "tc-session-cutover",
                    Map::new(),
                )
                .await
            })
        };
        fixture.old_fake.wait_for_request("tool_call").await?;

        let reloading = {
            let session = Arc::clone(&fixture.session);
            tokio::spawn(async move { session.reload().await })
        };
        wait_for_session_cutover(&fixture).await?;

        assert!(
            !reloading.is_finished(),
            "reload must stay pending until blocked old traffic drains"
        );
        assert!(
            fixture.old_host.is_running(),
            "old transport must remain live until after session cutover"
        );
        let new_host = fixture
            .session
            .host_extension_runner()
            .ok_or("missing replacement host handle")?;
        assert!(
            !Arc::ptr_eq(&new_host, &fixture.old_host),
            "host handle must point at the replacement"
        );
        assert!(
            Arc::ptr_eq(
                &fixture.session.extension_runner(),
                &(Arc::clone(&new_host) as Arc<dyn ExtensionRunner>)
            ) || fixture
                .session
                .extension_runner()
                .has_handlers("session_start"),
            "trait runner must be the replacement after cutover"
        );
        assert!(
            fixture
                .session
                .extension_runner()
                .has_handlers("session_start"),
            "post-cutover trait runner must expose replacement handlers"
        );

        let id = last_request_id(&fixture.old_fake, "tool_call")?;
        fixture
            .old_fake
            .emit(Frame {
                id,
                kind: FrameKind::Res,
                method: "tool_call".to_owned(),
                payload: json!({"block": true, "reason": "session-cutover"}),
            })
            .await;
        let verdict = tokio::time::timeout(Duration::from_secs(2), blocked)
            .await???
            .ok_or("blocked old hook must resolve against the still-live transport")?;
        assert!(
            verdict.block,
            "old host verdict must survive session cutover"
        );

        tokio::time::timeout(Duration::from_secs(2), reloading)
            .await??
            .map_err(|error| io::Error::other(error.to_string()))?;
        verify_replacement_lifecycle(&fixture).await
    }

    #[tokio::test]
    async fn reload_keeps_old_surfaces_when_replacement_start_fails() -> TestResult {
        let runtime = Arc::new(ModelRuntime::create_in_memory().await?);
        let (old_host, _old_fake) = make_host_runner(
            cutover_snapshot("oldTool", "oldProv"),
            Duration::from_secs(1),
        )
        .await?;
        let registration = old_host.register_providers_on(runtime.as_ref());
        assert!(registration.iter().all(|(_, result)| result.is_ok()));

        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(Arc::clone(&old_host) as Arc<dyn ExtensionRunner>);
        config.host_extension_runner = Some(Arc::clone(&old_host));
        config.model_runtime = Some(Arc::clone(&runtime));
        config.initial_active_tool_names = Some(vec!["oldTool".to_owned()]);
        let session = AgentSession::new(config)?;
        session.refresh_tool_registry(&super::super::tools::RefreshToolRegistryOptions {
            active_tool_names: Some(vec!["oldTool".to_owned()]),
            include_all_extension_tools: true,
        });

        session.set_reload_restart_factory(Some(Arc::new(|_paths, _cwd, _trusted| {
            Box::pin(async {
                Err(HostStartError::Spawn(
                    "injected replacement start failure".to_owned(),
                ))
            })
        })));

        let error = match session.reload().await {
            Ok(()) => return Err("failed replacement start unexpectedly succeeded".into()),
            Err(error) => error,
        };
        assert!(
            matches!(error, ExtensionBindError::HostRestart(_)),
            "reload must surface host restart failure: {error}"
        );
        assert!(
            session
                .host_extension_runner()
                .is_some_and(|host| Arc::ptr_eq(&host, &old_host)),
            "failed reload must keep the old host handle"
        );
        assert!(
            old_host.is_running(),
            "failed reload must not reap the old host"
        );
        assert!(
            session.get_tool("oldTool").is_some(),
            "failed reload must keep the old tool registry"
        );
        assert!(
            runtime.get_registered_provider_config("oldProv").is_some(),
            "failed reload must keep the old provider registration"
        );

        old_host.shutdown_once().await;
        Ok(())
    }

    #[tokio::test]
    async fn reload_reaping_before_cutover_would_fail_blocked_old_traffic() -> TestResult {
        // Mutation probe: if retire/reap ran before session cutover, a blocked
        // old-host hook admitted before reload could not complete.
        let runtime = Arc::new(ModelRuntime::create_in_memory().await?);
        let (old_host, old_fake) = make_host_runner(
            cutover_snapshot("oldTool", "oldProv"),
            Duration::from_secs(5),
        )
        .await?;
        let _ = old_host.register_providers_on(runtime.as_ref());

        let mut config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())?;
        config.extension_runner = Some(Arc::clone(&old_host) as Arc<dyn ExtensionRunner>);
        config.host_extension_runner = Some(Arc::clone(&old_host));
        config.model_runtime = Some(Arc::clone(&runtime));
        let session = AgentSession::new(config)?;

        let replacement_keep = Arc::new(StdMutex::new(None::<FakeHost>));
        let replacement_keep_for_factory = Arc::clone(&replacement_keep);
        let cutover_seen = Arc::new(AtomicBool::new(false));
        let cutover_seen_for_factory = Arc::clone(&cutover_seen);
        let old_host_for_factory = Arc::clone(&old_host);
        session.set_reload_restart_factory(Some(Arc::new(move |_paths, _cwd, _trusted| {
            let cutover_seen = Arc::clone(&cutover_seen_for_factory);
            let old_host = Arc::clone(&old_host_for_factory);
            let replacement_keep = Arc::clone(&replacement_keep_for_factory);
            Box::pin(async move {
                // Capture whether the old transport is still live when the
                // replacement is prepared (it must be).
                assert!(
                    old_host.is_running(),
                    "replacement factory must run while old host is still live"
                );
                let (runner, host) = make_host_runner(
                    cutover_snapshot("newTool", "newProv"),
                    Duration::from_secs(5),
                )
                .await
                .map_err(|error| HostStartError::Load(error.to_string()))?;
                *replacement_keep.lock().map_err(|_| {
                    HostStartError::Load("replacement keep lock poisoned".to_owned())
                })? = Some(host);
                cutover_seen.store(true, Ordering::SeqCst);
                Ok(runner)
            })
        })));

        old_fake.drop_method("tool_call");
        let frames_before = recorded_request_count(&old_fake)?;
        let blocked = {
            let old_host = Arc::clone(&old_host);
            tokio::spawn(async move {
                ExtensionRunner::emit_tool_call(
                    old_host.as_ref(),
                    "read",
                    "tc-mutation",
                    Map::new(),
                )
                .await
            })
        };
        old_fake.wait_for_request("tool_call").await?;
        assert!(
            recorded_request_count(&old_fake)? > frames_before,
            "blocked hook must reach the old transport"
        );

        let reloading = {
            let session = Arc::clone(&session);
            tokio::spawn(async move { session.reload().await })
        };

        // Poll until factory has prepared the replacement (restart complete)
        // while the blocked hook is still outstanding.
        tokio::time::timeout(Duration::from_secs(2), async {
            loop {
                if cutover_seen.load(Ordering::SeqCst) && old_host.is_running() {
                    break;
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .map_err(|_| io::Error::other("replacement never prepared while old host stayed live"))?;

        // Still blocked: answering after cutover proves the transport was not
        // reaped early.
        assert!(!blocked.is_finished(), "old hook must still be in flight");
        assert!(!reloading.is_finished(), "reload must wait for old drain");

        let id = last_request_id(&old_fake, "tool_call")?;
        old_fake
            .emit(Frame {
                id,
                kind: FrameKind::Res,
                method: "tool_call".to_owned(),
                payload: json!({"block": false}),
            })
            .await;
        let _ = tokio::time::timeout(Duration::from_secs(2), blocked).await???;
        tokio::time::timeout(Duration::from_secs(2), reloading)
            .await??
            .map_err(|error| io::Error::other(error.to_string()))?;
        assert!(!old_host.is_running(), "old host reaped after drain");
        // Keep the replacement FakeHost alive until after reload settles.
        let _ = replacement_keep
            .lock()
            .map_err(|_| io::Error::other("replacement keep lock poisoned"))?
            .take();
        Ok(())
    }
}