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

rig_core/test_utils/
tools.rs
