1use std::collections::BTreeMap;
6use std::path::Path;
7use std::sync::Mutex;
8
9use async_trait::async_trait;
10use rusqlite::{Connection, OptionalExtension};
11
12use super::contracts::ContinuityStore;
13use super::types::{
14 AgentIdentity, AgentRuntimeId, CheckpointVersion, ContinuityGeneration, ContinuityRecord,
15 ContinuityResolveState, ContinuityStoreError, FencingToken, SessionSnapshot,
16};
17
18pub struct LocalContinuityStore {
23 conn: Mutex<Connection>,
24}
25
26impl LocalContinuityStore {
27 pub fn open(path: impl AsRef<Path>) -> Result<Self, ContinuityStoreError> {
34 let conn =
35 Connection::open(path).map_err(|e| ContinuityStoreError::Io(format!("open: {e}")))?;
36 conn.execute_batch("PRAGMA journal_mode=WAL; PRAGMA busy_timeout=5000;")
37 .map_err(|e| ContinuityStoreError::Io(format!("pragma: {e}")))?;
38 conn.execute_batch(
39 "CREATE TABLE IF NOT EXISTS continuity_records (
40 identity TEXT PRIMARY KEY,
41 agent_runtime_id TEXT NOT NULL,
42 session_id TEXT NOT NULL,
43 generation INTEGER NOT NULL,
44 checkpoint_version INTEGER NOT NULL,
45 fencing_token INTEGER NOT NULL
46 );
47 CREATE TABLE IF NOT EXISTS session_snapshots (
48 session_id TEXT PRIMARY KEY,
49 identity TEXT NOT NULL,
50 generation INTEGER NOT NULL,
51 checkpoint_version INTEGER NOT NULL,
52 fencing_token INTEGER NOT NULL,
53 data BLOB NOT NULL
54 );",
55 )
56 .map_err(|e| ContinuityStoreError::Io(format!("schema: {e}")))?;
57 Ok(Self {
58 conn: Mutex::new(conn),
59 })
60 }
61
62 pub fn in_memory() -> Result<Self, ContinuityStoreError> {
68 let conn = Connection::open_in_memory()
69 .map_err(|e| ContinuityStoreError::Io(format!("in-memory open: {e}")))?;
70 conn.execute_batch(
71 "CREATE TABLE IF NOT EXISTS continuity_records (
72 identity TEXT PRIMARY KEY,
73 agent_runtime_id TEXT NOT NULL,
74 session_id TEXT NOT NULL,
75 generation INTEGER NOT NULL,
76 checkpoint_version INTEGER NOT NULL,
77 fencing_token INTEGER NOT NULL
78 );
79 CREATE TABLE IF NOT EXISTS session_snapshots (
80 session_id TEXT PRIMARY KEY,
81 identity TEXT NOT NULL,
82 generation INTEGER NOT NULL,
83 checkpoint_version INTEGER NOT NULL,
84 fencing_token INTEGER NOT NULL,
85 data BLOB NOT NULL
86 );",
87 )
88 .map_err(|e| ContinuityStoreError::Io(format!("schema: {e}")))?;
89 Ok(Self {
90 conn: Mutex::new(conn),
91 })
92 }
93
94 pub fn max_fencing_token(&self) -> Result<u64, ContinuityStoreError> {
108 let conn = self
109 .conn
110 .lock()
111 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
112 conn.query_row(
113 "SELECT COALESCE(MAX(t), 0) FROM (
114 SELECT MAX(fencing_token) AS t FROM continuity_records
115 UNION ALL
116 SELECT MAX(fencing_token) AS t FROM session_snapshots
117 )",
118 [],
119 |row| row.get::<_, u64>(0),
120 )
121 .map_err(|e| ContinuityStoreError::Io(format!("max_fencing_token: {e}")))
122 }
123}
124
125#[async_trait]
126impl ContinuityStore for LocalContinuityStore {
127 async fn resolve_many(
128 &self,
129 identities: &[AgentIdentity],
130 ) -> Result<BTreeMap<AgentIdentity, ContinuityResolveState>, ContinuityStoreError> {
131 let conn = self
132 .conn
133 .lock()
134 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
135 let mut map = BTreeMap::new();
136 for id in identities {
137 let mut stmt = conn
138 .prepare_cached(
139 "SELECT agent_runtime_id, session_id, generation, checkpoint_version
140 FROM continuity_records WHERE identity = ?1",
141 )
142 .map_err(|e| ContinuityStoreError::Io(format!("prepare: {e}")))?;
143 let row = stmt
144 .query_row(rusqlite::params![id.as_str()], |row| {
145 Ok((
146 row.get::<_, String>(0)?,
147 row.get::<_, String>(1)?,
148 row.get::<_, u64>(2)?,
149 row.get::<_, u64>(3)?,
150 ))
151 })
152 .optional()
153 .map_err(|e| ContinuityStoreError::Io(format!("query: {e}")))?;
154 match row {
155 Some((runtime_id, session_id_str, generation, cpv)) => {
156 let record = ContinuityRecord {
157 identity: id.clone(),
158 agent_runtime_id: AgentRuntimeId::parse(&runtime_id).map_err(|e| {
159 ContinuityStoreError::Corruption(format!(
160 "invalid runtime_id in store: {e}"
161 ))
162 })?,
163 session_id: meerkat_core::types::SessionId::parse(&session_id_str)
164 .map_err(|e| {
165 ContinuityStoreError::Corruption(format!(
166 "invalid session_id in store: {e}"
167 ))
168 })?,
169 generation: ContinuityGeneration::new(generation),
170 checkpoint_version: CheckpointVersion::new(cpv),
171 };
172 map.insert(id.clone(), ContinuityResolveState::Ready { record });
173 }
174 None => {
175 map.insert(id.clone(), ContinuityResolveState::Uninitialized);
176 }
177 }
178 }
179 Ok(map)
180 }
181
182 async fn load_session_snapshot(
183 &self,
184 session_id: &meerkat_core::types::SessionId,
185 ) -> Result<Option<SessionSnapshot>, ContinuityStoreError> {
186 let conn = self
187 .conn
188 .lock()
189 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
190 let mut stmt = conn
191 .prepare_cached("SELECT data FROM session_snapshots WHERE session_id = ?1")
192 .map_err(|e| ContinuityStoreError::Io(format!("prepare: {e}")))?;
193 let row = stmt
194 .query_row(rusqlite::params![session_id.to_string()], |row| {
195 row.get::<_, Vec<u8>>(0)
196 })
197 .optional()
198 .map_err(|e| ContinuityStoreError::Io(format!("query: {e}")))?;
199 Ok(row.map(|data| SessionSnapshot { data }))
200 }
201
202 async fn delete_session_snapshot_if_current_revision(
203 &self,
204 session_id: &meerkat_core::types::SessionId,
205 expected_current_revision: &str,
206 ) -> Result<bool, ContinuityStoreError> {
207 let mut conn = self
208 .conn
209 .lock()
210 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
211 let tx = conn
212 .transaction()
213 .map_err(|e| ContinuityStoreError::Io(format!("begin tx: {e}")))?;
214
215 let data = tx
216 .query_row(
217 "SELECT data FROM session_snapshots WHERE session_id = ?1",
218 rusqlite::params![session_id.to_string()],
219 |row| row.get::<_, Vec<u8>>(0),
220 )
221 .optional()
222 .map_err(|e| ContinuityStoreError::Io(format!("query snapshot: {e}")))?;
223
224 let Some(data) = data else {
225 return Ok(false);
226 };
227 let session: meerkat_core::Session = serde_json::from_slice(&data).map_err(|e| {
228 ContinuityStoreError::Io(format!(
229 "deserialize session snapshot for revision check: {e}"
230 ))
231 })?;
232 let current_revision = meerkat_core::session_store::session_projection_cas_token(&session)
233 .map_err(|e| ContinuityStoreError::Io(e.to_string()))?;
234 if current_revision != expected_current_revision {
235 return Ok(false);
236 }
237
238 let deleted = tx
239 .execute(
240 "DELETE FROM session_snapshots WHERE session_id = ?1",
241 rusqlite::params![session_id.to_string()],
242 )
243 .map_err(|e| ContinuityStoreError::Io(format!("delete snapshot: {e}")))?;
244 tx.commit()
245 .map_err(|e| ContinuityStoreError::Io(format!("commit snapshot delete: {e}")))?;
246 Ok(deleted > 0)
247 }
248
249 async fn save_session_snapshot(
250 &self,
251 identity: &AgentIdentity,
252 session_id: &meerkat_core::types::SessionId,
253 generation: ContinuityGeneration,
254 version: CheckpointVersion,
255 fencing_token: FencingToken,
256 snapshot: &SessionSnapshot,
257 ) -> Result<(), ContinuityStoreError> {
258 let conn = self
259 .conn
260 .lock()
261 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
262
263 let tx = conn
267 .unchecked_transaction()
268 .map_err(|e| ContinuityStoreError::Io(format!("begin tx: {e}")))?;
269
270 let mut stmt = tx
275 .prepare_cached(
276 "SELECT session_id, generation, fencing_token, checkpoint_version
277 FROM continuity_records WHERE identity = ?1",
278 )
279 .map_err(|e| ContinuityStoreError::Io(format!("prepare: {e}")))?;
280 let existing = stmt
281 .query_row(rusqlite::params![identity.as_str()], |row| {
282 Ok((
283 row.get::<_, String>(0)?,
284 row.get::<_, u64>(1)?,
285 row.get::<_, u64>(2)?,
286 row.get::<_, u64>(3)?,
287 ))
288 })
289 .optional()
290 .map_err(|e| ContinuityStoreError::Io(format!("query: {e}")))?;
291
292 drop(stmt);
294
295 let record_was_present = existing.is_some();
296 if let Some((current_session_id, current_generation, current_token, current_version)) =
297 existing
298 {
299 if current_session_id != session_id.to_string()
300 || current_generation != generation.get()
301 {
302 return Err(ContinuityStoreError::NotFound {
303 identity: identity.clone(),
304 });
305 }
306 if fencing_token.get() < current_token {
307 return Err(ContinuityStoreError::StaleFencingToken {
308 identity: identity.clone(),
309 presented: fencing_token,
310 current: FencingToken::new(current_token),
311 });
312 }
313 if version.get() <= current_version {
314 return Err(ContinuityStoreError::StaleCheckpointVersion {
315 identity: identity.clone(),
316 presented: version,
317 current: CheckpointVersion::new(current_version),
318 });
319 }
320 }
321
322 tx.execute(
324 "INSERT INTO session_snapshots (session_id, identity, generation, checkpoint_version, fencing_token, data)
325 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
326 ON CONFLICT(session_id) DO UPDATE SET
327 identity = excluded.identity,
328 generation = excluded.generation,
329 checkpoint_version = excluded.checkpoint_version,
330 fencing_token = excluded.fencing_token,
331 data = excluded.data",
332 rusqlite::params![
333 session_id.to_string(),
334 identity.as_str(),
335 generation.get(),
336 version.get(),
337 fencing_token.get(),
338 &snapshot.data,
339 ],
340 )
341 .map_err(|e| ContinuityStoreError::Io(format!("upsert snapshot: {e}")))?;
342
343 tx.execute(
347 "UPDATE continuity_records
348 SET checkpoint_version = ?1, fencing_token = ?2
349 WHERE identity = ?3 AND session_id = ?4 AND generation = ?5",
350 rusqlite::params![
351 version.get(),
352 fencing_token.get(),
353 identity.as_str(),
354 session_id.to_string(),
355 generation.get(),
356 ],
357 )
358 .map_err(|e| ContinuityStoreError::Io(format!("update continuity after snapshot: {e}")))?;
359 if record_was_present && tx.changes() == 0 {
360 return Err(ContinuityStoreError::NotFound {
361 identity: identity.clone(),
362 });
363 }
364
365 tx.commit()
366 .map_err(|e| ContinuityStoreError::Io(format!("commit tx: {e}")))?;
367
368 Ok(())
369 }
370
371 async fn upsert_continuity_record(
372 &self,
373 record: &ContinuityRecord,
374 fencing_token: FencingToken,
375 ) -> Result<(), ContinuityStoreError> {
376 let conn = self
377 .conn
378 .lock()
379 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
380
381 let mut stmt = conn
383 .prepare_cached("SELECT fencing_token FROM continuity_records WHERE identity = ?1")
384 .map_err(|e| ContinuityStoreError::Io(format!("prepare: {e}")))?;
385 let existing_token = stmt
386 .query_row(rusqlite::params![record.identity.as_str()], |row| {
387 row.get::<_, u64>(0)
388 })
389 .optional()
390 .map_err(|e| ContinuityStoreError::Io(format!("query: {e}")))?;
391
392 if let Some(current) = existing_token
393 && fencing_token.get() < current
394 {
395 return Err(ContinuityStoreError::StaleFencingToken {
396 identity: record.identity.clone(),
397 presented: fencing_token,
398 current: FencingToken::new(current),
399 });
400 }
401
402 conn.execute(
403 "INSERT INTO continuity_records (identity, agent_runtime_id, session_id, generation, checkpoint_version, fencing_token)
404 VALUES (?1, ?2, ?3, ?4, ?5, ?6)
405 ON CONFLICT(identity) DO UPDATE SET
406 agent_runtime_id = excluded.agent_runtime_id,
407 session_id = excluded.session_id,
408 generation = excluded.generation,
409 checkpoint_version = CASE
410 WHEN continuity_records.session_id = excluded.session_id
411 AND continuity_records.generation = excluded.generation
412 THEN MAX(continuity_records.checkpoint_version, excluded.checkpoint_version)
413 ELSE excluded.checkpoint_version
414 END,
415 fencing_token = excluded.fencing_token",
416 rusqlite::params![
417 record.identity.as_str(),
418 record.agent_runtime_id.as_str(),
419 record.session_id.to_string(),
420 record.generation.get(),
421 record.checkpoint_version.get(),
422 fencing_token.get(),
423 ],
424 )
425 .map_err(|e| ContinuityStoreError::Io(format!("upsert record: {e}")))?;
426
427 Ok(())
428 }
429
430 async fn delete_continuity_record(
431 &self,
432 identity: &AgentIdentity,
433 fencing_token: FencingToken,
434 ) -> Result<(), ContinuityStoreError> {
435 let conn = self
436 .conn
437 .lock()
438 .map_err(|e| ContinuityStoreError::Io(format!("lock: {e}")))?;
439
440 let tx = conn
445 .unchecked_transaction()
446 .map_err(|e| ContinuityStoreError::Io(format!("begin tx: {e}")))?;
447
448 let mut stmt = tx
450 .prepare_cached("SELECT fencing_token FROM continuity_records WHERE identity = ?1")
451 .map_err(|e| ContinuityStoreError::Io(format!("prepare: {e}")))?;
452 let existing_token = stmt
453 .query_row(rusqlite::params![identity.as_str()], |row| {
454 row.get::<_, u64>(0)
455 })
456 .optional()
457 .map_err(|e| ContinuityStoreError::Io(format!("query: {e}")))?;
458
459 drop(stmt);
461
462 if let Some(current) = existing_token
463 && fencing_token.get() < current
464 {
465 return Err(ContinuityStoreError::StaleFencingToken {
466 identity: identity.clone(),
467 presented: fencing_token,
468 current: FencingToken::new(current),
469 });
470 }
471
472 tx.execute(
474 "DELETE FROM session_snapshots WHERE identity = ?1",
475 rusqlite::params![identity.as_str()],
476 )
477 .map_err(|e| ContinuityStoreError::Io(format!("delete snapshots: {e}")))?;
478
479 tx.execute(
481 "DELETE FROM continuity_records WHERE identity = ?1",
482 rusqlite::params![identity.as_str()],
483 )
484 .map_err(|e| ContinuityStoreError::Io(format!("delete record: {e}")))?;
485
486 tx.commit()
487 .map_err(|e| ContinuityStoreError::Io(format!("commit tx: {e}")))?;
488
489 Ok(())
490 }
491}
492
493#[cfg(test)]
494#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
495mod tests {
496 use super::*;
497
498 fn record(
499 identity: &AgentIdentity,
500 session_id: &meerkat_core::types::SessionId,
501 ) -> ContinuityRecord {
502 ContinuityRecord {
503 identity: identity.clone(),
504 agent_runtime_id: AgentRuntimeId::parse("rt-001").unwrap(),
505 session_id: session_id.clone(),
506 generation: ContinuityGeneration::new(0),
507 checkpoint_version: CheckpointVersion::new(0),
508 }
509 }
510
511 #[tokio::test]
512 async fn delete_continuity_record_removes_record_and_snapshots_atomically() {
513 let store = LocalContinuityStore::in_memory().expect("in-memory store");
518 let identity = AgentIdentity::parse("triage:main").unwrap();
519 let session_id = meerkat_core::types::SessionId::new();
520
521 store
522 .upsert_continuity_record(&record(&identity, &session_id), FencingToken::new(1))
523 .await
524 .unwrap();
525 store
526 .save_session_snapshot(
527 &identity,
528 &session_id,
529 ContinuityGeneration::new(0),
530 CheckpointVersion::new(1),
531 FencingToken::new(1),
532 &SessionSnapshot {
533 data: vec![1, 2, 3],
534 },
535 )
536 .await
537 .unwrap();
538
539 assert!(
541 store
542 .load_session_snapshot(&session_id)
543 .await
544 .unwrap()
545 .is_some()
546 );
547
548 store
549 .delete_continuity_record(&identity, FencingToken::new(2))
550 .await
551 .unwrap();
552
553 let resolved = store
555 .resolve_many(std::slice::from_ref(&identity))
556 .await
557 .unwrap();
558 assert!(matches!(
559 resolved.get(&identity),
560 Some(ContinuityResolveState::Uninitialized)
561 ));
562 assert!(
564 store
565 .load_session_snapshot(&session_id)
566 .await
567 .unwrap()
568 .is_none()
569 );
570 }
571
572 #[tokio::test]
573 async fn delete_continuity_record_rejects_stale_fencing_token() {
574 let store = LocalContinuityStore::in_memory().expect("in-memory store");
575 let identity = AgentIdentity::parse("triage:main").unwrap();
576 let session_id = meerkat_core::types::SessionId::new();
577 store
578 .upsert_continuity_record(&record(&identity, &session_id), FencingToken::new(5))
579 .await
580 .unwrap();
581
582 let err = store
583 .delete_continuity_record(&identity, FencingToken::new(2))
584 .await
585 .expect_err("stale fencing token must be rejected");
586 assert!(matches!(
587 err,
588 ContinuityStoreError::StaleFencingToken { .. }
589 ));
590
591 let resolved = store
593 .resolve_many(std::slice::from_ref(&identity))
594 .await
595 .unwrap();
596 assert!(!matches!(
597 resolved.get(&identity),
598 Some(ContinuityResolveState::Uninitialized)
599 ));
600 }
601
602 #[tokio::test]
603 async fn max_fencing_token_recovers_high_water_across_tables_and_reopen() {
604 let dir = tempfile::tempdir().unwrap();
605 let path = dir.path().join("continuity.db");
606 let identity = AgentIdentity::parse("identity:parent-1").unwrap();
607 let session_id = meerkat_core::types::SessionId::new();
608 {
609 let store = LocalContinuityStore::open(&path).unwrap();
610 assert_eq!(store.max_fencing_token().unwrap(), 0, "empty store -> 0");
611 store
613 .upsert_continuity_record(&record(&identity, &session_id), FencingToken::new(1))
614 .await
615 .unwrap();
616 store
617 .save_session_snapshot(
618 &identity,
619 &session_id,
620 ContinuityGeneration::new(0),
621 CheckpointVersion::new(1),
622 FencingToken::new(1),
623 &SessionSnapshot {
624 data: vec![1, 2, 3],
625 },
626 )
627 .await
628 .unwrap();
629 store
632 .upsert_continuity_record(&record(&identity, &session_id), FencingToken::new(15))
633 .await
634 .unwrap();
635 assert_eq!(
636 store.max_fencing_token().unwrap(),
637 15,
638 "high-water = MAX over continuity_records (15) and session_snapshots (1)"
639 );
640 }
641 let store = LocalContinuityStore::open(&path).unwrap();
643 assert_eq!(
644 store.max_fencing_token().unwrap(),
645 15,
646 "high-water must persist across reopen"
647 );
648
649 let snap_only = LocalContinuityStore::in_memory().unwrap();
653 let sid = meerkat_core::types::SessionId::new();
654 snap_only
655 .save_session_snapshot(
656 &identity,
657 &sid,
658 ContinuityGeneration::new(0),
659 CheckpointVersion::new(1),
660 FencingToken::new(7),
661 &SessionSnapshot { data: vec![9] },
662 )
663 .await
664 .unwrap();
665 assert_eq!(
666 snap_only.max_fencing_token().unwrap(),
667 7,
668 "high-water must come from session_snapshots when no continuity record is present"
669 );
670 }
671
672 #[tokio::test]
676 async fn lease_fencing_resumes_above_high_water_on_restart() {
677 use super::super::contracts::LeaseProvider;
678 use super::super::local_lease::LocalLeaseProvider;
679 use super::super::types::LeaseAcquireResult;
680
681 let store = LocalContinuityStore::in_memory().unwrap();
682 let identity = AgentIdentity::parse("identity:parent-1").unwrap();
683 let session_id = meerkat_core::types::SessionId::new();
684 store
686 .upsert_continuity_record(&record(&identity, &session_id), FencingToken::new(15))
687 .await
688 .unwrap();
689
690 let high_water = store.max_fencing_token().unwrap();
692 assert_eq!(high_water, 15);
693 let provider = LocalLeaseProvider::with_floor(high_water);
694 let acquired = provider
695 .acquire_leases(std::slice::from_ref(&identity), "rt-restart")
696 .await
697 .unwrap();
698 let token = match acquired.get(&identity) {
699 Some(LeaseAcquireResult::Acquired(grant)) => grant.fencing_token,
700 _ => panic!("expected an acquired lease"),
701 };
702 assert!(
703 token.get() > high_water,
704 "resumed token {} must exceed the high-water {high_water}",
705 token.get()
706 );
707 store
709 .upsert_continuity_record(&record(&identity, &session_id), token)
710 .await
711 .expect("a token resumed above the high-water must be accepted");
712
713 let reset_provider = LocalLeaseProvider::with_floor(0);
715 let reset_acquired = reset_provider
716 .acquire_leases(std::slice::from_ref(&identity), "rt-reset")
717 .await
718 .unwrap();
719 let reset_token = match reset_acquired.get(&identity) {
720 Some(LeaseAcquireResult::Acquired(grant)) => grant.fencing_token,
721 _ => panic!("expected an acquired lease"),
722 };
723 let err = store
724 .upsert_continuity_record(&record(&identity, &session_id), reset_token)
725 .await
726 .expect_err("a reset-to-1 token must be rejected as stale");
727 assert!(matches!(
728 err,
729 ContinuityStoreError::StaleFencingToken { .. }
730 ));
731 }
732}