Skip to main content

ai_agent/
session.rs

1// Source: /data/home/swei/claudecode/openclaudecode/src/commands/session/session.tsx
2use crate::constants::env::system;
3use crate::types::Message;
4use serde::{Deserialize, Serialize};
5use std::path::PathBuf;
6use tokio::fs;
7
8/// Session metadata.
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct SessionMetadata {
11    pub id: String,
12    pub cwd: String,
13    pub model: String,
14    #[serde(rename = "createdAt")]
15    pub created_at: String,
16    #[serde(rename = "updatedAt")]
17    pub updated_at: String,
18    #[serde(rename = "messageCount")]
19    pub message_count: u32,
20    pub summary: Option<String>,
21    pub tag: Option<String>,
22}
23
24/// Session data on disk.
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct SessionData {
27    pub metadata: SessionMetadata,
28    pub messages: Vec<Message>,
29}
30
31/// Get the sessions directory path.
32pub fn get_sessions_dir() -> PathBuf {
33    let home = std::env::var(system::HOME)
34        .or_else(|_| std::env::var(system::USERPROFILE))
35        .unwrap_or_else(|_| "/tmp".to_string());
36    PathBuf::from(home).join(".open-agent-sdk").join("sessions")
37}
38
39/// Get the path for a specific session.
40pub fn get_session_path(session_id: &str) -> PathBuf {
41    get_sessions_dir().join(session_id)
42}
43
44/// Save session to disk.
45pub async fn save_session(
46    session_id: &str,
47    messages: Vec<Message>,
48    metadata: Option<SessionMetadata>,
49) -> Result<(), crate::error::AgentError> {
50    let dir = get_session_path(session_id);
51    fs::create_dir_all(&dir)
52        .await
53        .map_err(crate::error::AgentError::Io)?;
54
55    let cwd = metadata
56        .as_ref()
57        .and_then(|m| Some(m.cwd.clone()))
58        .unwrap_or_else(|| {
59            std::env::current_dir()
60                .unwrap_or_default()
61                .to_string_lossy()
62                .to_string()
63        });
64
65    let model = metadata
66        .as_ref()
67        .and_then(|m| Some(m.model.clone()))
68        .unwrap_or_else(|| "claude-sonnet-4-6".to_string());
69
70    let created_at = metadata
71        .as_ref()
72        .and_then(|m| Some(m.created_at.clone()))
73        .unwrap_or_else(|| chrono::Utc::now().to_rfc3339());
74
75    let summary = metadata.as_ref().and_then(|m| m.summary.clone());
76    let tag = metadata.as_ref().and_then(|m| m.tag.clone());
77
78    let data = SessionData {
79        metadata: SessionMetadata {
80            id: session_id.to_string(),
81            cwd,
82            model,
83            created_at: created_at.clone(),
84            updated_at: chrono::Utc::now().to_rfc3339(),
85            message_count: messages.len() as u32,
86            summary,
87            tag,
88        },
89        messages,
90    };
91
92    let path = dir.join("transcript.json");
93    let json = serde_json::to_string_pretty(&data).map_err(crate::error::AgentError::Json)?;
94    fs::write(&path, json)
95        .await
96        .map_err(crate::error::AgentError::Io)?;
97
98    Ok(())
99}
100
101/// Load session from disk.
102pub async fn load_session(
103    session_id: &str,
104) -> Result<Option<SessionData>, crate::error::AgentError> {
105    let path = get_session_path(session_id).join("transcript.json");
106
107    match fs::read_to_string(&path).await {
108        Ok(content) => {
109            let data: SessionData =
110                serde_json::from_str(&content).map_err(crate::error::AgentError::Json)?;
111            Ok(Some(data))
112        }
113        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(None),
114        Err(e) => Err(crate::error::AgentError::Io(e)),
115    }
116}
117
118/// List all sessions.
119pub async fn list_sessions() -> Result<Vec<SessionMetadata>, crate::error::AgentError> {
120    let dir = get_sessions_dir();
121
122    let mut entries = match fs::read_dir(&dir).await {
123        Ok(entries) => entries,
124        Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(vec![]),
125        Err(e) => return Err(crate::error::AgentError::Io(e)),
126    };
127
128    let mut sessions = Vec::new();
129
130    while let Some(entry) = entries
131        .next_entry()
132        .await
133        .map_err(crate::error::AgentError::Io)?
134    {
135        let entry_id = entry.file_name().to_string_lossy().to_string();
136        if let Ok(Some(data)) = load_session(&entry_id).await {
137            if let Some(metadata) = Some(data.metadata) {
138                sessions.push(metadata);
139            }
140        }
141    }
142
143    // Sort by updatedAt descending
144    sessions.sort_by(|a, b| b.updated_at.cmp(&a.updated_at));
145
146    Ok(sessions)
147}
148
149/// Fork a session (create a copy with a new ID).
150pub async fn fork_session(
151    source_session_id: &str,
152    new_session_id: Option<String>,
153) -> Result<Option<String>, crate::error::AgentError> {
154    let data = match load_session(source_session_id).await? {
155        Some(d) => d,
156        None => return Ok(None),
157    };
158
159    let fork_id = new_session_id.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
160
161    save_session(
162        &fork_id,
163        data.messages,
164        Some(SessionMetadata {
165            id: fork_id.clone(),
166            cwd: data.metadata.cwd,
167            model: data.metadata.model,
168            created_at: chrono::Utc::now().to_rfc3339(),
169            updated_at: chrono::Utc::now().to_rfc3339(),
170            message_count: data.metadata.message_count,
171            summary: Some(format!("Forked from session {}", source_session_id)),
172            tag: None,
173        }),
174    )
175    .await?;
176
177    Ok(Some(fork_id))
178}
179
180/// Get session messages.
181pub async fn get_session_messages(
182    session_id: &str,
183) -> Result<Vec<Message>, crate::error::AgentError> {
184    match load_session(session_id).await? {
185        Some(data) => Ok(data.messages),
186        None => Ok(vec![]),
187    }
188}
189
190/// Append a message to a session transcript.
191pub async fn append_to_session(
192    session_id: &str,
193    message: Message,
194) -> Result<(), crate::error::AgentError> {
195    let mut data = match load_session(session_id).await? {
196        Some(d) => d,
197        None => return Ok(()),
198    };
199
200    data.messages.push(message);
201    data.metadata.updated_at = chrono::Utc::now().to_rfc3339();
202    data.metadata.message_count = data.messages.len() as u32;
203
204    save_session(session_id, data.messages, Some(data.metadata)).await
205}
206
207/// Delete a session.
208pub async fn delete_session(session_id: &str) -> Result<bool, crate::error::AgentError> {
209    let path = get_session_path(session_id);
210
211    match fs::remove_dir_all(&path).await {
212        Ok(_) => Ok(true),
213        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
214        Err(e) => Err(crate::error::AgentError::Io(e)),
215    }
216}
217
218/// Get info about a specific session.
219pub async fn get_session_info(
220    session_id: &str,
221) -> Result<Option<SessionMetadata>, crate::error::AgentError> {
222    match load_session(session_id).await? {
223        Some(data) => Ok(Some(data.metadata)),
224        None => Ok(None),
225    }
226}
227
228/// Rename a session.
229pub async fn rename_session(session_id: &str, title: &str) -> Result<(), crate::error::AgentError> {
230    let mut data = match load_session(session_id).await? {
231        Some(d) => d,
232        None => return Ok(()),
233    };
234
235    data.metadata.summary = Some(title.to_string());
236    data.metadata.updated_at = chrono::Utc::now().to_rfc3339();
237
238    save_session(session_id, data.messages, Some(data.metadata)).await
239}
240
241/// Tag a session.
242pub async fn tag_session(
243    session_id: &str,
244    tag: Option<&str>,
245) -> Result<(), crate::error::AgentError> {
246    let mut data = match load_session(session_id).await? {
247        Some(d) => d,
248        None => return Ok(()),
249    };
250
251    data.metadata.tag = tag.map(|s| s.to_string());
252    data.metadata.updated_at = chrono::Utc::now().to_rfc3339();
253
254    save_session(session_id, data.messages, Some(data.metadata)).await
255}
256
257#[cfg(test)]
258mod tests {
259    use super::*;
260    use crate::types::MessageRole;
261
262    fn create_test_message(content: &str) -> Message {
263        Message {
264            role: MessageRole::User,
265            content: content.to_string(),
266            ..Default::default()
267        }
268    }
269
270    #[tokio::test]
271    async fn test_get_sessions_dir() {
272        let dir = get_sessions_dir();
273        assert!(dir.to_string_lossy().contains(".open-agent-sdk"));
274    }
275
276    #[tokio::test]
277    async fn test_save_and_load_session() {
278        let session_id = "test-session-1";
279        let messages = vec![create_test_message("Hello")];
280
281        // Save
282        save_session(session_id, messages.clone(), None)
283            .await
284            .unwrap();
285
286        // Load
287        let loaded = load_session(session_id).await.unwrap();
288        assert!(loaded.is_some());
289        assert_eq!(loaded.unwrap().messages.len(), 1);
290
291        // Cleanup
292        delete_session(session_id).await.unwrap();
293    }
294
295    #[tokio::test]
296    async fn test_load_nonexistent_session() {
297        let loaded = load_session("nonexistent-session").await.unwrap();
298        assert!(loaded.is_none());
299    }
300
301    #[tokio::test]
302    async fn test_fork_session() {
303        let source_id = "fork-source-test";
304        let messages = vec![
305            create_test_message("First"),
306            Message {
307                role: MessageRole::Assistant,
308                content: "Response".to_string(),
309                ..Default::default()
310            },
311        ];
312
313        // Save original
314        save_session(source_id, messages, None).await.unwrap();
315
316        // Fork
317        let fork_id = fork_session(source_id, None).await.unwrap();
318        assert!(fork_id.is_some());
319
320        // Verify fork has messages
321        let fork_messages = get_session_messages(fork_id.as_ref().unwrap())
322            .await
323            .unwrap();
324        assert_eq!(fork_messages.len(), 2);
325
326        // Cleanup
327        delete_session(source_id).await.unwrap();
328        delete_session(fork_id.as_ref().unwrap()).await.unwrap();
329    }
330
331    #[tokio::test]
332    async fn test_append_to_session() {
333        let session_id = "append-test-session";
334
335        // Create with initial message
336        save_session(session_id, vec![create_test_message("Initial")], None)
337            .await
338            .unwrap();
339
340        // Append
341        append_to_session(
342            session_id,
343            Message {
344                role: MessageRole::Assistant,
345                content: "Response".to_string(),
346                ..Default::default()
347            },
348        )
349        .await
350        .unwrap();
351
352        // Verify
353        let loaded = load_session(session_id).await.unwrap().unwrap();
354        assert_eq!(loaded.messages.len(), 2);
355
356        // Cleanup
357        delete_session(session_id).await.unwrap();
358    }
359
360    #[tokio::test]
361    async fn test_rename_session() {
362        let session_id = "rename-test-session";
363        save_session(session_id, vec![create_test_message("Test")], None)
364            .await
365            .unwrap();
366
367        rename_session(session_id, "My Session").await.unwrap();
368
369        let info = get_session_info(session_id).await.unwrap().unwrap();
370        assert_eq!(info.summary, Some("My Session".to_string()));
371
372        // Cleanup
373        delete_session(session_id).await.unwrap();
374    }
375
376    #[tokio::test]
377    async fn test_tag_session() {
378        let session_id = "tag-test-session";
379        save_session(session_id, vec![create_test_message("Test")], None)
380            .await
381            .unwrap();
382
383        tag_session(session_id, Some("important")).await.unwrap();
384
385        let info = get_session_info(session_id).await.unwrap().unwrap();
386        assert_eq!(info.tag, Some("important".to_string()));
387
388        // Cleanup
389        delete_session(session_id).await.unwrap();
390    }
391
392    #[tokio::test]
393    async fn test_delete_session() {
394        let session_id = "delete-test-session";
395        save_session(session_id, vec![create_test_message("Test")], None)
396            .await
397            .unwrap();
398
399        let result = delete_session(session_id).await.unwrap();
400        assert!(result);
401
402        // Should not exist now
403        let loaded = load_session(session_id).await.unwrap();
404        assert!(loaded.is_none());
405    }
406}