perf(memory_tree): batch summary fetch in hydrate_summary_inputs (#2948)

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

src/openhuman/memory_store/trees/store.rs

@@ -16,6 +16,8 @@
 //! writes populate it via [`insert_summary_tx`]; reads decode it when
 //! present.
 
+use std::collections::HashMap;
+
 use anyhow::{Context, Result};
 use chrono::{DateTime, TimeZone, Utc};
 use rusqlite::{params, Connection, OptionalExtension, Transaction};
@@ -469,6 +471,64 @@ pub fn get_summary(config: &Config, id: &str) -> Result<Option<SummaryNode>> {
     })
 }
 
+/// Defensive upper bound on the number of `?` placeholders per batched
+/// `SELECT … WHERE id IN (?,?,…)` query. SQLite's compile-time
+/// `SQLITE_MAX_VARIABLE_NUMBER` has been ≥ 32 766 since 3.32 — 500 leaves
+/// a ~65× safety margin. The current call-site (`hydrate_summary_inputs`)
+/// passes at most a single seal's fan-in (typically 5–20 ids), so the
+/// loop runs exactly once. The window exists so future callers with
+/// larger id slices do not blow up against a host with a lower
+/// compile-time SQLite cap. No volume reduction: all input ids in → all
+/// matching rows out; the merged `HashMap` is byte-identical to one
+/// giant query.
+const MAX_FETCH_BATCH: usize = 500;
+
+/// Fetch many summaries by id in a single SQL round-trip per
+/// [`MAX_FETCH_BATCH`] window. Replaces the per-id `get_summary` loop
+/// inside hot paths like `hydrate_summary_inputs` (sealing L≥1 levels)
+/// where N can grow to the seal fan-in and the loop fires on every seal
+/// during ingest. Soft-deleted rows are returned with `deleted = true`
+/// just like [`get_summary`]; missing ids are silently absent from the
+/// map so callers can preserve the existing
+/// "row missing → skip with warn" contract.
+pub fn get_summaries_batch(
+    config: &Config,
+    summary_ids: &[String],
+) -> Result<HashMap<String, SummaryNode>> {
+    if summary_ids.is_empty() {
+        return Ok(HashMap::new());
+    }
+    with_connection(config, |conn| {
+        let mut out: HashMap<String, SummaryNode> = HashMap::with_capacity(summary_ids.len());
+        for window in summary_ids.chunks(MAX_FETCH_BATCH) {
+            let placeholders = (1..=window.len())
+                .map(|i| format!("?{i}"))
+                .collect::<Vec<_>>()
+                .join(",");
+            let sql = format!(
+                "SELECT id, tree_id, tree_kind, level, parent_id,
+                        child_ids_json, content, token_count,
+                        entities_json, topics_json,
+                        time_range_start_ms, time_range_end_ms,
+                        score, sealed_at_ms, deleted, embedding
+                   FROM mem_tree_summaries
+                  WHERE id IN ({placeholders})"
+            );
+            let mut stmt = conn.prepare(&sql)?;
+            let params: Vec<&dyn rusqlite::ToSql> =
+                window.iter().map(|s| s as &dyn rusqlite::ToSql).collect();
+            let rows = stmt
+                .query_map(params.as_slice(), row_to_summary)?
+                .collect::<rusqlite::Result<Vec<_>>>()
+                .context("Failed to collect summaries batch")?;
+            for s in rows {
+                out.insert(s.id.clone(), s);
+            }
+        }
+        Ok(out)
+    })
+}
+
 /// List sealed summaries for a tree at a given level, ordered by
 /// `sealed_at` ascending. Skips tombstoned rows.
 pub fn list_summaries_at_level(

src/openhuman/memory_store/trees/store_tests.rs

@@ -287,3 +287,63 @@ fn list_stale_buffers_orders_by_age() {
     assert_eq!(only_oldest[0].level, 0);
     assert_eq!(only_oldest[0].tree_id, "tree-1");
 }
+
+// ── get_summaries_batch ────────────────────────────────────────────────
+//
+// Same shape as `chunks::store::get_chunks_batch` /
+// `score::store::get_scores_batch`: present ids decode through the same
+// `row_to_summary` path as the per-id `get_summary` and land in a
+// `HashMap` keyed by id; missing ids are silently absent so the
+// `hydrate_summary_inputs` "missing row → warn + skip" contract keeps
+// working without an extra Ok(None) sentinel.
+
+#[test]
+fn get_summaries_batch_returns_present_ids_in_map() {
+    let (_tmp, cfg) = test_config();
+    insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap();
+    let a = sample_summary("sum-a", "tree-1", 1);
+    let b = sample_summary("sum-b", "tree-1", 1);
+    with_connection(&cfg, |conn| {
+        let tx = conn.unchecked_transaction()?;
+        insert_summary_tx(&tx, &a, None, "test")?;
+        insert_summary_tx(&tx, &b, None, "test")?;
+        tx.commit()?;
+        Ok(())
+    })
+    .unwrap();
+
+    let ids = vec!["sum-a".to_string(), "sum-b".to_string()];
+    let map = get_summaries_batch(&cfg, &ids).unwrap();
+    assert_eq!(map.len(), 2);
+    // Each decoded row must match the per-id `get_summary` path bit-for-bit
+    // — same `row_to_summary` decoder under the hood, so the structs are
+    // equal including the deserialised JSON columns.
+    assert_eq!(map.get("sum-a").unwrap(), &a);
+    assert_eq!(map.get("sum-b").unwrap(), &b);
+}
+
+#[test]
+fn get_summaries_batch_empty_input_and_missing_ids() {
+    // Empty input: empty map (no SQL issued).
+    let (_tmp, cfg) = test_config();
+    let empty = get_summaries_batch(&cfg, &[]).unwrap();
+    assert!(empty.is_empty());
+
+    // Missing ids: silently absent so `hydrate_summary_inputs` can warn
+    // + skip without an extra `Ok(None)` sentinel per id.
+    insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap();
+    let a = sample_summary("sum-a", "tree-1", 1);
+    with_connection(&cfg, |conn| {
+        let tx = conn.unchecked_transaction()?;
+        insert_summary_tx(&tx, &a, None, "test")?;
+        tx.commit()?;
+        Ok(())
+    })
+    .unwrap();
+
+    let ids = vec!["sum-a".to_string(), "ghost:no-such".to_string()];
+    let map = get_summaries_batch(&cfg, &ids).unwrap();
+    assert_eq!(map.len(), 1);
+    assert_eq!(map.get("sum-a").unwrap(), &a);
+    assert!(map.get("ghost:no-such").is_none());
+}

src/openhuman/memory_tree/tree/bucket_seal.rs

@@ -821,17 +821,25 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result<Vec<Summ
 
 fn hydrate_summary_inputs(config: &Config, summary_ids: &[String]) -> Result<Vec<SummaryInput>> {
     use crate::openhuman::memory_store::content::read as content_read;
+    use crate::openhuman::memory_store::trees::store::get_summaries_batch;
+
+    // One batched `SELECT … WHERE id IN (?,?,…)` instead of N per-id
+    // `get_summary` round-trips. Body reads stay per-id because each
+    // summary's full markdown lives in its own on-disk file — batching
+    // there would mean concurrent file opens, not a single round-trip.
+    // Walking the caller's `summary_ids` slice (not the map) preserves
+    // input order, matching the previous per-id loop's semantics; the
+    // map lookup keys by id so the order of `HashMap`'s iteration is
+    // irrelevant.
+    let node_by_id = get_summaries_batch(config, summary_ids)?;
 
     let mut out: Vec<SummaryInput> = Vec::with_capacity(summary_ids.len());
     for id in summary_ids {
-        let node = match store::get_summary(config, id)? {
-            Some(n) => n,
-            None => {
-                log::warn!(
-                    "[tree::bucket_seal] hydrate_summary_inputs: missing summary {id} — skipping"
-                );
-                continue;
-            }
+        let Some(node) = node_by_id.get(id) else {
+            log::warn!(
+                "[tree::bucket_seal] hydrate_summary_inputs: missing summary {id} — skipping"
+            );
+            continue;
         };
         // Read the full body from disk — `node.content` is a ≤500-char preview
         // after the MD-on-disk migration. Higher-level seals (L2+) summarise

src/openhuman/memory_tree/tree/bucket_seal_tests.rs

@@ -686,3 +686,151 @@ fn scope_slug_non_gmail_uses_full_scope() {
         "slugifying the full 'gmail:...' scope must differ from the participants-only slug"
     );
 }
+
+/// `hydrate_summary_inputs` was rewritten to do one batched
+/// `get_summaries_batch` SELECT instead of N per-id `get_summary`
+/// round-trips. This test pins three behavioural invariants the per-id
+/// loop used to give us for free, and which the HashMap-walk now has to
+/// reproduce:
+///
+/// 1. **Input order preservation.** We iterate the caller's
+///    `summary_ids` slice (not the HashMap) so the `SummaryInput`s come
+///    out in the order the caller asked for, even though `HashMap`
+///    iteration is not insertion-ordered.
+/// 2. **Per-id field propagation by id, not by index.** Distinct field
+///    values per summary (content, token_count, score) prove the
+///    map.get() is keyed by id — not by enumerate().
+/// 3. **Missing id → silent skip + warn.** Mirrors the per-id
+///    `Ok(None)` → continue contract; the request does not error out.
+#[tokio::test]
+async fn hydrate_summary_inputs_batch_preserves_order_and_skips_missing_ids() {
+    use crate::openhuman::memory_store::content::atomic::stage_summary;
+    use crate::openhuman::memory_store::content::SummaryComposeInput;
+    use crate::openhuman::memory_store::content::SummaryTreeKind;
+    use crate::openhuman::memory_store::trees::store::insert_tree;
+    use crate::openhuman::memory_store::trees::types::{SummaryNode, Tree, TreeKind, TreeStatus};
+    use crate::openhuman::memory_tree::tree::store::insert_summary_tx;
+    use chrono::TimeZone;
+
+    let (_tmp, cfg) = test_config();
+    let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
+
+    let tree = Tree {
+        id: "tree-hydrate".into(),
+        kind: TreeKind::Source,
+        scope: "slack:#eng".into(),
+        root_id: None,
+        max_level: 0,
+        status: TreeStatus::Active,
+        created_at: ts,
+        last_sealed_at: None,
+    };
+    insert_tree(&cfg, &tree).unwrap();
+
+    // Two summaries with distinct content / token_count / score so we
+    // can prove the per-id HashMap lookup keys by id and not by
+    // enumerate() over `summary_ids`.
+    let sum_a = SummaryNode {
+        id: "sum-a".into(),
+        tree_id: tree.id.clone(),
+        tree_kind: TreeKind::Source,
+        level: 1,
+        parent_id: None,
+        child_ids: vec!["leaf-a".into()],
+        content: "BODY-A".into(),
+        token_count: 11,
+        entities: vec!["entity:alice".into()],
+        topics: vec!["#a".into()],
+        time_range_start: ts,
+        time_range_end: ts,
+        score: 0.11,
+        sealed_at: ts,
+        deleted: false,
+        embedding: None,
+    };
+    let sum_b = SummaryNode {
+        id: "sum-b".into(),
+        tree_id: tree.id.clone(),
+        tree_kind: TreeKind::Source,
+        level: 1,
+        parent_id: None,
+        child_ids: vec!["leaf-b".into()],
+        content: "BODY-B".into(),
+        token_count: 22,
+        entities: vec!["entity:bob".into()],
+        topics: vec!["#b".into()],
+        time_range_start: ts,
+        time_range_end: ts,
+        score: 0.22,
+        sealed_at: ts,
+        deleted: false,
+        embedding: None,
+    };
+
+    // Stage bodies to disk + record content pointers so
+    // `read_summary_body` (called per-id inside `hydrate_summary_inputs`)
+    // can resolve the path. Mirrors the production seal write path.
+    let content_root = cfg.memory_tree_content_root();
+    std::fs::create_dir_all(&content_root).expect("create content_root");
+    let stage = |n: &SummaryNode| {
+        stage_summary(
+            &content_root,
+            &SummaryComposeInput {
+                summary_id: &n.id,
+                tree_kind: SummaryTreeKind::Source,
+                tree_id: &tree.id,
+                tree_scope: &tree.scope,
+                level: n.level,
+                child_ids: &n.child_ids,
+                child_basenames: None,
+                child_count: n.child_ids.len(),
+                time_range_start: n.time_range_start,
+                time_range_end: n.time_range_end,
+                sealed_at: n.sealed_at,
+                body: &n.content,
+            },
+            "slack-eng",
+            None,
+        )
+        .unwrap()
+    };
+    let staged_a = stage(&sum_a);
+    let staged_b = stage(&sum_b);
+    crate::openhuman::memory_store::chunks::store::with_connection(&cfg, |conn| {
+        let tx = conn.unchecked_transaction()?;
+        insert_summary_tx(&tx, &sum_a, Some(&staged_a), "test")?;
+        insert_summary_tx(&tx, &sum_b, Some(&staged_b), "test")?;
+        tx.commit()?;
+        Ok(())
+    })
+    .unwrap();
+
+    // Interleaved order with a ghost id in the middle: if the function
+    // ever regresses to indexing by position into the batch result or
+    // returns the HashMap's iteration order, this assertion will catch
+    // it. `sum-b` deliberately comes first so a naive "iterate the map"
+    // implementation that happens to land on `sum-a` first would fail.
+    let ids = vec![
+        "sum-b".to_string(),
+        "ghost:no-such".to_string(),
+        "sum-a".to_string(),
+    ];
+    let out = hydrate_summary_inputs(&cfg, &ids).unwrap();
+
+    // Missing id silently skipped → 2 inputs, not 3.
+    assert_eq!(out.len(), 2, "ghost id must be skipped, not error");
+    // Input order preserved across the gap.
+    assert_eq!(out[0].id, "sum-b");
+    assert_eq!(out[1].id, "sum-a");
+    // Per-id field propagation: each input's score/token_count/content
+    // comes from its own row, not from its sibling.
+    assert_eq!(out[0].token_count, 22);
+    assert!((out[0].score - 0.22).abs() < 1e-6);
+    assert_eq!(out[0].content, "BODY-B");
+    assert_eq!(out[1].token_count, 11);
+    assert!((out[1].score - 0.11).abs() < 1e-6);
+    assert_eq!(out[1].content, "BODY-A");
+    // Entities and topics tracked per id too.
+    assert_eq!(out[0].entities, vec!["entity:bob".to_string()]);
+    assert_eq!(out[1].entities, vec!["entity:alice".to_string()]);
+}