Publish harness and TUI open-source

initial sync from the monorepo

crates/codegen/xai-grok-shell/src/agent/activity.rs

@@ -0,0 +1,411 @@
+//! Send-safe view of the agent's in-flight work, shared with the leader's
+//! auto-update checker and `RelaunchForUpdate` drain (`tokio::spawn` tasks
+//! that cannot read the `!Send` `MvpAgent` state on the `LocalSet`).
+//!
+//! The leader's `agent_busy` flag only counts IPC (Unix-socket) requests;
+//! relay (grok.com WebSocket) traffic is bridged straight into the agent's
+//! ACP stdin and never sets it, so a relay-driven leader (devbox / remote)
+//! always looked idle and got restarted mid-turn on every update —
+//! surfacing as "Subagent result channel dropped".
+//!
+//! [`AgentActivity::is_busy`] derives busyness from agent state regardless
+//! of transport, and [`AgentActivity::flush_all_sessions`] lets the shutdown
+//! path end session actors gracefully instead of aborting them via
+//! `LocalSet` drop.
+//!
+//! ## Lifecycle: entries expire with their actor, not with agent bookkeeping
+//!
+//! The agent only ever **registers** sessions (at handle creation). There is
+//! deliberately no unregister: an entry is live exactly while its actor
+//! holds the command receiver (`!cmd_tx.is_closed()`), and closed entries
+//! are purged opportunistically. This sidesteps a whole class of races
+//! between `MvpAgent`'s map bookkeeping and actor lifetime — an actor
+//! removed from the agent's map but still winding down stays visible to
+//! `is_busy`/`flush_all_sessions` until it actually exits, and a session id
+//! rebuilt with a fresh actor is just a second (distinct) entry.
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::Duration;
+
+use crate::session::pending_interaction::PendingInteractions;
+use crate::session::{SessionCommand, SessionHandle};
+
+/// How often [`AgentActivity::flush_all_sessions`] re-polls actors that have
+/// not yet exited.
+const FLUSH_POLL: Duration = Duration::from_millis(50);
+
+/// Per-session slice of state shared with the session actor (the same `Arc`s
+/// the actor mutates — see the matching `SessionHandle` fields).
+struct SessionActivityEntry {
+    id: String,
+    cmd_tx: tokio::sync::mpsc::UnboundedSender<SessionCommand>,
+    /// `Some` while a turn is running (relay- or IPC-driven alike).
+    current_prompt_id: Arc<Mutex<Option<String>>>,
+    /// Non-empty while a blocking reverse-request (permission / question /
+    /// plan approval) is parked.
+    pending_interactions: PendingInteractions,
+}
+
+impl SessionActivityEntry {
+    /// The actor still holds the command receiver.
+    fn is_live(&self) -> bool {
+        !self.cmd_tx.is_closed()
+    }
+
+    /// A running turn or a parked blocking interaction.
+    fn is_busy(&self) -> bool {
+        self.current_prompt_id
+            .lock()
+            .unwrap_or_else(|p| p.into_inner())
+            .is_some()
+            || !self
+                .pending_interactions
+                .lock()
+                .unwrap_or_else(|p| p.into_inner())
+                .is_empty()
+    }
+}
+
+#[derive(Default)]
+struct ActivityInner {
+    /// Self-expiring: entries are dead once the actor drops its receiver
+    /// (see module docs), and are purged whenever the list is locked.
+    sessions: Mutex<Vec<SessionActivityEntry>>,
+    /// Subagents currently initializing or running; kept in sync by
+    /// `SubagentCoordinator::sync_running_gauge`.
+    subagents: Arc<AtomicUsize>,
+}
+
+/// Cheap-to-clone, `Send + Sync` handle. See module docs.
+#[derive(Clone, Default)]
+pub struct AgentActivity {
+    inner: Arc<ActivityInner>,
+}
+
+impl AgentActivity {
+    /// Register a session's shared state at handle-creation time. No
+    /// unregister exists — the entry expires when the actor exits.
+    pub(crate) fn register_session(&self, id: &str, handle: &SessionHandle) {
+        self.lock_live_sessions().push(SessionActivityEntry {
+            id: id.to_string(),
+            cmd_tx: handle.cmd_tx.clone(),
+            current_prompt_id: handle.current_prompt_id.clone(),
+            pending_interactions: handle.pending_interactions.clone(),
+        });
+    }
+
+    /// Shared gauge of initializing + running subagents; handed to the
+    /// `SubagentCoordinator`, which recomputes it on every state change.
+    pub(crate) fn subagent_gauge(&self) -> Arc<AtomicUsize> {
+        self.inner.subagents.clone()
+    }
+
+    /// Whether the agent has live work: a running turn, a parked blocking
+    /// interaction, or an initializing/running subagent.
+    ///
+    /// Known sub-tick window: queued-but-not-started prompts
+    /// (`pending_inputs` in the actor) are not mirrored here, so a prompt
+    /// submitted exactly at a turn boundary can read as idle (the same
+    /// window `session_has_live_work` closes with an actor round-trip,
+    /// which a sync `Send` probe cannot do). The flush's quiesce loop
+    /// re-snapshots and still ends such an actor via its Shutdown arm.
+    pub fn is_busy(&self) -> bool {
+        self.inner.subagents.load(Ordering::Relaxed) > 0
+            || self.lock_live_sessions().iter().any(|e| e.is_busy())
+    }
+
+    /// Number of live registered sessions (diagnostics/tests).
+    pub fn session_count(&self) -> usize {
+        self.lock_live_sessions().len()
+    }
+
+    /// Send [`SessionCommand::Shutdown`] to every live session actor
+    /// (replay-buffer flush → hooks → memory save → actor returns) and wait
+    /// up to `grace` for the actors to exit, observed via
+    /// `cmd_tx.is_closed()`.
+    ///
+    /// This is a quiesce loop, not a one-shot broadcast: each poll
+    /// re-snapshots the registry and signals actors that appeared after the
+    /// flush started (deduped by channel identity, so a session id rebuilt
+    /// with a fresh actor gets its own signal), all against one deadline —
+    /// `grace` bounds the **total** shutdown delay.
+    ///
+    /// Call **before** cancelling the leader's root token so session state
+    /// is durable before the `LocalSet` drop aborts remaining tasks. Actors
+    /// that miss the grace are logged and abandoned.
+    pub async fn flush_all_sessions(&self, grace: Duration) {
+        let deadline = tokio::time::Instant::now() + grace;
+        // Every distinct channel signaled so far (id kept for logging).
+        let mut signaled: Vec<(String, tokio::sync::mpsc::UnboundedSender<SessionCommand>)> =
+            Vec::new();
+
+        loop {
+            let snapshot: Vec<_> = self
+                .lock_live_sessions()
+                .iter()
+                .map(|e| (e.id.clone(), e.cmd_tx.clone()))
+                .collect();
+            for (id, tx) in snapshot {
+                if !signaled.iter().any(|(_, s)| s.same_channel(&tx)) {
+                    tracing::info!(session_id = %id, "leader shutdown: flushing session");
+                    let _ = tx.send(SessionCommand::Shutdown);
+                    signaled.push((id, tx));
+                }
+            }
+
+            if signaled.iter().all(|(_, tx)| tx.is_closed()) {
+                return; // nothing to flush, or all actors exited
+            }
+            if tokio::time::Instant::now() >= deadline {
+                for (id, tx) in &signaled {
+                    if !tx.is_closed() {
+                        tracing::warn!(
+                            session_id = %id,
+                            "leader shutdown: session actor did not exit within grace; proceeding"
+                        );
+                    }
+                }
+                return;
+            }
+            tokio::time::sleep(FLUSH_POLL).await;
+        }
+    }
+
+    /// Lock the session list, dropping entries whose actor has exited.
+    ///
+    /// Purging happens only here, so in modes with no periodic reader (no
+    /// auto-update checker) a dead entry lingers until the next register —
+    /// bounded and tiny (a sender handle + two `Arc`s per entry).
+    fn lock_live_sessions(&self) -> std::sync::MutexGuard<'_, Vec<SessionActivityEntry>> {
+        let mut guard = self
+            .inner
+            .sessions
+            .lock()
+            .unwrap_or_else(|p| p.into_inner());
+        guard.retain(SessionActivityEntry::is_live);
+        guard
+    }
+
+    /// Register a synthetic session from raw parts (no full `SessionHandle`).
+    /// Returns the command receiver (the "actor" side) plus the shared
+    /// running-turn and pending-interaction slots.
+    #[cfg(test)]
+    pub(crate) fn register_for_test(
+        &self,
+        id: &str,
+    ) -> (
+        tokio::sync::mpsc::UnboundedReceiver<SessionCommand>,
+        Arc<Mutex<Option<String>>>,
+        PendingInteractions,
+    ) {
+        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
+        let current_prompt_id = Arc::new(Mutex::new(None));
+        let pending_interactions: PendingInteractions =
+            Arc::new(Mutex::new(std::collections::HashMap::new()));
+        self.lock_live_sessions().push(SessionActivityEntry {
+            id: id.to_string(),
+            cmd_tx,
+            current_prompt_id: current_prompt_id.clone(),
+            pending_interactions: pending_interactions.clone(),
+        });
+        (cmd_rx, current_prompt_id, pending_interactions)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Build a registered entry from raw parts without a full SessionHandle.
+    fn register_raw(
+        activity: &AgentActivity,
+        id: &str,
+    ) -> (
+        tokio::sync::mpsc::UnboundedReceiver<SessionCommand>,
+        Arc<Mutex<Option<String>>>,
+        PendingInteractions,
+    ) {
+        activity.register_for_test(id)
+    }
+
+    /// Simulated session actor: exits (dropping its receiver) `delay` after
+    /// receiving `Shutdown`; resolves to whether Shutdown was received.
+    fn spawn_actor(
+        mut rx: tokio::sync::mpsc::UnboundedReceiver<SessionCommand>,
+        delay: Duration,
+    ) -> tokio::task::JoinHandle<bool> {
+        tokio::spawn(async move {
+            while let Some(cmd) = rx.recv().await {
+                if matches!(cmd, SessionCommand::Shutdown) {
+                    tokio::time::sleep(delay).await;
+                    return true;
+                }
+            }
+            false
+        })
+    }
+
+    #[test]
+    fn idle_by_default() {
+        let activity = AgentActivity::default();
+        assert!(!activity.is_busy());
+    }
+
+    #[tokio::test]
+    async fn running_turn_marks_busy() {
+        let activity = AgentActivity::default();
+        let (_rx, prompt_id, _pending) = register_raw(&activity, "s1");
+        assert!(!activity.is_busy());
+
+        *prompt_id.lock().unwrap() = Some("prompt-1".to_string());
+        assert!(activity.is_busy());
+
+        *prompt_id.lock().unwrap() = None;
+        assert!(!activity.is_busy());
+    }
+
+    #[tokio::test]
+    async fn pending_interaction_marks_busy() {
+        let activity = AgentActivity::default();
+        let (_rx, _prompt_id, pending) = register_raw(&activity, "s1");
+
+        pending.lock().unwrap().insert(
+            "tc-1".to_string(),
+            crate::session::pending_interaction::PendingKind::Permission,
+        );
+        assert!(activity.is_busy());
+
+        pending.lock().unwrap().clear();
+        assert!(!activity.is_busy());
+    }
+
+    #[test]
+    fn subagent_gauge_marks_busy() {
+        let activity = AgentActivity::default();
+        let gauge = activity.subagent_gauge();
+        assert!(!activity.is_busy());
+        gauge.store(1, Ordering::Relaxed);
+        assert!(activity.is_busy());
+        gauge.store(0, Ordering::Relaxed);
+        assert!(!activity.is_busy());
+    }
+
+    /// An actor that is still running counts as busy even if the agent has
+    /// dropped its handle — liveness comes from the channel, not from agent
+    /// bookkeeping. Once the actor exits, the entry expires.
+    #[tokio::test]
+    async fn live_actor_counts_busy_until_it_exits() {
+        let activity = AgentActivity::default();
+        let (rx, prompt_id, _pending) = register_raw(&activity, "s1");
+        *prompt_id.lock().unwrap() = Some("prompt-1".to_string());
+        assert!(activity.is_busy());
+        assert_eq!(activity.session_count(), 1);
+
+        // Actor exits (receiver dropped) → entry expires, even though the
+        // shared prompt slot still says Some.
+        drop(rx);
+        assert!(!activity.is_busy());
+        assert_eq!(activity.session_count(), 0);
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn flush_sends_shutdown_and_waits_for_actor_exit() {
+        let activity = AgentActivity::default();
+        let (rx, _prompt_id, _pending) = register_raw(&activity, "s1");
+
+        // Simulated actor: exits (drops rx) when it receives Shutdown.
+        let actor = spawn_actor(rx, Duration::ZERO);
+
+        activity.flush_all_sessions(Duration::from_secs(5)).await;
+        assert!(actor.await.unwrap(), "actor should have received Shutdown");
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn flush_grace_bounds_total_delay_across_sessions() {
+        let activity = AgentActivity::default();
+        // One wedged actor (receiver kept open) and one healthy actor.
+        let (_wedged_rx, _p1, _i1) = register_raw(&activity, "wedged");
+        let (rx, _p2, _i2) = register_raw(&activity, "healthy");
+        let actor = spawn_actor(rx, Duration::ZERO);
+
+        // The wedged actor must not consume the healthy actor's budget, and
+        // the total wait must be ~one grace period, not one per session.
+        let start = tokio::time::Instant::now();
+        activity.flush_all_sessions(Duration::from_secs(2)).await;
+        let elapsed = start.elapsed();
+        assert!(elapsed >= Duration::from_secs(2));
+        assert!(
+            elapsed < Duration::from_secs(3),
+            "grace must be shared, not serial: {elapsed:?}"
+        );
+        assert!(actor.await.unwrap(), "healthy actor should get Shutdown");
+    }
+
+    /// A session id rebuilt with a fresh actor while the old actor is still
+    /// winding down: both channels must be signaled and awaited.
+    #[tokio::test(start_paused = true)]
+    async fn flush_awaits_both_channels_when_id_is_reused() {
+        let activity = AgentActivity::default();
+        let (old_rx, _p1, _i1) = register_raw(&activity, "s1");
+        let (new_rx, _p2, _i2) = register_raw(&activity, "s1");
+
+        let old_actor = spawn_actor(old_rx, Duration::from_millis(500));
+        let new_actor = spawn_actor(new_rx, Duration::ZERO);
+
+        activity.flush_all_sessions(Duration::from_secs(5)).await;
+        assert!(old_actor.is_finished(), "flush must wait for the old actor");
+        assert!(old_actor.await.unwrap());
+        assert!(new_actor.await.unwrap());
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn flush_signals_sessions_that_appear_mid_flush() {
+        let activity = AgentActivity::default();
+        // Actor 1: holds the flush open for a few polls, then exits.
+        let (rx1, _p1, _i1) = register_raw(&activity, "s1");
+        let actor1 = spawn_actor(rx1, Duration::from_millis(300));
+
+        // Actor 2 registers AFTER the flush has started (a relay-driven
+        // prompt racing the shutdown) — it must still receive Shutdown.
+        let activity_late = activity.clone();
+        let late = tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            let (mut rx2, _p2, _i2) = activity_late.register_for_test("s2");
+            while let Some(cmd) = rx2.recv().await {
+                if matches!(cmd, SessionCommand::Shutdown) {
+                    return true;
+                }
+            }
+            false
+        });
+
+        activity.flush_all_sessions(Duration::from_secs(5)).await;
+        assert!(actor1.await.unwrap());
+        assert!(
+            late.await.unwrap(),
+            "session registered mid-flush must receive Shutdown"
+        );
+    }
+
+    #[tokio::test(start_paused = true)]
+    async fn flush_gives_up_after_grace_when_actor_hangs() {
+        let activity = AgentActivity::default();
+        // Keep rx alive so the channel never closes (wedged actor).
+        let (_rx, _prompt_id, _pending) = register_raw(&activity, "s1");
+
+        let start = tokio::time::Instant::now();
+        activity.flush_all_sessions(Duration::from_secs(2)).await;
+        assert!(
+            start.elapsed() >= Duration::from_secs(2),
+            "flush should wait out the grace period"
+        );
+        // Returned rather than hanging forever — that's the assertion.
+    }
+
+    #[tokio::test]
+    async fn flush_with_no_sessions_is_noop() {
+        let activity = AgentActivity::default();
+        activity.flush_all_sessions(Duration::from_secs(1)).await;
+    }
+}

crates/codegen/xai-grok-shell/src/agent/chat_modes.rs

@@ -0,0 +1,334 @@
+//! grok.com chat-product model catalog: caches `/rest/modes` and maps modes to
+//! the `SessionModelState` returned by `load_chat_session` (the chat analogue of
+//! [`crate::agent::models::ModelsManager`]). NB: these "modes" populate the
+//! desktop MODEL picker, not the ACP session plan-modes in `LoadSessionResponse.modes`.
+use crate::auth::AuthManager;
+use crate::remote::chat_models_client::{
+    ChatModelsClient, ChatModelsError, ListModesResponse, Mode,
+};
+use agent_client_protocol as acp;
+use parking_lot::RwLock;
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+/// ~54 min, matching grok-web's refetch cadence.
+const CACHE_TTL: Duration = Duration::from_secs(54 * 60);
+/// Cold-miss budget on the `session/load` critical path (warm/stale served instantly).
+const COLD_FETCH_TIMEOUT: Duration = Duration::from_secs(2);
+const DEFAULT_LOCALE: &str = "en";
+/// Process-wide flag set by the pager when started with `--chat` so initialize
+/// and early UI seed the chat `/rest/modes` catalog instead of build models.
+pub const GROK_CHAT_MODE_ENV: &str = "GROK_CHAT_MODE";
+/// True when the process is a gateway light-frontend (`--chat`) agent.
+/// Hard-off in release builds so it can't be enabled via env.
+pub fn process_chat_mode_enabled() -> bool {
+    if true {
+        return false;
+    }
+    match std::env::var(GROK_CHAT_MODE_ENV) {
+        Ok(v) => {
+            let v = v.trim();
+            !v.is_empty() && v != "0" && !v.eq_ignore_ascii_case("false")
+        }
+        Err(_) => false,
+    }
+}
+#[derive(Clone)]
+struct CachedModes {
+    /// Keyed by identity; a mismatch is a miss so one user's modes never leak to another.
+    user_id: String,
+    locale: String,
+    fetched_at: Instant,
+    response: ListModesResponse,
+}
+/// Thread-safe, cheaply-cloneable manager. Cloning bumps the inner `Arc`.
+#[derive(Clone)]
+pub struct ChatModesManager {
+    inner: Arc<Inner>,
+}
+struct Inner {
+    auth: Arc<AuthManager>,
+    cache: RwLock<Option<CachedModes>>,
+    /// Single-flight guard so concurrent fetches coalesce.
+    fetch_lock: tokio::sync::Mutex<()>,
+}
+impl ChatModesManager {
+    pub fn new(auth: Arc<AuthManager>) -> Self {
+        Self {
+            inner: Arc::new(Inner {
+                auth,
+                cache: RwLock::new(None),
+                fetch_lock: tokio::sync::Mutex::new(()),
+            }),
+        }
+    }
+    /// The active grok.com identity, or `None` when unauthenticated. Modes are
+    /// per-identity (tier/ACL), so every cache key and store is gated on it.
+    fn current_user_id(&self) -> Option<String> {
+        self.inner.auth.current_or_expired().map(|a| a.user_id)
+    }
+    /// Chat model state for a `session/load` response. On missing auth or fetch
+    /// failure, serves last-good cache else empty — never the build catalog.
+    pub async fn model_state(&self) -> acp::SessionModelState {
+        let Some(user_id) = self.current_user_id() else {
+            return empty_state();
+        };
+        let locale = DEFAULT_LOCALE;
+        {
+            let guard = self.inner.cache.read();
+            if let Some(c) = guard.as_ref()
+                && c.user_id == user_id
+                && c.locale == locale
+            {
+                if c.fetched_at.elapsed() < CACHE_TTL {
+                    return modes_to_model_state(&c.response);
+                }
+                let stale = c.response.clone();
+                drop(guard);
+                self.spawn_refresh(user_id, locale);
+                return modes_to_model_state(&stale);
+            }
+        }
+        let _flight = self.inner.fetch_lock.lock().await;
+        {
+            let guard = self.inner.cache.read();
+            if let Some(c) = guard.as_ref()
+                && c.user_id == user_id
+                && c.locale == locale
+                && c.fetched_at.elapsed() < CACHE_TTL
+            {
+                return modes_to_model_state(&c.response);
+            }
+        }
+        match self.fetch(locale).await {
+            Ok(resp) if !resp.modes.is_empty() => {
+                if self.current_user_id().as_deref() != Some(user_id.as_str()) {
+                    return empty_state();
+                }
+                let mapped = modes_to_model_state(&resp);
+                if mapped.available_models.is_empty() {
+                    tracing::warn!(
+                        raw_modes = resp.modes.len(),
+                        "chat modes: fetch returned modes but none selectable after availability filter"
+                    );
+                }
+                self.store(user_id, locale.to_owned(), resp);
+                mapped
+            }
+            Ok(_) => empty_state(),
+            Err(err) => {
+                tracing::warn!(
+                    error = % err, "chat modes fetch failed; serving cache/empty"
+                );
+                let guard = self.inner.cache.read();
+                match guard.as_ref() {
+                    Some(c) if c.user_id == user_id => modes_to_model_state(&c.response),
+                    _ => empty_state(),
+                }
+            }
+        }
+    }
+    async fn fetch(&self, locale: &str) -> Result<ListModesResponse, ChatModelsError> {
+        let client = ChatModelsClient::new(self.inner.auth.clone());
+        match tokio::time::timeout(COLD_FETCH_TIMEOUT, client.list_modes(locale)).await {
+            Ok(result) => result,
+            Err(_elapsed) => Err(ChatModelsError::Timeout),
+        }
+    }
+    fn store(&self, user_id: String, locale: String, response: ListModesResponse) {
+        *self.inner.cache.write() = Some(CachedModes {
+            user_id,
+            locale,
+            fetched_at: Instant::now(),
+            response,
+        });
+    }
+    /// Best-effort stale refresh; skips if a fetch is already in flight.
+    fn spawn_refresh(&self, user_id: String, locale: &'static str) {
+        let me = self.clone();
+        tokio::spawn(async move {
+            let Ok(_flight) = me.inner.fetch_lock.try_lock() else {
+                return;
+            };
+            if me.current_user_id().as_deref() != Some(user_id.as_str()) {
+                return;
+            }
+            if let Ok(resp) = me.fetch(locale).await
+                && !resp.modes.is_empty()
+                && me.current_user_id().as_deref() == Some(user_id.as_str())
+            {
+                me.store(user_id, locale.to_owned(), resp);
+            }
+        });
+    }
+    /// Kick a background `/rest/modes` fill when auth is already present so
+    /// `--chat` initialize / first `session/new` hit a warm cache.
+    pub fn warm_in_background(&self) {
+        let Some(user_id) = self.current_user_id() else {
+            return;
+        };
+        self.spawn_refresh(user_id, DEFAULT_LOCALE);
+    }
+}
+fn empty_state() -> acp::SessionModelState {
+    acp::SessionModelState::new(acp::ModelId::from(String::new()), Vec::new())
+}
+/// Maps grok.com modes → `SessionModelState`: keeps only `available` modes,
+/// reconciles `current_model_id` (default → first available → empty, never
+/// out-of-set), and stashes `badgeText`/`iconHint`/`tags` in `_meta`.
+pub fn modes_to_model_state(resp: &ListModesResponse) -> acp::SessionModelState {
+    let available_models: Vec<acp::ModelInfo> = resp
+        .modes
+        .iter()
+        .filter(|m| m.is_available())
+        .map(mode_to_model_info)
+        .collect();
+    let current_model_id = reconcile_current(&resp.default_mode_id, &available_models);
+    acp::SessionModelState::new(current_model_id, available_models)
+}
+fn mode_to_model_info(m: &Mode) -> acp::ModelInfo {
+    let name = if m.title.trim().is_empty() {
+        m.id.clone()
+    } else {
+        m.title.clone()
+    };
+    acp::ModelInfo::new(acp::ModelId::from(m.id.clone()), name)
+        .description(if m.description.is_empty() {
+            None
+        } else {
+            Some(m.description.clone())
+        })
+        .meta(build_meta(m))
+}
+fn build_meta(m: &Mode) -> Option<acp::Meta> {
+    let mut map = serde_json::Map::new();
+    if let Some(badge) = m.badge_text.as_deref().filter(|s| !s.is_empty()) {
+        map.insert("badgeText".to_owned(), serde_json::json!(badge));
+    }
+    if !m.icon_hint.is_empty() {
+        map.insert("iconHint".to_owned(), serde_json::json!(m.icon_hint));
+    }
+    if !m.tags.is_empty() {
+        map.insert("tags".to_owned(), serde_json::json!(m.tags));
+    }
+    if map.is_empty() { None } else { Some(map) }
+}
+fn reconcile_current(default_mode_id: &str, available: &[acp::ModelInfo]) -> acp::ModelId {
+    let in_set = |id: &str| available.iter().any(|m| m.model_id.0.as_ref() == id);
+    if !default_mode_id.is_empty() && in_set(default_mode_id) {
+        acp::ModelId::from(default_mode_id.to_owned())
+    } else if let Some(first) = available.first() {
+        first.model_id.clone()
+    } else {
+        acp::ModelId::from(String::new())
+    }
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::remote::chat_models_client::ModeAvailability;
+    fn available(id: &str, title: &str) -> Mode {
+        Mode {
+            id: id.to_owned(),
+            title: title.to_owned(),
+            availability: ModeAvailability {
+                available: Some(serde_json::json!({})),
+                ..Default::default()
+            },
+            ..Default::default()
+        }
+    }
+    fn requires_upgrade(id: &str) -> Mode {
+        Mode {
+            id: id.to_owned(),
+            availability: ModeAvailability {
+                requires_upgrade: Some(serde_json::json!({ "message" : "Upgrade" })),
+                ..Default::default()
+            },
+            ..Default::default()
+        }
+    }
+    #[test]
+    fn filters_to_available_modes() {
+        let resp = ListModesResponse {
+            modes: vec![
+                available("auto", "Auto"),
+                requires_upgrade("heavy"),
+                available("fast", "Fast"),
+            ],
+            default_mode_id: "auto".to_owned(),
+        };
+        let state = modes_to_model_state(&resp);
+        let ids: Vec<String> = state
+            .available_models
+            .iter()
+            .map(|m| m.model_id.0.to_string())
+            .collect();
+        assert_eq!(ids, vec!["auto".to_string(), "fast".to_string()]);
+        assert_eq!(state.current_model_id.0.as_ref(), "auto");
+    }
+    #[test]
+    fn default_outside_filtered_set_falls_back_to_first_available() {
+        let resp = ListModesResponse {
+            modes: vec![requires_upgrade("heavy"), available("fast", "Fast")],
+            default_mode_id: "heavy".to_owned(),
+        };
+        let state = modes_to_model_state(&resp);
+        assert_eq!(state.current_model_id.0.as_ref(), "fast");
+        assert!(
+            state
+                .available_models
+                .iter()
+                .any(|m| m.model_id == state.current_model_id)
+        );
+    }
+    #[test]
+    fn empty_default_falls_back_to_first() {
+        let resp = ListModesResponse {
+            modes: vec![available("a", "A"), available("b", "B")],
+            default_mode_id: String::new(),
+        };
+        let state = modes_to_model_state(&resp);
+        assert_eq!(state.current_model_id.0.as_ref(), "a");
+    }
+    #[test]
+    fn no_available_modes_yields_empty_current() {
+        let resp = ListModesResponse {
+            modes: vec![requires_upgrade("heavy")],
+            default_mode_id: "heavy".to_owned(),
+        };
+        let state = modes_to_model_state(&resp);
+        assert!(state.available_models.is_empty());
+        assert_eq!(state.current_model_id.0.as_ref(), "");
+    }
+    #[test]
+    fn maps_fields_and_meta() {
+        let mut m = available("auto", "Auto");
+        m.description = "Picks the best model".to_owned();
+        m.badge_text = Some("New".to_owned());
+        m.icon_hint = "rocket".to_owned();
+        m.tags = vec!["TAG_PRIMARY".to_owned()];
+        let resp = ListModesResponse {
+            modes: vec![m],
+            default_mode_id: "auto".to_owned(),
+        };
+        let state = modes_to_model_state(&resp);
+        let info = &state.available_models[0];
+        assert_eq!(info.name, "Auto");
+        assert_eq!(info.description.as_deref(), Some("Picks the best model"));
+        let meta = info.meta.as_ref().unwrap();
+        assert_eq!(meta["badgeText"], serde_json::json!("New"));
+        assert_eq!(meta["iconHint"], serde_json::json!("rocket"));
+        assert_eq!(meta["tags"], serde_json::json!(["TAG_PRIMARY"]));
+    }
+    #[test]
+    fn name_falls_back_to_id_when_title_blank() {
+        let mut m = available("grok-4.5", "");
+        m.title = "   ".to_owned();
+        let resp = ListModesResponse {
+            modes: vec![m],
+            default_mode_id: String::new(),
+        };
+        let state = modes_to_model_state(&resp);
+        assert_eq!(state.available_models[0].name, "grok-4.5");
+    }
+}

crates/codegen/xai-grok-shell/src/agent/config_model_override_parse.rs

@@ -0,0 +1,639 @@
+//! Resilient parsing for `[model.<id>]` TOML overrides.
+//!
+//! A model entry must survive a bad field: warn and skip the field, never
+//! drop the model (managed configs must not lose catalog entries).
+//!
+//! Every table is deserialized through `serde_ignored`, so unknown fields
+//! warn on every path and [`ConfigModelOverride`] stays the single source of
+//! truth for the field set. When the whole-table parse fails, fields that
+//! fail to parse on their own are pruned (one warning each) and the table is
+//! parsed again. Non-table values are dropped with a warning.
+//!
+//! Warnings are retained on `Config::model_override_warnings` and surfaced by
+//! `grok inspect`.
+
+use indexmap::IndexMap;
+use serde::Serialize;
+
+use super::config::ConfigModelOverride;
+
+/// Category for a [`ModelOverrideWarning`].
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
+#[serde(rename_all = "kebab-case")]
+pub enum ModelOverrideWarningKind {
+    /// Field name not recognized; field ignored.
+    UnknownField,
+    /// Value failed to parse; field skipped.
+    InvalidValue,
+    /// Legacy alias given alongside its canonical key; alias skipped.
+    DuplicateAlias,
+    /// Entry value is not a TOML table; entry dropped.
+    NotATable,
+    /// Entry failed to parse even after skipping invalid fields; the model
+    /// keeps an empty override.
+    UnparseableEntry,
+}
+
+/// One skipped field or dropped entry from `[model.*]` parsing.
+#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ModelOverrideWarning {
+    /// `None` when the warning is about the `[model]` section itself.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model_key: Option<String>,
+    /// `None` for warnings about the entry as a whole.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub field: Option<String>,
+    pub kind: ModelOverrideWarningKind,
+    pub reason: String,
+}
+
+/// Result of [`parse_model_overrides`].
+pub(crate) struct ParsedModelOverrides {
+    pub models: IndexMap<String, ConfigModelOverride>,
+    pub warnings: Vec<ModelOverrideWarning>,
+}
+
+/// Parses every `[model.<id>]` entry in `raw_config`, returning the overrides
+/// and a warning for each skipped field or dropped entry.
+pub(crate) fn parse_model_overrides(raw_config: &toml::Value) -> ParsedModelOverrides {
+    let mut models = IndexMap::new();
+    let mut warnings = Vec::new();
+    let Some(section) = raw_config.get("model") else {
+        return ParsedModelOverrides { models, warnings };
+    };
+    let Some(table) = section.as_table() else {
+        warnings.push(ModelOverrideWarning {
+            model_key: None,
+            field: None,
+            kind: ModelOverrideWarningKind::NotATable,
+            reason: format!(
+                "`model` must be a table of [model.<id>] entries, got {}; all model overrides ignored",
+                section.type_str()
+            ),
+        });
+        return ParsedModelOverrides { models, warnings };
+    };
+    for (model_key, value) in table {
+        let Some(entry_table) = value.as_table() else {
+            warnings.push(ModelOverrideWarning {
+                model_key: Some(model_key.clone()),
+                field: None,
+                kind: ModelOverrideWarningKind::NotATable,
+                reason: format!(
+                    "expected a table like [model.\"{model_key}\"], got {}; entry dropped",
+                    value.type_str()
+                ),
+            });
+            continue;
+        };
+        let (entry, entry_warnings) = parse_model_override_table(model_key, entry_table.clone());
+        warnings.extend(entry_warnings);
+        models.insert(model_key.clone(), entry);
+    }
+    ParsedModelOverrides { models, warnings }
+}
+
+/// Logs the warnings when they differ from the previous parse, so a
+/// persistently broken config logs once per process instead of once per parse.
+pub(crate) fn log_model_override_warnings(warnings: &[ModelOverrideWarning]) {
+    use std::hash::{Hash as _, Hasher as _};
+    use std::sync::atomic::{AtomicU64, Ordering};
+
+    static LAST_LOGGED: AtomicU64 = AtomicU64::new(0);
+    // 0 means "no warnings"; real hashes are clamped to nonzero.
+    let hash = if warnings.is_empty() {
+        0
+    } else {
+        let mut hasher = std::collections::hash_map::DefaultHasher::new();
+        warnings.hash(&mut hasher);
+        hasher.finish().max(1)
+    };
+    if LAST_LOGGED.swap(hash, Ordering::Relaxed) == hash {
+        return;
+    }
+
+    for warning in warnings {
+        tracing::warn!(
+            model = warning.model_key.as_deref().unwrap_or("(section)"),
+            field = warning.field.as_deref().unwrap_or("(entry)"),
+            kind = ?warning.kind,
+            reason = %warning.reason,
+            "model_override: skipped invalid config"
+        );
+    }
+    if !warnings.is_empty() {
+        tracing::warn!(
+            warnings = warnings.len(),
+            "model_override: parsed with warnings; run `grok inspect` for details"
+        );
+    }
+}
+
+fn parse_model_override_table(
+    model_key: &str,
+    mut table: toml::map::Map<String, toml::Value>,
+) -> (ConfigModelOverride, Vec<ModelOverrideWarning>) {
+    let mut warnings = Vec::new();
+    dedupe_aliases(model_key, &mut table, &mut warnings);
+
+    // Unknown-field warnings come from whichever parse produces the returned
+    // entry, so both paths report them identically.
+    match deserialize_with_unknown_fields(table.clone()) {
+        Ok((entry, unknown)) => {
+            warnings.extend(unknown_field_warnings(model_key, unknown));
+            (entry, warnings)
+        }
+        Err(_) => {
+            prune_invalid_fields(model_key, &mut table, &mut warnings);
+            match deserialize_with_unknown_fields(table) {
+                Ok((entry, unknown)) => {
+                    warnings.extend(unknown_field_warnings(model_key, unknown));
+                    (entry, warnings)
+                }
+                Err(error) => {
+                    // Reachable only when fields conflict jointly, e.g. an
+                    // alias pair missing from `ALIASES`. Keep the model
+                    // rather than dropping it.
+                    warnings.push(ModelOverrideWarning {
+                        model_key: Some(model_key.to_owned()),
+                        field: None,
+                        kind: ModelOverrideWarningKind::UnparseableEntry,
+                        reason: format!(
+                            "failed to parse after skipping invalid fields ({error}); using empty override"
+                        ),
+                    });
+                    (ConfigModelOverride::default(), warnings)
+                }
+            }
+        }
+    }
+}
+
+/// `(canonical, legacy)` key pairs that serde rejects as duplicate fields
+/// when both appear in one table. Keep in sync with the `#[serde(alias)]`
+/// attributes on [`ConfigModelOverride`].
+const ALIASES: &[(&str, &str)] = &[("compactions_remaining", "send_compactions_remaining")];
+
+/// Removes one key of each [`ALIASES`] pair that appears twice in `table`.
+/// The canonical key wins; when its value doesn't parse, the legacy key is
+/// kept instead.
+fn dedupe_aliases(
+    model_key: &str,
+    table: &mut toml::map::Map<String, toml::Value>,
+    warnings: &mut Vec<ModelOverrideWarning>,
+) {
+    for &(canonical, legacy) in ALIASES {
+        if !(table.contains_key(canonical) && table.contains_key(legacy)) {
+            continue;
+        }
+        match field_parse_error(canonical, &table[canonical]) {
+            None => {
+                table.remove(legacy);
+                warnings.push(ModelOverrideWarning {
+                    model_key: Some(model_key.to_owned()),
+                    field: Some(legacy.to_owned()),
+                    kind: ModelOverrideWarningKind::DuplicateAlias,
+                    reason: format!("legacy alias of {canonical}; skipped in favor of {canonical}"),
+                });
+            }
+            Some(error) => {
+                table.remove(canonical);
+                warnings.push(ModelOverrideWarning {
+                    model_key: Some(model_key.to_owned()),
+                    field: Some(canonical.to_owned()),
+                    kind: ModelOverrideWarningKind::InvalidValue,
+                    reason: format!("{error}; skipped in favor of {legacy}"),
+                });
+            }
+        }
+    }
+}
+
+/// Deserializes `table`, also returning the unknown field names that serde
+/// would otherwise silently discard.
+fn deserialize_with_unknown_fields(
+    table: toml::map::Map<String, toml::Value>,
+) -> Result<(ConfigModelOverride, Vec<String>), toml::de::Error> {
+    let mut unknown = Vec::new();
+    let entry = serde_ignored::deserialize(toml::Value::Table(table), |path| {
+        unknown.push(path.to_string());
+    })?;
+    Ok((entry, unknown))
+}
+
+fn unknown_field_warnings(model_key: &str, unknown: Vec<String>) -> Vec<ModelOverrideWarning> {
+    unknown
+        .into_iter()
+        .map(|field| ModelOverrideWarning {
+            model_key: Some(model_key.to_owned()),
+            field: Some(field),
+            kind: ModelOverrideWarningKind::UnknownField,
+            reason: "unknown field".to_owned(),
+        })
+        .collect()
+}
+
+/// Removes each field that fails to parse on its own, one warning per field.
+/// Unknown fields stay; the follow-up parse reports them.
+fn prune_invalid_fields(
+    model_key: &str,
+    table: &mut toml::map::Map<String, toml::Value>,
+    warnings: &mut Vec<ModelOverrideWarning>,
+) {
+    table.retain(|field, value| match field_parse_error(field, value) {
+        None => true,
+        Some(error) => {
+            warnings.push(ModelOverrideWarning {
+                model_key: Some(model_key.to_owned()),
+                field: Some(field.to_owned()),
+                kind: ModelOverrideWarningKind::InvalidValue,
+                reason: error.to_string(),
+            });
+            false
+        }
+    });
+}
+
+/// Parses `field` in isolation, returning the error if it fails.
+fn field_parse_error(field: &str, value: &toml::Value) -> Option<toml::de::Error> {
+    let mut singleton = toml::map::Map::new();
+    singleton.insert(field.to_owned(), value.clone());
+    toml::Value::Table(singleton)
+        .try_into::<ConfigModelOverride>()
+        .err()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::sampling::ApiBackend;
+    use xai_grok_sampling_types::{
+        CompactionAtTokens, CompactionsRemaining, ReasoningEffort, ReasoningEffortOption,
+    };
+
+    fn parse_cfg(toml_str: &str) -> crate::agent::config::Config {
+        let raw: toml::Value = toml::from_str(toml_str).unwrap();
+        crate::agent::config::Config::new_from_toml_cfg(&raw).expect("config should parse")
+    }
+
+    fn parse_raw(
+        toml_str: &str,
+    ) -> (
+        IndexMap<String, ConfigModelOverride>,
+        Vec<ModelOverrideWarning>,
+    ) {
+        let raw: toml::Value = toml::from_str(toml_str).unwrap();
+        let ParsedModelOverrides { models, warnings } = parse_model_overrides(&raw);
+        (models, warnings)
+    }
+
+    #[test]
+    fn duplicate_compactions_keys_keeps_model() {
+        let cfg = parse_cfg(
+            r#"
+            [model."grok-4.5"]
+            model = "grok-4.5"
+            env_key = "ANTHROPIC_AUTH_TOKEN"
+            compactions_remaining = 1
+            send_compactions_remaining = true
+            "#,
+        );
+        let model = cfg
+            .config_models
+            .get("grok-4.5")
+            .expect("grok-4.5 must remain in catalog");
+        assert_eq!(
+            model.compactions_remaining,
+            Some(CompactionsRemaining::Fixed(1))
+        );
+        assert!(cfg.model_override_warnings.iter().any(|w| {
+            w.kind == ModelOverrideWarningKind::DuplicateAlias
+                && w.field.as_deref() == Some("send_compactions_remaining")
+        }));
+        let resolved = crate::agent::config::resolve_model_list(&cfg, None);
+        assert!(resolved.contains_key("grok-4.5"));
+    }
+
+    #[test]
+    fn legacy_alias_alone_parses_without_warning() {
+        let cfg = parse_cfg(
+            r#"
+            [model."grok-4.5"]
+            model = "grok-4.5"
+            send_compactions_remaining = 2
+            "#,
+        );
+        let model = cfg.config_models.get("grok-4.5").unwrap();
+        assert_eq!(
+            model.compactions_remaining,
+            Some(CompactionsRemaining::Fixed(2))
+        );
+        assert!(cfg.model_override_warnings.is_empty());
+    }
+
+    #[test]
+    fn invalid_reasoning_effort_skips_field_keeps_model() {
+        let cfg = parse_cfg(
+            r#"
+            [model."grok-4.5"]
+            model = "grok-4.5"
+            env_key = "ANTHROPIC_AUTH_TOKEN"
+            reasoning_effort = "not-a-level"
+            "#,
+        );
+        let model = cfg
+            .config_models
+            .get("grok-4.5")
+            .expect("grok-4.5 must remain in catalog");
+        assert_eq!(model.model.as_deref(), Some("grok-4.5"));
+        assert!(model.reasoning_effort.is_none());
+        assert!(cfg.model_override_warnings.iter().any(|w| {
+            w.kind == ModelOverrideWarningKind::InvalidValue
+                && w.field.as_deref() == Some("reasoning_effort")
+        }));
+    }
+
+    #[test]
+    fn unknown_field_warns_but_keeps_known_fields() {
+        let (models, warnings) = parse_raw(
+            r#"
+            [model."grok-4.5"]
+            model = "grok-4.5"
+            env_key = "TOKEN"
+            future_field = 1
+            "#,
+        );
+        let entry = models.get("grok-4.5").unwrap();
+        assert_eq!(entry.model.as_deref(), Some("grok-4.5"));
+        assert_eq!(
+            entry.env_key.as_ref().and_then(|k| k.primary()),
+            Some("TOKEN")
+        );
+        assert_eq!(
+            warnings,
+            vec![ModelOverrideWarning {
+                model_key: Some("grok-4.5".to_owned()),
+                field: Some("future_field".to_owned()),
+                kind: ModelOverrideWarningKind::UnknownField,
+                reason: "unknown field".to_owned(),
+            }]
+        );
+    }
+
+    /// An unknown field warns the same whether or not another field fails to
+    /// parse.
+    #[test]
+    fn unknown_field_warning_is_path_independent() {
+        let unknown_of = |toml_str: &str| {
+            let (_, warnings) = parse_raw(toml_str);
+            warnings
+                .into_iter()
+                .filter(|w| w.kind == ModelOverrideWarningKind::UnknownField)
+                .collect::<Vec<_>>()
+        };
+        let fast = unknown_of(
+            r#"
+            [model.m]
+            temprature = 0.5
+            "#,
+        );
+        let slow = unknown_of(
+            r#"
+            [model.m]
+            temprature = 0.5
+            reasoning_effort = "not-a-level"
+            "#,
+        );
+        assert_eq!(fast, slow);
+        assert_eq!(fast.len(), 1);
+        assert_eq!(fast[0].field.as_deref(), Some("temprature"));
+    }
+
+    #[test]
+    fn prune_skips_invalid_fields_and_keeps_the_rest() {
+        // A valid nested table survives an invalid sibling.
+        let (models, warnings) = parse_raw(
+            r#"
+            [model.m]
+            temperature = "hot"
+            [model.m.extra_headers]
+            x-team = "codegen"
+            "#,
+        );
+        let entry = models.get("m").unwrap();
+        assert_eq!(
+            entry.extra_headers.get("x-team").map(String::as_str),
+            Some("codegen")
+        );
+        assert!(entry.temperature.is_none());
+        assert!(warnings.iter().any(|w| {
+            w.kind == ModelOverrideWarningKind::InvalidValue
+                && w.field.as_deref() == Some("temperature")
+        }));
+
+        // All fields invalid: the model stays, with an empty override.
+        let (models, warnings) = parse_raw(
+            r#"
+            [model.m]
+            temperature = "hot"
+            max_retries = "many"
+            "#,
+        );
+        let entry = models.get("m").expect("model must remain in catalog");
+        assert!(entry.temperature.is_none());
+        assert!(entry.max_retries.is_none());
+        assert_eq!(warnings.len(), 2);
+        assert!(
+            warnings
+                .iter()
+                .all(|w| w.kind == ModelOverrideWarningKind::InvalidValue)
+        );
+    }
+
+    #[test]
+    fn invalid_canonical_key_falls_back_to_legacy_alias() {
+        let (models, warnings) = parse_raw(
+            r#"
+            [model.m]
+            compactions_remaining = "bad"
+            send_compactions_remaining = 2
+            "#,
+        );
+        let entry = models.get("m").unwrap();
+        assert_eq!(
+            entry.compactions_remaining,
+            Some(CompactionsRemaining::Fixed(2))
+        );
+        assert_eq!(warnings.len(), 1);
+        assert_eq!(warnings[0].kind, ModelOverrideWarningKind::InvalidValue);
+        assert_eq!(warnings[0].field.as_deref(), Some("compactions_remaining"));
+    }
+
+    #[test]
+    fn non_table_model_section_warns_and_is_ignored() {
+        let (models, warnings) = parse_raw(r#"model = "grok-4""#);
+        assert!(models.is_empty());
+        assert_eq!(warnings.len(), 1);
+        assert_eq!(warnings[0].kind, ModelOverrideWarningKind::NotATable);
+        assert_eq!(warnings[0].model_key, None);
+        assert_eq!(warnings[0].field, None);
+    }
+
+    #[test]
+    fn non_table_entry_is_dropped_with_warning() {
+        let (models, warnings) = parse_raw(
+            r#"
+            [model]
+            oops = 5
+            "#,
+        );
+        assert!(models.is_empty(), "a scalar cannot define a model");
+        assert_eq!(warnings.len(), 1);
+        assert_eq!(warnings[0].kind, ModelOverrideWarningKind::NotATable);
+        assert_eq!(warnings[0].model_key.as_deref(), Some("oops"));
+        assert_eq!(warnings[0].field, None);
+    }
+
+    /// Exhaustive literal (no `..`): a new struct field is a compile error
+    /// here until the drift-guard tests cover it.
+    fn fully_populated_override() -> ConfigModelOverride {
+        ConfigModelOverride {
+            model: Some("m".into()),
+            base_url: Some("https://example.com".into()),
+            name: Some("Model M".into()),
+            description: Some("desc".into()),
+            api_key: Some("key".into()),
+            env_key: Some(crate::agent::config::EnvKeys::single("ENV_KEY")),
+            api_base_url: Some("https://api.example.com".into()),
+            max_completion_tokens: Some(1024),
+            temperature: Some(0.5),
+            top_p: Some(0.9),
+            api_backend: Some(ApiBackend::Messages),
+            extra_headers: [("x-team".to_owned(), "codegen".to_owned())]
+                .into_iter()
+                .collect(),
+            context_window: Some(200_000),
+            auto_compact_threshold_percent: Some(80),
+            system_prompt_label: Some("label".into()),
+            use_concise: Some(true),
+            agent_type: Some("agent".into()),
+            inference_idle_timeout_secs: Some(60),
+            max_retries: Some(3),
+            hidden: Some(false),
+            supported_in_api: Some(true),
+            reasoning_effort: Some(ReasoningEffort::High),
+            supports_reasoning_effort: Some(true),
+            reasoning_efforts: vec![ReasoningEffortOption {
+                id: "deep".to_string(),
+                value: ReasoningEffort::High,
+                label: "Deep".to_string(),
+                description: Some("Deep reasoning".to_string()),
+                default: true,
+            }],
+            supports_backend_search: Some(false),
+            compactions_remaining: Some(CompactionsRemaining::Fixed(1)),
+            compaction_at_tokens: Some(CompactionAtTokens::Fixed(100_000)),
+            show_model_fingerprint: Some(true),
+            stream_tool_calls: Some(false),
+        }
+    }
+
+    fn parse_single_entry(
+        entry: toml::map::Map<String, toml::Value>,
+    ) -> (
+        IndexMap<String, ConfigModelOverride>,
+        Vec<ModelOverrideWarning>,
+    ) {
+        let mut model_table = toml::map::Map::new();
+        model_table.insert("m".to_owned(), toml::Value::Table(entry));
+        let mut root = toml::map::Map::new();
+        root.insert("model".to_owned(), toml::Value::Table(model_table));
+        let ParsedModelOverrides { models, warnings } =
+            parse_model_overrides(&toml::Value::Table(root));
+        (models, warnings)
+    }
+
+    #[test]
+    fn fully_populated_override_round_trips_without_warnings() {
+        let serialized = toml::Value::try_from(fully_populated_override()).unwrap();
+        let (models, warnings) = parse_single_entry(serialized.as_table().unwrap().clone());
+        assert_eq!(warnings, Vec::new(), "no field may be skipped or unknown");
+        let reparsed = toml::Value::try_from(models.get("m").unwrap()).unwrap();
+        assert_eq!(reparsed, serialized, "round-trip must be lossless");
+    }
+
+    /// Drift guard: every `#[serde(alias)]` on [`ConfigModelOverride`] must
+    /// have a matching `ALIASES` pair, and vice versa. An unregistered alias
+    /// would send both-keys configs to the empty-override fallback.
+    #[test]
+    fn every_struct_alias_is_registered_in_aliases() {
+        let source = include_str!("config.rs");
+        let start = source
+            .find("pub struct ConfigModelOverride {")
+            .expect("ConfigModelOverride definition in config.rs");
+        let block = &source[start..];
+        let block = &block[..block.find("\n}").expect("struct end")];
+
+        let mut found = Vec::new();
+        let mut rest = block;
+        while let Some(pos) = rest.find("#[serde(alias = \"") {
+            let after = &rest[pos + "#[serde(alias = \"".len()..];
+            let legacy = &after[..after.find('"').expect("closing quote")];
+            let field = &after[after.find("pub ").expect("field after alias") + 4..];
+            let canonical = &field[..field.find(':').expect("field type colon")];
+            found.push((canonical.to_owned(), legacy.to_owned()));
+            rest = after;
+        }
+        assert_eq!(
+            block.matches("alias").count(),
+            found.len(),
+            "an alias on ConfigModelOverride was not recognized; write it as \
+             `#[serde(alias = \"...\")]` on its own line, or update this scan"
+        );
+        found.sort();
+
+        let mut registered: Vec<(String, String)> = ALIASES
+            .iter()
+            .map(|&(c, l)| (c.to_owned(), l.to_owned()))
+            .collect();
+        registered.sort();
+        assert_eq!(
+            found, registered,
+            "#[serde(alias)] attributes on ConfigModelOverride and ALIASES must match"
+        );
+    }
+
+    /// Drift guard for `ALIASES`, in both directions: every pair must be a
+    /// real serde alias (a both-keys table fails a plain parse), and the
+    /// parser must resolve it to the canonical key with a single warning.
+    #[test]
+    fn every_aliases_pair_is_a_real_serde_alias_and_dedupes() {
+        let reference = toml::Value::try_from(fully_populated_override()).unwrap();
+        for &(canonical, legacy) in ALIASES {
+            let value = reference
+                .get(canonical)
+                .unwrap_or_else(|| panic!("{canonical} missing from fully_populated_override"));
+            let mut entry = toml::map::Map::new();
+            entry.insert(canonical.to_owned(), value.clone());
+            entry.insert(legacy.to_owned(), value.clone());
+            assert!(
+                toml::Value::Table(entry.clone())
+                    .try_into::<ConfigModelOverride>()
+                    .is_err(),
+                "{canonical}/{legacy} is not a serde alias pair; remove it from ALIASES"
+            );
+
+            let (models, warnings) = parse_single_entry(entry);
+            let parsed = toml::Value::try_from(models.get("m").unwrap()).unwrap();
+            assert_eq!(
+                parsed.get(canonical),
+                Some(value),
+                "canonical value must be retained"
+            );
+            assert_eq!(warnings.len(), 1);
+            assert_eq!(warnings[0].kind, ModelOverrideWarningKind::DuplicateAlias);
+            assert_eq!(warnings[0].field.as_deref(), Some(legacy));
+        }
+    }
+}

crates/codegen/xai-grok-shell/src/agent/ext_parsers.rs

@@ -0,0 +1,259 @@
+//! Wire-shape parsers for ext-notification params handled by `MvpAgent`.
+//!
+//! Pure parsing only (params JSON → `SessionCommand`); session lookup and
+//! command dispatch stay in `mvp_agent::ext_notification`.
+
+use crate::session::SessionCommand;
+
+/// Parse a `x.ai/queue/{remove,reorder,clear,edit,interject}` ext-notification's
+/// params into the corresponding [`SessionCommand`].
+/// `owner` is the resolved attribution (params `owner`/`clientIdentifier`) used
+/// to scope remove/clear to the requesting client's own items, and recorded as
+/// `last_editor` for in-place text edits. Returns `None` for unrecognized
+/// methods or for `edit` when `newText` is missing.
+pub(super) fn parse_queue_edit_command(
+    method: &str,
+    params: &serde_json::Value,
+    owner: Option<String>,
+) -> Option<SessionCommand> {
+    match method {
+        "x.ai/queue/remove" => {
+            let id = params.get("id").and_then(|v| v.as_str())?.to_string();
+            // The client supplies the version it last saw; the handler removes
+            // only on an exact match (stale = benign no-op + rebroadcast).
+            // Default 0 covers never-edited prompts (the common case).
+            let expected_version = params
+                .get("expectedVersion")
+                .and_then(|v| v.as_u64())
+                .unwrap_or(0);
+            Some(SessionCommand::RemoveQueuedPrompt {
+                id,
+                expected_version,
+                owner,
+            })
+        }
+        "x.ai/queue/reorder" => {
+            let ordered_ids = params
+                .get("orderedIds")
+                .and_then(|v| v.as_array())
+                .map(|arr| {
+                    arr.iter()
+                        .filter_map(|v| v.as_str().map(|s| s.to_string()))
+                        .collect::<Vec<_>>()
+                })
+                .unwrap_or_default();
+            Some(SessionCommand::ReorderQueue { ordered_ids })
+        }
+        "x.ai/queue/clear" => Some(SessionCommand::ClearQueue { owner }),
+        "x.ai/queue/interject" => {
+            let id = params.get("id").and_then(|v| v.as_str())?.to_string();
+            // The client supplies the version it last saw; the handler acts
+            // only on an exact match (stale = benign no-op + rebroadcast).
+            let expected_version = params
+                .get("expectedVersion")
+                .and_then(|v| v.as_u64())
+                .unwrap_or(0);
+            // Optional client-edited replacement text (atomic edit+interject).
+            // Blank overrides are dropped (degrade to the stored queue text) —
+            // never interject an empty prompt on a malformed client param.
+            let new_text = params
+                .get("newText")
+                .and_then(|v| v.as_str())
+                .filter(|s| !s.trim().is_empty())
+                .map(str::to_string);
+            Some(SessionCommand::InterjectQueuedPrompt {
+                id,
+                expected_version,
+                owner,
+                new_text,
+            })
+        }
+        "x.ai/queue/edit" => {
+            let id = params.get("id").and_then(|v| v.as_str())?.to_string();
+            let new_text = params.get("newText").and_then(|v| v.as_str())?.to_string();
+            // `owner` is the resolved attribution; for edit it represents the
+            // most recent editor (recorded as `last_editor`), not the original
+            // enqueuer.
+            Some(SessionCommand::EditQueuedPrompt {
+                id,
+                new_text,
+                editor: owner,
+            })
+        }
+        _ => None,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Each `x.ai/queue/*` ext-notification maps to the
+    /// correct versioned/idempotent `SessionCommand`.
+    #[test]
+    fn parse_queue_edit_command_maps_each_method() {
+        // remove: id + expectedVersion + owner.
+        let p = serde_json::json!({
+            "sessionId": "s1", "id": "p7", "expectedVersion": 3
+        });
+        match parse_queue_edit_command("x.ai/queue/remove", &p, Some("grok-tui".into())) {
+            Some(SessionCommand::RemoveQueuedPrompt {
+                id,
+                expected_version,
+                owner,
+            }) => {
+                assert_eq!(id, "p7");
+                assert_eq!(expected_version, 3);
+                assert_eq!(owner.as_deref(), Some("grok-tui"));
+            }
+            _ => panic!("expected RemoveQueuedPrompt"),
+        }
+
+        // remove without expectedVersion defaults to 0.
+        let p = serde_json::json!({ "sessionId": "s1", "id": "p8" });
+        match parse_queue_edit_command("x.ai/queue/remove", &p, None) {
+            Some(SessionCommand::RemoveQueuedPrompt {
+                expected_version, ..
+            }) => assert_eq!(expected_version, 0),
+            _ => panic!("expected RemoveQueuedPrompt"),
+        }
+
+        // reorder: orderedIds array.
+        let p = serde_json::json!({ "sessionId": "s1", "orderedIds": ["a", "b", "c"] });
+        match parse_queue_edit_command("x.ai/queue/reorder", &p, None) {
+            Some(SessionCommand::ReorderQueue { ordered_ids }) => {
+                assert_eq!(ordered_ids, vec!["a", "b", "c"]);
+            }
+            _ => panic!("expected ReorderQueue"),
+        }
+
+        // clear: owner-scoped.
+        match parse_queue_edit_command(
+            "x.ai/queue/clear",
+            &serde_json::json!({ "sessionId": "s1" }),
+            Some("grok-tui".into()),
+        ) {
+            Some(SessionCommand::ClearQueue { owner }) => {
+                assert_eq!(owner.as_deref(), Some("grok-tui"));
+            }
+            _ => panic!("expected ClearQueue"),
+        }
+
+        // edit: id + newText + editor (resolved via owner/clientIdentifier).
+        let p = serde_json::json!({
+            "sessionId": "s1", "id": "p9", "newText": "replacement text"
+        });
+        match parse_queue_edit_command("x.ai/queue/edit", &p, Some("grok-vscode".into())) {
+            Some(SessionCommand::EditQueuedPrompt {
+                id,
+                new_text,
+                editor,
+            }) => {
+                assert_eq!(id, "p9");
+                assert_eq!(new_text, "replacement text");
+                assert_eq!(editor.as_deref(), Some("grok-vscode"));
+            }
+            _ => panic!("expected EditQueuedPrompt"),
+        }
+
+        // edit without editor (no owner/clientIdentifier) → editor: None.
+        match parse_queue_edit_command(
+            "x.ai/queue/edit",
+            &serde_json::json!({ "sessionId": "s1", "id": "p9", "newText": "x" }),
+            None,
+        ) {
+            Some(SessionCommand::EditQueuedPrompt { editor, .. }) => {
+                assert!(editor.is_none());
+            }
+            _ => panic!("expected EditQueuedPrompt"),
+        }
+
+        // edit without newText → None (can't replace text we don't have).
+        assert!(
+            parse_queue_edit_command(
+                "x.ai/queue/edit",
+                &serde_json::json!({ "sessionId": "s1", "id": "p9" }),
+                None,
+            )
+            .is_none()
+        );
+
+        // edit without id → None (can't target an entry).
+        assert!(
+            parse_queue_edit_command(
+                "x.ai/queue/edit",
+                &serde_json::json!({ "sessionId": "s1", "newText": "x" }),
+                None,
+            )
+            .is_none()
+        );
+
+        // interject: id + expectedVersion + owner (mirrors remove).
+        let p = serde_json::json!({
+            "sessionId": "s1", "id": "p10", "expectedVersion": 2
+        });
+        match parse_queue_edit_command("x.ai/queue/interject", &p, Some("grok-tui".into())) {
+            Some(SessionCommand::InterjectQueuedPrompt {
+                id,
+                expected_version,
+                owner,
+                new_text,
+            }) => {
+                assert_eq!(id, "p10");
+                assert_eq!(expected_version, 2);
+                assert_eq!(owner.as_deref(), Some("grok-tui"));
+                assert_eq!(new_text, None, "newText absent → None");
+            }
+            _ => panic!("expected InterjectQueuedPrompt"),
+        }
+
+        // interject with newText (client-edited row) carries the override.
+        let p = serde_json::json!({
+            "sessionId": "s1", "id": "p10", "expectedVersion": 2, "newText": "edited"
+        });
+        match parse_queue_edit_command("x.ai/queue/interject", &p, None) {
+            Some(SessionCommand::InterjectQueuedPrompt { new_text, .. }) => {
+                assert_eq!(new_text.as_deref(), Some("edited"));
+            }
+            _ => panic!("expected InterjectQueuedPrompt"),
+        }
+
+        // Blank newText is dropped → degrades to the stored queue text.
+        let p = serde_json::json!({
+            "sessionId": "s1", "id": "p10", "expectedVersion": 2, "newText": "   "
+        });
+        match parse_queue_edit_command("x.ai/queue/interject", &p, None) {
+            Some(SessionCommand::InterjectQueuedPrompt { new_text, .. }) => {
+                assert_eq!(new_text, None, "blank override must be dropped");
+            }
+            _ => panic!("expected InterjectQueuedPrompt"),
+        }
+
+        // interject without expectedVersion defaults to 0.
+        match parse_queue_edit_command(
+            "x.ai/queue/interject",
+            &serde_json::json!({ "sessionId": "s1", "id": "p11" }),
+            None,
+        ) {
+            Some(SessionCommand::InterjectQueuedPrompt {
+                expected_version, ..
+            }) => assert_eq!(expected_version, 0),
+            _ => panic!("expected InterjectQueuedPrompt"),
+        }
+
+        // interject without id → None (can't target an entry).
+        assert!(
+            parse_queue_edit_command("x.ai/queue/interject", &serde_json::json!({}), None)
+                .is_none()
+        );
+
+        // unknown method → None.
+        assert!(
+            parse_queue_edit_command("x.ai/queue/bogus", &serde_json::json!({}), None).is_none()
+        );
+        // remove without id → None (can't target an entry).
+        assert!(
+            parse_queue_edit_command("x.ai/queue/remove", &serde_json::json!({}), None).is_none()
+        );
+    }
+}

crates/codegen/xai-grok-shell/src/agent/handlers/mod.rs

@@ -0,0 +1,3 @@
+pub(crate) mod model_switch;
+pub(crate) mod session;
+pub(crate) mod workspaces;

crates/codegen/xai-grok-shell/src/agent/handlers/model_switch.rs

@@ -0,0 +1,262 @@
+//! Applies a model switch to a session — the ungated path. `set_session_model`
+//! enforces the `allowed_models` gate before delegating here; internal callers
+//! (`new_session`, `load_session`) call `apply` directly.
+use crate::agent::config;
+use crate::agent::mvp_agent::{
+    MvpAgent, agent_name_after_model_switch, harnesses_are_compatible, resolve_required_agent_type,
+};
+use crate::session::SessionCommand;
+use agent_client_protocol::{self as acp};
+use tokio::sync::oneshot;
+use xai_grok_sampling_types::parse_reasoning_effort_meta;
+/// Apply a model switch to a session (no gate — `set_session_model` gates first).
+pub(crate) async fn apply(
+    agent: &MvpAgent,
+    args: acp::SetSessionModelRequest,
+) -> Result<acp::SetSessionModelResponse, acp::Error> {
+    tracing::info!("Received set session model request {args:?}");
+    xai_grok_telemetry::unified_log::info(
+        "model changed",
+        Some(args.session_id.0.as_ref()),
+        Some(serde_json::json!({ "model" : args.model_id.0.as_ref() })),
+    );
+    tracing::debug!("session_session_model::mvp_agent: {:?}", &args);
+    let effort_override = parse_reasoning_effort_meta(args.meta.as_ref());
+    let acp::SetSessionModelRequest {
+        session_id,
+        model_id,
+        ..
+    } = args;
+    let handle = agent
+        .session_handle_waiting_for_load(&session_id)
+        .await
+        .ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?;
+    let model = agent.resolve_model_id(&model_id)?;
+    let use_concise = model.info().use_concise;
+    let session_default = handle
+        .session_default_agent_profile
+        .as_deref()
+        .unwrap_or(&handle.agent_name);
+    let required_agent_type =
+        resolve_required_agent_type(Some(model.info().agent_type.as_str()), session_default);
+    let previous_model_id = handle.model_id.0.clone();
+    let mut pending_rebuild_definition: Option<xai_grok_agent::AgentDefinition> = None;
+    {
+        let required = &required_agent_type;
+        let turn_count = handle
+            .signals_handle
+            .snapshot()
+            .await
+            .map(|s| s.turn_count)
+            .unwrap_or(0);
+        let (agent_tx, agent_rx) = oneshot::channel();
+        let _ = handle.cmd_tx.send(SessionCommand::GetActiveAgent {
+            responds_to: agent_tx,
+        });
+        let active_agent_type = agent_rx.await.ok().flatten();
+        let is_mismatch = active_agent_type
+            .as_ref()
+            .is_some_and(|active| !harnesses_are_compatible(active, required));
+        tracing::info!(
+            session_id = % session_id.0, model_id = % model_id.0, ? required_agent_type,
+            ? active_agent_type, turn_count, is_mismatch,
+            "set_session_model: agent type compatibility check"
+        );
+        if is_mismatch && turn_count > 0 {
+            tracing::warn!(
+                session_id = % session_id.0, model_id = % model_id.0, active_agent = ?
+                active_agent_type, required_agent = % required, turn_count,
+                "set_session_model: agent type mismatch rejected"
+            );
+            xai_grok_telemetry::session_ctx::log_event(xai_grok_telemetry::events::ModelSwitched {
+                session_id: session_id.0.to_string(),
+                previous_model_id: previous_model_id.to_string(),
+                new_model_id: model_id.0.to_string(),
+                success: false,
+                error_code: Some(config::MODEL_SWITCH_INCOMPATIBLE_AGENT.to_string()),
+                required_agent_type: Some(required.clone()),
+                current_agent_type: active_agent_type.clone(),
+            });
+            let err_payload = config::ModelSwitchIncompatibleAgentError {
+                code: config::MODEL_SWITCH_INCOMPATIBLE_AGENT.to_string(),
+                active_agent_type: active_agent_type.unwrap_or_else(|| "unknown".to_owned()),
+                required_agent_type: required.clone(),
+                model_id: model_id.0.to_string(),
+                suggestion: "start_new_session".to_string(),
+            };
+            return Err(err_payload.into_acp_error());
+        }
+        if is_mismatch && turn_count == 0 {
+            let cwd = handle.tool_context.cwd.as_path();
+            let resolved = xai_grok_agent::discovery::by_name_in_cwd_with_plugins(
+                required,
+                cwd,
+                agent.plugin_registry_handle.snapshot().as_deref(),
+            );
+            match resolved {
+                Some(def) => {
+                    tracing::info!(
+                        session_id = % session_id.0, model_id = % model_id.0,
+                        required_agent_type = % required, agent_def_name = % def.name,
+                        "set_session_model: zero-turn harness switch — queued agent rebuild"
+                    );
+                    pending_rebuild_definition = Some(def);
+                }
+                None => {
+                    tracing::warn!(
+                        session_id = % session_id.0, model_id = % model_id.0,
+                        required_agent_type = % required,
+                        "set_session_model: zero-turn harness switch — could not resolve agent definition; proceeding with stale harness"
+                    );
+                }
+            }
+        }
+    }
+    let mut model_sampling =
+        agent.prepare_sampling_config_for_model(&model, handle.origin_client.clone());
+    if let Some(eff) = effort_override {
+        if agent
+            .models_manager
+            .model_supports_reasoning_effort(model_id.0.as_ref())
+        {
+            tracing::info!(
+                session_id = % session_id.0, effort = % eff,
+                "set_session_model: applying reasoning_effort override from meta"
+            );
+            model_sampling.reasoning_effort = Some(eff);
+        } else {
+            tracing::warn!(
+                session_id = % session_id.0, model_id = % model_id.0, effort = % eff,
+                "set_session_model: ignoring reasoning_effort override — model does not support it"
+            );
+        }
+    }
+    let applied_effort = model_sampling.reasoning_effort;
+    let gate_closed = !handle
+        .gateway_enabled
+        .load(std::sync::atomic::Ordering::Relaxed);
+    let apply_prompt_override = !gate_closed;
+    if gate_closed {
+        tracing::info!(
+            session_id = % session_id.0, model_id = % model_id.0,
+            "set_session_model: gateway gate closed, prompt override suppressed"
+        );
+        pending_rebuild_definition = None;
+    }
+    let did_rebuild = if let Some(def) = pending_rebuild_definition {
+        let (rebuild_tx, rebuild_rx) = oneshot::channel();
+        let _ = handle
+            .cmd_tx
+            .send(SessionCommand::RebuildAgentForDefinition {
+                definition: def,
+                responds_to: rebuild_tx,
+            });
+        let rebuild_result = rebuild_rx
+            .await
+            .map_err(|_| acp::Error::internal_error().data("rebuild_agent: actor closed"))?;
+        match rebuild_result {
+            Ok(()) => true,
+            Err(e) => {
+                tracing::error!(
+                    session_id = % session_id.0, model_id = % model_id.0, error = ? e,
+                    "set_session_model: zero-turn harness rebuild failed; aborting model switch"
+                );
+                xai_grok_telemetry::session_ctx::log_event(
+                    xai_grok_telemetry::events::ModelSwitched {
+                        session_id: session_id.0.to_string(),
+                        previous_model_id: previous_model_id.to_string(),
+                        new_model_id: model_id.0.to_string(),
+                        success: false,
+                        error_code: Some(config::MODEL_SWITCH_REBUILD_FAILED.to_string()),
+                        required_agent_type: Some(required_agent_type.clone()),
+                        current_agent_type: None,
+                    },
+                );
+                return Err(e);
+            }
+        }
+    } else {
+        false
+    };
+    let model_unchanged = previous_model_id == model_id.0;
+    let new_threshold = {
+        let cfg = agent.cfg.borrow();
+        let models = agent.models_manager.models();
+        let model = config::find_model_by_id(&models, model_sampling.model.as_str());
+        crate::util::config::resolve_auto_compact_threshold_percent(
+            &cfg,
+            model_sampling.model.as_str(),
+            model.map(|e| &e.info),
+        )
+    };
+    let (tx, rx) = oneshot::channel();
+    let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel {
+        sampling_config: model_sampling,
+        use_concise,
+        apply_prompt_override,
+        skip_prompt_rewrite: did_rebuild || model_unchanged,
+        auto_compact_threshold_percent: new_threshold,
+        responds_to: tx,
+    });
+    let updated_model = rx
+        .await
+        .map_err(|_| acp::Error::internal_error().data("failed to set session model"))?;
+    if let Some(handle) = agent.sessions.borrow_mut().get_mut(&session_id) {
+        handle.model_id = model_id.clone();
+        handle.reasoning_effort = applied_effort;
+        handle.agent_name =
+            agent_name_after_model_switch(did_rebuild, &required_agent_type, &handle.agent_name);
+    }
+    broadcast_model_changed(
+        agent,
+        &session_id,
+        model_id.0.as_ref(),
+        applied_effort.map(|eff| eff.to_string()),
+    );
+    xai_grok_telemetry::session_ctx::log_event(xai_grok_telemetry::events::ModelSwitched {
+        session_id: session_id.0.to_string(),
+        previous_model_id: previous_model_id.to_string(),
+        new_model_id: model_id.0.to_string(),
+        success: true,
+        error_code: None,
+        required_agent_type: Some(required_agent_type.clone()),
+        current_agent_type: None,
+    });
+    if agent.cfg.borrow().mode != config::AgentMode::Leader {
+        agent.models_manager.set_current_model_id(model_id);
+        agent
+            .models_manager
+            .set_current_reasoning_effort(applied_effort);
+    }
+    Ok(acp::SetSessionModelResponse::new().meta(
+        serde_json::json!({ "model" : updated_model, })
+            .as_object()
+            .cloned(),
+    ))
+}
+/// Broadcast a `ModelChanged` to every client subscribed to this session so
+/// followers mirror the new model. The originating client ignores its own echo
+/// (gated by `model_switch_pending`). Broadcast-only — no eventId, not persisted.
+fn broadcast_model_changed(
+    agent: &MvpAgent,
+    session_id: &acp::SessionId,
+    model_id: &str,
+    reasoning_effort: Option<String>,
+) {
+    let notification = crate::extensions::notification::SessionNotification {
+        session_id: session_id.clone(),
+        update: crate::extensions::notification::SessionUpdate::ModelChanged {
+            model_id: model_id.to_owned(),
+            reasoning_effort,
+        },
+        meta: None,
+    };
+    if let Ok(params) = serde_json::value::to_raw_value(&notification) {
+        agent
+            .gateway
+            .forward_fire_and_forget(acp::ExtNotification::new(
+                "x.ai/session_notification",
+                params.into(),
+            ));
+    }
+}

crates/codegen/xai-grok-shell/src/agent/handlers/session.rs

@@ -0,0 +1,291 @@
+//! Session meta-information handlers.
+//!
+//! Router pattern: single `handle()` dispatches by method name.
+//! Business logic delegates to pure functions or MvpAgent methods.
+
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+
+use agent_client_protocol::{self as acp};
+use serde::Deserialize;
+
+use super::super::mvp_agent::MvpAgent;
+use crate::session::persistence::{Summary, list_recent_summaries, list_summaries};
+use crate::session::{
+    AllSessionOverviewRequest, AllSessionOverviewResponse, ContextInfo, ExtMethodResult,
+    SessionCommand, SessionInfoData, SessionInfoResponse, SessionListRequest, SessionListResponse,
+};
+
+/// Mirrors the display title (`generated_title`, else `session_summary`) into
+/// `session_summary` so clients that only read that field show the same title
+/// as `display_title()` — including after a `/rename` that updated only
+/// `generated_title`. Mutates the response copy only; never persisted.
+fn backfill_session_summary(summary: &mut Summary) {
+    let display = summary.display_title().to_owned();
+    if !display.is_empty() && display != summary.session_summary {
+        summary.session_summary = display;
+    }
+}
+
+/// Router for x.ai/session/* and x.ai/session_summaries/* methods.
+pub async fn handle(
+    agent: &MvpAgent,
+    args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    match args.method.as_ref() {
+        "x.ai/session/info" => handle_session_info(agent, args).await,
+        "x.ai/session/close" => handle_session_close(agent, args).await,
+        "x.ai/session/list" => handle_session_list(agent, args).await,
+        "x.ai/sessions/list" => handle_roster_list(agent, args).await,
+        m if m.starts_with("x.ai/session_summaries/") => {
+            handle_session_summaries(agent, args).await
+        }
+        _ => Err(acp::Error::method_not_found()),
+    }
+}
+
+/// `x.ai/sessions/list` — the FleetView roster. Returns every
+/// resident session plus recently-touched on-disk `Dormant` sessions. Clients
+/// poll this while the dashboard is open and reconcile against the
+/// `x.ai/sessions/changed` broadcast.
+async fn handle_roster_list(
+    agent: &MvpAgent,
+    _args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    let sessions = agent.build_roster().await;
+    ExtMethodResult::success(crate::agent::roster::RosterListResponse { sessions })
+        .to_ext_response()
+        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct SessionInfoRequest {
+    session_id: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct RecentSessionsRequest {
+    limit: usize,
+}
+
+async fn handle_session_info(
+    agent: &MvpAgent,
+    args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    let req: SessionInfoRequest = serde_json::from_str(args.params.get())
+        .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
+
+    let session_id = req.session_id.or_else(|| {
+        agent
+            .sessions
+            .borrow()
+            .keys()
+            .next()
+            .map(|id| id.0.to_string())
+    });
+
+    let Some(session_id) = session_id else {
+        return ExtMethodResult::success(serde_json::json!({}))
+            .to_ext_response()
+            .map_err(|e| acp::Error::internal_error().data(e.to_string()));
+    };
+
+    let sid = acp::SessionId::new(session_id.clone());
+    let Some(session) = agent.sessions.borrow().get(&sid).cloned() else {
+        return ExtMethodResult::success(serde_json::json!({}))
+            .to_ext_response()
+            .map_err(|e| acp::Error::internal_error().data(e.to_string()));
+    };
+
+    let (tx, rx) = tokio::sync::oneshot::channel();
+    let _ = session
+        .cmd_tx
+        .send(SessionCommand::GetSessionInfo { responds_to: tx });
+    let info = rx.await.ok();
+
+    // Construct display data for `/session-info`.
+    let mut data = info.unwrap_or_else(|| SessionInfoData {
+        agent_name: None,
+        model: None,
+        model_display_name: None,
+        resolved_model_id: None,
+        model_fingerprint: None,
+        show_model_fingerprint: false,
+        api_backend: None,
+        conversation_id: None,
+        turns: 0,
+        turn_index: 0,
+        context: ContextInfo {
+            auto_compact_threshold_percent:
+                crate::util::config::DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT,
+            ..ContextInfo::default()
+        },
+    });
+
+    // Calculate the model's display name.
+    data.model_display_name = agent
+        .models_manager
+        .models()
+        .get(session.model_id.0.as_ref())
+        .and_then(|entry| entry.info.name.clone());
+
+    // Construct `SessionInfoResponse`.
+    let response = SessionInfoResponse {
+        session_id,
+        cwd: session.info.cwd.clone(),
+        data,
+    };
+
+    // Wrap `SessionInfoResponse` in `ExtMethodResult` and return it.
+    ExtMethodResult::success(serde_json::to_value(&response).unwrap_or_default())
+        .to_ext_response()
+        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
+}
+
+async fn handle_session_close(
+    agent: &MvpAgent,
+    args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    #[derive(Deserialize)]
+    #[serde(rename_all = "camelCase")]
+    struct CloseRequest {
+        session_id: String,
+    }
+
+    let req: CloseRequest = serde_json::from_str(args.params.get())
+        .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
+
+    let sid = acp::SessionId::new(req.session_id.clone());
+    let existed = agent.sessions.borrow().contains_key(&sid);
+    if existed {
+        // Explicit terminal close: shut the actor down and finalize the cloud
+        // replica (genuine session end). Distinct from a mere client disconnect,
+        // which detaches but keeps the session resumable and never finalizes
+        // (see `MvpAgent::handle_evict_sessions` / `close_session_explicit`).
+        agent.request_session_shutdown(&sid);
+        agent.close_session_explicit(&sid);
+        tracing::info!(session_id = %req.session_id, "session closed via x.ai/session/close");
+    } else {
+        tracing::debug!(session_id = %req.session_id, "session/close: session not found (already closed)");
+    }
+
+    ExtMethodResult::success(serde_json::json!({ "success": true }))
+        .to_ext_response()
+        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
+}
+
+async fn handle_session_summaries(
+    _agent: &MvpAgent,
+    args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    match args.method.as_ref() {
+        "x.ai/session_summaries/session_list" => {
+            let req = serde_json::from_str::<SessionListRequest>(args.params.get())?;
+            let cwd = req.workspace_directory.to_string_lossy().to_string();
+
+            let _timer = crate::instrumentation_timer!("session.list_sessions_for_workspace");
+
+            let mut summaries = list_summaries(Some(&cwd)).await.map_err(|e| {
+                acp::Error::internal_error().data(format!("failed to list sessions: {e}"))
+            })?;
+            for s in &mut summaries {
+                backfill_session_summary(s);
+            }
+
+            let value = serde_json::to_value(SessionListResponse {
+                session_summaries: summaries,
+            })
+            .map(|v| serde_json::value::to_raw_value(&v).map(Arc::from))
+            .expect("to work")
+            .expect("to work");
+
+            Ok(acp::ExtResponse::new(value))
+        }
+        "x.ai/session_summaries/workspace_list" => {
+            tracing::debug!("xai/session_summaries/workspace_list is working");
+            let _req = serde_json::from_str::<AllSessionOverviewRequest>(args.params.get())?;
+
+            let _timer = crate::instrumentation_timer!("session.list_sessions_for_load");
+
+            let summaries = list_summaries(None).await.map_err(|e| {
+                acp::Error::internal_error().data(format!("failed to list workspaces: {e}"))
+            })?;
+
+            summaries_to_overview_response(summaries)
+        }
+        "x.ai/session_summaries/workspace_list_recent" => {
+            let req = serde_json::from_str::<RecentSessionsRequest>(args.params.get())?;
+
+            let _timer = crate::instrumentation_timer!("session.list_sessions_recent");
+
+            let limit = req.limit.min(10_000);
+            let mut summaries = list_recent_summaries(limit).await.map_err(|e| {
+                acp::Error::internal_error().data(format!("failed to list workspaces: {e}"))
+            })?;
+            for s in &mut summaries {
+                backfill_session_summary(s);
+            }
+
+            let value = serde_json::to_value(&summaries)
+                .map(|v| serde_json::value::to_raw_value(&v).map(Arc::from))
+                .expect("to work")
+                .expect("to work");
+
+            Ok(acp::ExtResponse::new(value))
+        }
+        _ => Err(acp::Error::method_not_found()),
+    }
+}
+
+/// Group summaries by cwd and serialize into an [`AllSessionOverviewResponse`].
+fn summaries_to_overview_response(summaries: Vec<Summary>) -> Result<acp::ExtResponse, acp::Error> {
+    let mut by_cwd: BTreeMap<String, Vec<Summary>> = Default::default();
+    for mut s in summaries {
+        backfill_session_summary(&mut s);
+        by_cwd.entry(s.info.cwd.clone()).or_default().push(s);
+    }
+
+    let value = serde_json::to_value(AllSessionOverviewResponse {
+        all_sessions: by_cwd
+            .into_iter()
+            .map(|(k, v)| (PathBuf::from(k), v))
+            .collect(),
+    })
+    .map(|v| serde_json::value::to_raw_value(&v).map(Arc::from))
+    .expect("to work")
+    .expect("to work");
+
+    Ok(acp::ExtResponse::new(value))
+}
+// ── Merged session list (local + remote) ─────────────────────────────
+
+async fn handle_session_list(
+    agent: &MvpAgent,
+    args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    use crate::session::unified_list;
+
+    // Under chat mode `parse_list_req` REPLACES any client-sent `kind` facet
+    // (never union) so every list surface is conversations-only.
+    let req = unified_list::parse_list_req(args.params.get())
+        .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
+    tracing::debug!(
+        chat_mode_forced_kind = crate::agent::chat_modes::process_chat_mode_enabled(),
+        "session/list"
+    );
+
+    let registry_client = agent.session_registry_client();
+    let conversations_client = agent.conversations_client();
+    let result = unified_list::build_unified_list(
+        registry_client.as_ref(),
+        conversations_client.as_ref(),
+        req,
+    )
+    .await;
+
+    ExtMethodResult::success(unified_list::ext_list_response(result))
+        .to_ext_response()
+        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
+}

crates/codegen/xai-grok-shell/src/agent/handlers/workspaces.rs

@@ -0,0 +1,174 @@
+use agent_client_protocol::{self as acp};
+use serde::{Deserialize, Serialize};
+
+use super::super::mvp_agent::MvpAgent;
+use crate::remote::{ListWorkspacesPage, WsError, WsQuery};
+use crate::session::ExtMethodResult;
+
+const DEFAULT_PAGE_SIZE: i64 = 50;
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct WorkspacesListRequest {
+    #[serde(default)]
+    page_size: Option<i64>,
+    #[serde(default)]
+    page_token: Option<String>,
+    #[serde(default)]
+    query: Option<String>,
+    #[serde(default)]
+    kind: Option<String>,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+struct WorkspaceRow {
+    id: String,
+    name: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    kind: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    create_time: Option<String>,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+struct WorkspacesListResponse {
+    workspaces: Vec<WorkspaceRow>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    next_page_token: Option<String>,
+    #[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
+    meta: Option<WorkspacesMeta>,
+}
+
+#[derive(Debug, Serialize)]
+struct WorkspacesMeta {
+    #[serde(rename = "x.ai/partial")]
+    partial: PartialInfo,
+}
+
+#[derive(Debug, Serialize)]
+struct PartialInfo {
+    workspaces: bool,
+    reason: &'static str,
+}
+
+pub async fn handle(
+    agent: &MvpAgent,
+    args: &acp::ExtRequest,
+) -> Result<acp::ExtResponse, acp::Error> {
+    let req: WorkspacesListRequest = serde_json::from_str(args.params.get())
+        .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
+
+    let q = WsQuery {
+        // Clamp to a sane positive page size: a missing, zero, or negative
+        // `pageSize` falls back to the default rather than being forwarded
+        // verbatim to `/rest/workspaces`.
+        page_size: match req.page_size {
+            Some(n) if n > 0 => n,
+            _ => DEFAULT_PAGE_SIZE,
+        },
+        page_token: req.page_token,
+        query: req.query,
+        kind: req.kind,
+    };
+
+    let response = match agent.workspaces_client().list_workspaces(&q).await {
+        Ok(page) => success_response(page),
+        Err(WsError::NoOauth) => degraded_response("no_oauth"),
+        Err(e) => {
+            // Degrade to a partial result, but don't silently swallow the
+            // cause — log it so field failures are diagnosable.
+            tracing::warn!("workspaces/list fetch failed: {e}");
+            degraded_response("error")
+        }
+    };
+
+    ExtMethodResult::success(response)
+        .to_ext_response()
+        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
+}
+
+fn success_response(page: ListWorkspacesPage) -> WorkspacesListResponse {
+    WorkspacesListResponse {
+        workspaces: page
+            .workspaces
+            .into_iter()
+            .map(|w| WorkspaceRow {
+                id: w.workspace_id,
+                name: w.name,
+                kind: w.kind,
+                create_time: w.create_time,
+            })
+            .collect(),
+        next_page_token: page.next_page_token,
+        meta: None,
+    }
+}
+
+fn degraded_response(reason: &'static str) -> WorkspacesListResponse {
+    WorkspacesListResponse {
+        workspaces: Vec::new(),
+        next_page_token: None,
+        meta: Some(WorkspacesMeta {
+            partial: PartialInfo {
+                workspaces: true,
+                reason,
+            },
+        }),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::remote::Workspace;
+
+    #[test]
+    fn request_parses_camelcase_and_defaults_page_size() {
+        let req: WorkspacesListRequest =
+            serde_json::from_value(serde_json::json!({})).expect("empty params parse");
+        assert!(req.page_size.is_none());
+
+        let req: WorkspacesListRequest = serde_json::from_value(serde_json::json!({
+            "pageSize": 10,
+            "pageToken": "tok",
+            "query": "gpu",
+            "kind": "WORKSPACE_KIND_IMAGINE"
+        }))
+        .expect("full params parse");
+        assert_eq!(req.page_size, Some(10));
+        assert_eq!(req.page_token.as_deref(), Some("tok"));
+        assert_eq!(req.query.as_deref(), Some("gpu"));
+        assert_eq!(req.kind.as_deref(), Some("WORKSPACE_KIND_IMAGINE"));
+    }
+
+    #[test]
+    fn success_response_projects_grok_workspace_fields() {
+        let page = ListWorkspacesPage {
+            workspaces: vec![Workspace {
+                workspace_id: "ws_1".into(),
+                name: "Research".into(),
+                create_time: Some("2026-06-18T17:30:00Z".into()),
+                kind: Some("WORKSPACE_KIND_IMAGINE".into()),
+            }],
+            next_page_token: Some("tok2".into()),
+        };
+        let value = serde_json::to_value(success_response(page)).unwrap();
+        assert_eq!(value["workspaces"][0]["id"], "ws_1");
+        assert_eq!(value["workspaces"][0]["name"], "Research");
+        assert_eq!(value["workspaces"][0]["kind"], "WORKSPACE_KIND_IMAGINE");
+        assert_eq!(value["workspaces"][0]["createTime"], "2026-06-18T17:30:00Z");
+        assert_eq!(value["nextPageToken"], "tok2");
+        assert!(value.get("_meta").is_none());
+    }
+
+    #[test]
+    fn degraded_response_carries_partial_reason() {
+        let value = serde_json::to_value(degraded_response("no_oauth")).unwrap();
+        assert_eq!(value["workspaces"].as_array().unwrap().len(), 0);
+        assert!(value.get("nextPageToken").is_none());
+        assert_eq!(value["_meta"]["x.ai/partial"]["workspaces"], true);
+        assert_eq!(value["_meta"]["x.ai/partial"]["reason"], "no_oauth");
+    }
+}

crates/codegen/xai-grok-shell/src/agent/init.rs

@@ -0,0 +1,198 @@
+//! Agent bootstrap and lifecycle hooks.
+//!
+//! [`bootstrap`] runs the full init sequence (config resolution, process
+//! singletons, model catalog) and returns a resolved config + `ModelsManager`.
+//! [`update_telemetry_config`] re-initializes telemetry after auth changes.
+
+use std::sync::Arc;
+
+use indexmap::IndexMap;
+
+use crate::agent::config::{self, Config as AgentConfig, ModelEntry};
+use crate::agent::models::ModelsManager;
+use crate::auth::AuthManager;
+use crate::config::StorageMode;
+
+/// Resolve config, init process singletons, build the model catalog.
+///
+/// The `ModelsManager` is `Clone + Send`, so callers that need a handle
+/// for the config watcher can clone it before passing it to
+/// `MvpAgent::with_models`.
+pub fn bootstrap(
+    cfg: &AgentConfig,
+    auth_manager: &Arc<AuthManager>,
+    prefetched: Option<IndexMap<String, ModelEntry>>,
+) -> Result<(AgentConfig, ModelsManager), String> {
+    // Fail closed before any policy is read: a tampered managed policy must not run unmanaged.
+    crate::managed_config::managed_policy_gate()?;
+    let cfg = resolve_config(cfg, auth_manager);
+    cfg.validate_model_filters()?;
+    init_process(&cfg, auth_manager);
+    let models_manager = ModelsManager::from_config(&cfg, prefetched, auth_manager.clone())?;
+
+    // Refresh on every auth refresh — the FSEvents watcher can silently die after
+    // macOS sleep, stranding the catalog on bundled defaults.
+    models_manager.start_auth_refresh_watcher(auth_manager.refresh_notifier());
+
+    Ok((cfg, models_manager))
+}
+
+/// Print a `bootstrap`/`MvpAgent::new` config error and exit (process boundary).
+///
+/// Restores native stderr first: a managed-policy refusal on the ACP/server path reaches here
+/// while fd 2 may still point at the `/dev/null` the TUI's `redirect_native_stderr()` set, which
+/// would swallow the message. No-op when stderr was never redirected (headless).
+pub(crate) fn exit_on_config_error<T>(e: String) -> T {
+    xai_tty_utils::restore_native_stderr();
+    eprintln!("\nConfiguration error:\n\n    {e}\n");
+    std::process::exit(1);
+}
+
+/// Config transform: apply managed settings, fetch remote settings,
+/// resolve storage mode.
+fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig {
+    let mut cfg = cfg.clone();
+
+    if let Ok(layers) = crate::config::ConfigLayers::load()
+        && layers.has_managed()
+    {
+        let origins = crate::config::config_origins(&layers);
+        let managed_keys: Vec<&str> = origins
+            .iter()
+            .filter(|(_, s)| matches!(s, config::ConfigSource::ManagedConfig))
+            .map(|(k, _)| k.as_str())
+            .collect();
+        if !managed_keys.is_empty() {
+            tracing::info!(keys = ?managed_keys, "managed_config.toml fields");
+        }
+    }
+
+    let managed_enforced = crate::config::apply_managed_settings_features(&mut cfg);
+    let requirements_enforced = crate::config::apply_requirements(&mut cfg);
+
+    for e in managed_enforced.iter().chain(&requirements_enforced) {
+        tracing::info!(field = %e.path, value = %e.value, source = %e.source, "policy override");
+    }
+
+    // Fallback: if the client didn't pre-supply remote settings, fetch them
+    // now so remote-settings-gated features work regardless of which client
+    // spawned us.  Clients that already call `start_early_prefetch()` and
+    // thread the result into `cfg.remote_settings` skip this entirely.
+    if cfg.remote_settings.is_none()
+        && let Some(handle) =
+            crate::agent::models::start_early_prefetch(Some(cfg.grok_com_config.clone()))
+    {
+        match handle.join() {
+            Ok(result) => {
+                cfg.remote_settings = result.settings;
+                crate::util::config::set_remote_campaigns_from_settings(
+                    cfg.remote_settings.as_ref(),
+                );
+                tracing::info!("remote_settings fetched as shell-level fallback");
+            }
+            Err(_) => {
+                tracing::warn!("remote_settings fallback prefetch thread panicked");
+            }
+        }
+    }
+    crate::util::config::sync_campaign_fields(&mut cfg);
+    crate::agent::config::apply_remote_settings_side_effects(cfg.remote_settings.as_ref());
+
+    // env var > remote settings > Local. Skip remote settings for Generic (grok -p, subagents).
+    if cfg.storage_mode == StorageMode::Local
+        && cfg.mode != crate::agent::config::AgentMode::Generic
+    {
+        cfg.storage_mode = StorageMode::resolve(None, cfg.remote_settings.as_ref());
+    }
+    // Writeback talks to the code backend; requires grok.com auth.
+    if cfg.storage_mode == StorageMode::Writeback
+        && !auth_manager.current().is_some_and(|a| a.is_xai_auth())
+    {
+        tracing::info!("Writeback is disabled: requires auth with grok.com");
+        cfg.storage_mode = StorageMode::Local;
+    }
+
+    if let Some(rs) = cfg.remote_settings.as_ref()
+        && let Some(v) = rs.path_not_found_hints
+    {
+        cfg.path_not_found_hints = v;
+    }
+
+    cfg
+}
+
+/// Initialize process-level singletons (deployment sync, bundled files,
+/// telemetry). `Once`-guarded: only the first call takes effect.
+/// Telemetry user ID is updated separately via [`update_telemetry_config`].
+fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
+    use std::sync::Once;
+    static INIT: Once = Once::new();
+    INIT.call_once(|| {
+        if !cfg!(test) {
+            // Clear a logged-out team's files before the background sync runs.
+            crate::managed_config::clear_orphan();
+            crate::managed_config::spawn_sync(tokio_util::sync::CancellationToken::new());
+        }
+
+        let grok_home = crate::util::grok_home::grok_home();
+        crate::builtin::extract_bundled_files(&grok_home);
+
+        // Auto-register is gated (default off; env/remote settings enables). Kept out
+        // of extract_bundled_files so the gate can read the resolved
+        // remote_settings, which resolve_config has populated by now.
+        if cfg.resolve_official_marketplace_auto_register().value {
+            crate::extensions::marketplace::ensure_official_marketplace_source(&grok_home);
+        }
+
+        let telemetry_mode = cfg.resolve_telemetry_mode();
+        let trace_upload = cfg.resolve_trace_upload();
+        let feedback = cfg.resolve_feedback();
+        let feedback_url = cfg.endpoints.resolve_feedback_base_url();
+        let trace_upload_url = cfg.endpoints.resolve_trace_upload_url();
+        tracing::info!(
+            telemetry = %telemetry_mode,
+            trace_upload = %trace_upload,
+            feedback = %feedback,
+            feedback_url = %feedback_url,
+            feedback_url_custom = cfg.endpoints.feedback_base_url.is_some(),
+            trace_upload_url = %trace_upload_url,
+            trace_upload_url_custom = cfg.endpoints.trace_upload_url.is_some(),
+            trace_upload_bucket = cfg.endpoints.trace_upload_bucket.as_deref().unwrap_or("none"),
+            trace_upload_region = cfg.endpoints.trace_upload_region.as_deref().unwrap_or("none"),
+            "data capture config resolved",
+        );
+        if telemetry_mode.value.is_disabled() && trace_upload.value {
+            tracing::info!(
+                "Telemetry disabled but trace uploads enabled: \
+                 session artifacts will be uploaded, analytics events will not"
+            );
+        }
+        update_telemetry_config(cfg, auth_manager);
+    });
+}
+
+/// Apply current telemetry config + auth identity. Tears down the client
+/// when telemetry is disabled, so it's safe to call repeatedly.
+pub fn update_telemetry_config(config: &AgentConfig, auth_manager: &AuthManager) {
+    let grok_auth = auth_manager.current().filter(|a| a.is_xai_auth());
+    let user_id = grok_auth.as_ref().map(|a| a.user_id.clone());
+    let team_id = grok_auth.as_ref().and_then(|a| a.team_id.clone());
+    let subscription_tier = super::mvp_agent::resolve_subscription_tier_for_telemetry(
+        config
+            .remote_settings
+            .as_ref()
+            .and_then(|rs| rs.subscription_tier_display.clone()),
+        auth_manager.current_or_expired().as_ref(),
+    );
+    xai_grok_telemetry::client::init(
+        config.telemetry.clone(),
+        config.resolve_telemetry_mode().value,
+        user_id,
+        team_id,
+        config.endpoints.deployment_key.clone(),
+        crate::http::origin_client_info_from_env(),
+        xai_grok_version::VERSION.to_owned(),
+        subscription_tier,
+        crate::http::shared_client(),
+    );
+}

crates/codegen/xai-grok-shell/src/agent/mod.rs

@@ -0,0 +1,31 @@
+pub mod activity;
+pub mod app;
+pub mod auth_method;
+pub mod chat_modes;
+pub mod config;
+pub mod config_model_override_parse;
+mod ext_parsers;
+pub mod feedback_client;
+pub mod folder_trust;
+pub(crate) mod handlers;
+pub mod init;
+pub mod models;
+pub mod mvp_agent;
+pub(crate) mod proxy;
+pub mod relay;
+pub(crate) mod restore_code;
+pub mod roster;
+pub mod server;
+pub mod session_config;
+pub(crate) mod session_metrics;
+pub mod session_registry_client;
+pub(crate) mod subagent;
+pub(crate) mod subscription_check;
+pub(crate) mod update_chunk_merge;
+
+pub use mvp_agent::MvpAgent;
+pub use relay::{RelayConfig, RelayHandle, spawn_relay_connection};
+pub use server::{ServerConfig, run_agent_server};
+
+#[cfg(test)]
+mod storage_client_tests;

crates/codegen/xai-grok-shell/src/agent/mvp_agent/code_nav.rs

@@ -0,0 +1,261 @@
+//! Code-navigation eligibility gating and codebase-index management for [`MvpAgent`].
+//! Co-located child of `mvp_agent` (`use super::*`).
+
+use super::*;
+
+impl MvpAgent {
+    /// Parse the `x.ai/codeNavigation.enabled` capability from an initialize
+    /// request.  Returns `false` if the field is absent or not `true`.
+    pub(crate) fn parse_code_nav_capability(init: &acp::InitializeRequest) -> bool {
+        init.client_capabilities
+            .meta
+            .as_ref()
+            .and_then(|m| m.get("x.ai/codeNavigation"))
+            .and_then(|v| v.get("enabled"))
+            .and_then(|v| v.as_bool())
+            .unwrap_or(false)
+    }
+
+    /// Start (or reuse) the codebase index for an eligible code-nav request.
+    ///
+    /// Returns `Some((handle, was_newly_started))` on success or `None` when
+    /// config/git-root checks prevent starting.  The bool is the authoritative
+    /// "first spawn vs reuse" signal threaded up from `CodebaseIndexManager`.
+    ///
+    /// This is the narrow `pub(crate)` entry point for lazy index startup
+    /// from `extensions/code_nav.rs`.  Callers must verify eligibility with
+    /// [`code_nav_eligibility_for_request`] before calling this.
+    pub(crate) fn start_codebase_index_for_code_nav(
+        &self,
+        session_id: Option<&acp::SessionId>,
+        cwd: &std::path::Path,
+    ) -> Option<(std::sync::Arc<xai_codebase_graph::IndexManagerHandle>, bool)> {
+        let (handle, was_newly_started) = self.resolve_codebase_index(cwd)?;
+        // Pin the index to the requesting session so the Weak in
+        // CodebaseIndexManager doesn't orphan it immediately.
+        if let Some(sid) = session_id {
+            self.session_index_claims
+                .borrow_mut()
+                .insert(sid.clone(), std::sync::Arc::clone(&handle));
+        }
+        Some((handle, was_newly_started))
+    }
+
+    /// Core eligibility check — pure function that accepts explicit client
+    /// context rather than reading global agent state.
+    ///
+    /// This is the single place that applies all four gates.  Call it via
+    /// [`code_nav_eligibility_for_request`] (leader-mode safe) or
+    /// [`code_nav_eligibility`] (global state, non-leader use only).
+    pub(super) fn code_nav_eligibility_inner(
+        &self,
+        cwd: &std::path::Path,
+        client_type: ClientType,
+        code_nav_enabled: bool,
+    ) -> Result<(), CodeNavEligibility> {
+        use crate::agent::config::CodebaseIndexingSetting;
+
+        // Gate 1: client type
+        if !matches!(client_type, ClientType::GrokWeb) {
+            tracing::info!(
+                client_type = ?client_type,
+                gate = "client_type",
+                skip_reason = "client_not_web",
+                "code-nav eligibility check: skipping (client type not eligible)"
+            );
+            return Err(CodeNavEligibility::ClientNotWeb);
+        }
+
+        // Gate 2: capability advertised
+        if !code_nav_enabled {
+            tracing::info!(
+                gate = "capability",
+                skip_reason = "capability_not_advertised",
+                "code-nav eligibility check: skipping (x.ai/codeNavigation.enabled not advertised)"
+            );
+            return Err(CodeNavEligibility::CapabilityNotAdvertised);
+        }
+
+        // Gate 3: config
+        let setting = self.cfg.borrow().features.codebase_indexing.clone();
+        if let CodebaseIndexingSetting::Enabled(false) = &setting {
+            tracing::info!(
+                gate = "config",
+                skip_reason = "disabled_by_config",
+                "code-nav eligibility check: skipping (codebase_indexing disabled in config)"
+            );
+            return Err(CodeNavEligibility::DisabledByConfig);
+        }
+
+        // Gate 4: git root / config globs
+        let git_root = xai_grok_workspace::session::git::find_git_root_from_path(cwd).ok();
+        match &setting {
+            CodebaseIndexingSetting::Enabled(true) => {
+                if git_root.is_none() {
+                    tracing::info!(
+                        cwd = %cwd.display(),
+                        gate = "git_root",
+                        skip_reason = "not_git_repo",
+                        "code-nav eligibility check: skipping (not inside a git repo)"
+                    );
+                    return Err(CodeNavEligibility::NotGitRepo);
+                }
+            }
+            CodebaseIndexingSetting::Patterns(_) => {
+                let check_path = git_root.as_deref().unwrap_or(cwd);
+                if !setting.should_index(check_path) {
+                    tracing::info!(
+                        cwd = %cwd.display(),
+                        gate = "config_globs",
+                        skip_reason = "disabled_by_config",
+                        "code-nav eligibility check: skipping (not matched by config globs)"
+                    );
+                    return Err(CodeNavEligibility::DisabledByConfig);
+                }
+            }
+            CodebaseIndexingSetting::Enabled(false) => {} // handled above
+        }
+
+        Ok(())
+    }
+
+    /// Check eligibility using per-session context (leader-mode safe).
+    ///
+    /// When `session_id` is provided, reads the session's own client type
+    /// and code-nav capability — the values that were in effect when that
+    /// specific client created the session.  This is correct in leader mode
+    /// where multiple clients share one agent process and `initialize()` is
+    /// called once per connection; the global fields on `MvpAgent` reflect
+    /// only the **last** client to call `initialize()`.
+    ///
+    /// Falls back to global agent state when no session_id is given.
+    pub fn code_nav_eligibility_for_request(
+        &self,
+        session_id: Option<&acp::SessionId>,
+        cwd: &std::path::Path,
+    ) -> Result<(), CodeNavEligibility> {
+        let session_id = match session_id {
+            Some(sid) => sid,
+            // No session_id: per-client capability cannot be determined without a
+            // session.  Reject with SessionRequired rather than fall back to shared
+            // global state.  Callers must provide sessionId for x.ai/code/* requests.
+            None => return Err(CodeNavEligibility::SessionRequired),
+        };
+
+        let sessions = self.sessions.borrow();
+        let (client_type, code_nav_enabled) = if let Some(handle) = sessions.get(session_id) {
+            let ct = crate::http::client_type_from_origin(handle.origin_client.as_ref());
+            (ct, handle.code_nav_enabled)
+        } else {
+            // Session not found (evicted/unknown): reject rather than silently
+            // falling back to shared global state — that would reintroduce the
+            // last-client-wins bug for stale session IDs in leader mode.
+            return Err(CodeNavEligibility::SessionRequired);
+        };
+        drop(sessions);
+        self.code_nav_eligibility_inner(cwd, client_type, code_nav_enabled)
+    }
+
+    /// Check eligibility using the stored initialize_request context.
+    ///
+    /// **Not safe in leader mode** — reads the last `initialize()` call's
+    /// client_type and capability.  Prefer [`code_nav_eligibility_for_request`]
+    /// when a session_id is available.
+    pub fn code_nav_eligibility(&self, cwd: &std::path::Path) -> Result<(), CodeNavEligibility> {
+        let client_type = *self.client_type.borrow();
+        let code_nav_enabled = self.code_nav_enabled.get();
+        self.code_nav_eligibility_inner(cwd, client_type, code_nav_enabled)
+    }
+
+    /// Resolve and get-or-create the codebase index for `cwd`, applying config
+    /// and git-root eligibility checks.
+    ///
+    /// Returns `Some((handle, was_newly_started))` when an index is available,
+    /// `None` when config or git-root checks rule it out.  The bool is the
+    /// authoritative "was this a first spawn?" signal from the manager.
+    pub(super) fn resolve_codebase_index(
+        &self,
+        cwd: &std::path::Path,
+    ) -> Option<(std::sync::Arc<xai_codebase_graph::IndexManagerHandle>, bool)> {
+        use crate::agent::config::CodebaseIndexingSetting;
+
+        let setting = self.cfg.borrow().features.codebase_indexing.clone();
+        let git_root = xai_grok_workspace::session::git::find_git_root_from_path(cwd).ok();
+
+        match (&setting, &git_root) {
+            (CodebaseIndexingSetting::Enabled(false), _) => {
+                tracing::info!(
+                    cwd = %cwd.display(),
+                    skip_reason = "disabled_by_config",
+                    "code-nav: skipping index creation (disabled in config)"
+                );
+                return None;
+            }
+            (CodebaseIndexingSetting::Enabled(true), None) => {
+                tracing::info!(
+                    cwd = %cwd.display(),
+                    skip_reason = "not_git_repo",
+                    "code-nav: skipping index creation (not inside a git repo)"
+                );
+                return None;
+            }
+            (CodebaseIndexingSetting::Patterns(_), _) => {
+                let check_path = git_root.as_deref().unwrap_or(cwd);
+                if !setting.should_index(check_path) {
+                    tracing::info!(
+                        cwd = %cwd.display(),
+                        skip_reason = "disabled_by_config",
+                        "code-nav: skipping index creation (not matched by config globs)"
+                    );
+                    return None;
+                }
+            }
+            (CodebaseIndexingSetting::Enabled(true), Some(_)) => {}
+        }
+
+        let target = git_root.unwrap_or_else(|| cwd.to_path_buf());
+        // get_or_create returns the authoritative (handle, was_newly_started) pair.
+        // Log only on actual first spawn so reuse requests are not misleadingly
+        // labelled as "starting".
+        let (handle, was_newly_started) = self.get_or_create_codebase_index(target.clone());
+        if was_newly_started {
+            tracing::info!(
+                cwd = %cwd.display(),
+                index_target = %target.display(),
+                event = "index_first_spawn",
+                "code-nav: first lazy spawn of codebase index"
+            );
+        }
+        Some((handle, was_newly_started))
+    }
+
+    pub(super) fn indexed_roots_for(&self, cwd: &std::path::Path) -> Vec<String> {
+        if self.get_codebase_index(cwd).is_some() {
+            return vec![cwd.to_string_lossy().into_owned()];
+        }
+        if let Ok(git_root) = xai_grok_workspace::session::git::find_git_root_from_path(cwd)
+            && self.get_codebase_index(&git_root).is_some()
+        {
+            return vec![git_root.to_string_lossy().into_owned()];
+        }
+        Vec::new()
+    }
+
+    /// Returns `(handle, was_newly_started)` — the bool is the authoritative
+    /// "did this call spawn a new actor?" bit from `CodebaseIndexManager::get_or_create`.
+    pub(super) fn get_or_create_codebase_index(
+        &self,
+        cwd: PathBuf,
+    ) -> (std::sync::Arc<xai_codebase_graph::IndexManagerHandle>, bool) {
+        self.codebase_indexes.lock().get_or_create(cwd)
+    }
+
+    /// Get an existing codebase index for the given cwd.
+    /// Returns None if no index exists for this cwd.
+    pub fn get_codebase_index(
+        &self,
+        cwd: &std::path::Path,
+    ) -> Option<std::sync::Arc<xai_codebase_graph::IndexManagerHandle>> {
+        self.codebase_indexes.lock().get(cwd)
+    }
+}

crates/codegen/xai-grok-shell/src/agent/mvp_agent/folder_trust_prompt.rs

@@ -0,0 +1,453 @@
+//! Interactive folder-trust prompt: a dormant agent→GUI-client ACP round-trip
+//! (`x.ai/folder_trust/request`) that asks a GUI client (grok-desktop) to decide
+//! trust for an untrusted-with-configs workspace, then grants + reloads the
+//! now-trusted project servers without a restart.
+//!
+//! DORMANT in production: it only fires when the connected client advertised
+//! `x.ai/folderTrust.interactive` AND the folder-trust feature flag is on AND the
+//! verdict is [`xai_grok_workspace::folder_trust::TrustOutcome::Prompt`]. No
+//! client advertises the capability until the desktop UI ships — so this is
+//! inert by default even with the feature flag on. The TUI/headless clients never
+//! advertise it (they self-gate trust client-side), so they are never
+//! double-prompted. Co-located child of `mvp_agent` (`use super::*`).
+//!
+//! Post-grant reload scope: MCP, plugins, and each session's own project hooks
+//! are hot-reloaded in place — for EVERY session sharing the granted workspace
+//! (same `workspace_key`), each reloaded against its OWN cwd. Project LSP is NOT
+//! hot-reloaded — the LSP backend is baked into the agent's tool bridge at build
+//! time (one-shot startup coordinator, no in-place reconfigure API), so repo-local
+//! `.grok/lsp.json` servers start on the NEXT session open (the durable grant
+//! makes the re-spawn trusted). `lsp` is still REPORTED in the prompt's
+//! `configKinds` (it is a real reason the folder is gated) — only the post-grant
+//! hot-reload skips it.
+
+use super::*;
+
+/// Max wait for a GUI client's trust decision before giving up (fail-closed).
+/// Generous because it is a human decision, but bounds the detached task so a
+/// connected-but-silent client (modal left open / client bug) can't leak it for
+/// the whole connection lifetime.
+const TRUST_PROMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60);
+
+/// ACP `x.ai/folder_trust/request` payload (agent → GUI client). Serialized as
+/// `camelCase` for the ACP JSON-RPC wire format.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct FolderTrustRequest {
+    /// The session this prompt belongs to. REQUIRED for leader Tier-2 routing:
+    /// non-interaction reverse-requests are delivered to the driver keyed on
+    /// `params.sessionId`; omitting it makes the leader silently drop the message,
+    /// so the prompt would never reach the client.
+    pub session_id: String,
+    /// The session cwd whose workspace is being gated.
+    pub cwd: String,
+    /// Display path of the canonical workspace key (the trust grant's scope).
+    pub workspace: String,
+    /// Detected repo-local config kinds (e.g. `mcp`, `hooks`, `lsp`) — the
+    /// reasons the folder is gated — for the prompt UI. Display-only, NOT the
+    /// trust gate; derived from the same scan as the gate. `lsp` may appear: it
+    /// is a real reason to prompt, but project LSP applies on the NEXT session
+    /// open rather than hot-reloading on grant (see module docs).
+    pub config_kinds: Vec<String>,
+}
+
+/// Outcome of the trust prompt (GUI client → agent). Fail-closed: any value
+/// other than `"trust"` (including unknown strings, via `#[serde(other)]`)
+/// decodes to [`FolderTrustOutcome::Reject`], so only an explicit grant unblocks.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub(crate) enum FolderTrustOutcome {
+    Trust,
+    #[serde(other)]
+    Reject,
+}
+
+/// ACP `x.ai/folder_trust/request` response (GUI client → agent).
+#[derive(Debug, Clone, serde::Deserialize)]
+pub(crate) struct FolderTrustResponse {
+    pub outcome: FolderTrustOutcome,
+}
+
+impl MvpAgent {
+    /// Parse the `x.ai/folderTrust.interactive` capability from an initialize
+    /// request. Returns `false` if absent or not `true`. Mirrors
+    /// [`Self::parse_code_nav_capability`].
+    pub(crate) fn parse_interactive_trust_capability(init: &acp::InitializeRequest) -> bool {
+        init.client_capabilities
+            .meta
+            .as_ref()
+            .and_then(|m| m.get("x.ai/folderTrust"))
+            .and_then(|v| v.get("interactive"))
+            .and_then(|v| v.as_bool())
+            .unwrap_or(false)
+    }
+
+    /// Ask a GUI client to decide trust for `session_id`'s workspace, then grant
+    /// + reload on accept. DORMANT no-op unless the client advertised
+    /// `x.ai/folderTrust.interactive` AND [`folder_trust::prompt_warranted`]
+    /// (feature on + untrusted + repo configs present).
+    ///
+    /// Non-blocking: the session was already created with project servers GATED
+    /// (the untrusted resolve in `new_session`/`load_session`), so nothing
+    /// repo-local spawns while the prompt is open. The round-trip + reload run in
+    /// a detached `spawn_local` task, so the `new_session` response is not
+    /// delayed by the (potentially long) user decision. At most one outstanding
+    /// request per workspace per process (dedup), and the await is bounded by
+    /// [`TRUST_PROMPT_TIMEOUT`].
+    pub(crate) fn maybe_spawn_interactive_trust_prompt(
+        &self,
+        session_id: &acp::SessionId,
+        cwd: &std::path::Path,
+        remote: Option<&crate::util::config::RemoteSettings>,
+    ) {
+        if !self.interactive_trust_client.get() {
+            return;
+        }
+        if !folder_trust::prompt_warranted(cwd, remote) {
+            return;
+        }
+        let key = xai_grok_workspace::trust::workspace_key(cwd);
+        // Dedup: skip if this workspace was already prompted/decided (reconnect)
+        // or has a prompt in flight (concurrent same-workspace session). `insert`
+        // returns false when already present. Agent-owned set (no process
+        // global), captured into the task for release on failure/timeout.
+        let prompted = self.interactive_trust_prompted.clone();
+        if !prompted.borrow_mut().insert(key.clone()) {
+            return;
+        }
+
+        // Capture EVERY session sharing the GRANTED WORKSPACE (same
+        // `workspace_key` — the grant's actual scope, aligned with the dedup key),
+        // each with its OWN cwd, so a grant reloads every sibling against its own
+        // project config — exactly like the per-cwd `handle_reload_project_mcp_servers`
+        // / `broadcast_plugin_registry_to_sessions`. `&self` can't be borrowed
+        // across the `spawn_local` boundary, so capture owned clones now.
+        //
+        // INTENTIONAL fail-safe limitation: this is a one-time snapshot taken at
+        // prompt-spawn. A same-workspace session created WHILE the modal is open is
+        // deduped (no second prompt) and is not in this set, so the grant won't
+        // reload it — it stays GATED until its own next session (secure, never
+        // over-exposed). Re-querying at grant time would need the `sessions` map
+        // (a non-`Rc` `RefCell` field) shared into the detached task, which isn't
+        // available here; the fail-safe stale-session window is accepted instead.
+        let targets: Vec<ReloadTarget> = self
+            .sessions
+            .borrow()
+            .values()
+            .filter(|h| {
+                xai_grok_workspace::trust::workspace_key(std::path::Path::new(&h.info.cwd)) == key
+            })
+            .map(|h| ReloadTarget {
+                cmd_tx: h.cmd_tx.clone(),
+                initial_client_mcp_servers: h.initial_client_mcp_servers.clone(),
+                cwd: PathBuf::from(&h.info.cwd),
+            })
+            .collect();
+        if targets.is_empty() {
+            prompted.borrow_mut().remove(&key);
+            return;
+        }
+
+        let gateway = self.gateway.clone();
+        let plugin_handle = self.plugin_registry_handle.clone();
+        let managed_mcp_cache = self.managed_mcp_cache.clone();
+        let auth_manager = self.auth_manager.clone();
+        let can_fetch_managed = self.can_fetch_managed_mcps();
+        let proxy_url = self.cfg.borrow().endpoints.proxy_url();
+        let compat = self.cfg.borrow().compat_resolved;
+        let remote = remote.cloned();
+        let cwd = cwd.to_path_buf();
+        let workspace = key.display().to_string();
+        let config_kinds = folder_trust::detected_config_kinds(&cwd);
+        let session_id = session_id.0.to_string();
+        // Regression guard: every reverse-request must carry a
+        // non-empty sessionId or leader Tier-2 routing silently drops it.
+        debug_assert!(
+            !session_id.is_empty(),
+            "folder_trust reverse-request must carry a non-empty sessionId (design §5.4)"
+        );
+
+        tokio::task::spawn_local(async move {
+            let request = FolderTrustRequest {
+                session_id,
+                cwd: cwd.to_string_lossy().into_owned(),
+                workspace,
+                config_kinds,
+            };
+            // Non-panicking: a struct of String/Vec<String> can't fail to
+            // serialize, but avoid `expect` in prod — bail (and release the dedup
+            // key) on the impossible error rather than aborting the task thread.
+            let raw_params = match serde_json::value::to_raw_value(&request) {
+                Ok(p) => p,
+                Err(e) => {
+                    tracing::error!(error = %e, "folder trust: request serialization failed");
+                    prompted.borrow_mut().remove(&key);
+                    return;
+                }
+            };
+            let ext_request = acp::ExtRequest::new("x.ai/folder_trust/request", raw_params.into());
+
+            use agent_client_protocol::Client as _;
+            let outcome = match tokio::time::timeout(
+                TRUST_PROMPT_TIMEOUT,
+                gateway.ext_method(ext_request),
+            )
+            .await
+            {
+                // A decodable response carries the user's decision. An
+                // undecodable success payload is a client/protocol error, not a
+                // decision: stay gated (fail-closed) but release the dedup key so
+                // a later session can re-prompt — same as transport/timeout below.
+                Ok(Ok(raw)) => match serde_json::from_str::<FolderTrustResponse>(raw.0.get()) {
+                    Ok(r) => r.outcome,
+                    Err(e) => {
+                        tracing::debug!(error = %e, "folder trust: undecodable trust response; staying gated, releasing dedup key");
+                        prompted.borrow_mut().remove(&key);
+                        return;
+                    }
+                },
+                Ok(Err(e)) => {
+                    // Client disconnected / transport error: not a decision —
+                    // release the key so a later session can re-prompt.
+                    tracing::debug!(error = %e, "folder trust: client trust request failed");
+                    prompted.borrow_mut().remove(&key);
+                    return;
+                }
+                Err(_elapsed) => {
+                    // Connected but silent past the deadline: stay gated, release
+                    // the key so a future session may re-prompt.
+                    tracing::info!(
+                        cwd = %cwd.display(),
+                        "folder trust: no client decision before timeout; staying gated"
+                    );
+                    prompted.borrow_mut().remove(&key);
+                    return;
+                }
+            };
+            if outcome != FolderTrustOutcome::Trust {
+                // Decided "reject": keep the dedup key so the user is not
+                // re-prompted for this workspace on every reconnect.
+                tracing::info!(
+                    cwd = %cwd.display(),
+                    "folder trust: GUI client declined; workspace stays gated"
+                );
+                return;
+            }
+
+            // Re-check the dedup key before granting. `HooksAction::Untrust`
+            // removes this workspace's key (and revokes asynchronously) when the
+            // user untrusts. If that fired while the modal was open, the key is
+            // gone — honor the untrust and drop this now-stale "trust" rather
+            // than re-persisting a grant the user just revoked. The single-
+            // threaded LocalSet makes this check + grant atomic w.r.t. the
+            // untrust task (no await in between).
+            if !prompted.borrow().contains(&key) {
+                tracing::info!(
+                    cwd = %cwd.display(),
+                    "folder trust: workspace untrusted while prompt was open; ignoring stale grant"
+                );
+                return;
+            }
+
+            // Persist the grant, then flip the cached untrusted verdict to trusted
+            // (the `Some(false)` arm of `resolve_and_record` re-reads the store).
+            folder_trust::grant_folder_trust(&cwd);
+            folder_trust::resolve_and_record(&cwd, remote.as_ref(), false);
+
+            reload_project_servers_after_grant(ReloadAfterGrant {
+                gateway: &gateway,
+                targets,
+                plugin_handle: &plugin_handle,
+                managed_mcp_cache: &managed_mcp_cache,
+                auth_manager: &auth_manager,
+                can_fetch_managed,
+                proxy_url: &proxy_url,
+                compat: &compat,
+                prompt_cwd: &cwd,
+            })
+            .await;
+
+            tracing::info!(
+                cwd = %cwd.display(),
+                "folder trust: granted via GUI client; reloaded project servers"
+            );
+        });
+    }
+}
+
+/// One session to reload after a grant, with ITS OWN cwd (so the MCP merge +
+/// plugin build use the session's own project config — matching the per-cwd
+/// canonical reloaders, not the prompt's cwd).
+struct ReloadTarget {
+    cmd_tx: tokio::sync::mpsc::UnboundedSender<crate::session::SessionCommand>,
+    initial_client_mcp_servers: Vec<acp::McpServer>,
+    cwd: PathBuf,
+}
+
+/// Inputs for [`reload_project_servers_after_grant`], bundled to keep the
+/// orchestrator free of a long positional arg list.
+struct ReloadAfterGrant<'a> {
+    gateway: &'a GatewaySender,
+    /// Every session sharing the granted workspace, each with its own cwd.
+    targets: Vec<ReloadTarget>,
+    plugin_handle: &'a xai_grok_agent::plugins::SharedPluginRegistryHandle,
+    managed_mcp_cache: &'a crate::session::managed_mcp::ManagedMcpStateHandle,
+    auth_manager: &'a std::sync::Arc<AuthManager>,
+    can_fetch_managed: bool,
+    proxy_url: &'a str,
+    compat: &'a xai_grok_tools::types::CompatConfig,
+    /// The prompting session's cwd — used only for the client catalog push.
+    prompt_cwd: &'a std::path::Path,
+}
+
+/// Reload each granted-workspace session's now-trusted project servers in place
+/// (no restart), driving the canonical primitives the normal spawn/reload paths
+/// use — PER SESSION CWD, like `handle_reload_project_mcp_servers` /
+/// `broadcast_plugin_registry_to_sessions`: `fetch_managed_mcp_configs` +
+/// `merge_managed_mcp_servers` (`SessionCommand::UpdateMcpServers`), `build_for_cwd`
+/// (`SessionCommand::ReloadPlugins`), and `reload_hooks_impl`
+/// (`SessionCommand::ReloadHooks`), then push the refreshed MCP catalog. LSP is
+/// spawn-baked and applies on the next session open (see module docs). Caller
+/// must have granted + recorded trust first.
+async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) {
+    // Managed (gateway/Toolbox) servers must survive the re-merge; fetch them once
+    // (cwd-independent) via the shared helper (single-sources the auth-key dance
+    // with `MvpAgent::get_managed_mcp_configs`). The plugin MCP snapshot is also
+    // global, so it is fine to reuse across cwds for the merge.
+    let managed = if ctx.can_fetch_managed {
+        crate::session::managed_mcp::fetch_managed_mcp_configs(
+            ctx.managed_mcp_cache,
+            ctx.proxy_url,
+            ctx.auth_manager,
+        )
+        .await
+    } else {
+        vec![]
+    };
+    let plugin_snapshot = ctx.plugin_handle.snapshot();
+
+    for target in ctx.targets {
+        // Per-session cwd: a sibling session in a subdir of the granted workspace
+        // must get ITS OWN project config, not the prompt's.
+        let session_cwd = target.cwd.as_path();
+        // MCP: `merge_managed_mcp_servers` re-reads disk + runs
+        // `filter_untrusted_project_mcp`, which now KEEPS project servers because
+        // the cached verdict was flipped to trusted (same workspace key).
+        let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
+            target.initial_client_mcp_servers,
+            session_cwd,
+            &managed,
+            plugin_snapshot.as_deref(),
+            ctx.compat,
+        );
+        let (tx, _rx) = tokio::sync::oneshot::channel();
+        let _ = target
+            .cmd_tx
+            .send(crate::session::SessionCommand::UpdateMcpServers {
+                mcp_servers: merged,
+                respond_to: tx,
+            });
+        // Plugins (+ plugin-contributed hooks) built for this session's own cwd
+        // on the folder-trust verdict (mirrors `broadcast_plugin_registry_to_sessions`);
+        // the grant + resolve_and_record above flipped the cached verdict to trusted.
+        let disk_cfg =
+            crate::config::resolve_effective_plugins_config(session_cwd).to_discovery_config();
+        let project_trusted = folder_trust::project_scope_allowed(session_cwd);
+        // Session `_meta.pluginDirs` are re-merged by the receiving actor
+        // (`preserve_session_plugin_dirs` on `ReloadPlugins`).
+        let registry =
+            ctx.plugin_handle
+                .build_for_cwd(session_cwd, &disk_cfg, &[], project_trusted);
+        let _ = target
+            .cmd_tx
+            .send(crate::session::SessionCommand::ReloadPlugins { registry });
+        // The session's OWN project hooks (`.grok/hooks`, `.cursor/hooks.json`),
+        // which `ReloadPlugins` does NOT touch — re-discovered against the actor's
+        // own `session_info.cwd` on the now-trusted verdict by `reload_hooks_impl`.
+        let _ = target
+            .cmd_tx
+            .send(crate::session::SessionCommand::ReloadHooks);
+    }
+
+    // Push the refreshed MCP catalog (for the prompting session's cwd) so the
+    // client UI reflects the now-trusted repo-local servers.
+    let local = folder_trust::filter_untrusted_project_mcp(
+        ctx.prompt_cwd,
+        crate::util::config::load_mcp_servers(ctx.prompt_cwd, ctx.compat),
+    );
+    crate::extensions::mcp::notify_servers_updated(ctx.gateway, &managed, &local).await;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn init_with_meta(meta: Option<serde_json::Value>) -> acp::InitializeRequest {
+        // Production reads `client_capabilities.meta`, not top-level request meta.
+        let mut caps = acp::ClientCapabilities::new()
+            .fs(acp::FileSystemCapabilities::new())
+            .terminal(false);
+        if let Some(m) = meta
+            && let Some(map) = m.as_object().cloned()
+        {
+            caps = caps.meta(map);
+        }
+        acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_capabilities(caps)
+    }
+
+    #[test]
+    fn parse_interactive_trust_capability_present_and_true() {
+        let mut meta = serde_json::Map::new();
+        meta.insert(
+            "x.ai/folderTrust".to_string(),
+            serde_json::json!({ "interactive": true }),
+        );
+        let init = init_with_meta(Some(serde_json::Value::Object(meta)));
+        assert!(MvpAgent::parse_interactive_trust_capability(&init));
+    }
+
+    #[test]
+    fn parse_interactive_trust_capability_absent_returns_false() {
+        let init = init_with_meta(None);
+        assert!(!MvpAgent::parse_interactive_trust_capability(&init));
+    }
+
+    #[test]
+    fn parse_interactive_trust_capability_false_returns_false() {
+        let mut meta = serde_json::Map::new();
+        meta.insert(
+            "x.ai/folderTrust".to_string(),
+            serde_json::json!({ "interactive": false }),
+        );
+        let init = init_with_meta(Some(serde_json::Value::Object(meta)));
+        assert!(!MvpAgent::parse_interactive_trust_capability(&init));
+    }
+
+    #[test]
+    fn request_serializes_camel_case_with_session_id() {
+        let req = FolderTrustRequest {
+            session_id: "sess-1".into(),
+            cwd: "/repo".into(),
+            workspace: "/repo".into(),
+            config_kinds: vec!["mcp".into()],
+        };
+        let json = serde_json::to_value(&req).unwrap();
+        assert!(json.get("configKinds").is_some());
+        assert!(json.get("config_kinds").is_none());
+        // Leader Tier-2 routing reads `params.sessionId`; it must be present and
+        // non-empty (regression guard for the silently-dropped-in-leader bug).
+        assert_eq!(json["sessionId"], "sess-1");
+        assert!(!json["sessionId"].as_str().unwrap().is_empty());
+    }
+
+    #[test]
+    fn response_decodes_trust_reject_and_unknown_fail_closed() {
+        let trust: FolderTrustResponse = serde_json::from_str(r#"{"outcome":"trust"}"#).unwrap();
+        assert_eq!(trust.outcome, FolderTrustOutcome::Trust);
+        let reject: FolderTrustResponse = serde_json::from_str(r#"{"outcome":"reject"}"#).unwrap();
+        assert_eq!(reject.outcome, FolderTrustOutcome::Reject);
+        // Unknown outcome must fail closed to Reject (never silently "trust").
+        let unknown: FolderTrustResponse = serde_json::from_str(r#"{"outcome":"banana"}"#).unwrap();
+        assert_eq!(unknown.outcome, FolderTrustOutcome::Reject);
+    }
+}

crates/codegen/xai-grok-shell/src/agent/mvp_agent/heap_profile.rs

@@ -0,0 +1,352 @@
+//! Heap-profile monitor wiring for [`MvpAgent`].
+//!
+//! Full-reapply sites call [`MvpAgent::reconfigure_heap_profile_monitor`].
+//! K12 scoped kill-switch reconfigures only jemalloc fields (no wholesale
+//! `remote_settings` rewrite, no `re_resolve_runtime_fields` / telemetry re-init).
+
+use super::*;
+use crate::heap_profile::{SCOPED_KILL_SWITCH_INTERVAL, build_upload_handles};
+
+impl MvpAgent {
+    pub(super) fn reconfigure_heap_profile_monitor(&self) {
+        let zdr = self.is_data_collection_disabled();
+        let config = self.cfg.borrow().resolve_jemalloc_heap_profile(zdr);
+        let handles = self.heap_profile_upload_handles();
+        self.heap_profile_monitor
+            .borrow_mut()
+            .reconfigure(config, handles);
+    }
+
+    pub(super) fn heap_profile_set_session_id(&self, session_id: &str) {
+        self.heap_profile_monitor
+            .borrow_mut()
+            .set_session_id(session_id.to_owned());
+    }
+
+    fn heap_profile_upload_handles(&self) -> Option<crate::heap_profile::HeapProfileUploadHandles> {
+        let method = self.trace_upload_config_snapshot()?;
+        let bucket_url = self
+            .cfg
+            .borrow()
+            .endpoints
+            .resolve_trace_bucket_url()
+            .map(|r| r.value);
+        // Only direct GCS uploads need a bucket.
+        if bucket_url.is_none()
+            && matches!(
+                method,
+                crate::session::repo_changes::UploadMethod::Direct { .. }
+            )
+        {
+            tracing::debug!("no trace bucket configured; heap-profile uploads disabled");
+            return None;
+        }
+        Some(build_upload_handles(
+            Arc::clone(&self.auth_manager),
+            bucket_url,
+            method,
+        ))
+    }
+
+    /// Background poll + scoped kill-switch (agent entrypoints only).
+    /// Idempotent; skipped under `cfg!(test)`.
+    pub(super) fn spawn_heap_profile_monitor(&self) {
+        if cfg!(test) || self.heap_profile_started.replace(true) {
+            return;
+        }
+        self.reconfigure_heap_profile_monitor();
+        let agent_ref = LocalRef::new(self);
+        tokio::task::spawn_local(async move {
+            let mut last_kill_switch = tokio::time::Instant::now();
+            loop {
+                let poll_interval = {
+                    let mon = agent_ref.get().heap_profile_monitor.borrow();
+                    if mon.config().enabled {
+                        mon.config().poll_interval
+                    } else {
+                        std::time::Duration::from_secs(30)
+                    }
+                };
+                tokio::time::sleep(poll_interval).await;
+
+                let enabled = agent_ref
+                    .get()
+                    .heap_profile_monitor
+                    .borrow()
+                    .config()
+                    .enabled;
+                if !enabled {
+                    continue;
+                }
+
+                if last_kill_switch.elapsed() >= SCOPED_KILL_SWITCH_INTERVAL {
+                    let result = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(
+                        agent_ref.get().poll_scoped_jemalloc_kill_switch_once(),
+                    ))
+                    .await;
+                    if result.is_err() {
+                        tracing::error!(
+                            "heap_profile: scoped kill-switch tick panicked; continuing"
+                        );
+                    }
+                    last_kill_switch = tokio::time::Instant::now();
+                }
+
+                let result = futures::FutureExt::catch_unwind(std::panic::AssertUnwindSafe(
+                    agent_ref.get().heap_profile_poll_tick_once(),
+                ))
+                .await;
+                if result.is_err() {
+                    agent_ref
+                        .get()
+                        .heap_profile_monitor
+                        .borrow_mut()
+                        .clear_upload_in_flight();
+                    tracing::error!("heap_profile: poll tick panicked; continuing");
+                }
+            }
+        });
+    }
+
+    async fn heap_profile_poll_tick_once(&self) {
+        let pending = {
+            let mut mon = self.heap_profile_monitor.borrow_mut();
+            mon.begin_tick()
+        };
+        let Some(pending) = pending else {
+            return;
+        };
+        let threshold = pending.threshold;
+        let outcome = pending.execute().await;
+        self.heap_profile_monitor
+            .borrow_mut()
+            .finish_tick(threshold, outcome);
+    }
+
+    /// K12: fetch settings, reconfigure from jemalloc fields only.
+    ///
+    /// Also patches jemalloc fields on stored `remote_settings` so full-reapply
+    /// sites (`/new` → `reconfigure_heap_profile_monitor`) cannot re-enable
+    /// profiling from a stale enabled flag after a live kill-switch when the
+    /// subsequent wholesale refresh is skipped or fails.
+    pub(super) async fn poll_scoped_jemalloc_kill_switch_once(&self) {
+        if !self.heap_profile_monitor.borrow().config().enabled {
+            return;
+        }
+        let Ok(auth) = self.auth_manager.auth().await else {
+            tracing::debug!("heap_profile scoped poll skipped: not authenticated");
+            return;
+        };
+        let Some(settings) = self.fetch_remote_settings(auth).await else {
+            tracing::debug!("heap_profile scoped poll skipped: settings fetch failed");
+            return;
+        };
+
+        // Keep stored jemalloc knobs in sync with the live fetch without a
+        // wholesale remote_settings rewrite (no telemetry / announcements churn).
+        {
+            let mut cfg = self.cfg.borrow_mut();
+            if let Some(rs) = cfg.remote_settings.as_mut() {
+                rs.jemalloc_heap_profile_enabled = settings.jemalloc_heap_profile_enabled;
+                rs.jemalloc_heap_profile_thresholds_bytes =
+                    settings.jemalloc_heap_profile_thresholds_bytes.clone();
+                rs.jemalloc_heap_profile_poll_interval_secs =
+                    settings.jemalloc_heap_profile_poll_interval_secs;
+            }
+        }
+
+        let resolved = self
+            .cfg
+            .borrow()
+            .resolve_jemalloc_heap_profile_from_partial(
+                settings.jemalloc_heap_profile_enabled,
+                settings.jemalloc_heap_profile_thresholds_bytes.as_deref(),
+                settings.jemalloc_heap_profile_poll_interval_secs,
+                self.is_data_collection_disabled(),
+            );
+
+        let handles = self.heap_profile_upload_handles();
+        self.heap_profile_monitor
+            .borrow_mut()
+            .reconfigure(resolved, handles);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn cfg_with_remote(rs: crate::util::config::RemoteSettings) -> AgentConfig {
+        AgentConfig {
+            remote_settings: Some(rs),
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn scoped_resolve_does_not_mutate_stored_remote_settings() {
+        let cfg = cfg_with_remote(crate::util::config::RemoteSettings {
+            jemalloc_heap_profile_enabled: Some(true),
+            jemalloc_heap_profile_thresholds_bytes: Some(vec![1_000_000]),
+            jemalloc_heap_profile_poll_interval_secs: Some(30),
+            trace_upload_enabled: Some(true),
+            telemetry_mode: Some("all".into()),
+            ..Default::default()
+        });
+
+        let resolved = cfg.resolve_jemalloc_heap_profile_from_partial(
+            Some(false),
+            Some(&[1_000_000]),
+            Some(30),
+            false,
+        );
+        assert!(!resolved.enabled);
+
+        let stored = cfg.remote_settings.as_ref().unwrap();
+        assert_eq!(stored.jemalloc_heap_profile_enabled, Some(true));
+        assert_eq!(
+            stored.jemalloc_heap_profile_thresholds_bytes.as_deref(),
+            Some([1_000_000u64].as_slice())
+        );
+    }
+
+    /// Kill-switch poll patches stored jemalloc knobs so full reapply (`/new`)
+    /// cannot re-enable from a stale flag when wholesale refresh is skipped.
+    #[test]
+    fn kill_switch_patch_keeps_full_reapply_disabled() {
+        let mut cfg = cfg_with_remote(crate::util::config::RemoteSettings {
+            jemalloc_heap_profile_enabled: Some(true),
+            jemalloc_heap_profile_thresholds_bytes: Some(vec![1_000_000]),
+            jemalloc_heap_profile_poll_interval_secs: Some(30),
+            trace_upload_enabled: Some(true),
+            telemetry_mode: Some("all".into()),
+            ..Default::default()
+        });
+
+        // Simulate live kill-switch fetch: remote disabled jemalloc profiling.
+        if let Some(rs) = cfg.remote_settings.as_mut() {
+            rs.jemalloc_heap_profile_enabled = Some(false);
+            rs.jemalloc_heap_profile_thresholds_bytes = Some(vec![1_000_000]);
+            rs.jemalloc_heap_profile_poll_interval_secs = Some(30);
+        }
+
+        // Full reapply path reads stored fields (hooks available for gate check).
+        let free = crate::heap_profile::resolve_jemalloc_heap_profile(
+            cfg.remote_settings
+                .as_ref()
+                .and_then(|s| s.jemalloc_heap_profile_enabled),
+            cfg.remote_settings
+                .as_ref()
+                .and_then(|s| s.jemalloc_heap_profile_thresholds_bytes.as_deref()),
+            cfg.remote_settings
+                .as_ref()
+                .and_then(|s| s.jemalloc_heap_profile_poll_interval_secs),
+            false,
+            true,
+            true,
+        );
+        assert!(!free.enabled);
+        assert_eq!(
+            cfg.remote_settings
+                .as_ref()
+                .and_then(|s| s.jemalloc_heap_profile_enabled),
+            Some(false)
+        );
+    }
+
+    #[test]
+    fn scoped_partial_gates_and_free_resolve_agree() {
+        let thresholds = [2u64 * 1024 * 1024 * 1024];
+        let free = crate::heap_profile::resolve_jemalloc_heap_profile(
+            Some(true),
+            Some(&thresholds),
+            Some(15),
+            false,
+            true,
+            true,
+        );
+        assert!(free.enabled);
+        assert_eq!(free.poll_interval, std::time::Duration::from_secs(15));
+        assert_eq!(free.thresholds, thresholds);
+
+        let cfg = cfg_with_remote(crate::util::config::RemoteSettings {
+            trace_upload_enabled: Some(true),
+            telemetry_mode: Some("all".into()),
+            jemalloc_heap_profile_enabled: Some(false),
+            jemalloc_heap_profile_thresholds_bytes: Some(vec![100]),
+            ..Default::default()
+        });
+        assert!(
+            !cfg.resolve_jemalloc_heap_profile_from_partial(
+                Some(false),
+                Some(&[100]),
+                Some(30),
+                false
+            )
+            .enabled
+        );
+        assert!(
+            !cfg.resolve_jemalloc_heap_profile_from_partial(Some(true), Some(&[]), Some(30), false)
+                .enabled
+        );
+        assert!(
+            !cfg.resolve_jemalloc_heap_profile_from_partial(
+                Some(true),
+                Some(&[100]),
+                Some(30),
+                true
+            )
+            .enabled
+        );
+    }
+
+    #[test]
+    fn full_reapply_reads_stored_remote_jemalloc_fields() {
+        let thresholds = vec![100u64, 200];
+        let cfg = cfg_with_remote(crate::util::config::RemoteSettings {
+            jemalloc_heap_profile_enabled: Some(true),
+            jemalloc_heap_profile_thresholds_bytes: Some(thresholds.clone()),
+            jemalloc_heap_profile_poll_interval_secs: Some(45),
+            trace_upload_enabled: Some(true),
+            telemetry_mode: Some("all".into()),
+            ..Default::default()
+        });
+
+        let full = cfg.resolve_jemalloc_heap_profile(false);
+        // Without installed hooks, prof_available is false → disabled.
+        assert!(!full.enabled);
+        assert_eq!(full.poll_interval, std::time::Duration::from_secs(45));
+        assert_eq!(full.thresholds, vec![100, 200]);
+
+        assert!(
+            !cfg.resolve_jemalloc_heap_profile_from_partial(
+                Some(false),
+                Some(&[100, 200]),
+                Some(45),
+                false
+            )
+            .enabled
+        );
+        assert!(
+            !cfg.resolve_jemalloc_heap_profile_from_partial(
+                Some(true),
+                Some(&[100]),
+                Some(45),
+                true
+            )
+            .enabled
+        );
+
+        let free = crate::heap_profile::resolve_jemalloc_heap_profile(
+            Some(true),
+            Some(&thresholds),
+            Some(45),
+            false,
+            true,
+            true,
+        );
+        assert!(free.enabled);
+        assert_eq!(free.thresholds, vec![100, 200]);
+        assert_eq!(free.poll_interval, std::time::Duration::from_secs(45));
+    }
+}

crates/codegen/xai-grok-shell/src/agent/mvp_agent/prompt_response_meta_tests.rs

@@ -0,0 +1,147 @@
+use super::{PromptResponseMetaArgs, build_prompt_response_meta};
+use xai_grok_sampling_types::TokenUsage;
+
+/// Baseline args with no usage, cancellation, or structured output.
+fn args<'a>(
+    session_id: &'a str,
+    prompt_id: &'a str,
+    total_tokens: u64,
+    model_id: &'a str,
+) -> PromptResponseMetaArgs<'a> {
+    PromptResponseMetaArgs {
+        session_id,
+        prompt_id,
+        total_tokens,
+        model_id,
+        last_turn_usage: None,
+        prompt_usage: None,
+        cancellation_category: None,
+        cancel_trigger: None,
+        structured_output: None,
+    }
+}
+
+#[test]
+fn includes_baseline_keys_without_usage() {
+    let meta = build_prompt_response_meta(args("sess-1", "prompt-1", 42_000, "grok-4.5"));
+    assert_eq!(meta["sessionId"], "sess-1");
+    assert_eq!(meta["requestId"], "prompt-1");
+    assert_eq!(meta["promptId"], "prompt-1");
+    assert_eq!(meta["totalTokens"], 42_000);
+    assert_eq!(meta["modelId"], "grok-4.5");
+    // No per-turn keys when usage is absent.
+    assert!(meta.get("inputTokens").is_none());
+    assert!(meta.get("outputTokens").is_none());
+    assert!(meta.get("cachedReadTokens").is_none());
+}
+
+#[test]
+fn enriches_meta_with_camelcase_token_keys() {
+    let usage = TokenUsage {
+        prompt_tokens: 1500,
+        completion_tokens: 200,
+        total_tokens: 1700,
+        reasoning_tokens: 75,
+        cached_prompt_tokens: 1000,
+    };
+    let meta = build_prompt_response_meta(PromptResponseMetaArgs {
+        last_turn_usage: Some(&usage),
+        ..args("sess-1", "prompt-1", 1_700, "grok-4.5")
+    });
+    // Bot's _META_TOKEN_KEY_MAP expects exactly these camelCase keys.
+    assert_eq!(meta["inputTokens"], 1500);
+    assert_eq!(meta["outputTokens"], 200);
+    assert_eq!(meta["cachedReadTokens"], 1000);
+    // Reasoning tokens carried through for diagnostic visibility.
+    assert_eq!(meta["reasoningTokens"], 75);
+}
+
+#[test]
+fn preserves_zero_token_values() {
+    // Responses API hits with no cache return cached_prompt_tokens=0.
+    // The key is still emitted as 0 so the bot can distinguish "no cache
+    // hit" from "no usage data". (The bot's _merge_meta_usage requires
+    // the key to be present and integer-typed.)
+    let usage = TokenUsage {
+        prompt_tokens: 100,
+        completion_tokens: 10,
+        total_tokens: 110,
+        reasoning_tokens: 0,
+        cached_prompt_tokens: 0,
+    };
+    let meta = build_prompt_response_meta(PromptResponseMetaArgs {
+        last_turn_usage: Some(&usage),
+        ..args("s", "p", 110, "m")
+    });
+    assert_eq!(meta["cachedReadTokens"], 0);
+    assert_eq!(meta["reasoningTokens"], 0);
+}
+
+#[test]
+fn usage_object_lands_on_meta() {
+    let mut ledger = xai_chat_state::UsageLedger::default();
+    ledger.record_main_loop_call(
+        "m",
+        &TokenUsage {
+            prompt_tokens: 100,
+            completion_tokens: 10,
+            total_tokens: 999_999,
+            reasoning_tokens: 0,
+            cached_prompt_tokens: 0,
+        },
+        None,
+        None,
+    );
+    let meta = build_prompt_response_meta(PromptResponseMetaArgs {
+        prompt_usage: Some(crate::extensions::notification::PromptUsage::from(&ledger)),
+        ..args("s", "p", 110, "m")
+    });
+    assert_eq!(meta["usage"]["totalTokens"], 110);
+    assert_eq!(meta["usage"]["modelUsage"]["m"]["inputTokens"], 100);
+    assert!(
+        build_prompt_response_meta(args("s", "p", 0, "m"))
+            .get("usage")
+            .is_none()
+    );
+}
+
+#[test]
+fn cancel_trigger_lands_as_camelcase_meta_key() {
+    // A send-now cancelled turn's PromptResponse `_meta` carries `cancelTrigger: "send_now"`.
+    let meta = build_prompt_response_meta(PromptResponseMetaArgs {
+        cancel_trigger: Some("send_now".to_string()),
+        ..args("s", "p", 0, "m")
+    });
+    assert_eq!(meta["cancelTrigger"], "send_now");
+
+    // Absent for non-cancel completions — the key must not appear.
+    let none = build_prompt_response_meta(args("s", "p", 0, "m"));
+    assert!(none.get("cancelTrigger").is_none());
+}
+
+#[test]
+fn structured_output_maps_to_camelcase_meta_keys() {
+    // Success carries the validated value under `structuredOutput`; no error key.
+    let ok = build_prompt_response_meta(PromptResponseMetaArgs {
+        structured_output: Some(Ok(serde_json::json!({"name": "ada"}))),
+        ..args("s", "p", 0, "m")
+    });
+    assert_eq!(ok["structuredOutput"]["name"], "ada");
+    assert!(ok.get("structuredOutputError").is_none());
+
+    // Failure carries the message under `structuredOutputError`; no value key.
+    let err = build_prompt_response_meta(PromptResponseMetaArgs {
+        structured_output: Some(Err("output does not match the required schema".to_string())),
+        ..args("s", "p", 0, "m")
+    });
+    assert_eq!(
+        err["structuredOutputError"],
+        "output does not match the required schema"
+    );
+    assert!(err.get("structuredOutput").is_none());
+
+    // No schema requested → neither key present.
+    let none = build_prompt_response_meta(args("s", "p", 0, "m"));
+    assert!(none.get("structuredOutput").is_none());
+    assert!(none.get("structuredOutputError").is_none());
+}

crates/codegen/xai-grok-shell/src/agent/mvp_agent/session_lifecycle.rs

@@ -0,0 +1,406 @@
+//! Session lifecycle, roster deltas, and the idle-session supervisor for [`MvpAgent`].
+//! Co-located `#[path]`-style child of `mvp_agent` (`use super::*`) so the `impl`
+//! block keeps access to `MvpAgent`'s private fields.
+use super::*;
+impl MvpAgent {
+    /// Ask a live session actor to shut down.
+    pub(crate) fn request_session_shutdown(&self, id: &acp::SessionId) {
+        if let Some(handle) = self.sessions.borrow().get(id) {
+            let _ = handle.cmd_tx.send(SessionCommand::Shutdown);
+        }
+    }
+    /// Finalize the cloud session replica (fire-and-forget, "Hook 4").
+    ///
+    /// Marks the session **done** upstream, so this MUST only run on a genuine
+    /// session end — a terminal/explicit close (`x.ai/session/close`). It must
+    /// NOT run on a mere client disconnect or a dead-actor reap: those leave the
+    /// conversation resumable on disk, and finalizing would wrongly mark a still
+    /// running/resumable session "done".
+    pub(super) fn finalize_session_replica(&self, id: &acp::SessionId) {
+        #[cfg(test)]
+        self.finalize_spy.borrow_mut().push(id.0.to_string());
+        if let Some(client) = self.session_registry_client() {
+            let sid = id.0.to_string();
+            tokio::spawn(async move {
+                if let Err(e) = client.finalize(&sid).await {
+                    tracing::warn!(
+                        error = % e, "session registry finalize failed (non-fatal)"
+                    );
+                }
+            });
+        }
+    }
+    /// Remove a session and its thread handle **without** finalizing the cloud
+    /// replica.
+    ///
+    /// Used for dead-actor reaping and idle-unload: the conversation stays
+    /// resumable on disk, so it must NOT be marked "done" upstream. Genuine
+    /// terminal closes go through [`MvpAgent::close_session_explicit`]. Also
+    /// drops the `session_live_state` entry so that map stays bounded.
+    pub(crate) fn remove_session(&self, id: &acp::SessionId) {
+        self.sessions.borrow_mut().remove(id);
+        self.prompt_intake_locks.borrow_mut().remove(id);
+        self.session_threads.borrow_mut().remove(id);
+        self.session_index_claims.borrow_mut().remove(id);
+        self.require_gateway_sessions.borrow_mut().remove(id);
+        self.session_live_state.borrow_mut().remove(id);
+    }
+    /// Get-or-create the per-session prompt-intake lock (see
+    /// [`Self::prompt_intake_locks`]). Cheap clone of the shared `Rc`.
+    pub(super) fn prompt_intake_lock(
+        &self,
+        id: &acp::SessionId,
+    ) -> std::rc::Rc<tokio::sync::Mutex<()>> {
+        self.prompt_intake_locks
+            .borrow_mut()
+            .entry(id.clone())
+            .or_default()
+            .clone()
+    }
+    /// Close a session in response to an **explicit** terminal close
+    /// (`x.ai/session/close`). Finalizes the cloud replica (genuine session
+    /// end), then removes the session terminally as `Completed`.
+    pub(crate) fn close_session_explicit(&self, id: &acp::SessionId) {
+        self.finalize_session_replica(id);
+        self.remove_session_terminal(id, SessionLiveState::Completed);
+    }
+    /// Record the coarse lifecycle state for a session.
+    pub(super) fn set_session_live_state(&self, id: &acp::SessionId, state: SessionLiveState) {
+        self.session_live_state
+            .borrow_mut()
+            .insert(id.clone(), state);
+    }
+    /// Read the recorded lifecycle state for a session (test observability).
+    #[cfg(test)]
+    pub(super) fn session_live_state_for(&self, id: &acp::SessionId) -> Option<SessionLiveState> {
+        self.session_live_state.borrow().get(id).copied()
+    }
+    /// Roster-delta hook for a terminally removed session. Broadcasts an
+    /// `x.ai/sessions/changed` notification with the session in `removed` so
+    /// every attached dashboard drops the row promptly. Also
+    /// records the call site (and the terminal state) for test observability,
+    /// since the `session_live_state` entry is dropped on removal.
+    pub(super) fn record_roster_delta(&self, id: &acp::SessionId, final_state: SessionLiveState) {
+        #[cfg(test)]
+        self.roster_delta_spy
+            .borrow_mut()
+            .push((id.0.to_string(), final_state));
+        tracing::debug!(
+            session_id = % id.0, ? final_state, "roster delta: session removed"
+        );
+        self.emit_roster_changed(Vec::new(), vec![id.0.to_string()]);
+    }
+    /// Roster-delta hook for a newly-resident / changed session. Broadcasts an
+    /// `x.ai/sessions/changed` notification with the current entry in
+    /// `upserted` so dashboards add/refresh the row.
+    pub(crate) fn push_roster_delta_upserted(&self, id: &acp::SessionId) {
+        if let Some(entry) = self.resident_roster_entry(id) {
+            self.emit_roster_changed(vec![entry], Vec::new());
+        }
+    }
+    /// Emit an `x.ai/sessions/changed` upsert for a resident session with an
+    /// explicit `activity`, so every attached dashboard reflects a
+    /// turn-boundary transition (Working / Idle / NeedsInput) *immediately*
+    /// rather than waiting for the ≤1s roster poll (deltas are emitted
+    /// at turn-start/turn-end). Without this, a viewer client that holds no
+    /// local `AgentView` for the session only learns its activity from the
+    /// poll, so a turn driven by another client shows as `Idle` for up to a
+    /// poll interval — and not at all while that viewer's poll is dormant.
+    ///
+    /// The `activity` is supplied by the caller rather than read from
+    /// `resident_activity` because at turn-start the actor may not have
+    /// published `current_prompt_id` yet (it is set asynchronously once the
+    /// actor dequeues the `SessionCommand::Prompt`), so a natural read would
+    /// still observe `Idle`. The authoritative entry (cwd / worktree / model /
+    /// yolo) is built by `resident_roster_entry`, so it never diverges from
+    /// the polled entry; only the `activity` field is overridden.
+    pub(super) fn push_roster_activity_delta(
+        &self,
+        id: &acp::SessionId,
+        activity: crate::agent::roster::RosterActivity,
+    ) {
+        if let Some(mut entry) = self.resident_roster_entry(id) {
+            entry.activity = activity;
+            self.emit_roster_changed(vec![entry], Vec::new());
+        }
+    }
+    /// Fan an `x.ai/sessions/changed` delta out to every attached client.
+    ///
+    /// This is a roster-wide notification (no `sessionId`), so the leader IPC
+    /// server broadcasts it to all clients rather than routing by session (see
+    /// the `x.ai/sessions/changed` special-case in `leader/server.rs`).
+    pub(super) fn emit_roster_changed(
+        &self,
+        upserted: Vec<crate::agent::roster::RosterEntry>,
+        removed: Vec<String>,
+    ) {
+        if upserted.is_empty() && removed.is_empty() {
+            return;
+        }
+        let payload = crate::agent::roster::RosterChanged { upserted, removed };
+        if let Ok(params) = serde_json::value::to_raw_value(&payload) {
+            self.gateway
+                .forward_fire_and_forget(acp::ExtNotification::new(
+                    crate::agent::roster::SESSIONS_CHANGED_METHOD,
+                    params.into(),
+                ));
+        }
+    }
+    /// Coarse activity of a resident session for the dashboard status column.
+    ///
+    /// Precedence: a non-empty pending-interaction map →
+    /// `NeedsInput` (wins even over a running turn — a session awaiting a
+    /// permission *mid-turn* is "needs input"); else a running turn →
+    /// `Working`; else map the coarse `SessionLiveState`.
+    pub(super) fn resident_activity(
+        &self,
+        id: &acp::SessionId,
+    ) -> crate::agent::roster::RosterActivity {
+        use crate::agent::roster::RosterActivity;
+        let (needs_input, turn_running) = self
+            .sessions
+            .borrow()
+            .get(id)
+            .map(|h| {
+                let needs_input = h
+                    .pending_interactions
+                    .lock()
+                    .map(|g| !g.is_empty())
+                    .unwrap_or(false);
+                let turn_running = h
+                    .current_prompt_id
+                    .lock()
+                    .map(|g| g.is_some())
+                    .unwrap_or(false);
+                (needs_input, turn_running)
+            })
+            .unwrap_or((false, false));
+        if needs_input {
+            return RosterActivity::NeedsInput;
+        }
+        if turn_running {
+            return RosterActivity::Working;
+        }
+        match self.session_live_state.borrow().get(id).copied() {
+            Some(SessionLiveState::Completed) => RosterActivity::Completed,
+            Some(SessionLiveState::DeadFailed) => RosterActivity::Dead,
+            Some(SessionLiveState::Dormant) => RosterActivity::Dormant,
+            _ => RosterActivity::Idle,
+        }
+    }
+    /// Build a single roster entry for a resident session, or `None` if it is
+    /// not currently resident.
+    pub(super) fn resident_roster_entry(
+        &self,
+        id: &acp::SessionId,
+    ) -> Option<crate::agent::roster::RosterEntry> {
+        let session_id = id.0.to_string();
+        let (cwd, is_worktree, model_id, reasoning_effort, yolo) = {
+            let sessions = self.sessions.borrow();
+            let h = sessions.get(id)?;
+            (
+                h.display_cwd.clone().unwrap_or_else(|| h.info.cwd.clone()),
+                h.display_cwd.is_some(),
+                Some(h.model_id.0.to_string()),
+                h.reasoning_effort,
+                h.yolo_mode,
+            )
+        };
+        Some(crate::agent::roster::RosterEntry {
+            title: self
+                .resident_roster_titles
+                .borrow()
+                .get(&session_id)
+                .cloned(),
+            session_id,
+            cwd,
+            is_worktree,
+            model_id,
+            reasoning_effort,
+            yolo,
+            activity: self.resident_activity(id),
+            resident: true,
+            last_change_unix_ms: chrono::Utc::now().timestamp_millis(),
+            origin: crate::agent::roster::RosterOrigin::Local,
+        })
+    }
+    /// Snapshot all resident sessions as roster entries (synchronous; no disk).
+    pub(super) fn resident_roster_entries(&self) -> Vec<crate::agent::roster::RosterEntry> {
+        let ids: Vec<acp::SessionId> = self.sessions.borrow().keys().cloned().collect();
+        ids.iter()
+            .filter_map(|id| self.resident_roster_entry(id))
+            .collect()
+    }
+    /// Build the full roster: resident actors plus recently-touched on-disk
+    /// (`Dormant`) sessions. Resident wins on an id collision; hidden sessions
+    /// are excluded.
+    pub(crate) async fn build_roster(&self) -> Vec<crate::agent::roster::RosterEntry> {
+        let resident = self.resident_roster_entries();
+        let summaries = crate::session::persistence::list_recent_summaries(200)
+            .await
+            .unwrap_or_default();
+        let entries = crate::agent::roster::merge_roster(resident, summaries);
+        self.cache_resident_titles(&entries);
+        entries
+    }
+    /// Refresh `resident_roster_titles` from the freshly-built roster.
+    pub(super) fn cache_resident_titles(&self, entries: &[crate::agent::roster::RosterEntry]) {
+        *self.resident_roster_titles.borrow_mut() = entries
+            .iter()
+            .filter(|e| e.resident)
+            .filter_map(|e| Some((e.session_id.clone(), e.title.clone()?)))
+            .collect();
+    }
+    /// Terminally remove a session: emit the roster delta with its final state,
+    /// then drop it from all maps (no finalize — callers that need finalize do
+    /// it first, see `close_session_explicit`).
+    pub(super) fn remove_session_terminal(
+        &self,
+        id: &acp::SessionId,
+        final_state: SessionLiveState,
+    ) {
+        self.record_roster_delta(id, final_state);
+        self.remove_session(id);
+    }
+    /// Reap a session whose **resident** actor thread exited unexpectedly
+    /// (panic / load failure). Demotes it to `DeadFailed`, emits the roster
+    /// delta, and removes it WITHOUT finalize — the conversation persists on
+    /// disk and stays resumable (reaping a dead actor is harmless;
+    /// it demotes to Dormant).
+    pub(super) fn reap_dead_session(&self, id: &acp::SessionId) {
+        self.remove_session_terminal(id, SessionLiveState::DeadFailed);
+    }
+    /// Sweep `session_threads` for finished threads and clean them up.
+    ///
+    /// A finished thread has two distinct meanings, and conflating them
+    /// corrupts the `SessionLiveState` roster source:
+    ///
+    /// - **Still resident in `sessions`** → the actor exited unexpectedly while
+    ///   the session was hosted (panic / load failure). Reap as `DeadFailed`.
+    /// - **Not resident** (already idle-unloaded → `Dormant`, or explicitly
+    ///   closed) → this is the *expected* clean exit. The `SessionThread` was
+    ///   kept only so `drain_old_session_thread` could wait on it; now that it
+    ///   has finished there is nothing left to drain, so just drop the leftover
+    ///   `SessionThread`/state entries. Do **not** demote to `DeadFailed` and do
+    ///   **not** emit a second roster delta.
+    ///
+    /// `JoinHandle::is_finished()` is non-blocking and cannot distinguish a
+    /// clean exit from a panic on its own, which is exactly why the residency
+    /// check is required. Runs both opportunistically and from the join-handle
+    /// supervisor (`ensure_session_supervisor`).
+    pub(super) fn sweep_dead_sessions(&self) {
+        let dead: Vec<acp::SessionId> = self
+            .session_threads
+            .borrow()
+            .iter()
+            .filter(|(_, t)| t.is_finished())
+            .map(|(id, _)| id.clone())
+            .collect();
+        for id in dead {
+            if self.sessions.borrow().contains_key(&id) {
+                tracing::warn!(
+                    session_id = % id.0,
+                    "Resident session actor exited unexpectedly; reaping as DeadFailed"
+                );
+                self.reap_dead_session(&id);
+            } else {
+                self.session_threads.borrow_mut().remove(&id);
+                self.session_live_state.borrow_mut().remove(&id);
+                tracing::debug!(
+                    session_id = % id.0,
+                    "Reaped finished thread for non-resident session (clean exit)"
+                );
+            }
+        }
+    }
+    /// Start the join-handle supervisor. **Idempotent.**
+    ///
+    /// A single `spawn_local` task periodically reaps actor threads that have
+    /// exited (panicked or finished) so a dead actor never lingers as a roster
+    /// zombie. `std::thread::JoinHandle` is not awaitable, so we poll
+    /// `is_finished()` on a tick — the same mechanism `drain_old_session_thread`
+    /// and `sweep_dead_sessions` already use. A panicked actor is therefore
+    /// reaped within one [`SESSION_SUPERVISOR_TICK`].
+    ///
+    /// The sweep body is wrapped in `catch_unwind` so a single panicking sweep
+    /// can never terminate the loop (which would silently disable reaping for
+    /// the rest of the process). The task holds a `LocalRef` (raw pointer) to
+    /// `self` for the lifetime of the `LocalSet`; this is sound because the
+    /// agent owns the `LocalSet` and outlives it (same contract as
+    /// `start_subagent_coordinator`), and `LocalRef` is `!Send`.
+    pub(super) fn ensure_session_supervisor(&self) {
+        if self.supervisor_started.replace(true) {
+            return;
+        }
+        #[cfg(test)]
+        self.supervisor_spawn_count
+            .set(self.supervisor_spawn_count.get() + 1);
+        let agent_ref = LocalRef::new(self);
+        tokio::task::spawn_local(async move {
+            loop {
+                tokio::time::sleep(SESSION_SUPERVISOR_TICK).await;
+                let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
+                    agent_ref.get().sweep_dead_sessions();
+                }));
+                if result.is_err() {
+                    tracing::error!("session supervisor sweep panicked; continuing supervision");
+                }
+            }
+        });
+    }
+    /// Coarse "any work pending" check for the idle-unload stub.
+    /// Returns `true` while the session has work in flight.
+    ///
+    /// Three layers:
+    /// 1. **Fast path (sync):** the shared `current_prompt_id` slot, which the
+    ///    actor sets while a turn is running (`maybe_start_running_task`) and
+    ///    clears via its RAII guard. A poisoned lock is treated as busy → never
+    ///    unload.
+    /// 1b. **Parked plan-approval (sync):** the shared `pending_interactions`
+    ///    slot. The parked plan-approval resume re-park is the one outstanding work with no
+    ///    running turn, so it needs its own sync check (the same shared-`Arc`
+    ///    idiom as `current_prompt_id`) rather than the async round-trip below.
+    /// 2. **Queue check (async):** when no turn is running, the actor is between
+    ///    turns and responsive, so we ask it whether `pending_inputs` is
+    ///    non-empty (a prompt queued at the turn boundary). This closes the
+    ///    sub-tick window where `current_prompt_id` is momentarily `None` but a
+    ///    queued input is about to be drained. On timeout we keep the session
+    ///    resident (conservative).
+    ///
+    /// TODO(PR-4): once the aggregate `SessionActivity` signal exists, also
+    /// consult the autonomous background sources so a detached session is never
+    /// idle-unloaded (→ `Shutdown` → `KillOnDrop`) while they are live:
+    /// `monitor_event_buffer`, pending scheduler fires,
+    /// `ToolContext.background_tasks`, and background subagent sessions. Until
+    /// then those background-only sessions rely on the keep-resident default and
+    /// the `current_prompt_id` auto-wake turn being active.
+    ///
+    /// TODO(PR-4): this is also inherently a *check-then-act* across the
+    /// actor-thread boundary — work can arrive (a new `Prompt`/auto-wake) in the
+    /// gap between this `IsBusy` answer and the caller's subsequent `Shutdown`,
+    /// so an idle-unload can still race a just-arrived turn. The actor processes
+    /// its mailbox in order, so the lost work is bounded and recoverable on
+    /// reload; PR-4 closes the gap properly by gating the unload inside the
+    /// actor (a single `Unload`-if-idle command) rather than check-then-send.
+    pub(super) async fn session_has_live_work(&self, id: &acp::SessionId) -> bool {
+        let Some(handle) = self.sessions.borrow().get(id).cloned() else {
+            return false;
+        };
+        let turn_running = handle
+            .current_prompt_id
+            .lock()
+            .map(|g| g.is_some())
+            .unwrap_or(true);
+        if turn_running {
+            return true;
+        }
+        if crate::session::pending_interaction::has_parked_plan_approval(
+            &handle.pending_interactions,
+        ) {
+            return true;
+        }
+        tokio::time::timeout(IDLE_QUERY_TIMEOUT, handle.is_busy())
+            .await
+            .unwrap_or(true)
+    }
+}

crates/codegen/xai-grok-shell/src/agent/mvp_agent/subagent_coordinator.rs

@@ -0,0 +1,599 @@
+//! Subagent coordinator drain task and spawn-context construction for [`MvpAgent`].
+//! Co-located child of `mvp_agent` (`use super::*`); tested by `tests/subagent_spawn_context_tests.rs`.
+use super::*;
+impl MvpAgent {
+    /// Start the subagent coordinator drain task.
+    ///
+    /// Takes the `subagent_event_rx` receiver (once) and spawns a `spawn_local` task
+    /// that receives `SubagentRequest`s and delegates each to
+    /// `handle_subagent_request()` on its own `spawn_local` task.
+    ///
+    /// Uses `LocalRef` to reference `self` from
+    /// `spawn_local` closures. Idempotent: subsequent calls are no-ops.
+    pub(super) fn start_subagent_coordinator(&self) {
+        let Some(mut rx) = self.subagent_event_rx.borrow_mut().take() else {
+            return;
+        };
+        let agent_ref = LocalRef::new(self);
+        use crate::agent::subagent::{BlockWaitSlot, is_running, resolve_snapshot};
+        use xai_grok_tools::implementations::grok_build::task::types::{
+            SubagentCancelOutcome, SubagentCancelTarget, SubagentEvent,
+        };
+        tokio::task::spawn_local({
+            let agent_ref = agent_ref.clone();
+            async move {
+                while let Some(event) = rx.recv().await {
+                    match event {
+                        SubagentEvent::Spawn(boxed) => {
+                            let request = *boxed;
+                            let agent_ref = agent_ref.clone();
+                            tokio::task::spawn_local(async move {
+                                let this = agent_ref.get();
+                                let parent_sid = request.parent_session_id.clone();
+                                let mut ctx = this.build_subagent_spawn_context(&parent_sid);
+                                let parent_handle = {
+                                    let parent_sid_acp = acp::SessionId::new(parent_sid.clone());
+                                    this.sessions.borrow().get(&parent_sid_acp).cloned()
+                                };
+                                if let Some(handle) = parent_handle {
+                                    ctx.parent_mcp_pool = handle.snapshot_mcp_pool().await;
+                                    ctx.client_hooks = handle.snapshot_client_hooks().await;
+                                    let parent_tools = handle.snapshot_tool_definitions().await;
+                                    ctx.parent_tool_snapshot =
+                                        (!parent_tools.is_empty()).then_some(parent_tools);
+                                }
+                                crate::agent::subagent::handle_subagent_request(
+                                    request,
+                                    ctx,
+                                    &this.subagent_coordinator,
+                                    &this.gateway,
+                                )
+                                .await;
+                            });
+                        }
+                        SubagentEvent::Query(query) => {
+                            let agent_ref = agent_ref.clone();
+                            tokio::task::spawn_local(async move {
+                                let subagent_id = query.subagent_id;
+                                let block = query.block;
+                                let timeout_ms = query.timeout_ms;
+                                let slot: BlockWaitSlot = std::rc::Rc::new(
+                                    std::cell::RefCell::new(Some(query.respond_to)),
+                                );
+                                let send_via_slot =
+                                    |slot: &BlockWaitSlot, snap| match slot.borrow_mut().take() {
+                                        Some(tx) => tx.send(snap).is_ok(),
+                                        None => false,
+                                    };
+                                let lookup = {
+                                    let this = agent_ref.get();
+                                    let result =
+                                        this.subagent_coordinator.borrow().lookup(&subagent_id);
+                                    if block && result.is_some() {
+                                        this.subagent_coordinator
+                                            .borrow_mut()
+                                            .register_block_wait(&subagent_id, slot.clone());
+                                    }
+                                    this.subagent_coordinator
+                                        .borrow_mut()
+                                        .evict_stale_completed();
+                                    result
+                                };
+                                let snapshot = resolve_snapshot(lookup).await;
+                                let should_block =
+                                    block && snapshot.as_ref().is_some_and(is_running);
+                                if should_block {
+                                    let timeout_ms = timeout_ms.unwrap_or(30_000);
+                                    let deadline = tokio::time::Instant::now()
+                                        + tokio::time::Duration::from_millis(timeout_ms);
+                                    loop {
+                                        tokio::time::sleep(tokio::time::Duration::from_millis(200))
+                                            .await;
+                                        let receiver_gone =
+                                            slot.borrow().as_ref().is_none_or(|tx| tx.is_closed());
+                                        if receiver_gone {
+                                            let this = agent_ref.get();
+                                            let mut coord = this.subagent_coordinator.borrow_mut();
+                                            coord.clear_block_waited(&subagent_id);
+                                            coord.unregister_block_wait(&subagent_id, &slot);
+                                            return;
+                                        }
+                                        let lookup = {
+                                            let this = agent_ref.get();
+                                            this.subagent_coordinator.borrow().lookup(&subagent_id)
+                                        };
+                                        let snap = resolve_snapshot(lookup).await;
+                                        let still_running = snap.as_ref().is_some_and(is_running);
+                                        if !still_running || tokio::time::Instant::now() >= deadline
+                                        {
+                                            {
+                                                let this = agent_ref.get();
+                                                let mut coord =
+                                                    this.subagent_coordinator.borrow_mut();
+                                                if still_running {
+                                                    coord.clear_block_waited(&subagent_id);
+                                                }
+                                                coord.unregister_block_wait(&subagent_id, &slot);
+                                            }
+                                            if !send_via_slot(&slot, snap) && !still_running {
+                                                let this = agent_ref.get();
+                                                this.subagent_coordinator
+                                                    .borrow_mut()
+                                                    .clear_block_waited(&subagent_id);
+                                            }
+                                            return;
+                                        }
+                                    }
+                                } else {
+                                    let delivered = send_via_slot(&slot, snapshot);
+                                    if block {
+                                        let this = agent_ref.get();
+                                        let mut coord = this.subagent_coordinator.borrow_mut();
+                                        coord.unregister_block_wait(&subagent_id, &slot);
+                                        if !delivered {
+                                            coord.clear_block_waited(&subagent_id);
+                                        }
+                                    }
+                                }
+                            });
+                        }
+                        SubagentEvent::Cancel(request) => {
+                            let this = agent_ref.get();
+                            let outcome = {
+                                let mut coord = this.subagent_coordinator.borrow_mut();
+                                match request.target {
+                                    SubagentCancelTarget::SubagentId(ref subagent_id) => {
+                                        coord.mark_explicitly_killed(subagent_id);
+                                        coord.cancel_with_outcome(subagent_id)
+                                    }
+                                    SubagentCancelTarget::ParentPromptId(ref parent_prompt_id) => {
+                                        coord.cancel_by_parent_prompt_id(parent_prompt_id);
+                                        SubagentCancelOutcome::Cancelled
+                                    }
+                                }
+                            };
+                            let _ = request.respond_to.send(outcome);
+                        }
+                        SubagentEvent::ListActive(request) => {
+                            let this = agent_ref.get();
+                            let summaries = this
+                                .subagent_coordinator
+                                .borrow()
+                                .active_summaries_for(&request.parent_session_id);
+                            let _ = request.respond_to.send(summaries);
+                        }
+                        SubagentEvent::Completions(request) => {
+                            let this = agent_ref.get();
+                            let mut completions = this
+                                .subagent_coordinator
+                                .borrow_mut()
+                                .drain_pending_completions();
+                            completions.retain(|c| !request.suppress_ids.contains(&c.subagent_id));
+                            let _ = request.respond_to.send(completions);
+                        }
+                        SubagentEvent::Outstanding(request) => {
+                            let this = agent_ref.get();
+                            let reply = this
+                                .subagent_coordinator
+                                .borrow()
+                                .outstanding_reply_for_prompt(&request.prompt_id);
+                            let _ = request.respond_to.send(reply);
+                        }
+                        SubagentEvent::ClearUsageNotApplied(request) => {
+                            let this = agent_ref.get();
+                            this.subagent_coordinator
+                                .borrow_mut()
+                                .clear_subagent_usage_not_applied(&request.prompt_id);
+                        }
+                        SubagentEvent::MarkUsageNotApplied(request) => {
+                            let this = agent_ref.get();
+                            this.subagent_coordinator
+                                .borrow_mut()
+                                .mark_subagent_usage_not_applied(&request.prompt_id);
+                            let _ = request.respond_to.send(());
+                        }
+                        SubagentEvent::ValidateType(request) => {
+                            let agent_ref = agent_ref.clone();
+                            tokio::task::spawn_local(async move {
+                                let this = agent_ref.get();
+                                let ctx = this
+                                    .build_subagent_validation_context(&request.parent_session_id);
+                                let outcome = crate::agent::subagent::validate_subagent_type(
+                                    &request.subagent_type,
+                                    &ctx,
+                                );
+                                let _ = request.respond_to.send(outcome);
+                            });
+                        }
+                        SubagentEvent::DescribeType(request) => {
+                            let agent_ref = agent_ref.clone();
+                            tokio::task::spawn_local(async move {
+                                use xai_grok_tools::implementations::grok_build::task::types::SubagentDescribeOutcome;
+                                let this = agent_ref.get();
+                                let outcome = match this
+                                    .try_build_subagent_spawn_context(&request.parent_session_id)
+                                {
+                                    Some(ctx) => crate::agent::subagent::describe_subagent_type(
+                                        &request.subagent_type,
+                                        request.harness_agent_type.as_deref(),
+                                        &ctx,
+                                    ),
+                                    None => {
+                                        tracing::warn!(
+                                            parent_session_id = % request.parent_session_id,
+                                            subagent_type = % request.subagent_type,
+                                            "DescribeType for unknown/evicted parent session, replying Unavailable",
+                                        );
+                                        SubagentDescribeOutcome::Unavailable
+                                    }
+                                };
+                                let _ = request.respond_to.send(outcome);
+                            });
+                        }
+                    }
+                }
+            }
+        });
+        {
+            let (trace_tx, mut trace_rx) = tokio::sync::mpsc::unbounded_channel::<
+                crate::upload::turn::SyntheticTurnTraceRequest,
+            >();
+            self.subagent_coordinator.borrow_mut().synthetic_trace_tx = Some(trace_tx);
+            tokio::task::spawn_local({
+                let agent_ref = agent_ref.clone();
+                async move {
+                    while let Some(request) = trace_rx.recv().await {
+                        tokio::task::spawn_local({
+                            let agent_ref = agent_ref.clone();
+                            async move {
+                                handle_synthetic_turn_trace(agent_ref, request).await;
+                            }
+                        });
+                    }
+                }
+            });
+        }
+    }
+    /// Lightweight context for the `SubagentEvent::ValidateType` drain arm;
+    /// tolerates evicted parent sessions (returns built-in defaults + warns).
+    pub(super) fn build_subagent_validation_context(
+        &self,
+        parent_session_id: &str,
+    ) -> crate::agent::subagent::SubagentValidationContext {
+        let parent_sid = acp::SessionId::new(parent_session_id);
+        let (parent_cwd, allowed_subagent_types) = {
+            let sessions = self.sessions.borrow();
+            let ps = sessions.get(&parent_sid);
+            warn_on_missing_parent_session_for_validate_type(parent_session_id, ps.is_some());
+            (
+                ps.map(|h| std::path::PathBuf::from(&h.info.cwd))
+                    .unwrap_or_default(),
+                ps.and_then(|h| h.allowed_subagent_types.clone()),
+            )
+        };
+        let cli_agent_names: Vec<String> = {
+            let cfg = self.cfg.borrow();
+            cfg.cli_agents.iter().map(|d| d.name.clone()).collect()
+        };
+        crate::agent::subagent::SubagentValidationContext {
+            parent_cwd,
+            plugin_registry: self.plugin_registry_handle.snapshot(),
+            subagent_toggle: self.subagent_toggle.clone(),
+            allowed_subagent_types,
+            cli_agent_names,
+        }
+    }
+    /// Build a `SubagentSpawnContext` from the current agent state and the
+    /// parent session's shared resources.
+    ///
+    /// This is the ONLY subagent-related method on MvpAgent besides the
+    /// coordinator startup.
+    /// Build a spawn context for a real subagent spawn. The parent session is
+    /// guaranteed present here because the parent just issued the spawn request,
+    /// so a missing parent is a real invariant violation and panics. Read-only
+    /// callers that can race a parent teardown (e.g. `DescribeType`) must use
+    /// [`Self::try_build_subagent_spawn_context`] instead.
+    pub(super) fn build_subagent_spawn_context(
+        &self,
+        parent_session_id: &str,
+    ) -> crate::agent::subagent::SubagentSpawnContext {
+        self.try_build_subagent_spawn_context(parent_session_id)
+            .expect("parent session must exist when spawning subagents")
+    }
+    /// Fallible variant of [`Self::build_subagent_spawn_context`]: returns
+    /// `None` when the parent `SessionHandle` is absent (evicted / torn down)
+    /// instead of panicking, so read-only paths that can race a teardown can
+    /// fail open.
+    pub(super) fn try_build_subagent_spawn_context(
+        &self,
+        parent_session_id: &str,
+    ) -> Option<crate::agent::subagent::SubagentSpawnContext> {
+        let parent_sid = acp::SessionId::new(parent_session_id);
+        let (
+            parent_model_id,
+            parent_chat_state,
+            parent_cmd_tx,
+            parent_cwd,
+            yolo_mode,
+            parent_depth,
+            hunk_tracker_handle,
+            hunk_tracking_enabled,
+            fs,
+            terminal,
+            session_env,
+            parent_attribution_callback,
+            parent_agent_name,
+            parent_managed_mcp_proxy_base_url,
+        ) = {
+            let sessions = self.sessions.borrow();
+            let ps = sessions.get(&parent_sid);
+            (
+                ps.map(|h| h.model_id.clone())
+                    .unwrap_or_else(|| self.models_manager.current_model_id()),
+                ps.map(|h| h.chat_state_handle.clone()),
+                ps.map(|h| h.cmd_tx.clone()),
+                ps.map(|h| std::path::PathBuf::from(&h.info.cwd))
+                    .unwrap_or_default(),
+                ps.map(|h| h.yolo_mode).unwrap_or(self.default_yolo_mode),
+                ps.map(|h| h.tool_context.subagent_depth).unwrap_or(0),
+                ps.map(|h| h.tool_context.hunk_tracker_handle.clone())
+                    .unwrap_or_else(xai_hunk_tracker::HunkTrackerHandle::noop),
+                ps.map(|h| h.tool_context.hunk_tracking_enabled)
+                    .unwrap_or(false),
+                ps.map(|h| h.tool_context.fs.inner().clone())
+                    .unwrap_or_else(|| {
+                        let cwd = ps
+                            .map(|h| std::path::PathBuf::from(&h.info.cwd))
+                            .unwrap_or_default();
+                        std::sync::Arc::new(xai_grok_workspace::file_system::LocalFs::new(cwd))
+                    }),
+                ps.map(|h| h.tool_context.terminal.clone())
+                    .unwrap_or_else(|| {
+                        std::sync::Arc::new(crate::terminal::TerminalRunner::new(
+                            std::sync::Arc::new(self.gateway.clone()),
+                            parent_sid.clone(),
+                        ))
+                    }),
+                ps.map(|h| h.tool_context.session_env.clone())
+                    .unwrap_or_else(|| std::sync::Arc::new(std::collections::HashMap::new())),
+                ps.and_then(|h| h.attribution_callback.clone()),
+                ps.map(|h| h.agent_name.clone()),
+                ps.map(|h| h.managed_mcp_proxy_base_url.clone()),
+            )
+        };
+        let (
+            parent_workspace_ops,
+            parent_terminal_backend,
+            parent_notification_handle,
+            parent_scheduler_handle,
+        ) = {
+            let sessions = self.sessions.borrow();
+            sessions.get(&parent_sid).map(|ps| {
+                (
+                    ps.workspace_ops.clone(),
+                    ps.terminal_backend.clone(),
+                    ps.tools_notification_handle.clone(),
+                    ps.scheduler_handle.clone(),
+                )
+            })
+        }?;
+        let available_models = self.models_manager.models();
+        let parent_lsp = {
+            let sessions = self.sessions.borrow();
+            sessions
+                .get(&parent_sid)
+                .and_then(|h| h.tool_context.lsp.clone())
+        };
+        let am = self.auth_manager.clone();
+        let inference_idle_timeout_secs = {
+            let per_model = config::find_model_by_id(&available_models, parent_model_id.0.as_ref())
+                .and_then(|e| e.info.inference_idle_timeout_secs);
+            let cfg = self.cfg.borrow();
+            let remote = cfg
+                .remote_settings
+                .as_ref()
+                .and_then(|s| s.inference_idle_timeout_secs);
+            per_model.or(remote).unwrap_or(600).max(10)
+        };
+        let parent_hook_registry = {
+            let sessions = self.sessions.borrow();
+            sessions
+                .get(&parent_sid)
+                .and_then(|h| h.hook_registry.clone())
+        };
+        let parent_max_turns = {
+            let sessions = self.sessions.borrow();
+            sessions.get(&parent_sid).and_then(|h| h.max_turns)
+        };
+        let parent_model_agent_type =
+            config::find_model_by_id(&available_models, parent_model_id.0.as_ref())
+                .map(|e| e.info.agent_type.clone());
+        let ask_user_question_enabled = {
+            let sessions = self.sessions.borrow();
+            sessions
+                .get(&parent_sid)
+                .map(|h| h.ask_user_question_enabled)
+                .unwrap_or_else(|| self.cfg.borrow().resolve_ask_user_question().value)
+        };
+        let (gcs_upload_method, gcs_bucket_url) = match self.trace_upload_config_snapshot() {
+            Some(method) => {
+                use crate::session::repo_changes::UploadMethod;
+                let bucket = match &method {
+                    UploadMethod::Direct { .. } => self
+                        .cfg
+                        .borrow()
+                        .endpoints
+                        .resolve_trace_bucket_url()
+                        .map(|r| r.value),
+                    UploadMethod::Proxy { .. } => Some("proxy-managed".to_string()),
+                    UploadMethod::S3 { bucket, .. } => Some(format!("s3://{bucket}")),
+                };
+                match bucket {
+                    Some(url) => (Some(method), Some(url)),
+                    None => (None, None),
+                }
+            }
+            None => (None, None),
+        };
+        Some(crate::agent::subagent::SubagentSpawnContext {
+            lsp: parent_lsp,
+            gateway: self.gateway.clone(),
+            client_hooks: Default::default(),
+            sampling_config: self.sampling_config.borrow().clone(),
+            managed_mcp_proxy_base_url: parent_managed_mcp_proxy_base_url
+                .unwrap_or_else(|| self.cli_chat_proxy_base_url()),
+            alpha_test_key: self.alpha_test_key(),
+            auth_method_id: self
+                .auth_method_id
+                .load()
+                .as_deref()
+                .cloned()
+                .unwrap_or_else(|| acp::AuthMethodId::new("default")),
+            model_id: parent_model_id,
+            storage_mode: self.storage_mode,
+            auth: self.current_or_buffered_auth(),
+            parent_cwd: parent_cwd.clone(),
+            parent_session_id: parent_session_id.to_string(),
+            yolo_mode,
+            subagent_event_tx: self.subagent_event_tx.clone(),
+            parent_depth,
+            inference_idle_timeout_secs,
+            auto_compact_threshold_tiers:
+                crate::agent::subagent::AutoCompactThresholdTiers::capture(&self.cfg.borrow()),
+            hunk_tracker_handle,
+            hunk_tracking_enabled,
+            fs,
+            terminal,
+            session_env,
+            memory_config: self.memory_config.clone(),
+            web_search_sampling_config: self.prepare_web_search_sampling_config(),
+            web_fetch_config: self.prepare_web_fetch_config(),
+            image_gen_config: self.prepare_image_gen_config(),
+            video_gen_config: self.prepare_video_gen_config(),
+            app_builder_deployer_config: self.prepare_app_builder_deployer_config(),
+            write_file_enabled: self.cfg.borrow().resolve_write_file().value,
+            goal_enabled: self.cfg.borrow().resolve_goal().value,
+            ask_user_question_enabled,
+            parent_cmd_tx: parent_cmd_tx.clone(),
+            parent_session_info: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .map(|h| crate::session::info::Info {
+                        id: parent_sid.clone(),
+                        cwd: h.info.cwd.clone(),
+                    })
+            },
+            parent_chat_state,
+            parent_max_turns,
+            available_models,
+            subagent_model_overrides: self.subagent_model_overrides.clone(),
+            subagent_toggle: self.subagent_toggle.clone(),
+            subagent_roles: self.subagent_roles.clone(),
+            subagent_personas: self.subagent_personas.clone(),
+            persona_io_summaries: self.persona_io_summaries.clone(),
+            disable_web_search: self.cfg.borrow().disable_web_search,
+            todo_gate: self.cfg.borrow().todo_gate,
+            remote_settings: self.cfg.borrow().remote_settings.clone(),
+            laziness_debug_log: self.cfg.borrow().laziness_debug_log.clone(),
+            backend_tools_enabled: self.cfg.borrow().resolve_backend_tools().value,
+            respect_gitignore: self.cfg.borrow().respect_gitignore,
+            path_not_found_hints: self.cfg.borrow().path_not_found_hints,
+            plugin_registry: self.plugin_registry_handle.snapshot(),
+            models_manager: self.models_manager.clone(),
+            file_tool_overrides: {
+                let cfg = self.cfg.borrow();
+                let effective = cfg
+                    .toolset
+                    .resolve_file_toolset(cfg.remote_settings.as_ref());
+                if effective != crate::tools::FileToolset::Standard {
+                    effective.tool_configs(&cfg.toolset.hashline).ok()
+                } else {
+                    None
+                }
+            },
+            gcs_bucket_url,
+            agent_config: Some(self.cfg.borrow().clone()),
+            gcs_upload_method,
+            hook_registry: parent_hook_registry,
+            hook_workspace_root: String::new(),
+            permission_handle: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .map(|h| h.permission_handle.clone())
+            },
+            worktree_type: self.worktree_type,
+            api_key_provider: Some(Arc::new(crate::auth::manager::SharedAuthKeyProvider(
+                am.clone(),
+            ))),
+            image_description_model: self.resolve_image_description_model(),
+            workspace_ops: parent_workspace_ops.clone(),
+            auth_manager: am.clone(),
+            attribution_callback: parent_attribution_callback,
+            parent_agent_name,
+            parent_model_agent_type,
+            allowed_subagent_types: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .and_then(|h| h.allowed_subagent_types.clone())
+            },
+            parent_mcp_configs: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .map(|h| h.mcp_servers.clone())
+                    .unwrap_or_default()
+            },
+            managed_mcp_state: self.managed_mcp_cache.clone(),
+            parent_mcp_pool: None,
+            parent_tool_snapshot: None,
+            parent_skills: None,
+            parent_skills_config: self.cfg.borrow().skills.clone(),
+            parent_compat: self.cfg.borrow().compat_resolved,
+            auto_wake_delivered: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .and_then(|h| h.tool_context.auto_wake_delivered.clone())
+            },
+            synthetic_trace_tx: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .and_then(|h| h.tool_context.synthetic_trace_tx.clone())
+            },
+            task_output_tool_name: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .map(|h| h.tool_context.task_output_tool_name.clone())
+                    .unwrap_or_else(|| {
+                        xai_grok_tools::reminders::task_completion::DEFAULT_TASK_OUTPUT_TOOL
+                            .to_string()
+                    })
+            },
+            auto_wake_enabled: self.cfg.borrow().auto_wake_enabled,
+            goal_loop_active: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .map(|h| h.tool_context.goal_loop_active_gate.clone())
+                    .unwrap_or_else(|| {
+                        std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false))
+                    })
+            },
+            parent_blocking_wait_depth: {
+                let sessions = self.sessions.borrow();
+                sessions
+                    .get(&parent_sid)
+                    .map(|h| h.tool_context.blocking_wait_depth.clone())
+                    .unwrap_or_else(|| std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)))
+            },
+            parent_terminal_backend: parent_terminal_backend.clone(),
+            parent_notification_handle: parent_notification_handle.clone(),
+            parent_scheduler_handle: parent_scheduler_handle.clone(),
+        })
+    }
+}

crates/codegen/xai-grok-shell/src/agent/mvp_agent/tests/subagent_spawn_context_tests.rs

@@ -0,0 +1,135 @@
+//! Subagent spawn-context inheritance: a child session must inherit the parent's
+//! permission handle and goal-loop gate so policy and run-state can't be bypassed
+//! by delegating to a subagent.
+
+use super::{build_minimal_agent_for_tests, make_test_handle};
+use agent_client_protocol as acp;
+use xai_acp_lib::AcpAgentGatewaySender as GatewaySender;
+
+/// Subagents inherit the parent permission handle, so a managed `Read(**/.env)`
+/// deny still blocks the child — direct read and the `cat .env` shell equivalent.
+#[tokio::test]
+async fn subagent_spawn_context_inherits_parent_permission_handle() {
+    use xai_grok_workspace::permission::types::{
+        PatternMode, PermissionConfig, PermissionRule, RuleAction, ToolFilter,
+    };
+
+    let local = tokio::task::LocalSet::new();
+    local
+        .run_until(async {
+            let agent = build_minimal_agent_for_tests();
+            let sid = acp::SessionId::new("parent-permission");
+            let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
+            let gateway = GatewaySender::new(tx);
+            let cwd = xai_grok_paths::AbsPathBuf::new(std::path::PathBuf::from("/tmp"))
+                .expect("absolute cwd");
+            let (permission_handle, _events_rx) =
+                xai_grok_workspace::permission::spawn_permission_manager(
+                    sid.clone(),
+                    gateway,
+                    cwd,
+                    xai_grok_workspace::permission::types::ClientType::Generic,
+                    Some(PermissionConfig::new(vec![PermissionRule {
+                        action: RuleAction::Deny,
+                        tool: ToolFilter::Read,
+                        pattern: Some("**/.env".to_owned()),
+                        pattern_mode: PatternMode::Glob,
+                    }])),
+                    Vec::new(), // deny_read_globs
+                    Vec::new(),
+                    false,
+                    None,
+                );
+
+            let mut handle = make_test_handle("test-model", false, None);
+            handle.permission_handle = permission_handle;
+            agent.sessions.borrow_mut().insert(sid.clone(), handle);
+
+            let ctx = agent.build_subagent_spawn_context(sid.0.as_ref());
+            let inherited = ctx
+                .permission_handle
+                .expect("subagent context must inherit parent permission handle");
+
+            // Direct file read and the shell equivalent both hit the parent deny.
+            for access in [
+                xai_grok_workspace::permission::AccessKind::Read(Some(".env".into())),
+                xai_grok_workspace::permission::AccessKind::Bash("cat .env".into()),
+            ] {
+                let decision = inherited
+                    .request(
+                        access.clone(),
+                        acp::ToolCallUpdate::new(acp::ToolCallId::new("tc"), Default::default()),
+                        Some("child-session".to_owned()),
+                        Some("general-purpose".to_owned()),
+                        Some("permission inheritance regression".to_owned()),
+                    )
+                    .await;
+                assert!(
+                    matches!(
+                        decision,
+                        xai_grok_workspace::permission::Decision::PolicyDeny(_)
+                    ),
+                    "subagent-inherited handle must enforce parent deny for {access:?}, got {decision:?}"
+                );
+            }
+        })
+        .await;
+}
+
+/// A subagent shares the parent's `goal_loop_active_gate` Arc, so flipping the
+/// parent gate is observed through the child context (same allocation).
+#[tokio::test]
+async fn subagent_spawn_context_shares_parent_goal_loop_gate() {
+    use std::sync::atomic::Ordering::Relaxed;
+
+    let agent = build_minimal_agent_for_tests();
+    let sid = acp::SessionId::new("parent-goal");
+    let handle = make_test_handle("test-model", false, None);
+    // Clone the parent's live gate before the handle moves into `sessions`.
+    let parent_gate = handle.tool_context.goal_loop_active_gate.clone();
+    agent.sessions.borrow_mut().insert(sid.clone(), handle);
+
+    let ctx = agent.build_subagent_spawn_context(sid.0.as_ref());
+
+    // Flipping the parent gate must surface through the child flag (shared Arc).
+    assert!(!ctx.goal_loop_active.load(Relaxed));
+    parent_gate.store(true, Relaxed);
+    assert!(
+        ctx.goal_loop_active.load(Relaxed),
+        "subagent context must observe the parent's goal-loop gate (same Arc)"
+    );
+}
+
+/// A subagent inherits the parent session's `ask_user_question` gate, so
+/// `--no-ask-user` strips the tool from subagents too, while the default keeps it.
+#[tokio::test]
+async fn subagent_spawn_context_inherits_parent_ask_user_question_gate() {
+    let agent = build_minimal_agent_for_tests();
+
+    // Parent with the tool disabled (the `--no-ask-user` case) → child off.
+    let sid_off = acp::SessionId::new("parent-no-ask");
+    let mut handle_off = make_test_handle("test-model", false, None);
+    handle_off.ask_user_question_enabled = false;
+    agent
+        .sessions
+        .borrow_mut()
+        .insert(sid_off.clone(), handle_off);
+    let ctx_off = agent.build_subagent_spawn_context(sid_off.0.as_ref());
+    assert!(
+        !ctx_off.ask_user_question_enabled,
+        "subagent must inherit the parent's disabled ask_user_question gate (--no-ask-user)"
+    );
+
+    // Parent with the tool enabled (the default) → child on.
+    let sid_on = acp::SessionId::new("parent-ask");
+    let handle_on = make_test_handle("test-model", false, None);
+    agent
+        .sessions
+        .borrow_mut()
+        .insert(sid_on.clone(), handle_on);
+    let ctx_on = agent.build_subagent_spawn_context(sid_on.0.as_ref());
+    assert!(
+        ctx_on.ask_user_question_enabled,
+        "subagent must inherit the parent's enabled ask_user_question gate"
+    );
+}

crates/codegen/xai-grok-shell/src/agent/proxy.rs

@@ -0,0 +1,619 @@
+//! HTTP CONNECT proxy support for WebSocket connections.
+//!
+//! When running behind a corporate egress proxy,
+//! `tokio-tungstenite`'s `connect_async` cannot reach external
+//! hosts directly because it does not read the standard `HTTPS_PROXY` /
+//! `HTTP_PROXY` environment variables.
+//!
+//! This module provides:
+//! - [`resolve_proxy_for_host`]: reads proxy env vars and `NO_PROXY`, returning
+//!   the proxy URL to use for a given target host (or `None` for direct).
+//! - [`connect_via_proxy`]: opens a TCP connection to the proxy, sends an HTTP
+//!   CONNECT request to create a tunnel, wraps the result in TLS, and returns a
+//!   stream suitable for `tokio_tungstenite::client_async`.
+
+use std::sync::{Arc, OnceLock};
+
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tokio::net::TcpStream;
+use tokio_tungstenite::MaybeTlsStream;
+use tracing::debug;
+
+// ---------------------------------------------------------------------------
+// Environment-variable resolution
+// ---------------------------------------------------------------------------
+
+/// Read proxy configuration from the environment and decide whether `target_host`
+/// should be connected through a proxy.
+///
+/// Resolution order (matches `curl` / `reqwest` behaviour):
+/// 1. If `NO_PROXY` contains `target_host` (or a matching domain suffix / CIDR),
+///    return `None`.
+/// 2. If `HTTPS_PROXY` (or `https_proxy`) is set, return its value.
+/// 3. If `HTTP_PROXY` (or `http_proxy`) is set, return its value.
+/// 4. Otherwise return `None`.
+pub fn resolve_proxy_for_host(target_host: &str) -> Option<String> {
+    resolve_proxy_for_host_with(target_host, |key| std::env::var(key))
+}
+
+/// Testable inner implementation that accepts a custom env-var reader.
+fn resolve_proxy_for_host_with<F>(target_host: &str, env: F) -> Option<String>
+where
+    F: for<'a> Fn(&'a str) -> Result<String, std::env::VarError>,
+{
+    // Check NO_PROXY / no_proxy.
+    let no_proxy = env("NO_PROXY")
+        .or_else(|_| env("no_proxy"))
+        .unwrap_or_default();
+    if is_host_bypassed(target_host, &no_proxy) {
+        return None;
+    }
+
+    // HTTPS_PROXY takes precedence (our target is always wss://).
+    if let Ok(url) = env("HTTPS_PROXY").or_else(|_| env("https_proxy")) {
+        let url = url.trim().to_string();
+        if !url.is_empty() {
+            return Some(url);
+        }
+    }
+
+    // Fall back to HTTP_PROXY.
+    if let Ok(url) = env("HTTP_PROXY").or_else(|_| env("http_proxy")) {
+        let url = url.trim().to_string();
+        if !url.is_empty() {
+            return Some(url);
+        }
+    }
+
+    None
+}
+
+/// Check whether `host` is in the `no_proxy` list.
+///
+/// The `no_proxy` value is a comma-separated list of hostnames, domain
+/// suffixes (with or without a leading dot), IP addresses, or CIDR ranges.
+/// The special value `*` matches everything.
+fn is_host_bypassed(host: &str, no_proxy: &str) -> bool {
+    let host_lower = host.to_ascii_lowercase();
+    for entry in no_proxy.split(',') {
+        let entry = entry.trim().to_ascii_lowercase();
+        if entry.is_empty() {
+            continue;
+        }
+        // Wildcard — bypass all hosts.
+        if entry == "*" {
+            return true;
+        }
+        // Exact match.
+        if host_lower == entry {
+            return true;
+        }
+        // Domain suffix match: ".example.com" matches "foo.example.com".
+        // Also handle the common convention of omitting the leading dot:
+        // "example.com" in NO_PROXY should match "sub.example.com".
+        let matches_suffix = if entry.starts_with('.') {
+            host_lower.ends_with(entry.as_str())
+        } else {
+            host_lower.len() > entry.len()
+                && host_lower.ends_with(entry.as_str())
+                && host_lower.as_bytes()[host_lower.len() - entry.len() - 1] == b'.'
+        };
+        if matches_suffix {
+            return true;
+        }
+        // CIDR / IP matching is intentionally omitted here — our target host
+        // is always a DNS name, not an IP literal. Keeping this simple avoids
+        // pulling in a CIDR parsing dependency.
+    }
+    false
+}
+
+// ---------------------------------------------------------------------------
+// HTTP CONNECT tunnel
+// ---------------------------------------------------------------------------
+
+/// Establish a TLS-wrapped TCP stream through an HTTP CONNECT proxy.
+///
+/// Steps:
+/// 1. Parse the proxy URL to get host + port.
+/// 2. Open a TCP connection to the proxy and perform the CONNECT handshake.
+/// 3. Wrap the tunnel in TLS (using rustls with native root certificates).
+/// 4. Return the stream as `MaybeTlsStream<TcpStream>` so it is compatible
+///    with `tokio_tungstenite::client_async`.
+pub async fn connect_via_proxy(
+    proxy_url: &str,
+    target_host: &str,
+    target_port: u16,
+) -> anyhow::Result<MaybeTlsStream<TcpStream>> {
+    let stream = open_connect_tunnel(proxy_url, target_host, target_port).await?;
+    let tls_stream = tls_wrap(stream, target_host).await?;
+    Ok(MaybeTlsStream::Rustls(tls_stream))
+}
+
+/// Open a raw TCP tunnel through an HTTP CONNECT proxy (no TLS).
+///
+/// 1. Parse the proxy URL to get host + port.
+/// 2. Open a plain TCP connection to the proxy.
+/// 3. Send `CONNECT target_host:target_port HTTP/1.1\r\n\r\n`.
+/// 4. Read the proxy's response; expect `HTTP/1.x 200 …`.
+/// 5. Return the raw `TcpStream` positioned after the CONNECT response.
+async fn open_connect_tunnel(
+    proxy_url: &str,
+    target_host: &str,
+    target_port: u16,
+) -> anyhow::Result<TcpStream> {
+    // 1. Parse proxy URL.
+    let (proxy_host, proxy_port) = parse_proxy_url(proxy_url)?;
+
+    // 2. TCP connect to proxy.
+    let proxy_addr = format!("{proxy_host}:{proxy_port}");
+    debug!(proxy_addr = %proxy_addr, "Opening TCP to proxy");
+    let stream = TcpStream::connect(&proxy_addr)
+        .await
+        .map_err(|e| anyhow::anyhow!("Failed to connect to proxy at {proxy_addr}: {e}"))?;
+
+    // 3. Send HTTP CONNECT.
+    let connect_req = format!(
+        "CONNECT {target_host}:{target_port} HTTP/1.1\r\n\
+         Host: {target_host}:{target_port}\r\n\
+         \r\n"
+    );
+    let (reader_half, mut writer_half) = stream.into_split();
+    writer_half.write_all(connect_req.as_bytes()).await?;
+    writer_half.flush().await?;
+
+    // 4. Read the status line from the proxy.
+    let mut reader = BufReader::new(reader_half);
+    let mut status_line = String::new();
+    reader.read_line(&mut status_line).await?;
+    debug!(status_line = %status_line.trim(), "Proxy CONNECT response");
+
+    if !status_line.starts_with("HTTP/1.1 200") && !status_line.starts_with("HTTP/1.0 200") {
+        anyhow::bail!("Proxy CONNECT failed: {}", status_line.trim());
+    }
+
+    // Consume remaining response headers (until empty line).
+    loop {
+        let mut line = String::new();
+        reader.read_line(&mut line).await?;
+        if line.trim().is_empty() {
+            break;
+        }
+    }
+
+    // 5. Assert the BufReader's internal buffer is empty before reuniting.
+    // BufReader::read_line may have read ahead into its buffer. If extra
+    // bytes were consumed beyond the HTTP headers (e.g., from a proxy that
+    // eagerly forwards data or coalesced TCP segments), dropping them would
+    // corrupt the subsequent TLS handshake.
+    let remaining = reader.buffer();
+    if !remaining.is_empty() {
+        anyhow::bail!(
+            "Proxy sent {} unexpected byte(s) after CONNECT response headers",
+            remaining.len()
+        );
+    }
+
+    // 6. Reunite the split halves back into a TcpStream.
+    let stream = reader.into_inner().reunite(writer_half)?;
+    Ok(stream)
+}
+
+/// Lazily-initialized TLS client configuration.
+///
+/// Loading native root certificates involves syscalls (reading `/etc/ssl/certs/`
+/// or the macOS Keychain) and the cert store never changes at runtime.  We build
+/// the `ClientConfig` once and reuse it across all proxy connections / reconnects.
+///
+/// Stores `Ok(config)` on success or `Err(message)` if cert loading fails.
+static TLS_CONFIG: OnceLock<Result<Arc<rustls::ClientConfig>, String>> = OnceLock::new();
+
+/// Build (or return the cached) TLS client configuration.
+fn get_tls_config() -> anyhow::Result<Arc<rustls::ClientConfig>> {
+    let result = TLS_CONFIG.get_or_init(|| {
+        let mut root_store = rustls::RootCertStore::empty();
+        let cert_result = rustls_native_certs::load_native_certs();
+        if cert_result.certs.is_empty() {
+            let errors: Vec<_> = cert_result.errors.iter().map(|e| e.to_string()).collect();
+            return Err(format!(
+                "No native root certificates found. Errors: {}",
+                if errors.is_empty() {
+                    "(none)".to_string()
+                } else {
+                    errors.join("; ")
+                }
+            ));
+        }
+        for cert in cert_result.certs {
+            if let Err(e) = root_store.add(cert) {
+                tracing::warn!(error = %e, "Skipping unparseable native root certificate");
+            }
+        }
+
+        let config = rustls::ClientConfig::builder()
+            .with_root_certificates(root_store)
+            .with_no_client_auth();
+        Ok(Arc::new(config))
+    });
+
+    match result {
+        Ok(config) => Ok(config.clone()),
+        Err(msg) => anyhow::bail!("{msg}"),
+    }
+}
+
+/// Perform a TLS handshake over an existing TCP stream using rustls with
+/// native root certificates (cached via [`TLS_CONFIG`]).
+async fn tls_wrap(
+    stream: TcpStream,
+    server_name: &str,
+) -> anyhow::Result<tokio_rustls::client::TlsStream<TcpStream>> {
+    let tls_config = get_tls_config()?;
+    let connector = tokio_rustls::TlsConnector::from(tls_config);
+    let dns_name = rustls::pki_types::ServerName::try_from(server_name.to_string())
+        .map_err(|e| anyhow::anyhow!("Invalid TLS server name '{server_name}': {e}"))?;
+
+    let tls_stream = connector
+        .connect(dns_name, stream)
+        .await
+        .map_err(|e| anyhow::anyhow!("TLS handshake through proxy failed: {e}"))?;
+
+    Ok(tls_stream)
+}
+
+/// Parse a proxy URL into (host, port).
+///
+/// Accepted formats:
+/// - `http://host:port`
+/// - `http://host` (defaults to port 80)
+/// - `host:port`
+fn parse_proxy_url(url: &str) -> anyhow::Result<(String, u16)> {
+    // Strip scheme if present.
+    let without_scheme = url
+        .strip_prefix("http://")
+        .or_else(|| url.strip_prefix("https://"))
+        .unwrap_or(url);
+
+    // Strip trailing path/slash.
+    let authority = without_scheme.split('/').next().unwrap_or(without_scheme);
+
+    if let Some((host, port_str)) = authority.rsplit_once(':') {
+        let port: u16 = port_str
+            .parse()
+            .map_err(|_| anyhow::anyhow!("Invalid proxy port in '{url}'"))?;
+        Ok((host.to_string(), port))
+    } else {
+        // No port — default to 80 for HTTP proxies.
+        Ok((authority.to_string(), 80))
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
+    // ===== parse_proxy_url =====
+
+    #[test]
+    fn test_parse_proxy_url_with_scheme_and_port() {
+        let (host, port) = parse_proxy_url("http://proxy.example.com:3140").unwrap();
+        assert_eq!(host, "proxy.example.com");
+        assert_eq!(port, 3140);
+    }
+
+    #[test]
+    fn test_parse_proxy_url_without_scheme() {
+        let (host, port) = parse_proxy_url("proxy.example.com:8080").unwrap();
+        assert_eq!(host, "proxy.example.com");
+        assert_eq!(port, 8080);
+    }
+
+    #[test]
+    fn test_parse_proxy_url_without_port() {
+        let (host, port) = parse_proxy_url("http://proxy.example.com").unwrap();
+        assert_eq!(host, "proxy.example.com");
+        assert_eq!(port, 80);
+    }
+
+    #[test]
+    fn test_parse_proxy_url_with_trailing_slash() {
+        let (host, port) = parse_proxy_url("http://proxy.example.com:3140/").unwrap();
+        assert_eq!(host, "proxy.example.com");
+        assert_eq!(port, 3140);
+    }
+
+    #[test]
+    fn test_parse_proxy_url_https_scheme() {
+        let (host, port) = parse_proxy_url("https://secure-proxy:443").unwrap();
+        assert_eq!(host, "secure-proxy");
+        assert_eq!(port, 443);
+    }
+
+    #[test]
+    fn test_parse_proxy_url_multi_label_host() {
+        let (host, port) =
+            parse_proxy_url("http://http-proxy.services.internal.example:3128").unwrap();
+        assert_eq!(host, "http-proxy.services.internal.example");
+        assert_eq!(port, 3128);
+    }
+
+    #[test]
+    fn test_parse_proxy_url_invalid_port() {
+        assert!(parse_proxy_url("http://proxy:notaport").is_err());
+    }
+
+    // ===== is_host_bypassed =====
+
+    #[test]
+    fn test_bypass_exact_match() {
+        assert!(is_host_bypassed("localhost", "localhost,127.0.0.1"));
+    }
+
+    #[test]
+    fn test_bypass_domain_suffix_with_dot() {
+        assert!(is_host_bypassed(
+            "api.corp.example",
+            "localhost,.corp.example"
+        ));
+    }
+
+    #[test]
+    fn test_bypass_domain_suffix_without_dot() {
+        // Common convention: "example.com" in NO_PROXY matches "api.example.com".
+        assert!(is_host_bypassed("api.example.com", "localhost,example.com"));
+    }
+
+    #[test]
+    fn test_bypass_wildcard() {
+        assert!(is_host_bypassed("anything.example.com", "*"));
+    }
+
+    #[test]
+    fn test_no_bypass_when_not_listed() {
+        assert!(!is_host_bypassed(
+            "api.external.example",
+            "localhost,127.0.0.1,.corp.example,.internal.example"
+        ));
+    }
+
+    #[test]
+    fn test_bypass_case_insensitive() {
+        assert!(is_host_bypassed("API.Corp.EXAMPLE", ".corp.example"));
+    }
+
+    #[test]
+    fn test_bypass_empty_no_proxy() {
+        assert!(!is_host_bypassed("api.external.example", ""));
+    }
+
+    #[test]
+    fn test_bypass_spaces_in_entries() {
+        assert!(is_host_bypassed(
+            "foo.example.com",
+            " localhost , .example.com , .other.com "
+        ));
+    }
+
+    #[test]
+    fn test_bypass_cidr_not_matched_for_dns_names() {
+        // CIDR entries like 10.0.0.0/8 should not match DNS names.
+        assert!(!is_host_bypassed("api.external.example", "10.0.0.0/8"));
+    }
+
+    #[test]
+    fn test_bypass_combined_no_proxy_list() {
+        // A typical corporate NO_PROXY mixes loopback, private CIDRs, and domain suffixes.
+        let no_proxy = "localhost,127.0.0.1,10.0.0.0/8,.internal.example,.corp.example";
+        assert!(!is_host_bypassed("api.external.example", no_proxy));
+        assert!(is_host_bypassed("db.internal.example", no_proxy));
+        assert!(is_host_bypassed("git.corp.example", no_proxy));
+        assert!(is_host_bypassed("localhost", no_proxy));
+    }
+
+    // ===== resolve_proxy_for_host_with =====
+
+    #[test]
+    fn test_resolve_no_proxy_vars_set() {
+        let result = resolve_proxy_for_host_with("api.external.example", |_| {
+            Err(std::env::VarError::NotPresent)
+        });
+        assert_eq!(result, None);
+    }
+
+    #[test]
+    fn test_resolve_https_proxy_used() {
+        let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
+            "HTTPS_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
+            "NO_PROXY" => Err(std::env::VarError::NotPresent),
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
+    }
+
+    #[test]
+    fn test_resolve_http_proxy_fallback() {
+        let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
+            "HTTP_PROXY" => Ok("http://proxy.example.com:8080".to_string()),
+            "NO_PROXY" => Err(std::env::VarError::NotPresent),
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(result, Some("http://proxy.example.com:8080".to_string()));
+    }
+
+    #[test]
+    fn test_resolve_no_proxy_bypasses() {
+        let result = resolve_proxy_for_host_with("api.corp.example", |key| match key {
+            "HTTPS_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
+            "NO_PROXY" => Ok("localhost,.corp.example".to_string()),
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(result, None);
+    }
+
+    #[test]
+    fn test_resolve_https_proxy_takes_precedence() {
+        let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
+            "HTTPS_PROXY" => Ok("http://https-proxy.example.com:443".to_string()),
+            "HTTP_PROXY" => Ok("http://http-proxy.example.com:80".to_string()),
+            "NO_PROXY" => Err(std::env::VarError::NotPresent),
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(
+            result,
+            Some("http://https-proxy.example.com:443".to_string())
+        );
+    }
+
+    #[test]
+    fn test_resolve_lowercase_env_vars() {
+        let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
+            "https_proxy" => Ok("http://proxy.example.com:3128".to_string()),
+            "no_proxy" => Err(std::env::VarError::NotPresent),
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
+    }
+
+    #[test]
+    fn test_resolve_empty_proxy_ignored() {
+        let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
+            "HTTPS_PROXY" => Ok("  ".to_string()),
+            "HTTP_PROXY" => Ok("http://proxy.example.com:8080".to_string()),
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(result, Some("http://proxy.example.com:8080".to_string()));
+    }
+
+    #[test]
+    fn test_resolve_respects_no_proxy_when_proxy_set() {
+        let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
+            "HTTPS_PROXY" | "HTTP_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
+            "NO_PROXY" => {
+                Ok("localhost,127.0.0.1,10.0.0.0/8,.internal.example,.corp.example".to_string())
+            }
+            _ => Err(std::env::VarError::NotPresent),
+        });
+        assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
+    }
+
+    // ===== HTTP CONNECT tunnel (integration-style) =====
+
+    /// Helper: spawn a mock HTTP CONNECT proxy that accepts one connection.
+    ///
+    /// On receiving a CONNECT request, it validates the request format,
+    /// replies with `status_line`, and then echoes data (simulating a tunnel).
+    /// Returns the proxy's listen address.
+    async fn spawn_mock_proxy(status_line: &'static str) -> std::net::SocketAddr {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+
+        tokio::spawn(async move {
+            let (mut stream, _) = listener.accept().await.unwrap();
+
+            // Read CONNECT request (read until \r\n\r\n).
+            let mut buf = vec![0u8; 4096];
+            let mut total = 0;
+            loop {
+                let n = stream.read(&mut buf[total..]).await.unwrap();
+                if n == 0 {
+                    return;
+                }
+                total += n;
+                let so_far = std::str::from_utf8(&buf[..total]).unwrap_or("");
+                if so_far.contains("\r\n\r\n") {
+                    break;
+                }
+            }
+
+            let request = std::str::from_utf8(&buf[..total]).unwrap().to_string();
+            assert!(
+                request.contains("CONNECT ") && request.contains(" HTTP/1.1"),
+                "Expected CONNECT request, got: {request}"
+            );
+
+            // Reply with the provided status line.
+            stream.write_all(status_line.as_bytes()).await.unwrap();
+
+            // Echo loop (simulates the transparent tunnel).
+            let mut echo_buf = [0u8; 1024];
+            loop {
+                let n = match stream.read(&mut echo_buf).await {
+                    Ok(0) | Err(_) => break,
+                    Ok(n) => n,
+                };
+                if stream.write_all(&echo_buf[..n]).await.is_err() {
+                    break;
+                }
+            }
+        });
+
+        addr
+    }
+
+    /// Tests that `open_connect_tunnel` sends a correct CONNECT request,
+    /// parses the proxy's 200 response, and returns a usable tunnel stream.
+    #[tokio::test]
+    async fn test_open_connect_tunnel_success() {
+        let addr =
+            spawn_mock_proxy("HTTP/1.1 200 Connection Established\r\nServer: mock\r\n\r\n").await;
+        let proxy_url = format!("http://{addr}");
+
+        // Call the real function under test.
+        let mut stream = open_connect_tunnel(&proxy_url, "example.com", 443)
+            .await
+            .expect("tunnel should succeed");
+
+        // Verify the tunnel works by echoing data through it.
+        stream.write_all(b"hello tunnel").await.unwrap();
+        stream.flush().await.unwrap();
+
+        let mut response = vec![0u8; 12];
+        stream.read_exact(&mut response).await.unwrap();
+        assert_eq!(&response, b"hello tunnel");
+    }
+
+    /// Tests that `open_connect_tunnel` with a non-default port sends the
+    /// correct CONNECT target.
+    #[tokio::test]
+    async fn test_open_connect_tunnel_custom_port() {
+        let addr = spawn_mock_proxy("HTTP/1.1 200 OK\r\n\r\n").await;
+        let proxy_url = format!("http://{addr}");
+
+        let stream = open_connect_tunnel(&proxy_url, "internal.example.com", 8443).await;
+        assert!(stream.is_ok(), "tunnel should succeed for custom port");
+    }
+
+    /// Tests that `open_connect_tunnel` returns an error when the proxy
+    /// rejects the CONNECT request with a non-200 status.
+    #[tokio::test]
+    async fn test_open_connect_tunnel_proxy_rejects() {
+        let addr = spawn_mock_proxy("HTTP/1.1 403 Forbidden\r\n\r\n").await;
+        let proxy_url = format!("http://{addr}");
+
+        let result = open_connect_tunnel(&proxy_url, "blocked.example.com", 443).await;
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("403"),
+            "Error should mention 403: {err_msg}"
+        );
+    }
+
+    /// Tests that `open_connect_tunnel` returns an error when connecting
+    /// to a proxy that isn't listening.
+    #[tokio::test]
+    async fn test_open_connect_tunnel_proxy_unreachable() {
+        let result = open_connect_tunnel("http://127.0.0.1:1", "example.com", 443).await;
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("Failed to connect to proxy"),
+            "Error should mention proxy connection failure: {err_msg}"
+        );
+    }
+}

crates/codegen/xai-grok-shell/src/agent/restore_code.rs

@@ -0,0 +1,53 @@
+//! Thin wire-format adapter that wraps the shared
+//! [`xai_grok_workspace::session::git::build_restore_decision`] helper
+//! into the JSON shape emitted by `LoadSession` on `_meta.codeRestore`.
+use serde_json::Value;
+use xai_grok_workspace::session::git::{
+    CheckoutSessionOutcome, RestoreKind, build_restore_decision,
+};
+/// Build the `codeRestore` JSON meta, or `None` when no restore should
+/// be reported (no checkout AND no archive applied). The shared
+/// [`build_restore_decision`] is the source of truth; this function
+/// only adapts the result into the wire JSON shape used by the
+/// non-worktree path.
+pub(crate) fn build_code_restore_meta(
+    target_sha: &str,
+    outcome: &CheckoutSessionOutcome,
+    kind: RestoreKind,
+) -> Option<Value> {
+    let decision = build_restore_decision(Some(target_sha), outcome, kind);
+    let summary = decision.summary?;
+    Some(serde_json::json!(
+        { "restored" : decision.restored, "summary" : summary, "degree" : decision
+        .degree, }
+    ))
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    fn outcome(
+        checked_out: bool,
+        stash_ref: Option<&str>,
+        skipped: Option<&str>,
+    ) -> CheckoutSessionOutcome {
+        CheckoutSessionOutcome {
+            checked_out,
+            stash_ref: stash_ref.map(str::to_owned),
+            stash_skipped_reason: skipped.map(str::to_owned),
+        }
+    }
+    #[test]
+    fn checkout_failed_emits_restored_false_meta() {
+        let meta = build_code_restore_meta(
+            "0123456789abcdef",
+            &outcome(false, None, Some("MERGE_HEAD present")),
+            RestoreKind::RegistryOff,
+        )
+        .unwrap();
+        assert_eq!(meta["restored"], false);
+        assert!(meta["degree"].is_null());
+        let s = meta["summary"].as_str().unwrap();
+        assert!(s.contains("restore aborted"));
+        assert!(s.contains("MERGE_HEAD present"));
+    }
+}

crates/codegen/xai-grok-shell/src/agent/roster.rs

@@ -0,0 +1,311 @@
+//! Roster types for the multi-client FleetView dashboard.
+//!
+//! The roster is a list
+//! of dashboard-sized summaries of every session the leader hosts (resident
+//! actors) plus recently-touched on-disk (`Dormant`) sessions. Clients read it
+//! two ways:
+//!
+//!   - request/response `x.ai/sessions/list` → `{ "sessions": [RosterEntry, …] }`
+//!   - broadcast notification `x.ai/sessions/changed` →
+//!     `{ "upserted": [RosterEntry, …], "removed": ["sess-abc", …] }`
+//!
+//! The wire shape is intentionally small and current-state only — no event
+//! fold or materialized snapshot is required (the snapshot is deferred).
+
+use serde::{Deserialize, Serialize};
+use xai_grok_sampling_types::ReasoningEffort;
+
+use crate::session::persistence::Summary;
+
+/// Coarse activity of a session as rendered in the dashboard's status column.
+///
+/// Mirrors the design's `SessionActivity` at dashboard granularity. A full
+/// background-work breakdown (bg tasks / monitors / scheduler / subagents)
+/// lands with a richer `SessionActivity`; the dashboard only needs this
+/// coarse signal to pick a status glyph.
+#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq)]
+#[serde(rename_all = "snake_case")]
+pub enum RosterActivity {
+    /// A turn (user-originated or autonomous) is running.
+    Working,
+    /// Resident, no turn in flight.
+    Idle,
+    /// A permission / question / plan-approval is pending.
+    NeedsInput,
+    /// On disk, not resident.
+    Dormant,
+    /// Finished and resumable.
+    Completed,
+    /// Actor panicked / load failed.
+    Dead,
+}
+
+/// Where the session lives. Only `Local` is produced today; `Remote` is
+/// reserved for cross-machine roster aggregation.
+#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
+#[serde(rename_all = "snake_case", tag = "kind")]
+pub enum RosterOrigin {
+    Local,
+    Remote { host: String },
+}
+
+/// One dashboard row.
+#[derive(Serialize, Deserialize, Clone, Debug)]
+#[serde(rename_all = "camelCase")]
+pub struct RosterEntry {
+    pub session_id: String,
+    /// Generated/display title, if known. Clients fall back to cwd / id.
+    #[serde(default)]
+    pub title: Option<String>,
+    pub cwd: String,
+    pub is_worktree: bool,
+    #[serde(default)]
+    pub model_id: Option<String>,
+    /// Per-session reasoning effort for `model_id`. Carried alongside the model
+    /// so clients can render the session's effort in the roster without a
+    /// separate `model_state` fetch. `None` means "use the model/global
+    /// default" (or the session predates per-session effort persistence).
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub reasoning_effort: Option<ReasoningEffort>,
+    pub yolo: bool,
+    pub activity: RosterActivity,
+    /// `true` while a resident actor hosts the session (vs. read from disk).
+    pub resident: bool,
+    /// Best-effort last-change timestamp (unix millis). Used for sort order.
+    pub last_change_unix_ms: i64,
+    pub origin: RosterOrigin,
+}
+
+/// Response payload for `x.ai/sessions/list`.
+#[derive(Serialize, Deserialize, Clone, Debug, Default)]
+pub struct RosterListResponse {
+    pub sessions: Vec<RosterEntry>,
+}
+
+/// Params payload for the `x.ai/sessions/changed` broadcast notification.
+#[derive(Serialize, Deserialize, Clone, Debug, Default)]
+pub struct RosterChanged {
+    #[serde(default)]
+    pub upserted: Vec<RosterEntry>,
+    #[serde(default)]
+    pub removed: Vec<String>,
+}
+
+/// JSON-RPC method names for the roster API.
+pub const SESSIONS_LIST_METHOD: &str = "x.ai/sessions/list";
+pub const SESSIONS_CHANGED_METHOD: &str = "x.ai/sessions/changed";
+
+/// Merge live `resident` rows with on-disk `summaries` into the sorted roster.
+/// Pure, so it is unit-testable without disk or a live actor.
+///
+/// Resident rows own the live state but carry no title or last-active time, so
+/// each adopts those from its summary — except a `Working` row keeps its "now"
+/// timestamp. Summaries with no resident row become `Dormant`; keying by id
+/// dedups them. Hidden summaries are excluded.
+pub(crate) fn merge_roster(
+    mut entries: Vec<RosterEntry>,
+    summaries: Vec<Summary>,
+) -> Vec<RosterEntry> {
+    let mut by_id: std::collections::HashMap<String, Summary> = summaries
+        .into_iter()
+        .filter(|s| !s.is_hidden())
+        .map(|s| (s.info.id.0.to_string(), s))
+        .collect();
+
+    // Backfill resident rows; remove the summary so it isn't re-emitted below.
+    for entry in &mut entries {
+        let Some(summary) = by_id.remove(&entry.session_id) else {
+            continue;
+        };
+        if let Some(title) = summary.display_title_opt() {
+            entry.title = Some(title);
+        }
+        if entry.activity != RosterActivity::Working {
+            entry.last_change_unix_ms = summary.last_change_unix_ms();
+        }
+    }
+
+    // Remaining summaries have no resident row: emit them as dormant.
+    entries.extend(by_id.into_values().map(|summary| RosterEntry {
+        session_id: summary.info.id.0.to_string(),
+        title: summary.display_title_opt(),
+        cwd: summary.info.cwd.clone(),
+        is_worktree: summary.session_kind.as_deref() == Some("worktree")
+            || summary.source_workspace_dir.is_some(),
+        model_id: Some(summary.current_model_id.0.to_string()),
+        reasoning_effort: summary.reasoning_effort,
+        yolo: false,
+        activity: RosterActivity::Dormant,
+        resident: false,
+        last_change_unix_ms: summary.last_change_unix_ms(),
+        origin: RosterOrigin::Local,
+    }));
+
+    // Most-recently-changed first.
+    entries.sort_by(|a, b| b.last_change_unix_ms.cmp(&a.last_change_unix_ms));
+    entries
+}
+
+#[cfg(test)]
+mod merge_roster_tests {
+    use super::*;
+    use crate::session::info::Info;
+    use crate::session::persistence::default_model_id;
+    use agent_client_protocol as acp;
+
+    fn summary(id: &str, title: Option<&str>, last_active_ms: i64) -> Summary {
+        let mut s = Summary::new(
+            &Info {
+                id: acp::SessionId::new(id),
+                cwd: format!("/repo/{id}"),
+            },
+            default_model_id(),
+        )
+        .expect("summary");
+        s.generated_title = title.map(String::from);
+        s.last_active_at = chrono::DateTime::from_timestamp_millis(last_active_ms);
+        s
+    }
+
+    fn resident(id: &str, activity: RosterActivity, last_change_unix_ms: i64) -> RosterEntry {
+        RosterEntry {
+            session_id: id.to_string(),
+            title: None,
+            cwd: format!("/live/{id}"),
+            is_worktree: false,
+            model_id: Some("grok-4".into()),
+            reasoning_effort: None,
+            yolo: false,
+            activity,
+            resident: true,
+            last_change_unix_ms,
+            origin: RosterOrigin::Local,
+        }
+    }
+
+    #[test]
+    fn idle_resident_adopts_persisted_title_and_last_active() {
+        let now = 9_000;
+        let out = merge_roster(
+            vec![resident("a", RosterActivity::Idle, now)],
+            vec![summary("a", Some("Fix the roster"), 1_234)],
+        );
+        assert_eq!(out.len(), 1, "resident must not be duplicated as dormant");
+        assert_eq!(out[0].title.as_deref(), Some("Fix the roster"));
+        assert_eq!(out[0].last_change_unix_ms, 1_234, "idle adopts last-active");
+        assert!(out[0].resident);
+        assert_eq!(out[0].cwd, "/live/a", "live cwd is preserved");
+        assert_eq!(out[0].activity, RosterActivity::Idle);
+    }
+
+    #[test]
+    fn working_resident_keeps_now_but_adopts_title() {
+        let now = 9_000;
+        let out = merge_roster(
+            vec![resident("a", RosterActivity::Working, now)],
+            vec![summary("a", Some("Busy turn"), 1_234)],
+        );
+        assert_eq!(out[0].title.as_deref(), Some("Busy turn"));
+        assert_eq!(out[0].last_change_unix_ms, now, "Working stays 'now'");
+    }
+
+    #[test]
+    fn new_resident_without_summary_stays_titleless_now() {
+        let now = 9_000;
+        let out = merge_roster(vec![resident("a", RosterActivity::Idle, now)], vec![]);
+        assert_eq!(out[0].title, None);
+        assert_eq!(out[0].last_change_unix_ms, now);
+    }
+
+    #[test]
+    fn blank_persisted_title_leaves_row_untitled() {
+        let out = merge_roster(
+            vec![resident("a", RosterActivity::Idle, 9_000)],
+            vec![summary("a", Some("   "), 1_234)],
+        );
+        assert_eq!(out[0].title, None, "blank title normalizes to None");
+        assert_eq!(out[0].last_change_unix_ms, 1_234);
+    }
+
+    #[test]
+    fn dormant_sessions_are_emitted_and_sorted_after_residents() {
+        let out = merge_roster(
+            vec![resident("live", RosterActivity::Idle, 5_000)],
+            vec![
+                summary("live", Some("Live one"), 4_000),
+                summary("old", Some("Dormant one"), 1_000),
+                summary("new", Some("Newer dormant"), 8_000),
+            ],
+        );
+        let ids: Vec<&str> = out.iter().map(|e| e.session_id.as_str()).collect();
+        assert_eq!(
+            ids,
+            vec!["new", "live", "old"],
+            "sorted by last-change desc"
+        );
+        let dormant = out.iter().find(|e| e.session_id == "new").unwrap();
+        assert_eq!(dormant.activity, RosterActivity::Dormant);
+        assert!(!dormant.resident);
+    }
+
+    #[test]
+    fn duplicate_summaries_are_deduped() {
+        let out = merge_roster(
+            vec![],
+            vec![
+                summary("dup", Some("First"), 1_000),
+                summary("dup", Some("Second"), 2_000),
+            ],
+        );
+        assert_eq!(out.len(), 1, "duplicate ids collapse to one row");
+    }
+
+    #[test]
+    fn hidden_summaries_are_excluded() {
+        let mut hidden = summary("sub", Some("Subagent"), 5_000);
+        hidden.session_kind = Some("subagent".into());
+        let out = merge_roster(vec![], vec![hidden]);
+        assert!(out.is_empty(), "hidden/subagent summaries are dropped");
+    }
+
+    #[test]
+    fn dormant_row_carries_persisted_reasoning_effort() {
+        let mut s = summary("dorm", Some("Dormant"), 1_000);
+        s.reasoning_effort = Some(ReasoningEffort::Xhigh);
+        let out = merge_roster(vec![], vec![s]);
+        assert_eq!(out[0].reasoning_effort, Some(ReasoningEffort::Xhigh));
+    }
+
+    #[test]
+    fn resident_effort_is_taken_from_the_live_row_not_the_summary() {
+        // The live handle is authoritative for a resident session, so the
+        // resident row's effort must survive the summary backfill.
+        let mut live = resident("a", RosterActivity::Idle, 9_000);
+        live.reasoning_effort = Some(ReasoningEffort::High);
+        let mut s = summary("a", Some("Title"), 1_234);
+        s.reasoning_effort = Some(ReasoningEffort::Low);
+        let out = merge_roster(vec![live], vec![s]);
+        assert_eq!(out.len(), 1);
+        assert_eq!(out[0].reasoning_effort, Some(ReasoningEffort::High));
+    }
+
+    #[test]
+    fn reasoning_effort_serializes_as_camel_case_and_skips_when_none() {
+        let with_effort = RosterEntry {
+            reasoning_effort: Some(ReasoningEffort::Xhigh),
+            ..resident("a", RosterActivity::Idle, 1)
+        };
+        let json = serde_json::to_string(&with_effort).unwrap();
+        assert!(
+            json.contains("\"reasoningEffort\":\"xhigh\""),
+            "effort must be camelCase and snake_case-valued: {json}"
+        );
+
+        let without = resident("b", RosterActivity::Idle, 1);
+        let json = serde_json::to_string(&without).unwrap();
+        assert!(
+            !json.contains("reasoningEffort"),
+            "a None effort must not be serialized: {json}"
+        );
+    }
+}

crates/codegen/xai-grok-shell/src/agent/server.rs

@@ -0,0 +1,487 @@
+//! WebSocket server for remote agent connections.
+//!
+//! This module provides a WebSocket server that allows remote TUI clients to
+//! connect to a grok agent running on a different machine.
+//!
+//! The agent persists across WebSocket reconnections: a single MvpAgent instance
+//! is created on first connection and reused for all subsequent connections. This
+//! ensures that session actors (and any in-flight prompts) survive client
+//! disconnects — when a client reconnects and loads an existing session, ongoing
+//! work continues to stream to the new connection.
+
+use std::cell::RefCell;
+use std::net::SocketAddr;
+use std::rc::Rc;
+use std::sync::Arc;
+use std::thread;
+
+use axum::{
+    Router,
+    extract::{
+        ConnectInfo, Query, State,
+        ws::{Message, WebSocket, WebSocketUpgrade},
+    },
+    http::{HeaderMap, StatusCode},
+    response::{IntoResponse, Response},
+    routing::get,
+};
+use futures::{SinkExt, StreamExt};
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, simplex};
+use tokio::net::TcpListener;
+use tokio::sync::mpsc;
+use tokio::time::Duration;
+use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
+use tracing::{info, warn};
+
+use agent_client_protocol as acp;
+use xai_acp_lib::{
+    AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
+    AcpClientMessage, LineBufferedRead,
+};
+
+use crate::agent::config::{Config as AgentConfig, ModelEntry};
+use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking};
+use crate::agent::mvp_agent::MvpAgent;
+
+use indexmap::IndexMap;
+
+/// Swappable destination for the relay task.
+///
+/// Points at the current ACP connection's gateway sender. When no client is
+/// connected, the value is `None` and outbound messages are silently dropped
+/// (matching the old behaviour where the gateway channel's receiver was simply
+/// gone).
+type RelayDest = Rc<RefCell<Option<mpsc::UnboundedSender<AcpClientMessage>>>>;
+
+const MAX_BUFFER_SIZE: usize = 8 * 1024 * 1024;
+const KEEPALIVE_INTERVAL_SECS: u64 = 15;
+
+/// Configuration for the agent WebSocket server.
+#[derive(Debug, Clone)]
+pub struct ServerConfig {
+    /// Address to bind the server to
+    pub bind_addr: SocketAddr,
+    /// Secret token for client authentication (required)
+    pub secret: String,
+}
+
+/// Shared state for the WebSocket server.
+struct ServerState {
+    agent_config: AgentConfig,
+    secret: String,
+    /// Channel to send new WebSocket connections to the persistent agent thread.
+    /// Lazily initialised on first connection; protected by a tokio Mutex so the
+    /// axum handler (which is `Send`) can acquire it.
+    agent_conn_tx: tokio::sync::Mutex<Option<mpsc::UnboundedSender<NewConnectionChannels>>>,
+}
+
+/// Channels bridging a single WebSocket connection to the agent thread.
+struct NewConnectionChannels {
+    from_ws_rx: mpsc::UnboundedReceiver<String>,
+    to_ws_tx: mpsc::UnboundedSender<String>,
+}
+
+/// Query parameters for WebSocket connection.
+#[derive(Debug, serde::Deserialize, Default)]
+pub struct WsQueryParams {
+    #[serde(rename = "server-key")]
+    pub server_key: Option<String>,
+}
+
+/// Validate the bearer token from request headers or query parameters.
+fn validate_auth(headers: &HeaderMap, query: &WsQueryParams, expected_secret: &str) -> bool {
+    // Try Authorization header
+    if let Some(token) = headers
+        .get("authorization")
+        .and_then(|v| v.to_str().ok())
+        .and_then(|v| v.strip_prefix("Bearer "))
+    {
+        return token == expected_secret;
+    }
+
+    // Fall back to query parameter for browser connections
+    if let Some(ref key) = query.server_key {
+        return key == expected_secret;
+    }
+
+    false
+}
+
+/// WebSocket upgrade handler with authentication.
+async fn ws_handler(
+    ws: WebSocketUpgrade,
+    State(state): State<Arc<ServerState>>,
+    ConnectInfo(addr): ConnectInfo<SocketAddr>,
+    headers: HeaderMap,
+    Query(query): Query<WsQueryParams>,
+) -> Response {
+    // Validate secret token from header or query param
+    if !validate_auth(&headers, &query, &state.secret) {
+        warn!("Unauthorized connection attempt from {}", addr);
+        return (
+            StatusCode::UNAUTHORIZED,
+            "Invalid or missing authorization token",
+        )
+            .into_response();
+    }
+
+    info!("Authenticated WebSocket connection from {}", addr);
+    ws.on_upgrade(move |socket| handle_connection(socket, state, addr))
+}
+
+/// Handle an authenticated WebSocket connection.
+///
+/// On first connection, spawns a persistent agent thread that owns the MvpAgent.
+/// On subsequent connections (reconnects), sends new WS channels to the existing
+/// agent thread so that session actors can continue streaming to the new client.
+async fn handle_connection(ws: WebSocket, state: Arc<ServerState>, peer_addr: SocketAddr) {
+    info!("New WebSocket connection from {}", peer_addr);
+
+    let (mut ws_write, mut ws_read) = ws.split();
+
+    // Channels for bridging WS <-> Agent thread
+    let (to_agent_tx, to_agent_rx) = mpsc::unbounded_channel::<String>();
+    let (from_agent_tx, mut from_agent_rx) = mpsc::unbounded_channel::<String>();
+
+    // Ensure the persistent agent thread is running (lazy init on first connection).
+    // If the previous agent thread died (panic, etc.), clear the stale sender so we
+    // respawn a fresh one.
+    {
+        let mut agent_tx_guard = state.agent_conn_tx.lock().await;
+
+        // Check if existing sender is still alive (receiver not dropped)
+        if let Some(ref tx) = *agent_tx_guard
+            && tx.is_closed()
+        {
+            warn!("Persistent agent thread died — will respawn");
+            *agent_tx_guard = None;
+        }
+
+        if agent_tx_guard.is_none() {
+            let (conn_tx, conn_rx) = mpsc::unbounded_channel();
+
+            let agent_config = state.agent_config.clone();
+            let _agent_thread = thread::Builder::new()
+                .name("agent-persistent".to_string())
+                .spawn(move || {
+                    // Prefetch models before creating the runtime (blocking is OK here)
+                    let auth = agent_config.create_auth_manager().current();
+                    let fetch_auth =
+                        ModelFetchAuth::resolve(&agent_config.endpoints, auth.is_some());
+                    let prefetched_models = if auth.is_some()
+                        || agent_config.endpoints.has_custom_endpoint()
+                        || fetch_auth != ModelFetchAuth::Session
+                    {
+                        prefetch_models_blocking(&agent_config.endpoints, auth.as_ref(), fetch_auth)
+                    } else {
+                        None
+                    };
+
+                    info!("Prefetched models: {:?}", prefetched_models);
+
+                    let rt = tokio::runtime::Builder::new_current_thread()
+                        .enable_all()
+                        .build()
+                        .expect("Failed to create runtime for agent");
+
+                    let local_set = tokio::task::LocalSet::new();
+                    local_set.block_on(&rt, async move {
+                        run_persistent_agent(agent_config, conn_rx, prefetched_models).await
+                    });
+
+                    warn!("Persistent agent thread exiting");
+                });
+
+            *agent_tx_guard = Some(conn_tx);
+            info!("Persistent agent thread spawned");
+        }
+
+        // Send new WS channels to the agent thread
+        if let Some(ref tx) = *agent_tx_guard
+            && tx
+                .send(NewConnectionChannels {
+                    from_ws_rx: to_agent_rx,
+                    to_ws_tx: from_agent_tx,
+                })
+                .is_err()
+        {
+            warn!("Failed to send connection channels to agent thread");
+        }
+    }
+
+    // Task: Read from WS, send to agent thread
+    let read_task = tokio::spawn(async move {
+        while let Some(msg) = ws_read.next().await {
+            match msg {
+                Ok(Message::Text(text)) => {
+                    let text_str: &str = text.as_ref();
+                    let trimmed = text_str.trim_end_matches(['\r', '\n']);
+                    // Skip browser keepalive pings (non-JSON text)
+                    if trimmed == "ping" || trimmed.is_empty() {
+                        continue;
+                    }
+                    if to_agent_tx.send(trimmed.to_string()).is_err() {
+                        break;
+                    }
+                }
+                Ok(Message::Binary(bin)) => {
+                    if let Ok(s) = std::str::from_utf8(&bin) {
+                        let trimmed = s.trim_end_matches(['\r', '\n']);
+                        if trimmed == "ping" || trimmed.is_empty() {
+                            continue;
+                        }
+                        if to_agent_tx.send(trimmed.to_string()).is_err() {
+                            break;
+                        }
+                    }
+                }
+                Ok(Message::Close(frame)) => {
+                    if let Some(f) = frame {
+                        info!(
+                            "WebSocket close from {}: {} {}",
+                            peer_addr, f.code, f.reason
+                        );
+                    }
+                    break;
+                }
+                Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => {}
+                Err(e) => {
+                    warn!("WebSocket read error from {}: {:?}", peer_addr, e);
+                    break;
+                }
+            }
+        }
+    });
+
+    // Task: Read from agent thread, send to WS (with keepalive)
+    let write_task = tokio::spawn(async move {
+        let mut keepalive = tokio::time::interval(Duration::from_secs(KEEPALIVE_INTERVAL_SECS));
+
+        loop {
+            tokio::select! {
+                Some(msg) = from_agent_rx.recv() => {
+                    if ws_write.send(Message::Text(msg.into())).await.is_err() {
+                        break;
+                    }
+                }
+                _ = keepalive.tick() => {
+                    if ws_write.send(Message::Ping(vec![].into())).await.is_err() {
+                        break;
+                    }
+                }
+                else => break,
+            }
+        }
+    });
+
+    // Wait for either task to complete
+    tokio::select! {
+        _ = read_task => {}
+        _ = write_task => {}
+    }
+
+    info!("WebSocket connection ended for {}", peer_addr);
+}
+
+/// Run the persistent agent on a dedicated thread with LocalSet.
+///
+/// The MvpAgent is created **once** and reused across WebSocket reconnections.
+/// A persistent gateway channel ensures that session actors (which hold cloned
+/// `GatewaySender` handles) can always send notifications. A relay task forwards
+/// messages from the persistent channel to the *current* ACP connection's channel,
+/// so notifications reach whichever client is currently connected.
+async fn run_persistent_agent(
+    agent_config: AgentConfig,
+    mut connection_rx: mpsc::UnboundedReceiver<NewConnectionChannels>,
+    prefetched_models: Option<IndexMap<String, ModelEntry>>,
+) {
+    // Persistent gateway channel — the MvpAgent and all session actors hold
+    // clones of `gw_tx`. This channel survives across reconnections.
+    let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel::<AcpClientMessage>();
+    let gateway = GatewaySender::new(gw_tx);
+
+    // Create MvpAgent ONCE -- it persists for the lifetime of the server.
+    let auth_manager = Arc::new(agent_config.create_auth_manager());
+    // Proactive token refresh; runs until process exit.
+    auth_manager.start_proactive_refresh(tokio_util::sync::CancellationToken::new());
+    // Restore managed policy right before bootstrap reads it — the agent is created lazily here,
+    // so an earlier restore could go stale before the gate.
+    crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
+    let agent = Rc::new(
+        MvpAgent::new(gateway, &agent_config, auth_manager, prefetched_models)
+            .unwrap_or_else(crate::agent::init::exit_on_config_error),
+    );
+
+    let relay_dest: RelayDest = Rc::new(RefCell::new(None));
+
+    // Relay task: reads from the persistent gateway channel and forwards to
+    // whichever ACP connection is currently active.
+    let relay_dest_for_task = relay_dest.clone();
+    tokio::task::spawn_local(async move {
+        while let Some(msg) = gw_rx.recv().await {
+            let maybe_tx = relay_dest_for_task.borrow().clone();
+            if let Some(tx) = maybe_tx
+                && tx.send(msg).is_err()
+            {
+                // Connection's gateway receiver was dropped — clear it.
+                *relay_dest_for_task.borrow_mut() = None;
+            }
+            // If no connection, the message (and its response_tx) is dropped.
+            // The caller (session actor) gets a send error which is already
+            // handled with `let _ = ...`.
+        }
+    });
+
+    // Accept new connections in a loop
+    while let Some(channels) = connection_rx.recv().await {
+        info!("Agent thread: setting up new ACP connection (reconnect)");
+        setup_acp_connection(agent.clone(), channels, relay_dest.clone());
+    }
+
+    info!("Agent thread: connection channel closed, exiting");
+}
+
+/// Set up a new ACP connection for a WebSocket connection, reusing the existing
+/// MvpAgent. The relay destination is updated so that session actor notifications
+/// flow to the new client.
+fn setup_acp_connection(
+    agent: Rc<MvpAgent>,
+    channels: NewConnectionChannels,
+    relay_dest: RelayDest,
+) {
+    let NewConnectionChannels {
+        mut from_ws_rx,
+        to_ws_tx,
+    } = channels;
+
+    // Create new simplex IO streams for this ACP connection
+    let (agent_read_rx, mut agent_read_tx) = simplex(MAX_BUFFER_SIZE);
+    let (agent_write_rx, agent_write_tx) = simplex(MAX_BUFFER_SIZE);
+
+    let incoming = agent_read_rx.compat();
+    let outgoing = agent_write_tx.compat_write();
+
+    // Create a per-connection gateway channel for the GatewayReceiver.
+    // The relay task will forward persistent-channel messages here.
+    let (conn_gw_tx, conn_gw_rx) = tokio::sync::mpsc::unbounded_channel::<AcpClientMessage>();
+
+    // Point the relay at this new connection's channel
+    *relay_dest.borrow_mut() = Some(conn_gw_tx);
+
+    // Create new ACP connection reusing the same MvpAgent (via Rc clone).
+    // `Agent` is implemented for `Rc<T: Agent>` so this works.
+    let incoming = LineBufferedRead::spawn_local(incoming);
+    let (conn, handle_io) = acp::AgentSideConnection::new(agent, outgoing, incoming, |fut| {
+        tokio::task::spawn_local(fut);
+    });
+    tokio::task::spawn_local(
+        GatewayReceiver::new(conn_gw_rx, conn)
+            .with_on_meta(xai_file_utils::trace_context::span_from_meta_traceparent)
+            .run(),
+    );
+
+    // Task: Forward WS messages → agent (incoming ACP bytes)
+    tokio::task::spawn_local(async move {
+        while let Some(msg) = from_ws_rx.recv().await {
+            // Log messages that lack both `id` and `method` — the ACP layer
+            // only prints "received message with neither id nor method" without
+            // the payload, making debugging impossible.
+            if let Ok(v) = serde_json::from_str::<serde_json::Value>(&msg)
+                && v.get("id").is_none()
+                && v.get("method").is_none()
+            {
+                warn!(
+                    len = msg.len(),
+                    "incoming WS message has neither id nor method"
+                );
+            }
+            if agent_read_tx.write_all(msg.as_bytes()).await.is_err() {
+                break;
+            }
+            if agent_read_tx.write_all(b"\n").await.is_err() {
+                break;
+            }
+        }
+        // WS disconnected — the simplex writer is dropped, causing `handle_io`
+        // to complete. The GatewayReceiver for this connection will also stop.
+        // But the MvpAgent and session actors stay alive, ready for the next
+        // connection.
+    });
+
+    // Task: Forward agent messages → WS (outgoing ACP bytes)
+    tokio::task::spawn_local(async move {
+        let mut reader = BufReader::new(agent_write_rx);
+        let mut line = String::new();
+
+        loop {
+            line.clear();
+            match reader.read_line(&mut line).await {
+                Ok(0) => break,
+                Ok(_) => {
+                    let msg = line.trim_end_matches(['\r', '\n']);
+                    if !msg.is_empty() && to_ws_tx.send(msg.to_string()).is_err() {
+                        break;
+                    }
+                }
+                Err(_) => break,
+            }
+        }
+    });
+
+    // Run the ACP IO handler — fire-and-forget since we don't block the
+    // connection loop. It completes when the WS disconnects.
+    tokio::task::spawn_local(async move {
+        let _ = handle_io.await;
+        info!("ACP connection IO handler completed");
+    });
+}
+
+/// Run the agent WebSocket server.
+///
+/// This starts a WebSocket server that accepts authenticated connections from
+/// remote TUI clients. A single agent instance is shared across all connections
+/// (persisted across reconnections) so that in-flight session work survives
+/// client disconnects.
+///
+/// # Arguments
+/// * `config` - Server configuration (bind address and secret)
+/// * `agent_config` - Agent configuration to use for each connection
+///
+/// # Example
+/// ```ignore
+/// let server_config = ServerConfig {
+///     bind_addr: "0.0.0.0:9000".parse().unwrap(),
+///     secret: "my-secret-token".to_string(),
+/// };
+/// run_agent_server(server_config, agent_config).await?;
+/// ```
+pub async fn run_agent_server(
+    config: ServerConfig,
+    agent_config: AgentConfig,
+) -> anyhow::Result<()> {
+    let state = Arc::new(ServerState {
+        agent_config,
+        secret: config.secret,
+        agent_conn_tx: tokio::sync::Mutex::new(None),
+    });
+
+    let app = Router::new()
+        .route("/ws", get(ws_handler))
+        .with_state(state);
+
+    let listener = TcpListener::bind(config.bind_addr).await?;
+    info!("Agent server listening on ws://{}/ws", config.bind_addr);
+    info!(
+        "Clients should connect with: --remote ws://{}:{}/ws --secret <token>",
+        config.bind_addr.ip(),
+        config.bind_addr.port()
+    );
+
+    axum::serve(
+        listener,
+        app.into_make_service_with_connect_info::<SocketAddr>(),
+    )
+    .await?;
+
+    Ok(())
+}

crates/codegen/xai-grok-shell/src/agent/session_config.rs

@@ -0,0 +1,218 @@
+use agent_client_protocol as acp;
+use serde::Serialize;
+use xai_grok_sampling_types::{ReasoningEffort, ReasoningEffortOption};
+
+use crate::session::unified_list::SessionKind;
+
+pub(crate) const SELECTABLE_REASONING_EFFORTS: [ReasoningEffort; 5] = [
+    ReasoningEffort::Minimal,
+    ReasoningEffort::Low,
+    ReasoningEffort::Medium,
+    ReasoningEffort::High,
+    ReasoningEffort::Xhigh,
+];
+
+#[derive(Debug, Clone, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SessionConfigOption {
+    pub id: String,
+    pub category: String,
+    pub label: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub description: Option<String>,
+    pub selected: bool,
+}
+
+#[derive(Debug, Clone, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct GrokSessionDetail {
+    pub session_id: String,
+    pub kind: String,
+    pub cwd: String,
+    pub current_model_id: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub title: Option<String>,
+}
+
+impl GrokSessionDetail {
+    pub fn build(
+        session_id: String,
+        cwd: String,
+        current_model_id: String,
+        title: Option<String>,
+    ) -> Self {
+        Self {
+            session_id,
+            kind: SessionKind::Build.as_str().to_string(),
+            cwd,
+            current_model_id,
+            title,
+        }
+    }
+}
+
+fn effort_label(effort: ReasoningEffort) -> String {
+    match effort {
+        ReasoningEffort::None => "None",
+        ReasoningEffort::Minimal => "Minimal",
+        ReasoningEffort::Low => "Low",
+        ReasoningEffort::Medium => "Medium",
+        ReasoningEffort::High => "High",
+        ReasoningEffort::Xhigh => "X-High",
+    }
+    .to_string()
+}
+
+/// The built-in session-picker modes used when the model has no server list.
+/// Reproduces the historical five rows and their labels.
+pub(crate) fn legacy_session_effort_options() -> Vec<ReasoningEffortOption> {
+    SELECTABLE_REASONING_EFFORTS
+        .iter()
+        .map(|&effort| ReasoningEffortOption {
+            id: effort.as_str().to_string(),
+            value: effort,
+            label: effort_label(effort),
+            description: None,
+            default: false,
+        })
+        .collect()
+}
+
+pub(crate) fn build_session_config_options(
+    available_models: &[acp::ModelInfo],
+    current_model_id: &acp::ModelId,
+    effort_options: &[ReasoningEffortOption],
+    current_effort: Option<ReasoningEffort>,
+) -> Vec<SessionConfigOption> {
+    let mut options = Vec::with_capacity(available_models.len() + effort_options.len());
+
+    for model in available_models {
+        let label = if model.name.is_empty() {
+            model.model_id.0.to_string()
+        } else {
+            model.name.clone()
+        };
+        options.push(SessionConfigOption {
+            id: model.model_id.0.to_string(),
+            category: "model".to_string(),
+            label,
+            description: None,
+            selected: model.model_id == *current_model_id,
+        });
+    }
+
+    for effort in effort_options {
+        options.push(SessionConfigOption {
+            id: effort.id.clone(),
+            category: "mode".to_string(),
+            label: effort.label.clone(),
+            description: effort.description.clone(),
+            selected: Some(effort.value) == current_effort,
+        });
+    }
+
+    options
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn model(id: &'static str, name: &str) -> acp::ModelInfo {
+        acp::ModelInfo::new(acp::ModelId::new(id), name.to_string())
+    }
+
+    #[test]
+    fn options_have_one_selected_model_and_a_mode_per_effort() {
+        let models = [
+            model("grok-build", "Grok Build"),
+            model("grok-4.5", "Grok 4.5"),
+        ];
+        let current = acp::ModelId::from("grok-build");
+        let opts = build_session_config_options(
+            &models,
+            &current,
+            &legacy_session_effort_options(),
+            Some(ReasoningEffort::High),
+        );
+
+        let model_opts: Vec<_> = opts.iter().filter(|o| o.category == "model").collect();
+        assert_eq!(model_opts.len(), 2);
+        let selected_models: Vec<_> = model_opts.iter().filter(|o| o.selected).collect();
+        assert_eq!(selected_models.len(), 1);
+        assert_eq!(selected_models[0].id, "grok-build");
+
+        let mode_opts: Vec<_> = opts.iter().filter(|o| o.category == "mode").collect();
+        assert_eq!(mode_opts.len(), SELECTABLE_REASONING_EFFORTS.len());
+        let selected_modes: Vec<_> = mode_opts.iter().filter(|o| o.selected).collect();
+        assert_eq!(selected_modes.len(), 1);
+        assert_eq!(selected_modes[0].id, "high");
+        assert_eq!(selected_modes[0].label, "High");
+    }
+
+    #[test]
+    fn none_effort_is_not_a_user_selectable_mode() {
+        assert!(!SELECTABLE_REASONING_EFFORTS.contains(&ReasoningEffort::None));
+        let models = [model("grok-build", "Grok Build")];
+        let current = acp::ModelId::from("grok-build");
+        let opts = build_session_config_options(
+            &models,
+            &current,
+            &legacy_session_effort_options(),
+            Some(ReasoningEffort::None),
+        );
+        let modes: Vec<_> = opts.iter().filter(|o| o.category == "mode").collect();
+        assert!(modes.iter().all(|o| o.id != "none"));
+        assert!(modes.iter().all(|o| !o.selected));
+    }
+
+    #[test]
+    fn no_mode_options_when_model_lacks_effort_support() {
+        let models = [model("grok-build", "Grok Build")];
+        let current = acp::ModelId::from("grok-build");
+        let opts = build_session_config_options(&models, &current, &[], None);
+        assert_eq!(opts.len(), 1);
+        assert!(opts.iter().all(|o| o.category == "model"));
+    }
+
+    #[test]
+    fn model_label_falls_back_to_id_when_name_empty() {
+        let models = [model("grok-build", "")];
+        let current = acp::ModelId::from("grok-build");
+        let opts = build_session_config_options(&models, &current, &[], None);
+        assert_eq!(opts[0].label, "grok-build");
+    }
+
+    #[test]
+    fn session_config_option_serializes_camel_case() {
+        let opt = SessionConfigOption {
+            id: "grok-build".to_string(),
+            category: "model".to_string(),
+            label: "Grok Build".to_string(),
+            description: None,
+            selected: true,
+        };
+        let v = serde_json::to_value(&opt).expect("serialize");
+        assert_eq!(v["id"], "grok-build");
+        assert_eq!(v["category"], "model");
+        assert_eq!(v["label"], "Grok Build");
+        assert_eq!(v["selected"], true);
+        assert!(v.get("description").is_none());
+    }
+
+    #[test]
+    fn grok_session_detail_serializes_camel_case() {
+        let detail = GrokSessionDetail::build(
+            "sess-1".to_string(),
+            "/Users/me/xai".to_string(),
+            "grok-build".to_string(),
+            None,
+        );
+        let v = serde_json::to_value(&detail).expect("serialize");
+        assert_eq!(v["sessionId"], "sess-1");
+        assert_eq!(v["kind"], "build");
+        assert_eq!(v["cwd"], "/Users/me/xai");
+        assert_eq!(v["currentModelId"], "grok-build");
+        assert!(v.get("title").is_none());
+    }
+}

crates/codegen/xai-grok-shell/src/agent/session_metrics.rs

@@ -0,0 +1,10 @@
+//! Session lifecycle event structs.
+//!
+//! Re-exported from `xai-grok-telemetry` after the telemetry crate split.
+//! The structs themselves live in the telemetry crate; this module preserves
+//! the existing import path so nothing else in shell needs to change.
+
+pub(crate) use xai_grok_telemetry::session_metrics::{
+    DoomLoopRecovery, SessionStarted, TraceUploadAttempted, TraceUploadFailed, TraceUploadSkipped,
+    TraceUploadSucceeded, Turn, TurnCompletedLifecycle,
+};

crates/codegen/xai-grok-shell/src/agent/session_registry_client.rs

@@ -0,0 +1,642 @@
+//! REST client for the session replicas registry (cli-chat-proxy).
+//!
+//! Handles registering, updating, finalizing, searching, and downloading
+//! session replicas for cross-host session replication. Write methods
+//! (register/update/finalize) are fire-and-forget safe. Read methods
+//! (search/get/download_file) return typed results.
+
+use anyhow::{Context, Result};
+use reqwest::RequestBuilder;
+use serde::{Deserialize, Serialize};
+
+// ============================================================================
+// Request / response types (local — not in cli-chat-proxy since these
+// are only used by the agent, not consumed by other crates)
+// ============================================================================
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct RegisterRequest {
+    pub session_id: String,
+    pub cwd: String,
+    pub gcs_trace_prefix: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model_id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub repo_remote_url: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub repo_branch: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub repo_head_at_start: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub hostname: Option<String>,
+    /// Opaque per-machine device id (telemetry `agent_id()`) for machine disambiguation.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub device_id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub parent_session_id: Option<String>,
+    // --- Subagent-specific fields (optional, backward-compatible) ---
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub session_kind: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub subagent_type: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub subagent_persona: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub subagent_role: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub fork_context_source: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub subagent_depth: Option<u32>,
+}
+
+#[derive(Debug, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct UpdateRequest {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub summary: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub first_prompt: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub last_turn_number: Option<i32>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub repo_head_at_end: Option<String>,
+    /// Latest turn whose restore artifacts are confirmed durable.
+    /// Omitted from the wire when `None` — old servers ignore unknown fields.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub restorable_turn_number: Option<i32>,
+}
+
+// ============================================================================
+// Response types
+// ============================================================================
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SessionRecord {
+    pub session_id: String,
+    pub summary: String,
+    pub first_prompt: Option<String>,
+    pub model_id: Option<String>,
+    pub created_at: String,
+    pub updated_at: String,
+    pub last_turn_number: i32,
+    /// Present on servers that have applied the restorable-turn migration.
+    /// `None` when talking to an older server — callers should fall back to
+    /// `last_turn_number` in that case.
+    #[serde(default)]
+    pub restorable_turn_number: Option<i32>,
+    pub cwd: String,
+    pub repo_remote_url: Option<String>,
+    pub hostname: Option<String>,
+    pub status: String,
+    pub gcs_trace_prefix: String,
+    pub gcs_bucket: String,
+    #[serde(default)]
+    pub last_active_at: Option<String>,
+}
+
+impl From<crate::session::persistence::Summary> for SessionRecord {
+    fn from(s: crate::session::persistence::Summary) -> Self {
+        Self {
+            session_id: s.info.id.to_string(),
+            summary: s.session_summary,
+            first_prompt: None,
+            model_id: Some(s.current_model_id.to_string()),
+            created_at: s.created_at.to_rfc3339(),
+            updated_at: s.updated_at.to_rfc3339(),
+            last_turn_number: s.num_messages as i32,
+            restorable_turn_number: None,
+            cwd: s.info.cwd,
+            repo_remote_url: None,
+            hostname: None,
+            status: "local".to_string(),
+            gcs_trace_prefix: String::new(),
+            gcs_bucket: String::new(),
+            last_active_at: s.last_active_at.map(|t| t.to_rfc3339()),
+        }
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SearchResponse {
+    pub sessions: Vec<SessionRecord>,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct DownloadResponse {
+    pub download_url: String,
+    pub file: String,
+    pub turn: i32,
+}
+
+// ============================================================================
+// Client
+// ============================================================================
+
+#[derive(Clone)]
+pub struct SessionRegistryClient {
+    raw_client: reqwest::Client,
+    client: reqwest_middleware::ClientWithMiddleware,
+    base_url: String,
+    credentials: crate::util::grok_auth_credentials::GrokAuthCredentials,
+    session_id: Option<String>,
+}
+
+impl SessionRegistryClient {
+    pub fn new(base_url: impl Into<String>, user_token: impl Into<String>) -> Self {
+        let http_client = crate::http::shared_client();
+        Self {
+            raw_client: http_client.clone(),
+            client: reqwest_middleware::ClientBuilder::new(http_client).build(),
+            base_url: base_url.into(),
+            credentials: crate::util::grok_auth_credentials::GrokAuthCredentials::new(Some(
+                user_token.into(),
+            )),
+            session_id: None,
+        }
+    }
+
+    pub fn with_deployment_key(mut self, key: Option<String>) -> Self {
+        self.credentials.deployment_key = key;
+        self
+    }
+
+    pub fn with_alpha_test_key(mut self, key: Option<String>) -> Self {
+        self.credentials.alpha_test_key = key;
+        self
+    }
+
+    pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
+        self.session_id = Some(session_id.into());
+        self
+    }
+
+    /// Attach an `AuthManager` so the request signing and 401
+    /// recovery go through the consolidated auth path.
+    pub fn with_auth(mut self, auth_manager: std::sync::Arc<crate::auth::AuthManager>) -> Self {
+        let provider: std::sync::Arc<dyn xai_grok_auth::AuthCredentialProvider> =
+            std::sync::Arc::new(
+                crate::auth::credential_provider::ShellAuthCredentialProvider::new(
+                    auth_manager.clone(),
+                    self.credentials.deployment_key.clone(),
+                    self.credentials.alpha_test_key.clone(),
+                ),
+            );
+        self.credentials = self.credentials.with_auth_manager(auth_manager);
+        self.client = crate::http::with_auth_retry(self.raw_client.clone(), provider);
+        self
+    }
+
+    async fn send_authed(
+        &self,
+        builder: RequestBuilder,
+        op: &'static str,
+    ) -> Result<reqwest::Response> {
+        let builder = xai_file_utils::trace_context::inject_trace_context_into_request(builder);
+        let request = builder.build().context(op)?;
+        self.client.execute(request).await.map_err(|e| match e {
+            reqwest_middleware::Error::Middleware(e) => e.context(op),
+            reqwest_middleware::Error::Reqwest(e) => anyhow::Error::from(e).context(op),
+        })
+    }
+
+    /// Non-auth headers only -- the `Authorization` header lives in
+    /// `send_authed` so it picks up freshly-refreshed tokens.
+    fn add_common_headers(&self, builder: RequestBuilder) -> RequestBuilder {
+        builder
+    }
+
+    fn check_response(&self, response: reqwest::Response, op: &str) -> anyhow::Error {
+        if response.status() == reqwest::StatusCode::UNAUTHORIZED {
+            self.record_401_attribution(op);
+            anyhow::anyhow!("{op}: {}", self.credentials.auth_error_hint())
+        } else {
+            anyhow::anyhow!("{op} failed: {}", response.status())
+        }
+    }
+
+    /// Emit a single `auth 401 attribution` log entry tagged with
+    /// `consumer = "SessionRegistryClient.<op>"`. The op string is the
+    /// operation name passed to `check_response` (e.g.,
+    /// `"session register"`).
+    fn record_401_attribution(&self, op: &str) {
+        if let Some(manager) = self.credentials.auth_manager() {
+            let resolved = self.credentials.resolve();
+            let sent = resolved
+                .deployment_key
+                .clone()
+                .or(resolved.user_token.clone());
+            crate::auth::attribution::record_consumer_401(
+                manager.as_ref(),
+                self.session_id.as_deref(),
+                crate::auth::attribution::ConsumerKind::SessionRegistryClient,
+                op,
+                sent.as_deref(),
+            );
+        }
+    }
+
+    fn post(&self, url: &str) -> RequestBuilder {
+        self.add_common_headers(self.raw_client.post(url))
+    }
+
+    fn get(&self, url: &str) -> RequestBuilder {
+        self.add_common_headers(self.raw_client.get(url))
+    }
+
+    /// POST /v1/sessions/register (idempotent via ON CONFLICT)
+    pub async fn register(&self, req: &RegisterRequest) -> Result<()> {
+        let url = format!("{}/sessions/register", self.base_url);
+        let response = self
+            .send_authed(self.post(&url).json(req), "session register")
+            .await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session register"));
+        }
+        Ok(())
+    }
+
+    /// POST /v1/sessions/{id}/replicas/update
+    pub async fn update(&self, session_id: &str, req: &UpdateRequest) -> Result<()> {
+        let url = format!("{}/sessions/{}/replicas/update", self.base_url, session_id);
+        let response = self
+            .send_authed(self.post(&url).json(req), "session update")
+            .await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session update"));
+        }
+        Ok(())
+    }
+
+    /// POST /v1/sessions/{id}/replicas/finalize
+    pub async fn finalize(&self, session_id: &str) -> Result<()> {
+        let url = format!(
+            "{}/sessions/{}/replicas/finalize",
+            self.base_url, session_id
+        );
+        let response = self
+            .send_authed(self.post(&url), "session finalize")
+            .await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session finalize"));
+        }
+        Ok(())
+    }
+
+    /// GET /v1/sessions/search
+    pub async fn search(&self, query: Option<&str>, limit: i64) -> Result<Vec<SessionRecord>> {
+        let url = format!("{}/sessions/search", self.base_url);
+        let mut builder = self.get(&url).query(&[("limit", limit.to_string())]);
+        if let Some(q) = query {
+            builder = builder.query(&[("query", q)]);
+        }
+        let response = self.send_authed(builder, "session search").await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session search"));
+        }
+        let resp: SearchResponse = response.json().await.context("parse search response")?;
+        Ok(resp.sessions)
+    }
+
+    /// GET /v1/sessions/{id}/replicas
+    pub async fn get_session(&self, session_id: &str) -> Result<SessionRecord> {
+        let url = format!("{}/sessions/{}/replicas", self.base_url, session_id);
+        let response = self.send_authed(self.get(&url), "session get").await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session get"));
+        }
+        response.json().await.context("parse session response")
+    }
+
+    /// GET /v1/sessions/{id}/download — returns a signed GCS URL without downloading.
+    pub async fn get_download_url(
+        &self,
+        session_id: &str,
+        file: &str,
+        turn: i32,
+    ) -> Result<String> {
+        let url = format!("{}/sessions/{}/download", self.base_url, session_id);
+        let builder = self
+            .get(&url)
+            .query(&[("file", file), ("turn", &turn.to_string())]);
+        let response = self.send_authed(builder, "session download url").await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session download url"));
+        }
+        let resp: DownloadResponse = response.json().await.context("parse download response")?;
+        Ok(resp.download_url)
+    }
+
+    /// GET /v1/sessions/{id}/download — returns a signed URL, then streams to dest file.
+    pub async fn download_file(
+        &self,
+        session_id: &str,
+        file: &str,
+        turn: i32,
+        dest: &std::path::Path,
+    ) -> Result<()> {
+        let url = format!("{}/sessions/{}/download", self.base_url, session_id);
+        let builder = self
+            .get(&url)
+            .query(&[("file", file), ("turn", &turn.to_string())]);
+        let response = self.send_authed(builder, "session download").await?;
+        if !response.status().is_success() {
+            return Err(self.check_response(response, "session download"));
+        }
+        let resp: DownloadResponse = response.json().await.context("parse download response")?;
+
+        // Stream from the signed GCS URL directly to disk (archives can be hundreds of MB)
+        let mut gcs_response = self
+            .raw_client
+            .get(&resp.download_url)
+            .send()
+            .await
+            .context("download from GCS")?;
+        if !gcs_response.status().is_success() {
+            anyhow::bail!("GCS download failed: {}", gcs_response.status());
+        }
+        if let Some(parent) = dest.parent() {
+            tokio::fs::create_dir_all(parent).await?;
+        }
+        let mut out = tokio::fs::File::create(dest)
+            .await
+            .context("create dest file")?;
+        let chunk_timeout = std::time::Duration::from_secs(60);
+        loop {
+            match tokio::time::timeout(chunk_timeout, gcs_response.chunk()).await {
+                Ok(Ok(Some(chunk))) => {
+                    tokio::io::AsyncWriteExt::write_all(&mut out, &chunk)
+                        .await
+                        .context("write chunk to disk")?;
+                }
+                Ok(Ok(None)) => break,
+                Ok(Err(e)) => return Err(e).context("read GCS chunk"),
+                Err(_) => anyhow::bail!(
+                    "GCS download stalled: no data received for {chunk_timeout:?} \
+                     while downloading {file}"
+                ),
+            }
+        }
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // ── UpdateRequest wire shapes ────────────────────────────────────────────
+    //
+    // The writer split relies on two distinct update payloads being sent at
+    // different times:
+    //
+    //   1. Immediate post-turn: `last_turn_number` + `repo_head_at_end`
+    //   2. Artifact-ready:      `restorable_turn_number` only
+    //
+    // These tests verify that `skip_serializing_if = "Option::is_none"` does the
+    // right thing for each shape, so old servers silently ignore the new field and
+    // clients don't accidentally overwrite unrelated fields with nulls.
+
+    #[test]
+    fn immediate_turn_update_omits_restorable_field() {
+        let req = UpdateRequest {
+            summary: None,
+            first_prompt: None,
+            last_turn_number: Some(5),
+            repo_head_at_end: Some("abc123".into()),
+            restorable_turn_number: None,
+        };
+        let json = serde_json::to_value(&req).unwrap();
+        assert_eq!(json["lastTurnNumber"], 5);
+        assert_eq!(json["repoHeadAtEnd"], "abc123");
+        assert!(json.get("restorableTurnNumber").is_none());
+        assert!(json.get("summary").is_none());
+        assert!(json.get("firstPrompt").is_none());
+    }
+
+    #[test]
+    fn restorable_turn_update_omits_last_turn_and_head_fields() {
+        let req = UpdateRequest {
+            summary: None,
+            first_prompt: None,
+            last_turn_number: None,
+            repo_head_at_end: None,
+            restorable_turn_number: Some(5),
+        };
+        let json = serde_json::to_value(&req).unwrap();
+        assert_eq!(json["restorableTurnNumber"], 5);
+        assert!(json.get("lastTurnNumber").is_none());
+        assert!(json.get("repoHeadAtEnd").is_none());
+        assert!(json.get("summary").is_none());
+        assert!(json.get("firstPrompt").is_none());
+    }
+
+    #[test]
+    fn summary_update_omits_all_turn_fields() {
+        let req = UpdateRequest {
+            summary: Some("My session summary".into()),
+            first_prompt: None,
+            last_turn_number: None,
+            repo_head_at_end: None,
+            restorable_turn_number: None,
+        };
+        let json = serde_json::to_value(&req).unwrap();
+        assert_eq!(json["summary"], "My session summary");
+        assert!(json.get("lastTurnNumber").is_none());
+        assert!(json.get("restorableTurnNumber").is_none());
+        assert!(json.get("repoHeadAtEnd").is_none());
+    }
+
+    // Wire-contract tests: server reads the camelCase `deviceId` key.
+
+    fn minimal_register_request(device_id: Option<String>) -> RegisterRequest {
+        RegisterRequest {
+            session_id: "s1".into(),
+            cwd: "/x".into(),
+            gcs_trace_prefix: "t".into(),
+            model_id: None,
+            repo_remote_url: None,
+            repo_branch: None,
+            repo_head_at_start: None,
+            hostname: None,
+            device_id,
+            parent_session_id: None,
+            session_kind: None,
+            subagent_type: None,
+            subagent_persona: None,
+            subagent_role: None,
+            fork_context_source: None,
+            subagent_depth: None,
+        }
+    }
+
+    #[test]
+    fn register_request_serializes_device_id_as_camel_case() {
+        let req = minimal_register_request(Some("machine-uuid-123".into()));
+        let json = serde_json::to_value(&req).unwrap();
+        assert_eq!(json["deviceId"], "machine-uuid-123");
+        assert!(json.get("device_id").is_none());
+    }
+
+    #[test]
+    fn register_request_serializes_empty_device_id_as_present() {
+        let req = minimal_register_request(Some(String::new()));
+        let json = serde_json::to_value(&req).unwrap();
+        assert_eq!(json["deviceId"], "");
+    }
+
+    #[test]
+    fn register_request_omits_device_id_when_none() {
+        let req = minimal_register_request(None);
+        let json = serde_json::to_value(&req).unwrap();
+        assert!(json.get("deviceId").is_none());
+        assert!(json.get("device_id").is_none());
+    }
+
+    // ── SessionRecord backward compatibility ─────────────────────────────────
+    //
+    // Older servers do not include `restorable_turn_number` in their response.
+    // The field is `#[serde(default)]` so it must deserialize as `None` when
+    // absent, keeping new clients compatible with old servers.
+
+    #[test]
+    fn session_record_without_restorable_turn_deserializes_as_none() {
+        let json = serde_json::json!({
+            "sessionId": "sess-abc",
+            "summary": "hello",
+            "firstPrompt": null,
+            "modelId": null,
+            "createdAt": "2026-01-01T00:00:00Z",
+            "updatedAt": "2026-01-01T00:00:00Z",
+            "lastTurnNumber": 3,
+            "cwd": "/home/user/repo",
+            "repoRemoteUrl": null,
+            "hostname": null,
+            "status": "active",
+            "gcsTracePrefix": "sessions/sess-abc",
+            "gcsBucket": "my-bucket"
+        });
+        let record: SessionRecord = serde_json::from_value(json).unwrap();
+        assert_eq!(record.last_turn_number, 3);
+        assert_eq!(record.restorable_turn_number, None);
+    }
+
+    #[test]
+    fn session_record_with_restorable_turn_deserializes_correctly() {
+        let json = serde_json::json!({
+            "sessionId": "sess-xyz",
+            "summary": "hello",
+            "firstPrompt": null,
+            "modelId": null,
+            "createdAt": "2026-01-01T00:00:00Z",
+            "updatedAt": "2026-01-01T00:00:00Z",
+            "lastTurnNumber": 7,
+            "restorableTurnNumber": 6,
+            "cwd": "/home/user/repo",
+            "repoRemoteUrl": null,
+            "hostname": null,
+            "status": "active",
+            "gcsTracePrefix": "sessions/sess-xyz",
+            "gcsBucket": "my-bucket"
+        });
+        let record: SessionRecord = serde_json::from_value(json).unwrap();
+        assert_eq!(record.last_turn_number, 7);
+        assert_eq!(record.restorable_turn_number, Some(6));
+    }
+
+    /// Verify per-request auth resolve picks up rotated tokens.
+    #[tokio::test]
+    async fn session_registry_client_uses_active_auth_for_each_request() {
+        use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
+        use axum::{Router, response::IntoResponse, routing::post};
+        use chrono::{Duration, Utc};
+        use std::net::SocketAddr;
+        use std::sync::Arc;
+        use tokio::net::TcpListener;
+
+        let captured = Arc::new(parking_lot::Mutex::new(None::<String>));
+        let captured_for_handler = captured.clone();
+        let router = Router::new().route(
+            "/sessions/register",
+            post(move |headers: axum::http::HeaderMap, _body: String| {
+                let captured = captured_for_handler.clone();
+                async move {
+                    if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) {
+                        *captured.lock() = Some(auth.to_str().unwrap_or("").to_owned());
+                    }
+                    (axum::http::StatusCode::OK, "").into_response()
+                }
+            }),
+        );
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr: SocketAddr = listener.local_addr().unwrap();
+        tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
+
+        let dir = tempfile::tempdir().unwrap();
+        let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
+        am.hot_swap(GrokAuth {
+            key: "fresh-from-auth-manager".into(),
+            auth_mode: AuthMode::ApiKey,
+            create_time: Utc::now(),
+            user_id: "user-42".into(),
+            expires_at: Some(Utc::now() + Duration::hours(1)),
+            ..GrokAuth::test_default()
+        });
+
+        let client = SessionRegistryClient::new(format!("http://{addr}"), "STALE-build-time-token")
+            .with_auth(am);
+        let req = RegisterRequest {
+            session_id: "s1".into(),
+            cwd: "/x".into(),
+            gcs_trace_prefix: "t".into(),
+            model_id: None,
+            repo_remote_url: None,
+            repo_branch: None,
+            repo_head_at_start: None,
+            hostname: None,
+            device_id: None,
+            parent_session_id: None,
+            session_kind: None,
+            subagent_type: None,
+            subagent_persona: None,
+            subagent_role: None,
+            fork_context_source: None,
+            subagent_depth: None,
+        };
+        client.register(&req).await.unwrap();
+
+        let sent = captured.lock().clone().expect("server saw the request");
+        assert_eq!(
+            sent, "Bearer fresh-from-auth-manager",
+            "outgoing bearer must come from AuthManager (not the build-time token)"
+        );
+    }
+
+    // Verify the split-pointer invariant: last_turn_number can be ahead of
+    // restorable_turn_number (codebase best-effort means a turn may be "done"
+    // but not yet restorable if session-state upload is still in flight).
+    #[test]
+    fn session_record_allows_last_turn_ahead_of_restorable() {
+        let json = serde_json::json!({
+            "sessionId": "sess-lag",
+            "summary": "",
+            "firstPrompt": null,
+            "modelId": null,
+            "createdAt": "2026-01-01T00:00:00Z",
+            "updatedAt": "2026-01-01T00:00:00Z",
+            "lastTurnNumber": 10,
+            "restorableTurnNumber": 8,
+            "cwd": "/repo",
+            "repoRemoteUrl": null,
+            "hostname": null,
+            "status": "active",
+            "gcsTracePrefix": "sessions/sess-lag",
+            "gcsBucket": "bucket"
+        });
+        let record: SessionRecord = serde_json::from_value(json).unwrap();
+        assert!(record.last_turn_number > record.restorable_turn_number.unwrap_or(0));
+    }
+}

crates/codegen/xai-grok-shell/src/agent/storage_client_tests.rs

@@ -0,0 +1,337 @@
+//! Tests for StorageClient retry logic.
+//!
+//! Uses a local axum server to simulate various HTTP error scenarios
+//! and verify that the client handles retries correctly.
+
+use std::net::SocketAddr;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::time::{Duration, Instant};
+
+use axum::{
+    Router,
+    body::Body,
+    extract::State,
+    http::{HeaderMap, StatusCode},
+    response::{IntoResponse, Response},
+    routing::post,
+};
+use tokio::net::TcpListener;
+
+use xai_file_utils::storage_client::{RetryConfig, StorageClient};
+
+/// Shared state for tracking request counts in tests.
+#[derive(Clone, Default)]
+struct TestServerState {
+    request_count: Arc<AtomicU32>,
+    /// Number of 429 responses to return before succeeding
+    fail_count: Arc<AtomicU32>,
+    /// Optional Retry-After header value in seconds
+    retry_after_secs: Option<u32>,
+}
+
+impl TestServerState {
+    fn new(fail_count: u32) -> Self {
+        Self {
+            request_count: Arc::new(AtomicU32::new(0)),
+            fail_count: Arc::new(AtomicU32::new(fail_count)),
+            retry_after_secs: None,
+        }
+    }
+
+    fn with_retry_after(mut self, secs: u32) -> Self {
+        self.retry_after_secs = Some(secs);
+        self
+    }
+
+    fn get_request_count(&self) -> u32 {
+        self.request_count.load(Ordering::SeqCst)
+    }
+}
+
+/// Handler that returns 429 for the first N requests, then succeeds.
+async fn upload_handler_429(
+    State(state): State<TestServerState>,
+    _headers: HeaderMap,
+    _body: Body,
+) -> Response {
+    let count = state.request_count.fetch_add(1, Ordering::SeqCst);
+    let fail_count = state.fail_count.load(Ordering::SeqCst);
+
+    if count < fail_count {
+        let mut response = (
+            StatusCode::TOO_MANY_REQUESTS,
+            r#"{"error": "rate limited"}"#,
+        )
+            .into_response();
+
+        // Add Retry-After header if configured
+        if let Some(secs) = state.retry_after_secs {
+            response
+                .headers_mut()
+                .insert("retry-after", secs.to_string().parse().unwrap());
+        }
+
+        return response;
+    }
+
+    // Success response
+    (
+        StatusCode::OK,
+        r#"{"bucket": "test-bucket", "path": "test/path", "size": 100, "content_type": "application/json", "generation": 1}"#,
+    )
+        .into_response()
+}
+
+/// Handler that returns 500 for the first N requests, then succeeds.
+async fn upload_handler_500(
+    State(state): State<TestServerState>,
+    _headers: HeaderMap,
+    _body: Body,
+) -> Response {
+    let count = state.request_count.fetch_add(1, Ordering::SeqCst);
+    let fail_count = state.fail_count.load(Ordering::SeqCst);
+
+    if count < fail_count {
+        return (
+            StatusCode::INTERNAL_SERVER_ERROR,
+            r#"{"error": "internal server error"}"#,
+        )
+            .into_response();
+    }
+
+    // Success response
+    (
+        StatusCode::OK,
+        r#"{"bucket": "test-bucket", "path": "test/path", "size": 100, "content_type": "application/json", "generation": 1}"#,
+    )
+        .into_response()
+}
+
+/// Handler that always returns 400 (non-retryable).
+async fn upload_handler_400(
+    State(state): State<TestServerState>,
+    _headers: HeaderMap,
+    _body: Body,
+) -> Response {
+    state.request_count.fetch_add(1, Ordering::SeqCst);
+
+    (StatusCode::BAD_REQUEST, r#"{"error": "bad request"}"#).into_response()
+}
+
+/// Handler that always returns 429 (never succeeds).
+async fn upload_handler_always_429(
+    State(state): State<TestServerState>,
+    _headers: HeaderMap,
+    _body: Body,
+) -> Response {
+    state.request_count.fetch_add(1, Ordering::SeqCst);
+
+    let mut response = (
+        StatusCode::TOO_MANY_REQUESTS,
+        r#"{"error": "rate limited"}"#,
+    )
+        .into_response();
+
+    if let Some(secs) = state.retry_after_secs {
+        response
+            .headers_mut()
+            .insert("retry-after", secs.to_string().parse().unwrap());
+    }
+
+    response
+}
+
+/// Handler that always succeeds immediately.
+async fn upload_handler_success(
+    State(state): State<TestServerState>,
+    _headers: HeaderMap,
+    _body: Body,
+) -> Response {
+    state.request_count.fetch_add(1, Ordering::SeqCst);
+
+    (
+        StatusCode::OK,
+        r#"{"bucket": "test-bucket", "path": "test/path", "size": 100, "content_type": "application/json", "generation": 1}"#,
+    )
+        .into_response()
+}
+
+/// Start a test server with the given handler and return its address.
+async fn start_test_server<H, T>(state: TestServerState, handler: H) -> SocketAddr
+where
+    H: axum::handler::Handler<T, TestServerState> + Clone + Send + 'static,
+    T: 'static,
+{
+    let app = Router::new()
+        .route("/v1/storage", post(handler))
+        .with_state(state);
+
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+
+    tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+
+    // Give the server a moment to start
+    tokio::time::sleep(Duration::from_millis(10)).await;
+
+    addr
+}
+
+#[tokio::test]
+async fn test_upload_succeeds_on_first_try() {
+    let state = TestServerState::new(0);
+    let addr = start_test_server(state.clone(), upload_handler_success).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token");
+
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+
+    assert!(result.is_ok());
+    assert_eq!(state.get_request_count(), 1);
+}
+
+#[tokio::test]
+async fn test_upload_retries_on_429() {
+    let state = TestServerState::new(2); // Fail twice, then succeed
+    let addr = start_test_server(state.clone(), upload_handler_429).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token")
+        .with_retry_config(
+            RetryConfig::new()
+                .with_initial_delay(Duration::from_millis(10))
+                .with_max_retries(5),
+        );
+
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+
+    assert!(result.is_ok());
+    assert_eq!(state.get_request_count(), 3); // 2 failures + 1 success
+}
+
+#[tokio::test]
+async fn test_upload_retries_on_500() {
+    let state = TestServerState::new(2); // Fail twice, then succeed
+    let addr = start_test_server(state.clone(), upload_handler_500).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token")
+        .with_retry_config(
+            RetryConfig::new()
+                .with_initial_delay(Duration::from_millis(10))
+                .with_max_retries(5),
+        );
+
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+
+    assert!(result.is_ok());
+    assert_eq!(state.get_request_count(), 3); // 2 failures + 1 success
+}
+
+#[tokio::test]
+async fn test_upload_does_not_retry_on_400() {
+    let state = TestServerState::new(0);
+    let addr = start_test_server(state.clone(), upload_handler_400).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token")
+        .with_retry_config(
+            RetryConfig::new()
+                .with_initial_delay(Duration::from_millis(10))
+                .with_max_retries(5),
+        );
+
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+
+    assert!(result.is_err());
+    assert_eq!(state.get_request_count(), 1); // Only 1 request, no retries
+}
+
+#[tokio::test]
+async fn test_upload_respects_max_retries() {
+    let state = TestServerState::new(100); // Always fail
+    let addr = start_test_server(state.clone(), upload_handler_always_429).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token")
+        .with_retry_config(
+            RetryConfig::new()
+                .with_initial_delay(Duration::from_millis(10))
+                .with_max_retries(3),
+        );
+
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+
+    assert!(result.is_err());
+    // Should have 1 initial request + 3 retries = 4 total
+    assert_eq!(state.get_request_count(), 4);
+}
+
+#[tokio::test]
+async fn test_upload_respects_retry_after_header() {
+    let state = TestServerState::new(1).with_retry_after(1); // 1 second Retry-After
+    let addr = start_test_server(state.clone(), upload_handler_429).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token")
+        .with_retry_config(
+            RetryConfig::new()
+                .with_initial_delay(Duration::from_millis(10))
+                .with_max_retries(3),
+        );
+
+    let start = Instant::now();
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+    let elapsed = start.elapsed();
+
+    assert!(result.is_ok());
+    // Should have waited at least 1 second due to Retry-After header
+    assert!(
+        elapsed >= Duration::from_millis(900),
+        "Expected delay >= 900ms due to Retry-After, got {:?}",
+        elapsed
+    );
+}
+
+#[tokio::test]
+async fn test_exponential_backoff_increases_delay() {
+    let state = TestServerState::new(3); // Fail 3 times, then succeed
+    let addr = start_test_server(state.clone(), upload_handler_429).await;
+
+    let client = StorageClient::new(&format!("http://{}/v1", addr), "test-token")
+        .with_retry_config(
+            RetryConfig::new()
+                .with_initial_delay(Duration::from_millis(50))
+                .with_multiplier(2.0)
+                .with_jitter_factor(0.0) // No jitter for predictable timing
+                .with_max_retries(5),
+        );
+
+    let start = Instant::now();
+    let result = client
+        .upload("test/path", b"test content", "text/plain")
+        .await;
+    let elapsed = start.elapsed();
+
+    assert!(result.is_ok());
+
+    // With 3 failures before success:
+    // Delay 1: 50ms, Delay 2: 100ms, Delay 3: 200ms = 350ms minimum
+    assert!(
+        elapsed >= Duration::from_millis(300),
+        "Expected delay >= 300ms for exponential backoff, got {:?}",
+        elapsed
+    );
+}
+
+// ============================================================================

crates/codegen/xai-grok-shell/src/agent/subagent/coordinator_lifecycle.rs

@@ -0,0 +1,507 @@
+#![cfg_attr(rustfmt, rustfmt::skip)]
+#![allow(unused_imports)]
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use agent_client_protocol as acp;
+use tokio::sync::{Notify, mpsc, oneshot};
+use tokio_util::sync::CancellationToken;
+use crate::extensions::notification::{SessionNotification, SessionUpdate};
+use crate::session::{
+    self, SessionCommand, SessionHandle, SessionThread,
+    commands::{PromptCompletionKind, PromptTurnResult as SubagentPromptTurnResult},
+    fs_watch::FsWatchCapabilities, info::Info as SessionInfo,
+};
+use crate::terminal::AsyncTerminalRunner;
+use crate::tools::ToolContext;
+use crate::upload::trace::{
+    GCS_SCHEMA_VERSION, PromptMetadata, SubagentSpawnedRef, TurnResultMetadata,
+    local_sandbox_telemetry, upload_config, upload_metadata, upload_session_state,
+    upload_subagent_metadata, upload_turn_result,
+};
+use crate::upload::turn::{PromptTraceContext, complete_prompt_trace};
+use xai_acp_lib::AcpAgentGatewaySender as GatewaySender;
+use xai_grok_tools::implementations::grok_build::task::types::*;
+use xai_grok_workspace::file_system::AsyncFileSystem;
+use xai_hunk_tracker::HunkTrackerHandle;
+use super::*;
+impl SubagentCoordinator {
+    pub fn new() -> Self {
+        Self {
+            pending: HashMap::new(),
+            active: HashMap::new(),
+            completed: HashMap::new(),
+            completion_notify: Arc::new(Notify::new()),
+            pending_completions: Vec::new(),
+            is_turn_active: Arc::new(std::sync::atomic::AtomicBool::new(false)),
+            synthetic_trace_tx: None,
+            running_gauge: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
+            block_wait_slots: HashMap::new(),
+            subagent_usage_not_applied_prompts: std::collections::HashSet::new(),
+        }
+    }
+    pub fn mark_subagent_usage_not_applied(&mut self, prompt_id: &str) {
+        self.subagent_usage_not_applied_prompts.insert(prompt_id.to_string());
+    }
+    pub fn subagent_usage_not_applied(&self, prompt_id: &str) -> bool {
+        self.subagent_usage_not_applied_prompts.contains(prompt_id)
+    }
+    pub fn clear_subagent_usage_not_applied(&mut self, prompt_id: &str) {
+        self.subagent_usage_not_applied_prompts.remove(prompt_id);
+    }
+    pub fn parent_prompt_id_for(&self, subagent_id: &str) -> Option<String> {
+        self.active
+            .get(subagent_id)
+            .and_then(|t| t.parent_prompt_id.clone())
+            .or_else(|| {
+                self.pending.get(subagent_id).and_then(|p| p.parent_prompt_id.clone())
+            })
+    }
+    /// Rebind the running-subagent gauge, copying the current count so a
+    /// late rebind cannot under-report.
+    pub fn set_running_gauge(&mut self, gauge: Arc<std::sync::atomic::AtomicUsize>) {
+        gauge
+            .store(
+                self.pending.len() + self.active.len(),
+                std::sync::atomic::Ordering::Relaxed,
+            );
+        self.running_gauge = gauge;
+    }
+    /// Recompute the gauge from `pending` + `active` after every mutation of
+    /// either map — recomputing (rather than incrementing) prevents drift.
+    fn sync_running_gauge(&self) {
+        self.running_gauge
+            .store(
+                self.pending.len() + self.active.len(),
+                std::sync::atomic::Ordering::Relaxed,
+            );
+    }
+    /// Returns a handle to the completion [`Notify`].
+    #[cfg_attr(
+        not(test),
+        expect(
+            dead_code,
+            reason = "used from tests only; remove expect when wired in production"
+        )
+    )]
+    pub fn completion_notify(&self) -> Arc<Notify> {
+        Arc::clone(&self.completion_notify)
+    }
+    /// Returns a shared handle to the turn-active flag.
+    pub fn turn_active_flag(&self) -> Arc<std::sync::atomic::AtomicBool> {
+        Arc::clone(&self.is_turn_active)
+    }
+    /// Whether the model's turn is currently active.
+    #[cfg_attr(
+        not(test),
+        expect(
+            dead_code,
+            reason = "used from tests only; remove expect when wired in production"
+        )
+    )]
+    pub fn is_turn_active(&self) -> bool {
+        self.is_turn_active.load(std::sync::atomic::Ordering::Relaxed)
+    }
+    /// Pending + active turn-blocking subagent IDs for `prompt_id`.
+    /// Background children are excluded: they outlive the turn by design, so
+    /// the freeze drain must not wait on them (their spend reaches the session
+    /// ledger when they finish; the prompt report flags them via
+    /// `background_live`).
+    pub fn outstanding_for_prompt(&self, prompt_id: &str) -> Vec<String> {
+        let mut ids: Vec<String> = self
+            .pending
+            .values()
+            .filter(|p| {
+                p.parent_prompt_id.as_deref() == Some(prompt_id) && !p.run_in_background
+            })
+            .map(|p| p.subagent_id.clone())
+            .chain(
+                self
+                    .active
+                    .values()
+                    .filter(|t| {
+                        t.parent_prompt_id.as_deref() == Some(prompt_id)
+                            && !t.run_in_background
+                    })
+                    .map(|t| t.subagent_id.clone()),
+            )
+            .collect();
+        ids.sort();
+        ids
+    }
+    /// True while any background child of `prompt_id` is pending or active.
+    /// Their spend is missing from the prompt report (it lands on the session
+    /// ledger at completion), so the report is incomplete — without waiting.
+    pub fn background_live_for_prompt(&self, prompt_id: &str) -> bool {
+        self
+            .pending
+            .values()
+            .any(|p| {
+                p.parent_prompt_id.as_deref() == Some(prompt_id) && p.run_in_background
+            })
+            || self
+                .active
+                .values()
+                .any(|t| {
+                    t.parent_prompt_id.as_deref() == Some(prompt_id)
+                        && t.run_in_background
+                })
+    }
+    /// Record that a foreground child was auto-backgrounded (await budget
+    /// expired): it no longer blocks the turn, so the freeze drain must stop
+    /// waiting on it.
+    pub fn mark_backgrounded(&mut self, subagent_id: &str) {
+        if let Some(t) = self.active.values_mut().find(|t| t.subagent_id == subagent_id)
+        {
+            t.run_in_background = true;
+        }
+        if let Some(p) = self.pending.values_mut().find(|p| p.subagent_id == subagent_id)
+        {
+            p.run_in_background = true;
+        }
+    }
+    pub fn outstanding_reply_for_prompt(
+        &self,
+        prompt_id: &str,
+    ) -> xai_grok_tools::implementations::grok_build::task::types::SubagentOutstandingReply {
+        xai_grok_tools::implementations::grok_build::task::types::SubagentOutstandingReply {
+            live_ids: self.outstanding_for_prompt(prompt_id),
+            background_live: self.background_live_for_prompt(prompt_id),
+            subagent_usage_not_applied: self.subagent_usage_not_applied(prompt_id),
+        }
+    }
+    /// Drain all buffered completion summaries, returning them and clearing the buffer.
+    pub fn drain_pending_completions(&mut self) -> Vec<SubagentCompletionSummary> {
+        std::mem::take(&mut self.pending_completions)
+    }
+    /// Collect references to subagents spawned for a specific parent prompt.
+    /// Returns only the children whose `parent_prompt_id` matches, so the
+    /// parent turn's `turn_result.json` accurately reflects what was spawned
+    /// during that turn — not the entire coordinator lifetime.
+    pub fn spawned_refs_for_prompt(&self, prompt_id: &str) -> Vec<SubagentSpawnedRef> {
+        let mut refs: Vec<_> = self
+            .active
+            .values()
+            .filter(|t| t.parent_prompt_id.as_deref() == Some(prompt_id))
+            .map(|t| SubagentSpawnedRef {
+                subagent_id: t.subagent_id.clone(),
+                child_session_id: t.child_session_id.0.to_string(),
+                subagent_type: t.subagent_type.clone(),
+                description: t.description.clone(),
+                persona: t.persona.clone(),
+                resumed_from: t.resumed_from.clone(),
+            })
+            .chain(
+                self
+                    .completed
+                    .values()
+                    .filter(|c| c.parent_prompt_id.as_deref() == Some(prompt_id))
+                    .map(|c| SubagentSpawnedRef {
+                        subagent_id: c.subagent_id.clone(),
+                        child_session_id: c.child_session_id.clone(),
+                        subagent_type: c.subagent_type.clone(),
+                        description: c.description.clone(),
+                        persona: c.persona.clone(),
+                        resumed_from: c.resumed_from.clone(),
+                    }),
+            )
+            .collect();
+        refs.sort_by(|a, b| a.subagent_id.cmp(&b.subagent_id));
+        refs
+    }
+    /// Register a subagent as pending (initializing). Call this early,
+    /// before any blocking work like worktree creation, so that
+    /// `get_task_output` can report the subagent as initializing instead
+    /// of "not found".
+    pub fn insert_pending(&mut self, entry: PendingSubagent) {
+        self.pending.insert(entry.subagent_id.clone(), entry);
+        self.sync_running_gauge();
+    }
+    /// Remove a pending subagent without recording a failure.
+    /// Used by cancel flows where the subagent was intentionally stopped.
+    #[cfg(test)]
+    pub fn remove_pending(&mut self, id: &str) {
+        self.pending.remove(id);
+        self.sync_running_gauge();
+    }
+    /// Move a pending subagent directly to `completed` so it stays queryable via
+    /// `get_task_output`. `cancelled` stamps `"cancelled"` vs `"failed"`.
+    fn move_pending_to_terminal(&mut self, id: &str, error: &str, cancelled: bool) {
+        let Some(pending) = self.pending.remove(id) else {
+            return;
+        };
+        self.record_failure_completion(FailureCompletion {
+            subagent_id: pending.subagent_id,
+            subagent_type: pending.subagent_type,
+            description: pending.description,
+            parent_prompt_id: pending.parent_prompt_id,
+            parent_session_id: pending.parent_session_id,
+            persona: pending.persona,
+            started_at: pending.started_at,
+            error,
+            surface_completion: pending.surface_completion,
+            cancelled,
+        });
+    }
+    /// Move a pending subagent to `completed` as a failure so it stays queryable
+    /// via `get_task_output`.
+    pub fn move_pending_to_failed(&mut self, id: &str, error: &str) {
+        self.move_pending_to_terminal(id, error, false);
+    }
+    /// Like [`Self::move_pending_to_failed`] but stamps `"cancelled"` — a pending
+    /// subagent killed while initializing.
+    pub fn move_pending_to_cancelled(&mut self, id: &str, error: &str) {
+        self.move_pending_to_terminal(id, error, true);
+    }
+    /// Record a synthetic failure for a subagent that never reached `pending`.
+    pub fn record_pre_spawn_failure(
+        &mut self,
+        subagent_id: String,
+        subagent_type: String,
+        description: String,
+        parent_prompt_id: Option<String>,
+        parent_session_id: String,
+        error: &str,
+        surface_completion: bool,
+    ) {
+        self.record_failure_completion(FailureCompletion {
+            subagent_id,
+            subagent_type,
+            description,
+            parent_prompt_id,
+            parent_session_id,
+            persona: None,
+            started_at: std::time::Instant::now(),
+            error,
+            surface_completion,
+            cancelled: false,
+        });
+    }
+    /// Insert a synthetic failed entry, push a completion summary, notify waiters.
+    /// Clears any stale pending entry for the same id.
+    fn record_failure_completion(&mut self, c: FailureCompletion<'_>) {
+        self.pending.remove(&c.subagent_id);
+        self.sync_running_gauge();
+        let FailureCompletion {
+            subagent_id,
+            subagent_type,
+            description,
+            parent_prompt_id,
+            parent_session_id,
+            persona,
+            started_at,
+            error,
+            surface_completion,
+            cancelled,
+        } = c;
+        let result = SubagentResult {
+            success: false,
+            cancelled,
+            error: Some(error.to_string()),
+            subagent_id: subagent_id.clone(),
+            ..Default::default()
+        };
+        let summary_output = result.output.clone();
+        self.completed
+            .insert(
+                subagent_id.clone(),
+                CompletedSubagent {
+                    subagent_id: subagent_id.clone(),
+                    parent_session_id,
+                    parent_prompt_id,
+                    child_session_id: String::new(),
+                    description: description.clone(),
+                    subagent_type: subagent_type.clone(),
+                    persona,
+                    started_at,
+                    completed_at: std::time::Instant::now(),
+                    result,
+                    resumed_from: None,
+                    child_cwd: String::new(),
+                    worktree_path: None,
+                    snapshot_ref: None,
+                    effective_model_id: String::new(),
+                    block_waited: false,
+                    explicitly_killed: false,
+                },
+            );
+        if surface_completion {
+            self.pending_completions
+                .push(SubagentCompletionSummary {
+                    subagent_id,
+                    subagent_type,
+                    description,
+                    success: false,
+                    duration_ms: 0,
+                    tool_calls: 0,
+                    turns: 0,
+                    output: summary_output,
+                });
+        }
+        self.completion_notify.notify_waiters();
+    }
+    pub fn insert(&mut self, tracker: SubagentTracker) {
+        self.pending.remove(&tracker.subagent_id);
+        self.active.insert(tracker.subagent_id.clone(), tracker);
+        self.sync_running_gauge();
+    }
+    /// Move a finished subagent from `active` to `completed`.
+    /// Returns the tracker if it was active.
+    pub fn move_to_completed(
+        &mut self,
+        id: &str,
+        description: String,
+        subagent_type: String,
+        result: SubagentResult,
+    ) -> Option<SubagentTracker> {
+        let tracker = self.active.remove(id);
+        self.sync_running_gauge();
+        let started_at = tracker
+            .as_ref()
+            .map(|t| t.started_at)
+            .unwrap_or_else(std::time::Instant::now);
+        let parent_session_id = tracker
+            .as_ref()
+            .map(|t| t.parent_session_id.clone())
+            .unwrap_or_default();
+        let child_session_id = tracker
+            .as_ref()
+            .map(|t| t.child_session_id.0.to_string())
+            .unwrap_or_default();
+        let parent_prompt_id = tracker.as_ref().and_then(|t| t.parent_prompt_id.clone());
+        let persona = tracker.as_ref().and_then(|t| t.persona.clone());
+        let child_cwd = tracker
+            .as_ref()
+            .map(|t| t.child_cwd.clone())
+            .unwrap_or_default();
+        let worktree_path = tracker.as_ref().and_then(|t| t.worktree_path.clone());
+        let resumed_from = tracker.as_ref().and_then(|t| t.resumed_from.clone());
+        let effective_model_id = tracker
+            .as_ref()
+            .map(|t| t.effective_model_id.clone())
+            .unwrap_or_default();
+        let block_waited = tracker.as_ref().is_some_and(|t| t.block_waited);
+        let explicitly_killed = tracker.as_ref().is_some_and(|t| t.explicitly_killed);
+        let surface_completion = tracker.as_ref().is_none_or(|t| t.surface_completion);
+        self.completed
+            .insert(
+                id.to_string(),
+                CompletedSubagent {
+                    subagent_id: id.to_string(),
+                    parent_session_id,
+                    parent_prompt_id,
+                    child_session_id,
+                    description,
+                    subagent_type,
+                    persona,
+                    started_at,
+                    completed_at: std::time::Instant::now(),
+                    result,
+                    resumed_from,
+                    child_cwd,
+                    worktree_path,
+                    snapshot_ref: None,
+                    effective_model_id,
+                    block_waited,
+                    explicitly_killed,
+                },
+            );
+        let completed = self.completed.get(id).expect("just inserted");
+        let success = completed.result.success && !completed.result.cancelled;
+        {
+            let preview = crate::util::truncate(&completed.result.output, 200);
+            let level_fn = if success {
+                xai_grok_telemetry::unified_log::info
+            } else {
+                xai_grok_telemetry::unified_log::error
+            };
+            level_fn(
+                if success { "subagent completed" } else { "subagent failed" },
+                None,
+                Some(
+                    serde_json::json!(
+                        { "subagent_id" : & completed.subagent_id, "subagent_type" : &
+                        completed.subagent_type, "effective_model" : & completed
+                        .effective_model_id, "success" : success, "cancelled" : completed
+                        .result.cancelled, "duration_ms" : completed.result.duration_ms,
+                        "turns" : completed.result.turns, "tool_calls" : completed.result
+                        .tool_calls, "output_preview" : preview, "error" : & completed
+                        .result.error, }
+                    ),
+                ),
+            );
+        }
+        if surface_completion {
+            self.pending_completions
+                .push(SubagentCompletionSummary {
+                    subagent_id: id.to_string(),
+                    subagent_type: completed.subagent_type.clone(),
+                    description: completed.description.clone(),
+                    success,
+                    duration_ms: completed.result.duration_ms,
+                    tool_calls: completed.result.tool_calls,
+                    turns: completed.result.turns,
+                    output: completed.result.output.clone(),
+                });
+        }
+        self.completion_notify.notify_waiters();
+        tracker
+    }
+    /// Record the durable worktree snapshot ref on a completed subagent so
+    /// in-memory `resume_from` resolution can rehydrate the disposed worktree.
+    /// No-op if the entry was already evicted (the on-disk meta.json still has it).
+    pub fn set_completed_snapshot_ref(&mut self, id: &str, snapshot_ref: String) {
+        if let Some(completed) = self.completed.get_mut(id) {
+            completed.snapshot_ref = Some(snapshot_ref);
+        }
+    }
+    /// Cancel all active subagents that were launched by a specific parent turn,
+    /// including `run_in_background: true` subagents.
+    pub fn cancel_by_parent_prompt_id(&mut self, parent_prompt_id: &str) {
+        for tracker in self.active.values() {
+            if tracker.parent_prompt_id.as_deref() == Some(parent_prompt_id) {
+                Self::cancel_tracker(tracker);
+            }
+        }
+        for pending in self.pending.values() {
+            if pending.parent_prompt_id.as_deref() == Some(parent_prompt_id) {
+                pending.cancel_token.cancel();
+            }
+        }
+    }
+    /// Attempt to cancel a subagent. Returns a typed outcome covering all cases:
+    /// - Active → cancel it, return Cancelled
+    /// - Pending (initializing) → fire its spawn token, return Cancelled
+    /// - Already finished → return AlreadyFinished with terminal status
+    /// - Unknown ID → return NotFound
+    pub fn cancel_with_outcome(&mut self, subagent_id: &str) -> SubagentCancelOutcome {
+        if let Some(tracker) = self.active.get(subagent_id) {
+            Self::cancel_tracker(tracker);
+            return SubagentCancelOutcome::Cancelled;
+        }
+        if let Some(pending) = self.pending.get(subagent_id) {
+            pending.cancel_token.cancel();
+            return SubagentCancelOutcome::Cancelled;
+        }
+        if let Some(entry) = self.completed.get(subagent_id) {
+            return SubagentCancelOutcome::AlreadyFinished {
+                status: entry.result.status().to_string(),
+            };
+        }
+        SubagentCancelOutcome::NotFound
+    }
+    /// Internal: send Cancel + Shutdown to a tracked subagent.
+    fn cancel_tracker(tracker: &SubagentTracker) {
+        tracker.cancel_token.cancel();
+        let _ = tracker
+            .child_handle
+            .cmd_tx
+            .send(SessionCommand::Cancel {
+                cancel_subagents: true,
+                kill_background_tasks: true,
+                rewind_if_pristine: false,
+                trigger: None,
+            });
+        let _ = tracker.child_handle.cmd_tx.send(SessionCommand::Shutdown);
+    }
+}

crates/codegen/xai-grok-shell/src/agent/subagent/coordinator_query.rs

@@ -0,0 +1,354 @@
+#![cfg_attr(rustfmt, rustfmt::skip)]
+#![allow(unused_imports)]
+use std::collections::HashMap;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use agent_client_protocol as acp;
+use tokio::sync::{Notify, mpsc, oneshot};
+use tokio_util::sync::CancellationToken;
+use crate::extensions::notification::{SessionNotification, SessionUpdate};
+use crate::session::{
+    self, SessionCommand, SessionHandle, SessionThread,
+    commands::{PromptCompletionKind, PromptTurnResult as SubagentPromptTurnResult},
+    fs_watch::FsWatchCapabilities, info::Info as SessionInfo,
+};
+use crate::terminal::AsyncTerminalRunner;
+use crate::tools::ToolContext;
+use crate::upload::trace::{
+    GCS_SCHEMA_VERSION, PromptMetadata, SubagentSpawnedRef, TurnResultMetadata,
+    local_sandbox_telemetry, upload_config, upload_metadata, upload_session_state,
+    upload_subagent_metadata, upload_turn_result,
+};
+use crate::upload::turn::{PromptTraceContext, complete_prompt_trace};
+use xai_acp_lib::AcpAgentGatewaySender as GatewaySender;
+use xai_grok_tools::implementations::grok_build::task::types::*;
+use xai_grok_workspace::file_system::AsyncFileSystem;
+use xai_hunk_tracker::HunkTrackerHandle;
+use super::*;
+impl SubagentCoordinator {
+    /// Synchronous lookup of a subagent by ID.
+    ///
+    /// Returns a three-way result so the caller can drop the `RefCell` borrow
+    /// before awaiting the signals handle for running subagents.
+    ///
+    /// - `Ready` — completed/failed/cancelled snapshot, no async work needed.
+    /// - `NeedsSignals` — subagent is running; caller must await
+    ///   `resolve_snapshot()` after dropping the coordinator borrow.
+    /// - `None` — ID not found in active, completed, or pending maps.
+    pub(crate) fn lookup(&self, id: &str) -> Option<SnapshotLookup> {
+        if let Some(tracker) = self.active.get(id) {
+            return Some(
+                SnapshotLookup::NeedsSignals(RunningSnapshotSeed {
+                    subagent_id: tracker.subagent_id.clone(),
+                    description: tracker.description.clone(),
+                    subagent_type: tracker.subagent_type.clone(),
+                    started_at_epoch_ms: instant_to_epoch_ms(tracker.started_at),
+                    duration_ms: tracker.started_at.elapsed().as_millis() as u64,
+                    persona: tracker.persona.clone(),
+                    signals_handle: tracker.child_handle.signals_handle.clone(),
+                }),
+            );
+        }
+        if let Some(completed) = self.completed.get(id) {
+            let status = if completed.result.cancelled {
+                SubagentSnapshotStatus::Cancelled {
+                    reason: completed.result.error.clone(),
+                }
+            } else if completed.result.success {
+                SubagentSnapshotStatus::Completed {
+                    output: completed.result.output.to_string(),
+                    tool_calls: completed.result.tool_calls,
+                    turns: completed.result.turns,
+                    worktree_path: completed.result.worktree_path.clone(),
+                }
+            } else {
+                SubagentSnapshotStatus::Failed {
+                        error: completed
+                            .result
+                            .error
+                            .clone()
+                            .unwrap_or_else(|| "Unknown error".to_string()),
+                    }
+            };
+            return Some(
+                SnapshotLookup::Ready(SubagentSnapshot {
+                    subagent_id: completed.subagent_id.clone(),
+                    description: completed.description.clone(),
+                    subagent_type: completed.subagent_type.clone(),
+                    status,
+                    started_at_epoch_ms: instant_to_epoch_ms(completed.started_at),
+                    duration_ms: completed.result.duration_ms,
+                    persona: completed.persona.clone(),
+                }),
+            );
+        }
+        if let Some(pending) = self.pending.get(id) {
+            return Some(
+                SnapshotLookup::Ready(SubagentSnapshot {
+                    subagent_id: pending.subagent_id.clone(),
+                    description: pending.description.clone(),
+                    subagent_type: pending.subagent_type.clone(),
+                    status: SubagentSnapshotStatus::Initializing,
+                    started_at_epoch_ms: instant_to_epoch_ms(pending.started_at),
+                    duration_ms: pending.started_at.elapsed().as_millis() as u64,
+                    persona: pending.persona.clone(),
+                }),
+            );
+        }
+        None
+    }
+    /// Return `(parent_session_id, child_session_id)` for a given subagent.
+    ///
+    /// Checks active first, then completed. Returns `None` if not found.
+    pub(crate) fn session_ids_for(&self, id: &str) -> Option<(String, String)> {
+        if let Some(t) = self.active.get(id) {
+            return Some((t.parent_session_id.clone(), t.child_session_id.0.to_string()));
+        }
+        if let Some(c) = self.completed.get(id) {
+            return Some((c.parent_session_id.clone(), c.child_session_id.clone()));
+        }
+        None
+    }
+    /// Mark a subagent as block-waited so auto-wake is suppressed on completion.
+    pub(crate) fn mark_block_waited(&mut self, id: &str) {
+        if let Some(t) = self.active.get_mut(id) {
+            t.block_waited = true;
+        } else if let Some(c) = self.completed.get_mut(id) {
+            c.block_waited = true;
+        }
+    }
+    /// Clear the block-waited flag after a block timed out without receiving
+    /// the completion, so auto-wake can still fire when the subagent finishes.
+    pub(crate) fn clear_block_waited(&mut self, id: &str) {
+        if let Some(t) = self.active.get_mut(id) {
+            t.block_waited = false;
+        } else if let Some(c) = self.completed.get_mut(id) {
+            c.block_waited = false;
+        }
+    }
+    /// Whether a block-waiter already consumed this subagent's result.
+    pub(crate) fn is_block_waited(&self, id: &str) -> bool {
+        self.active.get(id).is_some_and(|t| t.block_waited)
+            || self.completed.get(id).is_some_and(|c| c.block_waited)
+    }
+    /// Register a live blocking-query reply slot and mark `block_waited`.
+    ///
+    /// The slot lets `block_wait_delivered_or_live` verify at completion
+    /// time that the waiter can still receive the result — the flag alone
+    /// can be stale when the waiting turn was cancelled moments before the
+    /// subagent finished.
+    pub(crate) fn register_block_wait(&mut self, id: &str, slot: BlockWaitSlot) {
+        self.mark_block_waited(id);
+        self.block_wait_slots.entry(id.to_string()).or_default().push(slot);
+    }
+    /// Drop a previously registered reply slot (query poll loop exited).
+    pub(crate) fn unregister_block_wait(&mut self, id: &str, slot: &BlockWaitSlot) {
+        if let Some(slots) = self.block_wait_slots.get_mut(id) {
+            slots.retain(|s| !std::rc::Rc::ptr_eq(s, slot));
+            if slots.is_empty() {
+                self.block_wait_slots.remove(id);
+            }
+        }
+    }
+    /// Decision-time gate for the completion auto-wake: returns true when
+    /// the result was already delivered to a blocking waiter, or a live
+    /// waiter is still parked and will receive it. When every registered
+    /// waiter is gone (receivers dropped by a cancelled turn), clears
+    /// `block_waited` and returns false so the auto-wake fires.
+    ///
+    /// This closes the race where the query poll loop clears the flag up to
+    /// one poll interval *after* the caller cancelled — the completion
+    /// handler could read the stale flag in that window and skip the wake.
+    /// Consumes the id's slot registrations (completion is terminal).
+    pub(crate) fn block_wait_delivered_or_live(&mut self, id: &str) -> bool {
+        let slots = self.block_wait_slots.remove(id).unwrap_or_default();
+        if !self.is_block_waited(id) {
+            return false;
+        }
+        let delivered_or_live = slots.is_empty()
+            || slots
+                .iter()
+                .any(|s| s.borrow().as_ref().is_none_or(|tx| !tx.is_closed()));
+        if !delivered_or_live {
+            self.clear_block_waited(id);
+        }
+        delivered_or_live
+    }
+    /// Mark a subagent as explicitly killed so auto-wake is suppressed on completion.
+    pub(crate) fn mark_explicitly_killed(&mut self, id: &str) {
+        if let Some(t) = self.active.get_mut(id) {
+            t.explicitly_killed = true;
+        } else if let Some(c) = self.completed.get_mut(id) {
+            c.explicitly_killed = true;
+        }
+    }
+    /// Whether the model explicitly killed this subagent via the kill tool.
+    pub(crate) fn is_explicitly_killed(&self, id: &str) -> bool {
+        self.active.get(id).is_some_and(|t| t.explicitly_killed)
+            || self.completed.get(id).is_some_and(|c| c.explicitly_killed)
+    }
+    /// Return fork provenance for a given subagent.
+    pub(crate) fn provenance_for(&self, id: &str) -> SubagentProvenance {
+        if let Some(t) = self.active.get(id) {
+            return SubagentProvenance {
+                fork_parent_prompt_id: t.parent_prompt_id.clone(),
+                resumed_from: t.resumed_from.clone(),
+            };
+        }
+        if let Some(c) = self.completed.get(id) {
+            return SubagentProvenance {
+                fork_parent_prompt_id: c.parent_prompt_id.clone(),
+                resumed_from: c.resumed_from.clone(),
+            };
+        }
+        SubagentProvenance::default()
+    }
+    /// Resolve a completed subagent scoped to the requesting parent session.
+    ///
+    /// Returns `None` if the subagent is not found, still active, or belongs
+    /// to a different parent session (prevents cross-session context bleed).
+    ///
+    /// Fast path: checks the in-memory `completed` map first. When that
+    /// misses (e.g. after TTL eviction), falls back to on-disk metadata
+    /// in `{parent_session_dir}/subagents/{id}/meta.json`.
+    pub(crate) fn resumable_source_for(
+        &self,
+        id: &str,
+        parent_session_id: &str,
+        parent_cwd: &Path,
+    ) -> Option<ResumeSourceData> {
+        if let Some(completed) = self.completed.get(id) {
+            if completed.parent_session_id != parent_session_id {
+                return None;
+            }
+            return Some(ResumeSourceData {
+                subagent_id: completed.subagent_id.clone(),
+                child_session_id: completed.child_session_id.clone(),
+                child_cwd: completed.child_cwd.clone(),
+                worktree_path: completed.worktree_path.clone(),
+                snapshot_ref: completed.snapshot_ref.clone(),
+                subagent_type: completed.subagent_type.clone(),
+                persona: completed.persona.clone(),
+                model_id: Some(completed.effective_model_id.clone()),
+            });
+        }
+        let parent_info = SessionInfo {
+            id: acp::SessionId::new(parent_session_id),
+            cwd: parent_cwd.to_string_lossy().to_string(),
+        };
+        let meta_path = session::persistence::session_dir(&parent_info)
+            .join("subagents")
+            .join(id)
+            .join("meta.json");
+        let data = std::fs::read_to_string(&meta_path).ok()?;
+        let meta: SubagentMeta = serde_json::from_str(&data).ok()?;
+        if meta.parent_session_id != parent_session_id {
+            return None;
+        }
+        match meta.status.as_str() {
+            "completed" | "failed" | "cancelled" => {}
+            _ => return None,
+        }
+        Some(ResumeSourceData {
+            subagent_id: meta.subagent_id,
+            child_session_id: meta.child_session_id,
+            child_cwd: meta.child_cwd.unwrap_or_default(),
+            worktree_path: meta.worktree_path.map(PathBuf::from),
+            snapshot_ref: meta.snapshot_ref,
+            subagent_type: meta.subagent_type,
+            persona: meta.persona,
+            model_id: meta.effective_model_id,
+        })
+    }
+    /// Check whether an ID refers to a currently-active (running) subagent.
+    pub(crate) fn is_active(&self, id: &str) -> bool {
+        self.active.contains_key(id)
+    }
+    /// Whether the coordinator still has this id in flight (spawning or running).
+    /// Orphan reconcile skips these — there is nothing stuck to heal.
+    pub(crate) fn is_active_or_pending(&self, id: &str) -> bool {
+        self.active.contains_key(id) || self.pending.contains_key(id)
+    }
+    /// The terminal `SubagentFinished` for an id the coordinator already holds in
+    /// `completed`, else `None`. Lets orphan reconcile re-emit a subagent's real
+    /// outcome when only its terminal meta write was lost (reconnect race: entry
+    /// in `completed` but the on-disk meta is still `running`) instead of
+    /// force-cancelling it and discarding the result.
+    pub(crate) fn completed_finish(&self, id: &str) -> Option<SessionUpdate> {
+        let c = self.completed.get(id)?;
+        let duration_ms = c
+            .completed_at
+            .saturating_duration_since(c.started_at)
+            .as_millis() as u64;
+        Some(SessionUpdate::SubagentFinished {
+            subagent_id: c.subagent_id.clone(),
+            child_session_id: c.child_session_id.clone(),
+            status: c.result.status().to_string(),
+            error: c.result.error.clone(),
+            tool_calls: c.result.tool_calls,
+            turns: c.result.turns,
+            duration_ms,
+            tokens_used: 0,
+            output: None,
+            will_wake: false,
+        })
+    }
+    /// TTL cleanup: remove completed entries older than 30 minutes.
+    pub fn evict_stale_completed(&mut self) {
+        let cutoff = std::time::Duration::from_secs(30 * 60);
+        self.completed.retain(|_, entry| entry.completed_at.elapsed() < cutoff);
+    }
+    /// Snapshot all currently-running subagents for compaction state context.
+    ///
+    /// Returns one `ActiveSubagentSummary` per entry in the `active` map.
+    /// Completed/failed/cancelled subagents are NOT included — they live in
+    /// the `completed` map and are irrelevant for post-compaction reminders
+    /// (the model already saw their tool results before compaction).
+    ///
+    /// The `elapsed_ms` field is computed from `started_at.elapsed()` at call
+    /// time, so the values are a snapshot of "right now" — appropriate for
+    /// compaction since it happens once and the reminder is static.
+    #[cfg(test)]
+    pub fn active_summaries(&self) -> Vec<ActiveSubagentSummary> {
+        self.active.values().map(tracker_to_summary).collect()
+    }
+    pub fn active_summaries_for(
+        &self,
+        parent_session_id: &str,
+    ) -> Vec<ActiveSubagentSummary> {
+        self.active
+            .values()
+            .filter(|t| t.parent_session_id == parent_session_id)
+            .map(tracker_to_summary)
+            .collect()
+    }
+    /// Return seeds for all running subagents belonging to `parent_session_id`.
+    ///
+    /// Each seed carries copied identity metadata plus a cloned
+    /// `SessionSignalsHandle` so the caller can resolve live progress
+    /// asynchronously after dropping the coordinator borrow.
+    ///
+    /// Returns an empty `Vec` if no active subagents match the given
+    /// parent session ID. Callers (e.g. the `x.ai/subagent/list_running`
+    /// ACP handler) should treat an empty result as a normal "no running
+    /// subagents" response, not an error.
+    pub(crate) fn list_running_for_parent(
+        &self,
+        parent_session_id: &str,
+    ) -> Vec<RunningSubagentListSeed> {
+        self.active
+            .values()
+            .filter(|t| t.parent_session_id == parent_session_id)
+            .map(|t| RunningSubagentListSeed {
+                subagent_id: t.subagent_id.clone(),
+                parent_session_id: t.parent_session_id.clone(),
+                child_session_id: t.child_session_id.0.to_string(),
+                subagent_type: t.subagent_type.clone(),
+                description: t.description.clone(),
+                started_at_epoch_ms: instant_to_epoch_ms(t.started_at),
+                duration_ms: t.started_at.elapsed().as_millis() as u64,
+                signals_handle: t.child_handle.signals_handle.clone(),
+            })
+            .collect()
+    }
+}

crates/codegen/xai-grok-shell/src/agent/subscription_check.rs

@@ -0,0 +1,191 @@
+//! Subscription check for paywall gate lift.
+//!
+//! Provides `single_check()` which queries `GET /user?include=subscription`
+//! for the live subscription tier from the backend, independent of the JWT.
+//! If a qualifying tier is detected, does a best-effort JWT refresh and
+//! settings re-fetch, then returns an `UnblockResult` so the agent can
+//! lift the gate.
+//!
+//! The pager drives the polling via `x.ai/auth/check_subscription`: the 5s
+//! paywall chain, the free-tier watch, the refocus check, and
+//! verify-before-paywall gate deferral (see the pager's `app::subscription`
+//! module).
+use crate::auth::AuthManager;
+use crate::auth::UserInfo;
+use crate::auth::manager::RefreshReason;
+use crate::auth::token_type::TokenType;
+use std::sync::Arc;
+use std::time::Duration;
+/// Subscription tiers that qualify for Grok Build access.
+/// Any active subscription qualifies -- the access gate in remote settings
+/// controls which tiers are actually allowed.
+const QUALIFYING_TIERS: &[&str] = &[
+    "SuperGrokPro",
+    "GrokPro",
+    "SuperGrokLite",
+    "XPremiumPlus",
+    "XPremium",
+    "XBasic",
+];
+/// Successful subscription check result: confirmed qualifying tier +
+/// optionally refreshed settings.
+pub(crate) struct UnblockResult {
+    pub(crate) new_tier: String,
+    pub(crate) settings: Option<crate::util::config::RemoteSettings>,
+}
+/// Fetch `/user?include=subscription` and return the parsed `UserInfo`.
+async fn fetch_user_info(
+    http_client: &reqwest::Client,
+    url: &str,
+    auth: &crate::auth::GrokAuth,
+    auth_manager: &AuthManager,
+    alpha_test_key: Option<&str>,
+) -> Result<UserInfo, &'static str> {
+    let request = http_client
+        .get(url)
+        .timeout(Duration::from_secs(10))
+        .header("Authorization", format!("Bearer {}", auth.key))
+        .header(
+            "X-XAI-Token-Auth",
+            auth_manager.grok_com_config().token_header.as_str(),
+        )
+        .header("x-grok-client-version", xai_grok_version::VERSION)
+        .header(
+            crate::http::CLIENT_MODE_HEADER,
+            crate::http::process_client_mode(),
+        );
+    let _ = alpha_test_key;
+    match request.send().await {
+        Ok(resp) if resp.status().is_success() => {
+            resp.json::<UserInfo>().await.map_err(|_| "parse")
+        }
+        Ok(_resp) => Err("http_status"),
+        Err(e) if e.is_timeout() => Err("timeout"),
+        Err(_) => Err("transport"),
+    }
+}
+/// Single-shot subscription check. Called by the pager every 5s while
+/// the paywall is shown (`x.ai/auth/check_subscription`).
+///
+/// Queries `/user?include=subscription` for the live tier. If a qualifying
+/// tier is found, does a best-effort JWT refresh + settings re-fetch and
+/// returns `Some(UnblockResult)`. Returns `None` if no qualifying
+/// subscription exists or the request fails.
+#[tracing::instrument(name = "paywall_check", skip_all, fields(user_id = %user_id))]
+pub(crate) async fn single_check(
+    auth_manager: Arc<AuthManager>,
+    proxy_base_url: &str,
+    alpha_test_key: Option<&str>,
+    user_id: &str,
+) -> Option<UnblockResult> {
+    let user_url = format!("{}/user?include=subscription", proxy_base_url);
+    let http_client = crate::http::shared_client();
+    let auth = auth_manager.current()?;
+    let user_info = match fetch_user_info(
+        &http_client,
+        &user_url,
+        &auth,
+        &auth_manager,
+        alpha_test_key,
+    )
+    .await
+    {
+        Ok(ui) => ui,
+        Err(kind) => {
+            xai_grok_telemetry::unified_log::warn(
+                "paywall_check_error",
+                None,
+                Some(serde_json::json!({ "user_id" : user_id, "kind" : kind })),
+            );
+            return None;
+        }
+    };
+    xai_grok_telemetry::unified_log::info(
+        "paywall_check_result",
+        None,
+        Some(serde_json::json!(
+            { "user_id" : user_id, "subscription_tier" : user_info.subscription_tier,
+            }
+        )),
+    );
+    let new_tier = match &user_info.subscription_tier {
+        Some(tier) if !tier.is_empty() => tier.clone(),
+        _ => return None,
+    };
+    if !QUALIFYING_TIERS.contains(&new_tier.as_str()) {
+        return None;
+    }
+    xai_grok_telemetry::unified_log::info(
+        "paywall_check_subscription_detected",
+        None,
+        Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier, })),
+    );
+    if let Err(e) = auth_manager
+        .refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
+        .await
+    {
+        xai_grok_telemetry::unified_log::warn(
+            "paywall_check_error",
+            None,
+            Some(serde_json::json!(
+                { "user_id" : user_id, "kind" : "refresh_failed", "detail" : e
+                .to_string(), }
+            )),
+        );
+    }
+    let settings = if crate::util::config::resolve_remote_fetch_enabled() {
+        let base_url = proxy_base_url.to_string();
+        let auth_for_settings = auth_manager.current().unwrap_or(auth);
+        let atk = alpha_test_key.map(str::to_string);
+        tokio::task::spawn_blocking(move || {
+            crate::remote::fetch_settings_blocking(&base_url, &auth_for_settings, atk.as_deref())
+        })
+        .await
+        .ok()
+        .flatten()
+    } else {
+        None
+    };
+    xai_grok_telemetry::unified_log::info(
+        "paywall_check_unblocked",
+        None,
+        Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier })),
+    );
+    Some(UnblockResult { new_tier, settings })
+}
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn qualifying_tiers_includes_all_paid_tiers() {
+        for tier in &[
+            "SuperGrokPro",
+            "GrokPro",
+            "SuperGrokLite",
+            "XPremiumPlus",
+            "XPremium",
+            "XBasic",
+        ] {
+            assert!(
+                QUALIFYING_TIERS.contains(tier),
+                "{tier} must be in QUALIFYING_TIERS"
+            );
+        }
+    }
+    #[test]
+    fn free_tier_is_not_qualifying() {
+        assert!(!QUALIFYING_TIERS.contains(&"Free"));
+    }
+    #[test]
+    fn empty_tier_is_not_qualifying() {
+        assert!(!QUALIFYING_TIERS.contains(&""));
+    }
+    /// The subscription check only returns `Some` when `/user` reports a
+    /// qualifying tier. Verify the tier matching is exact (no prefix match).
+    #[test]
+    fn partial_tier_name_is_not_qualifying() {
+        assert!(!QUALIFYING_TIERS.contains(&"Super"));
+        assert!(!QUALIFYING_TIERS.contains(&"Grok"));
+        assert!(!QUALIFYING_TIERS.contains(&"XPremium+"));
+    }
+}