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