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

rig_agent/test_utils/
tools.rs
