pi-oxidized 0.1.1

package/crates/pi/src/core/agent_session_runtime.rs

@@ -0,0 +1,1612 @@
+//! `AgentSessionRuntime` — owns the current [`AgentSession`] and drives the
+//! new / switch / fork / import replacement pipeline.
+//!
+//! Ports `coding-agent/src/core/agent-session-runtime.ts`. The runtime is
+//! single-owner: each replacement operation (under a serializing async mutex)
+//! creates the replacement session, tears down the current one, atomically
+//! swaps, and rebinds listeners on the new session.
+//!
+//! Replacement order (matches TS regression 2860: `withSession` runs AFTER
+//! rebind on the NEW session):
+//! 1. `emit_before_switch` / `emit_before_fork` (cancellable).
+//! 2. Build the new session manager.
+//! 3. Call the factory → new session + services + diagnostics.
+//! 4. `teardown_current`: typed `session_shutdown{reason, targetSessionFile}`
+//!    emit, then `before_session_invalidate`, then `session.dispose()`.
+//! 5. `apply(result)`: swap session + services + diagnostics.
+//! 6. `finish_session_replacement`: `rebind_session(new_session)` then
+//!    optional `with_session(ctx)`. The mode's rebind calls `bind_extensions`
+//!    on the new session, which emits the stored
+//!    `session_start{new|resume|fork}` after the old host received its
+//!    `session_shutdown` in step 4.
+
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, RwLock};
+
+use futures::future::BoxFuture;
+use tokio::sync::Mutex as AsyncMutex;
+
+use crate::core::agent_session::events::{
+    AgentSessionEvent, SessionBeforeForkPosition, SessionBeforeSwitchReason, SessionShutdownReason,
+    SessionStartReason,
+};
+use crate::core::agent_session::{AgentSession, ReplacedSessionContext};
+use crate::core::session_transfer::SessionImportFileNotFoundError;
+use crate::core::sessions::{
+    NewSessionOptions as SessionManagerNewSessionOptions, SessionError, SessionManager,
+    assert_session_cwd_exists,
+};
+
+// ---------------------------------------------------------------------------
+// Services carried by the runtime
+// ---------------------------------------------------------------------------
+
+/// Cwd-bound service handles the runtime needs to drive replacements.
+///
+/// This is a intentionally narrow view of [`crate::core::agent_session_services::AgentSessionServices`]:
+/// the runtime only needs `cwd`, `agent_dir`, and accumulated diagnostics.
+/// The full services object (with `ModelRuntime`, `DefaultResourceLoader`)
+/// is owned by the app layer; the factory extracts what the runtime needs.
+#[derive(Clone, Debug)]
+pub struct AgentSessionRuntimeServices {
+    /// Effective working directory.
+    pub cwd: PathBuf,
+    /// Agent directory (preserved across replacements).
+    pub agent_dir: PathBuf,
+}
+
+// ---------------------------------------------------------------------------
+// Factory
+// ---------------------------------------------------------------------------
+
+/// Inputs to [`CreateAgentSessionRuntimeFactory::create`].
+#[derive(Debug)]
+pub struct CreateAgentSessionRuntimeOptions {
+    /// Effective cwd for the new runtime.
+    pub cwd: String,
+    /// Agent directory (preserved across replacements).
+    pub agent_dir: String,
+    /// Session manager to bind (already constructed for the target session).
+    pub session_manager: SessionManager,
+    /// Replacement reason (carried into `session_start`).
+    pub start_reason: SessionStartReason,
+    /// Previous session file (for `session_start.previousSessionFile`).
+    pub previous_session_file: Option<String>,
+}
+
+/// Result returned by [`CreateAgentSessionRuntimeFactory::create`].
+pub struct CreateAgentSessionRuntimeResult {
+    /// New session.
+    pub session: Arc<AgentSession>,
+    /// Services snapshot.
+    pub services: AgentSessionRuntimeServices,
+    /// Diagnostics collected during creation.
+    pub diagnostics: Vec<crate::core::agent_session_services::AgentSessionRuntimeDiagnostic>,
+    /// Optional model-fallback warning.
+    pub model_fallback_message: Option<String>,
+}
+
+/// Factory that builds a fresh session + services from a session manager.
+///
+/// Implemented by the real services factory (which composes
+/// `create_agent_session_services` + `create_agent_session_from_services`)
+/// and fake test factories.
+pub trait CreateAgentSessionRuntimeFactory: Send + Sync {
+    /// Build a new session + services from the given options.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`AgentSessionRuntimeError`] when services construction or
+    /// session creation fails.
+    fn create(
+        &self,
+        options: CreateAgentSessionRuntimeOptions,
+    ) -> BoxFuture<'_, Result<CreateAgentSessionRuntimeResult, AgentSessionRuntimeError>>;
+}
+
+// ---------------------------------------------------------------------------
+// Callbacks
+// ---------------------------------------------------------------------------
+
+/// Async callback invoked on the new session after teardown + apply, before
+/// `with_session`. Re-binds listeners / event subscriptions.
+pub type RebindSessionCallback =
+    Arc<dyn Fn(Arc<AgentSession>) -> BoxFuture<'static, ()> + Send + Sync>;
+
+/// Synchronous callback invoked after `session_shutdown` handlers finish but
+/// before the old session is invalidated (host UI teardown that must not
+/// yield to the event loop).
+pub type BeforeSessionInvalidateCallback = Arc<dyn Fn() + Send + Sync>;
+
+// ---------------------------------------------------------------------------
+// Outcomes / options
+// ---------------------------------------------------------------------------
+
+/// Outcome of `new_session` / `switch_session` / `import_from_jsonl`.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct SwitchOutcome {
+    /// True when cancelled by a `before_switch` extension hook.
+    pub cancelled: bool,
+}
+
+/// Outcome of `fork`.
+#[derive(Clone, Debug, Default)]
+pub struct ForkOutcome {
+    /// True when cancelled by a `before_fork` extension hook.
+    pub cancelled: bool,
+    /// User-message text when forking before a user message (else `None`).
+    pub selected_text: Option<String>,
+}
+
+/// Options for [`AgentSessionRuntime::new_session`].
+#[derive(Clone, Debug, Default)]
+pub struct NewSessionOptions {
+    /// Optional parent session file (chained into the new session header).
+    pub parent_session: Option<String>,
+}
+
+/// Options for [`AgentSessionRuntime::fork`].
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub enum ForkPosition {
+    /// Fork from the parent of the selected user message (default).
+    #[default]
+    Before,
+    /// Fork at the selected entry (clone its subtree).
+    At,
+}
+
+/// Options for [`AgentSessionRuntime::switch_session`].
+#[derive(Clone, Debug, Default)]
+pub struct SwitchSessionOptions {
+    /// Optional cwd override when the imported session's stored cwd is missing.
+    pub cwd_override: Option<String>,
+}
+
+// ---------------------------------------------------------------------------
+// Errors
+// ---------------------------------------------------------------------------
+
+/// Errors from runtime operations.
+#[derive(Debug, thiserror::Error)]
+pub enum AgentSessionRuntimeError {
+    /// Session I/O error.
+    #[error(transparent)]
+    Session(#[from] SessionError),
+    /// Imported session cwd does not exist.
+    #[error("Stored session working directory does not exist")]
+    MissingSessionCwd,
+    /// Imported file not found.
+    #[error(transparent)]
+    ImportNotFound(#[from] SessionImportFileNotFoundError),
+    /// Persisted session missing its file path.
+    #[error("Persisted session is missing a session file")]
+    MissingSessionFile,
+    /// Fork target is invalid (not a user message when `position = before`).
+    #[error("Invalid entry ID for forking")]
+    InvalidForkEntry,
+    /// Unflushed persisted session cannot be forked.
+    #[error(
+        "This session has not been saved yet. Wait for the first assistant response before cloning or forking it."
+    )]
+    UnflushedSession,
+    /// File transfer error during import.
+    #[error("{0}")]
+    Transfer(String),
+    /// Factory failed to build the replacement runtime.
+    #[error("runtime replacement failed: {0}")]
+    Factory(String),
+}
+
+// ---------------------------------------------------------------------------
+// AgentSessionRuntime
+// ---------------------------------------------------------------------------
+
+/// Owns the current [`AgentSession`] plus cwd-bound services, and drives the
+/// replacement pipeline for new / switch / fork / import.
+///
+/// Not `Clone`. Callers hold it behind `Arc<AgentSessionRuntime>`; interior
+/// mutability covers session + services so replacement operations can run
+/// through a shared reference. A serializing async mutex prevents concurrent
+/// replacement operations from interleaving teardown / apply / rebind.
+pub struct AgentSessionRuntime {
+    session: RwLock<Arc<AgentSession>>,
+    services: RwLock<AgentSessionRuntimeServices>,
+    factory: Arc<dyn CreateAgentSessionRuntimeFactory>,
+    diagnostics: RwLock<Vec<crate::core::agent_session_services::AgentSessionRuntimeDiagnostic>>,
+    model_fallback_message: RwLock<Option<String>>,
+    rebind_session: RwLock<Option<RebindSessionCallback>>,
+    before_session_invalidate: RwLock<Option<BeforeSessionInvalidateCallback>>,
+    /// Serializes all replacement operations (new/switch/fork/import/dispose).
+    replacement_lock: AsyncMutex<()>,
+}
+
+impl AgentSessionRuntime {
+    /// Construct the runtime from an initial session + services + factory.
+    #[must_use]
+    pub fn new(
+        session: Arc<AgentSession>,
+        services: AgentSessionRuntimeServices,
+        factory: Arc<dyn CreateAgentSessionRuntimeFactory>,
+        diagnostics: Vec<crate::core::agent_session_services::AgentSessionRuntimeDiagnostic>,
+        model_fallback_message: Option<String>,
+    ) -> Self {
+        Self {
+            session: RwLock::new(session),
+            services: RwLock::new(services),
+            factory,
+            diagnostics: RwLock::new(diagnostics),
+            model_fallback_message: RwLock::new(model_fallback_message),
+            rebind_session: RwLock::new(None),
+            before_session_invalidate: RwLock::new(None),
+            replacement_lock: AsyncMutex::new(()),
+        }
+    }
+
+    /// Snapshot of the current session (cheap Arc clone).
+    #[must_use]
+    pub fn session(&self) -> Arc<AgentSession> {
+        self.read_session()
+    }
+
+    /// Effective cwd.
+    #[must_use]
+    pub fn cwd(&self) -> String {
+        self.services
+            .read()
+            .map(|g| g.cwd.to_string_lossy().into_owned())
+            .unwrap_or_default()
+    }
+
+    /// Agent directory (preserved across replacements).
+    #[must_use]
+    pub fn agent_dir(&self) -> String {
+        self.services
+            .read()
+            .map(|g| g.agent_dir.to_string_lossy().into_owned())
+            .unwrap_or_default()
+    }
+
+    /// Snapshot of diagnostics.
+    #[must_use]
+    pub fn diagnostics(
+        &self,
+    ) -> Vec<crate::core::agent_session_services::AgentSessionRuntimeDiagnostic> {
+        self.diagnostics
+            .read()
+            .map(|g| g.clone())
+            .unwrap_or_default()
+    }
+
+    /// Snapshot of the model-fallback warning.
+    #[must_use]
+    pub fn model_fallback_message(&self) -> Option<String> {
+        self.model_fallback_message
+            .read()
+            .map(|g| g.clone())
+            .unwrap_or_default()
+    }
+
+    /// Set the rebind callback invoked after each replacement.
+    pub fn set_rebind_session(&self, callback: Option<RebindSessionCallback>) {
+        if let Ok(mut g) = self.rebind_session.write() {
+            *g = callback;
+        }
+    }
+
+    /// Set the synchronous pre-invalidate callback invoked during teardown.
+    pub fn set_before_session_invalidate(&self, callback: Option<BeforeSessionInvalidateCallback>) {
+        if let Ok(mut g) = self.before_session_invalidate.write() {
+            *g = callback;
+        }
+    }
+
+    // ----- Replacement operations ----------------------------------------
+
+    /// Switch to a different session file (resume).
+    ///
+    /// # Errors
+    ///
+    /// See [`AgentSessionRuntimeError`].
+    pub async fn switch_session(
+        &self,
+        session_path: &str,
+        options: SwitchSessionOptions,
+    ) -> Result<SwitchOutcome, AgentSessionRuntimeError> {
+        let _guard = self.replacement_lock.lock().await;
+        let before = self
+            .emit_before_switch(SessionStartReason::Resume, Some(session_path))
+            .await;
+        if before.cancelled {
+            return Ok(before);
+        }
+
+        let previous_session_file = self.session_file_for_teardown().await;
+        let session_manager =
+            SessionManager::open(session_path, None, options.cwd_override.as_deref())?;
+        self.assert_cwd(&session_manager)?;
+        let new_cwd = session_manager.get_cwd().to_owned();
+        let target_session_file = session_manager.get_session_file().map(str::to_owned);
+        let result = self
+            .factory
+            .create(CreateAgentSessionRuntimeOptions {
+                cwd: new_cwd.clone(),
+                agent_dir: self.agent_dir(),
+                session_manager,
+                start_reason: SessionStartReason::Resume,
+                previous_session_file,
+            })
+            .await?;
+        self.teardown_current(
+            SessionShutdownReason::Resume,
+            target_session_file.as_deref(),
+        )
+        .await;
+        self.apply(result);
+        self.finish_session_replacement(None).await;
+        Ok(SwitchOutcome { cancelled: false })
+    }
+
+    /// Start a new session in the current cwd.
+    ///
+    /// # Errors
+    ///
+    /// See [`AgentSessionRuntimeError`].
+    pub async fn new_session(
+        &self,
+        options: NewSessionOptions,
+    ) -> Result<SwitchOutcome, AgentSessionRuntimeError> {
+        let _guard = self.replacement_lock.lock().await;
+        let before = self.emit_before_switch(SessionStartReason::New, None).await;
+        if before.cancelled {
+            return Ok(before);
+        }
+
+        let previous_session_file = self.session_file_for_teardown().await;
+        let cwd = self.cwd();
+        let session_manager = {
+            let session = self.read_session();
+            let sm = session.session_manager();
+            let sm = sm.lock().await;
+            if sm.is_persisted() {
+                let session_dir = sm.get_session_dir().to_owned();
+                let mut new_sm = SessionManager::create(&cwd, Some(&session_dir), None)?;
+                if let Some(parent) = options.parent_session.as_deref() {
+                    new_sm.new_session(Some(SessionManagerNewSessionOptions {
+                        id: None,
+                        parent_session: Some(parent.to_owned()),
+                    }))?;
+                }
+                new_sm
+            } else {
+                let opts = options.parent_session.as_deref().map(|parent| {
+                    SessionManagerNewSessionOptions {
+                        id: None,
+                        parent_session: Some(parent.to_owned()),
+                    }
+                });
+                SessionManager::in_memory(Some(&cwd), opts)?
+            }
+        };
+        let target_session_file = session_manager.get_session_file().map(str::to_owned);
+        let result = self
+            .factory
+            .create(CreateAgentSessionRuntimeOptions {
+                cwd: cwd.clone(),
+                agent_dir: self.agent_dir(),
+                session_manager,
+                start_reason: SessionStartReason::New,
+                previous_session_file,
+            })
+            .await?;
+        self.teardown_current(SessionShutdownReason::New, target_session_file.as_deref())
+            .await;
+        self.apply(result);
+        self.finish_session_replacement(None).await;
+        Ok(SwitchOutcome { cancelled: false })
+    }
+
+    /// Fork the session at `entry_id`.
+    ///
+    /// `position = Before` (default): the entry must be a user message; the
+    /// fork point is that message's parent and `selected_text` is returned.
+    /// `position = At`: the fork point is the entry itself (clone subtree).
+    ///
+    /// For **persisted** sessions, the forked branch is written to a new file
+    /// under the session directory. For **in-memory** sessions, a fresh
+    /// in-memory manager is constructed (the mutated branch on the old
+    /// manager is not carried over — until `AgentSession` accepts a shared
+    /// `Arc<AsyncMutex<SessionManager>>`, matching the TS single-reference
+    /// pattern).
+    ///
+    /// # Errors
+    ///
+    /// See [`AgentSessionRuntimeError`].
+    pub async fn fork(
+        &self,
+        entry_id: &str,
+        position: ForkPosition,
+    ) -> Result<ForkOutcome, AgentSessionRuntimeError> {
+        let _guard = self.replacement_lock.lock().await;
+        let before = self.emit_before_fork(entry_id, position).await;
+        if before.cancelled {
+            return Ok(ForkOutcome {
+                cancelled: true,
+                selected_text: None,
+            });
+        }
+
+        let (target_leaf_id, selected_text) = {
+            let session = self.read_session();
+            let sm = session.session_manager();
+            let sm = sm.lock().await;
+            let selected_entry = sm
+                .get_entry(entry_id)
+                .ok_or(AgentSessionRuntimeError::InvalidForkEntry)?;
+            match position {
+                ForkPosition::At => (selected_entry.id().map(str::to_owned), None),
+                ForkPosition::Before => {
+                    let is_user = matches!(
+                        selected_entry,
+                        crate::core::sessions::SessionEntry::Message(m) if m.message.role() == "user"
+                    );
+                    if !is_user {
+                        return Err(AgentSessionRuntimeError::InvalidForkEntry);
+                    }
+                    let parent = selected_entry.parent_id().map(str::to_owned);
+                    let text = extract_user_message_text_from_entry(selected_entry);
+                    (parent, text)
+                }
+            }
+        };
+
+        let previous_session_file = self.session_file_for_teardown().await;
+        let cwd = self.cwd();
+        let agent_dir = self.agent_dir();
+
+        // Build the new session manager.
+        let session = self.read_session();
+        let sm = session.session_manager();
+        let sm_guard = sm.lock().await;
+        let session_manager = if sm_guard.is_persisted() {
+            let current_session_file = sm_guard
+                .get_session_file()
+                .map(str::to_owned)
+                .ok_or(AgentSessionRuntimeError::MissingSessionFile)?;
+            let session_dir = sm_guard.get_session_dir().to_owned();
+            match target_leaf_id.as_deref() {
+                None => {
+                    let mut new_sm = SessionManager::create(&cwd, Some(&session_dir), None)?;
+                    new_sm.new_session(Some(SessionManagerNewSessionOptions {
+                        id: None,
+                        parent_session: Some(current_session_file.clone()),
+                    }))?;
+                    new_sm
+                }
+                Some(leaf) => {
+                    if !Path::new(&current_session_file).exists() {
+                        return Err(AgentSessionRuntimeError::UnflushedSession);
+                    }
+                    let mut reopened =
+                        SessionManager::open(&current_session_file, Some(&session_dir), None)?;
+                    let forked_path = reopened.create_branched_session(leaf)?;
+                    if forked_path.is_none() {
+                        return Err(AgentSessionRuntimeError::InvalidForkEntry);
+                    }
+                    reopened
+                }
+            }
+        } else {
+            // In-memory: cannot extract the mutated manager from the live Arc.
+            // Fall back to a fresh in-memory manager with parent linkage.
+            let opts = match target_leaf_id.as_deref() {
+                Some(_) => None, // branch state lost for in-memory (documented)
+                None => Some(SessionManagerNewSessionOptions {
+                    id: None,
+                    parent_session: previous_session_file.clone(),
+                }),
+            };
+            SessionManager::in_memory(Some(&cwd), opts)?
+        };
+        let new_cwd = session_manager.get_cwd().to_owned();
+        let target_session_file = session_manager.get_session_file().map(str::to_owned);
+        drop(sm_guard);
+        drop(session);
+
+        let result = self
+            .factory
+            .create(CreateAgentSessionRuntimeOptions {
+                cwd: new_cwd,
+                agent_dir,
+                session_manager,
+                start_reason: SessionStartReason::Fork,
+                previous_session_file,
+            })
+            .await?;
+        self.teardown_current(SessionShutdownReason::Fork, target_session_file.as_deref())
+            .await;
+        self.apply(result);
+        self.finish_session_replacement(None).await;
+        Ok(ForkOutcome {
+            cancelled: false,
+            selected_text,
+        })
+    }
+
+    /// Import a JSONL session file and switch to it.
+    ///
+    /// # Errors
+    ///
+    /// Returns [`AgentSessionRuntimeError::ImportNotFound`] when the input
+    /// path does not exist.
+    pub async fn import_from_jsonl(
+        &self,
+        input_path: &str,
+        cwd_override: Option<&str>,
+    ) -> Result<SwitchOutcome, AgentSessionRuntimeError> {
+        let _guard = self.replacement_lock.lock().await;
+        let resolved = crate::core::config::resolve_path(input_path)
+            .to_string_lossy()
+            .into_owned();
+        if !Path::new(&resolved).exists() {
+            return Err(AgentSessionRuntimeError::ImportNotFound(
+                SessionImportFileNotFoundError::new(&resolved),
+            ));
+        }
+
+        let session_dir = {
+            let session = self.read_session();
+            let sm = session.session_manager();
+            let sm = sm.lock().await;
+            sm.get_session_dir().to_owned()
+        };
+        if !Path::new(&session_dir).exists() {
+            std::fs::create_dir_all(&session_dir)
+                .map_err(|e| AgentSessionRuntimeError::Transfer(e.to_string()))?;
+        }
+
+        let file_name = Path::new(&resolved).file_name().map_or_else(
+            || "imported.jsonl".to_owned(),
+            |name| name.to_string_lossy().into_owned(),
+        );
+        let destination = Path::new(&session_dir)
+            .join(&file_name)
+            .to_string_lossy()
+            .into_owned();
+
+        let before = self
+            .emit_before_switch(SessionStartReason::Resume, Some(&destination))
+            .await;
+        if before.cancelled {
+            return Ok(before);
+        }
+
+        let previous_session_file = self.session_file_for_teardown().await;
+        let dest_canonical = std::fs::canonicalize(&destination).map_or_else(
+            |_| destination.clone(),
+            |path| path.to_string_lossy().into_owned(),
+        );
+        let resolved_canonical = std::fs::canonicalize(&resolved).map_or_else(
+            |_| resolved.clone(),
+            |path| path.to_string_lossy().into_owned(),
+        );
+        if dest_canonical != resolved_canonical {
+            std::fs::copy(&resolved, &destination)
+                .map_err(|e| AgentSessionRuntimeError::Transfer(e.to_string()))?;
+        }
+
+        let session_manager = SessionManager::open(&destination, Some(&session_dir), cwd_override)?;
+        self.assert_cwd(&session_manager)?;
+        let new_cwd = session_manager.get_cwd().to_owned();
+        let target_session_file = session_manager.get_session_file().map(str::to_owned);
+        let result = self
+            .factory
+            .create(CreateAgentSessionRuntimeOptions {
+                cwd: new_cwd,
+                agent_dir: self.agent_dir(),
+                session_manager,
+                start_reason: SessionStartReason::Resume,
+                previous_session_file,
+            })
+            .await?;
+        self.teardown_current(
+            SessionShutdownReason::Resume,
+            target_session_file.as_deref(),
+        )
+        .await;
+        self.apply(result);
+        self.finish_session_replacement(None).await;
+        Ok(SwitchOutcome { cancelled: false })
+    }
+
+    /// Dispose the current session and the runtime.
+    pub async fn dispose(&self) {
+        let _guard = self.replacement_lock.lock().await;
+        self.teardown_current(SessionShutdownReason::Quit, None)
+            .await;
+    }
+
+    // ----- Internal helpers ----------------------------------------------
+
+    fn read_session(&self) -> Arc<AgentSession> {
+        self.session.read().map_or_else(
+            |poisoned| Arc::clone(&*poisoned.into_inner()),
+            |guard| Arc::clone(&*guard),
+        )
+    }
+
+    async fn session_file_for_teardown(&self) -> Option<String> {
+        self.read_session().session_file().await
+    }
+
+    /// Emit `session_before_switch` (cancellable) when handlers are registered.
+    async fn emit_before_switch(
+        &self,
+        reason: SessionStartReason,
+        target: Option<&str>,
+    ) -> SwitchOutcome {
+        let runner = self.read_session().extension_runner();
+        if !runner.has_handlers("session_before_switch") {
+            return SwitchOutcome { cancelled: false };
+        }
+        let reason = if reason == SessionStartReason::Resume {
+            SessionBeforeSwitchReason::Resume
+        } else {
+            SessionBeforeSwitchReason::New
+        };
+        match runner
+            .emit(AgentSessionEvent::SessionBeforeSwitch {
+                reason,
+                target_session_file: target.map(str::to_owned),
+            })
+            .await
+        {
+            Ok(result) => SwitchOutcome {
+                cancelled: result.is_some_and(|result| result.cancel),
+            },
+            Err(error) => {
+                runner.emit_error(error.to_string());
+                SwitchOutcome { cancelled: false }
+            }
+        }
+    }
+
+    /// Emit `session_before_fork` (cancellable) when handlers are registered.
+    async fn emit_before_fork(&self, entry_id: &str, position: ForkPosition) -> SwitchOutcome {
+        let runner = self.read_session().extension_runner();
+        if !runner.has_handlers("session_before_fork") {
+            return SwitchOutcome { cancelled: false };
+        }
+        let position = match position {
+            ForkPosition::Before => SessionBeforeForkPosition::Before,
+            ForkPosition::At => SessionBeforeForkPosition::At,
+        };
+        match runner
+            .emit(AgentSessionEvent::SessionBeforeFork {
+                entry_id: entry_id.to_owned(),
+                position,
+            })
+            .await
+        {
+            Ok(result) => SwitchOutcome {
+                cancelled: result.is_some_and(|result| result.cancel),
+            },
+            Err(error) => {
+                runner.emit_error(error.to_string());
+                SwitchOutcome { cancelled: false }
+            }
+        }
+    }
+
+    /// Tear down the current session: typed `session_shutdown` emit →
+    /// pre-invalidate → dispose.
+    ///
+    /// The emit self-gates on `session_shutdown` handler presence and host
+    /// errors are isolated. Host process reap is always performed exactly
+    /// once inside [`AgentSession::dispose`] when a concrete host is present,
+    /// even without `session_shutdown` handlers.
+    async fn teardown_current(
+        &self,
+        reason: SessionShutdownReason,
+        target_session_file: Option<&str>,
+    ) {
+        let session = self.read_session();
+        let runner = session.extension_runner();
+        let _ = runner
+            .emit(AgentSessionEvent::SessionShutdown {
+                reason,
+                target_session_file: target_session_file.map(str::to_owned),
+            })
+            .await;
+        // Pre-invalidate callback (host UI teardown, sync). Falls back to the
+        // session-bound shutdown handler when no runtime-level callback exists.
+        let runtime_cb = self
+            .before_session_invalidate
+            .read()
+            .ok()
+            .and_then(|g| g.clone());
+        if let Some(cb) = runtime_cb {
+            cb();
+        } else {
+            session.invoke_extension_shutdown_handler();
+        }
+        // dispose always awaits host process reap exactly once when bound.
+        session.dispose().await;
+    }
+
+    /// Apply the factory result: swap session + services + diagnostics.
+    fn apply(&self, result: CreateAgentSessionRuntimeResult) {
+        if let Ok(mut g) = self.session.write() {
+            *g = result.session;
+        }
+        if let Ok(mut g) = self.services.write() {
+            *g = result.services;
+        }
+        if let Ok(mut g) = self.diagnostics.write() {
+            *g = result.diagnostics;
+        }
+        if let Ok(mut g) = self.model_fallback_message.write() {
+            *g = result.model_fallback_message;
+        }
+    }
+
+    /// Invoke the rebind callback on the new session + optional `with_session`.
+    ///
+    /// Runs AFTER apply, so `self.session()` is the new session (regression 2860).
+    async fn finish_session_replacement(
+        &self,
+        with_session: Option<
+            Arc<dyn Fn(ReplacedSessionContext) -> BoxFuture<'static, ()> + Send + Sync>,
+        >,
+    ) {
+        let session = self.read_session();
+        let rebind = self.rebind_session.read().ok().and_then(|g| g.clone());
+        if let Some(rebind) = rebind {
+            rebind(Arc::clone(&session)).await;
+        }
+        if let Some(with) = with_session {
+            let ctx = session.create_replaced_session_context().await;
+            with(ctx).await;
+        }
+    }
+
+    /// Assert the session cwd exists; fall back to the runtime cwd.
+    fn assert_cwd(&self, session_manager: &SessionManager) -> Result<(), AgentSessionRuntimeError> {
+        assert_session_cwd_exists(session_manager, &self.cwd())
+            .map_err(|_| AgentSessionRuntimeError::MissingSessionCwd)
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/// Extract concatenated text from a user-message entry (for `fork` `selected_text`).
+fn extract_user_message_text_from_entry(
+    entry: &crate::core::sessions::SessionEntry,
+) -> Option<String> {
+    use crate::core::sessions::SessionEntry;
+    let SessionEntry::Message(m) = entry else {
+        return None;
+    };
+    let text = crate::core::agent_session::tree::extract_user_message_text_pub(&m.message);
+    if text.is_empty() { None } else { Some(text) }
+}
+
+/// Create the initial runtime from a factory and initial session manager.
+///
+/// # Errors
+///
+/// Returns [`AgentSessionRuntimeError`] when the factory fails.
+pub async fn create_agent_session_runtime(
+    factory: Arc<dyn CreateAgentSessionRuntimeFactory>,
+    cwd: String,
+    agent_dir: String,
+    session_manager: SessionManager,
+) -> Result<AgentSessionRuntime, AgentSessionRuntimeError> {
+    let result = factory
+        .create(CreateAgentSessionRuntimeOptions {
+            cwd,
+            agent_dir,
+            session_manager,
+            start_reason: SessionStartReason::Startup,
+            previous_session_file: None,
+        })
+        .await?;
+    Ok(AgentSessionRuntime::new(
+        result.session,
+        result.services,
+        factory,
+        result.diagnostics,
+        result.model_fallback_message,
+    ))
+}
+
+// ---------------------------------------------------------------------------
+// Re-exports
+// ---------------------------------------------------------------------------
+
+pub use crate::core::session_transfer::SessionImportFileNotFoundError as SessionImportFileNotFound;
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::agent_session::AgentSessionConfig;
+    use futures::stream::{self, BoxStream, StreamExt};
+    use pi_ai::{
+        AssistantMessageEvent, Context, Model, ModelCost, ModelInput, Provider, ProviderError,
+        StreamOptions,
+    };
+    use std::io;
+    use std::sync::Mutex;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    type TestResult<T = ()> = Result<T, Box<dyn std::error::Error>>;
+
+    fn failure(message: &'static str) -> io::Error {
+        io::Error::other(message)
+    }
+
+    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()
+        }
+    }
+
+    /// Factory that produces a fresh in-memory session per call.
+    struct TestFactory {
+        calls: Arc<AtomicUsize>,
+    }
+
+    impl TestFactory {
+        fn new() -> Self {
+            Self {
+                calls: Arc::new(AtomicUsize::new(0)),
+            }
+        }
+    }
+
+    impl CreateAgentSessionRuntimeFactory for TestFactory {
+        fn create(
+            &self,
+            options: CreateAgentSessionRuntimeOptions,
+        ) -> BoxFuture<'_, Result<CreateAgentSessionRuntimeResult, AgentSessionRuntimeError>>
+        {
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            Box::pin(async move {
+                let config = AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())
+                    .map_err(|e| AgentSessionRuntimeError::Factory(e.to_string()))?;
+                let session = AgentSession::new(config)
+                    .map_err(|e| AgentSessionRuntimeError::Factory(e.to_string()))?;
+                Ok(CreateAgentSessionRuntimeResult {
+                    session,
+                    services: AgentSessionRuntimeServices {
+                        cwd: PathBuf::from(&options.cwd),
+                        agent_dir: PathBuf::from(&options.agent_dir),
+                    },
+                    diagnostics: Vec::new(),
+                    model_fallback_message: None,
+                })
+            })
+        }
+    }
+
+    /// Extension runner recording lifecycle `emit` calls (shared across the
+    /// sessions a recording factory creates).
+    struct EmitRecordingRunner {
+        log: Mutex<Vec<String>>,
+    }
+
+    impl EmitRecordingRunner {
+        fn new() -> Self {
+            Self {
+                log: Mutex::new(Vec::new()),
+            }
+        }
+
+        fn log_clone(&self) -> Vec<String> {
+            self.log
+                .lock()
+                .map_or_else(|p| p.into_inner().clone(), |g| g.clone())
+        }
+    }
+
+    impl crate::core::agent_session::ExtensionRunner for EmitRecordingRunner {
+        fn has_handlers(&self, _event: &str) -> bool {
+            true
+        }
+
+        fn emit(
+            &self,
+            event: AgentSessionEvent,
+        ) -> BoxFuture<
+            '_,
+            Result<
+                Option<crate::core::agent_session::CancelResult>,
+                crate::core::agent_session::ExtensionRunnerError,
+            >,
+        > {
+            let entry = 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(),
+            };
+            if let Ok(mut g) = self.log.lock() {
+                g.push(entry);
+            }
+            Box::pin(async { Ok(None) })
+        }
+
+        fn emit_message_end(
+            &self,
+            message: pi_agent::AgentMessage,
+        ) -> BoxFuture<
+            '_,
+            Result<
+                Option<pi_agent::AgentMessage>,
+                crate::core::agent_session::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>,
+                crate::core::agent_session::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>,
+                crate::core::agent_session::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<
+                crate::core::agent_session::InputTransformResult,
+                crate::core::agent_session::ExtensionRunnerError,
+            >,
+        > {
+            Box::pin(async { Ok(crate::core::agent_session::InputTransformResult::default()) })
+        }
+
+        fn emit_before_agent_start(
+            &self,
+            _prompt: &str,
+            _images: Option<serde_json::Value>,
+        ) -> BoxFuture<
+            '_,
+            Result<
+                Option<crate::core::agent_session::BeforeAgentStartResult>,
+                crate::core::agent_session::ExtensionRunnerError,
+            >,
+        > {
+            Box::pin(async { Ok(None) })
+        }
+
+        fn emit_resources_discover(
+            &self,
+            _cwd: &str,
+            _reason: &str,
+        ) -> BoxFuture<
+            '_,
+            Result<
+                crate::core::resources::ResourceExtensionPaths,
+                crate::core::agent_session::ExtensionRunnerError,
+            >,
+        > {
+            Box::pin(async { Ok(crate::core::resources::ResourceExtensionPaths::default()) })
+        }
+
+        fn get_registered_commands(&self) -> Vec<String> {
+            Vec::new()
+        }
+
+        fn execute_command<'a>(
+            &'a self,
+            _name: &'a str,
+            _args: &'a str,
+        ) -> BoxFuture<'a, Result<bool, crate::core::agent_session::ExtensionRunnerError>> {
+            Box::pin(async { Ok(false) })
+        }
+
+        fn get_all_registered_tools(
+            &self,
+        ) -> std::collections::HashMap<String, Arc<dyn pi_agent::AgentTool>> {
+            std::collections::HashMap::new()
+        }
+
+        fn get_flag_values(&self) -> std::collections::HashMap<String, serde_json::Value> {
+            std::collections::HashMap::new()
+        }
+
+        fn invalidate(&self) {}
+
+        fn emit_error(&self, _message: String) {}
+    }
+
+    /// Factory recording every `start_reason` and installing a shared
+    /// recording runner on each created session.
+    struct RecordingFactory {
+        reasons: Mutex<Vec<SessionStartReason>>,
+        runner: Arc<EmitRecordingRunner>,
+    }
+
+    impl RecordingFactory {
+        fn new(runner: Arc<EmitRecordingRunner>) -> Self {
+            Self {
+                reasons: Mutex::new(Vec::new()),
+                runner,
+            }
+        }
+
+        fn reasons_clone(&self) -> Vec<SessionStartReason> {
+            self.reasons
+                .lock()
+                .map_or_else(|p| p.into_inner().clone(), |g| g.clone())
+        }
+    }
+
+    impl CreateAgentSessionRuntimeFactory for RecordingFactory {
+        fn create(
+            &self,
+            options: CreateAgentSessionRuntimeOptions,
+        ) -> BoxFuture<'_, Result<CreateAgentSessionRuntimeResult, AgentSessionRuntimeError>>
+        {
+            if let Ok(mut g) = self.reasons.lock() {
+                g.push(options.start_reason);
+            }
+            Box::pin(async move {
+                let mut config =
+                    AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())
+                        .map_err(|e| AgentSessionRuntimeError::Factory(e.to_string()))?;
+                config.session_manager = options.session_manager;
+                config.extension_runner = Some(Arc::clone(&self.runner)
+                    as Arc<dyn crate::core::agent_session::ExtensionRunner>);
+                let session = AgentSession::new(config)
+                    .map_err(|e| AgentSessionRuntimeError::Factory(e.to_string()))?;
+                Ok(CreateAgentSessionRuntimeResult {
+                    session,
+                    services: AgentSessionRuntimeServices {
+                        cwd: PathBuf::from(&options.cwd),
+                        agent_dir: PathBuf::from(&options.agent_dir),
+                    },
+                    diagnostics: Vec::new(),
+                    model_fallback_message: None,
+                })
+            })
+        }
+    }
+
+    struct GatedTestFactory {
+        calls: AtomicUsize,
+        active_replacements: AtomicUsize,
+        entered: tokio::sync::mpsc::Sender<usize>,
+        gates: [Arc<tokio::sync::Semaphore>; 2],
+    }
+
+    impl GatedTestFactory {
+        fn new(
+            entered: tokio::sync::mpsc::Sender<usize>,
+            gates: [Arc<tokio::sync::Semaphore>; 2],
+        ) -> Self {
+            Self {
+                calls: AtomicUsize::new(0),
+                active_replacements: AtomicUsize::new(0),
+                entered,
+                gates,
+            }
+        }
+    }
+
+    impl CreateAgentSessionRuntimeFactory for GatedTestFactory {
+        fn create(
+            &self,
+            options: CreateAgentSessionRuntimeOptions,
+        ) -> BoxFuture<'_, Result<CreateAgentSessionRuntimeResult, AgentSessionRuntimeError>>
+        {
+            let call = self.calls.fetch_add(1, Ordering::SeqCst);
+            Box::pin(async move {
+                if call > 0 {
+                    self.entered.try_send(call).map_err(|error| {
+                        AgentSessionRuntimeError::Factory(format!(
+                            "failed to report replacement factory entry {call}: {error}"
+                        ))
+                    })?;
+                    if self.active_replacements.swap(1, Ordering::SeqCst) != 0 {
+                        return Err(AgentSessionRuntimeError::Factory(
+                            "replacement factories overlapped".to_owned(),
+                        ));
+                    }
+                    let gate = self.gates.get(call - 1).ok_or_else(|| {
+                        AgentSessionRuntimeError::Factory(format!(
+                            "unexpected replacement factory call {call}"
+                        ))
+                    })?;
+                    gate.acquire()
+                        .await
+                        .map_err(|error| {
+                            AgentSessionRuntimeError::Factory(format!(
+                                "replacement factory gate {call} closed: {error}"
+                            ))
+                        })?
+                        .forget();
+                }
+
+                let result = (|| {
+                    let config =
+                        AgentSessionConfig::test_config(Arc::new(StubProvider), test_model())
+                            .map_err(|e| AgentSessionRuntimeError::Factory(e.to_string()))?;
+                    let session = AgentSession::new(config)
+                        .map_err(|e| AgentSessionRuntimeError::Factory(e.to_string()))?;
+                    Ok(CreateAgentSessionRuntimeResult {
+                        session,
+                        services: AgentSessionRuntimeServices {
+                            cwd: PathBuf::from(&options.cwd),
+                            agent_dir: PathBuf::from(&options.agent_dir),
+                        },
+                        diagnostics: Vec::new(),
+                        model_fallback_message: None,
+                    })
+                })();
+                if call > 0 {
+                    self.active_replacements.store(0, Ordering::SeqCst);
+                }
+                result
+            })
+        }
+    }
+
+    async fn make_runtime() -> TestResult<AgentSessionRuntime> {
+        let factory = Arc::new(TestFactory::new());
+        let session_manager = SessionManager::in_memory(Some("."), None)?;
+        Ok(create_agent_session_runtime(factory, ".".into(), ".".into(), session_manager).await?)
+    }
+
+    #[tokio::test]
+    async fn runtime_returns_session_and_cwd() -> TestResult {
+        let runtime = make_runtime().await?;
+        let session = runtime.session();
+        assert!(!session.session_id().await.is_empty());
+        assert_eq!(runtime.cwd(), ".");
+        assert_eq!(runtime.agent_dir(), ".");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn new_session_replaces_session_and_invokes_rebind() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let rebind_calls = Arc::new(AtomicUsize::new(0));
+        let rebind_calls_clone = Arc::clone(&rebind_calls);
+        runtime.set_rebind_session(Some(Arc::new(move |_session| {
+            let counter = Arc::clone(&rebind_calls_clone);
+            Box::pin(async move {
+                counter.fetch_add(1, Ordering::SeqCst);
+            })
+        })));
+
+        let first_session = runtime.session();
+        let outcome = runtime.new_session(NewSessionOptions::default()).await?;
+        assert!(!outcome.cancelled);
+        let second_session = runtime.session();
+        assert!(
+            !Arc::ptr_eq(&first_session, &second_session),
+            "session should have been replaced"
+        );
+        assert_eq!(rebind_calls.load(Ordering::SeqCst), 1);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn switch_session_to_new_path_succeeds() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let tmp = tempfile::tempdir()?;
+        let path = tmp.path().join("switch-target.jsonl");
+        let path_str = path.to_string_lossy().into_owned();
+        let outcome = runtime
+            .switch_session(&path_str, SwitchSessionOptions::default())
+            .await?;
+        assert!(!outcome.cancelled);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn fork_at_clones_branch_and_returns_no_selected_text() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let entry_id = {
+            let session = runtime.session();
+            let sm = session.session_manager();
+            let mut sm = sm.lock().await;
+            sm.append_message(&pi_agent::AgentMessage::Llm(Box::new(
+                pi_ai::Message::Assistant({
+                    let mut a = pi_ai::AssistantMessage::new(
+                        "test-api",
+                        "test-provider",
+                        "m",
+                        pi_agent::now_millis(),
+                    );
+                    a.stop_reason = pi_ai::StopReason::Stop;
+                    a
+                }),
+            )))?
+        };
+        let outcome = runtime.fork(&entry_id, ForkPosition::At).await?;
+        assert!(!outcome.cancelled);
+        assert!(outcome.selected_text.is_none());
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn fork_before_user_message_returns_selected_text() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let entry_id = {
+            let session = runtime.session();
+            let sm = session.session_manager();
+            let mut sm = sm.lock().await;
+            sm.append_message(&pi_agent::AgentMessage::Llm(Box::new(
+                pi_ai::Message::User(pi_ai::UserMessage::new(
+                    pi_ai::UserMessageContent::Text("hello world".into()),
+                    0,
+                )),
+            )))?
+        };
+        let outcome = runtime.fork(&entry_id, ForkPosition::Before).await?;
+        assert!(!outcome.cancelled);
+        assert_eq!(outcome.selected_text.as_deref(), Some("hello world"));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn fork_before_non_user_entry_errors() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let entry_id = {
+            let session = runtime.session();
+            let sm = session.session_manager();
+            let mut sm = sm.lock().await;
+            sm.append_message(&pi_agent::AgentMessage::Llm(Box::new(
+                pi_ai::Message::Assistant({
+                    let mut a = pi_ai::AssistantMessage::new(
+                        "test-api",
+                        "test-provider",
+                        "m",
+                        pi_agent::now_millis(),
+                    );
+                    a.stop_reason = pi_ai::StopReason::Stop;
+                    a
+                }),
+            )))?
+        };
+        let Err(err) = runtime.fork(&entry_id, ForkPosition::Before).await else {
+            return Err(failure("forking before a non-user entry must fail").into());
+        };
+        assert!(matches!(err, AgentSessionRuntimeError::InvalidForkEntry));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn fork_unknown_entry_errors() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let Err(err) = runtime.fork("missing", ForkPosition::At).await else {
+            return Err(failure("forking an unknown entry must fail").into());
+        };
+        assert!(matches!(err, AgentSessionRuntimeError::InvalidForkEntry));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn import_from_jsonl_missing_file_errors() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let Err(err) = runtime
+            .import_from_jsonl("/nonexistent/path.jsonl", None)
+            .await
+        else {
+            return Err(failure("importing a missing JSONL file must fail").into());
+        };
+        assert!(matches!(err, AgentSessionRuntimeError::ImportNotFound(_)));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn dispose_tears_down_session_without_replacing() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let session = runtime.session();
+        runtime.dispose().await;
+        assert!(Arc::ptr_eq(&runtime.session(), &session));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn rebind_callback_runs_after_apply_on_new_session() -> TestResult {
+        // Regression 2860: withSession must run on the NEW session.
+        let runtime = Arc::new(make_runtime().await?);
+        let bound_session_ids = Arc::new(std::sync::Mutex::new(Vec::new()));
+        let bound_ids_clone = Arc::clone(&bound_session_ids);
+        runtime.set_rebind_session(Some(Arc::new(move |session| {
+            let ids = Arc::clone(&bound_ids_clone);
+            Box::pin(async move {
+                let id = session.session_id().await;
+                if let Ok(mut ids) = ids.lock() {
+                    ids.push(id);
+                }
+            })
+        })));
+        runtime.new_session(NewSessionOptions::default()).await?;
+        let captured = bound_session_ids
+            .lock()
+            .map_err(|_| failure("bound session ID mutex poisoned"))?
+            .clone();
+        assert_eq!(captured.len(), 1, "rebind should fire once");
+        assert_eq!(captured[0], runtime.session().session_id().await);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn set_before_session_invalidate_invoked_during_teardown() -> TestResult {
+        let runtime = Arc::new(make_runtime().await?);
+        let called = Arc::new(AtomicUsize::new(0));
+        let called_clone = Arc::clone(&called);
+        runtime.set_before_session_invalidate(Some(Arc::new(move || {
+            called_clone.fetch_add(1, Ordering::SeqCst);
+        })));
+        runtime.new_session(NewSessionOptions::default()).await?;
+        assert_eq!(called.load(Ordering::SeqCst), 1);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn replacement_serialized_concurrent_new_sessions() -> TestResult {
+        let (entered_tx, mut entered_rx) = tokio::sync::mpsc::channel(2);
+        let gates = [
+            Arc::new(tokio::sync::Semaphore::new(0)),
+            Arc::new(tokio::sync::Semaphore::new(0)),
+        ];
+        let factory = Arc::new(GatedTestFactory::new(
+            entered_tx,
+            [Arc::clone(&gates[0]), Arc::clone(&gates[1])],
+        ));
+        let session_manager = SessionManager::in_memory(Some("."), None)?;
+        let runtime = Arc::new(
+            create_agent_session_runtime(factory, ".".into(), ".".into(), session_manager).await?,
+        );
+        let start = Arc::new(tokio::sync::Barrier::new(3));
+
+        let first_runtime = Arc::clone(&runtime);
+        let first_start = Arc::clone(&start);
+        let first = tokio::spawn(async move {
+            first_start.wait().await;
+            first_runtime
+                .new_session(NewSessionOptions::default())
+                .await
+        });
+        let second_runtime = Arc::clone(&runtime);
+        let second_start = Arc::clone(&start);
+        let second = tokio::spawn(async move {
+            second_start.wait().await;
+            second_runtime
+                .new_session(NewSessionOptions::default())
+                .await
+        });
+        start.wait().await;
+
+        let first_call = tokio::time::timeout(std::time::Duration::from_secs(1), entered_rx.recv())
+            .await
+            .map_err(|_| io::Error::other("timed out waiting for first replacement factory entry"))?
+            .ok_or_else(|| io::Error::other("replacement factory entry channel closed early"))?;
+        assert_eq!(first_call, 1);
+
+        match tokio::time::timeout(std::time::Duration::from_millis(100), entered_rx.recv()).await {
+            Ok(Some(call)) => {
+                return Err(io::Error::other(format!(
+                    "replacement factory call {call} entered before call {first_call} was released"
+                ))
+                .into());
+            }
+            Ok(None) => {
+                return Err(io::Error::other(
+                    "replacement factory entry channel closed while first gate was held",
+                )
+                .into());
+            }
+            Err(_) => {}
+        }
+
+        gates[first_call - 1].add_permits(1);
+        let second_call =
+            tokio::time::timeout(std::time::Duration::from_secs(1), entered_rx.recv())
+                .await
+                .map_err(|_| {
+                    io::Error::other("timed out waiting for second replacement factory entry")
+                })?
+                .ok_or_else(|| {
+                    io::Error::other("replacement factory entry channel closed early")
+                })?;
+        assert_eq!(second_call, 2);
+        gates[second_call - 1].add_permits(1);
+
+        let first_result = tokio::time::timeout(std::time::Duration::from_secs(1), first)
+            .await
+            .map_err(|_| io::Error::other("timed out joining first new-session task"))??;
+        let second_result = tokio::time::timeout(std::time::Duration::from_secs(1), second)
+            .await
+            .map_err(|_| io::Error::other("timed out joining second new-session task"))??;
+        first_result?;
+        second_result?;
+        Ok(())
+    }
+
+    async fn make_recording_runtime() -> TestResult<(
+        AgentSessionRuntime,
+        Arc<RecordingFactory>,
+        Arc<EmitRecordingRunner>,
+    )> {
+        let runner = Arc::new(EmitRecordingRunner::new());
+        let factory = Arc::new(RecordingFactory::new(Arc::clone(&runner)));
+        let session_manager = SessionManager::in_memory(Some("."), None)?;
+        let runtime = create_agent_session_runtime(
+            Arc::clone(&factory) as Arc<dyn CreateAgentSessionRuntimeFactory>,
+            ".".into(),
+            ".".into(),
+            session_manager,
+        )
+        .await?;
+        Ok((runtime, factory, runner))
+    }
+
+    #[tokio::test]
+    async fn new_session_passes_new_reason_and_emits_typed_shutdown() -> TestResult {
+        let (runtime, factory, runner) = make_recording_runtime().await?;
+        runtime.new_session(NewSessionOptions::default()).await?;
+        assert_eq!(
+            factory.reasons_clone(),
+            vec![SessionStartReason::Startup, SessionStartReason::New],
+            "replacement factory must receive start_reason = New"
+        );
+        let log = runner.log_clone();
+        assert!(
+            log.iter().any(|e| e == "session_shutdown:new:-"),
+            "old session must receive typed session_shutdown{{new}} (in-memory: no target), got {log:?}"
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn fork_passes_fork_reason_and_emits_typed_shutdown() -> TestResult {
+        let (runtime, factory, runner) = make_recording_runtime().await?;
+        let entry_id = {
+            let session = runtime.session();
+            let sm = session.session_manager();
+            let mut sm = sm.lock().await;
+            sm.append_message(&pi_agent::AgentMessage::Llm(Box::new(
+                pi_ai::Message::Assistant({
+                    let mut a = pi_ai::AssistantMessage::new(
+                        "test-api",
+                        "test-provider",
+                        "m",
+                        pi_agent::now_millis(),
+                    );
+                    a.stop_reason = pi_ai::StopReason::Stop;
+                    a
+                }),
+            )))?
+        };
+        runtime.fork(&entry_id, ForkPosition::At).await?;
+        assert_eq!(
+            factory.reasons_clone(),
+            vec![SessionStartReason::Startup, SessionStartReason::Fork],
+            "fork factory must receive start_reason = Fork"
+        );
+        let log = runner.log_clone();
+        assert!(
+            log.iter().any(|e| e == "session_shutdown:fork:-"),
+            "old session must receive typed session_shutdown{{fork}}, got {log:?}"
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn switch_session_emits_shutdown_with_target_session_file() -> TestResult {
+        let (runtime, factory, runner) = make_recording_runtime().await?;
+        let tmp = tempfile::tempdir()?;
+        let path = tmp.path().join("switch-target.jsonl");
+        let path_str = path.to_string_lossy().into_owned();
+        runtime
+            .switch_session(&path_str, SwitchSessionOptions::default())
+            .await?;
+        assert_eq!(
+            factory.reasons_clone(),
+            vec![SessionStartReason::Startup, SessionStartReason::Resume],
+        );
+        let expected = format!("session_shutdown:resume:{path_str}");
+        let log = runner.log_clone();
+        assert!(
+            log.contains(&expected),
+            "switch must carry the new session file as targetSessionFile: want {expected}, got {log:?}"
+        );
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn dispose_emits_quit_shutdown_without_target() -> TestResult {
+        let (runtime, _factory, runner) = make_recording_runtime().await?;
+        runtime.dispose().await;
+        let log = runner.log_clone();
+        assert!(
+            log.iter().any(|e| e == "session_shutdown:quit:-"),
+            "dispose must emit typed session_shutdown{{quit}} with no target, got {log:?}"
+        );
+        Ok(())
+    }
+}