diff --git a/src/openhuman/agent/harness/archivist.rs b/src/openhuman/agent/harness/archivist.rs index 7e49ddc51..931277205 100644 --- a/src/openhuman/agent/harness/archivist.rs +++ b/src/openhuman/agent/harness/archivist.rs @@ -4,10 +4,17 @@ //! After each turn, the Archivist: //! 1. Inserts the turn into the FTS5 episodic table. //! 2. Manages conversation segments (boundary detection + lifecycle). -//! 3. On segment close: extracts events (heuristic) and updates user profile. +//! 3. On segment close: produces an LLM recap (soft-fallback to heuristic), +//! embeds the recap, extracts events, and updates user profile. //! 4. Extracts simple lessons from tool failures. -//! 5. (Phase 1 / #566) Pipes the turn into the memory tree as `conversations:agent` -//! when `config.learning.chat_to_tree_enabled` is true. +//! 5. (Phase 2 / #566) At segment close/flush, ingests the segment's raw prose +//! turns (user + assistant; tool-call JSON stripped) into the memory tree as +//! `source_id = "conversations:agent"` when +//! `config.learning.chat_to_tree_enabled` is true. The leaf is RAW PROSE — +//! the LLM recap is NEVER fed into the tree (evidence-vs-interpretation +//! policy). Each leaf carries episodic provenance stamped in `source_ref`. +//! 6. `flush_open_segment` force-closes the trailing open segment at session +//! end so the last segment always gets a recap + embedding + tree ingest. use crate::openhuman::agent::hooks::{PostTurnHook, TurnContext}; use crate::openhuman::config::Config; @@ -18,7 +25,16 @@ use crate::openhuman::memory::store::segments::{ self, BoundaryConfig, BoundaryDecision, ConversationSegment, }; use crate::openhuman::memory::tree::canonicalize::chat::{ChatBatch, ChatMessage}; +use crate::openhuman::memory::tree::chat::{ChatConsumer, ChatProvider}; use crate::openhuman::memory::tree::ingest; +use crate::openhuman::memory::tree::score::embed::{build_embedder_from_config, Embedder}; +use crate::openhuman::memory::tree::tree_source::summariser::llm::{ + LlmSummariser, LlmSummariserConfig, +}; +use crate::openhuman::memory::tree::tree_source::summariser::{ + Summariser, SummaryContext, SummaryInput, +}; +use crate::openhuman::memory::tree::tree_source::types::TreeKind; use async_trait::async_trait; use parking_lot::Mutex; use rusqlite::Connection; @@ -29,6 +45,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; /// Background Archivist that indexes turns into FTS5 episodic memory /// and manages conversation segmentation. +/// +/// Produces an LLM recap + embedding for each closed segment and flushes +/// the trailing open segment at session end. pub struct ArchivistHook { /// SQLite connection shared with UnifiedMemory. conn: Option>>, @@ -36,46 +55,138 @@ pub struct ArchivistHook { enabled: bool, /// Boundary detection configuration. boundary_config: BoundaryConfig, - /// Optional runtime config — used to gate the tree-ingest path. + /// Optional runtime config — used to gate the tree-ingest path and to + /// build the LLM chat provider + embedder. /// /// When `None`, the tree-ingest path is skipped. Set via /// [`ArchivistHook::with_config`] on the production path. config: Option, + /// Optional LLM provider for segment recap. When `None`, the + /// fallback heuristic summary is used instead. + chat_provider: Option>, + /// Optional embedder for segment recap vectors. When `None`, embedding + /// is skipped (segment is still summarised). + embedder: Option>, } impl ArchivistHook { /// Create an Archivist hook with a shared SQLite connection. /// - /// Tree-ingest is disabled by default; call [`Self::with_config`] to - /// enable it on the production path. + /// LLM recap and embedding are disabled by default; call + /// [`Self::with_config`] on the production path to wire them in. pub fn new(conn: Arc>, enabled: bool) -> Self { Self { conn: Some(conn), enabled, boundary_config: BoundaryConfig::default(), config: None, + chat_provider: None, + embedder: None, } } - /// Attach runtime config so the archivist can gate the tree-ingest path. + /// Attach runtime config so the archivist can gate the tree-ingest path + /// and build its LLM chat provider + embedder from config. /// - /// When `config.learning.chat_to_tree_enabled` is `true`, each completed - /// turn is also piped into the memory tree as `source="conversations:agent"`. + /// When `config.learning.chat_to_tree_enabled` is `true`, each closed + /// segment's raw prose turns are ingested into the memory tree as + /// `source_id="conversations:agent"` (one batch per segment, not per turn). + /// The chat provider is built via `build_chat_provider(config, Summarise)`; + /// the embedder via `build_embedder_from_config(config)`. Both are + /// soft-fallback: if construction fails, the fields stay `None` and the + /// archivist falls back to heuristic summary / no embedding. pub fn with_config(mut self, config: Config) -> Self { + // Build the LLM chat provider for segment recap. + let chat_provider: Option> = + match crate::openhuman::memory::tree::chat::build_chat_provider( + &config, + ChatConsumer::Summarise, + ) { + Ok(p) => { + tracing::debug!("[archivist] segment recap provider={} registered", p.name()); + Some(p) + } + Err(e) => { + tracing::warn!( + "[archivist] failed to build chat provider for recap (will use fallback): {e}" + ); + None + } + }; + + // Build the embedder for segment recap vectors. + let embedder: Option> = match build_embedder_from_config(&config) { + Ok(e) => { + tracing::debug!("[archivist] segment embed provider={} registered", e.name()); + Some(Arc::from(e)) + } + Err(e) => { + tracing::warn!( + "[archivist] failed to build embedder for segment recap (embedding skipped): {e}" + ); + None + } + }; + + self.chat_provider = chat_provider; + self.embedder = embedder; self.config = Some(config); self } - /// Create a disabled/no-op Archivist (when FTS5 is not enabled). + /// Create a disabled/no-op Archivist (when FTS5 is not available). pub fn disabled() -> Self { Self { conn: None, enabled: false, boundary_config: BoundaryConfig::default(), config: None, + chat_provider: None, + embedder: None, } } + /// Flush the currently-open segment for `session_id`, if any, by + /// force-closing it and running the same close path (recap + embed + + /// event extraction). This guarantees the trailing segment of a session + /// is always finalized even when no boundary-triggering turn arrives. + /// + /// Called at session end (see `Agent::spawn_session_memory_extraction` + /// in `session/turn.rs`). Safe to call multiple times — segment_close + /// is idempotent (only transitions `open → closed`). + pub async fn flush_open_segment(&self, session_id: &str) { + if !self.enabled { + return; + } + let Some(conn) = &self.conn else { + return; + }; + let now = Self::now_timestamp(); + tracing::debug!("[archivist] flush_open_segment: checking session={session_id}"); + let open_segment = match segments::open_segment_for_session(conn, session_id) { + Ok(seg) => seg, + Err(e) => { + tracing::warn!("[archivist] flush: failed to query open segment: {e}"); + return; + } + }; + let Some(segment) = open_segment else { + tracing::debug!("[archivist] flush: no open segment for session={session_id}"); + return; + }; + tracing::debug!( + "[archivist] flush: force-closing segment={} turn_count={}", + segment.segment_id, + segment.turn_count + ); + if let Err(e) = segments::segment_close(conn, &segment.segment_id, now) { + tracing::warn!("[archivist] flush: failed to close segment: {e}"); + return; + } + self.on_segment_closed(conn, &segment, session_id, now) + .await; + } + fn now_timestamp() -> f64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -85,18 +196,18 @@ impl ArchivistHook { /// Handle segment lifecycle for a new turn. /// - /// The close→extract→create path uses a SQLite transaction for the - /// close + create operations to ensure atomicity. Event extraction - /// runs between close and create (outside the transaction) because - /// it needs to re-acquire the connection lock via fts5 functions. - fn manage_segment( + /// Returns the closed segment (if any) so the caller can run + /// `on_segment_closed` asynchronously after this function returns. + /// Event extraction and recap run outside this function because they + /// are async and may re-acquire the connection lock. + fn manage_segment_sync( &self, conn: &Arc>, session_id: &str, timestamp: f64, user_message: &str, current_episodic_id: i64, - ) { + ) -> Option { let now = Self::now_timestamp(); // Check for an open segment for this session. @@ -104,7 +215,7 @@ impl ArchivistHook { Ok(seg) => seg, Err(e) => { tracing::warn!("[archivist] failed to query open segment: {e}"); - return; + return None; } }; @@ -121,6 +232,11 @@ impl ArchivistHook { match decision { BoundaryDecision::Continue => { + tracing::debug!( + "[archivist] segment={} continues (turn_count={})", + segment.segment_id, + segment.turn_count + ); if let Err(e) = segments::segment_append_turn( conn, &segment.segment_id, @@ -130,6 +246,7 @@ impl ArchivistHook { ) { tracing::warn!("[archivist] failed to append turn to segment: {e}"); } + None } BoundaryDecision::Boundary(reason) => { tracing::debug!( @@ -140,14 +257,9 @@ impl ArchivistHook { // Close the current segment. if let Err(e) = segments::segment_close(conn, &segment.segment_id, now) { tracing::warn!("[archivist] failed to close segment: {e}"); - return; + return None; } - // Extract events from the closed segment and update profile. - // This runs outside a transaction because it calls fts5 functions - // that re-acquire the connection lock. - self.on_segment_closed(conn, &segment, session_id, now); - // Create a new segment for the new topic. // The new segment starts at the current turn's episodic ID. let new_id = format!("seg-{}", uuid_v4()); @@ -162,12 +274,19 @@ impl ArchivistHook { ) { tracing::warn!("[archivist] failed to create new segment: {e}"); } + + // Return the closed segment so the caller can run + // on_segment_closed asynchronously. + Some(segment) } } } None => { // No open segment — create the first one using the current episodic ID. let segment_id = format!("seg-{}", uuid_v4()); + tracing::debug!( + "[archivist] creating first segment={segment_id} for session={session_id}" + ); if let Err(e) = segments::segment_create( conn, &segment_id, @@ -179,13 +298,20 @@ impl ArchivistHook { ) { tracing::warn!("[archivist] failed to create initial segment: {e}"); } + None } } } - /// Called when a segment is closed. Runs heuristic event extraction - /// and updates the user profile from extracted preferences/facts. - fn on_segment_closed( + /// Called when a segment is closed. + /// + /// Produces a segment recap (LLM if a chat provider is configured, + /// otherwise the heuristic fallback), embeds the recap, extracts + /// heuristic events, and updates the user profile. + /// + /// Soft-fallback contract (mirrors `LlmSummariser`): this function + /// never returns `Err`; all failures are logged and ignored. + async fn on_segment_closed( &self, conn: &Arc>, segment: &ConversationSegment, @@ -211,10 +337,14 @@ impl ArchivistHook { .collect(); if segment_entries.is_empty() { + tracing::debug!( + "[archivist] segment={} has no entries — skipping recap", + segment.segment_id + ); return; } - // Build segment text from user messages. + // Build segment text from user messages (for event extraction). let segment_text: String = segment_entries .iter() .filter(|e| e.role == "user") @@ -222,87 +352,158 @@ impl ArchivistHook { .collect::>() .join(". "); - if segment_text.is_empty() { - return; - } + // ── Segment recap (LLM or heuristic fallback) ──────────────────── + let (summary, _from_llm) = self + .summarize_entries(&segment_entries, &segment.segment_id, segment.turn_count) + .await; - // Generate a fallback summary from first and last content. - let first = segment_entries - .first() - .map(|e| e.content.as_str()) - .unwrap_or(""); - let last = segment_entries - .last() - .map(|e| e.content.as_str()) - .unwrap_or(first); - let summary = segments::fallback_summary(first, last, segment.turn_count); + // Persist the recap. if let Err(e) = segments::segment_set_summary(conn, &segment.segment_id, &summary, now) { tracing::warn!("[archivist] failed to set segment summary: {e}"); + } else { + tracing::debug!( + "[archivist] recap persisted segment={} summary_chars={}", + segment.segment_id, + summary.len() + ); } - // Extract events via heuristic patterns. - let extracted = events::extract_events_heuristic(&segment_text); - tracing::debug!( - "[archivist] extracted {} events from segment {}", - extracted.len(), - segment.segment_id - ); - - for (event_type, content) in &extracted { - let event_id = format!("evt-{}", uuid_v4()); - let event = EventRecord { - event_id, - segment_id: segment.segment_id.clone(), - session_id: session_id.to_string(), - namespace: segment.namespace.clone(), - event_type: event_type.clone(), - content: content.clone(), - subject: None, - timestamp_ref: None, - confidence: 0.6, - embedding: None, - source_turn_ids: None, - created_at: now, - }; - if let Err(e) = events::event_insert(conn, &event) { - tracing::warn!("[archivist] failed to insert event: {e}"); + // ── Finalize-time embedding ─────────────────────────────────────── + // Embed the recap only when the segment is being finalized (closed). + // Never embed per-turn or on an open segment — this is the single + // write point for segment_embeddings rows. + if let Some(ref embedder) = self.embedder { + let model_signature = embedder.name().to_string(); + tracing::debug!( + "[archivist] embedding recap segment={} model={}", + segment.segment_id, + model_signature + ); + match embedder.embed(&summary).await { + Ok(vec) => { + match segments::segment_embedding_upsert( + conn, + &segment.segment_id, + &model_signature, + &vec, + now, + ) { + Ok(()) => { + tracing::debug!( + "[archivist] embedding stored segment={} model={} dim={}", + segment.segment_id, + model_signature, + vec.len() + ); + } + Err(e) => { + tracing::warn!( + "[archivist] failed to persist segment embedding (non-fatal) segment={}: {e}", + segment.segment_id + ); + } + } + } + Err(e) => { + tracing::warn!( + "[archivist] embed call failed (non-fatal) segment={} model={}: {e}", + segment.segment_id, + model_signature + ); + } } + } else { + tracing::debug!( + "[archivist] no embedder — skipping segment embedding segment={}", + segment.segment_id + ); + } - // Update user profile from preference and fact events. - match event_type { - EventType::Preference => { - let key = extract_profile_key(content, "preference"); - let facet_id = format!("prf-{}", uuid_v4()); - if let Err(e) = profile::profile_upsert( - conn, - &facet_id, - &FacetType::Preference, - &key, - content, - 0.6, - Some(&segment.segment_id), - now, - ) { - tracing::warn!("[archivist] failed to upsert profile facet: {e}"); - } + // ── Heuristic event extraction ──────────────────────────────────── + if !segment_text.is_empty() { + let extracted = events::extract_events_heuristic(&segment_text); + tracing::debug!( + "[archivist] extracted {} events from segment {}", + extracted.len(), + segment.segment_id + ); + + for (event_type, content) in &extracted { + let event_id = format!("evt-{}", uuid_v4()); + let event = EventRecord { + event_id, + segment_id: segment.segment_id.clone(), + session_id: session_id.to_string(), + namespace: segment.namespace.clone(), + event_type: event_type.clone(), + content: content.clone(), + subject: None, + timestamp_ref: None, + confidence: 0.6, + embedding: None, + source_turn_ids: None, + created_at: now, + }; + if let Err(e) = events::event_insert(conn, &event) { + tracing::warn!("[archivist] failed to insert event: {e}"); } - EventType::Fact => { - let key = extract_profile_key(content, "fact"); - let facet_id = format!("prf-{}", uuid_v4()); - if let Err(e) = profile::profile_upsert( - conn, - &facet_id, - &FacetType::Context, - &key, - content, - 0.6, - Some(&segment.segment_id), - now, - ) { - tracing::warn!("[archivist] failed to upsert profile fact: {e}"); + + // Update user profile from preference and fact events. + match event_type { + EventType::Preference => { + let key = extract_profile_key(content, "preference"); + let facet_id = format!("prf-{}", uuid_v4()); + if let Err(e) = profile::profile_upsert( + conn, + &facet_id, + &FacetType::Preference, + &key, + content, + 0.6, + Some(&segment.segment_id), + now, + ) { + tracing::warn!("[archivist] failed to upsert profile facet: {e}"); + } } + EventType::Fact => { + let key = extract_profile_key(content, "fact"); + let facet_id = format!("prf-{}", uuid_v4()); + if let Err(e) = profile::profile_upsert( + conn, + &facet_id, + &FacetType::Context, + &key, + content, + 0.6, + Some(&segment.segment_id), + now, + ) { + tracing::warn!("[archivist] failed to upsert profile fact: {e}"); + } + } + _ => {} } - _ => {} + } + } + + // ── Phase 2: tree ingest at segment granularity ─────────────────── + // Gate: only when config is attached and chat_to_tree_enabled is true. + // Ingest the segment's raw prose turns (NOT the LLM recap) as one + // ChatBatch into the memory tree under `source_id="conversations:agent"`. + // Evidence-vs-interpretation: the tree must ingest raw prose and build + // its own summaries; feeding the recap would make the tree summarise + // a summary. Non-fatal: failures are logged and swallowed. + if let Some(ref cfg) = self.config { + if cfg.learning.chat_to_tree_enabled { + tracing::debug!( + "[archivist] piping segment into tree as conversations:agent \ + session={session_id} segment={} entries={}", + segment.segment_id, + segment_entries.len() + ); + self.pipe_segment_to_tree(cfg, segment, session_id, &segment_entries) + .await; } } } @@ -380,8 +581,11 @@ impl PostTurnHook for ArchivistHook { }, )?; - // Manage conversation segmentation. - self.manage_segment( + tracing::debug!("[archivist] episodic rows written: session={session_id}"); + + // Manage conversation segmentation (sync boundary detection + SQLite + // operations). Returns the just-closed segment when a boundary fired. + let closed_segment = self.manage_segment_sync( conn, session_id, timestamp, @@ -389,104 +593,367 @@ impl PostTurnHook for ArchivistHook { current_episodic_id, ); - // ── Phase 1 / #566: pipe turn into the memory tree ─────────────────── - // Gate: only when config is attached and chat_to_tree_enabled is true. - // Non-fatal: if tree-ingest fails, the episodic write already succeeded - // and the turn result is not affected. - if let Some(ref cfg) = self.config { - if cfg.learning.chat_to_tree_enabled { - tracing::debug!( - "[archivist] piping turn into tree as conversations:agent session={}", - session_id - ); - self.pipe_turn_to_tree(cfg, ctx, session_id, timestamp) - .await; - } + // Run async recap + embed + segment-tree ingest on the closed segment + // (if any). Per-turn tree ingest is intentionally absent — Phase 2 + // moves the tree write to segment granularity inside on_segment_closed. + if let Some(ref segment) = closed_segment { + let now = Self::now_timestamp(); + self.on_segment_closed(conn, segment, session_id, now).await; } - tracing::debug!("[archivist] turn indexed successfully"); + tracing::debug!("[archivist] turn indexed successfully: session={session_id}"); Ok(()) } } impl ArchivistHook { - /// Pipe the completed turn into the memory tree as `source="conversations:agent"`. + /// Shared summarize helper — the **single LLM summarizer** used by both + /// the finalize path (`on_segment_closed`) and the rolling-recap path + /// (`rolling_segment_recap`). /// - /// Tool-call JSON is stripped from the assistant text before ingest — only - /// the assistant's prose response flows into the tree (memory ingestion - /// policy: tool outputs must not reach memory). + /// Builds a prose corpus from `entries`, calls the `LlmSummariser` when a + /// `chat_provider` is configured, and falls back to the heuristic + /// `segments::fallback_summary` on any failure or when no provider is + /// wired in. Always returns a non-empty string. + /// + /// Invariants: + /// - NEVER mutates DB state (no `segment_set_summary`, no embedding). + /// - NEVER closes a segment. + /// - Safe to call on both open and closed segments. + /// Summarize a set of episodic entries into a recap string. + /// + /// Returns `(text, produced_by_llm)`. `produced_by_llm == false` means the + /// LLM was unavailable / failed / returned empty and `text` is the shallow + /// heuristic `fallback_summary` bookend stub. That stub is an acceptable + /// durable last-resort on the *finalize* path, but callers driving the + /// **live prompt** (rolling recap → compaction) must treat + /// `produced_by_llm == false` as "no real recap" and fall back to their + /// own strategy — the stub must never become live compaction text. + async fn summarize_entries( + &self, + entries: &[&EpisodicEntry], + segment_id: &str, + turn_count: i32, + ) -> (String, bool) { + if entries.is_empty() { + tracing::debug!( + "[archivist] summarize_entries: no entries for segment={segment_id} — \ + returning empty fallback" + ); + return (segments::fallback_summary("", "", turn_count), false); + } + + // Build a full prose corpus from ALL entries (user + assistant prose; + // tool-call JSON is already excluded because the archivist stores + // stripped prose in the `content` column). + let corpus_inputs: Vec = entries + .iter() + .filter(|e| !e.content.trim().is_empty()) + .map(|e| { + use crate::openhuman::memory::tree::types::approx_token_count; + let content = e.content.clone(); + let token_count = approx_token_count(&content); + let ts = chrono::DateTime::from_timestamp(e.timestamp as i64, 0) + .unwrap_or_else(chrono::Utc::now); + SummaryInput { + id: format!("{}-{}", e.role, e.timestamp as u64), + content, + token_count, + entities: Vec::new(), + topics: Vec::new(), + time_range_start: ts, + time_range_end: ts, + score: 0.5, + } + }) + .collect(); + + let summary_ctx = SummaryContext { + tree_id: segment_id, + tree_kind: TreeKind::Source, + target_level: 0, + token_budget: 2_000, + }; + + let first = entries.first().map(|e| e.content.as_str()).unwrap_or(""); + let last = entries.last().map(|e| e.content.as_str()).unwrap_or(first); + + if let Some(ref provider) = self.chat_provider { + let cfg = LlmSummariserConfig { + model: provider.name().to_string(), + structured_facet_extraction: false, + }; + let summariser = LlmSummariser::new(cfg, Arc::clone(provider)); + tracing::debug!( + "[archivist] summarize_entries: LLM recap segment={segment_id} \ + provider={} entries={}", + provider.name(), + entries.len() + ); + match summariser.summarise(&corpus_inputs, &summary_ctx).await { + Ok(output) if !output.content.is_empty() => { + tracing::debug!( + "[archivist] summarize_entries: LLM recap ok segment={segment_id} \ + chars={}", + output.content.len() + ); + (output.content, true) + } + Ok(_) => { + tracing::debug!( + "[archivist] summarize_entries: LLM returned empty — \ + heuristic fallback segment={segment_id}" + ); + (segments::fallback_summary(first, last, turn_count), false) + } + Err(e) => { + tracing::warn!( + "[archivist] summarize_entries: LLM recap failed (non-fatal) \ + segment={segment_id}: {e} — heuristic fallback" + ); + (segments::fallback_summary(first, last, turn_count), false) + } + } + } else { + tracing::debug!( + "[archivist] summarize_entries: no chat provider — \ + heuristic fallback segment={segment_id}" + ); + (segments::fallback_summary(first, last, turn_count), false) + } + } + + /// Produce a rolling recap of the **currently-open** segment for + /// `session_id` WITHOUT closing it, writing `segment_set_summary`, or + /// embedding. + /// + /// This is the Phase 1.5 "one summarizer" entry point. Both + /// `on_segment_closed` (finalize) and this function delegate to the same + /// [`Self::summarize_entries`] helper so the same LLM path is used in both + /// cases. The distinction is purely in what happens *after* the summary + /// string is produced: + /// + /// - **Finalize** (`on_segment_closed`): persists the summary via + /// `segment_set_summary`, embeds it, extracts events, pipes tree ingest. + /// - **Rolling** (this function): returns the summary string and does + /// nothing else — segment stays open, DB is untouched. + /// + /// Returns `None` when: + /// - The archivist is disabled or has no connection. + /// - There is no open segment for `session_id`. + /// - The open segment has no episodic entries. + /// - No real LLM recap was produced (LLM unavailable / failed / empty, so + /// only the heuristic bookend stub is available). The shallow stub is + /// deliberately NOT used as live compaction text. + /// + /// Callers must treat `None` as "recap unavailable" and fall back to + /// their own compaction strategy (e.g. `ProviderSummarizer`). + pub async fn rolling_segment_recap(&self, session_id: &str) -> Option { + if !self.enabled { + tracing::debug!( + "[archivist] rolling_segment_recap: archivist disabled \ + session={session_id} — returning None" + ); + return None; + } + let conn = self.conn.as_ref()?; + + // Find the currently-open segment for this session. + let open_segment = match segments::open_segment_for_session(conn, session_id) { + Ok(Some(seg)) => seg, + Ok(None) => { + tracing::debug!( + "[archivist] rolling_segment_recap: no open segment for \ + session={session_id} — returning None" + ); + return None; + } + Err(e) => { + tracing::warn!( + "[archivist] rolling_segment_recap: failed to query open segment \ + session={session_id}: {e} — returning None" + ); + return None; + } + }; + + // Gather the episodic entries for this session so far. + let all_entries = fts5::episodic_session_entries(conn, session_id).unwrap_or_default(); + + // Keep only entries within the open segment's time window (start → + // now, inclusive). An open segment has `end_timestamp = None`. + let segment_entries: Vec<&EpisodicEntry> = all_entries + .iter() + .filter(|e| e.timestamp >= open_segment.start_timestamp) + .collect(); + + if segment_entries.is_empty() { + tracing::debug!( + "[archivist] rolling_segment_recap: no entries in open segment={} \ + session={session_id} — returning None", + open_segment.segment_id + ); + return None; + } + + tracing::debug!( + "[archivist] rolling_segment_recap: summarizing open segment={} \ + entries={} session={session_id}", + open_segment.segment_id, + segment_entries.len() + ); + + let (recap, from_llm) = self + .summarize_entries( + &segment_entries, + &open_segment.segment_id, + open_segment.turn_count, + ) + .await; + + if !from_llm { + tracing::debug!( + "[archivist] rolling_segment_recap: only heuristic bookend stub \ + available (no real LLM recap) session={session_id} segment={} — \ + returning None so compaction falls back to ProviderSummarizer", + open_segment.segment_id + ); + return None; + } + + if recap.is_empty() { + tracing::debug!( + "[archivist] rolling_segment_recap: summarize_entries returned empty \ + session={session_id} segment={} — returning None", + open_segment.segment_id + ); + return None; + } + + tracing::debug!( + "[archivist] rolling_segment_recap: produced LLM recap chars={} \ + session={session_id} segment={}", + recap.len(), + open_segment.segment_id + ); + Some(recap) + } + + /// Pipe a closed segment's raw prose turns into the memory tree as + /// `source_id="conversations:agent"`. + /// + /// **Design contract (Phase 2):** + /// - ONE ingest per segment (not per turn) — the batch boundary is the + /// segment, so all turns land as a single ChatBatch. + /// - RAW PROSE only — the LLM recap (summary) is explicitly NOT ingested. + /// The tree must build its own summaries from evidence (raw turns); + /// feeding a summary-of-a-summary violates the evidence-vs-interpretation + /// policy. + /// - `source_id = "conversations:agent"` is a CONSTANT — a single shared + /// tree source for all agent chat sessions (never per-session or per-segment). + /// - Tool-call JSON is stripped from assistant entries so structured + /// payloads do not reach the tree (memory ingestion policy). + /// - Provenance is stamped on each `ChatMessage.source_ref` as + /// `agent://session/{session_id}/segment/{segment_id}#ep{start}-{end}` + /// so tree leaves can be traced back to episodic rows for drill-down and + /// deduplication. /// /// Failures are logged and swallowed; the episodic write is the source of /// truth. - async fn pipe_turn_to_tree( + async fn pipe_segment_to_tree( &self, config: &Config, - ctx: &TurnContext, + segment: &crate::openhuman::memory::store::segments::ConversationSegment, session_id: &str, - timestamp: f64, + entries: &[&fts5::EpisodicEntry], ) { use chrono::{TimeZone, Utc}; - // Build turn timestamps. The assistant message is offset by 1ms as in - // the episodic write so ordering is stable. - let user_ts = Utc - .timestamp_opt( - timestamp as i64, - ((timestamp.fract() * 1e9) as u32).min(999_999_999), - ) - .single() - .unwrap_or_else(Utc::now); - let asst_ts = Utc - .timestamp_opt( - (timestamp + 0.001) as i64, - (((timestamp + 0.001).fract() * 1e9) as u32).min(999_999_999), - ) - .single() - .unwrap_or(user_ts); + // Collect the episodic id span for provenance stamping. + // start_episodic_id comes from the segment record (set at creation); + // end_episodic_id is the latest turn id (may be None if only one turn). + let start_ep = segment.start_episodic_id; + let end_ep = segment.end_episodic_id.unwrap_or(start_ep); + let segment_id = &segment.segment_id; - // Strip tool-call JSON from the assistant response. - // Per memory ingestion policy, structured tool-call payloads must not - // flow into the tree — only the prose response is ingested. - let assistant_prose = strip_tool_calls_from_response(&ctx.assistant_response); + // The provenance URI embeds session + segment + episodic id span so + // tree leaves can be traced back to episodic_log rows without a + // full-text scan. + let provenance = + format!("agent://session/{session_id}/segment/{segment_id}#ep{start_ep}-{end_ep}"); + + // Build one ChatMessage per episodic entry (user + assistant; skip + // empties). Tool-call JSON is stripped from assistant content so only + // prose flows into the tree. + let messages: Vec = entries + .iter() + .filter_map(|e| { + let raw_text = if e.role == "assistant" { + strip_tool_calls_from_response(&e.content) + } else { + e.content.clone() + }; + let text = raw_text.trim().to_string(); + if text.is_empty() { + return None; + } + + // Convert the f64 Unix timestamp to DateTime. + let secs = e.timestamp as i64; + let nanos = ((e.timestamp.fract()) * 1e9) as u32; + let ts = Utc + .timestamp_opt(secs, nanos.min(999_999_999)) + .single() + .unwrap_or_else(Utc::now); + + Some(ChatMessage { + author: e.role.clone(), + timestamp: ts, + text, + source_ref: Some(provenance.clone()), + }) + }) + .collect(); + + if messages.is_empty() { + tracing::debug!( + "[archivist] pipe_segment_to_tree: no prose messages in segment={segment_id} — skipping" + ); + return; + } let batch = ChatBatch { platform: "agent".into(), + // channel_label carries session_id for human-readable context. channel_label: session_id.to_string(), - messages: vec![ - ChatMessage { - author: "user".into(), - timestamp: user_ts, - text: ctx.user_message.clone(), - source_ref: Some(format!("agent://session/{session_id}")), - }, - ChatMessage { - author: "assistant".into(), - timestamp: asst_ts, - text: assistant_prose, - source_ref: Some(format!("agent://session/{session_id}")), - }, - ], + messages, }; - // Use the session_id as the owner / identity tag. + // `source_id` is intentionally a CONSTANT — all agent sessions share + // one tree source so cross-session summarisation sees the full history. let source_id = "conversations:agent"; + // `owner` scopes the memory to the session; `tags` enable filtering. let owner = session_id; let tags = vec!["agent_chat".to_string()]; + tracing::debug!( + "[archivist] tree ingest start: source_id={source_id} session={session_id} \ + segment={segment_id} ep_span={start_ep}-{end_ep} provenance={provenance}" + ); + match ingest::ingest_chat(config, source_id, owner, tags, batch).await { Ok(result) => { tracing::debug!( - "[archivist] tree ingest ok: source_id={} chunks_written={} session={}", - source_id, - result.chunks_written, - session_id + "[archivist] tree ingest ok: source_id={source_id} \ + session={session_id} segment={segment_id} \ + chunks_written={} provenance={provenance}", + result.chunks_written ); } Err(e) => { tracing::warn!( - "[archivist] tree ingest failed (non-fatal): source_id={} session={} error={e}", - source_id, - session_id + "[archivist] tree ingest failed (non-fatal): source_id={source_id} \ + session={session_id} segment={segment_id} error={e}" ); } } @@ -530,6 +997,20 @@ fn strip_tool_calls_from_response(response: &str) -> String { } } + // Drop JSON / tool-use payload lines the XML strip above cannot catch + // (evidence-vs-interpretation policy: tool-call payloads must never reach + // tree ingest). + cleaned = cleaned + .lines() + .filter(|line| { + let l = line.trim(); + !(l.contains("\"tool_use\"") + || l.starts_with("{\"tool_calls\"") + || l.starts_with("\"tool_calls\"")) + }) + .collect::>() + .join("\n"); + // Trim and collapse runs of blank lines left by block removal. let trimmed = cleaned .lines() @@ -617,6 +1098,51 @@ fn rand_u32() -> u32 { hasher.finish() as u32 } +#[cfg(test)] +impl ArchivistHook { + /// Test-only constructor that injects a stub `ChatProvider` and `Embedder` + /// directly, bypassing `with_config`'s provider-build logic. Used by + /// Phase 1 tests to verify LLM recap and embedding paths without hitting + /// a real LLM or Ollama daemon. Exposed as `pub(crate)` so Phase 3 + /// STM recall integration tests can drive the full archivist path. + pub(crate) fn new_with_stubs( + conn: Arc>, + chat_provider: Arc, + embedder: Arc, + ) -> Self { + Self { + conn: Some(conn), + enabled: true, + boundary_config: BoundaryConfig::default(), + config: None, + chat_provider: Some(chat_provider), + embedder: Some(embedder), + } + } + + /// Test-only constructor that injects stub providers AND a `Config`, so the + /// Phase 2 segment-tree ingest path (gated by + /// `config.learning.chat_to_tree_enabled`) can be exercised hermetically. + /// + /// `config.learning.chat_to_tree_enabled` must be set to `true` by the caller + /// for the tree ingest to fire; the hook does NOT force it on. + pub(crate) fn new_with_stubs_and_config( + conn: Arc>, + chat_provider: Arc, + embedder: Arc, + config: Config, + ) -> Self { + Self { + conn: Some(conn), + enabled: true, + boundary_config: BoundaryConfig::default(), + config: Some(config), + chat_provider: Some(chat_provider), + embedder: Some(embedder), + } + } +} + #[cfg(test)] #[path = "archivist_tests.rs"] mod tests; diff --git a/src/openhuman/agent/harness/archivist_tests.rs b/src/openhuman/agent/harness/archivist_tests.rs index 044c0ae1d..775a734c2 100644 --- a/src/openhuman/agent/harness/archivist_tests.rs +++ b/src/openhuman/agent/harness/archivist_tests.rs @@ -1,6 +1,7 @@ use super::*; use crate::openhuman::agent::hooks::{ToolCallRecord, TurnContext}; use crate::openhuman::memory::store::{events as ev, fts5, segments as seg}; +use crate::openhuman::memory::tree::chat::ChatPrompt; fn setup_conn() -> Arc> { let conn = Connection::open_in_memory().unwrap(); @@ -233,3 +234,611 @@ async fn archivist_extracts_preference_event_on_boundary() { events.iter().map(|e| &e.content).collect::>() ); } + +// ── Phase 0: episodic_capture_enabled independent of learning.enabled ──────── + +/// When `learning.enabled = false` but `episodic_capture_enabled = true`, +/// the ArchivistHook (constructed directly, as builder.rs would produce) +/// must still write 2 episodic_log rows (user + assistant) and create/advance +/// a segment. This verifies the core contract: episodic capture runs +/// regardless of the learning inference stack toggle. +#[tokio::test] +async fn phase0_episodic_rows_and_segment_without_learning_enabled() { + let conn = setup_conn(); + // Simulate what builder.rs does when learning.enabled=false but + // episodic_capture_enabled=true: construct the hook directly with + // the SQLite conn, enabled=true. No config attached (no LLM recap + // or tree ingest — those are gated by learning.enabled / chat_to_tree_enabled). + let hook = ArchivistHook::new(conn.clone(), true); + + let session = "phase0-test-session"; + + hook.on_turn_complete(&TurnContext { + user_message: "Hello, what is Rust?".into(), + assistant_response: "Rust is a systems language.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Verify 2 episodic rows were written. + let entries = fts5::episodic_session_entries(&conn, session).unwrap(); + assert_eq!( + entries.len(), + 2, + "Expected 2 episodic rows (user + assistant), got {}", + entries.len() + ); + assert_eq!(entries[0].role, "user"); + assert_eq!(entries[1].role, "assistant"); + + // Verify a segment was created. + let open_seg = seg::open_segment_for_session(&conn, session) + .unwrap() + .expect("Expected an open segment after first turn"); + assert_eq!(open_seg.turn_count, 1); + + // Add a second turn to verify segment advances. + hook.on_turn_complete(&TurnContext { + user_message: "Tell me more about ownership.".into(), + assistant_response: "Ownership prevents data races.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 2, + }) + .await + .unwrap(); + + let entries2 = fts5::episodic_session_entries(&conn, session).unwrap(); + assert_eq!( + entries2.len(), + 4, + "Expected 4 episodic rows after 2 turns, got {}", + entries2.len() + ); + let open_seg2 = seg::open_segment_for_session(&conn, session) + .unwrap() + .expect("Expected an open segment after 2 turns"); + assert_eq!( + open_seg2.turn_count, 2, + "Segment should have 2 turns, got {}", + open_seg2.turn_count + ); +} + +// ── Phase 1: LLM recap + finalize-time embedding ───────────────────────────── + +/// Stub ChatProvider that returns a fixed recap string without hitting +/// any real LLM, so the test is hermetic. +struct StubChatProvider; + +#[async_trait::async_trait] +impl crate::openhuman::memory::tree::chat::ChatProvider for StubChatProvider { + fn name(&self) -> &str { + "stub:test" + } + + async fn chat_for_json(&self, _prompt: &ChatPrompt) -> anyhow::Result { + Ok("stub recap: discussed Rust ownership model".to_string()) + } + + async fn chat_for_text(&self, _prompt: &ChatPrompt) -> anyhow::Result { + Ok("stub recap: discussed Rust ownership model".to_string()) + } +} + +/// Stub Embedder that returns a fixed unit vector without hitting Ollama. +struct StubEmbedder; + +#[async_trait::async_trait] +impl crate::openhuman::memory::tree::score::embed::Embedder for StubEmbedder { + fn name(&self) -> &'static str { + "stub-embedder-v1" + } + + async fn embed(&self, _text: &str) -> anyhow::Result> { + // Return a simple 4-dim unit vector. + Ok(vec![0.5_f32, 0.5, 0.5, 0.5]) + } +} + +/// Build an ArchivistHook with stub provider + embedder injected directly. +/// Uses the test-only `new_with_stubs` constructor to bypass `with_config`. +fn hook_with_stubs(conn: Arc>) -> ArchivistHook { + ArchivistHook::new_with_stubs(conn, Arc::new(StubChatProvider), Arc::new(StubEmbedder)) +} + +/// When a segment closes, the LLM chat provider recap is used (verified by +/// a non-empty segment summary) and an embedding row is written to +/// `segment_embeddings`. +#[tokio::test] +async fn phase1_llm_recap_and_embedding_on_segment_close() { + let conn = setup_conn(); + let hook = hook_with_stubs(conn.clone()); + + let session = "phase1-recap-test"; + + // Turn 1 — opens first segment. + hook.on_turn_complete(&TurnContext { + user_message: "Tell me about Rust ownership".into(), + assistant_response: "Rust's ownership model prevents data races.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Turn 2 — continues same segment. + hook.on_turn_complete(&TurnContext { + user_message: "What about the borrow checker?".into(), + assistant_response: "The borrow checker enforces ownership rules at compile time.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 2, + }) + .await + .unwrap(); + + // Turn 3 — topic change triggers a boundary → closes first segment → recap + embed fire. + hook.on_turn_complete(&TurnContext { + user_message: "Completely different topic: what is async/await in Python?".into(), + assistant_response: "Python asyncio enables concurrent programming.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 3, + }) + .await + .unwrap(); + + // Verify segments exist. + let segments = seg::segments_by_namespace(&conn, "global", 10).unwrap(); + assert!( + segments.len() >= 2, + "Expected at least 2 segments (closed + open), got {}", + segments.len() + ); + + // Find the closed segment (has a summary). + let closed = segments + .iter() + .find(|s| s.summary.as_ref().map(|s| !s.is_empty()).unwrap_or(false)); + assert!( + closed.is_some(), + "Expected at least one closed segment with a non-empty summary" + ); + + let closed_seg = closed.unwrap(); + let summary = closed_seg.summary.as_ref().unwrap(); + // The stub provider returns a fixed string — verify it was persisted. + assert!( + summary.contains("stub recap"), + "Expected summary to contain 'stub recap', got: {:?}", + summary + ); + + // Verify an embedding row was written for the closed segment. + let embedding = + seg::segment_embedding_get(&conn, &closed_seg.segment_id, "stub-embedder-v1").unwrap(); + assert!( + embedding.is_some(), + "Expected an embedding row for segment={} model=stub-embedder-v1", + closed_seg.segment_id + ); + let vec = embedding.unwrap(); + assert_eq!(vec.len(), 4, "Expected 4-dim vector from stub embedder"); + for v in &vec { + assert!( + (*v - 0.5_f32).abs() < 1e-4, + "Expected vector components ≈ 0.5, got {v}" + ); + } +} + +/// `flush_open_segment` must force-close the trailing open segment and +/// trigger recap + embedding even without a boundary-triggering turn. +#[tokio::test] +async fn phase1_flush_open_segment_finalizes_trailing_segment() { + let conn = setup_conn(); + let hook = hook_with_stubs(conn.clone()); + + let session = "phase1-flush-test"; + + // Write 2 turns — stays in one open segment (no topic boundary fires). + for i in 1..=2 { + hook.on_turn_complete(&TurnContext { + user_message: format!("Question about Rust turn {i}"), + assistant_response: format!("Answer about Rust turn {i}"), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: i, + }) + .await + .unwrap(); + } + + // Confirm the segment is still open (no boundary fired). + let open_seg_before = seg::open_segment_for_session(&conn, session).unwrap(); + assert!( + open_seg_before.is_some(), + "Expected an open segment before flush" + ); + + // Flush — should force-close, recap, and embed. + hook.flush_open_segment(session).await; + + // Segment should now be closed (no open segment for this session). + let open_seg_after = seg::open_segment_for_session(&conn, session).unwrap(); + assert!( + open_seg_after.is_none(), + "Expected no open segment after flush_open_segment" + ); + + // The formerly-open segment should now have a summary. + let segments = seg::segments_by_namespace(&conn, "global", 10).unwrap(); + let flushed = segments.iter().find(|s| { + s.session_id == session && s.summary.as_ref().map(|s| !s.is_empty()).unwrap_or(false) + }); + assert!( + flushed.is_some(), + "Expected flushed segment to have a non-empty summary" + ); + + let seg_id = &flushed.unwrap().segment_id; + let embedding = seg::segment_embedding_get(&conn, seg_id, "stub-embedder-v1").unwrap(); + assert!( + embedding.is_some(), + "Expected embedding row for flushed segment={seg_id}" + ); +} + +// ── Phase 2: segment-granularity tree ingest ───────────────────────────────── +// +// The following tests verify: +// a) No per-turn tree write fires from on_turn_complete (no double-write). +// b) Exactly ONE tree ingest fires when a segment closes (not N per turn). +// c) The ingested batch contains all the segment's raw prose turns. +// d) The `source_id` is the constant "conversations:agent". +// e) Each leaf message carries session/segment/episodic-span provenance. +// f) The ingested content is raw prose, NOT the LLM recap. +// g) flush_open_segment also triggers tree ingest. + +use crate::openhuman::config::Config; +use crate::openhuman::memory::tree::store::{count_chunks, list_chunks, ListChunksQuery}; +use tempfile::TempDir; + +/// Build a Config that points at a temp workspace, suitable for tree-ingest tests. +/// The memory_tree DB and content dir are created under `tmp.path()`. +fn test_config_with_tree() -> (TempDir, Config) { + let tmp = TempDir::new().unwrap(); + let mut cfg = Config::default(); + cfg.workspace_dir = tmp.path().to_path_buf(); + // Disable embedding so ingest doesn't fail trying to contact Ollama. + cfg.memory_tree.embedding_endpoint = None; + cfg.memory_tree.embedding_model = None; + cfg.memory_tree.embedding_strict = false; + // Ensure the tree ingest gate is on. + cfg.learning.chat_to_tree_enabled = true; + (tmp, cfg) +} + +/// Build a hook that has both stub providers AND a real-enough Config wired in, +/// so the Phase 2 tree ingest path is exercised hermetically. +fn hook_with_stubs_and_tree_config(conn: Arc>, cfg: Config) -> ArchivistHook { + ArchivistHook::new_with_stubs_and_config( + conn, + Arc::new(StubChatProvider), + Arc::new(StubEmbedder), + cfg, + ) +} + +/// After a single turn (no segment boundary), the tree must have ZERO chunks — +/// the per-turn pipe_turn_to_tree path no longer exists. +#[tokio::test] +async fn phase2_no_per_turn_tree_write() { + let conn = setup_conn(); + let (_tmp, cfg) = test_config_with_tree(); + let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone()); + + let session = "phase2-no-per-turn"; + + // Single turn — no segment close fires, so no tree ingest should happen. + hook.on_turn_complete(&TurnContext { + user_message: "What is Rust?".into(), + assistant_response: "Rust is a systems programming language.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Segment is still open (no boundary fired) — tree must have 0 chunks. + let open_seg = seg::open_segment_for_session(&conn, session).unwrap(); + assert!( + open_seg.is_some(), + "Expected an open segment (no boundary should have fired)" + ); + + let chunk_count = count_chunks(&cfg).unwrap(); + assert_eq!( + chunk_count, 0, + "Expected 0 tree chunks after a single turn (no segment close): \ + per-turn tree write must not exist (Phase 2)" + ); +} + +/// When a segment closes (boundary triggered), exactly ONE tree ingest fires +/// for that segment containing all its turns — not one ingest per turn. +#[tokio::test] +async fn phase2_exactly_one_tree_ingest_per_segment_close() { + let conn = setup_conn(); + let (_tmp, cfg) = test_config_with_tree(); + let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone()); + + let session = "phase2-one-ingest"; + + // Turn 1 — opens first segment. + hook.on_turn_complete(&TurnContext { + user_message: "Tell me about Rust ownership".into(), + assistant_response: "Rust ownership prevents memory bugs.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Turn 2 — stays in same segment. + hook.on_turn_complete(&TurnContext { + user_message: "What about the borrow checker?".into(), + assistant_response: "The borrow checker enforces ownership at compile time.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 2, + }) + .await + .unwrap(); + + // No tree write yet — segment still open. + let pre_close_chunks = count_chunks(&cfg).unwrap(); + assert_eq!( + pre_close_chunks, 0, + "Expected 0 tree chunks before any segment close; got {pre_close_chunks}" + ); + + // Turn 3 — topic change triggers boundary → closes first segment → tree ingest fires. + hook.on_turn_complete(&TurnContext { + user_message: "Switching to a completely different topic: tell me about Python asyncio." + .into(), + assistant_response: "Python asyncio enables concurrent coroutines.".into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 3, + }) + .await + .unwrap(); + + // Segment closed → exactly one ingest for the closed segment (containing turns 1+2). + // The ingest packs the messages into one or more chunks (greedy packing), + // but chunks_written >= 1 confirms ingest happened. + let post_close_chunks = count_chunks(&cfg).unwrap(); + assert!( + post_close_chunks >= 1, + "Expected ≥ 1 tree chunk after segment close; got {post_close_chunks}" + ); + + // List the chunks and check they come from the constant source_id. + let chunks = list_chunks( + &cfg, + &ListChunksQuery { + source_id: Some("conversations:agent".to_string()), + ..Default::default() + }, + ) + .unwrap(); + assert!( + !chunks.is_empty(), + "Expected chunks under source_id='conversations:agent'" + ); +} + +/// The ingested leaf messages must carry the episodic-provenance `source_ref` +/// in the expected format: +/// `agent://session/{session_id}/segment/{segment_id}#ep{start}-{end}`. +/// +/// Also verifies that `source_id` is the constant `"conversations:agent"`. +#[tokio::test] +async fn phase2_provenance_stamped_on_leaf_and_source_id_is_constant() { + let conn = setup_conn(); + let (_tmp, cfg) = test_config_with_tree(); + let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone()); + + let session = "phase2-provenance"; + + // Two turns in the first segment. + for i in 1..=2 { + hook.on_turn_complete(&TurnContext { + user_message: format!("Ownership question {i}"), + assistant_response: format!("Ownership answer {i}"), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: i, + }) + .await + .unwrap(); + } + + // Force a segment close via flush_open_segment. + hook.flush_open_segment(session).await; + + // Retrieve the closed segment to extract its ID. + let all_segs = seg::segments_by_namespace(&conn, "global", 10).unwrap(); + let closed = all_segs + .iter() + .find(|s| { + s.session_id == session + && s.status != crate::openhuman::memory::store::segments::SegmentStatus::Open + }) + .expect("Expected a closed segment after flush"); + + let segment_id = &closed.segment_id; + let start_ep = closed.start_episodic_id; + let end_ep = closed.end_episodic_id.unwrap_or(start_ep); + + // Chunks should be present. + let chunks = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); + assert!( + !chunks.is_empty(), + "Expected tree chunks after flush_open_segment" + ); + + // source_id must be the constant — never per-session or per-segment. + for chunk in &chunks { + assert_eq!( + chunk.metadata.source_id, "conversations:agent", + "source_id must be the constant 'conversations:agent', got: {}", + chunk.metadata.source_id + ); + } + + // The source_ref on at least one chunk must contain the provenance pattern. + let expected_provenance = + format!("agent://session/{session}/segment/{segment_id}#ep{start_ep}-{end_ep}"); + let has_provenance = chunks.iter().any(|chunk| { + chunk + .metadata + .source_ref + .as_ref() + .map(|r| { + r.value + .contains(&format!("agent://session/{session}/segment/{segment_id}")) + }) + .unwrap_or(false) + }); + assert!( + has_provenance, + "Expected at least one chunk with source_ref containing provenance pattern \ + '{expected_provenance}'; found: {:?}", + chunks + .iter() + .map(|c| c.metadata.source_ref.as_ref().map(|r| r.value.as_str())) + .collect::>() + ); +} + +/// The ingested content must be the raw prose turns (user + assistant text), +/// NOT equal to the LLM recap text. The recap lives only in the STM segment +/// layer; the tree must ingest raw evidence so it can build its own summaries. +#[tokio::test] +async fn phase2_ingested_content_is_raw_prose_not_recap() { + let conn = setup_conn(); + let (_tmp, cfg) = test_config_with_tree(); + let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone()); + + let session = "phase2-raw-prose"; + + // The stub recap always returns "stub recap: discussed Rust ownership model". + // The raw user messages contain very different text. + let user_msg = "My specific question about lifetimes in Rust code"; + let asst_msg = "Lifetimes annotate how long references are valid in memory"; + + hook.on_turn_complete(&TurnContext { + user_message: user_msg.into(), + assistant_response: asst_msg.into(), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Flush to close the segment and trigger tree ingest. + hook.flush_open_segment(session).await; + + let chunks = list_chunks(&cfg, &ListChunksQuery::default()).unwrap(); + assert!( + !chunks.is_empty(), + "Expected tree chunks after flush_open_segment" + ); + + // The stub recap text must NOT appear in any chunk body. + let stub_recap_text = "stub recap: discussed Rust ownership model"; + for chunk in &chunks { + assert!( + !chunk.content.contains(stub_recap_text), + "Chunk content must NOT contain the recap text (evidence-vs-interpretation policy). \ + Found recap text in chunk id={}: {:?}", + chunk.id, + &chunk.content[..chunk.content.len().min(200)] + ); + } + + // The raw prose text MUST appear in at least one chunk. + let has_user_prose = chunks.iter().any(|c| c.content.contains("lifetimes")); + assert!( + has_user_prose, + "Expected at least one chunk body to contain raw prose from the turn \ + (keyword 'lifetimes'); found: {:?}", + chunks + .iter() + .map(|c| &c.content[..c.content.len().min(100)]) + .collect::>() + ); +} + +/// `flush_open_segment` must also trigger the tree ingest for the trailing +/// open segment (same as on_segment_closed at a topic boundary). +#[tokio::test] +async fn phase2_flush_also_triggers_tree_ingest() { + let conn = setup_conn(); + let (_tmp, cfg) = test_config_with_tree(); + let hook = hook_with_stubs_and_tree_config(conn.clone(), cfg.clone()); + + let session = "phase2-flush-tree"; + + // Two turns — no boundary fires, segment stays open. + for i in 1..=2 { + hook.on_turn_complete(&TurnContext { + user_message: format!("Rust borrowing question {i}"), + assistant_response: format!("Borrowing answer {i}"), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: i, + }) + .await + .unwrap(); + } + + // Confirm no tree chunks yet (segment still open). + let before = count_chunks(&cfg).unwrap(); + assert_eq!( + before, 0, + "Expected 0 tree chunks before flush; got {before}" + ); + + // Flush should close the segment and trigger tree ingest. + hook.flush_open_segment(session).await; + + let after = count_chunks(&cfg).unwrap(); + assert!( + after >= 1, + "Expected ≥ 1 tree chunk after flush_open_segment triggers segment ingest; got {after}" + ); +} diff --git a/src/openhuman/agent/harness/session/builder.rs b/src/openhuman/agent/harness/session/builder.rs index 7fbed4f37..55da8a666 100644 --- a/src/openhuman/agent/harness/session/builder.rs +++ b/src/openhuman/agent/harness/session/builder.rs @@ -17,7 +17,7 @@ use crate::openhuman::agent::host_runtime; use crate::openhuman::agent::memory_loader::{DefaultMemoryLoader, MemoryLoader}; use crate::openhuman::config::{Config, ContextConfig}; use crate::openhuman::context::prompt::SystemPromptBuilder; -use crate::openhuman::context::{ContextManager, ProviderSummarizer}; +use crate::openhuman::context::{ContextManager, ProviderSummarizer, SegmentRecapSummarizer}; use crate::openhuman::inference::provider::{self, Provider}; use crate::openhuman::memory::{self, Memory}; use crate::openhuman::security::SecurityPolicy; @@ -86,6 +86,8 @@ impl AgentBuilder { omit_profile: None, omit_memory_md: None, payload_summarizer: None, + archivist_hook: None, + unified_compaction_enabled: true, } } @@ -313,6 +315,39 @@ impl AgentBuilder { self } + /// Attach the production [`ArchivistHook`] instance so the session + /// turn loop can call [`ArchivistHook::flush_open_segment`] at + /// session-wind-down time, guaranteeing the trailing open segment is + /// always finalized with an LLM recap + embedding. + /// + /// Set from `build_session_agent_inner` when + /// `config.learning.episodic_capture_enabled` is `true` and a + /// SQLite connection is available. Callers that construct an `Agent` + /// directly (tests, CLI) can leave this `None` — flush is a no-op + /// when the hook is absent. + pub fn archivist_hook( + mut self, + hook: Option>, + ) -> Self { + self.archivist_hook = hook; + self + } + + /// Phase 1.5 — gate the unified compaction path. + /// + /// When `true` (the default) and an archivist hook is wired in via + /// [`Self::archivist_hook`], the session's `ContextManager` summarizer is + /// wrapped with a [`SegmentRecapSummarizer`] that routes autocompaction + /// through the archivist's rolling recap (one LLM summarizer, soft-fallback + /// to [`ProviderSummarizer`] when the recap is unavailable). + /// + /// When `false` the `ProviderSummarizer` is used directly and Phase 1.5 is + /// completely absent from the hot path — behaviour is identical to today's. + pub fn unified_compaction_enabled(mut self, enabled: bool) -> Self { + self.unified_compaction_enabled = enabled; + self + } + /// Validates the configuration and constructs a new `Agent` instance. /// /// This method is responsible for wiring together the provided components, @@ -377,7 +412,52 @@ impl AgentBuilder { // summarizer — every concern that touches "what's in the // model's context window" routes through this single handle. let context_config = self.context_config.unwrap_or_default(); - let summarizer = Arc::new(ProviderSummarizer::new(provider.clone())); + + // Phase 1.5 — unified compaction. + // + // When `unified_compaction_enabled` is true AND an archivist hook + // is wired in, wrap the inner `ProviderSummarizer` with a + // `SegmentRecapSummarizer`. The outer type: + // 1. Tries the rolling segment recap from the open segment. + // 2. Falls back to the inner `ProviderSummarizer` if unavailable. + // + // With the flag off OR no archivist, the plain `ProviderSummarizer` + // is used and Phase 1.5 is completely absent from the hot path + // — behaviour is identical to Phase 1. + let inner_summarizer: Arc = + Arc::new(ProviderSummarizer::new(provider.clone())); + let session_id_for_recap = self + .event_session_id + .clone() + .unwrap_or_else(|| "standalone".to_string()); + let summarizer: Arc = + if self.unified_compaction_enabled { + if let Some(ref archivist) = self.archivist_hook { + log::debug!( + "[agent::builder] unified_compaction_enabled=true — \ + wrapping summarizer with SegmentRecapSummarizer \ + session_id={session_id_for_recap}" + ); + Arc::new(SegmentRecapSummarizer::new( + Arc::clone(archivist), + session_id_for_recap, + inner_summarizer, + )) + } else { + log::debug!( + "[agent::builder] unified_compaction_enabled=true but \ + no archivist hook — using ProviderSummarizer" + ); + inner_summarizer + } + } else { + log::debug!( + "[agent::builder] unified_compaction_enabled=false — \ + using ProviderSummarizer (Phase 1.5 disabled)" + ); + inner_summarizer + }; + let context = ContextManager::new( &context_config, summarizer, @@ -462,6 +542,7 @@ impl AgentBuilder { omit_memory_md: self.omit_memory_md.unwrap_or(true), payload_summarizer: self.payload_summarizer, last_seen_integrations_hash: 0, + archivist_hook: self.archivist_hook, synthesized_tool_names: std::collections::HashSet::new(), }) } @@ -1015,6 +1096,46 @@ impl Agent { } } + // ── ArchivistHook — register independently of learning.enabled ────── + // + // Episodic capture (FTS5 index, segment lifecycle, LLM recap, embedding) + // is the system-of-record for chat turns and must stay active even when + // the inference stack (`reflection`, `stability_detector`) is disabled. + // Gated only on `config.learning.episodic_capture_enabled` (default: true) + // and on the memory backend exposing a SQLite connection. + let archivist_hook_arc: Option< + Arc, + > = if config.learning.episodic_capture_enabled { + match memory.sqlite_conn() { + Some(conn) => { + let hook = Arc::new( + crate::openhuman::agent::harness::archivist::ArchivistHook::new(conn, true) + .with_config(config.clone()), + ); + post_turn_hooks + .push(Arc::clone(&hook) + as Arc); + log::info!( + "[archivist] episodic capture hook registered (learning.enabled={})", + config.learning.enabled + ); + Some(hook) + } + None => { + log::warn!( + "[archivist] no SQLite connection available from memory backend — \ + episodic capture disabled" + ); + None + } + } + } else { + log::info!( + "[archivist] episodic_capture_enabled=false — archivist hook not registered" + ); + None + }; + // Resolve the per-agent delegation tool set and visible-tool // whitelist from the target definition (when we have one) or // fall back to the orchestrator's synthesis path. @@ -1347,6 +1468,8 @@ impl Agent { if let Some(ps) = payload_summarizer { builder = builder.payload_summarizer(ps); } + builder = builder.archivist_hook(archivist_hook_arc); + builder = builder.unified_compaction_enabled(config.learning.unified_compaction_enabled); builder.build() } } diff --git a/src/openhuman/agent/harness/session/turn.rs b/src/openhuman/agent/harness/session/turn.rs index ad1fe5d63..cdf85ee81 100644 --- a/src/openhuman/agent/harness/session/turn.rs +++ b/src/openhuman/agent/harness/session/turn.rs @@ -267,6 +267,12 @@ impl Agent { // error. The timestamp is bumped on every successful `load` (even // when the digest is empty) so an empty workspace doesn't get // re-queried every turn. + // + // Gate STM preemptive recall on the session turn index, independent + // of tree-prefetch success/failure. (Previously keyed off + // `last_tree_prefetch_at.is_none()`, which stays `None` when tree + // prefetch fails — re-firing STM recall on every later turn.) + let is_first_turn_for_stm = self.context.stats().session_memory_current_turn == 0; let now = std::time::Instant::now(); let context = if crate::openhuman::agent::tree_loader::should_prefetch( self.last_tree_prefetch_at, @@ -309,6 +315,70 @@ impl Agent { context }; + // ── Phase 3 STM preemptive recall ──────────────────────────── + // On the very first turn only, assemble a bounded cross-thread + // context block from the FTS5 episodic arm (keyword match) and the + // segment-embedding arm (cosine similarity). The block rides on the + // user message (NOT the system prompt) to keep the KV-cache prefix + // stable, exactly like the tree-context injection above. + // + // Gate: `learning.stm_recall_enabled` must be true AND this must + // be the first turn (STM is snapshot-frozen at session start). + // Failure is non-fatal — bare `context` passes through untouched. + let context = if is_first_turn_for_stm { + // Load config to check the gate. Use a cached load (cheap). + let stm_enabled = crate::openhuman::config::rpc::load_config_with_timeout() + .await + .map(|cfg| cfg.learning.stm_recall_enabled) + .unwrap_or(true); // default: enabled + + if stm_enabled { + if let Some(conn) = self.memory.sqlite_conn() { + use crate::openhuman::memory::stm_recall::recall::{stm_recall, StmRecallOpts}; + let opts = StmRecallOpts { + exclude_session: &self.event_session_id, + query: if user_message.trim().is_empty() { + None + } else { + Some(user_message) + }, + model_signature: None, + }; + match stm_recall(&conn, &opts, None) { + Ok(block) if !block.is_empty() => { + let stm_md = block.render(); + log::info!( + "[stm_recall] preemptive block injected: {} items, ~{} chars, fts5_candidates={}, dropped_dedup={}", + block.items.len(), + stm_md.chars().count(), + block.fts5_candidates, + block.dropped_dedup + ); + format!("{stm_md}{context}") + } + Ok(_) => { + log::debug!( + "[stm_recall] preemptive recall: no cross-thread context found" + ); + context + } + Err(e) => { + log::warn!("[stm_recall] preemptive recall failed (non-fatal): {e}"); + context + } + } + } else { + log::debug!("[stm_recall] preemptive recall skipped — no SQLite connection on memory backend"); + context + } + } else { + log::debug!("[stm_recall] preemptive recall skipped — stm_recall_enabled=false"); + context + } + } else { + context + }; + let enriched = if context.is_empty() { log::info!("[agent] no memory context found — using raw user message"); self.last_memory_context = None; @@ -904,7 +974,7 @@ impl Agent { // later), which is the right amount of retry behaviour for a // librarian task that's idempotent across reruns. if result.is_ok() && self.context.should_extract_session_memory() { - self.spawn_session_memory_extraction(); + self.spawn_session_memory_extraction().await; // Sibling pipeline (#1399): heuristic transcript ingestion // turns the just-written transcript into durable // conversational memory + reflections so a brand-new chat @@ -1687,7 +1757,29 @@ impl Agent { /// Gated by [`context_pipeline::SessionMemoryState::should_extract`] /// — see its docs for the threshold invariants. Safe to call from /// inside `turn()` after the turn body has settled. - pub(super) fn spawn_session_memory_extraction(&mut self) { + pub(super) async fn spawn_session_memory_extraction(&mut self) { + // ── Flush the trailing open segment before the session winds down ── + // + // The ArchivistHook manages per-turn segment lifecycle but cannot + // force-close the *last* open segment because there is no explicit + // "session end" event in the turn loop. `spawn_session_memory_extraction` + // is the closest available signal: it fires when the context manager + // decides the session has accumulated enough material to archive. + // + // GUARANTEE: the flush is *awaited* here (not fire-and-forget) so + // the trailing segment always receives its recap + embedding + tree + // ingest before the function returns, even during runtime wind-down. + // This honours the doc-comment guarantee on `flush_open_segment` in + // `archivist.rs`. No deadlock risk: no mutex guard is held across + // this await point. + if let Some(ref archivist) = self.archivist_hook { + let session_id = self.event_session_id.clone(); + log::debug!( + "[archivist] awaiting flush_open_segment for session={session_id} at session wind-down" + ); + archivist.flush_open_segment(&session_id).await; + } + let Some(registry) = harness::AgentDefinitionRegistry::global() else { log::debug!("[session_memory] registry not initialised — skipping extraction spawn"); return; diff --git a/src/openhuman/agent/harness/session/types.rs b/src/openhuman/agent/harness/session/types.rs index e44ad4a35..380a7b601 100644 --- a/src/openhuman/agent/harness/session/types.rs +++ b/src/openhuman/agent/harness/session/types.rs @@ -7,6 +7,7 @@ //! crate gaining field access. use crate::openhuman::agent::dispatcher::ToolDispatcher; +use crate::openhuman::agent::harness::archivist::ArchivistHook; use crate::openhuman::agent::hooks::PostTurnHook; use crate::openhuman::agent::memory_loader::MemoryLoader; use crate::openhuman::agent::progress::AgentProgress; @@ -161,6 +162,12 @@ pub struct Agent { /// dormant on session startup and only fires when integrations /// actually change mid-conversation. pub(super) last_seen_integrations_hash: u64, + /// Optional reference to the `ArchivistHook` registered in + /// `post_turn_hooks`. Kept separately so the turn loop can call + /// `flush_open_segment` at session-memory-extraction time (the + /// closest available signal to "session is ending") to finalize the + /// trailing open segment with an LLM recap + embedding. + pub(super) archivist_hook: Option>, /// Names of every tool currently in [`Agent::tools`] that was /// produced by [`crate::openhuman::tools::orchestrator_tools::collect_orchestrator_tools`] /// (i.e. `delegate_` skill tools and archetype-delegation @@ -220,6 +227,18 @@ pub struct AgentBuilder { /// to a `SubagentPayloadSummarizer` instance. pub(super) payload_summarizer: Option>, + /// Optional reference to the production `ArchivistHook`. Set when + /// `config.learning.episodic_capture_enabled` is true. Used to call + /// `flush_open_segment` at the closest available session-end signal. + pub(super) archivist_hook: Option>, + /// Phase 1.5 — when `true` AND `archivist_hook` is `Some`, the + /// `ContextManager`'s summarizer is wrapped with a + /// `SegmentRecapSummarizer` that routes compaction through the + /// archivist's rolling segment recap (one summarizer, soft-fallback). + /// When `false` (or archivist absent), the plain `ProviderSummarizer` + /// is used and Phase 1.5 is completely absent from the hot path. + /// Default: `true` (mirrors `LearningConfig::unified_compaction_enabled`). + pub(super) unified_compaction_enabled: bool, } impl Default for AgentBuilder { diff --git a/src/openhuman/config/schema/learning.rs b/src/openhuman/config/schema/learning.rs index 4a93b043d..7fe614c0b 100644 --- a/src/openhuman/config/schema/learning.rs +++ b/src/openhuman/config/schema/learning.rs @@ -76,6 +76,55 @@ pub struct LearningConfig { #[serde(default = "default_true")] pub stability_detector_enabled: bool, + /// Enable episodic capture (ArchivistHook) regardless of the master + /// `learning.enabled` toggle. + /// + /// Episodic capture is the system-of-record for chat turns + /// (`episodic_log` FTS5 table, conversation segmentation, segment + /// summaries with LLM recap, and segment embeddings). It must remain + /// active even when the inference stack + /// (reflection / stability-detector) is off. + /// + /// Default: `true`. Set to `false` to fully disable the Archivist. + /// + /// Override via `OPENHUMAN_LEARNING_EPISODIC_CAPTURE_ENABLED=0|1`. + #[serde(default = "default_true")] + pub episodic_capture_enabled: bool, + + /// Enable preemptive STM recall injection at session start and on-demand + /// `stm_recall_search` tool exposure. + /// + /// When enabled, a bounded cross-thread context block is assembled from + /// recent episodic entries (FTS5 keyword arm) and segment recaps (cosine + /// similarity arm) from OTHER sessions and injected into the first turn's + /// user message. The `stm_recall_search` tool is also registered in the + /// agent's tool list. + /// + /// Default: `true`. Set to `false` to fully disable STM recall. + /// + /// Override via `OPENHUMAN_LEARNING_STM_RECALL_ENABLED=0|1`. + #[serde(default = "default_true")] + pub stm_recall_enabled: bool, + + /// Use the rolling segment recap as the compaction text for evicted turns + /// (Phase 1.5 — unified compaction). + /// + /// When `true`, the [`ContextManager`]'s autocompaction summarizer is + /// wrapped with a `SegmentRecapSummarizer` that first tries to obtain the + /// current open segment's rolling recap from the `ArchivistHook` and uses + /// it as the replacement text for the evicted head. If the rolling recap + /// is unavailable (no archivist, no open segment, LLM failure, flag off) + /// the inner `ProviderSummarizer` runs as before — the live prompt is + /// NEVER left over-budget regardless of the recap path's health. + /// + /// Default: `true`. Set to `false` to revert to the standalone + /// `ProviderSummarizer` path (today's behaviour, Phase 1.5 completely + /// absent from the hot path). + /// + /// Override via `OPENHUMAN_LEARNING_UNIFIED_COMPACTION_ENABLED=0|1`. + #[serde(default = "default_true")] + pub unified_compaction_enabled: bool, + /// How often the periodic rebuild loop runs in seconds. Default: 1800 (30 minutes). #[serde(default = "default_rebuild_interval_secs")] pub rebuild_interval_secs: u64, @@ -111,6 +160,9 @@ impl Default for LearningConfig { chat_to_tree_enabled: default_true(), stability_detector_enabled: default_true(), rebuild_interval_secs: default_rebuild_interval_secs(), + episodic_capture_enabled: default_true(), + stm_recall_enabled: default_true(), + unified_compaction_enabled: default_true(), } } } diff --git a/src/openhuman/config/schema/load.rs b/src/openhuman/config/schema/load.rs index 70da827e5..1295f8ddc 100644 --- a/src/openhuman/config/schema/load.rs +++ b/src/openhuman/config/schema/load.rs @@ -1308,6 +1308,28 @@ impl Config { self.learning.min_turn_complexity = min; } } + if let Some(flag) = env.get("OPENHUMAN_LEARNING_EPISODIC_CAPTURE_ENABLED") { + if let Some(enabled) = + parse_env_bool("OPENHUMAN_LEARNING_EPISODIC_CAPTURE_ENABLED", flag.as_str()) + { + self.learning.episodic_capture_enabled = enabled; + } + } + if let Some(flag) = env.get("OPENHUMAN_LEARNING_STM_RECALL_ENABLED") { + if let Some(enabled) = + parse_env_bool("OPENHUMAN_LEARNING_STM_RECALL_ENABLED", flag.as_str()) + { + self.learning.stm_recall_enabled = enabled; + } + } + if let Some(flag) = env.get("OPENHUMAN_LEARNING_UNIFIED_COMPACTION_ENABLED") { + if let Some(enabled) = parse_env_bool( + "OPENHUMAN_LEARNING_UNIFIED_COMPACTION_ENABLED", + flag.as_str(), + ) { + self.learning.unified_compaction_enabled = enabled; + } + } // Phase 4 memory-tree embedding overrides (#710). Setting the env // var to an empty string explicitly clears the default — useful diff --git a/src/openhuman/context/mod.rs b/src/openhuman/context/mod.rs index b29d12551..95f17be6c 100644 --- a/src/openhuman/context/mod.rs +++ b/src/openhuman/context/mod.rs @@ -34,6 +34,7 @@ pub mod manager; pub mod microcompact; pub mod pipeline; pub mod prompt; +pub mod segment_recap_summarizer; pub mod session_memory; pub mod summarizer; pub mod tool_result_budget; @@ -49,6 +50,7 @@ pub use prompt::{ PromptSection, PromptTool, RuntimeSection, SafetySection, SystemPromptBuilder, ToolsSection, WorkspaceSection, }; +pub use segment_recap_summarizer::SegmentRecapSummarizer; pub use session_memory::{ SessionMemoryConfig, SessionMemoryState, ARCHIVIST_EXTRACTION_PROMPT, DEFAULT_MIN_TOKEN_GROWTH, DEFAULT_MIN_TOOL_CALLS, DEFAULT_MIN_TURNS_BETWEEN, diff --git a/src/openhuman/context/segment_recap_summarizer.rs b/src/openhuman/context/segment_recap_summarizer.rs new file mode 100644 index 000000000..4649e9603 --- /dev/null +++ b/src/openhuman/context/segment_recap_summarizer.rs @@ -0,0 +1,224 @@ +//! Phase 1.5 — segment-recap-backed compaction summarizer. +//! +//! [`SegmentRecapSummarizer`] wraps an inner [`Summarizer`] (normally a +//! [`super::ProviderSummarizer`]) and intercepts the autocompaction call to +//! try the rolling segment recap from the [`ArchivistHook`] first. +//! +//! ## Strategy +//! +//! When `AutocompactionRequested` fires, the manager calls +//! [`SegmentRecapSummarizer::summarize`]. That method: +//! +//! 1. Calls [`ArchivistHook::rolling_segment_recap`] for the current session. +//! 2. If a non-empty recap is returned, it replaces the evicted head of the +//! history with a single `[segment-recap]` system message containing that +//! text. The head/tail split follows the same "never break a tool-call +//! pair" rule as the inner summarizer. +//! 3. If the recap is `None`/empty (archivist absent, no open segment, LLM +//! fail, flag off) — soft-fallback: delegate to the inner summarizer +//! unchanged. The prompt is never left over-budget; the existing +//! compaction safety net always fires. +//! +//! ## Invariants (never violated by this code) +//! +//! - The `ArchivistHook` is only *read* through `rolling_segment_recap` — no +//! segment is closed, no `segment_set_summary` is written, no embedding is +//! produced. Those side-effects are finalize-only (Phase 1 owns them). +//! - Events/profile/tree derive from RAW episodic rows — this path never +//! touches those subsystems. +//! - On any error in the recap path, the inner summarizer runs. The +//! circuit-breaker logic in [`super::ContextManager`] is fed the result of +//! whichever path actually ran; the breaker sees a success if either path +//! succeeds. +//! - The history is either fully rewritten (success) or left completely +//! untouched (the inner summarizer also fails, in which case its `Err` is +//! returned — the breaker will nudge and eventually trip, preventing +//! infinite loops). + +use super::summarizer::{Summarizer, SummaryStats}; +use crate::openhuman::agent::harness::archivist::ArchivistHook; +use crate::openhuman::inference::provider::{ChatMessage, ConversationMessage}; +use anyhow::Result; +use async_trait::async_trait; +use std::sync::Arc; + +/// How many messages at the tail of the history are preserved verbatim when +/// the recap path fires (same default as [`super::ProviderSummarizer`] so the +/// two paths preserve the same recent context window). +const DEFAULT_KEEP_RECENT: usize = super::summarizer::DEFAULT_KEEP_RECENT; + +/// Compaction summarizer that prefers the rolling segment recap from the +/// archivist over a standalone LLM summarization call. +/// +/// Constructed by [`super::SessionBuilder`] when +/// `config.learning.unified_compaction_enabled` is `true` and an +/// `ArchivistHook` is wired in; otherwise the plain `ProviderSummarizer` +/// is used and this type is never instantiated. +pub struct SegmentRecapSummarizer { + /// The archivist that owns episodic state and can produce rolling recaps. + archivist: Arc, + /// Session ID needed to look up the open segment. + session_id: String, + /// Inner summarizer — the safety-net fallback (normally a + /// [`super::ProviderSummarizer`] wrapping the same provider the agent uses + /// for its normal turns). + inner: Arc, + /// How many tail messages to preserve verbatim when the recap path fires. + keep_recent: usize, +} + +impl SegmentRecapSummarizer { + /// Construct a new [`SegmentRecapSummarizer`]. + /// + /// * `archivist` — shared archivist handle (same Arc the agent holds). + /// * `session_id` — the session whose open segment will be recapped. + /// * `inner` — fallback summarizer; used when the recap is unavailable. + pub fn new( + archivist: Arc, + session_id: String, + inner: Arc, + ) -> Self { + Self { + archivist, + session_id, + inner, + keep_recent: DEFAULT_KEEP_RECENT, + } + } + + /// Override how many tail messages are preserved verbatim. Useful in + /// tests that want to exercise the recap path with short histories. + #[cfg(test)] + pub fn with_keep_recent(mut self, n: usize) -> Self { + self.keep_recent = n; + self + } +} + +#[async_trait] +impl Summarizer for SegmentRecapSummarizer { + async fn summarize( + &self, + history: &mut Vec, + model: &str, + ) -> Result { + // ── 1. Try the rolling segment recap ───────────────────────────── + let recap_opt = self.archivist.rolling_segment_recap(&self.session_id).await; + + match recap_opt { + Some(recap) if !recap.is_empty() => { + tracing::info!( + session_id = %self.session_id, + recap_chars = recap.len(), + history_len = history.len(), + keep_recent = self.keep_recent, + "[context::segment_recap] using rolling segment recap as compaction text" + ); + + let total = history.len(); + if total <= self.keep_recent { + tracing::debug!( + session_id = %self.session_id, + "[context::segment_recap] history below keep_recent — \ + nothing to compact (NoOp)" + ); + return Ok(SummaryStats::default()); + } + + // Use the same "never break a tool-call pair" split rule as + // ProviderSummarizer so the API invariant + // (AssistantToolCalls ↔ ToolResults) is preserved. + // Delegates to the canonical implementation in `summarizer` + // so the two compaction paths share a single definition. + let proposed_head = total - self.keep_recent; + let head_len = super::summarizer::snap_split_forward(history, proposed_head); + if head_len == 0 { + tracing::debug!( + session_id = %self.session_id, + "[context::segment_recap] split snapped to 0 — \ + falling back to inner summarizer" + ); + return self.inner.summarize(history, model).await; + } + + // Estimate bytes freed (same formula as ProviderSummarizer). + let approx_input_bytes: usize = history[..head_len] + .iter() + .map(|m| conversation_message_approx_bytes(m)) + .sum(); + + let summary_body = format!( + "[segment-recap] Summary of {head_len} earlier messages \ + (archivist rolling recap):\n\n{recap}" + ); + let summary_chars = summary_body.len(); + let approx_tokens_freed = (approx_input_bytes as u64) + .saturating_sub(summary_chars as u64) + .div_ceil(4); + + // Atomically rewrite the head in place — no partial mutation on + // failure because all failure paths returned early above. + let tail: Vec = history.drain(head_len..).collect(); + history.clear(); + history.push(ConversationMessage::Chat(ChatMessage::system(summary_body))); + history.extend(tail); + + tracing::info!( + session_id = %self.session_id, + messages_removed = head_len, + approx_tokens_freed, + summary_chars, + "[context::segment_recap] compaction via segment recap complete" + ); + + Ok(SummaryStats { + messages_removed: head_len, + approx_tokens_freed, + summary_chars, + }) + } + + // ── 2. Soft-fallback ───────────────────────────────────────── + // + // The rolling recap was unavailable (None, empty, LLM fail, no + // open segment, archivist disabled). Delegate to the inner + // summarizer so the prompt is NEVER left over-budget. This is the + // "always bounded" guarantee. + recap_result => { + let reason = match &recap_result { + None => "rolling_segment_recap returned None", + Some(s) if s.is_empty() => "rolling_segment_recap returned empty string", + _ => "unreachable", + }; + tracing::info!( + session_id = %self.session_id, + reason, + "[context::segment_recap] recap unavailable — \ + falling back to inner summarizer" + ); + self.inner.summarize(history, model).await + } + } + } +} + +/// Very rough byte count for a [`ConversationMessage`] — used only for the +/// "approx_tokens_freed" stat. Accuracy doesn't matter much (it's the same +/// rough accounting ProviderSummarizer uses). +fn conversation_message_approx_bytes(msg: &ConversationMessage) -> usize { + match msg { + ConversationMessage::Chat(m) => m.content.len(), + ConversationMessage::AssistantToolCalls { text, tool_calls } => { + text.as_deref().map_or(0, str::len) + + tool_calls + .iter() + .map(|tc| tc.arguments.len()) + .sum::() + } + ConversationMessage::ToolResults(results) => results.iter().map(|r| r.content.len()).sum(), + } +} + +#[cfg(test)] +#[path = "segment_recap_summarizer_tests.rs"] +mod tests; diff --git a/src/openhuman/context/segment_recap_summarizer_tests.rs b/src/openhuman/context/segment_recap_summarizer_tests.rs new file mode 100644 index 000000000..fb5cd14b2 --- /dev/null +++ b/src/openhuman/context/segment_recap_summarizer_tests.rs @@ -0,0 +1,562 @@ +//! Tests for Phase 1.5 — `SegmentRecapSummarizer`. +//! +//! Proves: +//! 1. Rolling recap summarizes the open segment WITHOUT closing it, writing +//! `segment_set_summary`, or producing an embedding row. +//! 2. When a rolling recap is available, the compaction replacement text +//! equals it (provenance/path), not a separately-generated summary. +//! 3. Soft-fallback: archivist absent / LLM stub failing / flag off → +//! compaction falls back to the inner summarizer; prompt stays bounded. +//! 4. Finalize path (Phase 1 `on_segment_closed`) still works unchanged +//! (recap persisted + embedded at close) — regression guard. + +use super::*; +use crate::openhuman::agent::harness::archivist::ArchivistHook; +use crate::openhuman::agent::hooks::{PostTurnHook as _, TurnContext}; +use crate::openhuman::context::summarizer::{Summarizer, SummaryStats}; +use crate::openhuman::inference::provider::{ChatMessage, ConversationMessage}; +use crate::openhuman::memory::store::{fts5, segments as seg}; +use crate::openhuman::memory::tree::chat::ChatPrompt; +use anyhow::Result; +use async_trait::async_trait; +use parking_lot::Mutex; +use rusqlite::Connection; +use std::sync::Arc; + +// ── Shared test infrastructure ──────────────────────────────────────────────── + +fn setup_conn() -> Arc> { + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(fts5::EPISODIC_INIT_SQL).unwrap(); + conn.execute_batch(seg::SEGMENTS_INIT_SQL).unwrap(); + conn.execute_batch(crate::openhuman::memory::store::events::EVENTS_INIT_SQL) + .unwrap(); + conn.execute_batch(crate::openhuman::memory::store::profile::PROFILE_INIT_SQL) + .unwrap(); + Arc::new(Mutex::new(conn)) +} + +/// Stub ChatProvider always returns a fixed recap string. +struct StubChatProvider; + +#[async_trait] +impl crate::openhuman::memory::tree::chat::ChatProvider for StubChatProvider { + fn name(&self) -> &str { + "stub:test" + } + async fn chat_for_json(&self, _prompt: &ChatPrompt) -> Result { + Ok("rolling recap: discussed memory safety in Rust".to_string()) + } + async fn chat_for_text(&self, _prompt: &ChatPrompt) -> Result { + Ok("rolling recap: discussed memory safety in Rust".to_string()) + } +} + +/// Stub ChatProvider that always fails — simulates LLM unavailability. +struct FailingChatProvider; + +#[async_trait] +impl crate::openhuman::memory::tree::chat::ChatProvider for FailingChatProvider { + fn name(&self) -> &str { + "stub:failing" + } + async fn chat_for_json(&self, _prompt: &ChatPrompt) -> Result { + anyhow::bail!("stub LLM unavailable") + } + async fn chat_for_text(&self, _prompt: &ChatPrompt) -> Result { + anyhow::bail!("stub LLM unavailable") + } +} + +/// Stub Embedder returns a fixed unit vector. +struct StubEmbedder; + +#[async_trait] +impl crate::openhuman::memory::tree::score::embed::Embedder for StubEmbedder { + fn name(&self) -> &'static str { + "stub-embedder-v1" + } + async fn embed(&self, _text: &str) -> Result> { + Ok(vec![0.5_f32, 0.5, 0.5, 0.5]) + } +} + +/// Inner mock summarizer that records call count and always succeeds. +struct RecordingSummarizer { + calls: std::sync::Mutex, +} + +impl RecordingSummarizer { + fn new() -> Arc { + Arc::new(Self { + calls: std::sync::Mutex::new(0), + }) + } + fn call_count(&self) -> usize { + *self.calls.lock().unwrap() + } +} + +#[async_trait] +impl Summarizer for RecordingSummarizer { + async fn summarize( + &self, + history: &mut Vec, + _model: &str, + ) -> Result { + *self.calls.lock().unwrap() += 1; + // Replace history with a single "inner fallback" message. + let removed = history.len(); + history.clear(); + history.push(ConversationMessage::Chat(ChatMessage::system( + "inner fallback summary", + ))); + Ok(SummaryStats { + messages_removed: removed, + approx_tokens_freed: 500, + summary_chars: 22, + }) + } +} + +fn user(s: &str) -> ConversationMessage { + ConversationMessage::Chat(ChatMessage::user(s)) +} + +// ── Test 1: rolling recap does NOT close the segment or write summary/embedding + +/// `rolling_segment_recap` must produce a non-empty string from the open +/// segment's entries without closing it, writing `segment_set_summary`, or +/// producing an embedding row. +#[tokio::test] +async fn rolling_recap_does_not_close_segment_or_write_summary_or_embedding() { + let conn = setup_conn(); + let hook = Arc::new(ArchivistHook::new_with_stubs( + conn.clone(), + Arc::new(StubChatProvider), + Arc::new(StubEmbedder), + )); + + let session = "p15-rolling-no-close"; + + // Write two turns into the open segment (no boundary fires). + for i in 1..=2u64 { + hook.on_turn_complete(&TurnContext { + user_message: format!("Turn {i} about Rust memory safety"), + assistant_response: format!("Answer {i} about ownership"), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: i as usize, + }) + .await + .unwrap(); + } + + // Verify the segment is still open before calling rolling_segment_recap. + let open_before = seg::open_segment_for_session(&conn, session).unwrap(); + assert!( + open_before.is_some(), + "Expected an open segment after 2 turns (no boundary should have fired)" + ); + + // Call rolling_segment_recap — must NOT close the segment or write DB state. + let recap = hook.rolling_segment_recap(session).await; + assert!( + recap.is_some(), + "Expected rolling_segment_recap to return Some for an open segment with entries" + ); + let recap_text = recap.unwrap(); + assert!( + !recap_text.is_empty(), + "Expected non-empty recap from rolling_segment_recap" + ); + + // Segment must STILL be open after the call. + let open_after = seg::open_segment_for_session(&conn, session).unwrap(); + assert!( + open_after.is_some(), + "Segment must still be open after rolling_segment_recap (no close side-effect)" + ); + + let seg_id = open_after.unwrap().segment_id; + + // No summary must have been written. + let all_segs = seg::segments_by_namespace(&conn, "global", 10).unwrap(); + let our_seg = all_segs.iter().find(|s| s.segment_id == seg_id).unwrap(); + assert!( + our_seg.summary.is_none() || our_seg.summary.as_deref() == Some(""), + "segment_set_summary must NOT have been called by rolling_segment_recap; \ + found: {:?}", + our_seg.summary + ); + + // No embedding must have been written. + let embedding = seg::segment_embedding_get(&conn, &seg_id, "stub-embedder-v1").unwrap(); + assert!( + embedding.is_none(), + "No embedding must be written by rolling_segment_recap \ + (finalize-only invariant); found an embedding for segment={seg_id}" + ); +} + +// ── Test 2: compaction uses recap text, not inner summarizer ───────────────── + +/// When the rolling recap is available, `SegmentRecapSummarizer::summarize` +/// must use it as the replacement text — the inner summarizer must NOT be +/// called and the history head must contain `[segment-recap]`. +#[tokio::test] +async fn compaction_uses_recap_text_not_inner_summarizer() { + let conn = setup_conn(); + let hook = Arc::new(ArchivistHook::new_with_stubs( + conn.clone(), + Arc::new(StubChatProvider), + Arc::new(StubEmbedder), + )); + + let session = "p15-compaction-recap"; + + // Write one turn so the open segment has entries. + hook.on_turn_complete(&TurnContext { + user_message: "Tell me about Rust ownership".into(), + assistant_response: "Ownership prevents data races.".into(), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + let inner = RecordingSummarizer::new(); + let recap_summ = SegmentRecapSummarizer::new( + Arc::clone(&hook), + session.to_string(), + inner.clone() as Arc, + ) + .with_keep_recent(1); // compact everything except the last message + + // Build a history that exceeds keep_recent so summarization fires. + let mut history = vec![ + user("message 1 (will be evicted)"), + user("message 2 (will be evicted)"), + user("message 3 — preserved tail"), + ]; + + let stats = recap_summ + .summarize(&mut history, "test-model") + .await + .expect("summarize must succeed"); + + // Inner summarizer must NOT have been called. + assert_eq!( + inner.call_count(), + 0, + "Inner summarizer must not be called when rolling recap is available" + ); + + // Stats must reflect a non-zero reduction. + assert!( + stats.messages_removed > 0, + "Expected some messages to be removed; got 0" + ); + + // The new head message must contain the recap text (not the inner fallback). + assert_eq!( + history.len(), + 2, + "Expected summary message + 1 preserved tail; got {}", + history.len() + ); + match &history[0] { + ConversationMessage::Chat(m) => { + assert!( + m.content.contains("[segment-recap]"), + "Expected summary message to contain '[segment-recap]'; got: {:?}", + m.content + ); + // Must also contain the actual recap text from the stub. + assert!( + m.content.contains("rolling recap"), + "Expected summary to contain recap text from stub provider; got: {:?}", + m.content + ); + } + other => panic!("Expected Chat message for summary, got: {:?}", other), + } + + // Tail must be preserved verbatim. + match &history[1] { + ConversationMessage::Chat(m) => { + assert_eq!(m.content, "message 3 — preserved tail"); + } + other => panic!("Expected Chat tail message, got: {:?}", other), + } +} + +// ── Test 3: soft-fallback when archivist absent ─────────────────────────────── + +/// When the archivist is disabled (no conn), `rolling_segment_recap` returns +/// `None` and `SegmentRecapSummarizer` must fall back to the inner summarizer. +/// The prompt must remain bounded (no panic, no over-budget, no error). +#[tokio::test] +async fn soft_fallback_when_archivist_absent() { + // Use a disabled archivist (no SQLite connection). + let disabled_hook = Arc::new(ArchivistHook::disabled()); + + let inner = RecordingSummarizer::new(); + let recap_summ = SegmentRecapSummarizer::new( + disabled_hook, + "no-session".to_string(), + inner.clone() as Arc, + ) + .with_keep_recent(1); + + let mut history = vec![user("evict me"), user("evict me too"), user("keep me")]; + + let stats = recap_summ + .summarize(&mut history, "test-model") + .await + .expect("summarize must not return Err — soft-fallback guarantees boundedness"); + + // Inner summarizer must have been called exactly once. + assert_eq!( + inner.call_count(), + 1, + "Inner summarizer must be called when archivist is absent" + ); + + // History must be bounded (inner replaced it). + assert_eq!( + history.len(), + 1, + "Inner summarizer must have reduced the history" + ); + match &history[0] { + ConversationMessage::Chat(m) => { + assert_eq!(m.content, "inner fallback summary"); + } + _ => panic!("Expected system summary from inner summarizer"), + } + + // Stats must be valid (not zeroes — inner mock returns non-zero values). + assert_eq!(stats.approx_tokens_freed, 500); +} + +// ── Test 4: soft-fallback when LLM stub fails ──────────────────────────────── + +/// Tier 3 — the bookend heuristic stub must NEVER become live compaction +/// text. Here there is no chat provider configured, so `summarize_entries` +/// produces only `segments::fallback_summary` (bookend stub) with +/// `produced_by_llm = false`. `rolling_segment_recap` must therefore return +/// `None`, and `SegmentRecapSummarizer` falls back to the inner summarizer. +/// (Option A: only a genuine summariser-produced recap may be compaction.) +#[tokio::test] +async fn bookend_stub_never_becomes_compaction_falls_back_to_inner() { + let conn = setup_conn(); + // `ArchivistHook::new` leaves `chat_provider = None` → summarize_entries + // takes the no-provider branch → bookend stub, produced_by_llm = false. + let hook = Arc::new(ArchivistHook::new(conn.clone(), true)); + + let session = "p15-bookend-stub-none"; + + hook.on_turn_complete(&TurnContext { + user_message: "Hello, what is Rust?".into(), + assistant_response: "Rust is a systems language.".into(), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Only the bookend stub is available → MUST be None. + let recap = hook.rolling_segment_recap(session).await; + assert!( + recap.is_none(), + "rolling_segment_recap must return None when only the bookend stub \ + is available — the stub must never be live compaction text" + ); + + let inner = RecordingSummarizer::new(); + let recap_summ = SegmentRecapSummarizer::new( + Arc::clone(&hook), + session.to_string(), + inner.clone() as Arc, + ) + .with_keep_recent(1); + + let mut history = vec![user("evict"), user("evict2"), user("keep")]; + recap_summ + .summarize(&mut history, "test-model") + .await + .expect("must not panic — inner summarizer is the safety net"); + + assert_eq!( + inner.call_count(), + 1, + "Inner summarizer must run when only the bookend stub is available \ + (stub must never be live compaction text)" + ); +} + +/// Tier 2 — when a chat provider IS configured but its call fails, +/// `LlmSummariser`'s soft-fallback yields an *inert clipped-content* recap +/// (the real conversation text, truncated — NOT the bookend stub). Per +/// Option A this is acceptable as live compaction text (real content, +/// strictly better than no compaction), so `rolling_segment_recap` returns +/// `Some` and the inner summarizer is NOT used. +#[tokio::test] +async fn failing_provider_yields_inert_clipped_recap_used_as_compaction() { + let conn = setup_conn(); + let hook = Arc::new(ArchivistHook::new_with_stubs( + conn.clone(), + Arc::new(FailingChatProvider), + Arc::new(StubEmbedder), + )); + + let session = "p15-inert-clipped-kept"; + + hook.on_turn_complete(&TurnContext { + user_message: "Explain ownership in Rust in detail.".into(), + assistant_response: "Ownership means each value has a single owner; \ + borrows are checked at compile time." + .into(), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: 1, + }) + .await + .unwrap(); + + // Provider present but failing → LlmSummariser inert fallback → real + // clipped content (not the bookend stub) → Some, treated as usable. + let recap = hook.rolling_segment_recap(session).await; + assert!( + recap.is_some(), + "Inert clipped-content recap (real text) is acceptable compaction \ + text — must be Some, not None" + ); + + let inner = RecordingSummarizer::new(); + let recap_summ = SegmentRecapSummarizer::new( + Arc::clone(&hook), + session.to_string(), + inner.clone() as Arc, + ) + .with_keep_recent(1); + + let mut history = vec![user("evict"), user("evict2"), user("keep")]; + recap_summ + .summarize(&mut history, "test-model") + .await + .expect("must not panic"); + + assert_eq!( + inner.call_count(), + 0, + "Inner summarizer must NOT run when an inert clipped-content recap \ + is available (real content, better than no compaction)" + ); +} + +// ── Test 5: flag off → inner summarizer runs, Phase 1.5 absent ─────────────── + +/// When `unified_compaction_enabled = false`, the `SegmentRecapSummarizer` +/// is NOT instantiated — the builder uses `ProviderSummarizer` directly. +/// This test verifies the flag-off path via the `RecordingSummarizer` fallback: +/// a `SegmentRecapSummarizer` with an otherwise-healthy archivist must still +/// fall through to the inner summarizer if the session has no entries yet +/// (no open segment entries → recap = None → inner). +#[tokio::test] +async fn no_entries_returns_none_and_inner_summarizer_fires() { + let conn = setup_conn(); + // Archivist with no turns written — no open segment, no entries. + let hook = Arc::new(ArchivistHook::new_with_stubs( + conn.clone(), + Arc::new(StubChatProvider), + Arc::new(StubEmbedder), + )); + + let inner = RecordingSummarizer::new(); + let recap_summ = SegmentRecapSummarizer::new( + Arc::clone(&hook), + "empty-session".to_string(), + inner.clone() as Arc, + ) + .with_keep_recent(1); + + let mut history = vec![user("a"), user("b"), user("c")]; + let _stats = recap_summ + .summarize(&mut history, "test-model") + .await + .expect("must not error"); + + // With no open segment, rolling_segment_recap returns None → + // inner summarizer fires. + assert_eq!( + inner.call_count(), + 1, + "Inner summarizer must run when session has no open segment entries" + ); +} + +// ── Test 6: Phase 1 regression guard ───────────────────────────────────────── + +/// `on_segment_closed` (finalize path) still works after Phase 1.5 refactor: +/// the shared `summarize_entries` helper must produce the same result, and +/// `segment_set_summary` + embedding must still fire at close time. +#[tokio::test] +async fn phase1_finalize_path_still_persists_summary_and_embedding() { + let conn = setup_conn(); + let hook = Arc::new(ArchivistHook::new_with_stubs( + conn.clone(), + Arc::new(StubChatProvider), + Arc::new(StubEmbedder), + )); + + let session = "p15-finalize-regression"; + + // Two turns in the first segment. + for i in 1..=2u64 { + hook.on_turn_complete(&TurnContext { + user_message: format!("Rust turn {i}"), + assistant_response: format!("Answer {i}"), + tool_calls: vec![], + turn_duration_ms: 50, + session_id: Some(session.into()), + iteration_count: i as usize, + }) + .await + .unwrap(); + } + + // Force-flush to trigger on_segment_closed (finalize path). + hook.flush_open_segment(session).await; + + // The closed segment must have a summary. + let all_segs = seg::segments_by_namespace(&conn, "global", 10).unwrap(); + let flushed = all_segs + .iter() + .find(|s| { + s.session_id == session && s.summary.as_ref().map(|s| !s.is_empty()).unwrap_or(false) + }) + .expect("Expected flushed segment to have a non-empty summary after Phase 1.5 refactor"); + + let summary = flushed.summary.as_ref().unwrap(); + // The stub always returns "rolling recap: discussed memory safety in Rust". + assert!( + summary.contains("rolling recap"), + "Expected summary to contain stub recap text after finalize; got: {summary:?}" + ); + + // Embedding must also still be written (finalize-only invariant). + let embedding = + seg::segment_embedding_get(&conn, &flushed.segment_id, "stub-embedder-v1").unwrap(); + assert!( + embedding.is_some(), + "Expected finalize-time embedding to still be written after Phase 1.5 refactor" + ); +} diff --git a/src/openhuman/context/summarizer.rs b/src/openhuman/context/summarizer.rs index 145040bf2..fb5607fd8 100644 --- a/src/openhuman/context/summarizer.rs +++ b/src/openhuman/context/summarizer.rs @@ -245,7 +245,11 @@ impl Summarizer for ProviderSummarizer { /// head length. Returns 0 when the adjustment would consume the entire /// history, meaning there is nothing we can safely summarize without /// breaking the API invariant. -fn snap_split_forward(history: &[ConversationMessage], proposed_head: usize) -> usize { +/// +/// Exported as `pub(super)` so sibling modules (e.g. +/// `segment_recap_summarizer`) can reuse the same invariant instead of +/// maintaining a separate copy. +pub(super) fn snap_split_forward(history: &[ConversationMessage], proposed_head: usize) -> usize { let mut head = proposed_head.min(history.len()); // If the message immediately *before* the split is an // AssistantToolCalls and the message *at* the split is its diff --git a/src/openhuman/memory/mod.rs b/src/openhuman/memory/mod.rs index 0fe86924b..5966a4689 100644 --- a/src/openhuman/memory/mod.rs +++ b/src/openhuman/memory/mod.rs @@ -13,6 +13,7 @@ pub mod ops; pub mod rpc_models; pub mod safety; pub mod schemas; +pub mod stm_recall; pub mod store; pub mod sync_status; pub mod tool_memory; diff --git a/src/openhuman/memory/stm_recall/constants.rs b/src/openhuman/memory/stm_recall/constants.rs new file mode 100644 index 000000000..88c90b7ab --- /dev/null +++ b/src/openhuman/memory/stm_recall/constants.rs @@ -0,0 +1,35 @@ +//! Tunable constants for STM recall — the STM/LTM knobs. +//! +//! Split out of `mod.rs` per the repo's "light `mod.rs`" rule (CLAUDE.md): +//! `mod.rs` stays export-focused; operational/tunable values live here and +//! are re-exported (`pub use constants::*;`) so existing `super::CONST` +//! references in `recall.rs` / `tool.rs` keep working unchanged. + +/// STM recency window in days. Segments or episodic entries older than this +/// are excluded — they belong in LTM (the memory tree). +pub const RECENCY_WINDOW_DAYS: f64 = 14.0; + +/// Hard cap on segments loaded for vector search (Arm 2). +/// Keeps the brute-force cosine pass bounded at single-user scale. +pub const RECENCY_WINDOW_MAX_SEGMENTS: usize = 100; + +/// Cosine similarity gate for Arm 2 (segment recaps). +/// Below this threshold a recap is excluded regardless of recency. +/// Range: [0.0, 1.0]; 0.65 is "medium gate" — confident topical overlap. +pub const COSINE_GATE: f32 = 0.65; + +/// Maximum segment recaps to include in the output block. +pub const MAX_SEGMENT_RECAPS: usize = 5; + +/// Maximum raw episodic turns to include in the output block. +pub const MAX_EPISODIC_TURNS: usize = 5; + +/// Approximate token budget for the entire STM block (chars / 4 ≈ tokens). +/// ~1500 tokens × 4 chars/token = 6000 chars. +pub const TOKEN_BUDGET: usize = 6_000; + +/// How many FTS5 candidates to fetch before applying the high-precision gate. +/// The gate is: only strong keyword matches survive — FTS5 rank threshold is +/// applied at the DB level via LIMIT; we over-fetch slightly and let dedup +/// finish trimming. +pub const FTS5_LIMIT: usize = 20; diff --git a/src/openhuman/memory/stm_recall/mod.rs b/src/openhuman/memory/stm_recall/mod.rs new file mode 100644 index 000000000..787d4ad07 --- /dev/null +++ b/src/openhuman/memory/stm_recall/mod.rs @@ -0,0 +1,47 @@ +//! Phase 3 — Bounded cross-thread STM recall. +//! +//! Assembles a bounded, recency-weighted context block from two arms: +//! +//! - **Arm 1** — FTS5 over not-yet-compacted recent episodic entries from +//! OTHER sessions. Reuses [`crate::openhuman::memory::store::fts5::episodic_cross_session_search`]. +//! When no user query is available (preemptive/session-start case), falls back to +//! a recency selection of recent non-current-session episodic turns. +//! +//! - **Arm 2** — Brute-force cosine nearest-neighbour over `segment_embeddings` +//! (per-model table from Phase 0+1). Single-user scale: no ANN index, no new +//! deps. Loads candidate vectors, computes cosine, top-k. Filters by +//! `model_signature` and excludes the current `session_id`. +//! +//! **Merge → dedup → bound:** +//! - Dedup: if a segment recap and its raw episodic rows (within +//! `start_episodic_id..=end_episodic_id`) both appear, prefer the recap and +//! drop the overlapping episodic hits. +//! - Recency-weight: scored by `updated_at` timestamp proximity. +//! - Hard-cap: tunable token budget ([`TOKEN_BUDGET`]) and top-k bounds +//! ([`MAX_SEGMENT_RECAPS`] + [`MAX_EPISODIC_TURNS`]). +//! +//! ## Tunable consts (all in this file, all documented) +//! - [`RECENCY_WINDOW_DAYS`] — how many days back to search (STM/LTM boundary) +//! - [`RECENCY_WINDOW_MAX_SEGMENTS`] — max segments to load for vector search +//! - [`COSINE_GATE`] — minimum similarity for Arm 2 (medium gate) +//! - [`MAX_SEGMENT_RECAPS`] — top-k segment recaps to include +//! - [`MAX_EPISODIC_TURNS`] — max raw episodic turns to include +//! - [`TOKEN_BUDGET`] — hard token budget (chars / 4 approx) +//! - [`FTS5_LIMIT`] — how many FTS5 candidates to fetch before gating +//! +//! ## Scope boundary +//! Does NOT traverse `tree::*` (`SummaryNode`, `memory_tree_*`). The memory +//! tree is LTM; this module is strictly STM (recent episodic + segment layer). + +pub mod recall; +pub mod tool; + +pub use recall::{stm_recall, StmRecallBlock, StmRecallOpts}; + +// ───────────────────────────────────────────────────────────────────────────── +// Tunable constants — the STM/LTM knobs. Moved to `constants.rs` per the +// repo's "light mod.rs" rule; re-exported so `super::CONST` keeps resolving. +// ───────────────────────────────────────────────────────────────────────────── + +mod constants; +pub use constants::*; diff --git a/src/openhuman/memory/stm_recall/recall.rs b/src/openhuman/memory/stm_recall/recall.rs new file mode 100644 index 000000000..6ddacc302 --- /dev/null +++ b/src/openhuman/memory/stm_recall/recall.rs @@ -0,0 +1,629 @@ +//! Core STM recall logic — two-arm hybrid retrieval. +//! +//! Arm 1: FTS5 episodic cross-session search (keyword, high-precision gate). +//! Arm 2: Brute-force cosine over `segment_embeddings` (vector, medium gate). + +use parking_lot::Mutex; +use rusqlite::{params, Connection}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use crate::openhuman::memory::store::fts5; +use crate::openhuman::memory::store::fts5::EpisodicEntry; + +use super::{ + COSINE_GATE, FTS5_LIMIT, MAX_EPISODIC_TURNS, MAX_SEGMENT_RECAPS, RECENCY_WINDOW_DAYS, + RECENCY_WINDOW_MAX_SEGMENTS, TOKEN_BUDGET, +}; + +/// A single item in the assembled STM recall block. +#[derive(Debug, Clone)] +pub enum StmItem { + /// A segment-level recap (compacted, from `segment_embeddings` arm). + SegmentRecap { + segment_id: String, + session_id: String, + summary: String, + /// `start_episodic_id..=end_episodic_id` span — used for dedup. + start_episodic_id: i64, + end_episodic_id: Option, + updated_at: f64, + cosine: f32, + }, + /// A raw episodic turn from Arm 1 (FTS5 keyword match). + EpisodicTurn { + id: Option, + session_id: String, + timestamp: f64, + role: String, + content: String, + }, +} + +impl StmItem { + fn timestamp(&self) -> f64 { + match self { + Self::SegmentRecap { updated_at, .. } => *updated_at, + Self::EpisodicTurn { timestamp, .. } => *timestamp, + } + } + + fn approx_chars(&self) -> usize { + match self { + Self::SegmentRecap { summary, .. } => summary.len() + 60, + Self::EpisodicTurn { content, role, .. } => content.len() + role.len() + 20, + } + } +} + +/// Options for a single STM recall pass. +#[derive(Debug, Clone, Default)] +pub struct StmRecallOpts<'a> { + /// The active session to exclude from all results. + pub exclude_session: &'a str, + /// Optional query text for Arm 1 (FTS5) and Arm 2 (embed). + /// When `None`, Arm 1 falls back to a recency selection and Arm 2 is skipped. + pub query: Option<&'a str>, + /// Model signature to filter segment embeddings (e.g. `"cloud:voyage-3:1024"`). + /// When `None`, any model signature is accepted (weaker but still useful). + pub model_signature: Option<&'a str>, +} + +/// The assembled STM recall block. +#[derive(Debug, Clone, Default)] +pub struct StmRecallBlock { + /// Deduplicated, recency-weighted items, bounded to [`TOKEN_BUDGET`]. + pub items: Vec, + /// Number of items dropped due to token-budget exhaustion. + pub dropped_budget: usize, + /// Number of episodic hits dropped because they fell inside a segment span (dedup). + pub dropped_dedup: usize, + /// Cosine arm — items retrieved before gating. + pub cosine_candidates: usize, + /// FTS5 arm — items retrieved before gating. + pub fts5_candidates: usize, +} + +impl StmRecallBlock { + /// `true` when the block has no usable content. + pub fn is_empty(&self) -> bool { + self.items.is_empty() + } + + /// Render the block into a markdown string suitable for injection into the + /// system prompt or a user-turn context block. + pub fn render(&self) -> String { + if self.items.is_empty() { + return String::new(); + } + let mut out = String::from("## Recent context from other conversations\n\n"); + out.push_str( + "The following snippets are from previous conversations in other chat threads. \ + They are provided for continuity — reference them when they are relevant to \ + the current request, but do not surface them unless asked.\n\n", + ); + + let mut seg_count = 0usize; + let mut ep_count = 0usize; + + for item in &self.items { + match item { + StmItem::SegmentRecap { + session_id, + summary, + updated_at, + .. + } => { + seg_count += 1; + let age_days = age_days_from_ts(*updated_at); + out.push_str(&format!( + "**Conversation recap** (thread `{session_id}`, ~{age_days:.0} days ago):\n{summary}\n\n" + )); + } + StmItem::EpisodicTurn { + session_id, + timestamp, + role, + content, + .. + } => { + ep_count += 1; + let age_days = age_days_from_ts(*timestamp); + out.push_str(&format!( + "**{role}** (thread `{session_id}`, ~{age_days:.0} days ago): {content}\n\n" + )); + } + } + } + + tracing::debug!( + "[stm_recall] rendered block: {} recaps + {} episodic turns, {} dropped_dedup, {} dropped_budget", + seg_count, + ep_count, + self.dropped_dedup, + self.dropped_budget + ); + out + } +} + +fn age_days_from_ts(ts: f64) -> f64 { + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64(); + (now - ts).max(0.0) / 86_400.0 +} + +fn now_secs() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs_f64() +} + +/// Brute-force cosine similarity between two equal-length float slices. +/// Returns 0.0 when either vector has zero magnitude (zero-padded / inert embedder). +pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f32 { + if a.len() != b.len() || a.is_empty() { + return 0.0; + } + let mut dot = 0.0_f32; + let mut norm_a = 0.0_f32; + let mut norm_b = 0.0_f32; + for (x, y) in a.iter().zip(b.iter()) { + dot += x * y; + norm_a += x * x; + norm_b += y * y; + } + let denom = norm_a.sqrt() * norm_b.sqrt(); + if denom < f32::EPSILON { + 0.0 + } else { + (dot / denom).clamp(-1.0, 1.0) + } +} + +/// Decode a raw BLOB from `segment_embeddings.vector` into `Vec`. +/// +/// A well-formed embedding blob is a whole number of little-endian `f32`s +/// (length a multiple of 4). A non-multiple-of-4 length means the blob is +/// truncated/corrupt; silently dropping the trailing bytes would yield a +/// wrong-length vector, so we reject it (empty → cosine treats it as a +/// non-match, which is the safe outcome). +fn decode_vector_blob(bytes: &[u8]) -> Vec { + if bytes.len() % 4 != 0 { + tracing::warn!( + "[stm_recall] decode_vector_blob: blob length {} is not a multiple of 4 — \ + discarding malformed embedding", + bytes.len() + ); + return Vec::new(); + } + bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() +} + +// ── Arm 2: vector search over segment_embeddings ────────────────────────────── + +/// A row pulled from `segment_embeddings` joined with `conversation_segments`. +struct SegmentEmbeddingRow { + segment_id: String, + session_id: String, + summary: Option, + start_episodic_id: i64, + end_episodic_id: Option, + updated_at: f64, + vector: Vec, + model_signature: String, +} + +/// Load candidate segment embeddings for the cosine pass. +/// +/// Applies: +/// - Recency window (last [`RECENCY_WINDOW_DAYS`] days) +/// - Exclude current `session_id` +/// - Optional `model_signature` filter +/// - Capped at [`RECENCY_WINDOW_MAX_SEGMENTS`] rows +fn load_segment_embedding_candidates( + conn: &Arc>, + exclude_session: &str, + model_signature: Option<&str>, + cutoff_ts: f64, +) -> anyhow::Result> { + let conn = conn.lock(); + + // We join `segment_embeddings` (se) with `conversation_segments` (cs) + // to get the session_id, summary, episodic span, and recency. + // Filter: cs.session_id != exclude_session, cs.updated_at >= cutoff_ts, + // cs.summary IS NOT NULL (only summarised segments have useful recaps), + // optionally se.model_signature = ? + let rows: Vec = if let Some(sig) = model_signature { + let mut stmt = conn.prepare( + "SELECT se.segment_id, cs.session_id, cs.summary, + cs.start_episodic_id, cs.end_episodic_id, + cs.updated_at, se.vector, se.model_signature + FROM segment_embeddings AS se + JOIN conversation_segments AS cs ON se.segment_id = cs.segment_id + WHERE cs.session_id != ?1 + AND cs.updated_at >= ?2 + AND cs.summary IS NOT NULL + AND se.model_signature = ?3 + ORDER BY cs.updated_at DESC + LIMIT ?4", + )?; + let collected: Vec = stmt + .query_map( + params![ + exclude_session, + cutoff_ts, + sig, + RECENCY_WINDOW_MAX_SEGMENTS as i64 + ], + |row| { + let vector_bytes: Vec = row.get(6)?; + Ok(SegmentEmbeddingRow { + segment_id: row.get(0)?, + session_id: row.get(1)?, + summary: row.get(2)?, + start_episodic_id: row.get(3)?, + end_episodic_id: row.get(4)?, + updated_at: row.get(5)?, + vector: decode_vector_blob(&vector_bytes), + model_signature: row.get(7)?, + }) + }, + )? + .collect::, _>>()?; + collected + } else { + // No model filter — accept any model (weaker but useful when no sig available). + let mut stmt = conn.prepare( + "SELECT se.segment_id, cs.session_id, cs.summary, + cs.start_episodic_id, cs.end_episodic_id, + cs.updated_at, se.vector, se.model_signature + FROM segment_embeddings AS se + JOIN conversation_segments AS cs ON se.segment_id = cs.segment_id + WHERE cs.session_id != ?1 + AND cs.updated_at >= ?2 + AND cs.summary IS NOT NULL + ORDER BY cs.updated_at DESC + LIMIT ?3", + )?; + let collected: Vec = stmt + .query_map( + params![ + exclude_session, + cutoff_ts, + RECENCY_WINDOW_MAX_SEGMENTS as i64 + ], + |row| { + let vector_bytes: Vec = row.get(6)?; + Ok(SegmentEmbeddingRow { + segment_id: row.get(0)?, + session_id: row.get(1)?, + summary: row.get(2)?, + start_episodic_id: row.get(3)?, + end_episodic_id: row.get(4)?, + updated_at: row.get(5)?, + vector: decode_vector_blob(&vector_bytes), + model_signature: row.get(7)?, + }) + }, + )? + .collect::, _>>()?; + collected + }; + + tracing::debug!( + "[stm_recall] arm2: loaded {} segment embedding candidates (model_sig={:?})", + rows.len(), + model_signature + ); + Ok(rows) +} + +/// Load recent episodic turns from other sessions (recency fallback for Arm 1 +/// when no query is available). +fn load_recent_episodic_other_sessions( + conn: &Arc>, + exclude_session: &str, + cutoff_ts: f64, + limit: usize, +) -> anyhow::Result> { + let conn = conn.lock(); + let mut stmt = conn.prepare( + "SELECT id, session_id, timestamp, role, content, lesson, tool_calls_json, cost_microdollars + FROM episodic_log + WHERE session_id != ?1 + AND timestamp >= ?2 + ORDER BY timestamp DESC + LIMIT ?3", + )?; + let rows = stmt + .query_map(params![exclude_session, cutoff_ts, limit as i64], |row| { + Ok(EpisodicEntry { + id: row.get(0)?, + session_id: row.get(1)?, + timestamp: row.get(2)?, + role: row.get(3)?, + content: row.get(4)?, + lesson: row.get(5)?, + tool_calls_json: row.get(6)?, + cost_microdollars: row.get::<_, i64>(7)? as u64, + }) + })? + .collect::, _>>()?; + tracing::debug!( + "[stm_recall] arm1 recency fallback: loaded {} episodic turns from other sessions", + rows.len() + ); + Ok(rows) +} + +/// Main entry point: run the two-arm STM recall and return an assembled block. +/// +/// When `opts.query` is `None` (preemptive/session-start case): +/// - Arm 1 uses a recency selection of recent other-session episodic turns. +/// - Arm 2 is skipped (no query vector to embed; caller may pre-compute one). +/// +/// When `opts.query` is `Some(q)`: +/// - Arm 1 runs FTS5 cross-session search with `q` and the high-precision gate. +/// - Arm 2 runs brute-force cosine against any pre-computed `query_embedding`. +/// +/// `query_embedding` is the embedding of `opts.query` (when provided). +/// The caller is responsible for producing it (avoids a blocking embed call here). +/// If `None`, Arm 2 is skipped regardless of `opts.query`. +pub fn stm_recall( + conn: &Arc>, + opts: &StmRecallOpts<'_>, + query_embedding: Option<&[f32]>, +) -> anyhow::Result { + let cutoff_ts = now_secs() - RECENCY_WINDOW_DAYS * 86_400.0; + let mut block = StmRecallBlock::default(); + + tracing::debug!( + "[stm_recall] starting recall exclude_session={} has_query={} has_embedding={} recency_days={}", + opts.exclude_session, + opts.query.is_some(), + query_embedding.is_some(), + RECENCY_WINDOW_DAYS + ); + + // ── Arm 2: vector search over segment_embeddings ────────────────────────── + let mut segment_items: Vec = Vec::new(); + let mut segment_spans: Vec<(i64, Option)> = Vec::new(); // for dedup + + if let Some(q_emb) = query_embedding { + if !q_emb.is_empty() { + let candidates = load_segment_embedding_candidates( + conn, + opts.exclude_session, + opts.model_signature, + cutoff_ts, + )?; + block.cosine_candidates = candidates.len(); + + let mut scored: Vec<(f32, SegmentEmbeddingRow)> = candidates + .into_iter() + .filter_map(|row| { + if row.vector.is_empty() { + tracing::debug!( + "[stm_recall] arm2: skipping segment {} — zero-length vector (inert embedder?)", + row.segment_id + ); + return None; + } + let cos = cosine_similarity(q_emb, &row.vector); + tracing::debug!( + "[stm_recall] arm2: segment={} session={} cosine={:.3} gate={}", + row.segment_id, + row.session_id, + cos, + COSINE_GATE + ); + if cos >= COSINE_GATE { + Some((cos, row)) + } else { + None + } + }) + .collect(); + + // Sort descending by cosine then recency + scored.sort_by(|a, b| { + b.0.partial_cmp(&a.0) + .unwrap_or(std::cmp::Ordering::Equal) + .then( + b.1.updated_at + .partial_cmp(&a.1.updated_at) + .unwrap_or(std::cmp::Ordering::Equal), + ) + }); + + for (cos, row) in scored.into_iter().take(MAX_SEGMENT_RECAPS) { + let summary = match row.summary { + Some(ref s) if !s.trim().is_empty() => s.clone(), + _ => continue, + }; + tracing::debug!( + "[stm_recall] arm2: accepting segment={} session={} cosine={:.3} model={}", + row.segment_id, + row.session_id, + cos, + row.model_signature + ); + segment_spans.push((row.start_episodic_id, row.end_episodic_id)); + segment_items.push(StmItem::SegmentRecap { + segment_id: row.segment_id, + session_id: row.session_id, + summary, + start_episodic_id: row.start_episodic_id, + end_episodic_id: row.end_episodic_id, + updated_at: row.updated_at, + cosine: cos, + }); + } + tracing::debug!( + "[stm_recall] arm2: {} recaps accepted after cosine gate", + segment_items.len() + ); + } + } + + // ── Arm 1: FTS5 or recency episodic ────────────────────────────────────── + let mut episodic_items: Vec = Vec::new(); + + let raw_episodic: Vec = if let Some(query) = opts.query { + // Keyword search path + let hits = fts5::episodic_cross_session_search( + conn, + query, + FTS5_LIMIT, + Some(opts.exclude_session), + )?; + block.fts5_candidates = hits.len(); + tracing::debug!( + "[stm_recall] arm1 FTS5: {} hits for query (exclude={})", + hits.len(), + opts.exclude_session + ); + hits + } else { + // Recency fallback — no query; pull most-recent turns from other sessions + let hits = + load_recent_episodic_other_sessions(conn, opts.exclude_session, cutoff_ts, FTS5_LIMIT)?; + block.fts5_candidates = hits.len(); + hits + }; + + // Dedup: drop episodic rows whose ID falls within any accepted segment's span. + for entry in raw_episodic { + let entry_id = entry.id.unwrap_or(-1); + + // Check if this episodic entry is covered by any accepted segment span. + let covered = segment_spans + .iter() + .any(|(start, end)| entry_id >= *start && end.map_or(false, |e| entry_id <= e)); + + if covered { + block.dropped_dedup += 1; + tracing::debug!( + "[stm_recall] arm1: dropping episodic id={} — covered by segment span", + entry_id + ); + continue; + } + + // Recency window applies consistently — including keyword/FTS mode. + // FTS5 keyword search is not time-bounded, so without this an + // older-than-window episodic hit could leak into STM. The recency + // window IS the STM/LTM boundary and must always hold. + if entry.timestamp < cutoff_ts { + continue; + } + + episodic_items.push(StmItem::EpisodicTurn { + id: entry.id, + session_id: entry.session_id, + timestamp: entry.timestamp, + role: entry.role, + content: entry.content, + }); + } + + tracing::debug!( + "[stm_recall] arm1: {} episodic items after dedup (dropped_dedup={})", + episodic_items.len(), + block.dropped_dedup + ); + + // ── Merge + cap ─────────────────────────────────────────────────────────── + // Recency-weight: sort each arm descending by timestamp. Interleave by picking + // the most-recent item across both arms. + segment_items.sort_by(|a, b| { + b.timestamp() + .partial_cmp(&a.timestamp()) + .unwrap_or(std::cmp::Ordering::Equal) + }); + episodic_items.sort_by(|a, b| { + b.timestamp() + .partial_cmp(&a.timestamp()) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + // Apply top-k caps before interleave + let seg_capped: Vec = segment_items.into_iter().take(MAX_SEGMENT_RECAPS).collect(); + let ep_capped: Vec = episodic_items + .into_iter() + .take(MAX_EPISODIC_TURNS) + .collect(); + + // Interleave: recency-first merge + let mut seg_iter = seg_capped.into_iter().peekable(); + let mut ep_iter = ep_capped.into_iter().peekable(); + let mut merged: Vec = Vec::new(); + + loop { + match (seg_iter.peek(), ep_iter.peek()) { + (None, None) => break, + (Some(_), None) => { + merged.extend(seg_iter.by_ref()); + break; + } + (None, Some(_)) => { + merged.extend(ep_iter.by_ref()); + break; + } + (Some(s), Some(e)) => { + if s.timestamp() >= e.timestamp() { + merged.push(seg_iter.next().expect("peek confirmed Some")); + } else { + merged.push(ep_iter.next().expect("peek confirmed Some")); + } + } + } + } + + // Apply token budget + let mut used_chars = 0usize; + let mut final_items: Vec = Vec::new(); + let mut dropped_budget = 0usize; + + for item in merged { + let chars = item.approx_chars(); + if used_chars + chars > TOKEN_BUDGET { + dropped_budget += 1; + tracing::debug!( + "[stm_recall] budget: dropping item (would exceed {TOKEN_BUDGET} chars)" + ); + continue; + } + used_chars += chars; + final_items.push(item); + } + + tracing::debug!( + "[stm_recall] final block: {} items, ~{} chars, {} dropped_budget, {} dropped_dedup, {} cosine_candidates, {} fts5_candidates", + final_items.len(), + used_chars, + dropped_budget, + block.dropped_dedup, + block.cosine_candidates, + block.fts5_candidates + ); + + block.items = final_items; + block.dropped_budget = dropped_budget; + Ok(block) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Unit tests +// ───────────────────────────────────────────────────────────────────────────── + +#[cfg(test)] +#[path = "recall_tests.rs"] +mod tests; diff --git a/src/openhuman/memory/stm_recall/recall_tests.rs b/src/openhuman/memory/stm_recall/recall_tests.rs new file mode 100644 index 000000000..1914fc7a9 --- /dev/null +++ b/src/openhuman/memory/stm_recall/recall_tests.rs @@ -0,0 +1,721 @@ +//! Unit + integration tests for Phase 3 STM recall. + +use super::*; +use crate::openhuman::agent::harness::archivist::ArchivistHook; +use crate::openhuman::agent::hooks::{PostTurnHook, TurnContext}; +use crate::openhuman::memory::store::events::EVENTS_INIT_SQL; +use crate::openhuman::memory::store::fts5; +use crate::openhuman::memory::store::profile::PROFILE_INIT_SQL; +use crate::openhuman::memory::store::segments::SEGMENTS_INIT_SQL; +use parking_lot::Mutex; +use rusqlite::{params, Connection}; +use std::sync::Arc; + +// ── helpers ────────────────────────────────────────────────────────────────── + +fn setup_conn() -> Arc> { + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch(fts5::EPISODIC_INIT_SQL).unwrap(); + conn.execute_batch(SEGMENTS_INIT_SQL).unwrap(); + conn.execute_batch(EVENTS_INIT_SQL).unwrap(); + conn.execute_batch(PROFILE_INIT_SQL).unwrap(); + Arc::new(Mutex::new(conn)) +} + +fn now_ts() -> f64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs_f64() +} + +/// Insert an episodic entry with an explicit timestamp. +fn insert_episodic( + conn: &Arc>, + session_id: &str, + ts: f64, + role: &str, + content: &str, +) -> i64 { + let c = conn.lock(); + c.execute( + "INSERT INTO episodic_log (session_id, timestamp, role, content, lesson, tool_calls_json, cost_microdollars) VALUES (?1,?2,?3,?4,NULL,NULL,0)", + params![session_id, ts, role, content], + ).unwrap(); + c.last_insert_rowid() +} + +/// Insert a segment with a summary and optional embedding. +fn insert_segment_with_embedding( + conn: &Arc>, + segment_id: &str, + session_id: &str, + start_id: i64, + end_id: i64, + summary: &str, + embedding: Option>, + updated_at: f64, + model_sig: &str, +) { + let c = conn.lock(); + c.execute( + "INSERT INTO conversation_segments + (segment_id, session_id, namespace, start_episodic_id, end_episodic_id, + start_timestamp, end_timestamp, turn_count, summary, status, created_at, updated_at) + VALUES (?1,?2,'global',?3,?4,?5,?5,2,?6,'summarised',?5,?5)", + params![segment_id, session_id, start_id, end_id, updated_at, summary], + ) + .unwrap(); + + if let Some(emb) = embedding { + let bytes: Vec = emb.iter().flat_map(|f| f.to_le_bytes()).collect(); + let dim = emb.len() as i64; + c.execute( + "INSERT INTO segment_embeddings (segment_id, model_signature, vector, dim, created_at) + VALUES (?1,?2,?3,?4,?5)", + params![segment_id, model_sig, bytes, dim, updated_at], + ) + .unwrap(); + } +} + +// ── cosine_similarity unit tests ────────────────────────────────────────────── + +#[test] +fn cosine_identical_vectors_returns_one() { + let v = vec![1.0_f32, 0.0, 0.0]; + assert!((cosine_similarity(&v, &v) - 1.0).abs() < 1e-6); +} + +#[test] +fn cosine_orthogonal_vectors_returns_zero() { + let a = vec![1.0_f32, 0.0, 0.0]; + let b = vec![0.0_f32, 1.0, 0.0]; + assert!((cosine_similarity(&a, &b)).abs() < 1e-6); +} + +#[test] +fn cosine_opposite_vectors_returns_minus_one() { + let a = vec![1.0_f32, 0.0, 0.0]; + let b = vec![-1.0_f32, 0.0, 0.0]; + assert!((cosine_similarity(&a, &b) + 1.0).abs() < 1e-6); +} + +#[test] +fn cosine_zero_vector_returns_zero_not_nan() { + let a = vec![0.0_f32, 0.0, 0.0]; + let b = vec![1.0_f32, 0.0, 0.0]; + let sim = cosine_similarity(&a, &b); + assert!(!sim.is_nan(), "cosine_similarity must not return NaN"); + assert_eq!(sim, 0.0); +} + +#[test] +fn cosine_mismatched_lengths_returns_zero() { + let a = vec![1.0_f32, 0.0]; + let b = vec![1.0_f32, 0.0, 0.0]; + assert_eq!(cosine_similarity(&a, &b), 0.0); +} + +#[test] +fn cosine_empty_vectors_returns_zero() { + assert_eq!(cosine_similarity(&[], &[]), 0.0); +} + +// ── gating threshold tests ──────────────────────────────────────────────────── + +#[test] +fn cosine_gate_const_is_reasonable() { + // Gate must be in (0.5, 1.0) — below 0.5 lets in noise, above 0.9 is too strict. + assert!(super::super::COSINE_GATE > 0.5 && super::super::COSINE_GATE < 0.9); +} + +#[test] +fn arm2_drops_below_gate_and_accepts_above() { + let conn = setup_conn(); + let now = now_ts(); + + // Build query embedding — unit vector along dim 0 + let mut q_emb = vec![0.0_f32; 8]; + q_emb[0] = 1.0; + + // High-match segment: unit vector along dim 0 (cos = 1.0) + let mut high_emb = vec![0.0_f32; 8]; + high_emb[0] = 1.0; + insert_episodic(&conn, "other-session", now - 100.0, "user", "seed turn"); + let id = insert_episodic( + &conn, + "other-session", + now - 90.0, + "assistant", + "high match reply", + ); + insert_segment_with_embedding( + &conn, + "seg-high", + "other-session", + id - 1, + id, + "This conversation covered high-match topics", + Some(high_emb), + now - 50.0, + "test:model:8", + ); + + // Low-match segment: orthogonal vector (cos = 0.0 < gate) + let mut low_emb = vec![0.0_f32; 8]; + low_emb[1] = 1.0; + let id2 = insert_episodic(&conn, "low-session", now - 200.0, "user", "unrelated"); + insert_segment_with_embedding( + &conn, + "seg-low", + "low-session", + id2, + id2, + "This is about something completely unrelated", + Some(low_emb), + now - 150.0, + "test:model:8", + ); + + let opts = StmRecallOpts { + exclude_session: "current-session", + query: Some("high match"), + model_signature: None, + }; + let block = stm_recall(&conn, &opts, Some(&q_emb)).unwrap(); + + let recap_ids: Vec<&str> = block + .items + .iter() + .filter_map(|it| { + if let StmItem::SegmentRecap { segment_id, .. } = it { + Some(segment_id.as_str()) + } else { + None + } + }) + .collect(); + + assert!( + recap_ids.contains(&"seg-high"), + "high-cosine segment must be accepted; got: {:?}", + recap_ids + ); + assert!( + !recap_ids.contains(&"seg-low"), + "low-cosine segment must be excluded by gate; got: {:?}", + recap_ids + ); +} + +// ── exclude-own-session tests ───────────────────────────────────────────────── + +#[test] +fn exclude_own_session_arm1_fts5() { + let conn = setup_conn(); + let now = now_ts(); + + // Other session — should appear + insert_episodic( + &conn, + "other-sess", + now - 100.0, + "user", + "Rust programming concepts", + ); + // Current session — must be excluded + insert_episodic( + &conn, + "current-sess", + now - 50.0, + "user", + "Rust programming today", + ); + + let opts = StmRecallOpts { + exclude_session: "current-sess", + query: Some("Rust programming"), + model_signature: None, + }; + let block = stm_recall(&conn, &opts, None).unwrap(); + + for item in &block.items { + if let StmItem::EpisodicTurn { session_id, .. } = item { + assert_ne!( + session_id, "current-sess", + "arm1 must never return current session items; got session_id={session_id}" + ); + } + } + // Must see the other session + let has_other = block.items.iter().any( + |it| matches!(it, StmItem::EpisodicTurn { session_id, .. } if session_id == "other-sess"), + ); + assert!(has_other, "arm1 must surface items from other sessions"); +} + +#[test] +fn exclude_own_session_arm2_vector() { + let conn = setup_conn(); + let now = now_ts(); + + let mut emb = vec![0.0_f32; 8]; + emb[0] = 1.0; + + // Insert segment from "current-session" — should be excluded + let id = insert_episodic( + &conn, + "current-session", + now - 100.0, + "user", + "current thread", + ); + insert_segment_with_embedding( + &conn, + "seg-current", + "current-session", + id, + id, + "Current session recap", + Some(emb.clone()), + now - 50.0, + "test:model:8", + ); + + // Segment from another session — should appear + let id2 = insert_episodic(&conn, "other-session", now - 200.0, "user", "other thread"); + insert_segment_with_embedding( + &conn, + "seg-other", + "other-session", + id2, + id2, + "Other session recap", + Some(emb.clone()), + now - 100.0, + "test:model:8", + ); + + let opts = StmRecallOpts { + exclude_session: "current-session", + query: None, + model_signature: None, + }; + let block = stm_recall(&conn, &opts, Some(&emb)).unwrap(); + + for item in &block.items { + if let StmItem::SegmentRecap { session_id, .. } = item { + assert_ne!( + session_id, "current-session", + "arm2 must never return current session recaps" + ); + } + } +} + +// ── dedup-by-episodic-span tests ────────────────────────────────────────────── + +#[test] +fn dedup_drops_episodic_row_inside_segment_span() { + let conn = setup_conn(); + let now = now_ts(); + + // Insert episodic rows for "other-session" + let id_start = insert_episodic( + &conn, + "other-session", + now - 300.0, + "user", + "Rust ownership", + ); + let id_end = insert_episodic( + &conn, + "other-session", + now - 290.0, + "assistant", + "Rust uses borrow checker", + ); + + // A high-similarity segment recap covers those episodic rows + let mut emb = vec![0.0_f32; 8]; + emb[0] = 1.0; + insert_segment_with_embedding( + &conn, + "seg-covers", + "other-session", + id_start, + id_end, + "Conversation about Rust ownership and borrow checker", + Some(emb.clone()), + now - 100.0, + "test:model:8", + ); + + let opts = StmRecallOpts { + exclude_session: "current", + query: Some("Rust ownership"), + model_signature: None, + }; + let block = stm_recall(&conn, &opts, Some(&emb)).unwrap(); + + // The segment recap must appear + let has_recap = block.items.iter().any( + |it| matches!(it, StmItem::SegmentRecap { segment_id, .. } if segment_id == "seg-covers"), + ); + assert!(has_recap, "segment recap must appear in output"); + + // The covered episodic rows must NOT appear (dedup) + for item in &block.items { + if let StmItem::EpisodicTurn { id, .. } = item { + assert!( + *id != Some(id_start) && *id != Some(id_end), + "episodic rows inside segment span must be deduplicated; id={id:?}" + ); + } + } + assert!( + block.dropped_dedup > 0, + "dropped_dedup must be > 0 when rows are inside a segment span" + ); +} + +// ── recency window bound test ───────────────────────────────────────────────── + +#[test] +fn recency_window_excludes_old_segments() { + let conn = setup_conn(); + let now = now_ts(); + + // Recent segment — within window + let id1 = insert_episodic(&conn, "recent-session", now - 100.0, "user", "recent"); + let emb_recent: Vec = (0..8).map(|i| if i == 0 { 1.0 } else { 0.0 }).collect(); + insert_segment_with_embedding( + &conn, + "seg-recent", + "recent-session", + id1, + id1, + "Recent segment recap", + Some(emb_recent.clone()), + now - 100.0, // recent + "test:model:8", + ); + + // Old segment — beyond RECENCY_WINDOW_DAYS + let old_ts = now - (super::super::RECENCY_WINDOW_DAYS + 2.0) * 86_400.0; + let id2 = insert_episodic(&conn, "old-session", old_ts, "user", "old content"); + insert_segment_with_embedding( + &conn, + "seg-old", + "old-session", + id2, + id2, + "Old segment recap", + Some(emb_recent.clone()), + old_ts, // older than window + "test:model:8", + ); + + let opts = StmRecallOpts { + exclude_session: "current", + query: None, + model_signature: None, + }; + let block = stm_recall(&conn, &opts, Some(&emb_recent)).unwrap(); + + let seg_ids: Vec<&str> = block + .items + .iter() + .filter_map(|it| { + if let StmItem::SegmentRecap { segment_id, .. } = it { + Some(segment_id.as_str()) + } else { + None + } + }) + .collect(); + + assert!( + !seg_ids.contains(&"seg-old"), + "old segment beyond recency window must be excluded; got: {:?}", + seg_ids + ); + assert!( + seg_ids.contains(&"seg-recent"), + "recent segment must appear; got: {:?}", + seg_ids + ); +} + +// ── token budget test ───────────────────────────────────────────────────────── + +#[test] +fn token_budget_limits_output_size() { + let conn = setup_conn(); + let now = now_ts(); + + // Insert many episodic turns from other sessions + for i in 0..50 { + let large_content = "X".repeat(300); // 300 chars each + insert_episodic( + &conn, + &format!("session-{i}"), + now - i as f64 * 60.0, + "user", + &large_content, + ); + } + + let opts = StmRecallOpts { + exclude_session: "current", + query: None, + model_signature: None, + }; + let block = stm_recall(&conn, &opts, None).unwrap(); + + let total_chars: usize = block.items.iter().map(|it| it.approx_chars()).sum(); + assert!( + total_chars <= super::super::TOKEN_BUDGET, + "total chars {} must not exceed budget {}", + total_chars, + super::super::TOKEN_BUDGET + ); +} + +// ── preemptive recency fallback (no query) ──────────────────────────────────── + +#[test] +fn preemptive_no_query_returns_recent_other_sessions() { + let conn = setup_conn(); + let now = now_ts(); + + // Insert turns from other sessions + insert_episodic( + &conn, + "session-a", + now - 300.0, + "user", + "Alpha session content", + ); + insert_episodic( + &conn, + "session-b", + now - 200.0, + "user", + "Beta session content", + ); + + // Also insert turns for the current session — must be excluded + insert_episodic( + &conn, + "current-session", + now - 100.0, + "user", + "Current session content", + ); + + let opts = StmRecallOpts { + exclude_session: "current-session", + query: None, + model_signature: None, + }; + let block = stm_recall(&conn, &opts, None).unwrap(); + + // Check that we got results from other sessions + let other_sessions: Vec<&str> = block + .items + .iter() + .filter_map(|it| { + if let StmItem::EpisodicTurn { session_id, .. } = it { + Some(session_id.as_str()) + } else { + None + } + }) + .collect(); + + assert!( + !other_sessions.is_empty() || block.items.is_empty(), // empty is OK if no rows + "preemptive fallback must only return other-session items" + ); + for sid in &other_sessions { + assert_ne!( + *sid, "current-session", + "current session must be excluded in preemptive mode" + ); + } +} + +// ── rendered block format ───────────────────────────────────────────────────── + +#[test] +fn render_produces_non_empty_markdown_when_items_present() { + let conn = setup_conn(); + let now = now_ts(); + + insert_episodic(&conn, "other-session", now - 100.0, "user", "Test content"); + + let opts = StmRecallOpts { + exclude_session: "current", + query: None, + model_signature: None, + }; + let block = stm_recall(&conn, &opts, None).unwrap(); + + if !block.items.is_empty() { + let rendered = block.render(); + assert!( + rendered.contains("## Recent context"), + "rendered block must contain heading" + ); + } +} + +#[test] +fn render_empty_block_returns_empty_string() { + let block = StmRecallBlock::default(); + assert!(block.render().is_empty()); + assert!(block.is_empty()); +} + +// ── end-to-end integration test ─────────────────────────────────────────────── +// Drive the real chain: completed turns → episodic rows → segment close +// (recap + embedding via the Phase 0+1 path using stub providers) → +// STM recall returns cross-thread recaps and excludes the current session. + +#[tokio::test] +async fn e2e_stm_recall_chain() { + use crate::openhuman::memory::tree::chat::ChatPrompt; + + let conn = setup_conn(); + + // ── Phase 0+1 stub providers ───────────────────────────────────────────── + // We use a stub chat provider that returns a fixed recap string, and the + // InertEmbedder that returns zero vectors. This exercises the real + // archivist code path (recap + segment_embedding_upsert) without + // requiring a live LLM or Ollama daemon. + + struct StubChat; + use crate::openhuman::memory::tree::chat::ChatProvider; + #[async_trait::async_trait] + impl ChatProvider for StubChat { + fn name(&self) -> &str { + "stub" + } + async fn chat_for_json(&self, _prompt: &ChatPrompt) -> anyhow::Result { + Ok("RECAP: stub LLM summary of the segment.".to_string()) + } + async fn chat_for_text(&self, _prompt: &ChatPrompt) -> anyhow::Result { + Ok("RECAP: stub LLM summary of the segment.".to_string()) + } + } + + use crate::openhuman::memory::tree::score::embed::InertEmbedder; + let chat_provider: Arc = + Arc::new(StubChat); + let embedder: Arc = + Arc::new(InertEmbedder::new()); + + let archivist = ArchivistHook::new_with_stubs(conn.clone(), chat_provider, embedder); + + // ── Turns for "other-thread" ───────────────────────────────────────────── + // Drive 25 turns on session "other-thread" — exceeds max_turns_per_segment (20) + // so a segment boundary fires, the segment closes, and recap + embedding happen. + + for i in 0..25 { + let ctx = TurnContext { + user_message: format!("User message {i} about Rust and memory safety"), + assistant_response: format!("Assistant response {i}: Rust ownership is great."), + tool_calls: vec![], + turn_duration_ms: 100, + session_id: Some("other-thread".to_string()), + iteration_count: i + 1, + }; + archivist.on_turn_complete(&ctx).await.unwrap(); + } + + // Force-flush any trailing open segment so we definitely get a recap. + archivist.flush_open_segment("other-thread").await; + + // ── Verify episodic rows were written ──────────────────────────────────── + let ep_rows = fts5::episodic_session_entries(&conn, "other-thread").unwrap(); + assert!( + ep_rows.len() >= 50, + "expected >=50 episodic rows (2 per turn × 25), got {}", + ep_rows.len() + ); + + // ── Verify segment embedding written ───────────────────────────────────── + let has_embedding = { + let c = conn.lock(); + let count: i64 = c + .query_row( + "SELECT COUNT(*) FROM segment_embeddings se + JOIN conversation_segments cs ON se.segment_id = cs.segment_id + WHERE cs.session_id = 'other-thread'", + [], + |r| r.get(0), + ) + .unwrap_or(0); + count > 0 + }; + assert!( + has_embedding, + "CRITICAL: Phase 0+1 did NOT write segment_embeddings for other-thread. \ + This means the archivist recap+embed path is broken. \ + STM recall Arm 2 would have no data to query." + ); + + // ── Now run STM recall from "current-thread" ───────────────────────────── + // InertEmbedder returns zero vectors, cosine of zero vectors = 0.0 < COSINE_GATE. + // So Arm 2 will find no hits (expected — inert embedder produces identical vectors). + // Arm 1 (FTS5 or recency) should still return episodic turns from other-thread. + + let opts = StmRecallOpts { + exclude_session: "current-thread", + query: Some("Rust memory safety"), + model_signature: None, + }; + let block = stm_recall(&conn, &opts, None).unwrap(); // no embedding for Arm 2 + + // With keyword "Rust memory safety" + other-thread has "Rust and memory safety" + // in the episodic log, Arm 1 should surface at least some results. + // (FTS5 porter-stems "safety" and "Rust" matches the stored content.) + + // Verify: nothing from current-thread + for item in &block.items { + match item { + StmItem::EpisodicTurn { session_id, .. } => { + assert_ne!( + session_id, "current-thread", + "STM recall must never return current-thread items" + ); + } + StmItem::SegmentRecap { session_id, .. } => { + assert_ne!( + session_id, "current-thread", + "STM recall must never return current-thread recaps" + ); + } + } + } + + // The FTS5 arm should have found the other-thread episodic rows + let fts5_hits = block.fts5_candidates; + assert!( + fts5_hits > 0, + "Arm 1 (FTS5) must have found candidates from other-thread for 'Rust memory safety' query; \ + fts5_candidates={}. This proves episodic rows are written and searchable.", + fts5_hits + ); + + // Verify block is well-formed + let rendered = block.render(); + if !block.items.is_empty() { + assert!( + rendered.contains("## Recent context"), + "rendered block must have heading" + ); + } +} diff --git a/src/openhuman/memory/stm_recall/tool.rs b/src/openhuman/memory/stm_recall/tool.rs new file mode 100644 index 000000000..3109ab061 --- /dev/null +++ b/src/openhuman/memory/stm_recall/tool.rs @@ -0,0 +1,200 @@ +//! Agent-callable tool for on-demand STM recall. +//! +//! The agent invokes `stm_recall_search[query]` mid-session to pull +//! cross-thread context by keyword. The tool: +//! +//! 1. Extracts the active session_id from the tool call arguments (or falls +//! back to a process-level default when not provided). +//! 2. Runs [`stm_recall`] with Arm 1 (FTS5 keyword) only — no embedding step +//! at call time to avoid blocking the hot path. +//! 3. Returns the rendered markdown block as a tool result. +//! +//! The tool is registered in `tools/ops.rs` alongside `MemoryRecallTool` when +//! `learning.stm_recall_enabled` is true. + +use crate::openhuman::memory::Memory; +use crate::openhuman::tools::traits::{Tool, ToolResult}; +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; + +use super::recall::{stm_recall, StmRecallOpts}; + +/// On-demand STM recall tool. +/// +/// Searches recent episodic memory from **other** chat threads using +/// keyword matching (FTS5). Vector similarity (Arm 2) is not run here — +/// it requires an embedding call that belongs in the session-start +/// preemptive path, not the on-demand agent invocation path. +pub struct StmRecallTool { + memory: Arc, + /// Session ID to exclude from results. Injected at construction so + /// the tool knows the "current thread" without the agent having to + /// pass it explicitly. + session_id: String, + /// Optional model signature for filtering segment embeddings. + model_signature: Option, +} + +impl StmRecallTool { + /// Create a new `StmRecallTool`. + /// + /// `session_id` — the current session's ID. Results from this session + /// are always excluded. + /// + /// `model_signature` — when `Some`, Arm 2 is filtered to this model. + /// Pass `None` to accept any model (Arm 2 is skipped in on-demand mode + /// anyway, but stored for future extension). + pub fn new( + memory: Arc, + session_id: String, + model_signature: Option, + ) -> Self { + Self { + memory, + session_id, + model_signature, + } + } +} + +#[async_trait] +impl Tool for StmRecallTool { + fn name(&self) -> &str { + "stm_recall_search" + } + + fn description(&self) -> &str { + "Search recent conversational context from other chat threads. \ + Use this when you need facts or context discussed in a previous conversation \ + that may be relevant to the current request. \ + Returns a bounded set of snippets and conversation recaps from other sessions." + } + + fn parameters_schema(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Keywords or phrase to search across recent other-session conversations" + } + }, + "required": ["query"] + }) + } + + async fn execute(&self, args: serde_json::Value) -> anyhow::Result { + let query = args + .get("query") + .and_then(|v| v.as_str()) + .ok_or_else(|| anyhow::anyhow!("Missing 'query' parameter"))? + .trim(); + if query.is_empty() { + return Err(anyhow::anyhow!("query cannot be empty")); + } + + tracing::debug!( + "[stm_recall_tool] on-demand recall query_len={} session={}", + query.chars().count(), + self.session_id + ); + + // Get SQLite connection via the Memory trait's sqlite_conn() hook. + let conn = match self.memory.sqlite_conn() { + Some(c) => c, + None => { + tracing::warn!( + "[stm_recall_tool] memory backend has no SQLite connection — stm_recall unavailable" + ); + return Ok(ToolResult::success( + "STM recall is not available (memory backend is not SQLite-backed).", + )); + } + }; + + let opts = StmRecallOpts { + exclude_session: &self.session_id, + query: Some(query), + model_signature: self.model_signature.as_deref(), + }; + + match stm_recall(&conn, &opts, None) { + Ok(block) => { + tracing::debug!( + "[stm_recall_tool] recall complete: {} items, {} fts5_candidates, {} dropped_dedup", + block.items.len(), + block.fts5_candidates, + block.dropped_dedup + ); + if block.is_empty() { + Ok(ToolResult::success( + "No relevant context found in recent other-session conversations.", + )) + } else { + Ok(ToolResult::success(block.render())) + } + } + Err(e) => { + tracing::warn!("[stm_recall_tool] stm_recall failed: {e}"); + Ok(ToolResult::error(format!("STM recall failed: {e}"))) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::openhuman::embeddings::NoopEmbedding; + use crate::openhuman::memory::UnifiedMemory; + use tempfile::TempDir; + + fn make_mem() -> (TempDir, Arc) { + let tmp = TempDir::new().unwrap(); + let mem = UnifiedMemory::new(tmp.path(), Arc::new(NoopEmbedding), None).unwrap(); + (tmp, Arc::new(mem)) + } + + #[tokio::test] + async fn tool_name_and_schema() { + let (_tmp, mem) = make_mem(); + let tool = StmRecallTool::new(mem, "s1".into(), None); + assert_eq!(tool.name(), "stm_recall_search"); + let schema = tool.parameters_schema(); + assert!(schema["properties"]["query"].is_object()); + assert_eq!(schema["required"][0].as_str(), Some("query")); + } + + #[tokio::test] + async fn tool_empty_query_returns_error() { + let (_tmp, mem) = make_mem(); + let tool = StmRecallTool::new(mem, "s1".into(), None); + let result = tool.execute(json!({"query": " "})).await; + assert!(result.is_err(), "empty query must return Err"); + } + + #[tokio::test] + async fn tool_missing_query_returns_error() { + let (_tmp, mem) = make_mem(); + let tool = StmRecallTool::new(mem, "s1".into(), None); + let result = tool.execute(json!({})).await; + assert!(result.is_err(), "missing query must return Err"); + } + + #[tokio::test] + async fn tool_returns_no_matches_when_empty_db() { + let (_tmp, mem) = make_mem(); + let tool = StmRecallTool::new(mem, "s1".into(), None); + let result = tool + .execute(json!({"query": "Rust ownership"})) + .await + .unwrap(); + assert!(!result.is_error); + assert!( + result.output().contains("No relevant context") || !result.output().is_empty(), + "empty db must return a no-match message, got: {}", + result.output() + ); + } +} diff --git a/src/openhuman/memory/store/memory_trait.rs b/src/openhuman/memory/store/memory_trait.rs index 83a525a4a..97f313f6a 100644 --- a/src/openhuman/memory/store/memory_trait.rs +++ b/src/openhuman/memory/store/memory_trait.rs @@ -398,6 +398,10 @@ impl Memory for UnifiedMemory { async fn health_check(&self) -> bool { self.workspace_dir.exists() && self.db_path.exists() } + + fn sqlite_conn(&self) -> Option>> { + Some(std::sync::Arc::clone(&self.conn)) + } } #[cfg(test)] diff --git a/src/openhuman/memory/traits.rs b/src/openhuman/memory/traits.rs index 6fb7d81e9..e28a44df3 100644 --- a/src/openhuman/memory/traits.rs +++ b/src/openhuman/memory/traits.rs @@ -5,7 +5,10 @@ //! types used for representing and organizing memories. use async_trait::async_trait; +use parking_lot::Mutex; +use rusqlite::Connection; use serde::{Deserialize, Serialize}; +use std::sync::Arc; /// Represents a single stored memory entry with associated metadata. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -151,6 +154,17 @@ pub trait Memory: Send + Sync { /// Performs a health check on the underlying storage system. async fn health_check(&self) -> bool; + + /// Return the shared SQLite connection when the backend is `UnifiedMemory`. + /// + /// Used by subsystems (e.g. `ArchivistHook`) that need direct SQLite + /// access for FTS5 / segment writes without going through the async + /// `Memory` trait. + /// + /// Default: `None`. Only `UnifiedMemory` overrides this. + fn sqlite_conn(&self) -> Option>> { + None + } } #[cfg(test)] diff --git a/src/openhuman/tools/ops.rs b/src/openhuman/tools/ops.rs index 5912cdbde..9aafb2662 100644 --- a/src/openhuman/tools/ops.rs +++ b/src/openhuman/tools/ops.rs @@ -228,6 +228,22 @@ pub fn all_tools_with_runtime( root_config.curl.timeout_secs, ))); + // Phase 3 STM recall — on-demand cross-thread episodic search tool. + // Feature-gated on `learning.stm_recall_enabled` (default true) so the + // tool surface and the preemptive prompt injection are enabled/disabled + // together. `session_id` is not known at tool-build time; exclude-own- + // session is enforced by the preemptive first-turn injection in turn.rs + // (the on-demand tool intentionally uses an empty exclude_session). + if root_config.learning.stm_recall_enabled { + tools.push(Box::new( + crate::openhuman::memory::stm_recall::tool::StmRecallTool::new( + memory.clone(), + String::new(), + None, + ), + )); + } + // gitbooks — answers questions about OpenHuman by calling the // GitBook MCP server. Two tools mirroring the upstream MCP tools. if root_config.gitbooks.enabled {