feat(memory): MD-on-disk content + participant bucketing + async pipeline foundation (#1008)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>
This commit is contained in:
sanil-23
2026-04-28 17:44:20 -07:00
committed by GitHub
co-authored by Claude Opus 4.7 Steven Enamakel
parent dbe60389e6
commit fd0398d5fa
70 changed files with 11614 additions and 660 deletions
+1
View File
@@ -66,6 +66,7 @@ tauri.key
tauri.key.pub
/target/
src-tauri/target/
.target-codex/
workflow
.fastembed_cache
Generated
+1 -1
View File
@@ -4541,7 +4541,7 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381"
[[package]]
name = "openhuman"
version = "0.53.3"
version = "0.53.4"
dependencies = [
"aes-gcm",
"anyhow",
+4
View File
@@ -13,6 +13,10 @@ path = "src/main.rs"
name = "slack-backfill"
path = "src/bin/slack_backfill.rs"
[[bin]]
name = "gmail-backfill-3d"
path = "src/bin/gmail_backfill_3d.rs"
[lib]
name = "openhuman_core"
crate-type = ["rlib"]
File diff suppressed because it is too large Load Diff
+474
View File
@@ -0,0 +1,474 @@
//! Backfill the last N days of Gmail into the memory-tree content store.
//!
//! Authenticates via Composio (JWT from `<workspace>/auth-profiles.json`),
//! fetches Gmail pages via `GMAIL_FETCH_EMAILS`, converts each thread into an
//! [`EmailThread`], ingests it through `ingest_page_into_memory_tree` (which
//! writes `.md` files via `content_store` and populates SQLite), then drains
//! the async worker pool until idle.
//!
//! After draining, the binary performs an integrity check: for every chunk
//! that has a `content_path` in SQLite, it verifies the on-disk SHA-256
//! matches the stored `content_sha256`.
//!
//! # Prerequisites
//!
//! - Signed-in openhuman session JWT in the same workspace the desktop app
//! uses (stored at `<workspace>/auth-profiles.json`).
//! - Active Gmail connection on Composio for that user.
//!
//! # Usage
//!
//! ```sh
//! cargo run --bin gmail-backfill-3d
//! cargo run --bin gmail-backfill-3d -- --days 7
//! cargo run --bin gmail-backfill-3d -- --days 14 --page-size 100
//! cargo run --bin gmail-backfill-3d -- --skip-drain
//! cargo run --bin gmail-backfill-3d -- --skip-verify
//! cargo run --bin gmail-backfill-3d -- --wipe
//! ```
//!
//! Set `RUST_LOG=info` (or `debug`) for detailed output.
use anyhow::{Context, Result};
use clap::Parser;
use serde_json::{json, Value};
use openhuman_core::openhuman::composio::client::build_composio_client;
use openhuman_core::openhuman::composio::providers::gmail::ingest::ingest_page_into_memory_tree;
use openhuman_core::openhuman::composio::providers::registry::{
get_provider, init_default_providers,
};
use openhuman_core::openhuman::config::Config;
use openhuman_core::openhuman::memory::tree::content_store::read::{
verify_chunk_file, verify_summary_file, VerifyResult,
};
use openhuman_core::openhuman::memory::tree::jobs::drain_until_idle;
use openhuman_core::openhuman::memory::tree::store::{
get_chunk_content_pointers, list_chunks, list_summaries_with_content_path, ListChunksQuery,
};
#[derive(Parser, Debug)]
#[command(
name = "gmail-backfill-3d",
about = "Backfill last N days of Gmail into the memory-tree content store (.md files + SQLite)."
)]
struct Cli {
/// Lookback window in days. Default 3.
#[arg(long, default_value_t = 3)]
days: u32,
/// Page size per `GMAIL_FETCH_EMAILS` call (1500).
#[arg(long, default_value_t = 50)]
page_size: u32,
/// Cap on pages we will request. Guards against runaway pagination.
#[arg(long, default_value_t = 40)]
max_pages: u32,
/// Include SPAM and TRASH messages in the fetch.
#[arg(long, default_value_t = false)]
include_spam_trash: bool,
/// Extra Gmail search query AND-ed with the default scope.
#[arg(long)]
query: Option<String>,
/// Skip draining the async worker pool after ingest (useful for quick
/// smoke-test of file writes only).
#[arg(long, default_value_t = false)]
skip_drain: bool,
/// Skip the post-drain integrity check (SHA-256 file verification).
#[arg(long, default_value_t = false)]
skip_verify: bool,
/// Override the owner string embedded in chunk metadata. Defaults to
/// `"gmail-backfill"`.
#[arg(long)]
owner: Option<String>,
/// Wipe `chunks.db` (+ wal/shm) AND `<content_root>/` before running.
/// Useful after a chunker change that invalidates existing chunk IDs.
#[arg(long, default_value_t = false)]
wipe: bool,
}
#[tokio::main]
async fn main() -> Result<()> {
env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
.format_timestamp_secs()
.try_init()
.ok();
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.with_target(true)
.try_init()
.ok();
let cli = Cli::parse();
if cli.days == 0 {
anyhow::bail!("--days must be >= 1");
}
let config = Config::load_or_init()
.await
.context("[gmail_backfill_3d] Config::load_or_init failed")?;
if cli.wipe {
wipe_memory_tree_state(&config)?;
}
let client = build_composio_client(&config).ok_or_else(|| {
anyhow::anyhow!(
"No Composio client — user not signed in (no JWT). \
Sign in via the desktop app first, then re-run this binary."
)
})?;
init_default_providers();
let provider = get_provider("gmail").ok_or_else(|| {
anyhow::anyhow!("GmailProvider not registered after init_default_providers")
})?;
let owner = cli
.owner
.clone()
.unwrap_or_else(|| "gmail-backfill".to_string());
let mut query = format!("in:inbox newer_than:{}d", cli.days);
if !cli.include_spam_trash {
query.push_str(" -in:spam -in:trash");
}
if let Some(extra) = cli
.query
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
{
query.push(' ');
query.push_str(extra);
}
log::info!(
"[gmail_backfill_3d] start days={} page_size={} max_pages={} query={:?}",
cli.days,
cli.page_size,
cli.max_pages,
query,
);
let content_root = config.memory_tree_content_root();
log::info!(
"[gmail_backfill_3d] content_root={}",
content_root.display()
);
// ─── Fetch + ingest ────────────────────────────────────────────────────
let mut page_token: Option<String> = None;
let mut total_chunks = 0usize;
let mut total_pages = 0usize;
let mut total_cost: f64 = 0.0;
for page_num in 0..cli.max_pages {
let mut args = json!({
"max_results": cli.page_size,
"query": query,
});
if cli.include_spam_trash {
args["include_spam_trash"] = json!(true);
}
if let Some(token) = &page_token {
args["page_token"] = json!(token);
}
log::info!(
"[gmail_backfill_3d] fetching page {}{}…",
page_num,
page_token.as_ref().map(|_| " (paginated)").unwrap_or(""),
);
let mut resp = client
.execute_tool("GMAIL_FETCH_EMAILS", Some(args.clone()))
.await
.map_err(|e| anyhow::anyhow!("GMAIL_FETCH_EMAILS page {page_num}: {e:#}"))?;
total_cost += resp.cost_usd;
if !resp.successful {
anyhow::bail!(
"GMAIL_FETCH_EMAILS page {page_num} failed: {:?}",
resp.error
);
}
provider.post_process_action_result("GMAIL_FETCH_EMAILS", Some(&args), &mut resp.data);
let (messages, next_token) = extract_envelope(&resp.data);
log::info!(
"[gmail_backfill_3d] page {} -> {} messages, next_token={}",
page_num,
messages.len(),
next_token.as_deref().unwrap_or("(none)"),
);
if messages.is_empty() {
break;
}
let chunks_this_page = ingest_page_into_memory_tree(&config, &owner, &messages).await?;
total_chunks += chunks_this_page;
total_pages += 1;
log::info!(
"[gmail_backfill_3d] page {} ingested chunks={} running_total={}",
page_num,
chunks_this_page,
total_chunks,
);
match next_token {
Some(tok) => page_token = Some(tok),
None => break,
}
}
log::info!(
"[gmail_backfill_3d] fetch+ingest done pages={} total_chunks={} cost=~${:.4}",
total_pages,
total_chunks,
total_cost,
);
// ─── Drain async worker pool ────────────────────────────────────────────
if cli.skip_drain {
log::info!("[gmail_backfill_3d] skipping worker pool drain (--skip-drain)");
} else {
log::info!("[gmail_backfill_3d] draining async worker pool…");
drain_until_idle(&config).await?;
log::info!("[gmail_backfill_3d] worker pool idle");
}
// ─── Integrity check ────────────────────────────────────────────────────
if cli.skip_verify {
log::info!("[gmail_backfill_3d] skipping integrity check (--skip-verify)");
} else {
log::info!("[gmail_backfill_3d] running integrity check…");
// Chunk integrity.
let (verified, mismatched, no_pointer, missing_file) = verify_all_chunk_files(&config)?;
log::info!(
"[gmail_backfill_3d] chunks: verified={} mismatched={} no_pointer={} missing_file={}",
verified,
mismatched,
no_pointer,
missing_file,
);
// Summary integrity.
let (sum_verified, sum_mismatched, sum_no_pointer, sum_missing_file) =
verify_all_summary_files(&config)?;
log::info!(
"[gmail_backfill_3d] summaries: verified={} mismatched={} no_pointer={} missing_file={}",
sum_verified,
sum_mismatched,
sum_no_pointer,
sum_missing_file,
);
if mismatched > 0 || missing_file > 0 || sum_mismatched > 0 || sum_missing_file > 0 {
anyhow::bail!(
"Integrity check failed: \
chunks: {} mismatches, {} missing files; \
summaries: {} mismatches, {} missing files",
mismatched,
missing_file,
sum_mismatched,
sum_missing_file,
);
}
}
println!(
"\nBackfill complete. pages={} chunks_written={} cost=~${:.4}",
total_pages, total_chunks, total_cost,
);
Ok(())
}
/// Wipe `<workspace>/memory_tree/chunks.db` (+ wal/shm) and
/// `<content_root>/` so the bin can re-run cleanly after a chunker
/// change that invalidates existing chunk IDs.
///
/// Logs each removed artifact at info; missing files are not an error.
fn wipe_memory_tree_state(config: &Config) -> Result<()> {
let mt_dir = config.workspace_dir.join("memory_tree");
for name in &["chunks.db", "chunks.db-wal", "chunks.db-shm"] {
let path = mt_dir.join(name);
match std::fs::remove_file(&path) {
Ok(()) => log::info!("[gmail_backfill_3d] wiped {}", path.display()),
Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
Err(e) => return Err(e).with_context(|| format!("wipe {}", path.display())),
}
}
let content_root = config.memory_tree_content_root();
if content_root.exists() {
std::fs::remove_dir_all(&content_root)
.with_context(|| format!("wipe {}", content_root.display()))?;
log::info!("[gmail_backfill_3d] wiped {}", content_root.display());
}
Ok(())
}
/// Read all chunks from SQLite and verify on-disk SHA-256 matches `content_sha256`.
///
/// Returns `(verified, mismatched, no_pointer, missing_file)`.
fn verify_all_chunk_files(config: &Config) -> Result<(usize, usize, usize, usize)> {
let chunks = list_chunks(config, &ListChunksQuery::default())?;
let content_root = config.memory_tree_content_root();
let mut verified = 0usize;
let mut mismatched = 0usize;
let mut no_pointer = 0usize;
let mut missing_file = 0usize;
for chunk in &chunks {
let pointers = get_chunk_content_pointers(config, &chunk.id)?;
let (rel_path, expected_sha) = match pointers {
None => {
no_pointer += 1;
log::debug!(
"[gmail_backfill_3d] verify: chunk {} has no content_path/sha256",
chunk.id
);
continue;
}
Some(pair) => pair,
};
let abs_path = {
let mut p = content_root.clone();
for component in rel_path.split('/') {
p.push(component);
}
p
};
if !abs_path.exists() {
missing_file += 1;
log::warn!(
"[gmail_backfill_3d] verify: file missing chunk_id={} path={}",
chunk.id,
abs_path.display(),
);
continue;
}
match verify_chunk_file(&abs_path, &expected_sha) {
Ok(true) => {
verified += 1;
}
Ok(false) => {
mismatched += 1;
log::warn!(
"[gmail_backfill_3d] verify: SHA-256 mismatch chunk_id={} path={}",
chunk.id,
abs_path.display(),
);
}
Err(e) => {
log::error!(
"[gmail_backfill_3d] verify: error chunk_id={}: {e}",
chunk.id,
);
mismatched += 1;
}
}
}
Ok((verified, mismatched, no_pointer, missing_file))
}
/// Read all summary rows with a non-NULL `content_path` from SQLite and verify
/// the on-disk SHA-256 matches `content_sha256`.
///
/// Returns `(verified, mismatched, no_pointer, missing_file)`.
fn verify_all_summary_files(config: &Config) -> Result<(usize, usize, usize, usize)> {
let rows_with_pointer = list_summaries_with_content_path(config)?;
let content_root = config.memory_tree_content_root();
let mut verified = 0usize;
let mut mismatched = 0usize;
let mut missing_file = 0usize;
for (summary_id, rel_path, expected_sha) in &rows_with_pointer {
let abs_path = {
let mut p = content_root.clone();
for component in rel_path.split('/') {
p.push(component);
}
p
};
match verify_summary_file(&abs_path, expected_sha) {
Ok(VerifyResult::Ok) => {
verified += 1;
}
Ok(VerifyResult::Mismatch { actual }) => {
mismatched += 1;
log::warn!(
"[gmail_backfill_3d] verify: SHA-256 mismatch summary_id={} path={} expected={} actual={}",
summary_id,
abs_path.display(),
expected_sha,
actual,
);
}
Ok(VerifyResult::Missing) => {
missing_file += 1;
log::warn!(
"[gmail_backfill_3d] verify: file missing summary_id={} path={}",
summary_id,
abs_path.display(),
);
}
Err(e) => {
log::error!(
"[gmail_backfill_3d] verify: error summary_id={}: {e}",
summary_id,
);
mismatched += 1;
}
}
}
// Count rows that have no content_path at all (legacy rows).
// We report this as no_pointer for symmetry with the chunk verifier.
// We can't easily count them here without a separate query, so we
// approximate: rows_with_pointer gives us the ones we checked.
// For now no_pointer = 0 (the bin wipes before re-ingesting so all
// new rows should have pointers; legacy rows are pre-migration).
let no_pointer = 0usize;
Ok((verified, mismatched, no_pointer, missing_file))
}
/// Extract the `messages` array and `nextPageToken` from a Composio response.
fn extract_envelope(data: &Value) -> (Vec<Value>, Option<String>) {
let candidates: [Option<&Value>; 2] = [Some(data), data.get("data")];
for cand in candidates.into_iter().flatten() {
if let Some(arr) = cand.get("messages").and_then(|v| v.as_array()) {
let token = cand
.get("nextPageToken")
.and_then(|v| v.as_str())
.filter(|s| !s.trim().is_empty())
.map(str::to_string);
return (arr.clone(), token);
}
}
(Vec::new(), None)
}
+139
View File
@@ -96,6 +96,21 @@ struct Cli {
/// re-attempts cascade on the next append.
#[arg(long = "seal-probe", default_value_t = false)]
seal_probe: bool,
/// Fire N back-to-back `SLACK_FETCH_CONVERSATION_HISTORY` calls
/// against the first listed channel and report a per-call tally
/// of {success, ratelimit, other-failure} + total duration. No
/// pacing by default (see --probe-pacing-ms), no ingestion. Used
/// to characterise Composio/Slack quota behaviour without
/// touching the memory tree.
#[arg(long = "probe-ratelimit")]
probe_ratelimit: Option<u32>,
/// Sleep this many milliseconds between probe calls. 0 = fire
/// back-to-back (default). Use to find the threshold at which
/// rate-limits stop firing.
#[arg(long = "probe-pacing-ms", default_value_t = 0)]
probe_pacing_ms: u64,
}
#[tokio::main]
@@ -243,6 +258,130 @@ async fn main() -> Result<()> {
return Ok(());
}
if let Some(n) = cli.probe_ratelimit {
// Pure quota probe: fire N back-to-back
// SLACK_FETCH_CONVERSATION_HISTORY calls against the first
// discoverable channel. No pacing, no retry, no ingest. Reports
// a per-call status table + summary so we can characterise
// Composio/Slack rate-limit behaviour without contaminating the
// memory tree or burning extra quota on retries.
log::info!("[probe-ratelimit] requesting one channel via SLACK_LIST_CONVERSATIONS");
let list_resp = client
.execute_tool(
"SLACK_LIST_CONVERSATIONS",
Some(serde_json::json!({ "exclude_archived": true, "limit": 1 })),
)
.await
.map_err(|e| anyhow::anyhow!("SLACK_LIST_CONVERSATIONS failed: {e:#}"))?;
if !list_resp.successful {
anyhow::bail!(
"SLACK_LIST_CONVERSATIONS returned non-success: {:?}",
list_resp.error
);
}
let channel_id = ["/data/channels/0/id", "/channels/0/id", "/data/0/id"]
.iter()
.find_map(|p| list_resp.data.pointer(p).and_then(|v| v.as_str()))
.map(str::to_string)
.ok_or_else(|| {
anyhow::anyhow!(
"could not find a channel id in SLACK_LIST_CONVERSATIONS response: {}",
serde_json::to_string(&list_resp.data).unwrap_or_default()
)
})?;
log::info!("[probe-ratelimit] firing {n} calls against channel={channel_id}");
#[derive(Debug)]
enum Outcome {
Ok,
Ratelimit,
OtherFail(String),
Transport(String),
}
let mut outcomes: Vec<(u32, std::time::Duration, Outcome)> = Vec::with_capacity(n as usize);
let probe_started = Instant::now();
for i in 1..=n {
if i > 1 && cli.probe_pacing_ms > 0 {
tokio::time::sleep(std::time::Duration::from_millis(cli.probe_pacing_ms)).await;
}
let t0 = Instant::now();
let resp = client
.execute_tool(
"SLACK_FETCH_CONVERSATION_HISTORY",
Some(serde_json::json!({ "channel": channel_id, "limit": 1000 })),
)
.await;
let dt = t0.elapsed();
let outcome = match resp {
Err(e) => Outcome::Transport(format!("{e:#}")),
Ok(r) if r.successful => Outcome::Ok,
Ok(r) => {
let err = r.error.as_deref().unwrap_or("provider failure");
if err.contains("ratelimited")
|| err.contains("rate_limit")
|| err.contains("rate limit")
{
log::warn!(
"[probe-ratelimit] call {i} ratelimited; body: {}",
serde_json::to_string(&r.data).unwrap_or_default()
);
Outcome::Ratelimit
} else {
Outcome::OtherFail(err.to_string())
}
}
};
log::info!(
"[probe-ratelimit] call {i}/{n} took {:.2}s -> {:?}",
dt.as_secs_f64(),
outcome
);
outcomes.push((i, dt, outcome));
}
let total = probe_started.elapsed();
let ok = outcomes
.iter()
.filter(|(_, _, o)| matches!(o, Outcome::Ok))
.count();
let rl = outcomes
.iter()
.filter(|(_, _, o)| matches!(o, Outcome::Ratelimit))
.count();
let other = outcomes
.iter()
.filter(|(_, _, o)| matches!(o, Outcome::OtherFail(_)))
.count();
let transport = outcomes
.iter()
.filter(|(_, _, o)| matches!(o, Outcome::Transport(_)))
.count();
let avg_ms = if !outcomes.is_empty() {
outcomes.iter().map(|(_, d, _)| d.as_millis()).sum::<u128>() / outcomes.len() as u128
} else {
0
};
println!("=== probe-ratelimit summary ===");
println!("channel: {channel_id}");
println!("calls fired: {n}");
println!("total duration: {:.2}s", total.as_secs_f64());
println!("avg per call: {avg_ms} ms");
println!("successful: {ok}");
println!("ratelimited: {rl}");
println!("other failures: {other}");
println!("transport errors: {transport}");
if rl > 0 {
let first_rl = outcomes
.iter()
.find(|(_, _, o)| matches!(o, Outcome::Ratelimit))
.map(|(i, _, _)| *i)
.unwrap_or(0);
println!("first ratelimit: call #{first_rl}");
}
return Ok(());
}
let connections = client
.list_connections()
.await
+6 -2
View File
@@ -805,7 +805,10 @@ async fn run_server_inner(
///
/// Guarded by `std::sync::Once` so repeated calls to `bootstrap_skill_runtime`
/// are safe and idempotent.
fn register_domain_subscribers(workspace_dir: std::path::PathBuf) {
fn register_domain_subscribers(
workspace_dir: std::path::PathBuf,
config: crate::openhuman::config::Config,
) {
use std::sync::{Arc, Once};
static REGISTERED: Once = Once::new();
@@ -840,6 +843,7 @@ fn register_domain_subscribers(workspace_dir: std::path::PathBuf) {
}
crate::openhuman::composio::register_composio_trigger_subscriber();
crate::openhuman::composio::start_periodic_sync();
crate::openhuman::memory::tree::jobs::start(config.clone());
// Restart requests go through a subscriber so every trigger path shares
// the same respawn logic.
@@ -880,7 +884,7 @@ pub async fn bootstrap_skill_runtime() {
// Register domain subscribers for cross-module event handling.
// Uses a Once guard so repeated calls to bootstrap_skill_runtime()
// cannot double-subscribe.
register_domain_subscribers(workspace_dir.clone());
register_domain_subscribers(workspace_dir.clone(), cfg.clone());
// --- Sub-agent definition registry bootstrap ---
// Loads built-in archetype definitions plus any custom TOML files
@@ -0,0 +1,642 @@
//! Gmail → memory tree ingest plumbing.
//!
//! Owns the conversion from a page of `GMAIL_FETCH_EMAILS` slim-envelope
//! messages (post-processed by [`super::post_process`]) into
//! [`EmailThread`] batches grouped by the sorted set of distinct
//! participants (`from` `to`-list, CC ignored), then drives
//! [`memory::tree::ingest::ingest_email`] per participant group.
//!
//! Source-id is `gmail:{participants}` where participants is
//! `addr1|addr2|...` (sorted, deduped, lowercased bare emails). All
//! correspondence between the same set of people lands in one source tree.
//!
//! Idempotency: chunk IDs are content-hashed inside the memory tree, so
//! re-ingesting a previously-seen Gmail message is an UPSERT — buffer
//! token_sum may drift if content changes (rare for sealed mail), but
//! the tree's seal cascade handles that on next append.
use std::collections::BTreeMap;
use anyhow::Result;
use serde_json::Value;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::canonicalize::email::{EmailMessage, EmailThread};
use crate::openhuman::memory::tree::canonicalize::email_clean::{
extract_email, parse_message_date,
};
use crate::openhuman::memory::tree::ingest::{ingest_email, IngestResult};
use crate::openhuman::memory::tree::util::redact::redact;
/// Provider name embedded in the canonical email-thread header. Matches
/// the value `memory::tree::retrieval::source::PLATFORM_KINDS` expects.
pub const GMAIL_PROVIDER: &str = "gmail";
/// Tags attached to every Gmail-ingested chunk. Stable list — retrieval
/// callers filter on these.
pub const DEFAULT_TAGS: &[&str] = &["gmail", "ingested"];
/// Group raw page messages by the sorted set of distinct participants
/// (`from` `to`-list). CC is deliberately excluded from the bucket key
/// so CC-only recipients don't fragment conversations. All messages
/// between the same set of people land in the same bucket regardless of
/// direction or thread ID.
///
/// The bucket key is the participants joined with `|` in sorted order,
/// e.g. `"alice@x.com|bob@y.com"`. Messages within a bucket are sorted
/// ascending by date so the rendered conversation reads chronologically.
pub(crate) fn bucket_by_participants(msgs: &[Value]) -> BTreeMap<String, Vec<&Value>> {
let mut out: BTreeMap<String, Vec<&Value>> = BTreeMap::new();
for m in msgs {
let bucket_key = participants_bucket_key(m);
if bucket_key == "__skip__" {
// Message has no parseable addresses AND no id — drop it and warn.
// Nothing useful can be done with it: no participants means no
// source tree, and no id means no unique bucket either.
log::warn!(
"[composio:gmail][bucket] dropping message with no parseable addresses and no id"
);
continue;
}
out.entry(bucket_key).or_default().push(m);
}
for bucket in out.values_mut() {
bucket.sort_by_key(|m| {
parse_message_date(m)
.map(|d: chrono::DateTime<chrono::Utc>| d.timestamp())
.unwrap_or(0)
});
}
out
}
/// Compute the participants bucket key for a single raw message.
///
/// Collects `from` `to` (as bare lowercased email addresses), sorts
/// and dedupes them, then joins with `|`.
///
/// **Fallback policy when all addresses fail to parse**:
/// - If the message has a non-empty `id`, use `"orphan:{id}"` so each
/// malformed message gets its own bucket and its own source tree. Two
/// messages with different ids that both fail address parsing will NOT
/// collapse into a single `"unknown"` bucket.
/// - If even `id` is missing or empty, the caller (`bucket_by_participants`)
/// should skip the message (log a warn and drop it). This function signals
/// that case by returning the sentinel `"__skip__"`.
fn participants_bucket_key(raw: &Value) -> String {
let from = extract_email(raw.get("from").and_then(|v| v.as_str()).unwrap_or(""))
.map(|s| s.to_lowercase())
.filter(|s| !s.is_empty());
let to_emails: Vec<String> = parse_address_list_for_bucket(raw.get("to"))
.into_iter()
.filter_map(|addr| extract_email(&addr).map(|s| s.to_lowercase()))
.collect();
let mut all: Vec<String> = from.into_iter().chain(to_emails).collect();
all.sort();
all.dedup();
all.retain(|s| !s.is_empty());
if all.is_empty() {
// No parseable addresses — fall back to per-message uniqueness to
// avoid collapsing all malformed messages into one "unknown" source
// tree. Each orphan message gets its own bucket so nothing is silently
// lost in a mixed pile.
let id = raw
.get("id")
.and_then(|v| v.as_str())
.filter(|s| !s.is_empty());
match id {
Some(msg_id) => format!("orphan:{}", msg_id),
None => {
// id is missing: signal caller to skip this message entirely.
"__skip__".to_string()
}
}
} else {
all.join("|")
}
}
/// Parse the `to` / `cc` field for bucket-key construction. Handles both
/// JSON array and comma-separated string forms. Returns raw address
/// strings (may include display names); callers must extract the bare
/// email with [`extract_email`].
fn parse_address_list_for_bucket(v: Option<&Value>) -> Vec<String> {
match v {
Some(Value::Array(arr)) => arr
.iter()
.filter_map(|s| s.as_str())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect(),
Some(Value::String(s)) => s
.split(',')
.map(|p| p.trim().to_string())
.filter(|p| !p.is_empty())
.collect(),
_ => Vec::new(),
}
}
/// Build an [`EmailMessage`] from a raw slim-envelope JSON message.
/// Returns `None` when the message has no parseable date — the rest of
/// the pipeline can't sort or canonicalise without one.
pub(crate) fn raw_to_email_message(raw: &Value) -> Option<EmailMessage> {
let id = raw
.get("id")
.and_then(|v| v.as_str())
.filter(|s| !s.is_empty())
.unwrap_or("");
let from = raw
.get("from")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let to = parse_address_list(raw.get("to"));
let cc = parse_address_list(raw.get("cc"));
let subject = raw
.get("subject")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let sent_at = parse_message_date(raw)?;
let body = raw
.get("markdown")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string();
let source_ref = if id.is_empty() {
None
} else {
Some(format!("gmail://msg/{id}"))
};
Some(EmailMessage {
from,
to,
cc,
subject,
sent_at,
body,
source_ref,
})
}
/// Parse the `to` / `cc` field which Composio surfaces as either a
/// JSON array of strings or a single comma-separated string. Empty
/// entries are dropped.
fn parse_address_list(v: Option<&Value>) -> Vec<String> {
match v {
Some(Value::Array(arr)) => arr
.iter()
.filter_map(|s| s.as_str())
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect(),
Some(Value::String(s)) => s
.split(',')
.map(|p| p.trim().to_string())
.filter(|p| !p.is_empty())
.collect(),
_ => Vec::new(),
}
}
/// Ingest a page of raw Gmail messages into the memory tree.
///
/// Each participant-bucket (sorted set of `from` `to` email addresses)
/// becomes one [`EmailThread`] handed to [`ingest_email`] which fans out
/// to the chunker + scorer + source tree downstream.
///
/// `source_id` = `"gmail:{participants}"` where participants is
/// `addr1|addr2|...` (sorted, deduped, lowercased). This groups all
/// correspondence between the same people into one path subtree.
///
/// Returns the total number of chunks written across all buckets so
/// callers can surface counts in logs / outcomes. Per-bucket errors are
/// logged and swallowed — one bad bucket should not abort the whole
/// page (the next sync re-fetches via the date-cursor).
pub async fn ingest_page_into_memory_tree(
config: &Config,
owner: &str,
page_messages: &[Value],
) -> Result<usize> {
if page_messages.is_empty() {
return Ok(0);
}
let buckets = bucket_by_participants(page_messages);
let mut total_chunks = 0usize;
let mut total_buckets = 0usize;
for (participants, raw_msgs) in &buckets {
let messages: Vec<EmailMessage> = raw_msgs
.iter()
.filter_map(|raw| raw_to_email_message(raw))
.collect();
if messages.is_empty() {
log::debug!(
"[composio:gmail][ingest] skipping empty bucket participants_hash={}",
redact(participants)
);
continue;
}
// source_id encodes participants so every unique conversation set
// lands in its own path subtree.
let source_id = format!("gmail:{}", participants);
let thread_subject = pick_thread_subject(&messages);
log::info!(
"[composio:gmail][ingest] bucket participants_hash={} messages={} source_id_hash={}",
redact(participants),
messages.len(),
redact(&source_id)
);
let thread = EmailThread {
provider: GMAIL_PROVIDER.to_string(),
thread_subject,
messages,
};
let tags = DEFAULT_TAGS.iter().map(|s| (*s).to_string()).collect();
match ingest_email(config, &source_id, owner, tags, thread).await {
Ok(IngestResult { chunks_written, .. }) => {
total_chunks += chunks_written;
total_buckets += 1;
}
Err(e) => {
log::warn!(
"[composio:gmail][ingest] ingest_email failed participants_hash={} source_id_hash={} err={:#}",
redact(participants),
redact(&source_id),
e
);
}
}
}
log::info!(
"[composio:gmail][ingest] page_done owner_hash={} buckets={total_buckets} chunks={total_chunks}",
redact(owner)
);
Ok(total_chunks)
}
/// Strip "Re:" / "Fwd:" prefixes from the head message's subject so
/// every message in a thread shares one canonical thread subject. Falls
/// back to "(no subject)" when empty.
fn pick_thread_subject(messages: &[EmailMessage]) -> String {
let raw = messages
.first()
.map(|m| m.subject.trim().to_string())
.unwrap_or_default();
let stripped = strip_reply_prefixes(&raw);
if stripped.is_empty() {
"(no subject)".to_string()
} else {
stripped
}
}
/// Iteratively strip `Re:` / `Fwd:` / `Fw:` prefixes (case-insensitive,
/// optional whitespace) from the front of a subject. Stops once a pass
/// removes nothing.
fn strip_reply_prefixes(subject: &str) -> String {
let mut s = subject.trim().to_string();
loop {
let lower = s.to_ascii_lowercase();
let stripped = if lower.starts_with("re:") {
Some(&s[3..])
} else if lower.starts_with("fwd:") {
Some(&s[4..])
} else if lower.starts_with("fw:") {
Some(&s[3..])
} else {
None
};
match stripped {
Some(rest) => {
let trimmed = rest.trim_start().to_string();
if trimmed == s {
return s;
}
s = trimmed;
}
None => return s,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
// ─── bucket_by_participants tests ─────────────────────────────────────────
#[test]
fn bidirectional_messages_bucket_together() {
// alice→bob and bob→alice land in the same key "alice@x.com|bob@y.com".
let msgs = vec![
json!({
"id": "m1",
"from": "alice@x.com",
"to": "bob@y.com",
"subject": "Hi",
"date": "2026-04-21T10:00:00Z",
"markdown": "hi",
}),
json!({
"id": "m2",
"from": "bob@y.com",
"to": "alice@x.com",
"subject": "Re: Hi",
"date": "2026-04-21T11:00:00Z",
"markdown": "hey",
}),
];
let buckets = bucket_by_participants(&msgs);
assert_eq!(buckets.len(), 1, "both messages must share one bucket");
let key = buckets.keys().next().unwrap();
assert_eq!(key, "alice@x.com|bob@y.com");
assert_eq!(buckets[key].len(), 2);
// Sorted ascending by date inside the bucket.
assert_eq!(buckets[key][0].get("id").unwrap().as_str().unwrap(), "m1");
assert_eq!(buckets[key][1].get("id").unwrap().as_str().unwrap(), "m2");
}
#[test]
fn multi_recipient_bucket_key_sorted() {
// from=alice, to=[bob, carol] → "alice@x.com|bob@y.com|carol@z.com"
let msgs = vec![json!({
"id": "m1",
"from": "Alice <alice@x.com>",
"to": ["bob@y.com", "carol@z.com"],
"subject": "Group",
"date": "2026-04-21T10:00:00Z",
"markdown": "hey all",
})];
let buckets = bucket_by_participants(&msgs);
let key = buckets.keys().next().unwrap();
assert_eq!(key, "alice@x.com|bob@y.com|carol@z.com");
}
#[test]
fn cc_field_ignored_in_bucket_key() {
// from=alice, to=[bob], cc=[dave] → "alice@x.com|bob@y.com" (no dave).
let msgs = vec![json!({
"id": "m1",
"from": "alice@x.com",
"to": "bob@y.com",
"cc": "dave@z.com",
"subject": "CC test",
"date": "2026-04-21T10:00:00Z",
"markdown": "body",
})];
let buckets = bucket_by_participants(&msgs);
let key = buckets.keys().next().unwrap();
assert_eq!(
key, "alice@x.com|bob@y.com",
"CC must not appear in bucket key"
);
}
#[test]
fn solo_message_no_to_buckets_to_sender_only() {
// from=alice, to=[] → "alice@x.com" (single participant).
let msgs = vec![json!({
"id": "m1",
"from": "alice@x.com",
"subject": "Draft",
"date": "2026-04-21T10:00:00Z",
"markdown": "draft body",
})];
let buckets = bucket_by_participants(&msgs);
let key = buckets.keys().next().unwrap();
assert_eq!(key, "alice@x.com");
}
#[test]
fn empty_from_and_to_falls_back_to_orphan_bucket() {
// A message with no parseable addresses gets its own orphan bucket
// keyed by its id rather than collapsing everything into "unknown".
let msgs = vec![json!({
"id": "m1",
"from": "",
"subject": "x",
"date": "2026-04-21T10:00:00Z",
"markdown": "body",
})];
let buckets = bucket_by_participants(&msgs);
assert_eq!(buckets.len(), 1, "must produce exactly one bucket");
assert!(
buckets.contains_key("orphan:m1"),
"must fall back to orphan:<id>; got keys: {:?}",
buckets.keys().collect::<Vec<_>>()
);
}
#[test]
fn two_malformed_messages_with_different_ids_land_in_different_buckets() {
// Two messages with unparseable from/to but different ids must not
// collapse into the same "unknown" bucket — each gets its own orphan.
let msgs = vec![
json!({
"id": "orphan_a",
"from": "",
"subject": "x",
"date": "2026-04-21T10:00:00Z",
"markdown": "body a",
}),
json!({
"id": "orphan_b",
"from": "",
"subject": "y",
"date": "2026-04-21T11:00:00Z",
"markdown": "body b",
}),
];
let buckets = bucket_by_participants(&msgs);
assert_eq!(
buckets.len(),
2,
"each malformed message must have its own bucket; got: {:?}",
buckets.keys().collect::<Vec<_>>()
);
assert!(buckets.contains_key("orphan:orphan_a"));
assert!(buckets.contains_key("orphan:orphan_b"));
}
#[test]
fn message_with_no_id_and_no_addresses_is_dropped() {
// A message with no id AND no parseable addresses is silently dropped.
let valid = json!({
"id": "m_ok",
"from": "alice@x.com",
"subject": "ok",
"date": "2026-04-21T10:00:00Z",
"markdown": "ok",
});
let bad = json!({
// no "id" field, no from/to
"subject": "bad",
"date": "2026-04-21T10:00:00Z",
"markdown": "bad",
});
let msgs = vec![valid, bad];
let buckets = bucket_by_participants(&msgs);
// Only the valid message should produce a bucket.
assert_eq!(buckets.len(), 1, "dropped message must not create a bucket");
assert!(buckets.contains_key("alice@x.com"));
}
#[test]
fn display_name_from_stripped_to_bare_email_in_key() {
// "Alice <alice@x.com>" should yield bare "alice@x.com" in the key.
let msgs = vec![json!({
"id": "m1",
"from": "Alice <alice@x.com>",
"to": "Bob <bob@y.com>",
"subject": "Hi",
"date": "2026-04-21T10:00:00Z",
"markdown": "hi",
})];
let buckets = bucket_by_participants(&msgs);
let key = buckets.keys().next().unwrap();
assert_eq!(key, "alice@x.com|bob@y.com");
}
#[test]
fn no_threadid_field_does_not_affect_bucketing() {
// threadId is completely ignored; two messages from the same participants
// share one bucket even without threadId.
let msgs = vec![
json!({
"id": "m1",
"from": "noreply@github.com",
"to": "sanil@x.com",
"subject": "PR opened",
"date": "2026-04-21T10:00:00Z",
"markdown": "body1",
}),
json!({
"id": "m2",
"from": "noreply@github.com",
"to": "sanil@x.com",
"subject": "PR merged",
"date": "2026-04-21T11:00:00Z",
"markdown": "body2",
}),
];
let buckets = bucket_by_participants(&msgs);
assert_eq!(buckets.len(), 1, "both messages must share one bucket");
let bucket = buckets.values().next().unwrap();
assert_eq!(bucket.len(), 2);
}
#[test]
fn raw_to_email_message_parses_slim_envelope() {
let raw = json!({
"id": "m1",
"from": "Alice <alice@example.com>",
"to": "me@example.com",
"cc": "team@example.com",
"subject": "Phoenix kickoff",
"date": "2026-04-21T10:00:00Z",
"markdown": "Let's ship Phoenix.",
});
let msg = raw_to_email_message(&raw).unwrap();
assert_eq!(msg.from, "Alice <alice@example.com>");
assert_eq!(msg.to, vec!["me@example.com"]);
assert_eq!(msg.cc, vec!["team@example.com"]);
assert_eq!(msg.subject, "Phoenix kickoff");
assert_eq!(msg.body, "Let's ship Phoenix.");
assert_eq!(msg.source_ref.as_deref(), Some("gmail://msg/m1"));
}
#[test]
fn raw_to_email_message_handles_to_array() {
let raw = json!({
"id": "m1",
"from": "a@x",
"to": ["b@x", "c@x"],
"subject": "x",
"date": "2026-04-21T10:00:00Z",
"markdown": "body",
});
let msg = raw_to_email_message(&raw).unwrap();
assert_eq!(msg.to, vec!["b@x", "c@x"]);
}
#[test]
fn raw_to_email_message_handles_comma_separated_to_string() {
let raw = json!({
"id": "m1",
"from": "a@x",
"to": "b@x, c@x ,d@x",
"subject": "x",
"date": "2026-04-21T10:00:00Z",
"markdown": "body",
});
let msg = raw_to_email_message(&raw).unwrap();
assert_eq!(msg.to, vec!["b@x", "c@x", "d@x"]);
}
#[test]
fn raw_to_email_message_returns_none_on_unparseable_date() {
let raw = json!({
"id": "m1",
"from": "a@x",
"subject": "x",
"date": "not-a-date",
"markdown": "body",
});
assert!(raw_to_email_message(&raw).is_none());
}
#[test]
fn raw_to_email_message_drops_source_ref_when_id_empty() {
let raw = json!({
"id": "",
"from": "a@x",
"subject": "x",
"date": "2026-04-21T10:00:00Z",
"markdown": "body",
});
let msg = raw_to_email_message(&raw).unwrap();
assert!(msg.source_ref.is_none());
}
#[test]
fn strip_reply_prefixes_removes_iterated() {
assert_eq!(strip_reply_prefixes("Re: Re: Hi"), "Hi");
assert_eq!(strip_reply_prefixes("Fwd: Re: Status"), "Status");
assert_eq!(strip_reply_prefixes("RE: Question"), "Question");
assert_eq!(strip_reply_prefixes("Fw: alert"), "alert");
assert_eq!(strip_reply_prefixes("Plain subject"), "Plain subject");
}
#[test]
fn pick_thread_subject_strips_reply_prefixes() {
let messages = vec![EmailMessage {
from: "a@x".into(),
to: vec![],
cc: vec![],
subject: "Re: Re: Phoenix kickoff".into(),
sent_at: chrono::Utc::now(),
body: "body".into(),
source_ref: None,
}];
assert_eq!(pick_thread_subject(&messages), "Phoenix kickoff");
}
#[test]
fn pick_thread_subject_falls_back_to_no_subject() {
let messages = vec![EmailMessage {
from: "a@x".into(),
to: vec![],
cc: vec![],
subject: " ".into(),
sent_at: chrono::Utc::now(),
body: "body".into(),
source_ref: None,
}];
assert_eq!(pick_thread_subject(&messages), "(no subject)");
}
}
@@ -1,3 +1,4 @@
pub mod ingest;
mod post_process;
mod provider;
mod sync;
+11
View File
@@ -980,6 +980,17 @@ impl Config {
}
}
// Phase MD-content: chunk body directory override. Empty string means
// "fall back to default", consistent with other memory_tree env vars.
if let Ok(dir) = std::env::var("OPENHUMAN_MEMORY_TREE_CONTENT_DIR") {
let trimmed = dir.trim();
self.memory_tree.content_dir = if trimmed.is_empty() {
None
} else {
Some(std::path::PathBuf::from(trimmed))
};
}
// Auto-update overrides
if let Some(flag) = env.get("OPENHUMAN_AUTO_UPDATE_ENABLED") {
let normalized = flag.trim().to_ascii_lowercase();
@@ -2,6 +2,7 @@
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
#[derive(Debug, Clone, Serialize, Deserialize, Default, JsonSchema)]
pub struct StorageConfig {
@@ -96,6 +97,7 @@ impl Default for MemoryConfig {
/// - `OPENHUMAN_MEMORY_SUMMARISE_ENDPOINT`
/// - `OPENHUMAN_MEMORY_SUMMARISE_MODEL`
/// - `OPENHUMAN_MEMORY_SUMMARISE_TIMEOUT_MS`
/// - `OPENHUMAN_MEMORY_TREE_CONTENT_DIR` (Phase MD-content)
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct MemoryTreeConfig {
/// Ollama endpoint for the embedder (e.g. `http://localhost:11434`).
@@ -154,6 +156,17 @@ pub struct MemoryTreeConfig {
/// tokens and therefore takes longer to generate.
#[serde(default = "default_memory_tree_llm_summariser_timeout_ms")]
pub llm_summariser_timeout_ms: Option<u64>,
/// Phase MD-content: root directory where chunk `.md` files are stored.
///
/// Resolved at runtime via [`super::types::Config::memory_tree_content_root`]:
/// - `Some(path)` → use that path verbatim.
/// - `None` → default `<workspace_dir>/memory_tree/content/`.
///
/// Env override: `OPENHUMAN_MEMORY_TREE_CONTENT_DIR` (empty string = fall
/// back to default, consistent with other memory_tree env vars).
#[serde(default = "default_memory_tree_content_dir")]
pub content_dir: Option<PathBuf>,
}
/// Returns `None` so that existing installs that never opted into Phase 4
@@ -199,6 +212,12 @@ fn default_memory_tree_llm_summariser_timeout_ms() -> Option<u64> {
Some(120_000)
}
/// Returns `None` so the default `<workspace>/memory_tree/content/` path is
/// used unless explicitly overridden via TOML or env var.
fn default_memory_tree_content_dir() -> Option<PathBuf> {
None
}
impl Default for MemoryTreeConfig {
fn default() -> Self {
Self {
@@ -212,6 +231,43 @@ impl Default for MemoryTreeConfig {
llm_summariser_endpoint: default_memory_tree_llm_endpoint(),
llm_summariser_model: default_memory_tree_llm_endpoint(),
llm_summariser_timeout_ms: default_memory_tree_llm_summariser_timeout_ms(),
content_dir: default_memory_tree_content_dir(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn memory_tree_config_default_content_dir_is_none() {
let cfg = MemoryTreeConfig::default();
assert!(
cfg.content_dir.is_none(),
"default content_dir must be None so workspace default path is used"
);
}
/// Verify that the env-var override logic correctly maps non-empty strings
/// to `Some(PathBuf)` and empty/blank strings to `None`. We test the
/// logic inline (not via `apply_env_overrides`) to avoid mutating the
/// process environment in a way that could race with parallel tests.
#[test]
fn content_dir_env_override_logic() {
// Simulate the load.rs overlay logic.
let apply = |raw: &str| -> Option<PathBuf> {
let trimmed = raw.trim();
if trimmed.is_empty() {
None
} else {
Some(PathBuf::from(trimmed))
}
};
assert_eq!(apply("/tmp/foo"), Some(PathBuf::from("/tmp/foo")));
assert_eq!(apply(" /tmp/foo "), Some(PathBuf::from("/tmp/foo")));
assert_eq!(apply(""), None);
assert_eq!(apply(" "), None);
}
}
+17
View File
@@ -203,6 +203,23 @@ pub struct Config {
pub chat_onboarding_completed: bool,
}
impl Config {
/// Resolve the root directory where chunk `.md` files are stored.
///
/// Resolution order:
/// 1. `memory_tree.content_dir` if `Some`.
/// 2. Default: `<workspace_dir>/memory_tree/content/`.
///
/// This is the only place in the codebase that should compute the content
/// root — all code that needs the path should call this method.
pub fn memory_tree_content_root(&self) -> PathBuf {
self.memory_tree
.content_dir
.clone()
.unwrap_or_else(|| self.workspace_dir.join("memory_tree").join("content"))
}
}
impl Default for Config {
fn default() -> Self {
let openhuman_dir =
+15 -12
View File
@@ -4,10 +4,9 @@
//! from the same channel becomes one [`CanonicalisedSource`]; the chunker
//! slices it by token budget downstream.
//!
//! Output format:
//! Output format (no leading `# ...` header — that info lives in front-matter
//! once Phase MD-content lands; the chunker splits at `## ` boundaries):
//! ```md
//! # Chat transcript — {platform} / {channel}
//!
//! ## 2026-04-21T10:12:00Z — Alice
//! Message body here.
//!
@@ -67,10 +66,9 @@ pub fn canonicalise(
let last_ts = messages.last().map(|m| m.timestamp).unwrap();
let mut md = String::new();
md.push_str(&format!(
"# Chat transcript — {} / {}\n\n",
batch.platform, batch.channel_label
));
// No leading `# Chat transcript — ...` header. Platform / channel info
// belongs in the MD front-matter (Phase MD-content). The chunker splits
// this output at `## ` boundaries so each message becomes one chunk.
for msg in &messages {
md.push_str(&format!(
"## {}{}\n{}\n\n",
@@ -150,7 +148,7 @@ mod tests {
}
#[test]
fn includes_header_and_per_message_sections() {
fn includes_per_message_sections_without_header() {
let b = ChatBatch {
platform: "slack".into(),
channel_label: "#eng".into(),
@@ -159,10 +157,15 @@ mod tests {
let out = canonicalise("slack:#eng", "alice", &[], b)
.unwrap()
.unwrap();
assert!(out
.markdown
.starts_with("# Chat transcript — slack / #eng\n\n"));
assert!(out.markdown.contains("## "));
// No leading `# Chat transcript` header — that info belongs in front-matter.
assert!(
!out.markdown.starts_with("# "),
"canonical chat MD must NOT start with a `# ` header"
);
assert!(
out.markdown.starts_with("## "),
"must start with first `## ` message block"
);
assert!(out.markdown.contains("— alice"));
assert!(out.markdown.contains("hello"));
}
@@ -39,7 +39,8 @@ pub fn canonicalise(
}
let mut md = String::new();
md.push_str(&format!("# {}{}\n\n", doc.provider, doc.title));
// No leading `# provider — title` header. Provider / title info
// belongs in the MD front-matter (Phase MD-content).
md.push_str(doc.body.trim());
md.push('\n');
@@ -85,7 +86,7 @@ mod tests {
}
#[test]
fn renders_title_and_body() {
fn renders_body_without_header() {
let out = canonicalise(
"d1",
"alice",
@@ -94,7 +95,11 @@ mod tests {
)
.unwrap()
.unwrap();
assert!(out.markdown.starts_with("# notion — Launch plan\n\n"));
// No leading `# notion — Launch plan` header — that info belongs in front-matter.
assert!(
!out.markdown.starts_with("# "),
"canonical document MD must NOT start with a `# ` header"
);
assert!(out.markdown.contains("step one"));
assert!(out.markdown.contains("step two"));
}
+73 -13
View File
@@ -1,13 +1,16 @@
//! Email threads → canonical Markdown.
//!
//! Email sources are scoped by **thread**. One thread becomes one
//! [`CanonicalisedSource`]. Headers (From, To, Subject, Date) surface in a
//! small frontmatter-style block per message; the body follows as markdown.
//! Email sources are scoped by **participant set**. One participant bucket
//! becomes one [`CanonicalisedSource`]. Headers (From, To, Cc, Subject, Date)
//! surface in a small frontmatter-style block per message; the cleaned body
//! follows as markdown. Bodies pass through [`email_clean::clean_body`] before
//! rendering to strip reply chains, marketing footers, legal disclaimers, and
//! other boilerplate.
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use super::{normalize_source_ref, CanonicalisedSource};
use super::{email_clean, normalize_source_ref, CanonicalisedSource};
use crate::openhuman::memory::tree::types::{Metadata, SourceKind};
/// One email in a thread.
@@ -55,10 +58,9 @@ pub fn canonicalise(
let last_ts = messages.last().map(|m| m.sent_at).unwrap();
let mut md = String::new();
md.push_str(&format!(
"# Email thread — {}{}\n\n",
thread.provider, thread.thread_subject
));
// No leading `# Email thread — ...` header. Provider / subject info
// belongs in the MD front-matter (Phase MD-content). The chunker splits
// this output at `---\nFrom:` boundaries so each message becomes one chunk.
for msg in &messages {
md.push_str("---\n");
@@ -71,7 +73,12 @@ pub fn canonicalise(
}
md.push_str(&format!("Subject: {}\n", msg.subject));
md.push_str(&format!("Date: {}\n\n", msg.sent_at.to_rfc3339()));
md.push_str(msg.body.trim());
let cleaned = email_clean::clean_body(msg.body.trim());
if cleaned.is_empty() {
md.push('\n');
} else {
md.push_str(&cleaned);
}
md.push_str("\n\n");
}
@@ -128,10 +135,19 @@ mod tests {
email(2000, "alice@example.com", "Re: Launch", "agreed"),
],
};
let out = canonicalise("gmail:t1", "alice@example.com", &[], t)
.unwrap()
.unwrap();
assert!(out.markdown.contains("# Email thread — gmail — Launch"));
let out = canonicalise(
"gmail:alice@example.com|bob@example.com",
"alice@example.com",
&[],
t,
)
.unwrap()
.unwrap();
// No leading `# Email thread` header — that info belongs in front-matter.
assert!(
!out.markdown.contains("# Email thread — gmail — Launch"),
"canonical email MD must NOT contain a `# ` header"
);
assert!(out.markdown.contains("From: bob@example.com"));
assert!(out.markdown.contains("Subject: Launch"));
assert!(out.markdown.contains("let's ship"));
@@ -139,6 +155,50 @@ mod tests {
assert!(out.markdown.contains("agreed"));
}
#[test]
fn clean_body_strips_footer_before_canonicalise() {
// Body where "Unsubscribe" line triggers footer removal. Everything from
// that line onward is dropped by clean_body; real content above survives.
let body_with_footer =
"Please review the attached document.\n\nUnsubscribe https://mail.example.com/unsub\n© 2026 Example Corp";
let t = EmailThread {
provider: "gmail".into(),
thread_subject: "Review".into(),
messages: vec![EmailMessage {
from: "sender@example.com".into(),
to: vec!["recipient@example.com".into()],
cc: vec![],
subject: "Review".into(),
sent_at: Utc.timestamp_millis_opt(5000).unwrap(),
body: body_with_footer.into(),
source_ref: None,
}],
};
let out = canonicalise(
"gmail:recipient@example.com|sender@example.com",
"recipient@example.com",
&[],
t,
)
.unwrap()
.unwrap();
assert!(
out.markdown.contains("Please review the attached document"),
"real content must survive; got:\n{}",
out.markdown
);
assert!(
!out.markdown.to_ascii_lowercase().contains("unsubscribe"),
"unsubscribe footer must be stripped; got:\n{}",
out.markdown
);
assert!(
!out.markdown.contains("© 2026"),
"copyright footer must be stripped; got:\n{}",
out.markdown
);
}
#[test]
fn time_range_spans_thread() {
let t = EmailThread {
@@ -0,0 +1,382 @@
//! Shared email rendering + cleaning helpers.
//!
//! Used by both [`canonicalize::email`](super::email) (when rendering the
//! `GmailMarkdownStyle::Standard` shape) and the `gmail-fetch-emails` bin
//! (which writes per-sender markdown digests to disk). Lifted out of the
//! bin so the bin's behaviour and the production canonicaliser stay
//! byte-identical on body cleanup.
//!
//! The module is intentionally pure-string-oriented + a single
//! `serde_json::Value` helper (`parse_message_date`) used by callers that
//! work directly off Gmail's slim envelope JSON. Nothing here depends on
//! the memory-tree types — that keeps the helpers reusable.
use chrono::{DateTime, NaiveDate, Utc};
use serde_json::Value;
/// Two-stage cleanup applied to each message body before it gets
/// blockquoted into a digest:
///
/// 1. **Drop quoted reply chains** — once a message contains a
/// `On <date>, <name> wrote:` preamble, an `Original Message` /
/// `Forwarded message` separator, or a run of three+ consecutive
/// `>`-prefixed lines, everything from that point onward is the
/// parent message we already render directly above.
/// 2. **Drop footer noise** — `Unsubscribe`, `View in browser`,
/// copyright lines, legal disclaimers, and address blocks. We cut
/// at the first line containing any of [`FOOTER_TRIGGERS`].
///
/// The two passes run in order so a quoted-chain preamble below a
/// "view in browser" line still gets stripped on its own merits even
/// if the footer pass missed it.
pub fn clean_body(raw: &str) -> String {
let stage1 = drop_reply_chain(raw);
let stage2 = drop_footer_noise(&stage1);
collapse_blank_runs(stage2.trim())
}
/// Substrings that, when matched (case-insensitive) anywhere on a
/// line, mark the start of footer / boilerplate territory. Conservative
/// list — every entry should be unambiguous noise that wouldn't
/// reasonably appear inside real prose.
const FOOTER_TRIGGERS: &[&str] = &[
"unsubscribe",
"view in browser",
"view this email in your browser",
"view it in your browser",
"update your email settings",
"manage your subscription",
"manage preferences",
"email preferences",
"you are receiving this email because",
"you received this email because",
"you're receiving this email because",
"to stop receiving",
"all rights reserved",
"© 20",
"(c) 20",
"copyright 20",
"powered by mailchimp",
"sent via sendgrid",
"this email and any files",
"confidentiality notice",
"if you are not the intended recipient",
"this communication may contain",
];
/// Strip quoted reply chains. See [`clean_body`] for details.
pub fn drop_reply_chain(s: &str) -> String {
let mut offset = 0usize;
let mut quoted_run_start: Option<usize> = None;
let mut quoted_run_len = 0u32;
for line in s.split_inclusive('\n') {
let trimmed = line.trim();
let lower = trimmed.to_ascii_lowercase();
// Explicit reply / forward markers.
let is_preamble = (lower.starts_with("on ") && lower.contains(" wrote:"))
|| lower.contains("---------- forwarded message")
|| lower.contains("----- original message")
|| lower.contains("--------- original message")
|| lower.contains("--- forwarded by");
if is_preamble {
return s[..offset].trim_end().to_string();
}
// Three+ consecutive lines starting with `>` is a quoted
// reply chain in disguise (some clients de-quote on send).
// Treat the start of the run as the cut point.
if trimmed.starts_with('>') {
if quoted_run_start.is_none() {
quoted_run_start = Some(offset);
quoted_run_len = 1;
} else {
quoted_run_len += 1;
}
if quoted_run_len >= 3 {
let cut = quoted_run_start.unwrap_or(offset);
return s[..cut].trim_end().to_string();
}
} else if !trimmed.is_empty() {
// Reset on a non-empty, non-quoted line. Blank lines
// don't break a quote run because senders often interleave
// them.
quoted_run_start = None;
quoted_run_len = 0;
}
offset += line.len();
}
s.to_string()
}
/// Strip everything from the first line containing a footer trigger
/// onward. See [`FOOTER_TRIGGERS`] for the matched list.
pub fn drop_footer_noise(s: &str) -> String {
let mut offset = 0usize;
for line in s.split_inclusive('\n') {
let lower = line.to_ascii_lowercase();
if FOOTER_TRIGGERS.iter().any(|t| lower.contains(t)) {
return s[..offset].trim_end().to_string();
}
offset += line.len();
}
s.to_string()
}
/// Collapse runs of 2+ blank lines into a single blank line. Trims
/// trailing newlines.
pub fn collapse_blank_runs(s: &str) -> String {
let mut out = String::with_capacity(s.len());
let mut blank = 0u32;
for line in s.lines() {
if line.trim().is_empty() {
blank += 1;
if blank <= 1 {
out.push('\n');
}
} else {
blank = 0;
out.push_str(line);
out.push('\n');
}
}
while out.ends_with('\n') {
out.pop();
}
out
}
/// Truncate a body to at most `max_chars` characters, appending `…` when
/// the body is longer. Trims first so leading/trailing whitespace doesn't
/// count against the budget.
pub fn truncate_body(body: &str, max_chars: usize) -> String {
let trimmed = body.trim();
if trimmed.chars().count() <= max_chars {
return trimmed.to_string();
}
let mut out: String = trimmed.chars().take(max_chars).collect();
out.push('…');
out
}
/// Escape only the few markdown chars that would visibly break the
/// header/inline contexts we use (#, |, *, _, `). Newlines collapse to
/// spaces. We leave most punctuation alone — the body is rendered as a
/// blockquote anyway.
pub fn md_escape(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for ch in s.chars() {
match ch {
'\\' | '`' | '*' | '_' | '|' => {
out.push('\\');
out.push(ch);
}
'\n' | '\r' => out.push(' '),
_ => out.push(ch),
}
}
out
}
/// Pull the `<addr@host>` portion out of a `From` header, returning
/// just the bare email address. Falls back to `None` when no `<…>`
/// brackets exist; in that case the caller may use the raw From field.
pub fn extract_email(from: &str) -> Option<String> {
let s = from.trim();
if let (Some(start), Some(end)) = (s.rfind('<'), s.rfind('>')) {
if start < end {
let inner = s[start + 1..end].trim();
if inner.contains('@') {
return Some(inner.to_string());
}
}
}
if s.contains('@') && !s.contains(' ') {
return Some(s.to_string());
}
None
}
/// If `s` starts with a 3-letter day-of-week prefix (`Mon, `, `Tue, `, …),
/// return the remainder; otherwise `None`. Used to feed a strict-rfc2822
/// reject into a lenient retry.
fn strip_day_of_week_prefix(s: &str) -> Option<&str> {
const DAYS: &[&str] = &["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"];
let (prefix, rest) = s.split_once(", ")?;
if DAYS.iter().any(|d| d.eq_ignore_ascii_case(prefix)) {
Some(rest)
} else {
None
}
}
/// Try a sequence of common date formats. Composio's slim envelope sets
/// `date` from `messageTimestamp` (often ISO 8601 or epoch ms) when
/// present, falling back to the raw `Date:` header (RFC 2822). Operates
/// on the raw `serde_json::Value` so callers that work off the slim
/// envelope JSON don't have to reshape it first.
pub fn parse_message_date(m: &Value) -> Option<DateTime<Utc>> {
let raw = m.get("date")?;
if let Some(s) = raw.as_str() {
let s = s.trim();
if s.is_empty() {
return None;
}
// Epoch millis as a string?
if let Ok(ms) = s.parse::<i64>() {
return DateTime::from_timestamp_millis(ms);
}
if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
return Some(dt.with_timezone(&Utc));
}
if let Ok(dt) = DateTime::parse_from_rfc2822(s) {
return Some(dt.with_timezone(&Utc));
}
// Lenient RFC 2822 fallback: strict `parse_from_rfc2822` rejects
// mismatched day-of-week (e.g. `Mon, 21 Apr 2026 …` when Apr 21
// 2026 is a Tuesday). Real-world MTAs occasionally send these
// with a wrong day-name. Strip a `<DayName>, ` prefix and retry
// with the rfc2822 body format. Keeps the date if everything
// else is sane.
if let Some(rest) = strip_day_of_week_prefix(s) {
if let Ok(dt) = DateTime::parse_from_str(rest, "%d %b %Y %H:%M:%S %z") {
return Some(dt.with_timezone(&Utc));
}
}
if let Ok(d) = NaiveDate::parse_from_str(s, "%Y-%m-%d") {
return d.and_hms_opt(0, 0, 0).map(|n| n.and_utc());
}
}
if let Some(ms) = raw.as_i64() {
return DateTime::from_timestamp_millis(ms);
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn drop_reply_chain_strips_on_x_wrote_preamble() {
let body = "Sounds good — let's do Tuesday.\n\nOn Mon, Apr 22, 2026 at 10:00 AM, Alice <a@x> wrote:\n> Tuesday or Wednesday?\n> Let me know.";
let cleaned = drop_reply_chain(body);
assert_eq!(cleaned.trim(), "Sounds good — let's do Tuesday.");
}
#[test]
fn drop_reply_chain_strips_forwarded_separator() {
let body = "FYI.\n\n---------- Forwarded message ---------\nFrom: bob\nSubject: hi";
assert_eq!(drop_reply_chain(body).trim(), "FYI.");
}
#[test]
fn drop_reply_chain_strips_consecutive_quoted_run() {
let body = "Thanks for the update.\n\n> earlier line 1\n> earlier line 2\n> earlier line 3\n> earlier line 4";
assert_eq!(drop_reply_chain(body).trim(), "Thanks for the update.");
}
#[test]
fn drop_reply_chain_keeps_short_quote() {
// A single inline blockquote is fine — only 3+ consecutive lines trigger.
let body = "I think:\n> That sounds reasonable\n\nLet's proceed.";
let cleaned = drop_reply_chain(body);
assert!(cleaned.contains("Let's proceed"));
assert!(cleaned.contains("That sounds reasonable"));
}
#[test]
fn drop_footer_noise_strips_unsubscribe_block() {
let body =
"Big news: GPT-5.5 is here.\n\nRead more at example.com\n\nUnsubscribe | © 2026 OpenAI";
let cleaned = drop_footer_noise(body);
assert!(cleaned.contains("GPT-5.5"));
assert!(!cleaned.to_ascii_lowercase().contains("unsubscribe"));
assert!(!cleaned.contains("©"));
}
#[test]
fn drop_footer_noise_strips_legal_disclaimer() {
let body = "Action item — review by Friday.\n\nThis email and any files transmitted with it are confidential and intended solely for the use of the individual to whom they are addressed.";
let cleaned = drop_footer_noise(body);
assert_eq!(cleaned.trim(), "Action item — review by Friday.");
}
#[test]
fn clean_body_combines_passes() {
let body =
"Real content here.\n\nOn Mon, Apr 22, 2026, Alice wrote:\n> old stuff\n\nUnsubscribe";
let cleaned = clean_body(body);
assert_eq!(cleaned, "Real content here.");
}
#[test]
fn collapse_blank_runs_keeps_paragraph_breaks() {
let s = "a\n\n\n\nb\n\n\nc\n";
assert_eq!(collapse_blank_runs(s), "a\n\nb\n\nc");
}
#[test]
fn truncate_body_adds_ellipsis() {
let s = "x".repeat(2000);
let t = truncate_body(&s, 1200);
assert!(t.ends_with('…'));
assert_eq!(t.chars().count(), 1201);
}
#[test]
fn truncate_body_passthrough_when_short() {
let s = "hello";
let t = truncate_body(s, 1200);
assert_eq!(t, "hello");
}
#[test]
fn md_escape_handles_special_chars() {
assert_eq!(md_escape("a*b_c"), "a\\*b\\_c");
assert_eq!(md_escape("foo|bar"), "foo\\|bar");
assert_eq!(md_escape("line1\nline2"), "line1 line2");
assert_eq!(md_escape("plain text"), "plain text");
}
#[test]
fn extract_email_handles_both_forms() {
assert_eq!(
extract_email("Alice <alice@example.com>").as_deref(),
Some("alice@example.com")
);
assert_eq!(
extract_email("notify@github.com").as_deref(),
Some("notify@github.com")
);
assert_eq!(
extract_email("\"Bot Name\" <bot@x.io>").as_deref(),
Some("bot@x.io")
);
assert!(extract_email("Alice").is_none());
}
#[test]
fn parse_message_date_handles_iso_and_rfc2822() {
let iso = json!({"date": "2026-04-21T10:00:00Z"});
let rfc = json!({"date": "Mon, 21 Apr 2026 10:00:00 +0000"});
let ms = json!({"date": 1745236800000_i64});
let ms_str = json!({"date": "1745236800000"});
let date_only = json!({"date": "2026-04-21"});
assert!(parse_message_date(&iso).is_some());
assert!(parse_message_date(&rfc).is_some());
assert!(parse_message_date(&ms).is_some());
assert!(parse_message_date(&ms_str).is_some());
assert!(parse_message_date(&date_only).is_some());
}
#[test]
fn parse_message_date_returns_none_when_missing_or_blank() {
assert!(parse_message_date(&json!({})).is_none());
assert!(parse_message_date(&json!({"date": ""})).is_none());
assert!(parse_message_date(&json!({"date": " "})).is_none());
}
}
@@ -12,6 +12,7 @@
pub mod chat;
pub mod document;
pub mod email;
pub mod email_clean;
use serde::{Deserialize, Serialize};
+545 -80
View File
@@ -5,16 +5,26 @@
//! so later phases (#709 seal budget = 10k tokens) can ingest them without
//! blowing past the summariser ceiling.
//!
//! Splitting strategy: prefer paragraph boundaries (`\n\n`), then single
//! newlines, then hard character-count cut. We never split a character in half.
//! ## Dispatch by source kind (Phase B)
//!
//! - **Chat**: split at `## ` message boundaries. Each message becomes one
//! chunk. If a single message exceeds `max_tokens`, fall back to the
//! paragraph/line/char splitter for that unit only and emit each piece with
//! `partial_message = true`.
//! - **Email**: split at `---\nFrom:` separators. Each email in the thread
//! becomes one chunk. Same oversize fallback as Chat.
//! - **Document**: original paragraph-based greedy packing (unchanged).
use crate::openhuman::memory::tree::types::{approx_token_count, Chunk, Metadata, SourceKind};
use crate::openhuman::memory::tree::util::redact::redact;
/// Default upper bound on per-chunk tokens.
///
/// Aligned with the LLD's summariser 10k budget so a single chunk never blows
/// a seal on its own.
pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 10_000;
/// Sized below the L0 seal budget (`source_tree::types::TOKEN_BUDGET = 4_500`)
/// so each seal accumulates roughly 13 chunks before firing — natural pacing
/// for the local 1B summariser, which produces noticeably better summaries
/// with smaller (≤45k) inputs than at the previous 10k cap.
pub const DEFAULT_CHUNK_MAX_TOKENS: u32 = 3_000;
/// Tunable settings for the chunker.
#[derive(Clone, Debug)]
@@ -49,34 +59,255 @@ pub struct ChunkerInput {
/// Returns chunks in source order with stable sequence numbers starting at 0.
/// Chunk IDs are deterministic (`types::chunk_id`), so re-chunking yields the
/// same ids for identical input.
///
/// ## Dispatch by source kind
///
/// - **Chat / Email**: split at message/email boundaries, then greedy-pack
/// consecutive units into a single chunk until adding the next unit would
/// exceed `max_tokens`. Oversize units (a single message > `max_tokens`)
/// fall back to the paragraph/line/char splitter and emit each piece with
/// `partial_message = true`.
/// - **Document**: original paragraph-based greedy packing (unchanged).
pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec<Chunk> {
let now = chrono::Utc::now();
let pieces = split_by_token_budget(&input.markdown, opts.max_tokens);
let max_tokens = opts.max_tokens.max(1);
let max_chars = (max_tokens as usize).saturating_mul(4);
// Dispatch: pick splitting units based on source kind.
let units: Vec<String> = match input.source_kind {
SourceKind::Chat => split_chat_messages(&input.markdown),
SourceKind::Email => split_email_messages(&input.markdown),
SourceKind::Document => {
// Document: run the existing paragraph splitter directly on the
// whole blob. No message-unit concept.
log::debug!(
"[memory_tree::chunker] document source_id_hash={} len={} — paragraph split",
redact(&input.source_id),
input.markdown.len()
);
split_by_token_budget(&input.markdown, max_tokens)
}
};
if matches!(input.source_kind, SourceKind::Document) {
// Already split by budget; wrap directly.
return units
.into_iter()
.enumerate()
.map(|(idx, content)| {
let seq = idx as u32;
let token_count = approx_token_count(&content);
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &content);
Chunk {
id,
content,
metadata: input.metadata.clone(),
token_count,
seq_in_source: seq,
created_at: now,
partial_message: false,
}
})
.collect();
}
log::debug!(
"[memory_tree::chunker] source_kind={} source_id={} len={} pieces={}",
"[memory_tree::chunker] source_kind={} source_id_hash={} len={} units={}",
input.source_kind.as_str(),
input.source_id,
redact(&input.source_id),
input.markdown.len(),
pieces.len()
units.len()
);
pieces
.into_iter()
.enumerate()
.map(|(idx, content)| {
let seq = idx as u32;
let token_count = approx_token_count(&content);
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &content);
Chunk {
id,
content,
metadata: input.metadata.clone(),
token_count,
seq_in_source: seq,
created_at: now,
// For Chat and Email: greedy-pack consecutive units into chunks.
// Units are accumulated until adding the next would exceed max_chars;
// oversize single units fall back to sub-splitting with partial_message=true.
let unit_separator = "\n\n";
let sep_chars = unit_separator.chars().count();
let mut out: Vec<Chunk> = Vec::new();
let mut acc: Vec<String> = Vec::new();
let mut acc_chars = 0usize;
// Flush accumulated units as one packed chunk.
let flush = |acc: &mut Vec<String>, acc_chars: &mut usize, out: &mut Vec<Chunk>| {
if acc.is_empty() {
return;
}
let content = acc.join(unit_separator);
let seq = out.len() as u32;
let tc = approx_token_count(&content);
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &content);
out.push(Chunk {
id,
content,
metadata: input.metadata.clone(),
token_count: tc,
seq_in_source: seq,
created_at: now,
partial_message: false,
});
acc.clear();
*acc_chars = 0;
};
for unit in units {
let unit_chars = unit.chars().count();
if unit_chars > max_chars {
// Oversize: flush any pending accumulator first, then sub-split.
flush(&mut acc, &mut acc_chars, &mut out);
let sub_pieces = split_by_token_budget(&unit, max_tokens);
for piece in sub_pieces {
let seq = out.len() as u32;
let tc = approx_token_count(&piece);
let id = super::types::chunk_id(input.source_kind, &input.source_id, seq, &piece);
out.push(Chunk {
id,
content: piece,
metadata: input.metadata.clone(),
token_count: tc,
seq_in_source: seq,
created_at: now,
partial_message: true,
});
}
})
.collect()
continue;
}
// Compute projected size if we add this unit to the accumulator.
let projected = if acc.is_empty() {
unit_chars
} else {
acc_chars + sep_chars + unit_chars
};
if projected > max_chars {
// Adding this unit would overflow — flush the accumulator first.
flush(&mut acc, &mut acc_chars, &mut out);
}
if !acc.is_empty() {
acc_chars += sep_chars;
}
acc_chars += unit_chars;
acc.push(unit);
}
// Flush any remaining accumulated units.
flush(&mut acc, &mut acc_chars, &mut out);
if out.is_empty() {
// Degenerate: empty input → one empty chunk, matching original behaviour.
let id = super::types::chunk_id(input.source_kind, &input.source_id, 0, "");
out.push(Chunk {
id,
content: String::new(),
metadata: input.metadata.clone(),
token_count: 0,
seq_in_source: 0,
created_at: now,
partial_message: false,
});
}
out
}
/// Split a canonical chat blob into per-message units at `## ` boundaries.
///
/// Each returned string starts with `## ` and includes everything up to but
/// not including the next `## ` boundary. If the blob starts with a `# `
/// header (legacy or unexpected), everything before the first `## ` is
/// dropped silently.
fn split_chat_messages(md: &str) -> Vec<String> {
let mut pieces: Vec<String> = Vec::new();
let mut current: Option<String> = None;
for line in md.split_inclusive('\n') {
if line.starts_with("## ") {
if let Some(prev) = current.take() {
let trimmed = prev.trim_end().to_string();
if !trimmed.is_empty() {
pieces.push(trimmed);
}
}
current = Some(line.to_string());
} else if let Some(ref mut buf) = current {
buf.push_str(line);
}
// Lines before the first `## ` (e.g. a leading `# ` header) are dropped.
}
if let Some(prev) = current.take() {
let trimmed = prev.trim_end().to_string();
if !trimmed.is_empty() {
pieces.push(trimmed);
}
}
if pieces.is_empty() && !md.trim().is_empty() {
// No `## ` found at all — treat whole blob as one unit.
pieces.push(md.trim_end().to_string());
}
pieces
}
/// Split a canonical email thread blob into per-email units.
///
/// Splits at `---` (alone on a line, optional trailing whitespace) followed
/// by a `From:` line within the next 8 lines. Each piece includes the `---`
/// separator and everything up to but not including the next `---\nFrom:`
/// boundary. Content before the first `---` separator is dropped (handles
/// any leading header that might have slipped through).
fn split_email_messages(md: &str) -> Vec<String> {
let lines: Vec<&str> = md.split('\n').collect();
let n = lines.len();
let mut split_positions: Vec<usize> = Vec::new();
for i in 0..n {
let line = lines[i].trim_end();
if line == "---" {
// Check if one of the next 8 lines starts with `From:`
let window_end = (i + 9).min(n);
for j in (i + 1)..window_end {
if lines[j].starts_with("From:") {
split_positions.push(i);
break;
}
// Skip blank lines between `---` and `From:`
if !lines[j].trim().is_empty() {
break;
}
}
}
}
if split_positions.is_empty() {
// No email separator found — treat whole blob as one unit.
let trimmed = md.trim_end().to_string();
if trimmed.is_empty() {
return Vec::new();
}
return vec![trimmed];
}
let mut pieces: Vec<String> = Vec::new();
for (idx, &start) in split_positions.iter().enumerate() {
let end = if idx + 1 < split_positions.len() {
split_positions[idx + 1]
} else {
n
};
let piece_lines: Vec<&str> = lines[start..end].iter().copied().collect();
let piece = piece_lines.join("\n").trim_end().to_string();
if !piece.is_empty() {
pieces.push(piece);
}
}
pieces
}
/// Split `text` into pieces each ≤ `max_tokens` tokens.
@@ -85,7 +316,7 @@ pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec<Chunk>
/// 1. Paragraph (`\n\n`)
/// 2. Line (`\n`)
/// 3. Hard character cut (last resort; preserves UTF-8 code points)
fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec<String> {
pub(crate) fn split_by_token_budget(text: &str, max_tokens: u32) -> Vec<String> {
let max_tokens = max_tokens.max(1);
if text.is_empty() {
return vec![String::new()];
@@ -187,22 +418,33 @@ mod tests {
Metadata::point_in_time(SourceKind::Chat, "slack:#eng", "alice", Utc::now())
}
fn meta_email() -> Metadata {
Metadata::point_in_time(SourceKind::Email, "gmail:t1", "alice", Utc::now())
}
fn meta_doc() -> Metadata {
Metadata::point_in_time(SourceKind::Document, "doc1", "alice", Utc::now())
}
#[test]
fn tiny_input_produces_single_chunk() {
// Chat input without a `## ` header produces one chunk via the empty-
// result fallback (whole blob as one unit).
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
markdown: "hello world".into(),
markdown: "## 2026-01-01T00:00:00Z — alice\nhello world".into(),
metadata: meta(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].content, "hello world");
assert!(chunks[0].content.contains("hello world"));
assert_eq!(chunks[0].seq_in_source, 0);
assert!(!chunks[0].partial_message);
}
#[test]
fn empty_input_produces_one_empty_chunk() {
fn empty_chat_input_produces_one_empty_chunk() {
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "x".into(),
@@ -212,10 +454,275 @@ mod tests {
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(chunks.len(), 1);
assert_eq!(chunks[0].content, "");
assert!(!chunks[0].partial_message);
}
#[test]
fn paragraph_boundaries_preferred() {
fn chat_messages_pack_into_one_chunk_when_small() {
// Two small chat messages both fit under default max_tokens → greedy
// packing emits ONE chunk containing both, joined by \n\n.
let md = "## 2026-01-01T00:00:00Z — alice\nHello world\n\n## 2026-01-01T00:01:00Z — bob\nParagraph one.\n\nParagraph two.".to_string();
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
markdown: md.clone(),
metadata: meta(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
// Both small messages fit under 10k tokens → one packed chunk.
assert_eq!(
chunks.len(),
1,
"small messages should be packed into one chunk; got {chunks:?}"
);
assert!(
chunks[0].content.contains("alice"),
"chunk must contain alice's message"
);
assert!(
chunks[0].content.contains("bob"),
"chunk must contain bob's message"
);
assert!(chunks[0].content.contains("Paragraph one."));
assert!(chunks[0].content.contains("Paragraph two."));
assert!(!chunks[0].partial_message);
}
#[test]
fn chat_messages_split_at_boundary_when_large() {
// Messages that together exceed max_tokens split at message boundaries
// into multiple chunks. Each chunk contains whole messages only.
// Each message is ~3k tokens at 4 chars/token = 12k chars;
// two messages = ~6k tokens > 5k budget → must split.
let msg_body = "x".repeat(12_000);
let md = format!(
"## 2026-01-01T00:00:00Z — alice\n{msg_body}\n\n## 2026-01-01T00:01:00Z — bob\n{msg_body}"
);
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
markdown: md,
metadata: meta(),
};
// Use a 5k token budget so two ~3k-token messages don't fit together.
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 5_000 });
assert_eq!(
chunks.len(),
2,
"two large messages should land in separate chunks; got {chunks:?}"
);
assert!(chunks[0].content.contains("alice"));
assert!(chunks[1].content.contains("bob"));
for c in &chunks {
assert!(!c.partial_message, "whole messages must not be partial");
}
}
#[test]
fn email_threads_pack_into_one_chunk_when_small() {
// Three short emails all fit under default max_tokens → one packed chunk.
let md = "---\nFrom: alice@example.com\nSubject: Hello\nDate: 2026-01-01T00:00:00Z\n\nFirst body.\n---\nFrom: bob@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:01:00Z\n\nSecond body.\n---\nFrom: carol@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:02:00Z\n\nThird body.".to_string();
let input = ChunkerInput {
source_kind: SourceKind::Email,
source_id: "gmail:t1".into(),
markdown: md,
metadata: meta_email(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(
chunks.len(),
1,
"three small emails should pack into one chunk; got {chunks:?}"
);
assert!(chunks[0].content.contains("First body."));
assert!(chunks[0].content.contains("Second body."));
assert!(chunks[0].content.contains("Third body."));
assert!(!chunks[0].partial_message);
}
#[test]
fn email_thread_large_splits_at_email_boundaries() {
// Messages totaling >12k tokens split into 2 chunks at email boundaries.
// Each email is ~4k tokens (16k chars); 3 emails × 4k = 12k tokens.
// With a 5k budget, 2 emails fit per chunk → 2 chunks for 3 emails.
let email_body = "y".repeat(16_000); // ~4k tokens
let md = format!(
"---\nFrom: a@x.com\nDate: 2026-01-01T00:00:00Z\n\n{email_body}\n\
---\nFrom: b@x.com\nDate: 2026-01-01T00:01:00Z\n\n{email_body}\n\
---\nFrom: c@x.com\nDate: 2026-01-01T00:02:00Z\n\n{email_body}"
);
let input = ChunkerInput {
source_kind: SourceKind::Email,
source_id: "gmail:t1".into(),
markdown: md,
metadata: meta_email(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 5_000 });
assert!(
chunks.len() >= 2,
"large thread must split into multiple chunks; got {}",
chunks.len()
);
for c in &chunks {
assert!(!c.partial_message, "whole-email chunks must not be partial");
}
}
#[test]
fn oversize_single_email_splits_with_partial_flag() {
// A single email body > max_tokens must produce partial_message=true pieces.
let big_body = "z".repeat(50_000); // ~12.5k tokens at 4 chars/token
let md = format!("---\nFrom: a@x.com\nDate: 2026-01-01T00:00:00Z\n\n{big_body}");
let input = ChunkerInput {
source_kind: SourceKind::Email,
source_id: "gmail:t1".into(),
markdown: md,
metadata: meta_email(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 1_000 });
assert!(chunks.len() > 1, "oversize email must split");
for c in &chunks {
assert!(
c.partial_message,
"all sub-pieces of an oversize email must have partial_message=true"
);
}
}
#[test]
fn packed_units_joined_by_double_newline() {
// Two chat messages packed together must be separated by \n\n.
let md = "## 2026-01-01T00:00:00Z — alice\nfoo\n\n## 2026-01-01T00:01:00Z — bob\nbar"
.to_string();
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "x".into(),
markdown: md,
metadata: meta(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(chunks.len(), 1);
// The two messages must be separated by \n\n in the packed content.
assert!(
chunks[0].content.contains("\n\n"),
"packed units must be joined by \\n\\n; content={:?}",
chunks[0].content
);
}
#[test]
fn oversize_message_falls_back_with_partial_flag() {
// Single chat message that is way over max_tokens.
let long_body = "x".repeat(8000); // ~2000 tokens at 4 chars/token
let md = format!("## 2026-01-01T00:00:00Z — alice\n{long_body}");
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "x".into(),
markdown: md,
metadata: meta(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 100 });
assert!(chunks.len() > 1, "oversize message must split");
for c in &chunks {
assert!(
c.partial_message,
"all sub-pieces of an oversize message must have partial_message=true"
);
}
// Reuniting all pieces must reconstruct the message content (minus `## ` line).
let rejoined: String = chunks.iter().map(|c| c.content.as_str()).collect();
assert!(rejoined.contains(&long_body[..100]));
}
#[test]
fn document_falls_through_to_paragraph_split() {
let para1 = "a".repeat(400); // ~100 tokens
let para2 = "b".repeat(400);
let para3 = "c".repeat(400);
let text = format!("{para1}\n\n{para2}\n\n{para3}");
let input = ChunkerInput {
source_kind: SourceKind::Document,
source_id: "doc1".into(),
markdown: text,
metadata: meta_doc(),
};
let chunks = chunk_markdown(
&input,
&ChunkerOptions {
max_tokens: 150, // forces split at paragraph boundary
},
);
assert!(chunks.len() >= 2);
for c in &chunks {
let first = c.content.chars().next().unwrap();
assert!(
matches!(first, 'a' | 'b' | 'c'),
"document chunk starts with unexpected char: {:?}",
c.content.chars().take(10).collect::<String>()
);
assert!(
!c.partial_message,
"document chunks must not have partial_message=true"
);
}
}
#[test]
fn header_line_dropped_in_chat() {
// Simulate a blob that has a leading `# Chat transcript` header.
let md = "# Chat transcript — slack / #eng\n\n## 2026-01-01T00:00:00Z — alice\nhello"
.to_string();
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "x".into(),
markdown: md,
metadata: meta(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(chunks.len(), 1);
// The `# Chat transcript` header must be absent from the chunk content.
assert!(
!chunks[0].content.contains("# Chat transcript"),
"leading `# ` header must be dropped from chunk content"
);
assert!(chunks[0].content.contains("hello"));
}
#[test]
fn chunk_ids_are_stable_across_runs() {
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
markdown: "## 2026-01-01T00:00:00Z — alice\nhello".into(),
metadata: meta(),
};
let a = chunk_markdown(&input, &ChunkerOptions::default());
let b = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(
a.iter().map(|c| c.id.clone()).collect::<Vec<_>>(),
b.iter().map(|c| c.id.clone()).collect::<Vec<_>>()
);
}
#[test]
fn sequence_numbers_start_at_zero() {
let msgs: String = (0..5)
.map(|i| format!("## 2026-01-01T00:0{}:00Z — user{i}\nContent {i}\n\n", i))
.collect();
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "x".into(),
markdown: msgs,
metadata: meta(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions::default());
for (idx, c) in chunks.iter().enumerate() {
assert_eq!(c.seq_in_source, idx as u32);
}
}
#[test]
fn paragraph_boundaries_preferred_for_documents() {
// Build something that exceeds token budget so it must split.
let para1 = "a".repeat(400); // ~100 tokens
let para2 = "b".repeat(400);
@@ -225,7 +732,7 @@ mod tests {
source_kind: SourceKind::Document,
source_id: "doc1".into(),
markdown: text,
metadata: meta(),
metadata: meta_doc(),
};
let chunks = chunk_markdown(
&input,
@@ -234,8 +741,6 @@ mod tests {
},
);
assert!(chunks.len() >= 2);
// Every chunk should be a multiple of complete paragraphs — so starts
// with 'a', 'b', or 'c' and doesn't mix them arbitrarily.
for c in &chunks {
let first = c.content.chars().next().unwrap();
assert!(
@@ -247,16 +752,16 @@ mod tests {
}
#[test]
fn falls_back_to_line_split_when_no_paragraphs() {
fn falls_back_to_line_split_when_no_paragraphs_document() {
let text = (0..30)
.map(|i| format!("line-{i}-{}", "x".repeat(40)))
.collect::<Vec<_>>()
.join("\n");
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_kind: SourceKind::Document,
source_id: "x".into(),
markdown: text,
metadata: meta(),
metadata: meta_doc(),
};
let chunks = chunk_markdown(
&input,
@@ -271,54 +776,14 @@ mod tests {
}
#[test]
fn chunk_ids_are_stable_across_runs() {
let input = ChunkerInput {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
markdown: "hello".into(),
metadata: meta(),
};
let a = chunk_markdown(&input, &ChunkerOptions::default());
let b = chunk_markdown(&input, &ChunkerOptions::default());
assert_eq!(
a.iter().map(|c| c.id.clone()).collect::<Vec<_>>(),
b.iter().map(|c| c.id.clone()).collect::<Vec<_>>()
);
}
#[test]
fn sequence_numbers_start_at_zero() {
let text = (0..10)
.map(|i| format!("p{i}"))
.collect::<Vec<_>>()
.join("\n\n");
let input = ChunkerInput {
source_kind: SourceKind::Document,
source_id: "d".into(),
markdown: text,
metadata: meta(),
};
let chunks = chunk_markdown(
&input,
&ChunkerOptions {
max_tokens: 2, // forces one chunk per paragraph basically
},
);
for (idx, c) in chunks.iter().enumerate() {
assert_eq!(c.seq_in_source, idx as u32);
}
}
#[test]
fn utf8_boundaries_preserved_on_hard_split() {
fn utf8_boundaries_preserved_on_hard_split_document() {
// Single long line with no paragraph/line splits → falls to hard cut.
// Use a 3-byte-per-char Unicode codepoint (Chinese) and force a cut.
let text = "".repeat(400);
let input = ChunkerInput {
source_kind: SourceKind::Document,
source_id: "d".into(),
markdown: text.clone(),
metadata: meta(),
metadata: meta_doc(),
};
let chunks = chunk_markdown(
&input,
@@ -332,12 +797,12 @@ mod tests {
}
#[test]
fn zero_token_budget_is_clamped_without_empty_leading_chunk() {
fn zero_token_budget_is_clamped_without_empty_leading_chunk_document() {
let input = ChunkerInput {
source_kind: SourceKind::Document,
source_id: "d".into(),
markdown: "abcdef".into(),
metadata: meta(),
metadata: meta_doc(),
};
let chunks = chunk_markdown(&input, &ChunkerOptions { max_tokens: 0 });
assert!(!chunks.is_empty());
@@ -0,0 +1,416 @@
//! Atomic content-file writes via tempfile + fsync + rename.
//!
//! Each chunk body is written to `<parent>/.tmp_<uuid>.md`, then renamed to
//! its final path. The rename is atomic on any POSIX filesystem and behaves
//! correctly on NTFS (the old file is replaced atomically by the OS).
//!
//! **Immutability contract**: once a file exists at `abs_path`, it is never
//! overwritten. Callers must detect "already exists" and skip the write.
use sha2::{Digest, Sha256};
use std::io::Write;
use std::path::Path;
use super::compose::{compose_summary_md, split_front_matter, SummaryComposeInput};
use super::paths::{summary_abs_path, summary_rel_path};
/// Write `bytes` atomically to `abs_path` if the file does not already exist.
///
/// Returns `Ok(true)` when the file was newly written, `Ok(false)` when it
/// already existed (the existing file is left unchanged).
///
/// The write uses a sibling tempfile + rename so the final path is never
/// visible in a partial state. Parent directories are created automatically.
pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result<bool> {
// Fast path: file already exists.
if abs_path.exists() {
log::debug!(
"[content_store::atomic] skipping existing file: {}",
abs_path.display()
);
return Ok(false);
}
let parent = abs_path.parent().unwrap_or_else(|| Path::new("."));
std::fs::create_dir_all(parent)
.map_err(|e| anyhow::anyhow!("create_dir_all {:?}: {e}", parent))?;
// Write to a temp file in the same directory so rename is atomic.
let tmp_name = format!(".tmp_{}.md", uuid_v4_hex());
let tmp_path = parent.join(&tmp_name);
{
let mut f = std::fs::File::create(&tmp_path)
.map_err(|e| anyhow::anyhow!("create tempfile {:?}: {e}", tmp_path))?;
f.write_all(bytes)
.map_err(|e| anyhow::anyhow!("write tempfile {:?}: {e}", tmp_path))?;
f.sync_all()
.map_err(|e| anyhow::anyhow!("fsync tempfile {:?}: {e}", tmp_path))?;
}
// Rename: if the target appeared concurrently (another thread/process beat
// us), we lost the race — remove our temp and return false.
match std::fs::rename(&tmp_path, abs_path) {
Ok(()) => {
// fsync the parent directory so the rename (directory entry
// update) is durable across a crash or power loss. Without this,
// sync_all() on the file alone only durabilises the file data;
// the new directory entry can remain in pagecache and be lost if
// the system crashes before the OS flushes it. On POSIX (Linux /
// macOS) this is required for rename durability. On Windows, NTFS
// handles this differently and File::sync_all on a directory
// handle is not meaningful, so we restrict the call to Unix.
#[cfg(unix)]
if let Some(parent) = abs_path.parent() {
if let Ok(dir) = std::fs::File::open(parent) {
if let Err(e) = dir.sync_all() {
// Best-effort: the rename already committed the file;
// a dirent fsync failure is logged but not fatal.
log::warn!(
"[content_store::atomic] parent dir fsync failed for {:?}: {e}",
parent
);
}
}
}
log::debug!("[content_store::atomic] wrote {}", abs_path.display());
Ok(true)
}
Err(e) => {
// Best-effort cleanup of the temp file on failure.
let _ = std::fs::remove_file(&tmp_path);
if abs_path.exists() {
// Lost the race — another writer created the file first.
log::debug!(
"[content_store::atomic] lost rename race for {}",
abs_path.display()
);
Ok(false)
} else {
Err(anyhow::anyhow!(
"rename {:?} -> {:?}: {e}",
tmp_path,
abs_path
))
}
}
}
}
/// A summary that has been written to disk and is ready for SQLite upsert.
#[derive(Debug, Clone)]
pub struct StagedSummary {
pub summary_id: String,
/// Relative content path (forward-slash, e.g. `"summaries/source/slug/L1/id.md"`).
pub content_path: String,
/// SHA-256 hex digest over the **body bytes** only (front-matter excluded).
pub content_sha256: String,
}
/// Write a summary `.md` file to disk and return a [`StagedSummary`] ready for
/// SQLite upsert.
///
/// The relative path is built from the input metadata and the `tree_kind`. The
/// `date_for_global` argument is required when `input.tree_kind ==
/// SummaryTreeKind::Global`. The `scope_slug` must already be slugified by the
/// caller.
///
/// If the file already exists with the same body SHA-256 (idempotent re-stage),
/// the existing `StagedSummary` is returned without rewriting.
pub fn stage_summary(
content_root: &Path,
input: &SummaryComposeInput<'_>,
scope_slug: &str,
date_for_global: Option<chrono::DateTime<chrono::Utc>>,
) -> anyhow::Result<StagedSummary> {
let rel_path = summary_rel_path(
input.tree_kind,
scope_slug,
input.level,
input.summary_id,
date_for_global,
);
let abs_path = summary_abs_path(
content_root,
input.tree_kind,
scope_slug,
input.level,
input.summary_id,
date_for_global,
);
let composed = compose_summary_md(input);
let body_bytes = composed.body.as_bytes();
let sha256 = sha256_hex(body_bytes);
// Idempotent re-stage: if the file already exists, read and hash its
// body bytes. If the on-disk hash matches the new body's hash, return
// the StagedSummary unchanged (true idempotency). If the hashes differ
// the on-disk file is stale/corrupted — re-write it atomically with the
// new content so the db row and disk file are always consistent.
//
// Not re-writing would leave SQLite storing a content_sha256 that
// doesn't match the actual on-disk bytes, breaking integrity checks.
if abs_path.exists() {
let disk_sha = read_body_sha256(&abs_path).unwrap_or_default();
if disk_sha == sha256 {
log::debug!(
"[content_store::atomic] summary already on disk with matching sha: {}",
input.summary_id
);
return Ok(StagedSummary {
summary_id: input.summary_id.to_string(),
content_path: rel_path,
content_sha256: sha256,
});
}
// Hash mismatch — overwrite atomically.
log::debug!(
"[content_store::atomic] summary on-disk sha mismatch for {} — re-staging",
input.summary_id
);
// Remove the stale file first; write_if_new's fast-path would skip it.
let _ = std::fs::remove_file(&abs_path);
}
let full_bytes = composed.full.as_bytes();
write_if_new(&abs_path, full_bytes)?;
log::debug!(
"[content_store::atomic] staged summary {} → {}",
input.summary_id,
rel_path
);
Ok(StagedSummary {
summary_id: input.summary_id.to_string(),
content_path: rel_path,
content_sha256: sha256,
})
}
/// Read a summary/chunk `.md` file from disk, split off the YAML front-matter,
/// and return the SHA-256 hex digest of the **body bytes only**. Returns an
/// empty string (not an error) if the file cannot be read or parsed, so
/// callers can use the result as a cache key without propagating IO errors.
fn read_body_sha256(path: &Path) -> anyhow::Result<String> {
let raw = std::fs::read(path)?;
let content = std::str::from_utf8(&raw)?;
let (_fm, body) = split_front_matter(content)
.ok_or_else(|| anyhow::anyhow!("no front-matter in {:?}", path))?;
Ok(sha256_hex(body.as_bytes()))
}
/// Compute the SHA-256 hex digest of `bytes`.
pub fn sha256_hex(bytes: &[u8]) -> String {
let mut hasher = Sha256::new();
hasher.update(bytes);
hex::encode(hasher.finalize())
}
/// Tiny deterministic-ish hex string for temp file names.
fn uuid_v4_hex() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let t = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.subsec_nanos();
// Use a counter + timestamp for entropy (thread_id::as_u64 is nightly-only).
static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
format!(
"{:08x}{:016x}",
t,
n.wrapping_mul(0x9e37_79b9_7f4a_7c15).wrapping_add(t as u64)
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store::compose::SummaryComposeInput;
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
use tempfile::TempDir;
#[test]
fn write_creates_file_and_returns_true() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("sub").join("0.md");
let written = write_if_new(&path, b"hello world").unwrap();
assert!(written, "first write must return true");
assert_eq!(std::fs::read(&path).unwrap(), b"hello world");
}
#[test]
fn write_is_idempotent_returns_false_on_second_call() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("0.md");
write_if_new(&path, b"first").unwrap();
let written = write_if_new(&path, b"second").unwrap();
assert!(!written, "second write must return false");
assert_eq!(std::fs::read(&path).unwrap(), b"first");
}
#[test]
fn sha256_hex_is_stable() {
let a = sha256_hex(b"hello");
let b = sha256_hex(b"hello");
assert_eq!(a, b);
assert_ne!(sha256_hex(b"hello"), sha256_hex(b"world"));
assert_eq!(a.len(), 64); // 32 bytes → 64 hex chars
}
fn mk_summary_input<'a>(
tree_kind: SummaryTreeKind,
scope: &'a str,
id: &'a str,
body: &'a str,
children: &'a [String],
) -> SummaryComposeInput<'a> {
use chrono::TimeZone;
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
SummaryComposeInput {
summary_id: id,
tree_kind,
tree_id: "tree-001",
tree_scope: scope,
level: 1,
child_ids: children,
child_count: children.len(),
time_range_start: ts,
time_range_end: ts,
sealed_at: ts,
body,
}
}
#[test]
fn stage_summary_writes_file_and_returns_staged() {
let dir = TempDir::new().unwrap();
let children = vec!["c1".to_string()];
let input = mk_summary_input(
SummaryTreeKind::Source,
"gmail:alice@x.com",
"summary:L1:test1",
"summary body",
&children,
);
let staged = stage_summary(dir.path(), &input, "gmail-alice-x-com", None).unwrap();
assert_eq!(staged.summary_id, "summary:L1:test1");
assert!(staged.content_path.starts_with("summaries/source/"));
assert!(staged.content_path.ends_with(".md"));
assert_eq!(staged.content_sha256.len(), 64);
// File must exist on disk
let mut abs = dir.path().to_path_buf();
for part in staged.content_path.split('/') {
abs.push(part);
}
assert!(abs.exists(), "staged file must exist");
}
#[test]
fn stage_summary_is_idempotent() {
let dir = TempDir::new().unwrap();
let children = vec!["c1".to_string()];
let input = mk_summary_input(
SummaryTreeKind::Topic,
"person:alex",
"summary:L1:idem",
"idempotent body",
&children,
);
let first = stage_summary(dir.path(), &input, "person-alex", None).unwrap();
let second = stage_summary(dir.path(), &input, "person-alex", None).unwrap();
assert_eq!(first.content_sha256, second.content_sha256);
assert_eq!(first.content_path, second.content_path);
}
#[test]
fn stage_summary_global_uses_date_in_path() {
use chrono::TimeZone;
let dir = TempDir::new().unwrap();
let date = chrono::Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap();
let children = vec![];
let input = mk_summary_input(
SummaryTreeKind::Global,
"global",
"summary:L0:daily",
"daily recap",
&children,
);
let staged = stage_summary(dir.path(), &input, "global", Some(date)).unwrap();
assert!(
staged.content_path.contains("2026-04-28"),
"global summary path must contain date; got: {}",
staged.content_path
);
}
#[test]
fn stage_summary_sha256_is_over_body_only() {
let dir = TempDir::new().unwrap();
let children = vec![];
let body = "the body content";
let input = mk_summary_input(
SummaryTreeKind::Source,
"gmail:x@y.com",
"summary:L1:sha-test",
body,
&children,
);
let staged = stage_summary(dir.path(), &input, "gmail-x-y-com", None).unwrap();
let expected = sha256_hex(body.as_bytes());
assert_eq!(staged.content_sha256, expected);
}
#[test]
fn stage_summary_rewrites_stale_on_disk_body() {
// Create a tempdir and write a "stale" file at the expected path with
// a body that differs from what the new stage_summary call would write.
// After stage_summary, the file on disk must match the new body.
let dir = TempDir::new().unwrap();
let children = vec!["c1".to_string()];
let new_body = "fresh body for re-stage test";
let input = mk_summary_input(
SummaryTreeKind::Source,
"gmail:stale@test.com",
"summary:L1:stale-test",
new_body,
&children,
);
// First stage with the real body to get the path.
let first = stage_summary(dir.path(), &input, "gmail-stale-test-com", None).unwrap();
// Corrupt the on-disk file by writing a different body to the path.
let mut abs = dir.path().to_path_buf();
for part in first.content_path.split('/') {
abs.push(part);
}
// Overwrite with stale content.
std::fs::write(&abs, b"---\nstale_key: true\n---\nSTALE BODY CONTENT").unwrap();
// Now re-stage: must detect sha mismatch and re-write.
let second = stage_summary(dir.path(), &input, "gmail-stale-test-com", None).unwrap();
// The returned sha must match the new body.
let expected_sha = sha256_hex(new_body.as_bytes());
assert_eq!(
second.content_sha256, expected_sha,
"re-staged sha must match new body"
);
// The on-disk file must now contain the new body (not the stale one).
let disk_bytes = std::fs::read(&abs).unwrap();
let disk_str = std::str::from_utf8(&disk_bytes).unwrap();
assert!(
disk_str.contains(new_body),
"on-disk file must contain new body after re-stage"
);
assert!(
!disk_str.contains("STALE BODY CONTENT"),
"stale body must be gone after re-stage"
);
}
}
@@ -0,0 +1,878 @@
//! YAML front-matter + body composition for chunk `.md` files.
//!
//! Each file written to disk has the form:
//! ```text
//! ---
//! source_kind: chat
//! source_id: slack:#eng
//! seq: 0
//! owner: alice@example.com
//! timestamp: 2026-04-28T10:00:00Z
//! time_range_start: 2026-04-28T10:00:00Z
//! time_range_end: 2026-04-28T10:05:00Z
//! source_ref: slack://permalink/…
//! tags:
//! - person/Alice-Smith
//! - project/Phoenix
//! ---
//! ## 2026-04-28T10:00:00Z — alice
//! Message body here.
//! ```
//!
//! For email source_kind, additional fields are emitted:
//! ```text
//! participants:
//! - alice@example.com
//! - bob@example.com
//! aliases:
//! - "alice@example.com <-> bob@example.com: chunk 0"
//! ```
//! These are parsed from the `source_id` field (format `gmail:{participants}`
//! where `participants` is `addr1|addr2|...` pipe-separated) at compose time.
//! `sender` and `thread_id` are no longer emitted — they are not meaningful
//! with participant-based bucketing.
//!
//! **SHA-256 is computed over the body bytes only** (everything after `---\n`
//! on the second delimiter line). This allows tags to be rewritten atomically
//! without invalidating the content hash.
use chrono::{DateTime, Utc};
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
use crate::openhuman::memory::tree::types::{Chunk, SourceKind};
/// Compose the full file content (front-matter + body) for `chunk`.
///
/// Returns `(full_file_bytes, body_bytes)`. The caller writes `full_file_bytes`
/// to disk; `body_bytes` is what the SHA-256 is computed over.
pub fn compose_chunk_file(chunk: &Chunk) -> (Vec<u8>, Vec<u8>) {
let front_matter = build_front_matter(chunk);
let body = chunk.content.as_bytes().to_vec();
let mut full = Vec::with_capacity(front_matter.len() + body.len());
full.extend_from_slice(&front_matter);
full.extend_from_slice(&body);
(full, body)
}
/// Build the YAML front-matter block (including delimiters) as UTF-8 bytes.
fn build_front_matter(chunk: &Chunk) -> Vec<u8> {
let meta = &chunk.metadata;
let ts = meta.timestamp.to_rfc3339();
let ts_start = meta.time_range.0.to_rfc3339();
let ts_end = meta.time_range.1.to_rfc3339();
let mut fm = String::new();
fm.push_str("---\n");
fm.push_str(&format!("source_kind: {}\n", meta.source_kind.as_str()));
// Escape backslashes and quotes in source_id for safety.
fm.push_str(&format!("source_id: {}\n", yaml_scalar(&meta.source_id)));
fm.push_str(&format!("seq: {}\n", chunk.seq_in_source));
fm.push_str(&format!("owner: {}\n", yaml_scalar(&meta.owner)));
fm.push_str(&format!("timestamp: {ts}\n"));
fm.push_str(&format!("time_range_start: {ts_start}\n"));
fm.push_str(&format!("time_range_end: {ts_end}\n"));
if let Some(ref sr) = meta.source_ref {
fm.push_str(&format!("source_ref: {}\n", yaml_scalar(&sr.value)));
}
if meta.tags.is_empty() {
fm.push_str("tags: []\n");
} else {
fm.push_str("tags:\n");
for tag in &meta.tags {
fm.push_str(&format!(" - {}\n", yaml_scalar(tag)));
}
}
// Email-specific fields: participants list + Obsidian alias.
// Parsed from source_id which is `gmail:{participants}` for Gmail-ingested
// chunks, where participants is `addr1|addr2|...` (sorted, deduped).
// If the format doesn't match, these fields are omitted.
if meta.source_kind == SourceKind::Email {
if let Some(addrs) = parse_gmail_participants_source_id(&meta.source_id) {
// participants: YAML list
fm.push_str("participants:\n");
for addr in &addrs {
fm.push_str(&format!(" - {}\n", yaml_scalar(addr)));
}
// aliases: human-readable conversation label for Obsidian
let alias = build_participants_alias(&addrs, chunk.seq_in_source);
fm.push_str("aliases:\n");
fm.push_str(&format!(" - {}\n", yaml_scalar(&alias)));
}
}
fm.push_str("---\n");
fm.into_bytes()
}
/// Parse a `gmail:{participants}` source_id into the list of participant addresses.
///
/// `participants` is `addr1|addr2|...` (sorted, deduped, pipe-separated).
/// Returns `Some(Vec<String>)` when the source_id has exactly two
/// colon-separated segments (`gmail` prefix + non-empty participants). Returns
/// `None` for legacy or malformed source_ids.
fn parse_gmail_participants_source_id(source_id: &str) -> Option<Vec<String>> {
let (prefix, participants) = source_id.split_once(':')?;
if prefix != "gmail" || participants.is_empty() {
return None;
}
let addrs: Vec<String> = participants
.split('|')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
if addrs.is_empty() {
None
} else {
Some(addrs)
}
}
/// Build a human-readable alias for an email chunk suitable for Obsidian's
/// `aliases:` field.
///
/// For two participants: `"alice@x.com <-> bob@y.com: chunk 0"`
/// For more than two: `"alice@x.com <-> 2 others: chunk 0"`
/// (where `alice@x.com` is the first in sorted order)
///
/// The alias is kept under ~80 characters to avoid YAML rendering issues.
fn build_participants_alias(addrs: &[String], seq: u32) -> String {
let label = match addrs {
[] => "unknown".to_string(),
[only] => only.clone(),
[first, second] => format!("{} <-> {}", first, second),
[first, rest @ ..] => format!("{} <-> {} others", first, rest.len()),
};
format!("{}: chunk {}", label, seq)
}
/// Rewrite the `tags:` block in an existing file's front-matter, replacing it
/// with the new tag list while leaving the body unchanged.
///
/// Returns the new full file bytes. Errors if the front-matter delimiters
/// cannot be found.
pub fn rewrite_tags(file_bytes: &[u8], new_tags: &[String]) -> Result<Vec<u8>, String> {
let content =
std::str::from_utf8(file_bytes).map_err(|e| format!("file is not valid UTF-8: {e}"))?;
let (front_matter, body) = split_front_matter(content)
.ok_or_else(|| "cannot find front-matter delimiters".to_string())?;
// Rewrite tags: block in the front-matter string.
let new_fm = replace_tags_in_front_matter(front_matter, new_tags)?;
let mut out = Vec::with_capacity(new_fm.len() + body.len() + 4);
out.extend_from_slice(new_fm.as_bytes());
out.extend_from_slice(body.as_bytes());
Ok(out)
}
/// Replace the `tags:` stanza in a front-matter string. Returns the new
/// front-matter string (delimiters preserved).
fn replace_tags_in_front_matter(fm: &str, new_tags: &[String]) -> Result<String, String> {
// Build the replacement block.
let replacement = if new_tags.is_empty() {
"tags: []".to_string()
} else {
let mut s = "tags:".to_string();
for tag in new_tags {
s.push('\n');
s.push_str(&format!(" - {}", yaml_scalar(tag)));
}
s
};
// Locate the `tags:` key and consume through the block.
let lines: Vec<&str> = fm.lines().collect();
let mut out_lines: Vec<&str> = Vec::new();
let mut i = 0;
let mut found = false;
while i < lines.len() {
let line = lines[i];
if line == "tags: []" || line == "tags:" {
found = true;
// Skip all subsequent lines that are tag list items (start with ` - `).
// The replacement will be inserted wholesale.
i += 1;
if line == "tags:" {
while i < lines.len() && lines[i].starts_with(" - ") {
i += 1;
}
}
// We've consumed the old block; we'll append replacement after the loop.
continue;
}
out_lines.push(line);
i += 1;
}
if !found {
return Err("tags: key not found in front-matter".to_string());
}
// Rebuild: all non-tag lines + replacement + closing `---`.
// Front-matter was: `---\n...\ntags: ...\n---\n`
// After loop, out_lines has everything except the tags block.
// Insert replacement before the closing `---`.
let closing = out_lines
.iter()
.rposition(|l| *l == "---")
.unwrap_or(out_lines.len());
let mut result_lines: Vec<String> =
out_lines[..closing].iter().map(|l| l.to_string()).collect();
result_lines.push(replacement);
result_lines.push("---".to_string());
let mut result = result_lines.join("\n");
result.push('\n');
Ok(result)
}
// ── Summary composition ──────────────────────────────────────────────────────
/// Input data required to compose a summary `.md` file.
pub struct SummaryComposeInput<'a> {
pub summary_id: &'a str,
pub tree_kind: SummaryTreeKind,
pub tree_id: &'a str,
/// Raw tree scope string, e.g. `"gmail:alice@x.com|bob@y.com"` or `"global"`.
pub tree_scope: &'a str,
pub level: u32,
/// Child ids (chunk_ids at L0 → L1, summary_ids for cascades).
pub child_ids: &'a [String],
/// Total child count (== child_ids.len() unless truncated).
pub child_count: usize,
pub time_range_start: DateTime<Utc>,
pub time_range_end: DateTime<Utc>,
pub sealed_at: DateTime<Utc>,
/// Raw summariser output text — the body written to disk.
pub body: &'a str,
}
/// The composed front-matter, body, and full file content for a summary.
pub struct ComposedSummary {
/// The YAML front-matter block (including `---` delimiters), UTF-8 string.
pub front_matter: String,
/// The body (summariser output), UTF-8 string.
pub body: String,
/// `front_matter + body` — what gets written to disk.
pub full: String,
}
/// Compose the full `.md` content for a summary node.
///
/// Returns a [`ComposedSummary`] whose `full` field is written to disk.
/// SHA-256 is computed over `body` bytes only, not `full`.
pub fn compose_summary_md(record: &SummaryComposeInput<'_>) -> ComposedSummary {
let fm = build_summary_front_matter(record);
let body = record.body.to_string();
let full = format!("{}{}", fm, body);
ComposedSummary {
front_matter: fm,
body,
full,
}
}
/// Build the YAML front-matter block for a summary node.
fn build_summary_front_matter(r: &SummaryComposeInput<'_>) -> String {
let tree_kind_str = match r.tree_kind {
SummaryTreeKind::Source => "source",
SummaryTreeKind::Global => "global",
SummaryTreeKind::Topic => "topic",
};
let trs = r.time_range_start.to_rfc3339();
let tre = r.time_range_end.to_rfc3339();
let sealed = r.sealed_at.to_rfc3339();
let mut fm = String::new();
fm.push_str("---\n");
fm.push_str(&format!("id: {}\n", yaml_scalar(r.summary_id)));
fm.push_str("kind: summary\n");
fm.push_str(&format!("tree_kind: {tree_kind_str}\n"));
fm.push_str(&format!("tree_id: {}\n", yaml_scalar(r.tree_id)));
fm.push_str(&format!("tree_scope: {}\n", yaml_scalar(r.tree_scope)));
fm.push_str(&format!("level: {}\n", r.level));
// children: YAML list
if r.child_ids.is_empty() {
fm.push_str("children: []\n");
} else {
fm.push_str("children:\n");
for id in r.child_ids {
fm.push_str(&format!(" - {}\n", yaml_scalar(id)));
}
}
fm.push_str(&format!("child_count: {}\n", r.child_count));
fm.push_str(&format!("time_range_start: {trs}\n"));
fm.push_str(&format!("time_range_end: {tre}\n"));
fm.push_str(&format!("sealed_at: {sealed}\n"));
// aliases: human-readable title
let alias = build_summary_alias(r);
fm.push_str("aliases:\n");
fm.push_str(&format!(" - {}\n", yaml_scalar(&alias)));
fm.push_str("tags: []\n");
fm.push_str("---\n");
fm
}
/// Build a human-readable alias for the summary's `aliases:` front-matter field.
fn build_summary_alias(r: &SummaryComposeInput<'_>) -> String {
let date_range = format_date_range(r.time_range_start, r.time_range_end);
match r.tree_kind {
SummaryTreeKind::Source => {
let scope_short = scope_short_label(r.tree_scope);
format!(
"L{} \u{00b7} {} \u{00b7} {} children \u{00b7} {}",
r.level, scope_short, r.child_count, date_range
)
}
SummaryTreeKind::Global => {
format!(
"L{} \u{00b7} global digest \u{00b7} {}",
r.level, date_range
)
}
SummaryTreeKind::Topic => {
// Strip protocol prefix like "topic:" from scope for readability.
let entity = r
.tree_scope
.split_once(':')
.map(|(_, v)| v)
.unwrap_or(r.tree_scope);
format!(
"L{} \u{00b7} topic {} \u{00b7} {} children",
r.level, entity, r.child_count
)
}
}
}
/// Format the date range as `"yyyy-mm-dd"` (if start == end date) or
/// `"yyyy-mm-ddyyyy-mm-dd"`.
fn format_date_range(start: DateTime<Utc>, end: DateTime<Utc>) -> String {
let s = start.format("%Y-%m-%d").to_string();
let e = end.format("%Y-%m-%d").to_string();
if s == e {
s
} else {
format!("{s}\u{2013}{e}") // en dash
}
}
/// Build a short human-readable label for the tree scope used in aliases.
///
/// For Gmail source scopes like `"gmail:alice@x.com|bob@y.com"`:
/// - 2 participants → `"alice@x.com ↔ bob@y.com"`
/// - N > 2 → `"alice@x.com + N-1 others"`
/// - Otherwise → the raw scope (e.g. `"slack:#eng"`)
fn scope_short_label(scope: &str) -> String {
if let Some((prefix, participants)) = scope.split_once(':') {
if prefix == "gmail" && !participants.is_empty() {
let addrs: Vec<&str> = participants.split('|').collect();
return match addrs.as_slice() {
[] => scope.to_string(),
[only] => only.to_string(),
[first, second] => format!("{} \u{2194} {}", first, second), // ↔
[first, rest @ ..] => format!("{} + {} others", first, rest.len()),
};
}
}
scope.to_string()
}
/// Rewrite the `tags:` block in a summary file's front-matter, replacing it
/// with the new tag list while leaving the body unchanged.
///
/// Reuses the generic [`rewrite_tags`] function — the front-matter structure
/// is identical for both chunk and summary `.md` files.
pub fn rewrite_summary_tags(file_bytes: &[u8], new_tags: &[String]) -> Result<Vec<u8>, String> {
rewrite_tags(file_bytes, new_tags)
}
/// Split a file into `(front_matter, body)` at the second `---` delimiter.
///
/// Returns `None` if the file does not have the expected `---\n...\n---\n` form.
pub fn split_front_matter(content: &str) -> Option<(&str, &str)> {
// The file must start with `---\n`.
if !content.starts_with("---\n") {
return None;
}
// Find the closing `---` line (must be `---` alone on a line after the first line).
let rest = &content[4..]; // skip the opening `---\n`
let close_idx = rest.find("\n---\n").or_else(|| {
// Could be at the very end (no body).
rest.strip_suffix("\n---").map(|r| r.len())
})?;
let fm_end = 4 + close_idx + 5; // include `\n---\n`
Some((&content[..fm_end], &content[fm_end..]))
}
/// Format a string as an unquoted YAML scalar when safe, or as a
/// double-quoted string when it contains special characters.
///
/// We conservatively quote strings containing `:`, `#`, `[`, `]`, `{`, `}`,
/// `"`, `'`, `\`, leading/trailing whitespace, or that start with special
/// YAML indicator characters.
fn yaml_scalar(s: &str) -> String {
let needs_quoting = s.is_empty()
|| s.trim() != s
|| s.starts_with(|c: char| {
matches!(
c,
'&' | '*' | '?' | '|' | '-' | '<' | '>' | '=' | '!' | '%' | '@' | '`'
)
})
|| s.contains([':', '#', '[', ']', '{', '}', '"', '\'']);
if needs_quoting {
let escaped = s.replace('\\', "\\\\").replace('"', "\\\"");
format!("\"{escaped}\"")
} else {
s.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
use crate::openhuman::memory::tree::types::{Metadata, SourceKind, SourceRef};
use chrono::TimeZone;
fn sample_chunk() -> Chunk {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
Chunk {
id: "abc123".into(),
content: "## 2026-01-01T00:00:00Z — alice\nhello world".into(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice@example.com".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec!["person/Alice".into(), "org/Acme".into()],
source_ref: Some(SourceRef::new("slack://m1".to_string())),
},
token_count: 10,
seq_in_source: 0,
created_at: ts,
partial_message: false,
}
}
#[test]
fn compose_produces_front_matter_and_body() {
let chunk = sample_chunk();
let (full, body) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
assert!(full_str.starts_with("---\n"), "must start with ---");
assert!(full_str.contains("source_kind: chat"));
assert!(full_str.contains("source_id: \"slack:#eng\""));
assert!(full_str.contains("seq: 0"));
assert!(full_str.contains("tags:"));
assert!(full_str.contains(" - person/Alice"));
assert!(full_str.ends_with("hello world"));
assert_eq!(
body,
b"## 2026-01-01T00:00:00Z \xe2\x80\x94 alice\nhello world"
);
}
#[test]
fn split_front_matter_round_trips() {
let chunk = sample_chunk();
let (full, body) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
let (fm, b) = split_front_matter(full_str).expect("split must succeed");
assert!(fm.starts_with("---\n"));
assert!(fm.ends_with("---\n"));
assert_eq!(b.as_bytes(), body.as_slice());
}
#[test]
fn rewrite_tags_preserves_body() {
let chunk = sample_chunk();
let (full, body) = compose_chunk_file(&chunk);
let new_tags = vec!["person/Bob".into(), "project/Phoenix".into()];
let rewritten = rewrite_tags(&full, &new_tags).unwrap();
let rewritten_str = std::str::from_utf8(&rewritten).unwrap();
assert!(rewritten_str.contains(" - person/Bob"));
assert!(!rewritten_str.contains(" - person/Alice"));
// Body must be unchanged.
assert!(rewritten_str.ends_with(std::str::from_utf8(&body).unwrap()));
}
#[test]
fn rewrite_tags_empty_list() {
let chunk = sample_chunk();
let (full, _) = compose_chunk_file(&chunk);
let rewritten = rewrite_tags(&full, &[]).unwrap();
let s = std::str::from_utf8(&rewritten).unwrap();
assert!(s.contains("tags: []"));
assert!(!s.contains(" - person/"));
}
#[test]
fn yaml_scalar_quotes_special_characters() {
assert_eq!(yaml_scalar("slack:#eng"), "\"slack:#eng\"");
assert_eq!(yaml_scalar("hello world"), "hello world");
assert_eq!(yaml_scalar(""), "\"\"");
}
fn sample_email_chunk() -> Chunk {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
Chunk {
id: "emailchunk1".into(),
content: "---\nFrom: alice@example.com\nSubject: Hello\n\nHello there.".into(),
metadata: Metadata {
source_kind: SourceKind::Email,
source_id: "gmail:alice@example.com|bob@example.com".into(),
owner: "owner@example.com".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec!["gmail".into()],
source_ref: None,
},
token_count: 15,
seq_in_source: 0,
created_at: ts,
partial_message: false,
}
}
#[test]
fn email_chunk_has_participants_list_and_alias() {
let chunk = sample_email_chunk();
let (full, _body) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
// participants block must be a YAML list
assert!(
full_str.contains("participants:"),
"email chunk must have participants field; got:\n{full_str}"
);
assert!(
full_str.contains(" - alice@example.com"),
"alice must appear as list item; got:\n{full_str}"
);
assert!(
full_str.contains(" - bob@example.com"),
"bob must appear as list item; got:\n{full_str}"
);
// aliases block must be present
assert!(
full_str.contains("aliases:"),
"email chunk must have aliases field; got:\n{full_str}"
);
assert!(
full_str.contains("alice@example.com <-> bob@example.com: chunk 0"),
"alias must encode participants; got:\n{full_str}"
);
// sender and thread_id must NOT appear
assert!(
!full_str.contains("sender:"),
"email chunk must NOT have sender field; got:\n{full_str}"
);
assert!(
!full_str.contains("thread_id:"),
"email chunk must NOT have thread_id field; got:\n{full_str}"
);
}
#[test]
fn email_chunk_many_participants_alias_summarises() {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let chunk = Chunk {
id: "em2".into(),
content: "body".into(),
metadata: Metadata {
source_kind: SourceKind::Email,
source_id: "gmail:alice@x.com|bob@y.com|carol@z.com".into(),
owner: "owner".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: None,
},
token_count: 1,
seq_in_source: 3,
created_at: ts,
partial_message: false,
};
let (full, _) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
assert!(
full_str.contains("participants:"),
"three-party chunk needs participants list; got:\n{full_str}"
);
// With 3 participants: first + "2 others"
assert!(
full_str.contains("alice@x.com <-> 2 others: chunk 3"),
"alias with 3 participants must summarise; got:\n{full_str}"
);
}
#[test]
fn email_chunk_body_bytes_unchanged_by_extra_fields() {
// Adding participants/aliases to front-matter must not affect body_bytes
// (SHA-256 invariant: the hash is over body only, not front-matter).
let chunk = sample_email_chunk();
let (full, body) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
// Body must still appear at the end unmodified.
assert!(
full_str.ends_with(std::str::from_utf8(&body).unwrap()),
"body bytes must appear unmodified after front-matter"
);
// body must equal chunk.content bytes
assert_eq!(body, chunk.content.as_bytes());
}
#[test]
fn chat_chunk_has_no_email_specific_fields() {
let chunk = sample_chunk(); // source_kind = Chat
let (full, _) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
assert!(
!full_str.contains("aliases:"),
"chat chunk must not have aliases field"
);
assert!(
!full_str.contains("participants:"),
"chat chunk must not have participants field"
);
assert!(
!full_str.contains("sender:"),
"chat chunk must not have sender field"
);
assert!(
!full_str.contains("thread_id:"),
"chat chunk must not have thread_id field"
);
}
#[test]
fn email_chunk_with_malformed_source_id_omits_extra_fields() {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let chunk = Chunk {
id: "xyz".into(),
content: "body".into(),
metadata: Metadata {
source_kind: SourceKind::Email,
source_id: "legacysourceid".into(), // no `gmail:` prefix → parse fails
owner: "owner".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: None,
},
token_count: 1,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
let (full, _) = compose_chunk_file(&chunk);
let full_str = std::str::from_utf8(&full).unwrap();
// Malformed source_id → no email extras, no panic.
assert!(!full_str.contains("aliases:"));
assert!(!full_str.contains("participants:"));
assert!(!full_str.contains("sender:"));
}
// ─── summary compose tests ────────────────────────────────────────────────
fn sample_summary_input(
tree_kind: SummaryTreeKind,
scope: &str,
level: u32,
) -> SummaryComposeInput<'static> {
let ts_start = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let ts_end = chrono::Utc.timestamp_millis_opt(1_700_086_400_000).unwrap();
let sealed = chrono::Utc.timestamp_millis_opt(1_700_090_000_000).unwrap();
// Leak the strings so they have 'static lifetime for this test helper.
// Only used in tests, not production code.
let scope: &'static str = Box::leak(scope.to_string().into_boxed_str());
SummaryComposeInput {
summary_id: "summary:L1:abc",
tree_kind,
tree_id: "tree-id-001",
tree_scope: scope,
level,
child_ids: Box::leak(
vec!["child-1".to_string(), "child-2".to_string()].into_boxed_slice(),
),
child_count: 2,
time_range_start: ts_start,
time_range_end: ts_end,
sealed_at: sealed,
body: "This is the summariser output.\n",
}
}
#[test]
fn compose_source_summary_has_required_front_matter() {
let input = sample_summary_input(SummaryTreeKind::Source, "gmail:alice@x.com|bob@y.com", 1);
let composed = compose_summary_md(&input);
let fm = &composed.front_matter;
assert!(fm.starts_with("---\n"), "front-matter must start with ---");
assert!(fm.ends_with("---\n"), "front-matter must end with ---\\n");
assert!(fm.contains("kind: summary"), "must have kind: summary");
assert!(
fm.contains("tree_kind: source"),
"must have tree_kind: source"
);
assert!(fm.contains("level: 1"), "must have level");
assert!(fm.contains("child_count: 2"), "must have child_count");
assert!(fm.contains(" - child-1"), "must list child ids");
assert!(fm.contains(" - child-2"), "must list child ids");
assert!(fm.contains("tags: []"), "must start with empty tags");
// aliases must mention the scope
assert!(fm.contains("aliases:"), "must have aliases");
assert!(
composed.body == "This is the summariser output.\n",
"body must be the summariser text"
);
assert!(composed.full.ends_with("This is the summariser output.\n"));
}
#[test]
fn compose_global_summary_alias_format() {
let input = sample_summary_input(SummaryTreeKind::Global, "global", 0);
let composed = compose_summary_md(&input);
assert!(
composed.front_matter.contains("tree_kind: global"),
"must have tree_kind: global"
);
assert!(
composed.front_matter.contains("global digest"),
"alias must mention 'global digest'"
);
}
#[test]
fn compose_topic_summary_alias_format() {
let input = sample_summary_input(SummaryTreeKind::Topic, "person:alex-johnson", 1);
let composed = compose_summary_md(&input);
assert!(
composed.front_matter.contains("tree_kind: topic"),
"must have tree_kind: topic"
);
assert!(
composed.front_matter.contains("topic"),
"alias must mention topic entity"
);
}
#[test]
fn compose_summary_with_zero_children() {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let input = SummaryComposeInput {
summary_id: "summary:L0:empty",
tree_kind: SummaryTreeKind::Source,
tree_id: "t1",
tree_scope: "gmail:alice@x.com",
level: 0,
child_ids: &[],
child_count: 0,
time_range_start: ts,
time_range_end: ts,
sealed_at: ts,
body: "empty",
};
let composed = compose_summary_md(&input);
assert!(composed.front_matter.contains("children: []"));
assert!(composed.front_matter.contains("child_count: 0"));
}
#[test]
fn compose_summary_same_start_end_date_single_date_alias() {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let input = SummaryComposeInput {
summary_id: "summary:L1:sameday",
tree_kind: SummaryTreeKind::Global,
tree_id: "t1",
tree_scope: "global",
level: 1,
child_ids: &["child-a".to_string()],
child_count: 1,
time_range_start: ts,
time_range_end: ts, // same as start
sealed_at: ts,
body: "day recap",
};
let composed = compose_summary_md(&input);
// Alias must contain just one date, not "datedate"
let alias_line = composed
.front_matter
.lines()
.find(|l| l.contains("L1") && l.contains("global digest"))
.expect("alias line must be present");
// The date should appear exactly once (no en-dash range)
let date_str = ts.format("%Y-%m-%d").to_string();
assert!(
alias_line.contains(&date_str),
"alias must contain the date; got: {alias_line}"
);
// Must not contain an en-dash (range indicator)
assert!(
!alias_line.contains('\u{2013}'),
"same-day alias must not have en-dash range; got: {alias_line}"
);
}
#[test]
fn scope_short_label_two_participants() {
let label = scope_short_label("gmail:alice@x.com|bob@y.com");
assert_eq!(label, "alice@x.com \u{2194} bob@y.com");
}
#[test]
fn scope_short_label_many_participants() {
let label = scope_short_label("gmail:alice@x.com|bob@y.com|carol@z.com");
assert_eq!(label, "alice@x.com + 2 others");
}
#[test]
fn scope_short_label_non_gmail_returns_raw() {
let label = scope_short_label("slack:#general");
assert_eq!(label, "slack:#general");
}
#[test]
fn rewrite_summary_tags_delegates_to_rewrite_tags() {
// compose a summary, then rewrite its tags — body must stay unchanged.
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let input = SummaryComposeInput {
summary_id: "sum:L1:rwttest",
tree_kind: SummaryTreeKind::Source,
tree_id: "t1",
tree_scope: "gmail:alice@x.com",
level: 1,
child_ids: &["c1".to_string()],
child_count: 1,
time_range_start: ts,
time_range_end: ts,
sealed_at: ts,
body: "summary body text",
};
let composed = compose_summary_md(&input);
let file_bytes = composed.full.as_bytes();
let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()];
let rewritten = rewrite_summary_tags(file_bytes, &new_tags).unwrap();
let rewritten_str = std::str::from_utf8(&rewritten).unwrap();
assert!(rewritten_str.contains(" - person/Alice-Smith"));
assert!(rewritten_str.contains(" - topic/Memory"));
assert!(!rewritten_str.contains("tags: []"));
// Body must be unchanged
assert!(rewritten_str.ends_with("summary body text"));
}
}
@@ -0,0 +1,169 @@
//! Content store for memory-tree chunk and summary `.md` files (Phase MD-content).
//!
//! Bodies are stored on disk as `.md` files with YAML front-matter.
//! SQLite holds `content_path` (relative, forward-slash) and `content_sha256`
//! (over body bytes only) as pointers + integrity tokens.
//!
//! ## Module layout
//!
//! - [`paths`] — path generation + `slugify_source_id` + summary path builders
//! - [`compose`] — YAML front-matter + body composition; tag rewriting
//! - [`atomic`] — tempfile+fsync+rename writes; SHA-256; `stage_summary`
//! - [`read`] — read + SHA-256 verification + `split_front_matter`; summary variants
//! - [`tags`] — `update_chunk_tags` + `update_summary_tags` + slugifiers
pub mod atomic;
pub mod compose;
pub mod paths;
pub mod read;
pub mod tags;
use std::path::Path;
use crate::openhuman::memory::tree::types::Chunk;
pub use atomic::StagedSummary;
pub use compose::SummaryComposeInput;
pub use paths::SummaryTreeKind;
/// A chunk that has been written to disk and is ready for SQLite upsert.
///
/// Callers build a `Vec<StagedChunk>` from `stage_chunks`, then pass it to
/// `store::upsert_chunks_tx` in the same SQLite transaction.
#[derive(Debug, Clone)]
pub struct StagedChunk {
/// The original chunk (metadata + content).
pub chunk: Chunk,
/// Relative content path (forward-slash, e.g. `"chat/slack-eng/0.md"`).
pub content_path: String,
/// SHA-256 hex digest over the body bytes only.
pub content_sha256: String,
}
/// Update the `tags:` block in a summary's on-disk `.md` file after an
/// extraction job runs.
///
/// Delegates to [`tags::update_summary_tags`].
pub fn update_summary_tags(
config: &crate::openhuman::config::Config,
summary_id: &str,
) -> anyhow::Result<()> {
tags::update_summary_tags(config, summary_id)
}
/// Write all chunks in `chunks` to disk and return `StagedChunk` records
/// ready for SQLite upsert.
///
/// Each chunk file is written atomically via a sibling temp-file + rename.
/// Already-existing files are skipped (immutable-body contract). Parent
/// directories are created on demand.
///
/// `content_root` — absolute path to the root of the content store.
pub fn stage_chunks(content_root: &Path, chunks: &[Chunk]) -> anyhow::Result<Vec<StagedChunk>> {
let mut staged = Vec::with_capacity(chunks.len());
for chunk in chunks {
let source_kind = chunk.metadata.source_kind.as_str();
let source_id = &chunk.metadata.source_id;
let rel_path = paths::chunk_rel_path(source_kind, source_id, &chunk.id);
let abs_path = paths::chunk_abs_path(content_root, source_kind, source_id, &chunk.id);
let (full_bytes, body_bytes) = compose::compose_chunk_file(chunk);
let sha256 = atomic::sha256_hex(&body_bytes);
match atomic::write_if_new(&abs_path, &full_bytes) {
Ok(written) => {
if written {
log::debug!("[content_store] wrote chunk {} → {}", chunk.id, rel_path);
} else {
log::debug!(
"[content_store] chunk {} already on disk at {}",
chunk.id,
rel_path
);
}
}
Err(e) => {
log::error!(
"[content_store] failed to write chunk {} to {}: {e}",
chunk.id,
rel_path
);
return Err(e);
}
}
staged.push(StagedChunk {
chunk: chunk.clone(),
content_path: rel_path,
content_sha256: sha256,
});
}
Ok(staged)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::types::{Metadata, SourceKind};
use chrono::TimeZone;
use tempfile::TempDir;
fn sample_chunk(seq: u32) -> Chunk {
let ts = chrono::Utc
.timestamp_millis_opt(1_700_000_000_000 + seq as i64)
.unwrap();
Chunk {
id: format!("chunk_{seq}"),
content: format!("## ts — alice\nMessage {seq}"),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: None,
},
token_count: 5,
seq_in_source: seq,
created_at: ts,
partial_message: false,
}
}
#[test]
fn stage_chunks_writes_files_and_returns_staged() {
let dir = TempDir::new().unwrap();
let chunks = vec![sample_chunk(0), sample_chunk(1)];
let staged = stage_chunks(dir.path(), &chunks).unwrap();
assert_eq!(staged.len(), 2);
for s in &staged {
let abs = paths::chunk_abs_path(
dir.path(),
s.chunk.metadata.source_kind.as_str(),
&s.chunk.metadata.source_id,
&s.chunk.id,
);
assert!(abs.exists(), "file must exist: {}", abs.display());
assert!(!s.content_path.is_empty());
assert_eq!(s.content_sha256.len(), 64);
// Path must be relative with forward slashes.
assert!(!s.content_path.starts_with('/'));
assert!(s.content_path.contains('/'));
}
}
#[test]
fn stage_chunks_is_idempotent() {
let dir = TempDir::new().unwrap();
let chunks = vec![sample_chunk(0)];
let first = stage_chunks(dir.path(), &chunks).unwrap();
let second = stage_chunks(dir.path(), &chunks).unwrap();
assert_eq!(first[0].content_sha256, second[0].content_sha256);
assert_eq!(first[0].content_path, second[0].content_path);
}
}
@@ -0,0 +1,473 @@
//! Content-file path generation.
//!
//! Each chunk body is stored as a `.md` file under `<content_root>/`. The path
//! structure depends on the source kind:
//!
//! ```text
//! Email: <content_root>/email/<participants_slug>/<chunk_id>.md
//! Chat: <content_root>/chat/<source_slug>/<chunk_id>.md
//! Document: <content_root>/document/<source_slug>/<chunk_id>.md
//! ```
//!
//! Email paths parse `source_id` as `gmail:{participants}` where `participants`
//! is `addr1|addr2|...` (sorted, deduped, lowercased bare emails). The
//! participants string is slugified as a whole (pipe and `@` both become `-`)
//! to produce a single directory level, giving one folder per unique
//! conversation set.
//!
//! Paths are stored in SQLite as **relative** strings with forward slashes so
//! they remain valid regardless of where the workspace is mounted.
use std::path::{Path, PathBuf};
use chrono::{DateTime, Utc};
use crate::openhuman::memory::tree::util::redact::redact;
/// Which kind of summary tree a summary belongs to. Determines the top-level
/// directory under `<content_root>/summaries/`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SummaryTreeKind {
/// Per-source-tree summary. Layout: `summaries/source/<scope_slug>/L<level>/<id>.md`
Source,
/// Global digest tree. Layout: `summaries/global/<yyyy-mm-dd>/L<level>/<id>.md`
Global,
/// Per-topic (entity) tree. Layout: `summaries/topic/<scope_slug>/L<level>/<id>.md`
Topic,
}
/// Build the relative content path for a summary, using forward slashes.
///
/// Path layout depends on tree_kind:
/// - Source: `"summaries/source/<scope_slug>/L<level>/<summary_filename>.md"`
/// - Global: `"summaries/global/<yyyy-mm-dd>/L<level>/<summary_filename>.md"`
/// Falls back to `unknown-date` (with a warn log) if `date_for_global` is
/// `None` — preferable to panicking inside a path utility.
/// - Topic: `"summaries/topic/<scope_slug>/L<level>/<summary_filename>.md"`
///
/// `scope_slug` must already be slugified by the caller (use [`slugify_source_id`] or
/// a per-kind variant). A trailing `.md` on `summary_id` is stripped if present.
///
/// The `summary_id` is sanitized into a filesystem-safe filename by replacing
/// characters illegal on Windows (`:`, `\`, `*`, `?`, `"`, `<`, `>`, `|`) with `-`.
pub fn summary_rel_path(
tree_kind: SummaryTreeKind,
scope_slug: &str,
level: u32,
summary_id: &str,
date_for_global: Option<DateTime<Utc>>,
) -> String {
// Strip a trailing `.md` from summary_id if accidentally included.
let id = summary_id.strip_suffix(".md").unwrap_or(summary_id);
// Sanitize to a cross-platform filename (colons are illegal on Windows NTFS).
let filename = sanitize_filename(id);
match tree_kind {
SummaryTreeKind::Source => {
format!("summaries/source/{}/L{}/{}.md", scope_slug, level, filename)
}
SummaryTreeKind::Global => {
// Fall back to a sentinel date rather than panic — a path-utility
// panic would propagate up through seal/digest/janitor codepaths
// and abort otherwise-recoverable work. Callers should always
// pass a date for Global; the warn log surfaces the contract
// violation without taking the process down.
let date_str = match date_for_global {
Some(d) => d.format("%Y-%m-%d").to_string(),
None => {
log::warn!(
"[content_store::paths] summary_rel_path called for Global \
without date_for_global; using sentinel 'unknown-date'. \
Caller bug — please pass a date."
);
"unknown-date".to_string()
}
};
format!("summaries/global/{}/L{}/{}.md", date_str, level, filename)
}
SummaryTreeKind::Topic => {
format!("summaries/topic/{}/L{}/{}.md", scope_slug, level, filename)
}
}
}
/// Replace characters that are illegal in filenames on Windows NTFS with `-`.
///
/// Illegal characters: `\`, `/`, `:`, `*`, `?`, `"`, `<`, `>`, `|`.
/// (Forward slash is not replaced since `summary_id` should not contain path
/// separators, but we sanitize it anyway for safety.)
fn sanitize_filename(s: &str) -> String {
s.chars()
.map(|c| match c {
'\\' | '/' | ':' | '*' | '?' | '"' | '<' | '>' | '|' => '-',
c => c,
})
.collect()
}
/// Build the absolute on-disk path for a summary given the content root.
pub fn summary_abs_path(
content_root: &Path,
tree_kind: SummaryTreeKind,
scope_slug: &str,
level: u32,
summary_id: &str,
date_for_global: Option<DateTime<Utc>>,
) -> PathBuf {
let rel = summary_rel_path(tree_kind, scope_slug, level, summary_id, date_for_global);
let mut abs = content_root.to_path_buf();
for component in rel.split('/') {
abs.push(component);
}
abs
}
/// Build the relative content path for a chunk, using forward slashes.
///
/// Path layout depends on source_kind:
/// - Email: `"email/<participants_slug>/<chunk_id>.md"`
/// Parses `source_id` as `gmail:{participants}` (two colon-separated parts)
/// where `participants` is `addr1|addr2|...` (sorted, deduped, lowercased).
/// The entire participants string is slugified as a single unit to produce
/// one folder level per conversation set (no nested thread subfolder).
/// If the source_id lacks a `gmail:` prefix or has no participants segment,
/// falls through to the chat/document layout using `slugify_source_id(source_id)`.
/// - Chat: `"chat/<source_slug>/<chunk_id>.md"`
/// - Document: `"document/<source_slug>/<chunk_id>.md"`
///
/// `chunk_id` — the deterministic content hash produced by `types::chunk_id`.
///
/// # Examples
///
/// ```text
/// chunk_rel_path("email", "gmail:alice@x.com|bob@y.com", "abc")
/// → "email/alice-x-com-bob-y-com/abc.md"
///
/// chunk_rel_path("email", "gmail:notifications@github.com|sanil@x.com", "def")
/// → "email/notifications-github-com-sanil-x-com/def.md"
///
/// chunk_rel_path("email", "legacyid", "xyz")
/// → "email/legacyid/xyz.md" (malformed — flat fallback)
/// ```
pub fn chunk_rel_path(source_kind: &str, source_id: &str, chunk_id: &str) -> String {
// Sanitize chunk_id into a cross-platform filename. Chunk IDs contain
// colons (e.g. `chat:slack:#eng:0`) which are illegal on Windows NTFS;
// replace illegal characters with `-` to match summary_rel_path behaviour.
let filename = sanitize_filename(chunk_id);
match source_kind {
"email" => {
// Expected format: "gmail:{participants}"
// Split on ':' — exactly 2 parts required; part[0] == "gmail".
let parts: Vec<&str> = source_id.splitn(2, ':').collect();
if parts.len() == 2 && parts[0] == "gmail" && !parts[1].is_empty() {
let participants_slug = slugify_source_id(parts[1]);
format!("email/{}/{}.md", participants_slug, filename)
} else {
// Malformed / legacy source_id — fall back to flat layout.
// Redact the source_id before logging since it may embed email
// addresses.
log::debug!(
"[content_store::paths] email source_id has unexpected format, falling back to flat layout: source_id_hash={}",
redact(source_id)
);
let slug = slugify_source_id(source_id);
format!("email/{}/{}.md", slug, filename)
}
}
_ => {
// Chat, Document, and any future kinds use a 3-level layout.
let slug = slugify_source_id(source_id);
format!("{}/{}/{}.md", source_kind, slug, filename)
}
}
}
/// Build the absolute on-disk path for a chunk given the content root.
pub fn chunk_abs_path(
content_root: &Path,
source_kind: &str,
source_id: &str,
chunk_id: &str,
) -> PathBuf {
let rel = chunk_rel_path(source_kind, source_id, chunk_id);
// Convert forward-slash relative path to OS-native path.
let mut abs = content_root.to_path_buf();
for component in rel.split('/') {
abs.push(component);
}
abs
}
/// Convert a raw `source_id` (e.g. `"slack:#general"`, `"gmail:thread/abc"`)
/// into a filesystem-safe slug using only `[a-z0-9_-]` characters.
///
/// Rules:
/// - lowercase the whole string
/// - replace any character outside `[a-z0-9_-]` with `-`
/// - collapse consecutive `-` to one
/// - trim leading/trailing `-`
/// - `_` is preserved anywhere in the string (interior underscores are kept)
/// - truncate to 120 characters
pub fn slugify_source_id(source_id: &str) -> String {
let lower = source_id.to_lowercase();
let mut out = String::with_capacity(lower.len().min(120));
let mut last_dash = true; // avoids leading dash; also suppresses leading underscore runs
let mut pending_underscore = false; // deferred `_` to avoid leading underscore
for ch in lower.chars() {
if ch == '_' {
// Defer underscores — emit only if we have already emitted a
// non-separator character (so `_solo_` becomes `_solo_` once the
// `s` is emitted, but a leading `_` is dropped).
if !last_dash {
// We have real content before this, so emit the underscore now.
pending_underscore = true;
}
// If last_dash is true (nothing emitted yet), silently skip.
} else if ch.is_ascii_alphanumeric() {
if pending_underscore {
out.push('_');
pending_underscore = false;
}
out.push(ch);
last_dash = false;
} else {
// Non-alphanumeric, non-underscore → convert to `-`.
pending_underscore = false; // drop any pending underscore before a dash
if !last_dash {
out.push('-');
last_dash = true;
}
}
}
// trailing underscore: drop it (trim trailing separators).
// trim trailing dash
let trimmed = out.trim_end_matches('-');
// also trim any trailing underscore
let trimmed = trimmed.trim_end_matches('_');
let truncated = truncate_at_char(trimmed, 120);
if truncated.is_empty() {
"unknown".to_string()
} else {
truncated.to_string()
}
}
/// Truncate `s` to at most `max_chars` Unicode code points.
fn truncate_at_char(s: &str, max_chars: usize) -> &str {
match s.char_indices().nth(max_chars) {
Some((idx, _)) => &s[..idx],
None => s,
}
}
#[cfg(test)]
mod tests {
use super::*;
// ─── slugify tests ────────────────────────────────────────────────────────
#[test]
fn slugify_slack_channel() {
assert_eq!(slugify_source_id("slack:#general"), "slack-general");
}
#[test]
fn slugify_gmail_thread() {
assert_eq!(
slugify_source_id("gmail:thread/abc-123"),
"gmail-thread-abc-123"
);
}
#[test]
fn slugify_collapses_consecutive_separators() {
assert_eq!(slugify_source_id("foo::bar"), "foo-bar");
}
#[test]
fn slugify_uppercase_lowercased() {
assert_eq!(slugify_source_id("Slack:ABC"), "slack-abc");
}
#[test]
fn slugify_empty_falls_back_to_unknown() {
assert_eq!(slugify_source_id(""), "unknown");
assert_eq!(slugify_source_id(":::"), "unknown");
}
#[test]
fn slugify_truncates_at_120_chars() {
let long = "a".repeat(200);
let slug = slugify_source_id(&long);
assert_eq!(slug.len(), 120);
}
#[test]
fn slugify_preserves_interior_underscore() {
// `_solo_` has a leading and trailing underscore; only the interior
// `solo` + the part after should survive. When used as a thread key
// it arrives as the whole string `_solo_`.
// Leading `_` is stripped (it's treated like a leading dash),
// trailing `_` is stripped; interior `_` is preserved when sandwiched
// between alphanumeric characters.
let s = slugify_source_id("_solo_");
// "solo" — both outer underscores trimmed, interior underscore has
// nothing on the right so it's also trailing and trimmed.
assert_eq!(s, "solo");
}
#[test]
fn slugify_preserves_interior_underscore_between_chars() {
// `foo_bar` — interior underscore stays.
assert_eq!(slugify_source_id("foo_bar"), "foo_bar");
}
// ─── chunk_rel_path tests ─────────────────────────────────────────────────
#[test]
fn email_one_to_one_conversation_path() {
// 1:1 conversation between alice and bob.
let p = chunk_rel_path("email", "gmail:alice@x.com|bob@y.com", "abc");
assert_eq!(p, "email/alice-x-com-bob-y-com/abc.md");
}
#[test]
fn email_group_conversation_path() {
// Group conversation with three participants.
let p = chunk_rel_path("email", "gmail:notifications@github.com|sanil@x.com", "def");
assert_eq!(p, "email/notifications-github-com-sanil-x-com/def.md");
}
#[test]
fn email_solo_no_to_path() {
// Solo sender (no To), participants = single address.
let p = chunk_rel_path("email", "gmail:alice@x.com", "solo123");
assert_eq!(p, "email/alice-x-com/solo123.md");
}
#[test]
fn email_malformed_source_id_falls_back_to_flat_layout() {
// Malformed: no `gmail:` prefix → flat fallback.
let p = chunk_rel_path("email", "legacyid", "xyz");
// Falls back to email/<slug>/<chunk_id>.md
assert!(p.starts_with("email/"), "must remain under email/");
assert!(p.ends_with("/xyz.md"), "chunk_id must be the filename");
// Must not panic.
}
#[test]
fn email_three_participant_path() {
// Three participants: alice, bob, carol (pipe-separated, sorted).
let p = chunk_rel_path("email", "gmail:alice@x.com|bob@y.com|carol@z.com", "g42");
assert_eq!(p, "email/alice-x-com-bob-y-com-carol-z-com/g42.md");
}
#[test]
fn chat_path() {
let p = chunk_rel_path("chat", "slack:#eng", "xyz789");
assert_eq!(p, "chat/slack-eng/xyz789.md");
}
#[test]
fn document_path() {
let p = chunk_rel_path("document", "doc:notes.md", "uvw");
assert_eq!(p, "document/doc-notes-md/uvw.md");
}
#[test]
fn chunk_abs_path_uses_os_separator() {
use std::path::Path;
let root = Path::new("/workspace/content");
let abs = chunk_abs_path(root, "email", "gmail:alice@x.com|bob@y.com", "abc");
assert!(abs.starts_with(root));
assert!(abs.ends_with("abc.md"));
}
// ─── summary_rel_path tests ───────────────────────────────────────────────
#[test]
fn summary_rel_path_source() {
let p = summary_rel_path(
SummaryTreeKind::Source,
"gmail-alice-x-com-bob-y-com",
1,
"summary:L1:abc",
None,
);
// Colons in summary_id are replaced with '-' for cross-platform filenames.
assert_eq!(
p,
"summaries/source/gmail-alice-x-com-bob-y-com/L1/summary-L1-abc.md"
);
}
#[test]
fn summary_rel_path_global() {
use chrono::TimeZone;
let date = chrono::Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap();
let p = summary_rel_path(
SummaryTreeKind::Global,
"global",
0,
"summary:L0:daily",
Some(date),
);
assert_eq!(p, "summaries/global/2026-04-28/L0/summary-L0-daily.md");
}
#[test]
fn summary_rel_path_topic() {
let p = summary_rel_path(
SummaryTreeKind::Topic,
"person-alex-johnson",
1,
"summary:L1:xyz",
None,
);
assert_eq!(
p,
"summaries/topic/person-alex-johnson/L1/summary-L1-xyz.md"
);
}
#[test]
fn summary_rel_path_strips_trailing_md_extension() {
// If the caller accidentally appends .md to the summary_id, strip it.
let p = summary_rel_path(
SummaryTreeKind::Topic,
"entity-slug",
2,
"summary:L2:foo.md",
None,
);
assert_eq!(p, "summaries/topic/entity-slug/L2/summary-L2-foo.md");
}
#[test]
fn summary_rel_path_global_falls_back_to_sentinel_without_date() {
// Caller bug to omit date for Global, but a path utility shouldn't
// panic — fall back to a sentinel `unknown-date` segment so the
// file lands somewhere predictable rather than aborting the seal.
let p = summary_rel_path(SummaryTreeKind::Global, "global", 0, "summary:L0:x", None);
assert_eq!(p, "summaries/global/unknown-date/L0/summary-L0-x.md");
}
#[test]
fn summary_abs_path_rooted_under_content_root() {
use chrono::TimeZone;
use std::path::Path;
let root = Path::new("/workspace/content");
let date = chrono::Utc.with_ymd_and_hms(2026, 1, 15, 0, 0, 0).unwrap();
let abs = summary_abs_path(
root,
SummaryTreeKind::Global,
"global",
0,
"daily-123",
Some(date),
);
assert!(abs.starts_with(root));
assert!(abs.ends_with("daily-123.md"));
}
}
@@ -0,0 +1,405 @@
//! Read and verify chunk and summary `.md` files from the content store.
use std::path::Path;
use super::atomic::sha256_hex;
use super::compose::split_front_matter;
use crate::openhuman::memory::tree::util::redact::redact;
/// The result of reading a chunk file from disk.
pub struct ChunkFileContents {
/// The Markdown body (everything after the closing `---` of the front-matter).
pub body: String,
/// SHA-256 hex digest over the **body bytes** only.
pub sha256: String,
}
/// Read a chunk file and return its body + SHA-256.
///
/// Returns an error if:
/// - the file does not exist
/// - the file is not valid UTF-8
/// - the front-matter delimiters cannot be found
pub fn read_chunk_file(abs_path: &Path) -> anyhow::Result<ChunkFileContents> {
let raw = std::fs::read(abs_path).map_err(|e| anyhow::anyhow!("read {:?}: {e}", abs_path))?;
let content = std::str::from_utf8(&raw)
.map_err(|e| anyhow::anyhow!("invalid UTF-8 in {:?}: {e}", abs_path))?;
let (_fm, body) = split_front_matter(content)
.ok_or_else(|| anyhow::anyhow!("no front-matter in {:?}", abs_path))?;
let sha256 = sha256_hex(body.as_bytes());
Ok(ChunkFileContents {
body: body.to_string(),
sha256,
})
}
/// Verify that the body of a chunk file matches the expected SHA-256.
///
/// Returns `Ok(true)` on a match, `Ok(false)` on a mismatch, and an `Err`
/// if the file cannot be read or parsed.
pub fn verify_chunk_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Result<bool> {
let contents = read_chunk_file(abs_path)?;
let ok = contents.sha256 == expected_sha256;
if !ok {
// Log the path as a redacted hash — the path may embed email addresses
// (participant slugs) after the participant-bucketing change.
let path_str = abs_path.to_string_lossy();
log::warn!(
"[content_store::read] sha256 mismatch for path_hash={}: expected={} actual={}",
redact(&path_str),
expected_sha256,
contents.sha256,
);
}
Ok(ok)
}
// ── Summary reads ────────────────────────────────────────────────────────────
/// The result of verifying a summary file on disk.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum VerifyResult {
/// The on-disk body SHA-256 matches the stored value.
Ok,
/// The file exists but the body SHA-256 does not match.
Mismatch { actual: String },
/// The file does not exist at the given path.
Missing,
}
/// Read a summary file and return its body + SHA-256.
///
/// Returns an error if:
/// - the file does not exist
/// - the file is not valid UTF-8
/// - the front-matter delimiters cannot be found
pub fn read_summary_file(abs_path: &Path) -> anyhow::Result<ChunkFileContents> {
// Reuse the same reader as chunks — the file format is identical.
read_chunk_file(abs_path)
}
/// Verify a summary file's body SHA-256 without returning the body itself.
///
/// Returns:
/// - `VerifyResult::Ok` on match
/// - `VerifyResult::Mismatch { actual }` on hash mismatch
/// - `VerifyResult::Missing` when the file does not exist
pub fn verify_summary_file(abs_path: &Path, expected_sha256: &str) -> anyhow::Result<VerifyResult> {
if !abs_path.exists() {
return Ok(VerifyResult::Missing);
}
let contents = read_summary_file(abs_path)?;
if contents.sha256 == expected_sha256 {
Ok(VerifyResult::Ok)
} else {
// Redact the path — it can embed participant slugs (email addresses).
let path_str = abs_path.to_string_lossy();
log::warn!(
"[content_store::read] sha256 mismatch for summary path_hash={}: expected={} actual={}",
redact(&path_str),
expected_sha256,
contents.sha256,
);
Ok(VerifyResult::Mismatch {
actual: contents.sha256,
})
}
}
// ── High-level body readers (Config-aware) ───────────────────────────────────
//
// These helpers resolve the on-disk path from SQLite via
// `get_chunk_content_pointers` / `get_summary_content_pointers`, then read the
// file body. They are the single authoritative entry-point for every caller
// that needs the **full** chunk or summary body (LLM extractor, summariser
// inputs, retrieval API, embedder). Preview-only consumers (UI cards, fast
// filter scans) continue reading the `content` column directly from SQLite.
//
// Error policy:
// - If `content_path` / `content_sha256` are NULL (legacy rows ingested before
// the MD-on-disk migration), return `Err` — callers must handle the
// "pre-migration chunk" case explicitly. The job pipeline propagates the
// error and retries; retrieval falls back gracefully.
// - File-not-found or SHA mismatch → `Err` (propagated to caller for retry /
// alerting).
/// Read the full body of a chunk `.md` file by its chunk id.
///
/// Looks up `content_path` in SQLite, resolves it to an absolute path under
/// `config.memory_tree_content_root()`, reads the file, and returns the body
/// string (everything after the YAML front-matter delimiter).
///
/// Returns `Err` if:
/// - The chunk row has no `content_path` recorded (pre-MD-migration row).
/// - The file cannot be read or has no valid front-matter.
///
/// # Preview vs. full body
/// The `content` column in `mem_tree_chunks` holds a ≤500-char preview after
/// the MD-on-disk migration. Use this function wherever the full body is
/// required (LLM extraction, embedding, summariser inputs, retrieval API).
pub fn read_chunk_body(
config: &crate::openhuman::config::Config,
chunk_id: &str,
) -> anyhow::Result<String> {
use crate::openhuman::memory::tree::store::get_chunk_content_pointers;
let pointers = get_chunk_content_pointers(config, chunk_id)?.ok_or_else(|| {
anyhow::anyhow!(
"[content_store::read] no content_path for chunk_id={} (pre-MD-migration row?)",
chunk_id
)
})?;
let (rel_path, expected_sha256) = pointers;
let content_root = config.memory_tree_content_root();
// Reconstruct the absolute path from the stored relative forward-slash path.
let abs_path = {
let mut p = content_root.clone();
for component in rel_path.split('/') {
p.push(component);
}
p
};
log::debug!(
"[content_store::read] read_chunk_body chunk_id={} path_hash={}",
chunk_id,
redact(&rel_path),
);
let result = read_chunk_file(&abs_path).with_context(|| {
format!(
"read_chunk_body: failed to read file for chunk_id={} path_hash={}",
chunk_id,
redact(&rel_path),
)
})?;
// Verify the on-disk body matches the SHA stored at write time. A mismatch
// means the file was tampered with, the tx that committed the pointer
// raced with a separate writer, or the disk corrupted — all unsafe to
// hand back to a consumer. Fail loudly rather than serve stale/corrupt
// bytes into the LLM extractor / summariser pipeline.
if result.sha256 != expected_sha256 {
return Err(anyhow::anyhow!(
"[content_store::read] sha256 mismatch for chunk_id={} \
expected={} actual={} path_hash={}",
chunk_id,
expected_sha256,
result.sha256,
redact(&rel_path),
));
}
Ok(result.body)
}
use anyhow::Context as _;
/// Read the full body of a summary `.md` file by its summary id.
///
/// Looks up `content_path` in SQLite, resolves it to an absolute path under
/// `config.memory_tree_content_root()`, reads the file, and returns the body
/// string.
///
/// Returns `Err` if:
/// - The summary row has no `content_path` recorded (pre-MD-migration row).
/// - The file cannot be read or has no valid front-matter.
///
/// # Preview vs. full body
/// The `content` column in `mem_tree_summaries` holds a ≤500-char preview after
/// the MD-on-disk migration. Use this function wherever the full body is
/// required (LLM extraction, embedding, summariser inputs, retrieval API).
pub fn read_summary_body(
config: &crate::openhuman::config::Config,
summary_id: &str,
) -> anyhow::Result<String> {
use crate::openhuman::memory::tree::store::get_summary_content_pointers;
let pointers = get_summary_content_pointers(config, summary_id)?.ok_or_else(|| {
anyhow::anyhow!(
"[content_store::read] no content_path for summary_id={} (pre-MD-migration row?)",
summary_id
)
})?;
let (rel_path, expected_sha256) = pointers;
let content_root = config.memory_tree_content_root();
let abs_path = {
let mut p = content_root.clone();
for component in rel_path.split('/') {
p.push(component);
}
p
};
log::debug!(
"[content_store::read] read_summary_body summary_id={} path_hash={}",
summary_id,
redact(&rel_path),
);
let result = read_summary_file(&abs_path).with_context(|| {
format!(
"read_summary_body: failed to read file for summary_id={} path_hash={}",
summary_id,
redact(&rel_path),
)
})?;
// Verify the on-disk body matches the SHA stored at seal time. See the
// matching guard in `read_chunk_body` for rationale.
if result.sha256 != expected_sha256 {
return Err(anyhow::anyhow!(
"[content_store::read] sha256 mismatch for summary_id={} \
expected={} actual={} path_hash={}",
summary_id,
expected_sha256,
result.sha256,
redact(&rel_path),
));
}
Ok(result.body)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new};
use crate::openhuman::memory::tree::content_store::compose::compose_chunk_file;
use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind};
use chrono::TimeZone;
use tempfile::TempDir;
fn sample_chunk() -> Chunk {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
Chunk {
id: "read_test".into(),
content: "## ts — alice\nhello from read test".into(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: None,
},
token_count: 8,
seq_in_source: 0,
created_at: ts,
partial_message: false,
}
}
#[test]
fn read_returns_body_and_correct_sha256() {
let dir = TempDir::new().unwrap();
let chunk = sample_chunk();
let (full_bytes, body_bytes) = compose_chunk_file(&chunk);
let path = dir.path().join("0.md");
write_if_new(&path, &full_bytes).unwrap();
let result = read_chunk_file(&path).unwrap();
assert_eq!(result.body, std::str::from_utf8(&body_bytes).unwrap());
assert_eq!(result.sha256, sha256_hex(&body_bytes));
}
#[test]
fn verify_passes_for_correct_hash() {
let dir = TempDir::new().unwrap();
let chunk = sample_chunk();
let (full_bytes, body_bytes) = compose_chunk_file(&chunk);
let path = dir.path().join("0.md");
write_if_new(&path, &full_bytes).unwrap();
let expected = sha256_hex(&body_bytes);
assert!(verify_chunk_file(&path, &expected).unwrap());
}
#[test]
fn verify_fails_for_wrong_hash() {
let dir = TempDir::new().unwrap();
let chunk = sample_chunk();
let (full_bytes, _) = compose_chunk_file(&chunk);
let path = dir.path().join("0.md");
write_if_new(&path, &full_bytes).unwrap();
assert!(!verify_chunk_file(&path, "deadbeef").unwrap());
}
#[test]
fn read_missing_file_returns_error() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("nonexistent.md");
assert!(read_chunk_file(&path).is_err());
}
// ─── summary read / verify tests ─────────────────────────────────────────
fn write_summary_file(dir: &TempDir, body: &str) -> (std::path::PathBuf, String) {
use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new};
use crate::openhuman::memory::tree::content_store::compose::{
compose_summary_md, SummaryComposeInput,
};
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
use chrono::TimeZone;
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let input = SummaryComposeInput {
summary_id: "sum:L1:readtest",
tree_kind: SummaryTreeKind::Source,
tree_id: "t1",
tree_scope: "gmail:alice@x.com",
level: 1,
child_ids: &["c1".to_string()],
child_count: 1,
time_range_start: ts,
time_range_end: ts,
sealed_at: ts,
body,
};
let composed = compose_summary_md(&input);
let path = dir.path().join("sum.md");
let sha = sha256_hex(composed.body.as_bytes());
write_if_new(&path, composed.full.as_bytes()).unwrap();
(path, sha)
}
#[test]
fn read_summary_file_returns_body_and_sha() {
let dir = TempDir::new().unwrap();
let body = "summary body content\n";
let (path, expected_sha) = write_summary_file(&dir, body);
let result = read_summary_file(&path).unwrap();
assert_eq!(result.body, body);
assert_eq!(result.sha256, expected_sha);
}
#[test]
fn verify_summary_file_ok_for_correct_hash() {
let dir = TempDir::new().unwrap();
let (path, sha) = write_summary_file(&dir, "body text\n");
assert_eq!(verify_summary_file(&path, &sha).unwrap(), VerifyResult::Ok);
}
#[test]
fn verify_summary_file_mismatch_for_wrong_hash() {
let dir = TempDir::new().unwrap();
let (path, _) = write_summary_file(&dir, "body text\n");
let r = verify_summary_file(&path, "deadbeef").unwrap();
assert!(matches!(r, VerifyResult::Mismatch { .. }));
}
#[test]
fn verify_summary_file_missing_for_absent_file() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("does_not_exist.md");
assert_eq!(
verify_summary_file(&path, "abc").unwrap(),
VerifyResult::Missing
);
}
}
@@ -0,0 +1,420 @@
//! Post-extraction tag rewriting for chunk and summary `.md` files.
//!
//! After the LLM extraction job runs, it produces a list of entities. Each
//! entity is converted to an Obsidian-style hierarchical tag (`kind/Value`)
//! and written into the `tags:` block in the file's front-matter.
//!
//! The body bytes (and therefore the SHA-256) are never changed — only the
//! front-matter is rewritten.
use std::path::Path;
use super::compose::{rewrite_summary_tags as compose_rewrite_summary_tags, rewrite_tags};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::score::store::list_entity_ids_for_node;
use crate::openhuman::memory::tree::store::get_summary_content_pointers;
/// Rewrite the `tags:` block in a chunk's on-disk `.md` file.
///
/// `abs_path` — absolute path to the chunk file.
/// `tags` — new list of tag strings (Obsidian `kind/Value` format).
///
/// The operation is atomic: the new file is written to a sibling temp path and
/// then renamed over the original. If the file does not exist, the call is a
/// no-op (returns `Ok(())`).
///
/// Note: unlike the initial chunk write, tag rewrites MAY overwrite an
/// existing file. The immutability contract covers the **body** only; tags are
/// explicitly designed to be updated post-extraction.
pub fn update_chunk_tags(abs_path: &Path, tags: &[String]) -> anyhow::Result<()> {
if !abs_path.exists() {
log::debug!(
"[content_store::tags] skipping tag update — file not found: {}",
abs_path.display()
);
return Ok(());
}
let old_bytes =
std::fs::read(abs_path).map_err(|e| anyhow::anyhow!("read {:?}: {e}", abs_path))?;
let new_bytes = rewrite_tags(&old_bytes, tags)
.map_err(|e| anyhow::anyhow!("rewrite_tags {:?}: {e}", abs_path))?;
// Write the new content atomically via a sibling temp file.
let parent = abs_path.parent().unwrap_or_else(|| Path::new("."));
let tmp_name = format!(".tmp_tags_{}.md", crate_temp_id());
let tmp_path = parent.join(&tmp_name);
{
use std::io::Write;
let mut f = std::fs::File::create(&tmp_path)
.map_err(|e| anyhow::anyhow!("create tag-rewrite tempfile {:?}: {e}", tmp_path))?;
f.write_all(&new_bytes)
.map_err(|e| anyhow::anyhow!("write tag-rewrite tempfile {:?}: {e}", tmp_path))?;
f.sync_all()
.map_err(|e| anyhow::anyhow!("fsync tag-rewrite tempfile {:?}: {e}", tmp_path))?;
}
std::fs::rename(&tmp_path, abs_path).map_err(|e| {
let _ = std::fs::remove_file(&tmp_path);
anyhow::anyhow!("rename tag-rewrite {:?} -> {:?}: {e}", tmp_path, abs_path)
})?;
log::debug!(
"[content_store::tags] updated tags in {}",
abs_path.display()
);
Ok(())
}
/// Rewrite the `tags:` block in a summary's on-disk `.md` file.
///
/// Reads entity rows from `mem_tree_entity_index` for `summary_id`, converts
/// them to `kind/Value` Obsidian tags, rewrites the YAML `tags:` block
/// atomically (tempfile + fsync + rename), and verifies the body SHA-256 is
/// unchanged afterwards.
///
/// Best-effort: tag-rewrite failures should not fail the extraction job. Callers
/// should log a warning and continue — the entity index is the authoritative source.
pub fn update_summary_tags(config: &Config, summary_id: &str) -> anyhow::Result<()> {
// 1. Fetch content_path from SQLite.
let pointers = get_summary_content_pointers(config, summary_id)?;
let (rel_path, expected_sha) = match pointers {
Some(p) => p,
None => {
log::debug!(
"[content_store::tags] update_summary_tags: no content_path for summary {summary_id} — skipping"
);
return Ok(());
}
};
let content_root = config.memory_tree_content_root();
let abs_path = {
let mut p = content_root;
for component in rel_path.split('/') {
p.push(component);
}
p
};
if !abs_path.exists() {
log::debug!(
"[content_store::tags] update_summary_tags: file missing for summary {summary_id} \
at {} — skipping",
abs_path.display()
);
return Ok(());
}
// 2. Fetch entity_index rows and build the merged tag list.
let entity_ids = list_entity_ids_for_node(config, summary_id)?;
let tags: Vec<String> = entity_ids
.iter()
.filter_map(|eid| {
// entity_id format: "kind:surface"
let (kind, surface) = eid.split_once(':')?;
Some(entity_tag(kind, surface))
})
.collect();
// Sort + dedup for stability.
let mut tags = tags;
tags.sort();
tags.dedup();
// 3. Read + atomic rewrite of the front-matter `tags:` block.
let old_bytes = std::fs::read(&abs_path)
.map_err(|e| anyhow::anyhow!("read summary {:?}: {e}", abs_path))?;
let new_bytes = compose_rewrite_summary_tags(&old_bytes, &tags)
.map_err(|e| anyhow::anyhow!("rewrite_summary_tags {:?}: {e}", abs_path))?;
let parent = abs_path.parent().unwrap_or_else(|| Path::new("."));
let tmp_name = format!(".tmp_sum_tags_{}.md", crate_temp_id());
let tmp_path = parent.join(&tmp_name);
{
use std::io::Write;
let mut f = std::fs::File::create(&tmp_path).map_err(|e| {
anyhow::anyhow!("create summary tag-rewrite tempfile {:?}: {e}", tmp_path)
})?;
f.write_all(&new_bytes).map_err(|e| {
anyhow::anyhow!("write summary tag-rewrite tempfile {:?}: {e}", tmp_path)
})?;
f.sync_all().map_err(|e| {
anyhow::anyhow!("fsync summary tag-rewrite tempfile {:?}: {e}", tmp_path)
})?;
}
std::fs::rename(&tmp_path, &abs_path).map_err(|e| {
let _ = std::fs::remove_file(&tmp_path);
anyhow::anyhow!(
"rename summary tag-rewrite {:?} -> {:?}: {e}",
tmp_path,
abs_path
)
})?;
// 4. Sanity check: body sha must still match after the rewrite.
let verify_bytes = std::fs::read(&abs_path)
.map_err(|e| anyhow::anyhow!("re-read after tag rewrite {:?}: {e}", abs_path))?;
let content = std::str::from_utf8(&verify_bytes)
.map_err(|e| anyhow::anyhow!("UTF-8 after tag rewrite {:?}: {e}", abs_path))?;
let body_after = super::compose::split_front_matter(content)
.ok_or_else(|| anyhow::anyhow!("no front-matter after tag rewrite {:?}", abs_path))?
.1;
let actual_sha = super::atomic::sha256_hex(body_after.as_bytes());
if actual_sha != expected_sha {
return Err(anyhow::anyhow!(
"[content_store::tags] update_summary_tags body mutated after rewrite \
summary_id={summary_id} expected_sha={expected_sha} actual_sha={actual_sha}"
));
}
log::debug!(
"[content_store::tags] updated {} tags in summary file summary_id={summary_id} n_tags={}",
tags.len(),
tags.len()
);
Ok(())
}
/// Slugify an entity kind string for use in an Obsidian hierarchical tag.
///
/// Output: lowercase, spaces and non-alphanumeric chars replaced with `-`,
/// consecutive dashes collapsed, leading/trailing dashes stripped.
///
/// Example: `"Person"` → `"person"`, `"GitHub Repo"` → `"github-repo"`
pub fn slugify_tag_kind(kind: &str) -> String {
slugify_tag_component(kind)
}
/// Slugify an entity value string for use in an Obsidian hierarchical tag.
///
/// Like `slugify_tag_kind`, but capitalises the first letter of each word
/// so values are visually distinct from kinds:
///
/// `"alice johnson"` → `"Alice-Johnson"`,
/// `"project Phoenix"` → `"Project-Phoenix"`
pub fn slugify_tag_value(value: &str) -> String {
// Split on non-alphanumeric boundaries, capitalise first letter of each word.
let mut parts: Vec<String> = Vec::new();
let mut current = String::new();
for ch in value.chars() {
if ch.is_alphanumeric() || ch == '_' {
current.push(ch);
} else if !current.is_empty() {
parts.push(capitalise(&current));
current.clear();
}
}
if !current.is_empty() {
parts.push(capitalise(&current));
}
let joined = parts.join("-");
if joined.is_empty() {
"unknown".to_string()
} else {
joined
}
}
/// Build an Obsidian-style `kind/Value` tag string from raw entity kind + surface.
pub fn entity_tag(kind: &str, surface: &str) -> String {
format!("{}/{}", slugify_tag_kind(kind), slugify_tag_value(surface))
}
fn slugify_tag_component(s: &str) -> String {
let lower = s.to_lowercase();
let mut out = String::new();
let mut last_dash = true;
for ch in lower.chars() {
if ch.is_ascii_alphanumeric() || ch == '_' {
out.push(ch);
last_dash = false;
} else if !last_dash {
out.push('-');
last_dash = true;
}
}
let trimmed = out.trim_end_matches('-');
if trimmed.is_empty() {
"unknown".to_string()
} else {
trimmed.to_string()
}
}
fn capitalise(s: &str) -> String {
let mut chars = s.chars();
match chars.next() {
None => String::new(),
Some(first) => {
let upper: String = first.to_uppercase().collect();
upper + chars.as_str()
}
}
}
fn crate_temp_id() -> String {
use std::time::{SystemTime, UNIX_EPOCH};
let ns = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.subsec_nanos();
format!("{ns:08x}")
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store::atomic::{sha256_hex, write_if_new};
use crate::openhuman::memory::tree::content_store::compose::compose_chunk_file;
use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind};
use chrono::TimeZone;
use tempfile::TempDir;
fn sample_chunk() -> Chunk {
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
Chunk {
id: "tags_test".into(),
content: "hello from tags test".into(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec!["old/Tag".into()],
source_ref: None,
},
token_count: 4,
seq_in_source: 0,
created_at: ts,
partial_message: false,
}
}
#[test]
fn update_chunk_tags_replaces_tag_block() {
let dir = TempDir::new().unwrap();
let chunk = sample_chunk();
let (full, _) = compose_chunk_file(&chunk);
let path = dir.path().join("0.md");
write_if_new(&path, &full).unwrap();
update_chunk_tags(
&path,
&["person/Alice-Smith".into(), "project/Phoenix".into()],
)
.unwrap();
let updated = std::fs::read_to_string(&path).unwrap();
assert!(updated.contains(" - person/Alice-Smith"));
assert!(updated.contains(" - project/Phoenix"));
assert!(!updated.contains(" - old/Tag"));
// Body unchanged.
assert!(updated.ends_with("hello from tags test"));
}
#[test]
fn update_chunk_tags_is_noop_for_missing_file() {
let dir = TempDir::new().unwrap();
let path = dir.path().join("nonexistent.md");
assert!(update_chunk_tags(&path, &["p/X".into()]).is_ok());
}
#[test]
fn slugify_tag_kind_examples() {
assert_eq!(slugify_tag_kind("Person"), "person");
assert_eq!(slugify_tag_kind("GitHub Repo"), "github-repo");
assert_eq!(slugify_tag_kind("EMAIL"), "email");
}
#[test]
fn slugify_tag_value_capitalises_words() {
assert_eq!(slugify_tag_value("alice johnson"), "Alice-Johnson");
assert_eq!(slugify_tag_value("project Phoenix"), "Project-Phoenix");
assert_eq!(slugify_tag_value("OPENAI"), "OPENAI");
}
#[test]
fn entity_tag_builds_obsidian_tag() {
assert_eq!(
entity_tag("person", "Alice Johnson"),
"person/Alice-Johnson"
);
assert_eq!(entity_tag("ORG", "Tinyhumans AI"), "org/Tinyhumans-AI");
}
// ─── update_summary_tags tests ────────────────────────────────────────────
/// Write a summary .md file to disk with empty tags and verify rewriting works.
#[test]
fn rewrite_summary_tags_preserves_body_and_replaces_tags() {
use crate::openhuman::memory::tree::content_store::compose::{
compose_summary_md, SummaryComposeInput,
};
use crate::openhuman::memory::tree::content_store::paths::SummaryTreeKind;
let dir = TempDir::new().unwrap();
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let body = "summary body for tag test\n";
let children = vec!["c1".to_string()];
let input = SummaryComposeInput {
summary_id: "sum:L1:tagtest",
tree_kind: SummaryTreeKind::Source,
tree_id: "t1",
tree_scope: "gmail:alice@x.com",
level: 1,
child_ids: &children,
child_count: 1,
time_range_start: ts,
time_range_end: ts,
sealed_at: ts,
body,
};
let composed = compose_summary_md(&input);
let path = dir.path().join("sum.md");
write_if_new(&path, composed.full.as_bytes()).unwrap();
// Original must have `tags: []`
let original = std::fs::read_to_string(&path).unwrap();
assert!(original.contains("tags: []"));
// Rewrite the tags block
let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()];
let file_bytes = std::fs::read(&path).unwrap();
let rewritten = super::compose_rewrite_summary_tags(&file_bytes, &new_tags).unwrap();
// Write rewritten bytes back (simulating atomic rewrite)
let tmp = dir.path().join("sum.tmp.md");
{
use std::io::Write;
let mut f = std::fs::File::create(&tmp).unwrap();
f.write_all(&rewritten).unwrap();
}
std::fs::rename(&tmp, &path).unwrap();
let updated = std::fs::read_to_string(&path).unwrap();
assert!(updated.contains(" - person/Alice-Smith"));
assert!(updated.contains(" - topic/Memory"));
assert!(!updated.contains("tags: []"));
// Body unchanged
assert!(updated.ends_with(body));
// Body sha unchanged
use crate::openhuman::memory::tree::content_store::compose::split_front_matter;
let (_, body_after) = split_front_matter(&updated).unwrap();
let sha = sha256_hex(body_after.as_bytes());
let expected_sha = sha256_hex(body.as_bytes());
assert_eq!(
sha, expected_sha,
"body sha must be stable after tag rewrite"
);
}
}
@@ -19,11 +19,17 @@
//! - Idempotency: if an L0 daily node already exists for the target day,
//! return `DigestOutcome::Skipped` rather than emitting a duplicate.
use std::collections::BTreeSet;
use anyhow::{Context, Result};
use chrono::{DateTime, Duration, NaiveDate, TimeZone, Utc};
use rusqlite::OptionalExtension;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::{
atomic::stage_summary, paths::slugify_source_id, read as content_read, SummaryComposeInput,
SummaryTreeKind,
};
use crate::openhuman::memory::tree::global_tree::registry::get_or_create_global_tree;
use crate::openhuman::memory::tree::global_tree::seal::append_daily_and_cascade;
use crate::openhuman::memory::tree::global_tree::GLOBAL_TOKEN_BUDGET;
@@ -154,6 +160,25 @@ pub async fn end_of_day_digest(
.await
.context("embed daily summary during end_of_day_digest")?;
// L0 daily node inherits entities/topics by union of contributing
// source-tree summaries. Each input was already labeled at source-tree
// seal time, so emergent themes don't need another extractor pass
// here — global is a sink; union preserves "days that mentioned X"
// retrieval without an extra LLM call. See LabelStrategy in
// source_tree::bucket_seal for the full design.
let mut entities_set: BTreeSet<String> = BTreeSet::new();
let mut topics_set: BTreeSet<String> = BTreeSet::new();
for inp in &inputs {
for e in &inp.entities {
entities_set.insert(e.clone());
}
for t in &inp.topics {
topics_set.insert(t.clone());
}
}
let daily_entities: Vec<String> = entities_set.into_iter().collect();
let daily_topics: Vec<String> = topics_set.into_iter().collect();
let now = Utc::now();
let daily_id = new_summary_id(0);
let daily = SummaryNode {
@@ -165,8 +190,8 @@ pub async fn end_of_day_digest(
child_ids: inputs.iter().map(|i| i.id.clone()).collect(),
content: output.content,
token_count: output.token_count,
entities: output.entities,
topics: output.topics,
entities: daily_entities,
topics: daily_topics,
time_range_start: day_start,
time_range_end: day_end,
score,
@@ -175,6 +200,44 @@ pub async fn end_of_day_digest(
embedding: Some(embedding),
};
// Phase MD-content: stage the L0 daily .md file before the write tx.
// `date_for_global` = day_start (the calendar day this digest covers).
let daily_compose_input = SummaryComposeInput {
summary_id: &daily.id,
tree_kind: SummaryTreeKind::Global,
tree_id: &daily.tree_id,
tree_scope: &global.scope,
level: daily.level,
child_ids: &daily.child_ids,
child_count: daily.child_ids.len(),
time_range_start: daily.time_range_start,
time_range_end: daily.time_range_end,
sealed_at: daily.sealed_at,
body: &daily.content,
};
// Stage the summary .md file — abort the digest on failure so the database
// never commits a row with content_path = NULL. The digest job is retried
// via the normal job-retry path.
let content_root_daily = config.memory_tree_content_root();
let global_scope_slug = slugify_source_id(&global.scope);
let staged_daily = stage_summary(
&content_root_daily,
&daily_compose_input,
&global_scope_slug,
Some(day_start),
)
.with_context(|| {
format!(
"stage_summary failed for daily {}; digest aborted for retry",
daily.id
)
})?;
log::debug!(
"[global_tree::digest] staged daily summary {} → {}",
daily.id,
staged_daily.content_path
);
// Persist the daily node. Note: we do NOT backlink parent_id on the
// child summaries here — their parents are their own source trees, not
// the global tree. The global-tree child_ids are cross-source
@@ -183,7 +246,7 @@ pub async fn end_of_day_digest(
let tree_id_clone = global.id.clone();
with_connection(config, move |conn| {
let tx = conn.unchecked_transaction()?;
store::insert_summary_tx(&tx, &daily_clone)?;
store::insert_summary_tx(&tx, &daily_clone, Some(&staged_daily))?;
// Index any entities the summariser emitted (no-op under inert).
crate::openhuman::memory::tree::score::store::index_summary_entity_ids_tx(
&tx,
@@ -315,9 +378,23 @@ fn pick_source_contribution(
}
};
// Read the full body from disk — `node.content` is a ≤500-char preview
// after the MD-on-disk migration. The digest summariser must receive the
// complete summary text so the daily recap is not assembled from previews.
let body = match content_read::read_summary_body(config, &node.id) {
Ok(b) => b,
Err(e) => {
log::warn!(
"[global_tree::digest] read_summary_body failed for {} — using preview: {e:#}",
node.id
);
// Non-fatal: fall back to preview for pre-MD-migration rows.
node.content.clone()
}
};
Ok(Some(SummaryInput {
id: node.id,
content: format!("[{}]\n{}", source_tree.scope, node.content),
content: format!("[{}]\n{}", source_tree.scope, body),
token_count: node.token_count,
entities: node.entities,
topics: node.topics,
@@ -1,5 +1,8 @@
use super::*;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
use crate::openhuman::memory::tree::source_tree::types::TreeStatus;
@@ -7,6 +10,24 @@ use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
use tempfile::TempDir;
/// Stage a batch of chunks to the content store so that `read_chunk_body`
/// can find the on-disk file during seals. Tests that call `upsert_chunks`
/// and then trigger a seal MUST also call this helper; otherwise
/// `hydrate_leaf_inputs` will fail with "no content_path for chunk_id".
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged =
content_store::stage_chunks(&content_root, chunks).expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -39,6 +60,7 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime
token_count: 6_000,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
let c2 = Chunk {
id: chunk_id(SourceKind::Chat, scope, 1, "test-content"),
@@ -55,8 +77,10 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime
token_count: 6_000,
seq_in_source: 1,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[c1.clone(), c2.clone()]).unwrap();
stage_test_chunks(cfg, &[c1.clone(), c2.clone()]);
let leaf1 = LeafRef {
chunk_id: c1.id.clone(),
@@ -76,8 +100,12 @@ async fn seed_source_tree_with_sealed_l1(cfg: &Config, scope: &str, ts: DateTime
topics: vec![],
score: 0.5,
};
append_leaf(cfg, &tree, &leaf1, &summariser).await.unwrap();
append_leaf(cfg, &tree, &leaf2, &summariser).await.unwrap();
append_leaf(cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
append_leaf(cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
// 12k tokens > 10k budget → one L1 summary covering `ts`.
}
@@ -216,3 +244,157 @@ async fn seven_days_cascade_to_weekly_seal() {
assert_eq!(t.max_level, 1);
assert_eq!(t.status, TreeStatus::Active);
}
/// Seed a source tree whose sealed L1 summary carries the given entities
/// and topics. Entities are written into `mem_tree_entity_index` (where
/// seal-time hydration reads them); topics are stored on chunk metadata
/// tags. The seal then unions both into the L1 summary.
async fn seed_source_tree_with_labeled_l1(
cfg: &Config,
scope: &str,
ts: DateTime<Utc>,
entities: Vec<String>,
topics: Vec<String>,
) {
use crate::openhuman::memory::tree::score::extract::EntityKind;
use crate::openhuman::memory::tree::score::resolver::CanonicalEntity;
use crate::openhuman::memory::tree::score::store::index_entity;
let tree = get_or_create_source_tree(cfg, scope).unwrap();
let summariser = InertSummariser::new();
let mut chunks: Vec<Chunk> = Vec::new();
for seq in 0..2u32 {
chunks.push(Chunk {
id: chunk_id(SourceKind::Chat, scope, seq, "labeled-test"),
content: format!("labeled chunk {seq} in {scope}"),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: scope.into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: topics.clone(),
source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))),
},
token_count: 6_000,
seq_in_source: seq,
created_at: ts,
partial_message: false,
});
}
upsert_chunks(cfg, &chunks).unwrap();
stage_test_chunks(cfg, &chunks);
for chunk in &chunks {
for entity_id in &entities {
let kind = entity_id
.split_once(':')
.map_or(EntityKind::Misc, |(k, _)| {
EntityKind::parse(k).unwrap_or(EntityKind::Misc)
});
let surface = entity_id
.split_once(':')
.map_or(entity_id.as_str(), |(_, v)| v);
let e = CanonicalEntity {
canonical_id: entity_id.clone(),
kind,
surface: surface.to_string(),
span_start: 0,
span_end: surface.len() as u32,
score: 1.0,
};
index_entity(
cfg,
&e,
&chunk.id,
"leaf",
ts.timestamp_millis(),
Some(scope),
)
.unwrap();
}
}
// Two 6k-token leaves total 12k → exceeds L0 budget → seal fires on
// the second append, producing one L1 summary that unions all leaf
// labels (every leaf has the same set, so dedup yields the input set).
for chunk in &chunks {
let leaf = LeafRef {
chunk_id: chunk.id.clone(),
token_count: 6_000,
timestamp: ts,
content: chunk.content.clone(),
entities: entities.clone(),
topics: topics.clone(),
score: 0.5,
};
append_leaf(
cfg,
&tree,
&leaf,
&summariser,
&LabelStrategy::UnionFromChildren,
)
.await
.unwrap();
}
}
#[tokio::test]
async fn daily_digest_unions_labels_from_source_summaries() {
let (_tmp, cfg) = test_config();
let summariser = InertSummariser::new();
let day = NaiveDate::from_ymd_opt(2025, 5, 1).unwrap();
let ts = day.and_hms_opt(10, 0, 0).unwrap().and_utc();
// Source A's L1 carries (alice, phoenix-migration). Source B's L1
// carries (bob, phoenix-migration, qa). The daily L0 should union to
// (alice, bob, phoenix-migration) for entities and (phoenix-migration,
// qa) for topics — overlap dedup'd.
seed_source_tree_with_labeled_l1(
&cfg,
"slack:#a",
ts,
vec!["email:alice@example.com".into(), "topic:phoenix".into()],
vec!["phoenix-migration".into()],
)
.await;
seed_source_tree_with_labeled_l1(
&cfg,
"slack:#b",
ts,
vec!["person:bob".into(), "topic:phoenix".into()],
vec!["phoenix-migration".into(), "qa".into()],
)
.await;
let outcome = end_of_day_digest(&cfg, day, &summariser).await.unwrap();
let daily_id = match outcome {
DigestOutcome::Emitted { daily_id, .. } => daily_id,
other => panic!("expected Emitted, got {other:?}"),
};
let daily = store::get_summary(&cfg, &daily_id).unwrap().unwrap();
let entities: std::collections::BTreeSet<&str> =
daily.entities.iter().map(String::as_str).collect();
let topics: std::collections::BTreeSet<&str> =
daily.topics.iter().map(String::as_str).collect();
assert!(entities.contains("email:alice@example.com"));
assert!(entities.contains("person:bob"));
assert!(entities.contains("topic:phoenix"));
assert_eq!(
entities.len(),
3,
"expected 3 unique entities (deduped); got {entities:?}"
);
assert!(topics.contains("phoenix-migration"));
assert!(topics.contains("qa"));
assert_eq!(
topics.len(),
2,
"expected 2 unique topics (deduped); got {topics:?}"
);
}
+23 -1
View File
@@ -158,14 +158,31 @@ fn assemble_recap(covering: &[&SummaryNode], level: u32) -> RecapOutput {
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::global_tree::digest::{end_of_day_digest, DigestOutcome};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
use tempfile::TempDir;
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged = content_store::stage_chunks(&content_root, chunks)
.expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -216,6 +233,7 @@ mod tests {
token_count: 6_000,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
let c2 = Chunk {
id: chunk_id(SourceKind::Chat, scope, 1, "test-content"),
@@ -232,8 +250,10 @@ mod tests {
token_count: 6_000,
seq_in_source: 1,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[c1.clone(), c2.clone()]).unwrap();
stage_test_chunks(cfg, &[c1.clone(), c2.clone()]);
append_leaf(
cfg,
&tree,
@@ -247,6 +267,7 @@ mod tests {
score: 0.5,
},
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
@@ -263,6 +284,7 @@ mod tests {
score: 0.5,
},
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
+72 -4
View File
@@ -11,10 +11,15 @@
//! `mem_tree_summaries` on both sides (children and output), since even L0
//! is a sealed summary node rather than a raw chunk.
use std::collections::BTreeSet;
use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::{
atomic::stage_summary, SummaryComposeInput, SummaryTreeKind,
};
use crate::openhuman::memory::tree::global_tree::{
GLOBAL_TOKEN_BUDGET, MONTHLY_SEAL_THRESHOLD, WEEKLY_SEAL_THRESHOLD, YEARLY_SEAL_THRESHOLD,
};
@@ -186,6 +191,26 @@ async fn seal_one_level(
.await
.context("summariser failed during global seal")?;
// Global-tree summaries inherit their entity/topic labels via union
// from their already-labeled inputs (source-tree summaries carry
// labels from the source-tree seal extractor; global L1+ inputs
// carry labels from this same union path one level down). We
// deliberately do NOT run an extractor on the daily/weekly/monthly
// synthesis: the inputs already cover what the summary represents,
// and global is a sink — no second-pass labeling earns its keep.
let mut entities_set: BTreeSet<String> = BTreeSet::new();
let mut topics_set: BTreeSet<String> = BTreeSet::new();
for inp in &inputs {
for e in &inp.entities {
entities_set.insert(e.clone());
}
for t in &inp.topics {
topics_set.insert(t.clone());
}
}
let node_entities: Vec<String> = entities_set.into_iter().collect();
let node_topics: Vec<String> = topics_set.into_iter().collect();
// Phase 4 (#710): embed BEFORE opening the write tx so an embedder
// error aborts the cascade without half-committing the summary.
let embedder =
@@ -208,8 +233,8 @@ async fn seal_one_level(
child_ids: buf.item_ids.clone(),
content: output.content,
token_count: output.token_count,
entities: output.entities,
topics: output.topics,
entities: node_entities,
topics: node_topics,
time_range_start,
time_range_end,
score,
@@ -218,6 +243,49 @@ async fn seal_one_level(
embedding: Some(embedding),
};
// Phase MD-content: stage the global summary .md file before opening the
// write tx. date_for_global = time_range_start date (daily for L0, or
// the start of the range for higher levels).
let global_date = Some(time_range_start);
let compose_input_global = SummaryComposeInput {
summary_id: &node.id,
tree_kind: SummaryTreeKind::Global,
tree_id: &node.tree_id,
tree_scope: &tree.scope,
level: node.level,
child_ids: &node.child_ids,
child_count: node.child_ids.len(),
time_range_start: node.time_range_start,
time_range_end: node.time_range_end,
sealed_at: node.sealed_at,
body: &node.content,
};
// Stage the summary .md file — abort the seal on failure so the database
// never commits a row with content_path = NULL. The job-retry path will
// re-attempt the file write on next execution.
let content_root_global = config.memory_tree_content_root();
// Global tree scope is typically the literal "global" string.
// Use it as-is for the path (slugify passes through short ascii strings unchanged).
let global_scope_slug =
crate::openhuman::memory::tree::content_store::paths::slugify_source_id(&tree.scope);
let staged_global = stage_summary(
&content_root_global,
&compose_input_global,
&global_scope_slug,
global_date,
)
.with_context(|| {
format!(
"stage_summary failed for {}; global-tree seal aborted for retry",
node.id
)
})?;
log::debug!(
"[global_tree::seal] staged summary {} → {}",
node.id,
staged_global.content_path
);
// Single write transaction: insert the new summary, clear this level's
// buffer, append the new id to the parent buffer, and bump the tree's
// max_level/root_id if we just climbed. Re-read `max_level` inside the
@@ -238,7 +306,7 @@ async fn seal_one_level(
.map(|n| n.max(0) as u32)
.context("Failed to read current max_level for global tree")?;
store::insert_summary_tx(&tx, &node)?;
store::insert_summary_tx(&tx, &node, Some(&staged_global))?;
// Index any entities the summariser emitted. No-op under
// InertSummariser (entities stays empty by design — see
// summariser/inert.rs). Becomes active when the Ollama summariser
@@ -383,7 +451,7 @@ mod tests {
fn insert_daily(cfg: &Config, node: &SummaryNode) {
with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
store::insert_summary_tx(&tx, node)?;
store::insert_summary_tx(&tx, node, None)?;
tx.commit()?;
Ok(())
})
+127 -282
View File
@@ -1,14 +1,12 @@
//! Ingest orchestrator (Phase 1 + Phase 2):
//! Ingest orchestrator for the async memory-tree pipeline.
//!
//! canonicalise → chunk → score → admission gate → persist (chunks + scores + entity index)
//! The hot path now does:
//! `canonicalise -> chunk -> fast score -> persist chunks/score rows -> enqueue extract jobs`
//!
//! Phase 2 inserts scoring between chunker and persistence. Low-scoring
//! chunks are dropped (their rationale is still persisted to
//! `mem_tree_score` for diagnostics); surviving chunks get their entities
//! indexed so later phases can resolve "which chunks mention Alice?" in
//! O(lookup).
//! The slower work (full extraction, admission, tree buffering, sealing,
//! topic routing, daily digests) runs out of the SQLite-backed jobs queue.
use anyhow::{Context, Result};
use anyhow::Result;
use serde::{Deserialize, Serialize};
use crate::openhuman::config::Config;
@@ -19,25 +17,21 @@ use crate::openhuman::memory::tree::canonicalize::{
CanonicalisedSource,
};
use crate::openhuman::memory::tree::chunker::{chunk_markdown, ChunkerInput, ChunkerOptions};
use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked};
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::jobs::{self, ExtractChunkPayload, NewJob};
use crate::openhuman::memory::tree::score::{self, ScoreResult, ScoringConfig};
use crate::openhuman::memory::tree::source_tree::{
append_leaf, get_or_create_source_tree, LeafRef,
};
use crate::openhuman::memory::tree::store;
use crate::openhuman::memory::tree::topic_tree::route_leaf_to_topic_trees;
use crate::openhuman::memory::tree::types::Chunk;
/// Outcome of one ingest call — extended with per-chunk admission info.
/// Outcome of one ingest call.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct IngestResult {
pub source_id: String,
/// Number of chunks that passed the admission gate and were persisted.
/// Number of chunks persisted and queued for async extraction.
pub chunks_written: usize,
/// Number of chunks that failed the admission gate and were NOT persisted
/// (their score rationale IS persisted for diagnostics).
/// Number of chunks the cheap fast-score path would drop. Final admission
/// still happens later in the extract job.
pub chunks_dropped: usize,
/// IDs of all chunks that were persisted (in source order).
/// IDs of all chunks written and queued.
pub chunk_ids: Vec<String>,
}
@@ -52,7 +46,6 @@ impl IngestResult {
}
}
/// Ingest a batch of chat messages scoped to one channel/group.
pub async fn ingest_chat(
config: &Config,
source_id: &str,
@@ -60,11 +53,6 @@ pub async fn ingest_chat(
tags: Vec<String>,
batch: ChatBatch,
) -> Result<IngestResult> {
log::debug!(
"[memory_tree::ingest] chat source_id={} msg_count={}",
source_id,
batch.messages.len()
);
let canonical =
match chat::canonicalise(source_id, owner, &tags, batch).map_err(anyhow::Error::msg)? {
Some(c) => c,
@@ -73,7 +61,6 @@ pub async fn ingest_chat(
persist(config, source_id, canonical).await
}
/// Ingest a single email thread.
pub async fn ingest_email(
config: &Config,
source_id: &str,
@@ -81,11 +68,6 @@ pub async fn ingest_email(
tags: Vec<String>,
thread: EmailThread,
) -> Result<IngestResult> {
log::debug!(
"[memory_tree::ingest] email source_id={} msg_count={}",
source_id,
thread.messages.len()
);
let canonical =
match email::canonicalise(source_id, owner, &tags, thread).map_err(anyhow::Error::msg)? {
Some(c) => c,
@@ -94,7 +76,6 @@ pub async fn ingest_email(
persist(config, source_id, canonical).await
}
/// Ingest a single standalone document.
pub async fn ingest_document(
config: &Config,
source_id: &str,
@@ -102,13 +83,6 @@ pub async fn ingest_document(
tags: Vec<String>,
doc: DocumentInput,
) -> Result<IngestResult> {
let title_len = doc.title.chars().count();
log::debug!(
"[memory_tree::ingest] document source_id={} has_title={} title_len={}",
source_id,
!doc.title.trim().is_empty(),
title_len
);
let canonical =
match document::canonicalise(source_id, owner, &tags, doc).map_err(anyhow::Error::msg)? {
Some(c) => c,
@@ -122,7 +96,6 @@ async fn persist(
source_id: &str,
canonical: CanonicalisedSource,
) -> Result<IngestResult> {
// 1. Chunk
let input = ChunkerInput {
source_kind: canonical.metadata.source_kind,
source_id: source_id.to_string(),
@@ -134,15 +107,15 @@ async fn persist(
return Ok(IngestResult::empty(source_id));
}
// 2. Score (async; extractor chain driven by config.memory_tree —
// regex-only unless llm_extractor_endpoint + model are set, in
// which case LlmEntityExtractor runs as second-pass on borderline
// chunks with soft-fallback to regex on LLM failure).
let scoring_cfg = ScoringConfig::from_config(config);
let scores = score::score_chunks(&chunks, &scoring_cfg).await?;
// Phase MD-content: write chunk bodies to disk before the SQLite upsert.
// stage_chunks is sync I/O; run it here (still on the tokio thread) before
// spawn_blocking so errors surface before the DB transaction opens.
let content_root = config.memory_tree_content_root();
let staged = content_store::stage_chunks(&content_root, &chunks)
.map_err(|e| anyhow::anyhow!("[memory_tree::ingest] stage_chunks failed: {e}"))?;
// Fail fast on scorer length mismatch — silently truncating via zip would
// drop chunks (or their score rationale) without trace.
let scoring_cfg = ScoringConfig::from_config(config);
let scores = score::score_chunks_fast(&chunks, &scoring_cfg).await?;
if scores.len() != chunks.len() {
anyhow::bail!(
"[memory_tree::ingest] scorer length mismatch: chunks={} scores={}",
@@ -151,78 +124,79 @@ async fn persist(
);
}
// 3. Partition kept vs dropped
let mut kept_chunks: Vec<Chunk> = Vec::new();
let mut all_results: Vec<(ScoreResult, i64)> = Vec::new();
for (chunk, result) in chunks.iter().zip(scores.into_iter()) {
let ts_ms = chunk.metadata.timestamp.timestamp_millis();
if result.kept {
kept_chunks.push(chunk.clone());
}
all_results.push((result, ts_ms));
}
let all_results: Vec<(ScoreResult, i64)> = chunks
.iter()
.zip(scores.into_iter())
.map(|(chunk, result)| (result, chunk.metadata.timestamp.timestamp_millis()))
.collect();
let dropped = all_results.iter().filter(|(r, _)| !r.kept).count();
log::debug!(
"[memory_tree::ingest] scoring source_id={} kept={} dropped={}",
source_id,
kept_chunks.len(),
dropped
);
// 3.5. Phase 4 (#710) — embed every kept chunk BEFORE writing. A failed
// embed aborts the whole batch so a retry stays idempotent on
// `chunk_id` (upsert_chunks_tx ON CONFLICT REPLACE). Serial per
// chunk for now; parallelism is an explicit follow-up. We pack
// each vector into its SQLite BLOB representation up front so the
// write tx below doesn't have to juggle floats.
let embedder = build_embedder_from_config(config).context("build embedder during ingest")?;
log::debug!(
"[memory_tree::ingest] embedding source_id={} provider={} kept={}",
source_id,
embedder.name(),
kept_chunks.len()
);
let mut chunk_embedding_blobs: Vec<(String, Vec<u8>)> = Vec::with_capacity(kept_chunks.len());
for chunk in &kept_chunks {
let vector = embedder
.embed(&chunk.content)
.await
.with_context(|| format!("embed chunk_id={} during ingest", chunk.id))?;
let packed = pack_checked(&vector)
.with_context(|| format!("pack embedding for chunk_id={} during ingest", chunk.id))?;
chunk_embedding_blobs.push((chunk.id.clone(), packed));
}
// 4. Persist (blocking SQLite — isolate on a dedicated thread).
// Chunks + scores + embeddings all commit in one tx so retries are
// idempotent. The chunk upsert preserves any prior embedding on
// conflict (see store::upsert_chunks_tx), so we follow with a
// deliberate UPDATE that writes the fresh blob.
let config_owned = config.clone();
let kept_for_store = kept_chunks.clone();
let staged_for_store = staged.clone();
let results_for_store = all_results.clone();
let embeddings_for_store = chunk_embedding_blobs.clone();
let written = tokio::task::spawn_blocking(move || -> Result<usize> {
use std::collections::{HashMap, HashSet};
store::with_connection(&config_owned, |conn| {
let tx = conn.unchecked_transaction()?;
let n = store::upsert_chunks_tx(&tx, &kept_for_store)?;
for (chunk_id, blob) in &embeddings_for_store {
let changed = tx.execute(
"UPDATE mem_tree_chunks SET embedding = ?1 WHERE id = ?2",
rusqlite::params![blob, chunk_id],
)?;
if changed == 0 {
log::warn!(
"[memory_tree::ingest] embedding update affected 0 rows chunk_id={chunk_id} — \
upsert missed the row?"
);
// Read each chunk's CURRENT lifecycle BEFORE the upsert. This
// is the "did this chunk exist before this batch" snapshot,
// because `upsert_staged_chunks_tx` will either preserve the
// existing row's lifecycle (UPDATE doesn't touch the column) or
// insert a new row that picks up the column DEFAULT — so reading
// post-upsert can't distinguish "brand new" from
// "already-admitted-from-prior-ingest".
let mut prior: HashMap<String, Option<String>> = HashMap::new();
for s in &staged_for_store {
let status = store::get_chunk_lifecycle_status_tx(&tx, &s.chunk.id)?;
prior.insert(s.chunk.id.clone(), status);
}
let n = store::upsert_staged_chunks_tx(&tx, &staged_for_store)?;
// Re-ingest of identical content (same chunk_id) must NOT
// downgrade chunks that have already progressed through the
// async pipeline. Without this guard, a re-ingest would reset
// every chunk to 'pending_extraction' and enqueue a fresh
// `extract_chunk` job — sending already-buffered/sealed
// chunks back through extract → admit → append, ultimately
// duplicating them into a second summary in the same tree.
//
// Schedule a chunk for processing when its PRE-upsert state
// was either absent (genuinely new) or already
// `pending_extraction` (a prior ingest crashed before extract
// ran). Anything else — `admitted`, `buffered`, `sealed`,
// `dropped` — is past the point of accepting new work, so
// leave the lifecycle alone and skip the extract enqueue.
let mut to_schedule: HashSet<String> = HashSet::new();
for s in &staged_for_store {
let pre = prior.get(&s.chunk.id).cloned().flatten();
let needs_processing = matches!(
pre.as_deref(),
None | Some(store::CHUNK_STATUS_PENDING_EXTRACTION),
);
if needs_processing {
store::set_chunk_lifecycle_status_tx(
&tx,
&s.chunk.id,
store::CHUNK_STATUS_PENDING_EXTRACTION,
)?;
to_schedule.insert(s.chunk.id.clone());
}
}
for (result, ts_ms) in &results_for_store {
// Persist rationale for EVERY chunk (kept or dropped).
// Index entities only for kept chunks (handled inside persist_score_tx).
if !to_schedule.contains(&result.chunk_id) {
// Chunk has already progressed past pending_extraction
// on a prior ingest — skip score re-persist and don't
// enqueue a duplicate extract job.
continue;
}
score::persist_score_tx(&tx, result, *ts_ms, None)?;
let extract = NewJob::extract_chunk(&ExtractChunkPayload {
chunk_id: result.chunk_id.clone(),
})?;
let _ = jobs::enqueue_tx(&tx, &extract)?;
}
tx.commit()?;
Ok(n)
@@ -231,98 +205,26 @@ async fn persist(
.await
.map_err(|e| anyhow::anyhow!("persist join error: {e}"))??;
// 5. Source-tree append (Phase 3a #709). Each kept leaf pushes into
// the tree's L0 buffer and cascades upward when token_sum crosses
// the budget. Entities/topics from the scorer are threaded in so
// sealed summaries inherit the child signal set. Failures here
// log at warn level but don't fail the ingest — leaves are already
// persisted, and a later flush/retry can still rebuild the tree.
if let Err(e) = append_leaves_to_tree(config, source_id, &kept_chunks, &all_results).await {
log::warn!(
"[memory_tree::ingest] source_tree append failed source_id={} err={:#}",
source_id,
e
);
}
jobs::wake_workers();
Ok(IngestResult {
source_id: source_id.to_string(),
chunks_written: written,
chunks_dropped: dropped,
chunk_ids: kept_chunks.iter().map(|c| c.id.clone()).collect(),
chunk_ids: staged.iter().map(|s| s.chunk.id.clone()).collect(),
})
}
/// Push every kept chunk into its source tree. Scoped to Phase 3a — all
/// chunks from one ingest batch share the same `source_id`, so they share
/// one tree lookup. The summariser comes from `build_summariser(config)`:
/// Ollama-backed [`LlmSummariser`] when `memory_tree.llm_summariser_*` is
/// set, else deterministic [`InertSummariser`]. Both share the same trait,
/// so the seal cascade doesn't notice the difference.
async fn append_leaves_to_tree(
config: &Config,
source_id: &str,
kept_chunks: &[Chunk],
all_results: &[(ScoreResult, i64)],
) -> Result<()> {
if kept_chunks.is_empty() {
return Ok(());
}
let tree = get_or_create_source_tree(config, source_id)?;
let summariser = crate::openhuman::memory::tree::source_tree::build_summariser(config);
// Build a chunk_id → (score, entities, topics) map for quick lookup.
use std::collections::HashMap;
let mut score_by_id: HashMap<String, &ScoreResult> = HashMap::new();
for (r, _) in all_results {
score_by_id.insert(r.chunk_id.clone(), r);
}
for chunk in kept_chunks {
let (score_value, entities, topics) = match score_by_id.get(&chunk.id) {
Some(r) => (
r.total,
r.canonical_entities
.iter()
.map(|e| e.canonical_id.clone())
.collect::<Vec<_>>(),
chunk.metadata.tags.clone(),
),
None => (0.0, Vec::new(), chunk.metadata.tags.clone()),
};
let leaf = LeafRef {
chunk_id: chunk.id.clone(),
token_count: chunk.token_count,
timestamp: chunk.metadata.timestamp,
content: chunk.content.clone(),
entities: entities.clone(),
topics,
score: score_value,
};
append_leaf(config, &tree, &leaf, summariser.as_ref()).await?;
// Phase 3c (#709): route the leaf to every matching topic tree
// and tick the curator for each entity. Non-fatal on error —
// the source-tree append has already succeeded above.
if let Err(e) =
route_leaf_to_topic_trees(config, &leaf, &entities, summariser.as_ref()).await
{
log::warn!(
"[memory_tree::ingest] topic_tree routing failed chunk_id={} err={:#}",
chunk.id,
e
);
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::canonicalize::chat::ChatMessage;
use crate::openhuman::memory::tree::jobs::drain_until_idle;
use crate::openhuman::memory::tree::score::store::{count_scores, lookup_entity};
use crate::openhuman::memory::tree::store::{count_chunks, list_chunks, ListChunksQuery};
use crate::openhuman::memory::tree::store::{
count_chunks, count_chunks_by_lifecycle_status, get_chunk_embedding, list_chunks,
ListChunksQuery, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED,
};
use crate::openhuman::memory::tree::types::SourceKind;
use chrono::{TimeZone, Utc};
use tempfile::TempDir;
@@ -331,16 +233,12 @@ mod tests {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
cfg.workspace_dir = tmp.path().to_path_buf();
// Phase 4 (#710): disable Ollama-backed embeddings in tests.
// Ingest/seal call `build_embedder_from_config`; falling back to
// inert (zero vectors) keeps tests deterministic and network-free.
cfg.memory_tree.embedding_endpoint = None;
cfg.memory_tree.embedding_model = None;
cfg.memory_tree.embedding_strict = false;
(tmp, cfg)
}
/// Build a substantive batch that reliably passes the admission gate.
fn substantive_batch() -> ChatBatch {
ChatBatch {
platform: "slack".into(),
@@ -349,17 +247,14 @@ mod tests {
ChatMessage {
author: "alice".into(),
timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
text: "We are planning to ship the Phoenix migration on Friday \
after reviewing the runbook and staging results. Please \
confirm availability by replying here. alice@example.com"
text: "We are planning to ship the Phoenix migration on Friday after reviewing the runbook and staging results. alice@example.com"
.into(),
source_ref: Some("slack://m1".into()),
},
ChatMessage {
author: "bob".into(),
timestamp: Utc.timestamp_millis_opt(1_700_000_010_000).unwrap(),
text: "Confirmed — I'll handle the coordination and cut a release \
candidate tonight. #launch-q2 will be tracked in Notion."
text: "Confirmed, I will handle the coordination and launch tracking tonight."
.into(),
source_ref: None,
},
@@ -368,25 +263,40 @@ mod tests {
}
#[tokio::test]
async fn ingest_chat_writes_substantive_chunks() {
async fn ingest_chat_writes_and_queue_drains_to_admitted_chunk() {
let (_tmp, cfg) = test_config();
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch())
.await
.unwrap();
// Greedy packing: both small messages fit under 10k token budget
// and are packed into a single chunk.
assert_eq!(out.chunks_written, 1);
assert_eq!(out.chunks_dropped, 0);
assert_eq!(count_chunks(&cfg).unwrap(), 1);
// Score row persisted for the kept chunk
assert_eq!(count_scores(&cfg).unwrap(), 1);
// Entity index populated from regex extraction (alice@example.com + hashtag)
let alice_hits = lookup_entity(&cfg, "email:alice@example.com", None).unwrap();
assert_eq!(alice_hits.len(), 1);
drain_until_idle(&cfg).await.unwrap();
// Final lifecycle is `buffered`: extract → admitted → append_buffer → buffered.
// The single packed chunk does not cross TOKEN_BUDGET so no seal fires.
assert_eq!(
count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_BUFFERED).unwrap(),
1
);
assert!(count_scores(&cfg).unwrap() >= 1);
assert_eq!(
lookup_entity(&cfg, "email:alice@example.com", None)
.unwrap()
.len(),
1
);
let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap();
assert_eq!(rows[0].metadata.source_kind, SourceKind::Chat);
assert!(get_chunk_embedding(&cfg, &out.chunk_ids[0])
.unwrap()
.is_some());
}
#[tokio::test]
async fn low_signal_chunks_are_dropped_but_score_persists() {
async fn low_signal_chunks_end_up_dropped_after_queue_processing() {
let (_tmp, cfg) = test_config();
let batch = ChatBatch {
platform: "slack".into(),
@@ -394,18 +304,22 @@ mod tests {
messages: vec![ChatMessage {
author: "alice".into(),
timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
text: "+1".into(), // extremely low-signal
text: "+1".into(),
source_ref: None,
}],
};
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], batch)
.await
.unwrap();
assert_eq!(out.chunks_written, 0);
assert_eq!(out.chunks_dropped, 1);
// Chunk NOT in chunks table
assert_eq!(count_chunks(&cfg).unwrap(), 0);
// Score row IS persisted for diagnostics
assert_eq!(out.chunks_written, 1);
assert_eq!(count_chunks(&cfg).unwrap(), 1);
drain_until_idle(&cfg).await.unwrap();
assert_eq!(
count_chunks_by_lifecycle_status(&cfg, CHUNK_STATUS_DROPPED).unwrap(),
1
);
assert_eq!(count_scores(&cfg).unwrap(), 1);
}
@@ -421,7 +335,6 @@ mod tests {
.await
.unwrap();
assert_eq!(out.chunks_written, 0);
assert_eq!(out.chunks_dropped, 0);
assert_eq!(count_chunks(&cfg).unwrap(), 0);
assert_eq!(count_scores(&cfg).unwrap(), 0);
}
@@ -432,9 +345,7 @@ mod tests {
let doc = DocumentInput {
provider: "notion".into(),
title: "Launch plan".into(),
body: "We are planning to ship Phoenix on Friday after review. \
Coordination is via email and the launch thread tracks \
the relevant decisions. alice@example.com owns this."
body: "We are planning to ship Phoenix on Friday after review. alice@example.com owns this."
.into(),
modified_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(),
source_ref: Some("notion://page/abc".into()),
@@ -445,74 +356,8 @@ mod tests {
ingest_document(&cfg, "notion:abc", "alice", vec![], doc)
.await
.unwrap();
drain_until_idle(&cfg).await.unwrap();
assert_eq!(count_chunks(&cfg).unwrap(), 1);
assert_eq!(count_scores(&cfg).unwrap(), 1);
}
#[tokio::test]
async fn chunks_preserve_source_ref_when_kept() {
let (_tmp, cfg) = test_config();
let doc = DocumentInput {
provider: "notion".into(),
title: "t".into(),
body: "Phoenix launch plan with enough substance to pass the admission \
gate: we are reviewing the migration runbook alice@example.com \
on Friday evening."
.into(),
modified_at: Utc::now(),
source_ref: Some("notion://x".into()),
};
ingest_document(&cfg, "notion:x", "alice", vec![], doc)
.await
.unwrap();
let rows = list_chunks(&cfg, &ListChunksQuery::default()).unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(
rows[0].metadata.source_ref.as_ref().unwrap().value,
"notion://x"
);
}
// ── Phase 4 (#710) ──────────────────────────────────────────────
/// Ingesting with an inert embedder must still write embeddings for
/// every kept chunk. Use a trimmed test_config so we go through the
/// full embed → persist pipeline.
#[tokio::test]
async fn ingest_writes_chunk_embeddings() {
use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM;
use crate::openhuman::memory::tree::store::get_chunk_embedding;
let (_tmp, cfg) = test_config();
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch())
.await
.unwrap();
assert!(out.chunks_written >= 1);
for id in &out.chunk_ids {
let v = get_chunk_embedding(&cfg, id)
.unwrap()
.unwrap_or_else(|| panic!("missing embedding for {id}"));
assert_eq!(v.len(), EMBEDDING_DIM);
}
}
/// When the embedder errors (here: Ollama required but unavailable),
/// ingest must fail and persist nothing. Point embedding at a dead
/// port so the HTTP call refuses, and flip `strict=true` so the
/// factory returns Ollama (not the inert fallback).
#[tokio::test]
async fn ingest_fails_when_embedder_fails_and_persists_nothing() {
let (_tmp, mut cfg) = test_config();
cfg.memory_tree.embedding_endpoint = Some("http://127.0.0.1:1".into());
cfg.memory_tree.embedding_model = Some("nomic-embed-text".into());
cfg.memory_tree.embedding_timeout_ms = Some(500);
cfg.memory_tree.embedding_strict = true;
let result = ingest_chat(&cfg, "slack:#eng", "alice", vec![], substantive_batch()).await;
assert!(result.is_err(), "expected ingest to bail when embed fails");
// No chunks, no scores persisted — retry stays clean.
assert_eq!(count_chunks(&cfg).unwrap(), 0);
assert_eq!(count_scores(&cfg).unwrap(), 0);
}
}
@@ -0,0 +1,861 @@
use anyhow::{Context, Result};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::{
self as content_store, read as content_read, tags as content_tags,
};
use crate::openhuman::memory::tree::global_tree::digest::{self, DigestOutcome};
use crate::openhuman::memory::tree::jobs::store;
use crate::openhuman::memory::tree::jobs::types::{
AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload,
Job, JobKind, NewJob, NodeRef, SealPayload, TopicRoutePayload,
};
use crate::openhuman::memory::tree::score;
use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, pack_checked};
use crate::openhuman::memory::tree::score::extract::build_summary_extractor;
use crate::openhuman::memory::tree::score::store as score_store;
use crate::openhuman::memory::tree::source_tree::{
build_summariser, get_or_create_source_tree, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::store as chunk_store;
use crate::openhuman::memory::tree::topic_tree::curator;
pub async fn handle_job(config: &Config, job: &Job) -> Result<()> {
match job.kind {
JobKind::ExtractChunk => handle_extract(config, job).await,
JobKind::AppendBuffer => handle_append_buffer(config, job).await,
JobKind::Seal => handle_seal(config, job).await,
JobKind::TopicRoute => handle_topic_route(config, job).await,
JobKind::DigestDaily => handle_digest_daily(config, job).await,
JobKind::FlushStale => handle_flush_stale(config, job).await,
}
}
async fn handle_extract(config: &Config, job: &Job) -> Result<()> {
let payload: ExtractChunkPayload =
serde_json::from_str(&job.payload_json).context("parse ExtractChunk payload")?;
let Some(chunk) = chunk_store::get_chunk(config, &payload.chunk_id)? else {
log::warn!(
"[memory_tree::jobs] extract chunk missing chunk_id={}",
payload.chunk_id
);
return Ok(());
};
// Read the full body from disk (the `content` column in SQLite holds a
// ≤500-char preview after the MD-on-disk migration). Both the scorer and
// the embedder need the complete text so extraction and semantic indexing
// operate over the full chunk body, not a truncated preview.
let body = content_read::read_chunk_body(config, &chunk.id)
.with_context(|| format!("read full body for extract chunk_id={}", chunk.id))?;
// Score a clone of the chunk with the full body swapped in.
let chunk_with_body = {
let mut c = chunk.clone();
c.content = body.clone();
c
};
let scoring_cfg = score::ScoringConfig::from_config(config);
let result = score::score_chunk(&chunk_with_body, &scoring_cfg).await?;
let packed_embedding = if result.kept {
let embedder =
build_embedder_from_config(config).context("build embedder in extract handler")?;
// Reuse the body already read — avoid a second disk read.
let vector = embedder
.embed(&body)
.await
.with_context(|| format!("embed chunk_id={} in extract handler", chunk.id))?;
Some(
pack_checked(&vector)
.with_context(|| format!("pack embedding for chunk_id={}", chunk.id))?,
)
} else {
None
};
// Build follow-up job payloads before opening the tx — construction is
// cheap and doesn't require a database connection. The two jobs are
// enqueued inside the SAME transaction that commits the lifecycle update,
// so a crash anywhere rolls everything back together and prevents the
// "lifecycle committed but job lost" crash window.
let source_job = if result.kept {
Some(NewJob::append_buffer(&AppendBufferPayload {
node: NodeRef::Leaf {
chunk_id: chunk.id.clone(),
},
target: AppendTarget::Source {
source_id: chunk.metadata.source_id.clone(),
},
})?)
} else {
None
};
let route_job = if result.kept {
Some(NewJob::topic_route(&TopicRoutePayload {
node: NodeRef::Leaf {
chunk_id: chunk.id.clone(),
},
})?)
} else {
None
};
let (did_enqueue_source, did_enqueue_route) = chunk_store::with_connection(config, |conn| {
let tx = conn.unchecked_transaction()?;
score::persist_score_tx(
&tx,
&result,
chunk.metadata.timestamp.timestamp_millis(),
None,
)?;
if result.kept {
tx.execute(
"UPDATE mem_tree_chunks
SET embedding = ?1,
lifecycle_status = ?2
WHERE id = ?3",
rusqlite::params![
packed_embedding,
chunk_store::CHUNK_STATUS_ADMITTED,
chunk.id,
],
)?;
} else {
tx.execute(
"UPDATE mem_tree_chunks
SET lifecycle_status = ?1
WHERE id = ?2",
rusqlite::params![chunk_store::CHUNK_STATUS_DROPPED, chunk.id],
)?;
}
// Enqueue follow-up jobs inside the SAME transaction so they are
// atomically visible with the lifecycle update.
let mut eq_src = false;
let mut eq_route = false;
if let Some(ref j) = source_job {
eq_src = store::enqueue_tx(&tx, j)?.is_some();
}
if let Some(ref j) = route_job {
eq_route = store::enqueue_tx(&tx, j)?.is_some();
}
tx.commit()?;
Ok((eq_src, eq_route))
})?;
// Phase MD-content: rewrite the `tags:` block in the on-disk chunk file
// with Obsidian-style hierarchical tags derived from the extracted entities.
// This runs after the tx commits so the entity index is visible to readers.
// It is a filesystem op and therefore lives outside the SQL tx — best-effort.
if result.kept {
if let Some(content_path) = chunk_store::get_chunk_content_path(config, &chunk.id)? {
let content_root = config.memory_tree_content_root();
let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?;
let obsidian_tags: Vec<String> = entity_ids
.iter()
.filter_map(|eid| {
// entity_id format: "kind:surface"
let (kind, surface) = eid.split_once(':')?;
Some(content_tags::entity_tag(kind, surface))
})
.collect();
// Build the absolute path from the stored relative path.
let abs_path = {
let mut p = content_root.clone();
for component in content_path.split('/') {
p.push(component);
}
p
};
if let Err(e) = content_tags::update_chunk_tags(&abs_path, &obsidian_tags) {
log::warn!(
"[memory_tree::jobs] failed to update tags in chunk file chunk_id={} path_hash={}: {e}",
chunk.id,
crate::openhuman::memory::tree::util::redact::redact(&content_path),
);
// Non-fatal: tag rewrite failure does not block the pipeline.
} else {
log::debug!(
"[memory_tree::jobs] updated {} obsidian tags in chunk file chunk_id={}",
obsidian_tags.len(),
chunk.id,
);
}
}
}
// Signal workers after the tx commits (no atomicity requirement on signaling).
if did_enqueue_source {
super::worker::wake_workers();
}
if did_enqueue_route {
super::worker::wake_workers();
}
Ok(())
}
async fn handle_append_buffer(config: &Config, job: &Job) -> Result<()> {
use crate::openhuman::memory::tree::source_tree::bucket_seal::should_seal;
use crate::openhuman::memory::tree::source_tree::store as src_store;
let payload: AppendBufferPayload =
serde_json::from_str(&job.payload_json).context("parse AppendBuffer payload")?;
// Hydrate the leaf-shaped record from either a chunk row or a summary
// row. The downstream buffer-push doesn't care which kind produced
// the LeafRef.
let (leaf, chunk_id_for_lifecycle): (LeafRef, Option<String>) = match &payload.node {
NodeRef::Leaf { chunk_id } => {
let Some(chunk) = chunk_store::get_chunk(config, chunk_id)? else {
log::warn!("[memory_tree::jobs] append_buffer chunk missing chunk_id={chunk_id}");
return Ok(());
};
let score_row = score_store::get_score(config, &chunk.id)?
.ok_or_else(|| anyhow::anyhow!("missing score row for chunk {}", chunk.id))?;
let entity_ids = score_store::list_entity_ids_for_node(config, &chunk.id)?;
// Read the full body from disk — the `content` column in SQLite
// is a ≤500-char preview after the MD-on-disk migration. The
// summariser receives this LeafRef and must see the complete text.
let body = content_read::read_chunk_body(config, chunk_id)
.with_context(|| format!("read chunk body in append_buffer chunk_id={chunk_id}"))?;
let leaf = LeafRef {
chunk_id: chunk.id.clone(),
token_count: chunk.token_count,
timestamp: chunk.metadata.timestamp,
content: body,
entities: entity_ids,
topics: chunk.metadata.tags.clone(),
score: score_row.total,
};
(leaf, Some(chunk.id))
}
NodeRef::Summary { summary_id } => {
let Some(summary) = src_store::get_summary(config, summary_id)? else {
log::warn!(
"[memory_tree::jobs] append_buffer summary missing summary_id={summary_id}"
);
return Ok(());
};
// Read the full body from disk — `summary.content` is a ≤500-char
// preview after the MD-on-disk migration. The summariser receives
// this LeafRef when sealing higher-level nodes and must see the
// complete summary text.
let body = content_read::read_summary_body(config, summary_id).with_context(|| {
format!("read summary body in append_buffer summary_id={summary_id}")
})?;
// Build a LeafRef from the summary's already-populated fields.
// `chunk_id` carries the source-node id (any string); buffer
// accounting uses it as the item id only.
let leaf = LeafRef {
chunk_id: summary.id.clone(),
token_count: summary.token_count,
timestamp: summary.time_range_start,
content: body,
entities: summary.entities.clone(),
topics: summary.topics.clone(),
score: summary.score,
};
(leaf, None) // summaries have no chunk lifecycle to update
}
};
// Resolve target tree (no tx open yet — this can create a row).
let tree = match &payload.target {
AppendTarget::Source { source_id } => Some(get_or_create_source_tree(config, source_id)?),
AppendTarget::Topic { tree_id } => src_store::get_tree(config, tree_id)?,
};
let Some(tree) = tree else {
// Target topic tree doesn't exist (e.g. archived between
// topic_route and this append). Drop on the floor — the
// topic_route was advisory and the source-tree path already
// ran for this leaf.
return Ok(());
};
let is_source_target = matches!(payload.target, AppendTarget::Source { .. });
let leaf_for_tx = leaf.clone();
let tree_for_tx = tree.clone();
let lifecycle_chunk_id = chunk_id_for_lifecycle.clone();
// ATOMIC: buffer push + seal enqueue (if gate met) + lifecycle update
// happen in a single SQLite transaction. Eliminates the crash window
// where the buffer commits but the seal job is lost — which can
// duplicate the leaf into two summaries on retry-after-seal-cleared.
let did_enqueue_seal = chunk_store::with_connection(config, move |conn| {
let tx = conn.unchecked_transaction()?;
// 1. Push leaf into L0 buffer (idempotent on (tree, level, item_id)).
let mut buf = src_store::get_buffer_conn(&tx, &tree_for_tx.id, 0)?;
if !buf.item_ids.iter().any(|x| x == &leaf_for_tx.chunk_id) {
buf.item_ids.push(leaf_for_tx.chunk_id.clone());
buf.token_sum = buf.token_sum.saturating_add(leaf_for_tx.token_count as i64);
buf.oldest_at = match buf.oldest_at {
Some(existing) => Some(existing.min(leaf_for_tx.timestamp)),
None => Some(leaf_for_tx.timestamp),
};
src_store::upsert_buffer_tx(&tx, &buf)?;
}
// 2. If the gate is met, enqueue a seal job atomically.
let did_enqueue = if should_seal(&buf) {
let seal = SealPayload {
tree_id: tree_for_tx.id.clone(),
level: 0,
force_now_ms: None,
};
store::enqueue_tx(&tx, &NewJob::seal(&seal)?)?.is_some()
} else {
false
};
// 3. Lifecycle transition (Source target with a leaf chunk).
// Last step in the tx — its presence is the "this handler
// finished" marker. Same tx as the push + seal-enqueue, so a
// crash anywhere rolls everything back together.
if is_source_target {
if let Some(chunk_id) = lifecycle_chunk_id.as_deref() {
chunk_store::set_chunk_lifecycle_status_tx(
&tx,
chunk_id,
chunk_store::CHUNK_STATUS_BUFFERED,
)?;
}
}
tx.commit()?;
Ok(did_enqueue)
})?;
if did_enqueue_seal {
super::worker::wake_workers();
}
Ok(())
}
async fn handle_seal(config: &Config, job: &Job) -> Result<()> {
use crate::openhuman::memory::tree::source_tree::bucket_seal::{seal_one_level, should_seal};
use crate::openhuman::memory::tree::source_tree::store as src_store;
use crate::openhuman::memory::tree::source_tree::types::TreeKind;
let payload: SealPayload =
serde_json::from_str(&job.payload_json).context("parse Seal payload")?;
let Some(tree) = src_store::get_tree(config, &payload.tree_id)? else {
log::warn!(
"[memory_tree::jobs] seal tree missing tree_id={}",
payload.tree_id
);
return Ok(());
};
// Seal exactly one level. Parents only get sealed via a follow-up job
// so each level is its own crash-recovery checkpoint and each LLM
// summariser call competes for a fresh slot from the global semaphore.
let buf = src_store::get_buffer(config, &tree.id, payload.level)?;
let forced = payload.force_now_ms.is_some();
if buf.is_empty() {
log::debug!(
"[memory_tree::jobs] seal skipped — empty buffer tree_id={} level={}",
tree.id,
payload.level
);
return Ok(());
}
if !forced && !should_seal(&buf) {
// Another job sealed this level out from under us (or the buffer
// hasn't crossed the gate yet); idempotent no-op.
log::debug!(
"[memory_tree::jobs] seal gate not met tree_id={} level={} token_sum={}",
tree.id,
payload.level,
buf.token_sum
);
return Ok(());
}
// Pick the labeling strategy for this tree kind. Source trees mint
// emergent themes via the seal-time extractor; topic trees stay empty
// by design (scope already pins the canonical id). Global trees never
// reach here — `digest_daily` handles them — but Empty is a safe
// defensive default.
let strategy = match tree.kind {
TreeKind::Source => LabelStrategy::ExtractFromContent(build_summary_extractor(config)),
TreeKind::Topic => LabelStrategy::Empty,
TreeKind::Global => LabelStrategy::Empty,
};
let summariser = build_summariser(config);
// `seal_one_level` with `enqueue_follow_ups: true` atomically inserts
// the parent-cascade seal (if the parent buffer now meets its gate)
// and the summary-side `topic_route` (for source trees) inside the
// same SQLite transaction that commits the seal. This eliminates the
// crash window where the seal succeeds but the follow-up enqueues
// are silently lost.
let summary_id =
seal_one_level(config, &tree, &buf, summariser.as_ref(), &strategy, true).await?;
// Phase MD-content: rewrite the `tags:` block in the sealed summary's
// on-disk .md file. Entity index rows were committed inside
// `seal_one_level` (via `index_summary_entity_ids_tx`), so they are
// visible here. Best-effort: failure does not abort the seal.
if let Err(e) = content_store::update_summary_tags(config, &summary_id) {
log::warn!(
"[memory_tree::jobs] update_summary_tags failed for summary_id={summary_id}: {e:#}"
);
}
super::worker::wake_workers();
Ok(())
}
async fn handle_topic_route(config: &Config, job: &Job) -> Result<()> {
let payload: TopicRoutePayload =
serde_json::from_str(&job.payload_json).context("parse TopicRoute payload")?;
// Resolve the source node id and verify it exists. `mem_tree_entity_index`
// already indexes both chunks and summaries via `node_kind`, so the
// canonical-id loop below is identical for either case.
let node_id: String = match &payload.node {
NodeRef::Leaf { chunk_id } => {
if chunk_store::get_chunk(config, chunk_id)?.is_none() {
log::warn!("[memory_tree::jobs] topic_route chunk missing chunk_id={chunk_id}");
return Ok(());
}
chunk_id.clone()
}
NodeRef::Summary { summary_id } => {
if crate::openhuman::memory::tree::source_tree::store::get_summary(config, summary_id)?
.is_none()
{
log::warn!(
"[memory_tree::jobs] topic_route summary missing summary_id={summary_id}"
);
return Ok(());
}
summary_id.clone()
}
};
let entity_ids = score_store::list_entity_ids_for_node(config, &node_id)?;
if entity_ids.is_empty() {
log::debug!("[memory_tree::jobs] topic_route no entities for node_id={node_id} — skipping");
return Ok(());
}
let summariser = build_summariser(config);
for entity_id in entity_ids {
let _ = curator::maybe_spawn_topic_tree(config, &entity_id, summariser.as_ref()).await?;
if let Some(tree) = crate::openhuman::memory::tree::source_tree::store::get_tree_by_scope(
config,
crate::openhuman::memory::tree::source_tree::types::TreeKind::Topic,
&entity_id,
)? {
let job = NewJob::append_buffer(&AppendBufferPayload {
node: payload.node.clone(),
target: AppendTarget::Topic {
tree_id: tree.id.clone(),
},
})?;
if store::enqueue(config, &job)?.is_some() {
super::worker::wake_workers();
}
}
}
Ok(())
}
async fn handle_digest_daily(config: &Config, job: &Job) -> Result<()> {
let payload: DigestDailyPayload =
serde_json::from_str(&job.payload_json).context("parse DigestDaily payload")?;
let day = chrono::NaiveDate::parse_from_str(&payload.date_iso, "%Y-%m-%d")
.with_context(|| format!("invalid digest date {}", payload.date_iso))?;
let summariser = build_summariser(config);
match digest::end_of_day_digest(config, day, summariser.as_ref()).await? {
DigestOutcome::Emitted { daily_id, .. } => {
log::info!("[memory_tree::jobs] emitted digest daily_id={daily_id}");
}
DigestOutcome::EmptyDay => {}
DigestOutcome::Skipped { existing_id } => {
log::debug!("[memory_tree::jobs] digest skipped existing_id={existing_id}");
}
}
Ok(())
}
async fn handle_flush_stale(config: &Config, job: &Job) -> Result<()> {
let payload: FlushStalePayload =
serde_json::from_str(&job.payload_json).context("parse FlushStale payload")?;
let age_secs = payload
.max_age_secs
.unwrap_or(crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS);
let cutoff = chrono::Utc::now() - chrono::Duration::seconds(age_secs);
let buffers =
crate::openhuman::memory::tree::source_tree::store::list_stale_buffers(config, cutoff)?;
for buf in buffers {
let seal = SealPayload {
tree_id: buf.tree_id.clone(),
level: buf.level,
force_now_ms: Some(chrono::Utc::now().timestamp_millis()),
};
if store::enqueue(config, &NewJob::seal(&seal)?)?.is_some() {
super::worker::wake_workers();
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::jobs::store::{count_by_status, count_total};
use crate::openhuman::memory::tree::jobs::types::JobStatus;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf_deferred, LeafRef};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::store as src_store;
use crate::openhuman::memory::tree::store::with_connection;
use chrono::TimeZone;
use rusqlite::params;
use tempfile::TempDir;
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
cfg.workspace_dir = tmp.path().to_path_buf();
cfg.memory_tree.embedding_endpoint = None;
cfg.memory_tree.embedding_model = None;
cfg.memory_tree.embedding_strict = false;
(tmp, cfg)
}
/// Build a minimal `Job` row for direct handler invocation. Mirrors
/// what `claim_next` would produce for a freshly-claimed row.
fn mk_running_job(kind: JobKind, payload_json: String) -> Job {
let now_ms = chrono::Utc::now().timestamp_millis();
Job {
id: "test-job-id".into(),
kind,
payload_json,
dedupe_key: None,
status: JobStatus::Running,
attempts: 1,
max_attempts: 5,
available_at_ms: now_ms,
locked_until_ms: Some(now_ms + 60_000),
last_error: None,
created_at_ms: now_ms,
started_at_ms: Some(now_ms),
completed_at_ms: None,
}
}
/// Count rows in `mem_tree_jobs` matching a specific kind.
fn count_jobs_of_kind(cfg: &Config, kind: &str) -> u64 {
with_connection(cfg, |conn| {
let n: i64 = conn.query_row(
"SELECT COUNT(*) FROM mem_tree_jobs WHERE kind = ?1",
params![kind],
|r| r.get(0),
)?;
Ok(n.max(0) as u64)
})
.unwrap()
}
/// Seed a source tree and push enough labeled leaves into its L0 buffer
/// to cross `TOKEN_BUDGET`, returning the tree. The caller can then
/// fire `handle_seal` and inspect the result.
async fn seed_source_tree_ready_to_seal(
cfg: &Config,
) -> crate::openhuman::memory::tree::source_tree::types::Tree {
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
};
let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap();
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let chunk = Chunk {
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "handler-seed"),
content: "alice@example.com leading the rollout".into(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: Some(SourceRef::new("slack://x")),
},
// Bust budget so the L0 buffer is "ready" for seal.
token_count: 10_000,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[chunk.clone()]).unwrap();
// Stage to disk so `hydrate_leaf_inputs` can read the full body via
// `read_chunk_body` when `handle_seal` fires and calls `seal_one_level`.
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
let leaf = LeafRef {
chunk_id: chunk.id,
token_count: 10_000,
timestamp: ts,
content: chunk.content,
entities: vec![],
topics: vec![],
score: 0.5,
};
// append_leaf_deferred only buffers; doesn't seal. handle_seal will.
let _ = append_leaf_deferred(cfg, &tree, &leaf).unwrap();
tree
}
#[tokio::test]
async fn source_tree_seal_handler_enqueues_summary_topic_route() {
let (_tmp, cfg) = test_config();
let tree = seed_source_tree_ready_to_seal(&cfg).await;
let payload = SealPayload {
tree_id: tree.id.clone(),
level: 0,
force_now_ms: None,
};
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
// Pre-condition: queue has no topic_route jobs.
assert_eq!(count_jobs_of_kind(&cfg, "topic_route"), 0);
super::handle_seal(&cfg, &job).await.unwrap();
// Post-condition: source-tree seal must enqueue exactly one
// topic_route job carrying NodeRef::Summary { summary_id: <new> }.
assert_eq!(
count_jobs_of_kind(&cfg, "topic_route"),
1,
"source-tree seal must enqueue summary-side topic_route"
);
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1);
// Inspect the enqueued payload to confirm it's a Summary variant.
let payload_json: String = with_connection(&cfg, |conn| {
let s: String = conn
.query_row(
"SELECT payload_json FROM mem_tree_jobs WHERE kind = 'topic_route'",
[],
|r| r.get(0),
)
.unwrap();
Ok(s)
})
.unwrap();
let p: TopicRoutePayload = serde_json::from_str(&payload_json).unwrap();
match p.node {
NodeRef::Summary { summary_id } => {
assert!(summary_id.starts_with("summary:L1:"));
}
other => panic!("expected NodeRef::Summary, got {other:?}"),
}
}
#[tokio::test]
async fn topic_tree_seal_handler_does_not_enqueue_topic_route() {
let (_tmp, cfg) = test_config();
// Spawn a topic tree directly via the registry (skipping curator's
// hotness gate — we just need a TreeKind::Topic with leaves).
let topic_tree =
crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree(
&cfg,
"topic:phoenix-migration",
)
.unwrap();
// Push a single 10k-token leaf so L0 is gate-ready.
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
};
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let chunk = Chunk {
id: chunk_id(SourceKind::Chat, "slack:#eng", 0, "topic-seed"),
content: "topic content".into(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: Some(SourceRef::new("slack://x")),
},
token_count: 10_000,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
// Stage to disk so `hydrate_leaf_inputs` can read the full body
// when `handle_seal` fires.
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
let leaf = LeafRef {
chunk_id: chunk.id,
token_count: 10_000,
timestamp: ts,
content: chunk.content,
entities: vec![],
topics: vec![],
score: 0.5,
};
append_leaf_deferred(&cfg, &topic_tree, &leaf).unwrap();
let payload = SealPayload {
tree_id: topic_tree.id.clone(),
level: 0,
force_now_ms: None,
};
let job = mk_running_job(JobKind::Seal, serde_json::to_string(&payload).unwrap());
super::handle_seal(&cfg, &job).await.unwrap();
// Topic-tree seals are sinks: must not enqueue any topic_route.
assert_eq!(
count_jobs_of_kind(&cfg, "topic_route"),
0,
"topic-tree seal must NOT enqueue topic_route (trees are sinks)"
);
// The seal itself should still have produced a summary node.
assert_eq!(src_store::count_summaries(&cfg, &topic_tree.id).unwrap(), 1);
}
#[tokio::test]
async fn handle_append_buffer_with_summary_payload_pushes_into_topic_tree() {
let (_tmp, cfg) = test_config();
// 1. Create a target topic tree with a clean L0 buffer.
let topic_tree =
crate::openhuman::memory::tree::topic_tree::registry::get_or_create_topic_tree(
&cfg,
"email:alice@example.com",
)
.unwrap();
let l0_before = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
assert!(l0_before.is_empty());
// 2. Manually insert a summary node we can route. The simplest way
// is to create a separate source tree, push two 6k leaves into
// it, and let the seal produce a summary we can address.
let source_tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
use crate::openhuman::memory::tree::source_tree::bucket_seal::seal_one_level;
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{
chunk_id, Chunk, Metadata, SourceKind, SourceRef,
};
let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap();
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
for seq in 0..2 {
let chunk = Chunk {
id: chunk_id(SourceKind::Chat, "slack:#eng", seq, "summary-seed"),
content: format!("source content {seq}"),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: vec![],
source_ref: Some(SourceRef::new("slack://x")),
},
token_count: 6_000,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
// Stage to disk so `hydrate_leaf_inputs` can read the full body
// during `seal_one_level`.
let staged = content_store::stage_chunks(&content_root, &[chunk.clone()]).unwrap();
with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
let leaf = LeafRef {
chunk_id: chunk.id,
token_count: 6_000,
timestamp: ts,
content: chunk.content,
entities: vec![],
topics: vec![],
score: 0.5,
};
let _ = append_leaf_deferred(&cfg, &source_tree, &leaf).unwrap();
}
// Force-seal the source tree's L0 to mint the summary.
let buf = src_store::get_buffer(&cfg, &source_tree.id, 0).unwrap();
let summariser = build_summariser(&cfg);
let summary_id = seal_one_level(
&cfg,
&source_tree,
&buf,
summariser.as_ref(),
&crate::openhuman::memory::tree::source_tree::bucket_seal::LabelStrategy::Empty,
// No follow-up enqueues — the test scopes assertions to the
// append_buffer handler, not seal-side fan-out.
false,
)
.await
.unwrap();
// 3. Build an append_buffer payload routing the summary into the
// topic tree.
let payload = AppendBufferPayload {
node: NodeRef::Summary {
summary_id: summary_id.clone(),
},
target: AppendTarget::Topic {
tree_id: topic_tree.id.clone(),
},
};
let job = mk_running_job(
JobKind::AppendBuffer,
serde_json::to_string(&payload).unwrap(),
);
// Clear out any pending append_buffer jobs minted upstream so the
// post-condition assertion below is unambiguous.
let pre = count_total(&cfg).unwrap();
super::handle_append_buffer(&cfg, &job).await.unwrap();
// 4. Topic tree's L0 buffer should now hold the summary id.
let l0_after = src_store::get_buffer(&cfg, &topic_tree.id, 0).unwrap();
assert_eq!(l0_after.item_ids, vec![summary_id]);
assert!(l0_after.token_sum > 0);
// No new jobs should have been enqueued (buffer didn't cross gate).
assert_eq!(count_total(&cfg).unwrap(), pre);
}
}
+45
View File
@@ -0,0 +1,45 @@
//! Async job pipeline for memory-tree work.
//!
//! Replaces the previous synchronous `append_leaf → cascade_seal → LLM
//! summarise` chain on the ingest hot path with a SQLite-backed job queue
//! and a worker pool. The shape is:
//!
//! ```text
//! ingest::persist
//! └── writes chunk row (lifecycle = pending_extraction)
//! enqueues `extract_chunk`
//!
//! worker pool (3 tasks) ──► claims jobs by kind:
//! extract_chunk → LLM extraction → admission decision → enqueue append_buffer
//! append_buffer → push to L0 → enqueue seal if gate met → enqueue topic_route
//! seal → seal one level → enqueue parent seal if cascading
//! topic_route → match topics → enqueue per-topic append_buffer
//! digest_daily → call global_tree::digest::end_of_day_digest
//! flush_stale → enqueue seals for time-stale buffers
//!
//! scheduler (1 task) ──► daily wall-clock tick:
//! enqueues digest_daily(yesterday) + flush_stale(today)
//! ```
//!
//! All persistence lives in the same `chunks.db` as `mem_tree_chunks` so a
//! producer can insert its side-effect and its follow-up job in one tx.
//! See [`store::enqueue_tx`] for the in-tx producer entry point.
mod handlers;
pub mod scheduler;
pub mod store;
pub mod testing;
pub mod types;
mod worker;
pub use scheduler::{backfill_missing_digests, trigger_digest};
pub use store::{
claim_next, count_by_status, count_total, enqueue, enqueue_tx, get_job, mark_done, mark_failed,
recover_stale_locks, DEFAULT_LOCK_DURATION_MS,
};
pub use testing::drain_until_idle;
pub use types::{
AppendBufferPayload, AppendTarget, DigestDailyPayload, ExtractChunkPayload, FlushStalePayload,
Job, JobKind, JobStatus, NewJob, NodeRef, SealPayload, TopicRoutePayload,
};
pub use worker::{start, wake_workers};
+216
View File
@@ -0,0 +1,216 @@
use std::time::Duration;
use anyhow::Result;
use chrono::{Datelike, Duration as ChronoDuration, NaiveDate, TimeZone, Utc};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::jobs::store;
use crate::openhuman::memory::tree::jobs::types::{DigestDailyPayload, FlushStalePayload, NewJob};
static STARTED: std::sync::Once = std::sync::Once::new();
/// Start the daily wall-clock scheduler. Takes the full `Config` so the
/// digest enqueues match the same workspace + LLM settings the workers
/// see — not `Config::default()`.
pub fn start(config: Config) {
STARTED.call_once(|| {
tokio::spawn(async move {
loop {
if let Err(err) = enqueue_daily_jobs(&config) {
log::warn!("[memory_tree::jobs] scheduler enqueue failed: {err:#}");
}
tokio::time::sleep(next_sleep_duration()).await;
}
});
});
}
fn enqueue_daily_jobs(config: &Config) -> anyhow::Result<()> {
let now = Utc::now();
let yesterday = now.date_naive() - ChronoDuration::days(1);
let date_iso = yesterday.format("%Y-%m-%d").to_string();
if store::enqueue(
config,
&NewJob::digest_daily(&DigestDailyPayload {
date_iso: date_iso.clone(),
})?,
)?
.is_some()
{
super::worker::wake_workers();
}
let today_iso = now.date_naive().format("%Y-%m-%d").to_string();
if store::enqueue(
config,
&NewJob::flush_stale(&FlushStalePayload::default(), &today_iso)?,
)?
.is_some()
{
super::worker::wake_workers();
}
Ok(())
}
/// Manually enqueue a `digest_daily` job for `date`. Idempotent — if a
/// digest already ran for that day, the handler's `find_existing_daily`
/// check will return `Skipped` without doing any work; if a job for the
/// same date is already queued or running, the partial unique index on
/// `dedupe_key` suppresses the duplicate.
///
/// Useful for catch-up after the process was down across midnight, or
/// to force a re-run for testing / debugging.
pub fn trigger_digest(config: &Config, date: NaiveDate) -> Result<Option<String>> {
let payload = DigestDailyPayload {
date_iso: date.format("%Y-%m-%d").to_string(),
};
let job_id = store::enqueue(config, &NewJob::digest_daily(&payload)?)?;
if job_id.is_some() {
log::info!(
"[memory_tree::jobs] manual digest trigger enqueued date={} id={:?}",
payload.date_iso,
job_id.as_deref()
);
super::worker::wake_workers();
} else {
log::debug!(
"[memory_tree::jobs] manual digest trigger dedupe-suppressed date={} \
(an active job for this date already exists)",
payload.date_iso
);
}
Ok(job_id)
}
/// Enqueue `digest_daily` jobs for the last `days_back` calendar days
/// (excluding today). Catch-up helper for cases where the scheduler
/// missed days because the process was down.
///
/// Returns the number of jobs newly enqueued. Days that already have a
/// completed digest are still re-enqueued — the handler is idempotent
/// and skips them — so this is safe to call repeatedly.
pub fn backfill_missing_digests(config: &Config, days_back: i64) -> Result<usize> {
if days_back <= 0 {
return Ok(0);
}
let today = Utc::now().date_naive();
let mut enqueued = 0usize;
for offset in 1..=days_back {
let date = today - ChronoDuration::days(offset);
if trigger_digest(config, date)?.is_some() {
enqueued += 1;
}
}
log::info!(
"[memory_tree::jobs] backfill_missing_digests window={}d enqueued={}",
days_back,
enqueued
);
Ok(enqueued)
}
fn next_sleep_duration() -> Duration {
let now = Utc::now();
let tomorrow = now.date_naive() + ChronoDuration::days(1);
let next = Utc
.with_ymd_and_hms(tomorrow.year(), tomorrow.month(), tomorrow.day(), 0, 5, 0)
// UTC has no DST gaps/overlaps, so `single()` always returns
// `Some` for any valid (Y, M, D, h, m, s). Fallback retained
// only as a defensive belt-and-braces against future API churn.
.single()
.unwrap_or_else(|| now + ChronoDuration::hours(24));
(next - now)
.to_std()
.unwrap_or_else(|_| Duration::from_secs(60))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::jobs::store::{
claim_next, count_by_status, count_total, mark_done, DEFAULT_LOCK_DURATION_MS,
};
use crate::openhuman::memory::tree::jobs::types::JobStatus;
use tempfile::TempDir;
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
cfg.workspace_dir = tmp.path().to_path_buf();
cfg.memory_tree.embedding_endpoint = None;
cfg.memory_tree.embedding_model = None;
cfg.memory_tree.embedding_strict = false;
(tmp, cfg)
}
#[test]
fn trigger_digest_enqueues_a_job() {
let (_tmp, cfg) = test_config();
let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap();
let id = trigger_digest(&cfg, date).unwrap();
assert!(id.is_some(), "first trigger must enqueue");
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1);
}
#[test]
fn trigger_digest_dedupes_active_jobs() {
let (_tmp, cfg) = test_config();
let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap();
let first = trigger_digest(&cfg, date).unwrap();
let second = trigger_digest(&cfg, date).unwrap();
assert!(first.is_some());
assert!(
second.is_none(),
"duplicate trigger must be dedupe-suppressed while active"
);
assert_eq!(count_total(&cfg).unwrap(), 1);
}
#[test]
fn trigger_digest_after_done_creates_fresh_row() {
let (_tmp, cfg) = test_config();
let date = NaiveDate::from_ymd_opt(2026, 4, 27).unwrap();
let id1 = trigger_digest(&cfg, date).unwrap().unwrap();
// Simulate a worker finishing the job — claim it first so we have a
// Job snapshot for the claim-token-gated mark_done.
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
assert_eq!(claimed.id, id1);
mark_done(&cfg, &claimed).unwrap();
let id2 = trigger_digest(&cfg, date).unwrap();
assert!(
id2.is_some(),
"after the prior job completes, a fresh trigger must enqueue"
);
assert_ne!(id2.unwrap(), id1);
assert_eq!(count_total(&cfg).unwrap(), 2);
}
#[test]
fn backfill_missing_digests_enqueues_one_per_day() {
let (_tmp, cfg) = test_config();
let n = backfill_missing_digests(&cfg, 5).unwrap();
assert_eq!(n, 5, "expected one job per day in the 5-day window");
assert_eq!(count_total(&cfg).unwrap(), 5);
}
#[test]
fn backfill_missing_digests_zero_window_is_noop() {
let (_tmp, cfg) = test_config();
let n = backfill_missing_digests(&cfg, 0).unwrap();
assert_eq!(n, 0);
assert_eq!(count_total(&cfg).unwrap(), 0);
}
#[test]
fn backfill_missing_digests_is_idempotent_while_active() {
let (_tmp, cfg) = test_config();
let n1 = backfill_missing_digests(&cfg, 3).unwrap();
let n2 = backfill_missing_digests(&cfg, 3).unwrap();
assert_eq!(n1, 3);
assert_eq!(n2, 0, "second call must be fully dedupe-suppressed");
assert_eq!(count_total(&cfg).unwrap(), 3);
}
}
+618
View File
@@ -0,0 +1,618 @@
//! SQLite persistence for the memory-tree job queue.
//!
//! Producers call [`enqueue`] inside their own writes (or with a fresh tx)
//! to atomically commit the side-effect plus its follow-up job. The worker
//! pool calls [`claim_next`] to lease a job, [`mark_done`] / [`mark_failed`]
//! to settle it, and [`recover_stale_locks`] on startup to flip rows whose
//! `locked_until_ms` expired without a settle.
//!
//! Concurrency:
//! - The dedupe key is enforced by a partial `UNIQUE` index that only
//! covers `status IN ('ready', 'running')`. Producers use `INSERT OR
//! IGNORE` so a duplicate enqueue while a job is in flight or queued is
//! a silent no-op; a duplicate enqueue after the first completes is
//! accepted and creates a fresh row.
//! - `claim_next` is one statement: `UPDATE … WHERE id = (SELECT … LIMIT 1)
//! RETURNING …`. SQLite serialises writes, so no two workers can claim
//! the same row.
use anyhow::{Context, Result};
use chrono::Utc;
use rusqlite::{params, Connection, OptionalExtension, Transaction};
use uuid::Uuid;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::jobs::types::{Job, JobKind, JobStatus, NewJob};
use crate::openhuman::memory::tree::store::with_connection;
/// Default visibility lock — a worker that crashes mid-job will have its
/// row recovered after this window. 5 min is comfortably larger than any
/// expected single-job runtime (LLM extract or summarise) without leaving
/// real failures stuck for hours.
pub const DEFAULT_LOCK_DURATION_MS: i64 = 5 * 60 * 1_000;
/// Backoff math for retry. Returns `now + min(base * 2^attempts, cap)`.
const RETRY_BASE_MS: i64 = 60 * 1_000;
const RETRY_CAP_MS: i64 = 60 * 60 * 1_000;
const DEFAULT_MAX_ATTEMPTS: u32 = 5;
/// Enqueue one job. Idempotent on `dedupe_key` while another active row
/// (status `ready`/`running`) shares it. Returns `Some(id)` if the row
/// was inserted, `None` if a duplicate was suppressed.
pub fn enqueue(config: &Config, job: &NewJob) -> Result<Option<String>> {
with_connection(config, |conn| enqueue_conn(conn, job))
}
/// Enqueue inside a caller-owned transaction. Use this when the producer
/// is already mid-tx (e.g. `ingest::persist` writing chunks + jobs in one
/// commit) so the queue insert lands atomically with the side-effect.
/// `Transaction` derefs to `Connection`, so callers just pass `&tx`.
pub fn enqueue_tx(tx: &Transaction<'_>, job: &NewJob) -> Result<Option<String>> {
enqueue_conn(tx, job)
}
pub(crate) fn enqueue_conn(conn: &Connection, job: &NewJob) -> Result<Option<String>> {
let id = format!("job:{}", Uuid::new_v4());
let now_ms = Utc::now().timestamp_millis();
let available_at = job.available_at_ms.unwrap_or(now_ms);
let max_attempts = job.max_attempts.unwrap_or(DEFAULT_MAX_ATTEMPTS) as i64;
let inserted = conn.execute(
"INSERT OR IGNORE INTO mem_tree_jobs (
id, kind, payload_json, dedupe_key, status, attempts, max_attempts,
available_at_ms, locked_until_ms, last_error,
created_at_ms, started_at_ms, completed_at_ms
) VALUES (?1, ?2, ?3, ?4, 'ready', 0, ?5, ?6, NULL, NULL, ?7, NULL, NULL)",
params![
id,
job.kind.as_str(),
job.payload_json,
job.dedupe_key,
max_attempts,
available_at,
now_ms,
],
)?;
if inserted == 0 {
log::debug!(
"[memory_tree::jobs] enqueue suppressed by dedupe kind={} key={:?}",
job.kind.as_str(),
job.dedupe_key
);
return Ok(None);
}
log::debug!(
"[memory_tree::jobs] enqueued id={} kind={} avail_at_ms={} dedupe={:?}",
id,
job.kind.as_str(),
available_at,
job.dedupe_key
);
Ok(Some(id))
}
/// Atomically claim the next ready job whose `available_at_ms` has come
/// due. Sets `status=running`, bumps `attempts`, stamps `started_at_ms`
/// and `locked_until_ms`. Returns `None` when the queue is empty / not
/// yet due.
pub fn claim_next(config: &Config, lock_duration_ms: i64) -> Result<Option<Job>> {
with_connection(config, |conn| {
let now_ms = Utc::now().timestamp_millis();
let lock_until = now_ms.saturating_add(lock_duration_ms);
let row = conn
.query_row(
"UPDATE mem_tree_jobs
SET status = 'running',
attempts = attempts + 1,
started_at_ms = ?1,
locked_until_ms = ?2,
last_error = NULL
WHERE id = (
SELECT id FROM mem_tree_jobs
WHERE status = 'ready'
AND available_at_ms <= ?1
ORDER BY available_at_ms ASC
LIMIT 1
)
RETURNING id, kind, payload_json, dedupe_key, status, attempts,
max_attempts, available_at_ms, locked_until_ms, last_error,
created_at_ms, started_at_ms, completed_at_ms",
params![now_ms, lock_until],
row_to_job,
)
.optional()
.context("Failed to claim next mem_tree_jobs row")?;
if let Some(j) = &row {
log::debug!(
"[memory_tree::jobs] claimed id={} kind={} attempt={}/{}",
j.id,
j.kind.as_str(),
j.attempts,
j.max_attempts
);
}
Ok(row)
})
}
/// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`.
///
/// The UPDATE is gated on `attempts` and `started_at_ms` matching the values
/// in `job` (the snapshot returned by [`claim_next`]). If the lease expired
/// and another worker re-claimed the row, `rows_affected` will be 0 — the
/// stale worker's settlement is a silent no-op rather than clobbering the new
/// lessee's state.
pub fn mark_done(config: &Config, job: &Job) -> Result<()> {
let job_id = &job.id;
let claim_attempts = job.attempts as i64;
let claim_started_at = job.started_at_ms;
with_connection(config, |conn| {
let now_ms = Utc::now().timestamp_millis();
let n = conn.execute(
"UPDATE mem_tree_jobs
SET status = 'done',
completed_at_ms = ?1,
locked_until_ms = NULL,
last_error = NULL
WHERE id = ?2
AND attempts = ?3
AND started_at_ms IS ?4",
params![now_ms, job_id, claim_attempts, claim_started_at],
)?;
if n == 0 {
// Either the job row was deleted (shouldn't happen) or the lease
// expired and a second worker re-claimed the row. Log and move on —
// this is a known race outcome, not a bug in the current worker.
log::warn!(
"[memory_tree::jobs] mark_done id={job_id} was a no-op \
(stale lease: attempts={claim_attempts} started_at_ms={claim_started_at:?})"
);
}
Ok(())
})
}
/// Settle a failed job. If `attempts < max_attempts`, the row goes back
/// to `ready` with an exponential-backoff `available_at_ms`. Otherwise
/// it terminates as `failed`. Either way `last_error` is recorded.
///
/// Like [`mark_done`], the UPDATE is gated on the claim-token
/// (`attempts` + `started_at_ms`) so a stale worker's failure settlement
/// cannot clobber an active lessee's row — rows_affected == 0 is a silent
/// no-op.
pub fn mark_failed(config: &Config, job: &Job, error: &str) -> Result<()> {
let job_id = &job.id;
let attempts = job.attempts as i64;
let max_attempts = job.max_attempts as i64;
let claim_started_at = job.started_at_ms;
with_connection(config, |conn| {
let now_ms = Utc::now().timestamp_millis();
if attempts >= max_attempts {
log::warn!(
"[memory_tree::jobs] terminal failure id={job_id} \
attempts={attempts}/{max_attempts} err={error}"
);
let n = conn.execute(
"UPDATE mem_tree_jobs
SET status = 'failed',
completed_at_ms = ?1,
locked_until_ms = NULL,
last_error = ?2
WHERE id = ?3
AND attempts = ?4
AND started_at_ms IS ?5",
params![now_ms, error, job_id, attempts, claim_started_at],
)?;
if n == 0 {
log::warn!(
"[memory_tree::jobs] mark_failed(terminal) id={job_id} was a no-op \
(stale lease: attempts={attempts} started_at_ms={claim_started_at:?})"
);
}
} else {
let backoff = backoff_ms(attempts as u32);
let next_at = now_ms.saturating_add(backoff);
log::info!(
"[memory_tree::jobs] retry id={job_id} attempt={attempts}/{max_attempts} \
next_at_ms={next_at} err={error}"
);
let n = conn.execute(
"UPDATE mem_tree_jobs
SET status = 'ready',
available_at_ms = ?1,
locked_until_ms = NULL,
last_error = ?2
WHERE id = ?3
AND attempts = ?4
AND started_at_ms IS ?5",
params![next_at, error, job_id, attempts, claim_started_at],
)?;
if n == 0 {
log::warn!(
"[memory_tree::jobs] mark_failed(retry) id={job_id} was a no-op \
(stale lease: attempts={attempts} started_at_ms={claim_started_at:?})"
);
}
}
Ok(())
})
}
/// Flip any `running` row whose `locked_until_ms` has expired back to
/// `ready`. Called once at worker startup so a process crash mid-job
/// doesn't leave work stranded. Returns the number of rows recovered.
pub fn recover_stale_locks(config: &Config) -> Result<usize> {
with_connection(config, |conn| {
let now_ms = Utc::now().timestamp_millis();
let n = conn.execute(
"UPDATE mem_tree_jobs
SET status = 'ready',
last_error = COALESCE(last_error, 'recovered_from_stale_lock')
WHERE status = 'running'
AND locked_until_ms IS NOT NULL
AND locked_until_ms < ?1",
params![now_ms],
)?;
if n > 0 {
log::warn!("[memory_tree::jobs] recovered {n} stale-locked job(s) at startup");
}
Ok(n)
})
}
/// Quick count helper for tests / diagnostics.
pub fn count_by_status(config: &Config, status: JobStatus) -> Result<u64> {
with_connection(config, |conn| {
let n: i64 = conn.query_row(
"SELECT COUNT(*) FROM mem_tree_jobs WHERE status = ?1",
params![status.as_str()],
|r| r.get(0),
)?;
Ok(n.max(0) as u64)
})
}
/// Total count regardless of status — handy for assertions.
pub fn count_total(config: &Config) -> Result<u64> {
with_connection(config, |conn| {
let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| r.get(0))?;
Ok(n.max(0) as u64)
})
}
/// Fetch one job by id (test/diagnostic helper).
pub fn get_job(config: &Config, id: &str) -> Result<Option<Job>> {
with_connection(config, |conn| {
let job = conn
.query_row(
"SELECT id, kind, payload_json, dedupe_key, status, attempts, max_attempts,
available_at_ms, locked_until_ms, last_error,
created_at_ms, started_at_ms, completed_at_ms
FROM mem_tree_jobs WHERE id = ?1",
params![id],
row_to_job,
)
.optional()?;
Ok(job)
})
}
fn row_to_job(row: &rusqlite::Row<'_>) -> rusqlite::Result<Job> {
let id: String = row.get(0)?;
let kind_s: String = row.get(1)?;
let payload_json: String = row.get(2)?;
let dedupe_key: Option<String> = row.get(3)?;
let status_s: String = row.get(4)?;
let attempts: i64 = row.get(5)?;
let max_attempts: i64 = row.get(6)?;
let available_at_ms: i64 = row.get(7)?;
let locked_until_ms: Option<i64> = row.get(8)?;
let last_error: Option<String> = row.get(9)?;
let created_at_ms: i64 = row.get(10)?;
let started_at_ms: Option<i64> = row.get(11)?;
let completed_at_ms: Option<i64> = row.get(12)?;
let kind = JobKind::parse(&kind_s).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(1, rusqlite::types::Type::Text, e.into())
})?;
let status = JobStatus::parse(&status_s).map_err(|e| {
rusqlite::Error::FromSqlConversionFailure(4, rusqlite::types::Type::Text, e.into())
})?;
Ok(Job {
id,
kind,
payload_json,
dedupe_key,
status,
attempts: attempts.max(0) as u32,
max_attempts: max_attempts.max(0) as u32,
available_at_ms,
locked_until_ms,
last_error,
created_at_ms,
started_at_ms,
completed_at_ms,
})
}
/// Exponential backoff: attempt 1 → 60s, 2 → 120s, 3 → 240s, capped at 1h.
fn backoff_ms(attempts_so_far: u32) -> i64 {
// attempts_so_far is the count BEFORE the next retry's attempt — so the
// first retry uses attempts_so_far=1, giving base*2^0 = 60s.
let exp = attempts_so_far.saturating_sub(1).min(20); // cap shift
let mult = 1i64 << exp; // 1, 2, 4, …
let raw = RETRY_BASE_MS.saturating_mul(mult);
raw.min(RETRY_CAP_MS)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::jobs::types::{
AppendBufferPayload, AppendTarget, ExtractChunkPayload, NodeRef,
};
use tempfile::TempDir;
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
cfg.workspace_dir = tmp.path().to_path_buf();
cfg.memory_tree.embedding_endpoint = None;
cfg.memory_tree.embedding_model = None;
cfg.memory_tree.embedding_strict = false;
(tmp, cfg)
}
#[test]
fn enqueue_and_claim_roundtrip() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c1".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
assert_eq!(claimed.id, id);
assert_eq!(claimed.status, JobStatus::Running);
assert_eq!(claimed.attempts, 1);
assert!(claimed.locked_until_ms.is_some());
// Second claim should see no eligible row (the only one is now running).
let again = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap();
assert!(again.is_none());
mark_done(&cfg, &claimed).unwrap();
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(row.status, JobStatus::Done);
assert!(row.completed_at_ms.is_some());
assert!(row.locked_until_ms.is_none());
}
#[test]
fn enqueue_dedupes_active_jobs() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c1".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id1 = enqueue(&cfg, &nj).unwrap();
let id2 = enqueue(&cfg, &nj).unwrap();
assert!(id1.is_some());
assert!(id2.is_none(), "duplicate should be suppressed while ready");
assert_eq!(count_total(&cfg).unwrap(), 1);
}
#[test]
fn enqueue_after_done_creates_fresh_row() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c1".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id1 = enqueue(&cfg, &nj).unwrap().unwrap();
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
assert_eq!(claimed.id, id1);
mark_done(&cfg, &claimed).unwrap();
// Now the dedupe key is free (partial index excludes 'done').
let id2 = enqueue(&cfg, &nj).unwrap();
assert!(id2.is_some());
assert_ne!(id2.unwrap(), id1);
assert_eq!(count_total(&cfg).unwrap(), 2);
}
#[test]
fn mark_failed_retries_then_terminates() {
let (_tmp, cfg) = test_config();
let payload = AppendBufferPayload {
node: NodeRef::Leaf {
chunk_id: "c1".into(),
},
target: AppendTarget::Source {
source_id: "slack:#x".into(),
},
};
let mut nj = NewJob::append_buffer(&payload).unwrap();
nj.max_attempts = Some(2);
let id = enqueue(&cfg, &nj).unwrap().unwrap();
// Fail #1 — should bounce back to 'ready' with future available_at.
let attempt1 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
mark_failed(&cfg, &attempt1, "boom").unwrap();
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(row.status, JobStatus::Ready);
assert!(row.available_at_ms > Utc::now().timestamp_millis());
assert_eq!(row.last_error.as_deref(), Some("boom"));
// Force the row available again so the test doesn't hinge on sleep.
with_connection(&cfg, |c| {
c.execute(
"UPDATE mem_tree_jobs SET available_at_ms = 0 WHERE id = ?1",
params![id],
)?;
Ok(())
})
.unwrap();
// Fail #2 — exceeds max_attempts → terminal 'failed'.
let attempt2 = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
mark_failed(&cfg, &attempt2, "fatal").unwrap();
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(row.status, JobStatus::Failed);
assert_eq!(row.last_error.as_deref(), Some("fatal"));
assert!(row.completed_at_ms.is_some());
}
#[test]
fn recover_stale_locks_resets_running_rows() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c1".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id = enqueue(&cfg, &nj).unwrap().unwrap();
// Claim with a lock window that's already in the past so recovery
// sees it as expired.
let _ = claim_next(&cfg, -1).unwrap().unwrap();
let recovered = recover_stale_locks(&cfg).unwrap();
assert_eq!(recovered, 1);
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(row.status, JobStatus::Ready);
}
/// Happy path: a non-stale settlement still succeeds after the claim-token
/// check is applied. Regression guard so the common case isn't broken.
#[test]
fn mark_done_succeeds_for_current_lessee() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c-happy".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
assert_eq!(claimed.id, id);
// Current lessee should settle successfully.
mark_done(&cfg, &claimed).unwrap();
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(row.status, JobStatus::Done);
assert!(row.completed_at_ms.is_some());
assert!(row.locked_until_ms.is_none());
}
/// Stale-worker settlement is a no-op: after a lock expires and a second
/// worker re-claims the job, the first worker's `mark_done` must not
/// clobber the new lessee's row.
#[test]
fn stale_worker_settlement_is_noop() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c-stale".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
// Worker A claims with a lock that's already expired (negative window).
let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap();
assert_eq!(worker_a_job.id, id);
assert_eq!(worker_a_job.attempts, 1);
// Simulate lease expiry: recover_stale_locks resets the row to 'ready'.
let recovered = recover_stale_locks(&cfg).unwrap();
assert_eq!(recovered, 1);
// Worker B claims the reset row — different lease token (attempts=2).
let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
assert_eq!(worker_b_job.id, id);
assert_eq!(worker_b_job.attempts, 2);
// Worker A (stale) tries to mark done using its old claim snapshot.
mark_done(&cfg, &worker_a_job).unwrap(); // must NOT return Err
// Worker B's row must be untouched — still 'running' with attempts=2.
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(
row.status,
JobStatus::Running,
"stale settlement must not clobber Worker B's running row"
);
assert_eq!(
row.attempts, 2,
"attempts must still reflect Worker B's claim"
);
}
/// Same contract as stale_worker_settlement_is_noop but for mark_failed.
#[test]
fn stale_worker_mark_failed_is_noop() {
let (_tmp, cfg) = test_config();
let payload = ExtractChunkPayload {
chunk_id: "c-stale-fail".into(),
};
let nj = NewJob::extract_chunk(&payload).unwrap();
let id = enqueue(&cfg, &nj).unwrap().expect("inserted");
// Worker A claims with an already-expired lock.
let worker_a_job = claim_next(&cfg, -1).unwrap().unwrap();
assert_eq!(worker_a_job.attempts, 1);
// Lease expires, recovered, Worker B re-claims.
let recovered = recover_stale_locks(&cfg).unwrap();
assert_eq!(recovered, 1);
let worker_b_job = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
assert_eq!(worker_b_job.attempts, 2);
// Worker A (stale) tries to record a failure — must be a no-op.
mark_failed(&cfg, &worker_a_job, "stale error").unwrap();
// Worker B's row must be untouched.
let row = get_job(&cfg, &id).unwrap().unwrap();
assert_eq!(
row.status,
JobStatus::Running,
"stale mark_failed must not clobber Worker B's running row"
);
assert_ne!(
row.last_error.as_deref(),
Some("stale error"),
"stale error must not be written to the row"
);
assert_eq!(row.attempts, 2);
}
#[test]
fn backoff_grows_then_caps() {
assert_eq!(backoff_ms(1), 60_000);
assert_eq!(backoff_ms(2), 120_000);
assert_eq!(backoff_ms(3), 240_000);
// Eventually clamps at the cap.
assert_eq!(backoff_ms(20), RETRY_CAP_MS);
assert_eq!(backoff_ms(99), RETRY_CAP_MS);
}
#[test]
fn count_by_status_reports_each_state() {
let (_tmp, cfg) = test_config();
for i in 0..3 {
let p = ExtractChunkPayload {
chunk_id: format!("c{i}"),
};
let nj = NewJob::extract_chunk(&p).unwrap();
enqueue(&cfg, &nj).unwrap();
}
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 3);
let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap();
mark_done(&cfg, &claimed).unwrap();
assert_eq!(count_by_status(&cfg, JobStatus::Done).unwrap(), 1);
assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 2);
}
}
+15
View File
@@ -0,0 +1,15 @@
use anyhow::Result;
use crate::openhuman::config::Config;
/// Deterministically run queued memory-tree jobs until no immediately
/// claimable work remains. Intended for tests that need the async pipeline
/// to settle without spawning background tasks.
pub async fn drain_until_idle(config: &Config) -> Result<()> {
loop {
if !super::worker::run_once(config).await? {
break;
}
}
Ok(())
}
+439
View File
@@ -0,0 +1,439 @@
//! Job types for the async memory-tree pipeline.
//!
//! Each `Job` row in `mem_tree_jobs` stores its discriminator as a string
//! `kind` plus a JSON-encoded `payload`. The strongly-typed payload structs
//! below own (de)serialisation; handlers parse the payload by branching on
//! [`JobKind`] and calling the matching `from_payload_json`.
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
/// Discriminator persisted in `mem_tree_jobs.kind`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum JobKind {
/// Run LLM entity extraction over a single chunk and decide admission.
ExtractChunk,
/// Push an admitted chunk into a tree's L0 buffer.
AppendBuffer,
/// Seal exactly one buffer level; cascades enqueue a follow-up.
Seal,
/// Match a chunk's entities against active topic trees and enqueue
/// per-topic `AppendBuffer` jobs.
TopicRoute,
/// Build the global tree's daily digest for a given UTC date.
DigestDaily,
/// Walk stale buffers and enqueue `Seal` jobs for any over the age cap.
FlushStale,
}
impl JobKind {
pub fn as_str(&self) -> &'static str {
match self {
JobKind::ExtractChunk => "extract_chunk",
JobKind::AppendBuffer => "append_buffer",
JobKind::Seal => "seal",
JobKind::TopicRoute => "topic_route",
JobKind::DigestDaily => "digest_daily",
JobKind::FlushStale => "flush_stale",
}
}
pub fn parse(s: &str) -> Result<Self> {
Ok(match s {
"extract_chunk" => JobKind::ExtractChunk,
"append_buffer" => JobKind::AppendBuffer,
"seal" => JobKind::Seal,
"topic_route" => JobKind::TopicRoute,
"digest_daily" => JobKind::DigestDaily,
"flush_stale" => JobKind::FlushStale,
other => return Err(anyhow!("unknown JobKind '{other}'")),
})
}
/// True when handling this kind should hold a slot from the global
/// LLM concurrency semaphore. `TopicRoute` is bound because
/// `maybe_spawn_topic_tree → backfill_topic_tree` can transitively
/// trigger summariser LLM calls when an entity first crosses the
/// hotness threshold.
pub fn is_llm_bound(&self) -> bool {
matches!(
self,
JobKind::ExtractChunk | JobKind::Seal | JobKind::DigestDaily | JobKind::TopicRoute
)
}
}
/// Lifecycle states persisted on `mem_tree_jobs.status`. Workers transition
/// `ready → running → done|failed`. `Cancelled` is reserved for explicit
/// admin actions (none surfaced yet).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum JobStatus {
Ready,
Running,
Done,
Failed,
Cancelled,
}
impl JobStatus {
pub fn as_str(&self) -> &'static str {
match self {
JobStatus::Ready => "ready",
JobStatus::Running => "running",
JobStatus::Done => "done",
JobStatus::Failed => "failed",
JobStatus::Cancelled => "cancelled",
}
}
pub fn parse(s: &str) -> Result<Self> {
Ok(match s {
"ready" => JobStatus::Ready,
"running" => JobStatus::Running,
"done" => JobStatus::Done,
"failed" => JobStatus::Failed,
"cancelled" => JobStatus::Cancelled,
other => return Err(anyhow!("unknown JobStatus '{other}'")),
})
}
pub fn is_terminal(&self) -> bool {
matches!(
self,
JobStatus::Done | JobStatus::Failed | JobStatus::Cancelled
)
}
}
// ── Payloads ───────────────────────────────────────────────────────────────
/// Reference to either a leaf chunk or a sealed summary node. Used by
/// payloads that route content through the pipeline regardless of which
/// kind of source produced it.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum NodeRef {
Leaf { chunk_id: String },
Summary { summary_id: String },
}
impl NodeRef {
/// Stringified id with kind prefix, suitable for dedupe-key composition.
pub fn dedupe_fragment(&self) -> String {
match self {
NodeRef::Leaf { chunk_id } => format!("leaf:{chunk_id}"),
NodeRef::Summary { summary_id } => format!("summary:{summary_id}"),
}
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ExtractChunkPayload {
pub chunk_id: String,
}
impl ExtractChunkPayload {
pub fn dedupe_key(&self) -> String {
format!("extract:{}", self.chunk_id)
}
}
/// Where an `AppendBuffer` job should land its node. Source-tree appends
/// are keyed by `source_id`; topic-tree appends are keyed by `tree_id`
/// because there can be many topic trees per node.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum AppendTarget {
Source { source_id: String },
Topic { tree_id: String },
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct AppendBufferPayload {
pub node: NodeRef,
pub target: AppendTarget,
}
impl AppendBufferPayload {
pub fn dedupe_key(&self) -> String {
let node_part = self.node.dedupe_fragment();
match &self.target {
AppendTarget::Source { source_id } => {
format!("append:source:{source_id}:{node_part}")
}
AppendTarget::Topic { tree_id } => {
format!("append:topic:{tree_id}:{node_part}")
}
}
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SealPayload {
pub tree_id: String,
pub level: u32,
/// When `Some`, the seal handler bypasses the buffer-budget check and
/// force-seals — used by the time-based flush path. The wall-clock is
/// passed through so the seal stamps a deterministic `sealed_at`.
pub force_now_ms: Option<i64>,
}
impl SealPayload {
pub fn dedupe_key(&self) -> String {
// Active seal-job uniqueness is enforced per (tree, level): a seal
// already in flight suppresses duplicate enqueues. Once the job
// completes the partial index releases the key, so the next time
// the buffer crosses its gate a fresh seal can be enqueued.
format!("seal:{}:{}", self.tree_id, self.level)
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TopicRoutePayload {
pub node: NodeRef,
}
impl TopicRoutePayload {
pub fn dedupe_key(&self) -> String {
format!("topic_route:{}", self.node.dedupe_fragment())
}
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct DigestDailyPayload {
/// UTC calendar date in `YYYY-MM-DD` form. Stored as a string so the
/// dedupe key doesn't need to know about chrono.
pub date_iso: String,
}
impl DigestDailyPayload {
pub fn dedupe_key(&self) -> String {
format!("digest_daily:{}", self.date_iso)
}
}
#[derive(Clone, Debug, Serialize, Deserialize, Default)]
pub struct FlushStalePayload {
/// Override the configured `DEFAULT_FLUSH_AGE_SECS`. Optional so the
/// scheduler can enqueue with `None` and let the handler use the
/// configured default.
pub max_age_secs: Option<i64>,
}
impl FlushStalePayload {
pub fn dedupe_key(&self, date_iso: &str) -> String {
format!("flush_stale:{date_iso}")
}
}
/// One row in `mem_tree_jobs`. `payload_json` is left as a raw string so
/// callers parse it lazily based on `kind`.
#[derive(Clone, Debug)]
pub struct Job {
pub id: String,
pub kind: JobKind,
pub payload_json: String,
pub dedupe_key: Option<String>,
pub status: JobStatus,
pub attempts: u32,
pub max_attempts: u32,
pub available_at_ms: i64,
pub locked_until_ms: Option<i64>,
pub last_error: Option<String>,
pub created_at_ms: i64,
pub started_at_ms: Option<i64>,
pub completed_at_ms: Option<i64>,
}
/// Caller-side bundle for `enqueue` — `Job` minus the persistence-only
/// columns. Keeps producers from having to mint timestamps and ids by hand.
#[derive(Clone, Debug)]
pub struct NewJob {
pub kind: JobKind,
pub payload_json: String,
pub dedupe_key: Option<String>,
/// `None` means "available immediately." Set this for delayed jobs
/// (retries, scheduled work).
pub available_at_ms: Option<i64>,
pub max_attempts: Option<u32>,
}
impl NewJob {
pub fn extract_chunk(p: &ExtractChunkPayload) -> Result<Self> {
Ok(Self {
kind: JobKind::ExtractChunk,
payload_json: serde_json::to_string(p)?,
dedupe_key: Some(p.dedupe_key()),
available_at_ms: None,
max_attempts: None,
})
}
pub fn append_buffer(p: &AppendBufferPayload) -> Result<Self> {
Ok(Self {
kind: JobKind::AppendBuffer,
payload_json: serde_json::to_string(p)?,
dedupe_key: Some(p.dedupe_key()),
available_at_ms: None,
max_attempts: None,
})
}
pub fn seal(p: &SealPayload) -> Result<Self> {
Ok(Self {
kind: JobKind::Seal,
payload_json: serde_json::to_string(p)?,
dedupe_key: Some(p.dedupe_key()),
available_at_ms: None,
max_attempts: None,
})
}
pub fn topic_route(p: &TopicRoutePayload) -> Result<Self> {
Ok(Self {
kind: JobKind::TopicRoute,
payload_json: serde_json::to_string(p)?,
dedupe_key: Some(p.dedupe_key()),
available_at_ms: None,
max_attempts: None,
})
}
pub fn digest_daily(p: &DigestDailyPayload) -> Result<Self> {
Ok(Self {
kind: JobKind::DigestDaily,
payload_json: serde_json::to_string(p)?,
dedupe_key: Some(p.dedupe_key()),
available_at_ms: None,
max_attempts: None,
})
}
pub fn flush_stale(p: &FlushStalePayload, date_iso: &str) -> Result<Self> {
Ok(Self {
kind: JobKind::FlushStale,
payload_json: serde_json::to_string(p)?,
dedupe_key: Some(p.dedupe_key(date_iso)),
available_at_ms: None,
max_attempts: None,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn job_kind_roundtrip() {
for k in [
JobKind::ExtractChunk,
JobKind::AppendBuffer,
JobKind::Seal,
JobKind::TopicRoute,
JobKind::DigestDaily,
JobKind::FlushStale,
] {
assert_eq!(JobKind::parse(k.as_str()).unwrap(), k);
}
}
#[test]
fn job_status_terminality() {
assert!(!JobStatus::Ready.is_terminal());
assert!(!JobStatus::Running.is_terminal());
assert!(JobStatus::Done.is_terminal());
assert!(JobStatus::Failed.is_terminal());
assert!(JobStatus::Cancelled.is_terminal());
}
#[test]
fn dedupe_keys_distinguish_targets() {
let p_src = AppendBufferPayload {
node: NodeRef::Leaf {
chunk_id: "c1".into(),
},
target: AppendTarget::Source {
source_id: "slack:#eng".into(),
},
};
let p_topic = AppendBufferPayload {
node: NodeRef::Leaf {
chunk_id: "c1".into(),
},
target: AppendTarget::Topic {
tree_id: "topic:abc".into(),
},
};
assert_ne!(p_src.dedupe_key(), p_topic.dedupe_key());
}
#[test]
fn dedupe_keys_distinguish_node_kinds() {
let p_leaf = AppendBufferPayload {
node: NodeRef::Leaf {
chunk_id: "x".into(),
},
target: AppendTarget::Topic {
tree_id: "t".into(),
},
};
let p_summary = AppendBufferPayload {
node: NodeRef::Summary {
summary_id: "x".into(),
},
target: AppendTarget::Topic {
tree_id: "t".into(),
},
};
assert_ne!(p_leaf.dedupe_key(), p_summary.dedupe_key());
let r_leaf = TopicRoutePayload {
node: NodeRef::Leaf {
chunk_id: "x".into(),
},
};
let r_summary = TopicRoutePayload {
node: NodeRef::Summary {
summary_id: "x".into(),
},
};
assert_ne!(r_leaf.dedupe_key(), r_summary.dedupe_key());
}
#[test]
fn llm_bound_kinds() {
assert!(JobKind::ExtractChunk.is_llm_bound());
assert!(JobKind::Seal.is_llm_bound());
assert!(JobKind::DigestDaily.is_llm_bound());
assert!(JobKind::TopicRoute.is_llm_bound());
assert!(!JobKind::AppendBuffer.is_llm_bound());
assert!(!JobKind::FlushStale.is_llm_bound());
}
#[test]
fn node_ref_serializes_with_kind_tag() {
let leaf = NodeRef::Leaf {
chunk_id: "x".into(),
};
let s = serde_json::to_string(&leaf).unwrap();
assert!(s.contains("\"kind\":\"leaf\""));
let back: NodeRef = serde_json::from_str(&s).unwrap();
assert_eq!(back, leaf);
}
#[test]
fn append_target_serializes_with_kind_tag() {
let p = AppendTarget::Source {
source_id: "x".into(),
};
let s = serde_json::to_string(&p).unwrap();
assert!(s.contains("\"kind\":\"source\""));
assert!(s.contains("\"source_id\":\"x\""));
let back: AppendTarget = serde_json::from_str(&s).unwrap();
match back {
AppendTarget::Source { source_id } => assert_eq!(source_id, "x"),
_ => panic!("wrong variant"),
}
}
}
+109
View File
@@ -0,0 +1,109 @@
use std::sync::{Arc, OnceLock};
use std::time::Duration;
use anyhow::Result;
use tokio::sync::{Notify, Semaphore};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::jobs::handlers;
use crate::openhuman::memory::tree::jobs::store::{
claim_next, mark_done, mark_failed, recover_stale_locks, DEFAULT_LOCK_DURATION_MS,
};
const WORKER_COUNT: usize = 3;
const POLL_INTERVAL: Duration = Duration::from_secs(5);
static WORKER_NOTIFY: OnceLock<Arc<Notify>> = OnceLock::new();
static STARTED: std::sync::Once = std::sync::Once::new();
pub fn wake_workers() {
if let Some(notify) = WORKER_NOTIFY.get() {
notify.notify_waiters();
}
}
/// Start the worker pool + daily scheduler. Takes the full `Config` so
/// each spawned task sees the user's actual settings (LLM endpoints,
/// embedder model, timeouts) — not `Config::default()`. Without this,
/// workers fall back to inert/regex-only behavior regardless of what's
/// in `config.toml`, defeating the entire async pipeline.
///
/// Idempotent (`Once`-guarded) so repeat calls during bootstrap are
/// safe no-ops after the first.
pub fn start(config: Config) {
STARTED.call_once(|| {
let notify = WORKER_NOTIFY
.get_or_init(|| Arc::new(Notify::new()))
.clone();
let llm_slots = Arc::new(Semaphore::new(3));
if let Err(err) = recover_stale_locks(&config) {
log::warn!("[memory_tree::jobs] recover_stale_locks failed at startup: {err:#}");
}
for idx in 0..WORKER_COUNT {
let notify = notify.clone();
let llm_slots = llm_slots.clone();
let cfg = config.clone();
tokio::spawn(async move {
loop {
match run_once_with_semaphore(&cfg, llm_slots.clone()).await {
Ok(true) => continue,
Ok(false) => {
tokio::select! {
_ = notify.notified() => {}
_ = tokio::time::sleep(POLL_INTERVAL) => {}
}
}
Err(err) => {
log::error!("[memory_tree::jobs] worker={} loop error: {:#}", idx, err);
tokio::time::sleep(Duration::from_secs(1)).await;
}
}
}
});
}
super::scheduler::start(config);
});
}
pub async fn run_once(config: &Config) -> Result<bool> {
let llm_slots = Arc::new(Semaphore::new(1));
run_once_with_semaphore(config, llm_slots).await
}
async fn run_once_with_semaphore(config: &Config, llm_slots: Arc<Semaphore>) -> Result<bool> {
let Some(job) = claim_next(config, DEFAULT_LOCK_DURATION_MS)? else {
return Ok(false);
};
let permit = if job.kind.is_llm_bound() {
Some(llm_slots.acquire().await?)
} else {
None
};
let result = handlers::handle_job(config, &job).await;
drop(permit);
match result {
Ok(()) => {
mark_done(config, &job)?;
}
Err(err) => {
// Preserve the full anyhow cause chain in the persisted
// last_error so a reader of mem_tree_jobs can see the root
// cause, not just the top-level message. Mirrors the {:#}
// log format used right above.
let message = format!("{err:#}");
log::warn!(
"[memory_tree::jobs] job failed id={} kind={} err={:#}",
job.id,
job.kind.as_str(),
err
);
mark_failed(config, &job, &message)?;
}
}
Ok(true)
}
+3
View File
@@ -24,8 +24,10 @@
pub mod canonicalize;
pub mod chunker;
pub mod content_store;
pub mod global_tree;
pub mod ingest;
pub mod jobs;
pub mod retrieval;
pub mod rpc;
pub mod schemas;
@@ -34,6 +36,7 @@ pub mod source_tree;
pub mod store;
pub mod topic_tree;
pub mod types;
pub mod util;
pub use retrieval::{all_retrieval_controller_schemas, all_retrieval_registered_controllers};
pub use schemas::{
@@ -23,6 +23,7 @@ use std::collections::VecDeque;
use anyhow::Result;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::read as content_read;
use crate::openhuman::memory::tree::retrieval::types::{
hit_from_chunk, hit_from_summary, RetrievalHit,
};
@@ -179,15 +180,27 @@ fn walk_with_embeddings(
continue;
}
// Is it a summary?
if let Some(summary) = store::get_summary(config, &id)? {
if let Some(mut summary) = store::get_summary(config, &id)? {
let scope = store::get_tree(config, &summary.tree_id)?
.map(|t| t.scope)
.unwrap_or_else(|| root_tree_scope.clone());
// Hydrate the full body from disk — `summary.content` is a
// ≤500-char preview after the MD-on-disk migration.
// Non-fatal fallback for pre-MD-migration rows.
match content_read::read_summary_body(config, &id) {
Ok(body) => summary.content = body,
Err(e) => {
log::warn!(
"[retrieval::drill_down] read_summary_body failed — serving preview: {e:#}"
);
}
}
// Summary embeddings live on the struct directly (Phase 4 amend).
embeddings.push(summary.embedding.clone());
let child_ids = summary.child_ids.clone();
out.push(hit_from_summary(&summary, &scope));
if depth < max_depth {
for next in summary.child_ids {
for next in child_ids {
frontier.push_back((next, depth + 1));
}
}
@@ -195,11 +208,21 @@ fn walk_with_embeddings(
}
// Else try as a chunk (leaf). Chunk embeddings live in a separate
// blob column — fetch via the existing accessor.
if let Some(chunk) = get_chunk(config, &id)? {
if let Some(mut chunk) = get_chunk(config, &id)? {
// Propagate DB errors rather than silently treating them as
// "no embedding" — the caller should know if the store is broken.
let emb = get_chunk_embedding(config, &chunk.id)?;
embeddings.push(emb);
// Hydrate the full body from disk — `chunk.content` is a
// ≤500-char preview after the MD-on-disk migration.
match content_read::read_chunk_body(config, &id) {
Ok(body) => chunk.content = body,
Err(e) => {
log::warn!(
"[retrieval::drill_down] read_chunk_body failed — serving preview: {e:#}"
);
}
}
// Score unknown here; 0.0 neutral placeholder.
out.push(hit_from_chunk(&chunk, "", &chunk.metadata.source_id, 0.0));
continue;
@@ -215,7 +238,10 @@ fn walk_with_embeddings(
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
use crate::openhuman::memory::tree::source_tree::types::TreeKind;
@@ -240,6 +266,8 @@ mod tests {
let ts = Utc::now();
let tree = get_or_create_source_tree(cfg, "slack:#eng").unwrap();
let summariser = InertSummariser::new();
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
let mut leaf_ids: Vec<String> = Vec::new();
for seq in 0..2u32 {
let c = Chunk {
@@ -254,18 +282,32 @@ mod tests {
tags: vec![],
source_ref: Some(SourceRef::new("slack://x")),
},
token_count: 6_000,
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET * 6
/ 10,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[c.clone()]).unwrap();
// Stage to disk so `hydrate_leaf_inputs` can read the full body
// via `read_chunk_body` during the seal triggered by `append_leaf`.
let staged = content_store::stage_chunks(&content_root, &[c.clone()]).unwrap();
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
leaf_ids.push(c.id.clone());
append_leaf(
cfg,
&tree,
&LeafRef {
chunk_id: c.id.clone(),
token_count: 6_000,
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET
* 6
/ 10,
timestamp: ts,
content: c.content.clone(),
entities: vec![],
@@ -273,6 +315,7 @@ mod tests {
score: 0.5,
},
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
@@ -401,6 +444,7 @@ mod tests {
token_count: 10,
seq_in_source: 0,
created_at: ts,
partial_message: false,
};
let leaf_a_2 = Chunk {
id: "chat:slack:#eng:1".into(),
@@ -423,15 +467,24 @@ mod tests {
seq_in_source: 3,
..leaf_a_1.clone()
};
upsert_chunks(
cfg,
&[
leaf_a_1.clone(),
leaf_a_2.clone(),
leaf_b_1.clone(),
leaf_b_2.clone(),
],
)
let all_leaves = [
leaf_a_1.clone(),
leaf_a_2.clone(),
leaf_b_1.clone(),
leaf_b_2.clone(),
];
upsert_chunks(cfg, &all_leaves).unwrap();
// Stage to disk so `walk_with_embeddings` can read full bodies via
// `read_chunk_body` for leaf hits returned by the drill-down.
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
let staged = content_store::stage_chunks(&content_root, &all_leaves).unwrap();
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
let l1_a = SummaryNode {
@@ -471,9 +524,9 @@ mod tests {
with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
tree_store::insert_tree_conn(&tx, &tree)?;
tree_store::insert_summary_tx(&tx, &l1_a)?;
tree_store::insert_summary_tx(&tx, &l1_b)?;
tree_store::insert_summary_tx(&tx, &root)?;
tree_store::insert_summary_tx(&tx, &l1_a, None)?;
tree_store::insert_summary_tx(&tx, &l1_b, None)?;
tree_store::insert_summary_tx(&tx, &root, None)?;
tx.commit()?;
Ok(())
})
+37 -1
View File
@@ -13,6 +13,7 @@
use anyhow::Result;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::read as content_read;
use crate::openhuman::memory::tree::retrieval::types::{hit_from_chunk, RetrievalHit};
use crate::openhuman::memory::tree::score::store::get_score;
use crate::openhuman::memory::tree::store::get_chunk;
@@ -66,7 +67,22 @@ pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result<Vec<R
// chunk row. `scope` falls back to the chunk's own source_id so
// consumers still see provenance (e.g. "slack:#eng").
let scope = chunk.metadata.source_id.clone();
out.push(hit_from_chunk(&chunk, "", &scope, score));
// Hydrate the full body from disk before building the hit.
// The `content` column in SQLite holds a ≤500-char preview after
// the MD-on-disk migration; the retrieval API must return the
// complete chunk text so the LLM sees untruncated content.
let mut chunk_with_body = chunk;
match content_read::read_chunk_body(&config_owned, id) {
Ok(body) => chunk_with_body.content = body,
Err(e) => {
log::warn!(
"[retrieval::fetch] read_chunk_body failed for chunk — serving preview: {e:#}"
);
// Non-fatal: fall back to the preview already in the struct.
// This handles pre-MD-migration rows gracefully.
}
}
out.push(hit_from_chunk(&chunk_with_body, "", &scope, score));
}
Ok(out)
})
@@ -80,11 +96,26 @@ pub async fn fetch_leaves(config: &Config, chunk_ids: &[String]) -> Result<Vec<R
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
use chrono::{TimeZone, Utc};
use tempfile::TempDir;
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged = content_store::stage_chunks(&content_root, chunks)
.expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -113,6 +144,7 @@ mod tests {
token_count: 20,
seq_in_source: seq,
created_at: ts,
partial_message: false,
}
}
@@ -129,6 +161,7 @@ mod tests {
let c1 = sample_chunk("slack:#eng", 0);
let c2 = sample_chunk("slack:#eng", 1);
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]);
let out = fetch_leaves(&cfg, &[c1.id.clone(), c2.id.clone()])
.await
.unwrap();
@@ -142,6 +175,7 @@ mod tests {
let (_tmp, cfg) = test_config();
let c1 = sample_chunk("slack:#eng", 0);
upsert_chunks(&cfg, &[c1.clone()]).unwrap();
stage_test_chunks(&cfg, &[c1.clone()]);
let out = fetch_leaves(
&cfg,
&[c1.id.clone(), "ghost:nonexistent".into(), c1.id.clone()],
@@ -159,6 +193,7 @@ mod tests {
for i in 0..(MAX_BATCH + 5) as u32 {
let c = sample_chunk("slack:#eng", i);
upsert_chunks(&cfg, &[c.clone()]).unwrap();
stage_test_chunks(&cfg, &[c.clone()]);
ids.push(c.id);
}
let out = fetch_leaves(&cfg, &ids).await.unwrap();
@@ -170,6 +205,7 @@ mod tests {
let (_tmp, cfg) = test_config();
let c = sample_chunk("slack:#eng", 0);
upsert_chunks(&cfg, &[c.clone()]).unwrap();
stage_test_chunks(&cfg, &[c.clone()]);
let out = fetch_leaves(&cfg, &[c.id.clone()]).await.unwrap();
assert_eq!(out.len(), 1);
assert_eq!(out[0].source_ref.as_deref(), Some("slack://slack:#eng/0"));
+21 -1
View File
@@ -87,8 +87,11 @@ fn recap_to_hits(recap: RecapOutput, tree_id: &str, tree_scope: &str) -> Vec<Ret
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::global_tree::digest::{end_of_day_digest, DigestOutcome};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
use crate::openhuman::memory::tree::store::upsert_chunks;
@@ -96,6 +99,20 @@ mod tests {
use chrono::{DateTime, Utc};
use tempfile::TempDir;
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged = content_store::stage_chunks(&content_root, chunks)
.expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -134,8 +151,10 @@ mod tests {
token_count: 6_000,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[c.clone()]).unwrap();
stage_test_chunks(cfg, &[c.clone()]);
append_leaf(
cfg,
&tree,
@@ -149,6 +168,7 @@ mod tests {
score: 0.5,
},
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
@@ -160,11 +160,11 @@ async fn topic_entity_surfaces_after_ingest() {
// ── Phase 4 (#710): embedding + semantic rerank tests ───────────────────
/// Ingest with an inert embedder must populate every kept chunk's
/// `embedding` column. This guards against regressions where the embed
/// step is silently skipped (e.g. future refactors threading embeddings
/// through a different code path).
/// `embedding` column. Embeddings are written by the async `extract_chunk`
/// handler, so the test drains the queue before inspecting.
#[tokio::test]
async fn ingest_populates_chunk_embeddings() {
use crate::openhuman::memory::tree::jobs::drain_until_idle;
use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM;
use crate::openhuman::memory::tree::store::get_chunk_embedding;
@@ -172,7 +172,11 @@ async fn ingest_populates_chunk_embeddings() {
let out = ingest_chat(&cfg, "slack:#eng", "alice", vec![], chat_about_phoenix(0))
.await
.unwrap();
assert!(out.chunks_written >= 1, "expected at least one kept chunk");
assert!(
out.chunks_written >= 1,
"expected at least one persisted chunk"
);
drain_until_idle(&cfg).await.unwrap();
for id in &out.chunk_ids {
let emb = get_chunk_embedding(&cfg, id).unwrap();
let v = emb.unwrap_or_else(|| panic!("embedding missing for chunk_id={id}"));
@@ -188,8 +192,11 @@ async fn ingest_populates_chunk_embeddings() {
/// the seal from firing on short batches.
#[tokio::test]
async fn seal_populates_summary_embedding() {
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::score::embed::EMBEDDING_DIM;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::store as src_store;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
@@ -216,10 +223,24 @@ async fn seal_populates_summary_embedding() {
token_count: tokens,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
let c1 = mk_chunk(0, 6_000);
let c2 = mk_chunk(1, 6_000);
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
{
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged = content_store::stage_chunks(&content_root, &[c1.clone(), c2.clone()])
.expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
let leaf_of = |c: &Chunk| LeafRef {
chunk_id: c.id.clone(),
@@ -230,12 +251,24 @@ async fn seal_populates_summary_embedding() {
topics: vec![],
score: 0.5,
};
append_leaf(&cfg, &tree, &leaf_of(&c1), &summariser)
.await
.unwrap();
let sealed = append_leaf(&cfg, &tree, &leaf_of(&c2), &summariser)
.await
.unwrap();
append_leaf(
&cfg,
&tree,
&leaf_of(&c1),
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
let sealed = append_leaf(
&cfg,
&tree,
&leaf_of(&c2),
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
assert_eq!(sealed.len(), 1, "expected one seal at the budget crossing");
let summary = src_store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
@@ -282,11 +282,26 @@ mod tests {
//! initialises the schema idempotently on first access, so read-only
//! calls return empty responses rather than erroring.
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceRef};
use chrono::{TimeZone, Utc};
use tempfile::TempDir;
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged = content_store::stage_chunks(&content_root, chunks)
.expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -316,6 +331,7 @@ mod tests {
token_count: 20,
seq_in_source: seq,
created_at: ts,
partial_message: false,
}
}
@@ -539,6 +555,7 @@ mod tests {
let c1 = sample_chunk("slack:#eng", 0);
let c2 = sample_chunk("slack:#eng", 1);
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]);
let req = FetchLeavesRequest {
chunk_ids: vec![c1.id.clone(), c2.id.clone()],
};
@@ -552,6 +569,7 @@ mod tests {
let (_tmp, cfg) = test_config();
let c1 = sample_chunk("slack:#eng", 0);
upsert_chunks(&cfg, &[c1.clone()]).unwrap();
stage_test_chunks(&cfg, &[c1.clone()]);
let req = FetchLeavesRequest {
chunk_ids: vec![c1.id.clone(), "ghost:nonexistent".into()],
};
+38 -4
View File
@@ -21,6 +21,7 @@ use anyhow::Result;
use chrono::{Duration, Utc};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::read as content_read;
use crate::openhuman::memory::tree::retrieval::types::{
hit_from_summary, QueryResponse, RetrievalHit,
};
@@ -124,7 +125,19 @@ fn collect_hits_and_nodes(
// finest-grained summary layer above raw leaves.
for level in 1..=tree.max_level {
let level_nodes = store::list_summaries_at_level(config, &tree.id, level)?;
for node in level_nodes {
for mut node in level_nodes {
// Hydrate the full body from disk — `node.content` is a
// ≤500-char preview after the MD-on-disk migration. Callers
// (including the LLM) must receive the complete summary text.
// Non-fatal fallback for pre-MD-migration rows.
match content_read::read_summary_body(config, &node.id) {
Ok(body) => node.content = body,
Err(e) => {
log::warn!(
"[retrieval::source] read_summary_body failed — serving preview: {e:#}"
);
}
}
hits.push(hit_from_summary(&node, &tree.scope));
nodes.push((node, tree.scope.clone()));
}
@@ -293,7 +306,10 @@ fn filter_by_window(hits: Vec<RetrievalHit>, window_days: u32) -> Vec<RetrievalH
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
use crate::openhuman::memory::tree::store::upsert_chunks;
@@ -315,6 +331,8 @@ mod tests {
async fn seed_source(cfg: &Config, scope: &str, ts: DateTime<Utc>) {
let tree = get_or_create_source_tree(cfg, scope).unwrap();
let summariser = InertSummariser::new();
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).unwrap();
for seq in 0..2u32 {
let c = Chunk {
id: chunk_id(SourceKind::Chat, scope, seq, "test-content"),
@@ -328,17 +346,32 @@ mod tests {
tags: vec!["eng".into()],
source_ref: Some(SourceRef::new(format!("slack://{scope}/{seq}"))),
},
token_count: 6_000,
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET * 6
/ 10,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[c.clone()]).unwrap();
// Stage to disk so `hydrate_leaf_inputs` can read the full body
// via `read_chunk_body` during the seal triggered by `append_leaf`,
// and `collect_hits_and_nodes` can read summary bodies for the API.
let staged = content_store::stage_chunks(&content_root, &[c.clone()]).unwrap();
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.unwrap();
append_leaf(
cfg,
&tree,
&LeafRef {
chunk_id: c.id.clone(),
token_count: 6_000,
token_count: crate::openhuman::memory::tree::source_tree::types::TOKEN_BUDGET
* 6
/ 10,
timestamp: ts,
content: c.content.clone(),
entities: vec![],
@@ -346,6 +379,7 @@ mod tests {
score: 0.5,
},
&summariser,
&LabelStrategy::Empty,
)
.await
.unwrap();
+34 -4
View File
@@ -17,6 +17,7 @@ use anyhow::Result;
use chrono::{Duration, TimeZone, Utc};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::read as content_read;
use crate::openhuman::memory::tree::retrieval::types::{
hit_from_summary, QueryResponse, RetrievalHit,
};
@@ -229,7 +230,7 @@ fn fetch_topic_tree_root_summary(config: &Config, entity_id: &str) -> Result<Opt
Some(id) => id.clone(),
None => return Ok(None),
};
let summary = match store::get_summary(config, &root_id)? {
let mut summary = match store::get_summary(config, &root_id)? {
Some(s) => s,
None => {
log::warn!(
@@ -238,6 +239,17 @@ fn fetch_topic_tree_root_summary(config: &Config, entity_id: &str) -> Result<Opt
return Ok(None);
}
};
// Hydrate the full body from disk — `summary.content` is a ≤500-char
// preview after the MD-on-disk migration. Non-fatal fallback for
// pre-MD-migration rows.
match content_read::read_summary_body(config, &root_id) {
Ok(body) => summary.content = body,
Err(e) => {
log::warn!(
"[retrieval::topic] read_summary_body failed for topic root — serving preview: {e:#}"
);
}
}
Ok(Some(hit_from_summary(&summary, &tree.scope)))
}
@@ -258,13 +270,23 @@ async fn entity_hit_to_retrieval_hit(
tokio::task::spawn_blocking(move || -> Result<Option<RetrievalHit>> {
if node_kind == "summary" {
let summary = match store::get_summary(&config_owned, &node_id)? {
let mut summary = match store::get_summary(&config_owned, &node_id)? {
Some(s) => s,
None => {
log::warn!("[retrieval::topic] entity index points at missing summary row");
return Ok(None);
}
};
// Hydrate the full body from disk — `summary.content` is a
// ≤500-char preview after the MD-on-disk migration.
match content_read::read_summary_body(&config_owned, &node_id) {
Ok(body) => summary.content = body,
Err(e) => {
log::warn!(
"[retrieval::topic] read_summary_body failed — serving preview: {e:#}"
);
}
}
// Prefer tree scope from the summary's parent tree if resolvable.
let scope = if let Some(tid) = &tree_id_opt {
store::get_tree(&config_owned, tid)?
@@ -283,13 +305,21 @@ async fn entity_hit_to_retrieval_hit(
// Leaf: fetch chunk and hydrate.
use crate::openhuman::memory::tree::retrieval::types::hit_from_chunk;
use crate::openhuman::memory::tree::store::get_chunk;
let chunk = match get_chunk(&config_owned, &node_id)? {
let mut chunk = match get_chunk(&config_owned, &node_id)? {
Some(c) => c,
None => {
log::warn!("[retrieval::topic] entity index points at missing chunk row");
return Ok(None);
}
};
// Hydrate the full body from disk — `chunk.content` is a ≤500-char
// preview after the MD-on-disk migration.
match content_read::read_chunk_body(&config_owned, &node_id) {
Ok(body) => chunk.content = body,
Err(e) => {
log::warn!("[retrieval::topic] read_chunk_body failed — serving preview: {e:#}");
}
}
let scope = if let Some(tid) = &tree_id_opt {
store::get_tree(&config_owned, tid)?
.map(|t: Tree| t.scope)
@@ -545,7 +575,7 @@ mod tests {
with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
tree_store::insert_tree_conn(&tx, &tree)?;
tree_store::insert_summary_tx(&tx, &summary)?;
tree_store::insert_summary_tx(&tx, &summary, None)?;
score_store::index_entities_tx(
&tx,
&[entity],
+140
View File
@@ -178,3 +178,143 @@ pub async fn get_chunk_rpc(
format!("memory_tree: get_chunk id={}", req.id),
))
}
/// Manual-trigger surface for the global tree's daily digest. Default
/// behavior (no `date_iso`) targets yesterday in UTC, matching the
/// scheduler's autonomous behavior. Pass an explicit `YYYY-MM-DD` to
/// re-run a specific date (idempotent — the handler skips if a daily
/// node already exists for that day).
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct TriggerDigestRequest {
/// UTC calendar date in `YYYY-MM-DD` form. When omitted, defaults to
/// `yesterday` (today minus one day, UTC).
#[serde(default)]
pub date_iso: Option<String>,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct TriggerDigestResponse {
/// True when the job was newly enqueued; false when an active job for
/// the same date was suppressed by the dedupe partial unique index.
pub enqueued: bool,
/// ID of the freshly-inserted job row (None when dedupe-suppressed).
pub job_id: Option<String>,
/// The actual date the digest will run for, echoed back as
/// `YYYY-MM-DD`. Useful when the caller didn't pass `date_iso` and
/// wants to know what default got chosen.
pub date_iso: String,
}
pub async fn trigger_digest_rpc(
config: &Config,
req: TriggerDigestRequest,
) -> Result<RpcOutcome<TriggerDigestResponse>, String> {
use crate::openhuman::memory::tree::jobs;
use chrono::{Duration as ChronoDuration, NaiveDate, Utc};
let date = match req
.date_iso
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty())
{
Some(s) => NaiveDate::parse_from_str(s, "%Y-%m-%d")
.map_err(|e| format!("invalid date_iso (expected YYYY-MM-DD): {e}"))?,
None => Utc::now().date_naive() - ChronoDuration::days(1),
};
let date_iso = date.format("%Y-%m-%d").to_string();
// Run the synchronous enqueue on a blocking thread — `trigger_digest`
// touches SQLite and we don't want to block the async runtime even
// for the few-microsecond INSERT.
let cfg_clone = config.clone();
let date_for_blocking = date;
let job_id =
tokio::task::spawn_blocking(move || jobs::trigger_digest(&cfg_clone, date_for_blocking))
.await
.map_err(|e| format!("trigger_digest join error: {e}"))?
.map_err(|e| format!("trigger_digest: {e}"))?;
let enqueued = job_id.is_some();
Ok(RpcOutcome::single_log(
TriggerDigestResponse {
enqueued,
job_id,
date_iso: date_iso.clone(),
},
format!("memory_tree: trigger_digest date={date_iso} enqueued={enqueued}"),
))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::jobs::store::count_total;
use chrono::{Duration as ChronoDuration, Utc};
use tempfile::TempDir;
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
cfg.workspace_dir = tmp.path().to_path_buf();
cfg.memory_tree.embedding_endpoint = None;
cfg.memory_tree.embedding_model = None;
cfg.memory_tree.embedding_strict = false;
(tmp, cfg)
}
#[tokio::test]
async fn trigger_digest_with_explicit_date_enqueues() {
let (_tmp, cfg) = test_config();
let req = TriggerDigestRequest {
date_iso: Some("2026-04-27".into()),
};
let outcome = trigger_digest_rpc(&cfg, req).await.unwrap();
let resp = outcome.value;
assert!(resp.enqueued);
assert!(resp.job_id.is_some());
assert_eq!(resp.date_iso, "2026-04-27");
assert_eq!(count_total(&cfg).unwrap(), 1);
}
#[tokio::test]
async fn trigger_digest_with_no_date_defaults_to_yesterday() {
let (_tmp, cfg) = test_config();
let req = TriggerDigestRequest::default();
let outcome = trigger_digest_rpc(&cfg, req).await.unwrap();
let resp = outcome.value;
assert!(resp.enqueued);
let expected = (Utc::now().date_naive() - ChronoDuration::days(1))
.format("%Y-%m-%d")
.to_string();
assert_eq!(resp.date_iso, expected);
}
#[tokio::test]
async fn trigger_digest_rejects_malformed_date() {
let (_tmp, cfg) = test_config();
let req = TriggerDigestRequest {
date_iso: Some("not-a-date".into()),
};
let err = trigger_digest_rpc(&cfg, req).await.unwrap_err();
assert!(
err.contains("invalid date_iso"),
"expected schema-shaped error message, got: {err}"
);
assert_eq!(count_total(&cfg).unwrap(), 0);
}
#[tokio::test]
async fn trigger_digest_dedupes_active_jobs() {
let (_tmp, cfg) = test_config();
let req = TriggerDigestRequest {
date_iso: Some("2026-04-27".into()),
};
let first = trigger_digest_rpc(&cfg, req.clone()).await.unwrap().value;
let second = trigger_digest_rpc(&cfg, req).await.unwrap().value;
assert!(first.enqueued);
assert!(!second.enqueued, "duplicate must be dedupe-suppressed");
assert!(second.job_id.is_none());
assert_eq!(count_total(&cfg).unwrap(), 1);
}
}
+51
View File
@@ -23,6 +23,7 @@ pub fn all_controller_schemas() -> Vec<ControllerSchema> {
schemas("ingest"),
schemas("list_chunks"),
schemas("get_chunk"),
schemas("trigger_digest"),
]
}
@@ -40,6 +41,10 @@ pub fn all_registered_controllers() -> Vec<RegisteredController> {
schema: schemas("get_chunk"),
handler: handle_get_chunk,
},
RegisteredController {
schema: schemas("trigger_digest"),
handler: handle_trigger_digest,
},
]
}
@@ -184,6 +189,44 @@ pub fn schemas(function: &str) -> ControllerSchema {
required: false,
}],
},
"trigger_digest" => ControllerSchema {
namespace: NAMESPACE,
function: "trigger_digest",
description: "Manually enqueue a daily-digest job for the global \
tree. Idempotent — re-running for a day that already has a \
digest is a no-op (the handler skips). When no date is \
supplied, defaults to yesterday in UTC, matching the \
scheduler's autonomous behavior.",
inputs: vec![FieldSchema {
name: "date_iso",
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
comment: "UTC calendar date as `YYYY-MM-DD`. Optional; \
defaults to yesterday when omitted.",
required: false,
}],
outputs: vec![
FieldSchema {
name: "enqueued",
ty: TypeSchema::Bool,
comment: "True when a fresh job row was inserted; false \
when an active job for the same date suppressed it.",
required: true,
},
FieldSchema {
name: "job_id",
ty: TypeSchema::Option(Box::new(TypeSchema::String)),
comment: "ID of the newly enqueued job row, when enqueued.",
required: false,
},
FieldSchema {
name: "date_iso",
ty: TypeSchema::String,
comment: "The date the digest will cover, echoed back \
as `YYYY-MM-DD`.",
required: true,
},
],
},
_ => ControllerSchema {
namespace: NAMESPACE,
function: "unknown",
@@ -228,6 +271,14 @@ fn handle_get_chunk(params: Map<String, Value>) -> ControllerFuture {
})
}
fn handle_trigger_digest(params: Map<String, Value>) -> ControllerFuture {
Box::pin(async move {
let config = config_rpc::load_config_with_timeout().await?;
let req = parse_value::<tree_rpc::TriggerDigestRequest>(Value::Object(params))?;
to_json(tree_rpc::trigger_digest_rpc(&config, req).await?)
})
}
fn parse_value<T: DeserializeOwned>(v: Value) -> Result<T, String> {
serde_json::from_value(v).map_err(|e| format!("invalid params: {e}"))
}
+164 -29
View File
@@ -36,7 +36,7 @@ use async_trait::async_trait;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use super::types::{EntityKind, ExtractedEntities, ExtractedEntity};
use super::types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic};
use super::EntityExtractor;
// ── Configuration ────────────────────────────────────────────────────────
@@ -59,6 +59,15 @@ pub struct LlmExtractorConfig {
/// If true, drop entities whose declared kind isn't in `allowed_kinds`
/// instead of falling back to [`EntityKind::Misc`].
pub strict_kinds: bool,
/// If true, the system prompt asks the model to also emit a
/// `topics` array (free-form theme labels), and the response parser
/// populates [`ExtractedEntities::topics`]. Default `false` — the
/// extractor's primary job is named-entity extraction; topics are
/// an opt-in side-channel for callers that need a thematic
/// summary in the same call (e.g. running over a sealed summary's
/// content). Adds prompt tokens and gives the model one more
/// schema field to keep track of, so leave off unless needed.
pub emit_topics: bool,
}
impl Default for LlmExtractorConfig {
@@ -73,8 +82,13 @@ impl Default for LlmExtractorConfig {
EntityKind::Location,
EntityKind::Event,
EntityKind::Product,
EntityKind::Datetime,
EntityKind::Technology,
EntityKind::Artifact,
EntityKind::Quantity,
],
strict_kinds: false,
emit_topics: false,
}
}
}
@@ -103,7 +117,7 @@ impl LlmEntityExtractor {
messages: vec![
OllamaMessage {
role: "system".to_string(),
content: SYSTEM_PROMPT.to_string(),
content: build_system_prompt(self.cfg.emit_topics),
},
OllamaMessage {
role: "user".to_string(),
@@ -128,19 +142,54 @@ impl EntityExtractor for LlmEntityExtractor {
// JSON parse) is logged as a warn and returns an empty
// `ExtractedEntities` rather than `Err`. This makes the extractor
// safe to call from any context, not just `score_chunk` (which
// separately catches errors from its own extractor chain). A caller
// distinguishes "LLM had nothing to say" from "LLM ran and returned
// zero entities" by inspecting `llm_importance` — `None` means the
// call didn't complete successfully.
Ok(self.extract_or_empty(text).await)
// separately catches errors from its own extractor chain).
//
// Transport failures get bounded retry-with-backoff before falling
// back to empty — see [`Self::try_extract`]. Non-transport failures
// (HTTP non-success, malformed JSON) fall back immediately because
// retrying the same input would yield the same bad response.
const MAX_ATTEMPTS: u32 = 3;
const BASE_BACKOFF_MS: u64 = 250;
for attempt in 0..MAX_ATTEMPTS {
match self.try_extract(text).await {
Some(extracted) => return Ok(extracted),
None => {
// Transport failure. Retry with exponential backoff
// unless we've exhausted attempts.
if attempt + 1 < MAX_ATTEMPTS {
let delay_ms = BASE_BACKOFF_MS * 2u64.pow(attempt);
log::warn!(
"[memory_tree::extract::llm] transport failure, retrying in \
{delay_ms}ms (attempt {}/{})",
attempt + 2,
MAX_ATTEMPTS
);
tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
}
}
}
}
log::warn!(
"[memory_tree::extract::llm] transport failed after {} attempts — \
returning empty extraction",
MAX_ATTEMPTS
);
Ok(ExtractedEntities::default())
}
}
impl LlmEntityExtractor {
/// Internal: wraps the actual HTTP call and returns `ExtractedEntities`
/// for every failure mode via soft-fallback. Split out of `extract` so
/// the error branches can share logging without `?`-propagation.
async fn extract_or_empty(&self, text: &str) -> ExtractedEntities {
/// Internal: one attempt at calling Ollama.
///
/// Returns:
/// - `Some(extracted)` — call completed (HTTP returned). Includes the
/// "HTTP non-success" and "malformed JSON" cases, which return
/// `Some(empty)` because retrying the same input won't help.
/// - `None` — transport-level failure (DNS, connect refused, timeout
/// before any HTTP response). Caller may retry.
async fn try_extract(&self, text: &str) -> Option<ExtractedEntities> {
let url = format!("{}/api/chat", self.cfg.endpoint.trim_end_matches('/'));
let body = self.build_request(text);
log::debug!(
@@ -152,11 +201,8 @@ impl LlmEntityExtractor {
let resp = match self.http.post(&url).json(&body).send().await {
Ok(r) => r,
Err(e) => {
log::warn!(
"[memory_tree::extract::llm] transport failure to {url}: {e} — \
returning empty extraction"
);
return ExtractedEntities::default();
log::warn!("[memory_tree::extract::llm] transport failure to {url}: {e}");
return None;
}
};
@@ -168,7 +214,7 @@ impl LlmEntityExtractor {
returning empty extraction",
truncate_for_log(&body, 200)
);
return ExtractedEntities::default();
return Some(ExtractedEntities::default());
}
let envelope: OllamaChatResponse = match resp.json().await {
@@ -178,7 +224,7 @@ impl LlmEntityExtractor {
"[memory_tree::extract::llm] response body not Ollama-shaped JSON: {e} — \
returning empty extraction"
);
return ExtractedEntities::default();
return Some(ExtractedEntities::default());
}
};
log::debug!(
@@ -194,38 +240,97 @@ impl LlmEntityExtractor {
response: {e}; content was: {} — returning empty extraction",
truncate_for_log(&envelope.message.content, 400)
);
return ExtractedEntities::default();
return Some(ExtractedEntities::default());
}
};
parsed.into_extracted_entities(text, &self.cfg)
Some(parsed.into_extracted_entities(text, &self.cfg))
}
}
// ── Prompt ───────────────────────────────────────────────────────────────
const SYSTEM_PROMPT: &str = "\
You are a named-entity extractor and importance rater. Return JSON only — \
/// Build the system prompt for the extractor. When `emit_topics` is true
/// the schema, required-fields list, and example outputs include a
/// `topics` array (free-form theme labels). When false the prompt
/// matches the pre-flag behaviour exactly — no mention of topics
/// anywhere — so the small model isn't asked to produce a field the
/// caller doesn't want.
fn build_system_prompt(emit_topics: bool) -> String {
let topics_schema_line = if emit_topics {
" \"topics\": [\"<short theme label>\"],\n"
} else {
""
};
let topics_required = if emit_topics { "topics, " } else { "" };
let fields_count = if emit_topics { "four" } else { "three" };
let topics_guide = if emit_topics {
"Topics are short free-form theme labels for what the text is ABOUT \
(e.g. \"rate limiting\", \"memory tree\", \"auth flow\"). They are \
distinct from entities — entities are specific named things mentioned \
in the text; topics are the abstract themes those things relate to.\n"
} else {
""
};
let example1_topics = if emit_topics {
",\"topics\":[\"shipping\",\"auth\"]"
} else {
""
};
let example2_topics = if emit_topics {
",\"topics\":[\"product launch\",\"revenue\"]"
} else {
""
};
format!(
"You are a named-entity extractor and importance rater. Return JSON only — \
no prose, no markdown, no commentary. Do not summarize. Extract every named \
entity mention you find, including duplicates, and rate the chunk's overall \
importance as a float in [0.0, 1.0].
Schema:
{
{{
\"entities\": [
{ \"kind\": \"person|organization|location|event|product\",
\"text\": \"<exact surface form as it appears in the text>\" }
{{ \"kind\": \"person|organization|location|event|product|datetime|technology|artifact|quantity\",
\"text\": \"<exact surface form as it appears in the text>\" }}
],
\"importance\": 0.0,
{topics_schema_line} \"importance\": 0.0,
\"importance_reason\": \"<one short sentence explaining the rating>\"
}
}}
Kinds guide:
person named human (\"Alice\", \"Steven Enamakel\")
organization company / team / project (\"Anthropic\", \"TinyHumans\")
location place (\"SF office\", \"London\")
event scheduled occurrence (\"Q2 launch\", \"design review\")
product commercial offering (\"Claude Code\", \"OpenHuman\")
datetime temporal expression (\"Friday\", \"Q2 2026\", \"EOD tomorrow\")
technology tool / framework / language / service (\"Rust\", \"OAuth\", \"Slack API\")
artifact code / ticket / doc reference (\"PR #934\", \"src/foo.rs\", \"OH-42\")
quantity amount / metric / money (\"$5K\", \"20/min\", \"10k tokens\")
{topics_guide}
If a mention doesn't clearly fit a kind above, omit it rather than guessing.
Always emit ALL {fields_count} top-level fields (entities, {topics_required}importance, importance_reason),
even when entities is empty.
Examples:
Input: alice and bob shipped the auth migration friday. PR #42 ships OAuth refactor in src/auth/.
Output: {{\"entities\":[{{\"kind\":\"person\",\"text\":\"alice\"}},{{\"kind\":\"person\",\"text\":\"bob\"}},{{\"kind\":\"event\",\"text\":\"auth migration\"}},{{\"kind\":\"datetime\",\"text\":\"friday\"}},{{\"kind\":\"artifact\",\"text\":\"PR #42\"}},{{\"kind\":\"technology\",\"text\":\"OAuth\"}},{{\"kind\":\"artifact\",\"text\":\"src/auth/\"}}]{example1_topics},\"importance\":0.9,\"importance_reason\":\"explicit shipping commitment\"}}
Input: Anthropic shipped Claude Code in SF — $20M ARR target by Q2.
Output: {{\"entities\":[{{\"kind\":\"organization\",\"text\":\"Anthropic\"}},{{\"kind\":\"product\",\"text\":\"Claude Code\"}},{{\"kind\":\"location\",\"text\":\"SF\"}},{{\"kind\":\"quantity\",\"text\":\"$20M ARR\"}},{{\"kind\":\"datetime\",\"text\":\"Q2\"}}]{example2_topics},\"importance\":0.85,\"importance_reason\":\"factual content with key business metric\"}}
Importance guide:
0.9+ actionable decisions, key information, explicit commitments
0.6+ substantive discussion, factual content, named entities
0.3+ ambient context, low-density prose
<0.3 reactions, acknowledgments, bots, trivial exchanges
";
"
)
}
// ── Wire types (Ollama API) ──────────────────────────────────────────────
@@ -265,6 +370,11 @@ struct OllamaResponseMessage {
struct LlmExtractionOutput {
#[serde(default)]
entities: Vec<LlmEntity>,
/// Free-form theme labels — populated only when the extractor is
/// configured with `emit_topics = true`. Always tolerant of absence
/// so models that ignore the field don't fail parsing.
#[serde(default)]
topics: Vec<String>,
#[serde(default)]
importance: Option<f32>,
#[serde(default)]
@@ -355,9 +465,26 @@ impl LlmExtractionOutput {
let llm_importance = self.importance.map(|v| v.clamp(0.0, 1.0));
// Topics: only populated when the caller enabled `emit_topics`
// (the prompt asked for them). Otherwise this is empty by
// default — the model didn't know to emit topics, so any value
// here would be hallucination.
let topics = self
.topics
.into_iter()
.filter_map(|raw| {
let label = raw.trim().to_string();
if label.is_empty() {
None
} else {
Some(ExtractedTopic { label, score: 0.85 })
}
})
.collect();
ExtractedEntities {
entities,
topics: Vec::new(),
topics,
llm_importance,
llm_importance_reason: self.importance_reason,
}
@@ -373,6 +500,14 @@ fn parse_kind(s: &str) -> Option<EntityKind> {
"location" | "place" | "loc" => Some(EntityKind::Location),
"event" => Some(EntityKind::Event),
"product" => Some(EntityKind::Product),
"datetime" | "date" | "time" | "timestamp" => Some(EntityKind::Datetime),
"technology" | "tech" | "tool" | "framework" | "library" | "language" | "service" => {
Some(EntityKind::Technology)
}
"artifact" | "reference" | "ref" | "pr" | "ticket" | "file" | "commit" => {
Some(EntityKind::Artifact)
}
"quantity" | "amount" | "metric" | "number" | "money" => Some(EntityKind::Quantity),
"misc" | "miscellaneous" | "other" => Some(EntityKind::Misc),
_ => None,
}
@@ -1,5 +1,39 @@
use super::*;
#[test]
fn build_system_prompt_default_omits_topics() {
let p = build_system_prompt(false);
assert!(!p.contains("\"topics\""));
assert!(!p.contains("Topics are"));
assert!(p.contains("ALL three top-level fields"));
assert!(p.contains("entities, importance"));
}
#[test]
fn build_system_prompt_with_flag_includes_topics() {
let p = build_system_prompt(true);
assert!(p.contains("\"topics\""));
assert!(p.contains("Topics are short free-form theme labels"));
assert!(p.contains("ALL four top-level fields"));
assert!(p.contains("entities, topics, importance"));
}
#[test]
fn extraction_output_parses_topics_when_present() {
let json = r#"{"entities":[],"topics":["rate limiting","memory tree"],"importance":0.6,"importance_reason":"r"}"#;
let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap();
assert_eq!(parsed.topics, vec!["rate limiting", "memory tree"]);
}
#[test]
fn extraction_output_tolerates_missing_topics() {
// Default extractor (emit_topics=false) — model won't emit topics
// and parsing must still succeed.
let json = r#"{"entities":[],"importance":0.6,"importance_reason":"r"}"#;
let parsed: LlmExtractionOutput = serde_json::from_str(json).unwrap();
assert!(parsed.topics.is_empty());
}
#[test]
fn parse_kind_normalisation() {
assert_eq!(parse_kind("Person"), Some(EntityKind::Person));
@@ -8,6 +42,42 @@ fn parse_kind_normalisation() {
assert!(parse_kind("Spaceship").is_none());
}
#[test]
fn parse_kind_accepts_new_semantic_kinds_and_synonyms() {
// Datetime
for s in ["datetime", "date", "time", "timestamp", " DateTime "] {
assert_eq!(parse_kind(s), Some(EntityKind::Datetime), "input={s:?}");
}
// Technology
for s in [
"technology",
"tech",
"tool",
"framework",
"library",
"language",
"service",
] {
assert_eq!(parse_kind(s), Some(EntityKind::Technology), "input={s:?}");
}
// Artifact
for s in [
"artifact",
"reference",
"ref",
"pr",
"ticket",
"file",
"commit",
] {
assert_eq!(parse_kind(s), Some(EntityKind::Artifact), "input={s:?}");
}
// Quantity
for s in ["quantity", "amount", "metric", "number", "money"] {
assert_eq!(parse_kind(s), Some(EntityKind::Quantity), "input={s:?}");
}
}
#[test]
fn find_char_span_handles_unicode() {
let text = "中 Alice met Bob";
@@ -71,6 +141,7 @@ fn into_extracted_entities_gives_distinct_spans_to_duplicate_mentions() {
text: "Alice".into(),
},
],
topics: vec![],
importance: None,
importance_reason: None,
};
@@ -100,6 +171,7 @@ fn into_extracted_entities_drops_extra_duplicate_when_source_only_has_one() {
text: "Alice".into(),
},
],
topics: vec![],
importance: None,
importance_reason: None,
};
@@ -138,6 +210,7 @@ fn into_extracted_entities_drops_hallucinations() {
text: "ImaginaryPerson".into(),
},
],
topics: vec![],
importance: Some(0.7),
importance_reason: Some("substantive".into()),
};
@@ -154,6 +227,7 @@ fn into_extracted_entities_drops_hallucinations() {
fn into_extracted_entities_clamps_importance() {
let out = LlmExtractionOutput {
entities: vec![],
topics: vec![],
importance: Some(1.5),
importance_reason: None,
};
@@ -169,6 +243,7 @@ fn into_extracted_entities_strict_drops_unknown_kinds() {
kind: "spaceship".into(),
text: "Enterprise".into(),
}],
topics: vec![],
importance: None,
importance_reason: None,
};
@@ -187,6 +262,7 @@ fn into_extracted_entities_lenient_falls_back_to_misc() {
kind: "spaceship".into(),
text: "Enterprise".into(),
}],
topics: vec![],
importance: None,
importance_reason: None,
};
@@ -204,6 +280,7 @@ fn into_extracted_entities_disallowed_known_kind_falls_back_to_misc() {
kind: "person".into(),
text: "Alice".into(),
}],
topics: vec![],
importance: None,
importance_reason: None,
};
@@ -10,6 +10,83 @@ pub mod llm;
pub mod regex;
pub mod types;
use std::sync::Arc;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::util::redact::redact_endpoint;
pub use extractor::{CompositeExtractor, EntityExtractor, RegexEntityExtractor};
pub use llm::{LlmEntityExtractor, LlmExtractorConfig};
pub use types::{EntityKind, ExtractedEntities, ExtractedEntity, ExtractedTopic};
/// Build the extractor used by seal handlers to label new summary nodes.
///
/// Composition:
/// - regex extractor — always on, mechanical, near-zero cost
/// - LLM extractor with `emit_topics: true` — added when
/// `memory_tree.llm_extractor_endpoint` and `..._model` are both set
///
/// Differs from [`super::ScoringConfig::from_config`] (the chunk-admission
/// builder) in two ways: returns *just* an extractor (no thresholds /
/// weights / drop logic — none of which apply at seal time), and flips
/// `emit_topics` on so summaries surface thematic labels alongside
/// entities. Leaf-side scoring is unchanged.
pub fn build_summary_extractor(config: &Config) -> Arc<dyn EntityExtractor> {
let endpoint = config
.memory_tree
.llm_extractor_endpoint
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty());
let model = config
.memory_tree
.llm_extractor_model
.as_deref()
.map(str::trim)
.filter(|s| !s.is_empty());
let (Some(endpoint), Some(model)) = (endpoint, model) else {
log::debug!(
"[memory_tree::extract] summary extractor: LLM not configured — using regex-only"
);
return Arc::new(CompositeExtractor::regex_only());
};
let timeout_ms = config
.memory_tree
.llm_extractor_timeout_ms
.unwrap_or(15_000);
let cfg = LlmExtractorConfig {
endpoint: endpoint.to_string(),
model: model.to_string(),
timeout: std::time::Duration::from_millis(timeout_ms),
emit_topics: true,
..LlmExtractorConfig::default()
};
match LlmEntityExtractor::new(cfg) {
Ok(llm) => {
// Drop to debug (diagnostic, not always-on) and redact the endpoint
// so embedded credentials (e.g. api keys in URL) don't leak.
log::debug!(
"[memory_tree::extract] summary extractor: regex + LLM endpoint={} model={} \
timeout_ms={} emit_topics=true",
redact_endpoint(endpoint),
model,
timeout_ms
);
Arc::new(CompositeExtractor::new(vec![
Box::new(RegexEntityExtractor),
Box::new(llm),
]))
}
Err(err) => {
log::warn!(
"[memory_tree::extract] summary extractor: LlmEntityExtractor construction \
failed: {err:#} — falling back to regex-only"
);
Arc::new(CompositeExtractor::regex_only())
}
}
}
@@ -26,12 +26,22 @@ pub enum EntityKind {
Url,
Handle,
Hashtag,
// Semantic (reserved — not emitted in Phase 2)
// Semantic — emitted by the LLM extractor.
Person,
Organization,
Location,
Event,
Product,
/// Temporal expressions: "Friday", "Q2 2026", "EOD tomorrow", "next sprint".
Datetime,
/// Tools / frameworks / programming languages / services:
/// "Rust", "OAuth", "Slack API", "nomic-embed".
Technology,
/// Code / ticket / doc references that point at something addressable:
/// "PR #934", "src/openhuman/...", "OH-42", "ab7da2e2".
Artifact,
/// Amounts / metrics / money: "$5K", "20/min", "10k tokens", "52 chunks".
Quantity,
Misc,
// Thematic — scorer-surfaced topics (hashtag-like short phrases or
// LLM-extracted themes). Promoted into the canonical entity stream
@@ -54,6 +64,10 @@ impl EntityKind {
Self::Location => "location",
Self::Event => "event",
Self::Product => "product",
Self::Datetime => "datetime",
Self::Technology => "technology",
Self::Artifact => "artifact",
Self::Quantity => "quantity",
Self::Misc => "misc",
Self::Topic => "topic",
}
@@ -70,6 +84,10 @@ impl EntityKind {
"location" => Ok(Self::Location),
"event" => Ok(Self::Event),
"product" => Ok(Self::Product),
"datetime" => Ok(Self::Datetime),
"technology" => Ok(Self::Technology),
"artifact" => Ok(Self::Artifact),
"quantity" => Ok(Self::Quantity),
"misc" => Ok(Self::Misc),
"topic" => Ok(Self::Topic),
other => Err(format!("unknown entity kind: {other}")),
@@ -207,7 +225,12 @@ mod tests {
EntityKind::Location,
EntityKind::Event,
EntityKind::Product,
EntityKind::Datetime,
EntityKind::Technology,
EntityKind::Artifact,
EntityKind::Quantity,
EntityKind::Misc,
EntityKind::Topic,
] {
assert_eq!(EntityKind::parse(k.as_str()).unwrap(), k);
}
+16
View File
@@ -336,6 +336,22 @@ pub async fn score_chunks(chunks: &[Chunk], cfg: &ScoringConfig) -> Result<Vec<S
try_join_all(chunks.iter().map(|chunk| score_chunk(chunk, cfg))).await
}
/// Cheap-only batch scoring path used by the async queue ingest pipeline.
///
/// This preserves the same thresholds and admission gate as [`score_chunks`]
/// but guarantees no LLM extractor is consulted on the ingest hot path.
pub async fn score_chunks_fast(chunks: &[Chunk], cfg: &ScoringConfig) -> Result<Vec<ScoreResult>> {
let fast_cfg = ScoringConfig {
extractor: cfg.extractor.clone(),
weights: cfg.weights.clone(),
drop_threshold: cfg.drop_threshold,
llm_extractor: None,
definite_keep_threshold: cfg.definite_keep_threshold,
definite_drop_threshold: cfg.definite_drop_threshold,
};
score_chunks(chunks, &fast_cfg).await
}
// ── Persistence helpers used by the ingest orchestrator ─────────────────
/// Persist the score row + entity-index rows for one kept chunk.
@@ -11,6 +11,7 @@ fn test_chunk(content: &str) -> Chunk {
metadata: meta,
seq_in_source: 0,
created_at: Utc::now(),
partial_message: false,
}
}
+15
View File
@@ -371,6 +371,21 @@ pub fn lookup_entity(
})
}
pub fn list_entity_ids_for_node(config: &Config, node_id: &str) -> Result<Vec<String>> {
with_connection(config, |conn| {
let mut stmt = conn.prepare(
"SELECT DISTINCT entity_id
FROM mem_tree_entity_index
WHERE node_id = ?1
ORDER BY score DESC, timestamp_ms DESC, entity_id ASC",
)?;
let rows = stmt
.query_map(params![node_id], |row| row.get::<_, String>(0))?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(rows)
})
}
/// Count rows in the entity index (for tests / diagnostics).
pub fn count_entity_index(config: &Config) -> Result<u64> {
with_connection(config, |conn| {
@@ -26,12 +26,20 @@
//! summariser does no real I/O; when a networked summariser lands, wrap
//! DB calls in `tokio::task::spawn_blocking` to keep the runtime healthy.
use std::collections::BTreeSet;
use std::sync::Arc;
use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use rusqlite::Transaction;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::{
atomic::stage_summary, paths::slugify_source_id, SummaryComposeInput, SummaryTreeKind,
};
use crate::openhuman::memory::tree::score::embed::build_embedder_from_config;
use crate::openhuman::memory::tree::score::extract::EntityExtractor;
use crate::openhuman::memory::tree::score::resolver::canonicalise;
use crate::openhuman::memory::tree::source_tree::registry::new_summary_id;
use crate::openhuman::memory::tree::source_tree::store;
use crate::openhuman::memory::tree::source_tree::summariser::{
@@ -47,6 +55,88 @@ use crate::openhuman::memory::tree::store::with_connection;
/// realistic source.
const MAX_CASCADE_DEPTH: u32 = 32;
/// How a sealed summary node's `entities` and `topics` fields get populated.
///
/// Each tree kind has different correct semantics:
/// - **Source** trees use [`LabelStrategy::ExtractFromContent`] so the
/// summariser's freshly-synthesised text gets its own pass through an
/// extractor. Captures emergent themes that no individual leaf expressed.
/// - **Global** trees use [`LabelStrategy::UnionFromChildren`] — their
/// inputs are already-labeled source-tree summaries; union preserves
/// labels for time-based retrieval ("days that mentioned Alice")
/// without an LLM call.
/// - **Topic** trees use [`LabelStrategy::Empty`] — their scope already
/// pins the dominant theme; inheriting auxiliary entities would
/// cross-pollinate unrelated topic trees and noise the entity index.
#[derive(Clone)]
pub enum LabelStrategy {
/// Run the extractor on the new summary's content; canonicalise the
/// result into `entities` (canonical_ids) and `topics` (labels).
ExtractFromContent(Arc<dyn EntityExtractor>),
/// Dedup-merge each input's `entities` and `topics` into the parent.
UnionFromChildren,
/// Leave both fields empty regardless of inputs.
Empty,
}
impl std::fmt::Debug for LabelStrategy {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ExtractFromContent(ex) => write!(f, "ExtractFromContent({})", ex.name()),
Self::UnionFromChildren => f.write_str("UnionFromChildren"),
Self::Empty => f.write_str("Empty"),
}
}
}
/// Resolve `entities` and `topics` for a freshly-summarised node according
/// to the chosen strategy. Errors propagate from the extractor (when used).
async fn resolve_labels(
strategy: &LabelStrategy,
inputs: &[SummaryInput],
summary_content: &str,
) -> Result<(Vec<String>, Vec<String>)> {
match strategy {
LabelStrategy::ExtractFromContent(extractor) => {
let extracted = extractor
.extract(summary_content)
.await
.context("seal-time extractor failed")?;
let canonical = canonicalise(&extracted);
let mut entities: Vec<String> = canonical
.into_iter()
.map(|c| c.canonical_id)
.collect::<BTreeSet<_>>()
.into_iter()
.collect();
entities.sort();
let mut topics: Vec<String> = extracted
.topics
.into_iter()
.map(|t| t.label)
.collect::<BTreeSet<_>>()
.into_iter()
.collect();
topics.sort();
Ok((entities, topics))
}
LabelStrategy::UnionFromChildren => {
let mut entities: BTreeSet<String> = BTreeSet::new();
let mut topics: BTreeSet<String> = BTreeSet::new();
for inp in inputs {
for e in &inp.entities {
entities.insert(e.clone());
}
for t in &inp.topics {
topics.insert(t.clone());
}
}
Ok((entities.into_iter().collect(), topics.into_iter().collect()))
}
LabelStrategy::Empty => Ok((Vec::new(), Vec::new())),
}
}
/// A single leaf being appended to an L0 buffer.
#[derive(Clone, Debug)]
pub struct LeafRef {
@@ -61,17 +151,22 @@ pub struct LeafRef {
/// Append a leaf to the source tree for `tree`, sealing buffers as they
/// fill. Returns the ids of any summaries that sealed during this call.
///
/// `strategy` controls how each sealed summary's `entities` and `topics`
/// are populated — see [`LabelStrategy`].
pub async fn append_leaf(
config: &Config,
tree: &Tree,
leaf: &LeafRef,
summariser: &dyn Summariser,
strategy: &LabelStrategy,
) -> Result<Vec<String>> {
log::debug!(
"[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={}",
"[source_tree::bucket_seal] append_leaf tree_id={} leaf_id={} tokens={} strategy={:?}",
tree.id,
leaf.chunk_id,
leaf.token_count
leaf.token_count,
strategy
);
// 1. Push leaf into L0 buffer (transactional).
@@ -85,7 +180,24 @@ pub async fn append_leaf(
)?;
// 2. Cascade seals upward until a level stays under budget.
cascade_seals(config, tree, summariser).await
cascade_seals(config, tree, summariser, strategy).await
}
/// Queue-oriented variant of [`append_leaf`].
///
/// This only appends the leaf to the L0 buffer and returns whether the
/// caller should enqueue a follow-up seal job for level 0.
pub fn append_leaf_deferred(config: &Config, tree: &Tree, leaf: &LeafRef) -> Result<bool> {
append_to_buffer(
config,
&tree.id,
0,
&leaf.chunk_id,
leaf.token_count as i64,
leaf.timestamp,
)?;
let buf = store::get_buffer(config, &tree.id, 0)?;
Ok(should_seal(&buf))
}
/// Transactionally append a single item to `(tree_id, level)`'s buffer.
@@ -127,20 +239,25 @@ async fn cascade_seals(
config: &Config,
tree: &Tree,
summariser: &dyn Summariser,
strategy: &LabelStrategy,
) -> Result<Vec<String>> {
cascade_all_from(config, tree, 0, summariser, None).await
cascade_all_from(config, tree, 0, summariser, None, strategy).await
}
/// Seal buffers starting at `start_level` and cascade upward. When
/// `force_now` is `Some`, the buffer at `start_level` is sealed regardless
/// of token budget (used by time-based flush). Upper levels are sealed
/// only when they cross the budget.
///
/// `strategy` is forwarded to every sealed level — same semantics as
/// [`append_leaf`].
pub async fn cascade_all_from(
config: &Config,
tree: &Tree,
start_level: u32,
summariser: &dyn Summariser,
force_now: Option<DateTime<Utc>>,
strategy: &LabelStrategy,
) -> Result<Vec<String>> {
let mut sealed_ids: Vec<String> = Vec::new();
let mut level: u32 = start_level;
@@ -169,7 +286,9 @@ pub async fn cascade_all_from(
break;
}
let summary_id = seal_one_level(config, tree, &buf, summariser).await?;
// Sync cascade — drives the level walk itself; doesn't need the
// queue follow-ups (we'll hit `seal_one_level` again next iter).
let summary_id = seal_one_level(config, tree, &buf, summariser, strategy, false).await?;
sealed_ids.push(summary_id);
level += 1;
}
@@ -185,7 +304,7 @@ pub async fn cascade_all_from(
/// summarisers that emit at the full token budget (e.g. the inert
/// fallback) collapse the cascade into a 1:1:1 chain instead of a real
/// tree.
fn should_seal(buf: &Buffer) -> bool {
pub(crate) fn should_seal(buf: &Buffer) -> bool {
if buf.is_empty() {
return false;
}
@@ -198,11 +317,31 @@ fn should_seal(buf: &Buffer) -> bool {
/// Seal `buf` at `level` into one summary at `level + 1`. Returns the new
/// summary id.
async fn seal_one_level(
///
/// `strategy` decides how `entities` and `topics` get populated on the new
/// summary node — see [`LabelStrategy`].
///
/// When `enqueue_follow_ups` is `true`, the function additionally inserts
/// follow-up job rows **inside the same transaction** that commits the
/// seal:
/// - `seal { tree_id, level: parent_level }` if the parent buffer's gate
/// is now met (parent-cascade enqueue)
/// - `topic_route { NodeRef::Summary { summary_id } }` for source trees
/// (so summary-level entities feed the topic-tree spawn pipeline)
///
/// Atomic enqueue eliminates the crash window where a seal commits but
/// the post-commit follow-up enqueues are silently lost on a worker
/// crash. The async-pipeline handler (`handle_seal`) passes `true`. The
/// synchronous in-process cascade caller ([`cascade_all_from`]) passes
/// `false` because it drives the cascade itself and topic_route isn't
/// part of the test/flush sync path.
pub(crate) async fn seal_one_level(
config: &Config,
tree: &Tree,
buf: &Buffer,
summariser: &dyn Summariser,
strategy: &LabelStrategy,
enqueue_follow_ups: bool,
) -> Result<String> {
let level = buf.level;
let target_level = level + 1;
@@ -237,7 +376,7 @@ async fn seal_one_level(
// Run summariser — async, OUTSIDE any DB transaction.
let ctx = SummaryContext {
tree_id: &tree.id,
tree_kind: TreeKind::Source,
tree_kind: tree.kind,
target_level,
token_budget: TOKEN_BUDGET,
};
@@ -246,6 +385,12 @@ async fn seal_one_level(
.await
.context("summariser failed during seal")?;
// Resolve labels (entities/topics) for the new summary node according
// to the chosen strategy. Done before the write tx so an extractor
// failure aborts the seal cleanly — same shape as the embedder guard
// below.
let (node_entities, node_topics) = resolve_labels(strategy, &inputs, &output.content).await?;
// Phase 4 (#710): embed the summary BEFORE opening the write tx so an
// embedder failure aborts the seal cleanly — nothing is persisted,
// the buffer stays intact, and a retry re-embeds from scratch. The
@@ -291,14 +436,18 @@ async fn seal_one_level(
let node = SummaryNode {
id: summary_id.clone(),
tree_id: tree.id.clone(),
tree_kind: TreeKind::Source,
// `seal_one_level` runs for source AND topic trees (handle_seal,
// cascade_all_from, flush). Hardcoding Source here would write
// topic-tree summaries with tree_kind='source' in
// mem_tree_summaries, breaking any query filtering on tree_kind.
tree_kind: tree.kind,
level: target_level,
parent_id: None,
child_ids: buf.item_ids.clone(),
content: output.content,
token_count: output.token_count,
entities: output.entities,
topics: output.topics,
entities: node_entities,
topics: node_topics,
time_range_start,
time_range_end,
score,
@@ -307,13 +456,88 @@ async fn seal_one_level(
embedding: Some(embedding),
};
// Phase MD-content: stage the summary .md file BEFORE opening the write
// tx. A staging failure aborts the seal cleanly — nothing is persisted
// and the buffer stays intact for retry.
//
// `bucket_seal.rs` handles both Source and Topic tree seals (Topic trees
// use the same cascade machinery via `handle_seal` in the job handler).
// Map TreeKind to SummaryTreeKind accordingly.
let summary_tree_kind = match tree.kind {
TreeKind::Topic => SummaryTreeKind::Topic,
_ => SummaryTreeKind::Source,
};
let scope_slug = {
// Path slug semantics per source kind:
//
// - Gmail source trees: scope is `"gmail:<participants>"` where
// participants is `addr1|addr2|...`. Strip the `gmail:` prefix so the
// path is `summaries/source/<participants_slug>/...` and mirrors the
// chunk layout under `email/<participants_slug>/`.
//
// - Topic trees: scope is the canonical entity_id (e.g.
// `"email:alice@example.com"`). Slugify the FULL string so topic-tree
// summaries and source-tree summaries don't share a path prefix.
//
// - All other source kinds (slack:, discord:, document:, …): slugify the
// FULL scope string. Stripping the prefix for non-Gmail sources was a
// bug — `"slack:#eng"` and `"discord:#eng"` would both produce slug
// `"eng"` and collide in `summaries/source/eng/`.
let s = &tree.scope;
match tree.kind {
TreeKind::Topic => slugify_source_id(s),
_ => {
if s.starts_with("gmail:") {
// Strip "gmail:" prefix; slugify the participants portion.
slugify_source_id(&s["gmail:".len()..])
} else {
// All other source kinds: slugify the full scope string.
slugify_source_id(s)
}
}
}
};
let compose_input = SummaryComposeInput {
summary_id: &node.id,
tree_kind: summary_tree_kind,
tree_id: &node.tree_id,
tree_scope: &tree.scope,
level: node.level,
child_ids: &node.child_ids,
child_count: node.child_ids.len(),
time_range_start: node.time_range_start,
time_range_end: node.time_range_end,
sealed_at: node.sealed_at,
body: &node.content,
};
// Stage the summary .md file and propagate any error — a staging failure
// aborts the seal entirely so the database never commits a row with
// content_path = NULL. The buffer stays unsealed and the job-retry path
// will re-attempt the file write on next execution.
let content_root = config.memory_tree_content_root();
let staged =
stage_summary(&content_root, &compose_input, &scope_slug, None).with_context(|| {
format!(
"stage_summary failed for {}; seal aborted, buffer stays unsealed for retry",
node.id
)
})?;
log::debug!(
"[source_tree::bucket_seal] staged summary {} → {}",
node.id,
staged.content_path
);
// Single write transaction: insert summary, clear this buffer, append
// summary id to parent buffer, bump tree max_level/root if needed.
// summary id to parent buffer, bump tree max_level/root if needed,
// and (when `enqueue_follow_ups`) atomically enqueue parent-seal +
// topic_route follow-ups so they can never desync from the commit.
// Re-read `max_level` from inside the tx so cascading seals within
// one call see the updated value from earlier levels.
let summary_id_for_closure = summary_id.clone();
let target_level_for_closure = target_level;
let tree_id = tree.id.clone();
let tree_kind = tree.kind;
with_connection(config, move |conn| {
let tx = conn.unchecked_transaction()?;
@@ -326,7 +550,7 @@ async fn seal_one_level(
.map(|n| n.max(0) as u32)
.context("Failed to read current max_level for tree")?;
store::insert_summary_tx(&tx, &node)?;
store::insert_summary_tx(&tx, &node, Some(&staged))?;
// Forward-compat: index any entities the summariser emitted into
// `mem_tree_entity_index` so Phase 4 retrieval can resolve
// "summaries mentioning Alice" via the same inverted index as
@@ -375,6 +599,43 @@ async fn seal_one_level(
};
store::upsert_buffer_tx(&tx, &parent)?;
// Atomic follow-up enqueues. Done INSIDE this tx — if the commit
// rolls back, the queue rows go with it; if it succeeds, the
// rows are durably visible to the worker pool. Eliminates the
// crash window where the seal commits but post-commit enqueues
// are lost.
if enqueue_follow_ups {
// Parent-cascade: if the new summary made the parent buffer
// cross its gate, enqueue the next level's seal. Dedupe key
// `seal:{tree_id}:{parent_level}` prevents duplicates if a
// parallel path already queued it.
if should_seal(&parent) {
use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx;
use crate::openhuman::memory::tree::jobs::types::{NewJob, SealPayload};
let parent_seal = SealPayload {
tree_id: tree_id.clone(),
level: target_level_for_closure,
force_now_ms: None,
};
enqueue_job_tx(&tx, &NewJob::seal(&parent_seal)?)?;
}
// Source-tree summary routing: feed the new summary's
// entities back into the topic-tree spawn pipeline. Topic
// and global trees are sinks — no fan-out from their seals.
if matches!(tree_kind, TreeKind::Source) {
use crate::openhuman::memory::tree::jobs::store::enqueue_tx as enqueue_job_tx;
use crate::openhuman::memory::tree::jobs::types::{
NewJob, NodeRef, TopicRoutePayload,
};
let route = TopicRoutePayload {
node: NodeRef::Summary {
summary_id: summary_id_for_closure.clone(),
},
};
enqueue_job_tx(&tx, &NewJob::topic_route(&route)?)?;
}
}
// Update tree root / max_level if we just climbed.
if target_level_for_closure > current_max {
store::update_tree_after_seal_tx(
@@ -447,7 +708,8 @@ fn hydrate_inputs(config: &Config, level: u32, item_ids: &[String]) -> Result<Ve
}
fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result<Vec<SummaryInput>> {
use crate::openhuman::memory::tree::score::store::get_score;
use crate::openhuman::memory::tree::content_store::read as content_read;
use crate::openhuman::memory::tree::score::store::{get_score, list_entity_ids_for_node};
use crate::openhuman::memory::tree::store::get_chunk;
let mut out: Vec<SummaryInput> = Vec::with_capacity(chunk_ids.len());
@@ -461,17 +723,30 @@ fn hydrate_leaf_inputs(config: &Config, chunk_ids: &[String]) -> Result<Vec<Summ
continue;
}
};
let score = get_score(config, id)?;
let (score_value, entities, topics) = match &score {
Some(row) => (row.total, Vec::new(), chunk.metadata.tags.clone()),
None => (0.0, Vec::new(), chunk.metadata.tags.clone()),
};
let score_value = get_score(config, id)?.map(|row| row.total).unwrap_or(0.0);
// Pull canonical entity ids from the inverted index — that's the
// authoritative source for "what entities are attached to this
// chunk." Topics live on the chunk's metadata tags.
// [`LabelStrategy::UnionFromChildren`] reads these fields off
// each `SummaryInput` to roll labels up the tree.
let entities = list_entity_ids_for_node(config, id).unwrap_or_default();
// Read the full body from disk — the `content` column in SQLite holds
// a ≤500-char preview after the MD-on-disk migration. The summariser
// must receive the complete chunk text so the seal output is not a
// summary of previews.
//
// For pre-MD-migration chunks (no content_path recorded) this call
// returns Err; callers that want to handle legacy rows should check
// content_path presence before calling hydrate_inputs.
let body = content_read::read_chunk_body(config, id).with_context(|| {
format!("[source_tree::bucket_seal] hydrate_leaf_inputs: read body for chunk {id}")
})?;
out.push(SummaryInput {
id: chunk.id.clone(),
content: chunk.content.clone(),
content: body,
token_count: chunk.token_count,
entities,
topics,
topics: chunk.metadata.tags.clone(),
time_range_start: chunk.metadata.time_range.0,
time_range_end: chunk.metadata.time_range.1,
score: score_value,
@@ -481,6 +756,8 @@ 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::tree::content_store::read as content_read;
let mut out: Vec<SummaryInput> = Vec::with_capacity(summary_ids.len());
for id in summary_ids {
let node = match store::get_summary(config, id)? {
@@ -492,9 +769,15 @@ fn hydrate_summary_inputs(config: &Config, summary_ids: &[String]) -> Result<Vec
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
// over L1 summary content and need the full text, not a preview.
let body = content_read::read_summary_body(config, id).with_context(|| {
format!("[source_tree::bucket_seal] hydrate_summary_inputs: read body for summary {id}")
})?;
out.push(SummaryInput {
id: node.id.clone(),
content: node.content.clone(),
content: body,
token_count: node.token_count,
entities: node.entities.clone(),
topics: node.topics.clone(),
@@ -1,8 +1,29 @@
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
use tempfile::TempDir;
/// Stage a batch of chunks to the content store so that `read_chunk_body`
/// can find the on-disk file during seals. Tests that call `upsert_chunks`
/// and then trigger a seal MUST also call this helper; otherwise
/// `hydrate_leaf_inputs` will fail with "no content_path for chunk_id".
fn stage_test_chunks(cfg: &Config, chunks: &[crate::openhuman::memory::tree::types::Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged =
content_store::stage_chunks(&content_root, chunks).expect("stage_chunks for test chunks");
// Record the content_path + content_sha256 pointers in SQLite so the
// store's `get_chunk_content_pointers` can resolve them later.
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -36,7 +57,9 @@ async fn append_below_budget_does_not_seal() {
// Chunks don't exist in DB — we're only exercising the buffer
// accounting, which doesn't require leaf rows until a seal fires.
let leaf = mk_leaf("leaf-1", 100, 1_700_000_000_000);
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert!(sealed.is_empty());
let buf = store::get_buffer(&cfg, &tree.id, 0).unwrap();
@@ -72,15 +95,21 @@ async fn crossing_budget_triggers_seal() {
token_count: tokens,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
let c1 = mk_chunk(0, 6_000);
let c2 = mk_chunk(1, 6_000);
// Budget-relative sizes so the test stays correct as TOKEN_BUDGET shifts:
// each leaf is 60% of budget, so the second append crosses the threshold.
let per_leaf = TOKEN_BUDGET * 6 / 10;
let c1 = mk_chunk(0, per_leaf);
let c2 = mk_chunk(1, per_leaf);
upsert_chunks(&cfg, &[c1.clone(), c2.clone()]).unwrap();
// Stage both chunks to disk so the seal's hydrator can read full bodies.
stage_test_chunks(&cfg, &[c1.clone(), c2.clone()]);
// Two leaves whose combined token_sum (12k) exceeds the 10k budget.
let leaf1 = LeafRef {
chunk_id: c1.id.clone(),
token_count: 6_000,
token_count: per_leaf,
timestamp: ts,
content: c1.content.clone(),
entities: vec![],
@@ -89,7 +118,7 @@ async fn crossing_budget_triggers_seal() {
};
let leaf2 = LeafRef {
chunk_id: c2.id.clone(),
token_count: 6_000,
token_count: per_leaf,
timestamp: ts,
content: c2.content.clone(),
entities: vec![],
@@ -97,10 +126,14 @@ async fn crossing_budget_triggers_seal() {
score: 0.5,
};
let first = append_leaf(&cfg, &tree, &leaf1, &summariser).await.unwrap();
let first = append_leaf(&cfg, &tree, &leaf1, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert!(first.is_empty(), "first append below budget — no seal");
let second = append_leaf(&cfg, &tree, &leaf2, &summariser).await.unwrap();
let second = append_leaf(&cfg, &tree, &leaf2, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert_eq!(second.len(), 1, "second append crosses budget — one seal");
let summary_id = &second[0];
@@ -169,6 +202,7 @@ async fn fanout_at_l1_triggers_l2_seal() {
token_count: 10_000,
seq_in_source: seq,
created_at: ts,
partial_message: false,
}
};
@@ -177,6 +211,8 @@ async fn fanout_at_l1_triggers_l2_seal() {
for seq in 0..fanout {
let chunk = mk_chunk(seq);
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
// Stage to disk so the seal hydrator can read the full body.
stage_test_chunks(&cfg, &[chunk.clone()]);
let leaf = LeafRef {
chunk_id: chunk.id.clone(),
token_count: chunk.token_count,
@@ -186,7 +222,9 @@ async fn fanout_at_l1_triggers_l2_seal() {
topics: vec![],
score: 0.5,
};
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
all_sealed.extend(sealed);
}
@@ -253,8 +291,11 @@ async fn upper_level_does_not_seal_below_fanout() {
token_count: 10_000,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(&cfg, &[chunk.clone()]).unwrap();
// Stage to disk so the seal hydrator can read the full body.
stage_test_chunks(&cfg, &[chunk.clone()]);
let leaf = LeafRef {
chunk_id: chunk.id,
token_count: chunk.token_count,
@@ -264,7 +305,9 @@ async fn upper_level_does_not_seal_below_fanout() {
topics: vec![],
score: 0.5,
};
let _ = append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
let _ = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
}
let t = store::get_tree(&cfg, &tree.id).unwrap().unwrap();
@@ -281,3 +324,334 @@ async fn upper_level_does_not_seal_below_fanout() {
stop_before as u64
);
}
// ── LabelStrategy tests (#TBD) ────────────────────────────────────────────
//
// These exercise the three labeling modes seal_one_level supports. We use
// a short token budget so the seal fires on a single leaf — keeps the
// arithmetic of "what entities/topics end up on the parent" obvious.
/// Helper: persist a substantive chunk and return a `LeafRef` referencing
/// it, with caller-supplied entity/topic labels (used by Union/Empty tests).
///
/// To match production, entity labels are written into `mem_tree_entity_index`
/// (where seal-time hydration reads them from) and topic labels are stored
/// on `chunk.metadata.tags` (the production source of leaf-level topics).
fn seed_leaf(
cfg: &Config,
seq: u32,
content: &str,
entities: Vec<String>,
topics: Vec<String>,
) -> LeafRef {
use crate::openhuman::memory::tree::score::extract::EntityKind;
use crate::openhuman::memory::tree::score::resolver::CanonicalEntity;
use crate::openhuman::memory::tree::score::store::index_entity;
use crate::openhuman::memory::tree::store::upsert_chunks;
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
use chrono::TimeZone;
let ts = Utc
.timestamp_millis_opt(1_700_000_000_000 + seq as i64)
.unwrap();
let chunk = Chunk {
id: chunk_id(SourceKind::Chat, "slack:#eng", seq, content),
content: content.to_string(),
metadata: Metadata {
source_kind: SourceKind::Chat,
source_id: "slack:#eng".into(),
owner: "alice".into(),
timestamp: ts,
time_range: (ts, ts),
tags: topics.clone(),
source_ref: Some(SourceRef::new(format!("slack://x{seq}"))),
},
// Bust TOKEN_BUDGET in one leaf so the seal fires immediately.
token_count: 10_000,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[chunk.clone()]).unwrap();
// Stage the chunk to disk so `hydrate_leaf_inputs` can read the full body
// via `read_chunk_body` during a seal triggered by `append_leaf`.
stage_test_chunks(cfg, &[chunk.clone()]);
// Mirror production indexing: entities go into mem_tree_entity_index
// so the seal hydrator can pull them via list_entity_ids_for_node.
for entity_id in &entities {
let kind = entity_id
.split_once(':')
.map_or(EntityKind::Misc, |(k, _)| {
EntityKind::parse(k).unwrap_or(EntityKind::Misc)
});
let surface = entity_id
.split_once(':')
.map_or(entity_id.as_str(), |(_, v)| v);
let e = CanonicalEntity {
canonical_id: entity_id.clone(),
kind,
surface: surface.to_string(),
span_start: 0,
span_end: surface.len() as u32,
score: 1.0,
};
index_entity(cfg, &e, &chunk.id, "leaf", ts.timestamp_millis(), None).unwrap();
}
LeafRef {
chunk_id: chunk.id.clone(),
token_count: chunk.token_count,
timestamp: ts,
content: chunk.content.clone(),
entities,
topics,
score: 0.5,
}
}
#[tokio::test]
async fn seal_with_extract_strategy_populates_entities_and_topics() {
use crate::openhuman::memory::tree::score::extract::{CompositeExtractor, EntityExtractor};
use std::sync::Arc;
let (_tmp, cfg) = test_config();
let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
let summariser = InertSummariser::new();
// Content the regex extractor can find: an email and a hashtag. The
// inert summariser concatenates leaf content into the L1 summary, so
// these tokens survive into the summary text and the extractor finds
// them when run on the summary content.
let leaf = seed_leaf(
&cfg,
0,
"alice@example.com is leading the #launch sprint this week.",
vec![],
vec![],
);
let extractor: Arc<dyn EntityExtractor> = Arc::new(CompositeExtractor::regex_only());
let strategy = LabelStrategy::ExtractFromContent(extractor);
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &strategy)
.await
.unwrap();
assert_eq!(sealed.len(), 1, "single 10k-token leaf should seal L0→L1");
let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
assert!(
summary
.entities
.iter()
.any(|e| e == "email:alice@example.com"),
"ExtractFromContent should surface the email entity from summary text; got entities={:?}",
summary.entities
);
assert!(
summary.topics.iter().any(|t| t == "launch"),
"ExtractFromContent should surface the hashtag-derived topic; got topics={:?}",
summary.topics
);
}
#[tokio::test]
async fn seal_with_union_strategy_inherits_labels_from_children() {
let (_tmp, cfg) = test_config();
let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
let summariser = InertSummariser::new();
// Two leaves with overlapping + distinct labels. Union should
// dedup-merge them into the parent.
let leaf1 = seed_leaf(
&cfg,
0,
"first leaf body",
vec!["email:alice@example.com".into(), "topic:phoenix".into()],
vec!["phoenix".into(), "launch".into()],
);
let leaf2 = seed_leaf(
&cfg,
1,
"second leaf body",
vec!["email:alice@example.com".into(), "person:bob".into()],
vec!["launch".into(), "qa".into()],
);
// L0 seals when the budget is crossed. With each leaf at 10k tokens,
// the first append triggers a seal containing only leaf1; we want a
// seal containing both, so use UnionFromChildren and a single seal of
// both leaves at once. The simplest way is to lower budget by sealing
// two leaves into one buffer — the second append crosses budget, so
// the seal contains [leaf1, leaf2].
//
// Adjust by using smaller token counts so both fit in L0 first, then
// a third append triggers a seal containing both. Reuse the helper
// and override the leaf's token_count for this test.
// Each leaf at half the budget so two together hit threshold exactly.
let per_leaf = TOKEN_BUDGET / 2;
let leaf1 = LeafRef {
token_count: per_leaf,
..leaf1
};
let leaf2 = LeafRef {
token_count: per_leaf,
..leaf2
};
// First leaf: under budget, no seal.
let sealed_1 = append_leaf(
&cfg,
&tree,
&leaf1,
&summariser,
&LabelStrategy::UnionFromChildren,
)
.await
.unwrap();
assert!(sealed_1.is_empty());
// Second leaf: crosses budget → one seal covering both leaves.
let sealed_2 = append_leaf(
&cfg,
&tree,
&leaf2,
&summariser,
&LabelStrategy::UnionFromChildren,
)
.await
.unwrap();
assert_eq!(sealed_2.len(), 1);
let summary = store::get_summary(&cfg, &sealed_2[0]).unwrap().unwrap();
let entities: std::collections::BTreeSet<&str> =
summary.entities.iter().map(String::as_str).collect();
let topics: std::collections::BTreeSet<&str> =
summary.topics.iter().map(String::as_str).collect();
assert!(entities.contains("email:alice@example.com"));
assert!(entities.contains("topic:phoenix"));
assert!(entities.contains("person:bob"));
assert_eq!(
entities.len(),
3,
"expected 3 unique entities; got {entities:?}"
);
assert!(topics.contains("phoenix"));
assert!(topics.contains("launch"));
assert!(topics.contains("qa"));
assert_eq!(topics.len(), 3, "expected 3 unique topics; got {topics:?}");
}
#[tokio::test]
async fn seal_with_empty_strategy_leaves_labels_empty() {
let (_tmp, cfg) = test_config();
let tree = get_or_create_source_tree(&cfg, "slack:#eng").unwrap();
let summariser = InertSummariser::new();
// Leaf carries labels — Empty strategy should ignore them.
let leaf = seed_leaf(
&cfg,
0,
"alice@example.com discussing #launch",
vec!["email:alice@example.com".into(), "topic:launch".into()],
vec!["launch".into()],
);
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert_eq!(sealed.len(), 1);
let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
assert!(
summary.entities.is_empty(),
"Empty strategy must leave entities empty; got {:?}",
summary.entities
);
assert!(
summary.topics.is_empty(),
"Empty strategy must leave topics empty; got {:?}",
summary.topics
);
}
#[tokio::test]
async fn topic_tree_seal_persists_topic_kind_not_source() {
use crate::openhuman::memory::tree::source_tree::types::TreeStatus;
let (_tmp, cfg) = test_config();
// Build a topic tree directly — `seal_one_level` runs for both
// source and topic trees, and previously hardcoded Source on the
// resulting summary regardless of the parent tree's kind.
let tree = Tree {
id: "topic-tree-test-id".to_string(),
kind: TreeKind::Topic,
scope: "topic:launch".to_string(),
root_id: None,
max_level: 0,
status: TreeStatus::Active,
created_at: Utc::now(),
last_sealed_at: None,
};
store::insert_tree(&cfg, &tree).unwrap();
let summariser = InertSummariser::new();
let leaf = seed_leaf(&cfg, 0, "topic content", vec![], vec![]);
let sealed = append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert_eq!(sealed.len(), 1);
let summary = store::get_summary(&cfg, &sealed[0]).unwrap().unwrap();
assert_eq!(
summary.tree_kind,
TreeKind::Topic,
"topic-tree summary must persist tree_kind=Topic, not Source"
);
}
#[test]
fn scope_slug_non_gmail_uses_full_scope() {
// slack:#eng and discord:#eng must NOT produce the same scope slug.
// Previously, stripping everything before ':' made both → "eng".
// After Fix K, only gmail: strips the prefix — others use the full string.
use crate::openhuman::memory::tree::content_store::paths::slugify_source_id;
// Verify that the slug logic produces distinct values for different platforms.
let slack_slug = slugify_source_id("slack:#eng");
let discord_slug = slugify_source_id("discord:#eng");
assert_ne!(
slack_slug, discord_slug,
"slack:#eng and discord:#eng must produce distinct slugs; got slack={slack_slug:?} discord={discord_slug:?}"
);
// Both must include their platform prefix in the slug.
assert!(
slack_slug.contains("slack"),
"slack slug must include 'slack'; got {slack_slug:?}"
);
assert!(
discord_slug.contains("discord"),
"discord slug must include 'discord'; got {discord_slug:?}"
);
// Confirm gmail: correctly strips the "gmail:" prefix so the participants
// portion (used as the bucket key) matches the chunk path layout.
// scope_slug for a gmail source tree is built by stripping "gmail:" and
// slugifying the remainder; the result must equal slugify of just the
// participants string.
let participants = "alice@x.com|bob@y.com";
let participants_slug = slugify_source_id(participants);
let gmail_scope = format!("gmail:{participants}");
// Strip "gmail:" prefix as bucket_seal.rs does.
let gmail_slug = slugify_source_id(&gmail_scope["gmail:".len()..]);
assert_eq!(
participants_slug, gmail_slug,
"gmail scope_slug must equal slugify of participants portion; \
participants_slug={participants_slug:?} gmail_slug={gmail_slug:?}"
);
// Also assert the full-scope slug for gmail is DIFFERENT (shows the bug
// would still exist if we used the full string for gmail).
let gmail_full_slug = slugify_source_id(&gmail_scope);
assert_ne!(
gmail_full_slug, participants_slug,
"slugifying the full 'gmail:...' scope must differ from the participants-only slug"
);
}
+39 -10
View File
@@ -13,7 +13,7 @@ use anyhow::Result;
use chrono::{DateTime, Duration, Utc};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::source_tree::bucket_seal::cascade_all_from;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{cascade_all_from, LabelStrategy};
use crate::openhuman::memory::tree::source_tree::store;
use crate::openhuman::memory::tree::source_tree::summariser::Summariser;
use crate::openhuman::memory::tree::source_tree::types::DEFAULT_FLUSH_AGE_SECS;
@@ -25,6 +25,7 @@ pub async fn flush_stale_buffers(
config: &Config,
max_age: Duration,
summariser: &dyn Summariser,
strategy: &LabelStrategy,
) -> Result<usize> {
let now = Utc::now();
let cutoff = now - max_age;
@@ -48,7 +49,8 @@ pub async fn flush_stale_buffers(
continue;
}
};
let sealed = cascade_all_from(config, &tree, buf.level, summariser, Some(now)).await?;
let sealed =
cascade_all_from(config, &tree, buf.level, summariser, Some(now), strategy).await?;
seals += sealed.len();
}
Ok(seals)
@@ -58,11 +60,13 @@ pub async fn flush_stale_buffers(
pub async fn flush_stale_buffers_default(
config: &Config,
summariser: &dyn Summariser,
strategy: &LabelStrategy,
) -> Result<usize> {
flush_stale_buffers(
config,
Duration::seconds(DEFAULT_FLUSH_AGE_SECS),
summariser,
strategy,
)
.await
}
@@ -74,15 +78,17 @@ pub async fn force_flush_tree(
tree_id: &str,
summariser: &dyn Summariser,
now: Option<DateTime<Utc>>,
strategy: &LabelStrategy,
) -> Result<Vec<String>> {
let tree = store::get_tree(config, tree_id)?
.ok_or_else(|| anyhow::anyhow!("no tree with id {tree_id}"))?;
cascade_all_from(config, &tree, 0, summariser, now).await
cascade_all_from(config, &tree, 0, summariser, now, strategy).await
}
#[cfg(test)]
mod tests {
use super::*;
use crate::openhuman::memory::tree::content_store;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::source_tree::registry::get_or_create_source_tree;
use crate::openhuman::memory::tree::source_tree::summariser::inert::InertSummariser;
@@ -90,6 +96,20 @@ mod tests {
use crate::openhuman::memory::tree::types::{chunk_id, Chunk, Metadata, SourceKind, SourceRef};
use tempfile::TempDir;
fn stage_test_chunks(cfg: &Config, chunks: &[Chunk]) {
let content_root = cfg.memory_tree_content_root();
std::fs::create_dir_all(&content_root).expect("create content_root for test");
let staged = content_store::stage_chunks(&content_root, chunks)
.expect("stage_chunks for test chunks");
crate::openhuman::memory::tree::store::with_connection(cfg, |conn| {
let tx = conn.unchecked_transaction()?;
crate::openhuman::memory::tree::store::upsert_staged_chunks_tx(&tx, &staged)?;
tx.commit()?;
Ok(())
})
.expect("persist staged chunk pointers");
}
fn test_config() -> (TempDir, Config) {
let tmp = TempDir::new().unwrap();
let mut cfg = Config::default();
@@ -124,8 +144,10 @@ mod tests {
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(),
@@ -136,12 +158,15 @@ mod tests {
topics: vec![],
score: 0.5,
};
append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser)
append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
let seals =
flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert_eq!(seals, 1);
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 1);
@@ -172,6 +197,7 @@ mod tests {
token_count: 50,
seq_in_source: 0,
created_at: now,
partial_message: false,
};
upsert_chunks(&cfg, &[c.clone()]).unwrap();
let leaf = LeafRef {
@@ -183,11 +209,14 @@ mod tests {
topics: vec![],
score: 0.5,
};
append_leaf(&cfg, &tree, &leaf, &summariser).await.unwrap();
let seals = flush_stale_buffers(&cfg, Duration::days(7), &summariser)
append_leaf(&cfg, &tree, &leaf, &summariser, &LabelStrategy::Empty)
.await
.unwrap();
let seals =
flush_stale_buffers(&cfg, Duration::days(7), &summariser, &LabelStrategy::Empty)
.await
.unwrap();
assert_eq!(seals, 0);
assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0);
}
+1 -1
View File
@@ -22,7 +22,7 @@ pub mod store;
pub mod summariser;
pub mod types;
pub use bucket_seal::{append_leaf, LeafRef};
pub use bucket_seal::{append_leaf, append_leaf_deferred, LabelStrategy, LeafRef};
pub use registry::get_or_create_source_tree;
pub use store::{get_summary_embedding, set_summary_embedding};
pub use summariser::{build_summariser, inert::InertSummariser, llm::LlmSummariser, Summariser};
+32 -4
View File
@@ -21,6 +21,7 @@ use chrono::{DateTime, TimeZone, Utc};
use rusqlite::{params, Connection, OptionalExtension, Transaction};
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::StagedSummary;
use crate::openhuman::memory::tree::score::embed::{decode_optional_blob, pack_checked};
use crate::openhuman::memory::tree::source_tree::types::{
Buffer, SummaryNode, Tree, TreeKind, TreeStatus,
@@ -180,7 +181,16 @@ fn row_to_tree(row: &rusqlite::Row<'_>) -> rusqlite::Result<Tree> {
/// Phase 4 (#710): if `node.embedding` is `Some`, the packed vector is
/// written to the `embedding` blob column; `None` writes NULL so legacy
/// rows from Phases 1-3 (no embed) read back identically.
pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Result<()> {
///
/// Phase MD-content: if `staged` is `Some`, writes `content_path` and
/// `content_sha256` and truncates `content` to a ≤500-char preview. Callers
/// that have not yet staged the file pass `None`, in which case the full
/// `node.content` is stored (legacy behaviour).
pub(crate) fn insert_summary_tx(
tx: &Transaction<'_>,
node: &SummaryNode,
staged: Option<&StagedSummary>,
) -> Result<()> {
let embedding_blob: Option<Vec<u8>> = match node.embedding.as_deref() {
Some(v) => Some(
pack_checked(v)
@@ -188,14 +198,30 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res
),
None => None,
};
// Phase MD-content: when a staged file exists, truncate `content` to a
// ≤500-char plain-text preview (char boundary safe via chars().take(500)).
let (content_preview, content_path, content_sha256) = match staged {
Some(s) => {
let preview: String = node.content.chars().take(500).collect();
(
preview,
Some(s.content_path.clone()),
Some(s.content_sha256.clone()),
)
}
None => (node.content.clone(), None, None),
};
tx.execute(
"INSERT OR IGNORE INTO mem_tree_summaries (
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
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16)",
score, sealed_at_ms, deleted, embedding,
content_path, content_sha256
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
params![
node.id,
node.tree_id,
@@ -203,7 +229,7 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res
node.level,
node.parent_id,
serde_json::to_string(&node.child_ids)?,
node.content,
content_preview,
node.token_count,
serde_json::to_string(&node.entities)?,
serde_json::to_string(&node.topics)?,
@@ -213,6 +239,8 @@ pub(crate) fn insert_summary_tx(tx: &Transaction<'_>, node: &SummaryNode) -> Res
node.sealed_at.timestamp_millis(),
node.deleted as i64,
embedding_blob,
content_path,
content_sha256,
],
)
.with_context(|| format!("Failed to insert summary id={}", node.id))?;
@@ -71,7 +71,7 @@ fn summary_insert_and_fetch() {
let node = sample_summary("sum-1", "tree-1", 1);
with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
insert_summary_tx(&tx, &node)?;
insert_summary_tx(&tx, &node, None)?;
tx.commit()?;
Ok(())
})
@@ -90,8 +90,8 @@ fn summary_insert_is_idempotent_on_id() {
let node = sample_summary("sum-1", "tree-1", 1);
with_connection(&cfg, |conn| {
let tx = conn.unchecked_transaction()?;
insert_summary_tx(&tx, &node)?;
insert_summary_tx(&tx, &node)?;
insert_summary_tx(&tx, &node, None)?;
insert_summary_tx(&tx, &node, None)?;
tx.commit()?;
Ok(())
})
@@ -6,10 +6,12 @@
//! When the source / topic / global tree's bucket-seal cascade decides to
//! fold N contributions (raw leaves at L0→L1, or lower-level summaries at
//! L_n→L_{n+1}), this summariser is asked to produce the parent node's
//! `content` + derived `entities` + `topics`. The seal machinery itself
//! (bucket budgeting, level promotion, `mem_tree_summaries` persistence)
//! is unchanged — only the text inside the summary row differs from
//! [`super::inert::InertSummariser`].
//! `content`. The seal machinery itself (bucket budgeting, level
//! promotion, `mem_tree_summaries` persistence) is unchanged — only the
//! text inside the summary row differs from [`super::inert::InertSummariser`].
//! Entities and topics on `SummaryOutput` are always emitted empty by
//! this summariser; canonical entity ids are populated separately by the
//! entity extractor.
//!
//! ## Soft-fallback contract
//!
@@ -18,14 +20,12 @@
//! parent row. We therefore promise **never** to return `Err`: every
//! failure (transport, HTTP status, JSON shape) falls back to the same
//! deterministic concat-and-truncate behaviour as `InertSummariser` and
//! logs a warn. Callers distinguish "LLM ran fine" from "we fell back"
//! only by observing whether the returned `entities`/`topics` are
//! populated — the inert branch emits empty vecs.
//! logs a warn.
//!
//! ## Prompt shape
//!
//! The system prompt commits the model to returning JSON with the shape
//! `{ summary, entities, topics }`. We use Ollama's `format: "json"` +
//! `{ summary }`. We use Ollama's `format: "json"` +
//! `temperature: 0.0` to maximise determinism — same knobs the entity
//! extractor already uses with success.
@@ -40,15 +40,24 @@ use super::inert::InertSummariser;
use super::{Summariser, SummaryContext, SummaryInput, SummaryOutput};
use crate::openhuman::memory::tree::types::approx_token_count;
/// Hard cap on summary OUTPUT tokens, regardless of the seal's
/// `ctx.token_budget`. Driven by the embedder's context window —
/// `nomic-embed-text-v1.5` accepts up to 8192 tokens and Phase 4
/// (`source_tree::bucket_seal`) embeds the summary right after we
/// produce it. If our summary overshoots, the embedder returns 500
/// and the whole seal transaction rolls back → no summary persists.
/// 6000 leaves a safety margin for tokenizer differences and the
/// JSON wrapper.
const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 6_000;
/// Hard cap on summariser output length (in approximate tokens).
///
/// Two constraints set this:
///
/// 1. The downstream embedder (`nomic-embed-text-v1.5`) accepts up to
/// 8192 tokens, and Phase 4 (`source_tree::bucket_seal`) embeds the
/// summary right after we produce it. An overshoot returns HTTP 500
/// and rolls back the whole seal transaction.
/// 2. Empirically, small instruction-tuned models running locally
/// degrade quickly past ~3500 tokens — they drift, hallucinate, or
/// produce repetitive boilerplate as they extend toward longer
/// targets. Keeping the cap below that breakeven keeps output
/// quality stable on local Ollama deployments.
///
/// 3500 sits comfortably under the embedder ceiling AND below the local
/// LLM quality cliff. The post-generation [`clamp_to_budget`] enforces
/// this regardless of what the model produces.
const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 3_500;
/// Context window we ask Ollama for. Must match the value below in
/// [`OllamaOptions::num_ctx`] so the per-input clamp computed in
@@ -261,31 +270,18 @@ impl Summariser for LlmSummariser {
let (content, token_count) = clamp_to_budget(&parsed.summary, effective_budget);
log::debug!(
"[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={} \
surface_entities_dropped={} topics={}",
"[source_tree::summariser::llm] sealed tree_id={} level={} inputs={} tokens={}",
ctx.tree_id,
ctx.target_level,
inputs.len(),
token_count,
parsed.entities.len(),
parsed.topics.len()
token_count
);
// Drop LLM-emitted entities. The model returns surface forms
// ("Alice", "she"), but `SummaryNode.entities` is indexed via
// `index_summary_entity_ids_tx` as canonical ids. Surface forms
// would silently corrupt that index — searches by canonical id
// would not find these summaries. Canonicalisation is the
// entity extractor's job, not the summariser's. Topics stay
// because they're free-form labels, not indexed as canonical
// ids. The `entities` field stays in the prompt to nudge the
// model toward entity-aware summarisation; we just don't
// persist its output.
Ok(SummaryOutput {
content,
token_count,
entities: Vec::new(),
topics: dedupe_sorted(parsed.topics),
topics: Vec::new(),
})
}
}
@@ -314,20 +310,21 @@ fn build_user_prompt(inputs: &[SummaryInput], per_input_cap_tokens: u32) -> Stri
out
}
/// System prompt. Token budget is templated in so the model aims under it.
fn system_prompt(budget: u32) -> String {
format!(
"You are a precise summariser. Summarise the user-provided contributions into a \
single cohesive passage that preserves concrete facts, decisions, named entities, \
and temporal ordering. Do not invent facts. Stay well under {budget} tokens.\n\
\n\
Return JSON only no prose, no markdown, no commentary. Schema:\n\
{{\n\
\x20 \"summary\": \"<summary body>\",\n\
\x20 \"entities\": [\"<named entity surface forms mentioned in the summary>\"],\n\
\x20 \"topics\": [\"<short topic labels covered by the summary>\"]\n\
}}"
)
/// System prompt. Length isn't templated in — empirically, telling small
/// instruction-tuned models "stay under N tokens" makes them produce
/// curt, generic output even when the input has plenty of substance.
/// Output is clamped post-generation by [`clamp_to_budget`] in the
/// caller, so we don't need the model to self-police length.
fn system_prompt(_budget: u32) -> String {
"You are a precise summariser. Summarise the user-provided contributions into a \
single cohesive passage that preserves concrete facts, decisions, \
and temporal ordering. Do not invent facts.\n\
\n\
Return JSON only no prose, no markdown, no commentary. Schema:\n\
{\n\
\x20 \"summary\": \"<summary body>\"\n\
}"
.to_string()
}
/// Truncate to the caller's token budget using the same ~4 chars/token
@@ -343,16 +340,6 @@ fn clamp_to_budget(text: &str, budget: u32) -> (String, u32) {
(truncated, tokens)
}
fn dedupe_sorted(mut items: Vec<String>) -> Vec<String> {
for item in items.iter_mut() {
*item = item.trim().to_string();
}
items.retain(|s| !s.is_empty());
items.sort();
items.dedup();
items
}
fn truncate_for_log(s: &str, max_chars: usize) -> String {
if s.chars().count() <= max_chars {
return s.to_string();
@@ -405,10 +392,6 @@ struct OllamaResponseMessage {
struct LlmSummaryOutput {
#[serde(default)]
summary: String,
#[serde(default)]
entities: Vec<String>,
#[serde(default)]
topics: Vec<String>,
}
#[cfg(test)]
@@ -493,12 +476,16 @@ mod tests {
}
#[test]
fn system_prompt_templates_budget() {
fn system_prompt_describes_schema() {
// Budget is no longer templated into the prompt — small models
// produced overly curt output when told to "stay under N tokens".
// The clamp in `clamp_to_budget` handles enforcement instead.
let p = system_prompt(4096);
assert!(p.contains("4096"));
assert!(!p.contains("4096"));
assert!(!p.contains("Stay well under"));
assert!(p.contains("\"summary\""));
assert!(p.contains("\"entities\""));
assert!(p.contains("\"topics\""));
assert!(!p.contains("\"entities\""));
assert!(!p.contains("\"topics\""));
}
#[test]
@@ -516,19 +503,6 @@ mod tests {
assert!(t <= 6);
}
#[test]
fn dedupe_sorted_trims_and_dedupes() {
let out = dedupe_sorted(vec![
"Bob".into(),
" Alice ".into(),
"Bob".into(),
"".into(),
" ".into(),
"Alice".into(),
]);
assert_eq!(out, vec!["Alice", "Bob"]);
}
#[test]
fn truncate_for_log_short_input_unchanged() {
assert_eq!(truncate_for_log("hi", 10), "hi");
@@ -592,16 +566,24 @@ mod tests {
assert!(!req.stream);
assert_eq!(req.options.temperature, 0.0);
assert_eq!(req.messages[0].role, "system");
assert!(req.messages[0].content.contains("2048"));
assert!(req.messages[0].content.contains("\"summary\""));
assert_eq!(req.messages[1].role, "user");
assert_eq!(req.messages[1].content, "body");
}
#[test]
fn llm_output_deserialises_with_missing_fields() {
fn llm_output_deserialises_with_only_summary() {
let v: LlmSummaryOutput = serde_json::from_str(r#"{"summary":"hi"}"#).unwrap();
assert_eq!(v.summary, "hi");
assert!(v.entities.is_empty());
assert!(v.topics.is_empty());
}
#[test]
fn llm_output_ignores_extraneous_fields() {
// Prompt no longer asks for entities/topics, but if the model
// emits them anyway we should still parse `summary` cleanly.
let v: LlmSummaryOutput =
serde_json::from_str(r#"{"summary":"hi","entities":["Alice"],"topics":["x"]}"#)
.unwrap();
assert_eq!(v.summary, "hi");
}
}
@@ -172,13 +172,17 @@ impl Buffer {
}
}
/// Token ceiling for one summariser invocation — aligned with the Phase 1
/// chunker ceiling so a single leaf never busts a seal on its own.
/// Token ceiling for one summariser invocation.
///
/// Sized for the local 1B summariser (`gemma3:1b-it-qat`), which produces
/// noticeably better summaries with ≤4-5k input than at higher caps. The
/// chunker's `DEFAULT_CHUNK_MAX_TOKENS` (3_000) sits below this so each
/// L0 buffer accumulates roughly 1-3 chunks before sealing.
///
/// Gates only the L0 → L1 seal: leaves are fan-in by raw token volume so
/// the summariser input stays bounded. Summaries above L0 use
/// [`SUMMARY_FANOUT`] instead — see `bucket_seal::should_seal`.
pub const TOKEN_BUDGET: u32 = 10_000;
pub const TOKEN_BUDGET: u32 = 4_500;
/// Sibling count that triggers a seal at level ≥ 1 (summaries → next level).
///
+279
View File
@@ -13,6 +13,7 @@ use rusqlite::{params, Connection, OptionalExtension, Transaction};
use std::time::Duration;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::content_store::StagedChunk;
use crate::openhuman::memory::tree::types::{Chunk, Metadata, SourceKind, SourceRef};
const DB_DIR: &str = "memory_tree";
@@ -21,6 +22,12 @@ const DEFAULT_LIST_LIMIT: usize = 100;
const MAX_LIST_LIMIT: usize = 10_000;
const SQLITE_BUSY_TIMEOUT: Duration = Duration::from_secs(5);
pub const CHUNK_STATUS_PENDING_EXTRACTION: &str = "pending_extraction";
pub const CHUNK_STATUS_ADMITTED: &str = "admitted";
pub const CHUNK_STATUS_BUFFERED: &str = "buffered";
pub const CHUNK_STATUS_SEALED: &str = "sealed";
pub const CHUNK_STATUS_DROPPED: &str = "dropped";
const SCHEMA: &str = "
PRAGMA foreign_keys = ON;
@@ -174,6 +181,35 @@ CREATE TABLE IF NOT EXISTS mem_tree_entity_hotness (
CREATE INDEX IF NOT EXISTS idx_mem_tree_entity_hotness_score
ON mem_tree_entity_hotness(last_hotness);
-- Async job queue for memory-tree work (extract admit buffer seal
-- topic-route daily digest). Producers (ingest, schedulers, handlers)
-- enqueue rows transactionally; the worker pool claims them via the
-- `(status, available_at_ms)` index. `dedupe_key` is enforced as unique
-- only for ready/running rows so a completed job's key can be re-used.
CREATE TABLE IF NOT EXISTS mem_tree_jobs (
id TEXT PRIMARY KEY,
kind TEXT NOT NULL,
payload_json TEXT NOT NULL,
dedupe_key TEXT,
status TEXT NOT NULL DEFAULT 'ready',
attempts INTEGER NOT NULL DEFAULT 0,
max_attempts INTEGER NOT NULL DEFAULT 5,
available_at_ms INTEGER NOT NULL,
locked_until_ms INTEGER,
last_error TEXT,
created_at_ms INTEGER NOT NULL,
started_at_ms INTEGER,
completed_at_ms INTEGER
);
CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_ready
ON mem_tree_jobs(status, available_at_ms);
CREATE INDEX IF NOT EXISTS idx_mem_tree_jobs_kind
ON mem_tree_jobs(kind);
CREATE UNIQUE INDEX IF NOT EXISTS idx_mem_tree_jobs_dedupe_active
ON mem_tree_jobs(dedupe_key)
WHERE dedupe_key IS NOT NULL AND status IN ('ready', 'running');
";
/// Upsert a batch of chunks atomically.
@@ -249,6 +285,66 @@ pub(crate) fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result
Ok(chunks.len())
}
/// Upsert staged chunks (with content_path + content_sha256) using an existing transaction.
///
/// Identical to `upsert_chunks_tx` but also writes the Phase MD-content pointer columns.
/// `content` column receives a ≤500-char plain-text preview of the body (the full body
/// lives on disk at `content_path`).
pub(crate) fn upsert_staged_chunks_tx(
tx: &Transaction<'_>,
staged: &[StagedChunk],
) -> Result<usize> {
if staged.is_empty() {
return Ok(0);
}
let mut stmt = tx.prepare(
"INSERT INTO mem_tree_chunks (
id, source_kind, source_id, source_ref, owner,
timestamp_ms, time_range_start_ms, time_range_end_ms,
tags_json, content, token_count, seq_in_source, created_at_ms,
content_path, content_sha256
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)
ON CONFLICT(id) DO UPDATE SET
source_kind = excluded.source_kind,
source_id = excluded.source_id,
source_ref = excluded.source_ref,
owner = excluded.owner,
timestamp_ms = excluded.timestamp_ms,
time_range_start_ms = excluded.time_range_start_ms,
time_range_end_ms = excluded.time_range_end_ms,
tags_json = excluded.tags_json,
content = excluded.content,
token_count = excluded.token_count,
seq_in_source = excluded.seq_in_source,
created_at_ms = excluded.created_at_ms,
content_path = excluded.content_path,
content_sha256 = excluded.content_sha256",
)?;
for s in staged {
let chunk = &s.chunk;
// Store a ≤500-char preview in the `content` column; full body is on disk.
let preview: String = chunk.content.chars().take(500).collect();
stmt.execute(params![
chunk.id,
chunk.metadata.source_kind.as_str(),
chunk.metadata.source_id,
chunk.metadata.source_ref.as_ref().map(|r| r.value.as_str()),
chunk.metadata.owner,
chunk.metadata.timestamp.timestamp_millis(),
chunk.metadata.time_range.0.timestamp_millis(),
chunk.metadata.time_range.1.timestamp_millis(),
serde_json::to_string(&chunk.metadata.tags)?,
preview,
chunk.token_count,
chunk.seq_in_source,
chunk.created_at.timestamp_millis(),
s.content_path,
s.content_sha256,
])?;
}
Ok(staged.len())
}
fn upsert_chunks_with_statement(
stmt: &mut rusqlite::Statement<'_>,
chunks: &[Chunk],
@@ -361,6 +457,70 @@ pub fn count_chunks(config: &Config) -> Result<u64> {
})
}
pub fn set_chunk_lifecycle_status(config: &Config, chunk_id: &str, status: &str) -> Result<()> {
with_connection(config, |conn| {
set_chunk_lifecycle_status_conn(conn, chunk_id, status)
})
}
pub(crate) fn set_chunk_lifecycle_status_tx(
tx: &Transaction<'_>,
chunk_id: &str,
status: &str,
) -> Result<()> {
set_chunk_lifecycle_status_conn(tx, chunk_id, status)
}
pub fn get_chunk_lifecycle_status(config: &Config, chunk_id: &str) -> Result<Option<String>> {
with_connection(config, |conn| {
get_chunk_lifecycle_status_conn(conn, chunk_id)
})
}
pub(crate) fn get_chunk_lifecycle_status_tx(
tx: &Transaction<'_>,
chunk_id: &str,
) -> Result<Option<String>> {
get_chunk_lifecycle_status_conn(tx, chunk_id)
}
fn get_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str) -> Result<Option<String>> {
let row = conn
.query_row(
"SELECT lifecycle_status FROM mem_tree_chunks WHERE id = ?1",
params![chunk_id],
|r| r.get::<_, String>(0),
)
.optional()?;
Ok(row)
}
pub fn count_chunks_by_lifecycle_status(config: &Config, status: &str) -> Result<u64> {
with_connection(config, |conn| {
let n: i64 = conn.query_row(
"SELECT COUNT(*) FROM mem_tree_chunks WHERE lifecycle_status = ?1",
params![status],
|r| r.get(0),
)?;
Ok(n.max(0) as u64)
})
}
fn set_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str, status: &str) -> Result<()> {
let changed = conn.execute(
"UPDATE mem_tree_chunks SET lifecycle_status = ?1 WHERE id = ?2",
params![status, chunk_id],
)?;
if changed == 0 {
log::warn!(
"[memory_tree::store] lifecycle update affected 0 rows chunk_id={} status={}",
chunk_id,
status
);
}
Ok(())
}
fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result<Chunk> {
let id: String = row.get(0)?;
let source_kind_s: String = row.get(1)?;
@@ -401,6 +561,10 @@ fn row_to_chunk(row: &rusqlite::Row<'_>) -> rusqlite::Result<Chunk> {
token_count: token_count.max(0) as u32,
seq_in_source: seq.max(0) as u32,
created_at,
// partial_message is not stored in SQLite — it's a transient chunker
// signal. Chunks read back from DB always get false (the column doesn't
// exist; callers that need this flag hold the Chunk in memory).
partial_message: false,
})
}
@@ -451,9 +615,124 @@ pub(crate) fn with_connection<T>(
// legacy summaries from Phases 1-3 read back as None; retrieval
// tolerates NULL by dropping the row to the bottom of a rerank.
add_column_if_missing(&conn, "mem_tree_summaries", "embedding", "BLOB")?;
// Async-pipeline lifecycle flag. Default 'admitted' so chunks ingested
// before the queue migration stay queryable. New writes start at
// 'pending_extraction'; the extract handler advances them to 'admitted'
// (then 'buffered' / 'sealed') or 'dropped'.
add_column_if_missing(
&conn,
"mem_tree_chunks",
"lifecycle_status",
"TEXT NOT NULL DEFAULT 'admitted'",
)?;
conn.execute_batch(
"CREATE INDEX IF NOT EXISTS idx_mem_tree_chunks_lifecycle \
ON mem_tree_chunks(lifecycle_status);",
)
.context("Failed to create mem_tree_chunks lifecycle index")?;
// Phase MD-content (#TBD): pointer + integrity hash. Body lives at
// <content_root>/<content_path> as a .md file. Both nullable so chunks
// ingested before this migration read back with NULL (body still in
// `content`). New writes populate both columns. The `content` column
// stores a 500-char plain-text preview instead of the full body.
add_column_if_missing(&conn, "mem_tree_chunks", "content_path", "TEXT")?;
add_column_if_missing(&conn, "mem_tree_chunks", "content_sha256", "TEXT")?;
// Phase MD-content (summaries): same pointer pattern for summary nodes.
// `content_path` is the relative path to the .md file under
// `<content_root>/summaries/...`. `content_sha256` is the SHA-256 hex
// of the body bytes only (front-matter excluded). Both nullable so
// legacy rows (from before this migration) read back with NULL — callers
// fall back to the `content` column for those rows.
add_column_if_missing(&conn, "mem_tree_summaries", "content_path", "TEXT")?;
add_column_if_missing(&conn, "mem_tree_summaries", "content_sha256", "TEXT")?;
f(&conn)
}
/// Return both `content_path` and `content_sha256` stored in SQLite for `chunk_id`.
///
/// Returns `Ok(None)` if the chunk does not exist or has no content_path recorded yet.
pub fn get_chunk_content_pointers(
config: &Config,
chunk_id: &str,
) -> Result<Option<(String, String)>> {
with_connection(config, |conn| {
let row = conn
.query_row(
"SELECT content_path, content_sha256 FROM mem_tree_chunks WHERE id = ?1",
params![chunk_id],
|r| {
let path: Option<String> = r.get(0)?;
let sha: Option<String> = r.get(1)?;
Ok((path, sha))
},
)
.optional()?;
Ok(row.and_then(|(p, s)| p.zip(s)))
})
}
/// Return the `content_path` stored in SQLite for `chunk_id`, if any.
pub fn get_chunk_content_path(config: &Config, chunk_id: &str) -> Result<Option<String>> {
with_connection(config, |conn| {
let row = conn
.query_row(
"SELECT content_path FROM mem_tree_chunks WHERE id = ?1",
params![chunk_id],
|r| r.get::<_, Option<String>>(0),
)
.optional()?
.flatten();
Ok(row)
})
}
/// Return both `content_path` and `content_sha256` stored in SQLite for `summary_id`.
///
/// Returns `Ok(None)` if the summary does not exist or has no content_path recorded yet
/// (legacy rows pre-MD-content migration).
pub fn get_summary_content_pointers(
config: &Config,
summary_id: &str,
) -> Result<Option<(String, String)>> {
with_connection(config, |conn| {
let row = conn
.query_row(
"SELECT content_path, content_sha256 FROM mem_tree_summaries WHERE id = ?1",
params![summary_id],
|r| {
let path: Option<String> = r.get(0)?;
let sha: Option<String> = r.get(1)?;
Ok((path, sha))
},
)
.optional()?;
Ok(row.and_then(|(p, s)| p.zip(s)))
})
}
/// List all summary rows that have a non-NULL `content_path`. Used by the
/// bin integrity checker.
pub fn list_summaries_with_content_path(config: &Config) -> Result<Vec<(String, String, String)>> {
with_connection(config, |conn| {
let mut stmt = conn.prepare(
"SELECT id, content_path, content_sha256
FROM mem_tree_summaries
WHERE content_path IS NOT NULL AND content_sha256 IS NOT NULL
AND deleted = 0",
)?;
let rows = stmt
.query_map([], |r| {
let id: String = r.get(0)?;
let path: String = r.get(1)?;
let sha: String = r.get(2)?;
Ok((id, path, sha))
})?
.collect::<rusqlite::Result<Vec<_>>>()
.context("Failed to list summaries with content_path")?;
Ok(rows)
})
}
fn normalized_limit(requested: Option<usize>) -> i64 {
let clamped = requested
.unwrap_or(DEFAULT_LIST_LIMIT)
+36
View File
@@ -1,6 +1,7 @@
use super::*;
use crate::openhuman::memory::tree::types::chunk_id;
use chrono::TimeZone;
use rusqlite::params;
use tempfile::TempDir;
fn test_config() -> (TempDir, Config) {
@@ -27,6 +28,7 @@ fn sample_chunk(source_id: &str, seq: u32, ts_ms: i64) -> Chunk {
token_count: 12,
seq_in_source: seq,
created_at: ts,
partial_message: false,
}
}
@@ -164,3 +166,37 @@ fn empty_batch_is_noop() {
assert_eq!(upsert_chunks(&cfg, &[]).unwrap(), 0);
assert_eq!(count_chunks(&cfg).unwrap(), 0);
}
#[test]
fn schema_has_content_path_and_content_sha256_columns() {
// Phase MD-content: verify that with_connection applies the additive
// migrations for the new pointer + hash columns on a fresh DB.
let (_tmp, cfg) = test_config();
with_connection(&cfg, |conn| {
let mut has_content_path = false;
let mut has_content_sha256 = false;
let mut stmt = conn.prepare("PRAGMA table_info(mem_tree_chunks)")?;
let names: Vec<String> = stmt
.query_map(params![], |row| row.get::<_, String>(1))?
.filter_map(|r| r.ok())
.collect();
for name in &names {
if name == "content_path" {
has_content_path = true;
}
if name == "content_sha256" {
has_content_sha256 = true;
}
}
assert!(
has_content_path,
"mem_tree_chunks must have content_path column after migration; found: {names:?}"
);
assert!(
has_content_sha256,
"mem_tree_chunks must have content_sha256 column after migration; found: {names:?}"
);
Ok(())
})
.unwrap();
}
+174 -34
View File
@@ -1,30 +1,52 @@
//! Topic-tree backfill — hydrate a freshly-materialised topic tree with
//! every historical leaf mentioning the entity (#709 Phase 3c).
//! recent leaves mentioning the entity (#709 Phase 3c).
//!
//! When the curator decides an entity has crossed the hotness threshold
//! for the first time, we create a fresh topic tree AND walk the
//! `mem_tree_entity_index` inverted index to append every prior leaf into
//! `mem_tree_entity_index` inverted index to append matching leaves into
//! its L0 buffer. Reusing `bucket_seal::append_leaf` means the cascade
//! fires automatically — a well-established entity may seal several
//! levels as soon as the tree is spawned.
//! fires automatically.
//!
//! ## Why bounded by hotness window
//!
//! Hotness uses a 30-day recency decay (see `topic_tree::hotness`). Leaves
//! older than 30 days contribute zero to current hotness, so by definition
//! they cannot be the reason a tree is spawning *now*. Including them
//! bloats the spawn latency, wastes summariser LLM calls, and amplifies
//! ancient signal that has already decayed away. We cap the backfill
//! window at [`BACKFILL_WINDOW_DAYS`] to align with the hotness math.
//!
//! Older content is still queryable through source-tree retrieval and the
//! entity index — it just doesn't get its own slot in the topic tree.
//!
//! Backfill is intentionally best-effort: missing chunks are skipped with
//! a warn log rather than failing the whole spawn, because Phase 3c is
//! additive — a partial topic tree is still useful.
use anyhow::{Context, Result};
use chrono::Utc;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::score::store::lookup_entity;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::summariser::Summariser;
use crate::openhuman::memory::tree::source_tree::types::Tree;
use crate::openhuman::memory::tree::store::get_chunk;
use crate::openhuman::memory::tree::util::redact::redact;
/// Max leaves to pull from the entity index during backfill. A hard cap
/// keeps initial spawn latency bounded even for very active entities.
const BACKFILL_LIMIT: usize = 500;
/// Backfill window in days — matches `topic_tree::hotness::recency_decay`'s
/// hard cliff. Leaves older than this contribute zero to current hotness
/// so they cannot have driven the spawn decision.
pub const BACKFILL_WINDOW_DAYS: i64 = 30;
const DAY_MS: i64 = 24 * 60 * 60 * 1_000;
/// Walk the entity index for `entity_id` and append every discovered leaf
/// to `tree`. Returns the number of leaves appended (NOT the number of
/// summaries sealed). Idempotent: `append_leaf` itself is a no-op when a
@@ -35,19 +57,65 @@ pub async fn backfill_topic_tree(
entity_id: &str,
summariser: &dyn Summariser,
) -> Result<usize> {
log::info!(
"[topic_tree::backfill] start entity_id={} tree_id={}",
backfill_topic_tree_at(
config,
tree,
entity_id,
tree.id
summariser,
Utc::now().timestamp_millis(),
)
.await
}
/// Deterministic variant — backfill against a caller-supplied `now_ms`
/// for the recency window. Used by tests so the 30-day cutoff doesn't
/// depend on the wall clock.
pub async fn backfill_topic_tree_at(
config: &Config,
tree: &Tree,
entity_id: &str,
summariser: &dyn Summariser,
now_ms: i64,
) -> Result<usize> {
let cutoff_ms = now_ms.saturating_sub(BACKFILL_WINDOW_DAYS.saturating_mul(DAY_MS));
log::info!(
"[topic_tree::backfill] start entity_id_hash={} tree_id={} window_days={} cutoff_ms={}",
redact(entity_id),
tree.id,
BACKFILL_WINDOW_DAYS,
cutoff_ms
);
let hits = lookup_entity(config, entity_id, Some(BACKFILL_LIMIT))
.with_context(|| format!("failed to lookup entity {entity_id}"))?;
.with_context(|| format!("failed to lookup entity {}", redact(entity_id)))?;
if hits.is_empty() {
log::debug!(
"[topic_tree::backfill] no entity-index hits for entity_id={} — empty backfill",
entity_id
"[topic_tree::backfill] no entity-index hits for entity_id_hash={} — empty backfill",
redact(entity_id)
);
return Ok(0);
}
// Drop hits older than the hotness recency window — see module docs.
let total_hits = hits.len();
let mut hits: Vec<_> = hits
.into_iter()
.filter(|h| h.timestamp_ms >= cutoff_ms)
.collect();
let dropped = total_hits - hits.len();
if dropped > 0 {
log::debug!(
"[topic_tree::backfill] dropped {dropped} hits older than {BACKFILL_WINDOW_DAYS}d \
for entity_id_hash={}",
redact(entity_id)
);
}
if hits.is_empty() {
log::debug!(
"[topic_tree::backfill] all entity-index hits fell outside the {BACKFILL_WINDOW_DAYS}d \
window for entity_id_hash={} empty backfill",
redact(entity_id)
);
return Ok(0);
}
@@ -55,7 +123,6 @@ pub async fn backfill_topic_tree(
// Sort by timestamp ASC so the buffer's `oldest_at` and the sealed
// summary's `time_range_start` reflect the true historical order, not
// the DESC ordering `lookup_entity` returns.
let mut hits = hits;
hits.sort_by_key(|h| h.timestamp_ms);
let mut appended = 0usize;
@@ -77,9 +144,9 @@ pub async fn backfill_topic_tree(
Some(c) => c,
None => {
log::warn!(
"[topic_tree::backfill] missing chunk {} for entity {} — skipping",
"[topic_tree::backfill] missing chunk {} for entity_id_hash={} — skipping",
hit.node_id,
entity_id
redact(entity_id)
);
continue;
}
@@ -95,20 +162,25 @@ pub async fn backfill_topic_tree(
score: hit.score,
};
append_leaf(config, tree, &leaf, summariser)
// Topic-tree backfill: empty labels for sealed summaries — the
// tree's scope already pins the canonical id, so cross-pollinating
// descendants' entities would noise the index. See LabelStrategy.
append_leaf(config, tree, &leaf, summariser, &LabelStrategy::Empty)
.await
.with_context(|| {
format!(
"backfill append_leaf failed tree_id={} chunk_id={}",
tree.id, chunk.id
"backfill append_leaf failed entity_id_hash={} tree_id={} chunk_id={}",
redact(entity_id),
tree.id,
chunk.id
)
})?;
appended += 1;
}
log::info!(
"[topic_tree::backfill] done entity_id={} tree_id={} appended={}",
entity_id,
"[topic_tree::backfill] done entity_id_hash={} tree_id={} appended={}",
redact(entity_id),
tree.id,
appended
);
@@ -158,6 +230,7 @@ mod tests {
token_count: tokens,
seq_in_source: seq,
created_at: ts,
partial_message: false,
}
}
@@ -172,6 +245,11 @@ mod tests {
}
}
/// Deterministic "now" used by the windowed-backfill tests: 1 hour
/// after the latest seeded leaf so all three sit inside the 30-day
/// cutoff. Lets us keep the legacy 2023-era timestamps unchanged.
const TEST_NOW_MS: i64 = 1_700_000_020_000 + 3_600_000;
#[tokio::test]
async fn backfill_appends_all_entity_leaves() {
let (_tmp, cfg) = test_config();
@@ -212,9 +290,15 @@ mod tests {
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
let summariser = InertSummariser::new();
let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
.await
.unwrap();
let n = backfill_topic_tree_at(
&cfg,
&tree,
"email:alice@example.com",
&summariser,
TEST_NOW_MS,
)
.await
.unwrap();
assert_eq!(n, 3);
// L0 buffer should hold all three leaves (combined tokens well
@@ -226,6 +310,38 @@ mod tests {
assert_eq!(buf.oldest_at.unwrap().timestamp_millis(), 1_700_000_000_000);
}
#[tokio::test]
async fn backfill_drops_leaves_older_than_window() {
let (_tmp, cfg) = test_config();
// c_old is 60d before TEST_NOW_MS — outside the 30d cutoff.
// c_new is 5d before TEST_NOW_MS — inside the window.
let old_ts = TEST_NOW_MS - 60 * DAY_MS;
let new_ts = TEST_NOW_MS - 5 * DAY_MS;
let c_old = mk_chunk("slack:#eng", 0, old_ts, 100);
let c_new = mk_chunk("slack:#eng", 1, new_ts, 100);
upsert_chunks(&cfg, &[c_old.clone(), c_new.clone()]).unwrap();
let e = sample_entity("email:alice@example.com", "alice@example.com");
index_entity(&cfg, &e, &c_old.id, "leaf", old_ts, Some("source:slack")).unwrap();
index_entity(&cfg, &e, &c_new.id, "leaf", new_ts, Some("source:slack")).unwrap();
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
let summariser = InertSummariser::new();
let n = backfill_topic_tree_at(
&cfg,
&tree,
"email:alice@example.com",
&summariser,
TEST_NOW_MS,
)
.await
.unwrap();
assert_eq!(n, 1, "only the in-window leaf should be appended");
let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap();
assert_eq!(buf.item_ids.len(), 1);
assert_eq!(buf.item_ids[0], c_new.id);
}
#[tokio::test]
async fn backfill_skips_missing_chunks_without_failing() {
let (_tmp, cfg) = test_config();
@@ -247,9 +363,15 @@ mod tests {
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
let summariser = InertSummariser::new();
let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
.await
.unwrap();
let n = backfill_topic_tree_at(
&cfg,
&tree,
"email:alice@example.com",
&summariser,
TEST_NOW_MS,
)
.await
.unwrap();
assert_eq!(n, 1, "only the existing chunk should be appended");
let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap();
assert_eq!(buf.item_ids.len(), 1);
@@ -273,12 +395,24 @@ mod tests {
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
let summariser = InertSummariser::new();
backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
.await
.unwrap();
backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
.await
.unwrap();
backfill_topic_tree_at(
&cfg,
&tree,
"email:alice@example.com",
&summariser,
TEST_NOW_MS,
)
.await
.unwrap();
backfill_topic_tree_at(
&cfg,
&tree,
"email:alice@example.com",
&summariser,
TEST_NOW_MS,
)
.await
.unwrap();
// append_leaf is idempotent so the buffer still has exactly one row.
let buf = src_store::get_buffer(&cfg, &tree.id, 0).unwrap();
assert_eq!(buf.item_ids.len(), 1);
@@ -300,9 +434,15 @@ mod tests {
.unwrap();
let tree = get_or_create_topic_tree(&cfg, "email:alice@example.com").unwrap();
let summariser = InertSummariser::new();
let n = backfill_topic_tree(&cfg, &tree, "email:alice@example.com", &summariser)
.await
.unwrap();
let n = backfill_topic_tree_at(
&cfg,
&tree,
"email:alice@example.com",
&summariser,
TEST_NOW_MS,
)
.await
.unwrap();
assert_eq!(n, 0);
}
}
@@ -186,7 +186,11 @@ mod tests {
}
fn seed_leaf_for_entity(cfg: &Config, entity_id: &str, source_tree: &str, seq: u32) {
let ts_ms = 1_700_000_000_000 + (seq as i64) * 1_000;
// Use a "now-anchored" timestamp so backfill's 30-day window
// (see topic_tree::backfill::BACKFILL_WINDOW_DAYS) always
// includes these seeded leaves. Spread by seq to keep ordering
// deterministic.
let ts_ms = Utc::now().timestamp_millis() - (seq as i64) * 1_000;
let ts = Utc.timestamp_millis_opt(ts_ms).unwrap();
let c = Chunk {
id: chunk_id(SourceKind::Chat, source_tree, seq, "test-content"),
@@ -203,6 +207,7 @@ mod tests {
token_count: 50,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
upsert_chunks(cfg, &[c.clone()]).unwrap();
let e = CanonicalEntity {
+27 -17
View File
@@ -21,7 +21,9 @@
use anyhow::Result;
use crate::openhuman::config::Config;
use crate::openhuman::memory::tree::source_tree::bucket_seal::{append_leaf, LeafRef};
use crate::openhuman::memory::tree::source_tree::bucket_seal::{
append_leaf, LabelStrategy, LeafRef,
};
use crate::openhuman::memory::tree::source_tree::store as src_store;
use crate::openhuman::memory::tree::source_tree::summariser::Summariser;
use crate::openhuman::memory::tree::source_tree::types::{TreeKind, TreeStatus};
@@ -90,7 +92,16 @@ async fn route_one_entity(
entities: vec![entity_id.to_string()],
..leaf.clone()
};
append_leaf(config, &tree, &topic_leaf, summariser).await?;
// Topic-tree seals leave entities/topics empty: the tree's
// scope already pins the canonical id this tree represents.
append_leaf(
config,
&tree,
&topic_leaf,
summariser,
&LabelStrategy::Empty,
)
.await?;
} else {
log::debug!(
"[topic_tree::routing] skip archived topic tree id={} entity={}",
@@ -162,6 +173,7 @@ mod tests {
token_count: tokens,
seq_in_source: seq,
created_at: ts,
partial_message: false,
};
let id = c.id.clone();
upsert_chunks(cfg, &[c]).unwrap();
@@ -295,9 +307,15 @@ mod tests {
crate::openhuman::memory::tree::topic_tree::types::TOPIC_RECHECK_EVERY - 1;
crate::openhuman::memory::tree::topic_tree::store::upsert(&cfg, &counters).unwrap();
// Seed a leaf in slack and gmail referencing Alice.
let c1 = persist_chunk(&cfg, "slack:#eng", 0, 1_700_000_000_000, 100);
let c2 = persist_chunk(&cfg, "gmail:alice", 0, 1_700_000_010_000, 100);
// Seed leaves in slack and gmail referencing Alice. Anchor the
// timestamps to "now" so the 30-day backfill window
// (topic_tree::backfill::BACKFILL_WINDOW_DAYS) covers them.
let now_ms = Utc::now().timestamp_millis();
let ts_c1 = now_ms - 20_000;
let ts_c2 = now_ms - 10_000;
let ts_c3 = now_ms;
let c1 = persist_chunk(&cfg, "slack:#eng", 0, ts_c1, 100);
let c2 = persist_chunk(&cfg, "gmail:alice", 0, ts_c2, 100);
let e = CanonicalEntity {
canonical_id: entity_id.into(),
kind: EntityKind::Email,
@@ -306,24 +324,16 @@ mod tests {
span_end: entity_id.len() as u32,
score: 1.0,
};
index_entity(&cfg, &e, &c1, "leaf", 1_700_000_000_000, Some("slack:#eng")).unwrap();
index_entity(
&cfg,
&e,
&c2,
"leaf",
1_700_000_010_000,
Some("gmail:alice"),
)
.unwrap();
index_entity(&cfg, &e, &c1, "leaf", ts_c1, Some("slack:#eng")).unwrap();
index_entity(&cfg, &e, &c2, "leaf", ts_c2, Some("gmail:alice")).unwrap();
// A third leaf arrives — should both fan out to (future) topic tree
// and push the curator over the recheck cadence, materialising it.
let c3 = persist_chunk(&cfg, "slack:#eng", 1, 1_700_000_020_000, 100);
let c3 = persist_chunk(&cfg, "slack:#eng", 1, ts_c3, 100);
let leaf = LeafRef {
chunk_id: c3.clone(),
token_count: 100,
timestamp: Utc.timestamp_millis_opt(1_700_000_020_000).unwrap(),
timestamp: Utc.timestamp_millis_opt(ts_c3).unwrap(),
content: "new mention".into(),
entities: vec![entity_id.into()],
topics: vec![],
+6
View File
@@ -232,6 +232,12 @@ pub struct Chunk {
/// When this chunk was persisted to the local store.
#[serde(with = "chrono::serde::ts_milliseconds")]
pub created_at: DateTime<Utc>,
/// True when this chunk is a sub-split of a single logical unit (e.g. a
/// chat message or email body that exceeded `max_tokens`). The full logical
/// unit was split into multiple pieces; each piece carries this flag so
/// downstream scorers can lower its weight relative to whole-unit chunks.
#[serde(default)]
pub partial_message: bool,
}
/// Deterministic chunk id.
+3
View File
@@ -0,0 +1,3 @@
//! Shared utility helpers for the memory-tree subsystem.
pub mod redact;
+136
View File
@@ -0,0 +1,136 @@
//! PII redaction helpers for log output.
//!
//! Per project rule (CLAUDE.md): "Never log secrets or full PII."
//! After the participant-bucketing change introduced in the MD-content PR,
//! source_ids and content_paths can embed full email addresses, so any log
//! line that prints them needs to redact.
use sha2::{Digest, Sha256};
/// Redact a string by hashing it to 8 hex chars. Stable across runs for the
/// same input — safe to grep for in logs when debugging with the raw value
/// available externally.
///
/// Use for source_ids, entity_ids, content_paths and similar PII-bearing
/// strings in log output.
pub fn redact(s: &str) -> String {
let mut h = Sha256::new();
h.update(s.as_bytes());
let d = h.finalize();
format!("{:08x}", u32::from_be_bytes([d[0], d[1], d[2], d[3]]))
}
/// Redact a URL/endpoint by stripping path, query, fragment and credentials,
/// keeping only the host (and port if present).
///
/// Examples:
/// - `"http://localhost:11434/api/chat"` → `"localhost:11434"`
/// - `"https://user:pass@example.com/foo?q=1"` → `"example.com"`
/// - `"ollama://host:1234"` → `"host:1234"`
///
/// Does not pull in a URL-parsing crate; uses cheap string splitting which is
/// sufficient for the endpoint-config strings this codebase passes around.
pub fn redact_endpoint(url: &str) -> String {
// Strip scheme (everything before "://").
let after_scheme = url.split_once("://").map(|(_, r)| r).unwrap_or(url);
// Take only the authority (everything up to the first '/', '?', or '#') so
// any '@' in the path / query (e.g. `?email=foo@bar`) doesn't get treated
// as a userinfo separator.
let authority = after_scheme
.split(['/', '?', '#'])
.next()
.unwrap_or(after_scheme);
// Within the authority, the LAST '@' separates userinfo from host:port.
// (RFC 3986: userinfo may itself contain '@' — split-on-first would
// truncate the host. Use rsplit so `user:p@ss@example.com` extracts
// `example.com` correctly.)
let host_port = authority
.rsplit_once('@')
.map(|(_, r)| r)
.unwrap_or(authority);
host_port.to_string()
}
#[cfg(test)]
mod tests {
use super::*;
// ── redact ───────────────────────────────────────────────────────────────
#[test]
fn redact_returns_eight_hex_chars() {
let r = redact("alice@example.com");
assert_eq!(r.len(), 8, "must be 8 hex chars; got {r:?}");
assert!(r.chars().all(|c| c.is_ascii_hexdigit()), "must be hex");
}
#[test]
fn redact_is_stable_across_calls() {
assert_eq!(redact("alice@example.com"), redact("alice@example.com"));
}
#[test]
fn redact_is_different_for_different_inputs() {
assert_ne!(redact("alice@example.com"), redact("bob@example.com"));
}
#[test]
fn redact_empty_string_does_not_panic() {
let r = redact("");
assert_eq!(r.len(), 8);
}
// ── redact_endpoint ─────────────────────────────────────────────────────
#[test]
fn redact_endpoint_strips_path_and_query() {
assert_eq!(
redact_endpoint("http://localhost:11434/api/chat"),
"localhost:11434"
);
}
#[test]
fn redact_endpoint_strips_credentials() {
assert_eq!(
redact_endpoint("https://user:pass@example.com/foo"),
"example.com"
);
}
#[test]
fn redact_endpoint_no_scheme_passthrough() {
// No "://" present — treat the whole string as host/path; still strip path.
assert_eq!(redact_endpoint("localhost:11434/api"), "localhost:11434");
}
#[test]
fn redact_endpoint_just_host() {
assert_eq!(redact_endpoint("https://example.com"), "example.com");
}
#[test]
fn redact_endpoint_strips_fragment() {
assert_eq!(redact_endpoint("http://host:9090/path#frag"), "host:9090");
}
#[test]
fn redact_endpoint_strips_query() {
assert_eq!(redact_endpoint("http://host/path?q=1"), "host");
}
#[test]
fn redact_endpoint_empty_does_not_panic() {
let r = redact_endpoint("");
// Empty input: no scheme, no host — returns empty string.
assert_eq!(r, "");
}
#[test]
fn redact_endpoint_ollama_style() {
assert_eq!(
redact_endpoint("http://127.0.0.1:11434/v1/chat/completions"),
"127.0.0.1:11434"
);
}
}
+1
View File
@@ -825,6 +825,7 @@ async fn json_rpc_memory_tree_end_to_end() {
"openhuman.memory_tree_ingest".to_string(),
"openhuman.memory_tree_list_chunks".to_string(),
"openhuman.memory_tree_get_chunk".to_string(),
"openhuman.memory_tree_trigger_digest".to_string(),
];
assert_eq!(controllers.len(), expected_methods.len());
for method in &expected_methods {