lash_core/tool_registry/
state.rs

1const PLUGIN_SOURCE_ID: &str = "plugins";
2
3#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
4#[serde(transparent)]
5pub struct ToolSourceHandle {
6    id: String,
7}
8
9impl ToolSourceHandle {
10    pub(crate) fn new(id: impl Into<String>) -> Self {
11        Self { id: id.into() }
12    }
13
14    pub(crate) fn as_str(&self) -> &str {
15        &self.id
16    }
17}
18
19#[derive(Clone, Debug, Serialize, Deserialize)]
20pub struct ToolStateEntry {
21    manifest: ToolManifest,
22    /// True when this tool was not resolvable from any registered source at
23    /// export time (e.g. a detached MCP server). Orphaned entries keep their
24    /// last-known manifest, surface as [`crate::ToolAvailability::Off`], and
25    /// rebind automatically when a source re-advertises the same (name, id).
26    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
27    orphaned: bool,
28}
29
30impl ToolStateEntry {
31    #[cfg(test)]
32    pub(crate) fn new(manifest: ToolManifest) -> Self {
33        Self {
34            manifest,
35            orphaned: false,
36        }
37    }
38
39    /// The manifest as exposed to callers. Orphaned entries are known to the
40    /// session but unavailable until their source returns, so their public
41    /// availability is always forced to `Off`.
42    pub fn manifest(&self) -> ToolManifest {
43        let mut manifest = self.manifest.clone();
44        if self.orphaned {
45            manifest.availability_override = Some(crate::ToolAvailability::Off);
46        }
47        manifest
48    }
49
50    fn stored_manifest(&self) -> &ToolManifest {
51        &self.manifest
52    }
53
54    pub fn is_orphaned(&self) -> bool {
55        self.orphaned
56    }
57}
58
59#[derive(Clone, Debug, Default)]
60pub struct ToolState {
61    generation: u64,
62    tools: Arc<BTreeMap<String, ToolStateEntry>>,
63}
64
65impl ToolState {
66    pub(crate) fn new(generation: u64, tools: BTreeMap<String, ToolStateEntry>) -> Self {
67        Self {
68            generation,
69            tools: Arc::new(tools),
70        }
71    }
72
73    pub fn generation(&self) -> u64 {
74        self.generation
75    }
76
77    pub fn with_generation(mut self, generation: u64) -> Self {
78        self.generation = generation;
79        self
80    }
81
82    pub fn tool_manifests(&self) -> Vec<ToolManifest> {
83        self.tools.values().map(ToolStateEntry::manifest).collect()
84    }
85
86    pub fn get(&self, name: &str) -> Option<&ToolStateEntry> {
87        self.tools.get(name)
88    }
89
90    pub fn manifest_mut(&mut self, name: &str) -> Option<&mut ToolManifest> {
91        Arc::make_mut(&mut self.tools)
92            .get_mut(name)
93            .map(|entry| &mut entry.manifest)
94    }
95
96    pub fn contains(&self, name: &str) -> bool {
97        self.tools.contains_key(name)
98    }
99
100    pub fn is_empty(&self) -> bool {
101        self.tools.is_empty()
102    }
103
104    pub fn len(&self) -> usize {
105        self.tools.len()
106    }
107
108    pub fn iter(&self) -> impl Iterator<Item = (&str, &ToolStateEntry)> {
109        self.tools
110            .iter()
111            .map(|(name, entry)| (name.as_str(), entry))
112    }
113
114    pub fn set_availability(
115        &mut self,
116        name: &str,
117        availability: Option<crate::ToolAvailability>,
118    ) -> Result<(), ReconfigureError> {
119        let Some(entry) = Arc::make_mut(&mut self.tools).get_mut(name) else {
120            return Err(ReconfigureError::Validation(format!(
121                "unknown tool `{name}`"
122            )));
123        };
124        entry.manifest.availability_override = availability;
125        Ok(())
126    }
127
128    pub fn retain(&mut self, mut keep: impl FnMut(&str, &ToolStateEntry) -> bool) {
129        Arc::make_mut(&mut self.tools).retain(|name, entry| keep(name, entry));
130    }
131
132    pub fn remove(&mut self, name: &str) -> Option<ToolStateEntry> {
133        Arc::make_mut(&mut self.tools).remove(name)
134    }
135
136    pub(crate) fn entries(&self) -> &BTreeMap<String, ToolStateEntry> {
137        self.tools.as_ref()
138    }
139}
140
141impl Serialize for ToolState {
142    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
143    where
144        S: serde::Serializer,
145    {
146        #[derive(Serialize)]
147        struct ToolStateRef<'a> {
148            generation: u64,
149            tools: &'a BTreeMap<String, ToolStateEntry>,
150        }
151
152        ToolStateRef {
153            generation: self.generation,
154            tools: self.tools.as_ref(),
155        }
156        .serialize(serializer)
157    }
158}
159
160impl<'de> Deserialize<'de> for ToolState {
161    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
162    where
163        D: serde::Deserializer<'de>,
164    {
165        #[derive(Deserialize)]
166        struct ToolStateOwned {
167            generation: u64,
168            tools: BTreeMap<String, ToolStateEntry>,
169        }
170
171        let owned = ToolStateOwned::deserialize(deserializer)?;
172        Ok(Self {
173            generation: owned.generation,
174            tools: Arc::new(owned.tools),
175        })
176    }
177}
178
179#[async_trait::async_trait]
180pub(crate) trait ToolSourceExecutor: Send + Sync + 'static {
181    fn id(&self) -> &str;
182    fn advertised_tools(&self) -> Vec<ToolManifest>;
183    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
184        self.advertised_tools()
185            .into_iter()
186            .find(|manifest| manifest.name == name)
187    }
188    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
189    async fn prepare_tool_call(
190        &self,
191        call: ToolPrepareCall<'_>,
192    ) -> Result<PreparedToolCall, ToolResult> {
193        Ok(PreparedToolCall::identity(call.pending))
194    }
195    async fn execute(
196        &self,
197        tool: &str,
198        args: &serde_json::Value,
199        context: &ToolContext<'_>,
200        progress: Option<&ProgressSender>,
201    ) -> ToolResult;
202}
lash_core/tool_registry/state.rs

lash_core/tool_registry/
state.rs
