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