Skip to main content

trustee_api/
state.rs

1//! Shared server state: per-user multi-session registry, broadcast channels, and auth state.
2//!
3//! ## Multi-Session Per User (MSU)
4//!
5//! Each authenticated user gets their own [`UserSessions`] containing N independent
6//! [`UserSessionEntry`] instances (default max 4). Each entry has:
7//! - An independent `Session` (workflow state, output, etc.)
8//! - A dedicated broadcast channel for WebSocket fan-out
9//! - Creation and last-active timestamps
10//!
11//! Sessions are keyed by user identity (`sub` claim from JWT, or `dev:email` for
12//! dev mode). Unauthenticated deployments use a single `"default"` key, preserving
13//! backward compatibility with single-user CLI operation.
14
15use std::sync::Arc;
16
17use dashmap::DashMap;
18use tokio::sync::{broadcast, mpsc, Mutex};
19use trustee_core::session::Session;
20use trustee_core::types::TuiMessage;
21
22use crate::auth::AuthState;
23
24// ---------------------------------------------------------------------------
25// Multi-session types
26// ---------------------------------------------------------------------------
27
28/// A single session with its own broadcast channel.
29pub struct UserSessionEntry {
30    /// The agent session, protected by a mutex.
31    pub session: Arc<Mutex<Session>>,
32    /// Broadcast sender for this session's WebSocket fan-out.
33    pub ws_tx: broadcast::Sender<String>,
34    /// When this session was created.
35    pub created_at: chrono::DateTime<chrono::Utc>,
36    /// Last time a command was submitted or state changed.
37    /// Updated on every /sessions/{id}/command and /sessions/{id}/cancel call.
38    pub last_active: Arc<Mutex<chrono::DateTime<chrono::Utc>>>,
39}
40
41/// All sessions belonging to one authenticated user.
42pub struct UserSessions {
43    /// session_id → session entry
44    pub sessions: DashMap<String, UserSessionEntry>,
45    /// Shared token store for all this user's sessions (MCP credential isolation).
46    pub token_store: Arc<pep::MemoryTokenStore>,
47    /// Which session_id is "active" for legacy /session/* routes.
48    pub active_session_id: Mutex<String>,
49}
50
51/// Summary of an active session for listing (serializable for API responses).
52#[derive(Debug, serde::Serialize)]
53pub struct SessionListItem {
54    pub session_id: String,
55    pub session_name: Option<String>,
56    pub workflow_state: String,
57    pub created_at: String,
58    pub last_active: String,
59    /// Number of handoff rotations this session has performed (Bug 5). The
60    /// in-memory registry key never changes across a rotation — only the
61    /// checkpoint-chain name does — so a card whose name changed without this
62    /// hint looks like it silently renamed. 0 = never rotated.
63    pub handoff_count: u32,
64}
65
66/// Errors from multi-session operations.
67#[derive(Debug)]
68pub enum SessionError {
69    /// User has reached max_sessions_per_user limit.
70    MaxSessionsReached(usize),
71    /// Session ID not found for this user.
72    NotFound(String),
73    /// Session is not Idle (cannot destroy/overwrite a running session).
74    NotIdle(String),
75}
76
77impl std::fmt::Display for SessionError {
78    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79        match self {
80            SessionError::MaxSessionsReached(n) => {
81                write!(f, "Maximum {} sessions per user reached", n)
82            }
83            SessionError::NotFound(id) => write!(f, "Session {} not found", id),
84            SessionError::NotIdle(state) => write!(f, "Session is not idle (state: {})", state),
85        }
86    }
87}
88
89impl std::error::Error for SessionError {}
90
91/// Top-level registry: user_key → user's session collection.
92pub type SessionRegistry = Arc<DashMap<String, UserSessions>>;
93
94// ---------------------------------------------------------------------------
95// ServerState
96// ---------------------------------------------------------------------------
97
98/// Cached per-user MCP tool loader (16C).
99///
100/// One entry per user hash. `loader: None` means the user's effective
101/// config has MCP disabled — agents run MCP-less via abk's
102/// `McpSource::Prebuilt(None)` (semantically identical to
103/// `[mcp] enabled = false` today, but with zero config re-parsing).
104pub struct McpLoaderEntry {
105    /// Built loader; `None` = MCP disabled for this user.
106    pub loader: Option<std::sync::Arc<abk::agent::McpToolLoader>>,
107    /// Content fingerprint of the effective `[mcp]` config (SHA-256, first
108    /// 8 bytes as u64). Content, never mtime — overlays are rewritten in place.
109    pub fingerprint: u64,
110    pub built_at: chrono::DateTime<chrono::Utc>,
111    /// Set when the last build failed; the message surfaces to the user's
112    /// next dispatch (fail loud). Other users are never affected.
113    pub degraded: Option<String>,
114    /// When `degraded` was set. Rebuilds are held back for
115    /// [`MCP_BUILD_RETRY_BACKOFF`] — a poison entry never sticks forever.
116    pub failed_at: Option<chrono::DateTime<chrono::Utc>>,
117}
118
119/// Minimum delay before retrying a failed MCP loader build (16C).
120pub const MCP_BUILD_RETRY_BACKOFF: std::time::Duration = std::time::Duration::from_secs(30);
121
122/// 16F: per-agent dispatch target — THQ agent_name → the agent-user that
123/// THQ-dispatched sessions must run AS.
124///
125/// Populated once at boot by the 16E discovery scan
126/// (`thq_dispatch::populate`); restart-only lifecycle, exactly like the
127/// enrollment lane itself.
128#[derive(Debug, Clone)]
129pub struct ThqDispatchEntry {
130    /// The agent-user's stable key = its Kanidm `sub` (16E sub-pin). This is
131    /// the identity impersonated on the inner dispatch: session bucket,
132    /// per-user home, MCP loader cache, and Cedar principal all resolve
133    /// through it.
134    pub user_key: String,
135    /// The agent's Kanidm service token from its per-user `.env` — exchanged
136    /// for a short-lived `role=agent` Bearer at dispatch time. `None` = the
137    /// agent has no provisioned credential (dispatch fails loud, 502).
138    pub service_token: Option<String>,
139    /// The issuer the token's own credential declares (`[mcp.credentials.*]
140    /// .issuer_url` from the agent's overlay). Kanidm accepts a token
141    /// exchange only on the origin the token was minted for — the issuer
142    /// MUST travel with the credential that owns the token (verified live
143    /// 2026-08-30: same token, 200 on its own vhost, 400 elsewhere).
144    pub issuer_url: Option<String>,
145}
146
147/// 16F: the service-account issuer declared by the shared config's first
148/// `[mcp.credentials.*]` entry with `type = "service-account"` (see
149/// `ServerState::service_issuer` for why this exists).
150pub fn service_issuers_from_config(config_toml: &str) -> Vec<String> {
151    let Ok(v) = config_toml.parse::<toml::Value>() else {
152        return Vec::new();
153    };
154    let Some(creds) = v.get("mcp").and_then(|m| m.get("credentials")).and_then(|c| c.as_table())
155    else {
156        return Vec::new();
157    };
158    let mut out = Vec::new();
159    for (_name, cred) in creds {
160        if cred.get("type").and_then(|t| t.as_str()) == Some("service-account") {
161            if let Some(issuer) = cred.get("issuer_url").and_then(|i| i.as_str()) {
162                if !out.iter().any(|e: &String| e == issuer) {
163                    out.push(issuer.to_string());
164                }
165            }
166        }
167    }
168    out
169}
170
171/// Shared state accessible by all axum handlers.
172#[derive(Clone)]
173pub struct ServerState {
174    /// Per-user multi-session registry (MSU).
175    pub sessions: SessionRegistry,
176    /// Broadcast sender for backward compat — delegates to the default user's channel.
177    pub ws_tx: broadcast::Sender<String>,
178    /// Auth state (None = auth disabled, all endpoints open).
179    pub auth: Option<Arc<AuthState>>,
180    /// Shared config TOML (all users share the same agent config).
181    pub config_toml: Option<String>,
182    /// Shared secrets (injected into every per-user session).
183    pub secrets: Option<std::collections::HashMap<String, String>>,
184    /// Shared build info (injected into every per-user session).
185    pub build_info: Option<trustee_core::types::BuildInfo>,
186    /// Global concurrency limiter — limits the number of simultaneous workflows
187    /// across all users. Default: 8 concurrent workflows.
188    pub workflow_semaphore: Arc<tokio::sync::Semaphore>,
189    /// Maximum number of concurrent sessions per user. Default: 4.
190    pub max_sessions_per_user: usize,
191    /// Whether per-user config overlays may override the `[llm]` section.
192    /// Default: `false` — overlays are limited to `[mcp]`.
193    pub allow_llm_overlay: bool,
194    /// Per-user McpToolLoader cache (16C), keyed by user hash (never the raw key).
195    pub mcp_loaders: Arc<DashMap<String, McpLoaderEntry>>,
196    /// Single-flight build locks per user hash (16C).
197    mcp_build_locks: Arc<DashMap<String, Arc<tokio::sync::Mutex<()>>>>,
198    /// 16F: THQ per-agent dispatch table (agent_name → target), boot-populated.
199    pub thq_dispatch: Arc<DashMap<String, ThqDispatchEntry>>,
200    /// 16F: cached impersonation Bearers per agent user_key
201    /// `(token, expires_at)` — expiry-buffered, re-minted on 401.
202    pub agent_dispatch_tokens: Arc<DashMap<String, (String, std::time::Instant)>>,
203}
204
205impl ServerState {
206    /// Create new shared state from a default session, broadcast sender, and optional auth.
207    pub fn new(
208        session: Session,
209        ws_tx: broadcast::Sender<String>,
210        auth: Option<Arc<AuthState>>,
211    ) -> Self {
212        let sessions = Arc::new(DashMap::new());
213
214        // Store the default user's UserSessions with an initial session
215        let token_store = Arc::new(pep::MemoryTokenStore::new());
216        let (ws_tx_entry, _) = broadcast::channel::<String>(256);
217
218        let now = chrono::Utc::now();
219        let initial_entry = UserSessionEntry {
220            session: Arc::new(Mutex::new(session)),
221            ws_tx: ws_tx_entry,
222            created_at: now,
223            last_active: Arc::new(Mutex::new(now)),
224        };
225
226        let user_sessions = UserSessions {
227            sessions: DashMap::new(),
228            token_store,
229            active_session_id: Mutex::new(String::new()),
230        };
231        user_sessions.sessions.insert("default".to_string(), initial_entry);
232
233        sessions.insert("default".to_string(), user_sessions);
234
235        Self {
236            sessions,
237            ws_tx,
238            auth,
239            config_toml: None,
240            secrets: None,
241            build_info: None,
242            workflow_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
243            max_sessions_per_user: 4,
244            allow_llm_overlay: false,
245            mcp_loaders: Arc::new(DashMap::new()),
246            mcp_build_locks: Arc::new(DashMap::new()),
247            thq_dispatch: Arc::new(DashMap::new()),
248            agent_dispatch_tokens: Arc::new(DashMap::new()),
249        }
250    }
251
252    pub fn with_config_toml(mut self, config_toml: String) -> Self {
253        self.config_toml = Some(config_toml);
254        self
255    }
256
257    /// 16F: the issuer URL Kanidm accepts SERVICE-ACCOUNT token exchange on.
258    ///
259    /// Kanidm binds exchange to the OAuth2 client's configured origin: the
260    /// `[oidc]` issuer host serves logins/validation fine but REJECTS token
261    /// exchange with `invalid_request`, while the `idp` vhost — the issuer
262    /// every working `[mcp.credentials.*]` service-account entry already
263    /// uses — accepts it (verified live 2026-08-30: same token, 200 vs 400).
264    /// The exchange therefore takes its issuer from the shared config's
265    /// first `service-account` credential; callers fall back to the auth
266    /// issuer when absent.
267    pub fn service_issuers(&self) -> Vec<String> {
268        service_issuers_from_config(self.config_toml.as_deref().unwrap_or(""))
269    }
270
271    pub fn with_secrets(mut self, secrets: std::collections::HashMap<String, String>) -> Self {
272        self.secrets = Some(secrets);
273        self
274    }
275
276    pub fn with_build_info(mut self, build_info: trustee_core::types::BuildInfo) -> Self {
277        self.build_info = Some(build_info);
278        self
279    }
280
281    pub fn with_max_concurrent_workflows(mut self, max: usize) -> Self {
282        self.workflow_semaphore = Arc::new(tokio::sync::Semaphore::new(max));
283        self
284    }
285
286    /// Set the max sessions per user.
287    pub fn with_max_sessions_per_user(mut self, max: usize) -> Self {
288        self.max_sessions_per_user = max;
289        self
290    }
291
292    /// Allow per-user config overlays to override the `[llm]` section.
293    /// Default: `false` (overlays limited to `[mcp]`).
294    pub fn with_allow_llm_overlay(mut self, allow: bool) -> Self {
295        self.allow_llm_overlay = allow;
296        self
297    }
298
299    // -----------------------------------------------------------------------
300    // MSU: Multi-session methods
301    // -----------------------------------------------------------------------
302
303    /// Create a new session for a user. Returns the session_id.
304    ///
305    /// Creates a fresh `Session::new()`, copies shared config, sets per-user
306    /// isolation, creates a broadcast channel, spawns a drain task, and inserts
307    /// into the user's session DashMap. The new session becomes the "active" one.
308    pub async fn create_session(
309        &self,
310        user_key: &str,
311        session_name: Option<String>,
312        identity: Option<String>,
313        activate: bool,
314    ) -> Result<String, SessionError> {
315        // Get or create the user's UserSessions entry
316        let user_sessions = self
317            .sessions
318            .entry(user_key.to_string())
319            .or_insert_with(|| UserSessions {
320                sessions: DashMap::new(),
321                token_store: Arc::new(pep::MemoryTokenStore::new()),
322                active_session_id: Mutex::new(String::new()),
323            });
324
325        // Check session limit
326        if user_sessions.sessions.len() >= self.max_sessions_per_user {
327            return Err(SessionError::MaxSessionsReached(self.max_sessions_per_user));
328        }
329
330        // Create new Session
331        let (mut session, workflow_rx) = Session::new();
332
333        // Copy shared config
334        if let Some(ref config_toml) = self.config_toml {
335            session.config_toml = Some(config_toml.clone());
336            session.parse_auto_handoff_config();
337            if let Ok(table) = config_toml.parse::<toml::Value>() {
338                if let Some(name) = table
339                    .get("agent")
340                    .and_then(|a| a.get("name"))
341                    .and_then(|n| n.as_str())
342                {
343                    session.agent_name = name.to_string();
344                }
345            }
346        }
347
348        session.secrets = self.secrets.clone();
349        session.build_info = self.build_info.clone();
350
351        // Per-user isolation
352        self.apply_user_isolation(&mut session, user_key);
353
354        // Apply session_name if provided
355        session.session_name = session_name;
356
357        // Apply agent identity if provided
358        session.identity = identity;
359
360        // Create broadcast channel
361        let (ws_tx_entry, _) = broadcast::channel::<String>(256);
362
363        // Generate session_id
364        let session_id = format!(
365            "session_{}_{}",
366            chrono::Utc::now().format("%Y_%m_%d_%H_%M"),
367            &uuid::Uuid::new_v4().to_string()[..8]
368        );
369
370        let now = chrono::Utc::now();
371
372        // Insert into user's sessions DashMap
373        user_sessions.sessions.insert(
374            session_id.clone(),
375            UserSessionEntry {
376                session: Arc::new(Mutex::new(session)),
377                ws_tx: ws_tx_entry.clone(),
378                created_at: now,
379                last_active: Arc::new(Mutex::new(now)),
380            },
381        );
382
383        // Set as active session only when the caller opts in.
384        //
385        // `resume_session` passes `activate: false` so resuming an arbitrary
386        // checkpoint never hijacks the caller's current active live session
387        // pointer. `new_session` / `create_session` (fresh-start paths) pass
388        // `activate: true` to keep the existing "new session becomes active"
389        // behavior. This keeps the per-user active pointer under UI control
390        // rather than letting any authenticated client overwrite it.
391        if activate {
392            *user_sessions.active_session_id.lock().await = session_id.clone();
393        }
394
395        // Spawn drain task
396        let session_arc = user_sessions
397            .sessions
398            .get(&session_id)
399            .map(|e| e.session.clone());
400        if let Some(session_arc) = session_arc {
401            self.spawn_user_drain_task(
402                session_id.clone(),
403                session_arc,
404                ws_tx_entry,
405                workflow_rx,
406            );
407        }
408
409        Ok(session_id)
410    }
411
412    /// Get a specific session by user_key + session_id.
413    /// Updates last_active on the session entry.
414    pub async fn get_session(
415        &self,
416        user_key: &str,
417        session_id: &str,
418    ) -> Option<(Arc<Mutex<Session>>, broadcast::Sender<String>)> {
419        let user_sessions = self.sessions.get(user_key)?;
420        let entry = user_sessions.sessions.get(session_id)?;
421
422        // Update last_active
423        let now = chrono::Utc::now();
424        *entry.last_active.lock().await = now;
425
426        Some((entry.session.clone(), entry.ws_tx.clone()))
427    }
428
429    /// Get a session by EITHER its live MSU registry key OR its
430    /// checkpoint/session identity (`session.session_id`).
431    ///
432    /// The web frontend tracks `currentSessionId` from the `ResumeInfo` WS
433    /// message, which carries the auto-derived checkpoint id
434    /// (`session_YYYY_MM_DD_HH_MM_uuid8`) — NOT the live MSU registry key
435    /// (`"default"` or the key from `create_session()`). External clients
436    /// like Torpi/THQ pass the live registry key. This resolver accepts both:
437    ///
438    /// 1. Try registry-key lookup first (precise, used by Torpi/THQ).
439    /// 2. Fall back to scanning the user's live sessions for one whose
440    ///    `session.session_id` matches the requested id (used by the
441    ///    embedded web UI after a command or resume).
442    ///
443    /// Returns `(live_registry_key, session_arc, ws_tx)`, or `None` if not
444    /// found. The live key is returned so callers that need to set it as
445    /// active (or otherwise reference the registry) use the real key.
446    pub async fn get_session_by_any_id(
447        &self,
448        user_key: &str,
449        id: &str,
450    ) -> Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> {
451        // Fast path: registry key match.
452        let user_sessions = self.sessions.get(user_key)?;
453        if let Some(entry) = user_sessions.sessions.get(id) {
454            // Update last_active
455            let now = chrono::Utc::now();
456            *entry.last_active.lock().await = now;
457            return Some((id.to_string(), entry.session.clone(), entry.ws_tx.clone()));
458        }
459
460        // Slow path: scan live sessions for a matching session.session_id.
461        for entry in user_sessions.sessions.iter() {
462            let session = entry.session.lock().await;
463            if session.session_id.as_deref() == Some(id) {
464                let key = entry.key().clone();
465                let ws_tx = entry.ws_tx.clone();
466                drop(session);
467                // Update last_active
468                let now = chrono::Utc::now();
469                *entry.last_active.lock().await = now;
470                return Some((key, entry.session.clone(), ws_tx));
471            }
472        }
473
474        None
475    }
476
477    /// Tombstone directory for destroyed-but-on-disk session chains.
478    ///
479    /// Destroy removes the live session but deliberately keeps checkpoints
480    /// ("checkpoint data is preserved"). Restoring must NOT resurrect a
481    /// chain the user explicitly terminated, so destroy writes a marker
482    /// here and [`Self::restore_session_from_checkpoint`] refuses to
483    /// hydrate tombstoned chains. Markers are tiny and intentionally
484    /// permanent.
485    fn tombstone_path(&self, user_key: &str, chain_id: &str) -> Option<std::path::PathBuf> {
486        if chain_id.is_empty() || !chain_id.starts_with("session_") {
487            return None;
488        }
489        let hash = trustee_core::user_hash(user_key);
490        dirs::home_dir().map(|home| {
491            home.join(".trustee")
492                .join("users")
493                .join(&hash)
494                .join("session_tombstones")
495                .join(chain_id)
496        })
497    }
498
499    /// Write a destroy tombstone for a checkpoint chain (best effort).
500    ///
501    /// Called from [`Self::destroy_session`] when the destroyed live
502    /// session was attached to an on-disk chain.
503    fn write_tombstone(&self, user_key: &str, chain_id: &str) {
504        if let Some(path) = self.tombstone_path(user_key, chain_id) {
505            if let Some(parent) = path.parent() {
506                let _ = std::fs::create_dir_all(parent);
507            }
508            let _ = std::fs::write(&path, chrono::Utc::now().to_rfc3339());
509        }
510    }
511
512    /// Look up a live session by any id; on a miss, transparently restore
513    /// the checkpoint chain from disk.
514    ///
515    /// Restart resilience: trustee-api holds live sessions in memory only,
516    /// so a process restart orphans every open console. When a client asks
517    /// for a chain id that is not in memory but EXISTS on disk
518    /// (checkpoints present, not tombstoned by destroy), a live session is
519    /// rebuilt with `resume_info` pointing at the latest checkpoint and
520    /// registered under the chain id. The next command continues the
521    /// conversation; the console's alive-check and WS reconnect succeed.
522    ///
523    /// Returns `(live_registry_key, session_arc, ws_tx)` — the key is the
524    /// chain id for restored sessions.
525    pub async fn get_or_restore_session_by_any_id(
526        &self,
527        user_key: &str,
528        id: &str,
529    ) -> Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> {
530        if let Some(found) = self.get_session_by_any_id(user_key, id).await {
531            return Some(found);
532        }
533        self.restore_session_from_checkpoint(user_key, id).await
534    }
535
536    /// Rebuild a live session from an on-disk checkpoint chain.
537    ///
538    /// Persona note: dispatched (THQ/xagent) sessions carry `identity:
539    /// None` by design — their persona is the per-user overlay config
540    /// ([lifecycle].system_template), which `apply_user_isolation` +
541    /// `get_user_config_and_home` re-apply here exactly as at creation.
542    /// Explicit per-request identity overrides (rare) are not persisted
543    /// and therefore not restored.
544    async fn restore_session_from_checkpoint(
545        &self,
546        user_key: &str,
547        chain_id: &str,
548    ) -> Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> {
549        // Never resurrect an explicitly destroyed chain.
550        if let Some(path) = self.tombstone_path(user_key, chain_id) {
551            if path.exists() {
552                return None;
553            }
554        }
555
556        // Restore needs the caller's resolved config (persona) and home.
557        let (config_toml, home_dir) = self.get_user_config_and_home(user_key);
558        let config_toml = config_toml?;
559
560        // The chain must exist on disk with at least one checkpoint;
561        // otherwise this id is genuinely unknown → 404 semantics.
562        let info = trustee_core::sessions::create_resume_info(
563            &config_toml,
564            chain_id,
565            home_dir.as_deref(),
566        )
567        .await
568        .ok()??;
569
570        // Single-flight per chain id: racing lookups dedupe on the DashMap
571        // entry — the first insert wins, the loser adopts the winner.
572        let user_sessions = self
573            .sessions
574            .entry(user_key.to_string())
575            .or_insert_with(|| UserSessions {
576                sessions: DashMap::new(),
577                token_store: Arc::new(pep::MemoryTokenStore::new()),
578                active_session_id: Mutex::new(String::new()),
579            });
580
581        if let Some(existing) = user_sessions.sessions.get(chain_id) {
582            let now = chrono::Utc::now();
583            *existing.last_active.lock().await = now;
584            return Some((
585                chain_id.to_string(),
586                existing.session.clone(),
587                existing.ws_tx.clone(),
588            ));
589        }
590
591        // Respect the per-user cap: evict the least-recently-active Idle
592        // live session to make room (restores are user-initiated
593        // reattaches, so refusing them would recreate the lost-session
594        // bug for heavy users). If nothing is evictable, refuse.
595        if user_sessions.sessions.len() >= self.max_sessions_per_user {
596            let mut evict: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
597            for entry in user_sessions.sessions.iter() {
598                let is_idle = entry.session.lock().await.workflow_state
599                    == trustee_core::types::WorkflowState::Idle;
600                if !is_idle {
601                    continue;
602                }
603                let la = entry.last_active.lock().await;
604                if evict.as_ref().is_none_or(|(_, t)| *la < *t) {
605                    evict = Some((entry.key().clone(), *la));
606                }
607            }
608            if let Some((victim, _)) = evict {
609                user_sessions.sessions.remove(&victim);
610            } else {
611                return None;
612            }
613        }
614
615        // Build the session exactly like create_session, then attach the
616        // on-disk chain identity and resume pointer.
617        let (mut session, workflow_rx) = Session::new();
618        session.config_toml = Some(config_toml.clone());
619        session.parse_auto_handoff_config();
620        if let Ok(table) = config_toml.parse::<toml::Value>() {
621            if let Some(name) = table
622                .get("agent")
623                .and_then(|a| a.get("name"))
624                .and_then(|n| n.as_str())
625            {
626                session.agent_name = name.to_string();
627            }
628        }
629        session.secrets = self.secrets.clone();
630        session.build_info = self.build_info.clone();
631        self.apply_user_isolation(&mut session, user_key);
632        session.session_id = Some(chain_id.to_string());
633        session.resume_info = Some(info);
634        session.session_name = None; // display name re-derives from metadata
635
636        let (ws_tx_entry, _) = broadcast::channel::<String>(256);
637        let now = chrono::Utc::now();
638
639        user_sessions.sessions.insert(
640            chain_id.to_string(),
641            UserSessionEntry {
642                session: Arc::new(Mutex::new(session)),
643                ws_tx: ws_tx_entry.clone(),
644                created_at: now,
645                last_active: Arc::new(Mutex::new(now)),
646            },
647        );
648
649        // Drain task: forward workflow output into the WS broadcast channel
650        // (same contract as create_session — without this the restored
651        // session would run silently).
652        let session_arc = user_sessions
653            .sessions
654            .get(chain_id)
655            .map(|e| e.session.clone())
656            .ok_or_else(|| SessionError::NotFound(chain_id.to_string()))
657            .ok()?;
658        self.spawn_user_drain_task(
659            chain_id.to_string(),
660            session_arc.clone(),
661            ws_tx_entry.clone(),
662            workflow_rx,
663        );
664
665        Some((chain_id.to_string(), session_arc, ws_tx_entry))
666    }
667
668    /// List all active sessions for a user, sorted by last_active desc.
669    pub async fn list_sessions(&self, user_key: &str) -> Vec<SessionListItem> {        let Some(user_sessions) = self.sessions.get(user_key) else {
670            return Vec::new();
671        };
672
673        let mut items = Vec::new();
674        for entry in user_sessions.sessions.iter() {
675            let session = entry.session.lock().await;
676            let workflow_state = match session.workflow_state {
677                trustee_core::types::WorkflowState::Idle => "Idle",
678                trustee_core::types::WorkflowState::Running => "Running",
679                trustee_core::types::WorkflowState::Cancelling => "Cancelling",
680            };
681            let last_active = entry.last_active.lock().await;
682            items.push(SessionListItem {
683                session_id: entry.key().clone(),
684                session_name: session.session_name.clone(),
685                workflow_state: workflow_state.to_string(),
686                created_at: entry.created_at.to_rfc3339(),
687                last_active: last_active.to_rfc3339(),
688                handoff_count: session.handoff_count,
689            });
690        }
691        drop(user_sessions);
692
693        // Sort by last_active descending
694        items.sort_by(|a, b| b.last_active.cmp(&a.last_active));
695        items
696    }
697
698    /// Destroy a session. The session must be Idle.
699    pub async fn destroy_session(
700        &self,
701        user_key: &str,
702        session_id: &str,
703    ) -> Result<(), SessionError> {
704        let user_sessions = self
705            .sessions
706            .get(user_key)
707            .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
708
709        // Check workflow state before removing; capture the on-disk chain
710        // id (if any) so destroy can tombstone it against later restores.
711        let chain_id_for_tombstone = {
712            let entry = user_sessions
713                .sessions
714                .get(session_id)
715                .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
716            let session = entry.session.lock().await;
717            if session.workflow_state != trustee_core::types::WorkflowState::Idle {
718                let state_str = match session.workflow_state {
719                    trustee_core::types::WorkflowState::Running => "Running",
720                    trustee_core::types::WorkflowState::Cancelling => "Cancelling",
721                    _ => "Unknown",
722                };
723                return Err(SessionError::NotIdle(state_str.to_string()));
724            }
725            session.session_id.clone()
726        };
727
728        // Remove from DashMap
729        user_sessions.sessions.remove(session_id);
730
731        // If the destroyed live session was attached to an on-disk chain,
732        // tombstone it so a later get_or_restore lookup does not resurrect
733        // a session the user explicitly terminated.
734        if let Some(chain_id) = chain_id_for_tombstone {
735            self.write_tombstone(user_key, &chain_id);
736        }
737
738        // If this was the active session, pick a new active
739        let mut active_id = user_sessions.active_session_id.lock().await;
740        if &*active_id == session_id {
741            // Pick the most recently active remaining session
742            let mut newest: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
743            for entry in user_sessions.sessions.iter() {
744                let la = entry.last_active.lock().await;
745                if newest.as_ref().map_or(true, |(_, t)| *la > *t) {
746                    newest = Some((entry.key().clone(), *la));
747                }
748            }
749            *active_id = newest.map(|(id, _)| id).unwrap_or_default();
750        }
751
752        Ok(())
753    }
754
755    /// Get or create the user's "active" session for legacy routes.
756    ///
757    /// Behavior:
758    /// 1. If user has no sessions → create one
759    /// 2. If active session exists → return it
760    /// 3. If active session was destroyed → create a new one
761    ///
762    /// Returns: (session_id, session_arc, ws_tx, token_store)
763    pub async fn ensure_active_session(
764        &self,
765        user_key: &str,
766    ) -> (
767        String,
768        Arc<Mutex<Session>>,
769        broadcast::Sender<String>,
770        Arc<pep::MemoryTokenStore>,
771    ) {
772        // Get or create user's UserSessions
773        let token_store = {
774            let user_sessions = self
775                .sessions
776                .entry(user_key.to_string())
777                .or_insert_with(|| UserSessions {
778                    sessions: DashMap::new(),
779                    token_store: Arc::new(pep::MemoryTokenStore::new()),
780                    active_session_id: Mutex::new(String::new()),
781                });
782            user_sessions.token_store.clone()
783        };
784
785        // Check if active session exists
786        let active_id = {
787            let user_sessions = self.sessions.get(user_key).unwrap();
788            let guard = user_sessions.active_session_id.lock().await;
789            guard.clone()
790        };
791
792        if !active_id.is_empty() {
793            if let Some((session, ws_tx)) = self.get_session(user_key, &active_id).await {
794                return (active_id, session, ws_tx, token_store);
795            }
796            // Active session was destroyed, fall through to create
797        }
798
799        // Need to create a new session
800        // For the "default" user, we may already have a "default" session entry
801        // from ServerState::new() — check for it
802        let existing_session: Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> = {
803            let user_sessions = self.sessions.get(user_key).unwrap();
804            let result = user_sessions.sessions.iter().next().map(|first| {
805                (
806                    first.key().clone(),
807                    first.session.clone(),
808                    first.ws_tx.clone(),
809                )
810            });
811            result
812        };
813        if let Some((id, session, ws_tx)) = existing_session {
814            let now = chrono::Utc::now();
815            if let Some(entry) = self.sessions.get(user_key) {
816                if let Some(e) = entry.sessions.get(&id) {
817                    *e.last_active.lock().await = now;
818                }
819                *entry.active_session_id.lock().await = id.clone();
820            }
821
822            return (id, session, ws_tx, token_store);
823        }
824
825        // Create a brand new session
826        let session_id = self
827            .create_session(user_key, None, None, true)
828            .await
829            .unwrap_or_else(|_| "default".to_string());
830
831        let (session, ws_tx) = self
832            .get_session(user_key, &session_id)
833            .await
834            .expect("just-created session must exist");
835
836        (session_id, session, ws_tx, token_store)
837    }
838
839    /// DEPRECATED: Use ensure_active_session() instead.
840    /// Kept for backward compatibility — same 3-tuple return type.
841    pub async fn ensure_user_session(
842        &self,
843        user_key: &str,
844    ) -> (Arc<Mutex<Session>>, broadcast::Sender<String>, Arc<pep::MemoryTokenStore>) {
845        let (_id, session, ws_tx, token_store) = self.ensure_active_session(user_key).await;
846        (session, ws_tx, token_store)
847    }
848
849    /// Set a session as the user's active session.
850    pub async fn set_active_session(&self, user_key: &str, session_id: &str) {
851        if let Some(user_sessions) = self.sessions.get(user_key) {
852            if user_sessions.sessions.contains_key(session_id) {
853                *user_sessions.active_session_id.lock().await = session_id.to_string();
854            }
855        }
856    }
857
858    // -----------------------------------------------------------------------
859    // Read-only helpers (no session creation side effects)
860    // -----------------------------------------------------------------------
861
862    /// Resolve a user's home_dir without creating an in-memory session.
863    ///
864    /// This is the read-only equivalent of the isolation logic in
865    /// `apply_user_isolation`. Used by endpoints that only need to read
866    /// checkpoint data from disk (history, session list, session detail)
867    /// and must NOT create ghost sessions as a side effect.
868    pub fn get_user_home_dir(&self, user_key: &str) -> Option<std::path::PathBuf> {
869        let hash = trustee_core::user_hash(user_key);
870        dirs::home_dir().map(|home| home.join(".trustee").join("users").join(&hash))
871    }
872
873    /// Resolve config_toml and home_dir without creating an in-memory session.
874    ///
875    /// Returns `(config_toml, home_dir)`. If config is not loaded,
876    /// config_toml will be None.
877    pub fn get_user_config_and_home(&self, user_key: &str) -> (Option<String>, Option<std::path::PathBuf>) {
878        (self.config_toml.clone(), self.get_user_home_dir(user_key))
879    }
880
881    // -----------------------------------------------------------------------
882    // Private helpers
883    // -----------------------------------------------------------------------
884
885    /// Apply per-user isolation: SHA-256 hash → home_dir + project_id.
886    ///
887    /// Hashing goes through the single consolidated [`trustee_core::user_hash`]
888    /// so the web path and the CLI path can never drift apart.
889    fn apply_user_isolation(&self, session: &mut Session, user_key: &str) {
890        let user_hash = trustee_core::user_hash(user_key);
891
892        // Set per-user home directory for checkpoint isolation
893        let user_home = if let Some(home) = dirs::home_dir() {
894            let user_home = home.join(".trustee").join("users").join(&user_hash);
895            session.home_dir = Some(user_home.clone());
896            Some(user_home)
897        } else {
898            None
899        };
900
901        session.project_id = Some(format!("web{}", &user_hash[..16]));
902
903        // ── Per-user .env (Task 2) ──────────────────────────────────────
904        //
905        // Load per-user secrets from ~/.trustee/users/{hash}/.env
906        // These are merged on top of shared secrets (per-user wins).
907        // They are NEVER set as process env vars — used only for ${VAR}
908        // substitution in the config TOML below.
909        let shared_secrets = session.secrets.clone().unwrap_or_default();
910        let mut merged_secrets = shared_secrets.clone();
911
912        if let Some(ref user_home) = user_home {
913            if let Ok(merged) = self.load_user_secrets(user_home, &merged_secrets) {
914                merged_secrets = merged;
915            }
916        }
917
918        // ── Per-user config overlay (Task 3) ────────────────────────────
919        if let Some(ref user_home) = user_home {
920            if let Some(merged) = self.merge_user_config(user_home, session.config_toml.as_deref()) {
921                session.config_toml = Some(merged);
922                tracing::debug!("Merged per-user config into session");
923            }
924        }
925
926        // ── ${VAR} substitution (Task 4) ────────────────────────────────
927        if let Some(ref mut config_toml) = session.config_toml {
928            substitute_env_vars(config_toml, &merged_secrets);
929        }
930
931        // ── Strip per-user secrets (Task 5) ────────────────────────────
932        session.secrets = Some(shared_secrets);
933    }
934
935    /// Load secrets from a user's ~/.trustee/users/{hash}/.env and merge
936    /// on top of `base` (user wins). Returns the merged map.
937    fn load_user_secrets(
938        &self,
939        user_home: &std::path::Path,
940        base: &std::collections::HashMap<String, String>,
941    ) -> std::io::Result<std::collections::HashMap<String, String>> {
942        let user_env_path = user_home.join(".env");
943        if !user_env_path.exists() {
944            return Ok(base.clone());
945        }
946        let content = std::fs::read_to_string(&user_env_path)?;
947        let mut merged = base.clone();
948        for line in content.lines() {
949            let line = line.trim();
950            if line.is_empty() || line.starts_with('#') {
951                continue;
952            }
953            if let Some((key, value)) = line.split_once('=') {
954                let key = key.trim().to_string();
955                let value = value
956                    .trim()
957                    .trim_matches('"')
958                    .trim_matches('\'')
959                    .to_string();
960                merged.insert(key, value);
961            }
962        }
963        tracing::debug!("Loaded per-user secrets from {}", user_env_path.display());
964        Ok(merged)
965    }
966
967    /// Deep-merge a user's per-user config overlay on top of the shared
968    /// config. Returns the merged TOML string, or None if no overlay exists.
969    ///
970    /// Overlay allowlist (task 16B): only allowlisted top-level sections of
971    /// the user overlay are merged; anything else is dropped loudly.
972    /// - `[mcp]` is always allowed — per-user MCP tool sets are the point of
973    ///   the overlay convention.
974    /// - `[llm]` is allowed only when the instance opted in via the
975    ///   `[users].allow_llm_overlay` knob (default **false**).
976    ///
977    /// Everything else (`[server]`, `[auth]`, `[storage]`, `[web]`, …) is
978    /// boot-time, instance-level config and must not be rewritable through a
979    /// per-user file. This is predictability hardening, not a security
980    /// boundary: those sections were never re-read from session config.
981    fn merge_user_config(
982        &self,
983        user_home: &std::path::Path,
984        shared_config: Option<&str>,
985    ) -> Option<String> {
986        let user_config_path = user_home.join("config").join("trustee.toml");
987        if !user_config_path.exists() {
988            return None;
989        }
990        let user_config_toml = std::fs::read_to_string(&user_config_path).ok()?;
991        let shared = shared_config
992            .unwrap_or("")
993            .parse::<toml::Value>()
994            .ok()?;
995        let overlay = user_config_toml.parse::<toml::Value>().ok()?;
996
997        let allowed = |section: &str| {
998            section == "mcp"
999                // 16E: [thq] is per-user identity config (read from the overlay
1000                // file at boot); keeping it in the merged session config is
1001                // inert but avoids a false "dropping" warn on every dispatch.
1002                || section == "thq"
1003                // xagent personas: an agent's [lifecycle].system_template IS
1004                // her identity — user-owned by design. The WASM extension
1005                // switch ([lifecycle].enabled) is stripped below: persona is
1006                // user-owned, extension control is not.
1007                || section == "lifecycle"
1008                || (self.allow_llm_overlay && section == "llm")
1009        };
1010        // Mask the user in logs: user_home's dir name IS the hash — never
1011        // log the raw key (keys may be emails). Log a short hash prefix.
1012        let dir_name = user_home
1013            .file_name()
1014            .and_then(|n| n.to_str())
1015            .unwrap_or("<unknown>");
1016        let masked_user = dir_name.get(..8).unwrap_or(dir_name);
1017
1018        let mut filtered_overlay = toml::map::Map::new();
1019        if let Some(table) = overlay.as_table() {
1020            for (section, value) in table {
1021                if section == "lifecycle" {
1022                    let mut v = value.clone();
1023                    if let Some(t) = v.as_table_mut() {
1024                        if t.remove("enabled").is_some() {
1025                            tracing::info!(
1026                                "user config overlay: [lifecycle].enabled is not user-settable — stripped for user {}",
1027                                masked_user
1028                            );
1029                        }
1030                    }
1031                    filtered_overlay.insert(section.clone(), v);
1032                } else if allowed(section) {
1033                    filtered_overlay.insert(section.clone(), value.clone());
1034                } else {
1035                    tracing::warn!(
1036                        "user config overlay: dropping non-allowlisted section [{}] for user {}",
1037                        section,
1038                        masked_user
1039                    );
1040                }
1041            }
1042        }
1043
1044        if filtered_overlay.is_empty() {
1045            // Nothing survived the allowlist — no-op, keep shared config as-is.
1046            return None;
1047        }
1048        let overlay = toml::Value::Table(filtered_overlay);
1049
1050        let mut shared = shared;
1051        deep_merge_toml(&mut shared, &overlay);
1052        let merged = toml::to_string(&shared).ok()?;
1053        tracing::debug!("Merged per-user config from {}", user_config_path.display());
1054        Some(merged)
1055    }
1056
1057    /// Resolve the fully-merged config TOML for a user WITHOUT creating a session.
1058    ///
1059    /// This is the read-only equivalent of the config resolution in
1060    /// `apply_user_isolation`: shared config + per-user overlay + ${VAR}
1061    /// substitution. Used by endpoints that need to inspect config (e.g.
1062    /// listing available LLM models) without spawning a ghost session.
1063    ///
1064    /// Returns None if no shared config is loaded.
1065    pub fn resolve_user_config(&self, user_key: &str) -> Option<String> {
1066        let config_toml = self.config_toml.clone()?;
1067
1068        // Resolve user home dir (same hash scheme as apply_user_isolation)
1069        let user_home = self.get_user_home_dir(user_key)?;
1070
1071        // Start from shared secrets; merge per-user .env on top
1072        let mut merged_secrets = self.secrets.clone().unwrap_or_default();
1073        if let Ok(merged) = self.load_user_secrets(&user_home, &merged_secrets) {
1074            merged_secrets = merged;
1075        }
1076
1077        // Merge per-user config overlay
1078        let mut resolved = config_toml;
1079        if let Some(merged) = self.merge_user_config(&user_home, Some(&resolved)) {
1080            resolved = merged;
1081        }
1082
1083        // Substitute ${VAR} from merged secrets
1084        substitute_env_vars(&mut resolved, &merged_secrets);
1085
1086        Some(resolved)
1087    }
1088
1089    /// Get or build the per-user MCP tool loader (16C).
1090    ///
1091    /// Cache semantics:
1092    /// - fingerprint = content hash of the effective `[mcp]` section
1093    ///   (shared + allowlist-filtered overlay + ${VAR} substitution, i.e.
1094    ///   exactly what `resolve_user_config` produces) — stale on ANY change.
1095    /// - hit + match → `Arc` clone, zero network I/O.
1096    /// - miss/stale → single-flight build (one build per user at a time;
1097    ///   late arrivals re-check and reuse the winner's entry).
1098    /// - `Ok(None)` = MCP disabled for this user → agent runs MCP-less via
1099    ///   abk `McpSource::Prebuilt(None)` (same semantics as
1100    ///   `[mcp] enabled = false`, no per-task re-evaluation).
1101    /// - build failure → cached degraded entry with
1102    ///   [`MCP_BUILD_RETRY_BACKOFF`]; the error is returned so THIS user's
1103    ///   dispatch fails loud while other users are unaffected.
1104    pub async fn get_or_build_mcp_loader(
1105        &self,
1106        user_key: &str,
1107        token_store: &Arc<pep::MemoryTokenStore>,
1108    ) -> Result<Option<std::sync::Arc<abk::agent::McpToolLoader>>, String> {
1109        let user_hash = trustee_core::user_hash(user_key);
1110
1111        // Effective per-user config — the fingerprint source of truth.
1112        let resolved = self.resolve_user_config(user_key);
1113        let fingerprint = fingerprint_mcp_section(resolved.as_deref());
1114
1115        // Fast path: fresh, non-degraded entry.
1116        if let Some(entry) = self.mcp_loaders.get(&user_hash) {
1117            if entry.degraded.is_none() {
1118                if entry.fingerprint == fingerprint {
1119                    return Ok(entry.loader.clone());
1120                }
1121            } else if let (Some(err), Some(failed_at)) = (&entry.degraded, entry.failed_at) {
1122                let backoff = chrono::Duration::from_std(MCP_BUILD_RETRY_BACKOFF)
1123                    .unwrap_or_else(|_| chrono::Duration::seconds(30));
1124                if chrono::Utc::now() < failed_at + backoff {
1125                    return Err(err.clone());
1126                }
1127            }
1128        }
1129
1130        // Single-flight: one build per user at a time.
1131        let lock = self
1132            .mcp_build_locks
1133            .entry(user_hash.clone())
1134            .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
1135            .clone();
1136        let _guard = lock.lock().await;
1137
1138        // Double-check: another task may have built while we waited.
1139        if let Some(entry) = self.mcp_loaders.get(&user_hash) {
1140            if entry.degraded.is_none() && entry.fingerprint == fingerprint {
1141                return Ok(entry.loader.clone());
1142            }
1143        }
1144
1145        match self
1146            .build_mcp_loader(&user_hash, resolved.as_deref(), fingerprint, token_store)
1147            .await
1148        {
1149            Ok(entry) => {
1150                self.mcp_loaders.insert(user_hash.clone(), entry);
1151                Ok(self.mcp_loaders.get(&user_hash).unwrap().loader.clone())
1152            }
1153            Err(err) => {
1154                tracing::warn!(
1155                    "MCP loader build FAILED for user {}; dispatch fails loud, retry after {:?}",
1156                    &user_hash[..8.min(user_hash.len())],
1157                    MCP_BUILD_RETRY_BACKOFF
1158                );
1159                self.mcp_loaders.insert(
1160                    user_hash,
1161                    McpLoaderEntry {
1162                        loader: None,
1163                        fingerprint,
1164                        built_at: chrono::Utc::now(),
1165                        degraded: Some(err.clone()),
1166                        failed_at: Some(chrono::Utc::now()),
1167                    },
1168                );
1169                Err(err)
1170            }
1171        }
1172    }
1173
1174    /// Build a loader entry from the effective config. No caching here —
1175    /// the caller owns insertion and degraded handling.
1176    async fn build_mcp_loader(
1177        &self,
1178        user_hash: &str,
1179        resolved: Option<&str>,
1180        fingerprint: u64,
1181        token_store: &Arc<pep::MemoryTokenStore>,
1182    ) -> Result<McpLoaderEntry, String> {
1183        let mcp_config: Option<abk::config::McpConfig> = match resolved {
1184            Some(toml_str) => {
1185                let value = toml_str
1186                    .parse::<toml::Value>()
1187                    .map_err(|e| format!("config parse failed: {}", e))?;
1188                match value.get("mcp") {
1189                    Some(section) => {
1190                        use serde::Deserialize as _;
1191                        Some(
1192                            abk::config::McpConfig::deserialize(section.clone())
1193                                .map_err(|e| format!("invalid [mcp] config: {}", e))?,
1194                        )
1195                    }
1196                    None => None,
1197                }
1198            }
1199            None => None,
1200        };
1201
1202        let loader = match mcp_config {
1203            Some(cfg) if cfg.enabled => {
1204                let built = abk::agent::McpToolLoader::with_token_store(
1205                    &cfg,
1206                    Some(token_store.clone() as std::sync::Arc<dyn pep::token_store::TokenStore>),
1207                )
1208                .await
1209                .map_err(|e| format!("MCP loader build failed: {}", e))?;
1210
1211                // THE parity-evidence log line (16C acceptance + migration task):
1212                // one INFO line per build; a task loop that reuses the cache
1213                // shows exactly one line per user per fingerprint.
1214                let servers: Vec<String> = built
1215                    .server_statuses
1216                    .iter()
1217                    .map(|s| {
1218                        if s.connected {
1219                            format!("{}(up,{}tools)", s.name, s.tool_count)
1220                        } else {
1221                            format!("{}(DOWN)", s.name)
1222                        }
1223                    })
1224                    .collect();
1225                tracing::info!(
1226                    "MCP loader built for user {}: servers=[{}] total_tools={}",
1227                    &user_hash[..8.min(user_hash.len())],
1228                    servers.join(", "),
1229                    built.tool_count
1230                );
1231                Some(std::sync::Arc::new(built))
1232            }
1233            _ => None,
1234        };
1235
1236        Ok(McpLoaderEntry {
1237            loader,
1238            fingerprint,
1239            built_at: chrono::Utc::now(),
1240            degraded: None,
1241            failed_at: None,
1242        })
1243    }
1244
1245    /// Spawn a background drain task for a specific session's workflow receiver.
1246    fn spawn_user_drain_task(
1247        &self,
1248        session_id: String,
1249        session: Arc<Mutex<Session>>,
1250        ws_tx: broadcast::Sender<String>,
1251        mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>,
1252    ) {
1253        // Bug 6: broadcast StateChanged only on actual transitions. Without
1254        // this, every WS message was accompanied by a duplicate StateChanged,
1255        // doubling traffic and re-triggering frontend updateState on hot
1256        // StreamDelta/ReasoningDelta bursts. New subscribers learn the current
1257        // state from the WS snapshot, so the initial None is safe.
1258        let mut last_broadcast_state: Option<String> = None;
1259        tokio::spawn(async move {
1260            while let Some(msg) = workflow_rx.recv().await {
1261                {
1262                    let mut session = session.lock().await;
1263                    session.handle_workflow_message(msg.clone());
1264
1265                    let state_str = match session.workflow_state {
1266                        trustee_core::types::WorkflowState::Idle => "Idle",
1267                        trustee_core::types::WorkflowState::Running => "Running",
1268                        trustee_core::types::WorkflowState::Cancelling => "Cancelling",
1269                    };
1270                    if last_broadcast_state.as_deref() != Some(state_str) {
1271                        last_broadcast_state = Some(state_str.to_string());
1272                        let state_msg = serde_json::json!({
1273                            "type": "StateChanged",
1274                            "state": state_str
1275                        });
1276                        let _ = ws_tx.send(state_msg.to_string());
1277                    }
1278                }
1279
1280                let json =
1281                    serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
1282                let _ = ws_tx.send(json);
1283            }
1284            tracing::debug!("Drain task ended for session: {}", session_id);
1285        });
1286    }
1287
1288    /// Spawn the default user's drain task (backward compatibility).
1289    /// Called during server startup for the initial session.
1290    pub fn spawn_drain_task(self, mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>) {
1291        // Get the default user's first session
1292        let default_user = self
1293            .sessions
1294            .get("default")
1295            .expect("default user must exist");
1296        let first_entry = default_user
1297            .sessions
1298            .iter()
1299            .next()
1300            .expect("default user must have at least one session");
1301        let session = first_entry.session.clone();
1302        let ws_tx = first_entry.ws_tx.clone();
1303        let session_id = first_entry.key().clone();
1304        drop(first_entry);
1305        drop(default_user);
1306
1307        // Bug 6: transition-only StateChanged broadcasts (see
1308        // spawn_user_drain_task). The task starts in Running state because
1309        // spawn_drain_task is only called after a command began executing.
1310        let mut last_broadcast_state = Some("Running".to_string());
1311        tokio::spawn(async move {
1312            while let Some(msg) = workflow_rx.recv().await {
1313                {
1314                    let mut session = session.lock().await;
1315                    session.handle_workflow_message(msg.clone());
1316
1317                    let state_str = match session.workflow_state {
1318                        trustee_core::types::WorkflowState::Idle => "Idle",
1319                        trustee_core::types::WorkflowState::Running => "Running",
1320                        trustee_core::types::WorkflowState::Cancelling => "Cancelling",
1321                    };
1322                    if last_broadcast_state.as_deref() != Some(state_str) {
1323                        last_broadcast_state = Some(state_str.to_string());
1324                        let state_msg = serde_json::json!({
1325                            "type": "StateChanged",
1326                            "state": state_str
1327                        });
1328                        let _ = ws_tx.send(state_msg.to_string());
1329                    }
1330                }
1331
1332                let json =
1333                    serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
1334                let _ = ws_tx.send(json);
1335            }
1336            tracing::debug!("Drain task ended for session: {}", session_id);
1337        });
1338    }
1339
1340    /// Resolve the user key from request headers.
1341    pub async fn resolve_user_key(&self, headers: &axum::http::HeaderMap) -> String {
1342        let Some(ref auth) = self.auth else {
1343            return "default".to_string();
1344        };
1345
1346        // Try Bearer header first
1347        if let Some(token) = headers
1348            .get(axum::http::header::AUTHORIZATION)
1349            .and_then(|v| v.to_str().ok())
1350            .and_then(|v| v.strip_prefix("Bearer "))
1351            .map(|s| s.to_string())
1352        {
1353            if token.starts_with("dev:") {
1354                let parts: Vec<&str> = token.splitn(4, ':').collect();
1355                if parts.len() >= 4 {
1356                    return format!("dev:{}", parts[1]);
1357                }
1358            }
1359            if let Ok(claims) = auth.validate_token(&token).await {
1360                return claims.sub;
1361            }
1362        }
1363
1364        // Try cookie
1365        let cookie_session_id = headers
1366            .get(axum::http::header::COOKIE)
1367            .and_then(|v| v.to_str().ok())
1368            .and_then(|cookies| {
1369                cookies
1370                    .split(';')
1371                    .map(|c| c.trim())
1372                    .find_map(|c| {
1373                        c.strip_prefix(&format!("{}=", auth.config.cookie_name))
1374                            .map(|s| s.to_string())
1375                    })
1376            });
1377
1378        if let Some(session_id) = cookie_session_id {
1379            if session_id.starts_with("dev:") {
1380                let parts: Vec<&str> = session_id.splitn(4, ':').collect();
1381                if parts.len() >= 4 {
1382                    return format!("dev:{}", parts[1]);
1383                }
1384            }
1385
1386            if let Ok(access_token) = auth.session_manager.get_token(&session_id).await {
1387                if let Ok(claims) = auth.validate_token(&access_token).await {
1388                    return claims.sub;
1389                }
1390            }
1391        }
1392
1393        "default".to_string()
1394    }
1395}
1396
1397// ---------------------------------------------------------------------------
1398// SerializableMessage (unchanged)
1399// ---------------------------------------------------------------------------
1400
1401/// Wrapper to serialize `TuiMessage` as JSON with a `type` discriminator.
1402struct SerializableMessage<'a>(&'a TuiMessage);
1403
1404impl<'a> serde::Serialize for SerializableMessage<'a> {
1405    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1406    where
1407        S: serde::Serializer,
1408    {
1409        use serde::ser::SerializeStruct;
1410
1411        match self.0 {
1412            TuiMessage::OutputLine(line) => {
1413                let mut s = serializer.serialize_struct("msg", 2)?;
1414                s.serialize_field("type", "OutputLine")?;
1415                s.serialize_field("line", line)?;
1416                s.end()
1417            }
1418            TuiMessage::StreamDelta(delta) => {
1419                let mut s = serializer.serialize_struct("msg", 2)?;
1420                s.serialize_field("type", "StreamDelta")?;
1421                s.serialize_field("delta", delta)?;
1422                s.end()
1423            }
1424            TuiMessage::ReasoningDelta(delta) => {
1425                let mut s = serializer.serialize_struct("msg", 2)?;
1426                s.serialize_field("type", "ReasoningDelta")?;
1427                s.serialize_field("delta", delta)?;
1428                s.end()
1429            }
1430            TuiMessage::WorkflowCompleted => {
1431                let mut s = serializer.serialize_struct("msg", 2)?;
1432                s.serialize_field("type", "WorkflowCompleted")?;
1433                s.serialize_field("state", "Idle")?;
1434                s.end()
1435            }
1436            TuiMessage::WorkflowError(err) => {
1437                let mut s = serializer.serialize_struct("msg", 2)?;
1438                s.serialize_field("type", "WorkflowError")?;
1439                s.serialize_field("error", err)?;
1440                s.end()
1441            }
1442            TuiMessage::ResumeInfo(info) => match info {
1443                Some(ri) => {
1444                    let mut s = serializer.serialize_struct("msg", 5)?;
1445                    s.serialize_field("type", "ResumeInfo")?;
1446                    s.serialize_field("state", "Idle")?;
1447                    s.serialize_field("session_id", &ri.session_id)?;
1448                    s.serialize_field("checkpoint_id", &ri.checkpoint_id)?;
1449                    s.serialize_field("iteration", &ri.iteration)?;
1450                    s.end()
1451                }
1452                None => {
1453                    let mut s = serializer.serialize_struct("msg", 2)?;
1454                    s.serialize_field("type", "ResumeInfo")?;
1455                    s.serialize_field("state", "Idle")?;
1456                    s.end()
1457                }
1458            },
1459            TuiMessage::TodoUpdate(content) => {
1460                let mut s = serializer.serialize_struct("msg", 2)?;
1461                s.serialize_field("type", "TodoUpdate")?;
1462                s.serialize_field("content", content)?;
1463                s.end()
1464            }
1465            TuiMessage::WorkflowCancelled => {
1466                let mut s = serializer.serialize_struct("msg", 2)?;
1467                s.serialize_field("type", "WorkflowCancelled")?;
1468                s.serialize_field("state", "Idle")?;
1469                s.end()
1470            }
1471            TuiMessage::HandoffReady(briefing) => {
1472                let mut s = serializer.serialize_struct("msg", 3)?;
1473                s.serialize_field("type", "HandoffReady")?;
1474                s.serialize_field("state", "Idle")?;
1475                s.serialize_field("briefing", briefing)?;
1476                s.end()
1477            }
1478            TuiMessage::SessionRotated { old, new } => {
1479                let mut s = serializer.serialize_struct("msg", 3)?;
1480                s.serialize_field("type", "SessionRotated")?;
1481                s.serialize_field("old", old)?;
1482                s.serialize_field("new", new)?;
1483                s.end()
1484            }
1485            TuiMessage::HandoffFailed => {
1486                let mut s = serializer.serialize_struct("msg", 2)?;
1487                s.serialize_field("type", "HandoffFailed")?;
1488                s.serialize_field("state", "Idle")?;
1489                s.end()
1490            }
1491            TuiMessage::ToolPending {
1492                tool_name,
1493                hint,
1494            } => {
1495                let mut s = serializer.serialize_struct("msg", 3)?;
1496                s.serialize_field("type", "ToolPending")?;
1497                s.serialize_field("tool_name", tool_name)?;
1498                s.serialize_field("hint", hint)?;
1499                s.end()
1500            }
1501            TuiMessage::ToolDone {
1502                tool_name,
1503                success,
1504                hint,
1505            } => {
1506                let mut s = serializer.serialize_struct("msg", 4)?;
1507                s.serialize_field("type", "ToolDone")?;
1508                s.serialize_field("tool_name", tool_name)?;
1509                s.serialize_field("success", success)?;
1510                s.serialize_field("hint", hint)?;
1511                s.end()
1512            }
1513            TuiMessage::ContextTokensUpdated(count) => {
1514                let mut s = serializer.serialize_struct("msg", 2)?;
1515                s.serialize_field("type", "ContextTokensUpdated")?;
1516                s.serialize_field("count", count)?;
1517                s.end()
1518            }
1519            TuiMessage::McpServerStatus {
1520                name,
1521                connected,
1522                tool_count,
1523                error,
1524            } => {
1525                let mut s = serializer.serialize_struct("msg", 5)?;
1526                s.serialize_field("type", "McpServerStatus")?;
1527                s.serialize_field("name", name)?;
1528                s.serialize_field("connected", connected)?;
1529                s.serialize_field("tool_count", tool_count)?;
1530                s.serialize_field("error", error)?;
1531                s.end()
1532            }
1533            TuiMessage::SessionTitleUpdated(title) => {
1534                let mut s = serializer.serialize_struct("msg", 2)?;
1535                s.serialize_field("type", "SessionTitleUpdated")?;
1536                s.serialize_field("title", title)?;
1537                s.end()
1538            }
1539        }
1540    }
1541}
1542
1543// ---------------------------------------------------------------------------
1544// Per-user config helpers
1545// ---------------------------------------------------------------------------
1546
1547/// Deep-merge a TOML overlay on top of a base value (in-place).
1548///
1549/// - Tables: recursively merge key-by-key (overlay wins on conflict).
1550/// - Arrays: overlay replaces base entirely (no merging).
1551/// - Scalars: overlay replaces base.
1552/// - If a key exists in overlay but not base, it's added.
1553fn deep_merge_toml(base: &mut toml::Value, overlay: &toml::Value) {
1554    match (base, overlay) {
1555        (toml::Value::Table(base_table), toml::Value::Table(overlay_table)) => {
1556            for (key, overlay_val) in overlay_table {
1557                match base_table.get_mut(key) {
1558                    Some(base_val) => {
1559                        // Both exist — recurse if both are tables, else replace
1560                        deep_merge_toml(base_val, overlay_val);
1561                    }
1562                    None => {
1563                        // Key only in overlay — insert
1564                        base_table.insert(key.clone(), overlay_val.clone());
1565                    }
1566                }
1567            }
1568        }
1569        // Non-table: overlay replaces base
1570        (base, overlay) => {
1571            *base = overlay.clone();
1572        }
1573    }
1574}
1575
1576/// Replace `${VAR_NAME}` references in a string with values from a secrets map.
1577///
1578/// Falls back to process environment if the variable is not in the map.
1579/// Variables not found in either are left as-is.
1580fn substitute_env_vars(s: &mut String, secrets: &std::collections::HashMap<String, String>) {
1581    // Simple state machine: scan for ${, read until }, replace.
1582    let mut result = String::with_capacity(s.len());
1583    let bytes = s.as_bytes();
1584    let mut i = 0;
1585
1586    while i < bytes.len() {
1587        if i + 1 < bytes.len() && bytes[i] == b'$' && bytes[i + 1] == b'{' {
1588            // Find closing }
1589            if let Some(end) = s[i + 2..].find('}') {
1590                let var_name = &s[i + 2..i + 2 + end];
1591                // Look up in per-user secrets first, then process env
1592                if let Some(value) = secrets.get(var_name) {
1593                    result.push_str(value);
1594                } else if let Ok(value) = std::env::var(var_name) {
1595                    result.push_str(&value);
1596                } else {
1597                    // Not found — leave as-is
1598                    result.push_str(&s[i..i + 2 + end + 1]);
1599                }
1600                i = i + 2 + end + 1;
1601            } else {
1602                // No closing } — copy as-is
1603                result.push('$');
1604                i += 1;
1605            }
1606        } else {
1607            result.push(bytes[i] as char);
1608            i += 1;
1609        }
1610    }
1611
1612    *s = result;
1613}
1614
1615// ---------------------------------------------------------------------------
1616// Tests (task 16B: overlay allowlist + consolidated user hash)
1617// ---------------------------------------------------------------------------
1618
1619/// Fingerprint the effective `[mcp]` section of a resolved config (16C).
1620///
1621/// Content hash (SHA-256, first 8 bytes as u64) — never mtime: overlays are
1622/// rewritten in place. Configs without `[mcp]` fingerprint to a stable
1623/// constant, so the disabled state caches too.
1624fn fingerprint_mcp_section(resolved: Option<&str>) -> u64 {
1625    use sha2::{Digest, Sha256};
1626    let section = resolved
1627        .and_then(|s| s.parse::<toml::Value>().ok())
1628        .and_then(|v| v.get("mcp").cloned());
1629    let bytes = match section {
1630        Some(v) => v.to_string().into_bytes(),
1631        None => b"<no-mcp>".to_vec(),
1632    };
1633    let digest = Sha256::digest(&bytes);
1634    u64::from_be_bytes(digest[..8].try_into().expect("sha256 digest >= 8 bytes"))
1635}
1636
1637#[cfg(test)]
1638mod tests {
1639    use super::*;
1640
1641    /// Default-knob ServerState for overlay tests.
1642    fn test_state() -> ServerState {
1643        let (session, _rx) = Session::new();
1644        let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
1645        ServerState::new(session, ws_tx, None)
1646    }
1647
1648    /// Unique temp dir with a `config/` subdir (std-only; no tempfile dep).
1649    fn temp_user_home(tag: &str) -> std::path::PathBuf {
1650        static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1651        let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1652        let dir = std::env::temp_dir().join(format!(
1653            "trustee-state-test-{}-{}-{}",
1654            tag,
1655            std::process::id(),
1656            n
1657        ));
1658        std::fs::create_dir_all(dir.join("config")).expect("create temp user home");
1659        dir
1660    }
1661
1662    fn parse(toml_str: &str) -> toml::Value {
1663        toml_str.parse::<toml::Value>().expect("valid test TOML")
1664    }
1665
1666    /// (a) Overlay with [mcp] + [server] + [auth]: [mcp] applied (user wins),
1667    /// shared [server] preserved untouched, overlay [auth] dropped entirely.
1668    #[test]
1669    fn overlay_allowlist_drops_non_allowlisted_sections() {
1670        let state = test_state();
1671        let home = temp_user_home("allowlist");
1672        std::fs::write(
1673            home.join("config").join("trustee.toml"),
1674            "[server]\nport = 1\n\n[auth]\nmode = \"kanidm\"\n\n[mcp]\nmode = \"user\"\n",
1675        )
1676        .expect("write overlay");
1677
1678        let shared = "[server]\nport = 8080\n\n[mcp]\nmode = \"shared\"\n";
1679        let merged = state
1680            .merge_user_config(&home, Some(shared))
1681            .expect("overlay has allowlisted content");
1682
1683        let merged_val = parse(&merged);
1684        let expected = parse("[server]\nport = 8080\n\n[mcp]\nmode = \"user\"\n");
1685        assert_eq!(merged_val, expected, "merged config must be shared + [mcp] overlay only");
1686        assert!(merged_val.get("auth").is_none(), "overlay [auth] must be dropped");
1687    }
1688
1689    /// (b) No overlay file → merge is a no-op (None), shared config untouched.
1690    #[test]
1691    fn no_overlay_file_returns_none() {
1692        let state = test_state();
1693        let home = temp_user_home("empty");
1694        assert!(state
1695            .merge_user_config(&home, Some("[server]\nport = 8080\n"))
1696            .is_none());
1697    }
1698
1699    /// Overlay whose sections are ALL non-allowlisted → no-op (None).
1700    #[test]
1701    fn overlay_with_no_allowlisted_sections_is_noop() {
1702        let state = test_state();
1703        let home = temp_user_home("all-dropped");
1704        std::fs::write(
1705            home.join("config").join("trustee.toml"),
1706            "[server]\nport = 1\n\n[storage]\npath = \"/tmp/x\"\n",
1707        )
1708        .expect("write overlay");
1709        assert!(state
1710            .merge_user_config(&home, Some("[server]\nport = 8080\n"))
1711            .is_none());
1712    }
1713
1714    /// (c) [llm] dropped when allow_llm_overlay=false (default), applied when
1715    /// the instance opted in via with_allow_llm_overlay(true).
1716    #[test]
1717    fn llm_overlay_dropped_by_default_and_kept_when_enabled() {
1718        let shared = "[llm]\nprovider = \"openai\"\n\n[mcp]\nmode = \"shared\"\n";
1719        let overlay = "[llm]\nprovider = \"anthropic\"\n";
1720
1721        // Default: knob false → [llm]-only overlay is a no-op.
1722        let state = test_state();
1723        assert!(!state.allow_llm_overlay);
1724        let home = temp_user_home("llm-off");
1725        std::fs::write(home.join("config").join("trustee.toml"), overlay).expect("write overlay");
1726        assert!(state.merge_user_config(&home, Some(shared)).is_none());
1727
1728        // Opted in: [llm] kept, user value wins; shared [mcp] untouched.
1729        let state = test_state().with_allow_llm_overlay(true);
1730        assert!(state.allow_llm_overlay);
1731        let home = temp_user_home("llm-on");
1732        std::fs::write(home.join("config").join("trustee.toml"), overlay).expect("write overlay");
1733        let merged = state
1734            .merge_user_config(&home, Some(shared))
1735            .expect("[llm] overlay applies when opted in");
1736        let merged_val = parse(&merged);
1737        assert_eq!(merged_val["llm"]["provider"].as_str(), Some("anthropic"));
1738        assert_eq!(merged_val["mcp"]["mode"].as_str(), Some("shared"));
1739    }
1740
1741    /// (d) Web-path home dir resolves through the consolidated
1742    /// trustee_core::user_hash (determinism/known-vectors covered in
1743    /// trustee-core's own tests).
1744    #[test]
1745    fn user_home_dir_uses_consolidated_hash() {
1746        let state = test_state();
1747        if let Some(home) = state.get_user_home_dir("farzan@example.com") {
1748            assert_eq!(
1749                home.file_name().and_then(|n| n.to_str()),
1750                Some(trustee_core::user_hash("farzan@example.com")).as_deref()
1751            );
1752            let users_root = dirs::home_dir().unwrap().join(".trustee").join("users");
1753            assert_eq!(home.parent(), Some(&users_root).map(|p| p.as_path()));
1754        }
1755        // dirs::home_dir() unavailable in sandbox → covered by core tests.
1756    }
1757
1758    /// xagent personas: [lifecycle] survives the overlay allowlist (persona
1759    /// is user-owned) but the WASM extension switch does not.
1760    #[test]
1761    fn lifecycle_overlay_survives_allowlist_enabled_stripped() {
1762        let state = test_state();
1763        let home = temp_user_home("lifecycle-persona");
1764        std::fs::write(
1765            home.join("config").join("trustee.toml"),
1766            "[lifecycle]\nsystem_template = \"You are Saman.\"\nenabled = true\n",
1767        )
1768        .expect("write overlay");
1769
1770        let merged = state
1771            .merge_user_config(
1772                &home,
1773                Some("[lifecycle]\nsystem_template = \"shared default\"\n"),
1774            )
1775            .expect("lifecycle overlay is allowlisted");
1776
1777        let v = merged.parse::<toml::Value>().expect("valid TOML");
1778        assert_eq!(
1779            v["lifecycle"]["system_template"].as_str(),
1780            Some("You are Saman."),
1781            "agent persona wins over shared default"
1782        );
1783        assert_eq!(
1784            v["lifecycle"].get("enabled"),
1785            None,
1786            "WASM extension switch is not user-settable"
1787        );
1788    }
1789
1790    /// 16D integration guard: dev-agent principals (user_key `agent-<name>`)
1791    /// get their OWN session bucket and home-dir namespace — one agent can
1792    /// never see another's sessions, and both go through the same
1793    /// user_hash isolation path as humans.
1794    #[tokio::test]
1795    async fn agent_principals_get_isolated_session_buckets() {
1796        let state = test_state();
1797        let key_a = "agent-farzan";
1798        let key_b = "agent-paydar";
1799
1800        let (sid_a, session_a, _tx_a, _ts_a) = state.ensure_active_session(key_a).await;
1801        let (_sid_b, _session_b, _tx_b, _ts_b) = state.ensure_active_session(key_b).await;
1802
1803        // Each agent resolves its own session; B cannot see A's by id.
1804        assert!(
1805            state.get_session_by_any_id(key_a, &sid_a).await.is_some(),
1806            "owner bucket resolves its own session"
1807        );
1808        assert!(
1809            state.get_session_by_any_id(key_b, &sid_a).await.is_none(),
1810            "cross-agent session access must be 404/None"
1811        );
1812        assert!(Arc::strong_count(&session_a) >= 1);
1813
1814        // Same isolation path as humans: home dir = users_root/user_hash(key).
1815        if let Some(home) = state.get_user_home_dir(key_a) {
1816            assert_eq!(
1817                home.file_name().and_then(|n| n.to_str()),
1818                Some(trustee_core::user_hash(key_a)).as_deref()
1819            );
1820        }
1821    }
1822}
1823
1824// ---------------------------------------------------------------------------
1825// 16C — per-user McpToolLoader cache
1826// ---------------------------------------------------------------------------
1827
1828#[cfg(test)]
1829mod mcp_loader_cache_tests {
1830    use super::*;
1831
1832    fn state_with_shared(shared: &str) -> ServerState {
1833        let (session, _rx) = Session::new();
1834        let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
1835        let mut state = ServerState::new(session, ws_tx, None);
1836        state.config_toml = Some(shared.to_string());
1837        state
1838    }
1839
1840    /// Deterministic throwaway user: real user-home path under
1841    /// ~/.trustee/users/{hash} (that IS the resolution path under test),
1842    /// with config/ subdir; caller cleans up.
1843    struct TempUser {
1844        key: String,
1845        home: std::path::PathBuf,
1846    }
1847
1848    impl TempUser {
1849        fn new(tag: &str) -> Self {
1850            let key = format!("16c-{tag}-{}@test.invalid", std::process::id());
1851            let home = dirs::home_dir()
1852                .expect("HOME available in test env")
1853                .join(".trustee")
1854                .join("users")
1855                .join(trustee_core::user_hash(&key));
1856            std::fs::create_dir_all(home.join("config")).expect("create user home");
1857            Self { key, home }
1858        }
1859
1860        fn write_overlay(&self, toml_str: &str) {
1861            std::fs::write(self.home.join("config").join("trustee.toml"), toml_str)
1862                .expect("write overlay");
1863        }
1864    }
1865
1866    impl Drop for TempUser {
1867        fn drop(&mut self) {
1868            let _ = std::fs::remove_dir_all(&self.home);
1869        }
1870    }
1871
1872    #[tokio::test]
1873    async fn no_mcp_config_caches_disabled_marker() {
1874        let state = state_with_shared("[server]\nport = 8080\n");
1875        let user = TempUser::new("nomcp");
1876        let ts = Arc::new(pep::MemoryTokenStore::new());
1877
1878        let first = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
1879        assert!(first.is_none(), "no [mcp] anywhere → disabled marker");
1880        assert_eq!(state.mcp_loaders.len(), 1, "exactly one cache entry");
1881
1882        // Second call is a fingerprint hit — same disabled marker, still one entry.
1883        let second = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
1884        assert!(second.is_none());
1885        assert_eq!(state.mcp_loaders.len(), 1);
1886    }
1887
1888    #[tokio::test]
1889    async fn concurrent_cold_builds_single_flight() {
1890        let state = Arc::new(state_with_shared("[server]\nport = 8080\n"));
1891        let user = TempUser::new("singleflight");
1892        let ts = Arc::new(pep::MemoryTokenStore::new());
1893
1894        let mut handles = Vec::new();
1895        for _ in 0..5 {
1896            let state = state.clone();
1897            let key = user.key.clone();
1898            let ts = ts.clone();
1899            handles.push(tokio::spawn(async move {
1900                state.get_or_build_mcp_loader(&key, &ts).await
1901            }));
1902        }
1903        for h in handles {
1904            h.await.unwrap().expect("all five succeed");
1905        }
1906        assert_eq!(state.mcp_loaders.len(), 1, "single-flight → one entry");
1907    }
1908
1909    #[tokio::test]
1910    async fn fingerprint_change_triggers_rebuild() {
1911        let state = state_with_shared("[server]\nport = 8080\n");
1912        let user = TempUser::new("fpchange");
1913        let ts = Arc::new(pep::MemoryTokenStore::new());
1914
1915        // v1: an enabled MCP server pointing at a blackhole port — abk keeps
1916        // the loader with a DOWN status (connect refused is fast on loopback).
1917        user.write_overlay("[mcp]\nenabled = true\n\n[[mcp.servers]]\nname = \"v1\"\nurl = \"http://127.0.0.1:9/sse\"\n");
1918        let v1 = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
1919        assert!(v1.is_some(), "enabled [mcp] → real loader");
1920        let fp1 = state
1921            .mcp_loaders
1922            .get(&trustee_core::user_hash(&user.key))
1923            .unwrap()
1924            .fingerprint;
1925
1926        // v2: different server URL → different content fingerprint → rebuild.
1927        std::thread::sleep(std::time::Duration::from_millis(5));
1928        user.write_overlay("[mcp]\nenabled = true\n\n[[mcp.servers]]\nname = \"v2\"\nurl = \"http://127.0.0.1:9/other\"\n");
1929        let v2 = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
1930        assert!(v2.is_some());
1931        let entry = state
1932            .mcp_loaders
1933            .get(&trustee_core::user_hash(&user.key))
1934            .unwrap();
1935        assert_ne!(
1936            entry.fingerprint, fp1,
1937            "fingerprint must change with content"
1938        );
1939        assert!(entry.degraded.is_none());
1940
1941        // The old Arc stays valid for in-flight sessions (no panic, no revoke).
1942        let _still_usable = v1.as_ref().unwrap().tool_count;
1943    }
1944
1945    #[tokio::test]
1946    async fn degraded_entry_fails_loud_within_backoff_and_isolates_users() {
1947        let state = state_with_shared("[server]\nport = 8080\n");
1948        let bad = TempUser::new("degraded-bad");
1949        let good = TempUser::new("degraded-good");
1950        let ts = Arc::new(pep::MemoryTokenStore::new());
1951
1952        // [mcp] present but not a table → McpConfig deserialize fails → Err.
1953        bad.write_overlay("[mcp]\nenabled = \"not-a-bool\"\n");
1954        let err = match state.get_or_build_mcp_loader(&bad.key, &ts).await {
1955            Ok(_) => panic!("invalid [mcp] must fail loud"),
1956            Err(e) => e,
1957        };
1958        assert!(
1959            err.contains("invalid [mcp]"),
1960            "surfaces the parse error: {err}"
1961        );
1962
1963        let entry = state
1964            .mcp_loaders
1965            .get(&trustee_core::user_hash(&bad.key))
1966            .unwrap();
1967        assert!(entry.degraded.is_some(), "poison entry recorded");
1968
1969        // Within backoff: fast-fail again (cached error).
1970        let err2 = match state.get_or_build_mcp_loader(&bad.key, &ts).await {
1971            Ok(_) => panic!("still within backoff"),
1972            Err(e) => e,
1973        };
1974        assert_eq!(err, err2, "same cached error");
1975
1976        // Other users are completely unaffected.
1977        let good_loader = state
1978            .get_or_build_mcp_loader(&good.key, &ts)
1979            .await
1980            .expect("other user unaffected");
1981        assert!(
1982            good_loader.is_none(),
1983            "good user has no [mcp] → disabled marker"
1984        );
1985    }
1986}
1987
1988#[cfg(test)]
1989mod restore_tests {
1990    use super::*;
1991
1992    fn fresh_state() -> ServerState {
1993        let (session, _rx) = Session::new();
1994        let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
1995        let mut state = ServerState::new(session, ws_tx, None);
1996        state.config_toml = Some("[agent]\nname = \"trustee\"\n".to_string());
1997        state
1998    }
1999
2000    /// Deterministic throwaway user (same pattern as mcp_loader_cache_tests).
2001    struct TempUser {
2002        key: String,
2003    }
2004
2005    impl TempUser {
2006        fn new(tag: &str) -> Self {
2007            let key = format!("restore-{tag}-{}@test.invalid", std::process::id());
2008            let home = dirs::home_dir()
2009                .expect("HOME available in test env")
2010                .join(".trustee")
2011                .join("users")
2012                .join(trustee_core::user_hash(&key));
2013            std::fs::create_dir_all(home).expect("create user home");
2014            Self { key }
2015        }
2016    }
2017
2018    #[tokio::test]
2019    async fn restore_unknown_chain_returns_none() {
2020        let state = fresh_state();
2021        let user = TempUser::new("unknown");
2022        let out = state
2023            .get_or_restore_session_by_any_id(&user.key, "session_2099_01_01_00_00_deadbeef")
2024            .await;
2025        assert!(out.is_none(), "unknown chain must 404, not restore");
2026    }
2027
2028    #[tokio::test]
2029    async fn destroy_tombstones_chain_id() {
2030        let state = fresh_state();
2031        let user = TempUser::new("tombstone");
2032        let sid = state
2033            .create_session(&user.key, None, None, true)
2034            .await
2035            .expect("create session");
2036
2037        // Attach the live session to an on-disk-style chain id.
2038        let (session, _ws) = state.get_session(&user.key, &sid).await.unwrap();
2039        {
2040            let mut s = session.lock().await;
2041            s.session_id = Some("session_2026_09_04_10_00_aaaabbbb".to_string());
2042        }
2043
2044        state.destroy_session(&user.key, &sid).await.expect("destroy");
2045
2046        let path = state
2047            .tombstone_path(&user.key, "session_2026_09_04_10_00_aaaabbbb")
2048            .expect("tombstone path");
2049        assert!(path.exists(), "destroy must write a chain tombstone");
2050
2051        // And a fresh live session with NO chain id must not tombstone anything.
2052        let sid2 = state
2053            .create_session(&user.key, None, None, true)
2054            .await
2055            .expect("create second session");
2056        state.destroy_session(&user.key, &sid2).await.expect("destroy 2");
2057        // (No assertion beyond "no panic": session_id was None → no new markers.)
2058    }
2059
2060    #[tokio::test]
2061    async fn restore_refuses_registry_key_shapes() {
2062        // Registry keys ("default", "web…") are not checkpoint chains; even
2063        // if create_resume_info could not find them anyway, the tombstone
2064        // guard short-circuits non-`session_` ids without touching disk.
2065        let state = fresh_state();
2066        let user = TempUser::new("shapes");
2067        assert!(state.tombstone_path(&user.key, "default").is_none());
2068        let out = state.get_or_restore_session_by_any_id(&user.key, "default").await;
2069        assert!(out.is_none());
2070    }
2071}