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}