1use 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
24pub struct UserSessionEntry {
30 pub session: Arc<Mutex<Session>>,
32 pub ws_tx: broadcast::Sender<String>,
34 pub created_at: chrono::DateTime<chrono::Utc>,
36 pub last_active: Arc<Mutex<chrono::DateTime<chrono::Utc>>>,
39}
40
41pub struct UserSessions {
43 pub sessions: DashMap<String, UserSessionEntry>,
45 pub token_store: Arc<pep::MemoryTokenStore>,
47 pub active_session_id: Mutex<String>,
49}
50
51#[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 pub handoff_count: u32,
64}
65
66#[derive(Debug)]
68pub enum SessionError {
69 MaxSessionsReached(usize),
71 NotFound(String),
73 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
91pub type SessionRegistry = Arc<DashMap<String, UserSessions>>;
93
94pub struct McpLoaderEntry {
105 pub loader: Option<std::sync::Arc<abk::agent::McpToolLoader>>,
107 pub fingerprint: u64,
110 pub built_at: chrono::DateTime<chrono::Utc>,
111 pub degraded: Option<String>,
114 pub failed_at: Option<chrono::DateTime<chrono::Utc>>,
117}
118
119pub const MCP_BUILD_RETRY_BACKOFF: std::time::Duration = std::time::Duration::from_secs(30);
121
122#[derive(Debug, Clone)]
129pub struct ThqDispatchEntry {
130 pub user_key: String,
135 pub service_token: Option<String>,
139 pub issuer_url: Option<String>,
145}
146
147pub 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#[derive(Clone)]
173pub struct ServerState {
174 pub sessions: SessionRegistry,
176 pub ws_tx: broadcast::Sender<String>,
178 pub auth: Option<Arc<AuthState>>,
180 pub config_toml: Option<String>,
182 pub secrets: Option<std::collections::HashMap<String, String>>,
184 pub build_info: Option<trustee_core::types::BuildInfo>,
186 pub workflow_semaphore: Arc<tokio::sync::Semaphore>,
189 pub max_sessions_per_user: usize,
191 pub allow_llm_overlay: bool,
194 pub mcp_loaders: Arc<DashMap<String, McpLoaderEntry>>,
196 mcp_build_locks: Arc<DashMap<String, Arc<tokio::sync::Mutex<()>>>>,
198 pub thq_dispatch: Arc<DashMap<String, ThqDispatchEntry>>,
200 pub agent_dispatch_tokens: Arc<DashMap<String, (String, std::time::Instant)>>,
203}
204
205impl ServerState {
206 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 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 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 pub fn with_max_sessions_per_user(mut self, max: usize) -> Self {
288 self.max_sessions_per_user = max;
289 self
290 }
291
292 pub fn with_allow_llm_overlay(mut self, allow: bool) -> Self {
295 self.allow_llm_overlay = allow;
296 self
297 }
298
299 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 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 if user_sessions.sessions.len() >= self.max_sessions_per_user {
327 return Err(SessionError::MaxSessionsReached(self.max_sessions_per_user));
328 }
329
330 let (mut session, workflow_rx) = Session::new();
332
333 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 self.apply_user_isolation(&mut session, user_key);
353
354 session.session_name = session_name;
356
357 session.identity = identity;
359
360 let (ws_tx_entry, _) = broadcast::channel::<String>(256);
362
363 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 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 if activate {
392 *user_sessions.active_session_id.lock().await = session_id.clone();
393 }
394
395 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 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 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 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 let user_sessions = self.sessions.get(user_key)?;
453 if let Some(entry) = user_sessions.sessions.get(id) {
454 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 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 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 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 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 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 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 if let Some(path) = self.tombstone_path(user_key, chain_id) {
551 if path.exists() {
552 return None;
553 }
554 }
555
556 let (config_toml, home_dir) = self.get_user_config_and_home(user_key);
558 let config_toml = config_toml?;
559
560 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 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 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 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; 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 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 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 items.sort_by(|a, b| b.last_active.cmp(&a.last_active));
695 items
696 }
697
698 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 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 user_sessions.sessions.remove(session_id);
730
731 if let Some(chain_id) = chain_id_for_tombstone {
735 self.write_tombstone(user_key, &chain_id);
736 }
737
738 let mut active_id = user_sessions.active_session_id.lock().await;
740 if &*active_id == session_id {
741 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 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 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 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 }
798
799 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 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 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 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 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 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 fn apply_user_isolation(&self, session: &mut Session, user_key: &str) {
890 let user_hash = trustee_core::user_hash(user_key);
891
892 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 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 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 if let Some(ref mut config_toml) = session.config_toml {
928 substitute_env_vars(config_toml, &merged_secrets);
929 }
930
931 session.secrets = Some(shared_secrets);
933 }
934
935 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 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 || section == "thq"
1003 || section == "lifecycle"
1008 || (self.allow_llm_overlay && section == "llm")
1009 };
1010 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 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 pub fn resolve_user_config(&self, user_key: &str) -> Option<String> {
1066 let config_toml = self.config_toml.clone()?;
1067
1068 let user_home = self.get_user_home_dir(user_key)?;
1070
1071 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 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_env_vars(&mut resolved, &merged_secrets);
1085
1086 Some(resolved)
1087 }
1088
1089 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 let resolved = self.resolve_user_config(user_key);
1113 let fingerprint = fingerprint_mcp_section(resolved.as_deref());
1114
1115 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 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 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 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 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 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 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 pub fn spawn_drain_task(self, mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>) {
1291 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 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 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 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 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
1397struct 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
1543fn 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 deep_merge_toml(base_val, overlay_val);
1561 }
1562 None => {
1563 base_table.insert(key.clone(), overlay_val.clone());
1565 }
1566 }
1567 }
1568 }
1569 (base, overlay) => {
1571 *base = overlay.clone();
1572 }
1573 }
1574}
1575
1576fn substitute_env_vars(s: &mut String, secrets: &std::collections::HashMap<String, String>) {
1581 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 if let Some(end) = s[i + 2..].find('}') {
1590 let var_name = &s[i + 2..i + 2 + end];
1591 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 result.push_str(&s[i..i + 2 + end + 1]);
1599 }
1600 i = i + 2 + end + 1;
1601 } else {
1602 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
1615fn 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 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 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 #[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 #[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 #[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 #[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 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 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 #[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 }
1757
1758 #[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 #[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 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 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#[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 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 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 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 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 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 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 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 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 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 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 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 }
2059
2060 #[tokio::test]
2061 async fn restore_refuses_registry_key_shapes() {
2062 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}