recursive/tools/
team_manage.rs

1//! Dynamic team management tools: `team_add_role`, `team_remove_role`, `team_list_roles`.
2//!
3//! These tools give a coordinator agent the ability to create and destroy
4//! specialist roles at runtime — analogous to Fake CC's `TeamCreateTool` and
5//! `TeamDeleteTool`.
6//!
7//! The shared state is an `Arc<tokio::sync::RwLock<AgentPool>>` passed into
8//! each tool constructor so all three tools see the same pool.
9
10use async_trait::async_trait;
11use serde_json::{json, Value};
12use std::sync::Arc;
13use tokio::sync::RwLock;
14
15use crate::error::{Error, Result};
16use crate::llm::ToolSpec;
17use crate::multi::{AgentPool, AgentRole};
18use crate::tools::{Tool, ToolSideEffect};
19
20// ---------------------------------------------------------------------------
21// TeamAddRole
22// ---------------------------------------------------------------------------
23
24/// Add or update a role in the shared AgentPool at runtime.
25pub struct TeamAddRole {
26    pool: Arc<RwLock<AgentPool>>,
27}
28
29impl TeamAddRole {
30    pub fn new(pool: Arc<RwLock<AgentPool>>) -> Self {
31        Self { pool }
32    }
33}
34
35#[async_trait]
36impl Tool for TeamAddRole {
37    fn spec(&self) -> ToolSpec {
38        ToolSpec {
39            name: "team_add_role".into(),
40            description: concat!(
41                "Add (or replace) a specialist role in the coordinator's agent pool. ",
42                "After calling this, spawn_worker can use the new role name. ",
43                "Use this to dynamically create specialists for specific subtasks."
44            )
45            .into(),
46            parameters: json!({
47                "type": "object",
48                "properties": {
49                    "name": {
50                        "type": "string",
51                        "description": "Unique role name (e.g. 'sql-expert', 'ui-designer')."
52                    },
53                    "system_prompt": {
54                        "type": "string",
55                        "description": "The system prompt that defines this role's personality and focus."
56                    },
57                    "max_steps": {
58                        "type": "integer",
59                        "description": "Maximum steps for this role's agent (default 30).",
60                        "default": 30
61                    },
62                    "allowed_tools": {
63                        "type": "array",
64                        "items": { "type": "string" },
65                        "description": "Optional tool allowlist. Empty / absent means all tools."
66                    }
67                },
68                "required": ["name", "system_prompt"]
69            }),
70        }
71    }
72
73    fn side_effect_class(&self) -> ToolSideEffect {
74        ToolSideEffect::ReadOnly
75    }
76
77    async fn execute(&self, arguments: Value) -> Result<String> {
78        let name = arguments["name"]
79            .as_str()
80            .ok_or_else(|| Error::BadToolArgs {
81                name: "team_add_role".into(),
82                message: "missing required parameter: name".to_string(),
83            })?
84            .to_string();
85
86        let system_prompt = arguments["system_prompt"]
87            .as_str()
88            .ok_or_else(|| Error::BadToolArgs {
89                name: "team_add_role".into(),
90                message: "missing required parameter: system_prompt".to_string(),
91            })?
92            .to_string();
93
94        let max_steps = arguments["max_steps"].as_i64().unwrap_or(30).clamp(1, 200) as usize;
95
96        let allowed_tools: Vec<String> = arguments["allowed_tools"]
97            .as_array()
98            .map(|arr| {
99                arr.iter()
100                    .filter_map(|v| v.as_str().map(String::from))
101                    .collect()
102            })
103            .unwrap_or_default();
104
105        let role = AgentRole {
106            name: name.clone(),
107            system_prompt,
108            max_steps,
109            allowed_tools,
110        };
111
112        self.pool.write().await.add_role(role);
113        Ok(format!("Role '{name}' added to team pool."))
114    }
115}
116
117// ---------------------------------------------------------------------------
118// TeamRemoveRole
119// ---------------------------------------------------------------------------
120
121/// Remove a role from the shared AgentPool.
122pub struct TeamRemoveRole {
123    pool: Arc<RwLock<AgentPool>>,
124}
125
126impl TeamRemoveRole {
127    pub fn new(pool: Arc<RwLock<AgentPool>>) -> Self {
128        Self { pool }
129    }
130}
131
132#[async_trait]
133impl Tool for TeamRemoveRole {
134    fn spec(&self) -> ToolSpec {
135        ToolSpec {
136            name: "team_remove_role".into(),
137            description: concat!(
138                "Remove a specialist role from the coordinator's agent pool. ",
139                "Use this to clean up roles that are no longer needed."
140            )
141            .into(),
142            parameters: json!({
143                "type": "object",
144                "properties": {
145                    "name": {
146                        "type": "string",
147                        "description": "The role name to remove."
148                    }
149                },
150                "required": ["name"]
151            }),
152        }
153    }
154
155    fn side_effect_class(&self) -> ToolSideEffect {
156        ToolSideEffect::ReadOnly
157    }
158
159    async fn execute(&self, arguments: Value) -> Result<String> {
160        let name = arguments["name"]
161            .as_str()
162            .ok_or_else(|| Error::BadToolArgs {
163                name: "team_remove_role".into(),
164                message: "missing required parameter: name".to_string(),
165            })?;
166
167        let removed = self.pool.write().await.remove_role(name);
168        if removed {
169            Ok(format!("Role '{name}' removed from team pool."))
170        } else {
171            Ok(format!("Role '{name}' not found in team pool."))
172        }
173    }
174}
175
176// ---------------------------------------------------------------------------
177// TeamListRoles
178// ---------------------------------------------------------------------------
179
180/// List all roles currently registered in the shared AgentPool.
181pub struct TeamListRoles {
182    pool: Arc<RwLock<AgentPool>>,
183}
184
185impl TeamListRoles {
186    pub fn new(pool: Arc<RwLock<AgentPool>>) -> Self {
187        Self { pool }
188    }
189}
190
191#[async_trait]
192impl Tool for TeamListRoles {
193    fn spec(&self) -> ToolSpec {
194        ToolSpec {
195            name: "team_list_roles".into(),
196            description: "List all specialist roles currently registered in the agent pool.".into(),
197            parameters: json!({
198                "type": "object",
199                "properties": {}
200            }),
201        }
202    }
203
204    fn side_effect_class(&self) -> ToolSideEffect {
205        ToolSideEffect::ReadOnly
206    }
207
208    async fn execute(&self, _arguments: Value) -> Result<String> {
209        let pool = self.pool.read().await;
210        let names = pool.role_names();
211        if names.is_empty() {
212            Ok("No roles registered in team pool.".to_string())
213        } else {
214            let mut sorted = names;
215            sorted.sort_unstable();
216            Ok(format!(
217                "Roles in team pool ({}):\n{}",
218                sorted.len(),
219                sorted.join("\n")
220            ))
221        }
222    }
223}
224
225// ---------------------------------------------------------------------------
226// Tests
227// ---------------------------------------------------------------------------
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use crate::llm::MockProvider;
233    use crate::{Config, LlmProvider};
234
235    fn mock_pool() -> Arc<RwLock<AgentPool>> {
236        let provider: Arc<dyn LlmProvider> = Arc::new(MockProvider::new(vec![]));
237        let config = Config::from_env().unwrap_or_else(|_| Config {
238            workspace: std::path::PathBuf::from("."),
239            api_base: String::new(),
240            api_key: None,
241            model: "mock".into(),
242            provider_type: "mock".into(),
243            preset: None,
244            max_steps: 30,
245            temperature: 0.0,
246            system_prompt: String::new(),
247            retry_max: 0,
248            retry_initial_backoff_secs: 1,
249            retry_max_backoff_secs: 10,
250            shell_timeout_secs: 30,
251            headless: false,
252            memory_summary_limit: 5,
253            thinking_budget: None,
254            session_name: None,
255        });
256        Arc::new(RwLock::new(AgentPool::new(provider, config)))
257    }
258
259    #[tokio::test]
260    async fn team_add_role_registers_role() {
261        let pool = mock_pool();
262        let tool = TeamAddRole::new(pool.clone());
263
264        let result = tool
265            .execute(json!({
266                "name": "sql-expert",
267                "system_prompt": "You are an expert in SQL databases."
268            }))
269            .await
270            .unwrap();
271
272        assert!(result.contains("sql-expert"));
273        assert!(result.contains("added"));
274        assert!(pool.read().await.get_role("sql-expert").is_some());
275    }
276
277    #[tokio::test]
278    async fn team_remove_role_removes_existing() {
279        let pool = mock_pool();
280        // Pre-populate
281        pool.write().await.add_role(AgentRole {
282            name: "test-role".into(),
283            system_prompt: "Test.".into(),
284            max_steps: 10,
285            allowed_tools: vec![],
286        });
287
288        let tool = TeamRemoveRole::new(pool.clone());
289        let result = tool.execute(json!({"name": "test-role"})).await.unwrap();
290
291        assert!(result.contains("removed"));
292        assert!(pool.read().await.get_role("test-role").is_none());
293    }
294
295    #[tokio::test]
296    async fn team_remove_role_missing_returns_not_found() {
297        let pool = mock_pool();
298        let tool = TeamRemoveRole::new(pool.clone());
299        let result = tool.execute(json!({"name": "nonexistent"})).await.unwrap();
300        assert!(result.contains("not found"));
301    }
302
303    #[tokio::test]
304    async fn team_list_roles_empty() {
305        let pool = mock_pool();
306        let tool = TeamListRoles::new(pool);
307        let result = tool.execute(json!({})).await.unwrap();
308        assert!(result.contains("No roles"));
309    }
310
311    #[tokio::test]
312    async fn team_list_roles_with_entries() {
313        let pool = mock_pool();
314        pool.write().await.add_role(AgentRole {
315            name: "alpha".into(),
316            system_prompt: "Alpha.".into(),
317            max_steps: 10,
318            allowed_tools: vec![],
319        });
320        pool.write().await.add_role(AgentRole {
321            name: "beta".into(),
322            system_prompt: "Beta.".into(),
323            max_steps: 10,
324            allowed_tools: vec![],
325        });
326
327        let tool = TeamListRoles::new(pool);
328        let result = tool.execute(json!({})).await.unwrap();
329        assert!(result.contains("alpha"));
330        assert!(result.contains("beta"));
331        assert!(result.contains("2"));
332    }
333
334    #[tokio::test]
335    async fn team_add_role_missing_name_errors() {
336        let pool = mock_pool();
337        let tool = TeamAddRole::new(pool);
338        let result = tool.execute(json!({"system_prompt": "Something"})).await;
339        assert!(result.is_err());
340    }
341}
recursive/tools/team_manage.rs

recursive/tools/
team_manage.rs
