claude_code_acp/mcp/tools/
task_output.rs

1//! TaskOutput tool for retrieving output from background tasks
2//!
3//! Retrieves output from running or completed tasks (background shells, agents, or remote sessions).
4
5use async_trait::async_trait;
6use serde::Deserialize;
7use serde_json::{Value, json};
8
9use super::base::Tool;
10use crate::mcp::registry::{ToolContext, ToolResult};
11
12/// Input parameters for TaskOutput
13#[derive(Debug, Deserialize)]
14struct TaskOutputInput {
15    /// The task ID to get output from
16    task_id: String,
17    /// Whether to wait for completion
18    #[serde(default = "default_block")]
19    block: bool,
20    /// Max wait time in milliseconds
21    #[serde(default = "default_timeout")]
22    timeout: u64,
23}
24
25fn default_block() -> bool {
26    true
27}
28
29fn default_timeout() -> u64 {
30    30000
31}
32
33/// TaskOutput tool for retrieving task results
34#[derive(Debug, Default)]
35pub struct TaskOutputTool;
36
37impl TaskOutputTool {
38    /// Create a new TaskOutput tool
39    pub fn new() -> Self {
40        Self
41    }
42}
43
44#[async_trait]
45impl Tool for TaskOutputTool {
46    fn name(&self) -> &str {
47        "TaskOutput"
48    }
49
50    fn description(&self) -> &str {
51        "Retrieves output from a running or completed task (background shell, agent, or remote session). \
52         Takes a task_id parameter identifying the task. Returns the task output along with status information. \
53         Use block=true (default) to wait for task completion. Use block=false for non-blocking check of current status."
54    }
55
56    fn input_schema(&self) -> Value {
57        json!({
58            "$schema": "http://json-schema.org/draft-07/schema#",
59            "type": "object",
60            "required": ["task_id"],
61            "additionalProperties": false,
62            "properties": {
63                "task_id": {
64                    "type": "string",
65                    "description": "The task ID to get output from"
66                },
67                "block": {
68                    "type": "boolean",
69                    "default": true,
70                    "description": "Whether to wait for completion"
71                },
72                "timeout": {
73                    "type": "number",
74                    "default": 30000,
75                    "minimum": 0,
76                    "maximum": 600_000,
77                    "description": "Max wait time in ms"
78                }
79            }
80        })
81    }
82
83    async fn execute(&self, input: Value, context: &ToolContext) -> ToolResult {
84        // Parse input
85        let params: TaskOutputInput = match serde_json::from_value(input) {
86            Ok(p) => p,
87            Err(e) => return ToolResult::error(format!("Invalid input: {}", e)),
88        };
89
90        // Validate timeout
91        if params.timeout > 600_000 {
92            return ToolResult::error("Timeout cannot exceed 600000ms (10 minutes)");
93        }
94
95        tracing::info!(
96            "TaskOutput request for task: {} (session: {})",
97            params.task_id,
98            context.session_id
99        );
100
101        // Note: Full implementation would:
102        // 1. Look up the task by ID from a task registry
103        // 2. If blocking, wait for completion up to timeout
104        // 3. Return task output and status
105        // 4. Handle different task types (shell, agent, remote)
106
107        let mut output = format!("Task output for: {}\n\n", params.task_id);
108        output.push_str(&format!("Blocking: {}\n", params.block));
109        output.push_str(&format!("Timeout: {}ms\n", params.timeout));
110        output.push_str("\nStatus: Task not found\n");
111        output.push_str(
112            "\nNote: TaskOutput tool requires task registry integration for full functionality. \
113             Use the /tasks command to see available task IDs.",
114        );
115
116        ToolResult::success(output).with_metadata(json!({
117            "task_id": params.task_id,
118            "status": "not_found",
119            "block": params.block,
120            "timeout": params.timeout
121        }))
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128    use tempfile::TempDir;
129
130    #[test]
131    fn test_task_output_properties() {
132        let tool = TaskOutputTool::new();
133        assert_eq!(tool.name(), "TaskOutput");
134        assert!(tool.description().contains("task"));
135        assert!(tool.description().contains("output"));
136    }
137
138    #[test]
139    fn test_task_output_input_schema() {
140        let tool = TaskOutputTool::new();
141        let schema = tool.input_schema();
142
143        assert_eq!(schema["type"], "object");
144        assert!(schema["properties"]["task_id"].is_object());
145        assert!(schema["properties"]["block"].is_object());
146        assert!(schema["properties"]["timeout"].is_object());
147
148        let required = schema["required"].as_array().unwrap();
149        assert!(required.contains(&json!("task_id")));
150    }
151
152    #[tokio::test]
153    async fn test_task_output_execute() {
154        let temp_dir = TempDir::new().unwrap();
155        let tool = TaskOutputTool::new();
156        let context = ToolContext::new("test-session", temp_dir.path());
157
158        let result = tool
159            .execute(
160                json!({
161                    "task_id": "test-task-123"
162                }),
163                &context,
164            )
165            .await;
166
167        assert!(!result.is_error);
168        assert!(result.content.contains("test-task-123"));
169        assert!(result.content.contains("not found"));
170    }
171
172    #[tokio::test]
173    async fn test_task_output_non_blocking() {
174        let temp_dir = TempDir::new().unwrap();
175        let tool = TaskOutputTool::new();
176        let context = ToolContext::new("test-session", temp_dir.path());
177
178        let result = tool
179            .execute(
180                json!({
181                    "task_id": "test-task-456",
182                    "block": false
183                }),
184                &context,
185            )
186            .await;
187
188        assert!(!result.is_error);
189        assert!(result.content.contains("Blocking: false"));
190    }
191
192    #[tokio::test]
193    async fn test_task_output_with_timeout() {
194        let temp_dir = TempDir::new().unwrap();
195        let tool = TaskOutputTool::new();
196        let context = ToolContext::new("test-session", temp_dir.path());
197
198        let result = tool
199            .execute(
200                json!({
201                    "task_id": "test-task-789",
202                    "timeout": 60000
203                }),
204                &context,
205            )
206            .await;
207
208        assert!(!result.is_error);
209        assert!(result.content.contains("Timeout: 60000ms"));
210    }
211
212    #[tokio::test]
213    async fn test_task_output_timeout_too_large() {
214        let temp_dir = TempDir::new().unwrap();
215        let tool = TaskOutputTool::new();
216        let context = ToolContext::new("test-session", temp_dir.path());
217
218        let result = tool
219            .execute(
220                json!({
221                    "task_id": "test-task",
222                    "timeout": 999_999
223                }),
224                &context,
225            )
226            .await;
227
228        assert!(result.is_error);
229        assert!(result.content.contains("600000ms"));
230    }
231}
claude_code_acp/mcp/tools/task_output.rs

claude_code_acp/mcp/tools/
task_output.rs
