clau_core/
session.rs

1use crate::{Error, Result};
2use serde::{Deserialize, Serialize};
3use std::sync::Arc;
4use tokio::sync::RwLock;
5use std::collections::HashMap;
6
7#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
8pub struct SessionId(String);
9
10impl SessionId {
11    pub fn new(id: impl Into<String>) -> Self {
12        Self(id.into())
13    }
14    
15    pub fn as_str(&self) -> &str {
16        &self.0
17    }
18}
19
20impl std::fmt::Display for SessionId {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        write!(f, "{}", self.0)
23    }
24}
25
26#[derive(Debug, Clone)]
27pub struct Session {
28    pub id: SessionId,
29    pub system_prompt: Option<String>,
30    pub metadata: HashMap<String, serde_json::Value>,
31}
32
33impl Session {
34    pub fn new(id: SessionId) -> Self {
35        Self {
36            id,
37            system_prompt: None,
38            metadata: HashMap::new(),
39        }
40    }
41    
42    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
43        self.system_prompt = Some(prompt.into());
44        self
45    }
46    
47    pub fn id(&self) -> &SessionId {
48        &self.id
49    }
50}
51
52#[derive(Clone)]
53pub struct SessionManager {
54    sessions: Arc<RwLock<HashMap<SessionId, Session>>>,
55}
56
57impl SessionManager {
58    pub fn new() -> Self {
59        Self {
60            sessions: Arc::new(RwLock::new(HashMap::new())),
61        }
62    }
63    
64    pub fn builder() -> SessionBuilder {
65        SessionBuilder::new()
66    }
67    
68    pub async fn create_session(&self) -> SessionBuilder {
69        SessionBuilder::with_manager(self.clone())
70    }
71    
72    pub async fn get(&self, id: &SessionId) -> Result<Option<Session>> {
73        let sessions = self.sessions.read().await;
74        Ok(sessions.get(id).cloned())
75    }
76    
77    pub async fn resume(&self, id: &SessionId) -> Result<Session> {
78        let sessions = self.sessions.read().await;
79        sessions.get(id).cloned()
80            .ok_or_else(|| Error::SessionNotFound(id.to_string()))
81    }
82    
83    pub async fn list(&self) -> Result<Vec<SessionId>> {
84        let sessions = self.sessions.read().await;
85        Ok(sessions.keys().cloned().collect())
86    }
87    
88    async fn store(&self, session: Session) -> Result<()> {
89        let mut sessions = self.sessions.write().await;
90        sessions.insert(session.id.clone(), session);
91        Ok(())
92    }
93}
94
95impl Default for SessionManager {
96    fn default() -> Self {
97        Self::new()
98    }
99}
100
101pub struct SessionBuilder {
102    session: Session,
103    manager: Option<SessionManager>,
104}
105
106impl SessionBuilder {
107    pub fn new() -> Self {
108        let id = SessionId::new(uuid::Uuid::new_v4().to_string());
109        Self {
110            session: Session::new(id),
111            manager: None,
112        }
113    }
114    
115    fn with_manager(manager: SessionManager) -> Self {
116        let id = SessionId::new(uuid::Uuid::new_v4().to_string());
117        Self {
118            session: Session::new(id),
119            manager: Some(manager),
120        }
121    }
122    
123    pub fn with_system_prompt(mut self, prompt: impl Into<String>) -> Self {
124        self.session.system_prompt = Some(prompt.into());
125        self
126    }
127    
128    pub fn with_metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
129        self.session.metadata.insert(key.into(), value);
130        self
131    }
132    
133    pub async fn build(self) -> Result<Session> {
134        if let Some(manager) = self.manager {
135            manager.store(self.session.clone()).await?;
136        }
137        Ok(self.session)
138    }
139}