diff --git a/Cargo.lock b/Cargo.lock index 4aef392..11b430d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1957,6 +1957,7 @@ dependencies = [ "anyhow", "async-trait", "axum", + "futures", "reqwest", "serde", "serde_json", diff --git a/README.md b/README.md index d543f32..1799f3e 100644 --- a/README.md +++ b/README.md @@ -67,21 +67,41 @@ let provider = Arc::new(tinymemory::remote::supermemory_provider(backend)); The remote adapter reaches only crates.io dependencies, so cargo resolves it without any `[patch]` entries. -**The embedded engine (TinyCortex) — three patch entries:** +**The embedded engine (TinyCortex) — vendor this repository as a submodule.** + +The remote recipe above works by git because the remote adapter reaches only +published crates. The embedded engine does not: it pulls `tinycortex`, +`tinycortex-api` and `tinyagents`, none of which are published, and +`tinycortex-api` takes `tinymemory-api` *by git*, which cargo will resolve as a +second copy of a crate this workspace also provides by path. Patching that away +needs the crates on disk, so the embedded path is a submodule dependency until +these crates are published: + +```sh +git submodule add https://github.com/tinyhumansai/tinymemory vendor/tinymemory +git -C vendor/tinymemory submodule update --init --recursive +``` ```toml [dependencies] -tinymemory = { git = "https://github.com/tinyhumansai/tinymemory", features = ["tinycortex"] } +tinymemory = { path = "vendor/tinymemory", features = ["tinycortex"] } -# The engine and its api are unpublished; without these, cargo resolves a -# second copy of each from the network and type identities split at the seam. +# All four are required. The first three are unpublished crates the engine +# needs; the fourth collapses `tinycortex-api`'s git dependency on +# `tinymemory-api` onto the copy in this tree — without it two distinct +# `tinymemory_api::MemoryEntry` types exist and the seam stops type-checking. [patch.crates-io] -tinycortex = { git = "https://github.com/tinyhumansai/tinycortex" } -tinycortex-api = { git = "https://github.com/tinyhumansai/tinycortex" } +tinycortex = { path = "vendor/tinymemory/vendor/tinycortex" } +tinycortex-api = { path = "vendor/tinymemory/vendor/tinycortex/api" } +tinyagents = { path = "vendor/tinymemory/vendor/tinyagents" } [patch."https://github.com/tinyhumansai/tinymemory"] -tinymemory-api = { git = "https://github.com/tinyhumansai/tinymemory" } +tinymemory-api = { path = "vendor/tinymemory/api" } ``` +This exact patch set is what the reference consumer in `examples/` and the +repository's own root manifest use; a build missing any of the four fails at +resolution, before compiling a line. + ```rust,ignore use std::sync::Arc; use tinymemory::tinycortex::{provider, InMemoryMemoryStore}; diff --git a/adapters/remote/Cargo.toml b/adapters/remote/Cargo.toml index 2690528..3f8fcbd 100644 --- a/adapters/remote/Cargo.toml +++ b/adapters/remote/Cargo.toml @@ -16,9 +16,14 @@ async-trait = "0.1" # The storage trait deliberately uses opaque backend errors. anyhow = "1" # Native self-hosted APIs are HTTP/JSON; multipart is required by Cognee. -reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls"] } +# `stream` is for `bytes_stream()`: response bodies are read against a byte +# cap rather than buffered whole, because the endpoint is operator-supplied +# and a broken or hostile one must not be able to OOM the host. +reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "stream"] } # Remote records are translated through a private, lossless envelope. serde = { version = "1", features = ["derive"] } +# Streaming a capped body needs a Stream combinator. +futures = "0.3" serde_json = "1" # Supermemory custom ids are bounded, so namespace/key identities use SHA-256. sha2 = "0.10" diff --git a/adapters/remote/src/common.rs b/adapters/remote/src/common.rs index 96b42d7..c8a8622 100644 --- a/adapters/remote/src/common.rs +++ b/adapters/remote/src/common.rs @@ -40,6 +40,56 @@ impl std::fmt::Debug for HttpClient { } } +/// Largest response body any hosted engine may return. +/// +/// The endpoint is operator-supplied (`SupermemoryMemory::api`, +/// `Mem0Memory::new`, `CogneeMemory::self_hosted` all take an arbitrary URL), +/// so a broken or hostile server must not be able to exhaust the host's +/// memory. 64 MiB is far above any real memory payload -- the largest thing +/// these APIs return is a page of records -- and far below a size that +/// threatens a process. +const MAX_RESPONSE_BYTES: u64 = 64 * 1024 * 1024; + +/// Read a response body, failing once it exceeds [`MAX_RESPONSE_BYTES`]. +/// +/// `Response::json()`/`text()` buffer the whole body before any size check, so +/// a server that omits or understates `Content-Length` (a chunked response, +/// say) could OOM the process despite a declared limit. Reading incrementally +/// enforces the cap while the bytes arrive. Same argument, and same shape, as +/// `tinymemory-sources`' `read_body_capped` -- that guard was written for the +/// web-page reader and simply had not been applied on this path. +async fn read_capped(response: reqwest::Response, path: &str) -> anyhow::Result> { + use futures::StreamExt; + if let Some(len) = response.content_length() { + if len > MAX_RESPONSE_BYTES { + anyhow::bail!( + "memory API {path} response exceeds {MAX_RESPONSE_BYTES}-byte limit \ + (Content-Length={len})" + ); + } + } + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| format!("memory API {path} body read failed"))?; + // Check BEFORE appending: one oversized chunk would otherwise be + // allocated in full before the limit is noticed, which is the + // allocation this cap exists to prevent. + let next_len = body + .len() + .checked_add(chunk.len()) + .context("memory API response length overflowed")?; + if next_len as u64 > MAX_RESPONSE_BYTES { + anyhow::bail!( + "memory API {path} response exceeds {MAX_RESPONSE_BYTES}-byte limit \ + (would reach {next_len} bytes)" + ); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + impl HttpClient { /// Builds a client that optionally authenticates with a bearer token. pub(crate) fn bearer(endpoint: &str, credential: Option<&str>) -> anyhow::Result { @@ -135,9 +185,8 @@ impl HttpClient { if !status.is_success() { return Err(self.status_error(path, status)); } - response - .json() - .await + let body = read_capped(response, path).await?; + serde_json::from_slice(&body) .with_context(|| format!("memory API {path} returned invalid JSON")) } @@ -153,10 +202,8 @@ impl HttpClient { if !status.is_success() { return Err(self.status_error(path, status)); } - response - .text() - .await - .context("memory API response was unreadable") + let body = read_capped(response, path).await?; + String::from_utf8(body).context("memory API response was not valid UTF-8") } /// Sends a request whose successful response body is not needed. diff --git a/adapters/remote/src/mem0.rs b/adapters/remote/src/mem0.rs index d92546d..747e94b 100644 --- a/adapters/remote/src/mem0.rs +++ b/adapters/remote/src/mem0.rs @@ -119,17 +119,52 @@ struct Mem0Dialect { } impl Mem0Dialect { + /// Largest listing this adapter will request in one call. + /// + /// Every exact-CRUD path here enumerates through [`Self::values`], so this + /// is the ceiling on the whole store, not on one page. + const LISTING_TOP_K: usize = 1000; + /// Fetches Mem0's administrative memory listing. + /// + /// # A hard ceiling, deliberately loud + /// + /// This is a single unpaginated request, and it is the ONLY enumeration + /// path in this adapter -- `get`, `list`, `count` and `export_page` all + /// route through it. Past the ceiling the results are not merely + /// incomplete, they are silently WRONG: `get(ns, key)` for a record beyond + /// the cut-off returns `Ok(None)`, which the contract defines as "no such + /// entry", so a caller reads "deleted" where the truth is "present but + /// past the window". + /// + /// Returning an error instead is the honest failure. A full response is + /// indistinguishable from a truncated one -- both are exactly `top_k` + /// items -- so this cannot detect truncation, only its own boundary, and + /// it refuses at that boundary rather than answering wrongly. Paginating + /// properly needs Mem0's paging parameters verified against a live + /// service; guessing them here would trade a loud failure for a quiet one. async fn values(&self) -> anyhow::Result> { + let top_k = Self::LISTING_TOP_K; let response: Value = self .client - .json(Method::GET, "memories?top_k=1000", None) + .json(Method::GET, &format!("memories?top_k={top_k}"), None) .await?; - Ok(response + let results = response .get("results") .and_then(Value::as_array) .cloned() - .unwrap_or_default()) + .unwrap_or_default(); + if results.len() >= top_k { + anyhow::bail!( + "mem0 returned {} memories, this adapter's unpaginated listing ceiling. \ + Exact reads (get/list/count/export) cannot be answered correctly beyond \ + it -- a record past the window would read as absent -- so the adapter \ + refuses rather than answering wrongly. Recall is unaffected (it queries \ + mem0's search API directly).", + results.len() + ); + } + Ok(results) } /// Decodes a Mem0 result containing TinyMemory-owned metadata. diff --git a/adapters/tinycortex/src/engine/mod.rs b/adapters/tinycortex/src/engine/mod.rs index a9a2c25..9a1ae2f 100644 --- a/adapters/tinycortex/src/engine/mod.rs +++ b/adapters/tinycortex/src/engine/mod.rs @@ -1024,6 +1024,7 @@ impl MemorySourceSink for TinycortexProvider { items: Vec, taint: MemoryTaint, ) -> Result { + let items_len = items.len(); let namespace = format!("source:{source_id}"); let mut outcome = IngestOutcome::default(); for item in items { @@ -1063,8 +1064,20 @@ impl MemorySourceSink for TinycortexProvider { outcome.written = outcome.written.saturating_add(1); outcome.ids.push(id); } - Err(_) => { - outcome.skipped = outcome.skipped.saturating_add(1); + // A write failure is NOT `skipped`. The contract defines that + // field as "units the driver recognised as already present" + // (`IngestOutcome::skipped`), so counting a failed write there + // reports a locked database, a full disk or a dead embedder as + // a successful no-op: the sync caller marks the items done and + // they are never written. Propagate instead — a partial batch + // has no truthful representation in `IngestOutcome`, and a + // caller that wants best-effort ingestion can catch this. + Err(error) => { + return Err(MemoryError::Other(anyhow::anyhow!( + "source ingest failed after {} of {} item(s) were written: {error}", + outcome.written, + items_len + ))); } } } diff --git a/core/src/store/namespace_store/documents.rs b/core/src/store/namespace_store/documents.rs index 9c317f4..7c39a6c 100644 --- a/core/src/store/namespace_store/documents.rs +++ b/core/src/store/namespace_store/documents.rs @@ -33,6 +33,19 @@ impl UnifiedMemory { if key.is_empty() { return Err("document key cannot be empty".to_string()); } + // Serialise writers of one key for the WHOLE operation. A deterministic + // document id stops two writers orphaning each other's chunks, but it + // does not ORDER them: the row write and the chunk replacement are + // separated by embedding, which awaits. Without this, writer A can + // update the row, await the embedder, and have B update the row and + // replace the chunks in between -- leaving B's content beside A's + // chunks, plus A's trailing chunks if A had more. The metadata-only + // path below takes the same lock: it writes the same row, so it must + // not interleave with a full write either. Same guard shape as the + // sync path's per-connection lock. + let _write_guard = Self::document_write_lock(&self.db_path, &namespace, &key) + .lock_owned() + .await; let existing_document_id = { let conn = self.conn.lock(); conn.query_row( @@ -46,11 +59,19 @@ impl UnifiedMemory { let document_id = input .document_id .or(existing_document_id) - .unwrap_or_else(|| { - let ts = Self::now_ts() as u64; - let short = &Uuid::new_v4().to_string()[..8]; - format!("{ts}_{short}") - }); + // Derived from (namespace, key), NOT random. The lookup above and + // the write below are separated by `.await`s, so two concurrent + // stores of a not-yet-existing key both miss and both mint an id. + // The ROW is safe -- `ON CONFLICT(namespace, key) DO UPDATE` keeps + // exactly one -- but that clause does not update `document_id`, + // and each writer has already written `vector_chunks` under ITS + // OWN id. The loser's chunks are then unreachable from the row, so + // `forget` (which deletes chunks by the row's document_id) leaves + // them behind and recall keeps returning content the caller + // deleted. A deterministic id makes both writers choose the same + // one, so the second write updates the first's chunks instead of + // orphaning them. + .unwrap_or_else(|| Self::derive_document_id(&namespace, &key)); let now = Self::now_ts(); let created_at = { let conn = self.conn.lock(); @@ -221,6 +242,19 @@ impl UnifiedMemory { if key.is_empty() { return Err("document key cannot be empty".to_string()); } + // Serialise writers of one key for the WHOLE operation. A deterministic + // document id stops two writers orphaning each other's chunks, but it + // does not ORDER them: the row write and the chunk replacement are + // separated by embedding, which awaits. Without this, writer A can + // update the row, await the embedder, and have B update the row and + // replace the chunks in between -- leaving B's content beside A's + // chunks, plus A's trailing chunks if A had more. The metadata-only + // path below takes the same lock: it writes the same row, so it must + // not interleave with a full write either. Same guard shape as the + // sync path's per-connection lock. + let _write_guard = Self::document_write_lock(&self.db_path, &namespace, &key) + .lock_owned() + .await; let existing_document_id = { let conn = self.conn.lock(); conn.query_row( @@ -659,8 +693,126 @@ impl UnifiedMemory { } Ok(json!({"deleted": deleted, "namespace": ns, "documentId": document_id })) } + + /// The write lock for one `(database, namespace, key)`. + /// + /// Process-global rather than per-instance: the same store file can be + /// opened by more than one `UnifiedMemory`, and a lock living on the + /// instance would not serialise those. Keyed by db path so two workspaces + /// never contend. + /// + /// The table only grows, bounded by the number of distinct keys this + /// process has written -- one `Arc` and an unlocked mutex each. + fn document_write_lock( + db_path: &std::path::Path, + namespace: &str, + key: &str, + ) -> std::sync::Arc> { + use std::collections::HashMap; + use std::sync::{Arc, Mutex, OnceLock}; + type Table = Mutex>>>; + static LOCKS: OnceLock = OnceLock::new(); + let table = LOCKS.get_or_init(|| Mutex::new(HashMap::new())); + let id = ( + db_path.to_string_lossy().into_owned(), + namespace.to_owned(), + key.to_owned(), + ); + // Recover from a poisoned table: it holds `Arc`s only, so a panicking + // writer leaves nothing torn, and refusing every later write would be + // a worse failure than continuing. + let mut table = table + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + Arc::clone(table.entry(id).or_default()) + } + + /// A document id derived from `(namespace, key)`. + /// + /// Deterministic so two concurrent first-writes of one key agree, which is + /// what keeps `vector_chunks` addressable from the row. Hashed rather than + /// concatenated so the id is a fixed-width opaque token whatever the + /// namespace or key contains; the zero byte is a domain separator, so + /// ("a","bc") and ("ab","c") cannot collide. + pub(crate) fn derive_document_id(namespace: &str, key: &str) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(namespace.as_bytes()); + hasher.update([0u8]); + hasher.update(key.as_bytes()); + format!("{:x}", hasher.finalize())[..32].to_string() + } } #[cfg(test)] #[path = "documents_tests.rs"] mod tests; + +#[cfg(test)] +mod document_id_tests { + use super::UnifiedMemory; + + /// Two concurrent first-writes of one key must choose the SAME document + /// id. If they do not, each writes `vector_chunks` under its own id, the + /// `ON CONFLICT(namespace, key)` row keeps only one of them, and the + /// loser's chunks outlive `forget` — deleted content stays recallable. + #[test] + fn the_id_is_derived_from_namespace_and_key_not_random() { + let a = UnifiedMemory::derive_document_id("notes", "q3-plan"); + let b = UnifiedMemory::derive_document_id("notes", "q3-plan"); + assert_eq!(a, b, "the same key must derive the same id"); + assert_ne!( + a, + UnifiedMemory::derive_document_id("notes", "q4-plan"), + "different keys must not collide" + ); + assert_ne!( + a, + UnifiedMemory::derive_document_id("other", "q3-plan"), + "the namespace must participate" + ); + } + + /// The guard must be per key, not global: two different keys writing at + /// once must not serialise, or every concurrent write in the process + /// queues behind one slow embedding. + #[test] + fn the_write_lock_is_per_key_and_shared_per_key() { + let db = std::path::Path::new("/w/memory/memory.db"); + let a1 = UnifiedMemory::document_write_lock(db, "notes", "k1"); + let a2 = UnifiedMemory::document_write_lock(db, "notes", "k1"); + let b = UnifiedMemory::document_write_lock(db, "notes", "k2"); + let other_ns = UnifiedMemory::document_write_lock(db, "other", "k1"); + let other_db = UnifiedMemory::document_write_lock( + std::path::Path::new("/w2/memory/memory.db"), + "notes", + "k1", + ); + assert!( + std::sync::Arc::ptr_eq(&a1, &a2), + "same key must share one lock" + ); + assert!( + !std::sync::Arc::ptr_eq(&a1, &b), + "different keys must not contend" + ); + assert!( + !std::sync::Arc::ptr_eq(&a1, &other_ns), + "the namespace must participate" + ); + assert!( + !std::sync::Arc::ptr_eq(&a1, &other_db), + "two workspaces must not contend" + ); + } + + /// The separator matters: without it ("a","bc") and ("ab","c") hash the + /// same bytes and two distinct records share one id. + #[test] + fn the_namespace_key_boundary_cannot_be_shifted() { + assert_ne!( + UnifiedMemory::derive_document_id("a", "bc"), + UnifiedMemory::derive_document_id("ab", "c") + ); + } +} diff --git a/core/src/store/namespace_store/init.rs b/core/src/store/namespace_store/init.rs index a8d0093..c3296f5 100644 --- a/core/src/store/namespace_store/init.rs +++ b/core/src/store/namespace_store/init.rs @@ -405,7 +405,7 @@ impl UnifiedMemory { if trimmed.is_empty() { return GLOBAL_NAMESPACE.to_string(); } - trimmed + let sanitized: String = trimmed .chars() .map(|ch| { if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '/' { @@ -414,7 +414,19 @@ impl UnifiedMemory { '_' } }) - .collect() + .collect(); + // `/` is kept so a namespace can be hierarchical, but a LEADING one + // makes the result an absolute path — and `Path::join` with an + // absolute path discards the base entirely, so + // `memory_dir/namespaces/` vanishes and the namespace addresses + // anywhere on the filesystem. `clear_namespace` calls + // `remove_dir_all` on that path. (`..` is already neutralised above: + // `.` is not in the allow-list, so it becomes `_`.) + let sanitized = sanitized.trim_start_matches('/').to_string(); + if sanitized.is_empty() { + return GLOBAL_NAMESPACE.to_string(); + } + sanitized } /// Resolved memory subdirectory for this store instance (e.g. @@ -474,6 +486,36 @@ mod tests { } } + /// A namespace beginning with `/` must not escape the workspace: + /// `Path::join` with an absolute path DISCARDS the base, so + /// `memory_dir/namespaces/` would vanish and `clear_namespace`'s + /// `remove_dir_all` would run against an arbitrary absolute path. + #[test] + fn a_namespace_cannot_escape_the_workspace() { + for hostile in [ + "/Users/me/Documents", + "//tmp/x", + "///etc", + "/", + "a/../../etc", + "../../etc", + ] { + let sanitized = UnifiedMemory::sanitize_namespace(hostile); + assert!( + !sanitized.starts_with('/'), + "{hostile:?} sanitized to {sanitized:?}, which is absolute" + ); + let dir = std::path::Path::new("/w/memory") + .join("namespaces") + .join(&sanitized); + assert!( + dir.starts_with("/w/memory/namespaces"), + "{hostile:?} escaped to {}", + dir.display() + ); + } + } + #[test] fn namespace_dir_uses_sanitized_namespace() { let tmp = TempDir::new().unwrap();