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(agent_id: &str, task: &str, label: Option<String>, parent_key: Option<SessionKey>) -> Self {
106        let now_ms = now_millis();
107        let run_id = generate_uuid();
108        Self {
109            key: generate_subagent_key(agent_id),
110            agent_id: agent_id.to_string(),
111            kind: SessionKind::Subagent,
112            status: SessionStatus::Active,
113            label,
114            task: Some(task.to_string()),
115            created_ms: now_ms,
116            finished_ms: None,
117            messages: Vec::new(),
118            run_id: Some(run_id),
119            parent_key,
120        }
121    }
122
123    /// Add a message to the session.
124    pub fn add_message(&mut self, role: &str, content: &str) {
125        self.messages.push(SessionMessage {
126            role: role.to_string(),
127            content: content.to_string(),
128            timestamp_ms: now_millis(),
129            tool_name: None,
130        });
131
132        // Keep only last 100 messages in memory
133        if self.messages.len() > 100 {
134            self.messages.remove(0);
135        }
136    }
137
138    /// Mark session as completed.
139    pub fn complete(&mut self) {
140        self.status = SessionStatus::Completed;
141        self.finished_ms = Some(now_millis());
142    }
143
144    /// Mark session as errored.
145    pub fn error(&mut self) {
146        self.status = SessionStatus::Error;
147        self.finished_ms = Some(now_millis());
148    }
149
150    /// Get runtime in seconds.
151    pub fn runtime_secs(&self) -> u64 {
152        let end = self.finished_ms.unwrap_or_else(now_millis);
153        (end - self.created_ms) / 1000
154    }
155}
156
157/// Global session manager.
158pub struct SessionManager {
159    sessions: HashMap<SessionKey, Session>,
160    /// Map labels to session keys for easy lookup.
161    labels: HashMap<String, SessionKey>,
162}
163
164impl SessionManager {
165    /// Create a new session manager.
166    pub fn new() -> Self {
167        Self {
168            sessions: HashMap::new(),
169            labels: HashMap::new(),
170        }
171    }
172
173    /// Create or get a main session.
174    pub fn get_or_create_main(&mut self, agent_id: &str) -> &Session {
175        let key = format!("agent:{}:main", agent_id);
176        self.sessions
177            .entry(key.clone())
178            .or_insert_with(|| Session::new_main(agent_id))
179    }
180
181    /// Spawn a sub-agent session.
182    pub fn spawn_subagent(
183        &mut self,
184        agent_id: &str,
185        task: &str,
186        label: Option<String>,
187        parent_key: Option<SessionKey>,
188    ) -> SessionKey {
189        let session = Session::new_subagent(agent_id, task, label.clone(), parent_key);
190        let key = session.key.clone();
191
192        if let Some(ref lbl) = label {
193            self.labels.insert(lbl.clone(), key.clone());
194        }
195
196        self.sessions.insert(key.clone(), session);
197        key
198    }
199
200    /// Get a session by key.
201    pub fn get(&self, key: &str) -> Option<&Session> {
202        self.sessions.get(key)
203    }
204
205    /// Get a session by label.
206    pub fn get_by_label(&self, label: &str) -> Option<&Session> {
207        self.labels.get(label).and_then(|k| self.sessions.get(k))
208    }
209
210    /// Get a mutable session by key.
211    pub fn get_mut(&mut self, key: &str) -> Option<&mut Session> {
212        self.sessions.get_mut(key)
213    }
214
215    /// List sessions with optional filters.
216    pub fn list(&self, kinds: Option<&[SessionKind]>, active_only: bool, limit: usize) -> Vec<&Session> {
217        let mut sessions: Vec<_> = self
218            .sessions
219            .values()
220            .filter(|s| {
221                let kind_match = kinds
222                    .map(|ks| ks.contains(&s.kind))
223                    .unwrap_or(true);
224                let active_match = !active_only || s.status == SessionStatus::Active;
225                kind_match && active_match
226            })
227            .collect();
228
229        // Sort by created time descending
230        sessions.sort_by(|a, b| b.created_ms.cmp(&a.created_ms));
231        sessions.truncate(limit);
232        sessions
233    }
234
235    /// Get message history for a session.
236    pub fn history(&self, key: &str, limit: usize, include_tools: bool) -> Option<Vec<&SessionMessage>> {
237        self.sessions.get(key).map(|s| {
238            s.messages
239                .iter()
240                .filter(|m| include_tools || m.role != "tool")
241                .rev()
242                .take(limit)
243                .collect::<Vec<_>>()
244                .into_iter()
245                .rev()
246                .collect()
247        })
248    }
249
250    /// Send a message to a session.
251    pub fn send_message(&mut self, key: &str, message: &str) -> Result<(), String> {
252        let session = self
253            .sessions
254            .get_mut(key)
255            .ok_or_else(|| format!("Session not found: {}", key))?;
256
257        if session.status != SessionStatus::Active {
258            return Err(format!("Session is not active: {:?}", session.status));
259        }
260
261        session.add_message("user", message);
262        Ok(())
263    }
264
265    /// Complete a session.
266    pub fn complete_session(&mut self, key: &str) -> Result<(), String> {
267        let session = self
268            .sessions
269            .get_mut(key)
270            .ok_or_else(|| format!("Session not found: {}", key))?;
271        session.complete();
272        Ok(())
273    }
274}
275
276impl Default for SessionManager {
277    fn default() -> Self {
278        Self::new()
279    }
280}
281
282/// Thread-safe session manager.
283pub type SharedSessionManager = Arc<Mutex<SessionManager>>;
284
285/// Global session manager instance.
286static SESSION_MANAGER: OnceLock<SharedSessionManager> = OnceLock::new();
287
288/// Get the global session manager.
289pub fn session_manager() -> &'static SharedSessionManager {
290    SESSION_MANAGER.get_or_init(|| Arc::new(Mutex::new(SessionManager::new())))
291}
292
293/// Spawn result returned to the agent.
294#[derive(Debug, Clone, Serialize, Deserialize)]
295#[serde(rename_all = "camelCase")]
296pub struct SpawnResult {
297    pub status: String,
298    pub run_id: String,
299    pub session_key: SessionKey,
300    pub message: String,
301}
302
303/// Get current time in milliseconds.
304fn now_millis() -> u64 {
305    SystemTime::now()
306        .duration_since(UNIX_EPOCH)
307        .unwrap_or_default()
308        .as_millis() as u64
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314
315    #[test]
316    fn test_session_creation() {
317        let session = Session::new_main("main");
318        assert_eq!(session.key, "agent:main:main");
319        assert_eq!(session.kind, SessionKind::Main);
320        assert_eq!(session.status, SessionStatus::Active);
321    }
322
323    #[test]
324    fn test_subagent_spawn() {
325        let mut manager = SessionManager::new();
326        let key = manager.spawn_subagent("main", "Research task", Some("research".to_string()), None);
327
328        assert!(key.contains("subagent"));
329
330        let session = manager.get(&key).unwrap();
331        assert_eq!(session.kind, SessionKind::Subagent);
332        assert_eq!(session.task, Some("Research task".to_string()));
333        assert_eq!(session.label, Some("research".to_string()));
334
335        // Should be findable by label
336        let by_label = manager.get_by_label("research").unwrap();
337        assert_eq!(by_label.key, key);
338    }
339
340    #[test]
341    fn test_message_history() {
342        let mut manager = SessionManager::new();
343        let key = manager.spawn_subagent("main", "Test", None, None);
344
345        manager.send_message(&key, "Hello").unwrap();
346
347        let session = manager.get_mut(&key).unwrap();
348        session.add_message("assistant", "Hi there!");
349
350        let history = manager.history(&key, 10, false).unwrap();
351        assert_eq!(history.len(), 2);
352        assert_eq!(history[0].content, "Hello");
353        assert_eq!(history[1].content, "Hi there!");
354    }
355
356    #[test]
357    fn test_session_listing() {
358        let mut manager = SessionManager::new();
359        manager.get_or_create_main("main");
360        manager.spawn_subagent("main", "Task 1", None, None);
361        manager.spawn_subagent("main", "Task 2", None, None);
362
363        let all = manager.list(None, false, 10);
364        assert_eq!(all.len(), 3);
365
366        let subagents = manager.list(Some(&[SessionKind::Subagent]), false, 10);
367        assert_eq!(subagents.len(), 2);
368    }
369}
rustyclaw_core/sessions.rs

rustyclaw_core/
sessions.rs
