rig_agent/test_utils/
tools.rs

1//! Tool helpers for deterministic tests.
2
3use std::sync::{Arc, Mutex};
4
5use serde::{Deserialize, Serialize};
6use serde_json::json;
7
8use rig_core::{
9    OneOrMany,
10    message::{ImageMediaType, ToolResultContent},
11    vector_store::{VectorSearchRequest, VectorStoreError, VectorStoreIndex, request::Filter},
12    wasm_compat::WasmCompatSend,
13};
14
15use crate::tool::{Tool, ToolContext, ToolErrorKind, ToolExecutionError, ToolOutput, ToolSet};
16
17/// Shared error type for mock tools.
18#[derive(Debug, thiserror::Error)]
19#[error("Mock tool error")]
20pub struct MockToolError;
21
22/// Arguments for arithmetic mock tools.
23#[derive(Deserialize)]
24pub struct MockOperationArgs {
25    x: i32,
26    y: i32,
27}
28
29/// A mock tool that adds `x` and `y`.
30#[derive(Deserialize, Serialize)]
31pub struct MockAddTool;
32
33impl Tool for MockAddTool {
34    const NAME: &'static str = "add";
35    type Error = MockToolError;
36    type Args = MockOperationArgs;
37    type Output = i32;
38
39    fn description(&self) -> String {
40        "Add x and y together".to_string()
41    }
42
43    fn parameters(&self) -> serde_json::Value {
44        json!({
45            "type": "object",
46            "properties": {
47                "x": {
48                    "type": "number",
49                    "description": "The first number to add"
50                },
51                "y": {
52                    "type": "number",
53                    "description": "The second number to add"
54                }
55            },
56            "required": ["x", "y"],
57        })
58    }
59
60    async fn call(
61        &self,
62        _context: &mut crate::tool::ToolContext,
63        args: Self::Args,
64    ) -> Result<Self::Output, Self::Error> {
65        Ok(args.x + args.y)
66    }
67}
68
69/// A caller-injected context value, like a session id or auth token carried in
70/// a [`ToolContext`](crate::tool::ToolContext).
71#[derive(Clone)]
72pub struct SessionId(pub String);
73
74/// A mock tool that records whatever it observed in its per-call
75/// [`ToolContext`], so tests can assert the context reached tool execution.
76///
77/// The single `call` method records `session:<id>` (or `no-session`).
78#[derive(Clone, Default)]
79pub struct MockContextProbeTool {
80    /// One entry per call, in call order — lets tests assert across multiple
81    /// tool-call rounds, not just the most recent.
82    seen: Arc<Mutex<Vec<String>>>,
83}
84
85impl MockContextProbeTool {
86    /// What the tool observed on its most recent call, if it has been called.
87    pub fn observed(&self) -> Option<String> {
88        self.seen
89            .lock()
90            .unwrap_or_else(|poisoned| poisoned.into_inner())
91            .last()
92            .cloned()
93    }
94
95    /// Everything the tool observed, one entry per call in call order.
96    pub fn observations(&self) -> Vec<String> {
97        self.seen
98            .lock()
99            .unwrap_or_else(|poisoned| poisoned.into_inner())
100            .clone()
101    }
102}
103
104impl Tool for MockContextProbeTool {
105    const NAME: &'static str = "context_probe";
106    type Error = rig::tool::ToolExecutionError;
107    type Args = serde_json::Value;
108    type Output = String;
109
110    fn description(&self) -> String {
111        "Records the SessionId observed in its call context".to_string()
112    }
113
114    fn parameters(&self) -> serde_json::Value {
115        json!({"type": "object", "properties": {}})
116    }
117
118    async fn call(
119        &self,
120        context: &mut ToolContext,
121        _args: Self::Args,
122    ) -> Result<Self::Output, ToolExecutionError> {
123        let observed = match context.get::<SessionId>() {
124            Some(session) => format!("session:{}", session.0),
125            None => "no-session".to_string(),
126        };
127        self.seen
128            .lock()
129            .unwrap_or_else(|poisoned| poisoned.into_inner())
130            .push(observed.clone());
131        Ok(observed)
132    }
133}
134
135/// A mock tool that subtracts `y` from `x`.
136#[derive(Deserialize, Serialize)]
137pub struct MockSubtractTool;
138
139impl Tool for MockSubtractTool {
140    const NAME: &'static str = "subtract";
141    type Error = MockToolError;
142    type Args = MockOperationArgs;
143    type Output = i32;
144
145    fn description(&self) -> String {
146        "Subtract y from x".to_string()
147    }
148
149    fn parameters(&self) -> serde_json::Value {
150        json!({
151            "type": "object",
152            "properties": {
153                "x": {
154                    "type": "number",
155                    "description": "The number to subtract from"
156                },
157                "y": {
158                    "type": "number",
159                    "description": "The number to subtract"
160                }
161            },
162            "required": ["x", "y"],
163        })
164    }
165
166    async fn call(
167        &self,
168        _context: &mut crate::tool::ToolContext,
169        args: Self::Args,
170    ) -> Result<Self::Output, Self::Error> {
171        Ok(args.x - args.y)
172    }
173}
174
175/// Create a [`ToolSet`] containing [`MockAddTool`] and [`MockSubtractTool`].
176pub fn mock_math_toolset() -> ToolSet {
177    let mut toolset = ToolSet::default();
178    toolset.add_tool(MockAddTool);
179    toolset.add_tool(MockSubtractTool);
180    toolset
181}
182
183/// A mock tool that returns a multiline string.
184#[derive(Deserialize, Serialize)]
185pub struct MockStringOutputTool;
186
187impl Tool for MockStringOutputTool {
188    const NAME: &'static str = "string_output";
189    type Error = MockToolError;
190    type Args = serde_json::Value;
191    type Output = String;
192
193    fn description(&self) -> String {
194        "Returns a multiline string".to_string()
195    }
196
197    fn parameters(&self) -> serde_json::Value {
198        json!({
199            "type": "object",
200            "properties": {}
201        })
202    }
203
204    async fn call(
205        &self,
206        _context: &mut crate::tool::ToolContext,
207        _args: Self::Args,
208    ) -> Result<Self::Output, Self::Error> {
209        Ok("Hello\nWorld".to_string())
210    }
211}
212
213/// A mock tool that returns explicit image content.
214#[derive(Deserialize, Serialize)]
215pub struct MockImageOutputTool;
216
217impl Tool for MockImageOutputTool {
218    const NAME: &'static str = "image_output";
219    type Error = MockToolError;
220    type Args = serde_json::Value;
221    type Output = ToolOutput;
222
223    fn description(&self) -> String {
224        "Returns an image".to_string()
225    }
226
227    fn parameters(&self) -> serde_json::Value {
228        json!({
229            "type": "object",
230            "properties": {}
231        })
232    }
233
234    async fn call(
235        &self,
236        _context: &mut crate::tool::ToolContext,
237        _args: Self::Args,
238    ) -> Result<Self::Output, Self::Error> {
239        Ok(ToolOutput::content(OneOrMany::one(
240            ToolResultContent::image_base64("base64data==", Some(ImageMediaType::PNG), None),
241        )))
242    }
243}
244
245/// A mock tool named `generate_test_image` that returns a 1x1 red PNG image payload.
246#[derive(Debug, Deserialize, Serialize)]
247pub struct MockImageGeneratorTool;
248
249impl Tool for MockImageGeneratorTool {
250    const NAME: &'static str = "generate_test_image";
251    type Error = MockToolError;
252    type Args = serde_json::Value;
253    type Output = ToolOutput;
254
255    fn description(&self) -> String {
256        "Generates a small test image (a 1x1 red pixel). Call this tool when asked to generate or show an image.".to_string()
257    }
258
259    fn parameters(&self) -> serde_json::Value {
260        json!({
261            "type": "object",
262            "properties": {},
263            "required": []
264        })
265    }
266
267    async fn call(
268        &self,
269        _context: &mut crate::tool::ToolContext,
270        _args: Self::Args,
271    ) -> Result<Self::Output, Self::Error> {
272        Ok(ToolOutput::content(OneOrMany::one(
273            ToolResultContent::image_base64(
274                "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==",
275                Some(ImageMediaType::PNG),
276                None,
277            ),
278        )))
279    }
280}
281
282/// A mock tool that returns a JSON object.
283#[derive(Deserialize, Serialize)]
284pub struct MockObjectOutputTool;
285
286impl Tool for MockObjectOutputTool {
287    const NAME: &'static str = "object_output";
288    type Error = MockToolError;
289    type Args = serde_json::Value;
290    type Output = serde_json::Value;
291
292    fn description(&self) -> String {
293        "Returns an object".to_string()
294    }
295
296    fn parameters(&self) -> serde_json::Value {
297        json!({
298            "type": "object",
299            "properties": {}
300        })
301    }
302
303    async fn call(
304        &self,
305        _context: &mut crate::tool::ToolContext,
306        _args: Self::Args,
307    ) -> Result<Self::Output, Self::Error> {
308        Ok(json!({
309            "status": "ok",
310            "count": 42
311        }))
312    }
313}
314
315/// A mock tool named `example_tool` that returns `"Example answer"`.
316pub struct MockExampleTool;
317
318impl Tool for MockExampleTool {
319    const NAME: &'static str = "example_tool";
320    type Error = MockToolError;
321    type Args = ();
322    type Output = String;
323
324    fn description(&self) -> String {
325        "A tool that returns some example text.".to_string()
326    }
327
328    fn parameters(&self) -> serde_json::Value {
329        json!({
330            "type": "object",
331            "properties": {},
332            "required": []
333        })
334    }
335
336    async fn call(
337        &self,
338        _context: &mut crate::tool::ToolContext,
339        _input: Self::Args,
340    ) -> Result<Self::Output, Self::Error> {
341        Ok("Example answer".to_string())
342    }
343}
344
345/// A mock tool that waits at a barrier before returning `"done"`.
346#[derive(Clone)]
347pub struct MockBarrierTool {
348    /// Barrier waited on during each tool call.
349    pub barrier: Arc<tokio::sync::Barrier>,
350}
351
352impl MockBarrierTool {
353    /// Create a barrier-backed tool.
354    pub fn new(barrier: Arc<tokio::sync::Barrier>) -> Self {
355        Self { barrier }
356    }
357}
358
359impl Tool for MockBarrierTool {
360    const NAME: &'static str = "barrier_tool";
361    type Error = MockToolError;
362    type Args = serde_json::Value;
363    type Output = String;
364
365    fn description(&self) -> String {
366        "Waits at a barrier to test concurrency".to_string()
367    }
368
369    fn parameters(&self) -> serde_json::Value {
370        json!({"type": "object", "properties": {}})
371    }
372
373    async fn call(
374        &self,
375        _context: &mut crate::tool::ToolContext,
376        _args: Self::Args,
377    ) -> Result<Self::Output, Self::Error> {
378        self.barrier.wait().await;
379        Ok("done".to_string())
380    }
381}
382
383/// A mock tool that notifies when started and waits for an explicit finish signal.
384#[derive(Clone)]
385pub struct MockControlledTool {
386    /// Notified when a tool call starts.
387    pub started: Arc<tokio::sync::Notify>,
388    /// Waited on before a tool call finishes.
389    pub allow_finish: Arc<tokio::sync::Notify>,
390}
391
392impl MockControlledTool {
393    /// Create a controlled tool from notification primitives.
394    pub fn new(started: Arc<tokio::sync::Notify>, allow_finish: Arc<tokio::sync::Notify>) -> Self {
395        Self {
396            started,
397            allow_finish,
398        }
399    }
400}
401
402impl Tool for MockControlledTool {
403    const NAME: &'static str = "controlled";
404    type Error = MockToolError;
405    type Args = serde_json::Value;
406    type Output = i32;
407
408    fn description(&self) -> String {
409        "Test tool".to_string()
410    }
411
412    fn parameters(&self) -> serde_json::Value {
413        json!({"type": "object", "properties": {}})
414    }
415
416    async fn call(
417        &self,
418        _context: &mut crate::tool::ToolContext,
419        _args: Self::Args,
420    ) -> Result<Self::Output, Self::Error> {
421        self.started.notify_one();
422        self.allow_finish.notified().await;
423        Ok(42)
424    }
425}
426
427/// A vector index that returns a predefined list of tool IDs from `top_n_ids`.
428pub struct MockToolIndex {
429    tool_ids: Vec<String>,
430}
431
432impl MockToolIndex {
433    /// Create a tool index that returns the given IDs in order.
434    pub fn new(tool_ids: impl IntoIterator<Item = impl Into<String>>) -> Self {
435        Self {
436            tool_ids: tool_ids.into_iter().map(Into::into).collect(),
437        }
438    }
439}
440
441impl VectorStoreIndex for MockToolIndex {
442    type Filter = Filter<serde_json::Value>;
443
444    async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
445        &self,
446        _req: VectorSearchRequest,
447    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
448        Ok(vec![])
449    }
450
451    async fn top_n_ids(
452        &self,
453        _req: VectorSearchRequest,
454    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
455        Ok(self
456            .tool_ids
457            .iter()
458            .enumerate()
459            .map(|(i, id)| (1.0 - (i as f64 * 0.1), id.clone()))
460            .collect())
461    }
462}
463
464/// A vector index that waits at a barrier before returning one tool ID.
465pub struct BarrierMockToolIndex {
466    barrier: Arc<tokio::sync::Barrier>,
467    tool_id: String,
468}
469
470impl BarrierMockToolIndex {
471    /// Create a barrier-backed tool index.
472    pub fn new(barrier: Arc<tokio::sync::Barrier>, tool_id: impl Into<String>) -> Self {
473        Self {
474            barrier,
475            tool_id: tool_id.into(),
476        }
477    }
478}
479
480impl VectorStoreIndex for BarrierMockToolIndex {
481    type Filter = Filter<serde_json::Value>;
482
483    async fn top_n<T: for<'a> Deserialize<'a> + WasmCompatSend>(
484        &self,
485        _req: VectorSearchRequest,
486    ) -> Result<Vec<(f64, String, T)>, VectorStoreError> {
487        Ok(vec![])
488    }
489
490    async fn top_n_ids(
491        &self,
492        _req: VectorSearchRequest,
493    ) -> Result<Vec<(f64, String)>, VectorStoreError> {
494        self.barrier.wait().await;
495        Ok(vec![(1.0, self.tool_id.clone())])
496    }
497}
498
499/// Error type for [`MockFailingTool`], carrying a fixed message.
500#[derive(Debug, thiserror::Error)]
501#[error("mock tool call failed")]
502pub struct MockFailure;
503
504/// A tool that always fails with a configured [`ToolErrorKind`]. Used to exercise structured
505/// tool-failure surfacing (timeout, not-found, rate-limited, …) without a live
506/// provider. Registered under the name `flaky_tool`.
507#[derive(Clone)]
508pub struct MockFailingTool {
509    kind: ToolErrorKind,
510}
511
512impl MockFailingTool {
513    /// A tool that fails with the given classification every call.
514    pub fn new(kind: ToolErrorKind) -> Self {
515        Self { kind }
516    }
517}
518
519impl Tool for MockFailingTool {
520    const NAME: &'static str = "flaky_tool";
521    type Error = MockFailure;
522    type Args = serde_json::Value;
523    type Output = String;
524
525    fn description(&self) -> String {
526        "A tool that always fails".to_string()
527    }
528
529    fn parameters(&self) -> serde_json::Value {
530        json!({ "type": "object", "properties": {} })
531    }
532
533    async fn call(
534        &self,
535        _context: &mut ToolContext,
536        _args: Self::Args,
537    ) -> Result<Self::Output, Self::Error> {
538        Err(MockFailure)
539    }
540
541    fn map_error(&self, error: Self::Error) -> ToolExecutionError {
542        let error = ToolExecutionError::new(self.kind, error.to_string()).with_source(error);
543        match self.kind {
544            ToolErrorKind::NotFound => error.with_http_status(404),
545            ToolErrorKind::RateLimited => error.with_http_status(429),
546            _ => error,
547        }
548    }
549}
550
551/// A tool failure with separate operator and model-visible feedback.
552#[derive(Clone)]
553pub struct MockHandledFailureTool;
554
555impl Tool for MockHandledFailureTool {
556    const NAME: &'static str = "lookup";
557    type Error = MockToolError;
558    type Args = serde_json::Value;
559    type Output = String;
560
561    fn description(&self) -> String {
562        "Looks up a record".to_string()
563    }
564
565    fn parameters(&self) -> serde_json::Value {
566        json!({ "type": "object", "properties": {} })
567    }
568
569    async fn call(
570        &self,
571        _context: &mut ToolContext,
572        _args: Self::Args,
573    ) -> Result<Self::Output, Self::Error> {
574        Err(MockToolError)
575    }
576
577    fn map_error(&self, error: Self::Error) -> ToolExecutionError {
578        ToolExecutionError::not_found("record id 42 is missing")
579            .with_http_status(404)
580            .with_model_feedback("no record found for id 42; try a different id")
581            .with_source(error)
582    }
583}
584
585/// A tool that refuses execution, distinct from a framework policy skip.
586#[derive(Clone)]
587pub struct MockDeniedTool;
588
589impl Tool for MockDeniedTool {
590    const NAME: &'static str = "guarded";
591    type Error = MockToolError;
592    type Args = serde_json::Value;
593    type Output = String;
594
595    fn description(&self) -> String {
596        "A tool with an internal authorization check".to_string()
597    }
598
599    fn parameters(&self) -> serde_json::Value {
600        json!({ "type": "object", "properties": {} })
601    }
602
603    async fn call(
604        &self,
605        _context: &mut ToolContext,
606        _args: Self::Args,
607    ) -> Result<Self::Output, Self::Error> {
608        Err(MockToolError)
609    }
610
611    fn map_error(&self, error: Self::Error) -> ToolExecutionError {
612        ToolExecutionError::refused("operator authorization policy rejected the request")
613            .with_model_feedback("access to this resource is not permitted")
614            .with_source(error)
615    }
616}
617
618/// Cloneable metadata a [`MockMetadataTool`] attaches to its result, used to
619/// verify that result metadata reaches hooks without being sent to the model.
620#[derive(Clone, Debug, PartialEq, Eq)]
621pub struct MockRequestId(pub String);
622
623/// A tool whose success carries a [`MockRequestId`] in its result metadata.
624/// Registered under the name `with_meta`.
625#[derive(Clone)]
626pub struct MockMetadataTool;
627
628impl Tool for MockMetadataTool {
629    const NAME: &'static str = "with_meta";
630    type Error = MockToolError;
631    type Args = serde_json::Value;
632    type Output = String;
633
634    fn description(&self) -> String {
635        "Succeeds and attaches request metadata".to_string()
636    }
637
638    fn parameters(&self) -> serde_json::Value {
639        json!({ "type": "object", "properties": {} })
640    }
641
642    async fn call(
643        &self,
644        context: &mut ToolContext,
645        _args: Self::Args,
646    ) -> Result<Self::Output, Self::Error> {
647        context.insert_result(MockRequestId("req-7".to_string()));
648        Ok("done".to_string())
649    }
650}
rig_agent/test_utils/tools.rs

rig_agent/test_utils/
tools.rs
