1use rusqlite::backup::Backup;
2use rusqlite::blob::Blob;
3use rusqlite::{
4 Connection, ErrorCode, MAIN_DB, OpenFlags, OptionalExtension, Transaction, TransactionBehavior,
5 params,
6};
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use sha2::{Digest, Sha256};
10use std::fs::{self, OpenOptions};
11use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
12use std::path::{Path, PathBuf};
13use std::time::{Duration, Instant};
14use tempfile::{NamedTempFile, tempfile};
15
16#[cfg(test)]
17use crate::archive_protocol::ArchiveAck;
18use crate::archive_protocol::{ArchiveOperation, INGEST_SCHEMA_VERSION, MAX_FRAME_BYTES};
19use crate::config;
20
21const SCHEMA_VERSION: i64 = 1;
22const ZSTD_LEVEL: i32 = 3;
23const COMPRESSION_PERCENT: u64 = 95;
24const ARCHIVE_BUSY_TIMEOUT: Duration = Duration::from_secs(5);
25const ARCHIVE_REF_PREFIX: &str = "yr_";
26const ARCHIVE_REF_LEN: usize = 35;
27
28const SCHEMA: &str = r"
29CREATE TABLE sessions (
30 id INTEGER PRIMARY KEY,
31 agent TEXT NOT NULL,
32 account TEXT NOT NULL,
33 source_session_id TEXT NOT NULL,
34 started_at_ms INTEGER,
35 UNIQUE (agent, account, source_session_id)
36);
37
38CREATE TABLE tool_calls (
39 id INTEGER PRIMARY KEY,
40 session_id INTEGER NOT NULL REFERENCES sessions(id),
41 source_call_id TEXT NOT NULL,
42 archive_ref TEXT NOT NULL UNIQUE CHECK (
43 length(archive_ref) = 35
44 AND substr(archive_ref, 1, 3) = 'yr_'
45 AND substr(archive_ref, 4) NOT GLOB '*[^0-9a-f]*'
46 ),
47 tool_name TEXT NOT NULL,
48 provider TEXT,
49 model TEXT,
50 working_directory TEXT,
51 started_at_ms INTEGER NOT NULL,
52 requires_streams INTEGER NOT NULL CHECK (requires_streams IN (0, 1)),
53 finished_at_ms INTEGER,
54 status TEXT NOT NULL CHECK (status IN ('started', 'finished')),
55 is_error INTEGER CHECK (is_error IN (0, 1)),
56 executed INTEGER CHECK (executed IN (0, 1)),
57 UNIQUE (session_id, source_call_id)
58);
59
60CREATE TABLE payloads (
61 sha256 BLOB PRIMARY KEY CHECK (length(sha256) = 32),
62 compression TEXT NOT NULL CHECK (compression IN ('none', 'zstd')),
63 uncompressed_byte_length INTEGER NOT NULL CHECK (uncompressed_byte_length >= 0),
64 body BLOB NOT NULL
65);
66
67CREATE TABLE snapshots (
68 tool_call_id INTEGER NOT NULL REFERENCES tool_calls(id) ON DELETE CASCADE,
69 subject TEXT NOT NULL CHECK (subject IN ('input', 'result', 'result_text', 'source_output', 'stdout', 'stderr')),
70 stage TEXT NOT NULL CHECK (stage IN ('before', 'after')),
71 media_type TEXT NOT NULL,
72 source_completeness TEXT CHECK (source_completeness IN ('complete', 'incomplete', 'unknown')),
73 captured_at_ms INTEGER NOT NULL,
74 payload_sha256 BLOB NOT NULL REFERENCES payloads(sha256),
75 PRIMARY KEY (tool_call_id, subject, stage),
76 CHECK (
77 (subject = 'result_text' AND stage = 'before' AND source_completeness IS NOT NULL)
78 OR (subject <> 'result_text' AND source_completeness IS NULL)
79 )
80);
81
82CREATE INDEX tool_calls_started_at_idx ON tool_calls(started_at_ms);
83CREATE INDEX tool_calls_tool_name_idx ON tool_calls(tool_name);
84CREATE INDEX snapshots_payload_idx ON snapshots(payload_sha256);
85";
86
87#[derive(Clone, Debug, Deserialize, Serialize)]
88#[serde(rename_all = "camelCase")]
89pub struct SessionIdentity {
90 pub agent: String,
91 pub account: String,
92 pub source_session_id: String,
93 pub started_at_ms: Option<i64>,
94}
95
96#[derive(Clone, Debug, Deserialize, Serialize)]
97#[serde(rename_all = "camelCase")]
98pub struct CallIdentity {
99 pub source_call_id: String,
100 pub tool_name: String,
101 pub provider: Option<String>,
102 pub model: Option<String>,
103 pub working_directory: Option<String>,
104 pub started_at_ms: i64,
105 pub requires_streams: bool,
106}
107
108#[derive(Clone, Debug)]
109pub struct ArchiveKey {
110 pub session: SessionIdentity,
111 pub source_call_id: String,
112}
113
114#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
115#[serde(rename_all = "snake_case")]
116pub enum SourceCompleteness {
117 Complete,
118 Incomplete,
119 Unknown,
120}
121
122impl SourceCompleteness {
123 #[must_use]
124 pub const fn as_str(self) -> &'static str {
125 match self {
126 Self::Complete => "complete",
127 Self::Incomplete => "incomplete",
128 Self::Unknown => "unknown",
129 }
130 }
131}
132
133#[derive(Clone, Copy, Debug, Eq, PartialEq)]
134pub enum SourceName {
135 Stdout,
136 Stderr,
137 SourceOutput,
138 ResultText,
139}
140
141impl SourceName {
142 #[must_use]
143 pub const fn as_str(self) -> &'static str {
144 match self {
145 Self::Stdout => "stdout",
146 Self::Stderr => "stderr",
147 Self::SourceOutput => "source_output",
148 Self::ResultText => "result_text",
149 }
150 }
151}
152
153pub struct VerifiedSource {
154 pub name: SourceName,
155 pub completeness: SourceCompleteness,
156 pub media_type: String,
157 pub body: std::fs::File,
158 pub byte_length: u64,
159}
160
161#[derive(Debug, Eq, PartialEq)]
162pub struct ArchiveStats {
163 pub sessions: i64,
164 pub calls: i64,
165 pub incomplete_calls: i64,
166 pub logical_payload_bytes: i64,
167 pub stored_payload_bytes: i64,
168 pub database_bytes: u64,
169 pub oldest_call_ms: Option<i64>,
170 pub newest_call_ms: Option<i64>,
171}
172
173#[derive(Debug, Default, Eq, PartialEq)]
174pub struct VerifyReport {
175 pub incomplete_calls: i64,
176 pub errors: Vec<String>,
177}
178
179pub struct Archive {
180 connection: Connection,
181 path: PathBuf,
182}
183
184pub(crate) struct PreparedArchiveOperation {
185 estimated_bytes: u64,
186 kind: PreparedOperationKind,
187}
188
189enum PreparedOperationKind {
190 BeginCall {
191 session: SessionIdentity,
192 call: CallIdentity,
193 before: PreparedPayload,
194 after: PreparedPayload,
195 captured_at_ms: i64,
196 },
197 ResultBefore {
198 session: SessionIdentity,
199 source_call_id: String,
200 result: PreparedPayload,
201 source_output: Option<PreparedSnapshot>,
202 captured_at_ms: i64,
203 },
204 ResultText {
205 session: SessionIdentity,
206 source_call_id: String,
207 text: PreparedPayload,
208 completeness: SourceCompleteness,
209 captured_at_ms: i64,
210 },
211 StageResult {
212 session: SessionIdentity,
213 source_call_id: String,
214 result: PreparedPayload,
215 is_error: bool,
216 captured_at_ms: i64,
217 },
218 FinishCall {
219 session: SessionIdentity,
220 source_call_id: String,
221 result: PreparedPayload,
222 is_error: bool,
223 require_pre_result: bool,
224 finished_at_ms: i64,
225 },
226 UpdateFinalResult {
227 session: SessionIdentity,
228 source_call_id: String,
229 result: PreparedPayload,
230 is_error: bool,
231 finished_at_ms: i64,
232 },
233 CaptureStreams {
234 session: SessionIdentity,
235 source_call_id: String,
236 captured_at_ms: i64,
237 stdout_before: PreparedSnapshot,
238 stderr_before: PreparedSnapshot,
239 stdout_after: PreparedSnapshot,
240 stderr_after: PreparedSnapshot,
241 },
242 PruneBefore {
243 timestamp_ms: i64,
244 },
245}
246
247struct PreparedSnapshot {
248 media_type: &'static str,
249 payload: PreparedPayload,
250}
251
252struct PreparedPayload {
253 sha: Vec<u8>,
254 compression: &'static str,
255 uncompressed_length: u64,
256 stored_length: u64,
257 stored: NamedTempFile,
258}
259
260#[derive(Debug)]
261pub(crate) enum BatchWriteError {
262 Busy(String),
263 Permanent(String),
264}
265
266impl BatchWriteError {
267 pub(crate) fn message(&self) -> &str {
268 match self {
269 Self::Busy(message) | Self::Permanent(message) => message,
270 }
271 }
272
273 pub(crate) const fn is_busy(&self) -> bool {
274 matches!(self, Self::Busy(_))
275 }
276}
277
278impl PreparedArchiveOperation {
279 pub(crate) const fn estimated_bytes(&self) -> u64 {
280 self.estimated_bytes
281 }
282}
283
284impl Archive {
285 pub(crate) fn open() -> Result<Self, String> {
291 let archive = config::load()?.archive;
292 Self::open_path_with_policy(archive.path, archive.is_default_path)
293 }
294
295 pub fn open_read_only() -> Result<Self, String> {
301 let path = archive_path()?;
302 require_private_archive_path(&path)?;
303 let connection = Connection::open_with_flags(&path, OpenFlags::SQLITE_OPEN_READ_ONLY)
304 .map_err(|error| {
305 format!(
306 "could not open archive {} read-only: {error}",
307 path.display()
308 )
309 })?;
310 connection
311 .busy_timeout(ARCHIVE_BUSY_TIMEOUT)
312 .map_err(|error| format!("could not set archive busy timeout: {error}"))?;
313 let version: i64 = connection
314 .pragma_query_value(None, "user_version", |row| row.get(0))
315 .map_err(|error| format!("could not read archive schema version: {error}"))?;
316 if version != SCHEMA_VERSION {
317 return Err(format!(
318 "archive schema version {version}, expected {SCHEMA_VERSION}"
319 ));
320 }
321 if !schema_is_current(&connection)? {
322 return Err(
323 "archive schema requires migration; open it with a writable YARP command first"
324 .to_owned(),
325 );
326 }
327 Ok(Self { connection, path })
328 }
329
330 pub fn open_path(path: PathBuf) -> Result<Self, String> {
336 Self::open_path_with_policy(path, false)
337 }
338
339 fn open_path_with_policy(
340 path: PathBuf,
341 repair_existing_directory: bool,
342 ) -> Result<Self, String> {
343 prepare_path(&path, repair_existing_directory)?;
344 let mut connection = Connection::open(&path)
345 .map_err(|error| format!("could not open archive {}: {error}", path.display()))?;
346 connection
347 .busy_timeout(ARCHIVE_BUSY_TIMEOUT)
348 .map_err(|error| format!("could not set archive busy timeout: {error}"))?;
349 connection
350 .pragma_update(None, "foreign_keys", "ON")
351 .map_err(|error| format!("could not enable archive foreign keys: {error}"))?;
352 enable_wal_mode(&connection)?;
353 connection
354 .pragma_update(None, "synchronous", "FULL")
355 .map_err(|error| format!("could not set archive sync mode: {error}"))?;
356
357 let version: i64 = connection
358 .pragma_query_value(None, "user_version", |row| row.get(0))
359 .map_err(|error| format!("could not read archive schema version: {error}"))?;
360 match version {
361 0 => initialize_schema(&mut connection)?,
362 SCHEMA_VERSION => ensure_current_schema(&mut connection, &path)?,
363 newer if newer > SCHEMA_VERSION => {
364 return Err(format!(
365 "archive schema version {newer} is newer than supported version {SCHEMA_VERSION}"
366 ));
367 }
368 older => {
369 return Err(format!(
370 "archive schema version {older} has no supported migration"
371 ));
372 }
373 }
374 set_file_mode(&path, 0o600)?;
375 Ok(Self { connection, path })
376 }
377
378 pub(crate) fn configure_for_broker(&self) -> Result<(), String> {
379 self.connection
380 .busy_timeout(Duration::from_millis(250))
381 .map_err(|error| format!("could not set archive broker busy timeout: {error}"))?;
382 self.connection
383 .pragma_update(None, "wal_autocheckpoint", 1000)
384 .map_err(|error| format!("could not set archive WAL checkpoint limit: {error}"))
385 }
386
387 #[expect(
388 clippy::too_many_lines,
389 reason = "the exhaustive protocol mapping keeps preparation rules in one place"
390 )]
391 pub(crate) fn prepare_operation(
392 &self,
393 operation: ArchiveOperation,
394 ) -> Result<PreparedArchiveOperation, String> {
395 if operation.schema_version() != INGEST_SCHEMA_VERSION {
396 return Err(format!(
397 "unsupported ingest schema version {}; expected {INGEST_SCHEMA_VERSION}",
398 operation.schema_version()
399 ));
400 }
401 let owner_uid = archive_owner_uid(&self.path)?;
402 let kind = match operation {
403 ArchiveOperation::BeginCall {
404 session,
405 call,
406 input_before,
407 input_after,
408 captured_at_ms,
409 ..
410 } => PreparedOperationKind::BeginCall {
411 session,
412 call,
413 before: prepare_payload(&mut Cursor::new(canonical_json(&input_before)?))?,
414 after: prepare_payload(&mut Cursor::new(canonical_json(&input_after)?))?,
415 captured_at_ms,
416 },
417 ArchiveOperation::ResultBefore {
418 session,
419 source_call_id,
420 result,
421 full_output_path,
422 captured_at_ms,
423 ..
424 } => PreparedOperationKind::ResultBefore {
425 session,
426 source_call_id,
427 result: prepare_payload(&mut Cursor::new(canonical_json(&result)?))?,
428 source_output: full_output_path
429 .as_deref()
430 .map(|path| prepare_path_snapshot(path, owner_uid))
431 .transpose()?,
432 captured_at_ms,
433 },
434 ArchiveOperation::ResultText {
435 session,
436 source_call_id,
437 text,
438 source_completeness,
439 captured_at_ms,
440 ..
441 } => PreparedOperationKind::ResultText {
442 session,
443 source_call_id,
444 text: prepare_payload(&mut Cursor::new(text.into_bytes()))?,
445 completeness: source_completeness,
446 captured_at_ms,
447 },
448 ArchiveOperation::StageResult {
449 session,
450 source_call_id,
451 result,
452 is_error,
453 captured_at_ms,
454 ..
455 } => PreparedOperationKind::StageResult {
456 session,
457 source_call_id,
458 result: prepare_payload(&mut Cursor::new(canonical_json(&result)?))?,
459 is_error,
460 captured_at_ms,
461 },
462 ArchiveOperation::FinishCall {
463 session,
464 source_call_id,
465 result,
466 is_error,
467 require_pre_result,
468 finished_at_ms,
469 ..
470 } => PreparedOperationKind::FinishCall {
471 session,
472 source_call_id,
473 result: prepare_payload(&mut Cursor::new(canonical_json(&result)?))?,
474 is_error,
475 require_pre_result,
476 finished_at_ms,
477 },
478 ArchiveOperation::UpdateFinalResult {
479 session,
480 source_call_id,
481 result,
482 is_error,
483 finished_at_ms,
484 ..
485 } => PreparedOperationKind::UpdateFinalResult {
486 session,
487 source_call_id,
488 result: prepare_payload(&mut Cursor::new(canonical_json(&result)?))?,
489 is_error,
490 finished_at_ms,
491 },
492 ArchiveOperation::CaptureStreams {
493 session,
494 source_call_id,
495 captured_at_ms,
496 stdout_before_path,
497 stderr_before_path,
498 stdout_after_path,
499 stderr_after_path,
500 ..
501 } => PreparedOperationKind::CaptureStreams {
502 session,
503 source_call_id,
504 captured_at_ms,
505 stdout_before: prepare_path_snapshot(&stdout_before_path, owner_uid)?,
506 stderr_before: prepare_path_snapshot(&stderr_before_path, owner_uid)?,
507 stdout_after: prepare_path_snapshot(&stdout_after_path, owner_uid)?,
508 stderr_after: prepare_path_snapshot(&stderr_after_path, owner_uid)?,
509 },
510 ArchiveOperation::CapturePassthroughStreams {
511 session,
512 source_call_id,
513 captured_at_ms,
514 stdout_path,
515 stderr_path,
516 ..
517 } => {
518 let stdout_before = prepare_path_snapshot(&stdout_path, owner_uid)?;
519 let stdout_after = prepare_path_snapshot(&stdout_path, owner_uid)?;
520 let stderr_before = prepare_path_snapshot(&stderr_path, owner_uid)?;
521 let stderr_after = prepare_path_snapshot(&stderr_path, owner_uid)?;
522 PreparedOperationKind::CaptureStreams {
523 session,
524 source_call_id,
525 captured_at_ms,
526 stdout_before,
527 stderr_before,
528 stdout_after,
529 stderr_after,
530 }
531 }
532 ArchiveOperation::PruneBefore { timestamp_ms, .. } => {
533 PreparedOperationKind::PruneBefore { timestamp_ms }
534 }
535 };
536 let estimated_bytes = prepared_kind_bytes(&kind);
537 Ok(PreparedArchiveOperation {
538 estimated_bytes,
539 kind,
540 })
541 }
542
543 pub(crate) fn apply_prepared_batch<'a>(
544 &mut self,
545 operations: impl IntoIterator<Item = &'a mut PreparedArchiveOperation>,
546 ) -> Result<Vec<Result<Option<String>, String>>, BatchWriteError> {
547 let mut operations = operations.into_iter().collect::<Vec<_>>();
548 let mut transaction = self
549 .connection
550 .transaction_with_behavior(TransactionBehavior::Immediate)
551 .map_err(|error| batch_error("could not start archive batch", &error))?;
552 let mut results = Vec::with_capacity(operations.len());
553 let mut pruned = false;
554 for operation in &mut operations {
555 let mut savepoint = transaction.savepoint().map_err(|error| {
556 batch_error("could not start archive request savepoint", &error)
557 })?;
558 match apply_prepared_operation(&savepoint, &mut operation.kind) {
559 Ok(result) => {
560 pruned |= matches!(&operation.kind, PreparedOperationKind::PruneBefore { .. });
561 savepoint.commit().map_err(|error| {
562 batch_error("could not release archive request savepoint", &error)
563 })?;
564 results.push(Ok(result));
565 }
566 Err(error) => {
567 savepoint.rollback().map_err(|rollback_error| {
568 batch_error(
569 "could not roll back failed archive request savepoint",
570 &rollback_error,
571 )
572 })?;
573 results.push(Err(error));
574 }
575 }
576 }
577 if pruned {
578 transaction
579 .execute_batch("PRAGMA incremental_vacuum")
580 .map_err(|error| batch_error("could not vacuum archive", &error))?;
581 }
582 transaction
583 .commit()
584 .map_err(|error| batch_error("could not commit archive batch", &error))?;
585 Ok(results)
586 }
587
588 pub(crate) fn checkpoint(&self) -> Result<(), String> {
589 self.connection
590 .execute_batch("PRAGMA wal_checkpoint(PASSIVE)")
591 .map_err(|error| format!("could not checkpoint archive WAL: {error}"))
592 }
593
594 pub fn begin_call(
600 &mut self,
601 session: &SessionIdentity,
602 call: &CallIdentity,
603 input_before: &Value,
604 input_after: &Value,
605 captured_at_ms: i64,
606 ) -> Result<String, String> {
607 let before = canonical_json(input_before)?;
608 let after = canonical_json(input_after)?;
609 let transaction = self.transaction()?;
610 let session_id = ensure_session(&transaction, session)?;
611 let call_id = ensure_call(&transaction, session_id, call)?;
612 let archive_ref = call_archive_ref(&transaction, call_id)?;
613 insert_snapshot_bytes(
614 &transaction,
615 call_id,
616 "input",
617 "before",
618 "application/json",
619 captured_at_ms,
620 &before,
621 )?;
622 insert_snapshot_bytes(
623 &transaction,
624 call_id,
625 "input",
626 "after",
627 "application/json",
628 captured_at_ms,
629 &after,
630 )?;
631 transaction
632 .commit()
633 .map_err(|error| format!("could not commit tool call: {error}"))?;
634 Ok(archive_ref)
635 }
636
637 pub fn result_before(
643 &mut self,
644 session: &SessionIdentity,
645 source_call_id: &str,
646 result: &Value,
647 full_output_path: Option<&Path>,
648 captured_at_ms: i64,
649 ) -> Result<(), String> {
650 let body = canonical_json(result)?;
651 let mut full_output = full_output_path
652 .map(|path| {
653 let file = OpenOptions::new().read(true).open(path).map_err(|error| {
654 format!(
655 "could not open source full output {}: {error}",
656 path.display()
657 )
658 })?;
659 let metadata = file.metadata().map_err(|error| {
660 format!(
661 "could not inspect source full output {}: {error}",
662 path.display()
663 )
664 })?;
665 if !metadata.is_file() {
666 return Err(format!(
667 "source full output {} is not a regular file",
668 path.display()
669 ));
670 }
671 Ok((file, metadata.len()))
672 })
673 .transpose()?;
674 let transaction = self.transaction()?;
675 let call_id = find_call(&transaction, session, source_call_id)?;
676 insert_snapshot_bytes(
677 &transaction,
678 call_id,
679 "result",
680 "before",
681 "application/json",
682 captured_at_ms,
683 &body,
684 )?;
685 if let Some((file, original_length)) = &mut full_output {
686 insert_snapshot_reader(
687 &transaction,
688 call_id,
689 "source_output",
690 "before",
691 stream_media_type(file)?,
692 captured_at_ms,
693 file,
694 )?;
695 let final_length = file
696 .metadata()
697 .map_err(|error| format!("could not recheck source full output: {error}"))?
698 .len();
699 if final_length != *original_length {
700 return Err(format!(
701 "source full output changed size while archiving: {original_length} to {final_length} bytes"
702 ));
703 }
704 }
705 transaction
706 .commit()
707 .map_err(|error| format!("could not commit pre-YARP result: {error}"))
708 }
709
710 pub fn result_text(
717 &mut self,
718 session: &SessionIdentity,
719 source_call_id: &str,
720 text: &str,
721 completeness: SourceCompleteness,
722 captured_at_ms: i64,
723 ) -> Result<String, String> {
724 let transaction = self.transaction()?;
725 let call_id = find_call(&transaction, session, source_call_id)?;
726 insert_snapshot_reader_with_completeness(
727 &transaction,
728 call_id,
729 "result_text",
730 "before",
731 "text/plain; charset=utf-8",
732 Some(completeness.as_str()),
733 captured_at_ms,
734 &mut Cursor::new(text.as_bytes()),
735 )?;
736 let archive_ref = call_archive_ref(&transaction, call_id)?;
737 transaction
738 .commit()
739 .map_err(|error| format!("could not commit result text: {error}"))?;
740 Ok(archive_ref)
741 }
742
743 pub fn archive_ref(&self, key: &ArchiveKey) -> Result<String, String> {
749 self.connection
750 .query_row(
751 "SELECT c.archive_ref
752 FROM tool_calls c
753 JOIN sessions s ON s.id = c.session_id
754 WHERE s.agent = ?1 AND s.account = ?2 AND s.source_session_id = ?3
755 AND c.source_call_id = ?4",
756 params![
757 key.session.agent,
758 key.session.account,
759 key.session.source_session_id,
760 key.source_call_id
761 ],
762 |row| row.get(0),
763 )
764 .map_err(|error| format!("could not resolve archive call reference: {error}"))
765 }
766
767 pub fn searchable_sources(&self, archive_ref: &str) -> Result<Vec<VerifiedSource>, String> {
774 validate_archive_ref(archive_ref)?;
775 let call_id: i64 = self
776 .connection
777 .query_row(
778 "SELECT id FROM tool_calls WHERE archive_ref = ?1",
779 [archive_ref],
780 |row| row.get(0),
781 )
782 .map_err(|error| format!("could not find archive reference {archive_ref}: {error}"))?;
783 let stdout = self.snapshot_exists(call_id, "stdout", "before")?;
784 let stderr = self.snapshot_exists(call_id, "stderr", "before")?;
785 if stdout || stderr {
786 let mut sources = Vec::with_capacity(2);
787 if stdout {
788 sources.push(self.verified_source(
789 call_id,
790 SourceName::Stdout,
791 SourceCompleteness::Complete,
792 )?);
793 }
794 if stderr {
795 sources.push(self.verified_source(
796 call_id,
797 SourceName::Stderr,
798 SourceCompleteness::Complete,
799 )?);
800 }
801 return Ok(sources);
802 }
803 if self.snapshot_exists(call_id, "result_text", "before")? {
804 let value: String = self
805 .connection
806 .query_row(
807 "SELECT source_completeness FROM snapshots
808 WHERE tool_call_id = ?1 AND subject = 'result_text' AND stage = 'before'",
809 [call_id],
810 |row| row.get(0),
811 )
812 .map_err(|error| format!("could not read result text completeness: {error}"))?;
813 let completeness = parse_completeness(&value)?;
814 return Ok(vec![self.verified_source(
815 call_id,
816 SourceName::ResultText,
817 completeness,
818 )?]);
819 }
820 if self.snapshot_exists(call_id, "source_output", "before")? {
821 return Ok(vec![self.verified_source(
822 call_id,
823 SourceName::SourceOutput,
824 SourceCompleteness::Complete,
825 )?]);
826 }
827 Err(format!(
828 "archive reference {archive_ref} has no searchable source"
829 ))
830 }
831
832 pub fn stage_result(
839 &mut self,
840 session: &SessionIdentity,
841 source_call_id: &str,
842 result: &Value,
843 is_error: bool,
844 captured_at_ms: i64,
845 ) -> Result<(), String> {
846 let body = canonical_json(result)?;
847 let transaction = self.transaction()?;
848 let call_id = find_call(&transaction, session, source_call_id)?;
849 require_call_status(&transaction, call_id, source_call_id, "started")?;
850 validate_completion(&transaction, call_id, true)?;
851 insert_snapshot_bytes(
852 &transaction,
853 call_id,
854 "result",
855 "after",
856 "application/json",
857 captured_at_ms,
858 &body,
859 )?;
860 transaction
861 .execute(
862 "UPDATE tool_calls SET is_error = ?1, executed = 1 WHERE id = ?2",
863 params![i64::from(is_error), call_id],
864 )
865 .map_err(|error| format!("could not stage tool result: {error}"))?;
866 transaction
867 .commit()
868 .map_err(|error| format!("could not commit staged tool result: {error}"))
869 }
870
871 pub fn finish_call(
877 &mut self,
878 session: &SessionIdentity,
879 source_call_id: &str,
880 result: &Value,
881 is_error: bool,
882 require_pre_result: bool,
883 finished_at_ms: i64,
884 ) -> Result<(), String> {
885 let body = canonical_json(result)?;
886 let transaction = self.transaction()?;
887 let call_id = find_call(&transaction, session, source_call_id)?;
888 validate_completion(&transaction, call_id, require_pre_result)?;
889 insert_snapshot_bytes(
890 &transaction,
891 call_id,
892 "result",
893 "after",
894 "application/json",
895 finished_at_ms,
896 &body,
897 )?;
898 transaction
899 .execute(
900 "UPDATE tool_calls
901 SET finished_at_ms = ?1, status = 'finished', is_error = ?2, executed = ?3
902 WHERE id = ?4",
903 params![
904 finished_at_ms,
905 i64::from(is_error),
906 i64::from(require_pre_result),
907 call_id
908 ],
909 )
910 .map_err(|error| format!("could not finish tool call: {error}"))?;
911 transaction
912 .commit()
913 .map_err(|error| format!("could not commit final tool result: {error}"))
914 }
915
916 pub fn update_final_result(
922 &mut self,
923 session: &SessionIdentity,
924 source_call_id: &str,
925 result: &Value,
926 is_error: bool,
927 finished_at_ms: i64,
928 ) -> Result<(), String> {
929 let body = canonical_json(result)?;
930 let transaction = self.transaction()?;
931 let call_id = find_call(&transaction, session, source_call_id)?;
932 let status: String = transaction
933 .query_row(
934 "SELECT status FROM tool_calls WHERE id = ?1",
935 [call_id],
936 |row| row.get(0),
937 )
938 .map_err(|error| format!("could not read final tool call status: {error}"))?;
939 if status == "finished" {
940 let expected_sha = Sha256::digest(&body).to_vec();
941 let stored: (Vec<u8>, bool, i64, bool) = transaction
942 .query_row(
943 "SELECT s.payload_sha256, c.is_error, c.finished_at_ms, c.executed
944 FROM tool_calls c
945 JOIN snapshots s ON s.tool_call_id = c.id
946 WHERE c.id = ?1 AND s.subject = 'result' AND s.stage = 'after'",
947 [call_id],
948 |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
949 )
950 .map_err(|error| format!("could not check finalized tool result: {error}"))?;
951 if stored == (expected_sha, is_error, finished_at_ms, true) {
952 return transaction
953 .commit()
954 .map_err(|error| format!("could not commit repeated final result: {error}"));
955 }
956 return Err(format!(
957 "tool call {source_call_id} was already finalized with different content or metadata"
958 ));
959 }
960 if status != "started" {
961 return Err(format!(
962 "tool call {source_call_id} has status {status}, expected started"
963 ));
964 }
965 validate_completion(&transaction, call_id, true)?;
966 let replaced_sha = replace_snapshot_bytes(
967 &transaction,
968 call_id,
969 "result",
970 "after",
971 "application/json",
972 finished_at_ms,
973 &body,
974 )?;
975 transaction
976 .execute(
977 "UPDATE tool_calls
978 SET finished_at_ms = ?1, status = 'finished', is_error = ?2, executed = 1
979 WHERE id = ?3",
980 params![finished_at_ms, i64::from(is_error), call_id],
981 )
982 .map_err(|error| format!("could not finalize staged tool call: {error}"))?;
983 transaction
984 .execute(
985 "DELETE FROM payloads
986 WHERE sha256 = ?1
987 AND NOT EXISTS (
988 SELECT 1 FROM snapshots WHERE snapshots.payload_sha256 = payloads.sha256
989 )",
990 [&replaced_sha],
991 )
992 .map_err(|error| format!("could not remove replaced final payload: {error}"))?;
993 transaction
994 .commit()
995 .map_err(|error| format!("could not commit final staged tool result: {error}"))
996 }
997
998 pub fn capture_streams(
1004 &mut self,
1005 key: &ArchiveKey,
1006 captured_at_ms: i64,
1007 stdout_before: &mut (impl Read + Seek),
1008 stderr_before: &mut (impl Read + Seek),
1009 stdout_after: &[u8],
1010 stderr_after: &[u8],
1011 ) -> Result<(), String> {
1012 let transaction = self.transaction()?;
1013 let call_id = find_call(&transaction, &key.session, &key.source_call_id)?;
1014 insert_snapshot_reader(
1015 &transaction,
1016 call_id,
1017 "stdout",
1018 "before",
1019 stream_media_type(stdout_before)?,
1020 captured_at_ms,
1021 stdout_before,
1022 )?;
1023 insert_snapshot_reader(
1024 &transaction,
1025 call_id,
1026 "stderr",
1027 "before",
1028 stream_media_type(stderr_before)?,
1029 captured_at_ms,
1030 stderr_before,
1031 )?;
1032 insert_snapshot_bytes(
1033 &transaction,
1034 call_id,
1035 "stdout",
1036 "after",
1037 stream_media_type(&mut Cursor::new(stdout_after))?,
1038 captured_at_ms,
1039 stdout_after,
1040 )?;
1041 insert_snapshot_bytes(
1042 &transaction,
1043 call_id,
1044 "stderr",
1045 "after",
1046 stream_media_type(&mut Cursor::new(stderr_after))?,
1047 captured_at_ms,
1048 stderr_after,
1049 )?;
1050 transaction
1051 .commit()
1052 .map_err(|error| format!("could not commit shell streams: {error}"))
1053 }
1054
1055 pub fn capture_passthrough_streams(
1061 &mut self,
1062 key: &ArchiveKey,
1063 captured_at_ms: i64,
1064 stdout: &mut (impl Read + Seek),
1065 stderr: &mut (impl Read + Seek),
1066 ) -> Result<(), String> {
1067 let transaction = self.transaction()?;
1068 let call_id = find_call(&transaction, &key.session, &key.source_call_id)?;
1069 let stdout_type = stream_media_type(stdout)?;
1070 insert_snapshot_reader(
1071 &transaction,
1072 call_id,
1073 "stdout",
1074 "before",
1075 stdout_type,
1076 captured_at_ms,
1077 stdout,
1078 )?;
1079 insert_snapshot_reader(
1080 &transaction,
1081 call_id,
1082 "stdout",
1083 "after",
1084 stdout_type,
1085 captured_at_ms,
1086 stdout,
1087 )?;
1088 let stderr_type = stream_media_type(stderr)?;
1089 insert_snapshot_reader(
1090 &transaction,
1091 call_id,
1092 "stderr",
1093 "before",
1094 stderr_type,
1095 captured_at_ms,
1096 stderr,
1097 )?;
1098 insert_snapshot_reader(
1099 &transaction,
1100 call_id,
1101 "stderr",
1102 "after",
1103 stderr_type,
1104 captured_at_ms,
1105 stderr,
1106 )?;
1107 transaction
1108 .commit()
1109 .map_err(|error| format!("could not commit passthrough shell streams: {error}"))
1110 }
1111
1112 pub fn restore_streams(
1118 &self,
1119 key: &ArchiveKey,
1120 mut stdout: impl Write,
1121 mut stderr: impl Write,
1122 ) -> Result<(), String> {
1123 let call_id: i64 = self
1124 .connection
1125 .query_row(
1126 "SELECT c.id
1127 FROM tool_calls c
1128 JOIN sessions s ON s.id = c.session_id
1129 WHERE s.agent = ?1 AND s.account = ?2 AND s.source_session_id = ?3
1130 AND c.source_call_id = ?4",
1131 params![
1132 key.session.agent,
1133 key.session.account,
1134 key.session.source_session_id,
1135 key.source_call_id
1136 ],
1137 |row| row.get(0),
1138 )
1139 .map_err(|error| format!("could not find tool call {}: {error}", key.source_call_id))?;
1140 let mut restored_stdout = self.verified_snapshot(call_id, "stdout", "before")?;
1141 let mut restored_stderr = self.verified_snapshot(call_id, "stderr", "before")?;
1142 io::copy(&mut restored_stdout, &mut stdout)
1143 .map_err(|error| format!("could not write restored stdout: {error}"))?;
1144 io::copy(&mut restored_stderr, &mut stderr)
1145 .map_err(|error| format!("could not write restored stderr: {error}"))?;
1146 Ok(())
1147 }
1148
1149 pub fn stats(&self) -> Result<ArchiveStats, String> {
1155 let row = self
1156 .connection
1157 .query_row(
1158 "SELECT
1159 (SELECT count(*) FROM sessions),
1160 (SELECT count(*) FROM tool_calls),
1161 (SELECT count(*) FROM tool_calls WHERE status = 'started'),
1162 (SELECT coalesce(sum(uncompressed_byte_length), 0) FROM payloads),
1163 (SELECT coalesce(sum(length(body)), 0) FROM payloads),
1164 (SELECT min(started_at_ms) FROM tool_calls),
1165 (SELECT max(started_at_ms) FROM tool_calls)",
1166 [],
1167 |row| {
1168 Ok((
1169 row.get(0)?,
1170 row.get(1)?,
1171 row.get(2)?,
1172 row.get(3)?,
1173 row.get(4)?,
1174 row.get(5)?,
1175 row.get(6)?,
1176 ))
1177 },
1178 )
1179 .map_err(|error| format!("could not read archive statistics: {error}"))?;
1180 let database_bytes = archive_file_bytes(&self.path)?;
1181 Ok(ArchiveStats {
1182 sessions: row.0,
1183 calls: row.1,
1184 incomplete_calls: row.2,
1185 logical_payload_bytes: row.3,
1186 stored_payload_bytes: row.4,
1187 oldest_call_ms: row.5,
1188 newest_call_ms: row.6,
1189 database_bytes,
1190 })
1191 }
1192
1193 pub fn verify(&self) -> Result<VerifyReport, String> {
1199 let mut report = VerifyReport::default();
1200 let integrity: String = self
1201 .connection
1202 .query_row("PRAGMA integrity_check", [], |row| row.get(0))
1203 .map_err(|error| format!("could not run SQLite integrity check: {error}"))?;
1204 if integrity != "ok" {
1205 report
1206 .errors
1207 .push(format!("SQLite integrity check: {integrity}"));
1208 }
1209 let foreign_keys: i64 = self
1210 .connection
1211 .query_row("SELECT count(*) FROM pragma_foreign_key_check", [], |row| {
1212 row.get(0)
1213 })
1214 .map_err(|error| format!("could not check archive foreign keys: {error}"))?;
1215 if foreign_keys > 0 {
1216 report
1217 .errors
1218 .push(format!("{foreign_keys} foreign key violation(s)"));
1219 }
1220 let version: i64 = self
1221 .connection
1222 .pragma_query_value(None, "user_version", |row| row.get(0))
1223 .map_err(|error| format!("could not read archive schema version: {error}"))?;
1224 if version != SCHEMA_VERSION {
1225 report.errors.push(format!(
1226 "archive schema version {version}, expected {SCHEMA_VERSION}"
1227 ));
1228 }
1229 report.incomplete_calls = self
1230 .connection
1231 .query_row(
1232 "SELECT count(*) FROM tool_calls WHERE status = 'started'",
1233 [],
1234 |row| row.get(0),
1235 )
1236 .map_err(|error| format!("could not count incomplete calls: {error}"))?;
1237 self.verify_constrained_values(&mut report)?;
1238 self.verify_snapshot_sets(&mut report)?;
1239 let orphaned: i64 = self
1240 .connection
1241 .query_row(
1242 "SELECT count(*) FROM payloads p WHERE NOT EXISTS (SELECT 1 FROM snapshots s WHERE s.payload_sha256 = p.sha256)",
1243 [],
1244 |row| row.get(0),
1245 )
1246 .map_err(|error| format!("could not check unreferenced payloads: {error}"))?;
1247 if orphaned > 0 {
1248 report
1249 .errors
1250 .push(format!("{orphaned} unreferenced payload(s)"));
1251 }
1252 self.verify_payloads(&mut report)?;
1253 verify_permissions(&self.path, &mut report)?;
1254 Ok(report)
1255 }
1256
1257 pub fn prune_before(&mut self, timestamp_ms: i64) -> Result<i64, String> {
1263 let transaction = self.transaction()?;
1264 let deleted = transaction
1265 .execute(
1266 "DELETE FROM tool_calls WHERE status = 'finished' AND finished_at_ms < ?1",
1267 [timestamp_ms],
1268 )
1269 .map_err(|error| format!("could not prune tool calls: {error}"))?;
1270 transaction
1271 .execute(
1272 "DELETE FROM payloads WHERE NOT EXISTS (SELECT 1 FROM snapshots WHERE snapshots.payload_sha256 = payloads.sha256)",
1273 [],
1274 )
1275 .map_err(|error| format!("could not prune unreferenced payloads: {error}"))?;
1276 transaction
1277 .commit()
1278 .map_err(|error| format!("could not commit archive prune: {error}"))?;
1279 self.connection
1280 .execute_batch("PRAGMA incremental_vacuum")
1281 .map_err(|error| format!("could not vacuum archive: {error}"))?;
1282 i64::try_from(deleted).map_err(|_| "pruned call count does not fit in i64".to_owned())
1283 }
1284
1285 fn verify_constrained_values(&self, report: &mut VerifyReport) -> Result<(), String> {
1286 let checks = [
1287 (
1288 "SELECT count(*) FROM tool_calls
1289 WHERE status NOT IN ('started', 'finished')",
1290 "tool call(s) with invalid status",
1291 ),
1292 (
1293 "SELECT count(*) FROM tool_calls
1294 WHERE requires_streams NOT IN (0, 1)",
1295 "tool call(s) with invalid requires_streams",
1296 ),
1297 (
1298 "SELECT count(*) FROM tool_calls
1299 WHERE is_error IS NOT NULL AND is_error NOT IN (0, 1)",
1300 "tool call(s) with invalid is_error",
1301 ),
1302 (
1303 "SELECT count(*) FROM tool_calls
1304 WHERE executed IS NOT NULL AND executed NOT IN (0, 1)",
1305 "tool call(s) with invalid executed",
1306 ),
1307 (
1308 "SELECT count(*) FROM tool_calls
1309 WHERE (status = 'started' AND (
1310 finished_at_ms IS NOT NULL
1311 OR (is_error IS NULL) != (executed IS NULL)
1312 OR executed = 0
1313 ))
1314 OR (status = 'finished' AND (finished_at_ms IS NULL OR is_error IS NULL OR executed IS NULL))",
1315 "tool call(s) with inconsistent lifecycle state",
1316 ),
1317 (
1318 "SELECT count(*) FROM snapshots
1319 WHERE subject NOT IN ('input', 'result', 'result_text', 'source_output', 'stdout', 'stderr')",
1320 "snapshot(s) with invalid subject",
1321 ),
1322 (
1323 "SELECT count(*) FROM snapshots
1324 WHERE (subject = 'result_text' AND (stage != 'before' OR source_completeness NOT IN ('complete', 'incomplete', 'unknown')))
1325 OR (subject != 'result_text' AND source_completeness IS NOT NULL)",
1326 "snapshot(s) with invalid source completeness",
1327 ),
1328 (
1329 "SELECT count(*) FROM snapshots WHERE stage NOT IN ('before', 'after')",
1330 "snapshot(s) with invalid stage",
1331 ),
1332 (
1333 "SELECT count(*) FROM snapshots
1334 WHERE (subject IN ('input', 'result') AND media_type != 'application/json')
1335 OR (subject = 'result_text' AND media_type != 'text/plain; charset=utf-8')
1336 OR (subject IN ('source_output', 'stdout', 'stderr') AND media_type NOT IN ('text/plain; charset=utf-8', 'application/octet-stream'))",
1337 "snapshot(s) with invalid media type",
1338 ),
1339 (
1340 "SELECT count(*) FROM tool_calls
1341 WHERE length(archive_ref) != 35
1342 OR substr(archive_ref, 1, 3) != 'yr_'
1343 OR substr(archive_ref, 4) GLOB '*[^0-9a-f]*'",
1344 "tool call(s) with invalid archive reference",
1345 ),
1346 (
1347 "SELECT count(*) FROM payloads
1348 WHERE length(sha256) != 32
1349 OR compression NOT IN ('none', 'zstd')
1350 OR uncompressed_byte_length < 0",
1351 "payload(s) with invalid constrained values",
1352 ),
1353 ];
1354 for (query, message) in checks {
1355 let invalid: i64 = self
1356 .connection
1357 .query_row(query, [], |row| row.get(0))
1358 .map_err(|error| format!("could not verify constrained values: {error}"))?;
1359 if invalid > 0 {
1360 report.errors.push(format!("{invalid} {message}"));
1361 }
1362 }
1363 Ok(())
1364 }
1365
1366 fn verify_snapshot_sets(&self, report: &mut VerifyReport) -> Result<(), String> {
1367 let checks = [
1368 (
1369 "SELECT count(*) FROM tool_calls c
1370 WHERE NOT EXISTS (SELECT 1 FROM snapshots s WHERE s.tool_call_id = c.id AND s.subject = 'input' AND s.stage = 'before')
1371 OR NOT EXISTS (SELECT 1 FROM snapshots s WHERE s.tool_call_id = c.id AND s.subject = 'input' AND s.stage = 'after')",
1372 "call(s) missing input snapshots",
1373 ),
1374 (
1375 "SELECT count(*) FROM tool_calls c
1376 WHERE c.status = 'finished'
1377 AND NOT EXISTS (SELECT 1 FROM snapshots s WHERE s.tool_call_id = c.id AND s.subject = 'result' AND s.stage = 'after')",
1378 "finished call(s) missing final results",
1379 ),
1380 (
1381 "SELECT count(*) FROM tool_calls c
1382 WHERE c.status = 'finished' AND c.executed = 1
1383 AND NOT EXISTS (
1384 SELECT 1 FROM snapshots s
1385 WHERE s.tool_call_id = c.id AND s.subject = 'result' AND s.stage = 'before'
1386 )",
1387 "executed call(s) missing pre-YARP results",
1388 ),
1389 (
1390 "SELECT count(*) FROM tool_calls c
1391 WHERE c.status = 'finished' AND c.requires_streams = 1 AND c.executed = 1
1392 AND 4 != (
1393 SELECT count(*) FROM snapshots s
1394 WHERE s.tool_call_id = c.id
1395 AND s.subject IN ('stdout', 'stderr')
1396 AND s.stage IN ('before', 'after')
1397 )",
1398 "shell call(s) missing complete stream snapshots",
1399 ),
1400 ];
1401 for (query, message) in checks {
1402 let missing: i64 = self
1403 .connection
1404 .query_row(query, [], |row| row.get(0))
1405 .map_err(|error| format!("could not verify snapshot sets: {error}"))?;
1406 if missing > 0 {
1407 report.errors.push(format!("{missing} {message}"));
1408 }
1409 }
1410 Ok(())
1411 }
1412
1413 fn snapshot_exists(&self, call_id: i64, subject: &str, stage: &str) -> Result<bool, String> {
1414 self.connection
1415 .query_row(
1416 "SELECT EXISTS(
1417 SELECT 1 FROM snapshots
1418 WHERE tool_call_id = ?1 AND subject = ?2 AND stage = ?3
1419 )",
1420 params![call_id, subject, stage],
1421 |row| row.get(0),
1422 )
1423 .map_err(|error| format!("could not inspect {subject}/{stage} snapshot: {error}"))
1424 }
1425
1426 fn verified_source(
1427 &self,
1428 call_id: i64,
1429 name: SourceName,
1430 completeness: SourceCompleteness,
1431 ) -> Result<VerifiedSource, String> {
1432 let subject = name.as_str();
1433 let (media_type, byte_length): (String, i64) = self
1434 .connection
1435 .query_row(
1436 "SELECT s.media_type, p.uncompressed_byte_length
1437 FROM snapshots s
1438 JOIN payloads p ON p.sha256 = s.payload_sha256
1439 WHERE s.tool_call_id = ?1 AND s.subject = ?2 AND s.stage = 'before'",
1440 params![call_id, subject],
1441 |row| Ok((row.get(0)?, row.get(1)?)),
1442 )
1443 .map_err(|error| format!("could not inspect {subject}/before snapshot: {error}"))?;
1444 let byte_length = u64::try_from(byte_length)
1445 .map_err(|_| format!("invalid {subject}/before byte length"))?;
1446 let body = self.verified_snapshot(call_id, subject, "before")?;
1447 Ok(VerifiedSource {
1448 name,
1449 completeness,
1450 media_type,
1451 body,
1452 byte_length,
1453 })
1454 }
1455
1456 fn verified_snapshot(
1457 &self,
1458 call_id: i64,
1459 subject: &str,
1460 stage: &str,
1461 ) -> Result<std::fs::File, String> {
1462 let (rowid, expected_sha, compression, expected_length): (i64, Vec<u8>, String, i64) = self
1463 .connection
1464 .query_row(
1465 "SELECT p.rowid, p.sha256, p.compression, p.uncompressed_byte_length
1466 FROM snapshots s
1467 JOIN payloads p ON p.sha256 = s.payload_sha256
1468 WHERE s.tool_call_id = ?1 AND s.subject = ?2 AND s.stage = ?3",
1469 params![call_id, subject, stage],
1470 |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
1471 )
1472 .map_err(|error| format!("could not find {subject}/{stage} snapshot: {error}"))?;
1473 let blob = self
1474 .connection
1475 .blob_open(MAIN_DB, "payloads", "body", rowid, true)
1476 .map_err(|error| format!("could not open {subject}/{stage} payload: {error}"))?;
1477 let mut reader: Box<dyn Read> = match compression.as_str() {
1478 "none" => Box::new(blob),
1479 "zstd" => Box::new(
1480 zstd::stream::read::Decoder::new(blob)
1481 .map_err(|error| format!("could not decode {subject}/{stage}: {error}"))?,
1482 ),
1483 unknown => return Err(format!("unknown payload compression {unknown}")),
1484 };
1485 let expected_length = u64::try_from(expected_length)
1486 .map_err(|_| format!("invalid {subject}/{stage} payload length {expected_length}"))?;
1487 let mut restored = tempfile().map_err(|error| {
1488 format!("could not create {subject}/{stage} restore spool: {error}")
1489 })?;
1490 let mut hasher = Sha256::new();
1491 let mut length = 0_u64;
1492 let mut buffer = vec![0_u8; 64 * 1024];
1493 loop {
1494 let count = reader
1495 .read(&mut buffer)
1496 .map_err(|error| format!("could not read {subject}/{stage} payload: {error}"))?;
1497 if count == 0 {
1498 break;
1499 }
1500 restored
1501 .write_all(&buffer[..count])
1502 .map_err(|error| format!("could not spool {subject}/{stage} payload: {error}"))?;
1503 hasher.update(&buffer[..count]);
1504 length = length.saturating_add(count as u64);
1505 }
1506 let actual_sha: [u8; 32] = hasher.finalize().into();
1507 if length != expected_length || actual_sha.as_slice() != expected_sha {
1508 return Err(format!(
1509 "{subject}/{stage} payload failed integrity verification"
1510 ));
1511 }
1512 restored.seek(SeekFrom::Start(0)).map_err(|error| {
1513 format!("could not rewind {subject}/{stage} restore spool: {error}")
1514 })?;
1515 Ok(restored)
1516 }
1517
1518 fn transaction(&mut self) -> Result<Transaction<'_>, String> {
1519 self.connection
1520 .transaction_with_behavior(TransactionBehavior::Immediate)
1521 .map_err(|error| format!("could not start archive transaction: {error}"))
1522 }
1523
1524 fn verify_payloads(&self, report: &mut VerifyReport) -> Result<(), String> {
1525 let mut statement = self
1526 .connection
1527 .prepare(
1528 "SELECT rowid, sha256, compression, uncompressed_byte_length, length(body) FROM payloads",
1529 )
1530 .map_err(|error| format!("could not prepare payload verification: {error}"))?;
1531 let rows = statement
1532 .query_map([], |row| {
1533 Ok((
1534 row.get::<_, i64>(0)?,
1535 row.get::<_, Vec<u8>>(1)?,
1536 row.get::<_, String>(2)?,
1537 row.get::<_, i64>(3)?,
1538 row.get::<_, i64>(4)?,
1539 ))
1540 })
1541 .map_err(|error| format!("could not list archive payloads: {error}"))?
1542 .collect::<Result<Vec<_>, _>>()
1543 .map_err(|error| format!("could not read archive payload metadata: {error}"))?;
1544 drop(statement);
1545
1546 for (rowid, expected_sha, compression, expected_length, _) in rows {
1547 let blob = self
1548 .connection
1549 .blob_open(MAIN_DB, "payloads", "body", rowid, true)
1550 .map_err(|error| format!("could not open payload {rowid}: {error}"))?;
1551 let mut reader: Box<dyn Read> = match compression.as_str() {
1552 "none" => Box::new(blob),
1553 "zstd" => Box::new(
1554 zstd::stream::read::Decoder::new(blob)
1555 .map_err(|error| format!("could not decode payload {rowid}: {error}"))?,
1556 ),
1557 unknown => {
1558 report
1559 .errors
1560 .push(format!("payload {rowid} has unknown compression {unknown}"));
1561 continue;
1562 }
1563 };
1564 let mut hasher = Sha256::new();
1565 let mut length = 0_u64;
1566 let mut buffer = vec![0_u8; 64 * 1024];
1567 loop {
1568 let count = reader
1569 .read(&mut buffer)
1570 .map_err(|error| format!("could not read payload {rowid}: {error}"))?;
1571 if count == 0 {
1572 break;
1573 }
1574 hasher.update(&buffer[..count]);
1575 length = length.saturating_add(count as u64);
1576 }
1577 if length != u64::try_from(expected_length).unwrap_or(u64::MAX) {
1578 report.errors.push(format!(
1579 "payload {rowid} length {length}, expected {expected_length}"
1580 ));
1581 }
1582 if hasher.finalize().as_slice() != expected_sha {
1583 report
1584 .errors
1585 .push(format!("payload {rowid} SHA-256 mismatch"));
1586 }
1587 }
1588 Ok(())
1589 }
1590}
1591
1592pub fn run_ingest(input: impl Read, output: impl Write) -> Result<(), String> {
1598 crate::archive_client::run_ingest_bridge(input, output)
1599}
1600
1601#[cfg(test)]
1602fn run_ingest_with_archive(
1603 mut archive: Archive,
1604 mut input: impl Read,
1605 mut output: impl Write,
1606) -> Result<(), String> {
1607 loop {
1608 let mut length_bytes = [0_u8; 8];
1609 match input.read(&mut length_bytes[..1]) {
1610 Ok(0) => return Ok(()),
1611 Ok(1) => {}
1612 Ok(_) => unreachable!("one-byte read returned more than one byte"),
1613 Err(error) => return Err(format!("could not read ingest frame length: {error}")),
1614 }
1615 input
1616 .read_exact(&mut length_bytes[1..])
1617 .map_err(|error| format!("truncated ingest frame length: {error}"))?;
1618 let length = u64::from_be_bytes(length_bytes);
1619 if length == 0 || length > MAX_FRAME_BYTES {
1620 return Err(format!("invalid ingest frame length {length}"));
1621 }
1622 let frame_length = usize::try_from(length)
1623 .map_err(|_| format!("ingest frame length {length} does not fit in memory"))?;
1624 let mut frame = vec![0_u8; frame_length];
1625 input
1626 .read_exact(&mut frame)
1627 .map_err(|error| format!("could not read ingest frame: {error}"))?;
1628 let operation: ArchiveOperation = serde_json::from_slice(&frame)
1629 .map_err(|error| format!("invalid ingest frame: {error}"))?;
1630 let request_id = operation.request_id();
1631 let result = apply_operation(&mut archive, operation);
1632 let ack = match result {
1633 Ok(archive_ref) => ArchiveAck::success(request_id, archive_ref),
1634 Err(error) => ArchiveAck::failure(request_id, error),
1635 };
1636 serde_json::to_writer(&mut output, &ack)
1637 .map_err(|error| format!("could not write ingest acknowledgement: {error}"))?;
1638 output
1639 .write_all(b"\n")
1640 .and_then(|()| output.flush())
1641 .map_err(|error| format!("could not flush ingest acknowledgement: {error}"))?;
1642 }
1643}
1644
1645#[cfg(test)]
1646fn apply_operation(
1647 archive: &mut Archive,
1648 operation: ArchiveOperation,
1649) -> Result<Option<String>, String> {
1650 let schema_version = operation.schema_version();
1651 if schema_version != INGEST_SCHEMA_VERSION {
1652 return Err(format!(
1653 "unsupported ingest schema version {schema_version}; expected {INGEST_SCHEMA_VERSION}"
1654 ));
1655 }
1656 match operation {
1657 ArchiveOperation::BeginCall {
1658 session,
1659 call,
1660 input_before,
1661 input_after,
1662 captured_at_ms,
1663 ..
1664 } => archive
1665 .begin_call(&session, &call, &input_before, &input_after, captured_at_ms)
1666 .map(Some),
1667 ArchiveOperation::ResultBefore {
1668 session,
1669 source_call_id,
1670 result,
1671 full_output_path,
1672 captured_at_ms,
1673 ..
1674 } => archive
1675 .result_before(
1676 &session,
1677 &source_call_id,
1678 &result,
1679 full_output_path.as_deref(),
1680 captured_at_ms,
1681 )
1682 .map(|()| None),
1683 ArchiveOperation::ResultText {
1684 session,
1685 source_call_id,
1686 text,
1687 source_completeness,
1688 captured_at_ms,
1689 ..
1690 } => archive
1691 .result_text(
1692 &session,
1693 &source_call_id,
1694 &text,
1695 source_completeness,
1696 captured_at_ms,
1697 )
1698 .map(Some),
1699 ArchiveOperation::StageResult {
1700 session,
1701 source_call_id,
1702 result,
1703 is_error,
1704 captured_at_ms,
1705 ..
1706 } => archive
1707 .stage_result(&session, &source_call_id, &result, is_error, captured_at_ms)
1708 .map(|()| None),
1709 ArchiveOperation::FinishCall {
1710 session,
1711 source_call_id,
1712 result,
1713 is_error,
1714 require_pre_result,
1715 finished_at_ms,
1716 ..
1717 } => archive
1718 .finish_call(
1719 &session,
1720 &source_call_id,
1721 &result,
1722 is_error,
1723 require_pre_result,
1724 finished_at_ms,
1725 )
1726 .map(|()| None),
1727 ArchiveOperation::UpdateFinalResult {
1728 session,
1729 source_call_id,
1730 result,
1731 is_error,
1732 finished_at_ms,
1733 ..
1734 } => archive
1735 .update_final_result(&session, &source_call_id, &result, is_error, finished_at_ms)
1736 .map(|()| None),
1737 ArchiveOperation::CaptureStreams { .. }
1738 | ArchiveOperation::CapturePassthroughStreams { .. }
1739 | ArchiveOperation::PruneBefore { .. } => {
1740 Err("operation is available only through the archive broker".to_owned())
1741 }
1742 }
1743}
1744
1745pub fn archive_path() -> Result<PathBuf, String> {
1751 config::load().map(|resolved| resolved.archive.path)
1752}
1753
1754fn enable_wal_mode(connection: &Connection) -> Result<(), String> {
1755 const ATTEMPT_TIMEOUT: Duration = Duration::from_millis(25);
1756 const RETRY_DELAY: Duration = Duration::from_millis(10);
1757 connection
1758 .busy_timeout(ATTEMPT_TIMEOUT)
1759 .map_err(|error| format!("could not set WAL retry timeout: {error}"))?;
1760 let deadline = Instant::now() + ARCHIVE_BUSY_TIMEOUT;
1761 loop {
1762 match connection.pragma_update(None, "journal_mode", "WAL") {
1763 Ok(()) => {
1764 connection
1765 .busy_timeout(ARCHIVE_BUSY_TIMEOUT)
1766 .map_err(|error| format!("could not restore archive busy timeout: {error}"))?;
1767 return Ok(());
1768 }
1769 Err(error) if is_lock_error(&error) && Instant::now() < deadline => {
1770 std::thread::sleep(RETRY_DELAY);
1771 }
1772 Err(error) => return Err(format!("could not enable archive WAL mode: {error}")),
1773 }
1774 }
1775}
1776
1777fn is_lock_error(error: &rusqlite::Error) -> bool {
1778 matches!(
1779 error,
1780 rusqlite::Error::SqliteFailure(inner, _)
1781 if matches!(inner.code, ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked)
1782 )
1783}
1784
1785fn initialize_schema(connection: &mut Connection) -> Result<(), String> {
1786 connection
1787 .pragma_update(None, "auto_vacuum", "INCREMENTAL")
1788 .map_err(|error| format!("could not enable incremental auto-vacuum: {error}"))?;
1789 let transaction = connection
1790 .transaction_with_behavior(TransactionBehavior::Exclusive)
1791 .map_err(|error| format!("could not start archive migration: {error}"))?;
1792 let version: i64 = transaction
1793 .pragma_query_value(None, "user_version", |row| row.get(0))
1794 .map_err(|error| format!("could not recheck archive schema version: {error}"))?;
1795 match version {
1796 0 => {
1797 transaction
1798 .execute_batch(SCHEMA)
1799 .map_err(|error| format!("could not create archive schema: {error}"))?;
1800 transaction
1801 .pragma_update(None, "user_version", SCHEMA_VERSION)
1802 .map_err(|error| format!("could not set archive schema version: {error}"))?;
1803 }
1804 SCHEMA_VERSION => {}
1805 other => {
1806 return Err(format!(
1807 "archive schema version changed to {other} while initializing; expected 0 or {SCHEMA_VERSION}"
1808 ));
1809 }
1810 }
1811 transaction
1812 .commit()
1813 .map_err(|error| format!("could not commit archive schema: {error}"))
1814}
1815
1816fn schema_is_current(connection: &Connection) -> Result<bool, String> {
1817 let archive_ref = table_has_column(connection, "tool_calls", "archive_ref")?;
1818 let source_completeness = table_has_column(connection, "snapshots", "source_completeness")?;
1819 if archive_ref != source_completeness {
1820 return Err("archive schema has a partial indexed-output migration".to_owned());
1821 }
1822 if !archive_ref {
1823 return Ok(false);
1824 }
1825 let snapshots_sql: String = connection
1826 .query_row(
1827 "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'snapshots'",
1828 [],
1829 |row| row.get(0),
1830 )
1831 .map_err(|error| format!("could not inspect snapshots schema: {error}"))?;
1832 Ok(snapshots_sql.contains("result_text") && snapshots_sql.contains("source_completeness"))
1833}
1834
1835fn table_has_column(connection: &Connection, table: &str, column: &str) -> Result<bool, String> {
1836 let query = format!("SELECT count(*) FROM pragma_table_info('{table}') WHERE name = ?1");
1837 connection
1838 .query_row(&query, [column], |row| row.get::<_, i64>(0))
1839 .map(|count| count == 1)
1840 .map_err(|error| format!("could not inspect {table}.{column}: {error}"))
1841}
1842
1843fn ensure_current_schema(connection: &mut Connection, path: &Path) -> Result<(), String> {
1844 if schema_is_current(connection)? {
1845 return Ok(());
1846 }
1847 create_migration_backup(connection, path)?;
1848 migrate_indexed_output_schema(connection)
1849}
1850
1851fn create_migration_backup(connection: &Connection, path: &Path) -> Result<(), String> {
1852 let parent = path
1853 .parent()
1854 .ok_or_else(|| "archive path has no parent for migration backup".to_owned())?;
1855 let required = archive_file_bytes(path)?.saturating_mul(2);
1856 let available = fs2::available_space(parent)
1857 .map_err(|error| format!("could not measure archive backup space: {error}"))?;
1858 if available < required {
1859 return Err(format!(
1860 "archive migration needs at least {required} free bytes, found {available}"
1861 ));
1862 }
1863 let mut backup_name = path
1864 .file_name()
1865 .ok_or_else(|| "archive path has no file name".to_owned())?
1866 .to_os_string();
1867 backup_name.push(".pre-indexed-output-v1.backup");
1868 let backup_path = parent.join(backup_name);
1869 if backup_path.exists() {
1870 return Err(format!(
1871 "archive migration backup already exists: {}",
1872 backup_path.display()
1873 ));
1874 }
1875 let temporary = NamedTempFile::new_in(parent)
1876 .map_err(|error| format!("could not create migration backup: {error}"))?;
1877 set_file_mode(temporary.path(), 0o600)?;
1878 let mut destination = Connection::open(temporary.path())
1879 .map_err(|error| format!("could not open migration backup: {error}"))?;
1880 let backup = Backup::new(connection, &mut destination)
1881 .map_err(|error| format!("could not start archive backup: {error}"))?;
1882 backup
1883 .run_to_completion(128, Duration::from_millis(10), None)
1884 .map_err(|error| format!("could not write archive backup: {error}"))?;
1885 drop(backup);
1886 let integrity: String = destination
1887 .query_row("PRAGMA integrity_check", [], |row| row.get(0))
1888 .map_err(|error| format!("could not verify migration backup: {error}"))?;
1889 if integrity != "ok" {
1890 return Err(format!(
1891 "migration backup integrity check failed: {integrity}"
1892 ));
1893 }
1894 destination
1895 .execute_batch("PRAGMA wal_checkpoint(TRUNCATE)")
1896 .map_err(|error| format!("could not checkpoint migration backup: {error}"))?;
1897 drop(destination);
1898 temporary
1899 .as_file()
1900 .sync_all()
1901 .map_err(|error| format!("could not sync migration backup: {error}"))?;
1902 temporary
1903 .persist(&backup_path)
1904 .map_err(|error| format!("could not install migration backup: {}", error.error))?;
1905 set_file_mode(&backup_path, 0o600)?;
1906 sync_directory(parent)
1907}
1908
1909#[expect(
1910 clippy::too_many_lines,
1911 reason = "the migration is one ordered transaction with shared invariants and rollback"
1912)]
1913fn migrate_indexed_output_schema(connection: &mut Connection) -> Result<(), String> {
1914 let transaction = connection
1915 .transaction_with_behavior(TransactionBehavior::Exclusive)
1916 .map_err(|error| format!("could not start indexed-output migration: {error}"))?;
1917 if table_has_column(&transaction, "tool_calls", "archive_ref")?
1918 || table_has_column(&transaction, "snapshots", "source_completeness")?
1919 {
1920 return Err("archive schema changed while preparing indexed-output migration".to_owned());
1921 }
1922 let calls_before: i64 = transaction
1923 .query_row("SELECT count(*) FROM tool_calls", [], |row| row.get(0))
1924 .map_err(|error| format!("could not count calls before migration: {error}"))?;
1925 let snapshots_before: i64 = transaction
1926 .query_row("SELECT count(*) FROM snapshots", [], |row| row.get(0))
1927 .map_err(|error| format!("could not count snapshots before migration: {error}"))?;
1928 transaction
1929 .execute_batch("ALTER TABLE tool_calls ADD COLUMN archive_ref TEXT")
1930 .map_err(|error| format!("could not add archive references: {error}"))?;
1931 transaction
1932 .execute(
1933 "UPDATE tool_calls SET archive_ref = 'yr_' || lower(hex(randomblob(16)))",
1934 [],
1935 )
1936 .map_err(|error| format!("could not backfill archive references: {error}"))?;
1937 for _ in 0..3 {
1938 let duplicates: i64 = transaction
1939 .query_row(
1940 "SELECT count(*) - count(DISTINCT archive_ref) FROM tool_calls",
1941 [],
1942 |row| row.get(0),
1943 )
1944 .map_err(|error| format!("could not check archive reference collisions: {error}"))?;
1945 if duplicates == 0 {
1946 break;
1947 }
1948 transaction
1949 .execute(
1950 "UPDATE tool_calls
1951 SET archive_ref = 'yr_' || lower(hex(randomblob(16)))
1952 WHERE id IN (
1953 SELECT later.id
1954 FROM tool_calls later
1955 JOIN tool_calls earlier
1956 ON earlier.archive_ref = later.archive_ref AND earlier.id < later.id
1957 )",
1958 [],
1959 )
1960 .map_err(|error| format!("could not retry archive reference collisions: {error}"))?;
1961 }
1962 let duplicates: i64 = transaction
1963 .query_row(
1964 "SELECT count(*) - count(DISTINCT archive_ref) FROM tool_calls",
1965 [],
1966 |row| row.get(0),
1967 )
1968 .map_err(|error| format!("could not verify archive reference uniqueness: {error}"))?;
1969 if duplicates != 0 {
1970 return Err("archive reference collisions remained after three retries".to_owned());
1971 }
1972 transaction
1973 .execute_batch(
1974 "CREATE TABLE tool_calls_new (
1975 id INTEGER PRIMARY KEY,
1976 session_id INTEGER NOT NULL REFERENCES sessions(id),
1977 source_call_id TEXT NOT NULL,
1978 archive_ref TEXT NOT NULL UNIQUE CHECK (
1979 length(archive_ref) = 35
1980 AND substr(archive_ref, 1, 3) = 'yr_'
1981 AND substr(archive_ref, 4) NOT GLOB '*[^0-9a-f]*'
1982 ),
1983 tool_name TEXT NOT NULL,
1984 provider TEXT,
1985 model TEXT,
1986 working_directory TEXT,
1987 started_at_ms INTEGER NOT NULL,
1988 requires_streams INTEGER NOT NULL CHECK (requires_streams IN (0, 1)),
1989 finished_at_ms INTEGER,
1990 status TEXT NOT NULL CHECK (status IN ('started', 'finished')),
1991 is_error INTEGER CHECK (is_error IN (0, 1)),
1992 executed INTEGER CHECK (executed IN (0, 1)),
1993 UNIQUE (session_id, source_call_id)
1994 );
1995 INSERT INTO tool_calls_new
1996 SELECT id, session_id, source_call_id, archive_ref, tool_name, provider, model,
1997 working_directory, started_at_ms, requires_streams, finished_at_ms, status,
1998 is_error, executed
1999 FROM tool_calls;
2000 CREATE TABLE snapshots_new (
2001 tool_call_id INTEGER NOT NULL REFERENCES tool_calls_new(id) ON DELETE CASCADE,
2002 subject TEXT NOT NULL CHECK (subject IN ('input', 'result', 'result_text', 'source_output', 'stdout', 'stderr')),
2003 stage TEXT NOT NULL CHECK (stage IN ('before', 'after')),
2004 media_type TEXT NOT NULL,
2005 source_completeness TEXT CHECK (source_completeness IN ('complete', 'incomplete', 'unknown')),
2006 captured_at_ms INTEGER NOT NULL,
2007 payload_sha256 BLOB NOT NULL REFERENCES payloads(sha256),
2008 PRIMARY KEY (tool_call_id, subject, stage),
2009 CHECK (
2010 (subject = 'result_text' AND stage = 'before' AND source_completeness IS NOT NULL)
2011 OR (subject <> 'result_text' AND source_completeness IS NULL)
2012 )
2013 );
2014 INSERT INTO snapshots_new (
2015 tool_call_id, subject, stage, media_type, source_completeness,
2016 captured_at_ms, payload_sha256
2017 )
2018 SELECT tool_call_id, subject, stage, media_type, NULL,
2019 captured_at_ms, payload_sha256
2020 FROM snapshots;
2021 DROP TABLE snapshots;
2022 DROP TABLE tool_calls;
2023 ALTER TABLE tool_calls_new RENAME TO tool_calls;
2024 ALTER TABLE snapshots_new RENAME TO snapshots;
2025 CREATE INDEX tool_calls_started_at_idx ON tool_calls(started_at_ms);
2026 CREATE INDEX tool_calls_tool_name_idx ON tool_calls(tool_name);
2027 CREATE INDEX snapshots_payload_idx ON snapshots(payload_sha256);",
2028 )
2029 .map_err(|error| format!("could not rebuild indexed-output tables: {error}"))?;
2030 let calls_after: i64 = transaction
2031 .query_row("SELECT count(*) FROM tool_calls", [], |row| row.get(0))
2032 .map_err(|error| format!("could not count calls after migration: {error}"))?;
2033 let snapshots_after: i64 = transaction
2034 .query_row("SELECT count(*) FROM snapshots", [], |row| row.get(0))
2035 .map_err(|error| format!("could not count snapshots after migration: {error}"))?;
2036 if calls_before != calls_after || snapshots_before != snapshots_after {
2037 return Err("archive row counts changed during migration".to_owned());
2038 }
2039 let invalid_refs: i64 = transaction
2040 .query_row(
2041 "SELECT count(*) FROM tool_calls
2042 WHERE length(archive_ref) != 35
2043 OR substr(archive_ref, 1, 3) != 'yr_'
2044 OR substr(archive_ref, 4) GLOB '*[^0-9a-f]*'",
2045 [],
2046 |row| row.get(0),
2047 )
2048 .map_err(|error| format!("could not verify migrated archive references: {error}"))?;
2049 let foreign_keys: i64 = transaction
2050 .query_row("SELECT count(*) FROM pragma_foreign_key_check", [], |row| {
2051 row.get(0)
2052 })
2053 .map_err(|error| format!("could not verify migrated foreign keys: {error}"))?;
2054 if invalid_refs != 0 || foreign_keys != 0 {
2055 return Err(format!(
2056 "archive migration verification failed: {invalid_refs} invalid references, {foreign_keys} foreign key violations"
2057 ));
2058 }
2059 transaction
2060 .commit()
2061 .map_err(|error| format!("could not commit indexed-output migration: {error}"))
2062}
2063
2064fn prepared_kind_bytes(kind: &PreparedOperationKind) -> u64 {
2065 let payload = |value: &PreparedPayload| value.stored_length;
2066 let snapshot = |value: &PreparedSnapshot| payload(&value.payload);
2067 match kind {
2068 PreparedOperationKind::BeginCall { before, after, .. } => payload(before) + payload(after),
2069 PreparedOperationKind::ResultBefore {
2070 result,
2071 source_output,
2072 ..
2073 } => payload(result) + source_output.as_ref().map_or(0, snapshot),
2074 PreparedOperationKind::ResultText { text, .. } => payload(text),
2075 PreparedOperationKind::StageResult { result, .. }
2076 | PreparedOperationKind::FinishCall { result, .. }
2077 | PreparedOperationKind::UpdateFinalResult { result, .. } => payload(result),
2078 PreparedOperationKind::CaptureStreams {
2079 stdout_before,
2080 stderr_before,
2081 stdout_after,
2082 stderr_after,
2083 ..
2084 } => {
2085 snapshot(stdout_before)
2086 + snapshot(stderr_before)
2087 + snapshot(stdout_after)
2088 + snapshot(stderr_after)
2089 }
2090 PreparedOperationKind::PruneBefore { .. } => 0,
2091 }
2092}
2093
2094fn prepare_payload(reader: &mut (impl Read + Seek)) -> Result<PreparedPayload, String> {
2095 reader
2096 .seek(SeekFrom::Start(0))
2097 .map_err(|error| format!("could not rewind payload for preparation: {error}"))?;
2098 let mut raw = NamedTempFile::new()
2099 .map_err(|error| format!("could not create prepared payload spool: {error}"))?;
2100 let mut hasher = Sha256::new();
2101 let mut uncompressed_length = 0_u64;
2102 let mut buffer = vec![0_u8; 64 * 1024];
2103 loop {
2104 let count = reader
2105 .read(&mut buffer)
2106 .map_err(|error| format!("could not read payload for preparation: {error}"))?;
2107 if count == 0 {
2108 break;
2109 }
2110 raw.write_all(&buffer[..count])
2111 .map_err(|error| format!("could not spool prepared payload: {error}"))?;
2112 hasher.update(&buffer[..count]);
2113 uncompressed_length = uncompressed_length.saturating_add(count as u64);
2114 if uncompressed_length > MAX_FRAME_BYTES {
2115 return Err(format!(
2116 "prepared payload is {uncompressed_length} bytes; maximum is {MAX_FRAME_BYTES}"
2117 ));
2118 }
2119 }
2120 raw.flush()
2121 .map_err(|error| format!("could not flush prepared payload: {error}"))?;
2122 raw.seek(SeekFrom::Start(0))
2123 .map_err(|error| format!("could not rewind prepared payload: {error}"))?;
2124 let mut compressed = NamedTempFile::new()
2125 .map_err(|error| format!("could not create compressed payload spool: {error}"))?;
2126 zstd::stream::copy_encode(&mut raw, &mut compressed, ZSTD_LEVEL)
2127 .map_err(|error| format!("could not compress prepared payload: {error}"))?;
2128 compressed
2129 .flush()
2130 .map_err(|error| format!("could not flush compressed prepared payload: {error}"))?;
2131 let compressed_length = compressed
2132 .as_file()
2133 .metadata()
2134 .map_err(|error| format!("could not stat compressed prepared payload: {error}"))?
2135 .len();
2136 let use_compressed = compressed_length.saturating_mul(100)
2137 <= uncompressed_length.saturating_mul(COMPRESSION_PERCENT);
2138 let (compression, stored_length, mut stored) = if use_compressed {
2139 ("zstd", compressed_length, compressed)
2140 } else {
2141 ("none", uncompressed_length, raw)
2142 };
2143 stored
2144 .seek(SeekFrom::Start(0))
2145 .map_err(|error| format!("could not rewind stored prepared payload: {error}"))?;
2146 Ok(PreparedPayload {
2147 sha: hasher.finalize().to_vec(),
2148 compression,
2149 uncompressed_length,
2150 stored_length,
2151 stored,
2152 })
2153}
2154
2155fn prepare_path_snapshot(path: &Path, expected_uid: u32) -> Result<PreparedSnapshot, String> {
2156 let before = fs::symlink_metadata(path)
2157 .map_err(|error| format!("could not inspect archive source file: {error}"))?;
2158 require_private_source(path, &before, expected_uid)?;
2159 let mut file = OpenOptions::new()
2160 .read(true)
2161 .open(path)
2162 .map_err(|error| format!("could not open archive source file: {error}"))?;
2163 let opened = file
2164 .metadata()
2165 .map_err(|error| format!("could not inspect open archive source file: {error}"))?;
2166 require_same_source(path, &before, &opened)?;
2167 let media_type = stream_media_type(&mut file)?;
2168 let payload = prepare_payload(&mut file)?;
2169 let after = fs::symlink_metadata(path)
2170 .map_err(|error| format!("could not recheck archive source file: {error}"))?;
2171 require_same_source(path, &opened, &after)?;
2172 Ok(PreparedSnapshot {
2173 media_type,
2174 payload,
2175 })
2176}
2177
2178#[cfg(unix)]
2179fn archive_owner_uid(path: &Path) -> Result<u32, String> {
2180 use std::os::unix::fs::MetadataExt as _;
2181 let parent = path
2182 .parent()
2183 .ok_or_else(|| format!("archive path {} has no parent", path.display()))?;
2184 parent
2185 .metadata()
2186 .map(|metadata| metadata.uid())
2187 .map_err(|error| {
2188 format!(
2189 "could not inspect archive directory {}: {error}",
2190 parent.display()
2191 )
2192 })
2193}
2194
2195#[cfg(not(unix))]
2196fn archive_owner_uid(_path: &Path) -> Result<u32, String> {
2197 Ok(0)
2198}
2199
2200#[cfg(unix)]
2201fn require_private_source(
2202 _path: &Path,
2203 metadata: &fs::Metadata,
2204 expected_uid: u32,
2205) -> Result<(), String> {
2206 use std::os::unix::fs::{MetadataExt as _, PermissionsExt as _};
2207 if metadata.file_type().is_symlink() || !metadata.is_file() {
2208 return Err("archive source is not a real regular file".to_owned());
2209 }
2210 if metadata.uid() != expected_uid {
2211 return Err("archive source has the wrong owner".to_owned());
2212 }
2213 if metadata.permissions().mode() & 0o077 != 0 {
2214 return Err("archive source is accessible by other users".to_owned());
2215 }
2216 if metadata.len() > MAX_FRAME_BYTES {
2217 return Err(format!(
2218 "archive source is {} bytes; maximum is {MAX_FRAME_BYTES}",
2219 metadata.len()
2220 ));
2221 }
2222 Ok(())
2223}
2224
2225#[cfg(not(unix))]
2226fn require_private_source(
2227 _path: &Path,
2228 metadata: &fs::Metadata,
2229 _expected_uid: u32,
2230) -> Result<(), String> {
2231 if !metadata.is_file() {
2232 return Err("archive source is not a regular file".to_owned());
2233 }
2234 Ok(())
2235}
2236
2237#[cfg(unix)]
2238fn require_same_source(
2239 _path: &Path,
2240 expected: &fs::Metadata,
2241 actual: &fs::Metadata,
2242) -> Result<(), String> {
2243 use std::os::unix::fs::MetadataExt as _;
2244 if (
2245 expected.dev(),
2246 expected.ino(),
2247 expected.len(),
2248 expected.mtime(),
2249 expected.mtime_nsec(),
2250 ) != (
2251 actual.dev(),
2252 actual.ino(),
2253 actual.len(),
2254 actual.mtime(),
2255 actual.mtime_nsec(),
2256 ) {
2257 return Err("archive source changed while being read".to_owned());
2258 }
2259 Ok(())
2260}
2261
2262#[cfg(not(unix))]
2263fn require_same_source(
2264 _path: &Path,
2265 expected: &fs::Metadata,
2266 actual: &fs::Metadata,
2267) -> Result<(), String> {
2268 if expected.len() != actual.len() {
2269 return Err("archive source changed while being read".to_owned());
2270 }
2271 Ok(())
2272}
2273
2274fn insert_prepared_payload(
2275 connection: &Connection,
2276 payload: &mut PreparedPayload,
2277) -> Result<Vec<u8>, String> {
2278 let exists = connection
2279 .query_row(
2280 "SELECT 1 FROM payloads WHERE sha256 = ?1",
2281 [&payload.sha],
2282 |_| Ok(()),
2283 )
2284 .optional()
2285 .map_err(|error| format!("could not check prepared payload digest: {error}"))?
2286 .is_some();
2287 if exists {
2288 return Ok(payload.sha.clone());
2289 }
2290 let stored_length = i64::try_from(payload.stored_length).map_err(|_| {
2291 format!(
2292 "payload is too large for SQLite: {} bytes",
2293 payload.stored_length
2294 )
2295 })?;
2296 let raw_length = i64::try_from(payload.uncompressed_length).map_err(|_| {
2297 format!(
2298 "payload is too large: {} bytes",
2299 payload.uncompressed_length
2300 )
2301 })?;
2302 connection
2303 .execute(
2304 "INSERT INTO payloads (sha256, compression, uncompressed_byte_length, body)
2305 VALUES (?1, ?2, ?3, zeroblob(?4))",
2306 params![payload.sha, payload.compression, raw_length, stored_length],
2307 )
2308 .map_err(|error| format!("could not insert prepared payload: {error}"))?;
2309 let rowid: i64 = connection
2310 .query_row(
2311 "SELECT rowid FROM payloads WHERE sha256 = ?1",
2312 [&payload.sha],
2313 |row| row.get(0),
2314 )
2315 .map_err(|error| format!("could not locate prepared payload body: {error}"))?;
2316 let mut blob = connection
2317 .blob_open(MAIN_DB, "payloads", "body", rowid, false)
2318 .map_err(|error| format!("could not open prepared payload body: {error}"))?;
2319 payload
2320 .stored
2321 .seek(SeekFrom::Start(0))
2322 .map_err(|error| format!("could not rewind prepared payload body: {error}"))?;
2323 copy_exact(&mut payload.stored, &mut blob, payload.stored_length)?;
2324 Ok(payload.sha.clone())
2325}
2326
2327#[expect(
2328 clippy::too_many_arguments,
2329 reason = "snapshot identity and source metadata are one strict database record"
2330)]
2331fn insert_prepared_snapshot(
2332 connection: &Connection,
2333 call_id: i64,
2334 subject: &str,
2335 stage: &str,
2336 media_type: &str,
2337 source_completeness: Option<&str>,
2338 captured_at_ms: i64,
2339 payload: &mut PreparedPayload,
2340) -> Result<(), String> {
2341 let sha = insert_prepared_payload(connection, payload)?;
2342 let inserted = connection
2343 .execute(
2344 "INSERT INTO snapshots (
2345 tool_call_id, subject, stage, media_type, source_completeness,
2346 captured_at_ms, payload_sha256
2347 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
2348 ON CONFLICT(tool_call_id, subject, stage) DO NOTHING",
2349 params![
2350 call_id,
2351 subject,
2352 stage,
2353 media_type,
2354 source_completeness,
2355 captured_at_ms,
2356 sha
2357 ],
2358 )
2359 .map_err(|error| format!("could not store {subject}/{stage} snapshot: {error}"))?;
2360 if inserted == 0 {
2361 let existing: (Vec<u8>, Option<String>) = connection
2362 .query_row(
2363 "SELECT payload_sha256, source_completeness FROM snapshots
2364 WHERE tool_call_id = ?1 AND subject = ?2 AND stage = ?3",
2365 params![call_id, subject, stage],
2366 |row| Ok((row.get(0)?, row.get(1)?)),
2367 )
2368 .map_err(|error| format!("could not check existing snapshot: {error}"))?;
2369 if existing != (sha, source_completeness.map(str::to_owned)) {
2370 return Err(format!(
2371 "snapshot {subject}/{stage} already exists with different content or completeness"
2372 ));
2373 }
2374 }
2375 Ok(())
2376}
2377
2378fn replace_prepared_snapshot(
2379 connection: &Connection,
2380 call_id: i64,
2381 subject: &str,
2382 stage: &str,
2383 media_type: &str,
2384 captured_at_ms: i64,
2385 payload: &mut PreparedPayload,
2386) -> Result<Vec<u8>, String> {
2387 let replaced_sha: Vec<u8> = connection
2388 .query_row(
2389 "SELECT payload_sha256 FROM snapshots
2390 WHERE tool_call_id = ?1 AND subject = ?2 AND stage = ?3",
2391 params![call_id, subject, stage],
2392 |row| row.get(0),
2393 )
2394 .map_err(|error| format!("could not find {subject}/{stage} snapshot: {error}"))?;
2395 let sha = insert_prepared_payload(connection, payload)?;
2396 let updated = connection
2397 .execute(
2398 "UPDATE snapshots
2399 SET media_type = ?1, captured_at_ms = ?2, payload_sha256 = ?3
2400 WHERE tool_call_id = ?4 AND subject = ?5 AND stage = ?6",
2401 params![media_type, captured_at_ms, sha, call_id, subject, stage],
2402 )
2403 .map_err(|error| format!("could not replace {subject}/{stage} snapshot: {error}"))?;
2404 if updated != 1 {
2405 return Err(format!("snapshot {subject}/{stage} does not exist"));
2406 }
2407 Ok(replaced_sha)
2408}
2409
2410#[expect(
2411 clippy::too_many_lines,
2412 reason = "the exhaustive operation mapping keeps one transaction policy"
2413)]
2414fn apply_prepared_operation(
2415 connection: &Connection,
2416 operation: &mut PreparedOperationKind,
2417) -> Result<Option<String>, String> {
2418 match operation {
2419 PreparedOperationKind::BeginCall {
2420 session,
2421 call,
2422 before,
2423 after,
2424 captured_at_ms,
2425 } => {
2426 let session_id = ensure_session(connection, session)?;
2427 let call_id = ensure_call(connection, session_id, call)?;
2428 let archive_ref = call_archive_ref(connection, call_id)?;
2429 insert_prepared_snapshot(
2430 connection,
2431 call_id,
2432 "input",
2433 "before",
2434 "application/json",
2435 None,
2436 *captured_at_ms,
2437 before,
2438 )?;
2439 insert_prepared_snapshot(
2440 connection,
2441 call_id,
2442 "input",
2443 "after",
2444 "application/json",
2445 None,
2446 *captured_at_ms,
2447 after,
2448 )?;
2449 Ok(Some(archive_ref))
2450 }
2451 PreparedOperationKind::ResultBefore {
2452 session,
2453 source_call_id,
2454 result,
2455 source_output,
2456 captured_at_ms,
2457 } => {
2458 let call_id = find_call(connection, session, source_call_id)?;
2459 insert_prepared_snapshot(
2460 connection,
2461 call_id,
2462 "result",
2463 "before",
2464 "application/json",
2465 None,
2466 *captured_at_ms,
2467 result,
2468 )?;
2469 if let Some(source) = source_output {
2470 insert_prepared_snapshot(
2471 connection,
2472 call_id,
2473 "source_output",
2474 "before",
2475 source.media_type,
2476 None,
2477 *captured_at_ms,
2478 &mut source.payload,
2479 )?;
2480 }
2481 Ok(None)
2482 }
2483 PreparedOperationKind::ResultText {
2484 session,
2485 source_call_id,
2486 text,
2487 completeness,
2488 captured_at_ms,
2489 } => {
2490 let call_id = find_call(connection, session, source_call_id)?;
2491 insert_prepared_snapshot(
2492 connection,
2493 call_id,
2494 "result_text",
2495 "before",
2496 "text/plain; charset=utf-8",
2497 Some(completeness.as_str()),
2498 *captured_at_ms,
2499 text,
2500 )?;
2501 call_archive_ref(connection, call_id).map(Some)
2502 }
2503 PreparedOperationKind::StageResult {
2504 session,
2505 source_call_id,
2506 result,
2507 is_error,
2508 captured_at_ms,
2509 } => {
2510 let call_id = find_call(connection, session, source_call_id)?;
2511 require_call_status(connection, call_id, source_call_id, "started")?;
2512 validate_completion(connection, call_id, true)?;
2513 insert_prepared_snapshot(
2514 connection,
2515 call_id,
2516 "result",
2517 "after",
2518 "application/json",
2519 None,
2520 *captured_at_ms,
2521 result,
2522 )?;
2523 connection
2524 .execute(
2525 "UPDATE tool_calls SET is_error = ?1, executed = 1 WHERE id = ?2",
2526 params![i64::from(*is_error), call_id],
2527 )
2528 .map_err(|error| format!("could not stage tool result: {error}"))?;
2529 Ok(None)
2530 }
2531 PreparedOperationKind::FinishCall {
2532 session,
2533 source_call_id,
2534 result,
2535 is_error,
2536 require_pre_result,
2537 finished_at_ms,
2538 } => {
2539 let call_id = find_call(connection, session, source_call_id)?;
2540 validate_completion(connection, call_id, *require_pre_result)?;
2541 insert_prepared_snapshot(
2542 connection,
2543 call_id,
2544 "result",
2545 "after",
2546 "application/json",
2547 None,
2548 *finished_at_ms,
2549 result,
2550 )?;
2551 connection
2552 .execute(
2553 "UPDATE tool_calls
2554 SET finished_at_ms = ?1, status = 'finished', is_error = ?2, executed = ?3
2555 WHERE id = ?4",
2556 params![
2557 *finished_at_ms,
2558 i64::from(*is_error),
2559 i64::from(*require_pre_result),
2560 call_id
2561 ],
2562 )
2563 .map_err(|error| format!("could not finish tool call: {error}"))?;
2564 Ok(None)
2565 }
2566 PreparedOperationKind::UpdateFinalResult {
2567 session,
2568 source_call_id,
2569 result,
2570 is_error,
2571 finished_at_ms,
2572 } => {
2573 let call_id = find_call(connection, session, source_call_id)?;
2574 let status: String = connection
2575 .query_row(
2576 "SELECT status FROM tool_calls WHERE id = ?1",
2577 [call_id],
2578 |row| row.get(0),
2579 )
2580 .map_err(|error| format!("could not read final tool call status: {error}"))?;
2581 if status == "finished" {
2582 let stored: (Vec<u8>, bool, i64, bool) = connection
2583 .query_row(
2584 "SELECT s.payload_sha256, c.is_error, c.finished_at_ms, c.executed
2585 FROM tool_calls c
2586 JOIN snapshots s ON s.tool_call_id = c.id
2587 WHERE c.id = ?1 AND s.subject = 'result' AND s.stage = 'after'",
2588 [call_id],
2589 |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
2590 )
2591 .map_err(|error| format!("could not check finalized tool result: {error}"))?;
2592 if stored == (result.sha.clone(), *is_error, *finished_at_ms, true) {
2593 return Ok(None);
2594 }
2595 return Err(format!(
2596 "tool call {source_call_id} was already finalized with different content or metadata"
2597 ));
2598 }
2599 if status != "started" {
2600 return Err(format!(
2601 "tool call {source_call_id} has status {status}, expected started"
2602 ));
2603 }
2604 validate_completion(connection, call_id, true)?;
2605 let replaced_sha = replace_prepared_snapshot(
2606 connection,
2607 call_id,
2608 "result",
2609 "after",
2610 "application/json",
2611 *finished_at_ms,
2612 result,
2613 )?;
2614 connection
2615 .execute(
2616 "UPDATE tool_calls
2617 SET finished_at_ms = ?1, status = 'finished', is_error = ?2, executed = 1
2618 WHERE id = ?3",
2619 params![*finished_at_ms, i64::from(*is_error), call_id],
2620 )
2621 .map_err(|error| format!("could not finalize staged tool call: {error}"))?;
2622 connection
2623 .execute(
2624 "DELETE FROM payloads
2625 WHERE sha256 = ?1
2626 AND NOT EXISTS (
2627 SELECT 1 FROM snapshots WHERE snapshots.payload_sha256 = payloads.sha256
2628 )",
2629 [&replaced_sha],
2630 )
2631 .map_err(|error| format!("could not remove replaced final payload: {error}"))?;
2632 Ok(None)
2633 }
2634 PreparedOperationKind::CaptureStreams {
2635 session,
2636 source_call_id,
2637 captured_at_ms,
2638 stdout_before,
2639 stderr_before,
2640 stdout_after,
2641 stderr_after,
2642 } => {
2643 let call_id = find_call(connection, session, source_call_id)?;
2644 for (subject, stage, snapshot) in [
2645 ("stdout", "before", stdout_before),
2646 ("stderr", "before", stderr_before),
2647 ("stdout", "after", stdout_after),
2648 ("stderr", "after", stderr_after),
2649 ] {
2650 insert_prepared_snapshot(
2651 connection,
2652 call_id,
2653 subject,
2654 stage,
2655 snapshot.media_type,
2656 None,
2657 *captured_at_ms,
2658 &mut snapshot.payload,
2659 )?;
2660 }
2661 Ok(None)
2662 }
2663 PreparedOperationKind::PruneBefore { timestamp_ms } => {
2664 let deleted = connection
2665 .execute(
2666 "DELETE FROM tool_calls WHERE status = 'finished' AND finished_at_ms < ?1",
2667 [*timestamp_ms],
2668 )
2669 .map_err(|error| format!("could not prune tool calls: {error}"))?;
2670 connection
2671 .execute(
2672 "DELETE FROM payloads WHERE NOT EXISTS (
2673 SELECT 1 FROM snapshots WHERE snapshots.payload_sha256 = payloads.sha256
2674 )",
2675 [],
2676 )
2677 .map_err(|error| format!("could not prune unreferenced payloads: {error}"))?;
2678 i64::try_from(deleted)
2679 .map(|count| Some(count.to_string()))
2680 .map_err(|_| "pruned call count does not fit in i64".to_owned())
2681 }
2682 }
2683}
2684
2685fn batch_error(context: &str, error: &rusqlite::Error) -> BatchWriteError {
2686 let message = format!("{context}: {error}");
2687 if is_lock_error(error) {
2688 BatchWriteError::Busy(message)
2689 } else {
2690 BatchWriteError::Permanent(message)
2691 }
2692}
2693
2694fn ensure_session(transaction: &Connection, session: &SessionIdentity) -> Result<i64, String> {
2695 transaction
2696 .execute(
2697 "INSERT INTO sessions (agent, account, source_session_id, started_at_ms)
2698 VALUES (?1, ?2, ?3, ?4)
2699 ON CONFLICT(agent, account, source_session_id) DO UPDATE SET
2700 started_at_ms = coalesce(sessions.started_at_ms, excluded.started_at_ms)",
2701 params![
2702 session.agent,
2703 session.account,
2704 session.source_session_id,
2705 session.started_at_ms
2706 ],
2707 )
2708 .map_err(|error| format!("could not store archive session: {error}"))?;
2709 transaction
2710 .query_row(
2711 "SELECT id FROM sessions WHERE agent = ?1 AND account = ?2 AND source_session_id = ?3",
2712 params![session.agent, session.account, session.source_session_id],
2713 |row| row.get(0),
2714 )
2715 .map_err(|error| format!("could not find archive session: {error}"))
2716}
2717
2718fn ensure_call(
2719 transaction: &Connection,
2720 session_id: i64,
2721 call: &CallIdentity,
2722) -> Result<i64, String> {
2723 let existing: Option<i64> = transaction
2724 .query_row(
2725 "SELECT id FROM tool_calls WHERE session_id = ?1 AND source_call_id = ?2",
2726 params![session_id, call.source_call_id],
2727 |row| row.get(0),
2728 )
2729 .optional()
2730 .map_err(|error| format!("could not check stored tool call: {error}"))?;
2731 if existing.is_none() {
2732 let mut inserted = false;
2733 for _ in 0..3 {
2734 let archive_ref = random_archive_ref(transaction)?;
2735 match transaction.execute(
2736 "INSERT INTO tool_calls (
2737 session_id, source_call_id, archive_ref, tool_name, provider, model,
2738 working_directory, started_at_ms, requires_streams, status
2739 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, 'started')",
2740 params![
2741 session_id,
2742 call.source_call_id,
2743 archive_ref,
2744 call.tool_name,
2745 call.provider,
2746 call.model,
2747 call.working_directory,
2748 call.started_at_ms,
2749 i64::from(call.requires_streams)
2750 ],
2751 ) {
2752 Ok(1) => {
2753 inserted = true;
2754 break;
2755 }
2756 Ok(_) => return Err("could not store tool call".to_owned()),
2757 Err(rusqlite::Error::SqliteFailure(inner, _))
2758 if inner.code == ErrorCode::ConstraintViolation => {}
2759 Err(error) => return Err(format!("could not store tool call: {error}")),
2760 }
2761 }
2762 if !inserted {
2763 return Err(
2764 "could not allocate a unique archive reference after three attempts".to_owned(),
2765 );
2766 }
2767 }
2768 let stored = transaction
2769 .query_row(
2770 "SELECT id, archive_ref, tool_name, provider, model, working_directory, started_at_ms, requires_streams
2771 FROM tool_calls WHERE session_id = ?1 AND source_call_id = ?2",
2772 params![session_id, call.source_call_id],
2773 |row| {
2774 Ok((
2775 row.get::<_, i64>(0)?,
2776 row.get::<_, String>(1)?,
2777 row.get::<_, String>(2)?,
2778 row.get::<_, Option<String>>(3)?,
2779 row.get::<_, Option<String>>(4)?,
2780 row.get::<_, Option<String>>(5)?,
2781 row.get::<_, i64>(6)?,
2782 row.get::<_, bool>(7)?,
2783 ))
2784 },
2785 )
2786 .map_err(|error| format!("could not find stored tool call: {error}"))?;
2787 validate_archive_ref(&stored.1)?;
2788 if stored.2 != call.tool_name
2789 || stored.3 != call.provider
2790 || stored.4 != call.model
2791 || stored.5 != call.working_directory
2792 || stored.6 != call.started_at_ms
2793 || stored.7 != call.requires_streams
2794 {
2795 return Err(format!(
2796 "tool call {} was already stored with different metadata",
2797 call.source_call_id
2798 ));
2799 }
2800 Ok(stored.0)
2801}
2802
2803fn random_archive_ref(transaction: &Connection) -> Result<String, String> {
2804 let archive_ref: String = transaction
2805 .query_row("SELECT 'yr_' || lower(hex(randomblob(16)))", [], |row| {
2806 row.get(0)
2807 })
2808 .map_err(|error| format!("could not generate archive reference: {error}"))?;
2809 validate_archive_ref(&archive_ref)?;
2810 Ok(archive_ref)
2811}
2812
2813fn validate_archive_ref(archive_ref: &str) -> Result<(), String> {
2814 if archive_ref.len() != ARCHIVE_REF_LEN
2815 || !archive_ref.starts_with(ARCHIVE_REF_PREFIX)
2816 || !archive_ref[ARCHIVE_REF_PREFIX.len()..]
2817 .bytes()
2818 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
2819 {
2820 return Err(
2821 "archive reference must be yr_ followed by 32 lowercase hexadecimal digits".to_owned(),
2822 );
2823 }
2824 Ok(())
2825}
2826
2827fn call_archive_ref(transaction: &Connection, call_id: i64) -> Result<String, String> {
2828 let archive_ref: String = transaction
2829 .query_row(
2830 "SELECT archive_ref FROM tool_calls WHERE id = ?1",
2831 [call_id],
2832 |row| row.get(0),
2833 )
2834 .map_err(|error| format!("could not read archive reference: {error}"))?;
2835 validate_archive_ref(&archive_ref)?;
2836 Ok(archive_ref)
2837}
2838
2839fn parse_completeness(value: &str) -> Result<SourceCompleteness, String> {
2840 match value {
2841 "complete" => Ok(SourceCompleteness::Complete),
2842 "incomplete" => Ok(SourceCompleteness::Incomplete),
2843 "unknown" => Ok(SourceCompleteness::Unknown),
2844 _ => Err(format!("invalid source completeness {value}")),
2845 }
2846}
2847
2848fn find_call(
2849 transaction: &Connection,
2850 session: &SessionIdentity,
2851 source_call_id: &str,
2852) -> Result<i64, String> {
2853 transaction
2854 .query_row(
2855 "SELECT c.id
2856 FROM tool_calls c
2857 JOIN sessions s ON s.id = c.session_id
2858 WHERE s.agent = ?1 AND s.account = ?2 AND s.source_session_id = ?3
2859 AND c.source_call_id = ?4",
2860 params![
2861 session.agent,
2862 session.account,
2863 session.source_session_id,
2864 source_call_id
2865 ],
2866 |row| row.get(0),
2867 )
2868 .map_err(|error| format!("could not find tool call {source_call_id}: {error}"))
2869}
2870
2871fn require_call_status(
2872 transaction: &Connection,
2873 call_id: i64,
2874 source_call_id: &str,
2875 expected: &str,
2876) -> Result<(), String> {
2877 let status: String = transaction
2878 .query_row(
2879 "SELECT status FROM tool_calls WHERE id = ?1",
2880 [call_id],
2881 |row| row.get(0),
2882 )
2883 .map_err(|error| format!("could not read tool call status: {error}"))?;
2884 if status != expected {
2885 return Err(format!(
2886 "tool call {source_call_id} has status {status}, expected {expected}"
2887 ));
2888 }
2889 Ok(())
2890}
2891
2892fn validate_completion(
2893 transaction: &Connection,
2894 call_id: i64,
2895 require_pre_result: bool,
2896) -> Result<(), String> {
2897 let requires_streams: bool = transaction
2898 .query_row(
2899 "SELECT requires_streams FROM tool_calls WHERE id = ?1",
2900 [call_id],
2901 |row| row.get(0),
2902 )
2903 .map_err(|error| format!("could not read tool call requirements: {error}"))?;
2904 if requires_streams && require_pre_result {
2905 let streams: i64 = transaction
2906 .query_row(
2907 "SELECT count(*) FROM snapshots
2908 WHERE tool_call_id = ?1
2909 AND subject IN ('stdout', 'stderr')
2910 AND stage IN ('before', 'after')",
2911 [call_id],
2912 |row| row.get(0),
2913 )
2914 .map_err(|error| format!("could not check shell stream snapshots: {error}"))?;
2915 if streams != 4 {
2916 return Err(format!(
2917 "shell tool call is missing stream snapshots: found {streams} of 4"
2918 ));
2919 }
2920 }
2921 if require_pre_result {
2922 let has_result: bool = transaction
2923 .query_row(
2924 "SELECT EXISTS(
2925 SELECT 1 FROM snapshots
2926 WHERE tool_call_id = ?1 AND subject = 'result' AND stage = 'before'
2927 )",
2928 [call_id],
2929 |row| row.get(0),
2930 )
2931 .map_err(|error| format!("could not check pre-YARP result snapshot: {error}"))?;
2932 if !has_result {
2933 return Err("executed tool call is missing its pre-YARP result".to_owned());
2934 }
2935 }
2936 Ok(())
2937}
2938
2939fn insert_snapshot_bytes(
2940 transaction: &Connection,
2941 call_id: i64,
2942 subject: &str,
2943 stage: &str,
2944 media_type: &str,
2945 captured_at_ms: i64,
2946 body: &[u8],
2947) -> Result<(), String> {
2948 insert_snapshot_reader(
2949 transaction,
2950 call_id,
2951 subject,
2952 stage,
2953 media_type,
2954 captured_at_ms,
2955 &mut Cursor::new(body),
2956 )
2957}
2958
2959fn replace_snapshot_bytes(
2960 transaction: &Connection,
2961 call_id: i64,
2962 subject: &str,
2963 stage: &str,
2964 media_type: &str,
2965 captured_at_ms: i64,
2966 body: &[u8],
2967) -> Result<Vec<u8>, String> {
2968 let replaced_sha: Vec<u8> = transaction
2969 .query_row(
2970 "SELECT payload_sha256 FROM snapshots
2971 WHERE tool_call_id = ?1 AND subject = ?2 AND stage = ?3",
2972 params![call_id, subject, stage],
2973 |row| row.get(0),
2974 )
2975 .map_err(|error| format!("could not find {subject}/{stage} snapshot: {error}"))?;
2976 let sha = insert_payload(transaction, &mut Cursor::new(body))?;
2977 let updated = transaction
2978 .execute(
2979 "UPDATE snapshots
2980 SET media_type = ?1, captured_at_ms = ?2, payload_sha256 = ?3
2981 WHERE tool_call_id = ?4 AND subject = ?5 AND stage = ?6",
2982 params![media_type, captured_at_ms, sha, call_id, subject, stage],
2983 )
2984 .map_err(|error| format!("could not replace {subject}/{stage} snapshot: {error}"))?;
2985 if updated != 1 {
2986 return Err(format!("snapshot {subject}/{stage} does not exist"));
2987 }
2988 Ok(replaced_sha)
2989}
2990
2991fn insert_snapshot_reader(
2992 transaction: &Connection,
2993 call_id: i64,
2994 subject: &str,
2995 stage: &str,
2996 media_type: &str,
2997 captured_at_ms: i64,
2998 reader: &mut (impl Read + Seek),
2999) -> Result<(), String> {
3000 insert_snapshot_reader_with_completeness(
3001 transaction,
3002 call_id,
3003 subject,
3004 stage,
3005 media_type,
3006 None,
3007 captured_at_ms,
3008 reader,
3009 )
3010}
3011
3012#[expect(
3013 clippy::too_many_arguments,
3014 reason = "snapshot identity and source metadata are one strict database record"
3015)]
3016fn insert_snapshot_reader_with_completeness(
3017 transaction: &Connection,
3018 call_id: i64,
3019 subject: &str,
3020 stage: &str,
3021 media_type: &str,
3022 source_completeness: Option<&str>,
3023 captured_at_ms: i64,
3024 reader: &mut (impl Read + Seek),
3025) -> Result<(), String> {
3026 let sha = insert_payload(transaction, reader)?;
3027 let inserted = transaction
3028 .execute(
3029 "INSERT INTO snapshots (
3030 tool_call_id, subject, stage, media_type, source_completeness,
3031 captured_at_ms, payload_sha256
3032 ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)
3033 ON CONFLICT(tool_call_id, subject, stage) DO NOTHING",
3034 params![
3035 call_id,
3036 subject,
3037 stage,
3038 media_type,
3039 source_completeness,
3040 captured_at_ms,
3041 sha
3042 ],
3043 )
3044 .map_err(|error| format!("could not store {subject}/{stage} snapshot: {error}"))?;
3045 if inserted == 0 {
3046 let existing: (Vec<u8>, Option<String>) = transaction
3047 .query_row(
3048 "SELECT payload_sha256, source_completeness FROM snapshots
3049 WHERE tool_call_id = ?1 AND subject = ?2 AND stage = ?3",
3050 params![call_id, subject, stage],
3051 |row| Ok((row.get(0)?, row.get(1)?)),
3052 )
3053 .map_err(|error| format!("could not check existing snapshot: {error}"))?;
3054 if existing != (sha, source_completeness.map(str::to_owned)) {
3055 return Err(format!(
3056 "snapshot {subject}/{stage} already exists with different content or completeness"
3057 ));
3058 }
3059 }
3060 Ok(())
3061}
3062
3063fn insert_payload(
3064 transaction: &Connection,
3065 reader: &mut (impl Read + Seek),
3066) -> Result<Vec<u8>, String> {
3067 reader
3068 .seek(SeekFrom::Start(0))
3069 .map_err(|error| format!("could not rewind payload: {error}"))?;
3070 let mut hasher = Sha256::new();
3071 let mut raw_length = 0_u64;
3072 let mut buffer = vec![0_u8; 64 * 1024];
3073 loop {
3074 let count = reader
3075 .read(&mut buffer)
3076 .map_err(|error| format!("could not read payload: {error}"))?;
3077 if count == 0 {
3078 break;
3079 }
3080 hasher.update(&buffer[..count]);
3081 raw_length = raw_length.saturating_add(count as u64);
3082 }
3083 let sha = hasher.finalize().to_vec();
3084 let exists = transaction
3085 .query_row("SELECT 1 FROM payloads WHERE sha256 = ?1", [&sha], |_| {
3086 Ok(())
3087 })
3088 .optional()
3089 .map_err(|error| format!("could not check payload digest: {error}"))?
3090 .is_some();
3091 if exists {
3092 return Ok(sha);
3093 }
3094
3095 reader
3096 .seek(SeekFrom::Start(0))
3097 .map_err(|error| format!("could not rewind payload for compression: {error}"))?;
3098 let mut compressed = NamedTempFile::new()
3099 .map_err(|error| format!("could not create compressed payload spool: {error}"))?;
3100 zstd::stream::copy_encode(&mut *reader, &mut compressed, ZSTD_LEVEL)
3101 .map_err(|error| format!("could not compress payload: {error}"))?;
3102 compressed
3103 .flush()
3104 .map_err(|error| format!("could not flush compressed payload: {error}"))?;
3105 let compressed_length = compressed
3106 .as_file()
3107 .metadata()
3108 .map_err(|error| format!("could not stat compressed payload: {error}"))?
3109 .len();
3110 let use_compressed =
3111 compressed_length.saturating_mul(100) <= raw_length.saturating_mul(COMPRESSION_PERCENT);
3112 let stored_length = if use_compressed {
3113 compressed_length
3114 } else {
3115 raw_length
3116 };
3117 let stored_length_i64 = i64::try_from(stored_length)
3118 .map_err(|_| format!("payload is too large for SQLite: {stored_length} bytes"))?;
3119 let raw_length_i64 = i64::try_from(raw_length)
3120 .map_err(|_| format!("payload is too large: {raw_length} bytes"))?;
3121 transaction
3122 .execute(
3123 "INSERT INTO payloads (sha256, compression, uncompressed_byte_length, body)
3124 VALUES (?1, ?2, ?3, zeroblob(?4))",
3125 params![
3126 sha,
3127 if use_compressed { "zstd" } else { "none" },
3128 raw_length_i64,
3129 stored_length_i64
3130 ],
3131 )
3132 .map_err(|error| format!("could not insert payload: {error}"))?;
3133 let rowid: i64 = transaction
3134 .query_row(
3135 "SELECT rowid FROM payloads WHERE sha256 = ?1",
3136 [&sha],
3137 |row| row.get(0),
3138 )
3139 .map_err(|error| format!("could not locate payload body: {error}"))?;
3140 let mut blob = transaction
3141 .blob_open(MAIN_DB, "payloads", "body", rowid, false)
3142 .map_err(|error| format!("could not open payload body: {error}"))?;
3143 if use_compressed {
3144 compressed
3145 .seek(SeekFrom::Start(0))
3146 .map_err(|error| format!("could not rewind compressed payload: {error}"))?;
3147 copy_exact(&mut compressed, &mut blob, stored_length)?;
3148 } else {
3149 reader
3150 .seek(SeekFrom::Start(0))
3151 .map_err(|error| format!("could not rewind raw payload: {error}"))?;
3152 copy_exact(reader, &mut blob, stored_length)?;
3153 }
3154 Ok(sha)
3155}
3156
3157fn copy_exact(reader: &mut impl Read, blob: &mut Blob<'_>, expected: u64) -> Result<(), String> {
3158 let copied =
3159 io::copy(reader, blob).map_err(|error| format!("could not write payload body: {error}"))?;
3160 if copied != expected {
3161 return Err(format!(
3162 "payload body length changed while storing it: expected {expected}, copied {copied}"
3163 ));
3164 }
3165 Ok(())
3166}
3167
3168fn canonical_json(value: &Value) -> Result<Vec<u8>, String> {
3169 serde_jcs::to_vec(value)
3170 .map_err(|error| format!("could not canonicalize JSON payload: {error}"))
3171}
3172
3173fn stream_media_type(reader: &mut (impl Read + Seek)) -> Result<&'static str, String> {
3174 reader
3175 .seek(SeekFrom::Start(0))
3176 .map_err(|error| format!("could not rewind stream for UTF-8 check: {error}"))?;
3177 let mut valid = true;
3178 let mut carry = Vec::new();
3179 let mut buffer = vec![0_u8; 64 * 1024];
3180 while valid {
3181 let count = reader
3182 .read(&mut buffer)
3183 .map_err(|error| format!("could not inspect stream encoding: {error}"))?;
3184 if count == 0 {
3185 break;
3186 }
3187 let mut chunk = std::mem::take(&mut carry);
3188 chunk.extend_from_slice(&buffer[..count]);
3189 if let Err(error) = std::str::from_utf8(&chunk) {
3190 if error.error_len().is_some() {
3191 valid = false;
3192 } else {
3193 carry.extend_from_slice(&chunk[error.valid_up_to()..]);
3194 if carry.len() > 3 {
3195 valid = false;
3196 }
3197 }
3198 }
3199 }
3200 if !carry.is_empty() {
3201 valid = false;
3202 }
3203 reader
3204 .seek(SeekFrom::Start(0))
3205 .map_err(|error| format!("could not rewind inspected stream: {error}"))?;
3206 Ok(if valid {
3207 "text/plain; charset=utf-8"
3208 } else {
3209 "application/octet-stream"
3210 })
3211}
3212
3213fn sync_directory(path: &Path) -> Result<(), String> {
3214 let directory = fs::File::open(path)
3215 .map_err(|error| format!("could not open {} for sync: {error}", path.display()))?;
3216 directory
3217 .sync_all()
3218 .map_err(|error| format!("could not sync {}: {error}", path.display()))
3219}
3220
3221fn archive_file_bytes(path: &Path) -> Result<u64, String> {
3222 let main = fs::metadata(path)
3223 .map_err(|error| format!("could not stat archive {}: {error}", path.display()))?
3224 .len();
3225 let companions = ["-wal", "-shm"].into_iter().map(|suffix| {
3226 let mut value = path.as_os_str().to_owned();
3227 value.push(suffix);
3228 PathBuf::from(value)
3229 });
3230 Ok(companions.fold(main, |total, companion| {
3231 total.saturating_add(fs::metadata(companion).map_or(0, |metadata| metadata.len()))
3232 }))
3233}
3234
3235fn prepare_path(path: &Path, repair_existing_directory: bool) -> Result<(), String> {
3236 let parent = path
3237 .parent()
3238 .ok_or_else(|| format!("archive path {} has no parent", path.display()))?;
3239 let parent_existed = parent.exists();
3240 create_archive_directory(parent)?;
3241 if !parent_existed || repair_existing_directory {
3242 set_file_mode(parent, 0o700)?;
3243 } else {
3244 require_private_directory(parent)?;
3245 }
3246 let mut options = OpenOptions::new();
3247 options.create(true).read(true).write(true);
3248 #[cfg(unix)]
3249 {
3250 use std::os::unix::fs::OpenOptionsExt as _;
3251 options.mode(0o600);
3252 }
3253 options
3254 .open(path)
3255 .map_err(|error| format!("could not create archive {}: {error}", path.display()))?;
3256 set_file_mode(path, 0o600)
3257}
3258
3259fn create_archive_directory(path: &Path) -> Result<(), String> {
3260 let mut builder = fs::DirBuilder::new();
3261 builder.recursive(true);
3262 #[cfg(unix)]
3263 {
3264 use std::os::unix::fs::DirBuilderExt as _;
3265 builder.mode(0o700);
3266 }
3267 builder.create(path).map_err(|error| {
3268 format!(
3269 "could not create archive directory {}: {error}",
3270 path.display()
3271 )
3272 })
3273}
3274
3275#[cfg(unix)]
3276fn require_private_archive_path(path: &Path) -> Result<(), String> {
3277 use std::os::unix::fs::PermissionsExt as _;
3278 require_private_directory(
3279 path.parent()
3280 .ok_or_else(|| format!("archive path {} has no parent", path.display()))?,
3281 )?;
3282 let metadata = fs::symlink_metadata(path)
3283 .map_err(|error| format!("could not inspect archive {}: {error}", path.display()))?;
3284 if metadata.file_type().is_symlink() || !metadata.is_file() {
3285 return Err(format!(
3286 "archive {} must be a regular non-symlink file",
3287 path.display()
3288 ));
3289 }
3290 let mode = metadata.permissions().mode() & 0o777;
3291 if mode & 0o077 != 0 {
3292 return Err(format!(
3293 "archive {} has mode {mode:o}, expected 600",
3294 path.display()
3295 ));
3296 }
3297 Ok(())
3298}
3299
3300#[cfg(not(unix))]
3301fn require_private_archive_path(_path: &Path) -> Result<(), String> {
3302 Ok(())
3303}
3304
3305#[cfg(unix)]
3306fn require_private_directory(path: &Path) -> Result<(), String> {
3307 use std::os::unix::fs::PermissionsExt as _;
3308 let mode = fs::metadata(path)
3309 .map_err(|error| format!("could not inspect {}: {error}", path.display()))?
3310 .permissions()
3311 .mode()
3312 & 0o777;
3313 if mode & 0o077 != 0 {
3314 return Err(format!(
3315 "archive directory {} has mode {mode:o}, expected 700; use a private directory",
3316 path.display()
3317 ));
3318 }
3319 Ok(())
3320}
3321
3322#[cfg(not(unix))]
3323fn require_private_directory(_path: &Path) -> Result<(), String> {
3324 Ok(())
3325}
3326
3327#[cfg(unix)]
3328fn set_file_mode(path: &Path, mode: u32) -> Result<(), String> {
3329 use std::os::unix::fs::PermissionsExt as _;
3330 fs::set_permissions(path, fs::Permissions::from_mode(mode))
3331 .map_err(|error| format!("could not set permissions on {}: {error}", path.display()))
3332}
3333
3334#[cfg(not(unix))]
3335fn set_file_mode(_path: &Path, _mode: u32) -> Result<(), String> {
3336 Ok(())
3337}
3338
3339#[cfg(unix)]
3340fn verify_permissions(path: &Path, report: &mut VerifyReport) -> Result<(), String> {
3341 use std::os::unix::fs::PermissionsExt as _;
3342 let directory = path
3343 .parent()
3344 .ok_or_else(|| format!("archive path {} has no parent", path.display()))?;
3345 let directory_mode = fs::metadata(directory)
3346 .map_err(|error| format!("could not stat archive directory: {error}"))?
3347 .permissions()
3348 .mode()
3349 & 0o777;
3350 if directory_mode != 0o700 {
3351 report.errors.push(format!(
3352 "archive directory mode is {directory_mode:o}, expected 700"
3353 ));
3354 }
3355 let file_mode = fs::metadata(path)
3356 .map_err(|error| format!("could not stat archive file: {error}"))?
3357 .permissions()
3358 .mode()
3359 & 0o777;
3360 if file_mode != 0o600 {
3361 report
3362 .errors
3363 .push(format!("archive file mode is {file_mode:o}, expected 600"));
3364 }
3365 Ok(())
3366}
3367
3368#[cfg(not(unix))]
3369fn verify_permissions(_path: &Path, _report: &mut VerifyReport) -> Result<(), String> {
3370 Ok(())
3371}
3372
3373#[cfg(test)]
3374mod tests {
3375 use super::*;
3376 #[cfg(unix)]
3377 use std::os::unix::fs::PermissionsExt as _;
3378 use tempfile::TempDir;
3379
3380 fn archive() -> (TempDir, Archive) {
3381 let directory = TempDir::new().expect("temp directory");
3382 let archive = Archive::open_path(directory.path().join("data/tool-calls.sqlite3"))
3383 .expect("open archive");
3384 (directory, archive)
3385 }
3386
3387 fn session() -> SessionIdentity {
3388 SessionIdentity {
3389 agent: "pi".to_owned(),
3390 account: "test".to_owned(),
3391 source_session_id: "session-1".to_owned(),
3392 started_at_ms: Some(10),
3393 }
3394 }
3395
3396 fn call() -> CallIdentity {
3397 call_with_id("call-1")
3398 }
3399
3400 fn call_with_id(source_call_id: &str) -> CallIdentity {
3401 CallIdentity {
3402 source_call_id: source_call_id.to_owned(),
3403 tool_name: "read".to_owned(),
3404 provider: Some("openai".to_owned()),
3405 model: Some("gpt".to_owned()),
3406 working_directory: Some("/tmp".to_owned()),
3407 started_at_ms: 20,
3408 requires_streams: false,
3409 }
3410 }
3411
3412 fn framed(value: &Value) -> Vec<u8> {
3413 let body = serde_json::to_vec(value).expect("json");
3414 let mut frame = Vec::new();
3415 frame.extend_from_slice(&(body.len() as u64).to_be_bytes());
3416 frame.extend_from_slice(&body);
3417 frame
3418 }
3419
3420 fn begin_operation(request_id: u64, source_call_id: &str, value: Value) -> ArchiveOperation {
3421 ArchiveOperation::BeginCall {
3422 request_id,
3423 schema_version: INGEST_SCHEMA_VERSION,
3424 session: session(),
3425 call: call_with_id(source_call_id),
3426 input_before: value.clone(),
3427 input_after: value,
3428 captured_at_ms: 20,
3429 }
3430 }
3431
3432 #[test]
3433 fn prepared_batch_commits_multiple_calls_once() {
3434 let (_directory, mut archive) = archive();
3435 archive.configure_for_broker().expect("broker policy");
3436 let busy_timeout: i64 = archive
3437 .connection
3438 .pragma_query_value(None, "busy_timeout", |row| row.get(0))
3439 .expect("busy timeout");
3440 let wal_autocheckpoint: i64 = archive
3441 .connection
3442 .pragma_query_value(None, "wal_autocheckpoint", |row| row.get(0))
3443 .expect("WAL checkpoint limit");
3444 assert_eq!(busy_timeout, 250);
3445 assert_eq!(wal_autocheckpoint, 1000);
3446 let mut first = archive
3447 .prepare_operation(begin_operation(
3448 1,
3449 "call-1",
3450 serde_json::json!({"value": 1}),
3451 ))
3452 .expect("prepare first");
3453 let mut second = archive
3454 .prepare_operation(begin_operation(
3455 2,
3456 "call-2",
3457 serde_json::json!({"value": 2}),
3458 ))
3459 .expect("prepare second");
3460 let results = archive
3461 .apply_prepared_batch([&mut first, &mut second])
3462 .expect("batch");
3463 assert_eq!(results.len(), 2);
3464 assert!(results.iter().all(Result::is_ok));
3465 assert_eq!(archive.stats().expect("stats").calls, 2);
3466 assert!(archive.verify().expect("verify").errors.is_empty());
3467 }
3468
3469 #[test]
3470 fn failed_savepoint_does_not_poison_shared_payload_in_next_request() {
3471 let (_directory, mut archive) = archive();
3472 archive
3473 .begin_call(
3474 &session(),
3475 &call_with_id("call-1"),
3476 &serde_json::json!({}),
3477 &serde_json::json!({}),
3478 20,
3479 )
3480 .expect("existing call");
3481 let shared = serde_json::json!({"shared": "payload"});
3482 let conflicting = ArchiveOperation::BeginCall {
3483 request_id: 1,
3484 schema_version: INGEST_SCHEMA_VERSION,
3485 session: session(),
3486 call: call_with_id("call-1"),
3487 input_before: serde_json::json!({}),
3488 input_after: shared.clone(),
3489 captured_at_ms: 20,
3490 };
3491 let mut first = archive
3492 .prepare_operation(conflicting)
3493 .expect("prepare conflict");
3494 let mut second = archive
3495 .prepare_operation(begin_operation(2, "call-2", shared))
3496 .expect("prepare second");
3497 let results = archive
3498 .apply_prepared_batch([&mut first, &mut second])
3499 .expect("batch commit");
3500 assert!(results[0].is_err());
3501 assert!(results[1].is_ok());
3502 assert_eq!(archive.stats().expect("stats").calls, 2);
3503 assert!(archive.verify().expect("verify").errors.is_empty());
3504 }
3505
3506 #[test]
3507 fn prepared_replay_is_idempotent_and_conflicts_on_changed_content() {
3508 let (_directory, mut archive) = archive();
3509 let mut first = archive
3510 .prepare_operation(begin_operation(
3511 1,
3512 "call-1",
3513 serde_json::json!({"value": 1}),
3514 ))
3515 .expect("prepare first");
3516 archive
3517 .apply_prepared_batch([&mut first])
3518 .expect("first commit");
3519 let mut replay = archive
3520 .prepare_operation(begin_operation(
3521 2,
3522 "call-1",
3523 serde_json::json!({"value": 1}),
3524 ))
3525 .expect("prepare replay");
3526 assert!(
3527 archive
3528 .apply_prepared_batch([&mut replay])
3529 .expect("replay commit")[0]
3530 .is_ok()
3531 );
3532 let mut conflict = archive
3533 .prepare_operation(begin_operation(
3534 3,
3535 "call-1",
3536 serde_json::json!({"value": 2}),
3537 ))
3538 .expect("prepare conflict");
3539 assert!(
3540 archive
3541 .apply_prepared_batch([&mut conflict])
3542 .expect("conflict batch")[0]
3543 .is_err()
3544 );
3545 assert_eq!(archive.stats().expect("stats").calls, 1);
3546 }
3547
3548 #[test]
3549 fn prepared_batch_covers_the_complete_shell_call_lifecycle() {
3550 let (directory, mut archive) = archive();
3551 let mut shell_call = call();
3552 shell_call.requires_streams = true;
3553 let source = directory.path().join("data/source-output");
3554 let stdout_before = directory.path().join("data/stdout-before");
3555 let stderr_before = directory.path().join("data/stderr-before");
3556 let stdout_after = directory.path().join("data/stdout-after");
3557 let stderr_after = directory.path().join("data/stderr-after");
3558 for (path, body) in [
3559 (&source, b"source output".as_slice()),
3560 (&stdout_before, b"stdout raw".as_slice()),
3561 (&stderr_before, b"stderr raw".as_slice()),
3562 (&stdout_after, b"stdout short".as_slice()),
3563 (&stderr_after, b"stderr short".as_slice()),
3564 ] {
3565 fs::write(path, body).expect("stream file");
3566 #[cfg(unix)]
3567 fs::set_permissions(path, fs::Permissions::from_mode(0o600)).expect("private stream");
3568 }
3569 let operations = [
3570 ArchiveOperation::BeginCall {
3571 request_id: 1,
3572 schema_version: INGEST_SCHEMA_VERSION,
3573 session: session(),
3574 call: shell_call,
3575 input_before: serde_json::json!({"cmd": "test"}),
3576 input_after: serde_json::json!({"cmd": "test"}),
3577 captured_at_ms: 20,
3578 },
3579 ArchiveOperation::CaptureStreams {
3580 request_id: 2,
3581 schema_version: INGEST_SCHEMA_VERSION,
3582 session: session(),
3583 source_call_id: "call-1".to_owned(),
3584 captured_at_ms: 30,
3585 stdout_before_path: stdout_before,
3586 stderr_before_path: stderr_before,
3587 stdout_after_path: stdout_after,
3588 stderr_after_path: stderr_after,
3589 },
3590 ArchiveOperation::ResultBefore {
3591 request_id: 3,
3592 schema_version: INGEST_SCHEMA_VERSION,
3593 session: session(),
3594 source_call_id: "call-1".to_owned(),
3595 result: serde_json::json!({"content": "raw"}),
3596 full_output_path: Some(source),
3597 captured_at_ms: 40,
3598 },
3599 ArchiveOperation::ResultText {
3600 request_id: 4,
3601 schema_version: INGEST_SCHEMA_VERSION,
3602 session: session(),
3603 source_call_id: "call-1".to_owned(),
3604 text: "exact visible text".to_owned(),
3605 source_completeness: SourceCompleteness::Complete,
3606 captured_at_ms: 41,
3607 },
3608 ArchiveOperation::StageResult {
3609 request_id: 5,
3610 schema_version: INGEST_SCHEMA_VERSION,
3611 session: session(),
3612 source_call_id: "call-1".to_owned(),
3613 result: serde_json::json!({"content": "staged"}),
3614 is_error: false,
3615 captured_at_ms: 50,
3616 },
3617 ArchiveOperation::UpdateFinalResult {
3618 request_id: 6,
3619 schema_version: INGEST_SCHEMA_VERSION,
3620 session: session(),
3621 source_call_id: "call-1".to_owned(),
3622 result: serde_json::json!({"content": "final"}),
3623 is_error: false,
3624 finished_at_ms: 60,
3625 },
3626 ];
3627 let mut prepared = operations
3628 .into_iter()
3629 .map(|operation| archive.prepare_operation(operation).expect("prepare"))
3630 .collect::<Vec<_>>();
3631 let results = archive
3632 .apply_prepared_batch(prepared.iter_mut())
3633 .expect("complete batch");
3634 assert!(results.iter().all(Result::is_ok));
3635 assert_eq!(archive.stats().expect("stats").incomplete_calls, 0);
3636 assert!(archive.verify().expect("verify").errors.is_empty());
3637 }
3638
3639 #[test]
3640 fn prepared_prune_uses_the_broker_transaction_path() {
3641 let (_directory, mut archive) = archive();
3642 archive
3643 .begin_call(
3644 &session(),
3645 &call(),
3646 &serde_json::json!({}),
3647 &serde_json::json!({}),
3648 20,
3649 )
3650 .expect("begin");
3651 archive
3652 .finish_call(
3653 &session(),
3654 "call-1",
3655 &serde_json::json!({"error": "blocked"}),
3656 true,
3657 false,
3658 30,
3659 )
3660 .expect("finish");
3661 let mut prune = archive
3662 .prepare_operation(ArchiveOperation::PruneBefore {
3663 request_id: 7,
3664 schema_version: INGEST_SCHEMA_VERSION,
3665 timestamp_ms: 31,
3666 })
3667 .expect("prepare prune");
3668 let result = archive
3669 .apply_prepared_batch([&mut prune])
3670 .expect("prune batch");
3671 assert_eq!(
3672 result[0].as_ref().expect("prune result").as_deref(),
3673 Some("1")
3674 );
3675 assert_eq!(archive.stats().expect("stats").calls, 0);
3676 }
3677
3678 #[test]
3679 fn result_before_replay_rejects_changed_full_output_path() {
3680 let (directory, mut archive) = archive();
3681 archive
3682 .begin_call(
3683 &session(),
3684 &call(),
3685 &serde_json::json!({}),
3686 &serde_json::json!({}),
3687 20,
3688 )
3689 .expect("begin");
3690 let source = directory.path().join("data/full-output");
3691 fs::write(&source, b"first").expect("source output");
3692 #[cfg(unix)]
3693 fs::set_permissions(&source, fs::Permissions::from_mode(0o600)).expect("private source");
3694 let operation = |request_id| ArchiveOperation::ResultBefore {
3695 request_id,
3696 schema_version: INGEST_SCHEMA_VERSION,
3697 session: session(),
3698 source_call_id: "call-1".to_owned(),
3699 result: serde_json::json!({"content": "visible"}),
3700 full_output_path: Some(source.clone()),
3701 captured_at_ms: 30,
3702 };
3703 let mut first = archive
3704 .prepare_operation(operation(1))
3705 .expect("prepare first");
3706 assert!(
3707 archive
3708 .apply_prepared_batch([&mut first])
3709 .expect("first batch")[0]
3710 .is_ok()
3711 );
3712 fs::write(&source, b"second").expect("changed source output");
3713 let mut replay = archive
3714 .prepare_operation(operation(2))
3715 .expect("prepare replay");
3716 assert!(
3717 archive
3718 .apply_prepared_batch([&mut replay])
3719 .expect("replay batch")[0]
3720 .is_err()
3721 );
3722 }
3723
3724 #[test]
3725 fn stores_and_verifies_a_complete_call() {
3726 let (_directory, mut archive) = archive();
3727 archive
3728 .begin_call(
3729 &session(),
3730 &call(),
3731 &serde_json::json!({"path": "file"}),
3732 &serde_json::json!({"path": "file"}),
3733 20,
3734 )
3735 .expect("begin");
3736 archive
3737 .result_before(
3738 &session(),
3739 "call-1",
3740 &serde_json::json!({"content": "raw"}),
3741 None,
3742 30,
3743 )
3744 .expect("before");
3745 archive
3746 .finish_call(
3747 &session(),
3748 "call-1",
3749 &serde_json::json!({"content": "pruned"}),
3750 false,
3751 true,
3752 40,
3753 )
3754 .expect("finish");
3755
3756 let stats = archive.stats().expect("stats");
3757 assert_eq!(stats.sessions, 1);
3758 assert_eq!(stats.calls, 1);
3759 assert_eq!(stats.incomplete_calls, 0);
3760 assert!(stats.logical_payload_bytes > 0);
3761 assert!(archive.verify().expect("verify").errors.is_empty());
3762 }
3763
3764 #[test]
3765 fn finalizes_the_staged_result_after_other_hooks() {
3766 let (_directory, mut archive) = archive();
3767 archive
3768 .begin_call(
3769 &session(),
3770 &call(),
3771 &serde_json::json!({}),
3772 &serde_json::json!({}),
3773 20,
3774 )
3775 .expect("begin");
3776 archive
3777 .result_before(
3778 &session(),
3779 "call-1",
3780 &serde_json::json!({"content": "before"}),
3781 None,
3782 30,
3783 )
3784 .expect("before");
3785 archive
3786 .stage_result(
3787 &session(),
3788 "call-1",
3789 &serde_json::json!({"content": "provisional"}),
3790 false,
3791 40,
3792 )
3793 .expect("stage");
3794 let staged_report = archive.verify().expect("verify staged call");
3795 assert_eq!(staged_report.incomplete_calls, 1);
3796 assert!(staged_report.errors.is_empty());
3797 archive
3798 .update_final_result(
3799 &session(),
3800 "call-1",
3801 &serde_json::json!({"content": "final"}),
3802 true,
3803 50,
3804 )
3805 .expect("update final");
3806 archive
3807 .update_final_result(
3808 &session(),
3809 "call-1",
3810 &serde_json::json!({"content": "final"}),
3811 true,
3812 50,
3813 )
3814 .expect("repeat committed final update");
3815
3816 let call_id: i64 = archive
3817 .connection
3818 .query_row("SELECT id FROM tool_calls", [], |row| row.get(0))
3819 .expect("call id");
3820 let mut restored = archive
3821 .verified_snapshot(call_id, "result", "after")
3822 .expect("final result");
3823 let mut body = Vec::new();
3824 restored.read_to_end(&mut body).expect("read final result");
3825 assert_eq!(body, br#"{"content":"final"}"#);
3826 let metadata: (i64, i64) = archive
3827 .connection
3828 .query_row(
3829 "SELECT is_error, finished_at_ms FROM tool_calls WHERE id = ?1",
3830 [call_id],
3831 |row| Ok((row.get(0)?, row.get(1)?)),
3832 )
3833 .expect("final metadata");
3834 assert_eq!(metadata, (1, 50));
3835 assert!(archive.verify().expect("verify").errors.is_empty());
3836 }
3837
3838 #[test]
3839 fn failed_finalization_keeps_the_staged_call_incomplete() {
3840 let (_directory, mut archive) = archive();
3841 archive
3842 .begin_call(
3843 &session(),
3844 &call(),
3845 &serde_json::json!({}),
3846 &serde_json::json!({}),
3847 20,
3848 )
3849 .expect("begin");
3850 archive
3851 .result_before(
3852 &session(),
3853 "call-1",
3854 &serde_json::json!({"content": "before"}),
3855 None,
3856 30,
3857 )
3858 .expect("before");
3859 archive
3860 .stage_result(
3861 &session(),
3862 "call-1",
3863 &serde_json::json!({"content": "provisional"}),
3864 false,
3865 40,
3866 )
3867 .expect("stage");
3868 archive
3869 .connection
3870 .execute_batch(
3871 "CREATE TEMP TRIGGER reject_finish
3872 BEFORE UPDATE OF status ON tool_calls
3873 WHEN NEW.status = 'finished'
3874 BEGIN SELECT RAISE(ABORT, 'simulated final write failure'); END;",
3875 )
3876 .expect("failure trigger");
3877
3878 let error = archive
3879 .update_final_result(
3880 &session(),
3881 "call-1",
3882 &serde_json::json!({"content": "final"}),
3883 false,
3884 50,
3885 )
3886 .expect_err("finalization failure");
3887 assert!(error.contains("simulated final write failure"));
3888 assert_eq!(archive.stats().expect("stats").incomplete_calls, 1);
3889 let stored: (String, Option<i64>) = archive
3890 .connection
3891 .query_row("SELECT status, finished_at_ms FROM tool_calls", [], |row| {
3892 Ok((row.get(0)?, row.get(1)?))
3893 })
3894 .expect("stored state");
3895 assert_eq!(stored, ("started".to_owned(), None));
3896 }
3897
3898 #[test]
3899 fn deduplicates_unchanged_snapshots() {
3900 let (_directory, mut archive) = archive();
3901 let value = serde_json::json!({"same": true});
3902 archive
3903 .begin_call(&session(), &call(), &value, &value, 20)
3904 .expect("begin");
3905 let payloads: i64 = archive
3906 .connection
3907 .query_row("SELECT count(*) FROM payloads", [], |row| row.get(0))
3908 .expect("count");
3909 assert_eq!(payloads, 1);
3910 }
3911
3912 #[test]
3913 fn duplicate_operations_are_idempotent_but_conflicts_fail() {
3914 let (_directory, mut archive) = archive();
3915 let value = serde_json::json!({"same": true});
3916 archive
3917 .begin_call(&session(), &call(), &value, &value, 20)
3918 .expect("begin");
3919 archive
3920 .begin_call(&session(), &call(), &value, &value, 20)
3921 .expect("retry");
3922 let error = archive
3923 .begin_call(
3924 &session(),
3925 &call(),
3926 &serde_json::json!({"different": true}),
3927 &value,
3928 20,
3929 )
3930 .expect_err("conflict");
3931 assert!(error.contains("different content"));
3932 }
3933
3934 #[test]
3935 fn marks_interrupted_calls_without_failing_verification() {
3936 let (_directory, mut archive) = archive();
3937 archive
3938 .begin_call(
3939 &session(),
3940 &call(),
3941 &serde_json::json!({}),
3942 &serde_json::json!({}),
3943 20,
3944 )
3945 .expect("begin");
3946 let report = archive.verify().expect("verify");
3947 assert_eq!(report.incomplete_calls, 1);
3948 assert!(report.errors.is_empty());
3949 }
3950
3951 #[test]
3952 fn stores_binary_and_text_streams() {
3953 let (_directory, mut archive) = archive();
3954 let archive_ref = archive
3955 .begin_call(
3956 &session(),
3957 &call(),
3958 &serde_json::json!({}),
3959 &serde_json::json!({}),
3960 20,
3961 )
3962 .expect("begin");
3963 let key = ArchiveKey {
3964 session: session(),
3965 source_call_id: "call-1".to_owned(),
3966 };
3967 let mut stdout = Cursor::new(b"hello\n".to_vec());
3968 let mut stderr = Cursor::new(vec![0xff, 0x00]);
3969 archive
3970 .capture_streams(&key, 30, &mut stdout, &mut stderr, b"hello\n", b"short")
3971 .expect("streams");
3972 let media: Vec<String> = archive
3973 .connection
3974 .prepare("SELECT media_type FROM snapshots WHERE subject IN ('stdout', 'stderr') ORDER BY subject, stage")
3975 .expect("prepare")
3976 .query_map([], |row| row.get(0))
3977 .expect("query")
3978 .collect::<Result<_, _>>()
3979 .expect("rows");
3980 assert!(media.contains(&"application/octet-stream".to_owned()));
3981 assert!(media.contains(&"text/plain; charset=utf-8".to_owned()));
3982
3983 archive
3984 .result_text(
3985 &session(),
3986 "call-1",
3987 "capped wrapped summary\n",
3988 SourceCompleteness::Unknown,
3989 40,
3990 )
3991 .expect("wrapped summary fallback");
3992 let sources = archive.searchable_sources(&archive_ref).expect("sources");
3993 assert_eq!(
3994 sources.iter().map(|source| source.name).collect::<Vec<_>>(),
3995 vec![SourceName::Stdout, SourceName::Stderr]
3996 );
3997 }
3998
3999 #[test]
4000 fn stores_exact_source_full_output_with_the_result() {
4001 let (directory, mut archive) = archive();
4002 let archive_ref = archive
4003 .begin_call(
4004 &session(),
4005 &call(),
4006 &serde_json::json!({}),
4007 &serde_json::json!({}),
4008 20,
4009 )
4010 .expect("begin");
4011 let full_output_path = directory.path().join("pi-full-output.log");
4012 let expected = b"complete output\n\xff";
4013 fs::write(&full_output_path, expected).expect("full output");
4014 archive
4015 .result_before(
4016 &session(),
4017 "call-1",
4018 &serde_json::json!({"content": "truncated"}),
4019 Some(&full_output_path),
4020 30,
4021 )
4022 .expect("result before");
4023
4024 let call_id: i64 = archive
4025 .connection
4026 .query_row("SELECT id FROM tool_calls", [], |row| row.get(0))
4027 .expect("call id");
4028 let mut restored = archive
4029 .verified_snapshot(call_id, "source_output", "before")
4030 .expect("source output");
4031 let mut actual = Vec::new();
4032 restored
4033 .read_to_end(&mut actual)
4034 .expect("read source output");
4035 assert_eq!(actual, expected);
4036
4037 archive
4038 .result_text(
4039 &session(),
4040 "call-1",
4041 "exact host text\n",
4042 SourceCompleteness::Incomplete,
4043 40,
4044 )
4045 .expect("result text");
4046 let mut sources = archive.searchable_sources(&archive_ref).expect("sources");
4047 assert_eq!(sources.len(), 1);
4048 assert_eq!(sources[0].name, SourceName::ResultText);
4049 let mut host_text = String::new();
4050 sources[0]
4051 .body
4052 .read_to_string(&mut host_text)
4053 .expect("result text body");
4054 assert_eq!(host_text, "exact host text\n");
4055 }
4056
4057 #[test]
4058 fn prunes_only_finished_calls() {
4059 let (_directory, mut archive) = archive();
4060 archive
4061 .begin_call(
4062 &session(),
4063 &call(),
4064 &serde_json::json!({}),
4065 &serde_json::json!({}),
4066 20,
4067 )
4068 .expect("begin");
4069 archive
4070 .finish_call(
4071 &session(),
4072 "call-1",
4073 &serde_json::json!({}),
4074 false,
4075 false,
4076 40,
4077 )
4078 .expect("finish");
4079 assert_eq!(archive.prune_before(50).expect("prune"), 1);
4080 let stats = archive.stats().expect("stats");
4081 assert_eq!(stats.calls, 0);
4082 assert_eq!(stats.logical_payload_bytes, 0);
4083 }
4084
4085 #[test]
4086 fn migrates_the_original_version_one_schema_in_place_with_a_private_backup() {
4087 let directory = TempDir::new().expect("temp directory");
4088 let path = directory.path().join("yarp/tool-calls.sqlite3");
4089 fs::create_dir(path.parent().expect("parent")).expect("archive directory");
4090 set_file_mode(path.parent().expect("parent"), 0o700).expect("private directory");
4091 let connection = Connection::open(&path).expect("old database");
4092 connection
4093 .execute_batch(include_str!("../tests/fixtures/archive-v1.sql"))
4094 .expect("old schema");
4095 drop(connection);
4096 set_file_mode(&path, 0o600).expect("private database");
4097
4098 let archive = Archive::open_path(path.clone()).expect("migrate archive");
4099 let archive_ref: String = archive
4100 .connection
4101 .query_row("SELECT archive_ref FROM tool_calls", [], |row| row.get(0))
4102 .expect("archive ref");
4103 validate_archive_ref(&archive_ref).expect("valid archive ref");
4104 assert!(
4105 table_has_column(&archive.connection, "snapshots", "source_completeness")
4106 .expect("schema column")
4107 );
4108 drop(archive);
4109
4110 let backup = path.with_file_name("tool-calls.sqlite3.pre-indexed-output-v1.backup");
4111 assert!(backup.is_file());
4112 #[cfg(unix)]
4113 assert_eq!(
4114 fs::metadata(&backup)
4115 .expect("backup metadata")
4116 .permissions()
4117 .mode()
4118 & 0o777,
4119 0o600
4120 );
4121 let reopened = Archive::open_path(path).expect("reopen migrated archive");
4122 let reopened_ref: String = reopened
4123 .connection
4124 .query_row("SELECT archive_ref FROM tool_calls", [], |row| row.get(0))
4125 .expect("stable archive ref");
4126 assert_eq!(archive_ref, reopened_ref);
4127 }
4128
4129 #[test]
4130 fn result_text_is_on_demand_and_canonical_source_selection_is_stable() {
4131 let (_directory, mut archive) = archive();
4132 let archive_ref = archive
4133 .begin_call(
4134 &session(),
4135 &call(),
4136 &serde_json::json!({}),
4137 &serde_json::json!({}),
4138 20,
4139 )
4140 .expect("begin");
4141 let repeated_ref = archive
4142 .result_text(
4143 &session(),
4144 "call-1",
4145 "first\nerror\nlast\n",
4146 SourceCompleteness::Incomplete,
4147 30,
4148 )
4149 .expect("result text");
4150 assert_eq!(archive_ref, repeated_ref);
4151 let mut sources = archive.searchable_sources(&archive_ref).expect("sources");
4152 assert_eq!(sources.len(), 1);
4153 assert_eq!(sources[0].name, SourceName::ResultText);
4154 assert_eq!(sources[0].completeness, SourceCompleteness::Incomplete);
4155 let mut body = String::new();
4156 sources[0]
4157 .body
4158 .read_to_string(&mut body)
4159 .expect("source body");
4160 assert_eq!(body, "first\nerror\nlast\n");
4161 assert!(archive.verify().expect("verify").errors.is_empty());
4162 }
4163
4164 #[test]
4165 fn reopens_the_current_schema_without_migration() {
4166 let directory = TempDir::new().expect("temp directory");
4167 let path = directory.path().join("yarp/tool-calls.sqlite3");
4168 drop(Archive::open_path(path.clone()).expect("create"));
4169 let reopened = Archive::open_path(path).expect("reopen");
4170 assert_eq!(reopened.stats().expect("stats").calls, 0);
4171 }
4172
4173 #[test]
4174 fn rejects_newer_schema_versions() {
4175 let directory = TempDir::new().expect("temp directory");
4176 let path = directory.path().join("yarp/tool-calls.sqlite3");
4177 fs::create_dir(path.parent().expect("parent")).expect("archive directory");
4178 set_file_mode(path.parent().expect("parent"), 0o700).expect("private directory");
4179 let connection = Connection::open(&path).expect("sqlite");
4180 connection
4181 .pragma_update(None, "user_version", SCHEMA_VERSION + 1)
4182 .expect("version");
4183 drop(connection);
4184 let error = Archive::open_path(path).err().expect("newer version error");
4185 assert!(error.contains("newer than supported"));
4186 }
4187
4188 #[test]
4189 fn compresses_large_repetitive_payloads() {
4190 let (_directory, mut archive) = archive();
4191 let value = serde_json::json!({"output": "x".repeat(32 * 1024)});
4192 archive
4193 .begin_call(&session(), &call(), &value, &value, 20)
4194 .expect("begin");
4195 let compression: String = archive
4196 .connection
4197 .query_row("SELECT compression FROM payloads", [], |row| row.get(0))
4198 .expect("compression");
4199 assert_eq!(compression, "zstd");
4200 }
4201
4202 #[test]
4203 fn keeps_small_payloads_uncompressed() {
4204 let (_directory, mut archive) = archive();
4205 archive
4206 .begin_call(
4207 &session(),
4208 &call(),
4209 &serde_json::json!({"value": "short"}),
4210 &serde_json::json!({"value": "short"}),
4211 20,
4212 )
4213 .expect("begin");
4214 let compression: String = archive
4215 .connection
4216 .query_row("SELECT compression FROM payloads", [], |row| row.get(0))
4217 .expect("compression");
4218 assert_eq!(compression, "none");
4219 }
4220
4221 #[test]
4222 fn synthetic_archive_stays_below_storage_budget() {
4223 let (_directory, mut archive) = archive();
4224 let value = serde_json::json!({"output": "archive fixture ".repeat(64 * 1024)});
4225 archive
4226 .begin_call(&session(), &call(), &value, &value, 20)
4227 .expect("begin");
4228 archive
4229 .result_before(&session(), "call-1", &value, None, 30)
4230 .expect("result before");
4231 archive
4232 .finish_call(&session(), "call-1", &value, false, true, 40)
4233 .expect("finish");
4234 let referenced_bytes: i64 = archive
4235 .connection
4236 .query_row(
4237 "SELECT sum(p.uncompressed_byte_length)
4238 FROM snapshots s JOIN payloads p ON p.sha256 = s.payload_sha256",
4239 [],
4240 |row| row.get(0),
4241 )
4242 .expect("logical snapshot bytes");
4243 let stats = archive.stats().expect("stats");
4244 assert!(stats.database_bytes < u64::try_from(referenced_bytes * 3 / 10).expect("budget"));
4245 }
4246
4247 #[test]
4248 fn detects_corrupt_compressed_payloads() {
4249 let (_directory, mut archive) = archive();
4250 let value = serde_json::json!({"output": "x".repeat(32 * 1024)});
4251 archive
4252 .begin_call(&session(), &call(), &value, &value, 20)
4253 .expect("begin");
4254 archive
4255 .connection
4256 .execute("UPDATE payloads SET body = x'00'", [])
4257 .expect("corrupt");
4258 let error = archive.verify().expect_err("corruption");
4259 assert!(error.contains("payload"));
4260 }
4261
4262 #[test]
4263 fn restore_rejects_corrupt_payloads_before_writing_output() {
4264 let (_directory, mut archive) = archive();
4265 let mut shell_call = call();
4266 shell_call.requires_streams = true;
4267 archive
4268 .begin_call(
4269 &session(),
4270 &shell_call,
4271 &serde_json::json!({}),
4272 &serde_json::json!({}),
4273 20,
4274 )
4275 .expect("begin");
4276 archive
4277 .capture_streams(
4278 &ArchiveKey {
4279 session: session(),
4280 source_call_id: "call-1".to_owned(),
4281 },
4282 30,
4283 &mut Cursor::new(b"raw stdout"),
4284 &mut Cursor::new(b"raw stderr"),
4285 b"after stdout",
4286 b"after stderr",
4287 )
4288 .expect("streams");
4289 archive
4290 .connection
4291 .execute(
4292 "UPDATE payloads SET body = ?1
4293 WHERE sha256 = (
4294 SELECT payload_sha256 FROM snapshots
4295 WHERE subject = 'stdout' AND stage = 'before'
4296 )",
4297 [b"BAD stdout".as_slice()],
4298 )
4299 .expect("corrupt stdout");
4300
4301 let mut stdout = Vec::new();
4302 let mut stderr = Vec::new();
4303 let error = archive
4304 .restore_streams(
4305 &ArchiveKey {
4306 session: session(),
4307 source_call_id: "call-1".to_owned(),
4308 },
4309 &mut stdout,
4310 &mut stderr,
4311 )
4312 .expect_err("corrupt restore");
4313 assert!(error.contains("integrity verification"));
4314 assert!(stdout.is_empty());
4315 assert!(stderr.is_empty());
4316 }
4317
4318 #[test]
4319 fn verification_rejects_values_inserted_without_check_constraints() {
4320 let (_directory, mut archive) = archive();
4321 archive
4322 .begin_call(
4323 &session(),
4324 &call(),
4325 &serde_json::json!({}),
4326 &serde_json::json!({}),
4327 20,
4328 )
4329 .expect("begin");
4330 archive
4331 .connection
4332 .pragma_update(None, "ignore_check_constraints", "ON")
4333 .expect("ignore constraints");
4334 archive
4335 .connection
4336 .execute(
4337 "INSERT INTO snapshots (
4338 tool_call_id, subject, stage, media_type, captured_at_ms, payload_sha256
4339 )
4340 SELECT tool_call_id, 'other', stage, media_type, captured_at_ms, payload_sha256
4341 FROM snapshots WHERE subject = 'input' AND stage = 'before'",
4342 [],
4343 )
4344 .expect("invalid snapshot");
4345
4346 let report = archive.verify().expect("verify");
4347 assert!(
4348 report
4349 .errors
4350 .iter()
4351 .any(|error| error.contains("invalid subject"))
4352 );
4353 }
4354
4355 #[test]
4356 fn concurrent_first_open_initializes_the_schema_once() {
4357 let directory = TempDir::new().expect("temp directory");
4358 let path = directory.path().join("yarp/tool-calls.sqlite3");
4359 let barrier = std::sync::Arc::new(std::sync::Barrier::new(16));
4360 let threads: Vec<_> = (0..16)
4361 .map(|_| {
4362 let path = path.clone();
4363 let barrier = std::sync::Arc::clone(&barrier);
4364 std::thread::spawn(move || {
4365 barrier.wait();
4366 Archive::open_path(path)
4367 .and_then(|archive| archive.stats())
4368 .expect("open archive")
4369 })
4370 })
4371 .collect();
4372 for thread in threads {
4373 assert_eq!(thread.join().expect("join").calls, 0);
4374 }
4375 assert!(
4376 Archive::open_path(path)
4377 .expect("reopen")
4378 .verify()
4379 .expect("verify")
4380 .errors
4381 .is_empty()
4382 );
4383 }
4384
4385 #[test]
4386 fn concurrent_writers_keep_every_call() {
4387 let directory = TempDir::new().expect("temp directory");
4388 let path = directory.path().join("yarp/tool-calls.sqlite3");
4389 drop(Archive::open_path(path.clone()).expect("create"));
4390 let threads: Vec<_> = (0..8)
4391 .map(|index| {
4392 let path = path.clone();
4393 std::thread::spawn(move || {
4394 let mut archive = Archive::open_path(path).expect("open writer");
4395 let call = call_with_id(&format!("call-{index}"));
4396 archive
4397 .begin_call(
4398 &session(),
4399 &call,
4400 &serde_json::json!({"index": index}),
4401 &serde_json::json!({"index": index}),
4402 20,
4403 )
4404 .expect("write call");
4405 })
4406 })
4407 .collect();
4408 for thread in threads {
4409 thread.join().expect("join");
4410 }
4411 let archive = Archive::open_path(path).expect("reopen");
4412 assert_eq!(archive.stats().expect("stats").calls, 8);
4413 assert_eq!(archive.verify().expect("verify").incomplete_calls, 8);
4414 }
4415
4416 #[test]
4417 fn refuses_to_finish_executed_calls_without_required_snapshots() {
4418 let (_directory, mut archive) = archive();
4419 archive
4420 .begin_call(
4421 &session(),
4422 &call(),
4423 &serde_json::json!({}),
4424 &serde_json::json!({}),
4425 20,
4426 )
4427 .expect("begin");
4428 let error = archive
4429 .finish_call(
4430 &session(),
4431 "call-1",
4432 &serde_json::json!({}),
4433 false,
4434 true,
4435 40,
4436 )
4437 .expect_err("missing pre-result");
4438 assert!(error.contains("missing its pre-YARP result"));
4439 assert_eq!(archive.stats().expect("stats").incomplete_calls, 1);
4440 }
4441
4442 #[test]
4443 fn finishes_preflight_rejected_shell_calls_without_streams() {
4444 let (_directory, mut archive) = archive();
4445 let mut shell_call = call();
4446 shell_call.requires_streams = true;
4447 archive
4448 .begin_call(
4449 &session(),
4450 &shell_call,
4451 &serde_json::json!({}),
4452 &serde_json::json!({}),
4453 20,
4454 )
4455 .expect("begin");
4456 archive
4457 .finish_call(
4458 &session(),
4459 "call-1",
4460 &serde_json::json!({"error": "blocked"}),
4461 true,
4462 false,
4463 40,
4464 )
4465 .expect("finish preflight rejection");
4466 assert!(archive.verify().expect("verify").errors.is_empty());
4467 }
4468
4469 #[test]
4470 fn requires_a_pre_result_for_executed_shell_calls() {
4471 let (_directory, mut archive) = archive();
4472 let mut shell_call = call();
4473 shell_call.requires_streams = true;
4474 archive
4475 .begin_call(
4476 &session(),
4477 &shell_call,
4478 &serde_json::json!({}),
4479 &serde_json::json!({}),
4480 20,
4481 )
4482 .expect("begin");
4483 archive
4484 .capture_streams(
4485 &ArchiveKey {
4486 session: session(),
4487 source_call_id: "call-1".to_owned(),
4488 },
4489 30,
4490 &mut Cursor::new(b"stdout"),
4491 &mut Cursor::new(b"stderr"),
4492 b"stdout",
4493 b"stderr",
4494 )
4495 .expect("streams");
4496 let error = archive
4497 .finish_call(
4498 &session(),
4499 "call-1",
4500 &serde_json::json!({}),
4501 false,
4502 true,
4503 40,
4504 )
4505 .expect_err("missing pre-result");
4506 assert!(error.contains("missing its pre-YARP result"));
4507 }
4508
4509 #[test]
4510 fn requires_every_shell_stream_before_finishing() {
4511 let (_directory, mut archive) = archive();
4512 let mut shell_call = call();
4513 shell_call.requires_streams = true;
4514 archive
4515 .begin_call(
4516 &session(),
4517 &shell_call,
4518 &serde_json::json!({}),
4519 &serde_json::json!({}),
4520 20,
4521 )
4522 .expect("begin");
4523 archive
4524 .result_before(&session(), "call-1", &serde_json::json!({}), None, 30)
4525 .expect("before result");
4526 let error = archive
4527 .finish_call(
4528 &session(),
4529 "call-1",
4530 &serde_json::json!({}),
4531 false,
4532 true,
4533 40,
4534 )
4535 .expect_err("missing streams");
4536 assert!(error.contains("found 0 of 4"));
4537 assert_eq!(archive.stats().expect("stats").incomplete_calls, 1);
4538 }
4539
4540 #[test]
4541 fn rejects_truncated_and_oversized_ingest_frames() {
4542 let (_directory, first_archive) = archive();
4543 let truncated = run_ingest_with_archive(first_archive, Cursor::new(vec![0, 1]), Vec::new())
4544 .expect_err("truncated length");
4545 assert!(truncated.contains("truncated ingest frame length"));
4546
4547 let (_directory, second_archive) = archive();
4548 let oversized = (MAX_FRAME_BYTES + 1).to_be_bytes().to_vec();
4549 let error = run_ingest_with_archive(second_archive, Cursor::new(oversized), Vec::new())
4550 .expect_err("oversized frame");
4551 assert!(error.contains("invalid ingest frame length"));
4552 }
4553
4554 #[test]
4555 fn rejects_unsupported_ingest_schema_versions() {
4556 let (_directory, archive) = archive();
4557 let operation = serde_json::json!({
4558 "operation": "begin_call",
4559 "requestId": 1,
4560 "schemaVersion": 2,
4561 "session": session(),
4562 "call": call(),
4563 "inputBefore": {},
4564 "inputAfter": {},
4565 "capturedAtMs": 20
4566 });
4567 let mut output = Vec::new();
4568 run_ingest_with_archive(archive, Cursor::new(framed(&operation)), &mut output)
4569 .expect("protocol response");
4570 let ack: Value = serde_json::from_slice(&output).expect("ack");
4571 assert_eq!(ack["ok"], false);
4572 assert!(
4573 ack["error"]
4574 .as_str()
4575 .is_some_and(|error| error.contains("unsupported ingest"))
4576 );
4577 }
4578
4579 #[test]
4580 fn ingest_protocol_handles_every_operation() {
4581 let (directory, archive) = archive();
4582 let full_output_path = directory.path().join("full-output.log");
4583 fs::write(&full_output_path, b"complete output").expect("full output");
4584 let operations = [
4585 serde_json::json!({
4586 "operation": "begin_call",
4587 "requestId": 1,
4588 "schemaVersion": 1,
4589 "session": session(),
4590 "call": call(),
4591 "inputBefore": {"path": "a"},
4592 "inputAfter": {"path": "a"},
4593 "capturedAtMs": 20
4594 }),
4595 serde_json::json!({
4596 "operation": "result_before",
4597 "requestId": 2,
4598 "schemaVersion": 1,
4599 "session": session(),
4600 "sourceCallId": "call-1",
4601 "result": {"content": "before"},
4602 "fullOutputPath": full_output_path,
4603 "capturedAtMs": 30
4604 }),
4605 serde_json::json!({
4606 "operation": "stage_result",
4607 "requestId": 3,
4608 "schemaVersion": 1,
4609 "session": session(),
4610 "sourceCallId": "call-1",
4611 "result": {"content": "after"},
4612 "isError": false,
4613 "capturedAtMs": 40
4614 }),
4615 serde_json::json!({
4616 "operation": "update_final_result",
4617 "requestId": 4,
4618 "schemaVersion": 1,
4619 "session": session(),
4620 "sourceCallId": "call-1",
4621 "result": {"content": "final"},
4622 "isError": false,
4623 "finishedAtMs": 50
4624 }),
4625 ];
4626 let input: Vec<u8> = operations.iter().flat_map(framed).collect();
4627 let mut output = Vec::new();
4628 run_ingest_with_archive(archive, Cursor::new(input), &mut output).expect("ingest");
4629 let acknowledgements = String::from_utf8(output).expect("UTF-8");
4630 assert_eq!(acknowledgements.lines().count(), 4);
4631 assert!(acknowledgements.contains("\"requestId\":4"));
4632 let reopened = Archive::open_path(directory.path().join("data/tool-calls.sqlite3"))
4633 .expect("reopen archive");
4634 let source_outputs: i64 = reopened
4635 .connection
4636 .query_row(
4637 "SELECT count(*) FROM snapshots WHERE subject = 'source_output'",
4638 [],
4639 |row| row.get(0),
4640 )
4641 .expect("source outputs");
4642 assert_eq!(source_outputs, 1);
4643 }
4644
4645 #[cfg(unix)]
4646 #[test]
4647 fn repairs_archive_permissions_on_open() {
4648 use std::os::unix::fs::PermissionsExt as _;
4649 let directory = TempDir::new().expect("temp directory");
4650 let data = directory.path().join("yarp");
4651 let path = data.join("tool-calls.sqlite3");
4652 drop(Archive::open_path(path.clone()).expect("create"));
4653 fs::set_permissions(&data, fs::Permissions::from_mode(0o777)).expect("directory mode");
4654 fs::set_permissions(&path, fs::Permissions::from_mode(0o666)).expect("file mode");
4655 let archive = Archive::open_path_with_policy(path, true).expect("repair");
4656 assert!(archive.verify().expect("verify").errors.is_empty());
4657 }
4658
4659 #[cfg(unix)]
4660 #[test]
4661 fn refuses_to_modify_an_unrelated_public_directory() {
4662 use std::os::unix::fs::PermissionsExt as _;
4663 let directory = TempDir::new().expect("temp directory");
4664 let public = directory.path().join("yarp");
4665 fs::create_dir(&public).expect("public directory");
4666 fs::set_permissions(&public, fs::Permissions::from_mode(0o777)).expect("public mode");
4667 let error = Archive::open_path(public.join("archive.sqlite3"))
4668 .err()
4669 .expect("insecure directory error");
4670 assert!(error.contains("use a private directory"));
4671 let mode = fs::metadata(public).expect("metadata").permissions().mode() & 0o777;
4672 assert_eq!(mode, 0o777);
4673 }
4674
4675 #[test]
4676 fn ingest_protocol_acknowledges_committed_frames() {
4677 let directory = TempDir::new().expect("temp directory");
4678 let archive = Archive::open_path(directory.path().join("yarp/tool-calls.sqlite3"))
4679 .expect("open archive");
4680 let operation = serde_json::json!({
4681 "operation": "begin_call",
4682 "requestId": 7,
4683 "schemaVersion": 1,
4684 "session": session(),
4685 "call": call(),
4686 "inputBefore": {"path": "a"},
4687 "inputAfter": {"path": "a"},
4688 "capturedAtMs": 20
4689 });
4690 let body = serde_json::to_vec(&operation).expect("json");
4691 let mut input = Vec::new();
4692 input.extend_from_slice(&(body.len() as u64).to_be_bytes());
4693 input.extend_from_slice(&body);
4694 let mut output = Vec::new();
4695 run_ingest_with_archive(archive, Cursor::new(input), &mut output).expect("ingest");
4696 let ack: Value = serde_json::from_slice(&output).expect("ack");
4697 assert_eq!(ack["requestId"], 7);
4698 assert_eq!(ack["ok"], true);
4699 }
4700}