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, 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 #[allow(dead_code)]
59 fn stored_manifest(&self) -> &ToolManifest {
60 &self.manifest
61 }
62
63 pub fn is_orphaned(&self) -> bool {
64 self.orphaned
65 }
66
67 pub fn is_member(&self) -> bool {
70 self.member && !self.orphaned
71 }
72}
73
74#[derive(Clone, Debug, Default)]
75pub struct ToolState {
76 generation: u64,
77 tools: Arc<BTreeMap<ToolId, ToolStateEntry>>,
78}
79
80impl ToolState {
81 pub(crate) fn new(generation: u64, tools: BTreeMap<ToolId, ToolStateEntry>) -> Self {
82 Self {
83 generation,
84 tools: Arc::new(tools),
85 }
86 }
87
88 pub fn generation(&self) -> u64 {
89 self.generation
90 }
91
92 pub fn with_generation(mut self, generation: u64) -> Self {
93 self.generation = generation;
94 self
95 }
96
97 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 pub fn manifest_mut(&mut self, id: &ToolId) -> Option<&mut ToolManifest> {
112 Arc::make_mut(&mut self.tools)
113 .get_mut(id)
114 .map(|entry| &mut entry.manifest)
115 }
116
117 pub fn contains(&self, id: &ToolId) -> bool {
118 self.tools.contains_key(id)
119 }
120
121 pub fn is_empty(&self) -> bool {
122 self.tools.is_empty()
123 }
124
125 pub fn len(&self) -> usize {
126 self.tools.len()
127 }
128
129 pub fn iter(&self) -> impl Iterator<Item = (&ToolId, &ToolStateEntry)> {
130 self.tools.iter()
131 }
132
133 pub fn set_membership(&mut self, id: &ToolId, present: bool) -> Result<(), ReconfigureError> {
137 let Some(entry) = Arc::make_mut(&mut self.tools).get_mut(id) else {
138 return Err(ReconfigureError::Validation(format!(
139 "unknown tool id `{id}`"
140 )));
141 };
142 entry.member = present;
143 Ok(())
144 }
145
146 pub fn retain(&mut self, mut keep: impl FnMut(&ToolId, &ToolStateEntry) -> bool) {
147 Arc::make_mut(&mut self.tools).retain(|id, entry| keep(id, entry));
148 }
149
150 pub fn remove(&mut self, id: &ToolId) -> Option<ToolStateEntry> {
151 Arc::make_mut(&mut self.tools).remove(id)
152 }
153
154 pub(crate) fn entries(&self) -> &BTreeMap<ToolId, ToolStateEntry> {
155 self.tools.as_ref()
156 }
157}
158
159impl Serialize for ToolState {
160 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
161 where
162 S: serde::Serializer,
163 {
164 #[derive(Serialize)]
165 struct ToolStateRef<'a> {
166 generation: u64,
167 tools: &'a BTreeMap<ToolId, ToolStateEntry>,
168 }
169
170 ToolStateRef {
171 generation: self.generation,
172 tools: self.tools.as_ref(),
173 }
174 .serialize(serializer)
175 }
176}
177
178impl<'de> Deserialize<'de> for ToolState {
179 fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
180 where
181 D: serde::Deserializer<'de>,
182 {
183 #[derive(Deserialize)]
184 struct ToolStateOwned {
185 generation: u64,
186 tools: BTreeMap<ToolId, ToolStateEntry>,
187 }
188
189 let owned = ToolStateOwned::deserialize(deserializer)?;
190 Ok(Self {
191 generation: owned.generation,
192 tools: Arc::new(owned.tools),
193 })
194 }
195}
196
197#[async_trait::async_trait]
198pub(crate) trait ToolSourceExecutor: Send + Sync + 'static {
199 fn id(&self) -> &str;
200 fn advertised_tools(&self) -> Vec<ToolManifest>;
201 fn resolve_manifest(&self, name: &str) -> Option<ToolManifest> {
202 self.advertised_tools()
203 .into_iter()
204 .find(|manifest| manifest.name == name)
205 }
206 fn resolve_manifest_by_id(&self, id: &ToolId) -> Option<ToolManifest> {
207 self.advertised_tools()
208 .into_iter()
209 .find(|manifest| manifest.id == *id)
210 }
211 fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>>;
212 fn resolve_contract_by_id(&self, id: &ToolId) -> Option<Arc<ToolContract>> {
213 let manifest = self.resolve_manifest_by_id(id)?;
214 self.resolve_contract(&manifest.name)
215 }
216 async fn prepare_tool_call(
217 &self,
218 call: ToolPrepareCall<'_>,
219 ) -> Result<PreparedToolCall, ToolResult> {
220 Ok(PreparedToolCall::identity(call.tool_id, call.pending))
221 }
222 async fn execute(
223 &self,
224 tool: &str,
225 args: &serde_json::Value,
226 context: &ToolContext<'_>,
227 progress: Option<&ProgressSender>,
228 ) -> ToolResult;
229 async fn execute_by_id(
230 &self,
231 tool_id: &ToolId,
232 args: &serde_json::Value,
233 context: &ToolContext<'_>,
234 progress: Option<&ProgressSender>,
235 ) -> ToolResult {
236 let Some(manifest) = self.resolve_manifest_by_id(tool_id) else {
237 return ToolResult::err_fmt(format_args!("Unknown tool id: {tool_id}"));
238 };
239 self.execute(&manifest.name, args, context, progress).await
240 }
241}