diff --git a/src/graph/agent.rs b/src/graph/agent.rs index e2b9608..c6560dc 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -2,13 +2,13 @@ use super::state::StateManager; use super::state_updates; use super::structured; use super::types::AgentNode; +use super::wall_clock; use crate::config::RequestContext; use crate::function::agents::run_agent_for_graph; use anyhow::{Context, Result}; use log::debug; use serde_json::Value; use std::collections::HashMap; -use std::time::Duration; use tokio::time::timeout; const DEFAULT_TIMEOUT_SECS: u64 = 300; @@ -33,7 +33,8 @@ impl AgentNodeExecutor { debug!("Agent '{}' graph inputs: {keys:?}", node.agent); } - let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS)); + let secs = node.timeout.unwrap_or(DEFAULT_TIMEOUT_SECS); + let bound = wall_clock(secs); // run_agent_for_graph takes the identity off the ctx; keep a handle so // a frontier peer is retired the moment the agent stops, not after @@ -46,22 +47,16 @@ impl AgentNodeExecutor { .map(|(id, _)| id.clone()), ); - let raw_result = timeout( - timeout_dur, - run_agent_for_graph(parent_ctx, &node.agent, &prompt, graph_inputs), - ) - .await; + let fut = run_agent_for_graph(parent_ctx, &node.agent, &prompt, graph_inputs); + let raw_result = match bound { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + }; if retire_peer_on_return && let Some((registry, id)) = &peer { registry.mark_finished(id); } let raw = raw_result - .with_context(|| { - format!( - "Agent '{}' timed out after {}s", - node.agent, - timeout_dur.as_secs() - ) - })? + .with_context(|| format!("Agent '{}' timed out after {}s", node.agent, secs))? .with_context(|| format!("Agent '{}' failed", node.agent))?; let output_value = match &node.output_schema { diff --git a/src/graph/executor.rs b/src/graph/executor.rs index 66a4822..937e778 100644 --- a/src/graph/executor.rs +++ b/src/graph/executor.rs @@ -9,6 +9,7 @@ use super::state::StateManager; use super::types::{EndNode, Graph, Node, NodeType}; use super::user_interaction::{ApprovalNodeExecutor, InputNodeExecutor}; use super::validator::{AgentValidationContext, GraphValidator}; +use super::wall_clock; use crate::config::{RenderMode, RequestContext}; use crate::supervisor::mailbox::{Inbox, PeerAssignment, PeerRegistry, graph_agent_id}; use crate::utils::AbortSignal; @@ -105,7 +106,7 @@ impl GraphExecutor { .unwrap_or_default(); let script_executor = ScriptExecutor::new(&base_dir).with_envs(agent_envs); let max_iterations = graph.settings.max_loop_iterations; - let graph_timeout = graph.settings.timeout.map(Duration::from_secs); + let graph_timeout = graph.settings.timeout.and_then(wall_clock); let max_concurrency = graph.settings.max_concurrency; let graph = Arc::new(graph); let start = Instant::now(); @@ -1221,6 +1222,40 @@ nodes: assert!(err.contains("sleeper"), "error should name frontier: {err}"); } + #[tokio::test] + async fn graph_timeout_zero_is_unbounded() { + if !cmd_available("bash") { + eprintln!("skipping: bash not available"); + return; + } + let ws = TestWorkspace::new(); + ws.write_script("sleeper.sh", "#!/bin/bash\nsleep 1.5\necho '{}'\n"); + + let yaml = r#" +name: zero_timeout_test +start: sleeper +settings: + timeout: 0 +nodes: + sleeper: + type: script + script: sleeper.sh + state_updates: {} + next: done + done: + type: end + output: "done" +"#; + let graph: Graph = serde_yaml::from_str(yaml).unwrap(); + let mut ctx = make_ctx(); + let abort = create_abort_signal(); + let result = GraphExecutor::new(graph, &ws.dir) + .execute(&mut ctx, abort) + .await; + + result.unwrap_or_else(|e| panic!("timeout: 0 should not bound the run: {e:#}")); + } + #[cfg(unix)] #[tokio::test] async fn background_job_survives_graph_node_execution() { diff --git a/src/graph/llm.rs b/src/graph/llm.rs index e07f927..b473db0 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -2,6 +2,7 @@ use super::state::StateManager; use super::state_updates; use super::structured; use super::types::LlmNode; +use super::wall_clock; use crate::client::{Model, ModelType, call_chat_completions}; use crate::config::prompts::DEFAULT_SKILL_INSTRUCTIONS; use crate::config::{ @@ -16,7 +17,6 @@ use log::warn; use serde_json::Value; use std::collections::HashSet; use std::sync::Arc; -use std::time::Duration; use tokio::time::timeout; #[derive(Debug, PartialEq, Eq)] @@ -185,15 +185,10 @@ async fn run( node.mcp_tools.clone().map(|map| (node_id.to_string(), map)), ); parent_ctx.refresh_mcp_tool_filters(); - let result = match node.timeout { - Some(secs) => match timeout( - Duration::from_secs(secs), - run_with_retries(node, &prompt, parent_ctx), - ) - .await - { + let result = match node.timeout.and_then(wall_clock) { + Some(d) => match timeout(d, run_with_retries(node, &prompt, parent_ctx)).await { Ok(r) => r, - Err(_) => Err(anyhow!("llm node timed out after {secs}s")), + Err(_) => Err(anyhow!("llm node timed out after {}s", d.as_secs())), }, None => run_with_retries(node, &prompt, parent_ctx).await, }; @@ -478,6 +473,7 @@ mod tests { use super::*; use serde_json::json; use std::collections::HashMap; + use std::time::Duration; fn manager_with(pairs: &[(&str, Value)]) -> StateManager { let mut map = HashMap::new(); @@ -731,4 +727,13 @@ mod tests { assert!(!is_transient(&anyhow!("hit max_iterations"))); assert!(!is_transient(&anyhow!("authentication failed"))); } + + #[test] + fn zero_timeout_resolves_to_no_bound() { + assert!(Some(0u64).and_then(wall_clock).is_none()); + assert_eq!( + Some(5u64).and_then(wall_clock), + Some(Duration::from_secs(5)) + ); + } } diff --git a/src/graph/map.rs b/src/graph/map.rs index 965b801..a0af2e6 100644 --- a/src/graph/map.rs +++ b/src/graph/map.rs @@ -1709,12 +1709,10 @@ nodes: assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); } - /// `tokio::time::timeout` polls the agent future first and a zero-duration - /// sleep still goes through the time driver, so `Elapsed` only fires if the - /// agent future is still pending when the timer is next polled. The - /// materialized agent's graph holds its script step open for seconds, so - /// the future is dropped mid-flight and the chain runner alone retires - /// the identity. + /// The materialized agent's graph holds its script step open for seconds, + /// so the 1s timer fires while that step is still pending: the agent + /// future is dropped mid-flight and the chain runner alone retires the + /// identity. #[tokio::test] #[serial] async fn run_item_chain_marks_identity_finished_when_agent_step_times_out() { @@ -1724,7 +1722,7 @@ nodes: } let _guard = TestConfigDirGuard::new(); materialize_probe_agent(5.0); - let h = Harness::new(&flagged_probe_graph(Some(0))); + let h = Harness::new(&flagged_probe_graph(Some(1))); let (registry, assignments) = h.provision("worker", 2); let peers = (Arc::clone(®istry), assignments[0].clone()); let mut state = item_state(json!(0)); @@ -1737,7 +1735,7 @@ nodes: "{chain}" ); assert!( - chain.contains(&format!("Agent '{PROBE_AGENT}' timed out after 0s")), + chain.contains(&format!("Agent '{PROBE_AGENT}' timed out after 1s")), "{chain}" ); assert!( @@ -1749,6 +1747,32 @@ nodes: assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); } + /// Positive twin of the timeout test above: with `timeout: 0` the agent + /// step has no wall-clock bound, so a hold longer than any accidental 0s + /// or 1s bound still runs to completion and hands its output back. + #[tokio::test] + #[serial] + async fn run_item_chain_probe_agent_with_zero_timeout_runs_to_completion() { + if !cmd_available("python3") { + eprintln!("skipping: python3 not available"); + return; + } + let _guard = TestConfigDirGuard::new(); + materialize_probe_agent(1.5); + let h = Harness::new(&flagged_probe_graph(Some(0))); + let peers = manual_peer("worker"); + let mut state = item_state(json!(0)); + let mut ctx = silent_ctx(); + + h.run("worker", Some(&peers), &mut state, &mut ctx) + .await + .unwrap_or_else(|e| panic!("chain failed: {e:#}")); + + assert_eq!(state.state().get("output"), Some(&json!("held"))); + assert!(peers.0.is_finished(&peers.1.0)); + assert!(ctx.peer_registry.is_none() && ctx.peer_assignment.is_none()); + } + #[tokio::test] #[serial] async fn run_item_chain_clears_peer_fields_after_flagged_agent_then_script() { diff --git a/src/graph/mod.rs b/src/graph/mod.rs index 0c5830f..3b1d72d 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -22,12 +22,18 @@ pub use dispatch::{ pub use executor::GraphExecutor; pub use parser::{GraphParser, agent_has_graph}; use serde_json::Value; +use std::time::Duration; pub use types::{Graph, NodeType}; pub const GRAPH_SCHEMA_VERSION: &str = "1.0"; pub const DEFAULT_MAX_LOOP_ITERATIONS: usize = 100; +/// A timeout of `0` disables the wall-clock bound. +pub(crate) fn wall_clock(secs: u64) -> Option { + (secs != 0).then(|| Duration::from_secs(secs)) +} + pub const MAX_STATE_SIZE_BYTES: usize = 32 * 1024; pub(in crate::graph) fn type_name(value: &Value) -> &'static str { @@ -40,3 +46,18 @@ pub(in crate::graph) fn type_name(value: &Value) -> &'static str { Value::Object(_) => "object", } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn wall_clock_zero_is_no_bound() { + assert!(wall_clock(0).is_none()); + } + + #[test] + fn wall_clock_nonzero_is_that_many_seconds() { + assert_eq!(wall_clock(7), Some(Duration::from_secs(7))); + } +} diff --git a/src/graph/rag.rs b/src/graph/rag.rs index 07a05f7..fbba2aa 100644 --- a/src/graph/rag.rs +++ b/src/graph/rag.rs @@ -1,11 +1,11 @@ use super::state::StateManager; use super::state_updates; use super::types::RagNode; +use super::wall_clock; use crate::config::RequestContext; use crate::utils::create_abort_signal; use anyhow::{Context, Result, anyhow}; use serde_json::{Map, Value}; -use std::time::Duration; use tokio::time::timeout; const DEFAULT_QUERY: &str = "{{initial_prompt}}"; @@ -34,18 +34,16 @@ impl RagNodeExecutor { let top_k = node.top_k.unwrap_or_else(|| rag.configured_top_k()); let rerank = rag.configured_reranker(); - let timeout_dur = Duration::from_secs(node.timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS)); + let secs = node.timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS); + let bound = wall_clock(secs); let abort = create_abort_signal(); - let (context, sources_str, _ids) = - timeout(timeout_dur, rag.search(&query, top_k, rerank, abort)) - .await - .with_context(|| { - format!( - "rag node '{node_id}' timed out after {}s", - timeout_dur.as_secs() - ) - })? - .with_context(|| format!("rag node '{node_id}' retrieval failed"))?; + let fut = rag.search(&query, top_k, rerank, abort); + let (context, sources_str, _ids) = match bound { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + } + .with_context(|| format!("rag node '{node_id}' timed out after {secs}s"))? + .with_context(|| format!("rag node '{node_id}' retrieval failed"))?; let output = build_rag_output(context, &sources_str); apply_state_updates(node, state_manager, &output); @@ -77,6 +75,7 @@ fn apply_state_updates(node: &RagNode, state_manager: &mut StateManager, output: mod tests { use super::*; use serde_json::json; + use std::time::Duration; #[test] fn build_rag_output_splits_bullet_sources_into_array() { @@ -106,4 +105,15 @@ mod tests { assert_eq!(out["sources"], json!(["plain/path"])); } + + /// Mirrors the resolution expression in `execute`. + fn resolve(timeout: Option) -> Option { + wall_clock(timeout.unwrap_or(DEFAULT_RAG_TIMEOUT_SECS)) + } + + #[test] + fn zero_timeout_resolves_to_no_bound() { + assert!(resolve(Some(0)).is_none()); + assert_eq!(resolve(None), Some(Duration::from_secs(120))); + } } diff --git a/src/graph/script.rs b/src/graph/script.rs index 688d6ad..a2d7363 100644 --- a/src/graph/script.rs +++ b/src/graph/script.rs @@ -1,5 +1,6 @@ use super::state::{StateManager, StateRepresentation}; use super::types::ScriptNode; +use super::wall_clock; use crate::config::paths; use crate::function::Language; use anyhow::{Context, Result, anyhow, bail}; @@ -8,7 +9,6 @@ use std::collections::HashMap; use std::env; use std::path::{Path, PathBuf}; use std::process::Stdio; -use std::time::Duration; use tokio::process::Command; use tokio::time::timeout; @@ -66,22 +66,25 @@ impl ScriptExecutor { } } - let timeout_dur = Duration::from_secs(node.timeout); - let output = timeout(timeout_dur, cmd.output()) - .await - .with_context(|| { - format!( - "Script '{}' timed out after {}s", - script_path.display(), - node.timeout - ) - })? - .with_context(|| { - format!( - "Failed to spawn script process for '{}'", - script_path.display() - ) - })?; + let bound = wall_clock(node.timeout); + let fut = cmd.output(); + let output = match bound { + Some(d) => timeout(d, fut).await, + None => Ok(fut.await), + } + .with_context(|| { + format!( + "Script '{}' timed out after {}s", + script_path.display(), + node.timeout + ) + })? + .with_context(|| { + format!( + "Failed to spawn script process for '{}'", + script_path.display() + ) + })?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); @@ -260,6 +263,33 @@ echo '{"quality": 0.85, "issues": 3, "_next": "approve"}' cleanup(&dir); } + #[tokio::test] + async fn zero_timeout_is_unbounded() { + if !cmd_available("bash") { + return; + } + let (dir, path) = write_script( + r#"#!/bin/bash +sleep 1.5 +echo '{"ok":true}' +"#, + "sh", + ); + let mut state = StateManager::new(HashMap::new()); + let executor = ScriptExecutor::new(&dir); + + executor + .execute( + &node_for(path.file_name().unwrap().to_str().unwrap(), 0), + &mut state, + ) + .await + .unwrap_or_else(|e| panic!("zero timeout should not bound the script: {e:#}")); + + assert_eq!(state.state().get("ok"), Some(&json!(true))); + cleanup(&dir); + } + #[tokio::test] async fn bash_script_can_read_state_from_env() { if !cmd_available("bash") {