diff --git a/crates/goose-mcp/src/developer/mod.rs b/crates/goose-mcp/src/developer/mod.rs index 3a3a433106a9..9540a1d3aca9 100644 --- a/crates/goose-mcp/src/developer/mod.rs +++ b/crates/goose-mcp/src/developer/mod.rs @@ -1,6 +1,7 @@ mod editor_models; mod lang; mod shell; +mod task_tracker; use anyhow::Result; use base64::Engine; @@ -39,6 +40,7 @@ use rmcp::object; use self::editor_models::{create_editor_model, EditorModel}; use self::shell::{expand_path, get_shell_config, is_absolute_path, normalize_line_endings}; +use self::task_tracker::TaskTracker; use indoc::indoc; use std::process::Stdio; use std::sync::{Arc, Mutex}; @@ -100,6 +102,7 @@ pub struct DeveloperRouter { file_history: Arc>>>, ignore_patterns: Arc, editor_model: Option, + task_tracker: TaskTracker, } impl Default for DeveloperRouter { @@ -345,6 +348,59 @@ impl DeveloperRouter { open_world_hint: Some(false), }); + let task_tracker_tool = Tool::new( + "task_tracker", + indoc! {r#" + The task tracker tool will help you keep track of tasks + ALWAYS to use this when starting a non trivial activity to help you plan, and when resuming or shifting activities + ALWAYS check the task tracker tasks, and update it as you go + This is an ESSENTIAL tool for breaking down your work into chunks and ensuring it is completed + Check the list often, and update it (one by one) as you complete tasks + When starting out, you SHOULD plan your tasks in advance in very short description for each + + use wip action when you start on one task at a time and done action when finished with it + + for example, + user: "build me a time machine". + task list: "establish a view of quantum physics", "solve causality paradoxes", "research negative energy", "test time machine" + + Each task has a status: to do, wip (work in progress), or done. + Note: Task descriptions must be 200 characters or less. Keep them concise + From time to time you may need to clear-tasks to start fresh + + Actions: + - list: Show all tasks and their status + - add: Add a new task + - wip: Mark a task as work in progress + - done: Mark a task as completed + - clear-tasks: Remove all tasks and start fresh + + Each task has a status: to do, wip (work in progress), or done. + "#}, + object!({ + "type": "object", + "required": ["action"], + "properties": { + "action": { + "type": "string", + "enum": ["list", "add", "wip", "done", "clear-tasks"], + "description": "The action to perform" + }, + "task": { + "type": "string", + "description": "The brief task description (required for add, wip, done actions)" + } + } + }), + ) + .annotate(ToolAnnotations { + title: Some("Task Tracker".to_string()), + read_only_hint: Some(false), + destructive_hint: Some(false), + idempotent_hint: Some(false), + open_world_hint: Some(false), + }); + // Get base instructions and working directory let cwd = std::env::current_dir().expect("should have a current working dir"); let os = std::env::consts::OS; @@ -505,12 +561,14 @@ impl DeveloperRouter { list_windows_tool, screen_capture_tool, image_processor_tool, + task_tracker_tool, ], prompts: Arc::new(load_prompt_files()), instructions, file_history: Arc::new(Mutex::new(HashMap::new())), ignore_patterns: Arc::new(ignore_patterns), editor_model, + task_tracker: TaskTracker::new(), } } @@ -1486,6 +1544,10 @@ impl DeveloperRouter { Content::image(data, "image/png").with_priority(0.0), ]) } + + async fn task_tracker(&self, params: Value) -> Result, ToolError> { + self.task_tracker.handle_request(params).await + } } fn recommend_read_range(path: &Path, total_lines: usize) -> Result, ToolError> { @@ -1532,6 +1594,7 @@ impl Router for DeveloperRouter { "list_windows" => this.list_windows(arguments).await, "screen_capture" => this.screen_capture(arguments).await, "image_processor" => this.image_processor(arguments).await, + "task_tracker" => this.task_tracker(arguments).await, _ => Err(ToolError::NotFound(format!("Tool {} not found", tool_name))), } }) @@ -1589,6 +1652,7 @@ impl Clone for DeveloperRouter { instructions: self.instructions.clone(), file_history: Arc::clone(&self.file_history), ignore_patterns: Arc::clone(&self.ignore_patterns), + task_tracker: self.task_tracker.clone(), editor_model: create_editor_model(), } } @@ -2056,6 +2120,7 @@ mod tests { file_history: Arc::new(Mutex::new(HashMap::new())), ignore_patterns: Arc::new(ignore_patterns), editor_model: None, + task_tracker: TaskTracker::new(), }; // Test basic file matching @@ -2107,6 +2172,7 @@ mod tests { file_history: Arc::new(Mutex::new(HashMap::new())), ignore_patterns: Arc::new(ignore_patterns), editor_model: None, + task_tracker: TaskTracker::new(), }; // Try to write to an ignored file @@ -2167,6 +2233,7 @@ mod tests { file_history: Arc::new(Mutex::new(HashMap::new())), ignore_patterns: Arc::new(ignore_patterns), editor_model: None, + task_tracker: TaskTracker::new(), }; // Create an ignored file diff --git a/crates/goose-mcp/src/developer/task_tracker.rs b/crates/goose-mcp/src/developer/task_tracker.rs new file mode 100644 index 000000000000..b9c0d4fdc2ef --- /dev/null +++ b/crates/goose-mcp/src/developer/task_tracker.rs @@ -0,0 +1,370 @@ +use mcp_core::handler::ToolError; +use rmcp::model::{Content, Role}; +use serde_json::Value; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +#[derive(Debug, Clone, PartialEq)] +pub enum TaskStatus { + Todo, + Wip, + Done, +} + +impl TaskStatus { + fn as_str(&self) -> &'static str { + match self { + TaskStatus::Todo => "to do", + TaskStatus::Wip => "wip", + TaskStatus::Done => "done", + } + } +} + +#[derive(Debug, Clone)] +pub struct TaskTracker { + tasks: Arc>>, +} + +impl TaskTracker { + pub fn new() -> Self { + Self { + tasks: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub fn list_tasks(&self) -> Vec { + let tasks = self.tasks.lock().unwrap(); + let mut task_list = Vec::new(); + for (task, status) in tasks.iter() { + task_list.push(format!("{} [{}]", task, status.as_str())); + } + task_list + } + + pub fn add_task(&self, task: String) -> String { + if task.len() > 200 { + return "Task description is too long (max 200 chars)".to_string(); + } + let mut tasks = self.tasks.lock().unwrap(); + tasks.insert(task.clone(), TaskStatus::Todo); + format!("Added task: {}", task) + } + + pub fn mark_task_wip(&self, task: String) -> String { + let mut tasks = self.tasks.lock().unwrap(); + if tasks.contains_key(&task) { + tasks.insert(task.clone(), TaskStatus::Wip); + format!("Marked as WIP: {}", task) + } else { + format!("Task not found: {}", task) + } + } + + pub fn mark_task_done(&self, task: String) -> String { + let mut tasks = self.tasks.lock().unwrap(); + if tasks.contains_key(&task) { + tasks.insert(task.clone(), TaskStatus::Done); + format!("Marked as done: {}", task) + } else { + format!("Task not found: {}", task) + } + } + + pub fn clear_tasks(&self) -> String { + let mut tasks = self.tasks.lock().unwrap(); + tasks.clear(); + "All tasks cleared".to_string() + } + + pub async fn handle_request(&self, params: Value) -> Result, ToolError> { + let action = + params + .get("action") + .and_then(|v| v.as_str()) + .ok_or(ToolError::InvalidParameters( + "The action string is required".to_string(), + ))?; + + match action { + "list" => { + let tasks = self.list_tasks(); + Ok(vec![ + Content::text(format!("Tasks:\n{}", tasks.join("\n"))) + .with_audience(vec![Role::Assistant]), + Content::text(format!("Tasks:\n{}", tasks.join("\n"))) + .with_audience(vec![Role::User]) + .with_priority(0.0), + ]) + } + "add" => { + let task = params.get("task").and_then(|v| v.as_str()).ok_or( + ToolError::InvalidParameters("The task string is required".to_string()), + )?; + + let result = self.add_task(task.to_string()); + + Ok(vec![Content::text(result)]) + } + "mark_wip" => { + let task = params.get("task").and_then(|v| v.as_str()).ok_or( + ToolError::InvalidParameters("The task string is required".to_string()), + )?; + + let result = self.mark_task_wip(task.to_string()); + + Ok(vec![Content::text(result)]) + } + "mark_done" => { + let task = params.get("task").and_then(|v| v.as_str()).ok_or( + ToolError::InvalidParameters("The task string is required".to_string()), + )?; + + let result = self.mark_task_done(task.to_string()); + + Ok(vec![Content::text(result)]) + } + "clear-tasks" => { + let result = self.clear_tasks(); + + Ok(vec![Content::text(result)]) + } + _ => Err(ToolError::InvalidParameters(format!( + "Unknown action '{}'", + action + ))), + } + } +} + +impl Default for TaskTracker { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_empty_task_list() { + let tracker = TaskTracker::new(); + let tasks = tracker.list_tasks(); + assert!(tasks.is_empty()); + } + + #[tokio::test] + async fn test_add_task() { + let tracker = TaskTracker::new(); + tracker.add_task("Write unit tests".to_string()); + + let tasks = tracker.list_tasks(); + assert_eq!(tasks.len(), 1); + assert!(tasks[0].contains("Write unit tests [to do]")); + } + + #[tokio::test] + async fn test_mark_task_wip() { + let tracker = TaskTracker::new(); + + // Add task first + tracker.add_task("Fix bug".to_string()); + + // Mark as WIP + tracker.mark_task_wip("Fix bug".to_string()); + + let tasks = tracker.list_tasks(); + assert!(tasks[0].contains("Fix bug [wip]")); + } + + #[tokio::test] + async fn test_mark_task_done() { + let tracker = TaskTracker::new(); + + // Add task first + tracker.add_task("Review PR".to_string()); + + // Mark as done + tracker.mark_task_done("Review PR".to_string()); + + let tasks = tracker.list_tasks(); + assert!(tasks[0].contains("Review PR [done]")); + } + + #[tokio::test] + async fn test_clear_tasks() { + let tracker = TaskTracker::new(); + + // Add some tasks + tracker.add_task("Task 1".to_string()); + tracker.add_task("Task 2".to_string()); + + // Clear all + tracker.clear_tasks(); + + let tasks = tracker.list_tasks(); + assert!(tasks.is_empty()); + } + + #[tokio::test] + async fn test_long_task_not_added() { + let tracker = TaskTracker::new(); + let long_task = "a".repeat(201); + + tracker.add_task(long_task); + + let tasks = tracker.list_tasks(); + assert!(tasks.is_empty()); // Task should not be added because it's too long + } + + #[tokio::test] + async fn test_mark_nonexistent_task() { + let tracker = TaskTracker::new(); + + // Try to mark non-existent task as WIP + tracker.mark_task_wip("Non-existent task".to_string()); + + let tasks = tracker.list_tasks(); + assert!(tasks.is_empty()); // No task should exist + } + + #[tokio::test] + async fn test_handle_request_list() { + let tracker = TaskTracker::new(); + tracker.add_task("Test task 1".to_string()); + tracker.add_task("Test task 2".to_string()); + + let params = serde_json::json!({"action": "list"}); + let result = tracker.handle_request(params).await.unwrap(); + + assert_eq!(result.len(), 2); + match &result[0].raw { + rmcp::model::RawContent::Text(text_content) => { + assert!(text_content.text.contains("Test task 1 [to do]")); + assert!(text_content.text.contains("Test task 2 [to do]")); + } + _ => panic!("Expected text content"), + } + } + + #[tokio::test] + async fn test_handle_request_add() { + let tracker = TaskTracker::new(); + + let params = serde_json::json!({"action": "add", "task": "New task"}); + let result = tracker.handle_request(params).await.unwrap(); + + assert_eq!(result.len(), 1); + match &result[0].raw { + rmcp::model::RawContent::Text(text_content) => { + assert_eq!(text_content.text, "Added task: New task"); + } + _ => panic!("Expected text content"), + } + + let tasks = tracker.list_tasks(); + assert_eq!(tasks.len(), 1); + assert!(tasks[0].contains("New task [to do]")); + } + + #[tokio::test] + async fn test_handle_request_mark_wip() { + let tracker = TaskTracker::new(); + tracker.add_task("Test task".to_string()); + + let params = serde_json::json!({"action": "mark_wip", "task": "Test task"}); + let result = tracker.handle_request(params).await.unwrap(); + + assert_eq!(result.len(), 1); + match &result[0].raw { + rmcp::model::RawContent::Text(text_content) => { + assert_eq!(text_content.text, "Marked as WIP: Test task"); + } + _ => panic!("Expected text content"), + } + + let tasks = tracker.list_tasks(); + assert!(tasks[0].contains("Test task [wip]")); + } + + #[tokio::test] + async fn test_handle_request_mark_done() { + let tracker = TaskTracker::new(); + tracker.add_task("Test task".to_string()); + + let params = serde_json::json!({"action": "mark_done", "task": "Test task"}); + let result = tracker.handle_request(params).await.unwrap(); + + assert_eq!(result.len(), 1); + match &result[0].raw { + rmcp::model::RawContent::Text(text_content) => { + assert_eq!(text_content.text, "Marked as done: Test task"); + } + _ => panic!("Expected text content"), + } + + let tasks = tracker.list_tasks(); + assert!(tasks[0].contains("Test task [done]")); + } + + #[tokio::test] + async fn test_handle_request_clear() { + let tracker = TaskTracker::new(); + tracker.add_task("Task 1".to_string()); + tracker.add_task("Task 2".to_string()); + + let params = serde_json::json!({"action": "clear-tasks"}); + let result = tracker.handle_request(params).await.unwrap(); + + assert_eq!(result.len(), 1); + match &result[0].raw { + rmcp::model::RawContent::Text(text_content) => { + assert_eq!(text_content.text, "All tasks cleared"); + } + _ => panic!("Expected text content"), + } + + let tasks = tracker.list_tasks(); + assert!(tasks.is_empty()); + } + + #[tokio::test] + async fn test_handle_request_missing_action() { + let tracker = TaskTracker::new(); + + let params = serde_json::json!({}); + let result = tracker.handle_request(params).await; + + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("The action string is required")); + } + + #[tokio::test] + async fn test_handle_request_unknown_action() { + let tracker = TaskTracker::new(); + + let params = serde_json::json!({"action": "unknown"}); + let result = tracker.handle_request(params).await; + + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("Unknown action")); + } + + #[tokio::test] + async fn test_handle_request_missing_task_parameter() { + let tracker = TaskTracker::new(); + + let params = serde_json::json!({"action": "add"}); + let result = tracker.handle_request(params).await; + + assert!(result.is_err()); + assert!(result + .unwrap_err() + .to_string() + .contains("The task string is required")); + } +}