diff --git a/src/client/common.rs b/src/client/common.rs index 3eb1892..e652499 100644 --- a/src/client/common.rs +++ b/src/client/common.rs @@ -523,14 +523,7 @@ pub async fn call_chat_completions( ctx.app.config.print_markdown(&text)?; } } - let mut tool_results = eval_tool_calls(ctx, tool_calls).await?; - if let Some(first) = tool_results.first_mut() { - first.thinking = thinking; - } - tool_results - .iter() - .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); - Ok((text, tool_results)) + finish_completion(ctx, text, tool_calls, thinking).await } Err(err) => Err(err), } @@ -592,14 +585,7 @@ pub async fn call_chat_completions_streaming( if !silent && !text.is_empty() && !text.ends_with('\n') { println!(); } - let mut tool_results = eval_tool_calls(ctx, tool_calls).await?; - if let Some(first) = tool_results.first_mut() { - first.thinking = thinking; - } - tool_results - .iter() - .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); - Ok((text, tool_results)) + finish_completion(ctx, text, tool_calls, thinking).await } Err(err) => { if !silent && !text.is_empty() { @@ -610,6 +596,72 @@ pub async fn call_chat_completions_streaming( } } +/// Streaming transport for graph `llm` nodes: accumulates the reply and the +/// transport itself renders nothing — no spinner, delta, or reasoning +/// (tool-call rendering inside `eval_tool_calls` is unchanged from the +/// non-streaming path). Because chunks arrive as they are generated, the +/// reqwest `read_timeout` bounds only stalls between chunks; the caller's +/// own deadline bounds the whole generation. Parity with the non-streaming +/// transport is kept deliberately: the handler's call-loop detection is off +/// (that path never had it), and a model the catalog marks `no_stream` +/// takes the non-streaming request, raced against the abort. An abort +/// observed after the provider call fails outright — partial output is +/// never returned as success. +pub async fn call_chat_completions_streaming_quiet( + input: &Input, + client: &dyn Client, + ctx: &mut RequestContext, + abort_signal: AbortSignal, +) -> Result<(String, Vec)> { + if client.model().no_stream() { + let output = tokio::select! { + ret = client.chat_completions(input.clone()) => ret?, + _ = wait_abort_signal(&abort_signal) => bail!("Aborted."), + }; + if abort_signal.aborted() { + bail!("Aborted."); + } + let ChatCompletionsOutput { + text, + tool_calls, + thinking, + } = output; + return finish_completion(ctx, text, tool_calls, thinking).await; + } + + // `_rx` outlives the call so `handler.done()` still finds a receiver. + let (tx, _rx) = unbounded_channel(); + let mut handler = SseHandler::new(tx, abort_signal.clone()); + handler.set_silent(true); + handler.set_call_loop_detection(false); + + let send_ret = client.chat_completions_streaming(input, &mut handler).await; + + if abort_signal.aborted() { + bail!("Aborted."); + } + send_ret?; + + let (text, tool_calls, thinking) = handler.take(); + finish_completion(ctx, text, tool_calls, thinking).await +} + +async fn finish_completion( + ctx: &mut RequestContext, + text: String, + tool_calls: Vec, + thinking: Vec, +) -> Result<(String, Vec)> { + let mut tool_results = eval_tool_calls(ctx, tool_calls).await?; + if let Some(first) = tool_results.first_mut() { + first.thinking = thinking; + } + tool_results + .iter() + .for_each(|res| ctx.tool_scope.tool_tracker.record_call(res.call.clone())); + Ok((text, tool_results)) +} + pub fn noop_prepare_rerank(_client: &T, _data: &RerankData) -> Result { bail!("The client doesn't support rerank api") } @@ -848,6 +900,8 @@ mod tests { use super::*; use super::super::access_token::{is_rejected, set_access_token}; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; fn catch_error_message(data: &Value, status: u16) -> String { catch_error(data, status).unwrap_err().to_string() @@ -1010,4 +1064,286 @@ mod tests { assert!(!should_retry_auth(&api_status_error(401), client)); assert!(!is_rejected(client, "at-1")); } + + /// Fake provider whose streaming path emits `PACED_DELTAS` text chunks + /// 60s apart, then one thinking block and one tool call, and whose + /// non-streaming path returns that same reply in one piece when + /// `nonstream_replies` is set (and bails otherwise). Exercises the real + /// trait defaults (build_client, prepare_completion_data, abort + /// `select!`) with no network. + struct PacedStreamClient { + config: AppConfig, + model: Model, + stall_forever: bool, + nonstream_replies: bool, + streaming_calls: AtomicUsize, + } + + const PACED_DELTAS: usize = 10; + const PACED_GAP: Duration = Duration::from_secs(60); + + fn paced_text() -> String { + (0..PACED_DELTAS).map(|i| format!("chunk{i} ")).collect() + } + + fn paced_tool_call() -> ToolCall { + ToolCall::new( + "lookup".to_string(), + json!({"q": 1}), + Some("call-1".to_string()), + ) + } + + fn paced_thinking() -> ThinkingBlock { + ThinkingBlock::Thinking { + thinking: "hmm".to_string(), + signature: "sig".to_string(), + } + } + + #[async_trait::async_trait] + impl Client for PacedStreamClient { + fn app_config(&self) -> &AppConfig { + &self.config + } + + fn extra_config(&self) -> Option<&ExtraConfig> { + None + } + + fn patch_config(&self) -> Option<&RequestPatch> { + None + } + + fn name(&self) -> &str { + "paced-stream" + } + + fn model(&self) -> &Model { + &self.model + } + + async fn chat_completions_inner( + &self, + _client: &ReqwestClient, + data: ChatCompletionsData, + ) -> Result { + assert!(!data.stream, "non-streaming path must not request a stream"); + if !self.nonstream_replies { + bail!("the quiet transport must not take the non-streaming path"); + } + if self.stall_forever { + std::future::pending::<()>().await; + } + Ok(ChatCompletionsOutput { + text: paced_text(), + tool_calls: vec![paced_tool_call()], + thinking: vec![paced_thinking()], + }) + } + + async fn chat_completions_streaming_inner( + &self, + _client: &ReqwestClient, + handler: &mut SseHandler, + data: ChatCompletionsData, + ) -> Result<()> { + self.streaming_calls.fetch_add(1, Ordering::SeqCst); + assert!(data.stream, "trait default must request a stream"); + for i in 0..PACED_DELTAS { + tokio::time::sleep(PACED_GAP).await; + handler.text(&format!("chunk{i} "))?; + } + if self.stall_forever { + std::future::pending::<()>().await; + } + handler.thinking_block(paced_thinking()); + handler.tool_call(paced_tool_call()) + } + } + + fn paced_client(stall_forever: bool) -> PacedStreamClient { + PacedStreamClient { + config: AppConfig::default(), + model: Model::default(), + stall_forever, + nonstream_replies: false, + streaming_calls: AtomicUsize::new(0), + } + } + + fn no_stream_model() -> Model { + let mut data = ModelData::new(""); + data.no_stream = true; + Model::from_config("", &[data]).remove(0) + } + + fn assert_paced_reply(text: &str, results: &[ToolResult]) { + assert_eq!(text, paced_text()); + assert_eq!(results.len(), 1); + assert_eq!(results[0].call.name, "lookup"); + assert_eq!(results[0].call.id.as_deref(), Some("call-1")); + assert!( + matches!(&results[0].thinking[..], [ThinkingBlock::Thinking { thinking, .. }] if thinking == "hmm") + ); + } + + fn quiet_ctx() -> RequestContext { + RequestContext::new( + Arc::new(crate::config::AppState::test_default()), + crate::config::WorkingMode::Cmd, + ) + } + + /// 600s of virtual generation completes and yields the same shape the + /// non-streaming transport would: full text, tool results with the + /// thinking stapled onto the first, and the tracker fed each call. The + /// fake ignores the reqwest client, so reqwest's `read_timeout` is not in + /// the loop here; the real-socket tests in graph/mod.rs pin that it + /// fires on a stall (`sse_read_timeout_stall_is_transient`) and not on a + /// slow but continuous stream (`sse_slow_but_continuous_stream_outlives_read_timeout`). + #[tokio::test(start_paused = true)] + async fn quiet_streaming_accumulates_across_a_long_generation() { + let mut ctx = quiet_ctx(); + let input = Input::from_str(&ctx, "hi", None).unwrap(); + let client = paced_client(false); + let started = tokio::time::Instant::now(); + + let (text, results) = + call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal()) + .await + .unwrap(); + + assert!(started.elapsed() >= PACED_GAP * PACED_DELTAS as u32); + assert_paced_reply(&text, &results); + assert!( + results[0].output["tool_call_error"].is_string(), + "an undeclared tool evaluates to a tool_call_error, not a panic: {}", + results[0].output + ); + } + + /// An abort mid-stream cancels the provider call promptly and fails the + /// call outright; the ten chunks already buffered are never returned. + #[tokio::test(start_paused = true)] + async fn quiet_streaming_abort_discards_partial_output() { + let mut ctx = quiet_ctx(); + let input = Input::from_str(&ctx, "hi", None).unwrap(); + let client = paced_client(true); + let abort = create_abort_signal(); + let trigger = abort.clone(); + tokio::spawn(async move { + tokio::time::sleep(PACED_GAP * PACED_DELTAS as u32 + Duration::from_secs(30)).await; + trigger.set_ctrlc(); + }); + let started = tokio::time::Instant::now(); + + let err = call_chat_completions_streaming_quiet(&input, &client, &mut ctx, abort) + .await + .expect_err("an aborted stream must not succeed"); + + assert_eq!(err.to_string(), "Aborted."); + assert!(started.elapsed() < PACED_GAP * (PACED_DELTAS as u32 + 1)); + } + + /// The quiet transport and the non-streaming transport hand the node the + /// same `(text, tool_results)` for the same provider reply. + #[tokio::test(start_paused = true)] + async fn quiet_streaming_matches_the_non_streaming_reply() { + let mut ctx = quiet_ctx(); + ctx.render_mode = RenderMode::Silent; + let input = Input::from_str(&ctx, "hi", None).unwrap(); + let mut client = paced_client(false); + client.nonstream_replies = true; + + let (streamed_text, streamed) = + call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal()) + .await + .unwrap(); + let (plain_text, plain) = call_chat_completions( + &input, + false, + false, + &client, + &mut ctx, + create_abort_signal(), + ) + .await + .unwrap(); + + assert_eq!(streamed_text, plain_text); + assert_paced_reply(&streamed_text, &streamed); + assert_paced_reply(&plain_text, &plain); + assert_eq!(streamed[0].output, plain[0].output); + } + + /// A model the catalog marks `no_stream` never sees a streaming request; + /// the quiet transport takes the non-streaming path and returns the same + /// shape. + #[tokio::test] + async fn quiet_transport_honours_no_stream_models() { + let mut ctx = quiet_ctx(); + let input = Input::from_str(&ctx, "hi", None).unwrap(); + let mut client = paced_client(false); + client.model = no_stream_model(); + client.nonstream_replies = true; + + let (text, results) = + call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal()) + .await + .unwrap(); + + assert_eq!(client.streaming_calls.load(Ordering::SeqCst), 0); + assert_paced_reply(&text, &results); + } + + /// The `no_stream` fallback is raced against the abort like the + /// streaming path: a hung non-streaming request is cancelled promptly. + #[tokio::test(start_paused = true)] + async fn quiet_transport_no_stream_fallback_honours_abort() { + let mut ctx = quiet_ctx(); + let input = Input::from_str(&ctx, "hi", None).unwrap(); + let mut client = paced_client(true); + client.model = no_stream_model(); + client.nonstream_replies = true; + let abort = create_abort_signal(); + let trigger = abort.clone(); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_secs(5)).await; + trigger.set_ctrlc(); + }); + let started = tokio::time::Instant::now(); + + let err = call_chat_completions_streaming_quiet(&input, &client, &mut ctx, abort) + .await + .expect_err("an aborted non-streaming fallback must not succeed"); + + assert_eq!(err.to_string(), "Aborted."); + assert!(started.elapsed() < Duration::from_secs(10)); + } + + /// The user-level `stream` toggle is a rendering preference; the quiet + /// transport still streams with it off. + #[tokio::test(start_paused = true)] + async fn quiet_transport_ignores_the_user_stream_toggle() { + let app = crate::config::AppState { + config: Arc::new(AppConfig { + stream: false, + ..AppConfig::default() + }), + ..crate::config::AppState::test_default() + }; + let mut ctx = RequestContext::new(Arc::new(app), crate::config::WorkingMode::Cmd); + let input = Input::from_str(&ctx, "hi", None).unwrap(); + assert!(!input.stream(), "fixture must have streaming disabled"); + let client = paced_client(false); + + let (text, results) = + call_chat_completions_streaming_quiet(&input, &client, &mut ctx, create_abort_signal()) + .await + .unwrap(); + + assert_eq!(client.streaming_calls.load(Ordering::SeqCst), 1); + assert_paced_reply(&text, &results); + } } diff --git a/src/client/stream.rs b/src/client/stream.rs index 79a2d87..8d37aa4 100644 --- a/src/client/stream.rs +++ b/src/client/stream.rs @@ -2,7 +2,7 @@ use super::{ApiStatusError, ThinkingBlock, ToolCall, catch_error}; use crate::utils::AbortSignal; use anyhow::{Context, Result, anyhow, bail}; -use eventsource_stream::Eventsource; +use eventsource_stream::{EventStreamError, Eventsource}; use futures_util::{Stream, StreamExt}; use reqwest::{RequestBuilder, header}; use serde_json::Value; @@ -18,6 +18,7 @@ pub struct SseHandler { max_call_repeats: usize, call_repeat_chain_len: usize, silent: bool, + call_loop_detection: bool, } impl SseHandler { @@ -32,6 +33,7 @@ impl SseHandler { max_call_repeats: 2, call_repeat_chain_len: 3, silent: false, + call_loop_detection: true, } } @@ -39,6 +41,13 @@ impl SseHandler { self.silent = silent; } + /// Loop detection guards an interactive stream; the non-streaming + /// transport never had it, so callers that must match that transport's + /// tool-call handling turn it off. + pub fn set_call_loop_detection(&mut self, enabled: bool) { + self.call_loop_detection = enabled; + } + pub fn text(&mut self, text: &str) -> Result<()> { if text.is_empty() { return Ok(()); @@ -73,7 +82,7 @@ impl SseHandler { } pub fn tool_call(&mut self, call: ToolCall) -> Result<()> { - if self.is_call_loop(&call) { + if self.call_loop_detection && self.is_call_loop(&call) { let loop_message = self.create_loop_detection_message(&call); return Err(anyhow!(loop_message)); } @@ -273,6 +282,11 @@ where break; } } + // Keep the typed reqwest error in the chain: a read_timeout stall + // is only recognised as transient via `reqwest::Error::is_timeout`. + Err(EventStreamError::Transport(err)) => { + return Err(err).context("Transport error"); + } Err(err) => { bail!("{err}"); } @@ -285,13 +299,12 @@ pub async fn json_stream(mut stream: S, mut handle: F) -> Result<()> where S: Stream> + Unpin, F: FnMut(&str) -> Result<()>, - E: std::error::Error, + E: std::error::Error + Send + Sync + 'static, { let mut parser = JsonStreamParser::default(); let mut unparsed_bytes = vec![]; while let Some(chunk_bytes) = stream.next().await { - let chunk_bytes = - chunk_bytes.map_err(|err| anyhow!("Failed to read json stream, {err}"))?; + let chunk_bytes = chunk_bytes.context("Failed to read json stream")?; unparsed_bytes.extend(chunk_bytes); match std::str::from_utf8(&unparsed_bytes) { Ok(text) => { @@ -430,6 +443,40 @@ mod tests { assert!(error_message.contains("test_function_loop")); } + /// With detection off, repeated identical calls all accumulate — the + /// non-streaming transport's behaviour, which the quiet transport mirrors. + #[test] + fn test_call_loop_detection_can_be_disabled() { + let (mut handler, _rx) = new_handler(); + handler.set_call_loop_detection(false); + let call = ToolCall::new("test_function_loop".to_string(), json!({"param": 1}), None); + + for _ in 0..5 { + handler.tool_call(call.clone()).unwrap(); + } + + let (_, calls, _) = handler.take(); + assert_eq!(calls.len(), 5); + } + + /// A transport failure keeps its typed error in the chain so callers can + /// classify it (a stall is only recognised as transient via the type). + #[tokio::test] + async fn json_stream_keeps_the_typed_transport_error() { + let stalled = std::io::Error::new(std::io::ErrorKind::TimedOut, "stall"); + let mut source = stream::iter(vec![Err::(stalled)]); + + let err = json_stream(&mut source, |_| Ok(())) + .await + .expect_err("a failed read must fail the stream"); + + assert!(err.chain().any(|c| { + c.downcast_ref::() + .is_some_and(|e| e.kind() == std::io::ErrorKind::TimedOut) + })); + assert!(format!("{err:#}").starts_with("Failed to read json stream")); + } + fn new_handler() -> (SseHandler, tokio::sync::mpsc::UnboundedReceiver) { let (sender, receiver) = tokio::sync::mpsc::unbounded_channel(); let abort_signal = crate::utils::create_abort_signal(); @@ -467,6 +514,43 @@ mod tests { assert!(handler.has_received_content()); } + /// A silent handler is the quiet transport's whole output gate: text, + /// tool calls, and thinking still accumulate for `take()`, but the + /// receiver sees no `Text` events, only the trailing `Done`. + #[test] + fn test_silent_handler_accumulates_without_emitting_text() { + let (mut handler, mut rx) = new_handler(); + handler.set_silent(true); + + handler.text("hel").unwrap(); + handler.text("lo").unwrap(); + handler.thinking_block(ThinkingBlock::Thinking { + thinking: "hmm".to_string(), + signature: "sig".to_string(), + }); + handler + .tool_call(ToolCall::new( + "lookup".to_string(), + json!({"q": 1}), + Some("call-1".to_string()), + )) + .unwrap(); + handler.done(); + + assert!(matches!(rx.try_recv(), Ok(SseEvent::Done))); + assert!(rx.try_recv().is_err(), "no events beyond Done"); + + let (text, calls, thinking) = handler.take(); + assert_eq!(text, "hello"); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].name, "lookup"); + assert_eq!(calls[0].id.as_deref(), Some("call-1")); + assert_eq!(thinking.len(), 1); + assert!( + matches!(&thinking[0], ThinkingBlock::Thinking { thinking, .. } if thinking == "hmm") + ); + } + fn split_chunks(text: &str) -> Vec> { let len = text.len(); let cut1 = random_range(1..len - 1); diff --git a/src/graph/agent.rs b/src/graph/agent.rs index 2f9468f..7faf413 100644 --- a/src/graph/agent.rs +++ b/src/graph/agent.rs @@ -7,6 +7,7 @@ use super::wall_clock; use crate::config::RequestContext; use crate::function::agents::run_agent_for_graph; use crate::supervisor::mailbox::{Inbox, PeerRegistry}; +use crate::utils::create_abort_signal; use anyhow::{Context, Error, Result, anyhow}; use log::{debug, warn}; use serde_json::Value; @@ -145,8 +146,14 @@ async fn attempt_and_extract( ) .await?; + // Agent nodes carry no graph abort; the session signal is the interrupt + // the executor bridges onto the graph abort anyway. + let abort = parent_ctx + .session_abort + .clone() + .unwrap_or_else(create_abort_signal); let output_value = match &node.output_schema { - Some(schema) => structured::extract(&raw, schema, parent_ctx) + Some(schema) => structured::extract(&raw, schema, parent_ctx, &abort) .await .with_context(|| { format!( @@ -733,6 +740,38 @@ mod tests { ); } + /// Agent-node extraction runs on the session abort: a fired signal + /// stops the extractor before it issues a request. + #[tokio::test] + async fn extraction_honours_a_fired_session_abort() { + let mut ctx = plain_ctx(); + let session = create_abort_signal(); + session.set_ctrlc(); + ctx.session_abort = Some(session); + let mut node = retryable_node(1); + node.output_schema = Some(json!({"type": "object"})); + let mut state = manager_with(&[]); + + let err = attempt_and_extract( + "test_node", + &node, + &mut state, + &mut ctx, + false, + &mut attempt_runner(|_ctx| boxed_attempt(async move { Ok("not json".to_string()) })), + ) + .await + .expect_err("an aborted session must stop the extractor"); + + let chain = format!("{err:#}"); + assert!( + chain.contains("output failed structured-output extraction"), + "{chain}" + ); + assert!(chain.contains("Aborted."), "{chain}"); + assert!(!chain.contains("Invalid model"), "{chain}"); + } + fn failure_capture_updates() -> HashMap { let mut u = HashMap::new(); u.insert("failure".into(), "{{output}}".into()); diff --git a/src/graph/llm.rs b/src/graph/llm.rs index 96d2f6b..614df7c 100644 --- a/src/graph/llm.rs +++ b/src/graph/llm.rs @@ -3,11 +3,12 @@ use super::state_updates; use super::structured; use super::types::LlmNode; use super::{is_transient_error, wall_clock}; -use crate::client::{Model, ModelType, call_chat_completions}; +use crate::client::{Model, ModelType, call_chat_completions_streaming_quiet}; use crate::config::prompts::DEFAULT_SKILL_INSTRUCTIONS; use crate::config::{ Input, RequestContext, Role, RoleLike, SkillPolicy, should_inject_skill_instructions, }; +use crate::function::ToolResult; use crate::function::agents::{GuardrailAction, check_pending_tasks_guardrail}; use crate::function::jobs::reap_jobs; use crate::function::skill::skill_function_declarations; @@ -16,6 +17,8 @@ use anyhow::{Context, Error, Result, anyhow, bail}; use log::warn; use serde_json::Value; use std::collections::HashSet; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use tokio::time::timeout; @@ -36,7 +39,7 @@ impl LlmNodeExecutor { abort: &AbortSignal, ) -> Result { let result = run(node_id, node, state_manager, parent_ctx, abort).await; - finish(node, state_manager, parent_ctx, result).await + finish(node, state_manager, parent_ctx, result, abort).await } } @@ -47,10 +50,11 @@ async fn finish( state_manager: &mut StateManager, parent_ctx: &mut RequestContext, result: Result, + abort: &AbortSignal, ) -> Result { let (output, failure_reason) = match result { 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, abort).await { Ok(value) => (value, None), Err(e) => { warn!("llm node structured extraction failed: {e:#}"); @@ -198,13 +202,14 @@ 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.and_then(wall_clock) { - 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, abort).await, - }; + let result = bounded_run( + node, + &prompt, + parent_ctx, + abort, + &mut default_completion_runner(), + ) + .await; parent_ctx.role = saved_role; let node_jobs = std::mem::replace(&mut parent_ctx.node_job_scope, saved_job_scope).unwrap_or_default(); @@ -251,15 +256,60 @@ fn restore_agent_skill_policy(ctx: &mut RequestContext, saved: Option = + Pin)>> + Send + 'a>>; + +/// The per-call runner the chat loop drives; production uses the quiet +/// streaming transport, tests substitute paced or aborting fakes. The runner +/// owns client construction and the before-hook so a failed `create_client` +/// never records a live `last_message`. +type CompletionRunner<'f> = dyn for<'a> FnMut(&'a Input, &'a mut RequestContext, AbortSignal) -> CompletionFuture<'a> + + Send + + 'f; + +fn default_completion_runner() +-> impl for<'a> FnMut(&'a Input, &'a mut RequestContext, AbortSignal) -> CompletionFuture<'a> + Send +{ + |input, ctx, abort| { + Box::pin(async move { + let client = input.create_client()?; + ctx.before_chat_completion(input)?; + call_chat_completions_streaming_quiet(input, client.as_ref(), ctx, abort).await + }) + } +} + +/// Applies the node's wall-clock bound around the whole retry loop. +async fn bounded_run( + node: &LlmNode, + prompt: &str, + ctx: &mut RequestContext, + abort: &AbortSignal, + runner: &mut CompletionRunner<'_>, +) -> Result { + match node.timeout.and_then(wall_clock) { + Some(d) => match timeout(d, run_with_retries(node, prompt, ctx, abort, runner)).await { + Ok(r) => r, + Err(_) => Err(anyhow!("llm node timed out after {}s", d.as_secs())), + }, + None => run_with_retries(node, prompt, ctx, abort, runner).await, + } +} + async fn run_with_retries( node: &LlmNode, prompt: &str, ctx: &mut RequestContext, abort: &AbortSignal, + runner: &mut CompletionRunner<'_>, ) -> Result { let mut last_err: Option = None; for attempt in 1..=node.max_attempts { - match run_chat_loop(node, prompt, ctx, abort).await { + match run_chat_loop(node, prompt, ctx, abort, runner).await { Ok(out) => return Ok(out), Err(e) if is_transient_error(&e) && attempt < node.max_attempts => { warn!("llm node attempt {attempt} failed (transient): {e:#}; retrying"); @@ -281,6 +331,7 @@ async fn run_chat_loop( prompt: &str, ctx: &mut RequestContext, abort: &AbortSignal, + runner: &mut CompletionRunner<'_>, ) -> Result { let abort = abort.clone(); let app_cfg = Arc::clone(&ctx.app.config); @@ -293,11 +344,12 @@ async fn run_chat_loop( if abort.aborted() { bail!("llm node aborted"); } - let client = input.create_client()?; - ctx.before_chat_completion(&input)?; - let (output, tool_results) = - call_chat_completions(&input, false, false, client.as_ref(), ctx, abort.clone()) - .await?; + let (output, tool_results) = runner(&input, ctx, abort.clone()).await?; + // Defence in depth: the transport bails on abort itself; this closes + // the window between its check and ours. + if abort.aborted() { + bail!("llm node aborted"); + } ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?; if !output.is_empty() { @@ -755,6 +807,165 @@ mod tests { assert!(!is_last_turn(u32::MAX, u32::MAX)); } + fn plain_ctx() -> RequestContext { + RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd) + } + + /// Identity funnel that pins a closure to the runner's higher-ranked + /// signature, so inline closures at call sites infer the right lifetimes. + fn completion_runner(f: F) -> F + where + F: for<'a> FnMut(&'a Input, &'a mut RequestContext, AbortSignal) -> CompletionFuture<'a> + + Send, + { + f + } + + fn boxed_completion<'a>( + fut: impl Future)>> + Send + 'a, + ) -> CompletionFuture<'a> { + Box::pin(fut) + } + + fn paced_reply(secs: u64) -> CompletionFuture<'static> { + boxed_completion(async move { + tokio::time::sleep(Duration::from_secs(secs)).await; + Ok(("done".to_string(), vec![])) + }) + } + + /// A 600s generation finishes under a 900s node timeout: the node's + /// own deadline is what bounds a reply. The runner is a fake, so this + /// does not exercise reqwest's `read_timeout`; the real-socket tests in + /// graph/mod.rs pin that it fires on a stall + /// (`sse_read_timeout_stall_is_transient`) and not on a slow but + /// continuous stream (`sse_slow_but_continuous_stream_outlives_read_timeout`). + #[tokio::test(start_paused = true)] + async fn node_timeout_above_generation_time_lets_the_reply_complete() { + let mut node = node_with(None); + node.timeout = Some(900); + let mut ctx = plain_ctx(); + let abort = create_abort_signal(); + + let out = bounded_run( + &node, + "user", + &mut ctx, + &abort, + &mut completion_runner(|_input, _ctx, _abort| paced_reply(600)), + ) + .await + .unwrap(); + + assert_eq!(out, "done"); + } + + #[tokio::test(start_paused = true)] + async fn node_timeout_below_generation_time_fails_transiently() { + let mut node = node_with(None); + node.timeout = Some(300); + let mut ctx = plain_ctx(); + let abort = create_abort_signal(); + + let err = bounded_run( + &node, + "user", + &mut ctx, + &abort, + &mut completion_runner(|_input, _ctx, _abort| paced_reply(600)), + ) + .await + .expect_err("a 600s reply must not survive a 300s node timeout"); + + assert_eq!(err.to_string(), "llm node timed out after 300s"); + assert!(is_transient_error(&err)); + } + + /// A typed stall (an `Elapsed` under the provider context, as the + /// transport surfaces it) is retried; the second attempt's reply wins. + #[tokio::test(start_paused = true)] + async fn run_with_retries_retries_a_typed_stall_once() { + let mut node = node_with(None); + node.max_attempts = 2; + let mut ctx = plain_ctx(); + let abort = create_abort_signal(); + let mut attempts = 0u32; + + let out = run_with_retries( + &node, + "user", + &mut ctx, + &abort, + &mut completion_runner(|_input, _ctx, _abort| { + attempts += 1; + if attempts == 1 { + boxed_completion(async { + let elapsed = tokio::time::timeout( + Duration::from_secs(1), + std::future::pending::<()>(), + ) + .await + .expect_err("pending future must time out"); + Err(Error::new(elapsed).context("Failed to call chat-completions api")) + }) + } else { + boxed_completion(async { Ok(("recovered".to_string(), vec![])) }) + } + }), + ) + .await + .unwrap(); + + assert_eq!(out, "recovered"); + assert_eq!(attempts, 2); + } + + /// The production runner builds the client before running the + /// before-hook, so a request that never leaves records no live + /// `last_message`. + #[tokio::test] + async fn default_completion_runner_leaves_no_last_message_when_create_client_fails() { + let mut ctx = plain_ctx(); + let input = Input::from_str(&ctx, "hi", None).unwrap(); + let mut runner = default_completion_runner(); + + let err = runner(&input, &mut ctx, create_abort_signal()) + .await + .expect_err("the test AppState has no model to build a client for"); + + assert!( + format!("{err:#}").contains("Invalid model"), + "expected create_client to fail, got: {err:#}" + ); + assert!(ctx.last_message.is_none()); + } + + /// A call that observes the abort but still returns text must not be + /// treated as a completed turn; the loop discards it and fails. + #[tokio::test] + async fn run_chat_loop_discards_output_from_a_call_aborted_midway() { + let node = node_with(None); + let mut ctx = plain_ctx(); + let abort = create_abort_signal(); + + let err = run_chat_loop( + &node, + "user", + &mut ctx, + &abort, + &mut completion_runner(|_input, _ctx, abort| { + boxed_completion(async move { + abort.set_ctrlc(); + Ok(("partial".to_string(), vec![])) + }) + }), + ) + .await + .expect_err("partial output from an aborted call is not a success"); + + assert_eq!(err.to_string(), "llm node aborted"); + } + /// 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. @@ -877,9 +1088,15 @@ mod tests { let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); ctx.agent = Some(Agent::test_new(AgentConfig::default())); - let outcome = finish(&node, &mut state, &mut ctx, Ok("not json".to_string())) - .await - .expect("a declared fallback turns the failure into a route"); + let outcome = finish( + &node, + &mut state, + &mut ctx, + Ok("not json".to_string()), + &create_abort_signal(), + ) + .await + .expect("a declared fallback turns the failure into a route"); assert_eq!(outcome, LlmExecutionOutcome::FellBack("fb".into())); let captured = state @@ -899,6 +1116,46 @@ mod tests { assert!(captured.contains("Invalid model"), "{captured}"); } + /// The extractor runs on the graph abort: an already-aborted graph + /// never issues the extraction request, and the failure text still + /// carries the structured-extraction prefix the fallback scripts anchor on. + #[tokio::test] + async fn structured_extraction_observes_the_graph_abort() { + let mut u = HashMap::new(); + u.insert("captured".into(), "{{output}}".into()); + let mut node = node_with_schema(Some(u), json!({"type": "object"})); + node.fallback = Some("fb".into()); + 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 outcome = finish( + &node, + &mut state, + &mut ctx, + Ok("not json".to_string()), + &abort, + ) + .await + .expect("a declared fallback turns the failure into a route"); + + assert_eq!(outcome, LlmExecutionOutcome::FellBack("fb".into())); + let captured = state + .state() + .get("captured") + .and_then(Value::as_str) + .expect("captured failure text is a string") + .to_string(); + assert!( + captured.starts_with("LLM node structured-extraction failed: "), + "{captured}" + ); + assert!(captured.contains("Aborted."), "{captured}"); + assert!(!captured.contains("Invalid model"), "{captured}"); + } + /// Without a fallback the extraction fault bails, wrapping the same /// extractor chain the state_updates text carries. #[tokio::test] @@ -911,9 +1168,15 @@ mod tests { let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd); ctx.agent = Some(Agent::test_new(AgentConfig::default())); - let err = finish(&node, &mut state, &mut ctx, Ok("not json".to_string())) - .await - .expect_err("no fallback means the failure bails"); + let err = finish( + &node, + &mut state, + &mut ctx, + Ok("not json".to_string()), + &create_abort_signal(), + ) + .await + .expect_err("no fallback means the failure bails"); let chain = format!("{err:#}"); assert!( diff --git a/src/graph/mod.rs b/src/graph/mod.rs index e14f808..4b0dafc 100644 --- a/src/graph/mod.rs +++ b/src/graph/mod.rs @@ -155,4 +155,89 @@ mod tests { assert!(err.is_connect()); assert!(is_transient_error(&Error::new(err))); } + + /// A `read_timeout` stall on the SSE path keeps its typed reqwest error, + /// so the classifier's typed branch fires rather than depending on the + /// rendered text. The server sends headers then never writes a body. + #[tokio::test] + async fn sse_read_timeout_stall_is_transient() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut sock, _) = listener.accept().await.unwrap(); + let mut buf = [0u8; 4096]; + let _ = sock.read(&mut buf).await; + sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n") + .await + .unwrap(); + std::future::pending::<()>().await; + }); + + let client = reqwest::Client::builder() + .no_proxy() + .read_timeout(Duration::from_millis(50)) + .build() + .unwrap(); + let err = crate::client::sse_stream(client.get(format!("http://{addr}/")), |_| Ok(false)) + .await + .expect_err("a body that never arrives must trip the read timeout"); + server.abort(); + + assert!( + err.chain().any(|c| { + c.downcast_ref::() + .is_some_and(|e| e.is_timeout()) + }), + "typed reqwest timeout missing from chain: {err:#}" + ); + assert!(is_transient_error(&err)); + } + + /// `read_timeout` resets on every read: a stream whose events keep + /// arriving well under the limit completes even though the whole body + /// takes several times longer than the limit. + #[tokio::test] + async fn sse_slow_but_continuous_stream_outlives_read_timeout() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + const EVENTS: usize = 15; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (mut sock, _) = listener.accept().await.unwrap(); + let mut buf = [0u8; 4096]; + let _ = sock.read(&mut buf).await; + sock.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\r\n") + .await + .unwrap(); + for i in 0..EVENTS { + tokio::time::sleep(Duration::from_millis(20)).await; + sock.write_all(format!("data: chunk{i}\n\n").as_bytes()) + .await + .unwrap(); + } + sock.write_all(b"data: [DONE]\n\n").await.unwrap(); + }); + + let client = reqwest::Client::builder() + .no_proxy() + .read_timeout(Duration::from_millis(100)) + .build() + .unwrap(); + let mut received = 0usize; + crate::client::sse_stream(client.get(format!("http://{addr}/")), |message| { + if message.data == "[DONE]" { + return Ok(true); + } + received += 1; + Ok(false) + }) + .await + .expect("events every 20ms must never trip a 100ms read timeout"); + server.abort(); + + assert_eq!(received, EVENTS); + } } diff --git a/src/graph/structured.rs b/src/graph/structured.rs index c1a4f15..a4f87af 100644 --- a/src/graph/structured.rs +++ b/src/graph/structured.rs @@ -1,6 +1,6 @@ -use crate::client::call_chat_completions; +use crate::client::call_chat_completions_streaming_quiet; use crate::config::{Input, RequestContext, Role, RoleLike}; -use crate::utils::create_abort_signal; +use crate::utils::AbortSignal; use anyhow::{Context, Result, bail}; use serde_json::Value; use std::sync::Arc; @@ -19,18 +19,24 @@ Rules: - If a field cannot be determined from the input, use `null` (when allowed) or your best inferred value. - Do NOT invent fields not present in the schema."; -pub async fn extract(raw: &str, schema: &Value, parent_ctx: &mut RequestContext) -> Result { +pub async fn extract( + raw: &str, + schema: &Value, + parent_ctx: &mut RequestContext, + abort: &AbortSignal, +) -> Result { if let Some(parsed) = try_parse_json(raw) { return Ok(parsed); } - extract_via_extractor(raw, schema, parent_ctx, false).await + extract_via_extractor(raw, schema, parent_ctx, abort, false).await } async fn extract_via_extractor( raw: &str, schema: &Value, parent_ctx: &mut RequestContext, + abort: &AbortSignal, is_repair: bool, ) -> Result { let role = build_extractor_role(parent_ctx); @@ -38,7 +44,7 @@ async fn extract_via_extractor( let saved_role = parent_ctx.role.clone(); parent_ctx.role = Some(role); - let result = run_one_shot(&prompt, parent_ctx).await; + let result = run_one_shot(&prompt, parent_ctx, abort).await; parent_ctx.role = saved_role; let output = result.context("Structured-output extractor LLM call failed")?; @@ -49,7 +55,12 @@ async fn extract_via_extractor( "Structured-output extractor failed to produce valid JSON after repair retry. \ Last response:\n{output}" ), - None => Box::pin(extract_via_extractor(&output, schema, parent_ctx, true)).await, + None => { + Box::pin(extract_via_extractor( + &output, schema, parent_ctx, abort, true, + )) + .await + } } } @@ -73,15 +84,21 @@ fn build_extractor_prompt(raw: &str, schema: &Value, is_repair: bool) -> String } } -async fn run_one_shot(prompt: &str, ctx: &mut RequestContext) -> Result { - let abort = create_abort_signal(); +async fn run_one_shot( + prompt: &str, + ctx: &mut RequestContext, + abort: &AbortSignal, +) -> Result { + if abort.aborted() { + bail!("Aborted."); + } let app_cfg = Arc::clone(&ctx.app.config); let role_for_input = ctx.role.clone(); let input = Input::from_str(ctx, prompt, role_for_input)?; let client = input.create_client()?; ctx.before_chat_completion(&input)?; let (output, tool_results) = - call_chat_completions(&input, false, false, client.as_ref(), ctx, abort).await?; + call_chat_completions_streaming_quiet(&input, client.as_ref(), ctx, abort.clone()).await?; ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?; Ok(output)