Skip to main content

imagegen_bridge_runtime/
sqlite_sessions.rs

1//! Durable session-to-thread bindings with a deliberately minimal schema.
2
3use 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/// Read-only status of the durable session database schema.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
16pub struct SessionSchemaStatus {
17    /// Whether the database file and migration table exist.
18    pub initialized: bool,
19    /// Highest applied migration when initialized.
20    pub version: Option<u32>,
21    /// Schema version required by this binary.
22    pub current_version: u32,
23}
24
25/// Inspects the session database without creating or migrating it.
26pub 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/// `SQLite`-backed provider session bindings.
75#[derive(Debug)]
76pub struct SqliteSessionBindingStore {
77    connection: Connection,
78    provider: String,
79}
80
81impl SqliteSessionBindingStore {
82    /// Opens the database and applies idempotent schema migrations.
83    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    /// Closes the database worker after pending calls complete.
120    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}