diff --git a/src/config/agent.rs b/src/config/agent.rs index 5dd3f64..c67cd49 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -633,6 +633,7 @@ impl Agent { Some(self.name().to_string()), tool_timeout, false, + None, )?; match value { Some(v) => Ok(v), diff --git a/src/function/mod.rs b/src/function/mod.rs index 39f399c..d80ae25 100644 --- a/src/function/mod.rs +++ b/src/function/mod.rs @@ -1684,6 +1684,7 @@ impl ToolCall { agent_name, ctx.app.config.tool_timeout, quiet, + ctx.session_abort.clone(), ) { Ok(Some(contents)) => serde_json::from_str(&contents) .ok() @@ -2293,6 +2294,7 @@ pub fn run_llm_function( agent_name: Option, tool_timeout: Option, quiet: bool, + abort: Option, ) -> Result> { let mut bin_dirs: Vec = vec![]; let mut command_name = cmd_name.clone(); @@ -2440,6 +2442,21 @@ pub fn run_llm_function( return Ok(Some(error_json.to_string())); } + if abort.as_ref().is_some_and(|a| a.aborted()) { + let _ = child.kill(); + let _ = child.wait(); + drop(stdout_thread); + drop(stderr_thread); + let tool_error_message = format!("Tool call '{command_name}' aborted"); + emit_tool_warning( + quiet, + &format!("⚠️ {tool_error_message} ⚠️"), + &tool_error_message, + ); + return Ok(Some( + json!({"tool_call_error": tool_error_message}).to_string(), + )); + } thread::sleep(Duration::from_millis(100)); }; let stdout_bytes = stdout_thread.join().unwrap_or_default(); @@ -4562,6 +4579,7 @@ mod tests { None, None, false, + None, ) .unwrap() .expect("nonzero exit must return an error payload"); @@ -4590,6 +4608,7 @@ mod tests { None, None, true, + None, ) .unwrap() .expect("nonzero exit must return an error payload"); @@ -4605,6 +4624,38 @@ mod tests { assert_eq!(json["output"], "partial-output\n"); } + /// An already-aborted signal is noticed on the first wait-loop poll, so a + /// tool that would run for seconds is killed and reported within one poll + /// interval instead of running to completion or `tool_timeout`. + #[cfg(unix)] + #[test] + fn run_llm_function_aborts_running_tool() { + if which::which("bash").is_err() { + eprintln!("skipping: bash not available"); + return; + } + let abort = create_abort_signal(); + abort.set_ctrlc(); + + let started = Instant::now(); + let result = run_llm_function( + "bash".into(), + vec!["-c".into(), "sleep 5".into()], + HashMap::new(), + None, + None, + true, + Some(abort), + ) + .unwrap() + .expect("an aborted tool must return an error payload"); + let elapsed = started.elapsed(); + + assert!(elapsed < Duration::from_secs(1), "took {elapsed:?}"); + let json: serde_json::Value = serde_json::from_str(&result).unwrap(); + assert_eq!(json["tool_call_error"], "Tool call 'bash' aborted"); + } + #[test] #[serial] fn tool_timeout_secs_prefers_env_then_config_then_default() { diff --git a/src/graph/executor.rs b/src/graph/executor.rs index a5b046b..2ddc2a3 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -580,7 +580,9 @@ pub(super) async fn step( Ok(StepResult::Continue(vec![next])) } NodeType::Llm(llm_node) => { - let outcome = LlmNodeExecutor::execute(current, llm_node, state, ctx).await?; + let outcome = + LlmNodeExecutor::execute(current, llm_node, state, ctx, step_ctx.abort_signal) + .await?; let targets = match outcome { LlmExecutionOutcome::Continue => { static_next_targets(node, current, "llm", step_ctx.branch_mode)? diff --git a/src/graph/llm.rs b/src/graph/llm.rs index bf00441..414442f 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -11,7 +11,7 @@ use crate::config::{ use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::function::jobs::reap_jobs; use crate::function::skill::skill_function_declarations; -use crate::utils::create_abort_signal; +use crate::utils::AbortSignal; use anyhow::{Context, Error, Result, anyhow, bail}; use log::warn; use serde_json::Value; @@ -33,8 +33,9 @@ impl LlmNodeExecutor { node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, + abort: &AbortSignal, ) -> Result { - let result = run(node_id, node, state_manager, parent_ctx).await; + let result = run(node_id, node, state_manager, parent_ctx, abort).await; let (output, failure_reason) = match result { Ok(raw) => match &node.output_schema { Some(schema) => match structured::extract(&raw, schema, parent_ctx).await { @@ -83,6 +84,7 @@ async fn run( node: &LlmNode, state_manager: &mut StateManager, parent_ctx: &mut RequestContext, + abort: &AbortSignal, ) -> Result { let mut instructions: Option = match &node.instructions { Some(s) => Some( @@ -186,11 +188,11 @@ async fn run( ); parent_ctx.refresh_mcp_tool_filters(); let result = match node.timeout.and_then(wall_clock) { - Some(d) => match timeout(d, run_with_retries(node, &prompt, parent_ctx)).await { + Some(d) => match timeout(d, run_with_retries(node, &prompt, parent_ctx, abort)).await { Ok(r) => r, Err(_) => Err(anyhow!("llm node timed out after {}s", d.as_secs())), }, - None => run_with_retries(node, &prompt, parent_ctx).await, + None => run_with_retries(node, &prompt, parent_ctx, abort).await, }; parent_ctx.role = saved_role; let node_jobs = @@ -242,10 +244,11 @@ async fn run_with_retries( node: &LlmNode, prompt: &str, ctx: &mut RequestContext, + abort: &AbortSignal, ) -> Result { let mut last_err: Option = None; for attempt in 1..=node.max_attempts { - match run_chat_loop(node, prompt, ctx).await { + match run_chat_loop(node, prompt, ctx, abort).await { Ok(out) => return Ok(out), Err(e) if is_transient(&e) && attempt < node.max_attempts => { warn!("llm node attempt {attempt} failed (transient): {e}; retrying"); @@ -262,8 +265,13 @@ pub(crate) fn is_last_turn(turn: u32, max_iterations: u32) -> bool { max_iterations > 0 && turn + 1 == max_iterations } -async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) -> Result { - let abort = create_abort_signal(); +async fn run_chat_loop( + node: &LlmNode, + prompt: &str, + ctx: &mut RequestContext, + abort: &AbortSignal, +) -> Result { + let abort = abort.clone(); let app_cfg = Arc::clone(&ctx.app.config); let role_for_input = ctx.role.clone(); let mut input = Input::from_str(ctx, prompt, role_for_input)?; @@ -271,6 +279,9 @@ async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) - let mut turn: u32 = 0; loop { + if abort.aborted() { + bail!("llm node aborted"); + } let client = input.create_client()?; ctx.before_chat_completion(&input)?; let (output, tool_results) = @@ -474,6 +485,8 @@ mod tests { use super::super::state_updates::OUTPUT_KEY; use super::super::types::*; use super::*; + use crate::config::{Agent, AgentConfig, AppState, WorkingMode}; + use crate::utils::create_abort_signal; use serde_json::json; use std::collections::HashMap; use std::time::Duration; @@ -754,4 +767,34 @@ mod tests { assert!(!is_last_turn(9, 0)); assert!(!is_last_turn(u32::MAX, 0)); } + + /// The turn-top abort check runs before any client is created, so an + /// already-aborted graph never issues a model call and the node fails + /// with the abort reason rather than a connection error. + #[tokio::test] + async fn run_chat_loop_bails_before_first_call_when_aborted() { + let node = node_with(None); + let mut state = manager_with(&[]); + let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); + ctx.agent = Some(Agent::test_new(AgentConfig::default())); + let abort = create_abort_signal(); + abort.set_ctrlc(); + + let err = LlmNodeExecutor::execute("think", &node, &mut state, &mut ctx, &abort) + .await + .expect_err("a pre-set abort must fail the node"); + + let chain = format!("{err:#}"); + assert!( + chain.contains( + "LLM node failed and no fallback declared: LLM call failed: llm node aborted" + ), + "{chain}" + ); + assert_eq!( + state.state().get(OUTPUT_KEY), + None, + "no output is recorded for an aborted node without state_updates" + ); + } }