11use {
22 anyhow:: Result ,
33 async_trait:: async_trait,
4- std:: { collections:: HashMap , sync:: Arc } ,
4+ std:: {
5+ collections:: HashMap ,
6+ sync:: { Arc , Mutex } ,
7+ } ,
58} ;
69
710/// Agent-callable tool.
@@ -29,17 +32,27 @@ pub enum ToolSource {
2932}
3033
3134/// Internal entry pairing a tool with its source metadata.
32- struct ToolEntry {
33- tool : Arc < dyn AgentTool > ,
34- source : ToolSource ,
35+ pub ( crate ) struct ToolEntry {
36+ pub ( crate ) tool : Arc < dyn AgentTool > ,
37+ pub ( crate ) source : ToolSource ,
3538}
3639
40+ /// Shared set of tools activated at runtime by [`ToolSearchTool`](crate::lazy_tools::ToolSearchTool).
41+ ///
42+ /// Uses `std::sync::Mutex` (not tokio) because the lock is held for
43+ /// microseconds — just a `HashMap` insert/lookup — and this keeps
44+ /// `list_schemas()` usable from sync contexts.
45+ pub ( crate ) type ActivatedTools = Arc < Mutex < HashMap < String , ToolEntry > > > ;
46+
3747/// Registry of available tools for an agent run.
3848///
3949/// Tools are stored as `Arc<dyn AgentTool>` so the registry can be cheaply
4050/// cloned (e.g. for sub-agents that need a filtered copy of the parent's tools).
4151pub struct ToolRegistry {
4252 tools : HashMap < String , ToolEntry > ,
53+ /// Tools activated at runtime via lazy tool discovery (`tool_search`).
54+ /// Always present (empty when lazy mode is not in use).
55+ pub ( crate ) activated : ActivatedTools ,
4356}
4457
4558impl Default for ToolRegistry {
@@ -52,6 +65,7 @@ impl ToolRegistry {
5265 pub fn new ( ) -> Self {
5366 Self {
5467 tools : HashMap :: new ( ) ,
68+ activated : Arc :: new ( Mutex :: new ( HashMap :: new ( ) ) ) ,
5569 }
5670 }
5771
@@ -112,49 +126,47 @@ impl ToolRegistry {
112126 before - self . tools . len ( )
113127 }
114128
115- pub fn get ( & self , name : & str ) -> Option < & dyn AgentTool > {
116- self . tools . get ( name) . map ( |e| e. tool . as_ref ( ) )
129+ pub fn get ( & self , name : & str ) -> Option < Arc < dyn AgentTool > > {
130+ if let Some ( e) = self . tools . get ( name) {
131+ return Some ( Arc :: clone ( & e. tool ) ) ;
132+ }
133+ let activated = self . activated . lock ( ) . unwrap_or_else ( |e| e. into_inner ( ) ) ;
134+ activated. get ( name) . map ( |e| Arc :: clone ( & e. tool ) )
117135 }
118136
119- /// Return a cloned tool handle by name.
120- pub fn get_arc ( & self , name : & str ) -> Option < Arc < dyn AgentTool > > {
121- self . tools . get ( name) . map ( |e| Arc :: clone ( & e . tool ) )
137+ /// Return the [`ToolSource`] for a tool by name.
138+ pub ( crate ) fn get_source ( & self , name : & str ) -> Option < ToolSource > {
139+ self . tools . get ( name) . map ( |e| e . source . clone ( ) )
122140 }
123141
124142 pub fn list_schemas ( & self ) -> Vec < serde_json:: Value > {
125- self . tools
126- . values ( )
127- . map ( |e| {
128- let mut schema = serde_json:: json!( {
129- "name" : e. tool. name( ) ,
130- "description" : e. tool. description( ) ,
131- "parameters" : e. tool. parameters_schema( ) ,
132- } ) ;
133- match & e. source {
134- ToolSource :: Builtin => {
135- schema[ "source" ] = serde_json:: json!( "builtin" ) ;
136- } ,
137- ToolSource :: Mcp { server } => {
138- schema[ "source" ] = serde_json:: json!( "mcp" ) ;
139- schema[ "mcpServer" ] = serde_json:: json!( server) ;
140- } ,
141- ToolSource :: Wasm { component_hash } => {
142- schema[ "source" ] = serde_json:: json!( "wasm" ) ;
143- schema[ "componentHash" ] =
144- serde_json:: json!( hex_component_hash( * component_hash) ) ;
145- } ,
146- }
147- schema
148- } )
149- . collect ( )
143+ let mut schemas: Vec < serde_json:: Value > =
144+ self . tools . values ( ) . map ( entry_to_schema) . collect ( ) ;
145+
146+ let activated = self . activated . lock ( ) . unwrap_or_else ( |e| e. into_inner ( ) ) ;
147+ for ( name, entry) in activated. iter ( ) {
148+ if !self . tools . contains_key ( name) {
149+ schemas. push ( entry_to_schema ( entry) ) ;
150+ }
151+ }
152+ schemas
150153 }
151154
152- /// List registered tool names.
155+ /// List registered tool names (static + activated) .
153156pub fn list_names ( & self ) -> Vec < String > {
154- self . tools . keys ( ) . cloned ( ) . collect ( )
157+ let mut names: Vec < String > = self . tools . keys ( ) . cloned ( ) . collect ( ) ;
158+ let activated = self . activated . lock ( ) . unwrap_or_else ( |e| e. into_inner ( ) ) ;
159+ for name in activated. keys ( ) {
160+ if !self . tools . contains_key ( name) {
161+ names. push ( name. clone ( ) ) ;
162+ }
163+ }
164+ names
155165 }
156166
157167 /// Clone the registry, excluding tools whose names start with `prefix`.
168+ ///
169+ /// Sub-agent registries get a fresh (empty) activated set.
158170pub fn clone_without_prefix ( & self , prefix : & str ) -> ToolRegistry {
159171 let tools = self
160172 . tools
@@ -167,7 +179,10 @@ impl ToolRegistry {
167179 } )
168180 } )
169181 . collect ( ) ;
170- ToolRegistry { tools }
182+ ToolRegistry {
183+ tools,
184+ activated : Arc :: new ( Mutex :: new ( HashMap :: new ( ) ) ) ,
185+ }
171186 }
172187
173188 /// Clone the registry, excluding all MCP-sourced tools.
@@ -183,7 +198,10 @@ impl ToolRegistry {
183198 } )
184199 } )
185200 . collect ( ) ;
186- ToolRegistry { tools }
201+ ToolRegistry {
202+ tools,
203+ activated : Arc :: new ( Mutex :: new ( HashMap :: new ( ) ) ) ,
204+ }
187205 }
188206
189207 /// Clone the registry, excluding tools whose names are in `exclude`.
@@ -199,7 +217,10 @@ impl ToolRegistry {
199217 } )
200218 } )
201219 . collect ( ) ;
202- ToolRegistry { tools }
220+ ToolRegistry {
221+ tools,
222+ activated : Arc :: new ( Mutex :: new ( HashMap :: new ( ) ) ) ,
223+ }
203224 }
204225
205226 /// Clone the registry keeping only tools that match `predicate`.
@@ -218,10 +239,35 @@ impl ToolRegistry {
218239 } )
219240 } )
220241 . collect ( ) ;
221- ToolRegistry { tools }
242+ ToolRegistry {
243+ tools,
244+ activated : Arc :: new ( Mutex :: new ( HashMap :: new ( ) ) ) ,
245+ }
222246 }
223247}
224248
249+ fn entry_to_schema ( e : & ToolEntry ) -> serde_json:: Value {
250+ let mut schema = serde_json:: json!( {
251+ "name" : e. tool. name( ) ,
252+ "description" : e. tool. description( ) ,
253+ "parameters" : e. tool. parameters_schema( ) ,
254+ } ) ;
255+ match & e. source {
256+ ToolSource :: Builtin => {
257+ schema[ "source" ] = serde_json:: json!( "builtin" ) ;
258+ } ,
259+ ToolSource :: Mcp { server } => {
260+ schema[ "source" ] = serde_json:: json!( "mcp" ) ;
261+ schema[ "mcpServer" ] = serde_json:: json!( server) ;
262+ } ,
263+ ToolSource :: Wasm { component_hash } => {
264+ schema[ "source" ] = serde_json:: json!( "wasm" ) ;
265+ schema[ "componentHash" ] = serde_json:: json!( hex_component_hash( * component_hash) ) ;
266+ } ,
267+ }
268+ schema
269+ }
270+
225271fn hex_component_hash ( component_hash : [ u8 ; 32 ] ) -> String {
226272 let mut output = String :: with_capacity ( component_hash. len ( ) * 2 ) ;
227273 for byte in component_hash {
@@ -409,13 +455,13 @@ mod tests {
409455 }
410456
411457 #[ test]
412- fn test_get_arc_returns_cloned_tool_handle ( ) {
458+ fn test_get_returns_cloned_tool_handle ( ) {
413459 let mut registry = ToolRegistry :: new ( ) ;
414460 registry. register ( Box :: new ( DummyTool {
415461 name : "exec" . to_string ( ) ,
416462 } ) ) ;
417- assert ! ( registry. get_arc ( "exec" ) . is_some( ) ) ;
418- assert ! ( registry. get_arc ( "missing" ) . is_none( ) ) ;
463+ assert ! ( registry. get ( "exec" ) . is_some( ) ) ;
464+ assert ! ( registry. get ( "missing" ) . is_none( ) ) ;
419465 }
420466
421467 #[ test]
0 commit comments