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:
Dark-Alex-17 committed 2026-09-14 10:50:34 -06:00
1 parent 22421a4e74
commit 00c3b60c2f
6 files changed
+877 -53

No files matched your search

+26 -9
View File
@@ -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<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) {
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<Value> {
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<String> {
let abort = create_abort_signal();
async fn run_one_shot(
prompt: &str,
ctx: &mut RequestContext,
abort: &AbortSignal,
) -> Result<String> {
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)