perf(memory_tree): batch tree fetch in flush_stale_buffers (#2976)

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

src/openhuman/memory_store/trees/store.rs

@@ -107,6 +107,72 @@ pub fn get_tree(config: &Config, id: &str) -> Result<Option<Tree>> {
     })
 }
 
+/// 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
+/// (`memory_tree::tree::flush::flush_stale_buffers`) passes one tree_id
+/// per stale L0 buffer, so a typical N (tens to low hundreds across all
+/// connected sources) runs the loop 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 TREES_MAX_FETCH_BATCH: usize = 500;
+
+/// Fetch many trees by id in a single SQL round-trip per
+/// [`TREES_MAX_FETCH_BATCH`] window. Replaces the per-id `get_tree`
+/// loop inside paths like the cron-driven `flush_stale_buffers`, where
+/// the previous code did one `SELECT … WHERE id = ?` per stale buffer.
+/// Missing ids are silently absent from the map so callers can preserve
+/// the existing "row missing → warn-and-skip" contract without an extra
+/// `Ok(None)` sentinel per id.
+pub fn get_trees_batch(config: &Config, tree_ids: &[String]) -> Result<HashMap<String, Tree>> {
+    if tree_ids.is_empty() {
+        return Ok(HashMap::new());
+    }
+    log::debug!(
+        "[tree::store] get_trees_batch: ids={} max_batch={TREES_MAX_FETCH_BATCH}",
+        tree_ids.len()
+    );
+    with_connection(config, |conn| {
+        let mut out: HashMap<String, Tree> = HashMap::with_capacity(tree_ids.len());
+        for window in tree_ids.chunks(TREES_MAX_FETCH_BATCH) {
+            // SAFETY (SQL injection): only the *count* of placeholders is
+            // interpolated into the query string — `?1,?2,…` are numbered
+            // bind slots, never the id values. Every id is bound via typed
+            // `rusqlite::ToSql` params below; nothing user-controlled is
+            // ever formatted into `sql`. Do NOT inline id values here.
+            let placeholders = (1..=window.len())
+                .map(|i| format!("?{i}"))
+                .collect::<Vec<_>>()
+                .join(",");
+            let sql = format!(
+                "SELECT id, kind, scope, root_id, max_level, status,
+                        created_at_ms, last_sealed_at_ms
+                   FROM mem_tree_trees
+                  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_tree)?
+                .collect::<rusqlite::Result<Vec<_>>>()
+                .context("Failed to collect trees batch")?;
+            for t in rows {
+                out.insert(t.id.clone(), t);
+            }
+        }
+        log::debug!(
+            "[tree::store] get_trees_batch: requested={} found={}",
+            tree_ids.len(),
+            out.len()
+        );
+        Ok(out)
+    })
+}
+
 /// List every tree of a given kind. Used by the global digest to enumerate
 /// source trees, and by diagnostics. Rows come back ordered by `created_at_ms`
 /// ASC so callers see a stable iteration order.

src/openhuman/memory_store/trees/store_tests.rs

@@ -288,6 +288,51 @@ fn list_stale_buffers_orders_by_age() {
     assert_eq!(only_oldest[0].tree_id, "tree-1");
 }
 
