ai_agent/tools/task_output/
mod.rs1use crate::error::AgentError;
7use crate::types::*;
8
9pub mod constants;
10pub use constants::TASK_OUTPUT_TOOL_NAME;
11
12pub 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 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
117async fn get_task_output(
123 task_id: &str,
124 block: bool,
125 timeout_ms: u64,
126) -> TaskOutputData {
127 if block {
134 let _timeout = timeout_ms;
136 }
138
139 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}