use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use crate::auth::TenantScope;
use crate::error::Error;
use super::{Confidentiality, Memory, MemoryEntry, MemoryQuery};
pub struct NamespacedMemory {
inner: Arc<dyn Memory>,
agent_name: String,
max_confidentiality: Option<Confidentiality>,
default_store_confidentiality: Confidentiality,
}
impl NamespacedMemory {
pub fn new(inner: Arc<dyn Memory>, agent_name: impl Into<String>) -> Self {
Self {
inner,
agent_name: agent_name.into(),
max_confidentiality: None,
default_store_confidentiality: Confidentiality::Public,
}
}
pub fn with_max_confidentiality(mut self, cap: Option<Confidentiality>) -> Self {
self.max_confidentiality = cap;
self
}
pub fn with_default_store_confidentiality(mut self, level: Confidentiality) -> Self {
self.default_store_confidentiality = level;
self
}
fn prefix_id(&self, id: &str) -> String {
format!("{}:{}", self.agent_name, id)
}
}
impl Memory for NamespacedMemory {
fn store(
&self,
scope: &TenantScope,
mut entry: MemoryEntry,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
entry.id = self.prefix_id(&entry.id);
entry.agent = self.agent_name.clone();
if entry.confidentiality < self.default_store_confidentiality {
entry.confidentiality = self.default_store_confidentiality;
}
let scope = scope.clone();
Box::pin(async move { self.inner.store(&scope, entry).await })
}
fn recall(
&self,
scope: &TenantScope,
query: MemoryQuery,
) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryEntry>, Error>> + Send + '_>> {
let mut query = MemoryQuery {
agent: Some(self.agent_name.clone()),
..query
};
if let Some(cap) = self.max_confidentiality {
query.max_confidentiality = Some(match query.max_confidentiality {
Some(existing) if existing < cap => existing,
_ => cap,
});
}
let prefix = format!("{}:", self.agent_name);
let scope = scope.clone();
Box::pin(async move {
let mut entries = self.inner.recall(&scope, query).await?;
for entry in &mut entries {
if let Some(stripped) = entry.id.strip_prefix(&prefix) {
entry.id = stripped.to_string();
}
}
Ok(entries)
})
}
fn update(
&self,
scope: &TenantScope,
id: &str,
content: String,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
let prefixed = self.prefix_id(id);
let scope = scope.clone();
Box::pin(async move { self.inner.update(&scope, &prefixed, content).await })
}
fn forget(
&self,
scope: &TenantScope,
id: &str,
) -> Pin<Box<dyn Future<Output = Result<bool, Error>> + Send + '_>> {
let prefixed = self.prefix_id(id);
let scope = scope.clone();
Box::pin(async move { self.inner.forget(&scope, &prefixed).await })
}
fn add_link(
&self,
scope: &TenantScope,
id: &str,
related_id: &str,
) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + '_>> {
let prefixed_id = self.prefix_id(id);
let prefixed_related = self.prefix_id(related_id);
let scope = scope.clone();
Box::pin(async move {
self.inner
.add_link(&scope, &prefixed_id, &prefixed_related)
.await
})
}
fn prune(
&self,
scope: &TenantScope,
min_strength: f64,
min_age: chrono::Duration,
_agent_prefix: Option<&str>,
) -> Pin<Box<dyn Future<Output = Result<usize, Error>> + Send + '_>> {
let scope = scope.clone();
let agent_name = self.agent_name.clone();
Box::pin(async move {
self.inner
.prune(&scope, min_strength, min_age, Some(&agent_name))
.await
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memory::in_memory::InMemoryStore;
use chrono::Utc;
use super::super::{Confidentiality, MemoryType};
fn test_scope() -> TenantScope {
TenantScope::default()
}
fn make_entry(id: &str, content: &str) -> MemoryEntry {
MemoryEntry {
id: id.into(),
agent: String::new(),
content: content.into(),
category: "fact".into(),
tags: vec![],
created_at: Utc::now(),
last_accessed: Utc::now(),
access_count: 0,
importance: 5,
memory_type: MemoryType::default(),
keywords: vec![],
summary: None,
strength: 1.0,
related_ids: vec![],
source_ids: vec![],
embedding: None,
confidentiality: Confidentiality::default(),
author_user_id: None,
author_tenant_id: None,
}
}
#[tokio::test]
async fn store_prefixes_id_and_agent() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "researcher");
ns.store(&test_scope(), make_entry("m1", "test data"))
.await
.unwrap();
let all = inner
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 1);
assert_eq!(all[0].id, "researcher:m1");
assert_eq!(all[0].agent, "researcher");
let ns_results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(ns_results[0].id, "m1"); }
#[tokio::test]
async fn recall_filters_by_agent() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns_a = NamespacedMemory::new(inner.clone(), "agent_a");
let ns_b = NamespacedMemory::new(inner.clone(), "agent_b");
ns_a.store(&test_scope(), make_entry("m1", "data from A"))
.await
.unwrap();
ns_b.store(&test_scope(), make_entry("m2", "data from B"))
.await
.unwrap();
let results = ns_a
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].content, "data from A");
let results = ns_b
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].content, "data from B");
}
#[tokio::test]
async fn namespace_forces_own_agent_even_with_explicit_override() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns_a = NamespacedMemory::new(inner.clone(), "agent_a");
let ns_b = NamespacedMemory::new(inner.clone(), "agent_b");
ns_a.store(&test_scope(), make_entry("m1", "from A"))
.await
.unwrap();
ns_b.store(&test_scope(), make_entry("m2", "from B"))
.await
.unwrap();
let results = ns_a
.recall(
&test_scope(),
MemoryQuery {
agent: Some(String::new()),
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].content, "from A");
let all = inner
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 2);
}
#[tokio::test]
async fn recall_then_update_roundtrip() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a");
ns.store(&test_scope(), make_entry("m1", "original"))
.await
.unwrap();
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results[0].id, "m1");
ns.update(
&test_scope(),
&results[0].id,
"updated via recall ID".into(),
)
.await
.unwrap();
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results[0].content, "updated via recall ID");
}
#[tokio::test]
async fn update_uses_prefixed_id() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a");
ns.store(&test_scope(), make_entry("m1", "original"))
.await
.unwrap();
ns.update(&test_scope(), "m1", "updated".into())
.await
.unwrap();
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results[0].content, "updated");
}
#[tokio::test]
async fn forget_uses_prefixed_id() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a");
ns.store(&test_scope(), make_entry("m1", "to delete"))
.await
.unwrap();
assert!(ns.forget(&test_scope(), "m1").await.unwrap());
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert!(results.is_empty());
}
#[tokio::test]
async fn add_link_delegates_with_prefix() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a");
ns.store(&test_scope(), make_entry("m1", "first"))
.await
.unwrap();
ns.store(&test_scope(), make_entry("m2", "second"))
.await
.unwrap();
ns.add_link(&test_scope(), "m1", "m2").await.unwrap();
let all = inner
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
let m1 = all.iter().find(|e| e.id == "agent_a:m1").unwrap();
let m2 = all.iter().find(|e| e.id == "agent_a:m2").unwrap();
assert!(m1.related_ids.contains(&"agent_a:m2".to_string()));
assert!(m2.related_ids.contains(&"agent_a:m1".to_string()));
}
#[tokio::test]
async fn max_confidentiality_caps_recall() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a")
.with_max_confidentiality(Some(Confidentiality::Public));
let mut public_entry = make_entry("m1", "public data");
public_entry.confidentiality = Confidentiality::Public;
ns.store(&test_scope(), public_entry).await.unwrap();
let mut confidential_entry = make_entry("m2", "confidential data");
confidential_entry.confidentiality = Confidentiality::Confidential;
confidential_entry.id = "agent_a:m2".into();
confidential_entry.agent = "agent_a".into();
inner
.store(&test_scope(), confidential_entry)
.await
.unwrap();
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].content, "public data");
}
#[tokio::test]
async fn no_confidentiality_cap_returns_all() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a");
let mut public_entry = make_entry("m1", "public data");
public_entry.confidentiality = Confidentiality::Public;
ns.store(&test_scope(), public_entry).await.unwrap();
let mut confidential_entry = make_entry("m2", "confidential data");
confidential_entry.confidentiality = Confidentiality::Confidential;
ns.store(&test_scope(), confidential_entry).await.unwrap();
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 2);
}
#[tokio::test]
async fn confidentiality_cap_uses_stricter_of_two() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a")
.with_max_confidentiality(Some(Confidentiality::Internal));
let mut public_entry = make_entry("m1", "public data");
public_entry.confidentiality = Confidentiality::Public;
ns.store(&test_scope(), public_entry).await.unwrap();
let mut internal_entry = make_entry("m2", "internal data");
internal_entry.confidentiality = Confidentiality::Internal;
ns.store(&test_scope(), internal_entry).await.unwrap();
let mut confidential_entry = make_entry("m3", "confidential data");
confidential_entry.confidentiality = Confidentiality::Confidential;
confidential_entry.id = "agent_a:m3".into();
confidential_entry.agent = "agent_a".into();
inner
.store(&test_scope(), confidential_entry)
.await
.unwrap();
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
max_confidentiality: Some(Confidentiality::Confidential),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 2);
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
max_confidentiality: Some(Confidentiality::Public),
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 1); }
#[tokio::test]
async fn default_store_confidentiality_upgrades_public() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "tg_agent")
.with_default_store_confidentiality(Confidentiality::Confidential);
let entry = make_entry("m1", "private chat data");
ns.store(&test_scope(), entry).await.unwrap();
let all = inner
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 1);
assert_eq!(all[0].confidentiality, Confidentiality::Confidential);
}
#[tokio::test]
async fn default_store_confidentiality_enforces_minimum_floor() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "tg_agent")
.with_default_store_confidentiality(Confidentiality::Confidential);
let mut entry = make_entry("m1", "internal data");
entry.confidentiality = Confidentiality::Internal;
ns.store(&test_scope(), entry).await.unwrap();
let all = inner
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 1);
assert_eq!(all[0].confidentiality, Confidentiality::Confidential);
}
#[tokio::test]
async fn default_store_confidentiality_preserves_higher_level() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "tg_agent")
.with_default_store_confidentiality(Confidentiality::Confidential);
let mut entry = make_entry("m1", "secret data");
entry.confidentiality = Confidentiality::Restricted;
ns.store(&test_scope(), entry).await.unwrap();
let all = inner
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 1);
assert_eq!(all[0].confidentiality, Confidentiality::Restricted);
}
#[tokio::test]
async fn prune_delegates_to_inner() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns = NamespacedMemory::new(inner.clone(), "agent_a");
let mut entry = make_entry("m1", "weak memory");
entry.strength = 0.01;
entry.created_at = Utc::now() - chrono::Duration::hours(48);
entry.last_accessed = Utc::now() - chrono::Duration::hours(48);
ns.store(&test_scope(), entry).await.unwrap();
let pruned = ns
.prune(&test_scope(), 0.1, chrono::Duration::hours(1), None)
.await
.unwrap();
assert_eq!(pruned, 1);
let results = ns
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert!(results.is_empty());
}
#[tokio::test]
async fn prune_scoped_to_own_namespace() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns_a = NamespacedMemory::new(inner.clone(), "agent_a");
let ns_b = NamespacedMemory::new(inner.clone(), "agent_b");
let mut weak_a = make_entry("m1", "weak from A");
weak_a.strength = 0.01;
weak_a.created_at = Utc::now() - chrono::Duration::hours(48);
weak_a.last_accessed = Utc::now() - chrono::Duration::hours(48);
ns_a.store(&test_scope(), weak_a).await.unwrap();
let mut weak_b = make_entry("m1", "weak from B");
weak_b.strength = 0.01;
weak_b.created_at = Utc::now() - chrono::Duration::hours(48);
weak_b.last_accessed = Utc::now() - chrono::Duration::hours(48);
ns_b.store(&test_scope(), weak_b).await.unwrap();
let pruned = ns_a
.prune(&test_scope(), 0.1, chrono::Duration::hours(1), None)
.await
.unwrap();
assert_eq!(pruned, 1, "should only prune agent_a's entry");
let a_results = ns_a
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert!(a_results.is_empty());
let b_results = ns_b
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(
b_results.len(),
1,
"agent_b's entry must survive agent_a's prune"
);
assert_eq!(b_results[0].content, "weak from B");
}
#[tokio::test]
async fn multi_tenant_prune_isolation() {
let shared: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let alice = NamespacedMemory::new(shared.clone(), "user:alice");
let bob = NamespacedMemory::new(shared.clone(), "user:bob");
let mut weak_alice = make_entry("m1", "alice weak");
weak_alice.strength = 0.01;
weak_alice.created_at = Utc::now() - chrono::Duration::hours(48);
weak_alice.last_accessed = Utc::now() - chrono::Duration::hours(48);
alice.store(&test_scope(), weak_alice).await.unwrap();
let mut strong_alice = make_entry("m2", "alice strong");
strong_alice.strength = 0.9;
alice.store(&test_scope(), strong_alice).await.unwrap();
let mut weak_bob = make_entry("m1", "bob weak");
weak_bob.strength = 0.01;
weak_bob.created_at = Utc::now() - chrono::Duration::hours(48);
weak_bob.last_accessed = Utc::now() - chrono::Duration::hours(48);
bob.store(&test_scope(), weak_bob).await.unwrap();
let mut strong_bob = make_entry("m2", "bob strong");
strong_bob.strength = 0.9;
bob.store(&test_scope(), strong_bob).await.unwrap();
let pruned = alice
.prune(&test_scope(), 0.1, chrono::Duration::hours(1), None)
.await
.unwrap();
assert_eq!(pruned, 1, "should only prune alice's weak entry");
let pruned = bob
.prune(&test_scope(), 0.1, chrono::Duration::hours(1), None)
.await
.unwrap();
assert_eq!(pruned, 1, "bob's weak entry must survive alice's prune");
let alice_results = alice
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(alice_results.len(), 1);
assert_eq!(alice_results[0].content, "alice strong");
let bob_results = bob
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(bob_results.len(), 1);
assert_eq!(bob_results[0].content, "bob strong");
}
#[tokio::test]
async fn prune_ignores_explicit_agent_prefix_override() {
let shared: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let alice = NamespacedMemory::new(shared.clone(), "user:alice");
let bob = NamespacedMemory::new(shared.clone(), "user:bob");
let mut weak_bob = make_entry("m1", "bob weak");
weak_bob.strength = 0.01;
weak_bob.created_at = Utc::now() - chrono::Duration::hours(48);
weak_bob.last_accessed = Utc::now() - chrono::Duration::hours(48);
bob.store(&test_scope(), weak_bob).await.unwrap();
let pruned = alice
.prune(
&test_scope(),
0.1,
chrono::Duration::hours(1),
Some("user:bob"),
)
.await
.unwrap();
assert_eq!(
pruned, 0,
"alice's prune must not affect bob even with explicit prefix"
);
let bob_results = bob
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(bob_results.len(), 1, "bob's entry must survive");
}
#[tokio::test]
async fn recall_ignores_explicit_agent_override() {
let inner: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let ns_a = NamespacedMemory::new(inner.clone(), "user:alice");
let ns_b = NamespacedMemory::new(inner.clone(), "user:bob");
ns_a.store(&test_scope(), make_entry("m1", "alice data"))
.await
.unwrap();
ns_b.store(&test_scope(), make_entry("m1", "bob data"))
.await
.unwrap();
let results = ns_a
.recall(
&test_scope(),
MemoryQuery {
agent: Some("user:bob".into()),
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].content, "alice data");
}
#[tokio::test]
async fn per_user_namespace_isolation() {
let shared: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let alice = NamespacedMemory::new(shared.clone(), "user:alice");
let bob = NamespacedMemory::new(shared.clone(), "user:bob");
alice
.store(&test_scope(), make_entry("m1", "Alice's deal notes"))
.await
.unwrap();
bob.store(&test_scope(), make_entry("m1", "Bob's pipeline review"))
.await
.unwrap();
let alice_results = alice
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(alice_results.len(), 1);
assert_eq!(alice_results[0].content, "Alice's deal notes");
assert_eq!(alice_results[0].id, "m1");
let bob_results = bob
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(bob_results.len(), 1);
assert_eq!(bob_results[0].content, "Bob's pipeline review");
let all = shared
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 2);
let ids: Vec<&str> = all.iter().map(|e| e.id.as_str()).collect();
assert!(ids.contains(&"user:alice:m1"));
assert!(ids.contains(&"user:bob:m1"));
}
#[tokio::test]
async fn per_user_can_coexist_with_shared_institutional_memory() {
let shared: Arc<dyn Memory> = Arc::new(InMemoryStore::new());
let mut institutional = make_entry("shared:playbook", "Always follow up within 24h");
institutional.agent = "shared".into();
institutional.id = "shared:playbook".into();
shared.store(&test_scope(), institutional).await.unwrap();
let alice = NamespacedMemory::new(shared.clone(), "user:alice");
alice
.store(&test_scope(), make_entry("m1", "Alice's note"))
.await
.unwrap();
let alice_results = alice
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(alice_results.len(), 1);
assert_eq!(alice_results[0].content, "Alice's note");
let all = shared
.recall(
&test_scope(),
MemoryQuery {
limit: 10,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(all.len(), 2);
}
}