Skip to main content

agent_sdk/tools/
context_tools.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use serde_json::json;
5
6use crate::error::SdkResult;
7use crate::traits::tool::{Tool, ToolDefinition};
8use crate::task::store::TaskStore;
9
10pub struct GetTaskContextTool {
11    pub task_store: Arc<TaskStore>,
12}
13
14#[async_trait]
15impl Tool for GetTaskContextTool {
16    fn definition(&self) -> ToolDefinition {
17        ToolDefinition {
18            name: "get_task_context".to_string(),
19            description: "Get the result and notes from a completed task.".to_string(),
20            parameters: json!({
21                "type": "object",
22                "properties": {
23                    "task_id": { "type": "string", "description": "The UUID of the completed task" }
24                },
25                "required": ["task_id"]
26            }),
27        }
28    }
29
30    async fn execute(&self, arguments: serde_json::Value) -> SdkResult<serde_json::Value> {
31        let task_id_str = arguments["task_id"].as_str().unwrap_or("");
32        if task_id_str.is_empty() {
33            return Ok(json!({"error": "Missing 'task_id' argument"}));
34        }
35
36        let task_id: uuid::Uuid = match task_id_str.parse() {
37            Ok(id) => id,
38            Err(_) => return Ok(json!({"error": "Invalid task_id format"})),
39        };
40
41        match self.task_store.read_task(task_id) {
42            Ok(task) => {
43                let status = format!("{:?}", task.status);
44                Ok(json!({
45                    "task_id": task_id_str,
46                    "title": task.title,
47                    "status": status,
48                    "target_file": task.target_file.to_string_lossy(),
49                    "result": task.result.as_ref().map(|r| json!({
50                        "notes": r.notes,
51                        "file_changes": r.file_changes.len(),
52                        "tokens_used": r.llm_tokens_used
53                    }))
54                }))
55            }
56            Err(_) => Ok(json!({ "error": format!("Task {} not found", task_id_str) })),
57        }
58    }
59}
60
61pub struct ListCompletedTasksTool {
62    pub task_store: Arc<TaskStore>,
63}
64
65#[async_trait]
66impl Tool for ListCompletedTasksTool {
67    fn definition(&self) -> ToolDefinition {
68        ToolDefinition {
69            name: "list_completed_tasks".to_string(),
70            description: "List all completed tasks with their IDs, titles, and target files.".to_string(),
71            parameters: json!({ "type": "object", "properties": {} }),
72        }
73    }
74
75    async fn execute(&self, _arguments: serde_json::Value) -> SdkResult<serde_json::Value> {
76        let tasks = self.task_store.list_tasks_in_dir("completed")?;
77        let summaries: Vec<serde_json::Value> = tasks
78            .iter()
79            .map(|t| {
80                json!({
81                    "task_id": t.id.to_string(),
82                    "title": t.title,
83                    "target_file": t.target_file.to_string_lossy(),
84                    "notes": t.result.as_ref().map(|r| r.notes.as_str()).unwrap_or("")
85                })
86            })
87            .collect();
88
89        Ok(json!({ "tasks": summaries, "count": summaries.len() }))
90    }
91}