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/// Opening delimiter for URI template variable placeholders.
58const TEMPLATE_VAR_START: char = '{';
59/// Closing delimiter for URI template variable placeholders.
60const TEMPLATE_VAR_END: char = '}';
61/// Empty string slice.
62const EMPTY_STR: &str = "";
63
64/// A single segment of a pre-compiled URI template.
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub enum TemplateSegment<'a> {
67    /// Literal substring that must match identically.
68    Literal(&'a str),
69    /// Named placeholder variable (e.g. `date` from `{date}`).
70    Variable(&'a str),
71}
72
73/// Pre-compiled representation of a URI template pattern.
74///
75/// Parses `{variable}` placeholders once at registration time so that
76/// [`matches`](Self::matches) and [`extract_borrowed`](Self::extract_borrowed)
77/// run in pure stack memory without re-scanning the template string or
78/// allocating heap `String`s.
79#[derive(Debug, Clone, PartialEq, Eq)]
80pub struct CompiledUriTemplate<'a> {
81    raw: &'a str,
82    segments: Vec<TemplateSegment<'a>>,
83    var_count: usize,
84}
85
86impl<'a> CompiledUriTemplate<'a> {
87    /// Compile a URI template string into its constituent segments.
88    #[must_use]
89    pub fn new(template: &'a str) -> Self {
90        if !template.contains(TEMPLATE_VAR_START) {
91            return Self {
92                raw: template,
93                segments: alloc::vec![TemplateSegment::Literal(template)],
94                var_count: 0,
95            };
96        }
97
98        let mut segments = Vec::new();
99        let mut rem = template;
100        let mut var_count = 0;
101
102        while let Some(start_idx) = rem.find(TEMPLATE_VAR_START) {
103            if start_idx > 0 {
104                segments.push(TemplateSegment::Literal(&rem[..start_idx]));
105            }
106            let after_start = &rem[start_idx + 1..];
107            if let Some(end_idx) = after_start.find(TEMPLATE_VAR_END) {
108                let var_name = &after_start[..end_idx];
109                segments.push(TemplateSegment::Variable(var_name));
110                var_count += 1;
111                rem = &after_start[end_idx + 1..];
112            } else {
113                segments.push(TemplateSegment::Literal(&rem[start_idx..]));
114                rem = EMPTY_STR;
115                break;
116            }
117        }
118        if !rem.is_empty() {
119            segments.push(TemplateSegment::Literal(rem));
120        }
121
122        Self {
123            raw: template,
124            segments,
125            var_count,
126        }
127    }
128
129    /// Return the original unparsed template string.
130    #[must_use]
131    pub const fn raw(&self) -> &'a str {
132        self.raw
133    }
134
135    /// Return the parsed template segments.
136    #[must_use]
137    pub fn segments(&self) -> &[TemplateSegment<'a>] {
138        &self.segments
139    }
140
141    /// Return the number of placeholder variables in this template.
142    #[must_use]
143    pub const fn var_count(&self) -> usize {
144        self.var_count
145    }
146
147    /// Check whether `uri` matches this template pattern without allocating heap strings.
148    #[must_use]
149    pub fn matches(&self, uri: &str) -> bool {
150        if self.var_count == 0 {
151            return self.raw == uri;
152        }
153        self.match_with_callback(uri, |_, _| {}).is_some()
154    }
155
156    /// Extract borrowed `(variable_name, value)` pairs from `uri`.
157    ///
158    /// Returns `None` if `uri` does not match this template.
159    #[must_use]
160    pub fn extract_borrowed<'u>(&self, uri: &'u str) -> Option<Vec<(&'a str, &'u str)>> {
161        if self.var_count == 0 {
162            return if self.raw == uri {
163                Some(Vec::new())
164            } else {
165                None
166            };
167        }
168        let mut vars = Vec::with_capacity(self.var_count);
169        self.match_with_callback(uri, |k, v| vars.push((k, v)))?;
170        Some(vars)
171    }
172
173    fn match_with_callback<'u>(
174        &self,
175        uri: &'u str,
176        mut on_var: impl FnMut(&'a str, &'u str),
177    ) -> Option<()> {
178        let mut u_rem = uri;
179        for (i, segment) in self.segments.iter().enumerate() {
180            match *segment {
181                TemplateSegment::Literal(lit) => {
182                    if !u_rem.starts_with(lit) {
183                        return None;
184                    }
185                    u_rem = &u_rem[lit.len()..];
186                }
187                TemplateSegment::Variable(var_name) => {
188                    let val_str = match self.segments.get(i + 1) {
189                        None => {
190                            let val = u_rem;
191                            u_rem = EMPTY_STR;
192                            val
193                        }
194                        Some(TemplateSegment::Literal(next_lit)) => {
195                            let val_end = u_rem.find(next_lit)?;
196                            let val = &u_rem[..val_end];
197                            u_rem = &u_rem[val_end..];
198                            val
199                        }
200                        Some(TemplateSegment::Variable(_)) => EMPTY_STR,
201                    };
202                    on_var(var_name, val_str);
203                }
204            }
205        }
206        if u_rem.is_empty() { Some(()) } else { None }
207    }
208}
209
210/// Helper to match an incoming URI against a URI template pattern with `{variable}` placeholders.
211///
212/// Returns `Some(map)` if the URI matches the pattern, mapping each `{variable}` name
213/// to its extracted value string. Returns `None` if the URI does not match.
214#[must_use]
215pub fn match_uri_template(
216    template: &str,
217    uri: &str,
218) -> Option<alloc::collections::BTreeMap<String, String>> {
219    let compiled = CompiledUriTemplate::new(template);
220    let mut map = alloc::collections::BTreeMap::new();
221    compiled.match_with_callback(uri, |k, v| {
222        map.insert(k.to_string(), v.to_string());
223    })?;
224    Some(map)
225}
226
227/// Type-erased future returned by [`ErasedResource::read_with_vars`].
228pub(crate) type BoxResourceFuture<'a> =
229    Pin<Box<dyn Future<Output = Result<ResourceOutput, ToolError>> + Send + 'a>>;
230
231/// Type-erased wrapper enabling heterogeneous resource storage.
232///
233/// This is an internal implementation detail of [`ResourceRegistry`]; callers
234/// interact with resources through the registry rather than this trait.
235pub(crate) trait ErasedResource: Send + Sync {
236    /// Execute `read` given pre-extracted borrowed `(key, value)` URI variables.
237    fn read_with_vars<'a>(
238        &'a self,
239        uri: &'a str,
240        vars: Vec<(&'static str, &'a str)>,
241    ) -> BoxResourceFuture<'a>;
242}
243
244impl<T: RustResource> ErasedResource for T {
245    fn read_with_vars<'a>(
246        &'a self,
247        uri: &'a str,
248        vars: Vec<(&'static str, &'a str)>,
249    ) -> BoxResourceFuture<'a> {
250        Box::pin(async move {
251            let deserializer = serde::de::value::MapDeserializer::new(
252                vars.into_iter()
253                    .map(|(k, v)| (k, serde::de::value::BorrowedStrDeserializer::new(v))),
254            );
255            let params: T::Params = serde::de::Deserialize::deserialize(deserializer).map_err(
256                |e: serde::de::value::Error| {
257                    ToolError::new(format!(
258                        "Failed to deserialize resource parameters from URI variables: {e}"
259                    ))
260                },
261            )?;
262            self.read(uri, params).await
263        })
264    }
265}
266
267/// A registered resource: its cached definition, pre-compiled URI template,
268/// and type-erased handler.
269struct RegisteredResource {
270    name: &'static str,
271    definition: ResourceDefinition,
272    template: CompiledUriTemplate<'static>,
273    erased: Box<dyn ErasedResource>,
274}
275
276/// A registry of resources and resource templates for dynamic dispatch.
277///
278/// Models MCP's `resources/list`, `resources/read`, and `resources/templates/list`
279/// primitives. Tools provide actionable commands; resources provide readable
280/// context (static documents, log files, configuration snapshots).
281///
282/// # Example
283///
284/// ```
285/// use llm_tool::{ResourceOutput, ResourceRegistry, RustResource, ToolError};
286///
287/// struct ConfigResource;
288///
289/// impl RustResource for ConfigResource {
290///     const NAME: &'static str = "config";
291///     const URI_TEMPLATE: &'static str = "file:///etc/app.conf";
292///     const DESCRIPTION: &'static str = "Application configuration";
293///     const MIME_TYPE: Option<&'static str> = Some("text/plain");
294///     type Params = ();
295///
296///     async fn read(
297///         &self,
298///         uri: &str,
299///         _params: Self::Params,
300///     ) -> Result<ResourceOutput, ToolError> {
301///         Ok(ResourceOutput::text(uri, Some("text/plain"), "debug=true"))
302///     }
303/// }
304///
305/// let mut reg = ResourceRegistry::new();
306/// reg.register(ConfigResource);
307/// assert_eq!(reg.len(), 1);
308/// assert!(reg.matches("file:///etc/app.conf"));
309/// ```
310#[derive(Default)]
311pub struct ResourceRegistry {
312    resources: Vec<RegisteredResource>,
313}
314
315impl core::fmt::Debug for ResourceRegistry {
316    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
317        let names: Vec<&str> = self.resources.iter().map(|r| r.name).collect();
318        f.debug_struct("ResourceRegistry")
319            .field("resource_count", &self.resources.len())
320            .field("resource_names", &names)
321            .finish()
322    }
323}
324
325impl ResourceRegistry {
326    /// Create a new, empty resource registry.
327    #[must_use]
328    pub const fn new() -> Self {
329        Self {
330            resources: Vec::new(),
331        }
332    }
333
334    /// Register a [`RustResource`].
335    ///
336    /// Replaces any existing registration with the same [`RustResource::NAME`].
337    pub fn register<R: RustResource + 'static>(&mut self, resource: R) -> &mut Self {
338        if let Some(pos) = self.resources.iter().position(|e| e.name == R::NAME) {
339            self.resources.remove(pos);
340        }
341        self.resources.push(RegisteredResource {
342            name: R::NAME,
343            definition: definition_of_resource(&resource),
344            template: CompiledUriTemplate::new(R::URI_TEMPLATE),
345            erased: Box::new(resource),
346        });
347        self
348    }
349
350    /// Register a [`RustResource`], consuming and returning `Self` for chaining.
351    #[must_use]
352    pub fn with_resource<R: RustResource + 'static>(mut self, resource: R) -> Self {
353        self.register(resource);
354        self
355    }
356
357    /// Remove a resource by name, returning `true` if it was present.
358    pub fn remove(&mut self, name: &str) -> bool {
359        if let Some(pos) = self.resources.iter().position(|e| e.name == name) {
360            self.resources.remove(pos);
361            true
362        } else {
363            false
364        }
365    }
366
367    /// Clear all registered resources.
368    pub fn clear(&mut self) {
369        self.resources.clear();
370    }
371
372    /// Collect [`ResourceDefinition`]s for all registered resources.
373    ///
374    /// Returns clones of the cached definitions computed at registration time.
375    #[must_use]
376    pub fn definitions(&self) -> Vec<ResourceDefinition> {
377        self.resources
378            .iter()
379            .map(|entry| entry.definition.clone())
380            .collect()
381    }
382
383    /// Number of registered resources.
384    #[must_use]
385    pub const fn len(&self) -> usize {
386        self.resources.len()
387    }
388
389    /// Whether the registry has no registered resources.
390    #[must_use]
391    pub const fn is_empty(&self) -> bool {
392        self.resources.is_empty()
393    }
394
395    /// Whether a resource with the given name is registered.
396    ///
397    /// Note that resources are *read* by URI (see [`matches`](Self::matches)),
398    /// not by name; this checks the registered resource **name** for parity
399    /// with [`ToolRegistry::contains`](crate::ToolRegistry::contains).
400    #[must_use]
401    pub fn contains(&self, name: &str) -> bool {
402        self.resources.iter().any(|entry| entry.name == name)
403    }
404
405    /// Whether any registered resource's URI template matches `uri`.
406    ///
407    /// This is the URI-keyed analog of [`contains`](Self::contains) and mirrors
408    /// what [`read`](Self::read) uses to select a resource.
409    #[must_use]
410    pub fn matches(&self, uri: &str) -> bool {
411        self.resources
412            .iter()
413            .any(|entry| entry.template.matches(uri))
414    }
415
416    /// Borrow the cached [`ResourceDefinition`] for a registered resource by name.
417    ///
418    /// Returns `None` if no resource named `name` is registered.
419    #[must_use]
420    pub fn definition(&self, name: &str) -> Option<&ResourceDefinition> {
421        self.resources
422            .iter()
423            .find(|entry| entry.name == name)
424            .map(|entry| &entry.definition)
425    }
426
427    /// Iterate over `(name, definition)` pairs for every registered resource.
428    ///
429    /// Yields clones of the cached definitions computed at registration time.
430    #[must_use]
431    pub fn iter(&self) -> ResourceDefinitions<'_> {
432        ResourceDefinitions {
433            inner: self.resources.iter(),
434        }
435    }
436
437    /// Read the first resource whose URI template matches `uri`.
438    ///
439    /// # Errors
440    ///
441    /// Returns [`ToolError::not_found`] if no registered resource's template
442    /// matches `uri` (carrying `error_kind = "not_registered"` metadata), or a
443    /// read error if URI-variable deserialization or reading fails.
444    pub async fn read(&self, uri: &str) -> Result<ResourceOutput, ToolError> {
445        for resource in &self.resources {
446            if let Some(vars) = resource.template.extract_borrowed(uri) {
447                return resource.erased.read_with_vars(uri, vars).await;
448            }
449        }
450        Err(ToolError::not_found(RegistryItem::Resource, uri))
451    }
452}
453
454/// Borrowing iterator over `(name, definition)` pairs, yielded by
455/// [`ResourceRegistry::iter`] and by `&ResourceRegistry`'s [`IntoIterator`] impl.
456///
457/// Each cached [`ResourceDefinition`] is cloned lazily as it is yielded.
458pub struct ResourceDefinitions<'a> {
459    inner: core::slice::Iter<'a, RegisteredResource>,
460}
461
462impl Iterator for ResourceDefinitions<'_> {
463    type Item = (&'static str, ResourceDefinition);
464
465    fn next(&mut self) -> Option<Self::Item> {
466        self.inner
467            .next()
468            .map(|entry| (entry.name, entry.definition.clone()))
469    }
470
471    fn size_hint(&self) -> (usize, Option<usize>) {
472        self.inner.size_hint()
473    }
474}
475
476impl ExactSizeIterator for ResourceDefinitions<'_> {
477    fn len(&self) -> usize {
478        self.inner.len()
479    }
480}
481
482/// Iterate over `(name, definition)` pairs for every registered resource.
483impl<'a> IntoIterator for &'a ResourceRegistry {
484    type Item = (&'static str, ResourceDefinition);
485    type IntoIter = ResourceDefinitions<'a>;
486
487    fn into_iter(self) -> Self::IntoIter {
488        self.iter()
489    }
490}
491
492#[cfg(all(test, feature = "std"))]
493mod tests {
494    use super::match_uri_template;
495
496    #[test]
497    fn exact_match_no_variables() {
498        let m = match_uri_template("config://app", "config://app").expect("should match");
499        assert!(m.is_empty());
500    }
501
502    #[test]
503    fn no_match_different_literal() {
504        assert!(match_uri_template("config://app", "config://other").is_none());
505    }
506
507    #[test]
508    fn single_trailing_variable_captures_rest() {
509        let m = match_uri_template("file:///{path}", "file:///etc/hosts").expect("should match");
510        assert_eq!(m.get("path").map(String::as_str), Some("etc/hosts"));
511    }
512
513    #[test]
514    fn multiple_variables() {
515        let m = match_uri_template(
516            "file:///logs/{date}/{app}.log",
517            "file:///logs/2024-01-01/server.log",
518        )
519        .expect("should match");
520        assert_eq!(m.get("date").map(String::as_str), Some("2024-01-01"));
521        assert_eq!(m.get("app").map(String::as_str), Some("server"));
522    }
523
524    #[test]
525    fn variable_stops_at_delimiter() {
526        let m = match_uri_template("x://{a}/{b}", "x://one/two").expect("should match");
527        assert_eq!(m.get("a").map(String::as_str), Some("one"));
528        assert_eq!(m.get("b").map(String::as_str), Some("two"));
529    }
530
531    #[test]
532    fn prefix_mismatch_returns_none() {
533        assert!(match_uri_template("x://{a}", "y://foo").is_none());
534    }
535
536    #[test]
537    fn unterminated_template_variable_returns_none() {
538        // '{' with no closing '}' cannot be parsed into a variable.
539        assert!(match_uri_template("x://{a", "x://foo").is_none());
540    }
541
542    #[test]
543    fn missing_delimiter_in_uri_returns_none() {
544        // Template expects '/' after {a}, but the URI has none.
545        assert!(match_uri_template("x://{a}/end", "x://noslash").is_none());
546    }
547
548    #[test]
549    fn trailing_literal_must_match() {
550        // {a} captures up to '.', then ".log" must match the remainder.
551        assert!(match_uri_template("f://{a}.log", "f://name.txt").is_none());
552    }
553
554    #[test]
555    fn empty_variable_value_is_allowed() {
556        let m = match_uri_template("a://{x}/b", "a:///b").expect("should match");
557        assert_eq!(m.get("x").map(String::as_str), Some(""));
558    }
559
560    #[test]
561    fn longer_uri_than_literal_template_returns_none() {
562        assert!(match_uri_template("a://b", "a://bc").is_none());
563    }
564
565    #[test]
566    fn adjacent_template_variables_match() {
567        let m = match_uri_template("x://{a}{b}", "x://hello").expect("should match");
568        assert_eq!(m.get("a").map(String::as_str), Some(""));
569        assert_eq!(m.get("b").map(String::as_str), Some("hello"));
570    }
571
572    #[tokio::test]
573    async fn erased_resource_deserialization_failure_on_non_string_param() {
574        use super::*;
575
576        struct IntResource;
577        #[derive(serde::Deserialize)]
578        struct IntParams {
579            count: u64,
580        }
581        impl RustResource for IntResource {
582            const NAME: &'static str = "int_res";
583            const URI_TEMPLATE: &'static str = "res://{count}";
584            const DESCRIPTION: &'static str = "Resource with integer param";
585            const MIME_TYPE: Option<&'static str> = None;
586            type Params = IntParams;
587            async fn read(
588                &self,
589                uri: &str,
590                params: Self::Params,
591            ) -> Result<ResourceOutput, ToolError> {
592                Ok(ResourceOutput::text(uri, None, format!("{}", params.count)))
593            }
594        }
595
596        let reg = ResourceRegistry::new().with_resource(IntResource);
597        let err = reg.read("res://42").await.unwrap_err();
598        assert!(
599            err.to_string()
600                .contains("Failed to deserialize resource parameters from URI variables")
601        );
602    }
603
604    #[test]
605    fn compiled_uri_template_segments_and_accessors() {
606        use super::{CompiledUriTemplate, TemplateSegment};
607
608        let static_tmpl = CompiledUriTemplate::new("config://app");
609        assert_eq!(static_tmpl.raw(), "config://app");
610        assert_eq!(static_tmpl.var_count(), 0);
611        assert_eq!(
612            static_tmpl.segments(),
613            &[TemplateSegment::Literal("config://app")]
614        );
615        assert!(static_tmpl.matches("config://app"));
616        assert!(!static_tmpl.matches("config://other"));
617        assert_eq!(
618            static_tmpl.extract_borrowed("config://app"),
619            Some(alloc::vec::Vec::new())
620        );
621
622        let dyn_tmpl = CompiledUriTemplate::new("file:///logs/{date}/{app}.log");
623        assert_eq!(dyn_tmpl.raw(), "file:///logs/{date}/{app}.log");
624        assert_eq!(dyn_tmpl.var_count(), 2);
625        assert_eq!(
626            dyn_tmpl.segments(),
627            &[
628                TemplateSegment::Literal("file:///logs/"),
629                TemplateSegment::Variable("date"),
630                TemplateSegment::Literal("/"),
631                TemplateSegment::Variable("app"),
632                TemplateSegment::Literal(".log"),
633            ]
634        );
635    }
636
637    #[test]
638    fn compiled_uri_template_extract_borrowed_zero_copy() {
639        use super::CompiledUriTemplate;
640
641        let tmpl = CompiledUriTemplate::new("res://{host}:{port}/{path}");
642        let uri = "res://localhost:8080/api/v1";
643        let extracted = tmpl.extract_borrowed(uri).expect("should match");
644        assert_eq!(
645            extracted,
646            alloc::vec![("host", "localhost"), ("port", "8080"), ("path", "api/v1"),]
647        );
648        // Verify that extracted values point directly into the input `uri` slice (zero copy)
649        let uri_start = uri.as_ptr() as usize;
650        let uri_end = uri_start + uri.len();
651        for (_, val) in extracted {
652            let val_ptr = val.as_ptr() as usize;
653            assert!(val_ptr >= uri_start && val_ptr + val.len() <= uri_end);
654        }
655    }
656}
llm_tool/rust_resource.rs

llm_tool/
rust_resource.rs
