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        let conn = SqliteConnection::open(path).await?;
42        ensure_schema(&conn).await?;
43        apply_pragmas(&conn, StoreBacking::File).await?;
44        Ok(Self {
45            conn,
46            clock,
47            artifact_cache: Mutex::new(BTreeMap::new()),
48            options,
49            commit_count: AtomicU64::new(0),
50            #[cfg(test)]
51            checkpoint_probe_count: AtomicUsize::new(0),
52            #[cfg(test)]
53            checkpoint_write_transaction_count: AtomicUsize::new(0),
54        })
55    }
56
57    /// Open the local database read-only. Used by export/resume call sites that
58    /// must never mutate the source.
59    pub async fn open_readonly(path: &Path) -> tokio_rusqlite::Result<Self> {
60        let conn = SqliteConnection::open_readonly(path).await?;
61        Ok(Self {
62            conn,
63            clock: Arc::new(lash_core::SystemClock),
64            artifact_cache: Mutex::new(BTreeMap::new()),
65            options: StoreOptions::default(),
66            commit_count: AtomicU64::new(0),
67            #[cfg(test)]
68            checkpoint_probe_count: AtomicUsize::new(0),
69            #[cfg(test)]
70            checkpoint_write_transaction_count: AtomicUsize::new(0),
71        })
72    }
73
74    pub async fn load_picker_info(&self) -> Option<SessionPickerInfo> {
75        self.conn
76            .call(|conn| {
77                let meta = conn
78                    .query_row(
79                        "SELECT session_id, cwd, relation_json
80                         FROM session_meta WHERE singleton = 1",
81                        [],
82                        |row| {
83                            let relation_json: Option<String> = row.get(2)?;
84                            let relation = relation_json
85                                .and_then(|json| serde_json::from_str(&json).ok())
86                                .unwrap_or_default();
87                            Ok((
88                                row.get::<_, String>(0)?,
89                                row.get::<_, Option<String>>(1)?,
90                                relation,
91                            ))
92                        },
93                    )
94                    .optional()?;
95                let Some((session_id, cwd, relation)) = meta else {
96                    return Ok(None);
97                };
98
99                let head_json: String = conn
100                    .query_row(
101                        "SELECT head_json FROM session_head WHERE singleton = 1",
102                        [],
103                        |row| row.get(0),
104                    )
105                    .optional()?
106                    .unwrap_or_else(|| "{}".to_string());
107                let head_meta =
108                    serde_json::from_str::<SessionHeadMeta>(&head_json).unwrap_or_default();
109                let graph = Self::load_session_graph_from_conn(conn, head_meta.leaf_node_id);
110
111                Ok(Some(SessionPickerInfo {
112                    session_id,
113                    cwd,
114                    relation,
115                    first_user_message: graph.first_user_message(),
116                    user_message_count: graph.user_message_count(),
117                }))
118            })
119            .await
120            .ok()
121            .flatten()
122    }
123
124    pub async fn memory() -> tokio_rusqlite::Result<Self> {
125        Self::memory_with_options(StoreOptions {
126            blob_profile: BuiltinBlobProfile::LowLatency,
127            gc_policy: StoreGcPolicy::default(),
128        })
129        .await
130    }
131
132    pub async fn memory_with_clock(
133        clock: Arc<dyn lash_core::Clock>,
134    ) -> tokio_rusqlite::Result<Self> {
135        Self::memory_with_options_and_clock(
136            StoreOptions {
137                blob_profile: BuiltinBlobProfile::LowLatency,
138                gc_policy: StoreGcPolicy::default(),
139            },
140            clock,
141        )
142        .await
143    }
144
145    pub async fn memory_with_options(options: StoreOptions) -> tokio_rusqlite::Result<Self> {
146        Self::memory_with_options_and_clock(options, Arc::new(lash_core::SystemClock)).await
147    }
148
149    pub async fn memory_with_options_and_clock(
150        options: StoreOptions,
151        clock: Arc<dyn lash_core::Clock>,
152    ) -> tokio_rusqlite::Result<Self> {
153        let conn = SqliteConnection::open_in_memory().await?;
154        ensure_schema(&conn).await?;
155        apply_pragmas(&conn, StoreBacking::Memory).await?;
156        Ok(Self {
157            conn,
158            clock,
159            artifact_cache: Mutex::new(BTreeMap::new()),
160            options,
161            commit_count: AtomicU64::new(0),
162            #[cfg(test)]
163            checkpoint_probe_count: AtomicUsize::new(0),
164            #[cfg(test)]
165            checkpoint_write_transaction_count: AtomicUsize::new(0),
166        })
167    }
168
169    #[cfg(test)]
170    pub(crate) fn checkpoint_claim_counts(&self) -> (usize, usize) {
171        (
172            self.checkpoint_probe_count
173                .load(std::sync::atomic::Ordering::Relaxed),
174            self.checkpoint_write_transaction_count
175                .load(std::sync::atomic::Ordering::Relaxed),
176        )
177    }
178
179    pub async fn save_session_head_meta(&self, meta: SessionHeadMeta) {
180        let head_json = encode_json(&meta);
181        let session_id = meta.session_id.clone();
182        let head_revision = meta.head_revision as i64;
183        let result = self
184            .conn
185            .call(move |conn| {
186                conn.execute(
187                    "INSERT OR REPLACE INTO session_head (singleton, session_id, head_json, head_revision)
188                     VALUES (1, ?1, ?2, ?3)",
189                    params![session_id, head_json, head_revision],
190                )
191            })
192            .await;
193        if let Err(err) = result {
194            tracing::warn!(error = %err, "failed to persist session head");
195        }
196    }
197
198    pub async fn load_session_head_meta(&self) -> Option<SessionHeadMeta> {
199        self.conn
200            .call(|conn| Ok(load_session_head_meta_from_conn(conn)))
201            .await
202            .ok()
203            .flatten()
204    }
205
206    pub async fn save_session_head(&self, head: SessionHead) {
207        self.replace_session_graph(&head.graph).await;
208        self.save_session_head_meta(session_head_meta(&head)).await;
209    }
210
211    pub async fn load_session_head(&self) -> Option<SessionHead> {
212        let meta = self.load_session_head_meta().await?;
213        let mut graph = self.load_session_graph().await;
214        graph.set_leaf_node_id(meta.leaf_node_id.clone());
215        Some(SessionHead {
216            session_id: meta.session_id,
217            head_revision: meta.head_revision,
218            agent_frames: meta.agent_frames,
219            current_agent_frame_id: meta.current_agent_frame_id,
220            graph,
221            config: meta.config,
222            checkpoint_ref: meta.checkpoint_ref,
223            token_ledger: merge_token_ledger_entries(self.load_usage_deltas().await),
224        })
225    }
226
227    pub async fn head_copy_from_store(&self, source: &Store) {
228        if let Some(head) = source.load_session_head().await {
229            if let Some(checkpoint_ref) = &head.checkpoint_ref
230                && let Some(record) = source
231                    .get_typed_blob::<SessionCheckpoint>(checkpoint_ref)
232                    .await
233            {
234                for blob_ref in [
235                    record.tool_state_ref.as_ref(),
236                    record.plugin_snapshot_ref.as_ref(),
237                ]
238                .into_iter()
239                .flatten()
240                {
241                    if let Some(blob) = source.get_blob(blob_ref).await {
242                        let descriptor = match record
243                            .tool_state_ref
244                            .as_ref()
245                            .filter(|candidate| *candidate == blob_ref)
246                        {
247                            Some(_) => BlobArtifactDescriptor::tool_state_snapshot(),
248                            None => BlobArtifactDescriptor::plugin_session_snapshot(),
249                        };
250                        let _ = self.put_artifact_blob(descriptor, &blob).await;
251                    }
252                }
253                if let Some(blob) = source.get_blob(checkpoint_ref).await {
254                    let _ = self
255                        .put_artifact_blob(BlobArtifactDescriptor::checkpoint_manifest(), &blob)
256                        .await;
257                }
258            }
259            self.replace_session_graph(&head.graph).await;
260            self.save_session_head_meta(session_head_meta(&head)).await;
261        }
262    }
263
264    pub async fn save_session_meta(&self, meta: SessionMeta) {
265        let relation_json = serde_json::to_string(&meta.relation).ok();
266        let session_id_for_log = meta.session_id.clone();
267        let result = self
268            .conn
269            .call(move |conn| {
270                conn.execute(
271                    "INSERT OR REPLACE INTO session_meta
272                     (singleton, session_id, session_name, created_at, model, cwd, relation_json)
273                     VALUES (1, ?1, ?2, ?3, ?4, ?5, ?6)",
274                    params![
275                        meta.session_id,
276                        meta.session_name,
277                        meta.created_at,
278                        meta.model,
279                        meta.cwd,
280                        relation_json
281                    ],
282                )
283            })
284            .await;
285        if let Err(err) = result {
286            tracing::warn!(
287                error = %err,
288                session_id = session_id_for_log,
289                "failed to persist session metadata"
290            );
291        }
292    }
293
294    pub async fn load_session_meta(&self) -> Option<SessionMeta> {
295        self.conn
296            .call(|conn| Ok(load_session_meta_from_conn(conn)))
297            .await
298            .ok()
299            .flatten()
300    }
301}
lash_sqlite_store/lifecycle.rs

lash_sqlite_store/
lifecycle.rs
