atomr_agents_tool/
registry.rs

1use std::sync::Arc;
2
3use atomr_agents_core::ToolSetId;
4use dashmap::DashMap;
5use semver::Version;
6
7use crate::toolset::ToolSet;
8
9/// Thread-safe registry of versioned tool sets. Multiple versions of
10/// the same id can coexist; consumers pin the version they want.
11#[derive(Default, Clone)]
12pub struct ToolSetRegistry {
13    inner: Arc<DashMap<(ToolSetId, Version), Arc<ToolSet>>>,
14}
15
16impl ToolSetRegistry {
17    pub fn new() -> Self {
18        Self::default()
19    }
20
21    pub fn publish(&self, ts: ToolSet) -> Arc<ToolSet> {
22        let key = (ts.id.clone(), ts.version.clone());
23        let arc = Arc::new(ts);
24        self.inner.insert(key, arc.clone());
25        arc
26    }
27
28    pub fn get(&self, id: &ToolSetId, version: &Version) -> Option<Arc<ToolSet>> {
29        self.inner
30            .get(&(id.clone(), version.clone()))
31            .map(|r| r.value().clone())
32    }
33
34    /// Highest version of `id`, or `None` if none published.
35    pub fn latest(&self, id: &ToolSetId) -> Option<Arc<ToolSet>> {
36        self.inner
37            .iter()
38            .filter(|r| r.key().0.as_str() == id.as_str())
39            .map(|r| r.value().clone())
40            .max_by(|a, b| a.version.cmp(&b.version))
41    }
42}
43
44#[cfg(test)]
45mod tests {
46    use super::*;
47    use crate::descriptor::{ToolDescriptor, ToolSchema};
48    use crate::r#trait::{DynTool, Tool};
49    use async_trait::async_trait;
50    use atomr_agents_core::{InvokeCtx, Result, ToolId, Value};
51
52    struct EchoTool {
53        descriptor: ToolDescriptor,
54    }
55    impl EchoTool {
56        fn new() -> Self {
57            Self {
58                descriptor: ToolDescriptor {
59                    id: ToolId::from("echo"),
60                    name: "echo".into(),
61                    description: "echo input".into(),
62                    schema: ToolSchema::empty_object(),
63                },
64            }
65        }
66    }
67    #[async_trait]
68    impl Tool for EchoTool {
69        fn descriptor(&self) -> &ToolDescriptor {
70            &self.descriptor
71        }
72        async fn invoke(&self, args: Value, _ctx: &InvokeCtx) -> Result<Value> {
73            Ok(args)
74        }
75    }
76
77    #[test]
78    fn publish_get_latest() {
79        let r = ToolSetRegistry::new();
80        let tools: Vec<DynTool> = vec![Arc::new(EchoTool::new())];
81        let ts1 = ToolSet::new("echos", Version::new(0, 1, 0), tools.clone());
82        let ts2 = ToolSet::new("echos", Version::new(0, 2, 0), tools);
83        r.publish(ts1);
84        r.publish(ts2);
85        let latest = r.latest(&ToolSetId::from("echos")).unwrap();
86        assert_eq!(latest.version, Version::new(0, 2, 0));
87        let pinned = r.get(&ToolSetId::from("echos"), &Version::new(0, 1, 0)).unwrap();
88        assert_eq!(pinned.version, Version::new(0, 1, 0));
89    }
90}
atomr_agents_tool/registry.rs

atomr_agents_tool/
registry.rs
