feat(graph): thread the graph abort into llm nodes and tool execution

LlmNodeExecutor::execute now takes the graph's abort signal and threads
it through run -> run_with_retries -> run_chat_loop, replacing the
loop-private signal, so an in-flight model call is interrupted and every
turn starts with an abort check that bails 'llm node aborted' before any
client is created.

run_llm_function gains an optional abort signal: the synchronous tool
wait loop kills the child and returns a 'Tool call ... aborted' error
within one poll interval when it is set. The tool dispatcher passes the
session abort; the agent _instructions runner passes None.
This commit is contained in:
Dark-Alex-17 committed 2026-09-11 10:56:28 -06:00
1 parent 1ea755e864
commit cce8381962
4 files changed
+105 -8

No files matched your search

+1
View File
@@ -633,6 +633,7 @@ impl Agent {
Some(self.name().to_string()), Some(self.name().to_string()),
tool_timeout, tool_timeout,
false, false,
None,
)?; )?;
match value { match value {
Some(v) => Ok(v), Some(v) => Ok(v),
+51
View File
@@ -1684,6 +1684,7 @@ impl ToolCall {
agent_name, agent_name,
ctx.app.config.tool_timeout, ctx.app.config.tool_timeout,
quiet, quiet,
ctx.session_abort.clone(),
) { ) {
Ok(Some(contents)) => serde_json::from_str(&contents) Ok(Some(contents)) => serde_json::from_str(&contents)
.ok() .ok()
@@ -2293,6 +2294,7 @@ pub fn run_llm_function(
agent_name: Option<String>, agent_name: Option<String>,
tool_timeout: Option<u64>, tool_timeout: Option<u64>,
quiet: bool, quiet: bool,
abort: Option<AbortSignal>,
) -> Result<Option<String>> { ) -> Result<Option<String>> {
let mut bin_dirs: Vec<PathBuf> = vec![]; let mut bin_dirs: Vec<PathBuf> = vec![];
let mut command_name = cmd_name.clone(); let mut command_name = cmd_name.clone();
@@ -2440,6 +2442,21 @@ pub fn run_llm_function(
return Ok(Some(error_json.to_string())); 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)); thread::sleep(Duration::from_millis(100));
}; };
let stdout_bytes = stdout_thread.join().unwrap_or_default(); let stdout_bytes = stdout_thread.join().unwrap_or_default();
@@ -4562,6 +4579,7 @@ mod tests {
None, None,
None, None,
false, false,
None,
) )
.unwrap() .unwrap()
.expect("nonzero exit must return an error payload"); .expect("nonzero exit must return an error payload");
@@ -4590,6 +4608,7 @@ mod tests {
None, None,
None, None,
true, true,
None,
) )
.unwrap() .unwrap()
.expect("nonzero exit must return an error payload"); .expect("nonzero exit must return an error payload");
@@ -4605,6 +4624,38 @@ mod tests {
assert_eq!(json["output"], "partial-output\n"); 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] #[test]
#[serial] #[serial]
fn tool_timeout_secs_prefers_env_then_config_then_default() { fn tool_timeout_secs_prefers_env_then_config_then_default() {
+3 -1
View File
@@ -580,7 +580,9 @@ pub(super) async fn step(
Ok(StepResult::Continue(vec![next])) Ok(StepResult::Continue(vec![next]))
} }
NodeType::Llm(llm_node) => { 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 { let targets = match outcome {
LlmExecutionOutcome::Continue => { LlmExecutionOutcome::Continue => {
static_next_targets(node, current, "llm", step_ctx.branch_mode)? static_next_targets(node, current, "llm", step_ctx.branch_mode)?
+50 -7
View File
@@ -11,7 +11,7 @@ use crate::config::{
use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail};
use crate::function::jobs::reap_jobs; use crate::function::jobs::reap_jobs;
use crate::function::skill::skill_function_declarations; 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 anyhow::{Context, Error, Result, anyhow, bail};
use log::warn; use log::warn;
use serde_json::Value; use serde_json::Value;
@@ -33,8 +33,9 @@ impl LlmNodeExecutor {
node: &LlmNode, node: &LlmNode,
state_manager: &mut StateManager, state_manager: &mut StateManager,
parent_ctx: &mut RequestContext, parent_ctx: &mut RequestContext,
abort: &AbortSignal,
) -> Result<LlmExecutionOutcome> { ) -> Result<LlmExecutionOutcome> {
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 { let (output, failure_reason) = match result {
Ok(raw) => match &node.output_schema { Ok(raw) => match &node.output_schema {
Some(schema) => match structured::extract(&raw, schema, parent_ctx).await { Some(schema) => match structured::extract(&raw, schema, parent_ctx).await {
@@ -83,6 +84,7 @@ async fn run(
node: &LlmNode, node: &LlmNode,
state_manager: &mut StateManager, state_manager: &mut StateManager,
parent_ctx: &mut RequestContext, parent_ctx: &mut RequestContext,
abort: &AbortSignal,
) -> Result<String> { ) -> Result<String> {
let mut instructions: Option<String> = match &node.instructions { let mut instructions: Option<String> = match &node.instructions {
Some(s) => Some( Some(s) => Some(
@@ -186,11 +188,11 @@ async fn run(
); );
parent_ctx.refresh_mcp_tool_filters(); parent_ctx.refresh_mcp_tool_filters();
let result = match node.timeout.and_then(wall_clock) { 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, Ok(r) => r,
Err(_) => Err(anyhow!("llm node timed out after {}s", d.as_secs())), 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; parent_ctx.role = saved_role;
let node_jobs = let node_jobs =
@@ -242,10 +244,11 @@ async fn run_with_retries(
node: &LlmNode, node: &LlmNode,
prompt: &str, prompt: &str,
ctx: &mut RequestContext, ctx: &mut RequestContext,
abort: &AbortSignal,
) -> Result<String> { ) -> Result<String> {
let mut last_err: Option<Error> = None; let mut last_err: Option<Error> = None;
for attempt in 1..=node.max_attempts { 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), Ok(out) => return Ok(out),
Err(e) if is_transient(&e) && attempt < node.max_attempts => { Err(e) if is_transient(&e) && attempt < node.max_attempts => {
warn!("llm node attempt {attempt} failed (transient): {e}; retrying"); 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 max_iterations > 0 && turn + 1 == max_iterations
} }
async fn run_chat_loop(node: &LlmNode, prompt: &str, ctx: &mut RequestContext) -> Result<String> { async fn run_chat_loop(
let abort = create_abort_signal(); node: &LlmNode,
prompt: &str,
ctx: &mut RequestContext,
abort: &AbortSignal,
) -> Result<String> {
let abort = abort.clone();
let app_cfg = Arc::clone(&ctx.app.config); let app_cfg = Arc::clone(&ctx.app.config);
let role_for_input = ctx.role.clone(); let role_for_input = ctx.role.clone();
let mut input = Input::from_str(ctx, prompt, role_for_input)?; 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; let mut turn: u32 = 0;
loop { loop {
if abort.aborted() {
bail!("llm node aborted");
}
let client = input.create_client()?; let client = input.create_client()?;
ctx.before_chat_completion(&input)?; ctx.before_chat_completion(&input)?;
let (output, tool_results) = let (output, tool_results) =
@@ -474,6 +485,8 @@ mod tests {
use super::super::state_updates::OUTPUT_KEY; use super::super::state_updates::OUTPUT_KEY;
use super::super::types::*; use super::super::types::*;
use super::*; use super::*;
use crate::config::{Agent, AgentConfig, AppState, WorkingMode};
use crate::utils::create_abort_signal;
use serde_json::json; use serde_json::json;
use std::collections::HashMap; use std::collections::HashMap;
use std::time::Duration; use std::time::Duration;
@@ -754,4 +767,34 @@ mod tests {
assert!(!is_last_turn(9, 0)); assert!(!is_last_turn(9, 0));
assert!(!is_last_turn(u32::MAX, 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"
);
}
} }