oxi_agent/tools/
search_cache.rs

1use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
2use crate::tools::typed::TypedTool;
3use async_trait::async_trait;
4use parking_lot::Mutex;
5use schemars::JsonSchema;
6use serde::Deserialize;
7use serde_json::{Value, json};
8use std::collections::HashMap;
9use std::sync::Arc;
10use tokio::sync::oneshot;
11
12// Re-export oxibrowser's SearchResult for all search tools.
13pub use oxibrowser::SearchResult;
14
15// ── Search cache ──────────────────────────────────────────────────
16
17/// In-memory cache for search results, keyed by search ID.
18#[derive(Debug)]
19pub struct SearchCache {
20    /// Map of search_id → (query, results).
21    entries: Mutex<HashMap<String, CachedSearch>>,
22    /// Maximum number of cached searches. Oldest arbitrary entry is evicted when full.
23    max_entries: usize,
24}
25
26#[derive(Debug, Clone)]
27struct CachedSearch {
28    query: String,
29    results: Vec<SearchResult>,
30}
31
32impl Default for SearchCache {
33    fn default() -> Self {
34        Self::new()
35    }
36}
37
38impl SearchCache {
39    /// Create a new empty cache with default capacity (64 entries).
40    pub fn new() -> Self {
41        Self::with_capacity(64)
42    }
43
44    /// Create a new cache with the given maximum capacity.
45    pub fn with_capacity(max_entries: usize) -> Self {
46        Self {
47            entries: Mutex::new(HashMap::new()),
48            max_entries,
49        }
50    }
51
52    /// Insert search results and return the generated search ID.
53    pub fn insert(&self, query: &str, results: Vec<SearchResult>) -> String {
54        let id = generate_search_id();
55        let cached = CachedSearch {
56            query: query.to_string(),
57            results,
58        };
59
60        let mut entries = self.entries.lock();
61
62        // Evict oldest entries if at capacity
63        while entries.len() >= self.max_entries {
64            // Simple eviction: remove a random entry
65            if let Some(key) = entries.keys().next().cloned() {
66                entries.remove(&key);
67            }
68        }
69
70        entries.insert(id.clone(), cached);
71        id
72    }
73
74    /// Retrieve cached search results by ID.
75    pub fn get(&self, search_id: &str) -> Option<(String, Vec<SearchResult>)> {
76        let entries = self.entries.lock();
77        entries
78            .get(search_id)
79            .map(|c| (c.query.clone(), c.results.clone()))
80    }
81}
82
83/// Generate a short unique search ID.
84fn generate_search_id() -> String {
85    let ts = std::time::SystemTime::now()
86        .duration_since(std::time::UNIX_EPOCH)
87        .unwrap_or_default()
88        .as_millis();
89    let rand_part: u32 = rng::random();
90    format!("{:x}{:06x}", ts, rand_part & 0xFFFFFF)
91}
92
93// ── GetSearchResultsTool ──────────────────────────────────────────
94
95/// Tool for retrieving cached search results by ID.
96pub struct GetSearchResultsTool {
97    cache: Arc<SearchCache>,
98}
99
100#[derive(Deserialize, JsonSchema)]
101#[allow(missing_docs)]
102pub struct GetSearchResultsArgs {
103    #[serde(rename = "searchId")]
104    search_id: String,
105}
106
107impl GetSearchResultsTool {
108    /// Create a new GetSearchResultsTool with the given cache.
109    pub fn new(cache: Arc<SearchCache>) -> Self {
110        Self { cache }
111    }
112}
113
114#[async_trait]
115impl AgentTool for GetSearchResultsTool {
116    fn name(&self) -> &str {
117        "get_search_results"
118    }
119
120    fn label(&self) -> &str {
121        "Get Search Results"
122    }
123
124    fn description(&self) -> &str {
125        "Retrieve previous search results by ID. Use this to look up results from a prior web_search call."
126    }
127
128    fn parameters_schema(&self) -> Value {
129        json!({
130            "type": "object",
131            "properties": {
132                "searchId": {
133                    "type": "string",
134                    "description": "The search ID returned by a previous web_search call"
135                }
136            },
137            "required": ["searchId"]
138        })
139    }
140
141    async fn execute(
142        &self,
143        _tool_call_id: &str,
144        params: Value,
145        _signal: Option<oneshot::Receiver<()>>,
146        _ctx: &ToolContext,
147    ) -> Result<AgentToolResult, ToolError> {
148        let args: GetSearchResultsArgs =
149            serde_json::from_value(params).map_err(|e| format!("invalid params: {e}"))?;
150        self.execute_typed(_tool_call_id, args, _signal, _ctx).await
151    }
152}
153
154#[async_trait]
155impl TypedTool for GetSearchResultsTool {
156    type Args = GetSearchResultsArgs;
157
158    async fn execute_typed(
159        &self,
160        _tool_call_id: &str,
161        args: Self::Args,
162        _signal: Option<oneshot::Receiver<()>>,
163        _ctx: &ToolContext,
164    ) -> Result<AgentToolResult, ToolError> {
165        let (query, results) = self
166            .cache
167            .get(&args.search_id)
168            .ok_or_else(|| format!("Search not found for ID: {}", args.search_id))?;
169        let mut output = format!("Cached results for: \"{}\"\n\n", query);
170        for (i, result) in results.iter().enumerate() {
171            output.push_str(&format!(
172                "{}. **{}**\n   {}\n   {}\n\n",
173                i + 1,
174                result.title,
175                result.url,
176                result.snippet
177            ));
178        }
179        let results_json: Vec<Value> = results.iter().map(|r| {
180            json!({"title": r.title, "url": r.url, "snippet": r.snippet, "source": r.source})
181        }).collect();
182        Ok(AgentToolResult::success(output).with_metadata(
183            json!({ "results": results_json, "query": query, "searchId": args.search_id }),
184        ))
185    }
186}
187#[allow(dead_code)]
188mod rng {
189    #[allow(unused_imports)]
190    use std::cell::Cell;
191    #[allow(unused_imports)]
192    use std::time::SystemTime;
193
194    thread_local! {
195        static SEED: Cell<u64> = const { Cell::new(0) };
196    }
197
198    pub fn random() -> u32 {
199        SEED.with(|s| {
200            let mut x = if s.get() == 0 {
201                let ns = SystemTime::now()
202                    .duration_since(SystemTime::UNIX_EPOCH)
203                    .unwrap_or_default()
204                    .as_nanos() as u64;
205                ns ^ (thread_id() as u64)
206            } else {
207                s.get()
208            };
209            x ^= x << 13;
210            x ^= x >> 7;
211            x ^= x << 17;
212            s.set(x);
213            (x & 0xFFFFFFFF) as u32
214        })
215    }
216
217    fn thread_id() -> usize {
218        thread_local! { static ANCHOR: () = const {  }; }
219        ANCHOR.with(|_| &ANCHOR as *const _ as usize)
220    }
221}
222
223// ── Tests ─────────────────────────────────────────────────────────
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    fn make_result(title: &str, url: &str, snippet: &str) -> SearchResult {
230        SearchResult {
231            title: title.to_string(),
232            url: url.to_string(),
233            snippet: snippet.to_string(),
234            source: "test".to_string(),
235            extra: None,
236        }
237    }
238
239    #[test]
240    fn test_cache_insert_and_get() {
241        let cache = SearchCache::new();
242        let results = vec![make_result("Test", "https://example.com", "Test snippet")];
243
244        let id = cache.insert("test query", results.clone());
245        let (query, retrieved) = cache.get(&id).unwrap();
246        assert_eq!(query, "test query");
247        assert_eq!(retrieved.len(), 1);
248        assert_eq!(retrieved[0].title, "Test");
249    }
250
251    #[test]
252    fn test_cache_miss() {
253        let cache = SearchCache::new();
254        assert!(cache.get("nonexistent").is_none());
255    }
256
257    #[test]
258    fn test_cache_eviction() {
259        let cache = SearchCache::with_capacity(3);
260
261        let id1 = cache.insert("q1", vec![]);
262        let id2 = cache.insert("q2", vec![]);
263        let id3 = cache.insert("q3", vec![]);
264        let _id4 = cache.insert("q4", vec![]);
265
266        // At least one of the first 3 should have been evicted
267        let found = [&id1, &id2, &id3]
268            .iter()
269            .filter(|id| cache.get(id).is_some())
270            .count();
271        assert!(found < 3);
272        assert!(cache.get(&_id4).is_some());
273    }
274
275    #[test]
276    fn test_generate_search_id_unique() {
277        let id1 = generate_search_id();
278        let id2 = generate_search_id();
279        assert_ne!(id1, id2);
280    }
281
282    #[tokio::test]
283    async fn test_get_search_results_tool() {
284        let cache = Arc::new(SearchCache::new());
285        let results = vec![make_result("Rust", "https://rust-lang.org", "A language")];
286        let id = cache.insert("rust lang", results);
287
288        let tool = GetSearchResultsTool::new(cache);
289        let result = tool
290            .execute(
291                "test",
292                json!({ "searchId": id }),
293                None,
294                &ToolContext::default(),
295            )
296            .await
297            .unwrap();
298
299        assert!(result.success);
300        assert!(result.output.contains("Rust"));
301        assert!(result.output.contains("rust-lang.org"));
302    }
303
304    #[tokio::test]
305    async fn test_get_search_results_not_found() {
306        let cache = Arc::new(SearchCache::new());
307        let tool = GetSearchResultsTool::new(cache);
308        let result = tool
309            .execute(
310                "test",
311                json!({ "searchId": "bad-id" }),
312                None,
313                &ToolContext::default(),
314            )
315            .await;
316
317        assert!(result.is_err());
318    }
319
320    #[test]
321    fn test_get_search_results_schema() {
322        let cache = Arc::new(SearchCache::new());
323        let tool = GetSearchResultsTool::new(cache);
324        let schema = tool.parameters_schema();
325        assert_eq!(schema["type"], "object");
326        assert!(schema["properties"]["searchId"].is_object());
327    }
328}
oxi_agent/tools/search_cache.rs

oxi_agent/tools/
search_cache.rs
