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    /// Clone only the inbound values, for nested tool execution.
269    pub(crate) fn inbound_only(&self) -> Self {
270        self.for_dispatch()
271    }
272}
273
274impl std::fmt::Debug for ToolContext {
275    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
276        f.debug_struct("ToolContext")
277            .field("inbound_entries", &self.inbound.len())
278            .field("inbound_types", &self.inbound.type_names())
279            .field("result_entries", &self.result.len())
280            .field("result_types", &self.result.type_names())
281            .finish()
282    }
283}
284
285/// A required typed value was missing from a [`ToolContext`].
286#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
287#[error("required tool context value of type `{0}` was not found")]
288pub struct MissingToolContext(pub &'static str);
289
290impl From<MissingToolContext> for ToolExecutionError {
291    fn from(error: MissingToolContext) -> Self {
292        ToolExecutionError::other(error.to_string()).with_source(error)
293    }
294}
295
296#[cfg(not(target_family = "wasm"))]
297const _: fn() = || {
298    fn assert_send_sync<T: Send + Sync>() {}
299    assert_send_sync::<ToolContext>();
300};
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    #[test]
307    fn context_separates_inbound_and_result_values() {
308        let mut context = ToolContext::new();
309        context.insert(42_u32);
310        context.insert_result("request-1".to_string());
311        assert_eq!(context.get::<u32>(), Some(&42));
312        assert_eq!(
313            context.result::<String>().map(String::as_str),
314            Some("request-1")
315        );
316
317        let next = context.for_dispatch();
318        assert_eq!(next.get::<u32>(), Some(&42));
319        assert!(next.result::<String>().is_none());
320    }
321
322    #[test]
323    fn missing_context_converts_into_a_tool_execution_error() {
324        fn require_value(context: &ToolContext) -> Result<u32, ToolExecutionError> {
325            Ok(*context.require::<u32>()?)
326        }
327
328        let error = require_value(&ToolContext::new()).unwrap_err();
329        assert!(error.is::<MissingToolContext>());
330        assert_eq!(
331            error.model_feedback(),
332            Some("required tool context value of type `u32` was not found")
333        );
334    }
335}
336
337#[cfg(test)]
338mod migrated_tests {
339    use super::*;
340
341    #[test]
342    fn insert_and_get_returns_value() {
343        let mut c = ToolContext::new();
344        assert_eq!(c.insert(42u32), None);
345        assert_eq!(c.get::<u32>(), Some(&42));
346    }
347    #[test]
348    fn get_missing_type_returns_none() {
349        assert_eq!(ToolContext::new().get::<u32>(), None);
350    }
351    #[test]
352    fn insert_overwrites_and_returns_previous() {
353        let mut c = ToolContext::new();
354        c.insert(1u32);
355        assert_eq!(c.insert(2u32), Some(1));
356        assert_eq!(c.get::<u32>(), Some(&2));
357    }
358    #[test]
359    fn different_types_are_independent() {
360        let mut c = ToolContext::new();
361        c.insert(42u32);
362        c.insert("hello".to_string());
363        assert_eq!(c.get::<u32>(), Some(&42));
364        assert_eq!(c.get::<String>().map(String::as_str), Some("hello"));
365    }
366    #[test]
367    fn contains_tracks_types() {
368        let mut c = ToolContext::new();
369        c.insert(42u32);
370        assert!(c.contains::<u32>());
371        assert!(!c.contains::<String>());
372    }
373    #[test]
374    fn clone_produces_independent_copy() {
375        let mut c = ToolContext::new();
376        c.insert(42u32);
377        let mut clone = c.clone();
378        clone.insert(99u32);
379        assert_eq!(c.get::<u32>(), Some(&42));
380        assert_eq!(clone.get::<u32>(), Some(&99));
381    }
382    #[test]
383    fn clone_deep_copies_heap_values() {
384        let mut c = ToolContext::new();
385        c.insert(vec![1u8, 2, 3]);
386        let mut clone = c.clone();
387        clone.get_mut::<Vec<u8>>().unwrap().push(4);
388        assert_eq!(c.get::<Vec<u8>>(), Some(&vec![1, 2, 3]));
389        assert_eq!(clone.get::<Vec<u8>>(), Some(&vec![1, 2, 3, 4]));
390    }
391    #[test]
392    fn clone_preserves_intentionally_shared_value_state() {
393        let shared = std::sync::Arc::new(std::sync::Mutex::new(1_u32));
394        let mut context = ToolContext::new();
395        context.insert(shared.clone());
396
397        let snapshot = context.for_dispatch();
398        *snapshot
399            .get::<std::sync::Arc<std::sync::Mutex<u32>>>()
400            .expect("shared value")
401            .lock()
402            .expect("shared value lock") = 2;
403
404        assert_eq!(*shared.lock().expect("shared value lock"), 2);
405        assert!(context.contains::<std::sync::Arc<std::sync::Mutex<u32>>>());
406    }
407    #[test]
408    fn empty_context_is_default_and_allocation_free() {
409        let c = ToolContext::default();
410        assert!(!c.contains::<u32>());
411        assert!(c.inbound.map.is_none());
412        assert!(c.result.map.is_none());
413    }
414    #[test]
415    fn get_mut_modifies_in_place() {
416        let mut c = ToolContext::new();
417        c.insert(42u32);
418        *c.get_mut::<u32>().unwrap() = 99;
419        assert_eq!(c.get::<u32>(), Some(&99));
420    }
421    #[test]
422    fn remove_returns_value_and_clears_entry() {
423        let mut c = ToolContext::new();
424        c.insert(42u32);
425        assert_eq!(c.remove::<u32>(), Some(42));
426        assert!(!c.contains::<u32>());
427    }
428    #[test]
429    fn remove_missing_type_returns_none() {
430        assert_eq!(ToolContext::new().remove::<u32>(), None);
431    }
432    #[test]
433    fn require_present_returns_value() {
434        let mut c = ToolContext::new();
435        c.insert(42u32);
436        assert_eq!(c.require::<u32>().copied(), Ok(42));
437    }
438    #[test]
439    fn require_missing_names_type() {
440        let e = ToolContext::new().require::<u32>().unwrap_err();
441        assert!(e.to_string().contains("u32"));
442    }
443    #[test]
444    fn result_metadata_round_trips_and_requires() {
445        #[derive(Clone, Debug, PartialEq)]
446        struct Id(u32);
447        let mut c = ToolContext::new();
448        c.insert_result(Id(7));
449        assert_eq!(c.result::<Id>(), Some(&Id(7)));
450        assert_eq!(c.require_result::<Id>(), Ok(&Id(7)));
451        assert!(c.get::<Id>().is_none());
452    }
453    #[test]
454    fn debug_reports_types_without_values() {
455        #[derive(Clone)]
456        struct Secret(&'static str);
457        let mut c = ToolContext::new();
458        c.insert(42u32);
459        c.insert_result(Secret("do-not-print"));
460        let d = format!("{c:?}");
461        assert!(d.contains("u32"));
462        assert!(d.contains("Secret"));
463        assert!(!d.contains("do-not-print"));
464        assert_eq!(c.result::<Secret>().map(|s| s.0), Some("do-not-print"));
465    }
466    #[test]
467    fn dispatch_snapshot_isolates_inbound_and_publishes_only_result_metadata() {
468        let mut c = ToolContext::new();
469        c.insert(7u32);
470        c.insert_result("old".to_string());
471        let mut d = c.for_dispatch();
472        assert_eq!(d.get::<u32>(), Some(&7));
473        assert!(d.result::<String>().is_none());
474        *d.get_mut::<u32>().expect("snapshot value") = 8;
475        d.insert_result("new".to_string());
476
477        c.accept_dispatch_result(d);
478        assert_eq!(c.get::<u32>(), Some(&7));
479        assert_eq!(c.result::<String>().map(String::as_str), Some("new"));
480    }
481    #[test]
482    fn many_distinct_types_round_trip_through_type_id_hasher() {
483        #[derive(Clone, PartialEq, Debug)]
484        struct A(u8);
485        #[derive(Clone, PartialEq, Debug)]
486        struct B(u16);
487        let mut c = ToolContext::new();
488        c.insert(A(1));
489        c.insert(B(2));
490        c.insert(3u32);
491        c.insert("four".to_string());
492        assert_eq!(c.get::<A>(), Some(&A(1)));
493        assert_eq!(c.get::<B>(), Some(&B(2)));
494        assert_eq!(c.get::<u32>(), Some(&3));
495        assert_eq!(c.get::<String>().map(String::as_str), Some("four"));
496    }
497}
rig_agent/tool/extensions.rs

rig_agent/tool/
extensions.rs
