From a21fb5b17ac24c62edb941ee18326337f5a0a2dc Mon Sep 17 00:00:00 2001 From: Shanu Date: Mon, 7 Sep 2026 13:22:04 +0530 Subject: [PATCH 1/3] Embed a source batch's documents together instead of one at a time `accept_source_items` wrote each item through `put_doc`, and the document upsert embedded that one document's chunks in its own provider request, so a 500-item connector pass paid 500 embedding round-trips: about 1.7 s each against the managed embedder, roughly 15 minutes, right at the host's `AcceptSourceItems` deadline. Add a batch write path to the store. `UnifiedMemory::upsert_documents` gates every document, chunks them all, embeds the chunk texts in requests of at most `EMBED_REQUEST_MAX_TEXTS` (64) texts across the whole batch, then writes each document under its per-key lock with the row, the chunk replacement and the new vectors in one transaction. The first write failure ends the batch as the last result entry, so the caller's per-item accounting is unchanged. `MemoryClient::put_docs` wraps it and queues graph extraction per written document, and `accept_source_items` converts its whole batch up front and makes one `put_docs` call, then feeds the memory tree once per written item as before. The single-document upsert is now the one-element case: it embeds before it writes and no longer holds the write lock across the provider round-trip. Closes #138 --- crates/tinymemory-core/src/store/client.rs | 29 ++ .../tinymemory-core/src/store/client_tests.rs | 41 +++ .../src/store/namespace_store/README.md | 2 +- .../src/store/namespace_store/documents.rs | 254 ++++++++++++------ .../store/namespace_store/documents_tests.rs | 203 ++++++++++++++ .../tinymemory-core/src/store/write_gate.rs | 55 +++- .../src/store/write_gate_tests.rs | 36 +++ .../tinymemory-tinycortex/src/engine/mod.rs | 42 ++- .../tests/full_provider_conformance.rs | 148 ++++++++++ 9 files changed, 718 insertions(+), 92 deletions(-) diff --git a/crates/tinymemory-core/src/store/client.rs b/crates/tinymemory-core/src/store/client.rs index e4808279..ed0be84b 100644 --- a/crates/tinymemory-core/src/store/client.rs +++ b/crates/tinymemory-core/src/store/client.rs @@ -186,6 +186,35 @@ impl MemoryClient { Ok(document_id) } + /// Store many documents at once — the batch form of [`Self::put_doc`]. + /// + /// The documents' chunks are embedded together, one provider request per + /// bounded group of chunk texts across the whole batch, instead of one + /// request per document (tinymemory#138). Each document is still gated, + /// written and queued for background graph extraction exactly as + /// `put_doc` does it. + /// + /// Documents are written in order and the first failure ends the batch: + /// the result holds one entry per document attempted, in input order, so a + /// failure is always the last entry and every document before it was + /// written and queued. + pub async fn put_docs( + &self, + inputs: Vec, + ) -> Vec> { + let results = self.inner.upsert_documents(inputs.clone()).await; + for (document, result) in inputs.into_iter().zip(&results) { + if let Ok(document_id) = result { + self.ingestion_queue.submit(IngestionJob { + document_id: document_id.clone(), + document, + config: MemoryIngestionConfig::default(), + }); + } + } + results + } + /// Store a document (DB row + markdown file) without vector embedding or /// graph extraction. Use this for high-frequency, ephemeral writes where /// the full pipeline would be too expensive (e.g. transient sync diff --git a/crates/tinymemory-core/src/store/client_tests.rs b/crates/tinymemory-core/src/store/client_tests.rs index 20f6aabf..afbe90cc 100644 --- a/crates/tinymemory-core/src/store/client_tests.rs +++ b/crates/tinymemory-core/src/store/client_tests.rs @@ -379,3 +379,44 @@ async fn ingest_doc_completes_and_stores_document() { // is exercised (no panic). let _ = result; } + +#[tokio::test] +async fn put_docs_writes_every_document_and_returns_ids_in_order() { + let (_tmp, client) = make_client(); + let results = client + .put_docs(vec![ + doc("batch", "k1", "one"), + doc("batch", "k2", "two"), + doc("batch", "k3", "three"), + ]) + .await; + + let ids: Vec = results + .into_iter() + .map(|result| result.expect("each document is written")) + .collect(); + assert_eq!(ids.len(), 3); + + let listed = client.list_documents(Some("batch")).await.unwrap(); + assert_eq!(listed["count"].as_u64(), Some(3)); + let by_key: std::collections::BTreeMap = listed["documents"] + .as_array() + .unwrap() + .iter() + .map(|document| { + ( + document["key"].as_str().unwrap().to_string(), + document["documentId"].as_str().unwrap().to_string(), + ) + }) + .collect(); + assert_eq!( + ids, + vec![ + by_key["k1"].clone(), + by_key["k2"].clone(), + by_key["k3"].clone() + ], + "ids come back in input order" + ); +} diff --git a/crates/tinymemory-core/src/store/namespace_store/README.md b/crates/tinymemory-core/src/store/namespace_store/README.md index 1718dcfb..d64f572f 100644 --- a/crates/tinymemory-core/src/store/namespace_store/README.md +++ b/crates/tinymemory-core/src/store/namespace_store/README.md @@ -11,7 +11,7 @@ tier; this directory is intentionally not migration staging. - **`mod.rs`** — declares the `UnifiedMemory` struct (connection + paths + embedder) and wires the submodules. - **`init.rs`** — constructor, `CREATE TABLE` bootstrap (docs, kv, graph, vector chunks, episodic FTS5, segments, events, profile), idempotent legacy-namespace migrations, plus path / namespace helpers (`sanitize_namespace`, `now_ts`, `namespace_dir`). -- **`documents.rs`** — `memory_docs` CRUD: `upsert_document` (chunks + embeds + writes markdown sidecar), `upsert_document_metadata_only` (light path), `list_documents`, `list_namespaces`, `delete_document`, `clear_namespace`. +- **`documents.rs`** — `memory_docs` CRUD: `upsert_document` (chunks + embeds + writes markdown sidecar), `upsert_documents` (the batch form: one embedding request per `EMBED_REQUEST_MAX_TEXTS` chunk texts across all the documents, then one transaction per document — tinymemory#138), `upsert_document_metadata_only` (light path), `list_documents`, `list_namespaces`, `delete_document`, `clear_namespace`. - **`kv.rs`** — global and namespace-scoped get/set/delete/list against `kv_global` / `kv_namespace`. - **`../../safety/`** — secret redaction/validation helpers. Document, KV, and episodic writes sanitize credentials before persistence and emit `[memory:safety]` diagnostics when a payload is rewritten. diff --git a/crates/tinymemory-core/src/store/namespace_store/documents.rs b/crates/tinymemory-core/src/store/namespace_store/documents.rs index 426c7e8c..80edaeda 100644 --- a/crates/tinymemory-core/src/store/namespace_store/documents.rs +++ b/crates/tinymemory-core/src/store/namespace_store/documents.rs @@ -1,8 +1,9 @@ //! Document CRUD against the `memory_docs` table. //! -//! Owns the upsert pipeline (with chunking + embedding), metadata-only writes -//! for high-frequency callers, list/delete/clear-namespace operations, and the -//! markdown sidecar files in `memory/namespaces//docs/`. +//! Owns the upsert pipeline (with chunking + embedding, batched across +//! documents), metadata-only writes for high-frequency callers, +//! list/delete/clear-namespace operations, and the markdown sidecar files in +//! `memory/namespaces//docs/`. use rusqlite::{params, OptionalExtension}; use serde_json::{json, Value}; @@ -14,11 +15,29 @@ use crate::store::types::{NamespaceDocumentInput, StoredMemoryDocument, GLOBAL_N use super::UnifiedMemory; +/// Token budget per vector chunk when a document is split for embedding. +const DOCUMENT_CHUNK_MAX_TOKENS: usize = 225; + +/// Upper bound on chunk texts sent to the embedding provider in one request +/// when a batch of documents is embedded together +/// ([`UnifiedMemory::upsert_documents_presanitized`]). +/// +/// Sized under every provider's per-request input cap (Cohere admits 96 texts, +/// the others more) and, at [`DOCUMENT_CHUNK_MAX_TOKENS`] per chunk, roughly +/// 14k tokens per request — under every provider's token budget — so a batch +/// of any length becomes a handful of bounded requests rather than one a +/// provider may refuse or time out. Still one to two orders of magnitude +/// fewer round-trips than the one-per-document write path it replaces. +pub(crate) const EMBED_REQUEST_MAX_TEXTS: usize = 64; + impl UnifiedMemory { /// Insert or update a document by `(namespace, key)`. Writes the markdown /// sidecar, replaces vector chunks, and embeds them with the configured /// provider. /// + /// The one-document case of [`Self::upsert_documents_presanitized`]; see it + /// for the write order and the failure contract. + /// /// **Takes already-sanitized input.** The host secret/PII write gate runs /// in [`crate::store::write_gate`], which owns this /// method's only call site; use `UnifiedMemory::upsert_document` instead @@ -27,6 +46,115 @@ impl UnifiedMemory { pub(crate) async fn upsert_document_presanitized( &self, input: NamespaceDocumentInput, + ) -> Result { + self.upsert_documents_presanitized(vec![input]) + .await + .pop() + .unwrap_or_else(|| Err("document upsert produced no result".to_string())) + } + + /// Insert or update many documents, embedding their chunks **together**: + /// one provider request per [`EMBED_REQUEST_MAX_TEXTS`] chunk texts across + /// the whole batch rather than one request per document (tinymemory#138). + /// A connector pass of several hundred small items used to pay one + /// embedding round-trip each; here it pays one per bounded group of texts. + /// + /// Order of operations: every document is chunked, every chunk text is + /// embedded (see [`Self::embed_chunk_texts`]), then the documents are + /// written one at a time in input order — each under its own per-key write + /// lock, with the row, the chunk replacement and the new vectors in ONE + /// transaction, so a reader never sees a row whose chunks are still being + /// replaced. + /// + /// The first document whose write fails ends the batch: the result holds + /// one entry per document attempted, in input order, with that failure as + /// its last entry, and the documents after it are left untouched. An + /// embedding failure is not a write failure — a request the provider + /// refuses leaves the chunks it covered vector-less (keyword-searchable and + /// re-embeddable), exactly as the single-document path always has, and the + /// batch carries on. + /// + /// **Takes already-sanitized input** — same contract as + /// [`Self::upsert_document_presanitized`]; go through + /// `UnifiedMemory::upsert_documents` instead. + pub(crate) async fn upsert_documents_presanitized( + &self, + inputs: Vec, + ) -> Vec> { + let chunked: Vec> = inputs + .iter() + .map(|input| Self::chunk_document_content(&input.content, DOCUMENT_CHUNK_MAX_TOKENS)) + .collect(); + let texts: Vec<&str> = chunked.iter().flatten().map(String::as_str).collect(); + let mut vectors = self.embed_chunk_texts(&texts).await.into_iter(); + + let mut results = Vec::with_capacity(inputs.len()); + for (input, chunks) in inputs.into_iter().zip(chunked) { + // `embed_chunk_texts` yields exactly one slot per text, in order, + // so the next `chunks.len()` slots are this document's. + let embeddings: Vec>> = vectors.by_ref().take(chunks.len()).collect(); + let result = self + .write_document_presanitized(input, chunks, embeddings) + .await; + let failed = result.is_err(); + results.push(result); + if failed { + break; + } + } + results + } + + /// Embed `texts` in requests of at most [`EMBED_REQUEST_MAX_TEXTS`], + /// returning one slot per input position. + /// + /// Failure handling keeps the per-chunk resilience the single-document + /// path always had: + /// * a request the provider refuses leaves every position it covered + /// `None` — logged, not propagated, because a vector-less chunk is + /// still keyword-searchable and re-embeddable while a failed write is + /// lost; + /// * a provider that returns fewer vectors than texts, or an empty vector + /// for a position (`NoopEmbedding`, NaN recovery), leaves those + /// positions `None` by position. + async fn embed_chunk_texts(&self, texts: &[&str]) -> Vec>> { + let mut out: Vec>> = Vec::with_capacity(texts.len()); + for request in texts.chunks(EMBED_REQUEST_MAX_TEXTS) { + log::debug!( + "[memory] batch-embedding {} chunk text(s) in one request", + request.len() + ); + match self.embedder.embed(request).await { + Ok(vectors) => { + let mut vectors = vectors + .into_iter() + .map(|vector| (!vector.is_empty()).then_some(vector)); + out.extend(request.iter().map(|_| vectors.next().flatten())); + } + Err(e) => { + log::warn!( + "[memory] batch embed failed for {} chunk text(s); storing them without vectors: {e}", + request.len() + ); + out.resize(out.len() + request.len(), None); + } + } + } + out + } + + /// Persist one document whose chunks and vectors were computed up front. + /// + /// Under the per-key write lock: resolve the document id and `created_at`, + /// write the markdown sidecar, then commit the `memory_docs` row, the chunk + /// replacement and the new `vector_chunks` rows in a single transaction. + /// `embeddings` is aligned to `chunks` by position; a missing or `None` + /// slot stores that chunk without a vector. + async fn write_document_presanitized( + &self, + input: NamespaceDocumentInput, + chunks: Vec, + mut embeddings: Vec>>, ) -> Result { let namespace = Self::sanitize_namespace(&input.namespace); // The logical (delimiter-preserving) namespace, PII-redacted the same @@ -43,16 +171,17 @@ 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 + // Serialise writers of one key for the WHOLE write. 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. + // does not ORDER them: the id / `created_at` lookups, the sidecar write + // (which awaits) and the transaction below must not interleave with + // another writer of the same key, or writer B's row can land between + // writer A's lookups and A's commit and be overwritten by content A + // resolved against stale state. 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. Embedding happens before this lock is taken, so no provider + // round-trip is ever awaited while holding it. let _write_guard = Self::document_write_lock(&self.db_path, &namespace, &key) .lock_owned() .await; @@ -113,6 +242,8 @@ impl UnifiedMemory { let tags_json = serde_json::to_string(&input.tags).map_err(|e| e.to_string())?; let metadata_json = input.metadata.to_string(); + // Computed once; only attached to chunks that actually got a vector. + let signature = self.embedder.signature(); { let conn = self.conn.lock(); let tx = conn @@ -161,77 +292,36 @@ impl UnifiedMemory { params![namespace, document_id], ) .map_err(|e| format!("clear vector chunks: {e}"))?; - tx.commit().map_err(|e| format!("commit tx: {e}"))?; - } - - let chunks = Self::chunk_document_content(&input.content, 225); - - // Embed every chunk in a SINGLE provider call rather than one - // round-trip per chunk. All providers implement the batch `embed` - // (`embed_one` is just a convenience wrapper around it), so a document - // that chunks into N pieces previously paid N sequential network - // round-trips on the write path; this collapses them to one. - // - // Result handling preserves the previous per-chunk resilience: - // * a failed batch (provider error) stores all chunks WITHOUT a - // vector — exactly what `embed_one(...).await.ok()` did per chunk; - // * a provider that returns fewer/empty vectors than chunks (e.g. - // `NoopEmbedding` returns an empty Vec, or a blank position from - // NaN recovery) leaves those chunks vector-less by position. - let mut embeddings: Vec>> = if chunks.is_empty() { - Vec::new() - } else { - let chunk_refs: Vec<&str> = chunks.iter().map(String::as_str).collect(); - log::debug!( - "[memory] batch-embedding {} chunk(s) for {namespace}/{document_id}", - chunk_refs.len() - ); - match self.embedder.embed(&chunk_refs).await { - Ok(vectors) => vectors - .into_iter() - .map(|v| (!v.is_empty()).then_some(v)) - .collect(), - Err(e) => { - log::warn!( - "[memory] batch embed failed for {} chunk(s) in {namespace}/{document_id}; storing without vectors: {e}", - chunks.len() - ); - Vec::new() - } + for (idx, chunk) in chunks.iter().enumerate() { + // Move the vector out by position so recall can exclude vectors + // produced by a different embedding model (cross-model cosine is + // meaningless) and guard against dimension mismatches. Missing + // positions (short/empty provider result) stay vector-less. + let embedded = embeddings.get_mut(idx).and_then(Option::take); + let dim = embedded.as_ref().map(|v| v.len() as i64); + let model_signature = embedded.as_ref().map(|_| signature.clone()); + let embedding = embedded.as_ref().map(|v| Self::vec_to_bytes(v)); + let chunk_id = format!("{document_id}:{idx}"); + tx.execute( + "INSERT OR REPLACE INTO vector_chunks + (namespace, document_id, chunk_id, text, embedding, metadata_json, created_at, updated_at, model_signature, dim) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", + params![ + namespace, + document_id, + chunk_id, + chunk, + embedding, + json!({"lancedb_table": format!("ns_{namespace}"), "chunk_index": idx}).to_string(), + now, + now, + model_signature, + dim + ], + ) + .map_err(|e| format!("insert vector chunk: {e}"))?; } - }; - - // Computed once; only attached to chunks that actually got a vector. - let signature = self.embedder.signature(); - for (idx, chunk) in chunks.iter().enumerate() { - // Move the vector out by position so recall can exclude vectors - // produced by a different embedding model (cross-model cosine is - // meaningless) and guard against dimension mismatches. Missing - // positions (short/empty provider result) stay vector-less. - let embedded = embeddings.get_mut(idx).and_then(Option::take); - let dim = embedded.as_ref().map(|v| v.len() as i64); - let model_signature = embedded.as_ref().map(|_| signature.clone()); - let embedding = embedded.as_ref().map(|v| Self::vec_to_bytes(v)); - let chunk_id = format!("{document_id}:{idx}"); - let conn = self.conn.lock(); - conn.execute( - "INSERT OR REPLACE INTO vector_chunks - (namespace, document_id, chunk_id, text, embedding, metadata_json, created_at, updated_at, model_signature, dim) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)", - params![ - namespace, - document_id, - chunk_id, - chunk, - embedding, - json!({"lancedb_table": format!("ns_{namespace}"), "chunk_index": idx}).to_string(), - now, - now, - model_signature, - dim - ], - ) - .map_err(|e| format!("insert vector chunk: {e}"))?; + tx.commit().map_err(|e| format!("commit tx: {e}"))?; } Ok(document_id) diff --git a/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs b/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs index 7716c69e..f965d3fc 100644 --- a/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs +++ b/crates/tinymemory-core/src/store/namespace_store/documents_tests.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use serde_json::json; use tempfile::TempDir; +use super::EMBED_REQUEST_MAX_TEXTS; use crate::store::{NamespaceDocumentInput, UnifiedMemory}; use tinymemory_api::host::NoopEmbedding; @@ -40,6 +41,18 @@ fn count_vector_chunks(memory: &UnifiedMemory, namespace: &str, document_id: &st .unwrap() } +/// Like [`count_vector_chunks`], counting only the chunks that carry a vector. +fn count_embedded_chunks(memory: &UnifiedMemory, namespace: &str, document_id: &str) -> i64 { + let conn = memory.conn.lock(); + conn.query_row( + "SELECT COUNT(*) FROM vector_chunks + WHERE namespace = ?1 AND document_id = ?2 AND embedding IS NOT NULL", + rusqlite::params![UnifiedMemory::sanitize_namespace(namespace), document_id], + |row| row.get(0), + ) + .unwrap() +} + #[tokio::test] async fn list_documents_without_namespace_returns_all_docs_across_namespaces() { let tmp = TempDir::new().unwrap(); @@ -421,6 +434,196 @@ async fn upsert_document_batch_embeds_all_chunks_in_one_call() { ); } +/// Embedder that records the size of every request and can refuse one of them +/// (by 0-based request index), so the batch path's request splitting and its +/// per-request failure isolation are observable. +struct RequestRecordingEmbedder { + requests: std::sync::Mutex>, + fail_request: Option, +} + +impl RequestRecordingEmbedder { + fn new(fail_request: Option) -> Arc { + Arc::new(Self { + requests: std::sync::Mutex::new(Vec::new()), + fail_request, + }) + } + + fn requests(&self) -> Vec { + self.requests.lock().unwrap().clone() + } +} + +#[async_trait::async_trait] +impl tinymemory_api::host::EmbeddingProvider for RequestRecordingEmbedder { + fn name(&self) -> &str { + "recording" + } + + fn model_id(&self) -> &str { + "recording-test" + } + + fn dimensions(&self) -> usize { + 3 + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + let index = { + let mut requests = self.requests.lock().unwrap(); + requests.push(texts.len()); + requests.len() - 1 + }; + if self.fail_request == Some(index) { + anyhow::bail!("provider refused request {index}"); + } + Ok(texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect()) + } +} + +/// tinymemory#138: a batch of documents must share embedding requests rather +/// than pay one round-trip per document. +#[tokio::test] +async fn upsert_documents_embeds_every_document_in_one_request() { + let tmp = TempDir::new().unwrap(); + let embedder = RequestRecordingEmbedder::new(None); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + // Each body chunks into several pieces, so the batch has many more chunks + // than documents — and still fits one request. + let long_body = "alpha ".repeat(400); + let inputs: Vec = ["doc-a", "doc-b", "doc-c"] + .iter() + .map(|key| make_doc_input("test:batch-many", key, key, &long_body)) + .collect(); + + let ids: Vec = memory + .upsert_documents(inputs) + .await + .into_iter() + .map(|result| result.unwrap()) + .collect(); + assert_eq!(ids.len(), 3); + + let mut total_chunks = 0; + for id in &ids { + let chunks = count_vector_chunks(&memory, "test:batch-many", id); + assert!( + chunks >= 3, + "each body should chunk into >=3 pieces, got {chunks}" + ); + assert_eq!( + count_embedded_chunks(&memory, "test:batch-many", id), + chunks, + "every chunk of every document must carry a vector" + ); + total_chunks += chunks; + } + assert_eq!( + embedder.requests(), + vec![total_chunks as usize], + "three documents' chunks must travel in ONE provider request, not one per document" + ); +} + +#[tokio::test] +async fn upsert_documents_splits_embedding_requests_at_the_request_cap() { + let tmp = TempDir::new().unwrap(); + let embedder = RequestRecordingEmbedder::new(None); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + // One chunk per document, one more document than a request may carry. + let inputs: Vec = (0..=EMBED_REQUEST_MAX_TEXTS) + .map(|n| make_doc_input("test:cap", &format!("doc-{n}"), "Doc", &format!("body {n}"))) + .collect(); + + let results = memory.upsert_documents(inputs).await; + assert_eq!(results.len(), EMBED_REQUEST_MAX_TEXTS + 1); + assert!(results.iter().all(Result::is_ok), "{results:?}"); + assert_eq!( + embedder.requests(), + vec![EMBED_REQUEST_MAX_TEXTS, 1], + "chunk texts must be sent in requests of at most EMBED_REQUEST_MAX_TEXTS" + ); +} + +#[tokio::test] +async fn upsert_documents_keeps_writing_when_one_embedding_request_is_refused() { + let tmp = TempDir::new().unwrap(); + // The second request (the lone overflow chunk) is refused. + let embedder = RequestRecordingEmbedder::new(Some(1)); + let memory = UnifiedMemory::new(tmp.path(), embedder.clone(), None).unwrap(); + + let inputs: Vec = (0..=EMBED_REQUEST_MAX_TEXTS) + .map(|n| { + make_doc_input( + "test:refused", + &format!("doc-{n}"), + "Doc", + &format!("body {n}"), + ) + }) + .collect(); + + let ids: Vec = memory + .upsert_documents(inputs) + .await + .into_iter() + .map(|result| result.expect("an embedding failure is not a write failure")) + .collect(); + assert_eq!(ids.len(), EMBED_REQUEST_MAX_TEXTS + 1); + assert_eq!(embedder.requests(), vec![EMBED_REQUEST_MAX_TEXTS, 1]); + + let first = &ids[0]; + assert_eq!(count_vector_chunks(&memory, "test:refused", first), 1); + assert_eq!( + count_embedded_chunks(&memory, "test:refused", first), + 1, + "documents covered by the accepted request keep their vectors" + ); + let last = ids.last().unwrap(); + assert_eq!( + count_vector_chunks(&memory, "test:refused", last), + 1, + "the document covered by the refused request is still written" + ); + assert_eq!( + count_embedded_chunks(&memory, "test:refused", last), + 0, + "…but its chunk is stored without a vector, like the single-document path" + ); +} + +#[tokio::test] +async fn upsert_documents_stops_at_the_first_write_failure() { + let tmp = TempDir::new().unwrap(); + let memory = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap(); + + let results = memory + .upsert_documents(vec![ + make_doc_input("test:stop", "doc-a", "Doc A", "A body"), + make_doc_input("test:stop", " ", "Blank key", "rejected by the store"), + make_doc_input("test:stop", "doc-c", "Doc C", "never attempted"), + ]) + .await; + + assert_eq!( + results.len(), + 2, + "the failing document is the last entry; nothing after it is attempted" + ); + assert!(results[0].is_ok()); + let err = results[1].as_ref().unwrap_err(); + assert!( + err.contains("document key cannot be empty"), + "the failing entry carries the store's own error, got {err:?}" + ); + let docs = memory.list_documents(Some("test:stop")).await.unwrap(); + assert_eq!(docs["count"].as_u64(), Some(1)); + assert_eq!(docs["documents"][0]["key"], "doc-a"); +} + #[tokio::test] async fn upsert_document_reuses_document_id_preserves_created_at_and_replaces_vector_chunks() { let tmp = TempDir::new().unwrap(); diff --git a/crates/tinymemory-core/src/store/write_gate.rs b/crates/tinymemory-core/src/store/write_gate.rs index 1f204c36..1f04ed25 100644 --- a/crates/tinymemory-core/src/store/write_gate.rs +++ b/crates/tinymemory-core/src/store/write_gate.rs @@ -18,7 +18,8 @@ //! `upsert_document_metadata_only` as inherent methods on `UnifiedMemory` with //! the **same names and signatures they always had**, so every existing caller //! is routed through the gate without a single call-site edit — which is also -//! what makes the "no bypass" claim below checkable rather than hopeful. +//! what makes the "no bypass" claim below checkable rather than hopeful. The +//! batch form, `upsert_documents`, is declared here for the same reason. //! //! # The gate, in order //! @@ -45,12 +46,14 @@ //! //! # No bypass //! -//! `upsert_document_presanitized` / `upsert_document_metadata_only_presanitized` -//! are `pub(crate)` and newly named, and this module holds their only call -//! sites — verify with: +//! `upsert_document_presanitized` / `upsert_documents_presanitized` / +//! `upsert_document_metadata_only_presanitized` are `pub(crate)` and newly +//! named, and this module holds their only call sites outside `documents.rs` +//! itself (where the one-document method is the one-element case of the batch +//! one) — verify with: //! //! ```text -//! rg 'upsert_document(_metadata_only)?_presanitized' src/ +//! rg 'upsert_documents?(_metadata_only)?_presanitized' src/ //! ``` //! //! Every other writer in the tree (`Memory::store_with_taint`, `MemoryClient`, @@ -142,6 +145,48 @@ impl UnifiedMemory { } } + /// Insert or update many documents, applying the host secret/PII write + /// gate to each and embedding their chunks together — one provider request + /// per bounded group of chunk texts across the batch rather than one per + /// document (tinymemory#138). + /// + /// Inputs are gated in order and the admitted prefix is written by + /// `Self::upsert_documents_presanitized`, which stops at the first write + /// failure; the first gate rejection, if any, ends the batch the same way. + /// The result holds one entry per document attempted, in input order, so a + /// failure is always the last entry and every document before it was + /// written. + /// + /// # Errors + /// + /// Per document, the same failure modes as [`Self::upsert_document`]. + pub async fn upsert_documents( + &self, + inputs: Vec, + ) -> Vec> { + let mut admitted = Vec::with_capacity(inputs.len()); + let mut rejection = None; + for input in inputs { + match gate(input, "document") { + GateOutcome::Admit(input) => admitted.push(*input), + GateOutcome::Reject(err) => { + rejection = Some(err); + break; + } + } + } + let mut results = self.upsert_documents_presanitized(admitted).await; + if let Some(err) = rejection { + // Only when the admitted prefix was written in full: a write failure + // in it is already the batch's last entry, and everything after a + // failure — the rejected document included — stays unattempted. + if results.iter().all(Result::is_ok) { + results.push(Err(err)); + } + } + results + } + /// Store a document without chunking, embedding, or graph extraction, /// applying the host secret/PII write gate first. /// diff --git a/crates/tinymemory-core/src/store/write_gate_tests.rs b/crates/tinymemory-core/src/store/write_gate_tests.rs index 9048d6f1..a1789695 100644 --- a/crates/tinymemory-core/src/store/write_gate_tests.rs +++ b/crates/tinymemory-core/src/store/write_gate_tests.rs @@ -174,3 +174,39 @@ async fn gate_canonicalizes_a_pii_like_key_and_keeps_the_row_addressable() { docs[0].key ); } + +#[tokio::test] +async fn gated_batch_upsert_redacts_each_document_and_stops_at_the_first_rejection() { + let (_tmp, memory) = fresh(); + + let mut secret_key = secret_doc("sk-1234567890123456789012345"); + secret_key.namespace = "safe".to_string(); + let results = memory + .upsert_documents(vec![ + secret_doc("first"), + secret_key, + secret_doc("never-attempted"), + ]) + .await; + + assert_eq!( + results.len(), + 2, + "the rejected document ends the batch as its last entry, got {results:?}" + ); + assert!(results[0].is_ok(), "{results:?}"); + let err = results[1].as_ref().unwrap_err(); + assert!( + err.contains("cannot contain secrets"), + "secret-like key must be refused, got {err:?}" + ); + + let docs = memory.load_documents_for_scope("safe").await.unwrap(); + assert_eq!(docs.len(), 1, "only the admitted prefix reaches storage"); + assert_eq!(docs[0].key, "first"); + assert!( + !docs[0].content.contains("BEGIN PRIVATE KEY"), + "a batched write must be redacted exactly like a single one, got {:?}", + docs[0].content + ); +} diff --git a/crates/tinymemory-tinycortex/src/engine/mod.rs b/crates/tinymemory-tinycortex/src/engine/mod.rs index de4a0b3d..7baacb11 100644 --- a/crates/tinymemory-tinycortex/src/engine/mod.rs +++ b/crates/tinymemory-tinycortex/src/engine/mod.rs @@ -2198,11 +2198,25 @@ impl MemorySourceSink for TinycortexProvider { // Resolved once: whether a source's items belong in the memory tree is a // property of the source, not of the item. let tree_scope = connector_tree_scope(source_kind, source_id); + + // Convert every item up front so the store can embed the whole batch + // together (tinymemory#138): `put_docs` pays one embedding round-trip + // per bounded group of chunk texts across ALL the items instead of one + // per item, which is what put a 500-item connector pass at ~15 minutes + // and over the host's slow-call deadline. An item that cannot be + // converted (a blank id, a shape the store rejects) ends the conversion + // where it stands: the items before it are still written and treed + // below, then its error is returned — the same end state as when each + // item was written as it was reached. + let mut inputs = Vec::with_capacity(items_len); + let mut tree_items = Vec::with_capacity(items_len); + let mut rejected = None; for item in items { if item.item_id.trim().is_empty() { - return Err(MemoryError::Invalid( + rejected = Some(MemoryError::Invalid( "source item_id must not be empty".to_string(), )); + break; } let title = if item.title.trim().is_empty() { item.item_id.clone() @@ -2235,8 +2249,25 @@ impl MemorySourceSink for TinycortexProvider { document_id: None, taint, }; - let input = Self::cross(&input, "convert source document")?; - match self.client.put_doc(input).await { + match Self::cross(&input, "convert source document") { + Ok(input) => { + inputs.push(input); + tree_items.push(tree_item); + } + Err(error) => { + rejected = Some(error); + break; + } + } + } + + // One store call for the whole batch. Its result holds one entry per + // document attempted, in order, with a failed write always last, so + // walking it in order keeps exactly the per-item accounting the old + // one-write-per-item loop had. + let results = self.client.put_docs(inputs).await; + for (result, tree_item) in results.into_iter().zip(tree_items) { + match result { Ok(id) => { outcome.written = outcome.written.saturating_add(1); outcome.ids.push(id); @@ -2296,7 +2327,10 @@ impl MemorySourceSink for TinycortexProvider { } } } - Ok(outcome) + match rejected { + Some(error) => Err(error), + None => Ok(outcome), + } } async fn forget_source(&self, source_id: &str) -> Result { diff --git a/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs b/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs index 3b8e8258..5db9eab0 100644 --- a/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs +++ b/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs @@ -4016,3 +4016,151 @@ async fn a_connector_sync_reaches_the_memory_tree_and_is_forgotten_with_its_sour (openhuman#6007)" ); } + +/// Embedder that records the size of every request, so a test can prove how +/// many provider round-trips a batch of source items paid for. +struct RequestCountingEmbedder { + requests: std::sync::Mutex>, +} + +#[async_trait::async_trait] +impl tinymemory_api::host::EmbeddingProvider for RequestCountingEmbedder { + fn name(&self) -> &str { + "counting" + } + + fn model_id(&self) -> &str { + "counting-test" + } + + fn dimensions(&self) -> usize { + 3 + } + + async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { + self.requests + .lock() + .expect("requests lock") + .push(texts.len()); + Ok(texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect()) + } +} + +/// A provider whose document store embeds through `embedder`. The store takes +/// its embedder directly, so this bypasses the process-global seam and leaves +/// what every other test in this binary sees untouched. +fn provider_with_embedder( + workspace: &std::path::Path, + embedder: Arc, +) -> TinycortexProvider { + tinymemory_core::embedding_host::set_embedding_host(Arc::new(NoopEmbeddingHost)); + let memory = tinymemory_core::store::UnifiedMemory::new(workspace, embedder, None) + .expect("open the workspace store"); + let client = Arc::new(tinymemory_core::store::MemoryClient::from_unified_memory( + memory, + )); + TinycortexProvider::new( + "tinycortex".into(), + provider_config(workspace, serde_json::Value::Null), + client, + ) +} + +/// tinymemory#138: a connector pass hands the sink hundreds of small items, and +/// each one used to pay its own embedding round-trip — about 1.7 s per item +/// against the managed embedder, so a 500-item pass ran into the host's +/// 15-minute slow-call deadline. The batch must share requests across items. +#[tokio::test(flavor = "multi_thread")] +async fn source_items_are_embedded_together_rather_than_one_request_per_item() { + use tinymemory_api::provider::types::SourceItem; + use tinymemory_api::provider::MemoryProvider; + use tinymemory_api::types::MemoryTaint; + + let workspace = tempfile::tempdir().expect("workspace"); + let embedder = Arc::new(RequestCountingEmbedder { + requests: std::sync::Mutex::new(Vec::new()), + }); + let provider = provider_with_embedder(workspace.path(), Arc::clone(&embedder)); + let config = provider_config(workspace.path(), serde_json::Value::Null); + let source = provider.as_sources().expect("SourceSink"); + + let items: Vec = (1..=5) + .map(|n| SourceItem { + item_id: format!("msg-{n}"), + title: format!("Message {n}"), + content: format!("Short message number {n} about the roadmap."), + mime: Some("text/plain".into()), + url: None, + updated_at_ms: Some(n), + tags: vec!["gmail".into()], + }) + .collect(); + let outcome = source + .accept_source_items("gmail:conn-1", "composio", items, MemoryTaint::ExternalSync) + .await + .expect("accept the batch"); + assert_eq!(outcome.written, 5); + assert_eq!(outcome.ids.len(), 5, "one id per written item, in order"); + + let requests = embedder.requests.lock().expect("requests lock").clone(); + assert_eq!( + requests, + vec![5], + "five one-chunk items must cost ONE embedding request, not five (tinymemory#138); \ + got {requests:?}" + ); + // The tree funnel still runs once per written item (openhuman#6007). + assert_eq!( + tinymemory_core::store::chunks::count_chunks(&config).expect("count chunks"), + 5, + "every written item must still reach the memory tree" + ); +} + +/// The batch keeps the sink's per-item accounting: an item the sink rejects +/// fails the call, after the items before it were written and none after it. +#[tokio::test(flavor = "multi_thread")] +async fn a_blank_source_item_id_fails_the_batch_after_the_items_before_it() { + use tinymemory_api::error::MemoryError; + use tinymemory_api::provider::types::SourceItem; + use tinymemory_api::provider::MemoryProvider; + use tinymemory_api::types::MemoryTaint; + + let workspace = tempfile::tempdir().expect("workspace"); + let provider = provider_over(workspace.path()); + let source = provider.as_sources().expect("SourceSink"); + let item = |id: &str| SourceItem { + item_id: id.into(), + title: "Item".into(), + content: "body".into(), + mime: None, + url: None, + updated_at_ms: None, + tags: Vec::new(), + }; + + let error = source + .accept_source_items( + "drive-1", + "drive", + vec![item("a"), item("b"), item(" "), item("d")], + MemoryTaint::ExternalSync, + ) + .await + .expect_err("a blank item id must fail the call"); + assert!( + matches!(&error, MemoryError::Invalid(message) if message.contains("item_id must not be empty")), + "got {error:?}" + ); + + let documents = provider.as_documents().expect("Documents"); + let listed = documents + .list_documents(Some("source:drive-1")) + .await + .expect("list documents"); + assert_eq!( + listed["count"].as_u64(), + Some(2), + "the items before the blank one are written; the ones after it are not" + ); +} From 958e45e38fdb142f770a44500a806c679815ad0a Mon Sep 17 00:00:00 2001 From: Shanu Date: Mon, 7 Sep 2026 13:36:05 +0530 Subject: [PATCH 2/3] Bump vendor/tinycortex to 79131f2 (reembed_backfill claim priority) Picks up tinyhumansai/tinycortex#169 (`claim_next` ranks `reembed_backfill` right after `seal`, so the only path that writes chunk vectors is no longer round-robined behind the whole `extract_chunk` backlog by the LLM-gate defer; tinycortex#168), the v0.1.2 release (crate version 0.1.1 -> 0.1.2, `dirs` 5 -> 6) and the dependabot bumps between the two pins. Both lockfiles follow the pin; the `dirs` bump lets cargo drop the second copies of `dirs` / `dirs-sys` / `redox_users` / `thiserror 1.x` and the `windows-sys 0.48` family they alone pulled in. --- Cargo.lock | 166 ++++----------------------- crates/tinymemory-module/Cargo.lock | 168 +++++----------------------- vendor/tinycortex | 2 +- 3 files changed, 50 insertions(+), 286 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e8ccb45a..e98615a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -321,34 +321,13 @@ dependencies = [ "crypto-common 0.2.2", ] -[[package]] -name = "dirs" -version = "5.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" -dependencies = [ - "dirs-sys 0.4.1", -] - [[package]] name = "dirs" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" dependencies = [ - "dirs-sys 0.5.0", -] - -[[package]] -name = "dirs-sys" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" -dependencies = [ - "libc", - "option-ext", - "redox_users 0.4.6", - "windows-sys 0.48.0", + "dirs-sys", ] [[package]] @@ -359,7 +338,7 @@ checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" dependencies = [ "libc", "option-ext", - "redox_users 0.5.2", + "redox_users", "windows-sys 0.61.2", ] @@ -1168,7 +1147,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror 2.0.20", + "thiserror", "tokio", "tracing", "web-time", @@ -1190,7 +1169,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.20", + "thiserror", "tinyvec", "tracing", "web-time", @@ -1295,17 +1274,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "redox_users" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" -dependencies = [ - "getrandom 0.2.17", - "libredox", - "thiserror 1.0.69", -] - [[package]] name = "redox_users" version = "0.5.2" @@ -1314,7 +1282,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1429,7 +1397,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" dependencies = [ "hashbrown 0.16.1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1772,39 +1740,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", "windows-sys 0.61.2", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.20", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] @@ -1820,13 +1768,13 @@ dependencies = [ [[package]] name = "tinycortex" -version = "0.1.1" +version = "0.1.2" dependencies = [ "anyhow", "async-trait", "block2", "chrono", - "dirs 5.0.1", + "dirs", "futures", "git2", "hex", @@ -1843,7 +1791,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror 2.0.20", + "thiserror", "tinycortex-api", "tinyinference", "tokio", @@ -1872,7 +1820,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "tokio", ] @@ -1920,7 +1868,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "uuid", ] @@ -1943,7 +1891,7 @@ dependencies = [ "async-trait", "axum", "chrono", - "dirs 6.0.0", + "dirs", "futures", "log", "parking_lot", @@ -1955,7 +1903,7 @@ dependencies = [ "serde_json", "sha2 0.11.0", "tempfile", - "thiserror 2.0.20", + "thiserror", "tinycortex", "tinycortex-api", "tinyinference", @@ -2317,7 +2265,7 @@ dependencies = [ "log", "rand 0.9.5", "sha1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -2606,22 +2554,13 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets 0.48.5", -] - [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -2633,67 +2572,34 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-targets" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" -dependencies = [ - "windows_aarch64_gnullvm 0.48.5", - "windows_aarch64_msvc 0.48.5", - "windows_i686_gnu 0.48.5", - "windows_i686_msvc 0.48.5", - "windows_x86_64_gnu 0.48.5", - "windows_x86_64_gnullvm 0.48.5", - "windows_x86_64_msvc 0.48.5", -] - [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", "windows_i686_gnullvm", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -2706,48 +2612,24 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_x86_64_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" diff --git a/crates/tinymemory-module/Cargo.lock b/crates/tinymemory-module/Cargo.lock index c4f8de80..bc426e5a 100644 --- a/crates/tinymemory-module/Cargo.lock +++ b/crates/tinymemory-module/Cargo.lock @@ -266,34 +266,13 @@ dependencies = [ "crypto-common 0.2.2", ] -[[package]] -name = "dirs" -version = "5.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44c45a9d03d6676652bcb5e724c7e988de1acad23a711b5217ab9cbecbec2225" -dependencies = [ - "dirs-sys 0.4.1", -] - [[package]] name = "dirs" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" dependencies = [ - "dirs-sys 0.5.0", -] - -[[package]] -name = "dirs-sys" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "520f05a5cbd335fae5a99ff7a6ab8627577660ee5cfd6a94a6a929b52ff0321c" -dependencies = [ - "libc", - "option-ext", - "redox_users 0.4.6", - "windows-sys 0.48.0", + "dirs-sys", ] [[package]] @@ -304,7 +283,7 @@ checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" dependencies = [ "libc", "option-ext", - "redox_users 0.5.2", + "redox_users", "windows-sys 0.61.2", ] @@ -1117,7 +1096,7 @@ dependencies = [ "rustc-hash", "rustls", "socket2", - "thiserror 2.0.20", + "thiserror", "tokio", "tracing", "web-time", @@ -1139,7 +1118,7 @@ dependencies = [ "rustls", "rustls-pki-types", "slab", - "thiserror 2.0.20", + "thiserror", "tinyvec", "tracing", "web-time", @@ -1209,17 +1188,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "redox_users" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba009ff324d1fc1b900bd1fdb31564febe58a8ccc8a6fdbb93b543d33b13ca43" -dependencies = [ - "getrandom 0.2.17", - "libredox", - "thiserror 1.0.69", -] - [[package]] name = "redox_users" version = "0.5.2" @@ -1228,7 +1196,7 @@ checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" dependencies = [ "getrandom 0.2.17", "libredox", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1342,7 +1310,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c51c9ae4df8a7fba42103df5c621fa3c37eccf3a3c650879e90fc48b11cc192c" dependencies = [ "hashbrown 0.16.1", - "thiserror 2.0.20", + "thiserror", ] [[package]] @@ -1707,33 +1675,13 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "thiserror" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" -dependencies = [ - "thiserror-impl 1.0.69", -] - [[package]] name = "thiserror" version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ - "thiserror-impl 2.0.20", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.69" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.119", + "thiserror-impl", ] [[package]] @@ -1757,7 +1705,7 @@ dependencies = [ "serde_json", "tar", "tempfile", - "thiserror 2.0.20", + "thiserror", "tinybus-macros", "tokio", "toml 0.8.23", @@ -1789,13 +1737,13 @@ dependencies = [ [[package]] name = "tinycortex" -version = "0.1.1" +version = "0.1.2" dependencies = [ "anyhow", "async-trait", "block2", "chrono", - "dirs 5.0.1", + "dirs", "futures", "git2", "hex", @@ -1812,7 +1760,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.10.9", - "thiserror 2.0.20", + "thiserror", "tinycortex-api", "tinyinference", "tokio", @@ -1841,7 +1789,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "tokio", ] @@ -1879,7 +1827,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "uuid", ] @@ -1890,7 +1838,7 @@ dependencies = [ "anyhow", "async-trait", "chrono", - "dirs 6.0.0", + "dirs", "futures", "log", "parking_lot", @@ -1901,7 +1849,7 @@ dependencies = [ "serde", "serde_json", "sha2 0.11.0", - "thiserror 2.0.20", + "thiserror", "tinycortex", "tinycortex-api", "tinyinference", @@ -2450,7 +2398,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -2512,22 +2460,13 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-sys" -version = "0.48.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" -dependencies = [ - "windows-targets 0.48.5", -] - [[package]] name = "windows-sys" version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -2539,67 +2478,34 @@ dependencies = [ "windows-link", ] -[[package]] -name = "windows-targets" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" -dependencies = [ - "windows_aarch64_gnullvm 0.48.5", - "windows_aarch64_msvc 0.48.5", - "windows_i686_gnu 0.48.5", - "windows_i686_msvc 0.48.5", - "windows_x86_64_gnu 0.48.5", - "windows_x86_64_gnullvm 0.48.5", - "windows_x86_64_msvc 0.48.5", -] - [[package]] name = "windows-targets" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", "windows_i686_gnullvm", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" - [[package]] name = "windows_aarch64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_i686_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" - [[package]] name = "windows_i686_gnu" version = "0.52.6" @@ -2612,48 +2518,24 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_x86_64_gnu" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_msvc" -version = "0.48.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" @@ -2787,7 +2669,7 @@ dependencies = [ "flate2", "indexmap", "memchr", - "thiserror 2.0.20", + "thiserror", "zopfli", ] diff --git a/vendor/tinycortex b/vendor/tinycortex index cb1c1639..79131f27 160000 --- a/vendor/tinycortex +++ b/vendor/tinycortex @@ -1 +1 @@ -Subproject commit cb1c1639c8a82221e1d8e8a9715b7640bb3f05d8 +Subproject commit 79131f275cf98f06421fefd9dc812c5a63a5a182 From 80a0e22fdce4a1bc3502378b12192075f277e3df Mon Sep 17 00:00:00 2001 From: Shanu Date: Mon, 7 Sep 2026 13:36:05 +0530 Subject: [PATCH 3/3] Count the graph-extraction jobs a batch write could not queue `put_docs` submits one background graph-extraction job per written document in a single burst, where `put_doc` spaced them out by one embedding round-trip each, so an existing backlog can make the bounded ingestion queue refuse some of them. The refusal is the queue's documented best-effort drop, but it was invisible to the caller. `put_docs` now returns a `BatchPutOutcome` whose `dropped_extractions` counts the refused jobs and logs the count, and `accept_source_items` logs it per source. The documents and their memory-tree chunks are stored either way; only the namespace graph extraction for those documents is skipped. Also count the memory-tree chunks in the new conformance test through the public chunks capability rather than the store internals. --- crates/tinymemory-core/src/store/client.rs | 50 +++++++++++++--- .../tinymemory-core/src/store/client_tests.rs | 57 ++++++++++++++++++- crates/tinymemory-core/src/store/mod.rs | 2 +- .../tinymemory-tinycortex/src/engine/mod.rs | 21 +++++-- .../tests/full_provider_conformance.rs | 9 ++- 5 files changed, 120 insertions(+), 19 deletions(-) diff --git a/crates/tinymemory-core/src/store/client.rs b/crates/tinymemory-core/src/store/client.rs index ed0be84b..d82f720f 100644 --- a/crates/tinymemory-core/src/store/client.rs +++ b/crates/tinymemory-core/src/store/client.rs @@ -27,6 +27,23 @@ use tinymemory_api::host::EmbeddingProvider; /// Reference-counted handle to a `MemoryClient`. pub type MemoryClientRef = Arc; +/// Outcome of [`MemoryClient::put_docs`]. +#[derive(Debug, Default)] +pub struct BatchPutOutcome { + /// One entry per document attempted, in input order. The first failure + /// ends the batch and is always the last entry, so every document before + /// it was written and queued. + pub results: Vec>, + /// Written documents whose background graph-extraction job the ingestion + /// queue refused (it was full, or its worker had shut down). Those + /// documents are stored and searchable; only their entity/relation + /// extraction was skipped — the same best-effort drop a refused + /// [`MemoryClient::put_doc`] submission makes, counted here because a + /// batch submits hundreds of jobs in one burst where the per-document + /// write path spaced them out. + pub dropped_extractions: usize, +} + /// Thread-safe container for an optional `MemoryClientRef`. /// /// Used for global state management where the memory client may or may not @@ -195,24 +212,39 @@ impl MemoryClient { /// `put_doc` does it. /// /// Documents are written in order and the first failure ends the batch: - /// the result holds one entry per document attempted, in input order, so a - /// failure is always the last entry and every document before it was - /// written and queued. - pub async fn put_docs( - &self, - inputs: Vec, - ) -> Vec> { + /// [`BatchPutOutcome::results`] holds one entry per document attempted, in + /// input order, so a failure is always the last entry and every document + /// before it was written and queued. A graph-extraction job the ingestion + /// queue refuses is not a document failure — the document is stored — but + /// it is counted in [`BatchPutOutcome::dropped_extractions`] and logged, + /// so a burst that overruns the queue is visible to the caller. + pub async fn put_docs(&self, inputs: Vec) -> BatchPutOutcome { let results = self.inner.upsert_documents(inputs.clone()).await; + let mut dropped_extractions = 0; for (document, result) in inputs.into_iter().zip(&results) { if let Ok(document_id) = result { - self.ingestion_queue.submit(IngestionJob { + let queued = self.ingestion_queue.submit(IngestionJob { document_id: document_id.clone(), document, config: MemoryIngestionConfig::default(), }); + if !queued { + dropped_extractions += 1; + } } } - results + if dropped_extractions > 0 { + log::warn!( + "[memory] graph extraction skipped for {dropped_extractions} of {} written \ + document(s): the ingestion queue refused the job(s); the documents \ + themselves are stored", + results.iter().filter(|result| result.is_ok()).count() + ); + } + BatchPutOutcome { + results, + dropped_extractions, + } } /// Store a document (DB row + markdown file) without vector embedding or diff --git a/crates/tinymemory-core/src/store/client_tests.rs b/crates/tinymemory-core/src/store/client_tests.rs index afbe90cc..e41a994f 100644 --- a/crates/tinymemory-core/src/store/client_tests.rs +++ b/crates/tinymemory-core/src/store/client_tests.rs @@ -383,15 +383,20 @@ async fn ingest_doc_completes_and_stores_document() { #[tokio::test] async fn put_docs_writes_every_document_and_returns_ids_in_order() { let (_tmp, client) = make_client(); - let results = client + let outcome = client .put_docs(vec![ doc("batch", "k1", "one"), doc("batch", "k2", "two"), doc("batch", "k3", "three"), ]) .await; + assert_eq!( + outcome.dropped_extractions, 0, + "an idle default-capacity queue accepts every graph-extraction job" + ); - let ids: Vec = results + let ids: Vec = outcome + .results .into_iter() .map(|result| result.expect("each document is written")) .collect(); @@ -420,3 +425,51 @@ async fn put_docs_writes_every_document_and_returns_ids_in_order() { "ids come back in input order" ); } + +/// A batch submits its graph-extraction jobs in one burst, so a queue that a +/// per-document trickle never overran can refuse some of them. The refusal is +/// the queue's documented best-effort drop, but it must be counted rather +/// than lost: the caller sees how many documents skipped extraction. +#[tokio::test] +async fn put_docs_counts_the_graph_jobs_a_full_ingestion_queue_refuses() { + use tinymemory_api::host::NoopEmbedding; + crate::test_seams::init(); + let tmp = TempDir::new().unwrap(); + let inner = Arc::new(UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap()); + let state = IngestionState::new(); + // A one-slot queue whose worker cannot drain: the test holds the singleton + // run lock the worker takes before it processes a job, so after the worker + // pulls the first job the slot refills once and every later job is refused. + let ingestion_queue = + ingestion_queue::start_worker_with_capacity(Arc::clone(&inner), state.clone(), 1); + let _worker_blocked = state.acquire().await; + let client = MemoryClient { + inner, + ingestion_queue, + }; + + let outcome = client + .put_docs(vec![ + doc("burst", "k1", "one"), + doc("burst", "k2", "two"), + doc("burst", "k3", "three"), + ]) + .await; + + assert!( + outcome.results.iter().all(Result::is_ok), + "a refused extraction job is not a document failure, got {:?}", + outcome.results + ); + assert!( + (1..=2).contains(&outcome.dropped_extractions), + "one job fits the single slot (two if the worker pulled the first before the \ + next submit); the rest are refused and counted, got {}", + outcome.dropped_extractions + ); + assert_eq!( + client.list_documents(Some("burst")).await.unwrap()["count"].as_u64(), + Some(3), + "refused extraction jobs do not affect the documents themselves" + ); +} diff --git a/crates/tinymemory-core/src/store/mod.rs b/crates/tinymemory-core/src/store/mod.rs index d4aa5b9a..656b6279 100644 --- a/crates/tinymemory-core/src/store/mod.rs +++ b/crates/tinymemory-core/src/store/mod.rs @@ -55,7 +55,7 @@ mod write_gate; pub use kinds::MemoryKind; pub use traits::{ObsidianFile, ObsidianRepresentable, VectorEmbeddable}; -pub use client::{MemoryClient, MemoryClientRef, MemoryState}; +pub use client::{BatchPutOutcome, MemoryClient, MemoryClientRef, MemoryState}; pub use factories::{ active_embedding_signature, create_memory, create_memory_for_migration, create_memory_with_local_ai, effective_embedding_settings, effective_memory_backend_name, diff --git a/crates/tinymemory-tinycortex/src/engine/mod.rs b/crates/tinymemory-tinycortex/src/engine/mod.rs index 7baacb11..c32f80e8 100644 --- a/crates/tinymemory-tinycortex/src/engine/mod.rs +++ b/crates/tinymemory-tinycortex/src/engine/mod.rs @@ -2261,12 +2261,25 @@ impl MemorySourceSink for TinycortexProvider { } } - // One store call for the whole batch. Its result holds one entry per + // One store call for the whole batch. Its results hold one entry per // document attempted, in order, with a failed write always last, so - // walking it in order keeps exactly the per-item accounting the old + // walking them in order keeps exactly the per-item accounting the old // one-write-per-item loop had. - let results = self.client.put_docs(inputs).await; - for (result, tree_item) in results.into_iter().zip(tree_items) { + let batch = self.client.put_docs(inputs).await; + if batch.dropped_extractions > 0 { + // Best-effort by the queue's contract: the documents and (below) + // their memory-tree chunks are stored; only the namespace graph + // extraction for these items was skipped. Named per source so an + // operator can tell which sync overran the queue. + log::warn!( + "[tinycortex:sources] graph extraction skipped for {} of {} written \ + item(s) of source `{source_id}`: the ingestion queue refused the \ + job(s); the documents and their memory-tree chunks are stored", + batch.dropped_extractions, + batch.results.iter().filter(|result| result.is_ok()).count() + ); + } + for (result, tree_item) in batch.results.into_iter().zip(tree_items) { match result { Ok(id) => { outcome.written = outcome.written.saturating_add(1); diff --git a/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs b/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs index 5db9eab0..1e6b239c 100644 --- a/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs +++ b/crates/tinymemory-tinycortex/tests/full_provider_conformance.rs @@ -4073,7 +4073,7 @@ fn provider_with_embedder( #[tokio::test(flavor = "multi_thread")] async fn source_items_are_embedded_together_rather_than_one_request_per_item() { use tinymemory_api::provider::types::SourceItem; - use tinymemory_api::provider::MemoryProvider; + use tinymemory_api::provider::{ChunkQuery, MemoryProvider}; use tinymemory_api::types::MemoryTaint; let workspace = tempfile::tempdir().expect("workspace"); @@ -4081,7 +4081,6 @@ async fn source_items_are_embedded_together_rather_than_one_request_per_item() { requests: std::sync::Mutex::new(Vec::new()), }); let provider = provider_with_embedder(workspace.path(), Arc::clone(&embedder)); - let config = provider_config(workspace.path(), serde_json::Value::Null); let source = provider.as_sources().expect("SourceSink"); let items: Vec = (1..=5) @@ -4110,8 +4109,12 @@ async fn source_items_are_embedded_together_rather_than_one_request_per_item() { got {requests:?}" ); // The tree funnel still runs once per written item (openhuman#6007). + let chunks = provider.as_chunks().expect("Chunks"); assert_eq!( - tinymemory_core::store::chunks::count_chunks(&config).expect("count chunks"), + chunks + .count_chunks(&ChunkQuery::default(), None) + .await + .expect("count chunks"), 5, "every written item must still reach the memory tree" );