use async_trait::async_trait;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct MemoryOwner {
session_id: crate::types::SessionId,
}
impl MemoryOwner {
pub fn canonical_session(session_id: crate::types::SessionId) -> Self {
Self { session_id }
}
pub fn session_id(&self) -> &crate::types::SessionId {
&self.session_id
}
fn includes(&self, metadata: &MemoryMetadata) -> bool {
metadata.session_id == self.session_id
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
pub struct MessageRange {
start: u64,
end: u64,
}
impl MessageRange {
pub fn new(start: u64, end: u64) -> Result<Self, MemoryStoreError> {
if start > end {
return Err(MemoryStoreError::SourceRange { start, end });
}
Ok(Self { start, end })
}
pub fn single(offset: u64) -> Self {
Self {
start: offset,
end: offset.saturating_add(1),
}
}
pub fn start(&self) -> u64 {
self.start
}
pub fn end(&self) -> u64 {
self.end
}
pub fn len(&self) -> u64 {
self.end - self.start
}
pub fn is_empty(&self) -> bool {
self.start == self.end
}
pub fn overlaps(&self, other: &MessageRange) -> bool {
!self.is_empty() && !other.is_empty() && self.start < other.end && other.start < self.end
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum MemorySource {
Compaction {
source_range: MessageRange,
},
}
impl MemorySource {
pub fn source_range(&self) -> Option<MessageRange> {
match self {
MemorySource::Compaction { source_range } => Some(*source_range),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoryMetadata {
pub session_id: crate::types::SessionId,
pub source: MemorySource,
pub indexed_at: crate::time_compat::SystemTime,
}
#[derive(Debug, Clone)]
pub struct MemoryResult {
pub content: String,
pub metadata: MemoryMetadata,
pub score: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct MemorySearchScope {
pub owner: MemoryOwner,
}
impl MemorySearchScope {
pub fn for_session(session_id: crate::types::SessionId) -> Self {
Self {
owner: MemoryOwner::canonical_session(session_id),
}
}
pub fn for_owner(owner: MemoryOwner) -> Self {
Self { owner }
}
pub fn session_id(&self) -> &crate::types::SessionId {
self.owner.session_id()
}
pub fn includes(&self, metadata: &MemoryMetadata) -> bool {
self.owner.includes(metadata)
}
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct MemoryIndexScope {
pub owner: MemoryOwner,
}
impl MemoryIndexScope {
pub fn for_session(session_id: crate::types::SessionId) -> Self {
Self {
owner: MemoryOwner::canonical_session(session_id),
}
}
pub fn for_owner(owner: MemoryOwner) -> Self {
Self { owner }
}
pub fn session_id(&self) -> &crate::types::SessionId {
self.owner.session_id()
}
pub fn includes(&self, metadata: &MemoryMetadata) -> bool {
self.owner.includes(metadata)
}
}
#[derive(Debug, Clone)]
pub struct MemoryIndexRequest {
scope: MemoryIndexScope,
content: crate::types::MemoryIndexableContent,
metadata: MemoryMetadata,
}
impl MemoryIndexRequest {
pub fn new(
scope: MemoryIndexScope,
content: crate::types::MemoryIndexableContent,
metadata: MemoryMetadata,
) -> Result<Self, MemoryStoreError> {
if !scope.includes(&metadata) {
return Err(MemoryStoreError::Scope(format!(
"memory metadata session {} is outside indexing scope {}",
metadata.session_id,
scope.session_id()
)));
}
Ok(Self {
scope,
content,
metadata,
})
}
pub fn scope(&self) -> &MemoryIndexScope {
&self.scope
}
pub fn content(&self) -> &crate::types::MemoryIndexableContent {
&self.content
}
pub fn indexable_text(&self) -> Option<&str> {
self.content.indexable_text()
}
pub fn metadata(&self) -> &MemoryMetadata {
&self.metadata
}
pub fn into_parts(
self,
) -> (
MemoryIndexScope,
crate::types::MemoryIndexableContent,
MemoryMetadata,
) {
(self.scope, self.content, self.metadata)
}
}
#[derive(Debug, Clone)]
pub struct MemoryIndexBatch {
scope: MemoryIndexScope,
requests: Vec<MemoryIndexRequest>,
}
impl MemoryIndexBatch {
pub fn new(
scope: MemoryIndexScope,
requests: Vec<MemoryIndexRequest>,
) -> Result<Self, MemoryStoreError> {
for request in &requests {
if request.scope() != &scope {
return Err(MemoryStoreError::Scope(format!(
"memory index request scope {} is outside batch scope {}",
request.scope().session_id(),
scope.session_id()
)));
}
}
Ok(Self { scope, requests })
}
pub fn single(request: MemoryIndexRequest) -> Self {
Self {
scope: request.scope.clone(),
requests: vec![request],
}
}
pub fn scope(&self) -> &MemoryIndexScope {
&self.scope
}
pub fn len(&self) -> usize {
self.requests.len()
}
pub fn is_empty(&self) -> bool {
self.requests.is_empty()
}
pub fn into_parts(self) -> (MemoryIndexScope, Vec<MemoryIndexRequest>) {
(self.scope, self.requests)
}
}
#[derive(Debug, Clone)]
pub struct MemoryIndexReceipt {
pub scope: MemoryIndexScope,
pub indexed_entries: usize,
}
#[derive(Debug)]
pub enum MemoryIndexDelivery {
NoStore {
scope: MemoryIndexScope,
},
Delivered(MemoryIndexReceipt),
Rejected {
scope: MemoryIndexScope,
attempted_entries: usize,
error: MemoryStoreError,
},
}
#[derive(Debug, Clone)]
pub struct MemoryScopeDropReceipt {
pub owner: MemoryOwner,
pub dropped_entries: usize,
}
#[derive(Debug, Clone, Copy)]
pub struct MemoryEnumerationRequest {
pub limit: usize,
pub offset: usize,
pub source_overlap: Option<MessageRange>,
pub indexed_after: Option<crate::time_compat::SystemTime>,
}
impl MemoryEnumerationRequest {
pub fn admits(&self, metadata: &MemoryMetadata) -> bool {
if let Some(range) = self.source_overlap {
match metadata.source.source_range() {
Some(source_range) if source_range.overlaps(&range) => {}
_ => return false,
}
}
if let Some(after) = self.indexed_after
&& metadata.indexed_at <= after
{
return false;
}
true
}
}
#[derive(Debug, Clone)]
pub struct MemoryEnumerationPage {
pub records: Vec<MemoryRecord>,
pub next_offset: Option<usize>,
}
#[derive(Debug, Clone)]
pub struct MemoryRecord {
pub content: String,
pub metadata: MemoryMetadata,
}
pub trait EmbeddingModel: Send + Sync {
fn dimension(&self) -> usize;
fn embed(&self, text: &str) -> Vec<f32>;
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct HnswParams {
pub max_nb_connection: usize,
pub max_layer: usize,
pub ef_construction: usize,
pub ef_search: usize,
}
impl Default for HnswParams {
fn default() -> Self {
Self {
max_nb_connection: 16,
max_layer: 16,
ef_construction: 200,
ef_search: 200,
}
}
}
#[derive(Clone)]
pub struct MemoryRankingPolicy {
embedding_model: std::sync::Arc<dyn EmbeddingModel>,
hnsw_params: HnswParams,
}
impl MemoryRankingPolicy {
pub fn new(
embedding_model: std::sync::Arc<dyn EmbeddingModel>,
hnsw_params: HnswParams,
) -> Self {
Self {
embedding_model,
hnsw_params,
}
}
pub fn embedding_model(&self) -> &std::sync::Arc<dyn EmbeddingModel> {
&self.embedding_model
}
pub fn hnsw_params(&self) -> HnswParams {
self.hnsw_params
}
pub fn dimension(&self) -> usize {
self.embedding_model.dimension()
}
pub fn embed(&self, text: &str) -> Vec<f32> {
self.embedding_model.embed(text)
}
}
impl std::fmt::Debug for MemoryRankingPolicy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("MemoryRankingPolicy")
.field("dimension", &self.embedding_model.dimension())
.field("hnsw_params", &self.hnsw_params)
.finish()
}
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
pub trait MemoryStore: Send + Sync {
async fn index_scoped(
&self,
request: MemoryIndexRequest,
) -> Result<MemoryIndexReceipt, MemoryStoreError> {
self.index_scoped_batch(MemoryIndexBatch::single(request))
.await
}
async fn index_scoped_batch(
&self,
batch: MemoryIndexBatch,
) -> Result<MemoryIndexReceipt, MemoryStoreError>;
async fn search(
&self,
scope: &MemorySearchScope,
query: &str,
limit: usize,
) -> Result<Vec<MemoryResult>, MemoryStoreError>;
async fn drop_scope(
&self,
owner: &MemoryOwner,
) -> Result<MemoryScopeDropReceipt, MemoryStoreError> {
let _ = owner;
Err(MemoryStoreError::Unsupported {
operation: "drop_scope",
})
}
async fn enumerate_scoped(
&self,
scope: &MemorySearchScope,
request: MemoryEnumerationRequest,
) -> Result<MemoryEnumerationPage, MemoryStoreError> {
let _ = (scope, request);
Err(MemoryStoreError::Unsupported {
operation: "enumerate_scoped",
})
}
}
#[derive(Debug, thiserror::Error)]
pub enum MemoryStoreError {
#[error("Scope error: {0}")]
Scope(String),
#[error("invalid memory source range: start {start} > end {end}")]
SourceRange { start: u64, end: u64 },
#[error("Embedding error: {0}")]
Embedding(String),
#[error("Storage error: {0}")]
Storage(String),
#[error("memory index lock poisoned")]
LockPoisoned,
#[error("memory point ID out of range")]
PointIdOutOfRange,
#[error("memory point ID overflow")]
PointIdOverflow,
#[error("memory store task join failed: {0}")]
TaskJoin(String),
#[error("memory text corruption at point {point_id}: stored bytes are not valid UTF-8")]
TextCorruption { point_id: i64 },
#[error(
"memory index/store divergence at point {point_id}: live index references a missing durable row"
)]
IndexDivergence { point_id: i64 },
#[error("memory scope index is poisoned pending rebuild from durable state")]
ScopePoisoned,
#[error(
"memory scope repair failed after partial index failure: {repair} (original failure: {original})"
)]
ScopeRepairFailed {
original: Box<MemoryStoreError>,
repair: Box<MemoryStoreError>,
},
#[error("memory store operation '{operation}' is unsupported by this store")]
Unsupported { operation: &'static str },
#[error("memory enumeration limit must be non-zero")]
EnumerationLimitZero,
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
}
impl MemoryStoreError {
pub fn error_code(&self) -> &'static str {
match self {
Self::Scope(_) => "memory_scope",
Self::SourceRange { .. } => "memory_source_range",
Self::Embedding(_) => "memory_embedding",
Self::Storage(_) => "memory_storage",
Self::LockPoisoned => "memory_lock_poisoned",
Self::PointIdOutOfRange => "memory_point_id_out_of_range",
Self::PointIdOverflow => "memory_point_id_overflow",
Self::TaskJoin(_) => "memory_task_join",
Self::TextCorruption { .. } => "memory_text_corruption",
Self::IndexDivergence { .. } => "memory_index_divergence",
Self::ScopePoisoned => "memory_scope_poisoned",
Self::ScopeRepairFailed { .. } => "memory_scope_repair_failed",
Self::Unsupported { .. } => "memory_unsupported",
Self::EnumerationLimitZero => "memory_enumeration_limit_zero",
Self::Io(_) => "memory_io",
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use crate::time_compat::{Duration, UNIX_EPOCH};
fn range(start: u64, end: u64) -> MessageRange {
MessageRange::new(start, end).unwrap()
}
#[test]
fn overlaps_is_half_open() {
assert!(range(0, 5).overlaps(&range(4, 6)));
assert!(range(4, 6).overlaps(&range(0, 5)));
assert!(range(0, 10).overlaps(&range(3, 4)));
assert!(range(3, 4).overlaps(&range(0, 10)));
assert!(!range(0, 5).overlaps(&range(5, 10)));
assert!(!range(5, 10).overlaps(&range(0, 5)));
assert!(!range(0, 2).overlaps(&range(7, 9)));
}
#[test]
fn empty_range_never_overlaps() {
assert!(!range(3, 3).overlaps(&range(0, 10)));
assert!(!range(0, 10).overlaps(&range(3, 3)));
assert!(!range(3, 3).overlaps(&range(3, 3)));
}
#[test]
fn unsupported_error_code_is_stable() {
assert_eq!(
MemoryStoreError::Unsupported {
operation: "drop_scope",
}
.error_code(),
"memory_unsupported"
);
}
fn metadata_at(
indexed_at: crate::time_compat::SystemTime,
source: MemorySource,
) -> MemoryMetadata {
MemoryMetadata {
session_id: crate::types::SessionId::new(),
source,
indexed_at,
}
}
#[test]
fn enumeration_request_admits_on_source_overlap() {
let request = MemoryEnumerationRequest {
limit: 10,
offset: 0,
source_overlap: Some(range(4, 6)),
indexed_after: None,
};
let overlapping = metadata_at(
UNIX_EPOCH,
MemorySource::Compaction {
source_range: range(0, 5),
},
);
let disjoint = metadata_at(
UNIX_EPOCH,
MemorySource::Compaction {
source_range: range(6, 9),
},
);
assert!(request.admits(&overlapping));
assert!(!request.admits(&disjoint));
}
#[test]
fn enumeration_request_indexed_after_is_strict() {
let boundary = UNIX_EPOCH + Duration::from_secs(100);
let request = MemoryEnumerationRequest {
limit: 10,
offset: 0,
source_overlap: None,
indexed_after: Some(boundary),
};
let at_boundary = metadata_at(
boundary,
MemorySource::Compaction {
source_range: range(0, 1),
},
);
let after_boundary = metadata_at(
boundary + Duration::from_secs(1),
MemorySource::Compaction {
source_range: range(0, 1),
},
);
let before_boundary = metadata_at(
UNIX_EPOCH,
MemorySource::Compaction {
source_range: range(0, 1),
},
);
assert!(!request.admits(&at_boundary));
assert!(request.admits(&after_boundary));
assert!(!request.admits(&before_boundary));
}
#[test]
fn enumeration_request_filters_compose() {
let request = MemoryEnumerationRequest {
limit: 10,
offset: 0,
source_overlap: Some(range(0, 5)),
indexed_after: Some(UNIX_EPOCH + Duration::from_secs(100)),
};
let both = metadata_at(
UNIX_EPOCH + Duration::from_secs(200),
MemorySource::Compaction {
source_range: range(2, 3),
},
);
let wrong_range = metadata_at(
UNIX_EPOCH + Duration::from_secs(200),
MemorySource::Compaction {
source_range: range(5, 9),
},
);
let too_early = metadata_at(
UNIX_EPOCH,
MemorySource::Compaction {
source_range: range(2, 3),
},
);
assert!(request.admits(&both));
assert!(!request.admits(&wrong_range));
assert!(!request.admits(&too_early));
}
struct MinimalStore;
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl MemoryStore for MinimalStore {
async fn index_scoped_batch(
&self,
batch: MemoryIndexBatch,
) -> Result<MemoryIndexReceipt, MemoryStoreError> {
let (scope, requests) = batch.into_parts();
Ok(MemoryIndexReceipt {
scope,
indexed_entries: requests.len(),
})
}
async fn search(
&self,
_scope: &MemorySearchScope,
_query: &str,
_limit: usize,
) -> Result<Vec<MemoryResult>, MemoryStoreError> {
Ok(Vec::new())
}
}
#[tokio::test]
async fn drop_scope_default_is_typed_unsupported() {
let store = MinimalStore;
let owner = MemoryOwner::canonical_session(crate::types::SessionId::new());
let error = store.drop_scope(&owner).await.unwrap_err();
assert!(matches!(
error,
MemoryStoreError::Unsupported {
operation: "drop_scope",
}
));
assert_eq!(error.error_code(), "memory_unsupported");
}
#[tokio::test]
async fn enumerate_scoped_default_is_typed_unsupported() {
let store = MinimalStore;
let scope = MemorySearchScope::for_session(crate::types::SessionId::new());
let error = store
.enumerate_scoped(
&scope,
MemoryEnumerationRequest {
limit: 10,
offset: 0,
source_overlap: None,
indexed_after: None,
},
)
.await
.unwrap_err();
assert!(matches!(
error,
MemoryStoreError::Unsupported {
operation: "enumerate_scoped",
}
));
assert_eq!(error.error_code(), "memory_unsupported");
}
}