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_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
101pub(crate) type BoxResourceFuture<'a> =
103 Pin<Box<dyn Future<Output = Result<ResourceOutput, ToolError>> + Send + 'a>>;
104
105pub(crate) trait ErasedResource: Send + Sync {
110 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
138struct RegisteredResource {
140 name: &'static str,
141 definition: ResourceDefinition,
142 erased: Box<dyn ErasedResource>,
143}
144
145#[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 #[must_use]
197 pub const fn new() -> Self {
198 Self {
199 resources: Vec::new(),
200 }
201 }
202
203 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 #[must_use]
220 pub fn with_resource<R: RustResource + 'static>(mut self, resource: R) -> Self {
221 self.register(resource);
222 self
223 }
224
225 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 pub fn clear(&mut self) {
237 self.resources.clear();
238 }
239
240 #[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 #[must_use]
253 pub const fn len(&self) -> usize {
254 self.resources.len()
255 }
256
257 #[must_use]
259 pub const fn is_empty(&self) -> bool {
260 self.resources.is_empty()
261 }
262
263 #[must_use]
269 pub fn contains(&self, name: &str) -> bool {
270 self.resources.iter().any(|entry| entry.name == name)
271 }
272
273 #[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 #[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 #[must_use]
299 pub fn iter(&self) -> ResourceDefinitions<'_> {
300 ResourceDefinitions {
301 inner: self.resources.iter(),
302 }
303 }
304
305 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
322pub 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
350impl<'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 assert!(match_uri_template("x://{a", "x://foo").is_none());
408 }
409
410 #[test]
411 fn missing_delimiter_in_uri_returns_none() {
412 assert!(match_uri_template("x://{a}/end", "x://noslash").is_none());
414 }
415
416 #[test]
417 fn trailing_literal_must_match() {
418 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}