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 schema(&self) -> ToolSchema {
37        ToolSchema::new(
38            "delegate_task",
39            "Delegate a task to a specialized sub-agent (agent_type: general|code|test|research). Runs to completion and returns its summary.",
40            vec![
41                ToolParameter::new("task_description", ParamType::String, true, "Task for the sub-agent", &["task", "description"]),
42                ToolParameter::new("agent_type", ParamType::String, false, "general|code|test|research", &["type"]),
43            ],
44        )
45    }
46
47    async fn call(&self, args: &Value, cancel: &AtomicBool) -> Result<String> {
48        let task = args
49            .get("task_description")
50            .and_then(|v| v.as_str())
51            .unwrap_or("");
52        if task.is_empty() {
53            return Ok(json!({"success": false, "error": "missing task_description"}).to_string());
54        }
55        let agent_type = args
56            .get("agent_type")
57            .and_then(|v| v.as_str())
58            .unwrap_or("general");
59        let registry = (self.sub_registry)();
60        // Sub-agent output is collected but not streamed to the main sink to
61        // avoid confusing the UI; the returned summary is what matters.
62        let sink = CollectingSink::new(false);
63        let summary = run_sub_agent(
64            agent_type,
65            task,
66            self.llm.as_ref(),
67            &registry,
68            &self.config,
69            &self.project_path,
70            &sink,
71            cancel,
72        )
73        .await?;
74        Ok(json!({
75            "success": true,
76            "agent_type": agent_type,
77            "result": summary,
78        })
79        .to_string())
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86    use crate::llm::types::{ChatMessage, LlmResponse};
87    use crate::stream::EventSink;
88    use std::sync::Mutex;
89
90    struct OneShot {
91        text: Mutex<Option<String>>,
92    }
93    #[async_trait::async_trait]
94    impl LlmClient for OneShot {
95        async fn chat_stream(
96            &self,
97            _m: &[ChatMessage],
98            _t: &[Value],
99            _e: &dyn EventSink,
100            _c: &AtomicBool,
101        ) -> Result<LlmResponse> {
102            Ok(LlmResponse {
103                text: self.text.lock().unwrap().take().unwrap_or_default(),
104                ..Default::default()
105            })
106        }
107        async fn validate(&self) -> Result<()> {
108            Ok(())
109        }
110    }
111
112    fn tmp() -> PathBuf {
113        let d = std::env::temp_dir().join(format!(
114            "aacode_deleg_{}_{}",
115            std::process::id(),
116            uuid::Uuid::new_v4().simple()
117        ));
118        std::fs::create_dir_all(&d).unwrap();
119        d
120    }
121
122    #[test]
123    fn delegate_runs_sub_agent() {
124        let rt = tokio::runtime::Runtime::new().unwrap();
125        rt.block_on(async {
126            let llm: Arc<dyn LlmClient> = Arc::new(OneShot {
127                text: Mutex::new(Some("sub result".into())),
128            });
129            let factory: SubRegistryFactory = Arc::new(|| ToolRegistry::new());
130            let tool = DelegateTaskTool {
131                llm,
132                config: AgentConfig::default(),
133                project_path: tmp(),
134                sub_registry: factory,
135            };
136            let cancel = AtomicBool::new(false);
137            let out = tool
138                .call(
139                    &json!({"task_description": "do sub", "agent_type": "code"}),
140                    &cancel,
141                )
142                .await
143                .unwrap();
144            let v: Value = serde_json::from_str(&out).unwrap();
145            assert_eq!(v["success"], true);
146            assert_eq!(v["result"], "sub result");
147        });
148    }
149
150    #[tokio::test]
151    async fn delegate_missing_task() {
152        let llm: Arc<dyn LlmClient> = Arc::new(OneShot {
153            text: Mutex::new(None),
154        });
155        let factory: SubRegistryFactory = Arc::new(|| ToolRegistry::new());
156        let tool = DelegateTaskTool {
157            llm,
158            config: AgentConfig::default(),
159            project_path: tmp(),
160            sub_registry: factory,
161        };
162        let cancel = AtomicBool::new(false);
163        let out = tool.call(&json!({}), &cancel).await.unwrap();
164        assert_eq!(
165            serde_json::from_str::<Value>(&out).unwrap()["success"],
166            false
167        );
168    }
169}
aacode_rs/tools/delegate.rs

aacode_rs/tools/
delegate.rs
