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