From 3b01ae8f18892139e2e0b24e07a782ac585f455e Mon Sep 17 00:00:00 2001 From: Alex Clarke Date: Thu, 3 Sep 2026 13:03:10 -0600 Subject: [PATCH] fix: let ctrl-c interrupt blocking job and agent collects --- src/config/request_context.rs | 9 ++++ src/function/agents.rs | 90 ++++++++++++++++++++++++++++++++++- src/function/jobs.rs | 87 ++++++++++++++++++++++++++++++++- src/main.rs | 1 + src/repl/mod.rs | 1 + src/utils/abort_signal.rs | 38 +++++++++++++++ 6 files changed, 223 insertions(+), 3 deletions(-) diff --git a/src/config/request_context.rs b/src/config/request_context.rs index 38172ca..edae565 100644 --- a/src/config/request_context.rs +++ b/src/config/request_context.rs @@ -338,6 +338,11 @@ pub struct RequestContext { pub last_continuation_response: Option, pub pending_prefill: Option, + /// The turn's abort signal, set by the REPL/cmd entry points. Lets blocking + /// tools (job__collect, agent__collect) notice a user interrupt -- SIGINT is + /// otherwise swallowed in cooked mode during tool execution. + pub session_abort: Option, + pub render_mode: RenderMode, } @@ -375,6 +380,7 @@ impl RequestContext { skill_registry: SkillRegistry::default(), last_continuation_response: None, pending_prefill: None, + session_abort: None, render_mode: RenderMode::default(), } } @@ -439,6 +445,7 @@ impl RequestContext { skill_registry: SkillRegistry::default(), last_continuation_response: None, pending_prefill: None, + session_abort: None, render_mode: RenderMode::default(), }; ctx.refresh_mcp_tool_filters(); @@ -492,6 +499,7 @@ impl RequestContext { skill_registry: self.skill_registry.clone(), last_continuation_response: None, pending_prefill: None, + session_abort: self.session_abort.clone(), render_mode: self.render_mode, } } @@ -541,6 +549,7 @@ impl RequestContext { skill_registry: SkillRegistry::default(), last_continuation_response: None, pending_prefill: None, + session_abort: None, render_mode: parent.render_mode, } } diff --git a/src/function/agents.rs b/src/function/agents.rs index 0743611..a03f49a 100644 --- a/src/function/agents.rs +++ b/src/function/agents.rs @@ -7,7 +7,7 @@ use crate::config::{ use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox}; use crate::supervisor::notification::agent_notification; use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor, TaskKind}; -use crate::utils::{AbortSignal, create_abort_signal, wait_abort_signal}; +use crate::utils::{AbortSignal, create_abort_signal, wait_abort_signal, wait_user_interrupt}; use crate::graph; use anyhow::{Context, Result, anyhow, bail}; @@ -1031,6 +1031,14 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result sup.abort_signal_for(id) }; + let interrupted = || { + json!({ + "status": "interrupted", + "id": id, + "message": format!("agent__collect was interrupted by the user; agent '{id}' is still running in the background. Collect it again later, cancel it with agent__cancel, or continue with other work."), + }) + }; + loop { let is_finished = { let sup = supervisor.read(); @@ -1041,6 +1049,10 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result break; } + if ctx.session_abort.as_ref().is_some_and(|s| s.aborted()) { + return Ok(interrupted()); + } + if let Some(queue) = ctx.root_escalation_queue() && queue.has_pending() { @@ -1068,10 +1080,18 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result tokio::select! { _ = time::sleep(Duration::from_millis(200)) => {} _ = wait_abort_signal(abort) => {} + _ = wait_user_interrupt(ctx.session_abort.as_ref()) => { + return Ok(interrupted()); + } } } None => { - time::sleep(Duration::from_millis(200)).await; + tokio::select! { + _ = time::sleep(Duration::from_millis(200)) => {} + _ = wait_user_interrupt(ctx.session_abort.as_ref()) => { + return Ok(interrupted()); + } + } } } } @@ -2557,6 +2577,72 @@ mod tests { ); } + #[test] + fn handle_collect_session_interrupt_returns_and_keeps_agent() { + run_async(async { + let mut ctx = ctx_with_supervisor(4, 3); + let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + let join_handle = tokio::spawn(async move { + let _ = rx.await; + Ok(AgentResult { + id: "a1".to_string(), + agent_name: "explore".to_string(), + output: "late output".to_string(), + exit_status: AgentExitStatus::Completed, + }) + }); + let handle = AgentHandle { + id: "a1".to_string(), + agent_name: "explore".to_string(), + depth: 1, + inbox: Arc::new(Inbox::new()), + abort_signal: create_abort_signal(), + join_handle, + child_supervisor: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let session = create_abort_signal(); + session.set_ctrlc(); + ctx.session_abort = Some(session.clone()); + + let result = time::timeout( + Duration::from_secs(2), + handle_collect(&mut ctx, &json!({"id": "a1"})), + ) + .await + .expect("interrupted collect must return promptly") + .unwrap(); + + assert_eq!(result["status"], "interrupted"); + assert_eq!(result["id"], "a1"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("still running") + ); + assert_eq!( + ctx.supervisor.as_ref().unwrap().read().is_finished("a1"), + Some(false) + ); + + session.reset(); + tx.send(()).unwrap(); + + let collected = handle_collect(&mut ctx, &json!({"id": "a1"})) + .await + .unwrap(); + assert_eq!(collected["status"], "completed"); + assert_eq!(collected["output"], "late output"); + }); + } + #[test] fn handle_collect_pending_escalations_early_out_keeps_handle() { let rt = tokio::runtime::Builder::new_current_thread() diff --git a/src/function/jobs.rs b/src/function/jobs.rs index 9e35e80..b191051 100644 --- a/src/function/jobs.rs +++ b/src/function/jobs.rs @@ -16,7 +16,9 @@ use crate::mcp::{ }; use crate::supervisor::notification::job_notification; use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor}; -use crate::utils::{create_abort_signal, muted_warning_text, temp_file, wait_abort_signal}; +use crate::utils::{ + create_abort_signal, muted_warning_text, temp_file, wait_abort_signal, wait_user_interrupt, +}; use anyhow::{Context, Result, anyhow, bail}; use indexmap::IndexMap; @@ -623,6 +625,14 @@ async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result { job.abort_signal.clone() }; + let interrupted = || { + json!({ + "status": "interrupted", + "id": id, + "message": format!("job__collect was interrupted by the user; job '{id}' is still running in the background. Collect it again later, cancel it with job__cancel, or continue with other work."), + }) + }; + loop { let is_finished = { let sup = supervisor.read(); @@ -633,6 +643,10 @@ async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result { break; } + if ctx.session_abort.as_ref().is_some_and(|s| s.aborted()) { + return Ok(interrupted()); + } + if let Some(queue) = ctx.root_escalation_queue() && queue.has_pending() { @@ -666,6 +680,9 @@ async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result { tokio::select! { _ = time::sleep(Duration::from_millis(200)) => {} _ = wait_abort_signal(&target_abort) => {} + _ = wait_user_interrupt(ctx.session_abort.as_ref()) => { + return Ok(interrupted()); + } } } @@ -2140,6 +2157,74 @@ mod tests { }); } + #[test] + fn handle_collect_session_interrupt_returns_and_keeps_job() { + run_async(async { + let mut ctx = ctx_with_job_supervisor(4); + let state = Arc::new(Mutex::new(JobState { + status: JobStatus::Running, + pgid: None, + })); + let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + let task_state = Arc::clone(&state); + let join_handle = tokio::spawn(async move { + let _ = rx.await; + task_state.lock().status = JobStatus::Completed; + Ok(JobResult { + output: json!("late result"), + exit_code: Some(0), + output_bytes_captured: 0, + }) + }); + let handle = JobHandle { + id: "j1".to_string(), + tool: "execute_command".to_string(), + started_at: Instant::now(), + join_handle, + abort_signal: create_abort_signal(), + state, + output_buf: Arc::new(Mutex::new(RingBuf::default())), + no_change_checks: 0, + last_check_state: None, + }; + ctx.supervisor + .as_ref() + .unwrap() + .write() + .register(handle) + .unwrap(); + + let session = create_abort_signal(); + session.set_ctrlc(); + ctx.session_abort = Some(session.clone()); + + let result = time::timeout( + Duration::from_secs(2), + handle_collect(&ctx, &json!({"id": "j1"})), + ) + .await + .expect("interrupted collect must return promptly") + .unwrap(); + + assert_eq!(result["status"], "interrupted"); + assert_eq!(result["id"], "j1"); + assert!( + result["message"] + .as_str() + .unwrap() + .contains("still running") + ); + assert!(ctx.supervisor.as_ref().unwrap().read().has_job("j1")); + + session.reset(); + tx.send(()).unwrap(); + + let collected = handle_collect(&ctx, &json!({"id": "j1"})).await.unwrap(); + assert_eq!(collected["status"], "completed"); + assert_eq!(collected["result"], "late result"); + }); + } + #[test] fn handle_collect_full_result_returns_uncapped() { run_async(async { diff --git a/src/main.rs b/src/main.rs index 7181aa3..6db2fa9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -564,6 +564,7 @@ async fn start_directive( abort_signal: AbortSignal, ) -> Result<()> { let app: Arc = Arc::clone(&ctx.app.config); + ctx.session_abort = Some(abort_signal.clone()); if graph::active_agent_graph_name(ctx).is_some() { ctx.before_chat_completion(&input)?; diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 847e896..a57efd9 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -1432,6 +1432,7 @@ async fn ask( if input.is_empty() { return Ok(()); } + ctx.session_abort = Some(abort_signal.clone()); if with_embeddings { input.use_embeddings(abort_signal.clone()).await?; } diff --git a/src/utils/abort_signal.rs b/src/utils/abort_signal.rs index 36f4bf1..eebca96 100644 --- a/src/utils/abort_signal.rs +++ b/src/utils/abort_signal.rs @@ -68,6 +68,29 @@ pub async fn wait_abort_signal(abort_signal: &AbortSignal) { } } +/// Completes when the user interrupts: a SIGINT arrives, or the session's turn +/// signal is (or becomes) aborted. Needed because tools run in cooked mode -- +/// tokio's global SIGINT handler (installed by the first `ctrl_c()` poll) +/// swallows the signal with nobody listening. On SIGINT, sets ctrl-c on the +/// session signal so the surrounding turn aborts too. With `None`, only the +/// SIGINT arm applies. +pub async fn wait_user_interrupt(session: Option<&AbortSignal>) { + let ctrl_c = async { + if tokio::signal::ctrl_c().await.is_err() { + std::future::pending::<()>().await; + } + }; + match session { + Some(signal) => { + tokio::select! { + _ = ctrl_c => signal.set_ctrlc(), + _ = wait_abort_signal(signal) => {} + } + } + None => ctrl_c.await, + } +} + pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result { if event::poll(Duration::from_millis(25))? && let Event::Key(key) = event::read()? @@ -86,3 +109,18 @@ pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result { } Ok(false) } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn wait_user_interrupt_returns_promptly_on_preset_session_signal() { + let signal = create_abort_signal(); + signal.set_ctrlc(); + + tokio::time::timeout(Duration::from_secs(1), wait_user_interrupt(Some(&signal))) + .await + .expect("must return promptly when the session signal is already set"); + } +}