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/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,