Skip to main content

harness_loop/
recall_layer.rs

1//! Cross-session recall wiring for [`crate::AgentLoop`].
2//!
3//! - [`SessionSearchTool`] — LLM-callable search over the recall store, three
4//!   shapes (discovery / scroll / browse). Owner is read from
5//!   `World.profile.extra["recall_owner"]` so it can only see the caller's own
6//!   sessions.
7//! - [`RecallGuide`] — optional. At session start, searches the store with the
8//!   task description and injects the top snippets. Off unless `.auto_inject()`.
9
10use async_trait::async_trait;
11use harness_core::{
12    Block, Context, Execution, Guide, GuideError, GuideId, GuideScope, RecallStore, Tool,
13    ToolError, ToolResult, ToolRisk, ToolSchema, World,
14};
15use serde_json::{Value, json};
16use std::sync::{Arc, OnceLock};
17
18/// Read the recall owner from the world profile (fallback "default").
19pub fn recall_owner(world: &World) -> String {
20    world
21        .profile
22        .extra
23        .get("recall_owner")
24        .and_then(|v| v.as_str())
25        .unwrap_or("default")
26        .to_string()
27}
28
29// ───── session_search tool ────────────────────────────────────────────────
30
31pub struct SessionSearchTool {
32    store: Arc<dyn RecallStore>,
33    schema: ToolSchema,
34}
35
36impl SessionSearchTool {
37    pub fn new(store: Arc<dyn RecallStore>) -> Self {
38        Self {
39            store,
40            schema: ToolSchema {
41                name: "session_search".into(),
42                description: "Search your own past sessions, or scroll inside one. \
43                    Three shapes: (1) pass `query` to find relevant past sessions \
44                    (returns snippet + surrounding messages); (2) pass `session_id` + \
45                    `around` to scroll messages near a point in a session; (3) pass \
46                    nothing to list your most recent sessions."
47                    .into(),
48                input: json!({
49                    "type": "object",
50                    "properties": {
51                        "query": {"type": "string", "description": "Search text. Shape 1 (discovery)."},
52                        "session_id": {"type": "string", "description": "Scroll within this session. Shape 2."},
53                        "around": {"type": "integer", "description": "Anchor message id for scroll. Shape 2."},
54                        "window": {"type": "integer", "default": 5, "description": "± messages around the anchor."},
55                        "limit": {"type": "integer", "default": 3, "minimum": 1, "maximum": 20}
56                    }
57                }),
58            },
59        }
60    }
61}
62
63#[async_trait]
64impl Tool for SessionSearchTool {
65    fn name(&self) -> &str {
66        &self.schema.name
67    }
68    fn schema(&self) -> &ToolSchema {
69        &self.schema
70    }
71    fn risk(&self) -> ToolRisk {
72        ToolRisk::ReadOnly
73    }
74    async fn invoke(&self, args: Value, world: &mut World) -> Result<ToolResult, ToolError> {
75        let owner = recall_owner(world);
76        let limit = args
77            .get("limit")
78            .and_then(|v| v.as_u64())
79            .unwrap_or(3)
80            .min(20) as usize;
81
82        let result = if let Some(q) = args
83            .get("query")
84            .and_then(|v| v.as_str())
85            .filter(|s| !s.is_empty())
86        {
87            match self.store.search(&owner, q, limit).await {
88                Ok(hits) => {
89                    json!({"mode": "discover", "query": q, "count": hits.len(), "results": hits})
90                }
91                Err(e) => return Ok(err_result(e)),
92            }
93        } else if let Some(sid) = args.get("session_id").and_then(|v| v.as_str()) {
94            let around = args.get("around").and_then(|v| v.as_i64()).unwrap_or(0);
95            let window = args.get("window").and_then(|v| v.as_u64()).unwrap_or(5) as usize;
96            match self.store.scroll(&owner, sid, around, window).await {
97                Ok(msgs) => json!({"mode": "scroll", "session_id": sid, "messages": msgs}),
98                Err(e) => return Ok(err_result(e)),
99            }
100        } else {
101            match self.store.recent(&owner, limit).await {
102                Ok(sessions) => json!({"mode": "browse", "sessions": sessions}),
103                Err(e) => return Ok(err_result(e)),
104            }
105        };
106        Ok(ToolResult {
107            ok: true,
108            content: result,
109            trace: None,
110        })
111    }
112}
113
114fn err_result(e: harness_core::RecallError) -> ToolResult {
115    ToolResult {
116        ok: false,
117        content: json!({"error": e.to_string()}),
118        trace: None,
119    }
120}
121
122// ───── RecallGuide (opt-in auto-inject) ───────────────────────────────────
123
124const RECALL_MARKER: &str = "[recall]\n";
125
126pub struct RecallGuide {
127    store: Arc<dyn RecallStore>,
128    top_k: usize,
129}
130
131static RECALL_GUIDE_ID: OnceLock<GuideId> = OnceLock::new();
132static RECALL_GUIDE_SCOPE: OnceLock<GuideScope> = OnceLock::new();
133
134impl RecallGuide {
135    pub fn new(store: Arc<dyn RecallStore>) -> Self {
136        Self { store, top_k: 3 }
137    }
138    pub fn with_top_k(mut self, k: usize) -> Self {
139        self.top_k = k;
140        self
141    }
142}
143
144#[async_trait]
145impl Guide for RecallGuide {
146    fn id(&self) -> &GuideId {
147        RECALL_GUIDE_ID.get_or_init(|| "recall".to_string())
148    }
149    fn kind(&self) -> Execution {
150        Execution::Inferential
151    }
152    fn scope(&self) -> &GuideScope {
153        RECALL_GUIDE_SCOPE.get_or_init(|| GuideScope::Always)
154    }
155    async fn apply(&self, ctx: &mut Context, world: &World) -> Result<(), GuideError> {
156        let owner = recall_owner(world);
157        let query = ctx.task.description.clone();
158        let hits = self
159            .store
160            .search(&owner, &query, self.top_k)
161            .await
162            .unwrap_or_default();
163        if hits.is_empty() {
164            return Ok(());
165        }
166        let mut text = String::from(RECALL_MARKER);
167        text.push_str("Possibly-relevant context from your past sessions:\n");
168        for h in &hits {
169            text.push_str(&format!("- ({}) {}\n", h.session.session_id, h.snippet));
170        }
171        ctx.guides.push(Block::Text(text));
172        Ok(())
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179    use harness_context::{FileRecall, default_world};
180    use harness_core::{RecallMessage, SessionMeta};
181
182    fn tmp_root() -> std::path::PathBuf {
183        use std::sync::atomic::{AtomicU64, Ordering};
184        static N: AtomicU64 = AtomicU64::new(0);
185        let n = N.fetch_add(1, Ordering::SeqCst);
186        let nanos = std::time::SystemTime::now()
187            .duration_since(std::time::UNIX_EPOCH)
188            .unwrap()
189            .as_nanos();
190        std::env::temp_dir().join(format!(
191            "harness-recall-tool-{}-{nanos}-{n}",
192            std::process::id()
193        ))
194    }
195
196    #[tokio::test]
197    async fn tool_discovery_scoped_to_owner() {
198        let root = tmp_root();
199        let store: Arc<dyn RecallStore> = Arc::new(FileRecall::open(&root).unwrap());
200        store
201            .ensure_session("alice", "s1", &SessionMeta::new("s1", 1))
202            .await
203            .unwrap();
204        store
205            .append(
206                "alice",
207                "s1",
208                &RecallMessage::new("user", "deploy the payment service", 1),
209            )
210            .await
211            .unwrap();
212
213        let tool = SessionSearchTool::new(store.clone());
214        let mut world = default_world(".");
215        world
216            .profile
217            .extra
218            .insert("recall_owner".into(), serde_json::json!("alice"));
219        let out = tool
220            .invoke(serde_json::json!({"query": "payment deploy"}), &mut world)
221            .await
222            .unwrap();
223        assert!(out.ok);
224        assert_eq!(out.content["count"], 1);
225
226        let mut bob = default_world(".");
227        bob.profile
228            .extra
229            .insert("recall_owner".into(), serde_json::json!("bob"));
230        let out2 = tool
231            .invoke(serde_json::json!({"query": "payment deploy"}), &mut bob)
232            .await
233            .unwrap();
234        assert_eq!(out2.content["count"], 0);
235
236        let _ = std::fs::remove_dir_all(&root);
237    }
238}