rig_agent/tool/
extensions.rs

1//! Typed context passed through tool execution.
2
3use std::any::{Any, TypeId, type_name};
4use std::collections::HashMap;
5use std::hash::{BuildHasherDefault, Hasher};
6
7use crate::tool::ToolExecutionError;
8use rig_core::wasm_compat::{WasmCompatSend, WasmCompatSync};
9
10type AnyMap = HashMap<TypeId, Box<dyn AnyClone>, BuildHasherDefault<IdHasher>>;
11
12#[derive(Default)]
13struct IdHasher(u64);
14
15impl Hasher for IdHasher {
16    fn write_u64(&mut self, id: u64) {
17        self.0 = id;
18    }
19
20    fn write(&mut self, bytes: &[u8]) {
21        for &byte in bytes {
22            self.0 = self.0.rotate_left(8) ^ u64::from(byte);
23        }
24    }
25
26    fn finish(&self) -> u64 {
27        self.0
28    }
29}
30
31trait AnyClone: Any + WasmCompatSend + WasmCompatSync {
32    fn clone_box(&self) -> Box<dyn AnyClone>;
33    fn as_any(&self) -> &dyn Any;
34    fn as_any_mut(&mut self) -> &mut dyn Any;
35    fn into_any(self: Box<Self>) -> Box<dyn Any>;
36    fn type_name(&self) -> &'static str;
37}
38
39impl<T> AnyClone for T
40where
41    T: Clone + WasmCompatSend + WasmCompatSync + 'static,
42{
43    fn clone_box(&self) -> Box<dyn AnyClone> {
44        Box::new(self.clone())
45    }
46
47    fn as_any(&self) -> &dyn Any {
48        self
49    }
50
51    fn as_any_mut(&mut self) -> &mut dyn Any {
52        self
53    }
54
55    fn into_any(self: Box<Self>) -> Box<dyn Any> {
56        self
57    }
58
59    fn type_name(&self) -> &'static str {
60        type_name::<T>()
61    }
62}
63
64impl Clone for Box<dyn AnyClone> {
65    fn clone(&self) -> Self {
66        (**self).clone_box()
67    }
68}
69
70/// Internal type map shared by tool contexts and hook scratchpads.
71#[derive(Default, Clone)]
72pub(crate) struct TypeMap {
73    map: Option<Box<AnyMap>>,
74}
75
76impl TypeMap {
77    pub(crate) const EMPTY: Self = Self { map: None };
78
79    pub(crate) fn insert<T>(&mut self, value: T) -> Option<T>
80    where
81        T: Clone + WasmCompatSend + WasmCompatSync + 'static,
82    {
83        self.map
84            .get_or_insert_with(Default::default)
85            .insert(TypeId::of::<T>(), Box::new(value))
86            .and_then(|previous| previous.into_any().downcast::<T>().ok())
87            .map(|value| *value)
88    }
89
90    pub(crate) fn get<T>(&self) -> Option<&T>
91    where
92        T: WasmCompatSend + WasmCompatSync + 'static,
93    {
94        self.map
95            .as_ref()
96            .and_then(|map| map.get(&TypeId::of::<T>()))
97            .and_then(|value| (**value).as_any().downcast_ref::<T>())
98    }
99
100    pub(crate) fn get_mut<T>(&mut self) -> Option<&mut T>
101    where
102        T: WasmCompatSend + WasmCompatSync + 'static,
103    {
104        self.map
105            .as_mut()
106            .and_then(|map| map.get_mut(&TypeId::of::<T>()))
107            .and_then(|value| (**value).as_any_mut().downcast_mut::<T>())
108    }
109
110    pub(crate) fn remove<T>(&mut self) -> Option<T>
111    where
112        T: WasmCompatSend + WasmCompatSync + 'static,
113    {
114        self.map
115            .as_mut()
116            .and_then(|map| map.remove(&TypeId::of::<T>()))
117            .and_then(|value| value.into_any().downcast::<T>().ok())
118            .map(|value| *value)
119    }
120
121    pub(crate) fn contains<T>(&self) -> bool
122    where
123        T: WasmCompatSend + WasmCompatSync + 'static,
124    {
125        self.map
126            .as_ref()
127            .is_some_and(|map| map.contains_key(&TypeId::of::<T>()))
128    }
129
130    pub(crate) fn len(&self) -> usize {
131        self.map.as_ref().map_or(0, |map| map.len())
132    }
133
134    fn type_names(&self) -> Vec<&'static str> {
135        self.map
136            .as_ref()
137            .map(|map| map.values().map(|value| (**value).type_name()).collect())
138            .unwrap_or_default()
139    }
140}
141
142/// Context passed to every tool execution.
143///
144/// Callers insert typed inbound values with [`insert`](Self::insert). Tools read
145/// those values with [`get`](Self::get) or [`require`](Self::require), and attach
146/// host-only result metadata with [`insert_result`](Self::insert_result). Result
147/// hooks inspect that metadata through [`result`](Self::result). Neither inbound
148/// values nor result metadata are sent to the model.
149///
150/// Registry, server, and agent dispatch clone inbound values once per call.
151/// Map-level mutations (inserting, replacing, or removing typed slots) affect
152/// only that execution. Value-level isolation follows each value's [`Clone`]
153/// semantics, so intentionally shared values such as `Arc<Mutex<_>>` continue
154/// to share their referent across dispatches. The dispatch surface returns
155/// result metadata without replacing the caller's inbound slots.
156#[derive(Default, Clone)]
157pub struct ToolContext {
158    inbound: TypeMap,
159    result: TypeMap,
160}
161
162impl ToolContext {
163    /// Create an empty context.
164    pub const fn new() -> Self {
165        Self {
166            inbound: TypeMap::EMPTY,
167            result: TypeMap::EMPTY,
168        }
169    }
170
171    /// Insert an inbound typed value, returning the displaced value if present.
172    pub fn insert<T>(&mut self, value: T) -> Option<T>
173    where
174        T: Clone + WasmCompatSend + WasmCompatSync + 'static,
175    {
176        self.inbound.insert(value)
177    }
178
179    /// Read an inbound typed value.
180    pub fn get<T>(&self) -> Option<&T>
181    where
182        T: WasmCompatSend + WasmCompatSync + 'static,
183    {
184        self.inbound.get::<T>()
185    }
186
187    /// Require an inbound typed value.
188    pub fn require<T>(&self) -> Result<&T, MissingToolContext>
189    where
190        T: WasmCompatSend + WasmCompatSync + 'static,
191    {
192        self.get::<T>().ok_or(MissingToolContext(type_name::<T>()))
193    }
194
195    /// Mutably access an inbound typed value.
196    pub fn get_mut<T>(&mut self) -> Option<&mut T>
197    where
198        T: WasmCompatSend + WasmCompatSync + 'static,
199    {
200        self.inbound.get_mut::<T>()
201    }
202
203    /// Remove an inbound typed value.
204    pub fn remove<T>(&mut self) -> Option<T>
205    where
206        T: WasmCompatSend + WasmCompatSync + 'static,
207    {
208        self.inbound.remove::<T>()
209    }
210
211    /// Attach host-only metadata to this execution's result.
212    pub fn insert_result<T>(&mut self, value: T) -> Option<T>
213    where
214        T: Clone + WasmCompatSend + WasmCompatSync + 'static,
215    {
216        self.result.insert(value)
217    }
218
219    /// Read host-only result metadata.
220    pub fn result<T>(&self) -> Option<&T>
221    where
222        T: WasmCompatSend + WasmCompatSync + 'static,
223    {
224        self.result.get::<T>()
225    }
226
227    /// Require host-only result metadata.
228    pub fn require_result<T>(&self) -> Result<&T, MissingToolContext>
229    where
230        T: WasmCompatSend + WasmCompatSync + 'static,
231    {
232        self.result::<T>()
233            .ok_or(MissingToolContext(type_name::<T>()))
234    }
235
236    /// Whether this context contains the inbound type `T`.
237    pub fn contains<T>(&self) -> bool
238    where
239        T: WasmCompatSend + WasmCompatSync + 'static,
240    {
241        self.inbound.contains::<T>()
242    }
243
244    /// Build a fresh execution context with the same inbound values and no
245    /// result metadata.
246    ///
247    /// Dispatch always runs against this snapshot. Mutating its typed slots does
248    /// not change the run-wide or caller-owned map. Values with shared/interior
249    /// state remain shared according to their [`Clone`] implementation.
250    pub(crate) fn for_dispatch(&self) -> Self {
251        Self {
252            inbound: self.inbound.clone(),
253            result: TypeMap::EMPTY,
254        }
255    }
256
257    /// Publish metadata produced by one dispatch while preserving the caller's
258    /// inbound values.
259    pub(crate) fn accept_dispatch_result(&mut self, dispatched: Self) {
260        self.result = dispatched.result;
261    }
262
263    /// Clear metadata from the previous dispatch before starting another one.
264    pub(crate) fn clear_dispatch_result(&mut self) {
265        self.result = TypeMap::EMPTY;
266    }
267}
268
269impl std::fmt::Debug for ToolContext {
270    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
271        f.debug_struct("ToolContext")
272            .field("inbound_entries", &self.inbound.len())
273            .field("inbound_types", &self.inbound.type_names())
274            .field("result_entries", &self.result.len())
275            .field("result_types", &self.result.type_names())
276            .finish()
277    }
278}
279
280/// A required typed value was missing from a [`ToolContext`].
281#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
282#[error("required tool context value of type `{0}` was not found")]
283pub struct MissingToolContext(pub &'static str);
284
285impl From<MissingToolContext> for ToolExecutionError {
286    fn from(error: MissingToolContext) -> Self {
287        ToolExecutionError::other(error.to_string()).with_source(error)
288    }
289}
290
291#[cfg(not(target_family = "wasm"))]
292const _: fn() = || {
293    fn assert_send_sync<T: Send + Sync>() {}
294    assert_send_sync::<ToolContext>();
295};
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300
301    #[test]
302    fn context_separates_inbound_and_result_values() {
303        let mut context = ToolContext::new();
304        context.insert(42_u32);
305        context.insert_result("request-1".to_string());
306        assert_eq!(context.get::<u32>(), Some(&42));
307        assert_eq!(
308            context.result::<String>().map(String::as_str),
309            Some("request-1")
310        );
311
312        let next = context.for_dispatch();
313        assert_eq!(next.get::<u32>(), Some(&42));
314        assert!(next.result::<String>().is_none());
315    }
316
317    #[test]
318    fn missing_context_converts_into_a_tool_execution_error() {
319        fn require_value(context: &ToolContext) -> Result<u32, ToolExecutionError> {
320            Ok(*context.require::<u32>()?)
321        }
322
323        let error = require_value(&ToolContext::new()).unwrap_err();
324        assert!(error.is::<MissingToolContext>());
325        assert_eq!(
326            error.model_feedback(),
327            Some("required tool context value of type `u32` was not found")
328        );
329    }
330}
331
332#[cfg(test)]
333mod migrated_tests {
334    use super::*;
335
336    #[test]
337    fn insert_and_get_returns_value() {
338        let mut c = ToolContext::new();
339        assert_eq!(c.insert(42u32), None);
340        assert_eq!(c.get::<u32>(), Some(&42));
341    }
342    #[test]
343    fn get_missing_type_returns_none() {
344        assert_eq!(ToolContext::new().get::<u32>(), None);
345    }
346    #[test]
347    fn insert_overwrites_and_returns_previous() {
348        let mut c = ToolContext::new();
349        c.insert(1u32);
350        assert_eq!(c.insert(2u32), Some(1));
351        assert_eq!(c.get::<u32>(), Some(&2));
352    }
353    #[test]
354    fn different_types_are_independent() {
355        let mut c = ToolContext::new();
356        c.insert(42u32);
357        c.insert("hello".to_string());
358        assert_eq!(c.get::<u32>(), Some(&42));
359        assert_eq!(c.get::<String>().map(String::as_str), Some("hello"));
360    }
361    #[test]
362    fn contains_tracks_types() {
363        let mut c = ToolContext::new();
364        c.insert(42u32);
365        assert!(c.contains::<u32>());
366        assert!(!c.contains::<String>());
367    }
368    #[test]
369    fn clone_produces_independent_copy() {
370        let mut c = ToolContext::new();
371        c.insert(42u32);
372        let mut clone = c.clone();
373        clone.insert(99u32);
374        assert_eq!(c.get::<u32>(), Some(&42));
375        assert_eq!(clone.get::<u32>(), Some(&99));
376    }
377    #[test]
378    fn clone_deep_copies_heap_values() {
379        let mut c = ToolContext::new();
380        c.insert(vec![1u8, 2, 3]);
381        let mut clone = c.clone();
382        clone.get_mut::<Vec<u8>>().unwrap().push(4);
383        assert_eq!(c.get::<Vec<u8>>(), Some(&vec![1, 2, 3]));
384        assert_eq!(clone.get::<Vec<u8>>(), Some(&vec![1, 2, 3, 4]));
385    }
386    #[test]
387    fn clone_preserves_intentionally_shared_value_state() {
388        let shared = std::sync::Arc::new(std::sync::Mutex::new(1_u32));
389        let mut context = ToolContext::new();
390        context.insert(shared.clone());
391
392        let snapshot = context.for_dispatch();
393        *snapshot
394            .get::<std::sync::Arc<std::sync::Mutex<u32>>>()
395            .expect("shared value")
396            .lock()
397            .expect("shared value lock") = 2;
398
399        assert_eq!(*shared.lock().expect("shared value lock"), 2);
400        assert!(context.contains::<std::sync::Arc<std::sync::Mutex<u32>>>());
401    }
402    #[test]
403    fn empty_context_is_default_and_allocation_free() {
404        let c = ToolContext::default();
405        assert!(!c.contains::<u32>());
406        assert!(c.inbound.map.is_none());
407        assert!(c.result.map.is_none());
408    }
409    #[test]
410    fn get_mut_modifies_in_place() {
411        let mut c = ToolContext::new();
412        c.insert(42u32);
413        *c.get_mut::<u32>().unwrap() = 99;
414        assert_eq!(c.get::<u32>(), Some(&99));
415    }
416    #[test]
417    fn remove_returns_value_and_clears_entry() {
418        let mut c = ToolContext::new();
419        c.insert(42u32);
420        assert_eq!(c.remove::<u32>(), Some(42));
421        assert!(!c.contains::<u32>());
422    }
423    #[test]
424    fn remove_missing_type_returns_none() {
425        assert_eq!(ToolContext::new().remove::<u32>(), None);
426    }
427    #[test]
428    fn require_present_returns_value() {
429        let mut c = ToolContext::new();
430        c.insert(42u32);
431        assert_eq!(c.require::<u32>().copied(), Ok(42));
432    }
433    #[test]
434    fn require_missing_names_type() {
435        let e = ToolContext::new().require::<u32>().unwrap_err();
436        assert!(e.to_string().contains("u32"));
437    }
438    #[test]
439    fn result_metadata_round_trips_and_requires() {
440        #[derive(Clone, Debug, PartialEq)]
441        struct Id(u32);
442        let mut c = ToolContext::new();
443        c.insert_result(Id(7));
444        assert_eq!(c.result::<Id>(), Some(&Id(7)));
445        assert_eq!(c.require_result::<Id>(), Ok(&Id(7)));
446        assert!(c.get::<Id>().is_none());
447    }
448    #[test]
449    fn debug_reports_types_without_values() {
450        #[derive(Clone)]
451        struct Secret(&'static str);
452        let mut c = ToolContext::new();
453        c.insert(42u32);
454        c.insert_result(Secret("do-not-print"));
455        let d = format!("{c:?}");
456        assert!(d.contains("u32"));
457        assert!(d.contains("Secret"));
458        assert!(!d.contains("do-not-print"));
459        assert_eq!(c.result::<Secret>().map(|s| s.0), Some("do-not-print"));
460    }
461    #[test]
462    fn dispatch_snapshot_isolates_inbound_and_publishes_only_result_metadata() {
463        let mut c = ToolContext::new();
464        c.insert(7u32);
465        c.insert_result("old".to_string());
466        let mut d = c.for_dispatch();
467        assert_eq!(d.get::<u32>(), Some(&7));
468        assert!(d.result::<String>().is_none());
469        *d.get_mut::<u32>().expect("snapshot value") = 8;
470        d.insert_result("new".to_string());
471
472        c.accept_dispatch_result(d);
473        assert_eq!(c.get::<u32>(), Some(&7));
474        assert_eq!(c.result::<String>().map(String::as_str), Some("new"));
475    }
476    #[test]
477    fn many_distinct_types_round_trip_through_type_id_hasher() {
478        #[derive(Clone, PartialEq, Debug)]
479        struct A(u8);
480        #[derive(Clone, PartialEq, Debug)]
481        struct B(u16);
482        let mut c = ToolContext::new();
483        c.insert(A(1));
484        c.insert(B(2));
485        c.insert(3u32);
486        c.insert("four".to_string());
487        assert_eq!(c.get::<A>(), Some(&A(1)));
488        assert_eq!(c.get::<B>(), Some(&B(2)));
489        assert_eq!(c.get::<u32>(), Some(&3));
490        assert_eq!(c.get::<String>().map(String::as_str), Some("four"));
491    }
492}
rig_agent/tool/extensions.rs

rig_agent/tool/
extensions.rs
