lash_sqlite_store/
lifecycle.rs

1//! [`Store`] open/memory lifecycle plus session head/meta accessors.
2//!
3//! This is one of the two reference modules (with `blobs.rs`) establishing the
4//! tokio-rusqlite translation pattern every other module follows:
5//!
6//! * The async public methods keep the *exact* the prior store signatures.
7//! * A read goes through `self.conn.call(move |c| { ... })`, where the closure
8//!   is a *synchronous* rusqlite body returning `rusqlite::Result<T>`.
9//! * A read-then-write goes through `self.conn.write(move |tx| { ... })`.
10//! * The shared `*_from_conn` helpers in `lib.rs` are synchronous and take a
11//!   `&rusqlite::Connection`, so they can be called from inside either closure.
12//! * Closures must be `'static` + `Send`: capture owned values (clone strings,
13//!   move them in), not borrows of `self`.
14
15use super::*;
16
17impl Store {
18    pub async fn open(path: &Path) -> tokio_rusqlite::Result<Self> {
19        Self::open_with_options(path, StoreOptions::default()).await
20    }
21
22    pub async fn open_with_clock(
23        path: &Path,
24        clock: Arc<dyn lash_core::Clock>,
25    ) -> tokio_rusqlite::Result<Self> {
26        Self::open_with_options_and_clock(path, StoreOptions::default(), clock).await
27    }
28
29    pub async fn open_with_options(
30        path: &Path,
31        options: StoreOptions,
32    ) -> tokio_rusqlite::Result<Self> {
33        Self::open_with_options_and_clock(path, options, Arc::new(lash_core::SystemClock)).await
34    }
35
36    pub async fn open_with_options_and_clock(
37        path: &Path,
38        options: StoreOptions,
39        clock: Arc<dyn lash_core::Clock>,
40    ) -> tokio_rusqlite::Result<Self> {
41        Self::open_with_options_clock_and_process_registry(path, options, clock, None).await
42    }
43
44    pub(crate) async fn open_with_options_clock_and_process_registry(
45        path: &Path,
46        options: StoreOptions,
47        clock: Arc<dyn lash_core::Clock>,
48        process_registry_path: Option<&Path>,
49    ) -> tokio_rusqlite::Result<Self> {
50        let conn = SqliteConnection::open(path).await?;
51        ensure_schema(&conn).await?;
52        apply_pragmas(&conn, StoreBacking::File).await?;
53        let process_registry_attached = if let Some(process_registry_path) = process_registry_path {
54            if !process_registry_path.exists() {
55                return Err(tokio_rusqlite::Error::Error(
56                    rusqlite::Error::SqliteFailure(
57                        rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CANTOPEN),
58                        Some(format!(
59                            "configured Lash process registry does not exist: {}",
60                            process_registry_path.display()
61                        )),
62                    ),
63                ));
64            }
65            let path = process_registry_path.to_string_lossy().into_owned();
66            conn.call(move |conn| {
67                conn.execute("ATTACH DATABASE ?1 AS process_registry", params![path])?;
68                let expected_version = crate::schema::PROCESS_SCHEMA_VERSION;
69                let deadline = std::time::Instant::now()
70                    + std::time::Duration::from_millis(crate::conn::BUSY_TIMEOUT_MS as u64);
71                loop {
72                    let version: i32 = conn.query_row(
73                        "PRAGMA process_registry.user_version",
74                        [],
75                        |row| row.get(0),
76                    )?;
77                    let has_processes = conn
78                        .query_row(
79                            "SELECT 1 FROM process_registry.sqlite_master
80                             WHERE type = 'table' AND name = 'processes'",
81                            [],
82                            |_| Ok(()),
83                        )
84                        .optional()?
85                        .is_some();
86                    if version == expected_version && has_processes {
87                        break;
88                    }
89                    if version == 0 && !has_processes && std::time::Instant::now() < deadline {
90                        std::thread::sleep(std::time::Duration::from_millis(10));
91                        continue;
92                    }
93                    return Err(rusqlite::Error::SqliteFailure(
94                        rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_MISMATCH),
95                        Some(format!(
96                            "configured database is not a Lash process registry: expected schema version {expected_version} with table `processes`, found version {version}"
97                        )),
98                    ));
99                }
100                Ok(())
101            })
102            .await?;
103            true
104        } else {
105            false
106        };
107        Ok(Self {
108            conn,
109            clock,
110            artifact_cache: Mutex::new(BTreeMap::new()),
111            options,
112            commit_count: AtomicU64::new(0),
113            process_registry_attached,
114            #[cfg(test)]
115            checkpoint_probe_count: AtomicUsize::new(0),
116            #[cfg(test)]
117            checkpoint_write_transaction_count: AtomicUsize::new(0),
118        })
119    }
120
121    /// Open the local database read-only. Used by export/resume call sites that
122    /// must never mutate the source.
123    pub async fn open_readonly(path: &Path) -> tokio_rusqlite::Result<Self> {
124        let conn = SqliteConnection::open_readonly(path).await?;
125        Ok(Self {
126            conn,
127            clock: Arc::new(lash_core::SystemClock),
128            artifact_cache: Mutex::new(BTreeMap::new()),
129            options: StoreOptions::default(),
130            commit_count: AtomicU64::new(0),
131            process_registry_attached: false,
132            #[cfg(test)]
133            checkpoint_probe_count: AtomicUsize::new(0),
134            #[cfg(test)]
135            checkpoint_write_transaction_count: AtomicUsize::new(0),
136        })
137    }
138
139    pub async fn load_picker_info(&self) -> Option<SessionPickerInfo> {
140        self.conn
141            .call(|conn| {
142                let meta = conn
143                    .query_row(
144                        "SELECT session_id, cwd, relation_json
145                         FROM session_meta WHERE singleton = 1",
146                        [],
147                        |row| {
148                            let relation_json: Option<String> = row.get(2)?;
149                            let relation = relation_json
150                                .and_then(|json| serde_json::from_str(&json).ok())
151                                .unwrap_or_default();
152                            Ok((
153                                row.get::<_, String>(0)?,
154                                row.get::<_, Option<String>>(1)?,
155                                relation,
156                            ))
157                        },
158                    )
159                    .optional()?;
160                let Some((session_id, cwd, relation)) = meta else {
161                    return Ok(None);
162                };
163
164                let head_json: String = conn
165                    .query_row(
166                        "SELECT head_json FROM session_head WHERE singleton = 1",
167                        [],
168                        |row| row.get(0),
169                    )
170                    .optional()?
171                    .unwrap_or_else(|| "{}".to_string());
172                let head_meta =
173                    serde_json::from_str::<SessionHeadMeta>(&head_json).unwrap_or_default();
174                let graph = Self::load_session_graph_from_conn(conn, head_meta.leaf_node_id);
175
176                Ok(Some(SessionPickerInfo {
177                    session_id,
178                    cwd,
179                    relation,
180                    first_user_message: graph.first_user_message(),
181                    user_message_count: graph.user_message_count(),
182                }))
183            })
184            .await
185            .ok()
186            .flatten()
187    }
188
189    pub async fn memory() -> tokio_rusqlite::Result<Self> {
190        Self::memory_with_options(StoreOptions {
191            blob_profile: BuiltinBlobProfile::LowLatency,
192            gc_policy: StoreGcPolicy::default(),
193        })
194        .await
195    }
196
197    pub async fn memory_with_clock(
198        clock: Arc<dyn lash_core::Clock>,
199    ) -> tokio_rusqlite::Result<Self> {
200        Self::memory_with_options_and_clock(
201            StoreOptions {
202                blob_profile: BuiltinBlobProfile::LowLatency,
203                gc_policy: StoreGcPolicy::default(),
204            },
205            clock,
206        )
207        .await
208    }
209
210    pub async fn memory_with_options(options: StoreOptions) -> tokio_rusqlite::Result<Self> {
211        Self::memory_with_options_and_clock(options, Arc::new(lash_core::SystemClock)).await
212    }
213
214    pub async fn memory_with_options_and_clock(
215        options: StoreOptions,
216        clock: Arc<dyn lash_core::Clock>,
217    ) -> tokio_rusqlite::Result<Self> {
218        let conn = SqliteConnection::open_in_memory().await?;
219        ensure_schema(&conn).await?;
220        apply_pragmas(&conn, StoreBacking::Memory).await?;
221        Ok(Self {
222            conn,
223            clock,
224            artifact_cache: Mutex::new(BTreeMap::new()),
225            options,
226            commit_count: AtomicU64::new(0),
227            process_registry_attached: false,
228            #[cfg(test)]
229            checkpoint_probe_count: AtomicUsize::new(0),
230            #[cfg(test)]
231            checkpoint_write_transaction_count: AtomicUsize::new(0),
232        })
233    }
234
235    #[cfg(test)]
236    pub(crate) fn checkpoint_claim_counts(&self) -> (usize, usize) {
237        (
238            self.checkpoint_probe_count
239                .load(std::sync::atomic::Ordering::Relaxed),
240            self.checkpoint_write_transaction_count
241                .load(std::sync::atomic::Ordering::Relaxed),
242        )
243    }
244
245    pub async fn save_session_head_meta(&self, meta: SessionHeadMeta) {
246        let head_json = encode_json(&meta);
247        let session_id = meta.session_id.clone();
248        let head_revision = meta.head_revision as i64;
249        let result = self
250            .conn
251            .call(move |conn| {
252                conn.execute(
253                    "INSERT OR REPLACE INTO session_head (singleton, session_id, head_json, head_revision)
254                     VALUES (1, ?1, ?2, ?3)",
255                    params![session_id, head_json, head_revision],
256                )
257            })
258            .await;
259        if let Err(err) = result {
260            tracing::warn!(error = %err, "failed to persist session head");
261        }
262    }
263
264    pub async fn load_session_head_meta(&self) -> Option<SessionHeadMeta> {
265        self.conn
266            .call(|conn| Ok(load_session_head_meta_from_conn(conn)))
267            .await
268            .ok()
269            .flatten()
270    }
271
272    pub async fn save_session_head(&self, head: SessionHead) {
273        self.replace_session_graph(&head.graph).await;
274        self.save_session_head_meta(session_head_meta(&head)).await;
275    }
276
277    pub async fn load_session_head(&self) -> Option<SessionHead> {
278        let meta = self.load_session_head_meta().await?;
279        let mut graph = self.load_session_graph().await;
280        graph.set_leaf_node_id(meta.leaf_node_id.clone());
281        Some(SessionHead {
282            session_id: meta.session_id,
283            head_revision: meta.head_revision,
284            agent_frames: meta.agent_frames,
285            current_agent_frame_id: meta.current_agent_frame_id,
286            graph,
287            config: meta.config,
288            checkpoint_ref: meta.checkpoint_ref,
289            token_ledger: merge_token_ledger_entries(self.load_usage_deltas().await),
290        })
291    }
292
293    pub async fn head_copy_from_store(&self, source: &Store) {
294        if let Some(head) = source.load_session_head().await {
295            if let Some(checkpoint_ref) = &head.checkpoint_ref
296                && let Some(record) = source
297                    .get_typed_blob::<SessionCheckpoint>(checkpoint_ref)
298                    .await
299            {
300                for blob_ref in [
301                    record.tool_state_ref.as_ref(),
302                    record.plugin_snapshot_ref.as_ref(),
303                ]
304                .into_iter()
305                .flatten()
306                {
307                    if let Some(blob) = source.get_blob(blob_ref).await {
308                        let descriptor = match record
309                            .tool_state_ref
310                            .as_ref()
311                            .filter(|candidate| *candidate == blob_ref)
312                        {
313                            Some(_) => BlobArtifactDescriptor::tool_state_snapshot(),
314                            None => BlobArtifactDescriptor::plugin_session_snapshot(),
315                        };
316                        let _ = self.put_artifact_blob(descriptor, &blob).await;
317                    }
318                }
319                if let Some(blob) = source.get_blob(checkpoint_ref).await {
320                    let _ = self
321                        .put_artifact_blob(BlobArtifactDescriptor::checkpoint_manifest(), &blob)
322                        .await;
323                }
324            }
325            self.replace_session_graph(&head.graph).await;
326            self.save_session_head_meta(session_head_meta(&head)).await;
327        }
328    }
329
330    pub async fn save_session_meta(&self, meta: SessionMeta) {
331        let relation_json = serde_json::to_string(&meta.relation).ok();
332        let session_id_for_log = meta.session_id.clone();
333        let result = self
334            .conn
335            .call(move |conn| {
336                conn.execute(
337                    "INSERT OR REPLACE INTO session_meta
338                     (singleton, session_id, session_name, created_at, model, cwd, relation_json)
339                     VALUES (1, ?1, ?2, ?3, ?4, ?5, ?6)",
340                    params![
341                        meta.session_id,
342                        meta.session_name,
343                        meta.created_at,
344                        meta.model,
345                        meta.cwd,
346                        relation_json
347                    ],
348                )
349            })
350            .await;
351        if let Err(err) = result {
352            tracing::warn!(
353                error = %err,
354                session_id = session_id_for_log,
355                "failed to persist session metadata"
356            );
357        }
358    }
359
360    pub async fn load_session_meta(&self) -> Option<SessionMeta> {
361        self.conn
362            .call(|conn| Ok(load_session_meta_from_conn(conn)))
363            .await
364            .ok()
365            .flatten()
366    }
367}
lash_sqlite_store/lifecycle.rs

lash_sqlite_store/
lifecycle.rs
