Skip to main content

agentmux_srv/backend/storage/filestore/
replace.rs

1// Copyright 2025-2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Multi-file replace and delete in one transaction (Phase 5a-2b,
5//! SPEC_AGENT_PANE_BOUNDED_LIVE_WINDOW_MIGRATION_2026_09_23.md §6.3.7 item 4).
6//!
7//! A transcript `output` has sidecars that describe its bytes (`output.idx`
8//! line offsets, `output.tsidx` receive times). Replacing or deleting
9//! `output` and its sidecars in separate steps leaves a window — or, after a
10//! crash, a permanent state — where a sidecar describes content that no
11//! longer exists, and a reader trusts it by coincidence of size. These do
12//! the whole change in one transaction.
13
14use rusqlite::{params, OptionalExtension};
15
16use super::core::FileStore;
17use super::{FileMeta, FileOpts};
18use crate::backend::storage::error::StoreError;
19
20/// Reads inside one [`FileStore::read_snapshot`]: every call sees the same
21/// moment of the database.
22pub struct SnapshotReader<'a> {
23    tx: &'a rusqlite::Transaction<'a>,
24}
25
26/// A file as a [`SnapshotReader`] sees it.
27#[derive(Debug, Clone)]
28pub struct SnapshotFile {
29    pub size: i64,
30    /// Valid generation; `None` when the file is not counted.
31    pub gen: Option<String>,
32    pub meta: FileMeta,
33}
34
35impl SnapshotReader<'_> {
36    pub fn file(&self, zone_id: &str, name: &str) -> Result<Option<SnapshotFile>, StoreError> {
37        let Some(row) = super::counter::read_row(self.tx, zone_id, name)? else { return Ok(None) };
38        let meta: String = self.tx.query_row(
39            "SELECT meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
40            params![zone_id, name],
41            |r| r.get(0),
42        )?;
43        Ok(Some(SnapshotFile {
44            size: row.size,
45            gen: row.counter().map(|(gen, _)| gen),
46            meta: serde_json::from_str(&meta).unwrap_or_default(),
47        }))
48    }
49
50    /// `[offset, offset + len)`, or `None` if any of it isn't stored.
51    pub fn bytes(&self, zone_id: &str, name: &str, offset: i64, len: i64) -> Result<Option<Vec<u8>>, StoreError> {
52        super::counter::read_bytes_exact(self.tx, zone_id, name, offset, len)
53    }
54}
55
56/// Write `data` at byte `offset` of a file's stored parts, inside `tx`,
57/// reading and rewriting only the parts the range touches. A part shorter
58/// than the range is extended (with zeros up to `offset`, which callers
59/// never leave as a gap). Leaves the file's recorded size to the caller.
60fn patch_range_in(tx: &rusqlite::Transaction<'_>, zone_id: &str, name: &str, offset: i64, data: &[u8]) -> Result<(), StoreError> {
61    if data.is_empty() {
62        return Ok(());
63    }
64    let pds = super::core::PART_DATA_SIZE as i64;
65    let end = offset + data.len() as i64;
66    for idx in (offset / pds)..=((end - 1) / pds) {
67        let part_start = idx * pds;
68        let mut part: Vec<u8> = tx
69            .query_row(
70                "SELECT data FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx = ?3",
71                params![zone_id, name, idx as i32],
72                |r| r.get(0),
73            )
74            .optional()?
75            .unwrap_or_default();
76        let from = offset.max(part_start);
77        let to = end.min(part_start + pds);
78        let (lo, hi) = ((from - part_start) as usize, (to - part_start) as usize);
79        if part.len() < hi {
80            part.resize(hi, 0);
81        }
82        part[lo..hi].copy_from_slice(&data[(from - offset) as usize..(to - offset) as usize]);
83        tx.execute(
84            "REPLACE INTO db_file_data (zoneid, name, partidx, data) VALUES (?1, ?2, ?3, ?4)",
85            params![zone_id, name, idx as i32, part],
86        )?;
87    }
88    Ok(())
89}
90
91/// Drop a file's stored bytes from `new_size` on, inside `tx`.
92fn truncate_parts_in(tx: &rusqlite::Transaction<'_>, zone_id: &str, name: &str, new_size: i64) -> Result<(), StoreError> {
93    let pds = super::core::PART_DATA_SIZE as i64;
94    let keep_parts = (new_size + pds - 1) / pds;
95    tx.execute(
96        "DELETE FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx >= ?3",
97        params![zone_id, name, keep_parts as i32],
98    )?;
99    let in_last = (new_size % pds) as usize;
100    if in_last > 0 {
101        let last = (keep_parts - 1) as i32;
102        let part: Option<Vec<u8>> = tx
103            .query_row(
104                "SELECT data FROM db_file_data WHERE zoneid = ?1 AND name = ?2 AND partidx = ?3",
105                params![zone_id, name, last],
106                |r| r.get(0),
107            )
108            .optional()?;
109        if let Some(mut part) = part.filter(|p| p.len() > in_last) {
110            part.truncate(in_last);
111            tx.execute(
112                "REPLACE INTO db_file_data (zoneid, name, partidx, data) VALUES (?1, ?2, ?3, ?4)",
113                params![zone_id, name, last, part],
114            )?;
115        }
116    }
117    Ok(())
118}
119
120/// See [`FileStore::derived_snapshot`].
121#[derive(Debug, Clone)]
122pub struct DerivedSnapshot {
123    pub output_size: i64,
124    /// `output`'s valid generation; `None` when it is not counted.
125    pub output_gen: Option<String>,
126    pub derived: Option<DerivedView>,
127}
128
129#[derive(Debug, Clone)]
130pub struct DerivedView {
131    pub size: i64,
132    /// The requested bytes; `None` if any of them isn't stored yet.
133    pub bytes: Option<Vec<u8>>,
134    pub meta: FileMeta,
135}
136
137impl FileStore {
138    /// Replace `name`'s content with `data` (creating the file if missing)
139    /// and delete the files in `drop`, in one transaction. The new content
140    /// starts a new counted epoch with a fresh `gen` (counter.rs).
141    #[track_caller]
142    pub fn replace_file(&self, zone_id: &str, name: &str, data: &[u8], drop: &[&str]) -> Result<(), StoreError> {
143        self.replace_inner(zone_id, name, data, None, drop, None).map(|_| ())
144    }
145
146    /// Replace `name`'s content with `data` (creating the file if missing)
147    /// and merge `meta` into its metadata (a null value removes a key), in
148    /// one transaction — so the metadata can describe the content without a
149    /// window where it describes other content.
150    #[track_caller]
151    #[cfg(test)]
152    pub fn put_file_with_meta(&self, zone_id: &str, name: &str, data: &[u8], meta: FileMeta) -> Result<(), StoreError> {
153        self.replace_inner(zone_id, name, data, Some(meta), &[], None).map(|_| ())
154    }
155
156    /// [`Self::put_file_with_meta`], only if `guard`'s valid generation is
157    /// still `expected_gen` — checked in the same transaction as the write.
158    /// For a file derived from `guard` (an index of it): if `guard` was
159    /// replaced while it was being derived, nothing is written and this
160    /// returns `false`. With `expected_gen` `None` (an uncounted `guard`, no
161    /// generation to compare) it always writes.
162    #[track_caller]
163    pub fn put_file_with_meta_if(
164        &self,
165        zone_id: &str,
166        name: &str,
167        data: &[u8],
168        meta: FileMeta,
169        guard: &str,
170        expected_gen: Option<&str>,
171    ) -> Result<bool, StoreError> {
172        self.replace_inner(zone_id, name, data, Some(meta), &[], Some((guard, expected_gen)))
173    }
174
175    /// Patch a file derived from `guard` (an index of it) in place, in one
176    /// transaction: write `head` at offset 0 and `tail` at `tail_at`, the file
177    /// ending where `tail` does — only if `guard` is still generation
178    /// `expected_gen` and `name` is still `expected_size` bytes, the state the
179    /// patch was computed from. Otherwise nothing is written and this returns
180    /// `false`. Only the parts the two ranges touch are written: extending
181    /// `output.idx` used to rewrite all of it (15 MB for a 1.9 M-line
182    /// transcript) for every few lines appended.
183    #[track_caller]
184    #[allow(clippy::too_many_arguments)]
185    pub fn patch_derived_if(
186        &self,
187        zone_id: &str,
188        name: &str,
189        head: &[u8],
190        tail_at: i64,
191        tail: &[u8],
192        guard: &str,
193        expected_gen: &str,
194        expected_size: i64,
195    ) -> Result<bool, StoreError> {
196        if tail_at < head.len() as i64 {
197            return Err(StoreError::Other(format!("{zone_id}/{name}: patch tail at {tail_at} overlaps its head")));
198        }
199        let now = agentmux_common::time::now_ms();
200        let written = self.write_txn(|tx| {
201            let current = super::counter::read_row(tx, zone_id, guard)?.and_then(|r| r.counter()).map(|(gen, _)| gen);
202            if current.as_deref() != Some(expected_gen) {
203                return Ok(false);
204            }
205            let size: Option<i64> = tx
206                .query_row(
207                    "SELECT size FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
208                    params![zone_id, name],
209                    |r| r.get(0),
210                )
211                .optional()?;
212            if size != Some(expected_size) || tail_at > expected_size {
213                return Ok(false);
214            }
215            patch_range_in(tx, zone_id, name, 0, head)?;
216            patch_range_in(tx, zone_id, name, tail_at, tail)?;
217            let new_size = tail_at + tail.len() as i64;
218            if new_size < expected_size {
219                truncate_parts_in(tx, zone_id, name, new_size)?;
220            }
221            tx.execute(
222                "UPDATE db_wave_file SET size = ?1, modts = ?2 WHERE zoneid = ?3 AND name = ?4",
223                params![new_size, now, zone_id, name],
224            )?;
225            Ok(true)
226        })?;
227        if written {
228            self.forget_cached(zone_id, &[name]);
229        }
230        Ok(written)
231    }
232
233    /// Delete several files of one zone in one transaction. Missing files
234    /// are not an error.
235    #[track_caller]
236    pub fn delete_files(&self, zone_id: &str, names: &[&str]) -> Result<(), StoreError> {
237        self.write_txn(|tx| {
238            for name in names {
239                tx.execute("DELETE FROM db_wave_file WHERE zoneid = ?1 AND name = ?2", params![zone_id, name])?;
240                tx.execute("DELETE FROM db_file_data WHERE zoneid = ?1 AND name = ?2", params![zone_id, name])?;
241            }
242            Ok(())
243        })?;
244        self.forget_cached(zone_id, names);
245        Ok(())
246    }
247
248    #[track_caller]
249    fn replace_inner(
250        &self,
251        zone_id: &str,
252        name: &str,
253        data: &[u8],
254        meta: Option<FileMeta>,
255        drop: &[&str],
256        guard: Option<(&str, Option<&str>)>,
257    ) -> Result<bool, StoreError> {
258        debug_assert!(!drop.contains(&name), "replace_file would drop the file it writes");
259        let now = agentmux_common::time::now_ms();
260        let opts_json = serde_json::to_string(&FileOpts::default())?;
261        let written = self.write_txn(|tx| {
262            if let Some((guard, Some(expected))) = guard {
263                let current = super::counter::read_row(tx, zone_id, guard)?.and_then(|r| r.counter()).map(|(gen, _)| gen);
264                if current.as_deref() != Some(expected) {
265                    return Ok(false);
266                }
267            }
268            let current_meta: Option<String> = tx
269                .query_row(
270                    "SELECT meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
271                    params![zone_id, name],
272                    |r| r.get(0),
273                )
274                .optional()?;
275            if current_meta.is_none() {
276                tx.execute(
277                    "INSERT INTO db_wave_file (zoneid, name, size, createdts, modts, opts, meta)
278                     VALUES (?1, ?2, 0, ?3, ?3, ?4, '{}')",
279                    params![zone_id, name, now, opts_json],
280                )?;
281            }
282            Self::replace_content_in(tx, zone_id, name, data, now)?;
283            if let Some(meta) = meta {
284                let mut merged: FileMeta =
285                    serde_json::from_str(current_meta.as_deref().unwrap_or("{}")).unwrap_or_default();
286                for (k, v) in meta {
287                    if v.is_null() {
288                        merged.remove(&k);
289                    } else {
290                        merged.insert(k, v);
291                    }
292                }
293                tx.execute(
294                    "UPDATE db_wave_file SET meta = ?1 WHERE zoneid = ?2 AND name = ?3",
295                    params![serde_json::to_string(&merged)?, zone_id, name],
296                )?;
297            }
298            for d in drop {
299                tx.execute("DELETE FROM db_wave_file WHERE zoneid = ?1 AND name = ?2", params![zone_id, d])?;
300                tx.execute("DELETE FROM db_file_data WHERE zoneid = ?1 AND name = ?2", params![zone_id, d])?;
301            }
302            Ok(true)
303        })?;
304        if written {
305            let mut forget = vec![name];
306            forget.extend_from_slice(drop);
307            self.forget_cached(zone_id, &forget);
308        }
309        Ok(written)
310    }
311
312    /// Bytes `[offset, offset + len)` of a file, read from the database, not
313    /// clamped to this process's cached size (which another process's append
314    /// can make stale); callers bound `len` by a size read from the database
315    /// (`line_state`). An error if any byte of the range isn't stored — a
316    /// file mid-write by a build without transactions claims bytes before it
317    /// holds them — so nothing is ever indexed or served as zeros.
318    #[track_caller]
319    pub fn read_bytes_db(&self, zone_id: &str, name: &str, offset: i64, len: i64) -> Result<Vec<u8>, StoreError> {
320        self.read_conn(|conn| super::counter::read_bytes_exact(conn, zone_id, name, offset, len))?.ok_or_else(|| {
321            StoreError::Other(format!("{zone_id}/{name}: bytes {offset}..{} not stored yet", offset + len))
322        })
323    }
324
325    /// Run `f` with a [`SnapshotReader`]: everything it reads comes from one
326    /// database snapshot — for a decision and the reads it licenses (an
327    /// index judged fresh, then its entries and the bytes they point at),
328    /// which must not straddle another instance's replace (Codex on #3634).
329    #[track_caller]
330    pub fn read_snapshot<T>(&self, f: impl FnOnce(&SnapshotReader<'_>) -> Result<T, StoreError>) -> Result<T, StoreError> {
331        self.read_txn(|tx| f(&SnapshotReader { tx }))
332    }
333
334    /// A file (`output`) and a file derived from it (`output.idx`), read in
335    /// ONE snapshot (Codex on #3634): `output`'s size and valid generation,
336    /// and the derived file's size, its first `head` bytes (all of it when
337    /// `head` is `None`) and its metadata. `None` when `output` doesn't exist.
338    #[track_caller]
339    pub fn derived_snapshot(
340        &self,
341        zone_id: &str,
342        output: &str,
343        derived: &str,
344        head: Option<i64>,
345    ) -> Result<Option<DerivedSnapshot>, StoreError> {
346        self.read_txn(|tx| {
347            let Some(row) = super::counter::read_row(tx, zone_id, output)? else { return Ok(None) };
348            let output_gen = row.counter().map(|(gen, _)| gen);
349            let derived_row: Option<(i64, String)> = tx
350                .query_row(
351                    "SELECT size, meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
352                    params![zone_id, derived],
353                    |r| Ok((r.get(0)?, r.get(1)?)),
354                )
355                .optional()?;
356            let derived = match derived_row {
357                None => None,
358                Some((size, meta)) => {
359                    let len = head.map_or(size, |h| h.min(size));
360                    Some(DerivedView {
361                        size,
362                        bytes: super::counter::read_bytes_exact(tx, zone_id, derived, 0, len)?,
363                        meta: serde_json::from_str(&meta).unwrap_or_default(),
364                    })
365                }
366            };
367            Ok(Some(DerivedSnapshot { output_size: row.size, output_gen, derived }))
368        })
369    }
370
371    /// A file's metadata, read from the database rather than the cache.
372    #[track_caller]
373    #[cfg(test)]
374    pub fn meta_db(&self, zone_id: &str, name: &str) -> Result<Option<FileMeta>, StoreError> {
375        let conn = self.lock_conn();
376        let meta: Option<String> = conn
377            .query_row(
378                "SELECT meta FROM db_wave_file WHERE zoneid = ?1 AND name = ?2",
379                params![zone_id, name],
380                |r| r.get(0),
381            )
382            .optional()?;
383        Ok(meta.map(|m| serde_json::from_str(&m).unwrap_or_default()))
384    }
385}