perf(memory_tree): concurrently seal document-subtree sibling batches (#3416)

Co-authored-by: M3gA-Mind <megamind@mahadao.com>

src/openhuman/memory_tree/tree/bucket_seal.rs

@@ -38,6 +38,7 @@ use std::sync::Arc;
 
 use anyhow::{Context, Result};
 use chrono::{DateTime, Utc};
+use futures::stream::{StreamExt, TryStreamExt};
 use rusqlite::Transaction;
 
 use crate::openhuman::config::Config;
@@ -1012,6 +1013,17 @@ pub const MERGE_LEVEL_BASE: u32 = 1_000;
 /// children. Independent of [`SUMMARY_FANOUT`] (which gates the merge tier).
 const DOC_SUBTREE_MAX_FANIN: usize = 32;
 
+/// Bound on how many sibling batches within a single document-subtree level
+/// are sealed concurrently. Each [`seal_explicit_children`] call is one
+/// independent summarise + embed round-trip pair, so overlapping them with a
+/// small pool collapses the serial `N × (LLM + embed RTT)` wall-time toward
+/// `N/CONCURRENCY × RTT` on the ingest path. Kept low: the SQLite write tx
+/// serialises on a single writer anyway, and networked summarise/embed
+/// providers (Voyage, cloud LLMs) rate-limit — so a large pool buys nothing
+/// and risks 429s. `buffered` (ordered) preserves the serial cascade's batch
+/// order, so the resulting tree shape is identical.
+const DOC_SUBTREE_SEAL_CONCURRENCY: usize = 8;
+
 /// Build (or re-build, for a new version) one document's subtree and merge
 /// its doc-root into the connection tree.
 ///
@@ -1067,21 +1079,61 @@ pub async fn seal_document_subtree(
             batch_by_count(&current_ids, DOC_SUBTREE_MAX_FANIN)
         };
 
-        let mut next_ids: Vec<String> = Vec::with_capacity(batches.len());
-        for batch in &batches {
-            let node = seal_explicit_children(
-                config,
-                tree,
+        // Within one level the batches are disjoint child sets, and
+        // `seal_explicit_children` touches no shared `(tree, level)` buffer
+        // and never advances the tree root (see its docstring) — so the
+        // sibling seals are independent and their summarise + embed
+        // round-trips can overlap. Run them with bounded concurrency.
+        //
+        // `buffered` (ordered), NOT `buffer_unordered`: the next level groups
+        // `next_ids` positionally via `batch_by_count`, so the output order
+        // must match the serial cascade to keep the tree shape identical.
+        // `try_collect` short-circuits on the first error, preserving the
+        // old `?`-on-first-failure behaviour. Levels stay sequential — level
+        // N+1 consumes level N's output — so only siblings overlap, not the
+        // cross-level fan-in. The blocking SQLite write inside each seal
+        // still serialises (single writer), which is fine: the win is the
+        // overlapped network RTTs, not the DB writes.
+        // Build the per-batch futures eagerly into a `Vec` (they don't run until
+        // polled) rather than lazily inside a `stream::iter(...).map(closure)`:
+        // the lazy form makes rustc try to infer a higher-ranked `FnOnce` over
+        // the borrowed `batch` lifetime, which fails ("implementation of
+        // `FnOnce` is not general enough") once this async fn is spawned. Eager
+        // collection ties each future to this stack frame's concrete lifetime.
+        if batches.len() > 1 {
+            log::debug!(
+                "[tree::bucket_seal] doc-subtree concurrent seal tree_id={} doc_id_hash={} level={} batches={} concurrency={}",
+                tree.id,
+                redact(doc_id),
                 current_level,
-                batch,
-                Some(doc_id),
-                version_ms,
-                strategy,
-            )
-            .await?;
-            next_ids.push(node.id.clone());
-            doc_root = Some(node);
+                batches.len(),
+                DOC_SUBTREE_SEAL_CONCURRENCY
+            );
         }
+        let batch_futures: Vec<_> = batches
+            .iter()
+            .map(|batch| {
+                seal_explicit_children(
+                    config,
+                    tree,
+                    current_level,
+                    batch,
+                    Some(doc_id),
+                    version_ms,
+                    strategy,
+                )
+            })
+            .collect();
+        let nodes: Vec<SummaryNode> = futures::stream::iter(batch_futures)
+            .buffered(DOC_SUBTREE_SEAL_CONCURRENCY)
+            .try_collect()
+            .await?;
+
+        let next_ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();
+        // The terminal level produces exactly one node (the loop breaks once
+        // `current_ids.len() <= 1`); preserving batch order means the last
+        // node is the doc-root, identical to the serial assignment.
+        doc_root = nodes.into_iter().next_back();
 
         current_level += 1;
         current_ids = next_ids;

src/openhuman/memory_tree/tree/bucket_seal_tests.rs

@@ -351,6 +351,121 @@ async fn seal_document_subtree_single_chunk_is_verbatim_passthrough_no_llm() {
     );
 }
 
+/// Multi-batch document subtree: a doc whose chunks span several level-0
+/// batches must seal its siblings **concurrently** (overlapping the
+/// summarise + embed round-trips) while still producing an identical
+/// doc-root. A `ChatProvider` probe records the peak number of in-flight
+/// summarise calls; with four independent level-0 batches the peak must
+/// exceed 1 (proving overlap) yet stay within `DOC_SUBTREE_SEAL_CONCURRENCY`
+/// (proving the bound is respected). This is the regression guard for the
+/// serial → bounded-concurrency change in `seal_document_subtree`.
+#[tokio::test]
+async fn seal_document_subtree_seals_sibling_batches_concurrently() {
+    use crate::openhuman::memory::chat::ChatPrompt;
+    use crate::openhuman::memory_store::chunks::store::upsert_chunks;
+    use crate::openhuman::memory_store::chunks::types::{
+        chunk_id, Chunk, Metadata, SourceKind, SourceRef,
+    };
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    /// Records concurrent in-flight `chat_for_json` calls and the peak; a
+    /// short sleep forces would-be-concurrent calls to actually overlap so
+    /// the peak reflects real scheduling, not luck.
+    struct ConcurrencyProbe {
+        in_flight: AtomicUsize,
+        peak: AtomicUsize,
+        calls: AtomicUsize,
+    }
+    #[async_trait::async_trait]
+    impl ChatProvider for ConcurrencyProbe {
+        fn name(&self) -> &str {
+            "test:concurrency-probe"
+        }
+        async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result<String> {
+            let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1;
+            self.peak.fetch_max(now, Ordering::SeqCst);
+            self.calls.fetch_add(1, Ordering::SeqCst);
+            tokio::time::sleep(std::time::Duration::from_millis(50)).await;
+            self.in_flight.fetch_sub(1, Ordering::SeqCst);
+            Ok("doc batch summary".to_string())
+        }
+    }
+
+    let (_tmp, cfg) = test_config();
+    let tree = get_or_create_source_tree(&cfg, "notion:conn1").unwrap();
+    let doc_id = "notion:conn1:pageM";
+
+    // Four chunks, each ~60% of the input budget, so the greedy token-budget
+    // batcher puts each in its OWN level-0 batch → four independent
+    // summarise calls that should overlap under bounded concurrency.
+    let per_chunk = INPUT_TOKEN_BUDGET * 6 / 10;
+    let ts = Utc::now();
+    let mut ids = Vec::new();
+    let mut chunks = Vec::new();
+    for seq in 0..4u32 {
+        let content = format!("substantive document body number {seq}");
+        let c = Chunk {
+            id: chunk_id(SourceKind::Document, doc_id, seq, &content),
+            content,
+            metadata: Metadata {
+                source_kind: SourceKind::Document,
+                source_id: doc_id.to_string(),
+                owner: "notion:conn1".into(),
+                timestamp: ts,
+                time_range: (ts, ts),
+                tags: vec!["notion".into()],
+                source_ref: Some(SourceRef::new("notion://page/pageM")),
+                path_scope: Some("notion:conn1".into()),
+            },
+            token_count: per_chunk,
+            seq_in_source: seq,
+            created_at: ts,
+            partial_message: false,
+        };
+        ids.push(c.id.clone());
+        chunks.push(c);
+    }
+    upsert_chunks(&cfg, &chunks).unwrap();
+    stage_test_chunks(&cfg, &chunks);
+
+    let probe = Arc::new(ConcurrencyProbe {
+        in_flight: AtomicUsize::new(0),
+        peak: AtomicUsize::new(0),
+        calls: AtomicUsize::new(0),
+    });
+    let provider: Arc<dyn ChatProvider> = probe.clone();
+
+    let doc_root_id = test_override::with_provider(provider, async {
+        seal_document_subtree(&cfg, &tree, doc_id, Some(7), &ids, &LabelStrategy::Empty)
+            .await
+            .unwrap()
+    })
+    .await;
+
+    // Correctness: the four chunks roll up under a single doc-root tagged
+    // with the document id + version — tree shape unchanged by batching.
+    let root = store::get_summary(&cfg, &doc_root_id).unwrap().unwrap();
+    assert_eq!(root.doc_id.as_deref(), Some(doc_id));
+    assert_eq!(root.version_ms, Some(7));
+    assert_eq!(
+        root.child_ids.len(),
+        4,
+        "all four level-0 siblings must fan into the doc-root"
+    );
+
+    // Concurrency: at least two level-0 siblings overlapped, and the pool
+    // bound was respected.
+    let peak = probe.peak.load(Ordering::SeqCst);
+    assert!(
+        peak >= 2,
+        "expected overlapping sibling seals (peak in-flight = {peak})"
+    );
+    assert!(
+        peak <= DOC_SUBTREE_SEAL_CONCURRENCY,
+        "concurrency must stay within the pool bound (peak = {peak})"
+    );
+}
+
 #[tokio::test]
 async fn crossing_budget_triggers_seal() {
     use crate::openhuman::memory_store::chunks::store::upsert_chunks;