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 tool 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<ToolId, ToolStateEntry>>,
63}
64
65impl ToolState {
66    pub(crate) fn new(generation: u64, tools: BTreeMap<ToolId, 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, id: &ToolId) -> Option<&ToolStateEntry> {
87        self.tools.get(id)
88    }
89
90    pub fn manifest_mut(&mut self, id: &ToolId) -> Option<&mut ToolManifest> {
91        Arc::make_mut(&mut self.tools)
92            .get_mut(id)
93            .map(|entry| &mut entry.manifest)
94    }
95
96    pub fn contains(&self, id: &ToolId) -> bool {
97        self.tools.contains_key(id)
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 = (&ToolId, &ToolStateEntry)> {
109        self.tools.iter()
110    }
111
112    pub fn set_availability(
113        &mut self,
114        id: &ToolId,
115        availability: Option<crate::ToolAvailability>,
116    ) -> Result<(), ReconfigureError> {
117        let Some(entry) = Arc::make_mut(&mut self.tools).get_mut(id) else {
118            return Err(ReconfigureError::Validation(format!(
119                "unknown tool id `{id}`"
120            )));
121        };
122        entry.manifest.availability_override = availability;
123        Ok(())
124    }
125
126    pub fn retain(&mut self, mut keep: impl FnMut(&ToolId, &ToolStateEntry) -> bool) {
127        Arc::make_mut(&mut self.tools).retain(|id, entry| keep(id, entry));
128    }
129
130    pub fn remove(&mut self, id: &ToolId) -> Option<ToolStateEntry> {
131        Arc::make_mut(&mut self.tools).remove(id)
132    }
133
134    pub(crate) fn entries(&self) -> &BTreeMap<ToolId, ToolStateEntry> {
135        self.tools.as_ref()
136    }
137}
138
139impl Serialize for ToolState {
140    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
141    where
142        S: serde::Serializer,
143    {
144        #[derive(Serialize)]
145        struct ToolStateRef<'a> {
146            generation: u64,
147            tools: &'a BTreeMap<ToolId, ToolStateEntry>,
148        }
149
150        ToolStateRef {
151            generation: self.generation,
152            tools: self.tools.as_ref(),
153        }
154        .serialize(serializer)
155    }
156}
157
158impl<'de> Deserialize<'de> for ToolState {
159    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
160    where
161        D: serde::Deserializer<'de>,
162    {
163        #[derive(Deserialize)]
164        struct ToolStateOwned {
165            generation: u64,
166            tools: BTreeMap<ToolId, ToolStateEntry>,
167        }
168
169        let owned = ToolStateOwned::deserialize(deserializer)?;
170        Ok(Self {
171            generation: owned.generation,
172            tools: Arc::new(owned.tools),
173        })
174    }
175}
176
177#[async_trait::async_trait]
178pub(crate) trait ToolSourceExecutor: Send + Sync + 'static {
179    fn id(&self) -> &str;
180    fn advertised_tools(&self) -> Vec<ToolManifest>;
181    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
182        self.advertised_tools()
183            .into_iter()
184            .find(|manifest| manifest.name == name)
185    }
186    fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
187        self.advertised_tools()
188            .into_iter()
189            .find(|manifest| manifest.id == *id)
190    }
191    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
192    fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
193        let manifest = self.resolve_manifest_by_id(id)?;
194        self.resolve_contract(&manifest.name)
195    }
196    async fn prepare_tool_call(
197        &self,
198        call: ToolPrepareCall<'_>,
199    ) -> Result<PreparedToolCall, ToolResult> {
200        Ok(PreparedToolCall::identity(call.tool_id, call.pending))
201    }
202    async fn execute(
203        &self,
204        tool: &str,
205        args: &serde_json::Value,
206        context: &ToolContext<'_>,
207        progress: Option<&ProgressSender>,
208    ) -> ToolResult;
209}
lash_core/tool_registry/state.rs

lash_core/tool_registry/
state.rs
