1use std::path::Path;
4
5use async_trait::async_trait;
6use imagegen_bridge_codex_app_server::SessionBindingStore;
7use imagegen_bridge_core::{BridgeError, ErrorCode};
8use serde::{Deserialize, Serialize};
9use tokio_rusqlite::{Connection, params, rusqlite::OpenFlags};
10
11const MAX_IDENTIFIER_BYTES: usize = 512;
12const CURRENT_SCHEMA_VERSION: u32 = 1;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub struct SessionSchemaStatus {
17 pub initialized: bool,
19 pub version: Option<u32>,
21 pub current_version: u32,
23}
24
25pub async fn inspect_sqlite_session_schema(
27 path: impl AsRef<Path>,
28) -> Result<SessionSchemaStatus, BridgeError> {
29 let path = path.as_ref();
30 let metadata = match tokio::fs::symlink_metadata(path).await {
31 Ok(metadata) => metadata,
32 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
33 return Ok(SessionSchemaStatus {
34 initialized: false,
35 version: None,
36 current_version: CURRENT_SCHEMA_VERSION,
37 });
38 }
39 Err(_) => return Err(session_error("could not inspect session database")),
40 };
41 if !metadata.file_type().is_file() {
42 return Err(session_error("session database must be a regular file"));
43 }
44 let connection = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)
45 .await
46 .map_err(|_| session_error("could not open session database read-only"))?;
47 let version = connection
48 .call(|connection| {
49 let table_exists: bool = connection.query_row(
50 "SELECT EXISTS(SELECT 1 FROM sqlite_master WHERE type='table' AND name='schema_migrations')",
51 [],
52 |row| row.get(0),
53 )?;
54 if !table_exists {
55 return Ok(None);
56 }
57 connection.query_row("SELECT MAX(version) FROM schema_migrations", [], |row| {
58 row.get::<_, Option<u32>>(0)
59 })
60 })
61 .await
62 .map_err(|_| session_error("could not inspect session database schema"))?;
63 connection
64 .close()
65 .await
66 .map_err(|_| session_error("could not close session database"))?;
67 Ok(SessionSchemaStatus {
68 initialized: version.is_some(),
69 version,
70 current_version: CURRENT_SCHEMA_VERSION,
71 })
72}
73
74#[derive(Debug)]
76pub struct SqliteSessionBindingStore {
77 connection: Connection,
78 provider: String,
79}
80
81impl SqliteSessionBindingStore {
82 pub async fn open(path: impl AsRef<Path>, provider: &str) -> Result<Self, BridgeError> {
84 validate_identifier(provider, "provider")?;
85 let connection = Connection::open(path)
86 .await
87 .map_err(|_| session_error("could not open session database"))?;
88 connection
89 .call(|connection| {
90 connection.execute_batch(
91 "PRAGMA foreign_keys = ON;
92 PRAGMA journal_mode = WAL;
93 PRAGMA synchronous = FULL;
94 CREATE TABLE IF NOT EXISTS schema_migrations (
95 version INTEGER PRIMARY KEY,
96 applied_at INTEGER NOT NULL
97 );
98 CREATE TABLE IF NOT EXISTS session_bindings (
99 provider TEXT NOT NULL,
100 session_key TEXT NOT NULL,
101 thread_id TEXT NOT NULL,
102 created_at INTEGER NOT NULL,
103 updated_at INTEGER NOT NULL,
104 PRIMARY KEY (provider, session_key)
105 );
106 INSERT OR IGNORE INTO schema_migrations(version, applied_at)
107 VALUES (1, unixepoch());",
108 )?;
109 Ok::<(), tokio_rusqlite::rusqlite::Error>(())
110 })
111 .await
112 .map_err(|_| session_error("could not migrate session database"))?;
113 Ok(Self {
114 connection,
115 provider: provider.to_owned(),
116 })
117 }
118
119 pub async fn close(self) -> Result<(), BridgeError> {
121 self.connection
122 .close()
123 .await
124 .map_err(|_| session_error("could not close session database"))
125 }
126}
127
128#[async_trait]
129impl SessionBindingStore for SqliteSessionBindingStore {
130 async fn get(&self, key: &str) -> Result<Option<String>, BridgeError> {
131 validate_identifier(key, "session key")?;
132 let provider = self.provider.clone();
133 let key = key.to_owned();
134 self.connection
135 .call(
136 move |connection| -> tokio_rusqlite::rusqlite::Result<Option<String>> {
137 let mut statement = connection.prepare_cached(
138 "SELECT thread_id FROM session_bindings
139 WHERE provider = ?1 AND session_key = ?2",
140 )?;
141 let mut rows = statement.query(params![provider, key])?;
142 rows.next()?.map(|row| row.get(0)).transpose()
143 },
144 )
145 .await
146 .map_err(|_| session_error("could not read session binding"))
147 }
148
149 async fn put(&self, key: &str, thread_id: &str) -> Result<(), BridgeError> {
150 validate_identifier(key, "session key")?;
151 validate_identifier(thread_id, "thread ID")?;
152 let provider = self.provider.clone();
153 let key = key.to_owned();
154 let thread_id = thread_id.to_owned();
155 self.connection
156 .call(move |connection| {
157 connection.execute(
158 "INSERT INTO session_bindings(
159 provider, session_key, thread_id, created_at, updated_at
160 ) VALUES (?1, ?2, ?3, unixepoch(), unixepoch())
161 ON CONFLICT(provider, session_key) DO UPDATE SET
162 thread_id = excluded.thread_id,
163 updated_at = unixepoch()",
164 params![provider, key, thread_id],
165 )?;
166 Ok::<(), tokio_rusqlite::rusqlite::Error>(())
167 })
168 .await
169 .map_err(|_| session_error("could not persist session binding"))
170 }
171
172 async fn delete(&self, key: &str) -> Result<(), BridgeError> {
173 validate_identifier(key, "session key")?;
174 let provider = self.provider.clone();
175 let key = key.to_owned();
176 self.connection
177 .call(move |connection| {
178 connection.execute(
179 "DELETE FROM session_bindings WHERE provider = ?1 AND session_key = ?2",
180 params![provider, key],
181 )?;
182 Ok::<(), tokio_rusqlite::rusqlite::Error>(())
183 })
184 .await
185 .map_err(|_| session_error("could not delete session binding"))
186 }
187}
188
189fn validate_identifier(value: &str, label: &str) -> Result<(), BridgeError> {
190 if value.trim().is_empty()
191 || value.len() > MAX_IDENTIFIER_BYTES
192 || value.chars().any(char::is_control)
193 {
194 return Err(session_error(format!("invalid {label}")));
195 }
196 Ok(())
197}
198
199fn session_error(message: impl Into<String>) -> BridgeError {
200 BridgeError::new(ErrorCode::Session, message)
201}
202
203#[cfg(test)]
204mod tests {
205 #![allow(clippy::unwrap_used)]
206
207 use super::*;
208
209 #[tokio::test]
210 async fn schema_inspection_is_read_only_and_reports_migration_version() {
211 let directory = tempfile::tempdir().unwrap();
212 let path = directory.path().join("sessions.sqlite3");
213 let missing = inspect_sqlite_session_schema(&path).await.unwrap();
214 assert!(!missing.initialized);
215 assert!(!path.exists());
216
217 SqliteSessionBindingStore::open(&path, "codex-app-server")
218 .await
219 .unwrap()
220 .close()
221 .await
222 .unwrap();
223 let initialized = inspect_sqlite_session_schema(&path).await.unwrap();
224 assert!(initialized.initialized);
225 assert_eq!(initialized.version, Some(CURRENT_SCHEMA_VERSION));
226 assert_eq!(initialized.current_version, CURRENT_SCHEMA_VERSION);
227 }
228
229 #[tokio::test]
230 async fn binding_survives_database_reopen() {
231 let directory = tempfile::tempdir().unwrap();
232 let path = directory.path().join("sessions.sqlite3");
233 let store = SqliteSessionBindingStore::open(&path, "codex-app-server")
234 .await
235 .unwrap();
236 store.put("gallery", "thread-1").await.unwrap();
237 store.close().await.unwrap();
238
239 let reopened = SqliteSessionBindingStore::open(&path, "codex-app-server")
240 .await
241 .unwrap();
242 assert_eq!(
243 reopened.get("gallery").await.unwrap().as_deref(),
244 Some("thread-1")
245 );
246 reopened.delete("gallery").await.unwrap();
247 assert_eq!(reopened.get("gallery").await.unwrap(), None);
248 }
249
250 #[tokio::test]
251 async fn providers_have_independent_namespaces() {
252 let directory = tempfile::tempdir().unwrap();
253 let path = directory.path().join("sessions.sqlite3");
254 let first = SqliteSessionBindingStore::open(&path, "first")
255 .await
256 .unwrap();
257 let second = SqliteSessionBindingStore::open(&path, "second")
258 .await
259 .unwrap();
260 first.put("same-key", "thread-a").await.unwrap();
261 second.put("same-key", "thread-b").await.unwrap();
262 assert_eq!(
263 first.get("same-key").await.unwrap().as_deref(),
264 Some("thread-a")
265 );
266 assert_eq!(
267 second.get("same-key").await.unwrap().as_deref(),
268 Some("thread-b")
269 );
270 }
271
272 #[tokio::test]
273 async fn rejects_empty_or_control_character_identifiers() {
274 let directory = tempfile::tempdir().unwrap();
275 let store = SqliteSessionBindingStore::open(
276 directory.path().join("sessions.sqlite3"),
277 "codex-app-server",
278 )
279 .await
280 .unwrap();
281 assert!(store.put("", "thread").await.is_err());
282 assert!(store.put("key", "thread\nsecret").await.is_err());
283 }
284}