1use std::fmt;
2use std::path::{Path, PathBuf};
3use std::sync::atomic::{AtomicU64, Ordering};
4use std::sync::{Arc, Mutex};
5use std::time::{SystemTime, UNIX_EPOCH};
6
7use rusqlite::{params, Connection, OptionalExtension};
8
9use crate::app_server::protocol::{AppItem, AppThread, ThreadStartParams, ThreadStatus};
10
11#[derive(Clone)]
12pub struct SqliteThreadStore {
13 connection: Arc<Mutex<Connection>>,
14 next_thread_id: Arc<AtomicU64>,
15}
16
17impl SqliteThreadStore {
18 pub fn in_memory() -> Result<Self, ThreadStoreError> {
19 let connection = Connection::open_in_memory().map_err(ThreadStoreError::sql)?;
20 Self::from_connection(connection)
21 }
22
23 pub fn open(path: impl AsRef<Path>) -> Result<Self, ThreadStoreError> {
24 let connection = Connection::open(path).map_err(ThreadStoreError::sql)?;
25 Self::from_connection(connection)
26 }
27
28 fn from_connection(connection: Connection) -> Result<Self, ThreadStoreError> {
29 let store = Self {
30 connection: Arc::new(Mutex::new(connection)),
31 next_thread_id: Arc::new(AtomicU64::new(1)),
32 };
33 store.migrate()?;
34 Ok(store)
35 }
36
37 pub fn create_thread(&self, params: ThreadStartParams) -> Result<AppThread, ThreadStoreError> {
38 let sequence = self.next_thread_id.fetch_add(1, Ordering::Relaxed);
39 let now = timestamp_millis();
40 let thread = AppThread {
41 id: format!("thread_{sequence}"),
42 title: params.title,
43 cwd: params.cwd,
44 model: params.model,
45 status: ThreadStatus::Idle,
46 archived: false,
47 ephemeral: params.ephemeral,
48 created_at_ms: now,
49 updated_at_ms: now + sequence as u128,
50 active_turn_id: None,
51 };
52 self.insert_thread(&thread)?;
53 Ok(thread)
54 }
55
56 pub fn get_thread(&self, thread_id: &str) -> Result<Option<AppThread>, ThreadStoreError> {
57 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
58 connection
59 .query_row(
60 "SELECT id, title, cwd, model, status, archived, ephemeral, created_at_ms, updated_at_ms, active_turn_id
61 FROM threads WHERE id = ?1",
62 params![thread_id],
63 row_to_thread,
64 )
65 .optional()
66 .map_err(ThreadStoreError::sql)
67 }
68
69 pub fn list_threads(&self, include_archived: bool) -> Result<Vec<AppThread>, ThreadStoreError> {
70 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
71 let sql = if include_archived {
72 "SELECT id, title, cwd, model, status, archived, ephemeral, created_at_ms, updated_at_ms, active_turn_id
73 FROM threads ORDER BY updated_at_ms DESC, id DESC"
74 } else {
75 "SELECT id, title, cwd, model, status, archived, ephemeral, created_at_ms, updated_at_ms, active_turn_id
76 FROM threads WHERE archived = 0 ORDER BY updated_at_ms DESC, id DESC"
77 };
78 let mut statement = connection.prepare(sql).map_err(ThreadStoreError::sql)?;
79 let rows = statement
80 .query_map([], row_to_thread)
81 .map_err(ThreadStoreError::sql)?;
82 let mut threads = Vec::new();
83 for row in rows {
84 threads.push(row.map_err(ThreadStoreError::sql)?);
85 }
86 Ok(threads)
87 }
88
89 pub fn archive_thread(&self, thread_id: &str) -> Result<(), ThreadStoreError> {
90 let now = timestamp_millis();
91 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
92 connection
93 .execute(
94 "UPDATE threads
95 SET archived = 1, status = 'archived', updated_at_ms = ?2
96 WHERE id = ?1",
97 params![thread_id, now as i64],
98 )
99 .map_err(ThreadStoreError::sql)?;
100 Ok(())
101 }
102
103 pub fn set_active_turn(
104 &self,
105 thread_id: &str,
106 active_turn_id: Option<&str>,
107 status: ThreadStatus,
108 ) -> Result<(), ThreadStoreError> {
109 let now = timestamp_millis();
110 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
111 connection
112 .execute(
113 "UPDATE threads
114 SET active_turn_id = ?2, status = ?3, updated_at_ms = ?4
115 WHERE id = ?1",
116 params![
117 thread_id,
118 active_turn_id,
119 thread_status_to_str(status),
120 now as i64
121 ],
122 )
123 .map_err(ThreadStoreError::sql)?;
124 Ok(())
125 }
126
127 pub fn append_item(
128 &self,
129 thread_id: &str,
130 turn_id: &str,
131 item: AppItem,
132 ) -> Result<(), ThreadStoreError> {
133 let payload_json = serde_json::to_string(&item).map_err(ThreadStoreError::json)?;
134 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
135 let sequence: i64 = connection
136 .query_row(
137 "SELECT COALESCE(MAX(sequence), 0) + 1 FROM thread_items WHERE thread_id = ?1",
138 params![thread_id],
139 |row| row.get(0),
140 )
141 .map_err(ThreadStoreError::sql)?;
142 connection
143 .execute(
144 "INSERT INTO thread_items (id, thread_id, turn_id, run_event_id, sequence, payload_json)
145 VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
146 params![
147 item.id,
148 thread_id,
149 turn_id,
150 item.run_event_id,
151 sequence,
152 payload_json
153 ],
154 )
155 .map_err(ThreadStoreError::sql)?;
156 Ok(())
157 }
158
159 pub fn replay_items(&self, thread_id: &str) -> Result<Vec<AppItem>, ThreadStoreError> {
160 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
161 let mut statement = connection
162 .prepare(
163 "SELECT payload_json FROM thread_items
164 WHERE thread_id = ?1
165 ORDER BY sequence ASC",
166 )
167 .map_err(ThreadStoreError::sql)?;
168 let rows = statement
169 .query_map(params![thread_id], |row| row.get::<_, String>(0))
170 .map_err(ThreadStoreError::sql)?;
171 let mut items = Vec::new();
172 for row in rows {
173 let payload_json = row.map_err(ThreadStoreError::sql)?;
174 items.push(serde_json::from_str(&payload_json).map_err(ThreadStoreError::json)?);
175 }
176 Ok(items)
177 }
178
179 fn insert_thread(&self, thread: &AppThread) -> Result<(), ThreadStoreError> {
180 let cwd = thread
181 .cwd
182 .as_ref()
183 .map(|path| path_to_string(path.as_path()));
184 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
185 connection
186 .execute(
187 "INSERT INTO threads (
188 id, title, cwd, model, status, archived, ephemeral,
189 created_at_ms, updated_at_ms, active_turn_id
190 )
191 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
192 params![
193 thread.id,
194 thread.title,
195 cwd,
196 thread.model,
197 thread_status_to_str(thread.status),
198 bool_to_i64(thread.archived),
199 bool_to_i64(thread.ephemeral),
200 thread.created_at_ms as i64,
201 thread.updated_at_ms as i64,
202 thread.active_turn_id,
203 ],
204 )
205 .map_err(ThreadStoreError::sql)?;
206 Ok(())
207 }
208
209 fn migrate(&self) -> Result<(), ThreadStoreError> {
210 let connection = self.connection.lock().map_err(ThreadStoreError::poisoned)?;
211 connection
212 .execute_batch(
213 r#"
214 CREATE TABLE IF NOT EXISTS threads (
215 id TEXT PRIMARY KEY,
216 title TEXT,
217 cwd TEXT,
218 model TEXT,
219 status TEXT NOT NULL,
220 archived INTEGER NOT NULL,
221 ephemeral INTEGER NOT NULL,
222 created_at_ms INTEGER NOT NULL,
223 updated_at_ms INTEGER NOT NULL,
224 active_turn_id TEXT
225 );
226
227 CREATE TABLE IF NOT EXISTS thread_items (
228 id TEXT PRIMARY KEY,
229 thread_id TEXT NOT NULL,
230 turn_id TEXT NOT NULL,
231 run_event_id TEXT NOT NULL,
232 sequence INTEGER NOT NULL,
233 payload_json TEXT NOT NULL
234 );
235
236 CREATE INDEX IF NOT EXISTS idx_thread_items_thread_sequence
237 ON thread_items(thread_id, sequence);
238 "#,
239 )
240 .map_err(ThreadStoreError::sql)?;
241 Ok(())
242 }
243}
244
245#[derive(Debug, Clone, PartialEq, Eq)]
246pub struct ThreadStoreError {
247 message: String,
248}
249
250impl ThreadStoreError {
251 fn sql(error: rusqlite::Error) -> Self {
252 Self {
253 message: error.to_string(),
254 }
255 }
256
257 fn json(error: serde_json::Error) -> Self {
258 Self {
259 message: error.to_string(),
260 }
261 }
262
263 fn poisoned<T>(_: std::sync::PoisonError<T>) -> Self {
264 Self {
265 message: "thread store lock poisoned".to_string(),
266 }
267 }
268}
269
270impl fmt::Display for ThreadStoreError {
271 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
272 formatter.write_str(&self.message)
273 }
274}
275
276impl std::error::Error for ThreadStoreError {}
277
278fn row_to_thread(row: &rusqlite::Row<'_>) -> rusqlite::Result<AppThread> {
279 let cwd: Option<String> = row.get("cwd")?;
280 let status: String = row.get("status")?;
281 let created_at_ms: i64 = row.get("created_at_ms")?;
282 let updated_at_ms: i64 = row.get("updated_at_ms")?;
283 Ok(AppThread {
284 id: row.get("id")?,
285 title: row.get("title")?,
286 cwd: cwd.map(PathBuf::from),
287 model: row.get("model")?,
288 status: thread_status_from_str(&status).map_err(|message| {
289 rusqlite::Error::FromSqlConversionFailure(
290 0,
291 rusqlite::types::Type::Text,
292 Box::new(ThreadStoreError { message }),
293 )
294 })?,
295 archived: row.get::<_, i64>("archived")? != 0,
296 ephemeral: row.get::<_, i64>("ephemeral")? != 0,
297 created_at_ms: created_at_ms as u128,
298 updated_at_ms: updated_at_ms as u128,
299 active_turn_id: row.get("active_turn_id")?,
300 })
301}
302
303fn thread_status_to_str(status: ThreadStatus) -> &'static str {
304 match status {
305 ThreadStatus::Idle => "idle",
306 ThreadStatus::Running => "running",
307 ThreadStatus::Archived => "archived",
308 }
309}
310
311fn thread_status_from_str(status: &str) -> Result<ThreadStatus, String> {
312 match status {
313 "idle" => Ok(ThreadStatus::Idle),
314 "running" => Ok(ThreadStatus::Running),
315 "archived" => Ok(ThreadStatus::Archived),
316 other => Err(format!("unknown thread status: {other}")),
317 }
318}
319
320fn path_to_string(path: &Path) -> String {
321 path.to_string_lossy().to_string()
322}
323
324fn bool_to_i64(value: bool) -> i64 {
325 if value {
326 1
327 } else {
328 0
329 }
330}
331
332fn timestamp_millis() -> u128 {
333 SystemTime::now()
334 .duration_since(UNIX_EPOCH)
335 .map(|duration| duration.as_millis())
336 .unwrap_or_default()
337}