Skip to main content

llm_tool/
rust_resource.rs

1//! Strongly-typed Rust resource trait and URI template matching.
2
3use alloc::{
4    borrow::Cow,
5    boxed::Box,
6    format,
7    string::{String, ToString},
8    vec::Vec,
9};
10use core::{future::Future, pin::Pin};
11
12use super::types::{RegistryItem, ResourceDefinition, ResourceOutput, ToolError};
13
14/// A custom resource or resource template implemented in Rust.
15pub trait RustResource: Send + Sync {
16    /// Strongly-typed parameters extracted from URI variables.
17    type Params: serde::de::DeserializeOwned + Send;
18
19    /// URI template pattern (e.g. `"file:///logs/{date}/{app}.log"` or `"config://app"`).
20    const URI_TEMPLATE: &'static str;
21
22    /// Unique resource name.
23    const NAME: &'static str;
24
25    /// Human-readable description.
26    const DESCRIPTION: &'static str;
27
28    /// Optional MIME type.
29    const MIME_TYPE: Option<&'static str>;
30
31    /// Return the resource description.
32    fn description(&self) -> Cow<'static, str> {
33        Cow::Borrowed(Self::DESCRIPTION)
34    }
35
36    /// Read the resource with URI and extracted parameters.
37    fn read(
38        &self,
39        uri: &str,
40        params: Self::Params,
41    ) -> impl Future<Output = Result<ResourceOutput, ToolError>> + Send;
42}
43
44/// Build a [`ResourceDefinition`] from any [`RustResource`] implementor.
45///
46/// Infallible: the definition is built purely from associated constants.
47#[must_use]
48pub fn definition_of_resource<T: RustResource>(resource: &T) -> ResourceDefinition {
49    ResourceDefinition {
50        uri_template: T::URI_TEMPLATE.to_string(),
51        name: T::NAME.to_string(),
52        description: resource.description().into_owned(),
53        mime_type: T::MIME_TYPE.map(ToString::to_string),
54    }
55}
56
57/// Helper to match an incoming URI against a URI template pattern with `{variable}` placeholders.
58///
59/// Returns `Some(map)` if the URI matches the pattern, mapping each `{variable}` name
60/// to its extracted value string. Returns `None` if the URI does not match.
61#[must_use]
62pub fn match_uri_template(
63    template: &str,
64    uri: &str,
65) -> Option<alloc::collections::BTreeMap<String, String>> {
66    let mut map = alloc::collections::BTreeMap::new();
67    let mut t_rem = template;
68    let mut u_rem = uri;
69
70    while let Some(start_idx) = t_rem.find('{') {
71        let prefix = &t_rem[..start_idx];
72        if !u_rem.starts_with(prefix) {
73            return None;
74        }
75        u_rem = &u_rem[prefix.len()..];
76        t_rem = &t_rem[start_idx + 1..];
77
78        let end_idx = t_rem.find('}')?;
79        let var_name = &t_rem[..end_idx];
80        t_rem = &t_rem[end_idx + 1..];
81
82        let val_str = if t_rem.is_empty() {
83            let val = u_rem;
84            u_rem = "";
85            val
86        } else {
87            let next_char = t_rem.chars().next()?;
88            let val_end = u_rem.find(next_char)?;
89            let val = &u_rem[..val_end];
90            u_rem = &u_rem[val_end..];
91            val
92        };
93
94        map.insert(var_name.to_string(), val_str.to_string());
95    }
96
97    if t_rem == u_rem { Some(map) } else { None }
98}
99
100/// Type-erased future returned by [`ErasedResource::read_erased`].
101pub(crate) type BoxResourceFuture<'a> =
102    Pin<Box<dyn Future<Output = Result<ResourceOutput, ToolError>> + Send + 'a>>;
103
104/// Type-erased wrapper enabling heterogeneous resource storage.
105///
106/// This is an internal implementation detail of [`ResourceRegistry`]; callers
107/// interact with resources through the registry rather than this trait.
108pub(crate) trait ErasedResource: Send + Sync {
109    /// Check if the incoming URI matches this resource's pattern, extract
110    /// variables, and execute `read`.
111    fn read_erased<'a>(&'a self, uri: &'a str) -> Option<BoxResourceFuture<'a>>;
112}
113
114impl<T: RustResource> ErasedResource for T {
115    fn read_erased<'a>(&'a self, uri: &'a str) -> Option<BoxResourceFuture<'a>> {
116        let params_map = match_uri_template(T::URI_TEMPLATE, uri)?;
117        Some(Box::pin(async move {
118            let deserializer = serde::de::value::MapDeserializer::new(
119                params_map
120                    .into_iter()
121                    .map(|(k, v)| (k, serde::de::value::StringDeserializer::new(v))),
122            );
123            let params: T::Params = serde::de::Deserialize::deserialize(deserializer).map_err(
124                |e: serde::de::value::Error| {
125                    ToolError::new(format!(
126                        "Failed to deserialize resource parameters from URI variables: {e}"
127                    ))
128                },
129            )?;
130            self.read(uri, params).await
131        }))
132    }
133}
134
135/// A registered resource: its cached definition plus the type-erased handler.
136struct RegisteredResource {
137    name: &'static str,
138    definition: ResourceDefinition,
139    erased: Box<dyn ErasedResource>,
140}
141
142/// A registry of resources and resource templates for dynamic dispatch.
143///
144/// Mirrors [`ToolRegistry`](crate::ToolRegistry) for resources: it stores
145/// type-erased [`RustResource`] implementations, caches each
146/// [`ResourceDefinition`] at registration time, and reads the first resource
147/// whose URI template matches an incoming URI, keeping the type-erasure
148/// machinery a private implementation detail.
149///
150/// Unlike [`ToolRegistry`](crate::ToolRegistry) and
151/// [`PromptRegistry`](crate::PromptRegistry),
152/// registration is infallible: a [`ResourceDefinition`] is built purely from
153/// associated constants and never serializes a JSON schema, so there is no
154/// `try_register` counterpart.
155#[derive(Default)]
156pub struct ResourceRegistry {
157    resources: Vec<RegisteredResource>,
158}
159
160impl core::fmt::Debug for ResourceRegistry {
161    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
162        let names: Vec<&str> = self.resources.iter().map(|r| r.name).collect();
163        f.debug_struct("ResourceRegistry")
164            .field("resource_count", &self.resources.len())
165            .field("resource_names", &names)
166            .finish()
167    }
168}
169
170impl ResourceRegistry {
171    /// Create an empty resource registry.
172    #[must_use]
173    pub const fn new() -> Self {
174        Self {
175            resources: Vec::new(),
176        }
177    }
178
179    /// Register a [`RustResource`]. Returns `&mut Self` for chaining.
180    pub fn register<R: RustResource + 'static>(&mut self, resource: R) -> &mut Self {
181        self.resources.push(RegisteredResource {
182            name: R::NAME,
183            definition: definition_of_resource(&resource),
184            erased: Box::new(resource),
185        });
186        self
187    }
188
189    /// Register a [`RustResource`], consuming and returning `Self` for chaining.
190    #[must_use]
191    pub fn with_resource<R: RustResource + 'static>(mut self, resource: R) -> Self {
192        self.register(resource);
193        self
194    }
195
196    /// Collect [`ResourceDefinition`]s for all registered resources.
197    ///
198    /// Returns clones of the cached definitions computed at registration time.
199    #[must_use]
200    pub fn definitions(&self) -> Vec<ResourceDefinition> {
201        self.resources
202            .iter()
203            .map(|entry| entry.definition.clone())
204            .collect()
205    }
206
207    /// Number of registered resources.
208    #[must_use]
209    pub const fn len(&self) -> usize {
210        self.resources.len()
211    }
212
213    /// Whether the registry has no registered resources.
214    #[must_use]
215    pub const fn is_empty(&self) -> bool {
216        self.resources.is_empty()
217    }
218
219    /// Whether a resource with the given name is registered.
220    ///
221    /// Note that resources are *read* by URI (see [`matches`](Self::matches)),
222    /// not by name; this checks the registered resource **name** for parity
223    /// with [`ToolRegistry::contains`](crate::ToolRegistry::contains).
224    #[must_use]
225    pub fn contains(&self, name: &str) -> bool {
226        self.resources.iter().any(|entry| entry.name == name)
227    }
228
229    /// Whether any registered resource's URI template matches `uri`.
230    ///
231    /// This is the URI-keyed analog of [`contains`](Self::contains) and mirrors
232    /// what [`read`](Self::read) uses to select a resource.
233    #[must_use]
234    pub fn matches(&self, uri: &str) -> bool {
235        self.resources
236            .iter()
237            .any(|entry| entry.erased.read_erased(uri).is_some())
238    }
239
240    /// Borrow the cached [`ResourceDefinition`] for a registered resource by name.
241    ///
242    /// Returns `None` if no resource named `name` is registered.
243    #[must_use]
244    pub fn definition(&self, name: &str) -> Option<&ResourceDefinition> {
245        self.resources
246            .iter()
247            .find(|entry| entry.name == name)
248            .map(|entry| &entry.definition)
249    }
250
251    /// Iterate over `(name, definition)` pairs for every registered resource.
252    ///
253    /// Yields clones of the cached definitions computed at registration time.
254    #[must_use]
255    pub fn iter(&self) -> ResourceDefinitions<'_> {
256        ResourceDefinitions {
257            inner: self.resources.iter(),
258        }
259    }
260
261    /// Read the first resource whose URI template matches `uri`.
262    ///
263    /// # Errors
264    ///
265    /// Returns [`ToolError::not_found`] if no registered resource's template
266    /// matches `uri` (carrying `error_kind = "not_registered"` metadata), or a
267    /// read error if URI-variable deserialization or reading fails.
268    pub async fn read(&self, uri: &str) -> Result<ResourceOutput, ToolError> {
269        for resource in &self.resources {
270            if let Some(fut) = resource.erased.read_erased(uri) {
271                return fut.await;
272            }
273        }
274        Err(ToolError::not_found(RegistryItem::Resource, uri))
275    }
276}
277
278/// Borrowing iterator over `(name, definition)` pairs, yielded by
279/// [`ResourceRegistry::iter`] and by `&ResourceRegistry`'s [`IntoIterator`] impl.
280///
281/// Each cached [`ResourceDefinition`] is cloned lazily as it is yielded.
282pub struct ResourceDefinitions<'a> {
283    inner: core::slice::Iter<'a, RegisteredResource>,
284}
285
286impl Iterator for ResourceDefinitions<'_> {
287    type Item = (&'static str, ResourceDefinition);
288
289    fn next(&mut self) -> Option<Self::Item> {
290        self.inner
291            .next()
292            .map(|entry| (entry.name, entry.definition.clone()))
293    }
294
295    fn size_hint(&self) -> (usize, Option<usize>) {
296        self.inner.size_hint()
297    }
298}
299
300impl ExactSizeIterator for ResourceDefinitions<'_> {
301    fn len(&self) -> usize {
302        self.inner.len()
303    }
304}
305
306/// Iterate over `(name, definition)` pairs for every registered resource.
307impl<'a> IntoIterator for &'a ResourceRegistry {
308    type Item = (&'static str, ResourceDefinition);
309    type IntoIter = ResourceDefinitions<'a>;
310
311    fn into_iter(self) -> Self::IntoIter {
312        self.iter()
313    }
314}
315
316#[cfg(all(test, feature = "std"))]
317mod tests {
318    use super::match_uri_template;
319
320    #[test]
321    fn exact_match_no_variables() {
322        let m = match_uri_template("config://app", "config://app").expect("should match");
323        assert!(m.is_empty());
324    }
325
326    #[test]
327    fn no_match_different_literal() {
328        assert!(match_uri_template("config://app", "config://other").is_none());
329    }
330
331    #[test]
332    fn single_trailing_variable_captures_rest() {
333        let m = match_uri_template("file:///{path}", "file:///etc/hosts").expect("should match");
334        assert_eq!(m.get("path").map(String::as_str), Some("etc/hosts"));
335    }
336
337    #[test]
338    fn multiple_variables() {
339        let m = match_uri_template(
340            "file:///logs/{date}/{app}.log",
341            "file:///logs/2024-01-01/server.log",
342        )
343        .expect("should match");
344        assert_eq!(m.get("date").map(String::as_str), Some("2024-01-01"));
345        assert_eq!(m.get("app").map(String::as_str), Some("server"));
346    }
347
348    #[test]
349    fn variable_stops_at_delimiter() {
350        let m = match_uri_template("x://{a}/{b}", "x://one/two").expect("should match");
351        assert_eq!(m.get("a").map(String::as_str), Some("one"));
352        assert_eq!(m.get("b").map(String::as_str), Some("two"));
353    }
354
355    #[test]
356    fn prefix_mismatch_returns_none() {
357        assert!(match_uri_template("x://{a}", "y://foo").is_none());
358    }
359
360    #[test]
361    fn unterminated_template_variable_returns_none() {
362        // '{' with no closing '}' cannot be parsed into a variable.
363        assert!(match_uri_template("x://{a", "x://foo").is_none());
364    }
365
366    #[test]
367    fn missing_delimiter_in_uri_returns_none() {
368        // Template expects '/' after {a}, but the URI has none.
369        assert!(match_uri_template("x://{a}/end", "x://noslash").is_none());
370    }
371
372    #[test]
373    fn trailing_literal_must_match() {
374        // {a} captures up to '.', then ".log" must match the remainder.
375        assert!(match_uri_template("f://{a}.log", "f://name.txt").is_none());
376    }
377
378    #[test]
379    fn empty_variable_value_is_allowed() {
380        let m = match_uri_template("a://{x}/b", "a:///b").expect("should match");
381        assert_eq!(m.get("x").map(String::as_str), Some(""));
382    }
383
384    #[test]
385    fn longer_uri_than_literal_template_returns_none() {
386        assert!(match_uri_template("a://b", "a://bc").is_none());
387    }
388}