starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-10-13 08:02:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: retry transient AI provider failures before output reaches readers (#11623)

Browse Source 
* feat: retry transient AI provider failures before output reaches readers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: keep completed AI streams, stop retrying prompt blocks, bound retry-after

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* refactor: share the AI send-error retry rule between agent and proxy

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: accept a clean end of stream from OpenAI-compatible providers

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: classify AI provider errors the way the reference clients do

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: count a Gemini function call as the end of its turn

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: count an Anthropic stop_reason as the end of its stream

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: accept a Gemini answer that reports usage and no output

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix: accept a Gemini stream that closes cleanly with no output

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Ruben Fiszel <ruben@windmill.dev>
This commit is contained in:
hugocasa
2026-10-09 20:00:50 +02:00
committed by GitHub
co-authored by Claude Opus 5.5 Ruben Fiszel
parent 02e4697157
commit d3ddef0b7e
13 changed files with 926 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/Cargo.lock
Generated
+1
View File
@@ -14935,6 +14935,7 @@ dependencies = [
"indexmap 2.14.2",
"lazy_static",
"mime_guess",
"rand 0.9.5",
"reqwest 0.13.5",
"serde",
"serde_json",
backend/windmill-ai/Cargo.toml
+1
View File
@@ -28,6 +28,7 @@ eventsource-stream.workspace = true
futures.workspace = true
http.workspace = true
mime_guess.workspace = true
rand.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
backend/windmill-ai/src/ai_bedrock.rs
+23
View File
@@ -477,6 +477,29 @@ pub async fn assume_role_with_oidc_token(
// Error Formatting
// ============================================================================
/// The error for a Bedrock stream that failed after it opened, retrying the exceptions the
/// Vercel AI SDK retries (`amazon-bedrock-stream-error.ts`) and a broken connection. The
/// SDK's own retries cover the opening call alone.
pub fn bedrock_stream_error<R: std::fmt::Debug>(
error: &aws_sdk_bedrockruntime::error::SdkError<
aws_sdk_bedrockruntime::types::error::ConverseStreamOutputError,
R,
>,
) -> Error {
let message = format!("Bedrock stream error: {}", format_bedrock_error(error));
match error.as_service_error() {
Some(e)
if !(e.is_internal_server_exception()
|| e.is_model_stream_error_exception()
|| e.is_service_unavailable_exception()
|| e.is_throttling_exception()) =>
{
Error::AIError(message)
}
_ => crate::retry::transient_error(message, None),
}
}
/// Format AWS SDK errors with detailed information
pub fn format_bedrock_error<E, R>(error: &aws_sdk_bedrockruntime::error::SdkError<E, R>) -> String
where
backend/windmill-ai/src/ai_google.rs
+56
View File
@@ -326,6 +326,17 @@ pub struct GeminiSSEEvent {
pub candidates: Option<Vec<GeminiSSECandidate>>,
#[serde(rename = "usageMetadata")]
pub usage_metadata: Option<GeminiUsageMetadata>,
#[serde(default)]
pub error: Option<crate::retry::StreamErrorBody>,
#[serde(rename = "promptFeedback", default)]
pub prompt_feedback: Option<GeminiPromptFeedback>,
}
/// Set when Gemini refuses the prompt itself, in which case no candidate follows.
#[derive(Deserialize, Debug)]
pub struct GeminiPromptFeedback {
#[serde(rename = "blockReason", default)]
pub block_reason: Option<String>,
}
// ============================================================================
@@ -550,10 +561,40 @@ pub fn openai_tools_to_gemini(
}
}
/// Follows a Gemini stream to tell a complete turn from one cut short.
#[derive(Default)]
pub struct GeminiTurn {
ended: bool,
has_output: bool,
}
impl GeminiTurn {
pub fn observe(&mut self, event: &GeminiParsedEvent) {
// A function call arrives whole and may come without a `finishReason`, so it
// counts too, as in gemini-cli's stream validation.
self.ended |= event.finish_reason.is_some() || !event.tool_calls.is_empty();
self.has_output |= event.text.is_some() || event.reasoning.is_some();
}
/// Whether the turn's final event arrived.
pub fn ended(&self) -> bool {
self.ended
}
/// Whether a stream that closed cleanly here holds the whole turn. Gemini 2.5 closes
/// its stream without a single event when thinking spends the whole
/// `maxOutputTokens`: an empty answer, which a retry would only get again, rather
/// than a cut one.
pub fn complete_at_close(&self) -> bool {
self.ended || !self.has_output
}
}
/// Parse one Gemini SSE data line into a [`GeminiParsedEvent`].
///
/// Returns `Ok(None)` for empty data or unrecognised payloads (e.g. `"[DONE]"`).
/// Logs a warning and returns `Ok(None)` on JSON parse errors rather than propagating.
/// An error the stream reports is returned as one.
pub fn parse_gemini_sse_event(data: &str) -> Result<Option<GeminiParsedEvent>, Error> {
if data.is_empty() || data == "[DONE]" {
return Ok(None);
@@ -567,6 +608,21 @@ pub fn parse_gemini_sse_event(data: &str) -> Result<Option<GeminiParsedEvent>, E
}
};
if let Some(error) = event.error {
return Err(error.into_error("Gemini"));
}
// `BLOCK_REASON_UNSPECIFIED` is the protobuf enum's placeholder, not a block.
if let Some(reason) = event
.prompt_feedback
.and_then(|f| f.block_reason)
.filter(|r| !matches!(r.as_str(), "" | "BLOCK_REASON_UNSPECIFIED"))
{
return Err(Error::AIError(format!(
"Gemini blocked the prompt: {reason}"
)));
}
let mut parsed = GeminiParsedEvent { usage: event.usage_metadata, ..Default::default() };
let Some(candidates) = event.candidates else {
backend/windmill-ai/src/lib.rs
+1
View File
@@ -11,6 +11,7 @@ pub mod model_context;
pub mod providers;
pub mod proxy;
pub mod query_builder;
pub mod retry;
pub mod sse;
pub mod types;
pub mod utils;
backend/windmill-ai/src/providers/bedrock.rs
+36 -19
View File
@@ -9,11 +9,12 @@
use super::{anthropic_model_rejects_sampling_params, REASONING_OFF_SENTINEL};
use crate::{
ai_bedrock::{
bedrock_model_supports_prompt_caching, bedrock_stream_event_is_block_stop,
bedrock_stream_event_to_reasoning_delta, bedrock_stream_event_to_text,
bedrock_stream_event_to_tool_delta, bedrock_stream_event_to_tool_delta_with_block_index,
bedrock_stream_event_to_tool_start, bedrock_stream_event_to_tool_start_with_block_index,
build_tool_config, create_inference_config, format_bedrock_error, json_to_document,
bedrock_model_supports_prompt_caching, bedrock_stream_error,
bedrock_stream_event_is_block_stop, bedrock_stream_event_to_reasoning_delta,
bedrock_stream_event_to_text, bedrock_stream_event_to_tool_delta,
bedrock_stream_event_to_tool_delta_with_block_index, bedrock_stream_event_to_tool_start,
bedrock_stream_event_to_tool_start_with_block_index, build_tool_config,
create_inference_config, format_bedrock_error, json_to_document,
openai_messages_to_bedrock, streaming_tool_calls_to_openai, BearerTokenProvider,
BedrockClient, StreamingToolCall,
},
@@ -24,6 +25,7 @@ use crate::{
image_handler::prepare_messages_for_api,
proxy::ProxyBuildArgs,
query_builder::{ParsedResponse, StreamEventSink},
retry::{openai_stream_error_chunk, truncated_stream_error},
types::{OpenAIMessage, StreamingEvent, TokenUsage, ToolDef},
};
use bytes::Bytes;
@@ -461,26 +463,31 @@ pub fn sdk_stream_to_sse(
async_stream::stream! {
let mut stream = stream;
let mut state = BedrockSseStreamState::new(id, model, created);
let mut message_stopped = false;
loop {
let failure = loop {
match stream.recv().await {
Ok(Some(event)) => {
if matches!(
event,
aws_sdk_bedrockruntime::types::ConverseStreamOutput::MessageStop(_)
) {
message_stopped = true;
}
for chunk in bedrock_sse_chunks_for_event(&event, &mut state) {
yield Ok(chunk);
}
}
Ok(None) => break,
Err(e) => {
yield Err(std::io::Error::new(
std::io::ErrorKind::Other,
e.to_string(),
));
break;
}
Ok(None) => break (!message_stopped).then(truncated_stream_error),
Err(_) if message_stopped => break None,
Err(e) => break Some(bedrock_stream_error(&e)),
}
}
};
yield Ok(Bytes::from("data: [DONE]\n\n"));
match failure {
Some(error) => yield Ok(Bytes::from(openai_stream_error_chunk(&error))),
None => yield Ok(Bytes::from("data: [DONE]\n\n")),
}
}
}
@@ -1081,11 +1088,19 @@ impl BedrockQueryBuilder {
// Claude reasoning block for the turn (only populated when thinking is on),
// attached to the first tool call for replay before toolUse.
let mut reasoning: Option<crate::ai_types::BedrockExtraContent> = None;
let mut message_stopped = false;
// Process stream events using shared parsing functions
loop {
match stream.recv().await {
Ok(Some(event)) => {
if matches!(
event,
aws_sdk_bedrockruntime::types::ConverseStreamOutput::MessageStop(_)
) {
message_stopped = true;
}
// Fold reasoning deltas into the turn's reasoning block (for
// replay before toolUse, required by Claude when thinking is
// on) and stream the readable summary as a thinking affordance.
@@ -1161,11 +1176,13 @@ impl BedrockQueryBuilder {
}
}
Ok(None) => break,
Err(e) => {
return Err(Error::internal_err(format!("Bedrock stream error: {}", e)));
}
Err(_) if message_stopped => break,
Err(e) => return Err(bedrock_stream_error(&e)),
}
}
if !message_stopped {
return Err(truncated_stream_error());
}
// Send tool call events to stream processor
if let Some(processor) = stream_event_sink.as_ref() {
backend/windmill-ai/src/providers/google_ai.rs
+28 -25
View File
@@ -5,17 +5,20 @@ use crate::{
parse_gemini_response, parse_gemini_sse_event, sanitize_schema_for_google,
GeminiFunctionDeclaration, GeminiGenerationConfig, GeminiImageContent, GeminiImageRequest,
GeminiImageResponse, GeminiInlineData, GeminiPart, GeminiPredictContent, GeminiTextRequest,
GeminiThinkingConfig, GeminiTool,
GeminiThinkingConfig, GeminiTool, GeminiTurn,
},
image_handler::{download_and_encode_s3_image, prepare_messages_for_api},
proxy::{ProxyBuildArgs, ProxyRequest},
query_builder::{BuildRequestArgs, ParsedResponse, QueryBuilder, StreamEventSink},
retry::{
openai_stream_error_chunk, send_with_retries, transient_error, truncated_stream_error,
},
sse::{GeminiSSEParser, SSEParser},
types::*,
};
use async_trait::async_trait;
use bytes::Bytes;
use eventsource_stream::Eventsource;
use eventsource_stream::{EventStreamError, Eventsource};
use futures::{stream::BoxStream, StreamExt};
use http::{header, HeaderMap, HeaderValue, Method, StatusCode};
use serde::Deserialize;
@@ -336,9 +339,7 @@ pub async fn handle_google_ai_chat_proxy(
) -> Result<GoogleAIProxyResponse, Error> {
let GoogleAIProxyRequest { request, model, stream } = build_google_ai_chat_proxy_request(args)?;
let response =
send_google_ai_proxy_request(client, request, "Failed to send request to Gemini API")
.await?;
let response = send_google_ai_proxy_request(client, request).await?;
if stream {
Ok(convert_streaming_response(response, &model))
@@ -356,8 +357,7 @@ pub async fn handle_google_ai_models_proxy(
args: &ProxyBuildArgs<'_>,
) -> Result<GoogleAIProxyResponse, Error> {
let request = build_google_ai_models_proxy_request(args);
let response =
send_google_ai_proxy_request(client, request, "Failed to fetch Gemini models").await?;
let response = send_google_ai_proxy_request(client, request).await?;
let gemini_resp: GeminiModelsResponse = response.json().await.map_err(|e| {
Error::internal_err(format!("Failed to parse Gemini models response: {}", e))
@@ -505,26 +505,13 @@ fn add_google_ai_auth_header(headers: &mut Vec<(String, String)>, api_key: &str,
async fn send_google_ai_proxy_request(
client: &reqwest::Client,
proxy_request: ProxyRequest,
send_error_message: &str,
) -> Result<reqwest::Response, Error> {
let mut request = client.request(proxy_request.method.clone(), &proxy_request.url);
for (header_name, header_value) in &proxy_request.headers {
request = request.header(header_name.as_str(), header_value.as_str());
}
let response = request
.body(proxy_request.body)
.send()
.await
.map_err(|e| Error::internal_err(format!("{}: {}", send_error_message, e)))?;
if let Err(e) = response.error_for_status_ref() {
let status = e.status().map(|s| s.to_string()).unwrap_or_default();
let body = response.text().await.unwrap_or_default();
return Err(Error::AIError(format!("{}: {}", status, body)));
}
Ok(response)
send_with_retries(request.body(proxy_request.body)).await
}
fn convert_streaming_response(response: reqwest::Response, model: &str) -> GoogleAIProxyResponse {
@@ -535,10 +522,15 @@ fn convert_streaming_response(response: reqwest::Response, model: &str) -> Googl
let openai_sse_stream = async_stream::stream! {
tokio::pin!(gemini_sse_stream);
let mut tool_call_index: usize = 0;
while let Some(event) = gemini_sse_stream.next().await {
let mut turn = GeminiTurn::default();
let failure = loop {
let Some(event) = gemini_sse_stream.next().await else {
break (!turn.complete_at_close()).then(truncated_stream_error);
};
match event {
Ok(event) => match parse_gemini_sse_event(&event.data) {
Ok(Some(parsed)) => {
turn.observe(&parsed);
for chunk in gemini_event_to_openai_sse_chunks(
&parsed, &id, &model, &mut tool_call_index,
) {
@@ -546,12 +538,23 @@ fn convert_streaming_response(response: reqwest::Response, model: &str) -> Googl
}
}
Ok(None) => {}
Err(e) => tracing::error!("Error parsing Gemini SSE event: {}", e),
Err(e) => break Some(e),
},
Err(EventStreamError::Transport(_)) if turn.ended() => break None,
Err(EventStreamError::Transport(e)) => {
break Some(transient_error(
format!("The connection to Gemini broke off mid-response: {e}"),
None,
));
}
Err(e) => tracing::error!("Error reading Gemini SSE stream: {}", e),
}
}
yield Ok::<Bytes, reqwest::Error>(Bytes::from("data: [DONE]\n\n"));
};
let last = match failure {
Some(error) => openai_stream_error_chunk(&error),
None => "data: [DONE]\n\n".to_string(),
};
yield Ok::<Bytes, reqwest::Error>(Bytes::from(last));
}
.boxed();
backend/windmill-ai/src/providers/other.rs
+1
View File
@@ -266,6 +266,7 @@ impl QueryBuilder for OtherQueryBuilder {
mut events_str,
stream_event_processor,
usage: openai_usage,
..
} = openai_sse_parser;
// Process streaming events with error handling
backend/windmill-ai/src/retry.rs
+474
View File
@@ -0,0 +1,474 @@
use std::time::Duration;
use reqwest::{header::HeaderMap, RequestBuilder, Response, StatusCode};
use windmill_common::error::{to_anyhow, Error};
/// How many times a provider call is retried after its first attempt.
pub const MAX_RETRIES: u32 = 3;
const BASE_DELAY: Duration = Duration::from_secs(1);
const MAX_BACKOFF: Duration = Duration::from_secs(10);
/// A provider that asks for a longer wait is out of capacity for longer than anyone should
/// block on (a spent daily quota, say), so its error is returned instead of waited out.
const MAX_RETRY_AFTER: Duration = Duration::from_secs(60);
/// A provider call that failed for a reason a later attempt may not meet: the provider was
/// overloaded, throttled or unreachable, or its response broke off part way.
#[derive(Debug)]
pub struct TransientAIError {
pub message: String,
/// The wait the provider asked for, when it named one.
pub retry_after: Option<Duration>,
}
impl std::fmt::Display for TransientAIError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(&self.message)
}
}
impl std::error::Error for TransientAIError {}
pub fn transient_error(message: impl Into<String>, retry_after: Option<Duration>) -> Error {
to_anyhow(TransientAIError { message: message.into(), retry_after }).into()
}
pub fn as_transient(error: &Error) -> Option<&TransientAIError> {
match error {
Error::Anyhow { error, .. } => error.downcast_ref::<TransientAIError>(),
_ => None,
}
}
/// The error to report once a transient failure is no longer retried, without the source
/// location `Error::Anyhow` renders.
pub fn into_final_error(error: Error) -> Error {
match as_transient(&error) {
Some(transient) => Error::AIError(transient.message.clone()),
None => error,
}
}
/// The statuses the OpenAI and Anthropic SDKs and the Vercel AI SDK retry.
fn is_retryable_status(status: u16) -> bool {
matches!(status, 408 | 409 | 429 | 500..=599)
}
/// The error for a provider's error response. The rules follow the OpenAI and Anthropic
/// SDKs (`x-should-retry`, then the status) and gemini-cli (a daily quota is spent until
/// tomorrow, and the wait Gemini asks for is in its body).
pub fn response_error(
status: StatusCode,
headers: &HeaderMap,
body: &str,
message: String,
) -> Error {
let retryable = match headers.get("x-should-retry").and_then(|v| v.to_str().ok()) {
Some("true") => true,
Some("false") => false,
_ => {
is_retryable_status(status.as_u16())
// OpenAI answers a spent balance with the 429 it also uses for rate limits.
&& !body.contains("insufficient_quota")
&& !google_quota_spent_for_the_day(body)
}
};
if retryable {
transient_error(
message,
retry_after(headers).or_else(|| google_retry_delay(body)),
)
} else {
Error::AIError(message)
}
}
/// The `details` of a Google API error body, where Gemini puts its quota and retry hints.
fn google_error_details(body: &str) -> Vec<serde_json::Value> {
serde_json::from_str::<serde_json::Value>(body)
.ok()
.and_then(|v| v.get("error")?.get("details")?.as_array().cloned())
.unwrap_or_default()
}
fn google_quota_spent_for_the_day(body: &str) -> bool {
google_error_details(body).iter().any(|detail| {
detail
.get("violations")
.and_then(|v| v.as_array())
.is_some_and(|violations| {
violations.iter().any(|violation| {
violation
.get("quotaId")
.and_then(|id| id.as_str())
.is_some_and(|id| id.contains("PerDay") || id.contains("Daily"))
})
})
})
}
/// The `RetryInfo.retryDelay` of a Google API error, a protobuf duration such as `"53s"`.
fn google_retry_delay(body: &str) -> Option<Duration> {
google_error_details(body).iter().find_map(|detail| {
let delay = detail.get("retryDelay")?.as_str()?.strip_suffix('s')?;
let seconds = delay
.parse::<f64>()
.ok()
.filter(|v| v.is_finite() && *v >= 0.0)?;
Some(Duration::try_from_secs_f64(seconds).unwrap_or(Duration::MAX))
})
}
/// The wait a provider asked for in its `retry-after-ms` or `retry-after` header. An
/// HTTP-date `retry-after` is ignored, which falls back to the computed backoff.
fn retry_after(headers: &HeaderMap) -> Option<Duration> {
let seconds = |name: &str, scale: f64| {
headers
.get(name)?
.to_str()
.ok()?
.trim()
.parse::<f64>()
.ok()
.filter(|v| v.is_finite() && *v >= 0.0)
.map(|v| Duration::try_from_secs_f64(v / scale).unwrap_or(Duration::MAX))
};
seconds("retry-after-ms", 1000.0).or_else(|| seconds("retry-after", 1.0))
}
/// The error object a provider sends inside a stream: OpenAI-compatible chunks carry
/// `type` and a string or status `code`, Anthropic `type`, Gemini a status `code` and
/// `status`.
#[derive(Debug, Default, serde::Deserialize)]
pub struct StreamErrorBody {
#[serde(default)]
pub message: Option<String>,
#[serde(default, rename = "type")]
pub kind: Option<String>,
#[serde(default)]
pub code: Option<serde_json::Value>,
#[serde(default)]
pub status: Option<String>,
}
impl StreamErrorBody {
/// The error for a failure an OpenAI-compatible or Gemini stream reported after opening
/// with a 200, classified as the Vercel AI SDK classifies OpenAI stream errors
/// (`openai-stream-error.ts`): the status `code` holds, else the status its code and type
/// name, else a server error.
pub fn into_error(self, provider: &str) -> Error {
let code = self.code_str();
let discriminator = [code.as_deref(), self.kind.as_deref()]
.into_iter()
.flatten()
.collect::<Vec<_>>()
.join(" ")
.to_lowercase();
let has = |terms: &[&str]| terms.iter().any(|term| discriminator.contains(term));
let status = match code.as_deref().and_then(|code| code.parse::<u16>().ok()) {
Some(status @ 400..=599) => status,
_ if has(&["insufficient_quota", "rate_limit"]) => 429,
_ if has(&["authentication"]) => 401,
_ if has(&["permission"]) => 403,
_ if has(&["not_found"]) => 404,
_ if has(&["invalid", "bad_request", "context_length"]) => 400,
_ if has(&["overload"]) => 503,
_ if has(&["timeout"]) => 504,
_ => 500,
};
let retryable = !has(&["insufficient_quota"]) && is_retryable_status(status);
self.into_classified_error(provider, retryable)
}
/// The error for a failure an Anthropic stream reported after opening with a 200. The
/// kinds retried are those the Vercel AI SDK retries, plus `timeout_error`, which
/// Anthropic documents as its 504; any other kind is a fault in the request or account.
pub fn into_anthropic_error(self) -> Error {
let retryable = matches!(
self.kind.as_deref(),
Some("api_error" | "overloaded_error" | "rate_limit_error" | "timeout_error")
);
self.into_classified_error("Anthropic", retryable)
}
fn code_str(&self) -> Option<String> {
match self.code.as_ref()? {
serde_json::Value::String(code) => Some(code.clone()),
serde_json::Value::Number(code) => Some(code.to_string()),
_ => None,
}
}
fn into_classified_error(self, provider: &str, retryable: bool) -> Error {
let kind = [self.kind.clone(), self.code_str(), self.status]
.into_iter()
.flatten()
.collect::<Vec<_>>()
.join("/");
let message = format!(
"{provider} stream error{}: {}",
(!kind.is_empty())
.then(|| format!(" ({kind})"))
.unwrap_or_default(),
self.message.as_deref().unwrap_or("no details")
);
if retryable {
transient_error(message, None)
} else {
Error::AIError(message)
}
}
}
/// The error for a stream that closed before the provider sent the event ending a
/// complete response: cut off by the provider or by a proxy in between.
pub fn truncated_stream_error() -> Error {
transient_error(
"The AI provider's response stream ended before the response was complete",
None,
)
}
/// An OpenAI-format stream chunk reporting `error`, which ends a response converted to that
/// format after it started streaming: the chat's OpenAI SDK throws on a chunk carrying
/// `error`, where a bare end of stream would read as a finished answer.
pub fn openai_stream_error_chunk(error: &Error) -> String {
let message = as_transient(error).map_or_else(|| error.to_string(), |t| t.message.clone());
format!(
"data: {}\n\n",
serde_json::json!({ "error": { "message": message } })
)
}
/// Sends a provider request, retrying a transient failure while nothing of the response
/// has been read: a failed connection or an error status. Once retries are spent the error
/// carries the provider's answer, as a single attempt's would.
pub async fn send_with_retries(request: RequestBuilder) -> Result<Response, Error> {
let mut backoff = Backoff::default();
loop {
// A streamed body cannot be replayed, so such a request is sent once.
let Some(attempt) = request.try_clone() else {
return send_once(request).await.map_err(into_final_error);
};
let error = match send_once(attempt).await {
Ok(response) => return Ok(response),
Err(error) => error,
};
let Some(delay) = as_transient(&error).and_then(|t| backoff.next_delay(t)) else {
return Err(into_final_error(error));
};
tracing::warn!(
"AI provider call failed ({error}); retrying in {:.1}s ({}/{MAX_RETRIES})",
delay.as_secs_f64(),
backoff.retries()
);
tokio::time::sleep(delay).await;
}
}
/// The error for a provider request that failed before any response arrived. The request
/// timeout is the caller's whole budget, so a timed-out call leaves no time for another
/// attempt; a connect timeout has its own, shorter one.
pub fn send_error(e: reqwest::Error, context: &str) -> Error {
let message = format!("{context}: {e}");
if e.is_builder() || (e.is_timeout() && !e.is_connect()) {
Error::AIError(message)
} else {
transient_error(message, None)
}
}
async fn send_once(request: RequestBuilder) -> Result<Response, Error> {
let response = request
.send()
.await
.map_err(|e| send_error(e, "Failed to reach the AI provider"))?;
let status = response.status();
if status.is_client_error() || status.is_server_error() {
let headers = response.headers().clone();
let body = response.text().await.unwrap_or_default();
return Err(response_error(status, &headers, &body, body.clone()));
}
Ok(response)
}
/// Counts the retries of one provider call and spaces them out.
#[derive(Debug, Default)]
pub struct Backoff {
retries: u32,
}
impl Backoff {
/// The wait before retrying a call that failed with `error`, or `None` once it should
/// fail instead.
pub fn next_delay(&mut self, error: &TransientAIError) -> Option<Duration> {
if self.retries >= MAX_RETRIES {
return None;
}
let delay = match error.retry_after {
Some(wait) if wait > MAX_RETRY_AFTER => return None,
// A zero wait would retry back to back, so it gets the computed backoff.
Some(wait) if !wait.is_zero() => wait,
// Jittered, so callers throttled together do not come back together.
_ => (BASE_DELAY * 2u32.pow(self.retries))
.min(MAX_BACKOFF)
.mul_f64(0.5 + rand::random::<f64>() * 0.5),
};
self.retries += 1;
Some(delay)
}
/// Retries granted so far.
pub fn retries(&self) -> u32 {
self.retries
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Error responses as providers send them, each with the retry decision of the client
/// the rule is taken from.
#[test]
fn classifies_error_responses() {
const GEMINI_QUOTA: &str = r#"{"error":{"code":429,"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.QuotaFailure","violations":[{"quotaId":"QUOTA_ID"}]},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"53s"}]}}"#;
let cases: &[(u16, &[(&str, &str)], &str, Option<Option<u64>>)] = &[
(
429,
&[],
r#"{"error":{"code":"rate_limit_exceeded"}}"#,
Some(None),
),
(
429,
&[],
r#"{"error":{"type":"insufficient_quota","code":"insufficient_quota"}}"#,
None,
),
(
529,
&[],
r#"{"type":"error","error":{"type":"overloaded_error"}}"#,
Some(None),
),
(501, &[], "", Some(None)),
(400, &[], "", None),
(401, &[], "", None),
(400, &[("x-should-retry", "true")], "", Some(None)),
(503, &[("x-should-retry", "false")], "", None),
(429, &[("retry-after", "2")], "", Some(Some(2))),
(
429,
&[],
&GEMINI_QUOTA.replace("QUOTA_ID", "GenerateRequestsPerMinutePerProjectPerModel"),
Some(Some(53)),
),
(
429,
&[],
&GEMINI_QUOTA.replace(
"QUOTA_ID",
"GenerateRequestsPerDayPerProjectPerModel-FreeTier",
),
None,
),
];
for (status, headers, body, expected) in cases {
let headers = headers
.iter()
.map(|(k, v)| (k.parse().unwrap(), v.parse().unwrap()))
.collect::<HeaderMap>();
let error = response_error(
StatusCode::from_u16(*status).unwrap(),
&headers,
body,
String::new(),
);
let got = as_transient(&error).map(|t| t.retry_after.map(|d| d.as_secs()));
assert_eq!(&got, expected, "{status} {headers:?} {body}");
}
}
/// In-stream error objects as providers send them, after the 200 opened the stream.
#[test]
fn classifies_stream_errors() {
let retried = |json: &str, anthropic: bool| {
let body: StreamErrorBody = serde_json::from_str(json).unwrap();
let error = if anthropic {
body.into_anthropic_error()
} else {
body.into_error("AI provider")
};
as_transient(&error).is_some()
};
for (json, expected) in [
(r#"{"message":"m","type":"server_error"}"#, true),
(
r#"{"message":"m","type":"requests","code":"rate_limit_exceeded"}"#,
true,
),
(
r#"{"message":"m","type":"insufficient_quota","code":"insufficient_quota"}"#,
false,
),
(
r#"{"message":"m","type":"invalid_request_error","code":"context_length_exceeded"}"#,
false,
),
// LiteLLM's proxy puts the upstream status in `code` as a string.
(r#"{"message":"m","type":"None","code":"400"}"#, false),
(r#"{"message":"m","code":"503"}"#, true),
(r#"{"code":503,"message":"m","status":"UNAVAILABLE"}"#, true),
(
r#"{"code":400,"message":"m","status":"INVALID_ARGUMENT"}"#,
false,
),
(r#"{"message":"m"}"#, true),
] {
assert_eq!(retried(json, false), expected, "{json}");
}
for (kind, expected) in [
("overloaded_error", true),
("api_error", true),
("rate_limit_error", true),
("timeout_error", true),
("billing_error", false),
("invalid_request_error", false),
("some_new_error", false),
] {
let json = format!(r#"{{"type":"{kind}","message":"m"}}"#);
assert_eq!(retried(&json, true), expected, "{kind}");
}
}
#[test]
fn honours_retry_after_up_to_a_bound() {
let mut headers = HeaderMap::new();
headers.insert("retry-after", "2".parse().unwrap());
assert_eq!(retry_after(&headers), Some(Duration::from_secs(2)));
headers.insert("retry-after-ms", "1500".parse().unwrap());
assert_eq!(retry_after(&headers), Some(Duration::from_millis(1500)));
headers.remove("retry-after-ms");
headers.insert("retry-after", "1e30".parse().unwrap());
assert_eq!(retry_after(&headers), Some(Duration::MAX));
let mut backoff = Backoff::default();
let asked = |secs| TransientAIError {
message: String::new(),
retry_after: Some(Duration::from_secs(secs)),
};
assert_eq!(backoff.next_delay(&asked(5)), Some(Duration::from_secs(5)));
assert_eq!(backoff.next_delay(&asked(3600)), None);
}
#[test]
fn stops_after_max_retries() {
let mut backoff = Backoff::default();
let error = TransientAIError { message: String::new(), retry_after: None };
for _ in 0..MAX_RETRIES {
let delay = backoff.next_delay(&error).expect("retry");
assert!(delay <= MAX_BACKOFF);
}
assert_eq!(backoff.next_delay(&error), None);
}
}
backend/windmill-ai/src/sse.rs
+231 -13
View File
@@ -1,6 +1,6 @@
use std::collections::HashMap;
use eventsource_stream::Eventsource;
use eventsource_stream::{EventStreamError, Eventsource};
use indexmap::IndexMap;
use reqwest::Response;
use serde::Deserialize;
@@ -8,12 +8,13 @@ use tokio_stream::StreamExt;
use windmill_common::{error::Error, utils::rd_string};
use crate::{
ai_google::{parse_gemini_sse_event, GeminiUsageMetadata},
ai_google::{parse_gemini_sse_event, GeminiTurn, GeminiUsageMetadata},
ai_types::UrlCitation,
ai_types::{
AnthropicExtraContent, ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall,
},
query_builder::StreamEventSink,
retry::{transient_error, truncated_stream_error, StreamErrorBody},
types::{StreamingEvent, TokenUsage},
};
@@ -43,6 +44,8 @@ pub struct OpenAIChoiceDelta {
#[derive(Deserialize)]
pub struct OpenAIChoice {
pub delta: Option<OpenAIChoiceDelta>,
#[serde(default)]
pub finish_reason: Option<String>,
}
/// Nested prompt token details returned by the Chat Completions API.
@@ -87,6 +90,8 @@ pub struct OpenAISSEEvent {
pub choices: Option<Vec<OpenAIChoice>>,
#[serde(default)]
pub usage: Option<OpenAIChatUsage>,
#[serde(default)]
pub error: Option<StreamErrorBody>,
}
lazy_static::lazy_static! {
@@ -99,6 +104,14 @@ lazy_static::lazy_static! {
pub trait SSEParser {
async fn parse_event_data(&mut self, data: &str) -> Result<(), Error>;
/// Whether the stream delivered the event that ends a complete response.
fn is_complete(&self) -> bool;
/// Whether a stream that closed cleanly holds a complete response.
fn is_complete_at_close(&self) -> bool {
self.is_complete()
}
async fn parse_events(&mut self, response: Response) -> Result<(), Error> {
let mut stream = response.bytes_stream().eventsource();
let mut consecutive_errors = 0;
@@ -113,6 +126,18 @@ pub trait SSEParser {
self.parse_event_data(&event.data).await?;
}
// The body ends right after a transport error, so it cannot be skipped
// like a malformed event. A timeout spent the caller's whole budget.
Err(EventStreamError::Transport(_)) if self.is_complete() => break,
Err(EventStreamError::Transport(e)) => {
let message =
format!("The connection to the AI provider broke off mid-response: {e}");
return Err(if e.is_timeout() {
Error::AIError(message)
} else {
transient_error(message, None)
});
}
Err(e) => {
consecutive_errors += 1;
tracing::error!("Failed to parse SSE event: {}", e);
@@ -130,6 +155,9 @@ pub trait SSEParser {
}
}
if !self.is_complete_at_close() {
return Err(truncated_stream_error());
}
Ok(())
}
}
@@ -145,6 +173,9 @@ pub struct OpenAISSEParser {
pub stream_event_processor: Box<dyn StreamEventSink>,
/// Token usage from final chunk (when stream_options.include_usage is true)
pub usage: Option<OpenAIChatUsage>,
/// Set by `[DONE]` or by a choice's `finish_reason`: OpenAI-compatible gateways do not
/// all send `[DONE]`.
complete: bool,
}
impl OpenAISSEParser {
@@ -156,13 +187,19 @@ impl OpenAISSEParser {
events_str: String::new(),
stream_event_processor,
usage: None,
complete: false,
}
}
}
impl SSEParser for OpenAISSEParser {
fn is_complete(&self) -> bool {
self.complete
}
async fn parse_event_data(&mut self, data: &str) -> Result<(), Error> {
if data == "[DONE]" {
self.complete = true;
return Ok(());
}
@@ -173,13 +210,21 @@ impl SSEParser for OpenAISSEParser {
.ok();
if let Some(event) = event {
if let Some(error) = event.error {
return Err(error.into_error("AI provider"));
}
// Extract usage from final chunk (when stream_options.include_usage is true)
if let Some(usage) = event.usage {
self.usage = Some(usage);
}
if let Some(mut choices) = event.choices.filter(|s| !s.is_empty()) {
if let Some(delta) = choices.remove(0).delta {
let choice = choices.remove(0);
if choice.finish_reason.is_some() {
self.complete = true;
}
if let Some(delta) = choice.delta {
if let Some(reasoning) = delta.reasoning_content.filter(|s| !s.is_empty()) {
self.accumulated_reasoning.push_str(&reasoning);
let event = StreamingEvent::ReasoningTokenDelta { content: reasoning };
@@ -345,6 +390,8 @@ pub enum AnthropicSSEEvent {
MessageDelta {
#[serde(default)]
usage: Option<AnthropicUsage>,
#[serde(default)]
delta: Option<AnthropicMessageDelta>,
},
#[serde(rename = "message_stop")]
MessageStop {},
@@ -353,12 +400,18 @@ pub enum AnthropicSSEEvent {
#[serde(rename = "error")]
Error {
#[serde(default)]
message: Option<String>,
error: StreamErrorBody,
},
#[serde(other)]
Unknown,
}
#[derive(Deserialize, Debug)]
pub struct AnthropicMessageDelta {
#[serde(default)]
pub stop_reason: Option<String>,
}
/// Tracks state of a content block during streaming
#[derive(Debug)]
#[allow(dead_code)]
@@ -391,6 +444,9 @@ pub struct AnthropicSSEParser {
/// can be replayed before `tool_use` (required by Claude when thinking is on).
pending_reasoning: Option<AnthropicExtraContent>,
reasoning_attached: bool,
/// Set by `message_stop` or by the `stop_reason` of the `message_delta` sent just
/// before it, so a gateway that drops `message_stop` does not fail a whole answer.
complete: bool,
}
impl AnthropicSSEParser {
@@ -407,11 +463,16 @@ impl AnthropicSSEParser {
usage: None,
pending_reasoning: None,
reasoning_attached: false,
complete: false,
}
}
}
impl SSEParser for AnthropicSSEParser {
fn is_complete(&self) -> bool {
self.complete
}
async fn parse_event_data(&mut self, data: &str) -> Result<(), Error> {
let event: Option<AnthropicSSEEvent> = serde_json::from_str(data)
.inspect_err(|e| {
@@ -583,9 +644,8 @@ impl SSEParser for AnthropicSSEParser {
.await?;
}
}
AnthropicSSEEvent::Error { message } => {
let error_msg = message.unwrap_or_else(|| "Unknown error".to_string());
tracing::error!("Anthropic streaming error: {}", error_msg);
AnthropicSSEEvent::Error { error } => {
return Err(error.into_anthropic_error());
}
AnthropicSSEEvent::MessageStart { message } => {
// The only event carrying the prompt-side counts. `message_delta`
@@ -595,7 +655,10 @@ impl SSEParser for AnthropicSSEParser {
self.usage = Some(usage);
}
}
AnthropicSSEEvent::MessageDelta { usage } => {
AnthropicSSEEvent::MessageDelta { usage, delta } => {
if delta.is_some_and(|d| d.stop_reason.is_some()) {
self.complete = true;
}
if let Some(u) = usage {
match &mut self.usage {
// Field by field, so the prompt counts from `message_start`
@@ -614,10 +677,9 @@ impl SSEParser for AnthropicSSEParser {
}
}
}
AnthropicSSEEvent::MessageStop {} => self.complete = true,
// Ignore other events
AnthropicSSEEvent::MessageStop {}
| AnthropicSSEEvent::Ping {}
| AnthropicSSEEvent::Unknown => {}
AnthropicSSEEvent::Ping {} | AnthropicSSEEvent::Unknown => {}
}
}
@@ -645,6 +707,7 @@ pub struct GeminiSSEParser {
pub annotations: Vec<UrlCitation>,
pub used_websearch: bool,
pub usage: Option<GeminiUsageMetadata>,
turn: GeminiTurn,
}
impl GeminiSSEParser {
@@ -659,15 +722,25 @@ impl GeminiSSEParser {
annotations: Vec::new(),
used_websearch: false,
usage: None,
turn: GeminiTurn::default(),
}
}
}
impl SSEParser for GeminiSSEParser {
fn is_complete(&self) -> bool {
self.turn.ended()
}
fn is_complete_at_close(&self) -> bool {
self.turn.complete_at_close()
}
async fn parse_event_data(&mut self, data: &str) -> Result<(), Error> {
let Some(parsed) = parse_gemini_sse_event(data)? else {
return Ok(());
};
self.turn.observe(&parsed);
if let Some(reasoning) = parsed.reasoning.filter(|s| !s.is_empty()) {
self.accumulated_reasoning.push_str(&reasoning);
@@ -797,6 +870,8 @@ impl OpenAIResponsesUsage {
pub struct OpenAIResponsesResponse {
#[serde(default)]
pub usage: Option<OpenAIResponsesUsage>,
#[serde(default)]
pub error: Option<StreamErrorBody>,
}
/// SSE event types for OpenAI Responses API streaming
@@ -833,6 +908,20 @@ pub enum OpenAIResponsesSSEEvent {
#[serde(rename = "response.incomplete")]
Incomplete { response: OpenAIResponsesResponse },
#[serde(rename = "response.failed")]
Failed { response: OpenAIResponsesResponse },
/// Flat (`code`, `message`) from OpenAI, nested under `error` from some compatible servers.
#[serde(rename = "error")]
Error {
#[serde(default)]
code: Option<serde_json::Value>,
#[serde(default)]
message: Option<String>,
#[serde(default)]
error: Option<StreamErrorBody>,
},
/// Response created
#[serde(rename = "response.created")]
Created {},
@@ -890,6 +979,7 @@ pub struct OpenAIResponsesSSEParser {
pub usage: Option<OpenAIResponsesUsage>,
/// Reasoning summary parts seen so far, to separate them as paragraphs
reasoning_summary_parts: usize,
response_ended: bool,
}
impl OpenAIResponsesSSEParser {
@@ -906,11 +996,16 @@ impl OpenAIResponsesSSEParser {
used_websearch: false,
usage: None,
reasoning_summary_parts: 0,
response_ended: false,
}
}
}
impl SSEParser for OpenAIResponsesSSEParser {
fn is_complete(&self) -> bool {
self.response_ended
}
async fn parse_event_data(&mut self, data: &str) -> Result<(), Error> {
let event: Option<OpenAIResponsesSSEEvent> = serde_json::from_str(data)
.inspect_err(|e| {
@@ -1010,11 +1105,22 @@ impl SSEParser for OpenAIResponsesSSEParser {
OpenAIResponsesSSEEvent::Completed { response }
| OpenAIResponsesSSEEvent::Incomplete { response } => {
self.response_ended = true;
if let Some(usage) = response.usage {
self.usage = Some(usage);
}
}
OpenAIResponsesSSEEvent::Failed { response } => {
return Err(response.error.unwrap_or_default().into_error("OpenAI"));
}
OpenAIResponsesSSEEvent::Error { code, message, error } => {
return Err(error
.unwrap_or(StreamErrorBody { code, message, ..Default::default() })
.into_error("OpenAI"));
}
OpenAIResponsesSSEEvent::ReasoningSummaryPartAdded {} => {
self.reasoning_summary_parts += 1;
if self.reasoning_summary_parts > 1 {
@@ -1037,9 +1143,10 @@ impl SSEParser for OpenAIResponsesSSEParser {
}
}
OpenAIResponsesSSEEvent::Done {} => self.response_ended = true,
// Ignore other event types
OpenAIResponsesSSEEvent::Done {}
| OpenAIResponsesSSEEvent::Created {}
OpenAIResponsesSSEEvent::Created {}
| OpenAIResponsesSSEEvent::InProgress {}
| OpenAIResponsesSSEEvent::OutputItemDone {}
| OpenAIResponsesSSEEvent::ContentPartAdded {}
@@ -1093,6 +1200,117 @@ mod tests {
}
}
fn sse_response(chunks: Vec<Result<&'static str, std::io::Error>>) -> Response {
let body = reqwest::Body::wrap_stream(futures::stream::iter(chunks));
Response::from(http::Response::new(body))
}
/// A stream that breaks off, or reports a failure inside its 200, must fail the call
/// rather than hand back what arrived as if it were the whole answer.
#[tokio::test]
async fn stream_failures_are_errors_not_partial_answers() {
use crate::retry::as_transient;
const STARTED: &str = concat!(
"data: {\"type\":\"message_start\"}\n\n",
"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hel\"}}\n\n",
);
let anthropic = |chunks| async move {
AnthropicSSEParser::new(Box::new(NoopSink))
.parse_events(sse_response(chunks))
.await
};
let overloaded = anthropic(vec![
Ok(STARTED),
Ok("data: {\"type\":\"error\",\"error\":{\"type\":\"overloaded_error\",\"message\":\"Overloaded\"}}\n\n"),
])
.await
.unwrap_err();
assert!(as_transient(&overloaded).is_some(), "{overloaded}");
let refused = anthropic(vec![
Ok(STARTED),
Ok("data: {\"type\":\"error\",\"error\":{\"type\":\"invalid_request_error\",\"message\":\"bad\"}}\n\n"),
])
.await
.unwrap_err();
assert!(as_transient(&refused).is_none(), "{refused}");
let truncated = anthropic(vec![Ok(STARTED)]).await.unwrap_err();
assert!(as_transient(&truncated).is_some(), "{truncated}");
let dropped = anthropic(vec![
Ok(STARTED),
Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
)),
])
.await
.unwrap_err();
assert!(as_transient(&dropped).is_some(), "{dropped}");
anthropic(vec![
Ok(STARTED),
Ok("data: {\"type\":\"message_stop\"}\n\n"),
])
.await
.unwrap();
anthropic(vec![
Ok(STARTED),
Ok("data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n"),
])
.await
.unwrap();
let truncated = OpenAISSEParser::new(Box::new(NoopSink))
.parse_events(sse_response(vec![Ok(
"data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n",
)]))
.await
.unwrap_err();
assert!(as_transient(&truncated).is_some(), "{truncated}");
let quota = OpenAIResponsesSSEParser::new(Box::new(NoopSink))
.parse_events(sse_response(vec![Ok(
"data: {\"type\":\"error\",\"error\":{\"type\":\"insufficient_quota\",\"code\":\"insufficient_quota\",\"message\":\"m\"}}\n\n",
)]))
.await
.unwrap_err();
assert!(as_transient(&quota).is_none(), "{quota}");
GeminiSSEParser::new(Box::new(NoopSink))
.parse_events(sse_response(vec![Ok(
"data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{}}}]}}]}\n\n",
)]))
.await
.unwrap();
let gemini = |chunks| async move {
GeminiSSEParser::new(Box::new(NoopSink))
.parse_events(sse_response(chunks))
.await
};
// Thinking spent the whole `maxOutputTokens`: a clean close with no event.
gemini(vec![]).await.unwrap();
let dropped = gemini(vec![Err(std::io::Error::new(
std::io::ErrorKind::ConnectionReset,
"reset",
))])
.await
.unwrap_err();
assert!(as_transient(&dropped).is_some(), "{dropped}");
let truncated = gemini(vec![Ok(
"data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Hi\"}]}}]}\n\n",
)])
.await
.unwrap_err();
assert!(as_transient(&truncated).is_some(), "{truncated}");
}
/// The prompt-side counts arrive only on `message_start` and the completion total
/// only on `message_delta`; a parser that keeps just the last one reports a request
/// with no input tokens at all.
backend/windmill-api/src/ai.rs
+5 -12
View File
@@ -40,6 +40,7 @@ use windmill_ai::proxy::{
fim::maybe_transform_fim_request, proxy_execution_mode, ProxyBuildArgs, ProxyExecutionMode,
ProxyRequest,
};
use windmill_ai::retry::send_with_retries;
use windmill_audit::{audit_oss::audit_log, ActionKind};
use windmill_common::db::UserDB;
use windmill_common::error::{to_anyhow, Error, Result};
@@ -1188,15 +1189,11 @@ async fn global_proxy(
}
};
let response = request.send().await.map_err(to_anyhow)?;
let response = send_with_retries(request).await;
audit_global_ai_request(&db, &authed).await?;
if response.error_for_status_ref().is_err() {
let err_msg = response.text().await.unwrap_or("".to_string());
return Err(Error::AIError(err_msg));
}
let response = response?;
let status_code = response.status();
let headers = response.headers().clone();
let stream = response.bytes_stream();
@@ -1576,7 +1573,7 @@ async fn proxy(
}
};
let response = request.send().await.map_err(to_anyhow)?;
let response = send_with_retries(request).await;
let mut tx = db.begin().await?;
@@ -1592,11 +1589,7 @@ async fn proxy(
.await?;
tx.commit().await?;
if response.error_for_status_ref().is_err() {
let err_msg = response.text().await.unwrap_or("".to_string());
return Err(Error::AIError(err_msg));
}
let response = response?;
let status_code = response.status();
let headers = response.headers().clone();
let is_sse = is_sse_response(&headers);
backend/windmill-worker/src/ai/stream_event_processor.rs
+17 -3
View File
@@ -1,3 +1,8 @@
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use async_trait::async_trait;
use windmill_ai::{query_builder::StreamEventSink, types::*};
use windmill_common::{error::Error, worker::Connection};
@@ -10,11 +15,13 @@ use crate::job_logger::append_result_stream;
pub struct StreamEventProcessor {
tx: Option<tokio::sync::mpsc::Sender<String>>,
pub handle: Option<tokio::task::JoinHandle<()>>,
/// Shared by every clone, so a sink lent to a parser counts into its owner.
streamed: Arc<AtomicUsize>,
}
impl Clone for StreamEventProcessor {
fn clone(&self) -> Self {
Self { tx: self.tx.clone(), handle: None }
Self { tx: self.tx.clone(), handle: None, streamed: self.streamed.clone() }
}
}
@@ -48,13 +55,19 @@ impl StreamEventProcessor {
}
});
Self { tx: Some(tx), handle: Some(handle) }
Self { tx: Some(tx), handle: Some(handle), streamed: Arc::default() }
}
/// Create a silent StreamEventProcessor that only accumulates events locally
/// without persisting them to the database. Used when streaming is disabled.
pub fn new_silent() -> Self {
Self { tx: None, handle: None }
Self { tx: None, handle: None, streamed: Arc::default() }
}
/// How many events have been streamed to the job's readers. A silent processor
/// streams none.
pub fn streamed_events(&self) -> usize {
self.streamed.load(Ordering::Relaxed)
}
pub fn to_handle(self) -> Option<tokio::task::JoinHandle<()>> {
@@ -79,6 +92,7 @@ impl StreamEventSink for StreamEventProcessor {
Ok(event_json) => {
let event_json = format!("{}\n", event_json);
events_str.push_str(&event_json);
self.streamed.fetch_add(1, Ordering::Relaxed);
if let Err(err) = tx
.send(event_json.clone())
backend/windmill-worker/src/ai_executor.rs
+52 -6
View File
@@ -43,6 +43,7 @@ use windmill_ai::{
common_outbound_headers, needs_unavailable_oauth_exchange, retain_effective_credentials,
},
query_builder::{BuildRequestArgs, ParsedResponse},
retry::{as_transient, into_final_error, response_error, send_error, Backoff, MAX_RETRIES},
types::*,
utils::{pinned_ai_client_for, should_use_structured_output_tool},
};
@@ -1659,10 +1660,17 @@ pub async fn run_agent(
.await;
let mut retried_context = false;
let mut backoff = Backoff::default();
let streamed_events = || {
stream_event_processor
.as_ref()
.map_or(0, |processor| processor.streamed_events())
};
let (parsed, request_message_count) = loop {
// How many messages this request carries, so the provider's prompt count can
// later be told apart from what the response and its tool results add.
let request_message_count = messages.context().len();
let streamed_before = streamed_events();
// Handle AWS Bedrock provider specially using the official SDK
let attempt: error::Result<_> = async {
@@ -1818,7 +1826,7 @@ pub async fn run_agent(
let resp = build_http_request(&endpoint, &auth_headers, request_body)
.send()
.await
.map_err(|e| Error::internal_err(format!("Failed to call API: {}", e)))?;
.map_err(|e| send_error(e, "Failed to call API"))?;
match resp.error_for_status_ref() {
Ok(_) => {
@@ -1829,6 +1837,7 @@ pub async fn run_agent(
}
Err(e) => {
let status = resp.status();
let headers = resp.headers().clone();
let text = resp
.text()
.await
@@ -1898,10 +1907,9 @@ pub async fn run_agent(
query_builder = create_chat_completions_query_builder(&credentials);
include_usage = true;
} else {
return Err(Error::internal_err(format!(
"API error calling {}: {} - {}",
endpoint, e, text
)));
let message =
format!("API error calling {}: {} - {}", endpoint, e, text);
return Err(response_error(status, &headers, &text, message));
}
}
}
@@ -1961,7 +1969,45 @@ pub async fn run_agent(
return Err(error);
}
}
Err(error) => return Err(error),
Err(error) => {
let Some(transient) = as_transient(&error) else {
return Err(error);
};
// Readers have no way to take back what they were shown, so a call
// that already streamed part of its answer fails as it is.
if streamed_events() != streamed_before {
append_logs(
&job.id,
&job.workspace_id,
"AI provider call failed after part of its answer had streamed; not retried.\n".to_string(),
conn,
)
.await;
return Err(into_final_error(error));
}
let Some(delay) = backoff.next_delay(transient) else {
return Err(into_final_error(error));
};
append_logs(
&job.id,
&job.workspace_id,
format!(
"AI provider call failed: {}. Retrying in {:.1}s (retry {}/{MAX_RETRIES}).\n",
transient.message,
delay.as_secs_f64(),
backoff.retries()
),
conn,
)
.await;
let mut cancelled = cancel_rx.clone();
tokio::select! {
_ = tokio::time::sleep(delay) => {}
Ok(_) = cancelled.wait_for(|cancelled| *cancelled) => {
return Err(Error::ExecutionErr("Job cancelled".to_string()));
}
}
}
}
};