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
24tokio::task_local! {
29 static USER_KEY_OVERRIDE: String;
37}
38
39pub(crate) fn effective_user_key(user_key: &str) -> String {
42 USER_KEY_OVERRIDE
43 .try_with(|v| v.clone())
44 .unwrap_or_else(|_| user_key.to_string())
45}
46
47pub(crate) async fn in_dispatch_scope<F: std::future::Future>(
50 user_key: &str,
51 fut: F,
52) -> F::Output {
53 USER_KEY_OVERRIDE.scope(user_key.to_string(), fut).await
54}
55
56pub fn apply_home_links() -> Vec<String> {
66 let Some(home) = dirs::home_dir() else {
67 return Vec::new();
68 };
69 let users = home.join(".trustee").join("users");
70 let map_path = users.join("home_links.json");
71 let Ok(text) = std::fs::read_to_string(&map_path) else {
72 return Vec::new();
73 };
74 let Ok(doc) = serde_json::from_str::<serde_json::Value>(&text) else {
75 tracing::warn!(
76 "home_links.json is not valid JSON — ignored: {}",
77 map_path.display()
78 );
79 return Vec::new();
80 };
81 let mut report = Vec::new();
82 if let Some(links) = doc.get("links").and_then(|v| v.as_array()) {
83 for l in links {
84 let (Some(from), Some(to)) = (
85 l.get("from_hash").and_then(|v| v.as_str()),
86 l.get("to_hash").and_then(|v| v.as_str()),
87 ) else {
88 continue;
89 };
90 let from_dir = users.join(from);
91 let to_dir = users.join(to);
92 if !from_dir.exists() {
93 continue; }
95 if to_dir.exists() {
96 let empty = std::fs::read_dir(&to_dir)
97 .map(|it| it.count() == 0)
98 .unwrap_or(false);
99 if !empty {
100 let msg = format!(
101 "home link {from} → {to}: target exists with content — NOT touched, reconcile manually"
102 );
103 tracing::warn!("home_links: {msg}");
104 report.push(msg);
105 continue;
106 }
107 let _ = std::fs::remove_dir(&to_dir);
108 }
109 match std::fs::rename(&from_dir, &to_dir) {
110 Ok(()) => {
111 let msg = format!(
112 "home link applied: {from} → {to} (config/sessions restored)"
113 );
114 tracing::info!("home_links: {msg}");
115 report.push(msg);
116 }
117 Err(e) => {
118 let msg = format!("home link {from} → {to} FAILED: {e}");
119 tracing::warn!("home_links: {msg}");
120 report.push(msg);
121 }
122 }
123 }
124 }
125 report
126}
127
128pub struct UserSessionEntry {
134 pub session: Arc<Mutex<Session>>,
136 pub ws_tx: broadcast::Sender<String>,
138 pub created_at: chrono::DateTime<chrono::Utc>,
140 pub last_active: Arc<Mutex<chrono::DateTime<chrono::Utc>>>,
143}
144
145pub struct UserSessions {
147 pub sessions: DashMap<String, UserSessionEntry>,
149 pub token_store: Arc<pep::MemoryTokenStore>,
151 pub active_session_id: Mutex<String>,
153}
154
155#[derive(Debug, serde::Serialize)]
157pub struct SessionListItem {
158 pub session_id: String,
159 pub session_name: Option<String>,
160 pub workflow_state: String,
161 pub created_at: String,
162 pub last_active: String,
163 pub handoff_count: u32,
168}
169
170#[derive(Debug)]
172pub enum SessionError {
173 MaxSessionsReached(usize),
175 NotFound(String),
177 NotIdle(String),
179}
180
181impl std::fmt::Display for SessionError {
182 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183 match self {
184 SessionError::MaxSessionsReached(n) => {
185 write!(f, "Maximum {} sessions per user reached", n)
186 }
187 SessionError::NotFound(id) => write!(f, "Session {} not found", id),
188 SessionError::NotIdle(state) => write!(f, "Session is not idle (state: {})", state),
189 }
190 }
191}
192
193impl std::error::Error for SessionError {}
194
195pub type SessionRegistry = Arc<DashMap<String, UserSessions>>;
197
198pub struct McpLoaderEntry {
209 pub loader: Option<std::sync::Arc<abk::agent::McpToolLoader>>,
211 pub fingerprint: u64,
214 pub built_at: chrono::DateTime<chrono::Utc>,
215 pub degraded: Option<String>,
218 pub failed_at: Option<chrono::DateTime<chrono::Utc>>,
221}
222
223pub const MCP_BUILD_RETRY_BACKOFF: std::time::Duration = std::time::Duration::from_secs(30);
225
226#[derive(Debug, Clone)]
233pub struct ThqDispatchEntry {
234 pub user_key: String,
239 pub service_token: Option<String>,
243 pub issuer_url: Option<String>,
249}
250
251pub fn service_issuers_from_config(config_toml: &str) -> Vec<String> {
255 let Ok(v) = config_toml.parse::<toml::Value>() else {
256 return Vec::new();
257 };
258 let Some(creds) = v.get("mcp").and_then(|m| m.get("credentials")).and_then(|c| c.as_table())
259 else {
260 return Vec::new();
261 };
262 let mut out = Vec::new();
263 for (_name, cred) in creds {
264 if cred.get("type").and_then(|t| t.as_str()) == Some("service-account") {
265 if let Some(issuer) = cred.get("issuer_url").and_then(|i| i.as_str()) {
266 if !out.iter().any(|e: &String| e == issuer) {
267 out.push(issuer.to_string());
268 }
269 }
270 }
271 }
272 out
273}
274
275#[derive(Clone)]
277pub struct ServerState {
278 pub sessions: SessionRegistry,
280 pub ws_tx: broadcast::Sender<String>,
282 pub auth: Option<Arc<AuthState>>,
284 pub config_toml: Option<String>,
286 pub secrets: Option<std::collections::HashMap<String, String>>,
288 pub build_info: Option<trustee_core::types::BuildInfo>,
290 pub workflow_semaphore: Arc<tokio::sync::Semaphore>,
293 pub max_sessions_per_user: usize,
295 pub allow_llm_overlay: bool,
298 pub mcp_loaders: Arc<DashMap<String, McpLoaderEntry>>,
300 mcp_build_locks: Arc<DashMap<String, Arc<tokio::sync::Mutex<()>>>>,
302 pub thq_dispatch: Arc<DashMap<String, ThqDispatchEntry>>,
304 pub agent_dispatch_tokens: Arc<DashMap<String, (String, std::time::Instant)>>,
307}
308
309impl ServerState {
310 pub fn new(
312 session: Session,
313 ws_tx: broadcast::Sender<String>,
314 auth: Option<Arc<AuthState>>,
315 ) -> Self {
316 apply_home_links();
319
320 let sessions = Arc::new(DashMap::new());
321
322 let token_store = Arc::new(pep::MemoryTokenStore::new());
324 let (ws_tx_entry, _) = broadcast::channel::<String>(256);
325
326 let now = chrono::Utc::now();
327 let initial_entry = UserSessionEntry {
328 session: Arc::new(Mutex::new(session)),
329 ws_tx: ws_tx_entry,
330 created_at: now,
331 last_active: Arc::new(Mutex::new(now)),
332 };
333
334 let user_sessions = UserSessions {
335 sessions: DashMap::new(),
336 token_store,
337 active_session_id: Mutex::new(String::new()),
338 };
339 user_sessions.sessions.insert("default".to_string(), initial_entry);
340
341 sessions.insert("default".to_string(), user_sessions);
342
343 Self {
344 sessions,
345 ws_tx,
346 auth,
347 config_toml: None,
348 secrets: None,
349 build_info: None,
350 workflow_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
351 max_sessions_per_user: 4,
352 allow_llm_overlay: false,
353 mcp_loaders: Arc::new(DashMap::new()),
354 mcp_build_locks: Arc::new(DashMap::new()),
355 thq_dispatch: Arc::new(DashMap::new()),
356 agent_dispatch_tokens: Arc::new(DashMap::new()),
357 }
358 }
359
360 pub fn with_config_toml(mut self, config_toml: String) -> Self {
361 self.config_toml = Some(config_toml);
362 self
363 }
364
365 pub fn service_issuers(&self) -> Vec<String> {
376 service_issuers_from_config(self.config_toml.as_deref().unwrap_or(""))
377 }
378
379 pub fn with_secrets(mut self, secrets: std::collections::HashMap<String, String>) -> Self {
380 self.secrets = Some(secrets);
381 self
382 }
383
384 pub fn with_build_info(mut self, build_info: trustee_core::types::BuildInfo) -> Self {
385 self.build_info = Some(build_info);
386 self
387 }
388
389 pub fn with_max_concurrent_workflows(mut self, max: usize) -> Self {
390 self.workflow_semaphore = Arc::new(tokio::sync::Semaphore::new(max));
391 self
392 }
393
394 pub fn with_max_sessions_per_user(mut self, max: usize) -> Self {
396 self.max_sessions_per_user = max;
397 self
398 }
399
400 pub fn with_allow_llm_overlay(mut self, allow: bool) -> Self {
403 self.allow_llm_overlay = allow;
404 self
405 }
406
407 pub async fn create_session(
417 &self,
418 user_key: &str,
419 session_name: Option<String>,
420 identity: Option<String>,
421 activate: bool,
422 ) -> Result<String, SessionError> {
423 let user_sessions = self
425 .sessions
426 .entry(user_key.to_string())
427 .or_insert_with(|| UserSessions {
428 sessions: DashMap::new(),
429 token_store: Arc::new(pep::MemoryTokenStore::new()),
430 active_session_id: Mutex::new(String::new()),
431 });
432
433 if user_sessions.sessions.len() >= self.max_sessions_per_user {
435 return Err(SessionError::MaxSessionsReached(self.max_sessions_per_user));
436 }
437
438 let (mut session, workflow_rx) = Session::new();
440
441 if let Some(ref config_toml) = self.config_toml {
443 session.config_toml = Some(config_toml.clone());
444 session.parse_auto_handoff_config();
445 if let Ok(table) = config_toml.parse::<toml::Value>() {
446 if let Some(name) = table
447 .get("agent")
448 .and_then(|a| a.get("name"))
449 .and_then(|n| n.as_str())
450 {
451 session.agent_name = name.to_string();
452 }
453 }
454 }
455
456 session.secrets = self.secrets.clone();
457 session.build_info = self.build_info.clone();
458
459 self.apply_user_isolation(&mut session, user_key);
461
462 session.session_name = session_name;
464
465 session.identity = identity;
467
468 let (ws_tx_entry, _) = broadcast::channel::<String>(256);
470
471 let session_id = format!(
473 "session_{}_{}",
474 chrono::Utc::now().format("%Y_%m_%d_%H_%M"),
475 &uuid::Uuid::new_v4().to_string()[..8]
476 );
477
478 let now = chrono::Utc::now();
479
480 user_sessions.sessions.insert(
482 session_id.clone(),
483 UserSessionEntry {
484 session: Arc::new(Mutex::new(session)),
485 ws_tx: ws_tx_entry.clone(),
486 created_at: now,
487 last_active: Arc::new(Mutex::new(now)),
488 },
489 );
490
491 if activate {
500 *user_sessions.active_session_id.lock().await = session_id.clone();
501 }
502
503 let session_arc = user_sessions
505 .sessions
506 .get(&session_id)
507 .map(|e| e.session.clone());
508 if let Some(session_arc) = session_arc {
509 self.spawn_user_drain_task(
510 session_id.clone(),
511 session_arc,
512 ws_tx_entry,
513 workflow_rx,
514 );
515 }
516
517 Ok(session_id)
518 }
519
520 pub async fn get_session(
523 &self,
524 user_key: &str,
525 session_id: &str,
526 ) -> Option<(Arc<Mutex<Session>>, broadcast::Sender<String>)> {
527 let user_sessions = self.sessions.get(user_key)?;
528 let entry = user_sessions.sessions.get(session_id)?;
529
530 let now = chrono::Utc::now();
532 *entry.last_active.lock().await = now;
533
534 Some((entry.session.clone(), entry.ws_tx.clone()))
535 }
536
537 pub async fn get_session_by_any_id(
555 &self,
556 user_key: &str,
557 id: &str,
558 ) -> Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> {
559 let user_sessions = self.sessions.get(user_key)?;
561 if let Some(entry) = user_sessions.sessions.get(id) {
562 let now = chrono::Utc::now();
564 *entry.last_active.lock().await = now;
565 return Some((id.to_string(), entry.session.clone(), entry.ws_tx.clone()));
566 }
567
568 for entry in user_sessions.sessions.iter() {
570 let session = entry.session.lock().await;
571 if session.session_id.as_deref() == Some(id) {
572 let key = entry.key().clone();
573 let ws_tx = entry.ws_tx.clone();
574 drop(session);
575 let now = chrono::Utc::now();
577 *entry.last_active.lock().await = now;
578 return Some((key, entry.session.clone(), ws_tx));
579 }
580 }
581
582 None
583 }
584
585 fn tombstone_path(&self, user_key: &str, chain_id: &str) -> Option<std::path::PathBuf> {
594 if chain_id.is_empty() || !chain_id.starts_with("session_") {
595 return None;
596 }
597 let user_key = effective_user_key(user_key);
598 let hash = trustee_core::user_hash(&user_key);
599 dirs::home_dir().map(|home| {
600 home.join(".trustee")
601 .join("users")
602 .join(&hash)
603 .join("session_tombstones")
604 .join(chain_id)
605 })
606 }
607
608 fn write_tombstone(&self, user_key: &str, chain_id: &str) {
613 if let Some(path) = self.tombstone_path(user_key, chain_id) {
614 if let Some(parent) = path.parent() {
615 let _ = std::fs::create_dir_all(parent);
616 }
617 let _ = std::fs::write(&path, chrono::Utc::now().to_rfc3339());
618 }
619 }
620
621 pub async fn get_or_restore_session_by_any_id(
635 &self,
636 user_key: &str,
637 id: &str,
638 ) -> Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> {
639 if let Some(found) = self.get_session_by_any_id(user_key, id).await {
640 return Some(found);
641 }
642 self.restore_session_from_checkpoint(user_key, id).await
643 }
644
645 async fn restore_session_from_checkpoint(
654 &self,
655 user_key: &str,
656 chain_id: &str,
657 ) -> Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> {
658 if let Some(path) = self.tombstone_path(user_key, chain_id) {
660 if path.exists() {
661 return None;
662 }
663 }
664
665 let (config_toml, home_dir) = self.get_user_config_and_home(user_key);
667 let config_toml = config_toml?;
668
669 let info = trustee_core::sessions::create_resume_info(
672 &config_toml,
673 chain_id,
674 home_dir.as_deref(),
675 )
676 .await
677 .ok()??;
678
679 let user_sessions = self
682 .sessions
683 .entry(user_key.to_string())
684 .or_insert_with(|| UserSessions {
685 sessions: DashMap::new(),
686 token_store: Arc::new(pep::MemoryTokenStore::new()),
687 active_session_id: Mutex::new(String::new()),
688 });
689
690 if let Some(existing) = user_sessions.sessions.get(chain_id) {
691 let now = chrono::Utc::now();
692 *existing.last_active.lock().await = now;
693 return Some((
694 chain_id.to_string(),
695 existing.session.clone(),
696 existing.ws_tx.clone(),
697 ));
698 }
699
700 if user_sessions.sessions.len() >= self.max_sessions_per_user {
705 let mut evict: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
706 for entry in user_sessions.sessions.iter() {
707 let is_idle = entry.session.lock().await.workflow_state
708 == trustee_core::types::WorkflowState::Idle;
709 if !is_idle {
710 continue;
711 }
712 let la = entry.last_active.lock().await;
713 if evict.as_ref().is_none_or(|(_, t)| *la < *t) {
714 evict = Some((entry.key().clone(), *la));
715 }
716 }
717 if let Some((victim, _)) = evict {
718 user_sessions.sessions.remove(&victim);
719 } else {
720 return None;
721 }
722 }
723
724 let (mut session, workflow_rx) = Session::new();
727 session.config_toml = Some(config_toml.clone());
728 session.parse_auto_handoff_config();
729 if let Ok(table) = config_toml.parse::<toml::Value>() {
730 if let Some(name) = table
731 .get("agent")
732 .and_then(|a| a.get("name"))
733 .and_then(|n| n.as_str())
734 {
735 session.agent_name = name.to_string();
736 }
737 }
738 session.secrets = self.secrets.clone();
739 session.build_info = self.build_info.clone();
740 self.apply_user_isolation(&mut session, user_key);
741 session.session_id = Some(chain_id.to_string());
742 session.resume_info = Some(info);
743 session.session_name = None; let (ws_tx_entry, _) = broadcast::channel::<String>(256);
746 let now = chrono::Utc::now();
747
748 user_sessions.sessions.insert(
749 chain_id.to_string(),
750 UserSessionEntry {
751 session: Arc::new(Mutex::new(session)),
752 ws_tx: ws_tx_entry.clone(),
753 created_at: now,
754 last_active: Arc::new(Mutex::new(now)),
755 },
756 );
757
758 let session_arc = user_sessions
762 .sessions
763 .get(chain_id)
764 .map(|e| e.session.clone())
765 .ok_or_else(|| SessionError::NotFound(chain_id.to_string()))
766 .ok()?;
767 self.spawn_user_drain_task(
768 chain_id.to_string(),
769 session_arc.clone(),
770 ws_tx_entry.clone(),
771 workflow_rx,
772 );
773
774 Some((chain_id.to_string(), session_arc, ws_tx_entry))
775 }
776
777 pub async fn list_sessions(&self, user_key: &str) -> Vec<SessionListItem> { let Some(user_sessions) = self.sessions.get(user_key) else {
779 return Vec::new();
780 };
781
782 let mut items = Vec::new();
783 for entry in user_sessions.sessions.iter() {
784 let session = entry.session.lock().await;
785 let workflow_state = match session.workflow_state {
786 trustee_core::types::WorkflowState::Idle => "Idle",
787 trustee_core::types::WorkflowState::Running => "Running",
788 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
789 };
790 let last_active = entry.last_active.lock().await;
791 items.push(SessionListItem {
792 session_id: entry.key().clone(),
793 session_name: session.session_name.clone(),
794 workflow_state: workflow_state.to_string(),
795 created_at: entry.created_at.to_rfc3339(),
796 last_active: last_active.to_rfc3339(),
797 handoff_count: session.handoff_count,
798 });
799 }
800 drop(user_sessions);
801
802 items.sort_by(|a, b| b.last_active.cmp(&a.last_active));
804 items
805 }
806
807 pub async fn destroy_session(
809 &self,
810 user_key: &str,
811 session_id: &str,
812 ) -> Result<(), SessionError> {
813 let user_sessions = self
814 .sessions
815 .get(user_key)
816 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
817
818 let chain_id_for_tombstone = {
821 let entry = user_sessions
822 .sessions
823 .get(session_id)
824 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
825 let session = entry.session.lock().await;
826 if session.workflow_state != trustee_core::types::WorkflowState::Idle {
827 let state_str = match session.workflow_state {
828 trustee_core::types::WorkflowState::Running => "Running",
829 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
830 _ => "Unknown",
831 };
832 return Err(SessionError::NotIdle(state_str.to_string()));
833 }
834 session.session_id.clone()
835 };
836
837 user_sessions.sessions.remove(session_id);
839
840 if let Some(chain_id) = chain_id_for_tombstone {
844 self.write_tombstone(user_key, &chain_id);
845 }
846
847 let mut active_id = user_sessions.active_session_id.lock().await;
849 if &*active_id == session_id {
850 let mut newest: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
852 for entry in user_sessions.sessions.iter() {
853 let la = entry.last_active.lock().await;
854 if newest.as_ref().map_or(true, |(_, t)| *la > *t) {
855 newest = Some((entry.key().clone(), *la));
856 }
857 }
858 *active_id = newest.map(|(id, _)| id).unwrap_or_default();
859 }
860
861 Ok(())
862 }
863
864 pub async fn ensure_active_session(
873 &self,
874 user_key: &str,
875 ) -> (
876 String,
877 Arc<Mutex<Session>>,
878 broadcast::Sender<String>,
879 Arc<pep::MemoryTokenStore>,
880 ) {
881 let token_store = {
883 let user_sessions = self
884 .sessions
885 .entry(user_key.to_string())
886 .or_insert_with(|| UserSessions {
887 sessions: DashMap::new(),
888 token_store: Arc::new(pep::MemoryTokenStore::new()),
889 active_session_id: Mutex::new(String::new()),
890 });
891 user_sessions.token_store.clone()
892 };
893
894 let active_id = {
896 let user_sessions = self.sessions.get(user_key).unwrap();
897 let guard = user_sessions.active_session_id.lock().await;
898 guard.clone()
899 };
900
901 if !active_id.is_empty() {
902 if let Some((session, ws_tx)) = self.get_session(user_key, &active_id).await {
903 return (active_id, session, ws_tx, token_store);
904 }
905 }
907
908 let existing_session: Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> = {
912 let user_sessions = self.sessions.get(user_key).unwrap();
913 let result = user_sessions.sessions.iter().next().map(|first| {
914 (
915 first.key().clone(),
916 first.session.clone(),
917 first.ws_tx.clone(),
918 )
919 });
920 result
921 };
922 if let Some((id, session, ws_tx)) = existing_session {
923 let now = chrono::Utc::now();
924 if let Some(entry) = self.sessions.get(user_key) {
925 if let Some(e) = entry.sessions.get(&id) {
926 *e.last_active.lock().await = now;
927 }
928 *entry.active_session_id.lock().await = id.clone();
929 }
930
931 return (id, session, ws_tx, token_store);
932 }
933
934 let session_id = self
936 .create_session(user_key, None, None, true)
937 .await
938 .unwrap_or_else(|_| "default".to_string());
939
940 let (session, ws_tx) = self
941 .get_session(user_key, &session_id)
942 .await
943 .expect("just-created session must exist");
944
945 (session_id, session, ws_tx, token_store)
946 }
947
948 pub async fn ensure_user_session(
951 &self,
952 user_key: &str,
953 ) -> (Arc<Mutex<Session>>, broadcast::Sender<String>, Arc<pep::MemoryTokenStore>) {
954 let (_id, session, ws_tx, token_store) = self.ensure_active_session(user_key).await;
955 (session, ws_tx, token_store)
956 }
957
958 pub async fn set_active_session(&self, user_key: &str, session_id: &str) {
960 if let Some(user_sessions) = self.sessions.get(user_key) {
961 if user_sessions.sessions.contains_key(session_id) {
962 *user_sessions.active_session_id.lock().await = session_id.to_string();
963 }
964 }
965 }
966
967 pub fn get_user_home_dir(&self, user_key: &str) -> Option<std::path::PathBuf> {
978 let user_key = effective_user_key(user_key);
979 let hash = trustee_core::user_hash(&user_key);
980 dirs::home_dir().map(|home| home.join(".trustee").join("users").join(&hash))
981 }
982
983 pub fn get_user_config_and_home(&self, user_key: &str) -> (Option<String>, Option<std::path::PathBuf>) {
988 (self.config_toml.clone(), self.get_user_home_dir(user_key))
989 }
990
991 fn apply_user_isolation(&self, session: &mut Session, user_key: &str) {
1000 let user_key = effective_user_key(user_key);
1001 let user_hash = trustee_core::user_hash(&user_key);
1002
1003 let user_home = if let Some(home) = dirs::home_dir() {
1005 let user_home = home.join(".trustee").join("users").join(&user_hash);
1006 session.home_dir = Some(user_home.clone());
1007 Some(user_home)
1008 } else {
1009 None
1010 };
1011
1012 session.project_id = Some(format!("web{}", &user_hash[..16]));
1013
1014 let shared_secrets = session.secrets.clone().unwrap_or_default();
1021 let mut merged_secrets = shared_secrets.clone();
1022
1023 if let Some(ref user_home) = user_home {
1024 if let Ok(merged) = self.load_user_secrets(user_home, &merged_secrets) {
1025 merged_secrets = merged;
1026 }
1027 }
1028
1029 if let Some(ref user_home) = user_home {
1031 if let Some(merged) = self.merge_user_config(user_home, session.config_toml.as_deref()) {
1032 session.config_toml = Some(merged);
1033 tracing::debug!("Merged per-user config into session");
1034 }
1035 }
1036
1037 if let Some(ref mut config_toml) = session.config_toml {
1039 substitute_env_vars(config_toml, &merged_secrets);
1040 }
1041
1042 session.secrets = Some(shared_secrets);
1044 }
1045
1046 fn load_user_secrets(
1049 &self,
1050 user_home: &std::path::Path,
1051 base: &std::collections::HashMap<String, String>,
1052 ) -> std::io::Result<std::collections::HashMap<String, String>> {
1053 let user_env_path = user_home.join(".env");
1054 if !user_env_path.exists() {
1055 return Ok(base.clone());
1056 }
1057 let content = std::fs::read_to_string(&user_env_path)?;
1058 let mut merged = base.clone();
1059 for line in content.lines() {
1060 let line = line.trim();
1061 if line.is_empty() || line.starts_with('#') {
1062 continue;
1063 }
1064 if let Some((key, value)) = line.split_once('=') {
1065 let key = key.trim().to_string();
1066 let value = value
1067 .trim()
1068 .trim_matches('"')
1069 .trim_matches('\'')
1070 .to_string();
1071 merged.insert(key, value);
1072 }
1073 }
1074 tracing::debug!("Loaded per-user secrets from {}", user_env_path.display());
1075 Ok(merged)
1076 }
1077
1078 fn merge_user_config(
1093 &self,
1094 user_home: &std::path::Path,
1095 shared_config: Option<&str>,
1096 ) -> Option<String> {
1097 let user_config_path = user_home.join("config").join("trustee.toml");
1098 if !user_config_path.exists() {
1099 return None;
1100 }
1101 let user_config_toml = std::fs::read_to_string(&user_config_path).ok()?;
1102 let shared = shared_config
1103 .unwrap_or("")
1104 .parse::<toml::Value>()
1105 .ok()?;
1106 let overlay = user_config_toml.parse::<toml::Value>().ok()?;
1107
1108 let allowed = |section: &str| {
1109 section == "mcp"
1110 || section == "thq"
1114 || section == "lifecycle"
1119 || (self.allow_llm_overlay && section == "llm")
1120 };
1121 let dir_name = user_home
1124 .file_name()
1125 .and_then(|n| n.to_str())
1126 .unwrap_or("<unknown>");
1127 let masked_user = dir_name.get(..8).unwrap_or(dir_name);
1128
1129 let mut filtered_overlay = toml::map::Map::new();
1130 if let Some(table) = overlay.as_table() {
1131 for (section, value) in table {
1132 if section == "lifecycle" {
1133 let mut v = value.clone();
1134 if let Some(t) = v.as_table_mut() {
1135 if t.remove("enabled").is_some() {
1136 tracing::info!(
1137 "user config overlay: [lifecycle].enabled is not user-settable — stripped for user {}",
1138 masked_user
1139 );
1140 }
1141 }
1142 filtered_overlay.insert(section.clone(), v);
1143 } else if allowed(section) {
1144 filtered_overlay.insert(section.clone(), value.clone());
1145 } else {
1146 tracing::warn!(
1147 "user config overlay: dropping non-allowlisted section [{}] for user {}",
1148 section,
1149 masked_user
1150 );
1151 }
1152 }
1153 }
1154
1155 if filtered_overlay.is_empty() {
1156 return None;
1158 }
1159 let overlay = toml::Value::Table(filtered_overlay);
1160
1161 let mut shared = shared;
1162 deep_merge_toml(&mut shared, &overlay);
1163 let merged = toml::to_string(&shared).ok()?;
1164 tracing::debug!("Merged per-user config from {}", user_config_path.display());
1165 Some(merged)
1166 }
1167
1168 pub fn resolve_user_config(&self, user_key: &str) -> Option<String> {
1177 let config_toml = self.config_toml.clone()?;
1178
1179 let user_home = self.get_user_home_dir(user_key)?;
1181
1182 let mut merged_secrets = self.secrets.clone().unwrap_or_default();
1184 if let Ok(merged) = self.load_user_secrets(&user_home, &merged_secrets) {
1185 merged_secrets = merged;
1186 }
1187
1188 let mut resolved = config_toml;
1190 if let Some(merged) = self.merge_user_config(&user_home, Some(&resolved)) {
1191 resolved = merged;
1192 }
1193
1194 substitute_env_vars(&mut resolved, &merged_secrets);
1196
1197 Some(resolved)
1198 }
1199
1200 pub async fn get_or_build_mcp_loader(
1216 &self,
1217 user_key: &str,
1218 token_store: &Arc<pep::MemoryTokenStore>,
1219 ) -> Result<Option<std::sync::Arc<abk::agent::McpToolLoader>>, String> {
1220 let user_key = effective_user_key(user_key);
1221 let user_hash = trustee_core::user_hash(&user_key);
1222
1223 let resolved = self.resolve_user_config(&user_key);
1225 let fingerprint = fingerprint_mcp_section(resolved.as_deref());
1226
1227 if let Some(entry) = self.mcp_loaders.get(&user_hash) {
1229 if entry.degraded.is_none() {
1230 if entry.fingerprint == fingerprint {
1231 return Ok(entry.loader.clone());
1232 }
1233 } else if let (Some(err), Some(failed_at)) = (&entry.degraded, entry.failed_at) {
1234 let backoff = chrono::Duration::from_std(MCP_BUILD_RETRY_BACKOFF)
1235 .unwrap_or_else(|_| chrono::Duration::seconds(30));
1236 if chrono::Utc::now() < failed_at + backoff {
1237 return Err(err.clone());
1238 }
1239 }
1240 }
1241
1242 let lock = self
1244 .mcp_build_locks
1245 .entry(user_hash.clone())
1246 .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
1247 .clone();
1248 let _guard = lock.lock().await;
1249
1250 if let Some(entry) = self.mcp_loaders.get(&user_hash) {
1252 if entry.degraded.is_none() && entry.fingerprint == fingerprint {
1253 return Ok(entry.loader.clone());
1254 }
1255 }
1256
1257 match self
1258 .build_mcp_loader(&user_hash, resolved.as_deref(), fingerprint, token_store)
1259 .await
1260 {
1261 Ok(entry) => {
1262 self.mcp_loaders.insert(user_hash.clone(), entry);
1263 Ok(self.mcp_loaders.get(&user_hash).unwrap().loader.clone())
1264 }
1265 Err(err) => {
1266 tracing::warn!(
1267 "MCP loader build FAILED for user {}; dispatch fails loud, retry after {:?}",
1268 &user_hash[..8.min(user_hash.len())],
1269 MCP_BUILD_RETRY_BACKOFF
1270 );
1271 self.mcp_loaders.insert(
1272 user_hash,
1273 McpLoaderEntry {
1274 loader: None,
1275 fingerprint,
1276 built_at: chrono::Utc::now(),
1277 degraded: Some(err.clone()),
1278 failed_at: Some(chrono::Utc::now()),
1279 },
1280 );
1281 Err(err)
1282 }
1283 }
1284 }
1285
1286 async fn build_mcp_loader(
1289 &self,
1290 user_hash: &str,
1291 resolved: Option<&str>,
1292 fingerprint: u64,
1293 token_store: &Arc<pep::MemoryTokenStore>,
1294 ) -> Result<McpLoaderEntry, String> {
1295 let mcp_config: Option<abk::config::McpConfig> = match resolved {
1296 Some(toml_str) => {
1297 let value = toml_str
1298 .parse::<toml::Value>()
1299 .map_err(|e| format!("config parse failed: {}", e))?;
1300 match value.get("mcp") {
1301 Some(section) => {
1302 use serde::Deserialize as _;
1303 Some(
1304 abk::config::McpConfig::deserialize(section.clone())
1305 .map_err(|e| format!("invalid [mcp] config: {}", e))?,
1306 )
1307 }
1308 None => None,
1309 }
1310 }
1311 None => None,
1312 };
1313
1314 let loader = match mcp_config {
1315 Some(cfg) if cfg.enabled => {
1316 let built = abk::agent::McpToolLoader::with_token_store(
1317 &cfg,
1318 Some(token_store.clone() as std::sync::Arc<dyn pep::token_store::TokenStore>),
1319 )
1320 .await
1321 .map_err(|e| format!("MCP loader build failed: {}", e))?;
1322
1323 let servers: Vec<String> = built
1327 .server_statuses
1328 .iter()
1329 .map(|s| {
1330 if s.connected {
1331 format!("{}(up,{}tools)", s.name, s.tool_count)
1332 } else {
1333 format!("{}(DOWN)", s.name)
1334 }
1335 })
1336 .collect();
1337 tracing::info!(
1338 "MCP loader built for user {}: servers=[{}] total_tools={}",
1339 &user_hash[..8.min(user_hash.len())],
1340 servers.join(", "),
1341 built.tool_count
1342 );
1343 Some(std::sync::Arc::new(built))
1344 }
1345 _ => None,
1346 };
1347
1348 Ok(McpLoaderEntry {
1349 loader,
1350 fingerprint,
1351 built_at: chrono::Utc::now(),
1352 degraded: None,
1353 failed_at: None,
1354 })
1355 }
1356
1357 fn spawn_user_drain_task(
1359 &self,
1360 session_id: String,
1361 session: Arc<Mutex<Session>>,
1362 ws_tx: broadcast::Sender<String>,
1363 mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>,
1364 ) {
1365 let mut last_broadcast_state: Option<String> = None;
1371 tokio::spawn(async move {
1372 while let Some(msg) = workflow_rx.recv().await {
1373 {
1374 let mut session = session.lock().await;
1375 session.handle_workflow_message(msg.clone());
1376
1377 let state_str = match session.workflow_state {
1378 trustee_core::types::WorkflowState::Idle => "Idle",
1379 trustee_core::types::WorkflowState::Running => "Running",
1380 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
1381 };
1382 if last_broadcast_state.as_deref() != Some(state_str) {
1383 last_broadcast_state = Some(state_str.to_string());
1384 let state_msg = serde_json::json!({
1385 "type": "StateChanged",
1386 "state": state_str
1387 });
1388 let _ = ws_tx.send(state_msg.to_string());
1389 }
1390 }
1391
1392 let json =
1393 serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
1394 let _ = ws_tx.send(json);
1395 }
1396 tracing::debug!("Drain task ended for session: {}", session_id);
1397 });
1398 }
1399
1400 pub fn spawn_drain_task(self, mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>) {
1403 let default_user = self
1405 .sessions
1406 .get("default")
1407 .expect("default user must exist");
1408 let first_entry = default_user
1409 .sessions
1410 .iter()
1411 .next()
1412 .expect("default user must have at least one session");
1413 let session = first_entry.session.clone();
1414 let ws_tx = first_entry.ws_tx.clone();
1415 let session_id = first_entry.key().clone();
1416 drop(first_entry);
1417 drop(default_user);
1418
1419 let mut last_broadcast_state = Some("Running".to_string());
1423 tokio::spawn(async move {
1424 while let Some(msg) = workflow_rx.recv().await {
1425 {
1426 let mut session = session.lock().await;
1427 session.handle_workflow_message(msg.clone());
1428
1429 let state_str = match session.workflow_state {
1430 trustee_core::types::WorkflowState::Idle => "Idle",
1431 trustee_core::types::WorkflowState::Running => "Running",
1432 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
1433 };
1434 if last_broadcast_state.as_deref() != Some(state_str) {
1435 last_broadcast_state = Some(state_str.to_string());
1436 let state_msg = serde_json::json!({
1437 "type": "StateChanged",
1438 "state": state_str
1439 });
1440 let _ = ws_tx.send(state_msg.to_string());
1441 }
1442 }
1443
1444 let json =
1445 serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
1446 let _ = ws_tx.send(json);
1447 }
1448 tracing::debug!("Drain task ended for session: {}", session_id);
1449 });
1450 }
1451
1452 pub async fn resolve_user_key(&self, headers: &axum::http::HeaderMap) -> String {
1454 let Some(ref auth) = self.auth else {
1455 return "default".to_string();
1456 };
1457
1458 if let Some(token) = headers
1460 .get(axum::http::header::AUTHORIZATION)
1461 .and_then(|v| v.to_str().ok())
1462 .and_then(|v| v.strip_prefix("Bearer "))
1463 .map(|s| s.to_string())
1464 {
1465 if token.starts_with("dev:") {
1466 let parts: Vec<&str> = token.splitn(4, ':').collect();
1467 if parts.len() >= 4 {
1468 return format!("dev:{}", parts[1]);
1469 }
1470 }
1471 if let Ok(claims) = auth.validate_token(&token).await {
1472 return claims.sub;
1473 }
1474 }
1475
1476 let cookie_session_id = headers
1478 .get(axum::http::header::COOKIE)
1479 .and_then(|v| v.to_str().ok())
1480 .and_then(|cookies| {
1481 cookies
1482 .split(';')
1483 .map(|c| c.trim())
1484 .find_map(|c| {
1485 c.strip_prefix(&format!("{}=", auth.config.cookie_name))
1486 .map(|s| s.to_string())
1487 })
1488 });
1489
1490 if let Some(session_id) = cookie_session_id {
1491 if session_id.starts_with("dev:") {
1492 let parts: Vec<&str> = session_id.splitn(4, ':').collect();
1493 if parts.len() >= 4 {
1494 return format!("dev:{}", parts[1]);
1495 }
1496 }
1497
1498 if let Ok(access_token) = auth.session_manager.get_token(&session_id).await {
1499 if let Ok(claims) = auth.validate_token(&access_token).await {
1500 return claims.sub;
1501 }
1502 }
1503 }
1504
1505 "default".to_string()
1506 }
1507}
1508
1509struct SerializableMessage<'a>(&'a TuiMessage);
1515
1516impl<'a> serde::Serialize for SerializableMessage<'a> {
1517 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1518 where
1519 S: serde::Serializer,
1520 {
1521 use serde::ser::SerializeStruct;
1522
1523 match self.0 {
1524 TuiMessage::OutputLine(line) => {
1525 let mut s = serializer.serialize_struct("msg", 2)?;
1526 s.serialize_field("type", "OutputLine")?;
1527 s.serialize_field("line", line)?;
1528 s.end()
1529 }
1530 TuiMessage::StreamDelta(delta) => {
1531 let mut s = serializer.serialize_struct("msg", 2)?;
1532 s.serialize_field("type", "StreamDelta")?;
1533 s.serialize_field("delta", delta)?;
1534 s.end()
1535 }
1536 TuiMessage::ReasoningDelta(delta) => {
1537 let mut s = serializer.serialize_struct("msg", 2)?;
1538 s.serialize_field("type", "ReasoningDelta")?;
1539 s.serialize_field("delta", delta)?;
1540 s.end()
1541 }
1542 TuiMessage::WorkflowCompleted => {
1543 let mut s = serializer.serialize_struct("msg", 2)?;
1544 s.serialize_field("type", "WorkflowCompleted")?;
1545 s.serialize_field("state", "Idle")?;
1546 s.end()
1547 }
1548 TuiMessage::WorkflowError(err) => {
1549 let mut s = serializer.serialize_struct("msg", 2)?;
1550 s.serialize_field("type", "WorkflowError")?;
1551 s.serialize_field("error", err)?;
1552 s.end()
1553 }
1554 TuiMessage::ResumeInfo(info) => match info {
1555 Some(ri) => {
1556 let mut s = serializer.serialize_struct("msg", 5)?;
1557 s.serialize_field("type", "ResumeInfo")?;
1558 s.serialize_field("state", "Idle")?;
1559 s.serialize_field("session_id", &ri.session_id)?;
1560 s.serialize_field("checkpoint_id", &ri.checkpoint_id)?;
1561 s.serialize_field("iteration", &ri.iteration)?;
1562 s.end()
1563 }
1564 None => {
1565 let mut s = serializer.serialize_struct("msg", 2)?;
1566 s.serialize_field("type", "ResumeInfo")?;
1567 s.serialize_field("state", "Idle")?;
1568 s.end()
1569 }
1570 },
1571 TuiMessage::TodoUpdate(content) => {
1572 let mut s = serializer.serialize_struct("msg", 2)?;
1573 s.serialize_field("type", "TodoUpdate")?;
1574 s.serialize_field("content", content)?;
1575 s.end()
1576 }
1577 TuiMessage::WorkflowCancelled => {
1578 let mut s = serializer.serialize_struct("msg", 2)?;
1579 s.serialize_field("type", "WorkflowCancelled")?;
1580 s.serialize_field("state", "Idle")?;
1581 s.end()
1582 }
1583 TuiMessage::HandoffReady(briefing) => {
1584 let mut s = serializer.serialize_struct("msg", 3)?;
1585 s.serialize_field("type", "HandoffReady")?;
1586 s.serialize_field("state", "Idle")?;
1587 s.serialize_field("briefing", briefing)?;
1588 s.end()
1589 }
1590 TuiMessage::SessionRotated { old, new } => {
1591 let mut s = serializer.serialize_struct("msg", 3)?;
1592 s.serialize_field("type", "SessionRotated")?;
1593 s.serialize_field("old", old)?;
1594 s.serialize_field("new", new)?;
1595 s.end()
1596 }
1597 TuiMessage::HandoffFailed => {
1598 let mut s = serializer.serialize_struct("msg", 2)?;
1599 s.serialize_field("type", "HandoffFailed")?;
1600 s.serialize_field("state", "Idle")?;
1601 s.end()
1602 }
1603 TuiMessage::ToolPending {
1604 tool_name,
1605 hint,
1606 } => {
1607 let mut s = serializer.serialize_struct("msg", 3)?;
1608 s.serialize_field("type", "ToolPending")?;
1609 s.serialize_field("tool_name", tool_name)?;
1610 s.serialize_field("hint", hint)?;
1611 s.end()
1612 }
1613 TuiMessage::ToolDone {
1614 tool_name,
1615 success,
1616 hint,
1617 } => {
1618 let mut s = serializer.serialize_struct("msg", 4)?;
1619 s.serialize_field("type", "ToolDone")?;
1620 s.serialize_field("tool_name", tool_name)?;
1621 s.serialize_field("success", success)?;
1622 s.serialize_field("hint", hint)?;
1623 s.end()
1624 }
1625 TuiMessage::ContextTokensUpdated(count) => {
1626 let mut s = serializer.serialize_struct("msg", 2)?;
1627 s.serialize_field("type", "ContextTokensUpdated")?;
1628 s.serialize_field("count", count)?;
1629 s.end()
1630 }
1631 TuiMessage::McpServerStatus {
1632 name,
1633 connected,
1634 tool_count,
1635 error,
1636 } => {
1637 let mut s = serializer.serialize_struct("msg", 5)?;
1638 s.serialize_field("type", "McpServerStatus")?;
1639 s.serialize_field("name", name)?;
1640 s.serialize_field("connected", connected)?;
1641 s.serialize_field("tool_count", tool_count)?;
1642 s.serialize_field("error", error)?;
1643 s.end()
1644 }
1645 TuiMessage::SessionTitleUpdated(title) => {
1646 let mut s = serializer.serialize_struct("msg", 2)?;
1647 s.serialize_field("type", "SessionTitleUpdated")?;
1648 s.serialize_field("title", title)?;
1649 s.end()
1650 }
1651 }
1652 }
1653}
1654
1655fn deep_merge_toml(base: &mut toml::Value, overlay: &toml::Value) {
1666 match (base, overlay) {
1667 (toml::Value::Table(base_table), toml::Value::Table(overlay_table)) => {
1668 for (key, overlay_val) in overlay_table {
1669 match base_table.get_mut(key) {
1670 Some(base_val) => {
1671 deep_merge_toml(base_val, overlay_val);
1673 }
1674 None => {
1675 base_table.insert(key.clone(), overlay_val.clone());
1677 }
1678 }
1679 }
1680 }
1681 (base, overlay) => {
1683 *base = overlay.clone();
1684 }
1685 }
1686}
1687
1688fn substitute_env_vars(s: &mut String, secrets: &std::collections::HashMap<String, String>) {
1693 let mut result = String::with_capacity(s.len());
1695 let bytes = s.as_bytes();
1696 let mut i = 0;
1697
1698 while i < bytes.len() {
1699 if i + 1 < bytes.len() && bytes[i] == b'$' && bytes[i + 1] == b'{' {
1700 if let Some(end) = s[i + 2..].find('}') {
1702 let var_name = &s[i + 2..i + 2 + end];
1703 if let Some(value) = secrets.get(var_name) {
1705 result.push_str(value);
1706 } else if let Ok(value) = std::env::var(var_name) {
1707 result.push_str(&value);
1708 } else {
1709 result.push_str(&s[i..i + 2 + end + 1]);
1711 }
1712 i = i + 2 + end + 1;
1713 } else {
1714 result.push('$');
1716 i += 1;
1717 }
1718 } else {
1719 result.push(bytes[i] as char);
1720 i += 1;
1721 }
1722 }
1723
1724 *s = result;
1725}
1726
1727fn fingerprint_mcp_section(resolved: Option<&str>) -> u64 {
1737 use sha2::{Digest, Sha256};
1738 let section = resolved
1739 .and_then(|s| s.parse::<toml::Value>().ok())
1740 .and_then(|v| v.get("mcp").cloned());
1741 let bytes = match section {
1742 Some(v) => v.to_string().into_bytes(),
1743 None => b"<no-mcp>".to_vec(),
1744 };
1745 let digest = Sha256::digest(&bytes);
1746 u64::from_be_bytes(digest[..8].try_into().expect("sha256 digest >= 8 bytes"))
1747}
1748
1749#[cfg(test)]
1750mod tests {
1751 use super::*;
1752
1753 fn test_state() -> ServerState {
1755 let (session, _rx) = Session::new();
1756 let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
1757 ServerState::new(session, ws_tx, None)
1758 }
1759
1760 #[tokio::test]
1761 async fn dispatch_scope_pins_user_key_resolution() {
1762 let state = test_state();
1766 let caller = "some-human-sub";
1767 let plain_home = state.get_user_home_dir(caller).unwrap();
1768 let pinned_key = "prof-1234";
1769 let pinned_home = in_dispatch_scope(pinned_key, async {
1770 state.get_user_home_dir(caller).unwrap()
1771 })
1772 .await;
1773 let users_root = dirs::home_dir().unwrap().join(".trustee").join("users");
1774 assert_eq!(plain_home, users_root.join(trustee_core::user_hash(caller)));
1775 assert_eq!(
1776 pinned_home,
1777 users_root.join(trustee_core::user_hash(pinned_key)),
1778 "dispatch scope must pin resolution to the binding key"
1779 );
1780 assert_ne!(plain_home, pinned_home);
1781 }
1782
1783 fn temp_user_home(tag: &str) -> std::path::PathBuf {
1785 static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1786 let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1787 let dir = std::env::temp_dir().join(format!(
1788 "trustee-state-test-{}-{}-{}",
1789 tag,
1790 std::process::id(),
1791 n
1792 ));
1793 std::fs::create_dir_all(dir.join("config")).expect("create temp user home");
1794 dir
1795 }
1796
1797 fn parse(toml_str: &str) -> toml::Value {
1798 toml_str.parse::<toml::Value>().expect("valid test TOML")
1799 }
1800
1801 #[test]
1804 fn overlay_allowlist_drops_non_allowlisted_sections() {
1805 let state = test_state();
1806 let home = temp_user_home("allowlist");
1807 std::fs::write(
1808 home.join("config").join("trustee.toml"),
1809 "[server]\nport = 1\n\n[auth]\nmode = \"kanidm\"\n\n[mcp]\nmode = \"user\"\n",
1810 )
1811 .expect("write overlay");
1812
1813 let shared = "[server]\nport = 8080\n\n[mcp]\nmode = \"shared\"\n";
1814 let merged = state
1815 .merge_user_config(&home, Some(shared))
1816 .expect("overlay has allowlisted content");
1817
1818 let merged_val = parse(&merged);
1819 let expected = parse("[server]\nport = 8080\n\n[mcp]\nmode = \"user\"\n");
1820 assert_eq!(merged_val, expected, "merged config must be shared + [mcp] overlay only");
1821 assert!(merged_val.get("auth").is_none(), "overlay [auth] must be dropped");
1822 }
1823
1824 #[test]
1826 fn no_overlay_file_returns_none() {
1827 let state = test_state();
1828 let home = temp_user_home("empty");
1829 assert!(state
1830 .merge_user_config(&home, Some("[server]\nport = 8080\n"))
1831 .is_none());
1832 }
1833
1834 #[test]
1836 fn overlay_with_no_allowlisted_sections_is_noop() {
1837 let state = test_state();
1838 let home = temp_user_home("all-dropped");
1839 std::fs::write(
1840 home.join("config").join("trustee.toml"),
1841 "[server]\nport = 1\n\n[storage]\npath = \"/tmp/x\"\n",
1842 )
1843 .expect("write overlay");
1844 assert!(state
1845 .merge_user_config(&home, Some("[server]\nport = 8080\n"))
1846 .is_none());
1847 }
1848
1849 #[test]
1852 fn llm_overlay_dropped_by_default_and_kept_when_enabled() {
1853 let shared = "[llm]\nprovider = \"openai\"\n\n[mcp]\nmode = \"shared\"\n";
1854 let overlay = "[llm]\nprovider = \"anthropic\"\n";
1855
1856 let state = test_state();
1858 assert!(!state.allow_llm_overlay);
1859 let home = temp_user_home("llm-off");
1860 std::fs::write(home.join("config").join("trustee.toml"), overlay).expect("write overlay");
1861 assert!(state.merge_user_config(&home, Some(shared)).is_none());
1862
1863 let state = test_state().with_allow_llm_overlay(true);
1865 assert!(state.allow_llm_overlay);
1866 let home = temp_user_home("llm-on");
1867 std::fs::write(home.join("config").join("trustee.toml"), overlay).expect("write overlay");
1868 let merged = state
1869 .merge_user_config(&home, Some(shared))
1870 .expect("[llm] overlay applies when opted in");
1871 let merged_val = parse(&merged);
1872 assert_eq!(merged_val["llm"]["provider"].as_str(), Some("anthropic"));
1873 assert_eq!(merged_val["mcp"]["mode"].as_str(), Some("shared"));
1874 }
1875
1876 #[test]
1880 fn user_home_dir_uses_consolidated_hash() {
1881 let state = test_state();
1882 if let Some(home) = state.get_user_home_dir("farzan@example.com") {
1883 assert_eq!(
1884 home.file_name().and_then(|n| n.to_str()),
1885 Some(trustee_core::user_hash("farzan@example.com")).as_deref()
1886 );
1887 let users_root = dirs::home_dir().unwrap().join(".trustee").join("users");
1888 assert_eq!(home.parent(), Some(&users_root).map(|p| p.as_path()));
1889 }
1890 }
1892
1893 #[test]
1896 fn lifecycle_overlay_survives_allowlist_enabled_stripped() {
1897 let state = test_state();
1898 let home = temp_user_home("lifecycle-persona");
1899 std::fs::write(
1900 home.join("config").join("trustee.toml"),
1901 "[lifecycle]\nsystem_template = \"You are Saman.\"\nenabled = true\n",
1902 )
1903 .expect("write overlay");
1904
1905 let merged = state
1906 .merge_user_config(
1907 &home,
1908 Some("[lifecycle]\nsystem_template = \"shared default\"\n"),
1909 )
1910 .expect("lifecycle overlay is allowlisted");
1911
1912 let v = merged.parse::<toml::Value>().expect("valid TOML");
1913 assert_eq!(
1914 v["lifecycle"]["system_template"].as_str(),
1915 Some("You are Saman."),
1916 "agent persona wins over shared default"
1917 );
1918 assert_eq!(
1919 v["lifecycle"].get("enabled"),
1920 None,
1921 "WASM extension switch is not user-settable"
1922 );
1923 }
1924
1925 #[tokio::test]
1930 async fn agent_principals_get_isolated_session_buckets() {
1931 let state = test_state();
1932 let key_a = "agent-farzan";
1933 let key_b = "agent-paydar";
1934
1935 let (sid_a, session_a, _tx_a, _ts_a) = state.ensure_active_session(key_a).await;
1936 let (_sid_b, _session_b, _tx_b, _ts_b) = state.ensure_active_session(key_b).await;
1937
1938 assert!(
1940 state.get_session_by_any_id(key_a, &sid_a).await.is_some(),
1941 "owner bucket resolves its own session"
1942 );
1943 assert!(
1944 state.get_session_by_any_id(key_b, &sid_a).await.is_none(),
1945 "cross-agent session access must be 404/None"
1946 );
1947 assert!(Arc::strong_count(&session_a) >= 1);
1948
1949 if let Some(home) = state.get_user_home_dir(key_a) {
1951 assert_eq!(
1952 home.file_name().and_then(|n| n.to_str()),
1953 Some(trustee_core::user_hash(key_a)).as_deref()
1954 );
1955 }
1956 }
1957}
1958
1959#[cfg(test)]
1964mod mcp_loader_cache_tests {
1965 use super::*;
1966
1967 fn state_with_shared(shared: &str) -> ServerState {
1968 let (session, _rx) = Session::new();
1969 let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
1970 let mut state = ServerState::new(session, ws_tx, None);
1971 state.config_toml = Some(shared.to_string());
1972 state
1973 }
1974
1975 struct TempUser {
1979 key: String,
1980 home: std::path::PathBuf,
1981 }
1982
1983 impl TempUser {
1984 fn new(tag: &str) -> Self {
1985 let key = format!("16c-{tag}-{}@test.invalid", std::process::id());
1986 let home = dirs::home_dir()
1987 .expect("HOME available in test env")
1988 .join(".trustee")
1989 .join("users")
1990 .join(trustee_core::user_hash(&key));
1991 std::fs::create_dir_all(home.join("config")).expect("create user home");
1992 Self { key, home }
1993 }
1994
1995 fn write_overlay(&self, toml_str: &str) {
1996 std::fs::write(self.home.join("config").join("trustee.toml"), toml_str)
1997 .expect("write overlay");
1998 }
1999 }
2000
2001 impl Drop for TempUser {
2002 fn drop(&mut self) {
2003 let _ = std::fs::remove_dir_all(&self.home);
2004 }
2005 }
2006
2007 #[tokio::test]
2008 async fn no_mcp_config_caches_disabled_marker() {
2009 let state = state_with_shared("[server]\nport = 8080\n");
2010 let user = TempUser::new("nomcp");
2011 let ts = Arc::new(pep::MemoryTokenStore::new());
2012
2013 let first = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2014 assert!(first.is_none(), "no [mcp] anywhere → disabled marker");
2015 assert_eq!(state.mcp_loaders.len(), 1, "exactly one cache entry");
2016
2017 let second = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2019 assert!(second.is_none());
2020 assert_eq!(state.mcp_loaders.len(), 1);
2021 }
2022
2023 #[tokio::test]
2024 async fn concurrent_cold_builds_single_flight() {
2025 let state = Arc::new(state_with_shared("[server]\nport = 8080\n"));
2026 let user = TempUser::new("singleflight");
2027 let ts = Arc::new(pep::MemoryTokenStore::new());
2028
2029 let mut handles = Vec::new();
2030 for _ in 0..5 {
2031 let state = state.clone();
2032 let key = user.key.clone();
2033 let ts = ts.clone();
2034 handles.push(tokio::spawn(async move {
2035 state.get_or_build_mcp_loader(&key, &ts).await
2036 }));
2037 }
2038 for h in handles {
2039 h.await.unwrap().expect("all five succeed");
2040 }
2041 assert_eq!(state.mcp_loaders.len(), 1, "single-flight → one entry");
2042 }
2043
2044 #[tokio::test]
2045 async fn fingerprint_change_triggers_rebuild() {
2046 let state = state_with_shared("[server]\nport = 8080\n");
2047 let user = TempUser::new("fpchange");
2048 let ts = Arc::new(pep::MemoryTokenStore::new());
2049
2050 user.write_overlay("[mcp]\nenabled = true\n\n[[mcp.servers]]\nname = \"v1\"\nurl = \"http://127.0.0.1:9/sse\"\n");
2053 let v1 = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2054 assert!(v1.is_some(), "enabled [mcp] → real loader");
2055 let fp1 = state
2056 .mcp_loaders
2057 .get(&trustee_core::user_hash(&user.key))
2058 .unwrap()
2059 .fingerprint;
2060
2061 std::thread::sleep(std::time::Duration::from_millis(5));
2063 user.write_overlay("[mcp]\nenabled = true\n\n[[mcp.servers]]\nname = \"v2\"\nurl = \"http://127.0.0.1:9/other\"\n");
2064 let v2 = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2065 assert!(v2.is_some());
2066 let entry = state
2067 .mcp_loaders
2068 .get(&trustee_core::user_hash(&user.key))
2069 .unwrap();
2070 assert_ne!(
2071 entry.fingerprint, fp1,
2072 "fingerprint must change with content"
2073 );
2074 assert!(entry.degraded.is_none());
2075
2076 let _still_usable = v1.as_ref().unwrap().tool_count;
2078 }
2079
2080 #[tokio::test]
2081 async fn degraded_entry_fails_loud_within_backoff_and_isolates_users() {
2082 let state = state_with_shared("[server]\nport = 8080\n");
2083 let bad = TempUser::new("degraded-bad");
2084 let good = TempUser::new("degraded-good");
2085 let ts = Arc::new(pep::MemoryTokenStore::new());
2086
2087 bad.write_overlay("[mcp]\nenabled = \"not-a-bool\"\n");
2089 let err = match state.get_or_build_mcp_loader(&bad.key, &ts).await {
2090 Ok(_) => panic!("invalid [mcp] must fail loud"),
2091 Err(e) => e,
2092 };
2093 assert!(
2094 err.contains("invalid [mcp]"),
2095 "surfaces the parse error: {err}"
2096 );
2097
2098 let entry = state
2099 .mcp_loaders
2100 .get(&trustee_core::user_hash(&bad.key))
2101 .unwrap();
2102 assert!(entry.degraded.is_some(), "poison entry recorded");
2103
2104 let err2 = match state.get_or_build_mcp_loader(&bad.key, &ts).await {
2106 Ok(_) => panic!("still within backoff"),
2107 Err(e) => e,
2108 };
2109 assert_eq!(err, err2, "same cached error");
2110
2111 let good_loader = state
2113 .get_or_build_mcp_loader(&good.key, &ts)
2114 .await
2115 .expect("other user unaffected");
2116 assert!(
2117 good_loader.is_none(),
2118 "good user has no [mcp] → disabled marker"
2119 );
2120 }
2121}
2122
2123#[cfg(test)]
2124mod restore_tests {
2125 use super::*;
2126
2127 fn fresh_state() -> ServerState {
2128 let (session, _rx) = Session::new();
2129 let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
2130 let mut state = ServerState::new(session, ws_tx, None);
2131 state.config_toml = Some("[agent]\nname = \"trustee\"\n".to_string());
2132 state
2133 }
2134
2135 struct TempUser {
2137 key: String,
2138 }
2139
2140 impl TempUser {
2141 fn new(tag: &str) -> Self {
2142 let key = format!("restore-{tag}-{}@test.invalid", std::process::id());
2143 let home = dirs::home_dir()
2144 .expect("HOME available in test env")
2145 .join(".trustee")
2146 .join("users")
2147 .join(trustee_core::user_hash(&key));
2148 std::fs::create_dir_all(home).expect("create user home");
2149 Self { key }
2150 }
2151 }
2152
2153 #[tokio::test]
2154 async fn restore_unknown_chain_returns_none() {
2155 let state = fresh_state();
2156 let user = TempUser::new("unknown");
2157 let out = state
2158 .get_or_restore_session_by_any_id(&user.key, "session_2099_01_01_00_00_deadbeef")
2159 .await;
2160 assert!(out.is_none(), "unknown chain must 404, not restore");
2161 }
2162
2163 #[tokio::test]
2164 async fn destroy_tombstones_chain_id() {
2165 let state = fresh_state();
2166 let user = TempUser::new("tombstone");
2167 let sid = state
2168 .create_session(&user.key, None, None, true)
2169 .await
2170 .expect("create session");
2171
2172 let (session, _ws) = state.get_session(&user.key, &sid).await.unwrap();
2174 {
2175 let mut s = session.lock().await;
2176 s.session_id = Some("session_2026_09_04_10_00_aaaabbbb".to_string());
2177 }
2178
2179 state.destroy_session(&user.key, &sid).await.expect("destroy");
2180
2181 let path = state
2182 .tombstone_path(&user.key, "session_2026_09_04_10_00_aaaabbbb")
2183 .expect("tombstone path");
2184 assert!(path.exists(), "destroy must write a chain tombstone");
2185
2186 let sid2 = state
2188 .create_session(&user.key, None, None, true)
2189 .await
2190 .expect("create second session");
2191 state.destroy_session(&user.key, &sid2).await.expect("destroy 2");
2192 }
2194
2195 #[tokio::test]
2196 async fn restore_refuses_registry_key_shapes() {
2197 let state = fresh_state();
2201 let user = TempUser::new("shapes");
2202 assert!(state.tombstone_path(&user.key, "default").is_none());
2203 let out = state.get_or_restore_session_by_any_id(&user.key, "default").await;
2204 assert!(out.is_none());
2205 }
2206}