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