1use std::sync::Arc;
16
17use dashmap::DashMap;
18use tokio::sync::{broadcast, mpsc, Mutex};
19use trustee_core::session::Session;
20use trustee_core::types::TuiMessage;
21
22use crate::auth::AuthState;
23
24pub struct UserSessionEntry {
30 pub session: Arc<Mutex<Session>>,
32 pub ws_tx: broadcast::Sender<String>,
34 pub created_at: chrono::DateTime<chrono::Utc>,
36 pub last_active: Arc<Mutex<chrono::DateTime<chrono::Utc>>>,
39}
40
41pub struct UserSessions {
43 pub sessions: DashMap<String, UserSessionEntry>,
45 pub token_store: Arc<pep::MemoryTokenStore>,
47 pub active_session_id: Mutex<String>,
49}
50
51#[derive(Debug, serde::Serialize)]
53pub struct SessionListItem {
54 pub session_id: String,
55 pub session_name: Option<String>,
56 pub workflow_state: String,
57 pub created_at: String,
58 pub last_active: String,
59}
60
61#[derive(Debug)]
63pub enum SessionError {
64 MaxSessionsReached(usize),
66 NotFound(String),
68 NotIdle(String),
70}
71
72impl std::fmt::Display for SessionError {
73 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74 match self {
75 SessionError::MaxSessionsReached(n) => {
76 write!(f, "Maximum {} sessions per user reached", n)
77 }
78 SessionError::NotFound(id) => write!(f, "Session {} not found", id),
79 SessionError::NotIdle(state) => write!(f, "Session is not idle (state: {})", state),
80 }
81 }
82}
83
84impl std::error::Error for SessionError {}
85
86pub type SessionRegistry = Arc<DashMap<String, UserSessions>>;
88
89#[derive(Clone)]
95pub struct ServerState {
96 pub sessions: SessionRegistry,
98 pub ws_tx: broadcast::Sender<String>,
100 pub auth: Option<Arc<AuthState>>,
102 pub config_toml: Option<String>,
104 pub secrets: Option<std::collections::HashMap<String, String>>,
106 pub build_info: Option<trustee_core::types::BuildInfo>,
108 pub workflow_semaphore: Arc<tokio::sync::Semaphore>,
111 pub max_sessions_per_user: usize,
113}
114
115impl ServerState {
116 pub fn new(
118 session: Session,
119 ws_tx: broadcast::Sender<String>,
120 auth: Option<Arc<AuthState>>,
121 ) -> Self {
122 let sessions = Arc::new(DashMap::new());
123
124 let token_store = Arc::new(pep::MemoryTokenStore::new());
126 let (ws_tx_entry, _) = broadcast::channel::<String>(256);
127
128 let now = chrono::Utc::now();
129 let initial_entry = UserSessionEntry {
130 session: Arc::new(Mutex::new(session)),
131 ws_tx: ws_tx_entry,
132 created_at: now,
133 last_active: Arc::new(Mutex::new(now)),
134 };
135
136 let user_sessions = UserSessions {
137 sessions: DashMap::new(),
138 token_store,
139 active_session_id: Mutex::new(String::new()),
140 };
141 user_sessions.sessions.insert("default".to_string(), initial_entry);
142
143 sessions.insert("default".to_string(), user_sessions);
144
145 Self {
146 sessions,
147 ws_tx,
148 auth,
149 config_toml: None,
150 secrets: None,
151 build_info: None,
152 workflow_semaphore: Arc::new(tokio::sync::Semaphore::new(8)),
153 max_sessions_per_user: 4,
154 }
155 }
156
157 pub fn with_config_toml(mut self, config_toml: String) -> Self {
158 self.config_toml = Some(config_toml);
159 self
160 }
161
162 pub fn with_secrets(mut self, secrets: std::collections::HashMap<String, String>) -> Self {
163 self.secrets = Some(secrets);
164 self
165 }
166
167 pub fn with_build_info(mut self, build_info: trustee_core::types::BuildInfo) -> Self {
168 self.build_info = Some(build_info);
169 self
170 }
171
172 pub fn with_max_concurrent_workflows(mut self, max: usize) -> Self {
173 self.workflow_semaphore = Arc::new(tokio::sync::Semaphore::new(max));
174 self
175 }
176
177 pub fn with_max_sessions_per_user(mut self, max: usize) -> Self {
179 self.max_sessions_per_user = max;
180 self
181 }
182
183 pub async fn create_session(
193 &self,
194 user_key: &str,
195 session_name: Option<String>,
196 ) -> Result<String, SessionError> {
197 let user_sessions = self
199 .sessions
200 .entry(user_key.to_string())
201 .or_insert_with(|| UserSessions {
202 sessions: DashMap::new(),
203 token_store: Arc::new(pep::MemoryTokenStore::new()),
204 active_session_id: Mutex::new(String::new()),
205 });
206
207 if user_sessions.sessions.len() >= self.max_sessions_per_user {
209 return Err(SessionError::MaxSessionsReached(self.max_sessions_per_user));
210 }
211
212 let (mut session, workflow_rx) = Session::new();
214
215 if let Some(ref config_toml) = self.config_toml {
217 session.config_toml = Some(config_toml.clone());
218 session.parse_auto_handoff_config();
219 if let Ok(table) = config_toml.parse::<toml::Value>() {
220 if let Some(name) = table
221 .get("agent")
222 .and_then(|a| a.get("name"))
223 .and_then(|n| n.as_str())
224 {
225 session.agent_name = name.to_string();
226 }
227 }
228 }
229
230 session.secrets = self.secrets.clone();
231 session.build_info = self.build_info.clone();
232
233 self.apply_user_isolation(&mut session, user_key);
235
236 session.session_name = session_name;
238
239 let (ws_tx_entry, _) = broadcast::channel::<String>(256);
241
242 let session_id = format!(
244 "session_{}_{}",
245 chrono::Utc::now().format("%Y_%m_%d_%H_%M"),
246 &uuid::Uuid::new_v4().to_string()[..8]
247 );
248
249 let now = chrono::Utc::now();
250
251 user_sessions.sessions.insert(
253 session_id.clone(),
254 UserSessionEntry {
255 session: Arc::new(Mutex::new(session)),
256 ws_tx: ws_tx_entry.clone(),
257 created_at: now,
258 last_active: Arc::new(Mutex::new(now)),
259 },
260 );
261
262 *user_sessions.active_session_id.lock().await = session_id.clone();
264
265 let session_arc = user_sessions
267 .sessions
268 .get(&session_id)
269 .map(|e| e.session.clone());
270 if let Some(session_arc) = session_arc {
271 self.spawn_user_drain_task(
272 session_id.clone(),
273 session_arc,
274 ws_tx_entry,
275 workflow_rx,
276 );
277 }
278
279 Ok(session_id)
280 }
281
282 pub async fn get_session(
285 &self,
286 user_key: &str,
287 session_id: &str,
288 ) -> Option<(Arc<Mutex<Session>>, broadcast::Sender<String>)> {
289 let user_sessions = self.sessions.get(user_key)?;
290 let entry = user_sessions.sessions.get(session_id)?;
291
292 let now = chrono::Utc::now();
294 *entry.last_active.lock().await = now;
295
296 Some((entry.session.clone(), entry.ws_tx.clone()))
297 }
298
299 pub async fn list_sessions(&self, user_key: &str) -> Vec<SessionListItem> {
301 let Some(user_sessions) = self.sessions.get(user_key) else {
302 return Vec::new();
303 };
304
305 let mut items = Vec::new();
306 for entry in user_sessions.sessions.iter() {
307 let session = entry.session.lock().await;
308 let workflow_state = match session.workflow_state {
309 trustee_core::types::WorkflowState::Idle => "Idle",
310 trustee_core::types::WorkflowState::Running => "Running",
311 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
312 };
313 let last_active = entry.last_active.lock().await;
314 items.push(SessionListItem {
315 session_id: entry.key().clone(),
316 session_name: session.session_name.clone(),
317 workflow_state: workflow_state.to_string(),
318 created_at: entry.created_at.to_rfc3339(),
319 last_active: last_active.to_rfc3339(),
320 });
321 }
322 drop(user_sessions);
323
324 items.sort_by(|a, b| b.last_active.cmp(&a.last_active));
326 items
327 }
328
329 pub async fn destroy_session(
331 &self,
332 user_key: &str,
333 session_id: &str,
334 ) -> Result<(), SessionError> {
335 let user_sessions = self
336 .sessions
337 .get(user_key)
338 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
339
340 {
342 let entry = user_sessions
343 .sessions
344 .get(session_id)
345 .ok_or_else(|| SessionError::NotFound(session_id.to_string()))?;
346 let session = entry.session.lock().await;
347 if session.workflow_state != trustee_core::types::WorkflowState::Idle {
348 let state_str = match session.workflow_state {
349 trustee_core::types::WorkflowState::Running => "Running",
350 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
351 _ => "Unknown",
352 };
353 return Err(SessionError::NotIdle(state_str.to_string()));
354 }
355 }
356
357 user_sessions.sessions.remove(session_id);
359
360 let mut active_id = user_sessions.active_session_id.lock().await;
362 if &*active_id == session_id {
363 let mut newest: Option<(String, chrono::DateTime<chrono::Utc>)> = None;
365 for entry in user_sessions.sessions.iter() {
366 let la = entry.last_active.lock().await;
367 if newest.as_ref().map_or(true, |(_, t)| *la > *t) {
368 newest = Some((entry.key().clone(), *la));
369 }
370 }
371 *active_id = newest.map(|(id, _)| id).unwrap_or_default();
372 }
373
374 Ok(())
375 }
376
377 pub async fn ensure_active_session(
386 &self,
387 user_key: &str,
388 ) -> (
389 String,
390 Arc<Mutex<Session>>,
391 broadcast::Sender<String>,
392 Arc<pep::MemoryTokenStore>,
393 ) {
394 let token_store = {
396 let user_sessions = self
397 .sessions
398 .entry(user_key.to_string())
399 .or_insert_with(|| UserSessions {
400 sessions: DashMap::new(),
401 token_store: Arc::new(pep::MemoryTokenStore::new()),
402 active_session_id: Mutex::new(String::new()),
403 });
404 user_sessions.token_store.clone()
405 };
406
407 let active_id = {
409 let user_sessions = self.sessions.get(user_key).unwrap();
410 let guard = user_sessions.active_session_id.lock().await;
411 guard.clone()
412 };
413
414 if !active_id.is_empty() {
415 if let Some((session, ws_tx)) = self.get_session(user_key, &active_id).await {
416 return (active_id, session, ws_tx, token_store);
417 }
418 }
420
421 let existing_session: Option<(String, Arc<Mutex<Session>>, broadcast::Sender<String>)> = {
425 let user_sessions = self.sessions.get(user_key).unwrap();
426 let result = user_sessions.sessions.iter().next().map(|first| {
427 (
428 first.key().clone(),
429 first.session.clone(),
430 first.ws_tx.clone(),
431 )
432 });
433 result
434 };
435 if let Some((id, session, ws_tx)) = existing_session {
436 let now = chrono::Utc::now();
437 if let Some(entry) = self.sessions.get(user_key) {
438 if let Some(e) = entry.sessions.get(&id) {
439 *e.last_active.lock().await = now;
440 }
441 *entry.active_session_id.lock().await = id.clone();
442 }
443
444 return (id, session, ws_tx, token_store);
445 }
446
447 let session_id = self
449 .create_session(user_key, None)
450 .await
451 .unwrap_or_else(|_| "default".to_string());
452
453 let (session, ws_tx) = self
454 .get_session(user_key, &session_id)
455 .await
456 .expect("just-created session must exist");
457
458 (session_id, session, ws_tx, token_store)
459 }
460
461 pub async fn ensure_user_session(
464 &self,
465 user_key: &str,
466 ) -> (Arc<Mutex<Session>>, broadcast::Sender<String>, Arc<pep::MemoryTokenStore>) {
467 let (_id, session, ws_tx, token_store) = self.ensure_active_session(user_key).await;
468 (session, ws_tx, token_store)
469 }
470
471 pub async fn set_active_session(&self, user_key: &str, session_id: &str) {
473 if let Some(user_sessions) = self.sessions.get(user_key) {
474 if user_sessions.sessions.contains_key(session_id) {
475 *user_sessions.active_session_id.lock().await = session_id.to_string();
476 }
477 }
478 }
479
480 fn apply_user_isolation(&self, session: &mut Session, user_key: &str) {
486 use sha2::{Digest, Sha256};
487 let mut hasher = Sha256::new();
488 hasher.update(user_key.as_bytes());
489 let hash_bytes = hasher.finalize();
490 let user_hash = format!(
491 "{:016x}",
492 u64::from_be_bytes(hash_bytes[..8].try_into().unwrap())
493 );
494
495 if let Some(home) = dirs::home_dir() {
496 session.home_dir = Some(home.join(".trustee").join("users").join(&user_hash));
497 }
498 session.project_id = Some(format!("web{}", &user_hash[..16]));
499 }
500
501 fn spawn_user_drain_task(
503 &self,
504 session_id: String,
505 session: Arc<Mutex<Session>>,
506 ws_tx: broadcast::Sender<String>,
507 mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>,
508 ) {
509 tokio::spawn(async move {
510 while let Some(msg) = workflow_rx.recv().await {
511 {
512 let mut session = session.lock().await;
513 session.handle_workflow_message(msg.clone());
514
515 let state_str = match session.workflow_state {
516 trustee_core::types::WorkflowState::Idle => "Idle",
517 trustee_core::types::WorkflowState::Running => "Running",
518 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
519 };
520 let state_msg = serde_json::json!({
521 "type": "StateChanged",
522 "state": state_str
523 });
524 let _ = ws_tx.send(state_msg.to_string());
525 }
526
527 let json =
528 serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
529 let _ = ws_tx.send(json);
530 }
531 tracing::debug!("Drain task ended for session: {}", session_id);
532 });
533 }
534
535 pub fn spawn_drain_task(self, mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>) {
538 let default_user = self
540 .sessions
541 .get("default")
542 .expect("default user must exist");
543 let first_entry = default_user
544 .sessions
545 .iter()
546 .next()
547 .expect("default user must have at least one session");
548 let session = first_entry.session.clone();
549 let ws_tx = first_entry.ws_tx.clone();
550 let session_id = first_entry.key().clone();
551 drop(first_entry);
552 drop(default_user);
553
554 tokio::spawn(async move {
555 while let Some(msg) = workflow_rx.recv().await {
556 {
557 let mut session = session.lock().await;
558 session.handle_workflow_message(msg.clone());
559
560 let state_str = match session.workflow_state {
561 trustee_core::types::WorkflowState::Idle => "Idle",
562 trustee_core::types::WorkflowState::Running => "Running",
563 trustee_core::types::WorkflowState::Cancelling => "Cancelling",
564 };
565 let state_msg = serde_json::json!({
566 "type": "StateChanged",
567 "state": state_str
568 });
569 let _ = ws_tx.send(state_msg.to_string());
570 }
571
572 let json =
573 serde_json::to_string(&SerializableMessage(&msg)).unwrap_or_default();
574 let _ = ws_tx.send(json);
575 }
576 tracing::debug!("Drain task ended for session: {}", session_id);
577 });
578 }
579
580 pub async fn resolve_user_key(&self, headers: &axum::http::HeaderMap) -> String {
582 let Some(ref auth) = self.auth else {
583 return "default".to_string();
584 };
585
586 if let Some(token) = headers
588 .get(axum::http::header::AUTHORIZATION)
589 .and_then(|v| v.to_str().ok())
590 .and_then(|v| v.strip_prefix("Bearer "))
591 .map(|s| s.to_string())
592 {
593 if token.starts_with("dev:") {
594 let parts: Vec<&str> = token.splitn(4, ':').collect();
595 if parts.len() >= 4 {
596 return format!("dev:{}", parts[1]);
597 }
598 }
599 if let Ok(claims) = auth.validate_token(&token).await {
600 return claims.sub;
601 }
602 }
603
604 let cookie_session_id = headers
606 .get(axum::http::header::COOKIE)
607 .and_then(|v| v.to_str().ok())
608 .and_then(|cookies| {
609 cookies
610 .split(';')
611 .map(|c| c.trim())
612 .find_map(|c| {
613 c.strip_prefix(&format!("{}=", auth.config.cookie_name))
614 .map(|s| s.to_string())
615 })
616 });
617
618 if let Some(session_id) = cookie_session_id {
619 if session_id.starts_with("dev:") {
620 let parts: Vec<&str> = session_id.splitn(4, ':').collect();
621 if parts.len() >= 4 {
622 return format!("dev:{}", parts[1]);
623 }
624 }
625
626 if let Ok(access_token) = auth.session_manager.get_token(&session_id).await {
627 if let Ok(claims) = auth.validate_token(&access_token).await {
628 return claims.sub;
629 }
630 }
631 }
632
633 "default".to_string()
634 }
635}
636
637struct SerializableMessage<'a>(&'a TuiMessage);
643
644impl<'a> serde::Serialize for SerializableMessage<'a> {
645 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
646 where
647 S: serde::Serializer,
648 {
649 use serde::ser::SerializeStruct;
650
651 match self.0 {
652 TuiMessage::OutputLine(line) => {
653 let mut s = serializer.serialize_struct("msg", 2)?;
654 s.serialize_field("type", "OutputLine")?;
655 s.serialize_field("line", line)?;
656 s.end()
657 }
658 TuiMessage::StreamDelta(delta) => {
659 let mut s = serializer.serialize_struct("msg", 2)?;
660 s.serialize_field("type", "StreamDelta")?;
661 s.serialize_field("delta", delta)?;
662 s.end()
663 }
664 TuiMessage::ReasoningDelta(delta) => {
665 let mut s = serializer.serialize_struct("msg", 2)?;
666 s.serialize_field("type", "ReasoningDelta")?;
667 s.serialize_field("delta", delta)?;
668 s.end()
669 }
670 TuiMessage::WorkflowCompleted => {
671 let mut s = serializer.serialize_struct("msg", 2)?;
672 s.serialize_field("type", "WorkflowCompleted")?;
673 s.serialize_field("state", "Idle")?;
674 s.end()
675 }
676 TuiMessage::WorkflowError(err) => {
677 let mut s = serializer.serialize_struct("msg", 2)?;
678 s.serialize_field("type", "WorkflowError")?;
679 s.serialize_field("error", err)?;
680 s.end()
681 }
682 TuiMessage::ResumeInfo(info) => match info {
683 Some(ri) => {
684 let mut s = serializer.serialize_struct("msg", 5)?;
685 s.serialize_field("type", "ResumeInfo")?;
686 s.serialize_field("state", "Idle")?;
687 s.serialize_field("session_id", &ri.session_id)?;
688 s.serialize_field("checkpoint_id", &ri.checkpoint_id)?;
689 s.serialize_field("iteration", &ri.iteration)?;
690 s.end()
691 }
692 None => {
693 let mut s = serializer.serialize_struct("msg", 2)?;
694 s.serialize_field("type", "ResumeInfo")?;
695 s.serialize_field("state", "Idle")?;
696 s.end()
697 }
698 },
699 TuiMessage::TodoUpdate(content) => {
700 let mut s = serializer.serialize_struct("msg", 2)?;
701 s.serialize_field("type", "TodoUpdate")?;
702 s.serialize_field("content", content)?;
703 s.end()
704 }
705 TuiMessage::WorkflowCancelled => {
706 let mut s = serializer.serialize_struct("msg", 2)?;
707 s.serialize_field("type", "WorkflowCancelled")?;
708 s.serialize_field("state", "Idle")?;
709 s.end()
710 }
711 TuiMessage::HandoffReady(_) => {
712 let mut s = serializer.serialize_struct("msg", 2)?;
713 s.serialize_field("type", "HandoffReady")?;
714 s.serialize_field("state", "Idle")?;
715 s.end()
716 }
717 TuiMessage::ToolPending {
718 tool_name,
719 hint,
720 } => {
721 let mut s = serializer.serialize_struct("msg", 3)?;
722 s.serialize_field("type", "ToolPending")?;
723 s.serialize_field("tool_name", tool_name)?;
724 s.serialize_field("hint", hint)?;
725 s.end()
726 }
727 TuiMessage::ToolDone {
728 tool_name,
729 success,
730 hint,
731 } => {
732 let mut s = serializer.serialize_struct("msg", 4)?;
733 s.serialize_field("type", "ToolDone")?;
734 s.serialize_field("tool_name", tool_name)?;
735 s.serialize_field("success", success)?;
736 s.serialize_field("hint", hint)?;
737 s.end()
738 }
739 TuiMessage::ContextTokensUpdated(count) => {
740 let mut s = serializer.serialize_struct("msg", 2)?;
741 s.serialize_field("type", "ContextTokensUpdated")?;
742 s.serialize_field("count", count)?;
743 s.end()
744 }
745 TuiMessage::McpServerStatus {
746 name,
747 connected,
748 tool_count,
749 error,
750 } => {
751 let mut s = serializer.serialize_struct("msg", 5)?;
752 s.serialize_field("type", "McpServerStatus")?;
753 s.serialize_field("name", name)?;
754 s.serialize_field("connected", connected)?;
755 s.serialize_field("tool_count", tool_count)?;
756 s.serialize_field("error", error)?;
757 s.end()
758 }
759 }
760 }
761}