Synced from monorepo

Synced from monorepo

Changes:
- Release a shell session's resources in one drop
- Make the tools blocking-wait cap client-configurable and self-describing
- Recognize API "exceeds budget" errors as context overflow
- Retry /btw on model overload
- Carry running background tasks and subagents across compaction
- Require round-trip time for SDK liveness checks
- Background-subagent completion reminders with a selectable delivery surface
- Make a PTY shell reap itself until it reaches the registry
- Recover the OS error code from a TLS-phase connection reset
- Consume the attached-client signal and report why idle is withheld
- Treat `.grok/sandbox.toml` edits as protected so auto mode prompts before writing
- Surface history/search in the Ctrl+. cheatsheet and keep it working in history view
- Delete sessions from the dashboard and welcome list
- Release a session's activity record when the session ends
- Stop charging auth-retry budget for fail-closed 401s; reset it across suspends
- Scope skills watches on project vendor roots
- Make [stop] cancel in-flight compaction
- Make the leader soak measure the leader, not its harness

Source-Revision: 8d69c91f02bcacf01e98d5aebbf2f92547c45738

crates/codegen/xai-grok-shell/CHANGELOG.md

@@ -1,5 +1,24 @@
 # Changelog
 
+# 0.2.117 — 2026-07-30
+
+## Features
+
+- **GROK_EXTRA_CA_BUNDLE** env var allows adding custom TLS root certificates.
+
+## Bug Fixes
+
+- **Stop command** now terminates all background subagents from prior turns.
+- **kill_task** tool now correctly reports when a task does not exist over ACP connections.
+- **get_task_output** no longer waits the full timeout for already-finished tasks over ACP.
+- **/usage** command and billing UI are hidden for enterprise auth setups.
+- **Plan approval** no longer starts Build when pressing Enter without notes in revise mode.
+
+## Performance
+
+- **Terminal resize** is much faster on long conversations in fullscreen mode.
+
+
 # 0.2.116 — 2026-07-30
 
 ## Features

crates/codegen/xai-grok-shell/Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 license = "Apache-2.0"
 name = "xai-grok-shell"
-version = "0.2.116"
+version = "0.2.117"
 edition.workspace = true
 
 [features]
@@ -12,7 +12,12 @@ dhat-heap = ["dep:dhat"]
 # load, and bench tests. Off by default; the tests/benches that use it declare
 # it via `required-features`.
 test-support = []
-default-bazel = ["test-support"]
+# Local Computer Hub workspace_server (own/attach + crash-restart). Requires
+local-workspace = []
+default-bazel = [
+    "local-workspace",
+    "test-support",
+]
 
 [dependencies]
 dunce = { workspace = true }
@@ -189,6 +194,8 @@ windows = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true }
+# Pre-main unified-log redirect (`test_support::redirect_unified_log_for_tests`).
+ctor = { workspace = true }
 filetime = { workspace = true }
 pretty_assertions = { workspace = true }
 tempfile = { workspace = true }
@@ -223,6 +230,10 @@ name = "fork_copy"
 harness = false
 required-features = ["test-support"]
 
+[[bench]]
+name = "skills_watcher_startup"
+harness = false
+
 [[test]]
 name = "test_leader_soak"
 required-features = ["test-support"]

crates/codegen/xai-grok-shell/benches/skills_watcher_startup.rs

