llm_tool/
rust_resource.rs1use 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
14pub trait RustResource: Send + Sync {
16 type Params: serde::de::DeserializeOwned + Send;
18
19 const URI_TEMPLATE: &'static str;
21
22 const NAME: &'static str;
24
25 const DESCRIPTION: &'static str;
27
28 const MIME_TYPE: Option<&'static str>;
30
31 fn description(&self) -> Cow<'static, str> {
33 Cow::Borrowed(Self::DESCRIPTION)
34 }
35
36 fn read(
38 &self,
39 uri: &str,
40 params: Self::Params,
41 ) -> impl Future<Output = Result<ResourceOutput, ToolError>> + Send;
42}
43
44#[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#[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
100pub(crate) type BoxResourceFuture<'a> =
102 Pin<Box<dyn Future<Output = Result<ResourceOutput, ToolError>> + Send + 'a>>;
103
104pub(crate) trait ErasedResource: Send + Sync {
109 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
135struct RegisteredResource {
137 name: &'static str,
138 definition: ResourceDefinition,
139 erased: Box<dyn ErasedResource>,
140}
141
142#[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 #[must_use]
173 pub const fn new() -> Self {
174 Self {
175 resources: Vec::new(),
176 }
177 }
178
179 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 #[must_use]
191 pub fn with_resource<R: RustResource + 'static>(mut self, resource: R) -> Self {
192 self.register(resource);
193 self
194 }
195
196 #[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 #[must_use]
209 pub const fn len(&self) -> usize {
210 self.resources.len()
211 }
212
213 #[must_use]
215 pub const fn is_empty(&self) -> bool {
216 self.resources.is_empty()
217 }
218
219 #[must_use]
225 pub fn contains(&self, name: &str) -> bool {
226 self.resources.iter().any(|entry| entry.name == name)
227 }
228
229 #[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 #[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 #[must_use]
255 pub fn iter(&self) -> ResourceDefinitions<'_> {
256 ResourceDefinitions {
257 inner: self.resources.iter(),
258 }
259 }
260
261 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
278pub 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
306impl<'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 assert!(match_uri_template("x://{a", "x://foo").is_none());
364 }
365
366 #[test]
367 fn missing_delimiter_in_uri_returns_none() {
368 assert!(match_uri_template("x://{a}/end", "x://noslash").is_none());
370 }
371
372 #[test]
373 fn trailing_literal_must_match() {
374 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}