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

lash_sqlite_store/
lifecycle.rs
