tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

Scheduler cleanup (#5571)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2025-11-08 13:19:18 -05:00
committed by GitHub
parent 25dfd768e5
commit e5a1474b4c
21 changed files with 1715 additions and 3902 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/schedule.rs
+7 -19
View File
@@ -1,21 +1,10 @@
use anyhow::{bail, Context, Result};
use base64::engine::{general_purpose::STANDARD as BASE64_STANDARD, Engine};
use goose::scheduler::{
get_default_scheduled_recipes_dir, get_default_scheduler_storage_path, ScheduledJob,
get_default_scheduled_recipes_dir, get_default_scheduler_storage_path, ScheduledJob, Scheduler,
SchedulerError,
};
use goose::scheduler_factory::SchedulerFactory;
use std::path::Path;
// Base64 decoding function - might be needed if recipe_source_arg can be base64
// For now, handle_schedule_add will assume it's a path.
async fn _decode_base64_recipe(source: &str) -> Result<String> {
let bytes = BASE64_STANDARD
.decode(source.as_bytes())
.with_context(|| "Recipe source is not a valid path and not valid Base64.")?;
String::from_utf8(bytes).with_context(|| "Decoded Base64 recipe source is not valid UTF-8.")
}
fn validate_cron_expression(cron: &str) -> Result<()> {
// Basic validation and helpful suggestions
if cron.trim().is_empty() {
@@ -84,7 +73,6 @@ pub async fn handle_schedule_add(
schedule_id, cron, recipe_source_arg
);
// Validate cron expression and provide helpful feedback
validate_cron_expression(&cron)?;
// The Scheduler's add_scheduled_job will handle copying the recipe from recipe_source_arg
@@ -102,7 +90,7 @@ pub async fn handle_schedule_add(
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let scheduler = SchedulerFactory::create(scheduler_storage_path)
let scheduler = Scheduler::new(scheduler_storage_path)
.await
.context("Failed to initialize scheduler")?;
@@ -148,11 +136,11 @@ pub async fn handle_schedule_add(
pub async fn handle_schedule_list() -> Result<()> {
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let scheduler = SchedulerFactory::create(scheduler_storage_path)
let scheduler = Scheduler::new(scheduler_storage_path)
.await
.context("Failed to initialize scheduler")?;
let jobs = scheduler.list_scheduled_jobs().await?;
let jobs = scheduler.list_scheduled_jobs().await;
if jobs.is_empty() {
println!("No scheduled jobs found.");
} else {
@@ -183,7 +171,7 @@ pub async fn handle_schedule_list() -> Result<()> {
pub async fn handle_schedule_remove(schedule_id: String) -> Result<()> {
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let scheduler = SchedulerFactory::create(scheduler_storage_path)
let scheduler = Scheduler::new(scheduler_storage_path)
.await
.context("Failed to initialize scheduler")?;
@@ -210,7 +198,7 @@ pub async fn handle_schedule_remove(schedule_id: String) -> Result<()> {
pub async fn handle_schedule_sessions(schedule_id: String, limit: Option<usize>) -> Result<()> {
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let scheduler = SchedulerFactory::create(scheduler_storage_path)
let scheduler = Scheduler::new(scheduler_storage_path)
.await
.context("Failed to initialize scheduler")?;
@@ -246,7 +234,7 @@ pub async fn handle_schedule_sessions(schedule_id: String, limit: Option<usize>)
pub async fn handle_schedule_run_now(schedule_id: String) -> Result<()> {
let scheduler_storage_path =
get_default_scheduler_storage_path().context("Failed to get scheduler storage path")?;
let scheduler = SchedulerFactory::create(scheduler_storage_path)
let scheduler = Scheduler::new(scheduler_storage_path)
.await
.context("Failed to initialize scheduler")?;
crates/goose-server/src/routes/schedule.rs
+35 -50
View File
@@ -16,8 +16,6 @@ pub struct CreateScheduleRequest {
id: String,
recipe_source: String,
cron: String,
#[serde(default)]
execution_mode: Option<String>, // "foreground" or "background"
}
#[derive(Deserialize, Serialize, utoipa::ToSchema)]
@@ -36,7 +34,6 @@ pub struct KillJobResponse {
message: String,
}
// Response for the inspect endpoint
#[derive(Serialize, utoipa::ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct InspectJobResponse {
@@ -51,15 +48,9 @@ pub struct RunNowResponse {
session_id: String,
}
// Query parameters for the sessions endpoint
#[derive(Deserialize, utoipa::ToSchema, utoipa::IntoParams)]
pub struct SessionsQuery {
#[serde(default = "default_limit")]
limit: u32,
}
fn default_limit() -> u32 {
50 // Default limit for sessions listed
limit: usize,
}
// Struct for the frontend session list
@@ -151,10 +142,7 @@ async fn list_schedules(
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
tracing::info!("Server: Calling scheduler.list_scheduled_jobs()");
let jobs = scheduler.list_scheduled_jobs().await.map_err(|e| {
eprintln!("Error listing schedules: {:?}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
let jobs = scheduler.list_scheduled_jobs().await;
Ok(Json(ListSchedulesResponse { jobs }))
}
@@ -213,39 +201,40 @@ async fn run_now_handler(
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let (recipe_display_name, recipe_version_opt) = match scheduler.list_scheduled_jobs().await {
Ok(jobs) => {
if let Some(job) = jobs.into_iter().find(|job| job.id == id) {
let recipe_display_name = std::path::Path::new(&job.source)
.file_name()
.and_then(|name| name.to_str())
.map(|s| s.to_string())
.unwrap_or_else(|| id.clone());
let (recipe_display_name, recipe_version_opt) = if let Some(job) = scheduler
.list_scheduled_jobs()
.await
.into_iter()
.find(|job| job.id == id)
{
let recipe_display_name = std::path::Path::new(&job.source)
.file_name()
.and_then(|name| name.to_str())
.map(|s| s.to_string())
.unwrap_or_else(|| id.clone());
let recipe_version_opt = tokio::fs::read_to_string(&job.source)
.await
let recipe_version_opt =
tokio::fs::read_to_string(&job.source)
.await
.ok()
.and_then(|content: String| {
goose::recipe::template_recipe::parse_recipe_content(
&content,
Some(
std::path::Path::new(&job.source)
.parent()
.unwrap_or_else(|| std::path::Path::new(""))
.to_string_lossy()
.to_string(),
),
)
.ok()
.and_then(|content| {
goose::recipe::template_recipe::parse_recipe_content(
&content,
Some(
std::path::Path::new(&job.source)
.parent()
.unwrap_or_else(|| std::path::Path::new(""))
.to_string_lossy()
.to_string(),
),
)
.ok()
.map(|(r, _)| r.version)
});
.map(|(r, _)| r.version)
});
(recipe_display_name, recipe_version_opt)
} else {
(id.clone(), None)
}
}
Err(_) => (id.clone(), None),
(recipe_display_name, recipe_version_opt)
} else {
(id.clone(), None)
};
let recipe_version_tag = recipe_version_opt.as_deref().unwrap_or("");
@@ -308,7 +297,7 @@ async fn sessions_handler(
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
match scheduler
.sessions(&schedule_id_param, query_params.limit as usize)
.sessions(&schedule_id_param, query_params.limit)
.await
{
Ok(session_tuples) => {
@@ -448,11 +437,7 @@ async fn update_schedule(
}
})?;
// Return the updated schedule
let jobs = scheduler.list_scheduled_jobs().await.map_err(|e| {
eprintln!("Error listing schedules after update: {:?}", e);
StatusCode::INTERNAL_SERVER_ERROR
})?;
let jobs = scheduler.list_scheduled_jobs().await;
let updated_job = jobs
.into_iter()
.find(|job| job.id == id)
crates/goose/src/agents/agent.rs
+1 -2
View File
@@ -39,7 +39,6 @@ use crate::permission::PermissionConfirmation;
use crate::providers::base::Provider;
use crate::providers::errors::ProviderError;
use crate::recipe::{Author, Recipe, Response, Settings, SubRecipe};
use crate::scheduler_trait::SchedulerTrait;
use crate::security::security_inspector::SecurityInspector;
use crate::tool_inspection::ToolInspectionManager;
use crate::tool_monitor::RepetitionInspector;
@@ -60,6 +59,7 @@ use super::platform_tools;
use super::tool_execution::{ToolCallResult, CHAT_MODE_TOOL_SKIPPED_RESPONSE, DECLINED_RESPONSE};
use crate::agents::subagent_task_config::TaskConfig;
use crate::conversation::message::{Message, MessageContent, SystemNotificationType, ToolRequest};
use crate::scheduler_trait::SchedulerTrait;
use crate::session::extension_data::{EnabledExtensionsState, ExtensionState};
use crate::session::{Session, SessionManager};
@@ -346,7 +346,6 @@ impl Agent {
Ok(tool_futures)
}
/// Set the scheduler service for this agent
pub async fn set_scheduler(&self, scheduler: Arc<dyn SchedulerTrait>) {
let mut scheduler_service = self.scheduler_service.lock().await;
*scheduler_service = Some(scheduler);
crates/goose/src/agents/schedule_tool.rs
+11 -35
View File
@@ -9,11 +9,10 @@ use crate::mcp_utils::ToolResult;
use chrono::Utc;
use rmcp::model::{Content, ErrorCode, ErrorData};
use super::Agent;
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(
@@ -62,34 +61,24 @@ impl Agent {
}
}
/// List all scheduled jobs
async fn handle_list_jobs(
&self,
scheduler: Arc<dyn SchedulerTrait>,
) -> ToolResult<Vec<Content>> {
match scheduler.list_scheduled_jobs().await {
Ok(jobs) => {
let jobs_json = serde_json::to_string_pretty(&jobs).map_err(|e| {
ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to serialize jobs: {}", e),
None,
)
})?;
Ok(vec![Content::text(format!(
"Scheduled Jobs:\n{}",
jobs_json
))])
}
Err(e) => Err(ErrorData::new(
let jobs = scheduler.list_scheduled_jobs().await;
let jobs_json = serde_json::to_string_pretty(&jobs).map_err(|e| {
ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!("Failed to list jobs: {}", e),
format!("Failed to serialize jobs: {}", e),
None,
)),
}
)
})?;
Ok(vec![Content::text(format!(
"Scheduled Jobs:\n{}",
jobs_json
))])
}
/// Create a new scheduled job from a recipe file
async fn handle_create_job(
&self,
scheduler: Arc<dyn SchedulerTrait>,
@@ -123,19 +112,6 @@ impl Agent {
.and_then(|v| v.as_str())
.unwrap_or("background");
// Validate execution_mode is either "foreground" or "background"
if execution_mode != "foreground" && execution_mode != "background" {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
format!(
"Invalid execution_mode: {}. Must be 'foreground' or 'background'",
execution_mode
),
None,
));
}
// Validate recipe file exists and is readable
if !std::path::Path::new(recipe_path).exists() {
return Err(ErrorData::new(
ErrorCode::INTERNAL_ERROR,
crates/goose/src/cron_test.rs
-56
View File
@@ -1,56 +0,0 @@
#[cfg(test)]
mod cron_parsing_tests {
use crate::scheduler::normalize_cron_expression;
use tokio_cron_scheduler::Job;
// Helper: drop the last field if we have 7 so tokio_cron_scheduler (6-field) can parse
fn to_tokio_spec(spec: &str) -> String {
let parts: Vec<&str> = spec.split_whitespace().collect();
if parts.len() == 7 {
parts[..6].join(" ")
} else {
spec.to_string()
}
}
#[test]
fn test_normalize_cron_expression() {
// 5-field → 7-field
assert_eq!(normalize_cron_expression("0 12 * * *"), "0 0 12 * * * *");
assert_eq!(normalize_cron_expression("*/5 * * * *"), "0 */5 * * * * *");
assert_eq!(normalize_cron_expression("0 0 * * 1"), "0 0 0 * * 1 *");
// 6-field → 7-field (append *)
assert_eq!(normalize_cron_expression("0 0 12 * * *"), "0 0 12 * * * *");
assert_eq!(
normalize_cron_expression("*/30 */5 * * * *"),
"*/30 */5 * * * * *"
);
// Weekday expressions (unchanged apart from 7-field format)
assert_eq!(normalize_cron_expression("0 * * * 1-5"), "0 0 * * * 1-5 *");
assert_eq!(
normalize_cron_expression("*/20 * * * 1-5"),
"0 */20 * * * 1-5 *"
);
}
#[tokio::test]
async fn test_cron_expression_formats() {
let samples = [
"0 0 * * *", // 5-field
"0 0 0 * * *", // 6-field
"*/5 * * * *", // 5-field
];
for expr in samples {
let norm = normalize_cron_expression(expr);
let tokio_spec = to_tokio_spec(&norm);
assert!(
Job::new_async(&tokio_spec, |_id, _l| Box::pin(async {})).is_ok(),
"failed to parse {} -> {}",
expr,
norm
);
}
}
}
crates/goose/src/execution/manager.rs
+2 -2
View File
@@ -1,7 +1,7 @@
use crate::agents::extension::PlatformExtensionContext;
use crate::agents::Agent;
use crate::config::paths::Paths;
use crate::scheduler_factory::SchedulerFactory;
use crate::scheduler::Scheduler;
use crate::scheduler_trait::SchedulerTrait;
use anyhow::Result;
use lru::LruCache;
@@ -35,7 +35,7 @@ impl AgentManager {
async fn new(max_sessions: Option<usize>) -> Result<Self> {
let schedule_file_path = Paths::data_dir().join("schedule.json");
let scheduler = SchedulerFactory::create(schedule_file_path).await?;
let scheduler = Scheduler::new(schedule_file_path).await?;
let capacity = NonZeroUsize::new(max_sessions.unwrap_or(DEFAULT_MAX_SESSION))
.unwrap_or_else(|| NonZeroUsize::new(100).unwrap());
crates/goose/src/lib.rs
-4
View File
@@ -14,7 +14,6 @@ pub mod providers;
pub mod recipe;
pub mod recipe_deeplink;
pub mod scheduler;
pub mod scheduler_factory;
pub mod scheduler_trait;
pub mod security;
pub mod session;
@@ -25,6 +24,3 @@ pub mod tool_inspection;
pub mod tool_monitor;
pub mod tracing;
pub mod utils;
#[cfg(test)]
mod cron_test;
crates/goose/src/scheduler.rs
+428 -1061
View File
File diff suppressed because it is too large Load Diff
crates/goose/src/scheduler_factory.rs
-26
View File
@@ -1,26 +0,0 @@
use std::path::PathBuf;
use std::sync::Arc;
use crate::scheduler::{Scheduler, SchedulerError};
use crate::scheduler_trait::SchedulerTrait;
/// Factory for creating scheduler instances
pub struct SchedulerFactory;
impl SchedulerFactory {
/// Create a scheduler instance
pub async fn create(storage_path: PathBuf) -> Result<Arc<dyn SchedulerTrait>, SchedulerError> {
tracing::info!("Creating scheduler");
let scheduler = Scheduler::new(storage_path).await?;
Ok(scheduler as Arc<dyn SchedulerTrait>)
}
/// Create a scheduler (for testing or explicit use)
pub async fn create_legacy(
storage_path: PathBuf,
) -> Result<Arc<dyn SchedulerTrait>, SchedulerError> {
tracing::info!("Creating scheduler (explicit)");
let scheduler = Scheduler::new(storage_path).await?;
Ok(scheduler as Arc<dyn SchedulerTrait>)
}
}
crates/goose/src/scheduler_trait.rs
+1 -21
View File
@@ -4,42 +4,22 @@ use chrono::{DateTime, Utc};
use crate::scheduler::{ScheduledJob, SchedulerError};
use crate::session::Session;
/// Common trait for all scheduler implementations
#[async_trait]
pub trait SchedulerTrait: Send + Sync {
/// Add a new scheduled job
async fn add_scheduled_job(&self, job: ScheduledJob) -> Result<(), SchedulerError>;
/// List all scheduled jobs
async fn list_scheduled_jobs(&self) -> Result<Vec<ScheduledJob>, SchedulerError>;
/// Remove a scheduled job by ID
async fn list_scheduled_jobs(&self) -> Vec<ScheduledJob>;
async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError>;
/// Pause a scheduled job
async fn pause_schedule(&self, id: &str) -> Result<(), SchedulerError>;
/// Unpause a scheduled job
async fn unpause_schedule(&self, id: &str) -> Result<(), SchedulerError>;
/// Run a job immediately
async fn run_now(&self, id: &str) -> Result<String, SchedulerError>;
/// Get sessions for a scheduled job
async fn sessions(
&self,
sched_id: &str,
limit: usize,
) -> Result<Vec<(String, Session)>, SchedulerError>;
/// Update a schedule's cron expression
async fn update_schedule(&self, sched_id: &str, new_cron: String)
-> Result<(), SchedulerError>;
/// Kill a running job
async fn kill_running_job(&self, sched_id: &str) -> Result<(), SchedulerError>;
/// Get information about a running job
async fn get_running_job_info(
&self,
sched_id: &str,
crates/goose/tests/agent.rs
+2 -2
View File
@@ -40,9 +40,9 @@ mod tests {
Ok(())
}
async fn list_scheduled_jobs(&self) -> Result<Vec<ScheduledJob>, SchedulerError> {
async fn list_scheduled_jobs(&self) -> Vec<ScheduledJob> {
let jobs = self.jobs.lock().await;
Ok(jobs.clone())
jobs.clone()
}
async fn remove_scheduled_job(&self, id: &str) -> Result<(), SchedulerError> {
crates/goose/tests/test_support.rs
-401
View File
@@ -1,401 +0,0 @@
#![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 goose::agents::Agent;
use goose::scheduler::{ScheduledJob, SchedulerError};
use goose::scheduler_trait::SchedulerTrait;
use goose::session::Session;
#[derive(Debug, Clone)]
pub enum MockBehavior {
Success,
NotFound(String),
AlreadyExists(String),
InternalError(String),
JobCurrentlyRunning(String),
}
#[derive(Clone)]
pub struct ConfigurableMockScheduler {
jobs: Arc<Mutex<HashMap<String, ScheduledJob>>>,
running_jobs: Arc<Mutex<HashSet<String>>>,
call_log: Arc<Mutex<Vec<String>>>,
behaviors: Arc<Mutex<HashMap<String, MockBehavior>>>,
#[allow(clippy::type_complexity)]
sessions_data: Arc<Mutex<HashMap<String, Vec<(String, Session)>>>>,
}
#[allow(dead_code)]
impl Default for ConfigurableMockScheduler {
fn default() -> Self {
Self::new()
}
}
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 get_calls(&self) -> Vec<String> {
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<Vec<ScheduledJob>, 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<String, SchedulerError> {
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<Vec<(String, Session)>, 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<Option<(String, DateTime<Utc>)>, 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<ConfigurableMockScheduler>,
}
impl Default for ScheduleToolTestBuilder {
fn default() -> Self {
Self::new()
}
}
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, Session)>) -> 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<ConfigurableMockScheduler>) {
let agent = Agent::new();
agent.set_scheduler(self.scheduler.clone()).await;
(agent, self.scheduler)
}
}
pub fn create_test_session_metadata(message_count: usize, working_dir: &str) -> Session {
Session {
id: "".to_string(),
working_dir: PathBuf::from(working_dir),
name: "Test session".to_string(),
user_set_name: false,
created_at: Default::default(),
schedule_id: Some("test_job".to_string()),
recipe: None,
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),
extension_data: Default::default(),
updated_at: Default::default(),
conversation: None,
message_count,
user_recipe_values: None,
session_type: Default::default(),
}
}