pi-oxidized 0.1.1

package/crates/pi/src/modes/rpc/server.rs

@@ -0,0 +1,3037 @@
+//! RPC server: JSONL stdin → command dispatch → JSONL stdout.
+//!
+//! Ports `.references/pi/packages/coding-agent/src/modes/rpc/rpc-mode.ts`
+//! (`runRpcMode`, `handleCommand`, `handleInputLine`, `shutdown`).
+//!
+//! # Dispatch contract
+//!
+//! - **prompt** — spawns the session run; writes `{success:true}` exactly once
+//!   via the preflight callback (before events), returns `None`; failure error
+//!   only if preflight never fired.
+//! - **All other commands** — await the session call and return the exact
+//!   [`RpcResponse`].
+//! - **Unknown command** — `error(id, type, "Unknown command: {type}")`; the
+//!   `id` is echoed.
+//! - **Malformed JSON** — `error(None, "parse", …)`; the `id` is NOT echoed.
+//! - **Backpressure** — awaited after every non-prompt response and after
+//!   parse errors.
+//!
+//! # Shutdown
+//!
+//! - **SIGTERM** → exit 143 (no flush).
+//! - **SIGHUP** (unix) → exit 129 (flush).
+//! - **stdin EOF** → exit 0 (flush).
+//! - **Extension shutdown handler** → exit 0 (flush), checked after each
+//!   command and after `agent_settled`.
+
+use std::io;
+use std::pin::Pin;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::{Arc, Mutex};
+
+use futures::Future;
+use futures::future::BoxFuture;
+use pi_agent::{AgentMessage, QueueMode};
+use pi_ai::{ImageContent, Model, ModelThinkingLevel};
+use serde::Serialize;
+use serde_json::Value;
+use tokio::io::AsyncRead;
+use tokio::sync::{Notify, mpsc, oneshot};
+use tokio_util::sync::CancellationToken;
+
+use pi_ext::client::{HostUiRequest, HostUiResponse};
+use pi_ext::protocol::{ExtensionErrorEvent, NotifyLevel, SlotPlacement};
+
+use crate::core::agent_session::events::AgentSessionEvent;
+use crate::core::agent_session::extension::{ExtensionBindings, ExtensionMode};
+use crate::core::agent_session::prompt::PreflightCallback;
+use crate::core::agent_session_runtime::{ForkOutcome, ForkPosition};
+use crate::core::compaction::CompactionResult;
+use crate::core::extension_host::{ExtensionUiEvent, HostExtensionRunner};
+use crate::core::output_guard as output_guard_mod;
+use crate::core::sessions::SessionEntry;
+
+use super::extension_ui::ExtensionUiProxy;
+use super::jsonl::{JsonlLineReader, serialize_json_line};
+use super::types::{
+    BashResult, CycleModelData, CycleThinkingLevelData, ForkMessage, RpcCommand,
+    RpcExtensionUiRequest, RpcExtensionUiResponse, RpcResponse, RpcResponseData, RpcSessionState,
+    RpcSessionTreeNode, RpcSlashCommand, SessionStats, StreamingBehavior,
+};
+
+/// Serialize a value as a JSONL line, falling back to an empty string on
+/// serialization failure (our types are infallible serializers).
+fn to_jsonl<T: Serialize>(value: &T) -> String {
+    serialize_json_line(value).unwrap_or_default()
+}
+
+// ---------------------------------------------------------------------------
+// RpcSink — output abstraction
+// ---------------------------------------------------------------------------
+
+/// Future returned by [`RpcSink`] methods.
+type SinkFut = Pin<Box<dyn Future<Output = io::Result<()>> + Send>>;
+
+/// Protocol-stdout write/backpressure/flush abstraction.
+///
+/// All futures are `'static + Send` so they can be `tokio::spawn`'d from sync
+/// closures (preflight callbacks, event listeners).
+pub trait RpcSink: Send + Sync {
+    /// Write `text` (including trailing `\n`) to the ordered stdout sink.
+    fn write_stdout(&self, text: String) -> SinkFut;
+    /// Wait until every previously accepted write has finished draining.
+    fn backpressure(&self) -> SinkFut;
+    /// Wait for drain, then flush the underlying sink.
+    fn flush(&self) -> SinkFut;
+}
+
+/// [`OutputGuard`](crate::core::output_guard)-backed sink for production.
+#[derive(Clone, Copy, Debug)]
+pub struct OutputGuardSink;
+
+impl RpcSink for OutputGuardSink {
+    fn write_stdout(&self, text: String) -> SinkFut {
+        Box::pin(async move {
+            output_guard_mod::write_raw_stdout(text)
+                .await
+                .map_err(|e| io::Error::other(e.to_string()))
+        })
+    }
+    fn backpressure(&self) -> SinkFut {
+        Box::pin(async move {
+            output_guard_mod::wait_for_raw_stdout_backpressure()
+                .await
+                .map_err(|e| io::Error::other(e.to_string()))
+        })
+    }
+    fn flush(&self) -> SinkFut {
+        Box::pin(async move {
+            output_guard_mod::flush_raw_stdout()
+                .await
+                .map_err(|e| io::Error::other(e.to_string()))
+        })
+    }
+}
+
+/// In-memory sink for deterministic tests.
+#[derive(Clone, Default)]
+pub struct BufferSink {
+    stdout: Arc<Mutex<Vec<u8>>>,
+}
+
+impl BufferSink {
+    /// Create an empty buffer sink.
+    #[must_use]
+    pub fn new() -> Self {
+        Self::default()
+    }
+    /// Return the accumulated stdout bytes as a string.
+    #[must_use]
+    pub fn stdout_string(&self) -> String {
+        String::from_utf8_lossy(
+            &self
+                .stdout
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner),
+        )
+        .into_owned()
+    }
+    /// Return accumulated stdout split into trimmed lines.
+    #[must_use]
+    pub fn stdout_lines(&self) -> Vec<String> {
+        self.stdout_string()
+            .lines()
+            .filter(|l| !l.is_empty())
+            .map(str::to_owned)
+            .collect()
+    }
+}
+
+impl RpcSink for BufferSink {
+    fn write_stdout(&self, text: String) -> SinkFut {
+        let buf = Arc::clone(&self.stdout);
+        Box::pin(async move {
+            buf.lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .extend_from_slice(text.as_bytes());
+            Ok(())
+        })
+    }
+    fn backpressure(&self) -> SinkFut {
+        Box::pin(async { Ok(()) })
+    }
+    fn flush(&self) -> SinkFut {
+        Box::pin(async { Ok(()) })
+    }
+}
+
+#[cfg(test)]
+#[derive(Clone, Default)]
+struct GatedSink {
+    stdout: Arc<Mutex<Vec<u8>>>,
+    writes_started: Arc<std::sync::atomic::AtomicUsize>,
+    started: Arc<Notify>,
+    release: Arc<Notify>,
+}
+
+#[cfg(test)]
+impl GatedSink {
+    async fn wait_for_write(&self, target: usize) {
+        loop {
+            let started = self.started.notified();
+            if self.writes_started.load(Ordering::SeqCst) >= target {
+                return;
+            }
+            started.await;
+        }
+    }
+}
+
+#[cfg(test)]
+impl RpcSink for GatedSink {
+    fn write_stdout(&self, text: String) -> SinkFut {
+        let sink = self.clone();
+        Box::pin(async move {
+            sink.writes_started.fetch_add(1, Ordering::SeqCst);
+            sink.started.notify_waiters();
+            sink.release.notified().await;
+            sink.stdout
+                .lock()
+                .unwrap_or_else(std::sync::PoisonError::into_inner)
+                .extend_from_slice(text.as_bytes());
+            Ok(())
+        })
+    }
+
+    fn backpressure(&self) -> SinkFut {
+        Box::pin(async { Ok(()) })
+    }
+
+    fn flush(&self) -> SinkFut {
+        Box::pin(async { Ok(()) })
+    }
+}
+
+// ---------------------------------------------------------------------------
+// RpcSessionHost — session + runtime abstraction
+// ---------------------------------------------------------------------------
+
+/// Callback the host invokes after a session replacement.
+pub type RebindCallback = Arc<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>;
+
+/// Result of model cycling.
+#[derive(Clone, Debug, PartialEq)]
+pub struct ModelCycleResult {
+    /// Model now active.
+    pub model: Model,
+    /// Effective thinking level after clamping.
+    pub thinking_level: ModelThinkingLevel,
+    /// Whether cycling was across scoped entries.
+    pub is_scoped: bool,
+}
+
+/// Abstracts the `AgentSessionRuntime` + `AgentSession` surface consumed by the
+/// 31 [`RpcCommand`] variants.
+///
+/// All async methods return [`BoxFuture<'static>`] so they can be `tokio::spawn`'d
+/// without borrowing the host; the preflight callback fires synchronously
+/// during the first poll of `prompt`, which is how the prompt-response
+/// frame is enqueued before any agent event in the central write FIFO.
+pub trait RpcSessionHost: Send + Sync {
+    // ---- Prompt lifecycle ----
+    /// Submit a user prompt. The `preflight` callback fires with `true` on
+    /// accept (before the run starts) or `false` on reject.
+    fn prompt(
+        &self,
+        message: String,
+        images: Vec<ImageContent>,
+        streaming_behavior: Option<StreamingBehavior>,
+        preflight: PreflightCallback,
+    ) -> BoxFuture<'static, Result<(), String>>;
+    /// Steer into the active turn.
+    fn steer(
+        &self,
+        message: String,
+        images: Vec<ImageContent>,
+    ) -> BoxFuture<'static, Result<(), String>>;
+    /// Queue a follow-up after the current turn finishes.
+    fn follow_up(
+        &self,
+        message: String,
+        images: Vec<ImageContent>,
+    ) -> BoxFuture<'static, Result<(), String>>;
+    /// Abort the active agent run + retry, wait for idle.
+    fn abort(&self) -> BoxFuture<'static, ()>;
+
+    // ---- State ----
+    /// Snapshot the current session for the `get_state` RPC.
+    fn get_state(&self) -> BoxFuture<'static, RpcSessionState>;
+
+    // ---- Model ----
+    /// List available models from the model runtime.
+    fn get_available_models(&self) -> BoxFuture<'static, Vec<Model>>;
+    /// Set the active model.
+    fn set_model(&self, model: Model) -> BoxFuture<'static, Result<(), String>>;
+    /// Cycle to the next model (scoped or all-available).
+    fn cycle_model(&self) -> BoxFuture<'static, Option<ModelCycleResult>>;
+    /// Set the thinking level. Resolves `true` only when the level change was
+    /// durably committed (or was already effective).
+    fn set_thinking_level(&self, level: ModelThinkingLevel) -> BoxFuture<'static, bool>;
+    /// Cycle to the next thinking level.
+    fn cycle_thinking_level(&self) -> BoxFuture<'static, Option<ModelThinkingLevel>>;
+
+    // ---- Queue modes ----
+    /// Set steering queue drain mode.
+    fn set_steering_mode(&self, mode: QueueMode) -> BoxFuture<'static, ()>;
+    /// Set follow-up queue drain mode.
+    fn set_follow_up_mode(&self, mode: QueueMode) -> BoxFuture<'static, ()>;
+
+    // ---- Compaction ----
+    /// Compact the session.
+    fn compact(
+        &self,
+        custom_instructions: Option<String>,
+    ) -> BoxFuture<'static, Result<CompactionResult, String>>;
+    /// Enable / disable auto-compaction.
+    fn set_auto_compaction(&self, enabled: bool) -> BoxFuture<'static, ()>;
+
+    // ---- Retry ----
+    /// Enable / disable auto-retry.
+    fn set_auto_retry(&self, enabled: bool) -> BoxFuture<'static, ()>;
+    /// Abort in-flight retry sleep.
+    fn abort_retry(&self) -> BoxFuture<'static, ()>;
+
+    // ---- Bash ----
+    /// Execute a bash command.
+    fn execute_bash(
+        &self,
+        command: String,
+        exclude_from_context: Option<bool>,
+    ) -> BoxFuture<'static, Result<BashResult, String>>;
+    /// Abort a running bash command.
+    fn abort_bash(&self) -> BoxFuture<'static, ()>;
+
+    // ---- Session data ----
+    /// Aggregate session statistics for `get_session_stats`.
+    fn get_session_stats(&self) -> BoxFuture<'static, SessionStats>;
+    /// Export the session to HTML. Returns the written file path.
+    fn export_to_html(
+        &self,
+        output_path: Option<String>,
+    ) -> BoxFuture<'static, Result<String, String>>;
+    /// Set the session display name.
+    fn set_session_name(&self, name: String) -> BoxFuture<'static, Result<(), String>>;
+
+    // ---- Session mutations ----
+    /// Start a new session. Returns `true` when cancelled by an extension hook.
+    fn new_session(
+        &self,
+        parent_session: Option<String>,
+    ) -> BoxFuture<'static, Result<bool, String>>;
+    /// Switch to another session file. Returns `true` when cancelled.
+    fn switch_session(&self, session_path: String) -> BoxFuture<'static, Result<bool, String>>;
+    /// Fork the session. Returns the fork outcome.
+    fn fork(
+        &self,
+        entry_id: String,
+        position: ForkPosition,
+    ) -> BoxFuture<'static, Result<ForkOutcome, String>>;
+
+    // ---- Session tree / entries ----
+    /// All session entries.
+    fn get_entries(&self) -> BoxFuture<'static, Vec<SessionEntry>>;
+    /// Current leaf entry id.
+    fn get_leaf_id(&self) -> BoxFuture<'static, Option<String>>;
+    /// Session tree (wire-facing nodes).
+    fn get_tree(&self) -> BoxFuture<'static, Vec<RpcSessionTreeNode>>;
+    /// Forkable user messages.
+    fn get_fork_messages(&self) -> BoxFuture<'static, Vec<ForkMessage>>;
+    /// Last assistant text content, if any.
+    fn get_last_assistant_text(&self) -> BoxFuture<'static, Option<String>>;
+    /// All agent messages.
+    fn get_messages(&self) -> BoxFuture<'static, Vec<AgentMessage>>;
+    /// Available slash commands (extension + prompt + skill).
+    fn get_commands(&self) -> BoxFuture<'static, Vec<RpcSlashCommand>>;
+
+    // ---- Subscribe / extensions ----
+    /// Subscribe to public session events. Returns an unsubscribe closure.
+    fn subscribe(
+        &self,
+        listener: Arc<dyn Fn(&AgentSessionEvent) + Send + Sync>,
+    ) -> Box<dyn Fn() + Send + Sync>;
+    /// Register a backpressure hook invoked when the agent has pending events.
+    fn register_backpressure_hook(
+        &self,
+        hook: Arc<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>,
+    ) -> Box<dyn Fn() + Send + Sync>;
+    /// Bind extensions for RPC mode.
+    fn bind_extensions_rpc(
+        &self,
+        bindings: ExtensionBindings,
+    ) -> BoxFuture<'static, Result<(), String>>;
+    /// Current concrete extension host, when this session uses one.
+    fn host_extension_runner(&self) -> Option<Arc<HostExtensionRunner>> {
+        None
+    }
+
+    // ---- Lifecycle ----
+    /// Dispose the current session and runtime.
+    fn dispose(&self) -> BoxFuture<'static, ()>;
+    /// Set the rebind callback invoked after session replacement.
+    fn set_rebind(&self, callback: Option<RebindCallback>);
+}
+
+// ---------------------------------------------------------------------------
+// ServerOutput types
+// ---------------------------------------------------------------------------
+
+/// Extension error notification (`type: "extension_error"`).
+#[derive(Clone, Debug, Serialize)]
+pub struct ExtensionErrorOutput {
+    #[serde(rename = "type")]
+    type_name: &'static str,
+    #[serde(rename = "extensionPath")]
+    extension_path: String,
+    event: String,
+    error: String,
+}
+
+impl ExtensionErrorOutput {
+    /// Create a new extension error output.
+    #[must_use]
+    pub fn new(
+        path: impl Into<String>,
+        event: impl Into<String>,
+        error: impl Into<String>,
+    ) -> Self {
+        Self {
+            type_name: "extension_error",
+            extension_path: path.into(),
+            event: event.into(),
+            error: error.into(),
+        }
+    }
+}
+
+/// Tagged stdout frame for test deserialization.
+#[derive(Debug, Serialize)]
+#[serde(untagged)]
+pub enum ServerOutput {
+    /// Command response.
+    Response(RpcResponse),
+    /// Raw agent event.
+    Event(Box<AgentSessionEvent>),
+    /// Extension UI request.
+    UiRequest(RpcExtensionUiRequest),
+    /// Extension error notification.
+    ExtensionError(ExtensionErrorOutput),
+}
+
+// ---------------------------------------------------------------------------
+// ServerState
+// ---------------------------------------------------------------------------
+
+type UnsubSlot = Mutex<Option<Box<dyn Fn() + Send + Sync>>>;
+
+fn lock_unsub(slot: &UnsubSlot) -> std::sync::MutexGuard<'_, Option<Box<dyn Fn() + Send + Sync>>> {
+    slot.lock()
+        .unwrap_or_else(std::sync::PoisonError::into_inner)
+}
+
+fn take_unsub(slot: &UnsubSlot) {
+    if let Some(unsub) = lock_unsub(slot).take() {
+        unsub();
+    }
+}
+
+/// Mutable server state shared between the event loop, event subscriber, and
+/// rebind callback.
+///
+/// All stdout writes go through `write_tx`. A dedicated writer actor consumes
+/// line and drain-barrier messages in FIFO order, preserving prompt-response
+/// ordering while allowing producers to await the actual server queue.
+enum WriteMessage {
+    Line(String),
+    Drain(oneshot::Sender<()>),
+}
+
+pub(crate) struct ServerState {
+    sink: Arc<dyn RpcSink>,
+    write_tx: mpsc::UnboundedSender<WriteMessage>,
+    proxy: ExtensionUiProxy,
+    shutdown_requested: Arc<AtomicBool>,
+    needs_rebind: Arc<AtomicBool>,
+    signal: Arc<Notify>,
+    unsubscribe_events: UnsubSlot,
+    unsubscribe_backpressure: UnsubSlot,
+    unsubscribe_extension_ui: UnsubSlot,
+}
+
+impl ServerState {
+    fn new(
+        sink: Arc<dyn RpcSink>,
+        write_tx: mpsc::UnboundedSender<WriteMessage>,
+        proxy: ExtensionUiProxy,
+    ) -> Self {
+        Self {
+            sink,
+            write_tx,
+            proxy,
+            shutdown_requested: Arc::new(AtomicBool::new(false)),
+            needs_rebind: Arc::new(AtomicBool::new(false)),
+            signal: Arc::new(Notify::new()),
+            unsubscribe_events: Mutex::new(None),
+            unsubscribe_backpressure: Mutex::new(None),
+            unsubscribe_extension_ui: Mutex::new(None),
+        }
+    }
+
+    /// Enqueue a JSONL line for ordered stdout emission (synchronous, FIFO).
+    fn enqueue(&self, line: String) {
+        let _ = self.write_tx.send(WriteMessage::Line(line));
+    }
+
+    async fn wait_for_output(&self) {
+        let (done_tx, done_rx) = oneshot::channel();
+        if self.write_tx.send(WriteMessage::Drain(done_tx)).is_ok() {
+            let _ = done_rx.await;
+        }
+    }
+
+    /// Rebind extensions + subscriptions on the current session.
+    ///
+    /// All stdout emission (events, extension errors) goes through `write_tx`
+    /// so the event loop drains them in FIFO order.
+    pub(crate) async fn rebind<H>(&self, host: &H)
+    where
+        H: RpcSessionHost + ?Sized,
+    {
+        take_unsub(&self.unsubscribe_events);
+        take_unsub(&self.unsubscribe_backpressure);
+        take_unsub(&self.unsubscribe_extension_ui);
+
+        let shutdown_flag = Arc::clone(&self.shutdown_requested);
+        let shutdown_signal = Arc::clone(&self.signal);
+        let error_tx = self.write_tx.clone();
+        let bindings = ExtensionBindings {
+            mode: Some(ExtensionMode::Rpc),
+            shutdown_handler: Some(Arc::new(move || {
+                shutdown_flag.store(true, Ordering::SeqCst);
+                shutdown_signal.notify_one();
+            })),
+            on_error: Some(Arc::new(move |path: &str, event: &str, error: &str| {
+                let output = ExtensionErrorOutput::new(path, event, error);
+                let _ = error_tx.send(WriteMessage::Line(to_jsonl(&output)));
+            })),
+            ..Default::default()
+        };
+        let _ = host.bind_extensions_rpc(bindings).await;
+
+        let event_tx = self.write_tx.clone();
+        let signal = Arc::clone(&self.signal);
+        let unsub = host.subscribe(Arc::new(move |event: &AgentSessionEvent| {
+            let _ = event_tx.send(WriteMessage::Line(to_jsonl(event)));
+            if matches!(event, AgentSessionEvent::AgentSettled) {
+                signal.notify_one();
+            }
+        }));
+        *lock_unsub(&self.unsubscribe_events) = Some(unsub);
+
+        let backpressure_tx = self.write_tx.clone();
+        let unsub_bp = host.register_backpressure_hook(Arc::new(move || {
+            let write_tx = backpressure_tx.clone();
+            Box::pin(async move {
+                let (done_tx, done_rx) = oneshot::channel();
+                if write_tx.send(WriteMessage::Drain(done_tx)).is_ok() {
+                    let _ = done_rx.await;
+                }
+            })
+        }));
+        *lock_unsub(&self.unsubscribe_backpressure) = Some(unsub_bp);
+
+        if let Some(runner) = host.host_extension_runner() {
+            let cancel = CancellationToken::new();
+            let cancel_on_unbind = cancel.clone();
+            *lock_unsub(&self.unsubscribe_extension_ui) = Some(Box::new(move || {
+                cancel_on_unbind.cancel();
+            }));
+
+            if let Some(requests) = runner.take_ui_requests() {
+                tokio::spawn(run_extension_dialog_bridge(
+                    Arc::clone(&runner),
+                    requests,
+                    self.proxy.clone(),
+                    self.write_tx.clone(),
+                    cancel.clone(),
+                ));
+            }
+            tokio::spawn(run_extension_event_bridge(
+                runner,
+                self.write_tx.clone(),
+                cancel,
+            ));
+        }
+    }
+
+    /// Cleanup before exit: cancel UI, unsubscribe, dispose, flush.
+    async fn cleanup<H>(&self, host: &H, exit_code: i32)
+    where
+        H: RpcSessionHost + ?Sized,
+    {
+        self.proxy.cancel_all();
+        take_unsub(&self.unsubscribe_events);
+        take_unsub(&self.unsubscribe_backpressure);
+        take_unsub(&self.unsubscribe_extension_ui);
+        host.dispose().await;
+        self.wait_for_output().await;
+        if exit_code != 143 {
+            let _ = self.sink.flush().await;
+        }
+    }
+}
+
+async fn run_extension_dialog_bridge(
+    runner: Arc<HostExtensionRunner>,
+    mut requests: mpsc::Receiver<HostUiRequest>,
+    proxy: ExtensionUiProxy,
+    write_tx: mpsc::UnboundedSender<WriteMessage>,
+    cancel: CancellationToken,
+) {
+    loop {
+        let request = tokio::select! {
+            () = cancel.cancelled() => break,
+            request = requests.recv() => match request {
+                Some(request) => request,
+                None => break,
+            },
+        };
+        let runner = Arc::clone(&runner);
+        let proxy = proxy.clone();
+        let write_tx = write_tx.clone();
+        let request_cancel = cancel.child_token();
+        tokio::spawn(async move {
+            bridge_extension_dialog(runner, request, proxy, write_tx, request_cancel).await;
+        });
+    }
+}
+
+async fn bridge_extension_dialog(
+    runner: Arc<HostExtensionRunner>,
+    request: HostUiRequest,
+    proxy: ExtensionUiProxy,
+    write_tx: mpsc::UnboundedSender<WriteMessage>,
+    cancel: CancellationToken,
+) {
+    let timeout_ms = match &request {
+        HostUiRequest::Select { request, .. } => request.options_meta.timeout_ms,
+        HostUiRequest::Confirm { request, .. } => request.options_meta.timeout_ms,
+        HostUiRequest::Input { request, .. } => request.options_meta.timeout_ms,
+        HostUiRequest::Editor { .. } => None,
+    };
+    let (rpc_request, response_rx) = proxy.create_dialog(|id| match &request {
+        HostUiRequest::Select { request, .. } => RpcExtensionUiRequest::Select {
+            id: id.to_owned(),
+            title: request.title.clone(),
+            options: request.options.clone(),
+            timeout: request.options_meta.timeout_ms,
+        },
+        HostUiRequest::Confirm { request, .. } => RpcExtensionUiRequest::Confirm {
+            id: id.to_owned(),
+            title: request.title.clone(),
+            message: request.message.clone(),
+            timeout: request.options_meta.timeout_ms,
+        },
+        HostUiRequest::Input { request, .. } => RpcExtensionUiRequest::Input {
+            id: id.to_owned(),
+            title: request.title.clone(),
+            placeholder: request.placeholder.clone(),
+            timeout: request.options_meta.timeout_ms,
+        },
+        HostUiRequest::Editor { request, .. } => RpcExtensionUiRequest::Editor {
+            id: id.to_owned(),
+            title: request.title.clone(),
+            prefill: request.prefill.clone(),
+        },
+    });
+    let rpc_id = rpc_request.id().to_owned();
+    let _ = write_tx.send(WriteMessage::Line(to_jsonl(&ServerOutput::UiRequest(
+        rpc_request,
+    ))));
+
+    let response = if let Some(timeout_ms) = timeout_ms {
+        tokio::select! {
+            () = cancel.cancelled() => None,
+            result = tokio::time::timeout(
+                std::time::Duration::from_millis(timeout_ms),
+                response_rx,
+            ) => result.ok().and_then(Result::ok),
+        }
+    } else {
+        tokio::select! {
+            () = cancel.cancelled() => None,
+            result = response_rx => result.ok(),
+        }
+    };
+    if response.is_none() {
+        let _ = proxy.route_response(RpcExtensionUiResponse::Cancelled { id: rpc_id });
+    }
+    let host_response = map_rpc_ui_response(&request, response);
+    let _ = runner.respond_ui(host_response).await;
+}
+
+fn map_rpc_ui_response(
+    request: &HostUiRequest,
+    response: Option<RpcExtensionUiResponse>,
+) -> HostUiResponse {
+    match request {
+        HostUiRequest::Select { id, .. } => HostUiResponse::Select {
+            id: *id,
+            value: match response {
+                Some(RpcExtensionUiResponse::Value { value, .. }) => Some(value),
+                _ => None,
+            },
+        },
+        HostUiRequest::Confirm { id, .. } => HostUiResponse::Confirm {
+            id: *id,
+            confirmed: match response {
+                Some(RpcExtensionUiResponse::Confirmed { confirmed, .. }) => confirmed,
+                _ => false,
+            },
+        },
+        HostUiRequest::Input { id, .. } => HostUiResponse::Input {
+            id: *id,
+            value: match response {
+                Some(RpcExtensionUiResponse::Value { value, .. }) => Some(value),
+                _ => None,
+            },
+        },
+        HostUiRequest::Editor { id, .. } => HostUiResponse::Editor {
+            id: *id,
+            value: match response {
+                Some(RpcExtensionUiResponse::Value { value, .. }) => Some(value),
+                _ => None,
+            },
+        },
+    }
+}
+
+async fn run_extension_event_bridge(
+    runner: Arc<HostExtensionRunner>,
+    write_tx: mpsc::UnboundedSender<WriteMessage>,
+    cancel: CancellationToken,
+) {
+    let mut events = runner.subscribe_ui();
+    let mut errors = runner.subscribe_errors();
+    for slot in runner.current_slots() {
+        let request = map_extension_ui_event(ExtensionUiEvent::Slot(slot));
+        let _ = write_tx.send(WriteMessage::Line(to_jsonl(&ServerOutput::UiRequest(
+            request,
+        ))));
+    }
+    loop {
+        tokio::select! {
+            () = cancel.cancelled() => break,
+            event = events.recv() => {
+                let event = match event {
+                    Ok(event) => event,
+                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
+                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                };
+                let request = map_extension_ui_event(event);
+                let _ = write_tx.send(WriteMessage::Line(to_jsonl(
+                    &ServerOutput::UiRequest(request),
+                )));
+            }
+            error = errors.recv() => {
+                let error = match error {
+                    Ok(error) => error,
+                    Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => continue,
+                    Err(tokio::sync::broadcast::error::RecvError::Closed) => break,
+                };
+                let output = map_extension_error_event(&error);
+                let _ = write_tx.send(WriteMessage::Line(to_jsonl(
+                    &ServerOutput::ExtensionError(output),
+                )));
+            }
+        }
+    }
+}
+
+fn map_extension_error_event(error: &ExtensionErrorEvent) -> ExtensionErrorOutput {
+    let mut extension_path = "<runtime>".to_owned();
+    let mut event = error.code.clone();
+    let mut message = error.message.clone();
+    let mut structured = false;
+
+    if let Some(data) = error.data.as_ref() {
+        if let Some(value) = data
+            .get("extensionPath")
+            .or_else(|| data.get("path"))
+            .and_then(Value::as_str)
+        {
+            value.clone_into(&mut extension_path);
+            structured = true;
+        }
+        if let Some(value) = data.get("event").and_then(Value::as_str) {
+            value.clone_into(&mut event);
+            structured = true;
+        }
+        if let Some(value) = data.get("error").and_then(Value::as_str) {
+            value.clone_into(&mut message);
+            structured = true;
+        }
+    }
+
+    if !structured
+        && let Some(rest) = error.message.strip_prefix('[')
+        && let Some((path, rest)) = rest.split_once("] ")
+        && let Some((event_name, error_message)) = rest.split_once(": ")
+        && !path.is_empty()
+        && !event_name.is_empty()
+    {
+        path.clone_into(&mut extension_path);
+        event_name.clone_into(&mut event);
+        error_message.clone_into(&mut message);
+    }
+
+    ExtensionErrorOutput::new(extension_path, event, message)
+}
+
+fn map_extension_ui_event(event: ExtensionUiEvent) -> RpcExtensionUiRequest {
+    match event {
+        ExtensionUiEvent::Notify(notification) => ExtensionUiProxy::notify(
+            &notification.message,
+            Some(match notification.level {
+                NotifyLevel::Info => super::types::NotifyType::Info,
+                NotifyLevel::Warning => super::types::NotifyType::Warning,
+                NotifyLevel::Error => super::types::NotifyType::Error,
+            }),
+        ),
+        ExtensionUiEvent::Slot(slot) => {
+            let lines = slot
+                .lines
+                .iter()
+                .map(|line| line.iter().map(|run| run.text.as_str()).collect::<String>())
+                .collect::<Vec<_>>();
+            let placement = match slot.placement {
+                SlotPlacement::Footer | SlotPlacement::BelowEditor => {
+                    Some(super::types::WidgetPlacement::BelowEditor)
+                }
+                _ => Some(super::types::WidgetPlacement::AboveEditor),
+            };
+            ExtensionUiProxy::set_widget(&slot.key, Some(&lines), placement)
+        }
+        ExtensionUiEvent::Dispose { key } => ExtensionUiProxy::set_widget(&key, None, None),
+    }
+}
+
+// ---------------------------------------------------------------------------
+// handle_command — dispatch one RpcCommand
+// ---------------------------------------------------------------------------
+
+/// Dispatch a single [`RpcCommand`].
+///
+/// Returns `Some(response)` for synchronous commands. Returns `None` for
+/// `prompt` (response emitted asynchronously via preflight).
+#[allow(clippy::too_many_lines)]
+pub(crate) async fn handle_command<H>(
+    command: &RpcCommand,
+    host: &H,
+    state: &ServerState,
+) -> Option<RpcResponse>
+where
+    H: RpcSessionHost + ?Sized,
+{
+    let id = command.id().map(str::to_owned);
+
+    match command {
+        RpcCommand::Prompt {
+            message,
+            images,
+            streaming_behavior,
+            ..
+        } => {
+            // Wait until preflight enqueues the prompt response before returning,
+            // so a subsequent command's response cannot overtake it in the FIFO.
+            spawn_prompt(
+                id,
+                message.clone(),
+                images.clone().unwrap_or_default(),
+                *streaming_behavior,
+                host,
+                state.write_tx.clone(),
+            )
+            .await;
+            None
+        }
+
+        RpcCommand::Steer {
+            message, images, ..
+        } => match host
+            .steer(message.clone(), images.clone().unwrap_or_default())
+            .await
+        {
+            Ok(()) => Some(RpcResponse::ok(id, "steer")),
+            Err(e) => Some(RpcResponse::err(id, "steer", e)),
+        },
+
+        RpcCommand::FollowUp {
+            message, images, ..
+        } => match host
+            .follow_up(message.clone(), images.clone().unwrap_or_default())
+            .await
+        {
+            Ok(()) => Some(RpcResponse::ok(id, "follow_up")),
+            Err(e) => Some(RpcResponse::err(id, "follow_up", e)),
+        },
+
+        RpcCommand::Abort { .. } => {
+            host.abort().await;
+            Some(RpcResponse::ok(id, "abort"))
+        }
+
+        RpcCommand::NewSession { parent_session, .. } => {
+            match host.new_session(parent_session.clone()).await {
+                Ok(cancelled) => {
+                    if !cancelled {
+                        state.rebind(host).await;
+                    }
+                    Some(RpcResponse::ok_data(
+                        id,
+                        "new_session",
+                        RpcResponseData::Cancelled { cancelled },
+                    ))
+                }
+                Err(e) => Some(RpcResponse::err(id, "new_session", e)),
+            }
+        }
+
+        RpcCommand::GetState { .. } => {
+            let s = host.get_state().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_state",
+                RpcResponseData::SessionState(s),
+            ))
+        }
+
+        RpcCommand::SetModel {
+            provider, model_id, ..
+        } => {
+            let models = host.get_available_models().await;
+            match models
+                .into_iter()
+                .find(|m| m.provider == *provider && m.id == *model_id)
+            {
+                Some(model) => match host.set_model(model.clone()).await {
+                    Ok(()) => Some(RpcResponse::ok_data(
+                        id,
+                        "set_model",
+                        RpcResponseData::Model(model),
+                    )),
+                    Err(e) => Some(RpcResponse::err(id, "set_model", e)),
+                },
+                None => Some(RpcResponse::err(
+                    id,
+                    "set_model",
+                    format!("Model not found: {provider}/{model_id}"),
+                )),
+            }
+        }
+
+        RpcCommand::CycleModel { .. } => {
+            let r = host.cycle_model().await;
+            let data = r.map(|x| {
+                RpcResponseData::CycleModel(Some(CycleModelData {
+                    model: x.model,
+                    thinking_level: x.thinking_level,
+                    is_scoped: x.is_scoped,
+                }))
+            });
+            Some(RpcResponse::ok_data(
+                id,
+                "cycle_model",
+                data.unwrap_or(RpcResponseData::CycleModel(None)),
+            ))
+        }
+
+        RpcCommand::GetAvailableModels { .. } => {
+            let models = host.get_available_models().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_available_models",
+                RpcResponseData::AvailableModels { models },
+            ))
+        }
+
+        RpcCommand::SetThinkingLevel { level, .. } => {
+            if host.set_thinking_level(*level).await {
+                Some(RpcResponse::ok(id, "set_thinking_level"))
+            } else {
+                Some(RpcResponse::err(
+                    id,
+                    "set_thinking_level",
+                    "Failed to persist thinking level change",
+                ))
+            }
+        }
+
+        RpcCommand::CycleThinkingLevel { .. } => {
+            let level = host.cycle_thinking_level().await;
+            let data = level.map(|l| {
+                RpcResponseData::CycleThinkingLevel(Some(CycleThinkingLevelData { level: l }))
+            });
+            Some(RpcResponse::ok_data(
+                id,
+                "cycle_thinking_level",
+                data.unwrap_or(RpcResponseData::CycleThinkingLevel(None)),
+            ))
+        }
+
+        RpcCommand::SetSteeringMode { mode, .. } => {
+            host.set_steering_mode(*mode).await;
+            Some(RpcResponse::ok(id, "set_steering_mode"))
+        }
+
+        RpcCommand::SetFollowUpMode { mode, .. } => {
+            host.set_follow_up_mode(*mode).await;
+            Some(RpcResponse::ok(id, "set_follow_up_mode"))
+        }
+
+        RpcCommand::Compact {
+            custom_instructions,
+            ..
+        } => match host.compact(custom_instructions.clone()).await {
+            Ok(result) => Some(RpcResponse::ok_data(
+                id,
+                "compact",
+                RpcResponseData::Compaction(result),
+            )),
+            Err(e) => Some(RpcResponse::err(id, "compact", e)),
+        },
+
+        RpcCommand::SetAutoCompaction { enabled, .. } => {
+            host.set_auto_compaction(*enabled).await;
+            Some(RpcResponse::ok(id, "set_auto_compaction"))
+        }
+
+        RpcCommand::SetAutoRetry { enabled, .. } => {
+            host.set_auto_retry(*enabled).await;
+            Some(RpcResponse::ok(id, "set_auto_retry"))
+        }
+
+        RpcCommand::AbortRetry { .. } => {
+            host.abort_retry().await;
+            Some(RpcResponse::ok(id, "abort_retry"))
+        }
+
+        RpcCommand::Bash {
+            command: cmd,
+            exclude_from_context,
+            ..
+        } => match host.execute_bash(cmd.clone(), *exclude_from_context).await {
+            Ok(result) => Some(RpcResponse::ok_data(
+                id,
+                "bash",
+                RpcResponseData::Bash(result),
+            )),
+            Err(e) => Some(RpcResponse::err(id, "bash", e)),
+        },
+
+        RpcCommand::AbortBash { .. } => {
+            host.abort_bash().await;
+            Some(RpcResponse::ok(id, "abort_bash"))
+        }
+
+        RpcCommand::GetSessionStats { .. } => {
+            let session_stats = host.get_session_stats().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_session_stats",
+                RpcResponseData::SessionStats(session_stats),
+            ))
+        }
+
+        RpcCommand::ExportHtml { output_path, .. } => {
+            match host.export_to_html(output_path.clone()).await {
+                Ok(path) => Some(RpcResponse::ok_data(
+                    id,
+                    "export_html",
+                    RpcResponseData::ExportHtml { path },
+                )),
+                Err(e) => Some(RpcResponse::err(id, "export_html", e)),
+            }
+        }
+
+        RpcCommand::SwitchSession { session_path, .. } => {
+            match host.switch_session(session_path.clone()).await {
+                Ok(cancelled) => {
+                    if !cancelled {
+                        state.rebind(host).await;
+                    }
+                    Some(RpcResponse::ok_data(
+                        id,
+                        "switch_session",
+                        RpcResponseData::Cancelled { cancelled },
+                    ))
+                }
+                Err(e) => Some(RpcResponse::err(id, "switch_session", e)),
+            }
+        }
+
+        RpcCommand::Fork { entry_id, .. } => {
+            match host.fork(entry_id.clone(), ForkPosition::Before).await {
+                Ok(outcome) => {
+                    if !outcome.cancelled {
+                        state.rebind(host).await;
+                    }
+                    Some(RpcResponse::ok_data(
+                        id,
+                        "fork",
+                        RpcResponseData::Fork {
+                            text: outcome.selected_text.unwrap_or_default(),
+                            cancelled: outcome.cancelled,
+                        },
+                    ))
+                }
+                Err(e) => Some(RpcResponse::err(id, "fork", e)),
+            }
+        }
+
+        RpcCommand::Clone { .. } => {
+            let leaf = host.get_leaf_id().await;
+            match leaf {
+                None => Some(RpcResponse::err(
+                    id,
+                    "clone",
+                    "Cannot clone session: no current entry selected",
+                )),
+                Some(leaf_id) => match host.fork(leaf_id, ForkPosition::At).await {
+                    Ok(outcome) => {
+                        if !outcome.cancelled {
+                            state.rebind(host).await;
+                        }
+                        Some(RpcResponse::ok_data(
+                            id,
+                            "clone",
+                            RpcResponseData::Cancelled {
+                                cancelled: outcome.cancelled,
+                            },
+                        ))
+                    }
+                    Err(e) => Some(RpcResponse::err(id, "clone", e)),
+                },
+            }
+        }
+
+        RpcCommand::GetForkMessages { .. } => {
+            let messages = host.get_fork_messages().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_fork_messages",
+                RpcResponseData::ForkMessages { messages },
+            ))
+        }
+
+        RpcCommand::GetEntries { since, .. } => {
+            let entries = host.get_entries().await;
+            let filtered = if let Some(since_id) = since {
+                match entries
+                    .iter()
+                    .position(|e| e.id() == Some(since_id.as_str()))
+                {
+                    None => {
+                        return Some(RpcResponse::err(
+                            id,
+                            "get_entries",
+                            format!("Entry not found: {since_id}"),
+                        ));
+                    }
+                    Some(i) => entries.into_iter().skip(i + 1).collect::<Vec<_>>(),
+                }
+            } else {
+                entries
+            };
+            let leaf_id = host.get_leaf_id().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_entries",
+                RpcResponseData::Entries {
+                    entries: filtered,
+                    leaf_id,
+                },
+            ))
+        }
+
+        RpcCommand::GetTree { .. } => {
+            let tree = host.get_tree().await;
+            let leaf_id = host.get_leaf_id().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_tree",
+                RpcResponseData::Tree { tree, leaf_id },
+            ))
+        }
+
+        RpcCommand::GetLastAssistantText { .. } => {
+            let text = host.get_last_assistant_text().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_last_assistant_text",
+                RpcResponseData::LastAssistantText { text },
+            ))
+        }
+
+        RpcCommand::SetSessionName { name, .. } => {
+            let trimmed = name.trim();
+            if trimmed.is_empty() {
+                return Some(RpcResponse::err(
+                    id,
+                    "set_session_name",
+                    "Session name cannot be empty",
+                ));
+            }
+            match host.set_session_name(trimmed.to_owned()).await {
+                Ok(()) => Some(RpcResponse::ok(id, "set_session_name")),
+                Err(e) => Some(RpcResponse::err(id, "set_session_name", e)),
+            }
+        }
+
+        RpcCommand::GetMessages { .. } => {
+            let messages = host.get_messages().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_messages",
+                RpcResponseData::Messages { messages },
+            ))
+        }
+
+        RpcCommand::GetCommands { .. } => {
+            let commands = host.get_commands().await;
+            Some(RpcResponse::ok_data(
+                id,
+                "get_commands",
+                RpcResponseData::Commands { commands },
+            ))
+        }
+
+        RpcCommand::Unknown {
+            command_type: ct, ..
+        } => Some(RpcResponse::err(id, ct, format!("Unknown command: {ct}"))),
+    }
+}
+
+/// Spawn the prompt run. Emits exactly one success on first preflight `true`,
+/// or one error if preflight never fired.
+///
+/// The returned future resolves once the prompt future has either:
+/// - fired preflight (success response already enqueued), or
+/// - failed without preflight (error response enqueued).
+///
+/// Waiting here keeps subsequent command responses behind the prompt response
+/// in the central write FIFO. Agent events still enqueue through the same
+/// channel after preflight, preserving response-before-events ordering.
+async fn spawn_prompt<H>(
+    id: Option<String>,
+    message: String,
+    images: Vec<ImageContent>,
+    streaming_behavior: Option<StreamingBehavior>,
+    host: &H,
+    write_tx: mpsc::UnboundedSender<WriteMessage>,
+) where
+    H: RpcSessionHost + ?Sized,
+{
+    let preflight_succeeded = Arc::new(AtomicBool::new(false));
+    let (preflight_tx, preflight_rx) = tokio::sync::oneshot::channel::<()>();
+    let preflight_notify = Arc::new(Mutex::new(Some(preflight_tx)));
+
+    let success_flag = Arc::clone(&preflight_succeeded);
+    let success_tx = write_tx.clone();
+    let success_id = id.clone();
+    let notify_on_preflight = Arc::clone(&preflight_notify);
+    let preflight: PreflightCallback = Arc::new(move |did_succeed: bool| {
+        if did_succeed
+            && success_flag
+                .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
+                .is_ok()
+        {
+            let response = RpcResponse::ok(success_id.clone(), "prompt");
+            let _ = success_tx.send(WriteMessage::Line(to_jsonl(&response)));
+        }
+        // Unblock the dispatcher after preflight so later commands cannot
+        // enqueue ahead of the prompt response.
+        if let Some(tx) = notify_on_preflight
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .take()
+        {
+            let _ = tx.send(());
+        }
+    });
+
+    let prompt_future = host.prompt(message, images, streaming_behavior, preflight);
+
+    let error_flag = Arc::clone(&preflight_succeeded);
+    let error_tx = write_tx;
+    let error_id = id;
+    let notify_on_error = Arc::clone(&preflight_notify);
+
+    tokio::spawn(async move {
+        if let Err(msg) = prompt_future.await
+            && !error_flag.load(Ordering::SeqCst)
+        {
+            let response = RpcResponse::err(error_id, "prompt", msg);
+            let _ = error_tx.send(WriteMessage::Line(to_jsonl(&response)));
+        }
+        // If preflight never fired (host bug / panic path), still release the
+        // dispatcher so the command loop cannot hang forever.
+        if let Some(tx) = notify_on_error
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner)
+            .take()
+        {
+            let _ = tx.send(());
+        }
+    });
+
+    let _ = preflight_rx.await;
+}
+
+// ---------------------------------------------------------------------------
+// Input line processing
+// ---------------------------------------------------------------------------
+
+/// Outcome of processing one input line.
+#[derive(Debug, PartialEq, Eq)]
+pub enum LineOutcome {
+    /// Command handled (or UI response routed).
+    Done,
+    /// Shutdown requested — stop reading.
+    Shutdown,
+}
+
+/// Parse and process one JSONL input line.
+pub(crate) async fn process_input_line<H>(line: &str, host: &H, state: &ServerState) -> LineOutcome
+where
+    H: RpcSessionHost + ?Sized,
+{
+    let parsed: Value = match serde_json::from_str(line) {
+        Ok(v) => v,
+        Err(e) => {
+            let response = RpcResponse::err(None, "parse", format!("Failed to parse command: {e}"));
+            state.enqueue(to_jsonl(&response));
+            return LineOutcome::Done;
+        }
+    };
+
+    // Route extension UI responses before command dispatch.
+    if parsed.get("type").and_then(Value::as_str) == Some("extension_ui_response") {
+        if let Ok(ui_resp) = serde_json::from_value::<RpcExtensionUiResponse>(parsed.clone()) {
+            let _ = state.proxy.route_response(ui_resp);
+        }
+        return LineOutcome::Done;
+    }
+
+    let command = match RpcCommand::parse_value(&parsed) {
+        Ok(command) => command,
+        Err(error) => {
+            let response = RpcResponse::err(
+                error.id,
+                "parse",
+                format!("Failed to parse command: {}", error.message),
+            );
+            state.enqueue(to_jsonl(&response));
+            return LineOutcome::Done;
+        }
+    };
+
+    if let Some(resp) = handle_command(&command, host, state).await {
+        state.enqueue(to_jsonl(&resp));
+    }
+
+    if state.shutdown_requested.load(Ordering::SeqCst) {
+        return LineOutcome::Shutdown;
+    }
+    LineOutcome::Done
+}
+
+// ---------------------------------------------------------------------------
+// stdin reader task
+// ---------------------------------------------------------------------------
+
+enum LineMsg {
+    Line(String),
+    ReadError(String),
+    Eof,
+}
+
+async fn stdin_reader<R>(input: R, tx: mpsc::Sender<LineMsg>)
+where
+    R: AsyncRead + Unpin + Send + 'static,
+{
+    let mut reader = JsonlLineReader::new(input);
+    loop {
+        match reader.next_line().await {
+            Ok(Some(line)) => {
+                if tx.send(LineMsg::Line(line)).await.is_err() {
+                    break;
+                }
+            }
+            Ok(None) => {
+                let _ = tx.send(LineMsg::Eof).await;
+                break;
+            }
+            Err(error) => {
+                let _ = tx.send(LineMsg::ReadError(error.to_string())).await;
+                break;
+            }
+        }
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Event loop (testable)
+// ---------------------------------------------------------------------------
+
+/// Run the RPC event loop until a shutdown condition is reached.
+///
+/// Returns the process exit code.
+///
+/// All stdout frames (responses, events, extension errors, prompt preflight)
+/// flow through one writer actor. Ordered drain barriers provide backpressure
+/// without making the input/signal loop responsible for stdout progress.
+pub async fn run_rpc_loop<H, R>(host: H, sink: Arc<dyn RpcSink>, input: R) -> i32
+where
+    H: RpcSessionHost,
+    R: AsyncRead + Unpin + Send + 'static,
+{
+    let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteMessage>();
+    let state = Arc::new(ServerState::new(
+        Arc::clone(&sink),
+        write_tx,
+        ExtensionUiProxy::new(),
+    ));
+    let writer_task = tokio::spawn(writer_actor(write_rx, Arc::clone(&sink)));
+
+    let rebind_signal = Arc::clone(&state.signal);
+    let rebind_flag = Arc::clone(&state.needs_rebind);
+    host.set_rebind(Some(Arc::new(move || {
+        let signal = Arc::clone(&rebind_signal);
+        let flag = Arc::clone(&rebind_flag);
+        Box::pin(async move {
+            flag.store(true, Ordering::SeqCst);
+            signal.notify_one();
+        })
+    })));
+
+    state.rebind(&host).await;
+
+    let (line_tx, mut line_rx) = mpsc::channel::<LineMsg>(64);
+    tokio::spawn(stdin_reader(input, line_tx));
+
+    let exit_code = loop {
+        tokio::select! {
+            biased;
+            () = state.signal.notified() => {
+                if state.shutdown_requested.load(Ordering::SeqCst) {
+                    break 0;
+                }
+                if state.needs_rebind.swap(false, Ordering::SeqCst) {
+                    state.rebind(&host).await;
+                }
+            }
+            msg = line_rx.recv() => {
+                match msg {
+                    Some(LineMsg::Line(line)) => {
+                        let outcome = process_input_line(&line, &host, &state).await;
+                        state.wait_for_output().await;
+                        if outcome == LineOutcome::Shutdown {
+                            break 0;
+                        }
+                    }
+                    Some(LineMsg::ReadError(error)) => {
+                        let response = RpcResponse::err(
+                            None,
+                            "transport",
+                            format!("Failed to read stdin: {error}"),
+                        );
+                        state.enqueue(to_jsonl(&response));
+                        break 1;
+                    }
+                    Some(LineMsg::Eof) | None => break 0,
+                }
+            }
+        }
+    };
+
+    state.cleanup(&host, exit_code).await;
+    writer_task.abort();
+    exit_code
+}
+
+// ---------------------------------------------------------------------------
+// Production entry point
+// ---------------------------------------------------------------------------
+
+/// Production entry point. Takes over stdout, reads stdin, dispatches
+/// commands, handles signals/EOF, and returns the exit code.
+///
+/// The caller (bootstrap) is responsible for `std::process::exit(code)`.
+pub async fn run_rpc_mode<H>(host: H) -> i32
+where
+    H: RpcSessionHost,
+{
+    let _ = output_guard_mod::take_over_stdout();
+    let sink: Arc<dyn RpcSink> = Arc::new(OutputGuardSink);
+
+    let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteMessage>();
+    let state = Arc::new(ServerState::new(
+        Arc::clone(&sink),
+        write_tx,
+        ExtensionUiProxy::new(),
+    ));
+    let writer_task = tokio::spawn(writer_actor(write_rx, Arc::clone(&sink)));
+
+    let rs = Arc::clone(&state.signal);
+    let rf = Arc::clone(&state.needs_rebind);
+    host.set_rebind(Some(Arc::new(move || {
+        let s = Arc::clone(&rs);
+        let f = Arc::clone(&rf);
+        Box::pin(async move {
+            f.store(true, Ordering::SeqCst);
+            s.notify_one();
+        })
+    })));
+    state.rebind(&host).await;
+
+    let (signal_tx, mut signal_rx) = mpsc::channel::<i32>(1);
+    spawn_signal_handlers(signal_tx);
+
+    let stdin = tokio::io::stdin();
+    let (line_tx, mut line_rx) = mpsc::channel::<LineMsg>(64);
+    tokio::spawn(stdin_reader(stdin, line_tx));
+
+    let exit_code = loop {
+        tokio::select! {
+            biased;
+            code = signal_rx.recv() => {
+                break code.unwrap_or(0);
+            }
+            () = state.signal.notified() => {
+                if state.shutdown_requested.load(Ordering::SeqCst) {
+                    break 0;
+                }
+                if state.needs_rebind.swap(false, Ordering::SeqCst) {
+                    state.rebind(&host).await;
+                }
+            }
+            msg = line_rx.recv() => {
+                match msg {
+                    Some(LineMsg::Line(line)) => {
+                        let outcome = process_input_line(&line, &host, &state).await;
+                        state.wait_for_output().await;
+                        if outcome == LineOutcome::Shutdown {
+                            break 0;
+                        }
+                    }
+                    Some(LineMsg::ReadError(error)) => {
+                        let response = RpcResponse::err(
+                            None,
+                            "transport",
+                            format!("Failed to read stdin: {error}"),
+                        );
+                        state.enqueue(to_jsonl(&response));
+                        break 1;
+                    }
+                    Some(LineMsg::Eof) | None => break 0,
+                }
+            }
+        }
+    };
+
+    state.cleanup(&host, exit_code).await;
+    writer_task.abort();
+    output_guard_mod::restore_stdout();
+    exit_code
+}
+
+async fn writer_actor(mut write_rx: mpsc::UnboundedReceiver<WriteMessage>, sink: Arc<dyn RpcSink>) {
+    while let Some(message) = write_rx.recv().await {
+        match message {
+            WriteMessage::Line(line) => {
+                let _ = sink.write_stdout(line).await;
+            }
+            WriteMessage::Drain(done) => {
+                let _ = sink.backpressure().await;
+                let _ = done.send(());
+            }
+        }
+    }
+}
+
+/// Spawn SIGTERM (→143) and SIGHUP (→129, unix-only) handlers.
+fn spawn_signal_handlers(tx: mpsc::Sender<i32>) {
+    #[cfg(unix)]
+    {
+        {
+            let tx = tx.clone();
+            tokio::spawn(async move {
+                use tokio::signal::unix::{SignalKind, signal};
+                if let Ok(mut sig) = signal(SignalKind::terminate()) {
+                    sig.recv().await;
+                    let _ = tx.send(143).await;
+                }
+            });
+        }
+        tokio::spawn(async move {
+            use tokio::signal::unix::{SignalKind, signal};
+            if let Ok(mut sig) = signal(SignalKind::hangup()) {
+                sig.recv().await;
+                // tx is moved here; SIGHUP is the last handler.
+                let _ = tx.send(129).await;
+            }
+        });
+    }
+    #[cfg(not(unix))]
+    {
+        let _ = tx;
+    }
+}
+
+// ===========================================================================
+// Tests
+// ===========================================================================
+
+#[cfg(test)]
+#[allow(clippy::unwrap_used, clippy::panic, clippy::expect_used)]
+mod tests {
+    use super::*;
+    use crate::core::agent_session::extension::ExtensionBindings;
+    use crate::core::agent_session_runtime::{ForkOutcome, ForkPosition};
+    use crate::core::compaction::CompactionResult;
+    use crate::core::sessions::SessionEntry;
+    use crate::modes::rpc::types::{
+        RpcSessionState, RpcSessionTreeNode, RpcSlashCommand, RpcSlashCommandSource, RpcSourceInfo,
+        RpcSourceOrigin, RpcSourceScope, SessionStats, SessionStatsTokens,
+    };
+    use pi_agent::QueueMode;
+    use pi_ai::{ImageContent, Model, ModelThinkingLevel};
+    use pi_ext::client::HostClient;
+    use pi_ext::protocol::{Frame, FrameKind, HelloAck, Method, decode_frame_str, encode_frame};
+    use std::task::{Context, Poll};
+    use tokio::io::ReadBuf;
+    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream};
+
+    // -----------------------------------------------------------------------
+    // FakeRpcHost
+    // -----------------------------------------------------------------------
+
+    #[derive(Clone)]
+    struct FakeConfig {
+        state: RpcSessionState,
+        models: Vec<Model>,
+        commands: Vec<RpcSlashCommand>,
+        cycle_model_result: Option<ModelCycleResult>,
+        cycle_thinking_result: Option<ModelThinkingLevel>,
+        set_thinking_result: bool,
+        compact_result: Option<Result<CompactionResult, String>>,
+        bash_result: Option<Result<BashResult, String>>,
+        fork_outcome: Result<ForkOutcome, String>,
+        leaf_id: Option<String>,
+        prompt_error: Option<String>,
+        session_op_cancelled: bool,
+    }
+
+    impl Default for FakeConfig {
+        fn default() -> Self {
+            Self {
+                state: test_state(),
+                models: vec![],
+                commands: vec![],
+                cycle_model_result: None,
+                cycle_thinking_result: None,
+                set_thinking_result: true,
+                compact_result: None,
+                bash_result: None,
+                fork_outcome: Ok(ForkOutcome::default()),
+                leaf_id: Some("leaf1".into()),
+                prompt_error: None,
+                session_op_cancelled: false,
+            }
+        }
+    }
+
+    fn test_state() -> RpcSessionState {
+        RpcSessionState {
+            model: None,
+            thinking_level: ModelThinkingLevel::Medium,
+            is_streaming: false,
+            is_compacting: false,
+            steering_mode: QueueMode::All,
+            follow_up_mode: QueueMode::OneAtATime,
+            session_file: None,
+            session_id: "test-session".into(),
+            session_name: None,
+            auto_compaction_enabled: true,
+            message_count: 0,
+            pending_message_count: 0,
+        }
+    }
+
+    fn test_stats() -> SessionStats {
+        SessionStats {
+            session_file: None,
+            session_id: "test-session".into(),
+            user_messages: 0,
+            assistant_messages: 0,
+            tool_calls: 0,
+            tool_results: 0,
+            total_messages: 0,
+            tokens: SessionStatsTokens {
+                input: 0,
+                output: 0,
+                cache_read: 0,
+                cache_write: 0,
+                total: 0,
+            },
+            cost: 0.0,
+            context_usage: None,
+        }
+    }
+
+    struct FailingInput;
+
+    impl AsyncRead for FailingInput {
+        fn poll_read(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+            _buf: &mut ReadBuf<'_>,
+        ) -> Poll<io::Result<()>> {
+            Poll::Ready(Err(io::Error::new(
+                io::ErrorKind::ConnectionReset,
+                "stdin transport failed",
+            )))
+        }
+    }
+
+    #[derive(Clone)]
+    struct FakeRpcHost {
+        cfg: Arc<Mutex<FakeConfig>>,
+        calls: Arc<Mutex<Vec<String>>>,
+        disposed: Arc<AtomicBool>,
+        events_tx: Arc<Mutex<Option<tokio::sync::mpsc::UnboundedSender<AgentSessionEvent>>>>,
+        bindings: Arc<Mutex<Option<ExtensionBindings>>>,
+        extension_runner: Arc<Mutex<Option<Arc<HostExtensionRunner>>>>,
+    }
+
+    impl FakeRpcHost {
+        fn new(cfg: FakeConfig) -> Self {
+            Self {
+                cfg: Arc::new(Mutex::new(cfg)),
+                calls: Arc::new(Mutex::new(Vec::new())),
+                disposed: Arc::new(AtomicBool::new(false)),
+                events_tx: Arc::new(Mutex::new(None)),
+                bindings: Arc::new(Mutex::new(None)),
+                extension_runner: Arc::new(Mutex::new(None)),
+            }
+        }
+
+        fn set_extension_runner(&self, runner: Arc<HostExtensionRunner>) {
+            *self.extension_runner.lock().unwrap() = Some(runner);
+        }
+        fn rec(&self, name: &str) {
+            self.calls.lock().unwrap().push(name.to_owned());
+        }
+    }
+
+    impl RpcSessionHost for FakeRpcHost {
+        fn prompt(
+            &self,
+            _msg: String,
+            _img: Vec<ImageContent>,
+            _sb: Option<StreamingBehavior>,
+            preflight: PreflightCallback,
+        ) -> BoxFuture<'static, Result<(), String>> {
+            self.rec("prompt");
+            let err = self.cfg.lock().unwrap().prompt_error.clone();
+            let events_tx = Arc::clone(&self.events_tx);
+            Box::pin(async move {
+                if let Some(e) = err {
+                    preflight(false);
+                    Err(e)
+                } else {
+                    preflight(true);
+                    // Emit an event synchronously immediately after preflight!
+                    if let Some(tx) = events_tx.lock().unwrap().as_ref() {
+                        let _ = tx.send(AgentSessionEvent::TurnStart);
+                    }
+                    Ok(())
+                }
+            })
+        }
+        fn steer(
+            &self,
+            _m: String,
+            _i: Vec<ImageContent>,
+        ) -> BoxFuture<'static, Result<(), String>> {
+            self.rec("steer");
+            Box::pin(async { Ok(()) })
+        }
+        fn follow_up(
+            &self,
+            _m: String,
+            _i: Vec<ImageContent>,
+        ) -> BoxFuture<'static, Result<(), String>> {
+            self.rec("follow_up");
+            Box::pin(async { Ok(()) })
+        }
+        fn abort(&self) -> BoxFuture<'static, ()> {
+            self.rec("abort");
+            Box::pin(async {})
+        }
+        fn get_state(&self) -> BoxFuture<'static, RpcSessionState> {
+            self.rec("get_state");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().state.clone() })
+        }
+        fn get_available_models(&self) -> BoxFuture<'static, Vec<Model>> {
+            self.rec("get_available_models");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().models.clone() })
+        }
+        fn set_model(&self, _m: Model) -> BoxFuture<'static, Result<(), String>> {
+            self.rec("set_model");
+            Box::pin(async { Ok(()) })
+        }
+        fn cycle_model(&self) -> BoxFuture<'static, Option<ModelCycleResult>> {
+            self.rec("cycle_model");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().cycle_model_result.clone() })
+        }
+        fn set_thinking_level(&self, _l: ModelThinkingLevel) -> BoxFuture<'static, bool> {
+            self.rec("set_thinking_level");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().set_thinking_result })
+        }
+        fn cycle_thinking_level(&self) -> BoxFuture<'static, Option<ModelThinkingLevel>> {
+            self.rec("cycle_thinking_level");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().cycle_thinking_result })
+        }
+        fn set_steering_mode(&self, _m: QueueMode) -> BoxFuture<'static, ()> {
+            self.rec("set_steering_mode");
+            Box::pin(async {})
+        }
+        fn set_follow_up_mode(&self, _m: QueueMode) -> BoxFuture<'static, ()> {
+            self.rec("set_follow_up_mode");
+            Box::pin(async {})
+        }
+        fn compact(
+            &self,
+            _ci: Option<String>,
+        ) -> BoxFuture<'static, Result<CompactionResult, String>> {
+            self.rec("compact");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move {
+                cfg.lock()
+                    .unwrap()
+                    .compact_result
+                    .clone()
+                    .unwrap_or(Ok(test_compaction_result()))
+            })
+        }
+        fn set_auto_compaction(&self, _e: bool) -> BoxFuture<'static, ()> {
+            self.rec("set_auto_compaction");
+            Box::pin(async {})
+        }
+        fn set_auto_retry(&self, _e: bool) -> BoxFuture<'static, ()> {
+            self.rec("set_auto_retry");
+            Box::pin(async {})
+        }
+        fn abort_retry(&self) -> BoxFuture<'static, ()> {
+            self.rec("abort_retry");
+            Box::pin(async {})
+        }
+        fn execute_bash(
+            &self,
+            _c: String,
+            _e: Option<bool>,
+        ) -> BoxFuture<'static, Result<BashResult, String>> {
+            self.rec("bash");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move {
+                cfg.lock()
+                    .unwrap()
+                    .bash_result
+                    .clone()
+                    .unwrap_or(Ok(test_bash_result()))
+            })
+        }
+        fn abort_bash(&self) -> BoxFuture<'static, ()> {
+            self.rec("abort_bash");
+            Box::pin(async {})
+        }
+        fn get_session_stats(&self) -> BoxFuture<'static, SessionStats> {
+            self.rec("get_session_stats");
+            Box::pin(async { test_stats() })
+        }
+        fn export_to_html(&self, _p: Option<String>) -> BoxFuture<'static, Result<String, String>> {
+            self.rec("export_html");
+            Box::pin(async { Ok("/tmp/out.html".into()) })
+        }
+        fn set_session_name(&self, _n: String) -> BoxFuture<'static, Result<(), String>> {
+            self.rec("set_session_name");
+            Box::pin(async { Ok(()) })
+        }
+        fn new_session(&self, _p: Option<String>) -> BoxFuture<'static, Result<bool, String>> {
+            self.rec("new_session");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { Ok(cfg.lock().unwrap().session_op_cancelled) })
+        }
+        fn switch_session(&self, _p: String) -> BoxFuture<'static, Result<bool, String>> {
+            self.rec("switch_session");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { Ok(cfg.lock().unwrap().session_op_cancelled) })
+        }
+        fn fork(
+            &self,
+            _e: String,
+            _p: ForkPosition,
+        ) -> BoxFuture<'static, Result<ForkOutcome, String>> {
+            self.rec("fork");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().fork_outcome.clone() })
+        }
+        fn get_entries(&self) -> BoxFuture<'static, Vec<SessionEntry>> {
+            self.rec("get_entries");
+            Box::pin(async { vec![] })
+        }
+        fn get_leaf_id(&self) -> BoxFuture<'static, Option<String>> {
+            self.rec("get_leaf_id");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().leaf_id.clone() })
+        }
+        fn get_tree(&self) -> BoxFuture<'static, Vec<RpcSessionTreeNode>> {
+            self.rec("get_tree");
+            Box::pin(async { vec![] })
+        }
+        fn get_fork_messages(&self) -> BoxFuture<'static, Vec<ForkMessage>> {
+            self.rec("get_fork_messages");
+            Box::pin(async { vec![] })
+        }
+        fn get_last_assistant_text(&self) -> BoxFuture<'static, Option<String>> {
+            self.rec("get_last_assistant_text");
+            Box::pin(async { None })
+        }
+        fn get_messages(&self) -> BoxFuture<'static, Vec<AgentMessage>> {
+            self.rec("get_messages");
+            Box::pin(async { vec![] })
+        }
+        fn get_commands(&self) -> BoxFuture<'static, Vec<RpcSlashCommand>> {
+            self.rec("get_commands");
+            let cfg = Arc::clone(&self.cfg);
+            Box::pin(async move { cfg.lock().unwrap().commands.clone() })
+        }
+        fn subscribe(
+            &self,
+            listener: Arc<dyn Fn(&AgentSessionEvent) + Send + Sync>,
+        ) -> Box<dyn Fn() + Send + Sync> {
+            self.rec("subscribe");
+            let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<AgentSessionEvent>();
+            *self.events_tx.lock().unwrap() = Some(tx);
+            tokio::spawn(async move {
+                while let Some(event) = rx.recv().await {
+                    listener(&event);
+                }
+            });
+            Box::new(|| {})
+        }
+        fn register_backpressure_hook(
+            &self,
+            _hook: Arc<dyn Fn() -> BoxFuture<'static, ()> + Send + Sync>,
+        ) -> Box<dyn Fn() + Send + Sync> {
+            self.rec("register_backpressure_hook");
+            Box::new(|| {})
+        }
+        fn bind_extensions_rpc(
+            &self,
+            bindings: ExtensionBindings,
+        ) -> BoxFuture<'static, Result<(), String>> {
+            self.rec("bind_extensions_rpc");
+            *self.bindings.lock().unwrap() = Some(bindings);
+            Box::pin(async { Ok(()) })
+        }
+        fn host_extension_runner(&self) -> Option<Arc<HostExtensionRunner>> {
+            self.extension_runner.lock().unwrap().clone()
+        }
+        fn dispose(&self) -> BoxFuture<'static, ()> {
+            self.rec("dispose");
+            self.disposed.store(true, Ordering::SeqCst);
+            Box::pin(async {})
+        }
+        fn set_rebind(&self, _cb: Option<RebindCallback>) {
+            self.rec("set_rebind");
+        }
+    }
+    fn test_compaction_result() -> CompactionResult {
+        CompactionResult {
+            summary: "Summary".into(),
+            first_kept_entry_id: "entry1".into(),
+            tokens_before: 1000,
+            estimated_tokens_after: Some(500),
+            details: None,
+            from_hook: None,
+        }
+    }
+    fn test_bash_result() -> BashResult {
+        BashResult {
+            output: "done".into(),
+            exit_code: Some(0),
+            cancelled: false,
+            truncated: false,
+            full_output_path: None,
+        }
+    }
+    // -----------------------------------------------------------------------
+    // Test helpers
+    // -----------------------------------------------------------------------
+
+    fn make_state(
+        sink: BufferSink,
+    ) -> (
+        Arc<ServerState>,
+        BufferSink,
+        mpsc::UnboundedReceiver<WriteMessage>,
+    ) {
+        let s = sink.clone();
+        let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteMessage>();
+        let state = Arc::new(ServerState::new(
+            Arc::new(sink) as Arc<dyn RpcSink>,
+            write_tx,
+            ExtensionUiProxy::new(),
+        ));
+        (state, s, write_rx)
+    }
+
+    /// Drain all pending write-channel frames into the sink.
+    async fn drain(sink: &BufferSink, write_rx: &mut mpsc::UnboundedReceiver<WriteMessage>) {
+        while let Ok(message) = write_rx.try_recv() {
+            match message {
+                WriteMessage::Line(line) => {
+                    let _ = sink.write_stdout(line).await;
+                }
+                WriteMessage::Drain(done) => {
+                    let _ = sink.backpressure().await;
+                    let _ = done.send(());
+                }
+            }
+        }
+    }
+
+    async fn dispatch(cmd_json: &str, cfg: FakeConfig) -> (Value, BufferSink) {
+        let host = FakeRpcHost::new(cfg);
+        let sink = BufferSink::new();
+        let (state, sink_clone, mut write_rx) = make_state(sink);
+        process_input_line(cmd_json, &host, &state).await;
+        drain(&sink_clone, &mut write_rx).await;
+        let lines = sink_clone.stdout_lines();
+        assert!(!lines.is_empty(), "expected at least one response line");
+        let resp: Value = serde_json::from_str(&lines[0]).unwrap();
+        (resp, sink_clone)
+    }
+
+    async fn dispatch_no_response(
+        cmd_json: &str,
+        cfg: FakeConfig,
+    ) -> (BufferSink, mpsc::UnboundedReceiver<WriteMessage>) {
+        let host = FakeRpcHost::new(cfg);
+        let sink = BufferSink::new();
+        let (state, sink_clone, mut write_rx) = make_state(sink);
+        process_input_line(cmd_json, &host, &state).await;
+        drain(&sink_clone, &mut write_rx).await;
+        (sink_clone, write_rx)
+    }
+
+    // -----------------------------------------------------------------------
+    // Unknown command
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn unknown_command_echoes_id_and_type() {
+        let (resp, _) = dispatch(
+            r#"{"type":"totally_unknown","id":"abc","foo":"bar"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(resp["type"], "response");
+        assert_eq!(resp["id"], "abc");
+        assert_eq!(resp["command"], "totally_unknown");
+        assert_eq!(resp["success"], false);
+        assert_eq!(resp["error"], "Unknown command: totally_unknown");
+    }
+
+    // -----------------------------------------------------------------------
+    // Malformed JSON
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn malformed_json_parse_error_no_id() {
+        let (resp, _) = dispatch("{not valid json", FakeConfig::default()).await;
+        assert_eq!(resp["command"], "parse");
+        assert_eq!(resp["success"], false);
+        assert!(resp.get("id").is_none() || resp["id"].is_null());
+        assert!(
+            resp["error"]
+                .as_str()
+                .unwrap()
+                .starts_with("Failed to parse")
+        );
+    }
+
+    #[tokio::test]
+    async fn field_validation_error_echoes_valid_id() {
+        let (response, _) = dispatch(
+            r#"{"type":"bash","id":"req-17","command":"true","excludeFromContext":"yes"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(response["id"], "req-17");
+        assert_eq!(response["command"], "parse");
+        assert_eq!(response["success"], false);
+        assert!(
+            response["error"]
+                .as_str()
+                .unwrap()
+                .contains("excludeFromContext")
+        );
+    }
+
+    // -----------------------------------------------------------------------
+    // Prompt preflight semantics
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn prompt_emits_success_once() {
+        let (sink, mut write_rx) = dispatch_no_response(
+            r#"{"type":"prompt","id":"p1","message":"hello"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        drain(&sink, &mut write_rx).await;
+        let lines = sink.stdout_lines();
+        assert_eq!(lines.len(), 1, "exactly one success response");
+        let resp: Value = serde_json::from_str(&lines[0]).unwrap();
+        assert_eq!(resp["command"], "prompt");
+        assert_eq!(resp["success"], true);
+    }
+
+    #[tokio::test]
+    async fn prompt_emits_error() {
+        let cfg = FakeConfig {
+            prompt_error: Some("No model".into()),
+            ..FakeConfig::default()
+        };
+        let (sink, mut write_rx) =
+            dispatch_no_response(r#"{"type":"prompt","id":"p2","message":"hi"}"#, cfg).await;
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        drain(&sink, &mut write_rx).await;
+        let lines = sink.stdout_lines();
+        assert_eq!(lines.len(), 1);
+        let resp: Value = serde_json::from_str(&lines[0]).unwrap();
+        assert_eq!(resp["success"], false);
+        assert_eq!(resp["error"], "No model");
+    }
+
+    // -----------------------------------------------------------------------
+    // Simple success responses
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn steer_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"steer","id":"s1","message":"left"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "steer");
+        assert_eq!(r["success"], true);
+    }
+
+    #[tokio::test]
+    async fn follow_up_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"follow_up","id":"f1","message":"next"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "follow_up");
+    }
+
+    #[tokio::test]
+    async fn abort_success() {
+        let (r, _) = dispatch(r#"{"type":"abort","id":"a1"}"#, FakeConfig::default()).await;
+        assert_eq!(r["command"], "abort");
+    }
+
+    #[tokio::test]
+    async fn set_thinking_level_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_thinking_level","id":"t1","level":"high"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "set_thinking_level");
+        assert_eq!(r["success"], true);
+    }
+
+    #[tokio::test]
+    async fn set_thinking_level_uncommitted_returns_error() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_thinking_level","id":"t2","level":"high"}"#,
+            FakeConfig {
+                set_thinking_result: false,
+                ..FakeConfig::default()
+            },
+        )
+        .await;
+        assert_eq!(r["command"], "set_thinking_level");
+        assert_eq!(r["success"], false);
+        assert_eq!(r["error"], "Failed to persist thinking level change");
+    }
+
+    #[tokio::test]
+    async fn set_steering_mode_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_steering_mode","id":"sm1","mode":"all"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "set_steering_mode");
+    }
+
+    #[tokio::test]
+    async fn set_follow_up_mode_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_follow_up_mode","id":"fm1","mode":"one-at-a-time"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "set_follow_up_mode");
+    }
+
+    #[tokio::test]
+    async fn set_auto_compaction_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_auto_compaction","id":"ac1","enabled":true}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "set_auto_compaction");
+    }
+
+    #[tokio::test]
+    async fn set_auto_retry_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_auto_retry","id":"ar1","enabled":false}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "set_auto_retry");
+    }
+
+    #[tokio::test]
+    async fn abort_retry_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"abort_retry","id":"abr1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["command"], "abort_retry");
+    }
+
+    #[tokio::test]
+    async fn abort_bash_success() {
+        let (r, _) = dispatch(r#"{"type":"abort_bash","id":"ab1"}"#, FakeConfig::default()).await;
+        assert_eq!(r["command"], "abort_bash");
+    }
+
+    // -----------------------------------------------------------------------
+    // Data responses
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn get_state_returns_snapshot() {
+        let (r, _) = dispatch(r#"{"type":"get_state","id":"gs1"}"#, FakeConfig::default()).await;
+        assert_eq!(r["data"]["sessionId"], "test-session");
+        assert_eq!(r["data"]["thinkingLevel"], "medium");
+    }
+
+    #[tokio::test]
+    async fn get_session_stats_returns_data() {
+        let (r, _) = dispatch(
+            r#"{"type":"get_session_stats","id":"sst1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["data"]["sessionId"], "test-session");
+    }
+
+    #[tokio::test]
+    async fn export_html_returns_path() {
+        let (r, _) = dispatch(
+            r#"{"type":"export_html","id":"eh1","outputPath":"/tmp/x.html"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["data"]["path"], "/tmp/out.html");
+    }
+
+    #[tokio::test]
+    async fn get_last_assistant_text_null() {
+        let (r, _) = dispatch(
+            r#"{"type":"get_last_assistant_text","id":"lat1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert!(r["data"]["text"].is_null());
+    }
+
+    #[tokio::test]
+    async fn get_messages_empty() {
+        let (r, _) = dispatch(
+            r#"{"type":"get_messages","id":"gm1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["data"]["messages"].as_array().unwrap().len(), 0);
+    }
+
+    #[tokio::test]
+    async fn get_commands_returns_complete_catalog() {
+        let source_info = RpcSourceInfo {
+            path: "/tmp/resource".into(),
+            source: "test".into(),
+            scope: RpcSourceScope::Temporary,
+            origin: RpcSourceOrigin::TopLevel,
+            base_dir: None,
+        };
+        let commands = [
+            ("ext-command", RpcSlashCommandSource::Extension),
+            ("deploy", RpcSlashCommandSource::Prompt),
+            ("skill:review", RpcSlashCommandSource::Skill),
+        ]
+        .into_iter()
+        .map(|(name, source)| RpcSlashCommand {
+            name: name.into(),
+            description: Some(format!("{name} description")),
+            source,
+            source_info: source_info.clone(),
+        })
+        .collect();
+        let (response, _) = dispatch(
+            r#"{"type":"get_commands","id":"gc1"}"#,
+            FakeConfig {
+                commands,
+                ..FakeConfig::default()
+            },
+        )
+        .await;
+        let catalog = response["data"]["commands"].as_array().unwrap();
+        assert_eq!(catalog.len(), 3);
+        assert_eq!(catalog[0]["source"], "extension");
+        assert_eq!(catalog[1]["source"], "prompt");
+        assert_eq!(catalog[2]["name"], "skill:review");
+        assert_eq!(catalog[2]["source"], "skill");
+    }
+
+    #[tokio::test]
+    async fn bash_returns_result() {
+        let (r, _) = dispatch(
+            r#"{"type":"bash","id":"b1","command":"echo hi"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["data"]["output"], "done");
+    }
+
+    #[tokio::test]
+    async fn compact_returns_result() {
+        let (r, _) = dispatch(r#"{"type":"compact","id":"c1"}"#, FakeConfig::default()).await;
+        assert_eq!(r["data"]["summary"], "Summary");
+    }
+
+    #[tokio::test]
+    async fn cycle_model_null() {
+        let (r, _) = dispatch(
+            r#"{"type":"cycle_model","id":"cm1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert!(r["data"].is_null());
+    }
+
+    #[tokio::test]
+    async fn cycle_thinking_null() {
+        let (r, _) = dispatch(
+            r#"{"type":"cycle_thinking_level","id":"ct1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert!(r["data"].is_null());
+    }
+
+    #[tokio::test]
+    async fn get_available_models_empty() {
+        let (r, _) = dispatch(
+            r#"{"type":"get_available_models","id":"gam1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["data"]["models"].as_array().unwrap().len(), 0);
+    }
+
+    #[tokio::test]
+    async fn get_fork_messages_empty() {
+        let (r, _) = dispatch(
+            r#"{"type":"get_fork_messages","id":"gfm1"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert!(r["data"]["messages"].is_array());
+    }
+
+    #[tokio::test]
+    async fn get_tree_empty() {
+        let (r, _) = dispatch(r#"{"type":"get_tree","id":"gt1"}"#, FakeConfig::default()).await;
+        assert!(r["data"]["tree"].is_array());
+    }
+
+    // -----------------------------------------------------------------------
+    // Error cases
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn set_model_not_found() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_model","id":"sm1","provider":"openai","modelId":"gpt-999"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["success"], false);
+        assert_eq!(r["error"], "Model not found: openai/gpt-999");
+    }
+    #[tokio::test]
+    async fn clone_no_leaf() {
+        let cfg = FakeConfig {
+            leaf_id: None,
+            ..Default::default()
+        };
+        let (r, _) = dispatch(r#"{"type":"clone","id":"cl1"}"#, cfg).await;
+        assert_eq!(
+            r["error"],
+            "Cannot clone session: no current entry selected"
+        );
+    }
+
+    #[tokio::test]
+    async fn set_session_name_empty() {
+        let (r, _) = dispatch(
+            r#"{"type":"set_session_name","id":"ssn1","name":"   "}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["error"], "Session name cannot be empty");
+    }
+
+    #[tokio::test]
+    async fn get_entries_since_not_found() {
+        let (r, _) = dispatch(
+            r#"{"type":"get_entries","id":"ge1","since":"nope"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["error"], "Entry not found: nope");
+    }
+
+    // -----------------------------------------------------------------------
+    // Cancelled mutations
+    // -----------------------------------------------------------------------
+    #[tokio::test]
+    async fn new_session_cancelled() {
+        let cfg = FakeConfig {
+            session_op_cancelled: true,
+            ..Default::default()
+        };
+        let (r, _) = dispatch(r#"{"type":"new_session","id":"ns1"}"#, cfg).await;
+        assert_eq!(r["data"]["cancelled"], true);
+    }
+
+    #[tokio::test]
+    async fn fork_returns_text() {
+        let cfg = FakeConfig {
+            fork_outcome: Ok(ForkOutcome {
+                cancelled: false,
+                selected_text: Some("fork here".into()),
+            }),
+            ..Default::default()
+        };
+        let (r, _) = dispatch(r#"{"type":"fork","id":"fk1","entryId":"e1"}"#, cfg).await;
+        assert_eq!(r["data"]["text"], "fork here");
+    }
+
+    #[tokio::test]
+    async fn clone_with_leaf() {
+        let (r, host) = dispatch(r#"{"type":"clone","id":"cl2"}"#, FakeConfig::default()).await;
+        let _ = r;
+        let _ = host;
+    }
+
+    #[tokio::test]
+    async fn switch_session_success() {
+        let (r, _) = dispatch(
+            r#"{"type":"switch_session","id":"sw1","sessionPath":"/tmp/s.jsonl"}"#,
+            FakeConfig::default(),
+        )
+        .await;
+        assert_eq!(r["data"]["cancelled"], false);
+    }
+
+    // -----------------------------------------------------------------------
+    // Extension UI routing
+    // -----------------------------------------------------------------------
+
+    struct RpcHostPeer {
+        read: BufReader<DuplexStream>,
+        write: DuplexStream,
+    }
+
+    impl RpcHostPeer {
+        async fn read_frame(&mut self) -> Result<Frame, Box<dyn std::error::Error>> {
+            let mut line = String::new();
+            self.read.read_line(&mut line).await?;
+            Ok(decode_frame_str(line.trim_end())?)
+        }
+
+        async fn write_frame(&mut self, frame: &Frame) -> Result<(), Box<dyn std::error::Error>> {
+            self.write.write_all(&encode_frame(frame)?).await?;
+            self.write.flush().await?;
+            Ok(())
+        }
+    }
+
+    async fn make_rpc_extension_runner()
+    -> Result<(Arc<HostExtensionRunner>, RpcHostPeer), Box<dyn std::error::Error>> {
+        let (client_stdout, host_stdout) = tokio::io::duplex(64 * 1024);
+        let (host_stdin, client_stdin) = tokio::io::duplex(64 * 1024);
+        let client = Arc::new(HostClient::connect_boxed(
+            Box::new(client_stdin),
+            Box::new(client_stdout),
+            Box::new(tokio::io::empty()),
+            None,
+        ));
+        let connect_client = Arc::clone(&client);
+        let connect = tokio::spawn(async move {
+            HostExtensionRunner::connect_with_cwd_and_trust(
+                connect_client,
+                Vec::new(),
+                "/workspace",
+                false,
+                std::time::Duration::from_secs(1),
+            )
+            .await
+        });
+        let mut peer = RpcHostPeer {
+            read: BufReader::new(host_stdin),
+            write: host_stdout,
+        };
+        let hello = peer.read_frame().await?;
+        peer.write_frame(&Frame::response(
+            hello.id,
+            Method::Hello,
+            serde_json::to_value(HelloAck::local())?,
+        ))
+        .await?;
+        let load = peer.read_frame().await?;
+        assert_eq!(load.method, "extensions.load");
+        peer.write_frame(&Frame::response(
+            load.id,
+            Method::Notify,
+            serde_json::json!({
+                "tools": [],
+                "commands": [],
+                "shortcuts": [],
+                "flags": [],
+                "renderers": [],
+                "providers": [],
+                "handlers": [],
+                "errors": [],
+                "terminalInput": false
+            }),
+        ))
+        .await?;
+        Ok((connect.await??, peer))
+    }
+
+    #[tokio::test]
+    async fn host_dialog_round_trips_through_rpc_stdout_and_stdin()
+    -> Result<(), Box<dyn std::error::Error>> {
+        let (runner, mut peer) = make_rpc_extension_runner().await?;
+        let host = FakeRpcHost::new(FakeConfig::default());
+        host.set_extension_runner(Arc::clone(&runner));
+        let sink = BufferSink::new();
+        let sink_arc = Arc::new(sink.clone()) as Arc<dyn RpcSink>;
+        let (write_tx, write_rx) = mpsc::unbounded_channel::<WriteMessage>();
+        let state = ServerState::new(sink_arc.clone(), write_tx, ExtensionUiProxy::new());
+        let writer = tokio::spawn(writer_actor(write_rx, sink_arc));
+        state.rebind(&host).await;
+
+        peer.write_frame(&Frame {
+            id: 901,
+            kind: FrameKind::Req,
+            method: Method::Select.as_str().to_owned(),
+            payload: serde_json::json!({
+                "title": "Pick",
+                "options": ["a", "b"],
+                "timeoutMs": 1000
+            }),
+        })
+        .await?;
+
+        let ui_request = tokio::time::timeout(std::time::Duration::from_secs(1), async {
+            loop {
+                state.wait_for_output().await;
+                if let Some(line) = sink.stdout_lines().last().cloned()
+                    && serde_json::from_str::<Value>(&line)
+                        .ok()
+                        .and_then(|value| {
+                            value.get("type").and_then(Value::as_str).map(str::to_owned)
+                        })
+                        .as_deref()
+                        == Some("extension_ui_request")
+                {
+                    break line;
+                }
+                tokio::task::yield_now().await;
+            }
+        })
+        .await?;
+        let request_json: Value = serde_json::from_str(&ui_request)?;
+        assert_eq!(request_json["method"], "select");
+        assert_eq!(request_json["title"], "Pick");
+        let rpc_id = request_json["id"].as_str().ok_or("missing RPC UI id")?;
+        let response = serde_json::json!({
+            "type": "extension_ui_response",
+            "id": rpc_id,
+            "value": "b"
+        })
+        .to_string();
+        assert_eq!(
+            process_input_line(&response, &host, &state).await,
+            LineOutcome::Done
+        );
+
+        let host_response =
+            tokio::time::timeout(std::time::Duration::from_secs(1), peer.read_frame()).await??;
+        assert_eq!(host_response.kind, FrameKind::Res);
+        assert_eq!(host_response.id, 901);
+        assert_eq!(host_response.method, "select");
+        assert_eq!(host_response.payload["value"], "b");
+
+        state.cleanup(&host, 0).await;
+        writer.abort();
+        runner.shutdown_once().await;
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn ui_response_routes_to_proxy() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = BufferSink::new();
+        let proxy = ExtensionUiProxy::new();
+        let (req, rx) = proxy.create_dialog(|id| RpcExtensionUiRequest::Select {
+            id: id.to_owned(),
+            title: "Pick".into(),
+            options: vec!["a".into()],
+            timeout: None,
+        });
+        let pending_id = req.id().to_owned();
+        let (write_tx, _write_rx) = mpsc::unbounded_channel::<WriteMessage>();
+        let state = ServerState::new(Arc::new(sink.clone()) as Arc<dyn RpcSink>, write_tx, proxy);
+        let resp_json =
+            format!(r#"{{"type":"extension_ui_response","id":"{pending_id}","value":"picked"}}"#);
+        process_input_line(&resp_json, &host, &state).await;
+        assert!(sink.stdout_lines().is_empty());
+        let resp = rx.await.unwrap();
+        match resp {
+            RpcExtensionUiResponse::Value { value, .. } => assert_eq!(value, "picked"),
+            _ => panic!(),
+        }
+    }
+
+    #[tokio::test]
+    async fn orphan_ui_response_dropped() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = BufferSink::new();
+        let (state, _, _) = make_state(sink.clone());
+        process_input_line(
+            r#"{"type":"extension_ui_response","id":"orphan","value":"x"}"#,
+            &host,
+            &state,
+        )
+        .await;
+        assert!(sink.stdout_lines().is_empty());
+    }
+
+    // -----------------------------------------------------------------------
+    // Queued prompt ordering
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn prompt_then_steer_ordering() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = BufferSink::new();
+        let (state, sink_clone, mut write_rx) = make_state(sink);
+        process_input_line(
+            r#"{"type":"prompt","id":"q1","message":"first"}"#,
+            &host,
+            &state,
+        )
+        .await;
+        process_input_line(
+            r#"{"type":"steer","id":"q2","message":"second"}"#,
+            &host,
+            &state,
+        )
+        .await;
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        drain(&sink_clone, &mut write_rx).await;
+        let lines = sink_clone.stdout_lines();
+        assert_eq!(lines.len(), 2);
+        let r1: Value = serde_json::from_str(&lines[0]).unwrap();
+        let r2: Value = serde_json::from_str(&lines[1]).unwrap();
+        assert_eq!(r1["command"], "prompt");
+        assert_eq!(r2["command"], "steer");
+    }
+    #[tokio::test]
+    async fn prompt_response_precedes_agent_events() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = BufferSink::new();
+        let (state, sink_clone, mut write_rx) = make_state(sink);
+
+        // Setup event subscriber exactly as the real rebind() does.
+        // This ensures the host's emitted events go into the same write_tx queue.
+        let event_tx = state.write_tx.clone();
+        let unsub = host.subscribe(Arc::new(move |event: &AgentSessionEvent| {
+            let _ = event_tx.send(WriteMessage::Line(to_jsonl(event)));
+        }));
+        *state.unsubscribe_events.lock().unwrap() = Some(unsub);
+
+        // Dispatch prompt. The FakeRpcHost will call preflight(true) and then
+        // IMMEDIATELY emit TurnStart.
+        process_input_line(
+            r#"{"type":"prompt","id":"q3","message":"test"}"#,
+            &host,
+            &state,
+        )
+        .await;
+
+        // Wait for the spawned prompt task to run.
+        tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+        drain(&sink_clone, &mut write_rx).await;
+
+        let lines = sink_clone.stdout_lines();
+        assert_eq!(
+            lines.len(),
+            2,
+            "Expected exactly 2 frames: response and event"
+        );
+
+        // First line MUST be the prompt success response (from preflight).
+        let r1: Value = serde_json::from_str(&lines[0]).unwrap();
+        assert_eq!(r1["type"], "response");
+        assert_eq!(r1["command"], "prompt");
+        assert_eq!(r1["success"], true);
+
+        // Second line MUST be the event (emitted right after preflight).
+        let r2: Value = serde_json::from_str(&lines[1]).unwrap();
+        assert_eq!(r2["type"], "turn_start");
+    }
+    // -----------------------------------------------------------------------
+    // Shutdown after extension handler
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn shutdown_flag_triggers_exit() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = BufferSink::new();
+        let (state, _, _) = make_state(sink);
+        // Simulate an extension invoking the RPC shutdown handler.
+        state.shutdown_requested.store(true, Ordering::SeqCst);
+        let outcome = process_input_line(r#"{"type":"get_state","id":"x"}"#, &host, &state).await;
+        assert_eq!(outcome, LineOutcome::Shutdown);
+    }
+
+    #[tokio::test]
+    async fn extension_shutdown_wakes_idle_loop() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let observer = host.clone();
+        let sink = Arc::new(BufferSink::new()) as Arc<dyn RpcSink>;
+        let (_input_writer, input_reader) = tokio::io::duplex(64);
+        let loop_task = tokio::spawn(run_rpc_loop(host, sink, input_reader));
+
+        tokio::time::timeout(std::time::Duration::from_secs(1), async {
+            loop {
+                let handler = observer
+                    .bindings
+                    .lock()
+                    .unwrap()
+                    .as_ref()
+                    .and_then(|bindings| bindings.shutdown_handler.clone());
+                if let Some(handler) = handler {
+                    handler();
+                    break;
+                }
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("extension bindings were not installed");
+
+        let code = tokio::time::timeout(std::time::Duration::from_secs(1), loop_task)
+            .await
+            .expect("idle RPC loop did not wake")
+            .expect("RPC loop task panicked");
+        assert_eq!(code, 0);
+    }
+
+    // -----------------------------------------------------------------------
+    // Event loop with EOF
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn loop_eof_exit_zero() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = Arc::new(BufferSink::new()) as Arc<dyn RpcSink>;
+        let input: &[u8] = b"";
+        let code = run_rpc_loop(host, sink, input).await;
+        assert_eq!(code, 0);
+    }
+
+    #[tokio::test]
+    async fn loop_stdin_read_error_is_protocol_visible_and_nonzero() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let buffer = Arc::new(BufferSink::new());
+        let sink = Arc::clone(&buffer) as Arc<dyn RpcSink>;
+        let code = run_rpc_loop(host, sink, FailingInput).await;
+        assert_eq!(code, 1);
+        let lines = buffer.stdout_lines();
+        assert_eq!(lines.len(), 1);
+        let response: Value = serde_json::from_str(&lines[0]).unwrap();
+        assert_eq!(response["type"], "response");
+        assert_eq!(response["command"], "transport");
+        assert_eq!(response["success"], false);
+        assert!(
+            response["error"]
+                .as_str()
+                .unwrap()
+                .contains("stdin transport failed")
+        );
+    }
+
+    #[tokio::test]
+    async fn loop_command_then_eof() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let buf = BufferSink::new();
+        let sink: Arc<dyn RpcSink> = Arc::new(buf.clone());
+        let input: &[u8] = b"{\"type\":\"get_state\",\"id\":\"l1\"}\n";
+        let code = run_rpc_loop(host, Arc::clone(&sink), input).await;
+        assert_eq!(code, 0);
+        assert_eq!(buf.stdout_lines().len(), 1);
+    }
+
+    #[tokio::test]
+    async fn command_dispatch_waits_for_writer_drain() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let observer = host.clone();
+        let gated = GatedSink::default();
+        let sink = Arc::new(gated.clone()) as Arc<dyn RpcSink>;
+        let input =
+            &b"{\"type\":\"get_state\",\"id\":\"one\"}\n{\"type\":\"get_state\",\"id\":\"two\"}\n"
+                [..];
+        let loop_task = tokio::spawn(run_rpc_loop(host, sink, input));
+
+        gated.wait_for_write(1).await;
+        assert_eq!(
+            observer
+                .calls
+                .lock()
+                .unwrap()
+                .iter()
+                .filter(|call| call.as_str() == "get_state")
+                .count(),
+            1,
+            "second request dispatched before first response drained"
+        );
+        gated.release.notify_one();
+
+        gated.wait_for_write(2).await;
+        gated.release.notify_one();
+        let code = tokio::time::timeout(std::time::Duration::from_secs(1), loop_task)
+            .await
+            .expect("RPC loop did not finish")
+            .expect("RPC loop task panicked");
+        assert_eq!(code, 0);
+        assert_eq!(
+            observer
+                .calls
+                .lock()
+                .unwrap()
+                .iter()
+                .filter(|call| call.as_str() == "get_state")
+                .count(),
+            2
+        );
+    }
+
+    #[tokio::test]
+    async fn loop_disposes_on_exit() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = Arc::new(BufferSink::new()) as Arc<dyn RpcSink>;
+        let input: &[u8] = b"";
+        let _ = run_rpc_loop(host, sink, input).await;
+        // host was moved; we can't check disposed flag after move.
+        // Instead, verify via a shared flag.
+    }
+
+    #[tokio::test]
+    async fn loop_disposes_host() {
+        let host = Arc::new(FakeRpcHost::new(FakeConfig::default()));
+        let disposed = {
+            let h = Arc::clone(&host);
+            // We can't easily pass Arc<FakeRpcHost> to run_rpc_loop since it
+            // takes H: RpcSessionHost by value. Instead, test rebind+dispose
+            // via direct calls.
+            let sink = Arc::new(BufferSink::new()) as Arc<dyn RpcSink>;
+            let (write_tx, _) = mpsc::unbounded_channel::<WriteMessage>();
+            let state = ServerState::new(sink, write_tx, ExtensionUiProxy::new());
+            state.rebind(&*h).await;
+            state.cleanup(&*h, 0).await;
+            h.disposed.load(Ordering::SeqCst)
+        };
+        assert!(disposed);
+    }
+
+    // -----------------------------------------------------------------------
+    // BufferSink ordering
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn buffer_sink_fifo_order() {
+        let sink = BufferSink::new();
+        sink.clone().write_stdout("a\n".into()).await.unwrap();
+        sink.clone().write_stdout("b\n".into()).await.unwrap();
+        sink.clone().write_stdout("c\n".into()).await.unwrap();
+        assert_eq!(sink.stdout_lines(), vec!["a", "b", "c"]);
+    }
+
+    // -----------------------------------------------------------------------
+    // No-id command
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn command_without_id() {
+        let (r, _) = dispatch(r#"{"type":"get_state"}"#, FakeConfig::default()).await;
+        assert!(r.get("id").is_none() || r["id"].is_null());
+    }
+
+    // -----------------------------------------------------------------------
+    // Rebind binds extensions
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn rebind_binds_and_subscribes() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = Arc::new(BufferSink::new()) as Arc<dyn RpcSink>;
+        let (write_tx, _) = mpsc::unbounded_channel::<WriteMessage>();
+        let state = ServerState::new(sink, write_tx, ExtensionUiProxy::new());
+        state.rebind(&host).await;
+        let calls = host.calls.lock().unwrap();
+        assert!(calls.contains(&"bind_extensions_rpc".to_owned()));
+        assert!(calls.contains(&"subscribe".to_owned()));
+        assert!(calls.contains(&"register_backpressure_hook".to_owned()));
+    }
+
+    #[tokio::test]
+    async fn rebind_extension_error_listener_preserves_structured_fields() {
+        let host = FakeRpcHost::new(FakeConfig::default());
+        let sink = Arc::new(BufferSink::new()) as Arc<dyn RpcSink>;
+        let (write_tx, mut write_rx) = mpsc::unbounded_channel::<WriteMessage>();
+        let state = ServerState::new(sink, write_tx, ExtensionUiProxy::new());
+        state.rebind(&host).await;
+
+        let listener = host
+            .bindings
+            .lock()
+            .unwrap()
+            .as_ref()
+            .and_then(|bindings| bindings.on_error.clone())
+            .expect("RPC rebind must install extension error listener");
+        listener("/workspace/ext.ts", "tool_call", "handler failed");
+
+        let WriteMessage::Line(line) = write_rx.recv().await.expect("extension error output")
+        else {
+            panic!("unexpected drain message");
+        };
+        let parsed: Value = serde_json::from_str(&line).unwrap();
+        assert_eq!(parsed["extensionPath"], "/workspace/ext.ts");
+        assert_eq!(parsed["event"], "tool_call");
+        assert_eq!(parsed["error"], "handler failed");
+    }
+
+    // -----------------------------------------------------------------------
+    // ExtensionErrorOutput serialization
+    // -----------------------------------------------------------------------
+
+    #[tokio::test]
+    async fn extension_error_serializes() {
+        let output = ExtensionErrorOutput::new("ext/path", "event_type", "boom");
+        let line = to_jsonl(&output);
+        let parsed: Value = serde_json::from_str(&line).unwrap();
+        assert_eq!(parsed["type"], "extension_error");
+        assert_eq!(parsed["extensionPath"], "ext/path");
+        assert_eq!(parsed["event"], "event_type");
+        assert_eq!(parsed["error"], "boom");
+    }
+    #[test]
+    fn host_extension_error_wire_fields_are_recovered() {
+        let output = map_extension_error_event(&ExtensionErrorEvent {
+            code: "extension_error".to_owned(),
+            message: "[/workspace/ext.ts] agent_start: handler failed".to_owned(),
+            retryable: false,
+            data: None,
+        });
+        let parsed: Value = serde_json::from_str(&to_jsonl(&output)).unwrap();
+        assert_eq!(parsed["extensionPath"], "/workspace/ext.ts");
+        assert_eq!(parsed["event"], "agent_start");
+        assert_eq!(parsed["error"], "handler failed");
+    }
+
+    #[test]
+    fn structured_host_extension_error_data_takes_precedence() {
+        let output = map_extension_error_event(&ExtensionErrorEvent {
+            code: "extension_error".to_owned(),
+            message: "legacy summary".to_owned(),
+            retryable: false,
+            data: Some(serde_json::json!({
+                "extensionPath": "/workspace/structured.ts",
+                "event": "tool_result",
+                "error": "structured failure",
+            })),
+        });
+        let parsed: Value = serde_json::from_str(&to_jsonl(&output)).unwrap();
+        assert_eq!(parsed["extensionPath"], "/workspace/structured.ts");
+        assert_eq!(parsed["event"], "tool_result");
+        assert_eq!(parsed["error"], "structured failure");
+    }
+}