aacode_rs/tools/
delegate.rs

1// Copyright (c) 2026 xiefujin <490021684@qq.com>
2// Licensed under GPL-3.0, see LICENSE file for full license terms.
3
4//! Delegation tool — delegate_task.
5//!
6//! `delegate_task` runs a specialized sub-agent (code/test/research) inline to
7//! completion and returns its summary. The sub-agent uses a registry WITHOUT
8//! delegation tools to prevent unbounded recursion.
9//!
10//! Ported from Python `core/main_agent.py` (delegate_task) + `core/sub_agent.py`.
11
12use super::registry::{Tool, ToolRegistry};
13use super::schema::{ParamType, ToolParameter, ToolSchema};
14use crate::agent::sub_agent::run_sub_agent;
15use crate::config::AgentConfig;
16use crate::error::Result;
17use crate::llm::LlmClient;
18use crate::stream::CollectingSink;
19use serde_json::{json, Value};
20use std::path::PathBuf;
21use std::sync::atomic::AtomicBool;
22use std::sync::Arc;
23
24/// Factory that builds a sub-agent registry (no delegation tools).
25pub type SubRegistryFactory = Arc<dyn Fn() -> ToolRegistry + Send + Sync>;
26
27pub struct DelegateTaskTool {
28    pub llm: Arc<dyn LlmClient>,
29    pub config: AgentConfig,
30    pub project_path: PathBuf,
31    pub sub_registry: SubRegistryFactory,
32}
33
34#[async_trait::async_trait]
35impl Tool for DelegateTaskTool {
36    fn lane(&self) -> Option<&'static str> {
37        Some("delegate")
38    }
39    fn schema(&self) -> ToolSchema {
40        ToolSchema::new(
41            "delegate_task",
42            "Delegate a task to a specialized sub-agent (agent_type: general|code|test|research). Runs to completion and returns its summary.",
43            vec![
44                ToolParameter::new("task_description", ParamType::String, true, "Task for the sub-agent", &["task", "description"]),
45                ToolParameter::new("agent_type", ParamType::String, false, "general|code|test|research", &["type"]),
46            ],
47        )
48    }
49
50    async fn call(&self, args: &Value, cancel: &AtomicBool) -> Result<String> {
51        let task = args
52            .get("task_description")
53            .and_then(|v| v.as_str())
54            .unwrap_or("");
55        if task.is_empty() {
56            return Ok(json!({"success": false, "error": "missing task_description"}).to_string());
57        }
58        let agent_type = args
59            .get("agent_type")
60            .and_then(|v| v.as_str())
61            .unwrap_or("general");
62        let registry = (self.sub_registry)();
63        // Sub-agent output is collected but not streamed to the main sink to
64        // avoid confusing the UI; the returned summary is what matters.
65        let sink = CollectingSink::new(false);
66        let summary = run_sub_agent(
67            agent_type,
68            task,
69            self.llm.as_ref(),
70            &registry,
71            &self.config,
72            &self.project_path,
73            &sink,
74            cancel,
75        )
76        .await?;
77        Ok(json!({
78            "success": true,
79            "agent_type": agent_type,
80            "result": summary,
81        })
82        .to_string())
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use crate::llm::types::{ChatMessage, LlmResponse};
90    use crate::stream::EventSink;
91    use std::sync::Mutex;
92
93    struct OneShot {
94        text: Mutex<Option<String>>,
95    }
96    #[async_trait::async_trait]
97    impl LlmClient for OneShot {
98        async fn chat_stream(
99            &self,
100            _m: &[ChatMessage],
101            _t: &[Value],
102            _e: &dyn EventSink,
103            _c: &AtomicBool,
104        ) -> Result<LlmResponse> {
105            Ok(LlmResponse {
106                text: self.text.lock().unwrap().take().unwrap_or_default(),
107                ..Default::default()
108            })
109        }
110        async fn validate(&self) -> Result<()> {
111            Ok(())
112        }
113    }
114
115    fn tmp() -> PathBuf {
116        let d = std::env::temp_dir().join(format!(
117            "aacode_deleg_{}_{}",
118            std::process::id(),
119            uuid::Uuid::new_v4().simple()
120        ));
121        std::fs::create_dir_all(&d).unwrap();
122        d
123    }
124
125    #[test]
126    fn delegate_runs_sub_agent() {
127        let rt = tokio::runtime::Runtime::new().unwrap();
128        rt.block_on(async {
129            let llm: Arc<dyn LlmClient> = Arc::new(OneShot {
130                text: Mutex::new(Some("sub result".into())),
131            });
132            let factory: SubRegistryFactory = Arc::new(|| ToolRegistry::new());
133            let tool = DelegateTaskTool {
134                llm,
135                config: AgentConfig::default(),
136                project_path: tmp(),
137                sub_registry: factory,
138            };
139            let cancel = AtomicBool::new(false);
140            let out = tool
141                .call(
142                    &json!({"task_description": "do sub", "agent_type": "code"}),
143                    &cancel,
144                )
145                .await
146                .unwrap();
147            let v: Value = serde_json::from_str(&out).unwrap();
148            assert_eq!(v["success"], true);
149            assert_eq!(v["result"], "sub result");
150        });
151    }
152
153    #[tokio::test]
154    async fn delegate_missing_task() {
155        let llm: Arc<dyn LlmClient> = Arc::new(OneShot {
156            text: Mutex::new(None),
157        });
158        let factory: SubRegistryFactory = Arc::new(|| ToolRegistry::new());
159        let tool = DelegateTaskTool {
160            llm,
161            config: AgentConfig::default(),
162            project_path: tmp(),
163            sub_registry: factory,
164        };
165        let cancel = AtomicBool::new(false);
166        let out = tool.call(&json!({}), &cancel).await.unwrap();
167        assert_eq!(
168            serde_json::from_str::<Value>(&out).unwrap()["success"],
169            false
170        );
171    }
172}
aacode_rs/tools/delegate.rs

aacode_rs/tools/
delegate.rs
