Skip to main content

pawan/tools/
task.rs

1//! Task tool: spawn an in-process subagent with restricted tools.
2//!
3//! This tool runs a child `PawanAgent` with a narrowed `ToolRegistry`, a smaller
4//! context window, and a hard timeout. Subagents are depth-limited (they cannot
5//! spawn other agents).
6
7use super::Tool;
8use crate::agent::backend::LlmBackend;
9use crate::agent::PawanAgent;
10use crate::config::PawanConfig;
11use crate::tools::{bash, batch, edit, file, git, lsp_tool, mise, native, ToolRegistry, ToolTier};
12use crate::{PawanError, Result};
13use async_trait::async_trait;
14use serde::Deserialize;
15use serde_json::{json, Value};
16use std::path::PathBuf;
17use std::sync::Arc;
18use std::time::Duration;
19use tokio::time::timeout;
20
21const DEFAULT_TIMEOUT_SECS: u64 = 300;
22
23#[derive(Debug, Clone, Deserialize)]
24struct TaskArgs {
25    agent: String,
26    assignment: String,
27    #[serde(default)]
28    context: Option<String>,
29    #[serde(default)]
30    model: Option<String>,
31    /// Timeout in seconds (default: 300).
32    #[serde(default)]
33    timeout: Option<u64>,
34}
35
36#[derive(Clone)]
37pub struct TaskTool {
38    workspace_root: PathBuf,
39}
40
41impl TaskTool {
42    pub fn new(workspace_root: PathBuf) -> Self {
43        Self { workspace_root }
44    }
45
46    fn known_agent_types() -> &'static [&'static str] {
47        &[
48            "explore",
49            "plan",
50            "task",
51            "reviewer",
52            "designer",
53            "librarian",
54        ]
55    }
56
57    fn validate_agent_type(agent: &str) -> std::result::Result<(), String> {
58        if Self::known_agent_types().iter().any(|t| *t == agent) {
59            Ok(())
60        } else {
61            Err(format!(
62                "unknown agent type '{agent}'. Valid types: {}",
63                Self::known_agent_types().join(", ")
64            ))
65        }
66    }
67
68    fn validate_assignment(assignment: &str) -> std::result::Result<(), String> {
69        if assignment.trim().is_empty() {
70            Err("assignment must be non-empty".to_string())
71        } else {
72            Ok(())
73        }
74    }
75
76    fn system_prompt_for(agent: &str) -> String {
77        match agent {
78            "explore" => "You are a read-only exploration subagent. Use only the allowed read/search tools to gather facts. Do not propose or apply code edits. Return concise findings with file paths and evidence.".to_string(),
79            "plan" => "You are an architecture subagent. Do not modify code. Make design decisions and propose an implementation plan with tradeoffs, invariants, and acceptance criteria.".to_string(),
80            "reviewer" => "You are a code review subagent. Do not modify code. Identify bugs, security issues, and quality concerns. Return a structured review report with severity and recommendations.".to_string(),
81            "designer" => "You are a UI/UX design subagent. If editing tools are available, you may implement UI changes carefully. Prioritize accessibility and consistency.".to_string(),
82            "librarian" => "You are a research subagent. Verify details from authoritative sources and the local codebase. Do not modify code. Return actionable guidance.".to_string(),
83            _ => "You are a subagent executing a delegated task. Follow the assignment precisely and return the final result. Do not spawn other agents.".to_string(),
84        }
85    }
86
87    fn build_user_prompt(context: Option<&str>, assignment: &str) -> String {
88        match context {
89            Some(ctx) if !ctx.trim().is_empty() => format!(
90                "{ctx}\n\n[Assignment]\n{assignment}\n\n[Constraints]\n- Subagent depth limit: you cannot spawn other agents.\n"
91            ),
92            _ => format!(
93                "[Assignment]\n{assignment}\n\n[Constraints]\n- Subagent depth limit: you cannot spawn other agents.\n"
94            ),
95        }
96    }
97
98    fn registry_for(agent: &str, workspace_root: &PathBuf) -> ToolRegistry {
99        use ToolTier::*;
100        let mut reg = ToolRegistry::new();
101
102        // Read/search tools
103        reg.register_with_tier(
104            Arc::new(file::ReadFileTool::new(workspace_root.clone())),
105            Core,
106        );
107        reg.register_with_tier(
108            Arc::new(file::ListDirectoryTool::new(workspace_root.clone())),
109            Standard,
110        );
111        reg.register_with_tier(
112            Arc::new(native::GlobSearchTool::new(workspace_root.clone())),
113            Core,
114        );
115        reg.register_with_tier(
116            Arc::new(native::GrepSearchTool::new(workspace_root.clone())),
117            Core,
118        );
119        reg.register_with_tier(
120            Arc::new(native::AstGrepTool::new(workspace_root.clone())),
121            Core,
122        );
123        reg.register_with_tier(
124            Arc::new(native::RipgrepTool::new(workspace_root.clone())),
125            Extended,
126        );
127        reg.register_with_tier(
128            Arc::new(native::FdTool::new(workspace_root.clone())),
129            Extended,
130        );
131
132        match agent {
133            "explore" | "plan" | "reviewer" | "librarian" => {
134                reg.register_with_tier(
135                    Arc::new(git::GitStatusTool::new(workspace_root.clone())),
136                    Standard,
137                );
138                reg.register_with_tier(
139                    Arc::new(git::GitDiffTool::new(workspace_root.clone())),
140                    Standard,
141                );
142                reg.register_with_tier(
143                    Arc::new(git::GitLogTool::new(workspace_root.clone())),
144                    Standard,
145                );
146                reg.register_with_tier(
147                    Arc::new(git::GitBlameTool::new(workspace_root.clone())),
148                    Standard,
149                );
150                reg
151            }
152            "task" | "designer" => {
153                reg.register_with_tier(Arc::new(bash::BashTool::new(workspace_root.clone())), Core);
154                reg.register_with_tier(
155                    Arc::new(file::WriteFileTool::new(workspace_root.clone())),
156                    Core,
157                );
158                reg.register_with_tier(
159                    Arc::new(edit::EditFileTool::new(workspace_root.clone())),
160                    Core,
161                );
162                reg.register_with_tier(
163                    Arc::new(edit::EditFileLinesTool::new(workspace_root.clone())),
164                    Standard,
165                );
166                reg.register_with_tier(
167                    Arc::new(edit::InsertAfterTool::new(workspace_root.clone())),
168                    Standard,
169                );
170                reg.register_with_tier(
171                    Arc::new(edit::AppendFileTool::new(workspace_root.clone())),
172                    Standard,
173                );
174
175                reg.register_with_tier(
176                    Arc::new(git::GitStatusTool::new(workspace_root.clone())),
177                    Standard,
178                );
179                reg.register_with_tier(
180                    Arc::new(git::GitDiffTool::new(workspace_root.clone())),
181                    Standard,
182                );
183                reg.register_with_tier(
184                    Arc::new(git::GitAddTool::new(workspace_root.clone())),
185                    Standard,
186                );
187                reg.register_with_tier(
188                    Arc::new(git::GitCommitTool::new(workspace_root.clone())),
189                    Standard,
190                );
191                reg.register_with_tier(
192                    Arc::new(git::GitLogTool::new(workspace_root.clone())),
193                    Standard,
194                );
195                reg.register_with_tier(
196                    Arc::new(git::GitBlameTool::new(workspace_root.clone())),
197                    Standard,
198                );
199                reg.register_with_tier(
200                    Arc::new(git::GitBranchTool::new(workspace_root.clone())),
201                    Standard,
202                );
203                reg.register_with_tier(
204                    Arc::new(git::GitCheckoutTool::new(workspace_root.clone())),
205                    Standard,
206                );
207                reg.register_with_tier(
208                    Arc::new(git::GitStashTool::new(workspace_root.clone())),
209                    Standard,
210                );
211
212                reg.register_with_tier(
213                    Arc::new(batch::BatchTool::new(workspace_root.clone())),
214                    Standard,
215                );
216                reg.register_with_tier(
217                    Arc::new(lsp_tool::LspTool::new(workspace_root.clone())),
218                    Extended,
219                );
220                reg.register_with_tier(
221                    Arc::new(mise::MiseTool::new(workspace_root.clone())),
222                    Extended,
223                );
224                reg.register_with_tier(
225                    Arc::new(native::SdTool::new(workspace_root.clone())),
226                    Extended,
227                );
228                reg.register_with_tier(
229                    Arc::new(native::ErdTool::new(workspace_root.clone())),
230                    Extended,
231                );
232                reg
233            }
234            _ => reg,
235        }
236    }
237
238    async fn run_subagent(
239        &self,
240        agent_type: &str,
241        assignment: &str,
242        context: Option<&str>,
243        model: Option<&str>,
244        timeout_secs: u64,
245        backend_override: Option<Box<dyn LlmBackend>>,
246    ) -> Result<Value> {
247        let mut config = PawanConfig::default();
248        config.system_prompt = Some(Self::system_prompt_for(agent_type));
249        config.max_context_tokens = 32_000;
250        config.max_tool_iterations = 20;
251        config.eruka.enabled = false;
252        if let Some(m) = model {
253            config.model = m.to_string();
254        }
255
256        let tools = Self::registry_for(agent_type, &self.workspace_root);
257        let prompt = Self::build_user_prompt(context, assignment);
258
259        let mut agent = PawanAgent::new(config, self.workspace_root.clone()).with_tools(tools);
260        if let Some(backend) = backend_override {
261            agent = agent.with_backend(backend);
262        }
263
264        let run = agent.execute(&prompt);
265        let response = match timeout(Duration::from_secs(timeout_secs), run).await {
266            Ok(res) => res.map_err(|e| PawanError::Tool(format!("subagent error: {e}")))?,
267            Err(_) => {
268                return Ok(json!({
269                    "agent": agent_type,
270                    "status": "error",
271                    "result": format!("subagent timeout after {timeout_secs}s"),
272                }));
273            }
274        };
275
276        Ok(json!({
277            "agent": agent_type,
278            "status": "completed",
279            "result": response.content,
280            "usage": {
281                "prompt_tokens": response.usage.prompt_tokens,
282                "completion_tokens": response.usage.completion_tokens,
283                "total_tokens": response.usage.total_tokens,
284                "reasoning_tokens": response.usage.reasoning_tokens,
285                "action_tokens": response.usage.action_tokens,
286            }
287        }))
288    }
289}
290
291#[async_trait]
292impl Tool for TaskTool {
293    fn name(&self) -> &str {
294        "task"
295    }
296
297    fn description(&self) -> &str {
298        "Spawn an in-process subagent with restricted tools to complete an assignment."
299    }
300
301    fn mutating(&self) -> bool {
302        true
303    }
304
305    fn parameters_schema(&self) -> Value {
306        json!({
307            "type": "object",
308            "properties": {
309                "agent": {"type": "string"},
310                "assignment": {"type": "string"},
311                "context": {"type": "string"},
312                "model": {"type": "string"},
313                "timeout": {"type": "integer"}
314            },
315            "required": ["agent", "assignment"]
316        })
317    }
318
319    async fn execute(&self, args: Value) -> Result<Value> {
320        let parsed: TaskArgs = serde_json::from_value(args)
321            .map_err(|e| PawanError::Tool(format!("invalid task args: {e}")))?;
322
323        Self::validate_agent_type(&parsed.agent).map_err(PawanError::Tool)?;
324        Self::validate_assignment(&parsed.assignment).map_err(PawanError::Tool)?;
325
326        let timeout_secs = parsed.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS);
327
328        self.run_subagent(
329            &parsed.agent,
330            &parsed.assignment,
331            parsed.context.as_deref(),
332            parsed.model.as_deref(),
333            timeout_secs,
334            None,
335        )
336        .await
337    }
338}
339
340#[cfg(test)]
341mod tests {
342    use super::*;
343    use crate::agent::backend::mock::{MockBackend, MockResponse};
344    use serde_json::json;
345
346    #[tokio::test]
347    async fn unknown_agent_type_rejects() {
348        let dir = tempfile::tempdir().unwrap();
349        let tool = TaskTool::new(dir.path().to_path_buf());
350        let err = tool
351            .execute(json!({"agent": "nope", "assignment": "hi"}))
352            .await
353            .unwrap_err();
354        assert!(err.to_string().contains("unknown agent type"));
355    }
356
357    #[tokio::test]
358    async fn timeout_returns_error_status() {
359        let dir = tempfile::tempdir().unwrap();
360        let tool = TaskTool::new(dir.path().to_path_buf());
361
362        let out = tool
363            .run_subagent(
364                "explore",
365                "This will time out immediately.",
366                None,
367                None,
368                0,
369                None,
370            )
371            .await
372            .unwrap();
373
374        assert_eq!(out["status"].as_str().unwrap(), "error");
375        assert!(out["result"].as_str().unwrap().contains("timeout"));
376    }
377
378    #[tokio::test]
379    async fn explore_agent_runs_and_returns_findings_with_mock_backend() {
380        let dir = tempfile::tempdir().unwrap();
381        let tool = TaskTool::new(dir.path().to_path_buf());
382
383        let backend = Box::new(MockBackend::new(vec![MockResponse::text(
384            "Findings: crates/pawan-core/src/lib.rs is the crate root.",
385        )]));
386
387        let out = tool
388            .run_subagent(
389                "explore",
390                "Explore the repo and return findings.",
391                Some("Context here"),
392                None,
393                5,
394                Some(backend),
395            )
396            .await
397            .unwrap();
398
399        assert_eq!(out["agent"].as_str().unwrap(), "explore");
400        assert_eq!(out["status"].as_str().unwrap(), "completed");
401        assert!(out["result"].as_str().unwrap().contains("Findings:"));
402    }
403}