1use std::path::Path;
2use std::sync::{Arc, Mutex};
3
4use rusqlite::{params, Connection, OptionalExtension};
5
6use crate::types::{Message, MessageRole};
7
8pub trait Session: Send + Sync {
9 fn session_id(&self) -> &str;
10 fn get_items(&self, limit: Option<usize>) -> SessionFuture<Vec<SessionItem>>;
11 fn add_items(&self, items: Vec<SessionItem>) -> SessionFuture<()>;
12 fn pop_item(&self) -> SessionFuture<Option<SessionItem>>;
13 fn clear(&self) -> SessionFuture<()>;
14}
15
16pub type SessionFuture<T> =
17 std::pin::Pin<Box<dyn std::future::Future<Output = Result<T, String>> + Send>>;
18
19#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
20#[serde(tag = "type", rename_all = "snake_case")]
21pub enum SessionItem {
22 User {
23 content: String,
24 },
25 Assistant {
26 content: String,
27 },
28 System {
29 content: String,
30 },
31 Tool {
32 content: String,
33 tool_call_id: String,
34 },
35}
36
37impl SessionItem {
38 pub fn to_message(&self) -> Message {
39 match self {
40 Self::User { content } => Message::user(content.clone()),
41 Self::Assistant { content } => Message::assistant(content.clone()),
42 Self::System { content } => Message::system(content.clone()),
43 Self::Tool {
44 content,
45 tool_call_id,
46 } => Message::tool(content.clone(), tool_call_id.clone()),
47 }
48 }
49
50 pub fn from_message(message: &Message) -> Option<Self> {
51 match message.role {
52 MessageRole::System => Some(Self::System {
53 content: message.content.clone(),
54 }),
55 MessageRole::User => Some(Self::User {
56 content: message.content.clone(),
57 }),
58 MessageRole::Assistant => Some(Self::Assistant {
59 content: message.content.clone(),
60 }),
61 MessageRole::Tool => Some(Self::Tool {
62 content: message.content.clone(),
63 tool_call_id: message.tool_call_id.clone().unwrap_or_default(),
64 }),
65 }
66 }
67}
68
69#[derive(Clone)]
70pub struct MemorySession {
71 session_id: Arc<String>,
72 items: Arc<Mutex<Vec<SessionItem>>>,
73}
74
75impl MemorySession {
76 pub fn new(session_id: impl Into<String>) -> Self {
77 Self {
78 session_id: Arc::new(session_id.into()),
79 items: Arc::new(Mutex::new(Vec::new())),
80 }
81 }
82}
83
84impl Session for MemorySession {
85 fn session_id(&self) -> &str {
86 self.session_id.as_str()
87 }
88
89 fn get_items(&self, limit: Option<usize>) -> SessionFuture<Vec<SessionItem>> {
90 let items = self.items.clone();
91 Box::pin(async move {
92 let items = items
93 .lock()
94 .map_err(|_| "session lock poisoned".to_string())?;
95 let values = match limit {
96 Some(limit) => items
97 .iter()
98 .rev()
99 .take(limit)
100 .cloned()
101 .collect::<Vec<_>>()
102 .into_iter()
103 .rev()
104 .collect(),
105 None => items.clone(),
106 };
107 Ok(values)
108 })
109 }
110
111 fn add_items(&self, new_items: Vec<SessionItem>) -> SessionFuture<()> {
112 let items = self.items.clone();
113 Box::pin(async move {
114 items
115 .lock()
116 .map_err(|_| "session lock poisoned".to_string())?
117 .extend(new_items);
118 Ok(())
119 })
120 }
121
122 fn pop_item(&self) -> SessionFuture<Option<SessionItem>> {
123 let items = self.items.clone();
124 Box::pin(async move {
125 Ok(items
126 .lock()
127 .map_err(|_| "session lock poisoned".to_string())?
128 .pop())
129 })
130 }
131
132 fn clear(&self) -> SessionFuture<()> {
133 let items = self.items.clone();
134 Box::pin(async move {
135 items
136 .lock()
137 .map_err(|_| "session lock poisoned".to_string())?
138 .clear();
139 Ok(())
140 })
141 }
142}
143
144pub trait SessionStore: Send + Sync {
145 fn session(&self, session_id: &str) -> Arc<dyn Session>;
146}
147
148#[derive(Clone)]
149pub struct SqliteSessionStore {
150 connection: Arc<Mutex<Connection>>,
151}
152
153impl SqliteSessionStore {
154 pub fn open_memory() -> Result<Self, String> {
155 Self::open(":memory:")
156 }
157
158 pub fn open(path: impl AsRef<Path>) -> Result<Self, String> {
159 let connection = Connection::open(path).map_err(sqlite_error)?;
160 connection
161 .execute_batch(
162 r#"
163 PRAGMA journal_mode=WAL;
164 CREATE TABLE IF NOT EXISTS session_items (
165 id INTEGER PRIMARY KEY AUTOINCREMENT,
166 session_id TEXT NOT NULL,
167 item_json TEXT NOT NULL
168 );
169 CREATE INDEX IF NOT EXISTS idx_session_items_session_id_id
170 ON session_items (session_id, id);
171 "#,
172 )
173 .map_err(sqlite_error)?;
174 Ok(Self {
175 connection: Arc::new(Mutex::new(connection)),
176 })
177 }
178
179 pub fn session(&self, session_id: &str) -> Arc<dyn Session> {
180 <Self as SessionStore>::session(self, session_id)
181 }
182}
183
184impl SessionStore for SqliteSessionStore {
185 fn session(&self, session_id: &str) -> Arc<dyn Session> {
186 Arc::new(SqliteSession {
187 session_id: Arc::new(session_id.to_string()),
188 connection: self.connection.clone(),
189 })
190 }
191}
192
193#[derive(Clone)]
194struct SqliteSession {
195 session_id: Arc<String>,
196 connection: Arc<Mutex<Connection>>,
197}
198
199impl Session for SqliteSession {
200 fn session_id(&self) -> &str {
201 self.session_id.as_str()
202 }
203
204 fn get_items(&self, limit: Option<usize>) -> SessionFuture<Vec<SessionItem>> {
205 let session_id = self.session_id.to_string();
206 let connection = self.connection.clone();
207 Box::pin(async move {
208 let connection = connection
209 .lock()
210 .map_err(|_| "sqlite session store lock poisoned".to_string())?;
211 let mut statement = if limit.is_some() {
212 connection
213 .prepare(
214 r#"
215 SELECT item_json
216 FROM (
217 SELECT id, item_json
218 FROM session_items
219 WHERE session_id = ?1
220 ORDER BY id DESC
221 LIMIT ?2
222 )
223 ORDER BY id ASC
224 "#,
225 )
226 .map_err(sqlite_error)?
227 } else {
228 connection
229 .prepare(
230 r#"
231 SELECT item_json
232 FROM session_items
233 WHERE session_id = ?1
234 ORDER BY id ASC
235 "#,
236 )
237 .map_err(sqlite_error)?
238 };
239 let mut rows = if let Some(limit) = limit {
240 statement
241 .query(params![session_id, limit as i64])
242 .map_err(sqlite_error)?
243 } else {
244 statement.query(params![session_id]).map_err(sqlite_error)?
245 };
246 let mut items = Vec::new();
247 while let Some(row) = rows.next().map_err(sqlite_error)? {
248 let item_json: String = row.get(0).map_err(sqlite_error)?;
249 items.push(serde_json::from_str(&item_json).map_err(json_error)?);
250 }
251 Ok(items)
252 })
253 }
254
255 fn add_items(&self, items: Vec<SessionItem>) -> SessionFuture<()> {
256 let session_id = self.session_id.to_string();
257 let connection = self.connection.clone();
258 Box::pin(async move {
259 let mut connection = connection
260 .lock()
261 .map_err(|_| "sqlite session store lock poisoned".to_string())?;
262 let transaction = connection.transaction().map_err(sqlite_error)?;
263 for item in items {
264 let item_json = serde_json::to_string(&item).map_err(json_error)?;
265 transaction
266 .execute(
267 "INSERT INTO session_items (session_id, item_json) VALUES (?1, ?2)",
268 params![session_id, item_json],
269 )
270 .map_err(sqlite_error)?;
271 }
272 transaction.commit().map_err(sqlite_error)?;
273 Ok(())
274 })
275 }
276
277 fn pop_item(&self) -> SessionFuture<Option<SessionItem>> {
278 let session_id = self.session_id.to_string();
279 let connection = self.connection.clone();
280 Box::pin(async move {
281 let mut connection = connection
282 .lock()
283 .map_err(|_| "sqlite session store lock poisoned".to_string())?;
284 let transaction = connection.transaction().map_err(sqlite_error)?;
285 let row = transaction
286 .query_row(
287 r#"
288 SELECT id, item_json
289 FROM session_items
290 WHERE session_id = ?1
291 ORDER BY id DESC
292 LIMIT 1
293 "#,
294 params![session_id],
295 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, String>(1)?)),
296 )
297 .optional()
298 .map_err(sqlite_error)?;
299 let Some((id, item_json)) = row else {
300 transaction.commit().map_err(sqlite_error)?;
301 return Ok(None);
302 };
303 transaction
304 .execute("DELETE FROM session_items WHERE id = ?1", params![id])
305 .map_err(sqlite_error)?;
306 transaction.commit().map_err(sqlite_error)?;
307 Ok(Some(serde_json::from_str(&item_json).map_err(json_error)?))
308 })
309 }
310
311 fn clear(&self) -> SessionFuture<()> {
312 let session_id = self.session_id.to_string();
313 let connection = self.connection.clone();
314 Box::pin(async move {
315 connection
316 .lock()
317 .map_err(|_| "sqlite session store lock poisoned".to_string())?
318 .execute(
319 "DELETE FROM session_items WHERE session_id = ?1",
320 params![session_id],
321 )
322 .map_err(sqlite_error)?;
323 Ok(())
324 })
325 }
326}
327
328fn sqlite_error(error: rusqlite::Error) -> String {
329 error.to_string()
330}
331
332fn json_error(error: serde_json::Error) -> String {
333 error.to_string()
334}
335
336pub async fn session_store_conformance(store: &dyn SessionStore) -> Result<(), String> {
337 let session = store.session("conformance-thread");
338 session.clear().await?;
339 session
340 .add_items(vec![SessionItem::User {
341 content: "hello".to_string(),
342 }])
343 .await?;
344 let same_session = store.session("conformance-thread");
345 let items = same_session.get_items(None).await?;
346 if items
347 != vec![SessionItem::User {
348 content: "hello".to_string(),
349 }]
350 {
351 return Err("session store did not persist appended items".to_string());
352 }
353 let popped = same_session.pop_item().await?;
354 if !matches!(popped, Some(SessionItem::User { content }) if content == "hello") {
355 return Err("session store pop_item returned unexpected item".to_string());
356 }
357 if !same_session.get_items(None).await?.is_empty() {
358 return Err("session store pop_item did not remove the item".to_string());
359 }
360 Ok(())
361}