From d30ba16ab23b0e0e6bf08929bf41d2180411f04e Mon Sep 17 00:00:00 2001 From: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Date: Mon, 16 Jun 2025 09:47:04 -0400 Subject: [PATCH 1/5] Platform scheduler tool initial work (#1) --- crates/goose-server/src/commands/agent.rs | 5 +- crates/goose-server/src/routes/schedule.rs | 10 +- crates/goose/src/agents/agent.rs | 546 +++++++++++++- crates/goose/src/agents/mod.rs | 5 + crates/goose/src/agents/platform_tools.rs | 43 ++ crates/goose/src/agents/private_tests.rs | 802 +++++++++++++++++++++ crates/goose/src/agents/test_support.rs | 413 +++++++++++ crates/goose/src/scheduler_test_support.rs | 21 + 8 files changed, 1842 insertions(+), 3 deletions(-) create mode 100644 crates/goose/src/agents/private_tests.rs create mode 100644 crates/goose/src/agents/test_support.rs create mode 100644 crates/goose/src/scheduler_test_support.rs diff --git a/crates/goose-server/src/commands/agent.rs b/crates/goose-server/src/commands/agent.rs index 619451a7cdff..104e3b8128e6 100644 --- a/crates/goose-server/src/commands/agent.rs +++ b/crates/goose-server/src/commands/agent.rs @@ -29,7 +29,10 @@ pub async fn run() -> Result<()> { .join("schedules.json"); let scheduler_instance = SchedulerFactory::create(schedule_file_path).await?; - app_state.set_scheduler(scheduler_instance).await; + app_state.set_scheduler(scheduler_instance.clone()).await; + + // NEW: Provide scheduler access to the agent + agent_ref.set_scheduler(scheduler_instance).await; let cors = CorsLayer::new() .allow_origin(Any) diff --git a/crates/goose-server/src/routes/schedule.rs b/crates/goose-server/src/routes/schedule.rs index cb1199ded702..9b4433c52ee4 100644 --- a/crates/goose-server/src/routes/schedule.rs +++ b/crates/goose-server/src/routes/schedule.rs @@ -93,6 +93,8 @@ fn parse_session_name_to_iso(session_name: &str) -> String { request_body = CreateScheduleRequest, responses( (status = 200, description = "Scheduled job created successfully", body = ScheduledJob), + (status = 400, description = "Invalid cron expression or recipe file"), + (status = 409, description = "Job ID already exists"), (status = 500, description = "Internal server error") ), tag = "schedule" @@ -128,7 +130,13 @@ async fn create_schedule( .await .map_err(|e| { eprintln!("Error creating schedule: {:?}", e); // Log error - StatusCode::INTERNAL_SERVER_ERROR + match e { + goose::scheduler::SchedulerError::JobNotFound(_) => StatusCode::NOT_FOUND, + goose::scheduler::SchedulerError::CronParseError(_) => StatusCode::BAD_REQUEST, + goose::scheduler::SchedulerError::RecipeLoadError(_) => StatusCode::BAD_REQUEST, + goose::scheduler::SchedulerError::JobIdExists(_) => StatusCode::CONFLICT, + _ => StatusCode::INTERNAL_SERVER_ERROR, + } })?; Ok(Json(job)) } diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 991603007c87..c0d77655cbe1 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -4,6 +4,7 @@ use std::pin::Pin; use std::sync::Arc; use anyhow::{anyhow, Result}; +use chrono::Utc; use futures::stream::BoxStream; use futures::{FutureExt, Stream, TryStreamExt}; use futures_util::stream; @@ -17,6 +18,7 @@ use crate::permission::PermissionConfirmation; use crate::providers::base::Provider; use crate::providers::errors::ProviderError; use crate::recipe::{Author, Recipe}; +use crate::scheduler_trait::SchedulerTrait; use crate::tool_monitor::{ToolCall, ToolMonitor}; use regex::Regex; use serde_json::Value; @@ -27,7 +29,8 @@ use crate::agents::extension::{ExtensionConfig, ExtensionError, ExtensionResult, use crate::agents::extension_manager::{get_parameter_names, ExtensionManager}; use crate::agents::platform_tools::{ PLATFORM_LIST_RESOURCES_TOOL_NAME, PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME, - PLATFORM_READ_RESOURCE_TOOL_NAME, PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME, + PLATFORM_MANAGE_SCHEDULE_TOOL_NAME, PLATFORM_READ_RESOURCE_TOOL_NAME, + PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME, }; use crate::agents::prompt_manager::PromptManager; use crate::agents::router_tool_selector::{ @@ -59,6 +62,7 @@ pub struct Agent { pub(super) tool_result_rx: ToolResultReceiver, pub(super) tool_monitor: Mutex>, pub(super) router_tool_selector: Mutex>>>, + pub(super) scheduler_service: Mutex>>, } #[derive(Clone, Debug)] @@ -85,6 +89,7 @@ impl Agent { tool_result_rx: Arc::new(Mutex::new(tool_rx)), tool_monitor: Mutex::new(None), router_tool_selector: Mutex::new(None), + scheduler_service: Mutex::new(None), } } @@ -103,6 +108,12 @@ impl Agent { monitor.reset(); } } + + /// Set the scheduler service for this agent + pub async fn set_scheduler(&self, scheduler: Arc) { + let mut scheduler_service = self.scheduler_service.lock().await; + *scheduler_service = Some(scheduler); + } } impl Default for Agent { @@ -203,6 +214,13 @@ impl Agent { } } + if tool_call.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME { + let result = self + .handle_schedule_management(tool_call.arguments, request_id.clone()) + .await; + return (request_id, Ok(ToolCallResult::from(result))); + } + if tool_call.name == PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME { let extension_name = tool_call .arguments @@ -362,6 +380,325 @@ impl Agent { (request_id, result) } + pub(crate) async fn handle_schedule_management( + &self, + arguments: serde_json::Value, + _request_id: String, + ) -> ToolResult> { + let scheduler = match self.scheduler_service.lock().await.as_ref() { + Some(s) => s.clone(), + None => { + return Err(ToolError::ExecutionError( + "Scheduler not available. This tool only works in server mode.".to_string(), + )) + } + }; + + let action = arguments + .get("action") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'action' parameter".to_string()))?; + + match action { + "list" => self.handle_list_jobs(scheduler).await, + "create" => self.handle_create_job(scheduler, arguments).await, + "run_now" => self.handle_run_now(scheduler, arguments).await, + "pause" => self.handle_pause_job(scheduler, arguments).await, + "unpause" => self.handle_unpause_job(scheduler, arguments).await, + "delete" => self.handle_delete_job(scheduler, arguments).await, + "kill" => self.handle_kill_job(scheduler, arguments).await, + "inspect" => self.handle_inspect_job(scheduler, arguments).await, + "sessions" => self.handle_list_sessions(scheduler, arguments).await, + _ => Err(ToolError::ExecutionError(format!( + "Unknown action: {}", + action + ))), + } + } + + async fn handle_list_jobs( + &self, + scheduler: Arc, + ) -> ToolResult> { + match scheduler.list_scheduled_jobs().await { + Ok(jobs) => { + let jobs_json = serde_json::to_string_pretty(&jobs).map_err(|e| { + ToolError::ExecutionError(format!("Failed to serialize jobs: {}", e)) + })?; + Ok(vec![Content::text(format!( + "Scheduled Jobs:\n{}", + jobs_json + ))]) + } + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to list jobs: {}", + e + ))), + } + } + + async fn handle_create_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let recipe_path = arguments + .get("recipe_path") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionError("Missing 'recipe_path' parameter".to_string()) + })?; + + let cron_expression = arguments + .get("cron_expression") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionError("Missing 'cron_expression' parameter".to_string()) + })?; + + // Validate recipe file exists and is readable + if !std::path::Path::new(recipe_path).exists() { + return Err(ToolError::ExecutionError(format!( + "Recipe file not found: {}", + recipe_path + ))); + } + + // Validate it's a valid recipe by trying to parse it + match std::fs::read_to_string(recipe_path) { + Ok(content) => { + if recipe_path.ends_with(".json") { + serde_json::from_str::(&content).map_err(|e| { + ToolError::ExecutionError(format!("Invalid JSON recipe: {}", e)) + })?; + } else { + serde_yaml::from_str::(&content).map_err(|e| { + ToolError::ExecutionError(format!("Invalid YAML recipe: {}", e)) + })?; + } + } + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Cannot read recipe file: {}", + e + ))) + } + } + + // Generate unique job ID + let job_id = format!("agent_created_{}", Utc::now().timestamp()); + + let job = crate::scheduler::ScheduledJob { + id: job_id.clone(), + source: recipe_path.to_string(), + cron: cron_expression.to_string(), + last_run: None, + currently_running: false, + paused: false, + current_session_id: None, + process_start_time: None, + }; + + match scheduler.add_scheduled_job(job).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully created scheduled job '{}' for recipe '{}' with cron expression '{}'", + job_id, recipe_path, cron_expression + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to create job: {}", + e + ))), + } + } + + async fn handle_run_now( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.run_now(job_id).await { + Ok(session_id) => Ok(vec![Content::text(format!( + "Successfully started job '{}'. Session ID: {}", + job_id, session_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to run job: {}", + e + ))), + } + } + + async fn handle_pause_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.pause_schedule(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully paused job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to pause job: {}", + e + ))), + } + } + + async fn handle_unpause_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.unpause_schedule(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully unpaused job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to unpause job: {}", + e + ))), + } + } + + async fn handle_delete_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.remove_scheduled_job(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully deleted job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to delete job: {}", + e + ))), + } + } + + async fn handle_kill_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.kill_running_job(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully killed running job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to kill job: {}", + e + ))), + } + } + + async fn handle_inspect_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.get_running_job_info(job_id).await { + Ok(Some((session_id, start_time))) => { + let duration = Utc::now().signed_duration_since(start_time); + Ok(vec![Content::text(format!( + "Job '{}' is currently running:\n- Session ID: {}\n- Started: {}\n- Duration: {} seconds", + job_id, session_id, start_time.to_rfc3339(), duration.num_seconds() + ))]) + } + Ok(None) => Ok(vec![Content::text(format!( + "Job '{}' is not currently running", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to inspect job: {}", + e + ))), + } + } + + async fn handle_list_sessions( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + let limit = arguments + .get("limit") + .and_then(|v| v.as_u64()) + .unwrap_or(50) as usize; + + match scheduler.sessions(job_id, limit).await { + Ok(sessions) => { + if sessions.is_empty() { + Ok(vec![Content::text(format!( + "No sessions found for job '{}'", + job_id + ))]) + } else { + let sessions_info: Vec = sessions + .into_iter() + .map(|(session_name, metadata)| { + format!( + "- Session: {} (Messages: {}, Working Dir: {})", + session_name, + metadata.message_count, + metadata.working_dir.display() + ) + }) + .collect(); + + Ok(vec![Content::text(format!( + "Sessions for job '{}':\n{}", + job_id, + sessions_info.join("\n") + ))]) + } + } + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to list sessions: {}", + e + ))), + } + } + pub async fn add_extension(&self, extension: ExtensionConfig) -> ExtensionResult<()> { match &extension { ExtensionConfig::Frontend { @@ -432,6 +769,7 @@ impl Agent { // Add platform tools prefixed_tools.push(platform_tools::search_available_extensions_tool()); prefixed_tools.push(platform_tools::manage_extensions_tool()); + prefixed_tools.push(platform_tools::manage_schedule_tool()); // Add resource tools if supported if extension_manager.supports_resources() { @@ -986,3 +1324,209 @@ impl Agent { Ok(recipe) } } + +#[cfg(test)] +mod schedule_tool_tests { + use super::*; + use crate::scheduler::{ScheduledJob, SchedulerError}; + use crate::scheduler_trait::SchedulerTrait; + use crate::session::storage::SessionMetadata; + use async_trait::async_trait; + use chrono::{DateTime, Utc}; + use mcp_core::{Content, ToolError}; + use serde_json::json; + use std::sync::Arc; + + // Mock scheduler for testing + struct MockScheduler { + jobs: tokio::sync::Mutex>, + } + + impl MockScheduler { + fn new() -> Self { + Self { + jobs: tokio::sync::Mutex::new(Vec::new()), + } + } + } + + #[async_trait] + impl SchedulerTrait for MockScheduler { + async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError> { + let mut jobs = self.jobs.lock().await; + jobs.push(job); + Ok(()) + } + + async fn list_scheduled_jobs(&self) -> Result, SchedulerError> { + let jobs = self.jobs.lock().await; + Ok(jobs.clone()) + } + + async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> { + let mut jobs = self.jobs.lock().await; + if let Some(pos) = jobs.iter().position(|job| job.id == id) { + jobs.remove(pos); + Ok(()) + } else { + Err(SchedulerError::JobNotFound(id.to_string())) + } + } + + async fn pause_schedule(&self, _id: &str) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn unpause_schedule(&self, _id: &str) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn run_now(&self, _id: &str) -> Result { + Ok("test_session_123".to_string()) + } + + async fn sessions( + &self, + _sched_id: &str, + _limit: usize, + ) -> Result, SchedulerError> { + Ok(vec![]) + } + + async fn update_schedule( + &self, + _sched_id: &str, + _new_cron: String, + ) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn kill_running_job(&self, _sched_id: &str) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn get_running_job_info( + &self, + _sched_id: &str, + ) -> Result)>, SchedulerError> { + Ok(None) + } + } + + #[tokio::test] + async fn test_schedule_management_tool_list() { + let agent = Agent::new(); + let mock_scheduler = Arc::new(MockScheduler::new()); + agent.set_scheduler(mock_scheduler.clone()).await; + + // Test list action + let arguments = json!({ + "action": "list" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Scheduled Jobs:")); + } + } + + #[tokio::test] + async fn test_schedule_management_tool_no_scheduler() { + let agent = Agent::new(); + // Don't set scheduler + + let arguments = json!({ + "action": "list" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Scheduler not available")); + } + } + + #[tokio::test] + async fn test_schedule_management_tool_run_now() { + let agent = Agent::new(); + let mock_scheduler = Arc::new(MockScheduler::new()); + agent.set_scheduler(mock_scheduler.clone()).await; + + let arguments = json!({ + "action": "run_now", + "job_id": "test_job" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Successfully started job 'test_job'")); + assert!(text_content.text.contains("test_session_123")); + } + } + + #[tokio::test] + async fn test_schedule_management_tool_dispatch() { + let agent = Agent::new(); + let mock_scheduler = Arc::new(MockScheduler::new()); + agent.set_scheduler(mock_scheduler.clone()).await; + + // Test that the tool is properly dispatched through dispatch_tool_call + let tool_call = mcp_core::tool::ToolCall { + name: PLATFORM_MANAGE_SCHEDULE_TOOL_NAME.to_string(), + arguments: json!({ + "action": "list" + }), + }; + + let (request_id, result) = agent + .dispatch_tool_call(tool_call, "test_dispatch".to_string()) + .await; + assert_eq!(request_id, "test_dispatch"); + assert!(result.is_ok()); + + let tool_result = result.unwrap(); + // The result should be a future that resolves to the tool output + let output = tool_result.result.await; + assert!(output.is_ok()); + + let content = output.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Scheduled Jobs:")); + } + } + + #[tokio::test] + async fn test_schedule_management_tool_in_list_tools() { + let agent = Agent::new(); + let tools = agent.list_tools(None).await; + + // Check that the schedule management tool is included in the list + let schedule_tool = tools + .iter() + .find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME); + assert!(schedule_tool.is_some()); + + let tool = schedule_tool.unwrap(); + assert!(tool + .description + .contains("Manage scheduled recipe execution")); + } +} diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 24511ac64580..0484d68e9b17 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -4,10 +4,15 @@ pub mod extension; pub mod extension_manager; mod large_response_handler; pub mod platform_tools; +#[cfg(test)] +mod private_tests; pub mod prompt_manager; mod reply_parts; mod router_tool_selector; mod router_tools; +#[cfg(test)] +mod test_support; + mod tool_execution; mod tool_router_index_manager; pub(crate) mod tool_vectordb; diff --git a/crates/goose/src/agents/platform_tools.rs b/crates/goose/src/agents/platform_tools.rs index 3456167cb5da..49165c7e0fd4 100644 --- a/crates/goose/src/agents/platform_tools.rs +++ b/crates/goose/src/agents/platform_tools.rs @@ -7,6 +7,7 @@ pub const PLATFORM_LIST_RESOURCES_TOOL_NAME: &str = "platform__list_resources"; pub const PLATFORM_SEARCH_AVAILABLE_EXTENSIONS_TOOL_NAME: &str = "platform__search_available_extensions"; pub const PLATFORM_MANAGE_EXTENSIONS_TOOL_NAME: &str = "platform__manage_extensions"; +pub const PLATFORM_MANAGE_SCHEDULE_TOOL_NAME: &str = "platform__manage_schedule"; pub fn read_resource_tool() -> Tool { Tool::new( @@ -112,3 +113,45 @@ pub fn manage_extensions_tool() -> Tool { }), ) } + +pub fn manage_schedule_tool() -> Tool { + Tool::new( + PLATFORM_MANAGE_SCHEDULE_TOOL_NAME.to_string(), + indoc! {r#" + Manage scheduled recipe execution for this Goose instance. + + Actions: + - "list": List all scheduled jobs + - "create": Create a new scheduled job from a recipe file + - "run_now": Execute a scheduled job immediately + - "pause": Pause a scheduled job + - "unpause": Resume a paused job + - "delete": Remove a scheduled job + - "kill": Terminate a currently running job + - "inspect": Get details about a running job + - "sessions": List execution history for a job + "#} + .to_string(), + json!({ + "type": "object", + "required": ["action"], + "properties": { + "action": { + "type": "string", + "enum": ["list", "create", "run_now", "pause", "unpause", "delete", "kill", "inspect", "sessions"] + }, + "job_id": {"type": "string", "description": "Job identifier for operations on existing jobs"}, + "recipe_path": {"type": "string", "description": "Path to recipe file for create action"}, + "cron_expression": {"type": "string", "description": "Cron expression for create action"}, + "limit": {"type": "integer", "description": "Limit for sessions list", "default": 50} + } + }), + Some(ToolAnnotations { + title: Some("Manage scheduled recipes".to_string()), + read_only_hint: false, + destructive_hint: true, // Can kill jobs + idempotent_hint: false, + open_world_hint: false, + }), + ) +} diff --git a/crates/goose/src/agents/private_tests.rs b/crates/goose/src/agents/private_tests.rs new file mode 100644 index 000000000000..cde948412a65 --- /dev/null +++ b/crates/goose/src/agents/private_tests.rs @@ -0,0 +1,802 @@ +#![cfg(test)] + +use mcp_core::{Content, ToolError}; +use serde_json::json; + +use crate::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; +use crate::agents::test_support::{ + create_temp_recipe, create_test_session_metadata, MockBehavior, ScheduleToolTestBuilder, +}; + +// Test all actions of the scheduler platform tool +#[tokio::test] +async fn test_schedule_tool_list_action() { + // Create a test builder with existing jobs + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .with_existing_job("job2", "0 0 * * * *") + .await + .build() + .await; + + // Test list action + let arguments = json!({ + "action": "list" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Scheduled Jobs:")); + assert!(text_content.text.contains("job1")); + assert!(text_content.text.contains("job2")); + } else { + panic!("Expected text content"); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"list_scheduled_jobs".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_list_action_empty() { + // Create a test builder with no jobs + let (agent, scheduler) = ScheduleToolTestBuilder::new().build().await; + + // Test list action + let arguments = json!({ + "action": "list" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Scheduled Jobs:")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"list_scheduled_jobs".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_list_action_error() { + // Create a test builder with a list error + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_scheduler_behavior( + "list_scheduled_jobs", + MockBehavior::InternalError("Database error".to_string()), + ) + .await + .build() + .await; + + // Test list action + let arguments = json!({ + "action": "list" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Failed to list jobs")); + assert!(msg.contains("Database error")); + } else { + panic!("Expected ExecutionError"); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"list_scheduled_jobs".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_create_action() { + let (agent, scheduler) = ScheduleToolTestBuilder::new().build().await; + + // Create a temporary recipe file + let temp_recipe = create_temp_recipe(true, "json"); + + // Test create action + let arguments = json!({ + "action": "create", + "recipe_path": temp_recipe.path.to_str().unwrap(), + "cron_expression": "*/5 * * * * *" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Successfully created scheduled job")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"add_scheduled_job".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_create_action_missing_params() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test create action with missing recipe_path + let arguments = json!({ + "action": "create", + "cron_expression": "*/5 * * * * *" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Missing 'recipe_path' parameter")); + } else { + panic!("Expected ExecutionError"); + } + + // Test create action with missing cron_expression + let temp_recipe = create_temp_recipe(true, "json"); + let arguments = json!({ + "action": "create", + "recipe_path": temp_recipe.path.to_str().unwrap() + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Missing 'cron_expression' parameter")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_create_action_nonexistent_recipe() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test create action with nonexistent recipe + let arguments = json!({ + "action": "create", + "recipe_path": "/nonexistent/recipe.json", + "cron_expression": "*/5 * * * * *" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Recipe file not found")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_create_action_invalid_recipe() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Create an invalid recipe file + let temp_recipe = create_temp_recipe(false, "json"); + + // Test create action with invalid recipe + let arguments = json!({ + "action": "create", + "recipe_path": temp_recipe.path.to_str().unwrap(), + "cron_expression": "*/5 * * * * *" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Invalid JSON recipe")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_create_action_scheduler_error() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_scheduler_behavior( + "add_scheduled_job", + MockBehavior::AlreadyExists("job1".to_string()), + ) + .await + .build() + .await; + + // Create a temporary recipe file + let temp_recipe = create_temp_recipe(true, "json"); + + // Test create action + let arguments = json!({ + "action": "create", + "recipe_path": temp_recipe.path.to_str().unwrap(), + "cron_expression": "*/5 * * * * *" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Failed to create job")); + assert!(msg.contains("job1")); + } else { + panic!("Expected ExecutionError"); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"add_scheduled_job".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_run_now_action() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test run_now action + let arguments = json!({ + "action": "run_now", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Successfully started job 'job1'")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"run_now".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_run_now_action_missing_job_id() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test run_now action with missing job_id + let arguments = json!({ + "action": "run_now" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Missing 'job_id' parameter")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_run_now_action_nonexistent_job() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_scheduler_behavior("run_now", MockBehavior::NotFound("nonexistent".to_string())) + .await + .build() + .await; + + // Test run_now action with nonexistent job + let arguments = json!({ + "action": "run_now", + "job_id": "nonexistent" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Failed to run job")); + assert!(msg.contains("nonexistent")); + } else { + panic!("Expected ExecutionError"); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"run_now".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_pause_action() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test pause action + let arguments = json!({ + "action": "pause", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Successfully paused job 'job1'")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"pause_schedule".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_pause_action_missing_job_id() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test pause action with missing job_id + let arguments = json!({ + "action": "pause" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Missing 'job_id' parameter")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_pause_action_running_job() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_scheduler_behavior( + "pause_schedule", + MockBehavior::JobCurrentlyRunning("job1".to_string()), + ) + .await + .build() + .await; + + // Test pause action with a running job + let arguments = json!({ + "action": "pause", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Failed to pause job")); + assert!(msg.contains("job1")); + } else { + panic!("Expected ExecutionError"); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"pause_schedule".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_unpause_action() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test unpause action + let arguments = json!({ + "action": "unpause", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Successfully unpaused job 'job1'")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"unpause_schedule".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_delete_action() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test delete action + let arguments = json!({ + "action": "delete", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Successfully deleted job 'job1'")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"remove_scheduled_job".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_kill_action() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .with_running_job("job1") + .await + .build() + .await; + + // Test kill action + let arguments = json!({ + "action": "kill", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Successfully killed running job 'job1'")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"kill_running_job".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_kill_action_not_running() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test kill action with a job that's not running + let arguments = json!({ + "action": "kill", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Failed to kill job")); + } else { + panic!("Expected ExecutionError"); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"kill_running_job".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_inspect_action_running() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .with_running_job("job1") + .await + .build() + .await; + + // Test inspect action + let arguments = json!({ + "action": "inspect", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Job 'job1' is currently running")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"get_running_job_info".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_inspect_action_not_running() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test inspect action with a job that's not running + let arguments = json!({ + "action": "inspect", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("Job 'job1' is not currently running")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"get_running_job_info".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_sessions_action() { + // Create test session metadata + let sessions = vec![ + ( + "1234567890_session1".to_string(), + create_test_session_metadata(5, "/tmp"), + ), + ( + "0987654321_session2".to_string(), + create_test_session_metadata(10, "/home"), + ), + ]; + + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .with_sessions_data("job1", sessions) + .await + .build() + .await; + + // Test sessions action + let arguments = json!({ + "action": "sessions", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Sessions for job 'job1'")); + assert!(text_content.text.contains("session1")); + assert!(text_content.text.contains("session2")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"sessions".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_sessions_action_with_limit() { + // Create test session metadata + let sessions = vec![ + ( + "1234567890_session1".to_string(), + create_test_session_metadata(5, "/tmp"), + ), + ( + "0987654321_session2".to_string(), + create_test_session_metadata(10, "/home"), + ), + ( + "5555555555_session3".to_string(), + create_test_session_metadata(15, "/usr"), + ), + ]; + + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .with_sessions_data("job1", sessions) + .await + .build() + .await; + + // Test sessions action with limit + let arguments = json!({ + "action": "sessions", + "job_id": "job1", + "limit": 2 + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"sessions".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_sessions_action_empty() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test sessions action with no sessions + let arguments = json!({ + "action": "sessions", + "job_id": "job1" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_ok()); + + let content = result.unwrap(); + assert_eq!(content.len(), 1); + if let Content::Text(text_content) = &content[0] { + assert!(text_content + .text + .contains("No sessions found for job 'job1'")); + } + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"sessions".to_string())); +} + +#[tokio::test] +async fn test_schedule_tool_unknown_action() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test unknown action + let arguments = json!({ + "action": "unknown_action" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Unknown action")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_dispatch() { + let (agent, scheduler) = ScheduleToolTestBuilder::new() + .with_existing_job("job1", "*/5 * * * * *") + .await + .build() + .await; + + // Test that the tool is properly dispatched through dispatch_tool_call + let tool_call = mcp_core::tool::ToolCall { + name: PLATFORM_MANAGE_SCHEDULE_TOOL_NAME.to_string(), + arguments: json!({ + "action": "list" + }), + }; + + let (request_id, result) = agent + .dispatch_tool_call(tool_call, "test_dispatch".to_string()) + .await; + assert_eq!(request_id, "test_dispatch"); + assert!(result.is_ok()); + + let tool_result = result.unwrap(); + // The result should be a future that resolves to the tool output + let output = tool_result.result.await; + assert!(output.is_ok()); + + // Verify the scheduler was called + let calls = scheduler.get_calls().await; + assert!(calls.contains(&"list_scheduled_jobs".to_string())); +} diff --git a/crates/goose/src/agents/test_support.rs b/crates/goose/src/agents/test_support.rs new file mode 100644 index 000000000000..48a6ce332f51 --- /dev/null +++ b/crates/goose/src/agents/test_support.rs @@ -0,0 +1,413 @@ +#![cfg(test)] + +use std::collections::{HashMap, HashSet}; +use std::path::PathBuf; +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use tempfile::TempDir; +use tokio::sync::Mutex; + +use crate::agents::Agent; +use crate::scheduler::{ScheduledJob, SchedulerError}; +use crate::scheduler_trait::SchedulerTrait; +use crate::session::storage::SessionMetadata; + +#[derive(Debug, Clone)] +pub enum MockBehavior { + Success, + NotFound(String), + AlreadyExists(String), + InternalError(String), + JobCurrentlyRunning(String), +} + +#[derive(Clone)] +pub struct ConfigurableMockScheduler { + jobs: Arc>>, + running_jobs: Arc>>, + call_log: Arc>>, + behaviors: Arc>>, + sessions_data: Arc>>>, +} + +impl ConfigurableMockScheduler { + pub fn new() -> Self { + Self { + jobs: Arc::new(Mutex::new(HashMap::new())), + running_jobs: Arc::new(Mutex::new(HashSet::new())), + call_log: Arc::new(Mutex::new(Vec::new())), + behaviors: Arc::new(Mutex::new(HashMap::new())), + sessions_data: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub async fn with_behavior(self, method: &str, behavior: MockBehavior) -> Self { + self.behaviors + .lock() + .await + .insert(method.to_string(), behavior); + self + } + + pub async fn with_existing_job(self, job: ScheduledJob) -> Self { + self.jobs.lock().await.insert(job.id.clone(), job); + self + } + + pub async fn with_running_job(self, job_id: &str) -> Self { + self.running_jobs.lock().await.insert(job_id.to_string()); + self + } + + pub async fn with_sessions_data( + self, + job_id: &str, + sessions: Vec<(String, SessionMetadata)>, + ) -> Self { + self.sessions_data + .lock() + .await + .insert(job_id.to_string(), sessions); + self + } + + pub async fn get_calls(&self) -> Vec { + self.call_log.lock().await.clone() + } + + async fn log_call(&self, method: &str) { + self.call_log.lock().await.push(method.to_string()); + } + + async fn get_behavior(&self, method: &str) -> MockBehavior { + self.behaviors + .lock() + .await + .get(method) + .cloned() + .unwrap_or(MockBehavior::Success) + } +} + +#[async_trait] +impl SchedulerTrait for ConfigurableMockScheduler { + async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError> { + self.log_call("add_scheduled_job").await; + + match self.get_behavior("add_scheduled_job").await { + MockBehavior::Success => { + let mut jobs = self.jobs.lock().await; + if jobs.contains_key(&job.id) { + return Err(SchedulerError::JobIdExists(job.id)); + } + jobs.insert(job.id.clone(), job); + Ok(()) + } + MockBehavior::AlreadyExists(id) => Err(SchedulerError::JobIdExists(id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(()), + } + } + + async fn list_scheduled_jobs(&self) -> Result, SchedulerError> { + self.log_call("list_scheduled_jobs").await; + + match self.get_behavior("list_scheduled_jobs").await { + MockBehavior::Success => { + let jobs = self.jobs.lock().await; + Ok(jobs.values().cloned().collect()) + } + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(vec![]), + } + } + + async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> { + self.log_call("remove_scheduled_job").await; + + match self.get_behavior("remove_scheduled_job").await { + MockBehavior::Success => { + let mut jobs = self.jobs.lock().await; + if jobs.remove(id).is_some() { + Ok(()) + } else { + Err(SchedulerError::JobNotFound(id.to_string())) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(()), + } + } + + async fn pause_schedule(&self, id: &str) -> Result<(), SchedulerError> { + self.log_call("pause_schedule").await; + + match self.get_behavior("pause_schedule").await { + MockBehavior::Success => { + let jobs = self.jobs.lock().await; + if jobs.contains_key(id) { + Ok(()) + } else { + Err(SchedulerError::JobNotFound(id.to_string())) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::JobCurrentlyRunning(job_id) => { + Err(SchedulerError::AnyhowError(anyhow::anyhow!( + "Cannot pause schedule '{}' while it's currently running", + job_id + ))) + } + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(()), + } + } + + async fn unpause_schedule(&self, id: &str) -> Result<(), SchedulerError> { + self.log_call("unpause_schedule").await; + + match self.get_behavior("unpause_schedule").await { + MockBehavior::Success => { + let jobs = self.jobs.lock().await; + if jobs.contains_key(id) { + Ok(()) + } else { + Err(SchedulerError::JobNotFound(id.to_string())) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(()), + } + } + + async fn run_now(&self, id: &str) -> Result { + self.log_call("run_now").await; + + match self.get_behavior("run_now").await { + MockBehavior::Success => { + let jobs = self.jobs.lock().await; + if jobs.contains_key(id) { + Ok(format!("{}_session_{}", id, chrono::Utc::now().timestamp())) + } else { + Err(SchedulerError::JobNotFound(id.to_string())) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok("mock_session_123".to_string()), + } + } + + async fn sessions( + &self, + sched_id: &str, + limit: usize, + ) -> Result, SchedulerError> { + self.log_call("sessions").await; + + match self.get_behavior("sessions").await { + MockBehavior::Success => { + let sessions_data = self.sessions_data.lock().await; + let sessions = sessions_data.get(sched_id).cloned().unwrap_or_default(); + Ok(sessions.into_iter().take(limit).collect()) + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(vec![]), + } + } + + async fn update_schedule( + &self, + sched_id: &str, + _new_cron: String, + ) -> Result<(), SchedulerError> { + self.log_call("update_schedule").await; + + match self.get_behavior("update_schedule").await { + MockBehavior::Success => { + let jobs = self.jobs.lock().await; + if jobs.contains_key(sched_id) { + Ok(()) + } else { + Err(SchedulerError::JobNotFound(sched_id.to_string())) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(()), + } + } + + async fn kill_running_job(&self, sched_id: &str) -> Result<(), SchedulerError> { + self.log_call("kill_running_job").await; + + match self.get_behavior("kill_running_job").await { + MockBehavior::Success => { + let running_jobs = self.running_jobs.lock().await; + if running_jobs.contains(sched_id) { + Ok(()) + } else { + Err(SchedulerError::AnyhowError(anyhow::anyhow!( + "Schedule '{}' is not currently running", + sched_id + ))) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(()), + } + } + + async fn get_running_job_info( + &self, + sched_id: &str, + ) -> Result)>, SchedulerError> { + self.log_call("get_running_job_info").await; + + match self.get_behavior("get_running_job_info").await { + MockBehavior::Success => { + let running_jobs = self.running_jobs.lock().await; + if running_jobs.contains(sched_id) { + Ok(Some((format!("{}_session", sched_id), Utc::now()))) + } else { + Ok(None) + } + } + MockBehavior::NotFound(job_id) => Err(SchedulerError::JobNotFound(job_id)), + MockBehavior::InternalError(msg) => Err(SchedulerError::SchedulerInternalError(msg)), + _ => Ok(None), + } + } +} + +// Helper for creating temp recipe files +pub struct TempRecipe { + pub path: PathBuf, + _temp_dir: TempDir, // Keep alive +} + +pub fn create_temp_recipe(valid: bool, format: &str) -> TempRecipe { + let temp_dir = tempfile::tempdir().unwrap(); + let filename = format!("test_recipe.{}", format); + let path = temp_dir.path().join(filename); + + let content = if valid { + match format { + "json" => { + r#"{ + "version": "1.0.0", + "title": "Test Recipe", + "description": "A test recipe", + "prompt": "Hello world" +}"# + } + "yaml" | "yml" => { + r#"version: "1.0.0" +title: "Test Recipe" +description: "A test recipe" +prompt: "Hello world" +"# + } + _ => panic!("Unsupported format: {}", format), + } + } else { + match format { + "json" => r#"{"invalid": json syntax"#, + "yaml" | "yml" => "invalid:\n - yaml: syntax: error", + _ => "invalid content", + } + }; + + std::fs::write(&path, content).unwrap(); + TempRecipe { + path, + _temp_dir: temp_dir, + } +} + +// Test builder for easy setup +pub struct ScheduleToolTestBuilder { + scheduler: Arc, +} + +impl ScheduleToolTestBuilder { + pub fn new() -> Self { + Self { + scheduler: Arc::new(ConfigurableMockScheduler::new()), + } + } + + pub async fn with_scheduler_behavior(self, method: &str, behavior: MockBehavior) -> Self { + { + let mut behaviors = self.scheduler.behaviors.lock().await; + behaviors.insert(method.to_string(), behavior); + } + self + } + + pub async fn with_existing_job(self, job_id: &str, cron: &str) -> Self { + let job = ScheduledJob { + id: job_id.to_string(), + source: "/tmp/test.json".to_string(), + cron: cron.to_string(), + last_run: None, + currently_running: false, + paused: false, + current_session_id: None, + process_start_time: None, + }; + { + let mut jobs = self.scheduler.jobs.lock().await; + jobs.insert(job.id.clone(), job); + } + self + } + + pub async fn with_running_job(self, job_id: &str) -> Self { + { + let mut running_jobs = self.scheduler.running_jobs.lock().await; + running_jobs.insert(job_id.to_string()); + } + self + } + + pub async fn with_sessions_data( + self, + job_id: &str, + sessions: Vec<(String, SessionMetadata)>, + ) -> Self { + { + let mut sessions_data = self.scheduler.sessions_data.lock().await; + sessions_data.insert(job_id.to_string(), sessions); + } + self + } + + pub async fn build(self) -> (Agent, Arc) { + let agent = Agent::new(); + agent.set_scheduler(self.scheduler.clone()).await; + (agent, self.scheduler) + } +} + +// Helper function to create test session metadata +pub fn create_test_session_metadata(message_count: usize, working_dir: &str) -> SessionMetadata { + SessionMetadata { + message_count, + working_dir: PathBuf::from(working_dir), + description: "Test session".to_string(), + schedule_id: Some("test_job".to_string()), + total_tokens: Some(100), + input_tokens: Some(50), + output_tokens: Some(50), + accumulated_total_tokens: Some(100), + accumulated_input_tokens: Some(50), + accumulated_output_tokens: Some(50), + } +} diff --git a/crates/goose/src/scheduler_test_support.rs b/crates/goose/src/scheduler_test_support.rs new file mode 100644 index 000000000000..0b5c9f82dd90 --- /dev/null +++ b/crates/goose/src/scheduler_test_support.rs @@ -0,0 +1,21 @@ +//! Test-only utilities for the scheduler +#![cfg(test)] + +use once_cell::sync::Lazy; +use std::sync::Arc; +use tokio::sync::Mutex; + +use crate::providers::base::Provider as GooseProvider; + +static TEST_PROVIDER: Lazy>>> = Lazy::new(|| Mutex::new(None)); + +/// Register a default provider for scheduler job executions when running under tests. +/// The provider will be used by [`Scheduler`] when no provider_override is supplied. +pub async fn set_test_provider(p: Arc) { + let mut guard = TEST_PROVIDER.lock().await; + *guard = Some(p); +} + +pub async fn get_test_provider() -> Option> { + TEST_PROVIDER.lock().await.clone() +} From 23cf2a80465cf84ca77ddc26f7e936b10e72adb9 Mon Sep 17 00:00:00 2001 From: tlongwell-block <109685178+tlongwell-block@users.noreply.github.com> Date: Wed, 18 Jun 2025 11:12:14 -0400 Subject: [PATCH 2/5] Update platform_tools.rs --- crates/goose/src/agents/platform_tools.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/goose/src/agents/platform_tools.rs b/crates/goose/src/agents/platform_tools.rs index 49165c7e0fd4..49acda3b7f23 100644 --- a/crates/goose/src/agents/platform_tools.rs +++ b/crates/goose/src/agents/platform_tools.rs @@ -142,7 +142,7 @@ pub fn manage_schedule_tool() -> Tool { }, "job_id": {"type": "string", "description": "Job identifier for operations on existing jobs"}, "recipe_path": {"type": "string", "description": "Path to recipe file for create action"}, - "cron_expression": {"type": "string", "description": "Cron expression for create action"}, + "cron_expression": {"type": "string", "description": "A six field cron expression for create action"}, "limit": {"type": "integer", "description": "Limit for sessions list", "default": 50} } }), From 0931e4e76e77bb890d705319a171e31f12b71756 Mon Sep 17 00:00:00 2001 From: Tyler Longwell Date: Wed, 18 Jun 2025 14:09:49 -0400 Subject: [PATCH 3/5] initial work on session_content action --- crates/goose/src/agents/agent.rs | 74 +++++++++++++++++ crates/goose/src/agents/platform_tools.rs | 6 +- crates/goose/src/agents/private_tests.rs | 96 +++++++++++++++++++++++ 3 files changed, 174 insertions(+), 2 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index f729dba22090..efa317726bce 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -429,6 +429,7 @@ impl Agent { "kill" => self.handle_kill_job(scheduler, arguments).await, "inspect" => self.handle_inspect_job(scheduler, arguments).await, "sessions" => self.handle_list_sessions(scheduler, arguments).await, + "session_content" => self.handle_session_content(arguments).await, _ => Err(ToolError::ExecutionError(format!( "Unknown action: {}", action @@ -719,6 +720,79 @@ impl Agent { } } + async fn handle_session_content( + &self, + arguments: serde_json::Value, + ) -> ToolResult> { + let session_id = arguments + .get("session_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionError("Missing 'session_id' parameter".to_string()) + })?; + + // Get the session file path + let session_path = crate::session::storage::get_path( + crate::session::storage::Identifier::Name(session_id.to_string()), + ); + + // Check if session file exists + if !session_path.exists() { + return Err(ToolError::ExecutionError(format!( + "Session '{}' not found", + session_id + ))); + } + + // Read session metadata + let metadata = match crate::session::storage::read_metadata(&session_path) { + Ok(metadata) => metadata, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to read session metadata: {}", + e + ))); + } + }; + + // Read session messages + let messages = match crate::session::storage::read_messages(&session_path) { + Ok(messages) => messages, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to read session messages: {}", + e + ))); + } + }; + + // Format the response with metadata and messages + let metadata_json = match serde_json::to_string_pretty(&metadata) { + Ok(json) => json, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to serialize metadata: {}", + e + ))); + } + }; + + let messages_json = match serde_json::to_string_pretty(&messages) { + Ok(json) => json, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to serialize messages: {}", + e + ))); + } + }; + + Ok(vec![Content::text(format!( + "Session '{}' Content:\n\nMetadata:\n{}\n\nMessages:\n{}", + session_id, metadata_json, messages_json + ))]) + } + pub async fn add_extension(&self, extension: ExtensionConfig) -> ExtensionResult<()> { match &extension { ExtensionConfig::Frontend { diff --git a/crates/goose/src/agents/platform_tools.rs b/crates/goose/src/agents/platform_tools.rs index 49acda3b7f23..01f59f958fd8 100644 --- a/crates/goose/src/agents/platform_tools.rs +++ b/crates/goose/src/agents/platform_tools.rs @@ -130,6 +130,7 @@ pub fn manage_schedule_tool() -> Tool { - "kill": Terminate a currently running job - "inspect": Get details about a running job - "sessions": List execution history for a job + - "session_content": Get the full content (messages) of a specific session "#} .to_string(), json!({ @@ -138,12 +139,13 @@ pub fn manage_schedule_tool() -> Tool { "properties": { "action": { "type": "string", - "enum": ["list", "create", "run_now", "pause", "unpause", "delete", "kill", "inspect", "sessions"] + "enum": ["list", "create", "run_now", "pause", "unpause", "delete", "kill", "inspect", "sessions", "session_content"] }, "job_id": {"type": "string", "description": "Job identifier for operations on existing jobs"}, "recipe_path": {"type": "string", "description": "Path to recipe file for create action"}, "cron_expression": {"type": "string", "description": "A six field cron expression for create action"}, - "limit": {"type": "integer", "description": "Limit for sessions list", "default": 50} + "limit": {"type": "integer", "description": "Limit for sessions list", "default": 50}, + "session_id": {"type": "string", "description": "Session identifier for session_content action"} } }), Some(ToolAnnotations { diff --git a/crates/goose/src/agents/private_tests.rs b/crates/goose/src/agents/private_tests.rs index cde948412a65..1402bb046fb9 100644 --- a/crates/goose/src/agents/private_tests.rs +++ b/crates/goose/src/agents/private_tests.rs @@ -748,6 +748,102 @@ async fn test_schedule_tool_sessions_action_empty() { assert!(calls.contains(&"sessions".to_string())); } +#[tokio::test] +async fn test_schedule_tool_session_content_action() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test with a non-existent session + let arguments = json!({ + "action": "session_content", + "session_id": "non_existent_session" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Session 'non_existent_session' not found")); + } else { + panic!("Expected ExecutionError"); + } +} + +#[tokio::test] +async fn test_schedule_tool_session_content_action_with_real_session() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Create a temporary session file in the proper session directory + let session_dir = crate::session::storage::ensure_session_dir().unwrap(); + let session_id = "test_session_real"; + let session_path = session_dir.join(format!("{}.jsonl", session_id)); + + // Create test metadata and messages + let metadata = create_test_session_metadata(2, "/tmp"); + let messages = vec![ + crate::message::Message::user().with_text("Hello"), + crate::message::Message::assistant().with_text("Hi there!"), + ]; + + // Save the session file + crate::session::storage::save_messages_with_metadata(&session_path, &metadata, &messages) + .unwrap(); + + // Test the session_content action + let arguments = json!({ + "action": "session_content", + "session_id": session_id + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + + // Clean up the test session file + let _ = std::fs::remove_file(&session_path); + + // Verify the result + assert!(result.is_ok()); + + if let Ok(content) = result { + assert_eq!(content.len(), 1); + if let mcp_core::Content::Text(text_content) = &content[0] { + assert!(text_content.text.contains("Session 'test_session_real' Content:")); + assert!(text_content.text.contains("Metadata:")); + assert!(text_content.text.contains("Messages:")); + assert!(text_content.text.contains("Hello")); + assert!(text_content.text.contains("Hi there!")); + assert!(text_content.text.contains("Test session")); + } else { + panic!("Expected text content"); + } + } else { + panic!("Expected successful result"); + } +} + +#[tokio::test] +async fn test_schedule_tool_session_content_action_missing_session_id() { + let (agent, _) = ScheduleToolTestBuilder::new().build().await; + + // Test session_content action with missing session_id + let arguments = json!({ + "action": "session_content" + }); + + let result = agent + .handle_schedule_management(arguments, "test_req".to_string()) + .await; + assert!(result.is_err()); + + if let Err(ToolError::ExecutionError(msg)) = result { + assert!(msg.contains("Missing 'session_id' parameter")); + } else { + panic!("Expected ExecutionError"); + } +} + #[tokio::test] async fn test_schedule_tool_unknown_action() { let (agent, _) = ScheduleToolTestBuilder::new().build().await; From 8a53c109c7749ac533cc5dafbac1937106762b80 Mon Sep 17 00:00:00 2001 From: Tyler Longwell Date: Wed, 18 Jun 2025 18:23:13 -0400 Subject: [PATCH 4/5] move testing to the appropriate place --- crates/goose/src/agents/agent.rs | 210 +----------------- crates/goose/src/agents/mod.rs | 4 - crates/goose/tests/agent.rs | 180 +++++++++++++++ .../{src/agents => tests}/private_tests.rs | 13 +- .../{src => tests}/scheduler_test_support.rs | 2 +- .../{src/agents => tests}/test_support.rs | 8 +- 6 files changed, 194 insertions(+), 223 deletions(-) rename crates/goose/{src/agents => tests}/private_tests.rs (98%) rename crates/goose/{src => tests}/scheduler_test_support.rs (92%) rename crates/goose/{src/agents => tests}/test_support.rs (98%) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index efa317726bce..c02d47f24079 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -196,7 +196,7 @@ impl Agent { /// Dispatch a single tool call to the appropriate client #[instrument(skip(self, tool_call, request_id), fields(input, output))] - pub(super) async fn dispatch_tool_call( + pub async fn dispatch_tool_call( &self, tool_call: mcp_core::tool::ToolCall, request_id: String, @@ -400,7 +400,7 @@ impl Agent { (request_id, result) } - pub(crate) async fn handle_schedule_management( + pub async fn handle_schedule_management( &self, arguments: serde_json::Value, _request_id: String, @@ -1463,209 +1463,3 @@ impl Agent { Ok(recipe) } } - -#[cfg(test)] -mod schedule_tool_tests { - use super::*; - use crate::scheduler::{ScheduledJob, SchedulerError}; - use crate::scheduler_trait::SchedulerTrait; - use crate::session::storage::SessionMetadata; - use async_trait::async_trait; - use chrono::{DateTime, Utc}; - use mcp_core::{Content, ToolError}; - use serde_json::json; - use std::sync::Arc; - - // Mock scheduler for testing - struct MockScheduler { - jobs: tokio::sync::Mutex>, - } - - impl MockScheduler { - fn new() -> Self { - Self { - jobs: tokio::sync::Mutex::new(Vec::new()), - } - } - } - - #[async_trait] - impl SchedulerTrait for MockScheduler { - async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError> { - let mut jobs = self.jobs.lock().await; - jobs.push(job); - Ok(()) - } - - async fn list_scheduled_jobs(&self) -> Result, SchedulerError> { - let jobs = self.jobs.lock().await; - Ok(jobs.clone()) - } - - async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> { - let mut jobs = self.jobs.lock().await; - if let Some(pos) = jobs.iter().position(|job| job.id == id) { - jobs.remove(pos); - Ok(()) - } else { - Err(SchedulerError::JobNotFound(id.to_string())) - } - } - - async fn pause_schedule(&self, _id: &str) -> Result<(), SchedulerError> { - Ok(()) - } - - async fn unpause_schedule(&self, _id: &str) -> Result<(), SchedulerError> { - Ok(()) - } - - async fn run_now(&self, _id: &str) -> Result { - Ok("test_session_123".to_string()) - } - - async fn sessions( - &self, - _sched_id: &str, - _limit: usize, - ) -> Result, SchedulerError> { - Ok(vec![]) - } - - async fn update_schedule( - &self, - _sched_id: &str, - _new_cron: String, - ) -> Result<(), SchedulerError> { - Ok(()) - } - - async fn kill_running_job(&self, _sched_id: &str) -> Result<(), SchedulerError> { - Ok(()) - } - - async fn get_running_job_info( - &self, - _sched_id: &str, - ) -> Result)>, SchedulerError> { - Ok(None) - } - } - - #[tokio::test] - async fn test_schedule_management_tool_list() { - let agent = Agent::new(); - let mock_scheduler = Arc::new(MockScheduler::new()); - agent.set_scheduler(mock_scheduler.clone()).await; - - // Test list action - let arguments = json!({ - "action": "list" - }); - - let result = agent - .handle_schedule_management(arguments, "test_req".to_string()) - .await; - assert!(result.is_ok()); - - let content = result.unwrap(); - assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { - assert!(text_content.text.contains("Scheduled Jobs:")); - } - } - - #[tokio::test] - async fn test_schedule_management_tool_no_scheduler() { - let agent = Agent::new(); - // Don't set scheduler - - let arguments = json!({ - "action": "list" - }); - - let result = agent - .handle_schedule_management(arguments, "test_req".to_string()) - .await; - assert!(result.is_err()); - - if let Err(ToolError::ExecutionError(msg)) = result { - assert!(msg.contains("Scheduler not available")); - } - } - - #[tokio::test] - async fn test_schedule_management_tool_run_now() { - let agent = Agent::new(); - let mock_scheduler = Arc::new(MockScheduler::new()); - agent.set_scheduler(mock_scheduler.clone()).await; - - let arguments = json!({ - "action": "run_now", - "job_id": "test_job" - }); - - let result = agent - .handle_schedule_management(arguments, "test_req".to_string()) - .await; - assert!(result.is_ok()); - - let content = result.unwrap(); - assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { - assert!(text_content - .text - .contains("Successfully started job 'test_job'")); - assert!(text_content.text.contains("test_session_123")); - } - } - - #[tokio::test] - async fn test_schedule_management_tool_dispatch() { - let agent = Agent::new(); - let mock_scheduler = Arc::new(MockScheduler::new()); - agent.set_scheduler(mock_scheduler.clone()).await; - - // Test that the tool is properly dispatched through dispatch_tool_call - let tool_call = mcp_core::tool::ToolCall { - name: PLATFORM_MANAGE_SCHEDULE_TOOL_NAME.to_string(), - arguments: json!({ - "action": "list" - }), - }; - - let (request_id, result) = agent - .dispatch_tool_call(tool_call, "test_dispatch".to_string()) - .await; - assert_eq!(request_id, "test_dispatch"); - assert!(result.is_ok()); - - let tool_result = result.unwrap(); - // The result should be a future that resolves to the tool output - let output = tool_result.result.await; - assert!(output.is_ok()); - - let content = output.unwrap(); - assert_eq!(content.len(), 1); - if let Content::Text(text_content) = &content[0] { - assert!(text_content.text.contains("Scheduled Jobs:")); - } - } - - #[tokio::test] - async fn test_schedule_management_tool_in_list_tools() { - let agent = Agent::new(); - let tools = agent.list_tools(None).await; - - // Check that the schedule management tool is included in the list - let schedule_tool = tools - .iter() - .find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME); - assert!(schedule_tool.is_some()); - - let tool = schedule_tool.unwrap(); - assert!(tool - .description - .contains("Manage scheduled recipe execution")); - } -} diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 0484d68e9b17..94b3bd0e971e 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -4,14 +4,10 @@ pub mod extension; pub mod extension_manager; mod large_response_handler; pub mod platform_tools; -#[cfg(test)] -mod private_tests; pub mod prompt_manager; mod reply_parts; mod router_tool_selector; mod router_tools; -#[cfg(test)] -mod test_support; mod tool_execution; mod tool_router_index_manager; diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index eb7ef2d59afa..f3785c73c55b 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -330,3 +330,183 @@ mod tests { .await } } + +#[cfg(test)] +mod schedule_tool_tests { + use super::*; + use goose::scheduler::{ScheduledJob, SchedulerError}; + use goose::scheduler_trait::SchedulerTrait; + use goose::session::storage::SessionMetadata; + use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; + use async_trait::async_trait; + use chrono::{DateTime, Utc}; + use std::sync::Arc; + + // Mock scheduler for testing + struct MockScheduler { + jobs: tokio::sync::Mutex>, + } + + impl MockScheduler { + fn new() -> Self { + Self { + jobs: tokio::sync::Mutex::new(Vec::new()), + } + } + } + + #[async_trait] + impl SchedulerTrait for MockScheduler { + async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError> { + let mut jobs = self.jobs.lock().await; + jobs.push(job); + Ok(()) + } + + async fn list_scheduled_jobs(&self) -> Result, SchedulerError> { + let jobs = self.jobs.lock().await; + Ok(jobs.clone()) + } + + async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> { + let mut jobs = self.jobs.lock().await; + if let Some(pos) = jobs.iter().position(|job| job.id == id) { + jobs.remove(pos); + Ok(()) + } else { + Err(SchedulerError::JobNotFound(id.to_string())) + } + } + + async fn pause_schedule(&self, _id: &str) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn unpause_schedule(&self, _id: &str) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn run_now(&self, _id: &str) -> Result { + Ok("test_session_123".to_string()) + } + + async fn sessions( + &self, + _sched_id: &str, + _limit: usize, + ) -> Result, SchedulerError> { + Ok(vec![]) + } + + async fn update_schedule( + &self, + _sched_id: &str, + _new_cron: String, + ) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn kill_running_job(&self, _sched_id: &str) -> Result<(), SchedulerError> { + Ok(()) + } + + async fn get_running_job_info( + &self, + _sched_id: &str, + ) -> Result)>, SchedulerError> { + Ok(None) + } + } + + #[tokio::test] + async fn test_schedule_management_tool_list() { + let agent = Agent::new(); + let mock_scheduler = Arc::new(MockScheduler::new()); + agent.set_scheduler(mock_scheduler.clone()).await; + + // Test that the schedule management tool is available in the tools list + let tools = agent.list_tools(None).await; + let schedule_tool = tools + .iter() + .find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME); + assert!(schedule_tool.is_some()); + + let tool = schedule_tool.unwrap(); + assert!(tool + .description + .contains("Manage scheduled recipe execution")); + } + + #[tokio::test] + async fn test_schedule_management_tool_no_scheduler() { + let agent = Agent::new(); + // Don't set scheduler - test that the tool still appears in the list + // but would fail if actually called (which we can't test directly through public API) + + let tools = agent.list_tools(None).await; + let schedule_tool = tools + .iter() + .find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME); + assert!(schedule_tool.is_some()); + } + + #[tokio::test] + async fn test_schedule_management_tool_in_platform_tools() { + let agent = Agent::new(); + let tools = agent.list_tools(Some("platform".to_string())).await; + + // Check that the schedule management tool is included in platform tools + let schedule_tool = tools + .iter() + .find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME); + assert!(schedule_tool.is_some()); + + let tool = schedule_tool.unwrap(); + assert!(tool + .description + .contains("Manage scheduled recipe execution")); + + // Verify the tool has the expected actions in its schema + if let Some(properties) = tool.input_schema.get("properties") { + if let Some(action_prop) = properties.get("action") { + if let Some(enum_values) = action_prop.get("enum") { + let actions: Vec = enum_values + .as_array() + .unwrap() + .iter() + .map(|v| v.as_str().unwrap().to_string()) + .collect(); + + // Check that our session_content action is included + assert!(actions.contains(&"session_content".to_string())); + assert!(actions.contains(&"list".to_string())); + assert!(actions.contains(&"create".to_string())); + assert!(actions.contains(&"sessions".to_string())); + } + } + } + } + + #[tokio::test] + async fn test_schedule_management_tool_schema_validation() { + let agent = Agent::new(); + let tools = agent.list_tools(None).await; + let schedule_tool = tools + .iter() + .find(|tool| tool.name == PLATFORM_MANAGE_SCHEDULE_TOOL_NAME); + assert!(schedule_tool.is_some()); + + let tool = schedule_tool.unwrap(); + + // Verify the tool schema has the session_id parameter for session_content action + if let Some(properties) = tool.input_schema.get("properties") { + assert!(properties.get("session_id").is_some()); + + if let Some(session_id_prop) = properties.get("session_id") { + assert_eq!(session_id_prop.get("type").unwrap().as_str().unwrap(), "string"); + assert!(session_id_prop.get("description").unwrap().as_str().unwrap() + .contains("Session identifier for session_content action")); + } + } + } +} diff --git a/crates/goose/src/agents/private_tests.rs b/crates/goose/tests/private_tests.rs similarity index 98% rename from crates/goose/src/agents/private_tests.rs rename to crates/goose/tests/private_tests.rs index 1402bb046fb9..b94b04767d39 100644 --- a/crates/goose/src/agents/private_tests.rs +++ b/crates/goose/tests/private_tests.rs @@ -3,8 +3,9 @@ use mcp_core::{Content, ToolError}; use serde_json::json; -use crate::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; -use crate::agents::test_support::{ +use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; +mod test_support; +use test_support::{ create_temp_recipe, create_test_session_metadata, MockBehavior, ScheduleToolTestBuilder, }; @@ -775,19 +776,19 @@ async fn test_schedule_tool_session_content_action_with_real_session() { let (agent, _) = ScheduleToolTestBuilder::new().build().await; // Create a temporary session file in the proper session directory - let session_dir = crate::session::storage::ensure_session_dir().unwrap(); + let session_dir = goose::session::storage::ensure_session_dir().unwrap(); let session_id = "test_session_real"; let session_path = session_dir.join(format!("{}.jsonl", session_id)); // Create test metadata and messages let metadata = create_test_session_metadata(2, "/tmp"); let messages = vec![ - crate::message::Message::user().with_text("Hello"), - crate::message::Message::assistant().with_text("Hi there!"), + goose::message::Message::user().with_text("Hello"), + goose::message::Message::assistant().with_text("Hi there!"), ]; // Save the session file - crate::session::storage::save_messages_with_metadata(&session_path, &metadata, &messages) + goose::session::storage::save_messages_with_metadata(&session_path, &metadata, &messages) .unwrap(); // Test the session_content action diff --git a/crates/goose/src/scheduler_test_support.rs b/crates/goose/tests/scheduler_test_support.rs similarity index 92% rename from crates/goose/src/scheduler_test_support.rs rename to crates/goose/tests/scheduler_test_support.rs index 0b5c9f82dd90..6893caada1ba 100644 --- a/crates/goose/src/scheduler_test_support.rs +++ b/crates/goose/tests/scheduler_test_support.rs @@ -5,7 +5,7 @@ use once_cell::sync::Lazy; use std::sync::Arc; use tokio::sync::Mutex; -use crate::providers::base::Provider as GooseProvider; +use goose::providers::base::Provider as GooseProvider; static TEST_PROVIDER: Lazy>>> = Lazy::new(|| Mutex::new(None)); diff --git a/crates/goose/src/agents/test_support.rs b/crates/goose/tests/test_support.rs similarity index 98% rename from crates/goose/src/agents/test_support.rs rename to crates/goose/tests/test_support.rs index 48a6ce332f51..cec7cdb2b55b 100644 --- a/crates/goose/src/agents/test_support.rs +++ b/crates/goose/tests/test_support.rs @@ -9,10 +9,10 @@ use chrono::{DateTime, Utc}; use tempfile::TempDir; use tokio::sync::Mutex; -use crate::agents::Agent; -use crate::scheduler::{ScheduledJob, SchedulerError}; -use crate::scheduler_trait::SchedulerTrait; -use crate::session::storage::SessionMetadata; +use goose::agents::Agent; +use goose::scheduler::{ScheduledJob, SchedulerError}; +use goose::scheduler_trait::SchedulerTrait; +use goose::session::storage::SessionMetadata; #[derive(Debug, Clone)] pub enum MockBehavior { From 6ee5c6978e1c4aabeb25c9aa9891bd0d60d1b8e5 Mon Sep 17 00:00:00 2001 From: Tyler Longwell Date: Wed, 18 Jun 2025 18:42:10 -0400 Subject: [PATCH 5/5] move handerls --- crates/goose/src/agents/agent.rs | 396 +-------------------- crates/goose/src/agents/mod.rs | 1 + crates/goose/src/agents/schedule_tool.rs | 420 +++++++++++++++++++++++ crates/goose/tests/agent.rs | 25 +- crates/goose/tests/private_tests.rs | 8 +- 5 files changed, 443 insertions(+), 407 deletions(-) create mode 100644 crates/goose/src/agents/schedule_tool.rs diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index c02d47f24079..3f56dacc230b 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -4,7 +4,6 @@ use std::pin::Pin; use std::sync::Arc; use anyhow::{anyhow, Result}; -use chrono::Utc; use futures::stream::BoxStream; use futures::{FutureExt, Stream, TryStreamExt}; use futures_util::stream; @@ -400,399 +399,6 @@ impl Agent { (request_id, result) } - pub async fn handle_schedule_management( - &self, - arguments: serde_json::Value, - _request_id: String, - ) -> ToolResult> { - let scheduler = match self.scheduler_service.lock().await.as_ref() { - Some(s) => s.clone(), - None => { - return Err(ToolError::ExecutionError( - "Scheduler not available. This tool only works in server mode.".to_string(), - )) - } - }; - - let action = arguments - .get("action") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'action' parameter".to_string()))?; - - match action { - "list" => self.handle_list_jobs(scheduler).await, - "create" => self.handle_create_job(scheduler, arguments).await, - "run_now" => self.handle_run_now(scheduler, arguments).await, - "pause" => self.handle_pause_job(scheduler, arguments).await, - "unpause" => self.handle_unpause_job(scheduler, arguments).await, - "delete" => self.handle_delete_job(scheduler, arguments).await, - "kill" => self.handle_kill_job(scheduler, arguments).await, - "inspect" => self.handle_inspect_job(scheduler, arguments).await, - "sessions" => self.handle_list_sessions(scheduler, arguments).await, - "session_content" => self.handle_session_content(arguments).await, - _ => Err(ToolError::ExecutionError(format!( - "Unknown action: {}", - action - ))), - } - } - - async fn handle_list_jobs( - &self, - scheduler: Arc, - ) -> ToolResult> { - match scheduler.list_scheduled_jobs().await { - Ok(jobs) => { - let jobs_json = serde_json::to_string_pretty(&jobs).map_err(|e| { - ToolError::ExecutionError(format!("Failed to serialize jobs: {}", e)) - })?; - Ok(vec![Content::text(format!( - "Scheduled Jobs:\n{}", - jobs_json - ))]) - } - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to list jobs: {}", - e - ))), - } - } - - async fn handle_create_job( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let recipe_path = arguments - .get("recipe_path") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - ToolError::ExecutionError("Missing 'recipe_path' parameter".to_string()) - })?; - - let cron_expression = arguments - .get("cron_expression") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - ToolError::ExecutionError("Missing 'cron_expression' parameter".to_string()) - })?; - - // Validate recipe file exists and is readable - if !std::path::Path::new(recipe_path).exists() { - return Err(ToolError::ExecutionError(format!( - "Recipe file not found: {}", - recipe_path - ))); - } - - // Validate it's a valid recipe by trying to parse it - match std::fs::read_to_string(recipe_path) { - Ok(content) => { - if recipe_path.ends_with(".json") { - serde_json::from_str::(&content).map_err(|e| { - ToolError::ExecutionError(format!("Invalid JSON recipe: {}", e)) - })?; - } else { - serde_yaml::from_str::(&content).map_err(|e| { - ToolError::ExecutionError(format!("Invalid YAML recipe: {}", e)) - })?; - } - } - Err(e) => { - return Err(ToolError::ExecutionError(format!( - "Cannot read recipe file: {}", - e - ))) - } - } - - // Generate unique job ID - let job_id = format!("agent_created_{}", Utc::now().timestamp()); - - let job = crate::scheduler::ScheduledJob { - id: job_id.clone(), - source: recipe_path.to_string(), - cron: cron_expression.to_string(), - last_run: None, - currently_running: false, - paused: false, - current_session_id: None, - process_start_time: None, - }; - - match scheduler.add_scheduled_job(job).await { - Ok(()) => Ok(vec![Content::text(format!( - "Successfully created scheduled job '{}' for recipe '{}' with cron expression '{}'", - job_id, recipe_path, cron_expression - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to create job: {}", - e - ))), - } - } - - async fn handle_run_now( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - match scheduler.run_now(job_id).await { - Ok(session_id) => Ok(vec![Content::text(format!( - "Successfully started job '{}'. Session ID: {}", - job_id, session_id - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to run job: {}", - e - ))), - } - } - - async fn handle_pause_job( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - match scheduler.pause_schedule(job_id).await { - Ok(()) => Ok(vec![Content::text(format!( - "Successfully paused job '{}'", - job_id - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to pause job: {}", - e - ))), - } - } - - async fn handle_unpause_job( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - match scheduler.unpause_schedule(job_id).await { - Ok(()) => Ok(vec![Content::text(format!( - "Successfully unpaused job '{}'", - job_id - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to unpause job: {}", - e - ))), - } - } - - async fn handle_delete_job( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - match scheduler.remove_scheduled_job(job_id).await { - Ok(()) => Ok(vec![Content::text(format!( - "Successfully deleted job '{}'", - job_id - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to delete job: {}", - e - ))), - } - } - - async fn handle_kill_job( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - match scheduler.kill_running_job(job_id).await { - Ok(()) => Ok(vec![Content::text(format!( - "Successfully killed running job '{}'", - job_id - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to kill job: {}", - e - ))), - } - } - - async fn handle_inspect_job( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - match scheduler.get_running_job_info(job_id).await { - Ok(Some((session_id, start_time))) => { - let duration = Utc::now().signed_duration_since(start_time); - Ok(vec![Content::text(format!( - "Job '{}' is currently running:\n- Session ID: {}\n- Started: {}\n- Duration: {} seconds", - job_id, session_id, start_time.to_rfc3339(), duration.num_seconds() - ))]) - } - Ok(None) => Ok(vec![Content::text(format!( - "Job '{}' is not currently running", - job_id - ))]), - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to inspect job: {}", - e - ))), - } - } - - async fn handle_list_sessions( - &self, - scheduler: Arc, - arguments: serde_json::Value, - ) -> ToolResult> { - let job_id = arguments - .get("job_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; - - let limit = arguments - .get("limit") - .and_then(|v| v.as_u64()) - .unwrap_or(50) as usize; - - match scheduler.sessions(job_id, limit).await { - Ok(sessions) => { - if sessions.is_empty() { - Ok(vec![Content::text(format!( - "No sessions found for job '{}'", - job_id - ))]) - } else { - let sessions_info: Vec = sessions - .into_iter() - .map(|(session_name, metadata)| { - format!( - "- Session: {} (Messages: {}, Working Dir: {})", - session_name, - metadata.message_count, - metadata.working_dir.display() - ) - }) - .collect(); - - Ok(vec![Content::text(format!( - "Sessions for job '{}':\n{}", - job_id, - sessions_info.join("\n") - ))]) - } - } - Err(e) => Err(ToolError::ExecutionError(format!( - "Failed to list sessions: {}", - e - ))), - } - } - - async fn handle_session_content( - &self, - arguments: serde_json::Value, - ) -> ToolResult> { - let session_id = arguments - .get("session_id") - .and_then(|v| v.as_str()) - .ok_or_else(|| { - ToolError::ExecutionError("Missing 'session_id' parameter".to_string()) - })?; - - // Get the session file path - let session_path = crate::session::storage::get_path( - crate::session::storage::Identifier::Name(session_id.to_string()), - ); - - // Check if session file exists - if !session_path.exists() { - return Err(ToolError::ExecutionError(format!( - "Session '{}' not found", - session_id - ))); - } - - // Read session metadata - let metadata = match crate::session::storage::read_metadata(&session_path) { - Ok(metadata) => metadata, - Err(e) => { - return Err(ToolError::ExecutionError(format!( - "Failed to read session metadata: {}", - e - ))); - } - }; - - // Read session messages - let messages = match crate::session::storage::read_messages(&session_path) { - Ok(messages) => messages, - Err(e) => { - return Err(ToolError::ExecutionError(format!( - "Failed to read session messages: {}", - e - ))); - } - }; - - // Format the response with metadata and messages - let metadata_json = match serde_json::to_string_pretty(&metadata) { - Ok(json) => json, - Err(e) => { - return Err(ToolError::ExecutionError(format!( - "Failed to serialize metadata: {}", - e - ))); - } - }; - - let messages_json = match serde_json::to_string_pretty(&messages) { - Ok(json) => json, - Err(e) => { - return Err(ToolError::ExecutionError(format!( - "Failed to serialize messages: {}", - e - ))); - } - }; - - Ok(vec![Content::text(format!( - "Session '{}' Content:\n\nMetadata:\n{}\n\nMessages:\n{}", - session_id, metadata_json, messages_json - ))]) - } - pub async fn add_extension(&self, extension: ExtensionConfig) -> ExtensionResult<()> { match &extension { ExtensionConfig::Frontend { @@ -825,7 +431,7 @@ impl Agent { let mut extension_manager = self.extension_manager.lock().await; extension_manager.add_extension(extension.clone()).await?; } - }; + } // If vector tool selection is enabled, index the tools let selector = self.router_tool_selector.lock().await.clone(); diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 94b3bd0e971e..bfd51217ae37 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -8,6 +8,7 @@ pub mod prompt_manager; mod reply_parts; mod router_tool_selector; mod router_tools; +mod schedule_tool; mod tool_execution; mod tool_router_index_manager; diff --git a/crates/goose/src/agents/schedule_tool.rs b/crates/goose/src/agents/schedule_tool.rs new file mode 100644 index 000000000000..1dac50bed51d --- /dev/null +++ b/crates/goose/src/agents/schedule_tool.rs @@ -0,0 +1,420 @@ +//! Schedule tool handlers for the Goose agent +//! +//! This module contains all the handlers for the schedule management platform tool, +//! including job creation, execution, monitoring, and session management. + +use std::sync::Arc; + +use chrono::Utc; +use mcp_core::{Content, ToolError, ToolResult}; + +use crate::recipe::Recipe; +use crate::scheduler_trait::SchedulerTrait; + +use super::Agent; + +impl Agent { + /// Handle schedule management tool calls + pub async fn handle_schedule_management( + &self, + arguments: serde_json::Value, + _request_id: String, + ) -> ToolResult> { + let scheduler = match self.scheduler_service.lock().await.as_ref() { + Some(s) => s.clone(), + None => { + return Err(ToolError::ExecutionError( + "Scheduler not available. This tool only works in server mode.".to_string(), + )) + } + }; + + let action = arguments + .get("action") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'action' parameter".to_string()))?; + + match action { + "list" => self.handle_list_jobs(scheduler).await, + "create" => self.handle_create_job(scheduler, arguments).await, + "run_now" => self.handle_run_now(scheduler, arguments).await, + "pause" => self.handle_pause_job(scheduler, arguments).await, + "unpause" => self.handle_unpause_job(scheduler, arguments).await, + "delete" => self.handle_delete_job(scheduler, arguments).await, + "kill" => self.handle_kill_job(scheduler, arguments).await, + "inspect" => self.handle_inspect_job(scheduler, arguments).await, + "sessions" => self.handle_list_sessions(scheduler, arguments).await, + "session_content" => self.handle_session_content(arguments).await, + _ => Err(ToolError::ExecutionError(format!( + "Unknown action: {}", + action + ))), + } + } + + /// List all scheduled jobs + async fn handle_list_jobs( + &self, + scheduler: Arc, + ) -> ToolResult> { + match scheduler.list_scheduled_jobs().await { + Ok(jobs) => { + let jobs_json = serde_json::to_string_pretty(&jobs).map_err(|e| { + ToolError::ExecutionError(format!("Failed to serialize jobs: {}", e)) + })?; + Ok(vec![Content::text(format!( + "Scheduled Jobs:\n{}", + jobs_json + ))]) + } + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to list jobs: {}", + e + ))), + } + } + + /// Create a new scheduled job from a recipe file + async fn handle_create_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let recipe_path = arguments + .get("recipe_path") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionError("Missing 'recipe_path' parameter".to_string()) + })?; + + let cron_expression = arguments + .get("cron_expression") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionError("Missing 'cron_expression' parameter".to_string()) + })?; + + // Validate recipe file exists and is readable + if !std::path::Path::new(recipe_path).exists() { + return Err(ToolError::ExecutionError(format!( + "Recipe file not found: {}", + recipe_path + ))); + } + + // Validate it's a valid recipe by trying to parse it + match std::fs::read_to_string(recipe_path) { + Ok(content) => { + if recipe_path.ends_with(".json") { + serde_json::from_str::(&content).map_err(|e| { + ToolError::ExecutionError(format!("Invalid JSON recipe: {}", e)) + })?; + } else { + serde_yaml::from_str::(&content).map_err(|e| { + ToolError::ExecutionError(format!("Invalid YAML recipe: {}", e)) + })?; + } + } + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Cannot read recipe file: {}", + e + ))) + } + } + + // Generate unique job ID + let job_id = format!("agent_created_{}", Utc::now().timestamp()); + + let job = crate::scheduler::ScheduledJob { + id: job_id.clone(), + source: recipe_path.to_string(), + cron: cron_expression.to_string(), + last_run: None, + currently_running: false, + paused: false, + current_session_id: None, + process_start_time: None, + }; + + match scheduler.add_scheduled_job(job).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully created scheduled job '{}' for recipe '{}' with cron expression '{}'", + job_id, recipe_path, cron_expression + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to create job: {}", + e + ))), + } + } + + /// Run a scheduled job immediately + async fn handle_run_now( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.run_now(job_id).await { + Ok(session_id) => Ok(vec![Content::text(format!( + "Successfully started job '{}'. Session ID: {}", + job_id, session_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to run job: {}", + e + ))), + } + } + + /// Pause a scheduled job + async fn handle_pause_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.pause_schedule(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully paused job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to pause job: {}", + e + ))), + } + } + + /// Resume a paused scheduled job + async fn handle_unpause_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.unpause_schedule(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully unpaused job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to unpause job: {}", + e + ))), + } + } + + /// Delete a scheduled job + async fn handle_delete_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.remove_scheduled_job(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully deleted job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to delete job: {}", + e + ))), + } + } + + /// Terminate a currently running job + async fn handle_kill_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.kill_running_job(job_id).await { + Ok(()) => Ok(vec![Content::text(format!( + "Successfully killed running job '{}'", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to kill job: {}", + e + ))), + } + } + + /// Get information about a running job + async fn handle_inspect_job( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + match scheduler.get_running_job_info(job_id).await { + Ok(Some((session_id, start_time))) => { + let duration = Utc::now().signed_duration_since(start_time); + Ok(vec![Content::text(format!( + "Job '{}' is currently running:\n- Session ID: {}\n- Started: {}\n- Duration: {} seconds", + job_id, session_id, start_time.to_rfc3339(), duration.num_seconds() + ))]) + } + Ok(None) => Ok(vec![Content::text(format!( + "Job '{}' is not currently running", + job_id + ))]), + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to inspect job: {}", + e + ))), + } + } + + /// List execution sessions for a job + async fn handle_list_sessions( + &self, + scheduler: Arc, + arguments: serde_json::Value, + ) -> ToolResult> { + let job_id = arguments + .get("job_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| ToolError::ExecutionError("Missing 'job_id' parameter".to_string()))?; + + let limit = arguments + .get("limit") + .and_then(|v| v.as_u64()) + .unwrap_or(50) as usize; + + match scheduler.sessions(job_id, limit).await { + Ok(sessions) => { + if sessions.is_empty() { + Ok(vec![Content::text(format!( + "No sessions found for job '{}'", + job_id + ))]) + } else { + let sessions_info: Vec = sessions + .into_iter() + .map(|(session_name, metadata)| { + format!( + "- Session: {} (Messages: {}, Working Dir: {})", + session_name, + metadata.message_count, + metadata.working_dir.display() + ) + }) + .collect(); + + Ok(vec![Content::text(format!( + "Sessions for job '{}':\n{}", + job_id, + sessions_info.join("\n") + ))]) + } + } + Err(e) => Err(ToolError::ExecutionError(format!( + "Failed to list sessions: {}", + e + ))), + } + } + + /// Get the full content (metadata and messages) of a specific session + async fn handle_session_content( + &self, + arguments: serde_json::Value, + ) -> ToolResult> { + let session_id = arguments + .get("session_id") + .and_then(|v| v.as_str()) + .ok_or_else(|| { + ToolError::ExecutionError("Missing 'session_id' parameter".to_string()) + })?; + + // Get the session file path + let session_path = crate::session::storage::get_path( + crate::session::storage::Identifier::Name(session_id.to_string()), + ); + + // Check if session file exists + if !session_path.exists() { + return Err(ToolError::ExecutionError(format!( + "Session '{}' not found", + session_id + ))); + } + + // Read session metadata + let metadata = match crate::session::storage::read_metadata(&session_path) { + Ok(metadata) => metadata, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to read session metadata: {}", + e + ))); + } + }; + + // Read session messages + let messages = match crate::session::storage::read_messages(&session_path) { + Ok(messages) => messages, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to read session messages: {}", + e + ))); + } + }; + + // Format the response with metadata and messages + let metadata_json = match serde_json::to_string_pretty(&metadata) { + Ok(json) => json, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to serialize metadata: {}", + e + ))); + } + }; + + let messages_json = match serde_json::to_string_pretty(&messages) { + Ok(json) => json, + Err(e) => { + return Err(ToolError::ExecutionError(format!( + "Failed to serialize messages: {}", + e + ))); + } + }; + + Ok(vec![Content::text(format!( + "Session '{}' Content:\n\nMetadata:\n{}\n\nMessages:\n{}", + session_id, metadata_json, messages_json + ))]) + } +} diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index f3785c73c55b..58979af1cb05 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -334,12 +334,12 @@ mod tests { #[cfg(test)] mod schedule_tool_tests { use super::*; + use async_trait::async_trait; + use chrono::{DateTime, Utc}; + use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; use goose::scheduler::{ScheduledJob, SchedulerError}; use goose::scheduler_trait::SchedulerTrait; use goose::session::storage::SessionMetadata; - use goose::agents::platform_tools::PLATFORM_MANAGE_SCHEDULE_TOOL_NAME; - use async_trait::async_trait; - use chrono::{DateTime, Utc}; use std::sync::Arc; // Mock scheduler for testing @@ -465,7 +465,7 @@ mod schedule_tool_tests { assert!(tool .description .contains("Manage scheduled recipe execution")); - + // Verify the tool has the expected actions in its schema if let Some(properties) = tool.input_schema.get("properties") { if let Some(action_prop) = properties.get("action") { @@ -476,7 +476,7 @@ mod schedule_tool_tests { .iter() .map(|v| v.as_str().unwrap().to_string()) .collect(); - + // Check that our session_content action is included assert!(actions.contains(&"session_content".to_string())); assert!(actions.contains(&"list".to_string())); @@ -497,14 +497,21 @@ mod schedule_tool_tests { assert!(schedule_tool.is_some()); let tool = schedule_tool.unwrap(); - + // Verify the tool schema has the session_id parameter for session_content action if let Some(properties) = tool.input_schema.get("properties") { assert!(properties.get("session_id").is_some()); - + if let Some(session_id_prop) = properties.get("session_id") { - assert_eq!(session_id_prop.get("type").unwrap().as_str().unwrap(), "string"); - assert!(session_id_prop.get("description").unwrap().as_str().unwrap() + assert_eq!( + session_id_prop.get("type").unwrap().as_str().unwrap(), + "string" + ); + assert!(session_id_prop + .get("description") + .unwrap() + .as_str() + .unwrap() .contains("Session identifier for session_content action")); } } diff --git a/crates/goose/tests/private_tests.rs b/crates/goose/tests/private_tests.rs index b94b04767d39..78e7cdd7658c 100644 --- a/crates/goose/tests/private_tests.rs +++ b/crates/goose/tests/private_tests.rs @@ -800,17 +800,19 @@ async fn test_schedule_tool_session_content_action_with_real_session() { let result = agent .handle_schedule_management(arguments, "test_req".to_string()) .await; - + // Clean up the test session file let _ = std::fs::remove_file(&session_path); // Verify the result assert!(result.is_ok()); - + if let Ok(content) = result { assert_eq!(content.len(), 1); if let mcp_core::Content::Text(text_content) = &content[0] { - assert!(text_content.text.contains("Session 'test_session_real' Content:")); + assert!(text_content + .text + .contains("Session 'test_session_real' Content:")); assert!(text_content.text.contains("Metadata:")); assert!(text_content.text.contains("Messages:")); assert!(text_content.text.contains("Hello"));