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_brace = t_rem.find('{').unwrap_or(t_rem.len());
88            let delimiter = &t_rem[..next_brace];
89            let val_end = u_rem.find(delimiter)?;
90            let val = &u_rem[..val_end];
91            u_rem = &u_rem[val_end..];
92            val
93        };
94
95        map.insert(var_name.to_string(), val_str.to_string());
96    }
97
98    if t_rem == u_rem { Some(map) } else { None }
99}
100
101/// Type-erased future returned by [`ErasedResource::read_erased`].
102pub(crate) type BoxResourceFuture<'a> =
103    Pin<Box<dyn Future<Output = Result<ResourceOutput, ToolError>> + Send + 'a>>;
104
105/// Type-erased wrapper enabling heterogeneous resource storage.
106///
107/// This is an internal implementation detail of [`ResourceRegistry`]; callers
108/// interact with resources through the registry rather than this trait.
109pub(crate) trait ErasedResource: Send + Sync {
110    /// Check if the incoming URI matches this resource's pattern, extract
111    /// variables, and execute `read`.
112    ///
113    /// Returns `None` if `uri` does not match this resource's pattern.
114    fn read_erased<'a>(&'a self, uri: &'a str) -> Option<BoxResourceFuture<'a>>;
115}
116
117impl<T: RustResource> ErasedResource for T {
118    fn read_erased<'a>(&'a self, uri: &'a str) -> Option<BoxResourceFuture<'a>> {
119        let params_map = match_uri_template(T::URI_TEMPLATE, uri)?;
120        Some(Box::pin(async move {
121            let deserializer = serde::de::value::MapDeserializer::new(
122                params_map
123                    .into_iter()
124                    .map(|(k, v)| (k, serde::de::value::StringDeserializer::new(v))),
125            );
126            let params: T::Params = serde::de::Deserialize::deserialize(deserializer).map_err(
127                |e: serde::de::value::Error| {
128                    ToolError::new(format!(
129                        "Failed to deserialize resource parameters from URI variables: {e}"
130                    ))
131                },
132            )?;
133            self.read(uri, params).await
134        }))
135    }
136}
137
138/// A registered resource: its cached definition plus the type-erased handler.
139struct RegisteredResource {
140    name: &'static str,
141    definition: ResourceDefinition,
142    erased: Box<dyn ErasedResource>,
143}
144
145/// A registry of resources and resource templates for dynamic dispatch.
146///
147/// Models MCP's `resources/list`, `resources/read`, and `resources/templates/list`
148/// primitives. Tools provide actionable commands; resources provide readable
149/// context (static documents, log files, configuration snapshots).
150///
151/// # Example
152///
153/// ```
154/// use llm_tool::{ResourceOutput, ResourceRegistry, RustResource, ToolError};
155///
156/// struct ConfigResource;
157///
158/// impl RustResource for ConfigResource {
159///     const NAME: &'static str = "config";
160///     const URI_TEMPLATE: &'static str = "file:///etc/app.conf";
161///     const DESCRIPTION: &'static str = "Application configuration";
162///     const MIME_TYPE: Option<&'static str> = Some("text/plain");
163///     type Params = ();
164///
165///     async fn read(
166///         &self,
167///         uri: &str,
168///         _params: Self::Params,
169///     ) -> Result<ResourceOutput, ToolError> {
170///         Ok(ResourceOutput::text(uri, Some("text/plain"), "debug=true"))
171///     }
172/// }
173///
174/// let mut reg = ResourceRegistry::new();
175/// reg.register(ConfigResource);
176/// assert_eq!(reg.len(), 1);
177/// assert!(reg.matches("file:///etc/app.conf"));
178/// ```
179#[derive(Default)]
180pub struct ResourceRegistry {
181    resources: Vec<RegisteredResource>,
182}
183
184impl core::fmt::Debug for ResourceRegistry {
185    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
186        let names: Vec<&str> = self.resources.iter().map(|r| r.name).collect();
187        f.debug_struct("ResourceRegistry")
188            .field("resource_count", &self.resources.len())
189            .field("resource_names", &names)
190            .finish()
191    }
192}
193
194impl ResourceRegistry {
195    /// Create a new, empty resource registry.
196    #[must_use]
197    pub const fn new() -> Self {
198        Self {
199            resources: Vec::new(),
200        }
201    }
202
203    /// Register a [`RustResource`].
204    ///
205    /// Replaces any existing registration with the same [`RustResource::NAME`].
206    pub fn register<R: RustResource + 'static>(&mut self, resource: R) -> &mut Self {
207        if let Some(pos) = self.resources.iter().position(|e| e.name == R::NAME) {
208            self.resources.remove(pos);
209        }
210        self.resources.push(RegisteredResource {
211            name: R::NAME,
212            definition: definition_of_resource(&resource),
213            erased: Box::new(resource),
214        });
215        self
216    }
217
218    /// Register a [`RustResource`], consuming and returning `Self` for chaining.
219    #[must_use]
220    pub fn with_resource<R: RustResource + 'static>(mut self, resource: R) -> Self {
221        self.register(resource);
222        self
223    }
224
225    /// Remove a resource by name, returning `true` if it was present.
226    pub fn remove(&mut self, name: &str) -> bool {
227        if let Some(pos) = self.resources.iter().position(|e| e.name == name) {
228            self.resources.remove(pos);
229            true
230        } else {
231            false
232        }
233    }
234
235    /// Clear all registered resources.
236    pub fn clear(&mut self) {
237        self.resources.clear();
238    }
239
240    /// Collect [`ResourceDefinition`]s for all registered resources.
241    ///
242    /// Returns clones of the cached definitions computed at registration time.
243    #[must_use]
244    pub fn definitions(&self) -> Vec<ResourceDefinition> {
245        self.resources
246            .iter()
247            .map(|entry| entry.definition.clone())
248            .collect()
249    }
250
251    /// Number of registered resources.
252    #[must_use]
253    pub const fn len(&self) -> usize {
254        self.resources.len()
255    }
256
257    /// Whether the registry has no registered resources.
258    #[must_use]
259    pub const fn is_empty(&self) -> bool {
260        self.resources.is_empty()
261    }
262
263    /// Whether a resource with the given name is registered.
264    ///
265    /// Note that resources are *read* by URI (see [`matches`](Self::matches)),
266    /// not by name; this checks the registered resource **name** for parity
267    /// with [`ToolRegistry::contains`](crate::ToolRegistry::contains).
268    #[must_use]
269    pub fn contains(&self, name: &str) -> bool {
270        self.resources.iter().any(|entry| entry.name == name)
271    }
272
273    /// Whether any registered resource's URI template matches `uri`.
274    ///
275    /// This is the URI-keyed analog of [`contains`](Self::contains) and mirrors
276    /// what [`read`](Self::read) uses to select a resource.
277    #[must_use]
278    pub fn matches(&self, uri: &str) -> bool {
279        self.resources
280            .iter()
281            .any(|entry| entry.erased.read_erased(uri).is_some())
282    }
283
284    /// Borrow the cached [`ResourceDefinition`] for a registered resource by name.
285    ///
286    /// Returns `None` if no resource named `name` is registered.
287    #[must_use]
288    pub fn definition(&self, name: &str) -> Option<&ResourceDefinition> {
289        self.resources
290            .iter()
291            .find(|entry| entry.name == name)
292            .map(|entry| &entry.definition)
293    }
294
295    /// Iterate over `(name, definition)` pairs for every registered resource.
296    ///
297    /// Yields clones of the cached definitions computed at registration time.
298    #[must_use]
299    pub fn iter(&self) -> ResourceDefinitions<'_> {
300        ResourceDefinitions {
301            inner: self.resources.iter(),
302        }
303    }
304
305    /// Read the first resource whose URI template matches `uri`.
306    ///
307    /// # Errors
308    ///
309    /// Returns [`ToolError::not_found`] if no registered resource's template
310    /// matches `uri` (carrying `error_kind = "not_registered"` metadata), or a
311    /// read error if URI-variable deserialization or reading fails.
312    pub async fn read(&self, uri: &str) -> Result<ResourceOutput, ToolError> {
313        for resource in &self.resources {
314            if let Some(fut) = resource.erased.read_erased(uri) {
315                return fut.await;
316            }
317        }
318        Err(ToolError::not_found(RegistryItem::Resource, uri))
319    }
320}
321
322/// Borrowing iterator over `(name, definition)` pairs, yielded by
323/// [`ResourceRegistry::iter`] and by `&ResourceRegistry`'s [`IntoIterator`] impl.
324///
325/// Each cached [`ResourceDefinition`] is cloned lazily as it is yielded.
326pub struct ResourceDefinitions<'a> {
327    inner: core::slice::Iter<'a, RegisteredResource>,
328}
329
330impl Iterator for ResourceDefinitions<'_> {
331    type Item = (&'static str, ResourceDefinition);
332
333    fn next(&mut self) -> Option<Self::Item> {
334        self.inner
335            .next()
336            .map(|entry| (entry.name, entry.definition.clone()))
337    }
338
339    fn size_hint(&self) -> (usize, Option<usize>) {
340        self.inner.size_hint()
341    }
342}
343
344impl ExactSizeIterator for ResourceDefinitions<'_> {
345    fn len(&self) -> usize {
346        self.inner.len()
347    }
348}
349
350/// Iterate over `(name, definition)` pairs for every registered resource.
351impl<'a> IntoIterator for &'a ResourceRegistry {
352    type Item = (&'static str, ResourceDefinition);
353    type IntoIter = ResourceDefinitions<'a>;
354
355    fn into_iter(self) -> Self::IntoIter {
356        self.iter()
357    }
358}
359
360#[cfg(all(test, feature = "std"))]
361mod tests {
362    use super::match_uri_template;
363
364    #[test]
365    fn exact_match_no_variables() {
366        let m = match_uri_template("config://app", "config://app").expect("should match");
367        assert!(m.is_empty());
368    }
369
370    #[test]
371    fn no_match_different_literal() {
372        assert!(match_uri_template("config://app", "config://other").is_none());
373    }
374
375    #[test]
376    fn single_trailing_variable_captures_rest() {
377        let m = match_uri_template("file:///{path}", "file:///etc/hosts").expect("should match");
378        assert_eq!(m.get("path").map(String::as_str), Some("etc/hosts"));
379    }
380
381    #[test]
382    fn multiple_variables() {
383        let m = match_uri_template(
384            "file:///logs/{date}/{app}.log",
385            "file:///logs/2024-01-01/server.log",
386        )
387        .expect("should match");
388        assert_eq!(m.get("date").map(String::as_str), Some("2024-01-01"));
389        assert_eq!(m.get("app").map(String::as_str), Some("server"));
390    }
391
392    #[test]
393    fn variable_stops_at_delimiter() {
394        let m = match_uri_template("x://{a}/{b}", "x://one/two").expect("should match");
395        assert_eq!(m.get("a").map(String::as_str), Some("one"));
396        assert_eq!(m.get("b").map(String::as_str), Some("two"));
397    }
398
399    #[test]
400    fn prefix_mismatch_returns_none() {
401        assert!(match_uri_template("x://{a}", "y://foo").is_none());
402    }
403
404    #[test]
405    fn unterminated_template_variable_returns_none() {
406        // '{' with no closing '}' cannot be parsed into a variable.
407        assert!(match_uri_template("x://{a", "x://foo").is_none());
408    }
409
410    #[test]
411    fn missing_delimiter_in_uri_returns_none() {
412        // Template expects '/' after {a}, but the URI has none.
413        assert!(match_uri_template("x://{a}/end", "x://noslash").is_none());
414    }
415
416    #[test]
417    fn trailing_literal_must_match() {
418        // {a} captures up to '.', then ".log" must match the remainder.
419        assert!(match_uri_template("f://{a}.log", "f://name.txt").is_none());
420    }
421
422    #[test]
423    fn empty_variable_value_is_allowed() {
424        let m = match_uri_template("a://{x}/b", "a:///b").expect("should match");
425        assert_eq!(m.get("x").map(String::as_str), Some(""));
426    }
427
428    #[test]
429    fn longer_uri_than_literal_template_returns_none() {
430        assert!(match_uri_template("a://b", "a://bc").is_none());
431    }
432}