+// ── get_trees_batch ────────────────────────────────────────────────────
+//
+// Same shape as `chunks::store::get_chunks_batch` /
+// `score::store::get_scores_batch`: present ids decode through the same
+// `row_to_tree` path as the per-id `get_tree` and land in a `HashMap`
+// keyed by id; missing ids are silently absent so the
+// `flush_stale_buffers` orphan-buffer warn-and-skip path keeps working
+// without an extra Ok(None) sentinel per id.
+
+#[test]
+fn get_trees_batch_returns_present_ids_in_map() {
+    let (_tmp, cfg) = test_config();
+    let a = sample_tree("tree-a", "slack:#eng");
+    let b = sample_tree("tree-b", "slack:#design");
+    insert_tree(&cfg, &a).unwrap();
+    insert_tree(&cfg, &b).unwrap();
+
+    let ids = vec!["tree-a".to_string(), "tree-b".to_string()];
+    let map = get_trees_batch(&cfg, &ids).unwrap();
+    assert_eq!(map.len(), 2);
+    // Each decoded row must match the per-id `get_tree` path bit-for-bit
+    // — same `row_to_tree` decoder under the hood, so the structs are
+    // equal including the parsed `kind` / `status` enums.
+    assert_eq!(map.get("tree-a").unwrap(), &a);
+    assert_eq!(map.get("tree-b").unwrap(), &b);
+}
+
+#[test]
+fn get_trees_batch_empty_input_and_missing_ids() {
+    // Empty input: empty map (no SQL issued).
+    let (_tmp, cfg) = test_config();
+    let empty = get_trees_batch(&cfg, &[]).unwrap();
+    assert!(empty.is_empty());
+
+    // Missing ids: silently absent so `flush_stale_buffers` can warn
+    // + skip without an extra `Ok(None)` sentinel per id.
+    let a = sample_tree("tree-a", "slack:#eng");
+    insert_tree(&cfg, &a).unwrap();
+    let ids = vec!["tree-a".to_string(), "ghost:no-such".to_string()];
+    let map = get_trees_batch(&cfg, &ids).unwrap();
+    assert_eq!(map.len(), 1);
+    assert_eq!(map.get("tree-a").unwrap(), &a);
+    assert!(map.get("ghost:no-such").is_none());
+}
+
 // ── get_summaries_batch ────────────────────────────────────────────────
 //
 // Same shape as `chunks::store::get_chunks_batch` /

src/openhuman/memory_tree/tree/flush.rs

