fix: let ctrl-c interrupt blocking job and agent collects
This commit is contained in:
6 files changed
+223
-3
No files matched your search
@@ -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
@@ -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
@@ -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 {
|
||||||
|
|||||||
@@ -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)?;
|
||||||
|
|||||||
@@ -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?;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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");
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user