Skip to main content

yarp_cli/
archive.rs

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    /// Open the configured archive and initialize schema version 1 when needed.
286    ///
287    /// # Errors
288    ///
289    /// Returns an error when the path, permissions, `SQLite` database, or schema is invalid.
290    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    /// Open the configured archive without performing any writes.
296    ///
297    /// # Errors
298    ///
299    /// Returns an error when the archive is missing, unreadable, or uses another schema version.
300    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    /// Open an archive at an explicit path.
331    ///
332    /// # Errors
333    ///
334    /// Returns an error when the path, permissions, `SQLite` database, or schema is invalid.
335    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    /// Store one call and its input snapshots before execution.
600    ///
601    /// # Errors
602    ///
603    /// Returns an error when JSON encoding or the `SQLite` transaction fails.
604    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    /// Store the result exposed before YARP result processing.
643    ///
644    /// # Errors
645    ///
646    /// Returns an error when the call is missing or the payload cannot be committed.
647    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    /// Store the exact host-exposed result text selected for post-result reduction.
716    ///
717    /// # Errors
718    ///
719    /// Returns an error when the call is missing, completeness is invalid, or the snapshot cannot
720    /// be committed.
721    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    /// Resolve the stable reference for one internally identified call.
749    ///
750    /// # Errors
751    ///
752    /// Returns an error when the call or reference is missing.
753    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    /// Resolve and verify the canonical recovery sources for one call reference.
773    ///
774    /// # Errors
775    ///
776    /// Returns an error when the reference is malformed, missing, or any selected payload fails
777    /// integrity verification.
778    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    /// Store the provisional post-YARP result while leaving the call incomplete.
838    ///
839    /// # Errors
840    ///
841    /// Returns an error when the call is missing, required snapshots are missing, or the payload
842    /// cannot be committed.
843    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    /// Store the finalized result and mark the call finished.
877    ///
878    /// # Errors
879    ///
880    /// Returns an error when the call is missing or the payload cannot be committed.
881    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    /// Finalize a staged call with the result after every Pi result hook has run.
922    ///
923    /// # Errors
924    ///
925    /// Returns an error when the call is not staged or the final result cannot be committed.
926    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    /// Store exact shell streams before and after pruning.
1004    ///
1005    /// # Errors
1006    ///
1007    /// Returns an error when a stream cannot be read or the transaction cannot commit.
1008    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    /// Store exact shell streams when a wrapper deliberately passed output through unchanged.
1061    ///
1062    /// # Errors
1063    ///
1064    /// Returns an error when a stream cannot be read or the transaction cannot commit.
1065    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    /// Restore the exact pre-pruning shell streams for one call.
1118    ///
1119    /// # Errors
1120    ///
1121    /// Returns an error when the call, snapshots, or payload bodies cannot be read.
1122    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    /// Read aggregate archive statistics without reading payload content.
1155    ///
1156    /// # Errors
1157    ///
1158    /// Returns an error when `SQLite` queries or filesystem metadata reads fail.
1159    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    /// Verify `SQLite` integrity, relationships, snapshots, payloads, and permissions.
1199    ///
1200    /// # Errors
1201    ///
1202    /// Returns an error when verification itself cannot read the archive.
1203    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    /// Delete finished calls older than the given Unix millisecond timestamp.
1263    ///
1264    /// # Errors
1265    ///
1266    /// Returns an error when deletion, payload cleanup, or vacuuming fails.
1267    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
1597/// Read framed archive operations and acknowledge each committed transaction.
1598///
1599/// # Errors
1600///
1601/// Returns an error when the archive, frame, operation, or acknowledgement is invalid.
1602pub 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
1750/// Resolve the archive path from the YARP configuration.
1751///
1752/// # Errors
1753///
1754/// Returns an error when configuration or path resolution fails.
1755pub 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}