Skip to main content

ai_agent/tools/task_output/
mod.rs

1// Source: ~/claudecode/openclaudecode/src/tools/TaskOutputTool/TaskOutputTool.tsx
2//! TaskOutput tool — retrieve output from background tasks.
3//!
4//! Supports both blocking (wait for completion) and non-blocking modes.
5
6use crate::error::AgentError;
7use crate::types::*;
8
9pub mod constants;
10pub use constants::TASK_OUTPUT_TOOL_NAME;
11
12/// TaskOutput tool — retrieve output from a completed or running background task.
13pub struct TaskOutputTool;
14
15impl TaskOutputTool {
16    pub fn new() -> Self {
17        Self
18    }
19
20    pub fn name(&self) -> &str {
21        TASK_OUTPUT_TOOL_NAME
22    }
23
24    pub fn description(&self) -> &str {
25        "Retrieve output from a running or completed background task (bash command, agent, etc.). Supports blocking wait for completion with configurable timeout."
26    }
27
28    pub fn user_facing_name(&self, _input: Option<&serde_json::Value>) -> String {
29        "TaskOutput".to_string()
30    }
31
32    pub fn get_tool_use_summary(&self, input: Option<&serde_json::Value>) -> Option<String> {
33        input.and_then(|inp| inp["task_id"].as_str().map(String::from))
34    }
35
36    pub fn render_tool_result_message(
37        &self,
38        content: &serde_json::Value,
39    ) -> Option<String> {
40        let text = content["content"].as_str()?;
41        let lines = text.lines().count();
42        Some(format!("{} lines", lines))
43    }
44
45    pub fn input_schema(&self) -> ToolInputSchema {
46        ToolInputSchema {
47            schema_type: "object".to_string(),
48            properties: serde_json::json!({
49                "task_id": {
50                    "type": "string",
51                    "description": "The task ID to get output from"
52                },
53                "block": {
54                    "type": "boolean",
55                    "description": "Whether to wait for completion. Default: true"
56                },
57                "timeout": {
58                    "type": "number",
59                    "description": "Max wait time in ms. Default: 30000, max: 600000"
60                }
61            }),
62            required: Some(vec!["task_id".to_string()]),
63        }
64    }
65
66    pub async fn execute(
67        &self,
68        input: serde_json::Value,
69        _context: &ToolContext,
70    ) -> Result<ToolResult, AgentError> {
71        let task_id = input["task_id"]
72            .as_str()
73            .ok_or_else(|| AgentError::Tool("Missing required parameter: task_id".to_string()))?;
74
75        let block = input["block"]
76            .as_bool()
77            .unwrap_or(true);
78        let timeout_ms = input["timeout"]
79            .as_u64()
80            .unwrap_or(30_000)
81            .min(600_000);
82
83        // In the SDK, there is no shared task registry.
84        // The task output is retrieved from disk or the task framework.
85        // For tasks managed by this Agent instance, output would be stored
86        // in the task framework state. For now, attempt disk-based retrieval.
87        let output = get_task_output(task_id, block, timeout_ms).await;
88
89        let result = serde_json::json!({
90            "retrieval_status": output.status,
91            "task": {
92                "task_id": task_id,
93                "task_type": output.task_type,
94                "status": output.status.clone(),
95                "description": output.description,
96                "output": output.content
97            }
98        });
99
100        Ok(ToolResult {
101            result_type: "text".to_string(),
102            tool_use_id: "".to_string(),
103            content: serde_json::to_string_pretty(&result).unwrap_or_default(),
104            is_error: Some(false),
105            was_persisted: None,
106        })
107    }
108}
109
110struct TaskOutputData {
111    status: String,
112    task_type: String,
113    description: String,
114    content: String,
115}
116
117/// Attempt to retrieve task output.
118///
119/// In the SDK context, tasks are managed externally.
120/// This reads from the task output file if available, otherwise
121/// reports the task as not found.
122async fn get_task_output(
123    task_id: &str,
124    block: bool,
125    timeout_ms: u64,
126) -> TaskOutputData {
127    // In the full CLI implementation, this would:
128    // 1. Look up the task in appState.tasks
129    // 2. If blocking, poll until complete or timeout
130    // 3. Read stdout/stderr from the task
131    // 4. Format output with task metadata
132
133    if block {
134        // Poll for task completion (SDK: no shared state, timeout immediately)
135        let _timeout = timeout_ms;
136        // In a full implementation with a task registry, would poll here
137    }
138
139    // Try reading from disk (task output files)
140    // This is the same pattern as the TS getTaskOutput()
141    // which reads from transcript.jsonl sidechain output
142    TaskOutputData {
143        status: "not_found".to_string(),
144        task_type: "unknown".to_string(),
145        description: format!("Task {} not found in local task registry", task_id),
146        content: format!(
147            "Task output not available for '{}'. In the SDK context, background task output is managed by the caller's task framework.",
148            task_id
149        ),
150    }
151}
152
153impl Default for TaskOutputTool {
154    fn default() -> Self {
155        Self::new()
156    }
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_task_output_tool_name() {
165        let tool = TaskOutputTool::new();
166        assert_eq!(tool.name(), TASK_OUTPUT_TOOL_NAME);
167    }
168
169    #[test]
170    fn test_task_output_tool_schema() {
171        let tool = TaskOutputTool::new();
172        let schema = tool.input_schema();
173        assert_eq!(schema.schema_type, "object");
174        assert!(schema.properties.get("task_id").is_some());
175        assert!(schema.properties.get("block").is_some());
176        assert!(schema.properties.get("timeout").is_some());
177    }
178
179    #[tokio::test]
180    async fn test_task_output_requires_task_id() {
181        let tool = TaskOutputTool::new();
182        let input = serde_json::json!({});
183        let context = ToolContext::default();
184        let result = tool.execute(input, &context).await;
185        assert!(result.is_err());
186    }
187
188    #[tokio::test]
189    async fn test_task_output_with_task_id() {
190        let tool = TaskOutputTool::new();
191        let input = serde_json::json!({
192            "task_id": "test-task-123",
193            "block": false
194        });
195        let context = ToolContext::default();
196        let result = tool.execute(input, &context).await;
197        assert!(result.is_ok());
198        let content = result.unwrap().content;
199        assert!(content.contains("test-task-123"));
200        assert!(content.contains("not_found"));
201    }
202
203    #[tokio::test]
204    async fn test_task_output_blocking_mode() {
205        let tool = TaskOutputTool::new();
206        let input = serde_json::json!({
207            "task_id": "blocking-task-456",
208            "block": true,
209            "timeout": 1000
210        });
211        let context = ToolContext::default();
212        let result = tool.execute(input, &context).await;
213        assert!(result.is_ok());
214    }
215
216    #[tokio::test]
217    async fn test_task_output_timeout_cap() {
218        let tool = TaskOutputTool::new();
219        let input = serde_json::json!({
220            "task_id": "timeout-task",
221            "timeout": 999_999
222        });
223        let context = ToolContext::default();
224        let result = tool.execute(input, &context).await;
225        assert!(result.is_ok());
226    }
227}