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