heartbit_core/agent/
context_recall.rs

1//! Per-run "restore-on-demand" store: indexes every tool output by
2//! `tool_call_id` so a pruned/compacted result can be restored exactly
3//! (`get`) or found semantically (`recall`). Reuses `InMemoryStore`'s
4//! BM25(+vector)->RRF retrieval.
5
6use chrono::Utc;
7
8use crate::auth::TenantScope;
9use crate::memory::in_memory::InMemoryStore;
10use crate::memory::{Confidentiality, Memory, MemoryEntry, MemoryQuery, MemoryType};
11
12/// Default cap on stored tool outputs (bounded so the store can't leak).
13const DEFAULT_MAX_ENTRIES: usize = 256;
14
15/// Max characters of head-content returned per recall hit (generous, so the
16/// snippet often answers the question without a follow-up fetch).
17// 280 ≈ one paragraph — generous enough to answer most queries without a full fetch.
18pub const SNIPPET_CHARS: usize = 280;
19
20/// One ranked match from `recall`: the ref to fetch, which tool produced it,
21/// and a head snippet.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct RecallHit {
24    /// The `tool_call_id` used to index this entry; pass to `get` for full restore.
25    pub r#ref: String,
26    /// Name of the tool that produced the output (first tag on the entry).
27    pub tool_name: String,
28    /// Leading `SNIPPET_CHARS` characters of the stored content.
29    pub snippet: String,
30}
31
32/// A per-run store of tool outputs, keyed by `tool_call_id`.
33pub struct ContextRecallStore {
34    inner: InMemoryStore,
35    scope: TenantScope,
36}
37
38impl Default for ContextRecallStore {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44impl ContextRecallStore {
45    /// Create a bounded store with the default capacity.
46    pub fn new() -> Self {
47        Self {
48            inner: InMemoryStore::new().with_max_entries(DEFAULT_MAX_ENTRIES),
49            scope: TenantScope::default(),
50        }
51    }
52
53    /// Index a tool result so it can later be restored by `id` or found by query.
54    /// Best-effort: a store error is swallowed (recall is a convenience, not a
55    /// correctness dependency).
56    pub async fn index(&self, tool_call_id: &str, tool_name: &str, content: &str) {
57        let entry = make_entry(tool_call_id, tool_name, content);
58        let _ = self.inner.store(&self.scope, entry).await;
59    }
60
61    /// Exact restore of a stored tool output by its `tool_call_id`.
62    pub async fn get(&self, tool_call_id: &str) -> Option<String> {
63        self.inner.get(tool_call_id).map(|e| e.content)
64    }
65
66    /// Semantically find stored tool outputs by `query` (BM25, or BM25+vector
67    /// when an embedder is configured). Returns ranked refs + head snippets;
68    /// the caller restores the full body via `get`/`fetch_full_output`.
69    pub async fn recall(&self, query: &str, limit: usize) -> Vec<RecallHit> {
70        let q = MemoryQuery {
71            text: Some(query.to_string()),
72            limit,
73            reinforce: false,
74            ..Default::default()
75        };
76        let entries = self.inner.recall(&self.scope, q).await.unwrap_or_default();
77        entries
78            .into_iter()
79            .map(|e| RecallHit {
80                r#ref: e.id,
81                tool_name: e.tags.first().cloned().unwrap_or_default(),
82                snippet: e.content.chars().take(SNIPPET_CHARS).collect(),
83            })
84            .collect()
85    }
86}
87
88/// Build a `MemoryEntry` for a tool output (defaults for the unused fields).
89fn make_entry(id: &str, tool_name: &str, content: &str) -> MemoryEntry {
90    let now = Utc::now();
91    MemoryEntry {
92        id: id.to_string(),
93        agent: "context_recall".into(),
94        content: content.to_string(),
95        category: "tool_output".into(),
96        tags: vec![tool_name.to_string()],
97        created_at: now,
98        last_accessed: now,
99        access_count: 0,
100        importance: 5,
101        memory_type: MemoryType::Episodic,
102        keywords: Vec::new(),
103        summary: None,
104        strength: 1.0,
105        related_ids: Vec::new(),
106        source_ids: Vec::new(),
107        embedding: None,
108        confidentiality: Confidentiality::Public,
109        author_user_id: None,
110        author_tenant_id: None,
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[tokio::test]
119    async fn index_then_get_roundtrips_exact_content() {
120        let store = ContextRecallStore::new();
121        store
122            .index("tc_1", "bash", "the full untruncated output")
123            .await;
124        assert_eq!(
125            store.get("tc_1").await.as_deref(),
126            Some("the full untruncated output")
127        );
128        assert_eq!(store.get("nope").await, None);
129    }
130
131    #[tokio::test]
132    async fn recall_ranks_a_matching_output_above_noise_and_caps_snippet() {
133        let store = ContextRecallStore::new();
134        store
135            .index(
136                "tc_match",
137                "bash",
138                "cargo test failed: assertion error in parser module",
139            )
140            .await;
141        store
142            .index(
143                "tc_noise",
144                "read",
145                "the quick brown fox jumps over the lazy dog",
146            )
147            .await;
148
149        let hits = store.recall("test failure parser", 5).await;
150        assert!(!hits.is_empty(), "expected at least one hit");
151        assert_eq!(
152            hits[0].r#ref, "tc_match",
153            "the matching output must rank first"
154        );
155        assert_eq!(hits[0].tool_name, "bash");
156        assert!(hits[0].snippet.chars().count() <= SNIPPET_CHARS);
157    }
158
159    #[tokio::test]
160    async fn recall_on_empty_store_is_empty() {
161        let store = ContextRecallStore::new();
162        assert!(store.recall("anything", 5).await.is_empty());
163    }
164
165    #[tokio::test]
166    async fn recall_snippet_is_truncated_to_the_cap() {
167        let store = ContextRecallStore::new();
168        let long = "alpha ".repeat(200); // ~1200 chars, all matching the query term
169        store.index("tc_long", "bash", &long).await;
170        let hits = store.recall("alpha", 1).await;
171        assert_eq!(hits.len(), 1);
172        assert_eq!(
173            hits[0].snippet.chars().count(),
174            SNIPPET_CHARS,
175            "a long output's snippet must be capped at exactly SNIPPET_CHARS"
176        );
177    }
178
179    #[tokio::test]
180    async fn pruned_then_restored_by_marker_ref_roundtrips() {
181        use crate::agent::pruner::{SessionPruneConfig, prune_old_tool_results};
182        use crate::llm::types::{ContentBlock, Message, Role};
183
184        // 1. A big tool output is produced and indexed under its id.
185        let store = ContextRecallStore::new();
186        let big = "RESTORE_ME ".repeat(500); // well above the default 200-byte prune cap
187        store.index("tc_abc", "bash", &big).await;
188
189        // 2. The pruner truncates it in a request view; the marker names the ref.
190        //    Build a message list where the ToolResult(id "tc_abc") is OLD enough
191        //    to be pruned (followed by enough recent messages to push it out of the
192        //    kept tail). With default keep_recent_n=2, last 4 messages are kept;
193        //    index 1 (ToolResult) falls outside the recent window and gets pruned.
194        let messages = vec![
195            Message::user("task"),
196            Message {
197                role: Role::User,
198                content: vec![ContentBlock::ToolResult {
199                    tool_use_id: "tc_abc".into(),
200                    content: big.clone(),
201                    is_error: false,
202                }],
203            },
204            Message::user("r1"),
205            Message::assistant("r2"),
206            Message::user("r3"),
207            Message::assistant("r4"),
208        ];
209        let (pruned, stats) = prune_old_tool_results(&messages, &SessionPruneConfig::default());
210        assert!(stats.did_prune());
211        let marker: String = pruned
212            .iter()
213            .flat_map(|m| m.content.iter())
214            .filter_map(|b| match b {
215                ContentBlock::ToolResult { content, .. } => Some(content.clone()),
216                _ => None,
217            })
218            .collect();
219        assert!(
220            marker.contains("tc_abc"),
221            "the pruned marker must name the ref"
222        );
223        assert!(
224            marker.len() < big.len(),
225            "the result was actually shortened"
226        );
227
228        // 3. Restore the exact original content by the ref from the store.
229        assert_eq!(
230            store.get("tc_abc").await.as_deref(),
231            Some(big.as_str()),
232            "the full content is restorable by the marker's ref"
233        );
234    }
235}
heartbit_core/agent/context_recall.rs

heartbit_core/agent/
context_recall.rs
