From a892cd420677d2bcee6a248eb75730e3759c0bdb Mon Sep 17 00:00:00 2001 From: mysma-9403 <64923976+mysma-9403@users.noreply.github.com> Date: Tue, 16 Jun 2026 02:15:33 +0200 Subject: [PATCH] perf(memory-sources): batch composio reconcile upserts into one config save (#3432) Co-authored-by: Steven Enamakel --- src/openhuman/memory_sources/reconcile.rs | 88 ++++++++--- src/openhuman/memory_sources/registry.rs | 174 ++++++++++++++++++---- 2 files changed, 214 insertions(+), 48 deletions(-) diff --git a/src/openhuman/memory_sources/reconcile.rs b/src/openhuman/memory_sources/reconcile.rs index e746d5f83..4c2769433 100644 --- a/src/openhuman/memory_sources/reconcile.rs +++ b/src/openhuman/memory_sources/reconcile.rs @@ -55,30 +55,22 @@ pub async fn ensure_composio_sources() -> Option> { } }; - let mut upserted = 0u32; - for target in &targets { - // Use a title-cased toolkit name plus the truncated connection id - // so distinct Gmail accounts don't all show as "Gmail connection". - let label = format!( - "{} · {}", - title_case(&target.toolkit), - short_id(&target.connection_id) - ); - match registry::upsert_composio_source(&target.toolkit, &target.connection_id, &label).await - { - Ok(_) => { - upserted += 1; - } - Err(e) => { - tracing::warn!( - toolkit = %target.toolkit, - connection_id = %target.connection_id, - error = %e, - "[memory_sources:reconcile] upsert failed" - ); - } + // Build the upsert targets up front, then apply them with a single config + // load + save via the batch path. The per-call upsert does its own + // load-modify-save, so the old loop cost 2N config round-trips for N + // connections; batching collapses that to 2. + let upsert_targets = build_upsert_targets(&targets); + let upserted = match registry::upsert_composio_sources_batch(&upsert_targets).await { + Ok(n) => n, + Err(e) => { + tracing::warn!( + targets = targets.len(), + error = %e, + "[memory_sources:reconcile] batch upsert failed" + ); + 0 } - } + }; if !targets.is_empty() { tracing::info!( @@ -104,6 +96,26 @@ pub async fn ensure_composio_sources() -> Option> { Some(targets.iter().map(|t| t.connection_id.clone()).collect()) } +/// Build the `(toolkit, connection_id, label)` upsert targets for a batch +/// reconcile from the scanned Composio sync targets. +/// +/// The label is a title-cased toolkit name plus the truncated connection id so +/// distinct accounts of the same toolkit (e.g. two Gmail logins) don't all show +/// as "Gmail connection". Pure (no I/O) so it can be unit-tested directly. +fn build_upsert_targets(targets: &[composio::SyncTarget]) -> Vec { + targets + .iter() + .map(|target| { + let label = format!( + "{} · {}", + title_case(&target.toolkit), + short_id(&target.connection_id) + ); + (target.toolkit.clone(), target.connection_id.clone(), label) + }) + .collect() +} + /// Apply conservative default caps in-place to every cap-less source. /// /// For a Composio source with no `max_items`/`sync_depth_days`, writes the @@ -317,6 +329,36 @@ mod tests { } } + fn sync_target(toolkit: &str, connection_id: &str) -> composio::SyncTarget { + composio::SyncTarget { + toolkit: toolkit.to_string(), + connection_id: connection_id.to_string(), + } + } + + #[test] + fn build_upsert_targets_formats_label_and_preserves_order() { + let targets = vec![ + sync_target("gmail", "ca_WaktIDFlZwXO"), + sync_target("slack", "short"), + ]; + let out = build_upsert_targets(&targets); + assert_eq!(out.len(), 2); + // (toolkit, connection_id, label) — toolkit/connection_id carried through verbatim. + assert_eq!(out[0].0, "gmail"); + assert_eq!(out[0].1, "ca_WaktIDFlZwXO"); + assert_eq!(out[0].2, "Gmail · IDFlZwXO"); + assert_eq!(out[1].0, "slack"); + assert_eq!(out[1].1, "short"); + assert_eq!(out[1].2, "Slack · short"); + } + + #[test] + fn build_upsert_targets_empty_is_empty() { + let out = build_upsert_targets(&[]); + assert!(out.is_empty()); + } + #[test] fn short_id_truncates_ascii() { assert_eq!(short_id("ca_WaktIDFlZwXO"), "IDFlZwXO"); diff --git a/src/openhuman/memory_sources/registry.rs b/src/openhuman/memory_sources/registry.rs index a2f480661..0fbd4e368 100644 --- a/src/openhuman/memory_sources/registry.rs +++ b/src/openhuman/memory_sources/registry.rs @@ -218,32 +218,26 @@ pub async fn remove_composio_source_by_connection_id(connection_id: &str) -> Res Ok(removed) } -/// Upsert a composio source — used by the auto-registration path. -/// If a source with the same `connection_id` already exists, updates -/// the label; otherwise inserts a new entry. -pub async fn upsert_composio_source( +/// Apply a single composio upsert to an in-memory source list. +/// +/// If a source with the same `connection_id` already exists, updates its label +/// and returns a clone of the updated entry with `was_insert = false`; +/// otherwise pushes a new entry (with conservative per-toolkit defaults) and +/// returns it with `was_insert = true`. +/// +/// Pure (no I/O) so the single-call path, the batch path, and unit tests all +/// share one find-or-push predicate and cannot drift. +fn upsert_composio_entry_in_place( + sources: &mut Vec, toolkit: &str, connection_id: &str, label: &str, -) -> Result { - let _guard = memory_sources_write_guard().await; - let mut config = config_rpc::load_config_with_timeout().await?; - - if let Some(existing) = config.memory_sources.iter_mut().find(|s| { +) -> (MemorySourceEntry, bool) { + if let Some(existing) = sources.iter_mut().find(|s| { s.kind == SourceKind::Composio && s.connection_id.as_deref() == Some(connection_id) }) { existing.label = label.to_string(); - let updated = existing.clone(); - config - .save() - .await - .map_err(|e| format!("failed to save config: {e:#}"))?; - tracing::debug!( - connection_id = %connection_id, - toolkit = %toolkit, - "[memory_sources] upserted composio source (update)" - ); - return Ok(updated); + return (existing.clone(), false); } let (default_max_items, default_sync_depth_days) = memory_sync_defaults_for_toolkit(toolkit); @@ -277,19 +271,100 @@ pub async fn upsert_composio_source( max_cost_per_sync_usd: None, sync_depth_days: default_sync_depth_days, }; - config.memory_sources.push(entry.clone()); + sources.push(entry.clone()); + (entry, true) +} + +/// Upsert a composio source — used by the auto-registration path. +/// If a source with the same `connection_id` already exists, updates +/// the label; otherwise inserts a new entry. +pub async fn upsert_composio_source( + toolkit: &str, + connection_id: &str, + label: &str, +) -> Result { + let _guard = memory_sources_write_guard().await; + let mut config = config_rpc::load_config_with_timeout().await?; + + let (entry, was_insert) = + upsert_composio_entry_in_place(&mut config.memory_sources, toolkit, connection_id, label); + + config + .save() + .await + .map_err(|e| format!("failed to save config: {e:#}"))?; + + if was_insert { + tracing::info!( + connection_id = %connection_id, + toolkit = %toolkit, + "[memory_sources] upserted composio source (insert)" + ); + } else { + tracing::debug!( + connection_id = %connection_id, + toolkit = %toolkit, + "[memory_sources] upserted composio source (update)" + ); + } + + Ok(entry) +} + +/// Target for a batch composio upsert: `(toolkit, connection_id, label)`. +pub type ComposioUpsertTarget = (String, String, String); + +/// Batch-upsert many composio sources with a **single** config load + save. +/// +/// The per-call [`upsert_composio_source`] does its own load-modify-save, so +/// calling it in a loop costs `2N` config I/O round-trips and — worse — cannot +/// be parallelised safely (concurrent calls each load the same snapshot and +/// their saves clobber each other). This batch path loads the config once, +/// applies every upsert in-memory via the shared [`upsert_composio_entry_in_place`] +/// predicate, and saves once. Returns the number of targets applied. +/// +/// Targets are applied in order; a later target sharing a `connection_id` with +/// an earlier one updates the same (already inserted/updated) entry's label, +/// matching the sequential single-call semantics. +pub async fn upsert_composio_sources_batch( + targets: &[ComposioUpsertTarget], +) -> Result { + if targets.is_empty() { + return Ok(0); + } + + let _guard = memory_sources_write_guard().await; + let mut config = config_rpc::load_config_with_timeout().await?; + + let mut inserted = 0u32; + let mut updated = 0u32; + for (toolkit, connection_id, label) in targets { + let (_entry, was_insert) = upsert_composio_entry_in_place( + &mut config.memory_sources, + toolkit, + connection_id, + label, + ); + if was_insert { + inserted += 1; + } else { + updated += 1; + } + } + config .save() .await .map_err(|e| format!("failed to save config: {e:#}"))?; tracing::info!( - connection_id = %connection_id, - toolkit = %toolkit, - "[memory_sources] upserted composio source (insert)" + targets = targets.len(), + inserted, + updated, + "[memory_sources] batch-upserted composio sources (single save)" ); - Ok(entry) + Ok(inserted + updated) } /// Partial update payload for a source entry. @@ -428,6 +503,55 @@ mod tests { assert_eq!(memory_sync_defaults_for_toolkit(""), (Some(30), Some(14))); } + #[test] + fn in_place_upsert_inserts_new_entry_with_toolkit_defaults() { + let mut sources: Vec = vec![]; + let (entry, was_insert) = + upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "Gmail · conn_a"); + assert!(was_insert, "first upsert for a connection must insert"); + assert_eq!(sources.len(), 1); + assert_eq!(entry.kind, SourceKind::Composio); + assert_eq!(entry.connection_id.as_deref(), Some("conn_a")); + assert_eq!(entry.toolkit.as_deref(), Some("gmail")); + assert_eq!(entry.label, "Gmail · conn_a"); + assert!(entry.enabled); + // Conservative gmail defaults are applied on insert. + assert_eq!(entry.max_items, Some(100)); + assert_eq!(entry.sync_depth_days, Some(30)); + } + + #[test] + fn in_place_upsert_updates_label_only_for_existing_connection() { + let mut sources: Vec = vec![]; + upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "old label"); + // User-customised a cap after the initial insert. + sources[0].max_items = Some(7); + + let (entry, was_insert) = + upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "new label"); + assert!(!was_insert, "second upsert for same connection must update"); + assert_eq!(sources.len(), 1, "must not duplicate the entry"); + assert_eq!(entry.label, "new label"); + assert_eq!( + entry.max_items, + Some(7), + "update must not clobber user-set caps" + ); + } + + #[test] + fn in_place_upsert_keeps_distinct_connections_separate() { + let mut sources: Vec = vec![]; + upsert_composio_entry_in_place(&mut sources, "gmail", "conn_a", "A"); + upsert_composio_entry_in_place(&mut sources, "gmail", "conn_b", "B"); + assert_eq!(sources.len(), 2); + let ids: Vec<_> = sources + .iter() + .filter_map(|s| s.connection_id.as_deref()) + .collect(); + assert!(ids.contains(&"conn_a") && ids.contains(&"conn_b")); + } + #[test] fn memory_source_patch_deserializes_partial() { let json = serde_json::json!({ "label": "New label", "enabled": false });