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

feat: implement proper task cancellation for scheduled jobs (#2731)

Browse Source
This commit is contained in:
Max Novich
2025-05-29 18:33:27 -07:00
committed by GitHub
parent a05029773d
commit bd430866e8
12 changed files with 2945 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/schedule.rs
+2
View File
@@ -34,6 +34,8 @@ pub async fn handle_schedule_add(
last_run: None,
currently_running: false,
paused: false,
current_session_id: None,
process_start_time: None,
};
let scheduler_storage_path =
crates/goose-server/src/openapi.rs
+4
View File
@@ -45,6 +45,8 @@ use utoipa::OpenApi;
super::routes::schedule::run_now_handler,
super::routes::schedule::pause_schedule,
super::routes::schedule::unpause_schedule,
super::routes::schedule::kill_running_job,
super::routes::schedule::inspect_running_job,
super::routes::schedule::sessions_handler
),
components(schemas(
@@ -95,6 +97,8 @@ use utoipa::OpenApi;
SessionMetadata,
super::routes::schedule::CreateScheduleRequest,
super::routes::schedule::UpdateScheduleRequest,
super::routes::schedule::KillJobResponse,
super::routes::schedule::InspectJobResponse,
goose::scheduler::ScheduledJob,
super::routes::schedule::RunNowResponse,
super::routes::schedule::ListSchedulesResponse,
crates/goose-server/src/routes/schedule.rs
+116
View File
@@ -31,6 +31,21 @@ pub struct ListSchedulesResponse {
jobs: Vec<ScheduledJob>,
}
// Response for the kill endpoint
#[derive(Serialize, utoipa::ToSchema)]
pub struct KillJobResponse {
message: String,
}
// Response for the inspect endpoint
#[derive(Serialize, utoipa::ToSchema)]
#[serde(rename_all = "camelCase")]
pub struct InspectJobResponse {
session_id: Option<String>,
process_start_time: Option<String>,
running_duration_seconds: Option<i64>,
}
// Response for the run_now endpoint
#[derive(Serialize, utoipa::ToSchema)]
pub struct RunNowResponse {
@@ -100,6 +115,8 @@ async fn create_schedule(
last_run: None,
currently_running: false,
paused: false,
current_session_id: None,
process_start_time: None,
};
scheduler
.add_scheduled_job(job.clone())
@@ -199,6 +216,17 @@ async fn run_now_handler(
eprintln!("Error running schedule '{}' now: {:?}", id, e);
match e {
goose::scheduler::SchedulerError::JobNotFound(_) => Err(StatusCode::NOT_FOUND),
goose::scheduler::SchedulerError::AnyhowError(ref err) => {
// Check if this is a cancellation error
if err.to_string().contains("was successfully cancelled") {
// Return a special session_id to indicate cancellation
Ok(Json(RunNowResponse {
session_id: "CANCELLED".to_string(),
}))
} else {
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
_ => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
@@ -389,6 +417,92 @@ async fn update_schedule(
Ok(Json(updated_job))
}
#[utoipa::path(
post,
path = "/schedule/{id}/kill",
responses(
(status = 200, description = "Running job killed successfully"),
),
tag = "schedule"
)]
#[axum::debug_handler]
pub async fn kill_running_job(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Path(id): Path<String>,
) -> Result<Json<KillJobResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
scheduler.kill_running_job(&id).await.map_err(|e| {
eprintln!("Error killing running job '{}': {:?}", id, e);
match e {
goose::scheduler::SchedulerError::JobNotFound(_) => StatusCode::NOT_FOUND,
goose::scheduler::SchedulerError::AnyhowError(_) => StatusCode::BAD_REQUEST,
_ => StatusCode::INTERNAL_SERVER_ERROR,
}
})?;
Ok(Json(KillJobResponse {
message: format!("Successfully killed running job '{}'", id),
}))
}
#[utoipa::path(
get,
path = "/schedule/{id}/inspect",
params(
("id" = String, Path, description = "ID of the schedule to inspect")
),
responses(
(status = 200, description = "Running job information", body = InspectJobResponse),
(status = 404, description = "Scheduled job not found"),
(status = 500, description = "Internal server error")
),
tag = "schedule"
)]
#[axum::debug_handler]
pub async fn inspect_running_job(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Path(id): Path<String>,
) -> Result<Json<InspectJobResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
let scheduler = state
.scheduler()
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
match scheduler.get_running_job_info(&id).await {
Ok(info) => {
if let Some((session_id, start_time)) = info {
let duration = chrono::Utc::now().signed_duration_since(start_time);
Ok(Json(InspectJobResponse {
session_id: Some(session_id),
process_start_time: Some(start_time.to_rfc3339()),
running_duration_seconds: Some(duration.num_seconds()),
}))
} else {
Ok(Json(InspectJobResponse {
session_id: None,
process_start_time: None,
running_duration_seconds: None,
}))
}
}
Err(e) => {
eprintln!("Error inspecting running job '{}': {:?}", id, e);
match e {
goose::scheduler::SchedulerError::JobNotFound(_) => Err(StatusCode::NOT_FOUND),
_ => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
}
}
pub fn routes(state: Arc<AppState>) -> Router {
Router::new()
.route("/schedule/create", post(create_schedule))
@@ -398,6 +512,8 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/schedule/{id}/run_now", post(run_now_handler)) // Corrected
.route("/schedule/{id}/pause", post(pause_schedule))
.route("/schedule/{id}/unpause", post(unpause_schedule))
.route("/schedule/{id}/kill", post(kill_running_job))
.route("/schedule/{id}/inspect", get(inspect_running_job))
.route("/schedule/{id}/sessions", get(sessions_handler)) // Corrected
.with_state(state)
}
crates/goose-server/ui/desktop/openapi.json
+1989 -152
View File
File diff suppressed because it is too large Load Diff
crates/goose/src/scheduler.rs
+291 -31
View File
@@ -20,6 +20,10 @@ use crate::recipe::Recipe;
use crate::session;
use crate::session::storage::SessionMetadata;
// Track running tasks with their abort handles
type RunningTasksMap = HashMap<String, tokio::task::AbortHandle>;
type JobsMap = HashMap<String, (JobId, ScheduledJob)>;
pub fn get_default_scheduler_storage_path() -> Result<PathBuf, io::Error> {
let strategy = choose_app_strategy(config::APP_STRATEGY.clone())
.map_err(|e| io::Error::new(io::ErrorKind::NotFound, e.to_string()))?;
@@ -111,11 +115,15 @@ pub struct ScheduledJob {
pub currently_running: bool,
#[serde(default)]
pub paused: bool,
#[serde(default)]
pub current_session_id: Option<String>,
#[serde(default)]
pub process_start_time: Option<DateTime<Utc>>,
}
async fn persist_jobs_from_arc(
storage_path: &Path,
jobs_arc: &Arc<Mutex<HashMap<String, (JobId, ScheduledJob)>>>,
jobs_arc: &Arc<Mutex<JobsMap>>,
) -> Result<(), SchedulerError> {
let jobs_guard = jobs_arc.lock().await;
let list: Vec<ScheduledJob> = jobs_guard.values().map(|(_, j)| j.clone()).collect();
@@ -129,8 +137,9 @@ async fn persist_jobs_from_arc(
pub struct Scheduler {
internal_scheduler: TokioJobScheduler,
jobs: Arc<Mutex<HashMap<String, (JobId, ScheduledJob)>>>,
jobs: Arc<Mutex<JobsMap>>,
storage_path: PathBuf,
running_tasks: Arc<Mutex<RunningTasksMap>>,
}
impl Scheduler {
@@ -140,11 +149,13 @@ impl Scheduler {
.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?;
@@ -208,17 +219,21 @@ impl Scheduler {
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();
let cron_task = Job::new_async(&stored_job.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
@@ -243,6 +258,7 @@ impl Scheduler {
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;
}
}
@@ -258,14 +274,37 @@ impl Scheduler {
);
}
}
// Pass None for provider_override in normal execution
let result = run_scheduled_job_internal(job_to_execute, None).await;
// 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;
}
}
@@ -282,12 +321,27 @@ impl Scheduler {
}
}
if let Err(e) = result {
tracing::error!(
"Scheduled job '{}' execution failed: {}",
&e.job_id,
e.error
);
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
);
}
}
})
})
@@ -328,12 +382,14 @@ impl Scheduler {
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();
let cron_task = Job::new_async(&job_to_load.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
@@ -358,6 +414,7 @@ impl Scheduler {
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;
}
}
@@ -373,14 +430,37 @@ impl Scheduler {
);
}
}
// Pass None for provider_override in normal execution
let result = run_scheduled_job_internal(job_to_execute, None).await;
// 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;
}
}
@@ -397,12 +477,30 @@ impl Scheduler {
}
}
if let Err(e) = result {
tracing::error!(
"Scheduled job '{}' execution failed: {}",
&e.job_id,
e.error
);
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
);
}
}
})
})
@@ -421,7 +519,7 @@ impl Scheduler {
// 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<'_, HashMap<String, (JobId, ScheduledJob)>>,
jobs_guard: &tokio::sync::MutexGuard<'_, JobsMap>,
) -> Result<(), SchedulerError> {
let list: Vec<ScheduledJob> = jobs_guard.values().map(|(_, j)| j.clone()).collect();
if let Some(parent) = self.storage_path.parent() {
@@ -523,14 +621,36 @@ impl Scheduler {
}
};
// Pass None for provider_override in normal execution
let run_result = run_scheduled_job_internal(job_to_run.clone(), None).await;
// 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
}
@@ -539,12 +659,24 @@ impl Scheduler {
self.persist_jobs().await?;
match run_result {
Ok(session_id) => Ok(session_id),
Err(e) => Err(SchedulerError::AnyhowError(anyhow!(
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
))),
}
}
@@ -608,12 +740,14 @@ impl Scheduler {
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();
let cron_task = Job::new_async(&new_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
@@ -641,6 +775,7 @@ impl Scheduler {
{
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;
}
}
@@ -657,7 +792,29 @@ impl Scheduler {
}
}
let result = run_scheduled_job_internal(job_to_execute, None).await;
// 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
{
@@ -666,6 +823,8 @@ impl Scheduler {
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;
}
}
@@ -682,12 +841,33 @@ impl Scheduler {
}
}
if let Err(e) = result {
tracing::error!(
"Scheduled job '{}' execution failed: {}",
&e.job_id,
e.error
);
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
);
}
}
})
})
@@ -709,6 +889,70 @@ impl Scheduler {
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<Option<(String, DateTime<Utc>)>, 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)]
@@ -720,6 +964,8 @@ struct JobExecutionError {
async fn run_scheduled_job_internal(
job: ScheduledJob,
provider_override: Option<Arc<dyn GooseProvider>>, // New optional parameter
jobs_arc: Option<Arc<Mutex<JobsMap>>>,
job_id: Option<String>,
) -> std::result::Result<String, JobExecutionError> {
tracing::info!("Executing job: {} (Source: {})", job.id, job.source);
@@ -811,6 +1057,15 @@ async fn run_scheduled_job_internal(
tracing::info!("Agent configured with provider for job '{}'", job.id);
let session_id_for_return = session::generate_session_id();
// 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_for_return.clone());
}
}
let session_file_path = crate::session::storage::get_path(
crate::session::storage::Identifier::Name(session_id_for_return.clone()),
);
@@ -843,6 +1098,9 @@ async fn run_scheduled_job_internal(
use futures::StreamExt;
while let Some(message_result) = stream.next().await {
// Check if the task has been cancelled
tokio::task::yield_now().await;
match message_result {
Ok(msg) => {
if msg.role == mcp_core::role::Role::Assistant {
@@ -1053,6 +1311,8 @@ mod tests {
last_run: None,
currently_running: false,
paused: false,
current_session_id: None,
process_start_time: None,
};
// Create the mock provider instance for the test
@@ -1061,7 +1321,7 @@ mod tests {
// Call run_scheduled_job_internal, passing the mock provider
let created_session_id =
run_scheduled_job_internal(dummy_job.clone(), Some(mock_provider_instance))
run_scheduled_job_internal(dummy_job.clone(), Some(mock_provider_instance), None, None)
.await
.expect("run_scheduled_job_internal failed");