perf(memory-sources): batch composio reconcile upserts into one config save (#3432)

Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>

src/openhuman/memory_sources/reconcile.rs

@@ -55,30 +55,22 @@ pub async fn ensure_composio_sources() -> Option<HashSet<String>> {
         }
     };
 
-    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<HashSet<String>> {
     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<registry::ComposioUpsertTarget> {
+    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");

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<MemorySourceEntry>,
     toolkit: &str,
     connection_id: &str,
     label: &str,
-) -> Result<MemorySourceEntry, String> {
-    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<MemorySourceEntry, String> {
+    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<u32, String> {
+    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<MemorySourceEntry> = 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<MemorySourceEntry> = 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<MemorySourceEntry> = 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 });