feat(graph): stream llm-node completions so node timeout governs long generations
Graph llm nodes (and the structured-output extractor they share with agent nodes) now obtain the model reply through a quiet streaming transport, `call_chat_completions_streaming_quiet`, instead of the non-streaming call. Because chunks arrive as they are generated, reqwest's client-level `read_timeout` (default 300s) bounds only a genuine stall between chunks and the node's own `timeout` becomes the effective bound on total generation. Previously zero bytes arrived until the whole reply finished, so a 300s read_timeout silently capped every llm node regardless of its `timeout` (seen live: code-reviewer's `synthesize` node dying on a 61-file review). The quiet transport renders nothing (silent SseHandler, no spinner), keeps tool-call handling identical to the non-streaming path (call-loop detection off), falls back to the non-streaming request for catalog `no_stream` models (raced against the abort), and never returns partial output as success. Supporting changes: - `sse_stream`/`json_stream` keep the typed `reqwest::Error` in the anyhow chain so `is_transient_error` still classifies read_timeout stalls as transient and `max_attempts` retries them (previously stringified). - `structured::extract` takes the graph/session AbortSignal and refuses to issue an extraction request for an already-aborted graph. - `run_chat_loop` gains an injectable CompletionRunner seam (owning client construction + the before-hook) so the timeout/retry/abort behaviour is unit-testable over a paused clock; a fake `impl Client` drives the real trait-default streaming for 600s of virtual time; real-socket tests pin that read_timeout fires on a stall but not on a slow continuous stream.
This commit is contained in:
6 files changed
+873
-49
No files matched your search
+352
-16
@@ -523,14 +523,7 @@ pub async fn call_chat_completions(
|
|||||||
ctx.app.config.print_markdown(&text)?;
|
ctx.app.config.print_markdown(&text)?;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
let mut tool_results = eval_tool_calls(ctx, tool_calls).await?;
|
finish_completion(ctx, text, tool_calls, thinking).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))
|
|
||||||
}
|
}
|
||||||
Err(err) => Err(err),
|
Err(err) => Err(err),
|
||||||
}
|
}
|
||||||
@@ -592,14 +585,7 @@ pub async fn call_chat_completions_streaming(
|
|||||||
if !silent && !text.is_empty() && !text.ends_with('\n') {
|
if !silent && !text.is_empty() && !text.ends_with('\n') {
|
||||||
println!();
|
println!();
|
||||||
}
|
}
|
||||||
let mut tool_results = eval_tool_calls(ctx, tool_calls).await?;
|
finish_completion(ctx, text, tool_calls, thinking).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))
|
|
||||||
}
|
}
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
if !silent && !text.is_empty() {
|
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<ToolResult>)> {
|
||||||
|
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<ToolCall>,
|
||||||
|
thinking: Vec<ThinkingBlock>,
|
||||||
|
) -> Result<(String, Vec<ToolResult>)> {
|
||||||
|
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<T>(_client: &T, _data: &RerankData) -> Result<RequestData> {
|
pub fn noop_prepare_rerank<T>(_client: &T, _data: &RerankData) -> Result<RequestData> {
|
||||||
bail!("The client doesn't support rerank api")
|
bail!("The client doesn't support rerank api")
|
||||||
}
|
}
|
||||||
@@ -848,6 +900,8 @@ mod tests {
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
use super::super::access_token::{is_rejected, set_access_token};
|
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 {
|
fn catch_error_message(data: &Value, status: u16) -> String {
|
||||||
catch_error(data, status).unwrap_err().to_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!(!should_retry_auth(&api_status_error(401), client));
|
||||||
assert!(!is_rejected(client, "at-1"));
|
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<ChatCompletionsOutput> {
|
||||||
|
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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
+89
-5
@@ -2,7 +2,7 @@ use super::{ApiStatusError, ThinkingBlock, ToolCall, catch_error};
|
|||||||
use crate::utils::AbortSignal;
|
use crate::utils::AbortSignal;
|
||||||
|
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use eventsource_stream::Eventsource;
|
use eventsource_stream::{EventStreamError, Eventsource};
|
||||||
use futures_util::{Stream, StreamExt};
|
use futures_util::{Stream, StreamExt};
|
||||||
use reqwest::{RequestBuilder, header};
|
use reqwest::{RequestBuilder, header};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
@@ -18,6 +18,7 @@ pub struct SseHandler {
|
|||||||
max_call_repeats: usize,
|
max_call_repeats: usize,
|
||||||
call_repeat_chain_len: usize,
|
call_repeat_chain_len: usize,
|
||||||
silent: bool,
|
silent: bool,
|
||||||
|
call_loop_detection: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl SseHandler {
|
impl SseHandler {
|
||||||
@@ -32,6 +33,7 @@ impl SseHandler {
|
|||||||
max_call_repeats: 2,
|
max_call_repeats: 2,
|
||||||
call_repeat_chain_len: 3,
|
call_repeat_chain_len: 3,
|
||||||
silent: false,
|
silent: false,
|
||||||
|
call_loop_detection: true,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -39,6 +41,13 @@ impl SseHandler {
|
|||||||
self.silent = silent;
|
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<()> {
|
pub fn text(&mut self, text: &str) -> Result<()> {
|
||||||
if text.is_empty() {
|
if text.is_empty() {
|
||||||
return Ok(());
|
return Ok(());
|
||||||
@@ -73,7 +82,7 @@ impl SseHandler {
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub fn tool_call(&mut self, call: ToolCall) -> Result<()> {
|
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);
|
let loop_message = self.create_loop_detection_message(&call);
|
||||||
return Err(anyhow!(loop_message));
|
return Err(anyhow!(loop_message));
|
||||||
}
|
}
|
||||||
@@ -273,6 +282,11 @@ where
|
|||||||
break;
|
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) => {
|
Err(err) => {
|
||||||
bail!("{err}");
|
bail!("{err}");
|
||||||
}
|
}
|
||||||
@@ -285,13 +299,12 @@ pub async fn json_stream<S, F, E>(mut stream: S, mut handle: F) -> Result<()>
|
|||||||
where
|
where
|
||||||
S: Stream<Item = Result<bytes::Bytes, E>> + Unpin,
|
S: Stream<Item = Result<bytes::Bytes, E>> + Unpin,
|
||||||
F: FnMut(&str) -> Result<()>,
|
F: FnMut(&str) -> Result<()>,
|
||||||
E: std::error::Error,
|
E: std::error::Error + Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
let mut parser = JsonStreamParser::default();
|
let mut parser = JsonStreamParser::default();
|
||||||
let mut unparsed_bytes = vec![];
|
let mut unparsed_bytes = vec![];
|
||||||
while let Some(chunk_bytes) = stream.next().await {
|
while let Some(chunk_bytes) = stream.next().await {
|
||||||
let chunk_bytes =
|
let chunk_bytes = chunk_bytes.context("Failed to read json stream")?;
|
||||||
chunk_bytes.map_err(|err| anyhow!("Failed to read json stream, {err}"))?;
|
|
||||||
unparsed_bytes.extend(chunk_bytes);
|
unparsed_bytes.extend(chunk_bytes);
|
||||||
match std::str::from_utf8(&unparsed_bytes) {
|
match std::str::from_utf8(&unparsed_bytes) {
|
||||||
Ok(text) => {
|
Ok(text) => {
|
||||||
@@ -430,6 +443,40 @@ mod tests {
|
|||||||
assert!(error_message.contains("test_function_loop"));
|
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::<Bytes, std::io::Error>(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::<std::io::Error>()
|
||||||
|
.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<SseEvent>) {
|
fn new_handler() -> (SseHandler, tokio::sync::mpsc::UnboundedReceiver<SseEvent>) {
|
||||||
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
|
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
|
||||||
let abort_signal = crate::utils::create_abort_signal();
|
let abort_signal = crate::utils::create_abort_signal();
|
||||||
@@ -467,6 +514,43 @@ mod tests {
|
|||||||
assert!(handler.has_received_content());
|
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<Vec<u8>> {
|
fn split_chunks(text: &str) -> Vec<Vec<u8>> {
|
||||||
let len = text.len();
|
let len = text.len();
|
||||||
let cut1 = random_range(1..len - 1);
|
let cut1 = random_range(1..len - 1);
|
||||||
|
|||||||
+40
-1
@@ -7,6 +7,7 @@ use super::wall_clock;
|
|||||||
use crate::config::RequestContext;
|
use crate::config::RequestContext;
|
||||||
use crate::function::agents::run_agent_for_graph;
|
use crate::function::agents::run_agent_for_graph;
|
||||||
use crate::supervisor::mailbox::{Inbox, PeerRegistry};
|
use crate::supervisor::mailbox::{Inbox, PeerRegistry};
|
||||||
|
use crate::utils::create_abort_signal;
|
||||||
use anyhow::{Context, Error, Result, anyhow};
|
use anyhow::{Context, Error, Result, anyhow};
|
||||||
use log::{debug, warn};
|
use log::{debug, warn};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
@@ -145,8 +146,14 @@ async fn attempt_and_extract(
|
|||||||
)
|
)
|
||||||
.await?;
|
.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 {
|
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
|
.await
|
||||||
.with_context(|| {
|
.with_context(|| {
|
||||||
format!(
|
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<String, String> {
|
fn failure_capture_updates() -> HashMap<String, String> {
|
||||||
let mut u = HashMap::new();
|
let mut u = HashMap::new();
|
||||||
u.insert("failure".into(), "{{output}}".into());
|
u.insert("failure".into(), "{{output}}".into());
|
||||||
|
|||||||
+281
-18
@@ -3,11 +3,12 @@ use super::state_updates;
|
|||||||
use super::structured;
|
use super::structured;
|
||||||
use super::types::LlmNode;
|
use super::types::LlmNode;
|
||||||
use super::{is_transient_error, wall_clock};
|
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::prompts::DEFAULT_SKILL_INSTRUCTIONS;
|
||||||
use crate::config::{
|
use crate::config::{
|
||||||
Input, RequestContext, Role, RoleLike, SkillPolicy, should_inject_skill_instructions,
|
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::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;
|
||||||
@@ -16,6 +17,8 @@ use anyhow::{Context, Error, Result, anyhow, bail};
|
|||||||
use log::warn;
|
use log::warn;
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use std::collections::HashSet;
|
use std::collections::HashSet;
|
||||||
|
use std::future::Future;
|
||||||
|
use std::pin::Pin;
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::time::timeout;
|
use tokio::time::timeout;
|
||||||
|
|
||||||
@@ -36,7 +39,7 @@ impl LlmNodeExecutor {
|
|||||||
abort: &AbortSignal,
|
abort: &AbortSignal,
|
||||||
) -> Result<LlmExecutionOutcome> {
|
) -> Result<LlmExecutionOutcome> {
|
||||||
let result = run(node_id, node, state_manager, parent_ctx, abort).await;
|
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,
|
state_manager: &mut StateManager,
|
||||||
parent_ctx: &mut RequestContext,
|
parent_ctx: &mut RequestContext,
|
||||||
result: Result<String>,
|
result: Result<String>,
|
||||||
|
abort: &AbortSignal,
|
||||||
) -> Result<LlmExecutionOutcome> {
|
) -> Result<LlmExecutionOutcome> {
|
||||||
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, abort).await {
|
||||||
Ok(value) => (value, None),
|
Ok(value) => (value, None),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
warn!("llm node structured extraction failed: {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)),
|
node.mcp_tools.clone().map(|map| (node_id.to_string(), map)),
|
||||||
);
|
);
|
||||||
parent_ctx.refresh_mcp_tool_filters();
|
parent_ctx.refresh_mcp_tool_filters();
|
||||||
let result = match node.timeout.and_then(wall_clock) {
|
let result = bounded_run(
|
||||||
Some(d) => match timeout(d, run_with_retries(node, &prompt, parent_ctx, abort)).await {
|
node,
|
||||||
Ok(r) => r,
|
&prompt,
|
||||||
Err(_) => Err(anyhow!("llm node timed out after {}s", d.as_secs())),
|
parent_ctx,
|
||||||
},
|
abort,
|
||||||
None => run_with_retries(node, &prompt, parent_ctx, abort).await,
|
&mut default_completion_runner(),
|
||||||
};
|
)
|
||||||
|
.await;
|
||||||
parent_ctx.role = saved_role;
|
parent_ctx.role = saved_role;
|
||||||
let node_jobs =
|
let node_jobs =
|
||||||
std::mem::replace(&mut parent_ctx.node_job_scope, saved_job_scope).unwrap_or_default();
|
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<SavedAgent
|
|||||||
agent.set_enabled_skills(saved.enabled_skills);
|
agent.set_enabled_skills(saved.enabled_skills);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// One boxed model call against the input and ctx. Boxed for the same
|
||||||
|
/// reason as `AttemptFuture` in agent.rs: llm nodes run inside
|
||||||
|
/// `tokio::spawn`ed map branches, where higher-ranked opaque futures trip
|
||||||
|
/// the `Send` auto-trait solver.
|
||||||
|
type CompletionFuture<'a> =
|
||||||
|
Pin<Box<dyn Future<Output = Result<(String, Vec<ToolResult>)>> + 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<String> {
|
||||||
|
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(
|
async fn run_with_retries(
|
||||||
node: &LlmNode,
|
node: &LlmNode,
|
||||||
prompt: &str,
|
prompt: &str,
|
||||||
ctx: &mut RequestContext,
|
ctx: &mut RequestContext,
|
||||||
abort: &AbortSignal,
|
abort: &AbortSignal,
|
||||||
|
runner: &mut CompletionRunner<'_>,
|
||||||
) -> 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, abort).await {
|
match run_chat_loop(node, prompt, ctx, abort, runner).await {
|
||||||
Ok(out) => return Ok(out),
|
Ok(out) => return Ok(out),
|
||||||
Err(e) if is_transient_error(&e) && attempt < node.max_attempts => {
|
Err(e) if is_transient_error(&e) && attempt < node.max_attempts => {
|
||||||
warn!("llm node attempt {attempt} failed (transient): {e:#}; retrying");
|
warn!("llm node attempt {attempt} failed (transient): {e:#}; retrying");
|
||||||
@@ -281,6 +331,7 @@ async fn run_chat_loop(
|
|||||||
prompt: &str,
|
prompt: &str,
|
||||||
ctx: &mut RequestContext,
|
ctx: &mut RequestContext,
|
||||||
abort: &AbortSignal,
|
abort: &AbortSignal,
|
||||||
|
runner: &mut CompletionRunner<'_>,
|
||||||
) -> Result<String> {
|
) -> Result<String> {
|
||||||
let abort = abort.clone();
|
let abort = abort.clone();
|
||||||
let app_cfg = Arc::clone(&ctx.app.config);
|
let app_cfg = Arc::clone(&ctx.app.config);
|
||||||
@@ -293,11 +344,12 @@ async fn run_chat_loop(
|
|||||||
if abort.aborted() {
|
if abort.aborted() {
|
||||||
bail!("llm node aborted");
|
bail!("llm node aborted");
|
||||||
}
|
}
|
||||||
let client = input.create_client()?;
|
let (output, tool_results) = runner(&input, ctx, abort.clone()).await?;
|
||||||
ctx.before_chat_completion(&input)?;
|
// Defence in depth: the transport bails on abort itself; this closes
|
||||||
let (output, tool_results) =
|
// the window between its check and ours.
|
||||||
call_chat_completions(&input, false, false, client.as_ref(), ctx, abort.clone())
|
if abort.aborted() {
|
||||||
.await?;
|
bail!("llm node aborted");
|
||||||
|
}
|
||||||
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
|
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
|
||||||
|
|
||||||
if !output.is_empty() {
|
if !output.is_empty() {
|
||||||
@@ -755,6 +807,165 @@ mod tests {
|
|||||||
assert!(!is_last_turn(u32::MAX, u32::MAX));
|
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) -> F
|
||||||
|
where
|
||||||
|
F: for<'a> FnMut(&'a Input, &'a mut RequestContext, AbortSignal) -> CompletionFuture<'a>
|
||||||
|
+ Send,
|
||||||
|
{
|
||||||
|
f
|
||||||
|
}
|
||||||
|
|
||||||
|
fn boxed_completion<'a>(
|
||||||
|
fut: impl Future<Output = Result<(String, Vec<ToolResult>)>> + 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
|
/// 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
|
/// already-aborted graph never issues a model call and the node fails
|
||||||
/// with the abort reason rather than a connection error.
|
/// with the abort reason rather than a connection error.
|
||||||
@@ -877,7 +1088,13 @@ mod tests {
|
|||||||
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
ctx.agent = Some(Agent::test_new(AgentConfig::default()));
|
ctx.agent = Some(Agent::test_new(AgentConfig::default()));
|
||||||
|
|
||||||
let outcome = finish(&node, &mut state, &mut ctx, Ok("not json".to_string()))
|
let outcome = finish(
|
||||||
|
&node,
|
||||||
|
&mut state,
|
||||||
|
&mut ctx,
|
||||||
|
Ok("not json".to_string()),
|
||||||
|
&create_abort_signal(),
|
||||||
|
)
|
||||||
.await
|
.await
|
||||||
.expect("a declared fallback turns the failure into a route");
|
.expect("a declared fallback turns the failure into a route");
|
||||||
|
|
||||||
@@ -899,6 +1116,46 @@ mod tests {
|
|||||||
assert!(captured.contains("Invalid model"), "{captured}");
|
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
|
/// Without a fallback the extraction fault bails, wrapping the same
|
||||||
/// extractor chain the state_updates text carries.
|
/// extractor chain the state_updates text carries.
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
@@ -911,7 +1168,13 @@ mod tests {
|
|||||||
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
let mut ctx = RequestContext::new(Arc::new(AppState::test_default()), WorkingMode::Cmd);
|
||||||
ctx.agent = Some(Agent::test_new(AgentConfig::default()));
|
ctx.agent = Some(Agent::test_new(AgentConfig::default()));
|
||||||
|
|
||||||
let err = finish(&node, &mut state, &mut ctx, Ok("not json".to_string()))
|
let err = finish(
|
||||||
|
&node,
|
||||||
|
&mut state,
|
||||||
|
&mut ctx,
|
||||||
|
Ok("not json".to_string()),
|
||||||
|
&create_abort_signal(),
|
||||||
|
)
|
||||||
.await
|
.await
|
||||||
.expect_err("no fallback means the failure bails");
|
.expect_err("no fallback means the failure bails");
|
||||||
|
|
||||||
|
|||||||
@@ -155,4 +155,89 @@ mod tests {
|
|||||||
assert!(err.is_connect());
|
assert!(err.is_connect());
|
||||||
assert!(is_transient_error(&Error::new(err)));
|
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::<reqwest::Error>()
|
||||||
|
.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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
+26
-9
@@ -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::config::{Input, RequestContext, Role, RoleLike};
|
||||||
use crate::utils::create_abort_signal;
|
use crate::utils::AbortSignal;
|
||||||
use anyhow::{Context, Result, bail};
|
use anyhow::{Context, Result, bail};
|
||||||
use serde_json::Value;
|
use serde_json::Value;
|
||||||
use std::sync::Arc;
|
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.
|
- 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.";
|
- Do NOT invent fields not present in the schema.";
|
||||||
|
|
||||||
pub async fn extract(raw: &str, schema: &Value, parent_ctx: &mut RequestContext) -> Result<Value> {
|
pub async fn extract(
|
||||||
|
raw: &str,
|
||||||
|
schema: &Value,
|
||||||
|
parent_ctx: &mut RequestContext,
|
||||||
|
abort: &AbortSignal,
|
||||||
|
) -> Result<Value> {
|
||||||
if let Some(parsed) = try_parse_json(raw) {
|
if let Some(parsed) = try_parse_json(raw) {
|
||||||
return Ok(parsed);
|
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(
|
async fn extract_via_extractor(
|
||||||
raw: &str,
|
raw: &str,
|
||||||
schema: &Value,
|
schema: &Value,
|
||||||
parent_ctx: &mut RequestContext,
|
parent_ctx: &mut RequestContext,
|
||||||
|
abort: &AbortSignal,
|
||||||
is_repair: bool,
|
is_repair: bool,
|
||||||
) -> Result<Value> {
|
) -> Result<Value> {
|
||||||
let role = build_extractor_role(parent_ctx);
|
let role = build_extractor_role(parent_ctx);
|
||||||
@@ -38,7 +44,7 @@ async fn extract_via_extractor(
|
|||||||
|
|
||||||
let saved_role = parent_ctx.role.clone();
|
let saved_role = parent_ctx.role.clone();
|
||||||
parent_ctx.role = Some(role);
|
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;
|
parent_ctx.role = saved_role;
|
||||||
|
|
||||||
let output = result.context("Structured-output extractor LLM call failed")?;
|
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. \
|
"Structured-output extractor failed to produce valid JSON after repair retry. \
|
||||||
Last response:\n{output}"
|
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<String> {
|
async fn run_one_shot(
|
||||||
let abort = create_abort_signal();
|
prompt: &str,
|
||||||
|
ctx: &mut RequestContext,
|
||||||
|
abort: &AbortSignal,
|
||||||
|
) -> Result<String> {
|
||||||
|
if abort.aborted() {
|
||||||
|
bail!("Aborted.");
|
||||||
|
}
|
||||||
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 input = Input::from_str(ctx, prompt, role_for_input)?;
|
let input = Input::from_str(ctx, prompt, role_for_input)?;
|
||||||
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) =
|
||||||
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)?;
|
ctx.after_chat_completion(app_cfg.as_ref(), &input, &output, &tool_results)?;
|
||||||
|
|
||||||
Ok(output)
|
Ok(output)
|
||||||
|
|||||||
Reference in new issue
Block a user