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 let mut leftovers: Vec<String> = Vec::new();
1309 collect_unresolved_placeholders(section, &mut leftovers);
1310 if !leftovers.is_empty() {
1311 return Err(format!(
1312 "unresolved {} placeholder(s) in [mcp]: {} — supply the \
1313 exact-case key in the profile secrets",
1314 leftovers.len(),
1315 leftovers.join(", ")
1316 ));
1317 }
1318 use serde::Deserialize as _;
1319 Some(
1320 abk::config::McpConfig::deserialize(section.clone())
1321 .map_err(|e| format!("invalid [mcp] config: {}", e))?,
1322 )
1323 }
1324 None => None,
1325 }
1326 }
1327 None => None,
1328 };
1329
1330 let loader = match mcp_config {
1331 Some(cfg) if cfg.enabled => {
1332 let built = abk::agent::McpToolLoader::with_token_store(
1333 &cfg,
1334 Some(token_store.clone() as std::sync::Arc<dyn pep::token_store::TokenStore>),
1335 )
1336 .await
1337 .map_err(|e| format!("MCP loader build failed: {}", e))?;
1338
1339 let servers: Vec<String> = built
1343 .server_statuses
1344 .iter()
1345 .map(|s| {
1346 if s.connected {
1347 format!("{}(up,{}tools)", s.name, s.tool_count)
1348 } else {
1349 format!("{}(DOWN)", s.name)
1350 }
1351 })
1352 .collect();
1353 tracing::info!(
1354 "MCP loader built for user {}: servers=[{}] total_tools={}",
1355 &user_hash[..8.min(user_hash.len())],
1356 servers.join(", "),
1357 built.tool_count
1358 );
1359 Some(std::sync::Arc::new(built))
1360 }
1361 _ => None,
1362 };
1363
1364 Ok(McpLoaderEntry {
1365 loader,
1366 fingerprint,
1367 built_at: chrono::Utc::now(),
1368 degraded: None,
1369 failed_at: None,
1370 })
1371 }
1372
1373 fn spawn_user_drain_task(
1375 &self,
1376 session_id: String,
1377 session: Arc<Mutex<Session>>,
1378 ws_tx: broadcast::Sender<String>,
1379 mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>,
1380 ) {
1381 let mut last_broadcast_state: Option<String> = None;
1387 tokio::spawn(async move {
1388 while let Some(msg) = workflow_rx.recv().await {
1389 {
1390 let mut session = session.lock().await;
1391 session.handle_workflow_message(msg.clone());
1392
1393 let state_str = match session.workflow_state {
1394 trustee_core::types::WorkflowState::Idle => "Idle",
1395 trustee_core::types::WorkflowState::Running => "Running",
1396 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
1397 };
1398 if last_broadcast_state.as_deref() != Some(state_str) {
1399 last_broadcast_state = Some(state_str.to_string());
1400 let state_msg = serde_json::json!({
1401 "type": "StateChanged",
1402 "state": state_str
1403 });
1404 let _ = ws_tx.send(state_msg.to_string());
1405 }
1406 }
1407
1408 let json =
1409 serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
1410 let _ = ws_tx.send(json);
1411 }
1412 tracing::debug!("Drain task ended for session: {}", session_id);
1413 });
1414 }
1415
1416 pub fn spawn_drain_task(self, mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>) {
1419 let default_user = self
1421 .sessions
1422 .get("default")
1423 .expect("default user must exist");
1424 let first_entry = default_user
1425 .sessions
1426 .iter()
1427 .next()
1428 .expect("default user must have at least one session");
1429 let session = first_entry.session.clone();
1430 let ws_tx = first_entry.ws_tx.clone();
1431 let session_id = first_entry.key().clone();
1432 drop(first_entry);
1433 drop(default_user);
1434
1435 let mut last_broadcast_state = Some("Running".to_string());
1439 tokio::spawn(async move {
1440 while let Some(msg) = workflow_rx.recv().await {
1441 {
1442 let mut session = session.lock().await;
1443 session.handle_workflow_message(msg.clone());
1444
1445 let state_str = match session.workflow_state {
1446 trustee_core::types::WorkflowState::Idle => "Idle",
1447 trustee_core::types::WorkflowState::Running => "Running",
1448 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
1449 };
1450 if last_broadcast_state.as_deref() != Some(state_str) {
1451 last_broadcast_state = Some(state_str.to_string());
1452 let state_msg = serde_json::json!({
1453 "type": "StateChanged",
1454 "state": state_str
1455 });
1456 let _ = ws_tx.send(state_msg.to_string());
1457 }
1458 }
1459
1460 let json =
1461 serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
1462 let _ = ws_tx.send(json);
1463 }
1464 tracing::debug!("Drain task ended for session: {}", session_id);
1465 });
1466 }
1467
1468 pub async fn resolve_user_key(&self, headers: &axum::http::HeaderMap) -> String {
1470 let Some(ref auth) = self.auth else {
1471 return "default".to_string();
1472 };
1473
1474 if let Some(token) = headers
1476 .get(axum::http::header::AUTHORIZATION)
1477 .and_then(|v| v.to_str().ok())
1478 .and_then(|v| v.strip_prefix("Bearer "))
1479 .map(|s| s.to_string())
1480 {
1481 if token.starts_with("dev:") {
1482 let parts: Vec<&str> = token.splitn(4, ':').collect();
1483 if parts.len() >= 4 {
1484 return format!("dev:{}", parts[1]);
1485 }
1486 }
1487 if let Ok(claims) = auth.validate_token(&token).await {
1488 return claims.sub;
1489 }
1490 }
1491
1492 let cookie_session_id = headers
1494 .get(axum::http::header::COOKIE)
1495 .and_then(|v| v.to_str().ok())
1496 .and_then(|cookies| {
1497 cookies
1498 .split(';')
1499 .map(|c| c.trim())
1500 .find_map(|c| {
1501 c.strip_prefix(&format!("{}=", auth.config.cookie_name))
1502 .map(|s| s.to_string())
1503 })
1504 });
1505
1506 if let Some(session_id) = cookie_session_id {
1507 if session_id.starts_with("dev:") {
1508 let parts: Vec<&str> = session_id.splitn(4, ':').collect();
1509 if parts.len() >= 4 {
1510 return format!("dev:{}", parts[1]);
1511 }
1512 }
1513
1514 if let Ok(access_token) = auth.session_manager.get_token(&session_id).await {
1515 if let Ok(claims) = auth.validate_token(&access_token).await {
1516 return claims.sub;
1517 }
1518 }
1519 }
1520
1521 "default".to_string()
1522 }
1523}
1524
1525struct SerializableMessage<'a>(&'a TuiMessage);
1531
1532impl<'a> serde::Serialize for SerializableMessage<'a> {
1533 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
1534 where
1535 S: serde::Serializer,
1536 {
1537 use serde::ser::SerializeStruct;
1538
1539 match self.0 {
1540 TuiMessage::OutputLine(line) => {
1541 let mut s = serializer.serialize_struct("msg", 2)?;
1542 s.serialize_field("type", "OutputLine")?;
1543 s.serialize_field("line", line)?;
1544 s.end()
1545 }
1546 TuiMessage::StreamDelta(delta) => {
1547 let mut s = serializer.serialize_struct("msg", 2)?;
1548 s.serialize_field("type", "StreamDelta")?;
1549 s.serialize_field("delta", delta)?;
1550 s.end()
1551 }
1552 TuiMessage::ReasoningDelta(delta) => {
1553 let mut s = serializer.serialize_struct("msg", 2)?;
1554 s.serialize_field("type", "ReasoningDelta")?;
1555 s.serialize_field("delta", delta)?;
1556 s.end()
1557 }
1558 TuiMessage::WorkflowCompleted => {
1559 let mut s = serializer.serialize_struct("msg", 2)?;
1560 s.serialize_field("type", "WorkflowCompleted")?;
1561 s.serialize_field("state", "Idle")?;
1562 s.end()
1563 }
1564 TuiMessage::WorkflowError(err) => {
1565 let mut s = serializer.serialize_struct("msg", 2)?;
1566 s.serialize_field("type", "WorkflowError")?;
1567 s.serialize_field("error", err)?;
1568 s.end()
1569 }
1570 TuiMessage::ResumeInfo(info) => match info {
1571 Some(ri) => {
1572 let mut s = serializer.serialize_struct("msg", 5)?;
1573 s.serialize_field("type", "ResumeInfo")?;
1574 s.serialize_field("state", "Idle")?;
1575 s.serialize_field("session_id", &ri.session_id)?;
1576 s.serialize_field("checkpoint_id", &ri.checkpoint_id)?;
1577 s.serialize_field("iteration", &ri.iteration)?;
1578 s.end()
1579 }
1580 None => {
1581 let mut s = serializer.serialize_struct("msg", 2)?;
1582 s.serialize_field("type", "ResumeInfo")?;
1583 s.serialize_field("state", "Idle")?;
1584 s.end()
1585 }
1586 },
1587 TuiMessage::TodoUpdate(content) => {
1588 let mut s = serializer.serialize_struct("msg", 2)?;
1589 s.serialize_field("type", "TodoUpdate")?;
1590 s.serialize_field("content", content)?;
1591 s.end()
1592 }
1593 TuiMessage::WorkflowCancelled => {
1594 let mut s = serializer.serialize_struct("msg", 2)?;
1595 s.serialize_field("type", "WorkflowCancelled")?;
1596 s.serialize_field("state", "Idle")?;
1597 s.end()
1598 }
1599 TuiMessage::HandoffReady(briefing) => {
1600 let mut s = serializer.serialize_struct("msg", 3)?;
1601 s.serialize_field("type", "HandoffReady")?;
1602 s.serialize_field("state", "Idle")?;
1603 s.serialize_field("briefing", briefing)?;
1604 s.end()
1605 }
1606 TuiMessage::SessionRotated { old, new } => {
1607 let mut s = serializer.serialize_struct("msg", 3)?;
1608 s.serialize_field("type", "SessionRotated")?;
1609 s.serialize_field("old", old)?;
1610 s.serialize_field("new", new)?;
1611 s.end()
1612 }
1613 TuiMessage::HandoffFailed => {
1614 let mut s = serializer.serialize_struct("msg", 2)?;
1615 s.serialize_field("type", "HandoffFailed")?;
1616 s.serialize_field("state", "Idle")?;
1617 s.end()
1618 }
1619 TuiMessage::ToolPending {
1620 tool_name,
1621 hint,
1622 } => {
1623 let mut s = serializer.serialize_struct("msg", 3)?;
1624 s.serialize_field("type", "ToolPending")?;
1625 s.serialize_field("tool_name", tool_name)?;
1626 s.serialize_field("hint", hint)?;
1627 s.end()
1628 }
1629 TuiMessage::ToolDone {
1630 tool_name,
1631 success,
1632 hint,
1633 } => {
1634 let mut s = serializer.serialize_struct("msg", 4)?;
1635 s.serialize_field("type", "ToolDone")?;
1636 s.serialize_field("tool_name", tool_name)?;
1637 s.serialize_field("success", success)?;
1638 s.serialize_field("hint", hint)?;
1639 s.end()
1640 }
1641 TuiMessage::ContextTokensUpdated(count) => {
1642 let mut s = serializer.serialize_struct("msg", 2)?;
1643 s.serialize_field("type", "ContextTokensUpdated")?;
1644 s.serialize_field("count", count)?;
1645 s.end()
1646 }
1647 TuiMessage::McpServerStatus {
1648 name,
1649 connected,
1650 tool_count,
1651 error,
1652 } => {
1653 let mut s = serializer.serialize_struct("msg", 5)?;
1654 s.serialize_field("type", "McpServerStatus")?;
1655 s.serialize_field("name", name)?;
1656 s.serialize_field("connected", connected)?;
1657 s.serialize_field("tool_count", tool_count)?;
1658 s.serialize_field("error", error)?;
1659 s.end()
1660 }
1661 TuiMessage::SessionTitleUpdated(title) => {
1662 let mut s = serializer.serialize_struct("msg", 2)?;
1663 s.serialize_field("type", "SessionTitleUpdated")?;
1664 s.serialize_field("title", title)?;
1665 s.end()
1666 }
1667 }
1668 }
1669}
1670
1671fn deep_merge_toml(base: &mut toml::Value, overlay: &toml::Value) {
1682 match (base, overlay) {
1683 (toml::Value::Table(base_table), toml::Value::Table(overlay_table)) => {
1684 for (key, overlay_val) in overlay_table {
1685 match base_table.get_mut(key) {
1686 Some(base_val) => {
1687 deep_merge_toml(base_val, overlay_val);
1689 }
1690 None => {
1691 base_table.insert(key.clone(), overlay_val.clone());
1693 }
1694 }
1695 }
1696 }
1697 (base, overlay) => {
1699 *base = overlay.clone();
1700 }
1701 }
1702}
1703
1704fn collect_unresolved_placeholders(value: &toml::Value, out: &mut Vec<String>) {
1711 match value {
1712 toml::Value::String(s) => {
1713 if let Some(start) = s.find("${") {
1714 if let Some(end_rel) = s[start + 2..].find('}') {
1715 let name = &s[start + 2..start + 2 + end_rel];
1716 let rendered = format!("${{{name}}}");
1717 if !out.contains(&rendered) {
1718 out.push(rendered);
1719 }
1720 }
1721 }
1722 }
1723 toml::Value::Table(t) => {
1724 for v in t.values() {
1725 collect_unresolved_placeholders(v, out);
1726 }
1727 }
1728 toml::Value::Array(a) => {
1729 for v in a {
1730 collect_unresolved_placeholders(v, out);
1731 }
1732 }
1733 _ => {}
1734 }
1735}
1736
1737fn substitute_env_vars(s: &mut String, secrets: &std::collections::HashMap<String, String>) {
1738 let mut result = String::with_capacity(s.len());
1740 let bytes = s.as_bytes();
1741 let mut i = 0;
1742 let mut unresolved: Vec<&str> = Vec::new();
1743
1744 while i < bytes.len() {
1745 if i + 1 < bytes.len() && bytes[i] == b'$' && bytes[i + 1] == b'{' {
1746 if let Some(end) = s[i + 2..].find('}') {
1748 let var_name = &s[i + 2..i + 2 + end];
1749 if let Some(value) = secrets.get(var_name) {
1751 result.push_str(value);
1752 } else if let Ok(value) = std::env::var(var_name) {
1753 result.push_str(&value);
1754 } else {
1755 result.push_str(&s[i..i + 2 + end + 1]);
1760 if !unresolved.contains(&var_name) {
1761 unresolved.push(var_name);
1762 }
1763 }
1764 i = i + 2 + end + 1;
1765 } else {
1766 result.push('$');
1768 i += 1;
1769 }
1770 } else {
1771 result.push(bytes[i] as char);
1772 i += 1;
1773 }
1774 }
1775
1776 if !unresolved.is_empty() {
1777 tracing::warn!(
1778 "ENV SUBSTITUTION LEFT {} UNRESOLVED: {} — the literal ${{…}} remains in the \
1779 effective config; supply the key in the profile secrets (exact-case KEY=VALUE) \
1780 or the environment",
1781 unresolved.len(),
1782 unresolved
1783 .iter()
1784 .map(|v| format!("${{{v}}}"))
1785 .collect::<Vec<_>>()
1786 .join(", ")
1787 );
1788 }
1789
1790 *s = result;
1791}
1792
1793fn fingerprint_mcp_section(resolved: Option<&str>) -> u64 {
1803 use sha2::{Digest, Sha256};
1804 let section = resolved
1805 .and_then(|s| s.parse::<toml::Value>().ok())
1806 .and_then(|v| v.get("mcp").cloned());
1807 let bytes = match section {
1808 Some(v) => v.to_string().into_bytes(),
1809 None => b"<no-mcp>".to_vec(),
1810 };
1811 let digest = Sha256::digest(&bytes);
1812 u64::from_be_bytes(digest[..8].try_into().expect("sha256 digest >= 8 bytes"))
1813}
1814
1815#[cfg(test)]
1816mod tests {
1817 use super::*;
1818
1819 fn test_state() -> ServerState {
1821 let (session, _rx) = Session::new();
1822 let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
1823 ServerState::new(session, ws_tx, None)
1824 }
1825
1826 #[tokio::test]
1827 async fn dispatch_scope_pins_user_key_resolution() {
1828 let state = test_state();
1832 let caller = "some-human-sub";
1833 let plain_home = state.get_user_home_dir(caller).unwrap();
1834 let pinned_key = "prof-1234";
1835 let pinned_home = in_dispatch_scope(pinned_key, async {
1836 state.get_user_home_dir(caller).unwrap()
1837 })
1838 .await;
1839 let users_root = dirs::home_dir().unwrap().join(".trustee").join("users");
1840 assert_eq!(plain_home, users_root.join(trustee_core::user_hash(caller)));
1841 assert_eq!(
1842 pinned_home,
1843 users_root.join(trustee_core::user_hash(pinned_key)),
1844 "dispatch scope must pin resolution to the binding key"
1845 );
1846 assert_ne!(plain_home, pinned_home);
1847 }
1848
1849 fn temp_user_home(tag: &str) -> std::path::PathBuf {
1851 static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
1852 let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1853 let dir = std::env::temp_dir().join(format!(
1854 "trustee-state-test-{}-{}-{}",
1855 tag,
1856 std::process::id(),
1857 n
1858 ));
1859 std::fs::create_dir_all(dir.join("config")).expect("create temp user home");
1860 dir
1861 }
1862
1863 fn parse(toml_str: &str) -> toml::Value {
1864 toml_str.parse::<toml::Value>().expect("valid test TOML")
1865 }
1866
1867 #[test]
1870 fn overlay_allowlist_drops_non_allowlisted_sections() {
1871 let state = test_state();
1872 let home = temp_user_home("allowlist");
1873 std::fs::write(
1874 home.join("config").join("trustee.toml"),
1875 "[server]\nport = 1\n\n[auth]\nmode = \"kanidm\"\n\n[mcp]\nmode = \"user\"\n",
1876 )
1877 .expect("write overlay");
1878
1879 let shared = "[server]\nport = 8080\n\n[mcp]\nmode = \"shared\"\n";
1880 let merged = state
1881 .merge_user_config(&home, Some(shared))
1882 .expect("overlay has allowlisted content");
1883
1884 let merged_val = parse(&merged);
1885 let expected = parse("[server]\nport = 8080\n\n[mcp]\nmode = \"user\"\n");
1886 assert_eq!(merged_val, expected, "merged config must be shared + [mcp] overlay only");
1887 assert!(merged_val.get("auth").is_none(), "overlay [auth] must be dropped");
1888 }
1889
1890 #[test]
1892 fn no_overlay_file_returns_none() {
1893 let state = test_state();
1894 let home = temp_user_home("empty");
1895 assert!(state
1896 .merge_user_config(&home, Some("[server]\nport = 8080\n"))
1897 .is_none());
1898 }
1899
1900 #[test]
1902 fn overlay_with_no_allowlisted_sections_is_noop() {
1903 let state = test_state();
1904 let home = temp_user_home("all-dropped");
1905 std::fs::write(
1906 home.join("config").join("trustee.toml"),
1907 "[server]\nport = 1\n\n[storage]\npath = \"/tmp/x\"\n",
1908 )
1909 .expect("write overlay");
1910 assert!(state
1911 .merge_user_config(&home, Some("[server]\nport = 8080\n"))
1912 .is_none());
1913 }
1914
1915 #[test]
1918 fn llm_overlay_dropped_by_default_and_kept_when_enabled() {
1919 let shared = "[llm]\nprovider = \"openai\"\n\n[mcp]\nmode = \"shared\"\n";
1920 let overlay = "[llm]\nprovider = \"anthropic\"\n";
1921
1922 let state = test_state();
1924 assert!(!state.allow_llm_overlay);
1925 let home = temp_user_home("llm-off");
1926 std::fs::write(home.join("config").join("trustee.toml"), overlay).expect("write overlay");
1927 assert!(state.merge_user_config(&home, Some(shared)).is_none());
1928
1929 let state = test_state().with_allow_llm_overlay(true);
1931 assert!(state.allow_llm_overlay);
1932 let home = temp_user_home("llm-on");
1933 std::fs::write(home.join("config").join("trustee.toml"), overlay).expect("write overlay");
1934 let merged = state
1935 .merge_user_config(&home, Some(shared))
1936 .expect("[llm] overlay applies when opted in");
1937 let merged_val = parse(&merged);
1938 assert_eq!(merged_val["llm"]["provider"].as_str(), Some("anthropic"));
1939 assert_eq!(merged_val["mcp"]["mode"].as_str(), Some("shared"));
1940 }
1941
1942 #[test]
1946 fn user_home_dir_uses_consolidated_hash() {
1947 let state = test_state();
1948 if let Some(home) = state.get_user_home_dir("farzan@example.com") {
1949 assert_eq!(
1950 home.file_name().and_then(|n| n.to_str()),
1951 Some(trustee_core::user_hash("farzan@example.com")).as_deref()
1952 );
1953 let users_root = dirs::home_dir().unwrap().join(".trustee").join("users");
1954 assert_eq!(home.parent(), Some(&users_root).map(|p| p.as_path()));
1955 }
1956 }
1958
1959 #[test]
1962 fn lifecycle_overlay_survives_allowlist_enabled_stripped() {
1963 let state = test_state();
1964 let home = temp_user_home("lifecycle-persona");
1965 std::fs::write(
1966 home.join("config").join("trustee.toml"),
1967 "[lifecycle]\nsystem_template = \"You are Saman.\"\nenabled = true\n",
1968 )
1969 .expect("write overlay");
1970
1971 let merged = state
1972 .merge_user_config(
1973 &home,
1974 Some("[lifecycle]\nsystem_template = \"shared default\"\n"),
1975 )
1976 .expect("lifecycle overlay is allowlisted");
1977
1978 let v = merged.parse::<toml::Value>().expect("valid TOML");
1979 assert_eq!(
1980 v["lifecycle"]["system_template"].as_str(),
1981 Some("You are Saman."),
1982 "agent persona wins over shared default"
1983 );
1984 assert_eq!(
1985 v["lifecycle"].get("enabled"),
1986 None,
1987 "WASM extension switch is not user-settable"
1988 );
1989 }
1990
1991 #[tokio::test]
1996 async fn agent_principals_get_isolated_session_buckets() {
1997 let state = test_state();
1998 let key_a = "agent-farzan";
1999 let key_b = "agent-paydar";
2000
2001 let (sid_a, session_a, _tx_a, _ts_a) = state.ensure_active_session(key_a).await;
2002 let (_sid_b, _session_b, _tx_b, _ts_b) = state.ensure_active_session(key_b).await;
2003
2004 assert!(
2006 state.get_session_by_any_id(key_a, &sid_a).await.is_some(),
2007 "owner bucket resolves its own session"
2008 );
2009 assert!(
2010 state.get_session_by_any_id(key_b, &sid_a).await.is_none(),
2011 "cross-agent session access must be 404/None"
2012 );
2013 assert!(Arc::strong_count(&session_a) >= 1);
2014
2015 if let Some(home) = state.get_user_home_dir(key_a) {
2017 assert_eq!(
2018 home.file_name().and_then(|n| n.to_str()),
2019 Some(trustee_core::user_hash(key_a)).as_deref()
2020 );
2021 }
2022 }
2023}
2024
2025#[cfg(test)]
2030mod mcp_loader_cache_tests {
2031 use super::*;
2032
2033 fn state_with_shared(shared: &str) -> ServerState {
2034 let (session, _rx) = Session::new();
2035 let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
2036 let mut state = ServerState::new(session, ws_tx, None);
2037 state.config_toml = Some(shared.to_string());
2038 state
2039 }
2040
2041 struct TempUser {
2045 key: String,
2046 home: std::path::PathBuf,
2047 }
2048
2049 impl TempUser {
2050 fn new(tag: &str) -> Self {
2051 let key = format!("16c-{tag}-{}@test.invalid", std::process::id());
2052 let home = dirs::home_dir()
2053 .expect("HOME available in test env")
2054 .join(".trustee")
2055 .join("users")
2056 .join(trustee_core::user_hash(&key));
2057 std::fs::create_dir_all(home.join("config")).expect("create user home");
2058 Self { key, home }
2059 }
2060
2061 fn write_overlay(&self, toml_str: &str) {
2062 std::fs::write(self.home.join("config").join("trustee.toml"), toml_str)
2063 .expect("write overlay");
2064 }
2065 }
2066
2067 impl Drop for TempUser {
2068 fn drop(&mut self) {
2069 let _ = std::fs::remove_dir_all(&self.home);
2070 }
2071 }
2072
2073 #[tokio::test]
2074 async fn no_mcp_config_caches_disabled_marker() {
2075 let state = state_with_shared("[server]\nport = 8080\n");
2076 let user = TempUser::new("nomcp");
2077 let ts = Arc::new(pep::MemoryTokenStore::new());
2078
2079 let first = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2080 assert!(first.is_none(), "no [mcp] anywhere → disabled marker");
2081 assert_eq!(state.mcp_loaders.len(), 1, "exactly one cache entry");
2082
2083 let second = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2085 assert!(second.is_none());
2086 assert_eq!(state.mcp_loaders.len(), 1);
2087 }
2088
2089 #[tokio::test]
2090 async fn concurrent_cold_builds_single_flight() {
2091 let state = Arc::new(state_with_shared("[server]\nport = 8080\n"));
2092 let user = TempUser::new("singleflight");
2093 let ts = Arc::new(pep::MemoryTokenStore::new());
2094
2095 let mut handles = Vec::new();
2096 for _ in 0..5 {
2097 let state = state.clone();
2098 let key = user.key.clone();
2099 let ts = ts.clone();
2100 handles.push(tokio::spawn(async move {
2101 state.get_or_build_mcp_loader(&key, &ts).await
2102 }));
2103 }
2104 for h in handles {
2105 h.await.unwrap().expect("all five succeed");
2106 }
2107 assert_eq!(state.mcp_loaders.len(), 1, "single-flight → one entry");
2108 }
2109
2110 #[tokio::test]
2111 async fn fingerprint_change_triggers_rebuild() {
2112 let state = state_with_shared("[server]\nport = 8080\n");
2113 let user = TempUser::new("fpchange");
2114 let ts = Arc::new(pep::MemoryTokenStore::new());
2115
2116 user.write_overlay("[mcp]\nenabled = true\n\n[[mcp.servers]]\nname = \"v1\"\nurl = \"http://127.0.0.1:9/sse\"\n");
2119 let v1 = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2120 assert!(v1.is_some(), "enabled [mcp] → real loader");
2121 let fp1 = state
2122 .mcp_loaders
2123 .get(&trustee_core::user_hash(&user.key))
2124 .unwrap()
2125 .fingerprint;
2126
2127 std::thread::sleep(std::time::Duration::from_millis(5));
2129 user.write_overlay("[mcp]\nenabled = true\n\n[[mcp.servers]]\nname = \"v2\"\nurl = \"http://127.0.0.1:9/other\"\n");
2130 let v2 = state.get_or_build_mcp_loader(&user.key, &ts).await.unwrap();
2131 assert!(v2.is_some());
2132 let entry = state
2133 .mcp_loaders
2134 .get(&trustee_core::user_hash(&user.key))
2135 .unwrap();
2136 assert_ne!(
2137 entry.fingerprint, fp1,
2138 "fingerprint must change with content"
2139 );
2140 assert!(entry.degraded.is_none());
2141
2142 let _still_usable = v1.as_ref().unwrap().tool_count;
2144 }
2145
2146 #[tokio::test]
2147 async fn degraded_entry_fails_loud_within_backoff_and_isolates_users() {
2148 let state = state_with_shared("[server]\nport = 8080\n");
2149 let bad = TempUser::new("degraded-bad");
2150 let good = TempUser::new("degraded-good");
2151 let ts = Arc::new(pep::MemoryTokenStore::new());
2152
2153 bad.write_overlay("[mcp]\nenabled = \"not-a-bool\"\n");
2155 let err = match state.get_or_build_mcp_loader(&bad.key, &ts).await {
2156 Ok(_) => panic!("invalid [mcp] must fail loud"),
2157 Err(e) => e,
2158 };
2159 assert!(
2160 err.contains("invalid [mcp]"),
2161 "surfaces the parse error: {err}"
2162 );
2163
2164 let entry = state
2165 .mcp_loaders
2166 .get(&trustee_core::user_hash(&bad.key))
2167 .unwrap();
2168 assert!(entry.degraded.is_some(), "poison entry recorded");
2169
2170 let err2 = match state.get_or_build_mcp_loader(&bad.key, &ts).await {
2172 Ok(_) => panic!("still within backoff"),
2173 Err(e) => e,
2174 };
2175 assert_eq!(err, err2, "same cached error");
2176
2177 let good_loader = state
2179 .get_or_build_mcp_loader(&good.key, &ts)
2180 .await
2181 .expect("other user unaffected");
2182 assert!(
2183 good_loader.is_none(),
2184 "good user has no [mcp] → disabled marker"
2185 );
2186 }
2187}
2188
2189#[cfg(test)]
2190mod restore_tests {
2191 use super::*;
2192
2193 fn fresh_state() -> ServerState {
2194 let (session, _rx) = Session::new();
2195 let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(16);
2196 let mut state = ServerState::new(session, ws_tx, None);
2197 state.config_toml = Some("[agent]\nname = \"trustee\"\n".to_string());
2198 state
2199 }
2200
2201 struct TempUser {
2203 key: String,
2204 }
2205
2206 impl TempUser {
2207 fn new(tag: &str) -> Self {
2208 let key = format!("restore-{tag}-{}@test.invalid", std::process::id());
2209 let home = dirs::home_dir()
2210 .expect("HOME available in test env")
2211 .join(".trustee")
2212 .join("users")
2213 .join(trustee_core::user_hash(&key));
2214 std::fs::create_dir_all(home).expect("create user home");
2215 Self { key }
2216 }
2217 }
2218
2219 #[tokio::test]
2220 async fn restore_unknown_chain_returns_none() {
2221 let state = fresh_state();
2222 let user = TempUser::new("unknown");
2223 let out = state
2224 .get_or_restore_session_by_any_id(&user.key, "session_2099_01_01_00_00_deadbeef")
2225 .await;
2226 assert!(out.is_none(), "unknown chain must 404, not restore");
2227 }
2228
2229 #[tokio::test]
2230 async fn destroy_tombstones_chain_id() {
2231 let state = fresh_state();
2232 let user = TempUser::new("tombstone");
2233 let sid = state
2234 .create_session(&user.key, None, None, true)
2235 .await
2236 .expect("create session");
2237
2238 let (session, _ws) = state.get_session(&user.key, &sid).await.unwrap();
2240 {
2241 let mut s = session.lock().await;
2242 s.session_id = Some("session_2026_09_04_10_00_aaaabbbb".to_string());
2243 }
2244
2245 state.destroy_session(&user.key, &sid).await.expect("destroy");
2246
2247 let path = state
2248 .tombstone_path(&user.key, "session_2026_09_04_10_00_aaaabbbb")
2249 .expect("tombstone path");
2250 assert!(path.exists(), "destroy must write a chain tombstone");
2251
2252 let sid2 = state
2254 .create_session(&user.key, None, None, true)
2255 .await
2256 .expect("create second session");
2257 state.destroy_session(&user.key, &sid2).await.expect("destroy 2");
2258 }
2260
2261 #[tokio::test]
2262 async fn restore_refuses_registry_key_shapes() {
2263 let state = fresh_state();
2267 let user = TempUser::new("shapes");
2268 assert!(state.tombstone_path(&user.key, "default").is_none());
2269 let out = state.get_or_restore_session_by_any_id(&user.key, "default").await;
2270 assert!(out.is_none());
2271 }
2272}
2273
2274 #[test]
2275 fn collect_unresolved_placeholders_walks_the_tree() {
2276 let section: toml::Value = toml::from_str(
2277 r#"
2278[credentials.farzan_service_account]
2279service_token = "${FARZAN_SERVICE_ACCOUNT}"
2280issuer_url = "https://ok.example"
2281
2282[[servers]]
2283name = "fame"
2284url = "https://fame.example"
2285credentials = "farzan_service_account"
2286
2287[[servers]]
2288name = "other"
2289url = "https://other.example"
2290token = "${OTHER_TOKEN}"
2291"#,
2292 )
2293 .unwrap();
2294 let mut out = Vec::new();
2295 collect_unresolved_placeholders(§ion, &mut out);
2296 out.sort();
2297 assert_eq!(
2298 out,
2299 vec!["${FARZAN_SERVICE_ACCOUNT}".to_string(), "${OTHER_TOKEN}".to_string()],
2300 "every leftover placeholder in the [mcp] tree must be collected, deduped"
2301 );
2302 let clean: toml::Value = toml::from_str("[x]\ny = \"z\"\n").unwrap();
2304 let mut out2 = Vec::new();
2305 collect_unresolved_placeholders(&clean, &mut out2);
2306 assert!(out2.is_empty());
2307 }
2308
2309 #[test]
2310 fn substitute_env_vars_resolves_and_keeps_leftovers_deterministically() {
2311 let mut secrets = std::collections::HashMap::new();
2312 secrets.insert("A_TOKEN".to_string(), "resolved-value".to_string());
2313 let mut s = "x = \"${A_TOKEN}\" y = \"${MISSING_TOKEN}\"".to_string();
2314 substitute_env_vars(&mut s, &secrets);
2315 assert!(
2316 s.contains("resolved-value"),
2317 "declared key resolves: {s}"
2318 );
2319 assert!(
2320 s.contains("${MISSING_TOKEN}"),
2321 "undeclared key keeps the literal (the loader guard turns this into a loud failure): {s}"
2322 );
2323 }