fix: let ctrl-c interrupt blocking job and agent collects

This commit is contained in:
Dark-Alex-17 committed 2026-09-03 13:03:10 -06:00
1 parent dc41e9d5df
commit 3b01ae8f18
6 files changed
+223 -3

No files matched your search

+9
View File
@@ -338,6 +338,11 @@ pub struct RequestContext {
pub last_continuation_response: Option<String>, pub last_continuation_response: Option<String>,
pub pending_prefill: Option<String>, pub pending_prefill: Option<String>,
/// 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<AbortSignal>,
pub render_mode: RenderMode, pub render_mode: RenderMode,
} }
@@ -375,6 +380,7 @@ impl RequestContext {
skill_registry: SkillRegistry::default(), skill_registry: SkillRegistry::default(),
last_continuation_response: None, last_continuation_response: None,
pending_prefill: None, pending_prefill: None,
session_abort: None,
render_mode: RenderMode::default(), render_mode: RenderMode::default(),
} }
} }
@@ -439,6 +445,7 @@ impl RequestContext {
skill_registry: SkillRegistry::default(), skill_registry: SkillRegistry::default(),
last_continuation_response: None, last_continuation_response: None,
pending_prefill: None, pending_prefill: None,
session_abort: None,
render_mode: RenderMode::default(), render_mode: RenderMode::default(),
}; };
ctx.refresh_mcp_tool_filters(); ctx.refresh_mcp_tool_filters();
@@ -492,6 +499,7 @@ impl RequestContext {
skill_registry: self.skill_registry.clone(), skill_registry: self.skill_registry.clone(),
last_continuation_response: None, last_continuation_response: None,
pending_prefill: None, pending_prefill: None,
session_abort: self.session_abort.clone(),
render_mode: self.render_mode, render_mode: self.render_mode,
} }
} }
@@ -541,6 +549,7 @@ impl RequestContext {
skill_registry: SkillRegistry::default(), skill_registry: SkillRegistry::default(),
last_continuation_response: None, last_continuation_response: None,
pending_prefill: None, pending_prefill: None,
session_abort: None,
render_mode: parent.render_mode, render_mode: parent.render_mode,
} }
} }
+88 -2
View File
@@ -7,7 +7,7 @@ use crate::config::{
use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox}; use crate::supervisor::mailbox::{Envelope, EnvelopePayload, Inbox};
use crate::supervisor::notification::agent_notification; use crate::supervisor::notification::agent_notification;
use crate::supervisor::{AgentExitStatus, AgentHandle, AgentResult, Supervisor, TaskKind}; 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 crate::graph;
use anyhow::{Context, Result, anyhow, bail}; use anyhow::{Context, Result, anyhow, bail};
@@ -1031,6 +1031,14 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
sup.abort_signal_for(id) 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 { loop {
let is_finished = { let is_finished = {
let sup = supervisor.read(); let sup = supervisor.read();
@@ -1041,6 +1049,10 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
break; break;
} }
if ctx.session_abort.as_ref().is_some_and(|s| s.aborted()) {
return Ok(interrupted());
}
if let Some(queue) = ctx.root_escalation_queue() if let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending() && queue.has_pending()
{ {
@@ -1068,10 +1080,18 @@ async fn handle_collect(ctx: &mut RequestContext, args: &Value) -> Result<Value>
tokio::select! { tokio::select! {
_ = time::sleep(Duration::from_millis(200)) => {} _ = time::sleep(Duration::from_millis(200)) => {}
_ = wait_abort_signal(abort) => {} _ = wait_abort_signal(abort) => {}
_ = wait_user_interrupt(ctx.session_abort.as_ref()) => {
return Ok(interrupted());
}
} }
} }
None => { 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] #[test]
fn handle_collect_pending_escalations_early_out_keeps_handle() { fn handle_collect_pending_escalations_early_out_keeps_handle() {
let rt = tokio::runtime::Builder::new_current_thread() let rt = tokio::runtime::Builder::new_current_thread()
+86 -1
View File
@@ -16,7 +16,9 @@ use crate::mcp::{
}; };
use crate::supervisor::notification::job_notification; use crate::supervisor::notification::job_notification;
use crate::supervisor::{JobHandle, JobResult, JobState, JobStatus, Supervisor}; 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 anyhow::{Context, Result, anyhow, bail};
use indexmap::IndexMap; use indexmap::IndexMap;
@@ -623,6 +625,14 @@ async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result<Value> {
job.abort_signal.clone() 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 { loop {
let is_finished = { let is_finished = {
let sup = supervisor.read(); let sup = supervisor.read();
@@ -633,6 +643,10 @@ async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result<Value> {
break; break;
} }
if ctx.session_abort.as_ref().is_some_and(|s| s.aborted()) {
return Ok(interrupted());
}
if let Some(queue) = ctx.root_escalation_queue() if let Some(queue) = ctx.root_escalation_queue()
&& queue.has_pending() && queue.has_pending()
{ {
@@ -666,6 +680,9 @@ async fn handle_collect(ctx: &RequestContext, args: &Value) -> Result<Value> {
tokio::select! { tokio::select! {
_ = time::sleep(Duration::from_millis(200)) => {} _ = time::sleep(Duration::from_millis(200)) => {}
_ = wait_abort_signal(&target_abort) => {} _ = 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] #[test]
fn handle_collect_full_result_returns_uncapped() { fn handle_collect_full_result_returns_uncapped() {
run_async(async { run_async(async {
+1
View File
@@ -564,6 +564,7 @@ async fn start_directive(
abort_signal: AbortSignal, abort_signal: AbortSignal,
) -> Result<()> { ) -> Result<()> {
let app: Arc<AppConfig> = Arc::clone(&ctx.app.config); let app: Arc<AppConfig> = Arc::clone(&ctx.app.config);
ctx.session_abort = Some(abort_signal.clone());
if graph::active_agent_graph_name(ctx).is_some() { if graph::active_agent_graph_name(ctx).is_some() {
ctx.before_chat_completion(&input)?; ctx.before_chat_completion(&input)?;
+1
View File
@@ -1432,6 +1432,7 @@ async fn ask(
if input.is_empty() { if input.is_empty() {
return Ok(()); return Ok(());
} }
ctx.session_abort = Some(abort_signal.clone());
if with_embeddings { if with_embeddings {
input.use_embeddings(abort_signal.clone()).await?; input.use_embeddings(abort_signal.clone()).await?;
} }
+38
View File
@@ -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<bool> { pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result<bool> {
if event::poll(Duration::from_millis(25))? if event::poll(Duration::from_millis(25))?
&& let Event::Key(key) = event::read()? && let Event::Key(key) = event::read()?
@@ -86,3 +109,18 @@ pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result<bool> {
} }
Ok(false) 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");
}
}