Skip to main content

trustee_api/
state.rs

1//! Shared server state: per-user multi-session registry, broadcast channels, and auth state.
2//!
3//! ## Multi-Session Per User (MSU)
4//!
5//! Each authenticated user gets their own [`UserSessions`] containing N independent
6//! [`UserSessionEntry`] instances (default max 4). Each entry has:
7//! - An independent `Session` (workflow state, output, etc.)
8//! - A dedicated broadcast channel for WebSocket fan-out
9//! - Creation and last-active timestamps
10//!
11//! Sessions are keyed by user identity (`sub` claim from JWT, or `dev:email` for
12//! dev mode). Unauthenticated deployments use a single `"default"` key, preserving
13//! backward compatibility with single-user CLI operation.
14
15use std::sync::Arc;
16
17use dashmap::DashMap;
18use tokio::sync::{broadcast, mpsc, Mutex};
19use trustee_core::session::Session;
20use trustee_core::types::TuiMessage;
21
22use crate::auth::AuthState;
23
24// ---------------------------------------------------------------------------
25// Multi-session types
26// ---------------------------------------------------------------------------
27
28/// A single session with its own broadcast channel.
29pub struct UserSessionEntry {
30    /// The agent session, protected by a mutex.
31    pub session: Arc<Mutex<Session>>,
32    /// Broadcast sender for this session's WebSocket fan-out.
33    pub ws_tx: broadcast::Sender<String>,
34    /// When this session was created.
35    pub created_at: chrono::DateTime<chrono::Utc>,
36    /// Last time a command was submitted or state changed.
37    /// Updated on every /sessions/{id}/command and /sessions/{id}/cancel call.
38    pub last_active: Arc<Mutex<chrono::DateTime<chrono::Utc>>>,
39}
40
41/// All sessions belonging to one authenticated user.
42pub struct UserSessions {
43    /// session_id → session entry
44    pub sessions: DashMap<String, UserSessionEntry>,
45    /// Shared token store for all this user's sessions (MCP credential isolation).
46    pub token_store: Arc<pep::MemoryTokenStore>,
47    /// Which session_id is "active" for legacy /session/* routes.
48    pub active_session_id: Mutex<String>,
49}
50
51/// Summary of an active session for listing (serializable for API responses).
52#[derive(Debug, serde::Serialize)]
53pub struct SessionListItem {
54    pub session_id: String,
55    pub session_name: Option<String>,
56    pub workflow_state: String,
57    pub created_at: String,
58    pub last_active: String,
59}
60
61/// Errors from multi-session operations.
62#[derive(Debug)]
63pub enum SessionError {
64    /// User has reached max_sessions_per_user limit.
65    MaxSessionsReached(usize),
66    /// Session ID not found for this user.
67    NotFound(String),
68    /// Session is not Idle (cannot destroy/overwrite a running session).
69    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
86/// Top-level registry: user_key → user's session collection.
87pub type SessionRegistry = Arc<DashMap<String, UserSessions>>;
88
89// ---------------------------------------------------------------------------
90// ServerState
91// ---------------------------------------------------------------------------
92
93/// Shared state accessible by all axum handlers.
94#[derive(Clone)]
95pub struct ServerState {
96    /// Per-user multi-session registry (MSU).
97    pub sessions: SessionRegistry,
98    /// Broadcast sender for backward compat — delegates to the default user's channel.
99    pub ws_tx: broadcast::Sender<String>,
100    /// Auth state (None = auth disabled, all endpoints open).
101    pub auth: Option<Arc<AuthState>>,
102    /// Shared config TOML (all users share the same agent config).
103    pub config_toml: Option<String>,
104    /// Shared secrets (injected into every per-user session).
105    pub secrets: Option<std::collections::HashMap<String, String>>,
106    /// Shared build info (injected into every per-user session).
107    pub build_info: Option<trustee_core::types::BuildInfo>,
108    /// Global concurrency limiter — limits the number of simultaneous workflows
109    /// across all users. Default: 8 concurrent workflows.
110    pub workflow_semaphore: Arc<tokio::sync::Semaphore>,
111    /// Maximum number of concurrent sessions per user. Default: 4.
112    pub max_sessions_per_user: usize,
113}
114
115impl ServerState {
116    /// Create new shared state from a default session, broadcast sender, and optional auth.
117    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        // Store the default user's UserSessions with an initial session
125        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    /// Set the max sessions per user.
178    pub fn with_max_sessions_per_user(mut self, max: usize) -> Self {
179        self.max_sessions_per_user = max;
180        self
181    }
182
183    // -----------------------------------------------------------------------
184    // MSU: Multi-session methods
185    // -----------------------------------------------------------------------
186
187    /// Create a new session for a user. Returns the session_id.
188    ///
189    /// Creates a fresh `Session::new()`, copies shared config, sets per-user
190    /// isolation, creates a broadcast channel, spawns a drain task, and inserts
191    /// into the user's session DashMap. The new session becomes the "active" one.
192    pub async fn create_session(
193        &self,
194        user_key: &str,
195        session_name: Option<String>,
196    ) -> Result<String, SessionError> {
197        // Get or create the user's UserSessions entry
198        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        // Check session limit
208        if user_sessions.sessions.len() >= self.max_sessions_per_user {
209            return Err(SessionError::MaxSessionsReached(self.max_sessions_per_user));
210        }
211
212        // Create new Session
213        let (mut session, workflow_rx) = Session::new();
214
215        // Copy shared config
216        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        // Per-user isolation
234        self.apply_user_isolation(&mut session, user_key);
235
236        // Apply session_name if provided
237        session.session_name = session_name;
238
239        // Create broadcast channel
240        let (ws_tx_entry, _) = broadcast::channel::<String>(256);
241
242        // Generate session_id
243        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        // Insert into user's sessions DashMap
252        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        // Set as active session
263        *user_sessions.active_session_id.lock().await = session_id.clone();
264
265        // Spawn drain task
266        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    /// Get a specific session by user_key + session_id.
283    /// Updates last_active on the session entry.
284    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        // Update last_active
293        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    /// List all active sessions for a user, sorted by last_active desc.
300    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        // Sort by last_active descending
325        items.sort_by(|a, b| b.last_active.cmp(&a.last_active));
326        items
327    }
328
329    /// Destroy a session. The session must be Idle.
330    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        // Check workflow state before removing
341        {
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        // Remove from DashMap
358        user_sessions.sessions.remove(session_id);
359
360        // If this was the active session, pick a new active
361        let mut active_id = user_sessions.active_session_id.lock().await;
362        if &*active_id == session_id {
363            // Pick the most recently active remaining session
364            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    /// Get or create the user's "active" session for legacy routes.
378    ///
379    /// Behavior:
380    /// 1. If user has no sessions → create one
381    /// 2. If active session exists → return it
382    /// 3. If active session was destroyed → create a new one
383    ///
384    /// Returns: (session_id, session_arc, ws_tx, token_store)
385    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        // Get or create user's UserSessions
395        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        // Check if active session exists
408        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            // Active session was destroyed, fall through to create
419        }
420
421        // Need to create a new session
422        // For the "default" user, we may already have a "default" session entry
423        // from ServerState::new() — check for it
424        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        // Create a brand new session
448        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    /// DEPRECATED: Use ensure_active_session() instead.
462    /// Kept for backward compatibility — same 3-tuple return type.
463    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    /// Set a session as the user's active session.
472    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    // -----------------------------------------------------------------------
481    // Private helpers
482    // -----------------------------------------------------------------------
483
484    /// Apply per-user isolation: SHA-256 hash → home_dir + project_id.
485    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    /// Spawn a background drain task for a specific session's workflow receiver.
502    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    /// Spawn the default user's drain task (backward compatibility).
536    /// Called during server startup for the initial session.
537    pub fn spawn_drain_task(self, mut workflow_rx: mpsc::UnboundedReceiver<TuiMessage>) {
538        // Get the default user's first session
539        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    /// Resolve the user key from request headers.
581    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        // Try Bearer header first
587        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        // Try cookie
605        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
637// ---------------------------------------------------------------------------
638// SerializableMessage (unchanged)
639// ---------------------------------------------------------------------------
640
641/// Wrapper to serialize `TuiMessage` as JSON with a `type` discriminator.
642struct 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}