Skip to main content

lash_core/tool_registry/
state.rs

1pub const PLUGIN_TOOL_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
19fn is_member_default() -> bool {
20    true
21}
22
23fn is_default_member(member: &bool) -> bool {
24    *member
25}
26
27#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
28pub struct ToolStateEntry {
29    manifest: ToolManifest,
30    /// True when this tool was not resolvable from any registered source at
31    /// export time (e.g. a detached MCP server). Orphaned entries keep their
32    /// last-known manifest, are excluded from the Tool Catalog (non-members
33    /// until their source returns), and rebind automatically when a source
34    /// re-advertises the same tool id.
35    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
36    orphaned: bool,
37    /// Catalog membership. Members are callable; non-members do not exist to
38    /// the model. Hosts toggle this via `set_tool_membership`.
39    #[serde(default = "is_member_default", skip_serializing_if = "is_default_member")]
40    member: bool,
41}
42
43impl ToolStateEntry {
44    #[cfg(test)]
45    pub(crate) fn new(manifest: ToolManifest) -> Self {
46        Self {
47            manifest,
48            orphaned: false,
49            member: true,
50        }
51    }
52
53    /// The stored manifest as exposed to callers.
54    pub fn manifest(&self) -> ToolManifest {
55        self.manifest.clone()
56    }
57
58    fn stored_manifest(&self) -> &ToolManifest {
59        &self.manifest
60    }
61
62    pub fn is_orphaned(&self) -> bool {
63        self.orphaned
64    }
65
66    /// Whether this entry is currently a Tool Catalog member. Orphaned entries
67    /// are never members.
68    pub fn is_member(&self) -> bool {
69        self.member && !self.orphaned
70    }
71}
72
73#[derive(Clone, Debug, Default)]
74pub struct ToolState {
75    generation: u64,
76    tools: Arc<BTreeMap<ToolId, ToolStateEntry>>,
77}
78
79impl ToolState {
80    pub(crate) fn new(generation: u64, tools: BTreeMap<ToolId, ToolStateEntry>) -> Self {
81        Self {
82            generation,
83            tools: Arc::new(tools),
84        }
85    }
86
87    pub fn generation(&self) -> u64 {
88        self.generation
89    }
90
91    #[cfg(any(test, feature = "testing"))]
92    pub(crate) fn with_generation(mut self, generation: u64) -> Self {
93        self.generation = generation;
94        self
95    }
96
97    /// Manifests for current Tool Catalog members. Orphaned and host-removed
98    /// entries are excluded (non-membership) but kept in state for rebind.
99    pub fn tool_manifests(&self) -> Vec<ToolManifest> {
100        self.tools
101            .values()
102            .filter(|entry| entry.is_member())
103            .map(ToolStateEntry::manifest)
104            .collect()
105    }
106
107    pub fn get(&self, id: &ToolId) -> Option<&ToolStateEntry> {
108        self.tools.get(id)
109    }
110
111    /// Edit a manifest in an explicit [`ToolRegistry::apply_state`] delta.
112    ///
113    /// Automatic rebuilds replace stored manifests with their live versions,
114    /// so this is not a persistent source-curation mechanism.
115    pub fn manifest_mut(&mut self, id: &ToolId) -> Option<&mut ToolManifest> {
116        Arc::make_mut(&mut self.tools)
117            .get_mut(id)
118            .map(|entry| &mut entry.manifest)
119    }
120
121    pub fn contains(&self, id: &ToolId) -> bool {
122        self.tools.contains_key(id)
123    }
124
125    pub fn is_empty(&self) -> bool {
126        self.tools.is_empty()
127    }
128
129    pub fn len(&self) -> usize {
130        self.tools.len()
131    }
132
133    pub fn iter(&self) -> impl Iterator<Item = (&ToolId, &ToolStateEntry)> {
134        self.tools.iter()
135    }
136
137    /// Toggle Tool Catalog membership for a tool. `present == false` removes
138    /// the tool from the catalog (non-membership) while keeping its state entry;
139    /// `present == true` restores membership.
140    pub fn set_membership(&mut self, id: &ToolId, present: bool) -> Result<(), ReconfigureError> {
141        let Some(entry) = Arc::make_mut(&mut self.tools).get_mut(id) else {
142            return Err(ReconfigureError::Validation(format!(
143                "unknown tool id `{id}`"
144            )));
145        };
146        entry.member = present;
147        Ok(())
148    }
149
150    /// Delete a tool in an explicit [`ToolRegistry::apply_state`] delta.
151    ///
152    /// Deletion intentionally removes the entry for that delta only. Use
153    /// [`Self::set_membership`] for curation that must survive a rebuild from
154    /// live sources.
155    pub fn remove(&mut self, id: &ToolId) -> Option<ToolStateEntry> {
156        Arc::make_mut(&mut self.tools).remove(id)
157    }
158
159    pub(crate) fn entries(&self) -> &BTreeMap<ToolId, ToolStateEntry> {
160        self.tools.as_ref()
161    }
162}
163
164impl Serialize for ToolState {
165    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
166    where
167        S: serde::Serializer,
168    {
169        #[derive(Serialize)]
170        struct ToolStateRef<'a> {
171            generation: u64,
172            tools: &'a BTreeMap<ToolId, ToolStateEntry>,
173        }
174
175        ToolStateRef {
176            generation: self.generation,
177            tools: self.tools.as_ref(),
178        }
179        .serialize(serializer)
180    }
181}
182
183impl<'de> Deserialize<'de> for ToolState {
184    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
185    where
186        D: serde::Deserializer<'de>,
187    {
188        #[derive(Deserialize)]
189        struct ToolStateOwned {
190            generation: u64,
191            tools: BTreeMap<ToolId, ToolStateEntry>,
192        }
193
194        let owned = ToolStateOwned::deserialize(deserializer)?;
195        Ok(Self {
196            generation: owned.generation,
197            tools: Arc::new(owned.tools),
198        })
199    }
200}
201
202#[async_trait::async_trait]
203pub(crate) trait ToolSourceExecutor: Send + Sync + 'static {
204    fn id(&self) -> &str;
205    fn advertised_tools(&self) -> Vec<ToolManifest>;
206    fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
207        self.advertised_tools()
208            .into_iter()
209            .find(|manifest| manifest.name == name)
210    }
211    fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
212        self.advertised_tools()
213            .into_iter()
214            .find(|manifest| manifest.id == *id)
215    }
216    fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
217    fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
218        let manifest = self.resolve_manifest_by_id(id)?;
219        self.resolve_contract(&manifest.name)
220    }
221    async fn prepare_tool_call(
222        &self,
223        call: ToolPrepareCall<'_>,
224    ) -> Result<PreparedToolCall, ToolResult> {
225        Ok(PreparedToolCall::identity(call.tool_id, call.pending))
226    }
227    async fn execute(
228        &self,
229        tool: &str,
230        args: &serde_json::Value,
231        context: &ToolContext<'_>,
232        progress: Option<&ProgressSender>,
233    ) -> ToolResult;
234    async fn execute_by_id(
235        &self,
236        tool_id: &ToolId,
237        args: &serde_json::Value,
238        context: &ToolContext<'_>,
239        progress: Option<&ProgressSender>,
240    ) -> ToolResult {
241        let Some(manifest) = self.resolve_manifest_by_id(tool_id) else {
242            return ToolResult::err_fmt(format_args!("Unknown tool id: {tool_id}"));
243        };
244        self.execute(&manifest.name, args, context, progress).await
245    }
246}