#[cfg(feature = "local")]
mod local;
#[cfg(feature = "client")]
mod remote;
#[cfg(all(feature = "client", feature = "local"))]
use crate::{MemoryAccess, MemorySource, secrets::contains_likely_secret};
use crate::{
MemoryKey, MemoryLimits, MemoryRecord, MemoryScan,
server::protocol::{self, ExportCursor, SyncReport},
};
#[cfg(feature = "local")]
pub use local::LocalMemoryStore;
#[cfg(feature = "client")]
pub use remote::{RemoteClientError, RemoteMemoryClient, RemoteToken};
#[cfg(all(feature = "client", feature = "local"))]
use std::path::PathBuf;
#[cfg(feature = "local")]
use std::time::{SystemTime, UNIX_EPOCH};
use std::{error::Error, future::Future};
use thiserror::Error;
pub trait MemoryStore: Clone + Send + Sync + 'static {
fn scan(
&self,
query: &str,
limit: usize,
) -> impl Future<Output = Result<MemoryScan, MemoryError>> + Send;
fn read(
&self,
ids: &[i64],
keys: &[MemoryKey],
) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send;
fn list(&self) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send;
fn put(
&self,
content: &str,
replacement: Option<MemoryKey>,
) -> impl Future<Output = Result<MemoryRecord, MemoryError>> + Send;
fn delete(&self, key: MemoryKey) -> impl Future<Output = Result<(), MemoryError>> + Send;
fn sync(
&self,
memories: &[MemoryRecord],
) -> impl Future<Output = Result<SyncReport, MemoryError>> + Send;
fn export_page(
&self,
namespaces: Option<&[String]>,
cursor: Option<&ExportCursor>,
limit: usize,
) -> impl Future<Output = Result<(Vec<MemoryRecord>, Option<ExportCursor>), MemoryError>> + Send;
fn export_all(
&self,
namespaces: Option<&[String]>,
) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
async move {
let mut cursor = None;
let mut records = Vec::new();
let mut content_bytes = 0usize;
loop {
let (page, next_cursor) = self
.export_page(
namespaces,
cursor.as_ref(),
protocol::MAX_EXPORT_PAGE_RECORDS,
)
.await?;
let next_record_count = records
.len()
.checked_add(page.len())
.ok_or(MemoryError::InvalidPagination)?;
let page_bytes = page.iter().try_fold(0usize, |total, record| {
total.checked_add(record.content.len())
});
content_bytes = content_bytes
.checked_add(page_bytes.ok_or(MemoryError::InvalidPagination)?)
.ok_or(MemoryError::InvalidPagination)?;
if next_record_count > MemoryLimits::PRODUCTION.records
|| content_bytes > MemoryLimits::PRODUCTION.total_content_bytes
|| next_cursor
.as_ref()
.is_some_and(|next| cursor.as_ref() == Some(next))
|| (page.is_empty() && next_cursor.is_some())
{
return Err(MemoryError::InvalidPagination);
}
records.extend(page);
match next_cursor {
Some(next) => cursor = Some(next),
None => return Ok(records),
}
}
}
}
}
#[cfg(all(feature = "client", feature = "local"))]
#[derive(Clone, Debug)]
pub enum SelectedMemoryStore {
Local(LocalMemoryStore),
Remote(RemoteMemoryClient),
}
#[cfg(all(feature = "client", feature = "local"))]
impl SelectedMemoryStore {
pub fn local(path: impl Into<PathBuf>) -> Self {
Self::Local(LocalMemoryStore::new(path))
}
pub const fn remote(client: RemoteMemoryClient) -> Self {
Self::Remote(client)
}
pub const fn source(&self) -> MemorySource {
match self {
Self::Local(_) => MemorySource::Local,
Self::Remote(_) => MemorySource::Remote,
}
}
pub async fn access(&self) -> Result<MemoryAccess, MemoryError> {
match self {
Self::Local(_) => Ok(MemoryAccess {
source: MemorySource::Local,
namespace: None,
role: None,
}),
Self::Remote(client) => Ok(MemoryAccess {
source: MemorySource::Remote,
namespace: Some(client.namespace().to_owned()),
role: Some(client.session().await?),
}),
}
}
}
#[cfg(all(feature = "client", feature = "local"))]
impl MemoryStore for SelectedMemoryStore {
fn scan(
&self,
query: &str,
limit: usize,
) -> impl Future<Output = Result<MemoryScan, MemoryError>> + Send {
async move {
match self {
Self::Local(store) => MemoryStore::scan(store, query, limit).await,
Self::Remote(client) => MemoryStore::scan(client, query, limit).await,
}
}
}
fn read(
&self,
ids: &[i64],
keys: &[MemoryKey],
) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
async move {
match self {
Self::Local(store) => MemoryStore::read(store, ids, keys).await,
Self::Remote(client) => MemoryStore::read(client, ids, keys).await,
}
}
}
fn list(&self) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
async move {
match self {
Self::Local(store) => MemoryStore::list(store).await,
Self::Remote(client) => MemoryStore::list(client).await,
}
}
}
fn put(
&self,
content: &str,
replacement: Option<MemoryKey>,
) -> impl Future<Output = Result<MemoryRecord, MemoryError>> + Send {
async move {
reject_unsafe(content)?;
match self {
Self::Local(store) => MemoryStore::put(store, content, replacement).await,
Self::Remote(client) => MemoryStore::put(client, content, replacement).await,
}
}
}
fn delete(&self, key: MemoryKey) -> impl Future<Output = Result<(), MemoryError>> + Send {
async move {
match self {
Self::Local(store) => MemoryStore::delete(store, key).await,
Self::Remote(client) => MemoryStore::delete(client, key).await,
}
}
}
fn sync(
&self,
memories: &[MemoryRecord],
) -> impl Future<Output = Result<SyncReport, MemoryError>> + Send {
async move {
for memory in memories {
reject_unsafe(&memory.content)?;
}
match self {
Self::Local(store) => MemoryStore::sync(store, memories).await,
Self::Remote(client) => MemoryStore::sync(client, memories).await,
}
}
}
fn export_page(
&self,
namespaces: Option<&[String]>,
cursor: Option<&ExportCursor>,
limit: usize,
) -> impl Future<Output = Result<(Vec<MemoryRecord>, Option<ExportCursor>), MemoryError>> + Send
{
async move {
match self {
Self::Local(store) => {
MemoryStore::export_page(store, namespaces, cursor, limit).await
}
Self::Remote(client) => {
MemoryStore::export_page(client, namespaces, cursor, limit).await
}
}
}
}
}
#[cfg(all(feature = "client", feature = "local"))]
fn reject_unsafe(content: &str) -> Result<(), MemoryError> {
if contains_likely_secret(content) {
return Err(MemoryError::SecretRejected);
}
Ok(())
}
#[cfg(feature = "local")]
fn current_time_ms() -> i64 {
let milliseconds = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
i64::try_from(milliseconds).unwrap_or(i64::MAX)
}
#[derive(Debug, Error)]
pub enum MemoryError {
#[error("memory content is empty")]
EmptyContent,
#[error("memory content exceeds the {maximum_bytes}-byte limit")]
ContentTooLarge {
maximum_bytes: usize,
},
#[error("memory query exceeds the {maximum_bytes}-byte limit")]
QueryTooLarge {
maximum_bytes: usize,
},
#[error("memory record capacity of {maximum} was reached")]
RecordCapacity {
maximum: usize,
},
#[error("memory content capacity of {maximum_bytes} bytes was reached")]
ContentCapacity {
maximum_bytes: usize,
},
#[error("memory storage capacity was reached")]
StorageCapacity,
#[error("memory content was rejected as a likely secret")]
SecretRejected,
#[error("an equivalent memory already exists")]
Duplicate,
#[error("memory was not found")]
NotFound,
#[error("memory changed since it was read")]
Conflict,
#[error("memories from other namespaces are read-only")]
RemoteReadOnly,
#[error("memory store returned invalid pagination")]
InvalidPagination,
#[error(
"memory schema version {found} is unsupported; this build supports version {supported}"
)]
UnsupportedSchemaVersion {
found: i64,
supported: i64,
},
#[error("memory backend operation failed")]
Backend {
#[source]
source: Box<dyn Error + Send + Sync>,
},
#[error("memory backend is temporarily unavailable")]
Unavailable {
#[source]
source: Box<dyn Error + Send + Sync>,
},
}
impl MemoryError {
pub fn backend(source: impl Error + Send + Sync + 'static) -> Self {
Self::Backend {
source: Box::new(source),
}
}
pub fn unavailable(source: impl Error + Send + Sync + 'static) -> Self {
Self::Unavailable {
source: Box::new(source),
}
}
pub fn is_retryable(&self) -> bool {
matches!(self, Self::Unavailable { .. })
}
}
#[cfg(feature = "client")]
impl From<RemoteClientError> for MemoryError {
fn from(source: RemoteClientError) -> Self {
match source {
error @ (RemoteClientError::Transport | RemoteClientError::Unavailable) => {
Self::unavailable(error)
}
RemoteClientError::ReadOnly | RemoteClientError::NamespaceMismatch => {
Self::RemoteReadOnly
}
RemoteClientError::Rejected { code } => match code {
protocol::RemoteErrorCode::QueryTooLarge => Self::QueryTooLarge {
maximum_bytes: MemoryLimits::PRODUCTION.query_bytes,
},
protocol::RemoteErrorCode::ContentTooLarge => Self::ContentTooLarge {
maximum_bytes: MemoryLimits::PRODUCTION.content_bytes,
},
protocol::RemoteErrorCode::RecordCapacity => Self::RecordCapacity {
maximum: MemoryLimits::PRODUCTION.records,
},
protocol::RemoteErrorCode::ContentCapacity => Self::ContentCapacity {
maximum_bytes: MemoryLimits::PRODUCTION.total_content_bytes,
},
protocol::RemoteErrorCode::Duplicate => Self::Duplicate,
protocol::RemoteErrorCode::NotFound => Self::NotFound,
protocol::RemoteErrorCode::Conflict => Self::Conflict,
protocol::RemoteErrorCode::Forbidden
| protocol::RemoteErrorCode::NamespaceMismatch => Self::RemoteReadOnly,
protocol::RemoteErrorCode::Unavailable => {
Self::unavailable(RemoteClientError::Rejected { code })
}
_ => Self::backend(RemoteClientError::Rejected { code }),
},
error => Self::backend(error),
}
}
}