From d3ddef0b7ec4705aa578c4eca24b96e9152485bd Mon Sep 17 00:00:00 2001 From: hugocasa Date: Fri, 9 Oct 2026 20:00:50 +0200 Subject: [PATCH] feat: retry transient AI provider failures before output reaches readers (#11623) * feat: retry transient AI provider failures before output reaches readers Co-Authored-By: Claude Opus 5.5 * fix: keep completed AI streams, stop retrying prompt blocks, bound retry-after Co-Authored-By: Claude Opus 5.5 * refactor: share the AI send-error retry rule between agent and proxy Co-Authored-By: Claude Opus 5.5 * fix: accept a clean end of stream from OpenAI-compatible providers Co-Authored-By: Claude Opus 5.5 * fix: classify AI provider errors the way the reference clients do Co-Authored-By: Claude Opus 5.5 * fix: count a Gemini function call as the end of its turn Co-Authored-By: Claude Opus 5.5 * fix: count an Anthropic stop_reason as the end of its stream Co-Authored-By: Claude Opus 5.5 * fix: accept a Gemini answer that reports usage and no output Co-Authored-By: Claude Opus 5.5 * fix: accept a Gemini stream that closes cleanly with no output Co-Authored-By: Claude Opus 5.5 --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Ruben Fiszel --- backend/Cargo.lock | 1 + backend/windmill-ai/Cargo.toml | 1 + backend/windmill-ai/src/ai_bedrock.rs | 23 + backend/windmill-ai/src/ai_google.rs | 56 +++ backend/windmill-ai/src/lib.rs | 1 + backend/windmill-ai/src/providers/bedrock.rs | 55 +- .../windmill-ai/src/providers/google_ai.rs | 53 +- backend/windmill-ai/src/providers/other.rs | 1 + backend/windmill-ai/src/retry.rs | 474 ++++++++++++++++++ backend/windmill-ai/src/sse.rs | 244 ++++++++- backend/windmill-api/src/ai.rs | 17 +- .../src/ai/stream_event_processor.rs | 20 +- backend/windmill-worker/src/ai_executor.rs | 58 ++- 13 files changed, 926 insertions(+), 78 deletions(-) create mode 100644 backend/windmill-ai/src/retry.rs diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 758627445d..64a152c962 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -14935,6 +14935,7 @@ dependencies = [ "indexmap 2.14.2", "lazy_static", "mime_guess", + "rand 0.9.5", "reqwest 0.13.5", "serde", "serde_json", diff --git a/backend/windmill-ai/Cargo.toml b/backend/windmill-ai/Cargo.toml index b1dea3ae54..44c4ac83c7 100644 --- a/backend/windmill-ai/Cargo.toml +++ b/backend/windmill-ai/Cargo.toml @@ -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 diff --git a/backend/windmill-ai/src/ai_bedrock.rs b/backend/windmill-ai/src/ai_bedrock.rs index a9318a1f3a..f55fcca415 100644 --- a/backend/windmill-ai/src/ai_bedrock.rs +++ b/backend/windmill-ai/src/ai_bedrock.rs @@ -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( + 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(error: &aws_sdk_bedrockruntime::error::SdkError) -> String where diff --git a/backend/windmill-ai/src/ai_google.rs b/backend/windmill-ai/src/ai_google.rs index bb8ae83e13..9f48edca8a 100644 --- a/backend/windmill-ai/src/ai_google.rs +++ b/backend/windmill-ai/src/ai_google.rs @@ -326,6 +326,17 @@ pub struct GeminiSSEEvent { pub candidates: Option>, #[serde(rename = "usageMetadata")] pub usage_metadata: Option, + #[serde(default)] + pub error: Option, + #[serde(rename = "promptFeedback", default)] + pub prompt_feedback: Option, +} + +/// 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, } // ============================================================================ @@ -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, Error> { if data.is_empty() || data == "[DONE]" { return Ok(None); @@ -567,6 +608,21 @@ pub fn parse_gemini_sse_event(data: &str) -> Result, 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 { diff --git a/backend/windmill-ai/src/lib.rs b/backend/windmill-ai/src/lib.rs index ccc40da432..634614a3cf 100644 --- a/backend/windmill-ai/src/lib.rs +++ b/backend/windmill-ai/src/lib.rs @@ -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; diff --git a/backend/windmill-ai/src/providers/bedrock.rs b/backend/windmill-ai/src/providers/bedrock.rs index 836244a666..5aa142a127 100644 --- a/backend/windmill-ai/src/providers/bedrock.rs +++ b/backend/windmill-ai/src/providers/bedrock.rs @@ -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 = 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() { diff --git a/backend/windmill-ai/src/providers/google_ai.rs b/backend/windmill-ai/src/providers/google_ai.rs index 97d0383ba4..bbf45e5e57 100644 --- a/backend/windmill-ai/src/providers/google_ai.rs +++ b/backend/windmill-ai/src/providers/google_ai.rs @@ -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 { 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 { 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 { 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::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::from(last)); } .boxed(); diff --git a/backend/windmill-ai/src/providers/other.rs b/backend/windmill-ai/src/providers/other.rs index 1ea54a7e79..fea8edbdac 100644 --- a/backend/windmill-ai/src/providers/other.rs +++ b/backend/windmill-ai/src/providers/other.rs @@ -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 diff --git a/backend/windmill-ai/src/retry.rs b/backend/windmill-ai/src/retry.rs new file mode 100644 index 0000000000..9de9d309da --- /dev/null +++ b/backend/windmill-ai/src/retry.rs @@ -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, +} + +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, retry_after: Option) -> 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::(), + _ => 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::from_str::(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 { + google_error_details(body).iter().find_map(|detail| { + let delay = detail.get("retryDelay")?.as_str()?.strip_suffix('s')?; + let seconds = delay + .parse::() + .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 { + let seconds = |name: &str, scale: f64| { + headers + .get(name)? + .to_str() + .ok()? + .trim() + .parse::() + .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, + #[serde(default, rename = "type")] + pub kind: Option, + #[serde(default)] + pub code: Option, + #[serde(default)] + pub status: Option, +} + +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::>() + .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::().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 { + 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::>() + .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 { + 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 { + 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 { + 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::() * 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>)] = &[ + ( + 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::(); + 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); + } +} diff --git a/backend/windmill-ai/src/sse.rs b/backend/windmill-ai/src/sse.rs index ada50afca5..bdaa75025f 100644 --- a/backend/windmill-ai/src/sse.rs +++ b/backend/windmill-ai/src/sse.rs @@ -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, + #[serde(default)] + pub finish_reason: Option, } /// Nested prompt token details returned by the Chat Completions API. @@ -87,6 +90,8 @@ pub struct OpenAISSEEvent { pub choices: Option>, #[serde(default)] pub usage: Option, + #[serde(default)] + pub error: Option, } 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, /// Token usage from final chunk (when stream_options.include_usage is true) pub usage: Option, + /// 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, + #[serde(default)] + delta: Option, }, #[serde(rename = "message_stop")] MessageStop {}, @@ -353,12 +400,18 @@ pub enum AnthropicSSEEvent { #[serde(rename = "error")] Error { #[serde(default)] - message: Option, + error: StreamErrorBody, }, #[serde(other)] Unknown, } +#[derive(Deserialize, Debug)] +pub struct AnthropicMessageDelta { + #[serde(default)] + pub stop_reason: Option, +} + /// 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, 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 = 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, pub used_websearch: bool, pub usage: Option, + 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, + #[serde(default)] + pub error: Option, } /// 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(default)] + message: Option, + #[serde(default)] + error: Option, + }, + /// Response created #[serde(rename = "response.created")] Created {}, @@ -890,6 +979,7 @@ pub struct OpenAIResponsesSSEParser { pub usage: Option, /// 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 = 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>) -> 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("a).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. diff --git a/backend/windmill-api/src/ai.rs b/backend/windmill-api/src/ai.rs index c401568868..cf7447ceba 100644 --- a/backend/windmill-api/src/ai.rs +++ b/backend/windmill-api/src/ai.rs @@ -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); diff --git a/backend/windmill-worker/src/ai/stream_event_processor.rs b/backend/windmill-worker/src/ai/stream_event_processor.rs index fe8a3b9b00..c4d17f080a 100644 --- a/backend/windmill-worker/src/ai/stream_event_processor.rs +++ b/backend/windmill-worker/src/ai/stream_event_processor.rs @@ -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>, pub handle: Option>, + /// Shared by every clone, so a sink lent to a parser counts into its owner. + streamed: Arc, } 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> { @@ -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()) diff --git a/backend/windmill-worker/src/ai_executor.rs b/backend/windmill-worker/src/ai_executor.rs index 7dd2382731..1b889aea34 100644 --- a/backend/windmill-worker/src/ai_executor.rs +++ b/backend/windmill-worker/src/ai_executor.rs @@ -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())); + } + } + } } };