claude_code_sdk_rust/
session_store.rs

1use async_trait::async_trait;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fmt;
5use std::path::{Component, Path, PathBuf};
6use std::sync::Arc;
7use tokio::sync::Mutex;
8
9use crate::error::Result;
10use crate::session_summary::fold_session_summary;
11
12const MAX_SANITIZED_LENGTH: usize = 200;
13
14#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
15pub struct SessionKey {
16    pub project_key: String,
17    pub session_id: String,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub subpath: Option<String>,
20}
21
22#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
23pub struct SessionListSubkeysKey {
24    pub project_key: String,
25    pub session_id: String,
26}
27
28pub type SessionStoreEntry = serde_json::Map<String, serde_json::Value>;
29
30#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
31pub struct SessionStoreListEntry {
32    pub session_id: String,
33    pub mtime: i64,
34}
35
36#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
37pub struct SessionSummaryEntry {
38    pub session_id: String,
39    pub mtime: i64,
40    pub data: serde_json::Map<String, serde_json::Value>,
41}
42
43#[async_trait]
44pub trait SessionStore: Send + Sync {
45    async fn append(&self, key: SessionKey, entries: Vec<SessionStoreEntry>) -> Result<()>;
46    async fn load(&self, key: SessionKey) -> Result<Option<Vec<SessionStoreEntry>>>;
47
48    fn supports_list_sessions(&self) -> bool {
49        false
50    }
51
52    async fn list_sessions(&self, _project_key: &str) -> Result<Vec<SessionStoreListEntry>> {
53        Err(crate::error::ClaudeSDKError::Session(
54            "SessionStore::list_sessions is not implemented".to_string(),
55        ))
56    }
57
58    async fn list_session_summaries(&self, _project_key: &str) -> Result<Vec<SessionSummaryEntry>> {
59        Err(crate::error::ClaudeSDKError::Session(
60            "SessionStore::list_session_summaries is not implemented".to_string(),
61        ))
62    }
63
64    async fn delete(&self, _key: SessionKey) -> Result<()> {
65        Ok(())
66    }
67
68    async fn list_subkeys(&self, _key: SessionListSubkeysKey) -> Result<Vec<String>> {
69        Ok(Vec::new())
70    }
71}
72
73#[derive(Clone)]
74pub struct SessionStoreHandle(Arc<dyn SessionStore>);
75
76impl SessionStoreHandle {
77    pub fn new<S>(store: S) -> Self
78    where
79        S: SessionStore + 'static,
80    {
81        Self(Arc::new(store))
82    }
83
84    pub async fn append(&self, key: SessionKey, entries: Vec<SessionStoreEntry>) -> Result<()> {
85        self.0.append(key, entries).await
86    }
87
88    pub async fn load(&self, key: SessionKey) -> Result<Option<Vec<SessionStoreEntry>>> {
89        self.0.load(key).await
90    }
91
92    pub fn supports_list_sessions(&self) -> bool {
93        self.0.supports_list_sessions()
94    }
95
96    pub async fn list_sessions(&self, project_key: &str) -> Result<Vec<SessionStoreListEntry>> {
97        self.0.list_sessions(project_key).await
98    }
99
100    pub async fn list_session_summaries(
101        &self,
102        project_key: &str,
103    ) -> Result<Vec<SessionSummaryEntry>> {
104        self.0.list_session_summaries(project_key).await
105    }
106
107    pub async fn delete(&self, key: SessionKey) -> Result<()> {
108        self.0.delete(key).await
109    }
110
111    pub async fn list_subkeys(&self, key: SessionListSubkeysKey) -> Result<Vec<String>> {
112        self.0.list_subkeys(key).await
113    }
114}
115
116impl fmt::Debug for SessionStoreHandle {
117    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
118        f.debug_tuple("SessionStoreHandle")
119            .field(&"<session_store>")
120            .finish()
121    }
122}
123
124#[derive(Debug, Clone, Default)]
125pub struct InMemorySessionStore {
126    state: Arc<Mutex<InMemorySessionStoreState>>,
127}
128
129#[derive(Debug, Default)]
130struct InMemorySessionStoreState {
131    store: HashMap<String, Vec<SessionStoreEntry>>,
132    mtimes: HashMap<String, i64>,
133    summaries: HashMap<(String, String), SessionSummaryEntry>,
134    last_mtime: i64,
135}
136
137impl InMemorySessionStore {
138    pub fn new() -> Self {
139        Self::default()
140    }
141
142    pub async fn get_entries(&self, key: SessionKey) -> Vec<SessionStoreEntry> {
143        let state = self.state.lock().await;
144        state
145            .store
146            .get(&key_to_string(&key))
147            .cloned()
148            .unwrap_or_default()
149    }
150
151    pub async fn size(&self) -> usize {
152        let state = self.state.lock().await;
153        state
154            .store
155            .keys()
156            .filter(|key| {
157                key.find('/')
158                    .is_some_and(|idx| !key[idx + 1..].contains('/'))
159            })
160            .count()
161    }
162
163    pub async fn clear(&self) {
164        let mut state = self.state.lock().await;
165        state.store.clear();
166        state.mtimes.clear();
167        state.summaries.clear();
168        state.last_mtime = 0;
169    }
170}
171
172#[async_trait]
173impl SessionStore for InMemorySessionStore {
174    fn supports_list_sessions(&self) -> bool {
175        true
176    }
177
178    async fn append(&self, key: SessionKey, entries: Vec<SessionStoreEntry>) -> Result<()> {
179        let mut state = self.state.lock().await;
180        let store_key = key_to_string(&key);
181        state
182            .store
183            .entry(store_key.clone())
184            .or_default()
185            .extend(entries.clone());
186        let next = next_mtime(state.last_mtime);
187        state.last_mtime = next;
188        if key.subpath.is_none() {
189            let summary_key = (key.project_key.clone(), key.session_id.clone());
190            let mut summary =
191                fold_session_summary(state.summaries.get(&summary_key), &key, &entries);
192            summary.mtime = next;
193            state.summaries.insert(summary_key, summary);
194        }
195        state.mtimes.insert(store_key, next);
196        Ok(())
197    }
198
199    async fn load(&self, key: SessionKey) -> Result<Option<Vec<SessionStoreEntry>>> {
200        let state = self.state.lock().await;
201        Ok(state.store.get(&key_to_string(&key)).cloned())
202    }
203
204    async fn list_sessions(&self, project_key: &str) -> Result<Vec<SessionStoreListEntry>> {
205        let state = self.state.lock().await;
206        let prefix = format!("{project_key}/");
207        let mut sessions = Vec::new();
208        for key in state.store.keys() {
209            let Some(rest) = key.strip_prefix(&prefix) else {
210                continue;
211            };
212            if !rest.contains('/') {
213                sessions.push(SessionStoreListEntry {
214                    session_id: rest.to_string(),
215                    mtime: *state.mtimes.get(key).unwrap_or(&0),
216                });
217            }
218        }
219        Ok(sessions)
220    }
221
222    async fn delete(&self, key: SessionKey) -> Result<()> {
223        let mut state = self.state.lock().await;
224        let store_key = key_to_string(&key);
225        state.store.remove(&store_key);
226        state.mtimes.remove(&store_key);
227        if key.subpath.is_none() {
228            state
229                .summaries
230                .remove(&(key.project_key.clone(), key.session_id.clone()));
231            let prefix = format!("{}/{}/", key.project_key, key.session_id);
232            let subkeys: Vec<_> = state
233                .store
234                .keys()
235                .filter(|candidate| candidate.starts_with(&prefix))
236                .cloned()
237                .collect();
238            for subkey in subkeys {
239                state.store.remove(&subkey);
240                state.mtimes.remove(&subkey);
241            }
242        }
243        Ok(())
244    }
245
246    async fn list_session_summaries(&self, project_key: &str) -> Result<Vec<SessionSummaryEntry>> {
247        let state = self.state.lock().await;
248        Ok(state
249            .summaries
250            .iter()
251            .filter(|((candidate_project, _), _)| candidate_project == project_key)
252            .map(|(_, summary)| summary.clone())
253            .collect())
254    }
255
256    async fn list_subkeys(&self, key: SessionListSubkeysKey) -> Result<Vec<String>> {
257        let state = self.state.lock().await;
258        let prefix = format!("{}/{}/", key.project_key, key.session_id);
259        Ok(state
260            .store
261            .keys()
262            .filter_map(|store_key| store_key.strip_prefix(&prefix).map(String::from))
263            .collect())
264    }
265}
266
267pub fn project_key_for_directory(directory: Option<&Path>) -> String {
268    sanitize_path(&canonicalize_path(
269        directory.unwrap_or_else(|| Path::new(".")),
270    ))
271}
272
273pub fn file_path_to_session_key(file_path: &Path, projects_dir: &Path) -> Option<SessionKey> {
274    let relative = file_path.strip_prefix(projects_dir).ok()?;
275    let parts = normal_components(relative);
276    if parts.len() < 2 {
277        return None;
278    }
279
280    let project_key = parts[0].clone();
281    let second = &parts[1];
282    if parts.len() == 2 && second.ends_with(".jsonl") {
283        return Some(SessionKey {
284            project_key,
285            session_id: second.trim_end_matches(".jsonl").to_string(),
286            subpath: None,
287        });
288    }
289
290    if parts.len() >= 4 {
291        let mut subpath_parts = parts[2..].to_vec();
292        if let Some(last) = subpath_parts.last_mut() {
293            if last.ends_with(".jsonl") {
294                *last = last.trim_end_matches(".jsonl").to_string();
295            }
296        }
297        return Some(SessionKey {
298            project_key,
299            session_id: second.clone(),
300            subpath: Some(subpath_parts.join("/")),
301        });
302    }
303
304    None
305}
306
307fn key_to_string(key: &SessionKey) -> String {
308    match &key.subpath {
309        Some(subpath) if !subpath.is_empty() => {
310            format!("{}/{}/{}", key.project_key, key.session_id, subpath)
311        }
312        _ => format!("{}/{}", key.project_key, key.session_id),
313    }
314}
315
316fn next_mtime(last_mtime: i64) -> i64 {
317    let now = chrono::Utc::now().timestamp_millis();
318    if now <= last_mtime {
319        last_mtime + 1
320    } else {
321        now
322    }
323}
324
325fn canonicalize_path(path: &Path) -> String {
326    let absolute = if path.is_absolute() {
327        PathBuf::from(path)
328    } else {
329        std::env::current_dir()
330            .unwrap_or_else(|_| PathBuf::from("."))
331            .join(path)
332    };
333    std::fs::canonicalize(&absolute)
334        .unwrap_or(absolute)
335        .to_string_lossy()
336        .to_string()
337}
338
339fn sanitize_path(name: &str) -> String {
340    let sanitized: String = name
341        .chars()
342        .map(|ch| if ch.is_ascii_alphanumeric() { ch } else { '-' })
343        .collect();
344    if sanitized.len() <= MAX_SANITIZED_LENGTH {
345        sanitized
346    } else {
347        format!(
348            "{}-{}",
349            &sanitized[..MAX_SANITIZED_LENGTH],
350            simple_hash(name)
351        )
352    }
353}
354
355fn simple_hash(value: &str) -> String {
356    let mut hash = 0i32;
357    for ch in value.chars() {
358        hash = hash.wrapping_mul(31).wrapping_add(ch as i32);
359    }
360    let mut n = hash.unsigned_abs();
361    if n == 0 {
362        return "0".to_string();
363    }
364    let mut out = Vec::new();
365    while n > 0 {
366        let digit = (n % 36) as u8;
367        out.push(match digit {
368            0..=9 => (b'0' + digit) as char,
369            _ => (b'a' + digit - 10) as char,
370        });
371        n /= 36;
372    }
373    out.iter().rev().collect()
374}
375
376fn normal_components(path: &Path) -> Vec<String> {
377    path.components()
378        .filter_map(|component| match component {
379            Component::Normal(value) => Some(value.to_string_lossy().to_string()),
380            _ => None,
381        })
382        .collect()
383}
384
385#[cfg(test)]
386mod tests {
387    use super::*;
388
389    fn entry(uuid: &str) -> SessionStoreEntry {
390        let mut entry = serde_json::Map::new();
391        entry.insert("type".to_string(), serde_json::json!("user"));
392        entry.insert("uuid".to_string(), serde_json::json!(uuid));
393        entry.insert(
394            "message".to_string(),
395            serde_json::json!({"content": format!("prompt {uuid}")}),
396        );
397        entry
398    }
399
400    #[test]
401    fn derives_session_keys_from_main_and_subagent_paths() {
402        let projects = Path::new("/tmp/claude/projects");
403
404        assert_eq!(
405            file_path_to_session_key(projects.join("proj/session-1.jsonl").as_path(), projects),
406            Some(SessionKey {
407                project_key: "proj".to_string(),
408                session_id: "session-1".to_string(),
409                subpath: None,
410            })
411        );
412        assert_eq!(
413            file_path_to_session_key(
414                projects
415                    .join("proj/session-1/subagents/agent-abc.jsonl")
416                    .as_path(),
417                projects
418            ),
419            Some(SessionKey {
420                project_key: "proj".to_string(),
421                session_id: "session-1".to_string(),
422                subpath: Some("subagents/agent-abc".to_string()),
423            })
424        );
425    }
426
427    #[test]
428    fn sanitizes_project_keys_like_python_sdk() {
429        assert_eq!(
430            sanitize_path("/Users/alice/my project"),
431            "-Users-alice-my-project"
432        );
433        let long = "a".repeat(MAX_SANITIZED_LENGTH + 1);
434        assert!(sanitize_path(&long).starts_with(&"a".repeat(MAX_SANITIZED_LENGTH)));
435        assert_eq!(simple_hash("abc"), "22ci");
436    }
437
438    #[tokio::test]
439    async fn in_memory_store_lists_loads_and_cascades_delete() {
440        let store = InMemorySessionStore::new();
441        let main = SessionKey {
442            project_key: "proj".to_string(),
443            session_id: "session".to_string(),
444            subpath: None,
445        };
446        let sub = SessionKey {
447            subpath: Some("subagents/agent-1".to_string()),
448            ..main.clone()
449        };
450
451        store.append(main.clone(), vec![entry("1")]).await.unwrap();
452        store.append(sub.clone(), vec![entry("2")]).await.unwrap();
453
454        assert_eq!(store.load(main.clone()).await.unwrap().unwrap().len(), 1);
455        assert_eq!(store.list_sessions("proj").await.unwrap().len(), 1);
456        let summaries = store.list_session_summaries("proj").await.unwrap();
457        assert_eq!(summaries.len(), 1);
458        assert_eq!(summaries[0].data["first_prompt"], "prompt 1");
459        assert_eq!(
460            store
461                .list_subkeys(SessionListSubkeysKey {
462                    project_key: "proj".to_string(),
463                    session_id: "session".to_string(),
464                })
465                .await
466                .unwrap(),
467            vec!["subagents/agent-1".to_string()]
468        );
469
470        store.delete(main.clone()).await.unwrap();
471        assert!(store.load(main).await.unwrap().is_none());
472        assert!(store.load(sub).await.unwrap().is_none());
473        assert!(store
474            .list_session_summaries("proj")
475            .await
476            .unwrap()
477            .is_empty());
478    }
479}
claude_code_sdk_rust/session_store.rs

claude_code_sdk_rust/
session_store.rs
