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_upgrade = archive_schema_backfill(&sessions_dataset, Table::Sessions)?;
446 let messages_upgrade = archive_schema_backfill(&messages_dataset, Table::Messages)?;
447 let parts_upgrade = archive_schema_backfill(&parts_dataset, Table::Parts)?;
448 let (sessions, sessions_inserted) = self
449 .import_clean_table(Table::Sessions, sessions_dataset, sessions_upgrade)
450 .await?;
451 let (messages, messages_inserted) = self
452 .import_clean_table(Table::Messages, messages_dataset, messages_upgrade)
453 .await?;
454 let (parts, parts_inserted) = self
455 .import_clean_table(Table::Parts, parts_dataset, parts_upgrade)
456 .await?;
457 Ok(LanceArchiveImport {
458 rows: LanceArchiveCounts {
459 sessions,
460 messages,
461 parts,
462 },
463 inserted: LanceArchiveCounts {
464 sessions: sessions_inserted,
465 messages: messages_inserted,
466 parts: parts_inserted,
467 },
468 })
469 }
470
471 async fn export_clean_table(&self, table: Table, dest: &Path) -> Result<(usize, u64)> {
472 let dataset = self.handle.dataset(table).await?;
473 let source_version = dataset.version_id();
474 let schema = export_schema(table);
475 let mut scan = dataset.scan();
476 scan.blob_handling(lance::datatypes::BlobHandling::AllBinary);
481 let mut stream = scan
482 .try_into_stream()
483 .await
484 .with_context(|| format!("failed to scan {} for archive export", table.as_str()))?;
485 let dest_uri = dest
486 .to_str()
487 .with_context(|| format!("archive path is not UTF-8: {}", dest.display()))?;
488
489 let mut rows = 0usize;
490 let mut wrote = false;
491 while let Some(batch) = stream.next().await {
492 let batch = batch
493 .with_context(|| format!("failed to read {} archive batch", table.as_str()))?;
494 rows += batch.num_rows();
495 let reader = RecordBatchIterator::new([Ok(batch.clone())], batch.schema());
496 let mut params = write_params_for_create();
497 if wrote {
498 params.mode = WriteMode::Append;
499 }
500 Dataset::write(reader, dest_uri, Some(params))
501 .await
502 .with_context(|| format!("failed to write {} archive table", table.as_str()))?;
503 wrote = true;
504 }
505
506 if !wrote {
507 let batch = RecordBatch::new_empty(schema.clone());
508 let reader = RecordBatchIterator::new([Ok(batch)], schema);
509 Dataset::write(reader, dest_uri, Some(write_params_for_create()))
510 .await
511 .with_context(|| {
512 format!("failed to write empty {} archive table", table.as_str())
513 })?;
514 }
515 Ok((rows, source_version))
516 }
517
518 async fn import_clean_table(
523 &self,
524 table: Table,
525 dataset: Dataset,
526 upgrade: Option<ColumnBackfill>,
527 ) -> Result<(usize, usize)> {
528 let _ = self.handle.dataset(table).await?;
532 self.merge_scanner(table, dataset.scan(), "archive import", upgrade)
533 .await
534 }
535
536 async fn merge_scanner(
542 &self,
543 table: Table,
544 mut scanner: lance::dataset::scanner::Scanner,
545 context: &'static str,
546 upgrade: Option<ColumnBackfill>,
547 ) -> Result<(usize, usize)> {
548 scanner.blob_handling(lance::datatypes::BlobHandling::AllBinary);
549 let mut stream = scanner
550 .try_into_stream()
551 .await
552 .with_context(|| format!("failed to scan {} for {context}", table.as_str()))?;
553 let mut rows = 0usize;
554 let mut inserted = 0usize;
555 while let Some(batch) = stream.next().await {
556 let batch = batch
557 .with_context(|| format!("failed to read {} {context} batch", table.as_str()))?;
558 let batch = match &upgrade {
559 Some(spec) => upgraded_batch(&batch, spec)?,
560 None => batch,
561 };
562 let row_count = batch.num_rows();
563 rows += row_count;
564 let stats = self
565 .handle
566 .merge_insert_stats(table, batch, row_count)
567 .await
568 .with_context(|| format!("failed to merge {} during {context}", table.as_str()))?;
569 inserted += (stats.num_inserted_rows + stats.num_updated_rows) as usize;
570 }
571 Ok((rows, inserted))
572 }
573
574 pub async fn all_session_message_counts(&self) -> Result<HashMap<String, usize>> {
587 self.all_session_row_counts(Table::Messages).await
588 }
589
590 pub async fn all_session_part_counts(&self) -> Result<HashMap<String, usize>> {
591 self.all_session_row_counts(Table::Parts).await
592 }
593
594 async fn all_session_row_counts(&self, table: Table) -> Result<HashMap<String, usize>> {
599 let scanner = self
600 .handle
601 .scan(table, ScanOpts::project_only(&["session_id"]))
602 .await?;
603 let mut stream = scanner.try_into_stream().await?;
604 let mut out: HashMap<String, usize> = HashMap::new();
605 while let Some(batch) = stream.next().await {
606 let batch = batch?;
607 let session_ids = batch
608 .column_by_name("session_id")
609 .context("scan projection dropped the session_id column")?
610 .as_any()
611 .downcast_ref::<StringArray>()
612 .context("session_id column is not Utf8")?;
613 for row in 0..batch.num_rows() {
614 if session_ids.is_null(row) {
615 continue;
616 }
617 let session_id = session_ids.value(row);
618 if let Some(count) = out.get_mut(session_id) {
619 *count += 1;
620 } else {
621 out.insert(session_id.to_owned(), 1);
622 }
623 }
624 }
625 Ok(out)
626 }
627
628 pub async fn plan_incremental_from(&self, source: &Store) -> Result<DeltaPlan> {
638 let (
639 source_ids,
640 dest_ids,
641 source_msg_counts,
642 dest_msg_counts,
643 source_part_counts,
644 dest_part_counts,
645 ) = tokio::try_join!(
646 source.collect_ids(Table::Sessions),
647 self.collect_ids(Table::Sessions),
648 source.all_session_message_counts(),
649 self.all_session_message_counts(),
650 source.all_session_part_counts(),
651 self.all_session_part_counts(),
652 )?;
653 let source_sessions = source_ids.len();
654 let mut plan = DeltaPlan {
655 source_sessions,
656 ..DeltaPlan::default()
657 };
658 for id in &source_ids {
659 if !dest_ids.contains(id) {
662 plan.sessions.append.push(id.clone());
663 }
664 let source_msgs = source_msg_counts.get(id).copied().unwrap_or(0);
665 let dest_msgs = dest_msg_counts.get(id).copied().unwrap_or(0);
666 if dest_msgs == 0 {
667 if source_msgs > 0 {
668 plan.messages.append.push(id.clone());
669 }
670 } else if source_msgs > dest_msgs {
671 plan.messages.merge.push(id.clone());
672 }
673 let source_parts = source_part_counts.get(id).copied().unwrap_or(0);
674 let dest_parts = dest_part_counts.get(id).copied().unwrap_or(0);
675 if dest_parts == 0 {
676 if source_parts > 0 {
677 plan.parts.append.push(id.clone());
678 }
679 } else if source_parts > dest_parts {
680 plan.parts.merge.push(id.clone());
681 }
682 }
683 Ok(plan)
684 }
685
686 pub async fn copy_delta_from(
697 &self,
698 source: &Store,
699 plan: &DeltaPlan,
700 ) -> Result<LanceArchiveImport> {
701 let ((sessions, sessions_inserted), (messages, messages_inserted), (parts, parts_inserted)) =
705 tokio::try_join!(
706 self.copy_table(
707 source,
708 Table::Sessions,
709 "id",
710 &plan.sessions,
711 plan.source_sessions,
712 ),
713 self.copy_table(
714 source,
715 Table::Messages,
716 "session_id",
717 &plan.messages,
718 plan.source_sessions,
719 ),
720 self.copy_table(
721 source,
722 Table::Parts,
723 "session_id",
724 &plan.parts,
725 plan.source_sessions,
726 ),
727 )?;
728 Ok(LanceArchiveImport {
729 rows: LanceArchiveCounts {
730 sessions,
731 messages,
732 parts,
733 },
734 inserted: LanceArchiveCounts {
735 sessions: sessions_inserted,
736 messages: messages_inserted,
737 parts: parts_inserted,
738 },
739 })
740 }
741
742 async fn copy_table(
747 &self,
748 source: &Store,
749 table: Table,
750 key_column: &'static str,
751 table_plan: &TablePlan,
752 source_sessions: usize,
753 ) -> Result<(usize, usize)> {
754 let _ = self.handle.dataset(table).await?;
758
759 let appended = self
760 .append_sessions(
761 source,
762 table,
763 key_column,
764 &table_plan.append,
765 source_sessions,
766 )
767 .await?;
768
769 let mut grown = 0usize;
772 for chunk in table_plan.merge.chunks(COPY_SESSION_IN_CHUNK) {
773 let values: Vec<ScalarValue> = chunk
774 .iter()
775 .map(|id| ScalarValue::String(id.clone()))
776 .collect();
777 let predicate = Predicate::In("session_id", values.clone());
778 let stream = Self::source_scan(source, table, Some(&predicate)).await?;
779 grown += match table {
780 Table::Messages => {
781 let present = Arc::new(self.present_message_pks(&values).await?);
782 self.append_filtered(table, stream, Self::message_keep(present))
783 .await?
784 }
785 Table::Parts => {
786 let present = Arc::new(self.present_part_pks(&values).await?);
787 self.append_filtered(table, stream, Self::part_keep(present))
788 .await?
789 }
790 Table::Sessions => 0,
791 };
792 }
793 Ok((appended + grown, appended + grown))
794 }
795
796 async fn append_sessions(
803 &self,
804 source: &Store,
805 table: Table,
806 key_column: &'static str,
807 session_ids: &[String],
808 source_sessions: usize,
809 ) -> Result<usize> {
810 if session_ids.is_empty() {
811 return Ok(0);
812 }
813 if session_ids.len() == source_sessions {
814 return self.append_scanner(source, table, None).await;
815 }
816 let mut rows = 0usize;
817 for chunk in session_ids.chunks(COPY_SESSION_IN_CHUNK) {
818 let predicate = in_predicate(key_column, chunk);
819 rows += self.append_scanner(source, table, Some(&predicate)).await?;
820 }
821 Ok(rows)
822 }
823
824 async fn append_scanner(
830 &self,
831 source: &Store,
832 table: Table,
833 predicate: Option<&Predicate>,
834 ) -> Result<usize> {
835 let make_source = || Self::source_scan(source, table, predicate);
836 let stats = self.handle.append_stream(table, make_source).await?;
837 Ok(stats.rows as usize)
838 }
839
840 pub async fn append_absent_rows(
844 &self,
845 source: &Store,
846 table: Table,
847 filter_column: &'static str,
848 values: &[String],
849 ) -> Result<usize> {
850 if values.is_empty() {
851 return Ok(0);
852 }
853 let _ = self.handle.dataset(table).await?;
854 let mut rows = 0usize;
855 for chunk in values.chunks(COPY_SESSION_IN_CHUNK) {
856 let predicate = in_predicate(filter_column, chunk);
857 rows += self.append_scanner(source, table, Some(&predicate)).await?;
858 }
859 Ok(rows)
860 }
861
862 async fn present_message_pks(
865 &self,
866 session_id_values: &[ScalarValue],
867 ) -> Result<HashSet<(String, String)>> {
868 if session_id_values.is_empty() {
869 return Ok(HashSet::new());
870 }
871 let batch = self
872 .handle
873 .scan_batch(
874 Table::Messages,
875 Some(&Predicate::In("session_id", session_id_values.to_vec())),
876 pk_columns(Table::Messages),
877 )
878 .await?;
879 let mut set = HashSet::with_capacity(batch.num_rows());
880 for row in 0..batch.num_rows() {
881 let sid = string(&batch, "session_id", row)?.context("session_id is null")?;
882 let mid = string(&batch, "id", row)?.context("message id is null")?;
883 set.insert((sid, mid));
884 }
885 Ok(set)
886 }
887
888 async fn present_part_pks(
891 &self,
892 session_id_values: &[ScalarValue],
893 ) -> Result<HashSet<(String, String, String)>> {
894 if session_id_values.is_empty() {
895 return Ok(HashSet::new());
896 }
897 let batch = self
898 .handle
899 .scan_batch(
900 Table::Parts,
901 Some(&Predicate::In("session_id", session_id_values.to_vec())),
902 pk_columns(Table::Parts),
903 )
904 .await?;
905 let mut set = HashSet::with_capacity(batch.num_rows());
906 for row in 0..batch.num_rows() {
907 let sid = string(&batch, "session_id", row)?.context("session_id is null")?;
908 let mid = string(&batch, "message_id", row)?.context("message_id is null")?;
909 let pid = string(&batch, "id", row)?.context("part id is null")?;
910 set.insert((sid, mid, pid));
911 }
912 Ok(set)
913 }
914
915 async fn append_filtered<S, K>(&self, table: Table, mut batches: S, keep: K) -> Result<usize>
922 where
923 S: Stream<Item = std::result::Result<RecordBatch, DataFusionError>> + Unpin,
924 K: Fn(&RecordBatch, usize) -> Result<bool>,
925 {
926 let mut kept_batches: Vec<RecordBatch> = Vec::new();
927 let mut kept = 0usize;
928 while let Some(batch) = batches.next().await {
929 let batch = batch?;
930 let mut mask = Vec::with_capacity(batch.num_rows());
931 for row in 0..batch.num_rows() {
932 mask.push(keep(&batch, row)?);
933 }
934 let selected = filter_record_batch(&batch, &BooleanArray::from(mask))?;
935 if selected.num_rows() > 0 {
936 kept += selected.num_rows();
937 kept_batches.push(selected);
938 }
939 }
940 self.handle.append_batches(table, kept_batches).await?;
941 Ok(kept)
942 }
943
944 async fn source_scan(
947 source: &Store,
948 table: Table,
949 predicate: Option<&Predicate>,
950 ) -> Result<SendableRecordBatchStream> {
951 let mut scanner = source
952 .handle
953 .scan(
954 table,
955 ScanOpts {
956 predicate,
957 projection: None,
958 },
959 )
960 .await?;
961 scanner.blob_handling(lance::datatypes::BlobHandling::AllBinary);
962 Ok(scanner
963 .try_into_stream()
964 .await
965 .with_context(|| format!("failed to scan {} for copy", table.as_str()))?
966 .into())
967 }
968
969 fn message_keep(
972 present: Arc<HashSet<(String, String)>>,
973 ) -> impl Fn(&RecordBatch, usize) -> Result<bool> {
974 move |batch, row| {
975 let sid = string(batch, "session_id", row)?.context("session_id is null")?;
976 let mid = string(batch, "id", row)?.context("message id is null")?;
977 Ok(!present.contains(&(sid, mid)))
978 }
979 }
980
981 fn part_keep(
984 present: Arc<HashSet<(String, String, String)>>,
985 ) -> impl Fn(&RecordBatch, usize) -> Result<bool> {
986 move |batch, row| {
987 let sid = string(batch, "session_id", row)?.context("session_id is null")?;
988 let mid = string(batch, "message_id", row)?.context("message_id is null")?;
989 let pid = string(batch, "id", row)?.context("part id is null")?;
990 Ok(!present.contains(&(sid, mid, pid)))
991 }
992 }
993
994 pub async fn upsert_sessions(&self, sessions: &[Session]) -> Result<()> {
999 if sessions.is_empty() {
1000 return Ok(());
1001 }
1002 let batches = sessions_batches(sessions)?;
1003 merge_insert_chunks(&self.handle, Table::Sessions, batches).await?;
1004 Ok(())
1005 }
1006
1007 async fn embed_message_rows(
1013 &self,
1014 rows: &[MessageBatchRow<'_>],
1015 present: &HashSet<(String, String)>,
1016 ) -> Result<Vec<Option<Vec<f32>>>> {
1017 let mut out = vec![None; rows.len()];
1018 let Some(embedder) = &self.embedder else {
1019 return Ok(out);
1020 };
1021 let targets: Vec<usize> = rows
1022 .iter()
1023 .enumerate()
1024 .filter(|(_, row)| {
1025 row.search_text.is_some()
1026 && !present.contains(&(
1027 row.message.session_id().to_owned(),
1028 row.message.id().to_owned(),
1029 ))
1030 })
1031 .map(|(index, _)| index)
1032 .collect();
1033 if targets.is_empty() {
1034 return Ok(out);
1035 }
1036 let backend = embedder.get().await?;
1037 let texts: Vec<&str> = targets
1038 .iter()
1039 .map(|&index| rows[index].search_text.unwrap_or_default())
1040 .collect();
1041 let total = targets.len();
1042 let mut done = 0usize;
1043 if let Some(progress) = &self.ingest_embed_progress {
1044 (progress.0)(done, total);
1045 }
1046 let vectors = crate::embed::embed_passages(
1047 backend.as_ref(),
1048 &texts,
1049 crate::embed::DEFAULT_BATCH_SIZE,
1050 |batch| {
1051 done += batch;
1052 if let Some(progress) = &self.ingest_embed_progress {
1053 (progress.0)(done, total);
1054 }
1055 },
1056 )?;
1057 for (&index, vector) in targets.iter().zip(vectors) {
1058 out[index] = Some(vector);
1059 }
1060 Ok(out)
1061 }
1062
1063 async fn upsert_session_batch(
1087 &self,
1088 substreams: Vec<CompletedSubstream>,
1089 ) -> Result<(Vec<RowOutcome>, BatchCounts)> {
1090 if substreams.is_empty() {
1091 return Ok((Vec::new(), BatchCounts::default()));
1092 }
1093
1094 let mut outcomes: Vec<RowOutcome> = Vec::with_capacity(substreams.len());
1095 let mut counts = BatchCounts::default();
1096
1097 let mut merged: Vec<CompletedSubstream> = Vec::with_capacity(substreams.len());
1101 let mut by_session_id: std::collections::HashMap<String, usize> =
1102 std::collections::HashMap::with_capacity(substreams.len());
1103 for substream in substreams {
1104 if let Some(&existing_idx) = by_session_id.get(&substream.session.id) {
1105 let existing = &merged[existing_idx];
1106 if existing.session.source_agent != substream.session.source_agent
1107 || existing.session.project != substream.session.project
1108 {
1109 let reason = if existing.session.source_agent != substream.session.source_agent
1114 {
1115 IngestError::ImmutableField {
1116 field: "source_agent",
1117 session_id: substream.session.id.clone(),
1118 stored: existing.session.source_agent.clone(),
1119 attempted: substream.session.source_agent.clone(),
1120 }
1121 } else {
1122 IngestError::ImmutableField {
1123 field: "project",
1124 session_id: substream.session.id.clone(),
1125 stored: (*existing.session.project).clone(),
1126 attempted: (*substream.session.project).clone(),
1127 }
1128 };
1129 let field = match &reason {
1130 IngestError::ImmutableField { field, .. } => Some(*field),
1131 };
1132 let reason_key = match field {
1133 Some("project") => DROP_REASON_IMMUTABLE_PROJECT,
1134 Some("source_agent") => DROP_REASON_IMMUTABLE_SOURCE_AGENT,
1135 _ => DROP_REASON_UNCATEGORIZED,
1136 };
1137 outcomes.extend(error_outcomes_for_substream(
1138 substream.session_index,
1139 &substream.session,
1140 &substream.messages,
1141 reason.to_string(),
1142 field,
1143 reason_key,
1144 ));
1145 continue;
1146 }
1147 let existing = &mut merged[existing_idx];
1152 let mut seen: std::collections::HashSet<String> = existing
1153 .messages
1154 .iter()
1155 .map(|m| m.message.id().to_owned())
1156 .collect();
1157 for msg in substream.messages {
1158 if seen.insert(msg.message.id().to_owned()) {
1159 existing.messages.push(msg);
1160 }
1161 }
1162 continue;
1163 }
1164 by_session_id.insert(substream.session.id.clone(), merged.len());
1165 merged.push(substream);
1166 }
1167
1168 let session_id_values: Vec<ScalarValue> = merged
1173 .iter()
1174 .map(|substream| ScalarValue::String(substream.session.id.clone()))
1175 .collect();
1176 let existing_sessions: std::collections::HashMap<String, Session> =
1177 if session_id_values.is_empty() {
1178 std::collections::HashMap::new()
1179 } else {
1180 let batch = self
1181 .handle
1182 .scan_batch(
1183 Table::Sessions,
1184 Some(&Predicate::In("id", session_id_values.clone())),
1185 &[],
1186 )
1187 .await?;
1188 let mut map = std::collections::HashMap::with_capacity(batch.num_rows());
1189 for row in 0..batch.num_rows() {
1190 let session = session_from_batch(&batch, row)?;
1191 map.insert(session.id.clone(), session);
1192 }
1193 map
1194 };
1195 let existing_message_pks = Arc::new(self.present_message_pks(&session_id_values).await?);
1196 let existing_part_pks = Arc::new(self.present_part_pks(&session_id_values).await?);
1197
1198 let mut writeable: Vec<CompletedSubstream> = Vec::with_capacity(merged.len());
1199 for substream in merged {
1200 if let Some(existing) = existing_sessions.get(&substream.session.id)
1201 && let Err(failure) = ensure_immutable_match(existing, &substream.session)
1202 {
1203 let field = match &failure {
1204 IngestError::ImmutableField { field, .. } => Some(*field),
1205 };
1206 let reason_key = match field {
1207 Some("project") => DROP_REASON_IMMUTABLE_PROJECT,
1208 Some("source_agent") => DROP_REASON_IMMUTABLE_SOURCE_AGENT,
1209 _ => DROP_REASON_UNCATEGORIZED,
1210 };
1211 outcomes.extend(error_outcomes_for_substream(
1212 substream.session_index,
1213 &substream.session,
1214 &substream.messages,
1215 failure.to_string(),
1216 field,
1217 reason_key,
1218 ));
1219 continue;
1220 }
1221 writeable.push(substream);
1222 }
1223
1224 if writeable.is_empty() {
1225 outcomes.sort_by_key(|outcome| outcome.index);
1226 return Ok((outcomes, counts));
1227 }
1228
1229 let sessions_owned: Vec<Session> = writeable
1238 .iter()
1239 .map(|substream| &substream.session)
1240 .filter(|session| !existing_sessions.contains_key(&session.id))
1241 .cloned()
1242 .collect();
1243 let mut seen_messages: HashSet<(String, String)> = HashSet::new();
1246 let message_rows: Vec<MessageBatchRow<'_>> = writeable
1247 .iter()
1248 .flat_map(|substream| {
1249 substream.messages.iter().map(|buffered| MessageBatchRow {
1250 message: &buffered.message,
1251 source_agent: &substream.session.source_agent,
1252 project: &substream.session.project,
1253 search_text: buffered.search_text.as_deref(),
1254 })
1255 })
1256 .filter(|row| {
1257 seen_messages.insert((
1258 row.message.session_id().to_owned(),
1259 row.message.id().to_owned(),
1260 ))
1261 })
1262 .collect();
1263 let mut seen_parts: HashSet<(String, String, String)> = HashSet::new();
1264 let part_rows: Vec<Part> = writeable
1265 .iter()
1266 .flat_map(|substream| {
1267 substream.messages.iter().flat_map(|buffered| {
1268 buffered
1269 .parts
1270 .iter()
1271 .map(|buffered_part| buffered_part.part.clone())
1272 })
1273 })
1274 .filter(|part| {
1275 seen_parts.insert((
1276 part.session_id.clone(),
1277 part.message_id.clone(),
1278 part.id.clone(),
1279 ))
1280 })
1281 .collect();
1282
1283 let message_vectors = self
1286 .embed_message_rows(&message_rows, &existing_message_pks)
1287 .await?;
1288
1289 let message_stream = tokio_stream::iter(
1290 messages_batches(&message_rows, &message_vectors)?
1291 .into_iter()
1292 .map(Ok::<_, DataFusionError>),
1293 );
1294 let part_stream = tokio_stream::iter(
1295 parts_batches(&part_rows)?
1296 .into_iter()
1297 .map(Ok::<_, DataFusionError>),
1298 );
1299 let (_messages_appended, _parts_appended) = tokio::try_join!(
1300 self.append_filtered(
1301 Table::Messages,
1302 message_stream,
1303 Self::message_keep(existing_message_pks.clone()),
1304 ),
1305 self.append_filtered(
1306 Table::Parts,
1307 part_stream,
1308 Self::part_keep(existing_part_pks.clone()),
1309 ),
1310 )?;
1311 if !sessions_owned.is_empty() {
1312 let session_batches = sessions_batches(&sessions_owned)?;
1313 merge_insert_chunks(&self.handle, Table::Sessions, session_batches).await?;
1314 }
1315
1316 for substream in &writeable {
1317 outcomes.extend(success_outcomes_for_substream(
1318 substream.session_index,
1319 &substream.session,
1320 &substream.messages,
1321 &existing_sessions,
1322 &existing_message_pks,
1323 &existing_part_pks,
1324 &mut counts,
1325 ));
1326 }
1327
1328 outcomes.sort_by_key(|outcome| outcome.index);
1329 Ok((outcomes, counts))
1330 }
1331
1332 pub async fn upsert_messages(
1333 &self,
1334 session: &Session,
1335 messages: &[MessageWrite<'_>],
1336 ) -> Result<()> {
1337 if messages.is_empty() {
1338 return Ok(());
1339 }
1340
1341 let rows = messages
1342 .iter()
1343 .map(|write| MessageBatchRow {
1344 message: write.message,
1345 source_agent: &session.source_agent,
1346 project: &session.project,
1347 search_text: write.search_text,
1348 })
1349 .collect::<Vec<_>>();
1350 let batches = messages_batches(&rows, &vec![None; rows.len()])?;
1351 merge_insert_chunks(&self.handle, Table::Messages, batches).await?;
1352 Ok(())
1353 }
1354
1355 pub async fn upsert_parts(&self, parts: &[Part]) -> Result<()> {
1356 if parts.is_empty() {
1357 return Ok(());
1358 }
1359 let batches = parts_batches(parts)?;
1360 merge_insert_chunks(&self.handle, Table::Parts, batches).await?;
1361 Ok(())
1362 }
1363
1364 pub async fn get_session(&self, session_id: &str) -> Result<Option<SessionWithMessages>> {
1365 let Some(session) = self.find_session(session_id).await? else {
1366 return Ok(None);
1367 };
1368 let messages = self.messages_for_session(session_id).await?;
1369 Ok(Some(SessionWithMessages { session, messages }))
1370 }
1371
1372 pub async fn session_ids(&self) -> Result<Vec<String>> {
1374 let batch = self
1375 .handle
1376 .scan_batch(Table::Sessions, None, &["id"])
1377 .await?;
1378 let mut ids = Vec::with_capacity(batch.num_rows());
1379 for row in 0..batch.num_rows() {
1380 if let Some(id) = string(&batch, "id", row)? {
1381 ids.push(id);
1382 }
1383 }
1384 Ok(ids)
1385 }
1386
1387 pub async fn child_sessions(&self, parent_session_id: &str) -> Result<Vec<Session>> {
1388 let batch = self
1389 .handle
1390 .scan_batch(
1391 Table::Sessions,
1392 Some(&Predicate::Eq(
1393 "parent_session_id",
1394 parent_session_id.into(),
1395 )),
1396 &[
1397 "id",
1398 "parent_session_id",
1399 "parent_message_id",
1400 "source_agent",
1401 "created_at",
1402 "project",
1403 "options",
1404 ],
1405 )
1406 .await?;
1407 let mut sessions = Vec::with_capacity(batch.num_rows());
1408 for row in 0..batch.num_rows() {
1409 sessions.push(session_from_batch(&batch, row)?);
1410 }
1411 sessions.sort_by(|left, right| left.id.cmp(&right.id));
1412 Ok(sessions)
1413 }
1414
1415 pub async fn session_last_message_ids(&self) -> Result<HashMap<String, String>> {
1428 let (session_ids, latest) = tokio::try_join!(self.collect_ids(Table::Sessions), async {
1429 let scanner = self
1430 .handle
1431 .scan(
1432 Table::Messages,
1433 ScanOpts::project_only(&["session_id", "id", "timestamp"]),
1434 )
1435 .await?;
1436 let mut stream = scanner.try_into_stream().await?;
1437 let mut latest: HashMap<String, (DateTime<Utc>, String)> = HashMap::new();
1438 while let Some(batch) = stream.next().await {
1439 let batch = batch?;
1440 let session_ids = batch
1441 .column_by_name("session_id")
1442 .context("scan projection dropped the session_id column")?
1443 .as_any()
1444 .downcast_ref::<StringArray>()
1445 .context("session_id column is not Utf8")?;
1446 for row in 0..batch.num_rows() {
1447 if session_ids.is_null(row) {
1448 continue;
1449 }
1450 let session_id = session_ids.value(row);
1451 let Some(id) = string(&batch, "id", row)? else {
1452 continue;
1453 };
1454 let timestamp = datetime(&batch, "timestamp", row)?;
1455 match latest.get_mut(session_id) {
1456 Some((stored_ts, stored_id))
1457 if timestamp > *stored_ts
1458 || (timestamp == *stored_ts
1459 && id.as_str() > stored_id.as_str()) =>
1460 {
1461 *stored_ts = timestamp;
1462 *stored_id = id;
1463 }
1464 None => {
1465 latest.insert(session_id.to_owned(), (timestamp, id));
1466 }
1467 _ => {}
1468 }
1469 }
1470 }
1471 Ok::<_, anyhow::Error>(latest)
1472 })?;
1473 Ok(latest
1474 .into_iter()
1475 .filter(|(session_id, _)| session_ids.contains(session_id))
1476 .map(|(session_id, (_, message_id))| (session_id, message_id))
1477 .collect())
1478 }
1479
1480 pub async fn session_view(
1489 &self,
1490 session_id: &str,
1491 params: SessionViewParams<'_>,
1492 ) -> Result<GetLookup<SessionPage>> {
1493 let Some(session) = self.find_session(session_id).await? else {
1494 return Ok(GetLookup::NotFound);
1495 };
1496 let mut rows: Vec<ScanRow> = match self.session_scan_rows_resident(session_id).await? {
1499 Some(rows) => rows.into_iter().filter(|row| row.text.is_some()).collect(),
1500 None => self
1501 .scan_conversational_messages(session_id)
1502 .await?
1503 .into_iter()
1504 .map(|row| ScanRow {
1505 id: row.message_id,
1506 role: row.role,
1507 timestamp: row.timestamp,
1508 text: Some(row.text.into_inner()),
1509 content: None,
1510 })
1511 .collect(),
1512 };
1513 rows.sort_by(|a, b| a.timestamp.cmp(&b.timestamp).then_with(|| a.id.cmp(&b.id)));
1514
1515 let size = |row: &ScanRow| row.text.as_deref().map_or(0, str::len);
1516 let total = rows.len();
1517 let at_pos = params
1523 .at_message_id
1524 .and_then(|at| rows.iter().position(|row| row.id == at));
1525 let (win_start, win_end) = match (at_pos, params.after_message_id, params.before_message_id)
1526 {
1527 (Some(pos), _, _) => {
1528 let n = page_by(&rows[pos..], params.limit, params.budget_bytes, size);
1529 (pos, pos + n)
1530 }
1531 (None, Some(after), _) if params.at_message_id.is_none() => {
1532 let pos = match rows.iter().position(|row| row.id == after) {
1533 Some(idx) => idx + 1,
1534 None => return Ok(GetLookup::UnknownAnchor),
1535 };
1536 let n = page_by(&rows[pos..], params.limit, params.budget_bytes, size);
1537 (pos, pos + n)
1538 }
1539 (None, None, Some(before)) if params.at_message_id.is_none() => {
1540 let pos = match rows.iter().position(|row| row.id == before) {
1541 Some(idx) => idx,
1542 None => return Ok(GetLookup::UnknownAnchor),
1543 };
1544 let n = page_tail(&rows[..pos], params.limit, params.budget_bytes, size);
1545 (pos - n, pos)
1546 }
1547 _ => match params.session_from {
1548 SessionFrom::Start => (0, page_by(&rows, params.limit, params.budget_bytes, size)),
1549 SessionFrom::End => {
1550 let n = page_tail(&rows, params.limit, params.budget_bytes, size);
1551 (total - n, total)
1552 }
1553 },
1554 };
1555 let emitted = &rows[win_start..win_end];
1556 let before_remaining = win_start;
1557 let after_remaining = total - win_end;
1558 let ids: Vec<String> = emitted.iter().map(|row| row.id.clone()).collect();
1559
1560 let mut parts_by_message = self.summary_parts_for_messages(session_id, &ids).await?;
1561 let messages = emitted
1562 .iter()
1563 .map(|row| RetrievedMessage {
1564 id: row.id.clone(),
1565 role: row.role,
1566 timestamp: row.timestamp,
1567 text: row.text.clone(),
1568 content: row.content.clone(),
1569 parts: parts_by_message
1570 .remove(&(session_id.to_owned(), row.id.clone()))
1571 .unwrap_or_default(),
1572 })
1573 .collect();
1574
1575 Ok(GetLookup::Found(SessionPage {
1576 session,
1577 messages,
1578 before_remaining,
1579 after_remaining,
1580 }))
1581 }
1582
1583 pub async fn message_view(
1589 &self,
1590 message_id: &str,
1591 params: MessageViewParams,
1592 ) -> Result<GetLookup<MessagePage>> {
1593 let Some(session_id) = self.session_id_for_message(message_id).await? else {
1594 return Ok(GetLookup::NotFound);
1595 };
1596 let Some(session) = self.find_session(&session_id).await? else {
1597 return Ok(GetLookup::NotFound);
1598 };
1599 let mut rows = match self.session_scan_rows_resident(&session_id).await? {
1604 Some(rows)
1605 if rows
1606 .iter()
1607 .any(|row| row.id == message_id && row.role != Role::System) =>
1608 {
1609 rows
1610 }
1611 _ => self.scan_all_messages(&session_id).await?,
1612 };
1613 rows.retain(|row| row.text.is_some() || row.id == message_id);
1618 rows.sort_by(|a, b| a.timestamp.cmp(&b.timestamp).then_with(|| a.id.cmp(&b.id)));
1619 let Some(target_pos) = rows.iter().position(|row| row.id == message_id) else {
1620 return Ok(GetLookup::NotFound);
1621 };
1622
1623 let start = target_pos.saturating_sub(params.context_before);
1624 let end = (target_pos + params.context_after + 1).min(rows.len());
1625 let window = &rows[start..end];
1626 let window_ids: Vec<String> = window.iter().map(|row| row.id.clone()).collect();
1627 let mut parts_by_message = self.parts_for_messages(&session_id, &window_ids).await?;
1630
1631 let all_parts = parts_by_message
1632 .remove(&(session_id.clone(), message_id.to_owned()))
1633 .unwrap_or_default();
1634 let part_count = page_by(&all_parts, 1000, params.budget_bytes, |part| {
1637 serde_json::to_string(part).map_or(0, |json| json.len())
1638 });
1639 let target_parts = all_parts[..part_count].to_vec();
1640 let target_parts_remaining = all_parts.len() - part_count;
1641
1642 let target_row = &rows[target_pos];
1643 let target = RetrievedMessage {
1644 id: target_row.id.clone(),
1645 role: target_row.role,
1646 timestamp: target_row.timestamp,
1647 text: target_row.text.clone(),
1648 content: target_row.content.clone(),
1649 parts: Vec::new(),
1651 };
1652 let siblings = window
1653 .iter()
1654 .enumerate()
1655 .filter(|(idx, _)| start + idx != target_pos)
1656 .map(|(_, row)| RetrievedMessage {
1657 id: row.id.clone(),
1658 role: row.role,
1659 timestamp: row.timestamp,
1660 text: row.text.clone(),
1661 content: row.content.clone(),
1662 parts: parts_by_message
1663 .get(&(session_id.clone(), row.id.clone()))
1664 .cloned()
1665 .unwrap_or_default(),
1666 })
1667 .collect();
1668
1669 Ok(GetLookup::Found(MessagePage {
1670 session,
1671 target,
1672 target_parts,
1673 target_parts_remaining,
1674 siblings,
1675 }))
1676 }
1677
1678 async fn scan_all_messages(&self, session_id: &str) -> Result<Vec<ScanRow>> {
1679 let batch = self
1680 .handle
1681 .scan_batch(
1682 Table::Messages,
1683 Some(&Predicate::Eq("session_id", session_id.into())),
1684 &["id", "timestamp", "role", "search_text", "content"],
1685 )
1686 .await?;
1687 let mut rows = Vec::with_capacity(batch.num_rows());
1688 for row in 0..batch.num_rows() {
1689 let id = string(&batch, "id", row)?.context("message id is null")?;
1690 let role =
1691 role_from_str(&string(&batch, "role", row)?.context("message role is null")?)?;
1692 let timestamp = datetime(&batch, "timestamp", row)?;
1693 rows.push(ScanRow {
1694 id,
1695 role,
1696 timestamp,
1697 text: string(&batch, "search_text", row)?,
1698 content: string(&batch, "content", row)?,
1699 });
1700 }
1701 Ok(rows)
1702 }
1703
1704 async fn session_scan_rows_resident(&self, session_id: &str) -> Result<Option<Vec<ScanRow>>> {
1717 let Some(map) = self.rowmap.load_full() else {
1718 return Ok(None);
1719 };
1720 if map.version() != self.messages_version().await? {
1721 return Ok(None);
1722 }
1723 let Some(rowids) = map.session_row_ids(session_id) else {
1724 return Ok(None);
1725 };
1726 let (entries, misses) = map.hydrate(&rowids);
1727 if !misses.is_empty() {
1728 return Ok(None);
1729 }
1730 let mut rows = Vec::with_capacity(entries.len());
1731 for entry in entries {
1732 let Ok(role) = role_from_str(&entry.role) else {
1733 return Ok(None);
1734 };
1735 let Some(timestamp) = DateTime::from_timestamp_micros(entry.timestamp_micros) else {
1736 return Ok(None);
1737 };
1738 rows.push(ScanRow {
1739 id: entry.message_id,
1740 role,
1741 timestamp,
1742 text: (!entry.search_text.is_empty()).then_some(entry.search_text),
1743 content: None,
1744 });
1745 }
1746 Ok(Some(rows))
1747 }
1748
1749 pub async fn scan_conversational_messages(
1753 &self,
1754 session_id: &str,
1755 ) -> Result<Vec<ConversationalRow>> {
1756 let filter = Predicate::And(vec![
1757 Predicate::Eq("session_id", session_id.into()),
1758 Predicate::IsNotNull("search_text"),
1759 ]);
1760 let batch = self
1761 .handle
1762 .scan_batch(
1763 Table::Messages,
1764 Some(&filter),
1765 &["id", "timestamp", "role", "search_text"],
1766 )
1767 .await?;
1768
1769 let mut rows = Vec::with_capacity(batch.num_rows());
1770 for row in 0..batch.num_rows() {
1771 let message_id = string(&batch, "id", row)?.context("message id is null")?;
1772 let role =
1773 role_from_str(&string(&batch, "role", row)?.context("message role is null")?)?;
1774 let timestamp = datetime(&batch, "timestamp", row)?;
1775 let text_str = string(&batch, "search_text", row)?.context(
1776 "search_text null after IsNotNull pushdown - storage invariant violated",
1777 )?;
1778 rows.push(ConversationalRow {
1779 session_id: session_id.to_owned(),
1780 message_id,
1781 role,
1782 timestamp,
1783 text: SearchText(text_str),
1784 });
1785 }
1786 rows.sort_by(|a, b| {
1787 a.timestamp
1788 .cmp(&b.timestamp)
1789 .then_with(|| a.message_id.cmp(&b.message_id))
1790 });
1791 Ok(rows)
1792 }
1793
1794 pub async fn session_id_for_message(&self, message_id: &str) -> Result<Option<String>> {
1801 if let Some(map) = self.rowmap.load_full()
1802 && let Some(session_id) = map.lookup_session_for_message(message_id)
1803 {
1804 return Ok(Some(session_id.to_owned()));
1805 }
1806 let batch = self
1807 .handle
1808 .scan_batch(
1809 Table::Messages,
1810 Some(&Predicate::Eq("id", message_id.into())),
1811 &["session_id"],
1812 )
1813 .await?;
1814 if batch.num_rows() == 0 {
1815 return Ok(None);
1816 }
1817 string(&batch, "session_id", 0)
1818 }
1819
1820 pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
1821 self.handle.row_counts().await
1822 }
1823
1824 pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
1827 self.handle.collect_ids(table).await
1828 }
1829
1830 pub async fn composite_pk_index(&self, table: Table) -> Result<(HashSet<Vec<String>>, usize)> {
1837 let pk = pk_columns(table);
1838 let scanner = self.handle.scan(table, ScanOpts::project_only(pk)).await?;
1839 let mut stream = scanner.try_into_stream().await?;
1840 let mut keys: HashSet<Vec<String>> = HashSet::new();
1841 let mut rows = 0usize;
1842 while let Some(batch) = stream.next().await {
1843 let batch = batch?;
1844 for row in 0..batch.num_rows() {
1845 rows += 1;
1846 keys.insert(composite_key(&batch, pk, row)?);
1847 }
1848 }
1849 Ok((keys, rows))
1850 }
1851
1852 pub async fn composite_pk_diff_against(
1859 &self,
1860 table: Table,
1861 present: &HashSet<Vec<String>>,
1862 ) -> Result<(usize, usize)> {
1863 let pk = pk_columns(table);
1864 let scanner = self.handle.scan(table, ScanOpts::project_only(pk)).await?;
1865 let mut stream = scanner.try_into_stream().await?;
1866 let (mut rows, mut absent) = (0usize, 0usize);
1867 while let Some(batch) = stream.next().await {
1868 let batch = batch?;
1869 for row in 0..batch.num_rows() {
1870 rows += 1;
1871 if !present.contains(&composite_key(&batch, pk, row)?) {
1872 absent += 1;
1873 }
1874 }
1875 }
1876 Ok((rows, absent))
1877 }
1878
1879 pub async fn dataset(&self, table: Table) -> Result<Arc<Dataset>> {
1883 Ok(Arc::new(self.handle.dataset(table).await?))
1884 }
1885
1886 pub async fn prewarm(&self, cache_dir: &Path) -> Result<()> {
1896 let messages = self.dataset(Table::Messages).await?;
1897 if let Err(error) = messages.prewarm_index(MESSAGES_VECTOR_INDEX).await {
1898 tracing::debug!(%error, "vector index prewarm skipped");
1899 }
1900 if let Err(error) = self.ensure_rowmap(cache_dir).await {
1903 tracing::warn!(%error, "rowmap build skipped; arms fall back to data-take resolution");
1904 }
1905 if let Err(error) = self
1908 .fts_search("pond", 1, &Predicate::And(Vec::new()))
1909 .await
1910 {
1911 tracing::debug!(%error, "fts index prewarm skipped");
1912 }
1913 self.prune_index_cache(cache_dir).await;
1914 Ok(())
1915 }
1916
1917 pub async fn prune_index_cache(&self, cache_dir: &Path) {
1920 self.handle.prune_index_cache(cache_dir).await;
1921 }
1922
1923 fn store_key(&self) -> String {
1926 crate::substrate::store_key(self.handle.location())
1927 }
1928
1929 const MAX_ROWMAP_DELTAS: usize = 16;
1931
1932 const ROW_META_COLUMNS: [&str; 7] = [
1937 "session_id",
1938 "id",
1939 "role",
1940 "project",
1941 "source_agent",
1942 "timestamp",
1943 "search_text",
1944 ];
1945
1946 pub async fn ensure_rowmap(&self, cache_dir: &Path) -> Result<()> {
1953 let version = self.messages_version().await?;
1954 if let Some(current) = self.rowmap.load_full()
1955 && current.version() == version
1956 {
1957 return Ok(());
1958 }
1959 std::fs::create_dir_all(cache_dir)
1960 .with_context(|| format!("create cache dir {}", cache_dir.display()))?;
1961 let store_key = self.store_key();
1962
1963 if let Some(chain) = discover_chain(cache_dir, &store_key)
1966 && chain.version() == version
1967 && let Ok(set) = RowMetaSet::open(&chain)
1968 {
1969 self.rowmap.store(Some(Arc::new(set)));
1970 Self::sweep_stale_rowmaps(cache_dir, &store_key, chain.base_version);
1971 return Ok(());
1972 }
1973 if let Some(set) = self
1974 .extend_rowmap_coordinated(cache_dir, &store_key, version)
1975 .await?
1976 {
1977 self.rowmap.store(Some(Arc::new(set)));
1978 }
1979 Ok(())
1980 }
1981
1982 pub fn open_cached_rowmap(&self, cache_dir: &Path) -> Option<Arc<RowMetaSet>> {
1989 let chain = discover_chain(cache_dir, &self.store_key())?;
1990 RowMetaSet::open(&chain).ok().map(Arc::new)
1991 }
1992
1993 pub async fn load_rowmap_if_present(&self, cache_dir: &Path) -> Result<()> {
1999 let version = self.messages_version().await?;
2000 if let Some(current) = self.rowmap.load_full()
2001 && current.version() == version
2002 {
2003 return Ok(());
2004 }
2005 if let Some(chain) = discover_chain(cache_dir, &self.store_key())
2006 && chain.version() == version
2007 && let Ok(set) = RowMetaSet::open(&chain)
2008 {
2009 self.rowmap.store(Some(Arc::new(set)));
2010 }
2011 Ok(())
2012 }
2013
2014 async fn extend_rowmap_coordinated(
2020 &self,
2021 cache_dir: &Path,
2022 store_key: &str,
2023 version: u64,
2024 ) -> Result<Option<RowMetaSet>> {
2025 let lock_path = cache_dir.join(format!("rowmetamap-{store_key}.lock"));
2026 let lock = std::fs::File::create(&lock_path)
2027 .with_context(|| format!("create rowmap build lock {}", lock_path.display()))?;
2028 match lock.try_lock() {
2029 Ok(()) => {}
2030 Err(std::fs::TryLockError::WouldBlock) => return Ok(None),
2031 Err(std::fs::TryLockError::Error(error)) => {
2032 return Err(error).context("lock rowmap build");
2033 }
2034 }
2035
2036 if let Some(chain) = discover_chain(cache_dir, store_key)
2040 && chain.version() == version
2041 && let Ok(set) = RowMetaSet::open(&chain)
2042 {
2043 return Ok(Some(set));
2044 }
2045
2046 Self::sweep_orphan_temps(cache_dir, store_key);
2049
2050 let chain = discover_chain(cache_dir, store_key);
2056 let existing = match &chain {
2057 Some(paths) => match RowMetaSet::open(paths) {
2058 Ok(set) => Some((paths, set)),
2059 Err(error) => {
2060 tracing::warn!(%error, store = store_key, "rowmap unreadable; purging and rebuilding");
2061 Self::purge_rowmaps(cache_dir, store_key);
2062 None
2063 }
2064 },
2065 None => None,
2066 };
2067 let delta = match &existing {
2070 Some((_, set)) => {
2071 self.collect_row_metas_delta(
2072 set.version(),
2073 set.max_row_id().unwrap_or(0),
2074 set.len(),
2075 )
2076 .await?
2077 }
2078 None => None,
2079 };
2080
2081 let base_version = match (&existing, delta) {
2082 (Some((paths, _)), Some(entries)) if paths.deltas.len() < Self::MAX_ROWMAP_DELTAS => {
2084 let path = RowMetaMap::delta_path(cache_dir, store_key, version);
2085 RowMetaMap::build(&path, version, entries)?;
2086 paths.base_version
2087 }
2088 (Some((_, set)), Some(entries)) => {
2092 let mut merged = set.merged_entries();
2093 merged.extend(entries);
2094 let path = RowMetaMap::path_for(cache_dir, store_key, version);
2095 RowMetaMap::build(&path, version, merged)?;
2096 version
2097 }
2098 _ => {
2100 let entries = self.collect_row_metas().await?;
2101 let path = RowMetaMap::path_for(cache_dir, store_key, version);
2102 RowMetaMap::build(&path, version, entries)?;
2103 version
2104 }
2105 };
2106
2107 let chain =
2108 discover_chain(cache_dir, store_key).context("rowmap chain missing after build")?;
2109 let set = RowMetaSet::open(&chain)?;
2110 Self::sweep_stale_rowmaps(cache_dir, store_key, base_version);
2111 Ok(Some(set))
2112 }
2113
2114 fn sweep_stale_rowmaps(cache_dir: &Path, store_key: &str, keep: u64) {
2119 let prefix = format!("rowmetamap-{store_key}-");
2120 let Ok(entries) = std::fs::read_dir(cache_dir) else {
2121 return;
2122 };
2123 for entry in entries.flatten() {
2124 let name = entry.file_name();
2125 let Some(rest) = name
2126 .to_str()
2127 .and_then(|name| name.strip_prefix(&prefix))
2128 .and_then(|rest| rest.strip_suffix(".rmm"))
2129 else {
2130 continue;
2131 };
2132 let version = rest
2133 .strip_prefix('v')
2134 .or_else(|| rest.strip_prefix('d'))
2135 .and_then(|digits| digits.parse::<u64>().ok());
2136 if let Some(version) = version
2137 && version < keep
2138 {
2139 let _ = std::fs::remove_file(entry.path());
2140 }
2141 }
2142 }
2143
2144 fn purge_rowmaps(cache_dir: &Path, store_key: &str) {
2149 let prefix = format!("rowmetamap-{store_key}-");
2150 let Ok(entries) = std::fs::read_dir(cache_dir) else {
2151 return;
2152 };
2153 for entry in entries.flatten() {
2154 if let Some(name) = entry.file_name().to_str()
2155 && name.starts_with(&prefix)
2156 && name.ends_with(".rmm")
2157 {
2158 let _ = std::fs::remove_file(entry.path());
2159 }
2160 }
2161 }
2162
2163 fn sweep_orphan_temps(cache_dir: &Path, store_key: &str) {
2167 let prefix = format!("rowmetamap-{store_key}-");
2168 let Ok(entries) = std::fs::read_dir(cache_dir) else {
2169 return;
2170 };
2171 for entry in entries.flatten() {
2172 let name = entry.file_name();
2173 let Some(name) = name.to_str() else { continue };
2174 if name.starts_with(&prefix) && name.contains(".tmp-") {
2175 let _ = std::fs::remove_file(entry.path());
2176 }
2177 }
2178 }
2179
2180 #[cfg(test)]
2181 pub(crate) fn rowmap_delta_count(&self) -> Option<usize> {
2182 self.rowmap.load_full().map(|set| set.delta_count())
2183 }
2184
2185 pub fn rowmap_snapshot(&self) -> Option<Arc<RowMetaSet>> {
2189 self.rowmap.load_full()
2190 }
2191
2192 async fn resolve_rowid_hits(
2196 &self,
2197 map: &RowMetaSet,
2198 hits: Vec<(u64, f32)>,
2199 ) -> Result<Vec<SearchHit>> {
2200 let mut resolved = Vec::with_capacity(hits.len());
2201 let mut misses: Vec<(u64, f32)> = Vec::new();
2202 for (rowid, score) in hits {
2203 match map.lookup(rowid) {
2204 Some((session_id, message_id)) => resolved.push(SearchHit {
2205 rowid: Some(rowid),
2206 key: MessageKey {
2207 session_id: session_id.to_owned(),
2208 message_id: message_id.to_owned(),
2209 },
2210 score,
2211 }),
2212 None => misses.push((rowid, score)),
2213 }
2214 }
2215 if !misses.is_empty() {
2218 let rowids: Vec<u64> = misses.iter().map(|(rowid, _)| *rowid).collect();
2219 let keys = self.message_keys_by_rowids(&rowids).await?;
2220 for ((rowid, score), key) in misses.into_iter().zip(keys) {
2221 resolved.push(SearchHit {
2222 rowid: Some(rowid),
2223 key,
2224 score,
2225 });
2226 }
2227 }
2228 Ok(resolved)
2229 }
2230
2231 async fn message_keys_by_rowids(&self, rowids: &[u64]) -> Result<Vec<MessageKey>> {
2234 let dataset = self.handle.dataset(Table::Messages).await?;
2235 let projection = ProjectionRequest::from_columns(["session_id", "id"], dataset.schema());
2236 let batch = dataset.take_rows(rowids, projection).await?;
2237 let mut keys = Vec::with_capacity(batch.num_rows());
2238 for row in 0..batch.num_rows() {
2239 keys.push(MessageKey {
2240 session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2241 message_id: string(&batch, "id", row)?.context("fts hit id is null")?,
2242 });
2243 }
2244 Ok(keys)
2245 }
2246
2247 pub async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
2249 self.handle.export_write(name, bytes).await
2250 }
2251
2252 pub async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
2254 self.handle.export_read(name).await
2255 }
2256
2257 pub fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2259 self.handle.export_local_path(name)
2260 }
2261
2262 pub async fn adapter_names(&self, include_subagents: bool) -> Result<Vec<String>> {
2268 let scanner = self
2269 .handle
2270 .scan(Table::Sessions, ScanOpts::project_only(&["source_agent"]))
2271 .await?;
2272 let mut stream = scanner.try_into_stream().await?;
2273 let mut names: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
2274 while let Some(batch) = stream.next().await {
2275 let batch = batch?;
2276 for row in 0..batch.num_rows() {
2277 let agent = string(&batch, "source_agent", row)?.unwrap_or_default();
2278 if !include_subagents && agent.contains('/') {
2279 continue;
2280 }
2281 names.insert(agent);
2282 }
2283 }
2284 Ok(names.into_iter().collect())
2285 }
2286
2287 pub async fn ingest_host_activity(&self) -> Result<Vec<HostActivity>> {
2297 let scanner = self
2298 .handle
2299 .scan(
2300 Table::Messages,
2301 ScanOpts::project_only(&["session_id", "timestamp", "options"]),
2302 )
2303 .await?;
2304 let mut stream = scanner.try_into_stream().await?;
2305 let mut hosts: BTreeMap<Option<String>, (HashSet<String>, DateTime<Utc>)> = BTreeMap::new();
2306 while let Some(batch) = stream.next().await {
2307 let batch = batch?;
2308 for row in 0..batch.num_rows() {
2309 let session_id =
2310 string(&batch, "session_id", row)?.context("session_id is null")?;
2311 let timestamp = datetime(&batch, "timestamp", row)?;
2312 let hostname = json_column(&batch, "options", row)?
2313 .and_then(|bytes| json_parse::<serde_json::Value>(&bytes).ok())
2314 .as_ref()
2315 .and_then(|options| options.pointer("/pond/ingest/host/hostname"))
2316 .and_then(serde_json::Value::as_str)
2317 .map(str::to_owned);
2318 let entry = hosts
2319 .entry(hostname)
2320 .or_insert_with(|| (HashSet::new(), timestamp));
2321 entry.0.insert(session_id);
2322 entry.1 = entry.1.max(timestamp);
2323 }
2324 }
2325 Ok(hosts
2326 .into_iter()
2327 .map(|(hostname, (sessions, last_message_at))| HostActivity {
2328 hostname,
2329 sessions: sessions.len(),
2330 last_message_at,
2331 })
2332 .collect())
2333 }
2334
2335 pub async fn write_embeddings(&self, rows: &[EmbeddedMessage]) -> Result<()> {
2340 if rows.is_empty() {
2341 return Ok(());
2342 }
2343 let batch = embedding_update_batch(rows)?;
2344 self.handle
2345 .merge_update(Table::Messages, batch, rows.len())
2346 .await?;
2347 Ok(())
2348 }
2349
2350 pub fn pending_embedding_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2353 try_stream! {
2354 let filter = Predicate::And(vec![
2360 Predicate::IsNull("embedding_model"),
2361 Predicate::IsNotNull("search_text"),
2362 ]);
2363 let projection: &[&str] = &["session_id", "id", "search_text"];
2364 let scanner = self
2365 .handle
2366 .scan(
2367 Table::Messages,
2368 ScanOpts::with_predicate_and_projection(&filter, projection),
2369 )
2370 .await?;
2371 let mut batches = scanner
2372 .try_into_stream()
2373 .await
2374 .context("failed to open messages stream")?;
2375 while let Some(batch) = batches.next().await {
2376 let batch = batch?;
2377 for row in 0..batch.num_rows() {
2378 yield PendingMessage {
2379 session_id: string(&batch, "session_id", row)?
2380 .context("session_id is null")?,
2381 id: string(&batch, "id", row)?.context("message id is null")?,
2382 search_text: string(&batch, "search_text", row)?
2383 .context("search_text is null")?,
2384 };
2385 }
2386 }
2387 }
2388 }
2389
2390 pub fn pending_or_stale_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2395 try_stream! {
2396 let filter = Predicate::And(vec![
2400 Predicate::IsNotNull("search_text"),
2401 Predicate::Or(vec![
2402 Predicate::IsNull("embedding_model"),
2403 Predicate::Ne("embedding_model", embed::model_id().into()),
2404 ]),
2405 ]);
2406 let projection: &[&str] = &["session_id", "id", "search_text"];
2407 let scanner = self
2408 .handle
2409 .scan(
2410 Table::Messages,
2411 ScanOpts::with_predicate_and_projection(&filter, projection),
2412 )
2413 .await?;
2414 let mut batches = scanner
2415 .try_into_stream()
2416 .await
2417 .context("failed to open pending-or-stale messages stream")?;
2418 while let Some(batch) = batches.next().await {
2419 let batch = batch?;
2420 for row in 0..batch.num_rows() {
2421 yield PendingMessage {
2422 session_id: string(&batch, "session_id", row)?
2423 .context("session_id is null")?,
2424 id: string(&batch, "id", row)?.context("message id is null")?,
2425 search_text: string(&batch, "search_text", row)?
2426 .context("search_text is null")?,
2427 };
2428 }
2429 }
2430 }
2431 }
2432
2433 pub async fn fts_search(
2438 &self,
2439 query: &str,
2440 limit: usize,
2441 filter: &Predicate,
2442 ) -> Result<Vec<SearchHit>> {
2443 let mut hits = if let Some(map) = self.rowmap.load_full() {
2444 let rowid_hits = self.fts_search_rowids(query, limit, filter).await?;
2445 self.resolve_rowid_hits(&map, rowid_hits).await?
2446 } else {
2447 self.fts_search_keys(query, limit, filter).await?
2448 };
2449 hits.sort_by(|left, right| {
2455 right
2456 .score
2457 .partial_cmp(&left.score)
2458 .unwrap_or(std::cmp::Ordering::Equal)
2459 .then_with(|| left.key.session_id.cmp(&right.key.session_id))
2460 .then_with(|| left.key.message_id.cmp(&right.key.message_id))
2461 });
2462 Ok(hits)
2463 }
2464
2465 async fn fts_scanner(
2469 &self,
2470 query: &str,
2471 limit: usize,
2472 filter: &Predicate,
2473 ) -> Result<lance::dataset::scanner::Scanner> {
2474 let mut scanner = self.handle.scanner(Table::Messages, Some(filter)).await?;
2475 scanner.full_text_search(
2476 FullTextSearchQuery::new(query.to_owned()).with_column("search_text".to_owned())?,
2477 )?;
2478 if self
2479 .handle
2480 .messages_fast_search_ready(MESSAGES_FTS_INDEX)
2481 .await?
2482 {
2483 scanner.fast_search();
2484 }
2485 scanner.disable_scoring_autoprojection();
2491 scanner.limit(Some(i64::try_from(limit).unwrap_or(i64::MAX)), None)?;
2492 Ok(scanner)
2493 }
2494
2495 async fn fts_search_keys(
2498 &self,
2499 query: &str,
2500 limit: usize,
2501 filter: &Predicate,
2502 ) -> Result<Vec<SearchHit>> {
2503 let mut scanner = self.fts_scanner(query, limit, filter).await?;
2504 scanner.project(&["session_id", "id", "_score"])?;
2505 let batch = scanner.try_into_batch().await?;
2506 let mut hits = Vec::with_capacity(batch.num_rows());
2507 for row in 0..batch.num_rows() {
2508 let key = MessageKey {
2509 session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2510 message_id: string(&batch, "id", row)?.context("fts hit id is null")?,
2511 };
2512 hits.push(SearchHit {
2513 rowid: None,
2514 key,
2515 score: float32(&batch, "_score", row)?,
2516 });
2517 }
2518 Ok(hits)
2519 }
2520
2521 pub async fn messages_version(&self) -> Result<u64> {
2524 Ok(self
2525 .handle
2526 .dataset(Table::Messages)
2527 .await?
2528 .version()
2529 .version)
2530 }
2531
2532 pub async fn collect_row_metas(&self) -> Result<Vec<RowMetaEntry>> {
2536 let mut scanner = self.handle.scanner(Table::Messages, None).await?;
2537 scanner.with_row_id();
2538 scanner.project(&Self::ROW_META_COLUMNS)?;
2539 let mut stream = scanner.try_into_stream().await?;
2540 let mut out = Vec::new();
2541 while let Some(batch) = stream.next().await {
2542 let batch = batch?;
2543 let rowids = uint64(&batch, "_rowid")?;
2544 for row in 0..batch.num_rows() {
2545 out.push(row_meta_entry(&batch, rowids.value(row), row)?);
2546 }
2547 }
2548 Ok(out)
2549 }
2550
2551 async fn collect_row_metas_delta(
2567 &self,
2568 base_version: u64,
2569 base_max_row_id: u64,
2570 base_row_count: usize,
2571 ) -> Result<Option<Vec<RowMetaEntry>>> {
2572 let dataset = self.handle.dataset(Table::Messages).await?;
2573 let Ok(old) = dataset.checkout_version(base_version).await else {
2574 return Ok(None);
2575 };
2576 if dataset.count_rows(None).await? < base_row_count {
2577 return Ok(None);
2578 }
2579 let old_ids: HashSet<u64> = old.get_fragments().iter().map(|f| f.id() as u64).collect();
2584 let added: Vec<_> = dataset
2585 .get_fragments()
2586 .iter()
2587 .filter(|fragment| !old_ids.contains(&(fragment.id() as u64)))
2588 .map(|fragment| fragment.metadata().clone())
2589 .collect();
2590 if added.is_empty() {
2591 return Ok(Some(Vec::new()));
2592 }
2593 let mut scanner = dataset.scan();
2594 scanner.with_fragments(added);
2595 scanner.with_row_id();
2596 scanner.project(&Self::ROW_META_COLUMNS)?;
2597 let mut stream = scanner.try_into_stream().await?;
2598 let mut out = Vec::new();
2599 while let Some(batch) = stream.next().await {
2600 let batch = batch?;
2601 let rowids = uint64(&batch, "_rowid")?;
2602 for row in 0..batch.num_rows() {
2603 let row_id = rowids.value(row);
2604 if row_id > base_max_row_id {
2605 out.push(row_meta_entry(&batch, row_id, row)?);
2606 }
2607 }
2608 }
2609 Ok(Some(out))
2610 }
2611
2612 async fn fts_search_rowids(
2615 &self,
2616 query: &str,
2617 limit: usize,
2618 filter: &Predicate,
2619 ) -> Result<Vec<(u64, f32)>> {
2620 let mut scanner = self.fts_scanner(query, limit, filter).await?;
2621 scanner.with_row_id();
2622 scanner.project(&["_score"])?;
2623 let batch = scanner.try_into_batch().await?;
2624 let rowids = uint64(&batch, "_rowid")?;
2625 let mut hits = Vec::with_capacity(batch.num_rows());
2626 for row in 0..batch.num_rows() {
2627 hits.push((rowids.value(row), float32(&batch, "_score", row)?));
2628 }
2629 Ok(hits)
2630 }
2631
2632 pub async fn searchable_in_scope(&self, filter: &Predicate) -> Result<usize> {
2639 if matches!(filter, Predicate::And(clauses) if clauses.is_empty())
2645 && let Some(count) = self.fts_num_docs().await?
2646 {
2647 return Ok(count);
2648 }
2649 let scope = Predicate::And(vec![Predicate::IsNotNull("search_text"), filter.clone()]);
2650 let dataset = self.handle.dataset(Table::Messages).await?;
2651 let count = dataset.count_rows(Some(scope.to_lance())).await?;
2652 Ok(count)
2653 }
2654
2655 async fn fts_num_docs(&self) -> Result<Option<usize>> {
2659 if !self.handle.messages_has_index(MESSAGES_FTS_INDEX).await? {
2660 return Ok(None);
2661 }
2662 let dataset = self.handle.dataset(Table::Messages).await?;
2663 let json = dataset.index_statistics(MESSAGES_FTS_INDEX).await?;
2664 let parsed: Value =
2665 serde_json::from_str(&json).context("failed to parse FTS index_statistics")?;
2666 let total: u64 = parsed["indices"]
2667 .as_array()
2668 .map(|segments| {
2669 segments
2670 .iter()
2671 .filter_map(|segment| segment["num_docs"].as_u64())
2672 .sum()
2673 })
2674 .unwrap_or(0);
2675 Ok(Some(usize::try_from(total).unwrap_or(usize::MAX)))
2676 }
2677
2678 pub async fn has_embeddings(&self) -> Result<bool> {
2687 if self
2688 .handle
2689 .messages_has_index(MESSAGES_VECTOR_INDEX)
2690 .await?
2691 {
2692 return Ok(true);
2693 }
2694 let scope = Predicate::IsNotNull("vector");
2695 let mut scanner = self
2696 .handle
2697 .scan(
2698 Table::Messages,
2699 ScanOpts::with_predicate_and_projection(&scope, &["id"]),
2700 )
2701 .await?;
2702 scanner.limit(Some(1), None)?;
2703 let batch = scanner.try_into_batch().await?;
2704 Ok(batch.num_rows() > 0)
2705 }
2706
2707 pub async fn sample_embedded_model(&self) -> Result<Option<String>> {
2713 let scope = Predicate::IsNotNull("embedding_model");
2714 let mut scanner = self
2715 .handle
2716 .scan(
2717 Table::Messages,
2718 ScanOpts::with_predicate_and_projection(&scope, &["embedding_model"]),
2719 )
2720 .await?;
2721 scanner.limit(Some(1), None)?;
2722 let batch = scanner.try_into_batch().await?;
2723 if batch.num_rows() == 0 {
2724 return Ok(None);
2725 }
2726 string(&batch, "embedding_model", 0)
2727 }
2728
2729 pub async fn embedding_model_swapped(&self) -> Result<bool> {
2734 Ok(self
2735 .sample_embedded_model()
2736 .await?
2737 .is_some_and(|model| model != crate::embed::model_id()))
2738 }
2739
2740 pub async fn vector_search(
2748 &self,
2749 query: &[f32],
2750 limit: usize,
2751 filter: &Predicate,
2752 search: Option<&config::SearchConfig>,
2753 ) -> Result<Vec<SearchHit>> {
2754 let mut hits = if let Some(map) = self.rowmap.load_full() {
2755 let rowid_hits = self
2756 .vector_search_rowids(query, limit, filter, search)
2757 .await?;
2758 self.resolve_rowid_hits(&map, rowid_hits).await?
2759 } else {
2760 self.vector_search_keys(query, limit, filter, search)
2761 .await?
2762 };
2763 hits.sort_by(|left, right| {
2769 left.score
2770 .partial_cmp(&right.score)
2771 .unwrap_or(std::cmp::Ordering::Equal)
2772 .then_with(|| left.key.session_id.cmp(&right.key.session_id))
2773 .then_with(|| left.key.message_id.cmp(&right.key.message_id))
2774 });
2775 Ok(hits)
2776 }
2777
2778 async fn vector_scanner(
2780 &self,
2781 query: &[f32],
2782 limit: usize,
2783 filter: &Predicate,
2784 search: Option<&config::SearchConfig>,
2785 ) -> Result<lance::dataset::scanner::Scanner> {
2786 let scope = embedded_scope(filter);
2787 let mut scanner = self.handle.scanner(Table::Messages, Some(&scope)).await?;
2788 let key = Float32Array::from(query.to_vec());
2789 scanner.nearest("vector", &key, limit)?;
2790 apply_vector_search_knobs(&mut scanner, search);
2791 if self
2792 .handle
2793 .messages_fast_search_ready(MESSAGES_VECTOR_INDEX)
2794 .await?
2795 {
2796 scanner.fast_search();
2797 }
2798 scanner.disable_scoring_autoprojection();
2799 Ok(scanner)
2800 }
2801
2802 async fn vector_search_rowids(
2805 &self,
2806 query: &[f32],
2807 limit: usize,
2808 filter: &Predicate,
2809 search: Option<&config::SearchConfig>,
2810 ) -> Result<Vec<(u64, f32)>> {
2811 let mut scanner = self.vector_scanner(query, limit, filter, search).await?;
2812 scanner.with_row_id();
2813 scanner.project(&["_distance"])?;
2814 let batch = scanner.try_into_batch().await?;
2815 let rowids = uint64(&batch, "_rowid")?;
2816 let mut hits = Vec::with_capacity(batch.num_rows());
2817 for row in 0..batch.num_rows() {
2818 hits.push((rowids.value(row), float32(&batch, "_distance", row)?));
2819 }
2820 Ok(hits)
2821 }
2822
2823 async fn vector_search_keys(
2826 &self,
2827 query: &[f32],
2828 limit: usize,
2829 filter: &Predicate,
2830 search: Option<&config::SearchConfig>,
2831 ) -> Result<Vec<SearchHit>> {
2832 let mut scanner = self.vector_scanner(query, limit, filter, search).await?;
2833 scanner.project(&["session_id", "id", "_distance"])?;
2834 let batch = scanner.try_into_batch().await?;
2835 let mut hits = Vec::with_capacity(batch.num_rows());
2836 for row in 0..batch.num_rows() {
2837 let key = MessageKey {
2838 session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2839 message_id: string(&batch, "id", row)?.context("message id is null")?,
2840 };
2841 hits.push(SearchHit {
2842 rowid: None,
2843 key,
2844 score: float32(&batch, "_distance", row)?,
2845 });
2846 }
2847 Ok(hits)
2848 }
2849
2850 pub async fn explain_vector_plan(
2853 &self,
2854 query: &[f32],
2855 limit: usize,
2856 filter: &Predicate,
2857 search: Option<&config::SearchConfig>,
2858 ) -> Result<String> {
2859 let scanner = self.vector_scanner(query, limit, filter, search).await?;
2863 scanner
2864 .explain_plan(true)
2865 .await
2866 .context("explain_plan failed")
2867 }
2868
2869 pub async fn explain_fts_plan(
2870 &self,
2871 query: &str,
2872 limit: usize,
2873 filter: &Predicate,
2874 ) -> Result<String> {
2875 let mut scanner = self.fts_scanner(query, limit, filter).await?;
2877 scanner.project(&["session_id", "id"])?;
2878 scanner
2879 .explain_plan(true)
2880 .await
2881 .context("explain_plan failed")
2882 }
2883
2884 pub async fn message_metas_by_rowids(&self, rowids: &[u64]) -> Result<Vec<MessageMeta>> {
2892 if rowids.is_empty() {
2893 return Ok(Vec::new());
2894 }
2895 let mut metas = Vec::with_capacity(rowids.len());
2896 let misses: Vec<u64> = if let Some(map) = self.rowmap.load_full() {
2897 let (hits, misses) = map.hydrate(rowids);
2898 metas.extend(hits.into_iter().map(|entry| MessageMeta {
2899 message_id: entry.message_id,
2900 session_id: entry.session_id,
2901 role: entry.role,
2902 project: entry.project,
2903 source_agent: entry.source_agent,
2904 timestamp:
2905 DateTime::from_timestamp_micros(entry.timestamp_micros).unwrap_or_default(),
2906 search_text: entry.search_text,
2907 }));
2908 misses
2909 } else {
2910 rowids.to_vec()
2911 };
2912 if !misses.is_empty() {
2913 metas.extend(self.message_metas_by_rowids_take(&misses).await?);
2914 }
2915 Ok(metas)
2916 }
2917
2918 async fn message_metas_by_rowids_take(&self, rowids: &[u64]) -> Result<Vec<MessageMeta>> {
2923 let dataset = self.handle.dataset(Table::Messages).await?;
2924 let projection = ProjectionRequest::from_columns(
2925 [
2926 "id",
2927 "session_id",
2928 "role",
2929 "project",
2930 "source_agent",
2931 "timestamp",
2932 "search_text",
2933 ],
2934 dataset.schema(),
2935 );
2936 let batch = dataset.take_rows(rowids, projection).await?;
2937 let mut metas = Vec::with_capacity(batch.num_rows());
2938 for row in 0..batch.num_rows() {
2939 metas.push(message_meta_from_batch(&batch, row)?);
2940 }
2941 Ok(metas)
2942 }
2943
2944 pub async fn message_metas_by_keys(&self, keys: &[MessageKey]) -> Result<Vec<MessageMeta>> {
2946 if keys.is_empty() {
2947 return Ok(Vec::new());
2948 }
2949 let wanted = keys.iter().cloned().collect::<HashSet<_>>();
2950 let session_ids = keys
2951 .iter()
2952 .map(|key| key.session_id.clone())
2953 .collect::<Vec<_>>();
2954 let message_ids = keys
2955 .iter()
2956 .map(|key| key.message_id.clone())
2957 .collect::<Vec<_>>();
2958 let predicate = Predicate::And(vec![
2959 in_predicate("session_id", &session_ids),
2960 in_predicate("id", &message_ids),
2961 ]);
2962 let batch = self
2963 .handle
2964 .scan_batch(
2965 Table::Messages,
2966 Some(&predicate),
2967 &[
2968 "id",
2969 "session_id",
2970 "role",
2971 "project",
2972 "source_agent",
2973 "timestamp",
2974 "search_text",
2975 ],
2976 )
2977 .await?;
2978 let mut metas = Vec::with_capacity(batch.num_rows());
2979 for row in 0..batch.num_rows() {
2980 let meta = message_meta_from_batch(&batch, row)?;
2983 if wanted.contains(&MessageKey {
2984 session_id: meta.session_id.clone(),
2985 message_id: meta.message_id.clone(),
2986 }) {
2987 metas.push(meta);
2988 }
2989 }
2990 Ok(metas)
2991 }
2992
2993 pub async fn session_message_counts(
3002 &self,
3003 session_ids: &[String],
3004 ) -> Result<BTreeMap<String, usize>> {
3005 if session_ids.is_empty() {
3006 return Ok(BTreeMap::new());
3007 }
3008 if let Some(map) = self.rowmap.load_full()
3015 && map.version() == self.messages_version().await?
3016 {
3017 return Ok(session_ids
3018 .iter()
3019 .map(|id| (id.clone(), map.lookup_count(id).unwrap_or(0)))
3020 .collect());
3021 }
3022 let predicate = in_predicate("session_id", session_ids);
3023 let scanner = self
3024 .handle
3025 .scan(
3026 Table::Messages,
3027 ScanOpts::with_predicate_and_projection(&predicate, &["session_id"]),
3028 )
3029 .await?;
3030 let mut stream = scanner
3031 .try_into_stream()
3032 .await
3033 .context("failed to open session_message_counts stream")?;
3034 let mut counts: BTreeMap<String, usize> =
3035 session_ids.iter().map(|id| (id.clone(), 0)).collect();
3036 while let Some(batch) = stream.next().await {
3037 let batch = batch.context("failed to read session_message_counts batch")?;
3038 let column = batch
3039 .column_by_name("session_id")
3040 .context("session_message_counts: session_id column missing")?
3041 .as_any()
3042 .downcast_ref::<StringArray>()
3043 .context("session_message_counts: session_id column is not Utf8")?;
3044 for value in column.iter().flatten() {
3045 if let Some(entry) = counts.get_mut(value) {
3046 *entry += 1;
3047 }
3048 }
3049 }
3050 Ok(counts)
3051 }
3052
3053 pub async fn unindexed_message_backlog(&self) -> Result<usize> {
3056 self.handle
3057 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
3058 .await
3059 }
3060
3061 pub async fn unindexed_vector_backlog(&self) -> Result<usize> {
3069 self.handle
3070 .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
3071 .await
3072 }
3073
3074 pub async fn embedding_progress(&self) -> Result<EmbeddingProgress> {
3078 let dataset = self.handle.dataset(Table::Messages).await?;
3079 let embedded = dataset
3085 .count_rows(Some(Predicate::IsNotNull("embedding_model").to_lance()))
3086 .await?;
3087 let backlog = self.embed_backlog_count().await?;
3094 Ok(EmbeddingProgress {
3095 embedded,
3096 total: embedded + backlog,
3097 backlog,
3098 model: embed::model_id(),
3099 })
3100 }
3101
3102 pub async fn embed_backlog_count(&self) -> Result<usize> {
3108 let dataset = self.handle.dataset(Table::Messages).await?;
3109 let filter = Predicate::And(vec![
3110 Predicate::IsNull("embedding_model"),
3111 Predicate::IsNotNull("search_text"),
3112 ]);
3113 Ok(dataset.count_rows(Some(filter.to_lance())).await?)
3114 }
3115
3116 pub async fn stale_embedding_count(&self) -> Result<usize> {
3120 let dataset = self.handle.dataset(Table::Messages).await?;
3121 dataset
3127 .count_rows(Some(
3128 Predicate::And(vec![
3129 Predicate::IsNotNull("embedding_model"),
3130 Predicate::Ne("embedding_model", embed::model_id().into()),
3131 ])
3132 .to_lance(),
3133 ))
3134 .await
3135 .map_err(Into::into)
3136 }
3137
3138 pub async fn optimize_indices(
3144 &self,
3145 progress: Option<OptimizeProgressFn>,
3146 maintenance: &MaintenancePolicy,
3147 ) -> Result<OptimizeOutcome> {
3148 let intents = pond_index_intents();
3149 let mut tables = Vec::with_capacity(3);
3150 for (table, intents) in intents.all() {
3151 let outcome = self
3152 .handle
3153 .optimize_table(table, intents, progress.as_ref(), maintenance)
3154 .await;
3155 tables.push(outcome);
3156 }
3157 Ok(OptimizeOutcome { tables })
3158 }
3159
3160 pub async fn build_indices_only(
3166 &self,
3167 progress: Option<OptimizeProgressFn>,
3168 ) -> Result<OptimizeOutcome> {
3169 let policy = pond_index_intents();
3170 let mut tables = Vec::with_capacity(3);
3171 for (table, intents) in policy.all() {
3172 let indices = self
3173 .handle
3174 .optimize_table_indices_only(table, intents, progress.as_ref())
3175 .await;
3176 tables.push(TableOptimizeOutcome {
3177 table,
3178 indices,
3179 compaction: PhaseOutcome::NotAttempted,
3180 });
3181 }
3182 Ok(OptimizeOutcome { tables })
3183 }
3184
3185 #[cfg(test)]
3186 async fn optimize_indices_with_vector_threshold(
3187 &self,
3188 vector_threshold: usize,
3189 ) -> Result<OptimizeOutcome> {
3190 let intents = pond_index_intents_with_vector_threshold(vector_threshold);
3191 let policy = MaintenancePolicy::always_compact();
3192 let mut tables = Vec::with_capacity(3);
3193 for (table, intents) in intents.all() {
3194 let outcome = self
3195 .handle
3196 .optimize_table(table, intents, None, &policy)
3197 .await;
3198 tables.push(outcome);
3199 }
3200 Ok(OptimizeOutcome { tables })
3201 }
3202
3203 #[cfg(test)]
3204 async fn optimize_indices_with_scalar_fold_threshold(
3205 &self,
3206 scalar_fold_row_threshold: usize,
3207 ) -> Result<OptimizeOutcome> {
3208 let intents = pond_index_intents();
3209 let policy = MaintenancePolicy::always_compact()
3210 .with_scalar_fold_row_threshold(scalar_fold_row_threshold);
3211 let mut tables = Vec::with_capacity(3);
3212 for (table, intents) in intents.all() {
3213 let outcome = self
3214 .handle
3215 .optimize_table(table, intents, None, &policy)
3216 .await;
3217 tables.push(outcome);
3218 }
3219 Ok(OptimizeOutcome { tables })
3220 }
3221
3222 pub async fn cleanup_old_versions(&self, older_than: chrono::Duration) -> Result<()> {
3228 for (table, _) in pond_index_intents().all() {
3229 self.handle
3230 .cleanup_table_versions(table, older_than)
3231 .await?;
3232 }
3233 Ok(())
3234 }
3235
3236 pub async fn rebuild_indices(
3237 &self,
3238 intent_name: Option<&str>,
3239 progress: Option<OptimizeProgressFn>,
3240 ) -> Result<()> {
3241 let policy = pond_index_intents();
3242 let mut matched = false;
3243 for (table, intents) in policy.all() {
3244 for intent in intents {
3245 if intent_name.is_none_or(|name| name == intent.name) {
3246 matched = true;
3247 self.handle
3248 .rebuild_index(table, intent, progress.as_ref())
3249 .await?;
3250 }
3251 }
3252 }
3253 if let Some(name) = intent_name
3254 && !matched
3255 {
3256 anyhow::bail!("unknown index intent {name:?}");
3257 }
3258 Ok(())
3259 }
3260
3261 pub async fn drop_index_by_name(&self, name: &str) -> Result<()> {
3268 let Some(owner) = self.handle.find_index_owner(name).await? else {
3269 anyhow::bail!("no index named {name:?} found on any table");
3270 };
3271 self.handle.drop_index(owner, name).await
3272 }
3273
3274 pub async fn index_status(&self) -> Result<Vec<IndexStatus>> {
3275 self.index_status_with(false).await
3276 }
3277
3278 pub async fn index_status_indexable(&self) -> Result<Vec<IndexStatus>> {
3283 self.index_status_with(true).await
3284 }
3285
3286 async fn index_status_with(&self, indexable_only: bool) -> Result<Vec<IndexStatus>> {
3287 let policy = pond_index_intents();
3288 let mut statuses = Vec::new();
3289 for (table, intents) in policy.all() {
3290 statuses.extend(
3291 self.handle
3292 .index_status(table, intents, indexable_only)
3293 .await?,
3294 );
3295 }
3296 Ok(statuses)
3297 }
3298
3299 pub async fn drop_vector_index(&self) -> Result<()> {
3303 match self
3304 .handle
3305 .drop_index(Table::Messages, MESSAGES_VECTOR_INDEX)
3306 .await
3307 {
3308 Ok(()) => Ok(()),
3309 Err(error) => {
3310 let msg = error.to_string();
3311 if msg.contains("not found") || msg.contains("does not exist") {
3312 Ok(())
3313 } else {
3314 Err(error)
3315 }
3316 }
3317 }
3318 }
3319
3320 pub async fn table_sizes(&self) -> Result<TableSizes> {
3323 self.handle.table_sizes().await
3324 }
3325
3326 pub async fn initialized(&self) -> Result<bool> {
3327 self.handle.initialized().await
3328 }
3329
3330 pub(crate) async fn find_session(&self, session_id: &str) -> Result<Option<Session>> {
3333 let batch = self
3334 .handle
3335 .scan_batch(
3336 Table::Sessions,
3337 Some(&Predicate::Eq("id", session_id.into())),
3338 &[],
3339 )
3340 .await?;
3341 if batch.num_rows() == 0 {
3342 Ok(None)
3343 } else {
3344 Ok(Some(session_from_batch(&batch, 0)?))
3345 }
3346 }
3347
3348 async fn messages_for_session(&self, session_id: &str) -> Result<Vec<MessageWithParts>> {
3349 let batch = self
3350 .handle
3351 .scan_batch(
3352 Table::Messages,
3353 Some(&Predicate::Eq("session_id", session_id.into())),
3354 &[
3355 "session_id",
3356 "id",
3357 "timestamp",
3358 "role",
3359 "content",
3360 "options",
3361 ],
3362 )
3363 .await?;
3364 let mut messages = Vec::with_capacity(batch.num_rows());
3365 for row in 0..batch.num_rows() {
3366 messages.push(message_from_batch(&batch, row)?);
3367 }
3368 messages.sort_by(|left, right| {
3369 left.timestamp()
3370 .cmp(&right.timestamp())
3371 .then_with(|| left.id().cmp(right.id()))
3372 });
3373
3374 let message_ids = messages
3375 .iter()
3376 .map(|message| message.id().to_owned())
3377 .collect::<Vec<_>>();
3378 let mut parts_by_message = self.parts_for_messages(session_id, &message_ids).await?;
3379
3380 Ok(messages
3381 .into_iter()
3382 .map(|message| {
3383 let key = (message.session_id().to_owned(), message.id().to_owned());
3384 let parts = parts_by_message.remove(&key).unwrap_or_default();
3385 MessageWithParts { message, parts }
3386 })
3387 .collect())
3388 }
3389
3390 pub async fn parts_for_messages(
3394 &self,
3395 session_id: &str,
3396 message_ids: &[String],
3397 ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3398 self.scan_parts(session_id, message_ids, None).await
3399 }
3400
3401 pub async fn summary_parts_for_messages(
3406 &self,
3407 session_id: &str,
3408 message_ids: &[String],
3409 ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3410 self.scan_parts(session_id, message_ids, Some(SUMMARY_PART_TYPES))
3411 .await
3412 }
3413
3414 async fn scan_parts(
3415 &self,
3416 session_id: &str,
3417 message_ids: &[String],
3418 part_types: Option<&[&str]>,
3419 ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3420 if message_ids.is_empty() {
3421 return Ok(BTreeMap::new());
3422 }
3423 let mut clauses = vec![
3424 Predicate::Eq("session_id", session_id.into()),
3425 in_predicate("message_id", message_ids),
3426 ];
3427 if let Some(types) = part_types {
3428 clauses.push(Predicate::In(
3429 "type",
3430 types.iter().map(|&t| t.into()).collect(),
3431 ));
3432 }
3433 let predicate = Predicate::And(clauses);
3434 let summarizing = part_types.is_some();
3439 let mut scanner = self
3440 .handle
3441 .scan(
3442 Table::Parts,
3443 ScanOpts::with_predicate_and_projection(
3444 &predicate,
3445 &[
3446 "session_id",
3447 "message_id",
3448 "id",
3449 "ordinal",
3450 "type",
3451 "provenance",
3452 "variant_data",
3453 "options",
3454 ],
3455 ),
3456 )
3457 .await?;
3458 if !summarizing {
3459 scanner.with_row_address();
3460 }
3461 let batch = scanner.try_into_batch().await.context("scan failed")?;
3462 let mut file_payloads = BTreeMap::<usize, FileData>::new();
3463 if !summarizing {
3464 let dataset = std::sync::Arc::new(self.handle.dataset(Table::Parts).await?);
3465 let row_addresses = uint64(&batch, "_rowaddr")?;
3466 let mut file_rows = Vec::<(usize, u64, Vec<u8>)>::new();
3467 for row in 0..batch.num_rows() {
3468 if string(&batch, "type", row)?.as_deref() == Some("file") {
3469 let variant_data = json_column(&batch, "variant_data", row)?
3470 .context("variant_data is null")?;
3471 file_rows.push((row, row_addresses.value(row), variant_data));
3472 }
3473 }
3474 if !file_rows.is_empty() {
3475 let addresses = file_rows
3476 .iter()
3477 .map(|(_, address, _)| *address)
3478 .collect::<Vec<_>>();
3479 let blobs = dataset.take_blobs_by_addresses(&addresses, "data").await?;
3480 for ((row, _, variant_data), blob) in file_rows.into_iter().zip(blobs) {
3481 let payload = file_data_from_blob(&variant_data, &blob.read().await?)?;
3485 file_payloads.insert(row, payload);
3486 }
3487 }
3488 }
3489 let mut parts_by_message = BTreeMap::<(String, String), Vec<Part>>::new();
3490 for row in 0..batch.num_rows() {
3491 let file_data = if summarizing {
3495 (string(&batch, "type", row)?.as_deref() == Some("file"))
3496 .then(|| FileData::Bytes(Vec::new()))
3497 } else {
3498 file_payloads.remove(&row)
3499 };
3500 let part = part_from_batch(&batch, row, file_data)?;
3501 parts_by_message
3502 .entry((part.session_id.clone(), part.message_id.clone()))
3503 .or_default()
3504 .push(part);
3505 }
3506 for parts in parts_by_message.values_mut() {
3507 parts.sort_by_key(|part| part.ordinal);
3508 }
3509 Ok(parts_by_message)
3510 }
3511}
3512
3513#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
3514#[serde(tag = "kind", content = "data", rename_all = "snake_case")]
3515pub enum IngestEvent {
3516 Session(Session),
3517 Message(Message),
3518 Part(Part),
3519}
3520
3521#[derive(Debug, Clone, PartialEq, Eq, Default)]
3529pub struct IngestSummary {
3530 pub inserted: usize,
3534 pub matched: usize,
3536 pub sessions_inserted: usize,
3538 pub messages_inserted_total: usize,
3541 pub messages_inserted_searchable: usize,
3545 pub parts_inserted: usize,
3547 pub sessions_matched: usize,
3549 pub messages_matched_total: usize,
3551 pub messages_matched_searchable: usize,
3553 pub parts_matched: usize,
3555 pub dropped_events: usize,
3565 pub dropped_sessions: usize,
3570 pub skipped_files: usize,
3573 pub skipped_empty: usize,
3578 pub skipped_fresh: usize,
3582 pub skipped_superseded: usize,
3586 pub storage_errors: usize,
3590 pub truncated_values: usize,
3593 pub drop_reasons: BTreeMap<&'static str, usize>,
3599}
3600
3601pub const DROP_REASON_DUPLICATE_MESSAGE_ID: &str = "duplicate_message_id";
3607pub const DROP_REASON_DUPLICATE_PART_KEY: &str = "duplicate_part_key";
3608pub const DROP_REASON_MESSAGE_BEFORE_SESSION: &str = "message_before_session";
3609pub const DROP_REASON_MESSAGE_SESSION_MISMATCH: &str = "message_session_mismatch";
3610pub const DROP_REASON_PART_BEFORE_MESSAGE: &str = "part_before_message";
3611pub const DROP_REASON_PART_MESSAGE_MISMATCH: &str = "part_message_mismatch";
3612pub const DROP_REASON_EMPTY_SOURCE_AGENT: &str = "empty_source_agent";
3613pub const DROP_REASON_PARENT_MESSAGE_WITHOUT_SESSION: &str = "parent_message_without_session";
3614pub const DROP_REASON_IMMUTABLE_PROJECT: &str = "immutable_project";
3615pub const DROP_REASON_IMMUTABLE_SOURCE_AGENT: &str = "immutable_source_agent";
3616pub const DROP_REASON_UNCATEGORIZED: &str = "uncategorized";
3617
3618#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
3626pub struct BatchCounts {
3627 pub sessions_inserted: usize,
3628 pub sessions_matched: usize,
3629 pub messages_inserted_total: usize,
3630 pub messages_inserted_searchable: usize,
3631 pub messages_matched_total: usize,
3632 pub messages_matched_searchable: usize,
3633 pub parts_inserted: usize,
3634 pub parts_matched: usize,
3635}
3636
3637impl IngestSummary {
3638 pub fn accepted(&self) -> usize {
3639 self.inserted + self.matched
3640 }
3641
3642 pub fn add_batch(&mut self, counts: &BatchCounts) {
3646 self.sessions_inserted += counts.sessions_inserted;
3647 self.sessions_matched += counts.sessions_matched;
3648 self.messages_inserted_total += counts.messages_inserted_total;
3649 self.messages_inserted_searchable += counts.messages_inserted_searchable;
3650 self.messages_matched_total += counts.messages_matched_total;
3651 self.messages_matched_searchable += counts.messages_matched_searchable;
3652 self.parts_inserted += counts.parts_inserted;
3653 self.parts_matched += counts.parts_matched;
3654 self.inserted +=
3655 counts.sessions_inserted + counts.messages_inserted_total + counts.parts_inserted;
3656 self.matched +=
3657 counts.sessions_matched + counts.messages_matched_total + counts.parts_matched;
3658 }
3659
3660 pub fn merge(&mut self, other: &Self) {
3664 self.inserted += other.inserted;
3665 self.matched += other.matched;
3666 self.sessions_inserted += other.sessions_inserted;
3667 self.messages_inserted_total += other.messages_inserted_total;
3668 self.messages_inserted_searchable += other.messages_inserted_searchable;
3669 self.parts_inserted += other.parts_inserted;
3670 self.sessions_matched += other.sessions_matched;
3671 self.messages_matched_total += other.messages_matched_total;
3672 self.messages_matched_searchable += other.messages_matched_searchable;
3673 self.parts_matched += other.parts_matched;
3674 self.dropped_events += other.dropped_events;
3675 self.dropped_sessions += other.dropped_sessions;
3676 self.skipped_files += other.skipped_files;
3677 self.skipped_empty += other.skipped_empty;
3678 self.skipped_fresh += other.skipped_fresh;
3679 self.skipped_superseded += other.skipped_superseded;
3680 self.storage_errors += other.storage_errors;
3681 self.truncated_values += other.truncated_values;
3682 for (key, value) in &other.drop_reasons {
3683 *self.drop_reasons.entry(key).or_insert(0) += value;
3684 }
3685 }
3686
3687 pub fn add_outcomes_errors_only(&mut self, outcomes: &[RowOutcome]) {
3692 for outcome in outcomes {
3693 if !matches!(outcome.status, OutcomeStatus::Error) {
3694 continue;
3695 }
3696 if outcome.kind == "session" {
3697 self.dropped_sessions += 1;
3698 } else {
3699 self.dropped_events += 1;
3700 }
3701 let reason = outcome
3702 .error
3703 .as_ref()
3704 .and_then(|error| error.reason_key)
3705 .unwrap_or(DROP_REASON_UNCATEGORIZED);
3706 *self.drop_reasons.entry(reason).or_insert(0) += 1;
3707 }
3708 }
3709
3710 pub fn add_outcomes(&mut self, outcomes: &[RowOutcome]) {
3711 for outcome in outcomes {
3712 match outcome.status {
3713 OutcomeStatus::Inserted => {
3714 self.inserted += 1;
3715 match outcome.kind {
3716 "session" => self.sessions_inserted += 1,
3717 "message" => {
3718 self.messages_inserted_total += 1;
3719 if outcome.searchable {
3720 self.messages_inserted_searchable += 1;
3721 }
3722 }
3723 "part" => self.parts_inserted += 1,
3724 _ => {}
3725 }
3726 }
3727 OutcomeStatus::Matched => {
3728 self.matched += 1;
3729 match outcome.kind {
3730 "session" => self.sessions_matched += 1,
3731 "message" => {
3732 self.messages_matched_total += 1;
3733 if outcome.searchable {
3734 self.messages_matched_searchable += 1;
3735 }
3736 }
3737 "part" => self.parts_matched += 1,
3738 _ => {}
3739 }
3740 }
3741 OutcomeStatus::Error => {
3742 if outcome.kind == "session" {
3748 self.dropped_sessions += 1;
3749 } else {
3750 self.dropped_events += 1;
3751 }
3752 let reason = outcome
3753 .error
3754 .as_ref()
3755 .and_then(|e| e.reason_key)
3756 .unwrap_or(DROP_REASON_UNCATEGORIZED);
3757 *self.drop_reasons.entry(reason).or_insert(0) += 1;
3758 }
3759 }
3760 }
3761 }
3762}
3763
3764#[derive(Debug, Clone, PartialEq)]
3769pub struct RowOutcome {
3770 pub index: usize,
3771 pub kind: &'static str,
3772 pub pk: Value,
3773 pub status: OutcomeStatus,
3774 pub error: Option<RowError>,
3775 pub searchable: bool,
3780}
3781
3782#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3783pub enum OutcomeStatus {
3784 Inserted,
3785 Matched,
3786 Error,
3787}
3788
3789#[derive(Debug, Clone, PartialEq, Eq)]
3792pub struct RowError {
3793 pub message: String,
3794 pub field: Option<&'static str>,
3795 pub reason: Option<&'static str>,
3796 pub reason_key: Option<&'static str>,
3801}
3802
3803#[derive(Debug)]
3807struct BufferedSession {
3808 index: usize,
3809 session: Session,
3810}
3811
3812#[derive(Debug)]
3813struct BufferedMessage {
3814 index: usize,
3815 message: Message,
3816 parts: Vec<BufferedPart>,
3817 search_text: Option<String>,
3818}
3819
3820#[derive(Debug)]
3821struct BufferedPart {
3822 index: usize,
3823 part: Part,
3824}
3825
3826#[derive(Debug, Default)]
3843pub struct IngestValidator {
3844 session: Option<BufferedSession>,
3845 current_message: Option<BufferedMessage>,
3846 current_parts: Vec<BufferedPart>,
3847 messages: Vec<BufferedMessage>,
3848 seen_message_ids: HashSet<String>,
3852 seen_part_keys: HashSet<(String, String)>,
3855 completed: Vec<CompletedSubstream>,
3859}
3860
3861#[derive(Debug)]
3863struct CompletedSubstream {
3864 session_index: usize,
3865 session: Session,
3866 messages: Vec<BufferedMessage>,
3867}
3868
3869fn ingest_host_stamp() -> Option<&'static Value> {
3874 static STAMP: std::sync::OnceLock<Option<Value>> = std::sync::OnceLock::new();
3875 STAMP
3876 .get_or_init(|| {
3877 let mut host = serde_json::Map::new();
3878 if let Ok(username) = whoami::username() {
3879 host.insert("username".to_owned(), username.into());
3880 }
3881 if let Ok(hostname) = whoami::hostname() {
3882 host.insert("hostname".to_owned(), hostname.into());
3883 }
3884 if let Ok(devicename) = whoami::devicename() {
3885 host.insert("device_name".to_owned(), devicename.into());
3886 }
3887 (!host.is_empty()).then(|| serde_json::json!({ "ingest": { "host": host } }))
3888 })
3889 .as_ref()
3890}
3891
3892impl IngestValidator {
3893 pub async fn push(
3899 &mut self,
3900 store: &Store,
3901 index: usize,
3902 event: IngestEvent,
3903 ) -> Result<Vec<RowOutcome>> {
3904 match event {
3905 IngestEvent::Session(session) => self.push_session(store, index, session).await,
3906 IngestEvent::Message(message) => Ok(self.push_message(index, message)),
3907 IngestEvent::Part(part) => Ok(self.push_part(index, part)),
3908 }
3909 }
3910
3911 pub async fn finish(&mut self, store: &Store) -> Result<(Vec<RowOutcome>, BatchCounts)> {
3916 self.close_current_substream();
3917 self.flush(store).await
3918 }
3919
3920 pub async fn flush(&mut self, store: &Store) -> Result<(Vec<RowOutcome>, BatchCounts)> {
3927 if self.completed.is_empty() {
3928 return Ok((Vec::new(), BatchCounts::default()));
3929 }
3930 let completed = std::mem::take(&mut self.completed);
3931 store.upsert_session_batch(completed).await
3932 }
3933
3934 pub fn pending_substreams(&self) -> usize {
3937 self.completed.len()
3938 }
3939
3940 async fn push_session(
3941 &mut self,
3942 _store: &Store,
3943 index: usize,
3944 mut session: Session,
3945 ) -> Result<Vec<RowOutcome>> {
3946 self.close_current_substream();
3950
3951 let trimmed = session.source_agent.trim();
3956 if trimmed.is_empty() {
3957 return Ok(vec![RowOutcome {
3958 index,
3959 kind: "session",
3960 pk: Value::String(session.id.clone()),
3961 status: OutcomeStatus::Error,
3962 error: Some(RowError {
3963 message: format!("session {} has empty source_agent after trim", session.id),
3964 field: Some("source_agent"),
3965 reason: None,
3966 reason_key: Some(DROP_REASON_EMPTY_SOURCE_AGENT),
3967 }),
3968 searchable: false,
3969 }]);
3970 }
3971 if trimmed.len() != session.source_agent.len() {
3972 session.source_agent = trimmed.to_owned();
3973 }
3974
3975 if session.parent_message_id.is_some() && session.parent_session_id.is_none() {
3976 return Ok(vec![RowOutcome {
3977 index,
3978 kind: "session",
3979 pk: Value::String(session.id.clone()),
3980 status: OutcomeStatus::Error,
3981 error: Some(RowError {
3982 message: format!(
3983 "session {} has parent_message_id without parent_session_id",
3984 session.id,
3985 ),
3986 field: Some("parent_message_id"),
3987 reason: None,
3988 reason_key: Some(DROP_REASON_PARENT_MESSAGE_WITHOUT_SESSION),
3989 }),
3990 searchable: false,
3991 }]);
3992 }
3993
3994 self.seen_message_ids.clear();
3995 self.seen_part_keys.clear();
3996 self.session = Some(BufferedSession { index, session });
3997 Ok(Vec::new())
3998 }
3999
4000 fn close_current_substream(&mut self) {
4001 self.flush_current_message();
4002 let Some(BufferedSession {
4003 index: session_index,
4004 session,
4005 }) = self.session.take()
4006 else {
4007 return;
4008 };
4009 let messages = std::mem::take(&mut self.messages);
4010 self.seen_message_ids.clear();
4011 self.seen_part_keys.clear();
4012 self.completed.push(CompletedSubstream {
4013 session_index,
4014 session,
4015 messages,
4016 });
4017 }
4018
4019 fn push_message(&mut self, index: usize, mut message: Message) -> Vec<RowOutcome> {
4020 let pk = Value::Array(vec![
4021 Value::String(message.session_id().to_owned()),
4022 Value::String(message.id().to_owned()),
4023 ]);
4024 let Some(session) = &self.session else {
4025 return vec![error_outcome(
4026 index,
4027 "message",
4028 pk,
4029 "first event in a session stream must be Session",
4030 None,
4031 DROP_REASON_MESSAGE_BEFORE_SESSION,
4032 )];
4033 };
4034 if message.session_id() != session.session.id {
4035 let msg = format!(
4036 "message {} references session {}, expected {}",
4037 message.id(),
4038 message.session_id(),
4039 session.session.id
4040 );
4041 return vec![error_outcome(
4042 index,
4043 "message",
4044 pk,
4045 &msg,
4046 Some("session_id"),
4047 DROP_REASON_MESSAGE_SESSION_MISMATCH,
4048 )];
4049 }
4050 if !self.seen_message_ids.insert(message.id().to_owned()) {
4051 let msg = format!("duplicate message id {} in session substream", message.id());
4055 return vec![error_outcome(
4056 index,
4057 "message",
4058 pk,
4059 &msg,
4060 None,
4061 DROP_REASON_DUPLICATE_MESSAGE_ID,
4062 )];
4063 }
4064 match ingest_host_stamp() {
4069 Some(stamp) => {
4070 message
4071 .options_mut()
4072 .insert("pond".to_owned(), stamp.clone());
4073 }
4074 None => {
4075 message.options_mut().remove("pond");
4076 }
4077 }
4078 self.flush_current_message();
4079 self.current_message = Some(BufferedMessage {
4080 index,
4081 message,
4082 parts: Vec::new(),
4083 search_text: None,
4084 });
4085 Vec::new()
4086 }
4087
4088 fn push_part(&mut self, index: usize, part: Part) -> Vec<RowOutcome> {
4089 let pk = Value::Array(vec![
4090 Value::String(part.session_id.clone()),
4091 Value::String(part.message_id.clone()),
4092 Value::String(part.id.clone()),
4093 ]);
4094 let Some(current) = &self.current_message else {
4095 return vec![error_outcome(
4096 index,
4097 "part",
4098 pk,
4099 "part event appeared before a message",
4100 None,
4101 DROP_REASON_PART_BEFORE_MESSAGE,
4102 )];
4103 };
4104 if part.session_id != current.message.session_id() {
4105 let msg = format!(
4106 "part {} references session {}, expected {}",
4107 part.id,
4108 part.session_id,
4109 current.message.session_id()
4110 );
4111 return vec![error_outcome(
4112 index,
4113 "part",
4114 pk,
4115 &msg,
4116 Some("session_id"),
4117 DROP_REASON_PART_MESSAGE_MISMATCH,
4118 )];
4119 }
4120 if part.message_id != current.message.id() {
4121 let msg = format!(
4122 "part {} references message {}, expected {}",
4123 part.id,
4124 part.message_id,
4125 current.message.id()
4126 );
4127 return vec![error_outcome(
4128 index,
4129 "part",
4130 pk,
4131 &msg,
4132 Some("message_id"),
4133 DROP_REASON_PART_MESSAGE_MISMATCH,
4134 )];
4135 }
4136 let part_key = (part.message_id.clone(), part.id.clone());
4137 if !self.seen_part_keys.insert(part_key) {
4138 let msg = format!(
4139 "duplicate part id {} for message {} in session substream",
4140 part.id, part.message_id
4141 );
4142 return vec![error_outcome(
4143 index,
4144 "part",
4145 pk,
4146 &msg,
4147 None,
4148 DROP_REASON_DUPLICATE_PART_KEY,
4149 )];
4150 }
4151 self.current_parts.push(BufferedPart { index, part });
4152 Vec::new()
4153 }
4154
4155 fn flush_current_message(&mut self) {
4156 let Some(mut buffered) = self.current_message.take() else {
4157 return;
4158 };
4159 let parts = std::mem::take(&mut self.current_parts);
4160 let mut canonical_parts = Vec::with_capacity(parts.len());
4161 for part in &parts {
4162 canonical_parts.push(part.part.clone());
4163 }
4164 buffered.search_text = search_text(&buffered.message, &canonical_parts);
4165 buffered.parts = parts;
4166 self.messages.push(buffered);
4167 }
4168}
4169
4170fn error_outcome(
4171 index: usize,
4172 kind: &'static str,
4173 pk: Value,
4174 message: &str,
4175 field: Option<&'static str>,
4176 reason_key: &'static str,
4177) -> RowOutcome {
4178 RowOutcome {
4179 index,
4180 kind,
4181 pk,
4182 status: OutcomeStatus::Error,
4183 error: Some(RowError {
4184 message: message.to_owned(),
4185 field,
4186 reason: None,
4187 reason_key: Some(reason_key),
4188 }),
4189 searchable: false,
4190 }
4191}
4192
4193fn error_outcomes_for_substream(
4198 session_index: usize,
4199 session: &Session,
4200 _messages: &[BufferedMessage],
4201 message: impl Into<String>,
4202 field: Option<&'static str>,
4203 reason_key: &'static str,
4204) -> Vec<RowOutcome> {
4205 let reason = field.map(|_| "immutable");
4206 vec![RowOutcome {
4207 index: session_index,
4208 kind: "session",
4209 pk: Value::String(session.id.clone()),
4210 status: OutcomeStatus::Error,
4211 error: Some(RowError {
4212 message: message.into(),
4213 field,
4214 reason,
4215 reason_key: Some(reason_key),
4216 }),
4217 searchable: false,
4218 }]
4219}
4220
4221fn success_outcomes_for_substream(
4227 session_index: usize,
4228 session: &Session,
4229 messages: &[BufferedMessage],
4230 existing_sessions: &std::collections::HashMap<String, Session>,
4231 existing_message_pks: &HashSet<(String, String)>,
4232 existing_part_pks: &HashSet<(String, String, String)>,
4233 counts: &mut BatchCounts,
4234) -> Vec<RowOutcome> {
4235 let session_was_present = existing_sessions.contains_key(&session.id);
4236 let session_status = if session_was_present {
4237 counts.sessions_matched += 1;
4238 UpsertStatus::Matched
4239 } else {
4240 counts.sessions_inserted += 1;
4241 UpsertStatus::Inserted
4242 };
4243
4244 let mut outcomes = Vec::with_capacity(1 + messages.len());
4245 outcomes.push(success_outcome(
4246 session_index,
4247 "session",
4248 Value::String(session.id.clone()),
4249 session_status,
4250 false,
4251 ));
4252 for buffered in messages {
4253 let key = (
4254 buffered.message.session_id().to_owned(),
4255 buffered.message.id().to_owned(),
4256 );
4257 let searchable = buffered.search_text.is_some();
4258 let message_status = if existing_message_pks.contains(&key) {
4259 counts.messages_matched_total += 1;
4260 if searchable {
4261 counts.messages_matched_searchable += 1;
4262 }
4263 UpsertStatus::Matched
4264 } else {
4265 counts.messages_inserted_total += 1;
4266 if searchable {
4267 counts.messages_inserted_searchable += 1;
4268 }
4269 UpsertStatus::Inserted
4270 };
4271 let pk = Value::Array(vec![Value::String(key.0), Value::String(key.1)]);
4272 outcomes.push(success_outcome(
4273 buffered.index,
4274 "message",
4275 pk,
4276 message_status,
4277 searchable,
4278 ));
4279 for part in &buffered.parts {
4280 let part_key = (
4281 part.part.session_id.clone(),
4282 part.part.message_id.clone(),
4283 part.part.id.clone(),
4284 );
4285 let part_status = if existing_part_pks.contains(&part_key) {
4286 counts.parts_matched += 1;
4287 UpsertStatus::Matched
4288 } else {
4289 counts.parts_inserted += 1;
4290 UpsertStatus::Inserted
4291 };
4292 let part_pk = Value::Array(vec![
4293 Value::String(part_key.0),
4294 Value::String(part_key.1),
4295 Value::String(part_key.2),
4296 ]);
4297 outcomes.push(success_outcome(
4298 part.index,
4299 "part",
4300 part_pk,
4301 part_status,
4302 false,
4303 ));
4304 }
4305 }
4306 outcomes
4307}
4308
4309fn success_outcome(
4310 index: usize,
4311 kind: &'static str,
4312 pk: Value,
4313 status: UpsertStatus,
4314 searchable: bool,
4315) -> RowOutcome {
4316 let status = match status {
4317 UpsertStatus::Inserted => OutcomeStatus::Inserted,
4318 UpsertStatus::Matched => OutcomeStatus::Matched,
4319 };
4320 RowOutcome {
4321 index,
4322 kind,
4323 pk,
4324 status,
4325 error: None,
4326 searchable,
4327 }
4328}
4329
4330#[derive(Debug, Clone, PartialEq, Eq)]
4331enum IngestError {
4332 ImmutableField {
4337 field: &'static str,
4338 session_id: String,
4339 stored: String,
4340 attempted: String,
4341 },
4342}
4343
4344impl std::fmt::Display for IngestError {
4345 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4346 match self {
4347 Self::ImmutableField {
4348 field,
4349 session_id,
4350 stored,
4351 attempted,
4352 } => write!(
4353 formatter,
4354 "session {session_id} {field} is immutable: stored {stored:?}, attempted {attempted:?}",
4355 ),
4356 }
4357 }
4358}
4359
4360impl std::error::Error for IngestError {}
4361
4362fn ensure_immutable_match(
4366 existing: &Session,
4367 incoming: &Session,
4368) -> std::result::Result<(), IngestError> {
4369 if existing.source_agent != incoming.source_agent {
4370 return Err(IngestError::ImmutableField {
4371 field: "source_agent",
4372 session_id: incoming.id.clone(),
4373 stored: existing.source_agent.clone(),
4374 attempted: incoming.source_agent.clone(),
4375 });
4376 }
4377 if existing.project != incoming.project {
4378 return Err(IngestError::ImmutableField {
4379 field: "project",
4380 session_id: incoming.id.clone(),
4381 stored: (*existing.project).clone(),
4382 attempted: (*incoming.project).clone(),
4383 });
4384 }
4385 Ok(())
4386}
4387
4388pub fn search_text(message: &Message, parts: &[Part]) -> Option<String> {
4389 use crate::wire::Provenance;
4390 let mut chunks: Vec<String> = Vec::new();
4391 for part in parts {
4392 if part.provenance != Provenance::Conversational {
4395 continue;
4396 }
4397 match (message.role(), &part.kind) {
4398 (Role::User | Role::Assistant, PartKind::Text { text }) => {
4399 if let Some(text) = text {
4400 chunks.push(text.to_string());
4401 }
4402 }
4403 (
4404 Role::User | Role::Assistant,
4405 PartKind::File {
4406 media_type,
4407 file_name,
4408 data,
4409 },
4410 ) => {
4411 if let Some(file_name) = file_name {
4412 chunks.push(file_name.clone());
4413 }
4414 if let Some(media_type) = media_type {
4415 chunks.push(media_type.clone());
4416 }
4417 if let FileData::Url(uri) = data {
4418 chunks.push(uri.clone());
4419 }
4420 }
4421 (
4422 Role::System | Role::Tool,
4423 PartKind::Text { .. }
4424 | PartKind::Reasoning { .. }
4425 | PartKind::File { .. }
4426 | PartKind::ToolCall { .. }
4427 | PartKind::ToolResult { .. }
4428 | PartKind::ToolApprovalRequest { .. }
4429 | PartKind::ToolApprovalResponse { .. },
4430 )
4431 | (
4432 Role::User | Role::Assistant,
4433 PartKind::Reasoning { .. }
4434 | PartKind::ToolCall { .. }
4435 | PartKind::ToolResult { .. }
4436 | PartKind::ToolApprovalRequest { .. }
4437 | PartKind::ToolApprovalResponse { .. },
4438 ) => {}
4439 }
4440 }
4441
4442 let text = chunks
4443 .into_iter()
4444 .filter(|chunk| !chunk.trim().is_empty())
4445 .collect::<Vec<_>>()
4446 .join("\n");
4447 if text.is_empty() { None } else { Some(text) }
4448}
4449
4450#[derive(Debug, Clone, PartialEq, Eq)]
4452pub struct SearchText(String);
4453
4454impl SearchText {
4455 pub fn as_str(&self) -> &str {
4456 &self.0
4457 }
4458
4459 pub fn into_inner(self) -> String {
4460 self.0
4461 }
4462}
4463
4464impl AsRef<str> for SearchText {
4465 fn as_ref(&self) -> &str {
4466 &self.0
4467 }
4468}
4469
4470#[derive(Debug, Clone, PartialEq)]
4471pub struct MessageWithParts {
4472 pub message: Message,
4473 pub parts: Vec<Part>,
4474}
4475
4476#[derive(Debug, Clone, PartialEq)]
4477pub struct SessionWithMessages {
4478 pub session: Session,
4479 pub messages: Vec<MessageWithParts>,
4480}
4481
4482#[derive(Debug, Clone)]
4483pub struct SessionViewParams<'a> {
4484 pub at_message_id: Option<&'a str>,
4487 pub after_message_id: Option<&'a str>,
4489 pub before_message_id: Option<&'a str>,
4491 pub limit: usize,
4492 pub budget_bytes: usize,
4493 pub session_from: SessionFrom,
4495}
4496
4497#[derive(Debug, Clone)]
4498pub struct MessageViewParams {
4499 pub context_before: usize,
4501 pub context_after: usize,
4503 pub budget_bytes: usize,
4504}
4505
4506#[derive(Debug, Clone, PartialEq)]
4512pub enum GetLookup<T> {
4513 NotFound,
4514 UnknownAnchor,
4515 Found(T),
4516}
4517
4518#[derive(Debug, Clone, PartialEq)]
4522pub struct SessionPage {
4523 pub session: Session,
4524 pub messages: Vec<RetrievedMessage>,
4525 pub before_remaining: usize,
4526 pub after_remaining: usize,
4527}
4528
4529#[derive(Debug, Clone, PartialEq)]
4533pub struct MessagePage {
4534 pub session: Session,
4535 pub target: RetrievedMessage,
4536 pub target_parts: Vec<Part>,
4537 pub target_parts_remaining: usize,
4538 pub siblings: Vec<RetrievedMessage>,
4539}
4540
4541#[derive(Debug, Clone, PartialEq)]
4542pub struct RetrievedMessage {
4543 pub id: String,
4544 pub role: Role,
4545 pub timestamp: DateTime<Utc>,
4546 pub text: Option<String>,
4547 pub content: Option<String>,
4548 pub parts: Vec<Part>,
4549}
4550
4551#[derive(Debug, Clone)]
4552struct ScanRow {
4553 id: String,
4554 role: Role,
4555 timestamp: DateTime<Utc>,
4556 text: Option<String>,
4557 content: Option<String>,
4558}
4559
4560#[derive(Debug, Clone)]
4563pub struct ConversationalRow {
4564 pub session_id: String,
4565 pub message_id: String,
4566 pub role: Role,
4567 pub timestamp: DateTime<Utc>,
4568 pub text: SearchText,
4569}
4570
4571fn page_by<T>(items: &[T], limit: usize, budget_bytes: usize, size: impl Fn(&T) -> usize) -> usize {
4576 let capped = items.len().min(limit.clamp(1, 1000));
4577 let mut acc = 0usize;
4578 let mut emitted = 0usize;
4579 for item in &items[..capped] {
4580 let next = acc.saturating_add(size(item));
4581 if emitted > 0 && next > budget_bytes {
4582 break;
4583 }
4584 acc = next;
4585 emitted += 1;
4586 }
4587 emitted
4588}
4589
4590fn page_tail<T>(
4595 items: &[T],
4596 limit: usize,
4597 budget_bytes: usize,
4598 size: impl Fn(&T) -> usize,
4599) -> usize {
4600 let cap = limit.clamp(1, 1000);
4601 let mut acc = 0usize;
4602 let mut emitted = 0usize;
4603 for item in items.iter().rev() {
4604 if emitted >= cap {
4605 break;
4606 }
4607 let next = acc.saturating_add(size(item));
4608 if emitted > 0 && next > budget_bytes {
4609 break;
4610 }
4611 acc = next;
4612 emitted += 1;
4613 }
4614 emitted
4615}
4616
4617fn role_from_str(value: &str) -> Result<Role> {
4618 match value {
4619 "system" => Ok(Role::System),
4620 "user" => Ok(Role::User),
4621 "assistant" => Ok(Role::Assistant),
4622 "tool" => Ok(Role::Tool),
4623 other => anyhow::bail!("unknown message role {other}"),
4624 }
4625}
4626
4627const MESSAGE_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] = &[
4640 (
4641 "session_id",
4642 BuiltinIndexType::BTree,
4643 MESSAGES_SESSION_ID_INDEX,
4644 ),
4645 (
4646 "source_agent",
4647 BuiltinIndexType::Bitmap,
4648 "messages_source_agent_bitmap",
4649 ),
4650];
4651
4652const PARTS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] = &[
4658 (
4659 "session_id",
4660 BuiltinIndexType::BTree,
4661 "parts_session_id_btree",
4662 ),
4663 (
4664 "message_id",
4665 BuiltinIndexType::BTree,
4666 "parts_message_id_btree",
4667 ),
4668 (
4669 "tool_name",
4670 BuiltinIndexType::BTree,
4671 "parts_tool_name_btree",
4672 ),
4673];
4674
4675const SESSIONS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] =
4678 &[("id", BuiltinIndexType::BTree, "sessions_id_btree")];
4679
4680const COPY_SESSION_IN_CHUNK: usize = 512;
4684
4685fn in_predicate(column: &'static str, values: &[String]) -> Predicate {
4686 Predicate::In(
4687 column,
4688 values.iter().cloned().map(ScalarValue::String).collect(),
4689 )
4690}
4691
4692fn embedded_scope(filter: &Predicate) -> Predicate {
4705 filter.clone()
4706}
4707
4708pub const DEFAULT_NPROBES: usize = 32;
4713
4714fn apply_vector_search_knobs(
4721 scanner: &mut lance::dataset::scanner::Scanner,
4722 search: Option<&config::SearchConfig>,
4723) {
4724 let nprobes = search
4725 .and_then(|cfg| cfg.nprobes)
4726 .unwrap_or(DEFAULT_NPROBES);
4727 scanner.nprobes(nprobes);
4728}
4729
4730pub(crate) const SESSIONS: &str = "sessions";
4734pub(crate) const MESSAGES: &str = "messages";
4735pub(crate) const PARTS: &str = "parts";
4736
4737pub const MESSAGES_SESSION_ID_INDEX: &str = "messages_session_id_btree";
4740
4741pub const MESSAGES_FTS_INDEX: &str = "messages_search_text_fts";
4744
4745pub const MESSAGES_VECTOR_INDEX: &str = "messages_vector_ivfpq";
4752
4753const IVF_SQ_NUM_BITS: u16 = 8;
4758const IVF_SQ_MAX_ITERS: usize = 15;
4759
4760pub fn pond_index_intents() -> IndexIntents {
4763 pond_index_intents_with_vector_threshold(VECTOR_INDEX_ACTIVATION_ROWS)
4764}
4765
4766pub(crate) fn pond_index_intents_with_vector_threshold(vector_threshold: usize) -> IndexIntents {
4770 let mut messages = Vec::with_capacity(MESSAGE_SCALAR_INDICES.len() + 2);
4771 messages.push(IndexIntent {
4772 name: MESSAGES_FTS_INDEX,
4773 column: "search_text",
4774 trigger: IndexTrigger::OnAnyRows,
4775 params: IndexParamsKind::InvertedFtsWord,
4776 });
4777 for (column, kind, name) in MESSAGE_SCALAR_INDICES {
4778 messages.push(IndexIntent {
4779 name,
4780 column,
4781 trigger: IndexTrigger::OnAnyRows,
4782 params: IndexParamsKind::Scalar(kind.clone()),
4783 });
4784 }
4785 messages.push(IndexIntent {
4786 name: MESSAGES_VECTOR_INDEX,
4787 column: "vector",
4788 trigger: IndexTrigger::OnNonNullCount {
4789 column: "vector",
4790 threshold: vector_threshold,
4791 },
4792 params: IndexParamsKind::IvfSqCosine {
4793 num_bits: IVF_SQ_NUM_BITS,
4794 max_iters: IVF_SQ_MAX_ITERS,
4795 },
4796 });
4797 let parts = PARTS_SCALAR_INDICES
4798 .iter()
4799 .map(|(column, kind, name)| IndexIntent {
4800 name,
4801 column,
4802 trigger: IndexTrigger::OnAnyRows,
4803 params: IndexParamsKind::Scalar(kind.clone()),
4804 })
4805 .collect();
4806 let sessions = SESSIONS_SCALAR_INDICES
4807 .iter()
4808 .map(|(column, kind, name)| IndexIntent {
4809 name,
4810 column,
4811 trigger: IndexTrigger::OnAnyRows,
4812 params: IndexParamsKind::Scalar(kind.clone()),
4813 })
4814 .collect();
4815 IndexIntents {
4816 sessions,
4817 messages,
4818 parts,
4819 }
4820}
4821
4822pub const DEFAULT_EMBEDDING_DIM: usize = 384;
4826
4827static EMBEDDING_DIM_RUNTIME: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4833
4834pub fn embedding_dim() -> usize {
4837 EMBEDDING_DIM_RUNTIME
4838 .get()
4839 .copied()
4840 .unwrap_or(DEFAULT_EMBEDDING_DIM)
4841}
4842
4843pub fn init_embedding_dim(dim: usize) {
4845 EMBEDDING_DIM_RUNTIME.get_or_init(|| dim);
4846}
4847
4848pub(crate) fn write_params_for_create() -> WriteParams {
4855 WriteParams {
4856 data_storage_version: Some(LanceFileVersion::V2_1),
4857 enable_v2_manifest_paths: true,
4858 enable_stable_row_ids: true,
4859 auto_cleanup: Some(AutoCleanupParams {
4860 interval: 20,
4861 older_than: chrono::TimeDelta::days(1),
4862 }),
4863 skip_auto_cleanup: true,
4864 ..WriteParams::default()
4865 }
4866}
4867
4868fn export_schema(table: Table) -> Arc<Schema> {
4869 match table {
4870 Table::Sessions => session_schema(),
4871 Table::Messages => message_schema(),
4872 Table::Parts => part_schema(),
4873 }
4874}
4875
4876fn archive_schema_backfill(dataset: &Dataset, table: Table) -> Result<Option<ColumnBackfill>> {
4881 use std::collections::BTreeSet;
4882 let expected = export_schema(table);
4883 let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
4884 let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
4885 let expected_names: BTreeSet<&str> = expected
4886 .fields()
4887 .iter()
4888 .map(|f| f.name().as_str())
4889 .collect();
4890 let extra: Vec<&str> = actual_names.difference(&expected_names).copied().collect();
4891 if !extra.is_empty() {
4892 anyhow::bail!(
4893 "{} archive table has columns {actual_names:?} but this pond build expects \
4894 {expected_names:?} - the archive was written by a newer pond; upgrade pond \
4895 to restore it",
4896 table.as_str(),
4897 );
4898 }
4899 let missing: Vec<Field> = expected
4900 .fields()
4901 .iter()
4902 .filter(|f| !actual_names.contains(f.name().as_str()))
4903 .map(|f| f.as_ref().clone())
4904 .collect();
4905 if missing.is_empty() {
4906 return Ok(None);
4907 }
4908 column_backfill(table.as_str(), &missing)
4909 .map(Some)
4910 .with_context(|| {
4911 format!(
4912 "{} archive predates a schema change this pond build cannot bridge; \
4913 restore it with the pond version that wrote it",
4914 table.as_str(),
4915 )
4916 })
4917}
4918
4919fn upgraded_batch(batch: &RecordBatch, spec: &ColumnBackfill) -> Result<RecordBatch> {
4925 let derived = (spec.mapper)(batch)?;
4926 let mut fields: Vec<Field> = batch
4927 .schema()
4928 .fields()
4929 .iter()
4930 .map(|field| field.as_ref().clone())
4931 .collect();
4932 let mut columns = batch.columns().to_vec();
4933 for (field, column) in spec.output_schema.fields().iter().zip(derived.columns()) {
4934 fields.push(field.as_ref().clone());
4935 columns.push(column.clone());
4936 }
4937 RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
4938 .context("failed to assemble upgraded batch")
4939}
4940
4941async fn open_archive_table(table: Table, source: &Path) -> Result<Dataset> {
4942 let source_uri = source
4943 .to_str()
4944 .with_context(|| format!("archive path is not UTF-8: {}", source.display()))?;
4945 Dataset::open(source_uri)
4946 .await
4947 .with_context(|| format!("failed to open {} archive table", table.as_str()))
4948}
4949
4950pub(crate) fn pk_columns(table: Table) -> &'static [&'static str] {
4956 match table {
4957 Table::Sessions => &["id"],
4958 Table::Messages => &["session_id", "id"],
4959 Table::Parts => &["session_id", "message_id", "id"],
4960 }
4961}
4962
4963fn composite_key(batch: &RecordBatch, pk: &[&str], row: usize) -> Result<Vec<String>> {
4965 pk.iter()
4966 .map(|column| string(batch, column, row)?.with_context(|| format!("{column} is null")))
4967 .collect()
4968}
4969
4970pub(crate) fn session_schema() -> Arc<Schema> {
4971 Arc::new(Schema::new(vec![
4972 primary_field("id", DataType::Utf8, false),
4973 Field::new("parent_session_id", DataType::Utf8, true),
4974 Field::new("parent_message_id", DataType::Utf8, true),
4975 Field::new("source_agent", DataType::Utf8, false),
4976 Field::new(
4977 "created_at",
4978 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4979 false,
4980 ),
4981 Field::new("project", DataType::Utf8, false),
4982 json_field("options", false),
4983 ]))
4984}
4985
4986pub(crate) fn message_schema() -> Arc<Schema> {
4987 Arc::new(Schema::new(vec![
4988 primary_field("session_id", DataType::Utf8, false),
4989 primary_field("id", DataType::Utf8, false),
4990 Field::new(
4991 "timestamp",
4992 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4993 false,
4994 ),
4995 Field::new("role", DataType::Utf8, false),
4996 Field::new("source_agent", DataType::Utf8, false),
4997 Field::new("project", DataType::Utf8, false),
4998 Field::new("content", DataType::Utf8, true),
4999 Field::new("search_text", DataType::Utf8, true),
5000 Field::new("vector", embedding_vector_type(), true),
5004 Field::new("embedding_model", DataType::Utf8, true),
5005 json_field("options", false),
5006 ]))
5007}
5008
5009pub(crate) type BackfillMapper = Box<dyn Fn(&RecordBatch) -> Result<RecordBatch> + Send + Sync>;
5011
5012pub(crate) struct ColumnBackfill {
5018 pub read_columns: Vec<String>,
5019 pub output_schema: Arc<Schema>,
5020 pub mapper: BackfillMapper,
5021}
5022
5023pub(crate) fn column_backfill(table_name: &str, missing: &[Field]) -> Result<ColumnBackfill> {
5026 match table_name {
5027 PARTS => parts_backfill(missing),
5028 other => anyhow::bail!(
5029 "table {other} is missing columns {:?} that this pond build cannot derive from \
5030 stored data - use the pond version that wrote it",
5031 missing.iter().map(|f| f.name()).collect::<Vec<_>>(),
5032 ),
5033 }
5034}
5035
5036fn parts_backfill(missing: &[Field]) -> Result<ColumnBackfill> {
5037 for field in missing {
5038 anyhow::ensure!(
5039 ["tool_name", "call_id", "is_failure"].contains(&field.name().as_str())
5040 && field.is_nullable(),
5041 "column {PARTS}.{} cannot be derived from stored data - use the pond version \
5042 that wrote it",
5043 field.name(),
5044 );
5045 }
5046 let output_schema = Arc::new(Schema::new(missing.to_vec()));
5047 let schema = output_schema.clone();
5048 let mapper = move |batch: &RecordBatch| -> Result<RecordBatch> {
5049 let rows = batch.num_rows();
5050 let mut names: Vec<Option<String>> = Vec::with_capacity(rows);
5051 let mut call_ids: Vec<Option<String>> = Vec::with_capacity(rows);
5052 let mut failures: Vec<Option<bool>> = Vec::with_capacity(rows);
5053 for row in 0..rows {
5054 let type_name = string(batch, "type", row)?.context("part type is null")?;
5055 let kind = match type_name.as_str() {
5058 "tool_call" | "tool_result" | "tool_approval_request" => {
5059 let body =
5060 json_column(batch, "variant_data", row)?.context("variant_data is null")?;
5061 match part_kind_from_json(&type_name, &body, None) {
5066 Ok(kind) => Some(kind),
5067 Err(error) => {
5068 tracing::warn!(
5069 session_id = string(batch, "session_id", row)?.as_deref(),
5070 part_id = string(batch, "id", row)?.as_deref(),
5071 error = %format!("{error:#}"),
5072 "stored tool part body failed to decode; its backfilled \
5073 cells stay NULL",
5074 );
5075 None
5076 }
5077 }
5078 }
5079 _ => None,
5080 };
5081 let (name, call_id, is_failure) = kind
5082 .as_ref()
5083 .map(tool_identity)
5084 .unwrap_or((None, None, None));
5085 names.push(name.map(str::to_owned));
5086 call_ids.push(call_id.map(str::to_owned));
5087 failures.push(is_failure);
5088 }
5089 let arrays: Vec<ArrayRef> = schema
5090 .fields()
5091 .iter()
5092 .map(|field| -> ArrayRef {
5093 match field.name().as_str() {
5094 "tool_name" => Arc::new(StringArray::from(names.clone())),
5095 "call_id" => Arc::new(StringArray::from(call_ids.clone())),
5096 _ => Arc::new(BooleanArray::from(failures.clone())),
5097 }
5098 })
5099 .collect();
5100 RecordBatch::try_new(schema.clone(), arrays).context("failed to build backfill batch")
5101 };
5102 Ok(ColumnBackfill {
5103 read_columns: vec![
5104 "session_id".to_owned(),
5105 "id".to_owned(),
5106 "type".to_owned(),
5107 "variant_data".to_owned(),
5108 ],
5109 output_schema,
5110 mapper: Box::new(mapper),
5111 })
5112}
5113
5114pub(crate) fn part_schema() -> Arc<Schema> {
5115 Arc::new(Schema::new(vec![
5116 primary_field("session_id", DataType::Utf8, false),
5117 primary_field("message_id", DataType::Utf8, false),
5118 primary_field("id", DataType::Utf8, false),
5119 Field::new("ordinal", DataType::Int32, false),
5120 Field::new("type", DataType::Utf8, false),
5121 Field::new("provenance", DataType::Utf8, false),
5124 Field::new("tool_name", DataType::Utf8, true),
5129 Field::new("call_id", DataType::Utf8, true),
5130 Field::new("is_failure", DataType::Boolean, true),
5131 json_field("variant_data", false),
5132 legacy_blob_field("data", true),
5133 json_field("options", false),
5134 ]))
5135}
5136
5137pub(crate) fn empty_batch(schema: Arc<Schema>) -> Result<RecordBatch> {
5138 let arrays = schema
5139 .fields()
5140 .iter()
5141 .map(|field| lance::deps::arrow_array::new_empty_array(field.data_type()))
5142 .collect();
5143 RecordBatch::try_new(schema, arrays).context("failed to build empty Lance batch")
5144}
5145
5146pub(crate) fn empty_reader(
5147 schema: Arc<Schema>,
5148) -> Result<
5149 RecordBatchIterator<
5150 std::vec::IntoIter<Result<RecordBatch, lance::deps::arrow_schema::ArrowError>>,
5151 >,
5152> {
5153 let batch = empty_batch(schema.clone())?;
5154 Ok(RecordBatchIterator::new(
5155 vec![Ok(batch)].into_iter(),
5156 schema,
5157 ))
5158}
5159
5160pub(crate) struct MessageBatchRow<'a> {
5161 pub message: &'a Message,
5162 pub source_agent: &'a str,
5163 pub project: &'a str,
5164 pub search_text: Option<&'a str>,
5165}
5166
5167fn embedding_vector_type() -> DataType {
5173 DataType::FixedSizeList(
5174 Arc::new(Field::new("item", DataType::Float16, true)),
5175 embedding_dim() as i32,
5176 )
5177}
5178
5179fn embedding_update_schema() -> Arc<Schema> {
5183 Arc::new(Schema::new(vec![
5184 primary_field("session_id", DataType::Utf8, false),
5185 primary_field("id", DataType::Utf8, false),
5186 Field::new("vector", embedding_vector_type(), true),
5187 Field::new("embedding_model", DataType::Utf8, true),
5188 ]))
5189}
5190
5191fn embedding_columns(vectors: &[Option<Vec<f32>>]) -> Result<(ArrayRef, ArrayRef)> {
5195 let dim = embedding_dim();
5196 if vectors.iter().all(Option::is_none) {
5199 return Ok((
5200 new_null_array(&embedding_vector_type(), vectors.len()),
5201 new_null_array(&DataType::Utf8, vectors.len()),
5202 ));
5203 }
5204 let mut builder = FixedSizeListBuilder::new(
5205 Float16Builder::with_capacity(vectors.len() * dim),
5206 dim as i32,
5207 )
5208 .with_field(Arc::new(Field::new("item", DataType::Float16, true)));
5209 let mut models: Vec<Option<&str>> = Vec::with_capacity(vectors.len());
5210 for vector in vectors {
5211 match vector {
5212 Some(values) => {
5213 if values.len() != dim {
5214 anyhow::bail!("inline embedding has dim {}, expected {dim}", values.len());
5215 }
5216 for value in values {
5217 builder.values().append_value(half::f16::from_f32(*value));
5218 }
5219 builder.append(true);
5220 models.push(Some(embed::model_id()));
5221 }
5222 None => {
5223 for _ in 0..dim {
5224 builder.values().append_null();
5225 }
5226 builder.append(false);
5227 models.push(None);
5228 }
5229 }
5230 }
5231 Ok((
5232 Arc::new(builder.finish()) as ArrayRef,
5233 Arc::new(StringArray::from(models)) as ArrayRef,
5234 ))
5235}
5236
5237pub(crate) fn embedding_update_batch(rows: &[EmbeddedMessage]) -> Result<RecordBatch> {
5240 let dim = embedding_dim();
5241 let mut flat = Vec::with_capacity(rows.len() * dim);
5242 for row in rows {
5243 if row.vector.len() != dim {
5244 anyhow::bail!(
5245 "embedding for message {} has dim {}, expected {dim}",
5246 row.id,
5247 row.vector.len(),
5248 );
5249 }
5250 flat.extend(row.vector.iter().map(|value| half::f16::from_f32(*value)));
5251 }
5252 let values = Float16Array::from(flat);
5253 let item_field = Arc::new(Field::new("item", DataType::Float16, true));
5254 let vectors = FixedSizeListArray::try_new(item_field, dim as i32, Arc::new(values), None)
5255 .context("failed to build embedding vector column")?;
5256
5257 RecordBatch::try_new(
5258 embedding_update_schema(),
5259 vec![
5260 Arc::new(StringArray::from(
5261 rows.iter()
5262 .map(|row| row.session_id.as_str())
5263 .collect::<Vec<_>>(),
5264 )),
5265 Arc::new(StringArray::from(
5266 rows.iter().map(|row| row.id.as_str()).collect::<Vec<_>>(),
5267 )),
5268 Arc::new(vectors),
5269 Arc::new(StringArray::from(vec![embed::model_id(); rows.len()])),
5270 ],
5271 )
5272 .context("failed to build embedding update batch")
5273}
5274
5275const COLUMN_BYTE_BUDGET: usize = 1 << 30;
5280
5281fn chunk_ranges(cells: &[usize]) -> Vec<std::ops::Range<usize>> {
5286 let mut chunks = Vec::new();
5287 let mut start = 0usize;
5288 let mut running = 0usize;
5289 for (index, &row) in cells.iter().enumerate() {
5290 if running + row > COLUMN_BYTE_BUDGET && index > start {
5291 chunks.push(start..index);
5292 start = index;
5293 running = 0;
5294 }
5295 running += row;
5296 }
5297 if start < cells.len() {
5298 chunks.push(start..cells.len());
5299 }
5300 chunks
5301}
5302
5303fn guard_cell(table: &str, pk: &str, bytes: usize) -> Result<()> {
5304 if bytes >= COLUMN_BYTE_BUDGET {
5305 anyhow::bail!(
5306 "{table} row {pk}: a {bytes}-byte text cell meets the per-cell ceiling and would \
5307 overflow Arrow's i32 offset buffer"
5308 );
5309 }
5310 Ok(())
5311}
5312
5313async fn merge_insert_chunks(
5314 handle: &Handle,
5315 table: Table,
5316 batches: Vec<RecordBatch>,
5317) -> Result<u64> {
5318 let mut inserted = 0u64;
5319 for batch in batches {
5320 let rows = batch.num_rows();
5321 inserted += handle.merge_insert(table, batch, rows).await?;
5322 }
5323 Ok(inserted)
5324}
5325
5326pub(crate) fn sessions_batches(sessions: &[Session]) -> Result<Vec<RecordBatch>> {
5327 let options = sessions
5328 .iter()
5329 .map(|session| json_bytes(&session.options))
5330 .collect::<Result<Vec<_>>>()?;
5331 let mut cells = Vec::with_capacity(sessions.len());
5332 for (session, encoded) in sessions.iter().zip(&options) {
5333 let columns = [
5334 session.id.len(),
5335 session.parent_session_id.as_deref().map_or(0, str::len),
5336 session.parent_message_id.as_deref().map_or(0, str::len),
5337 session.source_agent.len(),
5338 session.project.as_str().len(),
5339 encoded.len(),
5340 ];
5341 for bytes in columns {
5342 guard_cell("sessions", &session.id, bytes)?;
5343 }
5344 cells.push(columns.iter().sum());
5345 }
5346 chunk_ranges(&cells)
5347 .into_iter()
5348 .map(|range| sessions_chunk(&sessions[range.clone()], &options[range]))
5349 .collect()
5350}
5351
5352fn sessions_chunk(sessions: &[Session], options: &[Vec<u8>]) -> Result<RecordBatch> {
5353 let schema = session_schema();
5354 RecordBatch::try_new(
5355 schema.clone(),
5356 vec![
5357 Arc::new(StringArray::from(
5358 sessions
5359 .iter()
5360 .map(|session| session.id.as_str())
5361 .collect::<Vec<_>>(),
5362 )),
5363 Arc::new(StringArray::from(
5364 sessions
5365 .iter()
5366 .map(|session| session.parent_session_id.as_deref())
5367 .collect::<Vec<_>>(),
5368 )),
5369 Arc::new(StringArray::from(
5370 sessions
5371 .iter()
5372 .map(|session| session.parent_message_id.as_deref())
5373 .collect::<Vec<_>>(),
5374 )),
5375 Arc::new(StringArray::from(
5376 sessions
5377 .iter()
5378 .map(|session| session.source_agent.as_str())
5379 .collect::<Vec<_>>(),
5380 )),
5381 Arc::new(
5382 TimestampMicrosecondArray::from(
5383 sessions
5384 .iter()
5385 .map(|session| micros(session.created_at))
5386 .collect::<Vec<_>>(),
5387 )
5388 .with_timezone("UTC"),
5389 ),
5390 Arc::new(StringArray::from(
5391 sessions
5392 .iter()
5393 .map(|session| session.project.as_str())
5394 .collect::<Vec<_>>(),
5395 )),
5396 Arc::new(LargeBinaryArray::from_iter_values(
5397 options.iter().map(Vec::as_slice),
5398 )),
5399 ],
5400 )
5401 .context("failed to build session batch")
5402}
5403
5404pub(crate) fn messages_batches(
5407 rows: &[MessageBatchRow<'_>],
5408 vectors: &[Option<Vec<f32>>],
5409) -> Result<Vec<RecordBatch>> {
5410 debug_assert_eq!(rows.len(), vectors.len(), "vectors must align with rows");
5411 let options = rows
5412 .iter()
5413 .map(|row| json_bytes(row.message.options()))
5414 .collect::<Result<Vec<_>>>()?;
5415 let mut cells = Vec::with_capacity(rows.len());
5416 for (row, encoded) in rows.iter().zip(&options) {
5417 let columns = [
5418 row.message.session_id().len(),
5419 row.message.id().len(),
5420 row.message.role().as_str().len(),
5421 row.source_agent.len(),
5422 row.project.len(),
5423 row.message.system_content().map_or(0, str::len),
5424 row.search_text.map_or(0, str::len),
5425 encoded.len(),
5426 ];
5427 for bytes in columns {
5428 guard_cell("messages", row.message.id(), bytes)?;
5429 }
5430 cells.push(columns.iter().sum());
5431 }
5432 chunk_ranges(&cells)
5433 .into_iter()
5434 .map(|range| {
5435 messages_chunk(
5436 &rows[range.clone()],
5437 &options[range.clone()],
5438 &vectors[range],
5439 )
5440 })
5441 .collect()
5442}
5443
5444fn messages_chunk(
5445 rows: &[MessageBatchRow<'_>],
5446 options: &[Vec<u8>],
5447 vectors: &[Option<Vec<f32>>],
5448) -> Result<RecordBatch> {
5449 let schema = message_schema();
5450 let (vector_column, embedding_model) = embedding_columns(vectors)?;
5451 RecordBatch::try_new(
5452 schema.clone(),
5453 vec![
5454 Arc::new(StringArray::from(
5455 rows.iter()
5456 .map(|row| row.message.session_id())
5457 .collect::<Vec<_>>(),
5458 )),
5459 Arc::new(StringArray::from(
5460 rows.iter().map(|row| row.message.id()).collect::<Vec<_>>(),
5461 )),
5462 Arc::new(
5463 TimestampMicrosecondArray::from(
5464 rows.iter()
5465 .map(|row| micros(row.message.timestamp()))
5466 .collect::<Vec<_>>(),
5467 )
5468 .with_timezone("UTC"),
5469 ),
5470 Arc::new(StringArray::from(
5471 rows.iter()
5472 .map(|row| row.message.role().as_str())
5473 .collect::<Vec<_>>(),
5474 )),
5475 Arc::new(StringArray::from(
5476 rows.iter().map(|row| row.source_agent).collect::<Vec<_>>(),
5477 )),
5478 Arc::new(StringArray::from(
5479 rows.iter().map(|row| row.project).collect::<Vec<_>>(),
5480 )),
5481 Arc::new(StringArray::from(
5482 rows.iter()
5483 .map(|row| row.message.system_content())
5484 .collect::<Vec<_>>(),
5485 )),
5486 Arc::new(StringArray::from(
5487 rows.iter().map(|row| row.search_text).collect::<Vec<_>>(),
5488 )),
5489 vector_column,
5494 embedding_model,
5495 Arc::new(LargeBinaryArray::from_iter_values(
5496 options.iter().map(Vec::as_slice),
5497 )),
5498 ],
5499 )
5500 .context("failed to build message batch")
5501}
5502
5503pub(crate) fn parts_batches(parts: &[Part]) -> Result<Vec<RecordBatch>> {
5504 let variant_data = parts
5505 .iter()
5506 .map(|part| part_variant_json(&part.kind))
5507 .collect::<Result<Vec<_>>>()?;
5508 let options = parts
5509 .iter()
5510 .map(|part| json_bytes(&part.options))
5511 .collect::<Result<Vec<_>>>()?;
5512 let mut cells = Vec::with_capacity(parts.len());
5513 for ((part, variant), encoded) in parts.iter().zip(&variant_data).zip(&options) {
5516 let columns = [
5517 part.session_id.len(),
5518 part.message_id.len(),
5519 part.id.len(),
5520 part.kind.type_name().len(),
5521 part.provenance.as_str().len(),
5522 variant.len(),
5523 encoded.len(),
5524 ];
5525 for bytes in columns {
5526 guard_cell("parts", &part.id, bytes)?;
5527 }
5528 cells.push(columns.iter().sum());
5529 }
5530 chunk_ranges(&cells)
5531 .into_iter()
5532 .map(|range| {
5533 parts_chunk(
5534 &parts[range.clone()],
5535 &variant_data[range.clone()],
5536 &options[range],
5537 )
5538 })
5539 .collect()
5540}
5541
5542fn parts_chunk(
5543 parts: &[Part],
5544 variant_data: &[Vec<u8>],
5545 options: &[Vec<u8>],
5546) -> Result<RecordBatch> {
5547 let schema = part_schema();
5548 let blob_payloads: Vec<Option<&[u8]>> = parts
5552 .iter()
5553 .map(|part| match &part.kind {
5554 PartKind::File { data, .. } => Some(match data {
5555 FileData::String(value) => value.as_bytes(),
5556 FileData::Bytes(value) => value.as_slice(),
5557 FileData::Url(value) => value.as_bytes(),
5558 }),
5559 PartKind::Text { .. }
5560 | PartKind::Reasoning { .. }
5561 | PartKind::ToolCall { .. }
5562 | PartKind::ToolResult { .. }
5563 | PartKind::ToolApprovalRequest { .. }
5564 | PartKind::ToolApprovalResponse { .. } => None,
5565 })
5566 .collect();
5567 let blob_array = LargeBinaryArray::from_iter(blob_payloads);
5568
5569 let mut tool_names: Vec<Option<&str>> = Vec::with_capacity(parts.len());
5570 let mut call_ids: Vec<Option<&str>> = Vec::with_capacity(parts.len());
5571 let mut failures: Vec<Option<bool>> = Vec::with_capacity(parts.len());
5572 for part in parts {
5573 let (name, call_id, is_failure) = tool_identity(&part.kind);
5574 tool_names.push(name);
5575 call_ids.push(call_id);
5576 failures.push(is_failure);
5577 }
5578
5579 RecordBatch::try_new(
5580 schema.clone(),
5581 vec![
5582 Arc::new(StringArray::from(
5583 parts
5584 .iter()
5585 .map(|part| part.session_id.as_str())
5586 .collect::<Vec<_>>(),
5587 )),
5588 Arc::new(StringArray::from(
5589 parts
5590 .iter()
5591 .map(|part| part.message_id.as_str())
5592 .collect::<Vec<_>>(),
5593 )),
5594 Arc::new(StringArray::from(
5595 parts
5596 .iter()
5597 .map(|part| part.id.as_str())
5598 .collect::<Vec<_>>(),
5599 )),
5600 Arc::new(Int32Array::from(
5601 parts.iter().map(|part| part.ordinal).collect::<Vec<_>>(),
5602 )),
5603 Arc::new(StringArray::from(
5604 parts
5605 .iter()
5606 .map(|part| part.kind.type_name())
5607 .collect::<Vec<_>>(),
5608 )),
5609 Arc::new(StringArray::from(
5610 parts
5611 .iter()
5612 .map(|part| part.provenance.as_str())
5613 .collect::<Vec<_>>(),
5614 )),
5615 Arc::new(StringArray::from(tool_names)),
5616 Arc::new(StringArray::from(call_ids)),
5617 Arc::new(BooleanArray::from(failures)),
5618 Arc::new(LargeBinaryArray::from_iter_values(
5619 variant_data.iter().map(Vec::as_slice),
5620 )),
5621 Arc::new(blob_array),
5622 Arc::new(LargeBinaryArray::from_iter_values(
5623 options.iter().map(Vec::as_slice),
5624 )),
5625 ],
5626 )
5627 .context("failed to build parts batch")
5628}
5629
5630pub(crate) fn session_from_batch(batch: &RecordBatch, row: usize) -> Result<Session> {
5631 Ok(Session {
5632 id: string(batch, "id", row)?.context("session id is null")?,
5633 parent_session_id: string(batch, "parent_session_id", row)?,
5634 parent_message_id: string(batch, "parent_message_id", row)?,
5635 source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5636 created_at: datetime(batch, "created_at", row)?,
5637 project: crate::adapter::Extracted::from_stored(
5638 string(batch, "project", row)?.context("project is null")?,
5639 ),
5640 options: json_parse(&json_column(batch, "options", row)?.context("options is null")?)?,
5641 })
5642}
5643
5644pub struct RowmapOracle(pub Option<Arc<RowMetaSet>>);
5651
5652impl crate::adapter::SkipOracle for RowmapOracle {
5653 fn session_max_ts(&self, session_id: &str) -> Option<i64> {
5654 self.0.as_ref()?.lookup_max_ts(session_id)
5655 }
5656
5657 fn is_empty(&self) -> bool {
5658 self.0.as_ref().is_none_or(|set| set.is_empty())
5659 }
5660}
5661
5662fn row_meta_entry(batch: &RecordBatch, row_id: u64, row: usize) -> Result<RowMetaEntry> {
5663 Ok(RowMetaEntry {
5664 row_id,
5665 session_id: string(batch, "session_id", row)?.context("session_id is null")?,
5666 message_id: string(batch, "id", row)?.context("message id is null")?,
5667 role: string(batch, "role", row)?.context("role is null")?,
5668 project: string(batch, "project", row)?.context("project is null")?,
5669 source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5670 timestamp_micros: datetime(batch, "timestamp", row)?.timestamp_micros(),
5671 search_text: string(batch, "search_text", row)?.unwrap_or_default(),
5672 })
5673}
5674
5675pub(crate) fn message_meta_from_batch(batch: &RecordBatch, row: usize) -> Result<MessageMeta> {
5676 Ok(MessageMeta {
5677 message_id: string(batch, "id", row)?.context("id is null")?,
5678 session_id: string(batch, "session_id", row)?.context("session_id is null")?,
5679 role: string(batch, "role", row)?.context("role is null")?,
5680 project: string(batch, "project", row)?.context("project is null")?,
5681 source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5682 timestamp: datetime(batch, "timestamp", row)?,
5683 search_text: string(batch, "search_text", row)?.unwrap_or_default(),
5684 })
5685}
5686
5687pub(crate) fn message_from_batch(batch: &RecordBatch, row: usize) -> Result<Message> {
5688 let id = string(batch, "id", row)?.context("message id is null")?;
5689 let session_id = string(batch, "session_id", row)?.context("message session_id is null")?;
5690 let timestamp = datetime(batch, "timestamp", row)?;
5691 let options =
5692 json_parse(&json_column(batch, "options", row)?.context("message options is null")?)?;
5693
5694 match string(batch, "role", row)?
5695 .context("message role is null")?
5696 .as_str()
5697 {
5698 "system" => Ok(Message::System {
5699 id,
5700 session_id,
5701 timestamp,
5702 content: string(batch, "content", row)?.map(crate::adapter::Extracted::from_stored),
5709 options,
5710 }),
5711 "user" => Ok(Message::User {
5712 id,
5713 session_id,
5714 timestamp,
5715 options,
5716 }),
5717 "assistant" => Ok(Message::Assistant {
5718 id,
5719 session_id,
5720 timestamp,
5721 options,
5722 }),
5723 "tool" => Ok(Message::Tool {
5724 id,
5725 session_id,
5726 timestamp,
5727 options,
5728 }),
5729 other => anyhow::bail!("unknown message role {other}"),
5730 }
5731}
5732
5733pub(crate) fn part_from_batch(
5734 batch: &RecordBatch,
5735 row: usize,
5736 file_data: Option<FileData>,
5737) -> Result<Part> {
5738 let type_name = string(batch, "type", row)?.context("part type is null")?;
5739 let variant_data = json_column(batch, "variant_data", row)?.context("variant_data is null")?;
5740 let provenance = string(batch, "provenance", row)?.context("part provenance is null")?;
5741 Ok(Part {
5742 session_id: string(batch, "session_id", row)?.context("part session_id is null")?,
5743 message_id: string(batch, "message_id", row)?.context("part message_id is null")?,
5744 id: string(batch, "id", row)?.context("part id is null")?,
5745 ordinal: int32(batch, "ordinal", row)?,
5746 provenance: provenance_from_str(&provenance)?,
5747 options: json_parse(&json_column(batch, "options", row)?.context("part options is null")?)?,
5748 kind: part_kind_from_json(&type_name, &variant_data, file_data)?,
5749 })
5750}
5751
5752fn provenance_from_str(value: &str) -> Result<crate::wire::Provenance> {
5753 match value {
5754 "conversational" => Ok(crate::wire::Provenance::Conversational),
5755 "injected" => Ok(crate::wire::Provenance::Injected),
5756 other => anyhow::bail!("unknown part provenance {other}"),
5757 }
5758}
5759
5760fn file_data_from_blob(variant_data: &[u8], bytes: &[u8]) -> Result<FileData> {
5761 let kind = file_data_kind(variant_data)?;
5762 match kind.as_str() {
5763 "string" => {
5764 let text = std::str::from_utf8(bytes)
5765 .context("file string payload is not UTF-8")?
5766 .to_owned();
5767 Ok(FileData::String(text))
5768 }
5769 "bytes" => Ok(FileData::Bytes(bytes.to_vec())),
5770 "url" => Ok(FileData::Url(
5771 std::str::from_utf8(bytes)
5772 .context("file URL payload is not UTF-8")?
5773 .to_owned(),
5774 )),
5775 other => anyhow::bail!("unknown file data_kind {other}"),
5776 }
5777}
5778
5779fn file_data_kind(variant_data: &[u8]) -> Result<String> {
5780 let value = json_parse::<Value>(variant_data)?;
5781 value
5782 .get("data_kind")
5783 .and_then(Value::as_str)
5784 .map(str::to_owned)
5785 .context("file part variant_data missing data_kind")
5786}
5787
5788fn uint64<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a UInt64Array> {
5789 batch
5790 .column_by_name(name)
5791 .with_context(|| format!("missing column {name}"))?
5792 .as_any()
5793 .downcast_ref::<UInt64Array>()
5794 .with_context(|| format!("column {name} is not UInt64"))
5795}
5796
5797pub(crate) fn string(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<String>> {
5798 let array = batch
5799 .column_by_name(name)
5800 .with_context(|| format!("missing column {name}"))?
5801 .as_any()
5802 .downcast_ref::<StringArray>()
5803 .with_context(|| format!("column {name} is not Utf8"))?;
5804 if array.is_null(row) {
5805 Ok(None)
5806 } else {
5807 Ok(Some(array.value(row).to_owned()))
5808 }
5809}
5810
5811fn json_column(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<Vec<u8>>> {
5812 let column = batch
5816 .column_by_name(name)
5817 .with_context(|| format!("missing column {name}"))?;
5818 if let Some(array) = column.as_any().downcast_ref::<LargeBinaryArray>() {
5819 return if array.is_null(row) {
5820 Ok(None)
5821 } else {
5822 Ok(Some(
5823 lance_arrow::json::decode_json(array.value(row)).into_bytes(),
5824 ))
5825 };
5826 }
5827 if let Some(array) = column.as_any().downcast_ref::<StringArray>() {
5828 return if array.is_null(row) {
5829 Ok(None)
5830 } else {
5831 Ok(Some(array.value(row).as_bytes().to_vec()))
5832 };
5833 }
5834 if let Some(array) = column.as_any().downcast_ref::<LargeStringArray>() {
5835 return if array.is_null(row) {
5836 Ok(None)
5837 } else {
5838 Ok(Some(array.value(row).as_bytes().to_vec()))
5839 };
5840 }
5841 anyhow::bail!("column {name} is not a JSON-compatible array")
5842}
5843
5844fn int32(batch: &RecordBatch, name: &str, row: usize) -> Result<i32> {
5845 let array = batch
5846 .column_by_name(name)
5847 .with_context(|| format!("missing column {name}"))?
5848 .as_any()
5849 .downcast_ref::<Int32Array>()
5850 .with_context(|| format!("column {name} is not Int32"))?;
5851 Ok(array.value(row))
5852}
5853
5854pub(crate) fn float32(batch: &RecordBatch, name: &str, row: usize) -> Result<f32> {
5855 let array = batch
5856 .column_by_name(name)
5857 .with_context(|| format!("missing column {name}"))?
5858 .as_any()
5859 .downcast_ref::<Float32Array>()
5860 .with_context(|| format!("column {name} is not Float32"))?;
5861 Ok(array.value(row))
5862}
5863
5864pub(crate) fn datetime(batch: &RecordBatch, name: &str, row: usize) -> Result<DateTime<Utc>> {
5865 let array = batch
5866 .column_by_name(name)
5867 .with_context(|| format!("missing column {name}"))?
5868 .as_any()
5869 .downcast_ref::<TimestampMicrosecondArray>()
5870 .with_context(|| format!("column {name} is not timestamp_micros"))?;
5871 Utc.timestamp_micros(array.value(row))
5872 .single()
5873 .context("timestamp is out of range")
5874}
5875
5876fn primary_field(name: &str, data_type: DataType, nullable: bool) -> Field {
5877 Field::new(name, data_type, nullable).with_metadata(
5878 [(
5879 "lance-schema:unenforced-primary-key".to_owned(),
5880 "true".to_owned(),
5881 )]
5882 .into(),
5883 )
5884}
5885
5886fn legacy_blob_field(name: &str, nullable: bool) -> Field {
5896 Field::new(name, DataType::LargeBinary, nullable).with_metadata(
5897 [(lance_arrow::BLOB_META_KEY.to_owned(), "true".to_owned())]
5898 .into_iter()
5899 .collect(),
5900 )
5901}
5902
5903fn json_field(name: &str, nullable: bool) -> Field {
5910 lance_arrow::json::json_field(name, nullable)
5911}
5912
5913fn micros(timestamp: DateTime<Utc>) -> i64 {
5914 timestamp.timestamp_micros()
5915}
5916
5917fn json_bytes<T: Serialize>(value: &T) -> Result<Vec<u8>> {
5918 let text = serde_json::to_string(value).context("failed to serialize JSON field")?;
5926 lance_arrow::json::encode_json(&text)
5927 .map_err(|err| anyhow::anyhow!("failed to encode JSON field as JSONB: {err}"))
5928}
5929
5930fn json_parse<T: DeserializeOwned>(value: &[u8]) -> Result<T> {
5931 serde_json::from_slice(value).context("failed to parse JSON field")
5932}
5933
5934fn tool_identity(kind: &PartKind) -> (Option<&str>, Option<&str>, Option<bool>) {
5941 match kind {
5942 PartKind::ToolCall { name, call_id, .. } => (
5943 name.as_deref().map(String::as_str),
5944 call_id.as_deref().map(String::as_str),
5945 None,
5946 ),
5947 PartKind::ToolResult {
5948 name,
5949 call_id,
5950 is_failure,
5951 ..
5952 } => (
5953 name.as_deref().map(String::as_str),
5954 call_id.as_deref().map(String::as_str),
5955 Some(*is_failure),
5956 ),
5957 PartKind::ToolApprovalRequest { tool_call_id, .. } => {
5958 (None, Some(tool_call_id.as_str()), None)
5959 }
5960 PartKind::Text { .. }
5961 | PartKind::Reasoning { .. }
5962 | PartKind::File { .. }
5963 | PartKind::ToolApprovalResponse { .. } => (None, None, None),
5964 }
5965}
5966
5967fn part_variant_json(kind: &PartKind) -> Result<Vec<u8>> {
5968 if let PartKind::File {
5969 media_type,
5970 file_name,
5971 data,
5972 } = kind
5973 {
5974 let data_kind = match data {
5975 FileData::String(_) => "string",
5976 FileData::Bytes(_) => "bytes",
5977 FileData::Url(_) => "url",
5978 };
5979 return json_bytes(&serde_json::json!({
5980 "media_type": media_type,
5981 "file_name": file_name,
5982 "data_kind": data_kind,
5983 }));
5984 }
5985 let value = serde_json::to_value(kind)?;
5986 let mut object = value
5987 .as_object()
5988 .cloned()
5989 .context("part variant did not serialize to an object")?;
5990 object.remove("type");
5991 json_bytes(&object)
5992}
5993
5994fn part_kind_from_json(
5995 type_name: &str,
5996 variant_data: &[u8],
5997 file_data: Option<FileData>,
5998) -> Result<PartKind> {
5999 let mut value = json_parse::<Value>(variant_data)?;
6000 let object = value
6001 .as_object_mut()
6002 .context("part variant data is not an object")?;
6003 object.insert("type".to_owned(), Value::String(type_name.to_owned()));
6004 if let Some(data) = file_data {
6005 object.remove("data_kind");
6006 object.insert("data".to_owned(), serde_json::to_value(data)?);
6007 }
6008 serde_json::from_value(value).context("failed to parse part kind")
6009}
6010
6011#[cfg(test)]
6012mod tests {
6013 #![allow(clippy::expect_used, clippy::unwrap_used)]
6014
6015 use super::*;
6016 use crate::{
6017 adapter::Extracted,
6018 handlers::ingest_events,
6019 wire::{FileData, Message, Part, PartKind, ProviderOptions, Session},
6020 };
6021 use chrono::Utc;
6022 use serde_json::json;
6023 use tempfile::TempDir;
6024
6025 fn synthetic_session(id: &str) -> Session {
6026 Session {
6027 id: id.to_owned(),
6028 parent_session_id: None,
6029 parent_message_id: None,
6030 source_agent: "claude-code".to_owned(),
6031 created_at: Utc::now(),
6032 project: crate::adapter::Extracted::from_test_value("/tmp/pond".to_owned()),
6033 options: ProviderOptions::new(),
6034 }
6035 }
6036
6037 #[derive(Default)]
6040 struct CountingEmbedder {
6041 texts: std::sync::atomic::AtomicUsize,
6042 }
6043 impl crate::embed::Embedder for CountingEmbedder {
6044 fn device(&self) -> &str {
6045 "test"
6046 }
6047 fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
6048 self.texts
6049 .fetch_add(texts.len(), std::sync::atomic::Ordering::SeqCst);
6050 Ok(texts
6051 .iter()
6052 .map(|_| vec![0.0_f32; embedding_dim()])
6053 .collect())
6054 }
6055 }
6056
6057 fn conversational_events(session_id: &str, count: usize) -> Vec<IngestEvent> {
6060 let mut events = vec![IngestEvent::Session(synthetic_session(session_id))];
6061 for index in 0..count {
6062 events.push(IngestEvent::Message(Message::User {
6063 id: format!("msg-{index}"),
6064 session_id: session_id.to_owned(),
6065 timestamp: Utc::now(),
6066 options: ProviderOptions::new(),
6067 }));
6068 events.push(IngestEvent::Part(Part {
6069 session_id: session_id.to_owned(),
6070 id: format!("msg-{index}:0001"),
6071 message_id: format!("msg-{index}"),
6072 ordinal: 0,
6073 provenance: crate::wire::Provenance::Conversational,
6074 options: ProviderOptions::new(),
6075 kind: PartKind::Text {
6076 text: Some(Extracted::from_test_value(format!("body {index}"))),
6077 },
6078 }));
6079 }
6080 events
6081 }
6082
6083 #[tokio::test(flavor = "multi_thread")]
6087 async fn ingest_embeds_inline_in_the_birth_commit() -> anyhow::Result<()> {
6088 let plain = Store::open(&Url::parse("shared-memory://pond-test-inline-plain/")?).await?;
6089 ingest_events(&plain, conversational_events("01HXYINLINE000PLAIN", 5)).await?;
6090 assert!(
6091 !plain.has_embeddings().await?,
6092 "no embedder attached -> every vector null",
6093 );
6094 let plain_version = plain.messages_version().await?;
6095
6096 let backend = Arc::new(CountingEmbedder::default());
6097 let embedder = Arc::new(crate::embed::LazyEmbedder::from_loaded(
6098 backend.clone() as Arc<dyn crate::embed::Embedder>
6099 ));
6100 let store = Store::open(&Url::parse("shared-memory://pond-test-inline-embed/")?)
6101 .await?
6102 .with_embedder(embedder);
6103 ingest_events(&store, conversational_events("01HXYINLINE000EMBED", 5)).await?;
6104 assert!(
6105 store.has_embeddings().await?,
6106 "embedder attached -> vectors filled at ingest",
6107 );
6108 assert_eq!(
6109 backend.texts.load(std::sync::atomic::Ordering::SeqCst),
6110 5,
6111 "every conversational message embedded once",
6112 );
6113 assert_eq!(
6114 store.messages_version().await?,
6115 plain_version,
6116 "inline embed must ride the append, not add a separate commit",
6117 );
6118
6119 ingest_events(&store, conversational_events("01HXYINLINE000EMBED", 5)).await?;
6120 assert_eq!(
6121 backend.texts.load(std::sync::atomic::Ordering::SeqCst),
6122 5,
6123 "an idempotent re-sync embeds no already-present row",
6124 );
6125 Ok(())
6126 }
6127
6128 #[tokio::test(flavor = "multi_thread")]
6134 async fn composite_pk_index_counts_duplicates_by_composite_key() -> anyhow::Result<()> {
6135 async fn duplicates(store: &Store, table: Table) -> anyhow::Result<usize> {
6136 let (keys, rows) = store.composite_pk_index(table).await?;
6137 Ok(rows - keys.len())
6138 }
6139 let store = Store::open(&Url::parse("shared-memory://pond-test-dupcount/")?).await?;
6140 ingest_events(&store, conversational_events("01HXYDUP00000SESS1", 1)).await?;
6141 ingest_events(&store, conversational_events("01HXYDUP00000SESS2", 1)).await?;
6142 assert_eq!(
6143 duplicates(&store, Table::Messages).await?,
6144 0,
6145 "the same message id in two sessions is not a duplicate (composite PK)",
6146 );
6147 assert_eq!(duplicates(&store, Table::Sessions).await?, 0);
6148 assert_eq!(duplicates(&store, Table::Parts).await?, 0);
6149
6150 let message = Message::User {
6153 id: "dup-msg".to_owned(),
6154 session_id: "01HXYDUP00000SESS1".to_owned(),
6155 timestamp: Utc::now(),
6156 options: ProviderOptions::new(),
6157 };
6158 let row = MessageBatchRow {
6159 message: &message,
6160 source_agent: "claude-code",
6161 project: "/tmp",
6162 search_text: None,
6163 };
6164 let batches = messages_batches(&[row], &[None])?;
6165 store
6166 .handle
6167 .append_batches(Table::Messages, batches.clone())
6168 .await?;
6169 store
6170 .handle
6171 .append_batches(Table::Messages, batches)
6172 .await?;
6173 assert_eq!(
6174 duplicates(&store, Table::Messages).await?,
6175 1,
6176 "the doubled (session_id, id) row is exactly one duplicate",
6177 );
6178 Ok(())
6179 }
6180
6181 #[test]
6182 fn search_text_excludes_injected_parts() {
6183 use crate::wire::Provenance;
6184 let message = Message::User {
6185 id: "m1".to_owned(),
6186 session_id: "s1".to_owned(),
6187 timestamp: Utc::now(),
6188 options: ProviderOptions::new(),
6189 };
6190 let text_part = |id: &str, text: &str, provenance: Provenance| Part {
6191 session_id: "s1".to_owned(),
6192 id: id.to_owned(),
6193 message_id: "m1".to_owned(),
6194 ordinal: 0,
6195 provenance,
6196 options: ProviderOptions::new(),
6197 kind: PartKind::Text {
6198 text: Some(Extracted::from_test_value(text.to_owned())),
6199 },
6200 };
6201
6202 let conversational = search_text(
6205 &message,
6206 &[text_part(
6207 "p1",
6208 "real human prompt",
6209 Provenance::Conversational,
6210 )],
6211 );
6212 assert_eq!(conversational.as_deref(), Some("real human prompt"));
6213
6214 let injected = search_text(
6215 &message,
6216 &[text_part(
6217 "p2",
6218 "<task-notification>...</task-notification>",
6219 Provenance::Injected,
6220 )],
6221 );
6222 assert!(
6223 injected.is_none(),
6224 "a message whose only part is injected has null search_text"
6225 );
6226 }
6227
6228 #[test]
6229 fn parts_chunk_materializes_tool_identity_columns() -> anyhow::Result<()> {
6230 let part = |ordinal: i32, kind: PartKind| Part {
6231 session_id: "s1".to_owned(),
6232 id: format!("p{ordinal}"),
6233 message_id: "m1".to_owned(),
6234 ordinal,
6235 provenance: crate::wire::Provenance::Conversational,
6236 options: ProviderOptions::new(),
6237 kind,
6238 };
6239 let parts = vec![
6240 part(
6241 0,
6242 PartKind::ToolCall {
6243 call_id: Some(Extracted::from_test_value("c1".to_owned())),
6244 name: Some(Extracted::from_test_value("Bash".to_owned())),
6245 params: serde_json::json!({}),
6246 provider_executed: false,
6247 },
6248 ),
6249 part(
6250 1,
6251 PartKind::ToolResult {
6252 call_id: Some(Extracted::from_test_value("c1".to_owned())),
6253 name: Some(Extracted::from_test_value("Bash".to_owned())),
6254 is_failure: true,
6255 result: serde_json::json!({}),
6256 },
6257 ),
6258 part(
6260 2,
6261 PartKind::ToolCall {
6262 call_id: None,
6263 name: None,
6264 params: serde_json::json!({}),
6265 provider_executed: false,
6266 },
6267 ),
6268 part(
6269 3,
6270 PartKind::ToolApprovalRequest {
6271 approval_id: "a1".to_owned(),
6272 tool_call_id: "c1".to_owned(),
6273 },
6274 ),
6275 part(
6276 4,
6277 PartKind::Text {
6278 text: Some(Extracted::from_test_value("hi".to_owned())),
6279 },
6280 ),
6281 ];
6282 let payloads = vec![vec![1u8]; parts.len()];
6283 let batch = parts_chunk(&parts, &payloads, &payloads)?;
6284
6285 let tool_name = |row| string(&batch, "tool_name", row).unwrap();
6286 let call_id = |row| string(&batch, "call_id", row).unwrap();
6287 let is_failure = batch
6288 .column_by_name("is_failure")
6289 .expect("is_failure column present")
6290 .as_any()
6291 .downcast_ref::<BooleanArray>()
6292 .expect("is_failure is boolean");
6293
6294 assert_eq!(tool_name(0).as_deref(), Some("Bash"));
6295 assert_eq!(call_id(0).as_deref(), Some("c1"));
6296 assert!(is_failure.is_null(0), "tool_call has no failure flag");
6297
6298 assert_eq!(tool_name(1).as_deref(), Some("Bash"));
6299 assert_eq!(call_id(1).as_deref(), Some("c1"));
6300 assert!(is_failure.value(1), "tool_result failure flag materialized");
6301
6302 assert_eq!(tool_name(2), None, "absent name stays NULL");
6303 assert_eq!(call_id(2), None, "absent call_id stays NULL");
6304
6305 assert_eq!(tool_name(3), None);
6306 assert_eq!(
6307 call_id(3).as_deref(),
6308 Some("c1"),
6309 "approval request surfaces its tool_call_id as call_id",
6310 );
6311
6312 assert_eq!(tool_name(4), None);
6313 assert_eq!(call_id(4), None);
6314 assert!(is_failure.is_null(4));
6315 Ok(())
6316 }
6317
6318 #[test]
6319 fn chunk_ranges_splits_on_byte_budget() {
6320 assert!(chunk_ranges(&[]).is_empty());
6321 assert_eq!(chunk_ranges(&[10, 10, 10]), vec![0..3]);
6322
6323 let two_thirds = COLUMN_BYTE_BUDGET * 2 / 3;
6324 assert_eq!(
6325 chunk_ranges(&[two_thirds, two_thirds, two_thirds]),
6326 vec![0..1, 1..2, 2..3],
6327 );
6328
6329 assert_eq!(
6331 chunk_ranges(&[10, COLUMN_BYTE_BUDGET + 1, 10]),
6332 vec![0..1, 1..2, 2..3],
6333 );
6334 }
6335
6336 #[tokio::test]
6337 async fn ordering_violation_drops_only_the_offending_event() -> anyhow::Result<()> {
6338 let temp = TempDir::new()?;
6343 let store = Store::open_local(temp.path()).await?;
6344 let session = synthetic_session("ordering");
6345 let orphan_part = Part {
6346 session_id: session.id.clone(),
6347 id: "orphan-part".to_owned(),
6348 message_id: "missing-message".to_owned(),
6349 ordinal: 0,
6350 provenance: crate::wire::Provenance::Conversational,
6351 options: ProviderOptions::new(),
6352 kind: PartKind::Text {
6353 text: Some(Extracted::from_test_value("orphan".to_owned())),
6354 },
6355 };
6356 let valid_message = Message::User {
6357 id: "valid-message".to_owned(),
6358 session_id: session.id.clone(),
6359 timestamp: Utc::now(),
6360 options: ProviderOptions::new(),
6361 };
6362 let valid_part = Part {
6363 session_id: session.id.clone(),
6364 id: "valid-part".to_owned(),
6365 message_id: valid_message.id().to_owned(),
6366 ordinal: 0,
6367 provenance: crate::wire::Provenance::Conversational,
6368 options: ProviderOptions::new(),
6369 kind: PartKind::Text {
6370 text: Some(Extracted::from_test_value("kept".to_owned())),
6371 },
6372 };
6373
6374 let mut validator = IngestValidator::default();
6375 validator
6376 .push(&store, 0, IngestEvent::Session(session.clone()))
6377 .await?;
6378 let part_outcomes = validator
6379 .push(&store, 1, IngestEvent::Part(orphan_part))
6380 .await?;
6381 assert_eq!(part_outcomes.len(), 1);
6382 assert_eq!(part_outcomes[0].kind, "part");
6383 assert_eq!(part_outcomes[0].status, OutcomeStatus::Error);
6384 assert!(
6385 part_outcomes[0]
6386 .error
6387 .as_ref()
6388 .map(|e| e.message.contains("part event appeared before a message"))
6389 .unwrap_or(false),
6390 "error message must explain the ordering violation: {part_outcomes:?}"
6391 );
6392 validator
6393 .push(&store, 2, IngestEvent::Message(valid_message))
6394 .await?;
6395 validator
6396 .push(&store, 3, IngestEvent::Part(valid_part))
6397 .await?;
6398 validator.finish(&store).await?;
6399
6400 let (sessions, messages, parts) = store.row_counts().await?;
6401 assert_eq!(sessions, 1, "session committed despite the orphan part");
6402 assert_eq!(messages, 1, "valid message committed");
6403 assert_eq!(parts, 1, "valid part committed; the orphan was dropped");
6404
6405 Ok(())
6406 }
6407
6408 #[tokio::test]
6409 async fn resident_meta_map_hydration_matches_take_rows_fallback() -> anyhow::Result<()> {
6410 let temp = TempDir::new()?;
6414 let store = Store::open_local(temp.path()).await?;
6415 let session = synthetic_session("hydration-parity");
6416
6417 let messages = [
6418 (
6419 "m1",
6420 "the auth refactor landed cleanly",
6421 1_700_000_000_123_456_i64,
6422 ),
6423 (
6424 "m2",
6425 "balance handler now retries on rpc timeout",
6426 1_700_000_050_654_321,
6427 ),
6428 ];
6429 let mut validator = IngestValidator::default();
6430 validator
6431 .push(&store, 0, IngestEvent::Session(session.clone()))
6432 .await?;
6433 let mut seq = 1;
6434 for (mid, text, micros) in messages {
6435 let message = Message::User {
6436 id: mid.to_owned(),
6437 session_id: session.id.clone(),
6438 timestamp: DateTime::from_timestamp_micros(micros).unwrap(),
6439 options: ProviderOptions::new(),
6440 };
6441 validator
6442 .push(&store, seq, IngestEvent::Message(message))
6443 .await?;
6444 seq += 1;
6445 let part = Part {
6446 session_id: session.id.clone(),
6447 id: format!("{mid}-p0"),
6448 message_id: mid.to_owned(),
6449 ordinal: 0,
6450 provenance: crate::wire::Provenance::Conversational,
6451 options: ProviderOptions::new(),
6452 kind: PartKind::Text {
6453 text: Some(Extracted::from_test_value(text.to_owned())),
6454 },
6455 };
6456 validator.push(&store, seq, IngestEvent::Part(part)).await?;
6457 seq += 1;
6458 }
6459 validator.finish(&store).await?;
6460
6461 let rowids: Vec<u64> = store
6462 .collect_row_metas()
6463 .await?
6464 .into_iter()
6465 .map(|entry| entry.row_id)
6466 .collect();
6467 assert_eq!(rowids.len(), 2);
6468
6469 let sort_by_id = |mut metas: Vec<MessageMeta>| {
6470 metas.sort_by(|left, right| left.message_id.cmp(&right.message_id));
6471 metas
6472 };
6473
6474 let fallback = sort_by_id(store.message_metas_by_rowids(&rowids).await?);
6475
6476 store.ensure_rowmap(&temp.path().join("cache")).await?;
6479 let resident = sort_by_id(store.message_metas_by_rowids(&rowids).await?);
6480
6481 assert_eq!(
6482 resident, fallback,
6483 "resident-map hydration must match the take_rows fallback"
6484 );
6485 assert_eq!(
6486 resident[0].timestamp.timestamp_micros(),
6487 1_700_000_000_123_456
6488 );
6489 Ok(())
6490 }
6491
6492 #[tokio::test]
6493 async fn get_paths_serve_from_resident_map_and_fall_back_on_staleness() -> anyhow::Result<()> {
6494 let temp = TempDir::new()?;
6498 let store = Store::open_local(temp.path()).await?;
6499 let session = synthetic_session("map-served-get");
6500 let cache = temp.path().join("cache");
6501
6502 let seed = |mid: &str, text: &str, micros: i64| {
6503 let message = Message::User {
6504 id: mid.to_owned(),
6505 session_id: session.id.clone(),
6506 timestamp: DateTime::from_timestamp_micros(micros).unwrap(),
6507 options: ProviderOptions::new(),
6508 };
6509 let part = Part {
6510 session_id: session.id.clone(),
6511 id: format!("{mid}-p0"),
6512 message_id: mid.to_owned(),
6513 ordinal: 0,
6514 provenance: crate::wire::Provenance::Conversational,
6515 options: ProviderOptions::new(),
6516 kind: PartKind::Text {
6517 text: Some(Extracted::from_test_value(text.to_owned())),
6518 },
6519 };
6520 (message, part)
6521 };
6522 let mut validator = IngestValidator::default();
6523 validator
6524 .push(&store, 0, IngestEvent::Session(session.clone()))
6525 .await?;
6526 let mut seq = 1;
6527 for (mid, text, micros) in [
6528 ("m1", "first page text", 1_700_000_000_000_000_i64),
6529 ("m2", "second page text", 1_700_000_001_000_000),
6530 ] {
6531 let (message, part) = seed(mid, text, micros);
6532 validator
6533 .push(&store, seq, IngestEvent::Message(message))
6534 .await?;
6535 validator
6536 .push(&store, seq + 1, IngestEvent::Part(part))
6537 .await?;
6538 seq += 2;
6539 }
6540 validator.finish(&store).await?;
6541
6542 let params = SessionViewParams {
6543 at_message_id: None,
6544 after_message_id: None,
6545 before_message_id: None,
6546 limit: 50,
6547 budget_bytes: 1_000_000,
6548 session_from: SessionFrom::Start,
6549 };
6550 let page_rows = |lookup: GetLookup<SessionPage>| match lookup {
6551 GetLookup::Found(page) => page
6552 .messages
6553 .iter()
6554 .map(|message| (message.id.clone(), message.text.clone()))
6555 .collect::<Vec<_>>(),
6556 other => panic!("expected a page, got {other:?}"),
6557 };
6558 let view_params = MessageViewParams {
6559 context_before: 5,
6560 context_after: 5,
6561 budget_bytes: 1_000_000,
6562 };
6563 let view_rows = |lookup: GetLookup<MessagePage>| match lookup {
6564 GetLookup::Found(page) => (
6565 page.target.text.clone(),
6566 page.siblings
6567 .iter()
6568 .map(|sibling| sibling.id.clone())
6569 .collect::<Vec<_>>(),
6570 ),
6571 other => panic!("expected a message page, got {other:?}"),
6572 };
6573
6574 let scanned = page_rows(store.session_view(&session.id, params.clone()).await?);
6575 let scanned_view = view_rows(store.message_view("m2", view_params.clone()).await?);
6576 store.ensure_rowmap(&cache).await?;
6577 let resident = page_rows(store.session_view(&session.id, params.clone()).await?);
6578 assert_eq!(resident, scanned, "map-served page must match the scan");
6579 let resident_view = view_rows(store.message_view("m2", view_params.clone()).await?);
6580 assert_eq!(
6581 resident_view, scanned_view,
6582 "map-served message window must match the scan"
6583 );
6584 assert_eq!(
6585 store.session_id_for_message("m2").await?.as_deref(),
6586 Some(session.id.as_str()),
6587 "message id resolves from the resident map"
6588 );
6589
6590 let (message, part) = seed("m3", "post-map text", 1_700_000_002_000_000);
6594 ingest_events(
6595 &store,
6596 vec![
6597 IngestEvent::Session(session.clone()),
6598 IngestEvent::Message(message),
6599 IngestEvent::Part(part),
6600 ],
6601 )
6602 .await?;
6603
6604 let stale_map_page = page_rows(store.session_view(&session.id, params).await?);
6605 assert_eq!(stale_map_page.len(), 3, "stale map must not hide m3");
6606 assert_eq!(stale_map_page[2].0, "m3");
6607 assert_eq!(
6608 store.session_id_for_message("m3").await?.as_deref(),
6609 Some(session.id.as_str()),
6610 "post-map message id resolves via the scan fallback"
6611 );
6612
6613 ingest_events(
6617 &store,
6618 vec![
6619 IngestEvent::Session(session.clone()),
6620 IngestEvent::Message(Message::System {
6621 id: "sys-1".to_owned(),
6622 session_id: session.id.clone(),
6623 timestamp: DateTime::from_timestamp_micros(1_700_000_003_000_000).unwrap(),
6624 content: Some(Extracted::from_test_value("system says hello".to_owned())),
6625 options: ProviderOptions::new(),
6626 }),
6627 ],
6628 )
6629 .await?;
6630 store.ensure_rowmap(&cache).await?;
6631
6632 let target_page = match store.message_view("m2", view_params.clone()).await? {
6633 GetLookup::Found(page) => page,
6634 other => panic!("expected a message page, got {other:?}"),
6635 };
6636 assert_eq!(target_page.target.text.as_deref(), Some("second page text"));
6637 let sibling_ids: Vec<&str> = target_page
6638 .siblings
6639 .iter()
6640 .map(|sibling| sibling.id.as_str())
6641 .collect();
6642 assert_eq!(
6643 sibling_ids,
6644 vec!["m1", "m3"],
6645 "map-served siblings are the conversational neighbors"
6646 );
6647
6648 let system_page = match store.message_view("sys-1", view_params).await? {
6649 GetLookup::Found(page) => page,
6650 other => panic!("expected the system message page, got {other:?}"),
6651 };
6652 assert_eq!(
6653 system_page.target.content.as_deref(),
6654 Some("system says hello"),
6655 "system target must fall back to the scan and keep its content"
6656 );
6657 Ok(())
6658 }
6659
6660 #[tokio::test]
6661 async fn initialized_flips_only_after_first_ingest() -> anyhow::Result<()> {
6662 let temp = TempDir::new()?;
6667 let store = Store::open_local(temp.path()).await?;
6668 assert!(
6669 !store.initialized().await?,
6670 "fresh store has no parts table"
6671 );
6672
6673 let session = synthetic_session("initialized-probe");
6674 let message = Message::User {
6675 id: "message-1".to_owned(),
6676 session_id: session.id.clone(),
6677 timestamp: Utc::now(),
6678 options: ProviderOptions::new(),
6679 };
6680 let part = Part {
6681 session_id: session.id.clone(),
6682 id: "part-1".to_owned(),
6683 message_id: message.id().to_owned(),
6684 ordinal: 0,
6685 provenance: crate::wire::Provenance::Conversational,
6686 options: ProviderOptions::new(),
6687 kind: PartKind::Text {
6688 text: Some(Extracted::from_test_value("hello".to_owned())),
6689 },
6690 };
6691 let mut validator = IngestValidator::default();
6692 validator
6693 .push(&store, 0, IngestEvent::Session(session))
6694 .await?;
6695 validator
6696 .push(&store, 1, IngestEvent::Message(message))
6697 .await?;
6698 validator.push(&store, 2, IngestEvent::Part(part)).await?;
6699 validator.finish(&store).await?;
6700
6701 assert!(store.initialized().await?, "ingest creates the parts table");
6702 Ok(())
6703 }
6704
6705 #[tokio::test]
6706 async fn summary_parts_label_a_file_without_reading_its_blob() -> anyhow::Result<()> {
6707 let temp = TempDir::new()?;
6712 let store = Store::open_local(temp.path()).await?;
6713 let session = synthetic_session("file-summary");
6714 let message = Message::User {
6715 id: "m1".to_owned(),
6716 session_id: session.id.clone(),
6717 timestamp: Utc::now(),
6718 options: ProviderOptions::new(),
6719 };
6720 let blob = "file contents the summary must never read";
6721 let part = Part {
6722 session_id: session.id.clone(),
6723 id: "m1-p0".to_owned(),
6724 message_id: "m1".to_owned(),
6725 ordinal: 0,
6726 provenance: crate::wire::Provenance::Conversational,
6727 options: ProviderOptions::new(),
6728 kind: PartKind::File {
6729 media_type: Some("text/plain".to_owned()),
6730 file_name: Some("notes.txt".to_owned()),
6731 data: FileData::String(blob.to_owned()),
6732 },
6733 };
6734 let mut validator = IngestValidator::default();
6735 validator
6736 .push(&store, 0, IngestEvent::Session(session.clone()))
6737 .await?;
6738 validator
6739 .push(&store, 1, IngestEvent::Message(message))
6740 .await?;
6741 validator.push(&store, 2, IngestEvent::Part(part)).await?;
6742 validator.finish(&store).await?;
6743
6744 let key = (session.id.clone(), "m1".to_owned());
6745 let ids = ["m1".to_owned()];
6746
6747 let summarized = store.summary_parts_for_messages(&session.id, &ids).await?;
6748 let summary_part = &summarized.get(&key).expect("file part summarized")[0];
6749 let summary = crate::wire::PartSummary::for_kind(&summary_part.kind)
6750 .expect("a file part yields a summary");
6751 assert_eq!(summary.kind, "file");
6752 assert_eq!(summary.label.as_deref(), Some("notes.txt"));
6753 match &summary_part.kind {
6754 PartKind::File { data, .. } => assert!(
6755 matches!(data, FileData::Bytes(bytes) if bytes.is_empty()),
6756 "summary must carry the empty placeholder, not the file blob",
6757 ),
6758 other => panic!("expected a file part, got {other:?}"),
6759 }
6760
6761 let full = store.parts_for_messages(&session.id, &ids).await?;
6762 match &full.get(&key).expect("file part")[0].kind {
6763 PartKind::File {
6764 data: FileData::String(bytes),
6765 ..
6766 } => assert_eq!(bytes, blob, "full path must still hydrate the real blob"),
6767 other => panic!("expected a string-backed file part, got {other:?}"),
6768 }
6769 Ok(())
6770 }
6771
6772 #[tokio::test]
6773 async fn duplicate_message_id_drops_the_second_keeps_the_first() -> anyhow::Result<()> {
6774 let temp = TempDir::new()?;
6778 let store = Store::open_local(temp.path()).await?;
6779 let session = synthetic_session("duplicate-message");
6780 let first = Message::User {
6781 id: "message-1".to_owned(),
6782 session_id: session.id.clone(),
6783 timestamp: Utc::now(),
6784 options: ProviderOptions::new(),
6785 };
6786 let second = Message::Assistant {
6787 id: "message-1".to_owned(),
6788 session_id: session.id.clone(),
6789 timestamp: Utc::now(),
6790 options: ProviderOptions::new(),
6791 };
6792
6793 let mut validator = IngestValidator::default();
6794 validator
6795 .push(&store, 0, IngestEvent::Session(session.clone()))
6796 .await?;
6797 validator
6798 .push(&store, 1, IngestEvent::Message(first))
6799 .await?;
6800 let dup_outcomes = validator
6801 .push(&store, 2, IngestEvent::Message(second))
6802 .await?;
6803 assert_eq!(dup_outcomes.len(), 1);
6804 assert_eq!(dup_outcomes[0].status, OutcomeStatus::Error);
6805 assert!(
6806 dup_outcomes[0]
6807 .error
6808 .as_ref()
6809 .map(|e| e.message.contains("duplicate message id message-1"))
6810 .unwrap_or(false),
6811 "duplicate-id rejection must name the offending id: {dup_outcomes:?}"
6812 );
6813
6814 validator.finish(&store).await?;
6815 let (sessions, messages, _) = store.row_counts().await?;
6816 assert_eq!(sessions, 1, "session committed");
6817 assert_eq!(messages, 1, "only the first message committed");
6818
6819 Ok(())
6820 }
6821
6822 #[tokio::test]
6823 async fn ingest_stamps_host_provenance_on_messages_and_strips_spoofed_pond_key()
6824 -> anyhow::Result<()> {
6825 let temp = TempDir::new()?;
6829 let store = Store::open_local(temp.path()).await?;
6830 let session = synthetic_session("host-provenance");
6831 let mut spoofed = ProviderOptions::new();
6832 spoofed.insert("pond".to_owned(), json!({"ingest": {"host": "spoofed"}}));
6833 let message = Message::User {
6834 id: "message-1".to_owned(),
6835 session_id: session.id.clone(),
6836 timestamp: Utc::now(),
6837 options: spoofed,
6838 };
6839 let part = Part {
6840 session_id: session.id.clone(),
6841 id: "part-1".to_owned(),
6842 message_id: "message-1".to_owned(),
6843 ordinal: 0,
6844 provenance: crate::wire::Provenance::Conversational,
6845 options: ProviderOptions::new(),
6846 kind: PartKind::Text {
6847 text: Some(Extracted::from_test_value("hello".to_owned())),
6848 },
6849 };
6850
6851 let mut validator = IngestValidator::default();
6852 validator
6853 .push(&store, 0, IngestEvent::Session(session.clone()))
6854 .await?;
6855 validator
6856 .push(&store, 1, IngestEvent::Message(message))
6857 .await?;
6858 validator.push(&store, 2, IngestEvent::Part(part)).await?;
6859 validator.finish(&store).await?;
6860
6861 let stored = store
6862 .get_session(&session.id)
6863 .await?
6864 .expect("ingested session is readable");
6865 assert!(
6866 !stored.session.options.contains_key("pond"),
6867 "session rows are not stamped (attribution derives from messages)"
6868 );
6869 let stored_message = &stored.messages[0].message;
6870 match ingest_host_stamp() {
6871 Some(stamp) => {
6872 assert_eq!(
6873 stored_message.options().get("pond"),
6874 Some(stamp),
6875 "stored message carries the real stamp, never the spoof"
6876 );
6877 let host = stamp
6878 .pointer("/ingest/host")
6879 .and_then(Value::as_object)
6880 .expect("stamp shape is {ingest: {host: {..}}}");
6881 assert!(!host.is_empty(), "an all-empty stamp must be None instead");
6882 assert!(
6883 host.values()
6884 .all(|v| v.as_str().is_some_and(|s| !s.is_empty())),
6885 "stamp fields are omitted when unavailable, never empty: {host:?}"
6886 );
6887 }
6888 None => assert!(
6889 stored_message.options().get("pond").is_none(),
6890 "with no resolvable stamp the spoofed key is still stripped"
6891 ),
6892 }
6893 assert!(
6894 !stored.messages[0].parts[0].options.contains_key("pond"),
6895 "part rows are not stamped (covered by their message's stamp)"
6896 );
6897
6898 Ok(())
6899 }
6900
6901 #[tokio::test(flavor = "multi_thread")]
6909 async fn optimize_indices_compacts_parts_with_blob_column() -> anyhow::Result<()> {
6910 use crate::wire::{FileData, PartKind, Provenance};
6911 let temp = TempDir::new()?;
6912 let store = Store::open_local(temp.path()).await?;
6913
6914 let session = synthetic_session("compact-blob");
6915 store
6916 .upsert_sessions(std::slice::from_ref(&session))
6917 .await?;
6918
6919 let make_part = |idx: usize, kind: PartKind| Part {
6920 session_id: session.id.clone(),
6921 message_id: format!("msg-{idx}"),
6922 id: format!("part-{idx}"),
6923 ordinal: 0,
6924 provenance: Provenance::Conversational,
6925 options: ProviderOptions::new(),
6926 kind,
6927 };
6928
6929 let batch_a = vec![
6930 make_part(
6931 0,
6932 PartKind::File {
6933 media_type: Some("text/plain".to_owned()),
6934 file_name: Some("a.txt".to_owned()),
6935 data: FileData::Bytes(b"alpha".to_vec()),
6936 },
6937 ),
6938 make_part(
6939 1,
6940 PartKind::File {
6941 media_type: Some("text/plain".to_owned()),
6942 file_name: Some("b.txt".to_owned()),
6943 data: FileData::String("beta".to_owned()),
6944 },
6945 ),
6946 ];
6947 store.upsert_parts(&batch_a).await?;
6948
6949 let batch_b = vec![
6950 make_part(
6951 2,
6952 PartKind::File {
6953 media_type: Some("application/octet-stream".to_owned()),
6954 file_name: None,
6955 data: FileData::Url("https://example.com/file".to_owned()),
6956 },
6957 ),
6958 make_part(
6959 3,
6960 PartKind::File {
6961 media_type: Some("image/png".to_owned()),
6962 file_name: Some("c.png".to_owned()),
6963 data: FileData::Bytes(vec![0x89, 0x50, 0x4e, 0x47]),
6964 },
6965 ),
6966 ];
6967 store.upsert_parts(&batch_b).await?;
6968
6969 store
6970 .optimize_indices(None, &MaintenancePolicy::always_compact())
6971 .await?
6972 .into_result()?;
6973
6974 Ok(())
6975 }
6976
6977 #[tokio::test]
6978 async fn file_part_blob_v2_round_trips_through_get() -> anyhow::Result<()> {
6979 let temp = TempDir::new()?;
6980 let store = Store::open_local(temp.path()).await?;
6981 let session = synthetic_session("blob");
6982 let message = Message::User {
6983 id: "message-1".to_owned(),
6984 session_id: session.id.clone(),
6985 timestamp: Utc::now(),
6986 options: ProviderOptions::new(),
6987 };
6988 let part = Part {
6989 session_id: session.id.clone(),
6990 id: "part-1".to_owned(),
6991 message_id: message.id().to_owned(),
6992 ordinal: 0,
6993 provenance: crate::wire::Provenance::Conversational,
6994 options: ProviderOptions::new(),
6995 kind: PartKind::File {
6996 media_type: Some("text/plain".to_owned()),
6997 file_name: Some("payload.txt".to_owned()),
6998 data: FileData::Bytes(b"pond".to_vec()),
6999 },
7000 };
7001
7002 let mut validator = IngestValidator::default();
7003 validator
7004 .push(&store, 0, IngestEvent::Session(session.clone()))
7005 .await?;
7006 validator
7007 .push(&store, 1, IngestEvent::Message(message.clone()))
7008 .await?;
7009 validator
7010 .push(&store, 2, IngestEvent::Part(part.clone()))
7011 .await?;
7012 validator.finish(&store).await?;
7013
7014 let stored = store
7015 .get_session(&session.id)
7016 .await?
7017 .expect("session should exist");
7018 let stored_part = &stored.messages[0].parts[0];
7019 assert_eq!(stored_part, &part);
7020
7021 Ok(())
7022 }
7023
7024 fn base_session() -> Session {
7035 Session {
7036 id: "01HXY00000000001".to_owned(),
7037 parent_session_id: None,
7038 parent_message_id: None,
7039 source_agent: "claude-code".to_owned(),
7040 created_at: Utc::now(),
7041 project: crate::adapter::Extracted::from_test_value("/home/me/proj".to_owned()),
7042 options: ProviderOptions::new(),
7043 }
7044 }
7045
7046 fn count_status(outcomes: &[RowOutcome], target: OutcomeStatus) -> usize {
7047 outcomes
7048 .iter()
7049 .filter(|outcome| outcome.status == target)
7050 .count()
7051 }
7052
7053 #[tokio::test(flavor = "multi_thread")]
7054 async fn re_ingesting_a_session_with_unchanged_immutable_fields_is_idempotent()
7055 -> anyhow::Result<()> {
7056 let temp = TempDir::new()?;
7057 let store = Store::open_local(temp.path()).await?;
7058
7059 let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
7060 assert_eq!(count_status(&first, OutcomeStatus::Inserted), 1);
7061
7062 let mut again = base_session();
7063 again.options.insert("title".to_owned(), json!("renamed"));
7064 let second = ingest_events(&store, vec![IngestEvent::Session(again)]).await?;
7065 assert_eq!(
7066 count_status(&second, OutcomeStatus::Error),
7067 0,
7068 "options is mutable; the re-ingest must not surface an error: {second:?}",
7069 );
7070 assert_eq!(
7071 count_status(&second, OutcomeStatus::Matched),
7072 1,
7073 "unchanged immutable fields must match-insert via merge_insert",
7074 );
7075
7076 Ok(())
7077 }
7078
7079 #[tokio::test(flavor = "multi_thread")]
7080 async fn re_ingesting_with_changed_source_agent_is_rejected() -> anyhow::Result<()> {
7081 let temp = TempDir::new()?;
7082 let store = Store::open_local(temp.path()).await?;
7083
7084 let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
7085 assert_eq!(count_status(&first, OutcomeStatus::Error), 0);
7086
7087 let mut tampered = base_session();
7088 tampered.source_agent = "codex-cli".to_owned();
7089 let second = ingest_events(&store, vec![IngestEvent::Session(tampered)]).await?;
7090 assert_eq!(count_status(&second, OutcomeStatus::Error), 1);
7091 let err_row = second
7092 .iter()
7093 .find(|outcome| outcome.status == OutcomeStatus::Error)
7094 .expect("error outcome present");
7095 let err = err_row.error.as_ref().expect("error body present");
7096 assert_eq!(err.field, Some("source_agent"));
7097 assert_eq!(err.reason, Some("immutable"));
7098
7099 let stored = store
7101 .get_session(&base_session().id)
7102 .await?
7103 .expect("session row survives the rejected re-ingest");
7104 assert_eq!(stored.session.source_agent, "claude-code");
7105
7106 Ok(())
7107 }
7108
7109 #[tokio::test(flavor = "multi_thread")]
7110 async fn re_ingesting_with_changed_project_is_rejected() -> anyhow::Result<()> {
7111 let temp = TempDir::new()?;
7112 let store = Store::open_local(temp.path()).await?;
7113
7114 let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
7115 assert_eq!(count_status(&first, OutcomeStatus::Error), 0);
7116
7117 let mut tampered = base_session();
7118 tampered.project = crate::adapter::Extracted::from_test_value("/somewhere/else".to_owned());
7119 let second = ingest_events(&store, vec![IngestEvent::Session(tampered)]).await?;
7120 let err_row = second
7121 .iter()
7122 .find(|outcome| outcome.status == OutcomeStatus::Error)
7123 .expect("project change must surface an error outcome");
7124 assert_eq!(err_row.error.as_ref().unwrap().field, Some("project"));
7125
7126 let stored = store
7127 .get_session(&base_session().id)
7128 .await?
7129 .expect("session row survives");
7130 assert_eq!(
7131 stored.session.project.as_str(),
7132 "/home/me/proj",
7133 "stored project must remain the original",
7134 );
7135
7136 Ok(())
7137 }
7138
7139 #[tokio::test(flavor = "multi_thread")]
7140 async fn batched_flush_attributes_new_messages_on_existing_session() -> anyhow::Result<()> {
7141 use crate::wire::Provenance;
7149 let temp = TempDir::new()?;
7150 let store = Store::open_local(temp.path()).await?;
7151 let session = base_session();
7152
7153 let text_part = |part_id: &str, message_id: &str, body: &str| Part {
7154 session_id: session.id.clone(),
7155 id: part_id.to_owned(),
7156 message_id: message_id.to_owned(),
7157 ordinal: 0,
7158 provenance: Provenance::Conversational,
7159 options: ProviderOptions::new(),
7160 kind: PartKind::Text {
7161 text: Some(Extracted::from_test_value(body.to_owned())),
7162 },
7163 };
7164 let user_message = |id: &str| Message::User {
7165 id: id.to_owned(),
7166 session_id: session.id.clone(),
7167 timestamp: Utc::now(),
7168 options: ProviderOptions::new(),
7169 };
7170
7171 let mut validator = IngestValidator::default();
7173 validator
7174 .push(&store, 0, IngestEvent::Session(session.clone()))
7175 .await?;
7176 validator
7177 .push(&store, 1, IngestEvent::Message(user_message("m1")))
7178 .await?;
7179 validator
7180 .push(&store, 2, IngestEvent::Part(text_part("p1", "m1", "alpha")))
7181 .await?;
7182 validator
7183 .push(&store, 3, IngestEvent::Message(user_message("m2")))
7184 .await?;
7185 validator
7186 .push(&store, 4, IngestEvent::Part(text_part("p2", "m2", "beta")))
7187 .await?;
7188 let (_first_outcomes, first_counts) = validator.finish(&store).await?;
7189 assert_eq!(first_counts.sessions_inserted, 1);
7190 assert_eq!(first_counts.messages_inserted_total, 2);
7191 assert_eq!(first_counts.messages_inserted_searchable, 2);
7192
7193 let mut validator = IngestValidator::default();
7195 validator
7196 .push(&store, 0, IngestEvent::Session(session.clone()))
7197 .await?;
7198 for (idx, mid) in ["m3", "m4", "m5"].iter().enumerate() {
7199 let pid = format!("p{}", idx + 3);
7200 validator
7201 .push(&store, idx * 2 + 1, IngestEvent::Message(user_message(mid)))
7202 .await?;
7203 validator
7204 .push(
7205 &store,
7206 idx * 2 + 2,
7207 IngestEvent::Part(text_part(&pid, mid, "gamma")),
7208 )
7209 .await?;
7210 }
7211 let (second_outcomes, second_counts) = validator.finish(&store).await?;
7212
7213 assert_eq!(
7214 second_counts.sessions_inserted, 0,
7215 "existing session row must report as Matched, not Inserted",
7216 );
7217 assert_eq!(second_counts.sessions_matched, 1);
7218 assert_eq!(
7219 second_counts.messages_inserted_total, 3,
7220 "the three NEW messages must register as Inserted in BatchCounts",
7221 );
7222 assert_eq!(
7223 second_counts.messages_inserted_searchable, 3,
7224 "all three new messages carry conversational text -> searchable",
7225 );
7226 assert_eq!(second_counts.messages_matched_total, 0);
7227 assert_eq!(second_counts.parts_inserted, 3);
7228 assert_eq!(second_counts.parts_matched, 0);
7229
7230 let session_outcome = second_outcomes
7233 .iter()
7234 .find(|outcome| outcome.kind == "session")
7235 .expect("session-row outcome present");
7236 assert_eq!(session_outcome.status, OutcomeStatus::Matched);
7237 for outcome in &second_outcomes {
7238 if outcome.kind == "message" || outcome.kind == "part" {
7239 assert_eq!(
7240 outcome.status,
7241 OutcomeStatus::Inserted,
7242 "new row must be Inserted, got: {outcome:?}",
7243 );
7244 }
7245 }
7246 Ok(())
7247 }
7248
7249 async fn store_with_messages(
7253 temp: &TempDir,
7254 count: usize,
7255 ) -> anyhow::Result<(Store, Vec<MessageKey>)> {
7256 store_with_messages_at_threshold(temp, count, VECTOR_INDEX_ACTIVATION_ROWS).await
7257 }
7258
7259 async fn store_with_messages_at_threshold(
7262 temp: &TempDir,
7263 count: usize,
7264 _vector_threshold: usize,
7265 ) -> anyhow::Result<(Store, Vec<MessageKey>)> {
7266 let store = Store::open_local(temp.path()).await?;
7267 let sessions = 8.min(count.max(1));
7268 let mut events = Vec::new();
7269 for s in 0..sessions {
7270 events.push(IngestEvent::Session(Session {
7271 id: format!("session-{s}"),
7272 parent_session_id: None,
7273 parent_message_id: None,
7274 source_agent: "claude-code".to_owned(),
7275 created_at: Utc::now(),
7276 project: Extracted::from_test_value(format!("/proj/{}", s % 4)),
7277 options: ProviderOptions::new(),
7278 }));
7279 for i in (s..count).step_by(sessions) {
7280 let message_id = format!("msg-{i}");
7281 events.push(IngestEvent::Message(Message::User {
7282 id: message_id.clone(),
7283 session_id: format!("session-{s}"),
7284 timestamp: Utc::now(),
7285 options: ProviderOptions::new(),
7286 }));
7287 events.push(IngestEvent::Part(Part {
7288 session_id: format!("session-{s}"),
7289 id: format!("{message_id}-part"),
7290 message_id,
7291 ordinal: 0,
7292 provenance: crate::wire::Provenance::Conversational,
7293 options: ProviderOptions::new(),
7294 kind: PartKind::Text {
7295 text: Some(Extracted::from_test_value(format!("synthetic message {i}"))),
7296 },
7297 }));
7298 }
7299 }
7300 ingest_events(&store, events).await?;
7301 let keys = (0..count)
7302 .map(|i| MessageKey {
7303 session_id: format!("session-{}", i % sessions),
7304 message_id: format!("msg-{i}"),
7305 })
7306 .collect();
7307 Ok((store, keys))
7308 }
7309
7310 fn synthetic_vector(seed: usize) -> Vec<f32> {
7312 let mut state = (seed as u64)
7313 .wrapping_mul(0x9E37_79B9_7F4A_7C15)
7314 .wrapping_add(1);
7315 (0..embedding_dim())
7316 .map(|_| {
7317 state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
7318 #[allow(clippy::cast_precision_loss)]
7319 let unit = (state >> 33) as f32 / (1u64 << 31) as f32;
7320 unit - 1.0
7321 })
7322 .collect()
7323 }
7324
7325 fn embedded(keys: &[MessageKey]) -> Vec<EmbeddedMessage> {
7327 keys.iter()
7328 .enumerate()
7329 .map(|(seed, key)| EmbeddedMessage {
7330 session_id: key.session_id.clone(),
7331 id: key.message_id.clone(),
7332 vector: synthetic_vector(seed),
7333 })
7334 .collect()
7335 }
7336
7337 fn embedding_update_batch_with_model(
7338 rows: &[EmbeddedMessage],
7339 model: &str,
7340 ) -> Result<RecordBatch> {
7341 let mut batch = embedding_update_batch(rows)?;
7342 let columns = batch
7343 .columns()
7344 .iter()
7345 .take(3)
7346 .cloned()
7347 .chain(std::iter::once(
7348 Arc::new(StringArray::from(vec![model; rows.len()])) as _,
7349 ))
7350 .collect::<Vec<_>>();
7351 batch = RecordBatch::try_new(batch.schema(), columns)?;
7352 Ok(batch)
7353 }
7354
7355 #[tokio::test]
7356 async fn filtered_vector_scan_pushes_scalar_predicate_into_the_index() -> anyhow::Result<()> {
7357 let temp = TempDir::new()?;
7358 let (store, keys) = store_with_messages(&temp, 4).await?;
7362 store.write_embeddings(&embedded(&keys)).await?;
7363 store
7364 .optimize_indices(None, &MaintenancePolicy::always_compact())
7365 .await?
7366 .into_result()?;
7367
7368 let query = vec![0.01_f32; embedding_dim()];
7369 let plan = store
7370 .explain_vector_plan(
7371 &query,
7372 10,
7373 &Predicate::Eq("session_id", "session-3".into()),
7374 None,
7375 )
7376 .await?;
7377
7378 assert!(
7383 plan.contains("ScalarIndexQuery"),
7384 "expected a ScalarIndexQuery node in the plan:\n{plan}",
7385 );
7386 let predicate_postfiltered = plan
7387 .lines()
7388 .any(|line| line.contains("FilterExec") && line.contains("session_id"));
7389 assert!(
7390 !predicate_postfiltered,
7391 "the scalar predicate must not fall back to a FilterExec postfilter:\n{plan}",
7392 );
7393 Ok(())
7394 }
7395
7396 #[tokio::test]
7397 async fn vector_index_activates_when_threshold_is_crossed() -> anyhow::Result<()> {
7398 let temp = TempDir::new()?;
7399 let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7400
7401 store.write_embeddings(&embedded(&keys[..255])).await?;
7404 store
7405 .optimize_indices_with_vector_threshold(256)
7406 .await?
7407 .into_result()?;
7408 assert!(
7409 !store
7410 .handle
7411 .messages_index_names()
7412 .await?
7413 .iter()
7414 .any(|name| name == MESSAGES_VECTOR_INDEX),
7415 "IVF_SQ must not exist below the activation threshold",
7416 );
7417
7418 store.write_embeddings(&embedded(&keys[255..256])).await?;
7421 store
7422 .optimize_indices_with_vector_threshold(256)
7423 .await?
7424 .into_result()?;
7425 assert!(
7426 store
7427 .handle
7428 .messages_index_names()
7429 .await?
7430 .iter()
7431 .any(|name| name == MESSAGES_VECTOR_INDEX),
7432 "optimize must create the IVF_SQ once the threshold is crossed",
7433 );
7434
7435 let hits = store
7438 .vector_search(&synthetic_vector(0), 10, &Predicate::And(Vec::new()), None)
7439 .await?;
7440 assert!(
7441 hits.iter().any(|hit| hit.key == keys[0]),
7442 "an embedded row is retrievable via the index",
7443 );
7444 Ok(())
7445 }
7446
7447 #[tokio::test]
7448 async fn scalar_fold_batching_defers_tail_without_losing_rows() -> anyhow::Result<()> {
7449 let temp = TempDir::new()?;
7450 let (store, _keys) = store_with_messages(&temp, 300).await?;
7451
7452 store
7454 .optimize_indices_with_scalar_fold_threshold(0)
7455 .await?
7456 .into_result()?;
7457 assert_eq!(
7458 store
7459 .handle
7460 .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7461 .await?,
7462 0,
7463 "threshold 0 must fold the scalar index over every row",
7464 );
7465
7466 let new_messages = 5usize;
7468 let mut events = vec![IngestEvent::Session(Session {
7469 id: "session-new".to_owned(),
7470 parent_session_id: None,
7471 parent_message_id: None,
7472 source_agent: "claude-code".to_owned(),
7473 created_at: Utc::now(),
7474 project: Extracted::from_test_value("/proj/new".to_owned()),
7475 options: ProviderOptions::new(),
7476 })];
7477 for i in 0..new_messages {
7478 let message_id = format!("new-msg-{i}");
7479 events.push(IngestEvent::Message(Message::User {
7480 id: message_id.clone(),
7481 session_id: "session-new".to_owned(),
7482 timestamp: Utc::now(),
7483 options: ProviderOptions::new(),
7484 }));
7485 events.push(IngestEvent::Part(Part {
7486 session_id: "session-new".to_owned(),
7487 id: format!("{message_id}-part"),
7488 message_id,
7489 ordinal: 0,
7490 provenance: crate::wire::Provenance::Conversational,
7491 options: ProviderOptions::new(),
7492 kind: PartKind::Text {
7493 text: Some(Extracted::from_test_value(format!("new text {i}"))),
7494 },
7495 }));
7496 }
7497 ingest_events(&store, events).await?;
7498
7499 store
7502 .optimize_indices_with_scalar_fold_threshold(1_000_000)
7503 .await?
7504 .into_result()?;
7505 assert_eq!(
7506 store
7507 .handle
7508 .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7509 .await?,
7510 new_messages,
7511 "the scalar fold must be deferred, leaving the delta tail unindexed",
7512 );
7513 assert_eq!(
7514 store
7515 .handle
7516 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7517 .await?,
7518 0,
7519 "FTS must fold every run regardless of the scalar threshold",
7520 );
7521
7522 let session = store
7524 .get_session("session-new")
7525 .await?
7526 .expect("deferred session must still be retrievable");
7527 assert_eq!(
7528 session.messages.len(),
7529 new_messages,
7530 "get must read the unindexed scalar tail via scan",
7531 );
7532
7533 store
7535 .optimize_indices_with_scalar_fold_threshold(0)
7536 .await?
7537 .into_result()?;
7538 assert_eq!(
7539 store
7540 .handle
7541 .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7542 .await?,
7543 0,
7544 "a threshold-0 fold must consolidate the deferred tail",
7545 );
7546 Ok(())
7547 }
7548
7549 #[tokio::test]
7557 async fn fts_consolidation_rebuilds_instead_of_merging() -> anyhow::Result<()> {
7558 use crate::substrate::{OptimizeEvent, OptimizePhase};
7559
7560 type PhaseLog = std::sync::Arc<std::sync::Mutex<Vec<(OptimizePhase, Option<String>)>>>;
7561 let temp = TempDir::new()?;
7562 let (store, _keys) = store_with_messages(&temp, 300).await?;
7563 let phases: PhaseLog = std::sync::Arc::default();
7564 let sink = phases.clone();
7565 let progress: crate::substrate::OptimizeProgressFn = std::sync::Arc::new(move |event| {
7566 if let OptimizeEvent::PhaseStart { phase, detail, .. } = event {
7567 sink.lock().unwrap().push((phase, detail));
7568 }
7569 });
7570
7571 for round in 0..=crate::substrate::DELTA_MERGE_THRESHOLD {
7572 ingest_events(&store, conversational_events(&format!("grow-{round}"), 2)).await?;
7573 store
7574 .optimize_indices(Some(progress.clone()), &MaintenancePolicy::always_compact())
7575 .await?
7576 .into_result()?;
7577 }
7578
7579 assert!(
7580 phases.lock().unwrap().iter().any(|(phase, detail)| {
7581 matches!(phase, OptimizePhase::IndexRebuild)
7582 && detail.as_deref() == Some(MESSAGES_FTS_INDEX)
7583 }),
7584 "crossing the threshold must rebuild the FTS index, not merge it",
7585 );
7586 let fts_segments = store
7587 .handle
7588 .messages_index_names()
7589 .await?
7590 .into_iter()
7591 .filter(|name| name == MESSAGES_FTS_INDEX)
7592 .count();
7593 assert_eq!(fts_segments, 1, "the rebuild collapses the segment chain");
7594 assert_eq!(
7595 store
7596 .handle
7597 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7598 .await?,
7599 0,
7600 "the rebuilt index covers every fragment",
7601 );
7602 Ok(())
7603 }
7604
7605 #[tokio::test]
7609 async fn grown_session_flush_skips_the_sessions_merge() -> anyhow::Result<()> {
7610 let temp = TempDir::new()?;
7611 let store = Store::open_local(temp.path()).await?;
7612 ingest_events(&store, conversational_events("session-grow", 1)).await?;
7613 let sessions_before = store.handle.dataset(Table::Sessions).await?.version_id();
7614 let messages_before = store.handle.dataset(Table::Messages).await?.version_id();
7615
7616 ingest_events(&store, conversational_events("session-grow", 2)).await?;
7617 assert_eq!(
7618 store.handle.dataset(Table::Sessions).await?.version_id(),
7619 sessions_before,
7620 "an all-present sessions batch must skip the merge commit",
7621 );
7622 assert!(
7623 store.handle.dataset(Table::Messages).await?.version_id() > messages_before,
7624 "the grown message rows must still commit",
7625 );
7626 let session = store
7627 .get_session("session-grow")
7628 .await?
7629 .expect("session row must survive the skipped merge");
7630 assert_eq!(session.messages.len(), 2);
7631 Ok(())
7632 }
7633
7634 #[tokio::test]
7641 async fn fts_fold_skips_a_tail_with_no_indexable_text() -> anyhow::Result<()> {
7642 let temp = TempDir::new()?;
7643 let (store, _keys) = store_with_messages(&temp, 300).await?;
7644 store
7645 .optimize_indices_with_scalar_fold_threshold(0)
7646 .await?
7647 .into_result()?;
7648
7649 let tool_only_message = |i: usize| {
7650 let message_id = format!("tool-msg-{i}");
7651 [
7652 IngestEvent::Message(Message::Assistant {
7653 id: message_id.clone(),
7654 session_id: "session-toolonly".to_owned(),
7655 timestamp: Utc::now(),
7656 options: ProviderOptions::new(),
7657 }),
7658 IngestEvent::Part(Part {
7659 session_id: "session-toolonly".to_owned(),
7660 id: format!("{message_id}-part"),
7661 message_id,
7662 ordinal: 0,
7663 provenance: crate::wire::Provenance::Conversational,
7664 options: ProviderOptions::new(),
7665 kind: PartKind::ToolCall {
7666 call_id: Some(Extracted::from_test_value(format!("call-{i}"))),
7667 name: Some(Extracted::from_test_value("Bash".to_owned())),
7668 params: serde_json::json!({"command": "ls"}),
7669 provider_executed: false,
7670 },
7671 }),
7672 ]
7673 };
7674 let mut events = vec![IngestEvent::Session(synthetic_session("session-toolonly"))];
7675 events.extend((0..3).flat_map(tool_only_message));
7676 ingest_events(&store, events).await?;
7677
7678 store
7679 .optimize_indices_with_scalar_fold_threshold(0)
7680 .await?
7681 .into_result()?;
7682 assert_eq!(
7683 store
7684 .handle
7685 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7686 .await?,
7687 3,
7688 "an all-null tail must not fold into the FTS index",
7689 );
7690 let fts_status = |statuses: Vec<crate::substrate::IndexStatus>| {
7691 statuses
7692 .into_iter()
7693 .find(|status| status.intent_name == MESSAGES_FTS_INDEX)
7694 .expect("FTS status present")
7695 };
7696 assert_eq!(
7697 fts_status(store.index_status().await?).unindexed_rows,
7698 3,
7699 "the raw view counts uncovered rows",
7700 );
7701 assert_eq!(
7702 fts_status(store.index_status_indexable().await?).unindexed_rows,
7703 0,
7704 "the indexable view treats an all-null tail as nothing pending",
7705 );
7706
7707 ingest_events(&store, conversational_events("session-toolonly", 1)).await?;
7710 store
7711 .optimize_indices_with_scalar_fold_threshold(0)
7712 .await?
7713 .into_result()?;
7714 assert_eq!(
7715 store
7716 .handle
7717 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7718 .await?,
7719 0,
7720 "a fold with real text must consolidate the skipped rows too",
7721 );
7722 Ok(())
7723 }
7724
7725 #[tokio::test]
7731 async fn vector_search_covers_the_unindexed_tail() -> anyhow::Result<()> {
7732 let temp = TempDir::new()?;
7733 let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7734 store.write_embeddings(&embedded(&keys)).await?;
7735 store
7736 .optimize_indices_with_vector_threshold(256)
7737 .await?
7738 .into_result()?;
7739 assert_eq!(
7740 store
7741 .handle
7742 .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
7743 .await?,
7744 0,
7745 "the IVF must cover the whole base after the fold",
7746 );
7747
7748 let tail = MessageKey {
7750 session_id: "session-tail".to_owned(),
7751 message_id: "tail-msg".to_owned(),
7752 };
7753 ingest_events(
7754 &store,
7755 vec![
7756 IngestEvent::Session(Session {
7757 id: tail.session_id.clone(),
7758 parent_session_id: None,
7759 parent_message_id: None,
7760 source_agent: "claude-code".to_owned(),
7761 created_at: Utc::now(),
7762 project: Extracted::from_test_value("/proj/tail".to_owned()),
7763 options: ProviderOptions::new(),
7764 }),
7765 IngestEvent::Message(Message::User {
7766 id: tail.message_id.clone(),
7767 session_id: tail.session_id.clone(),
7768 timestamp: Utc::now(),
7769 options: ProviderOptions::new(),
7770 }),
7771 IngestEvent::Part(Part {
7772 session_id: tail.session_id.clone(),
7773 id: format!("{}-part", tail.message_id),
7774 message_id: tail.message_id.clone(),
7775 ordinal: 0,
7776 provenance: crate::wire::Provenance::Conversational,
7777 options: ProviderOptions::new(),
7778 kind: PartKind::Text {
7779 text: Some(Extracted::from_test_value("tail body".to_owned())),
7780 },
7781 }),
7782 ],
7783 )
7784 .await?;
7785 let tail_vector = synthetic_vector(9999);
7786 store
7787 .write_embeddings(&[EmbeddedMessage {
7788 session_id: tail.session_id.clone(),
7789 id: tail.message_id.clone(),
7790 vector: tail_vector.clone(),
7791 }])
7792 .await?;
7793 assert!(
7794 store
7795 .handle
7796 .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
7797 .await?
7798 > 0,
7799 "the appended row must be an unindexed vector tail (no fold ran)",
7800 );
7801
7802 let hits = store
7806 .vector_search(&tail_vector, 10, &Predicate::And(Vec::new()), None)
7807 .await?;
7808 assert!(
7809 hits.iter().any(|hit| hit.key == tail),
7810 "the unindexed tail row must be reachable via vector search (complete recall)",
7811 );
7812 Ok(())
7813 }
7814
7815 #[tokio::test]
7816 async fn model_swap_force_re_embeds_only_stale_rows_and_rebuilds_ivf_pq() -> anyhow::Result<()>
7817 {
7818 let temp = TempDir::new()?;
7819 let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7820 let old_rows = embedded(&keys);
7821 let old_batch = embedding_update_batch_with_model(&old_rows, "old-model")?;
7822 store
7823 .handle
7824 .merge_update(Table::Messages, old_batch, old_rows.len())
7825 .await?;
7826 store
7827 .optimize_indices_with_vector_threshold(256)
7828 .await?
7829 .into_result()?;
7830 assert!(
7831 store
7832 .handle
7833 .messages_index_names()
7834 .await?
7835 .iter()
7836 .any(|name| name == MESSAGES_VECTOR_INDEX),
7837 "IVF_SQ must exist before a model swap",
7838 );
7839 assert_eq!(store.stale_embedding_count().await?, keys.len());
7840
7841 store.drop_vector_index().await?;
7842 let mut pending = Vec::new();
7843 let stream = store.pending_or_stale_messages();
7844 tokio::pin!(stream);
7845 while let Some(row) = stream.next().await {
7846 pending.push(row?);
7847 }
7848 assert_eq!(
7849 pending.len(),
7850 keys.len(),
7851 "force stream should see stale rows"
7852 );
7853 store.write_embeddings(&embedded(&keys)).await?;
7854 assert_eq!(store.stale_embedding_count().await?, 0);
7855 store
7856 .optimize_indices_with_vector_threshold(256)
7857 .await?
7858 .into_result()?;
7859 assert!(
7860 store
7861 .handle
7862 .messages_index_names()
7863 .await?
7864 .iter()
7865 .any(|name| name == MESSAGES_VECTOR_INDEX),
7866 "optimize must rebuild IVF_SQ after force re-embed",
7867 );
7868
7869 let stream = store.pending_or_stale_messages();
7870 tokio::pin!(stream);
7871 assert!(stream.next().await.is_none(), "up-to-date rows are skipped");
7872 Ok(())
7873 }
7874
7875 #[tokio::test]
7876 async fn session_last_message_ids_come_from_durable_messages() -> anyhow::Result<()> {
7877 let temp = TempDir::new()?;
7878 let store = Store::open_local(temp.path()).await?;
7879 let session = synthetic_session("oracle");
7880 store
7881 .upsert_sessions(std::slice::from_ref(&session))
7882 .await?;
7883 let timestamp =
7884 chrono::DateTime::from_timestamp(1_700_000_000, 0).expect("valid timestamp");
7885 let message_a = Message::User {
7886 id: "oracle-a".to_owned(),
7887 session_id: session.id.clone(),
7888 timestamp,
7889 options: ProviderOptions::new(),
7890 };
7891 let message_b = Message::User {
7892 id: "oracle-b".to_owned(),
7893 session_id: session.id.clone(),
7894 timestamp,
7895 options: ProviderOptions::new(),
7896 };
7897 store
7898 .upsert_messages(
7899 &session,
7900 &[
7901 MessageWrite {
7902 message: &message_a,
7903 parts: &[],
7904 search_text: Some("a"),
7905 },
7906 MessageWrite {
7907 message: &message_b,
7908 parts: &[],
7909 search_text: Some("b"),
7910 },
7911 ],
7912 )
7913 .await?;
7914
7915 let empty_session = synthetic_session("session-row-only");
7916 store.upsert_sessions(&[empty_session]).await?;
7917
7918 let orphan = synthetic_session("messages-no-row");
7922 let orphan_message = Message::User {
7923 id: "orphan-a".to_owned(),
7924 session_id: orphan.id.clone(),
7925 timestamp,
7926 options: ProviderOptions::new(),
7927 };
7928 store
7929 .upsert_messages(
7930 &orphan,
7931 &[MessageWrite {
7932 message: &orphan_message,
7933 parts: &[],
7934 search_text: Some("a"),
7935 }],
7936 )
7937 .await?;
7938
7939 let map = store.session_last_message_ids().await?;
7940 assert_eq!(map.get("oracle").map(String::as_str), Some("oracle-b"));
7941 assert!(
7942 !map.contains_key("session-row-only"),
7943 "a session row without durable messages must not produce a freshness key",
7944 );
7945 assert!(
7946 !map.contains_key("messages-no-row"),
7947 "messages without a durable session row must not produce a freshness key",
7948 );
7949 Ok(())
7950 }
7951
7952 #[tokio::test]
7953 async fn embedding_progress_counts_embedded_and_eligible_rows() -> anyhow::Result<()> {
7954 let temp = TempDir::new()?;
7955 let (store, keys) = store_with_messages(&temp, 10).await?;
7956
7957 let before = store.embedding_progress().await?;
7958 assert_eq!(before.embedded, 0);
7959 assert_eq!(before.total, 10);
7960 assert_eq!(before.backlog, 10);
7961 assert_eq!(before.model, crate::embed::model_id());
7962
7963 store.write_embeddings(&embedded(&keys[..4])).await?;
7964 let partial = store.embedding_progress().await?;
7965 assert_eq!(partial.embedded, 4);
7966 assert_eq!(partial.total, 10);
7967 assert_eq!(partial.backlog, 6);
7968
7969 store.write_embeddings(&embedded(&keys[4..])).await?;
7970 let full = store.embedding_progress().await?;
7971 assert_eq!(full.embedded, 10);
7972 assert_eq!(full.total, 10);
7973 assert_eq!(full.backlog, 0);
7976 assert_eq!(full.backlog, store.embed_backlog_count().await?);
7977 Ok(())
7978 }
7979
7980 #[tokio::test]
7981 async fn load_rowmap_if_present_installs_published_chain_without_building() -> anyhow::Result<()>
7982 {
7983 let temp = TempDir::new()?;
7984 let (builder, _keys) = store_with_messages(&temp, 6).await?;
7985 let cache = temp.path().join("cache");
7986
7987 let reader = Store::open_local(temp.path()).await?;
7990 reader.load_rowmap_if_present(&cache).await?;
7991 assert!(reader.rowmap_snapshot().is_none());
7992
7993 builder.ensure_rowmap(&cache).await?;
7995 reader.load_rowmap_if_present(&cache).await?;
7996 assert!(reader.rowmap_snapshot().is_some());
7997 Ok(())
7998 }
7999
8000 #[tokio::test]
8001 async fn ensure_rowmap_layers_a_delta_on_new_ingest() -> anyhow::Result<()> {
8002 let temp = TempDir::new()?;
8003 let (store, _keys) = store_with_messages(&temp, 6).await?;
8004 let cache = temp.path().join("cache");
8005
8006 store.ensure_rowmap(&cache).await?;
8007 assert_eq!(
8008 store.rowmap_delta_count(),
8009 Some(0),
8010 "first build is a lone base"
8011 );
8012
8013 ingest_events(
8015 &store,
8016 vec![
8017 IngestEvent::Session(Session {
8018 id: "session-new".to_owned(),
8019 parent_session_id: None,
8020 parent_message_id: None,
8021 source_agent: "claude-code".to_owned(),
8022 created_at: Utc::now(),
8023 project: Extracted::from_test_value("/proj/new".to_owned()),
8024 options: ProviderOptions::new(),
8025 }),
8026 IngestEvent::Message(Message::User {
8027 id: "m-new".to_owned(),
8028 session_id: "session-new".to_owned(),
8029 timestamp: Utc::now(),
8030 options: ProviderOptions::new(),
8031 }),
8032 IngestEvent::Part(Part {
8033 session_id: "session-new".to_owned(),
8034 id: "m-new-part".to_owned(),
8035 message_id: "m-new".to_owned(),
8036 ordinal: 0,
8037 provenance: crate::wire::Provenance::Conversational,
8038 options: ProviderOptions::new(),
8039 kind: PartKind::Text {
8040 text: Some(Extracted::from_test_value("brand new message".to_owned())),
8041 },
8042 }),
8043 ],
8044 )
8045 .await?;
8046
8047 store.ensure_rowmap(&cache).await?;
8050 assert_eq!(
8051 store.rowmap_delta_count(),
8052 Some(1),
8053 "new ingest layered a delta"
8054 );
8055
8056 let counts = store
8058 .session_message_counts(&["session-new".to_owned()])
8059 .await?;
8060 assert_eq!(counts.get("session-new").copied(), Some(1));
8061 Ok(())
8062 }
8063
8064 #[tokio::test]
8071 async fn ensure_rowmap_rebuilds_when_base_manifest_reclaimed() -> anyhow::Result<()> {
8072 let temp = TempDir::new()?;
8073 let (store, _keys) = store_with_messages(&temp, 6).await?;
8074 let cache = temp.path().join("cache");
8075
8076 store.ensure_rowmap(&cache).await?;
8079 assert_eq!(store.rowmap_delta_count(), Some(0), "first build is a base");
8080 let base_version = store.messages_version().await?;
8081 let versions_dir = temp.path().join("messages.lance").join("_versions");
8082 let base_manifests: Vec<_> = std::fs::read_dir(&versions_dir)?
8083 .filter_map(|entry| entry.ok().map(|entry| entry.path()))
8084 .filter(|path| path.extension().is_some_and(|ext| ext == "manifest"))
8085 .collect();
8086 assert!(
8087 !base_manifests.is_empty(),
8088 "the base version has a manifest"
8089 );
8090
8091 ingest_events(
8094 &store,
8095 vec![
8096 IngestEvent::Session(Session {
8097 id: "session-after".to_owned(),
8098 parent_session_id: None,
8099 parent_message_id: None,
8100 source_agent: "claude-code".to_owned(),
8101 created_at: Utc::now(),
8102 project: Extracted::from_test_value("/proj/after".to_owned()),
8103 options: ProviderOptions::new(),
8104 }),
8105 IngestEvent::Message(Message::User {
8106 id: "m-after".to_owned(),
8107 session_id: "session-after".to_owned(),
8108 timestamp: Utc::now(),
8109 options: ProviderOptions::new(),
8110 }),
8111 IngestEvent::Part(Part {
8112 session_id: "session-after".to_owned(),
8113 id: "m-after-part".to_owned(),
8114 message_id: "m-after".to_owned(),
8115 ordinal: 0,
8116 provenance: crate::wire::Provenance::Conversational,
8117 options: ProviderOptions::new(),
8118 kind: PartKind::Text {
8119 text: Some(Extracted::from_test_value("after the base".to_owned())),
8120 },
8121 }),
8122 ],
8123 )
8124 .await?;
8125 assert!(
8126 store.messages_version().await? > base_version,
8127 "the new ingest advanced the dataset past the chain's base"
8128 );
8129
8130 for manifest in &base_manifests {
8133 std::fs::remove_file(manifest)?;
8134 }
8135
8136 let reopened = Store::open_local(temp.path()).await?;
8139 reopened.ensure_rowmap(&cache).await?;
8140 assert!(
8141 reopened.rowmap_snapshot().is_some(),
8142 "map rebuilt after the base manifest was reclaimed"
8143 );
8144 assert_eq!(
8145 reopened.rowmap_delta_count(),
8146 Some(0),
8147 "a reclaimed base forces a fresh full-scan base, not a stuck chain"
8148 );
8149
8150 let counts = reopened
8152 .session_message_counts(&["session-after".to_owned()])
8153 .await?;
8154 assert_eq!(counts.get("session-after").copied(), Some(1));
8155 Ok(())
8156 }
8157
8158 #[tokio::test]
8165 async fn ensure_rowmap_deltas_across_embedding_fragment_rewrite() -> anyhow::Result<()> {
8166 let temp = TempDir::new()?;
8167 let (store, keys) = store_with_messages(&temp, 6).await?;
8168 let cache = temp.path().join("cache");
8169 store.ensure_rowmap(&cache).await?;
8170 assert_eq!(store.rowmap_delta_count(), Some(0), "first build is a base");
8171
8172 store.write_embeddings(&embedded(&keys)).await?;
8175
8176 ingest_events(
8178 &store,
8179 vec![
8180 IngestEvent::Session(Session {
8181 id: "session-after".to_owned(),
8182 parent_session_id: None,
8183 parent_message_id: None,
8184 source_agent: "claude-code".to_owned(),
8185 created_at: Utc::now(),
8186 project: Extracted::from_test_value("/proj/after".to_owned()),
8187 options: ProviderOptions::new(),
8188 }),
8189 IngestEvent::Message(Message::User {
8190 id: "m-after".to_owned(),
8191 session_id: "session-after".to_owned(),
8192 timestamp: Utc::now(),
8193 options: ProviderOptions::new(),
8194 }),
8195 IngestEvent::Part(Part {
8196 session_id: "session-after".to_owned(),
8197 id: "m-after-part".to_owned(),
8198 message_id: "m-after".to_owned(),
8199 ordinal: 0,
8200 provenance: crate::wire::Provenance::Conversational,
8201 options: ProviderOptions::new(),
8202 kind: PartKind::Text {
8203 text: Some(Extracted::from_test_value("after embedding".to_owned())),
8204 },
8205 }),
8206 ],
8207 )
8208 .await?;
8209
8210 store.ensure_rowmap(&cache).await?;
8213 assert_eq!(
8214 store.rowmap_delta_count(),
8215 Some(1),
8216 "fragment rewrite + append must layer a delta, not full-rebuild"
8217 );
8218
8219 let counts = store
8222 .session_message_counts(&["session-after".to_owned(), "session-0".to_owned()])
8223 .await?;
8224 assert_eq!(counts.get("session-after").copied(), Some(1));
8225 assert_eq!(
8226 counts.get("session-0").copied(),
8227 Some(1),
8228 "a base row survived the rewrite without being double-counted"
8229 );
8230 Ok(())
8231 }
8232
8233 #[tokio::test]
8234 async fn rowmap_chain_compacts_and_stays_bounded() -> anyhow::Result<()> {
8235 let temp = TempDir::new()?;
8238 let (store, _keys) = store_with_messages(&temp, 4).await?;
8239 let cache = temp.path().join("cache");
8240 store.ensure_rowmap(&cache).await?;
8241
8242 let mut reached_cap = false;
8243 let mut compacted = false;
8244 for i in 0..(Store::MAX_ROWMAP_DELTAS + 2) {
8245 let session = format!("session-x{i}");
8246 ingest_events(
8247 &store,
8248 vec![
8249 IngestEvent::Session(Session {
8250 id: session.clone(),
8251 parent_session_id: None,
8252 parent_message_id: None,
8253 source_agent: "claude-code".to_owned(),
8254 created_at: Utc::now(),
8255 project: Extracted::from_test_value("/proj/x".to_owned()),
8256 options: ProviderOptions::new(),
8257 }),
8258 IngestEvent::Message(Message::User {
8259 id: format!("mx{i}"),
8260 session_id: session.clone(),
8261 timestamp: Utc::now(),
8262 options: ProviderOptions::new(),
8263 }),
8264 IngestEvent::Part(Part {
8265 session_id: session.clone(),
8266 id: format!("mx{i}-part"),
8267 message_id: format!("mx{i}"),
8268 ordinal: 0,
8269 provenance: crate::wire::Provenance::Conversational,
8270 options: ProviderOptions::new(),
8271 kind: PartKind::Text {
8272 text: Some(Extracted::from_test_value(format!("msg {i}"))),
8273 },
8274 }),
8275 ],
8276 )
8277 .await?;
8278 store.ensure_rowmap(&cache).await?;
8279 let deltas = store.rowmap_delta_count().unwrap();
8280 assert!(
8281 deltas <= Store::MAX_ROWMAP_DELTAS,
8282 "delta count {deltas} exceeded the cap",
8283 );
8284 if deltas == Store::MAX_ROWMAP_DELTAS {
8285 reached_cap = true;
8286 }
8287 if reached_cap && deltas < Store::MAX_ROWMAP_DELTAS {
8288 compacted = true;
8289 }
8290 }
8291 assert!(reached_cap, "deltas accumulated to the cap (append path)");
8292 assert!(compacted, "the chain compacted back into a base");
8293
8294 let mut rmm = 0;
8296 for entry in std::fs::read_dir(&cache)? {
8297 let name = entry?.file_name().into_string().unwrap_or_default();
8298 assert!(!name.contains(".tmp-"), "leaked build temp: {name}");
8299 if name.ends_with(".rmm") {
8300 rmm += 1;
8301 }
8302 }
8303 assert!(
8304 rmm <= Store::MAX_ROWMAP_DELTAS + 1,
8305 "files unbounded: {rmm}"
8306 );
8307 Ok(())
8308 }
8309
8310 #[tokio::test]
8311 async fn embed_backlog_count_tracks_eligible_unembedded_rows() -> anyhow::Result<()> {
8312 let temp = TempDir::new()?;
8313 let (store, keys) = store_with_messages(&temp, 10).await?;
8314
8315 assert_eq!(store.embed_backlog_count().await?, 10);
8318
8319 store.write_embeddings(&embedded(&keys[..4])).await?;
8320 assert_eq!(store.embed_backlog_count().await?, 6);
8321
8322 store.write_embeddings(&embedded(&keys[4..])).await?;
8323 assert_eq!(store.embed_backlog_count().await?, 0);
8324 Ok(())
8325 }
8326
8327 #[tokio::test]
8332 async fn unindexed_vector_backlog_over_states_when_embedded_rows_are_unfolded()
8333 -> anyhow::Result<()> {
8334 let temp = TempDir::new()?;
8335 let (store, keys) = store_with_messages(&temp, 10).await?;
8336 store.write_embeddings(&embedded(&keys)).await?;
8337
8338 assert_eq!(store.unindexed_vector_backlog().await?, 10);
8339 assert_eq!(store.embed_backlog_count().await?, 0);
8340 Ok(())
8341 }
8342
8343 #[tokio::test]
8344 async fn session_message_counts_returns_per_session_counts_with_zeros_for_unknown_sessions()
8345 -> anyhow::Result<()> {
8346 let temp = TempDir::new()?;
8349 let (store, _keys) = store_with_messages(&temp, 32).await?;
8350
8351 let mut requested: Vec<String> = (0..8).map(|s| format!("session-{s}")).collect();
8352 requested.push("session-unknown-a".to_owned());
8353 requested.push("session-unknown-b".to_owned());
8354 let counts = store.session_message_counts(&requested).await?;
8355
8356 assert_eq!(counts.len(), requested.len());
8359 for s in 0..8 {
8360 assert_eq!(
8361 counts.get(&format!("session-{s}")).copied(),
8362 Some(4),
8363 "session-{s} should have 4 messages",
8364 );
8365 }
8366 assert_eq!(counts.get("session-unknown-a").copied(), Some(0));
8367 assert_eq!(counts.get("session-unknown-b").copied(), Some(0));
8368
8369 let empty = store.session_message_counts(&[]).await?;
8371 assert!(empty.is_empty());
8372 Ok(())
8373 }
8374}