Skip to main content

pond/
sessions.rs

1//! The session datasets (spec.md#datasets): the three Lance tables, the
2//! `Store` facade, ingest validation, and `search_text` extraction.
3
4use std::{
5    collections::{BTreeMap, HashMap, HashSet},
6    path::Path,
7    sync::Arc,
8};
9
10use anyhow::{Context, Result};
11use arc_swap::ArcSwapOption;
12use arrow_select::filter::filter_record_batch;
13use async_stream::try_stream;
14use chrono::{DateTime, TimeZone, Utc};
15use lance::Dataset;
16use lance::dataset::{AutoCleanupParams, ProjectionRequest, WriteMode, WriteParams};
17use lance::deps::arrow_array::builder::{FixedSizeListBuilder, Float16Builder};
18use lance::deps::arrow_array::{
19    Array, ArrayRef, BooleanArray, FixedSizeListArray, Float16Array, Float32Array, Int32Array,
20    LargeBinaryArray, LargeStringArray, RecordBatch, RecordBatchIterator, StringArray,
21    TimestampMicrosecondArray, UInt64Array, new_null_array,
22};
23use lance::deps::arrow_schema::{DataType, Field, Schema, TimeUnit};
24use lance::deps::datafusion::error::DataFusionError;
25use lance::deps::datafusion::physical_plan::SendableRecordBatchStream;
26use lance::index::DatasetIndexExt;
27use lance_file::version::LanceFileVersion;
28use lance_index::scalar::{BuiltinIndexType, FullTextSearchQuery};
29use serde::{Deserialize, Serialize, de::DeserializeOwned};
30use serde_json::Value;
31use tokio_stream::{Stream, StreamExt};
32
33use crate::{
34    config, embed,
35    rowmap::{RowMetaEntry, RowMetaMap, RowMetaSet, discover_chain},
36    substrate::{
37        Handle, IndexIntent, IndexParamsKind, IndexStatus, IndexTrigger, MaintenancePolicy,
38        OptimizeProgressFn, PhaseOutcome, Predicate, ScalarValue, ScanOpts, Table,
39        TableOptimizeOutcome, TableSizes, VECTOR_INDEX_ACTIVATION_ROWS,
40    },
41    wire::{FileData, Message, Part, PartKind, Role, SUMMARY_PART_TYPES, Session, SessionFrom},
42};
43use url::Url;
44
45#[derive(Debug)]
46pub struct Store {
47    handle: Handle,
48    /// Resident per-message meta map for index-only hit resolution and in-memory
49    /// hydration (see [`crate::rowmap`]). `None` until [`Store::ensure_rowmap`]
50    /// builds it (local tests, pre-prewarm), where the arms fall back to a
51    /// data-projection scan and hydration to `take_rows`. `ArcSwap` so a
52    /// version-bump rebuild swaps it under concurrent searches.
53    rowmap: ArcSwapOption<RowMetaSet>,
54    /// Resident embedder for inline embed-at-ingest. `None` keeps ingest
55    /// writing null vectors (tests, search-only stores); the CLI write paths
56    /// attach one via [`Store::with_embedder`]. Lazy, so a store that never
57    /// ingests an embeddable row never loads the model.
58    embedder: Option<Arc<crate::embed::LazyEmbedder>>,
59    /// Observer for inline embed-at-ingest: `(embedded_so_far, total)` per
60    /// model batch within one flush. Lets the CLI keep its progress line
61    /// moving through the otherwise-opaque commit phase.
62    ingest_embed_progress: Option<IngestEmbedProgress>,
63}
64
65/// One ingest host's slice of a shared store (see
66/// [`Store::ingest_host_activity`]). `hostname: None` groups rows carrying
67/// no host stamp.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub struct HostActivity {
70    pub hostname: Option<String>,
71    pub sessions: usize,
72    pub last_message_at: DateTime<Utc>,
73}
74
75/// Callback wrapper for [`Store::with_ingest_embed_progress`]; a newtype so
76/// `Store` keeps its derived `Debug`.
77#[derive(Clone)]
78pub struct IngestEmbedProgress(pub Arc<dyn Fn(usize, usize) + Send + Sync>);
79
80impl std::fmt::Debug for IngestEmbedProgress {
81    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
82        f.write_str("IngestEmbedProgress")
83    }
84}
85
86#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
87pub struct LanceArchiveCounts {
88    pub sessions: usize,
89    pub messages: usize,
90    pub parts: usize,
91}
92
93#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
94pub struct LanceArchiveVersions {
95    pub sessions: u64,
96    pub messages: u64,
97    pub parts: u64,
98}
99
100#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
101pub struct LanceArchiveExport {
102    pub rows: LanceArchiveCounts,
103    pub source_versions: LanceArchiveVersions,
104}
105
106#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
107pub struct LanceArchiveImport {
108    pub rows: LanceArchiveCounts,
109    pub inserted: LanceArchiveCounts,
110}
111
112/// One table's slice of a store-to-store copy plan: which sessions' rows for
113/// that table can be **appended** versus **filtered-appended** (the `merge`
114/// bucket, which despite the name appends too) (spec.md#session-durable-copy).
115/// The choice is made per table by row presence on the destination, because the
116/// three tables are written by separate commits and an interrupted copy can
117/// leave them in different states (e.g. the small `sessions` table committed but
118/// `messages` not):
119/// - `append`: the destination has **zero** rows for the session in this table,
120///   so they cannot collide -> append (no merge join, no target probe; the
121///   bandwidth-bound fast path).
122/// - `merge`: the destination already has *some* rows but the source has more ->
123///   append only the rows still absent (filtered against the destination's PKs).
124#[derive(Debug, Clone, Default)]
125pub struct TablePlan {
126    pub append: Vec<String>,
127    pub merge: Vec<String>,
128}
129
130impl TablePlan {
131    pub fn is_empty(&self) -> bool {
132        self.append.is_empty() && self.merge.is_empty()
133    }
134}
135
136/// A store-to-store `pond copy` plan, decided per table (see [`TablePlan`]).
137/// `source_sessions` is the full source session count, kept so the caller can
138/// tell "destination already up to date" (empty plan, non-empty source) from
139/// "empty source", and so each table can recognize a from-empty/resumed run
140/// (`append.len() == source_sessions`) and skip the per-session `IN` filter.
141#[derive(Debug, Clone, Default)]
142pub struct DeltaPlan {
143    pub sessions: TablePlan,
144    pub messages: TablePlan,
145    pub parts: TablePlan,
146    pub source_sessions: usize,
147}
148
149impl DeltaPlan {
150    pub fn is_empty(&self) -> bool {
151        self.sessions.is_empty() && self.messages.is_empty() && self.parts.is_empty()
152    }
153
154    /// Sessions whose own row is absent on the destination - the "new" count for
155    /// the plan receipt. Sessions never grow in row count (one immutable row
156    /// each), so the `sessions` table only ever appends.
157    pub fn new_sessions(&self) -> usize {
158        self.sessions.append.len()
159    }
160
161    /// Distinct sessions touched by the copy across all three tables - the
162    /// figure the progress bar totals against.
163    pub fn total(&self) -> usize {
164        let mut seen = std::collections::HashSet::new();
165        for plan in [&self.sessions, &self.messages, &self.parts] {
166            seen.extend(plan.append.iter());
167            seen.extend(plan.merge.iter());
168        }
169        seen.len()
170    }
171}
172
173#[derive(Debug, Clone, Default)]
174pub struct IndexIntents {
175    pub sessions: Vec<IndexIntent>,
176    pub messages: Vec<IndexIntent>,
177    pub parts: Vec<IndexIntent>,
178}
179
180impl IndexIntents {
181    fn all(&self) -> [(Table, &[IndexIntent]); 3] {
182        [
183            (Table::Sessions, &self.sessions),
184            (Table::Messages, &self.messages),
185            (Table::Parts, &self.parts),
186        ]
187    }
188}
189
190/// A message awaiting embedding: its primary key plus the `search_text` to
191/// embed. The vector lives on the same `messages` row, so no denormalized
192/// filter columns are needed (spec.md#session-embed-from-canonical).
193#[derive(Debug, Clone, PartialEq)]
194pub struct PendingMessage {
195    pub session_id: String,
196    pub id: String,
197    pub search_text: String,
198}
199
200/// One embedded message: a primary key and the vector to store. `pond optimize`
201/// writes a batch of these into `messages.vector` keyed on `(session_id, id)`.
202#[derive(Debug, Clone, PartialEq)]
203pub struct EmbeddedMessage {
204    pub session_id: String,
205    pub id: String,
206    pub vector: Vec<f32>,
207}
208
209/// Message metadata used to hydrate search hits after retriever ranking.
210#[derive(Debug, Clone, PartialEq)]
211pub struct MessageMeta {
212    pub message_id: String,
213    pub session_id: String,
214    pub role: String,
215    pub project: String,
216    pub source_agent: String,
217    pub timestamp: DateTime<Utc>,
218    pub search_text: String,
219}
220
221#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
222pub struct MessageKey {
223    pub session_id: String,
224    pub message_id: String,
225}
226
227/// One retrieval-arm hit. `rowid` is `Some` when the row meta map (or its
228/// take_rows miss-fallback) resolved a stable row id, which lets hydration
229/// `take_rows` the exact row instead of re-finding it with an `IN`-predicate
230/// scan; `None` on the no-map fallback path (local tests, pre-prewarm).
231#[derive(Debug, Clone, PartialEq)]
232pub struct SearchHit {
233    pub rowid: Option<u64>,
234    pub key: MessageKey,
235    pub score: f32,
236}
237
238#[derive(Debug, Clone, Copy, PartialEq, Eq)]
239pub enum UpsertStatus {
240    Inserted,
241    Matched,
242}
243
244/// What one `Store::optimize_indices` or `Store::build_indices_only` pass did
245/// across every table. Each [`TableOptimizeOutcome`] reports phase-by-phase
246/// results so the CLI can render compaction-skipped (under writer contention)
247/// distinctly from index-build failure (real problem).
248#[derive(Debug, Default)]
249pub struct OptimizeOutcome {
250    pub tables: Vec<TableOptimizeOutcome>,
251}
252
253impl OptimizeOutcome {
254    /// True if any table's indices phase reported a non-conflict failure.
255    /// `SkippedConflict` is expected under contention and does not count.
256    pub fn any_indices_failed(&self) -> bool {
257        self.tables.iter().any(|t| t.indices.is_failed())
258    }
259
260    /// Treat any `Failed` phase as an error. Tests that don't run under
261    /// contention use this to keep their existing `.await?` style: a real
262    /// failure becomes an `Err`, while `SkippedConflict` is impossible there.
263    pub fn into_result(self) -> Result<Self> {
264        for table in &self.tables {
265            if let PhaseOutcome::Failed(error) = &table.indices {
266                anyhow::bail!(
267                    "indices phase failed on {}: {error:#}",
268                    table.table.as_str()
269                );
270            }
271            if let PhaseOutcome::Failed(error) = &table.compaction {
272                anyhow::bail!(
273                    "compaction phase failed on {}: {error:#}",
274                    table.table.as_str()
275                );
276            }
277        }
278        Ok(self)
279    }
280}
281
282#[derive(Debug, Clone, Copy, PartialEq, Eq)]
283pub struct RowTotals {
284    pub sessions: u64,
285    pub messages: u64,
286    pub parts: u64,
287}
288
289/// Embedding coverage for `pond status` / `pond optimize`. `total` is the count of
290/// `messages` rows that carry `search_text` (i.e. are eligible to embed); rows
291/// without `search_text` produce no vector. `embedded` is the subset of those
292/// already carrying a vector under the current [`embed::model_id()`]. `backlog`
293/// is the authoritative count still owed an embedding (`total - embedded` by
294/// construction), read live from the dataset rather than derived by subtracting
295/// the FTS `num_docs`, which over-counts deleted-but-unpurged docs and would
296/// otherwise report a phantom backlog that never clears.
297#[derive(Debug, Clone, Copy, PartialEq, Eq)]
298pub struct EmbeddingProgress {
299    pub embedded: usize,
300    pub total: usize,
301    pub backlog: usize,
302    pub model: &'static str,
303}
304
305#[derive(Debug, Clone, Copy)]
306pub struct MessageWrite<'a> {
307    pub message: &'a Message,
308    pub parts: &'a [Part],
309    pub search_text: Option<&'a str>,
310}
311
312impl Store {
313    /// Open against a local filesystem URL or a remote one for which the
314    /// caller has no extra options to pass (env vars suffice). CLI verbs
315    /// that load `[storage]` from config should call
316    /// [`Store::open_with_options`] instead so the same options flow into
317    /// every dataset open and write.
318    pub async fn open(location: &Url) -> Result<Self> {
319        Ok(Self {
320            handle: Handle::open(location).await?,
321            rowmap: ArcSwapOption::empty(),
322            embedder: None,
323            ingest_embed_progress: None,
324        })
325    }
326
327    /// Attach a resident embedder so [`Store::upsert_session_batch`] embeds
328    /// eligible messages inline, in the same append commit as the rows
329    /// (spec.md#session-embed-from-canonical). The CLI write paths set this;
330    /// every other open leaves it `None` and writes null vectors as before.
331    #[must_use]
332    pub fn with_embedder(mut self, embedder: Arc<crate::embed::LazyEmbedder>) -> Self {
333        self.embedder = Some(embedder);
334        self
335    }
336
337    /// Attach an inline-embed progress observer (see [`IngestEmbedProgress`]).
338    #[must_use]
339    pub fn with_ingest_embed_progress(mut self, progress: IngestEmbedProgress) -> Self {
340        self.ingest_embed_progress = Some(progress);
341        self
342    }
343
344    /// Live byte size of the shared Lance session caches (index + metadata).
345    /// Diagnostic only - walks the caches.
346    pub fn lance_cache_bytes(&self) -> u64 {
347        self.handle.lance_cache_bytes()
348    }
349
350    /// Open with object-store options (S3 creds, region, endpoint, ...)
351    /// threaded through Lance verbatim. Keys are the standard `object_store`
352    /// config names; pond does not parse them. Empty options + default caps
353    /// is equivalent to [`Store::open`]. Cache caps come from the `[runtime]`
354    /// config block via [`crate::substrate::RuntimeCaps`].
355    pub async fn open_with_options(
356        location: &Url,
357        storage_options: std::collections::HashMap<String, String>,
358        caps: crate::substrate::RuntimeCaps,
359    ) -> Result<Self> {
360        Ok(Self {
361            handle: Handle::open_with_options(location, storage_options, caps).await?,
362            rowmap: ArcSwapOption::empty(),
363            embedder: None,
364            ingest_embed_progress: None,
365        })
366    }
367
368    /// Like [`Self::open_with_options`], plus the on-disk `_indices/*` cache
369    /// rooted at `index_cache_dir` (see [`Handle::open_with_options_cached`]).
370    pub async fn open_with_options_cached(
371        location: &Url,
372        storage_options: std::collections::HashMap<String, String>,
373        caps: crate::substrate::RuntimeCaps,
374        index_cache_dir: Option<std::path::PathBuf>,
375    ) -> Result<Self> {
376        Ok(Self {
377            handle: Handle::open_with_options_cached(
378                location,
379                storage_options,
380                caps,
381                index_cache_dir,
382            )
383            .await?,
384            rowmap: ArcSwapOption::empty(),
385            embedder: None,
386            ingest_embed_progress: None,
387        })
388    }
389
390    /// Convenience for tests and CLI verbs holding a `&Path`: wraps the path in
391    /// a `file://...` URL via [`config::url_for_path`] before opening. Routes
392    /// through [`Store::open_with_options`] so the production policy is
393    /// applied; tests get the backend-aware local-FS defaults.
394    pub async fn open_local(path: impl AsRef<std::path::Path>) -> Result<Self> {
395        let url = config::url_for_path(path)?;
396        Self::open_with_options(
397            &url,
398            std::collections::HashMap::new(),
399            crate::substrate::RuntimeCaps::default(),
400        )
401        .await
402    }
403
404    /// Export clean, index-free Lance datasets into `dest`.
405    ///
406    /// This rewrites the visible rows of each table instead of copying the
407    /// dataset roots. The resulting manifests therefore contain no references
408    /// to the source store's `_indices`, while `messages.vector` and
409    /// `messages.embedding_model` remain ordinary data columns and are
410    /// preserved.
411    pub async fn export_clean_lance_datasets(&self, dest: &Path) -> Result<LanceArchiveExport> {
412        std::fs::create_dir_all(dest)
413            .with_context(|| format!("failed to create archive staging dir {}", dest.display()))?;
414        let (sessions, sessions_version) = self
415            .export_clean_table(Table::Sessions, &dest.join("sessions.lance"))
416            .await?;
417        let (messages, messages_version) = self
418            .export_clean_table(Table::Messages, &dest.join("messages.lance"))
419            .await?;
420        let (parts, parts_version) = self
421            .export_clean_table(Table::Parts, &dest.join("parts.lance"))
422            .await?;
423        Ok(LanceArchiveExport {
424            rows: LanceArchiveCounts {
425                sessions,
426                messages,
427                parts,
428            },
429            source_versions: LanceArchiveVersions {
430                sessions: sessions_version,
431                messages: messages_version,
432                parts: parts_version,
433            },
434        })
435    }
436
437    pub async fn import_clean_lance_datasets(&self, source: &Path) -> Result<LanceArchiveImport> {
438        let sessions_dataset =
439            open_archive_table(Table::Sessions, &source.join("sessions.lance")).await?;
440        let messages_dataset =
441            open_archive_table(Table::Messages, &source.join("messages.lance")).await?;
442        let parts_dataset = open_archive_table(Table::Parts, &source.join("parts.lance")).await?;
443        // Validate every table's schema before importing any: a
444        // mixed-compatibility archive must fail whole, not half-restore.
445        let sessions_upgrade = archive_schema_backfill(&sessions_dataset, Table::Sessions)?;
446        let messages_upgrade = archive_schema_backfill(&messages_dataset, Table::Messages)?;
447        let parts_upgrade = archive_schema_backfill(&parts_dataset, Table::Parts)?;
448        let (sessions, sessions_inserted) = self
449            .import_clean_table(Table::Sessions, sessions_dataset, sessions_upgrade)
450            .await?;
451        let (messages, messages_inserted) = self
452            .import_clean_table(Table::Messages, messages_dataset, messages_upgrade)
453            .await?;
454        let (parts, parts_inserted) = self
455            .import_clean_table(Table::Parts, parts_dataset, parts_upgrade)
456            .await?;
457        Ok(LanceArchiveImport {
458            rows: LanceArchiveCounts {
459                sessions,
460                messages,
461                parts,
462            },
463            inserted: LanceArchiveCounts {
464                sessions: sessions_inserted,
465                messages: messages_inserted,
466                parts: parts_inserted,
467            },
468        })
469    }
470
471    async fn export_clean_table(&self, table: Table, dest: &Path) -> Result<(usize, u64)> {
472        let dataset = self.handle.dataset(table).await?;
473        let source_version = dataset.version_id();
474        let schema = export_schema(table);
475        let mut scan = dataset.scan();
476        // The default scan projects blob columns as descriptor structs
477        // (position/size into the source's blob storage) - meaningless in an
478        // archive and unwritable at V2_1. `AllBinary` materializes the bytes
479        // so the rewritten table is self-contained.
480        scan.blob_handling(lance::datatypes::BlobHandling::AllBinary);
481        let mut stream = scan
482            .try_into_stream()
483            .await
484            .with_context(|| format!("failed to scan {} for archive export", table.as_str()))?;
485        let dest_uri = dest
486            .to_str()
487            .with_context(|| format!("archive path is not UTF-8: {}", dest.display()))?;
488
489        let mut rows = 0usize;
490        let mut wrote = false;
491        while let Some(batch) = stream.next().await {
492            let batch = batch
493                .with_context(|| format!("failed to read {} archive batch", table.as_str()))?;
494            rows += batch.num_rows();
495            let reader = RecordBatchIterator::new([Ok(batch.clone())], batch.schema());
496            let mut params = write_params_for_create();
497            if wrote {
498                params.mode = WriteMode::Append;
499            }
500            Dataset::write(reader, dest_uri, Some(params))
501                .await
502                .with_context(|| format!("failed to write {} archive table", table.as_str()))?;
503            wrote = true;
504        }
505
506        if !wrote {
507            let batch = RecordBatch::new_empty(schema.clone());
508            let reader = RecordBatchIterator::new([Ok(batch)], schema);
509            Dataset::write(reader, dest_uri, Some(write_params_for_create()))
510                .await
511                .with_context(|| {
512                    format!("failed to write empty {} archive table", table.as_str())
513                })?;
514        }
515        Ok((rows, source_version))
516    }
517
518    /// `upgrade` carries the batch-level backfill for a pre-upgrade archive:
519    /// an archive is a snapshot, so restore MUST NOT mutate it in place - the
520    /// missing cells derive at the read boundary through the same recipe the
521    /// store migration uses.
522    async fn import_clean_table(
523        &self,
524        table: Table,
525        dataset: Dataset,
526        upgrade: Option<ColumnBackfill>,
527    ) -> Result<(usize, usize)> {
528        // Force the destination table into existence up front: an empty
529        // archive table yields zero batches, so merge_insert alone would
530        // leave a lazily-created table (sessions or parts) missing on the destination.
531        let _ = self.handle.dataset(table).await?;
532        self.merge_scanner(table, dataset.scan(), "archive import", upgrade)
533            .await
534    }
535
536    /// Stream a prepared source `scanner` into this store's `table` via
537    /// `merge_insert_stats`, materializing blob bytes (not descriptor structs)
538    /// so the merge writes them into the destination's own schema. Shared by
539    /// the archive-restore and store-to-store copy paths; `context` names the
540    /// caller in error messages. Returns (rows scanned, rows inserted).
541    async fn merge_scanner(
542        &self,
543        table: Table,
544        mut scanner: lance::dataset::scanner::Scanner,
545        context: &'static str,
546        upgrade: Option<ColumnBackfill>,
547    ) -> Result<(usize, usize)> {
548        scanner.blob_handling(lance::datatypes::BlobHandling::AllBinary);
549        let mut stream = scanner
550            .try_into_stream()
551            .await
552            .with_context(|| format!("failed to scan {} for {context}", table.as_str()))?;
553        let mut rows = 0usize;
554        let mut inserted = 0usize;
555        while let Some(batch) = stream.next().await {
556            let batch = batch
557                .with_context(|| format!("failed to read {} {context} batch", table.as_str()))?;
558            let batch = match &upgrade {
559                Some(spec) => upgraded_batch(&batch, spec)?,
560                None => batch,
561            };
562            let row_count = batch.num_rows();
563            rows += row_count;
564            let stats = self
565                .handle
566                .merge_insert_stats(table, batch, row_count)
567                .await
568                .with_context(|| format!("failed to merge {} during {context}", table.as_str()))?;
569            inserted += (stats.num_inserted_rows + stats.num_updated_rows) as usize;
570        }
571        Ok((rows, inserted))
572    }
573
574    /// Per-session message count - the data-intrinsic freshness key for
575    /// incremental `pond copy`. pond is append-only (merge is
576    /// `WhenMatched::DoNothing`; no edits or deletes), so this count rises iff a
577    /// session gained messages, catching growth a `MAX(timestamp)` key would
578    /// miss when a new message shares the session's latest timestamp. The count
579    /// is source-authored and survives the copy unchanged, so it compares
580    /// soundly across two stores with independent clocks
581    /// (spec.md#session-durable-copy). Projects only
582    /// the one column it counts; resolves the `session_id` array once per batch
583    /// and allocates a key only on a session's first row. Distinct from
584    /// `session_message_counts`, which counts a supplied id list one query each;
585    /// this counts every session in a single scan.
586    pub async fn all_session_message_counts(&self) -> Result<HashMap<String, usize>> {
587        self.all_session_row_counts(Table::Messages).await
588    }
589
590    pub async fn all_session_part_counts(&self) -> Result<HashMap<String, usize>> {
591        self.all_session_row_counts(Table::Parts).await
592    }
593
594    /// Count rows per `session_id` across one table in a single scan, projecting
595    /// only the `session_id` column and allocating a key on a session's first
596    /// row. Both `messages` and `parts` lead their primary key with `session_id`
597    /// (`lance-table-creation-session-scoped-pk`).
598    async fn all_session_row_counts(&self, table: Table) -> Result<HashMap<String, usize>> {
599        let scanner = self
600            .handle
601            .scan(table, ScanOpts::project_only(&["session_id"]))
602            .await?;
603        let mut stream = scanner.try_into_stream().await?;
604        let mut out: HashMap<String, usize> = HashMap::new();
605        while let Some(batch) = stream.next().await {
606            let batch = batch?;
607            let session_ids = batch
608                .column_by_name("session_id")
609                .context("scan projection dropped the session_id column")?
610                .as_any()
611                .downcast_ref::<StringArray>()
612                .context("session_id column is not Utf8")?;
613            for row in 0..batch.num_rows() {
614                if session_ids.is_null(row) {
615                    continue;
616                }
617                let session_id = session_ids.value(row);
618                if let Some(count) = out.get_mut(session_id) {
619                    *count += 1;
620                } else {
621                    out.insert(session_id.to_owned(), 1);
622                }
623            }
624        }
625        Ok(out)
626    }
627
628    /// Plan an incremental store-to-store copy into `self` from `source`,
629    /// deciding **per table** whether each source session's rows can be appended
630    /// wholesale (the destination has none, so they cannot collide) or go to the
631    /// `merge` bucket - which the copy executes as a filtered append, keeping
632    /// only the rows still absent (the destination has some, source has more).
633    /// Reads both id-sets plus per-session message and part counts. Parts have
634    /// their own data-derived signal so a part added under an existing message
635    /// routes through the `merge` bucket instead of relying on the closing verify
636    /// to catch it (spec.md#session-movement-complete).
637    pub async fn plan_incremental_from(&self, source: &Store) -> Result<DeltaPlan> {
638        let (
639            source_ids,
640            dest_ids,
641            source_msg_counts,
642            dest_msg_counts,
643            source_part_counts,
644            dest_part_counts,
645        ) = tokio::try_join!(
646            source.collect_ids(Table::Sessions),
647            self.collect_ids(Table::Sessions),
648            source.all_session_message_counts(),
649            self.all_session_message_counts(),
650            source.all_session_part_counts(),
651            self.all_session_part_counts(),
652        )?;
653        let source_sessions = source_ids.len();
654        let mut plan = DeltaPlan {
655            source_sessions,
656            ..DeltaPlan::default()
657        };
658        for id in &source_ids {
659            // The `sessions` table holds one immutable row per session, so it
660            // only ever appends an absent id - a present row is identical.
661            if !dest_ids.contains(id) {
662                plan.sessions.append.push(id.clone());
663            }
664            let source_msgs = source_msg_counts.get(id).copied().unwrap_or(0);
665            let dest_msgs = dest_msg_counts.get(id).copied().unwrap_or(0);
666            if dest_msgs == 0 {
667                if source_msgs > 0 {
668                    plan.messages.append.push(id.clone());
669                }
670            } else if source_msgs > dest_msgs {
671                plan.messages.merge.push(id.clone());
672            }
673            let source_parts = source_part_counts.get(id).copied().unwrap_or(0);
674            let dest_parts = dest_part_counts.get(id).copied().unwrap_or(0);
675            if dest_parts == 0 {
676                if source_parts > 0 {
677                    plan.parts.append.push(id.clone());
678                }
679            } else if source_parts > dest_parts {
680                plan.parts.merge.push(id.clone());
681            }
682        }
683        Ok(plan)
684    }
685
686    /// Copy the planned delta from `source` into `self`, streaming the source
687    /// scan straight into the destination - no local staging copy. Each table
688    /// picks its primitive per session from its [`TablePlan`]
689    /// (spec.md#session-durable-copy): **append** the sessions whose rows are
690    /// absent here (cannot collide; one commit per scan, bandwidth-bound), then
691    /// for the grown ones append just their absent rows after filtering out the
692    /// ones already present. Append-only storage is what makes the append safe: a
693    /// re-run re-plans from current destination state, so an
694    /// interrupted-then-resumed copy never double-appends (landed rows are no
695    /// longer absent).
696    pub async fn copy_delta_from(
697        &self,
698        source: &Store,
699        plan: &DeltaPlan,
700    ) -> Result<LanceArchiveImport> {
701        // The three tables are independent Lance datasets with separate write
702        // locks, so copy them concurrently - mirrors the ingest path's
703        // three-table `try_join!` (see `upsert_session_batch`).
704        let ((sessions, sessions_inserted), (messages, messages_inserted), (parts, parts_inserted)) =
705            tokio::try_join!(
706                self.copy_table(
707                    source,
708                    Table::Sessions,
709                    "id",
710                    &plan.sessions,
711                    plan.source_sessions,
712                ),
713                self.copy_table(
714                    source,
715                    Table::Messages,
716                    "session_id",
717                    &plan.messages,
718                    plan.source_sessions,
719                ),
720                self.copy_table(
721                    source,
722                    Table::Parts,
723                    "session_id",
724                    &plan.parts,
725                    plan.source_sessions,
726                ),
727            )?;
728        Ok(LanceArchiveImport {
729            rows: LanceArchiveCounts {
730                sessions,
731                messages,
732                parts,
733            },
734            inserted: LanceArchiveCounts {
735                sessions: sessions_inserted,
736                messages: messages_inserted,
737                parts: parts_inserted,
738            },
739        })
740    }
741
742    /// Copy one table's slice of the plan: append the absent sessions, then
743    /// append the grown sessions' absent rows. Sequential within a table (one
744    /// write lock); `copy_delta_from` runs the three tables in parallel. Returns
745    /// (rows, inserted), equal since both paths append and neither dedups.
746    async fn copy_table(
747        &self,
748        source: &Store,
749        table: Table,
750        key_column: &'static str,
751        table_plan: &TablePlan,
752        source_sessions: usize,
753    ) -> Result<(usize, usize)> {
754        // Force the destination table into existence up front so a lazily
755        // created table (sessions or parts) is never left missing when its slice
756        // is empty - same reason as the archive import path.
757        let _ = self.handle.dataset(table).await?;
758
759        let appended = self
760            .append_sessions(
761                source,
762                table,
763                key_column,
764                &table_plan.append,
765                source_sessions,
766            )
767            .await?;
768
769        // `Sessions` never reaches here: its row is immutable, so a present
770        // session is identical and routes to neither append nor merge.
771        let mut grown = 0usize;
772        for chunk in table_plan.merge.chunks(COPY_SESSION_IN_CHUNK) {
773            let values: Vec<ScalarValue> = chunk
774                .iter()
775                .map(|id| ScalarValue::String(id.clone()))
776                .collect();
777            let predicate = Predicate::In("session_id", values.clone());
778            let stream = Self::source_scan(source, table, Some(&predicate)).await?;
779            grown += match table {
780                Table::Messages => {
781                    let present = Arc::new(self.present_message_pks(&values).await?);
782                    self.append_filtered(table, stream, Self::message_keep(present))
783                        .await?
784                }
785                Table::Parts => {
786                    let present = Arc::new(self.present_part_pks(&values).await?);
787                    self.append_filtered(table, stream, Self::part_keep(present))
788                        .await?
789                }
790                Table::Sessions => 0,
791            };
792        }
793        Ok((appended + grown, appended + grown))
794    }
795
796    /// Append one table's slice for the listed sessions. A from-empty or resumed
797    /// copy (`session_ids.len() == source_sessions`: every session's rows for
798    /// this table are absent on the destination) scans the source wholesale
799    /// under one commit; a partial copy chunks the `IN` predicate (btree-pushed)
800    /// but still commits once per chunk, not per scan batch. Returns rows
801    /// appended.
802    async fn append_sessions(
803        &self,
804        source: &Store,
805        table: Table,
806        key_column: &'static str,
807        session_ids: &[String],
808        source_sessions: usize,
809    ) -> Result<usize> {
810        if session_ids.is_empty() {
811            return Ok(0);
812        }
813        if session_ids.len() == source_sessions {
814            return self.append_scanner(source, table, None).await;
815        }
816        let mut rows = 0usize;
817        for chunk in session_ids.chunks(COPY_SESSION_IN_CHUNK) {
818            let predicate = in_predicate(key_column, chunk);
819            rows += self.append_scanner(source, table, Some(&predicate)).await?;
820        }
821        Ok(rows)
822    }
823
824    /// Append a prepared source scan into this store's `table` via
825    /// `Handle::append_stream`, materializing blob bytes (`AllBinary`) so the
826    /// write is self-contained. The closure is a *factory*: `append_stream`
827    /// rebuilds the one-shot scan stream on each OCC attempt. Returns rows
828    /// appended.
829    async fn append_scanner(
830        &self,
831        source: &Store,
832        table: Table,
833        predicate: Option<&Predicate>,
834    ) -> Result<usize> {
835        let make_source = || Self::source_scan(source, table, predicate);
836        let stats = self.handle.append_stream(table, make_source).await?;
837        Ok(stats.rows as usize)
838    }
839
840    /// Append source rows whose `filter_column` is in `values`. Absent rows
841    /// can't collide, so append is safe where the count-based plan would merge
842    /// (spec.md#session-durable-copy).
843    pub async fn append_absent_rows(
844        &self,
845        source: &Store,
846        table: Table,
847        filter_column: &'static str,
848        values: &[String],
849    ) -> Result<usize> {
850        if values.is_empty() {
851            return Ok(0);
852        }
853        let _ = self.handle.dataset(table).await?;
854        let mut rows = 0usize;
855        for chunk in values.chunks(COPY_SESSION_IN_CHUNK) {
856            let predicate = in_predicate(filter_column, chunk);
857            rows += self.append_scanner(source, table, Some(&predicate)).await?;
858        }
859        Ok(rows)
860    }
861
862    /// Stored message PKs for the given sessions. On the full `(session_id, id)`
863    /// PK, not `id` alone: a message id is unique only within its session.
864    async fn present_message_pks(
865        &self,
866        session_id_values: &[ScalarValue],
867    ) -> Result<HashSet<(String, String)>> {
868        if session_id_values.is_empty() {
869            return Ok(HashSet::new());
870        }
871        let batch = self
872            .handle
873            .scan_batch(
874                Table::Messages,
875                Some(&Predicate::In("session_id", session_id_values.to_vec())),
876                pk_columns(Table::Messages),
877            )
878            .await?;
879        let mut set = HashSet::with_capacity(batch.num_rows());
880        for row in 0..batch.num_rows() {
881            let sid = string(&batch, "session_id", row)?.context("session_id is null")?;
882            let mid = string(&batch, "id", row)?.context("message id is null")?;
883            set.insert((sid, mid));
884        }
885        Ok(set)
886    }
887
888    /// Stored part PKs for the given sessions, on the full
889    /// `(session_id, message_id, id)` PK (see [`Self::present_message_pks`]).
890    async fn present_part_pks(
891        &self,
892        session_id_values: &[ScalarValue],
893    ) -> Result<HashSet<(String, String, String)>> {
894        if session_id_values.is_empty() {
895            return Ok(HashSet::new());
896        }
897        let batch = self
898            .handle
899            .scan_batch(
900                Table::Parts,
901                Some(&Predicate::In("session_id", session_id_values.to_vec())),
902                pk_columns(Table::Parts),
903            )
904            .await?;
905        let mut set = HashSet::with_capacity(batch.num_rows());
906        for row in 0..batch.num_rows() {
907            let sid = string(&batch, "session_id", row)?.context("session_id is null")?;
908            let mid = string(&batch, "message_id", row)?.context("message_id is null")?;
909            let pid = string(&batch, "id", row)?.context("part id is null")?;
910            set.insert((sid, mid, pid));
911        }
912        Ok(set)
913    }
914
915    /// The filter-and-append write seam, shared by ingest and grown-session
916    /// copy. Collects only the kept rows (the absent delta), so it never stages
917    /// a full copy. The terminal `append_batches` retries only a commit
918    /// conflict (not a transient post-commit fault), so a lost-ack retry cannot
919    /// re-append the same rows; an interrupted write heals on the next
920    /// re-planned re-run instead (spec.md#lance-deterministic-pk).
921    async fn append_filtered<S, K>(&self, table: Table, mut batches: S, keep: K) -> Result<usize>
922    where
923        S: Stream<Item = std::result::Result<RecordBatch, DataFusionError>> + Unpin,
924        K: Fn(&RecordBatch, usize) -> Result<bool>,
925    {
926        let mut kept_batches: Vec<RecordBatch> = Vec::new();
927        let mut kept = 0usize;
928        while let Some(batch) = batches.next().await {
929            let batch = batch?;
930            let mut mask = Vec::with_capacity(batch.num_rows());
931            for row in 0..batch.num_rows() {
932                mask.push(keep(&batch, row)?);
933            }
934            let selected = filter_record_batch(&batch, &BooleanArray::from(mask))?;
935            if selected.num_rows() > 0 {
936                kept += selected.num_rows();
937                kept_batches.push(selected);
938            }
939        }
940        self.handle.append_batches(table, kept_batches).await?;
941        Ok(kept)
942    }
943
944    /// One-shot source scan with blob bytes materialized (`AllBinary`) so the
945    /// appended rows are self-contained.
946    async fn source_scan(
947        source: &Store,
948        table: Table,
949        predicate: Option<&Predicate>,
950    ) -> Result<SendableRecordBatchStream> {
951        let mut scanner = source
952            .handle
953            .scan(
954                table,
955                ScanOpts {
956                    predicate,
957                    projection: None,
958                },
959            )
960            .await?;
961        scanner.blob_handling(lance::datatypes::BlobHandling::AllBinary);
962        Ok(scanner
963            .try_into_stream()
964            .await
965            .with_context(|| format!("failed to scan {} for copy", table.as_str()))?
966            .into())
967    }
968
969    /// `append_filtered` keep for `messages`: a row survives iff its
970    /// `(session_id, id)` PK is absent from `present`.
971    fn message_keep(
972        present: Arc<HashSet<(String, String)>>,
973    ) -> impl Fn(&RecordBatch, usize) -> Result<bool> {
974        move |batch, row| {
975            let sid = string(batch, "session_id", row)?.context("session_id is null")?;
976            let mid = string(batch, "id", row)?.context("message id is null")?;
977            Ok(!present.contains(&(sid, mid)))
978        }
979    }
980
981    /// `append_filtered` keep for `parts`, on the full
982    /// `(session_id, message_id, id)` PK.
983    fn part_keep(
984        present: Arc<HashSet<(String, String, String)>>,
985    ) -> impl Fn(&RecordBatch, usize) -> Result<bool> {
986        move |batch, row| {
987            let sid = string(batch, "session_id", row)?.context("session_id is null")?;
988            let mid = string(batch, "message_id", row)?.context("message_id is null")?;
989            let pid = string(batch, "id", row)?.context("part id is null")?;
990            Ok(!present.contains(&(sid, mid, pid)))
991        }
992    }
993
994    /// Flat write path. Per-row insert/match truth is not synthesized here -
995    /// honest outcomes come from the pre-existence scan on
996    /// [`Self::upsert_session_batch`]; the CLI sync and wire ingest paths use
997    /// that, so these helpers only need to surface write failure.
998    pub async fn upsert_sessions(&self, sessions: &[Session]) -> Result<()> {
999        if sessions.is_empty() {
1000            return Ok(());
1001        }
1002        let batches = sessions_batches(sessions)?;
1003        merge_insert_chunks(&self.handle, Table::Sessions, batches).await?;
1004        Ok(())
1005    }
1006
1007    /// Embeddings aligned to `rows` so they ride the rows' birth append. A row
1008    /// is embedded iff it has `search_text` and is absent from `present` (an
1009    /// idempotent re-sync re-embeds nothing the append would drop); otherwise,
1010    /// and whenever no embedder is attached, the slot is `None` and the batch
1011    /// writes a null vector for `pond optimize` to fill later.
1012    async fn embed_message_rows(
1013        &self,
1014        rows: &[MessageBatchRow<'_>],
1015        present: &HashSet<(String, String)>,
1016    ) -> Result<Vec<Option<Vec<f32>>>> {
1017        let mut out = vec![None; rows.len()];
1018        let Some(embedder) = &self.embedder else {
1019            return Ok(out);
1020        };
1021        let targets: Vec<usize> = rows
1022            .iter()
1023            .enumerate()
1024            .filter(|(_, row)| {
1025                row.search_text.is_some()
1026                    && !present.contains(&(
1027                        row.message.session_id().to_owned(),
1028                        row.message.id().to_owned(),
1029                    ))
1030            })
1031            .map(|(index, _)| index)
1032            .collect();
1033        if targets.is_empty() {
1034            return Ok(out);
1035        }
1036        let backend = embedder.get().await?;
1037        let texts: Vec<&str> = targets
1038            .iter()
1039            .map(|&index| rows[index].search_text.unwrap_or_default())
1040            .collect();
1041        let total = targets.len();
1042        let mut done = 0usize;
1043        if let Some(progress) = &self.ingest_embed_progress {
1044            (progress.0)(done, total);
1045        }
1046        let vectors = crate::embed::embed_passages(
1047            backend.as_ref(),
1048            &texts,
1049            crate::embed::DEFAULT_BATCH_SIZE,
1050            |batch| {
1051                done += batch;
1052                if let Some(progress) = &self.ingest_embed_progress {
1053                    (progress.0)(done, total);
1054                }
1055            },
1056        )?;
1057        for (&index, vector) in targets.iter().zip(vectors) {
1058            out[index] = Some(vector);
1059        }
1060        Ok(out)
1061    }
1062
1063    /// Batched write path used by the adapter ingest loop and by the wire
1064    /// handler's final flush. Receives N completed substreams from the
1065    /// validator and:
1066    ///
1067    ///   1. Runs the immutable-fields check (spec.md#protocol) against the stored row
1068    ///      per session, sequentially. Sessions that fail produce one Error
1069    ///      outcome and are excluded from the write batch.
1070    ///   2. Deduplicates in-batch at the substream level: when two substreams
1071    ///      in the same batch share a `session_id` (Claude Code's subagent
1072    ///      files reuse their parent's id), the first occurrence wins. The
1073    ///      second is either *merged* (same `source_agent` + `project`:
1074    ///      messages/parts append, no duplicate rows) or *rejected*
1075    ///      (different `project` - the subagent-vs-parent case). Row-level
1076    ///      duplicates that slip past here are caught downstream by Lance's
1077    ///      `SourceDedupeBehavior::FirstSeen` in `substrate::merge_insert`
1078    ///      (invariant 17): this layer's job is preserving substream merge
1079    ///      semantics, not policing the PK uniqueness Lance handles itself.
1080    ///   3. Builds one combined `RecordBatch` per table (sessions, messages,
1081    ///      parts) across every valid substream.
1082    ///   4. Commits messages + parts first, then sessions. The session row is
1083    ///      the freshness-bearing row; writing it last makes a partial
1084    ///      non-atomic flush re-ingest and heal (spec.md#session-movement-complete).
1085    ///   5. Composes per-session [`RowOutcome`]s in original substream order.
1086    async fn upsert_session_batch(
1087        &self,
1088        substreams: Vec<CompletedSubstream>,
1089    ) -> Result<(Vec<RowOutcome>, BatchCounts)> {
1090        if substreams.is_empty() {
1091            return Ok((Vec::new(), BatchCounts::default()));
1092        }
1093
1094        let mut outcomes: Vec<RowOutcome> = Vec::with_capacity(substreams.len());
1095        let mut counts = BatchCounts::default();
1096
1097        // In-batch dedup. First occurrence of each session_id wins; later
1098        // occurrences either merge or get rejected. Iteration order preserves
1099        // original substream order so outcomes index correctly.
1100        let mut merged: Vec<CompletedSubstream> = Vec::with_capacity(substreams.len());
1101        let mut by_session_id: std::collections::HashMap<String, usize> =
1102            std::collections::HashMap::with_capacity(substreams.len());
1103        for substream in substreams {
1104            if let Some(&existing_idx) = by_session_id.get(&substream.session.id) {
1105                let existing = &merged[existing_idx];
1106                if existing.session.source_agent != substream.session.source_agent
1107                    || existing.session.project != substream.session.project
1108                {
1109                    // Subagent-vs-parent class. The first occurrence's
1110                    // metadata stays authoritative; this substream is
1111                    // rejected on the same immutable-field axis as the
1112                    // storage-side check.
1113                    let reason = if existing.session.source_agent != substream.session.source_agent
1114                    {
1115                        IngestError::ImmutableField {
1116                            field: "source_agent",
1117                            session_id: substream.session.id.clone(),
1118                            stored: existing.session.source_agent.clone(),
1119                            attempted: substream.session.source_agent.clone(),
1120                        }
1121                    } else {
1122                        IngestError::ImmutableField {
1123                            field: "project",
1124                            session_id: substream.session.id.clone(),
1125                            stored: (*existing.session.project).clone(),
1126                            attempted: (*substream.session.project).clone(),
1127                        }
1128                    };
1129                    let field = match &reason {
1130                        IngestError::ImmutableField { field, .. } => Some(*field),
1131                    };
1132                    let reason_key = match field {
1133                        Some("project") => DROP_REASON_IMMUTABLE_PROJECT,
1134                        Some("source_agent") => DROP_REASON_IMMUTABLE_SOURCE_AGENT,
1135                        _ => DROP_REASON_UNCATEGORIZED,
1136                    };
1137                    outcomes.extend(error_outcomes_for_substream(
1138                        substream.session_index,
1139                        &substream.session,
1140                        &substream.messages,
1141                        reason.to_string(),
1142                        field,
1143                        reason_key,
1144                    ));
1145                    continue;
1146                }
1147                // Same session, same metadata: merge messages. Dedup message
1148                // ids defensively (within one batch, the validator's seen
1149                // sets are per-substream so cross-substream dups can happen
1150                // legally if both files re-emit the same row).
1151                let existing = &mut merged[existing_idx];
1152                let mut seen: std::collections::HashSet<String> = existing
1153                    .messages
1154                    .iter()
1155                    .map(|m| m.message.id().to_owned())
1156                    .collect();
1157                for msg in substream.messages {
1158                    if seen.insert(msg.message.id().to_owned()) {
1159                        existing.messages.push(msg);
1160                    }
1161                }
1162                continue;
1163            }
1164            by_session_id.insert(substream.session.id.clone(), merged.len());
1165            merged.push(substream);
1166        }
1167
1168        // Pre-existence sweep: one scan per table keyed on the batch's
1169        // session_ids, capped at the substream count. Replaces the prior
1170        // N-sequential `find_session` calls and gives us honest per-row
1171        // Inserted/Matched attribution downstream (spec.md#adapter-integrity-additive-sync).
1172        let session_id_values: Vec<ScalarValue> = merged
1173            .iter()
1174            .map(|substream| ScalarValue::String(substream.session.id.clone()))
1175            .collect();
1176        let existing_sessions: std::collections::HashMap<String, Session> =
1177            if session_id_values.is_empty() {
1178                std::collections::HashMap::new()
1179            } else {
1180                let batch = self
1181                    .handle
1182                    .scan_batch(
1183                        Table::Sessions,
1184                        Some(&Predicate::In("id", session_id_values.clone())),
1185                        &[],
1186                    )
1187                    .await?;
1188                let mut map = std::collections::HashMap::with_capacity(batch.num_rows());
1189                for row in 0..batch.num_rows() {
1190                    let session = session_from_batch(&batch, row)?;
1191                    map.insert(session.id.clone(), session);
1192                }
1193                map
1194            };
1195        let existing_message_pks = Arc::new(self.present_message_pks(&session_id_values).await?);
1196        let existing_part_pks = Arc::new(self.present_part_pks(&session_id_values).await?);
1197
1198        let mut writeable: Vec<CompletedSubstream> = Vec::with_capacity(merged.len());
1199        for substream in merged {
1200            if let Some(existing) = existing_sessions.get(&substream.session.id)
1201                && let Err(failure) = ensure_immutable_match(existing, &substream.session)
1202            {
1203                let field = match &failure {
1204                    IngestError::ImmutableField { field, .. } => Some(*field),
1205                };
1206                let reason_key = match field {
1207                    Some("project") => DROP_REASON_IMMUTABLE_PROJECT,
1208                    Some("source_agent") => DROP_REASON_IMMUTABLE_SOURCE_AGENT,
1209                    _ => DROP_REASON_UNCATEGORIZED,
1210                };
1211                outcomes.extend(error_outcomes_for_substream(
1212                    substream.session_index,
1213                    &substream.session,
1214                    &substream.messages,
1215                    failure.to_string(),
1216                    field,
1217                    reason_key,
1218                ));
1219                continue;
1220            }
1221            writeable.push(substream);
1222        }
1223
1224        if writeable.is_empty() {
1225            outcomes.sort_by_key(|outcome| outcome.index);
1226            return Ok((outcomes, counts));
1227        }
1228
1229        // The sessions merge is insert-only (`WhenMatched::DoNothing`), so a
1230        // row already present would be probed and left untouched while still
1231        // paying a commit - Lance 7 writes a new empty manifest version even
1232        // when every row matches. Filter to the genuinely absent rows and skip
1233        // the merge outright when none are new: the steady-state flush (grown
1234        // sessions, no new ones) then commits 2 tables, not 3. Absent rows
1235        // keep merge (not append): two writers can race the same new session
1236        // id, and merge makes the loser's row a no-op instead of a duplicate.
1237        let sessions_owned: Vec<Session> = writeable
1238            .iter()
1239            .map(|substream| &substream.session)
1240            .filter(|session| !existing_sessions.contains_key(&session.id))
1241            .cloned()
1242            .collect();
1243        // Drop only in-batch duplicates here (spec.md#adapter-integrity-dedup);
1244        // `append_filtered` drops the rows already present on the destination.
1245        let mut seen_messages: HashSet<(String, String)> = HashSet::new();
1246        let message_rows: Vec<MessageBatchRow<'_>> = writeable
1247            .iter()
1248            .flat_map(|substream| {
1249                substream.messages.iter().map(|buffered| MessageBatchRow {
1250                    message: &buffered.message,
1251                    source_agent: &substream.session.source_agent,
1252                    project: &substream.session.project,
1253                    search_text: buffered.search_text.as_deref(),
1254                })
1255            })
1256            .filter(|row| {
1257                seen_messages.insert((
1258                    row.message.session_id().to_owned(),
1259                    row.message.id().to_owned(),
1260                ))
1261            })
1262            .collect();
1263        let mut seen_parts: HashSet<(String, String, String)> = HashSet::new();
1264        let part_rows: Vec<Part> = writeable
1265            .iter()
1266            .flat_map(|substream| {
1267                substream.messages.iter().flat_map(|buffered| {
1268                    buffered
1269                        .parts
1270                        .iter()
1271                        .map(|buffered_part| buffered_part.part.clone())
1272                })
1273            })
1274            .filter(|part| {
1275                seen_parts.insert((
1276                    part.session_id.clone(),
1277                    part.message_id.clone(),
1278                    part.id.clone(),
1279                ))
1280            })
1281            .collect();
1282
1283        // Embed before the append so the vector rides the message rows' birth
1284        // commit (spec.md#session-durable-copy: one append, no extra commit).
1285        let message_vectors = self
1286            .embed_message_rows(&message_rows, &existing_message_pks)
1287            .await?;
1288
1289        let message_stream = tokio_stream::iter(
1290            messages_batches(&message_rows, &message_vectors)?
1291                .into_iter()
1292                .map(Ok::<_, DataFusionError>),
1293        );
1294        let part_stream = tokio_stream::iter(
1295            parts_batches(&part_rows)?
1296                .into_iter()
1297                .map(Ok::<_, DataFusionError>),
1298        );
1299        let (_messages_appended, _parts_appended) = tokio::try_join!(
1300            self.append_filtered(
1301                Table::Messages,
1302                message_stream,
1303                Self::message_keep(existing_message_pks.clone()),
1304            ),
1305            self.append_filtered(
1306                Table::Parts,
1307                part_stream,
1308                Self::part_keep(existing_part_pks.clone()),
1309            ),
1310        )?;
1311        if !sessions_owned.is_empty() {
1312            let session_batches = sessions_batches(&sessions_owned)?;
1313            merge_insert_chunks(&self.handle, Table::Sessions, session_batches).await?;
1314        }
1315
1316        for substream in &writeable {
1317            outcomes.extend(success_outcomes_for_substream(
1318                substream.session_index,
1319                &substream.session,
1320                &substream.messages,
1321                &existing_sessions,
1322                &existing_message_pks,
1323                &existing_part_pks,
1324                &mut counts,
1325            ));
1326        }
1327
1328        outcomes.sort_by_key(|outcome| outcome.index);
1329        Ok((outcomes, counts))
1330    }
1331
1332    pub async fn upsert_messages(
1333        &self,
1334        session: &Session,
1335        messages: &[MessageWrite<'_>],
1336    ) -> Result<()> {
1337        if messages.is_empty() {
1338            return Ok(());
1339        }
1340
1341        let rows = messages
1342            .iter()
1343            .map(|write| MessageBatchRow {
1344                message: write.message,
1345                source_agent: &session.source_agent,
1346                project: &session.project,
1347                search_text: write.search_text,
1348            })
1349            .collect::<Vec<_>>();
1350        let batches = messages_batches(&rows, &vec![None; rows.len()])?;
1351        merge_insert_chunks(&self.handle, Table::Messages, batches).await?;
1352        Ok(())
1353    }
1354
1355    pub async fn upsert_parts(&self, parts: &[Part]) -> Result<()> {
1356        if parts.is_empty() {
1357            return Ok(());
1358        }
1359        let batches = parts_batches(parts)?;
1360        merge_insert_chunks(&self.handle, Table::Parts, batches).await?;
1361        Ok(())
1362    }
1363
1364    pub async fn get_session(&self, session_id: &str) -> Result<Option<SessionWithMessages>> {
1365        let Some(session) = self.find_session(session_id).await? else {
1366            return Ok(None);
1367        };
1368        let messages = self.messages_for_session(session_id).await?;
1369        Ok(Some(SessionWithMessages { session, messages }))
1370    }
1371
1372    /// Every session id currently in the store, unsorted.
1373    pub async fn session_ids(&self) -> Result<Vec<String>> {
1374        let batch = self
1375            .handle
1376            .scan_batch(Table::Sessions, None, &["id"])
1377            .await?;
1378        let mut ids = Vec::with_capacity(batch.num_rows());
1379        for row in 0..batch.num_rows() {
1380            if let Some(id) = string(&batch, "id", row)? {
1381                ids.push(id);
1382            }
1383        }
1384        Ok(ids)
1385    }
1386
1387    pub async fn child_sessions(&self, parent_session_id: &str) -> Result<Vec<Session>> {
1388        let batch = self
1389            .handle
1390            .scan_batch(
1391                Table::Sessions,
1392                Some(&Predicate::Eq(
1393                    "parent_session_id",
1394                    parent_session_id.into(),
1395                )),
1396                &[
1397                    "id",
1398                    "parent_session_id",
1399                    "parent_message_id",
1400                    "source_agent",
1401                    "created_at",
1402                    "project",
1403                    "options",
1404                ],
1405            )
1406            .await?;
1407        let mut sessions = Vec::with_capacity(batch.num_rows());
1408        for row in 0..batch.num_rows() {
1409            sessions.push(session_from_batch(&batch, row)?);
1410        }
1411        sessions.sort_by(|left, right| left.id.cmp(&right.id));
1412        Ok(sessions)
1413    }
1414
1415    /// `session_id -> last durable message id` for the sync freshness gate.
1416    /// Scans stored message data only, never Lance version history:
1417    /// `Dataset::versions()` is remote-manifest-bound on object stores, and a
1418    /// write timestamp can exist even when a non-atomic ingest did not commit
1419    /// the messages (spec.md#session-movement-complete).
1420    ///
1421    /// Only emits a key when the session row is ALSO durable. `upsert_session_batch`
1422    /// commits messages+parts before the session row, so a partial flush can leave
1423    /// a session whose messages are stored but whose session row is not; keying on
1424    /// messages alone would report it fresh and orphan the missing row. Intersecting
1425    /// with the sessions id-set forces a re-ingest that heals it
1426    /// (spec.md#session-movement-complete).
1427    pub async fn session_last_message_ids(&self) -> Result<HashMap<String, String>> {
1428        let (session_ids, latest) = tokio::try_join!(self.collect_ids(Table::Sessions), async {
1429            let scanner = self
1430                .handle
1431                .scan(
1432                    Table::Messages,
1433                    ScanOpts::project_only(&["session_id", "id", "timestamp"]),
1434                )
1435                .await?;
1436            let mut stream = scanner.try_into_stream().await?;
1437            let mut latest: HashMap<String, (DateTime<Utc>, String)> = HashMap::new();
1438            while let Some(batch) = stream.next().await {
1439                let batch = batch?;
1440                let session_ids = batch
1441                    .column_by_name("session_id")
1442                    .context("scan projection dropped the session_id column")?
1443                    .as_any()
1444                    .downcast_ref::<StringArray>()
1445                    .context("session_id column is not Utf8")?;
1446                for row in 0..batch.num_rows() {
1447                    if session_ids.is_null(row) {
1448                        continue;
1449                    }
1450                    let session_id = session_ids.value(row);
1451                    let Some(id) = string(&batch, "id", row)? else {
1452                        continue;
1453                    };
1454                    let timestamp = datetime(&batch, "timestamp", row)?;
1455                    match latest.get_mut(session_id) {
1456                        Some((stored_ts, stored_id))
1457                            if timestamp > *stored_ts
1458                                || (timestamp == *stored_ts
1459                                    && id.as_str() > stored_id.as_str()) =>
1460                        {
1461                            *stored_ts = timestamp;
1462                            *stored_id = id;
1463                        }
1464                        None => {
1465                            latest.insert(session_id.to_owned(), (timestamp, id));
1466                        }
1467                        _ => {}
1468                    }
1469                }
1470            }
1471            Ok::<_, anyhow::Error>(latest)
1472        })?;
1473        Ok(latest
1474            .into_iter()
1475            .filter(|(session_id, _)| session_ids.contains(session_id))
1476            .map(|(session_id, (_, message_id))| (session_id, message_id))
1477            .collect())
1478    }
1479
1480    /// Whole-session view for `pond_get_session` (spec.md#protocol).
1481    /// Always the conversational view (`search_text IS NOT NULL`) with one-line
1482    /// part summaries - full part bodies are reached by `message_id` scope, not
1483    /// here. The page is the window selected by the anchors (`after_message_id`
1484    /// pages forward, `before_message_id` pages backward) or, with neither,
1485    /// `session_from` (start/end); it is bounded by `limit` and a byte budget,
1486    /// never cutting mid-message. `before_remaining`/`after_remaining` drive the
1487    /// bidirectional page markers.
1488    pub async fn session_view(
1489        &self,
1490        session_id: &str,
1491        params: SessionViewParams<'_>,
1492    ) -> Result<GetLookup<SessionPage>> {
1493        let Some(session) = self.find_session(session_id).await? else {
1494            return Ok(GetLookup::NotFound);
1495        };
1496        let mut rows: Vec<ScanRow> = self
1497            .scan_conversational_messages(session_id)
1498            .await?
1499            .into_iter()
1500            .map(|row| ScanRow {
1501                id: row.message_id,
1502                role: row.role,
1503                timestamp: row.timestamp,
1504                text: Some(row.text.into_inner()),
1505                content: None,
1506            })
1507            .collect();
1508        rows.sort_by(|a, b| a.timestamp.cmp(&b.timestamp).then_with(|| a.id.cmp(&b.id)));
1509
1510        let size = |row: &ScanRow| row.text.as_deref().map_or(0, str::len);
1511        let total = rows.len();
1512        // Append-only stream: a real anchor never vanishes, so an unknown
1513        // anchor is a stale/mistyped client cursor, not "start over". The
1514        // inclusive `at` anchor is the resolution path's: the resolved message
1515        // is known to exist, but a non-conversational one has no row here, so
1516        // it degrades to the first page rather than erroring.
1517        let at_pos = params
1518            .at_message_id
1519            .and_then(|at| rows.iter().position(|row| row.id == at));
1520        let (win_start, win_end) = match (at_pos, params.after_message_id, params.before_message_id)
1521        {
1522            (Some(pos), _, _) => {
1523                let n = page_by(&rows[pos..], params.limit, params.budget_bytes, size);
1524                (pos, pos + n)
1525            }
1526            (None, Some(after), _) if params.at_message_id.is_none() => {
1527                let pos = match rows.iter().position(|row| row.id == after) {
1528                    Some(idx) => idx + 1,
1529                    None => return Ok(GetLookup::UnknownAnchor),
1530                };
1531                let n = page_by(&rows[pos..], params.limit, params.budget_bytes, size);
1532                (pos, pos + n)
1533            }
1534            (None, None, Some(before)) if params.at_message_id.is_none() => {
1535                let pos = match rows.iter().position(|row| row.id == before) {
1536                    Some(idx) => idx,
1537                    None => return Ok(GetLookup::UnknownAnchor),
1538                };
1539                let n = page_tail(&rows[..pos], params.limit, params.budget_bytes, size);
1540                (pos - n, pos)
1541            }
1542            _ => match params.session_from {
1543                SessionFrom::Start => (0, page_by(&rows, params.limit, params.budget_bytes, size)),
1544                SessionFrom::End => {
1545                    let n = page_tail(&rows, params.limit, params.budget_bytes, size);
1546                    (total - n, total)
1547                }
1548            },
1549        };
1550        let emitted = &rows[win_start..win_end];
1551        let before_remaining = win_start;
1552        let after_remaining = total - win_end;
1553        let ids: Vec<String> = emitted.iter().map(|row| row.id.clone()).collect();
1554
1555        let mut parts_by_message = self.summary_parts_for_messages(session_id, &ids).await?;
1556        let messages = emitted
1557            .iter()
1558            .map(|row| RetrievedMessage {
1559                id: row.id.clone(),
1560                role: row.role,
1561                timestamp: row.timestamp,
1562                text: row.text.clone(),
1563                content: row.content.clone(),
1564                parts: parts_by_message
1565                    .remove(&(session_id.to_owned(), row.id.clone()))
1566                    .unwrap_or_default(),
1567            })
1568            .collect();
1569
1570        Ok(GetLookup::Found(SessionPage {
1571            session,
1572            messages,
1573            before_remaining,
1574            after_remaining,
1575        }))
1576    }
1577
1578    /// Message-scope retrieval for `pond_get_message` (spec.md#protocol):
1579    /// the target with its full parts (budget-bounded) plus `context_before`
1580    /// conversational siblings before and `context_after` after it. `NotFound`
1581    /// when no stored message carries `message_id`. Sibling parts are carried
1582    /// for summarizing; the target's parts ride `target_parts`.
1583    pub async fn message_view(
1584        &self,
1585        message_id: &str,
1586        params: MessageViewParams,
1587    ) -> Result<GetLookup<MessagePage>> {
1588        let Some(session_id) = self.session_id_for_message(message_id).await? else {
1589            return Ok(GetLookup::NotFound);
1590        };
1591        let Some(session) = self.find_session(&session_id).await? else {
1592            return Ok(GetLookup::NotFound);
1593        };
1594        let mut rows = self.scan_all_messages(&session_id).await?;
1595        // Siblings are always the conversational view: in carrier-heavy sessions
1596        // the system/tool rows would otherwise fill the whole window and push
1597        // the actual conversation out of it. The target stays regardless of its
1598        // own role - the caller asked for that message.
1599        rows.retain(|row| row.text.is_some() || row.id == message_id);
1600        rows.sort_by(|a, b| a.timestamp.cmp(&b.timestamp).then_with(|| a.id.cmp(&b.id)));
1601        let Some(target_pos) = rows.iter().position(|row| row.id == message_id) else {
1602            return Ok(GetLookup::NotFound);
1603        };
1604
1605        let start = target_pos.saturating_sub(params.context_before);
1606        let end = (target_pos + params.context_after + 1).min(rows.len());
1607        let window = &rows[start..end];
1608        let window_ids: Vec<String> = window.iter().map(|row| row.id.clone()).collect();
1609        // The target's full parts (blobs included) ride the response; siblings
1610        // are only summarized, but they share this one window scan.
1611        let mut parts_by_message = self.parts_for_messages(&session_id, &window_ids).await?;
1612
1613        let all_parts = parts_by_message
1614            .remove(&(session_id.clone(), message_id.to_owned()))
1615            .unwrap_or_default();
1616        // Target parts are budget-bounded (no per-part pagination cursor). The
1617        // 1000 cap is the page_by hard ceiling; the budget is the real bound.
1618        let part_count = page_by(&all_parts, 1000, params.budget_bytes, |part| {
1619            serde_json::to_string(part).map_or(0, |json| json.len())
1620        });
1621        let target_parts = all_parts[..part_count].to_vec();
1622        let target_parts_remaining = all_parts.len() - part_count;
1623
1624        let target_row = &rows[target_pos];
1625        let target = RetrievedMessage {
1626            id: target_row.id.clone(),
1627            role: target_row.role,
1628            timestamp: target_row.timestamp,
1629            text: target_row.text.clone(),
1630            content: target_row.content.clone(),
1631            // Target structure is carried in full by `target_parts`.
1632            parts: Vec::new(),
1633        };
1634        let siblings = window
1635            .iter()
1636            .enumerate()
1637            .filter(|(idx, _)| start + idx != target_pos)
1638            .map(|(_, row)| RetrievedMessage {
1639                id: row.id.clone(),
1640                role: row.role,
1641                timestamp: row.timestamp,
1642                text: row.text.clone(),
1643                content: row.content.clone(),
1644                parts: parts_by_message
1645                    .get(&(session_id.clone(), row.id.clone()))
1646                    .cloned()
1647                    .unwrap_or_default(),
1648            })
1649            .collect();
1650
1651        Ok(GetLookup::Found(MessagePage {
1652            session,
1653            target,
1654            target_parts,
1655            target_parts_remaining,
1656            siblings,
1657        }))
1658    }
1659
1660    async fn scan_all_messages(&self, session_id: &str) -> Result<Vec<ScanRow>> {
1661        let batch = self
1662            .handle
1663            .scan_batch(
1664                Table::Messages,
1665                Some(&Predicate::Eq("session_id", session_id.into())),
1666                &["id", "timestamp", "role", "search_text", "content"],
1667            )
1668            .await?;
1669        let mut rows = Vec::with_capacity(batch.num_rows());
1670        for row in 0..batch.num_rows() {
1671            let id = string(&batch, "id", row)?.context("message id is null")?;
1672            let role =
1673                role_from_str(&string(&batch, "role", row)?.context("message role is null")?)?;
1674            let timestamp = datetime(&batch, "timestamp", row)?;
1675            rows.push(ScanRow {
1676                id,
1677                role,
1678                timestamp,
1679                text: string(&batch, "search_text", row)?,
1680                content: string(&batch, "content", row)?,
1681            });
1682        }
1683        Ok(rows)
1684    }
1685
1686    /// Conversational scan over one session: rows ordered by
1687    /// `(timestamp, id)`, `IsNotNull("search_text")` pushed down at the
1688    /// read seam (spec.md#search-prefilter-pushdown).
1689    pub async fn scan_conversational_messages(
1690        &self,
1691        session_id: &str,
1692    ) -> Result<Vec<ConversationalRow>> {
1693        let filter = Predicate::And(vec![
1694            Predicate::Eq("session_id", session_id.into()),
1695            Predicate::IsNotNull("search_text"),
1696        ]);
1697        let batch = self
1698            .handle
1699            .scan_batch(
1700                Table::Messages,
1701                Some(&filter),
1702                &["id", "timestamp", "role", "search_text"],
1703            )
1704            .await?;
1705
1706        let mut rows = Vec::with_capacity(batch.num_rows());
1707        for row in 0..batch.num_rows() {
1708            let message_id = string(&batch, "id", row)?.context("message id is null")?;
1709            let role =
1710                role_from_str(&string(&batch, "role", row)?.context("message role is null")?)?;
1711            let timestamp = datetime(&batch, "timestamp", row)?;
1712            let text_str = string(&batch, "search_text", row)?.context(
1713                "search_text null after IsNotNull pushdown - storage invariant violated",
1714            )?;
1715            rows.push(ConversationalRow {
1716                session_id: session_id.to_owned(),
1717                message_id,
1718                role,
1719                timestamp,
1720                text: SearchText(text_str),
1721            });
1722        }
1723        rows.sort_by(|a, b| {
1724            a.timestamp
1725                .cmp(&b.timestamp)
1726                .then_with(|| a.message_id.cmp(&b.message_id))
1727        });
1728        Ok(rows)
1729    }
1730
1731    /// Locate the session id for a stored message. Cheap when only the routing
1732    /// hint is needed - callers that need the messages use `scan_all_messages`.
1733    pub async fn session_id_for_message(&self, message_id: &str) -> Result<Option<String>> {
1734        let batch = self
1735            .handle
1736            .scan_batch(
1737                Table::Messages,
1738                Some(&Predicate::Eq("id", message_id.into())),
1739                &["session_id"],
1740            )
1741            .await?;
1742        if batch.num_rows() == 0 {
1743            return Ok(None);
1744        }
1745        string(&batch, "session_id", 0)
1746    }
1747
1748    pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
1749        self.handle.row_counts().await
1750    }
1751
1752    /// The primary-key (`id`) set for `table`. Powers storage verification
1753    /// (`pond copy --verify-only` and copy's closing check).
1754    pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
1755        self.handle.collect_ids(table).await
1756    }
1757
1758    /// This store's set of composite primary keys for `table`, plus the row
1759    /// count. `rows - keys.len()` is the duplicate count - zero is the invariant
1760    /// (the append path has no row-level dedup, so a non-zero count is a write
1761    /// anomaly the copy verify reports rather than calling "synced"), and the
1762    /// key set drives the verify's completeness membership. One scan over only
1763    /// the PK columns yields both, holding a single composite-PK set per table.
1764    pub async fn composite_pk_index(&self, table: Table) -> Result<(HashSet<Vec<String>>, usize)> {
1765        let pk = pk_columns(table);
1766        let scanner = self.handle.scan(table, ScanOpts::project_only(pk)).await?;
1767        let mut stream = scanner.try_into_stream().await?;
1768        let mut keys: HashSet<Vec<String>> = HashSet::new();
1769        let mut rows = 0usize;
1770        while let Some(batch) = stream.next().await {
1771            let batch = batch?;
1772            for row in 0..batch.num_rows() {
1773                rows += 1;
1774                keys.insert(composite_key(&batch, pk, row)?);
1775            }
1776        }
1777        Ok((keys, rows))
1778    }
1779
1780    /// Stream `table`'s composite primary keys and return `(rows_scanned, rows
1781    /// whose key is absent from `present`)`. Composite-keyed, not bare `id`: a
1782    /// message id replayed into a new session by a fork/compaction is matched
1783    /// per session, so a wholly-absent replayed session whose ids collide with
1784    /// present ones is counted missing - a bare-`id` check would false-negative
1785    /// it as "present". Streams the scanned side, holding only `present`.
1786    pub async fn composite_pk_diff_against(
1787        &self,
1788        table: Table,
1789        present: &HashSet<Vec<String>>,
1790    ) -> Result<(usize, usize)> {
1791        let pk = pk_columns(table);
1792        let scanner = self.handle.scan(table, ScanOpts::project_only(pk)).await?;
1793        let mut stream = scanner.try_into_stream().await?;
1794        let (mut rows, mut absent) = (0usize, 0usize);
1795        while let Some(batch) = stream.next().await {
1796            let batch = batch?;
1797            for row in 0..batch.num_rows() {
1798                rows += 1;
1799                if !present.contains(&composite_key(&batch, pk, row)?) {
1800                    absent += 1;
1801                }
1802            }
1803        }
1804        Ok((rows, absent))
1805    }
1806
1807    /// A point-in-time `Arc<Dataset>` for `table`, for registering as a
1808    /// DataFusion `LanceTableProvider` in `pond_sql`. Goes through the
1809    /// handle's freshness gate, so each query sees a current snapshot.
1810    pub async fn dataset(&self, table: Table) -> Result<Arc<Dataset>> {
1811        Ok(Arc::new(self.handle.dataset(table).await?))
1812    }
1813
1814    /// Page the heavy search indices in from storage so the first user query
1815    /// after process start never eats the cold S3 index load (spec.md#search).
1816    /// Vector via `prewarm_index` (loads the IVF_SQ partition storage). FTS is
1817    /// warmed with one synthetic query rather than Lance's full FTS
1818    /// `prewarm_index`, which would resident-set the whole inverted index and
1819    /// blow the server RAM budget - so we settle the term dictionary + a hot
1820    /// token and let real queries page their own postings. Best effort: a
1821    /// missing index (IVF_SQ below activation, or no FTS yet on an empty store)
1822    /// is logged, not fatal.
1823    pub async fn prewarm(&self, cache_dir: &Path) -> Result<()> {
1824        let messages = self.dataset(Table::Messages).await?;
1825        if let Err(error) = messages.prewarm_index(MESSAGES_VECTOR_INDEX).await {
1826            tracing::debug!(%error, "vector index prewarm skipped");
1827        }
1828        // Best-effort: on failure `rowmap` stays empty and the arms fall back to
1829        // the data-take path, so search still works (slower on a remote store).
1830        if let Err(error) = self.ensure_rowmap(cache_dir).await {
1831            tracing::warn!(%error, "rowmap build skipped; arms fall back to data-take resolution");
1832        }
1833        // Warm the FTS posting lists; the rowmap build above touched only the
1834        // data columns.
1835        if let Err(error) = self
1836            .fts_search("pond", 1, &Predicate::And(Vec::new()))
1837            .await
1838        {
1839            tracing::debug!(%error, "fts index prewarm skipped");
1840        }
1841        self.prune_index_cache(cache_dir).await;
1842        Ok(())
1843    }
1844
1845    /// Reclaim disk-index-cache entries for index versions the store has moved
1846    /// past (see `Handle::prune_index_cache`). Best-effort.
1847    pub async fn prune_index_cache(&self, cache_dir: &Path) {
1848        self.handle.prune_index_cache(cache_dir).await;
1849    }
1850
1851    /// Stable filesystem-safe cache key: same store URL -> same key, so sibling
1852    /// pond processes share one map file and distinct stores never collide.
1853    fn store_key(&self) -> String {
1854        crate::substrate::store_key(self.handle.location())
1855    }
1856
1857    /// Max delta segments before the chain is compacted into a fresh base.
1858    const MAX_ROWMAP_DELTAS: usize = 16;
1859
1860    /// Columns the resident meta map is built from. The full scan and the delta
1861    /// scan MUST project the same set in the same order - both feed
1862    /// [`row_meta_entry`], so a column added to one only would silently corrupt
1863    /// delta hydration.
1864    const ROW_META_COLUMNS: [&str; 7] = [
1865        "session_id",
1866        "id",
1867        "role",
1868        "project",
1869        "source_agent",
1870        "timestamp",
1871        "search_text",
1872    ];
1873
1874    /// Install the resident meta map covering the current `messages` version.
1875    /// Idempotent - a chain already at that version is kept. On a version bump
1876    /// it layers a delta segment (scanning only the new fragments), compacts the
1877    /// deltas locally once they pile up, and full-rebuilds the base only on a
1878    /// store compaction - all under a build `flock` so N local processes don't
1879    /// rescan the store at once.
1880    pub async fn ensure_rowmap(&self, cache_dir: &Path) -> Result<()> {
1881        let version = self.messages_version().await?;
1882        if let Some(current) = self.rowmap.load_full()
1883            && current.version() == version
1884        {
1885            return Ok(());
1886        }
1887        std::fs::create_dir_all(cache_dir)
1888            .with_context(|| format!("create cache dir {}", cache_dir.display()))?;
1889        let store_key = self.store_key();
1890
1891        // A sibling may already have published a chain at this version; install
1892        // it without rebuilding.
1893        if let Some(chain) = discover_chain(cache_dir, &store_key)
1894            && chain.version() == version
1895            && let Ok(set) = RowMetaSet::open(&chain)
1896        {
1897            self.rowmap.store(Some(Arc::new(set)));
1898            Self::sweep_stale_rowmaps(cache_dir, &store_key, chain.base_version);
1899            return Ok(());
1900        }
1901        if let Some(set) = self
1902            .extend_rowmap_coordinated(cache_dir, &store_key, version)
1903            .await?
1904        {
1905            self.rowmap.store(Some(Arc::new(set)));
1906        }
1907        Ok(())
1908    }
1909
1910    /// Open the newest locally cached rowmap chain regardless of the store's
1911    /// current version, without installing it. Read-only estimate seam for
1912    /// `pond status`: the chain is as-of this host's last sync - exactly the
1913    /// baseline "pending since then" wants - and a version-matched load would
1914    /// cost a remote manifest read. Never assigned to `self.rowmap`: searches
1915    /// must not hydrate from a possibly-stale map.
1916    pub fn open_cached_rowmap(&self, cache_dir: &Path) -> Option<Arc<RowMetaSet>> {
1917        let chain = discover_chain(cache_dir, &self.store_key())?;
1918        RowMetaSet::open(&chain).ok().map(Arc::new)
1919    }
1920
1921    /// Install an already-published rowmap chain for the current version if a
1922    /// sibling built one, without building it (no full scan, no build flock).
1923    /// For one-shot read commands (`pond search`): a warm sibling makes
1924    /// hydration resident; with no chain, search falls back to take_rows for
1925    /// that single invocation.
1926    pub async fn load_rowmap_if_present(&self, cache_dir: &Path) -> Result<()> {
1927        let version = self.messages_version().await?;
1928        if let Some(current) = self.rowmap.load_full()
1929            && current.version() == version
1930        {
1931            return Ok(());
1932        }
1933        if let Some(chain) = discover_chain(cache_dir, &self.store_key())
1934            && chain.version() == version
1935            && let Ok(set) = RowMetaSet::open(&chain)
1936        {
1937            self.rowmap.store(Some(Arc::new(set)));
1938        }
1939        Ok(())
1940    }
1941
1942    /// Extend the chain to `version` under the build `flock` (spec: lock the
1943    /// build only; atomic rename already prevents corruption). `None` when
1944    /// another local process holds the lock - this caller keeps its current map
1945    /// (or the take_rows fallback) until a later refresh opens what the winner
1946    /// published.
1947    async fn extend_rowmap_coordinated(
1948        &self,
1949        cache_dir: &Path,
1950        store_key: &str,
1951        version: u64,
1952    ) -> Result<Option<RowMetaSet>> {
1953        let lock_path = cache_dir.join(format!("rowmetamap-{store_key}.lock"));
1954        let lock = std::fs::File::create(&lock_path)
1955            .with_context(|| format!("create rowmap build lock {}", lock_path.display()))?;
1956        match lock.try_lock() {
1957            Ok(()) => {}
1958            Err(std::fs::TryLockError::WouldBlock) => return Ok(None),
1959            Err(std::fs::TryLockError::Error(error)) => {
1960                return Err(error).context("lock rowmap build");
1961            }
1962        }
1963
1964        // Re-check after acquiring: a sibling may have published `version`. An
1965        // open failure here (older MAGIC after an upgrade, or corruption) falls
1966        // through to the purge+rebuild below rather than erroring.
1967        if let Some(chain) = discover_chain(cache_dir, store_key)
1968            && chain.version() == version
1969            && let Ok(set) = RowMetaSet::open(&chain)
1970        {
1971            return Ok(Some(set));
1972        }
1973
1974        // Holding the lock makes us the only builder, so every build temp is a
1975        // dead orphan from a crashed build - clear them before writing ours.
1976        Self::sweep_orphan_temps(cache_dir, store_key);
1977
1978        // Validate any existing chain opens; an unreadable segment (an older
1979        // MAGIC after a pond upgrade, or a corrupt file) is purged so the build
1980        // below is a clean full rebuild instead of erroring forever or appending
1981        // a fresh delta onto an unreadable base. The opened set also feeds the
1982        // delta its high-water mark and row count (cheap mmap reads).
1983        let chain = discover_chain(cache_dir, store_key);
1984        let existing = match &chain {
1985            Some(paths) => match RowMetaSet::open(paths) {
1986                Ok(set) => Some((paths, set)),
1987                Err(error) => {
1988                    tracing::warn!(%error, store = store_key, "rowmap unreadable; purging and rebuilding");
1989                    Self::purge_rowmaps(cache_dir, store_key);
1990                    None
1991                }
1992            },
1993            None => None,
1994        };
1995        // A row-id-keyed append delta (None on a reclaimed base or net deletion)
1996        // decides the path.
1997        let delta = match &existing {
1998            Some((_, set)) => {
1999                self.collect_row_metas_delta(
2000                    set.version(),
2001                    set.max_row_id().unwrap_or(0),
2002                    set.len(),
2003                )
2004                .await?
2005            }
2006            None => None,
2007        };
2008
2009        let base_version = match (&existing, delta) {
2010            // Append with room: layer a new delta segment.
2011            (Some((paths, _)), Some(entries)) if paths.deltas.len() < Self::MAX_ROWMAP_DELTAS => {
2012                let path = RowMetaMap::delta_path(cache_dir, store_key, version);
2013                RowMetaMap::build(&path, version, entries)?;
2014                paths.base_version
2015            }
2016            // Append but the deltas are full: compact the existing segments
2017            // (read locally from their mmaps) plus this delta into a fresh base -
2018            // no full store re-read.
2019            (Some((_, set)), Some(entries)) => {
2020                let mut merged = set.merged_entries();
2021                merged.extend(entries);
2022                let path = RowMetaMap::path_for(cache_dir, store_key, version);
2023                RowMetaMap::build(&path, version, merged)?;
2024                version
2025            }
2026            // No chain, or a reclaimed base / deletion since it: full scan -> base.
2027            _ => {
2028                let entries = self.collect_row_metas().await?;
2029                let path = RowMetaMap::path_for(cache_dir, store_key, version);
2030                RowMetaMap::build(&path, version, entries)?;
2031                version
2032            }
2033        };
2034
2035        let chain =
2036            discover_chain(cache_dir, store_key).context("rowmap chain missing after build")?;
2037        let set = RowMetaSet::open(&chain)?;
2038        Self::sweep_stale_rowmaps(cache_dir, store_key, base_version);
2039        Ok(Some(set))
2040    }
2041
2042    /// Remove this store's segment files (`-v{V}` bases, `-d{V}` deltas) for
2043    /// versions strictly older than `keep` (best-effort). A newer file belongs
2044    /// to a sibling that advanced past us; unlinking a superseded file is safe
2045    /// even if a sibling has it mapped - Unix keeps the inode alive until unmap.
2046    fn sweep_stale_rowmaps(cache_dir: &Path, store_key: &str, keep: u64) {
2047        let prefix = format!("rowmetamap-{store_key}-");
2048        let Ok(entries) = std::fs::read_dir(cache_dir) else {
2049            return;
2050        };
2051        for entry in entries.flatten() {
2052            let name = entry.file_name();
2053            let Some(rest) = name
2054                .to_str()
2055                .and_then(|name| name.strip_prefix(&prefix))
2056                .and_then(|rest| rest.strip_suffix(".rmm"))
2057            else {
2058                continue;
2059            };
2060            let version = rest
2061                .strip_prefix('v')
2062                .or_else(|| rest.strip_prefix('d'))
2063                .and_then(|digits| digits.parse::<u64>().ok());
2064            if let Some(version) = version
2065                && version < keep
2066            {
2067                let _ = std::fs::remove_file(entry.path());
2068            }
2069        }
2070    }
2071
2072    /// Remove every segment file (`-v{V}` / `-d{V}`) for this store regardless of
2073    /// version - used when a discovered chain is unreadable (older MAGIC after an
2074    /// upgrade, or corruption) so the next build starts clean. Sound under the
2075    /// build lock; POSIX keeps any inode a sibling still has mapped alive.
2076    fn purge_rowmaps(cache_dir: &Path, store_key: &str) {
2077        let prefix = format!("rowmetamap-{store_key}-");
2078        let Ok(entries) = std::fs::read_dir(cache_dir) else {
2079            return;
2080        };
2081        for entry in entries.flatten() {
2082            if let Some(name) = entry.file_name().to_str()
2083                && name.starts_with(&prefix)
2084                && name.ends_with(".rmm")
2085            {
2086                let _ = std::fs::remove_file(entry.path());
2087            }
2088        }
2089    }
2090
2091    /// Remove abandoned build temp files (`*.tmp-*`) for this store. Best-effort,
2092    /// and only sound under the build lock - the holder is the sole builder, so
2093    /// any temp present is a crashed-build orphan, not a live write.
2094    fn sweep_orphan_temps(cache_dir: &Path, store_key: &str) {
2095        let prefix = format!("rowmetamap-{store_key}-");
2096        let Ok(entries) = std::fs::read_dir(cache_dir) else {
2097            return;
2098        };
2099        for entry in entries.flatten() {
2100            let name = entry.file_name();
2101            let Some(name) = name.to_str() else { continue };
2102            if name.starts_with(&prefix) && name.contains(".tmp-") {
2103                let _ = std::fs::remove_file(entry.path());
2104            }
2105        }
2106    }
2107
2108    #[cfg(test)]
2109    pub(crate) fn rowmap_delta_count(&self) -> Option<usize> {
2110        self.rowmap.load_full().map(|set| set.delta_count())
2111    }
2112
2113    /// The currently-installed resident meta map, if any. `pond sync` reads it
2114    /// (via [`RowmapOracle`]) as the freshness oracle (max timestamp per
2115    /// session); `None` falls back to re-reading every source.
2116    pub fn rowmap_snapshot(&self) -> Option<Arc<RowMetaSet>> {
2117        self.rowmap.load_full()
2118    }
2119
2120    /// Resolve index-only `(row_id, score)` hits to keys via the map; row ids the
2121    /// map lacks (appended since build) fall back to one `take_rows` batch. The
2122    /// caller re-sorts, so the misses appended at the end carry no order meaning.
2123    async fn resolve_rowid_hits(
2124        &self,
2125        map: &RowMetaSet,
2126        hits: Vec<(u64, f32)>,
2127    ) -> Result<Vec<SearchHit>> {
2128        let mut resolved = Vec::with_capacity(hits.len());
2129        let mut misses: Vec<(u64, f32)> = Vec::new();
2130        for (rowid, score) in hits {
2131            match map.lookup(rowid) {
2132                Some((session_id, message_id)) => resolved.push(SearchHit {
2133                    rowid: Some(rowid),
2134                    key: MessageKey {
2135                        session_id: session_id.to_owned(),
2136                        message_id: message_id.to_owned(),
2137                    },
2138                    score,
2139                }),
2140                None => misses.push((rowid, score)),
2141            }
2142        }
2143        // A miss still knows its rowid; carry it so hydration can take_rows it
2144        // alongside the hits the map resolved.
2145        if !misses.is_empty() {
2146            let rowids: Vec<u64> = misses.iter().map(|(rowid, _)| *rowid).collect();
2147            let keys = self.message_keys_by_rowids(&rowids).await?;
2148            for ((rowid, score), key) in misses.into_iter().zip(keys) {
2149                resolved.push(SearchHit {
2150                    rowid: Some(rowid),
2151                    key,
2152                    score,
2153                });
2154            }
2155        }
2156        Ok(resolved)
2157    }
2158
2159    /// Resolve stable row ids to `(session_id, id)` via `take_rows`, which
2160    /// returns rows in `rowids` order - the caller's `zip` relies on that.
2161    async fn message_keys_by_rowids(&self, rowids: &[u64]) -> Result<Vec<MessageKey>> {
2162        let dataset = self.handle.dataset(Table::Messages).await?;
2163        let projection = ProjectionRequest::from_columns(["session_id", "id"], dataset.schema());
2164        let batch = dataset.take_rows(rowids, projection).await?;
2165        let mut keys = Vec::with_capacity(batch.num_rows());
2166        for row in 0..batch.num_rows() {
2167            keys.push(MessageKey {
2168                session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2169                message_id: string(&batch, "id", row)?.context("fts hit id is null")?,
2170            });
2171        }
2172        Ok(keys)
2173    }
2174
2175    /// Write a `pond_sql` export artifact.
2176    pub async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
2177        self.handle.export_write(name, bytes).await
2178    }
2179
2180    /// Read a `pond_sql` export artifact back.
2181    pub async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
2182        self.handle.export_read(name).await
2183    }
2184
2185    /// Local filesystem path of an export artifact on `file://` installs.
2186    pub fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2187        self.handle.export_local_path(name)
2188    }
2189
2190    /// Distinct adapter names present in the corpus, sorted. Scans only the
2191    /// `source_agent` column of the small `sessions` table, so `pond status`
2192    /// gets its adapter count without touching the 2M-row `messages` table.
2193    /// `include_subagents=false` drops `source_agent` values containing `/`
2194    /// (e.g. `claude-code/general-purpose`).
2195    pub async fn adapter_names(&self, include_subagents: bool) -> Result<Vec<String>> {
2196        let scanner = self
2197            .handle
2198            .scan(Table::Sessions, ScanOpts::project_only(&["source_agent"]))
2199            .await?;
2200        let mut stream = scanner.try_into_stream().await?;
2201        let mut names: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
2202        while let Some(batch) = stream.next().await {
2203            let batch = batch?;
2204            for row in 0..batch.num_rows() {
2205                let agent = string(&batch, "source_agent", row)?.unwrap_or_default();
2206                if !include_subagents && agent.contains('/') {
2207                    continue;
2208                }
2209                names.insert(agent);
2210            }
2211        }
2212        Ok(names.into_iter().collect())
2213    }
2214
2215    /// Per-ingest-host activity: distinct sessions and newest message
2216    /// timestamp per `options.pond.ingest.host.hostname` stamp
2217    /// (spec.md#model-pond-options). The stamp is written per message at
2218    /// ingest, so this scans the `messages` table - acceptable because
2219    /// `pond status --hosts` is an explicit opt-in view, never the default
2220    /// status path. Rows without the stamp (wire-ingested, or predating it)
2221    /// group under the `None` host; a session synced from several hosts
2222    /// counts once under each. Answers "is every machine feeding the pond"
2223    /// on a shared store.
2224    pub async fn ingest_host_activity(&self) -> Result<Vec<HostActivity>> {
2225        let scanner = self
2226            .handle
2227            .scan(
2228                Table::Messages,
2229                ScanOpts::project_only(&["session_id", "timestamp", "options"]),
2230            )
2231            .await?;
2232        let mut stream = scanner.try_into_stream().await?;
2233        let mut hosts: BTreeMap<Option<String>, (HashSet<String>, DateTime<Utc>)> = BTreeMap::new();
2234        while let Some(batch) = stream.next().await {
2235            let batch = batch?;
2236            for row in 0..batch.num_rows() {
2237                let session_id =
2238                    string(&batch, "session_id", row)?.context("session_id is null")?;
2239                let timestamp = datetime(&batch, "timestamp", row)?;
2240                let hostname = json_column(&batch, "options", row)?
2241                    .and_then(|bytes| json_parse::<serde_json::Value>(&bytes).ok())
2242                    .as_ref()
2243                    .and_then(|options| options.pointer("/pond/ingest/host/hostname"))
2244                    .and_then(serde_json::Value::as_str)
2245                    .map(str::to_owned);
2246                let entry = hosts
2247                    .entry(hostname)
2248                    .or_insert_with(|| (HashSet::new(), timestamp));
2249                entry.0.insert(session_id);
2250                entry.1 = entry.1.max(timestamp);
2251            }
2252        }
2253        Ok(hosts
2254            .into_iter()
2255            .map(|(hostname, (sessions, last_message_at))| HostActivity {
2256                hostname,
2257                sessions: sessions.len(),
2258                last_message_at,
2259            })
2260            .collect())
2261    }
2262
2263    /// Write a batch of embeddings into `messages`: set `vector` and
2264    /// `embedding_model` on each row by `(session_id, id)`
2265    /// (spec.md#session-embed-from-canonical). The column update goes through the
2266    /// write seam and lands as a new manifest version (`append-only`).
2267    pub async fn write_embeddings(&self, rows: &[EmbeddedMessage]) -> Result<()> {
2268        if rows.is_empty() {
2269            return Ok(());
2270        }
2271        let batch = embedding_update_batch(rows)?;
2272        self.handle
2273            .merge_update(Table::Messages, batch, rows.len())
2274            .await?;
2275        Ok(())
2276    }
2277
2278    /// Stream the backlog of messages needing embedding: rows with `search_text`
2279    /// set whose `vector` is null (spec.md#session-embed-from-canonical).
2280    pub fn pending_embedding_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2281        try_stream! {
2282            // Filter on `embedding_model IS NULL`, not `vector IS NULL`: the two
2283            // are co-set (write_embeddings sets both, spec.md#session-embed-from-canonical),
2284            // but evaluating the predicate over the narrow model-id column reads
2285            // ~50x fewer bytes than scanning the Float16 vector column - the
2286            // difference between a whole-table vector decode and a cheap scan.
2287            let filter = Predicate::And(vec![
2288                Predicate::IsNull("embedding_model"),
2289                Predicate::IsNotNull("search_text"),
2290            ]);
2291            let projection: &[&str] = &["session_id", "id", "search_text"];
2292            let scanner = self
2293                .handle
2294                .scan(
2295                    Table::Messages,
2296                    ScanOpts::with_predicate_and_projection(&filter, projection),
2297                )
2298                .await?;
2299            let mut batches = scanner
2300                .try_into_stream()
2301                .await
2302                .context("failed to open messages stream")?;
2303            while let Some(batch) = batches.next().await {
2304                let batch = batch?;
2305                for row in 0..batch.num_rows() {
2306                    yield PendingMessage {
2307                        session_id: string(&batch, "session_id", row)?
2308                            .context("session_id is null")?,
2309                        id: string(&batch, "id", row)?.context("message id is null")?,
2310                        search_text: string(&batch, "search_text", row)?
2311                            .context("search_text is null")?,
2312                    };
2313                }
2314            }
2315        }
2316    }
2317
2318    /// Stream messages that are either never embedded or stale under the
2319    /// current model. `pond optimize --force-embed` feeds this to the same unconditional
2320    /// merge_update as the normal backlog; the filter makes that semantically
2321    /// equivalent to the conditional update in spec.md#session-embed-from-canonical.
2322    pub fn pending_or_stale_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2323        try_stream! {
2324            // `embedding_model IS NULL` (co-set with `vector IS NULL`, but a ~50x
2325            // narrower column read) for the never-embedded rows, OR a model
2326            // mismatch for the stale ones - both decided off the model-id column.
2327            let filter = Predicate::And(vec![
2328                Predicate::IsNotNull("search_text"),
2329                Predicate::Or(vec![
2330                    Predicate::IsNull("embedding_model"),
2331                    Predicate::Ne("embedding_model", embed::model_id().into()),
2332                ]),
2333            ]);
2334            let projection: &[&str] = &["session_id", "id", "search_text"];
2335            let scanner = self
2336                .handle
2337                .scan(
2338                    Table::Messages,
2339                    ScanOpts::with_predicate_and_projection(&filter, projection),
2340                )
2341                .await?;
2342            let mut batches = scanner
2343                .try_into_stream()
2344                .await
2345                .context("failed to open pending-or-stale messages stream")?;
2346            while let Some(batch) = batches.next().await {
2347                let batch = batch?;
2348                for row in 0..batch.num_rows() {
2349                    yield PendingMessage {
2350                        session_id: string(&batch, "session_id", row)?
2351                            .context("session_id is null")?,
2352                        id: string(&batch, "id", row)?.context("message id is null")?,
2353                        search_text: string(&batch, "search_text", row)?
2354                            .context("search_text is null")?,
2355                    };
2356                }
2357            }
2358        }
2359    }
2360
2361    /// BM25 full-text retriever over `messages.search_text`. With the row meta map
2362    /// loaded the scan is index-only (no data columns -> no `TakeExec`, no
2363    /// scattered GETs) and hits resolve through the map; otherwise it falls back
2364    /// to `fts_search_keys` so search works before prewarm.
2365    pub async fn fts_search(
2366        &self,
2367        query: &str,
2368        limit: usize,
2369        filter: &Predicate,
2370    ) -> Result<Vec<SearchHit>> {
2371        let mut hits = if let Some(map) = self.rowmap.load_full() {
2372            let rowid_hits = self.fts_search_rowids(query, limit, filter).await?;
2373            self.resolve_rowid_hits(&map, rowid_hits).await?
2374        } else {
2375            self.fts_search_keys(query, limit, filter).await?
2376        };
2377        // Stable secondary sort: Lance returns tied-BM25-score hits in fragment
2378        // order, which varies between runs and across calls with different pool
2379        // sizes. Without an explicit tiebreak the downstream session grouping and
2380        // rank for a tied target can flip session-to-session, making results
2381        // nondeterministic. Sort by `score desc`, then `(session_id, message_id)` asc.
2382        hits.sort_by(|left, right| {
2383            right
2384                .score
2385                .partial_cmp(&left.score)
2386                .unwrap_or(std::cmp::Ordering::Equal)
2387                .then_with(|| left.key.session_id.cmp(&right.key.session_id))
2388                .then_with(|| left.key.message_id.cmp(&right.key.message_id))
2389        });
2390        Ok(hits)
2391    }
2392
2393    /// Shared FTS-scan setup: scope filter, the `search_text` full-text query,
2394    /// `fast_search` only when the index has no unindexed tail (else Lance
2395    /// index-probes + flat-scans the tail), and `limit`. Callers set the projection.
2396    async fn fts_scanner(
2397        &self,
2398        query: &str,
2399        limit: usize,
2400        filter: &Predicate,
2401    ) -> Result<lance::dataset::scanner::Scanner> {
2402        let mut scanner = self.handle.scanner(Table::Messages, Some(filter)).await?;
2403        scanner.full_text_search(
2404            FullTextSearchQuery::new(query.to_owned()).with_column("search_text".to_owned())?,
2405        )?;
2406        if self
2407            .handle
2408            .messages_fast_search_ready(MESSAGES_FTS_INDEX)
2409            .await?
2410        {
2411            scanner.fast_search();
2412        }
2413        // Lance ships an autoprojection that silently appends `_score` to FTS
2414        // output when the projection omits it. That behavior is going away;
2415        // we opt into the future explicit-projection contract here so the
2416        // scanner stops emitting a per-call deprecation warning, and each caller
2417        // lists `_score` in its own projection.
2418        scanner.disable_scoring_autoprojection();
2419        scanner.limit(Some(i64::try_from(limit).unwrap_or(i64::MAX)), None)?;
2420        Ok(scanner)
2421    }
2422
2423    /// No-map FTS fallback: project the key columns plus `_score` directly,
2424    /// taking the `TakeExec` cost. Unsorted; `fts_search` applies the sort.
2425    async fn fts_search_keys(
2426        &self,
2427        query: &str,
2428        limit: usize,
2429        filter: &Predicate,
2430    ) -> Result<Vec<SearchHit>> {
2431        let mut scanner = self.fts_scanner(query, limit, filter).await?;
2432        scanner.project(&["session_id", "id", "_score"])?;
2433        let batch = scanner.try_into_batch().await?;
2434        let mut hits = Vec::with_capacity(batch.num_rows());
2435        for row in 0..batch.num_rows() {
2436            let key = MessageKey {
2437                session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2438                message_id: string(&batch, "id", row)?.context("fts hit id is null")?,
2439            };
2440            hits.push(SearchHit {
2441                rowid: None,
2442                key,
2443                score: float32(&batch, "_score", row)?,
2444            });
2445        }
2446        Ok(hits)
2447    }
2448
2449    /// Current `messages` dataset version - the key a `RowMetaMap` is built
2450    /// against (pond's stable row ids keep a built map valid until this advances).
2451    pub async fn messages_version(&self) -> Result<u64> {
2452        Ok(self
2453            .handle
2454            .dataset(Table::Messages)
2455            .await?
2456            .version()
2457            .version)
2458    }
2459
2460    /// Scan the hydration columns with row ids into a `Vec`, the input to
2461    /// `RowMetaMap::build`. One large sequential scan (few big reads), unlike the
2462    /// scattered per-hit take it replaces; `search_text` dominates the bytes.
2463    pub async fn collect_row_metas(&self) -> Result<Vec<RowMetaEntry>> {
2464        let mut scanner = self.handle.scanner(Table::Messages, None).await?;
2465        scanner.with_row_id();
2466        scanner.project(&Self::ROW_META_COLUMNS)?;
2467        let mut stream = scanner.try_into_stream().await?;
2468        let mut out = Vec::new();
2469        while let Some(batch) = stream.next().await {
2470            let batch = batch?;
2471            let rowids = uint64(&batch, "_rowid")?;
2472            for row in 0..batch.num_rows() {
2473                out.push(row_meta_entry(&batch, rowids.value(row), row)?);
2474            }
2475        }
2476        Ok(out)
2477    }
2478
2479    /// Row metas for the rows appended since the base segment - the input to a
2480    /// delta layered on a base whose high-water mark is `base_max_row_id` and
2481    /// which covers `base_row_count` rows. `None` (caller rebuilds the base from
2482    /// a full scan) when the chain can't be cheaply extended:
2483    /// - `base_version`'s manifest was reclaimed by the cleanup retention window
2484    ///   (spec.md#concurrency), so the version no longer resolves; or
2485    /// - the live row count dropped below the base: rows were deleted, and a
2486    ///   pure append can't remove the base's now-stale entries.
2487    ///
2488    /// Stable row ids (`enable_stable_row_ids`) make this an append: embedding's
2489    /// `merge_update` and compaction rewrite message fragments but preserve
2490    /// row_ids and never touch a ROW_META column, so existing base entries stay
2491    /// valid under that churn. Only genuine appends carry `row_id >
2492    /// base_max_row_id`; emitting just those keeps the delta disjoint from the
2493    /// base, which the per-segment count sums depend on.
2494    async fn collect_row_metas_delta(
2495        &self,
2496        base_version: u64,
2497        base_max_row_id: u64,
2498        base_row_count: usize,
2499    ) -> Result<Option<Vec<RowMetaEntry>>> {
2500        let dataset = self.handle.dataset(Table::Messages).await?;
2501        let Ok(old) = dataset.checkout_version(base_version).await else {
2502            return Ok(None);
2503        };
2504        if dataset.count_rows(None).await? < base_row_count {
2505            return Ok(None);
2506        }
2507        // Restrict the scan to fragments added since the base (recent churn -
2508        // not the untouched bulk). Rewritten/compacted fragments carry only
2509        // existing row_ids (<= base_max_row_id) and are filtered out row-wise;
2510        // genuine appends carry higher ids and are kept.
2511        let old_ids: HashSet<u64> = old.get_fragments().iter().map(|f| f.id() as u64).collect();
2512        let added: Vec<_> = dataset
2513            .get_fragments()
2514            .iter()
2515            .filter(|fragment| !old_ids.contains(&(fragment.id() as u64)))
2516            .map(|fragment| fragment.metadata().clone())
2517            .collect();
2518        if added.is_empty() {
2519            return Ok(Some(Vec::new()));
2520        }
2521        let mut scanner = dataset.scan();
2522        scanner.with_fragments(added);
2523        scanner.with_row_id();
2524        scanner.project(&Self::ROW_META_COLUMNS)?;
2525        let mut stream = scanner.try_into_stream().await?;
2526        let mut out = Vec::new();
2527        while let Some(batch) = stream.next().await {
2528            let batch = batch?;
2529            let rowids = uint64(&batch, "_rowid")?;
2530            for row in 0..batch.num_rows() {
2531                let row_id = rowids.value(row);
2532                if row_id > base_max_row_id {
2533                    out.push(row_meta_entry(&batch, row_id, row)?);
2534                }
2535            }
2536        }
2537        Ok(Some(out))
2538    }
2539
2540    /// Index-only FTS retriever: `_rowid` + `_score` only, so Lance inserts no
2541    /// `TakeExec` and issues no scattered GETs. `fts_search` resolves the row ids.
2542    async fn fts_search_rowids(
2543        &self,
2544        query: &str,
2545        limit: usize,
2546        filter: &Predicate,
2547    ) -> Result<Vec<(u64, f32)>> {
2548        let mut scanner = self.fts_scanner(query, limit, filter).await?;
2549        scanner.with_row_id();
2550        scanner.project(&["_score"])?;
2551        let batch = scanner.try_into_batch().await?;
2552        let rowids = uint64(&batch, "_rowid")?;
2553        let mut hits = Vec::with_capacity(batch.num_rows());
2554        for row in 0..batch.num_rows() {
2555            hits.push((rowids.value(row), float32(&batch, "_score", row)?));
2556        }
2557        Ok(hits)
2558    }
2559
2560    /// Count of searchable messages (non-null `search_text`) inside the
2561    /// caller's filter scope - the universe a search actually ran over.
2562    /// Powers the response's absence honesty (spec.md#search): "no relevant
2563    /// hits" only means something relative to how many messages were
2564    /// searchable at all, and 0 tells the caller their filters excluded
2565    /// everything before retrieval even started.
2566    pub async fn searchable_in_scope(&self, filter: &Predicate) -> Result<usize> {
2567        // Unfiltered: the FTS index already counts non-null search_text rows
2568        // (`num_docs`), and fast_search only searches those indexed docs - so
2569        // num_docs is exactly the universe a search ran over. Reading it avoids
2570        // the ~133 MB `IsNotNull(search_text)` column scan Lance pays per query
2571        // (no per-column null metadata). Filtered scopes fall back to the scan.
2572        if matches!(filter, Predicate::And(clauses) if clauses.is_empty())
2573            && let Some(count) = self.fts_num_docs().await?
2574        {
2575            return Ok(count);
2576        }
2577        let scope = Predicate::And(vec![Predicate::IsNotNull("search_text"), filter.clone()]);
2578        let dataset = self.handle.dataset(Table::Messages).await?;
2579        let count = dataset.count_rows(Some(scope.to_lance())).await?;
2580        Ok(count)
2581    }
2582
2583    /// Non-null `search_text` count read from the FTS index's `num_docs`
2584    /// statistic (summed across delta segments). `None` when the FTS index is
2585    /// absent (empty store) so the caller falls back to the count scan.
2586    async fn fts_num_docs(&self) -> Result<Option<usize>> {
2587        if !self.handle.messages_has_index(MESSAGES_FTS_INDEX).await? {
2588            return Ok(None);
2589        }
2590        let dataset = self.handle.dataset(Table::Messages).await?;
2591        let json = dataset.index_statistics(MESSAGES_FTS_INDEX).await?;
2592        let parsed: Value =
2593            serde_json::from_str(&json).context("failed to parse FTS index_statistics")?;
2594        let total: u64 = parsed["indices"]
2595            .as_array()
2596            .map(|segments| {
2597                segments
2598                    .iter()
2599                    .filter_map(|segment| segment["num_docs"].as_u64())
2600                    .sum()
2601            })
2602            .unwrap_or(0);
2603        Ok(Some(usize::try_from(total).unwrap_or(usize::MAX)))
2604    }
2605
2606    /// Whether any `messages` row carries a vector (spec.md#search) - the
2607    /// signal that lets the `vector` arm run instead of degrading to `fts`.
2608    /// The IVF index exists only once embeddings cross the activation
2609    /// threshold, so its presence proves embeddings exist via a resident
2610    /// manifest read - NOT an `IsNotNull("vector")` scan, which Lance cannot
2611    /// answer from stats and so reads the whole ~GB vector column from the
2612    /// store on every query. Below the threshold (no index yet) fall back to
2613    /// the prior `IsNotNull("vector")` `LIMIT 1` probe.
2614    pub async fn has_embeddings(&self) -> Result<bool> {
2615        if self
2616            .handle
2617            .messages_has_index(MESSAGES_VECTOR_INDEX)
2618            .await?
2619        {
2620            return Ok(true);
2621        }
2622        let scope = Predicate::IsNotNull("vector");
2623        let mut scanner = self
2624            .handle
2625            .scan(
2626                Table::Messages,
2627                ScanOpts::with_predicate_and_projection(&scope, &["id"]),
2628            )
2629            .await?;
2630        scanner.limit(Some(1), None)?;
2631        let batch = scanner.try_into_batch().await?;
2632        Ok(batch.num_rows() > 0)
2633    }
2634
2635    /// One embedded row's model id, or `None` when nothing is embedded yet. A
2636    /// `LIMIT 1` point read: the single-active-model invariant (see
2637    /// `has_embeddings`) means any embedded row's model is representative, so a
2638    /// model swap is detectable by comparing this to the configured model -
2639    /// without the full-column `stale_embedding_count` scan that ran every sync.
2640    pub async fn sample_embedded_model(&self) -> Result<Option<String>> {
2641        let scope = Predicate::IsNotNull("embedding_model");
2642        let mut scanner = self
2643            .handle
2644            .scan(
2645                Table::Messages,
2646                ScanOpts::with_predicate_and_projection(&scope, &["embedding_model"]),
2647            )
2648            .await?;
2649        scanner.limit(Some(1), None)?;
2650        let batch = scanner.try_into_batch().await?;
2651        if batch.num_rows() == 0 {
2652            return Ok(None);
2653        }
2654        string(&batch, "embedding_model", 0)
2655    }
2656
2657    /// Whether `messages` were embedded under a model id other than the
2658    /// configured one - a swap that requires re-embedding under the new model.
2659    /// One `LIMIT 1` read via [`Self::sample_embedded_model`]; the shared check
2660    /// behind the sync swap guard and the optimize embed stage.
2661    pub async fn embedding_model_swapped(&self) -> Result<bool> {
2662        Ok(self
2663            .sample_embedded_model()
2664            .await?
2665            .is_some_and(|model| model != crate::embed::model_id()))
2666    }
2667
2668    /// Vector kNN retriever over `messages.vector`, prefiltered by the caller's
2669    /// scalar predicate alone (spec.md#search-prefilter-pushdown) - see
2670    /// `embedded_scope` for why pond does NOT add `vector IS NOT NULL`. nprobes
2671    /// falls back to [`DEFAULT_NPROBES`] when `[search]` leaves it unset, so a
2672    /// default install never inherits Lance's unbounded "probe every partition"
2673    /// behavior on a remote store. No refine (see `apply_vector_search_knobs`).
2674    /// Index-only + map resolve when loaded, else key projection - see `fts_search`.
2675    pub async fn vector_search(
2676        &self,
2677        query: &[f32],
2678        limit: usize,
2679        filter: &Predicate,
2680        search: Option<&config::SearchConfig>,
2681    ) -> Result<Vec<SearchHit>> {
2682        let mut hits = if let Some(map) = self.rowmap.load_full() {
2683            let rowid_hits = self
2684                .vector_search_rowids(query, limit, filter, search)
2685                .await?;
2686            self.resolve_rowid_hits(&map, rowid_hits).await?
2687        } else {
2688            self.vector_search_keys(query, limit, filter, search)
2689                .await?
2690        };
2691        // Stable secondary sort: same reasoning as `fts_search` - IVF_SQ can
2692        // emit hits with effectively identical `_distance` in fragment-dependent
2693        // order, which makes RRF dedup-ranks nondeterministic for tied
2694        // neighbors. Sort by distance asc (smaller = more similar), then by
2695        // `(session_id, message_id)` asc.
2696        hits.sort_by(|left, right| {
2697            left.score
2698                .partial_cmp(&right.score)
2699                .unwrap_or(std::cmp::Ordering::Equal)
2700                .then_with(|| left.key.session_id.cmp(&right.key.session_id))
2701                .then_with(|| left.key.message_id.cmp(&right.key.message_id))
2702        });
2703        Ok(hits)
2704    }
2705
2706    /// Shared vector-scan setup: scope, `nearest`, knobs, `fast_search`.
2707    async fn vector_scanner(
2708        &self,
2709        query: &[f32],
2710        limit: usize,
2711        filter: &Predicate,
2712        search: Option<&config::SearchConfig>,
2713    ) -> Result<lance::dataset::scanner::Scanner> {
2714        let scope = embedded_scope(filter);
2715        let mut scanner = self.handle.scanner(Table::Messages, Some(&scope)).await?;
2716        let key = Float32Array::from(query.to_vec());
2717        scanner.nearest("vector", &key, limit)?;
2718        apply_vector_search_knobs(&mut scanner, search);
2719        if self
2720            .handle
2721            .messages_fast_search_ready(MESSAGES_VECTOR_INDEX)
2722            .await?
2723        {
2724            scanner.fast_search();
2725        }
2726        scanner.disable_scoring_autoprojection();
2727        Ok(scanner)
2728    }
2729
2730    /// Index-only vector retriever: `_rowid` + `_distance` only, so no `TakeExec`.
2731    /// `vector_search` resolves the row ids. Mirrors `fts_search_rowids`.
2732    async fn vector_search_rowids(
2733        &self,
2734        query: &[f32],
2735        limit: usize,
2736        filter: &Predicate,
2737        search: Option<&config::SearchConfig>,
2738    ) -> Result<Vec<(u64, f32)>> {
2739        let mut scanner = self.vector_scanner(query, limit, filter, search).await?;
2740        scanner.with_row_id();
2741        scanner.project(&["_distance"])?;
2742        let batch = scanner.try_into_batch().await?;
2743        let rowids = uint64(&batch, "_rowid")?;
2744        let mut hits = Vec::with_capacity(batch.num_rows());
2745        for row in 0..batch.num_rows() {
2746            hits.push((rowids.value(row), float32(&batch, "_distance", row)?));
2747        }
2748        Ok(hits)
2749    }
2750
2751    /// No-map vector fallback: project the key columns plus `_distance` directly.
2752    /// Unsorted; `vector_search` sorts. Mirrors `fts_search_keys`.
2753    async fn vector_search_keys(
2754        &self,
2755        query: &[f32],
2756        limit: usize,
2757        filter: &Predicate,
2758        search: Option<&config::SearchConfig>,
2759    ) -> Result<Vec<SearchHit>> {
2760        let mut scanner = self.vector_scanner(query, limit, filter, search).await?;
2761        scanner.project(&["session_id", "id", "_distance"])?;
2762        let batch = scanner.try_into_batch().await?;
2763        let mut hits = Vec::with_capacity(batch.num_rows());
2764        for row in 0..batch.num_rows() {
2765            let key = MessageKey {
2766                session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2767                message_id: string(&batch, "id", row)?.context("message id is null")?,
2768            };
2769            hits.push(SearchHit {
2770                rowid: None,
2771                key,
2772                score: float32(&batch, "_distance", row)?,
2773            });
2774        }
2775        Ok(hits)
2776    }
2777
2778    /// The DataFusion plan string for a filtered vector scan - the
2779    /// `search-prefilter-pushdown` regression guard reads it.
2780    pub async fn explain_vector_plan(
2781        &self,
2782        query: &[f32],
2783        limit: usize,
2784        filter: &Predicate,
2785        search: Option<&config::SearchConfig>,
2786    ) -> Result<String> {
2787        // Reuse the real retriever's builder so the explained plan can never
2788        // drift from what a query actually runs - notably the fast_search vs
2789        // flat-tail gate (`messages_fast_search_ready`).
2790        let scanner = self.vector_scanner(query, limit, filter, search).await?;
2791        scanner
2792            .explain_plan(true)
2793            .await
2794            .context("explain_plan failed")
2795    }
2796
2797    pub async fn explain_fts_plan(
2798        &self,
2799        query: &str,
2800        limit: usize,
2801        filter: &Predicate,
2802    ) -> Result<String> {
2803        // Same builder as `fts_search` so the explained plan matches execution.
2804        let mut scanner = self.fts_scanner(query, limit, filter).await?;
2805        scanner.project(&["session_id", "id"])?;
2806        scanner
2807            .explain_plan(true)
2808            .await
2809            .context("explain_plan failed")
2810    }
2811
2812    /// Hydrate search hits by stable row id (spec.md#search). Resolves each
2813    /// rowid from the resident meta map in memory (no object-store round-trip -
2814    /// Lance caches index/metadata but never data column values, so a `take_rows`
2815    /// re-reads `search_text` from storage every query). Rowids the map lacks
2816    /// (appended since it was built, or no map loaded) fall back to a single
2817    /// `take_rows` batch. The caller indexes the result by key, so order is
2818    /// irrelevant.
2819    pub async fn message_metas_by_rowids(&self, rowids: &[u64]) -> Result<Vec<MessageMeta>> {
2820        if rowids.is_empty() {
2821            return Ok(Vec::new());
2822        }
2823        let mut metas = Vec::with_capacity(rowids.len());
2824        let misses: Vec<u64> = if let Some(map) = self.rowmap.load_full() {
2825            let (hits, misses) = map.hydrate(rowids);
2826            metas.extend(hits.into_iter().map(|entry| MessageMeta {
2827                message_id: entry.message_id,
2828                session_id: entry.session_id,
2829                role: entry.role,
2830                project: entry.project,
2831                source_agent: entry.source_agent,
2832                timestamp:
2833                    DateTime::from_timestamp_micros(entry.timestamp_micros).unwrap_or_default(),
2834                search_text: entry.search_text,
2835            }));
2836            misses
2837        } else {
2838            rowids.to_vec()
2839        };
2840        if !misses.is_empty() {
2841            metas.extend(self.message_metas_by_rowids_take(&misses).await?);
2842        }
2843        Ok(metas)
2844    }
2845
2846    /// `take_rows` hydration of exactly `rowids` - the cache-miss fallback for
2847    /// rows the resident meta map lacks. `take_rows` coalesces the reads per
2848    /// fragment (Lance's own batching), so a scattered take is few requests, not
2849    /// one per row.
2850    async fn message_metas_by_rowids_take(&self, rowids: &[u64]) -> Result<Vec<MessageMeta>> {
2851        let dataset = self.handle.dataset(Table::Messages).await?;
2852        let projection = ProjectionRequest::from_columns(
2853            [
2854                "id",
2855                "session_id",
2856                "role",
2857                "project",
2858                "source_agent",
2859                "timestamp",
2860                "search_text",
2861            ],
2862            dataset.schema(),
2863        );
2864        let batch = dataset.take_rows(rowids, projection).await?;
2865        let mut metas = Vec::with_capacity(batch.num_rows());
2866        for row in 0..batch.num_rows() {
2867            metas.push(message_meta_from_batch(&batch, row)?);
2868        }
2869        Ok(metas)
2870    }
2871
2872    /// Hydrate search hits: fetch message metadata for `(session_id, message_id)` keys.
2873    pub async fn message_metas_by_keys(&self, keys: &[MessageKey]) -> Result<Vec<MessageMeta>> {
2874        if keys.is_empty() {
2875            return Ok(Vec::new());
2876        }
2877        let wanted = keys.iter().cloned().collect::<HashSet<_>>();
2878        let session_ids = keys
2879            .iter()
2880            .map(|key| key.session_id.clone())
2881            .collect::<Vec<_>>();
2882        let message_ids = keys
2883            .iter()
2884            .map(|key| key.message_id.clone())
2885            .collect::<Vec<_>>();
2886        let predicate = Predicate::And(vec![
2887            in_predicate("session_id", &session_ids),
2888            in_predicate("id", &message_ids),
2889        ]);
2890        let batch = self
2891            .handle
2892            .scan_batch(
2893                Table::Messages,
2894                Some(&predicate),
2895                &[
2896                    "id",
2897                    "session_id",
2898                    "role",
2899                    "project",
2900                    "source_agent",
2901                    "timestamp",
2902                    "search_text",
2903                ],
2904            )
2905            .await?;
2906        let mut metas = Vec::with_capacity(batch.num_rows());
2907        for row in 0..batch.num_rows() {
2908            // The IN x IN predicate is a cross-product, so the scan can return
2909            // pairs that were never asked for; keep only the wanted keys.
2910            let meta = message_meta_from_batch(&batch, row)?;
2911            if wanted.contains(&MessageKey {
2912                session_id: meta.session_id.clone(),
2913                message_id: meta.message_id.clone(),
2914            }) {
2915                metas.push(meta);
2916            }
2917        }
2918        Ok(metas)
2919    }
2920
2921    /// Total message count per session, for search session summaries. One
2922    /// `session_id IN (...)` scan projecting only `session_id`, aggregated in
2923    /// pond, instead of `N` concurrent `count_rows(session_id = X)` round-trips
2924    /// against `messages_session_id_btree`. Same wire shape for any backend,
2925    /// but one S3 operation instead of `N` on remote stores. Sessions with
2926    /// zero matching messages are present in the map with count `0` so the
2927    /// caller can distinguish "filter excluded everything" from "session
2928    /// missing from the response."
2929    pub async fn session_message_counts(
2930        &self,
2931        session_ids: &[String],
2932    ) -> Result<BTreeMap<String, usize>> {
2933        if session_ids.is_empty() {
2934            return Ok(BTreeMap::new());
2935        }
2936        // A version-matched resident map covers every current row, so its
2937        // per-session counts are authoritative (a session absent from it has 0
2938        // messages) - serve them with no scan. The version gate is load-bearing:
2939        // unlike meta hydration, a count cannot detect staleness by a row-id
2940        // miss, so a map that predates appended rows would undercount. A stale
2941        // or absent map falls through to the IN-scan.
2942        if let Some(map) = self.rowmap.load_full()
2943            && map.version() == self.messages_version().await?
2944        {
2945            return Ok(session_ids
2946                .iter()
2947                .map(|id| (id.clone(), map.lookup_count(id).unwrap_or(0)))
2948                .collect());
2949        }
2950        let predicate = in_predicate("session_id", session_ids);
2951        let scanner = self
2952            .handle
2953            .scan(
2954                Table::Messages,
2955                ScanOpts::with_predicate_and_projection(&predicate, &["session_id"]),
2956            )
2957            .await?;
2958        let mut stream = scanner
2959            .try_into_stream()
2960            .await
2961            .context("failed to open session_message_counts stream")?;
2962        let mut counts: BTreeMap<String, usize> =
2963            session_ids.iter().map(|id| (id.clone(), 0)).collect();
2964        while let Some(batch) = stream.next().await {
2965            let batch = batch.context("failed to read session_message_counts batch")?;
2966            let column = batch
2967                .column_by_name("session_id")
2968                .context("session_message_counts: session_id column missing")?
2969                .as_any()
2970                .downcast_ref::<StringArray>()
2971                .context("session_message_counts: session_id column is not Utf8")?;
2972            for value in column.iter().flatten() {
2973                if let Some(entry) = counts.get_mut(value) {
2974                    *entry += 1;
2975                }
2976            }
2977        }
2978        Ok(counts)
2979    }
2980
2981    /// Rows appended to `messages` since the FTS index was last optimized.
2982    /// A missing index reports the whole table; the query is manifest-only.
2983    pub async fn unindexed_message_backlog(&self) -> Result<usize> {
2984        self.handle
2985            .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
2986            .await
2987    }
2988
2989    /// Rows added or rewritten in `messages` since the IVF_SQ vector index
2990    /// was last folded; a missing index reports the whole table. Manifest-only,
2991    /// and an upper bound on the embed backlog (a row folds only after it
2992    /// embeds, so no unembedded row is ever folded): zero proves nothing is
2993    /// unembedded, but non-zero can be all-embedded rows the index has not
2994    /// absorbed yet - confirm with [`embed_backlog_count`](Self::embed_backlog_count)
2995    /// before acting on it.
2996    pub async fn unindexed_vector_backlog(&self) -> Result<usize> {
2997        self.handle
2998            .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
2999            .await
3000    }
3001
3002    /// Embedding coverage: how many `messages` rows carry a vector and how
3003    /// many are still eligible. Drives the `pond status` embeddings line and
3004    /// the `pond optimize` progress bar's known total.
3005    pub async fn embedding_progress(&self) -> Result<EmbeddingProgress> {
3006        let dataset = self.handle.dataset(Table::Messages).await?;
3007        // `embedded` counts `embedding_model IS NOT NULL`, not `vector`: the two
3008        // are co-set (spec.md#session-embed-from-canonical) so the count is
3009        // identical, but the model-id string column is ~50x narrower than the
3010        // Float16 vector (Lance 7.0.0 has no per-column null_count, so this is a
3011        // data-page read).
3012        let embedded = dataset
3013            .count_rows(Some(Predicate::IsNotNull("embedding_model").to_lance()))
3014            .await?;
3015        // `backlog` and `total` come from live, deletion-aware counts, not the
3016        // FTS `num_docs`: num_docs counts indexed docs incl. deleted-but-unpurged
3017        // ones, so `num_docs - embedded` reports a phantom backlog that survives
3018        // every embed. `embedded` (model present) + `backlog` (model absent,
3019        // search_text present) is exactly the live eligible set, since embedding
3020        // a row requires its search_text.
3021        let backlog = self.embed_backlog_count().await?;
3022        Ok(EmbeddingProgress {
3023            embedded,
3024            total: embedded + backlog,
3025            backlog,
3026            model: embed::model_id(),
3027        })
3028    }
3029
3030    /// Messages eligible but not yet embedded (`search_text` present,
3031    /// `embedding_model` null) - the exact set [`crate::embed::EmbedWorker`]
3032    /// processes. Read straight from the dataset so it is correct right after
3033    /// ingest, unlike the FTS `num_docs` `embedding_progress` shows (which lags
3034    /// until the index is rebuilt - the embed stage runs before that).
3035    pub async fn embed_backlog_count(&self) -> Result<usize> {
3036        let dataset = self.handle.dataset(Table::Messages).await?;
3037        let filter = Predicate::And(vec![
3038            Predicate::IsNull("embedding_model"),
3039            Predicate::IsNotNull("search_text"),
3040        ]);
3041        Ok(dataset.count_rows(Some(filter.to_lance())).await?)
3042    }
3043
3044    /// Count rows whose `embedding_model` is not the currently configured
3045    /// model AND whose `vector` is still populated - the signal `pond optimize`
3046    /// uses to detect a model swap and require `--force-embed`.
3047    pub async fn stale_embedding_count(&self) -> Result<usize> {
3048        let dataset = self.handle.dataset(Table::Messages).await?;
3049        // Same shape as the original (IsNotNull AND Ne), but the null check is on
3050        // the narrow model-id column, not the ~50x-wider Float16 vector: the two
3051        // are co-set (spec.md#session-embed-from-canonical), so `embedding_model
3052        // IS NOT NULL` equals `vector IS NOT NULL`, and the model-id page read is
3053        // far cheaper than the vector's.
3054        dataset
3055            .count_rows(Some(
3056                Predicate::And(vec![
3057                    Predicate::IsNotNull("embedding_model"),
3058                    Predicate::Ne("embedding_model", embed::model_id().into()),
3059                ])
3060                .to_lance(),
3061            ))
3062            .await
3063            .map_err(Into::into)
3064    }
3065
3066    /// Run the per-table maintenance cycle (compact + indices) across every
3067    /// table, never short-circuiting. spec.md#lance-index-maintenance: indices
3068    /// and compaction commit independently, so a hot writer that starves
3069    /// compaction on one table does not abort the index work the operator
3070    /// asked for on other tables (or even on the same table).
3071    pub async fn optimize_indices(
3072        &self,
3073        progress: Option<OptimizeProgressFn>,
3074        maintenance: &MaintenancePolicy,
3075    ) -> Result<OptimizeOutcome> {
3076        let intents = pond_index_intents();
3077        let mut tables = Vec::with_capacity(3);
3078        for (table, intents) in intents.all() {
3079            let outcome = self
3080                .handle
3081                .optimize_table(table, intents, progress.as_ref(), maintenance)
3082                .await;
3083            tables.push(outcome);
3084        }
3085        Ok(OptimizeOutcome { tables })
3086    }
3087
3088    /// Fold trailing fragments into existing indices across every table,
3089    /// without running compaction. Used by `pond optimize`'s tail so newly
3090    /// written vectors land in the FTS / IVF_SQ / btree / bitmap indices
3091    /// without paying the compaction retry budget while embed itself may
3092    /// still be writing in a sibling process.
3093    pub async fn build_indices_only(
3094        &self,
3095        progress: Option<OptimizeProgressFn>,
3096    ) -> Result<OptimizeOutcome> {
3097        let policy = pond_index_intents();
3098        let mut tables = Vec::with_capacity(3);
3099        for (table, intents) in policy.all() {
3100            let indices = self
3101                .handle
3102                .optimize_table_indices_only(table, intents, progress.as_ref())
3103                .await;
3104            tables.push(TableOptimizeOutcome {
3105                table,
3106                indices,
3107                compaction: PhaseOutcome::NotAttempted,
3108            });
3109        }
3110        Ok(OptimizeOutcome { tables })
3111    }
3112
3113    #[cfg(test)]
3114    async fn optimize_indices_with_vector_threshold(
3115        &self,
3116        vector_threshold: usize,
3117    ) -> Result<OptimizeOutcome> {
3118        let intents = pond_index_intents_with_vector_threshold(vector_threshold);
3119        let policy = MaintenancePolicy::always_compact();
3120        let mut tables = Vec::with_capacity(3);
3121        for (table, intents) in intents.all() {
3122            let outcome = self
3123                .handle
3124                .optimize_table(table, intents, None, &policy)
3125                .await;
3126            tables.push(outcome);
3127        }
3128        Ok(OptimizeOutcome { tables })
3129    }
3130
3131    #[cfg(test)]
3132    async fn optimize_indices_with_scalar_fold_threshold(
3133        &self,
3134        scalar_fold_row_threshold: usize,
3135    ) -> Result<OptimizeOutcome> {
3136        let intents = pond_index_intents();
3137        let policy = MaintenancePolicy::always_compact()
3138            .with_scalar_fold_row_threshold(scalar_fold_row_threshold);
3139        let mut tables = Vec::with_capacity(3);
3140        for (table, intents) in intents.all() {
3141            let outcome = self
3142                .handle
3143                .optimize_table(table, intents, None, &policy)
3144                .await;
3145            tables.push(outcome);
3146        }
3147        Ok(OptimizeOutcome { tables })
3148    }
3149
3150    /// Reclaim superseded data/index files across every indexed table (Lance
3151    /// `cleanup_old_versions`), without compaction. `pond optimize --rebuild`
3152    /// runs this after the rebuild so the index segments it just replaced are
3153    /// dropped immediately. The retention floor still protects versions a live
3154    /// reader may have pinned (spec.md#concurrency).
3155    pub async fn cleanup_old_versions(&self, older_than: chrono::Duration) -> Result<()> {
3156        for (table, _) in pond_index_intents().all() {
3157            self.handle
3158                .cleanup_table_versions(table, older_than)
3159                .await?;
3160        }
3161        Ok(())
3162    }
3163
3164    pub async fn rebuild_indices(
3165        &self,
3166        intent_name: Option<&str>,
3167        progress: Option<OptimizeProgressFn>,
3168    ) -> Result<()> {
3169        let policy = pond_index_intents();
3170        let mut matched = false;
3171        for (table, intents) in policy.all() {
3172            for intent in intents {
3173                if intent_name.is_none_or(|name| name == intent.name) {
3174                    matched = true;
3175                    self.handle
3176                        .rebuild_index(table, intent, progress.as_ref())
3177                        .await?;
3178                }
3179            }
3180        }
3181        if let Some(name) = intent_name
3182            && !matched
3183        {
3184            anyhow::bail!("unknown index intent {name:?}");
3185        }
3186        Ok(())
3187    }
3188
3189    /// Drop a named index from whichever table owns it. Used by `pond optimize
3190    /// --drop-index <name>` to clean up orphaned indices (e.g. after renaming
3191    /// an intent whose on-disk name no longer matches the policy). Finds the
3192    /// owning table via parallel `load_indices` lookups, then drops on just
3193    /// that table - so real I/O errors surface with the right context instead
3194    /// of being hidden behind "no such index" from the wrong table.
3195    pub async fn drop_index_by_name(&self, name: &str) -> Result<()> {
3196        let Some(owner) = self.handle.find_index_owner(name).await? else {
3197            anyhow::bail!("no index named {name:?} found on any table");
3198        };
3199        self.handle.drop_index(owner, name).await
3200    }
3201
3202    pub async fn index_status(&self) -> Result<Vec<IndexStatus>> {
3203        self.index_status_with(false).await
3204    }
3205
3206    /// Like [`Self::index_status`], but content indexes (FTS, IVF) report only
3207    /// the non-null - actually indexable - rows of their unindexed tail. The
3208    /// honest number for `pond status`; costs a tail-bounded scan, so the
3209    /// per-sync summary stays on the cheap manifest-only variant.
3210    pub async fn index_status_indexable(&self) -> Result<Vec<IndexStatus>> {
3211        self.index_status_with(true).await
3212    }
3213
3214    async fn index_status_with(&self, indexable_only: bool) -> Result<Vec<IndexStatus>> {
3215        let policy = pond_index_intents();
3216        let mut statuses = Vec::new();
3217        for (table, intents) in policy.all() {
3218            statuses.extend(
3219                self.handle
3220                    .index_status(table, intents, indexable_only)
3221                    .await?,
3222            );
3223        }
3224        Ok(statuses)
3225    }
3226
3227    /// Drop the IVF_SQ index on `messages.vector`. Used by `pond optimize
3228    /// --force-embed` before re-bootstrapping under a different model. Silent
3229    /// when the index does not exist.
3230    pub async fn drop_vector_index(&self) -> Result<()> {
3231        match self
3232            .handle
3233            .drop_index(Table::Messages, MESSAGES_VECTOR_INDEX)
3234            .await
3235        {
3236            Ok(()) => Ok(()),
3237            Err(error) => {
3238                let msg = error.to_string();
3239                if msg.contains("not found") || msg.contains("does not exist") {
3240                    Ok(())
3241                } else {
3242                    Err(error)
3243                }
3244            }
3245        }
3246    }
3247
3248    /// On-disk byte totals per dataset, sized through Lance's object store
3249    /// (spec.md#lance-chokepoints-storage) so `pond status` works on any backend.
3250    pub async fn table_sizes(&self) -> Result<TableSizes> {
3251        self.handle.table_sizes().await
3252    }
3253
3254    pub async fn initialized(&self) -> Result<bool> {
3255        self.handle.initialized().await
3256    }
3257
3258    /// Single Sessions-table lookup - no message hydration. The cheap sibling
3259    /// of [`Store::get_session`] for callers that only need session metadata.
3260    pub(crate) async fn find_session(&self, session_id: &str) -> Result<Option<Session>> {
3261        let batch = self
3262            .handle
3263            .scan_batch(
3264                Table::Sessions,
3265                Some(&Predicate::Eq("id", session_id.into())),
3266                &[],
3267            )
3268            .await?;
3269        if batch.num_rows() == 0 {
3270            Ok(None)
3271        } else {
3272            Ok(Some(session_from_batch(&batch, 0)?))
3273        }
3274    }
3275
3276    async fn messages_for_session(&self, session_id: &str) -> Result<Vec<MessageWithParts>> {
3277        let batch = self
3278            .handle
3279            .scan_batch(
3280                Table::Messages,
3281                Some(&Predicate::Eq("session_id", session_id.into())),
3282                &[
3283                    "session_id",
3284                    "id",
3285                    "timestamp",
3286                    "role",
3287                    "content",
3288                    "options",
3289                ],
3290            )
3291            .await?;
3292        let mut messages = Vec::with_capacity(batch.num_rows());
3293        for row in 0..batch.num_rows() {
3294            messages.push(message_from_batch(&batch, row)?);
3295        }
3296        messages.sort_by(|left, right| {
3297            left.timestamp()
3298                .cmp(&right.timestamp())
3299                .then_with(|| left.id().cmp(right.id()))
3300        });
3301
3302        let message_ids = messages
3303            .iter()
3304            .map(|message| message.id().to_owned())
3305            .collect::<Vec<_>>();
3306        let mut parts_by_message = self.parts_for_messages(session_id, &message_ids).await?;
3307
3308        Ok(messages
3309            .into_iter()
3310            .map(|message| {
3311                let key = (message.session_id().to_owned(), message.id().to_owned());
3312                let parts = parts_by_message.remove(&key).unwrap_or_default();
3313                MessageWithParts { message, parts }
3314            })
3315            .collect())
3316    }
3317
3318    /// Every part of these messages, full fidelity (file blobs included). The
3319    /// canonical read primitive - restore/export, verbatim mode, and the
3320    /// message-mode target all need the complete set.
3321    pub async fn parts_for_messages(
3322        &self,
3323        session_id: &str,
3324        message_ids: &[String],
3325    ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3326        self.scan_parts(session_id, message_ids, None).await
3327    }
3328
3329    /// Only the parts that yield a [`PartSummary`] ([`SUMMARY_PART_TYPES`]),
3330    /// skipping `text`/`reasoning` (and their blobs) that would summarize to
3331    /// nothing. For the summary-only reads (conversational/complete session
3332    /// views, search hits) - it never feeds restore/export.
3333    pub async fn summary_parts_for_messages(
3334        &self,
3335        session_id: &str,
3336        message_ids: &[String],
3337    ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3338        self.scan_parts(session_id, message_ids, Some(SUMMARY_PART_TYPES))
3339            .await
3340    }
3341
3342    async fn scan_parts(
3343        &self,
3344        session_id: &str,
3345        message_ids: &[String],
3346        part_types: Option<&[&str]>,
3347    ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3348        if message_ids.is_empty() {
3349            return Ok(BTreeMap::new());
3350        }
3351        let mut clauses = vec![
3352            Predicate::Eq("session_id", session_id.into()),
3353            in_predicate("message_id", message_ids),
3354        ];
3355        if let Some(types) = part_types {
3356            clauses.push(Predicate::In(
3357                "type",
3358                types.iter().map(|&t| t.into()).collect(),
3359            ));
3360        }
3361        let predicate = Predicate::And(clauses);
3362        // Summary reads (search hits, conversational view) need only the part
3363        // metadata in `variant_data` to build a `PartSummary` - never the file
3364        // blob - so they skip the `_rowaddr` + `take_blobs` round trip. Only
3365        // full-fidelity callers (restore/export/message-mode) read the blobs.
3366        let summarizing = part_types.is_some();
3367        let mut scanner = self
3368            .handle
3369            .scan(
3370                Table::Parts,
3371                ScanOpts::with_predicate_and_projection(
3372                    &predicate,
3373                    &[
3374                        "session_id",
3375                        "message_id",
3376                        "id",
3377                        "ordinal",
3378                        "type",
3379                        "provenance",
3380                        "variant_data",
3381                        "options",
3382                    ],
3383                ),
3384            )
3385            .await?;
3386        if !summarizing {
3387            scanner.with_row_address();
3388        }
3389        let batch = scanner.try_into_batch().await.context("scan failed")?;
3390        let mut file_payloads = BTreeMap::<usize, FileData>::new();
3391        if !summarizing {
3392            let dataset = std::sync::Arc::new(self.handle.dataset(Table::Parts).await?);
3393            let row_addresses = uint64(&batch, "_rowaddr")?;
3394            let mut file_rows = Vec::<(usize, u64, Vec<u8>)>::new();
3395            for row in 0..batch.num_rows() {
3396                if string(&batch, "type", row)?.as_deref() == Some("file") {
3397                    let variant_data = json_column(&batch, "variant_data", row)?
3398                        .context("variant_data is null")?;
3399                    file_rows.push((row, row_addresses.value(row), variant_data));
3400                }
3401            }
3402            if !file_rows.is_empty() {
3403                let addresses = file_rows
3404                    .iter()
3405                    .map(|(_, address, _)| *address)
3406                    .collect::<Vec<_>>();
3407                let blobs = dataset.take_blobs_by_addresses(&addresses, "data").await?;
3408                for ((row, _, variant_data), blob) in file_rows.into_iter().zip(blobs) {
3409                    // Legacy blob (lance-encoding:blob): payload is bytes; the
3410                    // url variant stored its URL as UTF-8 bytes, recovered via
3411                    // `file_data_from_blob`'s `data_kind = "url"` branch.
3412                    let payload = file_data_from_blob(&variant_data, &blob.read().await?)?;
3413                    file_payloads.insert(row, payload);
3414                }
3415            }
3416        }
3417        let mut parts_by_message = BTreeMap::<(String, String), Vec<Part>>::new();
3418        for row in 0..batch.num_rows() {
3419            // A summary discards file contents (`PartSummary::for_kind` reads
3420            // only `file_name`/`media_type` from `variant_data`); pass an empty
3421            // placeholder so `PartKind::File` still deserializes without a blob.
3422            let file_data = if summarizing {
3423                (string(&batch, "type", row)?.as_deref() == Some("file"))
3424                    .then(|| FileData::Bytes(Vec::new()))
3425            } else {
3426                file_payloads.remove(&row)
3427            };
3428            let part = part_from_batch(&batch, row, file_data)?;
3429            parts_by_message
3430                .entry((part.session_id.clone(), part.message_id.clone()))
3431                .or_default()
3432                .push(part);
3433        }
3434        for parts in parts_by_message.values_mut() {
3435            parts.sort_by_key(|part| part.ordinal);
3436        }
3437        Ok(parts_by_message)
3438    }
3439}
3440
3441#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
3442#[serde(tag = "kind", content = "data", rename_all = "snake_case")]
3443pub enum IngestEvent {
3444    Session(Session),
3445    Message(Message),
3446    Part(Part),
3447}
3448
3449/// Aggregate accounting for an ingest pass (CLI sync, adapter-driven).
3450/// The wire layer (`pond_ingest`) instead returns per-row results; the
3451/// aggregate is derived from those at the wire boundary.
3452///
3453/// Fields are bucketed by population so the summary never conflates "100
3454/// validator-rejected rows in 1 bad session" with "100 separate failures."
3455/// The shape is set by spec.md#adapter-integrity-event-ordering.
3456#[derive(Debug, Clone, PartialEq, Eq, Default)]
3457pub struct IngestSummary {
3458    /// Rows actually written to Lance, summed across all three tables.
3459    /// Use the per-table fields below for user-facing counts; this stays
3460    /// for `accepted()` and existing wire callers.
3461    pub inserted: usize,
3462    /// Rows that already existed (merge_insert no-op match).
3463    pub matched: usize,
3464    /// Session rows inserted this pass.
3465    pub sessions_inserted: usize,
3466    /// Message rows inserted this pass (total - includes tool calls,
3467    /// tool results, and other non-searchable messages).
3468    pub messages_inserted_total: usize,
3469    /// Subset of `messages_inserted_total` whose `search_text` is non-null
3470    /// (eligible for FTS + semantic indexing). The user-facing "messages"
3471    /// count in `pond sync` / `pond status` reads this field.
3472    pub messages_inserted_searchable: usize,
3473    /// Part rows inserted this pass.
3474    pub parts_inserted: usize,
3475    /// Session rows already-present (merge_insert matched).
3476    pub sessions_matched: usize,
3477    /// Message rows already-present (merge_insert matched), total.
3478    pub messages_matched_total: usize,
3479    /// Subset of `messages_matched_total` with `search_text`.
3480    pub messages_matched_searchable: usize,
3481    /// Part rows already-present.
3482    pub parts_matched: usize,
3483    /// Events the validator dropped under per-event-drop policy (ordering
3484    /// violation, orphan part, mismatched parent, adapter parse failure,
3485    /// duplicate-id collision, ...). Counted by event, not by session: a
3486    /// session with one bad part stays in this bucket as 1, not as "the
3487    /// whole substream." Per spec.md#adapter-integrity-dedup, adapters SHOULD dedupe their
3488    /// own emissions upstream when source replay is expected; the
3489    /// validator's in-batch HashSet is a safety net, not a feature
3490    /// adapters may rely on. If this bucket grows on a clean adapter,
3491    /// inspect `drop_reasons` for the top contributors.
3492    pub dropped_events: usize,
3493    /// Sessions whose Session-level invariants (immutable `source_agent` /
3494    /// `project` against the stored row) failed at flush time and
3495    /// whose substream got rejected wholesale. Always small relative to
3496    /// `inserted`; if not, there's a real problem to investigate.
3497    pub dropped_sessions: usize,
3498    /// Files the adapter couldn't decode at all (no Session header
3499    /// extractable: empty `.jsonl`, missing required field).
3500    pub skipped_files: usize,
3501    /// Files that produced no importable session and were benignly skipped:
3502    /// empty `.jsonl`, sidecar-only rows (e.g. an `ai-title`/`agent-name`
3503    /// metadata file), or an unextractable header. Never an error or a drop;
3504    /// the underlying cause is logged at `-vv` (debug) verbosity.
3505    pub skipped_empty: usize,
3506    /// Sessions short-circuited via the per-session staleness skip
3507    /// (spec.md#adapter-integrity-event-ordering): file `mtime` was at or before the wall-clock time
3508    /// pond last wrote that session's row, so re-decode was bypassed.
3509    pub skipped_fresh: usize,
3510    /// Legacy/source copies dropped because an authoritative copy of the same
3511    /// session was ingested from another source form this run (currently:
3512    /// opencode tree copies superseded by the DB). Counted, never silent.
3513    pub skipped_superseded: usize,
3514    /// Storage-layer failures whose retries were exhausted (commit
3515    /// conflicts, transient IO that didn't recover). Hard zero on healthy
3516    /// runs.
3517    pub storage_errors: usize,
3518    /// Oversized values truncated to a bounded sentinel at the seam
3519    /// (spec.md#adapter-bounded-values); the rest of each such record is intact.
3520    pub truncated_values: usize,
3521    /// Histogram of stable reason keys for the combined `dropped_events +
3522    /// dropped_sessions` populations. Keys are `&'static str` (see the
3523    /// `DROP_REASON_*` constants) so consumers can match by identity.
3524    /// Empty on a clean run. Used by `pond sync` to print the top reasons
3525    /// and by `benches/ingest_bench.rs` to bucket Partial drops by cause.
3526    pub drop_reasons: BTreeMap<&'static str, usize>,
3527}
3528
3529/// Stable reason keys for the `IngestSummary::drop_reasons` histogram and
3530/// the per-row `RowError::reason_key`. `&'static str` so consumers can
3531/// match by identity rather than prose. Adding a new variant: pick a short
3532/// snake_case identifier, route it from the validator/adapter, and update
3533/// the per-row outcome docs in `docs/spec.md#adapter-integrity-event-ordering`.
3534pub const DROP_REASON_DUPLICATE_MESSAGE_ID: &str = "duplicate_message_id";
3535pub const DROP_REASON_DUPLICATE_PART_KEY: &str = "duplicate_part_key";
3536pub const DROP_REASON_MESSAGE_BEFORE_SESSION: &str = "message_before_session";
3537pub const DROP_REASON_MESSAGE_SESSION_MISMATCH: &str = "message_session_mismatch";
3538pub const DROP_REASON_PART_BEFORE_MESSAGE: &str = "part_before_message";
3539pub const DROP_REASON_PART_MESSAGE_MISMATCH: &str = "part_message_mismatch";
3540pub const DROP_REASON_EMPTY_SOURCE_AGENT: &str = "empty_source_agent";
3541pub const DROP_REASON_PARENT_MESSAGE_WITHOUT_SESSION: &str = "parent_message_without_session";
3542pub const DROP_REASON_IMMUTABLE_PROJECT: &str = "immutable_project";
3543pub const DROP_REASON_IMMUTABLE_SOURCE_AGENT: &str = "immutable_source_agent";
3544pub const DROP_REASON_UNCATEGORIZED: &str = "uncategorized";
3545
3546/// Honest per-table outcome of one batched flush. Built from `merge_insert`'s
3547/// returned counts together with the pre-existence sets captured by
3548/// `upsert_session_batch`. Folded into a per-sync summary via
3549/// [`IngestSummary::add_batch`]. spec.md#adapter-integrity-additive-sync: matched
3550/// is a no-op write, so the inserted/matched split is informational - we still
3551/// surface it because both `pond sync` and `pond_ingest` clients reconcile
3552/// against "which rows landed this call."
3553#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
3554pub struct BatchCounts {
3555    pub sessions_inserted: usize,
3556    pub sessions_matched: usize,
3557    pub messages_inserted_total: usize,
3558    pub messages_inserted_searchable: usize,
3559    pub messages_matched_total: usize,
3560    pub messages_matched_searchable: usize,
3561    pub parts_inserted: usize,
3562    pub parts_matched: usize,
3563}
3564
3565impl IngestSummary {
3566    pub fn accepted(&self) -> usize {
3567        self.inserted + self.matched
3568    }
3569
3570    /// Sole writer of the per-table counters on the CLI batched flush path.
3571    /// The wire single-row path keeps using [`Self::add_outcomes`]; emitting
3572    /// both for the same rows would double-count.
3573    pub fn add_batch(&mut self, counts: &BatchCounts) {
3574        self.sessions_inserted += counts.sessions_inserted;
3575        self.sessions_matched += counts.sessions_matched;
3576        self.messages_inserted_total += counts.messages_inserted_total;
3577        self.messages_inserted_searchable += counts.messages_inserted_searchable;
3578        self.messages_matched_total += counts.messages_matched_total;
3579        self.messages_matched_searchable += counts.messages_matched_searchable;
3580        self.parts_inserted += counts.parts_inserted;
3581        self.parts_matched += counts.parts_matched;
3582        self.inserted +=
3583            counts.sessions_inserted + counts.messages_inserted_total + counts.parts_inserted;
3584        self.matched +=
3585            counts.sessions_matched + counts.messages_matched_total + counts.parts_matched;
3586    }
3587
3588    /// Sum every counter from `other` into `self`. Used by the multi-source
3589    /// `pond sync` loop so adding a new field to this struct doesn't silently
3590    /// drop on aggregation - the prior hand-rolled `+=` block grew bugs.
3591    pub fn merge(&mut self, other: &Self) {
3592        self.inserted += other.inserted;
3593        self.matched += other.matched;
3594        self.sessions_inserted += other.sessions_inserted;
3595        self.messages_inserted_total += other.messages_inserted_total;
3596        self.messages_inserted_searchable += other.messages_inserted_searchable;
3597        self.parts_inserted += other.parts_inserted;
3598        self.sessions_matched += other.sessions_matched;
3599        self.messages_matched_total += other.messages_matched_total;
3600        self.messages_matched_searchable += other.messages_matched_searchable;
3601        self.parts_matched += other.parts_matched;
3602        self.dropped_events += other.dropped_events;
3603        self.dropped_sessions += other.dropped_sessions;
3604        self.skipped_files += other.skipped_files;
3605        self.skipped_empty += other.skipped_empty;
3606        self.skipped_fresh += other.skipped_fresh;
3607        self.skipped_superseded += other.skipped_superseded;
3608        self.storage_errors += other.storage_errors;
3609        self.truncated_values += other.truncated_values;
3610        for (key, value) in &other.drop_reasons {
3611            *self.drop_reasons.entry(key).or_insert(0) += value;
3612        }
3613    }
3614
3615    /// Same dispatch as [`Self::add_outcomes`] but ignores
3616    /// `Inserted`/`Matched` rows. The CLI batched path drives those counters
3617    /// via [`Self::add_batch`] and uses this method to attribute per-row
3618    /// `Error` outcomes from the same flush.
3619    pub fn add_outcomes_errors_only(&mut self, outcomes: &[RowOutcome]) {
3620        for outcome in outcomes {
3621            if !matches!(outcome.status, OutcomeStatus::Error) {
3622                continue;
3623            }
3624            if outcome.kind == "session" {
3625                self.dropped_sessions += 1;
3626            } else {
3627                self.dropped_events += 1;
3628            }
3629            let reason = outcome
3630                .error
3631                .as_ref()
3632                .and_then(|error| error.reason_key)
3633                .unwrap_or(DROP_REASON_UNCATEGORIZED);
3634            *self.drop_reasons.entry(reason).or_insert(0) += 1;
3635        }
3636    }
3637
3638    pub fn add_outcomes(&mut self, outcomes: &[RowOutcome]) {
3639        for outcome in outcomes {
3640            match outcome.status {
3641                OutcomeStatus::Inserted => {
3642                    self.inserted += 1;
3643                    match outcome.kind {
3644                        "session" => self.sessions_inserted += 1,
3645                        "message" => {
3646                            self.messages_inserted_total += 1;
3647                            if outcome.searchable {
3648                                self.messages_inserted_searchable += 1;
3649                            }
3650                        }
3651                        "part" => self.parts_inserted += 1,
3652                        _ => {}
3653                    }
3654                }
3655                OutcomeStatus::Matched => {
3656                    self.matched += 1;
3657                    match outcome.kind {
3658                        "session" => self.sessions_matched += 1,
3659                        "message" => {
3660                            self.messages_matched_total += 1;
3661                            if outcome.searchable {
3662                                self.messages_matched_searchable += 1;
3663                            }
3664                        }
3665                        "part" => self.parts_matched += 1,
3666                        _ => {}
3667                    }
3668                }
3669                OutcomeStatus::Error => {
3670                    // Session-level rejection: exactly one session-kind Error
3671                    // outcome (see `error_outcomes_for_substream`). Per-event
3672                    // drop: one Error per message/part. The two populations
3673                    // are counted separately so the operator can tell a
3674                    // structural reject from a row-level skip.
3675                    if outcome.kind == "session" {
3676                        self.dropped_sessions += 1;
3677                    } else {
3678                        self.dropped_events += 1;
3679                    }
3680                    let reason = outcome
3681                        .error
3682                        .as_ref()
3683                        .and_then(|e| e.reason_key)
3684                        .unwrap_or(DROP_REASON_UNCATEGORIZED);
3685                    *self.drop_reasons.entry(reason).or_insert(0) += 1;
3686                }
3687            }
3688        }
3689    }
3690}
3691
3692/// Per-row outcome surfaced by [`IngestValidator`] (spec.md#protocol). One
3693/// row per input event from the request's `events` array. The validator
3694/// returns these in array order so the wire layer can pack them directly
3695/// into [`crate::wire::IngestResult`] entries.
3696#[derive(Debug, Clone, PartialEq)]
3697pub struct RowOutcome {
3698    pub index: usize,
3699    pub kind: &'static str,
3700    pub pk: Value,
3701    pub status: OutcomeStatus,
3702    pub error: Option<RowError>,
3703    /// True iff `kind == "message"` AND the underlying row carries
3704    /// `search_text`. Drives `IngestSummary::messages_inserted_searchable`
3705    /// so the CLI can show "searchable" message deltas distinct from raw
3706    /// inserts. Always false for session/part rows.
3707    pub searchable: bool,
3708}
3709
3710#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3711pub enum OutcomeStatus {
3712    Inserted,
3713    Matched,
3714    Error,
3715}
3716
3717/// Structured per-row error body. Mirrors the wire shape so the handler
3718/// can pass it straight through.
3719#[derive(Debug, Clone, PartialEq, Eq)]
3720pub struct RowError {
3721    pub message: String,
3722    pub field: Option<&'static str>,
3723    pub reason: Option<&'static str>,
3724    /// Stable key for histogramming - see `DROP_REASON_*` constants. The
3725    /// `reason` field above is human-prose; `reason_key` is the machine
3726    /// bucket. `None` means uncategorized; consumers attribute to
3727    /// `DROP_REASON_UNCATEGORIZED`.
3728    pub reason_key: Option<&'static str>,
3729}
3730
3731/// Buffered session events tagged with their input array index, so the
3732/// per-row outcomes can be re-attributed once `merge_insert` returns its
3733/// per-row Inserted/Matched stats.
3734#[derive(Debug)]
3735struct BufferedSession {
3736    index: usize,
3737    session: Session,
3738}
3739
3740#[derive(Debug)]
3741struct BufferedMessage {
3742    index: usize,
3743    message: Message,
3744    parts: Vec<BufferedPart>,
3745    search_text: Option<String>,
3746}
3747
3748#[derive(Debug)]
3749struct BufferedPart {
3750    index: usize,
3751    part: Part,
3752}
3753
3754/// State machine that turns the `events: Vec<IngestEvent>` array into a
3755/// flat `Vec<RowOutcome>` matching the array's index space. Buffers a whole
3756/// session substream so `merge_insert` runs once per substream (three
3757/// batches: sessions, messages, parts). A validation error on a single event
3758/// drops *that event* (one [`OutcomeStatus::Error`] outcome) and the substream
3759/// continues; only Session-level invariants (immutable source_agent / project
3760/// on re-write) drop the whole substream (spec.md#adapter-integrity-event-ordering).
3761///
3762/// Writes are batched at flush time. As complete substreams arrive (a new
3763/// `Session` event closes out the current one), they accumulate in
3764/// `completed` rather than each one calling `merge_insert` immediately.
3765/// The caller drains the buffer via [`Self::flush`] / [`Self::finish`],
3766/// at which point one batched 3-parallel-merge-insert covers all pending
3767/// substreams. This is the load-bearing perf change: per-substream commit
3768/// overhead dominated the ingest profile (see `benches/ingest_bench.rs`),
3769/// and amortizing it across N sessions cuts wall time materially.
3770#[derive(Debug, Default)]
3771pub struct IngestValidator {
3772    session: Option<BufferedSession>,
3773    current_message: Option<BufferedMessage>,
3774    current_parts: Vec<BufferedPart>,
3775    messages: Vec<BufferedMessage>,
3776    /// Message ids already buffered in the current substream. Duplicate ids
3777    /// drop the offending event in-line rather than failing the whole batch
3778    /// downstream.
3779    seen_message_ids: HashSet<String>,
3780    /// `(message_id, part_id)` keys already buffered in the current
3781    /// substream. Same in-line duplicate-drop policy as `seen_message_ids`.
3782    seen_part_keys: HashSet<(String, String)>,
3783    /// Substreams whose end-of-stream boundary has been observed but whose
3784    /// rows haven't been written yet. Flushed in batched mode by
3785    /// [`Self::flush`].
3786    completed: Vec<CompletedSubstream>,
3787}
3788
3789/// One closed substream ready for the batched flush path.
3790#[derive(Debug)]
3791struct CompletedSubstream {
3792    session_index: usize,
3793    session: Session,
3794    messages: Vec<BufferedMessage>,
3795}
3796
3797/// Ingest host provenance (`options.pond`, spec.md#model-pond-options),
3798/// computed once per process. An audit fact - "the process that inserted this
3799/// row" - not identity. Fallible lookups are omitted, never synthesized as
3800/// placeholders.
3801fn ingest_host_stamp() -> Option<&'static Value> {
3802    static STAMP: std::sync::OnceLock<Option<Value>> = std::sync::OnceLock::new();
3803    STAMP
3804        .get_or_init(|| {
3805            let mut host = serde_json::Map::new();
3806            if let Ok(username) = whoami::username() {
3807                host.insert("username".to_owned(), username.into());
3808            }
3809            if let Ok(hostname) = whoami::hostname() {
3810                host.insert("hostname".to_owned(), hostname.into());
3811            }
3812            if let Ok(devicename) = whoami::devicename() {
3813                host.insert("device_name".to_owned(), devicename.into());
3814            }
3815            (!host.is_empty()).then(|| serde_json::json!({ "ingest": { "host": host } }))
3816        })
3817        .as_ref()
3818}
3819
3820impl IngestValidator {
3821    /// Drive one input event through the validator. Returns the per-row
3822    /// outcomes the event triggered: empty when the event is just buffered,
3823    /// or N entries when a session substream just flushed (success or
3824    /// failure). `Err` is reserved for catastrophic storage failures that
3825    /// should fail the whole `pond_ingest` request.
3826    pub async fn push(
3827        &mut self,
3828        store: &Store,
3829        index: usize,
3830        event: IngestEvent,
3831    ) -> Result<Vec<RowOutcome>> {
3832        match event {
3833            IngestEvent::Session(session) => self.push_session(store, index, session).await,
3834            IngestEvent::Message(message) => Ok(self.push_message(index, message)),
3835            IngestEvent::Part(part) => Ok(self.push_part(index, part)),
3836        }
3837    }
3838
3839    /// Final flush at end-of-batch. Closes the in-flight substream and
3840    /// drains the pending-flush buffer. Returns the per-row outcomes (for
3841    /// the wire layer) alongside the honest per-table counts (for
3842    /// `IngestSummary::add_batch`).
3843    pub async fn finish(&mut self, store: &Store) -> Result<(Vec<RowOutcome>, BatchCounts)> {
3844        self.close_current_substream();
3845        self.flush(store).await
3846    }
3847
3848    /// Drain every completed substream into batched 3-parallel-merge_insert
3849    /// writes. Caller invokes this periodically (every N completed
3850    /// substreams) to keep memory bounded; in adapter-driven sync that
3851    /// happens via the BATCH_SIZE check in `ingest_adapter`. The current
3852    /// in-flight substream stays buffered - close it explicitly via
3853    /// [`Self::finish`] or by feeding the next Session event.
3854    pub async fn flush(&mut self, store: &Store) -> Result<(Vec<RowOutcome>, BatchCounts)> {
3855        if self.completed.is_empty() {
3856            return Ok((Vec::new(), BatchCounts::default()));
3857        }
3858        let completed = std::mem::take(&mut self.completed);
3859        store.upsert_session_batch(completed).await
3860    }
3861
3862    /// Number of fully-buffered substreams awaiting batched write. Used by
3863    /// the adapter caller to decide when to call [`Self::flush`].
3864    pub fn pending_substreams(&self) -> usize {
3865        self.completed.len()
3866    }
3867
3868    async fn push_session(
3869        &mut self,
3870        _store: &Store,
3871        index: usize,
3872        mut session: Session,
3873    ) -> Result<Vec<RowOutcome>> {
3874        // Close out the current substream (if any) - move it to the pending
3875        // buffer instead of writing immediately. The actual write happens
3876        // when the caller invokes `flush` / `finish`.
3877        self.close_current_substream();
3878
3879        // spec.md#datasets: `source_agent` is trimmed at ingest and rejected
3880        // if empty after trim. A Session event with empty source_agent is
3881        // dropped on the spot - the substream that would follow has nothing
3882        // to anchor on, so subsequent message/part events will also drop.
3883        let trimmed = session.source_agent.trim();
3884        if trimmed.is_empty() {
3885            return Ok(vec![RowOutcome {
3886                index,
3887                kind: "session",
3888                pk: Value::String(session.id.clone()),
3889                status: OutcomeStatus::Error,
3890                error: Some(RowError {
3891                    message: format!("session {} has empty source_agent after trim", session.id),
3892                    field: Some("source_agent"),
3893                    reason: None,
3894                    reason_key: Some(DROP_REASON_EMPTY_SOURCE_AGENT),
3895                }),
3896                searchable: false,
3897            }]);
3898        }
3899        if trimmed.len() != session.source_agent.len() {
3900            session.source_agent = trimmed.to_owned();
3901        }
3902
3903        if session.parent_message_id.is_some() && session.parent_session_id.is_none() {
3904            return Ok(vec![RowOutcome {
3905                index,
3906                kind: "session",
3907                pk: Value::String(session.id.clone()),
3908                status: OutcomeStatus::Error,
3909                error: Some(RowError {
3910                    message: format!(
3911                        "session {} has parent_message_id without parent_session_id",
3912                        session.id,
3913                    ),
3914                    field: Some("parent_message_id"),
3915                    reason: None,
3916                    reason_key: Some(DROP_REASON_PARENT_MESSAGE_WITHOUT_SESSION),
3917                }),
3918                searchable: false,
3919            }]);
3920        }
3921
3922        self.seen_message_ids.clear();
3923        self.seen_part_keys.clear();
3924        self.session = Some(BufferedSession { index, session });
3925        Ok(Vec::new())
3926    }
3927
3928    fn close_current_substream(&mut self) {
3929        self.flush_current_message();
3930        let Some(BufferedSession {
3931            index: session_index,
3932            session,
3933        }) = self.session.take()
3934        else {
3935            return;
3936        };
3937        let messages = std::mem::take(&mut self.messages);
3938        self.seen_message_ids.clear();
3939        self.seen_part_keys.clear();
3940        self.completed.push(CompletedSubstream {
3941            session_index,
3942            session,
3943            messages,
3944        });
3945    }
3946
3947    fn push_message(&mut self, index: usize, mut message: Message) -> Vec<RowOutcome> {
3948        let pk = Value::Array(vec![
3949            Value::String(message.session_id().to_owned()),
3950            Value::String(message.id().to_owned()),
3951        ]);
3952        let Some(session) = &self.session else {
3953            return vec![error_outcome(
3954                index,
3955                "message",
3956                pk,
3957                "first event in a session stream must be Session",
3958                None,
3959                DROP_REASON_MESSAGE_BEFORE_SESSION,
3960            )];
3961        };
3962        if message.session_id() != session.session.id {
3963            let msg = format!(
3964                "message {} references session {}, expected {}",
3965                message.id(),
3966                message.session_id(),
3967                session.session.id
3968            );
3969            return vec![error_outcome(
3970                index,
3971                "message",
3972                pk,
3973                &msg,
3974                Some("session_id"),
3975                DROP_REASON_MESSAGE_SESSION_MISMATCH,
3976            )];
3977        }
3978        if !self.seen_message_ids.insert(message.id().to_owned()) {
3979            // Keep same-substream duplicate ids visible in `dropped_events`;
3980            // adapters are expected to dedupe upstream (see claude-code's
3981            // per-file `seen_uuids`), so a hit here is worth investigating.
3982            let msg = format!("duplicate message id {} in session substream", message.id());
3983            return vec![error_outcome(
3984                index,
3985                "message",
3986                pk,
3987                &msg,
3988                None,
3989                DROP_REASON_DUPLICATE_MESSAGE_ID,
3990            )];
3991        }
3992        // `options.pond` is core-owned (spec.md#model-pond-options): stripped
3993        // and restamped at ingest so neither adapters nor wire clients can
3994        // spoof provenance. Matched rows are merge_insert no-ops, so re-ingest
3995        // never restamps stored rows.
3996        match ingest_host_stamp() {
3997            Some(stamp) => {
3998                message
3999                    .options_mut()
4000                    .insert("pond".to_owned(), stamp.clone());
4001            }
4002            None => {
4003                message.options_mut().remove("pond");
4004            }
4005        }
4006        self.flush_current_message();
4007        self.current_message = Some(BufferedMessage {
4008            index,
4009            message,
4010            parts: Vec::new(),
4011            search_text: None,
4012        });
4013        Vec::new()
4014    }
4015
4016    fn push_part(&mut self, index: usize, part: Part) -> Vec<RowOutcome> {
4017        let pk = Value::Array(vec![
4018            Value::String(part.session_id.clone()),
4019            Value::String(part.message_id.clone()),
4020            Value::String(part.id.clone()),
4021        ]);
4022        let Some(current) = &self.current_message else {
4023            return vec![error_outcome(
4024                index,
4025                "part",
4026                pk,
4027                "part event appeared before a message",
4028                None,
4029                DROP_REASON_PART_BEFORE_MESSAGE,
4030            )];
4031        };
4032        if part.session_id != current.message.session_id() {
4033            let msg = format!(
4034                "part {} references session {}, expected {}",
4035                part.id,
4036                part.session_id,
4037                current.message.session_id()
4038            );
4039            return vec![error_outcome(
4040                index,
4041                "part",
4042                pk,
4043                &msg,
4044                Some("session_id"),
4045                DROP_REASON_PART_MESSAGE_MISMATCH,
4046            )];
4047        }
4048        if part.message_id != current.message.id() {
4049            let msg = format!(
4050                "part {} references message {}, expected {}",
4051                part.id,
4052                part.message_id,
4053                current.message.id()
4054            );
4055            return vec![error_outcome(
4056                index,
4057                "part",
4058                pk,
4059                &msg,
4060                Some("message_id"),
4061                DROP_REASON_PART_MESSAGE_MISMATCH,
4062            )];
4063        }
4064        let part_key = (part.message_id.clone(), part.id.clone());
4065        if !self.seen_part_keys.insert(part_key) {
4066            let msg = format!(
4067                "duplicate part id {} for message {} in session substream",
4068                part.id, part.message_id
4069            );
4070            return vec![error_outcome(
4071                index,
4072                "part",
4073                pk,
4074                &msg,
4075                None,
4076                DROP_REASON_DUPLICATE_PART_KEY,
4077            )];
4078        }
4079        self.current_parts.push(BufferedPart { index, part });
4080        Vec::new()
4081    }
4082
4083    fn flush_current_message(&mut self) {
4084        let Some(mut buffered) = self.current_message.take() else {
4085            return;
4086        };
4087        let parts = std::mem::take(&mut self.current_parts);
4088        let mut canonical_parts = Vec::with_capacity(parts.len());
4089        for part in &parts {
4090            canonical_parts.push(part.part.clone());
4091        }
4092        buffered.search_text = search_text(&buffered.message, &canonical_parts);
4093        buffered.parts = parts;
4094        self.messages.push(buffered);
4095    }
4096}
4097
4098fn error_outcome(
4099    index: usize,
4100    kind: &'static str,
4101    pk: Value,
4102    message: &str,
4103    field: Option<&'static str>,
4104    reason_key: &'static str,
4105) -> RowOutcome {
4106    RowOutcome {
4107        index,
4108        kind,
4109        pk,
4110        status: OutcomeStatus::Error,
4111        error: Some(RowError {
4112            message: message.to_owned(),
4113            field,
4114            reason: None,
4115            reason_key: Some(reason_key),
4116        }),
4117        searchable: false,
4118    }
4119}
4120
4121/// Session-level rejection (immutable `source_agent` / `project` violation):
4122/// emit exactly one Error outcome on the Session row. The buffered messages
4123/// and parts of this substream are *not* surfaced as per-row errors - their
4124/// loss is implied by the single session-rejection (spec.md#adapter-integrity-event-ordering).
4125fn error_outcomes_for_substream(
4126    session_index: usize,
4127    session: &Session,
4128    _messages: &[BufferedMessage],
4129    message: impl Into<String>,
4130    field: Option<&'static str>,
4131    reason_key: &'static str,
4132) -> Vec<RowOutcome> {
4133    let reason = field.map(|_| "immutable");
4134    vec![RowOutcome {
4135        index: session_index,
4136        kind: "session",
4137        pk: Value::String(session.id.clone()),
4138        status: OutcomeStatus::Error,
4139        error: Some(RowError {
4140            message: message.into(),
4141            field,
4142            reason,
4143            reason_key: Some(reason_key),
4144        }),
4145        searchable: false,
4146    }]
4147}
4148
4149/// Batched-path success helper. Each row's Inserted/Matched status is read
4150/// from the pre-existence sets captured by `upsert_session_batch` before its
4151/// `merge_insert` calls, so the per-row outcome is honest (spec.md#adapter-integrity-additive-sync).
4152/// Also accumulates the per-table totals into `counts` so the CLI summary
4153/// gets the same truth without re-walking the outcomes.
4154fn success_outcomes_for_substream(
4155    session_index: usize,
4156    session: &Session,
4157    messages: &[BufferedMessage],
4158    existing_sessions: &std::collections::HashMap<String, Session>,
4159    existing_message_pks: &HashSet<(String, String)>,
4160    existing_part_pks: &HashSet<(String, String, String)>,
4161    counts: &mut BatchCounts,
4162) -> Vec<RowOutcome> {
4163    let session_was_present = existing_sessions.contains_key(&session.id);
4164    let session_status = if session_was_present {
4165        counts.sessions_matched += 1;
4166        UpsertStatus::Matched
4167    } else {
4168        counts.sessions_inserted += 1;
4169        UpsertStatus::Inserted
4170    };
4171
4172    let mut outcomes = Vec::with_capacity(1 + messages.len());
4173    outcomes.push(success_outcome(
4174        session_index,
4175        "session",
4176        Value::String(session.id.clone()),
4177        session_status,
4178        false,
4179    ));
4180    for buffered in messages {
4181        let key = (
4182            buffered.message.session_id().to_owned(),
4183            buffered.message.id().to_owned(),
4184        );
4185        let searchable = buffered.search_text.is_some();
4186        let message_status = if existing_message_pks.contains(&key) {
4187            counts.messages_matched_total += 1;
4188            if searchable {
4189                counts.messages_matched_searchable += 1;
4190            }
4191            UpsertStatus::Matched
4192        } else {
4193            counts.messages_inserted_total += 1;
4194            if searchable {
4195                counts.messages_inserted_searchable += 1;
4196            }
4197            UpsertStatus::Inserted
4198        };
4199        let pk = Value::Array(vec![Value::String(key.0), Value::String(key.1)]);
4200        outcomes.push(success_outcome(
4201            buffered.index,
4202            "message",
4203            pk,
4204            message_status,
4205            searchable,
4206        ));
4207        for part in &buffered.parts {
4208            let part_key = (
4209                part.part.session_id.clone(),
4210                part.part.message_id.clone(),
4211                part.part.id.clone(),
4212            );
4213            let part_status = if existing_part_pks.contains(&part_key) {
4214                counts.parts_matched += 1;
4215                UpsertStatus::Matched
4216            } else {
4217                counts.parts_inserted += 1;
4218                UpsertStatus::Inserted
4219            };
4220            let part_pk = Value::Array(vec![
4221                Value::String(part_key.0),
4222                Value::String(part_key.1),
4223                Value::String(part_key.2),
4224            ]);
4225            outcomes.push(success_outcome(
4226                part.index,
4227                "part",
4228                part_pk,
4229                part_status,
4230                false,
4231            ));
4232        }
4233    }
4234    outcomes
4235}
4236
4237fn success_outcome(
4238    index: usize,
4239    kind: &'static str,
4240    pk: Value,
4241    status: UpsertStatus,
4242    searchable: bool,
4243) -> RowOutcome {
4244    let status = match status {
4245        UpsertStatus::Inserted => OutcomeStatus::Inserted,
4246        UpsertStatus::Matched => OutcomeStatus::Matched,
4247    };
4248    RowOutcome {
4249        index,
4250        kind,
4251        pk,
4252        status,
4253        error: None,
4254        searchable,
4255    }
4256}
4257
4258#[derive(Debug, Clone, PartialEq, Eq)]
4259enum IngestError {
4260    /// spec.md#protocol: `Session.source_agent` and `Session.project` are
4261    /// immutable post-first-write because the denormalized copies on
4262    /// `messages` were stamped from the prior Session at first ingest.
4263    /// A re-write that changes either would silently desync.
4264    ImmutableField {
4265        field: &'static str,
4266        session_id: String,
4267        stored: String,
4268        attempted: String,
4269    },
4270}
4271
4272impl std::fmt::Display for IngestError {
4273    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4274        match self {
4275            Self::ImmutableField {
4276                field,
4277                session_id,
4278                stored,
4279                attempted,
4280            } => write!(
4281                formatter,
4282                "session {session_id} {field} is immutable: stored {stored:?}, attempted {attempted:?}",
4283            ),
4284        }
4285    }
4286}
4287
4288impl std::error::Error for IngestError {}
4289
4290/// Compare an incoming Session row against the stored row on the two
4291/// immutable fields (spec.md#protocol). The `Option<String>` `project` field
4292/// counts a NULL-vs-non-NULL change as a mismatch.
4293fn ensure_immutable_match(
4294    existing: &Session,
4295    incoming: &Session,
4296) -> std::result::Result<(), IngestError> {
4297    if existing.source_agent != incoming.source_agent {
4298        return Err(IngestError::ImmutableField {
4299            field: "source_agent",
4300            session_id: incoming.id.clone(),
4301            stored: existing.source_agent.clone(),
4302            attempted: incoming.source_agent.clone(),
4303        });
4304    }
4305    if existing.project != incoming.project {
4306        return Err(IngestError::ImmutableField {
4307            field: "project",
4308            session_id: incoming.id.clone(),
4309            stored: (*existing.project).clone(),
4310            attempted: (*incoming.project).clone(),
4311        });
4312    }
4313    Ok(())
4314}
4315
4316pub fn search_text(message: &Message, parts: &[Part]) -> Option<String> {
4317    use crate::wire::Provenance;
4318    let mut chunks: Vec<String> = Vec::new();
4319    for part in parts {
4320        // spec.md#search: only conversational parts contribute to the indexed
4321        // text; harness-injected scaffolding is excluded from search.
4322        if part.provenance != Provenance::Conversational {
4323            continue;
4324        }
4325        match (message.role(), &part.kind) {
4326            (Role::User | Role::Assistant, PartKind::Text { text }) => {
4327                if let Some(text) = text {
4328                    chunks.push(text.to_string());
4329                }
4330            }
4331            (
4332                Role::User | Role::Assistant,
4333                PartKind::File {
4334                    media_type,
4335                    file_name,
4336                    data,
4337                },
4338            ) => {
4339                if let Some(file_name) = file_name {
4340                    chunks.push(file_name.clone());
4341                }
4342                if let Some(media_type) = media_type {
4343                    chunks.push(media_type.clone());
4344                }
4345                if let FileData::Url(uri) = data {
4346                    chunks.push(uri.clone());
4347                }
4348            }
4349            (
4350                Role::System | Role::Tool,
4351                PartKind::Text { .. }
4352                | PartKind::Reasoning { .. }
4353                | PartKind::File { .. }
4354                | PartKind::ToolCall { .. }
4355                | PartKind::ToolResult { .. }
4356                | PartKind::ToolApprovalRequest { .. }
4357                | PartKind::ToolApprovalResponse { .. },
4358            )
4359            | (
4360                Role::User | Role::Assistant,
4361                PartKind::Reasoning { .. }
4362                | PartKind::ToolCall { .. }
4363                | PartKind::ToolResult { .. }
4364                | PartKind::ToolApprovalRequest { .. }
4365                | PartKind::ToolApprovalResponse { .. },
4366            ) => {}
4367        }
4368    }
4369
4370    let text = chunks
4371        .into_iter()
4372        .filter(|chunk| !chunk.trim().is_empty())
4373        .collect::<Vec<_>>()
4374        .join("\n");
4375    if text.is_empty() { None } else { Some(text) }
4376}
4377
4378/// Non-empty conversational text (spec.md#search).
4379#[derive(Debug, Clone, PartialEq, Eq)]
4380pub struct SearchText(String);
4381
4382impl SearchText {
4383    pub fn as_str(&self) -> &str {
4384        &self.0
4385    }
4386
4387    pub fn into_inner(self) -> String {
4388        self.0
4389    }
4390}
4391
4392impl AsRef<str> for SearchText {
4393    fn as_ref(&self) -> &str {
4394        &self.0
4395    }
4396}
4397
4398#[derive(Debug, Clone, PartialEq)]
4399pub struct MessageWithParts {
4400    pub message: Message,
4401    pub parts: Vec<Part>,
4402}
4403
4404#[derive(Debug, Clone, PartialEq)]
4405pub struct SessionWithMessages {
4406    pub session: Session,
4407    pub messages: Vec<MessageWithParts>,
4408}
4409
4410#[derive(Debug, Clone)]
4411pub struct SessionViewParams<'a> {
4412    /// Inclusive anchor: the page starts at this id. Set only by the handler's
4413    /// message-to-session resolution; takes precedence over the other anchors.
4414    pub at_message_id: Option<&'a str>,
4415    /// Page forward: messages strictly after this id.
4416    pub after_message_id: Option<&'a str>,
4417    /// Page backward: messages strictly before this id.
4418    pub before_message_id: Option<&'a str>,
4419    pub limit: usize,
4420    pub budget_bytes: usize,
4421    /// First-page end when no anchor is set.
4422    pub session_from: SessionFrom,
4423}
4424
4425#[derive(Debug, Clone)]
4426pub struct MessageViewParams {
4427    /// Conversational siblings before the target (`grep -B`).
4428    pub context_before: usize,
4429    /// Conversational siblings after the target (`grep -A`).
4430    pub context_after: usize,
4431    pub budget_bytes: usize,
4432}
4433
4434/// Outcome of a get lookup. Separates a missing target (the handler
4435/// maps it to `not_found`) from a stale/unknown pagination anchor (mapped to
4436/// `validation_failed`): the message stream is append-only, so an anchor that
4437/// was ever valid never disappears - an unknown one is always a client error,
4438/// never a reason to silently restart the page.
4439#[derive(Debug, Clone, PartialEq)]
4440pub enum GetLookup<T> {
4441    NotFound,
4442    UnknownAnchor,
4443    Found(T),
4444}
4445
4446/// Canonical retrieval result for `pond_get_session`: the stored session
4447/// plus the page of messages (each with its `Part`s) and a remaining count.
4448/// Protocol-shaping into `GetResult`/`MessageView` happens in the handler.
4449#[derive(Debug, Clone, PartialEq)]
4450pub struct SessionPage {
4451    pub session: Session,
4452    pub messages: Vec<RetrievedMessage>,
4453    pub before_remaining: usize,
4454    pub after_remaining: usize,
4455}
4456
4457/// Canonical retrieval result for `pond_get_message`. `target.parts` is
4458/// empty - the target's parts ride `target_parts` (paginated); `siblings` carry
4459/// their parts so the handler can summarize them.
4460#[derive(Debug, Clone, PartialEq)]
4461pub struct MessagePage {
4462    pub session: Session,
4463    pub target: RetrievedMessage,
4464    pub target_parts: Vec<Part>,
4465    pub target_parts_remaining: usize,
4466    pub siblings: Vec<RetrievedMessage>,
4467}
4468
4469#[derive(Debug, Clone, PartialEq)]
4470pub struct RetrievedMessage {
4471    pub id: String,
4472    pub role: Role,
4473    pub timestamp: DateTime<Utc>,
4474    pub text: Option<String>,
4475    pub content: Option<String>,
4476    pub parts: Vec<Part>,
4477}
4478
4479#[derive(Debug, Clone)]
4480struct ScanRow {
4481    id: String,
4482    role: Role,
4483    timestamp: DateTime<Utc>,
4484    text: Option<String>,
4485    content: Option<String>,
4486}
4487
4488/// One row of the conversational scan. `text` is non-empty by
4489/// `IsNotNull("search_text")` pushdown (spec.md#search).
4490#[derive(Debug, Clone)]
4491pub struct ConversationalRow {
4492    pub session_id: String,
4493    pub message_id: String,
4494    pub role: Role,
4495    pub timestamp: DateTime<Utc>,
4496    pub text: SearchText,
4497}
4498
4499/// Number of leading `items` that fit within `limit` and the byte budget,
4500/// sizing each by `size`. Always emits at least one (a single oversize item
4501/// never blocks its own page); the budget then stops the page at the next item
4502/// boundary.
4503fn page_by<T>(items: &[T], limit: usize, budget_bytes: usize, size: impl Fn(&T) -> usize) -> usize {
4504    let capped = items.len().min(limit.clamp(1, 1000));
4505    let mut acc = 0usize;
4506    let mut emitted = 0usize;
4507    for item in &items[..capped] {
4508        let next = acc.saturating_add(size(item));
4509        if emitted > 0 && next > budget_bytes {
4510            break;
4511        }
4512        acc = next;
4513        emitted += 1;
4514    }
4515    emitted
4516}
4517
4518/// Like `page_by` but counts from the tail: how many trailing items fit
4519/// `limit` and the byte budget, dropping oldest first. The last (newest) item
4520/// is always kept, so the returned count is >= 1 for a non-empty slice and the
4521/// emitted page (`items[len - n..]`) stays chronological.
4522fn page_tail<T>(
4523    items: &[T],
4524    limit: usize,
4525    budget_bytes: usize,
4526    size: impl Fn(&T) -> usize,
4527) -> usize {
4528    let cap = limit.clamp(1, 1000);
4529    let mut acc = 0usize;
4530    let mut emitted = 0usize;
4531    for item in items.iter().rev() {
4532        if emitted >= cap {
4533            break;
4534        }
4535        let next = acc.saturating_add(size(item));
4536        if emitted > 0 && next > budget_bytes {
4537            break;
4538        }
4539        acc = next;
4540        emitted += 1;
4541    }
4542    emitted
4543}
4544
4545fn role_from_str(value: &str) -> Result<Role> {
4546    match value {
4547        "system" => Ok(Role::System),
4548        "user" => Ok(Role::User),
4549        "assistant" => Ok(Role::Assistant),
4550        "tool" => Ok(Role::Tool),
4551        other => anyhow::bail!("unknown message role {other}"),
4552    }
4553}
4554
4555/// Scalar indexes on `messages` (spec.md#datasets): only columns whose index
4556/// type matches the predicate actually issued against them. `project` is
4557/// filtered solely by `LikeContains`/`Regex` (substring), which a BTree cannot
4558/// accelerate, and `role` is never filtered - both are deliberately unindexed
4559/// (substring lookup stays on the SQL `LIKE` path). There is no index on
4560/// `embedding_model`: pond's invariant is one active model at a time (a model
4561/// swap goes through `pond optimize --force-embed` which drops the IVF_SQ,
4562/// clears stale rows, and re-bootstraps), so the only embedding-state filter is
4563/// `vector IS NOT NULL`. `id` lookups are rare and full-scan. Do NOT add a
4564/// ZoneMap on `timestamp`: it prunes every zone for the tz-aware column, so
4565/// date filters return empty (#75) - an upstream `safe_coerce_scalar` tz drop
4566/// that no literal form escapes. Date bounds run as a refine over the arm pool.
4567const MESSAGE_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] = &[
4568    (
4569        "session_id",
4570        BuiltinIndexType::BTree,
4571        MESSAGES_SESSION_ID_INDEX,
4572    ),
4573    (
4574        "source_agent",
4575        BuiltinIndexType::Bitmap,
4576        "messages_source_agent_bitmap",
4577    ),
4578];
4579
4580/// Scalar indexes on `parts`: `(session_id, message_id)` is the hot-path lookup key for
4581/// `parts_for_messages` (hydration on every `get` and grouped search). `tool_name`
4582/// serves the #89 analytics filters; BTree, not Bitmap, despite the categorical
4583/// shape - prefix LIKE (`tool_name LIKE 'mcp__%'`) errors on bitmap indexes
4584/// (same upstream limitation documented for `messages.source_agent`).
4585const PARTS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] = &[
4586    (
4587        "session_id",
4588        BuiltinIndexType::BTree,
4589        "parts_session_id_btree",
4590    ),
4591    (
4592        "message_id",
4593        BuiltinIndexType::BTree,
4594        "parts_message_id_btree",
4595    ),
4596    (
4597        "tool_name",
4598        BuiltinIndexType::BTree,
4599        "parts_tool_name_btree",
4600    ),
4601];
4602
4603/// Scalar index on `sessions`: `id` is filtered by `find_session` on every
4604/// `get` and every grouped search.
4605const SESSIONS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] =
4606    &[("id", BuiltinIndexType::BTree, "sessions_id_btree")];
4607
4608/// Session ids per `session_id IN (...)` chunk in an incremental copy: large
4609/// enough to amortize per-scan setup, small enough to keep the pushed-down
4610/// predicate string and its btree lookup batch bounded.
4611const COPY_SESSION_IN_CHUNK: usize = 512;
4612
4613fn in_predicate(column: &'static str, values: &[String]) -> Predicate {
4614    Predicate::In(
4615        column,
4616        values.iter().cloned().map(ScalarValue::String).collect(),
4617    )
4618}
4619
4620/// The kNN prefilter is the caller's scalar filter alone - pond does NOT add
4621/// `vector IS NOT NULL`. That looks like a safe guard but it is a remote-read
4622/// trap: Lance v2 keeps no per-column null metadata, so `IsNotNull(vector)`
4623/// forces a full read of the ~3 GiB `vector` column from the object store on
4624/// every query (the ANN prefilter is evaluated as a `LanceScan` over the
4625/// column) - measured at ~57 s/query on the 2M-row S3 corpus, dwarfing the
4626/// IVF probe itself. It is also redundant: the IVF_SQ index only contains
4627/// embedded rows, and Lance's `_distance IS NOT NULL` post-filter (present in
4628/// both the ANN and brute-force branches of the plan) already drops any
4629/// null-vector row the brute-force tail might surface. So an empty caller
4630/// filter yields an empty prefilter and a pure index probe (spec.md#search,
4631/// spec.md#search-prefilter-pushdown).
4632fn embedded_scope(filter: &Predicate) -> Predicate {
4633    filter.clone()
4634}
4635
4636/// IVF `nprobes` applied when `[search].nprobes` is unset. Left unset, Lance
4637/// probes up to every partition (~num_rows/4096, ~500 on the 2M-row corpus),
4638/// one object-store read each - the dominant cost of a vector scan on a remote
4639/// store. 32 bounds the reads while keeping recall (benchmarked, spec.md#search).
4640pub const DEFAULT_NPROBES: usize = 32;
4641
4642/// Apply pond's vector-search tuning to a kNN scanner, defaulting any unset
4643/// `[search]` knob so a default install never inherits Lance's unbounded
4644/// probe-every-partition behavior. No refine: IVF_SQ's per-dimension codes are
4645/// precise enough to rank from the prewarmed partition, so pond never re-reads
4646/// exact vectors from the data files (the remote-store GET storm PQ+refine
4647/// incurred).
4648fn apply_vector_search_knobs(
4649    scanner: &mut lance::dataset::scanner::Scanner,
4650    search: Option<&config::SearchConfig>,
4651) {
4652    let nprobes = search
4653        .and_then(|cfg| cfg.nprobes)
4654        .unwrap_or(DEFAULT_NPROBES);
4655    scanner.nprobes(nprobes);
4656}
4657
4658// Bare logical table names: the lance-namespace Directory impl owns the
4659// `.lance` directory suffix (spec.md#lance-chokepoints-catalog). No consumer reconstructs
4660// a `.lance` path.
4661pub(crate) const SESSIONS: &str = "sessions";
4662pub(crate) const MESSAGES: &str = "messages";
4663pub(crate) const PARTS: &str = "parts";
4664
4665/// BTree index name on `messages.session_id` (spec.md#datasets). Stable so
4666/// index creation, status, and the scalar-fold gate name the same index.
4667pub const MESSAGES_SESSION_ID_INDEX: &str = "messages_session_id_btree";
4668
4669/// FTS index name on `messages.search_text`. Stable so status and index
4670/// creation name the same index.
4671pub const MESSAGES_FTS_INDEX: &str = "messages_search_text_fts";
4672
4673/// IVF_SQ index name on `messages.vector` (spec.md#search). Stable so the
4674/// activation check, optimize/append, and status all name the same index. The
4675/// literal keeps the historical `_ivfpq` suffix as a stable identifier:
4676/// renaming it would orphan the existing segment under a new name. A plain
4677/// `optimize` folds into whatever index type already exists, so switching an
4678/// existing IVF_PQ store to IVF_SQ needs `pond optimize --rebuild`.
4679pub const MESSAGES_VECTOR_INDEX: &str = "messages_vector_ivfpq";
4680
4681/// IVF_SQ tuning constants (spec.md#search):
4682/// - num_bits = 8 (per-dimension scalar quantization)
4683/// - max_iters = 15 (kmeans cap)
4684/// - cosine metric (e5 vectors are L2-normalized)
4685const IVF_SQ_NUM_BITS: u16 = 8;
4686const IVF_SQ_MAX_ITERS: usize = 15;
4687
4688/// Pond's production IndexIntents: the per-table intent set
4689/// `Store::open_with_options` registers with the substrate.
4690pub fn pond_index_intents() -> IndexIntents {
4691    pond_index_intents_with_vector_threshold(VECTOR_INDEX_ACTIVATION_ROWS)
4692}
4693
4694/// Same as [`pond_index_intents`] but with an overridable IVF_SQ activation
4695/// threshold. Used by tests that need to exercise the activation boundary
4696/// without writing 100k vectors.
4697pub(crate) fn pond_index_intents_with_vector_threshold(vector_threshold: usize) -> IndexIntents {
4698    let mut messages = Vec::with_capacity(MESSAGE_SCALAR_INDICES.len() + 2);
4699    messages.push(IndexIntent {
4700        name: MESSAGES_FTS_INDEX,
4701        column: "search_text",
4702        trigger: IndexTrigger::OnAnyRows,
4703        params: IndexParamsKind::InvertedFtsWord,
4704    });
4705    for (column, kind, name) in MESSAGE_SCALAR_INDICES {
4706        messages.push(IndexIntent {
4707            name,
4708            column,
4709            trigger: IndexTrigger::OnAnyRows,
4710            params: IndexParamsKind::Scalar(kind.clone()),
4711        });
4712    }
4713    messages.push(IndexIntent {
4714        name: MESSAGES_VECTOR_INDEX,
4715        column: "vector",
4716        trigger: IndexTrigger::OnNonNullCount {
4717            column: "vector",
4718            threshold: vector_threshold,
4719        },
4720        params: IndexParamsKind::IvfSqCosine {
4721            num_bits: IVF_SQ_NUM_BITS,
4722            max_iters: IVF_SQ_MAX_ITERS,
4723        },
4724    });
4725    let parts = PARTS_SCALAR_INDICES
4726        .iter()
4727        .map(|(column, kind, name)| IndexIntent {
4728            name,
4729            column,
4730            trigger: IndexTrigger::OnAnyRows,
4731            params: IndexParamsKind::Scalar(kind.clone()),
4732        })
4733        .collect();
4734    let sessions = SESSIONS_SCALAR_INDICES
4735        .iter()
4736        .map(|(column, kind, name)| IndexIntent {
4737            name,
4738            column,
4739            trigger: IndexTrigger::OnAnyRows,
4740            params: IndexParamsKind::Scalar(kind.clone()),
4741        })
4742        .collect();
4743    IndexIntents {
4744        sessions,
4745        messages,
4746        parts,
4747    }
4748}
4749
4750/// Default width of the `messages.vector` embedding column (spec.md#search):
4751/// matches [`embed::DEFAULT_MODEL_ID`] (`intfloat/multilingual-e5-small`,
4752/// 384). Used when `[embeddings].dim` is absent.
4753pub const DEFAULT_EMBEDDING_DIM: usize = 384;
4754
4755/// Process-wide vector dimension, seeded once at startup from `[embeddings].dim`
4756/// via [`init_embedding_dim`]. `OnceLock` (not `const`) so a temporary config
4757/// file can pick a different-dim model (e.g. e5-small at 384) for an experiment
4758/// without touching every site. Uninitialized -> [`DEFAULT_EMBEDDING_DIM`],
4759/// which keeps unit tests config-free.
4760static EMBEDDING_DIM_RUNTIME: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4761
4762/// The active embedding dimension. Returns whatever [`init_embedding_dim`]
4763/// installed, or [`DEFAULT_EMBEDDING_DIM`] when nothing has installed one.
4764pub fn embedding_dim() -> usize {
4765    EMBEDDING_DIM_RUNTIME
4766        .get()
4767        .copied()
4768        .unwrap_or(DEFAULT_EMBEDDING_DIM)
4769}
4770
4771/// Seed [`embedding_dim`] from config. First call wins.
4772pub fn init_embedding_dim(dim: usize) {
4773    EMBEDDING_DIM_RUNTIME.get_or_init(|| dim);
4774}
4775
4776/// Initial-`CREATE` write params for the namespace-mediated path. The
4777/// substrate seam stamps in `session`, `mode`, and `store_params`.
4778/// `auto_cleanup` is short; long-term recovery is `pond copy --to <file>`
4779/// snapshots plus deferred Lance tags (spec.md#session-durable-copy).
4780/// `skip_auto_cleanup` suppresses the per-commit hook so cleanup stays
4781/// operator-driven via `pond optimize` (one LIST per command instead of per write).
4782pub(crate) fn write_params_for_create() -> WriteParams {
4783    WriteParams {
4784        data_storage_version: Some(LanceFileVersion::V2_1),
4785        enable_v2_manifest_paths: true,
4786        enable_stable_row_ids: true,
4787        auto_cleanup: Some(AutoCleanupParams {
4788            interval: 20,
4789            older_than: chrono::TimeDelta::days(1),
4790        }),
4791        skip_auto_cleanup: true,
4792        ..WriteParams::default()
4793    }
4794}
4795
4796fn export_schema(table: Table) -> Arc<Schema> {
4797    match table {
4798        Table::Sessions => session_schema(),
4799        Table::Messages => message_schema(),
4800        Table::Parts => part_schema(),
4801    }
4802}
4803
4804/// Decide how an archive table's schema relates to this build's: identical ->
4805/// import verbatim (`None`); missing exactly a derivable set of nullable
4806/// columns (the archive predates an additive schema change) -> the recipe to
4807/// derive them per batch; anything else -> a hard error naming the version fix.
4808fn archive_schema_backfill(dataset: &Dataset, table: Table) -> Result<Option<ColumnBackfill>> {
4809    use std::collections::BTreeSet;
4810    let expected = export_schema(table);
4811    let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
4812    let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
4813    let expected_names: BTreeSet<&str> = expected
4814        .fields()
4815        .iter()
4816        .map(|f| f.name().as_str())
4817        .collect();
4818    let extra: Vec<&str> = actual_names.difference(&expected_names).copied().collect();
4819    if !extra.is_empty() {
4820        anyhow::bail!(
4821            "{} archive table has columns {actual_names:?} but this pond build expects \
4822             {expected_names:?} - the archive was written by a newer pond; upgrade pond \
4823             to restore it",
4824            table.as_str(),
4825        );
4826    }
4827    let missing: Vec<Field> = expected
4828        .fields()
4829        .iter()
4830        .filter(|f| !actual_names.contains(f.name().as_str()))
4831        .map(|f| f.as_ref().clone())
4832        .collect();
4833    if missing.is_empty() {
4834        return Ok(None);
4835    }
4836    column_backfill(table.as_str(), &missing)
4837        .map(Some)
4838        .with_context(|| {
4839            format!(
4840                "{} archive predates a schema change this pond build cannot bridge; \
4841                 restore it with the pond version that wrote it",
4842                table.as_str(),
4843            )
4844        })
4845}
4846
4847/// Extend an older-schema `batch` with the cells `spec` derives from it. The
4848/// derived columns append after the scanned ones as-is: a scan may
4849/// auto-convert JSON columns to their Arrow text form, and aligning batch
4850/// columns to the stored schema (by name, with that conversion) is the merge
4851/// path's job, not this function's.
4852fn upgraded_batch(batch: &RecordBatch, spec: &ColumnBackfill) -> Result<RecordBatch> {
4853    let derived = (spec.mapper)(batch)?;
4854    let mut fields: Vec<Field> = batch
4855        .schema()
4856        .fields()
4857        .iter()
4858        .map(|field| field.as_ref().clone())
4859        .collect();
4860    let mut columns = batch.columns().to_vec();
4861    for (field, column) in spec.output_schema.fields().iter().zip(derived.columns()) {
4862        fields.push(field.as_ref().clone());
4863        columns.push(column.clone());
4864    }
4865    RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
4866        .context("failed to assemble upgraded batch")
4867}
4868
4869async fn open_archive_table(table: Table, source: &Path) -> Result<Dataset> {
4870    let source_uri = source
4871        .to_str()
4872        .with_context(|| format!("archive path is not UTF-8: {}", source.display()))?;
4873    Dataset::open(source_uri)
4874        .await
4875        .with_context(|| format!("failed to open {} archive table", table.as_str()))
4876}
4877
4878/// The composite primary-key columns each table's schema declares
4879/// (spec.md#lance-table-creation-session-scoped-pk): a message/part id is
4880/// unique only within its session, so the key leads with `session_id`. The one
4881/// source of truth for the PK structure - kept beside the schemas it mirrors,
4882/// not in the schema-agnostic substrate seam.
4883pub(crate) fn pk_columns(table: Table) -> &'static [&'static str] {
4884    match table {
4885        Table::Sessions => &["id"],
4886        Table::Messages => &["session_id", "id"],
4887        Table::Parts => &["session_id", "message_id", "id"],
4888    }
4889}
4890
4891/// One scanned row's composite primary key as owned strings, in `pk` order.
4892fn composite_key(batch: &RecordBatch, pk: &[&str], row: usize) -> Result<Vec<String>> {
4893    pk.iter()
4894        .map(|column| string(batch, column, row)?.with_context(|| format!("{column} is null")))
4895        .collect()
4896}
4897
4898pub(crate) fn session_schema() -> Arc<Schema> {
4899    Arc::new(Schema::new(vec![
4900        primary_field("id", DataType::Utf8, false),
4901        Field::new("parent_session_id", DataType::Utf8, true),
4902        Field::new("parent_message_id", DataType::Utf8, true),
4903        Field::new("source_agent", DataType::Utf8, false),
4904        Field::new(
4905            "created_at",
4906            DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4907            false,
4908        ),
4909        Field::new("project", DataType::Utf8, false),
4910        json_field("options", false),
4911    ]))
4912}
4913
4914pub(crate) fn message_schema() -> Arc<Schema> {
4915    Arc::new(Schema::new(vec![
4916        primary_field("session_id", DataType::Utf8, false),
4917        primary_field("id", DataType::Utf8, false),
4918        Field::new(
4919            "timestamp",
4920            DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4921            false,
4922        ),
4923        Field::new("role", DataType::Utf8, false),
4924        Field::new("source_agent", DataType::Utf8, false),
4925        Field::new("project", DataType::Utf8, false),
4926        Field::new("content", DataType::Utf8, true),
4927        Field::new("search_text", DataType::Utf8, true),
4928        // The message's derived embedding (spec.md#session-embed-from-canonical):
4929        // filled inline at ingest when embedding is on, else null until a later
4930        // `pond optimize` embed pass; `vector` and `embedding_model` set together.
4931        Field::new("vector", embedding_vector_type(), true),
4932        Field::new("embedding_model", DataType::Utf8, true),
4933        json_field("options", false),
4934    ]))
4935}
4936
4937/// Derives one batch of backfill cells from a batch of stored columns.
4938pub(crate) type BackfillMapper = Box<dyn Fn(&RecordBatch) -> Result<RecordBatch> + Send + Sync>;
4939
4940/// One table's recipe for the additive-schema backfill: which stored columns
4941/// to read and how to derive the missing ones. The substrate open path feeds
4942/// this to `Dataset::add_columns`, so an existing store upgrades in place -
4943/// never by re-ingest, which cannot recover sessions whose sources are gone
4944/// (spec.md#session-durable-copy).
4945pub(crate) struct ColumnBackfill {
4946    pub read_columns: Vec<String>,
4947    pub output_schema: Arc<Schema>,
4948    pub mapper: BackfillMapper,
4949}
4950
4951/// Build the recipe deriving `missing` for `table_name`, or fail when any of
4952/// the columns cannot be derived from stored data.
4953pub(crate) fn column_backfill(table_name: &str, missing: &[Field]) -> Result<ColumnBackfill> {
4954    match table_name {
4955        PARTS => parts_backfill(missing),
4956        other => anyhow::bail!(
4957            "table {other} is missing columns {:?} that this pond build cannot derive from \
4958             stored data - use the pond version that wrote it",
4959            missing.iter().map(|f| f.name()).collect::<Vec<_>>(),
4960        ),
4961    }
4962}
4963
4964fn parts_backfill(missing: &[Field]) -> Result<ColumnBackfill> {
4965    for field in missing {
4966        anyhow::ensure!(
4967            ["tool_name", "call_id", "is_failure"].contains(&field.name().as_str())
4968                && field.is_nullable(),
4969            "column {PARTS}.{} cannot be derived from stored data - use the pond version \
4970             that wrote it",
4971            field.name(),
4972        );
4973    }
4974    let output_schema = Arc::new(Schema::new(missing.to_vec()));
4975    let schema = output_schema.clone();
4976    let mapper = move |batch: &RecordBatch| -> Result<RecordBatch> {
4977        let rows = batch.num_rows();
4978        let mut names: Vec<Option<String>> = Vec::with_capacity(rows);
4979        let mut call_ids: Vec<Option<String>> = Vec::with_capacity(rows);
4980        let mut failures: Vec<Option<bool>> = Vec::with_capacity(rows);
4981        for row in 0..rows {
4982            let type_name = string(batch, "type", row)?.context("part type is null")?;
4983            // Only tool parts carry identity; skipping the JSONB decode for
4984            // text/reasoning rows keeps the backfill pass cheap.
4985            let kind = match type_name.as_str() {
4986                "tool_call" | "tool_result" | "tool_approval_request" => {
4987                    let body =
4988                        json_column(batch, "variant_data", row)?.context("variant_data is null")?;
4989                    // An undecodable body degrades to NULL cells (spec.md#model-no-synthesis:
4990                    // a cell the stored record cannot justify stays NULL) rather than failing
4991                    // the migration, which re-runs on every open and would leave the store
4992                    // permanently un-openable over one bad row.
4993                    match part_kind_from_json(&type_name, &body, None) {
4994                        Ok(kind) => Some(kind),
4995                        Err(error) => {
4996                            tracing::warn!(
4997                                session_id = string(batch, "session_id", row)?.as_deref(),
4998                                part_id = string(batch, "id", row)?.as_deref(),
4999                                error = %format!("{error:#}"),
5000                                "stored tool part body failed to decode; its backfilled \
5001                                 cells stay NULL",
5002                            );
5003                            None
5004                        }
5005                    }
5006                }
5007                _ => None,
5008            };
5009            let (name, call_id, is_failure) = kind
5010                .as_ref()
5011                .map(tool_identity)
5012                .unwrap_or((None, None, None));
5013            names.push(name.map(str::to_owned));
5014            call_ids.push(call_id.map(str::to_owned));
5015            failures.push(is_failure);
5016        }
5017        let arrays: Vec<ArrayRef> = schema
5018            .fields()
5019            .iter()
5020            .map(|field| -> ArrayRef {
5021                match field.name().as_str() {
5022                    "tool_name" => Arc::new(StringArray::from(names.clone())),
5023                    "call_id" => Arc::new(StringArray::from(call_ids.clone())),
5024                    _ => Arc::new(BooleanArray::from(failures.clone())),
5025                }
5026            })
5027            .collect();
5028        RecordBatch::try_new(schema.clone(), arrays).context("failed to build backfill batch")
5029    };
5030    Ok(ColumnBackfill {
5031        read_columns: vec![
5032            "session_id".to_owned(),
5033            "id".to_owned(),
5034            "type".to_owned(),
5035            "variant_data".to_owned(),
5036        ],
5037        output_schema,
5038        mapper: Box::new(mapper),
5039    })
5040}
5041
5042pub(crate) fn part_schema() -> Arc<Schema> {
5043    Arc::new(Schema::new(vec![
5044        primary_field("session_id", DataType::Utf8, false),
5045        primary_field("message_id", DataType::Utf8, false),
5046        primary_field("id", DataType::Utf8, false),
5047        Field::new("ordinal", DataType::Int32, false),
5048        Field::new("type", DataType::Utf8, false),
5049        // spec.md#model-part-provenance: conversation vs harness-injected; search
5050        // reads this column to exclude injected scaffolding.
5051        Field::new("provenance", DataType::Utf8, false),
5052        // Materialized copies of the tool-part identity fields (#89): analytics
5053        // must run on narrow native columns - a JSON getter over `variant_data`
5054        // reads the whole multi-GB column, which times out on object stores.
5055        // NULL for non-tool parts and absent source fields (spec.md#model-no-synthesis).
5056        Field::new("tool_name", DataType::Utf8, true),
5057        Field::new("call_id", DataType::Utf8, true),
5058        Field::new("is_failure", DataType::Boolean, true),
5059        json_field("variant_data", false),
5060        legacy_blob_field("data", true),
5061        json_field("options", false),
5062    ]))
5063}
5064
5065pub(crate) fn empty_batch(schema: Arc<Schema>) -> Result<RecordBatch> {
5066    let arrays = schema
5067        .fields()
5068        .iter()
5069        .map(|field| lance::deps::arrow_array::new_empty_array(field.data_type()))
5070        .collect();
5071    RecordBatch::try_new(schema, arrays).context("failed to build empty Lance batch")
5072}
5073
5074pub(crate) fn empty_reader(
5075    schema: Arc<Schema>,
5076) -> Result<
5077    RecordBatchIterator<
5078        std::vec::IntoIter<Result<RecordBatch, lance::deps::arrow_schema::ArrowError>>,
5079    >,
5080> {
5081    let batch = empty_batch(schema.clone())?;
5082    Ok(RecordBatchIterator::new(
5083        vec![Ok(batch)].into_iter(),
5084        schema,
5085    ))
5086}
5087
5088pub(crate) struct MessageBatchRow<'a> {
5089    pub message: &'a Message,
5090    pub source_agent: &'a str,
5091    pub project: &'a str,
5092    pub search_text: Option<&'a str>,
5093}
5094
5095// Lance v7.0.0-beta.16's IVF_SQ build path (`rust/lance/src/index/vector/utils.rs`
5096// `infer_vector_element_type_impl`) accepts only Float16/Float32/Float64/UInt8/Int8;
5097// `FixedSizeBinary(2)`-backed `lance.bfloat16` is rejected. The format docs list
5098// BFloat16 as a future-supported embedding type; until the Rust IVF_SQ build
5099// path catches up, store as Float16 (half-precision, also 2 bytes/element).
5100fn embedding_vector_type() -> DataType {
5101    DataType::FixedSizeList(
5102        Arc::new(Field::new("item", DataType::Float16, true)),
5103        embedding_dim() as i32,
5104    )
5105}
5106
5107/// The partial-schema source for the embedding column update: the `messages`
5108/// primary key plus the two columns `pond optimize` fills. The field definitions
5109/// match `message_schema` exactly so Lance accepts it as a subset upsert.
5110fn embedding_update_schema() -> Arc<Schema> {
5111    Arc::new(Schema::new(vec![
5112        primary_field("session_id", DataType::Utf8, false),
5113        primary_field("id", DataType::Utf8, false),
5114        Field::new("vector", embedding_vector_type(), true),
5115        Field::new("embedding_model", DataType::Utf8, true),
5116    ]))
5117}
5118
5119/// The `messages` `vector` + `embedding_model` columns for an inline-embed
5120/// batch: `Some` rows carry the embedding and the current model id, `None` rows
5121/// are null in both. Returned aligned to `vectors` for [`messages_chunk`].
5122fn embedding_columns(vectors: &[Option<Vec<f32>>]) -> Result<(ArrayRef, ArrayRef)> {
5123    let dim = embedding_dim();
5124    // The common case (no embedder, or every row already present) is all-null:
5125    // build both columns with one bulk allocation instead of dim per-row appends.
5126    if vectors.iter().all(Option::is_none) {
5127        return Ok((
5128            new_null_array(&embedding_vector_type(), vectors.len()),
5129            new_null_array(&DataType::Utf8, vectors.len()),
5130        ));
5131    }
5132    let mut builder = FixedSizeListBuilder::new(
5133        Float16Builder::with_capacity(vectors.len() * dim),
5134        dim as i32,
5135    )
5136    .with_field(Arc::new(Field::new("item", DataType::Float16, true)));
5137    let mut models: Vec<Option<&str>> = Vec::with_capacity(vectors.len());
5138    for vector in vectors {
5139        match vector {
5140            Some(values) => {
5141                if values.len() != dim {
5142                    anyhow::bail!("inline embedding has dim {}, expected {dim}", values.len());
5143                }
5144                for value in values {
5145                    builder.values().append_value(half::f16::from_f32(*value));
5146                }
5147                builder.append(true);
5148                models.push(Some(embed::model_id()));
5149            }
5150            None => {
5151                for _ in 0..dim {
5152                    builder.values().append_null();
5153                }
5154                builder.append(false);
5155                models.push(None);
5156            }
5157        }
5158    }
5159    Ok((
5160        Arc::new(builder.finish()) as ArrayRef,
5161        Arc::new(StringArray::from(models)) as ArrayRef,
5162    ))
5163}
5164
5165/// Build the merge-update source batch for [`Store::write_embeddings`]: one row
5166/// per embedded message carrying `(session_id, id, vector, embedding_model)`.
5167pub(crate) fn embedding_update_batch(rows: &[EmbeddedMessage]) -> Result<RecordBatch> {
5168    let dim = embedding_dim();
5169    let mut flat = Vec::with_capacity(rows.len() * dim);
5170    for row in rows {
5171        if row.vector.len() != dim {
5172            anyhow::bail!(
5173                "embedding for message {} has dim {}, expected {dim}",
5174                row.id,
5175                row.vector.len(),
5176            );
5177        }
5178        flat.extend(row.vector.iter().map(|value| half::f16::from_f32(*value)));
5179    }
5180    let values = Float16Array::from(flat);
5181    let item_field = Arc::new(Field::new("item", DataType::Float16, true));
5182    let vectors = FixedSizeListArray::try_new(item_field, dim as i32, Arc::new(values), None)
5183        .context("failed to build embedding vector column")?;
5184
5185    RecordBatch::try_new(
5186        embedding_update_schema(),
5187        vec![
5188            Arc::new(StringArray::from(
5189                rows.iter()
5190                    .map(|row| row.session_id.as_str())
5191                    .collect::<Vec<_>>(),
5192            )),
5193            Arc::new(StringArray::from(
5194                rows.iter().map(|row| row.id.as_str()).collect::<Vec<_>>(),
5195            )),
5196            Arc::new(vectors),
5197            Arc::new(StringArray::from(vec![embed::model_id(); rows.len()])),
5198        ],
5199    )
5200    .context("failed to build embedding update batch")
5201}
5202
5203/// The runtime backstop against Arrow's 2 GiB `i32` offset wall: a flush batch
5204/// is split before the running total of its text columns reaches this, and a
5205/// single cell at or above it is rejected rather than left to panic inside
5206/// `StringArray::from` (spec.md#adapter-bounded-values).
5207const COLUMN_BYTE_BUDGET: usize = 1 << 30;
5208
5209/// Contiguous row ranges whose summed text-column byte cost each stays within
5210/// `COLUMN_BYTE_BUDGET`. Budgeting the all-column total bounds every individual
5211/// column too, since no single column's total can exceed it. `cells[i]` is row
5212/// `i`'s byte cost summed across every text column.
5213fn chunk_ranges(cells: &[usize]) -> Vec<std::ops::Range<usize>> {
5214    let mut chunks = Vec::new();
5215    let mut start = 0usize;
5216    let mut running = 0usize;
5217    for (index, &row) in cells.iter().enumerate() {
5218        if running + row > COLUMN_BYTE_BUDGET && index > start {
5219            chunks.push(start..index);
5220            start = index;
5221            running = 0;
5222        }
5223        running += row;
5224    }
5225    if start < cells.len() {
5226        chunks.push(start..cells.len());
5227    }
5228    chunks
5229}
5230
5231fn guard_cell(table: &str, pk: &str, bytes: usize) -> Result<()> {
5232    if bytes >= COLUMN_BYTE_BUDGET {
5233        anyhow::bail!(
5234            "{table} row {pk}: a {bytes}-byte text cell meets the per-cell ceiling and would \
5235             overflow Arrow's i32 offset buffer"
5236        );
5237    }
5238    Ok(())
5239}
5240
5241async fn merge_insert_chunks(
5242    handle: &Handle,
5243    table: Table,
5244    batches: Vec<RecordBatch>,
5245) -> Result<u64> {
5246    let mut inserted = 0u64;
5247    for batch in batches {
5248        let rows = batch.num_rows();
5249        inserted += handle.merge_insert(table, batch, rows).await?;
5250    }
5251    Ok(inserted)
5252}
5253
5254pub(crate) fn sessions_batches(sessions: &[Session]) -> Result<Vec<RecordBatch>> {
5255    let options = sessions
5256        .iter()
5257        .map(|session| json_bytes(&session.options))
5258        .collect::<Result<Vec<_>>>()?;
5259    let mut cells = Vec::with_capacity(sessions.len());
5260    for (session, encoded) in sessions.iter().zip(&options) {
5261        let columns = [
5262            session.id.len(),
5263            session.parent_session_id.as_deref().map_or(0, str::len),
5264            session.parent_message_id.as_deref().map_or(0, str::len),
5265            session.source_agent.len(),
5266            session.project.as_str().len(),
5267            encoded.len(),
5268        ];
5269        for bytes in columns {
5270            guard_cell("sessions", &session.id, bytes)?;
5271        }
5272        cells.push(columns.iter().sum());
5273    }
5274    chunk_ranges(&cells)
5275        .into_iter()
5276        .map(|range| sessions_chunk(&sessions[range.clone()], &options[range]))
5277        .collect()
5278}
5279
5280fn sessions_chunk(sessions: &[Session], options: &[Vec<u8>]) -> Result<RecordBatch> {
5281    let schema = session_schema();
5282    RecordBatch::try_new(
5283        schema.clone(),
5284        vec![
5285            Arc::new(StringArray::from(
5286                sessions
5287                    .iter()
5288                    .map(|session| session.id.as_str())
5289                    .collect::<Vec<_>>(),
5290            )),
5291            Arc::new(StringArray::from(
5292                sessions
5293                    .iter()
5294                    .map(|session| session.parent_session_id.as_deref())
5295                    .collect::<Vec<_>>(),
5296            )),
5297            Arc::new(StringArray::from(
5298                sessions
5299                    .iter()
5300                    .map(|session| session.parent_message_id.as_deref())
5301                    .collect::<Vec<_>>(),
5302            )),
5303            Arc::new(StringArray::from(
5304                sessions
5305                    .iter()
5306                    .map(|session| session.source_agent.as_str())
5307                    .collect::<Vec<_>>(),
5308            )),
5309            Arc::new(
5310                TimestampMicrosecondArray::from(
5311                    sessions
5312                        .iter()
5313                        .map(|session| micros(session.created_at))
5314                        .collect::<Vec<_>>(),
5315                )
5316                .with_timezone("UTC"),
5317            ),
5318            Arc::new(StringArray::from(
5319                sessions
5320                    .iter()
5321                    .map(|session| session.project.as_str())
5322                    .collect::<Vec<_>>(),
5323            )),
5324            Arc::new(LargeBinaryArray::from_iter_values(
5325                options.iter().map(Vec::as_slice),
5326            )),
5327        ],
5328    )
5329    .context("failed to build session batch")
5330}
5331
5332/// `vectors` is aligned to `rows` (same length): `Some` carries the inline
5333/// embedding for that row, `None` writes a null `vector`/`embedding_model`.
5334pub(crate) fn messages_batches(
5335    rows: &[MessageBatchRow<'_>],
5336    vectors: &[Option<Vec<f32>>],
5337) -> Result<Vec<RecordBatch>> {
5338    debug_assert_eq!(rows.len(), vectors.len(), "vectors must align with rows");
5339    let options = rows
5340        .iter()
5341        .map(|row| json_bytes(row.message.options()))
5342        .collect::<Result<Vec<_>>>()?;
5343    let mut cells = Vec::with_capacity(rows.len());
5344    for (row, encoded) in rows.iter().zip(&options) {
5345        let columns = [
5346            row.message.session_id().len(),
5347            row.message.id().len(),
5348            row.message.role().as_str().len(),
5349            row.source_agent.len(),
5350            row.project.len(),
5351            row.message.system_content().map_or(0, str::len),
5352            row.search_text.map_or(0, str::len),
5353            encoded.len(),
5354        ];
5355        for bytes in columns {
5356            guard_cell("messages", row.message.id(), bytes)?;
5357        }
5358        cells.push(columns.iter().sum());
5359    }
5360    chunk_ranges(&cells)
5361        .into_iter()
5362        .map(|range| {
5363            messages_chunk(
5364                &rows[range.clone()],
5365                &options[range.clone()],
5366                &vectors[range],
5367            )
5368        })
5369        .collect()
5370}
5371
5372fn messages_chunk(
5373    rows: &[MessageBatchRow<'_>],
5374    options: &[Vec<u8>],
5375    vectors: &[Option<Vec<f32>>],
5376) -> Result<RecordBatch> {
5377    let schema = message_schema();
5378    let (vector_column, embedding_model) = embedding_columns(vectors)?;
5379    RecordBatch::try_new(
5380        schema.clone(),
5381        vec![
5382            Arc::new(StringArray::from(
5383                rows.iter()
5384                    .map(|row| row.message.session_id())
5385                    .collect::<Vec<_>>(),
5386            )),
5387            Arc::new(StringArray::from(
5388                rows.iter().map(|row| row.message.id()).collect::<Vec<_>>(),
5389            )),
5390            Arc::new(
5391                TimestampMicrosecondArray::from(
5392                    rows.iter()
5393                        .map(|row| micros(row.message.timestamp()))
5394                        .collect::<Vec<_>>(),
5395                )
5396                .with_timezone("UTC"),
5397            ),
5398            Arc::new(StringArray::from(
5399                rows.iter()
5400                    .map(|row| row.message.role().as_str())
5401                    .collect::<Vec<_>>(),
5402            )),
5403            Arc::new(StringArray::from(
5404                rows.iter().map(|row| row.source_agent).collect::<Vec<_>>(),
5405            )),
5406            Arc::new(StringArray::from(
5407                rows.iter().map(|row| row.project).collect::<Vec<_>>(),
5408            )),
5409            Arc::new(StringArray::from(
5410                rows.iter()
5411                    .map(|row| row.message.system_content())
5412                    .collect::<Vec<_>>(),
5413            )),
5414            Arc::new(StringArray::from(
5415                rows.iter().map(|row| row.search_text).collect::<Vec<_>>(),
5416            )),
5417            // `vector` / `embedding_model` carry the inline embedding when one
5418            // was produced for the row, null otherwise (embedder disabled, or a
5419            // non-embeddable row); `pond optimize` fills any remaining nulls
5420            // (spec.md#session-embed-from-canonical).
5421            vector_column,
5422            embedding_model,
5423            Arc::new(LargeBinaryArray::from_iter_values(
5424                options.iter().map(Vec::as_slice),
5425            )),
5426        ],
5427    )
5428    .context("failed to build message batch")
5429}
5430
5431pub(crate) fn parts_batches(parts: &[Part]) -> Result<Vec<RecordBatch>> {
5432    let variant_data = parts
5433        .iter()
5434        .map(|part| part_variant_json(&part.kind))
5435        .collect::<Result<Vec<_>>>()?;
5436    let options = parts
5437        .iter()
5438        .map(|part| json_bytes(&part.options))
5439        .collect::<Result<Vec<_>>>()?;
5440    let mut cells = Vec::with_capacity(parts.len());
5441    // The blob column is a BinaryArray, exempt from the text-column bound
5442    // (spec.md#adapter-bounded-values); only the StringArray columns are budgeted.
5443    for ((part, variant), encoded) in parts.iter().zip(&variant_data).zip(&options) {
5444        let columns = [
5445            part.session_id.len(),
5446            part.message_id.len(),
5447            part.id.len(),
5448            part.kind.type_name().len(),
5449            part.provenance.as_str().len(),
5450            variant.len(),
5451            encoded.len(),
5452        ];
5453        for bytes in columns {
5454            guard_cell("parts", &part.id, bytes)?;
5455        }
5456        cells.push(columns.iter().sum());
5457    }
5458    chunk_ranges(&cells)
5459        .into_iter()
5460        .map(|range| {
5461            parts_chunk(
5462                &parts[range.clone()],
5463                &variant_data[range.clone()],
5464                &options[range],
5465            )
5466        })
5467        .collect()
5468}
5469
5470fn parts_chunk(
5471    parts: &[Part],
5472    variant_data: &[Vec<u8>],
5473    options: &[Vec<u8>],
5474) -> Result<RecordBatch> {
5475    let schema = part_schema();
5476    // Legacy blob (`legacy_blob_field`) is a plain LargeBinary; the URL
5477    // variant is stored as UTF-8 bytes and recovered through `variant_data`'s
5478    // `data_kind = "url"` discriminator (see `file_data_from_blob`).
5479    let blob_payloads: Vec<Option<&[u8]>> = parts
5480        .iter()
5481        .map(|part| match &part.kind {
5482            PartKind::File { data, .. } => Some(match data {
5483                FileData::String(value) => value.as_bytes(),
5484                FileData::Bytes(value) => value.as_slice(),
5485                FileData::Url(value) => value.as_bytes(),
5486            }),
5487            PartKind::Text { .. }
5488            | PartKind::Reasoning { .. }
5489            | PartKind::ToolCall { .. }
5490            | PartKind::ToolResult { .. }
5491            | PartKind::ToolApprovalRequest { .. }
5492            | PartKind::ToolApprovalResponse { .. } => None,
5493        })
5494        .collect();
5495    let blob_array = LargeBinaryArray::from_iter(blob_payloads);
5496
5497    let mut tool_names: Vec<Option<&str>> = Vec::with_capacity(parts.len());
5498    let mut call_ids: Vec<Option<&str>> = Vec::with_capacity(parts.len());
5499    let mut failures: Vec<Option<bool>> = Vec::with_capacity(parts.len());
5500    for part in parts {
5501        let (name, call_id, is_failure) = tool_identity(&part.kind);
5502        tool_names.push(name);
5503        call_ids.push(call_id);
5504        failures.push(is_failure);
5505    }
5506
5507    RecordBatch::try_new(
5508        schema.clone(),
5509        vec![
5510            Arc::new(StringArray::from(
5511                parts
5512                    .iter()
5513                    .map(|part| part.session_id.as_str())
5514                    .collect::<Vec<_>>(),
5515            )),
5516            Arc::new(StringArray::from(
5517                parts
5518                    .iter()
5519                    .map(|part| part.message_id.as_str())
5520                    .collect::<Vec<_>>(),
5521            )),
5522            Arc::new(StringArray::from(
5523                parts
5524                    .iter()
5525                    .map(|part| part.id.as_str())
5526                    .collect::<Vec<_>>(),
5527            )),
5528            Arc::new(Int32Array::from(
5529                parts.iter().map(|part| part.ordinal).collect::<Vec<_>>(),
5530            )),
5531            Arc::new(StringArray::from(
5532                parts
5533                    .iter()
5534                    .map(|part| part.kind.type_name())
5535                    .collect::<Vec<_>>(),
5536            )),
5537            Arc::new(StringArray::from(
5538                parts
5539                    .iter()
5540                    .map(|part| part.provenance.as_str())
5541                    .collect::<Vec<_>>(),
5542            )),
5543            Arc::new(StringArray::from(tool_names)),
5544            Arc::new(StringArray::from(call_ids)),
5545            Arc::new(BooleanArray::from(failures)),
5546            Arc::new(LargeBinaryArray::from_iter_values(
5547                variant_data.iter().map(Vec::as_slice),
5548            )),
5549            Arc::new(blob_array),
5550            Arc::new(LargeBinaryArray::from_iter_values(
5551                options.iter().map(Vec::as_slice),
5552            )),
5553        ],
5554    )
5555    .context("failed to build parts batch")
5556}
5557
5558pub(crate) fn session_from_batch(batch: &RecordBatch, row: usize) -> Result<Session> {
5559    Ok(Session {
5560        id: string(batch, "id", row)?.context("session id is null")?,
5561        parent_session_id: string(batch, "parent_session_id", row)?,
5562        parent_message_id: string(batch, "parent_message_id", row)?,
5563        source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5564        created_at: datetime(batch, "created_at", row)?,
5565        project: crate::adapter::Extracted::from_stored(
5566            string(batch, "project", row)?.context("project is null")?,
5567        ),
5568        options: json_parse(&json_column(batch, "options", row)?.context("options is null")?)?,
5569    })
5570}
5571
5572/// [`SkipOracle`](crate::adapter::SkipOracle) over the resident row-meta map:
5573/// `pond sync` reads each session's stored max message timestamp from memory, so
5574/// the staleness check costs zero S3 (the map is rebuilt from the store, so the
5575/// check stays deterministic with no local cursor). A `None` map (never
5576/// prewarmed, or the build failed) yields no watermark, so every source
5577/// re-reads - safe, just slower.
5578pub struct RowmapOracle(pub Option<Arc<RowMetaSet>>);
5579
5580impl crate::adapter::SkipOracle for RowmapOracle {
5581    fn session_max_ts(&self, session_id: &str) -> Option<i64> {
5582        self.0.as_ref()?.lookup_max_ts(session_id)
5583    }
5584
5585    fn is_empty(&self) -> bool {
5586        self.0.as_ref().is_none_or(|set| set.is_empty())
5587    }
5588}
5589
5590fn row_meta_entry(batch: &RecordBatch, row_id: u64, row: usize) -> Result<RowMetaEntry> {
5591    Ok(RowMetaEntry {
5592        row_id,
5593        session_id: string(batch, "session_id", row)?.context("session_id is null")?,
5594        message_id: string(batch, "id", row)?.context("message id is null")?,
5595        role: string(batch, "role", row)?.context("role is null")?,
5596        project: string(batch, "project", row)?.context("project is null")?,
5597        source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5598        timestamp_micros: datetime(batch, "timestamp", row)?.timestamp_micros(),
5599        search_text: string(batch, "search_text", row)?.unwrap_or_default(),
5600    })
5601}
5602
5603pub(crate) fn message_meta_from_batch(batch: &RecordBatch, row: usize) -> Result<MessageMeta> {
5604    Ok(MessageMeta {
5605        message_id: string(batch, "id", row)?.context("id is null")?,
5606        session_id: string(batch, "session_id", row)?.context("session_id is null")?,
5607        role: string(batch, "role", row)?.context("role is null")?,
5608        project: string(batch, "project", row)?.context("project is null")?,
5609        source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5610        timestamp: datetime(batch, "timestamp", row)?,
5611        search_text: string(batch, "search_text", row)?.unwrap_or_default(),
5612    })
5613}
5614
5615pub(crate) fn message_from_batch(batch: &RecordBatch, row: usize) -> Result<Message> {
5616    let id = string(batch, "id", row)?.context("message id is null")?;
5617    let session_id = string(batch, "session_id", row)?.context("message session_id is null")?;
5618    let timestamp = datetime(batch, "timestamp", row)?;
5619    let options =
5620        json_parse(&json_column(batch, "options", row)?.context("message options is null")?)?;
5621
5622    match string(batch, "role", row)?
5623        .context("message role is null")?
5624        .as_str()
5625    {
5626        "system" => Ok(Message::System {
5627            id,
5628            session_id,
5629            timestamp,
5630            // `content` is nullable in the schema; preserve the distinction
5631            // between "no content row stored" (`None`) and "empty string
5632            // stored" (`Some(extracted_empty)`). The value originally
5633            // came from a `Source` extraction at ingest time; rewrap via
5634            // the storage-internal `from_stored` so the type-system seal
5635            // for adapters stays intact.
5636            content: string(batch, "content", row)?.map(crate::adapter::Extracted::from_stored),
5637            options,
5638        }),
5639        "user" => Ok(Message::User {
5640            id,
5641            session_id,
5642            timestamp,
5643            options,
5644        }),
5645        "assistant" => Ok(Message::Assistant {
5646            id,
5647            session_id,
5648            timestamp,
5649            options,
5650        }),
5651        "tool" => Ok(Message::Tool {
5652            id,
5653            session_id,
5654            timestamp,
5655            options,
5656        }),
5657        other => anyhow::bail!("unknown message role {other}"),
5658    }
5659}
5660
5661pub(crate) fn part_from_batch(
5662    batch: &RecordBatch,
5663    row: usize,
5664    file_data: Option<FileData>,
5665) -> Result<Part> {
5666    let type_name = string(batch, "type", row)?.context("part type is null")?;
5667    let variant_data = json_column(batch, "variant_data", row)?.context("variant_data is null")?;
5668    let provenance = string(batch, "provenance", row)?.context("part provenance is null")?;
5669    Ok(Part {
5670        session_id: string(batch, "session_id", row)?.context("part session_id is null")?,
5671        message_id: string(batch, "message_id", row)?.context("part message_id is null")?,
5672        id: string(batch, "id", row)?.context("part id is null")?,
5673        ordinal: int32(batch, "ordinal", row)?,
5674        provenance: provenance_from_str(&provenance)?,
5675        options: json_parse(&json_column(batch, "options", row)?.context("part options is null")?)?,
5676        kind: part_kind_from_json(&type_name, &variant_data, file_data)?,
5677    })
5678}
5679
5680fn provenance_from_str(value: &str) -> Result<crate::wire::Provenance> {
5681    match value {
5682        "conversational" => Ok(crate::wire::Provenance::Conversational),
5683        "injected" => Ok(crate::wire::Provenance::Injected),
5684        other => anyhow::bail!("unknown part provenance {other}"),
5685    }
5686}
5687
5688fn file_data_from_blob(variant_data: &[u8], bytes: &[u8]) -> Result<FileData> {
5689    let kind = file_data_kind(variant_data)?;
5690    match kind.as_str() {
5691        "string" => {
5692            let text = std::str::from_utf8(bytes)
5693                .context("file string payload is not UTF-8")?
5694                .to_owned();
5695            Ok(FileData::String(text))
5696        }
5697        "bytes" => Ok(FileData::Bytes(bytes.to_vec())),
5698        "url" => Ok(FileData::Url(
5699            std::str::from_utf8(bytes)
5700                .context("file URL payload is not UTF-8")?
5701                .to_owned(),
5702        )),
5703        other => anyhow::bail!("unknown file data_kind {other}"),
5704    }
5705}
5706
5707fn file_data_kind(variant_data: &[u8]) -> Result<String> {
5708    let value = json_parse::<Value>(variant_data)?;
5709    value
5710        .get("data_kind")
5711        .and_then(Value::as_str)
5712        .map(str::to_owned)
5713        .context("file part variant_data missing data_kind")
5714}
5715
5716fn uint64<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a UInt64Array> {
5717    batch
5718        .column_by_name(name)
5719        .with_context(|| format!("missing column {name}"))?
5720        .as_any()
5721        .downcast_ref::<UInt64Array>()
5722        .with_context(|| format!("column {name} is not UInt64"))
5723}
5724
5725pub(crate) fn string(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<String>> {
5726    let array = batch
5727        .column_by_name(name)
5728        .with_context(|| format!("missing column {name}"))?
5729        .as_any()
5730        .downcast_ref::<StringArray>()
5731        .with_context(|| format!("column {name} is not Utf8"))?;
5732    if array.is_null(row) {
5733        Ok(None)
5734    } else {
5735        Ok(Some(array.value(row).to_owned()))
5736    }
5737}
5738
5739fn json_column(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<Vec<u8>>> {
5740    // Lance can return a `lance.json` column either as raw JSONB bytes
5741    // (LargeBinary) or auto-converted to the Arrow text form (Utf8 /
5742    // LargeUtf8), depending on the read path. Handle both.
5743    let column = batch
5744        .column_by_name(name)
5745        .with_context(|| format!("missing column {name}"))?;
5746    if let Some(array) = column.as_any().downcast_ref::<LargeBinaryArray>() {
5747        return if array.is_null(row) {
5748            Ok(None)
5749        } else {
5750            Ok(Some(
5751                lance_arrow::json::decode_json(array.value(row)).into_bytes(),
5752            ))
5753        };
5754    }
5755    if let Some(array) = column.as_any().downcast_ref::<StringArray>() {
5756        return if array.is_null(row) {
5757            Ok(None)
5758        } else {
5759            Ok(Some(array.value(row).as_bytes().to_vec()))
5760        };
5761    }
5762    if let Some(array) = column.as_any().downcast_ref::<LargeStringArray>() {
5763        return if array.is_null(row) {
5764            Ok(None)
5765        } else {
5766            Ok(Some(array.value(row).as_bytes().to_vec()))
5767        };
5768    }
5769    anyhow::bail!("column {name} is not a JSON-compatible array")
5770}
5771
5772fn int32(batch: &RecordBatch, name: &str, row: usize) -> Result<i32> {
5773    let array = batch
5774        .column_by_name(name)
5775        .with_context(|| format!("missing column {name}"))?
5776        .as_any()
5777        .downcast_ref::<Int32Array>()
5778        .with_context(|| format!("column {name} is not Int32"))?;
5779    Ok(array.value(row))
5780}
5781
5782pub(crate) fn float32(batch: &RecordBatch, name: &str, row: usize) -> Result<f32> {
5783    let array = batch
5784        .column_by_name(name)
5785        .with_context(|| format!("missing column {name}"))?
5786        .as_any()
5787        .downcast_ref::<Float32Array>()
5788        .with_context(|| format!("column {name} is not Float32"))?;
5789    Ok(array.value(row))
5790}
5791
5792pub(crate) fn datetime(batch: &RecordBatch, name: &str, row: usize) -> Result<DateTime<Utc>> {
5793    let array = batch
5794        .column_by_name(name)
5795        .with_context(|| format!("missing column {name}"))?
5796        .as_any()
5797        .downcast_ref::<TimestampMicrosecondArray>()
5798        .with_context(|| format!("column {name} is not timestamp_micros"))?;
5799    Utc.timestamp_micros(array.value(row))
5800        .single()
5801        .context("timestamp is out of range")
5802}
5803
5804fn primary_field(name: &str, data_type: DataType, nullable: bool) -> Field {
5805    Field::new(name, data_type, nullable).with_metadata(
5806        [(
5807            "lance-schema:unenforced-primary-key".to_owned(),
5808            "true".to_owned(),
5809        )]
5810        .into(),
5811    )
5812}
5813
5814// Legacy blob storage (`LargeBinary` + `lance-encoding:blob=true`). Blob v2's
5815// `Struct<data, uri>` extension requires `data_storage_version >= 2.2`, which
5816// is marked unstable in Lance docs (`format/file/versioning.md`) and at
5817// v7.0.0-beta.16 trips a `compact_files` bug: the AllBinary blob_handling
5818// path leaves the field as a 2-child struct but `BlobV2StructuralEncoder`
5819// allocated only one column_info, so the decoder's second `expect_next()`
5820// fires `"there were more fields in the schema than provided column
5821// indices / infos"`. Legacy blob writes `BlobLayout` pages, which compact
5822// handles correctly (covered by Lance's own `test_compact_blob_columns`).
5823fn legacy_blob_field(name: &str, nullable: bool) -> Field {
5824    Field::new(name, DataType::LargeBinary, nullable).with_metadata(
5825        [(lance_arrow::BLOB_META_KEY.to_owned(), "true".to_owned())]
5826            .into_iter()
5827            .collect(),
5828    )
5829}
5830
5831// Deliberately NO `lance-encoding:compression` metadata (#89): lance's
5832// miniblock zstd compresses tiny chunks independently and cannot reach the
5833// cross-row redundancy where the real ratio lives - measured 0% on real
5834// corpora vs 4.8-7.4x full-window, while values > 32 KiB are already
5835// compressed per-value by default. Tool analytics avoid the fat column via
5836// the materialized tool_name/call_id/is_failure columns instead.
5837fn json_field(name: &str, nullable: bool) -> Field {
5838    lance_arrow::json::json_field(name, nullable)
5839}
5840
5841fn micros(timestamp: DateTime<Utc>) -> i64 {
5842    timestamp.timestamp_micros()
5843}
5844
5845fn json_bytes<T: Serialize>(value: &T) -> Result<Vec<u8>> {
5846    // Write JSONB bytes (not plain UTF-8 JSON text) so the on-disk encoding
5847    // matches the `lance.json` extension contract. Lance's compact path
5848    // (`optimize.rs:908`) reads through `DatasetRecordBatchStream` which
5849    // applies `decode_json -> encode_json` on this column; with proper JSONB
5850    // on disk that roundtrip is idempotent, with plain UTF-8 it corrupts
5851    // (the analogous fix landed for `update.rs` in PR #6741 by switching to
5852    // `try_into_dfstream`; compact still goes through the adapter).
5853    let text = serde_json::to_string(value).context("failed to serialize JSON field")?;
5854    lance_arrow::json::encode_json(&text)
5855        .map_err(|err| anyhow::anyhow!("failed to encode JSON field as JSONB: {err}"))
5856}
5857
5858fn json_parse<T: DeserializeOwned>(value: &[u8]) -> Result<T> {
5859    serde_json::from_slice(value).context("failed to parse JSON field")
5860}
5861
5862/// The materialized `(tool_name, call_id, is_failure)` cells for one part -
5863/// shared by the ingest write path and the schema-migration backfill (which
5864/// reconstructs the `PartKind` via [`part_kind_from_json`]) so both derive
5865/// identical cells. The approval request's `tool_call_id` surfaces under the
5866/// one `call_id` column (the same correlation key the call/result pair
5867/// carries); absent source fields stay NULL (spec.md#model-no-synthesis).
5868fn tool_identity(kind: &PartKind) -> (Option<&str>, Option<&str>, Option<bool>) {
5869    match kind {
5870        PartKind::ToolCall { name, call_id, .. } => (
5871            name.as_deref().map(String::as_str),
5872            call_id.as_deref().map(String::as_str),
5873            None,
5874        ),
5875        PartKind::ToolResult {
5876            name,
5877            call_id,
5878            is_failure,
5879            ..
5880        } => (
5881            name.as_deref().map(String::as_str),
5882            call_id.as_deref().map(String::as_str),
5883            Some(*is_failure),
5884        ),
5885        PartKind::ToolApprovalRequest { tool_call_id, .. } => {
5886            (None, Some(tool_call_id.as_str()), None)
5887        }
5888        PartKind::Text { .. }
5889        | PartKind::Reasoning { .. }
5890        | PartKind::File { .. }
5891        | PartKind::ToolApprovalResponse { .. } => (None, None, None),
5892    }
5893}
5894
5895fn part_variant_json(kind: &PartKind) -> Result<Vec<u8>> {
5896    if let PartKind::File {
5897        media_type,
5898        file_name,
5899        data,
5900    } = kind
5901    {
5902        let data_kind = match data {
5903            FileData::String(_) => "string",
5904            FileData::Bytes(_) => "bytes",
5905            FileData::Url(_) => "url",
5906        };
5907        return json_bytes(&serde_json::json!({
5908            "media_type": media_type,
5909            "file_name": file_name,
5910            "data_kind": data_kind,
5911        }));
5912    }
5913    let value = serde_json::to_value(kind)?;
5914    let mut object = value
5915        .as_object()
5916        .cloned()
5917        .context("part variant did not serialize to an object")?;
5918    object.remove("type");
5919    json_bytes(&object)
5920}
5921
5922fn part_kind_from_json(
5923    type_name: &str,
5924    variant_data: &[u8],
5925    file_data: Option<FileData>,
5926) -> Result<PartKind> {
5927    let mut value = json_parse::<Value>(variant_data)?;
5928    let object = value
5929        .as_object_mut()
5930        .context("part variant data is not an object")?;
5931    object.insert("type".to_owned(), Value::String(type_name.to_owned()));
5932    if let Some(data) = file_data {
5933        object.remove("data_kind");
5934        object.insert("data".to_owned(), serde_json::to_value(data)?);
5935    }
5936    serde_json::from_value(value).context("failed to parse part kind")
5937}
5938
5939#[cfg(test)]
5940mod tests {
5941    #![allow(clippy::expect_used, clippy::unwrap_used)]
5942
5943    use super::*;
5944    use crate::{
5945        adapter::Extracted,
5946        handlers::ingest_events,
5947        wire::{FileData, Message, Part, PartKind, ProviderOptions, Session},
5948    };
5949    use chrono::Utc;
5950    use serde_json::json;
5951    use tempfile::TempDir;
5952
5953    fn synthetic_session(id: &str) -> Session {
5954        Session {
5955            id: id.to_owned(),
5956            parent_session_id: None,
5957            parent_message_id: None,
5958            source_agent: "claude-code".to_owned(),
5959            created_at: Utc::now(),
5960            project: crate::adapter::Extracted::from_test_value("/tmp/pond".to_owned()),
5961            options: ProviderOptions::new(),
5962        }
5963    }
5964
5965    /// Counts the texts handed to the backend so a test can assert how many rows
5966    /// were embedded.
5967    #[derive(Default)]
5968    struct CountingEmbedder {
5969        texts: std::sync::atomic::AtomicUsize,
5970    }
5971    impl crate::embed::Embedder for CountingEmbedder {
5972        fn device(&self) -> &str {
5973            "test"
5974        }
5975        fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
5976            self.texts
5977                .fetch_add(texts.len(), std::sync::atomic::Ordering::SeqCst);
5978            Ok(texts
5979                .iter()
5980                .map(|_| vec![0.0_f32; embedding_dim()])
5981                .collect())
5982        }
5983    }
5984
5985    /// A session with `count` conversational user messages, each carrying a text
5986    /// part so its `search_text` is non-null (hence embeddable).
5987    fn conversational_events(session_id: &str, count: usize) -> Vec<IngestEvent> {
5988        let mut events = vec![IngestEvent::Session(synthetic_session(session_id))];
5989        for index in 0..count {
5990            events.push(IngestEvent::Message(Message::User {
5991                id: format!("msg-{index}"),
5992                session_id: session_id.to_owned(),
5993                timestamp: Utc::now(),
5994                options: ProviderOptions::new(),
5995            }));
5996            events.push(IngestEvent::Part(Part {
5997                session_id: session_id.to_owned(),
5998                id: format!("msg-{index}:0001"),
5999                message_id: format!("msg-{index}"),
6000                ordinal: 0,
6001                provenance: crate::wire::Provenance::Conversational,
6002                options: ProviderOptions::new(),
6003                kind: PartKind::Text {
6004                    text: Some(Extracted::from_test_value(format!("body {index}"))),
6005                },
6006            }));
6007        }
6008        events
6009    }
6010
6011    /// Inline embed-at-ingest: with an embedder attached the vectors are filled
6012    /// in the message rows' birth append (no extra embed commit - the version
6013    /// matches a plain ingest), and a re-sync embeds no already-present row.
6014    #[tokio::test(flavor = "multi_thread")]
6015    async fn ingest_embeds_inline_in_the_birth_commit() -> anyhow::Result<()> {
6016        let plain = Store::open(&Url::parse("shared-memory://pond-test-inline-plain/")?).await?;
6017        ingest_events(&plain, conversational_events("01HXYINLINE000PLAIN", 5)).await?;
6018        assert!(
6019            !plain.has_embeddings().await?,
6020            "no embedder attached -> every vector null",
6021        );
6022        let plain_version = plain.messages_version().await?;
6023
6024        let backend = Arc::new(CountingEmbedder::default());
6025        let embedder = Arc::new(crate::embed::LazyEmbedder::from_loaded(
6026            backend.clone() as Arc<dyn crate::embed::Embedder>
6027        ));
6028        let store = Store::open(&Url::parse("shared-memory://pond-test-inline-embed/")?)
6029            .await?
6030            .with_embedder(embedder);
6031        ingest_events(&store, conversational_events("01HXYINLINE000EMBED", 5)).await?;
6032        assert!(
6033            store.has_embeddings().await?,
6034            "embedder attached -> vectors filled at ingest",
6035        );
6036        assert_eq!(
6037            backend.texts.load(std::sync::atomic::Ordering::SeqCst),
6038            5,
6039            "every conversational message embedded once",
6040        );
6041        assert_eq!(
6042            store.messages_version().await?,
6043            plain_version,
6044            "inline embed must ride the append, not add a separate commit",
6045        );
6046
6047        ingest_events(&store, conversational_events("01HXYINLINE000EMBED", 5)).await?;
6048        assert_eq!(
6049            backend.texts.load(std::sync::atomic::Ordering::SeqCst),
6050            5,
6051            "an idempotent re-sync embeds no already-present row",
6052        );
6053        Ok(())
6054    }
6055
6056    /// The verify's duplicate count (`rows - distinct composite PKs` from
6057    /// [`Store::composite_pk_index`]) keys on the composite PK, so the same
6058    /// message id in two different sessions is NOT a duplicate (the bare-id
6059    /// false-positive trap), while a genuinely doubled `(session_id, id)` row
6060    /// counts as one.
6061    #[tokio::test(flavor = "multi_thread")]
6062    async fn composite_pk_index_counts_duplicates_by_composite_key() -> anyhow::Result<()> {
6063        async fn duplicates(store: &Store, table: Table) -> anyhow::Result<usize> {
6064            let (keys, rows) = store.composite_pk_index(table).await?;
6065            Ok(rows - keys.len())
6066        }
6067        let store = Store::open(&Url::parse("shared-memory://pond-test-dupcount/")?).await?;
6068        ingest_events(&store, conversational_events("01HXYDUP00000SESS1", 1)).await?;
6069        ingest_events(&store, conversational_events("01HXYDUP00000SESS2", 1)).await?;
6070        assert_eq!(
6071            duplicates(&store, Table::Messages).await?,
6072            0,
6073            "the same message id in two sessions is not a duplicate (composite PK)",
6074        );
6075        assert_eq!(duplicates(&store, Table::Sessions).await?, 0);
6076        assert_eq!(duplicates(&store, Table::Parts).await?, 0);
6077
6078        // Inject a real duplicate via the low-level append (no dedup) - the
6079        // write anomaly the copy verify must catch.
6080        let message = Message::User {
6081            id: "dup-msg".to_owned(),
6082            session_id: "01HXYDUP00000SESS1".to_owned(),
6083            timestamp: Utc::now(),
6084            options: ProviderOptions::new(),
6085        };
6086        let row = MessageBatchRow {
6087            message: &message,
6088            source_agent: "claude-code",
6089            project: "/tmp",
6090            search_text: None,
6091        };
6092        let batches = messages_batches(&[row], &[None])?;
6093        store
6094            .handle
6095            .append_batches(Table::Messages, batches.clone())
6096            .await?;
6097        store
6098            .handle
6099            .append_batches(Table::Messages, batches)
6100            .await?;
6101        assert_eq!(
6102            duplicates(&store, Table::Messages).await?,
6103            1,
6104            "the doubled (session_id, id) row is exactly one duplicate",
6105        );
6106        Ok(())
6107    }
6108
6109    #[test]
6110    fn search_text_excludes_injected_parts() {
6111        use crate::wire::Provenance;
6112        let message = Message::User {
6113            id: "m1".to_owned(),
6114            session_id: "s1".to_owned(),
6115            timestamp: Utc::now(),
6116            options: ProviderOptions::new(),
6117        };
6118        let text_part = |id: &str, text: &str, provenance: Provenance| Part {
6119            session_id: "s1".to_owned(),
6120            id: id.to_owned(),
6121            message_id: "m1".to_owned(),
6122            ordinal: 0,
6123            provenance,
6124            options: ProviderOptions::new(),
6125            kind: PartKind::Text {
6126                text: Some(Extracted::from_test_value(text.to_owned())),
6127            },
6128        };
6129
6130        // A conversational part contributes; an injected one is excluded
6131        // (spec.md#search).
6132        let conversational = search_text(
6133            &message,
6134            &[text_part(
6135                "p1",
6136                "real human prompt",
6137                Provenance::Conversational,
6138            )],
6139        );
6140        assert_eq!(conversational.as_deref(), Some("real human prompt"));
6141
6142        let injected = search_text(
6143            &message,
6144            &[text_part(
6145                "p2",
6146                "<task-notification>...</task-notification>",
6147                Provenance::Injected,
6148            )],
6149        );
6150        assert!(
6151            injected.is_none(),
6152            "a message whose only part is injected has null search_text"
6153        );
6154    }
6155
6156    #[test]
6157    fn parts_chunk_materializes_tool_identity_columns() -> anyhow::Result<()> {
6158        let part = |ordinal: i32, kind: PartKind| Part {
6159            session_id: "s1".to_owned(),
6160            id: format!("p{ordinal}"),
6161            message_id: "m1".to_owned(),
6162            ordinal,
6163            provenance: crate::wire::Provenance::Conversational,
6164            options: ProviderOptions::new(),
6165            kind,
6166        };
6167        let parts = vec![
6168            part(
6169                0,
6170                PartKind::ToolCall {
6171                    call_id: Some(Extracted::from_test_value("c1".to_owned())),
6172                    name: Some(Extracted::from_test_value("Bash".to_owned())),
6173                    params: serde_json::json!({}),
6174                    provider_executed: false,
6175                },
6176            ),
6177            part(
6178                1,
6179                PartKind::ToolResult {
6180                    call_id: Some(Extracted::from_test_value("c1".to_owned())),
6181                    name: Some(Extracted::from_test_value("Bash".to_owned())),
6182                    is_failure: true,
6183                    result: serde_json::json!({}),
6184                },
6185            ),
6186            // Absent source fields stay NULL - never a synthesized sentinel.
6187            part(
6188                2,
6189                PartKind::ToolCall {
6190                    call_id: None,
6191                    name: None,
6192                    params: serde_json::json!({}),
6193                    provider_executed: false,
6194                },
6195            ),
6196            part(
6197                3,
6198                PartKind::ToolApprovalRequest {
6199                    approval_id: "a1".to_owned(),
6200                    tool_call_id: "c1".to_owned(),
6201                },
6202            ),
6203            part(
6204                4,
6205                PartKind::Text {
6206                    text: Some(Extracted::from_test_value("hi".to_owned())),
6207                },
6208            ),
6209        ];
6210        let payloads = vec![vec![1u8]; parts.len()];
6211        let batch = parts_chunk(&parts, &payloads, &payloads)?;
6212
6213        let tool_name = |row| string(&batch, "tool_name", row).unwrap();
6214        let call_id = |row| string(&batch, "call_id", row).unwrap();
6215        let is_failure = batch
6216            .column_by_name("is_failure")
6217            .expect("is_failure column present")
6218            .as_any()
6219            .downcast_ref::<BooleanArray>()
6220            .expect("is_failure is boolean");
6221
6222        assert_eq!(tool_name(0).as_deref(), Some("Bash"));
6223        assert_eq!(call_id(0).as_deref(), Some("c1"));
6224        assert!(is_failure.is_null(0), "tool_call has no failure flag");
6225
6226        assert_eq!(tool_name(1).as_deref(), Some("Bash"));
6227        assert_eq!(call_id(1).as_deref(), Some("c1"));
6228        assert!(is_failure.value(1), "tool_result failure flag materialized");
6229
6230        assert_eq!(tool_name(2), None, "absent name stays NULL");
6231        assert_eq!(call_id(2), None, "absent call_id stays NULL");
6232
6233        assert_eq!(tool_name(3), None);
6234        assert_eq!(
6235            call_id(3).as_deref(),
6236            Some("c1"),
6237            "approval request surfaces its tool_call_id as call_id",
6238        );
6239
6240        assert_eq!(tool_name(4), None);
6241        assert_eq!(call_id(4), None);
6242        assert!(is_failure.is_null(4));
6243        Ok(())
6244    }
6245
6246    #[test]
6247    fn chunk_ranges_splits_on_byte_budget() {
6248        assert!(chunk_ranges(&[]).is_empty());
6249        assert_eq!(chunk_ranges(&[10, 10, 10]), vec![0..3]);
6250
6251        let two_thirds = COLUMN_BYTE_BUDGET * 2 / 3;
6252        assert_eq!(
6253            chunk_ranges(&[two_thirds, two_thirds, two_thirds]),
6254            vec![0..1, 1..2, 2..3],
6255        );
6256
6257        // An oversized single row gets its own chunk, never an infinite loop.
6258        assert_eq!(
6259            chunk_ranges(&[10, COLUMN_BYTE_BUDGET + 1, 10]),
6260            vec![0..1, 1..2, 2..3],
6261        );
6262    }
6263
6264    #[tokio::test]
6265    async fn ordering_violation_drops_only_the_offending_event() -> anyhow::Result<()> {
6266        // Per-event drop semantics (spec.md#adapter-integrity-event-ordering): a Part with no preceding
6267        // Message is dropped on the spot, with one Error outcome surfaced. The
6268        // rest of the substream continues normally - subsequent valid messages
6269        // and parts get written.
6270        let temp = TempDir::new()?;
6271        let store = Store::open_local(temp.path()).await?;
6272        let session = synthetic_session("ordering");
6273        let orphan_part = Part {
6274            session_id: session.id.clone(),
6275            id: "orphan-part".to_owned(),
6276            message_id: "missing-message".to_owned(),
6277            ordinal: 0,
6278            provenance: crate::wire::Provenance::Conversational,
6279            options: ProviderOptions::new(),
6280            kind: PartKind::Text {
6281                text: Some(Extracted::from_test_value("orphan".to_owned())),
6282            },
6283        };
6284        let valid_message = Message::User {
6285            id: "valid-message".to_owned(),
6286            session_id: session.id.clone(),
6287            timestamp: Utc::now(),
6288            options: ProviderOptions::new(),
6289        };
6290        let valid_part = Part {
6291            session_id: session.id.clone(),
6292            id: "valid-part".to_owned(),
6293            message_id: valid_message.id().to_owned(),
6294            ordinal: 0,
6295            provenance: crate::wire::Provenance::Conversational,
6296            options: ProviderOptions::new(),
6297            kind: PartKind::Text {
6298                text: Some(Extracted::from_test_value("kept".to_owned())),
6299            },
6300        };
6301
6302        let mut validator = IngestValidator::default();
6303        validator
6304            .push(&store, 0, IngestEvent::Session(session.clone()))
6305            .await?;
6306        let part_outcomes = validator
6307            .push(&store, 1, IngestEvent::Part(orphan_part))
6308            .await?;
6309        assert_eq!(part_outcomes.len(), 1);
6310        assert_eq!(part_outcomes[0].kind, "part");
6311        assert_eq!(part_outcomes[0].status, OutcomeStatus::Error);
6312        assert!(
6313            part_outcomes[0]
6314                .error
6315                .as_ref()
6316                .map(|e| e.message.contains("part event appeared before a message"))
6317                .unwrap_or(false),
6318            "error message must explain the ordering violation: {part_outcomes:?}"
6319        );
6320        validator
6321            .push(&store, 2, IngestEvent::Message(valid_message))
6322            .await?;
6323        validator
6324            .push(&store, 3, IngestEvent::Part(valid_part))
6325            .await?;
6326        validator.finish(&store).await?;
6327
6328        let (sessions, messages, parts) = store.row_counts().await?;
6329        assert_eq!(sessions, 1, "session committed despite the orphan part");
6330        assert_eq!(messages, 1, "valid message committed");
6331        assert_eq!(parts, 1, "valid part committed; the orphan was dropped");
6332
6333        Ok(())
6334    }
6335
6336    #[tokio::test]
6337    async fn resident_meta_map_hydration_matches_take_rows_fallback() -> anyhow::Result<()> {
6338        // The resident meta map must hydrate hits identically to the take_rows
6339        // fallback - same fields, and the microsecond timestamp survives the
6340        // i64 round-trip through the mmap blob.
6341        let temp = TempDir::new()?;
6342        let store = Store::open_local(temp.path()).await?;
6343        let session = synthetic_session("hydration-parity");
6344
6345        let messages = [
6346            (
6347                "m1",
6348                "the auth refactor landed cleanly",
6349                1_700_000_000_123_456_i64,
6350            ),
6351            (
6352                "m2",
6353                "balance handler now retries on rpc timeout",
6354                1_700_000_050_654_321,
6355            ),
6356        ];
6357        let mut validator = IngestValidator::default();
6358        validator
6359            .push(&store, 0, IngestEvent::Session(session.clone()))
6360            .await?;
6361        let mut seq = 1;
6362        for (mid, text, micros) in messages {
6363            let message = Message::User {
6364                id: mid.to_owned(),
6365                session_id: session.id.clone(),
6366                timestamp: DateTime::from_timestamp_micros(micros).unwrap(),
6367                options: ProviderOptions::new(),
6368            };
6369            validator
6370                .push(&store, seq, IngestEvent::Message(message))
6371                .await?;
6372            seq += 1;
6373            let part = Part {
6374                session_id: session.id.clone(),
6375                id: format!("{mid}-p0"),
6376                message_id: mid.to_owned(),
6377                ordinal: 0,
6378                provenance: crate::wire::Provenance::Conversational,
6379                options: ProviderOptions::new(),
6380                kind: PartKind::Text {
6381                    text: Some(Extracted::from_test_value(text.to_owned())),
6382                },
6383            };
6384            validator.push(&store, seq, IngestEvent::Part(part)).await?;
6385            seq += 1;
6386        }
6387        validator.finish(&store).await?;
6388
6389        let rowids: Vec<u64> = store
6390            .collect_row_metas()
6391            .await?
6392            .into_iter()
6393            .map(|entry| entry.row_id)
6394            .collect();
6395        assert_eq!(rowids.len(), 2);
6396
6397        let sort_by_id = |mut metas: Vec<MessageMeta>| {
6398            metas.sort_by(|left, right| left.message_id.cmp(&right.message_id));
6399            metas
6400        };
6401
6402        let fallback = sort_by_id(store.message_metas_by_rowids(&rowids).await?);
6403
6404        // Build and install the resident meta map; the same call now hydrates
6405        // from memory (zero misses - the map covers the whole table).
6406        store.ensure_rowmap(&temp.path().join("cache")).await?;
6407        let resident = sort_by_id(store.message_metas_by_rowids(&rowids).await?);
6408
6409        assert_eq!(
6410            resident, fallback,
6411            "resident-map hydration must match the take_rows fallback"
6412        );
6413        assert_eq!(
6414            resident[0].timestamp.timestamp_micros(),
6415            1_700_000_000_123_456
6416        );
6417        Ok(())
6418    }
6419
6420    #[tokio::test]
6421    async fn initialized_flips_only_after_first_ingest() -> anyhow::Result<()> {
6422        // `open` eagerly creates only `messages`; `sessions` and `parts` are
6423        // lazy, so a configured-but-never-synced store reports uninitialized -
6424        // the signal `pond status` uses to render an empty state instead of
6425        // erroring on the first parts describe.
6426        let temp = TempDir::new()?;
6427        let store = Store::open_local(temp.path()).await?;
6428        assert!(
6429            !store.initialized().await?,
6430            "fresh store has no parts table"
6431        );
6432
6433        let session = synthetic_session("initialized-probe");
6434        let message = Message::User {
6435            id: "message-1".to_owned(),
6436            session_id: session.id.clone(),
6437            timestamp: Utc::now(),
6438            options: ProviderOptions::new(),
6439        };
6440        let part = Part {
6441            session_id: session.id.clone(),
6442            id: "part-1".to_owned(),
6443            message_id: message.id().to_owned(),
6444            ordinal: 0,
6445            provenance: crate::wire::Provenance::Conversational,
6446            options: ProviderOptions::new(),
6447            kind: PartKind::Text {
6448                text: Some(Extracted::from_test_value("hello".to_owned())),
6449            },
6450        };
6451        let mut validator = IngestValidator::default();
6452        validator
6453            .push(&store, 0, IngestEvent::Session(session))
6454            .await?;
6455        validator
6456            .push(&store, 1, IngestEvent::Message(message))
6457            .await?;
6458        validator.push(&store, 2, IngestEvent::Part(part)).await?;
6459        validator.finish(&store).await?;
6460
6461        assert!(store.initialized().await?, "ingest creates the parts table");
6462        Ok(())
6463    }
6464
6465    #[tokio::test]
6466    async fn summary_parts_label_a_file_without_reading_its_blob() -> anyhow::Result<()> {
6467        // The summary path (search hits, conversational view) labels a file from
6468        // `variant_data` (`file_name`/`media_type`) and must NOT fetch the blob -
6469        // `scan_parts` substitutes an empty placeholder so `PartKind::File` still
6470        // deserializes. The full path keeps reading the real bytes. Guards both.
6471        let temp = TempDir::new()?;
6472        let store = Store::open_local(temp.path()).await?;
6473        let session = synthetic_session("file-summary");
6474        let message = Message::User {
6475            id: "m1".to_owned(),
6476            session_id: session.id.clone(),
6477            timestamp: Utc::now(),
6478            options: ProviderOptions::new(),
6479        };
6480        let blob = "file contents the summary must never read";
6481        let part = Part {
6482            session_id: session.id.clone(),
6483            id: "m1-p0".to_owned(),
6484            message_id: "m1".to_owned(),
6485            ordinal: 0,
6486            provenance: crate::wire::Provenance::Conversational,
6487            options: ProviderOptions::new(),
6488            kind: PartKind::File {
6489                media_type: Some("text/plain".to_owned()),
6490                file_name: Some("notes.txt".to_owned()),
6491                data: FileData::String(blob.to_owned()),
6492            },
6493        };
6494        let mut validator = IngestValidator::default();
6495        validator
6496            .push(&store, 0, IngestEvent::Session(session.clone()))
6497            .await?;
6498        validator
6499            .push(&store, 1, IngestEvent::Message(message))
6500            .await?;
6501        validator.push(&store, 2, IngestEvent::Part(part)).await?;
6502        validator.finish(&store).await?;
6503
6504        let key = (session.id.clone(), "m1".to_owned());
6505        let ids = ["m1".to_owned()];
6506
6507        let summarized = store.summary_parts_for_messages(&session.id, &ids).await?;
6508        let summary_part = &summarized.get(&key).expect("file part summarized")[0];
6509        let summary = crate::wire::PartSummary::for_kind(&summary_part.kind)
6510            .expect("a file part yields a summary");
6511        assert_eq!(summary.kind, "file");
6512        assert_eq!(summary.label.as_deref(), Some("notes.txt"));
6513        match &summary_part.kind {
6514            PartKind::File { data, .. } => assert!(
6515                matches!(data, FileData::Bytes(bytes) if bytes.is_empty()),
6516                "summary must carry the empty placeholder, not the file blob",
6517            ),
6518            other => panic!("expected a file part, got {other:?}"),
6519        }
6520
6521        let full = store.parts_for_messages(&session.id, &ids).await?;
6522        match &full.get(&key).expect("file part")[0].kind {
6523            PartKind::File {
6524                data: FileData::String(bytes),
6525                ..
6526            } => assert_eq!(bytes, blob, "full path must still hydrate the real blob"),
6527            other => panic!("expected a string-backed file part, got {other:?}"),
6528        }
6529        Ok(())
6530    }
6531
6532    #[tokio::test]
6533    async fn duplicate_message_id_drops_the_second_keeps_the_first() -> anyhow::Result<()> {
6534        // Per-event drop: a duplicate message id within a substream drops the
6535        // *duplicate* and surfaces an Error outcome for it. The first wins; the
6536        // session still commits.
6537        let temp = TempDir::new()?;
6538        let store = Store::open_local(temp.path()).await?;
6539        let session = synthetic_session("duplicate-message");
6540        let first = Message::User {
6541            id: "message-1".to_owned(),
6542            session_id: session.id.clone(),
6543            timestamp: Utc::now(),
6544            options: ProviderOptions::new(),
6545        };
6546        let second = Message::Assistant {
6547            id: "message-1".to_owned(),
6548            session_id: session.id.clone(),
6549            timestamp: Utc::now(),
6550            options: ProviderOptions::new(),
6551        };
6552
6553        let mut validator = IngestValidator::default();
6554        validator
6555            .push(&store, 0, IngestEvent::Session(session.clone()))
6556            .await?;
6557        validator
6558            .push(&store, 1, IngestEvent::Message(first))
6559            .await?;
6560        let dup_outcomes = validator
6561            .push(&store, 2, IngestEvent::Message(second))
6562            .await?;
6563        assert_eq!(dup_outcomes.len(), 1);
6564        assert_eq!(dup_outcomes[0].status, OutcomeStatus::Error);
6565        assert!(
6566            dup_outcomes[0]
6567                .error
6568                .as_ref()
6569                .map(|e| e.message.contains("duplicate message id message-1"))
6570                .unwrap_or(false),
6571            "duplicate-id rejection must name the offending id: {dup_outcomes:?}"
6572        );
6573
6574        validator.finish(&store).await?;
6575        let (sessions, messages, _) = store.row_counts().await?;
6576        assert_eq!(sessions, 1, "session committed");
6577        assert_eq!(messages, 1, "only the first message committed");
6578
6579        Ok(())
6580    }
6581
6582    #[tokio::test]
6583    async fn ingest_stamps_host_provenance_on_messages_and_strips_spoofed_pond_key()
6584    -> anyhow::Result<()> {
6585        // spec.md#model-pond-options: `options.pond` is core-owned. A stored
6586        // message carries the process's host stamp (when resolvable) and never
6587        // a client-supplied value; session and part options stay untouched.
6588        let temp = TempDir::new()?;
6589        let store = Store::open_local(temp.path()).await?;
6590        let session = synthetic_session("host-provenance");
6591        let mut spoofed = ProviderOptions::new();
6592        spoofed.insert("pond".to_owned(), json!({"ingest": {"host": "spoofed"}}));
6593        let message = Message::User {
6594            id: "message-1".to_owned(),
6595            session_id: session.id.clone(),
6596            timestamp: Utc::now(),
6597            options: spoofed,
6598        };
6599        let part = Part {
6600            session_id: session.id.clone(),
6601            id: "part-1".to_owned(),
6602            message_id: "message-1".to_owned(),
6603            ordinal: 0,
6604            provenance: crate::wire::Provenance::Conversational,
6605            options: ProviderOptions::new(),
6606            kind: PartKind::Text {
6607                text: Some(Extracted::from_test_value("hello".to_owned())),
6608            },
6609        };
6610
6611        let mut validator = IngestValidator::default();
6612        validator
6613            .push(&store, 0, IngestEvent::Session(session.clone()))
6614            .await?;
6615        validator
6616            .push(&store, 1, IngestEvent::Message(message))
6617            .await?;
6618        validator.push(&store, 2, IngestEvent::Part(part)).await?;
6619        validator.finish(&store).await?;
6620
6621        let stored = store
6622            .get_session(&session.id)
6623            .await?
6624            .expect("ingested session is readable");
6625        assert!(
6626            !stored.session.options.contains_key("pond"),
6627            "session rows are not stamped (attribution derives from messages)"
6628        );
6629        let stored_message = &stored.messages[0].message;
6630        match ingest_host_stamp() {
6631            Some(stamp) => {
6632                assert_eq!(
6633                    stored_message.options().get("pond"),
6634                    Some(stamp),
6635                    "stored message carries the real stamp, never the spoof"
6636                );
6637                let host = stamp
6638                    .pointer("/ingest/host")
6639                    .and_then(Value::as_object)
6640                    .expect("stamp shape is {ingest: {host: {..}}}");
6641                assert!(!host.is_empty(), "an all-empty stamp must be None instead");
6642                assert!(
6643                    host.values()
6644                        .all(|v| v.as_str().is_some_and(|s| !s.is_empty())),
6645                    "stamp fields are omitted when unavailable, never empty: {host:?}"
6646                );
6647            }
6648            None => assert!(
6649                stored_message.options().get("pond").is_none(),
6650                "with no resolvable stamp the spoofed key is still stripped"
6651            ),
6652        }
6653        assert!(
6654            !stored.messages[0].parts[0].options.contains_key("pond"),
6655            "part rows are not stamped (covered by their message's stamp)"
6656        );
6657
6658        Ok(())
6659    }
6660
6661    /// Regression: compact_files on `parts` with the blob column tripped a
6662    /// Lance v7.0.0-beta.16 dispatch bug under `lance.blob.v2`. Two upsert
6663    /// batches give compact fragments to merge; every `FileData` variant
6664    /// exercises the blob round-trip. All-File batches sidestep a debug-only
6665    /// `debug_assert_eq!` in Lance's legacy blob encoder that trips when one
6666    /// write batch mixes null + valid rows in the blob column - benign in
6667    /// release, irrelevant to this regression's scope.
6668    #[tokio::test(flavor = "multi_thread")]
6669    async fn optimize_indices_compacts_parts_with_blob_column() -> anyhow::Result<()> {
6670        use crate::wire::{FileData, PartKind, Provenance};
6671        let temp = TempDir::new()?;
6672        let store = Store::open_local(temp.path()).await?;
6673
6674        let session = synthetic_session("compact-blob");
6675        store
6676            .upsert_sessions(std::slice::from_ref(&session))
6677            .await?;
6678
6679        let make_part = |idx: usize, kind: PartKind| Part {
6680            session_id: session.id.clone(),
6681            message_id: format!("msg-{idx}"),
6682            id: format!("part-{idx}"),
6683            ordinal: 0,
6684            provenance: Provenance::Conversational,
6685            options: ProviderOptions::new(),
6686            kind,
6687        };
6688
6689        let batch_a = vec![
6690            make_part(
6691                0,
6692                PartKind::File {
6693                    media_type: Some("text/plain".to_owned()),
6694                    file_name: Some("a.txt".to_owned()),
6695                    data: FileData::Bytes(b"alpha".to_vec()),
6696                },
6697            ),
6698            make_part(
6699                1,
6700                PartKind::File {
6701                    media_type: Some("text/plain".to_owned()),
6702                    file_name: Some("b.txt".to_owned()),
6703                    data: FileData::String("beta".to_owned()),
6704                },
6705            ),
6706        ];
6707        store.upsert_parts(&batch_a).await?;
6708
6709        let batch_b = vec![
6710            make_part(
6711                2,
6712                PartKind::File {
6713                    media_type: Some("application/octet-stream".to_owned()),
6714                    file_name: None,
6715                    data: FileData::Url("https://example.com/file".to_owned()),
6716                },
6717            ),
6718            make_part(
6719                3,
6720                PartKind::File {
6721                    media_type: Some("image/png".to_owned()),
6722                    file_name: Some("c.png".to_owned()),
6723                    data: FileData::Bytes(vec![0x89, 0x50, 0x4e, 0x47]),
6724                },
6725            ),
6726        ];
6727        store.upsert_parts(&batch_b).await?;
6728
6729        store
6730            .optimize_indices(None, &MaintenancePolicy::always_compact())
6731            .await?
6732            .into_result()?;
6733
6734        Ok(())
6735    }
6736
6737    #[tokio::test]
6738    async fn file_part_blob_v2_round_trips_through_get() -> anyhow::Result<()> {
6739        let temp = TempDir::new()?;
6740        let store = Store::open_local(temp.path()).await?;
6741        let session = synthetic_session("blob");
6742        let message = Message::User {
6743            id: "message-1".to_owned(),
6744            session_id: session.id.clone(),
6745            timestamp: Utc::now(),
6746            options: ProviderOptions::new(),
6747        };
6748        let part = Part {
6749            session_id: session.id.clone(),
6750            id: "part-1".to_owned(),
6751            message_id: message.id().to_owned(),
6752            ordinal: 0,
6753            provenance: crate::wire::Provenance::Conversational,
6754            options: ProviderOptions::new(),
6755            kind: PartKind::File {
6756                media_type: Some("text/plain".to_owned()),
6757                file_name: Some("payload.txt".to_owned()),
6758                data: FileData::Bytes(b"pond".to_vec()),
6759            },
6760        };
6761
6762        let mut validator = IngestValidator::default();
6763        validator
6764            .push(&store, 0, IngestEvent::Session(session.clone()))
6765            .await?;
6766        validator
6767            .push(&store, 1, IngestEvent::Message(message.clone()))
6768            .await?;
6769        validator
6770            .push(&store, 2, IngestEvent::Part(part.clone()))
6771            .await?;
6772        validator.finish(&store).await?;
6773
6774        let stored = store
6775            .get_session(&session.id)
6776            .await?
6777            .expect("session should exist");
6778        let stored_part = &stored.messages[0].parts[0];
6779        assert_eq!(stored_part, &part);
6780
6781        Ok(())
6782    }
6783
6784    //
6785    // `Session.source_agent` and `Session.project` are immutable
6786    // post-first-write because `messages` denormalizes them at first
6787    // ingest; a silent overwrite would desync the denormalized
6788    // copies. pond core's `IngestValidator` probes the existing session
6789    // before the merge_insert and emits a per-row `validation_failed`
6790    // outcome with the typed field name when either changes. Other Session
6791    // fields (options, parent_session_id, created_at, parent_message_id)
6792    // re-write idempotently via merge_insert.
6793
6794    fn base_session() -> Session {
6795        Session {
6796            id: "01HXY00000000001".to_owned(),
6797            parent_session_id: None,
6798            parent_message_id: None,
6799            source_agent: "claude-code".to_owned(),
6800            created_at: Utc::now(),
6801            project: crate::adapter::Extracted::from_test_value("/home/me/proj".to_owned()),
6802            options: ProviderOptions::new(),
6803        }
6804    }
6805
6806    fn count_status(outcomes: &[RowOutcome], target: OutcomeStatus) -> usize {
6807        outcomes
6808            .iter()
6809            .filter(|outcome| outcome.status == target)
6810            .count()
6811    }
6812
6813    #[tokio::test(flavor = "multi_thread")]
6814    async fn re_ingesting_a_session_with_unchanged_immutable_fields_is_idempotent()
6815    -> anyhow::Result<()> {
6816        let temp = TempDir::new()?;
6817        let store = Store::open_local(temp.path()).await?;
6818
6819        let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
6820        assert_eq!(count_status(&first, OutcomeStatus::Inserted), 1);
6821
6822        let mut again = base_session();
6823        again.options.insert("title".to_owned(), json!("renamed"));
6824        let second = ingest_events(&store, vec![IngestEvent::Session(again)]).await?;
6825        assert_eq!(
6826            count_status(&second, OutcomeStatus::Error),
6827            0,
6828            "options is mutable; the re-ingest must not surface an error: {second:?}",
6829        );
6830        assert_eq!(
6831            count_status(&second, OutcomeStatus::Matched),
6832            1,
6833            "unchanged immutable fields must match-insert via merge_insert",
6834        );
6835
6836        Ok(())
6837    }
6838
6839    #[tokio::test(flavor = "multi_thread")]
6840    async fn re_ingesting_with_changed_source_agent_is_rejected() -> anyhow::Result<()> {
6841        let temp = TempDir::new()?;
6842        let store = Store::open_local(temp.path()).await?;
6843
6844        let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
6845        assert_eq!(count_status(&first, OutcomeStatus::Error), 0);
6846
6847        let mut tampered = base_session();
6848        tampered.source_agent = "codex-cli".to_owned();
6849        let second = ingest_events(&store, vec![IngestEvent::Session(tampered)]).await?;
6850        assert_eq!(count_status(&second, OutcomeStatus::Error), 1);
6851        let err_row = second
6852            .iter()
6853            .find(|outcome| outcome.status == OutcomeStatus::Error)
6854            .expect("error outcome present");
6855        let err = err_row.error.as_ref().expect("error body present");
6856        assert_eq!(err.field, Some("source_agent"));
6857        assert_eq!(err.reason, Some("immutable"));
6858
6859        // The stored row stayed on the original adapter - no silent rewrite.
6860        let stored = store
6861            .get_session(&base_session().id)
6862            .await?
6863            .expect("session row survives the rejected re-ingest");
6864        assert_eq!(stored.session.source_agent, "claude-code");
6865
6866        Ok(())
6867    }
6868
6869    #[tokio::test(flavor = "multi_thread")]
6870    async fn re_ingesting_with_changed_project_is_rejected() -> anyhow::Result<()> {
6871        let temp = TempDir::new()?;
6872        let store = Store::open_local(temp.path()).await?;
6873
6874        let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
6875        assert_eq!(count_status(&first, OutcomeStatus::Error), 0);
6876
6877        let mut tampered = base_session();
6878        tampered.project = crate::adapter::Extracted::from_test_value("/somewhere/else".to_owned());
6879        let second = ingest_events(&store, vec![IngestEvent::Session(tampered)]).await?;
6880        let err_row = second
6881            .iter()
6882            .find(|outcome| outcome.status == OutcomeStatus::Error)
6883            .expect("project change must surface an error outcome");
6884        assert_eq!(err_row.error.as_ref().unwrap().field, Some("project"));
6885
6886        let stored = store
6887            .get_session(&base_session().id)
6888            .await?
6889            .expect("session row survives");
6890        assert_eq!(
6891            stored.session.project.as_str(),
6892            "/home/me/proj",
6893            "stored project must remain the original",
6894        );
6895
6896        Ok(())
6897    }
6898
6899    #[tokio::test(flavor = "multi_thread")]
6900    async fn batched_flush_attributes_new_messages_on_existing_session() -> anyhow::Result<()> {
6901        // Regression guard: re-ingesting an existing session with NEW
6902        // messages must surface as sessions_inserted=0, messages_inserted_*>0
6903        // on `BatchCounts`, and per-row outcomes must mark the new message
6904        // rows `Inserted` while the session row is `Matched`. The prior
6905        // implementation derived all per-row statuses from the batch-level
6906        // session inserted count, which silently flipped the new messages
6907        // into `Matched` (visible as "up to date" in the CLI bar tail).
6908        use crate::wire::Provenance;
6909        let temp = TempDir::new()?;
6910        let store = Store::open_local(temp.path()).await?;
6911        let session = base_session();
6912
6913        let text_part = |part_id: &str, message_id: &str, body: &str| Part {
6914            session_id: session.id.clone(),
6915            id: part_id.to_owned(),
6916            message_id: message_id.to_owned(),
6917            ordinal: 0,
6918            provenance: Provenance::Conversational,
6919            options: ProviderOptions::new(),
6920            kind: PartKind::Text {
6921                text: Some(Extracted::from_test_value(body.to_owned())),
6922            },
6923        };
6924        let user_message = |id: &str| Message::User {
6925            id: id.to_owned(),
6926            session_id: session.id.clone(),
6927            timestamp: Utc::now(),
6928            options: ProviderOptions::new(),
6929        };
6930
6931        // First pass: 2 messages land fresh.
6932        let mut validator = IngestValidator::default();
6933        validator
6934            .push(&store, 0, IngestEvent::Session(session.clone()))
6935            .await?;
6936        validator
6937            .push(&store, 1, IngestEvent::Message(user_message("m1")))
6938            .await?;
6939        validator
6940            .push(&store, 2, IngestEvent::Part(text_part("p1", "m1", "alpha")))
6941            .await?;
6942        validator
6943            .push(&store, 3, IngestEvent::Message(user_message("m2")))
6944            .await?;
6945        validator
6946            .push(&store, 4, IngestEvent::Part(text_part("p2", "m2", "beta")))
6947            .await?;
6948        let (_first_outcomes, first_counts) = validator.finish(&store).await?;
6949        assert_eq!(first_counts.sessions_inserted, 1);
6950        assert_eq!(first_counts.messages_inserted_total, 2);
6951        assert_eq!(first_counts.messages_inserted_searchable, 2);
6952
6953        // Second pass: same session id, 3 NEW messages.
6954        let mut validator = IngestValidator::default();
6955        validator
6956            .push(&store, 0, IngestEvent::Session(session.clone()))
6957            .await?;
6958        for (idx, mid) in ["m3", "m4", "m5"].iter().enumerate() {
6959            let pid = format!("p{}", idx + 3);
6960            validator
6961                .push(&store, idx * 2 + 1, IngestEvent::Message(user_message(mid)))
6962                .await?;
6963            validator
6964                .push(
6965                    &store,
6966                    idx * 2 + 2,
6967                    IngestEvent::Part(text_part(&pid, mid, "gamma")),
6968                )
6969                .await?;
6970        }
6971        let (second_outcomes, second_counts) = validator.finish(&store).await?;
6972
6973        assert_eq!(
6974            second_counts.sessions_inserted, 0,
6975            "existing session row must report as Matched, not Inserted",
6976        );
6977        assert_eq!(second_counts.sessions_matched, 1);
6978        assert_eq!(
6979            second_counts.messages_inserted_total, 3,
6980            "the three NEW messages must register as Inserted in BatchCounts",
6981        );
6982        assert_eq!(
6983            second_counts.messages_inserted_searchable, 3,
6984            "all three new messages carry conversational text -> searchable",
6985        );
6986        assert_eq!(second_counts.messages_matched_total, 0);
6987        assert_eq!(second_counts.parts_inserted, 3);
6988        assert_eq!(second_counts.parts_matched, 0);
6989
6990        // Per-row outcomes mirror the BatchCounts shape: the session row is
6991        // Matched, every new message + part row is Inserted.
6992        let session_outcome = second_outcomes
6993            .iter()
6994            .find(|outcome| outcome.kind == "session")
6995            .expect("session-row outcome present");
6996        assert_eq!(session_outcome.status, OutcomeStatus::Matched);
6997        for outcome in &second_outcomes {
6998            if outcome.kind == "message" || outcome.kind == "part" {
6999                assert_eq!(
7000                    outcome.status,
7001                    OutcomeStatus::Inserted,
7002                    "new row must be Inserted, got: {outcome:?}",
7003                );
7004            }
7005        }
7006        Ok(())
7007    }
7008
7009    /// Ingest `count` synthetic messages spread across a handful of sessions
7010    /// and projects, each with conversational `search_text`. Returns the store
7011    /// and the message keys in `msg-{i}` order; every `vector` starts null.
7012    async fn store_with_messages(
7013        temp: &TempDir,
7014        count: usize,
7015    ) -> anyhow::Result<(Store, Vec<MessageKey>)> {
7016        store_with_messages_at_threshold(temp, count, VECTOR_INDEX_ACTIVATION_ROWS).await
7017    }
7018
7019    /// Same as [`store_with_messages`] but tests optimize with a custom
7020    /// IVF_SQ activation threshold.
7021    async fn store_with_messages_at_threshold(
7022        temp: &TempDir,
7023        count: usize,
7024        _vector_threshold: usize,
7025    ) -> anyhow::Result<(Store, Vec<MessageKey>)> {
7026        let store = Store::open_local(temp.path()).await?;
7027        let sessions = 8.min(count.max(1));
7028        let mut events = Vec::new();
7029        for s in 0..sessions {
7030            events.push(IngestEvent::Session(Session {
7031                id: format!("session-{s}"),
7032                parent_session_id: None,
7033                parent_message_id: None,
7034                source_agent: "claude-code".to_owned(),
7035                created_at: Utc::now(),
7036                project: Extracted::from_test_value(format!("/proj/{}", s % 4)),
7037                options: ProviderOptions::new(),
7038            }));
7039            for i in (s..count).step_by(sessions) {
7040                let message_id = format!("msg-{i}");
7041                events.push(IngestEvent::Message(Message::User {
7042                    id: message_id.clone(),
7043                    session_id: format!("session-{s}"),
7044                    timestamp: Utc::now(),
7045                    options: ProviderOptions::new(),
7046                }));
7047                events.push(IngestEvent::Part(Part {
7048                    session_id: format!("session-{s}"),
7049                    id: format!("{message_id}-part"),
7050                    message_id,
7051                    ordinal: 0,
7052                    provenance: crate::wire::Provenance::Conversational,
7053                    options: ProviderOptions::new(),
7054                    kind: PartKind::Text {
7055                        text: Some(Extracted::from_test_value(format!("synthetic message {i}"))),
7056                    },
7057                }));
7058            }
7059        }
7060        ingest_events(&store, events).await?;
7061        let keys = (0..count)
7062            .map(|i| MessageKey {
7063                session_id: format!("session-{}", i % sessions),
7064                message_id: format!("msg-{i}"),
7065            })
7066            .collect();
7067        Ok((store, keys))
7068    }
7069
7070    /// A deterministic pseudo-random vector of the production dimension.
7071    fn synthetic_vector(seed: usize) -> Vec<f32> {
7072        let mut state = (seed as u64)
7073            .wrapping_mul(0x9E37_79B9_7F4A_7C15)
7074            .wrapping_add(1);
7075        (0..embedding_dim())
7076            .map(|_| {
7077                state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
7078                #[allow(clippy::cast_precision_loss)]
7079                let unit = (state >> 33) as f32 / (1u64 << 31) as f32;
7080                unit - 1.0
7081            })
7082            .collect()
7083    }
7084
7085    /// One [`EmbeddedMessage`] per key, vectors seeded by slice position.
7086    fn embedded(keys: &[MessageKey]) -> Vec<EmbeddedMessage> {
7087        keys.iter()
7088            .enumerate()
7089            .map(|(seed, key)| EmbeddedMessage {
7090                session_id: key.session_id.clone(),
7091                id: key.message_id.clone(),
7092                vector: synthetic_vector(seed),
7093            })
7094            .collect()
7095    }
7096
7097    fn embedding_update_batch_with_model(
7098        rows: &[EmbeddedMessage],
7099        model: &str,
7100    ) -> Result<RecordBatch> {
7101        let mut batch = embedding_update_batch(rows)?;
7102        let columns = batch
7103            .columns()
7104            .iter()
7105            .take(3)
7106            .cloned()
7107            .chain(std::iter::once(
7108                Arc::new(StringArray::from(vec![model; rows.len()])) as _,
7109            ))
7110            .collect::<Vec<_>>();
7111        batch = RecordBatch::try_new(batch.schema(), columns)?;
7112        Ok(batch)
7113    }
7114
7115    #[tokio::test]
7116    async fn filtered_vector_scan_pushes_scalar_predicate_into_the_index() -> anyhow::Result<()> {
7117        let temp = TempDir::new()?;
7118        // 4 messages cycle session-0..session-3, so `session-3` is a real
7119        // partition. Scalar-index pushdown is volume-independent: the planner
7120        // emits `ScalarIndexQuery` whenever the index exists.
7121        let (store, keys) = store_with_messages(&temp, 4).await?;
7122        store.write_embeddings(&embedded(&keys)).await?;
7123        store
7124            .optimize_indices(None, &MaintenancePolicy::always_compact())
7125            .await?
7126            .into_result()?;
7127
7128        let query = vec![0.01_f32; embedding_dim()];
7129        let plan = store
7130            .explain_vector_plan(
7131                &query,
7132                10,
7133                &Predicate::Eq("session_id", "session-3".into()),
7134                None,
7135            )
7136            .await?;
7137
7138        // The load-bearing assertion (spec.md#search-prefilter-pushdown): the predicate
7139        // is served by a scalar-index node, not a postfilter `FilterExec`. (A
7140        // `FilterExec` for the KNN-internal `_distance IS NOT NULL` is expected
7141        // and unrelated.)
7142        assert!(
7143            plan.contains("ScalarIndexQuery"),
7144            "expected a ScalarIndexQuery node in the plan:\n{plan}",
7145        );
7146        let predicate_postfiltered = plan
7147            .lines()
7148            .any(|line| line.contains("FilterExec") && line.contains("session_id"));
7149        assert!(
7150            !predicate_postfiltered,
7151            "the scalar predicate must not fall back to a FilterExec postfilter:\n{plan}",
7152        );
7153        Ok(())
7154    }
7155
7156    #[tokio::test]
7157    async fn vector_index_activates_when_threshold_is_crossed() -> anyhow::Result<()> {
7158        let temp = TempDir::new()?;
7159        let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7160
7161        // First batch: 255 vectors, one below threshold. Optimize does not
7162        // create the IVF_SQ because the trigger is not met.
7163        store.write_embeddings(&embedded(&keys[..255])).await?;
7164        store
7165            .optimize_indices_with_vector_threshold(256)
7166            .await?
7167            .into_result()?;
7168        assert!(
7169            !store
7170                .handle
7171                .messages_index_names()
7172                .await?
7173                .iter()
7174                .any(|name| name == MESSAGES_VECTOR_INDEX),
7175            "IVF_SQ must not exist below the activation threshold",
7176        );
7177
7178        // Next batch: one more vector. Total reaches 256; optimize creates
7179        // the IVF_SQ.
7180        store.write_embeddings(&embedded(&keys[255..256])).await?;
7181        store
7182            .optimize_indices_with_vector_threshold(256)
7183            .await?
7184            .into_result()?;
7185        assert!(
7186            store
7187                .handle
7188                .messages_index_names()
7189                .await?
7190                .iter()
7191                .any(|name| name == MESSAGES_VECTOR_INDEX),
7192            "optimize must create the IVF_SQ once the threshold is crossed",
7193        );
7194
7195        // The remaining 44 rows stay un-embedded; the IVF_SQ trains over the
7196        // non-null subset and a planted vector is retrievable.
7197        let hits = store
7198            .vector_search(&synthetic_vector(0), 10, &Predicate::And(Vec::new()), None)
7199            .await?;
7200        assert!(
7201            hits.iter().any(|hit| hit.key == keys[0]),
7202            "an embedded row is retrievable via the index",
7203        );
7204        Ok(())
7205    }
7206
7207    #[tokio::test]
7208    async fn scalar_fold_batching_defers_tail_without_losing_rows() -> anyhow::Result<()> {
7209        let temp = TempDir::new()?;
7210        let (store, _keys) = store_with_messages(&temp, 300).await?;
7211
7212        // Threshold 0 folds every family: the scalar index covers all 300 rows.
7213        store
7214            .optimize_indices_with_scalar_fold_threshold(0)
7215            .await?
7216            .into_result()?;
7217        assert_eq!(
7218            store
7219                .handle
7220                .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7221                .await?,
7222            0,
7223            "threshold 0 must fold the scalar index over every row",
7224        );
7225
7226        // Append one small session -> a new unindexed fragment on messages.
7227        let new_messages = 5usize;
7228        let mut events = vec![IngestEvent::Session(Session {
7229            id: "session-new".to_owned(),
7230            parent_session_id: None,
7231            parent_message_id: None,
7232            source_agent: "claude-code".to_owned(),
7233            created_at: Utc::now(),
7234            project: Extracted::from_test_value("/proj/new".to_owned()),
7235            options: ProviderOptions::new(),
7236        })];
7237        for i in 0..new_messages {
7238            let message_id = format!("new-msg-{i}");
7239            events.push(IngestEvent::Message(Message::User {
7240                id: message_id.clone(),
7241                session_id: "session-new".to_owned(),
7242                timestamp: Utc::now(),
7243                options: ProviderOptions::new(),
7244            }));
7245            events.push(IngestEvent::Part(Part {
7246                session_id: "session-new".to_owned(),
7247                id: format!("{message_id}-part"),
7248                message_id,
7249                ordinal: 0,
7250                provenance: crate::wire::Provenance::Conversational,
7251                options: ProviderOptions::new(),
7252                kind: PartKind::Text {
7253                    text: Some(Extracted::from_test_value(format!("new text {i}"))),
7254                },
7255            }));
7256        }
7257        ingest_events(&store, events).await?;
7258
7259        // A threshold far above the delta defers the scalar fold; FTS (not
7260        // scalar) still folds every run.
7261        store
7262            .optimize_indices_with_scalar_fold_threshold(1_000_000)
7263            .await?
7264            .into_result()?;
7265        assert_eq!(
7266            store
7267                .handle
7268                .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7269                .await?,
7270            new_messages,
7271            "the scalar fold must be deferred, leaving the delta tail unindexed",
7272        );
7273        assert_eq!(
7274            store
7275                .handle
7276                .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7277                .await?,
7278            0,
7279            "FTS must fold every run regardless of the scalar threshold",
7280        );
7281
7282        // The deferred rows are still fully retrievable - get scans the tail.
7283        let session = store
7284            .get_session("session-new")
7285            .await?
7286            .expect("deferred session must still be retrievable");
7287        assert_eq!(
7288            session.messages.len(),
7289            new_messages,
7290            "get must read the unindexed scalar tail via scan",
7291        );
7292
7293        // A later threshold-0 fold consolidates the deferred tail.
7294        store
7295            .optimize_indices_with_scalar_fold_threshold(0)
7296            .await?
7297            .into_result()?;
7298        assert_eq!(
7299            store
7300                .handle
7301                .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7302                .await?,
7303            0,
7304            "a threshold-0 fold must consolidate the deferred tail",
7305        );
7306        Ok(())
7307    }
7308
7309    /// At the delta-merge threshold the FTS index must REBUILD, never merge:
7310    /// Lance 7.0.0's inverted merge fails two ways on real segments (posting
7311    /// tail codec mismatch on empty segments; index-out-of-bounds panic in
7312    /// `InnerBuilder::merge_from`), so consolidation routes FTS to the
7313    /// from-scratch rebuild path. This drives real grow -> fold cycles past
7314    /// the threshold and asserts an IndexRebuild phase fired, the segment
7315    /// chain collapsed to one, and the rebuilt index covers everything.
7316    #[tokio::test]
7317    async fn fts_consolidation_rebuilds_instead_of_merging() -> anyhow::Result<()> {
7318        use crate::substrate::{OptimizeEvent, OptimizePhase};
7319
7320        type PhaseLog = std::sync::Arc<std::sync::Mutex<Vec<(OptimizePhase, Option<String>)>>>;
7321        let temp = TempDir::new()?;
7322        let (store, _keys) = store_with_messages(&temp, 300).await?;
7323        let phases: PhaseLog = std::sync::Arc::default();
7324        let sink = phases.clone();
7325        let progress: crate::substrate::OptimizeProgressFn = std::sync::Arc::new(move |event| {
7326            if let OptimizeEvent::PhaseStart { phase, detail, .. } = event {
7327                sink.lock().unwrap().push((phase, detail));
7328            }
7329        });
7330
7331        for round in 0..=crate::substrate::DELTA_MERGE_THRESHOLD {
7332            ingest_events(&store, conversational_events(&format!("grow-{round}"), 2)).await?;
7333            store
7334                .optimize_indices(Some(progress.clone()), &MaintenancePolicy::always_compact())
7335                .await?
7336                .into_result()?;
7337        }
7338
7339        assert!(
7340            phases.lock().unwrap().iter().any(|(phase, detail)| {
7341                matches!(phase, OptimizePhase::IndexRebuild)
7342                    && detail.as_deref() == Some(MESSAGES_FTS_INDEX)
7343            }),
7344            "crossing the threshold must rebuild the FTS index, not merge it",
7345        );
7346        let fts_segments = store
7347            .handle
7348            .messages_index_names()
7349            .await?
7350            .into_iter()
7351            .filter(|name| name == MESSAGES_FTS_INDEX)
7352            .count();
7353        assert_eq!(fts_segments, 1, "the rebuild collapses the segment chain");
7354        assert_eq!(
7355            store
7356                .handle
7357                .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7358                .await?,
7359            0,
7360            "the rebuilt index covers every fragment",
7361        );
7362        Ok(())
7363    }
7364
7365    /// A flush whose every session row already exists must not commit the
7366    /// sessions table at all: the merge is insert-only, so the commit would
7367    /// be an empty manifest version paid on every steady-state sync.
7368    #[tokio::test]
7369    async fn grown_session_flush_skips_the_sessions_merge() -> anyhow::Result<()> {
7370        let temp = TempDir::new()?;
7371        let store = Store::open_local(temp.path()).await?;
7372        ingest_events(&store, conversational_events("session-grow", 1)).await?;
7373        let sessions_before = store.handle.dataset(Table::Sessions).await?.version_id();
7374        let messages_before = store.handle.dataset(Table::Messages).await?.version_id();
7375
7376        ingest_events(&store, conversational_events("session-grow", 2)).await?;
7377        assert_eq!(
7378            store.handle.dataset(Table::Sessions).await?.version_id(),
7379            sessions_before,
7380            "an all-present sessions batch must skip the merge commit",
7381        );
7382        assert!(
7383            store.handle.dataset(Table::Messages).await?.version_id() > messages_before,
7384            "the grown message rows must still commit",
7385        );
7386        let session = store
7387            .get_session("session-grow")
7388            .await?
7389            .expect("session row must survive the skipped merge");
7390        assert_eq!(session.messages.len(), 2);
7391        Ok(())
7392    }
7393
7394    /// A tail whose every row has a null `search_text` (tool-call-only
7395    /// messages) must not fold into the FTS index: Lance 7.0.0 writes an
7396    /// empty delta segment for it and reads that segment back with a
7397    /// mismatched posting-tail codec, deterministically failing every later
7398    /// merge. The guard skips the fold; the rows stay in the flat-scanned
7399    /// tail and get carried into the next fold that has real text.
7400    #[tokio::test]
7401    async fn fts_fold_skips_a_tail_with_no_indexable_text() -> anyhow::Result<()> {
7402        let temp = TempDir::new()?;
7403        let (store, _keys) = store_with_messages(&temp, 300).await?;
7404        store
7405            .optimize_indices_with_scalar_fold_threshold(0)
7406            .await?
7407            .into_result()?;
7408
7409        let tool_only_message = |i: usize| {
7410            let message_id = format!("tool-msg-{i}");
7411            [
7412                IngestEvent::Message(Message::Assistant {
7413                    id: message_id.clone(),
7414                    session_id: "session-toolonly".to_owned(),
7415                    timestamp: Utc::now(),
7416                    options: ProviderOptions::new(),
7417                }),
7418                IngestEvent::Part(Part {
7419                    session_id: "session-toolonly".to_owned(),
7420                    id: format!("{message_id}-part"),
7421                    message_id,
7422                    ordinal: 0,
7423                    provenance: crate::wire::Provenance::Conversational,
7424                    options: ProviderOptions::new(),
7425                    kind: PartKind::ToolCall {
7426                        call_id: Some(Extracted::from_test_value(format!("call-{i}"))),
7427                        name: Some(Extracted::from_test_value("Bash".to_owned())),
7428                        params: serde_json::json!({"command": "ls"}),
7429                        provider_executed: false,
7430                    },
7431                }),
7432            ]
7433        };
7434        let mut events = vec![IngestEvent::Session(synthetic_session("session-toolonly"))];
7435        events.extend((0..3).flat_map(tool_only_message));
7436        ingest_events(&store, events).await?;
7437
7438        store
7439            .optimize_indices_with_scalar_fold_threshold(0)
7440            .await?
7441            .into_result()?;
7442        assert_eq!(
7443            store
7444                .handle
7445                .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7446                .await?,
7447            3,
7448            "an all-null tail must not fold into the FTS index",
7449        );
7450        let fts_status = |statuses: Vec<crate::substrate::IndexStatus>| {
7451            statuses
7452                .into_iter()
7453                .find(|status| status.intent_name == MESSAGES_FTS_INDEX)
7454                .expect("FTS status present")
7455        };
7456        assert_eq!(
7457            fts_status(store.index_status().await?).unindexed_rows,
7458            3,
7459            "the raw view counts uncovered rows",
7460        );
7461        assert_eq!(
7462            fts_status(store.index_status_indexable().await?).unindexed_rows,
7463            0,
7464            "the indexable view treats an all-null tail as nothing pending",
7465        );
7466
7467        // One real text row lands: the next fold indexes the whole tail,
7468        // carrying the previously skipped rows - none are stranded.
7469        ingest_events(&store, conversational_events("session-toolonly", 1)).await?;
7470        store
7471            .optimize_indices_with_scalar_fold_threshold(0)
7472            .await?
7473            .into_result()?;
7474        assert_eq!(
7475            store
7476                .handle
7477                .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7478                .await?,
7479            0,
7480            "a fold with real text must consolidate the skipped rows too",
7481        );
7482        Ok(())
7483    }
7484
7485    /// f3 recall guard for the vector arm: with the FTS/vector fold batched, a
7486    /// row can sit in an unindexed tail. `vector_search` must still return it -
7487    /// the retriever drops `fast_search` when a tail exists so Lance ANN-probes
7488    /// the indexed base AND brute-forces the tail. (The FTS arm is covered by
7489    /// `tests/integration/search.rs::fts_search_covers_the_unindexed_tail`.)
7490    #[tokio::test]
7491    async fn vector_search_covers_the_unindexed_tail() -> anyhow::Result<()> {
7492        let temp = TempDir::new()?;
7493        let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7494        store.write_embeddings(&embedded(&keys)).await?;
7495        store
7496            .optimize_indices_with_vector_threshold(256)
7497            .await?
7498            .into_result()?;
7499        assert_eq!(
7500            store
7501                .handle
7502                .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
7503                .await?,
7504            0,
7505            "the IVF must cover the whole base after the fold",
7506        );
7507
7508        // Append one embedded row without folding -> an unindexed vector tail.
7509        let tail = MessageKey {
7510            session_id: "session-tail".to_owned(),
7511            message_id: "tail-msg".to_owned(),
7512        };
7513        ingest_events(
7514            &store,
7515            vec![
7516                IngestEvent::Session(Session {
7517                    id: tail.session_id.clone(),
7518                    parent_session_id: None,
7519                    parent_message_id: None,
7520                    source_agent: "claude-code".to_owned(),
7521                    created_at: Utc::now(),
7522                    project: Extracted::from_test_value("/proj/tail".to_owned()),
7523                    options: ProviderOptions::new(),
7524                }),
7525                IngestEvent::Message(Message::User {
7526                    id: tail.message_id.clone(),
7527                    session_id: tail.session_id.clone(),
7528                    timestamp: Utc::now(),
7529                    options: ProviderOptions::new(),
7530                }),
7531                IngestEvent::Part(Part {
7532                    session_id: tail.session_id.clone(),
7533                    id: format!("{}-part", tail.message_id),
7534                    message_id: tail.message_id.clone(),
7535                    ordinal: 0,
7536                    provenance: crate::wire::Provenance::Conversational,
7537                    options: ProviderOptions::new(),
7538                    kind: PartKind::Text {
7539                        text: Some(Extracted::from_test_value("tail body".to_owned())),
7540                    },
7541                }),
7542            ],
7543        )
7544        .await?;
7545        let tail_vector = synthetic_vector(9999);
7546        store
7547            .write_embeddings(&[EmbeddedMessage {
7548                session_id: tail.session_id.clone(),
7549                id: tail.message_id.clone(),
7550                vector: tail_vector.clone(),
7551            }])
7552            .await?;
7553        assert!(
7554            store
7555                .handle
7556                .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
7557                .await?
7558                > 0,
7559            "the appended row must be an unindexed vector tail (no fold ran)",
7560        );
7561
7562        // Querying with the tail's own vector: fast_search is dropped (tail
7563        // present), so the brute-force over the tail surfaces the exact match.
7564        // Under the old index-only gate this row would be invisible.
7565        let hits = store
7566            .vector_search(&tail_vector, 10, &Predicate::And(Vec::new()), None)
7567            .await?;
7568        assert!(
7569            hits.iter().any(|hit| hit.key == tail),
7570            "the unindexed tail row must be reachable via vector search (complete recall)",
7571        );
7572        Ok(())
7573    }
7574
7575    #[tokio::test]
7576    async fn model_swap_force_re_embeds_only_stale_rows_and_rebuilds_ivf_pq() -> anyhow::Result<()>
7577    {
7578        let temp = TempDir::new()?;
7579        let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7580        let old_rows = embedded(&keys);
7581        let old_batch = embedding_update_batch_with_model(&old_rows, "old-model")?;
7582        store
7583            .handle
7584            .merge_update(Table::Messages, old_batch, old_rows.len())
7585            .await?;
7586        store
7587            .optimize_indices_with_vector_threshold(256)
7588            .await?
7589            .into_result()?;
7590        assert!(
7591            store
7592                .handle
7593                .messages_index_names()
7594                .await?
7595                .iter()
7596                .any(|name| name == MESSAGES_VECTOR_INDEX),
7597            "IVF_SQ must exist before a model swap",
7598        );
7599        assert_eq!(store.stale_embedding_count().await?, keys.len());
7600
7601        store.drop_vector_index().await?;
7602        let mut pending = Vec::new();
7603        let stream = store.pending_or_stale_messages();
7604        tokio::pin!(stream);
7605        while let Some(row) = stream.next().await {
7606            pending.push(row?);
7607        }
7608        assert_eq!(
7609            pending.len(),
7610            keys.len(),
7611            "force stream should see stale rows"
7612        );
7613        store.write_embeddings(&embedded(&keys)).await?;
7614        assert_eq!(store.stale_embedding_count().await?, 0);
7615        store
7616            .optimize_indices_with_vector_threshold(256)
7617            .await?
7618            .into_result()?;
7619        assert!(
7620            store
7621                .handle
7622                .messages_index_names()
7623                .await?
7624                .iter()
7625                .any(|name| name == MESSAGES_VECTOR_INDEX),
7626            "optimize must rebuild IVF_SQ after force re-embed",
7627        );
7628
7629        let stream = store.pending_or_stale_messages();
7630        tokio::pin!(stream);
7631        assert!(stream.next().await.is_none(), "up-to-date rows are skipped");
7632        Ok(())
7633    }
7634
7635    #[tokio::test]
7636    async fn session_last_message_ids_come_from_durable_messages() -> anyhow::Result<()> {
7637        let temp = TempDir::new()?;
7638        let store = Store::open_local(temp.path()).await?;
7639        let session = synthetic_session("oracle");
7640        store
7641            .upsert_sessions(std::slice::from_ref(&session))
7642            .await?;
7643        let timestamp =
7644            chrono::DateTime::from_timestamp(1_700_000_000, 0).expect("valid timestamp");
7645        let message_a = Message::User {
7646            id: "oracle-a".to_owned(),
7647            session_id: session.id.clone(),
7648            timestamp,
7649            options: ProviderOptions::new(),
7650        };
7651        let message_b = Message::User {
7652            id: "oracle-b".to_owned(),
7653            session_id: session.id.clone(),
7654            timestamp,
7655            options: ProviderOptions::new(),
7656        };
7657        store
7658            .upsert_messages(
7659                &session,
7660                &[
7661                    MessageWrite {
7662                        message: &message_a,
7663                        parts: &[],
7664                        search_text: Some("a"),
7665                    },
7666                    MessageWrite {
7667                        message: &message_b,
7668                        parts: &[],
7669                        search_text: Some("b"),
7670                    },
7671                ],
7672            )
7673            .await?;
7674
7675        let empty_session = synthetic_session("session-row-only");
7676        store.upsert_sessions(&[empty_session]).await?;
7677
7678        // Orphan: messages committed but the session row never was (the crash
7679        // window `upsert_session_batch`'s write order can leave). The gate must
7680        // NOT key on it, so the source re-ingests and heals the missing row.
7681        let orphan = synthetic_session("messages-no-row");
7682        let orphan_message = Message::User {
7683            id: "orphan-a".to_owned(),
7684            session_id: orphan.id.clone(),
7685            timestamp,
7686            options: ProviderOptions::new(),
7687        };
7688        store
7689            .upsert_messages(
7690                &orphan,
7691                &[MessageWrite {
7692                    message: &orphan_message,
7693                    parts: &[],
7694                    search_text: Some("a"),
7695                }],
7696            )
7697            .await?;
7698
7699        let map = store.session_last_message_ids().await?;
7700        assert_eq!(map.get("oracle").map(String::as_str), Some("oracle-b"));
7701        assert!(
7702            !map.contains_key("session-row-only"),
7703            "a session row without durable messages must not produce a freshness key",
7704        );
7705        assert!(
7706            !map.contains_key("messages-no-row"),
7707            "messages without a durable session row must not produce a freshness key",
7708        );
7709        Ok(())
7710    }
7711
7712    #[tokio::test]
7713    async fn embedding_progress_counts_embedded_and_eligible_rows() -> anyhow::Result<()> {
7714        let temp = TempDir::new()?;
7715        let (store, keys) = store_with_messages(&temp, 10).await?;
7716
7717        let before = store.embedding_progress().await?;
7718        assert_eq!(before.embedded, 0);
7719        assert_eq!(before.total, 10);
7720        assert_eq!(before.backlog, 10);
7721        assert_eq!(before.model, crate::embed::model_id());
7722
7723        store.write_embeddings(&embedded(&keys[..4])).await?;
7724        let partial = store.embedding_progress().await?;
7725        assert_eq!(partial.embedded, 4);
7726        assert_eq!(partial.total, 10);
7727        assert_eq!(partial.backlog, 6);
7728
7729        store.write_embeddings(&embedded(&keys[4..])).await?;
7730        let full = store.embedding_progress().await?;
7731        assert_eq!(full.embedded, 10);
7732        assert_eq!(full.total, 10);
7733        // The pending signal is the live un-embedded count and matches the
7734        // authoritative backlog - never derived from FTS num_docs.
7735        assert_eq!(full.backlog, 0);
7736        assert_eq!(full.backlog, store.embed_backlog_count().await?);
7737        Ok(())
7738    }
7739
7740    #[tokio::test]
7741    async fn load_rowmap_if_present_installs_published_chain_without_building() -> anyhow::Result<()>
7742    {
7743        let temp = TempDir::new()?;
7744        let (builder, _keys) = store_with_messages(&temp, 6).await?;
7745        let cache = temp.path().join("cache");
7746
7747        // No chain published yet: a load-only reader installs nothing and does
7748        // not build one.
7749        let reader = Store::open_local(temp.path()).await?;
7750        reader.load_rowmap_if_present(&cache).await?;
7751        assert!(reader.rowmap_snapshot().is_none());
7752
7753        // A sibling publishes the chain; the reader then installs it as-is.
7754        builder.ensure_rowmap(&cache).await?;
7755        reader.load_rowmap_if_present(&cache).await?;
7756        assert!(reader.rowmap_snapshot().is_some());
7757        Ok(())
7758    }
7759
7760    #[tokio::test]
7761    async fn ensure_rowmap_layers_a_delta_on_new_ingest() -> anyhow::Result<()> {
7762        let temp = TempDir::new()?;
7763        let (store, _keys) = store_with_messages(&temp, 6).await?;
7764        let cache = temp.path().join("cache");
7765
7766        store.ensure_rowmap(&cache).await?;
7767        assert_eq!(
7768            store.rowmap_delta_count(),
7769            Some(0),
7770            "first build is a lone base"
7771        );
7772
7773        // A new session's message bumps the version with a fresh fragment.
7774        ingest_events(
7775            &store,
7776            vec![
7777                IngestEvent::Session(Session {
7778                    id: "session-new".to_owned(),
7779                    parent_session_id: None,
7780                    parent_message_id: None,
7781                    source_agent: "claude-code".to_owned(),
7782                    created_at: Utc::now(),
7783                    project: Extracted::from_test_value("/proj/new".to_owned()),
7784                    options: ProviderOptions::new(),
7785                }),
7786                IngestEvent::Message(Message::User {
7787                    id: "m-new".to_owned(),
7788                    session_id: "session-new".to_owned(),
7789                    timestamp: Utc::now(),
7790                    options: ProviderOptions::new(),
7791                }),
7792                IngestEvent::Part(Part {
7793                    session_id: "session-new".to_owned(),
7794                    id: "m-new-part".to_owned(),
7795                    message_id: "m-new".to_owned(),
7796                    ordinal: 0,
7797                    provenance: crate::wire::Provenance::Conversational,
7798                    options: ProviderOptions::new(),
7799                    kind: PartKind::Text {
7800                        text: Some(Extracted::from_test_value("brand new message".to_owned())),
7801                    },
7802                }),
7803            ],
7804        )
7805        .await?;
7806
7807        // The refresh scans only the new fragment and layers a delta - not a
7808        // full rebuild.
7809        store.ensure_rowmap(&cache).await?;
7810        assert_eq!(
7811            store.rowmap_delta_count(),
7812            Some(1),
7813            "new ingest layered a delta"
7814        );
7815
7816        // The new session's count is served from the chain (base + delta sum).
7817        let counts = store
7818            .session_message_counts(&["session-new".to_owned()])
7819            .await?;
7820        assert_eq!(counts.get("session-new").copied(), Some(1));
7821        Ok(())
7822    }
7823
7824    /// Regression for the v0.10.0 sync death-spiral: the 1h cleanup retention can
7825    /// reclaim the dataset version the on-disk chain was last built at. The delta
7826    /// extender's `checkout_version(base)` then errored, the error nuked
7827    /// `ensure_rowmap`, and the oracle silently fell back to re-reading every
7828    /// source on every sync forever (the chain never advanced past the reclaimed
7829    /// base). A reclaimed base must degrade to a full rebuild, like compaction.
7830    #[tokio::test]
7831    async fn ensure_rowmap_rebuilds_when_base_manifest_reclaimed() -> anyhow::Result<()> {
7832        let temp = TempDir::new()?;
7833        let (store, _keys) = store_with_messages(&temp, 6).await?;
7834        let cache = temp.path().join("cache");
7835
7836        // Build the chain at the current version, then snapshot the manifests
7837        // that exist at-or-below it - these are exactly what cleanup reclaims.
7838        store.ensure_rowmap(&cache).await?;
7839        assert_eq!(store.rowmap_delta_count(), Some(0), "first build is a base");
7840        let base_version = store.messages_version().await?;
7841        let versions_dir = temp.path().join("messages.lance").join("_versions");
7842        let base_manifests: Vec<_> = std::fs::read_dir(&versions_dir)?
7843            .filter_map(|entry| entry.ok().map(|entry| entry.path()))
7844            .filter(|path| path.extension().is_some_and(|ext| ext == "manifest"))
7845            .collect();
7846        assert!(
7847            !base_manifests.is_empty(),
7848            "the base version has a manifest"
7849        );
7850
7851        // A new session bumps the version, so the on-disk chain now trails the
7852        // dataset and a refresh would normally delta from `base_version`.
7853        ingest_events(
7854            &store,
7855            vec![
7856                IngestEvent::Session(Session {
7857                    id: "session-after".to_owned(),
7858                    parent_session_id: None,
7859                    parent_message_id: None,
7860                    source_agent: "claude-code".to_owned(),
7861                    created_at: Utc::now(),
7862                    project: Extracted::from_test_value("/proj/after".to_owned()),
7863                    options: ProviderOptions::new(),
7864                }),
7865                IngestEvent::Message(Message::User {
7866                    id: "m-after".to_owned(),
7867                    session_id: "session-after".to_owned(),
7868                    timestamp: Utc::now(),
7869                    options: ProviderOptions::new(),
7870                }),
7871                IngestEvent::Part(Part {
7872                    session_id: "session-after".to_owned(),
7873                    id: "m-after-part".to_owned(),
7874                    message_id: "m-after".to_owned(),
7875                    ordinal: 0,
7876                    provenance: crate::wire::Provenance::Conversational,
7877                    options: ProviderOptions::new(),
7878                    kind: PartKind::Text {
7879                        text: Some(Extracted::from_test_value("after the base".to_owned())),
7880                    },
7881                }),
7882            ],
7883        )
7884        .await?;
7885        assert!(
7886            store.messages_version().await? > base_version,
7887            "the new ingest advanced the dataset past the chain's base"
7888        );
7889
7890        // Reclaim the base version's manifest exactly as `cleanup_old_versions`
7891        // would: `checkout_version(base_version)` can no longer resolve.
7892        for manifest in &base_manifests {
7893            std::fs::remove_file(manifest)?;
7894        }
7895
7896        // A fresh Store finds the trailing chain on disk, tries to delta from the
7897        // reclaimed base, and must fall back to a full rebuild - Ok, not Err.
7898        let reopened = Store::open_local(temp.path()).await?;
7899        reopened.ensure_rowmap(&cache).await?;
7900        assert!(
7901            reopened.rowmap_snapshot().is_some(),
7902            "map rebuilt after the base manifest was reclaimed"
7903        );
7904        assert_eq!(
7905            reopened.rowmap_delta_count(),
7906            Some(0),
7907            "a reclaimed base forces a fresh full-scan base, not a stuck chain"
7908        );
7909
7910        // The rebuilt base covers the post-base ingest, so the oracle is whole.
7911        let counts = reopened
7912            .session_message_counts(&["session-after".to_owned()])
7913            .await?;
7914        assert_eq!(counts.get("session-after").copied(), Some(1));
7915        Ok(())
7916    }
7917
7918    /// The steady-state hot path: embedding rewrites the message fragments every
7919    /// sync (merge_update on the `vector` column). Keying the delta off fragment
7920    /// identity made that rewrite force a full 2.1M-row rebuild every sync.
7921    /// Stable row ids preserve row_ids across the rewrite, so the refresh must
7922    /// layer a cheap append-only delta of just the new rows - and must NOT
7923    /// double-count the rewritten rows that still live in the base.
7924    #[tokio::test]
7925    async fn ensure_rowmap_deltas_across_embedding_fragment_rewrite() -> anyhow::Result<()> {
7926        let temp = TempDir::new()?;
7927        let (store, keys) = store_with_messages(&temp, 6).await?;
7928        let cache = temp.path().join("cache");
7929        store.ensure_rowmap(&cache).await?;
7930        assert_eq!(store.rowmap_delta_count(), Some(0), "first build is a base");
7931
7932        // Embedding rewrites every message fragment (new fragment ids, same
7933        // stable row_ids, untouched ROW_META columns).
7934        store.write_embeddings(&embedded(&keys)).await?;
7935
7936        // A new session appends one row on top of the rewritten fragments.
7937        ingest_events(
7938            &store,
7939            vec![
7940                IngestEvent::Session(Session {
7941                    id: "session-after".to_owned(),
7942                    parent_session_id: None,
7943                    parent_message_id: None,
7944                    source_agent: "claude-code".to_owned(),
7945                    created_at: Utc::now(),
7946                    project: Extracted::from_test_value("/proj/after".to_owned()),
7947                    options: ProviderOptions::new(),
7948                }),
7949                IngestEvent::Message(Message::User {
7950                    id: "m-after".to_owned(),
7951                    session_id: "session-after".to_owned(),
7952                    timestamp: Utc::now(),
7953                    options: ProviderOptions::new(),
7954                }),
7955                IngestEvent::Part(Part {
7956                    session_id: "session-after".to_owned(),
7957                    id: "m-after-part".to_owned(),
7958                    message_id: "m-after".to_owned(),
7959                    ordinal: 0,
7960                    provenance: crate::wire::Provenance::Conversational,
7961                    options: ProviderOptions::new(),
7962                    kind: PartKind::Text {
7963                        text: Some(Extracted::from_test_value("after embedding".to_owned())),
7964                    },
7965                }),
7966            ],
7967        )
7968        .await?;
7969
7970        // The refresh layers a delta of just the appended row, not a full
7971        // rebuild - despite every prior fragment having been rewritten.
7972        store.ensure_rowmap(&cache).await?;
7973        assert_eq!(
7974            store.rowmap_delta_count(),
7975            Some(1),
7976            "fragment rewrite + append must layer a delta, not full-rebuild"
7977        );
7978
7979        // Counts stay honest: the rewritten base rows are not re-emitted into the
7980        // delta, so nothing is double-counted across base + delta segments.
7981        let counts = store
7982            .session_message_counts(&["session-after".to_owned(), "session-0".to_owned()])
7983            .await?;
7984        assert_eq!(counts.get("session-after").copied(), Some(1));
7985        assert_eq!(
7986            counts.get("session-0").copied(),
7987            Some(1),
7988            "a base row survived the rewrite without being double-counted"
7989        );
7990        Ok(())
7991    }
7992
7993    #[tokio::test]
7994    async fn rowmap_chain_compacts_and_stays_bounded() -> anyhow::Result<()> {
7995        // Many version bumps (the remote-writers case) must not grow the chain
7996        // unboundedly: deltas cap at MAX, then compact into a fresh base.
7997        let temp = TempDir::new()?;
7998        let (store, _keys) = store_with_messages(&temp, 4).await?;
7999        let cache = temp.path().join("cache");
8000        store.ensure_rowmap(&cache).await?;
8001
8002        let mut reached_cap = false;
8003        let mut compacted = false;
8004        for i in 0..(Store::MAX_ROWMAP_DELTAS + 2) {
8005            let session = format!("session-x{i}");
8006            ingest_events(
8007                &store,
8008                vec![
8009                    IngestEvent::Session(Session {
8010                        id: session.clone(),
8011                        parent_session_id: None,
8012                        parent_message_id: None,
8013                        source_agent: "claude-code".to_owned(),
8014                        created_at: Utc::now(),
8015                        project: Extracted::from_test_value("/proj/x".to_owned()),
8016                        options: ProviderOptions::new(),
8017                    }),
8018                    IngestEvent::Message(Message::User {
8019                        id: format!("mx{i}"),
8020                        session_id: session.clone(),
8021                        timestamp: Utc::now(),
8022                        options: ProviderOptions::new(),
8023                    }),
8024                    IngestEvent::Part(Part {
8025                        session_id: session.clone(),
8026                        id: format!("mx{i}-part"),
8027                        message_id: format!("mx{i}"),
8028                        ordinal: 0,
8029                        provenance: crate::wire::Provenance::Conversational,
8030                        options: ProviderOptions::new(),
8031                        kind: PartKind::Text {
8032                            text: Some(Extracted::from_test_value(format!("msg {i}"))),
8033                        },
8034                    }),
8035                ],
8036            )
8037            .await?;
8038            store.ensure_rowmap(&cache).await?;
8039            let deltas = store.rowmap_delta_count().unwrap();
8040            assert!(
8041                deltas <= Store::MAX_ROWMAP_DELTAS,
8042                "delta count {deltas} exceeded the cap",
8043            );
8044            if deltas == Store::MAX_ROWMAP_DELTAS {
8045                reached_cap = true;
8046            }
8047            if reached_cap && deltas < Store::MAX_ROWMAP_DELTAS {
8048                compacted = true;
8049            }
8050        }
8051        assert!(reached_cap, "deltas accumulated to the cap (append path)");
8052        assert!(compacted, "the chain compacted back into a base");
8053
8054        // Files stay bounded and no build temps leak.
8055        let mut rmm = 0;
8056        for entry in std::fs::read_dir(&cache)? {
8057            let name = entry?.file_name().into_string().unwrap_or_default();
8058            assert!(!name.contains(".tmp-"), "leaked build temp: {name}");
8059            if name.ends_with(".rmm") {
8060                rmm += 1;
8061            }
8062        }
8063        assert!(
8064            rmm <= Store::MAX_ROWMAP_DELTAS + 1,
8065            "files unbounded: {rmm}"
8066        );
8067        Ok(())
8068    }
8069
8070    #[tokio::test]
8071    async fn embed_backlog_count_tracks_eligible_unembedded_rows() -> anyhow::Result<()> {
8072        let temp = TempDir::new()?;
8073        let (store, keys) = store_with_messages(&temp, 10).await?;
8074
8075        // Read straight from the dataset (no FTS index here), so it is correct
8076        // right after ingest - the case that lagged `embedding_progress`.
8077        assert_eq!(store.embed_backlog_count().await?, 10);
8078
8079        store.write_embeddings(&embedded(&keys[..4])).await?;
8080        assert_eq!(store.embed_backlog_count().await?, 6);
8081
8082        store.write_embeddings(&embedded(&keys[4..])).await?;
8083        assert_eq!(store.embed_backlog_count().await?, 0);
8084        Ok(())
8085    }
8086
8087    // The post-copy shape: every row embedded, none folded into IVF_SQ yet
8088    // (here: no index at all, which reports the whole table). The lag must
8089    // over-state - never under-state - so the embed gate can trust lag == 0
8090    // as "nothing unembedded" and must confirm lag > 0 with the exact count.
8091    #[tokio::test]
8092    async fn unindexed_vector_backlog_over_states_when_embedded_rows_are_unfolded()
8093    -> anyhow::Result<()> {
8094        let temp = TempDir::new()?;
8095        let (store, keys) = store_with_messages(&temp, 10).await?;
8096        store.write_embeddings(&embedded(&keys)).await?;
8097
8098        assert_eq!(store.unindexed_vector_backlog().await?, 10);
8099        assert_eq!(store.embed_backlog_count().await?, 0);
8100        Ok(())
8101    }
8102
8103    #[tokio::test]
8104    async fn session_message_counts_returns_per_session_counts_with_zeros_for_unknown_sessions()
8105    -> anyhow::Result<()> {
8106        // store_with_messages stripes `count` messages across 8 sessions
8107        // round-robin. 32 messages -> 4 per session, 0..8 deterministic.
8108        let temp = TempDir::new()?;
8109        let (store, _keys) = store_with_messages(&temp, 32).await?;
8110
8111        let mut requested: Vec<String> = (0..8).map(|s| format!("session-{s}")).collect();
8112        requested.push("session-unknown-a".to_owned());
8113        requested.push("session-unknown-b".to_owned());
8114        let counts = store.session_message_counts(&requested).await?;
8115
8116        // Map has an entry for every requested id (the contract): known
8117        // sessions hit 4, unknown sessions sit at 0.
8118        assert_eq!(counts.len(), requested.len());
8119        for s in 0..8 {
8120            assert_eq!(
8121                counts.get(&format!("session-{s}")).copied(),
8122                Some(4),
8123                "session-{s} should have 4 messages",
8124            );
8125        }
8126        assert_eq!(counts.get("session-unknown-a").copied(), Some(0));
8127        assert_eq!(counts.get("session-unknown-b").copied(), Some(0));
8128
8129        // Empty input is the documented zero-path.
8130        let empty = store.session_message_counts(&[]).await?;
8131        assert!(empty.is_empty());
8132        Ok(())
8133    }
8134}