1use 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 rowmap: ArcSwapOption<RowMetaSet>,
54 embedder: Option<Arc<crate::embed::LazyEmbedder>>,
59 ingest_embed_progress: Option<IngestEmbedProgress>,
63}
64
65#[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#[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#[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#[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 pub fn new_sessions(&self) -> usize {
158 self.sessions.append.len()
159 }
160
161 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#[derive(Debug, Clone, PartialEq)]
194pub struct PendingMessage {
195 pub session_id: String,
196 pub id: String,
197 pub search_text: String,
198}
199
200#[derive(Debug, Clone, PartialEq)]
203pub struct EmbeddedMessage {
204 pub session_id: String,
205 pub id: String,
206 pub vector: Vec<f32>,
207}
208
209#[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#[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#[derive(Debug, Default)]
249pub struct OptimizeOutcome {
250 pub tables: Vec<TableOptimizeOutcome>,
251}
252
253impl OptimizeOutcome {
254 pub fn any_indices_failed(&self) -> bool {
257 self.tables.iter().any(|t| t.indices.is_failed())
258 }
259
260 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#[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 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 #[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 #[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 pub fn lance_cache_bytes(&self) -> u64 {
347 self.handle.lance_cache_bytes()
348 }
349
350 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 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 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 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 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 let _ = self.handle.dataset(table).await?;
516 self.merge_scanner(table, dataset.scan(), "archive import")
517 .await
518 }
519
520 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 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 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 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 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 pub async fn copy_delta_from(
676 &self,
677 source: &Store,
678 plan: &DeltaPlan,
679 ) -> Result<LanceArchiveImport> {
680 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
1723 self.handle.collect_ids(table).await
1724 }
1725
1726 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 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 pub async fn dataset(&self, table: Table) -> Result<Arc<Dataset>> {
1779 Ok(Arc::new(self.handle.dataset(table).await?))
1780 }
1781
1782 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 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 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 pub async fn prune_index_cache(&self, cache_dir: &Path) {
1816 self.handle.prune_index_cache(cache_dir).await;
1817 }
1818
1819 fn store_key(&self) -> String {
1822 crate::substrate::store_key(self.handle.location())
1823 }
1824
1825 const MAX_ROWMAP_DELTAS: usize = 16;
1827
1828 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 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 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 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 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 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 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 Self::sweep_orphan_temps(cache_dir, store_key);
1945
1946 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 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 (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 (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 _ => {
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 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 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 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 pub fn rowmap_snapshot(&self) -> Option<Arc<RowMetaSet>> {
2085 self.rowmap.load_full()
2086 }
2087
2088 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 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 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 pub async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
2145 self.handle.export_write(name, bytes).await
2146 }
2147
2148 pub async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
2150 self.handle.export_read(name).await
2151 }
2152
2153 pub fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2155 self.handle.export_local_path(name)
2156 }
2157
2158 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 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 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 pub fn pending_embedding_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2249 try_stream! {
2250 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 pub fn pending_or_stale_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2291 try_stream! {
2292 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 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 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 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 scanner.disable_scoring_autoprojection();
2387 scanner.limit(Some(i64::try_from(limit).unwrap_or(i64::MAX)), None)?;
2388 Ok(scanner)
2389 }
2390
2391 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 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 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 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 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 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 pub async fn searchable_in_scope(&self, filter: &Predicate) -> Result<usize> {
2535 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 pub async fn embedding_progress(&self) -> Result<EmbeddingProgress> {
2972 let dataset = self.handle.dataset(Table::Messages).await?;
2973 let embedded = dataset
2979 .count_rows(Some(Predicate::IsNotNull("embedding_model").to_lance()))
2980 .await?;
2981 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 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 pub async fn stale_embedding_count(&self) -> Result<usize> {
3014 let dataset = self.handle.dataset(Table::Messages).await?;
3015 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 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 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 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 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 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 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 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 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 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 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 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 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#[derive(Debug, Clone, PartialEq, Eq, Default)]
3421pub struct IngestSummary {
3422 pub inserted: usize,
3426 pub matched: usize,
3428 pub sessions_inserted: usize,
3430 pub messages_inserted_total: usize,
3433 pub messages_inserted_searchable: usize,
3437 pub parts_inserted: usize,
3439 pub sessions_matched: usize,
3441 pub messages_matched_total: usize,
3443 pub messages_matched_searchable: usize,
3445 pub parts_matched: usize,
3447 pub dropped_events: usize,
3457 pub dropped_sessions: usize,
3462 pub skipped_files: usize,
3465 pub skipped_empty: usize,
3470 pub skipped_fresh: usize,
3474 pub storage_errors: usize,
3478 pub truncated_values: usize,
3481 pub drop_reasons: BTreeMap<&'static str, usize>,
3487}
3488
3489pub 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#[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 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 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 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 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#[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 pub searchable: bool,
3667}
3668
3669#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3670pub enum OutcomeStatus {
3671 Inserted,
3672 Matched,
3673 Error,
3674}
3675
3676#[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 pub reason_key: Option<&'static str>,
3688}
3689
3690#[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#[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 seen_message_ids: HashSet<String>,
3739 seen_part_keys: HashSet<(String, String)>,
3742 completed: Vec<CompletedSubstream>,
3746}
3747
3748#[derive(Debug)]
3750struct CompletedSubstream {
3751 session_index: usize,
3752 session: Session,
3753 messages: Vec<BufferedMessage>,
3754}
3755
3756fn 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 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 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 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 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 self.close_current_substream();
3837
3838 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 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 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
4080fn 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
4108fn 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 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
4249fn 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 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#[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 pub after_message_id: Option<&'a str>,
4373 pub before_message_id: Option<&'a str>,
4375 pub limit: usize,
4376 pub budget_bytes: usize,
4377 pub session_from: SessionFrom,
4379}
4380
4381#[derive(Debug, Clone)]
4382pub struct MessageViewParams {
4383 pub context_before: usize,
4385 pub context_after: usize,
4387 pub budget_bytes: usize,
4388}
4389
4390#[derive(Debug, Clone, PartialEq)]
4396pub enum GetLookup<T> {
4397 NotFound,
4398 UnknownAnchor,
4399 Found(T),
4400}
4401
4402#[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#[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#[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
4455fn 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
4474fn 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
4511const 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
4536const 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
4551const SESSIONS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] =
4554 &[("id", BuiltinIndexType::BTree, "sessions_id_btree")];
4555
4556const 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
4568fn embedded_scope(filter: &Predicate) -> Predicate {
4581 filter.clone()
4582}
4583
4584pub const DEFAULT_NPROBES: usize = 32;
4589
4590fn 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
4606pub(crate) const SESSIONS: &str = "sessions";
4610pub(crate) const MESSAGES: &str = "messages";
4611pub(crate) const PARTS: &str = "parts";
4612
4613pub const MESSAGES_SESSION_ID_INDEX: &str = "messages_session_id_btree";
4616
4617pub const MESSAGES_FTS_INDEX: &str = "messages_search_text_fts";
4620
4621pub const MESSAGES_VECTOR_INDEX: &str = "messages_vector_ivfpq";
4628
4629const IVF_SQ_NUM_BITS: u16 = 8;
4634const IVF_SQ_MAX_ITERS: usize = 15;
4635
4636pub fn pond_index_intents() -> IndexIntents {
4639 pond_index_intents_with_vector_threshold(VECTOR_INDEX_ACTIVATION_ROWS)
4640}
4641
4642pub(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
4698pub const DEFAULT_EMBEDDING_DIM: usize = 384;
4702
4703static EMBEDDING_DIM_RUNTIME: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4709
4710pub fn embedding_dim() -> usize {
4713 EMBEDDING_DIM_RUNTIME
4714 .get()
4715 .copied()
4716 .unwrap_or(DEFAULT_EMBEDDING_DIM)
4717}
4718
4719pub fn init_embedding_dim(dim: usize) {
4721 EMBEDDING_DIM_RUNTIME.get_or_init(|| dim);
4722}
4723
4724pub(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
4777pub(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
4790fn 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 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 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
4882fn embedding_vector_type() -> DataType {
4888 DataType::FixedSizeList(
4889 Arc::new(Field::new("item", DataType::Float16, true)),
4890 embedding_dim() as i32,
4891 )
4892}
4893
4894fn 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
4906fn embedding_columns(vectors: &[Option<Vec<f32>>]) -> Result<(ArrayRef, ArrayRef)> {
4910 let dim = embedding_dim();
4911 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
4952pub(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
4990const COLUMN_BYTE_BUDGET: usize = 1 << 30;
4995
4996fn 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
5119pub(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_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 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 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
5346pub 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: 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 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
5588fn 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 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 #[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 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 #[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 #[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 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 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 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 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 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 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 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 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 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 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 #[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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 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 #[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 #[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 #[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 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 #[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 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 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 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 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 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 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 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 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 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 #[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 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 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 for manifest in &base_manifests {
7538 std::fs::remove_file(manifest)?;
7539 }
7540
7541 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 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 #[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 store.write_embeddings(&embedded(&keys)).await?;
7580
7581 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 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 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 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 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 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 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 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 let empty = store.session_message_counts(&[]).await?;
7760 assert!(empty.is_empty());
7761 Ok(())
7762 }
7763}