diff --git a/src/memory/chunks/connection.rs b/src/memory/chunks/connection.rs index e1786a7..f670438 100644 --- a/src/memory/chunks/connection.rs +++ b/src/memory/chunks/connection.rs @@ -23,7 +23,8 @@ pub(crate) use super::connection_breaker::CB_THRESHOLD; use super::migrations::{migrate_legacy_embeddings_to_sidecar, purge_global_topic_trees}; use super::recovery::{ - is_corrupt_error, is_io_open_error, quarantine_corrupt_files, try_cleanup_stale_files, + is_corrupt_error, is_io_open_error, is_transient_cold_start, quarantine_corrupt_files, + try_cleanup_stale_files, }; use super::schema::SCHEMA; use super::{db_path_for, SQLITE_BUSY_TIMEOUT}; @@ -366,6 +367,40 @@ pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result Result { + // Bootstrap I/O errors (`CANTOPEN`, the `-shm`/`-wal` cold-start races, + // `IOERR_FSTAT`) come from a sibling connection creating or truncating the + // same file at the same instant — a fresh open a moment later succeeds. Give + // a transient failure a couple of short, backed-off retries before + // surfacing it, which is what removes the flaky parallel-open failure. + const MAX_ATTEMPTS: u32 = 3; + let mut attempt = 0; + loop { + match try_open_and_init(db_path, config) { + Ok(conn) => return Ok(conn), + Err(error) if attempt + 1 < MAX_ATTEMPTS && is_transient_cold_start(&error) => { + let backoff = std::time::Duration::from_millis(20 * (1 << attempt)); + // `tracing` is optional and only arrives with `sync`; the chunk + // store compiles in the dependency-light default build too, so + // the call has to be gated or that build cannot link. + #[cfg(feature = "sync")] + tracing::debug!( + db_path = %db_path.display(), + attempt = attempt + 1, + backoff_ms = backoff.as_millis() as u64, + "[chunks] transient cold-start opening chunk DB; retrying" + ); + std::thread::sleep(backoff); + attempt += 1; + } + Err(error) => return Err(error), + } + } +} + +/// One open+init attempt. Retried by [`open_and_init`] on a transient +/// cold-start failure. Does not clean up a partially created directory or file +/// on failure — a subsequent attempt simply retries against the same path. +fn try_open_and_init(db_path: &Path, config: &MemoryConfig) -> Result { let dir = db_path.parent().expect("db_path always has a parent"); std::fs::create_dir_all(dir) .with_context(|| format!("Failed to create chunk DB dir: {}", dir.display()))?; diff --git a/src/memory/chunks/embeddings.rs b/src/memory/chunks/embeddings.rs index 0e07e25..3ef5561 100644 --- a/src/memory/chunks/embeddings.rs +++ b/src/memory/chunks/embeddings.rs @@ -6,6 +6,7 @@ //! ported here) — callers pass vectors in. use super::connection::with_connection; +use super::signature::{format_signature, signature_in_clause, signature_variants}; use anyhow::{Context, Result}; use chrono::Utc; use rusqlite::{Connection, OptionalExtension}; @@ -20,11 +21,20 @@ pub(crate) fn active_embedding_dims(config: &MemoryConfig) -> usize { } /// Resolve the active embedding signature — the canonical key every per-model -/// sidecar read/write is scoped by. Derived from the configured model + dim so -/// a provider/model/dimension switch becomes a query-time filter rather than a -/// destructive rewrite. +/// sidecar read/write is scoped by. Derived from the configured provider, model +/// and dim so a provider/model/dimension switch becomes a query-time filter +/// rather than a destructive rewrite. +/// +/// This is the canonical `provider=…;model=…;dims=…` spelling, shared with the +/// namespace store. Rows written under the tree's older `{model}@{dims}` +/// spelling are still found, because per-signature reads match every variant +/// (see [`signature_variants`]) rather than one exact string. pub fn tree_active_signature(config: &MemoryConfig) -> String { - format!("{}@{}", config.embedding.model, config.embedding.dim) + format_signature( + &config.embedding.provider, + &config.embedding.model, + config.embedding.dim, + ) } /// Store a chunk's embedding under the active model signature (see @@ -292,7 +302,8 @@ pub(crate) fn validate_reembed_skip_key<'a>(label: &str, value: &'a str) -> Resu Ok(trimmed) } -/// Fetch a chunk embedding for exactly one provider/model/dimension signature. +/// Fetch a chunk embedding for one provider/model/dimension signature — under +/// any of its spellings (see [`signature_variants`]). /// /// Returns `Ok(None)` when no row exists for `(chunk_id, model_signature)` — /// absence is not an error. @@ -307,13 +318,22 @@ pub fn get_chunk_embedding_for_signature( chunk_id: &str, model_signature: &str, ) -> Result>> { + let variants = signature_variants(model_signature); with_connection(config, |conn| { + let mut params: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(variants.len() + 1); + params.push(&chunk_id as &dyn rusqlite::ToSql); + for variant in &variants { + params.push(variant as &dyn rusqlite::ToSql); + } let row: Option<(Vec, i64)> = conn .query_row( - "SELECT vector, dim - FROM mem_tree_chunk_embeddings - WHERE chunk_id = ?1 AND model_signature = ?2", - rusqlite::params![chunk_id, model_signature], + &format!( + "SELECT vector, dim + FROM mem_tree_chunk_embeddings + WHERE chunk_id = ?1 AND model_signature {}", + signature_in_clause(variants.len(), 2) + ), + params.as_slice(), |r| Ok((r.get(0)?, r.get(1)?)), ) .optional()?; @@ -374,26 +394,40 @@ fn embedding_from_blob(bytes: &[u8], dim: i64, label: &str) -> Result rusqlite::Result { + let variants = signature_variants(model_signature); + let sig_clause = signature_in_clause(variants.len(), 1); + let params: Vec<&dyn rusqlite::ToSql> = variants + .iter() + .map(|variant| variant as &dyn rusqlite::ToSql) + .collect(); conn.query_row( - "SELECT EXISTS( - SELECT 1 FROM mem_tree_chunks c - WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e - WHERE e.chunk_id = c.id AND e.model_signature = ?1) - AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk - WHERE sk.chunk_id = c.id AND sk.model_signature = ?1)) - OR EXISTS( - SELECT 1 FROM mem_tree_summaries s - WHERE s.deleted = 0 - AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e - WHERE e.summary_id = s.id AND e.model_signature = ?1) - AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk - WHERE sk.summary_id = s.id AND sk.model_signature = ?1))", - rusqlite::params![model_signature], + &format!( + "SELECT EXISTS( + SELECT 1 FROM mem_tree_chunks c + WHERE NOT EXISTS (SELECT 1 FROM mem_tree_chunk_embeddings e + WHERE e.chunk_id = c.id AND e.model_signature {sig_clause}) + AND NOT EXISTS (SELECT 1 FROM mem_tree_chunk_reembed_skipped sk + WHERE sk.chunk_id = c.id AND sk.model_signature {sig_clause})) + OR EXISTS( + SELECT 1 FROM mem_tree_summaries s + WHERE s.deleted = 0 + AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_embeddings e + WHERE e.summary_id = s.id AND e.model_signature {sig_clause}) + AND NOT EXISTS (SELECT 1 FROM mem_tree_summary_reembed_skipped sk + WHERE sk.summary_id = s.id AND sk.model_signature {sig_clause}))" + ), + params.as_slice(), |row| row.get(0), ) } @@ -427,6 +461,7 @@ pub fn get_chunk_embeddings_for_signature_batch( if chunk_ids.is_empty() { return Ok(HashMap::new()); } + let variants = signature_variants(model_signature); with_connection(config, |conn| { let mut out: HashMap> = HashMap::with_capacity(chunk_ids.len()); for window in chunk_ids.chunks(MAX_EMBEDDING_BATCH) { @@ -437,17 +472,20 @@ pub fn get_chunk_embeddings_for_signature_batch( "SELECT chunk_id, vector, dim FROM mem_tree_chunk_embeddings WHERE chunk_id IN ({placeholders}) - AND model_signature = ?{sig_idx}", - sig_idx = window.len() + 1, + AND model_signature {sig_clause}", + sig_clause = signature_in_clause(variants.len(), window.len() + 1), ); let mut stmt = conn .prepare(&sql) .context("prepare get_chunk_embeddings_for_signature_batch")?; - let mut params: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(window.len() + 1); + let mut params: Vec<&dyn rusqlite::ToSql> = + Vec::with_capacity(window.len() + variants.len()); for id in window { params.push(id as &dyn rusqlite::ToSql); } - params.push(&model_signature as &dyn rusqlite::ToSql); + for variant in &variants { + params.push(variant as &dyn rusqlite::ToSql); + } let rows = stmt .query_map(params.as_slice(), |row| { Ok(( diff --git a/src/memory/chunks/mod.rs b/src/memory/chunks/mod.rs index 167e4d9..62a2b00 100644 --- a/src/memory/chunks/mod.rs +++ b/src/memory/chunks/mod.rs @@ -51,6 +51,8 @@ mod produce_split; mod raw_refs; #[path = "semantic.rs"] mod semantic; +#[path = "signature.rs"] +mod signature; #[path = "store.rs"] mod store; #[path = "store_delete.rs"] @@ -98,6 +100,10 @@ pub use raw_refs::{ pub use recovery::{ is_io_open_error, is_transient_cold_start, recover_corrupt_db, try_cleanup_stale_files, }; +pub(crate) use signature::signature_in_clause; +pub use signature::{ + format_signature, parse_signature, signature_variants, signatures_equivalent, SignatureParts, +}; pub use store::{ claim_source_ingest_tx, count_chunks, count_chunks_by_lifecycle_status, count_raw_paths_ingested_with_prefix, delete_source_ingest, extraction_coverage, diff --git a/src/memory/chunks/recovery.rs b/src/memory/chunks/recovery.rs index 5e1fca9..127b319 100644 --- a/src/memory/chunks/recovery.rs +++ b/src/memory/chunks/recovery.rs @@ -43,6 +43,10 @@ const SQLITE_IOERR_SHMSIZE: i32 = 4874; const SQLITE_IOERR_SHMMAP: i32 = 5386; /// `IOERR_IN_PAGE` — an mmap-page I/O fault, also seen under WAL cold-start. const SQLITE_IOERR_IN_PAGE: i32 = 8714; +/// `IOERR_FSTAT` — an `fstat()` on the db/side-file failed, observed when a +/// sibling connection is creating or truncating the file at the same instant +/// (seen as flaky `Error code 1802: disk I/O error` under parallel cold opens). +const SQLITE_IOERR_FSTAT: i32 = 1802; /// True if `err` (or anything in its cause chain) is one of the SQLite codes /// that fire during cold-start WAL/SHM bootstrap races. @@ -67,6 +71,7 @@ pub fn is_transient_cold_start(err: &anyhow::Error) -> bool { | SQLITE_IOERR_SHMSIZE | SQLITE_IOERR_SHMMAP | SQLITE_IOERR_IN_PAGE + | SQLITE_IOERR_FSTAT ); } false diff --git a/src/memory/chunks/recovery_tests.rs b/src/memory/chunks/recovery_tests.rs index 6cfd33c..dd38dbe 100644 --- a/src/memory/chunks/recovery_tests.rs +++ b/src/memory/chunks/recovery_tests.rs @@ -147,3 +147,22 @@ fn recover_corrupt_db_is_noop_on_healthy_db() { .any(|e| e.file_name().to_string_lossy().contains(".corrupt-")); assert!(!quarantined, "no quarantine file should be created"); } + +#[test] +fn is_transient_cold_start_classifies_ioerr_fstat() { + // `IOERR_FSTAT` (1802) is the flaky `disk I/O error` a parallel cold open + // hits; it must be treated as transient so `open_and_init` retries instead + // of surfacing it. A non-transient sqlite error must stay non-transient. + use super::is_transient_cold_start; + let fstat = rusqlite::Error::SqliteFailure( + rusqlite::ffi::Error::new(1802), + Some("disk I/O error".to_string()), + ); + assert!(is_transient_cold_start(&anyhow::Error::from(fstat))); + + let constraint = rusqlite::Error::SqliteFailure( + rusqlite::ffi::Error::new(19), + Some("constraint".to_string()), + ); + assert!(!is_transient_cold_start(&anyhow::Error::from(constraint))); +} diff --git a/src/memory/chunks/signature.rs b/src/memory/chunks/signature.rs new file mode 100644 index 0000000..0946f9e --- /dev/null +++ b/src/memory/chunks/signature.rs @@ -0,0 +1,155 @@ +//! Embedding-signature format and equivalence. +//! +//! One logical vector space has had two spellings in the wild: +//! +//! * `provider={provider};model={model};dims={dims}` — the canonical form, and +//! what the namespace store has always written. +//! * `{model}@{dims}` — the form the tree sidecars used. +//! +//! Every per-signature read is an exact string match, so the two spellings +//! partition one vector space in half: a store that changed spelling stops +//! seeing its own prior vectors and silently scores against nothing. That is +//! not hypothetical — it left hundreds of live tree chunks unreadable after a +//! format change that was never a model change. +//! +//! So writes emit the canonical form and reads match *any* spelling of the same +//! space via [`signature_variants`]. Nothing is rewritten on disk: a legacy row +//! keeps its own signature and simply becomes visible again. The comparison +//! identity is `(model, dims)` — provider is checked only when both sides +//! declare one, because the legacy spelling cannot carry it. + +/// Render the canonical signature for one vector space. +/// +/// Delegates to the same formatter the namespace store's embedding providers +/// use ([`crate::memory::store::vectors::format_embedding_signature`], itself a +/// re-export from `tinyagents`), so the two stores cannot drift into +/// byte-different spellings of the same space again. +pub fn format_signature(provider: &str, model: &str, dims: usize) -> String { + tinyagents::harness::embeddings::format_embedding_signature(provider, model, dims) +} + +/// A signature decomposed into the parts that identify its vector space. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SignatureParts { + /// Embedding backend (`cloud`, `ollama`, …). `None` for the legacy + /// spelling, which never recorded it. + pub provider: Option, + /// Model identifier — the one part every spelling carries. + pub model: String, + /// Vector dimensionality, when the spelling records it. + pub dims: Option, +} + +/// Decompose either spelling. Returns `None` only for a blank signature. +/// +/// An unrecognised shape degrades to "all of it is the model name" rather than +/// failing: an unparseable signature still has to compare equal to itself, or a +/// store using some third spelling would lose every vector it ever wrote. +pub fn parse_signature(signature: &str) -> Option { + let trimmed = signature.trim(); + if trimmed.is_empty() { + return None; + } + if trimmed.contains('=') { + let mut provider = None; + let mut model = None; + let mut dims = None; + for field in trimmed.split(';') { + let Some((key, value)) = field.split_once('=') else { + continue; + }; + let value = value.trim(); + match key.trim() { + "provider" => provider = Some(value.to_string()), + "model" => model = Some(value.to_string()), + // `dim` is accepted alongside `dims` so a caller that spells the + // key either way still resolves to the same vector space. + "dims" | "dim" => dims = value.parse::().ok(), + _ => {} + } + } + return Some(SignatureParts { + provider, + model: model.unwrap_or_else(|| trimmed.to_string()), + dims, + }); + } + // Legacy `{model}@{dims}`. Split at the LAST `@` so a model id containing + // one keeps it. + if let Some((model, dims)) = trimmed.rsplit_once('@') { + if let Ok(dims) = dims.trim().parse::() { + return Some(SignatureParts { + provider: None, + model: model.trim().to_string(), + dims: Some(dims), + }); + } + } + Some(SignatureParts { + provider: None, + model: trimmed.to_string(), + dims: None, + }) +} + +/// Whether two signatures name the same vector space. +/// +/// Equal models, and equal dimensions whenever both spellings state one. A +/// provider mismatch disqualifies only when both sides declare a provider — +/// the legacy spelling declares none, and refusing to match it would defeat the +/// whole point. +pub fn signatures_equivalent(a: &str, b: &str) -> bool { + let (Some(left), Some(right)) = (parse_signature(a), parse_signature(b)) else { + return false; + }; + if left.model != right.model { + return false; + } + if let (Some(left_dims), Some(right_dims)) = (left.dims, right.dims) { + if left_dims != right_dims { + return false; + } + } + match (left.provider, right.provider) { + (Some(left_provider), Some(right_provider)) => left_provider == right_provider, + _ => true, + } +} + +/// Every spelling of `signature` a stored row might carry, `signature` first. +/// +/// Bind these into an `IN (…)` predicate (see [`signature_in_clause`]) instead +/// of `= ?`: that is what makes a read find rows written under the other +/// convention without rewriting them. +pub fn signature_variants(signature: &str) -> Vec { + let mut variants = vec![signature.trim().to_string()]; + let Some(parts) = parse_signature(signature) else { + return variants; + }; + let mut push = |candidate: String| { + if !variants.contains(&candidate) { + variants.push(candidate); + } + }; + if let Some(dims) = parts.dims { + push(format!("{}@{}", parts.model, dims)); + if let Some(provider) = parts.provider.as_deref() { + push(format_signature(provider, &parts.model, dims)); + } + } + variants +} + +/// Build `IN (?n, ?n+1, …)` for `count` variants starting at parameter +/// `first_index` (rusqlite parameters are 1-based). +pub(crate) fn signature_in_clause(count: usize, first_index: usize) -> String { + let placeholders = (0..count) + .map(|offset| format!("?{}", first_index + offset)) + .collect::>() + .join(", "); + format!("IN ({placeholders})") +} + +#[cfg(test)] +#[path = "signature_tests.rs"] +mod tests; diff --git a/src/memory/chunks/signature_tests.rs b/src/memory/chunks/signature_tests.rs new file mode 100644 index 0000000..2a354f8 --- /dev/null +++ b/src/memory/chunks/signature_tests.rs @@ -0,0 +1,117 @@ +//! Tests for signature parsing, equivalence, and variant expansion. + +use super::*; + +#[test] +fn parses_the_canonical_spelling() { + let parts = parse_signature("provider=cloud;model=embedding-v1;dims=1024").unwrap(); + assert_eq!(parts.provider.as_deref(), Some("cloud")); + assert_eq!(parts.model, "embedding-v1"); + assert_eq!(parts.dims, Some(1024)); +} + +#[test] +fn parses_the_legacy_spelling() { + let parts = parse_signature("embedding-v1@1024").unwrap(); + assert_eq!(parts.provider, None); + assert_eq!(parts.model, "embedding-v1"); + assert_eq!(parts.dims, Some(1024)); +} + +#[test] +fn a_model_id_containing_an_at_sign_keeps_it() { + // Split at the LAST `@`, so `openai@v2@1536` is model `openai@v2`. + let parts = parse_signature("openai@v2@1536").unwrap(); + assert_eq!(parts.model, "openai@v2"); + assert_eq!(parts.dims, Some(1536)); +} + +#[test] +fn an_unrecognised_shape_degrades_to_a_bare_model_name() { + let parts = parse_signature("some-third-convention").unwrap(); + assert_eq!(parts.model, "some-third-convention"); + assert_eq!(parts.dims, None); + // …and still compares equal to itself, so such a store keeps its vectors. + assert!(signatures_equivalent( + "some-third-convention", + "some-third-convention" + )); + assert!(parse_signature(" ").is_none()); +} + +#[test] +fn the_two_spellings_of_one_space_are_equivalent() { + assert!(signatures_equivalent( + "provider=cloud;model=embedding-v1;dims=1024", + "embedding-v1@1024" + )); + // …in both directions, and regardless of which side carries the provider. + assert!(signatures_equivalent( + "embedding-v1@1024", + "provider=cloud;model=embedding-v1;dims=1024" + )); +} + +#[test] +fn a_different_model_or_dimension_is_a_different_space() { + assert!(!signatures_equivalent( + "provider=cloud;model=embedding-v1;dims=1024", + "nomic-embed-text@1024" + )); + assert!(!signatures_equivalent( + "provider=cloud;model=embedding-v1;dims=1024", + "embedding-v1@768" + )); + // Two *declared* providers that disagree stay separate. + assert!(!signatures_equivalent( + "provider=cloud;model=embedding-v1;dims=1024", + "provider=ollama;model=embedding-v1;dims=1024" + )); +} + +#[test] +fn variants_cover_both_spellings_with_the_input_first() { + let variants = signature_variants("provider=cloud;model=embedding-v1;dims=1024"); + assert_eq!( + variants, + vec![ + "provider=cloud;model=embedding-v1;dims=1024".to_string(), + "embedding-v1@1024".to_string(), + ] + ); + + // From the legacy side the canonical spelling is unknowable (no provider), + // so the expansion is just the input — matching stays correct, it simply + // cannot invent a provider it was never told. + assert_eq!( + signature_variants("embedding-v1@1024"), + vec!["embedding-v1@1024".to_string()] + ); +} + +#[test] +fn variants_never_repeat_a_spelling() { + // `dim` (singular) parses, so the canonical rebuild would duplicate the + // legacy form if dedup were missing. + let variants = signature_variants("embedding-v1@1024"); + let mut sorted = variants.clone(); + sorted.sort(); + sorted.dedup(); + assert_eq!(sorted.len(), variants.len(), "variants must be unique"); +} + +#[test] +fn the_canonical_spelling_matches_the_namespace_store_byte_for_byte() { + // The whole point of the unification: one space, one string, both stores. + assert_eq!( + format_signature("cloud", "embedding-v1", 1024), + crate::memory::store::vectors::format_embedding_signature("cloud", "embedding-v1", 1024) + ); +} + +#[test] +fn in_clause_numbers_placeholders_from_the_given_index() { + assert_eq!(signature_in_clause(2, 1), "IN (?1, ?2)"); + assert_eq!(signature_in_clause(2, 5), "IN (?5, ?6)"); + assert_eq!(signature_in_clause(1, 3), "IN (?3)"); +} diff --git a/src/memory/chunks/store_embed_tests.rs b/src/memory/chunks/store_embed_tests.rs index 4d364a8..5207db9 100644 --- a/src/memory/chunks/store_embed_tests.rs +++ b/src/memory/chunks/store_embed_tests.rs @@ -453,3 +453,53 @@ fn extraction_coverage_reflects_indexed_fraction() { .unwrap(); assert!((extraction_coverage(&cfg).unwrap() - 1.0).abs() < 1e-6); } + +#[test] +fn a_vector_written_under_the_legacy_spelling_is_readable_under_the_active_signature() { + // The regression this exists for: the tree used to key sidecars by + // `{model}@{dims}`. After the convention was unified on the canonical + // `provider=…;model=…;dims=…` string, an exact-match read stopped seeing + // those rows — hundreds of live chunks scored against nothing, from a + // format change that was never a model change. Nothing is rewritten on + // disk; the read simply accepts both spellings. + let (_tmp, mut cfg) = test_config(); + cfg.embedding.provider = "cloud".into(); + cfg.embedding.model = "embedding-v1".into(); + cfg.embedding.dim = 3; + + let chunk = sample_chunk("gmail:inbox", 0, 1_700_000_000_000); + upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); + set_chunk_embedding_for_signature(&cfg, &chunk.id, "embedding-v1@3", &[0.1, 0.2, 0.3]).unwrap(); + + let active = tree_active_signature(&cfg); + assert_eq!(active, "provider=cloud;model=embedding-v1;dims=3"); + assert_eq!( + get_chunk_embedding(&cfg, &chunk.id).unwrap(), + Some(vec![0.1, 0.2, 0.3]), + "a legacy-spelled vector must be visible under the canonical signature" + ); + + // Batch read — the retrieval hot path — sees it too. + let batch = + get_chunk_embeddings_for_signature_batch(&cfg, std::slice::from_ref(&chunk.id), &active) + .unwrap(); + assert_eq!(batch.get(&chunk.id), Some(&vec![0.1, 0.2, 0.3])); + + // …and it counts as covered, so the re-embed chain does not pay to + // recompute a vector that is already on disk. + with_connection(&cfg, |conn| { + assert!( + !super::has_uncovered_reembed_work(conn, &active)?, + "a legacy-spelled vector must count as coverage" + ); + Ok(()) + }) + .unwrap(); + + // A genuinely different space is still separate. + assert_eq!( + get_chunk_embedding_for_signature(&cfg, &chunk.id, "provider=cloud;model=other-v2;dims=3") + .unwrap(), + None + ); +} diff --git a/src/memory/config.rs b/src/memory/config.rs index 74f2545..ab8220e 100644 --- a/src/memory/config.rs +++ b/src/memory/config.rs @@ -245,6 +245,11 @@ fn validate_budget(name: &str, budget: &SyncBudgetConfig) -> anyhow::Result<()> pub struct EmbeddingConfig { /// Vector dimension. OpenHuman fixes this at 768. pub dim: usize, + /// Embedding backend name (`cloud`, `ollama`, …). Part of the signature + /// every per-model sidecar row is keyed by, so two backends serving the + /// same model id never share a vector space. + #[serde(default = "default_embedding_provider")] + pub provider: String, /// Backend model identifier (default Ollama `nomic-embed-text`). pub model: String, /// When `true`, ingest fails if embeddings are unavailable instead of @@ -252,10 +257,16 @@ pub struct EmbeddingConfig { pub strict: bool, } +/// Matches the default `model` below: `nomic-embed-text` is served locally. +fn default_embedding_provider() -> String { + "ollama".to_string() +} + impl Default for EmbeddingConfig { fn default() -> Self { Self { dim: DEFAULT_EMBEDDING_DIM, + provider: default_embedding_provider(), model: "nomic-embed-text".to_string(), strict: false, } diff --git a/src/memory/sources/types.rs b/src/memory/sources/types.rs index 18b4da5..a97d271 100644 --- a/src/memory/sources/types.rs +++ b/src/memory/sources/types.rs @@ -263,7 +263,14 @@ impl MemorySourcePatch { if self.selector.is_some() && kind != SourceKind::WebPage { return reject("selector"); } - if matches!(self.max_items, Some(Some(_))) && kind != SourceKind::RssFeed { + // `max_items` is the per-run ingest cap. It applies to RSS feeds and to + // Composio connections — the host UI (`SourceSettingsPanel`) exposes it + // for both, and a Composio source is created with a toolkit default, so + // rejecting it on edit desynced the UI from the store. Other kinds have + // no per-run item cap. + if matches!(self.max_items, Some(Some(_))) + && !matches!(kind, SourceKind::RssFeed | SourceKind::Composio) + { return reject("max_items"); } if self.url.is_some() diff --git a/src/memory/sources/types_tests.rs b/src/memory/sources/types_tests.rs index 9d63bd3..ed9f5ec 100644 --- a/src/memory/sources/types_tests.rs +++ b/src/memory/sources/types_tests.rs @@ -242,3 +242,20 @@ pub(super) fn default_entry() -> MemorySourceEntry { sync_depth_days: None, } } + +#[test] +fn max_items_is_applicable_to_composio_and_rss_but_not_other_kinds() { + // The host UI exposes `max_items` for Composio sources and creates them with + // a toolkit default, so editing one must not be rejected — the regression + // this guards ("field 'max_items' is not applicable to source kind + // 'composio'"). RSS keeps it; kinds with no per-run item cap still reject. + let patch = || MemorySourcePatch { + max_items: Some(Some(100)), + ..Default::default() + }; + assert!(patch().validate_for_kind(SourceKind::Composio).is_ok()); + assert!(patch().validate_for_kind(SourceKind::RssFeed).is_ok()); + assert!(patch().validate_for_kind(SourceKind::Folder).is_err()); + assert!(patch().validate_for_kind(SourceKind::GithubRepo).is_err()); + assert!(patch().validate_for_kind(SourceKind::WebPage).is_err()); +} diff --git a/src/memory/sync/composio/gmail.rs b/src/memory/sync/composio/gmail.rs index 0b1d9f2..cf20a3e 100644 --- a/src/memory/sync/composio/gmail.rs +++ b/src/memory/sync/composio/gmail.rs @@ -1,6 +1,9 @@ //! Incremental Gmail synchronization through Composio. +use std::sync::Arc; + use async_trait::async_trait; +use chrono::{DateTime, Utc}; use serde_json::Value; use super::client::{ActionExecutor, ComposioClient}; @@ -8,6 +11,8 @@ use super::orchestrator::{ run_incremental_sync, IncrementalSource, PageFetch, SyncItem, SyncScope, }; use crate::memory::config::MemoryConfig; +use crate::memory::ingest::canonicalize::email::{self, EmailMessage, EmailThread}; +use crate::memory::ingest::canonicalize::email_clean; use crate::memory::sync::state::SyncState; use crate::memory::sync::traits::{ SkillDocument, SyncContext, SyncOutcome, SyncPipeline, SyncPipelineKind, @@ -16,7 +21,7 @@ use crate::memory::sync::traits::{ const ACTION_FETCH_EMAILS: &str = "GMAIL_FETCH_EMAILS"; pub struct GmailSyncPipeline { - client: ComposioClient, + executor: Arc, connection_id: String, max_pages: usize, page_size: usize, @@ -24,9 +29,24 @@ pub struct GmailSyncPipeline { } impl GmailSyncPipeline { + /// Sync through a plain Composio client. pub fn new(client: ComposioClient, connection_id: impl Into) -> Self { + Self::with_executor(Arc::new(client), connection_id) + } + + /// Sync through a caller-supplied executor. + /// + /// The seam exists for host-side response reshaping: the Gmail envelope + /// rewrite (verbose MIME payload → one slim record per message, body + /// pre-rendered into `markdown`) lives in the host, above this crate, so + /// wrapping the executor is the only way it can reach the fetched page + /// before [`document`](SyncPipeline) turns it into a stored document. + pub fn with_executor( + executor: Arc, + connection_id: impl Into, + ) -> Self { Self { - client, + executor, connection_id: connection_id.into(), max_pages: 10, // Gmail fetches full message payloads (`include_payload: true`), so a @@ -69,7 +89,14 @@ impl SyncPipeline for GmailSyncPipeline { _config: &MemoryConfig, context: &SyncContext, ) -> anyhow::Result { - run_incremental_sync(self, &self.client, &self.connection_id, _config, context).await + run_incremental_sync( + self, + self.executor.as_ref(), + &self.connection_id, + _config, + context, + ) + .await } } @@ -94,6 +121,15 @@ impl IncrementalSource for GmailSyncPipeline { true } + /// Gmail pages are capped by `max_results`, and full message payloads make + /// a page's size depend on what is *in* the mail — a handful of large + /// attachments is enough for the provider to refuse 25 messages it accepted + /// yesterday. Naming the argument lets the orchestrator halve it and retry + /// rather than leaving the source stuck. + fn page_size_arg_key(&self) -> Option<&'static str> { + Some("max_results") + } + fn arguments( &self, _scope: &SyncScope, @@ -153,7 +189,7 @@ impl IncrementalSource for GmailSyncPipeline { connection_id: connection_id.into(), document_id: format!("gmail:{id}"), title: message_title(&item.raw), - content: serde_json::to_string_pretty(&item.raw)?, + content: canonical_markdown(&item.raw, &id), toolkit: "gmail".into(), metadata: serde_json::json!({ "source": "composio-provider-incremental", @@ -219,6 +255,117 @@ fn message_title(message: &Value) -> String { .to_owned() } +/// Render one Gmail message as canonical Markdown — the same shape the memory +/// tree ingests — rather than the provider's raw JSON. +/// +/// Storing `to_string_pretty(&item.raw)` puts a MIME tree, `Received:` headers +/// and base64 part bodies into the document: the literal words of the mail are +/// either absent or split mid-token by the chunker, so recall can never match +/// them. Routing through [`email::canonicalise`] reuses the canonicaliser the +/// tree already uses — headers as a small block, body through +/// `email_clean::clean_body` (reply chains and footer boilerplate stripped). +fn canonical_markdown(message: &Value, id: &str) -> String { + let thread = email_thread(message, id); + match email::canonicalise(&format!("gmail:{id}"), "", &[], thread) { + Ok(Some(canonical)) => canonical.markdown, + // The thread built here always holds exactly one message, so an empty + // thread (`None`) and `Err` are unreachable in practice. Degrade to the + // bare body rather than dropping the message out of memory. + _ => message_body(message), + } +} + +/// Adapt one provider message into the canonicaliser's input shape. A Gmail +/// sync item is a single message, so the thread wraps exactly one. +fn email_thread(message: &Value, id: &str) -> EmailThread { + let subject = message_title(message); + EmailThread { + provider: "gmail".into(), + thread_subject: subject.clone(), + messages: vec![EmailMessage { + from: message_sender(message), + to: message_recipients(message), + cc: Vec::new(), + subject, + sent_at: message_sent_at(message), + body: message_body(message), + source_ref: Some(format!("gmail:{id}")), + list_unsubscribe: None, + }], + } +} + +/// Body text for one message, best rendering first. +/// +/// `markdown` is what the Gmail response reshaper pins onto each message (HTML +/// stripped, URLs shortened, footers removed). `messageText` is the provider's +/// own plain-text rendering, used when the reshape did not run. `snippet` is a +/// last resort: truncated, but real prose — unlike the raw payload. +fn message_body(message: &Value) -> String { + ["markdown", "markdownFormatted", "messageText", "snippet"] + .iter() + .find_map(|key| nonempty_str(message, key)) + .unwrap_or_default() +} + +/// Sender header, rendered as `From:` and used by the canonicaliser as the +/// participant key. +fn message_sender(message: &Value) -> String { + ["from", "sender"] + .iter() + .find_map(|key| nonempty_str(message, key)) + .unwrap_or_else(|| "unknown".to_owned()) +} + +/// Recipients arrive as one comma-joined header string (some responses use an +/// array); split them so the canonicaliser can render a `To:` line. +fn message_recipients(message: &Value) -> Vec { + match message.get("to") { + Some(Value::String(header)) => header + .split(',') + .map(str::trim) + .filter(|address| !address.is_empty()) + .map(str::to_owned) + .collect(), + Some(Value::Array(values)) => values + .iter() + .filter_map(Value::as_str) + .map(str::trim) + .filter(|address| !address.is_empty()) + .map(str::to_owned) + .collect(), + _ => Vec::new(), + } +} + +/// Send time, preferring the canonicaliser's own `Value`-level date parser (it +/// already knows `date`, `internalDate`, and epoch-ms-as-string) and falling +/// back to the sync cursor. The epoch is the last resort because it is +/// *deterministic*: a message the provider dated with nothing must not rewrite +/// its own content — and so re-chunk and re-embed — on every sync. +fn message_sent_at(message: &Value) -> DateTime { + email_clean::parse_message_date(message) + .or_else(|| { + item_cursor(message) + .as_deref() + .and_then(cursor_to_seconds) + .and_then(|seconds| DateTime::from_timestamp(seconds, 0)) + }) + .unwrap_or_else(|| DateTime::from_timestamp(0, 0).expect("epoch is a valid timestamp")) +} + +/// Read `key` as a trimmed, non-empty string. Unlike a plain `get(..).as_str()` +/// chain over a candidate list, a present-but-blank field falls through to the +/// next candidate instead of ending the search. +fn nonempty_str(message: &Value, key: &str) -> Option { + message + .get(key) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_owned) +} + fn cursor_to_seconds(cursor: &str) -> Option { if let Ok(milliseconds) = cursor.trim().parse::() { return Some(milliseconds / 1000); @@ -227,3 +374,7 @@ fn cursor_to_seconds(cursor: &str) -> Option { .ok() .map(|date| date.timestamp()) } + +#[cfg(test)] +#[path = "gmail_tests.rs"] +mod tests; diff --git a/src/memory/sync/composio/gmail_tests.rs b/src/memory/sync/composio/gmail_tests.rs new file mode 100644 index 0000000..d4de768 --- /dev/null +++ b/src/memory/sync/composio/gmail_tests.rs @@ -0,0 +1,101 @@ +//! Tests for the Gmail message → canonical Markdown adapter. + +use serde_json::json; + +use super::{canonical_markdown, message_body, message_recipients, message_sent_at}; + +/// One message in the shape the Gmail response reshaper emits: a slim envelope +/// whose body is pre-rendered into `markdown`. +fn slim_message() -> serde_json::Value { + json!({ + "id": "18f0abc", + "threadId": "18f0abc", + "subject": "Boulder visit", + "from": "Advising ", + "to": "me@example.com, second@example.com", + "date": "2026-05-02T09:15:00Z", + "labels": ["INBOX"], + "markdown": "The University of Colorado orientation is on May 20.\n\nOn Fri, 1 May 2026, someone wrote:\n> please ignore this quoted reply", + }) +} + +#[test] +fn canonical_markdown_renders_headers_and_cleaned_body() { + let content = canonical_markdown(&slim_message(), "18f0abc"); + + // The literal words of the mail — the thing recall has to match — are + // present as prose, and the headers are readable rather than a MIME tree. + assert!( + content.contains("The University of Colorado orientation is on May 20."), + "body text must survive canonicalisation: {content}" + ); + assert!( + content.contains("From: Advising "), + "{content}" + ); + assert!(content.contains("Subject: Boulder visit"), "{content}"); + assert!( + content.contains("To: me@example.com, second@example.com"), + "{content}" + ); + + // Canonicalisation is what strips the quoted reply chain. + assert!( + !content.contains("please ignore this quoted reply"), + "reply chain must be stripped by clean_body: {content}" + ); + + // Nothing JSON-shaped is left: this is the regression the fix exists for. + assert!( + !content.contains("\"markdown\""), + "raw JSON must not be stored: {content}" + ); + assert!( + !content.contains("threadId"), + "envelope keys must not be stored: {content}" + ); +} + +#[test] +fn body_falls_back_to_message_text_when_the_reshape_did_not_run() { + // No `markdown` field — the provider's own plain text is used instead. + let raw = json!({ + "id": "18f0def", + "subject": "Direct", + "messageText": "Plain provider text about Colorado.", + }); + assert_eq!(message_body(&raw), "Plain provider text about Colorado."); + assert!(canonical_markdown(&raw, "18f0def").contains("Plain provider text about Colorado.")); +} + +#[test] +fn body_skips_a_present_but_blank_field() { + // A blank `markdown` must not shadow a usable `messageText`: the candidate + // list falls through on emptiness, not just on absence. + let raw = json!({ "markdown": " ", "messageText": "real body" }); + assert_eq!(message_body(&raw), "real body"); +} + +#[test] +fn recipients_split_from_either_a_header_string_or_an_array() { + let joined = json!({ "to": "a@x.com, b@y.com" }); + assert_eq!(message_recipients(&joined), vec!["a@x.com", "b@y.com"]); + + let array = json!({ "to": ["a@x.com", " b@y.com "] }); + assert_eq!(message_recipients(&array), vec!["a@x.com", "b@y.com"]); + + assert!(message_recipients(&json!({})).is_empty()); +} + +#[test] +fn sent_at_reads_epoch_millis_and_is_deterministic_when_undated() { + // Gmail's `internalDate` is epoch millis as a string. + let dated = json!({ "internalDate": "1777712100000" }); + assert_eq!(message_sent_at(&dated).timestamp(), 1_777_712_100); + + // Undated messages must resolve to the same value every sync, otherwise the + // rendered `Date:` header changes and the document re-chunks forever. + let undated = json!({ "subject": "no date anywhere" }); + assert_eq!(message_sent_at(&undated), message_sent_at(&undated)); + assert_eq!(message_sent_at(&undated).timestamp(), 0); +} diff --git a/src/memory/sync/composio/mod.rs b/src/memory/sync/composio/mod.rs index 1af0a74..d645c4f 100644 --- a/src/memory/sync/composio/mod.rs +++ b/src/memory/sync/composio/mod.rs @@ -4,6 +4,7 @@ pub mod client; pub mod connect; pub mod gmail; pub mod orchestrator; +pub(crate) mod page_size; pub mod providers; pub use client::{ActionExecutor, ComposioClient, ExecuteError, ExecuteResponse}; diff --git a/src/memory/sync/composio/orchestrator.rs b/src/memory/sync/composio/orchestrator.rs index 01d0d06..0eddde4 100644 --- a/src/memory/sync/composio/orchestrator.rs +++ b/src/memory/sync/composio/orchestrator.rs @@ -4,6 +4,7 @@ use async_trait::async_trait; use serde_json::Value; use super::client::ActionExecutor; +use super::page_size::{apply_page_size, is_payload_too_large, shrink_page_size}; use crate::memory::config::MemoryConfig; use crate::memory::sync::state::SyncState; use crate::memory::sync::traits::{ @@ -90,6 +91,16 @@ pub trait IncrementalSource: Send + Sync { ) -> anyhow::Result> { Ok(vec![SyncScope::flat()]) } + /// Name of the argument that caps how many items one page requests + /// (`max_results` for Gmail), when the action has one. + /// + /// Returning `Some` opts the source into the too-large-page retry: a page + /// the provider refuses purely for size is re-requested with the cap + /// halved, instead of failing the whole run. A source whose action has no + /// such knob returns `None` and keeps the previous behaviour. + fn page_size_arg_key(&self) -> Option<&'static str> { + None + } fn arguments( &self, scope: &SyncScope, @@ -218,45 +229,96 @@ async fn run_pages( .then(|| source.depth_floor(config, state)) .flatten(); + // Once a page proves too large for the provider, every later page of this + // run asks for the smaller size straight away rather than paying a rejected + // round-trip to rediscover the same limit. + let mut page_size_override: Option = None; + 'scopes: for scope in scopes { let mut page_token = None; let mut scope_newest_cursor: Option = None; let mut scope_failed = false; - for page_index in 0..source.max_pages().max(1) { + 'pages: for page_index in 0..source.max_pages().max(1) { if state.budget_exhausted() { more_pending = true; break 'scopes; } - let response = match executor - .execute( - source.action(), - source.arguments(scope, config, state, page_token.as_deref()), - Some(connection_id), - ) - .await - { - Ok(response) => response, - Err(error) if source.tolerate_scope_errors() => { - if let Some(execute_error) = error.downcast_ref::() - { - state.record_requests(execute_error.attempts); + let mut arguments = source.arguments(scope, config, state, page_token.as_deref()); + apply_page_size( + &mut arguments, + source.page_size_arg_key(), + page_size_override, + ); + // The size a shrink most recently *tried*. Promoted to the run's + // sticky override only once the provider accepts a page at it — + // before that it names a size that may itself be refused. + let mut attempted_page_size: Option = None; + let response = loop { + let response = match executor + .execute(source.action(), arguments.clone(), Some(connection_id)) + .await + { + Ok(response) => response, + Err(error) if source.tolerate_scope_errors() => { + if let Some(execute_error) = + error.downcast_ref::() + { + state.record_requests(execute_error.attempts); + } + tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] scope fetch failed; continuing"); + scope_failed = true; + break 'pages; } - tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] scope fetch failed; continuing"); - scope_failed = true; - break; + Err(error) => { + if let Some(execute_error) = + error.downcast_ref::() + { + state.record_requests(execute_error.attempts); + } + return Err(error); + } + }; + // A completed provider round-trip is billable even when its + // envelope reports failure. Transport failures return before + // this point. + state.record_action(response.attempts, response.cost_usd); + if response.successful { + // Sticky for the rest of the run, and set HERE rather than + // at the shrink: with several halvings (25 → 12 → 6) an + // assignment per attempt leaves the last *rejected* size in + // the override on every step but the final one, and is + // correct at the end only because that step happens to be + // the accepted one. Recording the accepted size makes the + // intent independent of the retry order. + if let Some(accepted) = attempted_page_size { + page_size_override = Some(accepted); + } + break response; } - Err(error) => { - if let Some(execute_error) = error.downcast_ref::() + // A page refused purely for its size is the one provider + // failure a *smaller request* can fix, so shrink and retry + // instead of failing the run. Without this a single oversized + // page stops the source dead until someone notices: on one live + // workspace Gmail sync sat broken for nine days that way. + if is_payload_too_large(response.error.as_deref()) { + if state.budget_exhausted() { + more_pending = true; + break 'scopes; + } + if let Some(reduced) = + shrink_page_size(&mut arguments, source.page_size_arg_key()) { - state.record_requests(execute_error.attempts); + attempted_page_size = Some(reduced); + tracing::warn!( + toolkit = source.toolkit(), + connection_id, + scope = %scope.label, + reduced_page_size = reduced, + "[sync:orchestrator] provider refused the page as too large; retrying with a smaller page" + ); + continue; } - return Err(error); } - }; - // A completed provider round-trip is billable even when its envelope - // reports failure. Transport failures return before this point. - state.record_action(response.attempts, response.cost_usd); - if !response.successful { let error = anyhow::anyhow!( "{} provider failure: {}", source.toolkit(), @@ -267,10 +329,10 @@ async fn run_pages( if source.tolerate_scope_errors() { tracing::warn!(toolkit = source.toolkit(), connection_id, scope = %scope.label, %error, "[sync:orchestrator] provider rejected scope; continuing"); scope_failed = true; - break; + break 'pages; } return Err(error); - } + }; let fetched = source.extract_page(&response.data, page_token.as_deref()); let mut reached_cursor_boundary = false; @@ -447,3 +509,7 @@ fn now_ms() -> u64 { .unwrap_or_default() .as_millis() as u64 } + +#[cfg(test)] +#[path = "orchestrator_tests.rs"] +mod tests; diff --git a/src/memory/sync/composio/orchestrator_tests.rs b/src/memory/sync/composio/orchestrator_tests.rs new file mode 100644 index 0000000..3962e67 --- /dev/null +++ b/src/memory/sync/composio/orchestrator_tests.rs @@ -0,0 +1,239 @@ +//! Tests for the too-large-page retry. + +use std::sync::{Arc, Mutex}; + +use serde_json::json; + +use super::*; +use crate::memory::sync::composio::client::ExecuteResponse; +use crate::memory::sync::state::SyncStateStore; +use crate::memory::sync::traits::{SkillDocSink, SkillDocument, SyncEvent, SyncEventSink}; + +/// Executor that mimics a provider with a response-size ceiling: it refuses any +/// page asking for more than `accepts` items and records every size it was +/// asked for. +struct SizeLimitedExecutor { + accepts: u64, + requested: Mutex>, +} + +impl SizeLimitedExecutor { + fn new(accepts: u64) -> Self { + Self { + accepts, + requested: Mutex::new(Vec::new()), + } + } + + fn requested_sizes(&self) -> Vec { + self.requested.lock().unwrap().clone() + } +} + +#[async_trait] +impl ActionExecutor for SizeLimitedExecutor { + async fn execute( + &self, + _action: &str, + arguments: Value, + _connection_id: Option<&str>, + ) -> anyhow::Result { + let requested = arguments + .get("max_results") + .and_then(Value::as_u64) + .unwrap_or(0); + self.requested.lock().unwrap().push(requested); + + let mut response: ExecuteResponse = serde_json::from_value(json!({})).unwrap(); + if requested > self.accepts { + response.successful = false; + response.error = Some( + "413 {\"error\":{\"message\":\"The tool response payload is too large.\",\ + \"code\":1613,\"slug\":\"Upstream_PayloadTooLarge\"}}" + .to_string(), + ); + return Ok(response); + } + response.successful = true; + response.data = json!({ + "messages": [{ "id": format!("m{requested}"), "date": "1700000000000" }], + }); + Ok(response) + } +} + +/// Minimal paged source: one page, one item, page size declared as +/// `max_results` unless `declare_page_size` is off. +struct StubSource { + declare_page_size: bool, + initial_page_size: u64, +} + +#[async_trait] +impl IncrementalSource for StubSource { + fn toolkit(&self) -> &'static str { + "stub" + } + fn action(&self) -> &'static str { + "STUB_FETCH" + } + fn max_pages(&self) -> usize { + 1 + } + fn page_size_arg_key(&self) -> Option<&'static str> { + self.declare_page_size.then_some("max_results") + } + fn arguments( + &self, + _scope: &SyncScope, + _config: &MemoryConfig, + _state: &SyncState, + _page: Option<&str>, + ) -> Value { + json!({ "max_results": self.initial_page_size }) + } + fn extract_page(&self, data: &Value, _page: Option<&str>) -> PageFetch { + PageFetch { + items: data + .get("messages") + .and_then(Value::as_array) + .cloned() + .unwrap_or_default(), + next: None, + } + } + fn dedup_key(&self, item: &Value) -> Option { + item.get("id").and_then(Value::as_str).map(str::to_string) + } + fn sort_cursor(&self, item: &Value) -> Option { + item.get("date").and_then(Value::as_str).map(str::to_string) + } + async fn document( + &self, + _scope: &SyncScope, + connection_id: &str, + item: SyncItem, + _executor: &dyn ActionExecutor, + _state: &mut SyncState, + ) -> anyhow::Result { + Ok(SkillDocument { + namespace_skill_id: "stub".into(), + connection_id: connection_id.into(), + document_id: item.dedup_key, + title: "stub".into(), + content: "stub".into(), + toolkit: "stub".into(), + metadata: Value::Null, + }) + } +} + +#[derive(Default)] +struct NoopHost(Mutex>); + +#[async_trait] +impl SkillDocSink for NoopHost { + async fn store(&self, _: SkillDocument) -> anyhow::Result<()> { + Ok(()) + } + async fn delete(&self, _: &str, _: &str) -> anyhow::Result<()> { + Ok(()) + } +} + +#[async_trait] +impl SyncEventSink for NoopHost { + async fn emit(&self, _: SyncEvent) -> anyhow::Result<()> { + Ok(()) + } +} + +#[async_trait] +impl SyncStateStore for NoopHost { + async fn get(&self, namespace: &str, key: &str) -> anyhow::Result> { + Ok(self + .0 + .lock() + .unwrap() + .get(&format!("{namespace}:{key}")) + .cloned()) + } + async fn set(&self, namespace: &str, key: &str, value: &Value) -> anyhow::Result<()> { + self.0 + .lock() + .unwrap() + .insert(format!("{namespace}:{key}"), value.clone()); + Ok(()) + } +} + +fn context() -> SyncContext { + let host = Arc::new(NoopHost::default()); + SyncContext { + events: host.clone(), + documents: host.clone(), + state: host, + local_documents: None, + external_sources: None, + summariser: None, + } +} + +#[tokio::test] +async fn an_oversized_page_is_halved_until_the_provider_accepts_it() { + // The provider takes 6 at a time; the source asks for 25. + let executor = SizeLimitedExecutor::new(6); + let source = StubSource { + declare_page_size: true, + initial_page_size: 25, + }; + + let outcome = run_incremental_sync( + &source, + &executor, + "conn-1", + &MemoryConfig::new("/tmp/unused-orchestrator-tests"), + &context(), + ) + .await + .expect("a too-large page must not fail the run"); + + assert_eq!( + outcome.records_ingested, 1, + "the page is ingested after the retry" + ); + assert_eq!( + executor.requested_sizes(), + vec![25, 12, 6], + "each rejection halves the request until it fits" + ); +} + +#[tokio::test] +async fn a_source_without_a_page_size_argument_still_fails_fast() { + // No `page_size_arg_key` — there is nothing to shrink, so the old + // behaviour (surface the provider failure) must be preserved rather than + // looping on a request that can never change. + let executor = SizeLimitedExecutor::new(6); + let source = StubSource { + declare_page_size: false, + initial_page_size: 25, + }; + + let error = run_incremental_sync( + &source, + &executor, + "conn-1", + &MemoryConfig::new("/tmp/unused-orchestrator-tests"), + &context(), + ) + .await + .expect_err("an unshrinkable too-large page is still a failure"); + + assert!(error.to_string().contains("provider failure"), "{error}"); + assert_eq!( + executor.requested_sizes(), + vec![25], + "no pointless retry of an identical request" + ); +} diff --git a/src/memory/sync/composio/page_size.rs b/src/memory/sync/composio/page_size.rs new file mode 100644 index 0000000..d970493 --- /dev/null +++ b/src/memory/sync/composio/page_size.rs @@ -0,0 +1,79 @@ +//! Page-size retry for a provider that refuses a page for its size. +//! +//! A page rejected purely because the response is too big is the one provider +//! failure a *smaller request* can fix. The orchestrator halves the page-size +//! argument and retries rather than failing the source, so a single oversized +//! page cannot stop a sync dead — on one live workspace a Gmail sync sat broken +//! for nine days that way. + +use serde_json::Value; + +/// Smallest page a shrink will ask for. One item is the point past which a +/// too-large response is about that single item, not the batch size. +pub(super) const MIN_PAGE_SIZE: u64 = 1; + +/// Whether the provider refused a page for its *size* rather than for anything +/// about the request's content — the only failure a smaller page can fix. +/// +/// Matched on the error text because that is all the envelope carries: Composio +/// reports it as HTTP 413 with a `Upstream_PayloadTooLarge` slug, and other +/// backends phrase it as "payload too large" / "response too large". +pub(super) fn is_payload_too_large(error: Option<&str>) -> bool { + error.is_some_and(|error| { + let lower = error.to_ascii_lowercase(); + lower.contains("payloadtoolarge") + || lower.contains("payload_too_large") + || mentions_status_413(&lower) + || (lower.contains("too large") + && (lower.contains("payload") || lower.contains("response"))) + }) +} + +/// Whether `text` names HTTP 413 as a status code. +/// +/// The digits have to stand alone. An unanchored `contains("413")` also matches +/// a message id, an amount, or a timestamp that merely contains those three +/// digits, and every such match costs a shrink-and-retry cycle before the real +/// error is finally surfaced — on a failure that a smaller page was never going +/// to fix. +fn mentions_status_413(lower: &str) -> bool { + lower.match_indices("413").any(|(at, _)| { + let before_is_digit = lower[..at] + .chars() + .next_back() + .is_some_and(|c| c.is_ascii_digit()); + let after_is_digit = lower[at + 3..] + .chars() + .next() + .is_some_and(|c| c.is_ascii_digit()); + !before_is_digit && !after_is_digit + }) +} + +/// Pin the page-size argument to `size`, if the source declared one. +pub(super) fn apply_page_size(arguments: &mut Value, key: Option<&str>, size: Option) { + if let (Some(key), Some(size)) = (key, size) { + if let Some(slot) = arguments.get_mut(key) { + *slot = Value::from(size); + } + } +} + +/// Halve the page-size argument in place, returning the new value. +/// +/// `None` means retrying is pointless — the source declares no page-size +/// argument, this request does not carry it, or it is already at the floor. +pub(super) fn shrink_page_size(arguments: &mut Value, key: Option<&str>) -> Option { + let key = key?; + let current = arguments.get(key)?.as_u64()?; + if current <= MIN_PAGE_SIZE { + return None; + } + let reduced = (current / 2).max(MIN_PAGE_SIZE); + arguments[key] = Value::from(reduced); + Some(reduced) +} + +#[cfg(test)] +#[path = "page_size_tests.rs"] +mod tests; diff --git a/src/memory/sync/composio/page_size_tests.rs b/src/memory/sync/composio/page_size_tests.rs new file mode 100644 index 0000000..6328c7b --- /dev/null +++ b/src/memory/sync/composio/page_size_tests.rs @@ -0,0 +1,51 @@ +//! Tests for the page-size retry helpers. + +use serde_json::json; + +use super::*; + +#[test] +fn payload_too_large_is_told_apart_from_other_provider_errors() { + assert!(is_payload_too_large(Some( + "413 {\"slug\":\"Upstream_PayloadTooLarge\"}" + ))); + assert!(is_payload_too_large(Some("Response too large for tool"))); + assert!(!is_payload_too_large(Some("rate limit exceeded"))); + assert!(!is_payload_too_large(Some("invalid grant"))); + assert!(!is_payload_too_large(None)); +} + +#[test] +fn shrinking_stops_at_the_floor() { + let mut arguments = json!({ "max_results": 3 }); + assert_eq!( + shrink_page_size(&mut arguments, Some("max_results")), + Some(1) + ); + assert_eq!(arguments["max_results"], json!(1)); + // At one item per page there is nothing left to halve. + assert_eq!(shrink_page_size(&mut arguments, Some("max_results")), None); + // A source that declares no key, or a request that lacks it, cannot shrink. + assert_eq!(shrink_page_size(&mut arguments, None), None); + assert_eq!( + shrink_page_size(&mut json!({ "other": 10 }), Some("max_results")), + None + ); +} + +/// The digits have to stand alone. Every false positive here costs a +/// shrink-and-retry cycle on a failure a smaller page was never going to fix, +/// and delays the real error reaching the caller. +#[test] +fn a_number_that_merely_contains_413_is_not_a_status_code() { + assert!(is_payload_too_large(Some("HTTP 413 Payload Too Large"))); + assert!(is_payload_too_large(Some("upstream returned 413."))); + assert!(is_payload_too_large(Some("(413)"))); + + assert!(!is_payload_too_large(Some("message id 4130 not found"))); + assert!(!is_payload_too_large(Some("amount 1413 exceeds the cap"))); + assert!(!is_payload_too_large(Some("thread 94137 is archived"))); + assert!(!is_payload_too_large(Some( + "at 1782891413 the token expired" + ))); +} diff --git a/src/memory/tree/store/summaries.rs b/src/memory/tree/store/summaries.rs index 57c1acd..9eae16e 100644 --- a/src/memory/tree/store/summaries.rs +++ b/src/memory/tree/store/summaries.rs @@ -13,7 +13,9 @@ use rusqlite::{params, OptionalExtension, Transaction}; use super::common::{decode_signature_blob, ms_to_utc, pack_embedding_blob}; use super::types::{SummaryNode, TreeKind}; -use crate::memory::chunks::{tree_active_signature, with_connection}; +use crate::memory::chunks::{ + signature_in_clause, signature_variants, tree_active_signature, with_connection, +}; use crate::memory::config::MemoryConfig; use crate::memory::score::embed::decode_optional_blob; use crate::memory::store::content::StagedSummary; @@ -281,18 +283,28 @@ pub fn set_summary_embedding( Ok(1) } -/// Fetch a summary embedding for exactly one signature. +/// Fetch a summary embedding for one signature, under any of its spellings +/// (see `chunks::signature_variants`). pub fn get_summary_embedding_for_signature( config: &MemoryConfig, summary_id: &str, model_signature: &str, ) -> Result>> { + let variants = signature_variants(model_signature); with_connection(config, |conn| { + let mut bound: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(variants.len() + 1); + bound.push(&summary_id as &dyn rusqlite::ToSql); + for variant in &variants { + bound.push(variant as &dyn rusqlite::ToSql); + } let row: Option<(Option>, i64)> = conn .query_row( - "SELECT vector, dim FROM mem_tree_summary_embeddings - WHERE summary_id = ?1 AND model_signature = ?2", - params![summary_id, model_signature], + &format!( + "SELECT vector, dim FROM mem_tree_summary_embeddings + WHERE summary_id = ?1 AND model_signature {}", + signature_in_clause(variants.len(), 2) + ), + bound.as_slice(), |r| Ok((Some(r.get(0)?), r.get(1)?)), ) .optional()?; @@ -323,6 +335,7 @@ pub fn get_summary_embeddings_for_signature_batch( if summary_ids.is_empty() { return Ok(HashMap::new()); } + let variants = signature_variants(model_signature); with_connection(config, |conn| { let mut out: HashMap> = HashMap::with_capacity(summary_ids.len()); for window in summary_ids.chunks(MAX_EMBEDDING_BATCH) { @@ -331,15 +344,18 @@ pub fn get_summary_embeddings_for_signature_batch( .join(","); let sql = format!( "SELECT summary_id, vector, dim FROM mem_tree_summary_embeddings - WHERE summary_id IN ({placeholders}) AND model_signature = ?{sig_idx}", - sig_idx = window.len() + 1, + WHERE summary_id IN ({placeholders}) AND model_signature {sig_clause}", + sig_clause = signature_in_clause(variants.len(), window.len() + 1), ); let mut stmt = conn.prepare(&sql)?; - let mut bound: Vec<&dyn rusqlite::ToSql> = Vec::with_capacity(window.len() + 1); + let mut bound: Vec<&dyn rusqlite::ToSql> = + Vec::with_capacity(window.len() + variants.len()); for id in window { bound.push(id as &dyn rusqlite::ToSql); } - bound.push(&model_signature as &dyn rusqlite::ToSql); + for variant in &variants { + bound.push(variant as &dyn rusqlite::ToSql); + } let rows = stmt.query_map(bound.as_slice(), |row| { Ok(( row.get::<_, String>(0)?,