lash_core/tool_registry/
state.rs1pub 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 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
36 orphaned: bool,
37 #[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 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 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 pub fn with_generation(mut self, generation: u64) -> Self {
92 self.generation = generation;
93 self
94 }
95
96 pub fn tool_manifests(&self) -> Vec<ToolManifest> {
99 self.tools
100 .values()
101 .filter(|entry| entry.is_member())
102 .map(ToolStateEntry::manifest)
103 .collect()
104 }
105
106 pub fn get(&self, id: &ToolId) -> Option<&ToolStateEntry> {
107 self.tools.get(id)
108 }
109
110 pub fn manifest_mut(&mut self, id: &ToolId) -> Option<&mut ToolManifest> {
115 Arc::make_mut(&mut self.tools)
116 .get_mut(id)
117 .map(|entry| &mut entry.manifest)
118 }
119
120 pub fn contains(&self, id: &ToolId) -> bool {
121 self.tools.contains_key(id)
122 }
123
124 pub fn is_empty(&self) -> bool {
125 self.tools.is_empty()
126 }
127
128 pub fn len(&self) -> usize {
129 self.tools.len()
130 }
131
132 pub fn iter(&self) -> impl Iterator<Item = (&ToolId, &ToolStateEntry)> {
133 self.tools.iter()
134 }
135
136 pub fn set_membership(&mut self, id: &ToolId, present: bool) -> Result<(), ReconfigureError> {
140 let Some(entry) = Arc::make_mut(&mut self.tools).get_mut(id) else {
141 return Err(ReconfigureError::Validation(format!(
142 "unknown tool id `{id}`"
143 )));
144 };
145 entry.member = present;
146 Ok(())
147 }
148
149 pub fn remove(&mut self, id: &ToolId) -> Option<ToolStateEntry> {
155 Arc::make_mut(&mut self.tools).remove(id)
156 }
157
158 pub(crate) fn entries(&self) -> &BTreeMap<ToolId, ToolStateEntry> {
159 self.tools.as_ref()
160 }
161}
162
163impl Serialize for ToolState {
164 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
165 where
166 S: serde::Serializer,
167 {
168 #[derive(Serialize)]
169 struct ToolStateRef<'a> {
170 generation: u64,
171 tools: &'a BTreeMap<ToolId, ToolStateEntry>,
172 }
173
174 ToolStateRef {
175 generation: self.generation,
176 tools: self.tools.as_ref(),
177 }
178 .serialize(serializer)
179 }
180}
181
182impl<'de> Deserialize<'de> for ToolState {
183 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
184 where
185 D: serde::Deserializer<'de>,
186 {
187 #[derive(Deserialize)]
188 struct ToolStateOwned {
189 generation: u64,
190 tools: BTreeMap<ToolId, ToolStateEntry>,
191 }
192
193 let owned = ToolStateOwned::deserialize(deserializer)?;
194 Ok(Self {
195 generation: owned.generation,
196 tools: Arc::new(owned.tools),
197 })
198 }
199}
200
201#[async_trait::async_trait]
202pub(crate) trait ToolSourceExecutor: Send + Sync + 'static {
203 fn id(&self) -> &str;
204 fn advertised_tools(&self) -> Vec<ToolManifest>;
205 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
206 self.advertised_tools()
207 .into_iter()
208 .find(|manifest| manifest.name == name)
209 }
210 fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
211 self.advertised_tools()
212 .into_iter()
213 .find(|manifest| manifest.id == *id)
214 }
215 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
216 fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
217 let manifest = self.resolve_manifest_by_id(id)?;
218 self.resolve_contract(&manifest.name)
219 }
220 async fn prepare_tool_call(
221 &self,
222 call: ToolPrepareCall<'_>,
223 ) -> Result<PreparedToolCall, ToolResult> {
224 Ok(PreparedToolCall::identity(call.tool_id, call.pending))
225 }
226 async fn execute(
227 &self,
228 tool: &str,
229 args: &serde_json::Value,
230 context: &ToolContext<'_>,
231 progress: Option<&ProgressSender>,
232 ) -> ToolResult;
233 async fn execute_by_id(
234 &self,
235 tool_id: &ToolId,
236 args: &serde_json::Value,
237 context: &ToolContext<'_>,
238 progress: Option<&ProgressSender>,
239 ) -> ToolResult {
240 let Some(manifest) = self.resolve_manifest_by_id(tool_id) else {
241 return ToolResult::err_fmt(format_args!("Unknown tool id: {tool_id}"));
242 };
243 self.execute(&manifest.name, args, context, progress).await
244 }
245}