feat(memory): activate chat memory pipeline — Archivist wiring, LLM recaps, cross-thread STM recall, segment-granular tree, unified compaction (#2175)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
sanil-23
2026-05-19 01:30:59 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 7741c5815c
commit 341ce42ff1
20 changed files with 4070 additions and 168 deletions
File diff suppressed because it is too large Load Diff
@@ -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<Mutex<Connection>> {
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::<Vec<_>>()
);
}
// ── 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<String> {
Ok("stub recap: discussed Rust ownership model".to_string())
}
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> anyhow::Result<String> {
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<Vec<f32>> {
// 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<Mutex<Connection>>) -> 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<Mutex<Connection>>, 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::<Vec<_>>()
);
}
/// 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::<Vec<_>>()
);
}
/// `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}"
);
}
+125 -2
View File
@@ -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<Arc<crate::openhuman::agent::harness::archivist::ArchivistHook>>,
) -> 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<dyn crate::openhuman::context::Summarizer> =
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<dyn crate::openhuman::context::Summarizer> =
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<crate::openhuman::agent::harness::archivist::ArchivistHook>,
> = 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<dyn crate::openhuman::agent::hooks::PostTurnHook>);
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()
}
}
+94 -2
View File
@@ -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;
@@ -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<Arc<ArchivistHook>>,
/// Names of every tool currently in [`Agent::tools`] that was
/// produced by [`crate::openhuman::tools::orchestrator_tools::collect_orchestrator_tools`]
/// (i.e. `delegate_<toolkit>` skill tools and archetype-delegation
@@ -220,6 +227,18 @@ pub struct AgentBuilder {
/// to a `SubagentPayloadSummarizer` instance.
pub(super) payload_summarizer:
Option<Arc<dyn crate::openhuman::agent::harness::payload_summarizer::PayloadSummarizer>>,
/// 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<Arc<ArchivistHook>>,
/// 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 {
+52
View File
@@ -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(),
}
}
}
+22
View File
@@ -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
+2
View File
@@ -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,
@@ -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<ArchivistHook>,
/// 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<dyn Summarizer>,
/// 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<ArchivistHook>,
session_id: String,
inner: Arc<dyn Summarizer>,
) -> 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<ConversationMessage>,
model: &str,
) -> Result<SummaryStats> {
// ── 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<ConversationMessage> = 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::<usize>()
}
ConversationMessage::ToolResults(results) => results.iter().map(|r| r.content.len()).sum(),
}
}
#[cfg(test)]
#[path = "segment_recap_summarizer_tests.rs"]
mod tests;
@@ -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<Mutex<Connection>> {
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<String> {
Ok("rolling recap: discussed memory safety in Rust".to_string())
}
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> Result<String> {
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<String> {
anyhow::bail!("stub LLM unavailable")
}
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> Result<String> {
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<Vec<f32>> {
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<usize>,
}
impl RecordingSummarizer {
fn new() -> Arc<Self> {
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<ConversationMessage>,
_model: &str,
) -> Result<SummaryStats> {
*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<dyn Summarizer>,
)
.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<dyn Summarizer>,
)
.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<dyn Summarizer>,
)
.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<dyn Summarizer>,
)
.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<dyn Summarizer>,
)
.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"
);
}
+5 -1
View File
@@ -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
+1
View File
@@ -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;
@@ -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;
+47
View File
@@ -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::*;
+629
View File
@@ -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<i64>,
updated_at: f64,
cosine: f32,
},
/// A raw episodic turn from Arm 1 (FTS5 keyword match).
EpisodicTurn {
id: Option<i64>,
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<StmItem>,
/// 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<f32>`.
///
/// 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<f32> {
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<String>,
start_episodic_id: i64,
end_episodic_id: Option<i64>,
updated_at: f64,
vector: Vec<f32>,
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<Mutex<Connection>>,
exclude_session: &str,
model_signature: Option<&str>,
cutoff_ts: f64,
) -> anyhow::Result<Vec<SegmentEmbeddingRow>> {
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<SegmentEmbeddingRow> = 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<SegmentEmbeddingRow> = stmt
.query_map(
params![
exclude_session,
cutoff_ts,
sig,
RECENCY_WINDOW_MAX_SEGMENTS as i64
],
|row| {
let vector_bytes: Vec<u8> = 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::<Result<Vec<_>, _>>()?;
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<SegmentEmbeddingRow> = stmt
.query_map(
params![
exclude_session,
cutoff_ts,
RECENCY_WINDOW_MAX_SEGMENTS as i64
],
|row| {
let vector_bytes: Vec<u8> = 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::<Result<Vec<_>, _>>()?;
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<Mutex<Connection>>,
exclude_session: &str,
cutoff_ts: f64,
limit: usize,
) -> anyhow::Result<Vec<EpisodicEntry>> {
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::<Result<Vec<_>, _>>()?;
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<Mutex<Connection>>,
opts: &StmRecallOpts<'_>,
query_embedding: Option<&[f32]>,
) -> anyhow::Result<StmRecallBlock> {
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<StmItem> = Vec::new();
let mut segment_spans: Vec<(i64, Option<i64>)> = 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<StmItem> = Vec::new();
let raw_episodic: Vec<EpisodicEntry> = 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<StmItem> = segment_items.into_iter().take(MAX_SEGMENT_RECAPS).collect();
let ep_capped: Vec<StmItem> = 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<StmItem> = 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<StmItem> = 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;
@@ -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<Mutex<Connection>> {
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<Mutex<Connection>>,
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<Mutex<Connection>>,
segment_id: &str,
session_id: &str,
start_id: i64,
end_id: i64,
summary: &str,
embedding: Option<Vec<f32>>,
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<u8> = 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<f32> = (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<String> {
Ok("RECAP: stub LLM summary of the segment.".to_string())
}
async fn chat_for_text(&self, _prompt: &ChatPrompt) -> anyhow::Result<String> {
Ok("RECAP: stub LLM summary of the segment.".to_string())
}
}
use crate::openhuman::memory::tree::score::embed::InertEmbedder;
let chat_provider: Arc<dyn crate::openhuman::memory::tree::chat::ChatProvider> =
Arc::new(StubChat);
let embedder: Arc<dyn crate::openhuman::memory::tree::score::embed::Embedder> =
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"
);
}
}
+200
View File
@@ -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<dyn Memory>,
/// 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<String>,
}
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<dyn Memory>,
session_id: String,
model_signature: Option<String>,
) -> 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<ToolResult> {
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<dyn Memory>) {
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()
);
}
}
@@ -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<std::sync::Arc<parking_lot::Mutex<rusqlite::Connection>>> {
Some(std::sync::Arc::clone(&self.conn))
}
}
#[cfg(test)]
+14
View File
@@ -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<Arc<Mutex<Connection>>> {
None
}
}
#[cfg(test)]
+16
View File
@@ -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 {