1use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
10pub struct MemoryOwner {
11 session_id: crate::types::SessionId,
13}
14
15impl MemoryOwner {
16 pub fn canonical_session(session_id: crate::types::SessionId) -> Self {
17 Self { session_id }
18 }
19
20 pub fn session_id(&self) -> &crate::types::SessionId {
21 &self.session_id
22 }
23
24 fn includes(&self, metadata: &MemoryMetadata) -> bool {
25 metadata.session_id == self.session_id
26 }
27}
28
29#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
36pub struct MessageRange {
37 start: u64,
38 end: u64,
39}
40
41impl MessageRange {
42 pub fn new(start: u64, end: u64) -> Result<Self, MemoryStoreError> {
47 if start > end {
48 return Err(MemoryStoreError::SourceRange { start, end });
49 }
50 Ok(Self { start, end })
51 }
52
53 pub fn single(offset: u64) -> Self {
55 Self {
56 start: offset,
57 end: offset.saturating_add(1),
58 }
59 }
60
61 pub fn start(&self) -> u64 {
62 self.start
63 }
64
65 pub fn end(&self) -> u64 {
66 self.end
67 }
68
69 pub fn len(&self) -> u64 {
71 self.end - self.start
72 }
73
74 pub fn is_empty(&self) -> bool {
75 self.start == self.end
76 }
77
78 pub fn overlaps(&self, other: &MessageRange) -> bool {
83 !self.is_empty() && !other.is_empty() && self.start < other.end && other.start < self.end
84 }
85}
86
87#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
93#[serde(tag = "kind", rename_all = "snake_case")]
94pub enum MemorySource {
95 Compaction {
98 source_range: MessageRange,
100 },
101}
102
103impl MemorySource {
104 pub fn source_range(&self) -> Option<MessageRange> {
106 match self {
107 MemorySource::Compaction { source_range } => Some(*source_range),
108 }
109 }
110}
111
112#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct MemoryMetadata {
115 pub session_id: crate::types::SessionId,
117 pub source: MemorySource,
119 pub indexed_at: crate::time_compat::SystemTime,
121}
122
123#[derive(Debug, Clone)]
125pub struct MemoryResult {
126 pub content: String,
128 pub metadata: MemoryMetadata,
130 pub score: f32,
132}
133
134#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
136pub struct MemorySearchScope {
137 pub owner: MemoryOwner,
139}
140
141impl MemorySearchScope {
142 pub fn for_session(session_id: crate::types::SessionId) -> Self {
143 Self {
144 owner: MemoryOwner::canonical_session(session_id),
145 }
146 }
147
148 pub fn for_owner(owner: MemoryOwner) -> Self {
149 Self { owner }
150 }
151
152 pub fn session_id(&self) -> &crate::types::SessionId {
153 self.owner.session_id()
154 }
155
156 pub fn includes(&self, metadata: &MemoryMetadata) -> bool {
157 self.owner.includes(metadata)
158 }
159}
160
161#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
163pub struct MemoryIndexScope {
164 pub owner: MemoryOwner,
166}
167
168impl MemoryIndexScope {
169 pub fn for_session(session_id: crate::types::SessionId) -> Self {
170 Self {
171 owner: MemoryOwner::canonical_session(session_id),
172 }
173 }
174
175 pub fn for_owner(owner: MemoryOwner) -> Self {
176 Self { owner }
177 }
178
179 pub fn session_id(&self) -> &crate::types::SessionId {
180 self.owner.session_id()
181 }
182
183 pub fn includes(&self, metadata: &MemoryMetadata) -> bool {
184 self.owner.includes(metadata)
185 }
186}
187
188#[derive(Debug, Clone)]
197pub struct MemoryIndexRequest {
198 scope: MemoryIndexScope,
199 content: crate::types::MemoryIndexableContent,
200 metadata: MemoryMetadata,
201}
202
203impl MemoryIndexRequest {
204 pub fn new(
205 scope: MemoryIndexScope,
206 content: crate::types::MemoryIndexableContent,
207 metadata: MemoryMetadata,
208 ) -> Result<Self, MemoryStoreError> {
209 if !scope.includes(&metadata) {
210 return Err(MemoryStoreError::Scope(format!(
211 "memory metadata session {} is outside indexing scope {}",
212 metadata.session_id,
213 scope.session_id()
214 )));
215 }
216 Ok(Self {
217 scope,
218 content,
219 metadata,
220 })
221 }
222
223 pub fn scope(&self) -> &MemoryIndexScope {
224 &self.scope
225 }
226
227 pub fn content(&self) -> &crate::types::MemoryIndexableContent {
229 &self.content
230 }
231
232 pub fn indexable_text(&self) -> Option<&str> {
234 self.content.indexable_text()
235 }
236
237 pub fn metadata(&self) -> &MemoryMetadata {
238 &self.metadata
239 }
240
241 pub fn into_parts(
242 self,
243 ) -> (
244 MemoryIndexScope,
245 crate::types::MemoryIndexableContent,
246 MemoryMetadata,
247 ) {
248 (self.scope, self.content, self.metadata)
249 }
250}
251
252#[derive(Debug, Clone)]
254pub struct MemoryIndexBatch {
255 scope: MemoryIndexScope,
256 requests: Vec<MemoryIndexRequest>,
257}
258
259impl MemoryIndexBatch {
260 pub fn new(
261 scope: MemoryIndexScope,
262 requests: Vec<MemoryIndexRequest>,
263 ) -> Result<Self, MemoryStoreError> {
264 for request in &requests {
265 if request.scope() != &scope {
266 return Err(MemoryStoreError::Scope(format!(
267 "memory index request scope {} is outside batch scope {}",
268 request.scope().session_id(),
269 scope.session_id()
270 )));
271 }
272 }
273 Ok(Self { scope, requests })
274 }
275
276 pub fn single(request: MemoryIndexRequest) -> Self {
277 Self {
278 scope: request.scope.clone(),
279 requests: vec![request],
280 }
281 }
282
283 pub fn scope(&self) -> &MemoryIndexScope {
284 &self.scope
285 }
286
287 pub fn len(&self) -> usize {
288 self.requests.len()
289 }
290
291 pub fn is_empty(&self) -> bool {
292 self.requests.is_empty()
293 }
294
295 pub fn into_parts(self) -> (MemoryIndexScope, Vec<MemoryIndexRequest>) {
296 (self.scope, self.requests)
297 }
298}
299
300#[derive(Debug, Clone)]
302pub struct MemoryIndexReceipt {
303 pub scope: MemoryIndexScope,
304 pub indexed_entries: usize,
305}
306
307#[derive(Debug)]
309pub enum MemoryIndexDelivery {
310 NoStore {
311 scope: MemoryIndexScope,
312 },
313 Delivered(MemoryIndexReceipt),
314 Rejected {
315 scope: MemoryIndexScope,
316 attempted_entries: usize,
317 error: MemoryStoreError,
318 },
319}
320
321#[derive(Debug, Clone)]
323pub struct MemoryScopeDropReceipt {
324 pub owner: MemoryOwner,
326 pub dropped_entries: usize,
328}
329
330#[derive(Debug, Clone, Copy)]
337pub struct MemoryEnumerationRequest {
338 pub limit: usize,
340 pub offset: usize,
342 pub source_overlap: Option<MessageRange>,
345 pub indexed_after: Option<crate::time_compat::SystemTime>,
349}
350
351impl MemoryEnumerationRequest {
352 pub fn admits(&self, metadata: &MemoryMetadata) -> bool {
360 if let Some(range) = self.source_overlap {
361 match metadata.source.source_range() {
362 Some(source_range) if source_range.overlaps(&range) => {}
363 _ => return false,
364 }
365 }
366 if let Some(after) = self.indexed_after
367 && metadata.indexed_at <= after
368 {
369 return false;
370 }
371 true
372 }
373}
374
375#[derive(Debug, Clone)]
377pub struct MemoryEnumerationPage {
378 pub records: Vec<MemoryRecord>,
380 pub next_offset: Option<usize>,
383}
384
385#[derive(Debug, Clone)]
388pub struct MemoryRecord {
389 pub content: String,
391 pub metadata: MemoryMetadata,
393}
394
395pub trait EmbeddingModel: Send + Sync {
400 fn dimension(&self) -> usize;
405
406 fn embed(&self, text: &str) -> Vec<f32>;
408}
409
410#[derive(Debug, Clone, Copy, PartialEq, Eq)]
415pub struct HnswParams {
416 pub max_nb_connection: usize,
418 pub max_layer: usize,
420 pub ef_construction: usize,
422 pub ef_search: usize,
424}
425
426impl Default for HnswParams {
427 fn default() -> Self {
428 Self {
429 max_nb_connection: 16,
430 max_layer: 16,
431 ef_construction: 200,
432 ef_search: 200,
433 }
434 }
435}
436
437#[derive(Clone)]
440pub struct MemoryRankingPolicy {
441 embedding_model: std::sync::Arc<dyn EmbeddingModel>,
442 hnsw_params: HnswParams,
443}
444
445impl MemoryRankingPolicy {
446 pub fn new(
448 embedding_model: std::sync::Arc<dyn EmbeddingModel>,
449 hnsw_params: HnswParams,
450 ) -> Self {
451 Self {
452 embedding_model,
453 hnsw_params,
454 }
455 }
456
457 pub fn embedding_model(&self) -> &std::sync::Arc<dyn EmbeddingModel> {
459 &self.embedding_model
460 }
461
462 pub fn hnsw_params(&self) -> HnswParams {
464 self.hnsw_params
465 }
466
467 pub fn dimension(&self) -> usize {
469 self.embedding_model.dimension()
470 }
471
472 pub fn embed(&self, text: &str) -> Vec<f32> {
474 self.embedding_model.embed(text)
475 }
476}
477
478impl std::fmt::Debug for MemoryRankingPolicy {
479 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
480 f.debug_struct("MemoryRankingPolicy")
481 .field("dimension", &self.embedding_model.dimension())
482 .field("hnsw_params", &self.hnsw_params)
483 .finish()
484 }
485}
486
487#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
489#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
490pub trait MemoryStore: Send + Sync {
491 async fn index_scoped(
493 &self,
494 request: MemoryIndexRequest,
495 ) -> Result<MemoryIndexReceipt, MemoryStoreError> {
496 self.index_scoped_batch(MemoryIndexBatch::single(request))
497 .await
498 }
499
500 async fn index_scoped_batch(
505 &self,
506 batch: MemoryIndexBatch,
507 ) -> Result<MemoryIndexReceipt, MemoryStoreError>;
508
509 async fn search(
511 &self,
512 scope: &MemorySearchScope,
513 query: &str,
514 limit: usize,
515 ) -> Result<Vec<MemoryResult>, MemoryStoreError>;
516
517 async fn drop_scope(
523 &self,
524 owner: &MemoryOwner,
525 ) -> Result<MemoryScopeDropReceipt, MemoryStoreError> {
526 let _ = owner;
527 Err(MemoryStoreError::Unsupported {
528 operation: "drop_scope",
529 })
530 }
531
532 async fn enumerate_scoped(
540 &self,
541 scope: &MemorySearchScope,
542 request: MemoryEnumerationRequest,
543 ) -> Result<MemoryEnumerationPage, MemoryStoreError> {
544 let _ = (scope, request);
545 Err(MemoryStoreError::Unsupported {
546 operation: "enumerate_scoped",
547 })
548 }
549}
550
551#[derive(Debug, thiserror::Error)]
557pub enum MemoryStoreError {
558 #[error("Scope error: {0}")]
560 Scope(String),
561
562 #[error("invalid memory source range: start {start} > end {end}")]
564 SourceRange { start: u64, end: u64 },
565
566 #[error("Embedding error: {0}")]
568 Embedding(String),
569
570 #[error("Storage error: {0}")]
572 Storage(String),
573
574 #[error("memory index lock poisoned")]
576 LockPoisoned,
577
578 #[error("memory point ID out of range")]
580 PointIdOutOfRange,
581
582 #[error("memory point ID overflow")]
584 PointIdOverflow,
585
586 #[error("memory store task join failed: {0}")]
588 TaskJoin(String),
589
590 #[error("memory text corruption at point {point_id}: stored bytes are not valid UTF-8")]
594 TextCorruption { point_id: i64 },
595
596 #[error(
600 "memory index/store divergence at point {point_id}: live index references a missing durable row"
601 )]
602 IndexDivergence { point_id: i64 },
603
604 #[error("memory scope index is poisoned pending rebuild from durable state")]
608 ScopePoisoned,
609
610 #[error(
613 "memory scope repair failed after partial index failure: {repair} (original failure: {original})"
614 )]
615 ScopeRepairFailed {
616 original: Box<MemoryStoreError>,
617 repair: Box<MemoryStoreError>,
618 },
619
620 #[error("memory store operation '{operation}' is unsupported by this store")]
622 Unsupported { operation: &'static str },
623
624 #[error("memory enumeration limit must be non-zero")]
628 EnumerationLimitZero,
629
630 #[error("IO error: {0}")]
632 Io(#[from] std::io::Error),
633}
634
635impl MemoryStoreError {
636 pub fn error_code(&self) -> &'static str {
641 match self {
642 Self::Scope(_) => "memory_scope",
643 Self::SourceRange { .. } => "memory_source_range",
644 Self::Embedding(_) => "memory_embedding",
645 Self::Storage(_) => "memory_storage",
646 Self::LockPoisoned => "memory_lock_poisoned",
647 Self::PointIdOutOfRange => "memory_point_id_out_of_range",
648 Self::PointIdOverflow => "memory_point_id_overflow",
649 Self::TaskJoin(_) => "memory_task_join",
650 Self::TextCorruption { .. } => "memory_text_corruption",
651 Self::IndexDivergence { .. } => "memory_index_divergence",
652 Self::ScopePoisoned => "memory_scope_poisoned",
653 Self::ScopeRepairFailed { .. } => "memory_scope_repair_failed",
654 Self::Unsupported { .. } => "memory_unsupported",
655 Self::EnumerationLimitZero => "memory_enumeration_limit_zero",
656 Self::Io(_) => "memory_io",
657 }
658 }
659}
660
661#[cfg(test)]
662#[allow(clippy::unwrap_used, clippy::expect_used)]
663mod tests {
664 use super::*;
665 use crate::time_compat::{Duration, UNIX_EPOCH};
666
667 fn range(start: u64, end: u64) -> MessageRange {
668 MessageRange::new(start, end).unwrap()
669 }
670
671 #[test]
672 fn overlaps_is_half_open() {
673 assert!(range(0, 5).overlaps(&range(4, 6)));
675 assert!(range(4, 6).overlaps(&range(0, 5)));
676 assert!(range(0, 10).overlaps(&range(3, 4)));
678 assert!(range(3, 4).overlaps(&range(0, 10)));
679 assert!(!range(0, 5).overlaps(&range(5, 10)));
681 assert!(!range(5, 10).overlaps(&range(0, 5)));
682 assert!(!range(0, 2).overlaps(&range(7, 9)));
684 }
685
686 #[test]
687 fn empty_range_never_overlaps() {
688 assert!(!range(3, 3).overlaps(&range(0, 10)));
689 assert!(!range(0, 10).overlaps(&range(3, 3)));
690 assert!(!range(3, 3).overlaps(&range(3, 3)));
691 }
692
693 #[test]
694 fn unsupported_error_code_is_stable() {
695 assert_eq!(
696 MemoryStoreError::Unsupported {
697 operation: "drop_scope",
698 }
699 .error_code(),
700 "memory_unsupported"
701 );
702 }
703
704 fn metadata_at(
705 indexed_at: crate::time_compat::SystemTime,
706 source: MemorySource,
707 ) -> MemoryMetadata {
708 MemoryMetadata {
709 session_id: crate::types::SessionId::new(),
710 source,
711 indexed_at,
712 }
713 }
714
715 #[test]
716 fn enumeration_request_admits_on_source_overlap() {
717 let request = MemoryEnumerationRequest {
718 limit: 10,
719 offset: 0,
720 source_overlap: Some(range(4, 6)),
721 indexed_after: None,
722 };
723 let overlapping = metadata_at(
724 UNIX_EPOCH,
725 MemorySource::Compaction {
726 source_range: range(0, 5),
727 },
728 );
729 let disjoint = metadata_at(
730 UNIX_EPOCH,
731 MemorySource::Compaction {
732 source_range: range(6, 9),
733 },
734 );
735 assert!(request.admits(&overlapping));
736 assert!(!request.admits(&disjoint));
737 }
738
739 #[test]
740 fn enumeration_request_indexed_after_is_strict() {
741 let boundary = UNIX_EPOCH + Duration::from_secs(100);
742 let request = MemoryEnumerationRequest {
743 limit: 10,
744 offset: 0,
745 source_overlap: None,
746 indexed_after: Some(boundary),
747 };
748 let at_boundary = metadata_at(
749 boundary,
750 MemorySource::Compaction {
751 source_range: range(0, 1),
752 },
753 );
754 let after_boundary = metadata_at(
755 boundary + Duration::from_secs(1),
756 MemorySource::Compaction {
757 source_range: range(0, 1),
758 },
759 );
760 let before_boundary = metadata_at(
761 UNIX_EPOCH,
762 MemorySource::Compaction {
763 source_range: range(0, 1),
764 },
765 );
766 assert!(!request.admits(&at_boundary));
767 assert!(request.admits(&after_boundary));
768 assert!(!request.admits(&before_boundary));
769 }
770
771 #[test]
772 fn enumeration_request_filters_compose() {
773 let request = MemoryEnumerationRequest {
774 limit: 10,
775 offset: 0,
776 source_overlap: Some(range(0, 5)),
777 indexed_after: Some(UNIX_EPOCH + Duration::from_secs(100)),
778 };
779 let both = metadata_at(
780 UNIX_EPOCH + Duration::from_secs(200),
781 MemorySource::Compaction {
782 source_range: range(2, 3),
783 },
784 );
785 let wrong_range = metadata_at(
786 UNIX_EPOCH + Duration::from_secs(200),
787 MemorySource::Compaction {
788 source_range: range(5, 9),
789 },
790 );
791 let too_early = metadata_at(
792 UNIX_EPOCH,
793 MemorySource::Compaction {
794 source_range: range(2, 3),
795 },
796 );
797 assert!(request.admits(&both));
798 assert!(!request.admits(&wrong_range));
799 assert!(!request.admits(&too_early));
800 }
801
802 struct MinimalStore;
806
807 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
808 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
809 impl MemoryStore for MinimalStore {
810 async fn index_scoped_batch(
811 &self,
812 batch: MemoryIndexBatch,
813 ) -> Result<MemoryIndexReceipt, MemoryStoreError> {
814 let (scope, requests) = batch.into_parts();
815 Ok(MemoryIndexReceipt {
816 scope,
817 indexed_entries: requests.len(),
818 })
819 }
820
821 async fn search(
822 &self,
823 _scope: &MemorySearchScope,
824 _query: &str,
825 _limit: usize,
826 ) -> Result<Vec<MemoryResult>, MemoryStoreError> {
827 Ok(Vec::new())
828 }
829 }
830
831 #[tokio::test]
832 async fn drop_scope_default_is_typed_unsupported() {
833 let store = MinimalStore;
834 let owner = MemoryOwner::canonical_session(crate::types::SessionId::new());
835 let error = store.drop_scope(&owner).await.unwrap_err();
836 assert!(matches!(
837 error,
838 MemoryStoreError::Unsupported {
839 operation: "drop_scope",
840 }
841 ));
842 assert_eq!(error.error_code(), "memory_unsupported");
843 }
844
845 #[tokio::test]
846 async fn enumerate_scoped_default_is_typed_unsupported() {
847 let store = MinimalStore;
848 let scope = MemorySearchScope::for_session(crate::types::SessionId::new());
849 let error = store
850 .enumerate_scoped(
851 &scope,
852 MemoryEnumerationRequest {
853 limit: 10,
854 offset: 0,
855 source_overlap: None,
856 indexed_after: None,
857 },
858 )
859 .await
860 .unwrap_err();
861 assert!(matches!(
862 error,
863 MemoryStoreError::Unsupported {
864 operation: "enumerate_scoped",
865 }
866 ));
867 assert_eq!(error.error_code(), "memory_unsupported");
868 }
869}