diff --git a/src/openhuman/memory_store/trees/store.rs b/src/openhuman/memory_store/trees/store.rs index 94634b6c5..3ca4e54bf 100644 --- a/src/openhuman/memory_store/trees/store.rs +++ b/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> { }) } +/// 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> { + if summary_ids.is_empty() { + return Ok(HashMap::new()); + } + with_connection(config, |conn| { + let mut out: HashMap = 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::>() + .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::>>() + .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( diff --git a/src/openhuman/memory_store/trees/store_tests.rs b/src/openhuman/memory_store/trees/store_tests.rs index 9570e9364..7ab9625c9 100644 --- a/src/openhuman/memory_store/trees/store_tests.rs +++ b/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()); +} diff --git a/src/openhuman/memory_tree/tree/bucket_seal.rs b/src/openhuman/memory_tree/tree/bucket_seal.rs index 4d3865ff5..422959028 100644 --- a/src/openhuman/memory_tree/tree/bucket_seal.rs +++ b/src/openhuman/memory_tree/tree/bucket_seal.rs @@ -821,17 +821,25 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result Result> { 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 = 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 diff --git a/src/openhuman/memory_tree/tree/bucket_seal_tests.rs b/src/openhuman/memory_tree/tree/bucket_seal_tests.rs index 3ce1fe24f..8edd327ee 100644 --- a/src/openhuman/memory_tree/tree/bucket_seal_tests.rs +++ b/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()]); +}