1use super::{ReplayEntry, ReplayStore, ReplayStoreError};
48use crate::types::RunId;
49use async_trait::async_trait;
50use rusqlite::params;
51use rusqlite_isle::{AsyncIsle, AsyncIsleDriver, IsleError};
52use std::path::Path;
53
54const SCHEMA_SQL: &str = "\
55CREATE TABLE IF NOT EXISTS replay_log (\
56 seq INTEGER PRIMARY KEY AUTOINCREMENT, \
57 run_id TEXT NOT NULL, \
58 step_ref TEXT NOT NULL, \
59 input_hash TEXT NOT NULL, \
60 occurrence INTEGER NOT NULL, \
61 ctx_snapshot_json TEXT NOT NULL, \
62 step_output_json TEXT NOT NULL, \
63 created_at INTEGER NOT NULL, \
64 UNIQUE (run_id, step_ref, input_hash, occurrence)\
65);\
66CREATE INDEX IF NOT EXISTS ix_replay_run ON replay_log(run_id, seq);\
67";
68
69const CURRENT_SCHEMA_VERSION: i64 = 1;
72
73fn init_schema(conn: &mut rusqlite::Connection) -> rusqlite::Result<()> {
81 let user_version: i64 = conn.query_row("PRAGMA user_version", [], |r| r.get(0))?;
83
84 match user_version {
85 0 => {
86 let table_present: i64 = conn.query_row(
91 "SELECT COUNT(*) FROM sqlite_master \
92 WHERE type = 'table' AND name = 'replay_log'",
93 [],
94 |r| r.get(0),
95 )?;
96 let has_ctx_column: i64 = if table_present > 0 {
97 conn.query_row(
98 "SELECT COUNT(*) FROM pragma_table_info('replay_log') \
99 WHERE name = 'ctx_snapshot_json'",
100 [],
101 |r| r.get(0),
102 )?
103 } else {
104 0
105 };
106 let is_legacy = table_present > 0 && has_ctx_column == 0;
107 if is_legacy {
108 conn.execute("DROP TABLE replay_log", [])?;
109 }
110 conn.execute_batch(SCHEMA_SQL)?;
111 conn.execute_batch(&format!("PRAGMA user_version = {CURRENT_SCHEMA_VERSION}"))?;
114 Ok(())
115 }
116 v if v == CURRENT_SCHEMA_VERSION => {
117 conn.execute_batch(SCHEMA_SQL)
120 }
121 v => Err(rusqlite::Error::SqliteFailure(
122 rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_ERROR),
123 Some(format!(
124 "replay_log schema version {v} is newer than supported \
125 ({CURRENT_SCHEMA_VERSION}); running an older mse binary \
126 against a store written by a newer one is not supported"
127 )),
128 )),
129 }
130}
131
132pub struct SqliteReplayStore {
140 isle: AsyncIsle,
141}
142
143impl SqliteReplayStore {
144 pub async fn open(path: impl AsRef<Path>) -> Result<(Self, AsyncIsleDriver), ReplayStoreError> {
148 let (isle, driver) = AsyncIsle::spawn(path.as_ref().to_path_buf(), init_schema)
149 .await
150 .map_err(map_isle_err)?;
151 Ok((Self { isle }, driver))
152 }
153
154 pub async fn open_in_memory() -> Result<(Self, AsyncIsleDriver), ReplayStoreError> {
159 let (isle, driver) = AsyncIsle::open_in_memory(init_schema)
160 .await
161 .map_err(map_isle_err)?;
162 Ok((Self { isle }, driver))
163 }
164}
165
166fn map_isle_err(e: IsleError) -> ReplayStoreError {
167 ReplayStoreError::Other(format!("sqlite: {e}"))
168}
169
170#[async_trait]
171impl ReplayStore for SqliteReplayStore {
172 fn name(&self) -> &str {
173 "sqlite"
174 }
175
176 async fn append(&self, entry: ReplayEntry) -> Result<(), ReplayStoreError> {
177 let ReplayEntry {
178 run_id,
179 step_ref,
180 input_hash,
181 occurrence,
182 ctx_snapshot_json,
183 step_output_json,
184 created_at,
185 } = entry;
186 let run_id_for_err = run_id.clone();
187 let step_ref_for_err = step_ref.clone();
188 let input_hash_for_err = input_hash.clone();
189 let run_id_str = run_id.to_string();
190
191 self.isle
192 .call(move |conn| {
193 conn.execute(
194 "INSERT INTO replay_log \
195 (run_id, step_ref, input_hash, occurrence, ctx_snapshot_json, \
196 step_output_json, created_at) \
197 VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
198 params![
199 run_id_str,
200 step_ref,
201 input_hash,
202 occurrence as i64,
203 ctx_snapshot_json,
204 step_output_json,
205 created_at as i64,
206 ],
207 )?;
208 Ok(())
209 })
210 .await
211 .map_err(|e| match &e {
212 IsleError::Sqlite(rusqlite::Error::SqliteFailure(err, _))
213 if err.extended_code == rusqlite::ffi::SQLITE_CONSTRAINT_UNIQUE
214 || err.code == rusqlite::ErrorCode::ConstraintViolation =>
215 {
216 ReplayStoreError::Duplicate {
217 run_id: run_id_for_err.clone(),
218 step_ref: step_ref_for_err.clone(),
219 input_hash: input_hash_for_err.clone(),
220 occurrence,
221 }
222 }
223 _ => map_isle_err(e),
224 })
225 }
226
227 async fn list_by_run(&self, run_id: &RunId) -> Result<Vec<ReplayEntry>, ReplayStoreError> {
228 let run_id_str = run_id.to_string();
229 let run_id_owned = run_id.clone();
230 let rows: Vec<(String, String, i64, String, String, i64)> = self
231 .isle
232 .call(move |conn| {
233 let mut stmt = conn.prepare(
234 "SELECT step_ref, input_hash, occurrence, ctx_snapshot_json, \
235 step_output_json, created_at FROM replay_log \
236 WHERE run_id = ?1 ORDER BY seq ASC",
237 )?;
238 let iter = stmt.query_map(params![run_id_str], |row| {
239 Ok((
240 row.get::<_, String>(0)?,
241 row.get::<_, String>(1)?,
242 row.get::<_, i64>(2)?,
243 row.get::<_, String>(3)?,
244 row.get::<_, String>(4)?,
245 row.get::<_, i64>(5)?,
246 ))
247 })?;
248 let mut out = Vec::new();
249 for r in iter {
250 out.push(r?);
251 }
252 Ok(out)
253 })
254 .await
255 .map_err(map_isle_err)?;
256
257 Ok(rows
258 .into_iter()
259 .map(
260 |(
261 step_ref,
262 input_hash,
263 occurrence,
264 ctx_snapshot_json,
265 step_output_json,
266 created_at,
267 )| {
268 ReplayEntry {
269 run_id: run_id_owned.clone(),
270 step_ref,
271 input_hash,
272 occurrence: occurrence as u32,
273 ctx_snapshot_json,
274 step_output_json,
275 created_at: created_at as u64,
276 }
277 },
278 )
279 .collect())
280 }
281
282 async fn delete_from(
283 &self,
284 run_id: &RunId,
285 from_index: usize,
286 ) -> Result<usize, ReplayStoreError> {
287 let run_id_str = run_id.to_string();
288 let deleted: usize = self
289 .isle
290 .call(move |conn| {
291 let mut stmt = conn.prepare(
297 "SELECT seq FROM replay_log WHERE run_id = ?1 \
298 ORDER BY seq ASC LIMIT -1 OFFSET ?2",
299 )?;
300 let seqs: Vec<i64> = stmt
301 .query_map(params![run_id_str, from_index as i64], |row| {
302 row.get::<_, i64>(0)
303 })?
304 .collect::<rusqlite::Result<Vec<_>>>()?;
305 drop(stmt);
306
307 if seqs.is_empty() {
308 return Ok(0usize);
309 }
310
311 let placeholders = std::iter::repeat("?")
315 .take(seqs.len())
316 .collect::<Vec<_>>()
317 .join(",");
318 let sql = format!("DELETE FROM replay_log WHERE seq IN ({placeholders})");
319 let params_dyn: Vec<&dyn rusqlite::ToSql> =
320 seqs.iter().map(|s| s as &dyn rusqlite::ToSql).collect();
321 conn.execute(&sql, params_dyn.as_slice())?;
322 Ok(seqs.len())
323 })
324 .await
325 .map_err(map_isle_err)?;
326 Ok(deleted)
327 }
328}
329
330#[cfg(test)]
335mod tests {
336 use super::*;
337 use crate::core::ctx::Ctx;
338 use crate::store::replay::ReplayCursor;
339 use crate::types::StepId;
340 use serde_json::json;
341
342 fn mk_ctx() -> Ctx {
343 let mut ctx = Ctx::new(StepId::new(), 1, "step-a");
344 ctx.meta.observer.insert("k".into(), json!("v"));
345 ctx
346 }
347
348 #[tokio::test]
349 async fn sqlite_append_and_list() {
350 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
351 let run_id = RunId::new();
352 let ctx = mk_ctx();
353
354 let e0 = ReplayEntry::from_completion(
355 run_id.clone(),
356 "step-a",
357 "hash-a",
358 0,
359 &ctx,
360 &json!({ "n": 1 }),
361 )
362 .unwrap();
363 store.append(e0).await.unwrap();
364
365 let listed = store.list_by_run(&run_id).await.unwrap();
366 assert_eq!(listed.len(), 1);
367 assert_eq!(listed[0].step_ref, "step-a");
368 assert_eq!(listed[0].occurrence, 0);
369 assert_eq!(listed[0].decode_step_output().unwrap(), json!({ "n": 1 }));
370
371 let restored = listed[0].decode_ctx_snapshot().unwrap();
373 assert_eq!(restored.agent, ctx.agent);
374 assert_eq!(restored.attempt, ctx.attempt);
375 assert_eq!(restored.meta.observer.get("k"), Some(&json!("v")));
376
377 driver.shutdown().await.unwrap();
378 }
379
380 #[tokio::test]
381 async fn sqlite_unique_4_tuple_is_enforced() {
382 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
383 let run_id = RunId::new();
384 let ctx = mk_ctx();
385
386 let e0 = ReplayEntry::from_completion(
387 run_id.clone(),
388 "step-a",
389 "hash-a",
390 0,
391 &ctx,
392 &json!("first"),
393 )
394 .unwrap();
395 store.append(e0.clone()).await.unwrap();
396
397 let dup_err = store.append(e0).await.unwrap_err();
399 assert!(
400 matches!(dup_err, ReplayStoreError::Duplicate { .. }),
401 "same (run_id, step_ref, input_hash, occurrence) must collide"
402 );
403
404 let e1 = ReplayEntry::from_completion(
407 run_id.clone(),
408 "step-a",
409 "hash-a",
410 1,
411 &ctx,
412 &json!("second"),
413 )
414 .unwrap();
415 store
416 .append(e1)
417 .await
418 .expect("occurrence=1 must not collide with occurrence=0");
419
420 let listed = store.list_by_run(&run_id).await.unwrap();
421 assert_eq!(listed.len(), 2);
422 let cursor = ReplayCursor::from_entries(listed);
423 assert_eq!(cursor.find("step-a", "hash-a", 0), Some(json!("first")));
424 assert_eq!(cursor.find("step-a", "hash-a", 1), Some(json!("second")));
425
426 driver.shutdown().await.unwrap();
427 }
428
429 #[tokio::test]
430 async fn sqlite_list_by_run_returns_in_seq_order() {
431 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
432 let run_id = RunId::new();
433 let ctx = mk_ctx();
434
435 for (i, (step, occ)) in [("a", 0), ("b", 0), ("a", 1)].iter().enumerate() {
436 store
437 .append(
438 ReplayEntry::from_completion(
439 run_id.clone(),
440 *step,
441 "h",
442 *occ,
443 &ctx,
444 &json!({ "idx": i }),
445 )
446 .unwrap(),
447 )
448 .await
449 .unwrap();
450 }
451
452 let listed = store.list_by_run(&run_id).await.unwrap();
453 let steps: Vec<(String, u32)> = listed
454 .iter()
455 .map(|e| (e.step_ref.clone(), e.occurrence))
456 .collect();
457 assert_eq!(
458 steps,
459 vec![
460 ("a".to_string(), 0),
461 ("b".to_string(), 0),
462 ("a".to_string(), 1),
463 ]
464 );
465
466 driver.shutdown().await.unwrap();
467 }
468
469 #[tokio::test]
470 async fn sqlite_name() {
471 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
472 assert_eq!(store.name(), "sqlite");
473 driver.shutdown().await.unwrap();
474 }
475
476 fn read_user_version(path: &std::path::Path) -> i64 {
483 let conn = rusqlite::Connection::open(path).unwrap();
484 conn.query_row("PRAGMA user_version", [], |r| r.get(0))
485 .unwrap()
486 }
487
488 #[tokio::test]
489 async fn sqlite_fresh_open_stamps_current_schema_version() {
490 let tmp = tempfile::TempDir::new().unwrap();
491 let path = tmp.path().join("replay.sqlite");
492
493 let (store, driver) = SqliteReplayStore::open(&path).await.unwrap();
494
495 let run_id = RunId::new();
497 let ctx = mk_ctx();
498 let entry = ReplayEntry::from_completion(
499 run_id.clone(),
500 "step-a",
501 "hash-a",
502 0,
503 &ctx,
504 &json!({ "n": 1 }),
505 )
506 .unwrap();
507 store.append(entry).await.unwrap();
508 let listed = store.list_by_run(&run_id).await.unwrap();
509 assert_eq!(listed.len(), 1);
510 assert_eq!(listed[0].decode_step_output().unwrap(), json!({ "n": 1 }));
511
512 driver.shutdown().await.unwrap();
513
514 assert_eq!(read_user_version(&path), CURRENT_SCHEMA_VERSION);
515 }
516
517 #[tokio::test]
518 async fn sqlite_legacy_schema_is_dropped_and_rebuilt() {
519 let tmp = tempfile::TempDir::new().unwrap();
520 let path = tmp.path().join("replay.sqlite");
521
522 {
526 let conn = rusqlite::Connection::open(&path).unwrap();
527 conn.execute_batch(
528 "CREATE TABLE replay_log (\
529 seq INTEGER PRIMARY KEY AUTOINCREMENT, \
530 run_id TEXT NOT NULL, \
531 step_ref TEXT NOT NULL, \
532 input_hash TEXT NOT NULL, \
533 occurrence INTEGER NOT NULL, \
534 value_json TEXT NOT NULL, \
535 created_at INTEGER NOT NULL, \
536 UNIQUE (run_id, step_ref, input_hash, occurrence)\
537 );",
538 )
539 .unwrap();
540 conn.execute(
541 "INSERT INTO replay_log \
542 (run_id, step_ref, input_hash, occurrence, value_json, created_at) \
543 VALUES ('legacy-run', 'legacy-step', 'legacy-hash', 0, 'legacy-value', 0)",
544 [],
545 )
546 .unwrap();
547 }
550
551 assert_eq!(read_user_version(&path), 0, "seed sanity: user_version=0");
552
553 let (store, driver) = SqliteReplayStore::open(&path).await.unwrap();
556
557 let run_id = RunId::new();
560 let ctx = mk_ctx();
561 let entry = ReplayEntry::from_completion(
562 run_id.clone(),
563 "step-a",
564 "hash-a",
565 0,
566 &ctx,
567 &json!({ "n": 42 }),
568 )
569 .unwrap();
570 store.append(entry).await.unwrap();
571 let listed = store.list_by_run(&run_id).await.unwrap();
572 assert_eq!(listed.len(), 1);
573
574 driver.shutdown().await.unwrap();
575
576 assert_eq!(read_user_version(&path), CURRENT_SCHEMA_VERSION);
577
578 let conn = rusqlite::Connection::open(&path).unwrap();
580 let legacy_count: i64 = conn
581 .query_row(
582 "SELECT COUNT(*) FROM replay_log WHERE run_id = 'legacy-run'",
583 [],
584 |r| r.get(0),
585 )
586 .unwrap();
587 assert_eq!(legacy_count, 0, "legacy rows must be dropped");
588
589 let ctx_col_count: i64 = conn
591 .query_row(
592 "SELECT COUNT(*) FROM pragma_table_info('replay_log') \
593 WHERE name = 'ctx_snapshot_json'",
594 [],
595 |r| r.get(0),
596 )
597 .unwrap();
598 assert_eq!(ctx_col_count, 1, "ctx_snapshot_json column must be present");
599 }
600
601 #[tokio::test]
602 async fn sqlite_current_schema_open_is_idempotent() {
603 let tmp = tempfile::TempDir::new().unwrap();
604 let path = tmp.path().join("replay.sqlite");
605
606 let run_id = {
608 let (store, driver) = SqliteReplayStore::open(&path).await.unwrap();
609 let run_id = RunId::new();
610 let ctx = mk_ctx();
611 let entry = ReplayEntry::from_completion(
612 run_id.clone(),
613 "step-a",
614 "hash-a",
615 0,
616 &ctx,
617 &json!({ "n": 7 }),
618 )
619 .unwrap();
620 store.append(entry).await.unwrap();
621 driver.shutdown().await.unwrap();
622 run_id
623 };
624
625 assert_eq!(read_user_version(&path), CURRENT_SCHEMA_VERSION);
626
627 let (store, driver) = SqliteReplayStore::open(&path).await.unwrap();
630 let listed = store.list_by_run(&run_id).await.unwrap();
631 assert_eq!(listed.len(), 1, "prior rows must survive re-open");
632 assert_eq!(listed[0].step_ref, "step-a");
633 assert_eq!(listed[0].decode_step_output().unwrap(), json!({ "n": 7 }));
634 driver.shutdown().await.unwrap();
635
636 assert_eq!(read_user_version(&path), CURRENT_SCHEMA_VERSION);
638 }
639
640 #[tokio::test]
641 async fn sqlite_delete_from_truncates_and_returns_count() {
642 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
643 let run_id = RunId::new();
644 let ctx = mk_ctx();
645 for (step, occ) in [("a", 0), ("b", 0), ("c", 0), ("d", 0)] {
646 store
647 .append(
648 ReplayEntry::from_completion(
649 run_id.clone(),
650 step,
651 "h",
652 occ,
653 &ctx,
654 &json!({ "s": step }),
655 )
656 .unwrap(),
657 )
658 .await
659 .unwrap();
660 }
661
662 let dropped = store.delete_from(&run_id, 2).await.unwrap();
663 assert_eq!(dropped, 2);
664
665 let listed = store.list_by_run(&run_id).await.unwrap();
666 let refs: Vec<String> = listed.iter().map(|e| e.step_ref.clone()).collect();
667 assert_eq!(refs, vec!["a", "b"]);
668
669 driver.shutdown().await.unwrap();
670 }
671
672 #[tokio::test]
673 async fn sqlite_delete_from_past_length_is_noop() {
674 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
675 let run_id = RunId::new();
676 let ctx = mk_ctx();
677 store
678 .append(
679 ReplayEntry::from_completion(run_id.clone(), "a", "h", 0, &ctx, &json!(1)).unwrap(),
680 )
681 .await
682 .unwrap();
683
684 assert_eq!(store.delete_from(&run_id, 1).await.unwrap(), 0);
685 assert_eq!(store.delete_from(&run_id, 99).await.unwrap(), 0);
686 assert_eq!(store.list_by_run(&run_id).await.unwrap().len(), 1);
687
688 driver.shutdown().await.unwrap();
689 }
690
691 #[tokio::test]
692 async fn sqlite_delete_from_missing_run_is_noop() {
693 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
694 let run_id = RunId::new();
695 assert_eq!(store.delete_from(&run_id, 0).await.unwrap(), 0);
696 driver.shutdown().await.unwrap();
697 }
698
699 #[tokio::test]
700 async fn sqlite_delete_from_isolates_runs() {
701 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
702 let r1 = RunId::new();
703 let r2 = RunId::new();
704 let ctx = mk_ctx();
705 for step in ["a", "b"] {
706 store
707 .append(
708 ReplayEntry::from_completion(r1.clone(), step, "h", 0, &ctx, &json!(step))
709 .unwrap(),
710 )
711 .await
712 .unwrap();
713 store
714 .append(
715 ReplayEntry::from_completion(r2.clone(), step, "h", 0, &ctx, &json!(step))
716 .unwrap(),
717 )
718 .await
719 .unwrap();
720 }
721 assert_eq!(store.delete_from(&r1, 0).await.unwrap(), 2);
722 assert!(store.list_by_run(&r1).await.unwrap().is_empty());
723 assert_eq!(store.list_by_run(&r2).await.unwrap().len(), 2);
724 driver.shutdown().await.unwrap();
725 }
726
727 #[tokio::test]
728 async fn sqlite_delete_from_frees_slots_for_reappend() {
729 let (store, driver) = SqliteReplayStore::open_in_memory().await.unwrap();
730 let run_id = RunId::new();
731 let ctx = mk_ctx();
732 for step in ["a", "b"] {
733 store
734 .append(
735 ReplayEntry::from_completion(
736 run_id.clone(),
737 step,
738 "hash",
739 0,
740 &ctx,
741 &json!({ "v": step }),
742 )
743 .unwrap(),
744 )
745 .await
746 .unwrap();
747 }
748 assert_eq!(store.delete_from(&run_id, 1).await.unwrap(), 1);
749 store
750 .append(
751 ReplayEntry::from_completion(
752 run_id.clone(),
753 "b",
754 "hash",
755 0,
756 &ctx,
757 &json!({ "v": "b-fresh" }),
758 )
759 .unwrap(),
760 )
761 .await
762 .expect("re-append after delete_from must not collide with UNIQUE");
763 let listed = store.list_by_run(&run_id).await.unwrap();
764 assert_eq!(listed.len(), 2);
765 assert_eq!(listed[0].step_ref, "a");
766 assert_eq!(listed[1].step_ref, "b");
767 assert_eq!(
768 listed[1].decode_step_output().unwrap(),
769 json!({ "v": "b-fresh" })
770 );
771 driver.shutdown().await.unwrap();
772 }
773
774 #[tokio::test]
775 async fn sqlite_future_schema_version_is_rejected() {
776 let tmp = tempfile::TempDir::new().unwrap();
777 let path = tmp.path().join("replay.sqlite");
778
779 {
782 let conn = rusqlite::Connection::open(&path).unwrap();
783 conn.execute_batch(&format!(
784 "PRAGMA user_version = {}",
785 CURRENT_SCHEMA_VERSION + 1
786 ))
787 .unwrap();
788 }
789
790 let res = SqliteReplayStore::open(&path).await;
791 let err = res
792 .err()
793 .expect("future user_version must be rejected by init_schema");
794 let msg = err.to_string();
797 assert!(
798 msg.contains("newer than supported"),
799 "unexpected error message: {msg}"
800 );
801 }
802}