rustyclaw_core/
sessions.rs

1//! Session management for RustyClaw multi-agent support.
2//!
3//! Provides tools for spawning sub-agents, sending messages between sessions,
4//! and managing session state.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8use std::sync::{Arc, Mutex, OnceLock};
9use std::time::{SystemTime, UNIX_EPOCH};
10
11/// Session key format: agent:<agentId>:subagent:<uuid> or agent:<agentId>:main
12pub type SessionKey = String;
13
14/// Generate a unique session key for a sub-agent.
15fn generate_subagent_key(agent_id: &str) -> SessionKey {
16    let uuid = generate_uuid();
17    format!("agent:{}:subagent:{}", agent_id, uuid)
18}
19
20/// Generate a simple UUID-like string.
21fn generate_uuid() -> String {
22    let timestamp = SystemTime::now()
23        .duration_since(UNIX_EPOCH)
24        .unwrap_or_default()
25        .as_nanos();
26    format!("{:x}", timestamp)
27}
28
29/// Session status.
30#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
31#[serde(rename_all = "camelCase")]
32pub enum SessionStatus {
33    Active,
34    Completed,
35    Error,
36    Timeout,
37    Stopped,
38}
39
40/// Session kind.
41#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
42#[serde(rename_all = "camelCase")]
43pub enum SessionKind {
44    Main,
45    Subagent,
46    Cron,
47}
48
49/// A message in a session.
50#[derive(Debug, Clone, Serialize, Deserialize)]
51#[serde(rename_all = "camelCase")]
52pub struct SessionMessage {
53    pub role: String, // "user", "assistant", "system", "tool"
54    pub content: String,
55    pub timestamp_ms: u64,
56    #[serde(skip_serializing_if = "Option::is_none")]
57    pub tool_name: Option<String>,
58}
59
60/// A session record.
61#[derive(Debug, Clone, Serialize, Deserialize)]
62#[serde(rename_all = "camelCase")]
63pub struct Session {
64    pub key: SessionKey,
65    pub agent_id: String,
66    pub kind: SessionKind,
67    pub status: SessionStatus,
68    #[serde(skip_serializing_if = "Option::is_none")]
69    pub label: Option<String>,
70    #[serde(skip_serializing_if = "Option::is_none")]
71    pub task: Option<String>,
72    pub created_ms: u64,
73    #[serde(skip_serializing_if = "Option::is_none")]
74    pub finished_ms: Option<u64>,
75    /// Recent messages (limited for memory efficiency).
76    pub messages: Vec<SessionMessage>,
77    /// Run ID for sub-agents.
78    #[serde(skip_serializing_if = "Option::is_none")]
79    pub run_id: Option<String>,
80    /// Parent session key (for sub-agents).
81    #[serde(skip_serializing_if = "Option::is_none")]
82    pub parent_key: Option<SessionKey>,
83}
84
85impl Session {
86    /// Create a new main session.
87    pub fn new_main(agent_id: &str) -> Self {
88        let now_ms = now_millis();
89        Self {
90            key: format!("agent:{}:main", agent_id),
91            agent_id: agent_id.to_string(),
92            kind: SessionKind::Main,
93            status: SessionStatus::Active,
94            label: None,
95            task: None,
96            created_ms: now_ms,
97            finished_ms: None,
98            messages: Vec::new(),
99            run_id: None,
100            parent_key: None,
101        }
102    }
103
104    /// Create a new sub-agent session.
105    pub fn new_subagent(
106        agent_id: &str,
107        task: &str,
108        label: Option<String>,
109        parent_key: Option<SessionKey>,
110    ) -> Self {
111        let now_ms = now_millis();
112        let run_id = generate_uuid();
113        Self {
114            key: generate_subagent_key(agent_id),
115            agent_id: agent_id.to_string(),
116            kind: SessionKind::Subagent,
117            status: SessionStatus::Active,
118            label,
119            task: Some(task.to_string()),
120            created_ms: now_ms,
121            finished_ms: None,
122            messages: Vec::new(),
123            run_id: Some(run_id),
124            parent_key,
125        }
126    }
127
128    /// Add a message to the session.
129    pub fn add_message(&mut self, role: &str, content: &str) {
130        self.messages.push(SessionMessage {
131            role: role.to_string(),
132            content: content.to_string(),
133            timestamp_ms: now_millis(),
134            tool_name: None,
135        });
136
137        // Keep only last 100 messages in memory
138        if self.messages.len() > 100 {
139            self.messages.remove(0);
140        }
141    }
142
143    /// Mark session as completed.
144    pub fn complete(&mut self) {
145        self.status = SessionStatus::Completed;
146        self.finished_ms = Some(now_millis());
147    }
148
149    /// Mark session as errored.
150    pub fn error(&mut self) {
151        self.status = SessionStatus::Error;
152        self.finished_ms = Some(now_millis());
153    }
154
155    /// Get runtime in seconds.
156    pub fn runtime_secs(&self) -> u64 {
157        let end = self.finished_ms.unwrap_or_else(now_millis);
158        (end - self.created_ms) / 1000
159    }
160}
161
162/// Global session manager.
163pub struct SessionManager {
164    sessions: HashMap<SessionKey, Session>,
165    /// Map labels to session keys for easy lookup.
166    labels: HashMap<String, SessionKey>,
167}
168
169impl SessionManager {
170    /// Create a new session manager.
171    pub fn new() -> Self {
172        Self {
173            sessions: HashMap::new(),
174            labels: HashMap::new(),
175        }
176    }
177
178    /// Create or get a main session.
179    pub fn get_or_create_main(&mut self, agent_id: &str) -> &Session {
180        let key = format!("agent:{}:main", agent_id);
181        self.sessions
182            .entry(key.clone())
183            .or_insert_with(|| Session::new_main(agent_id))
184    }
185
186    /// Spawn a sub-agent session.
187    pub fn spawn_subagent(
188        &mut self,
189        agent_id: &str,
190        task: &str,
191        label: Option<String>,
192        parent_key: Option<SessionKey>,
193    ) -> SessionKey {
194        let session = Session::new_subagent(agent_id, task, label.clone(), parent_key);
195        let key = session.key.clone();
196
197        if let Some(ref lbl) = label {
198            self.labels.insert(lbl.clone(), key.clone());
199        }
200
201        self.sessions.insert(key.clone(), session);
202        key
203    }
204
205    /// Get a session by key.
206    pub fn get(&self, key: &str) -> Option<&Session> {
207        self.sessions.get(key)
208    }
209
210    /// Get a session by label.
211    pub fn get_by_label(&self, label: &str) -> Option<&Session> {
212        self.labels.get(label).and_then(|k| self.sessions.get(k))
213    }
214
215    /// Get a mutable session by key.
216    pub fn get_mut(&mut self, key: &str) -> Option<&mut Session> {
217        self.sessions.get_mut(key)
218    }
219
220    /// List sessions with optional filters.
221    pub fn list(
222        &self,
223        kinds: Option<&[SessionKind]>,
224        active_only: bool,
225        limit: usize,
226    ) -> Vec<&Session> {
227        let mut sessions: Vec<_> = self
228            .sessions
229            .values()
230            .filter(|s| {
231                let kind_match = kinds.map(|ks| ks.contains(&s.kind)).unwrap_or(true);
232                let active_match = !active_only || s.status == SessionStatus::Active;
233                kind_match && active_match
234            })
235            .collect();
236
237        // Sort by created time descending
238        sessions.sort_by(|a, b| b.created_ms.cmp(&a.created_ms));
239        sessions.truncate(limit);
240        sessions
241    }
242
243    /// Get message history for a session.
244    pub fn history(
245        &self,
246        key: &str,
247        limit: usize,
248        include_tools: bool,
249    ) -> Option<Vec<&SessionMessage>> {
250        self.sessions.get(key).map(|s| {
251            s.messages
252                .iter()
253                .filter(|m| include_tools || m.role != "tool")
254                .rev()
255                .take(limit)
256                .collect::<Vec<_>>()
257                .into_iter()
258                .rev()
259                .collect()
260        })
261    }
262
263    /// Send a message to a session.
264    pub fn send_message(&mut self, key: &str, message: &str) -> Result<(), String> {
265        let session = self
266            .sessions
267            .get_mut(key)
268            .ok_or_else(|| format!("Session not found: {}", key))?;
269
270        if session.status != SessionStatus::Active {
271            return Err(format!("Session is not active: {:?}", session.status));
272        }
273
274        session.add_message("user", message);
275        Ok(())
276    }
277
278    /// Complete a session.
279    pub fn complete_session(&mut self, key: &str) -> Result<(), String> {
280        let session = self
281            .sessions
282            .get_mut(key)
283            .ok_or_else(|| format!("Session not found: {}", key))?;
284        session.complete();
285        Ok(())
286    }
287}
288
289impl Default for SessionManager {
290    fn default() -> Self {
291        Self::new()
292    }
293}
294
295/// Thread-safe session manager.
296pub type SharedSessionManager = Arc<Mutex<SessionManager>>;
297
298/// Global session manager instance.
299static SESSION_MANAGER: OnceLock<SharedSessionManager> = OnceLock::new();
300
301/// Get the global session manager.
302pub fn session_manager() -> &'static SharedSessionManager {
303    SESSION_MANAGER.get_or_init(|| Arc::new(Mutex::new(SessionManager::new())))
304}
305
306/// Spawn result returned to the agent.
307#[derive(Debug, Clone, Serialize, Deserialize)]
308#[serde(rename_all = "camelCase")]
309pub struct SpawnResult {
310    pub status: String,
311    pub run_id: String,
312    pub session_key: SessionKey,
313    pub message: String,
314}
315
316/// Get current time in milliseconds.
317fn now_millis() -> u64 {
318    SystemTime::now()
319        .duration_since(UNIX_EPOCH)
320        .unwrap_or_default()
321        .as_millis() as u64
322}
323
324#[cfg(test)]
325mod tests {
326    use super::*;
327
328    #[test]
329    fn test_session_creation() {
330        let session = Session::new_main("main");
331        assert_eq!(session.key, "agent:main:main");
332        assert_eq!(session.kind, SessionKind::Main);
333        assert_eq!(session.status, SessionStatus::Active);
334    }
335
336    #[test]
337    fn test_subagent_spawn() {
338        let mut manager = SessionManager::new();
339        let key =
340            manager.spawn_subagent("main", "Research task", Some("research".to_string()), None);
341
342        assert!(key.contains("subagent"));
343
344        let session = manager.get(&key).unwrap();
345        assert_eq!(session.kind, SessionKind::Subagent);
346        assert_eq!(session.task, Some("Research task".to_string()));
347        assert_eq!(session.label, Some("research".to_string()));
348
349        // Should be findable by label
350        let by_label = manager.get_by_label("research").unwrap();
351        assert_eq!(by_label.key, key);
352    }
353
354    #[test]
355    fn test_message_history() {
356        let mut manager = SessionManager::new();
357        let key = manager.spawn_subagent("main", "Test", None, None);
358
359        manager.send_message(&key, "Hello").unwrap();
360
361        let session = manager.get_mut(&key).unwrap();
362        session.add_message("assistant", "Hi there!");
363
364        let history = manager.history(&key, 10, false).unwrap();
365        assert_eq!(history.len(), 2);
366        assert_eq!(history[0].content, "Hello");
367        assert_eq!(history[1].content, "Hi there!");
368    }
369
370    #[test]
371    fn test_session_listing() {
372        let mut manager = SessionManager::new();
373        manager.get_or_create_main("main");
374        manager.spawn_subagent("main", "Task 1", None, None);
375        manager.spawn_subagent("main", "Task 2", None, None);
376
377        let all = manager.list(None, false, 10);
378        assert_eq!(all.len(), 3);
379
380        let subagents = manager.list(Some(&[SessionKind::Subagent]), false, 10);
381        assert_eq!(subagents.len(), 2);
382    }
383
384    #[test]
385    fn test_subagent_appears_in_active_list() {
386        // This test verifies that spawned subagents show up when listing active sessions
387        // which is what the gateway uses to populate the sidebar
388        let mut manager = SessionManager::new();
389        
390        // Spawn a subagent
391        let key = manager.spawn_subagent("main", "Test task", Some("test".to_string()), None);
392        
393        // Verify it appears in active-only list (what sidebar uses)
394        let active = manager.list(Some(&[SessionKind::Subagent]), true, 10);
395        assert_eq!(active.len(), 1, "Subagent should appear in active list");
396        assert_eq!(active[0].key, key);
397        assert_eq!(active[0].status, SessionStatus::Active);
398        assert_eq!(active[0].label, Some("test".to_string()));
399    }
400}
rustyclaw_core/sessions.rs

rustyclaw_core/
sessions.rs
