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> = self
1497 .scan_conversational_messages(session_id)
1498 .await?
1499 .into_iter()
1500 .map(|row| ScanRow {
1501 id: row.message_id,
1502 role: row.role,
1503 timestamp: row.timestamp,
1504 text: Some(row.text.into_inner()),
1505 content: None,
1506 })
1507 .collect();
1508 rows.sort_by(|a, b| a.timestamp.cmp(&b.timestamp).then_with(|| a.id.cmp(&b.id)));
1509
1510 let size = |row: &ScanRow| row.text.as_deref().map_or(0, str::len);
1511 let total = rows.len();
1512 let at_pos = params
1518 .at_message_id
1519 .and_then(|at| rows.iter().position(|row| row.id == at));
1520 let (win_start, win_end) = match (at_pos, params.after_message_id, params.before_message_id)
1521 {
1522 (Some(pos), _, _) => {
1523 let n = page_by(&rows[pos..], params.limit, params.budget_bytes, size);
1524 (pos, pos + n)
1525 }
1526 (None, Some(after), _) if params.at_message_id.is_none() => {
1527 let pos = match rows.iter().position(|row| row.id == after) {
1528 Some(idx) => idx + 1,
1529 None => return Ok(GetLookup::UnknownAnchor),
1530 };
1531 let n = page_by(&rows[pos..], params.limit, params.budget_bytes, size);
1532 (pos, pos + n)
1533 }
1534 (None, None, Some(before)) if params.at_message_id.is_none() => {
1535 let pos = match rows.iter().position(|row| row.id == before) {
1536 Some(idx) => idx,
1537 None => return Ok(GetLookup::UnknownAnchor),
1538 };
1539 let n = page_tail(&rows[..pos], params.limit, params.budget_bytes, size);
1540 (pos - n, pos)
1541 }
1542 _ => match params.session_from {
1543 SessionFrom::Start => (0, page_by(&rows, params.limit, params.budget_bytes, size)),
1544 SessionFrom::End => {
1545 let n = page_tail(&rows, params.limit, params.budget_bytes, size);
1546 (total - n, total)
1547 }
1548 },
1549 };
1550 let emitted = &rows[win_start..win_end];
1551 let before_remaining = win_start;
1552 let after_remaining = total - win_end;
1553 let ids: Vec<String> = emitted.iter().map(|row| row.id.clone()).collect();
1554
1555 let mut parts_by_message = self.summary_parts_for_messages(session_id, &ids).await?;
1556 let messages = emitted
1557 .iter()
1558 .map(|row| RetrievedMessage {
1559 id: row.id.clone(),
1560 role: row.role,
1561 timestamp: row.timestamp,
1562 text: row.text.clone(),
1563 content: row.content.clone(),
1564 parts: parts_by_message
1565 .remove(&(session_id.to_owned(), row.id.clone()))
1566 .unwrap_or_default(),
1567 })
1568 .collect();
1569
1570 Ok(GetLookup::Found(SessionPage {
1571 session,
1572 messages,
1573 before_remaining,
1574 after_remaining,
1575 }))
1576 }
1577
1578 pub async fn message_view(
1584 &self,
1585 message_id: &str,
1586 params: MessageViewParams,
1587 ) -> Result<GetLookup<MessagePage>> {
1588 let Some(session_id) = self.session_id_for_message(message_id).await? else {
1589 return Ok(GetLookup::NotFound);
1590 };
1591 let Some(session) = self.find_session(&session_id).await? else {
1592 return Ok(GetLookup::NotFound);
1593 };
1594 let mut rows = self.scan_all_messages(&session_id).await?;
1595 rows.retain(|row| row.text.is_some() || row.id == message_id);
1600 rows.sort_by(|a, b| a.timestamp.cmp(&b.timestamp).then_with(|| a.id.cmp(&b.id)));
1601 let Some(target_pos) = rows.iter().position(|row| row.id == message_id) else {
1602 return Ok(GetLookup::NotFound);
1603 };
1604
1605 let start = target_pos.saturating_sub(params.context_before);
1606 let end = (target_pos + params.context_after + 1).min(rows.len());
1607 let window = &rows[start..end];
1608 let window_ids: Vec<String> = window.iter().map(|row| row.id.clone()).collect();
1609 let mut parts_by_message = self.parts_for_messages(&session_id, &window_ids).await?;
1612
1613 let all_parts = parts_by_message
1614 .remove(&(session_id.clone(), message_id.to_owned()))
1615 .unwrap_or_default();
1616 let part_count = page_by(&all_parts, 1000, params.budget_bytes, |part| {
1619 serde_json::to_string(part).map_or(0, |json| json.len())
1620 });
1621 let target_parts = all_parts[..part_count].to_vec();
1622 let target_parts_remaining = all_parts.len() - part_count;
1623
1624 let target_row = &rows[target_pos];
1625 let target = RetrievedMessage {
1626 id: target_row.id.clone(),
1627 role: target_row.role,
1628 timestamp: target_row.timestamp,
1629 text: target_row.text.clone(),
1630 content: target_row.content.clone(),
1631 parts: Vec::new(),
1633 };
1634 let siblings = window
1635 .iter()
1636 .enumerate()
1637 .filter(|(idx, _)| start + idx != target_pos)
1638 .map(|(_, row)| RetrievedMessage {
1639 id: row.id.clone(),
1640 role: row.role,
1641 timestamp: row.timestamp,
1642 text: row.text.clone(),
1643 content: row.content.clone(),
1644 parts: parts_by_message
1645 .get(&(session_id.clone(), row.id.clone()))
1646 .cloned()
1647 .unwrap_or_default(),
1648 })
1649 .collect();
1650
1651 Ok(GetLookup::Found(MessagePage {
1652 session,
1653 target,
1654 target_parts,
1655 target_parts_remaining,
1656 siblings,
1657 }))
1658 }
1659
1660 async fn scan_all_messages(&self, session_id: &str) -> Result<Vec<ScanRow>> {
1661 let batch = self
1662 .handle
1663 .scan_batch(
1664 Table::Messages,
1665 Some(&Predicate::Eq("session_id", session_id.into())),
1666 &["id", "timestamp", "role", "search_text", "content"],
1667 )
1668 .await?;
1669 let mut rows = Vec::with_capacity(batch.num_rows());
1670 for row in 0..batch.num_rows() {
1671 let id = string(&batch, "id", row)?.context("message id is null")?;
1672 let role =
1673 role_from_str(&string(&batch, "role", row)?.context("message role is null")?)?;
1674 let timestamp = datetime(&batch, "timestamp", row)?;
1675 rows.push(ScanRow {
1676 id,
1677 role,
1678 timestamp,
1679 text: string(&batch, "search_text", row)?,
1680 content: string(&batch, "content", row)?,
1681 });
1682 }
1683 Ok(rows)
1684 }
1685
1686 pub async fn scan_conversational_messages(
1690 &self,
1691 session_id: &str,
1692 ) -> Result<Vec<ConversationalRow>> {
1693 let filter = Predicate::And(vec![
1694 Predicate::Eq("session_id", session_id.into()),
1695 Predicate::IsNotNull("search_text"),
1696 ]);
1697 let batch = self
1698 .handle
1699 .scan_batch(
1700 Table::Messages,
1701 Some(&filter),
1702 &["id", "timestamp", "role", "search_text"],
1703 )
1704 .await?;
1705
1706 let mut rows = Vec::with_capacity(batch.num_rows());
1707 for row in 0..batch.num_rows() {
1708 let message_id = string(&batch, "id", row)?.context("message id is null")?;
1709 let role =
1710 role_from_str(&string(&batch, "role", row)?.context("message role is null")?)?;
1711 let timestamp = datetime(&batch, "timestamp", row)?;
1712 let text_str = string(&batch, "search_text", row)?.context(
1713 "search_text null after IsNotNull pushdown - storage invariant violated",
1714 )?;
1715 rows.push(ConversationalRow {
1716 session_id: session_id.to_owned(),
1717 message_id,
1718 role,
1719 timestamp,
1720 text: SearchText(text_str),
1721 });
1722 }
1723 rows.sort_by(|a, b| {
1724 a.timestamp
1725 .cmp(&b.timestamp)
1726 .then_with(|| a.message_id.cmp(&b.message_id))
1727 });
1728 Ok(rows)
1729 }
1730
1731 pub async fn session_id_for_message(&self, message_id: &str) -> Result<Option<String>> {
1734 let batch = self
1735 .handle
1736 .scan_batch(
1737 Table::Messages,
1738 Some(&Predicate::Eq("id", message_id.into())),
1739 &["session_id"],
1740 )
1741 .await?;
1742 if batch.num_rows() == 0 {
1743 return Ok(None);
1744 }
1745 string(&batch, "session_id", 0)
1746 }
1747
1748 pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
1749 self.handle.row_counts().await
1750 }
1751
1752 pub async fn collect_ids(&self, table: Table) -> Result<std::collections::HashSet<String>> {
1755 self.handle.collect_ids(table).await
1756 }
1757
1758 pub async fn composite_pk_index(&self, table: Table) -> Result<(HashSet<Vec<String>>, usize)> {
1765 let pk = pk_columns(table);
1766 let scanner = self.handle.scan(table, ScanOpts::project_only(pk)).await?;
1767 let mut stream = scanner.try_into_stream().await?;
1768 let mut keys: HashSet<Vec<String>> = HashSet::new();
1769 let mut rows = 0usize;
1770 while let Some(batch) = stream.next().await {
1771 let batch = batch?;
1772 for row in 0..batch.num_rows() {
1773 rows += 1;
1774 keys.insert(composite_key(&batch, pk, row)?);
1775 }
1776 }
1777 Ok((keys, rows))
1778 }
1779
1780 pub async fn composite_pk_diff_against(
1787 &self,
1788 table: Table,
1789 present: &HashSet<Vec<String>>,
1790 ) -> Result<(usize, usize)> {
1791 let pk = pk_columns(table);
1792 let scanner = self.handle.scan(table, ScanOpts::project_only(pk)).await?;
1793 let mut stream = scanner.try_into_stream().await?;
1794 let (mut rows, mut absent) = (0usize, 0usize);
1795 while let Some(batch) = stream.next().await {
1796 let batch = batch?;
1797 for row in 0..batch.num_rows() {
1798 rows += 1;
1799 if !present.contains(&composite_key(&batch, pk, row)?) {
1800 absent += 1;
1801 }
1802 }
1803 }
1804 Ok((rows, absent))
1805 }
1806
1807 pub async fn dataset(&self, table: Table) -> Result<Arc<Dataset>> {
1811 Ok(Arc::new(self.handle.dataset(table).await?))
1812 }
1813
1814 pub async fn prewarm(&self, cache_dir: &Path) -> Result<()> {
1824 let messages = self.dataset(Table::Messages).await?;
1825 if let Err(error) = messages.prewarm_index(MESSAGES_VECTOR_INDEX).await {
1826 tracing::debug!(%error, "vector index prewarm skipped");
1827 }
1828 if let Err(error) = self.ensure_rowmap(cache_dir).await {
1831 tracing::warn!(%error, "rowmap build skipped; arms fall back to data-take resolution");
1832 }
1833 if let Err(error) = self
1836 .fts_search("pond", 1, &Predicate::And(Vec::new()))
1837 .await
1838 {
1839 tracing::debug!(%error, "fts index prewarm skipped");
1840 }
1841 self.prune_index_cache(cache_dir).await;
1842 Ok(())
1843 }
1844
1845 pub async fn prune_index_cache(&self, cache_dir: &Path) {
1848 self.handle.prune_index_cache(cache_dir).await;
1849 }
1850
1851 fn store_key(&self) -> String {
1854 crate::substrate::store_key(self.handle.location())
1855 }
1856
1857 const MAX_ROWMAP_DELTAS: usize = 16;
1859
1860 const ROW_META_COLUMNS: [&str; 7] = [
1865 "session_id",
1866 "id",
1867 "role",
1868 "project",
1869 "source_agent",
1870 "timestamp",
1871 "search_text",
1872 ];
1873
1874 pub async fn ensure_rowmap(&self, cache_dir: &Path) -> Result<()> {
1881 let version = self.messages_version().await?;
1882 if let Some(current) = self.rowmap.load_full()
1883 && current.version() == version
1884 {
1885 return Ok(());
1886 }
1887 std::fs::create_dir_all(cache_dir)
1888 .with_context(|| format!("create cache dir {}", cache_dir.display()))?;
1889 let store_key = self.store_key();
1890
1891 if let Some(chain) = discover_chain(cache_dir, &store_key)
1894 && chain.version() == version
1895 && let Ok(set) = RowMetaSet::open(&chain)
1896 {
1897 self.rowmap.store(Some(Arc::new(set)));
1898 Self::sweep_stale_rowmaps(cache_dir, &store_key, chain.base_version);
1899 return Ok(());
1900 }
1901 if let Some(set) = self
1902 .extend_rowmap_coordinated(cache_dir, &store_key, version)
1903 .await?
1904 {
1905 self.rowmap.store(Some(Arc::new(set)));
1906 }
1907 Ok(())
1908 }
1909
1910 pub fn open_cached_rowmap(&self, cache_dir: &Path) -> Option<Arc<RowMetaSet>> {
1917 let chain = discover_chain(cache_dir, &self.store_key())?;
1918 RowMetaSet::open(&chain).ok().map(Arc::new)
1919 }
1920
1921 pub async fn load_rowmap_if_present(&self, cache_dir: &Path) -> Result<()> {
1927 let version = self.messages_version().await?;
1928 if let Some(current) = self.rowmap.load_full()
1929 && current.version() == version
1930 {
1931 return Ok(());
1932 }
1933 if let Some(chain) = discover_chain(cache_dir, &self.store_key())
1934 && chain.version() == version
1935 && let Ok(set) = RowMetaSet::open(&chain)
1936 {
1937 self.rowmap.store(Some(Arc::new(set)));
1938 }
1939 Ok(())
1940 }
1941
1942 async fn extend_rowmap_coordinated(
1948 &self,
1949 cache_dir: &Path,
1950 store_key: &str,
1951 version: u64,
1952 ) -> Result<Option<RowMetaSet>> {
1953 let lock_path = cache_dir.join(format!("rowmetamap-{store_key}.lock"));
1954 let lock = std::fs::File::create(&lock_path)
1955 .with_context(|| format!("create rowmap build lock {}", lock_path.display()))?;
1956 match lock.try_lock() {
1957 Ok(()) => {}
1958 Err(std::fs::TryLockError::WouldBlock) => return Ok(None),
1959 Err(std::fs::TryLockError::Error(error)) => {
1960 return Err(error).context("lock rowmap build");
1961 }
1962 }
1963
1964 if let Some(chain) = discover_chain(cache_dir, store_key)
1968 && chain.version() == version
1969 && let Ok(set) = RowMetaSet::open(&chain)
1970 {
1971 return Ok(Some(set));
1972 }
1973
1974 Self::sweep_orphan_temps(cache_dir, store_key);
1977
1978 let chain = discover_chain(cache_dir, store_key);
1984 let existing = match &chain {
1985 Some(paths) => match RowMetaSet::open(paths) {
1986 Ok(set) => Some((paths, set)),
1987 Err(error) => {
1988 tracing::warn!(%error, store = store_key, "rowmap unreadable; purging and rebuilding");
1989 Self::purge_rowmaps(cache_dir, store_key);
1990 None
1991 }
1992 },
1993 None => None,
1994 };
1995 let delta = match &existing {
1998 Some((_, set)) => {
1999 self.collect_row_metas_delta(
2000 set.version(),
2001 set.max_row_id().unwrap_or(0),
2002 set.len(),
2003 )
2004 .await?
2005 }
2006 None => None,
2007 };
2008
2009 let base_version = match (&existing, delta) {
2010 (Some((paths, _)), Some(entries)) if paths.deltas.len() < Self::MAX_ROWMAP_DELTAS => {
2012 let path = RowMetaMap::delta_path(cache_dir, store_key, version);
2013 RowMetaMap::build(&path, version, entries)?;
2014 paths.base_version
2015 }
2016 (Some((_, set)), Some(entries)) => {
2020 let mut merged = set.merged_entries();
2021 merged.extend(entries);
2022 let path = RowMetaMap::path_for(cache_dir, store_key, version);
2023 RowMetaMap::build(&path, version, merged)?;
2024 version
2025 }
2026 _ => {
2028 let entries = self.collect_row_metas().await?;
2029 let path = RowMetaMap::path_for(cache_dir, store_key, version);
2030 RowMetaMap::build(&path, version, entries)?;
2031 version
2032 }
2033 };
2034
2035 let chain =
2036 discover_chain(cache_dir, store_key).context("rowmap chain missing after build")?;
2037 let set = RowMetaSet::open(&chain)?;
2038 Self::sweep_stale_rowmaps(cache_dir, store_key, base_version);
2039 Ok(Some(set))
2040 }
2041
2042 fn sweep_stale_rowmaps(cache_dir: &Path, store_key: &str, keep: u64) {
2047 let prefix = format!("rowmetamap-{store_key}-");
2048 let Ok(entries) = std::fs::read_dir(cache_dir) else {
2049 return;
2050 };
2051 for entry in entries.flatten() {
2052 let name = entry.file_name();
2053 let Some(rest) = name
2054 .to_str()
2055 .and_then(|name| name.strip_prefix(&prefix))
2056 .and_then(|rest| rest.strip_suffix(".rmm"))
2057 else {
2058 continue;
2059 };
2060 let version = rest
2061 .strip_prefix('v')
2062 .or_else(|| rest.strip_prefix('d'))
2063 .and_then(|digits| digits.parse::<u64>().ok());
2064 if let Some(version) = version
2065 && version < keep
2066 {
2067 let _ = std::fs::remove_file(entry.path());
2068 }
2069 }
2070 }
2071
2072 fn purge_rowmaps(cache_dir: &Path, store_key: &str) {
2077 let prefix = format!("rowmetamap-{store_key}-");
2078 let Ok(entries) = std::fs::read_dir(cache_dir) else {
2079 return;
2080 };
2081 for entry in entries.flatten() {
2082 if let Some(name) = entry.file_name().to_str()
2083 && name.starts_with(&prefix)
2084 && name.ends_with(".rmm")
2085 {
2086 let _ = std::fs::remove_file(entry.path());
2087 }
2088 }
2089 }
2090
2091 fn sweep_orphan_temps(cache_dir: &Path, store_key: &str) {
2095 let prefix = format!("rowmetamap-{store_key}-");
2096 let Ok(entries) = std::fs::read_dir(cache_dir) else {
2097 return;
2098 };
2099 for entry in entries.flatten() {
2100 let name = entry.file_name();
2101 let Some(name) = name.to_str() else { continue };
2102 if name.starts_with(&prefix) && name.contains(".tmp-") {
2103 let _ = std::fs::remove_file(entry.path());
2104 }
2105 }
2106 }
2107
2108 #[cfg(test)]
2109 pub(crate) fn rowmap_delta_count(&self) -> Option<usize> {
2110 self.rowmap.load_full().map(|set| set.delta_count())
2111 }
2112
2113 pub fn rowmap_snapshot(&self) -> Option<Arc<RowMetaSet>> {
2117 self.rowmap.load_full()
2118 }
2119
2120 async fn resolve_rowid_hits(
2124 &self,
2125 map: &RowMetaSet,
2126 hits: Vec<(u64, f32)>,
2127 ) -> Result<Vec<SearchHit>> {
2128 let mut resolved = Vec::with_capacity(hits.len());
2129 let mut misses: Vec<(u64, f32)> = Vec::new();
2130 for (rowid, score) in hits {
2131 match map.lookup(rowid) {
2132 Some((session_id, message_id)) => resolved.push(SearchHit {
2133 rowid: Some(rowid),
2134 key: MessageKey {
2135 session_id: session_id.to_owned(),
2136 message_id: message_id.to_owned(),
2137 },
2138 score,
2139 }),
2140 None => misses.push((rowid, score)),
2141 }
2142 }
2143 if !misses.is_empty() {
2146 let rowids: Vec<u64> = misses.iter().map(|(rowid, _)| *rowid).collect();
2147 let keys = self.message_keys_by_rowids(&rowids).await?;
2148 for ((rowid, score), key) in misses.into_iter().zip(keys) {
2149 resolved.push(SearchHit {
2150 rowid: Some(rowid),
2151 key,
2152 score,
2153 });
2154 }
2155 }
2156 Ok(resolved)
2157 }
2158
2159 async fn message_keys_by_rowids(&self, rowids: &[u64]) -> Result<Vec<MessageKey>> {
2162 let dataset = self.handle.dataset(Table::Messages).await?;
2163 let projection = ProjectionRequest::from_columns(["session_id", "id"], dataset.schema());
2164 let batch = dataset.take_rows(rowids, projection).await?;
2165 let mut keys = Vec::with_capacity(batch.num_rows());
2166 for row in 0..batch.num_rows() {
2167 keys.push(MessageKey {
2168 session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2169 message_id: string(&batch, "id", row)?.context("fts hit id is null")?,
2170 });
2171 }
2172 Ok(keys)
2173 }
2174
2175 pub async fn export_write(&self, name: &str, bytes: &[u8]) -> Result<()> {
2177 self.handle.export_write(name, bytes).await
2178 }
2179
2180 pub async fn export_read(&self, name: &str) -> Result<Vec<u8>> {
2182 self.handle.export_read(name).await
2183 }
2184
2185 pub fn export_local_path(&self, name: &str) -> Option<std::path::PathBuf> {
2187 self.handle.export_local_path(name)
2188 }
2189
2190 pub async fn adapter_names(&self, include_subagents: bool) -> Result<Vec<String>> {
2196 let scanner = self
2197 .handle
2198 .scan(Table::Sessions, ScanOpts::project_only(&["source_agent"]))
2199 .await?;
2200 let mut stream = scanner.try_into_stream().await?;
2201 let mut names: std::collections::BTreeSet<String> = std::collections::BTreeSet::new();
2202 while let Some(batch) = stream.next().await {
2203 let batch = batch?;
2204 for row in 0..batch.num_rows() {
2205 let agent = string(&batch, "source_agent", row)?.unwrap_or_default();
2206 if !include_subagents && agent.contains('/') {
2207 continue;
2208 }
2209 names.insert(agent);
2210 }
2211 }
2212 Ok(names.into_iter().collect())
2213 }
2214
2215 pub async fn ingest_host_activity(&self) -> Result<Vec<HostActivity>> {
2225 let scanner = self
2226 .handle
2227 .scan(
2228 Table::Messages,
2229 ScanOpts::project_only(&["session_id", "timestamp", "options"]),
2230 )
2231 .await?;
2232 let mut stream = scanner.try_into_stream().await?;
2233 let mut hosts: BTreeMap<Option<String>, (HashSet<String>, DateTime<Utc>)> = BTreeMap::new();
2234 while let Some(batch) = stream.next().await {
2235 let batch = batch?;
2236 for row in 0..batch.num_rows() {
2237 let session_id =
2238 string(&batch, "session_id", row)?.context("session_id is null")?;
2239 let timestamp = datetime(&batch, "timestamp", row)?;
2240 let hostname = json_column(&batch, "options", row)?
2241 .and_then(|bytes| json_parse::<serde_json::Value>(&bytes).ok())
2242 .as_ref()
2243 .and_then(|options| options.pointer("/pond/ingest/host/hostname"))
2244 .and_then(serde_json::Value::as_str)
2245 .map(str::to_owned);
2246 let entry = hosts
2247 .entry(hostname)
2248 .or_insert_with(|| (HashSet::new(), timestamp));
2249 entry.0.insert(session_id);
2250 entry.1 = entry.1.max(timestamp);
2251 }
2252 }
2253 Ok(hosts
2254 .into_iter()
2255 .map(|(hostname, (sessions, last_message_at))| HostActivity {
2256 hostname,
2257 sessions: sessions.len(),
2258 last_message_at,
2259 })
2260 .collect())
2261 }
2262
2263 pub async fn write_embeddings(&self, rows: &[EmbeddedMessage]) -> Result<()> {
2268 if rows.is_empty() {
2269 return Ok(());
2270 }
2271 let batch = embedding_update_batch(rows)?;
2272 self.handle
2273 .merge_update(Table::Messages, batch, rows.len())
2274 .await?;
2275 Ok(())
2276 }
2277
2278 pub fn pending_embedding_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2281 try_stream! {
2282 let filter = Predicate::And(vec![
2288 Predicate::IsNull("embedding_model"),
2289 Predicate::IsNotNull("search_text"),
2290 ]);
2291 let projection: &[&str] = &["session_id", "id", "search_text"];
2292 let scanner = self
2293 .handle
2294 .scan(
2295 Table::Messages,
2296 ScanOpts::with_predicate_and_projection(&filter, projection),
2297 )
2298 .await?;
2299 let mut batches = scanner
2300 .try_into_stream()
2301 .await
2302 .context("failed to open messages stream")?;
2303 while let Some(batch) = batches.next().await {
2304 let batch = batch?;
2305 for row in 0..batch.num_rows() {
2306 yield PendingMessage {
2307 session_id: string(&batch, "session_id", row)?
2308 .context("session_id is null")?,
2309 id: string(&batch, "id", row)?.context("message id is null")?,
2310 search_text: string(&batch, "search_text", row)?
2311 .context("search_text is null")?,
2312 };
2313 }
2314 }
2315 }
2316 }
2317
2318 pub fn pending_or_stale_messages(&self) -> impl Stream<Item = Result<PendingMessage>> + '_ {
2323 try_stream! {
2324 let filter = Predicate::And(vec![
2328 Predicate::IsNotNull("search_text"),
2329 Predicate::Or(vec![
2330 Predicate::IsNull("embedding_model"),
2331 Predicate::Ne("embedding_model", embed::model_id().into()),
2332 ]),
2333 ]);
2334 let projection: &[&str] = &["session_id", "id", "search_text"];
2335 let scanner = self
2336 .handle
2337 .scan(
2338 Table::Messages,
2339 ScanOpts::with_predicate_and_projection(&filter, projection),
2340 )
2341 .await?;
2342 let mut batches = scanner
2343 .try_into_stream()
2344 .await
2345 .context("failed to open pending-or-stale messages stream")?;
2346 while let Some(batch) = batches.next().await {
2347 let batch = batch?;
2348 for row in 0..batch.num_rows() {
2349 yield PendingMessage {
2350 session_id: string(&batch, "session_id", row)?
2351 .context("session_id is null")?,
2352 id: string(&batch, "id", row)?.context("message id is null")?,
2353 search_text: string(&batch, "search_text", row)?
2354 .context("search_text is null")?,
2355 };
2356 }
2357 }
2358 }
2359 }
2360
2361 pub async fn fts_search(
2366 &self,
2367 query: &str,
2368 limit: usize,
2369 filter: &Predicate,
2370 ) -> Result<Vec<SearchHit>> {
2371 let mut hits = if let Some(map) = self.rowmap.load_full() {
2372 let rowid_hits = self.fts_search_rowids(query, limit, filter).await?;
2373 self.resolve_rowid_hits(&map, rowid_hits).await?
2374 } else {
2375 self.fts_search_keys(query, limit, filter).await?
2376 };
2377 hits.sort_by(|left, right| {
2383 right
2384 .score
2385 .partial_cmp(&left.score)
2386 .unwrap_or(std::cmp::Ordering::Equal)
2387 .then_with(|| left.key.session_id.cmp(&right.key.session_id))
2388 .then_with(|| left.key.message_id.cmp(&right.key.message_id))
2389 });
2390 Ok(hits)
2391 }
2392
2393 async fn fts_scanner(
2397 &self,
2398 query: &str,
2399 limit: usize,
2400 filter: &Predicate,
2401 ) -> Result<lance::dataset::scanner::Scanner> {
2402 let mut scanner = self.handle.scanner(Table::Messages, Some(filter)).await?;
2403 scanner.full_text_search(
2404 FullTextSearchQuery::new(query.to_owned()).with_column("search_text".to_owned())?,
2405 )?;
2406 if self
2407 .handle
2408 .messages_fast_search_ready(MESSAGES_FTS_INDEX)
2409 .await?
2410 {
2411 scanner.fast_search();
2412 }
2413 scanner.disable_scoring_autoprojection();
2419 scanner.limit(Some(i64::try_from(limit).unwrap_or(i64::MAX)), None)?;
2420 Ok(scanner)
2421 }
2422
2423 async fn fts_search_keys(
2426 &self,
2427 query: &str,
2428 limit: usize,
2429 filter: &Predicate,
2430 ) -> Result<Vec<SearchHit>> {
2431 let mut scanner = self.fts_scanner(query, limit, filter).await?;
2432 scanner.project(&["session_id", "id", "_score"])?;
2433 let batch = scanner.try_into_batch().await?;
2434 let mut hits = Vec::with_capacity(batch.num_rows());
2435 for row in 0..batch.num_rows() {
2436 let key = MessageKey {
2437 session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2438 message_id: string(&batch, "id", row)?.context("fts hit id is null")?,
2439 };
2440 hits.push(SearchHit {
2441 rowid: None,
2442 key,
2443 score: float32(&batch, "_score", row)?,
2444 });
2445 }
2446 Ok(hits)
2447 }
2448
2449 pub async fn messages_version(&self) -> Result<u64> {
2452 Ok(self
2453 .handle
2454 .dataset(Table::Messages)
2455 .await?
2456 .version()
2457 .version)
2458 }
2459
2460 pub async fn collect_row_metas(&self) -> Result<Vec<RowMetaEntry>> {
2464 let mut scanner = self.handle.scanner(Table::Messages, None).await?;
2465 scanner.with_row_id();
2466 scanner.project(&Self::ROW_META_COLUMNS)?;
2467 let mut stream = scanner.try_into_stream().await?;
2468 let mut out = Vec::new();
2469 while let Some(batch) = stream.next().await {
2470 let batch = batch?;
2471 let rowids = uint64(&batch, "_rowid")?;
2472 for row in 0..batch.num_rows() {
2473 out.push(row_meta_entry(&batch, rowids.value(row), row)?);
2474 }
2475 }
2476 Ok(out)
2477 }
2478
2479 async fn collect_row_metas_delta(
2495 &self,
2496 base_version: u64,
2497 base_max_row_id: u64,
2498 base_row_count: usize,
2499 ) -> Result<Option<Vec<RowMetaEntry>>> {
2500 let dataset = self.handle.dataset(Table::Messages).await?;
2501 let Ok(old) = dataset.checkout_version(base_version).await else {
2502 return Ok(None);
2503 };
2504 if dataset.count_rows(None).await? < base_row_count {
2505 return Ok(None);
2506 }
2507 let old_ids: HashSet<u64> = old.get_fragments().iter().map(|f| f.id() as u64).collect();
2512 let added: Vec<_> = dataset
2513 .get_fragments()
2514 .iter()
2515 .filter(|fragment| !old_ids.contains(&(fragment.id() as u64)))
2516 .map(|fragment| fragment.metadata().clone())
2517 .collect();
2518 if added.is_empty() {
2519 return Ok(Some(Vec::new()));
2520 }
2521 let mut scanner = dataset.scan();
2522 scanner.with_fragments(added);
2523 scanner.with_row_id();
2524 scanner.project(&Self::ROW_META_COLUMNS)?;
2525 let mut stream = scanner.try_into_stream().await?;
2526 let mut out = Vec::new();
2527 while let Some(batch) = stream.next().await {
2528 let batch = batch?;
2529 let rowids = uint64(&batch, "_rowid")?;
2530 for row in 0..batch.num_rows() {
2531 let row_id = rowids.value(row);
2532 if row_id > base_max_row_id {
2533 out.push(row_meta_entry(&batch, row_id, row)?);
2534 }
2535 }
2536 }
2537 Ok(Some(out))
2538 }
2539
2540 async fn fts_search_rowids(
2543 &self,
2544 query: &str,
2545 limit: usize,
2546 filter: &Predicate,
2547 ) -> Result<Vec<(u64, f32)>> {
2548 let mut scanner = self.fts_scanner(query, limit, filter).await?;
2549 scanner.with_row_id();
2550 scanner.project(&["_score"])?;
2551 let batch = scanner.try_into_batch().await?;
2552 let rowids = uint64(&batch, "_rowid")?;
2553 let mut hits = Vec::with_capacity(batch.num_rows());
2554 for row in 0..batch.num_rows() {
2555 hits.push((rowids.value(row), float32(&batch, "_score", row)?));
2556 }
2557 Ok(hits)
2558 }
2559
2560 pub async fn searchable_in_scope(&self, filter: &Predicate) -> Result<usize> {
2567 if matches!(filter, Predicate::And(clauses) if clauses.is_empty())
2573 && let Some(count) = self.fts_num_docs().await?
2574 {
2575 return Ok(count);
2576 }
2577 let scope = Predicate::And(vec![Predicate::IsNotNull("search_text"), filter.clone()]);
2578 let dataset = self.handle.dataset(Table::Messages).await?;
2579 let count = dataset.count_rows(Some(scope.to_lance())).await?;
2580 Ok(count)
2581 }
2582
2583 async fn fts_num_docs(&self) -> Result<Option<usize>> {
2587 if !self.handle.messages_has_index(MESSAGES_FTS_INDEX).await? {
2588 return Ok(None);
2589 }
2590 let dataset = self.handle.dataset(Table::Messages).await?;
2591 let json = dataset.index_statistics(MESSAGES_FTS_INDEX).await?;
2592 let parsed: Value =
2593 serde_json::from_str(&json).context("failed to parse FTS index_statistics")?;
2594 let total: u64 = parsed["indices"]
2595 .as_array()
2596 .map(|segments| {
2597 segments
2598 .iter()
2599 .filter_map(|segment| segment["num_docs"].as_u64())
2600 .sum()
2601 })
2602 .unwrap_or(0);
2603 Ok(Some(usize::try_from(total).unwrap_or(usize::MAX)))
2604 }
2605
2606 pub async fn has_embeddings(&self) -> Result<bool> {
2615 if self
2616 .handle
2617 .messages_has_index(MESSAGES_VECTOR_INDEX)
2618 .await?
2619 {
2620 return Ok(true);
2621 }
2622 let scope = Predicate::IsNotNull("vector");
2623 let mut scanner = self
2624 .handle
2625 .scan(
2626 Table::Messages,
2627 ScanOpts::with_predicate_and_projection(&scope, &["id"]),
2628 )
2629 .await?;
2630 scanner.limit(Some(1), None)?;
2631 let batch = scanner.try_into_batch().await?;
2632 Ok(batch.num_rows() > 0)
2633 }
2634
2635 pub async fn sample_embedded_model(&self) -> Result<Option<String>> {
2641 let scope = Predicate::IsNotNull("embedding_model");
2642 let mut scanner = self
2643 .handle
2644 .scan(
2645 Table::Messages,
2646 ScanOpts::with_predicate_and_projection(&scope, &["embedding_model"]),
2647 )
2648 .await?;
2649 scanner.limit(Some(1), None)?;
2650 let batch = scanner.try_into_batch().await?;
2651 if batch.num_rows() == 0 {
2652 return Ok(None);
2653 }
2654 string(&batch, "embedding_model", 0)
2655 }
2656
2657 pub async fn embedding_model_swapped(&self) -> Result<bool> {
2662 Ok(self
2663 .sample_embedded_model()
2664 .await?
2665 .is_some_and(|model| model != crate::embed::model_id()))
2666 }
2667
2668 pub async fn vector_search(
2676 &self,
2677 query: &[f32],
2678 limit: usize,
2679 filter: &Predicate,
2680 search: Option<&config::SearchConfig>,
2681 ) -> Result<Vec<SearchHit>> {
2682 let mut hits = if let Some(map) = self.rowmap.load_full() {
2683 let rowid_hits = self
2684 .vector_search_rowids(query, limit, filter, search)
2685 .await?;
2686 self.resolve_rowid_hits(&map, rowid_hits).await?
2687 } else {
2688 self.vector_search_keys(query, limit, filter, search)
2689 .await?
2690 };
2691 hits.sort_by(|left, right| {
2697 left.score
2698 .partial_cmp(&right.score)
2699 .unwrap_or(std::cmp::Ordering::Equal)
2700 .then_with(|| left.key.session_id.cmp(&right.key.session_id))
2701 .then_with(|| left.key.message_id.cmp(&right.key.message_id))
2702 });
2703 Ok(hits)
2704 }
2705
2706 async fn vector_scanner(
2708 &self,
2709 query: &[f32],
2710 limit: usize,
2711 filter: &Predicate,
2712 search: Option<&config::SearchConfig>,
2713 ) -> Result<lance::dataset::scanner::Scanner> {
2714 let scope = embedded_scope(filter);
2715 let mut scanner = self.handle.scanner(Table::Messages, Some(&scope)).await?;
2716 let key = Float32Array::from(query.to_vec());
2717 scanner.nearest("vector", &key, limit)?;
2718 apply_vector_search_knobs(&mut scanner, search);
2719 if self
2720 .handle
2721 .messages_fast_search_ready(MESSAGES_VECTOR_INDEX)
2722 .await?
2723 {
2724 scanner.fast_search();
2725 }
2726 scanner.disable_scoring_autoprojection();
2727 Ok(scanner)
2728 }
2729
2730 async fn vector_search_rowids(
2733 &self,
2734 query: &[f32],
2735 limit: usize,
2736 filter: &Predicate,
2737 search: Option<&config::SearchConfig>,
2738 ) -> Result<Vec<(u64, f32)>> {
2739 let mut scanner = self.vector_scanner(query, limit, filter, search).await?;
2740 scanner.with_row_id();
2741 scanner.project(&["_distance"])?;
2742 let batch = scanner.try_into_batch().await?;
2743 let rowids = uint64(&batch, "_rowid")?;
2744 let mut hits = Vec::with_capacity(batch.num_rows());
2745 for row in 0..batch.num_rows() {
2746 hits.push((rowids.value(row), float32(&batch, "_distance", row)?));
2747 }
2748 Ok(hits)
2749 }
2750
2751 async fn vector_search_keys(
2754 &self,
2755 query: &[f32],
2756 limit: usize,
2757 filter: &Predicate,
2758 search: Option<&config::SearchConfig>,
2759 ) -> Result<Vec<SearchHit>> {
2760 let mut scanner = self.vector_scanner(query, limit, filter, search).await?;
2761 scanner.project(&["session_id", "id", "_distance"])?;
2762 let batch = scanner.try_into_batch().await?;
2763 let mut hits = Vec::with_capacity(batch.num_rows());
2764 for row in 0..batch.num_rows() {
2765 let key = MessageKey {
2766 session_id: string(&batch, "session_id", row)?.context("session_id is null")?,
2767 message_id: string(&batch, "id", row)?.context("message id is null")?,
2768 };
2769 hits.push(SearchHit {
2770 rowid: None,
2771 key,
2772 score: float32(&batch, "_distance", row)?,
2773 });
2774 }
2775 Ok(hits)
2776 }
2777
2778 pub async fn explain_vector_plan(
2781 &self,
2782 query: &[f32],
2783 limit: usize,
2784 filter: &Predicate,
2785 search: Option<&config::SearchConfig>,
2786 ) -> Result<String> {
2787 let scanner = self.vector_scanner(query, limit, filter, search).await?;
2791 scanner
2792 .explain_plan(true)
2793 .await
2794 .context("explain_plan failed")
2795 }
2796
2797 pub async fn explain_fts_plan(
2798 &self,
2799 query: &str,
2800 limit: usize,
2801 filter: &Predicate,
2802 ) -> Result<String> {
2803 let mut scanner = self.fts_scanner(query, limit, filter).await?;
2805 scanner.project(&["session_id", "id"])?;
2806 scanner
2807 .explain_plan(true)
2808 .await
2809 .context("explain_plan failed")
2810 }
2811
2812 pub async fn message_metas_by_rowids(&self, rowids: &[u64]) -> Result<Vec<MessageMeta>> {
2820 if rowids.is_empty() {
2821 return Ok(Vec::new());
2822 }
2823 let mut metas = Vec::with_capacity(rowids.len());
2824 let misses: Vec<u64> = if let Some(map) = self.rowmap.load_full() {
2825 let (hits, misses) = map.hydrate(rowids);
2826 metas.extend(hits.into_iter().map(|entry| MessageMeta {
2827 message_id: entry.message_id,
2828 session_id: entry.session_id,
2829 role: entry.role,
2830 project: entry.project,
2831 source_agent: entry.source_agent,
2832 timestamp:
2833 DateTime::from_timestamp_micros(entry.timestamp_micros).unwrap_or_default(),
2834 search_text: entry.search_text,
2835 }));
2836 misses
2837 } else {
2838 rowids.to_vec()
2839 };
2840 if !misses.is_empty() {
2841 metas.extend(self.message_metas_by_rowids_take(&misses).await?);
2842 }
2843 Ok(metas)
2844 }
2845
2846 async fn message_metas_by_rowids_take(&self, rowids: &[u64]) -> Result<Vec<MessageMeta>> {
2851 let dataset = self.handle.dataset(Table::Messages).await?;
2852 let projection = ProjectionRequest::from_columns(
2853 [
2854 "id",
2855 "session_id",
2856 "role",
2857 "project",
2858 "source_agent",
2859 "timestamp",
2860 "search_text",
2861 ],
2862 dataset.schema(),
2863 );
2864 let batch = dataset.take_rows(rowids, projection).await?;
2865 let mut metas = Vec::with_capacity(batch.num_rows());
2866 for row in 0..batch.num_rows() {
2867 metas.push(message_meta_from_batch(&batch, row)?);
2868 }
2869 Ok(metas)
2870 }
2871
2872 pub async fn message_metas_by_keys(&self, keys: &[MessageKey]) -> Result<Vec<MessageMeta>> {
2874 if keys.is_empty() {
2875 return Ok(Vec::new());
2876 }
2877 let wanted = keys.iter().cloned().collect::<HashSet<_>>();
2878 let session_ids = keys
2879 .iter()
2880 .map(|key| key.session_id.clone())
2881 .collect::<Vec<_>>();
2882 let message_ids = keys
2883 .iter()
2884 .map(|key| key.message_id.clone())
2885 .collect::<Vec<_>>();
2886 let predicate = Predicate::And(vec![
2887 in_predicate("session_id", &session_ids),
2888 in_predicate("id", &message_ids),
2889 ]);
2890 let batch = self
2891 .handle
2892 .scan_batch(
2893 Table::Messages,
2894 Some(&predicate),
2895 &[
2896 "id",
2897 "session_id",
2898 "role",
2899 "project",
2900 "source_agent",
2901 "timestamp",
2902 "search_text",
2903 ],
2904 )
2905 .await?;
2906 let mut metas = Vec::with_capacity(batch.num_rows());
2907 for row in 0..batch.num_rows() {
2908 let meta = message_meta_from_batch(&batch, row)?;
2911 if wanted.contains(&MessageKey {
2912 session_id: meta.session_id.clone(),
2913 message_id: meta.message_id.clone(),
2914 }) {
2915 metas.push(meta);
2916 }
2917 }
2918 Ok(metas)
2919 }
2920
2921 pub async fn session_message_counts(
2930 &self,
2931 session_ids: &[String],
2932 ) -> Result<BTreeMap<String, usize>> {
2933 if session_ids.is_empty() {
2934 return Ok(BTreeMap::new());
2935 }
2936 if let Some(map) = self.rowmap.load_full()
2943 && map.version() == self.messages_version().await?
2944 {
2945 return Ok(session_ids
2946 .iter()
2947 .map(|id| (id.clone(), map.lookup_count(id).unwrap_or(0)))
2948 .collect());
2949 }
2950 let predicate = in_predicate("session_id", session_ids);
2951 let scanner = self
2952 .handle
2953 .scan(
2954 Table::Messages,
2955 ScanOpts::with_predicate_and_projection(&predicate, &["session_id"]),
2956 )
2957 .await?;
2958 let mut stream = scanner
2959 .try_into_stream()
2960 .await
2961 .context("failed to open session_message_counts stream")?;
2962 let mut counts: BTreeMap<String, usize> =
2963 session_ids.iter().map(|id| (id.clone(), 0)).collect();
2964 while let Some(batch) = stream.next().await {
2965 let batch = batch.context("failed to read session_message_counts batch")?;
2966 let column = batch
2967 .column_by_name("session_id")
2968 .context("session_message_counts: session_id column missing")?
2969 .as_any()
2970 .downcast_ref::<StringArray>()
2971 .context("session_message_counts: session_id column is not Utf8")?;
2972 for value in column.iter().flatten() {
2973 if let Some(entry) = counts.get_mut(value) {
2974 *entry += 1;
2975 }
2976 }
2977 }
2978 Ok(counts)
2979 }
2980
2981 pub async fn unindexed_message_backlog(&self) -> Result<usize> {
2984 self.handle
2985 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
2986 .await
2987 }
2988
2989 pub async fn unindexed_vector_backlog(&self) -> Result<usize> {
2997 self.handle
2998 .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
2999 .await
3000 }
3001
3002 pub async fn embedding_progress(&self) -> Result<EmbeddingProgress> {
3006 let dataset = self.handle.dataset(Table::Messages).await?;
3007 let embedded = dataset
3013 .count_rows(Some(Predicate::IsNotNull("embedding_model").to_lance()))
3014 .await?;
3015 let backlog = self.embed_backlog_count().await?;
3022 Ok(EmbeddingProgress {
3023 embedded,
3024 total: embedded + backlog,
3025 backlog,
3026 model: embed::model_id(),
3027 })
3028 }
3029
3030 pub async fn embed_backlog_count(&self) -> Result<usize> {
3036 let dataset = self.handle.dataset(Table::Messages).await?;
3037 let filter = Predicate::And(vec![
3038 Predicate::IsNull("embedding_model"),
3039 Predicate::IsNotNull("search_text"),
3040 ]);
3041 Ok(dataset.count_rows(Some(filter.to_lance())).await?)
3042 }
3043
3044 pub async fn stale_embedding_count(&self) -> Result<usize> {
3048 let dataset = self.handle.dataset(Table::Messages).await?;
3049 dataset
3055 .count_rows(Some(
3056 Predicate::And(vec![
3057 Predicate::IsNotNull("embedding_model"),
3058 Predicate::Ne("embedding_model", embed::model_id().into()),
3059 ])
3060 .to_lance(),
3061 ))
3062 .await
3063 .map_err(Into::into)
3064 }
3065
3066 pub async fn optimize_indices(
3072 &self,
3073 progress: Option<OptimizeProgressFn>,
3074 maintenance: &MaintenancePolicy,
3075 ) -> Result<OptimizeOutcome> {
3076 let intents = pond_index_intents();
3077 let mut tables = Vec::with_capacity(3);
3078 for (table, intents) in intents.all() {
3079 let outcome = self
3080 .handle
3081 .optimize_table(table, intents, progress.as_ref(), maintenance)
3082 .await;
3083 tables.push(outcome);
3084 }
3085 Ok(OptimizeOutcome { tables })
3086 }
3087
3088 pub async fn build_indices_only(
3094 &self,
3095 progress: Option<OptimizeProgressFn>,
3096 ) -> Result<OptimizeOutcome> {
3097 let policy = pond_index_intents();
3098 let mut tables = Vec::with_capacity(3);
3099 for (table, intents) in policy.all() {
3100 let indices = self
3101 .handle
3102 .optimize_table_indices_only(table, intents, progress.as_ref())
3103 .await;
3104 tables.push(TableOptimizeOutcome {
3105 table,
3106 indices,
3107 compaction: PhaseOutcome::NotAttempted,
3108 });
3109 }
3110 Ok(OptimizeOutcome { tables })
3111 }
3112
3113 #[cfg(test)]
3114 async fn optimize_indices_with_vector_threshold(
3115 &self,
3116 vector_threshold: usize,
3117 ) -> Result<OptimizeOutcome> {
3118 let intents = pond_index_intents_with_vector_threshold(vector_threshold);
3119 let policy = MaintenancePolicy::always_compact();
3120 let mut tables = Vec::with_capacity(3);
3121 for (table, intents) in intents.all() {
3122 let outcome = self
3123 .handle
3124 .optimize_table(table, intents, None, &policy)
3125 .await;
3126 tables.push(outcome);
3127 }
3128 Ok(OptimizeOutcome { tables })
3129 }
3130
3131 #[cfg(test)]
3132 async fn optimize_indices_with_scalar_fold_threshold(
3133 &self,
3134 scalar_fold_row_threshold: usize,
3135 ) -> Result<OptimizeOutcome> {
3136 let intents = pond_index_intents();
3137 let policy = MaintenancePolicy::always_compact()
3138 .with_scalar_fold_row_threshold(scalar_fold_row_threshold);
3139 let mut tables = Vec::with_capacity(3);
3140 for (table, intents) in intents.all() {
3141 let outcome = self
3142 .handle
3143 .optimize_table(table, intents, None, &policy)
3144 .await;
3145 tables.push(outcome);
3146 }
3147 Ok(OptimizeOutcome { tables })
3148 }
3149
3150 pub async fn cleanup_old_versions(&self, older_than: chrono::Duration) -> Result<()> {
3156 for (table, _) in pond_index_intents().all() {
3157 self.handle
3158 .cleanup_table_versions(table, older_than)
3159 .await?;
3160 }
3161 Ok(())
3162 }
3163
3164 pub async fn rebuild_indices(
3165 &self,
3166 intent_name: Option<&str>,
3167 progress: Option<OptimizeProgressFn>,
3168 ) -> Result<()> {
3169 let policy = pond_index_intents();
3170 let mut matched = false;
3171 for (table, intents) in policy.all() {
3172 for intent in intents {
3173 if intent_name.is_none_or(|name| name == intent.name) {
3174 matched = true;
3175 self.handle
3176 .rebuild_index(table, intent, progress.as_ref())
3177 .await?;
3178 }
3179 }
3180 }
3181 if let Some(name) = intent_name
3182 && !matched
3183 {
3184 anyhow::bail!("unknown index intent {name:?}");
3185 }
3186 Ok(())
3187 }
3188
3189 pub async fn drop_index_by_name(&self, name: &str) -> Result<()> {
3196 let Some(owner) = self.handle.find_index_owner(name).await? else {
3197 anyhow::bail!("no index named {name:?} found on any table");
3198 };
3199 self.handle.drop_index(owner, name).await
3200 }
3201
3202 pub async fn index_status(&self) -> Result<Vec<IndexStatus>> {
3203 self.index_status_with(false).await
3204 }
3205
3206 pub async fn index_status_indexable(&self) -> Result<Vec<IndexStatus>> {
3211 self.index_status_with(true).await
3212 }
3213
3214 async fn index_status_with(&self, indexable_only: bool) -> Result<Vec<IndexStatus>> {
3215 let policy = pond_index_intents();
3216 let mut statuses = Vec::new();
3217 for (table, intents) in policy.all() {
3218 statuses.extend(
3219 self.handle
3220 .index_status(table, intents, indexable_only)
3221 .await?,
3222 );
3223 }
3224 Ok(statuses)
3225 }
3226
3227 pub async fn drop_vector_index(&self) -> Result<()> {
3231 match self
3232 .handle
3233 .drop_index(Table::Messages, MESSAGES_VECTOR_INDEX)
3234 .await
3235 {
3236 Ok(()) => Ok(()),
3237 Err(error) => {
3238 let msg = error.to_string();
3239 if msg.contains("not found") || msg.contains("does not exist") {
3240 Ok(())
3241 } else {
3242 Err(error)
3243 }
3244 }
3245 }
3246 }
3247
3248 pub async fn table_sizes(&self) -> Result<TableSizes> {
3251 self.handle.table_sizes().await
3252 }
3253
3254 pub async fn initialized(&self) -> Result<bool> {
3255 self.handle.initialized().await
3256 }
3257
3258 pub(crate) async fn find_session(&self, session_id: &str) -> Result<Option<Session>> {
3261 let batch = self
3262 .handle
3263 .scan_batch(
3264 Table::Sessions,
3265 Some(&Predicate::Eq("id", session_id.into())),
3266 &[],
3267 )
3268 .await?;
3269 if batch.num_rows() == 0 {
3270 Ok(None)
3271 } else {
3272 Ok(Some(session_from_batch(&batch, 0)?))
3273 }
3274 }
3275
3276 async fn messages_for_session(&self, session_id: &str) -> Result<Vec<MessageWithParts>> {
3277 let batch = self
3278 .handle
3279 .scan_batch(
3280 Table::Messages,
3281 Some(&Predicate::Eq("session_id", session_id.into())),
3282 &[
3283 "session_id",
3284 "id",
3285 "timestamp",
3286 "role",
3287 "content",
3288 "options",
3289 ],
3290 )
3291 .await?;
3292 let mut messages = Vec::with_capacity(batch.num_rows());
3293 for row in 0..batch.num_rows() {
3294 messages.push(message_from_batch(&batch, row)?);
3295 }
3296 messages.sort_by(|left, right| {
3297 left.timestamp()
3298 .cmp(&right.timestamp())
3299 .then_with(|| left.id().cmp(right.id()))
3300 });
3301
3302 let message_ids = messages
3303 .iter()
3304 .map(|message| message.id().to_owned())
3305 .collect::<Vec<_>>();
3306 let mut parts_by_message = self.parts_for_messages(session_id, &message_ids).await?;
3307
3308 Ok(messages
3309 .into_iter()
3310 .map(|message| {
3311 let key = (message.session_id().to_owned(), message.id().to_owned());
3312 let parts = parts_by_message.remove(&key).unwrap_or_default();
3313 MessageWithParts { message, parts }
3314 })
3315 .collect())
3316 }
3317
3318 pub async fn parts_for_messages(
3322 &self,
3323 session_id: &str,
3324 message_ids: &[String],
3325 ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3326 self.scan_parts(session_id, message_ids, None).await
3327 }
3328
3329 pub async fn summary_parts_for_messages(
3334 &self,
3335 session_id: &str,
3336 message_ids: &[String],
3337 ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3338 self.scan_parts(session_id, message_ids, Some(SUMMARY_PART_TYPES))
3339 .await
3340 }
3341
3342 async fn scan_parts(
3343 &self,
3344 session_id: &str,
3345 message_ids: &[String],
3346 part_types: Option<&[&str]>,
3347 ) -> Result<BTreeMap<(String, String), Vec<Part>>> {
3348 if message_ids.is_empty() {
3349 return Ok(BTreeMap::new());
3350 }
3351 let mut clauses = vec![
3352 Predicate::Eq("session_id", session_id.into()),
3353 in_predicate("message_id", message_ids),
3354 ];
3355 if let Some(types) = part_types {
3356 clauses.push(Predicate::In(
3357 "type",
3358 types.iter().map(|&t| t.into()).collect(),
3359 ));
3360 }
3361 let predicate = Predicate::And(clauses);
3362 let summarizing = part_types.is_some();
3367 let mut scanner = self
3368 .handle
3369 .scan(
3370 Table::Parts,
3371 ScanOpts::with_predicate_and_projection(
3372 &predicate,
3373 &[
3374 "session_id",
3375 "message_id",
3376 "id",
3377 "ordinal",
3378 "type",
3379 "provenance",
3380 "variant_data",
3381 "options",
3382 ],
3383 ),
3384 )
3385 .await?;
3386 if !summarizing {
3387 scanner.with_row_address();
3388 }
3389 let batch = scanner.try_into_batch().await.context("scan failed")?;
3390 let mut file_payloads = BTreeMap::<usize, FileData>::new();
3391 if !summarizing {
3392 let dataset = std::sync::Arc::new(self.handle.dataset(Table::Parts).await?);
3393 let row_addresses = uint64(&batch, "_rowaddr")?;
3394 let mut file_rows = Vec::<(usize, u64, Vec<u8>)>::new();
3395 for row in 0..batch.num_rows() {
3396 if string(&batch, "type", row)?.as_deref() == Some("file") {
3397 let variant_data = json_column(&batch, "variant_data", row)?
3398 .context("variant_data is null")?;
3399 file_rows.push((row, row_addresses.value(row), variant_data));
3400 }
3401 }
3402 if !file_rows.is_empty() {
3403 let addresses = file_rows
3404 .iter()
3405 .map(|(_, address, _)| *address)
3406 .collect::<Vec<_>>();
3407 let blobs = dataset.take_blobs_by_addresses(&addresses, "data").await?;
3408 for ((row, _, variant_data), blob) in file_rows.into_iter().zip(blobs) {
3409 let payload = file_data_from_blob(&variant_data, &blob.read().await?)?;
3413 file_payloads.insert(row, payload);
3414 }
3415 }
3416 }
3417 let mut parts_by_message = BTreeMap::<(String, String), Vec<Part>>::new();
3418 for row in 0..batch.num_rows() {
3419 let file_data = if summarizing {
3423 (string(&batch, "type", row)?.as_deref() == Some("file"))
3424 .then(|| FileData::Bytes(Vec::new()))
3425 } else {
3426 file_payloads.remove(&row)
3427 };
3428 let part = part_from_batch(&batch, row, file_data)?;
3429 parts_by_message
3430 .entry((part.session_id.clone(), part.message_id.clone()))
3431 .or_default()
3432 .push(part);
3433 }
3434 for parts in parts_by_message.values_mut() {
3435 parts.sort_by_key(|part| part.ordinal);
3436 }
3437 Ok(parts_by_message)
3438 }
3439}
3440
3441#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
3442#[serde(tag = "kind", content = "data", rename_all = "snake_case")]
3443pub enum IngestEvent {
3444 Session(Session),
3445 Message(Message),
3446 Part(Part),
3447}
3448
3449#[derive(Debug, Clone, PartialEq, Eq, Default)]
3457pub struct IngestSummary {
3458 pub inserted: usize,
3462 pub matched: usize,
3464 pub sessions_inserted: usize,
3466 pub messages_inserted_total: usize,
3469 pub messages_inserted_searchable: usize,
3473 pub parts_inserted: usize,
3475 pub sessions_matched: usize,
3477 pub messages_matched_total: usize,
3479 pub messages_matched_searchable: usize,
3481 pub parts_matched: usize,
3483 pub dropped_events: usize,
3493 pub dropped_sessions: usize,
3498 pub skipped_files: usize,
3501 pub skipped_empty: usize,
3506 pub skipped_fresh: usize,
3510 pub skipped_superseded: usize,
3514 pub storage_errors: usize,
3518 pub truncated_values: usize,
3521 pub drop_reasons: BTreeMap<&'static str, usize>,
3527}
3528
3529pub const DROP_REASON_DUPLICATE_MESSAGE_ID: &str = "duplicate_message_id";
3535pub const DROP_REASON_DUPLICATE_PART_KEY: &str = "duplicate_part_key";
3536pub const DROP_REASON_MESSAGE_BEFORE_SESSION: &str = "message_before_session";
3537pub const DROP_REASON_MESSAGE_SESSION_MISMATCH: &str = "message_session_mismatch";
3538pub const DROP_REASON_PART_BEFORE_MESSAGE: &str = "part_before_message";
3539pub const DROP_REASON_PART_MESSAGE_MISMATCH: &str = "part_message_mismatch";
3540pub const DROP_REASON_EMPTY_SOURCE_AGENT: &str = "empty_source_agent";
3541pub const DROP_REASON_PARENT_MESSAGE_WITHOUT_SESSION: &str = "parent_message_without_session";
3542pub const DROP_REASON_IMMUTABLE_PROJECT: &str = "immutable_project";
3543pub const DROP_REASON_IMMUTABLE_SOURCE_AGENT: &str = "immutable_source_agent";
3544pub const DROP_REASON_UNCATEGORIZED: &str = "uncategorized";
3545
3546#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
3554pub struct BatchCounts {
3555 pub sessions_inserted: usize,
3556 pub sessions_matched: usize,
3557 pub messages_inserted_total: usize,
3558 pub messages_inserted_searchable: usize,
3559 pub messages_matched_total: usize,
3560 pub messages_matched_searchable: usize,
3561 pub parts_inserted: usize,
3562 pub parts_matched: usize,
3563}
3564
3565impl IngestSummary {
3566 pub fn accepted(&self) -> usize {
3567 self.inserted + self.matched
3568 }
3569
3570 pub fn add_batch(&mut self, counts: &BatchCounts) {
3574 self.sessions_inserted += counts.sessions_inserted;
3575 self.sessions_matched += counts.sessions_matched;
3576 self.messages_inserted_total += counts.messages_inserted_total;
3577 self.messages_inserted_searchable += counts.messages_inserted_searchable;
3578 self.messages_matched_total += counts.messages_matched_total;
3579 self.messages_matched_searchable += counts.messages_matched_searchable;
3580 self.parts_inserted += counts.parts_inserted;
3581 self.parts_matched += counts.parts_matched;
3582 self.inserted +=
3583 counts.sessions_inserted + counts.messages_inserted_total + counts.parts_inserted;
3584 self.matched +=
3585 counts.sessions_matched + counts.messages_matched_total + counts.parts_matched;
3586 }
3587
3588 pub fn merge(&mut self, other: &Self) {
3592 self.inserted += other.inserted;
3593 self.matched += other.matched;
3594 self.sessions_inserted += other.sessions_inserted;
3595 self.messages_inserted_total += other.messages_inserted_total;
3596 self.messages_inserted_searchable += other.messages_inserted_searchable;
3597 self.parts_inserted += other.parts_inserted;
3598 self.sessions_matched += other.sessions_matched;
3599 self.messages_matched_total += other.messages_matched_total;
3600 self.messages_matched_searchable += other.messages_matched_searchable;
3601 self.parts_matched += other.parts_matched;
3602 self.dropped_events += other.dropped_events;
3603 self.dropped_sessions += other.dropped_sessions;
3604 self.skipped_files += other.skipped_files;
3605 self.skipped_empty += other.skipped_empty;
3606 self.skipped_fresh += other.skipped_fresh;
3607 self.skipped_superseded += other.skipped_superseded;
3608 self.storage_errors += other.storage_errors;
3609 self.truncated_values += other.truncated_values;
3610 for (key, value) in &other.drop_reasons {
3611 *self.drop_reasons.entry(key).or_insert(0) += value;
3612 }
3613 }
3614
3615 pub fn add_outcomes_errors_only(&mut self, outcomes: &[RowOutcome]) {
3620 for outcome in outcomes {
3621 if !matches!(outcome.status, OutcomeStatus::Error) {
3622 continue;
3623 }
3624 if outcome.kind == "session" {
3625 self.dropped_sessions += 1;
3626 } else {
3627 self.dropped_events += 1;
3628 }
3629 let reason = outcome
3630 .error
3631 .as_ref()
3632 .and_then(|error| error.reason_key)
3633 .unwrap_or(DROP_REASON_UNCATEGORIZED);
3634 *self.drop_reasons.entry(reason).or_insert(0) += 1;
3635 }
3636 }
3637
3638 pub fn add_outcomes(&mut self, outcomes: &[RowOutcome]) {
3639 for outcome in outcomes {
3640 match outcome.status {
3641 OutcomeStatus::Inserted => {
3642 self.inserted += 1;
3643 match outcome.kind {
3644 "session" => self.sessions_inserted += 1,
3645 "message" => {
3646 self.messages_inserted_total += 1;
3647 if outcome.searchable {
3648 self.messages_inserted_searchable += 1;
3649 }
3650 }
3651 "part" => self.parts_inserted += 1,
3652 _ => {}
3653 }
3654 }
3655 OutcomeStatus::Matched => {
3656 self.matched += 1;
3657 match outcome.kind {
3658 "session" => self.sessions_matched += 1,
3659 "message" => {
3660 self.messages_matched_total += 1;
3661 if outcome.searchable {
3662 self.messages_matched_searchable += 1;
3663 }
3664 }
3665 "part" => self.parts_matched += 1,
3666 _ => {}
3667 }
3668 }
3669 OutcomeStatus::Error => {
3670 if outcome.kind == "session" {
3676 self.dropped_sessions += 1;
3677 } else {
3678 self.dropped_events += 1;
3679 }
3680 let reason = outcome
3681 .error
3682 .as_ref()
3683 .and_then(|e| e.reason_key)
3684 .unwrap_or(DROP_REASON_UNCATEGORIZED);
3685 *self.drop_reasons.entry(reason).or_insert(0) += 1;
3686 }
3687 }
3688 }
3689 }
3690}
3691
3692#[derive(Debug, Clone, PartialEq)]
3697pub struct RowOutcome {
3698 pub index: usize,
3699 pub kind: &'static str,
3700 pub pk: Value,
3701 pub status: OutcomeStatus,
3702 pub error: Option<RowError>,
3703 pub searchable: bool,
3708}
3709
3710#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3711pub enum OutcomeStatus {
3712 Inserted,
3713 Matched,
3714 Error,
3715}
3716
3717#[derive(Debug, Clone, PartialEq, Eq)]
3720pub struct RowError {
3721 pub message: String,
3722 pub field: Option<&'static str>,
3723 pub reason: Option<&'static str>,
3724 pub reason_key: Option<&'static str>,
3729}
3730
3731#[derive(Debug)]
3735struct BufferedSession {
3736 index: usize,
3737 session: Session,
3738}
3739
3740#[derive(Debug)]
3741struct BufferedMessage {
3742 index: usize,
3743 message: Message,
3744 parts: Vec<BufferedPart>,
3745 search_text: Option<String>,
3746}
3747
3748#[derive(Debug)]
3749struct BufferedPart {
3750 index: usize,
3751 part: Part,
3752}
3753
3754#[derive(Debug, Default)]
3771pub struct IngestValidator {
3772 session: Option<BufferedSession>,
3773 current_message: Option<BufferedMessage>,
3774 current_parts: Vec<BufferedPart>,
3775 messages: Vec<BufferedMessage>,
3776 seen_message_ids: HashSet<String>,
3780 seen_part_keys: HashSet<(String, String)>,
3783 completed: Vec<CompletedSubstream>,
3787}
3788
3789#[derive(Debug)]
3791struct CompletedSubstream {
3792 session_index: usize,
3793 session: Session,
3794 messages: Vec<BufferedMessage>,
3795}
3796
3797fn ingest_host_stamp() -> Option<&'static Value> {
3802 static STAMP: std::sync::OnceLock<Option<Value>> = std::sync::OnceLock::new();
3803 STAMP
3804 .get_or_init(|| {
3805 let mut host = serde_json::Map::new();
3806 if let Ok(username) = whoami::username() {
3807 host.insert("username".to_owned(), username.into());
3808 }
3809 if let Ok(hostname) = whoami::hostname() {
3810 host.insert("hostname".to_owned(), hostname.into());
3811 }
3812 if let Ok(devicename) = whoami::devicename() {
3813 host.insert("device_name".to_owned(), devicename.into());
3814 }
3815 (!host.is_empty()).then(|| serde_json::json!({ "ingest": { "host": host } }))
3816 })
3817 .as_ref()
3818}
3819
3820impl IngestValidator {
3821 pub async fn push(
3827 &mut self,
3828 store: &Store,
3829 index: usize,
3830 event: IngestEvent,
3831 ) -> Result<Vec<RowOutcome>> {
3832 match event {
3833 IngestEvent::Session(session) => self.push_session(store, index, session).await,
3834 IngestEvent::Message(message) => Ok(self.push_message(index, message)),
3835 IngestEvent::Part(part) => Ok(self.push_part(index, part)),
3836 }
3837 }
3838
3839 pub async fn finish(&mut self, store: &Store) -> Result<(Vec<RowOutcome>, BatchCounts)> {
3844 self.close_current_substream();
3845 self.flush(store).await
3846 }
3847
3848 pub async fn flush(&mut self, store: &Store) -> Result<(Vec<RowOutcome>, BatchCounts)> {
3855 if self.completed.is_empty() {
3856 return Ok((Vec::new(), BatchCounts::default()));
3857 }
3858 let completed = std::mem::take(&mut self.completed);
3859 store.upsert_session_batch(completed).await
3860 }
3861
3862 pub fn pending_substreams(&self) -> usize {
3865 self.completed.len()
3866 }
3867
3868 async fn push_session(
3869 &mut self,
3870 _store: &Store,
3871 index: usize,
3872 mut session: Session,
3873 ) -> Result<Vec<RowOutcome>> {
3874 self.close_current_substream();
3878
3879 let trimmed = session.source_agent.trim();
3884 if trimmed.is_empty() {
3885 return Ok(vec![RowOutcome {
3886 index,
3887 kind: "session",
3888 pk: Value::String(session.id.clone()),
3889 status: OutcomeStatus::Error,
3890 error: Some(RowError {
3891 message: format!("session {} has empty source_agent after trim", session.id),
3892 field: Some("source_agent"),
3893 reason: None,
3894 reason_key: Some(DROP_REASON_EMPTY_SOURCE_AGENT),
3895 }),
3896 searchable: false,
3897 }]);
3898 }
3899 if trimmed.len() != session.source_agent.len() {
3900 session.source_agent = trimmed.to_owned();
3901 }
3902
3903 if session.parent_message_id.is_some() && session.parent_session_id.is_none() {
3904 return Ok(vec![RowOutcome {
3905 index,
3906 kind: "session",
3907 pk: Value::String(session.id.clone()),
3908 status: OutcomeStatus::Error,
3909 error: Some(RowError {
3910 message: format!(
3911 "session {} has parent_message_id without parent_session_id",
3912 session.id,
3913 ),
3914 field: Some("parent_message_id"),
3915 reason: None,
3916 reason_key: Some(DROP_REASON_PARENT_MESSAGE_WITHOUT_SESSION),
3917 }),
3918 searchable: false,
3919 }]);
3920 }
3921
3922 self.seen_message_ids.clear();
3923 self.seen_part_keys.clear();
3924 self.session = Some(BufferedSession { index, session });
3925 Ok(Vec::new())
3926 }
3927
3928 fn close_current_substream(&mut self) {
3929 self.flush_current_message();
3930 let Some(BufferedSession {
3931 index: session_index,
3932 session,
3933 }) = self.session.take()
3934 else {
3935 return;
3936 };
3937 let messages = std::mem::take(&mut self.messages);
3938 self.seen_message_ids.clear();
3939 self.seen_part_keys.clear();
3940 self.completed.push(CompletedSubstream {
3941 session_index,
3942 session,
3943 messages,
3944 });
3945 }
3946
3947 fn push_message(&mut self, index: usize, mut message: Message) -> Vec<RowOutcome> {
3948 let pk = Value::Array(vec![
3949 Value::String(message.session_id().to_owned()),
3950 Value::String(message.id().to_owned()),
3951 ]);
3952 let Some(session) = &self.session else {
3953 return vec![error_outcome(
3954 index,
3955 "message",
3956 pk,
3957 "first event in a session stream must be Session",
3958 None,
3959 DROP_REASON_MESSAGE_BEFORE_SESSION,
3960 )];
3961 };
3962 if message.session_id() != session.session.id {
3963 let msg = format!(
3964 "message {} references session {}, expected {}",
3965 message.id(),
3966 message.session_id(),
3967 session.session.id
3968 );
3969 return vec![error_outcome(
3970 index,
3971 "message",
3972 pk,
3973 &msg,
3974 Some("session_id"),
3975 DROP_REASON_MESSAGE_SESSION_MISMATCH,
3976 )];
3977 }
3978 if !self.seen_message_ids.insert(message.id().to_owned()) {
3979 let msg = format!("duplicate message id {} in session substream", message.id());
3983 return vec![error_outcome(
3984 index,
3985 "message",
3986 pk,
3987 &msg,
3988 None,
3989 DROP_REASON_DUPLICATE_MESSAGE_ID,
3990 )];
3991 }
3992 match ingest_host_stamp() {
3997 Some(stamp) => {
3998 message
3999 .options_mut()
4000 .insert("pond".to_owned(), stamp.clone());
4001 }
4002 None => {
4003 message.options_mut().remove("pond");
4004 }
4005 }
4006 self.flush_current_message();
4007 self.current_message = Some(BufferedMessage {
4008 index,
4009 message,
4010 parts: Vec::new(),
4011 search_text: None,
4012 });
4013 Vec::new()
4014 }
4015
4016 fn push_part(&mut self, index: usize, part: Part) -> Vec<RowOutcome> {
4017 let pk = Value::Array(vec![
4018 Value::String(part.session_id.clone()),
4019 Value::String(part.message_id.clone()),
4020 Value::String(part.id.clone()),
4021 ]);
4022 let Some(current) = &self.current_message else {
4023 return vec![error_outcome(
4024 index,
4025 "part",
4026 pk,
4027 "part event appeared before a message",
4028 None,
4029 DROP_REASON_PART_BEFORE_MESSAGE,
4030 )];
4031 };
4032 if part.session_id != current.message.session_id() {
4033 let msg = format!(
4034 "part {} references session {}, expected {}",
4035 part.id,
4036 part.session_id,
4037 current.message.session_id()
4038 );
4039 return vec![error_outcome(
4040 index,
4041 "part",
4042 pk,
4043 &msg,
4044 Some("session_id"),
4045 DROP_REASON_PART_MESSAGE_MISMATCH,
4046 )];
4047 }
4048 if part.message_id != current.message.id() {
4049 let msg = format!(
4050 "part {} references message {}, expected {}",
4051 part.id,
4052 part.message_id,
4053 current.message.id()
4054 );
4055 return vec![error_outcome(
4056 index,
4057 "part",
4058 pk,
4059 &msg,
4060 Some("message_id"),
4061 DROP_REASON_PART_MESSAGE_MISMATCH,
4062 )];
4063 }
4064 let part_key = (part.message_id.clone(), part.id.clone());
4065 if !self.seen_part_keys.insert(part_key) {
4066 let msg = format!(
4067 "duplicate part id {} for message {} in session substream",
4068 part.id, part.message_id
4069 );
4070 return vec![error_outcome(
4071 index,
4072 "part",
4073 pk,
4074 &msg,
4075 None,
4076 DROP_REASON_DUPLICATE_PART_KEY,
4077 )];
4078 }
4079 self.current_parts.push(BufferedPart { index, part });
4080 Vec::new()
4081 }
4082
4083 fn flush_current_message(&mut self) {
4084 let Some(mut buffered) = self.current_message.take() else {
4085 return;
4086 };
4087 let parts = std::mem::take(&mut self.current_parts);
4088 let mut canonical_parts = Vec::with_capacity(parts.len());
4089 for part in &parts {
4090 canonical_parts.push(part.part.clone());
4091 }
4092 buffered.search_text = search_text(&buffered.message, &canonical_parts);
4093 buffered.parts = parts;
4094 self.messages.push(buffered);
4095 }
4096}
4097
4098fn error_outcome(
4099 index: usize,
4100 kind: &'static str,
4101 pk: Value,
4102 message: &str,
4103 field: Option<&'static str>,
4104 reason_key: &'static str,
4105) -> RowOutcome {
4106 RowOutcome {
4107 index,
4108 kind,
4109 pk,
4110 status: OutcomeStatus::Error,
4111 error: Some(RowError {
4112 message: message.to_owned(),
4113 field,
4114 reason: None,
4115 reason_key: Some(reason_key),
4116 }),
4117 searchable: false,
4118 }
4119}
4120
4121fn error_outcomes_for_substream(
4126 session_index: usize,
4127 session: &Session,
4128 _messages: &[BufferedMessage],
4129 message: impl Into<String>,
4130 field: Option<&'static str>,
4131 reason_key: &'static str,
4132) -> Vec<RowOutcome> {
4133 let reason = field.map(|_| "immutable");
4134 vec![RowOutcome {
4135 index: session_index,
4136 kind: "session",
4137 pk: Value::String(session.id.clone()),
4138 status: OutcomeStatus::Error,
4139 error: Some(RowError {
4140 message: message.into(),
4141 field,
4142 reason,
4143 reason_key: Some(reason_key),
4144 }),
4145 searchable: false,
4146 }]
4147}
4148
4149fn success_outcomes_for_substream(
4155 session_index: usize,
4156 session: &Session,
4157 messages: &[BufferedMessage],
4158 existing_sessions: &std::collections::HashMap<String, Session>,
4159 existing_message_pks: &HashSet<(String, String)>,
4160 existing_part_pks: &HashSet<(String, String, String)>,
4161 counts: &mut BatchCounts,
4162) -> Vec<RowOutcome> {
4163 let session_was_present = existing_sessions.contains_key(&session.id);
4164 let session_status = if session_was_present {
4165 counts.sessions_matched += 1;
4166 UpsertStatus::Matched
4167 } else {
4168 counts.sessions_inserted += 1;
4169 UpsertStatus::Inserted
4170 };
4171
4172 let mut outcomes = Vec::with_capacity(1 + messages.len());
4173 outcomes.push(success_outcome(
4174 session_index,
4175 "session",
4176 Value::String(session.id.clone()),
4177 session_status,
4178 false,
4179 ));
4180 for buffered in messages {
4181 let key = (
4182 buffered.message.session_id().to_owned(),
4183 buffered.message.id().to_owned(),
4184 );
4185 let searchable = buffered.search_text.is_some();
4186 let message_status = if existing_message_pks.contains(&key) {
4187 counts.messages_matched_total += 1;
4188 if searchable {
4189 counts.messages_matched_searchable += 1;
4190 }
4191 UpsertStatus::Matched
4192 } else {
4193 counts.messages_inserted_total += 1;
4194 if searchable {
4195 counts.messages_inserted_searchable += 1;
4196 }
4197 UpsertStatus::Inserted
4198 };
4199 let pk = Value::Array(vec![Value::String(key.0), Value::String(key.1)]);
4200 outcomes.push(success_outcome(
4201 buffered.index,
4202 "message",
4203 pk,
4204 message_status,
4205 searchable,
4206 ));
4207 for part in &buffered.parts {
4208 let part_key = (
4209 part.part.session_id.clone(),
4210 part.part.message_id.clone(),
4211 part.part.id.clone(),
4212 );
4213 let part_status = if existing_part_pks.contains(&part_key) {
4214 counts.parts_matched += 1;
4215 UpsertStatus::Matched
4216 } else {
4217 counts.parts_inserted += 1;
4218 UpsertStatus::Inserted
4219 };
4220 let part_pk = Value::Array(vec![
4221 Value::String(part_key.0),
4222 Value::String(part_key.1),
4223 Value::String(part_key.2),
4224 ]);
4225 outcomes.push(success_outcome(
4226 part.index,
4227 "part",
4228 part_pk,
4229 part_status,
4230 false,
4231 ));
4232 }
4233 }
4234 outcomes
4235}
4236
4237fn success_outcome(
4238 index: usize,
4239 kind: &'static str,
4240 pk: Value,
4241 status: UpsertStatus,
4242 searchable: bool,
4243) -> RowOutcome {
4244 let status = match status {
4245 UpsertStatus::Inserted => OutcomeStatus::Inserted,
4246 UpsertStatus::Matched => OutcomeStatus::Matched,
4247 };
4248 RowOutcome {
4249 index,
4250 kind,
4251 pk,
4252 status,
4253 error: None,
4254 searchable,
4255 }
4256}
4257
4258#[derive(Debug, Clone, PartialEq, Eq)]
4259enum IngestError {
4260 ImmutableField {
4265 field: &'static str,
4266 session_id: String,
4267 stored: String,
4268 attempted: String,
4269 },
4270}
4271
4272impl std::fmt::Display for IngestError {
4273 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
4274 match self {
4275 Self::ImmutableField {
4276 field,
4277 session_id,
4278 stored,
4279 attempted,
4280 } => write!(
4281 formatter,
4282 "session {session_id} {field} is immutable: stored {stored:?}, attempted {attempted:?}",
4283 ),
4284 }
4285 }
4286}
4287
4288impl std::error::Error for IngestError {}
4289
4290fn ensure_immutable_match(
4294 existing: &Session,
4295 incoming: &Session,
4296) -> std::result::Result<(), IngestError> {
4297 if existing.source_agent != incoming.source_agent {
4298 return Err(IngestError::ImmutableField {
4299 field: "source_agent",
4300 session_id: incoming.id.clone(),
4301 stored: existing.source_agent.clone(),
4302 attempted: incoming.source_agent.clone(),
4303 });
4304 }
4305 if existing.project != incoming.project {
4306 return Err(IngestError::ImmutableField {
4307 field: "project",
4308 session_id: incoming.id.clone(),
4309 stored: (*existing.project).clone(),
4310 attempted: (*incoming.project).clone(),
4311 });
4312 }
4313 Ok(())
4314}
4315
4316pub fn search_text(message: &Message, parts: &[Part]) -> Option<String> {
4317 use crate::wire::Provenance;
4318 let mut chunks: Vec<String> = Vec::new();
4319 for part in parts {
4320 if part.provenance != Provenance::Conversational {
4323 continue;
4324 }
4325 match (message.role(), &part.kind) {
4326 (Role::User | Role::Assistant, PartKind::Text { text }) => {
4327 if let Some(text) = text {
4328 chunks.push(text.to_string());
4329 }
4330 }
4331 (
4332 Role::User | Role::Assistant,
4333 PartKind::File {
4334 media_type,
4335 file_name,
4336 data,
4337 },
4338 ) => {
4339 if let Some(file_name) = file_name {
4340 chunks.push(file_name.clone());
4341 }
4342 if let Some(media_type) = media_type {
4343 chunks.push(media_type.clone());
4344 }
4345 if let FileData::Url(uri) = data {
4346 chunks.push(uri.clone());
4347 }
4348 }
4349 (
4350 Role::System | Role::Tool,
4351 PartKind::Text { .. }
4352 | PartKind::Reasoning { .. }
4353 | PartKind::File { .. }
4354 | PartKind::ToolCall { .. }
4355 | PartKind::ToolResult { .. }
4356 | PartKind::ToolApprovalRequest { .. }
4357 | PartKind::ToolApprovalResponse { .. },
4358 )
4359 | (
4360 Role::User | Role::Assistant,
4361 PartKind::Reasoning { .. }
4362 | PartKind::ToolCall { .. }
4363 | PartKind::ToolResult { .. }
4364 | PartKind::ToolApprovalRequest { .. }
4365 | PartKind::ToolApprovalResponse { .. },
4366 ) => {}
4367 }
4368 }
4369
4370 let text = chunks
4371 .into_iter()
4372 .filter(|chunk| !chunk.trim().is_empty())
4373 .collect::<Vec<_>>()
4374 .join("\n");
4375 if text.is_empty() { None } else { Some(text) }
4376}
4377
4378#[derive(Debug, Clone, PartialEq, Eq)]
4380pub struct SearchText(String);
4381
4382impl SearchText {
4383 pub fn as_str(&self) -> &str {
4384 &self.0
4385 }
4386
4387 pub fn into_inner(self) -> String {
4388 self.0
4389 }
4390}
4391
4392impl AsRef<str> for SearchText {
4393 fn as_ref(&self) -> &str {
4394 &self.0
4395 }
4396}
4397
4398#[derive(Debug, Clone, PartialEq)]
4399pub struct MessageWithParts {
4400 pub message: Message,
4401 pub parts: Vec<Part>,
4402}
4403
4404#[derive(Debug, Clone, PartialEq)]
4405pub struct SessionWithMessages {
4406 pub session: Session,
4407 pub messages: Vec<MessageWithParts>,
4408}
4409
4410#[derive(Debug, Clone)]
4411pub struct SessionViewParams<'a> {
4412 pub at_message_id: Option<&'a str>,
4415 pub after_message_id: Option<&'a str>,
4417 pub before_message_id: Option<&'a str>,
4419 pub limit: usize,
4420 pub budget_bytes: usize,
4421 pub session_from: SessionFrom,
4423}
4424
4425#[derive(Debug, Clone)]
4426pub struct MessageViewParams {
4427 pub context_before: usize,
4429 pub context_after: usize,
4431 pub budget_bytes: usize,
4432}
4433
4434#[derive(Debug, Clone, PartialEq)]
4440pub enum GetLookup<T> {
4441 NotFound,
4442 UnknownAnchor,
4443 Found(T),
4444}
4445
4446#[derive(Debug, Clone, PartialEq)]
4450pub struct SessionPage {
4451 pub session: Session,
4452 pub messages: Vec<RetrievedMessage>,
4453 pub before_remaining: usize,
4454 pub after_remaining: usize,
4455}
4456
4457#[derive(Debug, Clone, PartialEq)]
4461pub struct MessagePage {
4462 pub session: Session,
4463 pub target: RetrievedMessage,
4464 pub target_parts: Vec<Part>,
4465 pub target_parts_remaining: usize,
4466 pub siblings: Vec<RetrievedMessage>,
4467}
4468
4469#[derive(Debug, Clone, PartialEq)]
4470pub struct RetrievedMessage {
4471 pub id: String,
4472 pub role: Role,
4473 pub timestamp: DateTime<Utc>,
4474 pub text: Option<String>,
4475 pub content: Option<String>,
4476 pub parts: Vec<Part>,
4477}
4478
4479#[derive(Debug, Clone)]
4480struct ScanRow {
4481 id: String,
4482 role: Role,
4483 timestamp: DateTime<Utc>,
4484 text: Option<String>,
4485 content: Option<String>,
4486}
4487
4488#[derive(Debug, Clone)]
4491pub struct ConversationalRow {
4492 pub session_id: String,
4493 pub message_id: String,
4494 pub role: Role,
4495 pub timestamp: DateTime<Utc>,
4496 pub text: SearchText,
4497}
4498
4499fn page_by<T>(items: &[T], limit: usize, budget_bytes: usize, size: impl Fn(&T) -> usize) -> usize {
4504 let capped = items.len().min(limit.clamp(1, 1000));
4505 let mut acc = 0usize;
4506 let mut emitted = 0usize;
4507 for item in &items[..capped] {
4508 let next = acc.saturating_add(size(item));
4509 if emitted > 0 && next > budget_bytes {
4510 break;
4511 }
4512 acc = next;
4513 emitted += 1;
4514 }
4515 emitted
4516}
4517
4518fn page_tail<T>(
4523 items: &[T],
4524 limit: usize,
4525 budget_bytes: usize,
4526 size: impl Fn(&T) -> usize,
4527) -> usize {
4528 let cap = limit.clamp(1, 1000);
4529 let mut acc = 0usize;
4530 let mut emitted = 0usize;
4531 for item in items.iter().rev() {
4532 if emitted >= cap {
4533 break;
4534 }
4535 let next = acc.saturating_add(size(item));
4536 if emitted > 0 && next > budget_bytes {
4537 break;
4538 }
4539 acc = next;
4540 emitted += 1;
4541 }
4542 emitted
4543}
4544
4545fn role_from_str(value: &str) -> Result<Role> {
4546 match value {
4547 "system" => Ok(Role::System),
4548 "user" => Ok(Role::User),
4549 "assistant" => Ok(Role::Assistant),
4550 "tool" => Ok(Role::Tool),
4551 other => anyhow::bail!("unknown message role {other}"),
4552 }
4553}
4554
4555const MESSAGE_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] = &[
4568 (
4569 "session_id",
4570 BuiltinIndexType::BTree,
4571 MESSAGES_SESSION_ID_INDEX,
4572 ),
4573 (
4574 "source_agent",
4575 BuiltinIndexType::Bitmap,
4576 "messages_source_agent_bitmap",
4577 ),
4578];
4579
4580const PARTS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] = &[
4586 (
4587 "session_id",
4588 BuiltinIndexType::BTree,
4589 "parts_session_id_btree",
4590 ),
4591 (
4592 "message_id",
4593 BuiltinIndexType::BTree,
4594 "parts_message_id_btree",
4595 ),
4596 (
4597 "tool_name",
4598 BuiltinIndexType::BTree,
4599 "parts_tool_name_btree",
4600 ),
4601];
4602
4603const SESSIONS_SCALAR_INDICES: &[(&str, BuiltinIndexType, &str)] =
4606 &[("id", BuiltinIndexType::BTree, "sessions_id_btree")];
4607
4608const COPY_SESSION_IN_CHUNK: usize = 512;
4612
4613fn in_predicate(column: &'static str, values: &[String]) -> Predicate {
4614 Predicate::In(
4615 column,
4616 values.iter().cloned().map(ScalarValue::String).collect(),
4617 )
4618}
4619
4620fn embedded_scope(filter: &Predicate) -> Predicate {
4633 filter.clone()
4634}
4635
4636pub const DEFAULT_NPROBES: usize = 32;
4641
4642fn apply_vector_search_knobs(
4649 scanner: &mut lance::dataset::scanner::Scanner,
4650 search: Option<&config::SearchConfig>,
4651) {
4652 let nprobes = search
4653 .and_then(|cfg| cfg.nprobes)
4654 .unwrap_or(DEFAULT_NPROBES);
4655 scanner.nprobes(nprobes);
4656}
4657
4658pub(crate) const SESSIONS: &str = "sessions";
4662pub(crate) const MESSAGES: &str = "messages";
4663pub(crate) const PARTS: &str = "parts";
4664
4665pub const MESSAGES_SESSION_ID_INDEX: &str = "messages_session_id_btree";
4668
4669pub const MESSAGES_FTS_INDEX: &str = "messages_search_text_fts";
4672
4673pub const MESSAGES_VECTOR_INDEX: &str = "messages_vector_ivfpq";
4680
4681const IVF_SQ_NUM_BITS: u16 = 8;
4686const IVF_SQ_MAX_ITERS: usize = 15;
4687
4688pub fn pond_index_intents() -> IndexIntents {
4691 pond_index_intents_with_vector_threshold(VECTOR_INDEX_ACTIVATION_ROWS)
4692}
4693
4694pub(crate) fn pond_index_intents_with_vector_threshold(vector_threshold: usize) -> IndexIntents {
4698 let mut messages = Vec::with_capacity(MESSAGE_SCALAR_INDICES.len() + 2);
4699 messages.push(IndexIntent {
4700 name: MESSAGES_FTS_INDEX,
4701 column: "search_text",
4702 trigger: IndexTrigger::OnAnyRows,
4703 params: IndexParamsKind::InvertedFtsWord,
4704 });
4705 for (column, kind, name) in MESSAGE_SCALAR_INDICES {
4706 messages.push(IndexIntent {
4707 name,
4708 column,
4709 trigger: IndexTrigger::OnAnyRows,
4710 params: IndexParamsKind::Scalar(kind.clone()),
4711 });
4712 }
4713 messages.push(IndexIntent {
4714 name: MESSAGES_VECTOR_INDEX,
4715 column: "vector",
4716 trigger: IndexTrigger::OnNonNullCount {
4717 column: "vector",
4718 threshold: vector_threshold,
4719 },
4720 params: IndexParamsKind::IvfSqCosine {
4721 num_bits: IVF_SQ_NUM_BITS,
4722 max_iters: IVF_SQ_MAX_ITERS,
4723 },
4724 });
4725 let parts = PARTS_SCALAR_INDICES
4726 .iter()
4727 .map(|(column, kind, name)| IndexIntent {
4728 name,
4729 column,
4730 trigger: IndexTrigger::OnAnyRows,
4731 params: IndexParamsKind::Scalar(kind.clone()),
4732 })
4733 .collect();
4734 let sessions = SESSIONS_SCALAR_INDICES
4735 .iter()
4736 .map(|(column, kind, name)| IndexIntent {
4737 name,
4738 column,
4739 trigger: IndexTrigger::OnAnyRows,
4740 params: IndexParamsKind::Scalar(kind.clone()),
4741 })
4742 .collect();
4743 IndexIntents {
4744 sessions,
4745 messages,
4746 parts,
4747 }
4748}
4749
4750pub const DEFAULT_EMBEDDING_DIM: usize = 384;
4754
4755static EMBEDDING_DIM_RUNTIME: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
4761
4762pub fn embedding_dim() -> usize {
4765 EMBEDDING_DIM_RUNTIME
4766 .get()
4767 .copied()
4768 .unwrap_or(DEFAULT_EMBEDDING_DIM)
4769}
4770
4771pub fn init_embedding_dim(dim: usize) {
4773 EMBEDDING_DIM_RUNTIME.get_or_init(|| dim);
4774}
4775
4776pub(crate) fn write_params_for_create() -> WriteParams {
4783 WriteParams {
4784 data_storage_version: Some(LanceFileVersion::V2_1),
4785 enable_v2_manifest_paths: true,
4786 enable_stable_row_ids: true,
4787 auto_cleanup: Some(AutoCleanupParams {
4788 interval: 20,
4789 older_than: chrono::TimeDelta::days(1),
4790 }),
4791 skip_auto_cleanup: true,
4792 ..WriteParams::default()
4793 }
4794}
4795
4796fn export_schema(table: Table) -> Arc<Schema> {
4797 match table {
4798 Table::Sessions => session_schema(),
4799 Table::Messages => message_schema(),
4800 Table::Parts => part_schema(),
4801 }
4802}
4803
4804fn archive_schema_backfill(dataset: &Dataset, table: Table) -> Result<Option<ColumnBackfill>> {
4809 use std::collections::BTreeSet;
4810 let expected = export_schema(table);
4811 let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
4812 let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
4813 let expected_names: BTreeSet<&str> = expected
4814 .fields()
4815 .iter()
4816 .map(|f| f.name().as_str())
4817 .collect();
4818 let extra: Vec<&str> = actual_names.difference(&expected_names).copied().collect();
4819 if !extra.is_empty() {
4820 anyhow::bail!(
4821 "{} archive table has columns {actual_names:?} but this pond build expects \
4822 {expected_names:?} - the archive was written by a newer pond; upgrade pond \
4823 to restore it",
4824 table.as_str(),
4825 );
4826 }
4827 let missing: Vec<Field> = expected
4828 .fields()
4829 .iter()
4830 .filter(|f| !actual_names.contains(f.name().as_str()))
4831 .map(|f| f.as_ref().clone())
4832 .collect();
4833 if missing.is_empty() {
4834 return Ok(None);
4835 }
4836 column_backfill(table.as_str(), &missing)
4837 .map(Some)
4838 .with_context(|| {
4839 format!(
4840 "{} archive predates a schema change this pond build cannot bridge; \
4841 restore it with the pond version that wrote it",
4842 table.as_str(),
4843 )
4844 })
4845}
4846
4847fn upgraded_batch(batch: &RecordBatch, spec: &ColumnBackfill) -> Result<RecordBatch> {
4853 let derived = (spec.mapper)(batch)?;
4854 let mut fields: Vec<Field> = batch
4855 .schema()
4856 .fields()
4857 .iter()
4858 .map(|field| field.as_ref().clone())
4859 .collect();
4860 let mut columns = batch.columns().to_vec();
4861 for (field, column) in spec.output_schema.fields().iter().zip(derived.columns()) {
4862 fields.push(field.as_ref().clone());
4863 columns.push(column.clone());
4864 }
4865 RecordBatch::try_new(Arc::new(Schema::new(fields)), columns)
4866 .context("failed to assemble upgraded batch")
4867}
4868
4869async fn open_archive_table(table: Table, source: &Path) -> Result<Dataset> {
4870 let source_uri = source
4871 .to_str()
4872 .with_context(|| format!("archive path is not UTF-8: {}", source.display()))?;
4873 Dataset::open(source_uri)
4874 .await
4875 .with_context(|| format!("failed to open {} archive table", table.as_str()))
4876}
4877
4878pub(crate) fn pk_columns(table: Table) -> &'static [&'static str] {
4884 match table {
4885 Table::Sessions => &["id"],
4886 Table::Messages => &["session_id", "id"],
4887 Table::Parts => &["session_id", "message_id", "id"],
4888 }
4889}
4890
4891fn composite_key(batch: &RecordBatch, pk: &[&str], row: usize) -> Result<Vec<String>> {
4893 pk.iter()
4894 .map(|column| string(batch, column, row)?.with_context(|| format!("{column} is null")))
4895 .collect()
4896}
4897
4898pub(crate) fn session_schema() -> Arc<Schema> {
4899 Arc::new(Schema::new(vec![
4900 primary_field("id", DataType::Utf8, false),
4901 Field::new("parent_session_id", DataType::Utf8, true),
4902 Field::new("parent_message_id", DataType::Utf8, true),
4903 Field::new("source_agent", DataType::Utf8, false),
4904 Field::new(
4905 "created_at",
4906 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4907 false,
4908 ),
4909 Field::new("project", DataType::Utf8, false),
4910 json_field("options", false),
4911 ]))
4912}
4913
4914pub(crate) fn message_schema() -> Arc<Schema> {
4915 Arc::new(Schema::new(vec![
4916 primary_field("session_id", DataType::Utf8, false),
4917 primary_field("id", DataType::Utf8, false),
4918 Field::new(
4919 "timestamp",
4920 DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())),
4921 false,
4922 ),
4923 Field::new("role", DataType::Utf8, false),
4924 Field::new("source_agent", DataType::Utf8, false),
4925 Field::new("project", DataType::Utf8, false),
4926 Field::new("content", DataType::Utf8, true),
4927 Field::new("search_text", DataType::Utf8, true),
4928 Field::new("vector", embedding_vector_type(), true),
4932 Field::new("embedding_model", DataType::Utf8, true),
4933 json_field("options", false),
4934 ]))
4935}
4936
4937pub(crate) type BackfillMapper = Box<dyn Fn(&RecordBatch) -> Result<RecordBatch> + Send + Sync>;
4939
4940pub(crate) struct ColumnBackfill {
4946 pub read_columns: Vec<String>,
4947 pub output_schema: Arc<Schema>,
4948 pub mapper: BackfillMapper,
4949}
4950
4951pub(crate) fn column_backfill(table_name: &str, missing: &[Field]) -> Result<ColumnBackfill> {
4954 match table_name {
4955 PARTS => parts_backfill(missing),
4956 other => anyhow::bail!(
4957 "table {other} is missing columns {:?} that this pond build cannot derive from \
4958 stored data - use the pond version that wrote it",
4959 missing.iter().map(|f| f.name()).collect::<Vec<_>>(),
4960 ),
4961 }
4962}
4963
4964fn parts_backfill(missing: &[Field]) -> Result<ColumnBackfill> {
4965 for field in missing {
4966 anyhow::ensure!(
4967 ["tool_name", "call_id", "is_failure"].contains(&field.name().as_str())
4968 && field.is_nullable(),
4969 "column {PARTS}.{} cannot be derived from stored data - use the pond version \
4970 that wrote it",
4971 field.name(),
4972 );
4973 }
4974 let output_schema = Arc::new(Schema::new(missing.to_vec()));
4975 let schema = output_schema.clone();
4976 let mapper = move |batch: &RecordBatch| -> Result<RecordBatch> {
4977 let rows = batch.num_rows();
4978 let mut names: Vec<Option<String>> = Vec::with_capacity(rows);
4979 let mut call_ids: Vec<Option<String>> = Vec::with_capacity(rows);
4980 let mut failures: Vec<Option<bool>> = Vec::with_capacity(rows);
4981 for row in 0..rows {
4982 let type_name = string(batch, "type", row)?.context("part type is null")?;
4983 let kind = match type_name.as_str() {
4986 "tool_call" | "tool_result" | "tool_approval_request" => {
4987 let body =
4988 json_column(batch, "variant_data", row)?.context("variant_data is null")?;
4989 match part_kind_from_json(&type_name, &body, None) {
4994 Ok(kind) => Some(kind),
4995 Err(error) => {
4996 tracing::warn!(
4997 session_id = string(batch, "session_id", row)?.as_deref(),
4998 part_id = string(batch, "id", row)?.as_deref(),
4999 error = %format!("{error:#}"),
5000 "stored tool part body failed to decode; its backfilled \
5001 cells stay NULL",
5002 );
5003 None
5004 }
5005 }
5006 }
5007 _ => None,
5008 };
5009 let (name, call_id, is_failure) = kind
5010 .as_ref()
5011 .map(tool_identity)
5012 .unwrap_or((None, None, None));
5013 names.push(name.map(str::to_owned));
5014 call_ids.push(call_id.map(str::to_owned));
5015 failures.push(is_failure);
5016 }
5017 let arrays: Vec<ArrayRef> = schema
5018 .fields()
5019 .iter()
5020 .map(|field| -> ArrayRef {
5021 match field.name().as_str() {
5022 "tool_name" => Arc::new(StringArray::from(names.clone())),
5023 "call_id" => Arc::new(StringArray::from(call_ids.clone())),
5024 _ => Arc::new(BooleanArray::from(failures.clone())),
5025 }
5026 })
5027 .collect();
5028 RecordBatch::try_new(schema.clone(), arrays).context("failed to build backfill batch")
5029 };
5030 Ok(ColumnBackfill {
5031 read_columns: vec![
5032 "session_id".to_owned(),
5033 "id".to_owned(),
5034 "type".to_owned(),
5035 "variant_data".to_owned(),
5036 ],
5037 output_schema,
5038 mapper: Box::new(mapper),
5039 })
5040}
5041
5042pub(crate) fn part_schema() -> Arc<Schema> {
5043 Arc::new(Schema::new(vec![
5044 primary_field("session_id", DataType::Utf8, false),
5045 primary_field("message_id", DataType::Utf8, false),
5046 primary_field("id", DataType::Utf8, false),
5047 Field::new("ordinal", DataType::Int32, false),
5048 Field::new("type", DataType::Utf8, false),
5049 Field::new("provenance", DataType::Utf8, false),
5052 Field::new("tool_name", DataType::Utf8, true),
5057 Field::new("call_id", DataType::Utf8, true),
5058 Field::new("is_failure", DataType::Boolean, true),
5059 json_field("variant_data", false),
5060 legacy_blob_field("data", true),
5061 json_field("options", false),
5062 ]))
5063}
5064
5065pub(crate) fn empty_batch(schema: Arc<Schema>) -> Result<RecordBatch> {
5066 let arrays = schema
5067 .fields()
5068 .iter()
5069 .map(|field| lance::deps::arrow_array::new_empty_array(field.data_type()))
5070 .collect();
5071 RecordBatch::try_new(schema, arrays).context("failed to build empty Lance batch")
5072}
5073
5074pub(crate) fn empty_reader(
5075 schema: Arc<Schema>,
5076) -> Result<
5077 RecordBatchIterator<
5078 std::vec::IntoIter<Result<RecordBatch, lance::deps::arrow_schema::ArrowError>>,
5079 >,
5080> {
5081 let batch = empty_batch(schema.clone())?;
5082 Ok(RecordBatchIterator::new(
5083 vec![Ok(batch)].into_iter(),
5084 schema,
5085 ))
5086}
5087
5088pub(crate) struct MessageBatchRow<'a> {
5089 pub message: &'a Message,
5090 pub source_agent: &'a str,
5091 pub project: &'a str,
5092 pub search_text: Option<&'a str>,
5093}
5094
5095fn embedding_vector_type() -> DataType {
5101 DataType::FixedSizeList(
5102 Arc::new(Field::new("item", DataType::Float16, true)),
5103 embedding_dim() as i32,
5104 )
5105}
5106
5107fn embedding_update_schema() -> Arc<Schema> {
5111 Arc::new(Schema::new(vec![
5112 primary_field("session_id", DataType::Utf8, false),
5113 primary_field("id", DataType::Utf8, false),
5114 Field::new("vector", embedding_vector_type(), true),
5115 Field::new("embedding_model", DataType::Utf8, true),
5116 ]))
5117}
5118
5119fn embedding_columns(vectors: &[Option<Vec<f32>>]) -> Result<(ArrayRef, ArrayRef)> {
5123 let dim = embedding_dim();
5124 if vectors.iter().all(Option::is_none) {
5127 return Ok((
5128 new_null_array(&embedding_vector_type(), vectors.len()),
5129 new_null_array(&DataType::Utf8, vectors.len()),
5130 ));
5131 }
5132 let mut builder = FixedSizeListBuilder::new(
5133 Float16Builder::with_capacity(vectors.len() * dim),
5134 dim as i32,
5135 )
5136 .with_field(Arc::new(Field::new("item", DataType::Float16, true)));
5137 let mut models: Vec<Option<&str>> = Vec::with_capacity(vectors.len());
5138 for vector in vectors {
5139 match vector {
5140 Some(values) => {
5141 if values.len() != dim {
5142 anyhow::bail!("inline embedding has dim {}, expected {dim}", values.len());
5143 }
5144 for value in values {
5145 builder.values().append_value(half::f16::from_f32(*value));
5146 }
5147 builder.append(true);
5148 models.push(Some(embed::model_id()));
5149 }
5150 None => {
5151 for _ in 0..dim {
5152 builder.values().append_null();
5153 }
5154 builder.append(false);
5155 models.push(None);
5156 }
5157 }
5158 }
5159 Ok((
5160 Arc::new(builder.finish()) as ArrayRef,
5161 Arc::new(StringArray::from(models)) as ArrayRef,
5162 ))
5163}
5164
5165pub(crate) fn embedding_update_batch(rows: &[EmbeddedMessage]) -> Result<RecordBatch> {
5168 let dim = embedding_dim();
5169 let mut flat = Vec::with_capacity(rows.len() * dim);
5170 for row in rows {
5171 if row.vector.len() != dim {
5172 anyhow::bail!(
5173 "embedding for message {} has dim {}, expected {dim}",
5174 row.id,
5175 row.vector.len(),
5176 );
5177 }
5178 flat.extend(row.vector.iter().map(|value| half::f16::from_f32(*value)));
5179 }
5180 let values = Float16Array::from(flat);
5181 let item_field = Arc::new(Field::new("item", DataType::Float16, true));
5182 let vectors = FixedSizeListArray::try_new(item_field, dim as i32, Arc::new(values), None)
5183 .context("failed to build embedding vector column")?;
5184
5185 RecordBatch::try_new(
5186 embedding_update_schema(),
5187 vec![
5188 Arc::new(StringArray::from(
5189 rows.iter()
5190 .map(|row| row.session_id.as_str())
5191 .collect::<Vec<_>>(),
5192 )),
5193 Arc::new(StringArray::from(
5194 rows.iter().map(|row| row.id.as_str()).collect::<Vec<_>>(),
5195 )),
5196 Arc::new(vectors),
5197 Arc::new(StringArray::from(vec![embed::model_id(); rows.len()])),
5198 ],
5199 )
5200 .context("failed to build embedding update batch")
5201}
5202
5203const COLUMN_BYTE_BUDGET: usize = 1 << 30;
5208
5209fn chunk_ranges(cells: &[usize]) -> Vec<std::ops::Range<usize>> {
5214 let mut chunks = Vec::new();
5215 let mut start = 0usize;
5216 let mut running = 0usize;
5217 for (index, &row) in cells.iter().enumerate() {
5218 if running + row > COLUMN_BYTE_BUDGET && index > start {
5219 chunks.push(start..index);
5220 start = index;
5221 running = 0;
5222 }
5223 running += row;
5224 }
5225 if start < cells.len() {
5226 chunks.push(start..cells.len());
5227 }
5228 chunks
5229}
5230
5231fn guard_cell(table: &str, pk: &str, bytes: usize) -> Result<()> {
5232 if bytes >= COLUMN_BYTE_BUDGET {
5233 anyhow::bail!(
5234 "{table} row {pk}: a {bytes}-byte text cell meets the per-cell ceiling and would \
5235 overflow Arrow's i32 offset buffer"
5236 );
5237 }
5238 Ok(())
5239}
5240
5241async fn merge_insert_chunks(
5242 handle: &Handle,
5243 table: Table,
5244 batches: Vec<RecordBatch>,
5245) -> Result<u64> {
5246 let mut inserted = 0u64;
5247 for batch in batches {
5248 let rows = batch.num_rows();
5249 inserted += handle.merge_insert(table, batch, rows).await?;
5250 }
5251 Ok(inserted)
5252}
5253
5254pub(crate) fn sessions_batches(sessions: &[Session]) -> Result<Vec<RecordBatch>> {
5255 let options = sessions
5256 .iter()
5257 .map(|session| json_bytes(&session.options))
5258 .collect::<Result<Vec<_>>>()?;
5259 let mut cells = Vec::with_capacity(sessions.len());
5260 for (session, encoded) in sessions.iter().zip(&options) {
5261 let columns = [
5262 session.id.len(),
5263 session.parent_session_id.as_deref().map_or(0, str::len),
5264 session.parent_message_id.as_deref().map_or(0, str::len),
5265 session.source_agent.len(),
5266 session.project.as_str().len(),
5267 encoded.len(),
5268 ];
5269 for bytes in columns {
5270 guard_cell("sessions", &session.id, bytes)?;
5271 }
5272 cells.push(columns.iter().sum());
5273 }
5274 chunk_ranges(&cells)
5275 .into_iter()
5276 .map(|range| sessions_chunk(&sessions[range.clone()], &options[range]))
5277 .collect()
5278}
5279
5280fn sessions_chunk(sessions: &[Session], options: &[Vec<u8>]) -> Result<RecordBatch> {
5281 let schema = session_schema();
5282 RecordBatch::try_new(
5283 schema.clone(),
5284 vec![
5285 Arc::new(StringArray::from(
5286 sessions
5287 .iter()
5288 .map(|session| session.id.as_str())
5289 .collect::<Vec<_>>(),
5290 )),
5291 Arc::new(StringArray::from(
5292 sessions
5293 .iter()
5294 .map(|session| session.parent_session_id.as_deref())
5295 .collect::<Vec<_>>(),
5296 )),
5297 Arc::new(StringArray::from(
5298 sessions
5299 .iter()
5300 .map(|session| session.parent_message_id.as_deref())
5301 .collect::<Vec<_>>(),
5302 )),
5303 Arc::new(StringArray::from(
5304 sessions
5305 .iter()
5306 .map(|session| session.source_agent.as_str())
5307 .collect::<Vec<_>>(),
5308 )),
5309 Arc::new(
5310 TimestampMicrosecondArray::from(
5311 sessions
5312 .iter()
5313 .map(|session| micros(session.created_at))
5314 .collect::<Vec<_>>(),
5315 )
5316 .with_timezone("UTC"),
5317 ),
5318 Arc::new(StringArray::from(
5319 sessions
5320 .iter()
5321 .map(|session| session.project.as_str())
5322 .collect::<Vec<_>>(),
5323 )),
5324 Arc::new(LargeBinaryArray::from_iter_values(
5325 options.iter().map(Vec::as_slice),
5326 )),
5327 ],
5328 )
5329 .context("failed to build session batch")
5330}
5331
5332pub(crate) fn messages_batches(
5335 rows: &[MessageBatchRow<'_>],
5336 vectors: &[Option<Vec<f32>>],
5337) -> Result<Vec<RecordBatch>> {
5338 debug_assert_eq!(rows.len(), vectors.len(), "vectors must align with rows");
5339 let options = rows
5340 .iter()
5341 .map(|row| json_bytes(row.message.options()))
5342 .collect::<Result<Vec<_>>>()?;
5343 let mut cells = Vec::with_capacity(rows.len());
5344 for (row, encoded) in rows.iter().zip(&options) {
5345 let columns = [
5346 row.message.session_id().len(),
5347 row.message.id().len(),
5348 row.message.role().as_str().len(),
5349 row.source_agent.len(),
5350 row.project.len(),
5351 row.message.system_content().map_or(0, str::len),
5352 row.search_text.map_or(0, str::len),
5353 encoded.len(),
5354 ];
5355 for bytes in columns {
5356 guard_cell("messages", row.message.id(), bytes)?;
5357 }
5358 cells.push(columns.iter().sum());
5359 }
5360 chunk_ranges(&cells)
5361 .into_iter()
5362 .map(|range| {
5363 messages_chunk(
5364 &rows[range.clone()],
5365 &options[range.clone()],
5366 &vectors[range],
5367 )
5368 })
5369 .collect()
5370}
5371
5372fn messages_chunk(
5373 rows: &[MessageBatchRow<'_>],
5374 options: &[Vec<u8>],
5375 vectors: &[Option<Vec<f32>>],
5376) -> Result<RecordBatch> {
5377 let schema = message_schema();
5378 let (vector_column, embedding_model) = embedding_columns(vectors)?;
5379 RecordBatch::try_new(
5380 schema.clone(),
5381 vec![
5382 Arc::new(StringArray::from(
5383 rows.iter()
5384 .map(|row| row.message.session_id())
5385 .collect::<Vec<_>>(),
5386 )),
5387 Arc::new(StringArray::from(
5388 rows.iter().map(|row| row.message.id()).collect::<Vec<_>>(),
5389 )),
5390 Arc::new(
5391 TimestampMicrosecondArray::from(
5392 rows.iter()
5393 .map(|row| micros(row.message.timestamp()))
5394 .collect::<Vec<_>>(),
5395 )
5396 .with_timezone("UTC"),
5397 ),
5398 Arc::new(StringArray::from(
5399 rows.iter()
5400 .map(|row| row.message.role().as_str())
5401 .collect::<Vec<_>>(),
5402 )),
5403 Arc::new(StringArray::from(
5404 rows.iter().map(|row| row.source_agent).collect::<Vec<_>>(),
5405 )),
5406 Arc::new(StringArray::from(
5407 rows.iter().map(|row| row.project).collect::<Vec<_>>(),
5408 )),
5409 Arc::new(StringArray::from(
5410 rows.iter()
5411 .map(|row| row.message.system_content())
5412 .collect::<Vec<_>>(),
5413 )),
5414 Arc::new(StringArray::from(
5415 rows.iter().map(|row| row.search_text).collect::<Vec<_>>(),
5416 )),
5417 vector_column,
5422 embedding_model,
5423 Arc::new(LargeBinaryArray::from_iter_values(
5424 options.iter().map(Vec::as_slice),
5425 )),
5426 ],
5427 )
5428 .context("failed to build message batch")
5429}
5430
5431pub(crate) fn parts_batches(parts: &[Part]) -> Result<Vec<RecordBatch>> {
5432 let variant_data = parts
5433 .iter()
5434 .map(|part| part_variant_json(&part.kind))
5435 .collect::<Result<Vec<_>>>()?;
5436 let options = parts
5437 .iter()
5438 .map(|part| json_bytes(&part.options))
5439 .collect::<Result<Vec<_>>>()?;
5440 let mut cells = Vec::with_capacity(parts.len());
5441 for ((part, variant), encoded) in parts.iter().zip(&variant_data).zip(&options) {
5444 let columns = [
5445 part.session_id.len(),
5446 part.message_id.len(),
5447 part.id.len(),
5448 part.kind.type_name().len(),
5449 part.provenance.as_str().len(),
5450 variant.len(),
5451 encoded.len(),
5452 ];
5453 for bytes in columns {
5454 guard_cell("parts", &part.id, bytes)?;
5455 }
5456 cells.push(columns.iter().sum());
5457 }
5458 chunk_ranges(&cells)
5459 .into_iter()
5460 .map(|range| {
5461 parts_chunk(
5462 &parts[range.clone()],
5463 &variant_data[range.clone()],
5464 &options[range],
5465 )
5466 })
5467 .collect()
5468}
5469
5470fn parts_chunk(
5471 parts: &[Part],
5472 variant_data: &[Vec<u8>],
5473 options: &[Vec<u8>],
5474) -> Result<RecordBatch> {
5475 let schema = part_schema();
5476 let blob_payloads: Vec<Option<&[u8]>> = parts
5480 .iter()
5481 .map(|part| match &part.kind {
5482 PartKind::File { data, .. } => Some(match data {
5483 FileData::String(value) => value.as_bytes(),
5484 FileData::Bytes(value) => value.as_slice(),
5485 FileData::Url(value) => value.as_bytes(),
5486 }),
5487 PartKind::Text { .. }
5488 | PartKind::Reasoning { .. }
5489 | PartKind::ToolCall { .. }
5490 | PartKind::ToolResult { .. }
5491 | PartKind::ToolApprovalRequest { .. }
5492 | PartKind::ToolApprovalResponse { .. } => None,
5493 })
5494 .collect();
5495 let blob_array = LargeBinaryArray::from_iter(blob_payloads);
5496
5497 let mut tool_names: Vec<Option<&str>> = Vec::with_capacity(parts.len());
5498 let mut call_ids: Vec<Option<&str>> = Vec::with_capacity(parts.len());
5499 let mut failures: Vec<Option<bool>> = Vec::with_capacity(parts.len());
5500 for part in parts {
5501 let (name, call_id, is_failure) = tool_identity(&part.kind);
5502 tool_names.push(name);
5503 call_ids.push(call_id);
5504 failures.push(is_failure);
5505 }
5506
5507 RecordBatch::try_new(
5508 schema.clone(),
5509 vec![
5510 Arc::new(StringArray::from(
5511 parts
5512 .iter()
5513 .map(|part| part.session_id.as_str())
5514 .collect::<Vec<_>>(),
5515 )),
5516 Arc::new(StringArray::from(
5517 parts
5518 .iter()
5519 .map(|part| part.message_id.as_str())
5520 .collect::<Vec<_>>(),
5521 )),
5522 Arc::new(StringArray::from(
5523 parts
5524 .iter()
5525 .map(|part| part.id.as_str())
5526 .collect::<Vec<_>>(),
5527 )),
5528 Arc::new(Int32Array::from(
5529 parts.iter().map(|part| part.ordinal).collect::<Vec<_>>(),
5530 )),
5531 Arc::new(StringArray::from(
5532 parts
5533 .iter()
5534 .map(|part| part.kind.type_name())
5535 .collect::<Vec<_>>(),
5536 )),
5537 Arc::new(StringArray::from(
5538 parts
5539 .iter()
5540 .map(|part| part.provenance.as_str())
5541 .collect::<Vec<_>>(),
5542 )),
5543 Arc::new(StringArray::from(tool_names)),
5544 Arc::new(StringArray::from(call_ids)),
5545 Arc::new(BooleanArray::from(failures)),
5546 Arc::new(LargeBinaryArray::from_iter_values(
5547 variant_data.iter().map(Vec::as_slice),
5548 )),
5549 Arc::new(blob_array),
5550 Arc::new(LargeBinaryArray::from_iter_values(
5551 options.iter().map(Vec::as_slice),
5552 )),
5553 ],
5554 )
5555 .context("failed to build parts batch")
5556}
5557
5558pub(crate) fn session_from_batch(batch: &RecordBatch, row: usize) -> Result<Session> {
5559 Ok(Session {
5560 id: string(batch, "id", row)?.context("session id is null")?,
5561 parent_session_id: string(batch, "parent_session_id", row)?,
5562 parent_message_id: string(batch, "parent_message_id", row)?,
5563 source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5564 created_at: datetime(batch, "created_at", row)?,
5565 project: crate::adapter::Extracted::from_stored(
5566 string(batch, "project", row)?.context("project is null")?,
5567 ),
5568 options: json_parse(&json_column(batch, "options", row)?.context("options is null")?)?,
5569 })
5570}
5571
5572pub struct RowmapOracle(pub Option<Arc<RowMetaSet>>);
5579
5580impl crate::adapter::SkipOracle for RowmapOracle {
5581 fn session_max_ts(&self, session_id: &str) -> Option<i64> {
5582 self.0.as_ref()?.lookup_max_ts(session_id)
5583 }
5584
5585 fn is_empty(&self) -> bool {
5586 self.0.as_ref().is_none_or(|set| set.is_empty())
5587 }
5588}
5589
5590fn row_meta_entry(batch: &RecordBatch, row_id: u64, row: usize) -> Result<RowMetaEntry> {
5591 Ok(RowMetaEntry {
5592 row_id,
5593 session_id: string(batch, "session_id", row)?.context("session_id is null")?,
5594 message_id: string(batch, "id", row)?.context("message id is null")?,
5595 role: string(batch, "role", row)?.context("role is null")?,
5596 project: string(batch, "project", row)?.context("project is null")?,
5597 source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5598 timestamp_micros: datetime(batch, "timestamp", row)?.timestamp_micros(),
5599 search_text: string(batch, "search_text", row)?.unwrap_or_default(),
5600 })
5601}
5602
5603pub(crate) fn message_meta_from_batch(batch: &RecordBatch, row: usize) -> Result<MessageMeta> {
5604 Ok(MessageMeta {
5605 message_id: string(batch, "id", row)?.context("id is null")?,
5606 session_id: string(batch, "session_id", row)?.context("session_id is null")?,
5607 role: string(batch, "role", row)?.context("role is null")?,
5608 project: string(batch, "project", row)?.context("project is null")?,
5609 source_agent: string(batch, "source_agent", row)?.context("source_agent is null")?,
5610 timestamp: datetime(batch, "timestamp", row)?,
5611 search_text: string(batch, "search_text", row)?.unwrap_or_default(),
5612 })
5613}
5614
5615pub(crate) fn message_from_batch(batch: &RecordBatch, row: usize) -> Result<Message> {
5616 let id = string(batch, "id", row)?.context("message id is null")?;
5617 let session_id = string(batch, "session_id", row)?.context("message session_id is null")?;
5618 let timestamp = datetime(batch, "timestamp", row)?;
5619 let options =
5620 json_parse(&json_column(batch, "options", row)?.context("message options is null")?)?;
5621
5622 match string(batch, "role", row)?
5623 .context("message role is null")?
5624 .as_str()
5625 {
5626 "system" => Ok(Message::System {
5627 id,
5628 session_id,
5629 timestamp,
5630 content: string(batch, "content", row)?.map(crate::adapter::Extracted::from_stored),
5637 options,
5638 }),
5639 "user" => Ok(Message::User {
5640 id,
5641 session_id,
5642 timestamp,
5643 options,
5644 }),
5645 "assistant" => Ok(Message::Assistant {
5646 id,
5647 session_id,
5648 timestamp,
5649 options,
5650 }),
5651 "tool" => Ok(Message::Tool {
5652 id,
5653 session_id,
5654 timestamp,
5655 options,
5656 }),
5657 other => anyhow::bail!("unknown message role {other}"),
5658 }
5659}
5660
5661pub(crate) fn part_from_batch(
5662 batch: &RecordBatch,
5663 row: usize,
5664 file_data: Option<FileData>,
5665) -> Result<Part> {
5666 let type_name = string(batch, "type", row)?.context("part type is null")?;
5667 let variant_data = json_column(batch, "variant_data", row)?.context("variant_data is null")?;
5668 let provenance = string(batch, "provenance", row)?.context("part provenance is null")?;
5669 Ok(Part {
5670 session_id: string(batch, "session_id", row)?.context("part session_id is null")?,
5671 message_id: string(batch, "message_id", row)?.context("part message_id is null")?,
5672 id: string(batch, "id", row)?.context("part id is null")?,
5673 ordinal: int32(batch, "ordinal", row)?,
5674 provenance: provenance_from_str(&provenance)?,
5675 options: json_parse(&json_column(batch, "options", row)?.context("part options is null")?)?,
5676 kind: part_kind_from_json(&type_name, &variant_data, file_data)?,
5677 })
5678}
5679
5680fn provenance_from_str(value: &str) -> Result<crate::wire::Provenance> {
5681 match value {
5682 "conversational" => Ok(crate::wire::Provenance::Conversational),
5683 "injected" => Ok(crate::wire::Provenance::Injected),
5684 other => anyhow::bail!("unknown part provenance {other}"),
5685 }
5686}
5687
5688fn file_data_from_blob(variant_data: &[u8], bytes: &[u8]) -> Result<FileData> {
5689 let kind = file_data_kind(variant_data)?;
5690 match kind.as_str() {
5691 "string" => {
5692 let text = std::str::from_utf8(bytes)
5693 .context("file string payload is not UTF-8")?
5694 .to_owned();
5695 Ok(FileData::String(text))
5696 }
5697 "bytes" => Ok(FileData::Bytes(bytes.to_vec())),
5698 "url" => Ok(FileData::Url(
5699 std::str::from_utf8(bytes)
5700 .context("file URL payload is not UTF-8")?
5701 .to_owned(),
5702 )),
5703 other => anyhow::bail!("unknown file data_kind {other}"),
5704 }
5705}
5706
5707fn file_data_kind(variant_data: &[u8]) -> Result<String> {
5708 let value = json_parse::<Value>(variant_data)?;
5709 value
5710 .get("data_kind")
5711 .and_then(Value::as_str)
5712 .map(str::to_owned)
5713 .context("file part variant_data missing data_kind")
5714}
5715
5716fn uint64<'a>(batch: &'a RecordBatch, name: &str) -> Result<&'a UInt64Array> {
5717 batch
5718 .column_by_name(name)
5719 .with_context(|| format!("missing column {name}"))?
5720 .as_any()
5721 .downcast_ref::<UInt64Array>()
5722 .with_context(|| format!("column {name} is not UInt64"))
5723}
5724
5725pub(crate) fn string(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<String>> {
5726 let array = batch
5727 .column_by_name(name)
5728 .with_context(|| format!("missing column {name}"))?
5729 .as_any()
5730 .downcast_ref::<StringArray>()
5731 .with_context(|| format!("column {name} is not Utf8"))?;
5732 if array.is_null(row) {
5733 Ok(None)
5734 } else {
5735 Ok(Some(array.value(row).to_owned()))
5736 }
5737}
5738
5739fn json_column(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<Vec<u8>>> {
5740 let column = batch
5744 .column_by_name(name)
5745 .with_context(|| format!("missing column {name}"))?;
5746 if let Some(array) = column.as_any().downcast_ref::<LargeBinaryArray>() {
5747 return if array.is_null(row) {
5748 Ok(None)
5749 } else {
5750 Ok(Some(
5751 lance_arrow::json::decode_json(array.value(row)).into_bytes(),
5752 ))
5753 };
5754 }
5755 if let Some(array) = column.as_any().downcast_ref::<StringArray>() {
5756 return if array.is_null(row) {
5757 Ok(None)
5758 } else {
5759 Ok(Some(array.value(row).as_bytes().to_vec()))
5760 };
5761 }
5762 if let Some(array) = column.as_any().downcast_ref::<LargeStringArray>() {
5763 return if array.is_null(row) {
5764 Ok(None)
5765 } else {
5766 Ok(Some(array.value(row).as_bytes().to_vec()))
5767 };
5768 }
5769 anyhow::bail!("column {name} is not a JSON-compatible array")
5770}
5771
5772fn int32(batch: &RecordBatch, name: &str, row: usize) -> Result<i32> {
5773 let array = batch
5774 .column_by_name(name)
5775 .with_context(|| format!("missing column {name}"))?
5776 .as_any()
5777 .downcast_ref::<Int32Array>()
5778 .with_context(|| format!("column {name} is not Int32"))?;
5779 Ok(array.value(row))
5780}
5781
5782pub(crate) fn float32(batch: &RecordBatch, name: &str, row: usize) -> Result<f32> {
5783 let array = batch
5784 .column_by_name(name)
5785 .with_context(|| format!("missing column {name}"))?
5786 .as_any()
5787 .downcast_ref::<Float32Array>()
5788 .with_context(|| format!("column {name} is not Float32"))?;
5789 Ok(array.value(row))
5790}
5791
5792pub(crate) fn datetime(batch: &RecordBatch, name: &str, row: usize) -> Result<DateTime<Utc>> {
5793 let array = batch
5794 .column_by_name(name)
5795 .with_context(|| format!("missing column {name}"))?
5796 .as_any()
5797 .downcast_ref::<TimestampMicrosecondArray>()
5798 .with_context(|| format!("column {name} is not timestamp_micros"))?;
5799 Utc.timestamp_micros(array.value(row))
5800 .single()
5801 .context("timestamp is out of range")
5802}
5803
5804fn primary_field(name: &str, data_type: DataType, nullable: bool) -> Field {
5805 Field::new(name, data_type, nullable).with_metadata(
5806 [(
5807 "lance-schema:unenforced-primary-key".to_owned(),
5808 "true".to_owned(),
5809 )]
5810 .into(),
5811 )
5812}
5813
5814fn legacy_blob_field(name: &str, nullable: bool) -> Field {
5824 Field::new(name, DataType::LargeBinary, nullable).with_metadata(
5825 [(lance_arrow::BLOB_META_KEY.to_owned(), "true".to_owned())]
5826 .into_iter()
5827 .collect(),
5828 )
5829}
5830
5831fn json_field(name: &str, nullable: bool) -> Field {
5838 lance_arrow::json::json_field(name, nullable)
5839}
5840
5841fn micros(timestamp: DateTime<Utc>) -> i64 {
5842 timestamp.timestamp_micros()
5843}
5844
5845fn json_bytes<T: Serialize>(value: &T) -> Result<Vec<u8>> {
5846 let text = serde_json::to_string(value).context("failed to serialize JSON field")?;
5854 lance_arrow::json::encode_json(&text)
5855 .map_err(|err| anyhow::anyhow!("failed to encode JSON field as JSONB: {err}"))
5856}
5857
5858fn json_parse<T: DeserializeOwned>(value: &[u8]) -> Result<T> {
5859 serde_json::from_slice(value).context("failed to parse JSON field")
5860}
5861
5862fn tool_identity(kind: &PartKind) -> (Option<&str>, Option<&str>, Option<bool>) {
5869 match kind {
5870 PartKind::ToolCall { name, call_id, .. } => (
5871 name.as_deref().map(String::as_str),
5872 call_id.as_deref().map(String::as_str),
5873 None,
5874 ),
5875 PartKind::ToolResult {
5876 name,
5877 call_id,
5878 is_failure,
5879 ..
5880 } => (
5881 name.as_deref().map(String::as_str),
5882 call_id.as_deref().map(String::as_str),
5883 Some(*is_failure),
5884 ),
5885 PartKind::ToolApprovalRequest { tool_call_id, .. } => {
5886 (None, Some(tool_call_id.as_str()), None)
5887 }
5888 PartKind::Text { .. }
5889 | PartKind::Reasoning { .. }
5890 | PartKind::File { .. }
5891 | PartKind::ToolApprovalResponse { .. } => (None, None, None),
5892 }
5893}
5894
5895fn part_variant_json(kind: &PartKind) -> Result<Vec<u8>> {
5896 if let PartKind::File {
5897 media_type,
5898 file_name,
5899 data,
5900 } = kind
5901 {
5902 let data_kind = match data {
5903 FileData::String(_) => "string",
5904 FileData::Bytes(_) => "bytes",
5905 FileData::Url(_) => "url",
5906 };
5907 return json_bytes(&serde_json::json!({
5908 "media_type": media_type,
5909 "file_name": file_name,
5910 "data_kind": data_kind,
5911 }));
5912 }
5913 let value = serde_json::to_value(kind)?;
5914 let mut object = value
5915 .as_object()
5916 .cloned()
5917 .context("part variant did not serialize to an object")?;
5918 object.remove("type");
5919 json_bytes(&object)
5920}
5921
5922fn part_kind_from_json(
5923 type_name: &str,
5924 variant_data: &[u8],
5925 file_data: Option<FileData>,
5926) -> Result<PartKind> {
5927 let mut value = json_parse::<Value>(variant_data)?;
5928 let object = value
5929 .as_object_mut()
5930 .context("part variant data is not an object")?;
5931 object.insert("type".to_owned(), Value::String(type_name.to_owned()));
5932 if let Some(data) = file_data {
5933 object.remove("data_kind");
5934 object.insert("data".to_owned(), serde_json::to_value(data)?);
5935 }
5936 serde_json::from_value(value).context("failed to parse part kind")
5937}
5938
5939#[cfg(test)]
5940mod tests {
5941 #![allow(clippy::expect_used, clippy::unwrap_used)]
5942
5943 use super::*;
5944 use crate::{
5945 adapter::Extracted,
5946 handlers::ingest_events,
5947 wire::{FileData, Message, Part, PartKind, ProviderOptions, Session},
5948 };
5949 use chrono::Utc;
5950 use serde_json::json;
5951 use tempfile::TempDir;
5952
5953 fn synthetic_session(id: &str) -> Session {
5954 Session {
5955 id: id.to_owned(),
5956 parent_session_id: None,
5957 parent_message_id: None,
5958 source_agent: "claude-code".to_owned(),
5959 created_at: Utc::now(),
5960 project: crate::adapter::Extracted::from_test_value("/tmp/pond".to_owned()),
5961 options: ProviderOptions::new(),
5962 }
5963 }
5964
5965 #[derive(Default)]
5968 struct CountingEmbedder {
5969 texts: std::sync::atomic::AtomicUsize,
5970 }
5971 impl crate::embed::Embedder for CountingEmbedder {
5972 fn device(&self) -> &str {
5973 "test"
5974 }
5975 fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
5976 self.texts
5977 .fetch_add(texts.len(), std::sync::atomic::Ordering::SeqCst);
5978 Ok(texts
5979 .iter()
5980 .map(|_| vec![0.0_f32; embedding_dim()])
5981 .collect())
5982 }
5983 }
5984
5985 fn conversational_events(session_id: &str, count: usize) -> Vec<IngestEvent> {
5988 let mut events = vec![IngestEvent::Session(synthetic_session(session_id))];
5989 for index in 0..count {
5990 events.push(IngestEvent::Message(Message::User {
5991 id: format!("msg-{index}"),
5992 session_id: session_id.to_owned(),
5993 timestamp: Utc::now(),
5994 options: ProviderOptions::new(),
5995 }));
5996 events.push(IngestEvent::Part(Part {
5997 session_id: session_id.to_owned(),
5998 id: format!("msg-{index}:0001"),
5999 message_id: format!("msg-{index}"),
6000 ordinal: 0,
6001 provenance: crate::wire::Provenance::Conversational,
6002 options: ProviderOptions::new(),
6003 kind: PartKind::Text {
6004 text: Some(Extracted::from_test_value(format!("body {index}"))),
6005 },
6006 }));
6007 }
6008 events
6009 }
6010
6011 #[tokio::test(flavor = "multi_thread")]
6015 async fn ingest_embeds_inline_in_the_birth_commit() -> anyhow::Result<()> {
6016 let plain = Store::open(&Url::parse("shared-memory://pond-test-inline-plain/")?).await?;
6017 ingest_events(&plain, conversational_events("01HXYINLINE000PLAIN", 5)).await?;
6018 assert!(
6019 !plain.has_embeddings().await?,
6020 "no embedder attached -> every vector null",
6021 );
6022 let plain_version = plain.messages_version().await?;
6023
6024 let backend = Arc::new(CountingEmbedder::default());
6025 let embedder = Arc::new(crate::embed::LazyEmbedder::from_loaded(
6026 backend.clone() as Arc<dyn crate::embed::Embedder>
6027 ));
6028 let store = Store::open(&Url::parse("shared-memory://pond-test-inline-embed/")?)
6029 .await?
6030 .with_embedder(embedder);
6031 ingest_events(&store, conversational_events("01HXYINLINE000EMBED", 5)).await?;
6032 assert!(
6033 store.has_embeddings().await?,
6034 "embedder attached -> vectors filled at ingest",
6035 );
6036 assert_eq!(
6037 backend.texts.load(std::sync::atomic::Ordering::SeqCst),
6038 5,
6039 "every conversational message embedded once",
6040 );
6041 assert_eq!(
6042 store.messages_version().await?,
6043 plain_version,
6044 "inline embed must ride the append, not add a separate commit",
6045 );
6046
6047 ingest_events(&store, conversational_events("01HXYINLINE000EMBED", 5)).await?;
6048 assert_eq!(
6049 backend.texts.load(std::sync::atomic::Ordering::SeqCst),
6050 5,
6051 "an idempotent re-sync embeds no already-present row",
6052 );
6053 Ok(())
6054 }
6055
6056 #[tokio::test(flavor = "multi_thread")]
6062 async fn composite_pk_index_counts_duplicates_by_composite_key() -> anyhow::Result<()> {
6063 async fn duplicates(store: &Store, table: Table) -> anyhow::Result<usize> {
6064 let (keys, rows) = store.composite_pk_index(table).await?;
6065 Ok(rows - keys.len())
6066 }
6067 let store = Store::open(&Url::parse("shared-memory://pond-test-dupcount/")?).await?;
6068 ingest_events(&store, conversational_events("01HXYDUP00000SESS1", 1)).await?;
6069 ingest_events(&store, conversational_events("01HXYDUP00000SESS2", 1)).await?;
6070 assert_eq!(
6071 duplicates(&store, Table::Messages).await?,
6072 0,
6073 "the same message id in two sessions is not a duplicate (composite PK)",
6074 );
6075 assert_eq!(duplicates(&store, Table::Sessions).await?, 0);
6076 assert_eq!(duplicates(&store, Table::Parts).await?, 0);
6077
6078 let message = Message::User {
6081 id: "dup-msg".to_owned(),
6082 session_id: "01HXYDUP00000SESS1".to_owned(),
6083 timestamp: Utc::now(),
6084 options: ProviderOptions::new(),
6085 };
6086 let row = MessageBatchRow {
6087 message: &message,
6088 source_agent: "claude-code",
6089 project: "/tmp",
6090 search_text: None,
6091 };
6092 let batches = messages_batches(&[row], &[None])?;
6093 store
6094 .handle
6095 .append_batches(Table::Messages, batches.clone())
6096 .await?;
6097 store
6098 .handle
6099 .append_batches(Table::Messages, batches)
6100 .await?;
6101 assert_eq!(
6102 duplicates(&store, Table::Messages).await?,
6103 1,
6104 "the doubled (session_id, id) row is exactly one duplicate",
6105 );
6106 Ok(())
6107 }
6108
6109 #[test]
6110 fn search_text_excludes_injected_parts() {
6111 use crate::wire::Provenance;
6112 let message = Message::User {
6113 id: "m1".to_owned(),
6114 session_id: "s1".to_owned(),
6115 timestamp: Utc::now(),
6116 options: ProviderOptions::new(),
6117 };
6118 let text_part = |id: &str, text: &str, provenance: Provenance| Part {
6119 session_id: "s1".to_owned(),
6120 id: id.to_owned(),
6121 message_id: "m1".to_owned(),
6122 ordinal: 0,
6123 provenance,
6124 options: ProviderOptions::new(),
6125 kind: PartKind::Text {
6126 text: Some(Extracted::from_test_value(text.to_owned())),
6127 },
6128 };
6129
6130 let conversational = search_text(
6133 &message,
6134 &[text_part(
6135 "p1",
6136 "real human prompt",
6137 Provenance::Conversational,
6138 )],
6139 );
6140 assert_eq!(conversational.as_deref(), Some("real human prompt"));
6141
6142 let injected = search_text(
6143 &message,
6144 &[text_part(
6145 "p2",
6146 "<task-notification>...</task-notification>",
6147 Provenance::Injected,
6148 )],
6149 );
6150 assert!(
6151 injected.is_none(),
6152 "a message whose only part is injected has null search_text"
6153 );
6154 }
6155
6156 #[test]
6157 fn parts_chunk_materializes_tool_identity_columns() -> anyhow::Result<()> {
6158 let part = |ordinal: i32, kind: PartKind| Part {
6159 session_id: "s1".to_owned(),
6160 id: format!("p{ordinal}"),
6161 message_id: "m1".to_owned(),
6162 ordinal,
6163 provenance: crate::wire::Provenance::Conversational,
6164 options: ProviderOptions::new(),
6165 kind,
6166 };
6167 let parts = vec![
6168 part(
6169 0,
6170 PartKind::ToolCall {
6171 call_id: Some(Extracted::from_test_value("c1".to_owned())),
6172 name: Some(Extracted::from_test_value("Bash".to_owned())),
6173 params: serde_json::json!({}),
6174 provider_executed: false,
6175 },
6176 ),
6177 part(
6178 1,
6179 PartKind::ToolResult {
6180 call_id: Some(Extracted::from_test_value("c1".to_owned())),
6181 name: Some(Extracted::from_test_value("Bash".to_owned())),
6182 is_failure: true,
6183 result: serde_json::json!({}),
6184 },
6185 ),
6186 part(
6188 2,
6189 PartKind::ToolCall {
6190 call_id: None,
6191 name: None,
6192 params: serde_json::json!({}),
6193 provider_executed: false,
6194 },
6195 ),
6196 part(
6197 3,
6198 PartKind::ToolApprovalRequest {
6199 approval_id: "a1".to_owned(),
6200 tool_call_id: "c1".to_owned(),
6201 },
6202 ),
6203 part(
6204 4,
6205 PartKind::Text {
6206 text: Some(Extracted::from_test_value("hi".to_owned())),
6207 },
6208 ),
6209 ];
6210 let payloads = vec![vec![1u8]; parts.len()];
6211 let batch = parts_chunk(&parts, &payloads, &payloads)?;
6212
6213 let tool_name = |row| string(&batch, "tool_name", row).unwrap();
6214 let call_id = |row| string(&batch, "call_id", row).unwrap();
6215 let is_failure = batch
6216 .column_by_name("is_failure")
6217 .expect("is_failure column present")
6218 .as_any()
6219 .downcast_ref::<BooleanArray>()
6220 .expect("is_failure is boolean");
6221
6222 assert_eq!(tool_name(0).as_deref(), Some("Bash"));
6223 assert_eq!(call_id(0).as_deref(), Some("c1"));
6224 assert!(is_failure.is_null(0), "tool_call has no failure flag");
6225
6226 assert_eq!(tool_name(1).as_deref(), Some("Bash"));
6227 assert_eq!(call_id(1).as_deref(), Some("c1"));
6228 assert!(is_failure.value(1), "tool_result failure flag materialized");
6229
6230 assert_eq!(tool_name(2), None, "absent name stays NULL");
6231 assert_eq!(call_id(2), None, "absent call_id stays NULL");
6232
6233 assert_eq!(tool_name(3), None);
6234 assert_eq!(
6235 call_id(3).as_deref(),
6236 Some("c1"),
6237 "approval request surfaces its tool_call_id as call_id",
6238 );
6239
6240 assert_eq!(tool_name(4), None);
6241 assert_eq!(call_id(4), None);
6242 assert!(is_failure.is_null(4));
6243 Ok(())
6244 }
6245
6246 #[test]
6247 fn chunk_ranges_splits_on_byte_budget() {
6248 assert!(chunk_ranges(&[]).is_empty());
6249 assert_eq!(chunk_ranges(&[10, 10, 10]), vec![0..3]);
6250
6251 let two_thirds = COLUMN_BYTE_BUDGET * 2 / 3;
6252 assert_eq!(
6253 chunk_ranges(&[two_thirds, two_thirds, two_thirds]),
6254 vec![0..1, 1..2, 2..3],
6255 );
6256
6257 assert_eq!(
6259 chunk_ranges(&[10, COLUMN_BYTE_BUDGET + 1, 10]),
6260 vec![0..1, 1..2, 2..3],
6261 );
6262 }
6263
6264 #[tokio::test]
6265 async fn ordering_violation_drops_only_the_offending_event() -> anyhow::Result<()> {
6266 let temp = TempDir::new()?;
6271 let store = Store::open_local(temp.path()).await?;
6272 let session = synthetic_session("ordering");
6273 let orphan_part = Part {
6274 session_id: session.id.clone(),
6275 id: "orphan-part".to_owned(),
6276 message_id: "missing-message".to_owned(),
6277 ordinal: 0,
6278 provenance: crate::wire::Provenance::Conversational,
6279 options: ProviderOptions::new(),
6280 kind: PartKind::Text {
6281 text: Some(Extracted::from_test_value("orphan".to_owned())),
6282 },
6283 };
6284 let valid_message = Message::User {
6285 id: "valid-message".to_owned(),
6286 session_id: session.id.clone(),
6287 timestamp: Utc::now(),
6288 options: ProviderOptions::new(),
6289 };
6290 let valid_part = Part {
6291 session_id: session.id.clone(),
6292 id: "valid-part".to_owned(),
6293 message_id: valid_message.id().to_owned(),
6294 ordinal: 0,
6295 provenance: crate::wire::Provenance::Conversational,
6296 options: ProviderOptions::new(),
6297 kind: PartKind::Text {
6298 text: Some(Extracted::from_test_value("kept".to_owned())),
6299 },
6300 };
6301
6302 let mut validator = IngestValidator::default();
6303 validator
6304 .push(&store, 0, IngestEvent::Session(session.clone()))
6305 .await?;
6306 let part_outcomes = validator
6307 .push(&store, 1, IngestEvent::Part(orphan_part))
6308 .await?;
6309 assert_eq!(part_outcomes.len(), 1);
6310 assert_eq!(part_outcomes[0].kind, "part");
6311 assert_eq!(part_outcomes[0].status, OutcomeStatus::Error);
6312 assert!(
6313 part_outcomes[0]
6314 .error
6315 .as_ref()
6316 .map(|e| e.message.contains("part event appeared before a message"))
6317 .unwrap_or(false),
6318 "error message must explain the ordering violation: {part_outcomes:?}"
6319 );
6320 validator
6321 .push(&store, 2, IngestEvent::Message(valid_message))
6322 .await?;
6323 validator
6324 .push(&store, 3, IngestEvent::Part(valid_part))
6325 .await?;
6326 validator.finish(&store).await?;
6327
6328 let (sessions, messages, parts) = store.row_counts().await?;
6329 assert_eq!(sessions, 1, "session committed despite the orphan part");
6330 assert_eq!(messages, 1, "valid message committed");
6331 assert_eq!(parts, 1, "valid part committed; the orphan was dropped");
6332
6333 Ok(())
6334 }
6335
6336 #[tokio::test]
6337 async fn resident_meta_map_hydration_matches_take_rows_fallback() -> anyhow::Result<()> {
6338 let temp = TempDir::new()?;
6342 let store = Store::open_local(temp.path()).await?;
6343 let session = synthetic_session("hydration-parity");
6344
6345 let messages = [
6346 (
6347 "m1",
6348 "the auth refactor landed cleanly",
6349 1_700_000_000_123_456_i64,
6350 ),
6351 (
6352 "m2",
6353 "balance handler now retries on rpc timeout",
6354 1_700_000_050_654_321,
6355 ),
6356 ];
6357 let mut validator = IngestValidator::default();
6358 validator
6359 .push(&store, 0, IngestEvent::Session(session.clone()))
6360 .await?;
6361 let mut seq = 1;
6362 for (mid, text, micros) in messages {
6363 let message = Message::User {
6364 id: mid.to_owned(),
6365 session_id: session.id.clone(),
6366 timestamp: DateTime::from_timestamp_micros(micros).unwrap(),
6367 options: ProviderOptions::new(),
6368 };
6369 validator
6370 .push(&store, seq, IngestEvent::Message(message))
6371 .await?;
6372 seq += 1;
6373 let part = Part {
6374 session_id: session.id.clone(),
6375 id: format!("{mid}-p0"),
6376 message_id: mid.to_owned(),
6377 ordinal: 0,
6378 provenance: crate::wire::Provenance::Conversational,
6379 options: ProviderOptions::new(),
6380 kind: PartKind::Text {
6381 text: Some(Extracted::from_test_value(text.to_owned())),
6382 },
6383 };
6384 validator.push(&store, seq, IngestEvent::Part(part)).await?;
6385 seq += 1;
6386 }
6387 validator.finish(&store).await?;
6388
6389 let rowids: Vec<u64> = store
6390 .collect_row_metas()
6391 .await?
6392 .into_iter()
6393 .map(|entry| entry.row_id)
6394 .collect();
6395 assert_eq!(rowids.len(), 2);
6396
6397 let sort_by_id = |mut metas: Vec<MessageMeta>| {
6398 metas.sort_by(|left, right| left.message_id.cmp(&right.message_id));
6399 metas
6400 };
6401
6402 let fallback = sort_by_id(store.message_metas_by_rowids(&rowids).await?);
6403
6404 store.ensure_rowmap(&temp.path().join("cache")).await?;
6407 let resident = sort_by_id(store.message_metas_by_rowids(&rowids).await?);
6408
6409 assert_eq!(
6410 resident, fallback,
6411 "resident-map hydration must match the take_rows fallback"
6412 );
6413 assert_eq!(
6414 resident[0].timestamp.timestamp_micros(),
6415 1_700_000_000_123_456
6416 );
6417 Ok(())
6418 }
6419
6420 #[tokio::test]
6421 async fn initialized_flips_only_after_first_ingest() -> anyhow::Result<()> {
6422 let temp = TempDir::new()?;
6427 let store = Store::open_local(temp.path()).await?;
6428 assert!(
6429 !store.initialized().await?,
6430 "fresh store has no parts table"
6431 );
6432
6433 let session = synthetic_session("initialized-probe");
6434 let message = Message::User {
6435 id: "message-1".to_owned(),
6436 session_id: session.id.clone(),
6437 timestamp: Utc::now(),
6438 options: ProviderOptions::new(),
6439 };
6440 let part = Part {
6441 session_id: session.id.clone(),
6442 id: "part-1".to_owned(),
6443 message_id: message.id().to_owned(),
6444 ordinal: 0,
6445 provenance: crate::wire::Provenance::Conversational,
6446 options: ProviderOptions::new(),
6447 kind: PartKind::Text {
6448 text: Some(Extracted::from_test_value("hello".to_owned())),
6449 },
6450 };
6451 let mut validator = IngestValidator::default();
6452 validator
6453 .push(&store, 0, IngestEvent::Session(session))
6454 .await?;
6455 validator
6456 .push(&store, 1, IngestEvent::Message(message))
6457 .await?;
6458 validator.push(&store, 2, IngestEvent::Part(part)).await?;
6459 validator.finish(&store).await?;
6460
6461 assert!(store.initialized().await?, "ingest creates the parts table");
6462 Ok(())
6463 }
6464
6465 #[tokio::test]
6466 async fn summary_parts_label_a_file_without_reading_its_blob() -> anyhow::Result<()> {
6467 let temp = TempDir::new()?;
6472 let store = Store::open_local(temp.path()).await?;
6473 let session = synthetic_session("file-summary");
6474 let message = Message::User {
6475 id: "m1".to_owned(),
6476 session_id: session.id.clone(),
6477 timestamp: Utc::now(),
6478 options: ProviderOptions::new(),
6479 };
6480 let blob = "file contents the summary must never read";
6481 let part = Part {
6482 session_id: session.id.clone(),
6483 id: "m1-p0".to_owned(),
6484 message_id: "m1".to_owned(),
6485 ordinal: 0,
6486 provenance: crate::wire::Provenance::Conversational,
6487 options: ProviderOptions::new(),
6488 kind: PartKind::File {
6489 media_type: Some("text/plain".to_owned()),
6490 file_name: Some("notes.txt".to_owned()),
6491 data: FileData::String(blob.to_owned()),
6492 },
6493 };
6494 let mut validator = IngestValidator::default();
6495 validator
6496 .push(&store, 0, IngestEvent::Session(session.clone()))
6497 .await?;
6498 validator
6499 .push(&store, 1, IngestEvent::Message(message))
6500 .await?;
6501 validator.push(&store, 2, IngestEvent::Part(part)).await?;
6502 validator.finish(&store).await?;
6503
6504 let key = (session.id.clone(), "m1".to_owned());
6505 let ids = ["m1".to_owned()];
6506
6507 let summarized = store.summary_parts_for_messages(&session.id, &ids).await?;
6508 let summary_part = &summarized.get(&key).expect("file part summarized")[0];
6509 let summary = crate::wire::PartSummary::for_kind(&summary_part.kind)
6510 .expect("a file part yields a summary");
6511 assert_eq!(summary.kind, "file");
6512 assert_eq!(summary.label.as_deref(), Some("notes.txt"));
6513 match &summary_part.kind {
6514 PartKind::File { data, .. } => assert!(
6515 matches!(data, FileData::Bytes(bytes) if bytes.is_empty()),
6516 "summary must carry the empty placeholder, not the file blob",
6517 ),
6518 other => panic!("expected a file part, got {other:?}"),
6519 }
6520
6521 let full = store.parts_for_messages(&session.id, &ids).await?;
6522 match &full.get(&key).expect("file part")[0].kind {
6523 PartKind::File {
6524 data: FileData::String(bytes),
6525 ..
6526 } => assert_eq!(bytes, blob, "full path must still hydrate the real blob"),
6527 other => panic!("expected a string-backed file part, got {other:?}"),
6528 }
6529 Ok(())
6530 }
6531
6532 #[tokio::test]
6533 async fn duplicate_message_id_drops_the_second_keeps_the_first() -> anyhow::Result<()> {
6534 let temp = TempDir::new()?;
6538 let store = Store::open_local(temp.path()).await?;
6539 let session = synthetic_session("duplicate-message");
6540 let first = Message::User {
6541 id: "message-1".to_owned(),
6542 session_id: session.id.clone(),
6543 timestamp: Utc::now(),
6544 options: ProviderOptions::new(),
6545 };
6546 let second = Message::Assistant {
6547 id: "message-1".to_owned(),
6548 session_id: session.id.clone(),
6549 timestamp: Utc::now(),
6550 options: ProviderOptions::new(),
6551 };
6552
6553 let mut validator = IngestValidator::default();
6554 validator
6555 .push(&store, 0, IngestEvent::Session(session.clone()))
6556 .await?;
6557 validator
6558 .push(&store, 1, IngestEvent::Message(first))
6559 .await?;
6560 let dup_outcomes = validator
6561 .push(&store, 2, IngestEvent::Message(second))
6562 .await?;
6563 assert_eq!(dup_outcomes.len(), 1);
6564 assert_eq!(dup_outcomes[0].status, OutcomeStatus::Error);
6565 assert!(
6566 dup_outcomes[0]
6567 .error
6568 .as_ref()
6569 .map(|e| e.message.contains("duplicate message id message-1"))
6570 .unwrap_or(false),
6571 "duplicate-id rejection must name the offending id: {dup_outcomes:?}"
6572 );
6573
6574 validator.finish(&store).await?;
6575 let (sessions, messages, _) = store.row_counts().await?;
6576 assert_eq!(sessions, 1, "session committed");
6577 assert_eq!(messages, 1, "only the first message committed");
6578
6579 Ok(())
6580 }
6581
6582 #[tokio::test]
6583 async fn ingest_stamps_host_provenance_on_messages_and_strips_spoofed_pond_key()
6584 -> anyhow::Result<()> {
6585 let temp = TempDir::new()?;
6589 let store = Store::open_local(temp.path()).await?;
6590 let session = synthetic_session("host-provenance");
6591 let mut spoofed = ProviderOptions::new();
6592 spoofed.insert("pond".to_owned(), json!({"ingest": {"host": "spoofed"}}));
6593 let message = Message::User {
6594 id: "message-1".to_owned(),
6595 session_id: session.id.clone(),
6596 timestamp: Utc::now(),
6597 options: spoofed,
6598 };
6599 let part = Part {
6600 session_id: session.id.clone(),
6601 id: "part-1".to_owned(),
6602 message_id: "message-1".to_owned(),
6603 ordinal: 0,
6604 provenance: crate::wire::Provenance::Conversational,
6605 options: ProviderOptions::new(),
6606 kind: PartKind::Text {
6607 text: Some(Extracted::from_test_value("hello".to_owned())),
6608 },
6609 };
6610
6611 let mut validator = IngestValidator::default();
6612 validator
6613 .push(&store, 0, IngestEvent::Session(session.clone()))
6614 .await?;
6615 validator
6616 .push(&store, 1, IngestEvent::Message(message))
6617 .await?;
6618 validator.push(&store, 2, IngestEvent::Part(part)).await?;
6619 validator.finish(&store).await?;
6620
6621 let stored = store
6622 .get_session(&session.id)
6623 .await?
6624 .expect("ingested session is readable");
6625 assert!(
6626 !stored.session.options.contains_key("pond"),
6627 "session rows are not stamped (attribution derives from messages)"
6628 );
6629 let stored_message = &stored.messages[0].message;
6630 match ingest_host_stamp() {
6631 Some(stamp) => {
6632 assert_eq!(
6633 stored_message.options().get("pond"),
6634 Some(stamp),
6635 "stored message carries the real stamp, never the spoof"
6636 );
6637 let host = stamp
6638 .pointer("/ingest/host")
6639 .and_then(Value::as_object)
6640 .expect("stamp shape is {ingest: {host: {..}}}");
6641 assert!(!host.is_empty(), "an all-empty stamp must be None instead");
6642 assert!(
6643 host.values()
6644 .all(|v| v.as_str().is_some_and(|s| !s.is_empty())),
6645 "stamp fields are omitted when unavailable, never empty: {host:?}"
6646 );
6647 }
6648 None => assert!(
6649 stored_message.options().get("pond").is_none(),
6650 "with no resolvable stamp the spoofed key is still stripped"
6651 ),
6652 }
6653 assert!(
6654 !stored.messages[0].parts[0].options.contains_key("pond"),
6655 "part rows are not stamped (covered by their message's stamp)"
6656 );
6657
6658 Ok(())
6659 }
6660
6661 #[tokio::test(flavor = "multi_thread")]
6669 async fn optimize_indices_compacts_parts_with_blob_column() -> anyhow::Result<()> {
6670 use crate::wire::{FileData, PartKind, Provenance};
6671 let temp = TempDir::new()?;
6672 let store = Store::open_local(temp.path()).await?;
6673
6674 let session = synthetic_session("compact-blob");
6675 store
6676 .upsert_sessions(std::slice::from_ref(&session))
6677 .await?;
6678
6679 let make_part = |idx: usize, kind: PartKind| Part {
6680 session_id: session.id.clone(),
6681 message_id: format!("msg-{idx}"),
6682 id: format!("part-{idx}"),
6683 ordinal: 0,
6684 provenance: Provenance::Conversational,
6685 options: ProviderOptions::new(),
6686 kind,
6687 };
6688
6689 let batch_a = vec![
6690 make_part(
6691 0,
6692 PartKind::File {
6693 media_type: Some("text/plain".to_owned()),
6694 file_name: Some("a.txt".to_owned()),
6695 data: FileData::Bytes(b"alpha".to_vec()),
6696 },
6697 ),
6698 make_part(
6699 1,
6700 PartKind::File {
6701 media_type: Some("text/plain".to_owned()),
6702 file_name: Some("b.txt".to_owned()),
6703 data: FileData::String("beta".to_owned()),
6704 },
6705 ),
6706 ];
6707 store.upsert_parts(&batch_a).await?;
6708
6709 let batch_b = vec![
6710 make_part(
6711 2,
6712 PartKind::File {
6713 media_type: Some("application/octet-stream".to_owned()),
6714 file_name: None,
6715 data: FileData::Url("https://example.com/file".to_owned()),
6716 },
6717 ),
6718 make_part(
6719 3,
6720 PartKind::File {
6721 media_type: Some("image/png".to_owned()),
6722 file_name: Some("c.png".to_owned()),
6723 data: FileData::Bytes(vec![0x89, 0x50, 0x4e, 0x47]),
6724 },
6725 ),
6726 ];
6727 store.upsert_parts(&batch_b).await?;
6728
6729 store
6730 .optimize_indices(None, &MaintenancePolicy::always_compact())
6731 .await?
6732 .into_result()?;
6733
6734 Ok(())
6735 }
6736
6737 #[tokio::test]
6738 async fn file_part_blob_v2_round_trips_through_get() -> anyhow::Result<()> {
6739 let temp = TempDir::new()?;
6740 let store = Store::open_local(temp.path()).await?;
6741 let session = synthetic_session("blob");
6742 let message = Message::User {
6743 id: "message-1".to_owned(),
6744 session_id: session.id.clone(),
6745 timestamp: Utc::now(),
6746 options: ProviderOptions::new(),
6747 };
6748 let part = Part {
6749 session_id: session.id.clone(),
6750 id: "part-1".to_owned(),
6751 message_id: message.id().to_owned(),
6752 ordinal: 0,
6753 provenance: crate::wire::Provenance::Conversational,
6754 options: ProviderOptions::new(),
6755 kind: PartKind::File {
6756 media_type: Some("text/plain".to_owned()),
6757 file_name: Some("payload.txt".to_owned()),
6758 data: FileData::Bytes(b"pond".to_vec()),
6759 },
6760 };
6761
6762 let mut validator = IngestValidator::default();
6763 validator
6764 .push(&store, 0, IngestEvent::Session(session.clone()))
6765 .await?;
6766 validator
6767 .push(&store, 1, IngestEvent::Message(message.clone()))
6768 .await?;
6769 validator
6770 .push(&store, 2, IngestEvent::Part(part.clone()))
6771 .await?;
6772 validator.finish(&store).await?;
6773
6774 let stored = store
6775 .get_session(&session.id)
6776 .await?
6777 .expect("session should exist");
6778 let stored_part = &stored.messages[0].parts[0];
6779 assert_eq!(stored_part, &part);
6780
6781 Ok(())
6782 }
6783
6784 fn base_session() -> Session {
6795 Session {
6796 id: "01HXY00000000001".to_owned(),
6797 parent_session_id: None,
6798 parent_message_id: None,
6799 source_agent: "claude-code".to_owned(),
6800 created_at: Utc::now(),
6801 project: crate::adapter::Extracted::from_test_value("/home/me/proj".to_owned()),
6802 options: ProviderOptions::new(),
6803 }
6804 }
6805
6806 fn count_status(outcomes: &[RowOutcome], target: OutcomeStatus) -> usize {
6807 outcomes
6808 .iter()
6809 .filter(|outcome| outcome.status == target)
6810 .count()
6811 }
6812
6813 #[tokio::test(flavor = "multi_thread")]
6814 async fn re_ingesting_a_session_with_unchanged_immutable_fields_is_idempotent()
6815 -> anyhow::Result<()> {
6816 let temp = TempDir::new()?;
6817 let store = Store::open_local(temp.path()).await?;
6818
6819 let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
6820 assert_eq!(count_status(&first, OutcomeStatus::Inserted), 1);
6821
6822 let mut again = base_session();
6823 again.options.insert("title".to_owned(), json!("renamed"));
6824 let second = ingest_events(&store, vec![IngestEvent::Session(again)]).await?;
6825 assert_eq!(
6826 count_status(&second, OutcomeStatus::Error),
6827 0,
6828 "options is mutable; the re-ingest must not surface an error: {second:?}",
6829 );
6830 assert_eq!(
6831 count_status(&second, OutcomeStatus::Matched),
6832 1,
6833 "unchanged immutable fields must match-insert via merge_insert",
6834 );
6835
6836 Ok(())
6837 }
6838
6839 #[tokio::test(flavor = "multi_thread")]
6840 async fn re_ingesting_with_changed_source_agent_is_rejected() -> anyhow::Result<()> {
6841 let temp = TempDir::new()?;
6842 let store = Store::open_local(temp.path()).await?;
6843
6844 let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
6845 assert_eq!(count_status(&first, OutcomeStatus::Error), 0);
6846
6847 let mut tampered = base_session();
6848 tampered.source_agent = "codex-cli".to_owned();
6849 let second = ingest_events(&store, vec![IngestEvent::Session(tampered)]).await?;
6850 assert_eq!(count_status(&second, OutcomeStatus::Error), 1);
6851 let err_row = second
6852 .iter()
6853 .find(|outcome| outcome.status == OutcomeStatus::Error)
6854 .expect("error outcome present");
6855 let err = err_row.error.as_ref().expect("error body present");
6856 assert_eq!(err.field, Some("source_agent"));
6857 assert_eq!(err.reason, Some("immutable"));
6858
6859 let stored = store
6861 .get_session(&base_session().id)
6862 .await?
6863 .expect("session row survives the rejected re-ingest");
6864 assert_eq!(stored.session.source_agent, "claude-code");
6865
6866 Ok(())
6867 }
6868
6869 #[tokio::test(flavor = "multi_thread")]
6870 async fn re_ingesting_with_changed_project_is_rejected() -> anyhow::Result<()> {
6871 let temp = TempDir::new()?;
6872 let store = Store::open_local(temp.path()).await?;
6873
6874 let first = ingest_events(&store, vec![IngestEvent::Session(base_session())]).await?;
6875 assert_eq!(count_status(&first, OutcomeStatus::Error), 0);
6876
6877 let mut tampered = base_session();
6878 tampered.project = crate::adapter::Extracted::from_test_value("/somewhere/else".to_owned());
6879 let second = ingest_events(&store, vec![IngestEvent::Session(tampered)]).await?;
6880 let err_row = second
6881 .iter()
6882 .find(|outcome| outcome.status == OutcomeStatus::Error)
6883 .expect("project change must surface an error outcome");
6884 assert_eq!(err_row.error.as_ref().unwrap().field, Some("project"));
6885
6886 let stored = store
6887 .get_session(&base_session().id)
6888 .await?
6889 .expect("session row survives");
6890 assert_eq!(
6891 stored.session.project.as_str(),
6892 "/home/me/proj",
6893 "stored project must remain the original",
6894 );
6895
6896 Ok(())
6897 }
6898
6899 #[tokio::test(flavor = "multi_thread")]
6900 async fn batched_flush_attributes_new_messages_on_existing_session() -> anyhow::Result<()> {
6901 use crate::wire::Provenance;
6909 let temp = TempDir::new()?;
6910 let store = Store::open_local(temp.path()).await?;
6911 let session = base_session();
6912
6913 let text_part = |part_id: &str, message_id: &str, body: &str| Part {
6914 session_id: session.id.clone(),
6915 id: part_id.to_owned(),
6916 message_id: message_id.to_owned(),
6917 ordinal: 0,
6918 provenance: Provenance::Conversational,
6919 options: ProviderOptions::new(),
6920 kind: PartKind::Text {
6921 text: Some(Extracted::from_test_value(body.to_owned())),
6922 },
6923 };
6924 let user_message = |id: &str| Message::User {
6925 id: id.to_owned(),
6926 session_id: session.id.clone(),
6927 timestamp: Utc::now(),
6928 options: ProviderOptions::new(),
6929 };
6930
6931 let mut validator = IngestValidator::default();
6933 validator
6934 .push(&store, 0, IngestEvent::Session(session.clone()))
6935 .await?;
6936 validator
6937 .push(&store, 1, IngestEvent::Message(user_message("m1")))
6938 .await?;
6939 validator
6940 .push(&store, 2, IngestEvent::Part(text_part("p1", "m1", "alpha")))
6941 .await?;
6942 validator
6943 .push(&store, 3, IngestEvent::Message(user_message("m2")))
6944 .await?;
6945 validator
6946 .push(&store, 4, IngestEvent::Part(text_part("p2", "m2", "beta")))
6947 .await?;
6948 let (_first_outcomes, first_counts) = validator.finish(&store).await?;
6949 assert_eq!(first_counts.sessions_inserted, 1);
6950 assert_eq!(first_counts.messages_inserted_total, 2);
6951 assert_eq!(first_counts.messages_inserted_searchable, 2);
6952
6953 let mut validator = IngestValidator::default();
6955 validator
6956 .push(&store, 0, IngestEvent::Session(session.clone()))
6957 .await?;
6958 for (idx, mid) in ["m3", "m4", "m5"].iter().enumerate() {
6959 let pid = format!("p{}", idx + 3);
6960 validator
6961 .push(&store, idx * 2 + 1, IngestEvent::Message(user_message(mid)))
6962 .await?;
6963 validator
6964 .push(
6965 &store,
6966 idx * 2 + 2,
6967 IngestEvent::Part(text_part(&pid, mid, "gamma")),
6968 )
6969 .await?;
6970 }
6971 let (second_outcomes, second_counts) = validator.finish(&store).await?;
6972
6973 assert_eq!(
6974 second_counts.sessions_inserted, 0,
6975 "existing session row must report as Matched, not Inserted",
6976 );
6977 assert_eq!(second_counts.sessions_matched, 1);
6978 assert_eq!(
6979 second_counts.messages_inserted_total, 3,
6980 "the three NEW messages must register as Inserted in BatchCounts",
6981 );
6982 assert_eq!(
6983 second_counts.messages_inserted_searchable, 3,
6984 "all three new messages carry conversational text -> searchable",
6985 );
6986 assert_eq!(second_counts.messages_matched_total, 0);
6987 assert_eq!(second_counts.parts_inserted, 3);
6988 assert_eq!(second_counts.parts_matched, 0);
6989
6990 let session_outcome = second_outcomes
6993 .iter()
6994 .find(|outcome| outcome.kind == "session")
6995 .expect("session-row outcome present");
6996 assert_eq!(session_outcome.status, OutcomeStatus::Matched);
6997 for outcome in &second_outcomes {
6998 if outcome.kind == "message" || outcome.kind == "part" {
6999 assert_eq!(
7000 outcome.status,
7001 OutcomeStatus::Inserted,
7002 "new row must be Inserted, got: {outcome:?}",
7003 );
7004 }
7005 }
7006 Ok(())
7007 }
7008
7009 async fn store_with_messages(
7013 temp: &TempDir,
7014 count: usize,
7015 ) -> anyhow::Result<(Store, Vec<MessageKey>)> {
7016 store_with_messages_at_threshold(temp, count, VECTOR_INDEX_ACTIVATION_ROWS).await
7017 }
7018
7019 async fn store_with_messages_at_threshold(
7022 temp: &TempDir,
7023 count: usize,
7024 _vector_threshold: usize,
7025 ) -> anyhow::Result<(Store, Vec<MessageKey>)> {
7026 let store = Store::open_local(temp.path()).await?;
7027 let sessions = 8.min(count.max(1));
7028 let mut events = Vec::new();
7029 for s in 0..sessions {
7030 events.push(IngestEvent::Session(Session {
7031 id: format!("session-{s}"),
7032 parent_session_id: None,
7033 parent_message_id: None,
7034 source_agent: "claude-code".to_owned(),
7035 created_at: Utc::now(),
7036 project: Extracted::from_test_value(format!("/proj/{}", s % 4)),
7037 options: ProviderOptions::new(),
7038 }));
7039 for i in (s..count).step_by(sessions) {
7040 let message_id = format!("msg-{i}");
7041 events.push(IngestEvent::Message(Message::User {
7042 id: message_id.clone(),
7043 session_id: format!("session-{s}"),
7044 timestamp: Utc::now(),
7045 options: ProviderOptions::new(),
7046 }));
7047 events.push(IngestEvent::Part(Part {
7048 session_id: format!("session-{s}"),
7049 id: format!("{message_id}-part"),
7050 message_id,
7051 ordinal: 0,
7052 provenance: crate::wire::Provenance::Conversational,
7053 options: ProviderOptions::new(),
7054 kind: PartKind::Text {
7055 text: Some(Extracted::from_test_value(format!("synthetic message {i}"))),
7056 },
7057 }));
7058 }
7059 }
7060 ingest_events(&store, events).await?;
7061 let keys = (0..count)
7062 .map(|i| MessageKey {
7063 session_id: format!("session-{}", i % sessions),
7064 message_id: format!("msg-{i}"),
7065 })
7066 .collect();
7067 Ok((store, keys))
7068 }
7069
7070 fn synthetic_vector(seed: usize) -> Vec<f32> {
7072 let mut state = (seed as u64)
7073 .wrapping_mul(0x9E37_79B9_7F4A_7C15)
7074 .wrapping_add(1);
7075 (0..embedding_dim())
7076 .map(|_| {
7077 state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
7078 #[allow(clippy::cast_precision_loss)]
7079 let unit = (state >> 33) as f32 / (1u64 << 31) as f32;
7080 unit - 1.0
7081 })
7082 .collect()
7083 }
7084
7085 fn embedded(keys: &[MessageKey]) -> Vec<EmbeddedMessage> {
7087 keys.iter()
7088 .enumerate()
7089 .map(|(seed, key)| EmbeddedMessage {
7090 session_id: key.session_id.clone(),
7091 id: key.message_id.clone(),
7092 vector: synthetic_vector(seed),
7093 })
7094 .collect()
7095 }
7096
7097 fn embedding_update_batch_with_model(
7098 rows: &[EmbeddedMessage],
7099 model: &str,
7100 ) -> Result<RecordBatch> {
7101 let mut batch = embedding_update_batch(rows)?;
7102 let columns = batch
7103 .columns()
7104 .iter()
7105 .take(3)
7106 .cloned()
7107 .chain(std::iter::once(
7108 Arc::new(StringArray::from(vec![model; rows.len()])) as _,
7109 ))
7110 .collect::<Vec<_>>();
7111 batch = RecordBatch::try_new(batch.schema(), columns)?;
7112 Ok(batch)
7113 }
7114
7115 #[tokio::test]
7116 async fn filtered_vector_scan_pushes_scalar_predicate_into_the_index() -> anyhow::Result<()> {
7117 let temp = TempDir::new()?;
7118 let (store, keys) = store_with_messages(&temp, 4).await?;
7122 store.write_embeddings(&embedded(&keys)).await?;
7123 store
7124 .optimize_indices(None, &MaintenancePolicy::always_compact())
7125 .await?
7126 .into_result()?;
7127
7128 let query = vec![0.01_f32; embedding_dim()];
7129 let plan = store
7130 .explain_vector_plan(
7131 &query,
7132 10,
7133 &Predicate::Eq("session_id", "session-3".into()),
7134 None,
7135 )
7136 .await?;
7137
7138 assert!(
7143 plan.contains("ScalarIndexQuery"),
7144 "expected a ScalarIndexQuery node in the plan:\n{plan}",
7145 );
7146 let predicate_postfiltered = plan
7147 .lines()
7148 .any(|line| line.contains("FilterExec") && line.contains("session_id"));
7149 assert!(
7150 !predicate_postfiltered,
7151 "the scalar predicate must not fall back to a FilterExec postfilter:\n{plan}",
7152 );
7153 Ok(())
7154 }
7155
7156 #[tokio::test]
7157 async fn vector_index_activates_when_threshold_is_crossed() -> anyhow::Result<()> {
7158 let temp = TempDir::new()?;
7159 let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7160
7161 store.write_embeddings(&embedded(&keys[..255])).await?;
7164 store
7165 .optimize_indices_with_vector_threshold(256)
7166 .await?
7167 .into_result()?;
7168 assert!(
7169 !store
7170 .handle
7171 .messages_index_names()
7172 .await?
7173 .iter()
7174 .any(|name| name == MESSAGES_VECTOR_INDEX),
7175 "IVF_SQ must not exist below the activation threshold",
7176 );
7177
7178 store.write_embeddings(&embedded(&keys[255..256])).await?;
7181 store
7182 .optimize_indices_with_vector_threshold(256)
7183 .await?
7184 .into_result()?;
7185 assert!(
7186 store
7187 .handle
7188 .messages_index_names()
7189 .await?
7190 .iter()
7191 .any(|name| name == MESSAGES_VECTOR_INDEX),
7192 "optimize must create the IVF_SQ once the threshold is crossed",
7193 );
7194
7195 let hits = store
7198 .vector_search(&synthetic_vector(0), 10, &Predicate::And(Vec::new()), None)
7199 .await?;
7200 assert!(
7201 hits.iter().any(|hit| hit.key == keys[0]),
7202 "an embedded row is retrievable via the index",
7203 );
7204 Ok(())
7205 }
7206
7207 #[tokio::test]
7208 async fn scalar_fold_batching_defers_tail_without_losing_rows() -> anyhow::Result<()> {
7209 let temp = TempDir::new()?;
7210 let (store, _keys) = store_with_messages(&temp, 300).await?;
7211
7212 store
7214 .optimize_indices_with_scalar_fold_threshold(0)
7215 .await?
7216 .into_result()?;
7217 assert_eq!(
7218 store
7219 .handle
7220 .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7221 .await?,
7222 0,
7223 "threshold 0 must fold the scalar index over every row",
7224 );
7225
7226 let new_messages = 5usize;
7228 let mut events = vec![IngestEvent::Session(Session {
7229 id: "session-new".to_owned(),
7230 parent_session_id: None,
7231 parent_message_id: None,
7232 source_agent: "claude-code".to_owned(),
7233 created_at: Utc::now(),
7234 project: Extracted::from_test_value("/proj/new".to_owned()),
7235 options: ProviderOptions::new(),
7236 })];
7237 for i in 0..new_messages {
7238 let message_id = format!("new-msg-{i}");
7239 events.push(IngestEvent::Message(Message::User {
7240 id: message_id.clone(),
7241 session_id: "session-new".to_owned(),
7242 timestamp: Utc::now(),
7243 options: ProviderOptions::new(),
7244 }));
7245 events.push(IngestEvent::Part(Part {
7246 session_id: "session-new".to_owned(),
7247 id: format!("{message_id}-part"),
7248 message_id,
7249 ordinal: 0,
7250 provenance: crate::wire::Provenance::Conversational,
7251 options: ProviderOptions::new(),
7252 kind: PartKind::Text {
7253 text: Some(Extracted::from_test_value(format!("new text {i}"))),
7254 },
7255 }));
7256 }
7257 ingest_events(&store, events).await?;
7258
7259 store
7262 .optimize_indices_with_scalar_fold_threshold(1_000_000)
7263 .await?
7264 .into_result()?;
7265 assert_eq!(
7266 store
7267 .handle
7268 .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7269 .await?,
7270 new_messages,
7271 "the scalar fold must be deferred, leaving the delta tail unindexed",
7272 );
7273 assert_eq!(
7274 store
7275 .handle
7276 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7277 .await?,
7278 0,
7279 "FTS must fold every run regardless of the scalar threshold",
7280 );
7281
7282 let session = store
7284 .get_session("session-new")
7285 .await?
7286 .expect("deferred session must still be retrievable");
7287 assert_eq!(
7288 session.messages.len(),
7289 new_messages,
7290 "get must read the unindexed scalar tail via scan",
7291 );
7292
7293 store
7295 .optimize_indices_with_scalar_fold_threshold(0)
7296 .await?
7297 .into_result()?;
7298 assert_eq!(
7299 store
7300 .handle
7301 .unindexed_row_count(Table::Messages, MESSAGES_SESSION_ID_INDEX)
7302 .await?,
7303 0,
7304 "a threshold-0 fold must consolidate the deferred tail",
7305 );
7306 Ok(())
7307 }
7308
7309 #[tokio::test]
7317 async fn fts_consolidation_rebuilds_instead_of_merging() -> anyhow::Result<()> {
7318 use crate::substrate::{OptimizeEvent, OptimizePhase};
7319
7320 type PhaseLog = std::sync::Arc<std::sync::Mutex<Vec<(OptimizePhase, Option<String>)>>>;
7321 let temp = TempDir::new()?;
7322 let (store, _keys) = store_with_messages(&temp, 300).await?;
7323 let phases: PhaseLog = std::sync::Arc::default();
7324 let sink = phases.clone();
7325 let progress: crate::substrate::OptimizeProgressFn = std::sync::Arc::new(move |event| {
7326 if let OptimizeEvent::PhaseStart { phase, detail, .. } = event {
7327 sink.lock().unwrap().push((phase, detail));
7328 }
7329 });
7330
7331 for round in 0..=crate::substrate::DELTA_MERGE_THRESHOLD {
7332 ingest_events(&store, conversational_events(&format!("grow-{round}"), 2)).await?;
7333 store
7334 .optimize_indices(Some(progress.clone()), &MaintenancePolicy::always_compact())
7335 .await?
7336 .into_result()?;
7337 }
7338
7339 assert!(
7340 phases.lock().unwrap().iter().any(|(phase, detail)| {
7341 matches!(phase, OptimizePhase::IndexRebuild)
7342 && detail.as_deref() == Some(MESSAGES_FTS_INDEX)
7343 }),
7344 "crossing the threshold must rebuild the FTS index, not merge it",
7345 );
7346 let fts_segments = store
7347 .handle
7348 .messages_index_names()
7349 .await?
7350 .into_iter()
7351 .filter(|name| name == MESSAGES_FTS_INDEX)
7352 .count();
7353 assert_eq!(fts_segments, 1, "the rebuild collapses the segment chain");
7354 assert_eq!(
7355 store
7356 .handle
7357 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7358 .await?,
7359 0,
7360 "the rebuilt index covers every fragment",
7361 );
7362 Ok(())
7363 }
7364
7365 #[tokio::test]
7369 async fn grown_session_flush_skips_the_sessions_merge() -> anyhow::Result<()> {
7370 let temp = TempDir::new()?;
7371 let store = Store::open_local(temp.path()).await?;
7372 ingest_events(&store, conversational_events("session-grow", 1)).await?;
7373 let sessions_before = store.handle.dataset(Table::Sessions).await?.version_id();
7374 let messages_before = store.handle.dataset(Table::Messages).await?.version_id();
7375
7376 ingest_events(&store, conversational_events("session-grow", 2)).await?;
7377 assert_eq!(
7378 store.handle.dataset(Table::Sessions).await?.version_id(),
7379 sessions_before,
7380 "an all-present sessions batch must skip the merge commit",
7381 );
7382 assert!(
7383 store.handle.dataset(Table::Messages).await?.version_id() > messages_before,
7384 "the grown message rows must still commit",
7385 );
7386 let session = store
7387 .get_session("session-grow")
7388 .await?
7389 .expect("session row must survive the skipped merge");
7390 assert_eq!(session.messages.len(), 2);
7391 Ok(())
7392 }
7393
7394 #[tokio::test]
7401 async fn fts_fold_skips_a_tail_with_no_indexable_text() -> anyhow::Result<()> {
7402 let temp = TempDir::new()?;
7403 let (store, _keys) = store_with_messages(&temp, 300).await?;
7404 store
7405 .optimize_indices_with_scalar_fold_threshold(0)
7406 .await?
7407 .into_result()?;
7408
7409 let tool_only_message = |i: usize| {
7410 let message_id = format!("tool-msg-{i}");
7411 [
7412 IngestEvent::Message(Message::Assistant {
7413 id: message_id.clone(),
7414 session_id: "session-toolonly".to_owned(),
7415 timestamp: Utc::now(),
7416 options: ProviderOptions::new(),
7417 }),
7418 IngestEvent::Part(Part {
7419 session_id: "session-toolonly".to_owned(),
7420 id: format!("{message_id}-part"),
7421 message_id,
7422 ordinal: 0,
7423 provenance: crate::wire::Provenance::Conversational,
7424 options: ProviderOptions::new(),
7425 kind: PartKind::ToolCall {
7426 call_id: Some(Extracted::from_test_value(format!("call-{i}"))),
7427 name: Some(Extracted::from_test_value("Bash".to_owned())),
7428 params: serde_json::json!({"command": "ls"}),
7429 provider_executed: false,
7430 },
7431 }),
7432 ]
7433 };
7434 let mut events = vec![IngestEvent::Session(synthetic_session("session-toolonly"))];
7435 events.extend((0..3).flat_map(tool_only_message));
7436 ingest_events(&store, events).await?;
7437
7438 store
7439 .optimize_indices_with_scalar_fold_threshold(0)
7440 .await?
7441 .into_result()?;
7442 assert_eq!(
7443 store
7444 .handle
7445 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7446 .await?,
7447 3,
7448 "an all-null tail must not fold into the FTS index",
7449 );
7450 let fts_status = |statuses: Vec<crate::substrate::IndexStatus>| {
7451 statuses
7452 .into_iter()
7453 .find(|status| status.intent_name == MESSAGES_FTS_INDEX)
7454 .expect("FTS status present")
7455 };
7456 assert_eq!(
7457 fts_status(store.index_status().await?).unindexed_rows,
7458 3,
7459 "the raw view counts uncovered rows",
7460 );
7461 assert_eq!(
7462 fts_status(store.index_status_indexable().await?).unindexed_rows,
7463 0,
7464 "the indexable view treats an all-null tail as nothing pending",
7465 );
7466
7467 ingest_events(&store, conversational_events("session-toolonly", 1)).await?;
7470 store
7471 .optimize_indices_with_scalar_fold_threshold(0)
7472 .await?
7473 .into_result()?;
7474 assert_eq!(
7475 store
7476 .handle
7477 .unindexed_row_count(Table::Messages, MESSAGES_FTS_INDEX)
7478 .await?,
7479 0,
7480 "a fold with real text must consolidate the skipped rows too",
7481 );
7482 Ok(())
7483 }
7484
7485 #[tokio::test]
7491 async fn vector_search_covers_the_unindexed_tail() -> anyhow::Result<()> {
7492 let temp = TempDir::new()?;
7493 let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7494 store.write_embeddings(&embedded(&keys)).await?;
7495 store
7496 .optimize_indices_with_vector_threshold(256)
7497 .await?
7498 .into_result()?;
7499 assert_eq!(
7500 store
7501 .handle
7502 .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
7503 .await?,
7504 0,
7505 "the IVF must cover the whole base after the fold",
7506 );
7507
7508 let tail = MessageKey {
7510 session_id: "session-tail".to_owned(),
7511 message_id: "tail-msg".to_owned(),
7512 };
7513 ingest_events(
7514 &store,
7515 vec![
7516 IngestEvent::Session(Session {
7517 id: tail.session_id.clone(),
7518 parent_session_id: None,
7519 parent_message_id: None,
7520 source_agent: "claude-code".to_owned(),
7521 created_at: Utc::now(),
7522 project: Extracted::from_test_value("/proj/tail".to_owned()),
7523 options: ProviderOptions::new(),
7524 }),
7525 IngestEvent::Message(Message::User {
7526 id: tail.message_id.clone(),
7527 session_id: tail.session_id.clone(),
7528 timestamp: Utc::now(),
7529 options: ProviderOptions::new(),
7530 }),
7531 IngestEvent::Part(Part {
7532 session_id: tail.session_id.clone(),
7533 id: format!("{}-part", tail.message_id),
7534 message_id: tail.message_id.clone(),
7535 ordinal: 0,
7536 provenance: crate::wire::Provenance::Conversational,
7537 options: ProviderOptions::new(),
7538 kind: PartKind::Text {
7539 text: Some(Extracted::from_test_value("tail body".to_owned())),
7540 },
7541 }),
7542 ],
7543 )
7544 .await?;
7545 let tail_vector = synthetic_vector(9999);
7546 store
7547 .write_embeddings(&[EmbeddedMessage {
7548 session_id: tail.session_id.clone(),
7549 id: tail.message_id.clone(),
7550 vector: tail_vector.clone(),
7551 }])
7552 .await?;
7553 assert!(
7554 store
7555 .handle
7556 .unindexed_row_count(Table::Messages, MESSAGES_VECTOR_INDEX)
7557 .await?
7558 > 0,
7559 "the appended row must be an unindexed vector tail (no fold ran)",
7560 );
7561
7562 let hits = store
7566 .vector_search(&tail_vector, 10, &Predicate::And(Vec::new()), None)
7567 .await?;
7568 assert!(
7569 hits.iter().any(|hit| hit.key == tail),
7570 "the unindexed tail row must be reachable via vector search (complete recall)",
7571 );
7572 Ok(())
7573 }
7574
7575 #[tokio::test]
7576 async fn model_swap_force_re_embeds_only_stale_rows_and_rebuilds_ivf_pq() -> anyhow::Result<()>
7577 {
7578 let temp = TempDir::new()?;
7579 let (store, keys) = store_with_messages_at_threshold(&temp, 300, 256).await?;
7580 let old_rows = embedded(&keys);
7581 let old_batch = embedding_update_batch_with_model(&old_rows, "old-model")?;
7582 store
7583 .handle
7584 .merge_update(Table::Messages, old_batch, old_rows.len())
7585 .await?;
7586 store
7587 .optimize_indices_with_vector_threshold(256)
7588 .await?
7589 .into_result()?;
7590 assert!(
7591 store
7592 .handle
7593 .messages_index_names()
7594 .await?
7595 .iter()
7596 .any(|name| name == MESSAGES_VECTOR_INDEX),
7597 "IVF_SQ must exist before a model swap",
7598 );
7599 assert_eq!(store.stale_embedding_count().await?, keys.len());
7600
7601 store.drop_vector_index().await?;
7602 let mut pending = Vec::new();
7603 let stream = store.pending_or_stale_messages();
7604 tokio::pin!(stream);
7605 while let Some(row) = stream.next().await {
7606 pending.push(row?);
7607 }
7608 assert_eq!(
7609 pending.len(),
7610 keys.len(),
7611 "force stream should see stale rows"
7612 );
7613 store.write_embeddings(&embedded(&keys)).await?;
7614 assert_eq!(store.stale_embedding_count().await?, 0);
7615 store
7616 .optimize_indices_with_vector_threshold(256)
7617 .await?
7618 .into_result()?;
7619 assert!(
7620 store
7621 .handle
7622 .messages_index_names()
7623 .await?
7624 .iter()
7625 .any(|name| name == MESSAGES_VECTOR_INDEX),
7626 "optimize must rebuild IVF_SQ after force re-embed",
7627 );
7628
7629 let stream = store.pending_or_stale_messages();
7630 tokio::pin!(stream);
7631 assert!(stream.next().await.is_none(), "up-to-date rows are skipped");
7632 Ok(())
7633 }
7634
7635 #[tokio::test]
7636 async fn session_last_message_ids_come_from_durable_messages() -> anyhow::Result<()> {
7637 let temp = TempDir::new()?;
7638 let store = Store::open_local(temp.path()).await?;
7639 let session = synthetic_session("oracle");
7640 store
7641 .upsert_sessions(std::slice::from_ref(&session))
7642 .await?;
7643 let timestamp =
7644 chrono::DateTime::from_timestamp(1_700_000_000, 0).expect("valid timestamp");
7645 let message_a = Message::User {
7646 id: "oracle-a".to_owned(),
7647 session_id: session.id.clone(),
7648 timestamp,
7649 options: ProviderOptions::new(),
7650 };
7651 let message_b = Message::User {
7652 id: "oracle-b".to_owned(),
7653 session_id: session.id.clone(),
7654 timestamp,
7655 options: ProviderOptions::new(),
7656 };
7657 store
7658 .upsert_messages(
7659 &session,
7660 &[
7661 MessageWrite {
7662 message: &message_a,
7663 parts: &[],
7664 search_text: Some("a"),
7665 },
7666 MessageWrite {
7667 message: &message_b,
7668 parts: &[],
7669 search_text: Some("b"),
7670 },
7671 ],
7672 )
7673 .await?;
7674
7675 let empty_session = synthetic_session("session-row-only");
7676 store.upsert_sessions(&[empty_session]).await?;
7677
7678 let orphan = synthetic_session("messages-no-row");
7682 let orphan_message = Message::User {
7683 id: "orphan-a".to_owned(),
7684 session_id: orphan.id.clone(),
7685 timestamp,
7686 options: ProviderOptions::new(),
7687 };
7688 store
7689 .upsert_messages(
7690 &orphan,
7691 &[MessageWrite {
7692 message: &orphan_message,
7693 parts: &[],
7694 search_text: Some("a"),
7695 }],
7696 )
7697 .await?;
7698
7699 let map = store.session_last_message_ids().await?;
7700 assert_eq!(map.get("oracle").map(String::as_str), Some("oracle-b"));
7701 assert!(
7702 !map.contains_key("session-row-only"),
7703 "a session row without durable messages must not produce a freshness key",
7704 );
7705 assert!(
7706 !map.contains_key("messages-no-row"),
7707 "messages without a durable session row must not produce a freshness key",
7708 );
7709 Ok(())
7710 }
7711
7712 #[tokio::test]
7713 async fn embedding_progress_counts_embedded_and_eligible_rows() -> anyhow::Result<()> {
7714 let temp = TempDir::new()?;
7715 let (store, keys) = store_with_messages(&temp, 10).await?;
7716
7717 let before = store.embedding_progress().await?;
7718 assert_eq!(before.embedded, 0);
7719 assert_eq!(before.total, 10);
7720 assert_eq!(before.backlog, 10);
7721 assert_eq!(before.model, crate::embed::model_id());
7722
7723 store.write_embeddings(&embedded(&keys[..4])).await?;
7724 let partial = store.embedding_progress().await?;
7725 assert_eq!(partial.embedded, 4);
7726 assert_eq!(partial.total, 10);
7727 assert_eq!(partial.backlog, 6);
7728
7729 store.write_embeddings(&embedded(&keys[4..])).await?;
7730 let full = store.embedding_progress().await?;
7731 assert_eq!(full.embedded, 10);
7732 assert_eq!(full.total, 10);
7733 assert_eq!(full.backlog, 0);
7736 assert_eq!(full.backlog, store.embed_backlog_count().await?);
7737 Ok(())
7738 }
7739
7740 #[tokio::test]
7741 async fn load_rowmap_if_present_installs_published_chain_without_building() -> anyhow::Result<()>
7742 {
7743 let temp = TempDir::new()?;
7744 let (builder, _keys) = store_with_messages(&temp, 6).await?;
7745 let cache = temp.path().join("cache");
7746
7747 let reader = Store::open_local(temp.path()).await?;
7750 reader.load_rowmap_if_present(&cache).await?;
7751 assert!(reader.rowmap_snapshot().is_none());
7752
7753 builder.ensure_rowmap(&cache).await?;
7755 reader.load_rowmap_if_present(&cache).await?;
7756 assert!(reader.rowmap_snapshot().is_some());
7757 Ok(())
7758 }
7759
7760 #[tokio::test]
7761 async fn ensure_rowmap_layers_a_delta_on_new_ingest() -> anyhow::Result<()> {
7762 let temp = TempDir::new()?;
7763 let (store, _keys) = store_with_messages(&temp, 6).await?;
7764 let cache = temp.path().join("cache");
7765
7766 store.ensure_rowmap(&cache).await?;
7767 assert_eq!(
7768 store.rowmap_delta_count(),
7769 Some(0),
7770 "first build is a lone base"
7771 );
7772
7773 ingest_events(
7775 &store,
7776 vec![
7777 IngestEvent::Session(Session {
7778 id: "session-new".to_owned(),
7779 parent_session_id: None,
7780 parent_message_id: None,
7781 source_agent: "claude-code".to_owned(),
7782 created_at: Utc::now(),
7783 project: Extracted::from_test_value("/proj/new".to_owned()),
7784 options: ProviderOptions::new(),
7785 }),
7786 IngestEvent::Message(Message::User {
7787 id: "m-new".to_owned(),
7788 session_id: "session-new".to_owned(),
7789 timestamp: Utc::now(),
7790 options: ProviderOptions::new(),
7791 }),
7792 IngestEvent::Part(Part {
7793 session_id: "session-new".to_owned(),
7794 id: "m-new-part".to_owned(),
7795 message_id: "m-new".to_owned(),
7796 ordinal: 0,
7797 provenance: crate::wire::Provenance::Conversational,
7798 options: ProviderOptions::new(),
7799 kind: PartKind::Text {
7800 text: Some(Extracted::from_test_value("brand new message".to_owned())),
7801 },
7802 }),
7803 ],
7804 )
7805 .await?;
7806
7807 store.ensure_rowmap(&cache).await?;
7810 assert_eq!(
7811 store.rowmap_delta_count(),
7812 Some(1),
7813 "new ingest layered a delta"
7814 );
7815
7816 let counts = store
7818 .session_message_counts(&["session-new".to_owned()])
7819 .await?;
7820 assert_eq!(counts.get("session-new").copied(), Some(1));
7821 Ok(())
7822 }
7823
7824 #[tokio::test]
7831 async fn ensure_rowmap_rebuilds_when_base_manifest_reclaimed() -> anyhow::Result<()> {
7832 let temp = TempDir::new()?;
7833 let (store, _keys) = store_with_messages(&temp, 6).await?;
7834 let cache = temp.path().join("cache");
7835
7836 store.ensure_rowmap(&cache).await?;
7839 assert_eq!(store.rowmap_delta_count(), Some(0), "first build is a base");
7840 let base_version = store.messages_version().await?;
7841 let versions_dir = temp.path().join("messages.lance").join("_versions");
7842 let base_manifests: Vec<_> = std::fs::read_dir(&versions_dir)?
7843 .filter_map(|entry| entry.ok().map(|entry| entry.path()))
7844 .filter(|path| path.extension().is_some_and(|ext| ext == "manifest"))
7845 .collect();
7846 assert!(
7847 !base_manifests.is_empty(),
7848 "the base version has a manifest"
7849 );
7850
7851 ingest_events(
7854 &store,
7855 vec![
7856 IngestEvent::Session(Session {
7857 id: "session-after".to_owned(),
7858 parent_session_id: None,
7859 parent_message_id: None,
7860 source_agent: "claude-code".to_owned(),
7861 created_at: Utc::now(),
7862 project: Extracted::from_test_value("/proj/after".to_owned()),
7863 options: ProviderOptions::new(),
7864 }),
7865 IngestEvent::Message(Message::User {
7866 id: "m-after".to_owned(),
7867 session_id: "session-after".to_owned(),
7868 timestamp: Utc::now(),
7869 options: ProviderOptions::new(),
7870 }),
7871 IngestEvent::Part(Part {
7872 session_id: "session-after".to_owned(),
7873 id: "m-after-part".to_owned(),
7874 message_id: "m-after".to_owned(),
7875 ordinal: 0,
7876 provenance: crate::wire::Provenance::Conversational,
7877 options: ProviderOptions::new(),
7878 kind: PartKind::Text {
7879 text: Some(Extracted::from_test_value("after the base".to_owned())),
7880 },
7881 }),
7882 ],
7883 )
7884 .await?;
7885 assert!(
7886 store.messages_version().await? > base_version,
7887 "the new ingest advanced the dataset past the chain's base"
7888 );
7889
7890 for manifest in &base_manifests {
7893 std::fs::remove_file(manifest)?;
7894 }
7895
7896 let reopened = Store::open_local(temp.path()).await?;
7899 reopened.ensure_rowmap(&cache).await?;
7900 assert!(
7901 reopened.rowmap_snapshot().is_some(),
7902 "map rebuilt after the base manifest was reclaimed"
7903 );
7904 assert_eq!(
7905 reopened.rowmap_delta_count(),
7906 Some(0),
7907 "a reclaimed base forces a fresh full-scan base, not a stuck chain"
7908 );
7909
7910 let counts = reopened
7912 .session_message_counts(&["session-after".to_owned()])
7913 .await?;
7914 assert_eq!(counts.get("session-after").copied(), Some(1));
7915 Ok(())
7916 }
7917
7918 #[tokio::test]
7925 async fn ensure_rowmap_deltas_across_embedding_fragment_rewrite() -> anyhow::Result<()> {
7926 let temp = TempDir::new()?;
7927 let (store, keys) = store_with_messages(&temp, 6).await?;
7928 let cache = temp.path().join("cache");
7929 store.ensure_rowmap(&cache).await?;
7930 assert_eq!(store.rowmap_delta_count(), Some(0), "first build is a base");
7931
7932 store.write_embeddings(&embedded(&keys)).await?;
7935
7936 ingest_events(
7938 &store,
7939 vec![
7940 IngestEvent::Session(Session {
7941 id: "session-after".to_owned(),
7942 parent_session_id: None,
7943 parent_message_id: None,
7944 source_agent: "claude-code".to_owned(),
7945 created_at: Utc::now(),
7946 project: Extracted::from_test_value("/proj/after".to_owned()),
7947 options: ProviderOptions::new(),
7948 }),
7949 IngestEvent::Message(Message::User {
7950 id: "m-after".to_owned(),
7951 session_id: "session-after".to_owned(),
7952 timestamp: Utc::now(),
7953 options: ProviderOptions::new(),
7954 }),
7955 IngestEvent::Part(Part {
7956 session_id: "session-after".to_owned(),
7957 id: "m-after-part".to_owned(),
7958 message_id: "m-after".to_owned(),
7959 ordinal: 0,
7960 provenance: crate::wire::Provenance::Conversational,
7961 options: ProviderOptions::new(),
7962 kind: PartKind::Text {
7963 text: Some(Extracted::from_test_value("after embedding".to_owned())),
7964 },
7965 }),
7966 ],
7967 )
7968 .await?;
7969
7970 store.ensure_rowmap(&cache).await?;
7973 assert_eq!(
7974 store.rowmap_delta_count(),
7975 Some(1),
7976 "fragment rewrite + append must layer a delta, not full-rebuild"
7977 );
7978
7979 let counts = store
7982 .session_message_counts(&["session-after".to_owned(), "session-0".to_owned()])
7983 .await?;
7984 assert_eq!(counts.get("session-after").copied(), Some(1));
7985 assert_eq!(
7986 counts.get("session-0").copied(),
7987 Some(1),
7988 "a base row survived the rewrite without being double-counted"
7989 );
7990 Ok(())
7991 }
7992
7993 #[tokio::test]
7994 async fn rowmap_chain_compacts_and_stays_bounded() -> anyhow::Result<()> {
7995 let temp = TempDir::new()?;
7998 let (store, _keys) = store_with_messages(&temp, 4).await?;
7999 let cache = temp.path().join("cache");
8000 store.ensure_rowmap(&cache).await?;
8001
8002 let mut reached_cap = false;
8003 let mut compacted = false;
8004 for i in 0..(Store::MAX_ROWMAP_DELTAS + 2) {
8005 let session = format!("session-x{i}");
8006 ingest_events(
8007 &store,
8008 vec![
8009 IngestEvent::Session(Session {
8010 id: session.clone(),
8011 parent_session_id: None,
8012 parent_message_id: None,
8013 source_agent: "claude-code".to_owned(),
8014 created_at: Utc::now(),
8015 project: Extracted::from_test_value("/proj/x".to_owned()),
8016 options: ProviderOptions::new(),
8017 }),
8018 IngestEvent::Message(Message::User {
8019 id: format!("mx{i}"),
8020 session_id: session.clone(),
8021 timestamp: Utc::now(),
8022 options: ProviderOptions::new(),
8023 }),
8024 IngestEvent::Part(Part {
8025 session_id: session.clone(),
8026 id: format!("mx{i}-part"),
8027 message_id: format!("mx{i}"),
8028 ordinal: 0,
8029 provenance: crate::wire::Provenance::Conversational,
8030 options: ProviderOptions::new(),
8031 kind: PartKind::Text {
8032 text: Some(Extracted::from_test_value(format!("msg {i}"))),
8033 },
8034 }),
8035 ],
8036 )
8037 .await?;
8038 store.ensure_rowmap(&cache).await?;
8039 let deltas = store.rowmap_delta_count().unwrap();
8040 assert!(
8041 deltas <= Store::MAX_ROWMAP_DELTAS,
8042 "delta count {deltas} exceeded the cap",
8043 );
8044 if deltas == Store::MAX_ROWMAP_DELTAS {
8045 reached_cap = true;
8046 }
8047 if reached_cap && deltas < Store::MAX_ROWMAP_DELTAS {
8048 compacted = true;
8049 }
8050 }
8051 assert!(reached_cap, "deltas accumulated to the cap (append path)");
8052 assert!(compacted, "the chain compacted back into a base");
8053
8054 let mut rmm = 0;
8056 for entry in std::fs::read_dir(&cache)? {
8057 let name = entry?.file_name().into_string().unwrap_or_default();
8058 assert!(!name.contains(".tmp-"), "leaked build temp: {name}");
8059 if name.ends_with(".rmm") {
8060 rmm += 1;
8061 }
8062 }
8063 assert!(
8064 rmm <= Store::MAX_ROWMAP_DELTAS + 1,
8065 "files unbounded: {rmm}"
8066 );
8067 Ok(())
8068 }
8069
8070 #[tokio::test]
8071 async fn embed_backlog_count_tracks_eligible_unembedded_rows() -> anyhow::Result<()> {
8072 let temp = TempDir::new()?;
8073 let (store, keys) = store_with_messages(&temp, 10).await?;
8074
8075 assert_eq!(store.embed_backlog_count().await?, 10);
8078
8079 store.write_embeddings(&embedded(&keys[..4])).await?;
8080 assert_eq!(store.embed_backlog_count().await?, 6);
8081
8082 store.write_embeddings(&embedded(&keys[4..])).await?;
8083 assert_eq!(store.embed_backlog_count().await?, 0);
8084 Ok(())
8085 }
8086
8087 #[tokio::test]
8092 async fn unindexed_vector_backlog_over_states_when_embedded_rows_are_unfolded()
8093 -> anyhow::Result<()> {
8094 let temp = TempDir::new()?;
8095 let (store, keys) = store_with_messages(&temp, 10).await?;
8096 store.write_embeddings(&embedded(&keys)).await?;
8097
8098 assert_eq!(store.unindexed_vector_backlog().await?, 10);
8099 assert_eq!(store.embed_backlog_count().await?, 0);
8100 Ok(())
8101 }
8102
8103 #[tokio::test]
8104 async fn session_message_counts_returns_per_session_counts_with_zeros_for_unknown_sessions()
8105 -> anyhow::Result<()> {
8106 let temp = TempDir::new()?;
8109 let (store, _keys) = store_with_messages(&temp, 32).await?;
8110
8111 let mut requested: Vec<String> = (0..8).map(|s| format!("session-{s}")).collect();
8112 requested.push("session-unknown-a".to_owned());
8113 requested.push("session-unknown-b".to_owned());
8114 let counts = store.session_message_counts(&requested).await?;
8115
8116 assert_eq!(counts.len(), requested.len());
8119 for s in 0..8 {
8120 assert_eq!(
8121 counts.get(&format!("session-{s}")).copied(),
8122 Some(4),
8123 "session-{s} should have 4 messages",
8124 );
8125 }
8126 assert_eq!(counts.get("session-unknown-a").copied(), Some(0));
8127 assert_eq!(counts.get("session-unknown-b").copied(), Some(0));
8128
8129 let empty = store.session_message_counts(&[]).await?;
8131 assert!(empty.is_empty());
8132 Ok(())
8133 }
8134}