use std::collections::HashMap; use std::fs; use std::io; use std::path::{Path, PathBuf}; use std::sync::Arc; use anyhow::{anyhow, Result}; use async_trait::async_trait; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; use tokio_cron_scheduler::{job::JobId, Job, JobScheduler as TokioJobScheduler}; use crate::agents::AgentEvent; use crate::agents::{Agent, SessionConfig}; use crate::config::paths::Paths; use crate::config::Config; use crate::conversation::message::Message; use crate::conversation::Conversation; use crate::providers::base::Provider as GooseProvider; // Alias to avoid conflict in test section use crate::providers::create; use crate::recipe::Recipe; use crate::scheduler_trait::SchedulerTrait; use crate::session::{Session, SessionManager}; // Track running tasks with their abort handles type RunningTasksMap = HashMap; type JobsMap = HashMap; /// Normalize a cron string so that: /// 1. It is always in **quartz 7-field format** expected by Temporal /// (seconds minutes hours dom month dow year). /// 2. Five-field → prepend seconds `0` and append year `*`. /// Six-field → append year `*`. /// 3. Everything else returned unchanged (with a warning). pub fn normalize_cron_expression(src: &str) -> String { let mut parts: Vec<&str> = src.split_whitespace().collect(); match parts.len() { 5 => { // min hour dom mon dow → 0 min hour dom mon dow * parts.insert(0, "0"); parts.push("*"); } 6 => { // sec min hour dom mon dow → sec min hour dom mon dow * parts.push("*"); } 7 => { // already quartz – do nothing } _ => { tracing::warn!( "Unrecognised cron expression '{}': expected 5, 6 or 7 fields (got {}). Leaving unchanged.", src, parts.len() ); return src.to_string(); } } parts.join(" ") } pub fn get_default_scheduler_storage_path() -> Result { let data_dir = Paths::data_dir(); fs::create_dir_all(&data_dir)?; Ok(data_dir.join("schedules.json")) } pub fn get_default_scheduled_recipes_dir() -> Result { let data_dir = Paths::data_dir(); let recipes_dir = data_dir.join("scheduled_recipes"); fs::create_dir_all(&recipes_dir).map_err(SchedulerError::StorageError)?; tracing::debug!( "Created scheduled recipes directory at: {}", recipes_dir.display() ); Ok(recipes_dir) } #[derive(Debug)] pub enum SchedulerError { JobIdExists(String), JobNotFound(String), StorageError(io::Error), RecipeLoadError(String), AgentSetupError(String), PersistError(String), CronParseError(String), SchedulerInternalError(String), AnyhowError(anyhow::Error), } impl std::fmt::Display for SchedulerError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { SchedulerError::JobIdExists(id) => write!(f, "Job ID '{}' already exists.", id), SchedulerError::JobNotFound(id) => write!(f, "Job ID '{}' not found.", id), SchedulerError::StorageError(e) => write!(f, "Storage error: {}", e), SchedulerError::RecipeLoadError(e) => write!(f, "Recipe load error: {}", e), SchedulerError::AgentSetupError(e) => write!(f, "Agent setup error: {}", e), SchedulerError::PersistError(e) => write!(f, "Failed to persist schedules: {}", e), SchedulerError::CronParseError(e) => write!(f, "Invalid cron string: {}", e), SchedulerError::SchedulerInternalError(e) => { write!(f, "Scheduler internal error: {}", e) } SchedulerError::AnyhowError(e) => write!(f, "Scheduler operation failed: {}", e), } } } impl std::error::Error for SchedulerError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { SchedulerError::StorageError(e) => Some(e), SchedulerError::AnyhowError(e) => Some(e.as_ref()), _ => None, } } } impl From for SchedulerError { fn from(err: io::Error) -> Self { SchedulerError::StorageError(err) } } impl From for SchedulerError { fn from(err: serde_json::Error) -> Self { SchedulerError::PersistError(err.to_string()) } } impl From for SchedulerError { fn from(err: anyhow::Error) -> Self { SchedulerError::AnyhowError(err) } } #[derive(Clone, Serialize, Deserialize, Debug, utoipa::ToSchema)] pub struct ScheduledJob { pub id: String, pub source: String, pub cron: String, pub last_run: Option>, #[serde(default)] pub currently_running: bool, #[serde(default)] pub paused: bool, #[serde(default)] pub current_session_id: Option, #[serde(default)] pub process_start_time: Option>, #[serde(default)] pub execution_mode: Option, // "foreground" or "background" } async fn persist_jobs_from_arc( storage_path: &Path, jobs_arc: &Arc>, ) -> Result<(), SchedulerError> { let jobs_guard = jobs_arc.lock().await; let list: Vec = jobs_guard.values().map(|(_, j)| j.clone()).collect(); if let Some(parent) = storage_path.parent() { fs::create_dir_all(parent).map_err(SchedulerError::StorageError)?; } let data = serde_json::to_string_pretty(&list).map_err(SchedulerError::from)?; fs::write(storage_path, data).map_err(SchedulerError::StorageError)?; Ok(()) } pub struct Scheduler { internal_scheduler: TokioJobScheduler, jobs: Arc>, storage_path: PathBuf, running_tasks: Arc>, } impl Scheduler { pub async fn new(storage_path: PathBuf) -> Result, SchedulerError> { let internal_scheduler = TokioJobScheduler::new() .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; let jobs = Arc::new(Mutex::new(HashMap::new())); let running_tasks = Arc::new(Mutex::new(HashMap::new())); let arc_self = Arc::new(Self { internal_scheduler, jobs, storage_path, running_tasks, }); arc_self.load_jobs_from_storage().await?; arc_self .internal_scheduler .start() .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; Ok(arc_self) } pub async fn add_scheduled_job( &self, original_job_spec: ScheduledJob, ) -> Result<(), SchedulerError> { let mut jobs_guard = self.jobs.lock().await; if jobs_guard.contains_key(&original_job_spec.id) { return Err(SchedulerError::JobIdExists(original_job_spec.id.clone())); } let original_recipe_path = Path::new(&original_job_spec.source); if !original_recipe_path.exists() { return Err(SchedulerError::RecipeLoadError(format!( "Original recipe file not found: {}", original_job_spec.source ))); } if !original_recipe_path.is_file() { return Err(SchedulerError::RecipeLoadError(format!( "Original recipe source is not a file: {}", original_job_spec.source ))); } let scheduled_recipes_dir = get_default_scheduled_recipes_dir()?; let original_extension = original_recipe_path .extension() .and_then(|ext| ext.to_str()) .unwrap_or("yaml"); let destination_filename = format!("{}.{}", original_job_spec.id, original_extension); let destination_recipe_path = scheduled_recipes_dir.join(destination_filename); tracing::info!( "Copying recipe from {} to {}", original_recipe_path.display(), destination_recipe_path.display() ); fs::copy(original_recipe_path, &destination_recipe_path).map_err(|e| { SchedulerError::StorageError(io::Error::new( e.kind(), format!( "Failed to copy recipe from {} to {}: {}", original_job_spec.source, destination_recipe_path.display(), e ), )) })?; let mut stored_job = original_job_spec.clone(); stored_job.source = destination_recipe_path.to_string_lossy().into_owned(); stored_job.current_session_id = None; stored_job.process_start_time = None; tracing::info!("Updated job source path to: {}", stored_job.source); let job_for_task = stored_job.clone(); let jobs_arc_for_task = self.jobs.clone(); let storage_path_for_task = self.storage_path.clone(); let running_tasks_for_task = self.running_tasks.clone(); tracing::info!("Attempting to parse cron expression: '{}'", stored_job.cron); let normalized_cron = normalize_cron_expression(&stored_job.cron); // Convert from 7-field (Temporal format) to 6-field (tokio-cron-scheduler format) let tokio_cron = { let parts: Vec<&str> = normalized_cron.split_whitespace().collect(); if parts.len() == 7 { parts[..6].join(" ") } else { normalized_cron.clone() } }; if tokio_cron != stored_job.cron { tracing::info!( "Converted cron expression from '{}' to '{}' for tokio-cron-scheduler", stored_job.cron, tokio_cron ); } let cron_task = Job::new_async(&tokio_cron, move |_uuid, _l| { let task_job_id = job_for_task.id.clone(); let current_jobs_arc = jobs_arc_for_task.clone(); let local_storage_path = storage_path_for_task.clone(); let job_to_execute = job_for_task.clone(); // Clone for run_scheduled_job_internal let running_tasks_arc = running_tasks_for_task.clone(); Box::pin(async move { // Check if the job is paused before executing let should_execute = { let jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, current_job_in_map)) = jobs_map_guard.get(&task_job_id) { !current_job_in_map.paused } else { false } }; if !should_execute { tracing::info!("Skipping execution of paused job '{}'", &task_job_id); return; } let current_time = Utc::now(); let mut needs_persist = false; { let mut jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, current_job_in_map)) = jobs_map_guard.get_mut(&task_job_id) { current_job_in_map.last_run = Some(current_time); current_job_in_map.currently_running = true; current_job_in_map.process_start_time = Some(current_time); needs_persist = true; } } if needs_persist { if let Err(e) = persist_jobs_from_arc(&local_storage_path, ¤t_jobs_arc).await { tracing::error!( "Failed to persist last_run update for job {}: {}", &task_job_id, e ); } } // Spawn the job execution as an abortable task let job_task = tokio::spawn(run_scheduled_job_internal( job_to_execute.clone(), None, Some(current_jobs_arc.clone()), Some(task_job_id.clone()), )); // Store the abort handle at the scheduler level { let mut running_tasks_guard = running_tasks_arc.lock().await; running_tasks_guard.insert(task_job_id.clone(), job_task.abort_handle()); } // Wait for the job to complete or be aborted let result = job_task.await; // Remove the abort handle { let mut running_tasks_guard = running_tasks_arc.lock().await; running_tasks_guard.remove(&task_job_id); } // Update the job status after execution { let mut jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, current_job_in_map)) = jobs_map_guard.get_mut(&task_job_id) { current_job_in_map.currently_running = false; current_job_in_map.current_session_id = None; current_job_in_map.process_start_time = None; needs_persist = true; } } if needs_persist { if let Err(e) = persist_jobs_from_arc(&local_storage_path, ¤t_jobs_arc).await { tracing::error!( "Failed to persist running status update for job {}: {}", &task_job_id, e ); } } match result { Ok(Ok(_session_id)) => { tracing::info!("Scheduled job '{}' completed successfully", &task_job_id); } Ok(Err(e)) => { tracing::error!( "Scheduled job '{}' execution failed: {}", &e.job_id, e.error ); } Err(join_error) if join_error.is_cancelled() => { tracing::info!("Scheduled job '{}' was cancelled/killed", &task_job_id); } Err(join_error) => { tracing::error!( "Scheduled job '{}' task failed: {}", &task_job_id, join_error ); } } }) }) .map_err(|e| SchedulerError::CronParseError(e.to_string()))?; let job_uuid = self .internal_scheduler .add(cron_task) .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; jobs_guard.insert(stored_job.id.clone(), (job_uuid, stored_job)); // Pass the jobs_guard by reference for the initial persist after adding a job self.persist_jobs_to_storage_with_guard(&jobs_guard).await?; Ok(()) } async fn load_jobs_from_storage(self: &Arc) -> Result<(), SchedulerError> { if !self.storage_path.exists() { return Ok(()); } let data = fs::read_to_string(&self.storage_path)?; if data.trim().is_empty() { return Ok(()); } let list: Vec = serde_json::from_str(&data).map_err(|e| { SchedulerError::PersistError(format!("Failed to deserialize schedules.json: {}", e)) })?; let mut jobs_guard = self.jobs.lock().await; for job_to_load in list { if !Path::new(&job_to_load.source).exists() { tracing::warn!("Recipe file {} for scheduled job {} not found in shared store. Skipping job load.", job_to_load.source, job_to_load.id); continue; } let job_for_task = job_to_load.clone(); let jobs_arc_for_task = self.jobs.clone(); let storage_path_for_task = self.storage_path.clone(); let running_tasks_for_task = self.running_tasks.clone(); tracing::info!( "Loading job '{}' with cron expression: '{}'", job_to_load.id, job_to_load.cron ); let normalized_cron = normalize_cron_expression(&job_to_load.cron); // Convert from 7-field (Temporal format) to 6-field (tokio-cron-scheduler format) let tokio_cron = { let parts: Vec<&str> = normalized_cron.split_whitespace().collect(); if parts.len() == 7 { parts[..6].join(" ") } else { normalized_cron.clone() } }; if tokio_cron != job_to_load.cron { tracing::info!( "Converted cron expression from '{}' to '{}' for tokio-cron-scheduler", job_to_load.cron, tokio_cron ); } let cron_task = Job::new_async(&tokio_cron, move |_uuid, _l| { let task_job_id = job_for_task.id.clone(); let current_jobs_arc = jobs_arc_for_task.clone(); let local_storage_path = storage_path_for_task.clone(); let job_to_execute = job_for_task.clone(); // Clone for run_scheduled_job_internal let running_tasks_arc = running_tasks_for_task.clone(); Box::pin(async move { // Check if the job is paused before executing let should_execute = { let jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, stored_job)) = jobs_map_guard.get(&task_job_id) { !stored_job.paused } else { false } }; if !should_execute { tracing::info!("Skipping execution of paused job '{}'", &task_job_id); return; } let current_time = Utc::now(); let mut needs_persist = false; { let mut jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, stored_job)) = jobs_map_guard.get_mut(&task_job_id) { stored_job.last_run = Some(current_time); stored_job.currently_running = true; stored_job.process_start_time = Some(current_time); needs_persist = true; } } if needs_persist { if let Err(e) = persist_jobs_from_arc(&local_storage_path, ¤t_jobs_arc).await { tracing::error!( "Failed to persist last_run update for loaded job {}: {}", &task_job_id, e ); } } // Spawn the job execution as an abortable task let job_task = tokio::spawn(run_scheduled_job_internal( job_to_execute, None, Some(current_jobs_arc.clone()), Some(task_job_id.clone()), )); // Store the abort handle at the scheduler level { let mut running_tasks_guard = running_tasks_arc.lock().await; running_tasks_guard.insert(task_job_id.clone(), job_task.abort_handle()); } // Wait for the job to complete or be aborted let result = job_task.await; // Remove the abort handle { let mut running_tasks_guard = running_tasks_arc.lock().await; running_tasks_guard.remove(&task_job_id); } // Update the job status after execution { let mut jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, stored_job)) = jobs_map_guard.get_mut(&task_job_id) { stored_job.currently_running = false; stored_job.current_session_id = None; stored_job.process_start_time = None; needs_persist = true; } } if needs_persist { if let Err(e) = persist_jobs_from_arc(&local_storage_path, ¤t_jobs_arc).await { tracing::error!( "Failed to persist running status update for job {}: {}", &task_job_id, e ); } } match result { Ok(Ok(_session_id)) => { tracing::info!( "Scheduled job '{}' completed successfully", &task_job_id ); } Ok(Err(e)) => { tracing::error!( "Scheduled job '{}' execution failed: {}", &e.job_id, e.error ); } Err(join_error) if join_error.is_cancelled() => { tracing::info!("Scheduled job '{}' was cancelled/killed", &task_job_id); } Err(join_error) => { tracing::error!( "Scheduled job '{}' task failed: {}", &task_job_id, join_error ); } } }) }) .map_err(|e| SchedulerError::CronParseError(e.to_string()))?; let job_uuid = self .internal_scheduler .add(cron_task) .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; jobs_guard.insert(job_to_load.id.clone(), (job_uuid, job_to_load)); } Ok(()) } // Renamed and kept for direct use when a guard is already held (e.g. add/remove) async fn persist_jobs_to_storage_with_guard( &self, jobs_guard: &tokio::sync::MutexGuard<'_, JobsMap>, ) -> Result<(), SchedulerError> { let list: Vec = jobs_guard.values().map(|(_, j)| j.clone()).collect(); if let Some(parent) = self.storage_path.parent() { fs::create_dir_all(parent)?; } let data = serde_json::to_string_pretty(&list)?; fs::write(&self.storage_path, data)?; Ok(()) } // New function that locks and calls the helper, for run_now and potentially other places async fn persist_jobs(&self) -> Result<(), SchedulerError> { persist_jobs_from_arc(&self.storage_path, &self.jobs).await } pub async fn list_scheduled_jobs(&self) -> Vec { self.jobs .lock() .await .values() .map(|(_, j)| j.clone()) .collect() } pub async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> { let mut jobs_guard = self.jobs.lock().await; if let Some((job_uuid, scheduled_job)) = jobs_guard.remove(id) { self.internal_scheduler .remove(&job_uuid) .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; let recipe_path = Path::new(&scheduled_job.source); if recipe_path.exists() { fs::remove_file(recipe_path).map_err(SchedulerError::StorageError)?; } self.persist_jobs_to_storage_with_guard(&jobs_guard).await?; Ok(()) } else { Err(SchedulerError::JobNotFound(id.to_string())) } } pub async fn sessions( &self, sched_id: &str, limit: usize, ) -> Result, SchedulerError> { let all_sessions = SessionManager::list_sessions() .await .map_err(|e| SchedulerError::StorageError(io::Error::other(e)))?; let mut schedule_sessions: Vec<(String, Session)> = Vec::new(); for session in all_sessions { if session.schedule_id.as_deref() == Some(sched_id) { schedule_sessions.push((session.id.clone(), session)); } } schedule_sessions.sort_by(|a, b| b.0.cmp(&a.0)); let result_sessions: Vec<(String, Session)> = schedule_sessions.into_iter().take(limit).collect(); Ok(result_sessions) } pub async fn run_now(&self, sched_id: &str) -> Result { let job_to_run: ScheduledJob = { let mut jobs_guard = self.jobs.lock().await; match jobs_guard.get_mut(sched_id) { Some((_, job_def)) => { // Set the currently_running flag before executing job_def.currently_running = true; let job_clone = job_def.clone(); // Drop the guard before persisting to avoid borrow issues drop(jobs_guard); // Persist the change immediately self.persist_jobs().await?; job_clone } None => return Err(SchedulerError::JobNotFound(sched_id.to_string())), } }; // Spawn the job execution as an abortable task for run_now let job_task = tokio::spawn(run_scheduled_job_internal( job_to_run.clone(), None, Some(self.jobs.clone()), Some(sched_id.to_string()), )); // Store the abort handle for run_now jobs { let mut running_tasks_guard = self.running_tasks.lock().await; running_tasks_guard.insert(sched_id.to_string(), job_task.abort_handle()); } // Wait for the job to complete or be aborted let run_result = job_task.await; // Remove the abort handle { let mut running_tasks_guard = self.running_tasks.lock().await; running_tasks_guard.remove(sched_id); } // Clear the currently_running flag after execution { let mut jobs_guard = self.jobs.lock().await; if let Some((_tokio_job_id, job_in_map)) = jobs_guard.get_mut(sched_id) { job_in_map.currently_running = false; job_in_map.current_session_id = None; job_in_map.process_start_time = None; job_in_map.last_run = Some(Utc::now()); } // MutexGuard is dropped here } // Persist after the lock is released and update is made. self.persist_jobs().await?; match run_result { Ok(Ok(session_id)) => Ok(session_id), Ok(Err(e)) => Err(SchedulerError::AnyhowError(anyhow!( "Failed to execute job '{}' immediately: {}", sched_id, e.error ))), Err(join_error) if join_error.is_cancelled() => { tracing::info!("Run now job '{}' was cancelled/killed", sched_id); Err(SchedulerError::AnyhowError(anyhow!( "Job '{}' was successfully cancelled", sched_id ))) } Err(join_error) => Err(SchedulerError::AnyhowError(anyhow!( "Failed to execute job '{}' immediately: {}", sched_id, join_error ))), } } pub async fn pause_schedule(&self, sched_id: &str) -> Result<(), SchedulerError> { let mut jobs_guard = self.jobs.lock().await; match jobs_guard.get_mut(sched_id) { Some((_, job_def)) => { if job_def.currently_running { return Err(SchedulerError::AnyhowError(anyhow!( "Cannot pause schedule '{}' while it's currently running", sched_id ))); } job_def.paused = true; self.persist_jobs_to_storage_with_guard(&jobs_guard).await?; Ok(()) } None => Err(SchedulerError::JobNotFound(sched_id.to_string())), } } pub async fn unpause_schedule(&self, sched_id: &str) -> Result<(), SchedulerError> { let mut jobs_guard = self.jobs.lock().await; match jobs_guard.get_mut(sched_id) { Some((_, job_def)) => { job_def.paused = false; self.persist_jobs_to_storage_with_guard(&jobs_guard).await?; Ok(()) } None => Err(SchedulerError::JobNotFound(sched_id.to_string())), } } pub async fn update_schedule( &self, sched_id: &str, new_cron: String, ) -> Result<(), SchedulerError> { let mut jobs_guard = self.jobs.lock().await; match jobs_guard.get_mut(sched_id) { Some((job_uuid, job_def)) => { if job_def.currently_running { return Err(SchedulerError::AnyhowError(anyhow!( "Cannot edit schedule '{}' while it's currently running", sched_id ))); } if new_cron == job_def.cron { // No change needed return Ok(()); } // Remove the old job from the scheduler self.internal_scheduler .remove(job_uuid) .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; // Create new job with updated cron let job_for_task = job_def.clone(); let jobs_arc_for_task = self.jobs.clone(); let storage_path_for_task = self.storage_path.clone(); let running_tasks_for_task = self.running_tasks.clone(); tracing::info!( "Updating job '{}' with new cron expression: '{}'", sched_id, new_cron ); let normalized_cron = normalize_cron_expression(&new_cron); // Convert from 7-field (Temporal format) to 6-field (tokio-cron-scheduler format) let tokio_cron = { let parts: Vec<&str> = normalized_cron.split_whitespace().collect(); if parts.len() == 7 { parts[..6].join(" ") } else { normalized_cron.clone() } }; if tokio_cron != new_cron { tracing::info!( "Converted cron expression from '{}' to '{}' for tokio-cron-scheduler", new_cron, tokio_cron ); } let cron_task = Job::new_async(&tokio_cron, move |_uuid, _l| { let task_job_id = job_for_task.id.clone(); let current_jobs_arc = jobs_arc_for_task.clone(); let local_storage_path = storage_path_for_task.clone(); let job_to_execute = job_for_task.clone(); let running_tasks_arc = running_tasks_for_task.clone(); Box::pin(async move { // Check if the job is paused before executing let should_execute = { let jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, current_job_in_map)) = jobs_map_guard.get(&task_job_id) { !current_job_in_map.paused } else { false } }; if !should_execute { tracing::info!("Skipping execution of paused job '{}'", &task_job_id); return; } let current_time = Utc::now(); let mut needs_persist = false; { let mut jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, current_job_in_map)) = jobs_map_guard.get_mut(&task_job_id) { current_job_in_map.last_run = Some(current_time); current_job_in_map.currently_running = true; current_job_in_map.process_start_time = Some(current_time); needs_persist = true; } } if needs_persist { if let Err(e) = persist_jobs_from_arc(&local_storage_path, ¤t_jobs_arc).await { tracing::error!( "Failed to persist last_run update for job {}: {}", &task_job_id, e ); } } // Spawn the job execution as an abortable task let job_task = tokio::spawn(run_scheduled_job_internal( job_to_execute, None, Some(current_jobs_arc.clone()), Some(task_job_id.clone()), )); // Store the abort handle at the scheduler level { let mut running_tasks_guard = running_tasks_arc.lock().await; running_tasks_guard .insert(task_job_id.clone(), job_task.abort_handle()); } // Wait for the job to complete or be aborted let result = job_task.await; // Remove the abort handle { let mut running_tasks_guard = running_tasks_arc.lock().await; running_tasks_guard.remove(&task_job_id); } // Update the job status after execution { let mut jobs_map_guard = current_jobs_arc.lock().await; if let Some((_, current_job_in_map)) = jobs_map_guard.get_mut(&task_job_id) { current_job_in_map.currently_running = false; current_job_in_map.current_session_id = None; current_job_in_map.process_start_time = None; needs_persist = true; } } if needs_persist { if let Err(e) = persist_jobs_from_arc(&local_storage_path, ¤t_jobs_arc).await { tracing::error!( "Failed to persist running status update for job {}: {}", &task_job_id, e ); } } match result { Ok(Ok(_session_id)) => { tracing::info!( "Scheduled job '{}' completed successfully", &task_job_id ); } Ok(Err(e)) => { tracing::error!( "Scheduled job '{}' execution failed: {}", &e.job_id, e.error ); } Err(join_error) if join_error.is_cancelled() => { tracing::info!( "Scheduled job '{}' was cancelled/killed", &task_job_id ); } Err(join_error) => { tracing::error!( "Scheduled job '{}' task failed: {}", &task_job_id, join_error ); } } }) }) .map_err(|e| SchedulerError::CronParseError(e.to_string()))?; let new_job_uuid = self .internal_scheduler .add(cron_task) .await .map_err(|e| SchedulerError::SchedulerInternalError(e.to_string()))?; // Update the job UUID and cron expression *job_uuid = new_job_uuid; job_def.cron = new_cron; self.persist_jobs_to_storage_with_guard(&jobs_guard).await?; Ok(()) } None => Err(SchedulerError::JobNotFound(sched_id.to_string())), } } pub async fn kill_running_job(&self, sched_id: &str) -> Result<(), SchedulerError> { let mut jobs_guard = self.jobs.lock().await; match jobs_guard.get_mut(sched_id) { Some((_, job_def)) => { if !job_def.currently_running { return Err(SchedulerError::AnyhowError(anyhow!( "Schedule '{}' is not currently running", sched_id ))); } tracing::info!("Killing running job '{}'", sched_id); // Abort the running task if it exists { let mut running_tasks_guard = self.running_tasks.lock().await; if let Some(abort_handle) = running_tasks_guard.remove(sched_id) { abort_handle.abort(); tracing::info!("Aborted running task for job '{}'", sched_id); } else { tracing::warn!( "No abort handle found for job '{}' in running tasks map", sched_id ); } } // Mark the job as no longer running job_def.currently_running = false; job_def.current_session_id = None; job_def.process_start_time = None; self.persist_jobs_to_storage_with_guard(&jobs_guard).await?; tracing::info!("Successfully killed job '{}'", sched_id); Ok(()) } None => Err(SchedulerError::JobNotFound(sched_id.to_string())), } } pub async fn get_running_job_info( &self, sched_id: &str, ) -> Result)>, SchedulerError> { let jobs_guard = self.jobs.lock().await; match jobs_guard.get(sched_id) { Some((_, job_def)) => { if job_def.currently_running { if let (Some(session_id), Some(start_time)) = (&job_def.current_session_id, &job_def.process_start_time) { Ok(Some((session_id.clone(), *start_time))) } else { Ok(None) } } else { Ok(None) } } None => Err(SchedulerError::JobNotFound(sched_id.to_string())), } } } #[derive(Debug)] struct JobExecutionError { job_id: String, error: String, } async fn run_scheduled_job_internal( job: ScheduledJob, provider_override: Option>, jobs_arc: Option>>, job_id: Option, ) -> std::result::Result { tracing::info!("Executing job: {} (Source: {})", job.id, job.source); let recipe_path = Path::new(&job.source); let recipe_content = match fs::read_to_string(recipe_path) { Ok(content) => content, Err(e) => { return Err(JobExecutionError { job_id: job.id.clone(), error: format!("Failed to load recipe file '{}': {}", job.source, e), }); } }; let recipe: Recipe = { let extension = recipe_path .extension() .and_then(|os_str| os_str.to_str()) .unwrap_or("yaml") .to_lowercase(); match extension.as_str() { "json" | "jsonl" => { serde_json::from_str::(&recipe_content).map_err(|e| JobExecutionError { job_id: job.id.clone(), error: format!("Failed to parse JSON recipe '{}': {}", job.source, e), }) } "yaml" | "yml" => { serde_yaml::from_str::(&recipe_content).map_err(|e| JobExecutionError { job_id: job.id.clone(), error: format!("Failed to parse YAML recipe '{}': {}", job.source, e), }) } _ => Err(JobExecutionError { job_id: job.id.clone(), error: format!( "Unsupported recipe file extension '{}' for: {}", extension, job.source ), }), } }?; let agent: Agent = Agent::new(); let agent_provider: Arc; if let Some(provider) = provider_override { agent_provider = provider; } else { let global_config = Config::global(); let provider_name: String = match global_config.get_param("GOOSE_PROVIDER") { Ok(name) => name, Err(_) => return Err(JobExecutionError { job_id: job.id.clone(), error: "GOOSE_PROVIDER not configured globally. Run 'goose configure' or set env var." .to_string(), }), }; let model_name: String = match global_config.get_param("GOOSE_MODEL") { Ok(name) => name, Err(_) => return Err(JobExecutionError { job_id: job.id.clone(), error: "GOOSE_MODEL not configured globally. Run 'goose configure' or set env var." .to_string(), }), }; let model_config = crate::model::ModelConfig::new(model_name.as_str()).map_err(|e| JobExecutionError { job_id: job.id.clone(), error: format!("Model config error: {}", e), })?; agent_provider = create(&provider_name, model_config) .await .map_err(|e| JobExecutionError { job_id: job.id.clone(), error: format!( "Failed to create provider instance '{}': {}", provider_name, e ), })?; } if let Some(ref recipe_extensions) = recipe.extensions { for extension in recipe_extensions { agent .add_extension(extension.clone()) .await .map_err(|e| JobExecutionError { job_id: job.id.clone(), error: format!("Failed to add extension '{}': {}", extension.name(), e), })?; } } if let Err(e) = agent.update_provider(agent_provider).await { return Err(JobExecutionError { job_id: job.id.clone(), error: format!("Failed to set provider on agent: {}", e), }); } tracing::info!("Agent configured with provider for job '{}'", job.id); let execution_mode = job.execution_mode.as_deref().unwrap_or("background"); tracing::info!("Job '{}' running in {} mode", job.id, execution_mode); let current_dir = match std::env::current_dir() { Ok(cd) => cd, Err(e) => { return Err(JobExecutionError { job_id: job.id.clone(), error: format!("Failed to get current directory for job execution: {}", e), }); } }; // Create session upfront for both cases let session = match SessionManager::create_session( current_dir.clone(), if recipe.prompt.is_some() { format!("Scheduled job: {}", job.id) } else { "Empty job - no prompt".to_string() }, ) .await { Ok(s) => s, Err(e) => { return Err(JobExecutionError { job_id: job.id.clone(), error: format!("Failed to create session: {}", e), }); } }; // Update the job with the session ID if we have access to the jobs arc if let (Some(jobs_arc), Some(job_id_str)) = (jobs_arc.as_ref(), job_id.as_ref()) { let mut jobs_guard = jobs_arc.lock().await; if let Some((_, job_def)) = jobs_guard.get_mut(job_id_str) { job_def.current_session_id = Some(session.id.clone()); } } if let Some(ref prompt_text) = recipe.prompt { let mut conversation = Conversation::new_unvalidated(vec![Message::user().with_text(prompt_text.clone())]); let session_config = SessionConfig { id: session.id.clone(), working_dir: current_dir.clone(), schedule_id: Some(job.id.clone()), execution_mode: job.execution_mode.clone(), max_turns: None, retry_config: None, }; match agent .reply(conversation.clone(), Some(session_config.clone()), None) .await { Ok(mut stream) => { use futures::StreamExt; while let Some(message_result) = stream.next().await { tokio::task::yield_now().await; match message_result { Ok(AgentEvent::Message(msg)) => { if msg.role == rmcp::model::Role::Assistant { tracing::info!("[Job {}] Assistant: {:?}", job.id, msg.content); } conversation.push(msg); } Ok(AgentEvent::McpNotification(_)) => {} Ok(AgentEvent::ModelChange { .. }) => {} Ok(AgentEvent::HistoryReplaced(updated_conversation)) => { conversation = updated_conversation; } Err(e) => { tracing::error!( "[Job {}] Error receiving message from agent: {}", job.id, e ); break; } } } } Err(e) => { return Err(JobExecutionError { job_id: job.id.clone(), error: format!("Agent failed to reply for recipe '{}': {}", job.source, e), }); } } } else { tracing::warn!( "[Job {}] Recipe '{}' has no prompt to execute.", job.id, job.source ); } if let Err(e) = SessionManager::update_session(&session.id) .schedule_id(Some(job.id.clone())) .recipe(Some(recipe)) .apply() .await { tracing::error!("[Job {}] Failed to update session metadata: {}", job.id, e); } tracing::info!("Finished job: {}", job.id); Ok(session.id) } #[async_trait] impl SchedulerTrait for Scheduler { async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError> { self.add_scheduled_job(job).await } async fn list_scheduled_jobs(&self) -> Result, SchedulerError> { Ok(self.list_scheduled_jobs().await) } async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> { self.remove_scheduled_job(id).await } async fn pause_schedule(&self, id: &str) -> Result<(), SchedulerError> { self.pause_schedule(id).await } async fn unpause_schedule(&self, id: &str) -> Result<(), SchedulerError> { self.unpause_schedule(id).await } async fn run_now(&self, id: &str) -> Result { self.run_now(id).await } async fn sessions( &self, sched_id: &str, limit: usize, ) -> Result, SchedulerError> { self.sessions(sched_id, limit).await } async fn update_schedule( &self, sched_id: &str, new_cron: String, ) -> Result<(), SchedulerError> { self.update_schedule(sched_id, new_cron).await } async fn kill_running_job(&self, sched_id: &str) -> Result<(), SchedulerError> { self.kill_running_job(sched_id).await } async fn get_running_job_info( &self, sched_id: &str, ) -> Result)>, SchedulerError> { self.get_running_job_info(sched_id).await } } #[cfg(test)] mod tests { use super::*; use crate::recipe::Recipe; use crate::{ model::ModelConfig, providers::base::{ProviderMetadata, ProviderUsage, Usage}, providers::errors::ProviderError, }; use rmcp::model::Tool; use rmcp::model::{AnnotateAble, RawTextContent, Role}; use crate::conversation::message::{Message, MessageContent}; use std::env; use std::fs::{self, File}; use std::io::Write; use tempfile::tempdir; #[derive(Clone)] struct MockSchedulerTestProvider { model_config: ModelConfig, } #[async_trait::async_trait] impl GooseProvider for MockSchedulerTestProvider { fn metadata() -> ProviderMetadata { ProviderMetadata::new( "mock-scheduler-test", "Mock for Scheduler Test", "A mock provider for scheduler tests", // description "test-model", // default_model vec!["test-model"], // model_names "", // model_doc_link (empty string if not applicable) vec![], // config_keys (empty vec if none) ) } fn get_model_config(&self) -> ModelConfig { self.model_config.clone() } async fn complete_with_model( &self, _model_config: &ModelConfig, _system: &str, _messages: &[Message], _tools: &[Tool], ) -> Result<(Message, ProviderUsage), ProviderError> { Ok(( Message::new( Role::Assistant, Utc::now().timestamp(), vec![MessageContent::Text( RawTextContent { text: "Mocked scheduled response".to_string(), meta: None, } .no_annotation(), )], ), ProviderUsage::new("mock-scheduler-test".to_string(), Usage::default()), )) } } // This function is pub(super) making it visible to run_scheduled_job_internal (parent module) // when cfg(test) is active for the whole compilation unit. pub(super) fn create_scheduler_test_mock_provider( model_config: ModelConfig, ) -> Arc { Arc::new(MockSchedulerTestProvider { model_config }) } #[tokio::test] async fn test_scheduled_session_has_schedule_id() -> Result<(), Box> { // Set environment variables for the test env::set_var("GOOSE_PROVIDER", "test_provider"); env::set_var("GOOSE_MODEL", "test_model"); let temp_dir = tempdir()?; let recipe_dir = temp_dir.path().join("recipes_for_test_scheduler"); fs::create_dir_all(&recipe_dir)?; let _ = crate::session::session_manager::ensure_session_dir() .expect("Failed to ensure app session dir"); let schedule_id_str = "test_schedule_001_scheduler_check".to_string(); let recipe_filename = recipe_dir.join(format!("{}.json", schedule_id_str)); let dummy_recipe = Recipe { version: "1.0.0".to_string(), title: "Test Schedule ID Recipe".to_string(), description: "A recipe for testing schedule_id propagation.".to_string(), instructions: None, prompt: Some("This is a test prompt for a scheduled job.".to_string()), extensions: None, context: None, activities: None, author: None, parameters: None, settings: None, response: None, sub_recipes: None, retry: None, }; let mut recipe_file = File::create(&recipe_filename)?; writeln!( recipe_file, "{}", serde_json::to_string_pretty(&dummy_recipe)? )?; recipe_file.flush()?; drop(recipe_file); let dummy_job = ScheduledJob { id: schedule_id_str.clone(), source: recipe_filename.to_string_lossy().into_owned(), cron: "* * * * * * ".to_string(), // Runs every second for quick testing last_run: None, currently_running: false, paused: false, current_session_id: None, process_start_time: None, execution_mode: Some("background".to_string()), // Default for test }; let mock_model_config = ModelConfig::new_or_fail("test_model"); let mock_provider_instance = create_scheduler_test_mock_provider(mock_model_config); // Call run_scheduled_job_internal, passing the mock provider let created_session_id = run_scheduled_job_internal(dummy_job.clone(), Some(mock_provider_instance), None, None) .await .expect("run_scheduled_job_internal failed"); let session = SessionManager::get_session(&created_session_id, true).await?; let schedule_id = session.schedule_id.clone(); assert_eq!( schedule_id, Some(schedule_id_str.clone()), "Session metadata schedule_id ({:?}) does not match the job ID ({}). Session: {}", schedule_id, schedule_id_str, created_session_id ); // Check if messages were written using SessionManager let messages_in_session = session.conversation.unwrap_or_default(); assert!( !messages_in_session.is_empty(), "No messages were written to the session: {}", created_session_id ); // We expect at least a user prompt and an assistant response assert!( messages_in_session.len() >= 2, "Expected at least 2 messages (prompt + response), found {} in session: {}", messages_in_session.len(), created_session_id ); // Clean up environment variables env::remove_var("GOOSE_PROVIDER"); env::remove_var("GOOSE_MODEL"); Ok(()) } }