@@ -27,6 +27,8 @@ pub async fn flush_stale_buffers(
     max_age: Duration,
     strategy: &LabelStrategy,
 ) -> Result<usize> {
+    use crate::openhuman::memory_store::trees::store::get_trees_batch;
+
     let now = Utc::now();
     let cutoff = now - max_age;
     let stale = store::list_stale_buffers(config, cutoff)?;
@@ -36,20 +38,38 @@ pub async fn flush_stale_buffers(
         max_age
     );
 
+    // One batched `SELECT … WHERE id IN (?,?,…)` over the distinct
+    // tree_ids instead of N per-buffer `get_tree` round-trips. `stale`
+    // is filtered to L0 by `list_stale_buffers` so the same tree_id
+    // typically appears at most once, but we still de-dup defensively
+    // before hitting the DB — future widenings of the filter (e.g.
+    // dropping the `level = 0` clause) would otherwise re-pay the cost.
+    // Missing rows stay silently absent from the map; the
+    // `Some(t) => ... / None => warn-and-skip` orphan-buffer path below
+    // preserves the per-id `get_tree` `Ok(None)` semantics.
+    let distinct_tree_ids: Vec<String> = {
+        let mut seen = std::collections::HashSet::new();
+        let mut out = Vec::new();
+        for buf in &stale {
+            if seen.insert(buf.tree_id.clone()) {
+                out.push(buf.tree_id.clone());
+            }
+        }
+        out
+    };
+    let tree_by_id = get_trees_batch(config, &distinct_tree_ids)?;
+
     let mut seals: usize = 0;
     for buf in stale {
-        let tree = match store::get_tree(config, &buf.tree_id)? {
-            Some(t) => t,
-            None => {
-                log::warn!(
-                    "[tree::flush] orphan buffer tree_id={} level={}",
-                    buf.tree_id,
-                    buf.level
-                );
-                continue;
-            }
+        let Some(tree) = tree_by_id.get(&buf.tree_id) else {
+            log::warn!(
+                "[tree::flush] orphan buffer tree_id={} level={}",
+                buf.tree_id,
+                buf.level
+            );
+            continue;
         };
-        let sealed = cascade_all_from(config, &tree, buf.level, Some(now), strategy).await?;
+        let sealed = cascade_all_from(config, tree, buf.level, Some(now), strategy).await?;
         seals += sealed.len();
     }
     Ok(seals)
@@ -274,4 +294,97 @@ mod tests {
         assert_eq!(seals, 0);
         assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
     }
+
+    /// Regression for the `get_trees_batch` refactor: two distinct
+    /// trees each carry a stale L0 buffer; `flush_stale_buffers` must
+    /// seal both via a single batched tree-fetch. Pre-refactor this
+    /// path issued N per-buffer `get_tree(...)` round-trips; post-
+    /// refactor it's one `SELECT … WHERE id IN (?,?)` and a per-id
+    /// HashMap lookup. The interesting bit is the HashMap lookup
+    /// keying by tree_id (not by `enumerate()` position over the input
+    /// slice) — with two trees, swapping the lookup key would leak one
+    /// tree's `Tree` into the other tree's seal and either error out
+    /// or write summaries against the wrong `root_id`.
+    #[tokio::test]
+    async fn flush_seals_multiple_distinct_trees_via_batched_lookup() {
+        let (_tmp, cfg) = test_config();
+        let tree_a = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
+        let tree_b = get_or_create_source_tree(&cfg, "slack:#design").unwrap();
+        assert_ne!(
+            tree_a.id, tree_b.id,
+            "test prerequisite: distinct source_ids must yield distinct tree_ids"
+        );
+        let provider: Arc<dyn ChatProvider> =
+            Arc::new(StaticChatProvider::new("test summary content"));
+
+        let old_ts = Utc::now() - Duration::days(10);
+
+        // Plant a stale L0 leaf in each tree, backed by real chunks so
+        // the cascade can seal end-to-end.
+        for (src_id, content) in [
+            ("slack:#eng", "engineering content"),
+            ("slack:#design", "design content"),
+        ] {
+            let c = Chunk {
+                id: chunk_id(SourceKind::Chat, src_id, 0, content),
+                content: content.into(),
+                metadata: Metadata {
+                    source_kind: SourceKind::Chat,
+                    source_id: src_id.into(),
+                    owner: "alice".into(),
+                    timestamp: old_ts,
+                    time_range: (old_ts, old_ts),
+                    tags: vec![],
+                    source_ref: Some(SourceRef::new(&format!("slack://{src_id}"))),
+                },
+                token_count: 100,
+                seq_in_source: 0,
+                created_at: old_ts,
+                partial_message: false,
+            };
+            upsert_chunks(&cfg, &[c.clone()]).unwrap();
+            stage_test_chunks(&cfg, &[c.clone()]);
+            let leaf = LeafRef {
+                chunk_id: c.id.clone(),
+                token_count: 100,
+                timestamp: old_ts,
+                content: c.content.clone(),
+                entities: vec![],
+                topics: vec![],
+                score: 0.5,
+            };
+            let target_tree = if src_id == "slack:#eng" {
+                &tree_a
+            } else {
+                &tree_b
+            };
+            test_override::with_provider(Arc::clone(&provider), async {
+                append_leaf(&cfg, target_tree, &leaf, &LabelStrategy::Empty)
+                    .await
+                    .unwrap();
+            })
+            .await;
+        }
+
+        let seals = test_override::with_provider(provider, async {
+            flush_stale_buffers(&cfg, Duration::days(7), &LabelStrategy::Empty)
+                .await
+                .unwrap()
+        })
+        .await;
+
+        // Both trees sealed via a single batched fetch — the HashMap
+        // lookup correctly routed each buffer to its own `Tree`.
+        assert_eq!(seals, 2, "both stale trees must seal in one flush pass");
+        assert_eq!(
+            store::count_summaries(&cfg, &tree_a.id).unwrap(),
+            1,
+            "tree_a must own its summary (HashMap keyed by id, not position)"
+        );
+        assert_eq!(
+            store::count_summaries(&cfg, &tree_b.id).unwrap(),
+            1,
+            "tree_b must own its summary (HashMap keyed by id, not position)"
+        );
+    }
 }