@@ -0,0 +1,150 @@
+//! Skills file-watcher startup latency.
+//!
+//! Times OS watch registration for a project-tier `.claude` tree with a large
+//! `worktrees/` subtree (Bazel-like fan-out). Compares:
+//!
+//! - **scoped** — current `SkillsFileWatcher::start_with_dirs` (vendor root
+//!   non-recursive + skills/commands/workflows only)
+//! - **recursive_control** — full `RecursiveMode::Recursive` on `.claude`
+//!   (pre-fix project-tier behavior on Linux: one inotify wd per directory)
+//!
+//! Fixture sizes stay comparable across scenarios. Medians land under
+//! `target/criterion/skills_watcher_startup/`.
+//!
+//! ```text
+//! cargo bench -p xai-grok-shell --bench skills_watcher_startup
+//! # optional scale:
+//! GROK_SKILLS_WATCHER_BENCH_DIRS=12000 cargo bench -p xai-grok-shell --bench skills_watcher_startup
+//! ```
+//!
+//! On macOS, recursive FSEvents is cheap so both arms may be close. On Linux
+//! inotify, `recursive_control` scales with directory count; `scoped` stays flat.
+
+use std::fs;
+use std::path::{Path, PathBuf};
+use std::time::Duration;
+
+use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
+use notify::RecursiveMode;
+use notify_debouncer_mini::new_debouncer;
+use tempfile::TempDir;
+use xai_grok_shell::config::watcher::SkillsFileWatcher;
+
+/// Default dirs under `.claude/worktrees/` (override with env).
+const DEFAULT_WORKTREE_DIRS: usize = 6_000;
+
+fn worktree_dir_count() -> usize {
+    std::env::var("GROK_SKILLS_WATCHER_BENCH_DIRS")
+        .ok()
+        .and_then(|v| v.parse().ok())
+        .filter(|&n| n > 0)
+        .unwrap_or(DEFAULT_WORKTREE_DIRS)
+}
+
+/// Nested groups of 100 so the tree has width and depth.
+fn make_nested_dirs(base: &Path, count: usize) {
+    for i in 0..count {
+        let dir = base.join(format!("g{}", i / 100)).join(format!("d{i}"));
+        fs::create_dir_all(&dir).unwrap();
+    }
+}
+
+struct Fixture {
+    _root: TempDir,
+    project: PathBuf,
+    claude: PathBuf,
+    grok_home: PathBuf,
+}
+
+/// Project with a real skill and a fat `.claude/worktrees` tree.
+fn build_fixture(worktree_dirs: usize) -> Fixture {
+    let root = TempDir::new().unwrap();
+    let project = root.path().join("project");
+    let claude = project.join(".claude");
+    let skills = claude.join("skills").join("alpha");
+    fs::create_dir_all(&skills).unwrap();
+    fs::write(skills.join("SKILL.md"), "# alpha\n").unwrap();
+
+    let worktrees = claude.join("worktrees").join("wt1");
+    make_nested_dirs(&worktrees, worktree_dirs);
+
+    let grok_home = root.path().join("grok-home");
+    fs::create_dir_all(&grok_home).unwrap();
+
+    Fixture {
+        _root: root,
+        project,
+        claude,
+        grok_home,
+    }
+}
+
+fn start_scoped(fixture: &Fixture) -> SkillsFileWatcher {
+    let dirs = vec![fixture.claude.clone()];
+    let (watcher, _rx) = SkillsFileWatcher::start_with_dirs(
+        &dirs,
+        &fixture.grok_home,
+        Some(fixture.project.as_path()),
+    )
+    .expect("scoped skills watcher should start");
+    watcher
+}
+
+/// Pre-fix control: one recursive watch on the whole project `.claude`.
+fn start_recursive_control(
+    claude: &Path,
+) -> notify_debouncer_mini::Debouncer<notify::RecommendedWatcher> {
+    let mut debouncer = new_debouncer(Duration::from_secs(2), |_| {}).expect("debouncer");
+    debouncer
+        .watcher()
+        .watch(claude, RecursiveMode::Recursive)
+        .expect("recursive watch");
+    debouncer
+}
+
+fn bench_skills_watcher_startup(c: &mut Criterion) {
+    let n = worktree_dir_count();
+    let fixture = build_fixture(n);
+
+    eprintln!(
+        "skills_watcher_startup fixture: project={:?} worktree_dirs={n}",
+        fixture.project
+    );
+
+    let mut group = c.benchmark_group("skills_watcher_startup");
+    group.sample_size(20);
+    group.throughput(Throughput::Elements(n as u64));
+    group.warm_up_time(Duration::from_secs(1));
+    group.measurement_time(Duration::from_secs(8));
+
+    group.bench_function(BenchmarkId::new("scoped", n), |b| {
+        b.iter_batched(|| (), |()| start_scoped(&fixture), BatchSize::PerIteration);
+    });
+
+    group.bench_function(BenchmarkId::new("recursive_control", n), |b| {
+        b.iter_batched(
+            || (),
+            |()| start_recursive_control(&fixture.claude),
+            BatchSize::PerIteration,
+        );
+    });
+
+    // Tiny tree: both arms should be similar (fixed overhead check).
+    let tiny = build_fixture(0);
+    group.throughput(Throughput::Elements(1));
+    group.bench_function(BenchmarkId::new("scoped_tiny", 0), |b| {
+        b.iter_batched(|| (), |()| start_scoped(&tiny), BatchSize::PerIteration);
+    });
+    group.bench_function(BenchmarkId::new("recursive_control_tiny", 0), |b| {
+        b.iter_batched(
+            || (),
+            |()| start_recursive_control(&tiny.claude),
+            BatchSize::PerIteration,
+        );
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, bench_skills_watcher_startup);
+criterion_main!(benches);

crates/codegen/xai-grok-shell/changelogs/0.2.115.md

@@ -1,5 +1,9 @@
 # 0.2.115 — 2026-07-29
 
+## Features
+
+- **Delete sessions from the dashboard and welcome list.** On the dashboard, press `Ctrl+X` twice (or hover a settled row and click `[✗]` twice); in the welcome and `/resume` lists, press `d` then `y`.
+
 ## Bug Fixes
 
 - **Fixed chat history corruption** that could duplicate tool results or cause later 400 errors after repeated identical tool calls.

crates/codegen/xai-grok-shell/changelogs/0.2.117.json

@@ -0,0 +1,37 @@
+[
+  {
+    "category": "fixes",
+    "description": "**Stop command** now terminates all background subagents from prior turns.",
+    "breaking_change": false
+  },
+  {
+    "category": "features",
+    "description": "**GROK_EXTRA_CA_BUNDLE** env var allows adding custom TLS root certificates.",
+    "breaking_change": false
+  },
+  {
+    "category": "fixes",
+    "description": "**kill_task** tool now correctly reports when a task does not exist over ACP connections.",
+    "breaking_change": false
+  },
+  {
+    "category": "fixes",
+    "description": "**get_task_output** no longer waits the full timeout for already-finished tasks over ACP.",
+    "breaking_change": false
+  },
+  {
+    "category": "performance",
+    "description": "**Terminal resize** is much faster on long conversations in fullscreen mode.",
+    "breaking_change": false
+  },
+  {
+    "category": "fixes",
+    "description": "**/usage** command and billing UI are hidden for enterprise auth setups.",
+    "breaking_change": false
+  },
+  {
+    "category": "fixes",
+    "description": "**Plan approval** no longer starts Build when pressing Enter without notes in revise mode.",
+    "breaking_change": false
+  }
+]

crates/codegen/xai-grok-shell/changelogs/0.2.117.md

@@ -0,0 +1,18 @@
+# 0.2.117 — 2026-07-30
+
+## Features
+
+- **GROK_EXTRA_CA_BUNDLE** env var allows adding custom TLS root certificates.
+
+## Bug Fixes
+
+- **Stop command** now terminates all background subagents from prior turns.
+- **kill_task** tool now correctly reports when a task does not exist over ACP connections.
+- **get_task_output** no longer waits the full timeout for already-finished tasks over ACP.
+- **/usage** command and billing UI are hidden for enterprise auth setups.
+- **Plan approval** no longer starts Build when pressing Enter without notes in revise mode.
+
+## Performance
+
+- **Terminal resize** is much faster on long conversations in fullscreen mode.
+

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

@@ -267,8 +267,8 @@ async fn handle_session_list(
 ) -> 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.
+    // Under chat mode `parse_list_req` force-rewrites `kind` to conversations
+    // unless `local-workspace` is compiled in and the client sent chat/build.
     let req = unified_list::parse_list_req(args.params.get())
         .map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
     tracing::debug!(

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

@@ -1020,6 +1020,25 @@ impl acp::Agent for MvpAgent {
             .as_ref()
             .and_then(|m| m.get("modelId").and_then(|v| v.as_str()))
             .filter(|s| !s.is_empty());
+        #[cfg(all(feature = "local-workspace", unix))]
+        let pending_local_workspace = self
+            .start_own_local_workspace_if_needed(
+                &mut session_meta_for_stamp,
+                cwd.as_path(),
+            )
+            .await?;
+        #[cfg(all(feature = "local-workspace", not(unix)))]
+        {
+            use crate::gateway_bridge::local_workspace_supervisor::parse_local_workspace_intent;
+            use crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceIntent;
+            use crate::gateway_bridge::local_workspace_supervisor::SupervisorError;
+            if matches!(
+                parse_local_workspace_intent(session_meta_for_stamp.as_ref()),
+                Some(LocalWorkspaceIntent::Own { .. })
+            ) {
+                return Err(SupervisorError::UnsupportedPlatform.into_acp_error());
+            }
+        }
         #[allow(unused_variables)]
         let session_computer_sessions = resolve_session_computer_sessions(
             arguments.meta.as_ref(),
@@ -1052,6 +1071,15 @@ impl acp::Agent for MvpAgent {
             }
             None => acp::SessionId::new(uuid::Uuid::now_v7().to_string()),
         };
+        #[cfg(all(feature = "local-workspace", unix))]
+        let mut local_ws_reap_guard = self
+            .new_local_workspace_reap_guard(session_id.clone(), false);
+        #[cfg(all(feature = "local-workspace", unix))]
+        if let Some(handle) = pending_local_workspace {
+            self.register_local_workspace_supervisor(session_id.clone(), handle);
+            local_ws_reap_guard = self
+                .new_local_workspace_reap_guard(session_id.clone(), true);
+        }
         let mut session_timer = crate::instrumentation_timer!("session.new_session");
         session_timer.with_field("session_id", session_id.0.as_ref());
         session_timer.with_field("cwd", cwd.as_str());
@@ -1273,8 +1301,16 @@ impl acp::Agent for MvpAgent {
             };
             self.spawn_and_register_session(init, spawn_opts).await
         };
+        #[cfg(all(feature = "local-workspace", unix))]
+        if spawn_res.is_err() {
+            self.shutdown_gateway_bridge(&session_id);
+        }
         spawn_res?;
         tracing::debug!(session_id = %session_id.0, "new_session: spawn_session_actor");
+        #[cfg(feature = "local-workspace")]
+        if local_workspace_intent_present(arguments.meta.as_ref()) {
+            self.mark_local_workspace_bound(session_id.clone());
+        }
         self.maybe_spawn_interactive_trust_prompt(
             &session_id,
             cwd.as_path(),
@@ -1409,6 +1445,7 @@ impl acp::Agent for MvpAgent {
             );
             insert_applied_tool_overrides(obj, applied_tool_overrides.as_ref());
         }
+        #[cfg(all(feature = "local-workspace", unix))] local_ws_reap_guard.disarm();
         Ok(
             acp::NewSessionResponse::new(session_id)
                 .models(Some(models))
@@ -1947,7 +1984,7 @@ impl acp::Agent for MvpAgent {
         let persisted_model = summary.current_model_id.clone();
         let models = self.models_manager.models();
         let available = self.models_manager.available();
-        self.model_unavailable_sessions.borrow_mut().remove(session_id.0.as_ref());
+        self.session_registry.take_unavailable_model(&session_id);
         let resolved_catalog_key = resolve_catalog_key(&models, &persisted_model);
         tracing::debug!(
             session_id = %session_id.0,
@@ -2061,9 +2098,8 @@ impl acp::Agent for MvpAgent {
                     &reason,
                 )
                 .await;
-            self.model_unavailable_sessions
-                .borrow_mut()
-                .insert(session_id.0.to_string(), persisted_model.clone());
+            self.session_registry
+                .set_unavailable_model(&session_id, persisted_model.clone());
             fallback
         };
         tracing::debug!(
@@ -2266,10 +2302,8 @@ impl acp::Agent for MvpAgent {
             return Ok(acp::PromptResponse::new(acp::StopReason::EndTurn));
         }
         let latched_model = self
-            .model_unavailable_sessions
-            .borrow()
-            .get(arguments.session_id.0.as_ref())
-            .cloned();
+            .session_registry
+            .unavailable_model(&arguments.session_id);
         if let Some(unavailable_model) = latched_model {
             let models = self.models_manager.models();
             let available = self.models_manager.available();
@@ -2294,9 +2328,7 @@ impl acp::Agent for MvpAgent {
                     }),
                     ),
                 );
-                self.model_unavailable_sessions
-                    .borrow_mut()
-                    .remove(arguments.session_id.0.as_ref());
+                self.session_registry.take_unavailable_model(&arguments.session_id);
                 if let Err(e) = crate::agent::handlers::model_switch::apply(
                         self,
                         acp::SetSessionModelRequest::new(
@@ -3426,9 +3458,8 @@ impl acp::Agent for MvpAgent {
         let res = crate::agent::handlers::model_switch::apply(self, args).await;
         if res.is_ok()
             && let Some(unavailable) = self
-                .model_unavailable_sessions
-                .borrow_mut()
-                .remove(session_id.0.as_ref())
+                .session_registry
+                .take_unavailable_model(&session_id)
         {
             tracing::info!(
                 session_id = %session_id.0,
@@ -3488,6 +3519,10 @@ impl acp::Agent for MvpAgent {
                 let ops = self.resolve_workspace_ops()?;
                 crate::extensions::worktree::handle(self, &ops, &args).await
             }
+            #[cfg(feature = "local-workspace")]
+            "x.ai/session/add_local_workspace" => {
+                crate::extensions::session_admin::handle(self, &args).await
+            }
             "x.ai/session/rename" | "x.ai/session/delete"
             | "x.ai/session/update_mcp_servers" | "x.ai/session/fork"
             | "x.ai/internal/reload_all_mcp_servers"

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

@@ -739,7 +739,7 @@ impl MvpAgent {
     /// Most recently allocated turn number for `sid`, or `None` if the
     /// session has not started a turn yet.
     pub(crate) fn session_turn_number(&self, sid: &acp::SessionId) -> Option<u64> {
-        self.retained_resources.borrow().get(sid).and_then(|d| d.turn_number)
+        self.session_registry.turn_number(sid)
     }
     /// Return the current GrokAuth credentials, if authenticated and not expired.
     pub(crate) fn current_auth(&self) -> Option<crate::auth::GrokAuth> {
@@ -763,6 +763,596 @@ impl MvpAgent {
     pub(crate) fn alpha_test_key(&self) -> Option<String> {
         self.cfg.borrow().endpoints.alpha_test_key.clone()
     }
+    #[cfg(all(feature = "local-workspace", unix))]
+    /// Spawn owned `workspace_server` for chat+local `own` intent.
+    /// Mints `server_id` into `_meta` before handshake parse.
+    pub(crate) async fn start_own_local_workspace_if_needed(
+        &self,
+        meta: &mut Option<acp::Meta>,
+        session_cwd: &std::path::Path,
+    ) -> Result<
+        Option<crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle>,
+        acp::Error,
+    > {
+        use crate::gateway_bridge::local_workspace_supervisor::{
+            parse_local_workspace_intent, stamp_server_id_into_meta, start_own,
+            StartOwnConfig, LocalWorkspaceIntent,
+        };
+        let Some(LocalWorkspaceIntent::Own { cwd }) = parse_local_workspace_intent(
+            meta.as_ref(),
+        ) else {
+            return Ok(None);
+        };
+        let cwd = if cwd.as_os_str().is_empty() {
+            session_cwd.to_path_buf()
+        } else {
+            cwd
+        };
+        crate::gateway_bridge::local_workspace_supervisor::validate_cwd(&cwd)
+            .map_err(|e| e.into_acp_error())?;
+        let hub_url = {
+            let cfg = self.cfg.borrow();
+            crate::gateway_bridge::local_workspace_supervisor::resolve_hub_url(
+                cfg.hub.url.as_deref(),
+            )
+        };
+        let handle = start_own(StartOwnConfig {
+                cwd,
+                hub_url,
+                auth_config: None,
+                binary: None,
+                ready_timeout: crate::gateway_bridge::local_workspace_supervisor::READY_TIMEOUT,
+                allow_missing_auth: false,
+            })
+            .await
+            .map_err(|e| e.into_acp_error())?;
+        let meta_map = meta.get_or_insert_with(acp::Meta::new);
+        stamp_server_id_into_meta(meta_map, &handle.server_id);
+        Ok(Some(handle))
+    }
+    #[cfg(all(feature = "local-workspace", unix))]
+    pub(crate) fn register_local_workspace_supervisor(
+        &self,
+        session_id: acp::SessionId,
+        handle: crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle,
+    ) {
+        let server_id = handle.server_id.clone();
+        self.arm_local_workspace_watcher(session_id.clone(), handle);
+        tracing::info!(
+            session_id = %session_id.0,
+            server_id = %server_id,
+            "local_workspace_supervisor: registered own workspace_server"
+        );
+    }
+    #[cfg(all(feature = "local-workspace", unix))]
+    pub(crate) fn new_local_workspace_reap_guard(
+        &self,
+        session_id: acp::SessionId,
+        armed: bool,
+    ) -> LocalWorkspaceReapGuard {
+        LocalWorkspaceReapGuard {
+            supervisors: self.local_workspace_supervisors.clone(),
+            generations: self.local_workspace_generations.clone(),
+            session_id,
+            armed,
+        }
+    }
+    /// Prefer live supervisor `server_id` over the parse-time stamp (pre-bridge crash).
+    #[cfg(all(feature = "local-workspace", unix))]
+    pub(crate) fn refresh_sessions_from_supervisor(
+        &self,
+        session_id: &acp::SessionId,
+        sessions: Option<Vec<crate::gateway_bridge::ComputerSession>>,
+    ) -> Option<Vec<crate::gateway_bridge::ComputerSession>> {
+        use crate::gateway_bridge::ComputerSession;
+        let supervisors = self.local_workspace_supervisors.borrow();
+        let Some(handle) = supervisors.get(session_id) else {
+            return sessions;
+        };
+        let server_id = handle.server_id.clone();
+        let cwd = Some(handle.cwd.to_string_lossy().into_owned());
+        match sessions {
+            None => Some(vec![ComputerSession::ExistingWorkspace { server_id, cwd }]),
+            Some(mut list) => {
+                for session in &mut list {
+                    if let ComputerSession::ExistingWorkspace {
+                        server_id: sid,
+                        cwd: existing_cwd,
+                    } = session {
+                        *sid = server_id.clone();
+                        if existing_cwd.is_none() {
+                            *existing_cwd = cwd.clone();
+                        }
+                    }
+                }
+                Some(list)
+            }
+        }
+    }
+    /// Wait out an in-flight crash restart before refreshing handshake sessions.
+    #[cfg(all(feature = "local-workspace", unix))]
+    pub(crate) async fn await_refresh_sessions_from_supervisor(
+        &self,
+        session_id: &acp::SessionId,
+        sessions: Option<Vec<crate::gateway_bridge::ComputerSession>>,
+    ) -> Option<Vec<crate::gateway_bridge::ComputerSession>> {
+        const WAIT: std::time::Duration = std::time::Duration::from_secs(5);
+        let deadline = tokio::time::Instant::now() + WAIT;
+        loop {
+            let pending = self
+                .local_workspace_restart_pending
+                .borrow()
+                .contains(session_id);
+            let live = self
+                .local_workspace_supervisors
+                .borrow()
+                .contains_key(session_id);
+            if live || !pending {
+                break;
+            }
+            if tokio::time::Instant::now() >= deadline {
+                tracing::warn!(
+                    session_id = %session_id.0,
+                    "timed out waiting for local-workspace crash restart before handshake refresh"
+                );
+                break;
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(25)).await;
+        }
+        self.refresh_sessions_from_supervisor(session_id, sessions)
+    }
+    #[cfg(all(feature = "local-workspace", unix))]
+    fn arm_local_workspace_watcher(
+        &self,
+        session_id: acp::SessionId,
+        handle: crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle,
+    ) {
+        let cwd = handle.cwd.clone();
+        let sessions = self.session_registry.clone();
+        let supervisors = self.local_workspace_supervisors.clone();
+        let generations = self.local_workspace_generations.clone();
+        let sid = session_id.clone();
+        let hub_url = {
+            let cfg = self.cfg.borrow();
+            crate::gateway_bridge::local_workspace_supervisor::resolve_hub_url(
+                cfg.hub.url.as_deref(),
+            )
+        };
+        let auth_path = xai_grok_workspace::hub_auth::default_auth_path().ok();
+        let binary = crate::gateway_bridge::local_workspace_supervisor::resolve_workspace_server_bin()
+            .ok();
+        let agent_ref = LocalRef::new(self);
+        let generation = {
+            let mut gens = generations.borrow_mut();
+            let e = gens.entry(session_id.clone()).or_insert(0);
+            *e = e.saturating_add(1);
+            *e
+        };
+        self.local_workspace_supervisors.borrow_mut().insert(session_id.clone(), handle);
+        let mut supervisors_mut = self.local_workspace_supervisors.borrow_mut();
+        let Some(handle_mut) = supervisors_mut.get_mut(&session_id) else {
+            return;
+        };
+        let _ = handle_mut
+            .spawn_exit_watcher(move || {
+                let sessions = sessions.clone();
+                let supervisors = supervisors.clone();
+                let generations = generations.clone();
+                let sid = sid.clone();
+                let hub_url = hub_url.clone();
+                let auth_path = auth_path.clone();
+                let binary = binary.clone();
+                let cwd = cwd.clone();
+                let agent_ref = agent_ref.clone();
+                tokio::task::spawn_local(async move {
+                    if generations.borrow().get(&sid) != Some(&generation) {
+                        return;
+                    }
+                    agent_ref
+                        .get()
+                        .local_workspace_restart_pending
+                        .borrow_mut()
+                        .insert(sid.clone());
+                    let prev = supervisors.borrow_mut().remove(&sid);
+                    let Some(prev) = prev else {
+                        agent_ref
+                            .get()
+                            .local_workspace_restart_pending
+                            .borrow_mut()
+                            .remove(&sid);
+                        return;
+                    };
+                    let Some(binary) = binary else {
+                        tracing::warn!(
+                        session_id = %sid.0,
+                        "local workspace crash restart skipped: binary missing"
+                    );
+                        prev.shutdown().await;
+                        agent_ref
+                            .get()
+                            .local_workspace_restart_pending
+                            .borrow_mut()
+                            .remove(&sid);
+                        return;
+                    };
+                    let auth = auth_path
+                        .unwrap_or_else(|| std::path::PathBuf::from("/nonexistent"));
+                    let restart_count = prev.restart_count;
+                    let prev_cwd = prev.cwd.clone();
+                    prev.shutdown().await;
+                    match crate::gateway_bridge::local_workspace_supervisor::restart_own_from(
+                            restart_count,
+                            prev_cwd,
+                            &binary,
+                            &hub_url,
+                            &auth,
+                            false,
+                        )
+                        .await
+                    {
+                        Ok(new_handle) => {
+                            if generations.borrow().get(&sid) != Some(&generation) {
+                                new_handle.shutdown().await;
+                                agent_ref
+                                    .get()
+                                    .local_workspace_restart_pending
+                                    .borrow_mut()
+                                    .remove(&sid);
+                                return;
+                            }
+                            let new_id = new_handle.server_id.clone();
+                            let cwd_str = cwd.to_string_lossy().into_owned();
+                            agent_ref
+                                .get()
+                                .arm_local_workspace_watcher(sid.clone(), new_handle);
+                            agent_ref
+                                .get()
+                                .local_workspace_restart_pending
+                                .borrow_mut()
+                                .remove(&sid);
+                            let armed_generation = generations
+                                .borrow()
+                                .get(&sid)
+                                .copied();
+                            let bridge = sessions.bridge(&sid);
+                            if let Some(bridge) = bridge {
+                                let _ = crate::gateway_bridge::local_workspace_supervisor::push_computer_sessions_update(
+                                        &bridge,
+                                        new_id.clone(),
+                                        Some(cwd_str.clone()),
+                                        false,
+                                    )
+                                    .await;
+                                let _ = bridge.wait_until_ready().await;
+                                if generations.borrow().get(&sid)
+                                    != armed_generation.as_ref()
+                                {
+                                    tracing::debug!(
+                                    session_id = %sid.0,
+                                    "skip stale local-workspace session.update after superseded restart"
+                                );
+                                    return;
+                                }
+                                let _ = crate::gateway_bridge::local_workspace_supervisor::push_computer_sessions_update(
+                                        &bridge,
+                                        new_id,
+                                        Some(cwd_str),
+                                        false,
+                                    )
+                                    .await;
+                            }
+                        }
+                        Err(err) => {
+                            tracing::warn!(
+                            session_id = %sid.0,
+                            error = %err,
+                            "local workspace crash restart failed"
+                        );
+                            agent_ref
+                                .get()
+                                .local_workspace_restart_pending
+                                .borrow_mut()
+                                .remove(&sid);
+                        }
+                    }
+                });
+            });
+    }
+    /// add-only mid-session local workspace via ACP extension / session.update.
+    ///
+    /// Refuses if a local existing workspace is already bound (no remove until session end).
+    /// Own mode requires unix (supervisor spawn). Attach is platform-agnostic.
+    #[cfg(feature = "local-workspace")]
+    pub(crate) async fn add_local_workspace_mid_session(
+        &self,
+        session_id: &acp::SessionId,
+        mut meta: Option<acp::Meta>,
+        session_cwd: &std::path::Path,
+    ) -> Result<serde_json::Value, acp::Error> {
+        use crate::gateway_bridge::ComputerSession;
+        use crate::gateway_bridge::local_workspace_supervisor::{
+            parse_local_workspace_intent, LocalWorkspaceIntent, SupervisorError,
+        };
+        if self.local_workspace_already_bound(session_id) {
+            return Err(
+                acp::Error::invalid_params()
+                    .data(
+                        serde_json::json!({
+                "code": "local_workspace_already_bound",
+                "message": "local workspace already bound; remove is not supported until session end",
+            }),
+                    ),
+            );
+        }
+        self.mark_local_workspace_bound(session_id.clone());
+        let mut bind_guard = LocalWorkspaceBindGuard {
+            bound: self.local_workspace_bound.clone(),
+            session_id: session_id.clone(),
+            keep: false,
+        };
+        let Some(intent) = parse_local_workspace_intent(meta.as_ref()) else {
+            return Err(
+                acp::Error::invalid_params()
+                    .data(
+                        serde_json::json!({
+                "code": "local_workspace_intent_missing",
+                "message": "x.ai/local_workspace intent required for mid-session add",
+            }),
+                    ),
+            );
+        };
+        let mode = match &intent {
+            LocalWorkspaceIntent::Own { .. } => "own",
+            LocalWorkspaceIntent::Attach { .. } => "attach",
+        };
+        #[cfg_attr(not(unix), allow(unused_variables))]
+        let pending: Option<
+            crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle,
+        > = match intent {
+            LocalWorkspaceIntent::Own { .. } => {
+                #[cfg(unix)]
+                {
+                    self.start_own_local_workspace_if_needed(&mut meta, session_cwd)
+                        .await?
+                }
+                #[cfg(not(unix))]
+                {
+                    let _ = session_cwd;
+                    return Err(SupervisorError::UnsupportedPlatform.into_acp_error());
+                }
+            }
+            LocalWorkspaceIntent::Attach { cwd, .. } => {
+                if let Some(ref cwd) = cwd {
+                    crate::gateway_bridge::local_workspace_supervisor::validate_cwd(cwd)
+                        .map_err(|e| e.into_acp_error())?;
+                }
+                Self::ensure_attach_fs_only_advertised_tools()
+                    .map_err(|msg| {
+                        acp::Error::invalid_params()
+                            .data(
+                                serde_json::json!({
+                        "code": "local_workspace_fs_only_required",
+                        "message": msg,
+                    }),
+                            )
+                    })?;
+                None
+            }
+        };
+        let sessions = resolve_session_computer_sessions(meta.as_ref())?;
+        let Some(sessions) = sessions.filter(|s| !s.is_empty()) else {
+            return Err(
+                acp::Error::invalid_params()
+                    .data(
+                        serde_json::json!({
+                "code": "local_workspace_stamp_failed",
+                "message": "failed to resolve existing_workspace stamp for mid-session add",
+            }),
+                    ),
+            );
+        };
+        if !sessions
+            .iter()
+            .any(|s| matches!(s, ComputerSession::ExistingWorkspace { .. }))
+        {
+            return Err(
+                acp::Error::invalid_params()
+                    .data(
+                        serde_json::json!({
+                "code": "local_workspace_stamp_failed",
+                "message": "mid-session add did not produce existing_workspace",
+            }),
+                    ),
+            );
+        }
+        #[cfg(unix)]
+        let mut reap_guard = self
+            .new_local_workspace_reap_guard(session_id.clone(), false);
+        #[cfg(unix)]
+        if let Some(handle) = pending {
+            self.register_local_workspace_supervisor(session_id.clone(), handle);
+            reap_guard = self.new_local_workspace_reap_guard(session_id.clone(), true);
+        }
+        let Some(bridge) = self.gateway_bridge_for(session_id) else {
+            return Err(
+                acp::Error::invalid_params()
+                    .data(
+                        serde_json::json!({
+                "code": "gateway_bridge_missing",
+                "message": "session has no gateway bridge for session.update computer_sessions",
+            }),
+                    ),
+            );
+        };
+        match tokio::time::timeout(BRIDGE_READY_TIMEOUT, bridge.wait_until_ready()).await
+        {
+            Ok(Ok(())) => {}
+            Ok(Err(err)) => {
+                return Err(err.into_acp_error());
+            }
+            Err(_) => {
+                return Err(
+                    acp::Error::internal_error()
+                        .data(
+                            "gateway bridge not ready for mid-session add_local_workspace",
+                        ),
+                );
+            }
+        }
+        let sessions = sessions;
+        let server_id = match sessions.first() {
+            Some(ComputerSession::ExistingWorkspace { server_id, .. }) => {
+                server_id.clone()
+            }
+            _ => {
+                return Err(
+                    acp::Error::internal_error()
+                        .data("expected existing_workspace as first computer session"),
+                );
+            }
+        };
+        let cwd = match sessions.first() {
+            Some(ComputerSession::ExistingWorkspace { cwd, .. }) => cwd.clone(),
+            _ => None,
+        };
+        if let Some(ref stamped_cwd) = cwd {
+            crate::gateway_bridge::local_workspace_supervisor::validate_cwd(
+                    std::path::Path::new(stamped_cwd),
+                )
+                .map_err(|e| e.into_acp_error())?;
+        }
+        if let Err(err) = crate::gateway_bridge::local_workspace_supervisor::push_computer_sessions_update(
+                &bridge,
+                server_id.clone(),
+                cwd,
+                true,
+            )
+            .await
+        {
+            return Err(err.into_acp_error());
+        }
+        #[cfg(unix)] reap_guard.disarm();
+        bind_guard.keep = true;
+        Ok(serde_json::json!({
+            "ok": true,
+            "server_id": server_id,
+            "mode": mode,
+        }))
+    }
+    #[cfg(feature = "local-workspace")]
+    pub(crate) fn local_workspace_already_bound(
+        &self,
+        session_id: &acp::SessionId,
+    ) -> bool {
+        if self.local_workspace_bound.borrow().contains(session_id) {
+            return true;
+        }
+        #[cfg(unix)]
+        if self.local_workspace_supervisors.borrow().contains_key(session_id) {
+            return true;
+        }
+        false
+    }
+    #[cfg(feature = "local-workspace")]
+    pub(crate) fn mark_local_workspace_bound(&self, session_id: acp::SessionId) {
+        self.local_workspace_bound.borrow_mut().insert(session_id);
+    }
+    /// Operator-attested FS-only toolset for mid-session attach.
+    #[cfg(feature = "local-workspace")]
+    fn ensure_attach_fs_only_advertised_tools() -> Result<(), String> {
+        const ENV: &str = "GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS";
+        const ALLOW: &[&str] = &[
+            "workspace.fs_list",
+            "workspace.fs_exists",
+            "workspace.fs_read_file",
+            "workspace.fs_write_file",
+            "workspace.fs_delete_file",
+            "workspace.put_files",
+            "workspace.get_files",
+        ];
+        let Some(raw) = std::env::var(ENV)
+            .ok()
+            .map(|s| s.trim().to_string())
+            .filter(|s| !s.is_empty()) else {
+            return Err(
+                "attached workspace_server advertised toolset is uncheckable; refuse attach \
+                 (set GROK_CHAT_LOCAL_WORKSPACE_ADVERTISED_TOOLS to a comma-separated FS-only catalog)"
+                    .into(),
+            );
+        };
+        let ids: Vec<&str> = raw
+            .split(',')
+            .map(str::trim)
+            .filter(|s| !s.is_empty())
+            .collect();
+        if ids.is_empty() {
+            return Err(
+                "attached workspace_server advertised an empty toolset; refuse attach"
+                    .into(),
+            );
+        }
+        let forbidden: Vec<&str> = ids
+            .into_iter()
+            .filter(|id| !ALLOW.contains(id))
+            .collect();
+        if forbidden.is_empty() {
+            Ok(())
+        } else {
+            Err(
+                    format!(
+                "attached workspace_server advertises tools outside the FS-only allowlist: {}",
+                forbidden.join(", ")
+            ),
+                )
+        }
+    }
+    #[cfg(feature = "local-workspace")]
+    /// After chat+local stamp, wait for handshake success.
+    ///
+    /// Only fail-closed for `x.ai/local_workspace` intent (not generic
+    /// GatewayAttach). Handshake errors propagate; session + bridge are reaped
+    /// on failure / timeout.
+    pub(crate) async fn await_existing_workspace_handshake(
+        &self,
+        session_id: &acp::SessionId,
+        local_workspace_intent: bool,
+    ) -> Result<(), acp::Error> {
+        if !local_workspace_intent {
+            return Ok(());
+        }
+        #[cfg(feature = "local-workspace")]
+        self.mark_local_workspace_bound(session_id.clone());
+        let Some(bridge) = self.gateway_bridge_for(session_id) else {
+            return Ok(());
+        };
+        match tokio::time::timeout(BRIDGE_READY_TIMEOUT, bridge.wait_until_ready()).await
+        {
+            Ok(Ok(())) => Ok(()),
+            Ok(Err(err)) => {
+                tracing::warn!(
+                    session_id = %session_id.0,
+                    error = %err,
+                    kind = "existing_workspace_handshake_failed",
+                    "chat+local handshake failed; reaping session"
+                );
+                self.request_session_shutdown(session_id);
+                self.remove_session(session_id);
+                Err(err.into_acp_error())
+            }
+            Err(_) => {
+                tracing::warn!(
+                    session_id = %session_id.0,
+                    kind = "existing_workspace_handshake_timeout",
+                    "chat+local handshake timed out; reaping session"
+                );
+                self.request_session_shutdown(session_id);
+                self.remove_session(session_id);
+                Err(
+                    acp::Error::internal_error().data("gateway bridge connect timed out"),
+                )
+            }
+        }
+    }
     /// Build the process-lifetime local `WorkspaceOps` on first use.
     ///
     /// Deferred past ACP wiring so `initialize` can respond before folder-trust
@@ -1896,9 +2486,8 @@ impl MvpAgent {
         let instance = Self {
             sessions: RefCell::new(HashMap::new()),
             activity,
+            session_registry: SessionRegistry::default(),
             loading_sessions: RefCell::new(HashMap::new()),
-            retained_resources: RefCell::new(HashMap::new()),
-            session_threads: RefCell::new(HashMap::new()),
             resident_roster_titles: RefCell::new(HashMap::new()),
             initialize_request: OnceLock::new(),
             gateway,
@@ -1948,7 +2537,6 @@ impl MvpAgent {
             codebase_indexes: Arc::new(
                 parking_lot::Mutex::new(CodebaseIndexManager::new()),
             ),
-            resident_resources: RefCell::new(HashMap::new()),
             worktree_type,
             restore_code,
             session_registry_local,
@@ -1958,7 +2546,6 @@ impl MvpAgent {
                     crate::session::mcp_servers::McpState::new(vec![]),
                 ),
             ),
-            model_unavailable_sessions: RefCell::new(std::collections::HashMap::new()),
             subagent_event_tx,
             subagent_event_rx: RefCell::new(Some(subagent_event_rx)),
             subagent_presentation: RefCell::new(
@@ -1970,7 +2557,18 @@ impl MvpAgent {
                 std::sync::atomic::AtomicBool::new(false),
             ),
             workspace_ops: RefCell::new(None),
-            session_live_state: RefCell::new(HashMap::new()),
+            #[cfg(all(feature = "local-workspace", unix))]
+            local_workspace_supervisors: Rc::new(RefCell::new(HashMap::new())),
+            #[cfg(all(feature = "local-workspace", unix))]
+            local_workspace_generations: Rc::new(RefCell::new(HashMap::new())),
+            #[cfg(all(feature = "local-workspace", unix))]
+            local_workspace_restart_pending: Rc::new(
+                RefCell::new(std::collections::HashSet::new()),
+            ),
+            #[cfg(feature = "local-workspace")]
+            local_workspace_bound: Rc::new(
+                RefCell::new(std::collections::HashSet::new()),
+            ),
             supervisor_started: std::cell::Cell::new(false),
             settings_reapply_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)),
             post_auth_settings_in_flight: std::rc::Rc::new(std::cell::Cell::new(false)),
@@ -2074,7 +2672,7 @@ impl MvpAgent {
             }
             self.request_session_shutdown(&id);
             if self.take_session(&id).is_some() {
-                self.resident_resources.borrow_mut().remove(&id);
+                self.session_registry.clear_resident(&id);
                 self.set_session_live_state(&id, SessionLiveState::Dormant);
                 unloaded += 1;
                 tracing::debug!(session_id = %id.0, "idle session unloaded to disk on disconnect");
@@ -2095,10 +2693,13 @@ impl MvpAgent {
     /// Uses async polling (never blocks the `LocalSet` runtime) with a 5s deadline
     /// to handle slow shutdowns (e.g., embedding API timeouts).
     pub(super) async fn drain_old_session_thread(&self, session_id: &acp::SessionId) {
-        let thread = self.session_threads.borrow_mut().remove(session_id);
-        let Some(thread) = thread else { return };
-        if thread.is_finished() {
-            return;
+        match self.session_registry.thread_is_finished(session_id) {
+            None => return,
+            Some(true) => {
+                self.session_registry.clear_thread(session_id);
+                return;
+            }
+            Some(false) => {}
         }
         tracing::info!(
             session_id = %session_id.0,
@@ -2106,12 +2707,17 @@ impl MvpAgent {
         );
         let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
         loop {
-            if thread.is_finished() {
-                tracing::debug!(
-                    session_id = %session_id.0,
-                    "Old session thread finished cleanly"
-                );
-                return;
+            match self.session_registry.thread_is_finished(session_id) {
+                None => return,
+                Some(true) => {
+                    self.session_registry.clear_thread(session_id);
+                    tracing::debug!(
+                        session_id = %session_id.0,
+                        "Old session thread finished cleanly"
+                    );
+                    return;
+                }
+                Some(false) => {}
             }
             if tokio::time::Instant::now() >= deadline {
                 tracing::warn!(
@@ -2904,12 +3510,7 @@ impl MvpAgent {
     }
     /// Set a session's next trace turn number.
     pub(super) fn set_turn_number(&self, session_id: &acp::SessionId, next: u64) {
-        self
-            .retained_resources
-            .borrow_mut()
-            .entry(session_id.clone())
-            .or_default()
-            .turn_number = Some(next);
+        self.session_registry.set_turn_number(session_id, next);
     }
     /// Upload each drained harness trace turn as its own `turn_{N}` artifact,
     /// numbered from the same counter as model turns so subagents interleave
@@ -4102,9 +4703,7 @@ impl MvpAgent {
                 )
                 .await?
         };
-        self.session_threads
-            .borrow_mut()
-            .insert(session_info.id.clone(), session_thread);
+        self.session_registry.set_thread(&session_info.id, session_thread);
         tracing::debug!(session_id = %session_info.id.0, "spawn_session_on_thread complete");
         self.set_session_live_state(&session_info.id, SessionLiveState::IdleResident);
         self.ensure_session_supervisor();
@@ -4164,12 +4763,8 @@ impl MvpAgent {
                 }
             });
         }
-        self
-            .retained_resources
-            .borrow_mut()
-            .entry(session_info.id.clone())
-            .or_default()
-            .permission_event_receiver = Some(permission_events_rx);
+        self.session_registry
+            .set_permission_receiver(&session_info.id, permission_events_rx);
         if handle_display_cwd.is_some() {
             handle.display_cwd = handle_display_cwd;
         }
@@ -4203,16 +4798,56 @@ impl MvpAgent {
         &self,
         session_id: &acp::SessionId,
     ) -> Vec<PermissionEvent> {
-        let mut events = Vec::new();
-        let mut retained = self.retained_resources.borrow_mut();
-        if let Some(rx) = retained
-            .get_mut(session_id)
-            .and_then(|d| d.permission_event_receiver.as_mut())
-        {
-            while let Ok(event) = rx.try_recv() {
-                events.push(event);
-            }
-        }
-        events
+        self.session_registry.drain_permission_events(session_id)
+    }
+}
+/// Rollback guard for mid-session bind reservation.
+#[cfg(feature = "local-workspace")]
+struct LocalWorkspaceBindGuard {
+    bound: Rc<RefCell<std::collections::HashSet<acp::SessionId>>>,
+    session_id: acp::SessionId,
+    keep: bool,
+}
+#[cfg(feature = "local-workspace")]
+impl Drop for LocalWorkspaceBindGuard {
+    fn drop(&mut self) {
+        if !self.keep {
+            self.bound.borrow_mut().remove(&self.session_id);
+        }
+    }
+}
+/// Reap guard: if session/new fails after register, Drop kills the supervisor.
+#[cfg(all(feature = "local-workspace", unix))]
+pub(crate) struct LocalWorkspaceReapGuard {
+    supervisors: Rc<
+        RefCell<
+            HashMap<
+                acp::SessionId,
+                crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle,
+            >,
+        >,
+    >,
+    generations: Rc<RefCell<HashMap<acp::SessionId, u64>>>,
+    session_id: acp::SessionId,
+    armed: bool,
+}
+#[cfg(all(feature = "local-workspace", unix))]
+impl LocalWorkspaceReapGuard {
+    pub(crate) fn disarm(&mut self) {
+        self.armed = false;
+    }
+}
+#[cfg(all(feature = "local-workspace", unix))]
+impl Drop for LocalWorkspaceReapGuard {
+    fn drop(&mut self) {
+        if !self.armed {
+            return;
+        }
+        self.generations.borrow_mut().remove(&self.session_id);
+        if let Some(handle) = self.supervisors.borrow_mut().remove(&self.session_id) {
+            tokio::spawn(async move {
+                handle.shutdown().await;
+            });
+        }
     }
 }

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

@@ -34,11 +34,8 @@ impl MvpAgent {
         // Pin the index to the requesting session so the Weak in
         // CodebaseIndexManager doesn't orphan it immediately.
         if let Some(sid) = session_id {
-            self.resident_resources
-                .borrow_mut()
-                .entry(sid.clone())
-                .or_default()
-                .codebase_index = Some(std::sync::Arc::clone(&handle));
+            self.session_registry
+                .set_codebase_index(sid, std::sync::Arc::clone(&handle));
         }
         Some((handle, was_newly_started))
     }

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

@@ -183,6 +183,51 @@ pub(crate) fn jwt_claim_matches_user_subscription_tier(
         _ => jwt_claim.parse::<u64>().is_ok_and(|n| n != 0),
     }
 }
+/// ACP `_meta` key for chat+local workspace intent (pager stamps on chat create).
+#[cfg(feature = "local-workspace")]
+const LOCAL_WORKSPACE_META_KEY: &str = "x.ai/local_workspace";
+/// True when `_meta` carries a valid chat+local intent object
+/// (`mode` is `"own"` or `"attach"`).
+#[cfg(feature = "local-workspace")]
+fn local_workspace_intent_present(meta: Option<&acp::Meta>) -> bool {
+    meta.and_then(|m| m.get(LOCAL_WORKSPACE_META_KEY))
+        .and_then(|v| v.as_object())
+        .and_then(|o| o.get("mode"))
+        .and_then(|m| m.as_str())
+        .is_some_and(|mode| mode == "own" || mode == "attach")
+}
+/// valid local-workspace intent → ExistingWorkspace only.
+///
+/// `server_id` comes from the intent object, else `cloud_existing_workspace`.
+/// Never reads `envId` / never emits `SandboxEnvironment`.
+#[cfg(feature = "local-workspace")]
+fn parse_local_workspace_existing(
+    meta: Option<&acp::Meta>,
+) -> Option<crate::gateway_bridge::ComputerSession> {
+    use crate::gateway_bridge::ComputerSession;
+    let local = meta.and_then(|m| m.get(LOCAL_WORKSPACE_META_KEY))?;
+    let mode = local.get("mode").and_then(|v| v.as_str())?;
+    if mode != "own" && mode != "attach" {
+        return None;
+    }
+    let server_id = meta_non_empty_str(local, "server_id")
+        .or_else(|| {
+            meta
+                .and_then(|m| m.get(CLOUD_EXISTING_WORKSPACE_META_KEY))
+                .and_then(|w| meta_non_empty_str(w, "server_id"))
+        })?;
+    let cwd = meta_non_empty_str(local, "cwd")
+        .or_else(|| {
+            meta
+                .and_then(|m| m.get(CLOUD_EXISTING_WORKSPACE_META_KEY))
+                .and_then(|w| meta_non_empty_str(w, "cwd"))
+        });
+    Some(ComputerSession::ExistingWorkspace {
+        server_id,
+        cwd,
+    })
+}
+#[allow(dead_code)]
 fn parse_session_computer_sessions(_meta: Option<&acp::Meta>) -> Option<Vec<()>> {
     None
 }
@@ -666,16 +711,12 @@ pub struct MvpAgent {
     /// leader's auto-update checker, which cannot read the `!Send` maps. Expires
     /// when the actor exits. See [`crate::agent::activity::AgentActivity`].
     pub(crate) activity: crate::agent::activity::AgentActivity,
-    /// LEADER-SAFE(per-session): in-flight `session/load` guards. Lets a racing
-    /// `session/prompt` wait via [`Self::wait_for_in_flight_session_load`] instead
-    /// of failing "unknown session id"; the RAII guard's drop wakes waiters.
+    /// LEADER-SAFE(per-session).
+    session_registry: SessionRegistry,
+    /// A load guard rather than session state: it exists before the session.
     loading_sessions: RefCell<
         HashMap<acp::SessionId, tokio::sync::watch::Receiver<bool>>,
     >,
-    /// LEADER-SAFE(per-session): reclaimed at `remove_session`. See [`RetainedResources`].
-    retained_resources: RefCell<HashMap<acp::SessionId, RetainedResources>>,
-    /// LEADER-SAFE(per-session): keyed by SessionId. Mirrors `sessions` lifecycle.
-    session_threads: RefCell<HashMap<acp::SessionId, SessionThread>>,
     /// Title per resident session id, refreshed each `build_roster`. Lets the
     /// synchronous roster deltas reuse the title instead of emitting an empty
     /// one — `resident_roster_entry` can't read disk.
@@ -793,8 +834,6 @@ pub struct MvpAgent {
     /// LEADER-SAFE(shared): agent-level code-nav index manager, keyed by cwd,
     /// no per-client state.
     codebase_indexes: Arc<parking_lot::Mutex<CodebaseIndexManager>>,
-    /// LEADER-SAFE(per-session): reclaimed on removal / idle-unload. See [`ResidentResources`].
-    resident_resources: RefCell<HashMap<acp::SessionId, ResidentResources>>,
     /// Worktree creation type (resolved: local config > remote > default Linked).
     pub(crate) worktree_type: crate::util::config::WorktreeType,
     /// Restore codebase state on worktree resume (resolved: local config > remote > default false).
@@ -810,15 +849,6 @@ pub struct MvpAgent {
     agent_mcp_state: std::sync::Arc<
         tokio::sync::Mutex<crate::session::mcp_servers::McpState>,
     >,
-    /// Sessions whose persisted model was unavailable at `session/load` time
-    /// with no same-family fallback, keyed by session id → the unavailable
-    /// model id. Prompts to these sessions are blocked until either
-    /// (a) the model reappears in the catalog — the catalog can be
-    /// transiently degraded when a reconnect replays `session/load` (e.g.
-    /// fetch still in flight after a leader restart), so the prompt path
-    /// re-checks and self-heals — or (b) the user explicitly switches
-    /// models via `set_session_model`. Released by `remove_session`.
-    model_unavailable_sessions: RefCell<std::collections::HashMap<String, acp::ModelId>>,
     /// Unified sender for all subagent coordinator events.
     /// LEADER-SAFE(shared): channel is multi-producer, coordinator drains.
     subagent_event_tx: tokio::sync::mpsc::UnboundedSender<
@@ -888,13 +918,28 @@ pub struct MvpAgent {
     /// The agent never opens Computer Hub as a harness/client; remote cloud
     /// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`).
     workspace_ops: RefCell<Option<xai_grok_workspace::WorkspaceOps>>,
-    /// Per-session coarse lifecycle state (residency + turn-state).
-    /// Updated by `spawn_and_register_session` (→ `IdleResident`) and the
-    /// join-handle supervisor on actor exit (→ `DeadFailed`) / explicit close
-    /// (→ `Completed`). This is the roster's data source in PR-6; for now it
-    /// gives the supervisor an observable demotion signal.
-    /// LEADER-SAFE(per-session): keyed by SessionId.
-    session_live_state: RefCell<HashMap<acp::SessionId, SessionLiveState>>,
+    /// Per-session owned local `workspace_server` handles (chat+local `own`).
+    #[cfg(all(feature = "local-workspace", unix))]
+    local_workspace_supervisors: Rc<
+        RefCell<
+            HashMap<
+                acp::SessionId,
+                crate::gateway_bridge::local_workspace_supervisor::LocalWorkspaceHandle,
+            >,
+        >,
+    >,
+    /// Invalidates in-flight crash restarts when the session supervisor is reaped.
+    #[cfg(all(feature = "local-workspace", unix))]
+    local_workspace_generations: Rc<RefCell<HashMap<acp::SessionId, u64>>>,
+    /// Sessions whose own supervisor is mid crash-restart (map entry temporarily empty).
+    #[cfg(all(feature = "local-workspace", unix))]
+    local_workspace_restart_pending: Rc<
+        RefCell<std::collections::HashSet<acp::SessionId>>,
+    >,
+    /// Sessions that already have a local existing workspace (own or attach).
+    /// mid-session add refuses while this is set; cleared on session end.
+    #[cfg(feature = "local-workspace")]
+    local_workspace_bound: Rc<RefCell<std::collections::HashSet<acp::SessionId>>>,
     /// Idempotency guard: the join-handle supervisor task is spawned at most
     /// once (on the first `spawn_and_register_session`). See
     /// `ensure_session_supervisor`.
@@ -1297,10 +1342,12 @@ impl Drop for SessionLoadGuard<'_> {
 mod code_nav;
 mod folder_trust_prompt;
 mod heap_profile;
+mod session_registry;
 mod session_lifecycle;
 mod subagent_coordinator;
 mod agent_ops;
 mod acp_agent;
+use session_registry::SessionRegistry;
 pub(crate) use session_lifecycle::RegistrySnapshot;
 pub(super) use super::ext_parsers;
 /// Emit the `auth.lifecycle` login span with optional user id and error

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

@@ -26,12 +26,8 @@ impl MvpAgent {
         });
         let _ = handle.cmd_tx.send(SessionCommand::Shutdown);
         drop(handle);
-        let thread = self.session_threads.borrow_mut().remove(id);
         self.remove_session_terminal(id, SessionLiveState::Completed);
-        if let Some(thread) = thread {
-            self.session_threads.borrow_mut().insert(id.clone(), thread);
-            self.drain_old_session_thread(id).await;
-        }
+        self.drain_old_session_thread(id).await;
     }
     /// Finalize the cloud session replica (fire-and-forget, "Hook 4").
     ///
@@ -78,13 +74,7 @@ impl MvpAgent {
                 parent_session_id: id.0.to_string(),
             });
         self.take_session(id);
-        self.session_threads.borrow_mut().remove(id);
-        self.resident_resources.borrow_mut().remove(id);
-        self.retained_resources.borrow_mut().remove(id);
-        self.model_unavailable_sessions
-            .borrow_mut()
-            .remove(id.0.as_ref());
-        self.session_live_state.borrow_mut().remove(id);
+        self.session_registry.release(id);
         if let Some(ops) = self.workspace_ops.borrow().as_ref() {
             ops.end_local_session(id.0.as_ref());
         }
@@ -97,13 +87,7 @@ impl MvpAgent {
     /// Cancels therefore wait behind an intake preamble: keep preambles lean.
     /// Bridge cancels take their own path and stay unordered against this lock.
     pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> std::rc::Rc<tokio::sync::Mutex<()>> {
-        self.retained_resources
-            .borrow_mut()
-            .entry(id.clone())
-            .or_default()
-            .dispatch_lock
-            .get_or_insert_with(Default::default)
-            .clone()
+        self.session_registry.dispatch_lock(id)
     }
     /// Close a session in response to an **explicit** terminal close
     /// (`x.ai/session/close`). Finalizes the cloud replica (genuine session
@@ -114,14 +98,12 @@ impl MvpAgent {
     }
     /// 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);
+        self.session_registry.set_live(id, 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()
+        self.session_registry.live(id)
     }
     /// Roster-delta hook for a terminally removed session. Broadcasts an
     /// `x.ai/sessions/changed` notification with the session in `removed` so
@@ -231,7 +213,7 @@ impl MvpAgent {
         if turn_running {
             return RosterActivity::Working;
         }
-        match self.session_live_state.borrow().get(id).copied() {
+        match self.session_registry.live(id) {
             Some(SessionLiveState::Completed) => RosterActivity::Completed,
             Some(SessionLiveState::DeadFailed) => RosterActivity::Dead,
             Some(SessionLiveState::Dormant) => RosterActivity::Dormant,
@@ -339,13 +321,7 @@ impl MvpAgent {
     /// 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();
+        let dead = self.session_registry.finished_threads();
         for id in dead {
             if self.sessions.borrow().contains_key(&id) {
                 tracing::warn!(
@@ -354,8 +330,7 @@ impl MvpAgent {
                 );
                 self.reap_dead_session(&id);
             } else {
-                self.session_threads.borrow_mut().remove(&id);
-                self.session_live_state.borrow_mut().remove(&id);
+                self.session_registry.clear_exited_thread(&id);
                 tracing::debug!(
                     session_id = %id.0,
                     "Reaped finished thread for non-resident session (clean exit)"
@@ -453,8 +428,9 @@ impl MvpAgent {
             .await
             .unwrap_or(true)
     }
-    /// Entry counts for every collection [`Self::remove_session`] drains,
-    /// plus workspace bindings and shared coordinator state.
+    /// Entry counts `remove_session` and `x.ai/debug/agent` care about, which
+    /// includes maps outside [`SessionResources`]: the handle map, load guards,
+    /// rewind snapshots, subagents, and workspace bindings.
     pub(crate) async fn registry_snapshot(&self) -> RegistrySnapshot {
         let subagents =
             xai_grok_tools::implementations::grok_build::task::backend::ChannelBackend::new(
@@ -462,52 +438,29 @@ impl MvpAgent {
             )
             .registry_counts()
             .await;
-        let (resident_resources, session_index_claims, require_gateway_sessions) = {
-            let resident = self.resident_resources.borrow();
-            (
-                resident.len(),
-                resident
-                    .values()
-                    .filter(|r| r.codebase_index.is_some())
-                    .count(),
-                resident.values().filter(|r| r.require_gateway).count(),
-            )
-        };
-        let retained = self.retained_resources.borrow();
-        let retained_resources = retained.len();
-        let dispatch_locks = retained
-            .values()
-            .filter(|d| d.dispatch_lock.is_some())
-            .count();
-        let session_turn_numbers = retained
-            .values()
-            .filter(|d| d.turn_number.is_some())
-            .count();
-        let permission_event_receivers = retained
-            .values()
-            .filter(|d| d.permission_event_receiver.is_some())
-            .count();
-        drop(retained);
+        let workspace_ops = self.workspace_ops.borrow();
+        let workspace = workspace_ops
+            .as_ref()
+            .and_then(|ops| ops.workspace_handle());
+        let counts = self.session_registry.counts();
         RegistrySnapshot {
             sessions: self.sessions.borrow().len(),
-            session_threads: self.session_threads.borrow().len(),
-            resident_resources,
-            retained_resources,
-            dispatch_locks,
-            session_turn_numbers,
-            permission_event_receivers,
-            model_unavailable_sessions: self.model_unavailable_sessions.borrow().len(),
-            session_live_state: self.session_live_state.borrow().len(),
-            session_index_claims,
-            require_gateway_sessions,
+            loading_sessions: self.loading_sessions.borrow().len(),
+            session_threads: counts.session_threads,
+            resident_resources: counts.resident_resources,
+            retained_resources: counts.retained_resources,
+            dispatch_locks: counts.dispatch_locks,
+            session_turn_numbers: counts.session_turn_numbers,
+            permission_event_receivers: counts.permission_event_receivers,
+            model_unavailable_sessions: counts.model_unavailable_sessions,
+            session_live_state: counts.session_live_state,
+            session_index_claims: counts.session_index_claims,
+            require_gateway_sessions: counts.require_gateway_sessions,
             subagent_pending: subagents.pending,
             subagent_active: subagents.active,
             subagent_completed: subagents.completed,
-            workspace_bindings: self
-                .workspace_ops
-                .borrow()
-                .as_ref()
-                .and_then(|ops| ops.workspace_handle().map(|h| h.session_count())),
+            workspace_bindings: workspace.map(|h| h.session_count()),
+            workspace_activity_sessions: workspace.map(|h| h.activity_tracker().session_count()),
         }
     }
 }
@@ -516,6 +469,7 @@ impl MvpAgent {
 #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)]
 pub struct RegistrySnapshot {
     pub sessions: usize,
+    pub loading_sessions: usize,
     pub session_threads: usize,
     pub resident_resources: usize,
     pub retained_resources: usize,
@@ -530,4 +484,5 @@ pub struct RegistrySnapshot {
     pub subagent_active: usize,
     pub subagent_completed: usize,
     pub workspace_bindings: Option<usize>,
+    pub workspace_activity_sessions: Option<usize>,
 }

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

@@ -0,0 +1,238 @@
+//! Per-session resources and the registry that owns them. Distinct from
+//! `agent::session_registry_client`, which talks to the remote registry.
+use super::*;
+/// The map stays private so every caller goes through a named operation.
+#[derive(Clone, Default)]
+pub(super) struct SessionRegistry {
+    sessions: Rc<RefCell<HashMap<acp::SessionId, SessionResources>>>,
+}
+/// The per-session state this registry owns: retained, resident, thread, live,
+/// unavailable model, and bridge. Load guards, rewind snapshots, local
+/// workspaces, and the handle map are owned elsewhere, so a new field belongs
+/// here only if `release` should drop it with the rest.
+#[derive(Default)]
+struct SessionResources {
+    retained: Option<RetainedResources>,
+    /// Cleared at idle-unload; survives a reload rebuild.
+    resident: Option<ResidentResources>,
+    thread: Option<SessionThread>,
+    live: Option<SessionLiveState>,
+    unavailable_model: Option<acp::ModelId>,
+}
+#[derive(Default)]
+pub(super) struct SessionCounts {
+    pub(super) retained_resources: usize,
+    pub(super) resident_resources: usize,
+    pub(super) session_threads: usize,
+    pub(super) session_live_state: usize,
+    pub(super) model_unavailable_sessions: usize,
+    pub(super) dispatch_locks: usize,
+    pub(super) session_turn_numbers: usize,
+    pub(super) permission_event_receivers: usize,
+    pub(super) session_index_claims: usize,
+    pub(super) require_gateway_sessions: usize,
+}
+impl SessionRegistry {
+    /// Releases everything a closing session leaves behind, in one drop.
+    ///
+    /// A running actor thread stays: dropping its handle would detach it, and
+    /// nothing would track the memory it holds. The sweep reclaims it later.
+    pub(super) fn release(&self, id: &acp::SessionId) {
+        let mut entries = self.sessions.borrow_mut();
+        let Some(mut released) = entries.remove(id) else {
+            return;
+        };
+        let running = released.thread.take().filter(|t| !t.is_finished());
+        drop(released);
+        if running.is_some() {
+            entries.insert(
+                id.clone(),
+                SessionResources {
+                    thread: running,
+                    retained: None,
+                    resident: None,
+                    live: None,
+                    unavailable_model: None,
+                },
+            );
+        }
+    }
+    pub(super) fn set_thread(&self, id: &acp::SessionId, thread: SessionThread) {
+        let displaced = self.edit(id, |e| e.thread.replace(thread));
+        if displaced.is_some_and(|t| !t.is_finished()) {
+            tracing::warn!(session_id = %id.0, "session thread displaced while still running");
+        }
+    }
+    /// Drops the tracked thread. Returns nothing on purpose: handing a
+    /// `SessionThread` to a caller lets the last handle die in a local, which
+    /// detaches the thread with no record left for the sweep.
+    pub(super) fn clear_thread(&self, id: &acp::SessionId) {
+        self.clear(id, |e| e.thread = None);
+    }
+    /// `None` when no thread is tracked for the session.
+    #[cfg(test)]
+    pub(super) fn has_thread(&self, id: &acp::SessionId) -> bool {
+        self.with(id, |e| e.thread.is_some()).unwrap_or(false)
+    }
+    pub(super) fn thread_is_finished(&self, id: &acp::SessionId) -> Option<bool> {
+        self.with(id, |e| e.thread.as_ref().map(SessionThread::is_finished))
+            .flatten()
+    }
+    pub(super) fn finished_threads(&self) -> Vec<acp::SessionId> {
+        self.sessions
+            .borrow()
+            .iter()
+            .filter(|(_, e)| e.thread.as_ref().is_some_and(SessionThread::is_finished))
+            .map(|(id, _)| id.clone())
+            .collect()
+    }
+    pub(super) fn clear_exited_thread(&self, id: &acp::SessionId) {
+        self.clear(id, |e| {
+            e.thread = None;
+            e.live = None;
+        });
+    }
+    pub(super) fn set_live(&self, id: &acp::SessionId, state: SessionLiveState) {
+        self.edit(id, |e| e.live = Some(state));
+    }
+    pub(super) fn live(&self, id: &acp::SessionId) -> Option<SessionLiveState> {
+        self.with(id, |e| e.live).flatten()
+    }
+    pub(super) fn clear_resident(&self, id: &acp::SessionId) {
+        self.clear(id, |e| e.resident = None);
+    }
+    pub(super) fn set_unavailable_model(&self, id: &acp::SessionId, model: acp::ModelId) {
+        self.edit(id, |e| e.unavailable_model = Some(model));
+    }
+    pub(super) fn unavailable_model(&self, id: &acp::SessionId) -> Option<acp::ModelId> {
+        self.with(id, |e| e.unavailable_model.clone()).flatten()
+    }
+    pub(super) fn take_unavailable_model(&self, id: &acp::SessionId) -> Option<acp::ModelId> {
+        let model = self
+            .sessions
+            .borrow_mut()
+            .get_mut(id)
+            .and_then(|e| e.unavailable_model.take());
+        self.drop_if_empty(id);
+        model
+    }
+    pub(super) fn turn_number(&self, id: &acp::SessionId) -> Option<u64> {
+        self.with(id, |e| e.retained.as_ref()?.turn_number)
+            .flatten()
+    }
+    pub(super) fn set_turn_number(&self, id: &acp::SessionId, next: u64) {
+        self.edit(id, |e| {
+            e.retained.get_or_insert_default().turn_number = Some(next);
+        });
+    }
+    pub(super) fn dispatch_lock(&self, id: &acp::SessionId) -> Rc<tokio::sync::Mutex<()>> {
+        self.edit(id, |e| {
+            e.retained
+                .get_or_insert_default()
+                .dispatch_lock
+                .get_or_insert_with(Default::default)
+                .clone()
+        })
+    }
+    pub(super) fn set_permission_receiver(
+        &self,
+        id: &acp::SessionId,
+        rx: tokio::sync::mpsc::UnboundedReceiver<PermissionEvent>,
+    ) {
+        self.edit(id, |e| {
+            e.retained.get_or_insert_default().permission_event_receiver = Some(rx);
+        });
+    }
+    pub(super) fn drain_permission_events(&self, id: &acp::SessionId) -> Vec<PermissionEvent> {
+        let mut events = Vec::new();
+        let mut entries = self.sessions.borrow_mut();
+        if let Some(rx) = entries
+            .get_mut(id)
+            .and_then(|e| e.retained.as_mut())
+            .and_then(|r| r.permission_event_receiver.as_mut())
+        {
+            while let Ok(event) = rx.try_recv() {
+                events.push(event);
+            }
+        }
+        events
+    }
+    pub(super) fn set_codebase_index(
+        &self,
+        id: &acp::SessionId,
+        index: std::sync::Arc<xai_codebase_graph::IndexManagerHandle>,
+    ) {
+        self.edit(id, |e| {
+            e.resident.get_or_insert_default().codebase_index = Some(index);
+        });
+    }
+    /// Destructured so a new field has to be counted, or go unmeasured.
+    pub(super) fn counts(&self) -> SessionCounts {
+        let mut counts = SessionCounts::default();
+        for entry in self.sessions.borrow().values() {
+            let SessionResources {
+                retained,
+                resident,
+                thread,
+                live,
+                unavailable_model,
+            } = entry;
+            counts.retained_resources += usize::from(retained.is_some());
+            counts.resident_resources += usize::from(resident.is_some());
+            counts.session_threads += usize::from(thread.is_some());
+            counts.session_live_state += usize::from(live.is_some());
+            counts.model_unavailable_sessions += usize::from(unavailable_model.is_some());
+            if let Some(retained) = retained {
+                counts.dispatch_locks += usize::from(retained.dispatch_lock.is_some());
+                counts.session_turn_numbers += usize::from(retained.turn_number.is_some());
+                counts.permission_event_receivers +=
+                    usize::from(retained.permission_event_receiver.is_some());
+            }
+            if let Some(resident) = resident {
+                counts.session_index_claims += usize::from(resident.codebase_index.is_some());
+                counts.require_gateway_sessions += usize::from(resident.require_gateway);
+            }
+        }
+        counts
+    }
+    fn with<R>(&self, id: &acp::SessionId, f: impl FnOnce(&SessionResources) -> R) -> Option<R> {
+        self.sessions.borrow().get(id).map(f)
+    }
+    fn edit<R>(&self, id: &acp::SessionId, f: impl FnOnce(&mut SessionResources) -> R) -> R {
+        f(self.sessions.borrow_mut().entry(id.clone()).or_default())
+    }
+    fn clear(&self, id: &acp::SessionId, f: impl FnOnce(&mut SessionResources)) {
+        {
+            let mut entries = self.sessions.borrow_mut();
+            let Some(entry) = entries.get_mut(id) else {
+                return;
+            };
+            f(entry);
+        }
+        self.drop_if_empty(id);
+    }
+    fn drop_if_empty(&self, id: &acp::SessionId) {
+        let mut entries = self.sessions.borrow_mut();
+        if entries.get(id).is_some_and(SessionResources::is_empty) {
+            entries.remove(id);
+        }
+    }
+}
+impl SessionResources {
+    fn is_empty(&self) -> bool {
+        let Self {
+            retained,
+            resident,
+            thread,
+            live,
+            unavailable_model,
+        } = self;
+        let chat_vacant = true;
+        retained.is_none()
+            && resident.is_none()
+            && thread.is_none()
+            && live.is_none()
+            && unavailable_model.is_none()
+            && chat_vacant
+    }
+}

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

@@ -3110,6 +3110,290 @@ fn chat_new_session_model_state_matrix() {
         );
     }
 }
+/// valid `x.ai/local_workspace` → ExistingWorkspace only.
+/// Never reads `envId` / never emits SandboxEnvironment.
+#[cfg(feature = "local-workspace")]
+#[test]
+fn parse_session_computer_sessions_local_workspace_matrix() {
+    use crate::gateway_bridge::ComputerSession;
+    use serde_json::json;
+    fn existing(server_id: &str, cwd: Option<&str>) -> Vec<ComputerSession> {
+        vec![ComputerSession::ExistingWorkspace {
+            server_id: server_id.to_owned(),
+            cwd: cwd.map(str::to_owned),
+        }]
+    }
+    let cases: &[(&str, serde_json::Value, Option<Vec<ComputerSession>>)] = &[
+        (
+            "attach_server_id_on_local",
+            json!({
+                "x.ai/local_workspace": {
+                    "mode": "attach",
+                    "server_id": "lw-attach-1",
+                    "cwd": "/repo",
+                },
+                "envId": "env-must-be-ignored",
+            }),
+            Some(existing("lw-attach-1", Some("/repo"))),
+        ),
+        (
+            "attach_server_id_from_cloud_existing",
+            json!({
+                "x.ai/local_workspace": {
+                    "mode": "attach",
+                    "cwd": "/repo",
+                },
+                "x.ai/cloud_existing_workspace": {
+                    "server_id": "lw-attach-2",
+                    "cwd": "/repo-existing",
+                },
+                "envId": "env-must-be-ignored",
+            }),
+            Some(existing("lw-attach-2", Some("/repo"))),
+        ),
+        (
+            "own_with_server_id_ignores_envid",
+            json!({
+                "x.ai/local_workspace": {
+                    "mode": "own",
+                    "server_id": "lw-own-1",
+                    "cwd": "/Users/me/src",
+                },
+                "envId": "env-must-be-ignored",
+            }),
+            Some(existing("lw-own-1", Some("/Users/me/src"))),
+        ),
+        (
+            "own_without_server_id_no_sandbox_fallback",
+            json!({
+                "x.ai/local_workspace": {
+                    "mode": "own",
+                    "cwd": "/Users/me/src",
+                },
+                "envId": "env-must-be-ignored",
+            }),
+            None,
+        ),
+        (
+            "invalid_mode_falls_through_to_envid",
+            json!({
+                "x.ai/local_workspace": {
+                    "mode": "bogus",
+                    "server_id": "lw-x",
+                },
+                "envId": "env-prod",
+            }),
+            Some(vec![ComputerSession::SandboxEnvironment {
+                environment_id: Some("env-prod".to_owned()),
+            }]),
+        ),
+        (
+            "non_object_local_falls_through_to_envid",
+            json!({
+                "x.ai/local_workspace": "not-an-object",
+                "envId": "env-prod",
+            }),
+            Some(vec![ComputerSession::SandboxEnvironment {
+                environment_id: Some("env-prod".to_owned()),
+            }]),
+        ),
+    ];
+    for (label, meta, expected) in cases {
+        let got = parse_session_computer_sessions(meta.as_object());
+        assert_eq!(
+            got.as_deref(),
+            expected.as_deref(),
+            "[{label}] local_workspace match-table mismatch"
+        );
+    }
+}
+/// Local intent without resolvable server_id fails closed (no silent unstamped start).
+#[cfg(feature = "local-workspace")]
+#[test]
+fn resolve_local_workspace_missing_server_id_fails_closed() {
+    use serde_json::json;
+    let meta = json!({
+        "x.ai/session": { "kind": "chat" },
+        "x.ai/local_workspace": {
+            "mode": "own",
+            "cwd": "/repo",
+        }
+    });
+    let err = resolve_session_computer_sessions(meta.as_object())
+        .expect_err("own without server_id must fail closed");
+    assert_eq!(
+        err.data
+            .as_ref()
+            .and_then(|d| d.get("code"))
+            .and_then(|v| v.as_str()),
+        Some("local_workspace_server_id_missing")
+    );
+}
+/// Supervisor map + reap guard / shutdown_gateway_bridge tear down the entry.
+#[cfg(all(feature = "local-workspace", unix))]
+#[test]
+fn local_workspace_reap_guard_and_shutdown_clear_map() {
+    run_local_for_bridge_test(|| async {
+        let agent = build_minimal_agent_for_tests();
+        let sid = gateway_bridge_test_session_id();
+        {
+            let mut guard = agent.new_local_workspace_reap_guard(sid.clone(), true);
+            guard.disarm();
+        }
+        assert!(agent.local_workspace_supervisors.borrow().is_empty());
+        agent.shutdown_gateway_bridge(&sid);
+        assert!(
+            agent
+                .local_workspace_generations
+                .borrow()
+                .get(&sid)
+                .is_none()
+        );
+    });
+}
+/// Pre-bridge crash refresh rewrites handshake stamp from live supervisor id.
+#[cfg(all(feature = "local-workspace", unix))]
+#[test]
+fn refresh_sessions_from_supervisor_overrides_server_id() {
+    use crate::gateway_bridge::ComputerSession;
+    use crate::gateway_bridge::local_workspace_supervisor::test_start_ready_own;
+    run_local_for_bridge_test(|| async {
+        let agent = build_minimal_agent_for_tests();
+        let sid = gateway_bridge_test_session_id();
+        let original = Some(vec![ComputerSession::ExistingWorkspace {
+            server_id: "lw-stale".into(),
+            cwd: Some("/repo".into()),
+        }]);
+        let unchanged = agent.refresh_sessions_from_supervisor(&sid, original.clone());
+        assert!(matches!(
+            unchanged.as_ref().and_then(|v| v.first()),
+            Some(ComputerSession::ExistingWorkspace { server_id, .. }) if server_id == "lw-stale"
+        ));
+        let (_dir, handle) = test_start_ready_own().await;
+        let live_id = handle.server_id.clone();
+        agent.register_local_workspace_supervisor(sid.clone(), handle);
+        let refreshed = agent.refresh_sessions_from_supervisor(&sid, original);
+        match refreshed.as_ref().and_then(|v| v.first()) {
+            Some(ComputerSession::ExistingWorkspace { server_id, .. }) => {
+                assert_eq!(
+                    server_id, &live_id,
+                    "refresh must use live supervisor server_id"
+                );
+            }
+            other => panic!("expected ExistingWorkspace, got {other:?}"),
+        }
+        agent.shutdown_gateway_bridge(&sid);
+    });
+}
+/// start_own + register stamps server_id into meta and stores the handle.
+#[cfg(all(feature = "local-workspace", unix))]
+#[test]
+fn start_own_registers_and_stamps_server_id() {
+    use crate::gateway_bridge::local_workspace_supervisor::{
+        stamp_server_id_into_meta, test_start_ready_own,
+    };
+    run_local_for_bridge_test(|| async {
+        let agent = build_minimal_agent_for_tests();
+        let sid = gateway_bridge_test_session_id();
+        let (_dir, handle) = test_start_ready_own().await;
+        let server_id = handle.server_id.clone();
+        let mut meta = acp::Meta::new();
+        meta.insert(
+            "x.ai/local_workspace".into(),
+            serde_json::json!({"mode": "own", "cwd": "/tmp/repo"}),
+        );
+        stamp_server_id_into_meta(&mut meta, &server_id);
+        assert_eq!(
+            meta.get("x.ai/local_workspace")
+                .and_then(|v| v.get("server_id"))
+                .and_then(|v| v.as_str()),
+            Some(server_id.as_str())
+        );
+        agent.register_local_workspace_supervisor(sid.clone(), handle);
+        assert!(
+            agent
+                .local_workspace_supervisors
+                .borrow()
+                .contains_key(&sid),
+            "handle must be registered by SessionId"
+        );
+        assert!(
+            agent
+                .local_workspace_generations
+                .borrow()
+                .get(&sid)
+                .is_some_and(|g| *g >= 1),
+            "arm must bump generation"
+        );
+        agent.shutdown_gateway_bridge(&sid);
+    });
+}
+/// Armed reap guard removes a registered supervisor on drop (session/new failure).
+#[cfg(all(feature = "local-workspace", unix))]
+#[test]
+fn reap_guard_drop_removes_registered_supervisor() {
+    use crate::gateway_bridge::local_workspace_supervisor::test_start_ready_own;
+    run_local_for_bridge_test(|| async {
+        let agent = build_minimal_agent_for_tests();
+        let sid = gateway_bridge_test_session_id();
+        let (_dir, handle) = test_start_ready_own().await;
+        agent.register_local_workspace_supervisor(sid.clone(), handle);
+        assert!(
+            agent
+                .local_workspace_supervisors
+                .borrow()
+                .contains_key(&sid)
+        );
+        {
+            let _guard = agent.new_local_workspace_reap_guard(sid.clone(), true);
+        }
+        assert!(
+            agent
+                .local_workspace_supervisors
+                .borrow()
+                .get(&sid)
+                .is_none(),
+            "armed guard drop must reap supervisor"
+        );
+        assert!(
+            agent
+                .local_workspace_generations
+                .borrow()
+                .get(&sid)
+                .is_none(),
+            "armed guard drop must invalidate generation"
+        );
+    });
+}
+/// Shutdown generation invalidates a pending restart re-insert.
+#[cfg(all(feature = "local-workspace", unix))]
+#[test]
+fn shutdown_generation_invalidates_stale_restart() {
+    use crate::gateway_bridge::local_workspace_supervisor::test_start_ready_own;
+    run_local_for_bridge_test(|| async {
+        let agent = build_minimal_agent_for_tests();
+        let sid = gateway_bridge_test_session_id();
+        let (_dir, handle) = test_start_ready_own().await;
+        agent.register_local_workspace_supervisor(sid.clone(), handle);
+        let generation = *agent
+            .local_workspace_generations
+            .borrow()
+            .get(&sid)
+            .expect("generation after register");
+        agent.shutdown_gateway_bridge(&sid);
+        assert!(
+            agent.local_workspace_generations.borrow().get(&sid) != Some(&generation),
+            "shutdown must invalidate generation so stale restart cannot re-insert"
+        );
+        assert!(
+            agent
+                .local_workspace_supervisors
+                .borrow()
+                .get(&sid)
+                .is_none()
+        );
+    });
+}
 /// `spawn_gateway_bridge` uses `tokio::task::spawn_local`.
 fn run_local_for_bridge_test<F, Fut, T>(body: F) -> T
 where
@@ -3174,39 +3458,26 @@ async fn remove_session_releases_workspace_binding_and_side_maps() {
     .expect("bind_local_session must succeed");
     assert!(toolset_weak.upgrade().is_some());
     *agent.workspace_ops.borrow_mut() = Some(ops);
-    agent.model_unavailable_sessions.borrow_mut().insert(
-        sid.0.to_string(),
-        acp::ModelId::new(std::sync::Arc::from("gone-model")),
-    );
+    agent
+        .session_registry
+        .set_unavailable_model(&sid, acp::ModelId::new(std::sync::Arc::from("gone-model")));
     agent.set_turn_number(&sid, 3);
     let (_permission_tx, permission_rx) =
         tokio::sync::mpsc::unbounded_channel::<xai_grok_workspace::permission::PermissionEvent>();
     agent
-        .retained_resources
-        .borrow_mut()
-        .entry(sid.clone())
-        .or_default()
-        .permission_event_receiver = Some(permission_rx);
-    agent
-        .resident_resources
-        .borrow_mut()
-        .entry(sid.clone())
-        .or_default()
-        .require_gateway = true;
+        .session_registry
+        .set_permission_receiver(&sid, permission_rx);
+    agent.session_registry.mark_require_gateway(&sid);
     agent.remove_session(&sid);
     assert!(
         toolset_weak.upgrade().is_none(),
         "the workspace binding must release the toolset"
     );
-    assert!(
-        !agent
-            .model_unavailable_sessions
-            .borrow()
-            .contains_key(sid.0.as_ref())
-    );
-    assert!(!agent.resident_resources.borrow().contains_key(&sid));
-    assert!(
-        !agent.retained_resources.borrow().contains_key(&sid),
+    assert!(agent.session_registry.unavailable_model(&sid).is_none());
+    assert_eq!(agent.session_registry.counts().resident_resources, 0);
+    assert_eq!(
+        agent.session_registry.counts().retained_resources,
+        0,
         "retained per-session resources must be reclaimed on removal"
     );
 }
@@ -3484,6 +3755,46 @@ fn disconnect_keeps_resident_on_poisoned_lock() {
         );
     });
 }
+/// A wedged actor stays tracked. `remove_session` releases everything else but
+/// keeps a still-running thread, because dropping its handle would detach the
+/// thread and leave nothing for the supervisor sweep to find.
+#[test]
+fn remove_session_keeps_a_running_thread_tracked() {
+    run_local_for_bridge_test(|| async {
+        let agent = build_minimal_agent_for_tests();
+        let sid = acp::SessionId::new("sess-wedged");
+        let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
+        agent.session_registry.set_thread(
+            &sid,
+            crate::session::SessionThread::from_handle(std::thread::spawn(move || {
+                let _ = release_rx.recv();
+            })),
+        );
+        agent.set_turn_number(&sid, 1);
+        agent.remove_session(&sid);
+        assert!(
+            agent.session_registry.has_thread(&sid),
+            "a running actor thread must survive removal for the sweep"
+        );
+        assert_eq!(
+            agent.session_registry.counts().retained_resources,
+            0,
+            "everything except the running thread must be released"
+        );
+        drop(release_tx);
+        for _ in 0..100 {
+            agent.sweep_dead_sessions();
+            if !agent.session_registry.has_thread(&sid) {
+                break;
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        }
+        assert!(
+            !agent.session_registry.has_thread(&sid),
+            "the sweep must reclaim the thread once it exits"
+        );
+    });
+}
 /// Idle-unload stub (memory bound) + supervisor interaction: a *fully idle*
 /// session is unloaded to disk on disconnect (actor `Shutdown`, handle
 /// dropped) while the `SessionThread` is **retained** for
@@ -3498,8 +3809,8 @@ fn disconnect_unloads_idle_session_without_finalize() {
         agent.sessions.borrow_mut().insert(sid.clone(), handle);
         let mut observed = spawn_fake_actor(cmd_rx, false);
         let (release_tx, release_rx) = std::sync::mpsc::channel::<()>();
-        agent.session_threads.borrow_mut().insert(
-            sid.clone(),
+        agent.session_registry.set_thread(
+            &sid,
             crate::session::SessionThread::from_handle(std::thread::spawn(move || {
                 let _ = release_rx.recv();
             })),
@@ -3511,7 +3822,7 @@ fn disconnect_unloads_idle_session_without_finalize() {
             "idle session must be unloaded from the resident map on disconnect"
         );
         assert!(
-            agent.session_threads.borrow().contains_key(&sid),
+            agent.session_registry.has_thread(&sid),
             "idle-unload must keep the SessionThread for reconnect drain"
         );
         let shutdown = tokio::time::timeout(std::time::Duration::from_secs(1), observed.recv())
@@ -3534,13 +3845,13 @@ fn disconnect_unloads_idle_session_without_finalize() {
         drop(release_tx);
         let deadline = tokio::time::Instant::now() + (SESSION_SUPERVISOR_TICK * 6);
         while tokio::time::Instant::now() < deadline {
-            if !agent.session_threads.borrow().contains_key(&sid) {
+            if !agent.session_registry.has_thread(&sid) {
                 break;
             }
             tokio::time::sleep(std::time::Duration::from_millis(20)).await;
         }
         assert!(
-            !agent.session_threads.borrow().contains_key(&sid),
+            !agent.session_registry.has_thread(&sid),
             "supervisor must drop the finished kept thread"
         );
         assert!(
@@ -3710,7 +4021,7 @@ fn session_live_state_map_is_bounded_across_cycles() {
             agent.close_session_explicit(&sid);
         }
         assert_eq!(
-            agent.session_live_state.borrow().len(),
+            agent.session_registry.counts().session_live_state,
             0,
             "terminal closes must leave no residual live-state entries (bounded map)"
         );
@@ -3782,21 +4093,21 @@ fn supervisor_reaps_panicked_resident_actor() {
         let (handle, _tx, _rx) = make_live_session_handle(&sid, Some("turn-1"));
         agent.sessions.borrow_mut().insert(sid.clone(), handle);
         let panic_thread = std::thread::spawn(|| panic!("injected actor panic"));
-        agent.session_threads.borrow_mut().insert(
-            sid.clone(),
+        agent.session_registry.set_thread(
+            &sid,
             crate::session::SessionThread::from_handle(panic_thread),
         );
         agent.set_session_live_state(&sid, SessionLiveState::Working);
         agent.ensure_session_supervisor();
         let deadline = tokio::time::Instant::now() + (SESSION_SUPERVISOR_TICK * 6);
         while tokio::time::Instant::now() < deadline {
-            if !agent.session_threads.borrow().contains_key(&sid) {
+            if !agent.session_registry.has_thread(&sid) {
                 break;
             }
             tokio::time::sleep(std::time::Duration::from_millis(20)).await;
         }
         assert!(
-            !agent.session_threads.borrow().contains_key(&sid),
+            !agent.session_registry.has_thread(&sid),
             "supervisor must reap the dead thread"
         );
         assert!(

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

@@ -37,17 +37,11 @@ fn populate_and_evict(agent: &MvpAgent, i: usize) {
     }
 
     let (_ptx, prx) = tokio::sync::mpsc::unbounded_channel::<PermissionEvent>();
-    agent
-        .retained_resources
-        .borrow_mut()
-        .entry(sid.clone())
-        .or_default()
-        .permission_event_receiver = Some(prx);
+    agent.session_registry.set_permission_receiver(&sid, prx);
     agent.set_turn_number(&sid, i as u64);
-    agent.model_unavailable_sessions.borrow_mut().insert(
-        sid.0.to_string(),
-        acp::ModelId::new(std::sync::Arc::from("gone-model")),
-    );
+    agent
+        .session_registry
+        .set_unavailable_model(&sid, acp::ModelId::new(std::sync::Arc::from("gone-model")));
     agent.remove_session(&sid);
 }
 

crates/codegen/xai-grok-shell/src/auth/device_code.rs

@@ -21,6 +21,11 @@ const DEFAULT_DEVICE_POLL_INTERVAL_SECS: i32 = 5;
 const DEVICE_SLOW_DOWN_INCREMENT_SECS: u64 = 5;
 const MIN_DEVICE_CODE_EXPIRY_FALLBACK_SECS: i64 = 10 * 60;
 
+/// Only the 404 "no device endpoint" case is typed, because the login flow
+/// matches on it to fall back to loopback. Every other device-code failure
+/// stays a plain `anyhow` error: wrapping one in a `#[error(transparent)]`
+/// variant hides the `reqwest::Error` the login funnel classifies, because
+/// transparent forwards `source()` past the error it wraps.
 #[derive(Debug, Error)]
 pub enum DeviceCodeError {
     #[error(
@@ -28,14 +33,6 @@ pub enum DeviceCodeError {
          Try `grok login` or set XAI_API_KEY instead."
     )]
     NotEnabled,
-    #[error(transparent)]
-    Other(#[from] anyhow::Error),
-}
-
-impl From<reqwest::Error> for DeviceCodeError {
-    fn from(e: reqwest::Error) -> Self {
-        Self::Other(e.into())
-    }
 }
 
 // --- Public types ---
@@ -137,7 +134,7 @@ pub async fn request_device_code(
     client_id: &str,
     scopes: &[String],
     surface: ClientSurface,
-) -> Result<DeviceCode, DeviceCodeError> {
+) -> anyhow::Result<DeviceCode> {
     let client = crate::http::shared_client();
     let url = format!("{}/oauth2/device/code", issuer.trim_end_matches('/'));
     let scope_str = scopes.join(" ");
@@ -164,9 +161,9 @@ pub async fn request_device_code(
         let status = resp.status();
         let body = resp.text().await.unwrap_or_default();
         if status.as_u16() == 404 {
-            return Err(DeviceCodeError::NotEnabled);
+            anyhow::bail!(DeviceCodeError::NotEnabled);
         }
-        return Err(anyhow::anyhow!("Device code request failed (HTTP {status}): {body}").into());
+        anyhow::bail!("Device code request failed (HTTP {status}): {body}");
     }
 
     let server_resp: DeviceCodeResponse = resp.json().await?;
@@ -177,10 +174,7 @@ pub async fn request_device_code(
         .chars()
         .all(|c| c.is_ascii_alphanumeric() || c == '-')
     {
-        return Err(anyhow::anyhow!(
-            "Server returned invalid user_code format (expected [A-Z0-9-])"
-        )
-        .into());
+        anyhow::bail!("Server returned invalid user_code format (expected [A-Z0-9-])");
     }
 
     validate_verification_uri(&server_resp.verification_uri)?;

crates/codegen/xai-grok-shell/src/auth/flow.rs

@@ -7,7 +7,9 @@ use tokio::sync::{mpsc, oneshot};
 
 use crate::auth::config::LEGACY_AUTH_SCOPE;
 use crate::auth::{AuthManager, GrokAuth, GrokComConfig, parse_output};
+use crate::http::TransportFailureKind;
 use crate::util::grok_home;
+use xai_grok_telemetry::events::{LoginFailed, LoginFailureKind};
 
 pub type StderrCallback = Box<dyn Fn(&str)>;
 
@@ -451,6 +453,9 @@ pub async fn run_auth_flow_interactive(
     .await
 }
 
+/// Every interactive login returns through here, so reporting the failure here
+/// costs one event per attempt — a retried request, or the discovery cache
+/// background token refresh shares, can't inflate it. Never changes the result.
 async fn run_auth_flow_inner(
     auth_manager: &Arc<AuthManager>,
     grok_com_config: &GrokComConfig,
@@ -460,6 +465,64 @@ async fn run_auth_flow_inner(
     url_tx: Option<Rc<RefCell<Option<oneshot::Sender<AuthUrlInfo>>>>>,
     code_rx: Option<mpsc::Receiver<String>>,
     login_override: LoginTransportOverride,
+) -> anyhow::Result<(GrokAuth, bool)> {
+    let result = run_auth_flow_steps(
+        auth_manager,
+        grok_com_config,
+        reauth,
+        force_interactive,
+        on_stderr,
+        url_tx,
+        code_rx,
+        login_override,
+    )
+    .await;
+    if let Err(err) = &result
+        && let Some(event) = login_failure_event(err)
+    {
+        xai_grok_telemetry::session_ctx::log_event(event);
+    }
+    result
+}
+
+/// `None` when nothing in the chain failed over HTTP (the user backed out, the
+/// loopback listener couldn't bind, the id_token didn't validate) rather than
+/// inventing a transport verdict for it.
+fn login_failure_event(err: &anyhow::Error) -> Option<LoginFailed> {
+    let source = err
+        .chain()
+        .find_map(|cause| cause.downcast_ref::<reqwest::Error>())?;
+    Some(LoginFailed {
+        error_kind: failure_kind(
+            crate::http::TransportFailure::classify(source).kind,
+            source.is_decode(),
+        ),
+        os_error: crate::http::find_os_error_code(source),
+    })
+}
+
+/// A body that won't parse is a decode failure, not a transport one — even
+/// though `reqwest` also reports it as a body-phase error.
+fn failure_kind(transport: TransportFailureKind, is_decode: bool) -> LoginFailureKind {
+    if is_decode {
+        return LoginFailureKind::Decode;
+    }
+    match transport {
+        TransportFailureKind::Unreachable => LoginFailureKind::TransportConnect,
+        TransportFailureKind::Interrupted => LoginFailureKind::TransportInterrupted,
+        TransportFailureKind::Permanent => LoginFailureKind::TransportPermanent,
+    }
+}
+
+async fn run_auth_flow_steps(
+    auth_manager: &Arc<AuthManager>,
+    grok_com_config: &GrokComConfig,
+    reauth: bool,
+    force_interactive: bool,
+    on_stderr: Option<StderrCallback>,
+    url_tx: Option<Rc<RefCell<Option<oneshot::Sender<AuthUrlInfo>>>>>,
+    code_rx: Option<mpsc::Receiver<String>>,
+    login_override: LoginTransportOverride,
 ) -> anyhow::Result<(GrokAuth, bool)> {
     tracing::info!(
         has_oidc = grok_com_config.oidc.is_some(),
@@ -900,6 +963,41 @@ pub async fn run_cli_login(
     oauth: bool,
     device_auth: bool,
     devbox: bool,
+) -> anyhow::Result<()> {
+    // Devbox never reaches the login funnel, so it reports nothing and needs
+    // no telemetry client — and `AuthManager::new` is not free (it logs, and
+    // may rewrite auth.json to drop a stale scope).
+    if devbox {
+        let auth = super::devbox_login::run_devbox_login(config).await?;
+        return apply_post_login_config(auth).await;
+    }
+
+    // Agent bootstrap is what normally initializes the product telemetry
+    // client, and `grok login` never boots an agent, so without this every
+    // event this process emits is dropped before reaching a sink. One manager
+    // serves both the identity it reads and the login flow below.
+    let auth_manager = Arc::new(AuthManager::new(
+        &grok_home::grok_home(),
+        config.grok_com_config.clone(),
+    ));
+    crate::agent::init::update_telemetry_config(config, &auth_manager);
+
+    let result = run_cli_login_steps(config, &auth_manager, oauth, device_auth).await;
+
+    // Posts run on a spawned task and this process exits as soon as we return.
+    xai_grok_telemetry::session_ctx::drain_pending(CLI_TELEMETRY_DRAIN).await;
+    result
+}
+
+/// Returns as soon as the post lands (~1.7s cold), so the bound only bites on a
+/// black-holed network — where waiting out the HTTP client timeout would be worse.
+const CLI_TELEMETRY_DRAIN: std::time::Duration = std::time::Duration::from_secs(5);
+
+async fn run_cli_login_steps(
+    config: &crate::agent::config::Config,
+    auth_manager: &Arc<AuthManager>,
+    oauth: bool,
+    device_auth: bool,
 ) -> anyhow::Result<()> {
     let login_override = LoginTransportOverride::from_flags(oauth, device_auth);
 
@@ -908,15 +1006,11 @@ pub async fn run_cli_login(
     // supports the device flow. Without this guard, `grok login` on an
     // enterprise-OIDC deployment would wrongly enter the device branch (which
     // requires `oauth2`) and error.
-    let authenticated = if devbox {
-        super::devbox_login::run_devbox_login(config).await?
-    } else if cli_should_use_device(&config.grok_com_config, login_override).await {
+    let authenticated = if cli_should_use_device(&config.grok_com_config, login_override).await {
         if config.grok_com_config.oauth2.is_none() {
             // No OIDC and no oauth2 here, so `--oauth` can't help.
             anyhow::bail!("Sign-in is not available for this deployment. Set XAI_API_KEY instead.");
         }
-        let grok_home = grok_home::grok_home();
-        let auth_manager = Arc::new(AuthManager::new(&grok_home, config.grok_com_config.clone()));
         // Route through the shared inner flow (not `run_device_code_login`
         // directly) so the external auth provider and devbox auto-migration run
         // before the interactive device login. `force_interactive` skips the
@@ -925,7 +1019,7 @@ pub async fn run_cli_login(
         // Already resolved/logged above; pass `Preresolved(true)` so the inner flow
         // honors device without a second fetch or a duplicate `cli`-attributed log.
         let (auth, did_auth) = run_auth_flow_interactive(
-            &auth_manager,
+            auth_manager,
             &config.grok_com_config,
             None,
             None,
@@ -945,21 +1039,35 @@ pub async fn run_cli_login(
             );
         }
         // Loopback. `reauth=true` clears creds up front (legacy-scope hygiene),
-        // so abandoning logs you out — unlike the device branch above.
+        // so abandoning logs you out — unlike the device branch above. Calls
+        // `run_auth_flow` rather than `ensure_authenticated_with_override`,
+        // which would build a second `AuthManager`; with `reauth` set and no
+        // message prefix, the rest of that wrapper is a no-op.
         // Already resolved/logged above; pass `Preresolved(false)` so the inner
         // flow honors loopback without a duplicate `cli`-attributed log.
-        ensure_authenticated_with_override(
+        let (auth, did_auth) = run_auth_flow(
+            auth_manager,
             &config.grok_com_config,
             true,
             None,
+            None,
+            None,
             LoginTransportOverride::Preresolved(false),
         )
-        .await?
+        .await?;
+        if did_auth {
+            report_signed_in(&auth);
+        }
+        auth
     };
 
-    // Sync this principal's config now rather than waiting for the background
-    // tick. Stay quiet about absence/failure during login — confirm only when
-    // config was actually applied; `grok setup` reports the no-config case.
+    apply_post_login_config(authenticated).await
+}
+
+/// Sync this principal's config now rather than waiting for the background
+/// tick. Stay quiet about absence/failure during login — confirm only when
+/// config was actually applied; `grok setup` reports the no-config case.
+async fn apply_post_login_config(authenticated: GrokAuth) -> anyhow::Result<()> {
     let outcome = crate::managed_config::post_login_sync(Some(authenticated)).await;
     match outcome {
         crate::managed_config::ManagedConfigSync::Updated { is_team: true } => {
@@ -1070,6 +1178,42 @@ mod tests {
     use crate::env::EnvVarGuard;
     use chrono::Utc;
 
+    /// `os_error` and the reqwest classification are covered in
+    /// `xai-grok-http`; what's local is which `LoginFailureKind` each maps to,
+    /// and that a decode failure never reads as a transport one.
+    #[test]
+    fn failure_kinds_map_one_to_one() {
+        assert_eq!(
+            failure_kind(TransportFailureKind::Unreachable, false),
+            LoginFailureKind::TransportConnect
+        );
+        assert_eq!(
+            failure_kind(TransportFailureKind::Interrupted, false),
+            LoginFailureKind::TransportInterrupted
+        );
+        assert_eq!(
+            failure_kind(TransportFailureKind::Permanent, false),
+            LoginFailureKind::TransportPermanent
+        );
+        assert_eq!(
+            failure_kind(TransportFailureKind::Interrupted, true),
+            LoginFailureKind::Decode
+        );
+    }
+
+    /// A login that never reached the network is not a transport failure. The
+    /// positive path needs a real `reqwest::Error`, and building a client here
+    /// flips `jsonwebtoken` into its "no CryptoProvider" panic and breaks
+    /// unrelated auth tests, so classification is covered in `xai-grok-http`.
+    #[test]
+    fn non_http_login_failures_are_not_reported() {
+        let abandoned = anyhow::anyhow!("Login timed out after 10 minutes. Please try again.");
+        assert!(login_failure_event(&abandoned).is_none());
+
+        let nested = abandoned.context("Login failed. Please try again.");
+        assert!(login_failure_event(&nested).is_none());
+    }
+
     /// Run `f` with `GROK_LOGIN_DEVICE_FLOW` set to `value` (unset for `None`).
     /// `EnvVarGuard` serializes the process env and restores it on drop, so
     /// `resolve_device_flow` reads the env tier from a known state.

crates/codegen/xai-grok-shell/src/auth/manager.rs

@@ -18,7 +18,9 @@ pub(super) mod lock;
 mod sleep_gate;
 
 use lock::try_lock_auth_file_async;
-use sleep_gate::{GateRaise, InFlightGuard, SleepGate};
+use sleep_gate::{InFlightGuard, SleepGate};
+
+use crate::util::dual_clock::DualClock;
 
 use crate::auth::config::GrokComConfig;
 use crate::auth::error::AuthError;
@@ -98,13 +100,13 @@ const RELOAD_RETRY_BACKOFF: StdDuration = StdDuration::from_millis(50);
 struct ScopedRefreshFailure {
     token_key: String,
     error: crate::auth::error::RefreshTokenFailedError,
-    /// Two-clock timestamp (see [`GateRaise`]): the TTL below is *real* time,
+    /// Two-clock timestamp (see [`DualClock`]): the TTL below is *real* time,
     /// so it must keep counting across a system sleep. The monotonic clock
     /// pauses during suspend — with it alone, a failure cached just before
     /// sleep would still short-circuit `auth()` for a further
     /// [`PERMANENT_FAILURE_TTL`] of *awake* time after wake, exactly when the
     /// user comes back and expects a recovered session.
-    recorded_at: GateRaise,
+    recorded_at: DualClock,
 }
 
 /// Auto-expiry safety net for the recoverable reasons (`ClientRejected`,
@@ -202,11 +204,11 @@ pub struct AuthManager {
     /// manager so repeated 401s on the most-recent dead credential emit once.
     manual_auth: crate::auth::recovery::ManualAuthTracker,
     /// When the current unbroken run of dark-wake refresh deferrals began, on
-    /// two clocks (see [`GateRaise`]); `None` outside such a run. Bounds the
+    /// two clocks (see [`DualClock`]); `None` outside such a run. Bounds the
     /// deferral to [`sleep_gate::DARK_WAKE_DEFER_MAX`] so a machine stuck
     /// reporting dark wake can't defer refresh forever — see
     /// [`AuthManager::should_defer_for_dark_wake`].
-    dark_wake_defer_since: parking_lot::RwLock<Option<GateRaise>>,
+    dark_wake_defer_since: parking_lot::RwLock<Option<DualClock>>,
     /// Test-only override for [`AuthManager::is_dark_wake`]. `Some(_)` forces
     /// the dark-wake decision so the refresh-deferral path is unit-testable
     /// without a real macOS dark wake. `None` = consult the OS.
@@ -2085,7 +2087,7 @@ impl AuthManager {
         *self.permanent_failure.write() = Some(ScopedRefreshFailure {
             token_key,
             error,
-            recorded_at: GateRaise::now(),
+            recorded_at: DualClock::now(),
         });
     }
 
@@ -2117,7 +2119,7 @@ impl AuthManager {
     /// on disk) must be allowed to refresh — otherwise a hard-expired sibling
     /// AT strands a process that could still refresh a live RT.
     ///
-    /// TTL expiry is judged on *both* clocks (see [`GateRaise`]): the monotonic
+    /// TTL expiry is judged on *both* clocks (see [`DualClock`]): the monotonic
     /// clock pauses during a system suspend, so a wall-clock arm is required
     /// for the TTL to elapse across sleep. Without it, a recoverable failure
     /// cached just before the lid closes (e.g. a transient escalation while
@@ -2206,7 +2208,7 @@ impl AuthManager {
             // which the asserting test will surface loudly.
             let now_mono = std::time::Instant::now();
             let now_wall = std::time::SystemTime::now();
-            pf.recorded_at = GateRaise {
+            pf.recorded_at = DualClock {
                 mono: now_mono.checked_sub(past_ttl).unwrap_or(now_mono),
                 wall: now_wall.checked_sub(past_ttl).unwrap_or(now_wall),
             };

crates/codegen/xai-grok-shell/src/auth/manager/sleep_gate.rs

@@ -21,11 +21,12 @@
 
 use std::sync::Arc;
 use std::sync::atomic::Ordering;
-use std::time::{Duration as StdDuration, Instant, SystemTime};
+use std::time::{Duration as StdDuration, Instant};
 
 use parking_lot::RwLock;
 
 use super::AuthManager;
+use crate::util::dual_clock::DualClock;
 
 /// Max lifetime of the "system sleep imminent" gate. A wake event normally
 /// clears it; this is the safety bound so a *missed* wake event can never
@@ -40,7 +41,7 @@ pub(super) const SLEEP_GATE_MAX: StdDuration = StdDuration::from_secs(120);
 /// interactive Mac with no display, whose system video capability is never set
 /// — which would otherwise defer every refresh forever and reach the same
 /// logged-out state this guard prevents. Bounded on two clocks (see
-/// [`GateRaise`]) so it also survives the machine sleeping between dark wakes.
+/// [`DualClock`]) so it also survives the machine sleeping between dark wakes.
 ///
 /// The straddle risk of one forced refresh is far smaller than a guaranteed
 /// logout: requests only force through while the machine is busy enough to
@@ -65,57 +66,23 @@ pub(super) const SLEEP_ACK_MAX_WAIT: StdDuration = StdDuration::from_secs(20);
 #[cfg(not(target_os = "macos"))]
 pub(super) const SLEEP_ACK_MAX_WAIT: StdDuration = StdDuration::from_secs(3);
 
-/// When a gate was raised, captured on *two* clocks so the [`SLEEP_GATE_MAX`]
-/// backstop survives a system sleep.
-///
-/// `Instant` is monotonic but, on macOS (`mach_absolute_time`) and Linux
-/// (`CLOCK_MONOTONIC`), *pauses while the machine is asleep*. A gate raised just
-/// before a long sleep would therefore never auto-expire on the monotonic clock
-/// alone — the exact bug that let an expired token reach the server and 401.
-/// The wall clock (`SystemTime`) keeps advancing through sleep, so we expire the
-/// gate once *either* clock passes the bound:
-/// - the monotonic clock bounds elapsed *awake* time (immune to wall-clock
-///   jumps from NTP / manual changes), and
-/// - the wall clock bounds elapsed *real* time (immune to the sleep pause).
-#[derive(Clone, Copy)]
-pub(super) struct GateRaise {
-    /// Monotonic; pauses during sleep. Bounds elapsed *awake* time.
-    pub(super) mono: Instant,
-    /// Wall clock; advances through sleep. Bounds elapsed *real* time.
-    pub(super) wall: SystemTime,
-}
-
-impl GateRaise {
-    pub(super) fn now() -> Self {
-        Self {
-            mono: Instant::now(),
-            wall: SystemTime::now(),
-        }
-    }
-
-    /// Elapsed on each clock as `(monotonic, wall)`. Wall-clock elapsed is
-    /// clamped to zero if the clock ran backwards (NTP step / manual change) so
-    /// a backward jump can never *extend* the gate — the monotonic clock still
-    /// bounds it in that case.
-    pub(super) fn elapsed(&self) -> (StdDuration, StdDuration) {
-        (
-            self.mono.elapsed(),
-            self.wall.elapsed().unwrap_or(StdDuration::ZERO),
-        )
-    }
-}
-
 /// A gate `refresh_chain` consults to avoid *starting* an IdP refresh just
 /// before sleep. Only *defers* a not-yet-started refresh; an in-flight one is
 /// left to finish (see [`AuthManager::refresh_chain`]).
+///
+/// The raise timestamp is a [`DualClock`] so the [`SLEEP_GATE_MAX`] backstop
+/// survives the sleep itself: a gate raised just before a long sleep would
+/// never auto-expire on the monotonic clock alone — the exact bug that let
+/// an expired token reach the server and 401 — so the gate expires once
+/// *either* clock passes the bound.
 #[derive(Default)]
 pub(super) struct SleepGate {
-    pub(super) raised_at: RwLock<Option<GateRaise>>,
+    pub(super) raised_at: RwLock<Option<DualClock>>,
 }
 
 impl SleepGate {
     pub(super) fn raise(&self) {
-        *self.raised_at.write() = Some(GateRaise::now());
+        *self.raised_at.write() = Some(DualClock::now());
         xai_grok_telemetry::unified_log::warn("auth.sleep.gate_set", None, None);
     }
 
@@ -142,7 +109,7 @@ impl SleepGate {
     /// A stale gate (a missed/late wake event) is lazily lowered here so it can
     /// never permanently block refresh; this read can therefore have a side
     /// effect. The gate expires once *either* clock passes [`SLEEP_GATE_MAX`]
-    /// (see [`GateRaise`]): without the wall-clock arm, a gate raised before a
+    /// (see [`DualClock`]): without the wall-clock arm, a gate raised before a
     /// long sleep would never auto-expire, because the monotonic clock pauses
     /// while the machine is asleep.
     pub(super) fn is_gated(&self) -> bool {
@@ -355,7 +322,7 @@ impl AuthManager {
     /// in a dark wake — bounded so deferral can never be indefinite.
     ///
     /// Tracks when the current unbroken run of dark-wake deferrals began (on two
-    /// clocks; see [`GateRaise`]). While inside the [`DARK_WAKE_DEFER_MAX`]
+    /// clocks; see [`DualClock`]). While inside the [`DARK_WAKE_DEFER_MAX`]
     /// budget it returns `true` (defer). Once either clock passes the bound it
     /// forces one refresh through (`false`) and resets the clock, so a machine
     /// stuck reporting a continuous dark wake refreshes periodically instead of
@@ -375,7 +342,7 @@ impl AuthManager {
         }
         let Some(raise) = *run else {
             // First deferral of this dark-wake run: start the budget clock.
-            *run = Some(GateRaise::now());
+            *run = Some(DualClock::now());
             return true;
         };
         let (mono, wall) = raise.elapsed();

crates/codegen/xai-grok-shell/src/auth/manager_tests.rs

@@ -4353,7 +4353,7 @@ async fn dark_wake_defer_forces_refresh_after_max() {
     ) else {
         return; // machine/clock can't represent the backdate — skip
     };
-    *mgr.dark_wake_defer_since.write() = Some(super::sleep_gate::GateRaise { mono, wall });
+    *mgr.dark_wake_defer_since.write() = Some(crate::util::dual_clock::DualClock { mono, wall });
 
     assert_eq!(
         mgr.auth().await.unwrap().key,
@@ -4490,7 +4490,7 @@ async fn sleep_gate_auto_expires_after_max() {
     ) else {
         return; // machine/clock can't represent the backdate — not reproducible; skip
     };
-    *mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise { mono, wall });
+    *mgr.sleep_gate.raised_at.write() = Some(crate::util::dual_clock::DualClock { mono, wall });
 
     assert!(
         !mgr.is_sleep_gated(),
@@ -4522,7 +4522,7 @@ async fn sleep_gate_auto_expires_when_wall_clock_passes_during_sleep() {
     let Some(wall) = std::time::SystemTime::now().checked_sub(back) else {
         return; // clock can't represent the backdate — not reproducible; skip
     };
-    *mgr.sleep_gate.raised_at.write() = Some(super::sleep_gate::GateRaise {
+    *mgr.sleep_gate.raised_at.write() = Some(crate::util::dual_clock::DualClock {
         mono: Instant::now(),
         wall,
     });

crates/codegen/xai-grok-shell/src/extensions/feedback.rs

@@ -18,7 +18,7 @@ use crate::agent::MvpAgent;
 use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry};
 use crate::session::{
     ClientFeedbackInput, CommentDeleteRequest, CommentDeleteResponse, CommentRequest,
-    CommentResponse, FeedbackRequestDismiss, FeedbackResponse, SessionCommand,
+    CommentResponse, FeedbackRequestDismiss, FeedbackResponse, SessionCommand, SideQuestionError,
 };
 use crate::upload::gcs::WithAuth as _;
 use xai_file_utils::gcs::upload_bytes;
@@ -75,7 +75,20 @@ async fn handle_btw(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
         Ok(answer) => super::to_ext_response(Ok(serde_json::json!({
             "answer": answer,
         }))),
-        Err(e) => Err(acp::Error::internal_error().data(e)),
+        // Model errors take the canonical mapping: overload gets its short
+        // display copy there, rate limits keep the typed code + upgrade
+        // copy, auth failures surface as auth_required.
+        Err(SideQuestionError::Sampling(e)) => {
+            Err(crate::sampling::error::map_sampling_err_to_acp(e))
+        }
+        // Non-model failures are already readable sentences. Set `message`
+        // and leave `data` unset — `Display` appends JSON-encoded `data`,
+        // and `internal_error().data(e)` rendered as `Internal error: "…"`,
+        // which made capacity failures look like client bugs in the TUI.
+        Err(e) => Err(acp::Error::new(
+            acp::ErrorCode::InternalError.into(),
+            e.to_string(),
+        )),
     }
 }
 

crates/codegen/xai-grok-shell/src/extensions/notification.rs

@@ -418,6 +418,25 @@ pub struct HookRunEntryDto {
     pub output: Option<String>,
 }
 
+/// Why auto-compaction stopped before completing.
+#[derive(
+    Debug,
+    Clone,
+    Copy,
+    PartialEq,
+    Eq,
+    serde::Serialize,
+    serde::Deserialize,
+    strum::Display,
+    strum::EnumString,
+    strum::AsRefStr,
+)]
+#[serde(rename_all = "snake_case")]
+#[strum(serialize_all = "snake_case")]
+pub enum AutoCompactCancelReason {
+    UserCancelled,
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq)]
 #[serde(rename_all = "snake_case", tag = "sessionUpdate")]
 pub enum SessionUpdate {
@@ -483,7 +502,7 @@ pub enum SessionUpdate {
     /// Auto-compact was cancelled (user pressed Ctrl+C)
     AutoCompactCancelled {
         /// Reason for cancellation
-        reason: String,
+        reason: AutoCompactCancelReason,
     },
     /// Auto-continue completed after compaction
     /// This signals the TUI to flush pending agent messages and end the turn
@@ -1704,6 +1723,16 @@ mod tests {
         let update: SessionUpdate = serde_json::from_str(json).unwrap();
         assert_eq!(update, SessionUpdate::MemoryFlushStarted);
 
+        // AutoCompactCancelled (strenum reason)
+        let json = r#"{"sessionUpdate": "auto_compact_cancelled", "reason": "user_cancelled"}"#;
+        let update: SessionUpdate = serde_json::from_str(json).unwrap();
+        assert_eq!(
+            update,
+            SessionUpdate::AutoCompactCancelled {
+                reason: AutoCompactCancelReason::UserCancelled,
+            }
+        );
+
         // AutoCompactFailed (struct variant)
         let json = r#"{"sessionUpdate": "auto_compact_failed", "error": "oom"}"#;
         let update: SessionUpdate = serde_json::from_str(json).unwrap();

crates/codegen/xai-grok-shell/src/extensions/session_admin.rs

@@ -7,6 +7,7 @@
 //! - `x.ai/session/rename`                  rename a session locally + remote
 //! - `x.ai/session/delete`                  delete a session locally + remote
 //! - `x.ai/session/update_mcp_servers`      mid-session MCP server swap
+//! - `x.ai/session/add_local_workspace`     mid-session local workspace add-only (chat)
 //! - `x.ai/session/fork`                    fork a session into a new one
 //! - `x.ai/internal/reload_all_mcp_servers` config hot-reload, all sessions
 //! - `x.ai/internal/reload_project_mcp_servers` config hot-reload, cwd-scoped
@@ -39,6 +40,8 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
         "x.ai/session/rename" => handle_session_rename(agent, args).await,
         "x.ai/session/delete" => handle_session_delete(agent, args).await,
         "x.ai/session/update_mcp_servers" => handle_update_mcp_servers(agent, args).await,
+        #[cfg(feature = "local-workspace")]
+        "x.ai/session/add_local_workspace" => handle_add_local_workspace(agent, args).await,
         "x.ai/session/fork" => handle_session_fork(agent, args).await,
         "x.ai/internal/reload_all_mcp_servers" => handle_reload_all_mcp_servers(agent).await,
         "x.ai/internal/reload_project_mcp_servers" => {
@@ -371,6 +374,43 @@ async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) ->
         .map_err(|e| acp::Error::internal_error().data(e.to_string()))
 }
 
+// session/add_local_workspace (add-only; local-workspace feature)
+
+#[cfg(feature = "local-workspace")]
+async fn handle_add_local_workspace(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
+    #[derive(Deserialize)]
+    #[serde(rename_all = "camelCase")]
+    struct Params {
+        session_id: acp::SessionId,
+        #[serde(default)]
+        meta: Option<acp::Meta>,
+    }
+
+    let params: Params = parse_params(args)?;
+    let cwd = {
+        let sessions = agent.sessions.borrow();
+        let h = sessions
+            .get(&params.session_id)
+            .ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?;
+        std::path::PathBuf::from(&h.info.cwd)
+    };
+    // Gate on actual chat kind — not `requires_gateway` (true for non-chat
+    // GatewayAttach; false for unknown ids).
+    if !agent.is_chat_kind_session(&params.session_id) {
+        return Err(acp::Error::invalid_params().data(serde_json::json!({
+            "code": "local_workspace_chat_only",
+            "message": "x.ai/session/add_local_workspace is only available on chat-kind sessions",
+        })));
+    }
+
+    let result = agent
+        .add_local_workspace_mid_session(&params.session_id, params.meta, &cwd)
+        .await?;
+    ExtMethodResult::success(result)
+        .to_ext_response()
+        .map_err(|e| acp::Error::internal_error().data(e.to_string()))
+}
+
 // internal/reload_skills
 
 /// Reload skills for ALL active sessions. Called by the skills file watcher

crates/codegen/xai-grok-shell/src/sampling/error.rs

@@ -96,12 +96,25 @@ fn pushes_consumer_subscription_upsell(detail: &str) -> bool {
     d.contains("grok.com/supergrok") || d.contains("upgrade to a grok subscription")
 }
 
+/// User-facing copy for capacity/overload failures (stream `overloaded_error`,
+/// HTTP 529, proxy-wrapped 5xx). See [`SamplingError::is_overloaded`].
+pub const OVERLOADED_USER_MESSAGE: &str = "Model is temporarily overloaded. Try again in a moment.";
+
 /// Map a `SamplingError` to an ACP `Error` for client-facing responses.
 /// This stays in xai-grok-shell because it depends on `agent_client_protocol::Error`.
 pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
     use reqwest::StatusCode;
+    // Capacity/overload gets the same short copy on every surface. Message
+    // only, `data` deliberately unset: `Display` appends JSON-encoded `data`,
+    // and this string is meant for direct display.
+    if err.is_overloaded() {
+        return acp::Error::new(
+            acp::ErrorCode::InternalError.into(),
+            OVERLOADED_USER_MESSAGE,
+        );
+    }
     match err {
-        SamplingError::Auth(msg) => acp::Error::auth_required().data(msg),
+        SamplingError::Auth { message, .. } => acp::Error::auth_required().data(message),
         SamplingError::InvalidConfiguration(msg) => acp::Error::invalid_params().data(msg),
         SamplingError::Http(e) => {
             acp::Error::internal_error().data(format!("http client init failed: {e}"))
@@ -489,6 +502,31 @@ mod tests {
         );
     }
 
+    #[test]
+    fn overload_maps_to_display_message_without_data() {
+        let err = SamplingError::StreamError {
+            error_type: "overloaded_error".into(),
+            message: "Overloaded".into(),
+        };
+        let acp_err = map_sampling_err_to_acp(err);
+        assert_eq!(acp_err.code, acp::ErrorCode::InternalError);
+        assert_eq!(acp_err.message, OVERLOADED_USER_MESSAGE);
+        // Display appends JSON-encoded `data`; direct-display copy must not
+        // carry any.
+        assert_eq!(acp_err.data, None);
+
+        let err_529 = SamplingError::Api {
+            status: StatusCode::from_u16(529).expect("valid status"),
+            message: "capacity".into(),
+            model_metadata: None,
+            retry_after_secs: None,
+            should_retry: None,
+        };
+        let acp_529 = map_sampling_err_to_acp(err_529);
+        assert_eq!(acp_529.message, OVERLOADED_USER_MESSAGE);
+        assert_eq!(acp_529.data, None);
+    }
+
     #[test]
     fn rate_limit_error_uses_dedicated_code() {
         let err = SamplingError::Api {

crates/codegen/xai-grok-shell/src/session/acp_session.rs

@@ -63,7 +63,6 @@ use std::sync::Arc;
 use std::sync::OnceLock;
 use tokio::sync::{Mutex as TokioMutex, mpsc, oneshot};
 use tokio::time::{Duration, sleep};
-use tokio_retry::strategy::ExponentialBackoff;
 use xai_acp_lib::AcpAgentGatewaySender as GatewaySender;
 use xai_grok_agent::AgentDefinition;
 use xai_grok_agent::prompt::agents_md::LEGACY_AGENTS_MD_REMINDER_PREFIX;
@@ -92,16 +91,21 @@ mod compaction_segments;
 mod types;
 pub(crate) use types::*;
 pub use types::{TodoGateDecision, TodoGateReason};
+#[path = "acp_session_impl/auth_retry.rs"]
+mod auth_retry;
 #[path = "acp_session_impl/goal.rs"]
 mod goal;
-#[path = "acp_session_impl/interjection.rs"]
-mod interjection;
-#[path = "acp_session_impl/tool_calls.rs"]
-mod tool_calls;
 #[path = "acp_session_impl/turn.rs"]
 mod turn;
 #[path = "acp_session_impl/workflow.rs"]
 mod workflow_run;
+pub(crate) use auth_retry::{
+    AuthRetryDecision, AuthRetrySchedule, human_duration, pace_uncharged_resubmit,
+};
+#[path = "acp_session_impl/interjection.rs"]
+mod interjection;
+#[path = "acp_session_impl/tool_calls.rs"]
+mod tool_calls;
 pub(crate) use interjection::*;
 #[path = "acp_session_impl/laziness.rs"]
 mod laziness;
@@ -1809,6 +1813,9 @@ impl Drop for TurnMetrics {
 #[cfg(test)]
 #[path = "acp_session_tests/auth_error_no_retry_tests.rs"]
 mod auth_error_no_retry_tests;
+#[cfg(test)]
+#[path = "acp_session_tests/turn/auth_retry_budget_tests.rs"]
+mod auth_retry_budget_tests;
 /// Regression coverage for the auto-wake suppression sweep + shutdown
 /// drain. These exercise the helpers added to fix the trailing
 /// `<system-reminder>` chat history bug.

crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry.rs

@@ -0,0 +1,218 @@
+//! Per-turn retry policy for 401s that follow a *successful* auth recovery
+//! (fresh token minted, request to be re-sent).
+
+use tokio_retry::strategy::ExponentialBackoff;
+use xai_grok_sampling_types::SentCredential;
+
+use super::RecoveredStore;
+use crate::auth::AuthManager;
+use crate::util::dual_clock::DualClock;
+
+/// Pace an uncharged resubmit: wait (bounded) for a session-token refresh
+/// when that is the store the recovery minted into and nothing wire-valid
+/// has landed yet; otherwise floor-pace so the runaway guard can never be a
+/// burst of back-to-back requests. Auth policy lives here, not in the turn
+/// loop.
+pub(crate) async fn pace_uncharged_resubmit(
+    store: RecoveredStore,
+    auth_manager: Option<&std::sync::Arc<AuthManager>>,
+) {
+    match (store, auth_manager) {
+        (RecoveredStore::SessionToken, Some(am)) if am.current_wire_valid().is_none() => {
+            am.wait_for_token_refresh(AuthRetrySchedule::UNCHARGED_REFRESH_WAIT)
+                .await;
+        }
+        _ => tokio::time::sleep(AuthRetrySchedule::UNCHARGED_RESUBMIT_FLOOR).await,
+    }
+}
+
+/// Compact `2h3m` / `4m7s` / `12s` rendering for turn-failure messages.
+pub(crate) fn human_duration(d: std::time::Duration) -> String {
+    let total_secs = d.as_secs();
+    if total_secs < 60 {
+        return format!("{total_secs}s");
+    }
+    let mins = total_secs / 60;
+    if mins < 60 {
+        return format!("{mins}m{}s", total_secs % 60);
+    }
+    format!("{}h{}m", mins / 60, mins % 60)
+}
+
+/// Decision for one post-recovery 401 (see
+/// [`AuthRetrySchedule::on_recovered_401`]).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum AuthRetryDecision {
+    /// No credential was on the wire, so no slot is charged; resubmit after
+    /// the refresh lands. `resubmit` is the 1-indexed count since the last
+    /// successful response.
+    UnchargedResubmit { resubmit: u32 },
+    /// Charged one escalating slot: back off `delay`, then resubmit.
+    Backoff {
+        attempt: u32,
+        delay: std::time::Duration,
+    },
+    /// Per-incident budget exhausted by credentialed 401s — fail the turn.
+    Exhausted,
+    /// Runaway guard tripped: recovery kept succeeding while the server
+    /// rejected `rejections` credential-less requests without a single
+    /// successful response — fail the turn.
+    RunawayGuard { rejections: u32 },
+}
+
+/// Escalating retry budget for post-recovery 401s. The budget is
+/// per-incident (successes and suspend boundaries reset it, the latter
+/// capped) and only credentialed rejections charge it; the doc on each
+/// method carries its own invariant. One thing no reader can derive from
+/// the code:
+///
+/// **Delays must be 1s/2s/4s.** `ExponentialBackoff::from_millis(base)`
+/// raises `base` to the attempt number, so the base must stay small:
+/// `from_millis(1000)` yields 1000ⁿ ms = 1s → 16m40s → 11.57 days of
+/// silent hang (a past field incident). `from_millis(2).factor(500)`
+/// yields 1s, 2s, 4s.
+pub(crate) struct AuthRetrySchedule {
+    delays: std::iter::Take<ExponentialBackoff>,
+    /// Slots charged this incident.
+    attempt: u32,
+    /// 401s seen this incident, total and the subset that provably carried
+    /// a credential. Feeds the exhaustion message so "real credential
+    /// rejected" and "budget exhausted" cannot be conflated.
+    incident_rejections: u32,
+    incident_authenticated: u32,
+    /// Stamped by the incident's first charged 401; cleared by resets.
+    incident_started: Option<DualClock>,
+    /// Uncharged fail-closed rejections since the last successful response
+    /// (survives suspend resets).
+    uncharged_resubmits: u32,
+    /// Suspend-triggered resets since the last successful response.
+    suspend_resets: u32,
+}
+
+impl AuthRetrySchedule {
+    /// Consecutive credentialed post-recovery 401s tolerated per incident
+    /// before the turn fails.
+    pub(crate) const MAX_RETRIES: u32 = 3;
+    /// Runaway guard: uncharged (no-credential) rejections tolerated
+    /// without an intervening successful response (~one per 16-minute
+    /// sleep cycle ⇒ >13 h lid-closed survival).
+    pub(crate) const MAX_UNCHARGED_RESUBMITS: u32 = 50;
+    /// Suspend resets tolerated without an intervening successful response
+    /// (~8 sleep cycles of a continuously failing incident) before the
+    /// budget stops resetting and is allowed to exhaust.
+    pub(crate) const MAX_SUSPEND_RESETS: u32 = 8;
+    /// Bounded wait for the proactive refresh / wake nudge to land a
+    /// wire-valid token before an uncharged resubmit.
+    const UNCHARGED_REFRESH_WAIT: std::time::Duration = std::time::Duration::from_secs(15);
+    /// Floor pacing for uncharged resubmits with no refresh to wait on.
+    const UNCHARGED_RESUBMIT_FLOOR: std::time::Duration = std::time::Duration::from_secs(1);
+    /// Wall-vs-monotonic drift beyond which the machine must have slept:
+    /// well below a real sleep cycle (minutes), well above NTP step jitter.
+    const SUSPEND_DRIFT_MIN: std::time::Duration = std::time::Duration::from_secs(30);
+
+    pub(crate) fn new() -> Self {
+        Self {
+            delays: ExponentialBackoff::from_millis(2)
+                .factor(500)
+                .max_delay(std::time::Duration::from_secs(10))
+                .take(Self::MAX_RETRIES as usize),
+            attempt: 0,
+            incident_rejections: 0,
+            incident_authenticated: 0,
+            incident_started: None,
+            uncharged_resubmits: 0,
+            suspend_resets: 0,
+        }
+    }
+
+    /// Decision for one post-recovery 401. Charges a slot only when the
+    /// rejected request carried a credential (or its provenance is unknown
+    /// — fail closed toward terminating).
+    pub(crate) fn on_recovered_401(&mut self, credential: SentCredential) -> AuthRetryDecision {
+        self.on_recovered_401_at(credential, DualClock::now())
+    }
+
+    /// Clock-injected twin of [`Self::on_recovered_401`] for tests.
+    fn on_recovered_401_at(
+        &mut self,
+        credential: SentCredential,
+        now: DualClock,
+    ) -> AuthRetryDecision {
+        if credential.is_missing() {
+            self.uncharged_resubmits += 1;
+            if self.uncharged_resubmits > Self::MAX_UNCHARGED_RESUBMITS {
+                return AuthRetryDecision::RunawayGuard {
+                    rejections: self.uncharged_resubmits,
+                };
+            }
+            return AuthRetryDecision::UnchargedResubmit {
+                resubmit: self.uncharged_resubmits,
+            };
+        }
+        self.incident_started.get_or_insert(now);
+        self.incident_rejections += 1;
+        if credential == SentCredential::Sent {
+            self.incident_authenticated += 1;
+        }
+        match self.delays.next() {
+            Some(delay) => {
+                self.attempt += 1;
+                AuthRetryDecision::Backoff {
+                    attempt: self.attempt,
+                    delay,
+                }
+            }
+            None => AuthRetryDecision::Exhausted,
+        }
+    }
+
+    /// Close the open incident if it spans a suspend (wall elapsed outgrew
+    /// monotonic elapsed by [`Self::SUSPEND_DRIFT_MIN`]): separate wakes are
+    /// independent 401 events. Capped at [`Self::MAX_SUSPEND_RESETS`] per
+    /// success-free stretch so a fault that persists across wakes exhausts
+    /// instead of retrying forever. Returns whether a reset happened.
+    pub(crate) fn reset_if_incident_spans_suspend(&mut self) -> bool {
+        self.reset_if_incident_spans_suspend_at(DualClock::now())
+    }
+
+    /// Clock-injected twin of [`Self::reset_if_incident_spans_suspend`].
+    fn reset_if_incident_spans_suspend_at(&mut self, now: DualClock) -> bool {
+        let Some(started) = self.incident_started else {
+            return false;
+        };
+        if self.suspend_resets >= Self::MAX_SUSPEND_RESETS {
+            return false;
+        }
+        let (awake, total) = started.elapsed_between(now);
+        if total.saturating_sub(awake) < Self::SUSPEND_DRIFT_MIN {
+            return false;
+        }
+        let (uncharged, resets) = (self.uncharged_resubmits, self.suspend_resets);
+        *self = Self::new();
+        self.uncharged_resubmits = uncharged;
+        self.suspend_resets = resets + 1;
+        true
+    }
+
+    /// A successful model response ends every open failure narrative:
+    /// restart the escalating schedule and clear the success-free-stretch
+    /// counters (uncharged rejections, suspend resets).
+    pub(crate) fn reset_on_success(&mut self) {
+        *self = Self::new();
+    }
+
+    /// `(rejections, authenticated)` seen this incident, for the exhaustion
+    /// message.
+    pub(crate) fn incident_counts(&self) -> (u32, u32) {
+        (self.incident_rejections, self.incident_authenticated)
+    }
+
+    /// Uncharged fail-closed rejections since the last successful response.
+    pub(crate) fn uncharged_rejections(&self) -> u32 {
+        self.uncharged_resubmits
+    }
+}
+
+#[cfg(test)]
+#[path = "auth_retry_tests.rs"]
+mod tests;

crates/codegen/xai-grok-shell/src/session/acp_session_impl/auth_retry_tests.rs

@@ -0,0 +1,197 @@
+use std::time::Duration;
+
+use xai_grok_sampling_types::SentCredential;
+
+use super::{AuthRetryDecision, AuthRetrySchedule};
+use crate::util::dual_clock::DualClock;
+
+/// `now` shifted `wall_ahead` on the wall clock only — the signature a
+/// suspend leaves behind (monotonic pauses, wall keeps advancing).
+fn after_suspend(base: DualClock, wall_ahead: Duration) -> DualClock {
+    DualClock {
+        mono: base.mono,
+        wall: base.wall + wall_ahead,
+    }
+}
+
+/// Pins the exact schedule. Guards against the `from_millis(1000)` footgun
+/// (baseⁿ semantics), which produced field sleeps of 1s, 16m40s, and 11.57
+/// days.
+#[test]
+fn schedule_is_one_two_four_seconds_then_exhausted() {
+    let mut schedule = AuthRetrySchedule::new();
+    let steps: Vec<_> = (0..3)
+        .map(|_| schedule.on_recovered_401(SentCredential::Sent))
+        .collect();
+    assert_eq!(
+        steps,
+        vec![
+            AuthRetryDecision::Backoff {
+                attempt: 1,
+                delay: Duration::from_secs(1)
+            },
+            AuthRetryDecision::Backoff {
+                attempt: 2,
+                delay: Duration::from_secs(2)
+            },
+            AuthRetryDecision::Backoff {
+                attempt: 3,
+                delay: Duration::from_secs(4)
+            },
+        ],
+    );
+    assert_eq!(
+        schedule.on_recovered_401(SentCredential::Sent),
+        AuthRetryDecision::Exhausted,
+    );
+    assert_eq!(schedule.incident_counts(), (4, 4));
+}
+
+/// Unknown provenance charges like an authenticated 401 (fail closed toward
+/// terminating) but is not reported as a proven credential rejection.
+#[test]
+fn unknown_credential_charges_but_is_not_counted_authenticated() {
+    let mut schedule = AuthRetrySchedule::new();
+    assert_eq!(
+        schedule.on_recovered_401(SentCredential::Unknown),
+        AuthRetryDecision::Backoff {
+            attempt: 1,
+            delay: Duration::from_secs(1)
+        },
+    );
+    assert_eq!(schedule.incident_counts(), (1, 0));
+}
+
+/// The overnight-failure regression: a credential-less 401 never consumes a
+/// budget slot; only the runaway guard bounds it.
+#[test]
+fn missing_credential_never_charges_until_runaway_guard() {
+    let mut schedule = AuthRetrySchedule::new();
+    for i in 1..=AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS {
+        assert_eq!(
+            schedule.on_recovered_401(SentCredential::Missing),
+            AuthRetryDecision::UnchargedResubmit { resubmit: i },
+        );
+    }
+    assert_eq!(
+        schedule.on_recovered_401(SentCredential::Missing),
+        AuthRetryDecision::RunawayGuard {
+            rejections: AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS + 1
+        },
+    );
+    assert_eq!(
+        schedule.on_recovered_401(SentCredential::Sent),
+        AuthRetryDecision::Backoff {
+            attempt: 1,
+            delay: Duration::from_secs(1)
+        },
+        "the credentialed budget must be untouched throughout"
+    );
+}
+
+/// A success ends every open failure narrative: the escalating delays, the
+/// attempt numbering, and the runaway counter all restart (a 200 disproves
+/// the runaway premise, so a productive multi-day turn can never accumulate
+/// into the guard).
+#[test]
+fn success_resets_budget_and_uncharged_counter() {
+    let mut schedule = AuthRetrySchedule::new();
+    schedule.on_recovered_401(SentCredential::Sent);
+    schedule.on_recovered_401(SentCredential::Sent);
+    for _ in 0..AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS {
+        schedule.on_recovered_401(SentCredential::Missing);
+    }
+    schedule.reset_on_success();
+    assert_eq!(
+        schedule.on_recovered_401(SentCredential::Missing),
+        AuthRetryDecision::UnchargedResubmit { resubmit: 1 },
+    );
+    assert_eq!(
+        schedule.on_recovered_401(SentCredential::Sent),
+        AuthRetryDecision::Backoff {
+            attempt: 1,
+            delay: Duration::from_secs(1)
+        },
+    );
+}
+
+/// The uncharged counter survives a suspend reset (the guard spans sleep
+/// cycles — that is its point) while the charged budget restarts.
+#[test]
+fn suspend_reset_preserves_uncharged_counter() {
+    let mut schedule = AuthRetrySchedule::new();
+    let start = DualClock::now();
+    schedule.on_recovered_401_at(SentCredential::Missing, start);
+    schedule.on_recovered_401_at(SentCredential::Missing, start);
+    schedule.on_recovered_401_at(SentCredential::Sent, start);
+
+    let woke = after_suspend(start, Duration::from_secs(16 * 60));
+    assert!(schedule.reset_if_incident_spans_suspend_at(woke));
+    assert_eq!(
+        schedule.on_recovered_401_at(SentCredential::Missing, woke),
+        AuthRetryDecision::UnchargedResubmit { resubmit: 3 },
+    );
+    assert_eq!(
+        schedule.on_recovered_401_at(SentCredential::Sent, woke),
+        AuthRetryDecision::Backoff {
+            attempt: 1,
+            delay: Duration::from_secs(1)
+        },
+        "post-suspend 401 starts a fresh incident instead of exhausting"
+    );
+}
+
+/// Suspend resets are capped per success-free stretch: a fault that
+/// persists across wakes must eventually exhaust instead of retrying
+/// forever. A success re-arms the cap.
+#[test]
+fn suspend_resets_cap_without_success_and_rearm_on_success() {
+    let mut schedule = AuthRetrySchedule::new();
+    let mut now = DualClock::now();
+    for _ in 0..AuthRetrySchedule::MAX_SUSPEND_RESETS {
+        schedule.on_recovered_401_at(SentCredential::Sent, now);
+        now = after_suspend(now, Duration::from_secs(16 * 60));
+        assert!(schedule.reset_if_incident_spans_suspend_at(now));
+    }
+    schedule.on_recovered_401_at(SentCredential::Sent, now);
+    now = after_suspend(now, Duration::from_secs(16 * 60));
+    assert!(
+        !schedule.reset_if_incident_spans_suspend_at(now),
+        "reset {} must be refused: the budget is now allowed to exhaust",
+        AuthRetrySchedule::MAX_SUSPEND_RESETS + 1
+    );
+
+    schedule.reset_on_success();
+    schedule.on_recovered_401_at(SentCredential::Sent, now);
+    now = after_suspend(now, Duration::from_secs(16 * 60));
+    assert!(
+        schedule.reset_if_incident_spans_suspend_at(now),
+        "a success re-arms the suspend-reset cap"
+    );
+}
+
+/// No suspend, no reset: sub-threshold wall drift (NTP jitter) and a
+/// schedule with no open incident are both no-ops.
+#[test]
+fn suspend_reset_requires_open_incident_and_real_drift() {
+    let mut schedule = AuthRetrySchedule::new();
+    let start = DualClock::now();
+    assert!(
+        !schedule
+            .reset_if_incident_spans_suspend_at(after_suspend(start, Duration::from_secs(3600))),
+        "no open incident: nothing to reset"
+    );
+    schedule.on_recovered_401_at(SentCredential::Sent, start);
+    assert!(
+        !schedule.reset_if_incident_spans_suspend_at(after_suspend(start, Duration::from_secs(5))),
+        "5s wall drift is NTP-jitter territory, not a suspend"
+    );
+    assert_eq!(
+        schedule.on_recovered_401_at(SentCredential::Sent, start),
+        AuthRetryDecision::Backoff {
+            attempt: 2,
+            delay: Duration::from_secs(2)
+        },
+        "the failed reset checks must not charge the budget"
+    );
+}

crates/codegen/xai-grok-shell/src/session/acp_session_impl/recap.rs

@@ -4,6 +4,37 @@
 use super::*;
 
 use crate::remote::DEFAULT_CONTEXT_WINDOW;
+use crate::session::SideQuestionError;
+use xai_grok_sampling_types::SamplingError;
+
+/// Retry policy for the one-shot `/btw` model call: 3 attempts total
+/// (1 try + 2 retries), 500ms → 1s jittered backoff. Deliberately short —
+/// nothing like the sampler actor's budget — so a fleet-wide capacity event
+/// can't multiply side-question traffic into a retry storm.
+fn side_question_retry_policy() -> backon::ExponentialBuilder {
+    backon::ExponentialBuilder::default()
+        .with_max_times(2)
+        .with_min_delay(std::time::Duration::from_millis(500))
+        .with_max_delay(std::time::Duration::from_secs(1))
+        .with_jitter()
+}
+
+/// Whether a failed `/btw` attempt is worth retrying: overload only (not
+/// every retryable 5xx / stream glitch), minus the shared retry vetoes
+/// (`x-should-retry: false`, context length — see
+/// [`SamplingError::is_retry_vetoed`], also enforced by the sampler actor's
+/// `classify_error`).
+fn should_retry_side_question(e: &SamplingError) -> bool {
+    e.is_overloaded() && !e.is_retry_vetoed()
+}
+
+/// Clone the base `/btw` request and stamp a fresh `req_id`, so retried
+/// attempts never collide in logs. Everything else is byte-identical.
+fn build_side_question_attempt(base: &ConversationRequest) -> ConversationRequest {
+    let mut request = base.clone();
+    request.x_grok_req_id = Some(format!("xai-btw-{}", uuid::Uuid::new_v4()));
+    request
+}
 
 impl SessionActor {
     /// Handle a /btw side question — single-turn model call using the
@@ -18,7 +49,10 @@ impl SessionActor {
     ///
     /// Generates a unique btw session ID and persists the result to
     /// `btw_history.jsonl` in the session folder.
-    pub(super) async fn handle_side_question(&self, question: &str) -> Result<String, String> {
+    pub(super) async fn handle_side_question(
+        &self,
+        question: &str,
+    ) -> Result<String, SideQuestionError> {
         let btw_session_id = format!("btw-{}", uuid::Uuid::new_v4());
         let parent_session_id = self.session_info.id.to_string();
         let asked_at = chrono::Utc::now();
@@ -26,7 +60,7 @@ impl SessionActor {
         let sampling_client = self
             .prepare_chat_completion(false)
             .await
-            .map_err(|e| format!("failed to prepare client: {e}"))?;
+            .map_err(|e| SideQuestionError::PrepareClient(e.to_string()))?;
 
         // Full conversation snapshot including system prompt, tool calls, and results.
         // Strip reasoning/thinking blocks from assistant items so we don't send
@@ -86,7 +120,7 @@ impl SessionActor {
             .map(|c| c.model)
             .unwrap_or_default();
 
-        let persist = |answer: String, success: bool, error: Option<String>| {
+        let persist = |answer: String, success: bool, error: Option<String>, attempts: u32| {
             let _ = self.notifications.persistence_tx.send(PersistenceMsg::Btw(
                 crate::session::persistence::BtwEntry {
                     btw_session_id: btw_session_id.clone(),
@@ -97,6 +131,7 @@ impl SessionActor {
                     model: model.clone(),
                     success,
                     error,
+                    attempts,
                 },
             ));
         };
@@ -105,34 +140,56 @@ impl SessionActor {
         // `thinking` config via request_defaults for thinking-enabled models,
         // Anthropic requires temperature == 1 when thinking is enabled.
         // Leaving it None lets the provider defaults apply correctly.
-        let request = ConversationRequest {
+        //
+        // Built once; each attempt clones it and stamps a fresh req_id (the
+        // per-attempt clone is the cost of the owned-request API — retries
+        // are rare, so the success path pays exactly one clone).
+        let base_request = ConversationRequest {
             items,
             tools: tool_specs,
             model: Some(model.clone()),
             temperature: None,
             x_grok_conv_id: Some(btw_session_id.clone()),
-            x_grok_req_id: Some(format!("xai-btw-{}", uuid::Uuid::new_v4())),
             x_grok_session_id: Some(parent_session_id.clone()),
             x_grok_agent_id: Some(xai_grok_telemetry::id::agent_id()),
             ..Default::default()
         };
 
-        let response = sampling_client
-            .conversation_collect(request)
-            .await
-            .map_err(|e| {
-                let msg = format!("side question model call failed: {e}");
-                persist(String::new(), false, Some(msg.clone()));
-                msg
-            })?;
-        let content = response.assistant_text();
+        // conversation_collect is one-shot (no sampler-actor retry); /btw adds
+        // its own bounded overload-only retry (policy + predicate above).
+        use backon::Retryable as _;
+        let attempts = std::cell::Cell::new(1u32);
+        let result =
+            (|| sampling_client.conversation_collect(build_side_question_attempt(&base_request)))
+                .retry(side_question_retry_policy())
+                .when(should_retry_side_question)
+                .notify(|e: &SamplingError, backoff: std::time::Duration| {
+                    attempts.set(attempts.get() + 1);
+                    tracing::warn!(
+                        backoff_ms = backoff.as_millis() as u64,
+                        error = %e,
+                        "side question overload; retrying"
+                    );
+                })
+                .await;
 
-        if content.is_empty() {
-            persist(String::new(), false, Some("No response from model".into()));
-            return Err("No response from model".to_string());
+        match result {
+            Ok(response) => {
+                let content = response.assistant_text();
+                if content.is_empty() {
+                    let err = SideQuestionError::EmptyResponse;
+                    persist(String::new(), false, Some(err.to_string()), attempts.get());
+                    return Err(err);
+                }
+                persist(content.clone(), true, None, attempts.get());
+                Ok(content)
+            }
+            Err(e) => {
+                let err = SideQuestionError::from(e);
+                persist(String::new(), false, Some(err.to_string()), attempts.get());
+                Err(err)
+            }
         }
-        persist(content.clone(), true, None);
-        Ok(content)
     }
 
     /// Generate a session recap and broadcast it via
@@ -210,11 +267,10 @@ impl SessionActor {
         };
 
         let tag = self.reminder_wrapper_tag();
-        // Strip reasoning ONLY on the Anthropic Messages backend (it rejects
-        // thinking blocks without a `thinking` config). Every other backend
-        // keeps reasoning verbatim so the prefix matches the last turn and the
-        // provider's prefix KV cache stays warm. Mirrors compaction's
-        // `summary_strips_reasoning`.
+        // Strip reasoning only on the Messages backend (it rejects thinking
+        // blocks without a `thinking` config). Other backends keep reasoning
+        // verbatim so the prefix matches the last turn and the prefix KV
+        // cache stays warm. Mirrors compaction's `summary_strips_reasoning`.
         let strip_reasoning =
             sampling_client.api_backend() == crate::sampling::ApiBackend::Messages;
 
@@ -689,3 +745,102 @@ impl SessionActor {
         suggestion
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn api(status: u16, message: &str, should_retry: Option<bool>) -> SamplingError {
+        SamplingError::Api {
+            status: reqwest::StatusCode::from_u16(status).unwrap(),
+            message: message.into(),
+            model_metadata: None,
+            retry_after_secs: None,
+            should_retry,
+        }
+    }
+
+    #[test]
+    fn side_question_retries_overload_only() {
+        // Stream overload and its proxy-wrapped 500 shape retry; so does 529.
+        assert!(should_retry_side_question(&SamplingError::StreamError {
+            error_type: "overloaded_error".into(),
+            message: "Overloaded".into(),
+        }));
+        assert!(should_retry_side_question(&api(
+            500,
+            "stream error (overloaded_error): Overloaded",
+            None
+        )));
+        assert!(should_retry_side_question(&api(529, "capacity", None)));
+
+        // Server veto (`x-should-retry: false`) wins over overload.
+        assert!(!should_retry_side_question(&api(
+            529,
+            "capacity",
+            Some(false)
+        )));
+        // Deterministic context-length failures never retry, even on 529.
+        assert!(!should_retry_side_question(&api(
+            529,
+            "invalid_request_error: prompt is too long: 300000 tokens > 200000 maximum",
+            None
+        )));
+        // Rate limit and generic 5xx are not overload — no /btw retry.
+        assert!(!should_retry_side_question(&api(429, "slow down", None)));
+        assert!(!should_retry_side_question(&api(
+            503,
+            "upstream connect timeout",
+            None
+        )));
+    }
+
+    /// The wired policy: 3 attempts total, backoff within the configured
+    /// bounds (500ms + 1s base, jitter adds up to the current delay), and a
+    /// fresh request id stamped per attempt.
+    #[tokio::test(start_paused = true)]
+    async fn side_question_retry_wiring_caps_attempts_and_bounds_backoff() {
+        use backon::Retryable as _;
+
+        let calls = std::cell::Cell::new(0u32);
+        let start = tokio::time::Instant::now();
+        let result: Result<(), SamplingError> = (|| async {
+            calls.set(calls.get() + 1);
+            Err(SamplingError::StreamError {
+                error_type: "overloaded_error".into(),
+                message: "Overloaded".into(),
+            })
+        })
+        .retry(side_question_retry_policy())
+        .when(should_retry_side_question)
+        .await;
+
+        assert!(result.is_err());
+        assert_eq!(calls.get(), 3, "1 try + 2 retries");
+        // Base delays 500ms + 1s; jitter adds (0, delay) per sleep.
+        let elapsed = start.elapsed();
+        assert!(
+            elapsed >= std::time::Duration::from_millis(1_500),
+            "elapsed {elapsed:?} below minimum backoff"
+        );
+        assert!(
+            elapsed <= std::time::Duration::from_millis(3_100),
+            "elapsed {elapsed:?} above maximum backoff"
+        );
+    }
+
+    #[test]
+    fn side_question_attempts_get_fresh_request_ids() {
+        let base = ConversationRequest {
+            x_grok_conv_id: Some("btw-test".into()),
+            ..Default::default()
+        };
+        let a = build_side_question_attempt(&base);
+        let b = build_side_question_attempt(&base);
+        let (a_id, b_id) = (a.x_grok_req_id.unwrap(), b.x_grok_req_id.unwrap());
+        assert!(a_id.starts_with("xai-btw-"));
+        assert_ne!(a_id, b_id, "each attempt must get a fresh req_id");
+        // Everything except the request id is byte-identical to the base.
+        assert_eq!(a.x_grok_conv_id, base.x_grok_conv_id);
+    }
+}

crates/codegen/xai-grok-shell/src/session/acp_session_impl/sampler_turn.rs

@@ -925,7 +925,10 @@ impl SessionActor {
                         "auth recovery: sampler 401, devbox re-mint, retrying"
                     );
                     self.prepare_sampler_for_turn().await;
-                    return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit);
+                    return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit {
+                        credential: error.credential,
+                        store: RecoveredStore::SessionToken,
+                    });
                 }
                 Err(e) => {
                     tracing::warn!(
@@ -953,7 +956,10 @@ impl SessionActor {
                     None,
                 );
                 self.prepare_sampler_for_turn().await;
-                return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit);
+                return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit {
+                    credential: error.credential,
+                    store: RecoveredStore::SessionToken,
+                });
             }
             tracing::warn!(session_id = %self.session_info.id.0, "auth recovery: sampler 401, refresh failed");
             xai_grok_telemetry::unified_log::warn(
@@ -966,7 +972,10 @@ impl SessionActor {
             && self.try_provider_401_recovery(provider).await
         {
             self.prepare_sampler_for_turn().await;
-            return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit);
+            return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit {
+                credential: error.credential,
+                store: RecoveredStore::AuthProvider,
+            });
         }
         if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
             self.signals_handle().record_idle_timeout();
@@ -1172,8 +1181,8 @@ impl SessionActor {
                     SamplerFailureRecovery::CompactAndResubmit => {
                         Ok(SamplerTurnOutcome::CompactAndResubmit)
                     }
-                    SamplerFailureRecovery::RefreshAuthAndResubmit => {
-                        Ok(SamplerTurnOutcome::RefreshAuthAndResubmit)
+                    SamplerFailureRecovery::RefreshAuthAndResubmit { credential, store } => {
+                        Ok(SamplerTurnOutcome::RefreshAuthAndResubmit { credential, store })
                     }
                 }
             }

crates/codegen/xai-grok-shell/src/session/acp_session_impl/spawn.rs

@@ -1593,6 +1593,7 @@ pub(crate) async fn spawn_session_actor(
             tool_choice: compaction_tool_choice,
             prefire: crate::session::compaction_config::PrefireState::default(),
             prefix_released: std::sync::atomic::AtomicBool::new(false),
+            cancel: Default::default(),
         },
         memory: super::memory_state::SessionMemory {
             flush_config: memory_config.as_ref().map_or_else(

crates/codegen/xai-grok-shell/src/session/acp_session_impl/tasks_cancel.rs

@@ -248,6 +248,9 @@ impl SessionActor {
         trigger: Option<String>,
     ) {
         let suppress_task_wakes = trigger.as_deref() == Some("ctrl_c");
+        // Abort in-flight `/compact` or auto-compact generation (stream select +
+        // pre-replace guard). Safe when no compact is running.
+        self.compaction.cancel.request_cancel();
         if suppress_task_wakes {
             if let Some(gate) = &self.tool_context.task_wake_suppressed {
                 gate.set(true);

crates/codegen/xai-grok-shell/src/session/acp_session_impl/turn.rs

@@ -1,6 +1,7 @@
 //! Turn-execution concern for `SessionActor` (`handle_prompt`, turn-end,
 //! sampling loop).
 use super::*;
+use crate::util::dual_clock::DualClock;
 use xai_grok_tools::implementations::grok_build::LoopFireMode;
 /// Synthetic tool the model calls to return its schema-constrained final answer
 /// on backends that can't constrain output natively (Messages API). Intercepted
@@ -1887,6 +1888,7 @@ impl SessionActor {
         json_schema: Option<serde_json::Value>,
     ) -> Result<TurnOutcome, acp::Error> {
         let conv_turn_start = std::time::Instant::now();
+        let conv_turn_clock = DualClock::now();
         self.maybe_refresh_model_metadata_on_resume().await;
         self.maybe_compact_on_model_switch().await?;
         self.chat_state_handle
@@ -2193,50 +2195,140 @@ impl SessionActor {
                     return Err(error);
                 }
                 Ok(SamplerTurnOutcome::CompactAndResubmit) => {
-                    auth_retry_schedule.reset();
+                    auth_retry_schedule.reset_on_success();
                     continue;
                 }
-                Ok(SamplerTurnOutcome::RefreshAuthAndResubmit) => {
-                    if let Some((attempt, delay)) = auth_retry_schedule.next_delay() {
-                        let delay_ms = delay.as_millis() as u64;
-                        tracing::warn!(
-                            attempt,
-                            delay_ms,
-                            "auth 401 retry: backing off before resubmit"
-                        );
-                        xai_grok_telemetry::unified_log::warn(
-                            "shell.turn.auth_retry_backoff",
+                Ok(SamplerTurnOutcome::RefreshAuthAndResubmit { credential, store }) => {
+                    if auth_retry_schedule.reset_if_incident_spans_suspend() {
+                        tracing::info!("auth 401 retry: incident spanned a suspend; budget reset");
+                        xai_grok_telemetry::unified_log::info(
+                            "shell.turn.auth_retry_reset_after_suspend",
                             Some(self.session_info.id.0.as_ref()),
-                            Some(serde_json::json!({
-                                "loop_index": loop_index,
-                                "attempt": attempt,
-                                "max_retries": AuthRetrySchedule::MAX_RETRIES,
-                                "delay_ms": delay_ms,
-                            })),
+                            Some(serde_json::json!({ "loop_index": loop_index })),
                         );
-                        self.send_xai_notification(XaiSessionUpdate::RetryState(
-                            crate::extensions::notification::RetryState::Retrying {
-                                attempt,
-                                max_retries: AuthRetrySchedule::MAX_RETRIES,
-                                reason: "Re-authenticated after 401; retrying request".to_string(),
-                            },
-                        ))
-                        .await;
-                        sleep(delay).await;
-                        continue;
                     }
-                    let msg = format!(
-                        "Auth recovery succeeded but inference request was \
-                         still rejected (401) after {} retries",
-                        AuthRetrySchedule::MAX_RETRIES
-                    );
-                    tracing::error!(msg);
-                    return Err(acp::Error::internal_error().data(
-                        crate::sampling::error::error_data_with_status(msg, Some(401)),
-                    ));
+                    match auth_retry_schedule.on_recovered_401(credential) {
+                        AuthRetryDecision::UnchargedResubmit { resubmit } => {
+                            tracing::warn!(
+                                resubmit,
+                                "auth 401 retry: no credential was sent; resubmitting uncharged"
+                            );
+                            xai_grok_telemetry::unified_log::warn(
+                                "shell.turn.auth_resubmit_uncharged",
+                                Some(self.session_info.id.0.as_ref()),
+                                Some(serde_json::json!({
+                                    "loop_index": loop_index,
+                                    "resubmit": resubmit,
+                                    "max_resubmits": AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS,
+                                })),
+                            );
+                            self.send_xai_notification(XaiSessionUpdate::RetryState(
+                                crate::extensions::notification::RetryState::Retrying {
+                                    attempt: resubmit,
+                                    max_retries: AuthRetrySchedule::MAX_UNCHARGED_RESUBMITS,
+                                    reason: "Re-authenticated after 401 (request carried no \
+                                             credential); retrying request"
+                                        .to_string(),
+                                },
+                            ))
+                            .await;
+                            pace_uncharged_resubmit(store, self.auth_manager.as_ref()).await;
+                            continue;
+                        }
+                        AuthRetryDecision::Backoff { attempt, delay } => {
+                            let delay_ms = delay.as_millis() as u64;
+                            tracing::warn!(
+                                attempt,
+                                delay_ms,
+                                "auth 401 retry: backing off before resubmit"
+                            );
+                            xai_grok_telemetry::unified_log::warn(
+                                "shell.turn.auth_retry_backoff",
+                                Some(self.session_info.id.0.as_ref()),
+                                Some(serde_json::json!({
+                                    "loop_index": loop_index,
+                                    "attempt": attempt,
+                                    "max_retries": AuthRetrySchedule::MAX_RETRIES,
+                                    "delay_ms": delay_ms,
+                                })),
+                            );
+                            self.send_xai_notification(XaiSessionUpdate::RetryState(
+                                crate::extensions::notification::RetryState::Retrying {
+                                    attempt,
+                                    max_retries: AuthRetrySchedule::MAX_RETRIES,
+                                    reason: "Re-authenticated after 401; retrying request"
+                                        .to_string(),
+                                },
+                            ))
+                            .await;
+                            sleep(delay).await;
+                            continue;
+                        }
+                        decision @ (AuthRetryDecision::Exhausted
+                        | AuthRetryDecision::RunawayGuard { .. }) => {
+                            let (awake, wall, suspended) = conv_turn_clock.elapsed_split();
+                            let duration_note = if suspended >= std::time::Duration::from_secs(1) {
+                                format!(
+                                    " Turn ran {} wall-clock, {} of it suspended.",
+                                    human_duration(wall),
+                                    human_duration(suspended)
+                                )
+                            } else {
+                                format!(" Turn ran {} wall-clock.", human_duration(wall))
+                            };
+                            let (rejections, authenticated) = auth_retry_schedule.incident_counts();
+                            let uncharged = auth_retry_schedule.uncharged_rejections();
+                            let msg = match decision {
+                                AuthRetryDecision::RunawayGuard { rejections } => {
+                                    format!(
+                                        "Auth recovery kept succeeding but {rejections} requests \
+                                     were rejected (401) before a credential could be sent, \
+                                     with no successful response in between; stopping as a \
+                                     runaway guard.{duration_note}"
+                                    )
+                                }
+                                _ if authenticated == rejections => {
+                                    format!(
+                                        "Auth recovery succeeded but {rejections} authenticated \
+                                     inference requests were still rejected (401); giving up \
+                                     after {} retries.{duration_note}",
+                                        AuthRetrySchedule::MAX_RETRIES
+                                    )
+                                }
+                                _ => {
+                                    format!(
+                                        "Auth retry budget exhausted after {rejections} \
+                                     post-recovery 401s ({authenticated} provably carried a \
+                                     credential).{duration_note}"
+                                    )
+                                }
+                            };
+                            tracing::error!(msg);
+                            xai_grok_telemetry::unified_log::error(
+                                "shell.turn.auth_retry_exhausted",
+                                Some(self.session_info.id.0.as_ref()),
+                                Some(serde_json::json!({
+                                    "loop_index": loop_index,
+                                    "decision": match decision {
+                                        AuthRetryDecision::RunawayGuard { .. } => "runaway_guard",
+                                        _ => "exhausted",
+                                    },
+                                    "rejections": rejections,
+                                    "authenticated": authenticated,
+                                    "uncharged": uncharged,
+                                    "wall_secs": wall.as_secs(),
+                                    "awake_secs": awake.as_secs(),
+                                    "suspended_secs": suspended.as_secs(),
+                                })),
+                            );
+                            return Err(acp::Error::internal_error().data(
+                                crate::sampling::error::error_data_with_status(msg, Some(401)),
+                            ));
+                        }
+                    }
                 }
             };
-            auth_retry_schedule.reset();
+            auth_retry_schedule.reset_on_success();
             let model_elapsed_ms = model_timer.elapsed().as_millis() as u64;
             let usage = response.usage.as_ref();
             let prompt_tokens = usage.map(|u| u.prompt_tokens);
@@ -2747,87 +2839,6 @@ mod identical_tool_call_run_tests {
         assert!(!run.take_nudge());
     }
 }
-/// Backoff schedule for resubmits after a *successful* 401 auth recovery
-/// (fresh token minted, request to be re-sent).
-///
-/// Two hard-won invariants, both regressions from the silent-hang incident
-/// where a turn froze 16m40s and then 11.6 days (user-cancelled at 27min):
-///
-/// - **Delays must be 1s/2s/4s.** `tokio_retry::ExponentialBackoff::
-///   from_millis(base)` raises `base` to the attempt number, so the base must
-///   stay small: `from_millis(1000)` yields 1000ⁿ ms = 1s → 16m40s → 11.57
-///   days. `from_millis(2).factor(500)` yields 2ⁿ × 500ms = 1s, 2s, 4s.
-/// - **The schedule is per-incident, not per-turn.** A long turn can span
-///   several hourly gateway token rotations; each rotation is an independent
-///   401→refresh→retry event. Without `reset()` after a successful response,
-///   the third rotation of one turn would land on the last (largest) delay
-///   and the fourth would fail the turn outright.
-struct AuthRetrySchedule {
-    delays: std::iter::Take<ExponentialBackoff>,
-    attempt: u32,
-}
-impl AuthRetrySchedule {
-    /// Consecutive post-recovery 401s tolerated before the turn fails.
-    const MAX_RETRIES: u32 = 3;
-    fn new() -> Self {
-        Self {
-            delays: ExponentialBackoff::from_millis(2)
-                .factor(500)
-                .max_delay(std::time::Duration::from_secs(10))
-                .take(Self::MAX_RETRIES as usize),
-            attempt: 0,
-        }
-    }
-    /// Next `(attempt_number, delay)` (1-indexed), or `None` once exhausted.
-    fn next_delay(&mut self) -> Option<(u32, std::time::Duration)> {
-        let delay = self.delays.next()?;
-        self.attempt += 1;
-        Some((self.attempt, delay))
-    }
-    /// A successful model response closes the incident: restart the schedule
-    /// so the next token rotation starts back at the shortest delay.
-    fn reset(&mut self) {
-        *self = Self::new();
-    }
-}
-#[cfg(test)]
-mod auth_retry_schedule_tests {
-    use super::AuthRetrySchedule;
-    use std::time::Duration;
-    /// Pins the exact schedule. Guards against the `from_millis(1000)`
-    /// footgun (baseⁿ semantics): that spelling produced sleeps of 1s,
-    /// 16m40s, and 11.57 days, observed in the field as a silent
-    /// ~27-minute hang in `waiting_model` that the user had to cancel.
-    #[test]
-    fn schedule_is_one_two_four_seconds_then_exhausted() {
-        let mut schedule = AuthRetrySchedule::new();
-        let steps: Vec<_> = std::iter::from_fn(|| schedule.next_delay()).collect();
-        assert_eq!(
-            steps,
-            vec![
-                (1, Duration::from_secs(1)),
-                (2, Duration::from_secs(2)),
-                (3, Duration::from_secs(4)),
-            ],
-        );
-        assert_eq!(
-            schedule.next_delay(),
-            None,
-            "must exhaust after MAX_RETRIES"
-        );
-    }
-    /// Each successful response must restart the schedule: hourly token
-    /// rotations within one long turn are independent incidents, so they
-    /// must not escalate toward exhaustion (turn failure).
-    #[test]
-    fn reset_restarts_delays_and_attempt_numbering() {
-        let mut schedule = AuthRetrySchedule::new();
-        schedule.next_delay();
-        schedule.next_delay();
-        schedule.reset();
-        assert_eq!(schedule.next_delay(), Some((1, Duration::from_secs(1))));
-    }
-}
 #[cfg(test)]
 mod user_echo_broadcast_tests {
     use super::{UserEchoMode, user_echo_mode};

crates/codegen/xai-grok-shell/src/session/acp_session_impl/types.rs

@@ -12,6 +12,20 @@ pub(crate) enum McpReminderMode {
     Full,
 }
 
+/// Which credential store a successful 401 recovery minted into. An
+/// uncharged resubmit can only usefully wait on the store that recovered:
+/// waiting on the session token for a provider-key 401 blocks 15s for a
+/// refresh that is irrelevant to the rejected credential.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) enum RecoveredStore {
+    /// `AuthManager` session token (devbox re-mint, OIDC refresh) —
+    /// `wait_for_token_refresh` is meaningful.
+    SessionToken,
+    /// Auth-provider key minted into chat-state credentials — nothing to
+    /// wait on in the `AuthManager`; floor-pace instead.
+    AuthProvider,
+}
+
 /// Recovery decision returned by
 /// `SessionActor::handle_sampling_failure` for the sampler-based
 /// turn loop.
@@ -19,10 +33,15 @@ pub(crate) enum SamplerFailureRecovery {
     /// Compaction ran. The turn loop should rebuild the request from
     /// the compacted conversation and resubmit.
     CompactAndResubmit,
-    /// Auth 401 recovery succeeded (devbox re-mint, OIDC refresh, or auth
-    /// provider re-mint). The turn loop should resubmit once with the
-    /// fresh token.
-    RefreshAuthAndResubmit,
+    /// Auth 401 recovery succeeded; the turn loop should resubmit with the
+    /// fresh token. `credential` is the wire provenance of the rejected
+    /// request: a 401 for a request that carried no credential at all (a
+    /// fail-closed send) must not be charged against the per-incident
+    /// auth-retry budget.
+    RefreshAuthAndResubmit {
+        credential: xai_grok_sampling_types::SentCredential,
+        store: RecoveredStore,
+    },
 }
 
 /// Outcome of a single turn attempt via the sampler-based path.
@@ -36,8 +55,12 @@ pub(crate) enum SamplerTurnOutcome {
         Box<xai_grok_sampler::InferenceLatencyStats>,
     ),
     CompactAndResubmit,
-    /// Auth recovery succeeded; the outer loop should retry once.
-    RefreshAuthAndResubmit,
+    /// Auth recovery succeeded; the outer loop should retry. Mirrors
+    /// [`SamplerFailureRecovery::RefreshAuthAndResubmit`].
+    RefreshAuthAndResubmit {
+        credential: xai_grok_sampling_types::SentCredential,
+        store: RecoveredStore,
+    },
 }
 
 /// Outcome of `process_conversation_turn`, distinguishing normal completion from cancellation.

crates/codegen/xai-grok-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs

@@ -60,6 +60,7 @@ fn auth_error() -> xai_grok_sampler::SamplingErrorInfo {
         empty_response_context: None,
         doom_loop_triggers: None,
         doom_loop_aborted_at_chunk: None,
+        credential: xai_grok_sampling_types::SentCredential::Unknown,
     }
 }
 
@@ -196,7 +197,13 @@ async fn sampler_401_recovery_returns_refresh_and_retry() {
             let (actor, _rx) = make_actor_with_auth_manager(Some(am)).await;
             let result = actor.handle_sampling_failure(auth_error()).await;
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit {
+                        store: RecoveredStore::SessionToken,
+                        ..
+                    })
+                ),
                 "session-based auth with a working refresher must return RefreshAuthAndResubmit"
             );
             assert!(called.load(Ordering::SeqCst), "refresher must be invoked");
@@ -529,6 +536,7 @@ fn model_not_found_error() -> xai_grok_sampler::SamplingErrorInfo {
             empty_response_context: None,
             doom_loop_triggers: None,
             doom_loop_aborted_at_chunk: None,
+            credential: xai_grok_sampling_types::SentCredential::Unknown,
         }
 }
 
@@ -596,6 +604,7 @@ fn unauthorized_401_error() -> xai_grok_sampler::SamplingErrorInfo {
             empty_response_context: None,
             doom_loop_triggers: None,
             doom_loop_aborted_at_chunk: None,
+            credential: xai_grok_sampling_types::SentCredential::Unknown,
         }
 }
 
@@ -774,7 +783,10 @@ async fn sampler_401_session_method_with_stale_api_key_auth_type_still_recovers(
             let result = actor.handle_sampling_failure(auth_error()).await;
 
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. })
+                ),
                 "session-based method must recover even when auth_type transiently reads ApiKey"
             );
             assert!(
@@ -808,7 +820,10 @@ async fn sampler_401_oidc_method_with_stale_api_key_auth_type_still_recovers() {
             let result = actor.handle_sampling_failure(auth_error()).await;
 
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. })
+                ),
                 "oidc method must recover even when auth_type transiently reads ApiKey"
             );
             assert!(
@@ -1253,8 +1268,14 @@ async fn sampler_401_on_provider_model_remints_and_resubmits() {
 
             let result = actor.handle_sampling_failure(auth_error()).await;
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
-                "provider 401 must re-mint and resubmit"
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit {
+                        store: RecoveredStore::AuthProvider,
+                        ..
+                    })
+                ),
+                "provider 401 must re-mint and resubmit via the provider store"
             );
             let creds = actor.chat_state_handle.get_credentials().await;
             assert_eq!(
@@ -1289,7 +1310,10 @@ async fn sampler_non_auth_kind_401_on_provider_model_still_recovers() {
             error.kind = xai_grok_sampler::SamplingErrorKind::Api;
             let result = actor.handle_sampling_failure(error).await;
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. })
+                ),
                 "a non-Auth-kind 401 on a provider model must still recover via 4c"
             );
             let creds = actor.chat_state_handle.get_credentials().await;
@@ -1321,7 +1345,10 @@ async fn sampler_401_with_no_key_on_provider_model_mints_and_resubmits() {
 
             let result = actor.handle_sampling_failure(auth_error()).await;
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. })
+                ),
                 "an unauthenticated 401 on a provider model must mint and resubmit"
             );
             let creds = actor.chat_state_handle.get_credentials().await;
@@ -1364,7 +1391,10 @@ async fn sampler_401_on_provider_model_never_refreshes_session() {
 
             let result = actor.handle_sampling_failure(auth_error()).await;
             assert!(
-                matches!(result, Ok(SamplerFailureRecovery::RefreshAuthAndResubmit)),
+                matches!(
+                    result,
+                    Ok(SamplerFailureRecovery::RefreshAuthAndResubmit { .. })
+                ),
                 "the provider arm must recover"
             );
             assert!(

crates/codegen/xai-grok-shell/src/session/acp_session_tests/cancel_running_task_tests.rs

@@ -169,6 +169,7 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
                     tool_choice: crate::util::config::CompactionToolChoice::Auto,
                     prefire: crate::session::compaction_config::PrefireState::default(),
                     prefix_released: std::sync::atomic::AtomicBool::new(false),
+                    cancel: Default::default(),
                 },
                 memory: crate::session::memory_state::SessionMemory {
                     flush_config: crate::config::MemoryFlushConfig::default(),
@@ -635,6 +636,7 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
                     tool_choice: crate::util::config::CompactionToolChoice::Auto,
                     prefire: crate::session::compaction_config::PrefireState::default(),
                     prefix_released: std::sync::atomic::AtomicBool::new(false),
+                    cancel: Default::default(),
                 },
                 memory: crate::session::memory_state::SessionMemory {
                     flush_config: crate::config::MemoryFlushConfig::default(),
@@ -920,6 +922,7 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
                     tool_choice: crate::util::config::CompactionToolChoice::Auto,
                     prefire: crate::session::compaction_config::PrefireState::default(),
                     prefix_released: std::sync::atomic::AtomicBool::new(false),
+                    cancel: Default::default(),
                 },
                 memory: crate::session::memory_state::SessionMemory {
                     flush_config: crate::config::MemoryFlushConfig::default(),
@@ -2176,6 +2179,7 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
                     tool_choice: crate::util::config::CompactionToolChoice::Auto,
                     prefire: crate::session::compaction_config::PrefireState::default(),
                     prefix_released: std::sync::atomic::AtomicBool::new(false),
+                    cancel: Default::default(),
                 },
                 memory: crate::session::memory_state::SessionMemory {
                     flush_config: crate::config::MemoryFlushConfig::default(),

crates/codegen/xai-grok-shell/src/session/acp_session_tests/idle_resume_tests.rs

@@ -195,6 +195,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
                     tool_choice: crate::util::config::CompactionToolChoice::Auto,
                     prefire: crate::session::compaction_config::PrefireState::default(),
                     prefix_released: std::sync::atomic::AtomicBool::new(false),
+                    cancel: Default::default(),
                 },
                 memory: crate::session::memory_state::SessionMemory {
                     flush_config: crate::config::MemoryFlushConfig::default(),

crates/codegen/xai-grok-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs

@@ -121,6 +121,7 @@ async fn create_test_actor(
             tool_choice: crate::util::config::CompactionToolChoice::Auto,
             prefire: crate::session::compaction_config::PrefireState::default(),
             prefix_released: std::sync::atomic::AtomicBool::new(false),
+            cancel: Default::default(),
         },
         memory: crate::session::memory_state::SessionMemory {
             flush_config: crate::config::MemoryFlushConfig::default(),
@@ -562,6 +563,7 @@ async fn create_test_actor_with_memory(
             tool_choice: crate::util::config::CompactionToolChoice::Auto,
             prefire: crate::session::compaction_config::PrefireState::default(),
             prefix_released: std::sync::atomic::AtomicBool::new(false),
+            cancel: Default::default(),
         },
         memory: crate::session::memory_state::SessionMemory {
             flush_config: memory_config
@@ -1158,6 +1160,7 @@ fn api_error_with_context_window(context_window: u64) -> xai_grok_sampler::Sampl
         empty_response_context: None,
         doom_loop_triggers: None,
         doom_loop_aborted_at_chunk: None,
+        credential: xai_grok_sampling_types::SentCredential::Unknown,
     }
 }
 /// Primary scenario: remote settings shrinks the context window mid-session.
@@ -1345,6 +1348,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
                     tool_choice: crate::util::config::CompactionToolChoice::Auto,
                     prefire: crate::session::compaction_config::PrefireState::default(),
                     prefix_released: std::sync::atomic::AtomicBool::new(false),
+                    cancel: Default::default(),
                 },
                 memory: crate::session::memory_state::SessionMemory {
                     flush_config: crate::config::MemoryFlushConfig::default(),
@@ -1546,6 +1550,7 @@ async fn test_compact_on_error_noop_without_model_metadata() {
                 empty_response_context: None,
                 doom_loop_triggers: None,
                 doom_loop_aborted_at_chunk: None,
+                credential: xai_grok_sampling_types::SentCredential::Unknown,
             };
             assert!(!actor.should_compact_on_error(&err).await);
         })

crates/codegen/xai-grok-shell/src/session/acp_session_tests/memory_config_tests.rs

@@ -171,6 +171,7 @@ async fn create_test_actor_with_memory(
             tool_choice: crate::util::config::CompactionToolChoice::Auto,
             prefire: crate::session::compaction_config::PrefireState::default(),
             prefix_released: std::sync::atomic::AtomicBool::new(false),
+            cancel: Default::default(),
         },
         memory: crate::session::memory_state::SessionMemory {
             flush_config: memory_config

crates/codegen/xai-grok-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs

@@ -123,6 +123,7 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
             tool_choice: crate::util::config::CompactionToolChoice::Auto,
             prefire: crate::session::compaction_config::PrefireState::default(),
             prefix_released: std::sync::atomic::AtomicBool::new(false),
+            cancel: Default::default(),
         },
         memory: crate::session::memory_state::SessionMemory {
             flush_config: crate::config::MemoryFlushConfig::default(),
@@ -808,6 +809,7 @@ async fn failed_event_preserves_streaming_capture_for_takeout() {
                         empty_response_context: None,
                         doom_loop_triggers: None,
                         doom_loop_aborted_at_chunk: None,
+                        credential: xai_grok_sampling_types::SentCredential::Unknown,
                     },
                 })
                 .await;
@@ -1229,6 +1231,7 @@ async fn reasoning_only_doomloop_turn_captures_every_generation_as_segments() {
                 }),
                 doom_loop_triggers: None,
                 doom_loop_aborted_at_chunk: None,
+                credential: xai_grok_sampling_types::SentCredential::Unknown,
             };
             actor
                 .handle_sampling_event(SamplingEvent::Failed {

crates/codegen/xai-grok-shell/src/session/acp_session_tests/support.rs

@@ -283,6 +283,7 @@ pub(crate) async fn create_test_actor_ex(
             tool_choice: crate::util::config::CompactionToolChoice::Auto,
             prefire: crate::session::compaction_config::PrefireState::default(),
             prefix_released: std::sync::atomic::AtomicBool::new(false),
+            cancel: Default::default(),
         },
         memory: crate::session::memory_state::SessionMemory {
             flush_config: crate::config::MemoryFlushConfig::default(),

crates/codegen/xai-grok-shell/src/session/acp_session_tests/turn/auth_retry_budget_tests.rs

@@ -0,0 +1,308 @@
+//! Real-turn-loop tests against a mock server that 401s unauthenticated
+//! requests and 200s a fresh bearer: a fail-closed (credential-less) 401
+//! must not consume `AuthRetrySchedule` budget — the field failure mode
+//! where each sleep cycle burned one slot — while credentialed 401s must
+//! still exhaust after `MAX_RETRIES`.
+
+use super::support::*;
+use super::*;
+use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU32, Ordering};
+use std::time::Duration;
+use xai_grok_test_support::{MockInferenceServer, MockModelEntry};
+
+/// The token the mock server accepts and the refresher mints on success.
+const FRESH_TOKEN: &str = "refreshed-test-token";
+
+/// With `fail_pre_request`, mimics the post-wake sequence: pre-send
+/// (`PreRequest`) refreshes fail transiently so the send goes out
+/// fail-closed, while the 401-triggered recovery (`ServerRejected`)
+/// succeeds and mints [`FRESH_TOKEN`]. Otherwise always succeeds.
+struct WakeGapRefresher {
+    calls: Arc<AtomicU32>,
+    fail_pre_request: bool,
+}
+
+#[async_trait::async_trait]
+impl crate::auth::refresh::TokenRefresher for WakeGapRefresher {
+    async fn refresh(
+        &self,
+        reason: crate::auth::refresh::RefreshReason,
+    ) -> crate::auth::refresh::RefreshOutcome {
+        self.calls.fetch_add(1, Ordering::SeqCst);
+        if self.fail_pre_request && reason == crate::auth::refresh::RefreshReason::PreRequest {
+            return crate::auth::refresh::RefreshOutcome::TransientFailure {
+                message: "simulated post-wake network gap".to_string(),
+            };
+        }
+        crate::auth::refresh::RefreshOutcome::success(GrokAuth {
+            key: FRESH_TOKEN.to_string(),
+            auth_mode: AuthMode::Oidc,
+            refresh_token: Some("rt-new".into()),
+            expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
+            ..GrokAuth::test_default()
+        })
+    }
+}
+
+/// `(tempdir, manager)` with a hard-expired OIDC token, so the wire-valid
+/// resolver has nothing to stamp until the refresher succeeds. The tempdir
+/// must outlive the manager (auth.json path).
+fn expired_auth_manager(
+    refresher: Arc<dyn crate::auth::refresh::TokenRefresher>,
+) -> (tempfile::TempDir, Arc<AuthManager>) {
+    let dir = tempfile::tempdir().expect("tempdir");
+    let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
+    am.hot_swap(GrokAuth {
+        key: "initial-test-key".into(),
+        auth_mode: AuthMode::Oidc,
+        refresh_token: Some("rt".into()),
+        expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
+        ..GrokAuth::test_default()
+    });
+    am.set_refresher(refresher);
+    (dir, am)
+}
+
+fn drain_gateway(mut rx: tokio::sync::mpsc::UnboundedReceiver<xai_acp_lib::AcpClientMessage>) {
+    tokio::task::spawn_local(async move {
+        while let Some(msg) = rx.recv().await {
+            if let xai_acp_lib::AcpClientMessage::SessionNotification(args) = msg {
+                let _ = args.response_tx.send(Ok(()));
+            }
+        }
+    });
+}
+
+fn drain_persistence(mut rx: tokio::sync::mpsc::UnboundedReceiver<PersistenceMsg>) {
+    tokio::task::spawn_local(async move {
+        while let Some(msg) = rx.recv().await {
+            if let PersistenceMsg::FlushAndAck { respond_to } = msg {
+                let _ = respond_to.send(());
+            }
+        }
+    });
+}
+
+/// Actor wired for session-token auth against the mock server: real sampler,
+/// `cached_token` method, `NotByok` model facts (so the session-token gate is
+/// active against the loopback URL), and the supplied auth manager.
+async fn session_token_actor(
+    server: &MockInferenceServer,
+    auth_manager: Arc<AuthManager>,
+) -> Arc<SessionActor> {
+    let sampling_cfg = xai_grok_sampler::SamplerConfig {
+        base_url: server.url(),
+        model: "test".to_string(),
+        api_backend: xai_grok_sampler::ApiBackend::Responses,
+        context_window: 256_000,
+        max_retries: Some(0),
+        idle_timeout_secs: Some(30),
+        ..Default::default()
+    };
+    let (sampler_event_tx, sampler_event_rx) =
+        tokio::sync::mpsc::unbounded_channel::<xai_grok_sampler::SamplingEvent>();
+    let sampler_handle = xai_grok_sampler::SamplerActor::spawn(
+        sampling_cfg,
+        xai_grok_sampler::RetryPolicy {
+            max_retries: 0,
+            rate_limit_retry_threshold: 0,
+            ..Default::default()
+        },
+        sampler_event_tx,
+    );
+
+    let (gateway_tx, gateway_rx) = tokio::sync::mpsc::unbounded_channel();
+    drain_gateway(gateway_rx);
+    let (persistence_tx, persistence_rx) = tokio::sync::mpsc::unbounded_channel();
+    drain_persistence(persistence_rx);
+
+    let mut actor = create_test_actor(0, 256_000, 85, gateway_tx, persistence_tx).await;
+    actor.sampler_handle = sampler_handle;
+    actor.auth_manager = Some(auth_manager);
+    actor.auth_method_id = test_auth_method_id("cached_token");
+
+    let mut cfg = actor
+        .chat_state_handle
+        .get_sampling_config()
+        .await
+        .expect("test actor has sampling config");
+    cfg.base_url = server.url();
+    cfg.api_backend = xai_grok_sampling_types::ApiBackend::Responses;
+    cfg.model = "test".to_string();
+    actor.chat_state_handle.update_sampling_config(cfg);
+    let mut creds = actor.chat_state_handle.get_credentials().await;
+    creds.api_key = None;
+    creds.auth_type = xai_chat_state::AuthType::SessionToken;
+    actor.chat_state_handle.update_credentials(creds);
+
+    // Definite NotByok: the session-token gate must stay active against the
+    // loopback mock URL (an `Unknown` would demand a first-party host).
+    actor
+        .model_auth_memo
+        .replace(Some(crate::session::acp_session::ModelAuthMemo {
+            model_id: "test".to_string(),
+            facts: crate::agent::config::ModelAuthFacts {
+                byok: crate::agent::auth_method::ModelByok::NotByok,
+                auth_scheme: Default::default(),
+            },
+            provider: None,
+        }));
+
+    actor
+        .workspace_ops
+        .bind_local_session(
+            &actor.session_id_string(),
+            actor.tool_context.cwd.as_path().to_path_buf(),
+            actor.tool_context.hunk_tracker_handle.clone(),
+            actor.agent.borrow().tool_bridge().toolset(),
+            None,
+        )
+        .expect("bind_local_session");
+
+    let actor = Arc::new(actor);
+    {
+        let drainer = actor.clone();
+        let mut sampler_event_rx = sampler_event_rx;
+        tokio::task::spawn_local(async move {
+            while let Some(event) = sampler_event_rx.recv().await {
+                drainer.handle_sampling_event(event).await;
+            }
+        });
+    }
+    actor
+}
+
+async fn run_prompt(
+    actor: &Arc<SessionActor>,
+    prompt_id: &str,
+) -> Result<crate::session::commands::PromptTurnOk, acp::Error> {
+    let prompt_blocks = vec![acp::ContentBlock::Text(acp::TextContent::new(
+        "hello".to_string(),
+    ))];
+    tokio::time::timeout(
+        Duration::from_secs(60),
+        actor.handle_prompt(
+            prompt_id,
+            prompt_blocks,
+            PromptMode::Agent,
+            None,
+            None,
+            None,
+            None,
+            true,
+            None,
+            None,
+            None,
+        ),
+    )
+    .await
+    .expect("turn must finish within timeout")
+}
+
+/// The wake sequence: the resolver has nothing wire-valid, the send goes
+/// out with no `Authorization` header, the server 401s it, recovery lands a
+/// fresh token. The turn must survive and resubmit with the fresh bearer.
+#[tokio::test(flavor = "current_thread")]
+async fn fail_closed_401_is_uncharged_and_turn_survives() {
+    let local = tokio::task::LocalSet::new();
+    local
+        .run_until(async {
+            let server = MockInferenceServer::start_with_required_auth(
+                vec![MockModelEntry::new("test")],
+                FRESH_TOKEN,
+            )
+            .await
+            .expect("mock inference server");
+
+            let calls = Arc::new(AtomicU32::new(0));
+            // Pre-send refreshes fail like a post-wake network gap, so the
+            // first send goes out fail-closed; the 401-recovery refresh
+            // succeeds.
+            let refresher = Arc::new(WakeGapRefresher {
+                calls: calls.clone(),
+                fail_pre_request: true,
+            });
+            let (_dir, am) = expired_auth_manager(refresher);
+            let actor = session_token_actor(&server, am).await;
+
+            let outcome = run_prompt(&actor, "auth-retry-budget-fail-closed").await;
+            assert!(
+                outcome.is_ok(),
+                "fail-closed 401 must not fail the turn: {outcome:?}"
+            );
+
+            let inference: Vec<_> = server
+                .requests()
+                .into_iter()
+                .filter(|r| r.path.contains("/responses"))
+                .collect();
+            assert!(
+                inference.len() >= 2,
+                "expected the fail-closed send plus the resubmit; got {}",
+                inference.len()
+            );
+            assert_eq!(
+                inference[0].authorization, None,
+                "first send must carry no Authorization header"
+            );
+            assert_eq!(
+                inference.last().unwrap().authorization.as_deref(),
+                Some(&format!("Bearer {FRESH_TOKEN}") as &str),
+                "resubmit must carry the freshly refreshed bearer"
+            );
+            assert!(
+                calls.load(Ordering::SeqCst) >= 2,
+                "both the failing pre-flight and the recovery refresh must run"
+            );
+        })
+        .await;
+}
+
+/// Real credential rejections must still terminate: when every request
+/// carries a bearer the server rejects, the escalating budget exhausts after
+/// `MAX_RETRIES` and the failure names authenticated rejections — not a
+/// generic budget message. `start_paused` auto-advances the backoff ladder.
+#[tokio::test(flavor = "current_thread", start_paused = true)]
+async fn authenticated_401s_still_exhaust_after_three_retries() {
+    let local = tokio::task::LocalSet::new();
+    local
+        .run_until(async {
+            // The server only accepts a token the refresher never mints, so
+            // every authenticated send is rejected.
+            let server = MockInferenceServer::start_with_required_auth(
+                vec![MockModelEntry::new("test")],
+                "never-issued-token",
+            )
+            .await
+            .expect("mock inference server");
+
+            let refresher = Arc::new(WakeGapRefresher {
+                calls: Arc::new(AtomicU32::new(0)),
+                fail_pre_request: false,
+            });
+            let (_dir, am) = expired_auth_manager(refresher);
+            let actor = session_token_actor(&server, am).await;
+
+            let outcome = run_prompt(&actor, "auth-retry-budget-exhaust").await;
+            let err = outcome.expect_err("authenticated 401s must exhaust and fail the turn");
+            let rendered = serde_json::to_string(&err.data).unwrap_or_default();
+            assert!(
+                rendered.contains("authenticated inference requests were still rejected"),
+                "exhaustion must name authenticated rejections, got: {rendered}"
+            );
+
+            let authenticated = server
+                .requests()
+                .into_iter()
+                .filter(|r| r.path.contains("/responses"))
+                .filter(|r| r.authorization.as_deref() == Some(&format!("Bearer {FRESH_TOKEN}")))
+                .count();
+            assert_eq!(
+                authenticated, 4,
+                "initial send plus MAX_RETRIES resubmits, all authenticated"
+            );
+        })
+        .await;
+}

crates/codegen/xai-grok-shell/src/session/commands.rs

@@ -20,6 +20,20 @@ pub struct CancellationContext {
     /// `None` for graceful in-turn cancels and older clients.
     pub trigger: Option<String>,
 }
+/// Failure surface of a `/btw` side question. Kept typed until the ACP
+/// boundary so `handle_btw` can route model errors through the canonical
+/// [`map_sampling_err_to_acp`](crate::sampling::error::map_sampling_err_to_acp)
+/// (typed rate-limit / auth codes) instead of a flattened string.
+#[derive(Debug, thiserror::Error)]
+#[non_exhaustive]
+pub enum SideQuestionError {
+    #[error("side question model call failed: {0}")]
+    Sampling(#[from] xai_grok_sampling_types::SamplingError),
+    #[error("failed to prepare client: {0}")]
+    PrepareClient(String),
+    #[error("No response from model")]
+    EmptyResponse,
+}
 /// Prompt completion kind returned to the ACP layer.
 #[derive(Debug, Clone)]
 pub enum PromptCompletionKind {
@@ -668,7 +682,7 @@ pub enum SessionCommand {
     /// tool-free model call, and returns the response text.
     SideQuestion {
         question: String,
-        respond_to: oneshot::Sender<Result<String, String>>,
+        respond_to: oneshot::Sender<Result<String, SideQuestionError>>,
     },
     /// Generate a session recap (a short "where was I" summary) and broadcast
     /// it to clients via `SessionUpdate::SessionRecap`.

crates/codegen/xai-grok-shell/src/session/compaction.rs

@@ -189,6 +189,7 @@ impl SessionActor {
             .compaction_policy()
             .wall_clock_budget_secs;
         let hosted_tools = self.hosted_tools_for_turn();
+        let (cancel, _cancel_scope) = self.compaction.cancel.enter();
         match generate_session_compact(
             history,
             tools,
@@ -199,6 +200,7 @@ impl SessionActor {
             self.inference_idle_timeout,
             wall_clock_budget_secs,
             self.compaction.tool_choice,
+            &cancel,
         )
         .await
         {
@@ -601,6 +603,7 @@ impl SessionActor {
         self: &Arc<Self>,
         user_context: Option<String>,
     ) -> Result<(), acp::Error> {
+        let (_cancel, _cancel_scope) = self.compaction.cancel.enter();
         self.record_compaction_variant();
         let total_tokens = self.chat_state_handle.get_total_tokens().await;
         tracing::Span::current().record("pre_tokens", total_tokens as i64);
@@ -638,6 +641,16 @@ impl SessionActor {
         .await;
         Ok(())
     }
+    async fn emit_compact_cancelled(&self, auto_trigger: bool) -> Result<(), acp::Error> {
+        if auto_trigger {
+            use crate::extensions::notification::SessionUpdate as XaiSessionUpdate;
+            self.send_xai_notification(XaiSessionUpdate::AutoCompactCancelled {
+                reason: crate::extensions::notification::AutoCompactCancelReason::UserCancelled,
+            })
+            .await;
+        }
+        Err(crate::session::helpers::session_compact::CompactFailure::cancelled_error())
+    }
     /// Suppress AUTO compaction after a deterministic failure. Scope depends on
     /// the reason (see [`SuppressReason::suppress_state`]): size/schema sticky,
     /// credit until 200, auth until credentials recover, other clears next turn.
@@ -891,6 +904,7 @@ impl SessionActor {
         auto_continue: Option<crate::extensions::notification::AutoContinueInfo>,
         trigger: xai_grok_telemetry::events::CompactionTrigger,
     ) -> Result<(), acp::Error> {
+        let (cancel, _cancel_scope) = self.compaction.cancel.enter();
         let tokens_before = self.chat_state_handle.get_total_tokens().await;
         tracing::Span::current().record("compaction_tokens_before", tokens_before as i64);
         self.signals_handle().record_compaction(tokens_before);
@@ -1075,6 +1089,7 @@ impl SessionActor {
             self.inference_idle_timeout,
             wall_clock_budget_secs,
             self.compaction.tool_choice,
+            cancel.clone(),
         );
         let observer =
             crate::session::helpers::full_replace_compaction::ShellFullReplaceObserver::new(
@@ -1135,6 +1150,13 @@ impl SessionActor {
                     deterministic,
                     context_overflow,
                 }) => {
+                    if cancel.is_cancelled()
+                        || message.contains(
+                            crate::session::helpers::session_compact::COMPACT_CANCELLED_MSG,
+                        )
+                    {
+                        return self.emit_compact_cancelled(auto_trigger).await;
+                    }
                     if context_overflow {
                         let next_stage = match input_stage {
                             InputStage::Verbatim => Some(InputStage::VerbatimFitted),
@@ -1575,7 +1597,6 @@ impl SessionActor {
         let agents_md_reminder = self.agent.borrow().agents_md_user_reminder();
         let compaction_context = state_context.for_compaction();
         let compaction_state_context: &CompactionStateContext = &compaction_context;
-        self.persist_compaction_segment(&segment_messages, &generate_session_compact);
         let transcript_hint = self.transcript_hint();
         let summary_count = self
             .compaction
@@ -1620,7 +1641,7 @@ impl SessionActor {
                 user_message_prefix,
                 agents_md_reminder,
                 state_context: &state_context.for_compaction(),
-                compaction_summary: generate_session_compact,
+                compaction_summary: generate_session_compact.clone(),
                 system_reminder,
                 summary_before_recent: use_short_prompt,
                 transcript_hint,
@@ -1628,8 +1649,6 @@ impl SessionActor {
             })
         };
         let prompt_index_at_compaction = self.chat_state_handle.get_prompt_index().await;
-        self.chat_state_handle
-            .record_compaction_at(prompt_index_at_compaction);
         let original_user_info = self
             .chat_state_handle
             .get_conversation_item_at(1)
@@ -1645,6 +1664,12 @@ impl SessionActor {
                 }
                 _ => None,
             });
+        if cancel.is_cancelled() {
+            return self.emit_compact_cancelled(auto_trigger).await;
+        }
+        self.persist_compaction_segment(&segment_messages, &generate_session_compact);
+        self.chat_state_handle
+            .record_compaction_at(prompt_index_at_compaction);
         self.persist_compaction_checkpoint(
             &compacted_history,
             prompt_index_at_compaction,
@@ -2008,6 +2033,7 @@ impl SessionActor {
         trigger_info: AutoCompactTriggerInfo,
     ) -> Result<(), acp::Error> {
         use crate::extensions::notification::SessionUpdate as XaiSessionUpdate;
+        let (_cancel, _cancel_scope) = self.compaction.cancel.enter();
         self.record_compaction_variant();
         let tokens_before = self.chat_state_handle.get_total_tokens().await;
         tracing::Span::current().record("pre_tokens", tokens_before as i64);
@@ -2058,11 +2084,16 @@ impl SessionActor {
                 let span = tracing::Span::current();
                 span.record("success", false);
                 span.record("error", e.to_string().as_str());
-                if self
-                    .compaction
-                    .auto_compact_suppressed
-                    .load(std::sync::atomic::Ordering::Relaxed)
-                    == SUPPRESS_NONE
+                let cancelled = self.compaction.cancel.is_cancelled()
+                    || e.data.as_ref().and_then(|d| d.as_str()).is_some_and(|s| {
+                        s.contains(crate::session::helpers::session_compact::COMPACT_CANCELLED_MSG)
+                    });
+                if !cancelled
+                    && self
+                        .compaction
+                        .auto_compact_suppressed
+                        .load(std::sync::atomic::Ordering::Relaxed)
+                        == SUPPRESS_NONE
                 {
                     self.send_xai_notification(XaiSessionUpdate::AutoCompactFailed {
                         error: String::new(),
@@ -2328,6 +2359,7 @@ mod inline_auto_compact_flow_tests {
                 tool_choice: crate::util::config::CompactionToolChoice::Auto,
                 prefire: crate::session::compaction_config::PrefireState::default(),
                 prefix_released: std::sync::atomic::AtomicBool::new(false),
+                cancel: Default::default(),
             },
             memory: crate::session::memory_state::SessionMemory {
                 flush_config: crate::config::MemoryFlushConfig::default(),
@@ -3683,6 +3715,7 @@ mod inline_auto_compact_flow_tests {
             empty_response_context: None,
             doom_loop_triggers: None,
             doom_loop_aborted_at_chunk: None,
+            credential: xai_grok_sampling_types::SentCredential::Unknown,
         }
     }
     /// Primary scenario: remote settings shrinks the context window mid-session.
@@ -3739,6 +3772,7 @@ mod inline_auto_compact_flow_tests {
                     empty_response_context: None,
                     doom_loop_triggers: None,
                     doom_loop_aborted_at_chunk: None,
+                    credential: xai_grok_sampling_types::SentCredential::Unknown,
                 };
                 assert!(!actor.should_compact_on_error(&err).await);
             })

crates/codegen/xai-grok-shell/src/session/compaction_config.rs

@@ -6,6 +6,7 @@ use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::AtomicU8;
 use std::sync::atomic::AtomicU64;
+use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 
 /// Auto-compaction is gated whenever `auto_compact_suppressed` is not [`SUPPRESS_NONE`].
@@ -53,6 +54,58 @@ pub struct AsyncCompactionCache {
     pub pass1_latency_ms: u64,
 }
 
+/// Cancel gate for an in-flight compact / prefire sample.
+///
+/// Holder count (not a bool): prefire and compact can overlap. The first
+/// `enter` installs a token; nested enters reuse it; `in_flight` stays true
+/// until the last scope drops. A normal turn stop is a no-op when idle.
+#[derive(Default)]
+pub struct CompactCancelGate {
+    token: RefCell<tokio_util::sync::CancellationToken>,
+    holders: AtomicUsize,
+}
+
+/// Decrements the holder count when a compact/prefire scope ends.
+pub struct CompactCancelScope<'a>(&'a CompactCancelGate);
+
+impl Drop for CompactCancelScope<'_> {
+    fn drop(&mut self) {
+        self.0.end();
+    }
+}
+
+impl CompactCancelGate {
+    /// Start or join a compact/prefire scope. Nested callers share one token,
+    /// including a token already cancelled by stop, so overlapping prefire +
+    /// compact both observe the same abort. A later independent enter after
+    /// holders drain installs a fresh token.
+    pub fn enter(&self) -> (tokio_util::sync::CancellationToken, CompactCancelScope<'_>) {
+        let prev = self.holders.fetch_add(1, Ordering::AcqRel);
+        let token = if prev == 0 {
+            let token = tokio_util::sync::CancellationToken::new();
+            self.token.replace(token.clone());
+            token
+        } else {
+            self.token.borrow().clone()
+        };
+        (token, CompactCancelScope(self))
+    }
+
+    fn end(&self) {
+        self.holders.fetch_sub(1, Ordering::AcqRel);
+    }
+
+    pub fn request_cancel(&self) {
+        if self.holders.load(Ordering::Acquire) > 0 {
+            self.token.borrow().cancel();
+        }
+    }
+
+    pub fn is_cancelled(&self) -> bool {
+        self.holders.load(Ordering::Acquire) > 0 && self.token.borrow().is_cancelled()
+    }
+}
+
 /// Prefire two-pass state. `Default` so it drops into existing `CompactionConfig`
 /// struct literals with a single `prefire: PrefireState::default()` field.
 ///
@@ -152,6 +205,8 @@ pub struct CompactionConfig {
     pub prefire: PrefireState,
     /// Sticky once a forked session releases its inherited prefix under compaction pressure (see `run_compact_inner`), so it stops re-pinning it.
     pub prefix_released: AtomicBool,
+    /// User/stop cancel for the current compact generation.
+    pub cancel: CompactCancelGate,
 }
 
 #[cfg(test)]
@@ -209,3 +264,62 @@ mod prefire_state_tests {
         assert!(state.take().is_none());
     }
 }
+
+#[cfg(test)]
+mod compact_cancel_gate_tests {
+    use super::*;
+
+    #[test]
+    fn request_cancel_trips_shared_token() {
+        let gate = CompactCancelGate::default();
+        let (token, _scope) = gate.enter();
+        assert!(!token.is_cancelled());
+        gate.request_cancel();
+        assert!(token.is_cancelled());
+        assert!(gate.is_cancelled());
+    }
+
+    #[test]
+    fn request_cancel_is_noop_when_idle() {
+        let gate = CompactCancelGate::default();
+        gate.request_cancel();
+        let (token, _scope) = gate.enter();
+        assert!(!token.is_cancelled());
+        assert!(!gate.is_cancelled());
+    }
+
+    #[test]
+    fn nested_enter_keeps_in_flight_after_inner_drop() {
+        let gate = CompactCancelGate::default();
+        let (outer_tok, outer) = gate.enter();
+        let (inner_tok, inner) = gate.enter();
+        gate.request_cancel();
+        assert!(outer_tok.is_cancelled());
+        assert!(inner_tok.is_cancelled());
+        drop(inner);
+        assert!(gate.is_cancelled());
+        drop(outer);
+        assert!(!gate.is_cancelled());
+        let (next, _scope) = gate.enter();
+        assert!(!next.is_cancelled());
+    }
+
+    #[test]
+    fn join_while_cancelled_reuses_cancelled_token() {
+        let gate = CompactCancelGate::default();
+        let (_outer, outer) = gate.enter();
+        gate.request_cancel();
+        let (joined, joined_scope) = gate.enter();
+        assert!(
+            joined.is_cancelled(),
+            "nested enter during stop must keep sharing the cancelled token"
+        );
+        drop(joined_scope);
+        drop(outer);
+        let (next, _scope) = gate.enter();
+        assert!(
+            !next.is_cancelled(),
+            "fresh enter after scopes drain must not inherit the prior stop"
+        );
+    }
+}

crates/codegen/xai-grok-shell/src/session/helpers/full_replace_compaction.rs

@@ -72,6 +72,7 @@ pub(crate) struct ShellCompactionSampler {
     /// reasoning-runaway backstop; `0` disables it.
     wall_clock_budget_secs: u64,
     tool_choice: crate::util::config::CompactionToolChoice,
+    cancel: tokio_util::sync::CancellationToken,
     /// Full output of the most recent successful sample (for L5 telemetry).
     last_success: Mutex<Option<CompactOutput>>,
 }
@@ -89,6 +90,7 @@ impl ShellCompactionSampler {
         idle_timeout: Duration,
         wall_clock_budget_secs: u64,
         tool_choice: crate::util::config::CompactionToolChoice,
+        cancel: tokio_util::sync::CancellationToken,
     ) -> Self {
         Self {
             use_short_prompt,
@@ -101,6 +103,7 @@ impl ShellCompactionSampler {
             idle_timeout,
             wall_clock_budget_secs,
             tool_choice,
+            cancel,
             last_success: Mutex::new(None),
         }
     }
@@ -140,6 +143,7 @@ impl CompactionSampler for ShellCompactionSampler {
             self.idle_timeout,
             self.wall_clock_budget_secs,
             self.tool_choice,
+            &self.cancel,
         )
         .await
         {
@@ -170,6 +174,7 @@ fn compact_failure_to_sample_error(failure: CompactFailure) -> CompactionSampleE
     let (deterministic, err) = match failure {
         CompactFailure::Deterministic(err) => (true, err),
         CompactFailure::Transient(err) => (false, err),
+        CompactFailure::Cancelled => (true, CompactFailure::cancelled_error()),
     };
     let message = acp_error_message(&err);
     if deterministic {

crates/codegen/xai-grok-shell/src/session/helpers/session_compact.rs

@@ -45,6 +45,15 @@ pub(crate) enum CompactFailure {
     /// Failure may resolve on retry. The caller follows its existing
     /// N-attempt + backoff loop.
     Transient(acp::Error),
+    /// User/stop cancelled the in-flight compact. Do not retry or suppress AUTO.
+    Cancelled,
+}
+/// Stable error payload for a user-cancelled compact (pager + retry loop).
+pub(crate) const COMPACT_CANCELLED_MSG: &str = "compact cancelled";
+impl CompactFailure {
+    pub(crate) fn cancelled_error() -> acp::Error {
+        acp::Error::internal_error().data(COMPACT_CANCELLED_MSG)
+    }
 }
 pub(crate) use xai_grok_sampling_types::is_context_length_error;
 /// Classify an upstream `SamplingError` for the compaction retry loop.
@@ -58,7 +67,7 @@ pub(crate) use xai_grok_sampling_types::is_context_length_error;
 fn classify_sampling_error(err: SamplingError) -> CompactFailure {
     let acp_err = acp::Error::internal_error().data(format!("compact failed: {err}"));
     let deterministic = match &err {
-        SamplingError::Auth(_)
+        SamplingError::Auth { .. }
         | SamplingError::InvalidConfiguration(_)
         | SamplingError::Serialization(_)
         | SamplingError::IdleTimeout { .. } => true,
@@ -315,14 +324,75 @@ enum StreamStep<T> {
     Ended,
     IdleTimeout,
 }
-async fn next_stream_step<S, T>(stream: &mut S, idle_timeout: std::time::Duration) -> StreamStep<T>
+async fn next_stream_step<S, T>(
+    stream: &mut S,
+    idle_timeout: std::time::Duration,
+    cancel: &tokio_util::sync::CancellationToken,
+) -> Result<StreamStep<T>, CompactFailure>
 where
     S: futures_util::Stream<Item = T> + Unpin,
 {
-    match tokio::time::timeout(idle_timeout, stream.next()).await {
-        Ok(Some(item)) => StreamStep::Item(item),
-        Ok(None) => StreamStep::Ended,
-        Err(_) => StreamStep::IdleTimeout,
+    tokio::select! {
+        biased;
+        _ = cancel.cancelled() => Err(CompactFailure::Cancelled),
+        step = tokio::time::timeout(idle_timeout, stream.next()) => Ok(match step {
+            Ok(Some(item)) => StreamStep::Item(item),
+            Ok(None) => StreamStep::Ended,
+            Err(_) => StreamStep::IdleTimeout,
+        }),
+    }
+}
+/// Abort `fut` if stop wins while the compact HTTP stream is still opening.
+async fn await_unless_cancelled<F, T>(
+    cancel: &tokio_util::sync::CancellationToken,
+    fut: F,
+) -> Result<T, CompactFailure>
+where
+    F: std::future::Future<Output = T>,
+{
+    tokio::select! {
+        biased;
+        _ = cancel.cancelled() => Err(CompactFailure::Cancelled),
+        result = fut => Ok(result),
+    }
+}
+#[cfg(test)]
+mod compact_cancel_await_tests {
+    use super::*;
+    use std::time::Duration;
+    use tokio_util::sync::CancellationToken;
+    #[tokio::test]
+    async fn pre_cancelled_token_skips_fut() {
+        let cancel = CancellationToken::new();
+        cancel.cancel();
+        let err = await_unless_cancelled(&cancel, async {
+            panic!("fut must not run when already cancelled");
+        })
+        .await
+        .unwrap_err();
+        assert!(matches!(err, CompactFailure::Cancelled));
+    }
+    #[tokio::test]
+    async fn cancel_aborts_pending_open() {
+        let cancel = CancellationToken::new();
+        let cancel2 = cancel.clone();
+        tokio::spawn(async move {
+            tokio::time::sleep(Duration::from_millis(20)).await;
+            cancel2.cancel();
+        });
+        let started = std::time::Instant::now();
+        let err = await_unless_cancelled(&cancel, async {
+            tokio::time::sleep(Duration::from_secs(30)).await;
+            0u8
+        })
+        .await
+        .unwrap_err();
+        assert!(matches!(err, CompactFailure::Cancelled));
+        assert!(
+            started.elapsed() < Duration::from_secs(2),
+            "stop must abort stream-open wait, elapsed {:?}",
+            started.elapsed()
+        );
     }
 }
 /// Generates a summary of the conversation for compaction.
@@ -355,7 +425,11 @@ pub(crate) async fn generate_session_compact(
     idle_timeout: std::time::Duration,
     wall_clock_budget_secs: u64,
     tool_choice: crate::util::config::CompactionToolChoice,
+    cancel: &tokio_util::sync::CancellationToken,
 ) -> Result<CompactOutput, CompactFailure> {
+    if cancel.is_cancelled() {
+        return Err(CompactFailure::Cancelled);
+    }
     let num_messages = chat_history.len();
     let wire_tool_choice = match tool_choice {
         crate::util::config::CompactionToolChoice::Auto => ToolChoice::auto(),
@@ -392,7 +466,8 @@ pub(crate) async fn generate_session_compact(
                 num_messages = num_messages,
                 "Sending compact request (streaming)"
             );
-            let stream_result = client.chat_completion_stream(message).await;
+            let stream_result =
+                await_unless_cancelled(cancel, client.chat_completion_stream(message)).await?;
             let mut stream = match stream_result {
                 Ok((s, _metadata)) => s,
                 Err(e) => return Err(classify_sampling_error(e)),
@@ -404,7 +479,9 @@ pub(crate) async fn generate_session_compact(
             let mut last_progress_at = std::time::Instant::now();
             loop {
                 let idle_remaining = idle_timeout.saturating_sub(last_progress_at.elapsed());
-                let chunk_result = match next_stream_step(&mut stream, idle_remaining).await {
+                let chunk_result = match next_stream_step(&mut stream, idle_remaining, cancel)
+                    .await?
+                {
                     StreamStep::Item(item) => item,
                     StreamStep::Ended => break,
                     StreamStep::IdleTimeout => {
@@ -486,7 +563,9 @@ pub(crate) async fn generate_session_compact(
                 x_grok_agent_id: Some(xai_grok_telemetry::id::agent_id()),
                 ..Default::default()
             };
-            let stream_result = client.conversation_stream_responses(request).await;
+            let stream_result =
+                await_unless_cancelled(cancel, client.conversation_stream_responses(request))
+                    .await?;
             let mut stream = match stream_result {
                 Ok((s, _metadata, _doom_loop)) => s,
                 Err(e) => return Err(classify_sampling_error(e)),
@@ -498,7 +577,9 @@ pub(crate) async fn generate_session_compact(
             let mut last_progress_at = std::time::Instant::now();
             loop {
                 let idle_remaining = idle_timeout.saturating_sub(last_progress_at.elapsed());
-                let chunk_result = match next_stream_step(&mut stream, idle_remaining).await {
+                let chunk_result = match next_stream_step(&mut stream, idle_remaining, cancel)
+                    .await?
+                {
                     StreamStep::Item(item) => item,
                     StreamStep::Ended => break,
                     StreamStep::IdleTimeout => {
@@ -611,7 +692,9 @@ pub(crate) async fn generate_session_compact(
                 x_grok_agent_id: Some(xai_grok_telemetry::id::agent_id()),
                 ..Default::default()
             };
-            let stream_result = client.conversation_stream_messages(request).await;
+            let stream_result =
+                await_unless_cancelled(cancel, client.conversation_stream_messages(request))
+                    .await?;
             let mut stream = match stream_result {
                 Ok((s, _metadata)) => s,
                 Err(e) => return Err(classify_sampling_error(e)),
@@ -623,7 +706,9 @@ pub(crate) async fn generate_session_compact(
             let mut last_progress_at = std::time::Instant::now();
             loop {
                 let idle_remaining = idle_timeout.saturating_sub(last_progress_at.elapsed());
-                let chunk_result = match next_stream_step(&mut stream, idle_remaining).await {
+                let chunk_result = match next_stream_step(&mut stream, idle_remaining, cancel)
+                    .await?
+                {
                     StreamStep::Item(item) => item,
                     StreamStep::Ended => break,
                     StreamStep::IdleTimeout => {
@@ -771,9 +856,9 @@ mod classify_tests {
     }
     #[test]
     fn sampling_non_api_variants_classify_correctly() {
-        assert!(is_det(&classify_sampling_error(SamplingError::Auth(
-            "expired".into()
-        ))));
+        assert!(is_det(&classify_sampling_error(
+            SamplingError::auth_unknown("expired")
+        )));
         assert!(is_det(&classify_sampling_error(
             SamplingError::InvalidConfiguration("missing key")
         )));
@@ -1619,6 +1704,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_secs(30),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await
         .unwrap_or_else(|_| panic!("compaction must succeed"));
@@ -1710,6 +1796,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_secs(30),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await;
         let output = result
@@ -1772,6 +1859,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_secs(30),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await
         .unwrap_or_else(|_| panic!("compaction with tools must succeed"));
@@ -1786,6 +1874,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_secs(30),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await
         .unwrap_or_else(|_| panic!("compaction without tools must succeed"));
@@ -1919,6 +2008,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_secs(30),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await
         .unwrap_or_else(|_| panic!("Responses compaction with tools must succeed"));
@@ -1933,6 +2023,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_secs(30),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await
         .unwrap_or_else(|_| panic!("Responses compaction without tools must succeed"));
@@ -2012,6 +2103,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_millis(150),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await;
         match result {
@@ -2026,8 +2118,10 @@ mod reasoning_compaction_regression_tests {
                     "expected an idle-timeout transient failure, got: {data}"
                 );
             }
-            Err(CompactFailure::Deterministic(_)) => {
-                panic!("a stalled stream must be retryable (Transient), not Deterministic")
+            Err(CompactFailure::Deterministic(_) | CompactFailure::Cancelled) => {
+                panic!(
+                    "a stalled stream must be retryable (Transient), not Deterministic/Cancelled"
+                )
             }
             Ok(_) => panic!("a stalled stream must not produce a summary"),
         }
@@ -2091,6 +2185,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_millis(150),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await;
         match result {
@@ -2105,7 +2200,7 @@ mod reasoning_compaction_regression_tests {
                     "expected an idle-timeout transient failure, got: {data}"
                 );
             }
-            Err(CompactFailure::Deterministic(_)) => {
+            Err(CompactFailure::Deterministic(_) | CompactFailure::Cancelled) => {
                 panic!("a stalled stream must be retryable (Transient), not Deterministic")
             }
             Ok(_) => {
@@ -2169,6 +2264,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_millis(150),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await;
         match result {
@@ -2183,7 +2279,7 @@ mod reasoning_compaction_regression_tests {
                     "expected an idle-timeout transient failure, got: {data}"
                 );
             }
-            Err(CompactFailure::Deterministic(_)) => {
+            Err(CompactFailure::Deterministic(_) | CompactFailure::Cancelled) => {
                 panic!("a stalled stream must be retryable (Transient), not Deterministic")
             }
             Ok(_) => {
@@ -2244,6 +2340,7 @@ mod reasoning_compaction_regression_tests {
             std::time::Duration::from_millis(150),
             0,
             crate::util::config::CompactionToolChoice::Auto,
+            &tokio_util::sync::CancellationToken::new(),
         )
         .await;
         match result {

crates/codegen/xai-grok-shell/src/session/persistence.rs

@@ -73,6 +73,14 @@ pub struct BtwEntry {
     /// Error message if failed.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub error: Option<String>,
+    /// Model-call attempts made (1 = no retry). Entries written before this
+    /// field existed deserialize as 1.
+    #[serde(default = "default_btw_attempts")]
+    pub attempts: u32,
+}
+
+fn default_btw_attempts() -> u32 {
+    1
 }
 
 // Local feedback persistence types

crates/codegen/xai-grok-shell/src/session/unified_list/mod.rs

@@ -55,10 +55,18 @@ pub fn conversations_lane_active() -> bool {
 }
 /// Parse `x.ai/session/list` params and, under process-wide chat mode, force
 /// the conversations-only `kind` facet (see [`force_kind_chat`]).
+///
+/// Client-sent `kind` of `chat`/`build` is honored only behind
+/// `feature = "local-workspace"` (pager welcome Local history). Chat-only
+/// Desktop/ACP agents keep the force-rewrite so `kind: ["build"]` cannot
+/// surface Build rows.
 pub fn parse_list_req(raw: &str) -> Result<ListReq, serde_json::Error> {
     let mut req: ListReq = serde_json::from_str(raw)?;
     if crate::agent::chat_modes::process_chat_mode_enabled() {
-        force_kind_chat(&mut req);
+        let honor_client_kind = cfg!(feature = "local-workspace") && client_sent_kind_filter(&req);
+        if !honor_client_kind {
+            force_kind_chat(&mut req);
+        }
     }
     Ok(req)
 }
@@ -72,6 +80,23 @@ where
         CwdScope::WithSiblings
     })
 }
+fn client_sent_kind_filter(req: &ListReq) -> bool {
+    let Some(kind) = req
+        .meta
+        .as_ref()
+        .and_then(|m| m.get("x.ai/facetFilters"))
+        .and_then(|f| f.get("kind"))
+    else {
+        return false;
+    };
+    match kind {
+        serde_json::Value::Array(arr) if !arr.is_empty() => arr
+            .iter()
+            .any(|v| matches!(v.as_str(), Some("chat" | "build"))),
+        serde_json::Value::String(s) if s == "chat" || s == "build" => true,
+        _ => false,
+    }
+}
 #[derive(Debug, Default, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct ListReq {
@@ -172,6 +197,10 @@ fn value_list(v: &serde_json::Value) -> Vec<serde_json::Value> {
 }
 /// Rewrite `req` so the `kind` facet filter is exactly `["chat"]`.
 ///
+/// Used when process chat mode is on **and** the client omitted a recognized
+/// `kind` facet (see [`parse_list_req`]). Welcome history sends an explicit
+/// `kind` (`chat` / `build`) that must not be rewritten. Other facet filters
+/// and `_meta` keys are left untouched.
 pub fn force_kind_chat(req: &mut ListReq) {
     force_kind(req, SessionKind::Chat);
 }
@@ -201,7 +230,7 @@ pub async fn build_unified_list(
     conversations_client: Option<&ConversationsClient>,
     mut req: ListReq,
 ) -> UnifiedListResult {
-    if crate::agent::chat_modes::process_chat_mode_enabled() {
+    if crate::agent::chat_modes::process_chat_mode_enabled() && !client_sent_kind_filter(&req) {
         force_kind_chat(&mut req);
     }
     let reg = facet_registry();
@@ -936,16 +965,42 @@ mod tests {
             let _on = xai_grok_test_support::EnvGuard::set(GROK_CHAT_MODE_ENV, "1");
             let req = parse_list_req(&raw).expect("parse");
             let parsed = ParsedMeta::parse(req.meta.as_ref());
-            let expected = "build";
+            let expected_build = if cfg!(feature = "local-workspace") {
+                Some(&vec![serde_json::json!("build")])
+            } else {
+                Some(&vec![serde_json::json!("build")])
+            };
             assert_eq!(
                 parsed.facet_filters.get(KIND_FACET_KEY),
-                Some(&vec![serde_json::json!(expected)])
+                expected_build,
+                "client kind=build under process chat mode"
             );
             assert_eq!(
                 parsed.facet_filters.get("starred"),
                 Some(&vec![serde_json::json!(true)]),
                 "other facets pass through"
             );
+            let req = parse_list_req("{}").expect("parse");
+            let parsed = ParsedMeta::parse(req.meta.as_ref());
+            let expected = None;
+            assert_eq!(
+                parsed.facet_filters.get(KIND_FACET_KEY),
+                expected,
+                "absent client kind still forces chat under process chat mode"
+            );
+            for bad in [
+                serde_json::json!({ "_meta": { "x.ai/facetFilters": { "kind": [] } } }),
+                serde_json::json!({ "_meta": { "x.ai/facetFilters": { "kind": null } } }),
+                serde_json::json!({ "_meta": { "x.ai/facetFilters": { "kind": ["other"] } } }),
+            ] {
+                let req = parse_list_req(&bad.to_string()).expect("parse");
+                let parsed = ParsedMeta::parse(req.meta.as_ref());
+                assert_eq!(
+                    parsed.facet_filters.get(KIND_FACET_KEY),
+                    expected,
+                    "empty/null/unknown kind must still force chat: {bad}"
+                );
+            }
         }
     }
     /// Wire pin for the cross-crate `x.ai/partial` envelope the pager parses:

crates/codegen/xai-grok-shell/src/terminal/pty_session.rs

@@ -5,6 +5,10 @@
 //! on its own leaves them running, as a terminal does: their process groups are
 //! their own and nothing here holds a handle to them.
 
+// A panic here loses a shell: teardown paths run inside `Drop`, where an
+// unwind during another unwind aborts the process. Tests panic freely.
+#![cfg_attr(not(test), deny(clippy::unwrap_used, clippy::expect_used))]
+
 use std::collections::{HashMap, VecDeque};
 use std::io::{Read, Write};
 use std::sync::{Arc, LazyLock};
@@ -147,6 +151,53 @@ impl Shell {
     }
 }
 
+/// A shell not yet in the registry, where teardown would never find it.
+/// Dropping reaps it; [`Self::into_registered`] hands it to the registry
+/// instead.
+///
+/// Disarming leaves [`Shell::Reaped`] behind rather than an empty slot, so the
+/// guard has no state in which its own field is missing.
+struct UnregisteredShell(Shell);
+
+impl UnregisteredShell {
+    fn new(child: Box<dyn portable_pty::Child + Send + Sync>) -> Self {
+        Self(Shell::Running { child, group: None })
+    }
+
+    fn pid(&self) -> Option<u32> {
+        self.0.pid()
+    }
+
+    fn attach_group(&mut self, enrolled: Arc<xai_tty_utils::ProcessGroup>) {
+        self.0.attach_group(enrolled);
+    }
+
+    /// Disarms the guard. The caller must reach the registry without awaiting,
+    /// or it reopens the window this type closes.
+    fn into_registered(mut self) -> Shell {
+        self.disarm()
+    }
+
+    fn disarm(&mut self) -> Shell {
+        std::mem::replace(&mut self.0, Shell::Reaped(None))
+    }
+}
+
+impl Drop for UnregisteredShell {
+    fn drop(&mut self) {
+        let mut shell = self.disarm();
+        // `reap_now` blocks through its grace waits, so keep it off an async
+        // thread. Not the runtime's blocking pool though: a task still queued
+        // there at shutdown is dropped unrun, taking the group with it and
+        // leaving the scope holding a dead `Weak`.
+        if tokio::runtime::Handle::try_current().is_ok() {
+            std::thread::spawn(move || shell.reap_now());
+        } else {
+            shell.reap_now();
+        }
+    }
+}
+
 type PtyMap = HashMap<String, Arc<Mutex<PtySession>>>;
 
 static PTY_REGISTRY: LazyLock<Mutex<PtyMap>> = LazyLock::new(|| Mutex::new(HashMap::new()));
@@ -208,44 +259,34 @@ pub async fn create_pty(
     cmd.env("LANG", "en_US.UTF-8");
     cmd.env("LC_ALL", "en_US.UTF-8");
 
+    // Enrolled below, and reaped by the guard until it reaches the registry.
+    #[allow(clippy::disallowed_methods)]
     let child = pair
         .slave
         .spawn_command(cmd)
         .map_err(|e| TerminalExtError::Internal(format!("failed to spawn shell: {e}")))?;
 
-    // Until the session reaches the registry nothing else can reach this shell,
-    // so every failure below has to kill it here or it is orphaned.
-    let mut shell = Shell::Running { child, group: None };
+    let mut shell = UnregisteredShell::new(child);
     if let Some(pid) = shell.pid() {
-        match xai_tty_utils::global_process_scope().enroll_terminal_pid(pid) {
-            Ok(enrolled) => shell.attach_group(enrolled),
-            Err(e) => {
-                shell.reap_now();
-                return Err(TerminalExtError::Internal(format!(
-                    "failed to enroll shell: {e}"
-                )));
-            }
-        }
+        // `enroll_terminal_pid` reaps with a grace wait if it loses the close
+        // race, so run it off-task.
+        let enrolled = tokio::task::spawn_blocking(move || {
+            xai_tty_utils::global_process_scope().enroll_terminal_pid(pid)
+        })
+        .await
+        .map_err(|e| TerminalExtError::Internal(format!("enroll task failed: {e}")))?
+        .map_err(|e| TerminalExtError::Internal(format!("failed to enroll shell: {e}")))?;
+        shell.attach_group(enrolled);
     }
 
-    let reader = match pair.master.try_clone_reader() {
-        Ok(reader) => reader,
-        Err(e) => {
-            shell.reap_now();
-            return Err(TerminalExtError::Internal(format!(
-                "failed to clone pty reader: {e}"
-            )));
-        }
-    };
-    let writer = match pair.master.take_writer() {
-        Ok(writer) => writer,
-        Err(e) => {
-            shell.reap_now();
-            return Err(TerminalExtError::Internal(format!(
-                "failed to take pty writer: {e}"
-            )));
-        }
-    };
+    let reader = pair
+        .master
+        .try_clone_reader()
+        .map_err(|e| TerminalExtError::Internal(format!("failed to clone pty reader: {e}")))?;
+    let writer = pair
+        .master
+        .take_writer()
+        .map_err(|e| TerminalExtError::Internal(format!("failed to take pty writer: {e}")))?;
 
     let (input_tx, input_rx) = mpsc::channel(INPUT_CHANNEL_CAPACITY);
     spawn_pty_input_loop(writer, input_rx);
@@ -277,12 +318,24 @@ pub async fn create_pty(
             })
     });
 
-    let mut session = PtySession {
+    // Taking the lock first means nothing can await between disarming the guard
+    // and the insert that gives teardown another way to reach the shell.
+    let mut registry = PTY_REGISTRY.lock().await;
+
+    // The scope can close during the setup above, and teardown has already run
+    // by then: publishing here would advertise a shell it just killed.
+    if xai_tty_utils::global_process_scope().is_closed() {
+        return Err(TerminalExtError::Internal(
+            "process scope closed while the shell was starting".to_string(),
+        ));
+    }
+
+    let entry = Arc::new(Mutex::new(PtySession {
         master: Some(pair.master),
         input_tx: Some(input_tx),
         output_offset: 0,
         output_ring: VecDeque::with_capacity(OUTPUT_RING_BUFFER_SIZE),
-        shell,
+        shell: shell.into_registered(),
         cwd: resolved_cwd,
         name: resolved_name,
         created_at,
@@ -291,22 +344,9 @@ pub async fn create_pty(
         target_client_id,
         busy: false,
         gateway: gateway.clone(),
-    };
-
-    // The scope can close during the setup above, and teardown has already run
-    // by then: publishing here would advertise a shell it just killed.
-    if xai_tty_utils::global_process_scope().is_closed() {
-        session.shell.reap_now();
-        return Err(TerminalExtError::Internal(
-            "process scope closed while the shell was starting".to_string(),
-        ));
-    }
-
-    let entry = Arc::new(Mutex::new(session));
-    PTY_REGISTRY
-        .lock()
-        .await
-        .insert(pty_id.clone(), entry.clone());
+    }));
+    registry.insert(pty_id.clone(), entry.clone());
+    drop(registry);
 
     let pty_id_clone = pty_id.clone();
     tokio::task::spawn_local(run_pty_output_loop(reader, entry, pty_id_clone, gateway));
@@ -447,21 +487,11 @@ async fn run_pty_output_loop(
     );
 }
 
-/// Whether the PTY's controlling terminal has a foreground process group
-/// distinct from the shell itself — i.e. a command is actively running
-/// rather than the shell sitting idle at its prompt.
+/// Whether a command is running, rather than the shell sitting at its prompt.
 ///
-/// `process_group_leader()` issues `tcgetpgrp` on the master fd; an idle
-/// shell is its own foreground process group, so it matches the shell
-/// child's pid. When a command runs in the foreground the kernel reports
-/// that command's process group instead. Returns false when the value is
-/// unavailable (the shell exited or runs without job control).
-///
-/// Limitation: a shell that `exec`s a program in place keeps the same pid and
-/// pgid, so `tcgetpgrp` still matches the recorded child pid and the program
-/// reads as idle. Telling that apart from a real idle prompt needs per-OS
-/// process inspection, so a command launched the usual way (fork then exec) is
-/// detected while an `exec`-replaced shell is not.
+/// A shell that `exec`s a program in place keeps its pid and pgid, so that
+/// program reads as idle. Separating it from a real prompt needs per-OS
+/// process inspection.
 #[cfg(unix)]
 fn session_has_foreground_process(session: &PtySession) -> bool {
     let Some(foreground_pgid) = session
@@ -642,12 +672,8 @@ pub async fn close_all() {
     }
 }
 
-/// Resolve the shell binary and arguments for an interactive PTY session.
-///
-/// Priority: explicit `shell` param > `$SHELL` env > platform default.
-/// On Windows falls back to the `detect_windows_shell` cascade
-/// (pwsh > powershell.exe > Git Bash > cmd.exe, overridable via
-/// `GROK_SHELL`) since `$SHELL` is absent.
+/// Explicit `shell` param, then `$SHELL`, then the platform default. Windows
+/// has no `$SHELL`, so it uses the `detect_windows_shell` cascade.
 fn resolve_pty_shell(shell: Option<&str>) -> (String, Vec<String>) {
     if let Some(s) = shell {
         return (s.to_string(), vec![]);
@@ -712,12 +738,9 @@ pub struct PtyLoadResult {
     pub exit_code: Option<i32>,
 }
 
-/// Reconnect to a PTY. Replays the full ring buffer (with `isReplay: true`)
-/// so the client can reset its VTE emulator and feed all bytes from scratch.
-/// Exited PTYs are still loadable so the client can see final output.
-///
-/// Updates the stored `target_client_id` so that subsequent output
-/// notifications from the output loop are routed to the reconnecting client.
+/// Reconnect to a PTY, replaying the ring buffer so the client can reset its
+/// VTE emulator and feed all bytes from scratch. An exited PTY still loads, so
+/// its final output stays readable.
 pub async fn load(
     pty_id: &str,
     gateway: &GatewaySender,
@@ -1005,6 +1028,36 @@ mod tests {
             .await;
     }
 
+    /// Covers the failure paths in [`create_pty`], which reap by returning.
+    #[tokio::test]
+    async fn dropping_an_unregistered_shell_reaps_it() {
+        let pair = native_pty_system()
+            .openpty(PtySize {
+                rows: 24,
+                cols: 80,
+                pixel_width: 0,
+                pixel_height: 0,
+            })
+            .expect("openpty");
+        let mut cmd = CommandBuilder::new("/bin/sh");
+        cmd.arg("-c");
+        cmd.arg("sleep 300");
+        #[allow(clippy::disallowed_methods)]
+        let child = pair.slave.spawn_command(cmd).expect("spawn shell");
+        let pid = child.process_id().expect("shell pid") as i32;
+
+        drop(UnregisteredShell::new(child));
+
+        let deadline = std::time::Instant::now() + Duration::from_secs(5);
+        while unsafe { libc::kill(pid, 0) } == 0 {
+            assert!(
+                std::time::Instant::now() < deadline,
+                "shell {pid} survived the guard drop"
+            );
+            tokio::time::sleep(Duration::from_millis(50)).await;
+        }
+    }
+
     #[cfg(unix)]
     async fn wait_for_reported_pid(pty_id: &str) -> i32 {
         let deadline = std::time::Instant::now() + Duration::from_secs(10);

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

@@ -2,6 +2,15 @@ pub(crate) mod lsp_runtime;
 
 pub(crate) const TEST_MODEL: &str = "test-model";
 
+/// Keep this crate's unit-test binary from writing synthetic events into
+/// the real unified log; pre-main so the redirect beats the lazily-opened
+/// writer. Integration binaries under `tests/` isolate via `TestSandbox`
+/// homes instead.
+#[ctor::ctor]
+fn redirect_unified_log_for_tests() {
+    xai_grok_telemetry::unified_log::redirect_to_temp_for_tests();
+}
+
 /// Prepend the hermetic git binary (via `GIT_BIN_PATH`) to `PATH` so that
 /// `Command::new("git")` in test helpers resolves to the Bazel-provided
 /// static binary instead of relying on system-installed git.

crates/codegen/xai-grok-shell/src/util/dual_clock.rs

@@ -0,0 +1,52 @@
+//! One instant captured on two clocks, so elapsed time stays honest across
+//! a system suspend without trusting the wall clock alone.
+
+use std::time::{Duration, Instant, SystemTime};
+
+/// `Instant` is monotonic but, on macOS (`mach_absolute_time`) and Linux
+/// (`CLOCK_MONOTONIC`), *pauses while the machine is asleep*, so it alone
+/// under-reports any span containing a suspend. `SystemTime` keeps advancing
+/// through sleep but jumps with NTP steps and manual changes. Capturing both
+/// lets a caller bound elapsed *awake* time (mono), elapsed *real* time
+/// (wall), and their difference — which grows by exactly the suspended time.
+#[derive(Clone, Copy)]
+pub(crate) struct DualClock {
+    /// Monotonic; pauses during sleep. Bounds elapsed *awake* time.
+    pub(crate) mono: Instant,
+    /// Wall clock; advances through sleep. Bounds elapsed *real* time.
+    pub(crate) wall: SystemTime,
+}
+
+impl DualClock {
+    pub(crate) fn now() -> Self {
+        Self {
+            mono: Instant::now(),
+            wall: SystemTime::now(),
+        }
+    }
+
+    /// Elapsed on each clock as `(monotonic, wall)`. Wall elapsed clamps to
+    /// zero if the clock ran backwards (NTP step) so a backward jump can
+    /// never fabricate a suspend or inflate a duration.
+    pub(crate) fn elapsed_between(&self, now: DualClock) -> (Duration, Duration) {
+        (
+            now.mono.saturating_duration_since(self.mono),
+            now.wall.duration_since(self.wall).unwrap_or(Duration::ZERO),
+        )
+    }
+
+    /// [`Self::elapsed_between`] against the live clocks.
+    pub(crate) fn elapsed(&self) -> (Duration, Duration) {
+        self.elapsed_between(Self::now())
+    }
+
+    /// `(awake, total, suspended)` durations since this instant.
+    pub(crate) fn elapsed_split(&self) -> (Duration, Duration, Duration) {
+        let (awake, total) = self.elapsed();
+        (awake, total, total.saturating_sub(awake))
+    }
+}
+
+#[cfg(test)]
+#[path = "dual_clock_tests.rs"]
+mod tests;

crates/codegen/xai-grok-shell/src/util/dual_clock_tests.rs

@@ -0,0 +1,18 @@
+use std::time::Duration;
+
+use super::DualClock;
+
+/// A backward wall jump (NTP step) clamps to zero rather than underflowing,
+/// so it can never fabricate a suspend or inflate a duration.
+#[test]
+fn backward_wall_jump_clamps_to_zero() {
+    let start = DualClock::now();
+    let stepped_back = DualClock {
+        mono: start.mono + Duration::from_secs(5),
+        wall: start.wall - Duration::from_secs(60),
+    };
+    assert_eq!(
+        start.elapsed_between(stepped_back),
+        (Duration::from_secs(5), Duration::ZERO)
+    );
+}

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

@@ -1,4 +1,5 @@
 pub mod config;
+pub(crate) mod dual_clock;
 pub mod grok_auth_credentials;
 pub mod hooks;
 pub mod limits;

crates/codegen/xai-grok-shell/tests/test_leader_soak.rs

@@ -105,6 +105,8 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
     let server = xai_grok_test_support::MockInferenceServer::start()
         .await
         .unwrap();
+    // Measure the leader, not the harness's copy of every conversation.
+    server.set_keep_requests(false);
     let grok_home = TempDir::new().unwrap();
     let workdir = TempDir::new().unwrap();
 
@@ -207,8 +209,6 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
             let mut turns: u64 = 0;
             let mut baseline: Option<serde_json::Value> = None;
 
-            // Each cycle: 10 fresh clients, 2 sessions each, one scripted
-            // turn per session, then all disconnect.
             while tokio::time::Instant::now() < soak_deadline {
                 cycles += 1;
                 let mut clients = Vec::new();
@@ -282,7 +282,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
                 }
 
                 // An entry that never drains names itself here, one cycle
-                // after it leaks, while memory is still within its budget.
+                // after it leaks.
                 let counts = registry_counts(&mut bootstrap, 1000 + cycles).await;
                 assert_eq!(
                     counts["sessions"], 0,
@@ -332,7 +332,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
                      ({:.2} MB per cycle)",
                     net_bytes as f64 / measured as f64 / (1024.0 * 1024.0)
                 );
-                let max_per_cycle = env_u64("LEADER_SOAK_MAX_HEAP_BYTES_PER_CYCLE", 4 << 20) as i64;
+                let max_per_cycle = env_u64("LEADER_SOAK_MAX_HEAP_BYTES_PER_CYCLE", 1 << 20) as i64;
                 assert!(
                     per_cycle <= max_per_cycle,
                     "leader retained {per_cycle} heap bytes per cycle (bound {max_per_cycle})"

crates/codegen/xai-grok-shell/tests/test_registry_churn.rs

@@ -31,6 +31,7 @@ const RPC_TIMEOUT: Duration = Duration::from_secs(60);
 #[serde(deny_unknown_fields)]
 struct Counts {
     sessions: usize,
+    loading_sessions: usize,
     session_threads: usize,
     resident_resources: usize,
     retained_resources: usize,
@@ -45,6 +46,7 @@ struct Counts {
     subagent_active: usize,
     subagent_completed: usize,
     workspace_bindings: Option<usize>,
+    workspace_activity_sessions: Option<usize>,
 }
 struct AutoApproveClient;
 #[async_trait::async_trait(?Send)]
@@ -91,6 +93,19 @@ async fn read_counts(conn: &acp::ClientSideConnection) -> Counts {
     serde_json::from_value(resp["result"]["registries"].clone())
         .unwrap_or_else(|e| panic!("x.ai/debug/agent: bad registries payload: {e}\n{resp}"))
 }
+/// Counts read once the actor threads are reaped. Nothing signals a thread
+/// exit, so this polls; both ends settle, so neither catches one mid-exit.
+async fn settled_counts(conn: &acp::ClientSideConnection) -> Counts {
+    let mut counts = read_counts(conn).await;
+    for _ in 0..100 {
+        if counts.session_threads == 0 {
+            break;
+        }
+        tokio::time::sleep(Duration::from_millis(100)).await;
+        counts = read_counts(conn).await;
+    }
+    counts
+}
 async fn new_session(conn: &acp::ClientSideConnection, cwd: &std::path::Path) -> acp::SessionId {
     tokio::time::timeout(
         RPC_TIMEOUT,
@@ -252,21 +267,29 @@ fn session_churn_returns_registry_snapshot_to_baseline() {
     agent_rt.block_on(local.run_until(async move {
         let client_conn = connect_and_auth().await;
         churn_one(&client_conn, workdir.path(), 0).await;
-        let baseline = read_counts(&client_conn).await;
+        let baseline = settled_counts(&client_conn).await;
         assert_eq!(
             baseline.sessions, 0,
             "warmup session must be fully removed before baseline"
         );
         assert_eq!(
-            (baseline.resident_resources, baseline.retained_resources),
-            (0, 0),
+            (
+                baseline.resident_resources,
+                baseline.retained_resources,
+                baseline.loading_sessions
+            ),
+            (0, 0, 0),
             "warmup must leave no per-session resource entries, including \
              entries holding no resources"
         );
         assert_eq!(
-            baseline.workspace_bindings,
-            Some(0),
-            "warmup must have built the local workspace and released its binding"
+            (
+                baseline.workspace_bindings,
+                baseline.workspace_activity_sessions
+            ),
+            (Some(0), Some(0)),
+            "warmup must have built the local workspace and released both its \
+             binding and its activity record"
         );
         assert_eq!(
             (
@@ -295,7 +318,7 @@ fn session_churn_returns_registry_snapshot_to_baseline() {
         }))
         .await;
         futures::future::join_all(concurrent.iter().map(|sid| close_session(conn, sid))).await;
-        let after = read_counts(&client_conn).await;
+        let after = settled_counts(&client_conn).await;
         assert_eq!(
             after, baseline,
             "session churn must return every registry count to baseline \

crates/codegen/xai-grok-shell/tests/test_sampling_client.rs

@@ -896,7 +896,7 @@ async fn test_chat_completions_401_unauthorized() {
     let result = client.conversation_stream(request).await;
     assert!(result.is_err());
 
-    if let Err(SamplingError::Auth(_)) = result {
+    if let Err(SamplingError::Auth { .. }) = result {
         // Expected
     } else {
         panic!("Expected Auth error");
@@ -939,7 +939,7 @@ async fn test_responses_api_401_unauthorized() {
     let result = client.conversation_stream_responses(request).await;
     assert!(result.is_err());
 
-    if let Err(SamplingError::Auth(_)) = result {
+    if let Err(SamplingError::Auth { .. }) = result {
         // Expected
     } else {
         panic!("Expected Auth error");