rig_core/tool/
extensions.rs

1//! Type-erased extension maps for tool execution.
2//!
3//! This module provides two type-maps built on the same storage primitive:
4//!
5//! - [`ToolCallExtensions`] — attached by a *caller* before a run and read by a
6//!   tool at call time (auth tokens, session IDs, …). It travels *into* a tool.
7//! - [`ToolResultExtensions`] — attached by a *tool* to its structured
8//!   [`ToolExecutionResult`](crate::tool::ToolExecutionResult) and read by hooks,
9//!   tracing, and policies. It travels *out of* a tool and is **never** sent to
10//!   the model.
11//!
12//! Both wrap the private [`TypeMap`], which follows the `AnyClone` pattern from
13//! the [`http::Extensions`] type in the `http` crate, including the no-op
14//! [`IdHasher`] over `TypeId` keys.
15//!
16//! [`http::Extensions`]: https://docs.rs/http/latest/http/struct.Extensions.html
17
18use std::any::{Any, TypeId, type_name};
19use std::collections::HashMap;
20use std::hash::{BuildHasherDefault, Hasher};
21
22use crate::wasm_compat::{WasmCompatSend, WasmCompatSync};
23
24type AnyMap = HashMap<TypeId, Box<dyn AnyClone>, BuildHasherDefault<IdHasher>>;
25
26// --- IdHasher (modeled after `http::Extensions`) ---
27
28/// Hasher for the `TypeId` keys of the type-map.
29///
30/// A `TypeId` is already a high-quality hash, so the keys do not need to be
31/// re-hashed with the default `SipHash`. The fast path stores the `u64` that
32/// `TypeId`'s `Hash` impl writes directly. The byte-oriented [`write`](Hasher::write)
33/// fallback keeps this correct (never panicking) if a future `TypeId`
34/// representation hashes via a different `Hasher` method — a poor hash only
35/// costs extra probe work, never correctness, because the map still compares
36/// full `TypeId` keys.
37#[derive(Default)]
38struct IdHasher(u64);
39
40impl Hasher for IdHasher {
41    #[inline]
42    fn write_u64(&mut self, id: u64) {
43        self.0 = id;
44    }
45
46    #[inline]
47    fn write(&mut self, bytes: &[u8]) {
48        // Fallback path: fold the bytes so we never panic regardless of how
49        // `TypeId` chooses to hash itself.
50        for &byte in bytes {
51            self.0 = self.0.rotate_left(8) ^ u64::from(byte);
52        }
53    }
54
55    #[inline]
56    fn finish(&self) -> u64 {
57        self.0
58    }
59}
60
61// --- AnyClone helper trait (modeled after `http::Extensions`) ---
62//
63// `WasmCompatSend`/`WasmCompatSync` are `Send`/`Sync` on every non-wasm target
64// and unconstrained on `wasm`, so a single trait definition covers both targets
65// (no `cfg`-duplicated copies). On non-wasm this transitively makes
66// `dyn AnyClone` — and therefore both extension maps — `Send + Sync`, which is
67// required because they are borrowed across `.await` points in async tool
68// execution. The `assert_send_sync` checks below pin that property.
69
70trait AnyClone: Any + WasmCompatSend + WasmCompatSync {
71    fn clone_box(&self) -> Box<dyn AnyClone>;
72    fn as_any(&self) -> &dyn Any;
73    fn as_any_mut(&mut self) -> &mut dyn Any;
74    fn into_any(self: Box<Self>) -> Box<dyn Any>;
75    fn type_name(&self) -> &'static str;
76}
77
78impl<T: Clone + WasmCompatSend + WasmCompatSync + 'static> AnyClone for T {
79    fn clone_box(&self) -> Box<dyn AnyClone> {
80        Box::new(self.clone())
81    }
82    fn as_any(&self) -> &dyn Any {
83        self
84    }
85    fn as_any_mut(&mut self) -> &mut dyn Any {
86        self
87    }
88    fn into_any(self: Box<Self>) -> Box<dyn Any> {
89        self
90    }
91    fn type_name(&self) -> &'static str {
92        type_name::<T>()
93    }
94}
95
96impl Clone for Box<dyn AnyClone> {
97    fn clone(&self) -> Self {
98        // Explicit deref to dispatch via the trait object's vtable, not the
99        // blanket `AnyClone` impl on `Box` itself (which would recurse).
100        (**self).clone_box()
101    }
102}
103
104// --- TypeMap: shared storage for the two public extension maps ---
105
106/// A `TypeId`-keyed map of cloneable, type-erased values.
107///
108/// The storage primitive shared by [`ToolCallExtensions`] and
109/// [`ToolResultExtensions`]. Uses `Option<Box<HashMap>>` internally so that an
110/// empty map (the common case) requires zero allocation.
111#[derive(Default, Clone)]
112pub(crate) struct TypeMap {
113    map: Option<Box<AnyMap>>,
114}
115
116impl TypeMap {
117    /// A `'static`, allocation-free empty map, so dispatch layers can hand out a
118    /// borrow of a default without owning a fresh value.
119    pub(crate) const EMPTY: TypeMap = TypeMap { map: None };
120
121    pub(crate) fn insert<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(
122        &mut self,
123        val: T,
124    ) -> Option<T> {
125        self.map
126            .get_or_insert_with(Default::default)
127            .insert(TypeId::of::<T>(), Box::new(val))
128            .and_then(|prev| prev.into_any().downcast::<T>().ok())
129            .map(|boxed| *boxed)
130    }
131
132    pub(crate) fn get<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<&T> {
133        self.map
134            .as_ref()
135            .and_then(|map| map.get(&TypeId::of::<T>()))
136            // Explicit deref to dispatch via the trait object's vtable, not the
137            // blanket `AnyClone` impl that `Box` itself satisfies.
138            .and_then(|boxed| (**boxed).as_any().downcast_ref::<T>())
139    }
140
141    pub(crate) fn require<T: WasmCompatSend + WasmCompatSync + 'static>(
142        &self,
143    ) -> Result<&T, MissingExtension> {
144        self.get::<T>().ok_or(MissingExtension(type_name::<T>()))
145    }
146
147    pub(crate) fn get_mut<T: WasmCompatSend + WasmCompatSync + 'static>(
148        &mut self,
149    ) -> Option<&mut T> {
150        self.map
151            .as_mut()
152            .and_then(|map| map.get_mut(&TypeId::of::<T>()))
153            .and_then(|boxed| (**boxed).as_any_mut().downcast_mut::<T>())
154    }
155
156    pub(crate) fn remove<T: WasmCompatSend + WasmCompatSync + 'static>(&mut self) -> Option<T> {
157        self.map
158            .as_mut()
159            .and_then(|map| map.remove(&TypeId::of::<T>()))
160            .and_then(|boxed| boxed.into_any().downcast::<T>().ok())
161            .map(|boxed| *boxed)
162    }
163
164    pub(crate) fn contains<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> bool {
165        self.map
166            .as_ref()
167            .is_some_and(|map| map.contains_key(&TypeId::of::<T>()))
168    }
169
170    pub(crate) fn len(&self) -> usize {
171        self.map.as_ref().map_or(0, |map| map.len())
172    }
173
174    /// The type names currently stored, for `Debug` output.
175    fn type_names(&self) -> Vec<&'static str> {
176        self.map
177            .as_ref()
178            .map(|map| map.values().map(|v| (**v).type_name()).collect())
179            .unwrap_or_default()
180    }
181}
182
183/// Write a `Debug` representation of a [`TypeMap`]-backed extension map.
184fn debug_type_map(f: &mut std::fmt::Formatter<'_>, name: &str, map: &TypeMap) -> std::fmt::Result {
185    let mut dbg = f.debug_struct(name);
186    dbg.field("entries", &map.len());
187    if map.len() > 0 {
188        dbg.field("types", &map.type_names());
189    }
190    dbg.finish()
191}
192
193// --- ToolCallExtensions ---
194
195/// Per-call runtime extensions supplied to a tool *by its caller*.
196///
197/// A type-map that allows callers to attach arbitrary typed values and tools to
198/// extract them. Tools that don't need any extensions ignore them.
199///
200/// Inspired by [`http::Extensions`](https://docs.rs/http/latest/http/struct.Extensions.html).
201/// Backed by the shared internal `TypeMap`, so empty extensions (the common case when no
202/// caller-provided values are needed) require zero allocation.
203///
204/// Tools receive these by shared reference, so a tool reads values with
205/// [`get`](Self::get) / [`require`](Self::require) / [`contains`](Self::contains);
206/// the mutating accessors ([`insert`](Self::insert), [`get_mut`](Self::get_mut),
207/// [`remove`](Self::remove)) are owner-side, used by the caller that builds the
208/// extensions before a run.
209///
210/// Values are keyed by [`TypeId`], so [`get`](Self::get) returns `None` both
211/// when nothing was inserted under that type *and* when a different type was
212/// inserted. For tools that genuinely require a value, prefer
213/// [`require`](Self::require), which returns a descriptive error instead of a
214/// silent `None`.
215///
216/// The result-side counterpart — metadata a tool attaches to its output — is
217/// [`ToolResultExtensions`](crate::tool::ToolResultExtensions).
218///
219/// # Example
220/// ```
221/// use rig_core::tool::ToolCallExtensions;
222///
223/// let mut extensions = ToolCallExtensions::new();
224/// assert_eq!(extensions.insert(42u32), None); // no prior value
225/// assert_eq!(extensions.get::<u32>(), Some(&42));
226/// assert_eq!(extensions.insert(7u32), Some(42)); // returns the displaced value
227/// ```
228#[derive(Default, Clone)]
229pub struct ToolCallExtensions {
230    inner: TypeMap,
231}
232
233impl ToolCallExtensions {
234    /// Shared empty instance. Lets dispatch layers that need a default value
235    /// hand out a `'static` reference instead of constructing (and having to
236    /// own) a fresh value just to borrow it.
237    pub(crate) const EMPTY: ToolCallExtensions = ToolCallExtensions {
238        inner: TypeMap::EMPTY,
239    };
240
241    /// Create an empty set of extensions.
242    pub const fn new() -> Self {
243        Self::EMPTY
244    }
245
246    /// Insert a typed value, returning the previous value of the same type if
247    /// one was present (mirroring [`http::Extensions::insert`] and
248    /// [`HashMap::insert`]).
249    ///
250    /// [`http::Extensions::insert`]: https://docs.rs/http/latest/http/struct.Extensions.html#method.insert
251    pub fn insert<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(
252        &mut self,
253        val: T,
254    ) -> Option<T> {
255        self.inner.insert(val)
256    }
257
258    /// Get a reference to a value by type. Returns `None` if not present.
259    pub fn get<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<&T> {
260        self.inner.get::<T>()
261    }
262
263    /// Get a reference to a value by type, returning a descriptive error instead
264    /// of `None` when it is absent.
265    ///
266    /// Prefer this over [`get`](Self::get) for tools that *require* an extension
267    /// value (auth tokens, session IDs, …): the error names the missing type,
268    /// turning a silent `None` into an actionable failure.
269    pub fn require<T: WasmCompatSend + WasmCompatSync + 'static>(
270        &self,
271    ) -> Result<&T, MissingExtension> {
272        self.inner.require::<T>()
273    }
274
275    /// Get a mutable reference to a value by type. Returns `None` if not present.
276    pub fn get_mut<T: WasmCompatSend + WasmCompatSync + 'static>(&mut self) -> Option<&mut T> {
277        self.inner.get_mut::<T>()
278    }
279
280    /// Remove a value by type, returning it if present.
281    pub fn remove<T: WasmCompatSend + WasmCompatSync + 'static>(&mut self) -> Option<T> {
282        self.inner.remove::<T>()
283    }
284
285    /// Check if a value of the given type is present.
286    pub fn contains<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> bool {
287        self.inner.contains::<T>()
288    }
289
290    /// Number of values currently stored.
291    pub fn len(&self) -> usize {
292        self.inner.len()
293    }
294
295    /// Whether no values are stored.
296    pub fn is_empty(&self) -> bool {
297        self.len() == 0
298    }
299}
300
301impl std::fmt::Debug for ToolCallExtensions {
302    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
303        debug_type_map(f, "ToolCallExtensions", &self.inner)
304    }
305}
306
307// --- ToolResultExtensions ---
308
309/// Type-erased metadata a *tool* attaches to its structured result.
310///
311/// The result-side counterpart of [`ToolCallExtensions`]. A tool builds these on
312/// its [`ToolReturn`](crate::tool::ToolReturn) (or directly on a
313/// [`ToolExecutionResult`](crate::tool::ToolExecutionResult)) to carry
314/// provider-specific or application-specific values — raw HTTP headers, a
315/// provider response id, retry hints — alongside the model-visible output.
316///
317/// These are **machine-visible only**: they are surfaced to hooks (via
318/// [`StepEvent::ToolResult`](crate::agent::StepEvent::ToolResult)), tracing, and
319/// policies, but are **never** serialized into the tool result the model sees.
320/// This is the same "not sent to the model" contract as Pydantic AI's tool-return
321/// `metadata` and the Vercel AI SDK's `toolMetadata`.
322///
323/// Backed by the same internal `TypeMap` as [`ToolCallExtensions`], so an empty set (the
324/// common case) allocates nothing.
325///
326/// # Example
327/// ```
328/// use rig_core::tool::ToolResultExtensions;
329///
330/// #[derive(Clone, Debug, PartialEq)]
331/// struct HttpStatus(u16);
332///
333/// let mut extensions = ToolResultExtensions::new();
334/// extensions.insert(HttpStatus(404));
335/// assert_eq!(extensions.get::<HttpStatus>(), Some(&HttpStatus(404)));
336/// ```
337#[derive(Default, Clone)]
338pub struct ToolResultExtensions {
339    inner: TypeMap,
340}
341
342impl ToolResultExtensions {
343    /// Create an empty set of result extensions (allocation-free).
344    pub const fn new() -> Self {
345        Self {
346            inner: TypeMap::EMPTY,
347        }
348    }
349
350    /// Insert a typed value, returning the previous value of the same type if
351    /// one was present.
352    pub fn insert<T: Clone + WasmCompatSend + WasmCompatSync + 'static>(
353        &mut self,
354        val: T,
355    ) -> Option<T> {
356        self.inner.insert(val)
357    }
358
359    /// Get a reference to a value by type. Returns `None` if not present.
360    pub fn get<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> Option<&T> {
361        self.inner.get::<T>()
362    }
363
364    /// Get a reference to a value by type, returning a descriptive error naming
365    /// the missing type instead of a silent `None`.
366    pub fn require<T: WasmCompatSend + WasmCompatSync + 'static>(
367        &self,
368    ) -> Result<&T, MissingExtension> {
369        self.inner.require::<T>()
370    }
371
372    /// Get a mutable reference to a value by type. Returns `None` if not present.
373    pub fn get_mut<T: WasmCompatSend + WasmCompatSync + 'static>(&mut self) -> Option<&mut T> {
374        self.inner.get_mut::<T>()
375    }
376
377    /// Remove a value by type, returning it if present.
378    pub fn remove<T: WasmCompatSend + WasmCompatSync + 'static>(&mut self) -> Option<T> {
379        self.inner.remove::<T>()
380    }
381
382    /// Check if a value of the given type is present.
383    pub fn contains<T: WasmCompatSend + WasmCompatSync + 'static>(&self) -> bool {
384        self.inner.contains::<T>()
385    }
386
387    /// Number of values currently stored.
388    pub fn len(&self) -> usize {
389        self.inner.len()
390    }
391
392    /// Whether no values are stored.
393    pub fn is_empty(&self) -> bool {
394        self.len() == 0
395    }
396}
397
398impl std::fmt::Debug for ToolResultExtensions {
399    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
400        debug_type_map(f, "ToolResultExtensions", &self.inner)
401    }
402}
403
404/// Error returned by [`ToolCallExtensions::require`] and
405/// [`ToolResultExtensions::require`] when the requested value is not present.
406#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
407#[error("required tool extension of type `{0}` was not found")]
408pub struct MissingExtension(pub &'static str);
409
410// Both extension maps must stay `Send + Sync` on native targets: the agent loop
411// borrows them across `.await` while executing tools. This fails to compile if a
412// future change (e.g. relaxing the `AnyClone` bounds) drops the property.
413#[cfg(not(target_family = "wasm"))]
414const _: fn() = || {
415    fn assert_send_sync<T: Send + Sync>() {}
416    assert_send_sync::<ToolCallExtensions>();
417    assert_send_sync::<ToolResultExtensions>();
418};
419
420#[cfg(test)]
421mod tests {
422    use super::*;
423
424    #[test]
425    fn insert_and_get_returns_value() {
426        let mut extensions = ToolCallExtensions::new();
427        assert_eq!(extensions.insert(42u32), None);
428        assert_eq!(extensions.get::<u32>(), Some(&42));
429    }
430
431    #[test]
432    fn get_missing_type_returns_none() {
433        let extensions = ToolCallExtensions::new();
434        assert_eq!(extensions.get::<u32>(), None);
435    }
436
437    #[test]
438    fn insert_overwrites_and_returns_previous() {
439        let mut extensions = ToolCallExtensions::new();
440        assert_eq!(extensions.insert(1u32), None);
441        assert_eq!(extensions.insert(2u32), Some(1));
442        assert_eq!(extensions.get::<u32>(), Some(&2));
443    }
444
445    #[test]
446    fn different_types_are_independent() {
447        let mut extensions = ToolCallExtensions::new();
448        extensions.insert(42u32);
449        extensions.insert("hello".to_string());
450        assert_eq!(extensions.get::<u32>(), Some(&42));
451        assert_eq!(extensions.get::<String>(), Some(&"hello".to_string()));
452    }
453
454    #[test]
455    fn contains_returns_true_when_present() {
456        let mut extensions = ToolCallExtensions::new();
457        extensions.insert(42u32);
458        assert!(extensions.contains::<u32>());
459        assert!(!extensions.contains::<String>());
460    }
461
462    #[test]
463    fn clone_produces_independent_copy() {
464        let mut extensions = ToolCallExtensions::new();
465        extensions.insert(42u32);
466        let mut cloned = extensions.clone();
467        cloned.insert(99u32);
468        assert_eq!(extensions.get::<u32>(), Some(&42));
469        assert_eq!(cloned.get::<u32>(), Some(&99));
470    }
471
472    #[test]
473    fn clone_deep_copies_inner_value() {
474        // Insert a heap-allocated value, clone the extensions, then mutate the
475        // clone's inner value in place. A shallow clone (sharing the boxed
476        // value) would let this mutation leak back into the original; a correct
477        // `clone_box` deep-copies, so the original stays unchanged.
478        let mut extensions = ToolCallExtensions::new();
479        extensions.insert(vec![1u8, 2, 3]);
480        let mut cloned = extensions.clone();
481        cloned.get_mut::<Vec<u8>>().unwrap().push(4);
482        assert_eq!(extensions.get::<Vec<u8>>(), Some(&vec![1, 2, 3]));
483        assert_eq!(cloned.get::<Vec<u8>>(), Some(&vec![1, 2, 3, 4]));
484    }
485
486    #[test]
487    fn debug_shows_entry_count_and_types() {
488        let mut extensions = ToolCallExtensions::new();
489        extensions.insert(42u32);
490        extensions.insert("hi".to_string());
491        let debug = format!("{:?}", extensions);
492        assert!(debug.contains("entries: 2"));
493        assert!(debug.contains("u32"));
494        assert!(debug.contains("String"));
495    }
496
497    #[test]
498    fn empty_extensions_is_default() {
499        let extensions = ToolCallExtensions::default();
500        assert!(!extensions.contains::<u32>());
501    }
502
503    #[test]
504    fn empty_extensions_has_no_allocation() {
505        let extensions = ToolCallExtensions::new();
506        assert!(extensions.inner.map.is_none());
507    }
508
509    #[test]
510    fn get_mut_modifies_in_place() {
511        let mut extensions = ToolCallExtensions::new();
512        extensions.insert(42u32);
513        if let Some(val) = extensions.get_mut::<u32>() {
514            *val = 99;
515        }
516        assert_eq!(extensions.get::<u32>(), Some(&99));
517    }
518
519    #[test]
520    fn remove_returns_value_and_clears_entry() {
521        let mut extensions = ToolCallExtensions::new();
522        extensions.insert(42u32);
523        assert_eq!(extensions.remove::<u32>(), Some(42));
524        assert!(!extensions.contains::<u32>());
525    }
526
527    #[test]
528    fn remove_missing_type_returns_none() {
529        let mut extensions = ToolCallExtensions::new();
530        assert_eq!(extensions.remove::<u32>(), None);
531    }
532
533    #[test]
534    fn require_present_returns_value() {
535        let mut extensions = ToolCallExtensions::new();
536        extensions.insert(42u32);
537        assert_eq!(extensions.require::<u32>().copied(), Ok(42));
538    }
539
540    #[test]
541    fn require_missing_names_the_type() {
542        let extensions = ToolCallExtensions::new();
543        let err = extensions.require::<u32>().unwrap_err();
544        assert!(err.to_string().contains("u32"));
545    }
546
547    #[test]
548    fn len_and_is_empty_track_entries() {
549        let mut extensions = ToolCallExtensions::new();
550        assert!(extensions.is_empty());
551        assert_eq!(extensions.len(), 0);
552        extensions.insert(1u32);
553        extensions.insert("two".to_string());
554        assert!(!extensions.is_empty());
555        assert_eq!(extensions.len(), 2);
556        extensions.remove::<u32>();
557        assert_eq!(extensions.len(), 1);
558    }
559
560    #[test]
561    fn many_distinct_types_round_trip_through_id_hasher() {
562        // Guards the custom IdHasher: distinct TypeId keys must not collide in a
563        // way that corrupts lookups. Insert several heterogeneous types and
564        // confirm each is retrievable with its own value.
565        #[derive(Clone, PartialEq, Debug)]
566        struct A(u8);
567        #[derive(Clone, PartialEq, Debug)]
568        struct B(u16);
569        #[derive(Clone, PartialEq, Debug)]
570        struct C(u32);
571        #[derive(Clone, PartialEq, Debug)]
572        struct D(u64);
573
574        let mut extensions = ToolCallExtensions::new();
575        extensions.insert(A(1));
576        extensions.insert(B(2));
577        extensions.insert(C(3));
578        extensions.insert(D(4));
579        extensions.insert(5u8);
580        extensions.insert(6i64);
581        extensions.insert("seven".to_string());
582        extensions.insert(8.0f64);
583
584        assert_eq!(extensions.len(), 8);
585        assert_eq!(extensions.get::<A>(), Some(&A(1)));
586        assert_eq!(extensions.get::<B>(), Some(&B(2)));
587        assert_eq!(extensions.get::<C>(), Some(&C(3)));
588        assert_eq!(extensions.get::<D>(), Some(&D(4)));
589        assert_eq!(extensions.get::<u8>(), Some(&5));
590        assert_eq!(extensions.get::<i64>(), Some(&6));
591        assert_eq!(extensions.get::<String>(), Some(&"seven".to_string()));
592        assert_eq!(extensions.get::<f64>(), Some(&8.0));
593    }
594
595    // --- ToolResultExtensions: shares the TypeMap, so a couple of checks
596    // suffice to confirm the wrapper delegates correctly and stays independent
597    // of the call-side map. ---
598
599    #[test]
600    fn result_extensions_round_trip_and_require() {
601        #[derive(Clone, Debug, PartialEq)]
602        struct RawHeaders(Vec<(String, String)>);
603
604        let mut extensions = ToolResultExtensions::new();
605        assert!(extensions.is_empty());
606        extensions.insert(RawHeaders(vec![("x-req-id".into(), "abc".into())]));
607        assert_eq!(
608            extensions.get::<RawHeaders>(),
609            Some(&RawHeaders(vec![("x-req-id".into(), "abc".into())]))
610        );
611        assert!(extensions.require::<u32>().is_err());
612        assert_eq!(extensions.len(), 1);
613        let debug = format!("{extensions:?}");
614        assert!(debug.contains("ToolResultExtensions"));
615        assert!(debug.contains("entries: 1"));
616    }
617}
rig_core/tool/extensions.rs

rig_core/tool/
extensions.rs
