Skip to main content

pond/
sessions.rs

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