From 00224d1441ec397a0002e7cb64e5fcfcd340e39b Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Wed, 1 Jul 2026 18:09:44 +0300 Subject: [PATCH] split agent maintenance modules --- .beads/interactions.jsonl | 1 + .beads/issues.jsonl | 1 + klbr-core/src/agent.rs | 1007 +---------------------------- klbr-core/src/agent/compaction.rs | 375 +++++++++++ klbr-core/src/agent/recall.rs | 67 ++ klbr-core/src/agent/reflection.rs | 213 ++++++ klbr-core/src/agent/tests.rs | 364 +++++++++++ 7 files changed, 1032 insertions(+), 996 deletions(-) create mode 100644 klbr-core/src/agent/compaction.rs create mode 100644 klbr-core/src/agent/recall.rs create mode 100644 klbr-core/src/agent/reflection.rs create mode 100644 klbr-core/src/agent/tests.rs diff --git a/.beads/interactions.jsonl b/.beads/interactions.jsonl index 4cedded..889fa85 100644 --- a/.beads/interactions.jsonl +++ b/.beads/interactions.jsonl @@ -98,3 +98,4 @@ {"id":"int-706d5ee3","kind":"field_change","created_at":"2026-07-01T09:35:51.422527514Z","actor":"dawn","issue_id":"klbr-6m2","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Split App.svelte into ConnectionRail, ChatHeader, Composer, AddDaemonModal, ReflinkPopover, ResolutionTraces, and shared UI types; verified svelte-check, build, diff check, and local Vite smoke."}} {"id":"int-6a5bb5cc","kind":"field_change","created_at":"2026-07-01T09:48:37.465114937Z","actor":"dawn","issue_id":"klbr-wj2","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Added real folgezettel usage trace report command with reuse, revision, and review flags."}} {"id":"int-f26ee4d7","kind":"field_change","created_at":"2026-07-01T14:33:43.153698998Z","actor":"dawn","issue_id":"klbr-pic","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Split turn streaming/tool execution into agent/turn.rs and moved local_send, wait_and_continue, and restart runtime behavior into tool modules while preserving workspace tests."}} +{"id":"int-b41fd926","kind":"field_change","created_at":"2026-07-01T15:09:11.257989689Z","actor":"dawn","issue_id":"klbr-9lt","extra":{"field":"status","new_value":"closed","old_value":"in_progress","reason":"Moved passive recall, compaction, reflection, and agent tests into agent submodules; agent.rs now owns only the outer runtime loop and shared helpers."}} diff --git a/.beads/issues.jsonl b/.beads/issues.jsonl index da372b4..566e84c 100644 --- a/.beads/issues.jsonl +++ b/.beads/issues.jsonl @@ -41,6 +41,7 @@ {"_type":"issue","id":"klbr-wmz.2","title":"Route runtime passive recall through the canonical memory pipeline","description":"Bench retrieval now goes through MemoryPipeline, but runtime passive recall in klbr-core/src/agent.rs still embeds the prompt, calls MemoryStore::get_searchable, runs retrieval::retrieve_exact over legacy memories, then injects Context memory packets. That bypasses fts, exact refs, markdown notes, graph expansion, and the canonical lane/lifecycle path the docs describe.","design":"Avoid duplicating retrieval logic in agent.rs. Either make MemoryPipeline usable by AgentRuntime or extract a shared retrieval facade that both MemoryPipeline and runtime passive recall call.","acceptance_criteria":"Runtime passive recall uses the same lane-aware canonical retrieval and context packet assembly policy as the production pipeline; recalled packets can include fts/exact/dense/graph candidates from refs and markdown notes; archived/tombstoned/suppressed refs do not leak; klbr-core/src/instructions.md matches the actual memory packet format; tests or a focused integration fixture cover passive recall from a markdown note and from an explicit ref.","notes":"Runtime passive recall now calls MemoryPipeline::retrieve_evidence and injects shared EvidencePacket XML via Context::inject_evidence_packets; no-model runtime packet fixture covers turn-window expansion. Remaining acceptance is blocked on klbr-wmz.1 because dense search still starts from legacy memory rows before ref mapping.","status":"closed","priority":1,"issue_type":"feature","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-26T17:53:19Z","created_by":"dawn","updated_at":"2026-06-26T19:44:55Z","started_at":"2026-06-26T19:37:32Z","closed_at":"2026-06-26T19:44:55Z","close_reason":"Completed: runtime passive recall now uses MemoryPipeline/EvidencePlanner packets, instructions document current packet XML, and fixtures cover turn-window recall plus explicit markdown refs; dense canonical dependency completed in klbr-wmz.1.","labels":["architecture","memory","retrieval","runtime"],"dependencies":[{"issue_id":"klbr-wmz.2","depends_on_id":"klbr-wmz","type":"parent-child","created_at":"2026-06-26T20:53:19Z","created_by":"dawn","metadata":"{}"},{"issue_id":"klbr-wmz.2","depends_on_id":"klbr-wmz.1","type":"blocks","created_at":"2026-06-26T22:37:53Z","created_by":"dawn","metadata":"{}"}],"dependency_count":1,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-wmz.1","title":"Move dense retrieval onto canonical refs and embedding_items","description":"Current source still has dense retrieval on the legacy memory-row surface: klbr-core/src/pipeline.rs search_dense calls MemoryStore::get_searchable and retrieval::retrieve_exact, while fts/exact retrieval uses canonical refs and promptable_text. The schema already has refs, promptable_text, markdown_note_chunks, and embedding_items, so dense retrieval should not be the odd path out.","design":"Prefer a ref-native embedding index backed by embedding_items. Backfill embeddings from promptable_text, keep memory-id aliases as compatibility aliases, and make klbr/full versus dense-only profiles exercise the same canonical identity layer as fts and exact retrieval.","acceptance_criteria":"Dense candidate generation works over canonical ref ids for memories, turn chunks, markdown note chunks, episode notes, profile notes, and procedural notes; lane and lifecycle filtering come from refs/ref_metadata instead of memory tags alone; benchmark traces return canonical refs for dense hits; regression tests cover a markdown-note-only hit and a tombstoned/suppressed ref not leaking through dense search.","status":"closed","priority":1,"issue_type":"feature","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-26T17:53:11Z","created_by":"dawn","updated_at":"2026-06-26T19:43:27Z","started_at":"2026-06-26T19:38:02Z","closed_at":"2026-06-26T19:43:27Z","close_reason":"Completed: dense candidate generation now lazily backfills embedding_items from active promptable refs, scores canonical ref embeddings directly, and tests markdown-note dense hits plus suppressed-ref filtering.","labels":["architecture","memory","refs","retrieval"],"dependencies":[{"issue_id":"klbr-wmz.1","depends_on_id":"klbr-wmz","type":"parent-child","created_at":"2026-06-26T20:53:11Z","created_by":"dawn","metadata":"{}"}],"dependency_count":0,"dependent_count":1,"comment_count":0} {"_type":"issue","id":"klbr-wmz","title":"Finish memory architecture follow-through","description":"Tracks the remaining memory architecture work identified from docs/memory-arch.md, docs/memory-benches.md, docs/memory-implementation-status.md, and current klbr-core/klbr-bench source. Current status says pipeline, typed memory packets, markdown notes, edge mirroring, lifecycle projection, and benchmark runner integration exist; this epic is for gaps still present in source/docs.","acceptance_criteria":"Close when the child issues are complete, docs/memory-implementation-status.md is updated from current verification, and the architecture docs no longer point at missing or stale follow-up work.","status":"closed","priority":1,"issue_type":"epic","owner":"90008@klbr.net","created_at":"2026-06-26T17:52:54Z","created_by":"dawn","updated_at":"2026-06-26T23:45:35Z","closed_at":"2026-06-26T23:45:35Z","close_reason":"All 19 memory architecture child issues are closed; docs/status were updated from current verification; remaining official evaluator run is tracked separately as external blocked klbr-1yn.","labels":["architecture","memory"],"dependency_count":0,"dependent_count":0,"comment_count":0} +{"_type":"issue","id":"klbr-9lt","title":"Split agent maintenance modules out of agent.rs","description":"agent.rs still contains compaction, reflection, passive-recall helpers, and their tests after the turn runtime split. Move compaction and reflection into agent submodules and relocate tests with their owning code so agent.rs only owns Agent construction, the outer event loop, interrupt dispatch, and thin shared helpers.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-07-01T15:05:10Z","created_by":"dawn","updated_at":"2026-07-01T15:09:11Z","started_at":"2026-07-01T15:05:11Z","closed_at":"2026-07-01T15:09:11Z","close_reason":"Moved passive recall, compaction, reflection, and agent tests into agent submodules; agent.rs now owns only the outer runtime loop and shared helpers.","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-pic","title":"Split agent turn and tool runtime out of agent.rs","description":"agent.rs is still too large after the previous run_turn decomposition. Move the turn streaming/tool iteration machinery and tool-specific runtime behaviors such as wait_and_continue/local_send/restart handling out of the main Agent file while preserving behavior and tests.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-07-01T14:27:20Z","created_by":"dawn","updated_at":"2026-07-01T14:33:43Z","started_at":"2026-07-01T14:27:24Z","closed_at":"2026-07-01T14:33:43Z","close_reason":"Split turn streaming/tool execution into agent/turn.rs and moved local_send, wait_and_continue, and restart runtime behavior into tool modules while preserving workspace tests.","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-u03","title":"bridge: refactor hardcoded /home/mayer paths and make binary path configurable","description":"antigravity-bridge/src/auth.rs contains hardcoded /home/mayer home directories and hardcoded binary path BINARY_PATH. Use dirs::home_dir() and make the binary path configurable via environment variables or configuration.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:26Z","created_by":"dawn","updated_at":"2026-06-30T21:59:42Z","started_at":"2026-06-30T21:58:02Z","closed_at":"2026-06-30T21:59:42Z","close_reason":"Replaced hardcoded /home/mayer paths in antigravity bridge auth with dirs::home_dir()-based helpers, added ANTIGRAVITY_BRIDGE_CONFIG_DIR for token/cache files, and made the agy binary path configurable via ANTIGRAVITY_BRIDGE_BINARY_PATH or AGY_BINARY_PATH with ~/.local/bin/agy fallback. Verified cargo fmt --check and cargo test -p antigravity-bridge (9 passed).","dependency_count":0,"dependent_count":0,"comment_count":0} {"_type":"issue","id":"klbr-noi","title":"nix: fix broken flake.nix referencing non-existent default.nix","description":"flake.nix references ./default.nix to build the default package, but default.nix does not exist in the repository. Refactor flake.nix to output the derivation directly using nix-cargo-integration outputs.","status":"closed","priority":2,"issue_type":"task","assignee":"dawn","owner":"90008@klbr.net","created_at":"2026-06-30T21:42:19Z","created_by":"dawn","updated_at":"2026-06-30T22:01:41Z","started_at":"2026-06-30T22:00:06Z","closed_at":"2026-06-30T22:01:41Z","close_reason":"Replaced the missing ./default.nix package with config.nci.outputs.\"klbr-daemon\".packages.release. Verified nix eval .#packages.x86_64-linux.default.name returns klbr-daemon and nix flake show succeeds without the default.nix error.","dependency_count":0,"dependent_count":0,"comment_count":0} diff --git a/klbr-core/src/agent.rs b/klbr-core/src/agent.rs index d2dca56..4ac868c 100644 --- a/klbr-core/src/agent.rs +++ b/klbr-core/src/agent.rs @@ -1,26 +1,30 @@ use anyhow::Result; -use std::collections::HashSet; use tokio::sync::{broadcast, mpsc}; use crate::MetricsSnapshot; -use crate::config::MemoryConfig; use crate::{ config::Config, context::Context, evidence::EvidencePacket, - harness_block::{self, HarnessBlock, DEFAULT_FORMAT_STYLE}, + harness_block::{self, DEFAULT_FORMAT_STYLE}, interrupt::Interrupt, memory::MemoryStore, - models::{LlmClient, LlmEvent, Message, Role}, - pipeline::{BenchQuery, ContextBudget, MemoryPipeline}, + models::{LlmClient, Role}, router::{RouteDecision, Router}, tools::{self, ToolContext}, - AgentEvent, CompactionRecord, + AgentEvent, }; use crate::tools::format_duration_human; +use compaction::compact; +use recall::recall_evidence_packets; +use reflection::build_reflection_messages; + +mod compaction; +mod recall; +mod reflection; mod turn; const MAX_CONSECUTIVE_STREAM_FAILURES: usize = 5; @@ -713,994 +717,5 @@ fn refresh_context_soul(memory: &MemoryStore, runtime_soul: &str, ctx: &mut Cont save_context_snapshot(memory, ctx); } -async fn recall_evidence_packets( - config: &MemoryConfig, - memory: &MemoryStore, - llm: &LlmClient, - query: &str, - reference_time: Option, - already_recalled: &HashSet, - profile: &str, - output: &broadcast::Sender, -) -> Result> { - let pipeline = - MemoryPipeline::new(memory.clone(), llm.clone(), config.clone()).with_profile(profile); - let budget = ContextBudget { - max_tokens: config.top_k.max(1).saturating_mul(900), - top_k: config.candidate_k.max(config.top_k).max(1), - graph_depth: 1, - }; - let retrieval = pipeline - .retrieve_evidence( - BenchQuery { - query_id: "runtime-passive-recall".to_string(), - text: query.to_string(), - reference_time, - op_hint: None, - }, - budget, - ) - .await?; - let selected = retrieval - .packets - .into_iter() - .filter(|packet| { - packet - .refs - .iter() - .any(|ref_id| !already_recalled.contains(ref_id)) - }) - .take(config.top_k) - .collect::>(); - - if !retrieval.candidates.is_empty() || !selected.is_empty() || !retrieval.route.archival_allowed - { - let _ = output.send(AgentEvent::Status(format!( - "memory packets: {} candidates, {} selected, route={}, archival={}", - retrieval.candidates.len(), - selected.len(), - retrieval.route.reason, - retrieval.route.archival_allowed - ))); - } - - Ok(selected) -} - -async fn compact( - tool_ctx: &ToolContext, - ctx: &mut Context, - keep: usize, - target_tokens: usize, - output: &broadcast::Sender, -) -> Result<()> { - let _ = output.send(AgentEvent::CompactionStarted); - let _ = output.send(AgentEvent::Status("compacting...".into())); - let result = compact_inner(tool_ctx, ctx, keep, target_tokens, output).await; - let _ = output.send(AgentEvent::CompactionDone); - result -} - -async fn compact_inner( - tool_ctx: &ToolContext, - ctx: &mut Context, - keep: usize, - target_tokens: usize, - output: &broadcast::Sender, -) -> Result<()> { - // run reflection before draining so the agent can curate memories - if let Err(e) = reflect(tool_ctx, ctx, output).await { - tracing::warn!(err = %e, "reflection failed"); - } - - let cut = ctx.drain_prefix_len(keep, target_tokens); - if cut == 0 { - let _ = output.send(AgentEvent::Status("nothing to compact".into())); - return Ok(()); - } - - let drained = ctx.turns[..cut].to_vec(); - let compacted_messages = compaction_source_messages(&drained); - if compacted_messages.is_empty() { - let drained_now = ctx.drain_oldest(keep, target_tokens); - debug_assert_eq!(drained_now.len(), drained.len()); - save_context_snapshot(&tool_ctx.memory, ctx); - let _ = output.send(AgentEvent::Status("nothing to compact".into())); - return Ok(()); - } - - let source_memory_ids = recalled_memory_ids(&compacted_messages); - let prompt = build_compaction_messages(&ctx.system, &compacted_messages); - - let _ = output.send(AgentEvent::Status("forming recollection...".into())); - let _ = output.send(AgentEvent::CompactionToken( - COMPACTION_ASSISTANT_PREFILL.to_string(), - )); - let recollection_completion = stream_compaction_recollection(tool_ctx, &prompt, output).await?; - let recollection = complete_compaction_recollection(&recollection_completion); - if recollection.trim().is_empty() { - return Err(anyhow::anyhow!( - "compaction recollection model returned empty output" - )); - } - - let now = unix_timestamp(); - let _ = output.send(AgentEvent::Status("storing recollection...".into())); - let emb = tool_ctx.llm.embed(&recollection).await?; - let recollection_id = tool_ctx - .memory - .store_with_metadata(&crate::mvp::MemoryRecordInput { - memory_id: None, - namespace: "default".to_string(), - layer: crate::mvp::MemoryLayer::L1, - text: recollection.clone(), - event_time: now, - ingest_time: now, - embedding_model: tool_ctx.llm.config.embedder.model.clone(), - embedding_dim: emb.len(), - embedding_version: "runtime".to_string(), - status: crate::mvp::MemoryStatus::Active, - source_ref: Some(format!("system:compaction:{now}")), - tags: vec!["compaction_recollection".to_string()], - pinned: false, - embedding: emb, - })?; - - let drained_now = ctx.drain_oldest(keep, target_tokens); - debug_assert_eq!(drained_now.len(), drained.len()); - let kept_turns = ctx.turn_count(); - if !recollection.is_empty() { - for source_id in source_memory_ids { - let _ = tool_ctx.memory.add_edge(&crate::mvp::MemoryEdgeInput { - from_memory_id: recollection_id, - to_memory_id: source_id, - edge_type: crate::mvp::MemoryEdgeType::DerivedFrom, - metadata: serde_json::json!({ - "source": "recalled_memory_in_compacted_context" - }), - }); - } - } - - if !recollection.is_empty() { - if ctx.turns.first().map(|m| &m.role) == Some(&Role::Assistant) { - let first = &mut ctx.turns[0]; - let new_content = if first.content_str().is_empty() { - recollection.clone() - } else { - format!("{}\n\n{}", recollection, first.content_str()) - }; - first.content = Some(new_content); - } else { - ctx.turns - .insert(0, Message::assistant(recollection.clone())); - } - } - let record = CompactionRecord { - recollection: recollection.clone(), - compacted: compaction_transcript(&compacted_messages), - memory_id: recollection_id, - drained_turns: drained.len(), - kept_turns, - }; - match serde_json::to_string(&record) { - Ok(content) => { - if let Err(e) = tool_ctx.memory.log_turn("compaction", &content, None) { - tracing::warn!(err = %e, "failed to persist compaction record"); - } - } - Err(e) => tracing::warn!(err = %e, "failed to serialize compaction record"), - } - let _ = output.send(AgentEvent::CompactionRecord(record)); - - save_context_snapshot(&tool_ctx.memory, ctx); - - Ok(()) -} - -async fn stream_compaction_recollection( - tool_ctx: &ToolContext, - prompt: &[Message], - output: &broadcast::Sender, -) -> Result { - let compact_registry = tools::memory_tools(); - let mut messages = prompt.to_vec(); - - for _ in 0..20 { - let (tok_tx, mut tok_rx) = mpsc::channel(512); - let llm = tool_ctx.compaction_llm.clone(); - let prompt = messages.clone(); - let definitions = compact_registry.definitions(); - let stream_task = - tokio::spawn(async move { llm.stream(&prompt, &definitions, tok_tx).await }); - let mut iteration_response = String::new(); - let mut iteration_thinking = String::new(); - let mut tool_calls = vec![]; - - while let Some(ev) = tok_rx.recv().await { - match ev { - LlmEvent::Token(token) => { - iteration_response.push_str(&token); - let _ = output.send(AgentEvent::CompactionToken(token)); - } - LlmEvent::ThinkToken(token) => { - iteration_thinking.push_str(&token); - let _ = output.send(AgentEvent::CompactionThinkToken(token)); - } - LlmEvent::ToolCalls(calls) => { - tool_calls = calls; - } - LlmEvent::Usage(_) => {} - LlmEvent::ToolCallStarted => {} - } - } - - match stream_task.await { - Ok(Ok(())) => {} - Ok(Err(e)) => return Err(e), - Err(e) => { - return Err(anyhow::anyhow!( - "compaction recollection stream task panicked/cancelled: {e}" - )) - } - } - - if tool_calls.is_empty() { - return Ok(iteration_response); - } - - let content = (!iteration_response.is_empty()).then_some(iteration_response); - let reasoning = (!iteration_thinking.is_empty()).then_some(iteration_thinking); - messages.push(Message::with_tool_calls( - tool_calls.clone(), - content, - reasoning, - )); - - for call in &tool_calls { - let name = call.function.name.clone(); - let args = call.function.arguments.clone(); - let _ = output.send(AgentEvent::ToolCall { - name: name.clone(), - args, - }); - let result = compact_registry.execute(call, tool_ctx).await; - let _ = output.send(AgentEvent::ToolResult { - name, - content: result.clone(), - }); - messages.push(Message::tool_result(&call.id, result)); - } - } - - Err(anyhow::anyhow!( - "compaction recollection exceeded tool iteration limit" - )) -} - -const COMPACTION_ASSISTANT_PREFILL: &str = "from this conversation —"; - -fn build_compaction_messages(system: &[Message], drained: &[Message]) -> Vec { - let mut messages = Vec::with_capacity(system.len() + drained.len() + 2); - messages.extend(system.iter().cloned()); - messages.extend(drained.iter().cloned()); - messages.push(Message::user(build_compaction_prompt())); - messages.push(Message::assistant(COMPACTION_ASSISTANT_PREFILL)); - messages -} - -fn build_compaction_prompt() -> String { - include_str!("compaction.md").to_string() -} - -fn complete_compaction_recollection(completion: &str) -> String { - let completion = completion.trim_end(); - if completion.trim().is_empty() { - return String::new(); - } - - let trimmed_start = completion.trim_start(); - if trimmed_start.starts_with(COMPACTION_ASSISTANT_PREFILL) { - return trimmed_start.to_string(); - } - - let separator = completion - .chars() - .next() - .is_some_and(char::is_whitespace) - .then_some("") - .unwrap_or(" "); - format!("{COMPACTION_ASSISTANT_PREFILL}{separator}{completion}") -} - -fn compaction_source_messages(messages: &[Message]) -> Vec { - messages - .iter() - .filter(|message| !is_reflection_maintenance_message(message)) - .cloned() - .collect() -} - -fn is_reflection_maintenance_message(message: &Message) -> bool { - message.role == Role::Reflection - || message.content.as_deref().is_some_and(|content| { - let trimmed = content.trim_start(); - trimmed.starts_with("") - || trimmed.starts_with("[[automated_reflection]]") - }) -} - -fn compaction_transcript(messages: &[Message]) -> String { - messages - .iter() - .filter_map(compaction_message_line) - .collect::>() - .join("\n") -} - -fn compaction_message_line(message: &Message) -> Option { - let mut parts = Vec::new(); - if let Some(reasoning) = message.reasoning_content.as_deref() { - let reasoning = reasoning.trim(); - if !reasoning.is_empty() { - parts.push(harness_block::format_reflection_block( - HarnessBlock::Reasoning, - reasoning, - None, - DEFAULT_FORMAT_STYLE, - )); - } - } - if let Some(content) = message.content.as_deref() { - let content = content.trim(); - if !content.is_empty() { - parts.push(content.to_string()); - } - } - if let Some(calls) = message - .tool_calls - .as_ref() - .filter(|calls| !calls.is_empty()) - { - let calls = calls - .iter() - .map(|call| { - let args = call.function.arguments.trim(); - if args.is_empty() { - call.function.name.clone() - } else { - format!("{}({args})", call.function.name) - } - }) - .collect::>() - .join(" | "); - parts.push(harness_block::format_reflection_block( - HarnessBlock::ToolCalls, - &calls, - None, - DEFAULT_FORMAT_STYLE, - )); - } - if let Some(tool_call_id) = message.tool_call_id.as_deref() { - parts.push(harness_block::format_reflection_block( - HarnessBlock::ToolCallId, - tool_call_id, - None, - DEFAULT_FORMAT_STYLE, - )); - } - - if parts.is_empty() { - None - } else { - Some(format!("{}: {}", message.role, parts.join(" "))) - } -} - -fn recalled_memory_ids(messages: &[Message]) -> Vec { - let mut ids = Vec::new(); - for message in messages { - let Some(content) = message.content.as_deref() else { - continue; - }; - // Bracket style - for part in content.split("[id:").skip(1) { - let Some((raw_id, _)) = part.split_once(']') else { - continue; - }; - let Ok(id) = raw_id.trim().parse::() else { - continue; - }; - if !ids.contains(&id) { - ids.push(id); - } - } - // XML style - for part in content.split("id=\"").skip(1) { - let Some((raw_id, _)) = part.split_once('"') else { - continue; - }; - let Ok(id) = raw_id.trim().parse::() else { - continue; - }; - if !ids.contains(&id) { - ids.push(id); - } - } - } - ids -} - -/// reflection loop: let the agent review and curate its memories -async fn reflect( - tool_ctx: &ToolContext, - ctx: &Context, - output: &broadcast::Sender, -) -> Result<()> { - let reflect_registry = tools::memory_tools(); - let mut msgs = build_reflection_messages(tool_ctx, ctx); - let mut reflection_log = Vec::new(); - - let _ = output.send(AgentEvent::ReflectStarted); - - for _ in 0..100 { - let (tok_tx, mut tok_rx) = mpsc::channel(512); - let llm2 = tool_ctx.llm.clone(); - let msgs_snap = msgs.clone(); - let defs_snap = reflect_registry.definitions(); - let stream_task = - tokio::spawn(async move { llm2.stream(&msgs_snap, &defs_snap, tok_tx).await }); - - let mut tool_calls = vec![]; - let mut stream_failed = false; - let mut iteration_thinking = String::new(); - let mut iteration_response = String::new(); - - while let Some(ev) = tok_rx.recv().await { - match ev { - LlmEvent::ToolCalls(calls) => tool_calls = calls, - LlmEvent::Token(t) => { - iteration_response.push_str(&t); - let _ = output.send(AgentEvent::Token(t)); - } - LlmEvent::ThinkToken(t) => { - iteration_thinking.push_str(&t); - let _ = output.send(AgentEvent::ThinkToken(t)); - } - _ => {} - } - } - match stream_task.await { - Ok(Ok(())) => {} - Ok(Err(e)) => { - let msg = format!("reflection llm stream failed: {e}"); - tracing::error!(%msg); - let _ = output.send(AgentEvent::Error(msg)); - stream_failed = true; - } - Err(e) => { - let msg = format!("reflection llm stream task panicked/cancelled: {e}"); - tracing::error!(%msg); - let _ = output.send(AgentEvent::Error(msg)); - stream_failed = true; - } - } - - if (tool_calls.is_empty() && iteration_response.is_empty()) || stream_failed { - break; - } - - let reasoning = if iteration_thinking.is_empty() { - None - } else { - reflection_log.push(harness_block::format_reflection_block( - HarnessBlock::Reasoning, - &iteration_thinking, - None, - DEFAULT_FORMAT_STYLE, - )); - Some(iteration_thinking) - }; - let content = if iteration_response.is_empty() { - None - } else { - reflection_log.push(harness_block::format_reflection_block( - HarnessBlock::Response, - &iteration_response, - None, - DEFAULT_FORMAT_STYLE, - )); - Some(iteration_response) - }; - - if !tool_calls.is_empty() { - msgs.push(Message::with_tool_calls( - tool_calls.clone(), - content, - reasoning, - )); - for call in &tool_calls { - let name = call.function.name.clone(); - let args = call.function.arguments.clone(); - reflection_log.push(harness_block::format_reflection_block( - HarnessBlock::ToolCall, - &args, - Some(&name), - DEFAULT_FORMAT_STYLE, - )); - let _ = output.send(AgentEvent::ToolCall { - name: name.clone(), - args: args.clone(), - }); - let result = reflect_registry.execute(call, tool_ctx).await; - reflection_log.push(harness_block::format_reflection_block( - HarnessBlock::ToolResult, - &result, - Some(&name), - DEFAULT_FORMAT_STYLE, - )); - let _ = output.send(AgentEvent::ToolResult { - name: name.clone(), - content: result.clone(), - }); - msgs.push(Message::tool_result(&call.id, &result)); - } - } else { - msgs.push(Message { - role: Role::Assistant, - content, - reasoning_content: reasoning, - tool_calls: None, - tool_call_id: None, - }); - break; - } - } - - if !reflection_log.is_empty() { - if let Err(e) = tool_ctx - .memory - .log_turn("reflection", &reflection_log.join("\n\n"), None) - { - tracing::warn!(err = %e, "failed to persist reflection record"); - } - } - let _ = output.send(AgentEvent::ReflectDone); - Ok(()) -} - -fn build_reflection_messages(tool_ctx: &ToolContext, ctx: &Context) -> Vec { - let pinned = tool_ctx.memory.pinned_memory_entries().unwrap_or_default(); - let unpinned = tool_ctx.memory.recent_unpinned(20).unwrap_or_default(); - - let pinned_text = if pinned.is_empty() { - "(none)".to_string() - } else { - pinned - .iter() - .map(|memory| { - harness_block::format_memory( - HarnessBlock::PinnedMemory, - memory.id, - &memory.tags, - &[], - false, - &memory.content, - DEFAULT_FORMAT_STYLE, - ) - }) - .collect::>() - .join("\n") - }; - - let unpinned_text = if unpinned.is_empty() { - "(none)".to_string() - } else { - unpinned - .iter() - .map(|(id, s, tags)| { - harness_block::format_memory( - HarnessBlock::UnpinnedMemory, - *id, - tags, - &[], - false, - s, - DEFAULT_FORMAT_STYLE, - ) - }) - .collect::>() - .join("\n") - }; - let reflection_prompt = - build_reflection_prompt(&pinned_text, &unpinned_text, chrono::Utc::now()); - - let mut messages = ctx.as_messages_with_refs(&tool_ctx.memory); - messages.push(Message::user(reflection_prompt)); - messages -} - -fn build_reflection_prompt( - pinned_text: &str, - unpinned_text: &str, - datetime: chrono::DateTime, -) -> String { - format!( - include_str!("reflection.md"), - pinned_text = pinned_text, - unpinned_text = unpinned_text, - date = datetime.date_naive(), - time = datetime.time(), - ) -} - #[cfg(test)] -mod tests { - use std::collections::HashSet; - - use super::{ - build_compaction_messages, build_compaction_prompt, build_reflection_prompt, - compaction_source_messages, compaction_transcript, complete_compaction_recollection, - needs_assistant_response, recall_evidence_packets, COMPACTION_ASSISTANT_PREFILL, - }; - use crate::models::Role; - use anyhow::Result; - use chrono::{TimeZone, Utc}; - use tempfile::NamedTempFile; - use tokio::sync::broadcast; - - use crate::{ - config::MemoryConfig, - context::Context, - memory::{MarkdownNoteInput, MemoryLane, MemoryStore}, - models::{LlmClient, Message, ModelsConfig, ToolCall, ToolCallFunction}, - pipeline::{BenchSession, BenchTurn, MemoryPipeline}, - }; - - fn tool_call(id: &str, name: &str) -> ToolCall { - ToolCall { - id: id.to_string(), - kind: "function".to_string(), - function: ToolCallFunction { - name: name.to_string(), - arguments: "{}".to_string(), - }, - } - } - - #[tokio::test] - async fn runtime_packet_recall_uses_pipeline_packet_window() -> Result<()> { - let tmp = NamedTempFile::new()?; - let store = MemoryStore::open(tmp.path().to_str().unwrap(), 4)?; - let llm = LlmClient::new(ModelsConfig::default()); - let mut config = MemoryConfig::default(); - config.top_k = 2; - config.candidate_k = 5; - - let pipeline = MemoryPipeline::new(store.clone(), llm.clone(), config.clone()) - .with_profile("raw-turns/fts-only"); - pipeline - .observe_session(BenchSession { - session_id: "play-session".to_string(), - timestamp: Some(1), - turns: vec![ - BenchTurn { - role: "user".to_string(), - content: "I recently went to a play at the local community theater." - .to_string(), - timestamp: Some(1), - }, - BenchTurn { - role: "assistant".to_string(), - content: "The play was The Glass Menagerie.".to_string(), - timestamp: Some(2), - }, - ], - }) - .await?; - - let (tx, _) = broadcast::channel(8); - let packets = recall_evidence_packets( - &config, - &store, - &llm, - "what play did I recently go to at the community theater?", - Some(3), - &HashSet::new(), - "raw-turns/fts-only", - &tx, - ) - .await?; - - assert!(packets - .iter() - .flat_map(|packet| &packet.bodies) - .any(|body| body.contains("The Glass Menagerie"))); - - let mut ctx = Context::new("soul", &[]); - let recalled_index = ctx.inject_evidence_packets(&packets, 900).unwrap(); - let rendered = ctx.turns[recalled_index] - .content - .as_deref() - .unwrap_or_default(); - assert!(rendered.contains("")); - assert!(rendered.contains("kind=\"turn_window\"")); - assert!(rendered.contains("The Glass Menagerie")); - let recalled_refs = ctx.passively_recalled_refs(); - assert!(packets - .iter() - .flat_map(|packet| &packet.refs) - .any(|ref_id| recalled_refs.contains(ref_id))); - - Ok(()) - } - - #[tokio::test] - async fn runtime_packet_recall_resolves_explicit_markdown_refs() -> Result<()> { - let tmp = NamedTempFile::new()?; - let store = MemoryStore::open(tmp.path().to_str().unwrap(), 4)?; - let llm = LlmClient::new(ModelsConfig::default()); - let mut config = MemoryConfig::default(); - config.top_k = 2; - config.candidate_k = 5; - - let note = store.upsert_markdown_note(&MarkdownNoteInput { - note_ref: Some("n_runtime_note".to_string()), - lane: MemoryLane::Semantic, - kind: "semantic_note".to_string(), - title: "runtime note".to_string(), - path: None, - body: "the explicit markdown answer is blue quartz".to_string(), - sources: vec![], - follow: None, - entities: vec![], - status: "active".to_string(), - frontmatter: serde_json::json!({}), - })?; - let chunk_alias = note - .chunk_refs - .first() - .expect("markdown note should create a chunk alias"); - - let (tx, _) = broadcast::channel(8); - let packets = recall_evidence_packets( - &config, - &store, - &llm, - &format!("what did [{}] say?", chunk_alias), - Some(3), - &HashSet::new(), - "raw-turns/fts-only", - &tx, - ) - .await?; - - assert!(packets - .iter() - .any(|packet| packet.signals.sources.contains(&"exact".to_string()))); - assert!(packets - .iter() - .flat_map(|packet| &packet.bodies) - .any(|body| body.contains("blue quartz"))); - - Ok(()) - } - - #[tokio::test] - async fn runtime_packet_recall_does_not_use_discourse_cue_filter() -> Result<()> { - let tmp = NamedTempFile::new()?; - let store = MemoryStore::open(tmp.path().to_str().unwrap(), 4)?; - let llm = LlmClient::new(ModelsConfig::default()); - let mut config = MemoryConfig::default(); - config.top_k = 2; - config.candidate_k = 5; - - let pipeline = MemoryPipeline::new(store.clone(), llm.clone(), config.clone()) - .with_profile("raw-turns/fts-only"); - pipeline - .observe_session(BenchSession { - session_id: "old-session".to_string(), - timestamp: Some(1), - turns: vec![BenchTurn { - role: "user".to_string(), - content: "yes do that was an old archival phrase".to_string(), - timestamp: Some(1), - }], - }) - .await?; - - let (tx, _) = broadcast::channel(8); - let packets = recall_evidence_packets( - &config, - &store, - &llm, - "yes do that", - Some(3), - &HashSet::new(), - "raw-turns/fts-only", - &tx, - ) - .await?; - - assert!(packets - .iter() - .flat_map(|packet| &packet.bodies) - .any(|body| body.contains("old archival phrase"))); - Ok(()) - } - - #[test] - fn reflection_prompt_is_automated_instruction() { - let prompt = build_reflection_prompt( - "(none)", - "(none)", - Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap(), - ); - - assert!(prompt.starts_with("")); - assert!(prompt.contains("## goal")); - assert!(prompt.contains("not a human user message")); - assert!(prompt.contains("current state / carry-forward")); - assert!(prompt.contains("errors and fixes")); - assert!(prompt.contains("contradictions and corrections")); - assert!(!prompt.contains("real soul instructions")); - } - - #[test] - fn compaction_messages_reuse_current_system_prefix() { - let system = vec![Message::system("runtime soul with pinned block")]; - let drained = vec![ - Message::user("old request"), - Message::assistant("old answer"), - ]; - - let messages = build_compaction_messages(&system, &drained); - - assert_eq!(messages.len(), 5); - assert_eq!(messages[0].role, "system"); - assert_eq!( - messages[0].content.as_deref(), - Some("runtime soul with pinned block") - ); - assert_eq!(messages[1].content.as_deref(), Some("old request")); - assert_eq!(messages[2].content.as_deref(), Some("old answer")); - assert_eq!(messages[3].role, "user"); - let maintenance_prompt = messages[3].content.as_deref().unwrap_or_default(); - assert!(maintenance_prompt.contains("a recollection of the context")); - assert!(maintenance_prompt.contains("carry forward what mattered")); - assert_eq!(messages[4].role, "assistant"); - assert_eq!( - messages[4].content.as_deref(), - Some(COMPACTION_ASSISTANT_PREFILL) - ); - } - - #[test] - fn compaction_prompt_is_reflection_instruction() { - let prompt = build_compaction_prompt(); - - assert!(prompt.contains("carry forward what mattered from this conversation")); - assert!(prompt.contains("dissent too")); - assert!(!prompt.contains("your dissent too")); - assert!(prompt.contains("no tools for this")); - assert!(!prompt.contains("{pinned_memories}")); - } - - #[test] - fn compaction_recollection_completes_assistant_prefill() { - assert_eq!( - complete_compaction_recollection("what mattered was the exact phrasing"), - "from this conversation — what mattered was the exact phrasing" - ); - assert_eq!( - complete_compaction_recollection(" carried forward a live task"), - "from this conversation — carried forward a live task" - ); - assert_eq!( - complete_compaction_recollection("from this conversation — already whole"), - "from this conversation — already whole" - ); - assert_eq!(complete_compaction_recollection(" "), ""); - } - - #[test] - fn compaction_transcript_keeps_tool_call_context() { - let messages = vec![ - Message::with_tool_calls( - vec![ToolCall { - id: "call_1".to_string(), - kind: "function".to_string(), - function: ToolCallFunction { - name: "read_file".to_string(), - arguments: "{\"path\":\"klbr-core/src/agent.rs\"}".to_string(), - }, - }], - Some("checking the file".to_string()), - Some("need context".to_string()), - ), - Message::tool_result("call_1", "file contents"), - ]; - - let transcript = compaction_transcript(&messages); - - assert!(transcript - .contains("assistant: \nneed context\n checking the file")); - assert!(transcript - .contains("read_file({\"path\":\"klbr-core/src/agent.rs\"})")); - assert!(transcript.contains("tool: file contents call_1")); - } - - #[test] - fn compaction_source_messages_excludes_reflection_prompts() { - let messages = vec![ - Message::user("real old turn"), - Message::assistant("[[automated_reflection]]\n\nreflect over this"), - Message::assistant( - "\n\nreflect over this", - ), - Message { - role: Role::Reflection, - content: Some("logged reflection".to_string()), - reasoning_content: None, - tool_calls: None, - tool_call_id: None, - }, - ]; - - let source = compaction_source_messages(&messages); - - assert_eq!(source.len(), 1); - assert_eq!(source[0].content.as_deref(), Some("real old turn")); - } - - #[test] - fn wait_timeout_result_does_not_need_another_assistant_response() { - let mut ctx = Context::new("soul", &[]); - ctx.push_assistant_tool_calls(vec![tool_call("call_1", "wait_and_continue")], None, None); - ctx.push_tool_result( - "call_1", - "waited 1min; no incoming event arrived; current time: 2026-06-14 23:07:36 +03:00; suspended", - ); - - assert!(!needs_assistant_response(&ctx)); - } - - #[test] - fn wait_tool_error_still_needs_assistant_response() { - let mut ctx = Context::new("soul", &[]); - ctx.push_assistant_tool_calls(vec![tool_call("call_1", "wait_and_continue")], None, None); - ctx.push_tool_result( - "call_1", - "error: duration is required (either 'timeout_ms' or 'seconds' must be specified)", - ); - - assert!(needs_assistant_response(&ctx)); - } - - #[test] - fn non_wait_tool_result_still_needs_assistant_response() { - let mut ctx = Context::new("soul", &[]); - ctx.push_assistant_tool_calls(vec![tool_call("call_1", "remember")], None, None); - ctx.push_tool_result("call_1", "stored"); - - assert!(needs_assistant_response(&ctx)); - } - - #[test] - fn test_format_duration_human() { - use std::time::Duration; - assert_eq!( - super::format_duration_human(Duration::from_millis(500)), - "500ms" - ); - assert_eq!(super::format_duration_human(Duration::from_secs(5)), "5sec"); - assert_eq!( - super::format_duration_human(Duration::from_secs(65)), - "1min 5sec" - ); - assert_eq!( - super::format_duration_human(Duration::from_secs(3665)), - "1hour 1min 5sec" - ); - } -} +mod tests; diff --git a/klbr-core/src/agent/compaction.rs b/klbr-core/src/agent/compaction.rs new file mode 100644 index 0000000..448ed1b --- /dev/null +++ b/klbr-core/src/agent/compaction.rs @@ -0,0 +1,375 @@ +use anyhow::Result; +use tokio::sync::{broadcast, mpsc}; + +use super::{reflection::reflect, save_context_snapshot, unix_timestamp}; +use crate::{ + context::Context, + harness_block::{self, HarnessBlock, DEFAULT_FORMAT_STYLE}, + models::{LlmEvent, Message, Role}, + tools::{self, ToolContext}, + AgentEvent, CompactionRecord, +}; + +pub(super) async fn compact( + tool_ctx: &ToolContext, + ctx: &mut Context, + keep: usize, + target_tokens: usize, + output: &broadcast::Sender, +) -> Result<()> { + let _ = output.send(AgentEvent::CompactionStarted); + let _ = output.send(AgentEvent::Status("compacting...".into())); + let result = compact_inner(tool_ctx, ctx, keep, target_tokens, output).await; + let _ = output.send(AgentEvent::CompactionDone); + result +} + +async fn compact_inner( + tool_ctx: &ToolContext, + ctx: &mut Context, + keep: usize, + target_tokens: usize, + output: &broadcast::Sender, +) -> Result<()> { + // run reflection before draining so the agent can curate memories + if let Err(e) = reflect(tool_ctx, ctx, output).await { + tracing::warn!(err = %e, "reflection failed"); + } + + let cut = ctx.drain_prefix_len(keep, target_tokens); + if cut == 0 { + let _ = output.send(AgentEvent::Status("nothing to compact".into())); + return Ok(()); + } + + let drained = ctx.turns[..cut].to_vec(); + let compacted_messages = compaction_source_messages(&drained); + if compacted_messages.is_empty() { + let drained_now = ctx.drain_oldest(keep, target_tokens); + debug_assert_eq!(drained_now.len(), drained.len()); + save_context_snapshot(&tool_ctx.memory, ctx); + let _ = output.send(AgentEvent::Status("nothing to compact".into())); + return Ok(()); + } + + let source_memory_ids = recalled_memory_ids(&compacted_messages); + let prompt = build_compaction_messages(&ctx.system, &compacted_messages); + + let _ = output.send(AgentEvent::Status("forming recollection...".into())); + let _ = output.send(AgentEvent::CompactionToken( + COMPACTION_ASSISTANT_PREFILL.to_string(), + )); + let recollection_completion = stream_compaction_recollection(tool_ctx, &prompt, output).await?; + let recollection = complete_compaction_recollection(&recollection_completion); + if recollection.trim().is_empty() { + return Err(anyhow::anyhow!( + "compaction recollection model returned empty output" + )); + } + + let now = unix_timestamp(); + let _ = output.send(AgentEvent::Status("storing recollection...".into())); + let emb = tool_ctx.llm.embed(&recollection).await?; + let recollection_id = tool_ctx + .memory + .store_with_metadata(&crate::mvp::MemoryRecordInput { + memory_id: None, + namespace: "default".to_string(), + layer: crate::mvp::MemoryLayer::L1, + text: recollection.clone(), + event_time: now, + ingest_time: now, + embedding_model: tool_ctx.llm.config.embedder.model.clone(), + embedding_dim: emb.len(), + embedding_version: "runtime".to_string(), + status: crate::mvp::MemoryStatus::Active, + source_ref: Some(format!("system:compaction:{now}")), + tags: vec!["compaction_recollection".to_string()], + pinned: false, + embedding: emb, + })?; + + let drained_now = ctx.drain_oldest(keep, target_tokens); + debug_assert_eq!(drained_now.len(), drained.len()); + let kept_turns = ctx.turn_count(); + if !recollection.is_empty() { + for source_id in source_memory_ids { + let _ = tool_ctx.memory.add_edge(&crate::mvp::MemoryEdgeInput { + from_memory_id: recollection_id, + to_memory_id: source_id, + edge_type: crate::mvp::MemoryEdgeType::DerivedFrom, + metadata: serde_json::json!({ + "source": "recalled_memory_in_compacted_context" + }), + }); + } + } + + if !recollection.is_empty() { + if ctx.turns.first().map(|m| &m.role) == Some(&Role::Assistant) { + let first = &mut ctx.turns[0]; + let new_content = if first.content_str().is_empty() { + recollection.clone() + } else { + format!("{}\n\n{}", recollection, first.content_str()) + }; + first.content = Some(new_content); + } else { + ctx.turns + .insert(0, Message::assistant(recollection.clone())); + } + } + let record = CompactionRecord { + recollection: recollection.clone(), + compacted: compaction_transcript(&compacted_messages), + memory_id: recollection_id, + drained_turns: drained.len(), + kept_turns, + }; + match serde_json::to_string(&record) { + Ok(content) => { + if let Err(e) = tool_ctx.memory.log_turn("compaction", &content, None) { + tracing::warn!(err = %e, "failed to persist compaction record"); + } + } + Err(e) => tracing::warn!(err = %e, "failed to serialize compaction record"), + } + let _ = output.send(AgentEvent::CompactionRecord(record)); + + save_context_snapshot(&tool_ctx.memory, ctx); + + Ok(()) +} + +async fn stream_compaction_recollection( + tool_ctx: &ToolContext, + prompt: &[Message], + output: &broadcast::Sender, +) -> Result { + let compact_registry = tools::memory_tools(); + let mut messages = prompt.to_vec(); + + for _ in 0..20 { + let (tok_tx, mut tok_rx) = mpsc::channel(512); + let llm = tool_ctx.compaction_llm.clone(); + let prompt = messages.clone(); + let definitions = compact_registry.definitions(); + let stream_task = + tokio::spawn(async move { llm.stream(&prompt, &definitions, tok_tx).await }); + let mut iteration_response = String::new(); + let mut iteration_thinking = String::new(); + let mut tool_calls = vec![]; + + while let Some(ev) = tok_rx.recv().await { + match ev { + LlmEvent::Token(token) => { + iteration_response.push_str(&token); + let _ = output.send(AgentEvent::CompactionToken(token)); + } + LlmEvent::ThinkToken(token) => { + iteration_thinking.push_str(&token); + let _ = output.send(AgentEvent::CompactionThinkToken(token)); + } + LlmEvent::ToolCalls(calls) => { + tool_calls = calls; + } + LlmEvent::Usage(_) => {} + LlmEvent::ToolCallStarted => {} + } + } + + match stream_task.await { + Ok(Ok(())) => {} + Ok(Err(e)) => return Err(e), + Err(e) => { + return Err(anyhow::anyhow!( + "compaction recollection stream task panicked/cancelled: {e}" + )) + } + } + + if tool_calls.is_empty() { + return Ok(iteration_response); + } + + let content = (!iteration_response.is_empty()).then_some(iteration_response); + let reasoning = (!iteration_thinking.is_empty()).then_some(iteration_thinking); + messages.push(Message::with_tool_calls( + tool_calls.clone(), + content, + reasoning, + )); + + for call in &tool_calls { + let name = call.function.name.clone(); + let args = call.function.arguments.clone(); + let _ = output.send(AgentEvent::ToolCall { + name: name.clone(), + args, + }); + let result = compact_registry.execute(call, tool_ctx).await; + let _ = output.send(AgentEvent::ToolResult { + name, + content: result.clone(), + }); + messages.push(Message::tool_result(&call.id, result)); + } + } + + Err(anyhow::anyhow!( + "compaction recollection exceeded tool iteration limit" + )) +} + +pub(super) const COMPACTION_ASSISTANT_PREFILL: &str = "from this conversation —"; + +pub(super) fn build_compaction_messages(system: &[Message], drained: &[Message]) -> Vec { + let mut messages = Vec::with_capacity(system.len() + drained.len() + 2); + messages.extend(system.iter().cloned()); + messages.extend(drained.iter().cloned()); + messages.push(Message::user(build_compaction_prompt())); + messages.push(Message::assistant(COMPACTION_ASSISTANT_PREFILL)); + messages +} + +pub(super) fn build_compaction_prompt() -> String { + include_str!("../compaction.md").to_string() +} + +pub(super) fn complete_compaction_recollection(completion: &str) -> String { + let completion = completion.trim_end(); + if completion.trim().is_empty() { + return String::new(); + } + + let trimmed_start = completion.trim_start(); + if trimmed_start.starts_with(COMPACTION_ASSISTANT_PREFILL) { + return trimmed_start.to_string(); + } + + let separator = completion + .chars() + .next() + .is_some_and(char::is_whitespace) + .then_some("") + .unwrap_or(" "); + format!("{COMPACTION_ASSISTANT_PREFILL}{separator}{completion}") +} + +pub(super) fn compaction_source_messages(messages: &[Message]) -> Vec { + messages + .iter() + .filter(|message| !is_reflection_maintenance_message(message)) + .cloned() + .collect() +} + +fn is_reflection_maintenance_message(message: &Message) -> bool { + message.role == Role::Reflection + || message.content.as_deref().is_some_and(|content| { + let trimmed = content.trim_start(); + trimmed.starts_with("") + || trimmed.starts_with("[[automated_reflection]]") + }) +} + +pub(super) fn compaction_transcript(messages: &[Message]) -> String { + messages + .iter() + .filter_map(compaction_message_line) + .collect::>() + .join("\n") +} + +fn compaction_message_line(message: &Message) -> Option { + let mut parts = Vec::new(); + if let Some(reasoning) = message.reasoning_content.as_deref() { + let reasoning = reasoning.trim(); + if !reasoning.is_empty() { + parts.push(harness_block::format_reflection_block( + HarnessBlock::Reasoning, + reasoning, + None, + DEFAULT_FORMAT_STYLE, + )); + } + } + if let Some(content) = message.content.as_deref() { + let content = content.trim(); + if !content.is_empty() { + parts.push(content.to_string()); + } + } + if let Some(calls) = message + .tool_calls + .as_ref() + .filter(|calls| !calls.is_empty()) + { + let calls = calls + .iter() + .map(|call| { + let args = call.function.arguments.trim(); + if args.is_empty() { + call.function.name.clone() + } else { + format!("{}({args})", call.function.name) + } + }) + .collect::>() + .join(" | "); + parts.push(harness_block::format_reflection_block( + HarnessBlock::ToolCalls, + &calls, + None, + DEFAULT_FORMAT_STYLE, + )); + } + if let Some(tool_call_id) = message.tool_call_id.as_deref() { + parts.push(harness_block::format_reflection_block( + HarnessBlock::ToolCallId, + tool_call_id, + None, + DEFAULT_FORMAT_STYLE, + )); + } + + if parts.is_empty() { + None + } else { + Some(format!("{}: {}", message.role, parts.join(" "))) + } +} + +fn recalled_memory_ids(messages: &[Message]) -> Vec { + let mut ids = Vec::new(); + for message in messages { + let Some(content) = message.content.as_deref() else { + continue; + }; + // Bracket style + for part in content.split("[id:").skip(1) { + let Some((raw_id, _)) = part.split_once(']') else { + continue; + }; + let Ok(id) = raw_id.trim().parse::() else { + continue; + }; + if !ids.contains(&id) { + ids.push(id); + } + } + // XML style + for part in content.split("id=\"").skip(1) { + let Some((raw_id, _)) = part.split_once('"') else { + continue; + }; + let Ok(id) = raw_id.trim().parse::() else { + continue; + }; + if !ids.contains(&id) { + ids.push(id); + } + } + } + ids +} diff --git a/klbr-core/src/agent/recall.rs b/klbr-core/src/agent/recall.rs new file mode 100644 index 0000000..a7757bb --- /dev/null +++ b/klbr-core/src/agent/recall.rs @@ -0,0 +1,67 @@ +use std::collections::HashSet; + +use anyhow::Result; +use tokio::sync::broadcast; + +use crate::{ + config::MemoryConfig, + evidence::EvidencePacket, + memory::MemoryStore, + models::LlmClient, + pipeline::{BenchQuery, ContextBudget, MemoryPipeline}, + AgentEvent, +}; + +pub(super) async fn recall_evidence_packets( + config: &MemoryConfig, + memory: &MemoryStore, + llm: &LlmClient, + query: &str, + reference_time: Option, + already_recalled: &HashSet, + profile: &str, + output: &broadcast::Sender, +) -> Result> { + let pipeline = + MemoryPipeline::new(memory.clone(), llm.clone(), config.clone()).with_profile(profile); + let budget = ContextBudget { + max_tokens: config.top_k.max(1).saturating_mul(900), + top_k: config.candidate_k.max(config.top_k).max(1), + graph_depth: 1, + }; + let retrieval = pipeline + .retrieve_evidence( + BenchQuery { + query_id: "runtime-passive-recall".to_string(), + text: query.to_string(), + reference_time, + op_hint: None, + }, + budget, + ) + .await?; + let selected = retrieval + .packets + .into_iter() + .filter(|packet| { + packet + .refs + .iter() + .any(|ref_id| !already_recalled.contains(ref_id)) + }) + .take(config.top_k) + .collect::>(); + + if !retrieval.candidates.is_empty() || !selected.is_empty() || !retrieval.route.archival_allowed + { + let _ = output.send(AgentEvent::Status(format!( + "memory packets: {} candidates, {} selected, route={}, archival={}", + retrieval.candidates.len(), + selected.len(), + retrieval.route.reason, + retrieval.route.archival_allowed + ))); + } + + Ok(selected) +} diff --git a/klbr-core/src/agent/reflection.rs b/klbr-core/src/agent/reflection.rs new file mode 100644 index 0000000..716b5f4 --- /dev/null +++ b/klbr-core/src/agent/reflection.rs @@ -0,0 +1,213 @@ +use anyhow::Result; +use tokio::sync::{broadcast, mpsc}; + +use crate::{ + context::Context, + harness_block::{self, HarnessBlock, DEFAULT_FORMAT_STYLE}, + models::{LlmEvent, Message, Role}, + tools::{self, ToolContext}, + AgentEvent, +}; + +/// reflection loop: let the agent review and curate its memories +pub(super) async fn reflect( + tool_ctx: &ToolContext, + ctx: &Context, + output: &broadcast::Sender, +) -> Result<()> { + let reflect_registry = tools::memory_tools(); + let mut msgs = build_reflection_messages(tool_ctx, ctx); + let mut reflection_log = Vec::new(); + + let _ = output.send(AgentEvent::ReflectStarted); + + for _ in 0..100 { + let (tok_tx, mut tok_rx) = mpsc::channel(512); + let llm2 = tool_ctx.llm.clone(); + let msgs_snap = msgs.clone(); + let defs_snap = reflect_registry.definitions(); + let stream_task = + tokio::spawn(async move { llm2.stream(&msgs_snap, &defs_snap, tok_tx).await }); + + let mut tool_calls = vec![]; + let mut stream_failed = false; + let mut iteration_thinking = String::new(); + let mut iteration_response = String::new(); + + while let Some(ev) = tok_rx.recv().await { + match ev { + LlmEvent::ToolCalls(calls) => tool_calls = calls, + LlmEvent::Token(t) => { + iteration_response.push_str(&t); + let _ = output.send(AgentEvent::Token(t)); + } + LlmEvent::ThinkToken(t) => { + iteration_thinking.push_str(&t); + let _ = output.send(AgentEvent::ThinkToken(t)); + } + _ => {} + } + } + match stream_task.await { + Ok(Ok(())) => {} + Ok(Err(e)) => { + let msg = format!("reflection llm stream failed: {e}"); + tracing::error!(%msg); + let _ = output.send(AgentEvent::Error(msg)); + stream_failed = true; + } + Err(e) => { + let msg = format!("reflection llm stream task panicked/cancelled: {e}"); + tracing::error!(%msg); + let _ = output.send(AgentEvent::Error(msg)); + stream_failed = true; + } + } + + if (tool_calls.is_empty() && iteration_response.is_empty()) || stream_failed { + break; + } + + let reasoning = if iteration_thinking.is_empty() { + None + } else { + reflection_log.push(harness_block::format_reflection_block( + HarnessBlock::Reasoning, + &iteration_thinking, + None, + DEFAULT_FORMAT_STYLE, + )); + Some(iteration_thinking) + }; + let content = if iteration_response.is_empty() { + None + } else { + reflection_log.push(harness_block::format_reflection_block( + HarnessBlock::Response, + &iteration_response, + None, + DEFAULT_FORMAT_STYLE, + )); + Some(iteration_response) + }; + + if !tool_calls.is_empty() { + msgs.push(Message::with_tool_calls( + tool_calls.clone(), + content, + reasoning, + )); + for call in &tool_calls { + let name = call.function.name.clone(); + let args = call.function.arguments.clone(); + reflection_log.push(harness_block::format_reflection_block( + HarnessBlock::ToolCall, + &args, + Some(&name), + DEFAULT_FORMAT_STYLE, + )); + let _ = output.send(AgentEvent::ToolCall { + name: name.clone(), + args: args.clone(), + }); + let result = reflect_registry.execute(call, tool_ctx).await; + reflection_log.push(harness_block::format_reflection_block( + HarnessBlock::ToolResult, + &result, + Some(&name), + DEFAULT_FORMAT_STYLE, + )); + let _ = output.send(AgentEvent::ToolResult { + name: name.clone(), + content: result.clone(), + }); + msgs.push(Message::tool_result(&call.id, &result)); + } + } else { + msgs.push(Message { + role: Role::Assistant, + content, + reasoning_content: reasoning, + tool_calls: None, + tool_call_id: None, + }); + break; + } + } + + if !reflection_log.is_empty() { + if let Err(e) = tool_ctx + .memory + .log_turn("reflection", &reflection_log.join("\n\n"), None) + { + tracing::warn!(err = %e, "failed to persist reflection record"); + } + } + let _ = output.send(AgentEvent::ReflectDone); + Ok(()) +} + +pub(super) fn build_reflection_messages(tool_ctx: &ToolContext, ctx: &Context) -> Vec { + let pinned = tool_ctx.memory.pinned_memory_entries().unwrap_or_default(); + let unpinned = tool_ctx.memory.recent_unpinned(20).unwrap_or_default(); + + let pinned_text = if pinned.is_empty() { + "(none)".to_string() + } else { + pinned + .iter() + .map(|memory| { + harness_block::format_memory( + HarnessBlock::PinnedMemory, + memory.id, + &memory.tags, + &[], + false, + &memory.content, + DEFAULT_FORMAT_STYLE, + ) + }) + .collect::>() + .join("\n") + }; + + let unpinned_text = if unpinned.is_empty() { + "(none)".to_string() + } else { + unpinned + .iter() + .map(|(id, s, tags)| { + harness_block::format_memory( + HarnessBlock::UnpinnedMemory, + *id, + tags, + &[], + false, + s, + DEFAULT_FORMAT_STYLE, + ) + }) + .collect::>() + .join("\n") + }; + let reflection_prompt = + build_reflection_prompt(&pinned_text, &unpinned_text, chrono::Utc::now()); + + let mut messages = ctx.as_messages_with_refs(&tool_ctx.memory); + messages.push(Message::user(reflection_prompt)); + messages +} + +pub(super) fn build_reflection_prompt( + pinned_text: &str, + unpinned_text: &str, + datetime: chrono::DateTime, +) -> String { + format!( + include_str!("../reflection.md"), + pinned_text = pinned_text, + unpinned_text = unpinned_text, + date = datetime.date_naive(), + time = datetime.time(), + ) +} diff --git a/klbr-core/src/agent/tests.rs b/klbr-core/src/agent/tests.rs new file mode 100644 index 0000000..343bdea --- /dev/null +++ b/klbr-core/src/agent/tests.rs @@ -0,0 +1,364 @@ +use std::collections::HashSet; + +use super::{ + compaction::{ + build_compaction_messages, build_compaction_prompt, compaction_source_messages, + compaction_transcript, complete_compaction_recollection, COMPACTION_ASSISTANT_PREFILL, + }, + needs_assistant_response, + recall::recall_evidence_packets, + reflection::build_reflection_prompt, +}; +use crate::{models::Role, tools::format_duration_human}; +use anyhow::Result; +use chrono::{TimeZone, Utc}; +use tempfile::NamedTempFile; +use tokio::sync::broadcast; + +use crate::{ + config::MemoryConfig, + context::Context, + memory::{MarkdownNoteInput, MemoryLane, MemoryStore}, + models::{LlmClient, Message, ModelsConfig, ToolCall, ToolCallFunction}, + pipeline::{BenchSession, BenchTurn, MemoryPipeline}, +}; + +fn tool_call(id: &str, name: &str) -> ToolCall { + ToolCall { + id: id.to_string(), + kind: "function".to_string(), + function: ToolCallFunction { + name: name.to_string(), + arguments: "{}".to_string(), + }, + } +} + +#[tokio::test] +async fn runtime_packet_recall_uses_pipeline_packet_window() -> Result<()> { + let tmp = NamedTempFile::new()?; + let store = MemoryStore::open(tmp.path().to_str().unwrap(), 4)?; + let llm = LlmClient::new(ModelsConfig::default()); + let mut config = MemoryConfig::default(); + config.top_k = 2; + config.candidate_k = 5; + + let pipeline = MemoryPipeline::new(store.clone(), llm.clone(), config.clone()) + .with_profile("raw-turns/fts-only"); + pipeline + .observe_session(BenchSession { + session_id: "play-session".to_string(), + timestamp: Some(1), + turns: vec![ + BenchTurn { + role: "user".to_string(), + content: "I recently went to a play at the local community theater." + .to_string(), + timestamp: Some(1), + }, + BenchTurn { + role: "assistant".to_string(), + content: "The play was The Glass Menagerie.".to_string(), + timestamp: Some(2), + }, + ], + }) + .await?; + + let (tx, _) = broadcast::channel(8); + let packets = recall_evidence_packets( + &config, + &store, + &llm, + "what play did I recently go to at the community theater?", + Some(3), + &HashSet::new(), + "raw-turns/fts-only", + &tx, + ) + .await?; + + assert!(packets + .iter() + .flat_map(|packet| &packet.bodies) + .any(|body| body.contains("The Glass Menagerie"))); + + let mut ctx = Context::new("soul", &[]); + let recalled_index = ctx.inject_evidence_packets(&packets, 900).unwrap(); + let rendered = ctx.turns[recalled_index] + .content + .as_deref() + .unwrap_or_default(); + assert!(rendered.contains("")); + assert!(rendered.contains("kind=\"turn_window\"")); + assert!(rendered.contains("The Glass Menagerie")); + let recalled_refs = ctx.passively_recalled_refs(); + assert!(packets + .iter() + .flat_map(|packet| &packet.refs) + .any(|ref_id| recalled_refs.contains(ref_id))); + + Ok(()) +} + +#[tokio::test] +async fn runtime_packet_recall_resolves_explicit_markdown_refs() -> Result<()> { + let tmp = NamedTempFile::new()?; + let store = MemoryStore::open(tmp.path().to_str().unwrap(), 4)?; + let llm = LlmClient::new(ModelsConfig::default()); + let mut config = MemoryConfig::default(); + config.top_k = 2; + config.candidate_k = 5; + + let note = store.upsert_markdown_note(&MarkdownNoteInput { + note_ref: Some("n_runtime_note".to_string()), + lane: MemoryLane::Semantic, + kind: "semantic_note".to_string(), + title: "runtime note".to_string(), + path: None, + body: "the explicit markdown answer is blue quartz".to_string(), + sources: vec![], + follow: None, + entities: vec![], + status: "active".to_string(), + frontmatter: serde_json::json!({}), + })?; + let chunk_alias = note + .chunk_refs + .first() + .expect("markdown note should create a chunk alias"); + + let (tx, _) = broadcast::channel(8); + let packets = recall_evidence_packets( + &config, + &store, + &llm, + &format!("what did [{}] say?", chunk_alias), + Some(3), + &HashSet::new(), + "raw-turns/fts-only", + &tx, + ) + .await?; + + assert!(packets + .iter() + .any(|packet| packet.signals.sources.contains(&"exact".to_string()))); + assert!(packets + .iter() + .flat_map(|packet| &packet.bodies) + .any(|body| body.contains("blue quartz"))); + + Ok(()) +} + +#[tokio::test] +async fn runtime_packet_recall_does_not_use_discourse_cue_filter() -> Result<()> { + let tmp = NamedTempFile::new()?; + let store = MemoryStore::open(tmp.path().to_str().unwrap(), 4)?; + let llm = LlmClient::new(ModelsConfig::default()); + let mut config = MemoryConfig::default(); + config.top_k = 2; + config.candidate_k = 5; + + let pipeline = MemoryPipeline::new(store.clone(), llm.clone(), config.clone()) + .with_profile("raw-turns/fts-only"); + pipeline + .observe_session(BenchSession { + session_id: "old-session".to_string(), + timestamp: Some(1), + turns: vec![BenchTurn { + role: "user".to_string(), + content: "yes do that was an old archival phrase".to_string(), + timestamp: Some(1), + }], + }) + .await?; + + let (tx, _) = broadcast::channel(8); + let packets = recall_evidence_packets( + &config, + &store, + &llm, + "yes do that", + Some(3), + &HashSet::new(), + "raw-turns/fts-only", + &tx, + ) + .await?; + + assert!(packets + .iter() + .flat_map(|packet| &packet.bodies) + .any(|body| body.contains("old archival phrase"))); + Ok(()) +} + +#[test] +fn reflection_prompt_is_automated_instruction() { + let prompt = build_reflection_prompt( + "(none)", + "(none)", + Utc.with_ymd_and_hms(2026, 4, 28, 12, 0, 0).unwrap(), + ); + + assert!(prompt.starts_with("")); + assert!(prompt.contains("## goal")); + assert!(prompt.contains("not a human user message")); + assert!(prompt.contains("current state / carry-forward")); + assert!(prompt.contains("errors and fixes")); + assert!(prompt.contains("contradictions and corrections")); + assert!(!prompt.contains("real soul instructions")); +} + +#[test] +fn compaction_messages_reuse_current_system_prefix() { + let system = vec![Message::system("runtime soul with pinned block")]; + let drained = vec![ + Message::user("old request"), + Message::assistant("old answer"), + ]; + + let messages = build_compaction_messages(&system, &drained); + + assert_eq!(messages.len(), 5); + assert_eq!(messages[0].role, "system"); + assert_eq!( + messages[0].content.as_deref(), + Some("runtime soul with pinned block") + ); + assert_eq!(messages[1].content.as_deref(), Some("old request")); + assert_eq!(messages[2].content.as_deref(), Some("old answer")); + assert_eq!(messages[3].role, "user"); + let maintenance_prompt = messages[3].content.as_deref().unwrap_or_default(); + assert!(maintenance_prompt.contains("a recollection of the context")); + assert!(maintenance_prompt.contains("carry forward what mattered")); + assert_eq!(messages[4].role, "assistant"); + assert_eq!( + messages[4].content.as_deref(), + Some(COMPACTION_ASSISTANT_PREFILL) + ); +} + +#[test] +fn compaction_prompt_is_reflection_instruction() { + let prompt = build_compaction_prompt(); + + assert!(prompt.contains("carry forward what mattered from this conversation")); + assert!(prompt.contains("dissent too")); + assert!(!prompt.contains("your dissent too")); + assert!(prompt.contains("no tools for this")); + assert!(!prompt.contains("{pinned_memories}")); +} + +#[test] +fn compaction_recollection_completes_assistant_prefill() { + assert_eq!( + complete_compaction_recollection("what mattered was the exact phrasing"), + "from this conversation — what mattered was the exact phrasing" + ); + assert_eq!( + complete_compaction_recollection(" carried forward a live task"), + "from this conversation — carried forward a live task" + ); + assert_eq!( + complete_compaction_recollection("from this conversation — already whole"), + "from this conversation — already whole" + ); + assert_eq!(complete_compaction_recollection(" "), ""); +} + +#[test] +fn compaction_transcript_keeps_tool_call_context() { + let messages = vec![ + Message::with_tool_calls( + vec![ToolCall { + id: "call_1".to_string(), + kind: "function".to_string(), + function: ToolCallFunction { + name: "read_file".to_string(), + arguments: "{\"path\":\"klbr-core/src/agent.rs\"}".to_string(), + }, + }], + Some("checking the file".to_string()), + Some("need context".to_string()), + ), + Message::tool_result("call_1", "file contents"), + ]; + + let transcript = compaction_transcript(&messages); + + assert!( + transcript.contains("assistant: \nneed context\n checking the file") + ); + assert!(transcript + .contains("read_file({\"path\":\"klbr-core/src/agent.rs\"})")); + assert!(transcript.contains("tool: file contents call_1")); +} + +#[test] +fn compaction_source_messages_excludes_reflection_prompts() { + let messages = vec![ + Message::user("real old turn"), + Message::assistant("[[automated_reflection]]\n\nreflect over this"), + Message::assistant("\n\nreflect over this"), + Message { + role: Role::Reflection, + content: Some("logged reflection".to_string()), + reasoning_content: None, + tool_calls: None, + tool_call_id: None, + }, + ]; + + let source = compaction_source_messages(&messages); + + assert_eq!(source.len(), 1); + assert_eq!(source[0].content.as_deref(), Some("real old turn")); +} + +#[test] +fn wait_timeout_result_does_not_need_another_assistant_response() { + let mut ctx = Context::new("soul", &[]); + ctx.push_assistant_tool_calls(vec![tool_call("call_1", "wait_and_continue")], None, None); + ctx.push_tool_result( + "call_1", + "waited 1min; no incoming event arrived; current time: 2026-06-14 23:07:36 +03:00; suspended", + ); + + assert!(!needs_assistant_response(&ctx)); +} + +#[test] +fn wait_tool_error_still_needs_assistant_response() { + let mut ctx = Context::new("soul", &[]); + ctx.push_assistant_tool_calls(vec![tool_call("call_1", "wait_and_continue")], None, None); + ctx.push_tool_result( + "call_1", + "error: duration is required (either 'timeout_ms' or 'seconds' must be specified)", + ); + + assert!(needs_assistant_response(&ctx)); +} + +#[test] +fn non_wait_tool_result_still_needs_assistant_response() { + let mut ctx = Context::new("soul", &[]); + ctx.push_assistant_tool_calls(vec![tool_call("call_1", "remember")], None, None); + ctx.push_tool_result("call_1", "stored"); + + assert!(needs_assistant_response(&ctx)); +} + +#[test] +fn test_format_duration_human() { + use std::time::Duration; + assert_eq!(format_duration_human(Duration::from_millis(500)), "500ms"); + assert_eq!(format_duration_human(Duration::from_secs(5)), "5sec"); + assert_eq!(format_duration_human(Duration::from_secs(65)), "1min 5sec"); + assert_eq!( + format_duration_human(Duration::from_secs(3665)), + "1hour 1min 5sec" + ); +} -- 2.51.2