From 27b56cf2e79dd272a8c1176c848ce15a53cfcb96 Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Mon, 24 Aug 2026 22:29:11 +0800 Subject: [PATCH 01/24] feat(core)!: own operation retries, timeouts, and the stream lifecycle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move retry and timeout from the HTTP layer into the user operations, mirroring the AI SDK's prepareRetries/retryWithExponentialBackoff RespectingRetryHeaders: a generic backoff loop that does not know ApiCallError, plus a wrapper that fills should_retry and get_delay_ms. RetryError { maxRetriesExceeded | errorNotRetryable } preserves the full attempt history; Full Jitter drives the exponential branch and a server retry-after(-ms) hint is honored exactly. Providers declare their policy through retry_config(); Core executes it around do_generate/do_stream. AbortSignal is a thin CancellationToken wrapper carrying only caller cancellation; total/first-chunk/chunk deadlines live in the operation future, with chunk timers reset only on semantic output (isOutputChunk parity). The consumption wrapper is also the stream's terminal fuse: Finish and non-recoverable errors end the stream under every configuration, while malformed frames surface as Err items and the stream continues (is_recoverable_stream_error), including through text()/consume() aggregation. VideoModel splits into do_start/do_status with a Core-owned poll loop matching the AI SDK generate-video flow: each phase retries independently, do_start — billable — mints one idempotency key per logical start outside the retry closure (a caller-supplied key wins), and the poll budget paces the loop while a hung status GET is bounded by the per-exchange response guard in provider-utils. VideoCallOptions defaults n and provider_options on deserialization for typed binding structs. Recording distinguishes operation attempts from exchange indices, records composite (Router/MoA) children as steps, and widens sensitive- key redaction to camelCase/kebab-case token keys. --- aimux-core/Cargo.toml | 10 +- aimux-core/src/abort_signal.rs | 37 + aimux-core/src/embedding_model.rs | 43 +- aimux-core/src/error.rs | 274 +++++-- aimux-core/src/files_model.rs | 6 +- aimux-core/src/generate.rs | 710 +++++++++++++++++- aimux-core/src/image_model.rs | 42 +- aimux-core/src/language_model.rs | 6 + aimux-core/src/lib.rs | 40 +- aimux-core/src/moa.rs | 249 +++++- aimux-core/src/openai_output.rs | 58 +- aimux-core/src/options.rs | 27 +- aimux-core/src/recording.rs | 455 +++++++++-- aimux-core/src/replay.rs | 5 + aimux-core/src/reranking_model.rs | 43 +- aimux-core/src/retry.rs | 644 ++++++++++++++++ aimux-core/src/router.rs | 301 +++++++- aimux-core/src/search_model.rs | 39 +- aimux-core/src/shared.rs | 56 +- aimux-core/src/speech_model.rs | 43 +- aimux-core/src/timeout.rs | 263 +++++++ aimux-core/src/trace/hash.rs | 4 + aimux-core/src/trace/layer.rs | 4 + aimux-core/src/transcription_model.rs | 62 +- aimux-core/src/video_model.rs | 549 +++++++++++++- aimux-core/tests/error_value_golden_test.rs | 148 ++-- aimux-core/tests/recording_e2e_test.rs | 13 +- .../tests/stream_timeout_matrix_test.rs | 150 ++++ 28 files changed, 3906 insertions(+), 375 deletions(-) create mode 100644 aimux-core/src/abort_signal.rs create mode 100644 aimux-core/src/retry.rs create mode 100644 aimux-core/src/timeout.rs create mode 100644 aimux-core/tests/stream_timeout_matrix_test.rs diff --git a/aimux-core/Cargo.toml b/aimux-core/Cargo.toml index afda2ec1..5cc8a1e1 100644 --- a/aimux-core/Cargo.toml +++ b/aimux-core/Cargo.toml @@ -17,11 +17,19 @@ futures = { workspace = true } async-stream = { workspace = true } tracing = { workspace = true } ts-rs = { workspace = true } +# Deliberately narrower than the workspace-wide `full`: core's production +# code uses only timers (`time`) and `select!` (`macros`). Keeping the +# declaration at the real API surface documents what Core actually asks +# of the runtime (RFC-0016 had removed core's tokio dependency; RFC-0031 +# reintroduces it for operation deadlines and retry backoff). +tokio = { version = "1", features = ["time", "macros"] } tokio-util = "0.7" +httpdate = { workspace = true } +rand = { workspace = true } base64 = "0.22" [dev-dependencies] -tokio = { workspace = true, features = ["macros", "rt"] } +tokio = { workspace = true, features = ["macros", "rt", "test-util"] } [lints] workspace = true diff --git a/aimux-core/src/abort_signal.rs b/aimux-core/src/abort_signal.rs new file mode 100644 index 00000000..c4ee0189 --- /dev/null +++ b/aimux-core/src/abort_signal.rs @@ -0,0 +1,37 @@ +//! Cooperative caller cancellation for Core operations. + +use tokio_util::sync::CancellationToken; + +/// A cancellation signal analogous to the Web `AbortSignal`. +/// +/// Timeouts are deliberately not encoded as cancellation reasons. Core owns +/// timeout deadlines and returns [`crate::AiMuxError::Timeout`] directly; +/// this type represents only caller-requested cancellation. +#[derive(Debug, Clone, Default)] +pub struct AbortSignal { + token: CancellationToken, +} + +impl AbortSignal { + /// Create a fresh signal. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Cancel with the default reason. + pub fn abort(&self) { + self.token.cancel(); + } + + /// Whether cancellation has been requested. + #[must_use] + pub fn is_aborted(&self) -> bool { + self.token.is_cancelled() + } + + /// Resolve when cancellation is requested. + pub fn cancelled(&self) -> impl std::future::Future + Send + 'static { + self.token.clone().cancelled_owned() + } +} diff --git a/aimux-core/src/embedding_model.rs b/aimux-core/src/embedding_model.rs index 9ffc8d01..7f180704 100644 --- a/aimux-core/src/embedding_model.rs +++ b/aimux-core/src/embedding_model.rs @@ -1,4 +1,4 @@ -//! The `EmbeddingModel` trait — the provider-facing interface for text embeddings. +//! The `EmbeddingModel` trait — the provider-facing interface for text embeddings. //! //! Aligned with Vercel AI SDK `EmbeddingModelV4` //! (`reference/ai/packages/provider/src/embedding-model/v4/`). @@ -11,9 +11,8 @@ use serde::{Deserialize, Serialize}; use ts_rs::TS; use crate::error::AiMuxError; -use crate::shared::{ - AbortSignal, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning, -}; +use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning}; +use crate::{AbortSignal, retry, timeout}; /// A single embedding vector. /// @@ -36,6 +35,12 @@ pub struct EmbeddingCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional provider-specific options, keyed by provider name. pub provider_options: Option, @@ -49,6 +54,8 @@ impl EmbeddingCallOptions { Self { values: vec![value.into()], abort_signal: None, + max_retries: None, + timeout: None, provider_options: None, headers: None, } @@ -121,6 +128,10 @@ pub trait EmbeddingModel: Send + Sync { /// Provider-specific model ID, e.g. `"text-embedding-3-small"`. fn model_id(&self) -> &str; + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Limit of how many embeddings can be generated in a single API call. /// /// `None` means the model has no fixed limit. The TS spec allows this to @@ -136,3 +147,27 @@ pub trait EmbeddingModel: Send + Sync { async fn do_embed(&self, options: &EmbeddingCallOptions) -> Result; } + +/// User-facing embedding operation with Core-owned retry and timeout. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +pub async fn embed( + model: &dyn EmbeddingModel, + options: EmbeddingCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + timeout::run( + retries.retry(|| model.do_embed(&options)), + abort_signal.as_ref(), + timeout, + ) + .await +} diff --git a/aimux-core/src/error.rs b/aimux-core/src/error.rs index e9d81741..d52c7e1b 100644 --- a/aimux-core/src/error.rs +++ b/aimux-core/src/error.rs @@ -1,5 +1,9 @@ //! Error types for aimux-core. +use std::collections::HashMap; +use std::sync::LazyLock; +use std::time::SystemTime; + use serde::{Deserialize, Serialize}; use thiserror::Error; use ts_rs::TS; @@ -15,18 +19,22 @@ use ts_rs::TS; /// body read) are `ApiCall` errors too — no response arrived, so /// `status_code` is `None` and `is_retryable` is `true`, exactly as the AI /// SDK's `handleFetchError` builds an `APICallError` with no `statusCode` -/// and `isRetryable: true`. `message` holds the provider's text **only** — -/// the status is a field, never baked into the string; `Display` composes -/// the human-readable form from the fields at print time, so nothing -/// downstream has to parse it back out. +/// and `isRetryable: true`. `message` holds the provider's text verbatim when +/// one is available; transport and parsing failures use a library message +/// with the source detail appended. The HTTP status is always a field, never +/// baked into the string; `Display` composes the human-readable form from the +/// fields at print time, so nothing downstream has to parse it back out. /// -/// `APICallError` fields not carried in this round: `url` / -/// `requestBodyValues` / `responseHeaders` (the request context is not -/// available at the error-construction sites today; all fields are -/// `#[serde(default)]`, so adding them later is not a breaking change). -#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)] +/// Every producer must provide the sanitized request URL and request values. +/// This keeps transport and response failures self-contained and matches the +/// AI SDK's `APICallError` contract. +#[derive(Debug, Clone, Serialize, Deserialize, TS)] #[ts(export)] pub struct ApiCallError { + /// Sanitized request URL. Required for every API-derived failure. + pub url: String, + /// Sanitized request values used to create the API request. + pub request_body_values: serde_json::Value, /// HTTP status of the response, when it came from one /// (`APICallError.statusCode`). Always the *observed* status: the HTTP /// layer fills it for every response-derived error; errors built without @@ -38,23 +46,21 @@ pub struct ApiCallError { /// status. Our normalized take on `APICallError.data`. #[serde(default)] pub provider_code: Option, - /// The provider's message, verbatim. No status prefix. + /// Human-readable failure text. Provider text is verbatim when available; + /// locally detected transport/parse failures include their source detail. + /// Never includes an HTTP status prefix. pub message: String, /// The raw response body, verbatim (`APICallError.responseBody`) — the /// lossless evidence when `message`/`provider_code` are extractions. /// `None` when the error did not come from a response body. #[serde(default)] pub response_body: Option, - /// Provider-assigned request id, from the `x-request-id` / `request-id` - /// response header when one was sent. Filled by the shared HTTP layer. + /// Sanitized response headers. #[serde(default)] - pub request_id: Option, - /// Retry hint in milliseconds, distilled from the `retry-after-ms` / - /// `retry-after` response headers on a 429. The classification lives in - /// `status_code`; this is the matching action input. + pub response_headers: Option>, + /// Parsed provider error data. #[serde(default)] - #[ts(type = "number | null")] - pub retry_after_ms: Option, + pub data: Option, /// Whether retrying can help (`APICallError.isRetryable`) — stored at /// construction, exactly like the AI SDK: the response path computes it /// from the status (408/409/429 or 5xx), and the transport path (no response, so @@ -63,6 +69,82 @@ pub struct ApiCallError { pub is_retryable: bool, } +impl ApiCallError { + /// Create an API failure with its required request context. + pub fn new( + message: impl Into, + url: impl Into, + request_body_values: serde_json::Value, + ) -> Self { + Self { + url: url.into(), + request_body_values, + status_code: None, + provider_code: None, + message: message.into(), + response_body: None, + response_headers: None, + data: None, + is_retryable: false, + } + } +} + +/// Why the retry wrapper stopped retrying. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, TS)] +#[serde(rename_all = "camelCase")] +#[ts(export)] +pub enum RetryErrorReason { + /// Every permitted retry attempt failed with a retryable error. + MaxRetriesExceeded, + /// A later attempt produced a non-retryable error. + ErrorNotRetryable, +} + +/// Complete error history for a retried model operation. +#[derive(Debug, Clone, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct RetryError { + pub reason: RetryErrorReason, + pub errors: Vec, +} + +static EMPTY_RETRY_HISTORY_ERROR: LazyLock = + LazyLock::new(|| AiMuxError::Other("Retry failed without recorded errors".into())); + +impl RetryError { + /// The final attempt error, or a stable invalid-history error when a + /// deserialized or manually constructed value has no attempts. + #[must_use] + pub fn last_error(&self) -> &AiMuxError { + match self.errors.last() { + Some(error) => error, + None => &EMPTY_RETRY_HISTORY_ERROR, + } + } +} + +impl std::fmt::Display for RetryError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let attempts = self.errors.len(); + let Some(last_error) = self.errors.last() else { + return f.write_str("Retry failed without recorded errors"); + }; + match self.reason { + RetryErrorReason::MaxRetriesExceeded => { + write!( + f, + "Failed after {attempts} attempts. Last error: {last_error}" + ) + } + RetryErrorReason::ErrorNotRetryable => write!( + f, + "Failed after {attempts} attempts with non-retryable error: '{last_error}'" + ), + } + } +} + impl std::fmt::Display for ApiCallError { /// `HTTP {status}: {message}` when a status is known (just `HTTP {status}` /// if the response had no body), else `{message}`. `provider_code` and @@ -79,19 +161,16 @@ impl std::fmt::Display for ApiCallError { /// Unified error type for all aimux operations. /// /// Variants are cut by *what the caller does about it*, not by where the -/// failure came from. `ApiCall` carries the [`ApiCallError`] detail inline -/// (unboxed, like async-openai's `ApiError(ApiError)`); the size guard in -/// `error_value_golden_test` keeps the enum under clippy's -/// `result_large_err` threshold. -/// -/// Constructed as a named struct literal everywhere -/// (`ApiCallError { status_code: .., ..Default::default() }`), the same -/// shape as the AI SDK's named-options constructor. +/// failure came from. `ApiCallError` is boxed only to keep the Rust enum +/// compact; serde and every binding still observe the same object shape. #[derive(Debug, Clone, Serialize, Deserialize, TS, Error)] #[ts(export)] pub enum AiMuxError { #[error("API call error: {0}")] - ApiCall(ApiCallError), + ApiCall(Box), + + #[error("{0}")] + Retry(RetryError), #[error("JSON parse error: {0}")] JsonParse(String), @@ -142,8 +221,8 @@ pub enum AiMuxError { #[error("request timed out: {0}")] Timeout(String), - #[error("request aborted")] - Aborted, + #[error("{0}")] + Aborted(String), #[error("{0}")] Other(String), @@ -163,7 +242,30 @@ impl From for AiMuxError { } } +/// The AI SDK's status-based retryability rule: 408 (timeout), 409 +/// (conflict), 429 (rate limit) and every 5xx are worth retrying; everything +/// else is not. One definition — response handlers and providers must not +/// hand-roll this formula. +#[must_use] +pub fn is_retryable_status(status: u16) -> bool { + matches!(status, 408 | 409 | 429) || status >= 500 +} + impl AiMuxError { + /// Returns `true` for a malformed individual stream frame that does not + /// prove the underlying transport has failed. Stream reducers may report + /// these errors and continue polling for later, independently framed data. + #[must_use] + pub fn is_recoverable_stream_error(&self) -> bool { + matches!(self, Self::JsonParse(_) | Self::InvalidResponseData(_)) + } + + /// Convert caller cancellation into the public error model. + #[must_use] + pub fn from_abort_signal(_signal: &crate::AbortSignal) -> Self { + Self::Aborted("request aborted".into()) + } + /// Returns `true` if the error is transient and can be retried. /// /// Reads the stored [`ApiCallError::is_retryable`] field, filled at @@ -181,26 +283,44 @@ impl AiMuxError { } /// Returns the retry-after delay hint (in milliseconds) carried by this - /// error, if any — distilled from the `retry-after-ms` / `retry-after` - /// response headers into [`ApiCallError::retry_after_ms`]. Returns `None` - /// for errors that don't advertise a delay, in which case the retry - /// strategy falls back to exponential backoff. + /// error's response headers, if any. /// /// Mirrors the header-consulting behaviour of /// `retryWithExponentialBackoffRespectingRetryHeaders` in the TS SDK. #[must_use] pub fn retry_after_hint(&self) -> Option { - match self { - AiMuxError::ApiCall(d) => d.retry_after_ms.map(|ms| ms as i64), - _ => None, - } + let AiMuxError::ApiCall(detail) = self else { + return None; + }; + let headers = detail.response_headers.as_ref()?; + headers + .get("retry-after-ms") + .and_then(|value| value.parse::().ok()) + .filter(|value| !value.is_nan()) + .map(|value| value as i64) + .or_else(|| { + headers.get("retry-after").and_then(|value| { + value + .parse::() + .ok() + .filter(|value| !value.is_nan()) + .map(|seconds| (seconds * 1_000.0) as i64) + .or_else(|| { + httpdate::parse_http_date(value).ok().map(|date| { + match date.duration_since(SystemTime::now()) { + Ok(duration) => duration.as_millis() as i64, + Err(error) => -(error.duration().as_millis() as i64), + } + }) + }) + }) + }) } /// Returns the HTTP status code carried by this error, if any. /// - /// This is the *observed* status only — the field the HTTP layer filled - /// when a response actually arrived (`parse_provider_error`, - /// `send_with_retry_raw`). An error built without an HTTP exchange (a + /// This is the *observed* status only — the field the selected response + /// handler filled when a response actually arrived. An error built without an HTTP exchange (a /// missing API key, a bare constructor, a transport failure) reports /// `None`: no status was seen, so none is invented. /// @@ -222,34 +342,29 @@ impl AiMuxError { mod tests { use super::*; + fn api_error(message: &str) -> ApiCallError { + ApiCallError::new(message, "https://example.test", serde_json::json!({})) + } + /// A bare constructor saw no HTTP exchange, so it reports no status — /// nothing is invented. Response-derived errors get theirs from the field /// (filled by the HTTP layer). #[test] fn bare_constructors_report_no_status() { assert_eq!( - AiMuxError::ApiCall(ApiCallError { - message: "no api key".into(), - ..Default::default() - }) - .status_code(), + AiMuxError::ApiCall(Box::new(api_error("no api key"))).status_code(), None ); assert_eq!( - AiMuxError::ApiCall(ApiCallError { - message: "reset".into(), + AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }) + ..api_error("reset") + })) .status_code(), None ); assert_eq!( - AiMuxError::ApiCall(ApiCallError { - message: "boom".into(), - ..Default::default() - }) - .status_code(), + AiMuxError::ApiCall(Box::new(api_error("boom"))).status_code(), None ); } @@ -269,21 +384,16 @@ mod tests { #[test] fn status_code_reads_the_field_not_the_message() { assert_eq!( - AiMuxError::ApiCall(ApiCallError { + AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(403), - message: "forbidden".into(), - ..Default::default() - }) + ..api_error("forbidden") + })) .status_code(), Some(403) ); // A message that merely *looks* like an HTTP prefix carries no status. assert_eq!( - AiMuxError::ApiCall(ApiCallError { - message: "HTTP 403: forbidden".into(), - ..Default::default() - }) - .status_code(), + AiMuxError::ApiCall(Box::new(api_error("HTTP 403: forbidden"))).status_code(), None ); assert_eq!( @@ -301,37 +411,43 @@ mod tests { fn timeout_is_not_retryable_transport_is() { assert!(!AiMuxError::Timeout("total timeout".into()).is_retryable()); assert!( - AiMuxError::ApiCall(ApiCallError { - message: "connection reset".into(), + AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }) + ..api_error("connection reset") + })) .is_retryable() ); assert!(!AiMuxError::JsonParse("parse".into()).is_retryable()); } + #[test] + fn empty_retry_history_round_trips_without_panicking() { + let json = r#"{"reason":"maxRetriesExceeded","errors":[]}"#; + let error: RetryError = serde_json::from_str(json).unwrap(); + + assert!(error.errors.is_empty()); + assert_eq!(error.to_string(), "Retry failed without recorded errors"); + assert_eq!( + error.last_error().to_string(), + "Retry failed without recorded errors" + ); + assert_eq!(serde_json::to_string(&error).unwrap(), json); + } + /// A 429 is retryable and carries its hint wherever it lives — the /// classification is the `status_code` field, not a variant. #[test] fn rate_limit_is_read_from_the_field() { - let err = AiMuxError::ApiCall(ApiCallError { + let err = AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(429), - message: "quota exceeded".into(), - retry_after_ms: Some(7), + response_headers: Some(HashMap::from([("retry-after-ms".into(), "7".into())])), is_retryable: true, - ..Default::default() - }); + ..api_error("quota exceeded") + })); assert!(err.is_retryable()); assert_eq!(err.retry_after_hint(), Some(7)); // A bare ApiCall error (no stored verdict) is not retryable. - assert!( - !AiMuxError::ApiCall(ApiCallError { - message: "boom".into(), - ..Default::default() - }) - .is_retryable() - ); + assert!(!AiMuxError::ApiCall(Box::new(api_error("boom"))).is_retryable()); let json = serde_json::to_string(&err).unwrap(); let back: AiMuxError = serde_json::from_str(&json).unwrap(); diff --git a/aimux-core/src/files_model.rs b/aimux-core/src/files_model.rs index 1473729c..d2fcb112 100644 --- a/aimux-core/src/files_model.rs +++ b/aimux-core/src/files_model.rs @@ -1,4 +1,4 @@ -//! The `Files` trait — the provider-facing interface for file management. +//! The `Files` trait — the provider-facing interface for file management. //! //! Aligned with Vercel AI SDK `FilesV4` //! (`reference/ai/packages/provider/src/files/v4/`). @@ -10,10 +10,10 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use ts_rs::TS; +use crate::AbortSignal; use crate::error::AiMuxError; use crate::shared::{ - AbortSignal, FileBytes, SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, - Warning, + FileBytes, SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, Warning, }; /// File data accepted by [`Files::upload_file`]. diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs index fd4af99b..9059cecc 100644 --- a/aimux-core/src/generate.rs +++ b/aimux-core/src/generate.rs @@ -10,7 +10,7 @@ use std::collections::HashMap; use std::pin::Pin; use std::time::Instant; -use futures::Stream; +use futures::{Stream, StreamExt}; use serde::{Deserialize, Serialize}; use serde_json::Value; use tracing::Instrument; @@ -29,6 +29,19 @@ use crate::result::{ use crate::stream_part::StreamPart; use crate::tool::Tool; use crate::types::{FinishReason, ReasoningEffort, Usage, Warning}; +use crate::{AbortSignal, retry, timeout}; + +/// Matches the AI SDK's `isOutputChunk`: only chunks containing model output +/// start or reset the first/chunk output timers. +fn is_output_chunk(part: &StreamPart) -> bool { + match part { + StreamPart::TextDelta { delta, .. } + | StreamPart::ReasoningDelta { delta, .. } + | StreamPart::ToolInputDelta { delta, .. } => !delta.is_empty(), + StreamPart::ToolCall { .. } | StreamPart::File { .. } => true, + _ => false, + } +} // ───────────────────────────────────────────────────────────────────────────── // User-facing options @@ -77,7 +90,7 @@ pub struct GenerateTextOptions { /// directly; the Node binding bridges a JS `AbortSignal` natively. #[serde(skip)] #[ts(skip)] - pub abort_signal: Option, + pub abort_signal: Option, /// Emit raw provider stream chunks as `StreamPart::Raw` (debugging aid). pub include_raw_chunks: Option, } @@ -229,17 +242,55 @@ impl StreamTextResult { /// Returns the stream's error part or transport error, propagated from the /// underlying model stream. pub async fn text(self) -> Result { - use futures::StreamExt; let mut result = String::new(); + let mut saw_output = false; + let mut saw_finish = false; + let mut provider_error: Option = None; let mut stream = self.stream; while let Some(part) = stream.next().await { - match part? { + let part = match part { + Ok(part) => part, + // A recoverable frame error is data on this stream (core keeps + // the stream alive past it); fold it into the same first-error + // slot as a provider error event and keep consuming so the + // trailing Finish still lands (usage, recording). + Err(error) if error.is_recoverable_stream_error() => { + if provider_error.is_none() { + provider_error = Some(error); + } + continue; + } + Err(error) => return Err(error), + }; + if is_output_chunk(&part) { + saw_output = true; + } + match part { StreamPart::TextDelta { delta, .. } => result.push_str(&delta), - StreamPart::Finish { .. } => break, - StreamPart::Error { error } => return Err(error), + StreamPart::Finish { .. } => { + saw_finish = true; + break; + } + // Keep consuming after a provider error event: a trailing + // Finish still carries usage for the recording layer. The + // error is returned once the stream reaches its end. + StreamPart::Error { error } if provider_error.is_none() => { + provider_error = Some(error); + } _ => {} } } + if let Some(error) = provider_error { + return Err(error); + } + // AI SDK semantics: an incomplete stream with no output at all is an + // error; an incomplete stream with partial output retains the partial + // result. + if !saw_finish && !saw_output { + return Err(AiMuxError::InvalidResponseData( + "No output generated. The model stream ended without a finish chunk.".into(), + )); + } Ok(result) } @@ -254,8 +305,6 @@ impl StreamTextResult { /// Returns the stream's error part or transport error, propagated from the /// underlying model stream. pub async fn consume(self) -> Result { - use futures::StreamExt; - let mut text = String::new(); let mut reasoning: Vec = Vec::new(); let mut reasoning_text_buf = String::new(); @@ -273,9 +322,29 @@ impl StreamTextResult { let mut response: Option = None; let mut response_content_parts: Vec = Vec::new(); + let mut saw_output = false; + let mut saw_finish = false; + let mut provider_error: Option = None; let mut stream = self.stream; while let Some(part) = stream.next().await { - match part? { + let part = match part { + Ok(part) => part, + // A recoverable frame error is data on this stream (core keeps + // the stream alive past it); fold it into the same first-error + // slot as a provider error event and keep consuming so the + // trailing Finish still lands (usage, recording). + Err(error) if error.is_recoverable_stream_error() => { + if provider_error.is_none() { + provider_error = Some(error); + } + continue; + } + Err(error) => return Err(error), + }; + if is_output_chunk(&part) { + saw_output = true; + } + match part { StreamPart::TextDelta { delta, .. } => { text.push_str(&delta); // Accumulate for response_messages lazily (see Finish below). @@ -372,9 +441,15 @@ impl StreamTextResult { finish_reason = fr; usage = u.clone(); finish_provider_metadata = pm; + saw_finish = true; break; } - StreamPart::Error { error } => return Err(error), + // Keep consuming after a provider error event: a trailing + // Finish still carries usage for the recording layer. The + // error is returned once the stream reaches its end. + StreamPart::Error { error } if provider_error.is_none() => { + provider_error = Some(error); + } StreamPart::ResponseMetadata { id, timestamp, @@ -390,6 +465,24 @@ impl StreamTextResult { } } + if let Some(error) = provider_error { + return Err(error); + } + if !saw_finish { + // AI SDK semantics: an incomplete stream with no output at all is + // an error; an incomplete stream with partial output retains the + // partial result but must not claim a normal stop. + if !saw_output { + return Err(AiMuxError::InvalidResponseData( + "No output generated. The model stream ended without a finish chunk.".into(), + )); + } + finish_reason = FinishReason { + unified: crate::types::FinishReasonUnified::Other, + raw: None, + }; + } + // Build response_content_parts in provider order: // reasoning (added during loop) → text → tool_calls. if !text.is_empty() { @@ -479,6 +572,14 @@ pub async fn generate_text( // 2. Build CallOptions. let mut call_options = options.into_call_options(lm_prompt); + let operation_timeout = + timeout::OperationTimeout::new(call_options.timeout.unwrap_or_default())?; + let abort_signal = call_options.abort_signal.clone(); + let retries = retry::prepare_retries( + call_options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); // 2a. RFC-0023: 关闭时零成本(M2 评审)——仅在录制开启时生成 call_id // 并绑定 recorder 快照;传输封闭由层 A 收尾声明(P1 无层 B,barrier @@ -486,10 +587,7 @@ pub async fn generate_text( let context = crate::recording::recorder().map(|recorder| { let call_id = crate::recording::new_call_id(); call_options.call_id = Some(call_id.clone()); - let ctx = crate::recording::RecordingContext { - call_id: call_id.clone(), - recorder, - }; + let ctx = crate::recording::RecordingContext::new(call_id.clone(), recorder); call_options.recording_context = Some(ctx.clone()); ctx.recorder.record_input( &ctx.call_id, @@ -539,7 +637,13 @@ pub async fn generate_text( model = %model.model_id(), modality = "text", ); - let result = match do_generate_with_logging(model, &call_options, span, started).await { + let operation = retries.retry(|| { + if let Some(context) = &context { + let _ = context.start_attempt(); + } + do_generate_with_logging(model, &call_options, span.clone(), started) + }); + let result = match timeout::run(operation, abort_signal.as_ref(), operation_timeout).await { Ok(r) => r, Err(e) => { if let (Some(rec), Some(call_id)) = (&recorder, &call_id) { @@ -796,9 +900,10 @@ fn extract_reasoning_signature(provider_metadata: Option<&Value>) -> Option, @@ -810,6 +915,22 @@ pub async fn stream_text( // 2. Build CallOptions. let mut call_options = options.into_call_options(lm_prompt); + let stream_timeout = call_options.timeout.unwrap_or_default(); + let operation_timeout = timeout::OperationTimeout::new(stream_timeout)?; + // Armed at operation start: providers await their first SSE event inside + // do_stream, so the first-chunk budget must already be counting there — + // a 200-then-silence server is otherwise unbounded when only + // first_chunk_ms is configured. + let first_chunk_deadline = stream_timeout + .first_chunk_ms + .map(|duration_ms| timeout::TimeoutDeadline::from_now("First chunk", duration_ms)) + .transpose()?; + let abort_signal = call_options.abort_signal.clone(); + let retries = retry::prepare_retries( + call_options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); // 2a. RFC-0023: 关闭时零成本(M2 评审)——仅在录制开启时生成 call_id // 并绑定 recorder 快照;传输封闭由层 A 收尾声明(P1 无层 B,barrier @@ -817,10 +938,7 @@ pub async fn stream_text( let context = crate::recording::recorder().map(|recorder| { let call_id = crate::recording::new_call_id(); call_options.call_id = Some(call_id.clone()); - let ctx = crate::recording::RecordingContext { - call_id: call_id.clone(), - recorder, - }; + let ctx = crate::recording::RecordingContext::new(call_id.clone(), recorder); call_options.recording_context = Some(ctx.clone()); ctx.recorder.record_input( &ctx.call_id, @@ -868,22 +986,104 @@ pub async fn stream_text( model = %model.model_id(), modality = "text", ); - let result: StreamResult = match model.do_stream(&call_options).instrument(span).await { - Ok(r) => r, - Err(e) => { - if let (Some(rec), Some(call_id)) = (&recorder, &call_id) { - rec.record_outcome(call_id, &crate::recording::OutcomeRecord::from_error(&e)); - } - return Err(e); + let operation = retries.retry(|| { + if let Some(context) = &context { + let _ = context.start_attempt(); } - }; + // Call-level retry boundary: once do_stream returns a stream, retry + // no longer recreates the provider request; the returned stream is + // consumed by the stream-step timeout logic below. + model.do_stream(&call_options).instrument(span.clone()) + }); + let setup_deadline = timeout::min_deadline(operation_timeout.deadline(), first_chunk_deadline); + let result: StreamResult = + match timeout::run_until(operation, abort_signal.as_ref(), setup_deadline).await { + Ok(r) => r, + Err(e) => { + if let (Some(rec), Some(call_id)) = (&recorder, &call_id) { + rec.record_outcome(call_id, &crate::recording::OutcomeRecord::from_error(&e)); + } + return Err(e); + } + }; // 解构避免部分 move(result 各字段去向不同)。 let StreamResult { stream, request_body, response_headers, } = result; - // 录制开启时才包装(终结时写 outcome + 传输封闭);关闭时零成本透传。 + + // The consumption phase keeps observing the same deadlines: the + // first-chunk budget armed at operation start continues until the first + // output chunk. + let operation_deadline = operation_timeout.deadline(); + // Always wrapped, even with no abort signal or deadlines armed: this + // loop is also the stream's terminal fuse — `Finish` and non-recoverable + // errors must end the stream under every configuration, or a provider + // that keeps the connection open after `Finish` hangs a consumer + // reading to end-of-stream. + let stream: Pin> + Send>> = { + let mut stream = stream; + let abort_signal = abort_signal; + let chunk_ms = stream_timeout.chunk_ms; + Box::pin(async_stream::stream! { + let mut chunk_deadline = first_chunk_deadline; + + loop { + let next = tokio::select! { + biased; + () = timeout::wait_for_abort(abort_signal.as_ref()) => { + yield Err(AiMuxError::from_abort_signal( + abort_signal.as_ref().expect("abort future only resolves for a signal"), + )); + break; + } + () = timeout::wait_for_deadline(operation_deadline) => { + yield Err(operation_deadline + .expect("operation deadline future only resolves for a deadline") + .error()); + break; + } + () = timeout::wait_for_deadline(chunk_deadline) => { + yield Err(chunk_deadline + .expect("chunk deadline future only resolves for a deadline") + .error()); + break; + } + item = stream.next() => item, + }; + + let Some(item) = next else { break }; + // A malformed individual frame does not prove that the + // transport is dead: SSE framing lets a later event still be + // valid. Transport/Core errors remain terminal so we do not + // keep polling a failed source indefinitely. + let terminal = matches!(&item, Ok(StreamPart::Finish { .. })) + || matches!(&item, Err(error) if !error.is_recoverable_stream_error()); + if let Ok(part) = &item + && is_output_chunk(part) + { + chunk_deadline = match chunk_ms + .map(|duration_ms| { + timeout::TimeoutDeadline::from_now("Chunk", duration_ms) + }) + .transpose() + { + Ok(deadline) => deadline, + Err(error) => { + yield Err(error); + break; + } + }; + } + + yield item; + if terminal { + break; + } + } + }) + }; let stream = crate::recording::RecordingOutcomeStream::new( stream, recorder.clone(), @@ -1036,3 +1236,451 @@ pub async fn stream_text_as_openai( stream_options, )) } + +#[cfg(test)] +mod operation_retry_tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use async_trait::async_trait; + + use super::*; + use crate::{ApiCallError, LanguageModel, prelude::StreamPart}; + + enum StreamBehavior { + Normal, + FinishThenPending, + ProviderErrorThenFinish, + ParseErrorThenEvent, + TransportErrorThenPending, + NeverReturns, + RetryableFirstError, + NonRetryableFirstError, + } + + struct StreamModel { + behavior: StreamBehavior, + calls: AtomicUsize, + } + + impl StreamModel { + fn new(behavior: StreamBehavior) -> Self { + Self { + behavior, + calls: AtomicUsize::new(0), + } + } + + fn api_error(status: u16, retryable: bool) -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_headers: Some(HashMap::from([("retry-after-ms".into(), "0".into())])), + is_retryable: retryable, + ..ApiCallError::new( + "stream setup failed", + "https://example.test/stream", + serde_json::json!({}), + ) + })) + } + + fn single_event_stream() -> StreamResult { + StreamResult { + stream: Box::pin(futures::stream::iter([Ok(StreamPart::TextDelta { + id: "text-1".into(), + delta: "hello".into(), + provider_metadata: None, + })])), + request_body: None, + response_headers: None, + } + } + + fn finish_stream() -> StreamResult { + let finish = StreamPart::Finish { + finish_reason: crate::types::FinishReason { + unified: crate::types::FinishReasonUnified::Stop, + raw: Some("stop".into()), + }, + usage: Usage::default(), + provider_metadata: None, + }; + StreamResult { + stream: Box::pin( + futures::stream::iter([Ok(finish)]).chain(futures::stream::pending()), + ), + request_body: None, + response_headers: None, + } + } + + fn provider_error_then_finish_stream() -> StreamResult { + let finish = StreamPart::Finish { + finish_reason: crate::types::FinishReason { + unified: crate::types::FinishReasonUnified::Error, + raw: None, + }, + usage: Usage::default(), + provider_metadata: None, + }; + StreamResult { + stream: Box::pin( + futures::stream::iter([ + Ok(StreamPart::Error { + error: AiMuxError::Other("provider error".into()), + }), + Ok(finish), + ]) + .chain(futures::stream::pending()), + ), + request_body: None, + response_headers: None, + } + } + + fn parse_error_then_event_stream() -> StreamResult { + StreamResult { + stream: Box::pin(futures::stream::iter([ + Err(AiMuxError::JsonParse("malformed SSE data".into())), + Ok(StreamPart::TextDelta { + id: "text-1".into(), + delta: "after-error".into(), + provider_metadata: None, + }), + ])), + request_body: None, + response_headers: None, + } + } + } + + #[async_trait] + impl LanguageModel for StreamModel { + fn provider(&self) -> &str { + "test" + } + + fn model_id(&self) -> &str { + "stream-model" + } + + async fn do_generate(&self, _options: &CallOptions) -> Result { + Err(AiMuxError::Other("unused".into())) + } + + async fn do_stream(&self, _options: &CallOptions) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst); + match self.behavior { + StreamBehavior::Normal => Ok(Self::single_event_stream()), + StreamBehavior::FinishThenPending => Ok(Self::finish_stream()), + StreamBehavior::ProviderErrorThenFinish => { + Ok(Self::provider_error_then_finish_stream()) + } + StreamBehavior::ParseErrorThenEvent => Ok(Self::parse_error_then_event_stream()), + StreamBehavior::TransportErrorThenPending => Ok(StreamResult { + stream: Box::pin( + futures::stream::iter([Err(AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..ApiCallError::new( + "connection reset", + "https://example.test/stream", + serde_json::json!({}), + ) + })))]) + .chain(futures::stream::pending()), + ), + request_body: None, + response_headers: None, + }), + StreamBehavior::NeverReturns => std::future::pending().await, + StreamBehavior::RetryableFirstError if attempt == 0 => { + Err(Self::api_error(429, true)) + } + StreamBehavior::RetryableFirstError => Ok(Self::single_event_stream()), + StreamBehavior::NonRetryableFirstError => Err(Self::api_error(400, false)), + } + } + } + + fn retry_options() -> GenerateTextOptions { + GenerateTextOptions { + max_retries: Some(2), + ..Default::default() + } + } + + #[tokio::test] + async fn stream_setup_preserves_the_peeked_first_event() { + let model = StreamModel::new(StreamBehavior::Normal); + let mut result = stream_text(&model, "hello", retry_options()).await.unwrap(); + let first = result.stream.next().await.unwrap().unwrap(); + assert!(matches!(first, StreamPart::TextDelta { delta, .. } if delta == "hello")); + assert!(result.stream.next().await.is_none()); + } + + #[tokio::test] + async fn first_chunk_timeout_starts_before_the_returned_stream_is_polled() { + let model = StreamModel::new(StreamBehavior::Normal); + let options = GenerateTextOptions { + max_retries: Some(0), + timeout: Some(crate::options::TimeoutConfiguration { + first_chunk_ms: Some(1), + ..Default::default() + }), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + + assert!(matches!( + result.stream.next().await.unwrap(), + Err(AiMuxError::Timeout(message)) if message == "First chunk timeout of 1ms exceeded" + )); + } + + #[tokio::test] + async fn caller_abort_still_cancels_the_returned_stream() { + let model = StreamModel::new(StreamBehavior::Normal); + let abort_signal = AbortSignal::new(); + let options = GenerateTextOptions { + max_retries: Some(0), + abort_signal: Some(abort_signal.clone()), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + abort_signal.abort(); + + assert!(matches!( + result.stream.next().await.unwrap(), + Err(AiMuxError::Aborted(message)) if message == "request aborted" + )); + } + + #[tokio::test] + async fn terminal_stream_part_stops_stream_consumption() { + let model = StreamModel::new(StreamBehavior::FinishThenPending); + let options = GenerateTextOptions { + max_retries: Some(0), + timeout: Some(crate::options::TimeoutConfiguration { + total_ms: Some(50), + ..Default::default() + }), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::Finish { .. } + )); + tokio::time::sleep(std::time::Duration::from_millis(60)).await; + assert!(result.stream.next().await.is_none()); + } + + #[tokio::test] + async fn provider_error_part_does_not_hide_the_final_finish() { + let model = StreamModel::new(StreamBehavior::ProviderErrorThenFinish); + let abort_signal = AbortSignal::new(); + let options = GenerateTextOptions { + max_retries: Some(0), + abort_signal: Some(abort_signal.clone()), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::Error { .. } + )); + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::Finish { finish_reason, .. } + if finish_reason.unified == crate::types::FinishReasonUnified::Error + )); + abort_signal.abort(); + assert!(result.stream.next().await.is_none()); + } + + #[tokio::test] + async fn timeout_wrapper_continues_after_a_malformed_stream_frame() { + let model = StreamModel::new(StreamBehavior::ParseErrorThenEvent); + let options = GenerateTextOptions { + max_retries: Some(0), + timeout: Some(crate::options::TimeoutConfiguration { + total_ms: Some(1_000), + ..Default::default() + }), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + assert!(matches!( + result.stream.next().await.unwrap(), + Err(AiMuxError::JsonParse(_)) + )); + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::TextDelta { delta, .. } if delta == "after-error" + )); + } + + #[tokio::test] + async fn abort_wrapper_continues_after_a_malformed_stream_frame() { + let model = StreamModel::new(StreamBehavior::ParseErrorThenEvent); + let options = GenerateTextOptions { + max_retries: Some(0), + abort_signal: Some(AbortSignal::new()), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + assert!(matches!( + result.stream.next().await.unwrap(), + Err(AiMuxError::JsonParse(_)) + )); + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::TextDelta { delta, .. } if delta == "after-error" + )); + } + + #[tokio::test] + async fn first_chunk_timeout_bounds_the_stream_setup_phase() { + // Providers await their first SSE event inside do_stream; a + // 200-then-silence server must be bounded by first_chunk_ms alone. + let model = StreamModel::new(StreamBehavior::NeverReturns); + let options = GenerateTextOptions { + max_retries: Some(0), + timeout: Some(crate::options::TimeoutConfiguration { + first_chunk_ms: Some(5), + ..Default::default() + }), + ..Default::default() + }; + let error = stream_text(&model, "hello", options).await.unwrap_err(); + assert!( + matches!(error, AiMuxError::Timeout(ref message) if message == "First chunk timeout of 5ms exceeded"), + "got {error:?}" + ); + } + + #[tokio::test] + async fn transport_error_ends_the_wrapped_stream() { + let model = StreamModel::new(StreamBehavior::TransportErrorThenPending); + let abort_signal = AbortSignal::new(); + let options = GenerateTextOptions { + max_retries: Some(0), + abort_signal: Some(abort_signal.clone()), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + + assert!(matches!( + result.stream.next().await.unwrap(), + Err(AiMuxError::ApiCall(_)) + )); + // Without treating Err as terminal this next() would hang on the + // pending tail of the dead source stream. + assert!(result.stream.next().await.is_none()); + } + + #[tokio::test] + async fn retryable_first_stream_error_retries_before_returning_stream() { + let model = StreamModel::new(StreamBehavior::RetryableFirstError); + let mut result = stream_text(&model, "hello", retry_options()).await.unwrap(); + assert_eq!(model.calls.load(Ordering::SeqCst), 2); + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::TextDelta { delta, .. } if delta == "hello" + )); + } + + #[tokio::test] + async fn non_retryable_first_stream_error_is_returned_unchanged() { + let model = StreamModel::new(StreamBehavior::NonRetryableFirstError); + let error = stream_text(&model, "hello", retry_options()) + .await + .unwrap_err(); + assert!( + matches!(error, AiMuxError::ApiCall(ref detail) if detail.status_code == Some(400)) + ); + assert_eq!(model.calls.load(Ordering::SeqCst), 1); + } +} + +#[cfg(test)] +mod stream_aggregation_tests { + use super::*; + use crate::types::FinishReasonUnified; + + fn result_from(parts: Vec>) -> StreamTextResult { + StreamTextResult { + stream: Box::pin(futures::stream::iter(parts)), + request_body: None, + response_headers: None, + } + } + + fn delta(text: &str) -> StreamPart { + StreamPart::TextDelta { + id: "text-1".into(), + delta: text.into(), + provider_metadata: None, + } + } + + fn finish() -> StreamPart { + StreamPart::Finish { + finish_reason: FinishReason { + unified: FinishReasonUnified::Stop, + raw: Some("stop".into()), + }, + usage: Usage::default(), + provider_metadata: None, + } + } + + #[tokio::test] + async fn truncated_stream_with_partial_output_does_not_claim_a_normal_stop() { + let result = result_from(vec![Ok(delta("par"))]).consume().await.unwrap(); + assert_eq!(result.text, "par"); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Other); + + // text() retains the partial result (AI SDK semantics). + let text = result_from(vec![Ok(delta("par"))]).text().await.unwrap(); + assert_eq!(text, "par"); + } + + #[tokio::test] + async fn empty_incomplete_stream_is_an_error() { + let error = result_from(vec![]).consume().await.unwrap_err(); + assert!( + matches!(&error, AiMuxError::InvalidResponseData(message) if message.contains("without a finish chunk")) + ); + let error = result_from(vec![]).text().await.unwrap_err(); + assert!(matches!(error, AiMuxError::InvalidResponseData(_))); + } + + #[tokio::test] + async fn stream_with_finish_still_reports_the_provider_finish_reason() { + let result = result_from(vec![Ok(delta("hi")), Ok(finish())]) + .consume() + .await + .unwrap(); + assert_eq!(result.finish_reason.unified, FinishReasonUnified::Stop); + } + + #[tokio::test] + async fn provider_error_still_wins_after_draining_a_trailing_finish() { + let parts = vec![ + Ok(delta("partial")), + Ok(StreamPart::Error { + error: AiMuxError::Other("provider error".into()), + }), + Ok(finish()), + ]; + let error = result_from(parts).consume().await.unwrap_err(); + assert!(matches!(error, AiMuxError::Other(message) if message == "provider error")); + } +} diff --git a/aimux-core/src/image_model.rs b/aimux-core/src/image_model.rs index faccf1cd..2a1ae941 100644 --- a/aimux-core/src/image_model.rs +++ b/aimux-core/src/image_model.rs @@ -1,4 +1,4 @@ -//! The `ImageModel` trait — the provider-facing interface for image generation. +//! The `ImageModel` trait — the provider-facing interface for image generation. //! //! Aligned with Vercel AI SDK `ImageModelV4` //! (`reference/ai/packages/provider/src/image-model/v4/`). @@ -9,9 +9,9 @@ use ts_rs::TS; use crate::error::AiMuxError; use crate::shared::{ - AbortSignal, AspectRatio, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Size, - Warning, + AspectRatio, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Size, Warning, }; +use crate::{AbortSignal, retry, timeout}; /// An image file used for image editing or variation generation. /// @@ -112,6 +112,12 @@ pub struct ImageCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional HTTP headers to send with the request. pub headers: Option, } @@ -129,6 +135,8 @@ impl ImageCallOptions { mask: None, provider_options: SharedProviderOptions::new(), abort_signal: None, + max_retries: None, + timeout: None, headers: None, } } @@ -184,6 +192,10 @@ pub trait ImageModel: Send + Sync { /// Provider-specific model ID, e.g. `"dall-e-3"`. fn model_id(&self) -> &str; + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Limit of how many images can be generated in a single API call. /// /// `None` means there is no fixed limit. The TS spec allows this to be a @@ -196,3 +208,27 @@ pub trait ImageModel: Send + Sync { /// Naming: the `do_` prefix prevents accidental direct usage by users. async fn do_generate(&self, options: &ImageCallOptions) -> Result; } + +/// User-facing image generation with Core-owned retry and timeout. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +pub async fn generate_image( + model: &dyn ImageModel, + options: ImageCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + timeout::run( + retries.retry(|| model.do_generate(&options)), + abort_signal.as_ref(), + timeout, + ) + .await +} diff --git a/aimux-core/src/language_model.rs b/aimux-core/src/language_model.rs index 55315e84..c6a37f50 100644 --- a/aimux-core/src/language_model.rs +++ b/aimux-core/src/language_model.rs @@ -34,6 +34,12 @@ pub trait LanguageModel: Send + Sync { /// Model identifier, e.g. `"gpt-4o"`. fn model_id(&self) -> &str; + /// Provider/model retry settings used when the caller does not override + /// the retry count. + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Generate a complete (non-streaming) response. async fn do_generate(&self, options: &CallOptions) -> Result; diff --git a/aimux-core/src/lib.rs b/aimux-core/src/lib.rs index 27d6943b..cbd5fc80 100644 --- a/aimux-core/src/lib.rs +++ b/aimux-core/src/lib.rs @@ -9,12 +9,15 @@ //! ────────────── ───────────────────── //! generate_text() ──► LanguageModel::do_generate() //! stream_text() ──► LanguageModel::do_stream() +//! embed() ──► EmbeddingModel::do_embed() +//! generate_image() ──► ImageModel::do_generate() +//! ... ──► ... //! ``` //! -//! The user calls `generate_text` / `stream_text` (free functions). -//! These convert user messages into `LanguageModelPrompt`, build `CallOptions`, -//! and call `LanguageModel::do_generate` / `do_stream` on the provider implementation. +//! Users call the free operation functions; providers implement the `do_*` +//! methods that perform one provider attempt. +mod abort_signal; pub mod composite; pub mod content; pub mod embedding_model; @@ -37,12 +40,14 @@ pub mod recording; pub mod replay; pub mod reranking_model; pub mod result; +pub mod retry; pub mod router; pub mod search_model; pub mod session; pub mod shared; pub mod speech_model; pub mod stream_part; +mod timeout; pub mod tool; pub mod trace; pub mod transcription_model; @@ -52,9 +57,12 @@ pub mod video_model; /// Convenience re-exports. pub mod prelude { + pub use crate::abort_signal::AbortSignal; pub use crate::composite::ChildModel; pub use crate::content::ContentPart; - pub use crate::embedding_model::{EmbeddingCallOptions, EmbeddingModel, EmbeddingResult}; + pub use crate::embedding_model::{ + EmbeddingCallOptions, EmbeddingModel, EmbeddingResult, embed, + }; pub use crate::error::{AiMuxError, ApiCallError}; pub use crate::files_model::{Files, UploadFileCallOptions, UploadFileResult}; pub use crate::generate::{ @@ -62,7 +70,7 @@ pub mod prelude { generate_object, generate_text, generate_text_as_openai, stream_text, stream_text_as_openai, }; - pub use crate::image_model::{ImageCallOptions, ImageModel, ImageResult}; + pub use crate::image_model::{ImageCallOptions, ImageModel, ImageResult, generate_image}; pub use crate::language_model::LanguageModel; pub use crate::language_model_message::LanguageModelPrompt; pub use crate::message::{MessageContent, ModelMessage, ModelPrompt, Role}; @@ -78,35 +86,43 @@ pub mod prelude { }; pub use crate::options::{CallOptions, ResponseFormat, ToolChoice}; pub use crate::provider::Provider; - pub use crate::reranking_model::{RerankingCallOptions, RerankingModel, RerankingResult}; + pub use crate::reranking_model::{ + RerankingCallOptions, RerankingModel, RerankingResult, rerank, + }; pub use crate::result::{GenerateResult, StreamResult}; pub use crate::router::{ FallbackPolicy, Router, RouterConfig, RouterModel, RuleRouter, WeightedRouter, }; pub use crate::search_model::{ - SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, + SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, search, }; pub use crate::session::{ SessionCall, SessionInferer, SessionSource, SessionStore, SessionView, init_session_infer, init_session_store, list_sessions, session_calls, }; pub use crate::shared::{ - AbortSignal, AspectRatio, FileBytes, FileData, SharedHeaders, SharedProviderMetadata, + AspectRatio, FileBytes, FileData, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, SharedProviderReference, Size, }; - pub use crate::speech_model::{SpeechCallOptions, SpeechModel, SpeechResult}; + pub use crate::speech_model::{SpeechCallOptions, SpeechModel, SpeechResult, generate_speech}; pub use crate::stream_part::StreamPart; pub use crate::tool::{FunctionTool, ProviderTool, Tool, ToolCall, ToolResult}; pub use crate::transcription_model::{ - TranscriptionCallOptions, TranscriptionModel, TranscriptionResult, + AudioChunk, InputAudioFormat, TranscriptionCallOptions, TranscriptionModel, + TranscriptionResult, TranscriptionStreamOptions, TranscriptionStreamPart, + TranscriptionStreamResult, stream_transcribe, transcribe, }; pub use crate::types::{FinishReason, FinishReasonUnified, ReasoningEffort, Usage, Warning}; - pub use crate::video_model::{VideoCallOptions, VideoModel, VideoResult}; + pub use crate::video_model::{ + VideoCallOptions, VideoModel, VideoOperationStart, VideoOperationStatus, VideoPollConfig, + VideoPollOptions, VideoResult, generate_video, + }; } // Root-level re-exports for convenience. +pub use abort_signal::AbortSignal; pub use embedding_model::EmbeddingModel; -pub use error::{AiMuxError, ApiCallError}; +pub use error::{AiMuxError, ApiCallError, RetryError, RetryErrorReason}; pub use files_model::Files; pub use image_model::ImageModel; pub use language_model::LanguageModel; diff --git a/aimux-core/src/moa.rs b/aimux-core/src/moa.rs index 2248605b..5813f0ef 100644 --- a/aimux-core/src/moa.rs +++ b/aimux-core/src/moa.rs @@ -18,6 +18,7 @@ use crate::error::AiMuxError; use crate::language_model::LanguageModel; use crate::options::{CallOptions, ToolChoice}; use crate::result::{GenerateResult, StreamResult}; +use crate::retry; use crate::stream_part::StreamPart; use crate::types::{Usage, Warning}; @@ -120,11 +121,25 @@ impl MoaModel { return Ok((Vec::new(), Usage::default(), Vec::new())); } let ref_opts = self.reference_options(options); - let results = join_all( - self.references - .iter() - .map(|m| async { m.do_generate(&ref_opts).await }), - ) + let results = join_all(self.references.iter().enumerate().map(|(i, m)| { + let child_opts = + ref_opts.for_step(format!("moa.ref[{i}]:{}/{}", m.provider(), m.model_id())); + async move { + let retries = retry::prepare_retries( + child_opts.max_retries, + m.retry_config(), + child_opts.abort_signal.clone(), + ); + retries + .retry(|| { + if let Some(context) = &child_opts.recording_context { + let _ = context.start_attempt(); + } + m.do_generate(&child_opts) + }) + .await + } + })) .await; let mut usage = Usage::default(); @@ -168,6 +183,15 @@ impl LanguageModel for MoaModel { &self.config.model_id } + fn retry_config(&self) -> retry::RetryConfig { + // Retrying the composite would rerun the entire reference fanout; + // references and aggregator are retried independently instead. + retry::RetryConfig { + max_retries: 0, + ..retry::RetryConfig::default() + } + } + async fn do_generate(&self, options: &CallOptions) -> Result { // 1. Fan out references (non-streaming). let (texts, ref_usage, warnings) = self.run_references_nonstream(options).await?; @@ -182,7 +206,24 @@ impl LanguageModel for MoaModel { agg_opts.prompt = agg_prompt; // 3. Run the aggregator. - let mut agg = self.aggregator.do_generate(&agg_opts).await?; + let retries = retry::prepare_retries( + agg_opts.max_retries, + self.aggregator.retry_config(), + agg_opts.abort_signal.clone(), + ); + let agg_opts = agg_opts.for_step(format!( + "moa.aggregator:{}/{}", + self.aggregator.provider(), + self.aggregator.model_id() + )); + let mut agg = retries + .retry(|| { + if let Some(context) = &agg_opts.recording_context { + let _ = context.start_attempt(); + } + self.aggregator.do_generate(&agg_opts) + }) + .await?; // 4. Fold reference usage + drop warnings into the aggregator result. agg.usage = add_usage(agg.usage, &ref_usage); @@ -207,7 +248,24 @@ impl LanguageModel for MoaModel { // 3. Aggregator streams; we emit our own StreamStart and add reference // usage onto its Finish. We swallow the aggregator's StreamStart // (we've already emitted ours). - let agg = self.aggregator.do_stream(&agg_opts).await?; + let retries = retry::prepare_retries( + agg_opts.max_retries, + self.aggregator.retry_config(), + agg_opts.abort_signal.clone(), + ); + let agg_opts = agg_opts.for_step(format!( + "moa.aggregator:{}/{}", + self.aggregator.provider(), + self.aggregator.model_id() + )); + let agg = retries + .retry(|| { + if let Some(context) = &agg_opts.recording_context { + let _ = context.start_attempt(); + } + self.aggregator.do_stream(&agg_opts) + }) + .await?; let mut agg_stream = agg.stream; let stream = async_stream::stream! { @@ -223,10 +281,16 @@ impl LanguageModel for MoaModel { }); } Ok(other) => yield Ok(other), - // Stream errors are terminal — relay and stop (the user's - // original request was NOT the aggregator body MoA - // synthesized, so relaying more would be noise). - Err(e) => { yield Err(e); break; } + Err(e) => { + let recoverable = e.is_recoverable_stream_error(); + yield Err(e); + if !recoverable { + // A transport/Core failure is terminal. Retrying or + // replaying the synthesized aggregator request after + // output has escaped would duplicate visible data. + break; + } + } } } }; @@ -255,6 +319,7 @@ mod tests { use async_trait::async_trait; use futures::stream; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; /// A mock child that returns fixed text with fixed usage. `fail` forces an /// error. Used as both reference and aggregator. @@ -263,6 +328,8 @@ mod tests { text: String, fail: bool, usage: Usage, + retry_failures: usize, + calls: Arc, } #[async_trait] @@ -274,9 +341,13 @@ mod tests { self.name } async fn do_generate(&self, _options: &CallOptions) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; if self.fail { return Err(AiMuxError::Other(format!("{} failed", self.name))); } + if attempt <= self.retry_failures { + return Err(retryable_error(self.name)); + } Ok(GenerateResult { content: vec![GenerateContent::Text { text: self.text.clone(), @@ -298,9 +369,13 @@ mod tests { }) } async fn do_stream(&self, _options: &CallOptions) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; if self.fail { return Err(AiMuxError::Other(format!("{} failed", self.name))); } + if attempt <= self.retry_failures { + return Err(retryable_error(self.name)); + } let parts: Vec> = vec![ Ok(StreamPart::StreamStart { warnings: vec![] }), Ok(StreamPart::TextDelta { @@ -325,6 +400,41 @@ mod tests { } } + fn retryable_error(name: &str) -> AiMuxError { + AiMuxError::ApiCall(Box::new(crate::ApiCallError { + status_code: Some(503), + response_headers: Some(std::collections::HashMap::from([( + "retry-after-ms".into(), + "0".into(), + )])), + is_retryable: true, + ..crate::ApiCallError::new( + format!("{name} retryable failure"), + "https://example.test", + serde_json::json!({}), + ) + })) + } + + fn retry_child( + name: &'static str, + text: &str, + retry_failures: usize, + ) -> (ChildModel, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + ( + Arc::new(MockChild { + name, + text: text.into(), + fail: false, + usage: Usage::default(), + retry_failures, + calls: calls.clone(), + }), + calls, + ) + } + fn mk(name: &'static str, text: &str, fail: bool, total: u32) -> ChildModel { Arc::new(MockChild { name, @@ -341,6 +451,8 @@ mod tests { }, raw: None, }, + retry_failures: 0, + calls: Arc::new(AtomicUsize::new(0)), }) } @@ -448,6 +560,52 @@ mod tests { assert_eq!(finish_usage.unwrap().input_tokens.total, Some(17)); } + #[tokio::test] + async fn retries_references_and_aggregator_independently_without_replaying_fanout() { + let (ref_a, ref_a_calls) = retry_child("ref-a", "A", 1); + let (ref_b, ref_b_calls) = retry_child("ref-b", "B", 0); + let (aggregator, aggregator_calls) = retry_child("aggregator", "final", 1); + let moa = MoaModel::new(vec![ref_a, ref_b], aggregator, MoaConfig::default()); + assert_eq!(moa.retry_config().max_retries, 0); + + let result = crate::generate::generate_text( + &moa, + "hello", + crate::generate::GenerateTextOptions { + max_retries: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(result.text, "final"); + assert_eq!(ref_a_calls.load(Ordering::SeqCst), 2); + assert_eq!(ref_b_calls.load(Ordering::SeqCst), 1); + assert_eq!(aggregator_calls.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn retries_aggregator_stream_setup_without_replaying_references() { + let (reference, reference_calls) = retry_child("reference", "A", 0); + let (aggregator, aggregator_stream_calls) = retry_child("aggregator", "final", 1); + let moa = MoaModel::new(vec![reference], aggregator, MoaConfig::default()); + + let _result = crate::generate::stream_text( + &moa, + "hello", + crate::generate::GenerateTextOptions { + max_retries: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(reference_calls.load(Ordering::SeqCst), 1); + assert_eq!(aggregator_stream_calls.load(Ordering::SeqCst), 2); + } + #[tokio::test] async fn stream_best_effort_emits_drop_warning_in_start() { let refs = vec![mk("ref-a", "A", false, 5), mk("ref-b", "B", true, 0)]; @@ -576,6 +734,73 @@ mod tests { ); } + struct RecoverableFrameErrorAggregator; + + #[async_trait] + impl LanguageModel for RecoverableFrameErrorAggregator { + fn provider(&self) -> &str { + "mock" + } + + fn model_id(&self) -> &str { + "agg-recoverable-error" + } + + async fn do_generate(&self, _options: &CallOptions) -> Result { + unimplemented!() + } + + async fn do_stream(&self, _options: &CallOptions) -> Result { + Ok(StreamResult { + stream: Box::pin(stream::iter([ + Ok(StreamPart::StreamStart { warnings: vec![] }), + Err(AiMuxError::JsonParse("malformed SSE data".into())), + Ok(StreamPart::TextDelta { + id: "t1".into(), + delta: "after-error".into(), + provider_metadata: None, + }), + Ok(StreamPart::Finish { + finish_reason: FinishReason { + unified: FinishReasonUnified::Stop, + raw: Some("stop".into()), + }, + usage: Usage::default(), + provider_metadata: None, + }), + ])), + request_body: None, + response_headers: None, + }) + } + } + + #[tokio::test] + async fn stream_continues_after_recoverable_aggregator_frame_error() { + let refs = vec![mk("ref-a", "A", false, 5)]; + let moa = MoaModel::new( + refs, + Arc::new(RecoverableFrameErrorAggregator), + MoaConfig::default(), + ); + let result = moa.do_stream(&opts_with_prompt()).await.unwrap(); + let outcomes: Vec<_> = result.stream.collect().await; + let error_index = outcomes + .iter() + .position(|outcome| matches!(outcome, Err(AiMuxError::JsonParse(_)))) + .expect("parse error should be relayed"); + + assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::TextDelta { delta, .. }) if delta == "after-error" + ))); + assert!(outcomes.iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::Finish { finish_reason, .. }) + if finish_reason.unified == FinishReasonUnified::Stop + ))); + } + /// Like `mk` but with full per-field usage (for G4 accumulation testing). fn mk_full( name: &'static str, @@ -605,6 +830,8 @@ mod tests { }, raw: None, }, + retry_failures: 0, + calls: Arc::new(AtomicUsize::new(0)), }) } diff --git a/aimux-core/src/openai_output.rs b/aimux-core/src/openai_output.rs index d7f72249..a2be4c3c 100644 --- a/aimux-core/src/openai_output.rs +++ b/aimux-core/src/openai_output.rs @@ -419,6 +419,10 @@ pub fn to_chat_completion_stream( while let Some(part_result) = stream.next().await { let part = match part_result { Ok(p) => p, + Err(e) if e.is_recoverable_stream_error() => { + yield Err(e); + continue; + } Err(e) => { // Emit error as a content delta + finish, then stop. if let Some(chunk) = state.error_chunk(&e) { @@ -1548,10 +1552,11 @@ mod tests { provider_metadata: None, }), Ok(StreamPart::Error { - error: AiMuxError::ApiCall(crate::error::ApiCallError { - message: "something went wrong".to_string(), - ..Default::default() - }), + error: AiMuxError::ApiCall(Box::new(crate::error::ApiCallError::new( + "something went wrong", + "https://example.test/v1", + serde_json::json!({}), + ))), }), ]; @@ -1586,6 +1591,51 @@ mod tests { assert_eq!(last.choices[0].finish_reason.as_deref(), Some("stop")); } + #[tokio::test] + async fn recoverable_frame_error_does_not_truncate_chat_completion_stream() { + let parts: Vec> = vec![ + Ok(StreamPart::StreamStart { warnings: vec![] }), + Err(AiMuxError::JsonParse("malformed SSE data".into())), + Ok(StreamPart::TextDelta { + id: "0".to_string(), + delta: "after-error".to_string(), + provider_metadata: None, + }), + Ok(StreamPart::Finish { + finish_reason: FinishReason { + unified: FinishReasonUnified::Stop, + raw: Some("stop".into()), + }, + usage: Usage::default(), + provider_metadata: None, + }), + ]; + + let result = to_chat_completion_stream( + Box::pin(futures::stream::iter(parts)), + "gpt-4o", + OpenAiStreamOptions::default(), + ); + let outcomes: Vec<_> = result.stream.collect().await; + let error_index = outcomes + .iter() + .position(|outcome| matches!(outcome, Err(AiMuxError::JsonParse(_)))) + .expect("parse error should be preserved as an error item"); + + assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( + outcome, + Ok(chunk) + if chunk.choices.first().and_then(|choice| choice.delta.content.as_deref()) + == Some("after-error") + ))); + assert!(outcomes.iter().any(|outcome| matches!( + outcome, + Ok(chunk) + if chunk.choices.first().and_then(|choice| choice.finish_reason.as_deref()) + == Some("stop") + ))); + } + #[tokio::test] async fn test_stream_completes_without_finish() { // Stream ends without Finish — should still emit a final chunk. diff --git a/aimux-core/src/options.rs b/aimux-core/src/options.rs index 4123b003..3d289139 100644 --- a/aimux-core/src/options.rs +++ b/aimux-core/src/options.rs @@ -8,8 +8,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use ts_rs::TS; +use crate::AbortSignal; use crate::language_model_message::LanguageModelPrompt; -use crate::shared::AbortSignal; pub use crate::tool::{FunctionTool, ProviderTool, Tool, ToolChoice}; use crate::types::ReasoningEffort; @@ -43,7 +43,15 @@ pub struct TimeoutConfiguration { // cannot reach 2^53 (~285k years), so precision is never at stake. #[ts(type = "number | null")] pub total_ms: Option, + /// Timeout for one generation step, including that step's attempts and + /// retry backoff, in milliseconds. Aimux currently has one step. + #[ts(type = "number | null")] + pub step_ms: Option, /// Timeout waiting for the first stream chunk (streaming only). + /// + /// Counted from operation start, so it also bounds stream establishment + /// and any retries before the first semantic output: it is the + /// user-perceived time-to-first-output budget, not a per-attempt timer. #[ts(type = "number | null")] pub first_chunk_ms: Option, /// Maximum idle time between consecutive stream chunks (streaming only). @@ -112,7 +120,7 @@ pub struct CallOptions { pub body_overrides: Option, /// Per-call retry count override. `None` uses the provider's configured - /// `RetryConfig.max_retries`. `Some(0)` disables retries. + /// Core operation retry. `Some(0)` disables retries. pub max_retries: Option, /// Per-call timeout configuration (total / first-chunk / chunk idle). @@ -159,6 +167,21 @@ pub struct CallOptions { } impl CallOptions { + /// Clone these options for one composite child step (Router child, MoA + /// reference or aggregator): the recording context, when present, is + /// replaced by a child context labeled `step`, so the child's exchanges + /// and retry attempts are recorded as their own step instead of blending + /// into the parent operation's. + #[must_use] + pub fn for_step(&self, step: impl Into) -> Self { + let mut child = self.clone(); + child.recording_context = self + .recording_context + .as_ref() + .map(|context| context.child(step.into())); + child + } + /// Create a `CallOptions` with the given prompt and all other fields set /// to their defaults (`None` / `ToolChoice::Auto`). /// diff --git a/aimux-core/src/recording.rs b/aimux-core/src/recording.rs index 0490ec01..dcd3965b 100644 --- a/aimux-core/src/recording.rs +++ b/aimux-core/src/recording.rs @@ -8,7 +8,7 @@ //! 区别于 HTTP 请求级 ID 与跨服务 trace)。 //! - **默认关闭**:不调 `init_recording`,热路径 = 1 读锁 + clone(次 ns 级)。 //! - **隐私受控**:api_key / Authorization / cookie 系恒脱敏(contains 式,含 `x-goog-api-key`); -//! token 不再 contains(曾误伤 `max_output_tokens` 等用量字段),仅精确脱敏 `x-amz-security-token`; +//! token 不再 contains(曾误伤 `max_output_tokens` 等用量字段),仅精确脱敏凭据字段; //! `InputRecord.options` 序列化前递归脱敏。 //! - **completion barrier**:outcome 与全部 exchange(流式含终结)齐才写行。 //! - **专用 writer thread + oneshot flush**:同步 `flush()` 阻塞至落盘,不依赖运行时。 @@ -20,7 +20,7 @@ use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet, VecDeque}; use std::io::{BufWriter, Write}; use std::path::PathBuf; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::sync::mpsc::{Receiver, SyncSender, sync_channel}; use std::sync::{Arc, Mutex, RwLock}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -29,7 +29,7 @@ use ts_rs::TS; // ── 数据模型(三层 + call_id 关联;schema 版本)──────────────────────────── /// 录制格式版本(用于未来字段迁移与绑定层兼容)。 -pub const RECORDING_SCHEMA: u32 = 1; +pub const RECORDING_SCHEMA: u32 = 2; /// 一次完整调用的录制记录(三层 + call_id 关联)。 #[derive(Debug, Clone, Serialize, Deserialize, TS)] @@ -168,8 +168,17 @@ impl ProviderRecord { #[derive(Debug, Clone, Serialize, Deserialize, TS)] #[ts(export)] pub struct HttpExchange { - /// 第几次重试(0=首次);per-attempt 递增。 + /// Composite step this exchange belongs to (e.g. `router[0]:openai/gpt-4o` + /// or `moa.ref[1]:...`). `None` for a plain, non-composite operation. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub step: Option, + /// Core operation attempt, starting at 1. Attempt numbers are unique + /// across the whole call, including composite child steps, so + /// `(attempt, exchange_index)` alone still identifies an exchange. pub attempt: u32, + /// HTTP exchange within the operation attempt, starting at 1. + #[serde(default)] + pub exchange_index: u32, pub request: HttpRecord, /// None = 请求失败未获响应。 pub response: Option, @@ -243,6 +252,10 @@ pub struct OutcomeRecord { pub status: OutcomeStatus, pub finish_reason: Option, pub error: Option, + /// Lossless structured domain error. In particular, `RetryError.errors` + /// keeps the complete attempt history rather than only its display text. + #[serde(default)] + pub error_value: Option, /// 序列化的 Usage。 pub usage: Option, } @@ -257,6 +270,7 @@ impl OutcomeRecord { .ok() .and_then(|v| v.as_str().map(std::string::ToString::to_string)), error: None, + error_value: None, usage: serde_json::to_value(&r.usage).ok(), } } @@ -268,6 +282,7 @@ impl OutcomeRecord { status: OutcomeStatus::Error, finish_reason: None, error: Some(e.to_string()), + error_value: serde_json::to_value(e).ok(), usage: None, } } @@ -332,6 +347,7 @@ pub trait Recorder: Send + Sync { &self, call_id: &str, attempt: u32, + exchange_index: u32, response: &ResponseRecord, error: Option, ); @@ -369,6 +385,65 @@ static RECORDER: RwLock>> = RwLock::new(None); pub struct RecordingContext { pub call_id: String, pub recorder: Arc, + attempt_allocator: Arc, + current_attempt: Arc, + exchange_index: Arc, + step: Option>, +} + +impl RecordingContext { + #[must_use] + pub fn new(call_id: impl Into, recorder: Arc) -> Self { + Self { + call_id: call_id.into(), + recorder, + attempt_allocator: Arc::new(AtomicU32::new(0)), + current_attempt: Arc::new(AtomicU32::new(0)), + exchange_index: Arc::new(AtomicU32::new(0)), + step: None, + } + } + + /// Derive a context for one composite child step (Router child, MoA + /// reference, or aggregator). The allocator stays shared with the parent + /// so attempt numbers are unique across the whole call, while the current + /// attempt and exchange counter are local to the child. An interleaved + /// sibling cannot relabel an exchange that belongs to an in-flight attempt. + #[must_use] + pub fn child(&self, step: impl Into) -> Self { + Self { + call_id: self.call_id.clone(), + recorder: self.recorder.clone(), + attempt_allocator: self.attempt_allocator.clone(), + current_attempt: Arc::new(AtomicU32::new(0)), + exchange_index: Arc::new(AtomicU32::new(0)), + step: Some(Arc::from(step.into())), + } + } + + /// The composite step label for exchanges recorded through this context. + #[must_use] + pub fn step(&self) -> Option { + self.step.as_deref().map(str::to_owned) + } + + /// Begin the next Core operation attempt and reset its exchange counter. + #[must_use] + pub fn start_attempt(&self) -> u32 { + let attempt = self.attempt_allocator.fetch_add(1, Ordering::AcqRel) + 1; + self.current_attempt.store(attempt, Ordering::Release); + self.exchange_index.store(0, Ordering::Release); + attempt + } + + /// Allocate the next exchange identity for the current attempt. + /// Direct provider-SPI calls that bypass Core are recorded as attempt 1. + #[must_use] + pub fn next_exchange(&self) -> (u32, u32) { + let attempt = self.current_attempt.load(Ordering::Acquire).max(1); + let exchange_index = self.exchange_index.fetch_add(1, Ordering::AcqRel) + 1; + (attempt, exchange_index) + } } impl std::fmt::Debug for RecordingContext { @@ -402,10 +477,7 @@ pub fn init_recording_from_env() -> bool { /// 从当前全局 recorder 生成一次调用的 context(关闭时 None)。 /// 层 A 入口:先读一次,再用 `context && ctx.start()` 记录 ①+②。 pub fn context(call_id: impl Into) -> Option { - recorder().map(|recorder| RecordingContext { - call_id: call_id.into(), - recorder, - }) + recorder().map(|recorder| RecordingContext::new(call_id, recorder)) } /// 热路径检查(关闭时 None,≈1 读锁 + 1 次 Arc clone)。 @@ -429,33 +501,36 @@ pub fn new_call_id() -> String { format!("call-{ns}-{}", CALL_SEQ.fetch_add(1, Ordering::Relaxed)) } -// ── 脱敏(contains 式;与 logging.rs 同规则)────────────────────────────── +// ── 脱敏（recording / logging / public error context 共用）───────────── /// 敏感键判断:受保护头/参数名(值将恒脱敏)。 /// -/// needle 集合与 `aimux_provider_utils::logging::is_sensitive_key` 对齐 -/// (`authorization`/`api-key`/`apikey`/`key`),录制侧额外覆盖 -/// `cookie`/`set-cookie`(logging 不脱敏 cookie)。其中 `key` 取 **exact** 匹配 +/// 这是 recording、logging 和 public error context 的唯一敏感键策略。 +/// 覆盖 `authorization`/`api-key`/`apikey`/`key`/`cookie`/`set-cookie`。 +/// 其中 `key` 取 **exact** 匹配 /// 而非 contains——避免误伤 `X-Key`/`monkey`/`keyboard`/`keyword` 等含 "key" /// 子串的非凭据名(既有 `redact_json` 测试即要求 `X-Key` 值保留)。 /// /// **token 不再用 contains**——教训:`contains("token")` 会误伤 /// `max_output_tokens`/`prompt_tokens`/`completion_tokens` 这类正常用量字段名 -/// (LLM 用量统计,非凭据),导致录制里这些值被无端替换成 `[REDACTED]`。改为只 -/// 精确匹配 AWS Bedrock sigv4 真实写入的凭据头 `x-amz-security-token`(name 已 -/// lowercase 归一化,直接 `==` 比较即可)。其余 needle 维持 contains,以覆盖 +/// (LLM 用量统计,非凭据),导致录制里这些值被无端替换成 `[REDACTED]`。改为 +/// 精确匹配常见 bearer-token 字段和 AWS Bedrock sigv4 的 +/// `x-amz-security-token`(name 已 lowercase 归一化)。其余 needle 维持 contains,以覆盖 /// `x-goog-api-key`/`proxy-authorization` 等变体。 #[must_use] pub fn is_sensitive_key(name: &str) -> bool { let n = name.to_ascii_lowercase(); + // Separator-insensitive form so `accessToken`, `access-token`, and + // `x-access-token` all classify like `access_token`. A `token` SUFFIX is + // the credential shape; plural usage fields (`max_tokens`, + // `input_tokens`) end in `tokens` and stay loggable. + let squashed: String = n.chars().filter(|c| *c != '-' && *c != '_').collect(); n == "cookie" || n == "set-cookie" || n == "key" - || n == "x-amz-security-token" + || squashed.ends_with("token") || n.contains("authorization") - || n.contains("api-key") - || n.contains("api_key") - || n.contains("apikey") + || squashed.contains("apikey") } /// 递归脱敏(JSON 中含敏感键的项值替换为 `[REDACTED]`)。 @@ -506,6 +581,7 @@ enum RecordEvent { ExchangeUpdate { call_id: String, attempt: u32, + exchange_index: u32, response: ResponseRecord, error: Option, }, @@ -725,12 +801,14 @@ impl Recorder for JsonlRecorder { &self, call_id: &str, attempt: u32, + exchange_index: u32, response: &ResponseRecord, error: Option, ) { self.send_ev(RecordEvent::ExchangeUpdate { call_id: call_id.to_string(), attempt, + exchange_index, response: response.clone(), error, }); @@ -833,6 +911,7 @@ fn writer_loop( RecordEvent::ExchangeUpdate { call_id, attempt, + exchange_index, response, error, } => { @@ -840,7 +919,7 @@ fn writer_loop( // 不静默 patch 第一条。 let rec = entry_or_init(&mut pending, &call_id); if !matches!( - apply_exchange_update(rec, attempt, response, error), + apply_exchange_update(rec, attempt, exchange_index, response, error), UpdateMatch::Patched ) { mark_inconsistent(&inconsistent, &call_id); @@ -961,7 +1040,7 @@ fn insert_exchange(rec: &mut Recording, exchange: HttpExchange) -> bool { if let Some(existing) = rec .exchanges .iter_mut() - .find(|e| e.attempt == exchange.attempt) + .find(|e| e.attempt == exchange.attempt && e.exchange_index == exchange.exchange_index) { merge_exchange(existing, &exchange); true @@ -991,6 +1070,7 @@ fn merge_exchange(existing: &mut HttpExchange, new: &HttpExchange) { fn apply_exchange_update( rec: &mut Recording, attempt: u32, + exchange_index: u32, response: ResponseRecord, error: Option, ) -> UpdateMatch { @@ -998,7 +1078,7 @@ fn apply_exchange_update( .exchanges .iter() .enumerate() - .filter(|(_, e)| e.attempt == attempt) + .filter(|(_, e)| e.attempt == attempt && e.exchange_index == exchange_index) .map(|(i, _)| i) .collect(); match matches.len() { @@ -1352,6 +1432,7 @@ impl Recorder for RingRecorder { &self, call_id: &str, attempt: u32, + exchange_index: u32, response: &ResponseRecord, error: Option, ) { @@ -1359,7 +1440,7 @@ impl Recorder for RingRecorder { // C4-7:要求恰好一个匹配;0 或 >1 → 标记 inconsistent,不静默 patch 第一条。 let rec = inner.entry_or_init_bounded(call_id); if !matches!( - apply_exchange_update(rec, attempt, response.clone(), error), + apply_exchange_update(rec, attempt, exchange_index, response.clone(), error), UpdateMatch::Patched ) { inner.mark_inconsistent(call_id); @@ -1422,6 +1503,9 @@ pub struct RecordingOutcomeStream { call_id: String, /// 已记录终结。 recorded: bool, + /// Provider error 事件先到时挂起的 outcome:尾随的 Finish 仍携带 usage + /// (计费证据),在终结时合并写入,而不是在 error 事件处立即定稿丢掉它。 + pending: Option, } impl RecordingOutcomeStream { @@ -1432,6 +1516,17 @@ impl RecordingOutcomeStream { recorder, call_id: call_id.into(), recorded: false, + pending: None, + } + } + + /// 写出挂起的 provider error outcome(若有);首个错误定性不变。 + fn flush_pending(&mut self) -> bool { + if let Some(pending) = self.pending.take() { + self.record(&pending); + true + } else { + false } } @@ -1447,12 +1542,11 @@ impl RecordingOutcomeStream { } } -impl futures::Stream for RecordingOutcomeStream +impl futures::Stream for RecordingOutcomeStream where - S: futures::Stream> + Unpin, - E: std::fmt::Display, + S: futures::Stream> + Unpin, { - type Item = Result; + type Item = Result; fn poll_next( mut self: std::pin::Pin<&mut Self>, @@ -1468,47 +1562,77 @@ where usage, .. } => { - let outcome = OutcomeRecord { - status: OutcomeStatus::Success, - finish_reason: serde_json::to_value(finish_reason.unified) - .ok() - .and_then(|v| v.as_str().map(std::string::ToString::to_string)), - error: None, - usage: serde_json::to_value(usage).ok(), - }; - self.as_mut().get_mut().record(&outcome); + let this = self.as_mut().get_mut(); + if let Some(mut pending) = this.pending.take() { + // Provider error 已定性;尾随 Finish 补上 usage。 + pending.usage = serde_json::to_value(usage).ok(); + this.record(&pending); + } else { + let outcome = OutcomeRecord { + status: OutcomeStatus::Success, + finish_reason: serde_json::to_value(finish_reason.unified) + .ok() + .and_then(|v| v.as_str().map(std::string::ToString::to_string)), + error: None, + error_value: None, + usage: serde_json::to_value(usage).ok(), + }; + this.record(&outcome); + } } crate::stream_part::StreamPart::Error { error } => { - let outcome = OutcomeRecord { - status: OutcomeStatus::Error, - finish_reason: None, - error: Some(error.to_string()), - usage: None, - }; - self.as_mut().get_mut().record(&outcome); + let this = self.as_mut().get_mut(); + if !this.recorded && this.pending.is_none() { + this.pending = Some(OutcomeRecord { + status: OutcomeStatus::Error, + finish_reason: None, + error: Some(error.to_string()), + error_value: serde_json::to_value(error).ok(), + usage: None, + }); + } } _ => {} } std::task::Poll::Ready(Some(Ok(part))) } std::task::Poll::Ready(Some(Err(e))) => { - let outcome = OutcomeRecord { - status: OutcomeStatus::Error, - finish_reason: None, - error: Some(e.to_string()), - usage: None, - }; - this.record(&outcome); + if e.is_recoverable_stream_error() { + if !this.recorded && this.pending.is_none() { + this.pending = Some(OutcomeRecord { + status: OutcomeStatus::Error, + finish_reason: None, + error: Some(e.to_string()), + error_value: serde_json::to_value(&e).ok(), + usage: None, + }); + } + return std::task::Poll::Ready(Some(Err(e))); + } + // 挂起的 provider error 先定性(首个错误胜出,与旧行为一致)。 + if !this.flush_pending() { + let outcome = OutcomeRecord { + status: OutcomeStatus::Error, + finish_reason: None, + error: Some(e.to_string()), + error_value: serde_json::to_value(&e).ok(), + usage: None, + }; + this.record(&outcome); + } std::task::Poll::Ready(Some(Err(e))) } std::task::Poll::Ready(None) => { - let outcome = OutcomeRecord { - status: OutcomeStatus::Incomplete, - finish_reason: None, - error: None, - usage: None, - }; - this.record(&outcome); + if !this.flush_pending() { + let outcome = OutcomeRecord { + status: OutcomeStatus::Incomplete, + finish_reason: None, + error: None, + error_value: None, + usage: None, + }; + this.record(&outcome); + } std::task::Poll::Ready(None) } std::task::Poll::Pending => std::task::Poll::Pending, @@ -1518,11 +1642,12 @@ where impl Drop for RecordingOutcomeStream { fn drop(&mut self) { - if !self.recorded { + if !self.recorded && !self.flush_pending() { let outcome = OutcomeRecord { status: OutcomeStatus::Cancelled, finish_reason: None, error: None, + error_value: None, usage: None, }; self.record(&outcome); @@ -1537,6 +1662,128 @@ mod tests { use super::*; use crate::generate::GenerateTextOptions; + struct CaptureRecorder(std::sync::Mutex>); + + impl Recorder for CaptureRecorder { + fn record_input(&self, _: &str, _: &CallOptions, _: &str, _: &str) {} + fn record_provider(&self, _: &str, _: &ProviderRecord) {} + fn record_exchange(&self, _: &str, _: &HttpExchange) {} + fn record_exchange_update( + &self, + _: &str, + _: u32, + _: u32, + _: &ResponseRecord, + _: Option, + ) { + } + fn record_outcome(&self, _: &str, outcome: &OutcomeRecord) { + self.0.lock().unwrap().push(outcome.clone()); + } + fn flush(&self) {} + } + + type PartResult = Result; + type CapturedStream = RecordingOutcomeStream< + futures::stream::Chain< + futures::stream::Iter>, + futures::stream::Pending, + >, + >; + + fn capture_stream(parts: Vec, recorder: &Arc) -> CapturedStream { + use futures::StreamExt as _; + let dyn_recorder: Arc = recorder.clone(); + RecordingOutcomeStream::new( + futures::stream::iter(parts).chain(futures::stream::pending()), + Some(dyn_recorder), + "call-1", + ) + } + + fn provider_error_part() -> Result { + Ok(crate::stream_part::StreamPart::Error { + error: crate::AiMuxError::Other("provider error".into()), + }) + } + + #[tokio::test] + async fn provider_error_outcome_keeps_usage_from_trailing_finish() { + use futures::StreamExt as _; + + let recorder = Arc::new(CaptureRecorder(std::sync::Mutex::new(Vec::new()))); + let finish = Ok(crate::stream_part::StreamPart::Finish { + finish_reason: crate::types::FinishReason { + unified: crate::types::FinishReasonUnified::Error, + raw: None, + }, + usage: crate::types::Usage::default(), + provider_metadata: None, + }); + let mut stream = capture_stream(vec![provider_error_part(), finish], &recorder); + let _ = stream.next().await; + let _ = stream.next().await; + drop(stream); + + let outcomes = recorder.0.lock().unwrap(); + assert_eq!(outcomes.len(), 1); + assert_eq!(outcomes[0].status, OutcomeStatus::Error); + assert_eq!(outcomes[0].error.as_deref(), Some("provider error")); + assert!(outcomes[0].usage.is_some(), "trailing Finish usage kept"); + } + + #[tokio::test] + async fn provider_error_outcome_is_flushed_when_the_stream_is_dropped_early() { + use futures::StreamExt as _; + + let recorder = Arc::new(CaptureRecorder(std::sync::Mutex::new(Vec::new()))); + let mut stream = capture_stream(vec![provider_error_part()], &recorder); + let _ = stream.next().await; + drop(stream); + + let outcomes = recorder.0.lock().unwrap(); + assert_eq!(outcomes.len(), 1); + assert_eq!(outcomes[0].status, OutcomeStatus::Error); + assert_eq!(outcomes[0].error.as_deref(), Some("provider error")); + } + + #[tokio::test] + async fn recoverable_frame_error_waits_for_the_trailing_finish() { + use futures::StreamExt as _; + + let recorder = Arc::new(CaptureRecorder(std::sync::Mutex::new(Vec::new()))); + let finish = Ok(crate::stream_part::StreamPart::Finish { + finish_reason: crate::types::FinishReason { + unified: crate::types::FinishReasonUnified::Stop, + raw: Some("stop".into()), + }, + usage: crate::types::Usage::default(), + provider_metadata: None, + }); + let mut stream = capture_stream( + vec![ + Err(crate::AiMuxError::JsonParse("bad frame".into())), + finish, + ], + &recorder, + ); + + assert!(matches!( + stream.next().await, + Some(Err(crate::AiMuxError::JsonParse(_))) + )); + assert!(recorder.0.lock().unwrap().is_empty()); + assert!(matches!( + stream.next().await, + Some(Ok(crate::stream_part::StreamPart::Finish { .. })) + )); + + let outcomes = recorder.0.lock().unwrap(); + assert_eq!(outcomes.len(), 1); + assert_eq!(outcomes[0].status, OutcomeStatus::Error); + assert!(outcomes[0].usage.is_some()); + } + fn sample_options() -> CallOptions { GenerateTextOptions { temperature: Some(0.7), @@ -1574,6 +1821,29 @@ mod tests { assert!(v.get("abort_signal").is_none()); } + #[test] + fn sensitive_key_matching_is_separator_and_case_insensitive() { + for key in [ + "token", + "access_token", + "accessToken", + "access-token", + "x-access-token", + "refreshToken", + "id_token", + "x-amz-security-token", + "Authorization", + "x-api-key", + "apiKey", + ] { + assert!(is_sensitive_key(key), "{key} must be redacted"); + } + // Usage counters carry no credentials and must stay loggable. + for key in ["max_tokens", "input_tokens", "output_tokens", "model"] { + assert!(!is_sensitive_key(key), "{key} must stay visible"); + } + } + #[test] fn redact_json_hides_sensitive_values_everywhere() { let v = serde_json::json!({ @@ -1584,7 +1854,11 @@ mod tests { "x-amz-security-token": "sts-tok" }, "provider_options": { "headers": { "Cookie": "s=1" } }, - "body_overrides": { "api_key": "sk-secret" }, + "body_overrides": { + "api_key": "sk-secret", + "token": "bearer-secret", + "access_token": "oauth-secret" + }, // 用量字段名含 "token" 子串,但非凭据——contains("token") 曾误伤。 "usage": { "max_output_tokens": 4096, @@ -1600,6 +1874,8 @@ mod tests { assert_eq!(r["headers"]["X-Key"], "ok"); assert_eq!(r["provider_options"]["headers"]["Cookie"], "[REDACTED]"); assert_eq!(r["body_overrides"]["api_key"], "[REDACTED]"); + assert_eq!(r["body_overrides"]["token"], "[REDACTED]"); + assert_eq!(r["body_overrides"]["access_token"], "[REDACTED]"); // 含 "token" 子串的用量字段不再被误脱敏(回归 contains("token"))。 assert_eq!(r["usage"]["max_output_tokens"], 4096); assert_eq!(r["usage"]["prompt_tokens"], 10); @@ -1634,6 +1910,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, ); @@ -1669,6 +1946,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }; assert!(!rec.ready(), "empty-prompt placeholder must not finalize"); @@ -1689,6 +1967,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, ); @@ -1746,6 +2025,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, }) @@ -1784,6 +2064,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, ); @@ -1836,7 +2117,9 @@ mod tests { rec.record_exchange( "call-x2", &HttpExchange { + step: None, attempt: 0, + exchange_index: 0, request: HttpRecord { method: "post".into(), url: "u".into(), @@ -1855,6 +2138,7 @@ mod tests { rec.record_exchange_update( "call-x2", 0, + 0, &ResponseRecord { status: 200, headers: vec![], @@ -1870,6 +2154,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: None, error: None, + error_value: None, usage: None, }, ); @@ -1894,7 +2179,9 @@ mod tests { rec.record_exchange( "call-errp", &HttpExchange { + step: None, attempt: 3, + exchange_index: 0, request: HttpRecord { method: "post".into(), url: "u".into(), @@ -1920,6 +2207,7 @@ mod tests { rec.record_exchange_update( "call-errp", 3, + 0, &ResponseRecord { status: 200, headers: vec![], @@ -1936,6 +2224,7 @@ mod tests { status: OutcomeStatus::Error, finish_reason: None, error: Some("mid-stream".into()), + error_value: None, usage: None, }, ); @@ -1981,6 +2270,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: None, error: None, + error_value: None, usage: None, }, ); @@ -2030,7 +2320,9 @@ mod tests { fn sample_exchange() -> HttpExchange { HttpExchange { + step: None, attempt: 0, + exchange_index: 0, request: HttpRecord { method: "post".into(), url: "https://api.openai.com/v1/chat/completions".into(), @@ -2058,6 +2350,7 @@ mod tests { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, } } @@ -2396,7 +2689,9 @@ mod tests { fn skeleton_at(attempt: u32) -> HttpExchange { HttpExchange { + step: None, attempt, + exchange_index: 1, request: HttpRecord { method: "post".into(), url: "u".into(), @@ -2424,6 +2719,40 @@ mod tests { ) } + #[test] + fn interleaved_child_attempts_keep_local_attempt_and_exchange_counters() { + let recorder: Arc = + Arc::new(CaptureRecorder(std::sync::Mutex::new(Vec::new()))); + let root = RecordingContext::new("call-1", recorder); + let _ = root.start_attempt(); // root operation attempt = 1 + + let a = root.child("router[0]:mock/a"); + let b = root.child("router[1]:mock/b"); + assert_eq!(a.step().as_deref(), Some("router[0]:mock/a")); + assert_eq!(root.step(), None); + + // Attempt numbers stay unique across the whole call (shared counter), + // so (attempt, exchange_index) still identifies an exchange without + // keying on step. + let a1 = a.start_attempt(); + let b1 = b.start_attempt(); + let mut attempts = vec![a1, b1]; + attempts.sort_unstable(); + attempts.dedup(); + assert_eq!(attempts.len(), 2); + + // Starting b after a must not relabel a's in-flight exchange. + assert_eq!(a.next_exchange(), (a1, 1)); + assert_eq!(b.next_exchange(), (b1, 1)); + assert_eq!(root.next_exchange(), (1, 1)); + + // Each child owns its exchange counter: one child's start_attempt + // reset cannot race another child's in-flight attempt. + let a2 = a.start_attempt(); // resets only a's local exchange counter + assert_eq!(a.next_exchange(), (a2, 1)); + assert_eq!(b.next_exchange(), (b1, 2)); + } + #[test] fn insert_exchange_and_apply_update_match_semantics() { let mut rec = empty_recording("c"); @@ -2447,7 +2776,7 @@ mod tests { // apply:1 匹配 → Patched。 assert_eq!( - apply_exchange_update(&mut rec, 1, resp_status(200), None), + apply_exchange_update(&mut rec, 1, 1, resp_status(200), None), UpdateMatch::Patched ); assert_eq!( @@ -2464,14 +2793,14 @@ mod tests { // apply:0 匹配 → NotFound。 assert_eq!( - apply_exchange_update(&mut rec, 9, resp_status(200), None), + apply_exchange_update(&mut rec, 9, 0, resp_status(200), None), UpdateMatch::NotFound ); // apply:2 匹配(手动塞重复)→ Ambiguous,不 patch 任何一条。 rec.exchanges.push(skeleton_at(1)); assert_eq!( - apply_exchange_update(&mut rec, 1, resp_status(503), Some("amb".into())), + apply_exchange_update(&mut rec, 1, 1, resp_status(503), Some("amb".into())), UpdateMatch::Ambiguous ); assert!( @@ -2506,7 +2835,7 @@ mod tests { let ring = RingRecorder::new(); ring.record_input("c", &sample_options(), "openai", "gpt-4o"); // 无骨架直接 update → 0 匹配 → 标记 inconsistent,不静默丢弃。 - ring.record_exchange_update("c", 0, &resp_status(200), None); + ring.record_exchange_update("c", 0, 0, &resp_status(200), None); assert!(ring.inconsistent_call_ids().contains(&"c".to_string())); assert_eq!( ring.pending_count(), diff --git a/aimux-core/src/replay.rs b/aimux-core/src/replay.rs index 2ab384ea..46b88965 100644 --- a/aimux-core/src/replay.rs +++ b/aimux-core/src/replay.rs @@ -1102,7 +1102,9 @@ mod tests { provider_options: None, }, exchanges: vec![HttpExchange { + step: None, attempt: 0, + exchange_index: 0, request: HttpRecord { method: "post".into(), url: "https://api.openai.com/v1/chat/completions".into(), @@ -1127,6 +1129,7 @@ mod tests { status: crate::recording::OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, complete: true, @@ -1835,7 +1838,9 @@ mod tests { "usage": {"prompt_tokens":5,"completion_tokens":3,"total_tokens":8} }); rec.exchanges.push(HttpExchange { + step: None, attempt: 1, + exchange_index: 1, request: HttpRecord { method: "post".into(), url: "https://api.openai.com/v1/chat/completions".into(), diff --git a/aimux-core/src/reranking_model.rs b/aimux-core/src/reranking_model.rs index 0b2a4658..a07c2051 100644 --- a/aimux-core/src/reranking_model.rs +++ b/aimux-core/src/reranking_model.rs @@ -1,4 +1,4 @@ -//! The `RerankingModel` trait — the provider-facing interface for reranking. +//! The `RerankingModel` trait — the provider-facing interface for reranking. //! //! Aligned with Vercel AI SDK `RerankingModelV4` //! (`reference/ai/packages/provider/src/reranking-model/v4/`). @@ -8,9 +8,8 @@ use serde::{Deserialize, Serialize}; use ts_rs::TS; use crate::error::AiMuxError; -use crate::shared::{ - AbortSignal, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning, -}; +use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning}; +use crate::{AbortSignal, retry, timeout}; /// Documents to rerank: either a list of texts or a list of JSON objects. /// @@ -44,6 +43,12 @@ pub struct RerankingCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional provider-specific options, keyed by provider name. pub provider_options: Option, @@ -59,6 +64,8 @@ impl RerankingCallOptions { query: query.into(), top_n: None, abort_signal: None, + max_retries: None, + timeout: None, provider_options: None, headers: None, } @@ -128,6 +135,10 @@ pub trait RerankingModel: Send + Sync { /// Provider-specific model ID, e.g. `"rerank-english-v3.0"`. fn model_id(&self) -> &str; + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Rerank a list of documents using the query. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -136,3 +147,27 @@ pub trait RerankingModel: Send + Sync { options: &RerankingCallOptions, ) -> Result; } + +/// User-facing reranking with Core-owned retry and timeout. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +pub async fn rerank( + model: &dyn RerankingModel, + options: RerankingCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + timeout::run( + retries.retry(|| model.do_rerank(&options)), + abort_signal.as_ref(), + timeout, + ) + .await +} diff --git a/aimux-core/src/retry.rs b/aimux-core/src/retry.rs new file mode 100644 index 00000000..dd97a147 --- /dev/null +++ b/aimux-core/src/retry.rs @@ -0,0 +1,644 @@ +//! Operation-level retry, aligned with the AI SDK retry utilities. + +use std::future::Future; +use std::time::Duration; + +use rand::Rng; + +use crate::AbortSignal; +use crate::error::{AiMuxError, RetryError, RetryErrorReason}; + +const DEFAULT_MAX_RETRIES: u32 = 2; +const INITIAL_DELAY_MS: u64 = 2_000; +const BACKOFF_FACTOR: u32 = 2; + +/// Default retry settings for model operations. +/// +/// Per-call `max_retries` overrides only [`Self::max_retries`]; the configured +/// delay and backoff factor remain in effect. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RetryConfig { + /// Maximum number of retries after the initial attempt. + pub max_retries: u32, + /// Initial delay between attempts. + pub initial_delay: Duration, + /// Multiplier applied to the delay after each retry. + pub backoff_factor: u32, +} + +impl Default for RetryConfig { + fn default() -> Self { + Self { + max_retries: DEFAULT_MAX_RETRIES, + initial_delay: Duration::from_millis(INITIAL_DELAY_MS), + backoff_factor: BACKOFF_FACTOR, + } + } +} + +/// A prepared operation retry function and its resolved retry count. +/// +/// This is the Rust representation of the AI SDK's +/// `prepareRetries()` result: `{ maxRetries, retry }`. +#[derive(Debug, Clone)] +pub struct PreparedRetries { + /// The effective maximum after applying the default. + pub max_retries: u32, + initial_delay_ms: u64, + backoff_factor: u64, + abort_signal: Option, +} + +/// Bind retry settings, applying a per-call retry-count override. +#[must_use] +pub fn prepare_retries( + max_retries: Option, + mut config: RetryConfig, + abort_signal: Option, +) -> PreparedRetries { + if let Some(max_retries) = max_retries { + config.max_retries = max_retries; + } + PreparedRetries { + max_retries: config.max_retries, + initial_delay_ms: u64::try_from(config.initial_delay.as_millis()).unwrap_or(u64::MAX), + backoff_factor: u64::from(config.backoff_factor), + abort_signal, + } +} + +impl PreparedRetries { + /// Retry a complete operation using API-call retryability and response hints. + /// + /// # Errors + /// + /// Returns the operation error, retry exhaustion, or caller cancellation. + pub async fn retry(&self, operation: F) -> Result + where + F: FnMut() -> Fut, + Fut: Future>, + { + retry_with_exponential_backoff_respecting_retry_headers( + operation, + self.max_retries, + self.initial_delay_ms, + self.backoff_factor, + self.abort_signal.as_ref(), + ) + .await + } +} + +async fn retry_with_exponential_backoff_respecting_retry_headers( + operation: F, + max_retries: u32, + initial_delay_ms: u64, + backoff_factor: u64, + abort_signal: Option<&AbortSignal>, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ + retry_with_exponential_backoff( + operation, + max_retries, + initial_delay_ms, + backoff_factor, + abort_signal, + |error: &AiMuxError| matches!(error, AiMuxError::ApiCall(detail) if detail.is_retryable), + |error: &AiMuxError, exponential_delay| { + retry_delay_with_jitter(error, exponential_delay, |maximum| { + rand::thread_rng().gen_range(0..=maximum) + }) + }, + ) + .await +} + +/// Generic exponential retry primitive. The caller supplies retry +/// classification and delay selection. +/// +/// # Errors +/// +/// Returns the operation error, retry exhaustion, or caller cancellation. +pub(crate) async fn retry_with_exponential_backoff( + mut operation: F, + max_retries: u32, + initial_delay_ms: u64, + backoff_factor: u64, + abort_signal: Option<&AbortSignal>, + mut should_retry: ShouldRetry, + mut get_delay_ms: GetDelay, +) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, + ShouldRetry: FnMut(&AiMuxError) -> bool, + GetDelay: FnMut(&AiMuxError, u64) -> u64, +{ + let mut errors = Vec::new(); + let mut exponential_delay = initial_delay_ms; + + loop { + let result = match abort_signal { + Some(signal) => { + if signal.is_aborted() { + return Err(AiMuxError::from_abort_signal(signal)); + } + tokio::select! { + biased; + () = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), + result = async { operation().await } => result, + } + } + None => operation().await, + }; + + match result { + Ok(value) => return Ok(value), + Err(error @ (AiMuxError::Aborted(_) | AiMuxError::Timeout(_))) => { + return Err(error); + } + // A safe inner exchange may exhaust its own retries inside a + // larger non-idempotent operation. Pass that exhaustion through + // so the enclosing operation is not submitted again. + Err(error @ AiMuxError::Retry(_)) => return Err(error), + Err(error) if max_retries == 0 => return Err(error), + Err(error) => { + tracing::warn!( + target: "aimux_core::retry", + attempt = errors.len() + 1, + error = %error, + "operation attempt failed" + ); + errors.push(error); + let try_number = errors.len() as u32; + + if try_number > max_retries { + tracing::error!( + target: "aimux_core::retry", + attempts = errors.len(), + error = %errors.last().expect("retry history is non-empty"), + "operation retry exhausted" + ); + return Err(AiMuxError::Retry(RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors, + })); + } + + let error = errors.last().expect("attempt error was just appended"); + if should_retry(error) { + let retry_delay = get_delay_ms(error, exponential_delay); + delay(Duration::from_millis(retry_delay), abort_signal).await?; + exponential_delay = exponential_delay.saturating_mul(backoff_factor); + continue; + } + + if try_number == 1 { + tracing::error!( + target: "aimux_core::retry", + error = %error, + "operation failed without retry" + ); + return Err(errors.pop().expect("the first attempt error exists")); + } + + tracing::error!( + target: "aimux_core::retry", + attempts = errors.len(), + error = %error, + "operation stopped on non-retryable error" + ); + return Err(AiMuxError::Retry(RetryError { + reason: RetryErrorReason::ErrorNotRetryable, + errors, + })); + } + } + } +} + +/// Abort-aware delay between attempts (also used by video poll pacing). +/// +/// # Errors +/// +/// Returns [`AiMuxError::Aborted`] when the caller cancels during the delay. +pub(crate) async fn delay( + duration: Duration, + abort_signal: Option<&AbortSignal>, +) -> Result<(), AiMuxError> { + match abort_signal { + Some(signal) => { + tokio::select! { + biased; + () = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), + () = tokio::time::sleep(duration) => Ok(()), + } + } + None => { + tokio::time::sleep(duration).await; + Ok(()) + } + } +} + +fn retry_header_delay(error: &AiMuxError, exponential_delay: u64) -> Option { + let parsed = u64::try_from(error.retry_after_hint()?).ok()?; + + (parsed < 60_000 || parsed < exponential_delay).then_some(parsed) +} + +fn retry_delay_with_jitter( + error: &AiMuxError, + exponential_delay: u64, + jitter: impl FnOnce(u64) -> u64, +) -> u64 { + retry_header_delay(error, exponential_delay) + .unwrap_or_else(|| jitter(exponential_delay).min(exponential_delay)) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + use crate::ApiCallError; + + fn api_error(retryable: bool) -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: retryable, + ..ApiCallError::new("retry", "https://example.test", serde_json::json!({})) + })) + } + + fn api_error_with_headers(status: u16, headers: &[(&str, &str)]) -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_headers: Some( + headers + .iter() + .map(|(name, value)| ((*name).to_owned(), (*value).to_owned())) + .collect::>(), + ), + is_retryable: true, + ..ApiCallError::new("retry", "https://example.test", serde_json::json!({})) + })) + } + + fn should_retry(error: &AiMuxError) -> bool { + error.is_retryable() + } + + fn unchanged_delay(_error: &AiMuxError, delay: u64) -> u64 { + delay + } + + #[tokio::test] + async fn prepared_retries_resolves_the_default_and_binds_abort() { + assert_eq!( + prepare_retries(None, RetryConfig::default(), None).max_retries, + 2 + ); + + let signal = AbortSignal::new(); + signal.abort(); + let retries = prepare_retries(Some(4), RetryConfig::default(), Some(signal)); + let attempts = AtomicUsize::new(0); + + let error = retries + .retry(|| { + attempts.fetch_add(1, Ordering::SeqCst); + async { Ok::<(), AiMuxError>(()) } + }) + .await + .unwrap_err(); + + assert_eq!(retries.max_retries, 4); + assert!(matches!(error, AiMuxError::Aborted(_))); + assert_eq!(attempts.load(Ordering::SeqCst), 0); + } + + #[test] + fn per_call_max_preserves_the_model_delay_and_backoff() { + let retries = prepare_retries( + Some(4), + RetryConfig { + max_retries: 7, + initial_delay: Duration::from_millis(123), + backoff_factor: 3, + }, + None, + ); + + assert_eq!(retries.max_retries, 4); + assert_eq!(retries.initial_delay_ms, 123); + assert_eq!(retries.backoff_factor, 3); + } + + #[tokio::test] + async fn first_non_retryable_error_is_not_wrapped() { + let error = retry_with_exponential_backoff( + || async { Err::<(), _>(AiMuxError::InvalidArgument("bad".into())) }, + 2, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::InvalidArgument(_))); + } + + #[tokio::test] + async fn zero_retries_returns_the_original_error() { + let attempts = AtomicUsize::new(0); + let error = retry_with_exponential_backoff( + || { + attempts.fetch_add(1, Ordering::SeqCst); + async { Err::<(), _>(api_error(true)) } + }, + 0, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::ApiCall(_))); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn exhausted_attempts_preserve_the_complete_history() { + let attempts = AtomicUsize::new(0); + let error = retry_with_exponential_backoff( + || { + attempts.fetch_add(1, Ordering::SeqCst); + async { Err::<(), _>(api_error(true)) } + }, + 2, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + let AiMuxError::Retry(error) = error else { + panic!("expected retry error") + }; + assert_eq!(error.reason, RetryErrorReason::MaxRetriesExceeded); + assert_eq!(error.errors.len(), 3); + assert_eq!(attempts.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn later_non_retryable_error_preserves_both_attempts() { + let attempts = AtomicUsize::new(0); + let error = retry_with_exponential_backoff( + || { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + async move { Err::<(), _>(api_error(attempt == 0)) } + }, + 2, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + let AiMuxError::Retry(error) = error else { + panic!("expected retry error") + }; + assert_eq!(error.reason, RetryErrorReason::ErrorNotRetryable); + assert_eq!(error.errors.len(), 2); + } + + #[tokio::test] + async fn retry_error_is_never_nested() { + let inner = RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors: vec![api_error(true), api_error(true)], + }; + let error = retry_with_exponential_backoff( + || { + let error = inner.clone(); + async move { Err::<(), _>(AiMuxError::Retry(error)) } + }, + 2, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + + assert!(matches!(error, AiMuxError::Retry(_))); + let AiMuxError::Retry(error) = error else { + unreachable!() + }; + assert!( + error + .errors + .iter() + .all(|error| !matches!(error, AiMuxError::Retry(_))) + ); + } + + #[tokio::test] + async fn timeout_and_abort_are_never_wrapped() { + for source in [ + AiMuxError::Timeout("Total timeout of 1ms exceeded".into()), + AiMuxError::Aborted("request aborted".into()), + ] { + let expected = source.to_string(); + let error = retry_with_exponential_backoff( + || { + let source = source.clone(); + async move { Err::<(), _>(source) } + }, + 2, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + assert_eq!(error.to_string(), expected); + assert!(!matches!(error, AiMuxError::Retry(_))); + } + } + + #[tokio::test] + async fn max_retries_check_precedes_final_retryability() { + let attempts = AtomicUsize::new(0); + let error = retry_with_exponential_backoff( + || { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + async move { Err::<(), _>(api_error(attempt < 2)) } + }, + 2, + 0, + 2, + None, + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + let AiMuxError::Retry(error) = error else { + panic!("expected retry error") + }; + assert_eq!(error.reason, RetryErrorReason::MaxRetriesExceeded); + assert_eq!(error.errors.len(), 3); + } + + #[test] + fn retry_headers_are_exact_and_only_exponential_delay_is_jittered() { + let milliseconds = api_error_with_headers(503, &[("retry-after-ms", "1500")]); + assert_eq!(retry_header_delay(&milliseconds, 2_000), Some(1_500)); + assert_eq!( + retry_delay_with_jitter(&milliseconds, 2_000, |_| panic!("hint must not jitter")), + 1_500 + ); + + let seconds = api_error_with_headers(503, &[("retry-after", "1.25")]); + assert_eq!(retry_header_delay(&seconds, 2_000), Some(1_250)); + + let unreasonable = api_error_with_headers(503, &[("retry-after-ms", "120000")]); + assert_eq!(retry_header_delay(&unreasonable, 2_000), None); + assert_eq!(retry_header_delay(&unreasonable, u64::MAX), Some(120_000)); + assert_eq!(retry_delay_with_jitter(&unreasonable, 2_000, |_| 0), 0); + assert_eq!( + retry_delay_with_jitter(&unreasonable, 2_000, |max| max), + 2_000 + ); + + let nan_then_seconds = + api_error_with_headers(503, &[("retry-after-ms", "NaN"), ("retry-after", "1")]); + assert_eq!(retry_header_delay(&nan_then_seconds, 2_000), Some(1_000)); + } + + #[tokio::test] + async fn abort_during_backoff_returns_abort_without_history() { + let signal = AbortSignal::new(); + let abort_after_attempt = signal.clone(); + let error = retry_with_exponential_backoff( + || async { Err::<(), _>(api_error(true)) }, + 2, + 60_000, + 2, + Some(&signal), + move |_error: &AiMuxError| { + abort_after_attempt.abort(); + true + }, + unchanged_delay, + ) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::Aborted(message) if message == "request aborted")); + } + + #[tokio::test] + async fn pre_aborted_signal_does_not_start_an_attempt() { + let signal = AbortSignal::new(); + signal.abort(); + let attempts = AtomicUsize::new(0); + + let error = retry_with_exponential_backoff( + || { + attempts.fetch_add(1, Ordering::SeqCst); + async { Ok::<(), AiMuxError>(()) } + }, + 2, + 0, + 2, + Some(&signal), + should_retry, + unchanged_delay, + ) + .await + .unwrap_err(); + + assert!(matches!(error, AiMuxError::Aborted(message) if message == "request aborted")); + assert_eq!(attempts.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn abort_during_attempt_drops_the_attempt_future() { + struct DropMarker(Arc); + + impl Drop for DropMarker { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + let signal = AbortSignal::new(); + let dropped = Arc::new(AtomicUsize::new(0)); + let marker = dropped.clone(); + let future = retry_with_exponential_backoff( + move || { + let marker = DropMarker(marker.clone()); + async move { + let _marker = marker; + std::future::pending::>().await + } + }, + 2, + 0, + 2, + Some(&signal), + should_retry, + unchanged_delay, + ); + tokio::pin!(future); + assert!(matches!( + futures::poll!(future.as_mut()), + std::task::Poll::Pending + )); + + signal.abort(); + let error = future.await.unwrap_err(); + + assert!(matches!(error, AiMuxError::Aborted(message) if message == "request aborted")); + assert_eq!(dropped.load(Ordering::SeqCst), 1); + } + + #[test] + fn retry_error_messages_match_the_ai_sdk() { + let max = RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors: vec![api_error(true), api_error(true), api_error(true)], + }; + assert_eq!( + max.to_string(), + "Failed after 3 attempts. Last error: API call error: retry" + ); + + let non_retryable = RetryError { + reason: RetryErrorReason::ErrorNotRetryable, + errors: vec![api_error(true), api_error(false)], + }; + assert_eq!( + non_retryable.to_string(), + "Failed after 2 attempts with non-retryable error: 'API call error: retry'" + ); + } +} diff --git a/aimux-core/src/router.rs b/aimux-core/src/router.rs index 0b352fdf..3a0745d7 100644 --- a/aimux-core/src/router.rs +++ b/aimux-core/src/router.rs @@ -15,6 +15,7 @@ use crate::language_model::LanguageModel; use crate::language_model_message::LanguageModelPrompt; use crate::options::CallOptions; use crate::result::{GenerateResult, StreamResult}; +use crate::retry; // ───────────────────────────────────────────────────────────────────────────── // Router trait @@ -135,10 +136,9 @@ impl Default for RouterConfig { /// A composite model that routes each call to one child and (optionally) falls /// back across the rest on error. /// -/// Streaming does not fall back: by the time `do_stream` returns, `StreamStart` -/// has been emitted to the user and a retry would duplicate tokens. This -/// matches Hermes / LiteLLM (route first, delegate second, no mid-stream -/// fallback). See RFC-0021 §3.3. +/// Streaming retries setup on the routed child, but does not fall back or retry +/// after setup: once stream parts are visible to the user, either would risk +/// duplicating tokens. See RFC-0021 §3.3. pub struct RouterModel { models: Vec, router: Box, @@ -177,7 +177,21 @@ impl RouterModel { if i == exclude { continue; } - match m.do_generate(options).await { + let retries = retry::prepare_retries( + options.max_retries, + m.retry_config(), + options.abort_signal.clone(), + ); + let child_options = options.for_step(self.step_label(i)); + match retries + .retry(|| { + if let Some(context) = &child_options.recording_context { + let _ = context.start_attempt(); + } + m.do_generate(&child_options) + }) + .await + { Ok(r) => return Ok(r), Err(e) => last_err = Some(e), } @@ -185,6 +199,11 @@ impl RouterModel { Err(last_err.expect("seeded primary_err makes last_err always Some")) } + fn step_label(&self, idx: usize) -> String { + let m = &self.models[idx]; + format!("router[{idx}]:{}/{}", m.provider(), m.model_id()) + } + /// Validate a `Router`-returned index before indexing `self.models`. A /// buggy/hostile `Router` (user-implementable trait) must not panic the /// process — surface it as `InvalidArgument` instead. @@ -210,10 +229,34 @@ impl LanguageModel for RouterModel { &self.config.model_id } + fn retry_config(&self) -> retry::RetryConfig { + // Retrying the composite would rerun routing and previously attempted + // children; each child is retried at its own execution boundary. + retry::RetryConfig { + max_retries: 0, + ..retry::RetryConfig::default() + } + } + async fn do_generate(&self, options: &CallOptions) -> Result { let raw = self.router.route(&options.prompt, &self.models)?; let idx = self.check_index(raw, "router")?; - match self.models[idx].do_generate(options).await { + let model = &self.models[idx]; + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + options.abort_signal.clone(), + ); + let child_options = options.for_step(self.step_label(idx)); + match retries + .retry(|| { + if let Some(context) = &child_options.recording_context { + let _ = context.start_attempt(); + } + model.do_generate(&child_options) + }) + .await + { Ok(r) => Ok(r), Err(e) => { if self.fallback == FallbackPolicy::OnError { @@ -228,27 +271,47 @@ impl LanguageModel for RouterModel { async fn do_stream(&self, options: &CallOptions) -> Result { let raw = self.router.route(&options.prompt, &self.models)?; let idx = self.check_index(raw, "router")?; - // Route first, delegate second. No mid-stream fallback (StreamStart is - // already emitted by the time we'd want to retry). - self.models[idx].do_stream(options).await + let model = &self.models[idx]; + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + options.abort_signal.clone(), + ); + let child_options = options.for_step(self.step_label(idx)); + // Only setup can be retried. The returned stream is passed through, so + // failures after setup never invoke another child or duplicate output. + retries + .retry(|| { + if let Some(context) = &child_options.recording_context { + let _ = context.start_attempt(); + } + model.do_stream(&child_options) + }) + .await } } #[cfg(test)] mod tests { use super::*; + use crate::recording::{ + HttpExchange, OutcomeRecord, ProviderRecord, Recorder, RecordingContext, ResponseRecord, + }; use crate::result::GenerateResult; use crate::stream_part::StreamPart; use crate::types::{FinishReason, FinishReasonUnified, Usage}; use async_trait::async_trait; use futures::stream; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; /// A mock child that either succeeds with `text` or always fails. struct MockChild { name: &'static str, text: String, fail: bool, + retry_failures: usize, + calls: Arc, } #[async_trait] @@ -260,9 +323,13 @@ mod tests { self.name } async fn do_generate(&self, _options: &CallOptions) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; if self.fail { return Err(AiMuxError::Other(format!("{} always fails", self.name))); } + if attempt <= self.retry_failures { + return Err(retryable_error(self.name)); + } Ok(GenerateResult { content: vec![crate::result::GenerateContent::Text { text: self.text.clone(), @@ -281,9 +348,13 @@ mod tests { }) } async fn do_stream(&self, _options: &CallOptions) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst) + 1; if self.fail { return Err(AiMuxError::Other(format!("{} always fails", self.name))); } + if attempt <= self.retry_failures { + return Err(retryable_error(self.name)); + } let parts: Vec> = vec![ Ok(StreamPart::StreamStart { warnings: vec![] }), Ok(StreamPart::TextDelta { @@ -308,11 +379,60 @@ mod tests { } } + fn retryable_error(name: &str) -> AiMuxError { + AiMuxError::ApiCall(Box::new(crate::ApiCallError { + status_code: Some(503), + response_headers: Some(std::collections::HashMap::from([( + "retry-after-ms".into(), + "0".into(), + )])), + is_retryable: true, + ..crate::ApiCallError::new( + format!("{name} retryable failure"), + "https://example.test", + serde_json::json!({}), + ) + })) + } + + fn retry_child( + name: &'static str, + text: &str, + retry_failures: usize, + ) -> (ChildModel, Arc) { + let calls = Arc::new(AtomicUsize::new(0)); + ( + Arc::new(MockChild { + name, + text: text.into(), + fail: false, + retry_failures, + calls: calls.clone(), + }), + calls, + ) + } + + struct CountingRuleRouter(Arc); + + impl Router for CountingRuleRouter { + fn route( + &self, + prompt: &LanguageModelPrompt, + models: &[ChildModel], + ) -> Result { + self.0.fetch_add(1, Ordering::SeqCst); + RuleRouter.route(prompt, models) + } + } + fn child(name: &'static str, text: &str, fail: bool) -> ChildModel { Arc::new(MockChild { name, text: text.into(), fail, + retry_failures: 0, + calls: Arc::new(AtomicUsize::new(0)), }) } @@ -534,4 +654,167 @@ mod tests { .unwrap(); assert_eq!(r.text, "works through trait object"); } + + #[tokio::test] + async fn retries_each_child_before_falling_back_without_rerouting() { + let (primary, primary_calls) = retry_child("primary", "primary", usize::MAX); + let (backup, backup_calls) = retry_child("backup", "backup", 1); + let route_calls = Arc::new(AtomicUsize::new(0)); + let router = RouterModel::new( + vec![primary, backup], + Box::new(CountingRuleRouter(route_calls.clone())), + FallbackPolicy::OnError, + RouterConfig::default(), + ); + assert_eq!(router.retry_config().max_retries, 0); + + let result = crate::generate::generate_text( + &router, + "hello", + crate::generate::GenerateTextOptions { + max_retries: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(result.text, "backup"); + assert_eq!(route_calls.load(Ordering::SeqCst), 1); + assert_eq!(primary_calls.load(Ordering::SeqCst), 3); + assert_eq!(backup_calls.load(Ordering::SeqCst), 2); + } + + /// Children must receive a step-labeled child recording context so their + /// exchanges and retry attempts are recorded as nested steps. + #[tokio::test] + async fn children_receive_step_labeled_recording_contexts() { + struct StepProbe { + seen: std::sync::Mutex>>, + } + + #[async_trait] + impl LanguageModel for StepProbe { + fn provider(&self) -> &str { + "mock" + } + fn model_id(&self) -> &str { + "probe" + } + async fn do_generate( + &self, + options: &CallOptions, + ) -> Result { + self.seen.lock().unwrap().push( + options + .recording_context + .as_ref() + .and_then(RecordingContext::step), + ); + Err(AiMuxError::Other("probe fails".into())) + } + async fn do_stream(&self, _: &CallOptions) -> Result { + unreachable!() + } + } + + struct NoopRecorder; + impl Recorder for NoopRecorder { + fn record_input(&self, _: &str, _: &CallOptions, _: &str, _: &str) {} + fn record_provider(&self, _: &str, _: &ProviderRecord) {} + fn record_exchange(&self, _: &str, _: &HttpExchange) {} + fn record_exchange_update( + &self, + _: &str, + _: u32, + _: u32, + _: &ResponseRecord, + _: Option, + ) { + } + fn record_outcome(&self, _: &str, _: &OutcomeRecord) {} + fn flush(&self) {} + } + + let probe = Arc::new(StepProbe { + seen: std::sync::Mutex::new(Vec::new()), + }); + let router = RouterModel::new( + vec![probe.clone()], + Box::new(CountingRuleRouter(Arc::new(AtomicUsize::new(0)))), + FallbackPolicy::OnError, + RouterConfig::default(), + ); + + let options = CallOptions { + recording_context: Some(RecordingContext::new("call-1", Arc::new(NoopRecorder))), + ..CallOptions::default() + }; + let _ = router.do_generate(&options).await; + + let seen = probe.seen.lock().unwrap(); + assert!(!seen.is_empty()); + assert_eq!(seen[0].as_deref(), Some("router[0]:mock/probe")); + } + + /// Even when every child fails and the caller asked for per-call retries, + /// the composite operation must run exactly once: the exhausted child's + /// `Retry` error passes through the outer retry loop instead of replaying + /// routing and fallback. + #[tokio::test] + async fn all_children_failing_does_not_replay_the_composite() { + let (primary, primary_calls) = retry_child("primary", "primary", usize::MAX); + let (backup, backup_calls) = retry_child("backup", "backup", usize::MAX); + let route_calls = Arc::new(AtomicUsize::new(0)); + let router = RouterModel::new( + vec![primary, backup], + Box::new(CountingRuleRouter(route_calls.clone())), + FallbackPolicy::OnError, + RouterConfig::default(), + ); + + let error = crate::generate::generate_text( + &router, + "hello", + crate::generate::GenerateTextOptions { + max_retries: Some(2), + ..Default::default() + }, + ) + .await + .unwrap_err(); + + assert!(matches!(error, crate::AiMuxError::Retry(_)), "{error:?}"); + assert_eq!(route_calls.load(Ordering::SeqCst), 1); + assert_eq!(primary_calls.load(Ordering::SeqCst), 3); + assert_eq!(backup_calls.load(Ordering::SeqCst), 3); + } + + #[tokio::test] + async fn retries_stream_setup_on_selected_child_without_fallback_or_rerouting() { + let (primary, primary_stream_calls) = retry_child("primary", "primary", 1); + let (backup, backup_stream_calls) = retry_child("backup", "backup", 0); + let route_calls = Arc::new(AtomicUsize::new(0)); + let router = RouterModel::new( + vec![primary, backup], + Box::new(CountingRuleRouter(route_calls.clone())), + FallbackPolicy::OnError, + RouterConfig::default(), + ); + + let _result = crate::generate::stream_text( + &router, + "hello", + crate::generate::GenerateTextOptions { + max_retries: Some(2), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(route_calls.load(Ordering::SeqCst), 1); + assert_eq!(primary_stream_calls.load(Ordering::SeqCst), 2); + assert_eq!(backup_stream_calls.load(Ordering::SeqCst), 0); + } } diff --git a/aimux-core/src/search_model.rs b/aimux-core/src/search_model.rs index 36dfe68f..cd03f7b1 100644 --- a/aimux-core/src/search_model.rs +++ b/aimux-core/src/search_model.rs @@ -10,8 +10,9 @@ use serde::{Deserialize, Serialize}; use ts_rs::TS; use crate::error::AiMuxError; -use crate::shared::{AbortSignal, SharedHeaders, SharedProviderMetadata, SharedProviderOptions}; +use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions}; use crate::types::Warning; +use crate::{AbortSignal, retry, timeout}; /// Options passed to [`SearchModel::do_search`]. #[derive(Debug, Clone, Serialize, Deserialize, TS)] @@ -44,6 +45,12 @@ pub struct SearchCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional provider-specific options, keyed by provider name. pub provider_options: Option, @@ -62,6 +69,8 @@ impl SearchCallOptions { include_domains: None, exclude_domains: None, abort_signal: None, + max_retries: None, + timeout: None, provider_options: None, headers: None, } @@ -141,8 +150,36 @@ pub trait SearchModel: Send + Sync { /// `"tavily-search"`; others accept endpoint-specific names). fn model_id(&self) -> &str; + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Execute a search query and return results. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. async fn do_search(&self, options: &SearchCallOptions) -> Result; } + +/// User-facing search with Core-owned retry and timeout. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +pub async fn search( + model: &dyn SearchModel, + options: SearchCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + timeout::run( + retries.retry(|| model.do_search(&options)), + abort_signal.as_ref(), + timeout, + ) + .await +} diff --git a/aimux-core/src/shared.rs b/aimux-core/src/shared.rs index 416252e5..705f3c94 100644 --- a/aimux-core/src/shared.rs +++ b/aimux-core/src/shared.rs @@ -12,11 +12,11 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; use serde_json::Value; -use tokio_util::sync::CancellationToken; use ts_rs::TS; -// Re-export the existing `Warning` so providers can import everything they -// need from `crate::shared` in one place. +// Keep the established shared-module paths while the canonical cancellation +// implementation lives in its focused module. +pub use crate::abort_signal::AbortSignal; pub use crate::types::Warning; /// Additional HTTP headers sent with a provider request. @@ -80,56 +80,6 @@ pub enum FileData { Text { text: String }, } -/// A cancellation signal analogous to the Web `AbortSignal`. -/// -/// Event-driven: backed by `tokio_util::sync::CancellationToken` (itself a -/// `tokio::sync::Notify`). Consumers can poll [`is_aborted`](Self::is_aborted) -/// synchronously, or await [`cancelled`](Self::cancelled) for prompt, -/// notification-based wakeup — no polling loop needed. -/// -/// This type is `Send + Sync` and cheap to clone. -#[derive(Debug, Clone)] -pub struct AbortSignal { - token: CancellationToken, -} - -impl Default for AbortSignal { - fn default() -> Self { - Self::new() - } -} - -impl AbortSignal { - /// Create a fresh, un-aborted signal. - #[must_use] - pub fn new() -> Self { - Self { - token: CancellationToken::new(), - } - } - - /// Request cancellation. All clones observe the aborted state and any - /// pending [`cancelled`](Self::cancelled) futures resolve. - pub fn abort(&self) { - self.token.cancel(); - } - - /// Returns `true` once [`abort`](Self::abort) has been called. - #[must_use] - pub fn is_aborted(&self) -> bool { - self.token.is_cancelled() - } - - /// A future that resolves as soon as the signal is aborted. - /// - /// If the signal is already aborted, the future resolves immediately on - /// the first poll. Suitable for `tokio::select!` arms. - pub fn cancelled(&self) -> impl std::future::Future + Send + 'static { - let token = self.token.clone(); - async move { token.cancelled_owned().await } - } -} - /// Image/video size in `{width}x{height}` format (e.g. `"1024x1024"`). /// /// Newtype that enforces the `WxH` format used by V4 image `size` and video diff --git a/aimux-core/src/speech_model.rs b/aimux-core/src/speech_model.rs index 77f60410..5add06e9 100644 --- a/aimux-core/src/speech_model.rs +++ b/aimux-core/src/speech_model.rs @@ -1,4 +1,4 @@ -//! The `SpeechModel` trait — the provider-facing interface for text-to-speech. +//! The `SpeechModel` trait — the provider-facing interface for text-to-speech. //! //! Aligned with Vercel AI SDK `SpeechModelV4` //! (`reference/ai/packages/provider/src/speech-model/v4/`). @@ -8,9 +8,8 @@ use serde::{Deserialize, Serialize}; use ts_rs::TS; use crate::error::AiMuxError; -use crate::shared::{ - AbortSignal, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning, -}; +use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning}; +use crate::{AbortSignal, retry, timeout}; /// Generated audio: a base64-encoded string or raw binary bytes. /// @@ -60,6 +59,12 @@ pub struct SpeechCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional HTTP headers to send with the request. pub headers: Option, } @@ -76,6 +81,8 @@ impl SpeechCallOptions { language: None, provider_options: None, abort_signal: None, + max_retries: None, + timeout: None, headers: None, } } @@ -141,8 +148,36 @@ pub trait SpeechModel: Send + Sync { /// Provider-specific model ID, e.g. `"tts-1"`. fn model_id(&self) -> &str; + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Generate speech audio from text. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. async fn do_generate(&self, options: &SpeechCallOptions) -> Result; } + +/// User-facing speech generation with Core-owned retry and timeout. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +pub async fn generate_speech( + model: &dyn SpeechModel, + options: SpeechCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + timeout::run( + retries.retry(|| model.do_generate(&options)), + abort_signal.as_ref(), + timeout, + ) + .await +} diff --git a/aimux-core/src/timeout.rs b/aimux-core/src/timeout.rs new file mode 100644 index 00000000..d46dc679 --- /dev/null +++ b/aimux-core/src/timeout.rs @@ -0,0 +1,263 @@ +//! Core-owned operation deadlines and caller cancellation. + +use std::future::{Future, pending}; + +use tokio::time::Instant; + +use crate::AbortSignal; +use crate::error::AiMuxError; +use crate::options::TimeoutConfiguration; + +/// Deadlines shared by the establishment and streaming phases of one model +/// operation. No timer task is spawned: the future currently driving the +/// operation observes the deadline directly. +#[derive(Debug, Clone, Copy)] +pub(crate) struct OperationTimeout { + total: Option, + step: Option, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct TimeoutDeadline { + pub(crate) at: Instant, + label: &'static str, + duration_ms: u64, +} + +impl TimeoutDeadline { + pub(crate) fn from_now(label: &'static str, duration_ms: u64) -> Result { + Self::after(Instant::now(), label, duration_ms) + } + + fn after(start: Instant, label: &'static str, duration_ms: u64) -> Result { + let at = start + .checked_add(std::time::Duration::from_millis(duration_ms)) + .ok_or_else(|| { + AiMuxError::InvalidArgument(format!( + "{label} timeout of {duration_ms}ms exceeds the supported range" + )) + })?; + Ok(Self { + at, + label, + duration_ms, + }) + } + + #[must_use] + pub(crate) fn error(self) -> AiMuxError { + AiMuxError::Timeout(format!( + "{} timeout of {}ms exceeded", + self.label, self.duration_ms + )) + } +} + +impl OperationTimeout { + pub(crate) fn new(configuration: TimeoutConfiguration) -> Result { + let now = Instant::now(); + let total = configuration + .total_ms + .map(|duration_ms| TimeoutDeadline::after(now, "Total", duration_ms)) + .transpose()?; + let step = configuration + .step_ms + .map(|duration_ms| TimeoutDeadline::after(now, "Step", duration_ms)) + .transpose()?; + + // Validate streaming-only durations here as well, before a provider + // operation starts. Their actual deadlines begin at stream-specific + // points and are constructed again with the same checked helper. + if let Some(duration_ms) = configuration.first_chunk_ms { + TimeoutDeadline::after(now, "First chunk", duration_ms)?; + } + if let Some(duration_ms) = configuration.chunk_ms { + TimeoutDeadline::after(now, "Chunk", duration_ms)?; + } + + Ok(Self { total, step }) + } + + /// The first active operation deadline. Aimux currently has one model + /// step, so total and step begin together. + #[must_use] + pub(crate) fn deadline(self) -> Option { + match (self.total, self.step) { + (Some(total), Some(step)) if step.at < total.at => Some(step), + (Some(total), _) => Some(total), + (None, step) => step, + } + } +} + +/// Run an operation while directly observing caller cancellation and its +/// total/step deadline. Dropping `operation` cancels the in-flight Rust HTTP +/// future; timeout is not represented by mutating an AbortSignal. +pub(crate) async fn run( + operation: impl Future>, + abort_signal: Option<&AbortSignal>, + timeout: OperationTimeout, +) -> Result { + run_until(operation, abort_signal, timeout.deadline()).await +} + +/// The earlier of two optional deadlines. +pub(crate) fn min_deadline( + a: Option, + b: Option, +) -> Option { + match (a, b) { + (Some(a), Some(b)) => Some(if b.at < a.at { b } else { a }), + (Some(deadline), None) | (None, Some(deadline)) => Some(deadline), + (None, None) => None, + } +} + +/// [`run`] with an explicit deadline, for callers that race additional +/// phase-specific deadlines (the stream setup phase also observes +/// `first_chunk_ms`). +pub(crate) async fn run_until( + operation: impl Future>, + abort_signal: Option<&AbortSignal>, + deadline: Option, +) -> Result { + tokio::select! { + biased; + () = wait_for_abort(abort_signal) => { + Err(AiMuxError::Aborted("request aborted".into())) + } + () = wait_for_deadline(deadline) => { + Err(deadline.expect("deadline future only resolves for a deadline").error()) + } + result = operation => result, + } +} + +pub(crate) async fn wait_for_abort(signal: Option<&AbortSignal>) { + match signal { + Some(signal) => signal.cancelled().await, + None => pending().await, + } +} + +pub(crate) async fn wait_for_deadline(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline.at).await, + None => pending().await, + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use super::*; + use crate::retry; + + fn always_retry(_error: &AiMuxError) -> bool { + true + } + + fn unchanged_delay(_error: &AiMuxError, delay: u64) -> u64 { + delay + } + + #[tokio::test] + async fn total_timeout_cancels_the_operation_without_mutating_abort_signal() { + let timeout = OperationTimeout::new(TimeoutConfiguration { + total_ms: Some(1), + ..Default::default() + }) + .unwrap(); + let result = run( + std::future::pending::>(), + None, + timeout, + ) + .await; + assert!( + matches!(result, Err(AiMuxError::Timeout(message)) if message == "Total timeout of 1ms exceeded") + ); + } + + #[tokio::test] + async fn caller_abort_remains_distinct_from_timeout() { + let caller = AbortSignal::new(); + let timeout = OperationTimeout::new(TimeoutConfiguration { + total_ms: Some(60_000), + ..Default::default() + }) + .unwrap(); + caller.abort(); + let error = run( + std::future::pending::>(), + Some(&caller), + timeout, + ) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::Aborted(message) if message == "request aborted")); + } + + #[tokio::test] + async fn total_timeout_includes_retry_backoff() { + let attempts = AtomicUsize::new(0); + let retry = retry::retry_with_exponential_backoff( + || { + attempts.fetch_add(1, Ordering::SeqCst); + async { Err::<(), _>(AiMuxError::Other("retry".into())) } + }, + 2, + 60_000, + 2, + None, + always_retry, + unchanged_delay, + ); + let timeout = OperationTimeout::new(TimeoutConfiguration { + total_ms: Some(1), + ..Default::default() + }) + .unwrap(); + + let error = run(retry, None, timeout).await.unwrap_err(); + + assert!( + matches!(error, AiMuxError::Timeout(message) if message == "Total timeout of 1ms exceeded") + ); + assert_eq!(attempts.load(Ordering::SeqCst), 1); + } + + #[test] + fn largest_timeout_durations_never_panic() { + let configurations = [ + TimeoutConfiguration { + total_ms: Some(u64::MAX), + ..Default::default() + }, + TimeoutConfiguration { + step_ms: Some(u64::MAX), + ..Default::default() + }, + TimeoutConfiguration { + first_chunk_ms: Some(u64::MAX), + ..Default::default() + }, + TimeoutConfiguration { + chunk_ms: Some(u64::MAX), + ..Default::default() + }, + ]; + + for configuration in configurations { + // `Instant` ranges differ by platform, so this value may be + // representable; the contract is either success or a typed error. + let result = std::panic::catch_unwind(|| OperationTimeout::new(configuration)) + .expect("timeout validation must not panic"); + assert!(matches!( + result, + Ok(_) | Err(AiMuxError::InvalidArgument(_)) + )); + } + } +} diff --git a/aimux-core/src/trace/hash.rs b/aimux-core/src/trace/hash.rs index 22943dbe..cda39d34 100644 --- a/aimux-core/src/trace/hash.rs +++ b/aimux-core/src/trace/hash.rs @@ -24,6 +24,10 @@ pub fn mix(mut x: u64) -> u64 { #[must_use] pub fn hash64(key: u64, data: &[u8]) -> u64 { let mut h = key ^ 0x9e37_79b9_7f4a_7c15; + // Clippy 1.98 suggests `as_chunks::<8>()`, which is stable only since + // Rust 1.88; the workspace MSRV is 1.85. `unknown_lints` keeps this + // buildable on toolchains older than the lint itself. + #[allow(unknown_lints, clippy::chunks_exact_to_as_chunks)] let mut chunks = data.chunks_exact(8); for c in &mut chunks { let w = u64::from_le_bytes(c.try_into().unwrap()); diff --git a/aimux-core/src/trace/layer.rs b/aimux-core/src/trace/layer.rs index 5e2cac10..a5b3b6e3 100644 --- a/aimux-core/src/trace/layer.rs +++ b/aimux-core/src/trace/layer.rs @@ -517,6 +517,10 @@ impl LanguageModel for TraceLayer { self.inner.model_id() } + fn retry_config(&self) -> crate::retry::RetryConfig { + self.inner.retry_config() + } + /// RFC-0023 §3.3: transparent decorators must forward the inner snapshot /// (otherwise recording sees the decorator's minimal record). fn config_snapshot(&self) -> crate::recording::ProviderRecord { diff --git a/aimux-core/src/transcription_model.rs b/aimux-core/src/transcription_model.rs index 556e0980..efeb2c75 100644 --- a/aimux-core/src/transcription_model.rs +++ b/aimux-core/src/transcription_model.rs @@ -1,4 +1,4 @@ -//! The `TranscriptionModel` trait — the provider-facing interface for +//! The `TranscriptionModel` trait — the provider-facing interface for //! speech-to-text. //! //! Aligned with Vercel AI SDK `TranscriptionModelV4` @@ -6,7 +6,8 @@ //! //! This is the only non-chat model type with an optional streaming method, //! [`TranscriptionModel::do_stream`]. The default implementation returns -//! [`AiMuxError::UnsupportedFunctionality`]; providers override it as needed. +//! [`AiMuxError::UnsupportedFunctionality`]; providers override it and users +//! enter through [`stream_transcribe`]. use std::pin::Pin; @@ -16,9 +17,8 @@ use serde::{Deserialize, Serialize}; use ts_rs::TS; use crate::error::AiMuxError; -use crate::shared::{ - AbortSignal, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning, -}; +use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning}; +use crate::{AbortSignal, retry, timeout}; /// Audio input: raw bytes or a base64-encoded string. #[derive(Debug, Clone, Serialize, Deserialize, TS)] @@ -72,6 +72,12 @@ pub struct TranscriptionCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional HTTP headers to send with the request. pub headers: Option, } @@ -84,6 +90,8 @@ impl TranscriptionCallOptions { media_type: media_type.into(), provider_options: None, abort_signal: None, + max_retries: None, + timeout: None, headers: None, } } @@ -287,6 +295,10 @@ pub trait TranscriptionModel: Send + Sync { /// Provider-specific model ID, e.g. `"whisper-1"`. fn model_id(&self) -> &str; + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig::default() + } + /// Generate a transcript. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. @@ -310,3 +322,43 @@ pub trait TranscriptionModel: Send + Sync { ))) } } + +/// User-facing non-streaming transcription with Core-owned retry and timeout. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, timeout, or caller abort. +pub async fn transcribe( + model: &dyn TranscriptionModel, + options: TranscriptionCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + timeout::run( + retries.retry(|| model.do_generate(&options)), + abort_signal.as_ref(), + timeout, + ) + .await +} + +/// Start user-facing live transcription. +/// +/// The live audio stream cannot be replayed, so session setup is attempted +/// once. The provider session applies the supplied abort signal and streaming +/// timeouts throughout connect, send, and receive. +/// +/// # Errors +/// +/// Returns the provider's session-establishment error. +pub async fn stream_transcribe( + model: &dyn TranscriptionModel, + options: TranscriptionStreamOptions, +) -> Result { + model.do_stream(options).await +} diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs index aeb7a34a..7a9c267a 100644 --- a/aimux-core/src/video_model.rs +++ b/aimux-core/src/video_model.rs @@ -1,4 +1,4 @@ -//! The `VideoModel` trait — the provider-facing interface for video generation. +//! The `VideoModel` trait — the provider-facing interface for video generation. //! //! Aligned with Vercel AI SDK `VideoModelV4` //! (`reference/ai/packages/provider/src/video-model/v4/`). @@ -7,11 +7,13 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; use ts_rs::TS; +use std::time::{Duration, Instant}; + use crate::error::AiMuxError; use crate::shared::{ - AbortSignal, AspectRatio, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Size, - Warning, + AspectRatio, SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Size, Warning, }; +use crate::{AbortSignal, retry, timeout}; /// A video or image file used for video editing or image-to-video generation. /// @@ -84,7 +86,7 @@ pub enum VideoData { Binary { data: Vec, media_type: String }, } -/// Options passed to [`VideoModel::do_generate`]. +/// Options passed to [`VideoModel::do_start`] and [`VideoModel::do_status`]. /// /// Aligned with V4 `VideoModelV4CallOptions`. #[derive(Debug, Clone, Serialize, Deserialize, TS)] @@ -95,6 +97,9 @@ pub struct VideoCallOptions { /// Number of videos to generate. Default `1`; most models only support /// `n = 1` due to computational cost. + // serde default: typed binding structs omit unset fields, and a missing + // `n` must mean 1, not a hard deserialization failure at the FFI boundary. + #[serde(default = "default_video_n")] pub n: u32, /// Aspect ratio, in `{width}:{height}` format (e.g. `"16:9"`). @@ -126,6 +131,7 @@ pub struct VideoCallOptions { pub generate_audio: Option, /// Additional provider-specific options, keyed by provider name. + #[serde(default)] pub provider_options: SharedProviderOptions, /// Abort signal for cancelling the operation. @@ -133,10 +139,24 @@ pub struct VideoCallOptions { #[ts(skip)] pub abort_signal: Option, + /// Per-call retry override. `None` uses the model default. + pub max_retries: Option, + + /// Per-call poll pacing override for the start/status flow. Unset fields + /// fall back to the model's [`VideoModel::poll_config`]. + pub poll: Option, + + /// Per-call operation timeout. + pub timeout: Option, + /// Additional HTTP headers to send with the request. pub headers: Option, } +fn default_video_n() -> u32 { + 1 +} + impl VideoCallOptions { /// Create options with a prompt and `n = 1`, all other fields defaulted. pub fn new(prompt: impl Into) -> Self { @@ -154,12 +174,29 @@ impl VideoCallOptions { generate_audio: None, provider_options: SharedProviderOptions::new(), abort_signal: None, + max_retries: None, + poll: None, + timeout: None, headers: None, } } } -/// The result of [`VideoModel::do_generate`]. +/// Per-call poll pacing for video generation (AI SDK `generateVideo` `poll`). +#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, TS)] +#[ts(export)] +pub struct VideoPollOptions { + /// Delay between consecutive status checks, in milliseconds. + // `number`, not the `bigint` ts-rs infers from u64: the JS bindings pass + // options through `JSON.stringify`, which throws on BigInt. + #[ts(type = "number | null")] + pub interval_ms: Option, + /// Maximum total time to wait for completion, in milliseconds. + #[ts(type = "number | null")] + pub timeout_ms: Option, +} + +/// The final result of a video generation operation. /// /// Aligned with V4 `VideoModelV4Result`. #[derive(Debug, Clone, Serialize, Deserialize, TS)] @@ -190,9 +227,64 @@ pub struct VideoResponse { pub headers: Option, } +/// Result of starting an asynchronous video generation via +/// [`VideoModel::do_start`]. +/// +/// Aligned with V4 `VideoModelV4OperationStartResult`. +#[derive(Debug, Clone)] +pub struct VideoOperationStart { + /// JSON-serializable opaque reference passed to [`VideoModel::do_status`] + /// to check the status of the generation (e.g. a task ID or poll URL). + pub operation: serde_json::Value, + + /// Warnings for the start call, e.g. unsupported features. + pub warnings: Vec, + + /// Additional provider-specific metadata from the start call. + pub provider_metadata: Option, + + /// Response information for telemetry and debugging. + pub response: VideoResponse, +} + +/// Status of an asynchronous video generation, from [`VideoModel::do_status`]. +/// +/// Aligned with V4 `VideoModelV4OperationStatusResult`; its `error` arm maps +/// to `Err(AiMuxError)` (a terminally failed task should be non-retryable). +#[derive(Debug, Clone)] +pub enum VideoOperationStatus { + /// The generation is still in progress; poll again later. + Pending, + /// The generation is complete. + Completed(VideoResult), +} + +/// Pacing for the Core-owned status poll loop. +#[derive(Debug, Clone, Copy)] +pub struct VideoPollConfig { + /// Delay between consecutive `do_status` calls. + pub interval: Duration, + /// Maximum total time to wait for the generation to complete. + pub timeout: Duration, +} + +impl Default for VideoPollConfig { + fn default() -> Self { + // AI SDK generate-video defaults (intervalMs: 5000, timeoutMs: 600_000). + Self { + interval: Duration::from_secs(5), + timeout: Duration::from_secs(600), + } + } +} + /// The unified video generation model trait (provider-facing). /// -/// Aligned with V4 `VideoModelV4`. +/// Aligned with V4 `VideoModelV4`'s asynchronous `doStart`/`doStatus` flow. +/// Every current provider API is task-based (create, then poll), so the +/// synchronous `doGenerate` arm is not offered: the polling loop lives in +/// [`generate_video`], and each phase is retried independently against the +/// same operation reference. #[async_trait] pub trait VideoModel: Send + Sync { /// Specification version (always `"v4"`). @@ -206,13 +298,454 @@ pub trait VideoModel: Send + Sync { /// Provider-specific model ID, e.g. `"kling-video"`. fn model_id(&self) -> &str; + fn retry_config(&self) -> retry::RetryConfig { + retry::RetryConfig::default() + } + + /// Poll pacing for this model. Defaults to the AI SDK values; providers + /// with configurable polling should surface their configuration here. + fn poll_config(&self) -> VideoPollConfig { + VideoPollConfig::default() + } + /// Limit of how many videos can be generated in a single API call. /// /// `None` means no fixed limit. Most video models only support `1`. fn max_videos_per_call(&self) -> Option; - /// Generate an array of videos. + /// Start an asynchronous video generation and return an opaque operation + /// reference for [`Self::do_status`]. /// /// Naming: the `do_` prefix prevents accidental direct usage by users. - async fn do_generate(&self, options: &VideoCallOptions) -> Result; + async fn do_start(&self, options: &VideoCallOptions) + -> Result; + + /// Check the status of a generation started with [`Self::do_start`]. + /// + /// A task that failed server-side should be reported as a non-retryable + /// `Err`, not `Pending`, so the poll loop stops immediately. + async fn do_status( + &self, + operation: &serde_json::Value, + options: &VideoCallOptions, + ) -> Result; +} + +/// User-facing video generation with Core-owned retry, polling, and timeout. +/// +/// Orchestration (AI SDK `generate-video`): `do_start` is retried as one +/// unit, then `do_status` is polled — each poll retried independently — +/// against the same operation reference, so a transient poll failure never +/// re-creates the billed task. +/// +/// # Errors +/// +/// Returns the provider failure, retry exhaustion, poll or operation +/// timeout, or caller abort. +pub async fn generate_video( + model: &dyn VideoModel, + options: VideoCallOptions, +) -> Result { + let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?; + let abort_signal = options.abort_signal.clone(); + timeout::run( + start_and_poll(model, options), + abort_signal.as_ref(), + timeout, + ) + .await +} + +async fn start_and_poll( + model: &dyn VideoModel, + options: VideoCallOptions, +) -> Result { + let abort_signal = options.abort_signal.clone(); + let retries = retry::prepare_retries( + options.max_retries, + model.retry_config(), + abort_signal.clone(), + ); + + // `do_start` is billable: mint one idempotency key per logical start, + // OUTSIDE the retry closure, so providers that honor this header can + // deduplicate a replay where the first attempt succeeded but its response + // was lost. A caller-supplied key wins (AI SDK generate-video parity). + let mut start_options = options.clone(); + let headers = start_options.headers.get_or_insert_with(SharedHeaders::new); + if !headers + .keys() + .any(|k| k.eq_ignore_ascii_case("idempotency-key")) + { + headers.insert( + "idempotency-key".to_string(), + format!("aimux_vid_{:016x}", rand::random::()), + ); + } + + let start = retries.retry(|| model.do_start(&start_options)).await?; + + let mut poll = model.poll_config(); + if let Some(overrides) = options.poll { + if let Some(ms) = overrides.interval_ms { + poll.interval = Duration::from_millis(ms); + } + if let Some(ms) = overrides.timeout_ms { + poll.timeout = Duration::from_millis(ms); + } + } + let poll_started = Instant::now(); + loop { + let elapsed = poll_started.elapsed(); + if elapsed >= poll.timeout { + return Err(AiMuxError::Timeout(format!( + "Video generation timed out after {:?}.", + poll.timeout + ))); + } + retry::delay( + poll.interval.min(poll.timeout - elapsed), + abort_signal.as_ref(), + ) + .await?; + + if poll_started.elapsed() >= poll.timeout { + return Err(AiMuxError::Timeout(format!( + "Video generation timed out after {:?}.", + poll.timeout + ))); + } + + // Matching the AI SDK, the poll budget paces the loop between status + // checks; a hung status GET is already bounded by the per-exchange + // response guard in provider-utils, and the retry count is finite. + let status = retries + .retry(|| model.do_status(&start.operation, &options)) + .await?; + match status { + VideoOperationStatus::Pending => {} + VideoOperationStatus::Completed(mut result) => { + // Start-call warnings/metadata precede the completion's own. + let mut warnings = start.warnings; + warnings.append(&mut result.warnings); + result.warnings = warnings; + if let Some(start_meta) = start.provider_metadata { + let merged = result + .provider_metadata + .get_or_insert_with(Default::default); + for (provider, value) in start_meta { + merged.entry(provider).or_insert(value); + } + } + return Ok(result); + } + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + use std::sync::atomic::{AtomicU32, Ordering}; + + use super::*; + use crate::error::ApiCallError; + + /// `do_status` behavior per call index, cycled through in order. + enum StatusStep { + RetryableFailure, + Pending, + Complete, + } + + struct ScriptedVideoModel { + starts: AtomicU32, + statuses: AtomicU32, + script: Vec, + start_failures: u32, + start_idempotency_keys: Mutex>>, + status_idempotency_keys: Mutex>>, + } + + impl ScriptedVideoModel { + fn new(script: Vec) -> Self { + Self { + starts: AtomicU32::new(0), + statuses: AtomicU32::new(0), + script, + start_failures: 0, + start_idempotency_keys: Mutex::new(Vec::new()), + status_idempotency_keys: Mutex::new(Vec::new()), + } + } + + fn with_start_failures(mut self, failures: u32) -> Self { + self.start_failures = failures; + self + } + } + + fn idempotency_key(options: &VideoCallOptions) -> Option { + options.headers.as_ref().and_then(|headers| { + headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("idempotency-key")) + .map(|(_, value)| value.clone()) + }) + } + + #[async_trait] + impl VideoModel for ScriptedVideoModel { + fn provider(&self) -> &str { + "test" + } + + fn model_id(&self) -> &str { + "scripted" + } + + fn retry_config(&self) -> crate::retry::RetryConfig { + crate::retry::RetryConfig { + initial_delay: Duration::from_millis(1), + ..crate::retry::RetryConfig::default() + } + } + + fn max_videos_per_call(&self) -> Option { + Some(1) + } + + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { + let attempt = self.starts.fetch_add(1, Ordering::SeqCst); + let key = idempotency_key(options); + self.start_idempotency_keys + .lock() + .expect("start key capture lock should not be poisoned") + .push(key.clone()); + // The Core-minted idempotency key must reach the start request. + assert!( + key.is_some(), + "do_start should receive an idempotency-key header" + ); + if attempt < self.start_failures { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(503), + is_retryable: true, + ..ApiCallError::new( + "start unavailable", + "https://test/start", + serde_json::json!({}), + ) + }))); + } + Ok(VideoOperationStart { + operation: serde_json::json!({ "task_id": "t-1" }), + warnings: vec![Warning::Other { + message: "from start".to_string(), + }], + provider_metadata: None, + response: VideoResponse::default(), + }) + } + + async fn do_status( + &self, + operation: &serde_json::Value, + options: &VideoCallOptions, + ) -> Result { + assert_eq!(operation["task_id"], "t-1"); + self.status_idempotency_keys + .lock() + .expect("status key capture lock should not be poisoned") + .push(idempotency_key(options)); + let i = self.statuses.fetch_add(1, Ordering::SeqCst) as usize; + match self.script.get(i).unwrap_or(&StatusStep::Complete) { + StatusStep::RetryableFailure => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(429), + is_retryable: true, + ..ApiCallError::new("rate limited", "https://test", serde_json::json!({})) + }))), + StatusStep::Pending => Ok(VideoOperationStatus::Pending), + StatusStep::Complete => Ok(VideoOperationStatus::Completed(VideoResult { + videos: vec![VideoData::Url { + url: "https://cdn/video.mp4".to_string(), + media_type: "video/mp4".to_string(), + }], + warnings: vec![Warning::Other { + message: "from status".to_string(), + }], + provider_metadata: None, + response: VideoResponse::default(), + })), + } + } + } + + fn fast_poll_options() -> VideoCallOptions { + let mut options = VideoCallOptions::new("a cat"); + options.poll = Some(VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(2_000), + }); + options + } + + /// The core guarantee of the do_start/do_status split: a transient poll + /// failure retries the status check against the same operation and never + /// re-creates the billed task. + #[tokio::test] + async fn transient_status_failure_never_restarts_the_task() { + let model = ScriptedVideoModel::new(vec![ + StatusStep::Pending, + StatusStep::RetryableFailure, + StatusStep::Complete, + ]); + let result = generate_video(&model, fast_poll_options()).await.unwrap(); + assert_eq!(model.starts.load(Ordering::SeqCst), 1); + assert_eq!(model.statuses.load(Ordering::SeqCst), 3); + assert!( + model + .status_idempotency_keys + .lock() + .unwrap() + .iter() + .all(Option::is_none), + "the Core-minted start key must not be forwarded to status calls" + ); + assert_eq!(result.videos.len(), 1); + } + + #[tokio::test] + async fn start_retry_reuses_one_idempotency_key() { + let model = ScriptedVideoModel::new(vec![StatusStep::Complete]).with_start_failures(1); + let mut options = fast_poll_options(); + options.max_retries = Some(1); + + let result = generate_video(&model, options).await.unwrap(); + + assert_eq!(model.starts.load(Ordering::SeqCst), 2); + assert_eq!(result.videos.len(), 1); + let keys = model + .start_idempotency_keys + .lock() + .expect("start key capture lock should not be poisoned"); + assert_eq!(keys.len(), 2); + assert!(keys[0].is_some()); + assert_eq!(keys[0], keys[1]); + assert_eq!( + model + .status_idempotency_keys + .lock() + .expect("status key capture lock should not be poisoned") + .as_slice(), + &[None] + ); + } + + #[tokio::test] + async fn caller_idempotency_key_is_preserved_for_start_and_status() { + let model = ScriptedVideoModel::new(vec![StatusStep::Complete]); + let mut options = fast_poll_options(); + options + .headers + .get_or_insert_with(SharedHeaders::new) + .insert("Idempotency-Key".to_string(), "caller-key".to_string()); + + generate_video(&model, options).await.unwrap(); + + assert_eq!( + model + .start_idempotency_keys + .lock() + .expect("start key capture lock should not be poisoned") + .as_slice(), + &[Some("caller-key".to_string())] + ); + assert_eq!( + model + .status_idempotency_keys + .lock() + .expect("status key capture lock should not be poisoned") + .as_slice(), + &[Some("caller-key".to_string())] + ); + } + + #[tokio::test] + async fn start_warnings_precede_completion_warnings() { + let model = ScriptedVideoModel::new(vec![StatusStep::Complete]); + let result = generate_video(&model, fast_poll_options()).await.unwrap(); + let messages: Vec = result.warnings.iter().map(|w| format!("{w:?}")).collect(); + assert!(messages[0].contains("from start"), "{messages:?}"); + assert!(messages[1].contains("from status"), "{messages:?}"); + } + + #[tokio::test] + async fn poll_timeout_fails_a_generation_that_never_completes() { + let model = ScriptedVideoModel::new(vec![]); + // Script that always reports Pending. + struct AlwaysPending(ScriptedVideoModel); + #[async_trait] + impl VideoModel for AlwaysPending { + fn provider(&self) -> &str { + "test" + } + fn model_id(&self) -> &str { + "pending" + } + fn max_videos_per_call(&self) -> Option { + Some(1) + } + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { + self.0.do_start(options).await + } + async fn do_status( + &self, + _operation: &serde_json::Value, + _options: &VideoCallOptions, + ) -> Result { + Ok(VideoOperationStatus::Pending) + } + } + let model = AlwaysPending(model); + let mut options = VideoCallOptions::new("a cat"); + options.poll = Some(VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(20), + }); + let error = generate_video(&model, options).await.unwrap_err(); + assert!( + matches!(error, AiMuxError::Timeout(_)), + "expected Timeout, got {error:?}" + ); + } + + #[tokio::test] + async fn poll_timeout_reached_during_delay_skips_status_call() { + let model = ScriptedVideoModel::new(vec![StatusStep::Complete]); + let mut options = VideoCallOptions::new("a cat"); + options.poll = Some(VideoPollOptions { + interval_ms: Some(100), + timeout_ms: Some(10), + }); + + let error = generate_video(&model, options).await.unwrap_err(); + + assert!(matches!(error, AiMuxError::Timeout(_))); + assert_eq!(model.starts.load(Ordering::SeqCst), 1); + assert_eq!(model.statuses.load(Ordering::SeqCst), 0); + assert!( + model + .status_idempotency_keys + .lock() + .expect("status key capture lock should not be poisoned") + .is_empty() + ); + } } diff --git a/aimux-core/tests/error_value_golden_test.rs b/aimux-core/tests/error_value_golden_test.rs index 92bb608b..b3161f78 100644 --- a/aimux-core/tests/error_value_golden_test.rs +++ b/aimux-core/tests/error_value_golden_test.rs @@ -5,7 +5,11 @@ //! cross-language contract: a field rename, a variant rename or a shape change //! breaks every binding, so it must break this test first. -use aimux_core::{AiMuxError, ApiCallError}; +use aimux_core::{AiMuxError, ApiCallError, RetryError, RetryErrorReason}; + +fn api_error(message: &str) -> ApiCallError { + ApiCallError::new(message, "https://example.test/v1", serde_json::json!({})) +} fn golden(err: &AiMuxError, expected: &str) { let json = serde_json::to_string(err).unwrap(); @@ -16,55 +20,53 @@ fn golden(err: &AiMuxError, expected: &str) { assert_eq!(back.to_string(), err.to_string()); } -/// `ApiCall` carries the full `ApiCallError` field set (unboxed, like -/// async-openai's `ApiError`); the classification is the `status_code` +/// `ApiCall` carries the full `ApiCallError` field set. The Rust payload is +/// boxed only to keep the enum compact; the classification is `status_code` /// field and the retry verdict the stored `is_retryable`. #[test] fn error_value_snapshots_api_call_shapes() { golden( - &AiMuxError::ApiCall(ApiCallError { + &AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(500), provider_code: Some("server_error".into()), - message: "boom".into(), - response_body: None, is_retryable: true, - ..Default::default() - }), - r#"{"ApiCall":{"status_code":500,"provider_code":"server_error","message":"boom","response_body":null,"request_id":null,"retry_after_ms":null,"is_retryable":true}}"#, + ..api_error("boom") + })), + r#"{"ApiCall":{"url":"https://example.test/v1","request_body_values":{},"status_code":500,"provider_code":"server_error","message":"boom","response_body":null,"response_headers":null,"data":null,"is_retryable":true}}"#, ); golden( - &AiMuxError::ApiCall(ApiCallError { + &AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(500), provider_code: Some("server_error".into()), - message: "boom".into(), response_body: Some(r#"{"error":{"message":"boom","type":"server_error"}}"#.into()), is_retryable: true, - ..Default::default() - }), - r#"{"ApiCall":{"status_code":500,"provider_code":"server_error","message":"boom","response_body":"{\"error\":{\"message\":\"boom\",\"type\":\"server_error\"}}","request_id":null,"retry_after_ms":null,"is_retryable":true}}"#, + ..api_error("boom") + })), + r#"{"ApiCall":{"url":"https://example.test/v1","request_body_values":{},"status_code":500,"provider_code":"server_error","message":"boom","response_body":"{\"error\":{\"message\":\"boom\",\"type\":\"server_error\"}}","response_headers":null,"data":null,"is_retryable":true}}"#, ); // A transport failure (no response arrived): no status, retryable — // exactly the AI SDK's handleFetchError shape. golden( - &AiMuxError::ApiCall(ApiCallError { - message: "connection reset".into(), + &AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }), - r#"{"ApiCall":{"status_code":null,"provider_code":null,"message":"connection reset","response_body":null,"request_id":null,"retry_after_ms":null,"is_retryable":true}}"#, + ..api_error("connection reset") + })), + r#"{"ApiCall":{"url":"https://example.test/v1","request_body_values":{},"status_code":null,"provider_code":null,"message":"connection reset","response_body":null,"response_headers":null,"data":null,"is_retryable":true}}"#, ); // A 429 is an ApiCall error whose classification is the status field — - // there is no RateLimited variant; the hint rides `retry_after_ms`. + // there is no RateLimited variant; the hint remains in response headers. golden( - &AiMuxError::ApiCall(ApiCallError { + &AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(429), provider_code: Some("rate_limit_exceeded".into()), - message: "slow down".into(), - retry_after_ms: Some(2500), + response_headers: Some(std::collections::HashMap::from([( + "retry-after-ms".into(), + "2500".into(), + )])), is_retryable: true, - ..Default::default() - }), - r#"{"ApiCall":{"status_code":429,"provider_code":"rate_limit_exceeded","message":"slow down","response_body":null,"request_id":null,"retry_after_ms":2500,"is_retryable":true}}"#, + ..api_error("slow down") + })), + r#"{"ApiCall":{"url":"https://example.test/v1","request_body_values":{},"status_code":429,"provider_code":"rate_limit_exceeded","message":"slow down","response_body":null,"response_headers":{"retry-after-ms":"2500"},"data":null,"is_retryable":true}}"#, ); golden( &AiMuxError::TokenExpired("expired".into()), @@ -72,6 +74,26 @@ fn error_value_snapshots_api_call_shapes() { ); } +#[test] +fn error_value_snapshot_retry_history() { + golden( + &AiMuxError::Retry(RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors: vec![ + AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..api_error("first") + })), + AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..api_error("second") + })), + ], + }), + r#"{"Retry":{"reason":"maxRetriesExceeded","errors":[{"ApiCall":{"url":"https://example.test/v1","request_body_values":{},"status_code":null,"provider_code":null,"message":"first","response_body":null,"response_headers":null,"data":null,"is_retryable":true}},{"ApiCall":{"url":"https://example.test/v1","request_body_values":{},"status_code":null,"provider_code":null,"message":"second","response_body":null,"response_headers":null,"data":null,"is_retryable":true}}]}}"#, + ); +} + #[test] fn error_value_snapshots_plain_variants() { for (err, expected) in [ @@ -120,7 +142,10 @@ fn error_value_snapshots_plain_variants() { AiMuxError::Timeout("total timeout".into()), r#"{"Timeout":"total timeout"}"#, ), - (AiMuxError::Aborted, r#""Aborted""#), + ( + AiMuxError::Aborted("request aborted".into()), + r#"{"Aborted":"request aborted"}"#, + ), (AiMuxError::Other("misc".into()), r#"{"Other":"misc"}"#), ] { golden(&err, expected); @@ -128,20 +153,17 @@ fn error_value_snapshots_plain_variants() { } /// The variant set is a cross-language contract of its own: bindings switch on -/// the wire tag. Adding or removing one is a breaking change (13 variants — +/// the wire tag. Adding or removing one is a breaking change (14 variants — /// the per-status avatars `Auth`/`ModelNotFound`/`RateLimited` are gone, and /// `Http`/`Provider` folded into `ApiCall`: a failed exchange is an `ApiCall` /// error classified by `status_code`, transport failures included). #[test] -fn variant_set_is_exactly_thirteen() { +fn variant_set_is_exactly_fourteen() { let all = [ - AiMuxError::ApiCall(ApiCallError { - message: "x".into(), - ..Default::default() - }), - AiMuxError::ApiCall(ApiCallError { - is_retryable: true, - ..Default::default() + AiMuxError::ApiCall(Box::new(api_error("x"))), + AiMuxError::Retry(RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors: vec![AiMuxError::ApiCall(Box::new(api_error("x")))], }), AiMuxError::JsonParse("x".into()), AiMuxError::InvalidResponseData("x".into()), @@ -158,7 +180,7 @@ fn variant_set_is_exactly_thirteen() { provider_id: "x".into(), }, AiMuxError::Timeout("x".into()), - AiMuxError::Aborted, + AiMuxError::Aborted("request aborted".into()), AiMuxError::Other("x".into()), ]; // The wire tag IS the variant name (externally-tagged serde JSON). @@ -184,6 +206,7 @@ fn variant_set_is_exactly_thirteen() { "NoSuchModel", "NoSuchProvider", "Other", + "Retry", "Timeout", "TokenExpired", "Tool", @@ -193,21 +216,12 @@ fn variant_set_is_exactly_thirteen() { ); } -/// Field additions must stay *deserialization*-compatible: a payload written -/// before the structured fields existed still loads (the `#[serde(default)]` -/// contract, M6). Serialization always emits the full current shape. +/// Request context is required. Payloads from the pre-context schema are +/// deliberately rejected instead of silently fabricating an empty URL/body. #[test] -fn structured_fields_deserialize_from_pre_field_payloads() { +fn api_call_requires_request_context() { let old = r#"{"ApiCall":{"message":"boom"}}"#; - let err: AiMuxError = serde_json::from_str(old).unwrap(); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.status_code, None); - assert_eq!(detail.provider_code, None); - assert_eq!(detail.message, "boom"); - assert!(!detail.is_retryable); - assert_eq!(err.status_code(), None); + assert!(serde_json::from_str::(old).is_err()); // The removed variants no longer deserialize — a deliberate breaking // change pinned here so it cannot happen silently a second time. @@ -239,11 +253,10 @@ fn structured_fields_deserialize_from_pre_field_payloads() { /// unchanged while no consumer has to parse it back out (H1). #[test] fn status_lives_in_the_field_and_display_composes_it() { - let err = AiMuxError::ApiCall(ApiCallError { + let err = AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(429), - message: "quota exceeded".into(), - ..Default::default() - }); + ..api_error("quota exceeded") + })); assert_eq!(err.status_code(), Some(429)); let AiMuxError::ApiCall(ref detail) = err else { panic!("expected ApiCall, got {err:?}") @@ -255,22 +268,18 @@ fn status_lives_in_the_field_and_display_composes_it() { assert_eq!(err.to_string(), "API call error: HTTP 429: quota exceeded"); // Without a status there is nothing to compose. - let err = AiMuxError::ApiCall(ApiCallError { - message: "plain failure".into(), - ..Default::default() - }); + let err = AiMuxError::ApiCall(Box::new(api_error("plain failure"))); assert_eq!(err.to_string(), "API call error: plain failure"); } /// `provider_code` is machine-readable and stays out of the human text. #[test] fn provider_code_is_readable_and_not_displayed() { - let err = AiMuxError::ApiCall(ApiCallError { + let err = AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(400), provider_code: Some("invalid_request".into()), - message: "bad input".into(), - ..Default::default() - }); + ..api_error("bad input") + })); let AiMuxError::ApiCall(ref detail) = err else { panic!("expected ApiCall, got {err:?}") }; @@ -283,19 +292,14 @@ fn provider_code_is_readable_and_not_displayed() { #[test] fn display_strings_are_unchanged_by_the_field_shape() { assert_eq!( - AiMuxError::ApiCall(ApiCallError { - message: "boom".into(), - ..Default::default() - }) - .to_string(), + AiMuxError::ApiCall(Box::new(api_error("boom"))).to_string(), "API call error: boom" ); assert_eq!( - AiMuxError::ApiCall(ApiCallError { - message: "reset".into(), + AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }) + ..api_error("reset") + })) .to_string(), "API call error: reset" ); @@ -306,7 +310,7 @@ fn display_strings_are_unchanged_by_the_field_shape() { } /// `AiMuxError` rides in every `Result`. `ApiCall` carries its -/// detail inline (unboxed, async-openai `ApiError` style); this guard pins the +/// detail boxed; this guard pins the /// size so growth is a deliberate decision, and keeps it under clippy's /// `result_large_err` threshold (128 bytes). #[test] diff --git a/aimux-core/tests/recording_e2e_test.rs b/aimux-core/tests/recording_e2e_test.rs index 989ac30f..28e46cf8 100644 --- a/aimux-core/tests/recording_e2e_test.rs +++ b/aimux-core/tests/recording_e2e_test.rs @@ -39,12 +39,13 @@ impl LanguageModel for EchoModel { _options: &CallOptions, ) -> Result { if self.gen_fails { - return Err(aimux_core::error::AiMuxError::ApiCall( - aimux_core::error::ApiCallError { - message: "gen boom".into(), - ..Default::default() - }, - )); + return Err(aimux_core::error::AiMuxError::ApiCall(Box::new( + aimux_core::error::ApiCallError::new( + "gen boom", + "https://example.test/v1", + serde_json::json!({}), + ), + ))); } Ok(GenerateResult { content: vec![GenerateContent::Text { diff --git a/aimux-core/tests/stream_timeout_matrix_test.rs b/aimux-core/tests/stream_timeout_matrix_test.rs new file mode 100644 index 00000000..9d839ab8 --- /dev/null +++ b/aimux-core/tests/stream_timeout_matrix_test.rs @@ -0,0 +1,150 @@ +//! RFC-0031 §13 acceptance: chunk-timer semantics (item 17) and +//! no-lingering-timer-state after drop (item 21b). +//! +//! Paused tokio time makes every assertion exact: auto-advance jumps to the +//! earliest armed deadline, so elapsed times distinguish "reset" from +//! "not reset" deterministically. + +use std::time::Duration; + +use aimux_core::AbortSignal; +use aimux_core::error::AiMuxError; +use aimux_core::generate::{GenerateTextOptions, stream_text}; +use aimux_core::language_model::LanguageModel; +use aimux_core::options::{CallOptions, TimeoutConfiguration}; +use aimux_core::result::{GenerateResult, StreamResult}; +use aimux_core::stream_part::StreamPart; +use async_trait::async_trait; +use futures::StreamExt; + +/// Streams: TextDelta at t=0, then after 600ms one more part (semantic or +/// not), then pends forever. The follow-up part is the probe: only semantic +/// output may reset the chunk timer. +struct ProbeModel { + second_part_is_output: bool, +} + +fn delta(text: &str) -> StreamPart { + StreamPart::TextDelta { + id: "1".into(), + delta: text.into(), + provider_metadata: None, + } +} + +#[async_trait] +impl LanguageModel for ProbeModel { + fn provider(&self) -> &str { + "mock-probe" + } + + fn model_id(&self) -> &str { + "probe-1" + } + + async fn do_generate(&self, _options: &CallOptions) -> Result { + unimplemented!("streaming-only mock") + } + + async fn do_stream(&self, _options: &CallOptions) -> Result { + let second_is_output = self.second_part_is_output; + Ok(StreamResult { + stream: Box::pin(async_stream::stream! { + yield Ok(delta("first")); + tokio::time::sleep(Duration::from_millis(600)).await; + if second_is_output { + yield Ok(delta("second")); + } else { + // Non-semantic per the AI SDK's isOutputChunk: must NOT + // reset the chunk timer. + yield Ok(StreamPart::TextStart { + id: "1".into(), + provider_metadata: None, + }); + } + futures::future::pending::<()>().await; + }), + request_body: None, + response_headers: None, + }) + } +} + +fn chunk_options() -> GenerateTextOptions { + GenerateTextOptions { + max_retries: Some(0), + timeout: Some(TimeoutConfiguration { + chunk_ms: Some(1_000), + ..Default::default() + }), + ..Default::default() + } +} + +/// Drives the probe stream to the chunk timeout and returns when it fired. +async fn time_to_chunk_timeout(second_part_is_output: bool) -> Duration { + let model = ProbeModel { + second_part_is_output, + }; + let mut result = stream_text(&model, "hello", chunk_options()).await.unwrap(); + let start = tokio::time::Instant::now(); + loop { + match result.stream.next().await.expect("stream must not end") { + Ok(_) => {} + Err(AiMuxError::Timeout(message)) => { + assert_eq!(message, "Chunk timeout of 1000ms exceeded"); + return start.elapsed(); + } + Err(other) => panic!("unexpected error: {other}"), + } + } +} + +#[tokio::test(start_paused = true)] +async fn non_semantic_part_does_not_reset_the_chunk_timer() { + // TextStart at t=600 must not push the deadline: timeout fires 1000ms + // after the only semantic output at t=0, not 1600ms. + assert_eq!( + time_to_chunk_timeout(false).await, + Duration::from_millis(1_000) + ); +} + +#[tokio::test(start_paused = true)] +async fn semantic_output_resets_the_chunk_timer() { + // The second TextDelta at t=600 re-arms the deadline to t=1600. + assert_eq!( + time_to_chunk_timeout(true).await, + Duration::from_millis(1_600) + ); +} + +#[tokio::test(start_paused = true)] +async fn dropped_operation_leaves_no_armed_timer_state() { + // §8.0/§13-21b: deadlines live inside the operation future, nothing is + // spawned. Dropping the stream mid-flight must leave no timer that could + // later fire and mutate the caller's signal (the rejected abort_after + // design would trip this). + let caller = AbortSignal::new(); + let model = ProbeModel { + second_part_is_output: false, + }; + let options = GenerateTextOptions { + max_retries: Some(0), + abort_signal: Some(caller.clone()), + timeout: Some(TimeoutConfiguration { + total_ms: Some(5_000), + chunk_ms: Some(1_000), + ..Default::default() + }), + ..Default::default() + }; + let mut result = stream_text(&model, "hello", options).await.unwrap(); + let first = result.stream.next().await.expect("first part"); + assert!(first.is_ok()); + drop(result); + + // Sail far past every configured deadline. + tokio::time::sleep(Duration::from_secs(60)).await; + assert!(!caller.is_aborted()); +} From 1ac7b54c0971a3f34f7d81da6feb93cb6f6b10a5 Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Mon, 24 Aug 2026 22:29:11 +0800 Subject: [PATCH 02/24] feat(provider-utils)!: single-exchange helpers and response handlers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace send/send_timed/send_stream_timed/send_with_retry_raw and the JSON-path ErrorStructure with AI SDK-shaped single-exchange helpers (post_json_to_api, post_form_data_to_api, post_to_api, get_from_api) dispatching to typed success/failure response handlers. Every call performs exactly one fetch attempt and records exactly one exchange; retry, operation timeouts, and backoff belong to Core. Non-streaming exchanges keep a per-exchange 30s whole-response guard; streaming handlers are exempt because their body outlives the call. Failure handling keeps the provider's error through every path: error bodies are read best-effort with a 64 KiB public cap (enforced after lossy UTF-8 decoding) and a 1 MiB parse bound so oversize-but-valid error JSON still reaches the provider's mapper; a handler failure on a 429/503 keeps the status's retryability; a transport failure mid-body reports no HTTP status; in-band stream errors built via stream_error_api_call are redacted and URL-sanitized at the shared entry. WebSocket handshake rejections classify by HTTP status like any other exchange. The shared HTTP client is keyed by runtime — pooled connections are driven by tasks on the runtime that made the request, so a process-wide client let a finished runtime leave dead pooled connections behind (the old HTTP-layer retry masked this as ~1-in-6 flaky wiremock failures). The cache is capped, evicting wholesale on overflow since runtimes leave no drop signal. --- aimux-provider-utils/Cargo.toml | 8 +- .../src/extract_response_headers.rs | 28 + aimux-provider-utils/src/get_from_api.rs | 25 + .../src/handle_fetch_error.rs | 21 + aimux-provider-utils/src/http.rs | 2352 ++++------------- aimux-provider-utils/src/lib.rs | 35 +- aimux-provider-utils/src/logging.rs | 89 +- aimux-provider-utils/src/multipart.rs | 27 +- aimux-provider-utils/src/post_to_api.rs | 261 ++ .../src/read_response_with_size_limit.rs | 134 + aimux-provider-utils/src/response.rs | 166 -- aimux-provider-utils/src/response_handler.rs | 425 +++ aimux-provider-utils/src/retry.rs | 325 +-- aimux-provider-utils/src/ws.rs | 84 +- .../tests/api_helpers_test.rs | 330 +++ .../tests/http_recording_test.rs | 660 ++--- aimux-provider-utils/tests/http_send_test.rs | 615 ----- aimux-provider-utils/tests/logging_test.rs | 229 +- .../tests/response_handler.rs | 138 - .../tests/retry_with_exponential_backoff.rs | 492 ---- .../tests/structured_error_fields.rs | 307 --- 21 files changed, 2136 insertions(+), 4615 deletions(-) create mode 100644 aimux-provider-utils/src/extract_response_headers.rs create mode 100644 aimux-provider-utils/src/get_from_api.rs create mode 100644 aimux-provider-utils/src/handle_fetch_error.rs create mode 100644 aimux-provider-utils/src/post_to_api.rs create mode 100644 aimux-provider-utils/src/read_response_with_size_limit.rs delete mode 100644 aimux-provider-utils/src/response.rs create mode 100644 aimux-provider-utils/src/response_handler.rs create mode 100644 aimux-provider-utils/tests/api_helpers_test.rs delete mode 100644 aimux-provider-utils/tests/http_send_test.rs delete mode 100644 aimux-provider-utils/tests/response_handler.rs delete mode 100644 aimux-provider-utils/tests/retry_with_exponential_backoff.rs delete mode 100644 aimux-provider-utils/tests/structured_error_fields.rs diff --git a/aimux-provider-utils/Cargo.toml b/aimux-provider-utils/Cargo.toml index 2123d518..50f05eb8 100644 --- a/aimux-provider-utils/Cargo.toml +++ b/aimux-provider-utils/Cargo.toml @@ -4,7 +4,7 @@ version.workspace = true edition.workspace = true license.workspace = true description = "Shared utilities for aimux provider implementations" -keywords = ["llm", "http", "retry", "backoff"] +keywords = ["llm", "http", "ai", "api"] categories = ["api-bindings", "asynchronous"] documentation = "https://docs.rs/aimux-provider-utils" @@ -13,18 +13,17 @@ aimux-core = { workspace = true } aimux-stream = { workspace = true } reqwest = { workspace = true } +http = "1" serde = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } -thiserror = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } bytes = { workspace = true } futures = { workspace = true } +async-stream = { workspace = true } url = { workspace = true } -httpdate = { workspace = true } chrono = "0.4" -rand = { workspace = true } # WebSocket client (RFC-0028 realtime transcription). Optional: enabled via # the `ws` feature so consumers without realtime needs compile it out. @@ -37,7 +36,6 @@ ws = ["dep:tokio-tungstenite"] aimux-core = { workspace = true } tokio = { workspace = true, features = ["test-util", "macros", "rt", "time"] } serial_test = { workspace = true } -httpdate = { workspace = true } wiremock = "0.6" [lints] diff --git a/aimux-provider-utils/src/extract_response_headers.rs b/aimux-provider-utils/src/extract_response_headers.rs new file mode 100644 index 00000000..c8243574 --- /dev/null +++ b/aimux-provider-utils/src/extract_response_headers.rs @@ -0,0 +1,28 @@ +//! Response header extraction for public error context. + +use std::collections::HashMap; + +/// Extract response headers as ordered pairs while redacting sensitive values. +/// Recording and public error context share this one policy. +#[must_use] +pub fn extract_response_header_pairs( + headers: &reqwest::header::HeaderMap, +) -> Vec<(String, String)> { + headers + .iter() + .map(|(name, value)| { + let value = if aimux_core::recording::is_sensitive_key(name.as_str()) { + "[REDACTED]".to_string() + } else { + value.to_str().unwrap_or_default().to_string() + }; + (name.as_str().to_ascii_lowercase(), value) + }) + .collect() +} + +/// Extract response headers while redacting sensitive values. +#[must_use] +pub fn extract_response_headers(headers: &reqwest::header::HeaderMap) -> HashMap { + extract_response_header_pairs(headers).into_iter().collect() +} diff --git a/aimux-provider-utils/src/get_from_api.rs b/aimux-provider-utils/src/get_from_api.rs new file mode 100644 index 00000000..4684303c --- /dev/null +++ b/aimux-provider-utils/src/get_from_api.rs @@ -0,0 +1,25 @@ +//! Single-exchange GET helper aligned with AI SDK. + +use aimux_core::AiMuxError; + +use crate::http::{HttpBody, HttpMethod, HttpRequest}; +use crate::response_handler::{ResponseHandler, ResponseHandlerOutput}; + +/// GET one URL and dispatch its response to exactly one response handler. +/// +/// # Errors +/// +/// Returns transport, response-handler, or caller-abort failures. +pub async fn get_from_api( + request: HttpRequest, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + crate::post_to_api::call_to_api( + request.prepare(HttpMethod::Get, HttpBody::Empty), + serde_json::json!({}), + successful_response_handler, + failed_response_handler, + ) + .await +} diff --git a/aimux-provider-utils/src/handle_fetch_error.rs b/aimux-provider-utils/src/handle_fetch_error.rs new file mode 100644 index 00000000..4a935e32 --- /dev/null +++ b/aimux-provider-utils/src/handle_fetch_error.rs @@ -0,0 +1,21 @@ +//! Fetch error normalization. + +use aimux_core::AiMuxError; + +/// Attach request context to transport failures without changing other errors. +#[must_use] +pub fn handle_fetch_error( + error: AiMuxError, + url: &str, + request_body_values: &serde_json::Value, +) -> AiMuxError { + match error { + AiMuxError::ApiCall(mut detail) => { + detail.url = url.to_string(); + detail.request_body_values = request_body_values.clone(); + AiMuxError::ApiCall(detail) + } + AiMuxError::Aborted(_) | AiMuxError::Timeout(_) => error, + other => other, + } +} diff --git a/aimux-provider-utils/src/http.rs b/aimux-provider-utils/src/http.rs index 9d717f63..7e1943a0 100644 --- a/aimux-provider-utils/src/http.rs +++ b/aimux-provider-utils/src/http.rs @@ -1,691 +1,499 @@ -//! HTTP 层 — 共享 client + 请求/响应抽象。 +//! One shared HTTP client and one-exchange request transport. //! -//! 本模块是 aimux 的 HTTP 边界：provider 给出纯数据的 [`HttpRequest`]，本 -//! 模块负责执行（含连接池、超时、retry），返回 [`HttpResponse`]（非流式）或 -//! [`HttpStreamResponse`]（流式）。**reqwest 类型完全不外泄到 provider**—— -//! provider 不持有 `Client`、不构造 `RequestBuilder`、不碰 `reqwest::Response`。 -//! -//! 三个职责： -//! - **连接池**：`shared_client()` / `shared_streaming_client()` 返回共享单例 -//! （`Result<&'static Client, AiMuxError>`——构建失败粘性返回错误），TLS 会话 -//! 与连接池全仓复用（RFC-0009 §4.1）。替代散落各处的 `Client::new()`。 -//! - **超时**：非流式带 30s 整体超时；流式禁用整体超时（LLM 流式时长取决于 -//! 生成长度，固定超时会误杀长生成，RFC-0009 §4.3）。 -//! - **retry**：408/409/429/5xx 重试 + Full Jitter 退避（RFC-0009 §4.2）。retry 是本 -//! 模块内部逻辑——provider 不感知重试发生。 +//! Retry and operation/stream timeouts belong to `aimux-core`. This module +//! sends exactly one request and leaves successful/failed body interpretation +//! to the response handlers selected by the provider. use std::collections::HashMap; use std::pin::Pin; -use std::sync::OnceLock; +use std::sync::{Mutex, OnceLock}; use std::task::{Context, Poll}; -use std::time::{Duration, Instant, SystemTime}; +use std::time::{Duration, Instant}; use bytes::Bytes; -use futures::stream::{BoxStream, Stream, StreamExt}; +use futures::{Stream, StreamExt, stream::BoxStream}; use reqwest::Client; use serde::{Deserialize, Serialize}; -use aimux_core::AiMuxError; -use aimux_core::ApiCallError; -use aimux_core::options::TimeoutConfiguration; -use aimux_core::shared::AbortSignal; - -use crate::logging::{auto_init_from_env, body_logging_enabled, redact_body}; -use crate::response::{ErrorStructure, parse_provider_error}; -use crate::retry::{RetryConfig, get_retry_delay_ms_with_jitter, parse_retry_after}; - -// ════════════════════════════════════════════════════════════════════════════ -// 连接池 & 超时配置 -// ════════════════════════════════════════════════════════════════════════════ - -/// 连接池配置（参考 catcher `PoolConfig` 字段设计）。 -#[derive(Debug, Clone)] -pub struct PoolConfig { - /// 每个主机的最大空闲连接数。catcher 默认 10。 - pub max_idle_per_host: usize, - /// 空闲连接超时（秒）。catcher 默认 30 — 防 retry 复用死连接。 - pub idle_timeout_secs: u64, - /// 是否启用 TCP keepalive。 - pub keep_alive: bool, - /// TCP keepalive 间隔（秒）。catcher 默认 20 — 更快发现死连接。 - pub keep_alive_interval_secs: u64, -} - -impl Default for PoolConfig { - fn default() -> Self { - Self { - max_idle_per_host: 10, - idle_timeout_secs: 30, - keep_alive: true, - keep_alive_interval_secs: 20, - } - } -} - -/// 超时配置（参考 catcher `HttpClientConfig` 的两个超时字段）。 -#[derive(Debug, Clone, Copy)] -pub struct TimeoutConfig { - /// 建连超时（毫秒）。catcher 默认 10_000。 - pub connect_timeout_ms: u64, - /// 整体响应超时（毫秒）。传 0 禁用整体超时（流式用）。 - /// catcher 默认 30_000。 - pub response_timeout_ms: u64, -} - -impl Default for TimeoutConfig { - fn default() -> Self { - Self { - connect_timeout_ms: 10_000, - response_timeout_ms: 30_000, - } - } -} +use aimux_core::recording::{ + HttpExchange, HttpRecord, RecordingContext, ResponseRecord, TimingRecord, +}; +use aimux_core::{AiMuxError, ApiCallError}; -impl TimeoutConfig { - /// 流式超时配置：仅 connect timeout，禁用整体超时。 - /// - /// LLM 流式时长取决于生成长度 / max_tokens，固定整体超时会误杀长 - /// 生成请求（RFC-0009 §4.3）。仅保留 connect timeout 守护建连阶段。 - #[must_use] - pub fn streaming() -> Self { - Self { - connect_timeout_ms: 10_000, - response_timeout_ms: 0, - } - } -} +use crate::logging::{body_logging_enabled, redact_body}; -/// Proxy configuration (M6, RFC-0016). All fields optional; when all `None`, -/// behaviour is unchanged (relies on reqwest's automatic `HTTP_PROXY` / -/// `HTTPS_PROXY` / `ALL_PROXY` / `NO_PROXY` env var support). -/// -/// Set via [`init_proxy`] **before** the first `generate_text` / `stream_text` -/// call (the shared client is lazily initialised on first use and locked for -/// the process lifetime). +/// Process-wide proxy configuration, fixed before the shared client is built. #[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ProxyConfig { - /// HTTP proxy URL (e.g. `"http://proxy:8080"`). pub http_url: Option, - /// HTTPS proxy URL. pub https_url: Option, - /// Proxy URL for all protocols (applied if `http_url` / `https_url` are unset). pub all_url: Option, - /// No-proxy whitelist: comma-separated host patterns to bypass the proxy - /// (e.g. `"localhost,127.0.0.1,.internal.team"`). Mirrors reqwest's `NoProxy::from_string`. pub no_proxy: Option, } -/// Global proxy config, read once when the shared client is first built. static GLOBAL_PROXY: OnceLock = OnceLock::new(); - -/// Set the global proxy configuration (M6). Must be called before the first -/// `generate_text` / `stream_text` call; a no-op if the shared client is -/// already initialised (returns `false`). -/// -/// `Default::default()` (all `None`) resets to env-proxy-only behaviour. +// One client (and so one connection pool) PER RUNTIME, not per process: +// pooled connections are driven by tasks spawned onto the runtime that +// made the request, so when that runtime shuts down its pooled +// connections become unusable while staying checked in — and the OS can +// recycle their ports to fresh servers, handing later requests a dead +// connection. Production processes run a single runtime and still get +// exactly one client. Runtimes leave no drop signal, so dead entries are +// undetectable; the map is instead bounded by `SHARED_CLIENT_CAP` — see +// `shared_client`. +static SHARED: OnceLock, Client>>> = OnceLock::new(); + +// Hosts that churn short-lived runtimes (a test binary, an FFI embedder +// creating a runtime per call) would otherwise grow the map without bound. +// Eviction is safe: correctness only requires never *reusing* a dead +// runtime's pool, and an evicted live runtime simply rebuilds a fresh +// client on its next request. +const SHARED_CLIENT_CAP: usize = 8; + +/// Set proxy configuration before the first HTTP operation. pub fn init_proxy(config: ProxyConfig) -> bool { GLOBAL_PROXY.set(config).is_ok() } -/// Read the global proxy config (or a default if unset). fn global_proxy() -> ProxyConfig { GLOBAL_PROXY.get().cloned().unwrap_or_default() } -/// 全局共享的 reqwest::Client（非流式，带 30s 整体超时）。构建失败（如 TLS -/// 后端初始化失败、系统资源耗尽）以错误字符串粘性保存——之后每次访问都返回 -/// 同一错误，而非在 release `panic=abort` 下杀死宿主进程（issue #115）。 -static SHARED: OnceLock> = OnceLock::new(); - -/// 全局共享的流式 reqwest::Client（无整体超时）。同上的粘性错误语义。 -static SHARED_STREAMING: OnceLock> = OnceLock::new(); - -/// 把构建失败映射为可诊断错误（无 HTTP 交换，status 留空、不可重试）。 -fn client_init_error(source: &str) -> AiMuxError { - AiMuxError::ApiCall(ApiCallError { - message: format!("shared HTTP client initialization failed: {source}"), - ..Default::default() - }) -} - -/// 获取（或惰性初始化）共享 reqwest Client（带 30s 整体超时，非流式用）。 -/// -/// 成功返回 `&'static Client`——provider **拿引用即用**，不持有、不 clone、 -/// 不传参；首次构建失败（及之后的每次调用）返回 [`AiMuxError::ApiCall`] -/// （message 带 "client initialization failed"，status 为 None）。 -/// -/// # Errors -/// -/// Returns a sticky, non-retryable [`AiMuxError::ApiCall`] (no HTTP status) -/// when the shared client could not be built, e.g. TLS backend or resource -/// initialization failures in restricted environments. -pub fn shared_client() -> Result<&'static Client, AiMuxError> { - SHARED - .get_or_init(|| { - build_client( - PoolConfig::default(), - TimeoutConfig::default(), - global_proxy(), - ) - }) - .as_ref() - .map_err(|e| client_init_error(e)) -} - -/// 获取（或惰性初始化）流式共享 reqwest Client（无整体超时，流式用）。 +/// Return the single shared client. It has a connect timeout but no +/// client-wide response timeout: the shared client also serves streaming +/// exchanges, so the 30s non-streaming response bound lives per-exchange in +/// `call_to_api`, and Core owns the operation deadline. /// /// # Errors /// -/// Same sticky-failure semantics as [`shared_client`]: a non-retryable -/// [`AiMuxError::ApiCall`] without an HTTP status when the build failed. -pub fn shared_streaming_client() -> Result<&'static Client, AiMuxError> { - SHARED_STREAMING - .get_or_init(|| { - build_client( - PoolConfig::default(), - TimeoutConfig::streaming(), - global_proxy(), - ) - }) - .as_ref() - .map_err(|e| client_init_error(e)) -} - -/// 用给定配置构建一个 reqwest Client。构建失败返回错误字符串（reqwest 仅 -/// 提供 Display），由调用方决定映射——不再 `expect`（issue #115：受限环境 -/// 下 TLS/资源初始化失败不应在 panic=abort 产物中直接终止宿主进程）。 -fn build_client( - pool: PoolConfig, - timeout: TimeoutConfig, - proxy: ProxyConfig, -) -> Result { - let mut b = Client::builder() - .connect_timeout(Duration::from_millis(timeout.connect_timeout_ms)) - .pool_max_idle_per_host(pool.max_idle_per_host) - .pool_idle_timeout(Some(Duration::from_secs(pool.idle_timeout_secs))); - if pool.keep_alive { - b = b.tcp_keepalive(Some(Duration::from_secs(pool.keep_alive_interval_secs))); - } - if timeout.response_timeout_ms > 0 { - b = b.timeout(Duration::from_millis(timeout.response_timeout_ms)); - } - b = apply_proxy(b, &proxy); - b.build().map_err(|e| e.to_string()) -} - -/// Apply proxy configuration to a reqwest client builder (by-value chain). -fn apply_proxy(b: reqwest::ClientBuilder, proxy: &ProxyConfig) -> reqwest::ClientBuilder { - use reqwest::Proxy as ReqwestProxy; - // Determine effective proxy URLs (all_url is a fallback). +/// Returns an initialization error if the shared client cannot be built. +pub fn shared_client() -> Result { + let key = tokio::runtime::Handle::try_current().ok().map(|h| h.id()); + let mut clients = SHARED + .get_or_init(|| Mutex::new(HashMap::new())) + .lock() + .expect("shared HTTP client mutex poisoned"); + if let Some(client) = clients.get(&key) { + return Ok(client.clone()); + } + let client = build_client(global_proxy()).map_err(|error| { + AiMuxError::Other(format!("shared HTTP client initialization failed: {error}")) + })?; + if clients.len() >= SHARED_CLIENT_CAP { + // Which entries are dead is unknowable, so evict them all; in-flight + // requests hold their own `Client` clone and are unaffected. + clients.clear(); + } + clients.insert(key, client.clone()); + Ok(client) +} + +fn build_client(proxy: ProxyConfig) -> Result { + let builder = Client::builder() + .connect_timeout(Duration::from_millis(10_000)) + .pool_max_idle_per_host(10) + .pool_idle_timeout(Some(Duration::from_secs(30))) + .tcp_keepalive(Some(Duration::from_secs(20))); + apply_proxy(builder, &proxy) + .build() + .map_err(|error| error.to_string()) +} + +fn apply_proxy(mut builder: reqwest::ClientBuilder, proxy: &ProxyConfig) -> reqwest::ClientBuilder { let http = proxy.http_url.as_deref().or(proxy.all_url.as_deref()); let https = proxy.https_url.as_deref().or(proxy.all_url.as_deref()); - let mut b = b; if let Some(url) = http - && let Ok(p) = ReqwestProxy::http(url) + && let Ok(reqwest_proxy) = reqwest::Proxy::http(url) { - b = apply_no_proxy(b, p, &proxy.no_proxy); + builder = apply_no_proxy(builder, reqwest_proxy, &proxy.no_proxy); } if let Some(url) = https - && let Ok(p) = ReqwestProxy::https(url) + && let Ok(reqwest_proxy) = reqwest::Proxy::https(url) { - b = apply_no_proxy(b, p, &proxy.no_proxy); + builder = apply_no_proxy(builder, reqwest_proxy, &proxy.no_proxy); } - b + builder } -/// Attach a no-proxy whitelist to a proxy if configured, then register it. +// Kept separate so both proxy schemes use precisely the same no-proxy rule. fn apply_no_proxy( - b: reqwest::ClientBuilder, + builder: reqwest::ClientBuilder, proxy: reqwest::Proxy, no_proxy: &Option, ) -> reqwest::ClientBuilder { match no_proxy.as_deref().map(reqwest::NoProxy::from_string) { - Some(Some(np)) => b.proxy(proxy.no_proxy(Some(np))), - _ => b.proxy(proxy), - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// 请求 / 响应抽象（纯数据，不依赖 reqwest） -// ════════════════════════════════════════════════════════════════════════════ - -/// HTTP 方法。 -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum HttpMethod { - Get, - Post, - Put, - Delete, - Patch, -} - -impl HttpMethod { - /// 方法名（小写），用于日志。 - #[must_use] - pub fn as_str(&self) -> &'static str { - match self { - HttpMethod::Get => "get", - HttpMethod::Post => "post", - HttpMethod::Put => "put", - HttpMethod::Delete => "delete", - HttpMethod::Patch => "patch", - } + Some(Some(no_proxy)) => builder.proxy(proxy.no_proxy(Some(no_proxy))), + _ => builder.proxy(proxy), } } -/// HTTP 请求体（纯数据）。 +/// Explicit raw request body used only by `post_to_api`. #[derive(Debug, Clone)] pub enum HttpBody { - /// JSON 请求体（自动设 `Content-Type: application/json`）。 Json(serde_json::Value), - /// 原始字节，带 content-type（multipart / 二进制上传）。 - /// `MultipartForm::finish()` 返回的 `(Vec, String)` 直接归入此变体。 Bytes(Vec, String), - /// 无请求体。 Empty, } -/// Per-request timeout limits (milliseconds). `None` disables a limit. -/// -/// This is the HTTP-layer view of [`TimeoutConfiguration`] (from -/// `CallOptions.timeout`). The HTTP layer enforces all three: -/// - `total_ms` — covers the whole call: connect + response (+ retries and, -/// for streaming, the entire stream). -/// - `first_chunk_ms` — streaming only: time waiting for the first chunk. -/// - `chunk_ms` — streaming only: max idle time between chunks. -#[derive(Debug, Clone, Copy, Default)] -pub struct RequestTimeout { - pub total_ms: Option, - pub first_chunk_ms: Option, - pub chunk_ms: Option, +#[derive(Debug, Clone, Copy)] +pub(crate) enum HttpMethod { + Get, + Post, } -impl From for RequestTimeout { - fn from(t: TimeoutConfiguration) -> Self { - Self { - total_ms: t.total_ms, - first_chunk_ms: t.first_chunk_ms, - chunk_ms: t.chunk_ms, +impl HttpMethod { + fn as_str(self) -> &'static str { + match self { + Self::Get => "GET", + Self::Post => "POST", } } } -/// HTTP 请求描述（纯数据，不依赖 reqwest）。 -/// -/// provider 构造此结构后交给 [`send`] / [`send_stream`] 执行。retry 时本 -/// 模块按需从此结构重建 `RequestBuilder`——provider 不参与重建。 +/// Metadata shared by `post_json_to_api`, `post_form_data_to_api`, +/// `post_to_api`, and `get_from_api`. Method and body are fixed by the helper +/// signature rather than supplied as runtime fields. #[derive(Debug, Clone)] pub struct HttpRequest { - pub method: HttpMethod, pub url: String, pub headers: Vec<(String, String)>, - pub body: HttpBody, - /// Optional signal used to cancel the request while it is connecting. - pub abort_signal: Option, - /// RFC-0023 recording correlation id (copied from `CallOptions.call_id` by - /// text-generation providers; `None` for other modalities). + pub abort_signal: Option, pub call_id: Option, - /// RFC-0023 recording context (R7 快照绑定)。层 A 入口构造, - /// text-generation provider 从 `CallOptions.recording_context` 复制。 - pub recording_context: Option, + pub recording_context: Option, } -/// HTTP 响应（非流式，纯数据）。 -/// -/// `body` 为完整响应体字节；provider 用 `serde_json::from_slice` 等解析。 -#[derive(Debug)] -pub struct HttpResponse { - pub status: u16, - pub headers: HashMap, - pub body: Bytes, +#[derive(Debug, Clone)] +pub(crate) struct PreparedRequest { + method: HttpMethod, + body: HttpBody, + pub(crate) url: String, + headers: Vec<(String, String)>, + pub(crate) abort_signal: Option, + call_id: Option, + recording_context: Option, } -/// HTTP 流式响应。 -/// -/// `body` 为字节流（`BoxStream>`，不依赖 reqwest）。 -/// 传给 `SseStream::new` 即可解析 SSE——`SseStream` 是泛型的，能直接接收。 -pub struct HttpStreamResponse { - pub status: u16, - pub headers: HashMap, - pub body: BoxStream<'static, Result>, +impl HttpRequest { + pub(crate) fn prepare(self, method: HttpMethod, body: HttpBody) -> PreparedRequest { + PreparedRequest { + method, + body, + url: self.url, + headers: self.headers, + abort_signal: self.abort_signal, + call_id: self.call_id, + recording_context: self.recording_context, + } + } } -// ════════════════════════════════════════════════════════════════════════════ -// 执行（含 retry，内部逻辑） -// ════════════════════════════════════════════════════════════════════════════ - -/// 发送非流式请求，带 retry。内部用 `shared_client()`（30s 整体超时）。 -/// -/// retry 策略（RFC-0009 §4.2）： -/// - 网络错误 → `ApiCall`（`is_retryable: true`,无 status——传输失败,可重试） -/// - 所有非 2xx → `parse_provider_error`；429 另读 `retry-after` / -/// `retry-after-ms` header 存入 `retry_after_ms` 字段（缺失即 `None`） -/// - 408/409/429/5xx 可重试；其他 4xx **立即返回不重试** -/// -/// 退避用 Full Jitter（`delay ∈ [0, base)`），防并发 429 惊群。 -/// -/// # Errors -/// -/// Returns `ApiCall` for transport failures and non-2xx responses (after the -/// configured retries), `Aborted` when the abort signal fires, and `Timeout` -/// when reading the body exceeds the 30s overall deadline. -pub async fn send( - request: HttpRequest, - retry_config: RetryConfig, - error_structure: &ErrorStructure, -) -> Result { - auto_init_from_env(); - let client = shared_client()?; - let started = Instant::now(); - let (resp, attempt) = - send_with_retry_raw(client, &request, retry_config, error_structure).await?; +fn api_call_error(request: &PreparedRequest, message: impl Into) -> ApiCallError { + ApiCallError::new( + message, + sanitized_request_url(&request.url), + crate::logging::redact_request_values(&request.body), + ) +} - let status = resp.status().as_u16(); - let headers = collect_headers(resp.headers()); - // The response body can be large/slow; keep honoring abort while reading. - let body = match &request.abort_signal { - Some(signal) => { - let bytes = resp.bytes(); - tokio::select! { - biased; - _ = signal.cancelled() => { - record_failed_exchange(&request, attempt, started.elapsed().as_millis() as u64, "aborted after 2xx", None); - record_transport_closed(&request); - return Err(AiMuxError::Aborted); - } - b = bytes => b.map_err(|e| { - record_failed_exchange(&request, attempt, started.elapsed().as_millis() as u64, &e.to_string(), None); - record_transport_closed(&request); - AiMuxError::ApiCall(ApiCallError { message: e.to_string(), is_retryable: true, ..Default::default() }) - })?, - } +struct ExchangeGuard<'a> { + request: &'a PreparedRequest, + attempt: u32, + exchange_index: u32, + started: Instant, + armed: bool, +} + +impl<'a> ExchangeGuard<'a> { + fn new( + request: &'a PreparedRequest, + attempt: u32, + exchange_index: u32, + started: Instant, + ) -> Self { + Self { + request, + attempt, + exchange_index, + started, + armed: true, } - None => resp.bytes().await.map_err(|e| { + } + + fn fail(&mut self, error: &str) { + if self.armed { record_failed_exchange( - &request, - attempt, - started.elapsed().as_millis() as u64, - &e.to_string(), - None, + self.request, + self.attempt, + self.exchange_index, + self.started.elapsed().as_millis() as u64, + error, ); - record_transport_closed(&request); - AiMuxError::ApiCall(ApiCallError { - message: e.to_string(), - is_retryable: true, - ..Default::default() - }) - })?, - }; + record_transport_closed(self.request); + self.armed = false; + } + } + + fn disarm(&mut self) { + self.armed = false; + } +} + +impl Drop for ExchangeGuard<'_> { + fn drop(&mut self) { + self.fail("exchange future cancelled"); + } +} - if body_logging_enabled() { +/// Execute exactly one HTTP exchange. Core owns retry and operation timeout. +pub(crate) async fn send_request_once( + request: &PreparedRequest, +) -> Result { + let started = Instant::now(); + let (attempt, exchange_index) = request + .recording_context + .as_ref() + .map(RecordingContext::next_exchange) + .unwrap_or((1, 1)); + let mut exchange = ExchangeGuard::new(request, attempt, exchange_index, started); + let client = match shared_client() { + Ok(client) => client, + Err(error) => { + exchange.fail(&error.to_string()); + return Err(error); + } + }; + tracing::debug!( + target: "aimux_provider_utils::http", + method = request.method.as_str(), + url = %sanitized_request_url(&request.url), + host = %request_host(&request.url), + body_size = request_body_size(request), + header_count = request.headers.len(), + call_id = request.call_id.as_deref().unwrap_or(""), + "request" + ); + if body_logging_enabled() + && let HttpBody::Json(value) = &request.body + { tracing::trace!( target: "aimux_provider_utils::http", - body = %redact_body(&String::from_utf8_lossy(&body)), - "response_body" + body = %redact_body(&value.to_string()), + "request_body" ); } - // RFC-0023: 成功 attempt 的 exchange(body 补全;attempt 来自重试循环)。 - // 门控提前:录制上下文不存在时零成本(S5 性能——不构造 exchange)。 - if request.recording_context.is_none() { - return Ok(HttpResponse { - status, - headers, - body, - }); - } - record_exchange( - &request, - aimux_core::recording::HttpExchange { - attempt, - request: to_http_record(&request), - response: Some(aimux_core::recording::ResponseRecord { - status, - headers: headers - .iter() - .map(|(k, v)| { - if aimux_core::recording::is_sensitive_key(k) { - (k.clone(), "[REDACTED]".to_string()) - } else { - (k.clone(), v.clone()) - } - }) - .collect(), - body: { - let s = String::from_utf8_lossy(&body); - Some(truncate_utf8(&s, RECORD_BODY_CAP).to_string()) - }, - stream_chunks: None, - ttfb_ms: None, - }), - timing: aimux_core::recording::TimingRecord { - latency_ms: started.elapsed().as_millis() as u64, - ttfb_ms: None, - }, - error: None, - finalized: true, - }, + let response = match send_one_request(&client, request).await { + Ok(response) => response, + Err(error) => { + exchange.fail(&error.to_string()); + return Err(error); + } + }; + exchange.disarm(); + let latency_ms = started.elapsed().as_millis() as u64; + tracing::debug!( + target: "aimux_provider_utils::http", + status = response.status().as_u16(), + latency_ms, + "response" ); - // 非流式:exchange 即终点 → 传输封闭(与层 A 幂等)。 - record_transport_closed(&request); + Ok(observe_response( + response, + request, + started, + latency_ms, + attempt, + exchange_index, + )) +} - Ok(HttpResponse { - status, - headers, - body, +async fn send_one_request( + client: &Client, + request: &PreparedRequest, +) -> Result { + if let Some(signal) = &request.abort_signal { + if signal.is_aborted() { + return Err(AiMuxError::from_abort_signal(signal)); + } + tokio::select! { + biased; + () = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), + response = build_request_builder(client, request)?.send() => { + response.map_err(|error| AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..api_call_error(request, error.to_string()) + }))) + } + } + } else { + build_request_builder(client, request)? + .send() + .await + .map_err(|error| { + AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..api_call_error(request, error.to_string()) + })) + }) + } +} + +fn build_request_builder( + client: &Client, + request: &PreparedRequest, +) -> Result { + let mut builder = match request.method { + HttpMethod::Get => client.get(&request.url), + HttpMethod::Post => client.post(&request.url), + }; + for (name, value) in &request.headers { + let header_name = reqwest::header::HeaderName::try_from(name) + .map_err(|_| AiMuxError::InvalidArgument(format!("invalid header name: {name}")))?; + let header_value = reqwest::header::HeaderValue::try_from(value).map_err(|_| { + AiMuxError::InvalidArgument(format!("invalid header value for {name}: {value}")) + })?; + builder = builder.header(header_name, header_value); + } + Ok(match &request.body { + HttpBody::Json(value) => builder.json(value), + HttpBody::Bytes(bytes, content_type) => builder + .header(reqwest::header::CONTENT_TYPE, content_type) + .body(bytes.clone()), + HttpBody::Empty => builder, }) } -/// 发送流式请求，返回字节流。内部用 `shared_streaming_client()`（无整体超时）。 -/// -/// retry 仅覆盖**建连阶段**——`.send()` 返回 200 后立即返回字节流，流中途 -/// 出错不重试（已吐 token 后重试会重复内容，RFC-0009 §5）。 +/// Abort-aware provider polling delay. This is not the operation retry delay; +/// Core owns exponential backoff. /// /// # Errors /// -/// Returns `ApiCall` when establishing the connection (including retries) -/// fails; transport errors on the byte stream surface as `Err` items in the -/// returned stream. -pub async fn send_stream( - request: HttpRequest, - retry_config: RetryConfig, - error_structure: &ErrorStructure, -) -> Result { - auto_init_from_env(); - let client = shared_streaming_client()?; - let started = Instant::now(); - let (resp, attempt) = - send_with_retry_raw(client, &request, retry_config, error_structure).await?; - - let status = resp.status().as_u16(); - let headers = collect_headers(resp.headers()); - tracing::debug!( - target: "aimux_provider_utils::http", - status = status, - "stream_connected" - ); +/// Returns [`AiMuxError::Aborted`] if the caller cancels during the delay. +pub async fn sleep_or_abort( + duration: Duration, + abort_signal: Option<&aimux_core::AbortSignal>, +) -> Result<(), AiMuxError> { + match abort_signal { + Some(signal) => tokio::select! { + biased; + () = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), + () = tokio::time::sleep(duration) => Ok(()), + }, + None => { + tokio::time::sleep(duration).await; + Ok(()) + } + } +} - // 流式录制上下文:仅调用点带 call_id 且录制开启时启用。 - let stream_rec = if request.recording_context.is_some() { - let resp_headers: Vec<(String, String)> = headers - .iter() - .map(|(k, v)| { - if aimux_core::recording::is_sensitive_key(k) { - (k.clone(), "[REDACTED]".to_string()) - } else { - (k.clone(), v.clone()) - } - }) - .collect(); - // 骨架 exchange(attempt 源自重试循环;finalized=false 等流终结补全)。 +fn observe_response( + response: reqwest::Response, + request: &PreparedRequest, + started: Instant, + latency_ms: u64, + attempt: u32, + exchange_index: u32, +) -> reqwest::Response { + let status = response.status(); + let version = response.version(); + let headers = response.headers().clone(); + let recording = request.recording_context.as_ref().map(|context| { + let redacted_headers = redacted_response_headers(&headers); + let response = ResponseRecord { + status: status.as_u16(), + headers: redacted_headers.clone(), + body: None, + stream_chunks: Some(0), + ttfb_ms: Some(latency_ms), + }; record_exchange( - &request, - aimux_core::recording::HttpExchange { + request, + HttpExchange { + step: request + .recording_context + .as_ref() + .and_then(RecordingContext::step), attempt, - request: to_http_record(&request), - response: Some(aimux_core::recording::ResponseRecord { - status, - headers: resp_headers.clone(), - body: None, - stream_chunks: None, - ttfb_ms: None, - }), - timing: aimux_core::recording::TimingRecord { - latency_ms: started.elapsed().as_millis() as u64, - ttfb_ms: None, + exchange_index, + request: to_http_record(request), + response: Some(response), + timing: TimingRecord { + latency_ms, + ttfb_ms: Some(latency_ms), }, error: None, finalized: false, }, ); - Some(StreamRec { - request: request.clone(), + ResponseRecording { + context: context.clone(), attempt, - status, - headers: resp_headers, - ttfb_ms: None, + exchange_index, + status: status.as_u16(), + headers: redacted_headers, + ttfb_ms: Some(latency_ms), body: ByteAccumulator::new(RECORD_BODY_CAP), - chunks: 0usize, - }) - } else { - None - }; - - // Provider request id, captured from the connect-phase headers: mid-stream - // transport errors are built after `resp` is consumed, so the header must - // be read now (`x-request-id` wins over `request-id`, as in the - // non-streaming path). - let request_id = ["x-request-id", "request-id"].iter().find_map(|k| { - headers - .iter() - .find(|(h, _)| h.eq_ignore_ascii_case(k)) - .map(|(_, v)| v.clone()) + chunks: 0, + finalized: false, + } }); - let body = ObservedByteStream { - inner: resp - .bytes_stream() - .map(move |item| { - item.map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: e.to_string(), - request_id: request_id.clone(), - is_retryable: true, - ..Default::default() - }) - }) - }) - .boxed(), - started: false, - start: started, - chunks: 0u64, + let observed = ObservedBody { + inner: response.bytes_stream().boxed(), + started, + chunks: 0, done: false, - record: stream_rec, - } - .boxed(); - - Ok(HttpStreamResponse { - status, - headers, - body, - }) -} - -/// 流式录制上下文(骨架 exchange 的补全数据)。 -struct StreamRec { - request: HttpRequest, - attempt: u32, - status: u16, - headers: Vec<(String, String)>, - ttfb_ms: Option, - /// 原始字节累积(跨 chunk 边界 UTF-8 安全;受 RECORD_BODY_CAP 上限)。 - body: ByteAccumulator, + recording, + }; + let mut rebuilt = http::Response::builder() + .status(status) + .version(version) + .body(reqwest::Body::wrap_stream(observed)) + .expect("status and version came from a valid reqwest response"); + *rebuilt.headers_mut() = headers; + rebuilt.into() +} + +struct ObservedBody { + inner: BoxStream<'static, Result>, + started: Instant, chunks: usize, -} - -impl StreamRec { - /// 终结时补全 response(transport 结束时对 attempt 定位更新)。 - fn finalize(&mut self, error: Option) { - // 流结束时统一做一次 UTF-8 解码:跨 chunk 边界的多字节字符不再被 - // lossy 拆成 U+FFFD(A6)。 - let body = self.body.decode(); - let response = aimux_core::recording::ResponseRecord { - status: self.status, - headers: self.headers.clone(), - body, - stream_chunks: Some(self.chunks), - ttfb_ms: self.ttfb_ms, - }; - // 一次 patch:补 body/status/ttfb + 可选 error,保持 attempt 唯一(S1)。 - update_exchange_response(&self.request, self.attempt, response, error); - record_transport_closed(&self.request); - } -} - -/// 流式响应体观测包装：发出「首字节」与「流结束」两个 debug 事件 -/// （RFC-0014 §4.2 stream 行）。零成本——事件只在开启时格式化。 -struct ObservedByteStream { - inner: BoxStream<'static, Result>, - started: bool, - start: Instant, - chunks: u64, done: bool, - /// RFC-0023:存在时累积原始 SSE 文本并终结补全。 - record: Option, + recording: Option, } -impl Stream for ObservedByteStream { - type Item = Result; +impl Stream for ObservedBody { + type Item = Result; - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.as_mut().get_mut(); - if this.done { + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.done { return Poll::Ready(None); } - let result = this.inner.as_mut().poll_next(cx); + let result = self.inner.as_mut().poll_next(context); match &result { Poll::Ready(Some(Ok(bytes))) => { - if !this.started { - this.started = true; - let ttfb = this.start.elapsed().as_millis() as u64; - tracing::debug!( - target: "aimux_provider_utils::http", - ttfb_ms = ttfb, - "stream_first_byte" - ); - if let Some(rec) = &mut this.record { - rec.ttfb_ms = Some(ttfb); - } - } - this.chunks += 1; - if let Some(rec) = &mut this.record { - // 累积原始字节(不在每个 chunk 上单独 lossy 解码,避免跨 - // chunk 边界拆坏多字节字符;A6)。结尾 finalize 统一解码。 - rec.body.push(bytes); - rec.chunks = this.chunks as usize; + self.chunks += 1; + let chunks = self.chunks; + if let Some(recording) = &mut self.recording { + recording.body.push(bytes); + recording.chunks = chunks; } } - Poll::Ready(Some(Err(e))) => { - this.done = true; - tracing::debug!( - target: "aimux_provider_utils::http", - chunks = this.chunks, - duration_ms = this.start.elapsed().as_millis() as u64, - "stream_end" - ); - if let Some(rec) = &mut this.record { - rec.finalize(Some(e.to_string())); + Poll::Ready(Some(Err(error))) => { + self.done = true; + if let Some(recording) = &mut self.recording { + recording.finalize(Some(error.to_string())); } } Poll::Ready(None) => { - this.done = true; + self.done = true; tracing::debug!( target: "aimux_provider_utils::http", - chunks = this.chunks, - duration_ms = this.start.elapsed().as_millis() as u64, - "stream_end" + chunks = self.chunks, + duration_ms = self.started.elapsed().as_millis() as u64, + "response_body_end" ); - if let Some(rec) = &mut this.record { - rec.finalize(None); + if let Some(recording) = &mut self.recording { + recording.finalize(None); } } Poll::Pending => {} @@ -694,879 +502,200 @@ impl Stream for ObservedByteStream { } } -impl Drop for ObservedByteStream { +impl Drop for ObservedBody { fn drop(&mut self) { - // 消费方提前放弃流:仍补全已累积的 body,避免骨架悬空。 - if self.record.is_some() - && !self.done - && let Some(rec) = &mut self.record + if !self.done + && let Some(recording) = &mut self.recording { - rec.finalize(Some("abandoned".to_string())); + recording.finalize(Some("response body abandoned".into())); } } } -// ════════════════════════════════════════════════════════════════════════════ -// Timed variants — per-call timeout support (RFC-0016 H3) -// ════════════════════════════════════════════════════════════════════════════ - -/// Send a non-streaming request with an optional per-call total timeout. -/// -/// The timeout covers the **entire** call: connect, response body, and any -/// retry backoff. On expiry the call fails with `AiMuxError::Timeout` (no -/// retry is attempted after the deadline). -/// -/// # Errors -/// -/// Returns `InvalidArgument` for a malformed timeout, `Timeout` when the -/// total deadline expires, and the inner retry loop's `ApiCall`/`Aborted` -/// errors. -pub async fn send_timed( - request: HttpRequest, - retry_config: RetryConfig, - error_structure: &ErrorStructure, - timeout: Option, -) -> Result { - let timeout = timeout.unwrap_or_default(); - validate_timeout(&timeout)?; +struct ResponseRecording { + context: RecordingContext, + attempt: u32, + exchange_index: u32, + status: u16, + headers: Vec<(String, String)>, + ttfb_ms: Option, + body: ByteAccumulator, + chunks: usize, + finalized: bool, +} - let fut = send(request, retry_config, error_structure); - match timeout.total_ms { - Some(0) => Err(AiMuxError::Timeout("total timeout".to_string())), - Some(ms) => Ok(tokio::time::timeout(Duration::from_millis(ms), fut) - .await - .map_err(|_| AiMuxError::Timeout(format!("total timeout after {ms}ms")))??), - None => fut.await, +impl ResponseRecording { + fn finalize(&mut self, error: Option) { + if self.finalized { + return; + } + self.finalized = true; + self.context.recorder.record_exchange_update( + &self.context.call_id, + self.attempt, + self.exchange_index, + &ResponseRecord { + status: self.status, + headers: self.headers.clone(), + body: self.body.decode(), + stream_chunks: Some(self.chunks), + ttfb_ms: self.ttfb_ms, + }, + error, + ); + self.context + .recorder + .record_transport_closed(&self.context.call_id); } } -/// Send a streaming request with optional per-call timeouts. -/// -/// `total_ms` covers connect (+ retries); the returned byte stream is then -/// wrapped so that `first_chunk_ms` / `chunk_ms` / the remaining `total_ms` -/// are enforced on the chunks themselves. Expired limits surface as -/// `AiMuxError::Timeout` items in the stream (after which the stream ends). -/// -/// # Errors -/// -/// Returns `InvalidArgument` for a malformed timeout and `Timeout` when the -/// connect-phase deadline expires; expired chunk limits surface as `Timeout` -/// items inside the returned stream. -pub async fn send_stream_timed( - request: HttpRequest, - retry_config: RetryConfig, - error_structure: &ErrorStructure, - timeout: Option, -) -> Result { - let timeout = timeout.unwrap_or_default(); - validate_timeout(&timeout)?; - let start = Instant::now(); - let resp = match timeout.total_ms { - Some(0) => return Err(AiMuxError::Timeout("total timeout".to_string())), - Some(ms) => tokio::time::timeout( - Duration::from_millis(ms), - send_stream(request.clone(), retry_config, error_structure), - ) - .await - .map_err(|_| AiMuxError::Timeout(format!("total timeout after {ms}ms")))??, - None => send_stream(request.clone(), retry_config, error_structure).await?, - }; - // Wrap the body whenever timeouts OR an abort signal are in play: the - // wrapper is what makes the body phase honor abort, so bypassing it when - // only an abort signal is set would drop body cancellation (RFC-0016 - // review S1). - if timeout.total_ms.is_none() - && timeout.first_chunk_ms.is_none() - && timeout.chunk_ms.is_none() - && request.abort_signal.is_none() - { - return Ok(resp); - } +const RECORD_BODY_CAP: usize = 1 << 20; - Ok(HttpStreamResponse { - status: resp.status, - headers: resp.headers, - body: TimeoutBodyStream { - inner: resp.body, - start, - total_ms: timeout.total_ms, - first_chunk_ms: timeout.first_chunk_ms, - chunk_ms: timeout.chunk_ms, - abort_signal: request.abort_signal, - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - } - .boxed(), - }) -} - -/// Reject timeout values whose deadline cannot be represented (e.g. -/// `u64::MAX` ms on platforms whose `Instant` range is narrower) — adding -/// such a duration to an `Instant` would panic. -fn validate_timeout(timeout: &RequestTimeout) -> Result<(), AiMuxError> { - let now = Instant::now(); - for ms in [timeout.total_ms, timeout.first_chunk_ms, timeout.chunk_ms] - .into_iter() - .flatten() - { - if now.checked_add(Duration::from_millis(ms)).is_none() { - return Err(AiMuxError::InvalidArgument(format!( - "timeout value {ms}ms is too large" - ))); - } - } - Ok(()) -} - -/// Which deadline fired (for error messages). -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum TimeoutKind { - FirstChunk, - ChunkIdle, - Total, -} - -impl TimeoutKind { - fn message(self) -> &'static str { - match self { - TimeoutKind::FirstChunk => "first chunk timeout", - TimeoutKind::ChunkIdle => "chunk idle timeout", - TimeoutKind::Total => "total timeout", - } - } -} - -/// Stream wrapper enforcing first-chunk / chunk-idle / total deadlines plus -/// abort-signal cancellation. -/// -/// Polling is bounded by the earliest applicable deadline: before the first -/// item, `first_chunk_ms` (from request start) applies; afterwards, -/// `chunk_ms` is a **sliding window** reset on every chunk. `total_ms` caps -/// everything from request start. An aborted signal ends the stream with -/// `AiMuxError::Aborted` (abort wins over deadlines). After an error item -/// the stream yields `None` (fused). -/// -/// The pending `tokio::time::Sleep` and the abort waiter are stored in the -/// stream (not locals): dropping a `Sleep` unregisters its timer entry, and -/// the abort waiter is a `Notify` future whose waker must stay registered -/// while the stream is Pending. -struct TimeoutBodyStream { - inner: BoxStream<'static, Result>, - start: Instant, - total_ms: Option, - first_chunk_ms: Option, - chunk_ms: Option, - abort_signal: Option, - first: bool, - last_chunk_at: Option, - sleep: Option>>, - abort_wait: Option + Send>>>, - done: bool, -} - -impl TimeoutBodyStream { - /// The earliest applicable deadline and what it is, or `None` when no - /// limit is set. - /// - /// `checked_add` guards the chunk-idle case: `last_chunk_at` moves - /// forward with every chunk, so a value that passed entry validation - /// (relative to that moment) could still overflow later. On overflow the - /// limit is treated as total-immediate (the earliest possible deadline). - fn next_deadline(&self) -> Option<(Instant, TimeoutKind)> { - let phase = if self.first { - self.first_chunk_ms.and_then(|ms| { - self.start - .checked_add(Duration::from_millis(ms)) - .map(|d| (d, TimeoutKind::FirstChunk)) - }) - } else { - self.chunk_ms.and_then(|ms| { - self.last_chunk_at - .and_then(|t| t.checked_add(Duration::from_millis(ms))) - .map(|d| (d, TimeoutKind::ChunkIdle)) - }) - }; - let total = self.total_ms.and_then(|ms| { - self.start - .checked_add(Duration::from_millis(ms)) - .map(|d| (d, TimeoutKind::Total)) - }); - match (phase, total) { - (Some((pd, pk)), Some((td, _))) => { - if pd <= td { - Some((pd, pk)) - } else { - Some((td, TimeoutKind::Total)) - } - } - (Some(x), None) | (None, Some(x)) => Some(x), - // (None, None): either nothing applies in the current phase - // (→ no deadline) or an applicable limit overflowed its - // representable range (→ treat as already expired). - (None, None) => { - let phase_configured = if self.first { - self.first_chunk_ms.is_some() - } else { - self.chunk_ms.is_some() - }; - (phase_configured || self.total_ms.is_some()) - .then_some((self.start, TimeoutKind::Total)) - } - } - } -} - -impl Stream for TimeoutBodyStream { - type Item = Result; - - fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.as_mut().get_mut(); - - if this.done { - return Poll::Ready(None); - } - - // Abort fast path — also catches signals aborted before the call. - if let Some(sig) = &this.abort_signal - && sig.is_aborted() - { - this.done = true; - this.sleep = None; - return Poll::Ready(Some(Err(AiMuxError::Aborted))); - } - - // Lazy-create the event-driven abort waiter so an abort wakes us - // while Pending (the signal's `Notify` stores the waker). - if this.abort_wait.is_none() - && let Some(signal) = this.abort_signal.as_ref() - { - this.abort_wait = Some(Box::pin(signal.cancelled())); - } - - let Some((deadline, kind)) = this.next_deadline() else { - // No timeout configured: pass through (abort still honored). - if let Some(wait) = this.abort_wait.as_mut() - && wait.as_mut().poll(cx).is_ready() - { - this.done = true; - return Poll::Ready(Some(Err(AiMuxError::Aborted))); - } - let result = this.inner.as_mut().poll_next(cx); - match &result { - Poll::Ready(Some(Ok(_))) => { - this.first = false; - this.last_chunk_at = Some(Instant::now()); - } - Poll::Ready(Some(Err(_))) | Poll::Ready(None) => { - this.done = true; - } - Poll::Pending => {} - } - return result; - }; - - if Instant::now() >= deadline { - this.done = true; - this.sleep = None; - return Poll::Ready(Some(Err(AiMuxError::Timeout(kind.message().to_string())))); - } - - // (Re)create the timer if missing or pointed at a different deadline. - let needs_reset = match &this.sleep { - Some(s) => s.deadline() != deadline.into(), - None => true, - }; - if needs_reset { - this.sleep = Some(Box::pin(tokio::time::sleep_until(deadline.into()))); - } - - { - // Abort wins over any deadline. - if let Some(wait) = this.abort_wait.as_mut() - && wait.as_mut().poll(cx).is_ready() - { - this.done = true; - this.sleep = None; - return Poll::Ready(Some(Err(AiMuxError::Aborted))); - } - if let Some(sleep) = this.sleep.as_mut() - && sleep.as_mut().poll(cx).is_ready() - { - this.done = true; - this.sleep = None; - return Poll::Ready(Some(Err(AiMuxError::Timeout(kind.message().to_string())))); - } - match this.inner.as_mut().poll_next(cx) { - Poll::Ready(Some(Ok(bytes))) => { - this.first = false; - this.last_chunk_at = Some(Instant::now()); - // Reset the idle window: the next poll re-creates the - // timer with a fresh `now + chunk_ms` deadline. - this.sleep = None; - Poll::Ready(Some(Ok(bytes))) - } - Poll::Ready(Some(Err(e))) => { - this.done = true; - this.sleep = None; - Poll::Ready(Some(Err(e))) - } - Poll::Ready(None) => { - this.done = true; - this.sleep = None; - Poll::Ready(None) - } - Poll::Pending => Poll::Pending, - } - } - } -} - -// ════════════════════════════════════════════════════════════════════════════ -// retry 核心 -// ════════════════════════════════════════════════════════════════════════════ - -/// Sleep that stays responsive to cancellation (RFC-0016 §7.6 R2/R4, §7.7). -/// -/// Contract: -/// - an already-aborted signal fails immediately (no sleep); -/// - abort wins over the sleep when both are ready (`biased` select); -/// - `None` signal behaves exactly like `tokio::time::sleep`; -/// - cancellation is exactly [`AiMuxError::Aborted`]. -/// -/// Shared by the retry backoff here and by provider polling waits. -/// -/// # Errors -/// -/// Returns `AiMuxError::Aborted` when the cancellation signal fires before or -/// during the sleep. -pub async fn sleep_or_abort( - duration: Duration, - signal: Option<&AbortSignal>, -) -> Result<(), AiMuxError> { - match signal { - Some(signal) => { - tokio::select! { - biased; - _ = signal.cancelled() => Err(AiMuxError::Aborted), - _ = tokio::time::sleep(duration) => Ok(()), - } - } - None => { - tokio::time::sleep(duration).await; - Ok(()) - } - } -} - -/// Read an error-response body honoring abort (RFC-0016 review S2). -/// -/// Body reads for non-2xx failures are awaited before the next retry -/// decision; without this, an abort during the read would be ignored until -/// the next attempt. -/// -/// The body is **streamed** (chunk-by-chunk) and only the first -/// [`MAX_ERROR_BODY_BYTES`] are retained — the total byte count is still -/// counted for the truncation marker — so a misbehaving provider cannot force -/// unbounded memory growth or flood logs/FFI envelopes with a huge error -/// payload (audit finding, rounds 1–2). -const MAX_ERROR_BODY_BYTES: usize = 64 * 1024; - -async fn read_error_body( - resp: reqwest::Response, - request: &HttpRequest, -) -> Result { - let read = async { - let mut stream = resp.bytes_stream(); - let mut collected: Vec = Vec::new(); - let mut total: usize = 0; - while let Some(chunk) = stream.next().await { - let chunk = match chunk { - Ok(c) => c, - // The connection died mid-read: surface what we have. - Err(_) => break, - }; - total = total.saturating_add(chunk.len()); - if collected.len() < MAX_ERROR_BODY_BYTES { - let take = (MAX_ERROR_BODY_BYTES - collected.len()).min(chunk.len()); - collected.extend_from_slice(&chunk[..take]); - } - } - let mut s = String::from_utf8_lossy(&collected).into_owned(); - // Truncate when either the raw body exceeded the budget (marker must - // warn that content was dropped) or lossy decoding expanded the string - // (invalid bytes become 3-byte U+FFFD) beyond the budget. - if total > MAX_ERROR_BODY_BYTES || s.len() > MAX_ERROR_BODY_BYTES { - // Truncate on a UTF-8 char boundary so the marker never splits a - // multi-byte character. `0` is always a char boundary, so the - // loop terminates. - let mut cut = s.len().min(MAX_ERROR_BODY_BYTES); - while !s.is_char_boundary(cut) { - cut -= 1; - } - s.truncate(cut); - s.push_str(&format!("…(truncated, full error body {total} bytes)")); - } - s - }; - match &request.abort_signal { - Some(signal) => { - let text = read; - tokio::select! { - biased; - _ = signal.cancelled() => Err(AiMuxError::Aborted), - t = text => Ok(t), - } - } - None => Ok(read.await), - } -} - -// ── RFC-0023 层 B 录制支持(经 http 咽喉点;per-attempt + 流式累积)────── - -/// 录制 body 上限(超出截断,防内存膨胀)。 -const RECORD_BODY_CAP: usize = 1 << 20; // 1 MiB - -/// UTF-8 安全截断:最多保留前 `cap` 字节(不切在多字节字符中间)。 -fn truncate_utf8(s: &str, cap: usize) -> &str { - if s.len() <= cap { - return s; - } - let mut end = cap; - while end > 0 && !s.is_char_boundary(end) { - end -= 1; - } - &s[..end] -} - -/// 有界字节累积器:流式录制按原始字节累积,流末统一解码。 -/// -/// 解决每个网络 chunk 单独 `String::from_utf8_lossy` 导致跨 chunk 边界的 -/// 多字节字符(中文/emoji)被拆成 U+FFFD 的问题(A6)。累积受 `cap` 限制, -/// 超出即截断并标记;截断处可能切在多字节字符中间,解码时去掉末尾替换符 -/// 并追加 `…(truncated)` 标记。仅用于流式 SSE 录制(非流式 body 不经此路径)。 struct ByteAccumulator { - buf: Vec, - cap: usize, + bytes: Vec, truncated: bool, } impl ByteAccumulator { - fn new(cap: usize) -> Self { + fn new(capacity: usize) -> Self { Self { - buf: Vec::new(), - cap, + bytes: Vec::with_capacity(capacity.min(8 * 1024)), truncated: false, } } - /// 追加一个网络 chunk(受 cap 限制,超出标记截断)。 fn push(&mut self, bytes: &[u8]) { - if self.truncated { - return; - } - let remaining = self.cap.saturating_sub(self.buf.len()); - if remaining == 0 { + let remaining = RECORD_BODY_CAP.saturating_sub(self.bytes.len()); + if bytes.len() > remaining { + self.bytes.extend_from_slice(&bytes[..remaining]); self.truncated = true; - return; - } - if bytes.len() <= remaining { - self.buf.extend_from_slice(bytes); } else { - // 截断到 cap(可能切在多字节字符中间;decode 时处理边界)。 - self.buf.extend_from_slice(&bytes[..remaining]); - self.truncated = true; + self.bytes.extend_from_slice(bytes); } } - /// 终结解码:统一一次 UTF-8 解码(跨 chunk 多字节字符不破坏)。 - /// 截断时去掉末尾的 U+FFFD(来自 cap 处的不完整字符)并追加标记。 fn decode(&self) -> Option { - if self.buf.is_empty() { + if self.bytes.is_empty() { return None; } - let mut s = String::from_utf8_lossy(&self.buf).into_owned(); + let mut value = String::from_utf8_lossy(&self.bytes).into_owned(); if self.truncated { - while s.ends_with('\u{FFFD}') { - s.pop(); + while value.ends_with('\u{FFFD}') { + value.pop(); } - s.push_str("…(truncated)"); + value.push_str("…(truncated)"); } - Some(s) + Some(value) } } -/// 把请求转成录制侧的 HttpRecord(敏感头脱敏;body 截断)。 -fn to_http_record(request: &HttpRequest) -> aimux_core::recording::HttpRecord { +fn to_http_record(request: &PreparedRequest) -> HttpRecord { use aimux_core::recording::is_sensitive_key; let headers = request .headers .iter() - .map(|(k, v)| { - if is_sensitive_key(k) { - (k.clone(), "[REDACTED]".to_string()) - } else { - (k.clone(), v.clone()) - } + .map(|(name, value)| { + ( + name.clone(), + if is_sensitive_key(name) { + "[REDACTED]".into() + } else { + value.clone() + }, + ) }) .collect(); let body = match &request.body { - HttpBody::Json(v) => serde_json::to_string(v).unwrap_or_default(), - HttpBody::Bytes(b, _) => String::from_utf8_lossy(b).into_owned(), - HttpBody::Empty => String::new(), - }; - let body = if body.is_empty() { - None - } else { - Some(truncate_utf8(&body, RECORD_BODY_CAP).to_string()) - }; - aimux_core::recording::HttpRecord { - method: request.method.as_str().to_string(), - url: request.url.clone(), + HttpBody::Json(value) => { + serde_json::to_string(&crate::logging::redact_error_context(value.clone())).ok() + } + HttpBody::Bytes(bytes, _) => Some(String::from_utf8_lossy(bytes).into_owned()), + HttpBody::Empty => None, + } + .map(|body| truncate_utf8(&body, RECORD_BODY_CAP).to_owned()); + HttpRecord { + method: request.method.as_str().into(), + url: sanitized_request_url(&request.url), headers, body, } } -/// 录制单条 exchange(层 B 入队;call_id 缺失或录制关闭时零开销)。 -fn record_exchange(request: &HttpRequest, exchange: aimux_core::recording::HttpExchange) { - if let Some(ctx) = &request.recording_context { - ctx.recorder.record_exchange(&ctx.call_id, &exchange); - } -} - -/// 补全流式 response(同 attempt 定位,令 finalized=true)。 -fn update_exchange_response( - request: &HttpRequest, - attempt: u32, - response: aimux_core::recording::ResponseRecord, - error: Option, -) { - if let Some(ctx) = &request.recording_context { - ctx.recorder - .record_exchange_update(&ctx.call_id, attempt, &response, error); - } -} - -/// 声明传输层封闭(层 B;P2 语义:不再有 exchange)。 -fn record_transport_closed(request: &HttpRequest) { - if let Some(ctx) = &request.recording_context { - ctx.recorder.record_transport_closed(&ctx.call_id); +fn record_exchange(request: &PreparedRequest, exchange: HttpExchange) { + if let Some(context) = &request.recording_context { + context + .recorder + .record_exchange(&context.call_id, &exchange); } } -/// 收集响应头并脱敏(录制 ResponseRecord.headers 用;敏感头值替换为 [REDACTED])。 -fn redacted_response_headers(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { - use aimux_core::recording::is_sensitive_key; - headers - .iter() - .map(|(k, v)| { - let key = k.to_string(); - if is_sensitive_key(&key) { - (key, "[REDACTED]".to_string()) - } else { - (key, v.to_str().unwrap_or("").to_string()) - } - }) - .collect() -} - -/// 由非 2xx 响应构造结构化 [`ResponseRecord`](aimux_core::recording::ResponseRecord) -/// (用于失败 attempt 的录制;A5)。`body` 已由 `read_error_body` 截断到 64KiB, -/// 这里再做一次 `RECORD_BODY_CAP` 安全截断(边界对齐 UTF-8 字符)。 -fn error_response_record( - status: u16, - headers: Vec<(String, String)>, - body: &str, -) -> aimux_core::recording::ResponseRecord { - aimux_core::recording::ResponseRecord { - status, - headers, - body: if body.is_empty() { - None - } else { - Some(truncate_utf8(body, RECORD_BODY_CAP).to_string()) - }, - stream_chunks: None, - ttfb_ms: None, +fn record_transport_closed(request: &PreparedRequest) { + if let Some(context) = &request.recording_context { + context.recorder.record_transport_closed(&context.call_id); } } -/// 录制失败 attempt 的 exchange(终态,error 携带原因)。 -/// -/// `response`:`Some` 表示已收到合法 HTTP 响应(4xx/5xx,含 429)——结构化 -/// status/headers/body 一并录制;`None` 表示 DNS/连接/TLS/abort 等无 HTTP -/// 响应的传输层失败。`error` 字符串始终保留(日志/诊断用),录制结构完整(A5)。 fn record_failed_exchange( - request: &HttpRequest, + request: &PreparedRequest, attempt: u32, + exchange_index: u32, latency_ms: u64, error: &str, - response: Option, ) { record_exchange( request, - aimux_core::recording::HttpExchange { + HttpExchange { + step: request + .recording_context + .as_ref() + .and_then(RecordingContext::step), attempt, + exchange_index, request: to_http_record(request), - response, - timing: aimux_core::recording::TimingRecord { + response: None, + timing: TimingRecord { latency_ms, ttfb_ms: None, }, - error: Some(error.to_string()), + error: Some(error.to_owned()), finalized: true, }, ); } -/// retry 核心：反复 `.send()` 直到拿到 2xx 响应或耗尽重试。 -/// -/// 这是 http 层内部函数——`reqwest::Response` 不外泄。每次重试从 `&request` -/// 重建 `RequestBuilder`（HttpRequest 是纯数据，可重复读）。 -async fn send_with_retry_raw( - client: &Client, - request: &HttpRequest, - retry_config: RetryConfig, - error_structure: &ErrorStructure, -) -> Result<(reqwest::Response, u32), AiMuxError> { - let span = tracing::debug_span!( - target: "aimux_provider_utils::http", - "http_request", - method = %request.method.as_str(), - host = %request_host(&request.url), - attempt = 0u64, - ); - let _enter = span.enter(); - - let mut last_error = AiMuxError::Other("no attempts made".to_string()); - let mut last_status: Option = None; - let mut exponential_delay_ms = retry_config.initial_delay.as_millis() as i64; - - for attempt in 0..=retry_config.max_retries { - let attempt_start = Instant::now(); - let resp = send_request(client, request).await; - let latency_ms = attempt_start.elapsed().as_millis() as u64; - - match resp { - Ok(resp) => { - let status_code = resp.status().as_u16(); - last_status = Some(status_code); - tracing::debug!( - target: "aimux_provider_utils::http", - status = status_code, - latency_ms = latency_ms, - "response" - ); - if resp.status().is_success() { - return Ok((resp, attempt)); - } - // Facts read from the response *before* the body consumes it: - // request id, retry hint, redacted headers for recording. - let request_id = ["x-request-id", "request-id"] - .iter() - .find_map(|k| resp.headers().get(*k)) - .and_then(|v| v.to_str().ok()) - .map(str::to_owned); - // Retry hint: stored only for 429 (current core contract). - // Missing/invalid/negative header → None; never fabricate a - // local fallback as provider data (RFC-0009 §4.2). - let retry_after_ms = (status_code == 429) - .then(|| { - parse_retry_after( - resp.headers() - .get("retry-after-ms") - .and_then(|v| v.to_str().ok()), - resp.headers() - .get("retry-after") - .and_then(|v| v.to_str().ok()), - SystemTime::now(), - ) - }) - .flatten() - .and_then(|ms| u64::try_from(ms).ok()); - let resp_headers = redacted_response_headers(resp.headers()); - let body = read_error_body(resp, request).await?; - // Every non-2xx goes through the provider error parser so - // extracted message/provider_code/raw body survive; the - // observed request id and real retry hint are attached to the - // same detail (no rebuilt error losing parsed fields). - let mut err = parse_provider_error(status_code, &body, error_structure); - if let AiMuxError::ApiCall(d) = &mut err { - d.request_id = request_id; - d.retry_after_ms = retry_after_ms; - } - record_failed_exchange( - request, - attempt, - latency_ms, - &err.to_string(), - Some(error_response_record(status_code, resp_headers, &body)), - ); - if err.is_retryable() { - // 408/409/429/5xx: flow into the shared backoff below. - last_error = err; - } else { - // Ordinary 4xx: return after one attempt. - return Err(err); - } - } - Err(e) => { - // 传输层失败(DNS/连接/TLS):无 HTTP 响应,response 保持 None(A5)。 - record_failed_exchange(request, attempt, latency_ms, &e.to_string(), None); - last_error = e; - } - } - - if !last_error.is_retryable() || attempt == retry_config.max_retries { - // 用户主动取消不算故障，不落 error 日志（RFC-0014 §4.2 failed 行）。 - if !matches!(last_error, AiMuxError::Aborted) { - tracing::error!( - target: "aimux_provider_utils::http", - status = last_status.unwrap_or(0), - attempts = attempt + 1, - reason = %error_variant(&last_error), - "failed" - ); - } - return Err(last_error); - } - - let hint = last_error.retry_after_hint(); - let delay_ms = { - let mut rng = rand::thread_rng(); - get_retry_delay_ms_with_jitter(hint, exponential_delay_ms, &mut rng) - }; - tracing::warn!( - target: "aimux_provider_utils::http", - attempt = attempt + 1, - max_retries = retry_config.max_retries, - status = last_status.unwrap_or(0), - delay_ms = delay_ms.max(0) as u64, - reason = %error_variant(&last_error), - "retry" - ); - span.record("attempt", attempt + 1); - // Backoff must also be abortable: aborting during a long backoff - // window (e.g. a large Retry-After) must not be delayed until the - // next attempt. - let delay = Duration::from_millis(delay_ms.max(0) as u64); - sleep_or_abort(delay, request.abort_signal.as_ref()).await?; - exponential_delay_ms = - exponential_delay_ms.saturating_mul(retry_config.backoff_factor as i64); - } - - Err(last_error) +fn redacted_response_headers(headers: &reqwest::header::HeaderMap) -> Vec<(String, String)> { + crate::extract_response_headers::extract_response_header_pairs(headers) } -/// Send one HTTP attempt, cancelling the in-flight connection if requested. -async fn send_request( - client: &Client, - request: &HttpRequest, -) -> Result { - tracing::debug!( - target: "aimux_provider_utils::http", - method = %request.method.as_str(), - url = %request_url_no_query(&request.url), - body_size = request_body_size(request), - header_count = request.headers.len(), - "request" - ); - if body_logging_enabled() - && let HttpBody::Json(value) = &request.body - { - tracing::trace!( - target: "aimux_provider_utils::http", - body = %redact_body(&value.to_string()), - "request_body" - ); +fn truncate_utf8(value: &str, maximum: usize) -> &str { + if value.len() <= maximum { + return value; } - - if let Some(signal) = &request.abort_signal { - if signal.is_aborted() { - return Err(AiMuxError::Aborted); - } - - let response = build_request_builder(client, request)?.send(); - tokio::select! { - biased; - _ = signal.cancelled() => Err(AiMuxError::Aborted), - result = response => result.map_err(|e| AiMuxError::ApiCall(ApiCallError { message: e.to_string(), is_retryable: true, ..Default::default() })), - } - } else { - build_request_builder(client, request)? - .send() - .await - .map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: e.to_string(), - is_retryable: true, - ..Default::default() - }) - }) + let mut end = maximum; + while end > 0 && !value.is_char_boundary(end) { + end -= 1; } + &value[..end] } -/// 把纯数据的 [`HttpRequest`] 转成 `reqwest::RequestBuilder`。 -/// -/// 每次重试调用一次。`request` 以引用传入，不消费——支持重试重建。 -/// -/// 若任一 header 的 name/value 无法转成合法的 reqwest header（含非法字节）， -/// 返回 [`AiMuxError::InvalidArgument`]——这通常意味着鉴权 header 等关键 -/// header 被丢弃，调用者必须感知，而不是静默跳过。 -fn build_request_builder( - client: &Client, - request: &HttpRequest, -) -> Result { - let mut builder = match request.method { - HttpMethod::Get => client.get(&request.url), - HttpMethod::Post => client.post(&request.url), - HttpMethod::Put => client.put(&request.url), - HttpMethod::Delete => client.delete(&request.url), - HttpMethod::Patch => client.patch(&request.url), - }; - for (name, value) in &request.headers { - let header_name = reqwest::header::HeaderName::try_from(name) - .map_err(|_| AiMuxError::InvalidArgument(format!("invalid header name: {name}")))?; - let header_value = reqwest::header::HeaderValue::try_from(value).map_err(|_| { - AiMuxError::InvalidArgument(format!("invalid header value for {name}: {value}")) - })?; - builder = builder.header(header_name, header_value); +fn request_body_size(request: &PreparedRequest) -> u64 { + match &request.body { + HttpBody::Json(value) => serde_json::to_vec(value) + .map(|bytes| bytes.len() as u64) + .unwrap_or_default(), + HttpBody::Bytes(bytes, _) => bytes.len() as u64, + HttpBody::Empty => 0, } - let builder = match &request.body { - HttpBody::Json(value) => builder.json(value), - HttpBody::Bytes(bytes, content_type) => builder - .header("Content-Type", content_type) - .body(bytes.clone()), - HttpBody::Empty => builder, - }; - Ok(builder) -} - -/// 从 `reqwest::header::HeaderMap` 提取 `HashMap`。 -fn collect_headers(headers: &reqwest::header::HeaderMap) -> HashMap { - headers - .iter() - .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) - .collect() } -// ════════════════════════════════════════════════════════════════════════════ -// 日志辅助（RFC-0014 §4.2/§4.3：URL 去 query、header 值永不落日志） -// ════════════════════════════════════════════════════════════════════════════ - -/// 提取 host（不含 query/path）。解析失败回退 "unknown"。 fn request_host(url: &str) -> String { url::Url::parse(url) .ok() - .and_then(|u| u.host_str().map(str::to_owned)) - .unwrap_or_else(|| "unknown".to_owned()) + .and_then(|url| url.host_str().map(str::to_owned)) + .unwrap_or_else(|| "unknown".into()) } -/// URL 去掉 query string（query 可能含 key，RFC-0014 §4.3）。 -fn request_url_no_query(url: &str) -> String { +pub(crate) fn sanitized_request_url(url: &str) -> String { match url::Url::parse(url) { Ok(mut parsed) => { parsed.set_query(None); @@ -1575,516 +704,3 @@ fn request_url_no_query(url: &str) -> String { Err(_) => url.split('?').next().unwrap_or(url).to_owned(), } } - -/// 请求体大小（字节）。JSON 体序列化一次求长度——只在 debug 事件启用时求值。 -fn request_body_size(request: &HttpRequest) -> u64 { - match &request.body { - HttpBody::Json(value) => serde_json::to_vec(value) - .map(|b| b.len() as u64) - .unwrap_or(0), - HttpBody::Bytes(bytes, _) => bytes.len() as u64, - HttpBody::Empty => 0, - } -} - -/// `AiMuxError` 变体的短名（日志 reason 字段用）。 -fn error_variant(e: &AiMuxError) -> &'static str { - match e { - AiMuxError::ApiCall(_) => "api_call", - AiMuxError::JsonParse(_) => "json_parse", - AiMuxError::InvalidResponseData(_) => "invalid_response_data", - AiMuxError::Tool(_) => "tool", - AiMuxError::InvalidArgument(_) => "invalid_argument", - AiMuxError::InvalidPrompt(_) => "invalid_prompt", - AiMuxError::TokenExpired(_) => "token_expired", - AiMuxError::UnsupportedFunctionality(_) => "unsupported_functionality", - AiMuxError::NoSuchModel { .. } => "no_such_model", - AiMuxError::NoSuchProvider { .. } => "no_such_provider", - AiMuxError::Timeout(_) => "timeout", - AiMuxError::Aborted => "aborted", - AiMuxError::Other(_) => "other", - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn shared_client_is_stable_handle() { - let a = shared_client().expect("client init in test env"); - let b = shared_client().expect("client init in test env"); - assert!(std::ptr::eq(a, b)); - } - - #[test] - fn shared_streaming_client_is_stable_handle() { - let a = shared_streaming_client().expect("client init in test env"); - let b = shared_streaming_client().expect("client init in test env"); - assert!(std::ptr::eq(a, b)); - } - - #[test] - fn shared_and_streaming_are_distinct() { - let a = shared_client().expect("client init in test env"); - let b = shared_streaming_client().expect("client init in test env"); - assert!(!std::ptr::eq(a, b)); - } - - #[test] - fn client_init_error_is_non_retryable_api_call() { - // Shape lock for the #115 mapping: no HTTP exchange, no status, and a - // message that names the failure source (so users can tell an init - // failure apart from a provider-side API call). - let e = client_init_error("simulated TLS backend failure"); - match e { - AiMuxError::ApiCall(call) => { - assert_eq!(call.status_code, None); - assert!(!call.is_retryable); - assert!(call.message.contains("client initialization failed")); - assert!(call.message.contains("simulated TLS backend failure")); - } - other => panic!("expected ApiCall, got {other:?}"), - } - } - - #[test] - fn streaming_config_has_no_response_timeout() { - let t = TimeoutConfig::streaming(); - assert_eq!(t.response_timeout_ms, 0); - assert_ne!(t.connect_timeout_ms, 0); - } - - #[test] - fn default_config_has_response_timeout() { - let t = TimeoutConfig::default(); - assert_ne!(t.response_timeout_ms, 0); - } - - #[test] - fn http_request_is_clone_for_retry() { - // HttpRequest 必须可 Clone —— retry 时从 &request 重建 RequestBuilder。 - let req = HttpRequest { - method: HttpMethod::Post, - url: "https://example.com".to_string(), - headers: vec![("Authorization".to_string(), "Bearer x".to_string())], - body: HttpBody::Json(serde_json::json!({"q": "hi"})), - - abort_signal: None, - call_id: None, - recording_context: None, - }; - let _clone = req.clone(); - assert_eq!(req.method, HttpMethod::Post); - } - - #[test] - fn build_request_builder_rejects_invalid_header_name() { - // A header name containing a space/CR is not a valid token and must be - // rejected rather than silently dropped (the header could be an auth - // header, which the caller must know about). - let client = Client::new(); - let req = HttpRequest { - method: HttpMethod::Get, - url: "https://example.com".to_string(), - headers: vec![("Invalid Header\r\n".to_string(), "secret".to_string())], - body: HttpBody::Empty, - - abort_signal: None, - call_id: None, - recording_context: None, - }; - let err = build_request_builder(&client, &req).unwrap_err(); - assert!(matches!(err, AiMuxError::InvalidArgument(_)), "got {err:?}"); - assert!(err.to_string().contains("invalid header name")); - } - - #[test] - fn build_request_builder_rejects_invalid_header_value() { - // A header value containing a raw CR is invalid and must be rejected. - let client = Client::new(); - let req = HttpRequest { - method: HttpMethod::Get, - url: "https://example.com".to_string(), - headers: vec![( - "Authorization".to_string(), - "Bearer evil\r\nX-Injected: yes".to_string(), - )], - body: HttpBody::Empty, - - abort_signal: None, - call_id: None, - recording_context: None, - }; - let err = build_request_builder(&client, &req).unwrap_err(); - assert!(matches!(err, AiMuxError::InvalidArgument(_)), "got {err:?}"); - assert!(err.to_string().contains("invalid header value")); - } - - #[test] - fn build_request_builder_accepts_valid_headers() { - let client = Client::new(); - let req = HttpRequest { - method: HttpMethod::Post, - url: "https://example.com".to_string(), - headers: vec![ - ("Authorization".to_string(), "Bearer token".to_string()), - ("X-Custom".to_string(), "value".to_string()), - ], - body: HttpBody::Json(serde_json::json!({"q": "hi"})), - - abort_signal: None, - call_id: None, - recording_context: None, - }; - assert!(build_request_builder(&client, &req).is_ok()); - } - - #[tokio::test] - async fn timeout_stream_yields_after_first_chunk_deadline() { - // Inner stream never produces anything → first-chunk deadline fires. - let inner = futures::stream::pending::>().boxed(); - let mut timed = TimeoutBodyStream { - inner, - start: Instant::now(), - total_ms: None, - first_chunk_ms: Some(200), - chunk_ms: None, - abort_signal: None, - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - let item = tokio::time::timeout(Duration::from_secs(5), timed.next()) - .await - .expect("first-chunk deadline must fire within 5s") - .expect("stream must yield an error"); - assert!(matches!(item, Err(AiMuxError::Timeout(_)))); - assert!( - item.unwrap_err() - .to_string() - .contains("first chunk timeout") - ); - } - - #[tokio::test] - async fn timeout_stream_enforces_chunk_idle_deadline() { - // First chunk arrives immediately; the second never comes → the - // chunk-idle deadline fires on the second poll. - let inner = futures::stream::iter(vec![Ok(Bytes::from("first"))]) - .chain(futures::stream::pending::>()) - .boxed(); - let mut timed = TimeoutBodyStream { - inner, - start: Instant::now(), - total_ms: None, - first_chunk_ms: None, - chunk_ms: Some(200), - abort_signal: None, - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - - let first = timed.next().await.expect("first chunk").unwrap(); - assert_eq!(&first[..], b"first"); - - let second = tokio::time::timeout(Duration::from_secs(5), timed.next()) - .await - .expect("chunk idle timeout must fire within 5s") - .expect("stream must yield an error"); - assert!(matches!(second, Err(AiMuxError::Timeout(_)))); - assert!( - second - .unwrap_err() - .to_string() - .contains("chunk idle timeout") - ); - } - - #[tokio::test] - async fn timeout_stream_early_chunks_pass_within_budget() { - let inner = futures::stream::iter(vec![Ok(Bytes::from("a")), Ok(Bytes::from("b"))]).boxed(); - let mut timed = TimeoutBodyStream { - inner, - start: Instant::now(), - total_ms: None, - first_chunk_ms: Some(1000), - chunk_ms: Some(1000), - abort_signal: None, - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - let a = timed.next().await.unwrap().unwrap(); - let b = timed.next().await.unwrap().unwrap(); - assert_eq!(&a[..], b"a"); - assert_eq!(&b[..], b"b"); - assert!(timed.next().await.is_none()); - } - - #[tokio::test] - async fn timeout_stream_is_fused_after_timeout() { - // After a timeout error item the stream must yield None, not an - // endless stream of errors (RFC-0016 review P0). - let inner = futures::stream::pending::>().boxed(); - let mut timed = TimeoutBodyStream { - inner, - start: Instant::now(), - total_ms: None, - first_chunk_ms: Some(100), - chunk_ms: None, - abort_signal: None, - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - let first = tokio::time::timeout(Duration::from_secs(5), timed.next()) - .await - .expect("must fire") - .expect("must be an error"); - assert!(matches!(first, Err(AiMuxError::Timeout(_)))); - let second = tokio::time::timeout(Duration::from_millis(500), timed.next()) - .await - .expect("stream must end right after the timeout error"); - assert!(second.is_none()); - } - - #[tokio::test] - async fn abort_wakes_pending_timeout_stream() { - // Abort while the inner stream is Pending must end the stream - // promptly (event-driven, no 50ms polling). No deadline configured — - // this is the "abort without timeout" streaming path (RFC-0016 - // review S1): the wrapper must still honor the signal. - let signal = AbortSignal::new(); - let inner = futures::stream::pending::>().boxed(); - let mut timed = TimeoutBodyStream { - inner, - start: Instant::now(), - total_ms: None, - first_chunk_ms: None, - chunk_ms: None, - abort_signal: Some(signal.clone()), - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - - let handle = tokio::spawn(async move { - let item = timed.next().await; - (item, timed) - }); - tokio::time::sleep(Duration::from_millis(50)).await; - signal.abort(); - - let (item, mut timed) = tokio::time::timeout(Duration::from_secs(2), handle) - .await - .expect("abort must wake the pending stream within 2s") - .expect("task must finish cleanly"); - let item = item.expect("stream must yield an item"); - assert!(matches!(item, Err(AiMuxError::Aborted)), "got {item:?}"); - // Fused after abort. - let next = timed.next().await; - assert!(next.is_none(), "stream must be fused after abort"); - } - - #[tokio::test] - async fn abort_before_first_poll_fails_fast() { - let signal = AbortSignal::new(); - signal.abort(); - let inner = futures::stream::pending::>().boxed(); - let mut timed = TimeoutBodyStream { - inner, - start: Instant::now(), - total_ms: None, - first_chunk_ms: None, - chunk_ms: None, - abort_signal: Some(signal), - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - let item = timed.next().await.expect("must yield an error"); - assert!(matches!(item, Err(AiMuxError::Aborted))); - assert!( - timed.next().await.is_none(), - "stream must be fused after abort" - ); - } - - #[test] - fn validate_timeout_rejects_overflowing_values() { - // Platform-dependent: on platforms whose `Instant` range is narrower - // (e.g. Unix), huge values must be rejected with InvalidArgument and - // must never panic; on wider ranges (Windows) they are accepted. - let now = Instant::now(); - let huge = RequestTimeout { - total_ms: Some(u64::MAX), - ..Default::default() - }; - match now.checked_add(Duration::from_millis(u64::MAX)) { - None => { - let err = validate_timeout(&huge).unwrap_err(); - assert!(matches!(err, AiMuxError::InvalidArgument(_))); - } - Some(_) => { - assert!(validate_timeout(&huge).is_ok()); - } - } - assert!(validate_timeout(&RequestTimeout::default()).is_ok()); - } - - #[test] - fn next_deadline_kind_selection() { - let now = Instant::now(); - let base = TimeoutBodyStream { - inner: futures::stream::pending().boxed(), - start: now, - total_ms: Some(10_000), - first_chunk_ms: Some(1_000), - chunk_ms: None, - abort_signal: None, - first: true, - last_chunk_at: None, - sleep: None, - abort_wait: None, - done: false, - }; - let (d, k) = base.next_deadline().unwrap(); - assert_eq!(k, TimeoutKind::FirstChunk); - assert_eq!(d, now + Duration::from_millis(1_000)); - - // Total wins when it is earlier. - let s2 = TimeoutBodyStream { - total_ms: Some(500), - first_chunk_ms: Some(1_000), - ..base - }; - let (d2, k2) = s2.next_deadline().unwrap(); - assert_eq!(k2, TimeoutKind::Total); - assert_eq!(d2, now + Duration::from_millis(500)); - } - - #[test] - fn byte_accumulator_preserves_multibyte_across_chunk_split() { - // "中" = E4 B8 AD(3 字节)。拆成两个 chunk 各含半个字符, - // 验证累积后统一解码不产生 U+FFFD(A6 回归)。 - let mut acc = ByteAccumulator::new(RECORD_BODY_CAP); - let ch = "中"; - acc.push(&ch.as_bytes()[..1]); // E4(不完整首字节) - acc.push(&ch.as_bytes()[1..]); // B8 AD(补全该字符) - let decoded = acc.decode().unwrap(); - assert_eq!(decoded, "中", "跨 chunk 多字节字符必须完整保留"); - assert!(!decoded.contains('\u{FFFD}'), "不得出现替换符"); - } - - #[test] - fn byte_accumulator_preserves_emoji_across_chunk_split() { - // 😀 = F0 9F 98 80(4 字节)。拆成 2+2 各半。 - let mut acc = ByteAccumulator::new(RECORD_BODY_CAP); - let b = "😀".as_bytes(); - acc.push(&b[..2]); - acc.push(&b[2..]); - let decoded = acc.decode().unwrap(); - assert_eq!(decoded, "😀"); - assert!(!decoded.contains('\u{FFFD}'), "不得出现替换符"); - } - - #[test] - fn byte_accumulator_truncates_at_cap_and_marks() { - // cap 远小于内容 → 截断到 cap 并追加标记;不 panic。 - let mut acc = ByteAccumulator::new(4); - acc.push(b"0123456789"); - let decoded = acc.decode().unwrap(); - assert_eq!(decoded, "0123…(truncated)"); - } - - #[test] - fn byte_accumulator_truncates_multibyte_without_dangling_replacement() { - // cap 切在多字节字符中间:"中文"(E4 B8 AD E6 96 87)cap=4 → - // E4 B8 AD E6(末尾 E6 是 "文" 的不完整首字节)。 - let mut acc = ByteAccumulator::new(4); - acc.push("中文".as_bytes()); - let decoded = acc.decode().unwrap(); - assert!(decoded.ends_with("…(truncated)"), "应标记截断: {decoded}"); - assert!( - !decoded.contains('\u{FFFD}'), - "截断处不得残留替换符: {decoded}" - ); - assert!(decoded.starts_with("中"), "完整字符保留: {decoded}"); - } - - #[test] - fn byte_accumulator_empty_decodes_to_none() { - let acc = ByteAccumulator::new(RECORD_BODY_CAP); - assert!(acc.decode().is_none()); - } - - // ── M6 ProxyConfig tests ──────────────────────────────────────────────── - - #[test] - fn proxy_config_default_is_no_proxy() { - let config = ProxyConfig::default(); - assert!(config.http_url.is_none()); - assert!(config.https_url.is_none()); - assert!(config.all_url.is_none()); - assert!(config.no_proxy.is_none()); - } - - #[test] - fn apply_proxy_with_all_url_sets_http_and_https() { - let config = ProxyConfig { - all_url: Some("http://proxy.example:8080".into()), - ..Default::default() - }; - let _ = apply_proxy(reqwest::Client::builder(), &config); - } - - #[test] - fn apply_proxy_with_no_proxy_whitelist_builds() { - let config = ProxyConfig { - https_url: Some("http://proxy.example:8080".into()), - no_proxy: Some("localhost,127.0.0.1,.internal.team".into()), - ..Default::default() - }; - let _ = apply_proxy(reqwest::Client::builder(), &config); - } - - #[test] - fn apply_proxy_default_is_noop() { - let config = ProxyConfig::default(); - let _ = apply_proxy(reqwest::Client::builder(), &config); - } - - #[test] - fn apply_proxy_invalid_url_is_silently_skipped() { - let config = ProxyConfig { - http_url: Some("not-a-url".into()), - ..Default::default() - }; - let _ = apply_proxy(reqwest::Client::builder(), &config); - } - - #[test] - fn init_proxy_returns_true_first_time() { - // init_proxy uses a global OnceLock; the first call should succeed. - // Note: once set, it cannot be reset (process-wide), so this test - // asserts only the return value, not the actual client behavior. - let result = init_proxy(ProxyConfig::default()); - // Either true (first call) or false (already set by another test) — - // both are valid. We just assert it doesn't panic. - let _ = result; - } -} diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index d7d4f412..0f6b4de0 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -2,15 +2,21 @@ //! //! Shared utilities for provider implementations. //! -//! Provides HTTP helpers, API key loading, header management, URL utilities, -//! and retry logic — the Rust equivalents of `@ai-sdk/provider-utils`. +//! Provides one-exchange HTTP helpers, response handlers, API key loading, +//! header management, and URL utilities — the Rust equivalents of +//! `@ai-sdk/provider-utils`. Operation retry and timeout live in `aimux-core`. pub mod api_key; +pub mod extract_response_headers; +pub mod get_from_api; +pub mod handle_fetch_error; pub mod headers; pub mod http; pub mod logging; pub mod multipart; -pub mod response; +pub mod post_to_api; +pub mod read_response_with_size_limit; +pub mod response_handler; pub mod retry; pub mod url; /// WebSocket client for realtime provider APIs (RFC-0028). Empty unless the @@ -19,20 +25,17 @@ pub mod url; pub mod ws; pub use api_key::load_api_key; +pub use get_from_api::get_from_api; pub use headers::with_user_agent_suffix; -pub use http::{ - HttpBody, HttpMethod, HttpRequest, HttpResponse, HttpStreamResponse, PoolConfig, ProxyConfig, - RequestTimeout, TimeoutConfig, init_proxy, send, send_stream, send_stream_timed, send_timed, - shared_client, shared_streaming_client, sleep_or_abort, -}; -pub use logging::{body_logging_enabled, init_logging, redact_body}; +pub use http::{HttpBody, HttpRequest, ProxyConfig, init_proxy, shared_client, sleep_or_abort}; +pub use logging::{body_logging_enabled, init_logging, redact_body, redact_error_context}; pub use multipart::{MultipartForm, media_type_to_extension}; -pub use response::{ - DEFAULT_ERROR_STRUCTURE, ErrorStructure, error_for_status, parse_provider_error, - parse_stream_error, -}; -pub use retry::{ - RetryConfig, get_retry_delay_ms, get_retry_delay_ms_with_jitter, parse_retry_after, - retry_with_exponential_backoff, retry_with_exponential_backoff_respecting_retry_headers, +pub use post_to_api::{post_form_data_to_api, post_json_to_api, post_to_api}; +pub use response_handler::{ + ProviderErrorParts, ResponseHandler, ResponseHandlerInput, ResponseHandlerOutput, + create_binary_response_handler, create_event_source_response_handler, + create_json_error_response_handler, create_json_response_handler, + create_status_code_error_response_handler, stream_error_api_call, }; +pub use retry::RetryConfig; pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt}; diff --git a/aimux-provider-utils/src/logging.rs b/aimux-provider-utils/src/logging.rs index 29f19ade..3b28790e 100644 --- a/aimux-provider-utils/src/logging.rs +++ b/aimux-provider-utils/src/logging.rs @@ -61,7 +61,7 @@ pub fn init_logging(level: &str) { init_once(Some(level)); } -/// Lazy env-driven auto-init, called at the HTTP throat (`send`/`send_stream`). +/// Lazy env-driven auto-init, called by the shared API-call primitive. /// /// Cheap when already initialized (`Once::is_completed` is a single atomic /// load). Registers a subscriber only when an `AIMUX_LOG*` env var is present @@ -117,44 +117,82 @@ fn truncate(body: &str) -> &str { } /// Redact a body for trace logging: JSON object keys whose lowercased name -/// contains `authorization` / `api-key` / `apikey` / `key` / `token` have -/// their (string) values replaced with `***`. Non-JSON bodies pass through -/// unchanged. Always truncated to [`BODY_LOG_LIMIT`]. +/// match the shared recording sensitivity policy have their values replaced +/// with `[REDACTED]`. Non-JSON bodies pass through unchanged. Always truncated +/// to [`BODY_LOG_LIMIT`]. #[must_use] pub fn redact_body(body: &str) -> String { let truncated = truncate(body); match serde_json::from_str::(truncated) { - Ok(value) => { - serde_json::to_string(&redact_value(value)).unwrap_or_else(|_| truncated.to_owned()) - } + Ok(value) => serde_json::to_string(&redact_error_context(value)) + .unwrap_or_else(|_| truncated.to_owned()), Err(_) => truncated.to_owned(), } } -fn redact_value(value: serde_json::Value) -> serde_json::Value { +const ERROR_CONTEXT_MAX_DEPTH: usize = 8; +const ERROR_CONTEXT_MAX_ITEMS: usize = 128; +const ERROR_CONTEXT_MAX_STRING: usize = 4096; + +/// Redact and bound structured values before exposing them in an API error. +#[must_use] +pub fn redact_error_context(value: serde_json::Value) -> serde_json::Value { + redact_value(value, 0) +} + +fn redact_value(value: serde_json::Value, depth: usize) -> serde_json::Value { use serde_json::Value; + if depth >= ERROR_CONTEXT_MAX_DEPTH { + return Value::String("[MAX_DEPTH]".to_owned()); + } match value { Value::Object(map) => { - let mut out = serde_json::Map::with_capacity(map.len()); - for (k, v) in map { - if is_sensitive_key(&k) { - out.insert(k, Value::String("***".to_owned())); + let mut out = serde_json::Map::with_capacity(map.len().min(ERROR_CONTEXT_MAX_ITEMS)); + for (index, (key, value)) in map.into_iter().enumerate() { + if index == ERROR_CONTEXT_MAX_ITEMS { + out.insert("__truncated__".into(), Value::Bool(true)); + break; + } + if aimux_core::recording::is_sensitive_key(&key) { + out.insert(key, Value::String("[REDACTED]".to_owned())); } else { - out.insert(k, redact_value(v)); + out.insert(key, redact_value(value, depth + 1)); } } Value::Object(out) } - Value::Array(items) => Value::Array(items.into_iter().map(redact_value).collect()), + Value::Array(items) => { + let original_len = items.len(); + let mut items: Vec<_> = items + .into_iter() + .take(ERROR_CONTEXT_MAX_ITEMS) + .map(|value| redact_value(value, depth + 1)) + .collect(); + if original_len > ERROR_CONTEXT_MAX_ITEMS { + items.push(serde_json::json!({ + "__truncated_items__": original_len - ERROR_CONTEXT_MAX_ITEMS + })); + } + Value::Array(items) + } + Value::String(value) if value.len() > ERROR_CONTEXT_MAX_STRING => { + Value::String(format!("[STRING {} bytes]", value.len())) + } other => other, } } -fn is_sensitive_key(key: &str) -> bool { - let k = key.to_ascii_lowercase(); - ["authorization", "api-key", "apikey", "key", "token"] - .iter() - .any(|needle| k.contains(needle)) +/// Sanitize request values before storing them in a public `ApiCallError`. +#[must_use] +pub fn redact_request_values(body: &crate::http::HttpBody) -> serde_json::Value { + match body { + crate::http::HttpBody::Json(value) => redact_error_context(value.clone()), + crate::http::HttpBody::Bytes(bytes, content_type) => serde_json::json!({ + "content_type": content_type, + "length": bytes.len(), + }), + crate::http::HttpBody::Empty => serde_json::json!({}), + } } /// Small helper used by tests to capture formatted output. @@ -236,11 +274,11 @@ mod tests { assert!(!out.contains("sk-secret"), "api_key value leaked: {out}"); assert!(!out.contains("Bearer xyz"), "authorization leaked: {out}"); assert!( - out.contains("\"api_key\":\"***\""), + out.contains("\"api_key\":\"[REDACTED]\""), "api_key not redacted: {out}" ); assert!( - out.contains("\"authorization\":\"***\""), + out.contains("\"authorization\":\"[REDACTED]\""), "authorization not redacted: {out}" ); assert!( @@ -256,6 +294,15 @@ mod tests { assert!(out.len() <= BODY_LOG_LIMIT); } + #[test] + fn truncation_respects_utf8_char_boundaries() { + // A 3-byte char straddling the limit must be dropped whole, not split + // (a byte-index slice would panic mid-char). + let body = format!("{}中中", "x".repeat(BODY_LOG_LIMIT - 1)); + let out = redact_body(&body); + assert_eq!(out, "x".repeat(BODY_LOG_LIMIT - 1)); + } + #[test] fn redact_non_json_passthrough() { let body = "not json at all"; diff --git a/aimux-provider-utils/src/multipart.rs b/aimux-provider-utils/src/multipart.rs index 1c46d886..65432f92 100644 --- a/aimux-provider-utils/src/multipart.rs +++ b/aimux-provider-utils/src/multipart.rs @@ -14,6 +14,7 @@ use aimux_core::AiMuxError; pub struct MultipartForm { boundary: String, parts: Vec, + values: serde_json::Map, } impl MultipartForm { @@ -27,6 +28,7 @@ impl MultipartForm { Self { boundary, parts: Vec::new(), + values: serde_json::Map::new(), } } @@ -48,6 +50,10 @@ impl MultipartForm { ); self.parts.extend_from_slice(value.as_bytes()); self.parts.extend_from_slice(b"\r\n"); + self.values.insert( + name.to_string(), + serde_json::Value::String(value.to_string()), + ); Ok(self) } @@ -81,17 +87,34 @@ impl MultipartForm { .extend_from_slice(format!("Content-Type: {media_type}\r\n\r\n").as_bytes()); self.parts.extend_from_slice(data); self.parts.extend_from_slice(b"\r\n"); + self.values.insert( + name.to_string(), + serde_json::json!({ + "filename": filename, + "mediaType": media_type, + "size": data.len(), + }), + ); Ok(self) } /// Finalize the body, returning the raw bytes and the content-type header /// value. #[must_use] - pub fn finish(mut self) -> (Vec, String) { + pub fn finish(self) -> (Vec, String) { + let (bytes, content_type, _) = self.into_parts(); + (bytes, content_type) + } + + pub(crate) fn into_parts(mut self) -> (Vec, String, serde_json::Value) { self.parts .extend_from_slice(format!("--{}--\r\n", self.boundary).as_bytes()); let content_type = format!("multipart/form-data; boundary={}", self.boundary); - (self.parts, content_type) + ( + self.parts, + content_type, + serde_json::Value::Object(self.values), + ) } } diff --git a/aimux-provider-utils/src/post_to_api.rs b/aimux-provider-utils/src/post_to_api.rs new file mode 100644 index 00000000..5a268d20 --- /dev/null +++ b/aimux-provider-utils/src/post_to_api.rs @@ -0,0 +1,261 @@ +//! Single-exchange POST helpers aligned with AI SDK. + +use std::time::Duration; + +use aimux_core::{AiMuxError, ApiCallError}; + +use crate::handle_fetch_error::handle_fetch_error; +use crate::http::{HttpBody, HttpMethod, HttpRequest, PreparedRequest}; +use crate::multipart::MultipartForm; +use crate::response_handler::{ResponseHandler, ResponseHandlerInput, ResponseHandlerOutput}; + +/// POST a JSON request. The body type is fixed by this signature. +/// +/// # Errors +/// +/// Returns transport, response-handler, or caller-abort failures. +pub async fn post_json_to_api( + request: HttpRequest, + body: serde_json::Value, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + let request_body_values = crate::logging::redact_error_context(body.clone()); + call_to_api( + request.prepare(HttpMethod::Post, HttpBody::Json(body)), + request_body_values, + successful_response_handler, + failed_response_handler, + ) + .await +} + +/// POST a multipart/form-data request. Binary fields are represented in +/// error context by filename, media type, and byte length rather than bytes. +/// +/// # Errors +/// +/// Returns transport, response-handler, or caller-abort failures. +pub async fn post_form_data_to_api( + request: HttpRequest, + form_data: MultipartForm, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + let (content, content_type, values) = form_data.into_parts(); + call_to_api( + request.prepare(HttpMethod::Post, HttpBody::Bytes(content, content_type)), + crate::logging::redact_error_context(values), + successful_response_handler, + failed_response_handler, + ) + .await +} + +/// POST an explicitly prepared raw body. +/// +/// # Errors +/// +/// Returns transport, response-handler, or caller-abort failures. +pub async fn post_to_api( + request: HttpRequest, + body: HttpBody, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + let request_body_values = crate::logging::redact_request_values(&body); + call_to_api( + request.prepare(HttpMethod::Post, body), + request_body_values, + successful_response_handler, + failed_response_handler, + ) + .await +} + +/// Default whole-exchange response timeout for non-streaming exchanges +/// (connect through fully read body), matching the pre-0.4 shared-client +/// behavior. Streaming exchanges are exempt: their body outlives the call. +/// Tripping it yields a retryable error — a hung exchange is a transport +/// fault, and the Core operation deadline (`AiMuxError::Timeout`) remains +/// the non-retryable outer bound. +const EXCHANGE_RESPONSE_TIMEOUT: Duration = Duration::from_secs(30); + +pub(crate) async fn call_to_api( + request: PreparedRequest, + request_body_values: serde_json::Value, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + call_to_api_with_deadline( + request, + request_body_values, + successful_response_handler, + failed_response_handler, + EXCHANGE_RESPONSE_TIMEOUT, + ) + .await +} + +pub(crate) async fn call_to_api_with_deadline( + request: PreparedRequest, + request_body_values: serde_json::Value, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, + response_timeout: Duration, +) -> Result, AiMuxError> { + let url = crate::http::sanitized_request_url(&request.url); + if successful_response_handler.is_streaming() { + return execute_exchange( + request, + request_body_values, + successful_response_handler, + failed_response_handler, + ) + .await; + } + let exchange = execute_exchange( + request, + request_body_values.clone(), + successful_response_handler, + failed_response_handler, + ); + match tokio::time::timeout(response_timeout, exchange).await { + Ok(result) => result, + Err(_) => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..ApiCallError::new( + format!( + "Request exceeded the {}s exchange response timeout", + response_timeout.as_secs() + ), + url, + request_body_values, + ) + }))), + } +} + +async fn execute_exchange( + request: PreparedRequest, + request_body_values: serde_json::Value, + successful_response_handler: ResponseHandler, + failed_response_handler: ResponseHandler, +) -> Result, AiMuxError> { + crate::logging::auto_init_from_env(); + let url = crate::http::sanitized_request_url(&request.url); + let response = crate::http::send_request_once(&request) + .await + .map_err(|error| handle_fetch_error(error, &url, &request_body_values))?; + let status = response.status().as_u16(); + let response_headers = + crate::extract_response_headers::extract_response_headers(response.headers()); + let input = ResponseHandlerInput { + url: url.clone(), + request_body_values: request_body_values.clone(), + response, + abort_signal: request.abort_signal.clone(), + }; + + if !(200..300).contains(&status) { + return match failed_response_handler.handle(input).await { + Ok(output) => Err(output.value), + Err( + error @ (AiMuxError::ApiCall(_) | AiMuxError::Aborted(_) | AiMuxError::Timeout(_)), + ) => Err(error), + Err(error) => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_headers: Some(response_headers), + // A handler that cannot parse a 429/503 body must not erase + // the status's retryability — retry classification comes from + // the HTTP status, not from the parse outcome. + is_retryable: aimux_core::error::is_retryable_status(status), + ..ApiCallError::new( + format!("Failed to process error response: {error}"), + url, + request_body_values, + ) + }))), + }; + } + + match successful_response_handler.handle(input).await { + Ok(output) => Ok(output), + Err(error @ (AiMuxError::ApiCall(_) | AiMuxError::Aborted(_) | AiMuxError::Timeout(_))) => { + Err(error) + } + Err(error) => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_headers: Some(response_headers), + ..ApiCallError::new( + format!("Failed to process successful response: {error}"), + url, + request_body_values, + ) + }))), + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use wiremock::matchers::method; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + use super::*; + use crate::response_handler::{ + create_event_source_response_handler, create_json_response_handler, + create_status_code_error_response_handler, + }; + + /// A response slower than the exchange deadline must fail as a retryable + /// transport error — not hang, and not use the non-retryable + /// `AiMuxError::Timeout` reserved for Core operation deadlines. + #[tokio::test] + async fn slow_response_trips_the_exchange_timeout_as_retryable() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"ok": true})) + .set_delay(Duration::from_millis(500)), + ) + .mount(&server) + .await; + + let request = HttpRequest { + url: server.uri(), + headers: vec![], + abort_signal: None, + call_id: None, + recording_context: None, + }; + let error = call_to_api_with_deadline( + request.prepare(HttpMethod::Post, HttpBody::Json(serde_json::json!({}))), + serde_json::json!({}), + create_json_response_handler::(), + create_status_code_error_response_handler(), + Duration::from_millis(50), + ) + .await + .unwrap_err(); + + match error { + AiMuxError::ApiCall(detail) => { + assert!(detail.is_retryable, "exchange timeout must be retryable"); + assert!(detail.message.contains("exchange response timeout")); + } + other => panic!("expected retryable ApiCall, got {other:?}"), + } + } + + /// Streaming handlers are exempt from the exchange response timeout. + #[test] + fn event_source_handler_is_marked_streaming() { + let handler = create_event_source_response_handler::(); + assert!(handler.is_streaming()); + let plain = create_json_response_handler::(); + assert!(!plain.is_streaming()); + } +} diff --git a/aimux-provider-utils/src/read_response_with_size_limit.rs b/aimux-provider-utils/src/read_response_with_size_limit.rs new file mode 100644 index 00000000..74909c17 --- /dev/null +++ b/aimux-provider-utils/src/read_response_with_size_limit.rs @@ -0,0 +1,134 @@ +//! Bounded response reads. + +use bytes::Bytes; +use futures::StreamExt; + +use aimux_core::{AiMuxError, ApiCallError}; + +/// Default maximum buffered response size (2 GiB, matching AI SDK). +pub const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024; + +/// Read a response incrementally and fail before unbounded allocation. +/// +/// # Errors +/// +/// Returns an API-call error when the body exceeds `max_bytes` or cannot be +/// read, and an abort error when the caller cancels. +pub async fn read_response_with_size_limit( + response: reqwest::Response, + url: &str, + request_body_values: &serde_json::Value, + max_bytes: usize, + abort_signal: Option<&aimux_core::AbortSignal>, +) -> Result { + let status = response.status().as_u16(); + let response_headers = + crate::extract_response_headers::extract_response_headers(response.headers()); + if let Some(length) = response.content_length() + && length > max_bytes as u64 + { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_headers: Some(response_headers), + is_retryable: aimux_core::error::is_retryable_status(status), + ..ApiCallError::new( + format!("Response exceeded maximum size of {max_bytes} bytes"), + url, + request_body_values.clone(), + ) + }))); + } + + let mut stream = response.bytes_stream(); + let mut bytes = Vec::new(); + loop { + let next = match abort_signal { + Some(signal) => tokio::select! { + biased; + () = signal.cancelled() => return Err(AiMuxError::from_abort_signal(signal)), + next = stream.next() => next, + }, + None => stream.next().await, + }; + let Some(chunk) = next else { + break; + }; + let chunk = chunk.map_err(|error| { + AiMuxError::ApiCall(Box::new(ApiCallError { + // Transport failure mid-body: retryable, and carrying NO + // status — the exchange died before a complete response, and + // an "HTTP 200: failed to read body" would contradict the + // transport-error contract (pre-response transport path and + // the SSE body path both report status-less errors). + status_code: None, + response_headers: Some(response_headers.clone()), + is_retryable: true, + ..ApiCallError::new( + format!("Failed to read response body: {error}"), + url, + request_body_values.clone(), + ) + })) + })?; + if bytes.len().saturating_add(chunk.len()) > max_bytes { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_headers: Some(response_headers), + is_retryable: aimux_core::error::is_retryable_status(status), + ..ApiCallError::new( + format!("Response exceeded maximum size of {max_bytes} bytes"), + url, + request_body_values.clone(), + ) + }))); + } + bytes.extend_from_slice(&chunk); + } + Ok(Bytes::from(bytes)) +} + +/// Best-effort read of an error body: keep the first `max_bytes`, mark +/// truncation instead of failing, and surface what was collected even when +/// the connection dies mid-read. The provider's error text is diagnosis +/// evidence — replacing it with a size-limit error would destroy the actual +/// failure reason. +/// +/// # Errors +/// +/// Returns an abort error when the caller cancels; never fails on read or +/// size problems. +pub(crate) async fn read_error_body_truncated( + response: reqwest::Response, + max_bytes: usize, + abort_signal: Option<&aimux_core::AbortSignal>, +) -> Result<(Vec, bool), AiMuxError> { + let mut stream = response.bytes_stream(); + let mut bytes = Vec::new(); + let mut truncated = false; + loop { + let next = match abort_signal { + Some(signal) => tokio::select! { + biased; + () = signal.cancelled() => return Err(AiMuxError::from_abort_signal(signal)), + next = stream.next() => next, + }, + None => stream.next().await, + }; + let Some(chunk) = next else { + break; + }; + // A dead connection mid-error-body: surface what we have. + let Ok(chunk) = chunk else { + break; + }; + let remaining = max_bytes.saturating_sub(bytes.len()); + if chunk.len() > remaining { + bytes.extend_from_slice(&chunk[..remaining]); + truncated = true; + // Stop reading: draining an unbounded error body is a DoS vector. + break; + } + bytes.extend_from_slice(&chunk); + } + Ok((bytes, truncated)) +} diff --git a/aimux-provider-utils/src/response.rs b/aimux-provider-utils/src/response.rs deleted file mode 100644 index 01d77458..00000000 --- a/aimux-provider-utils/src/response.rs +++ /dev/null @@ -1,166 +0,0 @@ -//! Response handling helpers. - -use aimux_core::{AiMuxError, ApiCallError}; - -/// Provider error structure (allows each provider to define its own error JSON shape). -pub struct ErrorStructure { - /// JSON path to the error message (e.g. `["error", "message"]`). - pub message_path: &'static [&'static str], - /// JSON path to the error type (e.g. `["error", "type"]`). - pub type_path: &'static [&'static str], -} - -/// Default OpenAI-compatible error structure: `{ "error": { "message": "...", "type": "..." } }`. -pub const DEFAULT_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "type"], -}; - -/// Parse an HTTP error response into an `AiMuxError`. -/// -/// Every status becomes an `ApiCall` error; the classification is the -/// `status_code` field (401 auth, 404 model, 429 rate limit, …), read by the -/// caller the way the AI SDK reads `APICallError.statusCode`. -/// Retryability (408/409/429/5xx) is stored in `is_retryable` at -/// construction, as the AI SDK stores `APICallError.isRetryable`. -#[must_use] -pub fn parse_provider_error(status: u16, body: &str, structure: &ErrorStructure) -> AiMuxError { - // Empty: `ApiCallError`'s Display shows the status on its own. - let mut message = String::new(); - let mut provider_code = None; - - if let Ok(val) = serde_json::from_str::(body) { - // Try to extract message. - let mut current = &val; - for key in structure.message_path { - if let Some(v) = current.get(key) { - current = v; - } else { - current = &val; - break; - } - } - if let Some(msg) = current.as_str() { - message = msg.to_string(); - } else if !body.is_empty() { - // The configured message path didn't land on a string value (e.g. - // xAI's flat responses-API shape `{"code","error"}` where `error` - // is itself a string, not an object). Surface the raw body so the - // provider error still carries the upstream payload instead of a - // bare "HTTP {status}" placeholder. - message = body.to_string(); - } - - // Try to extract type. - let mut current = &val; - for key in structure.type_path { - if let Some(v) = current.get(key) { - current = v; - } else { - current = &val; - break; - } - } - if let Some(t) = current.as_str() { - provider_code = Some(t.to_string()); - } - } else if !body.is_empty() { - message = body.to_string(); - } - - // Raw evidence (`APICallError.responseBody`): the AI SDK keeps the body on - // every error branch — parsed, unparseable or empty — so consumers never - // have to ask whether it is there. Same here: whenever a body arrived, it - // is kept verbatim, even when `message` is that same body. - let response_body = (!body.is_empty()).then(|| body.to_string()); - error_for_status(status, provider_code, message, None, response_body) -} - -/// Build the `ApiCall` error for a failed HTTP response. -/// -/// There is no status → variant mapping any more: the classification *is* -/// the `status_code` field, read by consumers exactly as the AI SDK reads -/// `APICallError.statusCode`. The retry hint is stored only when the provider -/// actually sent one (`retry-after-ms`/`retry-after`) — no fallback is -/// fabricated; a hint-less 429 gets the retry loop's exponential backoff. -#[must_use] -pub fn error_for_status( - status: u16, - provider_code: Option, - message: String, - retry_after_ms: Option, - response_body: Option, -) -> AiMuxError { - AiMuxError::ApiCall(ApiCallError { - status_code: Some(status), - provider_code, - message, - response_body, - retry_after_ms, - // Stored at construction (`APICallError.isRetryable`), aligned with - // the AI SDK response policy: 408/409/429/5xx retry, other 4xx don't - // (RFC-0016 alignment goal). - is_retryable: matches!(status, 408 | 409 | 429) || status >= 500, - ..Default::default() - }) -} - -/// Convert a stream-error JSON object (the `{"message","type","code",…}` payload -/// carried by an SSE `error` event) into an `AiMuxError`. -/// -/// Shares [`error_for_status`] with [`parse_provider_error`], so the same status -/// yields the same variant — and the same retryability — on both paths (M3). -/// -/// `code` is a *machine-readable provider code*: OpenAI-shaped payloads put a -/// string there (`"rate_limit_exceeded"`), Google-shaped ones an HTTP status. -/// Only a number in the HTTP range is read as a status; a string code lands in -/// `provider_code`. Reading a string `code` as a status was the M3 bug, and so -/// was defaulting a missing one to 500 — a mid-stream error arrives on a -/// *successful* response, so "no status" is the truth. -#[must_use] -pub fn parse_stream_error(err_obj: &serde_json::Value) -> AiMuxError { - let str_field = |k: &str| { - err_obj - .get(k) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - }; - let message = str_field("message").unwrap_or_else(|| "Unknown stream error".to_string()); - let code = err_obj.get("code"); - let status = code - .and_then(serde_json::Value::as_u64) - .filter(|c| (100..=599).contains(c)) - .map(|c| c as u16); - let provider_code = code - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string) - .or_else(|| str_field("type")); - // Retry hint from the payload when the provider sends one (ms wins over - // whole seconds); absent means absent — nothing is fabricated. - let retry_after_ms = err_obj - .get("retry_after_ms") - .and_then(serde_json::Value::as_u64) - .or_else(|| { - err_obj - .get("retry_after") - .and_then(serde_json::Value::as_u64) - .map(|s| s * 1000) - }); - - match status { - Some(status) => error_for_status( - status, - provider_code, - message, - retry_after_ms, - Some(err_obj.to_string()), - ), - None => AiMuxError::ApiCall(ApiCallError { - provider_code, - message, - response_body: Some(err_obj.to_string()), - retry_after_ms, - ..Default::default() - }), - } -} diff --git a/aimux-provider-utils/src/response_handler.rs b/aimux-provider-utils/src/response_handler.rs new file mode 100644 index 00000000..7c86ffc0 --- /dev/null +++ b/aimux-provider-utils/src/response_handler.rs @@ -0,0 +1,425 @@ +//! Successful and failed response handlers, aligned with AI SDK. + +use std::future::Future; +use std::pin::Pin; + +use bytes::Bytes; +use futures::{StreamExt, stream::BoxStream}; +use serde::de::DeserializeOwned; + +use aimux_core::{AiMuxError, ApiCallError}; + +use crate::extract_response_headers::extract_response_headers; +use crate::read_response_with_size_limit::{ + DEFAULT_MAX_DOWNLOAD_SIZE, read_response_with_size_limit, +}; + +/// Input supplied to one response handler. +pub struct ResponseHandlerInput { + pub url: String, + pub request_body_values: serde_json::Value, + pub response: reqwest::Response, + pub abort_signal: Option, +} + +/// A parsed value plus response metadata. +#[derive(Debug)] +pub struct ResponseHandlerOutput { + pub value: T, + pub raw_value: Option, + pub response_headers: Option>, +} + +type HandlerFuture = + Pin, AiMuxError>> + Send>>; +type HandlerFn = dyn FnOnce(ResponseHandlerInput) -> HandlerFuture + Send; + +/// Provider-specific fields extracted from a failed response body. +pub struct ProviderErrorParts { + pub message: String, + pub provider_code: Option, +} + +/// One-shot async response handler. Exactly one of the successful/failed +/// handlers passed to an API call is consumed. +pub struct ResponseHandler { + handler: Box>, + /// Whether the handler keeps the response body open as a stream. A + /// streaming exchange must not be bounded by the per-exchange response + /// timeout in `call_to_api`. + streaming: bool, +} + +impl ResponseHandler { + /// Create a handler from an async closure. + pub fn new(handler: F) -> Self + where + F: FnOnce(ResponseHandlerInput) -> Fut + Send + 'static, + Fut: Future, AiMuxError>> + Send + 'static, + { + Self { + handler: Box::new(move |input| Box::pin(handler(input))), + streaming: false, + } + } + + /// Mark the handler as streaming: it hands the response body onward as a + /// stream instead of reading it to completion, so the per-exchange + /// response timeout must not cover body consumption. + #[must_use] + pub fn streaming(mut self) -> Self { + self.streaming = true; + self + } + + pub(crate) fn is_streaming(&self) -> bool { + self.streaming + } + + /// Consume the handler and process one response. + /// + /// # Errors + /// + /// Returns the handler's contextual API, parse, or caller-abort failure. + pub async fn handle( + self, + input: ResponseHandlerInput, + ) -> Result, AiMuxError> { + (self.handler)(input).await + } +} + +/// Parse a successful JSON response using the endpoint's response type. +#[must_use] +pub fn create_json_response_handler() -> ResponseHandler +where + T: DeserializeOwned + Send + 'static, +{ + ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let headers = extract_response_headers(input.response.headers()); + let body = read_response_with_size_limit( + input.response, + &input.url, + &input.request_body_values, + DEFAULT_MAX_DOWNLOAD_SIZE, + input.abort_signal.as_ref(), + ) + .await?; + let raw_value = serde_json::from_slice::(&body).map_err(|error| { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_body: Some(String::from_utf8_lossy(&body).into_owned()), + response_headers: Some(headers.clone()), + ..ApiCallError::new( + format!("Invalid JSON response: {error}"), + input.url.clone(), + input.request_body_values.clone(), + ) + })) + })?; + let value = serde_json::from_value(raw_value.clone()).map_err(|error| { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_body: Some(String::from_utf8_lossy(&body).into_owned()), + response_headers: Some(headers.clone()), + ..ApiCallError::new( + format!("Invalid JSON response: {error}"), + input.url, + input.request_body_values, + ) + })) + })?; + Ok(ResponseHandlerOutput { + value, + raw_value: Some(raw_value), + response_headers: Some(headers), + }) + }) +} + +/// Public error bodies are capped at this size: `ApiCallError` crosses the +/// FFI as a serialized string and is persisted by recordings, so the body +/// must stay bounded. +const ERROR_BODY_PUBLIC_CAP: usize = 64 * 1024; + +/// Parse bound for error bodies: enough headroom that an oversize-but-valid +/// provider error JSON still reaches the mapper (a body truncated mid-JSON +/// can never parse), while staying bounded — draining an unbounded error +/// body is a DoS vector. +const ERROR_BODY_PARSE_CAP: usize = 1024 * 1024; + +/// Decode an error body for the public `response_body`, enforcing +/// [`ERROR_BODY_PUBLIC_CAP`] *after* lossy UTF-8 decoding: U+FFFD expands an +/// invalid byte to 3 bytes, so capping raw bytes alone can still produce a +/// ~192 KiB string. +fn error_body_string(body: &[u8], read_truncated: bool) -> String { + let mut value = String::from_utf8_lossy(body).into_owned(); + let mut truncated = read_truncated; + if value.len() > ERROR_BODY_PUBLIC_CAP { + let mut end = ERROR_BODY_PUBLIC_CAP; + while !value.is_char_boundary(end) { + end -= 1; + } + value.truncate(end); + truncated = true; + } + if truncated { + // A cut can land inside a replacement run; drop partial noise before + // the marker. + while value.ends_with('\u{FFFD}') { + value.pop(); + } + value.push_str("…(truncated)"); + } + value +} + +/// Parse an error JSON and map its data to a provider message and code. +#[must_use] +pub fn create_json_error_response_handler( + error_to_message_and_code: F, +) -> ResponseHandler +where + F: Fn(&serde_json::Value) -> ProviderErrorParts + Send + 'static, +{ + ResponseHandler::new(move |input| async move { + let status = input.response.status().as_u16(); + let status_text = input + .response + .status() + .canonical_reason() + .unwrap_or_default() + .to_string(); + let headers = extract_response_headers(input.response.headers()); + // Best-effort, truncating read: an oversize or dying error body must + // not replace the provider's actual error message. + let (body, truncated) = crate::read_response_with_size_limit::read_error_body_truncated( + input.response, + ERROR_BODY_PARSE_CAP, + input.abort_signal.as_ref(), + ) + .await?; + let response_body = error_body_string(&body, truncated); + let parsed_data = serde_json::from_slice::(&body).ok(); + let parts = parsed_data + .as_ref() + .map(error_to_message_and_code) + .unwrap_or_else(|| ProviderErrorParts { + message: status_text.clone(), + provider_code: None, + }); + let message = if parts.message.trim().is_empty() { + status_text + } else { + parts.message + }; + let error = AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + provider_code: parts.provider_code, + response_body: Some(response_body), + response_headers: Some(headers.clone()), + data: parsed_data.map(crate::logging::redact_error_context), + is_retryable: aimux_core::error::is_retryable_status(status), + ..ApiCallError::new(message, input.url, input.request_body_values) + })); + Ok(ResponseHandlerOutput { + value: error, + raw_value: None, + response_headers: Some(headers), + }) + }) +} + +/// Build the `ApiCallError` for an error event delivered *inside* a stream. +/// +/// The event rides on an already-successful response, so no HTTP status is +/// fabricated: `status_code` is set only when the payload itself +/// carried a numeric HTTP status, and retryability derives from that status +/// alone. A `retry_after_ms` / `retry_after` field in the payload is +/// surfaced through the response-headers channel so +/// `AiMuxError::retry_after_hint()` observes it. +/// +/// Callers reach this from inside a live stream holding the *raw* request +/// body and URL they built — unlike the exchange path, where `post_*_to_api` +/// redacts before any handler runs. Redaction and URL sanitization therefore +/// happen here, so credentials and data URLs never leak into the public +/// error (or the recordings that persist it). +#[must_use] +pub fn stream_error_api_call( + message: impl Into, + provider_code: Option, + status_code: Option, + error_payload: &serde_json::Value, + url: impl Into, + request_body_values: serde_json::Value, + mut response_headers: std::collections::HashMap, +) -> AiMuxError { + // Real headers win; the payload hint fills in only when absent (an + // in-stream rate-limit error rides on a 200 whose headers carry none). + if !response_headers.contains_key("retry-after-ms") + && !response_headers.contains_key("retry-after") + { + if let Some(ms) = error_payload + .get("retry_after_ms") + .and_then(serde_json::Value::as_f64) + { + response_headers.insert("retry-after-ms".into(), ms.to_string()); + } else if let Some(seconds) = error_payload + .get("retry_after") + .and_then(serde_json::Value::as_f64) + { + response_headers.insert("retry-after".into(), seconds.to_string()); + } + } + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code, + provider_code, + response_body: Some(error_payload.to_string()), + response_headers: Some(response_headers), + data: Some(crate::logging::redact_error_context(error_payload.clone())), + is_retryable: status_code.is_some_and(aimux_core::error::is_retryable_status), + ..ApiCallError::new( + message, + crate::http::sanitized_request_url(&url.into()), + crate::logging::redact_error_context(request_body_values), + ) + })) +} + +/// Read a successful response as bytes. +#[must_use] +pub fn create_binary_response_handler() -> ResponseHandler { + ResponseHandler::new(|input| async move { + let headers = extract_response_headers(input.response.headers()); + let value = read_response_with_size_limit( + input.response, + &input.url, + &input.request_body_values, + DEFAULT_MAX_DOWNLOAD_SIZE, + input.abort_signal.as_ref(), + ) + .await?; + Ok(ResponseHandlerOutput { + value, + raw_value: None, + response_headers: Some(headers), + }) + }) +} + +/// Parse an SSE response and deserialize each `data:` field as `T`. +#[must_use] +pub fn create_event_source_response_handler() +-> ResponseHandler>> +where + T: DeserializeOwned + Send + 'static, +{ + ResponseHandler::new(|input| async move { + let headers = extract_response_headers(input.response.headers()); + let output_headers = headers.clone(); + let url = input.url; + let request_body_values = input.request_body_values; + let signal = input.abort_signal; + let body_url = url.clone(); + let body_request_body_values = request_body_values.clone(); + let body = input.response.bytes_stream().map(move |result| { + result.map_err(|error| { + AiMuxError::ApiCall(Box::new(ApiCallError { + response_headers: Some(headers.clone()), + is_retryable: true, + ..ApiCallError::new( + error.to_string(), + body_url.clone(), + body_request_body_values.clone(), + ) + })) + }) + }); + + let stream_headers = output_headers.clone(); + let stream_url = url.clone(); + let stream_request_body_values = request_body_values.clone(); + let value: BoxStream<'static, Result> = Box::pin(async_stream::stream! { + let events = aimux_stream::SseStream::new(body); + futures::pin_mut!(events); + loop { + let item = match signal.as_ref() { + Some(signal) => tokio::select! { + biased; + () = signal.cancelled() => { + yield Err(AiMuxError::from_abort_signal(signal)); + break; + } + item = events.next() => item, + }, + None => events.next().await, + }; + let Some(item) = item else { + break; + }; + match item { + Ok(event) if event.data == "[DONE]" => continue, + Ok(event) => yield serde_json::from_str::(&event.data).map_err(AiMuxError::from), + Err(aimux_stream::SseError::Stream(message)) => { + // Preserve response transport failures as ApiCallError + // items. Framing/parser failures remain JsonParse below. + yield Err(AiMuxError::ApiCall(Box::new(ApiCallError { + response_headers: Some(stream_headers.clone()), + is_retryable: true, + ..ApiCallError::new( + message, + stream_url.clone(), + stream_request_body_values.clone(), + ) + }))); + break; + } + Err(error) => yield Err(AiMuxError::JsonParse(error.to_string())), + } + } + }); + Ok(ResponseHandlerOutput { + value, + raw_value: None, + response_headers: Some(output_headers), + }) + }) + .streaming() +} + +/// Build an error using only status, headers and body. +#[must_use] +pub fn create_status_code_error_response_handler() -> ResponseHandler { + ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let message = input + .response + .status() + .canonical_reason() + .unwrap_or_default() + .to_string(); + let headers = extract_response_headers(input.response.headers()); + // Same best-effort contract as the JSON error handler: an oversize + // or dying body must not replace the provider's actual error text + // with a size-limit failure. + let (body, truncated) = crate::read_response_with_size_limit::read_error_body_truncated( + input.response, + ERROR_BODY_PUBLIC_CAP, + input.abort_signal.as_ref(), + ) + .await?; + let error = AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + response_body: Some(error_body_string(&body, truncated)), + response_headers: Some(headers.clone()), + is_retryable: aimux_core::error::is_retryable_status(status), + ..ApiCallError::new(message, input.url, input.request_body_values) + })); + Ok(ResponseHandlerOutput { + value: error, + raw_value: None, + response_headers: Some(headers), + }) + }) +} diff --git a/aimux-provider-utils/src/retry.rs b/aimux-provider-utils/src/retry.rs index 41a7cf12..aba3b982 100644 --- a/aimux-provider-utils/src/retry.rs +++ b/aimux-provider-utils/src/retry.rs @@ -1,324 +1,3 @@ -//! Retry with exponential backoff. -//! -//! Two flavours are provided: -//! - [`retry_with_exponential_backoff`] — a plain exponential-backoff retry -//! that retries on `AiMuxError::is_retryable()` errors. -//! - [`retry_with_exponential_backoff_respecting_retry_headers`] — the same, -//! but consults a `retry-after` hint carried by the error (e.g. from a 429 -//! `retry-after-ms` / `retry-after` response header) and uses it in -//! preference to the exponential delay when the hint is reasonable. -//! -//! The header-parsing and delay-selection logic is exposed as the pure helpers -//! [`parse_retry_after`] and [`get_retry_delay_ms`] so it can be unit-tested -//! independently of the async retry loop. These mirror the TS SDK's -//! `getRetryDelayInMs` / `retryWithExponentialBackoffRespectingRetryHeaders`. +//! Compatibility exports for retry configuration. -use std::time::{Duration, SystemTime}; - -use aimux_core::AiMuxError; -use rand::Rng; - -/// Retry a fallible async operation with exponential backoff. -/// -/// - `max_retries`: maximum number of retry attempts (0 = no retry). -/// - `initial_delay`: initial delay between retries (doubled each time). -/// - Only retries on `AiMuxError::is_retryable()` errors. -/// -/// # Errors -/// -/// Returns the operation's last error once retries are exhausted (or -/// immediately for non-retryable errors). -pub async fn retry_with_exponential_backoff( - max_retries: u32, - initial_delay: Duration, - mut f: F, -) -> Result -where - F: FnMut() - -> std::pin::Pin> + Send>>, - T: Send, -{ - let mut last_error = AiMuxError::Other("no attempts made".to_string()); - let mut delay = initial_delay; - - for attempt in 0..=max_retries { - match f().await { - Ok(result) => return Ok(result), - Err(e) => { - last_error = e; - if !last_error.is_retryable() || attempt == max_retries { - return Err(last_error); - } - tokio::time::sleep(delay).await; - delay = delay.saturating_mul(2); - } - } - } - - Err(last_error) -} - -/// Configuration for [`retry_with_exponential_backoff_respecting_retry_headers`]. -/// -/// Mirrors the TS options `{ maxRetries, initialDelayInMs, backoffFactor }`. -#[derive(Debug, Clone, Copy)] -pub struct RetryConfig { - /// Maximum number of retry attempts (0 = no retry). Default 2. - pub max_retries: u32, - /// Initial delay between retries. Default 2000ms. - pub initial_delay: Duration, - /// Backoff factor applied to the delay after each retry. Default 2. - pub backoff_factor: u32, -} - -impl Default for RetryConfig { - fn default() -> Self { - Self { - max_retries: 2, - initial_delay: Duration::from_millis(2000), - backoff_factor: 2, - } - } -} - -/// Retry a fallible async operation with exponential backoff, respecting -/// `retry-after` hints carried by the error when they are reasonable -/// (0 <= delay < 60s, or shorter than the exponential backoff would be). -/// -/// Only retries on `AiMuxError::is_retryable()` errors. The delay for each -/// retry is chosen by [`get_retry_delay_ms`], fed from -/// [`AiMuxError::retry_after_hint`]. -/// -/// # Errors -/// -/// Returns the operation's last error once retries are exhausted (or -/// immediately for non-retryable errors). -pub async fn retry_with_exponential_backoff_respecting_retry_headers( - config: RetryConfig, - mut f: F, -) -> Result -where - F: FnMut() - -> std::pin::Pin> + Send>>, - T: Send, -{ - let mut last_error = AiMuxError::Other("no attempts made".to_string()); - let mut exponential_delay_ms = config.initial_delay.as_millis() as i64; - - for attempt in 0..=config.max_retries { - match f().await { - Ok(result) => return Ok(result), - Err(e) => { - last_error = e; - if !last_error.is_retryable() || attempt == config.max_retries { - return Err(last_error); - } - - let hint = last_error.retry_after_hint(); - let delay_ms = { - let mut rng = rand::thread_rng(); - get_retry_delay_ms_with_jitter(hint, exponential_delay_ms, &mut rng) - }; - tokio::time::sleep(Duration::from_millis(delay_ms.max(0) as u64)).await; - - exponential_delay_ms = - exponential_delay_ms.saturating_mul(config.backoff_factor as i64); - } - } - } - - Err(last_error) -} - -/// Choose the retry delay (in ms) given an optional `retry-after` hint and the -/// current exponential-backoff delay. -/// -/// Mirrors the TS `getRetryDelayInMs` "reasonable delay" check: the hint is -/// used when it is present, non-negative, and either shorter than 60 seconds -/// or shorter than the exponential backoff would be. Otherwise the exponential -/// backoff delay is used. -#[must_use] -pub fn get_retry_delay_ms(hint: Option, exponential_delay_ms: i64) -> i64 { - match hint { - // Use the hint when it is non-negative AND (shorter than 60s OR shorter - // than the exponential backoff would be). Otherwise fall back. - Some(ms) if ms >= 0 && (ms < 60_000 || ms < exponential_delay_ms) => ms, - _ => exponential_delay_ms, - } -} - -/// 在 [`get_retry_delay_ms`] 基础上叠加 Full Jitter（参考 catcher -/// `DecorrelatedJitter`，即 AWS Full Jitter）。 -/// -/// `delay = random(0, base)`，其中 `base` 仍优先采用 `retry-after` hint， -/// 回退指数退避。防并发 429 惊群，且不丢 retry-after 语义（RFC-0009 §4.2）。 -/// -/// `base <= 0` 时返回 0（`gen_range(0..0)` 会 panic，此处提前保护）。 -pub fn get_retry_delay_ms_with_jitter( - hint: Option, - exponential_delay_ms: i64, - rng: &mut impl Rng, -) -> i64 { - let base = get_retry_delay_ms(hint, exponential_delay_ms); - if base <= 0 { - return 0; - } - rng.gen_range(0..base) -} - -/// Parse a `retry-after` hint (in milliseconds) from response header values. -/// -/// `retry_after_ms_header` is the more precise `retry-after-ms` header (used by -/// e.g. OpenAI); `retry_after_header` is the standard `retry-after` header, -/// which may be either a number of seconds or an HTTP-date. `retry-after-ms` -/// takes precedence when both are present and parseable. -/// -/// `now` is the reference instant used to compute the delay for HTTP-date -/// values. Returns `None` when no header is present or none parse to a usable -/// value. Mirrors the TS `getRetryDelayInMs` header-reading branch. -#[must_use] -pub fn parse_retry_after( - retry_after_ms_header: Option<&str>, - retry_after_header: Option<&str>, - now: SystemTime, -) -> Option { - let mut ms: Option = None; - - // retry-after-ms is more precise than retry-after and used by e.g. OpenAI. - if let Some(raw) = retry_after_ms_header - && let Ok(v) = raw.trim().parse::() - && v.is_finite() - { - ms = Some(v as i64); - } - - // About the Retry-After header: - // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After - if ms.is_none() - && let Some(raw) = retry_after_header - { - let trimmed = raw.trim(); - // First try to parse as a number of seconds. - if let Ok(seconds) = trimmed.parse::() { - if seconds.is_finite() { - ms = Some((seconds * 1000.0) as i64); - } - } else { - // Otherwise try to parse as an HTTP date. - if let Ok(target) = httpdate::parse_http_date(trimmed) { - if let Ok(duration) = target.duration_since(now) { - ms = Some(duration.as_millis() as i64); - } else if let Ok(duration) = now.duration_since(target) { - // Date is in the past → non-positive delay. - ms = Some(-(duration.as_millis() as i64)); - } - } - } - } - - ms -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn get_delay_uses_hint_when_reasonable() { - assert_eq!(get_retry_delay_ms(Some(3000), 2000), 3000); - } - - #[test] - fn get_delay_uses_exponential_when_hint_too_long() { - assert_eq!(get_retry_delay_ms(Some(70_000), 2000), 2000); - } - - #[test] - fn get_delay_uses_exponential_when_hint_negative() { - assert_eq!(get_retry_delay_ms(Some(-1000), 2000), 2000); - } - - #[test] - fn get_delay_uses_exponential_when_no_hint() { - assert_eq!(get_retry_delay_ms(None, 2000), 2000); - } - - #[test] - fn get_delay_uses_hint_when_shorter_than_exponential_even_if_over_60s() { - // 70000ms hint but exponential is even larger → hint wins. - assert_eq!(get_retry_delay_ms(Some(70_000), 80_000), 70_000); - } - - #[test] - fn jitter_returns_zero_when_base_is_zero() { - // base == 0 → must return 0 (gen_range(0..0) would panic). - let mut rng = rand::thread_rng(); - assert_eq!(get_retry_delay_ms_with_jitter(None, 0, &mut rng), 0); - } - - #[test] - fn jitter_returns_zero_when_base_negative() { - let mut rng = rand::thread_rng(); - assert_eq!(get_retry_delay_ms_with_jitter(None, -100, &mut rng), 0); - } - - #[test] - fn jitter_stays_within_full_jitter_bounds() { - // Full Jitter: delay ∈ [0, base). base here = exponential 2000 (no hint). - let mut rng = rand::thread_rng(); - for _ in 0..1000 { - let d = get_retry_delay_ms_with_jitter(None, 2000, &mut rng); - assert!((0..2000).contains(&d), "delay {d} out of [0, 2000)"); - } - } - - #[test] - fn jitter_uses_retry_after_hint_as_upper_bound() { - // hint 3000 → base 3000 → delay ∈ [0, 3000). - let mut rng = rand::thread_rng(); - for _ in 0..1000 { - let d = get_retry_delay_ms_with_jitter(Some(3000), 2000, &mut rng); - assert!((0..3000).contains(&d), "delay {d} out of [0, 3000)"); - } - } - - #[test] - fn parse_retry_after_ms_header() { - assert_eq!( - parse_retry_after(Some("3000"), None, SystemTime::now()), - Some(3000) - ); - } - - #[test] - fn parse_retry_after_seconds_header() { - assert_eq!( - parse_retry_after(None, Some("5"), SystemTime::now()), - Some(5000) - ); - } - - #[test] - fn parse_retry_after_prefers_ms_over_seconds() { - assert_eq!( - parse_retry_after(Some("3000"), Some("10"), SystemTime::now()), - Some(3000) - ); - } - - #[test] - fn parse_retry_after_invalid_falls_back_to_none() { - assert_eq!( - parse_retry_after(Some("invalid"), Some("not-a-number"), SystemTime::now()), - None - ); - } - - #[test] - fn parse_retry_after_negative_ms() { - assert_eq!( - parse_retry_after(Some("-1000"), None, SystemTime::now()), - Some(-1000) - ); - } -} +pub use aimux_core::retry::RetryConfig; diff --git a/aimux-provider-utils/src/ws.rs b/aimux-provider-utils/src/ws.rs index 36fe0ab8..c6891f22 100644 --- a/aimux-provider-utils/src/ws.rs +++ b/aimux-provider-utils/src/ws.rs @@ -8,9 +8,9 @@ //! Design notes (RFC-0028 §3.1): //! - **Every await point is abort/timeout covered** — `connect`, `send`, and //! event receives all `select!` against the abort token and the timeout -//! timers. This is the WS analogue of the HTTP layer's `send_timed` / -//! `TimeoutBodyStream` pattern (RFC-0016 R1–R4 precedent: a select on the -//! loop alone does not cover the send path). +//! timers. This is the WS analogue of the HTTP API-call primitive and +//! Core's semantic stream timeout (RFC-0016 R1–R4 precedent: a select on +//! the loop alone does not cover the send path). //! - **Backpressure is socket-level**: tungstenite's `send().await` drives //! flush and pends while the socket write buffer is full. //! - **No proxy support**: tokio-tungstenite has no proxy parameter; WS @@ -24,9 +24,9 @@ use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; +use aimux_core::AbortSignal; use aimux_core::error::{AiMuxError, ApiCallError}; use aimux_core::options::TimeoutConfiguration; -use aimux_core::shared::AbortSignal; /// A request to open a WebSocket connection. pub struct WebSocketRequest { @@ -55,6 +55,7 @@ pub enum WsMessage { /// A connected WebSocket with abort/timeout enforcement built in. pub struct WsConnection { stream: WebSocketStream>, + url: String, abort: Option, /// Deadline for the first event (connect + session ack). Cleared after /// the first event arrives. @@ -74,11 +75,22 @@ async fn abort_future(abort: &Option) { } } -fn ws_error(msg: impl std::fmt::Display) -> AiMuxError { - AiMuxError::ApiCall(ApiCallError { - message: msg.to_string(), - ..Default::default() - }) +fn abort_error(abort: &Option) -> AiMuxError { + abort + .as_ref() + .map(AiMuxError::from_abort_signal) + .unwrap_or_else(|| AiMuxError::Aborted("request aborted".into())) +} + +fn ws_error(url: &str, msg: impl std::fmt::Display) -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + is_retryable: true, + ..ApiCallError::new( + msg.to_string(), + crate::http::sanitized_request_url(url), + serde_json::json!({}), + ) + })) } enum ConnectError { @@ -86,6 +98,9 @@ enum ConnectError { Tungstenite(tokio_tungstenite::tungstenite::Error), } +// Clippy 1.98 flags the 136-byte tungstenite error variant; this internal +// error is consumed immediately by the one caller, so boxing buys nothing. +#[allow(clippy::result_large_err)] async fn connect_with_timeout( request: tokio_tungstenite::tungstenite::http::Request<()>, timeout: Option, @@ -152,21 +167,43 @@ pub async fn ws_connect(req: &WebSocketRequest) -> Result { - return Err(AiMuxError::Aborted); + return Err(abort_error(&req.abort_signal)); } res = connect_with_timeout(http_req, connect_deadline.map(|d| d - tokio::time::Instant::now())) => match res { Ok(v) => v, Err(ConnectError::Timeout) => { return Err(AiMuxError::Timeout("websocket connect timed out".into())); } + // An HTTP handshake rejection carries a real status: keep it and + // classify retryability by the shared rule instead of the blanket + // transport `is_retryable = true` (401/403 must not be retried). + Err(ConnectError::Tungstenite(tokio_tungstenite::tungstenite::Error::Http( + response, + ))) => { + let status = response.status().as_u16(); + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + is_retryable: aimux_core::error::is_retryable_status(status), + response_body: response + .body() + .as_ref() + .map(|body| String::from_utf8_lossy(body).into_owned()), + ..ApiCallError::new( + format!("websocket connect rejected: HTTP {status}"), + crate::http::sanitized_request_url(&req.url), + serde_json::json!({}), + ) + }))); + } Err(ConnectError::Tungstenite(e)) => { - return Err(ws_error(format!("websocket connect failed: {e}"))); + return Err(ws_error(&req.url, format!("websocket connect failed: {e}"))); } }, }; Ok(WsConnection { stream, + url: req.url.clone(), abort: req.abort_signal.clone(), // The REMAINING first-chunk budget (anchored before connect) — not a // fresh full window, so connect+ack can never exceed first_chunk_ms. @@ -196,12 +233,12 @@ impl WsConnection { pub async fn send_text(&mut self, text: &str) -> Result<(), AiMuxError> { tokio::select! { biased; - _ = abort_future(&self.abort) => Err(AiMuxError::Aborted), - _ = total_deadline_future(&self.total_deadline), if self.total_deadline.is_some() => { + _ = abort_future(&self.abort) => Err(abort_error(&self.abort)), + _ = deadline_future(self.total_deadline), if self.total_deadline.is_some() => { Err(AiMuxError::Timeout("websocket send exceeded total timeout".into())) } res = self.stream.send(Message::Text(text.to_string())) => { - res.map_err(|e| ws_error(format!("websocket send failed: {e}"))) + res.map_err(|e| ws_error(&self.url, format!("websocket send failed: {e}"))) } } } @@ -216,12 +253,12 @@ impl WsConnection { pub async fn send_binary(&mut self, bytes: &[u8]) -> Result<(), AiMuxError> { tokio::select! { biased; - _ = abort_future(&self.abort) => Err(AiMuxError::Aborted), - _ = total_deadline_future(&self.total_deadline), if self.total_deadline.is_some() => { + _ = abort_future(&self.abort) => Err(abort_error(&self.abort)), + _ = deadline_future(self.total_deadline), if self.total_deadline.is_some() => { Err(AiMuxError::Timeout("websocket send exceeded total timeout".into())) } res = self.stream.send(Message::Binary(bytes.to_vec())) => { - res.map_err(|e| ws_error(format!("websocket send failed: {e}"))) + res.map_err(|e| ws_error(&self.url, format!("websocket send failed: {e}"))) } } } @@ -240,7 +277,7 @@ impl WsConnection { loop { tokio::select! { biased; - _ = abort_future(&self.abort) => return Some(Err(AiMuxError::Aborted)), + _ = abort_future(&self.abort) => return Some(Err(abort_error(&self.abort))), _ = deadline_future(self.first_chunk_deadline), if self.first_chunk_deadline.is_some() => { return Some(Err(AiMuxError::Timeout("timed out waiting for first websocket event".into()))); } @@ -271,12 +308,12 @@ impl WsConnection { let detail = frame .map(|f| format!(" (code {}: {})", u16::from(f.code), f.reason)) .unwrap_or_default(); - return Some(Err(ws_error(format!( + return Some(Err(ws_error(&self.url, format!( "websocket closed by peer{detail}" )))); } Some(Err(e)) => { - return Some(Err(ws_error(format!("websocket error: {e}")))); + return Some(Err(ws_error(&self.url, format!("websocket error: {e}")))); } }, } @@ -324,10 +361,3 @@ async fn deadline_future(deadline: Option) { None => pending().await, } } - -async fn total_deadline_future(deadline: &Option) { - match deadline { - Some(d) => tokio::time::sleep_until(*d).await, - None => pending().await, - } -} diff --git a/aimux-provider-utils/tests/api_helpers_test.rs b/aimux-provider-utils/tests/api_helpers_test.rs new file mode 100644 index 00000000..7d272869 --- /dev/null +++ b/aimux-provider-utils/tests/api_helpers_test.rs @@ -0,0 +1,330 @@ +//! Integration coverage for the single-exchange API helpers. + +use futures::StreamExt; +use serde::Deserialize; +use serde_json::json; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use aimux_core::{AbortSignal, AiMuxError}; +use aimux_provider_utils::{ + HttpBody, HttpRequest, ProviderErrorParts, create_binary_response_handler, + create_event_source_response_handler, create_json_error_response_handler, + create_json_response_handler, get_from_api, post_json_to_api, post_to_api, +}; + +#[derive(Debug, Deserialize, PartialEq, Eq)] +struct Reply { + value: String, +} + +fn request(url: String) -> HttpRequest { + HttpRequest { + url, + headers: vec![("authorization".into(), "Bearer test".into())], + abort_signal: None, + call_id: None, + recording_context: None, + } +} + +fn failed_response_handler() -> aimux_provider_utils::ResponseHandler { + create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + ProviderErrorParts { + message: error + .get("message") + .and_then(serde_json::Value::as_str) + .unwrap_or("request failed") + .to_owned(), + provider_code: error + .get("code") + .and_then(serde_json::Value::as_str) + .map(str::to_owned), + } + }) +} + +#[tokio::test] +async fn post_json_dispatches_typed_success_handler() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/call")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("x-test", "yes") + .set_body_json(json!({"value": "ok"})), + ) + .expect(1) + .mount(&server) + .await; + + let output = post_json_to_api( + request(format!("{}/call", server.uri())), + json!({"prompt": "hello"}), + create_json_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap(); + + assert_eq!(output.value, Reply { value: "ok".into() }); + assert_eq!(output.raw_value, Some(json!({"value": "ok"}))); + assert_eq!( + output.response_headers.unwrap().get("x-test"), + Some(&"yes".into()) + ); +} + +#[tokio::test] +async fn failed_handler_preserves_provider_context_and_redacts_secrets() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/call")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("set-cookie", "session=secret") + .insert_header("retry-after", "1") + .set_body_json(json!({ + "error": { + "message": "slow down", + "code": "rate_limit", + "token": "must-not-leak" + } + })), + ) + .expect(1) + .mount(&server) + .await; + + let error = post_json_to_api( + request(format!("{}/call?api_key=secret", server.uri())), + json!({"api_key": "secret", "prompt": "hello"}), + create_json_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap_err(); + + let AiMuxError::ApiCall(error) = error else { + panic!("expected ApiCallError") + }; + assert_eq!(error.status_code, Some(429)); + assert_eq!(error.message, "slow down"); + assert_eq!(error.provider_code.as_deref(), Some("rate_limit")); + assert!(error.is_retryable); + assert!(!error.url.contains("api_key")); + assert_eq!(error.request_body_values["api_key"], "[REDACTED]"); + assert_eq!( + error.response_headers.as_ref().unwrap().get("set-cookie"), + Some(&"[REDACTED]".into()) + ); + assert_eq!(error.data.as_ref().unwrap()["error"]["token"], "[REDACTED]"); +} + +#[tokio::test] +async fn invalid_success_json_is_contextual_api_call_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_string("not json")) + .expect(1) + .mount(&server) + .await; + + let error = post_json_to_api( + request(server.uri()), + json!({"prompt": "hello"}), + create_json_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap_err(); + + assert!( + matches!(error, AiMuxError::ApiCall(ref e) if e.status_code == Some(200) + && e.response_body.as_deref() == Some("not json") + && e.message.starts_with("Invalid JSON response:")) + ); +} + +#[tokio::test] +async fn binary_handler_accepts_an_empty_success_body() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_bytes(Vec::::new())) + .expect(1) + .mount(&server) + .await; + + let output = post_to_api( + request(server.uri()), + HttpBody::Empty, + create_binary_response_handler(), + failed_response_handler(), + ) + .await + .unwrap(); + + assert!(output.value.is_empty()); +} + +#[tokio::test] +async fn event_source_handler_deserializes_each_data_event() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_string( + "data: {\"value\":\"one\"}\n\ndata: {\"value\":\"two\"}\n\ndata: [DONE]\n\n", + )) + .expect(1) + .mount(&server) + .await; + + let output = post_json_to_api( + request(server.uri()), + json!({}), + create_event_source_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap(); + let values = output + .value + .collect::>() + .await + .into_iter() + .collect::, _>>() + .unwrap(); + + assert_eq!( + values, + vec![ + Reply { + value: "one".into() + }, + Reply { + value: "two".into() + } + ] + ); +} + +#[tokio::test] +async fn raw_post_and_get_share_the_same_dispatch_contract() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/binary")) + .respond_with(ResponseTemplate::new(200).set_body_bytes([1, 2, 3])) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/json")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "get"}))) + .expect(1) + .mount(&server) + .await; + + let binary = post_to_api( + request(format!("{}/binary", server.uri())), + HttpBody::Bytes(vec![9], "application/octet-stream".into()), + create_binary_response_handler(), + failed_response_handler(), + ) + .await + .unwrap(); + assert_eq!(binary.value.as_ref(), &[1, 2, 3]); + + let json = get_from_api( + request(format!("{}/json", server.uri())), + create_json_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap(); + assert_eq!(json.value.value, "get"); +} + +#[tokio::test] +async fn caller_abort_is_not_reclassified_as_transport_failure() { + let signal = AbortSignal::new(); + signal.abort(); + let mut request = request("http://127.0.0.1:1/unreachable".into()); + request.abort_signal = Some(signal); + + let error = post_json_to_api( + request, + json!({}), + create_json_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap_err(); + + assert!(matches!(error, AiMuxError::Aborted(_))); +} + +#[tokio::test] +async fn provider_utils_performs_exactly_one_exchange() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({ + "error": {"message": "retry me", "code": "unavailable"} + }))) + .expect(1) + .mount(&server) + .await; + + let error = post_json_to_api( + request(server.uri()), + json!({}), + create_json_response_handler::(), + failed_response_handler(), + ) + .await + .unwrap_err(); + assert!(matches!(error, AiMuxError::ApiCall(ref e) if e.is_retryable)); +} + +#[tokio::test] +async fn oversize_error_body_is_truncated_not_replaced() { + let server = MockServer::start().await; + // >64 KiB error body: the provider's bytes must survive (truncated), + // not be replaced by a size-limit error. + let huge = format!( + "{{\"error\":{{\"message\":\"real provider failure\",\"detail\":\"{}\"}}}}", + "x".repeat(70 * 1024) + ); + Mock::given(method("POST")) + .and(path("/call")) + .respond_with( + ResponseTemplate::new(400) + .insert_header("content-type", "application/json") + .set_body_string(huge), + ) + .expect(1) + .mount(&server) + .await; + + let error = post_json_to_api::( + request(format!("{}/call", server.uri())), + json!({"input": 1}), + create_json_response_handler(), + failed_response_handler(), + ) + .await + .unwrap_err(); + + let AiMuxError::ApiCall(detail) = error else { + panic!("expected ApiCall"); + }; + assert_eq!(detail.status_code, Some(400)); + assert!( + !detail.message.contains("exceeded maximum size"), + "size-limit error must not replace the provider error: {}", + detail.message + ); + let body = detail.response_body.as_deref().expect("body kept"); + assert!(body.starts_with("{\"error\""), "provider bytes kept"); + assert!(body.ends_with("…(truncated)"), "truncation marked"); + assert!(body.len() < 70 * 1024, "body actually truncated"); +} diff --git a/aimux-provider-utils/tests/http_recording_test.rs b/aimux-provider-utils/tests/http_recording_test.rs index 145433ce..07460e58 100644 --- a/aimux-provider-utils/tests/http_recording_test.rs +++ b/aimux-provider-utils/tests/http_recording_test.rs @@ -1,540 +1,282 @@ -//! RFC-0023 层 B(http 咽喉点)录制集成测试。 -//! -//! 覆盖:send 成功/失败 exchange、send_stream 骨架 + 终结补全、脱敏、 -//! transport_closed → barrier 写出。全局录制器单例,用例 `#[serial]` 串行。 +//! Recording coverage for the single-exchange HTTP throat. -use futures::StreamExt; -use serial_test::serial; use std::sync::Arc; use std::time::Duration; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::recording::{ self, JsonlRecorder, OutcomeRecord, OutcomeStatus, Recorder, Recording, }; use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, RetryConfig, send, send_stream, + HttpRequest, ProviderErrorParts, create_json_error_response_handler, + create_json_response_handler, post_json_to_api, }; +use serde::Deserialize; +use serde_json::json; +use serial_test::serial; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; -fn fast_config() -> RetryConfig { - RetryConfig { - max_retries: 1, - initial_delay: Duration::from_millis(1), - backoff_factor: 2, - } +#[derive(Deserialize)] +struct Reply { + #[allow(dead_code)] + ok: bool, +} + +fn failed() -> aimux_provider_utils::ResponseHandler { + create_json_error_response_handler(|_| ProviderErrorParts { + message: "request failed".into(), + provider_code: None, + }) +} + +fn recorder(tag: &str) -> (Arc, std::path::PathBuf) { + let directory = std::env::temp_dir().join(format!("aimux-rfc31-{tag}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&directory); + (Arc::new(JsonlRecorder::new(&directory)), directory) } -fn json_post(url: &str, call_id: Option<&str>) -> HttpRequest { - // 模拟层 A 已注册全局 recorder:测试里 init_recording(Some(rec_arc)) 后, - // 录制上下文由全局 recorder 现取(模拟 layerA)。call_id 与之一致。 - let recording_context = call_id.and_then(|_| { - aimux_core::recording::recorder().map(|recorder| aimux_core::recording::RecordingContext { - call_id: call_id.unwrap().to_string(), - recorder, - }) - }); +fn request_with_signal( + url: String, + call_id: &str, + recorder: Arc, + abort_signal: Option, +) -> HttpRequest { HttpRequest { - method: HttpMethod::Post, - url: url.to_string(), + url, headers: vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ("Authorization".to_string(), "Bearer secret".to_string()), - ("X-Custom".to_string(), "value".to_string()), + ("authorization".into(), "Bearer secret".into()), + ("x-custom".into(), "visible".into()), ], - body: HttpBody::Json(serde_json::json!({"q": "hi"})), - abort_signal: None, - call_id: call_id.map(std::string::ToString::to_string), - recording_context, + abort_signal, + call_id: Some(call_id.into()), + recording_context: Some(aimux_core::recording::RecordingContext::new( + call_id, recorder, + )), } } -fn recorder(tag: &str) -> (JsonlRecorder, std::path::PathBuf) { - let dir = std::env::temp_dir().join(format!("aimux-brec-{tag}-{}", std::process::id())); - let _ = std::fs::remove_dir_all(&dir); - let rec = JsonlRecorder::new(&dir); - (rec, dir) +fn request(url: String, call_id: &str, recorder: Arc) -> HttpRequest { + request_with_signal(url, call_id, recorder, None) } -fn wait_line(dir: &std::path::Path) -> Recording { - let path = dir.join("recordings.jsonl"); - let deadline = std::time::Instant::now() + Duration::from_secs(5); - loop { - if let Ok(s) = std::fs::read_to_string(&path) - && let Some(l) = s.lines().next() - && !l.trim().is_empty() - && let Ok(r) = serde_json::from_str::(l.trim()) - { - return r; - } - if std::time::Instant::now() > deadline { - panic!("no recording line at {}", path.display()); - } - std::thread::sleep(Duration::from_millis(20)); - } -} - -/// 模拟层 A 的完整生命周期(input + outcome + transport closed,barrier 所需)。 -/// 生产中层 A 总是先 `record_input`(RFC-0023 §3.2)——input 是 completion -/// barrier 的必要条件,缺失时记录以 incomplete 落盘。 -fn mimic_layer_a(rec: &dyn Recorder, call_id: &str) { +fn finish_call(recorder: &dyn Recorder, call_id: &str) { use aimux_core::content::ContentPart; use aimux_core::language_model_message::LanguageModelPromptMessage; use aimux_core::message::Role; use aimux_core::options::CallOptions; - let options = CallOptions::new(vec![LanguageModelPromptMessage { - role: Role::User, - content: vec![ContentPart::text("hi")], - ..Default::default() - }]); - rec.record_input(call_id, &options, "openai", "gpt-4o"); - rec.record_outcome( + recorder.record_input( + call_id, + &CallOptions::new(vec![LanguageModelPromptMessage { + role: Role::User, + content: vec![ContentPart::text("hi")], + ..Default::default() + }]), + "openai", + "gpt-test", + ); + recorder.record_outcome( call_id, &OutcomeRecord { status: OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, ); - rec.record_transport_closed(call_id); + recorder.record_transport_closed(call_id); + recorder.flush(); } -#[tokio::test] -#[serial] -async fn send_records_success_exchange_and_closing() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .mount(&server) - .await; - - let (recorder, dir) = recorder("send-ok"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - let url = format!("{}/v1/chat", server.uri()); - let req = json_post(&url, Some("call-b1")); - let resp = send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - assert_eq!(resp.status, 200); - - mimic_layer_a(rec_arc.as_ref(), "call-b1"); - rec_arc.flush(); - let rec = wait_line(&dir); - - assert_eq!(rec.call_id, "call-b1"); - assert_eq!(rec.exchanges.len(), 1); - let ex = &rec.exchanges[0]; - assert!(ex.finalized, "non-stream exchange final"); - // Authorization 脱敏;自定义头保留。 - let auth = ex - .request - .headers - .iter() - .find(|(k, _)| k == "Authorization") - .unwrap(); - assert_eq!(auth.1, "[REDACTED]"); - assert!( - ex.request - .headers - .iter() - .any(|(k, v)| k == "X-Custom" && v == "value") - ); - // 响应体。 - let resp_rec = ex.response.as_ref().unwrap(); - assert!(resp_rec.body.as_deref().unwrap().contains("ok")); - assert!(rec.complete); - - aimux_core::recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); -} - -#[tokio::test] -#[serial] -async fn send_records_failure_without_closing() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(500).set_body_string("server boom")) - .mount(&server) - .await; - - let (recorder, dir) = recorder("b-fail"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - let url = format!("{}/v1/chat", server.uri()); - let req = json_post(&url, Some("call-b2")); - let err = send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE).await; - assert!(err.is_err()); - - // 失败无 closed(等下骨架等待),层 A 的 outcome 若有则写出 ▼ - mimic_layer_a(rec_arc.as_ref(), "call-b2"); - rec_arc.flush(); - let rec = wait_line(&dir); - - // A5:5xx 是合法 HTTP 响应,结构化录制 status/headers/body(不仅 error)。 - let ex_fail = rec - .exchanges - .iter() - .find(|e| e.error.is_some()) - .expect("失败应有 error 字段"); - assert!( - ex_fail.error.as_deref().unwrap().contains("HTTP 500"), - "error 字符串保留用于诊断: {:?}", - ex_fail.error - ); - let resp = ex_fail - .response - .as_ref() - .expect("5xx 应有结构化 response (A5)"); - assert_eq!(resp.status, 500, "5xx response 应保留 status"); - assert_eq!( - resp.body.as_deref(), - Some("server boom"), - "5xx response body 应结构化录制" - ); - // 失败也是完整调用(wire 已终结:有结构化 response + error 字段)。 - assert!(rec.complete, "失败调用也应标记 complete"); - aimux_core::recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); +fn read_recording(directory: &std::path::Path) -> Recording { + let path = directory.join("recordings.jsonl"); + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + if let Ok(text) = std::fs::read_to_string(&path) + && let Some(line) = text.lines().next() + && let Ok(recording) = serde_json::from_str(line) + { + return recording; + } + assert!( + std::time::Instant::now() < deadline, + "recording was not flushed" + ); + std::thread::sleep(Duration::from_millis(20)); + } } #[tokio::test] #[serial] -async fn stream_finalizes_skeleton_with_accumulated_body() { - let sse = "data: {\"delta\":\"hel\"}\n\n\ - data: {\"delta\":\"lo\"}\n\n\ - data: [DONE]\n\n"; +async fn one_helper_call_records_one_finalized_exchange() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/v1/stream")) + .and(path("/call")) .respond_with( ResponseTemplate::new(200) - .set_body_string(sse) - .insert_header("content-type", "text/event-stream"), + .insert_header("set-cookie", "session=secret") + .set_body_json(json!({"ok": true})), ) + .expect(1) .mount(&server) .await; - let (recorder, dir) = recorder("b-stream"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - let url = format!("{}/v1/stream", server.uri()); - let req = json_post(&url, Some("call-b3")); - let resp = send_stream(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - // 消费完整流。 - let _: Vec<_> = resp.body.collect().await; - - mimic_layer_a(rec_arc.as_ref(), "call-b3"); // 层 A outcome+closed - rec_arc.flush(); - let rec = wait_line(&dir); - - assert_eq!(rec.exchanges.len(), 1); - let ex = &rec.exchanges[0]; - assert!(ex.finalized); - let r = ex.response.as_ref().unwrap(); - let stream_text = r.body.as_deref().unwrap_or(""); - assert!( - stream_text.contains(r#"data: {"delta":"hel"}"#), - "原始 SSE 应被累积: {r:?}" - ); - assert!(r.stream_chunks.unwrap_or(0) >= 1, "至少一个 chunk 计数"); - assert!(r.stream_chunks.unwrap_or(0) >= 1); - assert!(rec.complete); - aimux_core::recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); -} - -#[tokio::test] -#[serial] -async fn abandoned_stream_still_finalizes() { - let sse = "data: {\"delta\":\"x\"}\n\n"; - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/stream")) - .respond_with(ResponseTemplate::new(200).set_body_string(sse)) - .mount(&server) - .await; - let (recorder, dir) = recorder("b-abandon"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - let url = format!("{}/v1/stream", server.uri()); - let req = json_post(&url, Some("call-b4")); - let resp = send_stream(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - drop(resp.body); // 提前放弃 → ObservedByteStream Drop 兜底标记 "abandoned" - - mimic_layer_a(rec_arc.as_ref(), "call-b4"); - rec_arc.flush(); - let rec = wait_line(&dir); - let ex = &rec.exchanges[0]; - assert!(ex.finalized, "Drop 应兜底 finalized(即使无完整 body)"); - aimux_core::recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); -} - -#[tokio::test] -#[serial] -async fn retry_records_per_attempt_success_after_429() { - let server = MockServer::start().await; - // 首次 429,第二次 200。 - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(429).set_body_string("rate limited")) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .mount(&server) - .await; - - let (recorder, dir) = recorder("b-retry"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - // send_with_retry_raw 里 max_retries=2 会先 429 重试再成功。 - let url = format!("{}/v1/chat", server.uri()); - let req = json_post(&url, Some("call-retry")); - let resp = send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - assert_eq!(resp.status, 200); - - mimic_layer_a(rec_arc.as_ref(), "call-retry"); - rec_arc.flush(); - let rec = wait_line(&dir); - // 至少 2 exchange:attempt 0 失败(429)+ attempt 1 成功。 - assert!( - rec.exchanges.len() >= 2, - "expected retry exchanges, got {}", - rec.exchanges.len() - ); - // 第一条是 429:有 error 且结构化 response(A5:不再丢失 status/headers/body)。 - let ex429 = rec - .exchanges - .iter() - .find(|e| e.error.is_some()) - .expect("attempt 0 should be a 429 failure"); - assert!( - ex429.error.as_deref().unwrap().contains("HTTP 429"), - "429 error 字符串保留: {:?}", - ex429.error - ); - let r429 = ex429 - .response - .as_ref() - .expect("429 应有结构化 response (A5)"); - assert_eq!(r429.status, 429, "429 response 应保留 status"); + let (recorder, directory) = recorder("success"); + recording::init_recording(Some(recorder.clone())); + + post_json_to_api( + request(format!("{}/call", server.uri()), "call-1", recorder.clone()), + json!({"api_key": "secret", "prompt": "hi"}), + create_json_response_handler::(), + failed(), + ) + .await + .unwrap(); + finish_call(recorder.as_ref(), "call-1"); + + let recording = read_recording(&directory); + assert_eq!(recording.exchanges.len(), 1); + let exchange = &recording.exchanges[0]; + assert!(exchange.finalized); + assert_eq!(exchange.attempt, 1); assert_eq!( - r429.body.as_deref(), - Some("rate limited"), - "429 response body 应结构化录制" + exchange + .request + .headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("authorization")) + .unwrap() + .1, + "[REDACTED]" ); - // 最后一条成功,status 200。 + let response = exchange.response.as_ref().unwrap(); + assert_eq!(response.status, 200); + assert!(response.body.as_deref().unwrap().contains("true")); assert_eq!( - rec.exchanges - .last() - .unwrap() - .response - .as_ref() + response + .headers + .iter() + .find(|(name, _)| name.eq_ignore_ascii_case("set-cookie")) .unwrap() - .status, - 200 + .1, + "[REDACTED]" ); recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(directory); } #[tokio::test] #[serial] -async fn utf8_body_truncates_without_panic_at_char_boundary() { - // 构造恰好跨 1MiB 边界的中文(UTF-8 多字节)。 - let big = "中".repeat(400_000); // 400k * 3 bytes = 1.2MiB > 1MiB cap - let server = MockServer::start().await; - let big_for_mock = big.clone(); - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_string(big_for_mock)) - .mount(&server) - .await; - - let (recorder, dir) = recorder("b-utf8"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - let url = format!("{}/v1/chat", server.uri()); - let req = json_post(&url, Some("call-utf8")); - // 不应 panic(UTF-8 安全截断)。 - let big_len = big.len(); - let resp = send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - assert_eq!(resp.body.len(), big_len); // 完整 body 返回给调用方 - - mimic_layer_a(rec_arc.as_ref(), "call-utf8"); - rec_arc.flush(); - let rec = wait_line(&dir); - let rbody = rec.exchanges[0] - .response - .as_ref() - .unwrap() - .body - .as_deref() - .unwrap(); - assert!(rbody.len() <= 1 << 20, "recorded body must respect cap"); - // 必须是完整 UTF-8(无 panic 已证明);中文可安全转回。 - assert!(rbody.chars().count() > 0); +async fn cancelled_exchange_is_recorded_as_finalized_failure() { + let signal = aimux_core::AbortSignal::new(); + let (recorder, directory) = recorder("cancelled"); + recording::init_recording(Some(recorder.clone())); + + let request = request_with_signal( + "http://127.0.0.1:9/call".into(), + "call-cancelled", + recorder.clone(), + Some(signal.clone()), + ); + signal.abort(); + let result = post_json_to_api( + request, + json!({"prompt": "hi"}), + create_json_response_handler::(), + failed(), + ) + .await; + assert!(matches!(result, Err(aimux_core::AiMuxError::Aborted(_)))); + + finish_call(recorder.as_ref(), "call-cancelled"); + let recording = read_recording(&directory); + assert_eq!(recording.exchanges.len(), 1); + let exchange = &recording.exchanges[0]; + assert!(exchange.finalized); + assert_eq!(exchange.attempt, 1); + assert_eq!(exchange.exchange_index, 1); + assert!(exchange.response.is_none()); + assert!(exchange.error.is_some()); recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(directory); } - #[tokio::test] #[serial] -async fn security_token_header_redacted_usage_keys_preserved() { - // header 脱敏收窄:x-amz-security-token(AWS Bedrock sigv4 真实凭据头)恒脱敏; - // contains("token") 已移除,故 max_output_tokens 等用量键不再被误伤。 - // URL query 不再脱敏(provider 均 header 鉴权),原样录制。 +async fn failed_http_response_is_still_one_observed_exchange() { let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({ + "error": {"message": "unavailable"} + }))) + .expect(1) .mount(&server) .await; - - let (recorder, dir) = recorder("b-urlredact"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - // URL 仅携带非敏感 query(原样录制,不再脱敏);header 携带: - // - x-amz-security-token(Bedrock sigv4 凭据)→ 必须脱敏 - // - max_output_tokens(用量字段,含 "token" 子串)→ 必须保留 - // - Authorization → 必须脱敏(contains authorization 回归保护) - let url = format!("{}/v1/chat?model=gpt&version=1", server.uri()); - let recording_context = - aimux_core::recording::recorder().map(|recorder| aimux_core::recording::RecordingContext { - call_id: "call-redact".to_string(), - recorder, - }); - let req = HttpRequest { - method: HttpMethod::Post, - url, - headers: vec![ - ("Content-Type".to_string(), "application/json".to_string()), - ("x-amz-security-token".to_string(), "STS-TOKEN".to_string()), - ("max_output_tokens".to_string(), "4096".to_string()), - ("Authorization".to_string(), "Bearer SECRET".to_string()), - ], - body: HttpBody::Json(serde_json::json!({"q": "hi"})), - abort_signal: None, - call_id: Some("call-redact".to_string()), - recording_context, - }; - let resp = send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - assert_eq!(resp.status, 200); - - mimic_layer_a(rec_arc.as_ref(), "call-redact"); - rec_arc.flush(); - let rec = wait_line(&dir); - let ex = &rec.exchanges[0]; - let recorded_url = &ex.request.url; - - // URL query 原样保留(不再脱敏):非敏感 query 完整落盘。 - assert!( - recorded_url.ends_with("?model=gpt&version=1"), - "url should be recorded verbatim (no query redaction): {recorded_url}" - ); - - // x-amz-security-token(AWS 凭据头)脱敏。 - let sec_token = ex - .request - .headers - .iter() - .find(|(k, _)| k.eq_ignore_ascii_case("x-amz-security-token")) - .expect("x-amz-security-token header recorded"); - assert_eq!(sec_token.1, "[REDACTED]"); - - // max_output_tokens 含 "token" 子串但非凭据——contains("token") 收窄后保留。 - let usage_key = ex - .request - .headers - .iter() - .find(|(k, _)| k == "max_output_tokens") - .expect("max_output_tokens header recorded"); + let (recorder, directory) = recorder("failure"); + recording::init_recording(Some(recorder.clone())); + + let result = post_json_to_api( + request(server.uri(), "call-2", recorder.clone()), + json!({}), + create_json_response_handler::(), + failed(), + ) + .await; + assert!(result.is_err()); + finish_call(recorder.as_ref(), "call-2"); + + let recording = read_recording(&directory); + assert_eq!(recording.exchanges.len(), 1); assert_eq!( - usage_key.1, "4096", - "max_output_tokens must NOT be redacted (contains(token) narrowed): {:?}", - ex.request.headers - ); - - // Authorization 仍按 contains(authorization) 脱敏(回归保护)。 - let auth = ex - .request - .headers - .iter() - .find(|(k, _)| k.eq_ignore_ascii_case("Authorization")) - .expect("Authorization header recorded"); - assert_eq!(auth.1, "[REDACTED]"); - - // catch-all:整条录制序列化后不得出现明文凭据。 - let dump = serde_json::to_string(&rec).unwrap_or_default(); - assert!( - !dump.contains("STS-TOKEN"), - "STS token leaked in recording: {dump}" - ); - assert!( - !dump.contains("SECRET"), - "secret leaked in recording: {dump}" + recording.exchanges[0].response.as_ref().unwrap().status, + 503 ); + assert!(recording.exchanges[0].finalized); recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(directory); } #[tokio::test] #[serial] -async fn stream_error_patches_single_exchange_not_duplicate() { - // SSE 首帧后一个 error → provider 侧 Error(流中途)。 - let sse = r#"data: {"choices":[{"delta":{"content":"a"}}]}\n\ndata: [DONE]\n\n"#; +async fn redirect_chain_is_recorded_as_one_logical_exchange() { + // RFC-0031 §13-9: hops of an automatic redirect chain are transport + // detail; recording sees a single exchange carrying the final response. let server = MockServer::start().await; Mock::given(method("POST")) - .and(path("/v1/stream")) - .respond_with(ResponseTemplate::new(200).set_body_string(sse)) + .and(path("/call")) + .respond_with(ResponseTemplate::new(302).insert_header("location", "/moved")) + .expect(1) .mount(&server) .await; - let (recorder, dir) = recorder("b-serr"); - let rec_arc: Arc = Arc::new(recorder); - recording::init_recording(Some(rec_arc.clone())); - - let url = format!("{}/v1/stream", server.uri()); - let req = json_post(&url, Some("call-serr")); - let resp = send_stream(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - let _: Vec<_> = resp.body.collect().await; - - mimic_layer_a(rec_arc.as_ref(), "call-serr"); - rec_arc.flush(); - let rec = wait_line(&dir); - // 只有 1 条 exchange(骨架 patch,不新增)。status 应保留 200。 - assert_eq!(rec.exchanges.len(), 1, "must not duplicate attempt"); - assert_eq!(rec.exchanges[0].response.as_ref().unwrap().status, 200); - assert!(rec.exchanges[0].finalized); + Mock::given(method("GET")) + .and(path("/moved")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) + .expect(1) + .mount(&server) + .await; + let (recorder, directory) = recorder("redirect"); + recording::init_recording(Some(recorder.clone())); + + post_json_to_api( + request(format!("{}/call", server.uri()), "call-3", recorder.clone()), + json!({"prompt": "hi"}), + create_json_response_handler::(), + failed(), + ) + .await + .unwrap(); + finish_call(recorder.as_ref(), "call-3"); + + let recording = read_recording(&directory); + assert_eq!(recording.exchanges.len(), 1); + let exchange = &recording.exchanges[0]; + assert!(exchange.finalized); + assert_eq!(exchange.attempt, 1); + assert_eq!(exchange.exchange_index, 1); + assert_eq!(exchange.response.as_ref().unwrap().status, 200); recording::init_recording(None); - let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(directory); } diff --git a/aimux-provider-utils/tests/http_send_test.rs b/aimux-provider-utils/tests/http_send_test.rs deleted file mode 100644 index b80ed458..00000000 --- a/aimux-provider-utils/tests/http_send_test.rs +++ /dev/null @@ -1,615 +0,0 @@ -//! Integration tests for `http::send` / `http::send_stream` — verifies that -//! the HTTP layer's retry + status handling is correctly wired (RFC-0009). -//! -//! Uses wiremock to spin up a local mock server. Providers never touch reqwest -//! types — they hand `http::send` a pure-data `HttpRequest` and get back an -//! `HttpResponse` (or `HttpStreamResponse`). - -use std::time::{Duration, Instant}; - -use futures::StreamExt; -use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; - -use aimux_core::AiMuxError; -use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, RetryConfig, send, send_stream, -}; - -/// Short config for tests: 1ms initial delay (jitter → 0ms effective), -/// up to 2 retries. -fn fast_config() -> RetryConfig { - RetryConfig { - max_retries: 2, - initial_delay: Duration::from_millis(1), - backoff_factor: 2, - } -} - -/// Build a JSON POST request to `url`. -fn json_post(url: &str) -> HttpRequest { - HttpRequest { - method: HttpMethod::Post, - url: url.to_string(), - headers: vec![("Content-Type".to_string(), "application/json".to_string())], - body: HttpBody::Json(serde_json::json!({"q": "hi"})), - abort_signal: None, - call_id: None, - recording_context: None, - } -} - -#[tokio::test] -async fn returns_response_on_success() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ - "ok": true - }))) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - - assert_eq!(resp.status, 200); - let body: serde_json::Value = serde_json::from_slice(&resp.body).unwrap(); - assert_eq!(body["ok"], true); -} - -#[tokio::test] -async fn retries_on_429_then_succeeds() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(429).insert_header("retry-after-ms", "1")) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - - assert_eq!(resp.status, 200); -} - -#[tokio::test] -async fn retries_on_500_then_succeeds() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(500).set_body_string("Internal Server Error")) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - - assert_eq!(resp.status, 200); -} - -/// §9.2: 408/409 now enter the shared retry loop (AI SDK response policy); -/// one integration case is enough — status-policy unit coverage in -/// structured_error_fields handles 409. -#[tokio::test] -async fn retries_on_408_then_succeeds() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(408).set_body_string("Request Timeout")) - .up_to_n_times(1) - .mount(&server) - .await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - - assert_eq!(resp.status, 200); -} - -#[tokio::test] -async fn does_not_retry_on_401() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(401).set_body_json(serde_json::json!({ - "error": { "message": "Invalid API key", "type": "invalid_request_error" } - }))) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - - assert!( - matches!(err, AiMuxError::ApiCall(ref d) if d.status_code == Some(401)), - "got {err:?}" - ); -} - -#[tokio::test] -async fn does_not_retry_on_404() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(404).set_body_json(serde_json::json!({ - "error": { "message": "model not found", "type": "not_found" } - }))) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - - assert!( - matches!(err, AiMuxError::ApiCall(ref d) if d.status_code == Some(404)), - "got {err:?}" - ); -} - -#[tokio::test] -async fn exhausts_retries_on_persistent_429() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(429).insert_header("retry-after-ms", "1")) - .expect(3) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - - assert!( - matches!(err, AiMuxError::ApiCall(ref d) if d.status_code == Some(429)), - "got {err:?}" - ); -} - -#[tokio::test] -async fn exhausts_retries_on_persistent_500() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(503).set_body_string("Service Unavailable")) - .expect(3) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - - assert!( - matches!(err, AiMuxError::ApiCall(ref d) if d.status_code == Some(503)), - "got {err:?}" - ); -} - -#[tokio::test] -async fn send_stream_returns_byte_stream() { - // send_stream should return a byte stream that the caller can consume - // independently of reqwest types. - let sse_body = "data: {\"choices\":[]}\n\n"; - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("Content-Type", "text/event-stream") - .set_body_string(sse_body), - ) - .expect(1) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = send_stream(json_post(&url), fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap(); - - assert_eq!(resp.status, 200); - - // Consume the byte stream — it yields Bytes, not reqwest types. - let mut collected = Vec::new(); - let mut stream = resp.body; - while let Some(chunk) = stream.next().await { - collected.extend_from_slice(&chunk.unwrap()); - } - assert_eq!(collected, sse_body.as_bytes()); -} - -// ───────────────────────────────────────────────────────────────────────────── -// RFC-0016 H1/H3: abort + per-call timeout -// ───────────────────────────────────────────────────────────────────────────── - -use aimux_core::shared::AbortSignal; -use aimux_provider_utils::{RequestTimeout, send_stream_timed, send_timed}; - -#[tokio::test] -async fn abort_before_send_fails_immediately() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .expect(0) // must never be reached - .mount(&server) - .await; - - let signal = AbortSignal::new(); - signal.abort(); - - let mut req = json_post(&format!("{}/v1/chat", server.uri())); - req.abort_signal = Some(signal); - - let err = send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - assert!(matches!(err, AiMuxError::Aborted), "got {err:?}"); -} - -#[tokio::test] -async fn abort_mid_request_cancels() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"ok": true})) - .set_delay(Duration::from_millis(300)), - ) - .mount(&server) - .await; - - let signal = AbortSignal::new(); - let signal_clone = signal.clone(); - - let url = format!("{}/v1/chat", server.uri()); - let request = json_post(&url); - let handle = tokio::spawn(async move { - let mut req = request; - req.abort_signal = Some(signal_clone); - send(req, fast_config(), &DEFAULT_ERROR_STRUCTURE).await - }); - - tokio::time::sleep(Duration::from_millis(50)).await; - let started = Instant::now(); - signal.abort(); - - // Bounded join: if cancellation regresses to hanging, the outer timeout - // fails the test instead of hanging the suite. - let result = tokio::time::timeout(Duration::from_secs(2), handle) - .await - .expect("cancelled call must finish within 2s") - .expect("task must not panic"); - assert!(matches!(result, Err(AiMuxError::Aborted)), "got {result:?}"); - assert!( - started.elapsed() < Duration::from_millis(250), - "abort should cancel well before the 300ms server delay" - ); -} - -#[tokio::test] -async fn send_timed_total_timeout_fails() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"ok": true})) - .set_delay(Duration::from_millis(500)), - ) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send_timed( - json_post(&url), - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - total_ms: Some(100), - ..Default::default() - }), - ) - .await - .unwrap_err(); - - assert!(matches!(err, AiMuxError::Timeout(_)), "got {err:?}"); - assert!(err.to_string().contains("total timeout"), "got {err:?}"); -} - -#[tokio::test] -async fn send_timed_total_timeout_covers_retries() { - // A retryable 429 whose response takes longer than the total budget: - // the total deadline must bound the whole call (connect + response + - // retry backoff), not just one attempt. - // - // The response is delayed past the budget on purpose: a fast 429 would - // race with the jittered retry backoff (Full Jitter ∈ [0, retry-after)), - // making this test flaky depending on the random delay drawn. - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(429) - .insert_header("retry-after-ms", "200") - .set_delay(Duration::from_millis(500)), - ) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send_timed( - json_post(&url), - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - total_ms: Some(100), - ..Default::default() - }), - ) - .await - .unwrap_err(); - - assert!(matches!(err, AiMuxError::Timeout(_)), "got {err:?}"); -} - -#[tokio::test] -async fn send_timed_within_budget_succeeds() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"ok": true})) - .set_delay(Duration::from_millis(50)), - ) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = send_timed( - json_post(&url), - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - total_ms: Some(5000), - ..Default::default() - }), - ) - .await - .unwrap(); - - assert_eq!(resp.status, 200); -} - -#[tokio::test] -async fn send_stream_timed_first_chunk_timeout() { - // wiremock's `set_delay` delays the whole response (headers included), so - // this test covers "the first-byte budget includes response-header - // latency" — the first-chunk timer counts from request start. The - // body-pending wakeup path itself is covered by the unit tests - // (`timeout_stream_yields_after_first_chunk_deadline` and - // `timeout_stream_enforces_chunk_idle_deadline`), which use a pending - // inner stream. - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("Content-Type", "text/event-stream") - .set_body_string("data: {\"choices\":[]}\n\n") - .set_delay(Duration::from_millis(500)), - ) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let resp = match send_stream_timed( - json_post(&url), - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - first_chunk_ms: Some(100), - ..Default::default() - }), - ) - .await - { - Ok(r) => r, - Err(e) => panic!("unexpected error: {e:?}"), - }; - - let mut stream = resp.body; - let first = stream.next().await.expect("stream must yield an error"); - let err = first.unwrap_err(); - assert!(matches!(err, AiMuxError::Timeout(_)), "got {err:?}"); - assert!( - err.to_string().contains("first chunk timeout"), - "got {err:?}" - ); -} - -#[tokio::test] -async fn send_stream_timed_total_timeout_on_connect() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .set_body_string("data: x\n\n") - .set_delay(Duration::from_millis(500)), - ) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = match send_stream_timed( - json_post(&url), - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - total_ms: Some(100), - ..Default::default() - }), - ) - .await - { - Ok(_) => panic!("expected timeout, got success"), - Err(e) => e, - }; - - assert!(matches!(err, AiMuxError::Timeout(_)), "got {err:?}"); -} - -#[tokio::test] -async fn send_timed_zero_total_fails_immediately() { - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) - .expect(0) // must never be reached - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let err = send_timed( - json_post(&url), - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - total_ms: Some(0), - ..Default::default() - }), - ) - .await - .unwrap_err(); - assert!(matches!(err, AiMuxError::Timeout(_)), "got {err:?}"); -} - -#[tokio::test] -async fn abort_wins_over_total_timeout() { - // Abort fires strictly before the 100ms total deadline, so the call must - // report Aborted — verifying the documented "abort wins" semantics for - // the connect phase (biased select inside send_request). A true - // same-instant tie is not exercised here (that depends on tokio's - // internal poll order); see RFC-0016 §7.6 S5. - let server = MockServer::start().await; - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"ok": true})) - .set_delay(Duration::from_millis(500)), - ) - .mount(&server) - .await; - - let signal = AbortSignal::new(); - let signal_clone = signal.clone(); - let url = format!("{}/v1/chat", server.uri()); - let request = json_post(&url); - let handle = tokio::spawn(async move { - let mut req = request; - req.abort_signal = Some(signal_clone); - send_timed( - req, - fast_config(), - &DEFAULT_ERROR_STRUCTURE, - Some(RequestTimeout { - total_ms: Some(100), - ..Default::default() - }), - ) - .await - }); - - // Abort just before the 100ms deadline so both fire around the same time. - tokio::time::sleep(Duration::from_millis(90)).await; - signal.abort(); - - let result = tokio::time::timeout(Duration::from_secs(2), handle) - .await - .expect("must finish") - .expect("task must not panic"); - assert!(matches!(result, Err(AiMuxError::Aborted)), "got {result:?}"); -} - -/// Regression (audit round 3, B1): a large *valid-UTF-8* error body must be -/// truncated with a marker — a silent truncation would make callers believe -/// the response body was complete. -#[tokio::test] -async fn truncates_large_error_body_with_marker() { - let server = MockServer::start().await; - let big_body = "x".repeat(100 * 1024); - Mock::given(method("POST")) - .and(path("/v1/chat")) - .respond_with(ResponseTemplate::new(429).set_body_string(big_body)) - .up_to_n_times(3) - .mount(&server) - .await; - - let url = format!("{}/v1/chat", server.uri()); - let no_retry = RetryConfig { - max_retries: 0, - ..RetryConfig::default() - }; - let err = send(json_post(&url), no_retry, &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - - match err { - AiMuxError::ApiCall(detail) => { - assert!( - detail.message.contains("(truncated"), - "large error body must be marked as truncated, got len={}", - detail.message.len() - ); - assert!( - detail.message.len() < 70 * 1024, - "truncated message must stay bounded, got {} bytes", - detail.message.len() - ); - } - other => panic!("expected RateLimited, got {other:?}"), - } -} diff --git a/aimux-provider-utils/tests/logging_test.rs b/aimux-provider-utils/tests/logging_test.rs index 931d9747..2b361cb5 100644 --- a/aimux-provider-utils/tests/logging_test.rs +++ b/aimux-provider-utils/tests/logging_test.rs @@ -1,227 +1,104 @@ -//! RFC-0014 logging integration tests: capture formatted output with a -//! `tracing` subscriber and assert the retry/failed chain, redaction, and -//! `AIMUX_LOG_BODY` body logging behave as designed. +//! Logging tests for the single-exchange HTTP throat. use std::sync::{Arc, Mutex}; -use std::time::Duration; -use aimux_provider_utils::http::{HttpBody, HttpMethod, HttpRequest, send, send_stream}; use aimux_provider_utils::logging::CaptureWriter; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::retry::RetryConfig; -use futures::StreamExt; +use aimux_provider_utils::{ + HttpRequest, ProviderErrorParts, create_json_error_response_handler, + create_json_response_handler, post_json_to_api, +}; +use serde::Deserialize; +use serde_json::json; use tracing_subscriber::EnvFilter; use tracing_subscriber::fmt; use wiremock::matchers::method; use wiremock::{Mock, MockServer, ResponseTemplate}; -fn retry_cfg() -> RetryConfig { - RetryConfig { - max_retries: 1, - initial_delay: Duration::from_millis(1), - backoff_factor: 2, - } +#[derive(Deserialize)] +struct Reply { + #[allow(dead_code)] + ok: bool, } fn request(url: String) -> HttpRequest { HttpRequest { - method: HttpMethod::Post, url, - headers: vec![( - "Authorization".to_string(), - "Bearer secret-token".to_string(), - )], - body: HttpBody::Json(serde_json::json!({ - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi"}], - "api_key": "sk-must-not-leak", - })), + headers: vec![("authorization".into(), "Bearer secret-token".into())], abort_signal: None, call_id: None, recording_context: None, } } -/// 429 → retry (WARN) → 429 → failed (ERROR)：断言事件与字段、断言 header/body -/// 在默认级别下不泄露。 -#[tokio::test] -async fn retry_chain_logged_with_redaction() { - let captured = Arc::new(Mutex::new(Vec::new())); - let subscriber = fmt() - .with_ansi(false) - .with_env_filter(EnvFilter::try_new("aimux_provider_utils=warn").unwrap()) - .with_writer(CaptureWriter(captured.clone())) - .finish(); - let _guard = tracing::subscriber::set_default(subscriber); - - let server = MockServer::start().await; - Mock::given(method("POST")) - .respond_with(ResponseTemplate::new(429).set_body_string("rate limited")) - .expect(2) - .mount(&server) - .await; - - let res = send( - request(format!("{}/v1/test", server.uri())), - retry_cfg(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await; - assert!(res.is_err()); - - let output = String::from_utf8(captured.lock().unwrap().clone()).unwrap(); - assert!(output.contains("retry"), "missing retry event:\n{output}"); - assert!(output.contains("failed"), "missing failed event:\n{output}"); - assert!( - output.contains("status=429"), - "missing status field:\n{output}" - ); - // The classification is the logged status; the variant name is just - // "provider". - assert!( - output.contains("reason=api_call"), - "missing reason field:\n{output}" - ); - assert!( - !output.contains("Bearer secret-token"), - "header value leaked:\n{output}" - ); - assert!( - !output.contains("sk-must-not-leak"), - "request body leaked at warn level:\n{output}" - ); +fn failed() -> aimux_provider_utils::ResponseHandler { + create_json_error_response_handler(|_| ProviderErrorParts { + message: "request failed".into(), + provider_code: None, + }) } -/// debug 级别：request/response 摘要事件出现，URL 无 query、header 只记数量。 -#[tokio::test] -async fn debug_level_shows_request_response_summary() { +fn capture(level: &str) -> (Arc>>, tracing::dispatcher::DefaultGuard) { let captured = Arc::new(Mutex::new(Vec::new())); let subscriber = fmt() .with_ansi(false) - .with_env_filter(EnvFilter::try_new("aimux_provider_utils=debug").unwrap()) + .with_env_filter(EnvFilter::try_new(level).unwrap()) .with_writer(CaptureWriter(captured.clone())) .finish(); - let _guard = tracing::subscriber::set_default(subscriber); - - let server = MockServer::start().await; - Mock::given(method("POST")) - .respond_with(ResponseTemplate::new(200).set_body_string(r#"{"ok":true}"#)) - .mount(&server) - .await; - - let mut url = format!("{}/v1/test", server.uri()); - url.push_str("?api_key=should-not-appear"); - let res = send(request(url), retry_cfg(), &DEFAULT_ERROR_STRUCTURE).await; - assert!(res.is_ok()); - - let output = String::from_utf8(captured.lock().unwrap().clone()).unwrap(); - assert!( - output.contains("request"), - "missing request event:\n{output}" - ); - assert!( - output.contains("response"), - "missing response event:\n{output}" - ); - assert!( - output.contains("status=200"), - "missing status=200:\n{output}" - ); - assert!( - !output.contains("api_key=should-not-appear"), - "URL query leaked:\n{output}" - ); - assert!( - output.contains("header_count=1"), - "header_count missing:\n{output}" - ); + let guard = tracing::subscriber::set_default(subscriber); + (captured, guard) } -/// trace + AIMUX_LOG_BODY=1：请求体出现但 api_key 已打码。 #[tokio::test] -async fn body_logging_redacts_secrets() { - // SAFETY: single test touching this env var; serialized implicitly by - // being the only test that reads it. - unsafe { std::env::set_var("AIMUX_LOG_BODY", "1") }; - let captured = Arc::new(Mutex::new(Vec::new())); - let subscriber = fmt() - .with_ansi(false) - .with_env_filter(EnvFilter::try_new("aimux_provider_utils=trace").unwrap()) - .with_writer(CaptureWriter(captured.clone())) - .finish(); - let _guard = tracing::subscriber::set_default(subscriber); - +async fn debug_summary_sanitizes_url_and_never_logs_header_values() { + let (captured, _guard) = capture("aimux_provider_utils=debug"); let server = MockServer::start().await; Mock::given(method("POST")) - .respond_with(ResponseTemplate::new(200).set_body_string(r#"{"ok":true}"#)) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) .mount(&server) .await; - let res = send( - request(format!("{}/v1/test", server.uri())), - retry_cfg(), - &DEFAULT_ERROR_STRUCTURE, + post_json_to_api( + request(format!("{}?api_key=must-not-appear", server.uri())), + json!({"api_key": "must-not-appear"}), + create_json_response_handler::(), + failed(), ) - .await; - assert!(res.is_ok()); + .await + .unwrap(); let output = String::from_utf8(captured.lock().unwrap().clone()).unwrap(); - assert!( - output.contains("request_body"), - "missing request_body trace:\n{output}" - ); - assert!( - output.contains("\"api_key\":\"***\""), - "api_key not redacted:\n{output}" - ); - assert!( - !output.contains("sk-must-not-leak"), - "secret leaked in body log:\n{output}" - ); - unsafe { std::env::remove_var("AIMUX_LOG_BODY") }; + assert!(output.contains("request")); + assert!(output.contains("response")); + assert!(output.contains("status=200")); + assert!(output.contains("header_count=1")); + assert!(!output.contains("must-not-appear")); + assert!(!output.contains("Bearer secret-token")); } -/// 流式路径：stream_connected / stream_first_byte / stream_end 事件。 #[tokio::test] -async fn stream_events_are_emitted() { - let captured = Arc::new(Mutex::new(Vec::new())); - let subscriber = fmt() - .with_ansi(false) - .with_env_filter(EnvFilter::try_new("aimux_provider_utils=debug").unwrap()) - .with_writer(CaptureWriter(captured.clone())) - .finish(); - let _guard = tracing::subscriber::set_default(subscriber); - +async fn trace_body_uses_shared_sensitive_key_policy() { + // SAFETY: this integration-test process owns this environment variable. + unsafe { std::env::set_var("AIMUX_LOG_BODY", "1") }; + let (captured, _guard) = capture("aimux_provider_utils=trace"); let server = MockServer::start().await; Mock::given(method("POST")) - .respond_with(ResponseTemplate::new(200).set_body_string("data: hello\n\n")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"ok": true}))) .mount(&server) .await; - let resp = send_stream( - request(format!("{}/v1/stream", server.uri())), - retry_cfg(), - &DEFAULT_ERROR_STRUCTURE, + post_json_to_api( + request(server.uri()), + json!({"api_key": "must-not-appear", "prompt": "hello"}), + create_json_response_handler::(), + failed(), ) .await - .expect("stream connect failed"); - // 消费完整字节流，触发 first_byte / end 事件。 - let mut body = resp.body; - while let Some(chunk) = body.next().await { - assert!(chunk.is_ok()); - } + .unwrap(); let output = String::from_utf8(captured.lock().unwrap().clone()).unwrap(); - assert!( - output.contains("stream_connected"), - "missing connect event:\n{output}" - ); - assert!( - output.contains("stream_first_byte"), - "missing ttfb event:\n{output}" - ); - assert!( - output.contains("stream_end"), - "missing end event:\n{output}" - ); + assert!(output.contains("request_body")); + assert!(output.contains("[REDACTED]")); + assert!(!output.contains("must-not-appear")); + // SAFETY: paired with the set above. + unsafe { std::env::remove_var("AIMUX_LOG_BODY") }; } diff --git a/aimux-provider-utils/tests/response_handler.rs b/aimux-provider-utils/tests/response_handler.rs deleted file mode 100644 index 5c35d2f8..00000000 --- a/aimux-provider-utils/tests/response_handler.rs +++ /dev/null @@ -1,138 +0,0 @@ -//! Rust translation of `packages/provider-utils/src/response-handler.test.ts`. -//! -//! The TS response-handler tests exercise `createJsonResponseHandler`, -//! `createJsonErrorResponseHandler`, `createBinaryResponseHandler`, and -//! `createStatusCodeErrorResponseHandler` against the Web `Response` object, -//! including oversized-response size limiting and binary handling. -//! -//! The Rust equivalent works one layer down: providers fetch a response, then -//! call [`aimux_provider_utils::parse_provider_error`] with the HTTP status and -//! response body text to produce an [`aimux_core::AiMuxError`]. The cases that -//! depend on the streaming `Response` object (size-limit cancellation, binary -//! array buffers, the `rawValue`/`value` split) have no Rust counterpart and -//! are skipped with comments. The translatable cases — status-code → error -//! mapping, message extraction via an error structure, empty/non-JSON bodies, -//! and custom JSON paths — are covered below. - -use aimux_core::AiMuxError; -use aimux_provider_utils::{DEFAULT_ERROR_STRUCTURE, ErrorStructure, parse_provider_error}; - -const TEST_URL: &str = "test-url"; - -// --------------------------------------------------------------------------- -// createStatusCodeErrorResponseHandler -// --------------------------------------------------------------------------- - -#[test] -fn status_code_error_uses_status_text_for_non_json_body() { - // TS: "should create error with status text and response body". - // - // The TS handler sets `message` to `response.statusText` ("Not Found") and - // `responseBody` to the body text. Rust's `parse_provider_error` uses the - // body text as the message when it isn't JSON matching the error - // structure; a 404 maps to `ApiCall` (404). - let body = "Error message"; - let err = parse_provider_error(404, body, &DEFAULT_ERROR_STRUCTURE); - - assert!( - matches!(err, AiMuxError::ApiCall(ref m) if m.status_code == Some(404) && m.message == "Error message") - ); -} - -#[test] -fn status_code_500_maps_to_provider_error() { - // A 5xx-style status maps to the generic `Provider` variant. - let body = "Internal Server Error"; - let err = parse_provider_error(500, body, &DEFAULT_ERROR_STRUCTURE); - assert!(matches!(err, AiMuxError::ApiCall(_))); -} - -// --------------------------------------------------------------------------- -// createJsonErrorResponseHandler -// --------------------------------------------------------------------------- - -#[test] -fn json_error_response_extracts_message_from_default_structure() { - // TS: createJsonErrorResponseHandler parses the body JSON and maps - // `error.message` into the APICallError message. Rust's - // `parse_provider_error` with the default structure (`["error","message"]`) - // does the same; a 429 maps to `ApiCall` (429). - let body = r#"{"error":{"message":"Rate limit reached for requests","type":"requests"}}"#; - let err = parse_provider_error(429, body, &DEFAULT_ERROR_STRUCTURE); - assert!(matches!(err, AiMuxError::ApiCall(ref d) if d.status_code == Some(429))); -} - -#[test] -fn json_error_response_extracts_message_from_custom_structure() { - // Some providers nest the error differently; a custom `ErrorStructure` - // routes the message out of a different JSON path. The `Provider` arm - // prepends "HTTP {status}: " to the extracted message. - let structure = ErrorStructure { - message_path: &["error", "msg"], - type_path: &["error", "kind"], - }; - let body = r#"{"error":{"msg":"boom","kind":"failure"}}"#; - let err = parse_provider_error(500, body, &structure); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.message, "boom"); - assert_eq!(err.to_string(), "API call error: HTTP 500: boom"); -} - -#[test] -fn empty_body_falls_back_to_http_status_message() { - // TS: when the response body is empty, the handler uses `response.statusText` - // as the message. Here there is no body text at all, so the status carried in - // `ApiCallError` is the whole of what Display has to show. - let err = parse_provider_error(503, "", &DEFAULT_ERROR_STRUCTURE); - assert_eq!(err.status_code(), Some(503)); - assert!(err.to_string().contains("503"), "{err}"); -} - -#[test] -fn non_json_body_is_used_as_the_message() { - // A non-empty, non-JSON body is used verbatim as the message. - let err = parse_provider_error(401, "plain text problem", &DEFAULT_ERROR_STRUCTURE); - assert!( - matches!(err, AiMuxError::ApiCall(ref m) if m.status_code == Some(401) && m.message == "plain text problem") - ); -} - -#[test] -fn auth_status_maps_to_auth_variant() { - let body = r#"{"error":{"message":"Invalid API key","type":"auth"}}"#; - let err = parse_provider_error(401, body, &DEFAULT_ERROR_STRUCTURE); - assert!( - matches!(err, AiMuxError::ApiCall(ref m) if m.status_code == Some(401) && m.message == "Invalid API key") - ); -} - -// --------------------------------------------------------------------------- -// SKIPPED TS cases (no Rust counterpart) -// --------------------------------------------------------------------------- -// The following TS cases rely on the streaming Web `Response` object and a -// download size limit (`readResponseWithSizeLimit` / DEFAULT_MAX_DOWNLOAD_SIZE) -// that the Rust implementation does not model — providers consume an -// already-fetched `(status, body)` pair via `parse_provider_error`. They are -// intentionally not translated: -// -// - createJsonResponseHandler "should return both parsed value and rawValue" -// (Rust has no value/rawValue split; callers parse the body themselves.) -// - createJsonResponseHandler "should reject oversized responses before -// reading the body" (no size-limit gate in the Rust fetch path.) -// - createJsonErrorResponseHandler "should reject oversized responses before -// reading the body" (same.) -// - createBinaryResponseHandler "should handle binary response successfully" -// (Rust has no binary response handler; binary is read directly by callers.) -// - createBinaryResponseHandler "should throw APICallError when response body -// is null" (no binary handler to assert an empty-body error against.) -// - createStatusCodeErrorResponseHandler "should reject oversized responses -// before reading the body" (no size-limit gate.) - -#[test] -fn url_constant_is_unused_but_documented() { - // Kept so `TEST_URL` (mirroring the TS test's `testUrl` constant) is - // referenced; the Rust parse_provider_error API does not take a URL. - assert_eq!(TEST_URL, "test-url"); -} diff --git a/aimux-provider-utils/tests/retry_with_exponential_backoff.rs b/aimux-provider-utils/tests/retry_with_exponential_backoff.rs deleted file mode 100644 index 5851dc8f..00000000 --- a/aimux-provider-utils/tests/retry_with_exponential_backoff.rs +++ /dev/null @@ -1,492 +0,0 @@ -//! Rust translation of -//! `packages/ai/src/util/retry-with-exponential-backoff.test.ts`. -//! -//! The TS suite exercises `retryWithExponentialBackoffRespectingRetryHeaders`, -//! which retries on retryable errors and chooses each retry's delay from the -//! error's `retry-after-ms` / `retry-after` response headers when reasonable, -//! falling back to exponential backoff otherwise. -//! -//! In Rust, the equivalent is -//! [`aimux_provider_utils::retry_with_exponential_backoff_respecting_retry_headers`]. -//! The header-parsing and delay-selection logic lives in the pure helpers -//! [`aimux_provider_utils::parse_retry_after`] and -//! [`aimux_provider_utils::get_retry_delay_ms`]; the end-to-end timing behaviour -//! is driven through the async retry loop with a paused tokio clock -//! (`#[tokio::test(start_paused = true)]` + `tokio::time::advance`). -//! -//! The TS tests are written against `APICallError` carrying `responseHeaders`. -//! Rust's rate-limit error carries the parsed hint in the detail's -//! delay hint (exposed via `AiMuxError::retry_after_hint()`), so a -//! 429 `Provider` error stands in for an `APICallError` with a `retry-after-ms` -//! header. Errors without a hint (e.g. `ApiCall`, `Http`) stand in for -//! `APICallError` with no retry headers. - -use std::sync::Arc; -use std::sync::atomic::{AtomicU32, Ordering}; -use std::time::{Duration, SystemTime}; - -use aimux_core::{AiMuxError, ApiCallError}; - -/// A retryable transient failure: a 5xx `ApiCall` error. `is_retryable` is -/// stored at construction, as `error_for_status` does for a live 5xx. -fn server_error(msg: &str) -> AiMuxError { - AiMuxError::ApiCall(ApiCallError { - status_code: Some(500), - message: msg.into(), - is_retryable: true, - ..Default::default() - }) -} - -fn rate_limited(ms: u64) -> AiMuxError { - AiMuxError::ApiCall(ApiCallError { - status_code: Some(429), - retry_after_ms: Some(ms), - is_retryable: true, - ..Default::default() - }) -} -use aimux_provider_utils::{ - RetryConfig, get_retry_delay_ms, parse_retry_after, - retry_with_exponential_backoff_respecting_retry_headers, -}; - -/// Helper: build a boxed future-producing closure that counts attempts via a -/// shared counter and fails on the attempt numbers for which `fail_with` -/// returns `Some(error)`. -#[allow(clippy::type_complexity)] -fn failing_then_success( - counter: Arc, - fail_with: impl Fn(u32) -> Option + Send + Sync + 'static, -) -> Box< - dyn FnMut() -> std::pin::Pin< - Box> + Send>, - > + Send, -> { - // Erase to a shared closure so each spawned future can hold its own clone. - let fail_with: Arc Option + Send + Sync> = Arc::new(fail_with); - Box::new(move || { - let counter = counter.clone(); - let fail_with = fail_with.clone(); - Box::pin(async move { - let n = counter.fetch_add(1, Ordering::SeqCst) + 1; - if let Some(err) = fail_with(n) { - Err(err) - } else { - Ok("success".to_string()) - } - }) - }) -} - -/// Advance paused time and yield so the spawned retry task can process woken -/// timers and make progress. -async fn advance_and_yield(dur: Duration) { - tokio::time::advance(dur).await; - // Let the spawned task run any woken timers to completion. - for _ in 0..4 { - tokio::task::yield_now().await; - } -} - -// =========================================================================== -// Pure helper tests — mirror the TS `getRetryDelayInMs` edge cases that the -// TS suite asserts on indirectly through time advances. -// =========================================================================== - -#[test] -fn get_delay_uses_hint_when_reasonable() { - // TS: "should use rate limit header delay when present and reasonable" - assert_eq!(get_retry_delay_ms(Some(3000), 2000), 3000); -} - -#[test] -fn get_delay_uses_exponential_when_hint_too_long() { - // TS: "should use exponential backoff when rate limit delay is too long" - // (retry-after-ms 70000 is >= 60s and >= the 2000ms exponential delay) - assert_eq!(get_retry_delay_ms(Some(70_000), 2000), 2000); -} - -#[test] -fn get_delay_falls_back_when_negative() { - // TS: "should fall back to exponential backoff when rate limit delay is - // negative". The detail stores the delay as u64 so it cannot carry a - // negative value end-to-end; the negative branch is covered here via the - // pure helper that the retry loop uses. - assert_eq!(get_retry_delay_ms(Some(-1000), 2000), 2000); -} - -#[test] -fn get_delay_falls_back_when_no_hint() { - // TS: "should fall back to exponential backoff when no rate limit headers" - assert_eq!(get_retry_delay_ms(None, 2000), 2000); -} - -#[test] -fn parse_retry_after_ms_header_value() { - // TS: "should use rate limit header delay when present and reasonable" - assert_eq!( - parse_retry_after(Some("3000"), None, SystemTime::now()), - Some(3000) - ); -} - -#[test] -fn parse_retry_after_seconds_header() { - // TS: "should parse retry-after header in seconds" (5s -> 5000ms) - assert_eq!( - parse_retry_after(None, Some("5"), SystemTime::now()), - Some(5000) - ); -} - -#[test] -fn parse_retry_after_prefers_ms_over_seconds() { - // TS: "should prefer retry-after-ms over retry-after when both present" - assert_eq!( - parse_retry_after(Some("3000"), Some("10"), SystemTime::now()), - Some(3000) - ); -} - -#[test] -fn parse_retry_after_invalid_headers_yield_none() { - // TS: "should handle invalid rate limit header values" - assert_eq!( - parse_retry_after(Some("invalid"), Some("not-a-number"), SystemTime::now()), - None - ); -} - -#[test] -fn parse_retry_after_negative_ms() { - // TS: "should fall back to exponential backoff when rate limit delay is - // negative" — the parse step still surfaces the negative value; the - // "reasonable" clamp happens in get_retry_delay_ms (covered above). - assert_eq!( - parse_retry_after(Some("-1000"), None, SystemTime::now()), - Some(-1000) - ); -} - -#[test] -fn parse_retry_after_openai_seconds_header() { - // TS: "should handle OpenAI 429 response with retry-after header" (30s) - assert_eq!( - parse_retry_after(None, Some("30"), SystemTime::now()), - Some(30_000) - ); -} - -#[test] -fn parse_retry_after_http_date() { - // TS: "should handle retry-after header with HTTP date format" - let now = SystemTime::UNIX_EPOCH + Duration::from_millis(1_700_000_000_000); - let target = now + Duration::from_millis(5000); - let date_header = httpdate::fmt_http_date(target); - assert_eq!(parse_retry_after(None, Some(&date_header), now), Some(5000)); -} - -// =========================================================================== -// End-to-end retry tests — mirror the TS timing assertions via a paused clock. -// =========================================================================== - -#[tokio::test(start_paused = true)] -async fn uses_rate_limit_header_delay_when_reasonable() { - // TS: "should use rate limit header delay when present and reasonable" - // retry-after-ms 3000 -> the retry must wait ~3000ms, not the 2000ms - // exponential default. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(rate_limited(3000)) - } else { - None - } - }); - - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - RetryConfig::default(), - closure, - )); - tokio::task::yield_now().await; // run attempt 1, park on 3000ms sleep - - // Full Jitter (RFC-0009 §4.2): delay ∈ [0, 3000). Advancing past 3000ms - // guarantees the retry fired; it may fire earlier with jitter. - advance_and_yield(Duration::from_millis(3000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - let result = handle.await.unwrap().unwrap(); - assert_eq!(result, "success"); -} - -#[tokio::test(start_paused = true)] -async fn uses_exponential_backoff_when_delay_too_long() { - // TS: "should use exponential backoff when rate limit delay is too long" - // retry-after-ms 70000 (>= 60s) -> fall back to the 2000ms exponential delay. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(rate_limited(70_000)) - } else { - None - } - }); - - let config = RetryConfig { - initial_delay: Duration::from_millis(2000), - ..RetryConfig::default() - }; - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - config, closure, - )); - tokio::task::yield_now().await; - - // Full Jitter: delay ∈ [0, 2000). Advancing past 2000ms guarantees the retry. - advance_and_yield(Duration::from_millis(2000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn falls_back_to_exponential_when_no_rate_limit_headers() { - // TS: "should fall back to exponential backoff when no rate limit headers" - // A retryable error with no retry-after hint uses the 2000ms exponential delay. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(server_error("Temporary error")) - } else { - None - } - }); - - let config = RetryConfig { - initial_delay: Duration::from_millis(2000), - ..RetryConfig::default() - }; - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - config, closure, - )); - tokio::task::yield_now().await; - - // Full Jitter (RFC-0009 §4.2): delay ∈ [0, 2000). Advancing past 2000ms - // guarantees the retry fired; it may fire earlier with jitter. - advance_and_yield(Duration::from_millis(2000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn handles_anthropic_429_with_retry_after_ms() { - // TS: "should handle Anthropic 429 response with retry-after-ms header" - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(rate_limited(5000)) - } else { - None - } - }); - - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - RetryConfig::default(), - closure, - )); - tokio::task::yield_now().await; - - // Full Jitter: delay ∈ [0, 5000). Advancing past 5000ms guarantees the retry. - advance_and_yield(Duration::from_millis(5000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn multiple_retries_with_exponential_progression() { - // TS: "should handle multiple retries with exponential backoff progression" - // attempt 1 -> retry-after-ms 5000 (header wins, 5000 < 60s) - // attempt 2 -> retry-after-ms 2000 (header wins, 2000 < 60s); the - // exponential delay would be 4000ms by then, but the 2000ms header still - // wins per get_retry_delay_ms. The TS test advances 4000ms for the second - // retry, which covers the 2000ms delay. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| match n { - 1 => Some(rate_limited(5000)), - 2 => Some(rate_limited(2000)), - _ => None, - }); - - let config = RetryConfig { - max_retries: 3, - ..RetryConfig::default() - }; - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - config, closure, - )); - tokio::task::yield_now().await; - - // First retry uses the 5000ms header delay. - advance_and_yield(Duration::from_millis(5000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - // Second retry: header delay is 2000ms; advancing 4000ms covers it. - advance_and_yield(Duration::from_millis(4000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 3); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn retries_on_gateway_internal_server_error() { - // TS: "should retry on GatewayInternalServerError" — a retryable 5xx-style - // error with no retry-after hint uses exponential backoff. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(server_error("Internal server error")) - } else { - None - } - }); - - let config = RetryConfig { - initial_delay: Duration::from_millis(2000), - ..RetryConfig::default() - }; - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - config, closure, - )); - tokio::task::yield_now().await; - - // Full Jitter (RFC-0009 §4.2): delay ∈ [0, 2000). Advancing past 2000ms - // guarantees the retry fired; it may fire earlier with jitter. - advance_and_yield(Duration::from_millis(2000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn retries_on_gateway_rate_limit_error() { - // TS: "should retry on GatewayRateLimitError" — a retryable rate-limit - // error. Modelled as a 429 Provider error carrying a reasonable hint. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(rate_limited(2000)) - } else { - None - } - }); - - let config = RetryConfig { - initial_delay: Duration::from_millis(2000), - ..RetryConfig::default() - }; - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - config, closure, - )); - tokio::task::yield_now().await; - - // Full Jitter (RFC-0009 §4.2): delay ∈ [0, 2000). Advancing past 2000ms - // guarantees the retry fired; it may fire earlier with jitter. - advance_and_yield(Duration::from_millis(2000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn does_not_retry_on_non_retryable_auth_error() { - // TS: "should not retry on non-retryable GatewayAuthenticationError" - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |_| { - Some(AiMuxError::ApiCall(ApiCallError { - status_code: Some(401), - message: "Invalid API key".into(), - ..Default::default() - })) - }); - - let result = - retry_with_exponential_backoff_respecting_retry_headers(RetryConfig::default(), closure) - .await; - - // Auth is not retryable: a single attempt, no delay, error propagates. - assert_eq!(counter.load(Ordering::SeqCst), 1); - let err = result.unwrap_err(); - assert!(matches!(err, AiMuxError::ApiCall(ref m) if m.message == "Invalid API key")); -} - -#[tokio::test(start_paused = true)] -async fn uses_retry_after_hint_from_wrapped_error() { - // TS: "should use retry-after headers from APICallError cause" — the TS - // unwraps a GatewayInternalServerError's `cause` (an APICallError with - // retry-after-ms) to find the hint. Rust's flat `AiMuxError` has no cause - // chain, so the hint is modelled directly on the outer error (a - // `RateLimited` carrying the 3000ms hint). - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |n| { - if n == 1 { - Some(rate_limited(3000)) - } else { - None - } - }); - - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - RetryConfig::default(), - closure, - )); - tokio::task::yield_now().await; - - // Full Jitter: delay ∈ [0, 3000). Advancing past 3000ms guarantees the retry. - advance_and_yield(Duration::from_millis(3000)).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - - assert_eq!(handle.await.unwrap().unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn succeeds_on_first_attempt_without_delay() { - // Sanity check (no direct TS counterpart, but guards the happy path): a - // function that succeeds immediately must not sleep. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |_| None); - - let result = - retry_with_exponential_backoff_respecting_retry_headers(RetryConfig::default(), closure) - .await; - assert_eq!(counter.load(Ordering::SeqCst), 1); - assert_eq!(result.unwrap(), "success"); -} - -#[tokio::test(start_paused = true)] -async fn gives_up_after_max_retries() { - // TS counterpart: the suite's `maxRetries` config bounds attempts. A - // persistently-retryable error must surface after `max_retries + 1` tries. - let counter = Arc::new(AtomicU32::new(0)); - let closure = failing_then_success(counter.clone(), |_| Some(server_error("always fails"))); - - let config = RetryConfig { - max_retries: 2, - initial_delay: Duration::from_millis(10), - backoff_factor: 2, - }; - let handle = tokio::spawn(retry_with_exponential_backoff_respecting_retry_headers( - config, closure, - )); - - // max_retries=2 -> 3 attempts total (1 + 2 retries). - // Delays: 10ms, 20ms. Advance past both. - advance_and_yield(Duration::from_millis(10)).await; - advance_and_yield(Duration::from_millis(20)).await; - advance_and_yield(Duration::from_millis(50)).await; - - let result = handle.await.unwrap(); - assert!(result.is_err()); - assert_eq!(counter.load(Ordering::SeqCst), 3); -} diff --git a/aimux-provider-utils/tests/structured_error_fields.rs b/aimux-provider-utils/tests/structured_error_fields.rs deleted file mode 100644 index 53b5f586..00000000 --- a/aimux-provider-utils/tests/structured_error_fields.rs +++ /dev/null @@ -1,307 +0,0 @@ -//! The two entry points that turn an upstream failure into an `AiMuxError` -//! (`parse_provider_error` for a completed HTTP response, `parse_stream_error` -//! for an SSE `error` event) must agree with each other and must put the -//! structure in *fields*, not in the message string. - -use aimux_core::AiMuxError; -use aimux_provider_utils::parse_provider_error; -use aimux_provider_utils::response::{DEFAULT_ERROR_STRUCTURE, parse_stream_error}; -use serde_json::json; - -/// H4: the error `type` used to be extracted and thrown away (`_error_type`). -/// It is the provider's machine-readable code and now survives as a field. -#[test] -fn parse_provider_error_fills_status_and_provider_code() { - let body = r#"{"error":{"message":"boom","type":"server_error"}}"#; - let err = parse_provider_error(500, body, &DEFAULT_ERROR_STRUCTURE); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.status_code, Some(500)); - assert_eq!(detail.provider_code.as_deref(), Some("server_error")); - // The message is the provider's text alone; the status is not written into it. - assert_eq!(detail.message, "boom"); - assert_eq!(err.to_string(), "API call error: HTTP 500: boom"); - // H1: the status is read from the field, not re-parsed out of the message. - assert_eq!(err.status_code(), Some(500)); - // The raw body survives alongside the extraction. - assert_eq!(detail.response_body.as_deref(), Some(body)); -} - -/// The AI SDK keeps `responseBody` on every error branch, so a consumer never -/// has to ask whether it is populated. Ours does the same: a body that could -/// not be parsed (and therefore *became* the message) is still kept verbatim, -/// and only a genuinely empty body yields `None`. -#[test] -fn response_body_survives_every_parse_outcome() { - // Unparseable body: it becomes the message *and* stays as raw evidence. - let body = "upstream exploded"; - let err = parse_provider_error(500, body, &DEFAULT_ERROR_STRUCTURE); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.message, body); - assert_eq!(detail.response_body.as_deref(), Some(body)); - - // JSON that the configured path does not resolve to a string: same rule. - let body = r#"{"code":400,"error":"bad request"}"#; - let err = parse_provider_error(400, body, &DEFAULT_ERROR_STRUCTURE); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.response_body.as_deref(), Some(body)); - - // No body at all: nothing to keep. - let err = parse_provider_error(500, "", &DEFAULT_ERROR_STRUCTURE); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.response_body, None); - assert_eq!(err.to_string(), "API call error: HTTP 500"); -} - -/// The 429 detail keeps the provider code, so "retry helps" -/// (`rate_limit_exceeded`) and "retry is pointless" (`insufficient_quota`) -/// are finally distinguishable without string matching. -#[test] -fn rate_limited_keeps_the_provider_code() { - let body = r#"{"error":{"message":"quota exhausted","type":"insufficient_quota"}}"#; - let err = parse_provider_error(429, body, &DEFAULT_ERROR_STRUCTURE); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.provider_code.as_deref(), Some("insufficient_quota")); - assert_eq!(detail.message, "quota exhausted"); - assert_eq!(detail.response_body.as_deref(), Some(body)); - // Display composes from the fields; no header means no hint — the retry - // loop's exponential backoff covers it, nothing is fabricated. - assert_eq!(err.to_string(), "API call error: HTTP 429: quota exhausted"); - assert_eq!(err.retry_after_hint(), None); -} - -/// A 401/404 is a `Provider` error classified by the status field — there is -/// no variant to imply anything, and nothing is baked into the message. -#[test] -fn auth_and_not_found_are_classified_by_the_field() { - let body = - r#"{"error":{"message":"Incorrect API key provided","type":"invalid_request_error"}}"#; - let err = parse_provider_error(401, body, &DEFAULT_ERROR_STRUCTURE); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!( - detail.provider_code.as_deref(), - Some("invalid_request_error") - ); - assert_eq!(err.status_code(), Some(401)); - assert_eq!( - err.to_string(), - "API call error: HTTP 401: Incorrect API key provided" - ); - assert!(!err.is_retryable()); - - // Empty body: Display composes the status from the field on its own. - let err = parse_provider_error(404, "", &DEFAULT_ERROR_STRUCTURE); - assert_eq!(err.to_string(), "API call error: HTTP 404"); - assert_eq!(err.status_code(), Some(404)); -} - -/// `request_id` set on the detail rides the serialized error. -#[test] -fn request_id_rides_the_detail() { - let mut err = parse_provider_error(500, "boom", &DEFAULT_ERROR_STRUCTURE); - if let AiMuxError::ApiCall(d) = &mut err { - d.request_id = Some("req_abc123".into()); - } - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.request_id.as_deref(), Some("req_abc123")); - let json = serde_json::to_string(&err).unwrap(); - let back: AiMuxError = serde_json::from_str(&json).unwrap(); - let AiMuxError::ApiCall(ref detail) = back else { - panic!("expected ApiCall, got {back:?}") - }; - assert_eq!(detail.request_id.as_deref(), Some("req_abc123")); -} - -#[test] -fn parse_provider_error_fills_rate_limit_fields() { - let body = r#"{"error":{"message":"slow down","type":"rate_limit_exceeded"}}"#; - let err = parse_provider_error(429, body, &DEFAULT_ERROR_STRUCTURE); - // A 429 is an ApiCall error; retryability comes from the stored field, - // and without a retry-after header there is no hint to store. - assert_eq!(err.status_code(), Some(429)); - assert!(err.is_retryable()); - assert_eq!(err.retry_after_hint(), None); -} - -/// M3: an OpenAI-shaped stream error carries a *string* `code`. Reading it as -/// an HTTP status produced nonsense statuses (and, when it parsed as 429, a -/// hardcoded 1000ms retry). It is a provider code, and there is no status — -/// the response itself was a 200. -#[test] -fn stream_error_string_code_is_a_provider_code_not_a_status() { - let err = parse_stream_error(&json!({ - "message": "The server had an error processing your request.", - "type": "server_error", - "param": null, - "code": "internal_error", - })); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.provider_code.as_deref(), Some("internal_error")); - assert_eq!(detail.status_code, None); - assert_eq!(err.status_code(), None); -} - -#[test] -fn stream_error_falls_back_to_type_when_code_is_absent() { - let err = parse_stream_error(&json!({"message": "boom", "type": "server_error"})); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.provider_code.as_deref(), Some("server_error")); -} - -/// A numeric `code` in the HTTP range (Google-shaped payloads) *is* a status, -/// and then the stream path maps it exactly like the response path. -#[test] -fn stream_error_numeric_code_is_a_status() { - for status in [401u16, 404, 429, 500] { - let stream = parse_stream_error(&json!({"message": "boom", "code": status})); - let response = parse_provider_error( - status, - r#"{"error":{"message":"boom"}}"#, - &DEFAULT_ERROR_STRUCTURE, - ); - assert_eq!( - std::mem::discriminant(&stream), - std::mem::discriminant(&response), - "status {status} maps to different variants on the two paths" - ); - assert_eq!(stream.to_string(), response.to_string(), "status {status}"); - } -} - -/// M3: the retry hint is no longer hardcoded when the payload carries one. -#[test] -fn stream_error_reads_the_retry_hint_from_the_payload() { - let ms = parse_stream_error(&json!({"message": "slow", "code": 429, "retry_after_ms": 250})); - assert_eq!(ms.retry_after_hint(), Some(250)); - - let secs = parse_stream_error(&json!({"message": "slow", "code": 429, "retry_after": 3})); - assert_eq!(secs.retry_after_hint(), Some(3000)); - - // No hint in the payload: none is stored — nothing is fabricated. - let none = parse_stream_error(&json!({"message": "slow", "code": 429})); - assert_eq!(none.retry_after_hint(), None); -} - -/// Audit §2.2-3: a given status must imply the same retryability whichever -/// entry point produced the error, and the re-mapping helpers must not flip it -/// on the way through. -#[test] -fn retryability_is_invariant_across_the_conversion_paths() { - for status in [400u16, 401, 403, 404, 422, 429, 500, 502, 503] { - let response = parse_provider_error( - status, - r#"{"error":{"message":"boom","type":"some_code"}}"#, - &DEFAULT_ERROR_STRUCTURE, - ); - let stream = parse_stream_error(&json!({"message": "boom", "code": status})); - assert_eq!( - response.is_retryable(), - stream.is_retryable(), - "status {status}: response path {response:?} vs stream path {stream:?}" - ); - assert_eq!(response.status_code(), Some(status)); - assert_eq!(stream.status_code(), Some(status)); - } -} - -/// §9.1 retry policy: 408/409 join 429/5xx as retryable; a representative -/// ordinary 4xx is not. -#[test] -fn retry_policy_includes_408_and_409() { - for status in [408u16, 409, 429, 500, 503] { - assert!( - parse_provider_error(status, "", &DEFAULT_ERROR_STRUCTURE).is_retryable(), - "status {status} must be retryable" - ); - } - assert!(!parse_provider_error(400, "", &DEFAULT_ERROR_STRUCTURE).is_retryable()); -} - -/// §9.1 shared non-2xx regression, through the real HTTP loop: every non-2xx -/// goes through the provider error parser (extracted message/code + raw body -/// survive), the observed request id is attached to the same detail, and a -/// hint-less 429 keeps `retry_after_ms=None` — never a fabricated 1000. -#[tokio::test] -async fn http_loop_preserves_parsed_fields_on_429_and_500() { - use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; - use wiremock::matchers::method; - use wiremock::{Mock, MockServer, ResponseTemplate}; - - let body = r#"{"error":{"message":"boom","type":"some_code"}}"#; - for status in [429u16, 500] { - let server = MockServer::start().await; - Mock::given(method("POST")) - .respond_with( - ResponseTemplate::new(status) - .insert_header("x-request-id", "req_42") - .set_body_string(body), - ) - .mount(&server) - .await; - - let request = HttpRequest { - method: HttpMethod::Post, - url: server.uri(), - headers: vec![], - body: HttpBody::Json(serde_json::json!({})), - abort_signal: None, - call_id: None, - recording_context: None, - }; - let no_retry = RetryConfig { - max_retries: 0, - ..RetryConfig::default() - }; - let err = send(request, no_retry, &DEFAULT_ERROR_STRUCTURE) - .await - .unwrap_err(); - let AiMuxError::ApiCall(ref detail) = err else { - panic!("expected ApiCall, got {err:?}") - }; - assert_eq!(detail.status_code, Some(status)); - assert_eq!(detail.message, "boom", "status {status}"); - assert_eq!(detail.provider_code.as_deref(), Some("some_code")); - assert_eq!(detail.response_body.as_deref(), Some(body)); - assert_eq!(detail.request_id.as_deref(), Some("req_42")); - assert_eq!(detail.retry_after_ms, None, "no header means no hint"); - assert!(detail.is_retryable); - } -} - -/// A 403 needs no re-mapping helper any more: the observed status *is* the -/// classification, and a message that merely mentions a status is not one. -#[test] -fn observed_status_is_the_classification() { - let err = parse_provider_error( - 403, - r#"{"error":{"message":"AccessDenied"}}"#, - &DEFAULT_ERROR_STRUCTURE, - ); - assert_eq!(err.status_code(), Some(403)); - assert!(matches!(err, AiMuxError::ApiCall(ref m) if m.message == "AccessDenied")); - - // A 500 whose *message* happens to mention 403 stays a 500. - let err = parse_provider_error( - 500, - r#"{"error":{"message":"HTTP 403: not really"}}"#, - &DEFAULT_ERROR_STRUCTURE, - ); - assert_eq!(err.status_code(), Some(500)); -} From abb4291ba8389658136ba473cc876e86dcdb5d0b Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Mon, 24 Aug 2026 22:29:11 +0800 Subject: [PATCH 03/24] refactor(providers)!: migrate to Core retry and the handler API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every provider moves from the deleted send*/ErrorStructure surface to the single-exchange helpers with typed success/failure handlers; retry policy is declared through retry_config() and executed by Core. In-band stream errors route through stream_error_api_call (or redact their request context) so credentials and data URLs never reach the public error. list_models exchanges — not Core user operations — apply the Core retry primitive locally with the provider's configured settings. Failed-response mappers now match each provider's documented error format instead of assuming OpenAI's {error:{message,code}} shape: ElevenLabs detail.status/detail.message (FastAPI variants included), AssemblyAI's plain {error: string}, Deepgram err_code/err_msg and category/message, fal's FastAPI detail, Hume's top-level message/code. provider_code extraction no longer JSON-serializes string codes into quoted values. --- aimux-providers/Cargo.toml | 1 + aimux-providers/src/anthropic/files.rs | 17 +- aimux-providers/src/anthropic/mod.rs | 54 +-- aimux-providers/src/anthropic/model.rs | 25 +- aimux-providers/src/anthropic/stream.rs | 156 +++++--- aimux-providers/src/anthropic_aws/mod.rs | 40 +-- aimux-providers/src/anthropic_aws/model.rs | 18 +- aimux-providers/src/assemblyai.rs | 128 ++++--- aimux-providers/src/aws_polly.rs | 104 +++--- aimux-providers/src/azure/model.rs | 53 ++- aimux-providers/src/azure/responses.rs | 76 ++-- aimux-providers/src/bedrock/embedding.rs | 22 +- aimux-providers/src/bedrock/image.rs | 39 +- aimux-providers/src/bedrock/mod.rs | 63 ++-- aimux-providers/src/bedrock/model.rs | 111 ++++-- aimux-providers/src/bedrock/reranking.rs | 32 +- aimux-providers/src/bedrock/sigv4.rs | 2 +- aimux-providers/src/black_forest_labs.rs | 128 ++++--- aimux-providers/src/cartesia.rs | 52 ++- aimux-providers/src/catalogue.rs | 44 ++- aimux-providers/src/codex.rs | 116 +++--- aimux-providers/src/cohere/embedding.rs | 20 +- aimux-providers/src/cohere/mod.rs | 54 +-- aimux-providers/src/cohere/model.rs | 146 +++++--- aimux-providers/src/cohere/reranking.rs | 23 +- aimux-providers/src/cohere/types.rs | 14 +- aimux-providers/src/dataforseo.rs | 41 ++- aimux-providers/src/deepgram.rs | 54 ++- aimux-providers/src/elevenlabs.rs | 70 +++- aimux-providers/src/exa_ai.rs | 39 +- aimux-providers/src/fal.rs | 335 ++++++++++-------- aimux-providers/src/firecrawl.rs | 38 +- aimux-providers/src/gladia.rs | 112 +++--- aimux-providers/src/google/embedding.rs | 41 +-- aimux-providers/src/google/files.rs | 94 ++--- aimux-providers/src/google/image.rs | 40 +-- aimux-providers/src/google/mod.rs | 89 +++-- aimux-providers/src/google/model.rs | 163 +++------ aimux-providers/src/google/types.rs | 9 + aimux-providers/src/google/video.rs | 161 +++++---- aimux-providers/src/google_pse.rs | 50 +-- aimux-providers/src/huggingface.rs | 2 +- aimux-providers/src/huggingface/responses.rs | 206 +++++++---- aimux-providers/src/hume.rs | 49 ++- aimux-providers/src/jina_ai.rs | 37 +- aimux-providers/src/klingai.rs | 263 ++++++++------ aimux-providers/src/linkup.rs | 39 +- aimux-providers/src/lmnt.rs | 36 +- aimux-providers/src/luma.rs | 123 ++++--- aimux-providers/src/mistral/embedding.rs | 20 +- aimux-providers/src/mistral/mod.rs | 96 +++-- aimux-providers/src/mistral/model.rs | 110 +++--- aimux-providers/src/open_responses.rs | 170 +++++---- aimux-providers/src/openai/embedding.rs | 26 +- aimux-providers/src/openai/files.rs | 16 +- aimux-providers/src/openai/image.rs | 38 +- aimux-providers/src/openai/mod.rs | 68 +++- aimux-providers/src/openai/model.rs | 175 ++++----- aimux-providers/src/openai/responses/mod.rs | 65 ++-- .../src/openai/responses/responses_convert.rs | 162 +++++---- aimux-providers/src/openai/speech.rs | 20 +- aimux-providers/src/openai/transcription.rs | 52 +-- aimux-providers/src/parallel_ai.rs | 39 +- aimux-providers/src/prodia.rs | 187 ++++++---- aimux-providers/src/provider.rs | 11 +- aimux-providers/src/recraft.rs | 80 +++-- aimux-providers/src/replay.rs | 7 +- aimux-providers/src/replicate.rs | 215 ++++++----- aimux-providers/src/revai.rs | 160 +++++---- aimux-providers/src/runwayml.rs | 225 ++++++------ aimux-providers/src/searxng.rs | 27 +- aimux-providers/src/serper.rs | 40 ++- aimux-providers/src/stability.rs | 110 ++++-- aimux-providers/src/tavily.rs | 39 +- aimux-providers/src/tinyfish.rs | 48 +-- aimux-providers/src/vertex/anthropic_model.rs | 125 +++---- aimux-providers/src/vertex/embedding.rs | 43 ++- aimux-providers/src/vertex/image.rs | 41 +-- aimux-providers/src/vertex/mod.rs | 48 +-- aimux-providers/src/vertex/model.rs | 124 +++---- aimux-providers/src/vertex/transcription.rs | 36 +- aimux-providers/src/vertex/video.rs | 161 +++++---- aimux-providers/src/voyage/embedding.rs | 19 +- aimux-providers/src/voyage/mod.rs | 14 + aimux-providers/src/voyage/reranking.rs | 21 +- aimux-providers/src/xai/mod.rs | 207 ++++++++++- aimux-providers/src/xai/model.rs | 185 +++------- aimux-providers/src/xai/responses/mod.rs | 149 +++----- aimux-providers/src/you_com.rs | 48 +-- .../tests/anthropic_aws_model_test.rs | 6 +- .../tests/bedrock_embedding_test.rs | 6 + .../tests/bedrock_reranking_test.rs | 4 + .../tests/cassette_multimodal_test.rs | 4 + aimux-providers/tests/codex_test.rs | 41 ++- .../tests/cohere_embedding_test.rs | 8 + aimux-providers/tests/cohere_model_test.rs | 39 +- .../tests/cohere_reranking_test.rs | 4 + aimux-providers/tests/e2e_test.rs | 10 +- aimux-providers/tests/fal_video_test.rs | 27 +- .../tests/google_embedding_test.rs | 8 + aimux-providers/tests/google_video_test.rs | 255 ++++++++++++- .../tests/huggingface_responses_test.rs | 42 ++- aimux-providers/tests/jina_ai_test.rs | 4 + aimux-providers/tests/klingai_video_test.rs | 29 +- aimux-providers/tests/list_models_test.rs | 7 +- .../tests/mistral_embedding_test.rs | 4 + .../tests/openai_embedding_test.rs | 6 + aimux-providers/tests/openai_model_test.rs | 275 +++++++++++--- .../tests/openai_transcription_stream_test.rs | 12 +- aimux-providers/tests/prodia_video_test.rs | 21 +- aimux-providers/tests/provider_error_test.rs | 116 ++---- aimux-providers/tests/replicate_video_test.rs | 21 +- aimux-providers/tests/runwayml_test.rs | 29 +- .../tests/stream_error_semantics_test.rs | 139 ++++++++ .../tests/vertex_anthropic_test.rs | 4 +- .../tests/vertex_embedding_test.rs | 12 +- aimux-providers/tests/vertex_model_test.rs | 6 +- aimux-providers/tests/vertex_video_test.rs | 21 +- .../tests/voyage_embedding_test.rs | 8 + .../tests/voyage_reranking_test.rs | 4 + 120 files changed, 4993 insertions(+), 3235 deletions(-) create mode 100644 aimux-providers/tests/stream_error_semantics_test.rs diff --git a/aimux-providers/Cargo.toml b/aimux-providers/Cargo.toml index 9f2c9bef..942abdb5 100644 --- a/aimux-providers/Cargo.toml +++ b/aimux-providers/Cargo.toml @@ -23,6 +23,7 @@ serde_json = { workspace = true } reqwest = { workspace = true } tokio = { workspace = true } futures = { workspace = true } +bytes = { workspace = true } tracing = { workspace = true } async-stream = { workspace = true } base64 = "0.22" diff --git a/aimux-providers/src/anthropic/files.rs b/aimux-providers/src/anthropic/files.rs index c75a4126..fcc90d57 100644 --- a/aimux-providers/src/anthropic/files.rs +++ b/aimux-providers/src/anthropic/files.rs @@ -14,8 +14,7 @@ use aimux_core::error::AiMuxError; use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData, UploadFileResult}; use aimux_core::shared::FileBytes; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send}; +use aimux_provider_utils::{HttpBody, HttpRequest}; use super::AnthropicConfig; @@ -57,7 +56,7 @@ struct AnthropicFilesResponse { /// An Anthropic Files interface for uploading files. /// /// Aligned with TS `AnthropicFiles`. Does **not** hold an HTTP client — -/// `http::send` uses the process-wide shared `Client` internally (RFC-0009 §4.1). +/// the `aimux-provider-utils` API helpers use the process-wide shared `Client` internally (RFC-0009 §4.1). pub struct AnthropicFiles { config: AnthropicConfig, } @@ -131,24 +130,22 @@ impl Files for AnthropicFiles { // and its content-type are carried by `HttpBody::Bytes` — the HTTP layer // sets `Content-Type` from it, so it is intentionally not added to the // header list above. - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Bytes(body, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body, content_type), + aimux_provider_utils::create_json_response_handler(), + super::anthropic_failed_response_handler(), ) .await?; - let data: AnthropicFilesResponse = - serde_json::from_slice::(&resp.body)?; + let data: AnthropicFilesResponse = resp.value; // Build provider metadata. let mut metadata = serde_json::Map::new(); diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs index 670856d2..a9c7a17d 100644 --- a/aimux-providers/src/anthropic/mod.rs +++ b/aimux-providers/src/anthropic/mod.rs @@ -19,6 +19,21 @@ use aimux_core::provider::Provider; use aimux_provider_utils::{RetryConfig, load_api_key}; use serde_json::Value; +pub(crate) fn anthropic_failed_response_handler() +-> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error.get("type").and_then(Value::as_str).map(str::to_owned), + } + }) +} + /// The bare (unversioned) Anthropic API URL. const ANTHROPIC_API_URL: &str = "https://api.anthropic.com"; @@ -41,7 +56,7 @@ pub struct AnthropicConfig { pub name: String, /// Extra headers merged into every request. pub headers: Option>, - /// 重试配置。默认 `RetryConfig::default()`（max_retries=2）。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, /// Provider 级请求体覆盖（RFC-0017）。在标准请求体之后 deep-merge。 pub body_overrides: Option, @@ -315,24 +330,23 @@ impl Provider for AnthropicProvider { let mut header_list: Vec<(String, String)> = headers.into_iter().collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers: header_list, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: header_list.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + anthropic_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { @@ -345,7 +359,7 @@ impl Provider for AnthropicProvider { #[serde(default)] display_name: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .data .into_iter() diff --git a/aimux-providers/src/anthropic/model.rs b/aimux-providers/src/anthropic/model.rs index 49829172..f09e8356 100644 --- a/aimux-providers/src/anthropic/model.rs +++ b/aimux-providers/src/anthropic/model.rs @@ -17,7 +17,6 @@ use super::AnthropicConfig; use super::convert::build_request_body_with_warnings; use super::stream::{BodyEncoding, anthropic_generate_core, anthropic_stream_core}; use super::tool_name_mapping::ToolNameMapping; -use aimux_provider_utils::RetryConfig; /// An Anthropic language model (e.g. `claude-sonnet-4-20250514`). pub struct AnthropicModel { @@ -91,6 +90,10 @@ impl LanguageModel for AnthropicModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -121,16 +124,13 @@ impl LanguageModel for AnthropicModel { let req = build_request_body_with_warnings(&self.model_id, &options, false)?; let endpoint = self.endpoint(); let build_headers = self.make_header_builder(options.headers.as_ref(), req.betas); - let retry_config = resolve_anthropic_retry(&self.config.retry_config, options.max_retries); anthropic_generate_core( &endpoint, - retry_config, req.body, req.warnings, build_headers, BodyEncoding::Json, options.abort_signal.clone(), - options.timeout.map(Into::into), options.recording_context.clone(), &ToolNameMapping::new(options.tools.as_deref()), ) @@ -142,16 +142,13 @@ impl LanguageModel for AnthropicModel { let req = build_request_body_with_warnings(&self.model_id, &options, true)?; let endpoint = self.endpoint(); let build_headers = self.make_header_builder(options.headers.as_ref(), req.betas); - let retry_config = resolve_anthropic_retry(&self.config.retry_config, options.max_retries); anthropic_stream_core( &endpoint, - retry_config, req.body, req.warnings, build_headers, BodyEncoding::Json, options.abort_signal.clone(), - options.timeout.map(Into::into), options.recording_context.clone(), ToolNameMapping::new(options.tools.as_deref()), ) @@ -159,20 +156,6 @@ impl LanguageModel for AnthropicModel { } } -/// Resolve per-call max_retries override (RFC-0017). RetryConfig is Copy. -fn resolve_anthropic_retry( - provider: &RetryConfig, - max_retries_override: Option, -) -> RetryConfig { - match max_retries_override { - Some(n) => RetryConfig { - max_retries: n, - ..*provider - }, - None => *provider, - } -} - /// Merge provider-level body_overrides into per-call options (RFC-0017). fn merge_anthropic_body_overrides( options: &CallOptions, diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs index 28f232f9..66f9bd17 100644 --- a/aimux-providers/src/anthropic/stream.rs +++ b/aimux-providers/src/anthropic/stream.rs @@ -20,23 +20,51 @@ use std::sync::atomic::{AtomicU64, Ordering}; use futures::StreamExt; +use aimux_core::AbortSignal; use aimux_core::error::AiMuxError; -use aimux_core::error::ApiCallError; use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; -use aimux_core::shared::AbortSignal; use aimux_core::stream_part::StreamPart; use aimux_core::types::Warning; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, TokenUsage, Usage}; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RequestTimeout, RetryConfig, send_stream_timed, send_timed, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::{HttpBody, HttpRequest}; use serde_json::{Value, json}; use super::convert::parse_stop_reason; use super::tool_name_mapping::ToolNameMapping; -use super::types::{AnthropicResponse, ContentBlock, StreamEvent}; +use super::types::{AnthropicResponse, ContentBlock, StreamErrorData, StreamEvent}; + +pub(crate) fn anthropic_stream_error( + error: &StreamErrorData, + url: &str, + request_body_values: serde_json::Value, + response_headers: HashMap, +) -> AiMuxError { + // Documented Anthropic error types map to their documented statuses; an + // unknown type carries no status and is not retried (fabricating a 500 + // here was the M3 bug). + let status_code = match error.error_type.as_deref() { + Some("rate_limit_error") => Some(429), + Some("overloaded_error") => Some(529), + Some("api_error") => Some(500), + _ => None, + }; + let data = json!({ + "type": "error", + "error": { + "type": error.error_type, + "message": error.message, + } + }); + aimux_provider_utils::stream_error_api_call( + error.message.clone(), + error.error_type.clone(), + status_code, + &data, + url, + request_body_values, + response_headers, + ) +} /// How the request body is sent over the wire. #[derive(Debug, Clone, Copy)] @@ -66,7 +94,7 @@ fn build_anthropic_request( body_encoding: BodyEncoding, abort_signal: Option, recording_context: Option, -) -> Result { +) -> Result<(HttpRequest, HttpBody), AiMuxError> { // Serialize once; the Bytes path sends these exact bytes and the closure // signs over them, guaranteeing signature/body agreement. let body_bytes = serde_json::to_vec(body).map_err(|e| AiMuxError::JsonParse(e.to_string()))?; @@ -77,15 +105,16 @@ fn build_anthropic_request( BodyEncoding::Bytes => HttpBody::Bytes(body_bytes, "application/json".to_string()), }; - Ok(HttpRequest { - method: HttpMethod::Post, - url: endpoint.to_string(), - headers, - body: http_body, - abort_signal, - call_id: recording_context.as_ref().map(|c| c.call_id.clone()), - recording_context, - }) + Ok(( + HttpRequest { + url: endpoint.to_string(), + headers, + abort_signal, + call_id: recording_context.as_ref().map(|c| c.call_id.clone()), + recording_context, + }, + http_body, + )) } /// Read a string field, dropping absent / non-string values. @@ -681,17 +710,15 @@ pub(crate) fn parse_anthropic_content( #[allow(clippy::too_many_arguments)] // core plumbing: endpoint/retry/body/warnings/auth/encoding/abort/timeout pub(crate) async fn anthropic_generate_core( endpoint: &str, - retry_config: RetryConfig, body: serde_json::Value, warnings: Vec, build_headers: impl Fn(&[u8], &str) -> Result, AiMuxError>, body_encoding: BodyEncoding, abort_signal: Option, - timeout: Option, recording_context: Option, tool_names: &ToolNameMapping, ) -> Result { - let request = build_anthropic_request( + let (request, request_body) = build_anthropic_request( endpoint, &body, &build_headers, @@ -699,9 +726,15 @@ pub(crate) async fn anthropic_generate_core( abort_signal, recording_context, )?; - let resp = send_timed(request, retry_config, &DEFAULT_ERROR_STRUCTURE, timeout).await?; + let resp = aimux_provider_utils::post_to_api( + request, + request_body, + aimux_provider_utils::create_json_response_handler(), + super::anthropic_failed_response_handler(), + ) + .await?; - let data: AnthropicResponse = serde_json::from_slice(&resp.body)?; + let data: AnthropicResponse = resp.value; let content = parse_anthropic_content(&data.content, tool_names); @@ -769,17 +802,15 @@ enum BlockState { #[allow(clippy::too_many_arguments)] // core plumbing: endpoint/retry/body/warnings/auth/encoding/abort/timeout pub(crate) async fn anthropic_stream_core( endpoint: &str, - retry_config: RetryConfig, body: serde_json::Value, warnings: Vec, build_headers: impl Fn(&[u8], &str) -> Result, AiMuxError>, body_encoding: BodyEncoding, abort_signal: Option, - timeout: Option, recording_context: Option, tool_names: ToolNameMapping, ) -> Result { - let request = build_anthropic_request( + let (request, request_body) = build_anthropic_request( endpoint, &body, &build_headers, @@ -787,16 +818,37 @@ pub(crate) async fn anthropic_stream_core( abort_signal, recording_context, )?; - let resp = send_stream_timed(request, retry_config, &DEFAULT_ERROR_STRUCTURE, timeout).await?; + let resp = aimux_provider_utils::post_to_api( + request, + request_body, + aimux_provider_utils::create_event_source_response_handler::(), + super::anthropic_failed_response_handler(), + ) + .await?; - let response_headers = resp.headers; - let sse_stream = SseStream::new(resp.body); + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(StreamEvent::Error { error })) = first_event.as_ref() { + return Err(anthropic_stream_error( + error, + endpoint, + body.clone(), + response_headers, + )); + } + let stream_error_url = endpoint.to_string(); + let stream_request_body = body.clone(); + let stream_response_headers = response_headers.clone(); let stream = async_stream::stream! { // First part: StreamStart. yield Ok(StreamPart::StreamStart { warnings }); - let mut sse = sse_stream; + let mut sse = futures::stream::iter(first_event.into_iter()).chain(sse_stream); let mut blocks: HashMap = HashMap::new(); let mut final_usage = Usage::default(); let mut final_finish_reason: Option = None; @@ -808,9 +860,9 @@ pub(crate) async fn anthropic_stream_core( while let Some(event) = sse.next().await { match event { - Ok(sse_event) => { - match serde_json::from_str::(&sse_event.data) { - Ok(StreamEvent::MessageStart { message }) => { + Ok(stream_event) => { + match stream_event { + StreamEvent::MessageStart { message } => { if let Some(usage) = &message.usage { // RFC-0015 P0-2: full input side incl. cache // fields + raw (Anthropic reports cache only @@ -826,7 +878,7 @@ pub(crate) async fn anthropic_stream_core( response_meta_emitted = true; } } - Ok(StreamEvent::ContentBlockStart { index, content_block }) => { + StreamEvent::ContentBlockStart { index, content_block } => { match content_block { ContentBlock::Text { .. } => { blocks.insert(index, BlockState::Text { started: false }); @@ -921,7 +973,7 @@ pub(crate) async fn anthropic_stream_core( } } } - Ok(StreamEvent::ContentBlockDelta { index, delta }) => { + StreamEvent::ContentBlockDelta { index, delta } => { if let Some(text) = delta.text { // Start the text segment on the first delta. The // text id is the stringified content-block @@ -1022,7 +1074,7 @@ pub(crate) async fn anthropic_stream_core( } } } - Ok(StreamEvent::ContentBlockStop { index }) => { + StreamEvent::ContentBlockStop { index } => { // Removing the block releases the borrow before any // yield. if let Some(state) = blocks.remove(&index) { @@ -1080,7 +1132,7 @@ pub(crate) async fn anthropic_stream_core( } } } - Ok(StreamEvent::MessageDelta { delta, usage }) => { + StreamEvent::MessageDelta { delta, usage } => { if let Some(reason) = delta.stop_reason { final_finish_reason = Some(parse_stop_reason(&reason)); } @@ -1101,19 +1153,19 @@ pub(crate) async fn anthropic_stream_core( }; } } - Ok(StreamEvent::MessageStop) => break, - Ok(StreamEvent::Error { error }) => { + StreamEvent::MessageStop => break, + StreamEvent::Error { error } => { // Surface Anthropic in-stream errors (e.g. // overloaded_error) as a `StreamPart::Error` and // stop the stream, mirroring the TS "forward error // chunks" / "forward overloaded error" behaviour. yield Ok(StreamPart::Error { - error: AiMuxError::ApiCall(ApiCallError { - provider_code: error.error_type, - message: error.message, - response_body: Some(sse_event.data.clone()), - ..Default::default() - }), + error: anthropic_stream_error( + &error, + &stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), }); // Finish is the final-chunk contract: it must // still terminate the stream after an in-stream @@ -1122,15 +1174,15 @@ pub(crate) async fn anthropic_stream_core( stream_errored = true; break; } - Ok(_) | Err(_) => {} + _ => {} } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - stream_errored = true; - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs index 038f5c2f..474f21c6 100644 --- a/aimux-providers/src/anthropic_aws/mod.rs +++ b/aimux-providers/src/anthropic_aws/mod.rs @@ -46,8 +46,7 @@ pub struct AnthropicAwsProviderConfig { pub workspace_id: Option, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`（max_retries=2）。 - /// 取代之前硬编码的 `RetryConfig::default()`,让 per-call `max_retries` 生效。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -248,24 +247,23 @@ impl Provider for AnthropicAwsProvider { } } - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: headers.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + crate::anthropic::anthropic_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { @@ -278,7 +276,7 @@ impl Provider for AnthropicAwsProvider { #[serde(default)] display_name: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .data .into_iter() diff --git a/aimux-providers/src/anthropic_aws/model.rs b/aimux-providers/src/anthropic_aws/model.rs index 76e91edb..ec2d4d43 100644 --- a/aimux-providers/src/anthropic_aws/model.rs +++ b/aimux-providers/src/anthropic_aws/model.rs @@ -33,7 +33,7 @@ pub struct AnthropicAwsConfig { pub workspace_id: Option, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`（max_retries=2）。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -129,6 +129,10 @@ impl LanguageModel for AnthropicAwsModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; // M2b: record identity + credential source + endpoint config. Only the @@ -160,19 +164,13 @@ impl LanguageModel for AnthropicAwsModel { let req = build_request_body_with_warnings(&self.model_id, options, false)?; let endpoint = self.endpoint(); let build_headers = self.make_header_builder(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); anthropic_generate_core( &endpoint, - retry_config, req.body, req.warnings, build_headers, BodyEncoding::Bytes, options.abort_signal.clone(), - options.timeout.map(Into::into), options.recording_context.clone(), &ToolNameMapping::new(options.tools.as_deref()), ) @@ -183,19 +181,13 @@ impl LanguageModel for AnthropicAwsModel { let req = build_request_body_with_warnings(&self.model_id, options, true)?; let endpoint = self.endpoint(); let build_headers = self.make_header_builder(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); anthropic_stream_core( &endpoint, - retry_config, req.body, req.warnings, build_headers, BodyEncoding::Bytes, options.abort_signal.clone(), - options.timeout.map(Into::into), options.recording_context.clone(), ToolNameMapping::new(options.tools.as_deref()), ) diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs index a0b88fc0..2e812147 100644 --- a/aimux-providers/src/assemblyai.rs +++ b/aimux-providers/src/assemblyai.rs @@ -16,17 +16,47 @@ use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; use aimux_core::error::ApiCallError; +use aimux_core::retry; use aimux_core::shared::{SharedProviderMetadata, Warning}; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, + HttpBody, HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash, }; +/// AssemblyAI errors: `{"error": "..."}` where `error` is a plain string +/// (e.g. `"Authentication error, API token missing/invalid."`). No machine +/// code is documented. +fn assemblyai_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { + let error = data.get("error"); + let message = error + .and_then(Value::as_str) + .or_else(|| { + error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + }) + .or_else(|| data.get("message").and_then(Value::as_str)) + .unwrap_or("AssemblyAI request failed") + .to_string(); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: error + .and_then(|value| value.get("code")) + .and_then(|value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }), + } +} + +fn assemblyai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(assemblyai_error_parts) +} + // ── Config ────────────────────────────────────────────────────────────────── #[derive(Debug, Clone)] @@ -207,23 +237,22 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Step 1: Upload audio. - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.upload_url(), headers: header_list.clone(), - body: HttpBody::Bytes(audio_bytes, "application/octet-stream".to_string()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(audio_bytes, "application/octet-stream".to_string()), + aimux_provider_utils::create_json_response_handler(), + assemblyai_failed_response_handler(), ) .await?; - let upload: AssemblyAIUploadResponse = serde_json::from_slice(&resp.body)?; + let upload: AssemblyAIUploadResponse = resp.value; // Step 2: Submit transcript request. let mut body = Map::new(); @@ -254,23 +283,27 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { } } - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.transcript_url(), headers: header_list.clone(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + assemblyai_failed_response_handler(), ) .await?; - let submit: AssemblyAISubmitResponse = serde_json::from_slice(&resp.body)?; + let submit: AssemblyAISubmitResponse = resp.value; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Step 3: Poll for completion. let mut raw_body: Value; @@ -282,26 +315,30 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { ) .await?; - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: self.transcript_status_url(&submit.id), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - - response_headers = resp.headers; - - raw_body = serde_json::from_slice(&resp.body)?; - let parsed: AssemblyAITranscriptResponse = serde_json::from_value(raw_body.clone())?; + let poll_url = self.transcript_status_url(&submit.id); + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list.clone(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::< + AssemblyAITranscriptResponse, + >(), + assemblyai_failed_response_handler(), + ) + }) + .await?; + + response_headers = resp.response_headers.unwrap_or_default(); + + raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; if parsed.status == "completed" { // Build segments from words (timestamps are in milliseconds). @@ -408,16 +445,19 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { } if parsed.status == "error" { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some("error".to_string()), - message: format!( - "Transcription failed: {}", - parsed.error.unwrap_or_else(|| "Unknown error".to_string()) - ), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + response_body: Some(raw_body.to_string()), + ..ApiCallError::new( + format!( + "Transcription failed: {}", + parsed.error.unwrap_or_else(|| "Unknown error".to_string()) + ), + poll_url, + serde_json::json!({}), + ) + }))); } } } diff --git a/aimux-providers/src/aws_polly.rs b/aimux-providers/src/aws_polly.rs index 4ebd3ba2..349ca003 100644 --- a/aimux-providers/src/aws_polly.rs +++ b/aimux-providers/src/aws_polly.rs @@ -27,8 +27,7 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::response::{ErrorStructure, error_for_status}; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::{HttpBody, HttpRequest}; use crate::bedrock::sigv4::{AwsCredentials, sign_request}; @@ -61,10 +60,23 @@ const SUPPORTED_OUTPUT_FORMATS: &[&str] = &[ /// AWS error shape: `{"__type": "...", "message": "..."}` (no `error` wrapper). /// Passed to `send` so the http layer extracts the right fields on non-2xx. -const AWS_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["message"], - type_path: &["__type"], -}; +fn aws_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(aws_error_parts) +} + +fn aws_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("message") + .and_then(Value::as_str) + .unwrap_or("Amazon Polly request failed") + .to_string(), + provider_code: data + .get("__type") + .and_then(Value::as_str) + .map(str::to_string), + } +} // ── Config ─────────────────────────────────────────────────────────────────── @@ -266,25 +278,24 @@ impl SpeechModel for AwsPollySpeechModel { &extra_headers, ); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: signed.headers, - body: HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &AWS_ERROR_STRUCTURE, + HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), + aimux_provider_utils::create_binary_response_handler(), + aws_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); @@ -470,48 +481,6 @@ fn parse_polly_provider_options( }) } -// ── Error handling ─────────────────────────────────────────────────────────── - -/// Parse an AWS Polly error response into an [`AiMuxError`]. -/// -/// AWS errors are JSON objects of the form -/// `{"__type": "...", "message": "..."}`. Authentication failures surface as -/// HTTP 401 or 403; both surface as [`AiMuxError::ApiCall`] with the status in `status_code`. -/// -/// Live requests no longer call this — the shared http layer (`send`) maps -/// non-2xx responses itself via [`AWS_ERROR_STRUCTURE`]. Retained (and -/// exercised below) for unit-testing the AWS-specific 401/403 → `Auth` mapping, -/// which the generic handler does not reproduce. -#[allow(dead_code)] -fn parse_polly_error(status: u16, body: &str) -> AiMuxError { - let message = extract_aws_error_message(body); - let provider_code = serde_json::from_str::(body) - .ok() - .and_then(|v| v.get("__type").and_then(|t| t.as_str()).map(String::from)); - error_for_status(status, provider_code, message, None, Some(body.to_string())) -} - -/// Extract a human-readable message from an AWS error JSON body. -/// -/// Only used by [`parse_polly_error`] (which is itself test-only after the -/// migration to the shared http layer), so marked `#[allow(dead_code)]`. -#[allow(dead_code)] -fn extract_aws_error_message(body: &str) -> String { - if let Ok(val) = serde_json::from_str::(body) { - if let Some(msg) = val.get("message").and_then(|v| v.as_str()) { - return msg.to_string(); - } - if let Some(t) = val.get("__type").and_then(|v| v.as_str()) { - return t.to_string(); - } - } - if body.is_empty() { - "HTTP error".to_string() - } else { - body.to_string() - } -} - #[cfg(test)] mod tests { use super::*; @@ -551,16 +520,25 @@ mod tests { } #[test] - fn parse_error_maps_401_to_auth() { - let body = r#"{"__type":"UnrecognizedClientException","message":"The security token included in the request is invalid."}"#; - let err = parse_polly_error(401, body); - assert!(matches!(err, ref e if e.status_code() == Some(401))); + fn extracts_aws_error_fields() { + let data = serde_json::json!({ + "__type": "UnrecognizedClientException", + "message": "The security token included in the request is invalid." + }); + let parts = aws_error_parts(&data); + assert_eq!( + parts.provider_code.as_deref(), + Some("UnrecognizedClientException") + ); + assert_eq!( + parts.message, + "The security token included in the request is invalid." + ); } #[test] - fn parse_error_keeps_the_observed_403() { - let body = r#"{"__type":"AccessDeniedException","message":"User is not authorized."}"#; - let err = parse_polly_error(403, body); - assert_eq!(err.status_code(), Some(403)); + fn missing_aws_message_uses_provider_fallback() { + let parts = aws_error_parts(&serde_json::json!({ "__type": "AccessDeniedException" })); + assert_eq!(parts.message, "Amazon Polly request failed"); } } diff --git a/aimux-providers/src/azure/model.rs b/aimux-providers/src/azure/model.rs index ca6ddc9a..eb802284 100644 --- a/aimux-providers/src/azure/model.rs +++ b/aimux-providers/src/azure/model.rs @@ -87,7 +87,7 @@ pub struct AzureConfig { /// api_key 来源(RFC-0023):`None` = explicit (api-key 或 token provider); /// `Some("env:VAR")` = 环境变量。不存明文。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -324,24 +324,23 @@ impl Provider for AzureProvider { } } - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: headers.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + crate::openai::openai_failed_response_handler(), + ) + }) + .await?; // Azure response: { data: [{ id, model, ... }] } #[derive(serde::Deserialize)] @@ -357,7 +356,7 @@ impl Provider for AzureProvider { #[serde(default, rename = "modelName")] model_name: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .data .into_iter() @@ -479,6 +478,10 @@ impl LanguageModel for AzureModel { &self.deployment } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -502,10 +505,6 @@ impl LanguageModel for AzureModel { async fn do_generate(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()).await?; - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); execute_generate( &self.endpoint(), &headers, @@ -513,17 +512,12 @@ impl LanguageModel for AzureModel { options, "azure", &crate::openai::OpenAICompatProfile::full(), - &retry_config, ) .await } async fn do_stream(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()).await?; - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); execute_stream( &self.endpoint(), &headers, @@ -531,7 +525,6 @@ impl LanguageModel for AzureModel { options, "azure", &crate::openai::OpenAICompatProfile::full(), - &retry_config, ) .await } diff --git a/aimux-providers/src/azure/responses.rs b/aimux-providers/src/azure/responses.rs index 7535e6b3..d70063a4 100644 --- a/aimux-providers/src/azure/responses.rs +++ b/aimux-providers/src/azure/responses.rs @@ -35,12 +35,7 @@ use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateResult, StreamResult}; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed, with_user_agent_suffix, - without_trailing_slash, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::{HttpRequest, with_user_agent_suffix, without_trailing_slash}; use crate::azure::{AzureAuth, AzureConfig}; use crate::openai::responses::convert::build_responses_request_body; @@ -241,6 +236,10 @@ impl LanguageModel for AzureResponsesModel { &self.deployment } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -271,39 +270,36 @@ impl LanguageModel for AzureResponsesModel { Self::apply_file_id_prefixes(&mut body); let provider_key = provider_key().to_string(); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + crate::openai::openai_failed_response_handler(), ) .await?; - let status = resp.status; - let response_headers = resp.headers; - - let data: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let raw_body = resp + .raw_value + .as_ref() + .map(ToString::to_string) + .unwrap_or_default(); + let data = resp.value; build_responses_generate_result( &data, - status, - &String::from_utf8_lossy(&resp.body), + &raw_body, request_result.warnings, provider_key, + endpoint, body, response_headers, ) @@ -329,40 +325,38 @@ impl LanguageModel for AzureResponsesModel { .and_then(serde_json::Value::as_bool) == Some(true); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + crate::openai::openai_failed_response_handler(), ) .await?; - let status = resp.status; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let mut sse_stream = SseStream::new(resp.body); - let first_event = sse_stream.next().await; + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; let stream = build_responses_event_stream( first_event, sse_stream, - status, provider_key, warnings, store_flag, + endpoint, + body.clone(), + response_headers.clone(), )?; Ok(StreamResult { diff --git a/aimux-providers/src/bedrock/embedding.rs b/aimux-providers/src/bedrock/embedding.rs index 45d9809e..fa721fae 100644 --- a/aimux-providers/src/bedrock/embedding.rs +++ b/aimux-providers/src/bedrock/embedding.rs @@ -20,8 +20,7 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::{HttpBody, HttpRequest}; use super::BedrockAuth; use super::model::BedrockConfig; @@ -29,7 +28,7 @@ use super::sigv4::sign_request; /// An Amazon Bedrock embedding model (e.g. `"amazon.titan-embed-text-v2:0"`). /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct BedrockEmbeddingModel { model_id: String, @@ -103,6 +102,10 @@ impl EmbeddingModel for BedrockEmbeddingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_embeddings_per_call(&self) -> Option { if is_cohere_embedding_model(&self.model_id) { Some(96) @@ -196,26 +199,25 @@ impl EmbeddingModel for BedrockEmbeddingModel { let url = self.endpoint(); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url, headers, - body: HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), + aimux_provider_utils::create_json_response_handler(), + super::bedrock_failed_response_handler(), ) .await?; // Capture response headers (needed for token count extraction). - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let raw_value: Value = resp.value; // Extract embeddings based on response format. let embeddings: Vec> = if raw_value.get("embedding").is_some() { diff --git a/aimux-providers/src/bedrock/image.rs b/aimux-providers/src/bedrock/image.rs index e348ed26..7257642d 100644 --- a/aimux-providers/src/bedrock/image.rs +++ b/aimux-providers/src/bedrock/image.rs @@ -17,8 +17,7 @@ use aimux_core::image_model::{ }; use aimux_core::shared::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::{HttpBody, HttpRequest}; use super::sigv4::sign_request; use super::{BedrockAuth, BedrockConfig}; @@ -34,7 +33,7 @@ fn get_max_images_per_call(model_id: &str) -> u32 { /// An Amazon Bedrock image generation model. /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct BedrockImageModel { model_id: String, @@ -104,6 +103,9 @@ impl ImageModel for BedrockImageModel { fn model_id(&self) -> &str { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } fn max_images_per_call(&self) -> Option { Some(get_max_images_per_call(&self.model_id)) } @@ -280,23 +282,23 @@ impl ImageModel for BedrockImageModel { let url = self.endpoint(); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, - url, + url: url.clone(), headers, - body: HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), + aimux_provider_utils::create_json_response_handler::(), + super::bedrock_failed_response_handler(), ) .await?; - let rh = resp.headers; - let rb: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let rb = resp.value; // Handle moderated/blocked requests if let Some(s) = rb.get("status").and_then(|v| v.as_str()) @@ -314,13 +316,16 @@ impl ImageModel for BedrockImageModel { .join(", ") }) .unwrap_or_else(|| "Unknown".into()); - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some(s.to_string()), - message: format!("Amazon Bedrock request was moderated: {reasons_str}"), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + response_body, + ..ApiCallError::new( + format!("Amazon Bedrock request was moderated: {reasons_str}"), + url, + serde_json::json!({}), + ) + }))); } let images: Vec = rb diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs index 157661cb..6c0d5b39 100644 --- a/aimux-providers/src/bedrock/mod.rs +++ b/aimux-providers/src/bedrock/mod.rs @@ -24,6 +24,7 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_provider_utils::{RetryConfig, without_trailing_slash}; +use serde_json::Value; pub use embedding::BedrockEmbeddingModel; pub use image::BedrockImageModel; @@ -31,6 +32,27 @@ pub use model::{BedrockConfig, BedrockModel}; pub use reranking::BedrockRerankingModel; pub use sigv4::AwsCredentials; +pub(crate) fn bedrock_failed_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::create_json_error_response_handler(|data| { + // Converse errors may be a top-level AWS error object or be wrapped + // as `{ "error": { ... } }` by compatible gateways. + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("__type")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} + /// Authentication method for the Bedrock provider. #[derive(Debug, Clone)] pub enum BedrockAuth { @@ -47,7 +69,7 @@ pub struct BedrockProviderConfig { pub auth: BedrockAuth, /// AWS region, used for constructing model ARNs and agent-runtime URLs. pub region: String, - /// 重试配置(M1b)。默认 `RetryConfig::default()`。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, @@ -172,7 +194,7 @@ impl BedrockProviderConfig { /// Amazon Bedrock provider — creates [`BedrockModel`] instances. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct BedrockProvider { config: BedrockProviderConfig, @@ -247,6 +269,7 @@ impl BedrockProvider { self.config.region.clone(), self.config.auth.clone(), ) + .with_retry_config(self.config.retry_config) } } @@ -290,23 +313,23 @@ impl Provider for BedrockProvider { let mut headers = signed.headers; headers.push(("Accept".to_string(), "application/json".to_string())); - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send, - }; - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url, - headers, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: headers.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + bedrock_failed_response_handler(), + ) + }) + .await?; // AWS response: { modelSummaries: [{ modelId, modelName, ... }] } #[derive(serde::Deserialize)] @@ -321,7 +344,7 @@ impl Provider for BedrockProvider { #[serde(default, rename = "modelName")] name: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .summaries .into_iter() diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs index 86937a17..1e17c927 100644 --- a/aimux-providers/src/bedrock/model.rs +++ b/aimux-providers/src/bedrock/model.rs @@ -8,7 +8,8 @@ use std::collections::HashMap; use async_trait::async_trait; -use futures::StreamExt; +use bytes::Bytes; +use futures::{StreamExt, stream::BoxStream}; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; @@ -19,19 +20,67 @@ use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usa use serde_json::json; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send_stream_timed, send_timed, -}; +use aimux_provider_utils::{HttpBody, HttpRequest, RetryConfig}; use super::BedrockAuth; use super::convert::{build_request_body, convert_usage, map_finish_reason}; use super::sigv4::sign_request; use super::types::{BedrockContentBlock, BedrockConverseResponse}; +fn bedrock_event_stream_response_handler() +-> aimux_provider_utils::ResponseHandler>> { + aimux_provider_utils::ResponseHandler::new(|input| async move { + let headers = aimux_provider_utils::extract_response_headers::extract_response_headers( + input.response.headers(), + ); + let output_headers = headers.clone(); + let url = input.url; + let request_body_values = input.request_body_values; + let signal = input.abort_signal; + let stream = input.response.bytes_stream().map(move |result| { + result.map_err(|error| { + AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + response_headers: Some(headers.clone()), + is_retryable: true, + ..aimux_core::ApiCallError::new( + error.to_string(), + url.clone(), + request_body_values.clone(), + ) + })) + }) + }); + let value: BoxStream<'static, Result> = match signal { + Some(signal) => Box::pin(async_stream::stream! { + futures::pin_mut!(stream); + loop { + tokio::select! { + biased; + () = signal.cancelled() => { + yield Err(AiMuxError::from_abort_signal(&signal)); + break; + } + item = stream.next() => match item { + Some(item) => yield item, + None => break, + } + } + } + }), + None => Box::pin(stream), + }; + Ok(aimux_provider_utils::ResponseHandlerOutput { + value, + raw_value: None, + response_headers: Some(output_headers), + }) + }) + .streaming() +} + /// An Amazon Bedrock language model (e.g. `anthropic.claude-3-5-sonnet-20240620-v1:0`). /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct BedrockModel { model_id: String, @@ -43,7 +92,7 @@ pub struct BedrockModel { pub struct BedrockConfig { pub base_url: String, pub auth: BedrockAuth, - /// 重试配置(M1b)。默认 `RetryConfig::default()`。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 /// SigV4 记 access-key 来源;BearerToken 记 bearer-token 来源。不存明文。 @@ -113,6 +162,10 @@ impl LanguageModel for BedrockModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; // M2b: record identity + credential source + region (encoded in @@ -144,32 +197,24 @@ impl LanguageModel for BedrockModel { let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = self.endpoint(false); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url, headers, - body: HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), + aimux_provider_utils::create_json_response_handler(), + super::bedrock_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: BedrockConverseResponse = - serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let data: BedrockConverseResponse = resp.value; // Extract content from response.output.message.content let mut content = Vec::new(); @@ -218,36 +263,29 @@ impl LanguageModel for BedrockModel { let body_str = serde_json::to_string(&body).unwrap_or_default(); let url = self.endpoint(true); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url, headers, - body: HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), + bedrock_event_stream_response_handler(), + super::bedrock_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); // Bedrock converse-stream returns binary AWS event stream format. // We read the full body and decode it, then emit stream parts. // (For true streaming we'd decode incrementally, but the Bedrock event // stream codec requires buffering whole frames anyway.) let mut buf: Vec = Vec::new(); - let mut body_stream = resp.body; + let mut body_stream = resp.value; while let Some(chunk) = body_stream.next().await { match chunk { Ok(bytes) => buf.extend_from_slice(&bytes), @@ -649,8 +687,3 @@ fn extract_content(block: &BedrockContentBlock, content: &mut Vec Self { + self.retry_config = retry_config; + self + } + fn endpoint(&self) -> String { format!("{}/rerank", self.base_url) } @@ -129,6 +135,10 @@ impl RerankingModel for BedrockRerankingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.retry_config + } + async fn do_rerank( &self, options: &RerankingCallOptions, @@ -199,28 +209,26 @@ impl RerankingModel for BedrockRerankingModel { let url = self.endpoint(); let headers = self.build_headers(&body_str, &url, options.headers.as_ref())?; - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url, headers, - body: HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()), + aimux_provider_utils::create_json_response_handler::(), + super::bedrock_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: BedrockRerankingResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let ranking: Vec = data .results diff --git a/aimux-providers/src/bedrock/sigv4.rs b/aimux-providers/src/bedrock/sigv4.rs index e20a825c..e915d95c 100644 --- a/aimux-providers/src/bedrock/sigv4.rs +++ b/aimux-providers/src/bedrock/sigv4.rs @@ -61,7 +61,7 @@ pub fn sign_request( // The signed host must match the wire host: keep the port whenever the // URL carries a non-default one (Url::port() is Some only for non-default // ports). Signing a bare host diverges from the Host header reqwest would - // otherwise send, so gateways and proxies on non-standard ports reject the + // otherwise so gateways and proxies on non-standard ports reject the // signature (observed as 502 when an env proxy forwarded a signed // loopback request whose host lacked the port). let host = match parsed.port() { diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs index e0557c55..dfbe2dcf 100644 --- a/aimux-providers/src/black_forest_labs.rs +++ b/aimux-providers/src/black_forest_labs.rs @@ -17,12 +17,33 @@ use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; +use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; + +fn bfl_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let detail = data.get("detail"); + let message = detail + .and_then(Value::as_str) + .map(str::to_string) + .or_else(|| { + detail + .filter(|value| !value.is_null()) + .map(Value::to_string) + }) + .or_else(|| { + data.get("message") + .and_then(Value::as_str) + .map(str::to_string) + }) + .unwrap_or_else(|| "Unknown Black Forest Labs error".to_string()); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: None, + } + }) +} const DEFAULT_POLL_INTERVAL_MS: u64 = 500; const DEFAULT_POLL_TIMEOUT_MS: u64 = 60000; @@ -287,22 +308,21 @@ impl ImageModel for BlackForestLabsImageModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.submit_endpoint(), headers: header_list.clone(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + bfl_failed_response_handler(), ) .await?; - let submit_body: Value = serde_json::from_slice(&resp.body)?; + let submit_body: Value = resp.value; let poll_url = submit_body .get("polling_url") @@ -316,6 +336,11 @@ impl ImageModel for BlackForestLabsImageModel { .and_then(|v| v.as_str()) .unwrap_or("") .to_string(); + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Poll for result let poll_interval = bfl_opts @@ -343,22 +368,24 @@ impl ImageModel for BlackForestLabsImageModel { let mut result_duration = None; for _ in 0..max_attempts { - let pr = send( - HttpRequest { - method: HttpMethod::Get, - url: poll_url_with_id.to_string(), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - let pv: Value = serde_json::from_slice(&pr.body)?; + let pr = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url_with_id.to_string(), + headers: header_list.clone(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + bfl_failed_response_handler(), + ) + }) + .await?; + let response_body = pr.raw_value.as_ref().map(ToString::to_string); + let pv = pr.value; let poll_status = pv .get("status") @@ -384,13 +411,17 @@ impl ImageModel for BlackForestLabsImageModel { break; } if poll_status == "Error" || poll_status == "Failed" { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(pr.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some(poll_status.to_string()), message: "Black Forest Labs generation failed.".into(), - response_body: Some(String::from_utf8_lossy(&pr.body).into_owned()), - ..Default::default() - })); + response_body, + ..ApiCallError::new( + "Black Forest Labs generation failed.", + poll_url_with_id.to_string(), + serde_json::json!({}), + ) + }))); } sleep_or_abort( std::time::Duration::from_millis(poll_interval), @@ -406,22 +437,23 @@ impl ImageModel for BlackForestLabsImageModel { })?; // Download image - let ir = send( - HttpRequest { - method: HttpMethod::Get, - url: image_url, - headers: vec![], - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - let image_bytes = ir.body.to_vec(); + let ir = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: image_url.clone(), + headers: vec![], + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_binary_response_handler(), + bfl_failed_response_handler(), + ) + }) + .await?; + let image_bytes = ir.value.to_vec(); let download_headers: HashMap = HashMap::new(); // Build provider metadata diff --git a/aimux-providers/src/cartesia.rs b/aimux-providers/src/cartesia.rs index 619baea2..240828c2 100644 --- a/aimux-providers/src/cartesia.rs +++ b/aimux-providers/src/cartesia.rs @@ -29,8 +29,27 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send}; +use aimux_provider_utils::{HttpBody, HttpRequest, load_api_key}; + +fn cartesia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let title = data + .get("title") + .and_then(Value::as_str) + .unwrap_or("Cartesia error"); + let message = data + .get("message") + .and_then(Value::as_str) + .unwrap_or("request failed"); + aimux_provider_utils::ProviderErrorParts { + message: format!("{title}: {message}"), + provider_code: data + .get("error_code") + .and_then(Value::as_str) + .map(str::to_string), + } + }) +} // ── Constants ──────────────────────────────────────────────────────────────── @@ -180,25 +199,24 @@ impl SpeechModel for CartesiaSpeechModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers.into_iter().collect(), - body: HttpBody::Json(Value::Object(body.clone())), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body.clone()), + aimux_provider_utils::create_binary_response_handler(), + cartesia_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); @@ -756,27 +774,25 @@ impl TranscriptionModel for CartesiaTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers.into_iter().collect(), - body: HttpBody::Bytes(body_bytes, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_bytes, content_type), + aimux_provider_utils::create_json_response_handler::(), + cartesia_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); + let response_headers = resp.response_headers.unwrap_or_default(); - let parsed: CartesiaTranscriptionResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; let segments: Vec = parsed .words diff --git a/aimux-providers/src/catalogue.rs b/aimux-providers/src/catalogue.rs index 60831325..4e18bea4 100644 --- a/aimux-providers/src/catalogue.rs +++ b/aimux-providers/src/catalogue.rs @@ -412,35 +412,47 @@ async fn fetch_text(url: &str) -> Result { .timeout(Duration::from_secs(60)) .build() .map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("get_model_specs: build client: {e}"), + AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }) + ..ApiCallError::new( + format!("get_model_specs: build client: {e}"), + url, + serde_json::json!({}), + ) + })) })?; let resp = client.get(url).send().await.map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("get_model_specs: fetch {url}: {e}"), + AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }) + ..ApiCallError::new( + format!("get_model_specs: fetch {url}: {e}"), + url, + serde_json::json!({}), + ) + })) })?; if !resp.status().is_success() { let status = resp.status().as_u16(); - return Err(AiMuxError::ApiCall(ApiCallError { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(status), - message: format!("get_model_specs: fetch {url}"), response_body: resp.text().await.ok().filter(|b| !b.is_empty()), is_retryable: status == 429 || status >= 500, - ..Default::default() - })); + ..ApiCallError::new( + format!("get_model_specs: fetch {url}"), + url, + serde_json::json!({}), + ) + }))); } resp.text().await.map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("get_model_specs: read body: {e}"), + AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - }) + ..ApiCallError::new( + format!("get_model_specs: read body: {e}"), + url, + serde_json::json!({}), + ) + })) }) } diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs index d90c079b..e47efd16 100644 --- a/aimux-providers/src/codex.rs +++ b/aimux-providers/src/codex.rs @@ -30,18 +30,13 @@ use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; use aimux_core::error::AiMuxError; -use aimux_core::error::ApiCallError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::provider::Provider; use aimux_core::result::{GenerateResult, StreamResult}; use aimux_core::types::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send, send_stream_timed, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::{HttpRequest, RetryConfig}; use crate::openai::responses::responses_convert::{ build_header_list, build_responses_event_stream, build_responses_generate_result, @@ -149,7 +144,7 @@ impl CodexConfig { self } - /// Override the retry configuration. + /// Override the retry settings used by Core model operations. #[must_use] pub fn with_retry_config(mut self, config: RetryConfig) -> Self { self.openai = self.openai.with_retry_config(config); @@ -214,7 +209,7 @@ impl Provider for CodexProvider { let runtime = crate::openai::model::execute_list_models( &config.base_url, &headers, - &config.retry_config, + config.retry_config, ) .await?; Ok(runtime) @@ -297,39 +292,31 @@ impl CodexModel { let headers = self.inner().build_headers(opts.headers.as_ref()); let (body, warnings) = self.body(&opts, true)?; - let resp = send_stream_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: opts.abort_signal.clone(), call_id: opts.call_id.clone(), recording_context: opts.recording_context.clone(), }, - self.config.openai.retry_config, - &DEFAULT_ERROR_STRUCTURE, - opts.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + crate::openai::openai_failed_response_handler(), ) .await .map_err(|e| self.map_subscription_401(e))?; - let response_headers = resp.headers; - let mut sse = SseStream::new(resp.body); + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse = resp.value; // (parsed response object, raw event payload) of the terminal event. let mut completed: Option<(Value, String)> = None; let mut failure: Option = None; while let Some(event) = sse.next().await { match event { - Ok(ev) => { - let data: Value = match serde_json::from_str(&ev.data) { - Ok(v) => v, - Err(e) => { - failure = Some(e.into()); - break; - } - }; + Ok(data) => { // The endpoint emits `data:`-only SSE frames — the event // type lives in the JSON payload, not the SSE event field. let etype = data.get("type").and_then(|t| t.as_str()).unwrap_or(""); @@ -338,8 +325,9 @@ impl CodexModel { // The completed event nests the full response // object under `response` — the generate-result // parser expects the response object itself. + let raw = data.to_string(); let obj = data.get("response").cloned().unwrap_or(data); - completed = Some((obj, ev.data)); + completed = Some((obj, raw)); break; } "response.failed" | "error" => { @@ -350,25 +338,34 @@ impl CodexModel { .and_then(|e| e.get("message")) .and_then(|m| m.as_str()) .unwrap_or("subscription response failed"); - failure = Some(AiMuxError::ApiCall(ApiCallError { - // Provider-declared in-band failure: keep the - // observed 2xx envelope status (§2.2). - status_code: Some(resp.status), - provider_code: err_obj + // Provider-declared in-band failure: keep the + // observed 2xx envelope status (§2.2). The shared + // helper redacts the raw request context. + failure = Some(aimux_provider_utils::stream_error_api_call( + message, + err_obj .and_then(|e| e.get("code").or_else(|| e.get("type"))) .and_then(|c| c.as_str()) .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(ev.data.clone()), - ..Default::default() - })); + Some(200), + &data, + endpoint.clone(), + body.clone(), + response_headers.clone(), + )); break; } _ => {} } } + Err(e) if e.is_recoverable_stream_error() => { + // A malformed event does not invalidate later SSE frames. + // Retain the first parse error in case the stream ends + // without a terminal response object. + failure.get_or_insert(e); + } Err(e) => { - failure = Some(AiMuxError::InvalidResponseData(e.to_string())); + failure = Some(e); break; } } @@ -377,10 +374,10 @@ impl CodexModel { match completed { Some((data, raw)) => build_responses_generate_result( &data, - resp.status, &raw, warnings, "openai".to_string(), + endpoint, body, response_headers, ), @@ -398,34 +395,37 @@ impl CodexModel { let headers = self.inner().build_headers(opts.headers.as_ref()); let (body, warnings) = self.body(&opts, true)?; - let resp = send_stream_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: opts.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.openai.retry_config, - &DEFAULT_ERROR_STRUCTURE, - opts.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + crate::openai::openai_failed_response_handler(), ) .await .map_err(|e| self.map_subscription_401(e))?; - let status = resp.status; - let response_headers = resp.headers; - let mut sse_stream = SseStream::new(resp.body); - let first_event = sse_stream.next().await; + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; let stream = build_responses_event_stream( first_event, sse_stream, - status, "openai".to_string(), warnings, false, + endpoint, + body.clone(), + response_headers.clone(), )?; Ok(StreamResult { @@ -446,6 +446,10 @@ impl LanguageModel for CodexModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.openai.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -530,25 +534,21 @@ pub async fn codex_refresh_at( "client_id": client_id, }); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: token_url.to_string(), headers: vec![("content-type".to_string(), "application/json".to_string())], - body: HttpBody::Json(body), abort_signal: None, call_id: None, recording_context: None, }, - RetryConfig { - max_retries: 0, - ..RetryConfig::default() - }, - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + crate::openai::openai_failed_response_handler(), ) .await?; - let data: Value = serde_json::from_slice(&resp.body)?; + let data: Value = resp.value; let access_token = data .get("access_token") .and_then(|v| v.as_str()) diff --git a/aimux-providers/src/cohere/embedding.rs b/aimux-providers/src/cohere/embedding.rs index 26e3dfc4..d33be7ac 100644 --- a/aimux-providers/src/cohere/embedding.rs +++ b/aimux-providers/src/cohere/embedding.rs @@ -16,8 +16,7 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::CohereConfig; @@ -62,6 +61,10 @@ impl EmbeddingModel for CohereEmbeddingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_embeddings_per_call(&self) -> Option { Some(96) } @@ -104,25 +107,24 @@ impl EmbeddingModel for CohereEmbeddingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + super::cohere_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body)?; + let raw_value: Value = resp.value; // Extract embeddings: response.embeddings.float let embeddings: Vec> = raw_value diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs index e4e1070c..73480d78 100644 --- a/aimux-providers/src/cohere/mod.rs +++ b/aimux-providers/src/cohere/mod.rs @@ -17,6 +17,21 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; +use serde_json::Value; + +pub(crate) fn cohere_failed_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: None, + } + }) +} /// Configuration for the Cohere provider. #[derive(Debug, Clone)] @@ -25,7 +40,7 @@ pub struct CohereConfig { pub base_url: String, /// api_key 来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -135,24 +150,23 @@ impl Provider for CohereProvider { ), ("Content-Type".to_string(), "application/json".to_string()), ]; - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: headers.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + cohere_failed_response_handler(), + ) + }) + .await?; // Cohere v2: { models: [{ name, endpoints, ... }] } #[derive(serde::Deserialize)] struct Resp { @@ -165,7 +179,7 @@ impl Provider for CohereProvider { #[serde(default)] endpoints: Option>, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .models .into_iter() diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs index 509a0ee9..fdbda3f8 100644 --- a/aimux-providers/src/cohere/model.rs +++ b/aimux-providers/src/cohere/model.rs @@ -17,20 +17,12 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::CohereConfig; use super::convert::{build_request_body, parse_finish_reason}; use super::types::{ChatResponse, StreamEvent, TokenPair}; -/// Cohere error structure: `{ "message": "..." }` (flat, no `error` wrapper). -const COHERE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["message"], - type_path: &[], -}; - /// A Cohere language model. pub struct CohereModel { model_id: String, @@ -106,6 +98,10 @@ impl LanguageModel for CohereModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -126,33 +122,26 @@ impl LanguageModel for CohereModel { let body_result = build_request_body(&self.model_id, options, false); let body = body_result.body.clone(); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &COHERE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler(), + super::cohere_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let data: ChatResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let data: ChatResponse = resp.value; // Build content array. let mut content = Vec::new(); @@ -272,35 +261,46 @@ impl LanguageModel for CohereModel { let body_result = build_request_body(&self.model_id, options, true); let body = body_result.body.clone(); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &COHERE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + super::cohere_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let sse_stream = SseStream::new(resp.body); + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + match first_event.as_ref() { + Some(Ok(event)) if event.event_type == "error" => { + return Err(cohere_stream_error( + event, + &endpoint, + &body, + &response_headers, + )); + } + _ => {} + } let stream_warnings = body_result.warnings; + let stream_body = body.clone(); + let stream_response_headers = response_headers.clone(); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: stream_warnings }); @@ -310,7 +310,9 @@ impl LanguageModel for CohereModel { let mut is_reasoning = false; let mut stream_errored = false; - let mut sse_iter = Box::pin(sse_stream); + let mut sse_iter = Box::pin( + futures::stream::iter(first_event.into_iter()).chain(sse_stream), + ); while let Some(event) = sse_iter.next().await { if stream_errored { @@ -318,20 +320,7 @@ impl LanguageModel for CohereModel { } match event { - Ok(sse_event) => { - // Parse the JSON payload. - let parsed: StreamEvent = match serde_json::from_str(&sse_event.data) { - Ok(e) => e, - Err(e) => { - // Unparsable chunk — emit Error. - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - stream_errored = true; - break; - } - }; - + Ok(parsed) => { match parsed.event_type.as_str() { "message-start" => { yield Ok(StreamPart::ResponseMetadata { @@ -532,17 +521,30 @@ impl LanguageModel for CohereModel { } } + "error" => { + yield Ok(StreamPart::Error { + error: cohere_stream_error( + &parsed, + &endpoint, + &stream_body, + &stream_response_headers, + ), + }); + stream_errored = true; + break; + } + // citation-start, citation-end, and any unknown // event types are silently consumed. _ => {} } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - stream_errored = true; - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } @@ -576,6 +578,36 @@ impl LanguageModel for CohereModel { } } +fn cohere_stream_error( + event: &StreamEvent, + url: &str, + request_body: &Value, + response_headers: &std::collections::HashMap, +) -> AiMuxError { + let status = event.status_code; + let provider_code = event.code.as_ref().and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: status, + provider_code, + response_body: serde_json::to_string(event).ok(), + response_headers: Some(response_headers.clone()), + data: serde_json::to_value(event).ok(), + is_retryable: status.is_some_and(aimux_core::error::is_retryable_status), + ..aimux_core::ApiCallError::new( + event + .message + .clone() + .unwrap_or_else(|| "Cohere stream error".to_string()), + url, + aimux_provider_utils::logging::redact_error_context(request_body.clone()), + ) + })) +} + #[cfg(test)] mod tests { use super::*; diff --git a/aimux-providers/src/cohere/reranking.rs b/aimux-providers/src/cohere/reranking.rs index 35c91c66..d110f699 100644 --- a/aimux-providers/src/cohere/reranking.rs +++ b/aimux-providers/src/cohere/reranking.rs @@ -16,8 +16,7 @@ use aimux_core::reranking_model::{ }; use aimux_core::types::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::CohereConfig; @@ -103,6 +102,10 @@ impl RerankingModel for CohereRerankingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + async fn do_rerank( &self, options: &RerankingCallOptions, @@ -147,28 +150,26 @@ impl RerankingModel for CohereRerankingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + super::cohere_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: CohereRerankingResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let ranking: Vec = data .results diff --git a/aimux-providers/src/cohere/types.rs b/aimux-providers/src/cohere/types.rs index 1dd4cfea..1ade2918 100644 --- a/aimux-providers/src/cohere/types.rs +++ b/aimux-providers/src/cohere/types.rs @@ -73,7 +73,7 @@ pub struct TokenPair { // always has a `type` field matching the event name. We parse as a generic // `Value` and dispatch on the `type` field in the model code. -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Serialize)] pub struct StreamEvent { #[serde(rename = "type")] pub event_type: String, @@ -83,9 +83,15 @@ pub struct StreamEvent { pub id: Option, #[serde(default)] pub delta: Option, + #[serde(default)] + pub message: Option, + #[serde(default)] + pub status_code: Option, + #[serde(default)] + pub code: Option, } -#[derive(Debug, Deserialize, Default)] +#[derive(Debug, Deserialize, Serialize, Default)] pub struct StreamDelta { #[serde(default)] pub message: Option, @@ -95,7 +101,7 @@ pub struct StreamDelta { pub usage: Option, } -#[derive(Debug, Deserialize, Default)] +#[derive(Debug, Deserialize, Serialize, Default)] pub struct StreamMessage { /// Content can be {type:"text",text:""} or {type:"thinking",thinking:""}. #[serde(default)] @@ -110,7 +116,7 @@ pub struct StreamMessage { pub role: Option, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Serialize)] pub struct StreamUsage { pub tokens: TokenPair, } diff --git a/aimux-providers/src/dataforseo.rs b/aimux-providers/src/dataforseo.rs index ebc6e6aa..e6b39136 100644 --- a/aimux-providers/src/dataforseo.rs +++ b/aimux-providers/src/dataforseo.rs @@ -22,10 +22,7 @@ use aimux_core::search_model::{ }; use aimux_core::shared::SharedHeaders; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, without_trailing_slash}; /// Provider canonical name. const PROVIDER_NAME: &str = "dataforseo"; @@ -40,10 +37,22 @@ const DEFAULT_DEPTH: u32 = 10; const MAX_CREDITS: u32 = 1; /// DataForSEO error response structure: `{ status_code, status_message }`. -const DATAFORSEO_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["status_message"], - type_path: &["status_code"], -}; +fn dataforseo_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("status_message") + .and_then(Value::as_str) + .unwrap_or("DataForSEO request failed") + .to_string(), + provider_code: data.get("status_code").and_then(|value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }), + } + }) +} // ── Config ─────────────────────────────────────────────────────────────────── @@ -262,28 +271,26 @@ impl SearchModel for DataforseoSearchModel { .into_iter() .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DATAFORSEO_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler::(), + dataforseo_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: DataforseoResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; // Flatten tasks[].result[].organic[] preserving provider order. let results: Vec = data diff --git a/aimux-providers/src/deepgram.rs b/aimux-providers/src/deepgram.rs index c86b9cf0..8666f084 100644 --- a/aimux-providers/src/deepgram.rs +++ b/aimux-providers/src/deepgram.rs @@ -22,10 +22,40 @@ use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpBody, HttpRequest, load_api_key, without_trailing_slash}; + +/// Deepgram errors: `{"err_code": "...", "err_msg": "...", "request_id": ...}` +/// on most endpoints; some return `{"category": "...", "message": "...", +/// "details": ...}` instead (https://developers.deepgram.com/docs/errors). +fn deepgram_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { + let message = data + .get("err_msg") + .and_then(Value::as_str) + .or_else(|| data.get("message").and_then(Value::as_str)) + .or_else(|| { + data.get("error") + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + }) + .or_else(|| data.get("error").and_then(Value::as_str)) + .unwrap_or("Deepgram request failed") + .to_string(); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: data + .get("err_code") + .or_else(|| data.get("category")) + .and_then(|value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }), + } +} + +fn deepgram_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(deepgram_error_parts) +} // ── Config ────────────────────────────────────────────────────────────────── @@ -294,27 +324,25 @@ impl TranscriptionModel for DeepgramTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: headers.into_iter().collect(), - body: HttpBody::Bytes(audio_bytes, options.media_type.clone()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(audio_bytes, options.media_type.clone()), + aimux_provider_utils::create_json_response_handler::(), + deepgram_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); + let response_headers = resp.response_headers.unwrap_or_default(); - let parsed: DeepgramResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; let channel = parsed .results diff --git a/aimux-providers/src/elevenlabs.rs b/aimux-providers/src/elevenlabs.rs index e5e19fc9..f82d115e 100644 --- a/aimux-providers/src/elevenlabs.rs +++ b/aimux-providers/src/elevenlabs.rs @@ -24,8 +24,45 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send}; +use aimux_provider_utils::{HttpBody, HttpRequest, load_api_key}; + +/// ElevenLabs errors: `{"detail": {"status": "invalid_api_key", "message": ...}}` +/// where `detail.status` is the machine code. FastAPI validation errors carry +/// `detail` as a plain string or a `[{loc, msg, type}]` list instead. +fn elevenlabs_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { + let detail = data.get("detail"); + let message = detail + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .or_else(|| detail.and_then(Value::as_str)) + .or_else(|| { + detail + .and_then(Value::as_array) + .and_then(|items| items.first()) + .and_then(|item| item.get("msg")) + .and_then(Value::as_str) + }) + .or_else(|| { + data.get("error") + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + }) + .or_else(|| data.get("error").and_then(Value::as_str)) + .or_else(|| data.get("message").and_then(Value::as_str)) + .unwrap_or("ElevenLabs request failed") + .to_string(); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: detail + .and_then(|value| value.get("status")) + .and_then(Value::as_str) + .map(str::to_owned), + } +} + +fn elevenlabs_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(elevenlabs_error_parts) +} // ── Config ─────────────────────────────────────────────────────────────────── @@ -168,27 +205,26 @@ impl SpeechModel for ElevenLabsSpeechModel { format!("{}?{}", self.endpoint(&voice_id), qs.join("&")) }; - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(Value::Object(body.clone())), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body.clone()), + aimux_provider_utils::create_binary_response_handler(), + elevenlabs_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let audio_bytes = resp.body.to_vec(); + let response_headers = resp.response_headers.unwrap_or_default(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); @@ -640,29 +676,27 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Bytes(body_bytes, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_bytes, content_type), + aimux_provider_utils::create_json_response_handler::(), + elevenlabs_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); - - let parsed: ElevenLabsTranscriptionResponse = serde_json::from_value(raw_body.clone())?; + let response_headers = resp.response_headers.unwrap_or_default(); + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; let segments: Vec = parsed .words diff --git a/aimux-providers/src/exa_ai.rs b/aimux-providers/src/exa_ai.rs index 420480a5..dbcca33e 100644 --- a/aimux-providers/src/exa_ai.rs +++ b/aimux-providers/src/exa_ai.rs @@ -15,10 +15,24 @@ use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn exa_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("Exa request failed") + .to_string(), + provider_code: error + .and_then(|value| value.get("code").or_else(|| value.get("type"))) + .and_then(Value::as_str) + .map(str::to_string), + } + }) +} const MODEL_ID: &str = "exa-search"; @@ -167,24 +181,23 @@ impl SearchModel for ExaAiSearchModel { let body = build_request_body(options); let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + exa_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let parsed: ExaResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value; + let parsed: ExaResponse = resp.value; Ok(SearchResult { results: map_results(parsed.results), @@ -193,7 +206,7 @@ impl SearchModel for ExaAiSearchModel { warnings: Vec::new(), response: Some(SearchResponse { headers: Some(response_headers), - body: Some(serde_json::from_slice(&resp.body).unwrap_or(Value::Null)), + body: response_body, }), }) } diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs index 6fc0e2d2..24b902a4 100644 --- a/aimux-providers/src/fal.rs +++ b/aimux-providers/src/fal.rs @@ -13,16 +13,50 @@ use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::error::AiMuxError; +use aimux_core::retry; use aimux_core::shared::Warning; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; + +/// fal errors are FastAPI-style: `{"detail": "..."}` or +/// `{"detail": [{"loc": [...], "msg": "...", "type": "..."}]}` where `type` +/// is the machine code (https://fal.ai/docs/model-apis/errors). +fn fal_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { + let detail = data.get("detail"); + let first_detail = detail + .and_then(Value::as_array) + .and_then(|items| items.first()); + let message = detail + .and_then(Value::as_str) + .or_else(|| { + first_detail + .and_then(|item| item.get("msg")) + .and_then(Value::as_str) + }) + .or_else(|| { + data.get("error") + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + }) + .or_else(|| data.get("error").and_then(Value::as_str)) + .or_else(|| data.get("message").and_then(Value::as_str)) + .unwrap_or("Fal request failed") + .to_string(); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: first_detail + .and_then(|item| item.get("type")) + .and_then(Value::as_str) + .map(str::to_owned), + } +} + +fn fal_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(fal_error_parts) +} // ── Config ────────────────────────────────────────────────────────────────── @@ -225,81 +259,86 @@ impl TranscriptionModel for FalTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); // Submit job. - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.submit_url(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + fal_failed_response_handler(), ) .await?; - let job: FalJobResponse = serde_json::from_slice(&resp.body)?; + let job: FalJobResponse = resp.value; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Poll for result. let raw_body: Value; + let parsed: FalTranscriptionResponse; let response_headers: HashMap; loop { - // fal returns 400/404 while a queued request is still in progress — - // `send` surfaces those as errors, so catch them and keep polling. - let resp = match send( - HttpRequest { - method: HttpMethod::Get, - url: self.poll_url(&job.request_id), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await - { - Ok(resp) => resp, - Err(AiMuxError::ApiCall(ref d)) if d.status_code == Some(404) => { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - continue; - } - // The queue answers 400 until the job registers; read the status - // from the field rather than sniffing the message text. - Err(AiMuxError::ApiCall(d)) if d.status_code == Some(400) => { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - continue; - } - Err(e) => return Err(e), + // Fal returns 400/404 while a queued request is still registering. + // Normalize that state inside the retry attempt so a preceding 5xx + // cannot wrap the pending response in RetryError. + let resp = retries + .retry(|| { + let request = aimux_provider_utils::get_from_api( + HttpRequest { + url: self.poll_url(&job.request_id), + headers: headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::< + FalTranscriptionResponse, + >(), + fal_failed_response_handler(), + ); + async move { + match request.await { + Ok(response) => Ok(Some(response)), + Err(AiMuxError::ApiCall(detail)) + if matches!(detail.status_code, Some(400 | 404)) => + { + Ok(None) + } + Err(error) => Err(error), + } + } + }) + .await?; + let Some(resp) = resp else { + sleep_or_abort( + std::time::Duration::from_millis(100), + options.abort_signal.as_ref(), + ) + .await?; + continue; }; - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; + response_headers = resp.response_headers.unwrap_or_default(); + raw_body = resp.raw_value.unwrap_or(Value::Null); + parsed = resp.value; break; } - let parsed: FalTranscriptionResponse = serde_json::from_value(raw_body.clone())?; - let segments: Vec = parsed .chunks .as_ref() @@ -517,27 +556,31 @@ impl ImageModel for FalImageModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + fal_failed_response_handler(), ) .await?; - let rh = resp.headers; - let rb: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers.unwrap_or_default(); + let rb: Value = resp.value; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); let target_images: Vec = if let Some(i) = rb.get("images").and_then(|v| v.as_array()) { @@ -551,22 +594,23 @@ impl ImageModel for FalImageModel { let mut downloaded: Vec> = Vec::new(); for img in &target_images { if let Some(url) = img.get("url").and_then(|v| v.as_str()) { - let ir = send( - HttpRequest { - method: HttpMethod::Get, - url: url.to_string(), - headers: vec![], - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - downloaded.push(ir.body.to_vec()); + let ir = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.to_string(), + headers: vec![], + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_binary_response_handler(), + fal_failed_response_handler(), + ) + }) + .await?; + downloaded.push(ir.value.to_vec()); } } @@ -633,7 +677,8 @@ impl ImageModel for FalImageModel { // ════════════════════════════════════════════════════════════════════════════ use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoModel, VideoResponse, VideoResult, + VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoModel, VideoOperationStart, + VideoOperationStatus, VideoResponse, VideoResult, }; /// Fal video generation model — implements `VideoModel`. @@ -641,7 +686,8 @@ use aimux_core::video_model::{ /// Aligned with Vercel AI SDK `FalVideoModel` /// (`reference/ai/packages/fal/src/fal-video-model.ts`). /// -/// Uses the same async queue pattern as the transcription model. +/// Uses the same async queue pattern as the transcription model; the status +/// polling is driven by Core via `do_status`. pub struct FalVideoModel { model_id: String, config: FalConfig, @@ -723,7 +769,10 @@ impl VideoModel for FalVideoModel { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { let warnings: Vec = Vec::new(); let mut body = Map::new(); @@ -749,76 +798,84 @@ impl VideoModel for FalVideoModel { let headers = self.build_headers(options.headers.as_ref()); // Submit. - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.submit_url(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + fal_failed_response_handler(), ) .await?; - let job: FalJobResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let job: FalJobResponse = resp.value; - // Poll. - let raw_body: Value; - let response_headers: HashMap; - loop { - let resp = match send( - HttpRequest { - method: HttpMethod::Get, - url: self.poll_url(&job.request_id), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await - { - Ok(resp) => resp, - Err(AiMuxError::ApiCall(ref d)) if d.status_code == Some(404) => { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - continue; - } - // The queue answers 400 until the job registers; read the status - // from the field rather than sniffing the message text. - Err(AiMuxError::ApiCall(d)) if d.status_code == Some(400) => { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - continue; - } - Err(e) => return Err(e), - }; + Ok(VideoOperationStart { + operation: json!({ "request_id": job.request_id }), + warnings, + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: response_headers, + }, + }) + } - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; - break; - } + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let request_id = operation + .get("request_id") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "fal operation reference is missing request_id".to_string(), + ) + })?; + + let headers = self.build_headers(options.headers.as_ref()); + + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: self.poll_url(request_id), + headers: headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + fal_failed_response_handler(), + ) + .await; + + // Fal returns 400/404 while a queued request is still registering or + // running; that is Pending, not a failure. + let resp = match resp { + Ok(response) => response, + Err(AiMuxError::ApiCall(detail)) if matches!(detail.status_code, Some(400 | 404)) => { + return Ok(VideoOperationStatus::Pending); + } + Err(error) => return Err(error), + }; + + let response_headers = resp.response_headers.unwrap_or_default(); + let raw_body = resp.value; // Extract video URL from response. let videos: Vec = if let Some(video) = raw_body @@ -841,15 +898,15 @@ impl VideoModel for FalVideoModel { )); }; - Ok(VideoResult { + Ok(VideoOperationStatus::Completed(VideoResult { videos, - warnings, + warnings: Vec::new(), provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), headers: Some(response_headers), }, - }) + })) } } diff --git a/aimux-providers/src/firecrawl.rs b/aimux-providers/src/firecrawl.rs index b4ed8211..27f31253 100644 --- a/aimux-providers/src/firecrawl.rs +++ b/aimux-providers/src/firecrawl.rs @@ -15,18 +15,23 @@ use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; const MODEL_ID: &str = "firecrawl-search"; /// Firecrawl-specific error structure: `{ "success": false, "error": "..." }`. -const FIRECRAWL_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error"], - type_path: &[], -}; +fn firecrawl_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("error") + .and_then(Value::as_str) + .unwrap_or("Firecrawl request failed") + .to_string(), + provider_code: None, + } + }) +} /// Configuration for the Firecrawl provider. #[derive(Debug, Clone)] @@ -183,24 +188,23 @@ impl SearchModel for FirecrawlSearchModel { let body = build_request_body(options); let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &FIRECRAWL_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + firecrawl_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let parsed: FirecrawlResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value; + let parsed: FirecrawlResponse = resp.value; Ok(SearchResult { results: map_results(parsed.data.web), @@ -209,7 +213,7 @@ impl SearchModel for FirecrawlSearchModel { warnings: Vec::new(), response: Some(SearchResponse { headers: Some(response_headers), - body: Some(serde_json::from_slice(&resp.body).unwrap_or(Value::Null)), + body: response_body, }), }) } diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs index 646b2375..79cb91a1 100644 --- a/aimux-providers/src/gladia.rs +++ b/aimux-providers/src/gladia.rs @@ -15,17 +15,37 @@ use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::retry; use aimux_core::shared::{SharedProviderMetadata, Warning}; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, MultipartForm, RetryConfig, load_api_key, - media_type_to_extension, send, sleep_or_abort, without_trailing_slash, + HttpBody, HttpRequest, MultipartForm, load_api_key, media_type_to_extension, sleep_or_abort, + without_trailing_slash, }; +fn gladia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("Gladia request failed") + .to_string(), + provider_code: error.and_then(|value| value.get("code")).and_then( + |value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }, + ), + } + }) +} + // ── Config ────────────────────────────────────────────────────────────────── #[derive(Debug, Clone)] @@ -206,26 +226,25 @@ impl TranscriptionModel for GladiaTranscriptionModel { form.file("audio", &filename, &options.media_type, &audio_bytes)?; let (body_bytes, content_type) = form.finish(); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.upload_url(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Bytes(body_bytes, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_bytes, content_type), + aimux_provider_utils::create_json_response_handler(), + gladia_failed_response_handler(), ) .await?; - let upload: GladiaUploadResponse = serde_json::from_slice(&resp.body)?; + let upload: GladiaUploadResponse = resp.value; // Step 2: Initiate transcription. let mut body = Map::new(); @@ -243,26 +262,30 @@ impl TranscriptionModel for GladiaTranscriptionModel { } } - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.initiate_url(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + gladia_failed_response_handler(), ) .await?; - let init: GladiaInitResponse = serde_json::from_slice(&resp.body)?; + let init: GladiaInitResponse = resp.value; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Step 3: Poll for result. let mut raw_body: Value; @@ -274,28 +297,30 @@ impl TranscriptionModel for GladiaTranscriptionModel { ) .await?; - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: init.result_url.clone(), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: init.result_url.clone(), + headers: headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::( + ), + gladia_failed_response_handler(), + ) + }) + .await?; - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; - let parsed: GladiaResultResponse = serde_json::from_value(raw_body.clone())?; + response_headers = resp.response_headers.unwrap_or_default(); + raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; if parsed.status == "done" { let result = parsed.result.ok_or_else(|| { @@ -343,13 +368,16 @@ impl TranscriptionModel for GladiaTranscriptionModel { } if parsed.status == "error" { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some(parsed.status.clone()), - message: "Transcription job failed".to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + response_body: Some(raw_body.to_string()), + ..ApiCallError::new( + "Transcription job failed", + init.result_url.clone(), + serde_json::json!({}), + ) + }))); } } } diff --git a/aimux-providers/src/google/embedding.rs b/aimux-providers/src/google/embedding.rs index 3e7b7615..2fb55b42 100644 --- a/aimux-providers/src/google/embedding.rs +++ b/aimux-providers/src/google/embedding.rs @@ -18,20 +18,13 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::GoogleConfig; -/// Google-specific error structure: `{ "error": { "message": "..." } }`. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// A Google Gemini embedding model (e.g. `"gemini-embedding-001"`). /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct GoogleEmbeddingModel { model_id: String, @@ -76,6 +69,10 @@ impl EmbeddingModel for GoogleEmbeddingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_embeddings_per_call(&self) -> Option { Some(100) } @@ -123,25 +120,24 @@ impl EmbeddingModel for GoogleEmbeddingModel { self.config.base_url, self.model_id ); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + super::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let raw_value: Value = resp.value; // Single embedding: response.embedding.values let embedding = raw_value @@ -206,25 +202,24 @@ impl EmbeddingModel for GoogleEmbeddingModel { self.config.base_url, self.model_id ); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + super::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let raw_value: Value = resp.value; // Batch embeddings: response.embeddings[].values let embeddings: Vec> = raw_value diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs index e4d68a15..0ab8c39e 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -25,17 +25,11 @@ use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData, Uplo use aimux_core::shared::FileBytes; use aimux_core::types::Warning; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send, sleep_or_abort}; +use aimux_provider_utils::{HttpBody, HttpRequest, sleep_or_abort}; use super::GoogleConfig; /// Google-specific error structure: `{ "error": { "message": "..." } }`. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// Google provider-specific file upload options. #[derive(Debug, Clone, Default)] struct GoogleFilesUploadOptions { @@ -115,7 +109,7 @@ struct UploadResponse { /// A Google Files interface for uploading files. /// -/// Aligned with TS `GoogleFiles`. Does **not** hold an HTTP client — `http::send` +/// Aligned with TS `GoogleFiles`. Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers /// uses the process-wide shared `Client` internally (RFC-0009 §4.1). pub struct GoogleFiles { config: GoogleConfig, @@ -195,31 +189,42 @@ impl Files for GoogleFiles { )); init_headers.push(("Content-Type".to_string(), "application/json".to_string())); - let init_resp = send( + let init_resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.init_endpoint(), headers: init_headers, - body: HttpBody::Json(init_body_value), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + init_body_value, + aimux_provider_utils::ResponseHandler::new(|input| async move { + let headers = + aimux_provider_utils::extract_response_headers::extract_response_headers( + input.response.headers(), + ); + Ok(aimux_provider_utils::ResponseHandlerOutput { + value: (), + raw_value: None, + response_headers: Some(headers), + }) + }), + super::google_failed_response_handler(), ) .await .map_err(|e| match e { - AiMuxError::ApiCall(d) => AiMuxError::ApiCall(ApiCallError { + AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError { message: format!("Failed to initiate resumable upload: {}", d.message), - ..d - }), + ..*d + })), e => e, })?; let upload_url = init_resp - .headers + .response_headers + .as_ref() + .expect("header-only response handler always supplies headers") .get("x-goog-upload-url") .cloned() .ok_or_else(|| { @@ -237,33 +242,30 @@ impl Files for GoogleFiles { ), ]; - let upload_resp = send( + let upload_request_url = upload_url.clone(); + let upload_resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: upload_url, headers: upload_headers, - body: HttpBody::Bytes(file_bytes, media_type.clone()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + HttpBody::Bytes(file_bytes, media_type.clone()), + aimux_provider_utils::create_json_response_handler::(), + super::google_failed_response_handler(), ) .await .map_err(|e| match e { - AiMuxError::ApiCall(d) => AiMuxError::ApiCall(ApiCallError { + AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError { message: format!("Failed to upload file data: {}", d.message), - ..d - }), + ..*d + })), e => e, })?; - let upload_result: UploadResponse = - serde_json::from_slice(&upload_resp.body).map_err(AiMuxError::from)?; - - let mut file = upload_result.file; + let mut file = upload_resp.value.file; // Step 3: Poll if file is PROCESSING. let poll_interval_ms = google_options.poll_interval_ms.unwrap_or(2000); @@ -276,9 +278,9 @@ impl Files for GoogleFiles { // Seed evidence from the upload response so a file that is already // FAILED (never polled) still carries the observed status + raw body. - let mut last_poll_body: Option = - Some(String::from_utf8_lossy(&upload_resp.body).into_owned()); - let mut last_poll_status: Option = Some(upload_resp.status); + let mut last_poll_body = upload_resp.raw_value.map(|value| value.to_string()); + let mut last_poll_status: Option = Some(200); + let mut last_poll_url = upload_request_url; while file.state == "PROCESSING" { if start_time.elapsed() > Duration::from_millis(poll_timeout_ms) { return Err(AiMuxError::Timeout(format!( @@ -295,37 +297,39 @@ impl Files for GoogleFiles { let poll_url = format!("{}/{}", self.config.base_url, file.name); - let poll_resp = send( + let poll_resp = aimux_provider_utils::get_from_api( HttpRequest { - method: HttpMethod::Get, - url: poll_url, + url: poll_url.clone(), headers: poll_header_list.clone(), - body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + super::google_failed_response_handler(), ) .await?; - file = serde_json::from_slice::(&poll_resp.body) - .map_err(AiMuxError::from)?; - last_poll_body = Some(String::from_utf8_lossy(&poll_resp.body).into_owned()); - last_poll_status = Some(poll_resp.status); + file = poll_resp.value; + last_poll_body = poll_resp.raw_value.map(|value| value.to_string()); + last_poll_status = Some(200); + last_poll_url = poll_url; } if file.state == "FAILED" { // Provider-declared job failure inside a 2xx envelope: stays ApiCall. - return Err(AiMuxError::ApiCall(ApiCallError { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { status_code: last_poll_status, provider_code: Some("FAILED".to_string()), message: format!("File processing failed for {}", file.name), response_body: last_poll_body, - ..Default::default() - })); + ..ApiCallError::new( + format!("File processing failed for {}", file.name), + last_poll_url, + serde_json::json!({}), + ) + }))); } // Build provider metadata. diff --git a/aimux-providers/src/google/image.rs b/aimux-providers/src/google/image.rs index 27dbf148..102b1b7a 100644 --- a/aimux-providers/src/google/image.rs +++ b/aimux-providers/src/google/image.rs @@ -19,17 +19,11 @@ use aimux_core::image_model::{ }; use aimux_core::shared::{SharedProviderMetadata, Warning}; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::GoogleConfig; /// Google error structure: `{ "error": { "message": "...", "status": "..." } }`. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// Returns `true` if the model ID is a Gemini image model. fn is_gemini_model(model_id: &str) -> bool { model_id.starts_with("gemini-") @@ -44,7 +38,7 @@ pub struct GoogleImageSettings { /// A Google image generation model (Imagen or Gemini). /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct GoogleImageModel { model_id: String, @@ -174,25 +168,24 @@ impl GoogleImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list = build_header_list(&headers); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.predict_endpoint(), headers: header_list, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + super::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let response_body: Value = serde_json::from_slice(&resp.body)?; + let response_body: Value = resp.value; let images = extract_imagen_images(&response_body); let provider_metadata = extract_imagen_metadata(&response_body); @@ -326,25 +319,24 @@ impl GoogleImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list = build_header_list(&headers); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generate_content_endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + super::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let response_body: Value = serde_json::from_slice(&resp.body)?; + let response_body: Value = resp.value; let (images, provider_metadata, usage) = extract_gemini_result(&response_body); @@ -372,6 +364,10 @@ impl ImageModel for GoogleImageModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_images_per_call(&self) -> Option { Some(self.max_images()) } diff --git a/aimux-providers/src/google/mod.rs b/aimux-providers/src/google/mod.rs index e63dab20..9bcef2b9 100644 --- a/aimux-providers/src/google/mod.rs +++ b/aimux-providers/src/google/mod.rs @@ -29,6 +29,53 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; +use serde_json::Value; + +pub(crate) fn google_stream_error( + error: &types::GoogleError, + url: &str, + request_body_values: Value, + response_headers: std::collections::HashMap, +) -> AiMuxError { + let status_code = error + .code + .and_then(|status| u16::try_from(status).ok()) + .filter(|status| (400..=599).contains(status)); + let data = serde_json::json!({ + "error": { + "code": error.code, + "message": error.message, + "status": error.status, + } + }); + aimux_provider_utils::stream_error_api_call( + error.message.clone(), + error.status.clone(), + status_code, + &data, + url, + request_body_values, + response_headers, + ) +} + +pub(crate) fn google_failed_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("status") + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} /// Configuration for the Google Gemini provider. #[derive(Debug, Clone)] @@ -39,8 +86,7 @@ pub struct GoogleConfig { /// (config_snapshot 记为 "explicit");`Some("env:VAR")` = 来自环境变量。 /// 不存明文之外的信息,仅用于回放重建。 pub api_key_source: Option, - /// 重试配置。默认 `RetryConfig::default()`（max_retries=2）。M1b:取代 - /// 之前硬编码的 `RetryConfig::default()`,让 per-call `max_retries` 生效。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -90,7 +136,7 @@ impl GoogleConfig { /// Google Gemini provider — creates `GoogleModel` instances. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct GoogleProvider { config: GoogleConfig, @@ -180,24 +226,23 @@ impl Provider for GoogleProvider { let mut header_list: Vec<(String, String)> = headers.into_iter().collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers: header_list, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: header_list.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + google_failed_response_handler(), + ) + }) + .await?; // Gemini response: { models: [{ name: "models/gemini-...", displayName, ... }] } #[derive(serde::Deserialize)] @@ -211,7 +256,7 @@ impl Provider for GoogleProvider { #[serde(default)] display_name: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .models .into_iter() diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs index f036046f..3a0d31b9 100644 --- a/aimux-providers/src/google/model.rs +++ b/aimux-providers/src/google/model.rs @@ -6,7 +6,7 @@ use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; @@ -16,30 +16,17 @@ use aimux_core::types::{ FinishReason, FinishReasonUnified, ProviderMetadata, ResponseMetadata, Usage, }; -use aimux_provider_utils::response::{ErrorStructure, error_for_status}; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::GoogleConfig; use super::convert::{ build_request_body_with_warnings, convert_usage, extract_sources, parse_finish_reason, }; -use super::types::{ - Candidate, GenerateContentResponse, GoogleErrorEnvelope, GoogleUsageMetadata, StreamChunk, -}; - -/// Google-specific error structure: `{ "error": { "message": "..." } }`. -/// -/// Google uses `code` (HTTP status as a number) and `status` (a string like -/// `INVALID_ARGUMENT`) instead of OpenAI's `type` field. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; +use super::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; /// A Google Gemini language model. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct GoogleModel { model_id: String, @@ -110,6 +97,10 @@ impl LanguageModel for GoogleModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -129,32 +120,24 @@ impl LanguageModel for GoogleModel { async fn do_generate(&self, options: &CallOptions) -> Result { let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generate_endpoint(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &GOOGLE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler(), + super::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: GenerateContentResponse = - serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let data: GenerateContentResponse = resp.value; let candidate = data.candidates.into_iter().next().ok_or_else(|| { AiMuxError::InvalidResponseData("no candidates in response".to_string()) @@ -209,36 +192,45 @@ impl LanguageModel for GoogleModel { async fn do_stream(&self, options: &CallOptions) -> Result { let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let endpoint = self.stream_endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.stream_endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &GOOGLE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + super::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let sse_stream = SseStream::new(resp.body); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(GoogleStreamEvent::Error(error))) = first_event.as_ref() { + return Err(super::google_stream_error( + &error.error, + &endpoint, + body.clone(), + response_headers, + )); + } + let stream_error_url = endpoint.clone(); + let stream_request_body = body.clone(); + let stream_error_headers = response_headers.clone(); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: tool_warnings }); - let mut sse_stream = sse_stream; + let mut sse_stream = futures::stream::iter(first_event.into_iter()).chain(sse_stream); let mut text_id: Option = None; let mut reasoning_id: Option = None; let mut block_counter = 0usize; @@ -272,32 +264,7 @@ impl LanguageModel for GoogleModel { break; } match event { - Ok(sse_event) => { - // Parse as generic Value first to surface errors. - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - stream_errored = true; - break; - } - }; - - // Try to parse as a StreamChunk. The Google stream - // is a sequence of independent JSON objects (one per - // SSE event), not OpenAI-style delta fragments. - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - stream_errored = true; - break; - } - }; + Ok(GoogleStreamEvent::Chunk(chunk)) => { // Emit ResponseMetadata from the first chunk that has // a responseId (matches TS behaviour). @@ -630,13 +597,25 @@ impl LanguageModel for GoogleModel { Some(parse_finish_reason(reason, has_tool_calls)); } } - Err(e) => { + Ok(GoogleStreamEvent::Error(error)) => { yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), + error: super::google_stream_error( + &error.error, + &stream_error_url, + stream_request_body.clone(), + stream_error_headers.clone(), + ), }); stream_errored = true; break; } + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } + } } } @@ -918,42 +897,6 @@ fn set_provider_metadata(item: &mut GenerateContent, meta: ProviderMetadata) { } } -/// Parse a Google error JSON envelope into an `AiMuxError`. Used for -/// mid-stream error objects (the Gemini SSE stream can carry `{ "error": … }` -/// events). Kept for parity with the OpenAI provider even though the common -/// error path goes through `parse_provider_error`. -#[allow(dead_code)] -fn parse_google_error_body(body: &str) -> AiMuxError { - if let Ok(env) = serde_json::from_str::(body) { - let msg = env.error.message; - let provider_code = env.error.status; - return match env.error.code { - Some(code) => error_for_status( - code as u16, - provider_code, - msg, - None, - Some(body.to_string()), - ), - None => AiMuxError::ApiCall(ApiCallError { - provider_code, - message: msg, - response_body: Some(body.to_string()), - ..Default::default() - }), - }; - } - AiMuxError::ApiCall(ApiCallError { - message: body.to_string(), - response_body: Some(body.to_string()), - ..Default::default() - }) -} - -// Suppress unused warning for GoogleUsageMetadata (re-exported via convert). -#[allow(unused_imports)] -use GoogleUsageMetadata as _GoogleUsageMetadata; - #[cfg(test)] mod tests { use super::*; diff --git a/aimux-providers/src/google/types.rs b/aimux-providers/src/google/types.rs index 41f6f684..239a6ab2 100644 --- a/aimux-providers/src/google/types.rs +++ b/aimux-providers/src/google/types.rs @@ -117,6 +117,15 @@ pub struct StreamChunk { pub model_version: Option, } +/// One Google SSE event. Error envelopes must be attempted before the very +/// permissive chunk shape, whose fields are all optional. +#[derive(Debug, Deserialize)] +#[serde(untagged)] +pub enum GoogleStreamEvent { + Error(GoogleErrorEnvelope), + Chunk(Box), +} + // ── Provider error response ────────────────────────────────────────────────── /// Google error envelope: `{ "error": { "code": 400, "message": "...", "status": "INVALID_ARGUMENT" } }`. diff --git a/aimux-providers/src/google/video.rs b/aimux-providers/src/google/video.rs index 7b909dc5..2e40c7a5 100644 --- a/aimux-providers/src/google/video.rs +++ b/aimux-providers/src/google/video.rs @@ -5,7 +5,7 @@ //! //! Uses Google's Long Running Operations API: //! 1. POST `{base_url}/models/{model}:predictLongRunning` → returns operation name -//! 2. GET `{base_url}/{operation_name}` to poll until `done: true` +//! 2. GET `{base_url}/{operation_name}` — polled by Core via `do_status` until `done: true` //! 3. Return video URL(s) from operation result use std::collections::HashMap; @@ -14,24 +14,18 @@ use async_trait::async_trait; use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; -use aimux_core::shared::Warning; use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoModel, VideoResponse, VideoResult, + VideoCallOptions, VideoData, VideoModel, VideoOperationStart, VideoOperationStatus, + VideoResponse, VideoResult, }; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send, sleep_or_abort}; +use aimux_provider_utils::HttpRequest; use super::GoogleConfig; -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// A Google video generation model. /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct GoogleVideoModel { model_id: String, @@ -75,13 +69,17 @@ impl VideoModel for GoogleVideoModel { fn model_id(&self) -> &str { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } fn max_videos_per_call(&self) -> Option { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { - let warnings: Vec = Vec::new(); - + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { let mut instances = vec![json!({"prompt": options.prompt})]; if let Some(ref image) = options.image && let aimux_core::video_model::VideoFile::Url { url, .. } = image @@ -119,23 +117,23 @@ impl VideoModel for GoogleVideoModel { let url = self.predict_url(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: header_list.clone(), - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + super::google_failed_response_handler(), ) .await?; - let predict_response: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let predict_response: Value = resp.value; let operation_name = predict_response .get("name") .and_then(|v| v.as_str()) @@ -146,56 +144,75 @@ impl VideoModel for GoogleVideoModel { })? .to_string(); - // Poll for completion. - let mut raw_body: Value; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - let poll_url = self.operation_url(&operation_name); - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: poll_url, - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, - ) - .await?; - - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; - // Check the in-band error first: a terminal response may carry both - // done:true and an error object (provider-declared failure). - if let Some(err) = raw_body.get("error") { - let msg = err - .get("message") + Ok(VideoOperationStart { + operation: json!({ "operation_name": operation_name }), + warnings: Vec::new(), + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: response_headers, + }, + }) + } + + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let operation_name = operation + .get("operation_name") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "google operation reference is missing operation_name".to_string(), + ) + })?; + + let headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + let poll_url = self.operation_url(operation_name); + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list, + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + super::google_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let raw_body: Value = resp.value; + // Check the in-band error first: a terminal response may carry both + // done:true and an error object (provider-declared failure). + if let Some(err) = raw_body.get("error") { + let msg = err + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("Unknown error"); + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: err + .get("status") .and_then(|v| v.as_str()) - .unwrap_or("Unknown error"); - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: err - .get("status") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: msg.to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } - if raw_body.get("done").and_then(serde_json::Value::as_bool) == Some(true) { - break; - } + .map(std::string::ToString::to_string), + response_body, + ..ApiCallError::new(msg, poll_url, serde_json::json!({})) + }))); + } + if raw_body.get("done").and_then(serde_json::Value::as_bool) != Some(true) { + return Ok(VideoOperationStatus::Pending); } // Extract videos from response. @@ -224,15 +241,15 @@ impl VideoModel for GoogleVideoModel { )); } - Ok(VideoResult { + Ok(VideoOperationStatus::Completed(VideoResult { videos, - warnings, + warnings: Vec::new(), provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), headers: Some(response_headers), }, - }) + })) } } diff --git a/aimux-providers/src/google_pse.rs b/aimux-providers/src/google_pse.rs index 2759e191..58e9fce3 100644 --- a/aimux-providers/src/google_pse.rs +++ b/aimux-providers/src/google_pse.rs @@ -15,24 +15,36 @@ use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; /// Fixed model ID for the Google PSE search model. const MODEL_ID: &str = "google-pse-search"; /// Google PSE error response structure: `{ "error": { "code", "message" } }`. -const GOOGLE_PSE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "code"], -}; +fn google_pse_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("Google PSE request failed") + .to_string(), + provider_code: error.and_then(|value| value.get("code")).and_then( + |value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }, + ), + } + }) +} /// Configuration for the Google PSE provider. #[derive(Debug, Clone)] @@ -214,10 +226,7 @@ impl SearchModel for GooglePseSearchModel { .unwrap_or_default(); let mut url = url::Url::parse(&self.endpoint()).map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("invalid google_pse endpoint: {e}"), - ..Default::default() - }) + AiMuxError::InvalidArgument(format!("invalid google_pse endpoint: {e}")) })?; { let mut qp = url.query_pairs_mut(); @@ -229,27 +238,24 @@ impl SearchModel for GooglePseSearchModel { } } - let resp = send( + let resp = aimux_provider_utils::get_from_api( HttpRequest { - method: HttpMethod::Get, url: url.to_string(), headers, - body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_PSE_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + google_pse_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: GooglePseResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; Ok(SearchResult { results: map_results(data.items), diff --git a/aimux-providers/src/huggingface.rs b/aimux-providers/src/huggingface.rs index 83d1e054..8cf265a4 100644 --- a/aimux-providers/src/huggingface.rs +++ b/aimux-providers/src/huggingface.rs @@ -123,7 +123,7 @@ impl Provider for HuggingFaceProvider { let runtime = crate::openai::model::execute_list_models( &config.base_url, &headers, - &config.retry_config, + config.retry_config, ) .await?; Ok(runtime) diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs index fa853e03..069525f2 100644 --- a/aimux-providers/src/huggingface/responses.rs +++ b/aimux-providers/src/huggingface/responses.rs @@ -35,15 +35,92 @@ use aimux_core::types::{ FinishReason, FinishReasonUnified, ResponseMetadata, TokenUsage, Usage, Warning, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::HuggingFaceConfig; use crate::openai::responses::responses_convert::build_header_list; const PROVIDER_NAME: &str = "huggingface"; +fn huggingface_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("Hugging Face request failed") + .to_string(), + provider_code: error + .and_then(|value| value.get("code").or_else(|| value.get("type"))) + .and_then(Value::as_str) + .map(str::to_string), + } + }) +} + +fn huggingface_successful_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let url = input.url.clone(); + let request_body_values = input.request_body_values.clone(); + let output = aimux_provider_utils::create_json_response_handler::() + .handle(input) + .await?; + if let Some(error) = output.value.get("error").filter(|value| !value.is_null()) { + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("Hugging Face request failed"); + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + provider_code: error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str) + .map(str::to_string), + response_body: Some(output.value.to_string()), + response_headers: output.response_headers.clone(), + ..ApiCallError::new(message, url, request_body_values) + }))); + } + Ok(output) + }) +} + +fn huggingface_stream_error( + event: &Value, + url: &str, + request_body_values: Value, + response_headers: HashMap, +) -> Option { + let error = event.get("error").filter(|value| !value.is_null())?; + let message = error + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) + .unwrap_or("Hugging Face stream failed"); + let status_code = error + .get("status") + .or_else(|| error.get("code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) + .filter(|status| (400..=599).contains(status)); + let provider_code = error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str) + .map(str::to_owned); + Some(aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + event, + url, + request_body_values, + response_headers, + )) +} + // ───────────────────────────────────────────────────────────────────────────── // Model // ───────────────────────────────────────────────────────────────────────────── @@ -91,6 +168,10 @@ impl LanguageModel for HuggingFaceResponsesModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.openai_config().retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { // M2b: HuggingFace wraps OpenAIConfig — reuse the OpenAI snapshot helper. crate::openai::config_snapshot_from_config( @@ -105,49 +186,26 @@ impl LanguageModel for HuggingFaceResponsesModel { let body = request.body; let headers = self.build_headers(options.headers.as_ref()); - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.0.retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + huggingface_successful_response_handler(), + huggingface_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let response: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - // Check for an error field in the response body. - if let Some(error) = response.get("error") - && !error.is_null() - { - let message = error - .get("message") - .and_then(|v| v.as_str()) - .unwrap_or("Unknown error"); - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: error - .get("code") - .or_else(|| error.get("type")) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } + let response: Value = resp.value; - let content = build_generate_content(&response); + let content = build_generate_content(&response)?; let finish_reason = parse_finish_reason(&response); let usage = convert_usage(response.get("usage")); @@ -187,25 +245,40 @@ impl LanguageModel for HuggingFaceResponsesModel { let body = request.body; let headers = self.build_headers(options.headers.as_ref()); - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.0.retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + huggingface_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let sse_stream = SseStream::new(resp.body); + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(event)) = first_event.as_ref() + && let Some(error) = huggingface_stream_error( + event, + &self.endpoint(), + body.clone(), + response_headers.clone(), + ) + { + return Err(error); + } + let stream_error_url = self.endpoint(); + let stream_request_body = body.clone(); + let stream_response_headers = response_headers.clone(); let warnings = request.warnings.clone(); let stream = async_stream::stream! { @@ -218,30 +291,20 @@ impl LanguageModel for HuggingFaceResponsesModel { }; let mut response_id: Option = None; let mut usage_raw: Option = None; - let mut stream_errored = false; - - let mut event_iter = sse_stream; + let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); while let Some(event) = event_iter.next().await { - if stream_errored { - break; - } - match event { - Ok(sse_event) => { - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { error: e.into() }); - stream_errored = true; - continue; - } - }; - + Ok(parsed) => { + if let Some(error) = huggingface_stream_error( + &parsed, + &stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ) { + yield Ok(StreamPart::Error { error }); + break; + } let chunk_type = parsed .get("type") .and_then(|v| v.as_str()) @@ -484,15 +547,12 @@ impl LanguageModel for HuggingFaceResponsesModel { } } } - Err(e) => { - finish_reason = FinishReason { - unified: FinishReasonUnified::Error, - raw: None, - }; - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - stream_errored = true; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } @@ -1024,7 +1084,7 @@ pub fn prepare_responses_tools( /// Build the `GenerateContent` vector from a Responses API JSON body. /// /// Mirrors the TS `doGenerate` output-array processing. -fn build_generate_content(response: &Value) -> Vec { +fn build_generate_content(response: &Value) -> Result, AiMuxError> { let mut content = Vec::new(); let mut source_id_counter = 0u32; @@ -1169,7 +1229,7 @@ fn build_generate_content(response: &Value) -> Vec { } } - content + Ok(content) } /// Parse the finish reason from a Responses API JSON body. diff --git a/aimux-providers/src/hume.rs b/aimux-providers/src/hume.rs index 6dcb8828..440877fb 100644 --- a/aimux-providers/src/hume.rs +++ b/aimux-providers/src/hume.rs @@ -26,8 +26,40 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send}; +use aimux_provider_utils::{HttpRequest, load_api_key}; + +/// Hume errors carry a top-level `message` with a `code` such as `"E0101"` +/// (https://dev.hume.ai/docs/resources/errors); the nested `{error: +/// {message, code}}` shape matches the AI SDK's Hume error schema. +fn hume_error_parts(data: &Value) -> aimux_provider_utils::ProviderErrorParts { + let error = data.get("error"); + let message = data + .get("message") + .and_then(Value::as_str) + .or_else(|| { + error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + }) + .or_else(|| error.and_then(Value::as_str)) + .unwrap_or("Hume request failed") + .to_string(); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: data + .get("code") + .or_else(|| error.and_then(|value| value.get("code"))) + .and_then(|value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }), + } +} + +fn hume_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(hume_error_parts) +} // ── Config ─────────────────────────────────────────────────────────────────── @@ -160,25 +192,24 @@ impl SpeechModel for HumeSpeechModel { .into_iter() .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(Value::Object(body.clone())), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body.clone()), + aimux_provider_utils::create_binary_response_handler(), + hume_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); diff --git a/aimux-providers/src/jina_ai.rs b/aimux-providers/src/jina_ai.rs index 6cac25ff..bed4ea8b 100644 --- a/aimux-providers/src/jina_ai.rs +++ b/aimux-providers/src/jina_ai.rs @@ -20,10 +20,7 @@ use aimux_core::reranking_model::{ RerankingResult, }; use aimux_core::types::Warning; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; /// Jina AI error response structure: `{ "detail": "...", "code": "..." }`. /// @@ -32,10 +29,18 @@ use aimux_provider_utils::{ /// `ErrorResponse` (`detail` + optional `code`) and the FastAPI /// `HTTPValidationError` (`detail`) shapes are covered, since both surface /// the message under `detail`. -const JINA_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["detail"], - type_path: &["code"], -}; +fn jina_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("detail") + .and_then(Value::as_str) + .unwrap_or("Jina AI request failed") + .to_string(), + provider_code: data.get("code").and_then(Value::as_str).map(str::to_string), + } + }) +} /// Configuration for the Jina AI provider. #[derive(Debug, Clone)] @@ -244,28 +249,26 @@ impl RerankingModel for JinaAiRerankingModel { .into_iter() .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &JINA_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler::(), + jina_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: JinaRerankingResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let model = data.model; let ranking: Vec = data diff --git a/aimux-providers/src/klingai.rs b/aimux-providers/src/klingai.rs index 89511b18..92836902 100644 --- a/aimux-providers/src/klingai.rs +++ b/aimux-providers/src/klingai.rs @@ -5,7 +5,8 @@ //! //! KlingAI uses an async task pattern: //! 1. POST to `/v1/videos/text2video` (or `image2video`) → returns task `id` + `task_id` -//! 2. GET `/v1/videos/text2video/{id}/{task_id}` to poll until `succeeded` +//! 2. GET `/v1/videos/text2video/{id}/{task_id}` — polled by Core via `do_status` +//! until `succeeded` //! 3. Return the video URL from the result use std::collections::HashMap; @@ -15,16 +16,29 @@ use serde::Deserialize; use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; -use aimux_core::shared::Warning; use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoModel, VideoResponse, VideoResult, + VideoCallOptions, VideoData, VideoModel, VideoOperationStart, VideoOperationStatus, + VideoResponse, VideoResult, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn klingai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("message") + .and_then(Value::as_str) + .unwrap_or("Kling AI request failed") + .to_string(), + provider_code: data.get("code").and_then(|value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }), + } + }) +} // ── Config ────────────────────────────────────────────────────────────────── @@ -222,8 +236,10 @@ impl VideoModel for KlingAIVideoModel { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { - let warnings: Vec = Vec::new(); + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { let mode = detect_mode(&self.model_id).to_string(); let model_name = get_api_model_name(&self.model_id, &mode); @@ -274,34 +290,42 @@ impl VideoModel for KlingAIVideoModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit task. - let resp = send( + let submit_url = self.submit_endpoint(&mode); + let request_body = Value::Object(body); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.submit_endpoint(&mode), - headers: header_list.clone(), - body: HttpBody::Json(Value::Object(body)), + url: submit_url.clone(), + headers: header_list, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + request_body.clone(), + aimux_provider_utils::create_json_response_handler::(), + klingai_failed_response_handler(), ) .await?; - let task: KlingAITaskResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let task: KlingAITaskResponse = resp.value; if task.code != 0 { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some(task.code.to_string()), - message: task - .message - .unwrap_or_else(|| format!("KlingAI error code: {}", task.code)), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + response_body, + ..ApiCallError::new( + task.message + .unwrap_or_else(|| format!("KlingAI error code: {}", task.code)), + submit_url, + // In-band failure inside a 2xx envelope: the body is still + // raw here (i2v mode embeds base64 image data), so redact + // before it lands in the public error. + aimux_provider_utils::redact_error_context(request_body), + ) + }))); } let task_data = task.data.ok_or_else(|| { @@ -317,96 +341,121 @@ impl VideoModel for KlingAIVideoModel { ) })?; - // Poll for completion. - let videos: Vec; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: self.poll_endpoint(&mode, &id, &task_id), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; + Ok(VideoOperationStart { + // `mode` and `id` are needed alongside `task_id` to rebuild the + // poll endpoint in `do_status`. + operation: json!({ "mode": mode, "id": id, "task_id": task_id }), + warnings: Vec::new(), + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: response_headers, + }, + }) + } - response_headers = resp.headers; + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let field = |name: &str| { + operation.get(name).and_then(Value::as_str).ok_or_else(|| { + AiMuxError::InvalidArgument(format!( + "klingai operation reference is missing {name}" + )) + }) + }; + let mode = field("mode")?; + let id = field("id")?; + let task_id = field("task_id")?; - let result: KlingAITaskResult = serde_json::from_slice(&resp.body)?; + let headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers.into_iter().collect(); - if result.code != 0 { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: Some(result.code.to_string()), - message: result + let poll_url = self.poll_endpoint(mode, id, task_id); + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list, + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + klingai_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let result: KlingAITaskResult = resp.value; + + if result.code != 0 { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(result.code.to_string()), + response_body: response_body.clone(), + ..ApiCallError::new( + result .message .unwrap_or_else(|| format!("KlingAI error code: {}", result.code)), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } + poll_url.clone(), + serde_json::json!({}), + ) + }))); + } - if let Some(data) = result.data - && let Some(task_status) = data.task_status.as_deref() - { - if task_status == "succeeded" { - videos = data - .task_result - .and_then(|r| r.videos) - .unwrap_or_default() - .into_iter() - .filter_map(|v| { - v.get("url") - .and_then(|u| u.as_str()) - .map(|url| VideoData::Url { - url: url.to_string(), - media_type: "video/mp4".to_string(), - }) - }) - .collect(); - if videos.is_empty() { - return Err(AiMuxError::InvalidResponseData(format!( - "KlingAI task {task_id} succeeded without video URLs" - ))); - } - break; - } - if task_status == "failed" { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: Some("failed".to_string()), - message: "KlingAI video generation failed".to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + if let Some(data) = result.data + && let Some(task_status) = data.task_status.as_deref() + { + if task_status == "succeeded" { + let videos: Vec = data + .task_result + .and_then(|r| r.videos) + .unwrap_or_default() + .into_iter() + .filter_map(|v| { + v.get("url") + .and_then(|u| u.as_str()) + .map(|url| VideoData::Url { + url: url.to_string(), + media_type: "video/mp4".to_string(), + }) + }) + .collect(); + if videos.is_empty() { + return Err(AiMuxError::InvalidResponseData(format!( + "KlingAI task {task_id} succeeded without video URLs" + ))); } + return Ok(VideoOperationStatus::Completed(VideoResult { + videos, + warnings: Vec::new(), + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: Some(response_headers), + }, + })); + } + if task_status == "failed" { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some("failed".to_string()), + response_body, + ..ApiCallError::new( + "KlingAI video generation failed", + poll_url, + serde_json::json!({}), + ) + }))); } } - let timestamp = chrono::Utc::now().to_rfc3339(); - - Ok(VideoResult { - videos, - warnings, - provider_metadata: None, - response: VideoResponse { - timestamp: Some(timestamp), - model_id: Some(self.model_id.clone()), - headers: Some(response_headers), - }, - }) + Ok(VideoOperationStatus::Pending) } } diff --git a/aimux-providers/src/linkup.rs b/aimux-providers/src/linkup.rs index 7f2f2d50..e112ad19 100644 --- a/aimux-providers/src/linkup.rs +++ b/aimux-providers/src/linkup.rs @@ -19,10 +19,25 @@ use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn linkup_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} /// Fixed model ID for the Linkup search model. const MODEL_ID: &str = "linkup-search"; @@ -224,28 +239,26 @@ impl SearchModel for LinkupSearchModel { .into_iter() .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler::(), + linkup_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: LinkupResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; // Prefer `results` (searchResults); fall back to `sources` // (sourcedAnswer), which also carries an `answer`. diff --git a/aimux-providers/src/lmnt.rs b/aimux-providers/src/lmnt.rs index 3d5bd625..1c149abe 100644 --- a/aimux-providers/src/lmnt.rs +++ b/aimux-providers/src/lmnt.rs @@ -24,8 +24,27 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send}; +use aimux_provider_utils::{HttpRequest, load_api_key}; + +fn lmnt_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("LMNT request failed") + .to_string(), + provider_code: error.and_then(|value| value.get("code")).and_then( + |value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }, + ), + } + }) +} // ── Config ─────────────────────────────────────────────────────────────────── @@ -151,25 +170,24 @@ impl SpeechModel for LMNTSpeechModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers.into_iter().collect(), - body: HttpBody::Json(Value::Object(body.clone())), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body.clone()), + aimux_provider_utils::create_binary_response_handler(), + lmnt_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs index a40364f4..a775ae55 100644 --- a/aimux-providers/src/luma.rs +++ b/aimux-providers/src/luma.rs @@ -16,12 +16,26 @@ use aimux_core::error::ApiCallError; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; +use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; + +fn luma_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let message = data + .get("detail") + .and_then(Value::as_array) + .and_then(|items| items.first()) + .and_then(|item| item.get("msg")) + .and_then(Value::as_str) + .unwrap_or("Unknown Luma error") + .to_string(); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: None, + } + }) +} const DEFAULT_POLL_INTERVAL_MS: u64 = 500; const DEFAULT_MAX_POLL_ATTEMPTS: u64 = 120; @@ -303,23 +317,22 @@ impl ImageModel for LumaImageModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generations_url(None), headers: header_list.clone(), - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + luma_failed_response_handler(), ) .await?; - let rh = resp.headers; - let submit_body: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers.unwrap_or_default(); + let submit_body: Value = resp.value; let generation_id = submit_body .get("id") @@ -328,26 +341,33 @@ impl ImageModel for LumaImageModel { AiMuxError::InvalidResponseData("missing id in Luma response".to_string()) })? .to_string(); + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Poll for completion let mut image_url = None; for _ in 0..max_poll_attempts { - let pr = send( - HttpRequest { - method: HttpMethod::Get, - url: self.generations_url(Some(&generation_id)), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - let pv: Value = serde_json::from_slice(&pr.body)?; + let pr = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: self.generations_url(Some(&generation_id)), + headers: header_list.clone(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + luma_failed_response_handler(), + ) + }) + .await?; + let response_body = pr.raw_value.as_ref().map(ToString::to_string); + let pv = pr.value; let state = pv.get("state").and_then(|v| v.as_str()).unwrap_or(""); if state == "completed" { @@ -364,13 +384,17 @@ impl ImageModel for LumaImageModel { break; } if state == "failed" { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(pr.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some(state.to_string()), message: "Image generation failed.".into(), - response_body: Some(String::from_utf8_lossy(&pr.body).into_owned()), - ..Default::default() - })); + response_body, + ..ApiCallError::new( + "Image generation failed.", + self.generations_url(Some(&generation_id)), + serde_json::json!({}), + ) + }))); } sleep_or_abort( std::time::Duration::from_millis(poll_interval), @@ -387,22 +411,23 @@ impl ImageModel for LumaImageModel { })?; // Download image - let ir = send( - HttpRequest { - method: HttpMethod::Get, - url: image_url, - headers: vec![], - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - let image_bytes = ir.body.to_vec(); + let ir = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: image_url.clone(), + headers: vec![], + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_binary_response_handler(), + aimux_provider_utils::create_status_code_error_response_handler(), + ) + }) + .await?; + let image_bytes = ir.value.to_vec(); Ok(ImageResult { images: ImageOutputs::Binary(vec![image_bytes]), diff --git a/aimux-providers/src/mistral/embedding.rs b/aimux-providers/src/mistral/embedding.rs index 301aa3e6..37678140 100644 --- a/aimux-providers/src/mistral/embedding.rs +++ b/aimux-providers/src/mistral/embedding.rs @@ -16,8 +16,7 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::MistralConfig; @@ -62,6 +61,10 @@ impl EmbeddingModel for MistralEmbeddingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_embeddings_per_call(&self) -> Option { Some(32) } @@ -98,25 +101,24 @@ impl EmbeddingModel for MistralEmbeddingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + super::mistral_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body)?; + let raw_value: Value = resp.value; let embeddings: Vec> = raw_value .get("data") diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs index f9d3b2f5..57d38538 100644 --- a/aimux-providers/src/mistral/mod.rs +++ b/aimux-providers/src/mistral/mod.rs @@ -17,6 +17,63 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; +use serde_json::Value; + +pub(crate) fn mistral_failed_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: data + .get("code") + .or_else(|| data.get("type")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} + +pub(crate) fn mistral_stream_error( + error: &Value, + url: &str, + request_body_values: Value, + response_headers: std::collections::HashMap, +) -> AiMuxError { + let status_code = error + .get("status_code") + .or_else(|| error.get("status")) + .or_else(|| error.get("code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) + .filter(|status| (400..=599).contains(status)); + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("Mistral stream failed before any output was generated") + .to_owned(); + let provider_code = + error + .get("code") + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(code) => Some(code.clone()), + Value::Number(code) => Some(code.to_string()), + _ => None, + }); + aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + error, + url, + request_body_values, + response_headers, + ) +} /// Configuration for the Mistral provider. #[derive(Debug, Clone)] @@ -25,7 +82,7 @@ pub struct MistralConfig { pub base_url: String, /// api_key 来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -128,24 +185,23 @@ impl Provider for MistralProvider { ), ("Content-Type".to_string(), "application/json".to_string()), ]; - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, config.retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: headers.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + mistral_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { #[serde(default)] @@ -157,7 +213,7 @@ impl Provider for MistralProvider { #[serde(default)] owned_by: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .data .into_iter() diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs index c23d7393..e6accd7e 100644 --- a/aimux-providers/src/mistral/model.rs +++ b/aimux-providers/src/mistral/model.rs @@ -23,21 +23,12 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::{ErrorStructure, parse_stream_error}; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::MistralConfig; use super::convert::{build_request_body, parse_finish_reason}; use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; -/// Mistral error structure: `{ "message": "...", "type": "..." }` (flat, no -/// `error` wrapper). -const MISTRAL_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["message"], - type_path: &["type"], -}; - /// An Mistral language model. pub struct MistralModel { model_id: String, @@ -210,6 +201,10 @@ impl LanguageModel for MistralModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; ProviderRecord { @@ -229,33 +224,26 @@ impl LanguageModel for MistralModel { async fn do_generate(&self, options: &CallOptions) -> Result { let body = build_request_body(&self.model_id, options, false); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &MISTRAL_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler(), + super::mistral_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let data: ChatCompletionResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let data: ChatCompletionResponse = resp.value; let choice = data.choices.into_iter().next().ok_or_else(|| { @@ -336,43 +324,47 @@ impl LanguageModel for MistralModel { async fn do_stream(&self, options: &CallOptions) -> Result { let body = build_request_body(&self.model_id, options, true); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &MISTRAL_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + super::mistral_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let mut sse_stream = SseStream::new(resp.body); + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse_stream = resp.value; // Peek at the first SSE event to detect early errors. - let first_event = sse_stream.next().await; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; if let Some(Ok(ref event)) = first_event - && let Ok(val) = serde_json::from_str::(&event.data) - && let Some(err_obj) = val.get("error") + && let Some(err_obj) = event.get("error") { - return Err(parse_stream_error(err_obj)); + return Err(super::mistral_stream_error( + err_obj, + &self.endpoint(), + body.clone(), + response_headers.clone(), + )); } + let stream_error_url = self.endpoint(); + let stream_error_body = body.clone(); + let stream_response_headers = response_headers.clone(); + let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: vec![] }); @@ -395,23 +387,16 @@ impl LanguageModel for MistralModel { } match event { - Ok(sse_event) => { - if sse_event.data == "[DONE]" { - break; - } - - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { error: e.into() }); - stream_errored = true; - break; - } - }; + Ok(parsed) => { if let Some(err_obj) = parsed.get("error") { yield Ok(StreamPart::Error { - error: parse_stream_error(err_obj), + error: super::mistral_stream_error( + err_obj, + &stream_error_url, + stream_error_body.clone(), + stream_response_headers.clone(), + ), }); stream_errored = true; break; @@ -420,9 +405,8 @@ impl LanguageModel for MistralModel { let chunk: StreamChunk = match serde_json::from_value(parsed) { Ok(c) => c, Err(e) => { - yield Ok(StreamPart::Error { error: e.into() }); - stream_errored = true; - break; + yield Err(e.into()); + continue; } }; @@ -580,12 +564,12 @@ impl LanguageModel for MistralModel { } } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - stream_errored = true; - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs index 09d2210b..731bade9 100644 --- a/aimux-providers/src/open_responses.rs +++ b/aimux-providers/src/open_responses.rs @@ -28,11 +28,58 @@ use aimux_core::types::{ FinishReason, FinishReasonUnified, ReasoningEffort, ResponseMetadata, Usage, Warning, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send_stream_timed, send_timed, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; + +fn open_responses_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("Open Responses request failed") + .to_string(), + provider_code: error + .and_then(|value| value.get("code").or_else(|| value.get("type"))) + .and_then(Value::as_str) + .map(str::to_string), + } + }) +} + +fn open_responses_successful_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let url = input.url.clone(); + let request_body_values = input.request_body_values.clone(); + let output = aimux_provider_utils::create_json_response_handler::() + .handle(input) + .await?; + if let Some(error) = output + .value + .get("error") + .and_then(Value::as_object) + .filter(|error| !error.is_empty()) + { + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("Open Responses request failed"); + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status), + provider_code: error + .get("code") + .or_else(|| error.get("type")) + .and_then(Value::as_str) + .map(str::to_string), + response_body: Some(output.value.to_string()), + response_headers: output.response_headers.clone(), + ..ApiCallError::new(message, url, request_body_values) + }))); + } + Ok(output) + }) +} // == Config == @@ -267,47 +314,24 @@ impl LanguageModel for OpenResponsesModel { let (body, warnings) = build_request_body(&self.model_id, options, &self.config.provider_options_name); - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.config.url.clone(), headers: headers.into_iter().collect(), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + open_responses_successful_response_handler(), + open_responses_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw: Value = serde_json::from_slice(&resp.body)?; - - // Check for response.error first (surfaces before the no-output fallback). - if let Some(error) = raw.get("error").and_then(|e| e.as_object()) - && !error.is_empty() - { - let message = error - .get("message") - .and_then(|m| m.as_str()) - .unwrap_or("unknown error"); - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: error - .get("code") - .or_else(|| error.get("type")) - .and_then(|c| c.as_str()) - .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } + let raw: Value = resp.value; // Check for null/missing output. let output = raw.get("output"); @@ -442,48 +466,57 @@ impl LanguageModel for OpenResponsesModel { b }; - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.config.url.clone(), headers: headers.into_iter().collect(), - body: HttpBody::Json(stream_body), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + stream_body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + open_responses_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let mut sse_stream = SseStream::new(resp.body); + let mut sse_stream = resp.value; // Peek at the first SSE event to detect early errors. - let first_event = sse_stream.next().await; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; if let Some(Ok(ref event)) = first_event - && let Ok(val) = serde_json::from_str::(&event.data) - && let Some(err_obj) = val.get("error") + && let Some(err_obj) = event.get("error") { let message = err_obj .get("message") .and_then(|m| m.as_str()) .unwrap_or("stream error"); - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: err_obj - .get("code") - .or_else(|| err_obj.get("type")) - .and_then(|c| c.as_str()) - .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(event.data.clone()), - ..Default::default() - })); + let status_code = err_obj + .get("status") + .or_else(|| err_obj.get("code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) + .filter(|status| (400..=599).contains(status)); + let provider_code = err_obj + .get("code") + .or_else(|| err_obj.get("type")) + .and_then(|c| c.as_str()) + .map(std::string::ToString::to_string); + return Err(aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + event, + self.config.url.clone(), + stream_body, + response_headers.clone(), + )); } let stream = async_stream::stream! { @@ -506,19 +539,7 @@ impl LanguageModel for OpenResponsesModel { while let Some(event) = event_iter.next().await { match event { - Ok(sse_event) => { - if sse_event.data == "[DONE]" { - break; - } - - let chunk: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { error: e.into() }); - break; - } - }; - + Ok(chunk) => { let chunk_type = chunk .get("type") .and_then(|t| t.as_str()) @@ -782,11 +803,12 @@ impl LanguageModel for OpenResponsesModel { } } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/openai/embedding.rs b/aimux-providers/src/openai/embedding.rs index 2d06b3ca..a9358926 100644 --- a/aimux-providers/src/openai/embedding.rs +++ b/aimux-providers/src/openai/embedding.rs @@ -16,15 +16,14 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send}; +use aimux_provider_utils::HttpRequest; use super::OpenAIConfig; /// An OpenAI-compatible embedding model. /// /// Works with any OpenAI-compatible `/embeddings` endpoint. Does **not** hold an -/// HTTP client — `http::send` uses the shared `Client` internally (RFC-0009 §4.1). +/// HTTP client — the `aimux-provider-utils` API helpers use the shared `Client` internally (RFC-0009 §4.1). pub struct OpenAIEmbeddingModel { model_id: String, config: OpenAIConfig, @@ -77,6 +76,10 @@ impl EmbeddingModel for OpenAIEmbeddingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_embeddings_per_call(&self) -> Option { Some(2048) } @@ -110,27 +113,26 @@ impl EmbeddingModel for OpenAIEmbeddingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), ) .await?; // `send` retries 408/409/429/5xx and returns an error for non-2xx, so an `Ok` // response here is guaranteed to be 2xx — no manual is_success() check. - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body)?; + let raw_value: Value = resp.value; // Extract embeddings: response.data[].embedding // The embedding field can be a JSON array of floats (default) or a @@ -220,6 +222,10 @@ fn decode_base64_embedding(s: &str) -> Option> { if bytes.len() % 4 != 0 { return None; } + // Clippy 1.98 suggests `as_chunks::<4>()`, which is stable only since + // Rust 1.88; the workspace MSRV is 1.85. `unknown_lints` keeps this + // buildable on toolchains older than the lint itself. + #[allow(unknown_lints, clippy::chunks_exact_to_as_chunks)] Some( bytes .chunks_exact(4) diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs index 99cb8873..b75f62b1 100644 --- a/aimux-providers/src/openai/files.rs +++ b/aimux-providers/src/openai/files.rs @@ -14,8 +14,7 @@ use aimux_core::error::AiMuxError; use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData, UploadFileResult}; use aimux_core::shared::FileBytes; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send}; +use aimux_provider_utils::{HttpBody, HttpRequest}; use super::OpenAIConfig; @@ -87,7 +86,7 @@ struct OpenAIFilesResponse { /// An OpenAI Files interface for uploading files. /// -/// Aligned with TS `OpenAIFiles`. Does **not** hold an HTTP client — `http::send` +/// Aligned with TS `OpenAIFiles`. Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers /// uses the process-wide shared `Client` internally (RFC-0009 §4.1). pub struct OpenAIFiles { config: OpenAIConfig, @@ -168,23 +167,22 @@ impl Files for OpenAIFiles { // and its content-type are carried by `HttpBody::Bytes` — the HTTP layer // sets `Content-Type` from it, so it is intentionally not added to the // header list above. - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Bytes(body, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body, content_type), + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), ) .await?; - let data: OpenAIFilesResponse = serde_json::from_slice::(&resp.body)?; + let data: OpenAIFilesResponse = resp.value; // Build provider metadata. let mut metadata = serde_json::Map::new(); diff --git a/aimux-providers/src/openai/image.rs b/aimux-providers/src/openai/image.rs index cbd2d42b..8e588dbb 100644 --- a/aimux-providers/src/openai/image.rs +++ b/aimux-providers/src/openai/image.rs @@ -20,8 +20,7 @@ use aimux_core::image_model::{ }; use aimux_core::shared::{SharedProviderMetadata, Warning}; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send}; +use aimux_provider_utils::{HttpBody, HttpRequest}; use super::OpenAIConfig; @@ -119,6 +118,10 @@ impl ImageModel for OpenAIImageModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_images_per_call(&self) -> Option { Some(get_max_images_per_call(&self.model_id)) } @@ -151,48 +154,49 @@ impl ImageModel for OpenAIImageModel { let (form_body, content_type) = build_edit_multipart(&self.model_id, options, &openai_options); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.edits_endpoint(), headers: build_header_list(&headers), - // Content-Type is set by the http layer from the `Bytes` body. - body: HttpBody::Bytes(form_body, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(form_body, content_type), + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), ) .await?; - let val: Value = serde_json::from_slice(&resp.body)?; - (val, resp.headers, None) + let val: Value = resp.value; + (val, resp.response_headers.unwrap_or_default(), None) } else { // ── Generation path: JSON body ── let openai_options = parse_generation_provider_options(&options.provider_options); let body = build_generation_body(&self.model_id, options, &openai_options); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generations_endpoint(), headers: build_header_list(&headers), - body: HttpBody::Json(Value::Object(body.clone())), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body.clone()), + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), ) .await?; - let val: Value = serde_json::from_slice(&resp.body)?; - (val, resp.headers, Some(Value::Object(body))) + let val: Value = resp.value; + ( + val, + resp.response_headers.unwrap_or_default(), + Some(Value::Object(body)), + ) }; let images = extract_images(&body_value); diff --git a/aimux-providers/src/openai/mod.rs b/aimux-providers/src/openai/mod.rs index 4941a0a2..220758c8 100644 --- a/aimux-providers/src/openai/mod.rs +++ b/aimux-providers/src/openai/mod.rs @@ -30,6 +30,63 @@ use aimux_core::provider::Provider; use aimux_provider_utils::{RetryConfig, load_api_key, without_trailing_slash}; use serde_json::Value; +pub(crate) fn openai_failed_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("code") + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }), + } + }) +} + +pub(crate) fn openai_stream_error( + error: &Value, + url: &str, + request_body_values: Value, + response_headers: std::collections::HashMap, +) -> AiMuxError { + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("OpenAI stream failed before any output was generated") + .to_owned(); + let code = error.get("code").or_else(|| error.get("type")); + let provider_code = code.and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + // Only a numeric HTTP status in the payload is a status; a string code + // ("invalid_api_key") must not be laundered into a retryable 500. + let status_code = code + .and_then(Value::as_u64) + .filter(|status| (400..=599).contains(status)) + .map(|status| status as u16); + + aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + error, + url, + request_body_values, + response_headers, + ) +} + /// 描述 OpenAI 兼容厂商的差异。 /// /// 薄封装填这个结构，共享的请求构造和响应解析读它决定行为。 @@ -163,9 +220,7 @@ pub struct OpenAIConfig { /// 厂商能力差异描述。默认 `full()`（支持全部能力）。 /// 薄封装用 `with_profile()` 设置差异。 pub profile: OpenAICompatProfile, - /// 重试配置。默认 `RetryConfig::default()`（max_retries=2）。 - /// 测试中可用 `.with_retry_config(RetryConfig { max_retries: 0, .. })` - /// 关闭重试（RFC-0009 §4.2）。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, /// Provider 级请求体覆盖（RFC-0017）。在标准请求体 + 内置厂商 override /// 之后 deep-merge。per-call 的 `CallOptions.body_overrides` 在此之后 @@ -231,7 +286,7 @@ impl OpenAIConfig { self } - /// 设置重试配置。传入 `max_retries: 0` 可关闭重试。 + /// Set the retry configuration. Pass `max_retries: 0` to disable retries. #[must_use] pub fn with_retry_config(mut self, config: RetryConfig) -> Self { self.retry_config = config; @@ -265,7 +320,7 @@ impl OpenAIConfig { /// OpenAI provider — creates `OpenAIModel` instances. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use /// the process-wide shared `Client` internally (RFC-0009 §4.1). pub struct OpenAIProvider { config: OpenAIConfig, @@ -358,8 +413,7 @@ impl Provider for OpenAIProvider { Box::pin(async move { let headers = model::build_auth_headers(&config); let runtime = - model::execute_list_models(&config.base_url, &headers, &config.retry_config) - .await?; + model::execute_list_models(&config.base_url, &headers, config.retry_config).await?; Ok(runtime) }) } diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs index ea06317f..29604a43 100644 --- a/aimux-providers/src/openai/model.rs +++ b/aimux-providers/src/openai/model.rs @@ -2,7 +2,7 @@ //! //! The HTTP request/response handling lives in the free functions //! [`execute_generate`] and [`execute_stream`], which take an endpoint URL, a -//! header map and a model id. They call `http::send` / `http::send_stream` — +//! header map and a model id. They call the `aimux-provider-utils` API helpers — //! **no `reqwest` types cross this boundary**. This lets other providers that //! speak the OpenAI chat-completions wire format (notably Azure OpenAI) reuse //! the conversion + streaming logic while supplying their own URL and auth. @@ -20,11 +20,7 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::{DEFAULT_ERROR_STRUCTURE, parse_stream_error}; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send_stream_timed, send_timed, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::OpenAIConfig; use super::convert::{RequestBodyResult, build_request_body_with_warnings, parse_finish_reason}; @@ -32,7 +28,7 @@ use super::types::{ChatCompletionResponse, StreamChunk, UsageResponse}; /// An OpenAI-compatible language model. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct OpenAIModel { model_id: String, @@ -86,26 +82,6 @@ pub fn build_auth_headers(config: &super::OpenAIConfig) -> HashMap, -) -> RetryConfig { - match max_retries_override { - Some(n) => RetryConfig { - max_retries: n, - ..*provider - }, - None => *provider, - } -} - /// Merge provider-level `body_overrides` into the per-call options (RFC-0017). /// /// Provider-level overrides are applied first (lower priority); per-call @@ -224,13 +200,16 @@ impl LanguageModel for OpenAIModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { super::config_snapshot_from_config(&self.config.provider, &self.model_id, &self.config) } async fn do_generate(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()); - let retry_config = resolve_retry_config(&self.config.retry_config, options.max_retries); let options = merge_body_overrides(options, &self.config.body_overrides); execute_generate( &self.endpoint(), @@ -239,14 +218,12 @@ impl LanguageModel for OpenAIModel { &options, &self.config.provider, &self.config.profile, - &retry_config, ) .await } async fn do_stream(&self, options: &CallOptions) -> Result { let headers = self.build_headers(options.headers.as_ref()); - let retry_config = resolve_retry_config(&self.config.retry_config, options.max_retries); let options = merge_body_overrides(options, &self.config.body_overrides); execute_stream( &self.endpoint(), @@ -255,7 +232,6 @@ impl LanguageModel for OpenAIModel { &options, &self.config.provider, &self.config.profile, - &retry_config, ) .await } @@ -297,35 +273,32 @@ pub async fn execute_generate( options: &CallOptions, provider: &str, profile: &super::OpenAICompatProfile, - retry_config: &RetryConfig, ) -> Result { let request_result = build_request_body_with_warnings(model_id, options, false, provider, profile)?; let body = request_result.body; - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: endpoint.to_string(), headers: build_header_list(headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - *retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + super::openai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); // Parse the raw body once: the `Value` keeps the provider's original // fields (incl. vendor-specific usage fields) for `Usage.raw` (M10). - let response_value: Value = serde_json::from_slice(&resp.body)?; - let data: ChatCompletionResponse = serde_json::from_value(response_value.clone())?; + let response_value = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let choice = data .choices @@ -460,7 +433,6 @@ pub async fn execute_stream( options: &CallOptions, provider: &str, profile: &super::OpenAICompatProfile, - retry_config: &RetryConfig, ) -> Result { let request_result = build_request_body_with_warnings(model_id, options, true, provider, profile)?; @@ -468,36 +440,43 @@ pub async fn execute_stream( // they are emitted in `StreamStart` below instead of being dropped. let RequestBodyResult { body, warnings } = request_result; - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: endpoint.to_string(), headers: build_header_list(headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - *retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + super::openai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let mut sse_stream = SseStream::new(resp.body); + let mut sse_stream = resp.value; - // Peek at the first SSE event to detect early errors (before any output). - // The TS SDK rejects the doStream promise when the very first chunk is an - // error. We replicate that by reading one event here. - let first_event = sse_stream.next().await; + // Aimux checks only the first SSE event before returning the stream. The + // baseline OpenAI provider scans until semantic output; this narrower peek + // still keeps an immediately reported provider error inside Core's + // operation-retry boundary (RFC-0031 §8.3). A normal event is chained back + // below and is never consumed. + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; if let Some(Ok(ref event)) = first_event - && let Ok(val) = serde_json::from_str::(&event.data) - && let Some(err_obj) = val.get("error") + && let Some(err_obj) = event.get("error") { - return Err(parse_stream_error(err_obj)); + return Err(super::openai_stream_error( + err_obj, + endpoint, + body.clone(), + response_headers.clone(), + )); } // Capture the provider's stream usage key before entering the async stream @@ -509,6 +488,9 @@ pub async fn execute_stream( // M2 (RFC-0016): capture whether raw chunks should be emitted — the // borrowed `options` cannot be moved into the generator. let emit_raw_chunks = options.include_raw_chunks == Some(true); + let stream_error_url = endpoint.to_owned(); + let stream_error_body = body.clone(); + let stream_response_headers = response_headers.clone(); let stream = async_stream::stream! { // First part: StreamStart. @@ -540,21 +522,7 @@ pub async fn execute_stream( } match event { - Ok(sse_event) => { - if sse_event.data == "[DONE]" { - break; - } - - // Parse as generic Value first to detect errors. - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - // Unparsable chunk — emit Error, then finish. - yield Ok(StreamPart::Error { error: e.into() }); - stream_errored = true; - break; - } - }; + Ok(parsed) => { // M2 (RFC-0016): emit the raw provider chunk for debugging // before it is consumed below. JSON payloads only — the @@ -568,7 +536,12 @@ pub async fn execute_stream( // Check for mid-stream error. if let Some(err_obj) = parsed.get("error") { yield Ok(StreamPart::Error { - error: parse_stream_error(err_obj), + error: super::openai_stream_error( + err_obj, + &stream_error_url, + stream_error_body.clone(), + stream_response_headers.clone(), + ), }); stream_errored = true; break; @@ -593,9 +566,8 @@ pub async fn execute_stream( let chunk: StreamChunk = match serde_json::from_value(parsed) { Ok(c) => c, Err(e) => { - yield Ok(StreamPart::Error { error: e.into() }); - stream_errored = true; - break; + yield Err(e.into()); + continue; } }; @@ -804,12 +776,12 @@ pub async fn execute_stream( } } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - stream_errored = true; - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } @@ -928,29 +900,34 @@ struct ModelEntry { pub async fn execute_list_models( base_url: &str, headers: &HashMap, - retry_config: &RetryConfig, + retry_config: aimux_core::retry::RetryConfig, ) -> Result, AiMuxError> { // Strip a trailing slash so we don't get `//models`. let base = base_url.trim_end_matches('/'); let url = format!("{base}/models"); - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers: build_header_list(headers), - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - *retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + // `list_models` is not a Core user operation, so nothing above this call + // retries it — apply the Core retry primitive here (a catalogue GET is + // idempotent) so a transient 429/503/transport failure behaves like every + // other exchange instead of failing on the first hiccup. + let retries = aimux_core::retry::prepare_retries(None, retry_config, None); + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: build_header_list(headers), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + super::openai_failed_response_handler(), + ) + }) + .await?; - let parsed: ModelsListResponse = serde_json::from_slice(&resp.body)?; + let parsed: ModelsListResponse = resp.value; Ok(parsed .data diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs index ad458ee3..89b96e6e 100644 --- a/aimux-providers/src/openai/responses/mod.rs +++ b/aimux-providers/src/openai/responses/mod.rs @@ -43,16 +43,14 @@ use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateResult, StreamResult}; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::OpenAIConfig; use responses_convert::build_header_list; /// An OpenAI Responses API language model. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use /// the process-wide shared `Client` internally (RFC-0009 §4.1). /// /// Created via `OpenAIResponsesProvider` or @@ -113,6 +111,10 @@ impl LanguageModel for OpenAIResponsesModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { super::config_snapshot_from_config(&self.config.provider, &self.model_id, &self.config) } @@ -123,34 +125,36 @@ impl LanguageModel for OpenAIResponsesModel { let body = request_result.body; let provider_key = self.provider_options_name().to_string(); - let resp = send_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + super::openai_failed_response_handler(), ) .await?; - let status = resp.status; - let response_headers = resp.headers; - - let data: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let raw_body = resp + .raw_value + .as_ref() + .map(ToString::to_string) + .unwrap_or_default(); + let data = resp.value; responses_convert::build_responses_generate_result( &data, - status, - &String::from_utf8_lossy(&resp.body), + &raw_body, request_result.warnings, provider_key, + endpoint, body, response_headers, ) @@ -173,35 +177,38 @@ impl LanguageModel for OpenAIResponsesModel { .and_then(serde_json::Value::as_bool) == Some(true); - let resp = send_stream_timed( + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + super::openai_failed_response_handler(), ) .await?; - let status = resp.status; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let mut sse_stream = SseStream::new(resp.body); - let first_event = sse_stream.next().await; + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; let stream = responses_convert::build_responses_event_stream( first_event, sse_stream, - status, provider_key, warnings, store_flag, + endpoint, + body.clone(), + response_headers.clone(), )?; Ok(StreamResult { diff --git a/aimux-providers/src/openai/responses/responses_convert.rs b/aimux-providers/src/openai/responses/responses_convert.rs index ede77526..7b8203c8 100644 --- a/aimux-providers/src/openai/responses/responses_convert.rs +++ b/aimux-providers/src/openai/responses/responses_convert.rs @@ -24,7 +24,6 @@ use aimux_core::error::ApiCallError; use aimux_core::result::{GenerateContent, GenerateResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage, Warning}; -use aimux_stream::{SseError, SseEvent}; use super::convert::{convert_responses_usage, map_responses_finish_reason, parse_usage}; use super::types::ResponsesUsage; @@ -71,10 +70,10 @@ pub fn build_header_list(headers: &HashMap) -> Vec<(String, Stri /// fields such as `output` are missing. pub fn build_responses_generate_result( data: &Value, - status: u16, raw_body: &str, request_warnings: Vec, provider_key: String, + request_url: String, body: Value, response_headers: HashMap, ) -> Result { @@ -91,15 +90,21 @@ pub fn build_responses_generate_result( .or_else(|| err_obj.get("code")) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string); - return Err(AiMuxError::ApiCall(ApiCallError { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { // Provider-declared in-band failure: keep the observed 2xx - // envelope status and the full raw body (§2.2). - status_code: Some(status), + // envelope status and the full raw body (§2.2). The request body + // is raw here (the exchange-path redaction ran on its own copy), + // so redact before it lands in the public error. + status_code: Some(200), provider_code, - message: message.to_string(), response_body: Some(raw_body.to_string()), - ..Default::default() - })); + response_headers: Some(response_headers.clone()), + ..ApiCallError::new( + message, + request_url, + aimux_provider_utils::redact_error_context(body.clone()), + ) + }))); } let output = data.get("output").and_then(|v| v.as_array()); @@ -318,9 +323,7 @@ pub fn build_responses_generate_result( // -- Streaming SSE event reducer --------------------------------------------- /// A tool call being streamed (tracked by `output_index`). -#[allow(dead_code)] struct OngoingToolCall { - tool_name: String, tool_call_id: String, } @@ -375,31 +378,39 @@ fn generate_source_id() -> String { /// `function_call_arguments.delta`, `custom_tool_call_input.delta`, /// `reasoning_summary_part.added/done` and `reasoning_summary_text.delta`). /// -/// The caller performs the HTTP send (`send_stream`) and hands the peeked +/// The caller performs the API call and hands the peeked /// `first_event` plus the remainder `sse_stream` to this reducer; an early /// `error` / `response.failed` surfaces as a clean `Err` here. /// /// # Errors /// -/// An early `error` / `response.failed` event yields an `ApiCall` error item -/// in the returned stream; malformed events yield parse-error items. +/// An early `error` / `response.failed` event returns an `ApiCall` setup error; +/// malformed events yield parse-error items and do not end the stream. +// This reducer is shared by OpenAI, Azure, and Codex. Keeping the wire inputs +// explicit is clearer than introducing a second context object used nowhere +// else. +#[allow(clippy::too_many_arguments)] pub fn build_responses_event_stream( - first_event: Option>, + first_event: Option>, sse_stream: S, - status: u16, provider_key: String, warnings: Vec, store_flag: bool, + request_url: String, + request_body: Value, + response_headers: HashMap, ) -> Result where - S: Stream> + Unpin + Send + 'static, + S: Stream> + Unpin + Send + 'static, { // Peek at the first SSE event to detect early errors (before any output). - if let Some(Ok(ref event)) = first_event - && let Ok(val) = serde_json::from_str::(&event.data) - { + if let Some(Ok(ref val)) = first_event { let etype = val.get("type").and_then(|v| v.as_str()).unwrap_or(""); if etype == "error" || etype == "response.failed" { + let error_value = val + .get("response") + .and_then(|response| response.get("error")) + .or_else(|| val.get("error")); let message = val .get("response") .and_then(|r| r.get("error")) @@ -411,21 +422,30 @@ where .and_then(|v| v.as_str()) }) .unwrap_or("Responses API stream error"); - return Err(AiMuxError::ApiCall(ApiCallError { - // Mid-stream provider failure arrives on a successful HTTP - // response: keep the observed 2xx status (§2.2). - status_code: Some(status), - provider_code: val - .get("response") - .and_then(|r| r.get("error")) - .or_else(|| val.get("error")) - .and_then(|e| e.get("type").or_else(|| e.get("code"))) - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(event.data.clone()), - ..Default::default() - })); + let status_code = val + .get("status") + .or_else(|| val.get("status_code")) + .or_else(|| error_value.and_then(|error| error.get("status"))) + .or_else(|| error_value.and_then(|error| error.get("code"))) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) + .filter(|status| (400..=599).contains(status)); + let provider_code = val + .get("response") + .and_then(|r| r.get("error")) + .or_else(|| val.get("error")) + .and_then(|e| e.get("type").or_else(|| e.get("code"))) + .and_then(|v| v.as_str()) + .map(std::string::ToString::to_string); + return Err(aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + val, + request_url.clone(), + request_body.clone(), + response_headers.clone(), + )); } } @@ -451,27 +471,8 @@ where futures::stream::iter(first_event.into_iter()).chain(sse_stream); while let Some(event) = event_iter.next().await { - if stream_errored { - break; - } - match event { - Ok(sse_event) => { - if sse_event.data == "[DONE]" { - break; - } - - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::JsonParse(e.to_string()), - }); - stream_errored = true; - break; - } - }; - + Ok(parsed) => { let etype = parsed .get("type") .and_then(|v| v.as_str()) @@ -539,7 +540,6 @@ where ongoing_tool_calls.insert( output_index, OngoingToolCall { - tool_name: name.clone(), tool_call_id: call_id.clone(), }, ); @@ -566,7 +566,6 @@ where ongoing_tool_calls.insert( output_index, OngoingToolCall { - tool_name: name.clone(), tool_call_id: call_id.clone(), }, ); @@ -972,19 +971,27 @@ where .and_then(|e| e.get("message")) .and_then(|v| v.as_str()) .unwrap_or("Responses API stream failed"); + // In-band failure inside the SSE stream: the shared + // helper redacts the raw request context. yield Ok(StreamPart::Error { - error: AiMuxError::ApiCall(ApiCallError { - status_code: Some(status), - provider_code: resp_obj + error: aimux_provider_utils::stream_error_api_call( + message, + resp_obj .get("error") .and_then(|e| e.get("type").or_else(|| e.get("code"))) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(sse_event.data.clone()), - ..Default::default() - }), + Some(200), + &parsed, + request_url.clone(), + request_body.clone(), + response_headers.clone(), + ), }); + // A terminal error ends this stream; waiting + // for more events can hang on a source that + // keeps the connection open. + break; } } } @@ -1002,18 +1009,21 @@ where .and_then(|v| v.as_str()) .unwrap_or("Responses API stream error"); yield Ok(StreamPart::Error { - error: AiMuxError::ApiCall(ApiCallError { - status_code: Some(status), - provider_code: parsed + error: aimux_provider_utils::stream_error_api_call( + message, + parsed .get("error") .and_then(|e| e.get("type").or_else(|| e.get("code"))) .and_then(|v| v.as_str()) .map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(sse_event.data.clone()), - ..Default::default() - }), + Some(200), + &parsed, + request_url.clone(), + request_body.clone(), + response_headers.clone(), + ), }); + break; } _ => { @@ -1024,12 +1034,12 @@ where } } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - stream_errored = true; - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/openai/speech.rs b/aimux-providers/src/openai/speech.rs index d7a10a6e..a7a6467e 100644 --- a/aimux-providers/src/openai/speech.rs +++ b/aimux-providers/src/openai/speech.rs @@ -21,8 +21,7 @@ use aimux_core::speech_model::{ AudioData, SpeechCallOptions, SpeechModel, SpeechRequest, SpeechResponse, SpeechResult, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send}; +use aimux_provider_utils::HttpRequest; use super::OpenAIConfig; @@ -84,6 +83,10 @@ impl SpeechModel for OpenAISpeechModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + async fn do_generate(&self, options: &SpeechCallOptions) -> Result { let (body, warnings) = build_request_body_and_warnings(options, &self.model_id)?; @@ -94,27 +97,26 @@ impl SpeechModel for OpenAISpeechModel { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body.clone())), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body.clone()), + aimux_provider_utils::create_binary_response_handler(), + super::openai_failed_response_handler(), ) .await?; // send() returns Ok only for 2xx responses; non-2xx (incl. 408/409/429/5xx // after exhausting retries) is mapped to an AiMuxError internally. - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); diff --git a/aimux-providers/src/openai/transcription.rs b/aimux-providers/src/openai/transcription.rs index 602aa940..321ad529 100644 --- a/aimux-providers/src/openai/transcription.rs +++ b/aimux-providers/src/openai/transcription.rs @@ -30,10 +30,7 @@ use aimux_core::transcription_model::{ TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, MultipartForm, media_type_to_extension, send, -}; +use aimux_provider_utils::{HttpBody, HttpRequest, MultipartForm, media_type_to_extension}; use super::OpenAIConfig; @@ -269,6 +266,10 @@ impl TranscriptionModel for OpenAITranscriptionModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + async fn do_generate( &self, options: &TranscriptionCallOptions, @@ -342,27 +343,25 @@ impl TranscriptionModel for OpenAITranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Bytes(body_bytes, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - self.config.retry_config, - &DEFAULT_ERROR_STRUCTURE, + HttpBody::Bytes(body_bytes, content_type), + aimux_provider_utils::create_json_response_handler::(), + super::openai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); + let response_headers = resp.response_headers.unwrap_or_default(); - let parsed: OpenAITranscriptionResponse = serde_json::from_slice(&resp.body)?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; // Map language name to ISO 639-1 code. let language = parsed @@ -495,7 +494,7 @@ impl TranscriptionModel for OpenAITranscriptionModel { // surface from do_stream's Result (clear contract: connect failure = // Err, not an in-stream error part). let req = WebSocketRequest { - url: ws_url, + url: ws_url.clone(), headers: std::mem::take(&mut header_list), subprotocols: Vec::new(), abort_signal: abort.clone(), @@ -514,6 +513,8 @@ impl TranscriptionModel for OpenAITranscriptionModel { let mut audio = options.audio; let mut audio_done = false; let mut committed = false; + let error_url = ws_url; + let error_request = session_update.clone(); let stream = async_stream::stream! { yield Ok(TranscriptionStreamPart::StreamStart { warnings: vec![] }); @@ -580,11 +581,12 @@ impl TranscriptionModel for OpenAITranscriptionModel { // Peer closed. If we never got `completed`, // this is a truncated session — surface it // rather than ending silently. - yield Err(AiMuxError::ApiCall( - aimux_core::error::ApiCallError { - message: "realtime transcription socket closed before completion".into(), - ..Default::default() - })); + yield Err(AiMuxError::ApiCall(Box::new( + aimux_core::error::ApiCallError::new( + "realtime transcription socket closed before completion", + error_url.clone(), + error_request.clone(), + )))); break; } Some(Err(e)) => { @@ -645,11 +647,15 @@ impl TranscriptionModel for OpenAITranscriptionModel { let msg = value.pointer("/error/message") .and_then(|v| v.as_str()) .unwrap_or("realtime transcription error"); - yield Err(AiMuxError::ApiCall( + yield Err(AiMuxError::ApiCall(Box::new( aimux_core::error::ApiCallError { - message: msg.to_string(), - ..Default::default() - })); + response_body: Some(value.to_string()), + ..aimux_core::error::ApiCallError::new( + msg, + error_url.clone(), + error_request.clone(), + ) + }))); ws.close().await; break; } diff --git a/aimux-providers/src/parallel_ai.rs b/aimux-providers/src/parallel_ai.rs index a554b488..d66c9873 100644 --- a/aimux-providers/src/parallel_ai.rs +++ b/aimux-providers/src/parallel_ai.rs @@ -19,10 +19,25 @@ use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn parallel_ai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} /// Fixed model ID for the Parallel AI search model. const MODEL_ID: &str = "parallel-search"; @@ -183,27 +198,25 @@ impl SearchModel for ParallelAiSearchModel { .into_iter() .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler::(), + parallel_ai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: ParallelAiResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; Ok(SearchResult { results: map_results(data.results), diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs index 38e8f24b..0817dd08 100644 --- a/aimux-providers/src/prodia.rs +++ b/aimux-providers/src/prodia.rs @@ -16,11 +16,36 @@ use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; use aimux_core::shared::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn prodia_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let detail = data.get("detail"); + let message = detail + .and_then(Value::as_str) + .map(str::to_string) + .or_else(|| { + detail + .filter(|value| !value.is_null()) + .map(Value::to_string) + }) + .or_else(|| { + data.get("error") + .and_then(Value::as_str) + .map(str::to_string) + }) + .or_else(|| { + data.get("message") + .and_then(Value::as_str) + .map(str::to_string) + }) + .unwrap_or_else(|| "Unknown Prodia error".to_string()); + aimux_provider_utils::ProviderErrorParts { + message, + provider_code: None, + } + }) +} /// Configuration for the Prodia provider. #[derive(Debug, Clone)] @@ -232,23 +257,22 @@ impl ImageModel for ProdiaImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_binary_response_handler(), + prodia_failed_response_handler(), ) .await?; - let rh = resp.headers; + let rh = resp.response_headers.unwrap_or_default(); let content_type = rh.get("content-type").cloned().unwrap_or_default(); // Extract boundary @@ -264,7 +288,7 @@ impl ImageModel for ProdiaImageModel { )) })?; - let body_bytes = resp.body.to_vec(); + let body_bytes = resp.value.to_vec(); // Parse multipart let parts = parse_multipart(&body_bytes, &boundary); @@ -314,7 +338,8 @@ impl ImageModel for ProdiaImageModel { // ════════════════════════════════════════════════════════════════════════════ use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoModel, VideoResponse, VideoResult, + VideoCallOptions, VideoData, VideoModel, VideoOperationStart, VideoOperationStatus, + VideoResponse, VideoResult, }; /// Prodia video generation model — implements `VideoModel`. @@ -361,7 +386,10 @@ impl VideoModel for ProdiaVideoModel { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { let warnings: Vec = Vec::new(); let mut config_obj = Map::new(); @@ -378,23 +406,23 @@ impl VideoModel for ProdiaVideoModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit job. - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: format!("{}/job", self.config.base_url), - headers: header_list.clone(), - body: HttpBody::Json(body), + headers: header_list, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + prodia_failed_response_handler(), ) .await?; - let job: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let job: Value = resp.value; let job_id = job .get("job") .and_then(|v| v.as_str()) @@ -405,50 +433,75 @@ impl VideoModel for ProdiaVideoModel { })? .to_string(); - // Poll for completion. - let mut raw_body: Value; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: format!("{}/job/{}", self.config.base_url, job_id), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; - let status_str = raw_body - .get("status") - .and_then(|v| v.as_str()) - .unwrap_or(""); - if status_str == "done" || status_str == "failed" { - if status_str == "failed" { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: Some(status_str.to_string()), - message: "Prodia video generation failed".to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } - break; + Ok(VideoOperationStart { + operation: json!({ "job_id": job_id }), + warnings, + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: response_headers, + }, + }) + } + + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let job_id = operation + .get("job_id") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "prodia operation reference is missing job_id".to_string(), + ) + })?; + + let headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let poll_url = format!("{}/job/{}", self.config.base_url, job_id); + + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list, + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + prodia_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let raw_body = resp.value; + let status_str = raw_body + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or(""); + match status_str { + "done" => {} + // A terminally failed job must be a non-retryable error, not + // Pending, so the Core poll loop stops immediately. + "failed" => { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(status_str.to_string()), + response_body, + ..ApiCallError::new( + "Prodia video generation failed", + poll_url, + serde_json::json!({}), + ) + }))); } + // queued / running / unknown — keep polling. + _ => return Ok(VideoOperationStatus::Pending), } // Extract video URL. @@ -464,15 +517,15 @@ impl VideoModel for ProdiaVideoModel { media_type: "video/mp4".to_string(), }]; - Ok(VideoResult { + Ok(VideoOperationStatus::Completed(VideoResult { videos, - warnings, + warnings: Vec::new(), provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), headers: Some(response_headers), }, - }) + })) } } diff --git a/aimux-providers/src/provider.rs b/aimux-providers/src/provider.rs index aee64572..0c1c8fea 100644 --- a/aimux-providers/src/provider.rs +++ b/aimux-providers/src/provider.rs @@ -25,7 +25,6 @@ use serde_json::Value; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::RetryConfig; use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIProvider}; @@ -490,10 +489,7 @@ fn build_resolved_config( config = config.with_project(project.clone()); } if let Some(max_retries) = entry.max_retries { - config = config.with_retry_config(RetryConfig { - max_retries, - ..RetryConfig::default() - }); + config.retry_config.max_retries = max_retries; } if let Some(overrides) = &entry.body_overrides { config = config.with_body_overrides(overrides.clone()); @@ -514,10 +510,7 @@ fn build_resolved_config( config = config.with_project(project); } if let Some(max_retries) = opts.max_retries { - config = config.with_retry_config(RetryConfig { - max_retries, - ..RetryConfig::default() - }); + config.retry_config.max_retries = max_retries; } if let Some(overrides) = opts.body_overrides { config = config.with_body_overrides(overrides); diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs index f3b820ed..d0366bdc 100644 --- a/aimux-providers/src/recraft.rs +++ b/aimux-providers/src/recraft.rs @@ -19,11 +19,27 @@ use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; use aimux_core::provider::Provider; +use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn recraft_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} const DEFAULT_BASE_URL: &str = "https://external.api.recraft.ai/v1"; const ENV_VAR: &str = "RECRAFT_API_TOKEN"; @@ -210,7 +226,8 @@ fn build_generation_body( /// - If neither is present, returns an empty [`ImageOutputs::Base64`]. async fn extract_images( response: &Value, - abort_signal: Option, + retries: &retry::PreparedRetries, + abort_signal: Option, ) -> Result { let items = response.get("data").and_then(|d| d.as_array()); @@ -240,22 +257,23 @@ async fn extract_images( let mut binaries = Vec::with_capacity(urls.len()); for url in &urls { - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: url.clone(), - headers: vec![], - body: HttpBody::Empty, - - abort_signal: abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - binaries.push(resp.body.to_vec()); + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: vec![], + + abort_signal: abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_binary_response_handler(), + recraft_failed_response_handler(), + ) + }) + .await?; + binaries.push(resp.value.to_vec()); } Ok(ImageOutputs::Binary(binaries)) } @@ -293,26 +311,30 @@ impl ImageModel for RecraftImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generations_endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + recraft_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - let value: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let value: Value = resp.value; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); - let images = extract_images(&value, options.abort_signal.clone()).await?; + let images = extract_images(&value, &retries, options.abort_signal.clone()).await?; Ok(ImageResult { images, diff --git a/aimux-providers/src/replay.rs b/aimux-providers/src/replay.rs index 9d10bcba..1ee92f10 100644 --- a/aimux-providers/src/replay.rs +++ b/aimux-providers/src/replay.rs @@ -21,7 +21,6 @@ use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; use aimux_core::recording::ProviderRecord; -use aimux_provider_utils::RetryConfig; use crate::openai::{OpenAICompatProfile, OpenAIConfig, OpenAIProvider}; use crate::provider::ProviderOptions; @@ -83,10 +82,7 @@ pub fn rebuild_provider( config = config.with_project(project); } if let Some(max_retries) = opts.max_retries { - config = config.with_retry_config(RetryConfig { - max_retries, - ..RetryConfig::default() - }); + config.retry_config.max_retries = max_retries; } if let Some(overrides) = opts.body_overrides { config = config.with_body_overrides(overrides); @@ -440,6 +436,7 @@ mod tests { status: aimux_core::recording::OutcomeStatus::Success, finish_reason: Some("stop".into()), error: None, + error_value: None, usage: None, }, ); diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs index b8ca85a2..82253168 100644 --- a/aimux-providers/src/replicate.rs +++ b/aimux-providers/src/replicate.rs @@ -13,12 +13,23 @@ use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; +use aimux_core::retry; use aimux_core::shared::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn replicate_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("detail") + .or_else(|| data.get("error")) + .and_then(Value::as_str) + .unwrap_or("Unknown Replicate error") + .to_string(), + provider_code: None, + } + }) +} /// Configuration for the Replicate provider. #[derive(Debug, Clone)] @@ -271,24 +282,28 @@ impl ImageModel for ReplicateImageModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + replicate_failed_response_handler(), ) .await?; - let rh = resp.headers; - let rb: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers.unwrap_or_default(); + let rb: Value = resp.value; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Extract output (string or array of strings) let urls: Vec = match &rb["output"] { @@ -303,22 +318,23 @@ impl ImageModel for ReplicateImageModel { // Download images let mut downloaded: Vec> = Vec::new(); for url in &urls { - let ir = send( - HttpRequest { - method: HttpMethod::Get, - url: url.clone(), - headers: vec![], - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - downloaded.push(ir.body.to_vec()); + let ir = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: vec![], + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_binary_response_handler(), + replicate_failed_response_handler(), + ) + }) + .await?; + downloaded.push(ir.value.to_vec()); } Ok(ImageResult { @@ -340,7 +356,8 @@ impl ImageModel for ReplicateImageModel { // ════════════════════════════════════════════════════════════════════════════ use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoModel, VideoResponse, VideoResult, + VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoModel, VideoOperationStart, + VideoOperationStatus, VideoResponse, VideoResult, }; /// Replicate video generation model — implements `VideoModel`. @@ -348,7 +365,8 @@ use aimux_core::video_model::{ /// Aligned with Vercel AI SDK `ReplicateVideoModel` /// (`reference/ai/packages/replicate/src/replicate-video-model.ts`). /// -/// Uses the predictions API: POST to create, GET to poll. +/// Uses the predictions API: POST to create (`do_start`); Core drives the +/// status polling via `do_status`. pub struct ReplicateVideoModel { model_id: String, config: ReplicateConfig, @@ -405,7 +423,10 @@ impl VideoModel for ReplicateVideoModel { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { let warnings: Vec = Vec::new(); let mut input = Map::new(); @@ -428,23 +449,23 @@ impl VideoModel for ReplicateVideoModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit prediction. - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: format!("{}/predictions", self.config.base_url), - headers: header_list.clone(), - body: HttpBody::Json(body), + headers: header_list, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + replicate_failed_response_handler(), ) .await?; - let prediction: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let prediction: Value = resp.value; let prediction_id = prediction .get("id") .and_then(|v| v.as_str()) @@ -453,51 +474,75 @@ impl VideoModel for ReplicateVideoModel { })? .to_string(); - // Poll for completion. - let mut raw_body: Value; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; - - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: format!("{}/predictions/{}", self.config.base_url, prediction_id), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; - let status_str = raw_body - .get("status") - .and_then(|v| v.as_str()) - .unwrap_or(""); - match status_str { - "succeeded" => break, - "failed" | "canceled" => { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: Some(status_str.to_string()), - message: format!("Replicate prediction {status_str}"), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } - _ => {} + Ok(VideoOperationStart { + operation: json!({ "prediction_id": prediction_id }), + warnings, + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: response_headers, + }, + }) + } + + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let prediction_id = operation + .get("prediction_id") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "replicate operation reference is missing prediction_id".to_string(), + ) + })?; + + let headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let poll_url = format!("{}/predictions/{}", self.config.base_url, prediction_id); + + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list, + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + replicate_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let raw_body = resp.value; + let status_str = raw_body + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or(""); + match status_str { + "succeeded" => {} + // A terminally failed prediction must be a non-retryable error, + // not Pending, so the Core poll loop stops immediately. + "failed" | "canceled" => { + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(status_str.to_string()), + response_body, + ..ApiCallError::new( + format!("Replicate prediction {status_str}"), + poll_url, + serde_json::json!({}), + ) + }))); } + // starting / processing / unknown — keep polling. + _ => return Ok(VideoOperationStatus::Pending), } // Extract video from output. @@ -525,15 +570,15 @@ impl VideoModel for ReplicateVideoModel { )); } - Ok(VideoResult { + Ok(VideoOperationStatus::Completed(VideoResult { videos, - warnings, + warnings: Vec::new(), provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), headers: Some(response_headers), }, - }) + })) } } diff --git a/aimux-providers/src/revai.rs b/aimux-providers/src/revai.rs index 0c2dfdfd..29fe0282 100644 --- a/aimux-providers/src/revai.rs +++ b/aimux-providers/src/revai.rs @@ -16,17 +16,37 @@ use serde_json::{Value, json}; use aimux_core::error::AiMuxError; use aimux_core::error::ApiCallError; +use aimux_core::retry; use aimux_core::shared::Warning; use aimux_core::transcription_model::{ AudioInput, TranscriptionCallOptions, TranscriptionModel, TranscriptionRequest, TranscriptionResponse, TranscriptionResult, TranscriptionSegment, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, MultipartForm, RetryConfig, load_api_key, - media_type_to_extension, send, sleep_or_abort, without_trailing_slash, + HttpRequest, MultipartForm, load_api_key, media_type_to_extension, sleep_or_abort, + without_trailing_slash, }; +fn revai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error"); + aimux_provider_utils::ProviderErrorParts { + message: error + .and_then(|value| value.get("message")) + .and_then(Value::as_str) + .unwrap_or("Rev.ai request failed") + .to_string(), + provider_code: error.and_then(|value| value.get("code")).and_then( + |value| match value { + Value::String(s) => Some(s.clone()), + Value::Number(n) => Some(n.to_string()), + _ => None, + }, + ), + } + }) +} + // ── Config ────────────────────────────────────────────────────────────────── #[derive(Debug, Clone)] @@ -204,40 +224,42 @@ impl TranscriptionModel for RevaiTranscriptionModel { let mut form = MultipartForm::new(); form.file("media", &filename, &options.media_type, &audio_bytes)?; form.text("config", &config)?; - let (body_bytes, content_type) = form.finish(); - let headers = self.build_headers(options.headers.as_ref()); + let submit_url = self.jobs_url(); // Submit job. - let resp = send( + let resp = aimux_provider_utils::post_form_data_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.jobs_url(), + url: submit_url.clone(), headers: headers .iter() .map(|(k, v)| (k.clone(), v.clone())) .collect(), - body: HttpBody::Bytes(body_bytes, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + form, + aimux_provider_utils::create_json_response_handler::(), + revai_failed_response_handler(), ) .await?; - let submit_response: RevaiJobResponse = serde_json::from_slice(&resp.body)?; + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let submit_response: RevaiJobResponse = resp.value; if submit_response.status.as_deref() == Some("failed") { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some("failed".to_string()), - message: "Failed to submit transcription job to Rev.ai".to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + response_body, + ..ApiCallError::new( + "Failed to submit transcription job to Rev.ai", + submit_url, + serde_json::json!({}), + ) + }))); } let job_id = submit_response.id.ok_or_else(|| { @@ -246,6 +268,11 @@ impl TranscriptionModel for RevaiTranscriptionModel { ) })?; let submission_language = submit_response.language; + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); // Poll for completion. let job_status: RevaiJobResponse; @@ -256,66 +283,69 @@ impl TranscriptionModel for RevaiTranscriptionModel { ) .await?; - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: self.job_status_url(&job_id), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - - let poll: RevaiJobResponse = serde_json::from_slice(&resp.body)?; + let poll_url = self.job_status_url(&job_id); + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + revai_failed_response_handler(), + ) + }) + .await?; + + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let poll: RevaiJobResponse = resp.value; if poll.status.as_deref() == Some("transcribed") { job_status = poll; break; } if poll.status.as_deref() == Some("failed") { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), provider_code: Some("failed".to_string()), - message: "Transcription job failed".to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); + response_body, + ..ApiCallError::new("Transcription job failed", poll_url, serde_json::json!({})) + }))); } } let _ = job_status; // Fetch transcript. - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: self.transcript_url(&job_id), - headers: headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: self.transcript_url(&job_id), + headers: headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + revai_failed_response_handler(), + ) + }) + .await?; - let response_headers = resp.headers; - let raw_body: Value = serde_json::from_slice(&resp.body)?; - let parsed: RevaiTranscriptResponse = serde_json::from_value(raw_body.clone())?; + let response_headers = resp.response_headers.unwrap_or_default(); + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; // Process monologues to extract segments and text. let mut segments: Vec = Vec::new(); diff --git a/aimux-providers/src/runwayml.rs b/aimux-providers/src/runwayml.rs index 477123de..1f2d126a 100644 --- a/aimux-providers/src/runwayml.rs +++ b/aimux-providers/src/runwayml.rs @@ -18,16 +18,11 @@ use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; use aimux_core::provider::Provider; -use aimux_core::shared::Warning; use aimux_core::video_model::{ VideoCallOptions, VideoData, VideoFile, VideoFileData, VideoFrameType, VideoModel, - VideoResponse, VideoResult, -}; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, sleep_or_abort, - without_trailing_slash, + VideoOperationStart, VideoOperationStatus, VideoPollConfig, VideoResponse, VideoResult, }; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; const PROVIDER_NAME: &str = "runwayml"; const DEFAULT_BASE_URL: &str = "https://api.dev.runwayml.com"; @@ -35,10 +30,18 @@ const ENV_VAR: &str = "RUNWAYML_API_SECRET"; const RUNWAY_VERSION: &str = "2024-11-06"; /// RunwayML returns errors as a flat `{"error": ""}` object. -const RUNWAYML_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error"], - type_path: &[], -}; +fn runwayml_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("error") + .and_then(Value::as_str) + .unwrap_or("RunwayML request failed") + .to_string(), + provider_code: None, + } + }) +} // ── Config ────────────────────────────────────────────────────────────────── @@ -208,9 +211,17 @@ impl VideoModel for RunwaymlVideoModel { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { - let warnings: Vec = Vec::new(); + fn poll_config(&self) -> VideoPollConfig { + VideoPollConfig { + interval: self.config.poll_interval, + timeout: self.config.timeout, + } + } + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { // Image-to-video when an input image (or a first frame) is provided. let image_input: Option = options .image @@ -256,112 +267,118 @@ impl VideoModel for RunwaymlVideoModel { let submit_url = format!("{}{submit_path}", self.config.base_url); // Submit the task. - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: submit_url, - headers: header_list.clone(), - body: HttpBody::Json(Value::Object(body)), + headers: header_list, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &RUNWAYML_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + runwayml_failed_response_handler(), ) .await?; - let task: RunwaymlTaskCreationResponse = serde_json::from_slice(&resp.body)?; - let task_id = task.id; - - // Poll for completion. - let poll_url = format!("{}/v1/tasks/{}", self.config.base_url, task_id); - let deadline = tokio::time::Instant::now() + self.config.timeout; - // Assigned before the only `break` path that reaches the code after the - // loop, so they are always initialized before being read. - let mut response_headers: HashMap; - let final_output: Vec; - - loop { - if tokio::time::Instant::now() >= deadline { - return Err(AiMuxError::Timeout(format!( - "{PROVIDER_NAME} task {task_id} polling timed out after {:?}", - self.config.timeout - ))); - } - - sleep_or_abort(self.config.poll_interval, options.abort_signal.as_ref()).await?; - - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: poll_url.clone(), - headers: header_list.clone(), - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &RUNWAYML_ERROR_STRUCTURE, - ) - .await?; - - response_headers = resp.headers; - - let task_details: RunwaymlTaskDetailsResponse = serde_json::from_slice(&resp.body)?; - - let status_str = task_details.status.clone().unwrap_or_default(); - match status_str.as_str() { - "SUCCEEDED" => { - final_output = - task_details - .output - .filter(|o| !o.is_empty()) - .ok_or_else(|| { - AiMuxError::InvalidResponseData(format!( - "{PROVIDER_NAME} task {task_id} succeeded without output" - )) - })?; - break; - } - "FAILED" | "CANCELLED" => { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: Some(status_str.clone()), - message: format!( - "{PROVIDER_NAME} task {task_id} failed with status {status_str}" - ), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } - // PENDING / THROTTLED / RUNNING / unknown — keep polling. - _ => continue, - } - } - - let videos: Vec = final_output - .into_iter() - .map(|url| VideoData::Url { - url, - media_type: "video/mp4".to_string(), - }) - .collect(); - - let timestamp = chrono::Utc::now().to_rfc3339(); + let response_headers = resp.response_headers; + let task: RunwaymlTaskCreationResponse = resp.value; - Ok(VideoResult { - videos, - warnings, + Ok(VideoOperationStart { + operation: json!({ "task_id": task.id }), + warnings: Vec::new(), provider_metadata: None, response: VideoResponse { - timestamp: Some(timestamp), + timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), - headers: Some(response_headers), + headers: response_headers, }, }) } + + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let task_id = operation + .get("task_id") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument(format!( + "{PROVIDER_NAME} operation reference is missing task_id" + )) + })?; + + let headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers.into_iter().collect(); + let poll_url = format!("{}/v1/tasks/{}", self.config.base_url, task_id); + + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list, + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + runwayml_failed_response_handler(), + ) + .await?; + + let response_headers = resp.response_headers.clone().unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let task_details: RunwaymlTaskDetailsResponse = resp.value; + + let status_str = task_details.status.clone().unwrap_or_default(); + match status_str.as_str() { + "SUCCEEDED" => { + let final_output = + task_details + .output + .filter(|o| !o.is_empty()) + .ok_or_else(|| { + AiMuxError::InvalidResponseData(format!( + "{PROVIDER_NAME} task {task_id} succeeded without output" + )) + })?; + + let videos: Vec = final_output + .into_iter() + .map(|url| VideoData::Url { + url, + media_type: "video/mp4".to_string(), + }) + .collect(); + + Ok(VideoOperationStatus::Completed(VideoResult { + videos, + warnings: Vec::new(), + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: Some(response_headers), + }, + })) + } + // A terminally failed task must be a non-retryable error, not + // Pending, so the Core poll loop stops immediately. + "FAILED" | "CANCELLED" => Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: Some(status_str.clone()), + response_body, + ..ApiCallError::new( + format!("{PROVIDER_NAME} task {task_id} failed with status {status_str}"), + poll_url, + serde_json::json!({}), + ) + }))), + // PENDING / THROTTLED / RUNNING / unknown — keep polling. + _ => Ok(VideoOperationStatus::Pending), + } + } } diff --git a/aimux-providers/src/searxng.rs b/aimux-providers/src/searxng.rs index 7993c5be..d6fe6042 100644 --- a/aimux-providers/src/searxng.rs +++ b/aimux-providers/src/searxng.rs @@ -15,14 +15,12 @@ use serde::Deserialize; use serde_json::Value; use aimux_core::error::AiMuxError; -use aimux_core::error::ApiCallError; use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; +use aimux_provider_utils::HttpRequest; use aimux_provider_utils::without_trailing_slash; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; /// Fixed model ID for the SearXNG search model. const MODEL_ID: &str = "searxng-search"; @@ -164,36 +162,29 @@ impl SearchModel for SearxngSearchModel { }) .unwrap_or_default(); - let mut url = url::Url::parse(&self.endpoint()).map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("invalid searxng endpoint: {e}"), - ..Default::default() - }) - })?; + let mut url = url::Url::parse(&self.endpoint()) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid searxng endpoint: {e}")))?; url.query_pairs_mut() .append_pair("q", &options.query) .append_pair("format", "json"); - let resp = send( + let resp = aimux_provider_utils::get_from_api( HttpRequest { - method: HttpMethod::Get, url: url.to_string(), headers, - body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + aimux_provider_utils::create_status_code_error_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: SearxngResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; Ok(SearchResult { results: map_results(data.results), diff --git a/aimux-providers/src/serper.rs b/aimux-providers/src/serper.rs index 8c3496e8..14b78c6e 100644 --- a/aimux-providers/src/serper.rs +++ b/aimux-providers/src/serper.rs @@ -15,10 +15,25 @@ use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn serper_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} const MODEL_ID: &str = "serper-search"; @@ -161,24 +176,23 @@ impl SearchModel for SerperSearchModel { let body = build_request_body(options); let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + serper_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let parsed: SerperResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value; + let parsed: SerperResponse = resp.value; Ok(SearchResult { results: map_results(parsed.organic), @@ -187,7 +201,7 @@ impl SearchModel for SerperSearchModel { warnings: Vec::new(), response: Some(SearchResponse { headers: Some(response_headers), - body: Some(serde_json::from_slice(&resp.body).unwrap_or(Value::Null)), + body: response_body, }), }) } diff --git a/aimux-providers/src/stability.rs b/aimux-providers/src/stability.rs index a93b5964..e0324563 100644 --- a/aimux-providers/src/stability.rs +++ b/aimux-providers/src/stability.rs @@ -12,6 +12,7 @@ use std::collections::HashMap; use async_trait::async_trait; +use bytes::Bytes; use serde_json::Value; use aimux_core::Provider; @@ -20,10 +21,8 @@ use aimux_core::image_model::{ ImageCallOptions, ImageModel, ImageOutputs, ImageResponse, ImageResult, }; use aimux_core::shared::Warning; -use aimux_provider_utils::response::ErrorStructure; use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, MultipartForm, RetryConfig, load_api_key, send, - without_trailing_slash, + HttpBody, HttpRequest, MultipartForm, load_api_key, without_trailing_slash, }; /// Stability error response structure: `{ "id": "...", "name": "...", "errors": ["..."] }`. @@ -31,10 +30,70 @@ use aimux_provider_utils::{ /// The human-readable summary lives in the `name` field (e.g. `"unauthorized"`, /// `"bad_request"`); the detailed messages are in the `errors` array, which the /// shared error parser cannot index into, so we surface `name` as the message. -const STABILITY_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["name"], - type_path: &["name"], -}; +fn stability_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let name = data + .get("name") + .and_then(Value::as_str) + .unwrap_or("Stability request failed"); + aimux_provider_utils::ProviderErrorParts { + message: name.to_string(), + provider_code: Some(name.to_string()), + } + }) +} + +enum StabilitySuccess { + Json(Value), + Binary(Bytes), +} + +fn stability_successful_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let headers = aimux_provider_utils::extract_response_headers::extract_response_headers( + input.response.headers(), + ); + let is_json = input + .response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.starts_with("application/json")); + let bytes = + aimux_provider_utils::read_response_with_size_limit::read_response_with_size_limit( + input.response, + &input.url, + &input.request_body_values, + aimux_provider_utils::read_response_with_size_limit::DEFAULT_MAX_DOWNLOAD_SIZE, + input.abort_signal.as_ref(), + ) + .await?; + let value = if is_json { + let value = serde_json::from_slice(&bytes).map_err(|error| { + AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + response_headers: Some(headers.clone()), + response_body: Some(String::from_utf8_lossy(&bytes).into_owned()), + ..aimux_core::ApiCallError::new( + format!("Invalid JSON response: {error}"), + input.url, + input.request_body_values, + ) + })) + })?; + StabilitySuccess::Json(value) + } else { + StabilitySuccess::Binary(bytes) + }; + Ok(aimux_provider_utils::ResponseHandlerOutput { + value, + raw_value: None, + response_headers: Some(headers), + }) + }) +} /// Map a Stability model ID to its generate-endpoint sub-path. /// @@ -267,37 +326,34 @@ impl ImageModel for StabilityImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let resp = send( + let resp = aimux_provider_utils::post_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Bytes(body_bytes, content_type), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &STABILITY_ERROR_STRUCTURE, + HttpBody::Bytes(body_bytes, content_type), + stability_successful_response_handler(), + stability_failed_response_handler(), ) .await?; - let rh = resp.headers; - let response_content_type = rh.get("content-type").cloned().unwrap_or_default(); - - // The image is returned either as raw binary (Accept: image/*) or as - // base64 inside a JSON body (Accept: application/json). - let image_bytes = if response_content_type.starts_with("application/json") { - let v: Value = serde_json::from_slice(&resp.body)?; - let b64 = v.get("image").and_then(|i| i.as_str()).ok_or_else(|| { - AiMuxError::InvalidResponseData("Stability response missing `image` field".into()) - })?; - base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64).map_err( - |e| AiMuxError::InvalidResponseData(format!("invalid base64 image: {e}")), - )? - } else { - resp.body.to_vec() + let rh = resp.response_headers.unwrap_or_default(); + let image_bytes = match resp.value { + StabilitySuccess::Json(v) => { + let b64 = v.get("image").and_then(|i| i.as_str()).ok_or_else(|| { + AiMuxError::InvalidResponseData( + "Stability response missing `image` field".into(), + ) + })?; + base64::Engine::decode(&base64::engine::general_purpose::STANDARD, b64).map_err( + |e| AiMuxError::InvalidResponseData(format!("invalid base64 image: {e}")), + )? + } + StabilitySuccess::Binary(bytes) => bytes.to_vec(), }; Ok(ImageResult { diff --git a/aimux-providers/src/tavily.rs b/aimux-providers/src/tavily.rs index 4b072c86..1dd39c55 100644 --- a/aimux-providers/src/tavily.rs +++ b/aimux-providers/src/tavily.rs @@ -14,18 +14,24 @@ use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; const MODEL_ID: &str = "tavily-search"; /// Tavily-specific error structure: `{ "detail": { "error": "..." } }`. -const TAVILY_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["detail", "error"], - type_path: &[], -}; +fn tavily_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("detail") + .and_then(|detail| detail.get("error")) + .and_then(Value::as_str) + .unwrap_or("Tavily request failed") + .to_string(), + provider_code: None, + } + }) +} /// Configuration for the Tavily provider. #[derive(Debug, Clone)] @@ -185,24 +191,23 @@ impl SearchModel for TavilySearchModel { let body = build_request_body(options); let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &TAVILY_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + tavily_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let parsed: TavilyResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value; + let parsed: TavilyResponse = resp.value; Ok(SearchResult { results: map_results(parsed.results), @@ -211,7 +216,7 @@ impl SearchModel for TavilySearchModel { warnings: Vec::new(), response: Some(SearchResponse { headers: Some(response_headers), - body: Some(serde_json::from_slice(&resp.body).unwrap_or(Value::Null)), + body: response_body, }), }) } diff --git a/aimux-providers/src/tinyfish.rs b/aimux-providers/src/tinyfish.rs index 34039a21..6530f1ea 100644 --- a/aimux-providers/src/tinyfish.rs +++ b/aimux-providers/src/tinyfish.rs @@ -13,17 +13,32 @@ use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; use aimux_core::shared::SharedHeaders; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn tinyfish_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} /// Provider canonical name. const PROVIDER_NAME: &str = "tinyfish"; @@ -204,38 +219,31 @@ impl SearchModel for TinyfishSearchModel { .into_iter() .collect(); - let mut url = url::Url::parse(&self.endpoint()).map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("invalid tinyfish endpoint: {e}"), - ..Default::default() - }) - })?; + let mut url = url::Url::parse(&self.endpoint()) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid tinyfish endpoint: {e}")))?; url.query_pairs_mut() .append_pair("query", &options.query) .append_pair("count", &count.to_string()); - let resp = send( + let resp = aimux_provider_utils::get_from_api( HttpRequest { - method: HttpMethod::Get, url: url.to_string(), headers, - body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + tinyfish_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: TinyfishSearchResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let results: Vec = data.results.into_iter().map(map_result).collect(); diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs index 0d90ebac..ac14bd5d 100644 --- a/aimux-providers/src/vertex/anthropic_model.rs +++ b/aimux-providers/src/vertex/anthropic_model.rs @@ -20,18 +20,14 @@ use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send_stream_timed, send_timed, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::{HttpRequest, RetryConfig}; use crate::anthropic::convert::{build_request_body_with_warnings, parse_stop_reason}; use crate::anthropic::stream::stream_parts_for_result_block; @@ -46,13 +42,6 @@ const ANTHROPIC_VERTEX_VERSION: &str = "vertex-2023-10-16"; /// Google/Vertex error structure: `{ "error": { "message": "...", "status": "..." } }`. /// -/// `rawPredict` HTTP-level errors come back in the Vertex AI error shape; the -/// success body is Anthropic-shaped and is parsed separately. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// Configuration for an Anthropic-on-Vertex model instance. /// /// `base_url` is the Vertex AI base URL *without* a `/publishers/{publisher}` @@ -67,13 +56,13 @@ pub struct VertexAnthropicConfig { pub auth: VertexAuth, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`（max_retries=2）。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } /// An Anthropic Claude language model served via Vertex AI `rawPredict`. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct VertexAnthropicModel { model_id: String, @@ -153,6 +142,10 @@ impl LanguageModel for VertexAnthropicModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; // M2b: record identity + credential source + auth kind. Never serialize @@ -183,30 +176,22 @@ impl LanguageModel for VertexAnthropicModel { let body = self.wrap_raw_predict_body(req.body); let url = self.generate_endpoint(); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url, + url: url.clone(), headers, - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &GOOGLE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), ) .await?; - let data: AnthropicResponse = - serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let data: AnthropicResponse = resp.value; let content = crate::anthropic::stream::parse_anthropic_content( &data.content, @@ -248,36 +233,43 @@ impl LanguageModel for VertexAnthropicModel { let tool_names = ToolNameMapping::new(options.tools.as_deref()); let url = self.stream_endpoint(); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url, + url: url.clone(), headers, - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &GOOGLE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + crate::google::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let sse_stream = SseStream::new(resp.body); + let response_headers = resp.response_headers.unwrap_or_default(); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(StreamEvent::Error { error })) = first_event.as_ref() { + return Err(crate::anthropic::stream::anthropic_stream_error( + error, + &url, + body.clone(), + response_headers, + )); + } + let stream_error_url = url; + let stream_request_body = body.clone(); + let stream_response_headers = response_headers.clone(); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings }); - let mut sse = sse_stream; + let mut sse = futures::stream::iter(first_event.into_iter()).chain(sse_stream); let mut blocks: HashMap = HashMap::new(); let mut final_usage = Usage::default(); let mut final_finish_reason: Option = None; @@ -286,9 +278,9 @@ impl LanguageModel for VertexAnthropicModel { while let Some(event) = sse.next().await { match event { - Ok(sse_event) => { - match serde_json::from_str::(&sse_event.data) { - Ok(StreamEvent::MessageStart { message }) => { + Ok(stream_event) => { + match stream_event { + StreamEvent::MessageStart { message } => { if let Some(usage) = &message.usage { // RFC-0015 P0-2: full input side incl. // cache fields + raw (same as Anthropic). @@ -304,7 +296,7 @@ impl LanguageModel for VertexAnthropicModel { response_meta_emitted = true; } } - Ok(StreamEvent::ContentBlockStart { index, content_block }) => { + StreamEvent::ContentBlockStart { index, content_block } => { match content_block { ContentBlock::Text { .. } => { blocks.insert(index, BlockState::Text { started: false }); @@ -378,7 +370,7 @@ impl LanguageModel for VertexAnthropicModel { } } } - Ok(StreamEvent::ContentBlockDelta { index, delta }) => { + StreamEvent::ContentBlockDelta { index, delta } => { if let Some(text) = delta.text { let start_id: Option = match blocks.get_mut(&index) { Some(BlockState::Text { started: false }) => { @@ -445,7 +437,7 @@ impl LanguageModel for VertexAnthropicModel { }); } } - Ok(StreamEvent::ContentBlockStop { index }) => { + StreamEvent::ContentBlockStop { index } => { if let Some(state) = blocks.remove(&index) { match state { BlockState::Text { started: true } => { @@ -489,7 +481,7 @@ impl LanguageModel for VertexAnthropicModel { } } } - Ok(StreamEvent::MessageDelta { delta, usage }) => { + StreamEvent::MessageDelta { delta, usage } => { if let Some(reason) = delta.stop_reason { final_finish_reason = Some(parse_stop_reason(&reason)); } @@ -506,26 +498,27 @@ impl LanguageModel for VertexAnthropicModel { .output_tokens; } } - Ok(StreamEvent::MessageStop) => break, - Ok(StreamEvent::Error { error }) => { + StreamEvent::MessageStop => break, + StreamEvent::Error { error } => { yield Ok(StreamPart::Error { - error: AiMuxError::ApiCall(ApiCallError { - provider_code: error.error_type, - message: error.message, - response_body: Some(sse_event.data.clone()), - ..Default::default() - }), + error: crate::anthropic::stream::anthropic_stream_error( + &error, + &stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), }); return; } - Ok(_) | Err(_) => {} + _ => {} } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - return; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/vertex/embedding.rs b/aimux-providers/src/vertex/embedding.rs index 0a90e69d..45bb6dd7 100644 --- a/aimux-providers/src/vertex/embedding.rs +++ b/aimux-providers/src/vertex/embedding.rs @@ -13,27 +13,22 @@ use std::collections::HashMap; use async_trait::async_trait; use serde_json::{Map, Value, json}; +use crate::google::google_failed_response_handler; + use aimux_core::embedding_model::{ EmbeddingCallOptions, EmbeddingModel, EmbeddingResponse, EmbeddingResult, EmbeddingUsage, }; use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::VertexAuth; use super::model::VertexConfig; -/// Google-specific error structure: `{ "error": { "message": "..." } }`. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// A Google Vertex AI embedding model (e.g. `"textembedding-gecko@001"`). /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct VertexEmbeddingModel { model_id: String, @@ -81,6 +76,10 @@ impl EmbeddingModel for VertexEmbeddingModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn max_embeddings_per_call(&self) -> Option { if uses_embed_content_endpoint(&self.model_id) { Some(1) @@ -142,25 +141,24 @@ impl EmbeddingModel for VertexEmbeddingModel { self.config.base_url, self.model_id ); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let raw_value: Value = resp.value; let embedding = raw_value .get("embedding") @@ -224,25 +222,24 @@ impl EmbeddingModel for VertexEmbeddingModel { let url = format!("{}/models/{}:predict", self.config.base_url, self.model_id); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_value: Value = serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let raw_value: Value = resp.value; // Batch: response.predictions[].embeddings.values let (embeddings, total_tokens): (Vec>, u32) = raw_value diff --git a/aimux-providers/src/vertex/image.rs b/aimux-providers/src/vertex/image.rs index a4ab3c47..c0d88d23 100644 --- a/aimux-providers/src/vertex/image.rs +++ b/aimux-providers/src/vertex/image.rs @@ -12,6 +12,8 @@ use std::collections::HashMap; use async_trait::async_trait; use serde_json::{Map, Value, json}; +use crate::google::google_failed_response_handler; + use aimux_core::error::AiMuxError; use aimux_core::image_model::{ ImageCallOptions, ImageFile, ImageFileData, ImageModel, ImageOutputs, ImageResponse, @@ -19,23 +21,17 @@ use aimux_core::image_model::{ }; use aimux_core::shared::Warning; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; use super::{VertexAuth, VertexConfig}; -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - fn is_gemini_model(model_id: &str) -> bool { model_id.starts_with("gemini-") } /// A Google Vertex AI image generation model. /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct VertexImageModel { model_id: String, @@ -197,23 +193,22 @@ impl VertexImageModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.predict_endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), ) .await?; - let rh = resp.headers; - let rb: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers.unwrap_or_default(); + let rb: Value = resp.value; let images: Vec = rb .get("predictions") @@ -342,23 +337,22 @@ impl VertexImageModel { let body = json!({ "contents": [{ "role": "user", "parts": parts }], "generationConfig": Value::Object(gc) }); let headers = self.build_headers(options.headers.as_ref()); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generate_content_endpoint(), headers, - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + google_failed_response_handler(), ) .await?; - let rh = resp.headers; - let rb: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers.unwrap_or_default(); + let rb: Value = resp.value; let mut images: Vec = Vec::new(); if let Some(candidates) = rb.get("candidates").and_then(|c| c.as_array()) { @@ -428,6 +422,9 @@ impl ImageModel for VertexImageModel { fn model_id(&self) -> &str { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } fn max_images_per_call(&self) -> Option { if is_gemini_model(&self.model_id) { Some(10) diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs index a5d3ef6c..0745a693 100644 --- a/aimux-providers/src/vertex/mod.rs +++ b/aimux-providers/src/vertex/mod.rs @@ -55,8 +55,7 @@ pub struct VertexProviderConfig { pub auth: VertexAuth, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`（max_retries=2）。 - /// 取代之前硬编码的 `RetryConfig::default()`,让 per-call `max_retries` 生效。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } @@ -176,7 +175,7 @@ fn build_base_url(project: &str, location: &str, _endpoint: bool) -> String { /// Google Vertex AI provider — creates [`VertexModel`] instances. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct VertexProvider { config: VertexProviderConfig, @@ -303,7 +302,8 @@ impl VertexProvider { location, self.config.auth.clone(), self.config.base_url.clone(), - )) + ) + .with_retry_config(self.config.retry_config)) } /// Create a video generation model instance for the given Vertex AI model @@ -331,7 +331,8 @@ impl VertexProvider { location, self.config.auth.clone(), self.config.base_url.clone(), - )) + ) + .with_retry_config(self.config.retry_config)) } } @@ -372,24 +373,23 @@ impl Provider for VertexProvider { } headers.push(("Content-Type".to_string(), "application/json".to_string())); - use aimux_provider_utils::{ - DEFAULT_ERROR_STRUCTURE, HttpBody, HttpMethod, HttpRequest, send_timed, - }; - let resp = send_timed( - HttpRequest { - method: HttpMethod::Get, - url, - headers, - body: HttpBody::Empty, - abort_signal: None, - call_id: None, - recording_context: None, - }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - None, - ) - .await?; + use aimux_provider_utils::HttpRequest; + // Retry rationale: see `openai::model::execute_list_models`. + let resp = aimux_core::retry::prepare_retries(None, retry_config, None) + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: url.clone(), + headers: headers.clone(), + abort_signal: None, + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { @@ -402,7 +402,7 @@ impl Provider for VertexProvider { #[serde(default)] display_name: Option, } - let parsed: Resp = serde_json::from_slice(&resp.body)?; + let parsed: Resp = resp.value; let runtime: Vec = parsed .models .into_iter() diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs index 335d8e72..64230c28 100644 --- a/aimux-providers/src/vertex/model.rs +++ b/aimux-providers/src/vertex/model.rs @@ -17,25 +17,15 @@ use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, send_stream_timed, send_timed, -}; -use aimux_stream::SseStream; +use aimux_provider_utils::{HttpRequest, RetryConfig}; use crate::google::convert::{ build_request_body, convert_usage, extract_sources, parse_finish_reason, }; -use crate::google::types::{Candidate, GenerateContentResponse, GoogleUsageMetadata, StreamChunk}; +use crate::google::types::{Candidate, GenerateContentResponse, GoogleStreamEvent}; use super::VertexAuth; -/// Google-specific error structure: `{ "error": { "message": "..." } }`. -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// Configuration for a Vertex model instance (cloned from the provider). #[derive(Debug, Clone)] pub struct VertexConfig { @@ -43,13 +33,13 @@ pub struct VertexConfig { pub auth: VertexAuth, /// 凭证来源(RFC-0023):`None` = explicit;`Some("env:VAR")` = 环境变量。 pub api_key_source: Option, - /// 重试配置(M1b)。默认 `RetryConfig::default()`（max_retries=2）。 + /// Retry settings used by Core model operations. pub retry_config: RetryConfig, } /// A Google Vertex AI language model. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct VertexModel { model_id: String, @@ -133,6 +123,10 @@ impl LanguageModel for VertexModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { use aimux_core::recording::ProviderRecord; // M2b: record identity + credential source + auth kind. Never serialize @@ -160,32 +154,24 @@ impl LanguageModel for VertexModel { async fn do_generate(&self, options: &CallOptions) -> Result { let body = build_request_body(&self.model_id, options); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.generate_endpoint(), headers, - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &GOOGLE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: GenerateContentResponse = - serde_json::from_slice(&resp.body).map_err(AiMuxError::from)?; + let data: GenerateContentResponse = resp.value; let candidate = data.candidates.into_iter().next().ok_or_else(|| { AiMuxError::InvalidResponseData("no candidates in response".to_string()) @@ -240,38 +226,47 @@ impl LanguageModel for VertexModel { async fn do_stream(&self, options: &CallOptions) -> Result { let body = build_request_body(&self.model_id, options); let headers = self.build_headers(options.headers.as_ref()); - let retry_config = crate::openai::model::resolve_retry_config( - &self.config.retry_config, - options.max_retries, - ); - - let resp = send_stream_timed( + let endpoint = self.stream_endpoint(); + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.stream_endpoint(), + url: endpoint.clone(), headers, - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - retry_config, - &GOOGLE_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + aimux_provider_utils::create_event_source_response_handler::(), + crate::google::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); // Same source as the non-stream path: the response `Date` header. let response_timestamp = response_headers.get("date").cloned(); - let sse_stream = SseStream::new(resp.body); + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(GoogleStreamEvent::Error(error))) = first_event.as_ref() { + return Err(crate::google::google_stream_error( + &error.error, + &endpoint, + body.clone(), + response_headers, + )); + } + let stream_error_url = endpoint; + let stream_request_body = body.clone(); + let stream_error_headers = response_headers.clone(); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings: vec![] }); - let mut sse_stream = sse_stream; + let mut sse_stream = futures::stream::iter(first_event.into_iter()).chain(sse_stream); let mut text_id: Option = None; let mut block_counter = 0usize; let mut final_usage: Usage = Usage::default(); @@ -305,28 +300,7 @@ impl LanguageModel for VertexModel { break; } match event { - Ok(sse_event) => { - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - stream_errored = true; - break; - } - }; - - let chunk: StreamChunk = match serde_json::from_value(parsed) { - Ok(c) => c, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - stream_errored = true; - break; - } - }; + Ok(GoogleStreamEvent::Chunk(chunk)) => { if !response_metadata_emitted && let Some(id) = &chunk.response_id { @@ -557,13 +531,25 @@ impl LanguageModel for VertexModel { Some(parse_finish_reason(reason, has_tool_calls)); } } - Err(e) => { + Ok(GoogleStreamEvent::Error(error)) => { yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), + error: crate::google::google_stream_error( + &error.error, + &stream_error_url, + stream_request_body.clone(), + stream_error_headers.clone(), + ), }); stream_errored = true; break; } + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } + } } } @@ -662,7 +648,3 @@ fn extract_content_from_candidate(candidate: &Candidate) -> (Vec Self { + self.retry_config = retry_config; + self + } + fn auth_header(&self) -> HashMap { let mut headers = HashMap::new(); match &self.auth { @@ -188,6 +188,10 @@ impl TranscriptionModel for VertexTranscriptionModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.retry_config + } + async fn do_generate( &self, options: &TranscriptionCallOptions, @@ -252,27 +256,25 @@ impl TranscriptionModel for VertexTranscriptionModel { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url, headers: header_list, - body: HttpBody::Json(request_body.clone()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + request_body.clone(), + aimux_provider_utils::create_json_response_handler::(), + crate::google::google_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); + let response_headers = resp.response_headers.unwrap_or_default(); - let parsed: GoogleVertexResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; let results = parsed.results.unwrap_or_default(); diff --git a/aimux-providers/src/vertex/video.rs b/aimux-providers/src/vertex/video.rs index ee06cd40..85bc8063 100644 --- a/aimux-providers/src/vertex/video.rs +++ b/aimux-providers/src/vertex/video.rs @@ -5,7 +5,7 @@ //! //! Uses the Long Running Operations API: //! 1. POST `{base_url}/models/{model}:predictLongRunning` → returns operation name -//! 2. GET operation to poll until `done: true` +//! 2. GET operation — polled by Core via `do_status` until `done: true` //! 3. Return video URL(s) use std::collections::HashMap; @@ -14,24 +14,18 @@ use async_trait::async_trait; use serde_json::{Map, Value, json}; use aimux_core::error::{AiMuxError, ApiCallError}; -use aimux_core::shared::Warning; use aimux_core::video_model::{ - VideoCallOptions, VideoData, VideoModel, VideoResponse, VideoResult, + VideoCallOptions, VideoData, VideoModel, VideoOperationStart, VideoOperationStatus, + VideoResponse, VideoResult, }; -use aimux_provider_utils::response::ErrorStructure; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send, sleep_or_abort}; +use aimux_provider_utils::{HttpRequest, RetryConfig}; use super::VertexAuth; -const GOOGLE_ERROR_STRUCTURE: ErrorStructure = ErrorStructure { - message_path: &["error", "message"], - type_path: &["error", "status"], -}; - /// A Google Vertex AI video generation model. /// -/// Does **not** hold an HTTP client — `http::send` uses the process-wide shared +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the process-wide shared /// `Client` internally (RFC-0009 §4.1). pub struct VertexVideoModel { model_id: String, @@ -39,6 +33,7 @@ pub struct VertexVideoModel { location: String, auth: VertexAuth, base_url: String, + retry_config: RetryConfig, } impl VertexVideoModel { @@ -56,9 +51,15 @@ impl VertexVideoModel { location, auth, base_url, + retry_config: RetryConfig::default(), } } + pub(crate) fn with_retry_config(mut self, retry_config: RetryConfig) -> Self { + self.retry_config = retry_config; + self + } + fn build_headers(&self, extra: Option<&HashMap>) -> HashMap { let mut h = HashMap::new(); match &self.auth { @@ -113,13 +114,17 @@ impl VideoModel for VertexVideoModel { fn model_id(&self) -> &str { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.retry_config + } fn max_videos_per_call(&self) -> Option { Some(1) } - async fn do_generate(&self, options: &VideoCallOptions) -> Result { - let warnings: Vec = Vec::new(); - + async fn do_start( + &self, + options: &VideoCallOptions, + ) -> Result { let mut instances = vec![json!({"prompt": options.prompt})]; if let Some(ref image) = options.image && let aimux_core::video_model::VideoFile::Url { url, .. } = image @@ -152,23 +157,23 @@ impl VideoModel for VertexVideoModel { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.predict_url(), headers: header_list.clone(), - body: HttpBody::Json(body), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), ) .await?; - let predict_response: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let predict_response: Value = resp.value; let operation_name = predict_response .get("name") .and_then(|v| v.as_str()) @@ -179,55 +184,75 @@ impl VideoModel for VertexVideoModel { })? .to_string(); - // Poll for completion. - let mut raw_body: Value; - let mut response_headers: HashMap; - loop { - sleep_or_abort( - std::time::Duration::from_millis(100), - options.abort_signal.as_ref(), - ) - .await?; + Ok(VideoOperationStart { + operation: json!({ "operation_name": operation_name }), + warnings: Vec::new(), + provider_metadata: None, + response: VideoResponse { + timestamp: Some(chrono::Utc::now().to_rfc3339()), + model_id: Some(self.model_id.clone()), + headers: response_headers, + }, + }) + } - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: self.operation_url(&operation_name), - headers: header_list.clone(), - body: HttpBody::Empty, + async fn do_status( + &self, + operation: &Value, + options: &VideoCallOptions, + ) -> Result { + let operation_name = operation + .get("operation_name") + .and_then(Value::as_str) + .ok_or_else(|| { + AiMuxError::InvalidArgument( + "google.vertex operation reference is missing operation_name".to_string(), + ) + })?; - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, - ) - .await?; + let headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + + let poll_url = self.operation_url(operation_name); + let resp = aimux_provider_utils::get_from_api( + HttpRequest { + url: poll_url.clone(), + headers: header_list, + + abort_signal: options.abort_signal.clone(), + call_id: None, + recording_context: None, + }, + aimux_provider_utils::create_json_response_handler::(), + crate::google::google_failed_response_handler(), + ) + .await?; - response_headers = resp.headers; - raw_body = serde_json::from_slice(&resp.body)?; - // Check the in-band error first: a terminal response may carry both - // done:true and an error object (provider-declared failure). - if let Some(err) = raw_body.get("error") { - let msg = err - .get("message") + let response_headers = resp.response_headers.unwrap_or_default(); + let response_body = resp.raw_value.as_ref().map(ToString::to_string); + let raw_body: Value = resp.value; + // Check the in-band error first: a terminal response may carry both + // done:true and an error object (provider-declared failure). + if let Some(err) = raw_body.get("error") { + let msg = err + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or("Unknown error"); + return Err(AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(200), + provider_code: err + .get("status") .and_then(|v| v.as_str()) - .unwrap_or("Unknown error"); - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: err - .get("status") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: msg.to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } - if raw_body.get("done").and_then(serde_json::Value::as_bool) == Some(true) { - break; - } + .map(std::string::ToString::to_string), + response_body, + ..ApiCallError::new(msg, poll_url, serde_json::json!({})) + }))); + } + if raw_body.get("done").and_then(serde_json::Value::as_bool) != Some(true) { + return Ok(VideoOperationStatus::Pending); } let videos: Vec = raw_body @@ -255,15 +280,15 @@ impl VideoModel for VertexVideoModel { )); } - Ok(VideoResult { + Ok(VideoOperationStatus::Completed(VideoResult { videos, - warnings, + warnings: Vec::new(), provider_metadata: None, response: VideoResponse { timestamp: Some(chrono::Utc::now().to_rfc3339()), model_id: Some(self.model_id.clone()), headers: Some(response_headers), }, - }) + })) } } diff --git a/aimux-providers/src/voyage/embedding.rs b/aimux-providers/src/voyage/embedding.rs index 7c8c3f95..a78c313f 100644 --- a/aimux-providers/src/voyage/embedding.rs +++ b/aimux-providers/src/voyage/embedding.rs @@ -19,10 +19,9 @@ use aimux_core::embedding_model::{ use aimux_core::error::AiMuxError; use aimux_core::shared::SharedProviderOptions; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; -use super::VoyageConfig; +use super::{VoyageConfig, voyage_failed_response_handler}; /// A Voyage embedding model (e.g. `"voyage-3.5"`). pub struct VoyageEmbeddingModel { @@ -105,25 +104,25 @@ impl EmbeddingModel for VoyageEmbeddingModel { // error internally using the shared error structure. `HttpBody::Json` // sets `Content-Type: application/json`, so it is intentionally not // added to the header list above. - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(Value::Object(body)), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + Value::Object(body), + aimux_provider_utils::create_json_response_handler(), + voyage_failed_response_handler(), ) .await?; - let response_headers: HashMap = resp.headers.clone(); + let response_headers: HashMap = + resp.response_headers.unwrap_or_default().clone(); - let raw_value: Value = serde_json::from_slice::(&resp.body)?; + let raw_value: Value = resp.value; // Extract embeddings: sort data by index, then map to embedding arrays. let embeddings: Vec> = raw_value diff --git a/aimux-providers/src/voyage/mod.rs b/aimux-providers/src/voyage/mod.rs index c5157ac0..7ac94c98 100644 --- a/aimux-providers/src/voyage/mod.rs +++ b/aimux-providers/src/voyage/mod.rs @@ -13,6 +13,20 @@ use aimux_core::error::AiMuxError; use aimux_core::provider::Provider; use aimux_provider_utils::{load_api_key, without_trailing_slash}; +pub(crate) fn voyage_failed_response_handler() -> aimux_provider_utils::ResponseHandler +{ + aimux_provider_utils::create_json_error_response_handler(|data| { + aimux_provider_utils::ProviderErrorParts { + message: data + .get("detail") + .and_then(serde_json::Value::as_str) + .unwrap_or("Voyage request failed") + .to_owned(), + provider_code: None, + } + }) +} + /// Configuration for the Voyage AI provider. #[derive(Debug, Clone)] pub struct VoyageConfig { diff --git a/aimux-providers/src/voyage/reranking.rs b/aimux-providers/src/voyage/reranking.rs index acc39278..f74a9f8a 100644 --- a/aimux-providers/src/voyage/reranking.rs +++ b/aimux-providers/src/voyage/reranking.rs @@ -16,10 +16,9 @@ use aimux_core::reranking_model::{ }; use aimux_core::types::Warning; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send}; +use aimux_provider_utils::HttpRequest; -use super::VoyageConfig; +use super::{VoyageConfig, voyage_failed_response_handler}; /// Voyage provider-specific reranking options. #[derive(Debug, Clone, Default)] @@ -148,28 +147,26 @@ impl RerankingModel for VoyageRerankingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: header_list, - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body.clone(), + aimux_provider_utils::create_json_response_handler::(), + voyage_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; + let response_headers = resp.response_headers.unwrap_or_default(); - let raw_body: Value = serde_json::from_slice(&resp.body)?; - - let data: VoyageRerankingResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let ranking: Vec = data .data diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs index d00075a4..fdd623c0 100644 --- a/aimux-providers/src/xai/mod.rs +++ b/aimux-providers/src/xai/mod.rs @@ -18,7 +18,9 @@ pub use responses::XaiResponsesModel; use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::provider::Provider; -use aimux_provider_utils::{RetryConfig, load_api_key}; +use aimux_provider_utils::load_api_key; +use serde::de::DeserializeOwned; +use serde_json::Value; use crate::openai::OpenAIConfig; @@ -26,6 +28,188 @@ const DEFAULT_BASE_URL: &str = "https://api.x.ai/v1"; const ENV_VAR: &str = "XAI_API_KEY"; const PROVIDER_NAME: &str = "xai"; +pub(crate) fn xai_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + if let (Some(code), Some(message)) = ( + data.get("code").and_then(Value::as_str), + data.get("error").and_then(Value::as_str), + ) { + return aimux_provider_utils::ProviderErrorParts { + message: format!("{code}: {message}"), + provider_code: Some(code.to_owned()), + }; + } + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("code") + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }), + } + }) +} + +pub(crate) fn xai_stream_error( + event: &Value, + url: &str, + request_body_values: Value, + response_headers: std::collections::HashMap, +) -> AiMuxError { + let error = event.get("error").unwrap_or(event); + let message = error + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) + .or_else(|| event.get("message").and_then(Value::as_str)) + .unwrap_or("xAI stream failed"); + let status_code = event + .get("status") + .or_else(|| event.get("code")) + .or_else(|| error.get("status")) + .or_else(|| error.get("code")) + .and_then(Value::as_u64) + .and_then(|status| u16::try_from(status).ok()) + .filter(|status| (400..=599).contains(status)); + let provider_code = event + .get("code") + .or_else(|| error.get("code")) + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + aimux_provider_utils::stream_error_api_call( + message, + provider_code, + status_code, + event, + url, + request_body_values, + response_headers, + ) +} + +pub(crate) fn xai_successful_response_handler() -> aimux_provider_utils::ResponseHandler +where + T: DeserializeOwned + Send + 'static, +{ + aimux_provider_utils::ResponseHandler::new(|input| async move { + let status = input.response.status().as_u16(); + let url = input.url.clone(); + let request_body_values = input.request_body_values.clone(); + let output = aimux_provider_utils::create_json_response_handler::() + .handle(input) + .await?; + let headers = output.response_headers.clone().unwrap_or_default(); + let raw = output.value; + if let Some(error) = raw.get("error").filter(|value| !value.is_null()) { + let message = error + .as_str() + .or_else(|| error.get("message").and_then(Value::as_str)) + .unwrap_or("xAI request failed"); + let provider_code = raw + .get("code") + .or_else(|| error.get("code")) + .or_else(|| error.get("type")) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + return Err(AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + provider_code, + response_body: Some(raw.to_string()), + response_headers: Some(headers), + ..aimux_core::ApiCallError::new(message, url, request_body_values) + }))); + } + let value = serde_json::from_value(raw.clone()).map_err(|error| { + AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + response_body: Some(raw.to_string()), + response_headers: Some(headers.clone()), + ..aimux_core::ApiCallError::new( + format!("Invalid JSON response: {error}"), + url, + request_body_values, + ) + })) + })?; + Ok(aimux_provider_utils::ResponseHandlerOutput { + value, + raw_value: Some(raw), + response_headers: output.response_headers, + }) + }) +} + +/// xAI occasionally returns a JSON error document with a successful status +/// where an SSE response was requested. Classify that response before handing +/// the body to the standard typed event-source handler. +pub(crate) fn xai_event_source_response_handler() +-> aimux_provider_utils::ResponseHandler>> +where + T: DeserializeOwned + Send + 'static, +{ + aimux_provider_utils::ResponseHandler::new(|input| async move { + let is_json = input + .response + .headers() + .get(reqwest::header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .is_some_and(|value| value.contains("application/json")); + if !is_json { + return aimux_provider_utils::create_event_source_response_handler::() + .handle(input) + .await; + } + + let status = input.response.status().as_u16(); + let url = input.url.clone(); + let request_body_values = input.request_body_values.clone(); + let output = aimux_provider_utils::create_json_response_handler::() + .handle(input) + .await?; + let headers = output.response_headers.unwrap_or_default(); + let raw = output.value; + let error = raw.get("error").filter(|value| !value.is_null()); + let message = error + .and_then(|value| { + value + .as_str() + .or_else(|| value.get("message").and_then(Value::as_str)) + }) + .unwrap_or("Expected an event stream but received JSON"); + let provider_code = raw + .get("code") + .or_else(|| error.and_then(|value| value.get("code"))) + .or_else(|| error.and_then(|value| value.get("type"))) + .and_then(|value| match value { + Value::String(value) => Some(value.clone()), + Value::Number(value) => Some(value.to_string()), + _ => None, + }); + Err(AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError { + status_code: Some(status), + provider_code, + response_body: Some(raw.to_string()), + response_headers: Some(headers), + ..aimux_core::ApiCallError::new(message, url, request_body_values) + }))) + }) + .streaming() +} + /// Configuration for the xAI provider (wraps [`OpenAIConfig`]). #[derive(Clone)] pub struct XAIConfig(OpenAIConfig); @@ -74,16 +258,11 @@ impl XAIConfig { pub(crate) fn base_url(&self) -> &str { &self.0.base_url } - - /// Get the retry config. - pub(crate) fn retry_config(&self) -> RetryConfig { - self.0.retry_config - } } /// xAI provider — creates [`XaiModel`] instances pointed at xAI. /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct XAIProvider { config: XAIConfig, @@ -98,7 +277,7 @@ impl XAIProvider { /// Create a model instance for the given xAI model id (e.g. `"grok-2"`). /// /// Clones the provider config so the model inherits the same - /// `api_key_source` / `retry_config` (M2b: previously reconstructed with + /// `api_key_source` / `max_retries` (M2b: previously reconstructed with /// `XAIConfig::new`, which dropped the credential source). #[must_use] pub fn model(&self, model_id: &str) -> XaiModel { @@ -138,14 +317,14 @@ impl Provider for XAIProvider { let config = OpenAIConfig::new(self.config.api_key()) .with_base_url(self.config.base_url()) .with_provider(PROVIDER_NAME); + // The freshly built OpenAIConfig carries the default retry settings — + // use the user's configured ones from the wrapped config instead. + let retry_config = self.config.0.retry_config; Box::pin(async move { let headers = crate::openai::model::build_auth_headers(&config); - let runtime = crate::openai::model::execute_list_models( - &config.base_url, - &headers, - &config.retry_config, - ) - .await?; + let runtime = + crate::openai::model::execute_list_models(&config.base_url, &headers, retry_config) + .await?; Ok(runtime) }) } diff --git a/aimux-providers/src/xai/model.rs b/aimux-providers/src/xai/model.rs index 45885b26..a9b635aa 100644 --- a/aimux-providers/src/xai/model.rs +++ b/aimux-providers/src/xai/model.rs @@ -13,16 +13,14 @@ use async_trait::async_trait; use futures::StreamExt; use serde_json::Value; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::convert::{build_request_body_with_warnings, convert_xai_usage, parse_finish_reason}; use super::types::{XaiChatResponse, XaiStreamChunk}; @@ -40,7 +38,7 @@ fn generate_source_id() -> String { /// An xAI language model (Grok). /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct XaiModel { model_id: String, @@ -94,6 +92,10 @@ impl LanguageModel for XaiModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.openai_config().retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { // M2b: xAI wraps OpenAIConfig — reuse the OpenAI snapshot helper with // xAI's own provider name. api_key_source/profile come from the inner @@ -110,54 +112,26 @@ impl LanguageModel for XaiModel { let request_result = build_request_body_with_warnings(&self.model_id, options, false)?; let body = request_result.body; - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.retry_config(), - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + super::xai_successful_response_handler::(), + super::xai_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_value: Value = serde_json::from_slice(&resp.body)?; - - // Check for 200-status error (xAI sometimes returns errors with 200). - if let Some(error_msg) = raw_value.get("error").and_then(|v| v.as_str()) { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: raw_value - .get("code") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: error_msg.to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } + let response_headers = resp.response_headers.unwrap_or_default(); - let data: XaiChatResponse = serde_json::from_value(raw_value.clone())?; - - // Handle error field - if let Some(error_msg) = &data.error { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: data.code.clone(), - message: error_msg.clone(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } + let _raw_value = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let choice = data .choices @@ -281,79 +255,41 @@ impl LanguageModel for XaiModel { let request_result = build_request_body_with_warnings(&self.model_id, options, true)?; let body = request_result.body; let warnings = request_result.warnings; + let endpoint = self.endpoint(); - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.retry_config(), - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + super::xai_event_source_response_handler::(), + super::xai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - // Check if the response is JSON (not SSE) — xAI sometimes returns - // errors with 200 status and content-type application/json. - let content_type = response_headers - .get("content-type") - .map(std::string::String::as_str) - .unwrap_or(""); - if content_type.contains("application/json") { - // Collect the (non-SSE) JSON body and check for an error object. - let mut buf = Vec::new(); - let mut body_stream = resp.body; - while let Some(chunk) = body_stream.next().await { - if let Ok(bytes) = chunk { - buf.extend_from_slice(&bytes); - } - } - let val: Value = serde_json::from_slice(&buf)?; - if let Some(err_msg) = val.get("error").and_then(|v| v.as_str()) { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: val - .get("code") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: err_msg.to_string(), - response_body: Some(String::from_utf8_lossy(&buf).into_owned()), - ..Default::default() - })); - } - return Err(AiMuxError::InvalidResponseData( - "expected SSE stream but got JSON response without an error object".to_string(), - )); - } + let response_headers = resp.response_headers.unwrap_or_default(); - let mut sse_stream = SseStream::new(resp.body); + let mut sse_stream = resp.value; // Peek at the first SSE event to detect early errors. - let first_event = sse_stream.next().await; - if let Some(Ok(ref event)) = first_event - && let Ok(val) = serde_json::from_str::(&event.data) + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(ref val)) = first_event + && val.get("error").is_some_and(|value| !value.is_null()) { - // Check for error in the first chunk (200-status error). - if let Some(err_msg) = val.get("error").and_then(|v| v.as_str()) { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: val - .get("code") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: err_msg.to_string(), - response_body: Some(event.data.clone()), - ..Default::default() - })); - } + return Err(super::xai_stream_error( + val, + &endpoint, + body.clone(), + response_headers.clone(), + )); } let messages_for_dedup = body @@ -361,6 +297,9 @@ impl LanguageModel for XaiModel { .and_then(|m| m.as_array()) .cloned() .unwrap_or_default(); + let stream_error_url = endpoint.clone(); + let stream_request_body = body.clone(); + let stream_response_headers = response_headers.clone(); let stream = async_stream::stream! { // First part: StreamStart. @@ -384,33 +323,16 @@ impl LanguageModel for XaiModel { while let Some(event) = event_iter.next().await { match event { - Ok(sse_event) => { - if sse_event.data == "[DONE]" { - break; - } - - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - break; - } - }; - + Ok(parsed) => { // Check for mid-stream error. - if let Some(err_obj) = parsed.get("error") { - let msg = err_obj - .as_str() - .unwrap_or("Unknown stream error") - .to_string(); + if parsed.get("error").is_some() { yield Ok(StreamPart::Error { - error: AiMuxError::ApiCall(ApiCallError { - message: msg, - response_body: Some(sse_event.data.clone()), - ..Default::default() - }), + error: super::xai_stream_error( + &parsed, + &stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), }); break; } @@ -418,10 +340,8 @@ impl LanguageModel for XaiModel { let chunk: XaiStreamChunk = match serde_json::from_value(parsed) { Ok(c) => c, Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - break; + yield Err(AiMuxError::from(e)); + continue; } }; @@ -616,11 +536,12 @@ impl LanguageModel for XaiModel { } } } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs index d2a6f065..a0c8a543 100644 --- a/aimux-providers/src/xai/responses/mod.rs +++ b/aimux-providers/src/xai/responses/mod.rs @@ -22,16 +22,14 @@ use async_trait::async_trait; use futures::StreamExt; use serde_json::{Value, json}; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::options::CallOptions; use aimux_core::result::{GenerateContent, GenerateResult, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::types::{FinishReason, FinishReasonUnified, ResponseMetadata, Usage}; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, send_stream_timed, send_timed}; -use aimux_stream::SseStream; +use aimux_provider_utils::HttpRequest; use super::super::XAIConfig; use crate::openai::responses::responses_convert::build_header_list; @@ -49,7 +47,7 @@ fn generate_source_id() -> String { /// An xAI Responses language model (Grok). /// -/// Does **not** hold an HTTP client — `http::send` / `http::send_stream` use the +/// Does **not** hold an HTTP client — the `aimux-provider-utils` API helpers use the /// process-wide shared `Client` internally (RFC-0009 §4.1). pub struct XaiResponsesModel { model_id: String, @@ -91,6 +89,10 @@ impl LanguageModel for XaiResponsesModel { &self.model_id } + fn retry_config(&self) -> aimux_core::retry::RetryConfig { + self.config.openai_config().retry_config + } + fn config_snapshot(&self) -> aimux_core::recording::ProviderRecord { // M2b: reuse the OpenAI snapshot helper with xAI's provider name. crate::openai::config_snapshot_from_config( @@ -106,42 +108,25 @@ impl LanguageModel for XaiResponsesModel { let body = request_result.body; let provider_tool_names = request_result.provider_tool_names; - let resp = send_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, url: self.endpoint(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.retry_config(), - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + super::xai_successful_response_handler::(), + super::xai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - let raw_value: Value = serde_json::from_slice(&resp.body)?; - - // Check for 200-status error. - if let Some(error_msg) = raw_value.get("error").and_then(|v| v.as_str()) { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: raw_value - .get("code") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: error_msg.to_string(), - response_body: Some(String::from_utf8_lossy(&resp.body).into_owned()), - ..Default::default() - })); - } + let response_headers = resp.response_headers.unwrap_or_default(); - let data: types::XaiResponsesResponse = serde_json::from_value(raw_value.clone())?; + let _raw_value = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let mut content: Vec = Vec::new(); let mut has_function_call = false; @@ -385,59 +370,43 @@ impl LanguageModel for XaiResponsesModel { let body = request_result.body; let warnings = request_result.warnings; let provider_tool_names = request_result.provider_tool_names; + let endpoint = self.endpoint(); - let resp = send_stream_timed( + let resp = aimux_provider_utils::post_json_to_api( HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), + url: endpoint.clone(), headers: build_header_list(&headers), - body: HttpBody::Json(body.clone()), abort_signal: options.abort_signal.clone(), call_id: options.call_id.clone(), recording_context: options.recording_context.clone(), }, - self.config.retry_config(), - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + body.clone(), + super::xai_event_source_response_handler::(), + super::xai_failed_response_handler(), ) .await?; - let response_headers = resp.headers; - - // Check for JSON error (200-status). - let content_type = response_headers - .get("content-type") - .map(std::string::String::as_str) - .unwrap_or(""); - if content_type.contains("application/json") { - // Collect the (non-SSE) JSON body and check for an error object. - let mut buf = Vec::new(); - let mut body_stream = resp.body; - while let Some(chunk) = body_stream.next().await { - if let Ok(bytes) = chunk { - buf.extend_from_slice(&bytes); - } - } - let val: Value = serde_json::from_slice(&buf)?; - if let Some(err_msg) = val.get("error").and_then(|v| v.as_str()) { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: val - .get("code") - .and_then(|v| v.as_str()) - .map(std::string::ToString::to_string), - message: err_msg.to_string(), - response_body: Some(String::from_utf8_lossy(&buf).into_owned()), - ..Default::default() - })); - } - return Err(AiMuxError::InvalidResponseData( - "expected SSE stream but got JSON response without an error object".to_string(), + let response_headers = resp.response_headers.unwrap_or_default(); + + let mut sse_stream = resp.value; + let first_event = match sse_stream.next().await { + Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error), + first_event => first_event, + }; + if let Some(Ok(ref event)) = first_event + && types::event_type(event) == "error" + { + return Err(super::xai_stream_error( + event, + &endpoint, + body.clone(), + response_headers.clone(), )); } - - let sse_stream = SseStream::new(resp.body); + let stream_error_url = endpoint; + let stream_request_body = body.clone(); + let stream_response_headers = response_headers.clone(); let stream = async_stream::stream! { yield Ok(StreamPart::StreamStart { warnings }); @@ -459,25 +428,11 @@ impl LanguageModel for XaiResponsesModel { // Track ongoing function calls by output_index. let mut ongoing_tool_calls: HashMap = HashMap::new(); // output_index -> (tool_call_id, tool_name) - let mut event_iter = sse_stream; + let mut event_iter = futures::stream::iter(first_event.into_iter()).chain(sse_stream); while let Some(event) = event_iter.next().await { match event { - Ok(sse_event) => { - if sse_event.data == "[DONE]" { - break; - } - - let parsed: Value = match serde_json::from_str(&sse_event.data) { - Ok(v) => v, - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::from(e), - }); - break; - } - }; - + Ok(parsed) => { let event_type = types::event_type(&parsed); // ── response.created / response.in_progress ── @@ -665,14 +620,13 @@ impl LanguageModel for XaiResponsesModel { // ── error event ── if event_type == "error" { - let message = parsed.get("message").and_then(|v| v.as_str()).unwrap_or("Unknown error"); yield Ok(StreamPart::Error { - error: AiMuxError::ApiCall(ApiCallError { - provider_code: parsed.get("code").and_then(|v| v.as_str()).map(std::string::ToString::to_string), - message: message.to_string(), - response_body: Some(sse_event.data.clone()), - ..Default::default() - }), + error: super::xai_stream_error( + &parsed, + &stream_error_url, + stream_request_body.clone(), + stream_response_headers.clone(), + ), }); continue; } @@ -942,11 +896,12 @@ impl LanguageModel for XaiResponsesModel { // All other event types (web_search_call.in_progress, etc.) are ignored. } - Err(e) => { - yield Ok(StreamPart::Error { - error: AiMuxError::InvalidResponseData(e.to_string()), - }); - break; + Err(error) => { + let recoverable = error.is_recoverable_stream_error(); + yield Err(error); + if !recoverable { + return; + } } } } diff --git a/aimux-providers/src/you_com.rs b/aimux-providers/src/you_com.rs index 25dd5cdd..6e58dc7e 100644 --- a/aimux-providers/src/you_com.rs +++ b/aimux-providers/src/you_com.rs @@ -14,17 +14,32 @@ use async_trait::async_trait; use serde::Deserialize; use serde_json::Value; -use aimux_core::error::{AiMuxError, ApiCallError}; +use aimux_core::error::AiMuxError; use aimux_core::provider::Provider; use aimux_core::search_model::{ SearchCallOptions, SearchModel, SearchResponse, SearchResult, SearchResultItem, }; use aimux_core::shared::SharedHeaders; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{ - HttpBody, HttpMethod, HttpRequest, RetryConfig, load_api_key, send, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +fn you_com_failed_response_handler() -> aimux_provider_utils::ResponseHandler { + aimux_provider_utils::create_json_error_response_handler(|data| { + let error = data.get("error").unwrap_or(data); + aimux_provider_utils::ProviderErrorParts { + message: error + .get("message") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + provider_code: error + .get("type") + .or_else(|| error.get("code")) + .and_then(Value::as_str) + .map(str::to_owned), + } + }) +} /// Provider canonical name. const PROVIDER_NAME: &str = "you_com"; @@ -206,38 +221,31 @@ impl SearchModel for YouComSearchModel { .into_iter() .collect(); - let mut url = url::Url::parse(&self.endpoint()).map_err(|e| { - AiMuxError::ApiCall(ApiCallError { - message: format!("invalid you_com endpoint: {e}"), - ..Default::default() - }) - })?; + let mut url = url::Url::parse(&self.endpoint()) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid you_com endpoint: {e}")))?; url.query_pairs_mut() .append_pair("query", &options.query) .append_pair("count", &count.to_string()); - let resp = send( + let resp = aimux_provider_utils::get_from_api( HttpRequest { - method: HttpMethod::Get, url: url.to_string(), headers, - body: HttpBody::Empty, abort_signal: options.abort_signal.clone(), call_id: None, recording_context: None, }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + you_com_failed_response_handler(), ) .await?; // Capture response headers. - let response_headers = resp.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers.unwrap_or_default(); - let data: YoucomSearchResponse = serde_json::from_value(raw_body.clone())?; + let raw_body = resp.raw_value.unwrap_or(Value::Null); + let data = resp.value; let results: Vec = data.results.into_iter().map(map_result).collect(); diff --git a/aimux-providers/tests/anthropic_aws_model_test.rs b/aimux-providers/tests/anthropic_aws_model_test.rs index 318fea40..b612f692 100644 --- a/aimux-providers/tests/anthropic_aws_model_test.rs +++ b/aimux-providers/tests/anthropic_aws_model_test.rs @@ -23,8 +23,6 @@ use aimux_core::types::FinishReasonUnified; use aimux_providers::anthropic_aws::{AnthropicAwsAuth, AnthropicAwsConfig, AnthropicAwsModel}; -use aimux_provider_utils::RetryConfig; - // ── Helpers ────────────────────────────────────────────────────────────────── fn test_prompt() -> LanguageModelPrompt { @@ -48,7 +46,7 @@ fn make_model(server: &MockServer) -> AnthropicAwsModel { api_version: "2023-06-01".to_string(), workspace_id: None, api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }, ) } @@ -410,7 +408,7 @@ async fn anthropic_aws_sigv4_auth() { api_version: "2023-06-01".to_string(), workspace_id: None, api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }, ); diff --git a/aimux-providers/tests/bedrock_embedding_test.rs b/aimux-providers/tests/bedrock_embedding_test.rs index 056feace..3ccce5bd 100644 --- a/aimux-providers/tests/bedrock_embedding_test.rs +++ b/aimux-providers/tests/bedrock_embedding_test.rs @@ -41,6 +41,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -277,6 +279,8 @@ async fn should_pass_output_dimension_cohere_v4() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -358,6 +362,8 @@ async fn should_pass_nova_embedding_dimension() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); diff --git a/aimux-providers/tests/bedrock_reranking_test.rs b/aimux-providers/tests/bedrock_reranking_test.rs index 9c4a6143..8ce074a2 100644 --- a/aimux-providers/tests/bedrock_reranking_test.rs +++ b/aimux-providers/tests/bedrock_reranking_test.rs @@ -54,6 +54,8 @@ fn text_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: Some(bedrock_provider_options()), headers: None, + max_retries: None, + timeout: None, } } @@ -70,6 +72,8 @@ fn json_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: Some(bedrock_provider_options()), headers: None, + max_retries: None, + timeout: None, } } diff --git a/aimux-providers/tests/cassette_multimodal_test.rs b/aimux-providers/tests/cassette_multimodal_test.rs index 4c50b31a..bfa3b39f 100644 --- a/aimux-providers/tests/cassette_multimodal_test.rs +++ b/aimux-providers/tests/cassette_multimodal_test.rs @@ -95,6 +95,8 @@ async fn cassette_openai_embedding_documents() { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, }; let result = model.do_embed(&opts).await.expect("embed should succeed"); @@ -268,6 +270,8 @@ async fn cassette_xai_image_generation() { provider_options: std::collections::HashMap::new(), abort_signal: None, headers: None, + max_retries: None, + timeout: None, }; let result = model diff --git a/aimux-providers/tests/codex_test.rs b/aimux-providers/tests/codex_test.rs index beaf23b5..7ce2e17d 100644 --- a/aimux-providers/tests/codex_test.rs +++ b/aimux-providers/tests/codex_test.rs @@ -19,7 +19,7 @@ use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; use aimux_core::options::{CallOptions, Tool, ToolChoice}; -use aimux_core::result::StreamResult; +use aimux_core::result::{GenerateContent, StreamResult}; use aimux_core::stream_part::StreamPart; use aimux_core::tool::FunctionTool; use aimux_core::types::FinishReasonUnified; @@ -347,6 +347,43 @@ async fn subscription_generate_forces_stream_and_assembles_result() { ); } +#[tokio::test] +async fn subscription_generate_continues_after_a_malformed_sse_frame() { + let server = MockServer::start().await; + let events = vec![ + "data: {unparsable}\n\n".to_string(), + subscription_completed_event(), + ]; + mock_sse_response(&server, &sse_body_strings(&events)).await; + + let config = CodexConfig::subscription("account-token").with_base_url(server.uri()); + let result = CodexProvider::new(config) + .model("gpt-5.2-codex") + .do_generate(&default_options(test_prompt())) + .await + .expect("a later response.completed event should still be assembled"); + + assert!(matches!( + result.content.first(), + Some(GenerateContent::Text { text, .. }) if text == "answer text" + )); +} + +#[tokio::test] +async fn subscription_generate_returns_parse_error_without_a_completed_event() { + let server = MockServer::start().await; + mock_sse_response(&server, "data: {unparsable}\n\n").await; + + let config = CodexConfig::subscription("account-token").with_base_url(server.uri()); + let error = CodexProvider::new(config) + .model("gpt-5.2-codex") + .do_generate(&default_options(test_prompt())) + .await + .unwrap_err(); + + assert!(matches!(error, AiMuxError::JsonParse(_))); +} + #[tokio::test] async fn subscription_stream_forces_store_false() { let server = MockServer::start().await; @@ -509,7 +546,7 @@ async fn codex_refresh_rejects_bad_grant() { ) .await .expect_err("400 must fail"); - // parse_provider_error maps non-401/403 4xx to Provider (not retryable). + // The endpoint-specific failed-response handler keeps this 4xx non-retryable. assert!(matches!(err, AiMuxError::ApiCall(_))); assert!(!err.is_retryable()); } diff --git a/aimux-providers/tests/cohere_embedding_test.rs b/aimux-providers/tests/cohere_embedding_test.rs index 17e7dc1e..67427c7c 100644 --- a/aimux-providers/tests/cohere_embedding_test.rs +++ b/aimux-providers/tests/cohere_embedding_test.rs @@ -43,6 +43,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -191,6 +193,8 @@ async fn should_pass_input_type_setting() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -221,6 +225,8 @@ async fn should_pass_output_dimension_setting() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -256,6 +262,8 @@ async fn should_pass_headers() { abort_signal: None, provider_options: None, headers: Some(request_headers), + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); diff --git a/aimux-providers/tests/cohere_model_test.rs b/aimux-providers/tests/cohere_model_test.rs index e70f66ce..54472adc 100644 --- a/aimux-providers/tests/cohere_model_test.rs +++ b/aimux-providers/tests/cohere_model_test.rs @@ -11,6 +11,7 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::content::ContentPart; +use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; @@ -668,8 +669,14 @@ async fn should_send_streaming_request_body() { #[tokio::test] async fn should_handle_unparsable_stream_parts() { let server = MockServer::start().await; - let sse = "event: foo-message\ndata: {unparsable}\n\n"; - mock_sse_response(&server, sse).await; + let sse = cohere_sse_body(&[ + "{unparsable}", + r#"{"type":"content-start","index":0,"delta":{"message":{"content":{"type":"text","text":""}}}}"#, + r#"{"type":"content-delta","index":0,"delta":{"message":{"content":{"text":"after-error"}}}}"#, + r#"{"type":"content-end","index":0}"#, + r#"{"type":"message-end","delta":{"finish_reason":"COMPLETE","usage":{"billed_units":{"input_tokens":1,"output_tokens":1},"tokens":{"input_tokens":1,"output_tokens":1}}}}"#, + ]); + mock_sse_response(&server, &sse).await; let config = CohereConfig::new("test-api-key").with_base_url(server.uri()); let provider = CohereProvider::new(config); @@ -680,16 +687,30 @@ async fn should_handle_unparsable_stream_parts() { .await .expect("do_stream should succeed"); - let parts = collect_stream(result).await; + let outcomes: Vec<_> = result.stream.collect().await; - // Should have StreamStart, Error, Finish. assert!(matches!( - parts.first(), - Some(StreamPart::StreamStart { .. }) + outcomes.first(), + Some(Ok(StreamPart::StreamStart { .. })) )); - assert!(parts.iter().any(|p| matches!(p, StreamPart::Error { .. }))); - let finish = parts.last().expect("should have finish"); - assert!(matches!(finish, StreamPart::Finish { .. })); + let error_index = outcomes + .iter() + .position(|outcome| { + matches!( + outcome, + Err(AiMuxError::JsonParse(_) | AiMuxError::InvalidResponseData(_)) + ) + }) + .expect("malformed frame should surface as a parse error item"); + assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::TextDelta { delta, .. }) if delta == "after-error" + ))); + assert!(outcomes.iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::Finish { finish_reason, .. }) + if finish_reason.unified == FinishReasonUnified::Stop + ))); } /// TS: "should handle 401 auth error" diff --git a/aimux-providers/tests/cohere_reranking_test.rs b/aimux-providers/tests/cohere_reranking_test.rs index 92450b6d..66c31f0c 100644 --- a/aimux-providers/tests/cohere_reranking_test.rs +++ b/aimux-providers/tests/cohere_reranking_test.rs @@ -55,6 +55,8 @@ fn text_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: Some(cohere_provider_options()), headers: None, + max_retries: None, + timeout: None, } } @@ -71,6 +73,8 @@ fn json_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: Some(cohere_provider_options()), headers: None, + max_retries: None, + timeout: None, } } diff --git a/aimux-providers/tests/e2e_test.rs b/aimux-providers/tests/e2e_test.rs index 53b8ad48..3f9e6e56 100644 --- a/aimux-providers/tests/e2e_test.rs +++ b/aimux-providers/tests/e2e_test.rs @@ -7,6 +7,7 @@ //! (`google_model_test.rs` / `mistral_model_test.rs` / `cohere_model_test.rs`). use aimux_core::content::ContentPart; +use aimux_core::error::AiMuxError; use aimux_core::generate::{GenerateTextOptions, generate_text, stream_text}; use aimux_core::message::{MessageContent, ModelMessage, Role}; use aimux_core::stream_part::StreamPart; @@ -305,8 +306,13 @@ async fn e2e_anthropic_error_429() { assert!(result.is_err()); let err = result.unwrap_err(); assert!( - matches!(err, ref e if e.status_code() == Some(429)), - "expected RateLimited error, got {err:?}" + matches!( + err, + AiMuxError::Retry(ref retry) + if retry.errors.len() == 3 + && retry.errors.iter().all(|attempt| attempt.status_code() == Some(429)) + ), + "expected RetryError with three 429 attempts, got {err:?}" ); } diff --git a/aimux-providers/tests/fal_video_test.rs b/aimux-providers/tests/fal_video_test.rs index ffbb08da..dd7de0c5 100644 --- a/aimux-providers/tests/fal_video_test.rs +++ b/aimux-providers/tests/fal_video_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the Fal video model tests. +//! Rust translation of the Fal video model tests. //! Source: `reference/ai/packages/fal/src/fal-video-model.test.ts` use std::collections::HashMap; @@ -7,11 +7,20 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_core::video_model::{VideoCallOptions, generate_video}; use aimux_providers::{FalConfig, FalProvider}; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(prompt: &str) -> VideoCallOptions { - VideoCallOptions::new(prompt) + let mut o = VideoCallOptions::new(prompt); + o.poll = fast_poll(); + o } async fn mock_queue_and_result(server: &MockServer, result: &Value) { @@ -37,7 +46,9 @@ async fn should_generate_video() { let provider = FalProvider::new(config); let model = provider.video("fal-ai/kling-video"); - let r = model.do_generate(&options("A cat playing")).await.unwrap(); + let r = generate_video(&model, options("A cat playing")) + .await + .unwrap(); assert_eq!(r.videos.len(), 1); } @@ -51,7 +62,9 @@ async fn should_pass_prompt() { let provider = FalProvider::new(config); let model = provider.video("fal-ai/kling-video"); - model.do_generate(&options("A cat playing")).await.unwrap(); + generate_video(&model, options("A cat playing")) + .await + .unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); @@ -77,7 +90,7 @@ async fn should_pass_headers() { rh.insert("Custom-Request-Header".to_string(), "req-val".to_string()); opts.headers = Some(rh); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; @@ -96,7 +109,7 @@ async fn should_include_response_data() { let provider = FalProvider::new(config); let model = provider.video("fal-ai/kling-video"); - let r = model.do_generate(&options("test")).await.unwrap(); + let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!(r.response.model_id, Some("fal-ai/kling-video".to_string())); } diff --git a/aimux-providers/tests/google_embedding_test.rs b/aimux-providers/tests/google_embedding_test.rs index 31fc08df..33854bf8 100644 --- a/aimux-providers/tests/google_embedding_test.rs +++ b/aimux-providers/tests/google_embedding_test.rs @@ -42,6 +42,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -155,6 +157,8 @@ async fn should_pass_output_dimensionality() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -188,6 +192,8 @@ async fn should_pass_task_type() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -222,6 +228,8 @@ async fn should_pass_headers() { abort_signal: None, provider_options: None, headers: Some(request_headers), + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); diff --git a/aimux-providers/tests/google_video_test.rs b/aimux-providers/tests/google_video_test.rs index 4da7d0f8..a965f646 100644 --- a/aimux-providers/tests/google_video_test.rs +++ b/aimux-providers/tests/google_video_test.rs @@ -1,14 +1,34 @@ -//! Rust translation of the Google video model tests. +//! Rust translation of the Google video model tests. //! Source: `reference/ai/packages/google/src/google-video-model.test.ts` -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aimux_core::AiMuxError; +use aimux_core::video_model::{VideoCallOptions, generate_video}; use aimux_providers::{GoogleConfig, GoogleProvider}; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(p: &str) -> VideoCallOptions { - VideoCallOptions::new(p) + let mut o = VideoCallOptions::new(p); + o.poll = fast_poll(); + o +} + +fn request_header<'a>(request: &'a wiremock::Request, name: &str) -> Option<&'a str> { + request + .headers + .get(name) + .and_then(|value| value.to_str().ok()) } async fn mock_predict_and_poll(server: &MockServer, result: &Value) { @@ -37,10 +57,231 @@ async fn should_generate_video() { let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); let provider = GoogleProvider::new(config); let model = provider.video("veo-3.0-generate-001"); - let r = model.do_generate(&options("A cat")).await.unwrap(); + let r = generate_video(&model, options("A cat")).await.unwrap(); assert_eq!(r.videos.len(), 1); } +#[tokio::test] +async fn poll_retry_does_not_submit_a_second_generation() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/models/veo-3.0-generate-001:predictLongRunning")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"name": "operations/test-op"})), + ) + .mount(&server) + .await; + + let poll_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&poll_attempts); + Mock::given(method("GET")) + .and(path("/operations/test-op")) + .respond_with(move |_: &wiremock::Request| { + if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "0") + .set_body_json(json!({ + "error": {"message": "try again", "status": "UNAVAILABLE"} + })) + } else { + ResponseTemplate::new(200).set_body_json(json!({ + "done": true, + "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]} + })) + } + }) + .mount(&server) + .await; + + let config = GoogleConfig::new("test-api-key") + .with_base_url(server.uri()) + .with_retry_config(aimux_provider_utils::RetryConfig { + max_retries: 1, + ..Default::default() + }); + let provider = GoogleProvider::new(config); + let model = provider.video("veo-3.0-generate-001"); + let mut opts = options("A cat"); + opts.max_retries = Some(1); + + let result = generate_video(&model, opts).await.unwrap(); + assert_eq!(result.videos.len(), 1); + assert_eq!(poll_attempts.load(Ordering::SeqCst), 2); + + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path().ends_with(":predictLongRunning")) + .count(), + 1 + ); + assert!( + requests + .iter() + .find(|request| request.url.path().ends_with(":predictLongRunning")) + .and_then(|request| request_header(request, "idempotency-key")) + .is_some(), + "the start request should receive a generated idempotency key" + ); + assert!( + requests + .iter() + .filter(|request| request.url.path() == "/operations/test-op") + .all(|request| request_header(request, "idempotency-key").is_none()), + "the generated start key must not leak to status requests" + ); +} + +#[tokio::test] +async fn start_retry_reuses_the_same_idempotency_key() { + let server = MockServer::start().await; + let start_attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&start_attempts); + + Mock::given(method("POST")) + .and(path("/models/veo-3.0-generate-001:predictLongRunning")) + .respond_with(move |_: &wiremock::Request| { + if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "0") + .set_body_json(json!({ + "error": {"message": "try again", "status": "UNAVAILABLE"} + })) + } else { + ResponseTemplate::new(200).set_body_json(json!({"name": "operations/test-op"})) + } + }) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/operations/test-op")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "done": true, + "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]} + }))) + .mount(&server) + .await; + + let config = GoogleConfig::new("test-api-key") + .with_base_url(server.uri()) + .with_retry_config(aimux_provider_utils::RetryConfig { + max_retries: 1, + ..Default::default() + }); + let provider = GoogleProvider::new(config); + let model = provider.video("veo-3.0-generate-001"); + let mut opts = options("A cat"); + opts.max_retries = Some(1); + + generate_video(&model, opts).await.unwrap(); + + let requests = server.received_requests().await.unwrap(); + let starts: Vec<_> = requests + .iter() + .filter(|request| request.url.path().ends_with(":predictLongRunning")) + .collect(); + assert_eq!(starts.len(), 2); + let first_key = request_header(starts[0], "idempotency-key"); + assert!(first_key.is_some()); + assert_eq!(first_key, request_header(starts[1], "idempotency-key")); + assert!( + requests + .iter() + .filter(|request| request.url.path() == "/operations/test-op") + .all(|request| request_header(request, "idempotency-key").is_none()) + ); +} + +#[tokio::test] +async fn poll_deadline_reached_during_delay_does_not_issue_status_request() { + let server = MockServer::start().await; + let result = + json!({"done": true, "response": {"videos": [{"gcsUri": "gs://bucket/video.mp4"}]}}); + mock_predict_and_poll(&server, &result).await; + let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); + let provider = GoogleProvider::new(config); + let model = provider.video("veo-3.0-generate-001"); + let mut opts = options("A cat"); + opts.poll = Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(100), + timeout_ms: Some(10), + }); + + let error = generate_video(&model, opts).await.unwrap_err(); + + assert!(matches!(error, AiMuxError::Timeout(_))); + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path().ends_with(":predictLongRunning")) + .count(), + 1 + ); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/operations/test-op") + .count(), + 0 + ); +} + +#[tokio::test] +async fn poll_retry_exhaustion_does_not_submit_a_second_generation() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/models/veo-3.0-generate-001:predictLongRunning")) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"name": "operations/test-op"})), + ) + .mount(&server) + .await; + + Mock::given(method("GET")) + .and(path("/operations/test-op")) + .respond_with( + ResponseTemplate::new(503) + .insert_header("retry-after-ms", "0") + .set_body_json(json!({ + "error": {"message": "still unavailable", "status": "UNAVAILABLE"} + })), + ) + .mount(&server) + .await; + + let config = GoogleConfig::new("test-api-key") + .with_base_url(server.uri()) + .with_retry_config(aimux_provider_utils::RetryConfig { + max_retries: 1, + ..Default::default() + }); + let provider = GoogleProvider::new(config); + let model = provider.video("veo-3.0-generate-001"); + + let error = generate_video(&model, options("A cat")).await.unwrap_err(); + assert!(matches!(error, AiMuxError::Retry(_))); + + let requests = server.received_requests().await.unwrap(); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path().ends_with(":predictLongRunning")) + .count(), + 1 + ); + assert_eq!( + requests + .iter() + .filter(|request| request.url.path() == "/operations/test-op") + .count(), + 2 + ); +} + #[tokio::test] async fn should_pass_prompt() { let server = MockServer::start().await; @@ -50,7 +291,7 @@ async fn should_pass_prompt() { let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); let provider = GoogleProvider::new(config); let model = provider.video("veo-3.0-generate-001"); - model.do_generate(&options("A cat")).await.unwrap(); + generate_video(&model, options("A cat")).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["instances"][0]["prompt"], "A cat"); @@ -68,7 +309,7 @@ async fn should_pass_aspect_ratio_and_duration() { let mut opts = options("test"); opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); opts.duration = Some(5); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["parameters"]["aspectRatio"], "16:9"); @@ -84,7 +325,7 @@ async fn should_include_response_data() { let config = GoogleConfig::new("test-api-key").with_base_url(server.uri()); let provider = GoogleProvider::new(config); let model = provider.video("veo-3.0-generate-001"); - let r = model.do_generate(&options("test")).await.unwrap(); + let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!( r.response.model_id, diff --git a/aimux-providers/tests/huggingface_responses_test.rs b/aimux-providers/tests/huggingface_responses_test.rs index d52742ef..cd5420d0 100644 --- a/aimux-providers/tests/huggingface_responses_test.rs +++ b/aimux-providers/tests/huggingface_responses_test.rs @@ -24,6 +24,7 @@ use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::content::ContentPart; +use aimux_core::error::AiMuxError; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; @@ -757,6 +758,15 @@ async fn should_handle_streaming_errors() { r#"{"type":"response.output_item.added","output_index":0,"item":{"id":"msg_test","type":"message","role":"assistant"},"sequence_number":1}"#, ), "data:invalid json}\n\n".to_string(), + sse_event( + r#"{"type":"response.output_text.delta","item_id":"msg_test","output_index":0,"content_index":0,"delta":"after-error","sequence_number":2}"#, + ), + sse_event( + r#"{"type":"response.output_item.done","output_index":0,"item":{"id":"msg_test","type":"message","role":"assistant","status":"completed"},"sequence_number":3}"#, + ), + sse_event( + r#"{"type":"response.completed","response":{"id":"resp_test","status":"completed","incomplete_details":null,"usage":null},"sequence_number":4}"#, + ), ]); mock_sse(&server, chunks).await; @@ -768,21 +778,25 @@ async fn should_handle_streaming_errors() { .await .expect("should succeed"); - let parts = collect_stream(result).await; - - let has_error = parts.iter().any(|p| matches!(p, StreamPart::Error { .. })); - assert!(has_error, "should have an Error part"); - - let finish = parts + let outcomes: Vec<_> = result.stream.collect().await; + let error_index = outcomes .iter() - .find(|p| matches!(p, StreamPart::Finish { .. })) - .expect("should have a Finish part"); - match finish { - StreamPart::Finish { finish_reason, .. } => { - assert_eq!(finish_reason.unified, FinishReasonUnified::Error); - } - _ => unreachable!(), - } + .position(|outcome| { + matches!( + outcome, + Err(AiMuxError::JsonParse(_) | AiMuxError::InvalidResponseData(_)) + ) + }) + .expect("malformed frame should surface as a parse error item"); + assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::TextDelta { delta, .. }) if delta == "after-error" + ))); + assert!(outcomes.iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::Finish { finish_reason, .. }) + if finish_reason.unified == FinishReasonUnified::Stop + ))); } /// TS: doStream › "should send correct streaming request" diff --git a/aimux-providers/tests/jina_ai_test.rs b/aimux-providers/tests/jina_ai_test.rs index d5e38965..7729fd34 100644 --- a/aimux-providers/tests/jina_ai_test.rs +++ b/aimux-providers/tests/jina_ai_test.rs @@ -68,6 +68,8 @@ fn text_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -84,6 +86,8 @@ fn json_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } diff --git a/aimux-providers/tests/klingai_video_test.rs b/aimux-providers/tests/klingai_video_test.rs index 765eee0a..2b08ddde 100644 --- a/aimux-providers/tests/klingai_video_test.rs +++ b/aimux-providers/tests/klingai_video_test.rs @@ -1,4 +1,4 @@ -//! Rust translation of the KlingAI video model tests. +//! Rust translation of the KlingAI video model tests. //! Source: `reference/ai/packages/klingai/src/klingai-video-model.test.ts` use std::collections::HashMap; @@ -7,11 +7,20 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_core::video_model::{VideoCallOptions, VideoModel, generate_video}; use aimux_providers::{KlingAIConfig, KlingAIProvider}; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(prompt: &str) -> VideoCallOptions { - VideoCallOptions::new(prompt) + let mut o = VideoCallOptions::new(prompt); + o.poll = fast_poll(); + o } async fn mock_task_and_result( @@ -62,7 +71,9 @@ async fn should_generate_video_from_prompt() { let provider = KlingAIProvider::new(config); let model = provider.video("kling-v2.1-master-t2v"); - let result = model.do_generate(&options("A cat playing")).await.unwrap(); + let result = generate_video(&model, options("A cat playing")) + .await + .unwrap(); assert_eq!(result.videos.len(), 1); match &result.videos[0] { @@ -83,7 +94,9 @@ async fn should_pass_model_name_and_prompt() { let provider = KlingAIProvider::new(config); let model = provider.video("kling-v2.1-master-t2v"); - model.do_generate(&options("A cat playing")).await.unwrap(); + generate_video(&model, options("A cat playing")) + .await + .unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); @@ -110,7 +123,7 @@ async fn should_pass_headers() { rh.insert("Custom-Request-Header".to_string(), "req-val".to_string()); opts.headers = Some(rh); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; @@ -134,7 +147,7 @@ async fn should_pass_seed_duration_and_aspect_ratio() { opts.duration = Some(5); opts.aspect_ratio = Some(aimux_core::shared::AspectRatio::new(16, 9)); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); @@ -153,7 +166,7 @@ async fn should_include_response_data() { let provider = KlingAIProvider::new(config); let model = provider.video("kling-v2.1-master-t2v"); - let result = model.do_generate(&options("test")).await.unwrap(); + let result = generate_video(&model, options("test")).await.unwrap(); assert!(result.response.timestamp.is_some()); assert_eq!( diff --git a/aimux-providers/tests/list_models_test.rs b/aimux-providers/tests/list_models_test.rs index f716ea8a..67a7326d 100644 --- a/aimux-providers/tests/list_models_test.rs +++ b/aimux-providers/tests/list_models_test.rs @@ -174,7 +174,12 @@ async fn list_models_malformed_response() { let config = OpenAIConfig::new("test-key").with_base_url(format!("{}/v1", server.uri())); let provider = OpenAIProvider::new(config); let err = provider.list_models().await.unwrap_err(); - assert!(matches!(err, aimux_core::AiMuxError::JsonParse(_))); + assert!(matches!( + err, + aimux_core::AiMuxError::ApiCall(ref detail) + if detail.status_code == Some(200) + && detail.message.starts_with("Invalid JSON response:") + )); } #[tokio::test] diff --git a/aimux-providers/tests/mistral_embedding_test.rs b/aimux-providers/tests/mistral_embedding_test.rs index 7b694c38..7d594075 100644 --- a/aimux-providers/tests/mistral_embedding_test.rs +++ b/aimux-providers/tests/mistral_embedding_test.rs @@ -43,6 +43,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -191,6 +193,8 @@ async fn should_pass_headers() { abort_signal: None, provider_options: None, headers: Some(request_headers), + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); diff --git a/aimux-providers/tests/openai_embedding_test.rs b/aimux-providers/tests/openai_embedding_test.rs index 9df9d21a..1c7bc3ea 100644 --- a/aimux-providers/tests/openai_embedding_test.rs +++ b/aimux-providers/tests/openai_embedding_test.rs @@ -57,6 +57,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -246,6 +248,8 @@ async fn should_pass_dimensions_setting() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -299,6 +303,8 @@ async fn should_pass_headers() { abort_signal: None, provider_options: None, headers: Some(request_headers), + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); diff --git a/aimux-providers/tests/openai_model_test.rs b/aimux-providers/tests/openai_model_test.rs index 92117bb5..16e413ac 100644 --- a/aimux-providers/tests/openai_model_test.rs +++ b/aimux-providers/tests/openai_model_test.rs @@ -4,7 +4,8 @@ //! - `packages/openai/src/chat/openai-chat-language-model.test.ts` //! `describe('doGenerate')` response-parsing cases → `do_generate` tests //! - `packages/openai/src/chat/openai-chat-language-model.test.ts` -//! `describe('doStream')` streaming cases → `do_stream` tests +//! `describe('doStream')` streaming cases → `do_stream` tests, plus +//! Aimux's intentional first-event retry deviation from RFC-0031 §8.3. //! //! Tests that depend on features absent from the Rust data model //! (`providerOptions`, reasoning/text token breakdown, annotations/sources, @@ -16,13 +17,17 @@ //! or SSE response, creates an `OpenAIModel` pointing at the mock, calls //! `do_generate` / `do_stream`, and asserts on the result. +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + use futures::StreamExt; use serde_json::{Value, json}; use wiremock::matchers::{method, path}; -use wiremock::{Mock, MockServer, ResponseTemplate}; +use wiremock::{Mock, MockServer, Request, ResponseTemplate}; use aimux_core::content::ContentPart; use aimux_core::error::AiMuxError; +use aimux_core::generate::{GenerateTextOptions, generate_text, stream_text}; use aimux_core::language_model::LanguageModel; use aimux_core::language_model_message::{LanguageModelPrompt, LanguageModelPromptMessage}; use aimux_core::message::Role; @@ -1384,6 +1389,164 @@ mod do_stream { } } + /// RFC-0031 §13.19b/c: a retryable error in the first SSE event is an + /// attempt failure, and the first normal event from the retry is put back + /// at the head of the returned stream rather than being swallowed. + #[tokio::test] + async fn first_sse_429_retries_and_preserves_the_retry_first_event() { + let server = MockServer::start().await; + let attempts = Arc::new(AtomicUsize::new(0)); + let responder_attempts = Arc::clone(&attempts); + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(move |_request: &Request| { + let body = if responder_attempts.fetch_add(1, Ordering::SeqCst) == 0 { + sse_body(&[&sse_event( + r#"{"error":{"message":"slow down","type":"rate_limit_error","code":429}}"#, + )]) + } else { + sse_body(&[&sse_event( + r#"{"id":"chatcmpl-retry","object":"chat.completion.chunk","created":1702657020,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"hello"},"finish_reason":"stop"}]}"#, + )]) + }; + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .insert_header("retry-after-ms", "0") + .set_body_string(body) + }) + .mount(&server) + .await; + + let provider = + OpenAIProvider::new(OpenAIConfig::new("test-api-key").with_base_url(server.uri())); + let model = provider.model("gpt-3.5-turbo"); + let result = stream_text( + &model, + "Hello", + GenerateTextOptions { + max_retries: Some(1), + ..Default::default() + }, + ) + .await + .expect("the second stream attempt should succeed"); + let parts = collect_stream(StreamResult { + stream: result.stream, + request_body: result.request_body, + response_headers: result.response_headers, + }) + .await; + + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert_eq!(text_deltas(&parts), vec!["hello"]); + } + + #[tokio::test] + async fn first_sse_body_transport_error_enters_core_retry() { + use std::time::Duration; + + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::TcpListener; + + async fn read_request(socket: &mut tokio::net::TcpStream) { + let mut request = Vec::new(); + let mut buffer = [0_u8; 4096]; + let header_end = loop { + let count = socket.read(&mut buffer).await.unwrap(); + assert_ne!(count, 0, "client closed before sending request headers"); + request.extend_from_slice(&buffer[..count]); + if let Some(offset) = request.windows(4).position(|part| part == b"\r\n\r\n") { + break offset + 4; + } + }; + let headers = String::from_utf8_lossy(&request[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + let (name, value) = line.split_once(':')?; + name.eq_ignore_ascii_case("content-length") + .then(|| value.trim().parse::().ok()) + .flatten() + }) + .unwrap_or(0); + while request.len() - header_end < content_length { + let count = socket.read(&mut buffer).await.unwrap(); + if count == 0 { + break; + } + request.extend_from_slice(&buffer[..count]); + } + } + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let attempts = Arc::new(AtomicUsize::new(0)); + let observed_attempts = Arc::clone(&attempts); + let server = tokio::spawn(async move { + for attempt in 0..2 { + let (mut socket, _) = listener.accept().await.unwrap(); + observed_attempts.fetch_add(1, Ordering::SeqCst); + read_request(&mut socket).await; + if attempt == 0 { + socket + .write_all( + b"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: 10000\r\nconnection: close\r\n\r\ndata: {\"id\":\"partial", + ) + .await + .unwrap(); + } else { + let body = sse_body(&[&sse_event( + r#"{"id":"chatcmpl-retry","object":"chat.completion.chunk","created":1702657020,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"retried"},"finish_reason":"stop"}]}"#, + )]); + let response = format!( + "HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}", + body.len() + ); + socket.write_all(response.as_bytes()).await.unwrap(); + } + socket.shutdown().await.unwrap(); + } + }); + + let provider = OpenAIProvider::new( + OpenAIConfig::new("test-api-key") + .with_base_url(base_url) + .with_retry_config(aimux_core::retry::RetryConfig { + max_retries: 1, + initial_delay: Duration::ZERO, + backoff_factor: 2, + }), + ); + let result = stream_text( + &provider.model("gpt-3.5-turbo"), + "Hello", + GenerateTextOptions { + max_retries: Some(1), + ..Default::default() + }, + ) + .await + .expect("the body transport failure should be retried"); + + let attempt_count = attempts.load(Ordering::SeqCst); + if attempt_count == 2 { + server.await.unwrap(); + } else { + server.abort(); + } + assert_eq!( + attempt_count, 2, + "first body transport error was not retried" + ); + let parts = collect_stream(StreamResult { + stream: result.stream, + request_body: result.request_body, + response_headers: result.response_headers, + }) + .await; + assert_eq!(text_deltas(&parts), vec!["retried"]); + } + // ── should forward error stream parts after output has started ──────────── /// TS: "should forward error stream parts after output has started" @@ -1440,7 +1603,15 @@ mod do_stream { #[tokio::test] async fn should_handle_unparsable_stream_parts() { let server = MockServer::start().await; - let body = format!("{}{}", sse_event("{unparsable}"), "data: [DONE]\n\n"); + let body = sse_body(&[ + &sse_event("{unparsable}"), + &sse_event( + r#"{"id":"chatcmpl-after-error","object":"chat.completion.chunk","created":1,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"after-error"},"finish_reason":null}]}"#, + ), + &sse_event( + r#"{"id":"chatcmpl-after-error","object":"chat.completion.chunk","created":1,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}"#, + ), + ]); mock_sse_response(&server, &body).await; let config = OpenAIConfig::new("test-api-key").with_base_url(server.uri()); @@ -1451,23 +1622,30 @@ mod do_stream { .do_stream(&default_options(test_prompt())) .await .expect("do_stream should succeed"); - let parts = collect_stream(result).await; + let outcomes: Vec<_> = result.stream.collect().await; - // Should have: stream-start, Error, Finish (with error finish reason) - assert!(matches!(parts[0], StreamPart::StreamStart { .. })); - - let error_part = parts.iter().find(|p| matches!(p, StreamPart::Error { .. })); - assert!(error_part.is_some(), "should have Error part"); - - let finish = parts + assert!(matches!( + outcomes.first(), + Some(Ok(StreamPart::StreamStart { .. })) + )); + let error_index = outcomes .iter() - .find(|p| matches!(p, StreamPart::Finish { .. })); - match finish { - Some(StreamPart::Finish { finish_reason, .. }) => { - assert_eq!(finish_reason.unified, FinishReasonUnified::Error); - } - other => panic!("expected Finish, got {other:?}"), - } + .position(|outcome| { + matches!( + outcome, + Err(AiMuxError::JsonParse(_) | AiMuxError::InvalidResponseData(_)) + ) + }) + .expect("malformed frame should surface as a parse error item"); + assert!(outcomes[error_index + 1..].iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::TextDelta { delta, .. }) if delta == "after-error" + ))); + assert!(outcomes.iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::Finish { finish_reason, .. }) + if finish_reason.unified == FinishReasonUnified::Stop + ))); } // ── should expose the raw response headers ──────────────────────────────── @@ -1696,18 +1874,17 @@ mod do_stream { ); } - /// RFC-0016 M2: an unparsable chunk emits only `Error` (no `Raw`), and - /// the stream breaks there. + /// RFC-0016 M2: an unparsable chunk emits an error item (no `Raw`) without + /// preventing later independently framed SSE data from being reduced. #[tokio::test] - async fn raw_chunks_unparsable_chunk_emits_error_only() { + async fn raw_chunks_unparsable_chunk_emits_error_and_continues() { let server = MockServer::start().await; let body = sse_body(&[ &sse_event( r#"{"id":"c1","object":"chat.completion.chunk","created":1,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"Hi"},"finish_reason":null}]}"#, ), - // Unparsable chunk: no Raw for it, Error only, stream breaks. + // Unparsable chunk: no Raw for it, but the stream remains usable. "data: not-json\n\n", - // Never reached (stream broke at the unparsable chunk). &sse_event( r#"{"id":"c2","object":"chat.completion.chunk","created":1,"model":"gpt-3.5-turbo","choices":[{"index":0,"delta":{"content":"late"},"finish_reason":null}]}"#, ), @@ -1726,28 +1903,33 @@ mod do_stream { .do_stream(&options) .await .expect("do_stream should succeed"); - let parts = collect_stream(result).await; + let outcomes: Vec<_> = result.stream.collect().await; - // Only the content chunk emitted Raw; the unparsable chunk did not, - // and the stream broke before the late chunk. - let raw: Vec<&Value> = parts + // Each valid JSON chunk emits Raw; the unparsable chunk does not. + let raw: Vec<&Value> = outcomes .iter() - .filter_map(|p| match p { - StreamPart::Raw { raw_value } => Some(raw_value), - _ => None, + .filter_map(|outcome| match outcome { + Ok(p) => match p { + StreamPart::Raw { raw_value } => Some(raw_value), + _ => None, + }, + Err(_) => None, }) .collect(); - assert_eq!(raw.len(), 1, "no Raw for the unparsable chunk"); + assert_eq!(raw.len(), 2, "no Raw for the unparsable chunk"); assert_eq!(raw[0]["choices"][0]["delta"]["content"], json!("Hi")); + assert_eq!(raw[1]["choices"][0]["delta"]["content"], json!("late")); - // The unparsable chunk surfaces as Error. - let err_pos = parts + let err_pos = outcomes .iter() - .position(|p| matches!(p, StreamPart::Error { .. })) - .expect("unparsable chunk must surface as Error part"); + .position(|outcome| matches!(outcome, Err(AiMuxError::JsonParse(_)))) + .expect("unparsable chunk must surface as a JSON parse error item"); assert!( - !matches!(&parts[err_pos - 1], StreamPart::Raw { .. }), - "no Raw may precede the Error of an unparsable chunk" + outcomes[err_pos + 1..].iter().any(|outcome| matches!( + outcome, + Ok(StreamPart::TextDelta { delta, .. }) if delta == "late" + )), + "a later valid event must still be reduced" ); } @@ -1817,8 +1999,8 @@ mod do_stream { // RFC-0016 H1/H3: abort + per-call timeout (provider-level) // ════════════════════════════════════════════════════════════════════════════ +use aimux_core::AbortSignal; use aimux_core::options::TimeoutConfiguration; -use aimux_core::shared::AbortSignal; use std::time::Duration; /// TS: timeout — total_ms bounds the whole call. @@ -1847,14 +2029,15 @@ async fn total_timeout_aborts_slow_generate() { let provider = OpenAIProvider::new(config); let model = provider.model("gpt-3.5-turbo"); - let mut options = default_options(test_prompt()); - options.timeout = Some(TimeoutConfiguration { - total_ms: Some(100), + let options = GenerateTextOptions { + timeout: Some(TimeoutConfiguration { + total_ms: Some(100), + ..Default::default() + }), ..Default::default() - }); + }; - let err = model - .do_generate(&options) + let err = generate_text(&model, "Hello", options) .await .expect_err("slow response must be cut by total timeout"); assert!(matches!(err, AiMuxError::Timeout(_)), "got {err:?}"); @@ -1898,7 +2081,7 @@ async fn abort_signal_cancels_in_flight_generate() { let result = handle.await.expect("task must finish"); assert!( - matches!(result, Err(AiMuxError::Aborted)), + matches!(result, Err(AiMuxError::Aborted(_))), "aborted call must fail with Aborted, got {result:?}" ); } @@ -1937,7 +2120,7 @@ async fn abort_before_send_fails_fast() { .do_generate(&options) .await .expect_err("pre-aborted signal must fail fast"); - assert!(matches!(err, AiMuxError::Aborted), "got {err:?}"); + assert!(matches!(err, AiMuxError::Aborted(_)), "got {err:?}"); } // ════════════════════════════════════════════════════════════════════════════ diff --git a/aimux-providers/tests/openai_transcription_stream_test.rs b/aimux-providers/tests/openai_transcription_stream_test.rs index f6e60dd7..7ed19a2f 100644 --- a/aimux-providers/tests/openai_transcription_stream_test.rs +++ b/aimux-providers/tests/openai_transcription_stream_test.rs @@ -16,8 +16,8 @@ use futures::{SinkExt, StreamExt}; use tokio::net::TcpListener; use tokio_tungstenite::tungstenite::Message; +use aimux_core::AbortSignal; use aimux_core::error::AiMuxError; -use aimux_core::shared::AbortSignal; use aimux_core::transcription_model::{ AudioChunk, InputAudioFormat, TranscriptionModel, TranscriptionStreamOptions, TranscriptionStreamPart, @@ -344,7 +344,9 @@ async fn stream_abort_mid_session() { }); let parts = collect(result).await; assert!( - parts.iter().any(|p| matches!(p, Err(AiMuxError::Aborted))), + parts + .iter() + .any(|p| matches!(p, Err(AiMuxError::Aborted(_)))), "expected Aborted in parts: {parts:?}" ); } @@ -424,7 +426,7 @@ async fn stream_abort_during_connect() { .unwrap_err(); // Abort semantics exactly: distinct from timeout and from IO failure. assert!( - matches!(err, AiMuxError::Aborted), + matches!(err, AiMuxError::Aborted(_)), "connect-phase abort must surface Aborted, got {err:?}" ); assert!(!matches!(&err, AiMuxError::Timeout(_))); @@ -463,6 +465,7 @@ async fn stream_connect_timeout_fires() { timeout: Some(aimux_core::options::TimeoutConfiguration { first_chunk_ms: Some(300), total_ms: None, + step_ms: None, chunk_ms: None, }), }; @@ -517,6 +520,7 @@ async fn stream_first_chunk_timeout_fires() { // loopback but far below the server's 2s hold. first_chunk_ms: Some(500), total_ms: None, + step_ms: None, chunk_ms: None, }), }; @@ -670,6 +674,7 @@ async fn stream_chunk_idle_timeout_fires() { timeout: Some(aimux_core::options::TimeoutConfiguration { first_chunk_ms: None, total_ms: None, + step_ms: None, chunk_ms: Some(400), }), }; @@ -733,6 +738,7 @@ async fn stream_total_timeout_fires() { timeout: Some(aimux_core::options::TimeoutConfiguration { first_chunk_ms: None, total_ms: Some(700), + step_ms: None, chunk_ms: Some(300), }), }; diff --git a/aimux-providers/tests/prodia_video_test.rs b/aimux-providers/tests/prodia_video_test.rs index 3f8cf851..e485ae7e 100644 --- a/aimux-providers/tests/prodia_video_test.rs +++ b/aimux-providers/tests/prodia_video_test.rs @@ -1,15 +1,24 @@ //! Rust translation of the Prodia video model tests. //! Source: `reference/ai/packages/prodia/src/prodia-video-model.test.ts` -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_core::video_model::{VideoCallOptions, generate_video}; use aimux_providers::{ProdiaConfig, ProdiaProvider}; use serde_json::{Value, json}; use std::collections::HashMap; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(p: &str) -> VideoCallOptions { - VideoCallOptions::new(p) + let mut o = VideoCallOptions::new(p); + o.poll = fast_poll(); + o } async fn mock_job_and_result(server: &MockServer, result: &Value) { @@ -36,7 +45,7 @@ async fn should_generate_video() { let config = ProdiaConfig::new("test-api-key").with_base_url(server.uri()); let provider = ProdiaProvider::new(config); let model = provider.video("svd-xt"); - let r = model.do_generate(&options("A cat")).await.unwrap(); + let r = generate_video(&model, options("A cat")).await.unwrap(); assert_eq!(r.videos.len(), 1); } @@ -51,7 +60,7 @@ async fn should_pass_model_type_and_prompt() { let config = ProdiaConfig::new("test-api-key").with_base_url(server.uri()); let provider = ProdiaProvider::new(config); let model = provider.video("svd-xt"); - model.do_generate(&options("A cat")).await.unwrap(); + generate_video(&model, options("A cat")).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["type"], "svd-xt"); @@ -77,7 +86,7 @@ async fn should_pass_headers() { let mut rh = HashMap::new(); rh.insert("Custom-Request-Header".to_string(), "req-val".to_string()); opts.headers = Some(rh); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; assert_eq!(h.get("x-prodia-key").unwrap(), "test-api-key"); @@ -95,7 +104,7 @@ async fn should_include_response_data() { let config = ProdiaConfig::new("test-api-key").with_base_url(server.uri()); let provider = ProdiaProvider::new(config); let model = provider.video("svd-xt"); - let r = model.do_generate(&options("test")).await.unwrap(); + let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!(r.response.model_id, Some("svd-xt".to_string())); } diff --git a/aimux-providers/tests/provider_error_test.rs b/aimux-providers/tests/provider_error_test.rs index bd258596..a4a2857c 100644 --- a/aimux-providers/tests/provider_error_test.rs +++ b/aimux-providers/tests/provider_error_test.rs @@ -9,17 +9,17 @@ //! cases: 529 status, first-chunk overloaded, mid-stream overloaded). //! - `packages/provider-utils/src/response-handler.test.ts` //! (status-code → error-variant mapping exercised end-to-end through the -//! providers' `do_generate` / `do_stream`, which call -//! `parse_provider_error` under the hood). +//! providers' `do_generate` / `do_stream`, which select explicit +//! failed-response handlers). //! - `packages/openai/src/openai-error.test.ts` & -//! `packages/anthropic/src/anthropic-error.test.ts` (error-structure -//! message extraction for both providers' JSON shapes). +//! `packages/anthropic/src/anthropic-error.test.ts` (provider-specific +//! failed-response message extraction for both JSON shapes). //! //! Each test stands up a `wiremock` server returning an error status code + //! error JSON body (or, for the stream cases, an SSE body whose first/mid //! chunk is an error JSON) and asserts the resulting `AiMuxError` variant. //! -//! Status → error mapping (mirrors `parse_provider_error`): every status +//! Status → error mapping: every status //! yields `AiMuxError::ApiCall` with the observed code in `status_code`; //! 408/409/429/5xx are stored retryable (`is_retryable`), other 4xx are not. @@ -187,7 +187,7 @@ mod openai_generate_errors { } /// TS openai-error.test.ts: OpenRouter nests a stringified JSON error - /// inside `error.message`. `parse_provider_error` must keep that string + /// inside `error.message`. The OpenAI failed-response handler must keep that string /// verbatim (it is the message, not a nested structure to drill into). #[tokio::test] async fn openrouter_resource_exhausted_message_kept_verbatim() { @@ -205,11 +205,10 @@ mod openai_generate_errors { .await; let result = model(&server).do_generate(&options()).await; - // 429 → RateLimited; the nested message is carried but RateLimited - // only exposes retry_after_ms, so we just assert the variant. + // The provider message is retained inside the ApiCall detail. assert!( - matches!(result, Err(ref e) if e.status_code() == Some(429)), - "expected RateLimited for OpenRouter 429, got {result:?}" + matches!(result, Err(AiMuxError::ApiCall(ref e)) if e.status_code == Some(429) && e.message == nested), + "expected ApiCall for OpenRouter 429, got {result:?}" ); } } @@ -233,7 +232,7 @@ mod openai_stream_errors { } /// A non-success HTTP status on the stream endpoint makes `do_stream` - /// itself return `Err` via `parse_provider_error` (the stream is never + /// itself return `Err` via the failed-response handler (the stream is never /// started). 500 → Provider. #[tokio::test] async fn http_500_status_rejects_do_stream() { @@ -469,20 +468,6 @@ mod anthropic_generate_errors { "expected ApiCall error carrying 'Overloaded', got {result:?}" ); } - - /// TS anthropic-error.test.ts: the overloaded error structure - /// `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` - /// is parsed and `error.message` is extracted. - #[tokio::test] - async fn overloaded_error_structure_message_extracted() { - use aimux_provider_utils::{DEFAULT_ERROR_STRUCTURE, parse_provider_error}; - let body = r#"{"type":"error","error":{"details":null,"type":"overloaded_error","message":"Overloaded"}}"#; - let err = parse_provider_error(529, body, &DEFAULT_ERROR_STRUCTURE); - assert!( - matches!(err, AiMuxError::ApiCall(ref m) if m.message.contains("Overloaded")), - "expected Provider carrying 'Overloaded', got {err:?}" - ); - } } // ════════════════════════════════════════════════════════════════════════════ @@ -505,7 +490,7 @@ mod anthropic_stream_errors { /// TS: "should throw an api error when the server is returning a 529 /// overloaded error" (doStream). A non-success HTTP status makes - /// `do_stream` return `Err` via `parse_provider_error`. + /// `do_stream` return `Err` via the failed-response handler. #[tokio::test] async fn http_529_status_rejects_do_stream() { let server = MockServer::start().await; @@ -527,8 +512,8 @@ mod anthropic_stream_errors { /// overloaded error" (doStream). /// /// The stream returns 200 + SSE whose first event is an `error` event. - /// In the Rust model the error surfaces as a `StreamPart::Error` early in - /// the stream. + /// The provider peeks that event before returning the stream so the Core + /// operation retry can treat it as an attempt failure. #[tokio::test] async fn first_stream_chunk_is_overloaded_error() { let server = MockServer::start().await; @@ -547,21 +532,16 @@ mod anthropic_stream_errors { .mount(&server) .await; - let stream = model(&server).do_stream(&options()).await.unwrap(); - let parts = collect_stream(stream).await; + let error = model(&server) + .do_stream(&options()) + .await + .expect_err("first error event must reject do_stream"); assert!( - parts.iter().any(|p| matches!(p, - StreamPart::Error { error: AiMuxError::ApiCall(m) } if m.message == "Overloaded")), - "expected a StreamPart::Error carrying 'Overloaded', got {parts:?}" - ); - // Finish is the final-chunk contract: the stream must terminate with - // an error-finish even when the very first event is an error. - assert!( - matches!(parts.last(), - Some(StreamPart::Finish { finish_reason, .. }) - if matches!(finish_reason.unified, - aimux_core::prelude::FinishReasonUnified::Error)), - "expected the last part to be Finish with unified=Error, got {parts:?}" + matches!(error, AiMuxError::ApiCall(ref detail) + if detail.status_code == Some(529) + && detail.message == "Overloaded" + && detail.is_retryable), + "expected retryable 529 ApiCall error, got {error:?}" ); } @@ -664,55 +644,3 @@ mod anthropic_stream_errors { ); } } - -// ════════════════════════════════════════════════════════════════════════════ -// Error-structure parsing (openai-error.test.ts + anthropic-error.test.ts) -// ════════════════════════════════════════════════════════════════════════════ - -mod error_structure_parsing { - use super::*; - use aimux_provider_utils::{DEFAULT_ERROR_STRUCTURE, parse_provider_error}; - - /// TS openai-error.test.ts: the OpenAI error schema parses - /// `{"error":{"message":"...","code":429}}` and `error.message` is the - /// surfaced message. - #[test] - fn openai_error_message_extracted_from_default_structure() { - let body = r#"{"error":{"message":"Resource has been exhausted","type":"requests","param":null,"code":429}}"#; - let err = parse_provider_error(429, body, &DEFAULT_ERROR_STRUCTURE); - // 429 → RateLimited (variant only; message is not carried on the - // RateLimited variant in the Rust error model). - assert!(matches!(err, ref e if e.status_code() == Some(429))); - } - - /// TS openai-error.test.ts: OpenRouter nests a stringified JSON object - /// inside `error.message`. `parse_provider_error` must keep that string - /// verbatim rather than trying to drill into it. - #[test] - fn openrouter_nested_message_kept_verbatim() { - let nested = "{\n \"error\": {\n \"code\": 429,\n \"message\": \"Resource has been exhausted (e.g. check quota).\",\n \"status\": \"RESOURCE_EXHAUSTED\"\n }\n}\n"; - let body = format!( - r#"{{"error":{{"message":{msg},"code":429}}}}"#, - msg = serde_json::to_string(nested).unwrap() - ); - let err = parse_provider_error(500, &body, &DEFAULT_ERROR_STRUCTURE); - // 500 → Provider; the (stringified) nested JSON is the message. - assert!( - matches!(err, AiMuxError::ApiCall(ref m) if m.message.contains("RESOURCE_EXHAUSTED")), - "expected Provider carrying the OpenRouter nested message, got {err:?}" - ); - } - - /// TS anthropic-error.test.ts: the overloaded error structure - /// `{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}` - /// parses with `error.message` = "Overloaded". - #[test] - fn anthropic_overloaded_error_message_extracted() { - let body = r#"{"type":"error","error":{"details":null,"type":"overloaded_error","message":"Overloaded"}}"#; - let err = parse_provider_error(529, body, &DEFAULT_ERROR_STRUCTURE); - assert!( - matches!(err, AiMuxError::ApiCall(ref m) if m.message.contains("Overloaded")), - "expected Provider carrying 'Overloaded', got {err:?}" - ); - } -} diff --git a/aimux-providers/tests/replicate_video_test.rs b/aimux-providers/tests/replicate_video_test.rs index 5cc06fd4..41eb8162 100644 --- a/aimux-providers/tests/replicate_video_test.rs +++ b/aimux-providers/tests/replicate_video_test.rs @@ -1,15 +1,24 @@ //! Rust translation of the Replicate video model tests. //! Source: `reference/ai/packages/replicate/src/replicate-video-model.test.ts` -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_core::video_model::{VideoCallOptions, generate_video}; use aimux_providers::{ReplicateConfig, ReplicateProvider}; use serde_json::{Value, json}; use std::collections::HashMap; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(p: &str) -> VideoCallOptions { - VideoCallOptions::new(p) + let mut o = VideoCallOptions::new(p); + o.poll = fast_poll(); + o } async fn mock_predict_and_result(server: &MockServer, output: &Value) { @@ -35,7 +44,7 @@ async fn should_generate_video() { let config = ReplicateConfig::new("test-api-key").with_base_url(server.uri()); let provider = ReplicateProvider::new(config); let model = provider.video("wan-lab/wan-2.1-t2v-14b"); - let r = model.do_generate(&options("A cat")).await.unwrap(); + let r = generate_video(&model, options("A cat")).await.unwrap(); assert_eq!(r.videos.len(), 1); } @@ -46,7 +55,7 @@ async fn should_pass_model_and_prompt() { let config = ReplicateConfig::new("test-api-key").with_base_url(server.uri()); let provider = ReplicateProvider::new(config); let model = provider.video("wan-lab/wan-2.1-t2v-14b"); - model.do_generate(&options("A cat")).await.unwrap(); + generate_video(&model, options("A cat")).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["model"], "wan-lab/wan-2.1-t2v-14b"); @@ -68,7 +77,7 @@ async fn should_pass_headers() { let mut rh = HashMap::new(); rh.insert("Custom-Request-Header".to_string(), "req-val".to_string()); opts.headers = Some(rh); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; assert_eq!(h.get("authorization").unwrap(), "Token test-api-key"); @@ -82,7 +91,7 @@ async fn should_include_response_data() { let config = ReplicateConfig::new("test-api-key").with_base_url(server.uri()); let provider = ReplicateProvider::new(config); let model = provider.video("test-model"); - let r = model.do_generate(&options("test")).await.unwrap(); + let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!(r.response.model_id, Some("test-model".to_string())); } diff --git a/aimux-providers/tests/runwayml_test.rs b/aimux-providers/tests/runwayml_test.rs index 9ccb9f37..23e2ab71 100644 --- a/aimux-providers/tests/runwayml_test.rs +++ b/aimux-providers/tests/runwayml_test.rs @@ -12,14 +12,23 @@ use serde_json::{Value, json}; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_core::video_model::{VideoCallOptions, VideoModel, generate_video}; use aimux_providers::{RunwaymlConfig, RunwaymlProvider}; const MODEL_ID: &str = "gen3a_turbo"; const VIDEO_URL: &str = "https://cdn.runwayml.com/video.mp4"; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(prompt: &str) -> VideoCallOptions { - VideoCallOptions::new(prompt) + let mut o = VideoCallOptions::new(prompt); + o.poll = fast_poll(); + o } /// A config pointed at the mock server with a short poll interval so tests run @@ -80,7 +89,9 @@ async fn should_generate_video_from_prompt() { let provider = RunwaymlProvider::new(config(server.uri())); let model = provider.video(MODEL_ID); - let result = model.do_generate(&options("A cat playing")).await.unwrap(); + let result = generate_video(&model, options("A cat playing")) + .await + .unwrap(); assert_eq!(result.videos.len(), 1); match &result.videos[0] { @@ -101,7 +112,9 @@ async fn should_post_to_text_to_video_endpoint() { let provider = RunwaymlProvider::new(config(server.uri())); let model = provider.video(MODEL_ID); - model.do_generate(&options("A cat playing")).await.unwrap(); + generate_video(&model, options("A cat playing")) + .await + .unwrap(); let requests = server.received_requests().await.unwrap(); // The first request is the task submission. @@ -120,7 +133,7 @@ async fn should_send_auth_and_version_headers() { let provider = RunwaymlProvider::new(config(server.uri())); let model = provider.video(MODEL_ID); - model.do_generate(&options("test")).await.unwrap(); + generate_video(&model, options("test")).await.unwrap(); let requests = server.received_requests().await.unwrap(); // Submit request carries the bearer token and the required version header. @@ -164,7 +177,7 @@ async fn should_pass_custom_headers() { request_headers.insert("Custom-Request-Header".to_string(), "req-val".to_string()); opts.headers = Some(request_headers); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; @@ -193,7 +206,7 @@ async fn should_return_error_when_task_failed() { let provider = RunwaymlProvider::new(config(server.uri())); let model = provider.video(MODEL_ID); - let result = model.do_generate(&options("test")).await; + let result = generate_video(&model, options("test")).await; assert!(result.is_err()); let err = result.unwrap_err().to_string(); assert!( @@ -214,7 +227,7 @@ async fn should_return_error_on_submit_failure() { let provider = RunwaymlProvider::new(config(server.uri())); let model = provider.video(MODEL_ID); - let result = model.do_generate(&options("test")).await; + let result = generate_video(&model, options("test")).await; assert!(result.is_err()); } diff --git a/aimux-providers/tests/stream_error_semantics_test.rs b/aimux-providers/tests/stream_error_semantics_test.rs new file mode 100644 index 00000000..734ef9df --- /dev/null +++ b/aimux-providers/tests/stream_error_semantics_test.rs @@ -0,0 +1,139 @@ +//! Stream-error classification and termination semantics (review follow-ups). +//! +//! Locks three behaviors: +//! - no fabricated HTTP 500 for non-numeric provider error codes (the "M3 +//! bug"): such errors carry no status and are not retryable; +//! - payload `retry_after_ms` hints surface through `retry_after_hint()`; +//! - the shared Responses reducer ends its stream right after a terminal +//! error event instead of waiting on a source that stays open. + +use std::collections::HashMap; + +use futures::StreamExt; +use serde_json::json; + +use aimux_core::AiMuxError; +use aimux_core::stream_part::StreamPart; +use aimux_core::types::FinishReasonUnified; +use aimux_provider_utils::stream_error_api_call; +use aimux_providers::openai::responses::responses_convert::build_responses_event_stream; + +fn as_api_call(error: &AiMuxError) -> &aimux_core::ApiCallError { + match error { + AiMuxError::ApiCall(detail) => detail, + other => panic!("expected ApiCall, got {other:?}"), + } +} + +#[test] +fn string_error_code_is_not_a_retryable_500() { + let payload = json!({"message": "Incorrect API key", "code": "invalid_api_key"}); + let error = stream_error_api_call( + "Incorrect API key", + Some("invalid_api_key".into()), + None, + &payload, + "https://example.test", + json!({}), + HashMap::new(), + ); + let detail = as_api_call(&error); + assert_eq!(detail.status_code, None); + assert!(!error.is_retryable()); +} + +#[test] +fn numeric_status_keeps_status_based_retryability() { + let payload = json!({"message": "rate limited", "code": 429}); + let error = stream_error_api_call( + "rate limited", + Some("429".into()), + Some(429), + &payload, + "https://example.test", + json!({}), + HashMap::new(), + ); + assert_eq!(as_api_call(&error).status_code, Some(429)); + assert!(error.is_retryable()); +} + +#[test] +fn payload_retry_after_ms_surfaces_as_hint() { + let payload = json!({"message": "rate limited", "code": 429, "retry_after_ms": 15000}); + let error = stream_error_api_call( + "rate limited", + None, + Some(429), + &payload, + "https://example.test", + json!({}), + HashMap::new(), + ); + assert_eq!(error.retry_after_hint(), Some(15000)); +} + +#[test] +fn real_retry_after_header_wins_over_payload() { + let payload = json!({"message": "rate limited", "retry_after_ms": 15000}); + let headers = HashMap::from([("retry-after-ms".to_string(), "7".to_string())]); + let error = stream_error_api_call( + "rate limited", + None, + Some(429), + &payload, + "https://example.test", + json!({}), + HashMap::new(), + ); + assert_eq!(error.retry_after_hint(), Some(15000)); + let error = stream_error_api_call( + "rate limited", + None, + Some(429), + &payload, + "https://example.test", + json!({}), + headers, + ); + assert_eq!(error.retry_after_hint(), Some(7)); +} + +#[tokio::test] +async fn responses_stream_ends_after_terminal_error_event() { + // First event is benign so the peek passes; the mid-stream error event is + // followed by a source that never closes (heartbeat-style server). + let first = json!({"type": "response.created", "response": {"id": "resp_1"}}); + let error_event = + json!({"type": "error", "error": {"message": "boom", "type": "server_error"}}); + let source = futures::stream::iter(vec![Ok(error_event)]).chain(futures::stream::pending()); + + let stream = build_responses_event_stream( + Some(Ok(first)), + source, + "openai".to_string(), + Vec::new(), + false, + "https://example.test".to_string(), + json!({}), + HashMap::new(), + ) + .expect("stream setup"); + + let parts: Vec<_> = stream.collect().await; + assert!( + parts + .iter() + .any(|part| matches!(part, Ok(StreamPart::Error { .. }))), + "error part yielded" + ); + let last = parts.last().expect("stream not empty"); + assert!( + matches!( + last, + Ok(StreamPart::Finish { finish_reason, .. }) + if finish_reason.unified == FinishReasonUnified::Error + ), + "stream ends with an error Finish, got {last:?}" + ); +} diff --git a/aimux-providers/tests/vertex_anthropic_test.rs b/aimux-providers/tests/vertex_anthropic_test.rs index c7549560..d8e5fc85 100644 --- a/aimux-providers/tests/vertex_anthropic_test.rs +++ b/aimux-providers/tests/vertex_anthropic_test.rs @@ -27,8 +27,6 @@ use aimux_providers::vertex::{ VertexAnthropicModel, VertexAuth, VertexProvider, VertexProviderConfig, }; -use aimux_provider_utils::RetryConfig; - // ── Constants ───────────────────────────────────────────────────────────────── const MODEL_ID: &str = "claude-sonnet-4-20250514"; @@ -63,7 +61,7 @@ fn make_model(server: &MockServer) -> VertexAnthropicModel { location: Some("us-central1".to_string()), auth: VertexAuth::BearerToken("test-token".to_string()), api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }; VertexProvider::new(config) .anthropic_model(MODEL_ID) diff --git a/aimux-providers/tests/vertex_embedding_test.rs b/aimux-providers/tests/vertex_embedding_test.rs index 883b645b..856c7dd5 100644 --- a/aimux-providers/tests/vertex_embedding_test.rs +++ b/aimux-providers/tests/vertex_embedding_test.rs @@ -12,8 +12,6 @@ use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel}; use aimux_providers::{VertexAuth, VertexProvider, VertexProviderConfig}; -use aimux_provider_utils::RetryConfig; - const TEST_VALUES: &[&str] = &["test text one", "test text two"]; fn mock_provider_options() -> HashMap { @@ -70,6 +68,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: Some(mock_provider_options()), headers: None, + max_retries: None, + timeout: None, } } @@ -80,7 +80,7 @@ fn test_provider(base_url: String) -> VertexProvider { location: Some("us-central1".to_string()), auth: VertexAuth::BearerToken("test-token".to_string()), api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }; VertexProvider::new(config) } @@ -205,6 +205,8 @@ async fn should_accept_google_vertex_key() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -238,6 +240,8 @@ async fn should_pass_task_type_only() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -291,7 +295,7 @@ fn gemini_embedding_2_max_per_call() { location: Some("us-central1".to_string()), auth: VertexAuth::BearerToken("test".to_string()), api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }; let provider = VertexProvider::new(config); let model = provider.embedding_model("gemini-embedding-2"); diff --git a/aimux-providers/tests/vertex_model_test.rs b/aimux-providers/tests/vertex_model_test.rs index fbb16c41..33d093e6 100644 --- a/aimux-providers/tests/vertex_model_test.rs +++ b/aimux-providers/tests/vertex_model_test.rs @@ -25,8 +25,6 @@ use aimux_core::types::FinishReasonUnified; use aimux_providers::vertex::{VertexAuth, VertexConfig, VertexModel, VertexProviderConfig}; -use aimux_provider_utils::RetryConfig; - // ── Helpers ────────────────────────────────────────────────────────────────── fn test_prompt() -> LanguageModelPrompt { @@ -48,7 +46,7 @@ fn make_model(server: &MockServer) -> VertexModel { base_url: server.uri(), auth: VertexAuth::BearerToken("test-token".to_string()), api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }, ) } @@ -360,7 +358,7 @@ async fn vertex_api_key_auth() { base_url: server.uri(), auth: VertexAuth::ApiKey("test-api-key".to_string()), api_key_source: None, - retry_config: RetryConfig::default(), + retry_config: aimux_provider_utils::RetryConfig::default(), }, ); diff --git a/aimux-providers/tests/vertex_video_test.rs b/aimux-providers/tests/vertex_video_test.rs index 16a445f6..b1abe890 100644 --- a/aimux-providers/tests/vertex_video_test.rs +++ b/aimux-providers/tests/vertex_video_test.rs @@ -1,15 +1,24 @@ //! Rust translation of the Google Vertex video model tests. //! Source: `reference/ai/packages/google-vertex/src/google-vertex-video-model.test.ts` -use aimux_core::video_model::{VideoCallOptions, VideoModel}; +use aimux_core::video_model::{VideoCallOptions, generate_video}; use aimux_providers::{VertexProvider, VertexProviderConfig}; use serde_json::{Value, json}; use std::collections::HashMap; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +fn fast_poll() -> Option { + Some(aimux_core::video_model::VideoPollOptions { + interval_ms: Some(1), + timeout_ms: Some(10_000), + }) +} + fn options(p: &str) -> VideoCallOptions { - VideoCallOptions::new(p) + let mut o = VideoCallOptions::new(p); + o.poll = fast_poll(); + o } fn make_provider(server_uri: &str) -> VertexProvider { @@ -41,7 +50,7 @@ async fn should_generate_video() { mock_predict_and_poll(&server, &result).await; let provider = make_provider(&server.uri()); let model = provider.video("veo-3.0-generate-001").unwrap(); - let r = model.do_generate(&options("A cat")).await.unwrap(); + let r = generate_video(&model, options("A cat")).await.unwrap(); assert_eq!(r.videos.len(), 1); } @@ -53,7 +62,7 @@ async fn should_pass_prompt() { mock_predict_and_poll(&server, &result).await; let provider = make_provider(&server.uri()); let model = provider.video("veo-3.0-generate-001").unwrap(); - model.do_generate(&options("A cat")).await.unwrap(); + generate_video(&model, options("A cat")).await.unwrap(); let requests = server.received_requests().await.unwrap(); let body: Value = serde_json::from_slice(&requests[0].body).unwrap(); assert_eq!(body["instances"][0]["prompt"], "A cat"); @@ -71,7 +80,7 @@ async fn should_pass_headers() { let mut rh = HashMap::new(); rh.insert("Custom-Header".to_string(), "val".to_string()); opts.headers = Some(rh); - model.do_generate(&opts).await.unwrap(); + generate_video(&model, opts).await.unwrap(); let requests = server.received_requests().await.unwrap(); let h = &requests[0].headers; assert_eq!(h.get("authorization").unwrap(), "Bearer test-token"); @@ -86,7 +95,7 @@ async fn should_include_response_data() { mock_predict_and_poll(&server, &result).await; let provider = make_provider(&server.uri()); let model = provider.video("veo-3.0-generate-001").unwrap(); - let r = model.do_generate(&options("test")).await.unwrap(); + let r = generate_video(&model, options("test")).await.unwrap(); assert!(r.response.timestamp.is_some()); assert_eq!( r.response.model_id, diff --git a/aimux-providers/tests/voyage_embedding_test.rs b/aimux-providers/tests/voyage_embedding_test.rs index a400e4a6..8c32565e 100644 --- a/aimux-providers/tests/voyage_embedding_test.rs +++ b/aimux-providers/tests/voyage_embedding_test.rs @@ -40,6 +40,8 @@ fn default_options(values: Vec) -> EmbeddingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -192,6 +194,8 @@ async fn should_pass_input_type_setting() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -222,6 +226,8 @@ async fn should_pass_output_dimension_setting() { abort_signal: None, provider_options: Some(provider_options), headers: None, + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); @@ -256,6 +262,8 @@ async fn should_pass_headers() { abort_signal: None, provider_options: None, headers: Some(request_headers), + max_retries: None, + timeout: None, }; let _ = model.do_embed(&options).await.expect("should succeed"); diff --git a/aimux-providers/tests/voyage_reranking_test.rs b/aimux-providers/tests/voyage_reranking_test.rs index cec175e9..43432c76 100644 --- a/aimux-providers/tests/voyage_reranking_test.rs +++ b/aimux-providers/tests/voyage_reranking_test.rs @@ -42,6 +42,8 @@ fn text_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } @@ -58,6 +60,8 @@ fn json_docs_opts(query: &str, top_n: u32) -> RerankingCallOptions { abort_signal: None, provider_options: None, headers: None, + max_retries: None, + timeout: None, } } From 80ea1ba09a1d67e07c51dfba6a0e140f948b8f8a Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Mon, 24 Aug 2026 22:29:11 +0800 Subject: [PATCH 04/24] feat(ffi)!: align the C ABI with the pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expose the retry/timeout error context (retryable flag, retry_ms hint, provider code/message, response body, sanitized url/request values, response headers, provider data) and the video start/status split across the C surface; aimux_error_request_id is removed — request ids ride in response_headers. Both stream pumps keep the stream alive across recoverable frame errors: the plain pump forwards them as StreamPart::Error data frames, while the chunk-typed OpenAI pump skips them — that wire is typed as ChatCompletionChunk and an injected error object would kill every typed consumer. The header's timeout keys and code range are pinned by the header/exports contract test; the C examples consume the new error getters. --- aimux-ffi/aimux-error.h | 44 ++- aimux-ffi/aimux-ffi.h | 4 +- aimux-ffi/src/lib.rs | 503 ++++++++++++++++++++++--- aimux-ffi/src/transcription_session.rs | 9 +- aimux-ffi/tests/common/mod.rs | 6 +- aimux-ffi/tests/exports_smoke_test.rs | 66 ++-- bindings/c/example.c | 1 - bindings/c/example.cpp | 8 +- 8 files changed, 537 insertions(+), 104 deletions(-) diff --git a/aimux-ffi/aimux-error.h b/aimux-ffi/aimux-error.h index 40fff096..5cc92bf0 100644 --- a/aimux-ffi/aimux-error.h +++ b/aimux-ffi/aimux-error.h @@ -6,10 +6,11 @@ * trailing out-parameter, which remains at its documented sentinel on failure. * * Every non-NULL error has one non-zero `aimux_error_code_t` and one message. - * Codes 1..13 come from `AiMuxError`, 100..105 from `RecordingError`, and - * 200..206 identify failures detected while crossing the C ABI. Higher-level - * bindings reconstruct their native error types from that code; they map all - * 200..206 codes to the language's existing argument/state/invariant error. + * Codes 1..14 come from `AiMuxError`, 100..105 from + * `RecordingError`, and 200..206 identify failures detected while crossing + * the C ABI. Higher-level bindings reconstruct their native error types from + * that code; they map all 200..206 codes to the language's existing + * argument/state/invariant error. * * Strings returned by getters are owned by the caller and must be released * with `aimux_free_string()`. Release the error itself exactly once with @@ -39,7 +40,7 @@ typedef struct aimux_error aimux_error_t; typedef enum aimux_error_code { AIMUX_OK = 0, - /* AiMuxError: 1..13. */ + /* AiMuxError: 1..14. */ AIMUX_E_OTHER = 1, AIMUX_E_JSON_PARSE = 2, AIMUX_E_INVALID_RESPONSE_DATA = 3, @@ -53,6 +54,10 @@ typedef enum aimux_error_code { AIMUX_E_API_CALL = 11, AIMUX_E_TIMEOUT = 12, AIMUX_E_ABORTED = 13, + /* Reclaims the slot the pre-unification `Other` vacated: the opaque-pointer + * ABI break means no pre-unification caller can link, so nothing can + * misread it. */ + AIMUX_E_RETRY = 14, /* RecordingError: 100..105. */ AIMUX_E_RECORDING_INIT = 100, @@ -104,10 +109,17 @@ int64_t aimux_error_retry_ms(const aimux_error_t *error); char *aimux_error_provider_code(const aimux_error_t *error); /** Failure text without Aimux's composed prefix. */ char *aimux_error_provider_message(const aimux_error_t *error); -/** Provider request id. */ -char *aimux_error_request_id(const aimux_error_t *error); /** Raw provider response body. */ char *aimux_error_response_body(const aimux_error_t *error); +/** Sanitized request URL of the failed call. */ +char *aimux_error_url(const aimux_error_t *error); +/** Sanitized request body values, as a JSON string. */ +char *aimux_error_request_body_values(const aimux_error_t *error); +/** Sanitized response headers, as one JSON object string of + * name → value pairs, e.g. {"retry-after-ms":"1500"}. */ +char *aimux_error_response_headers(const aimux_error_t *error); +/** Parsed provider error data, as a JSON string. */ +char *aimux_error_provider_data(const aimux_error_t *error); /* AIMUX_E_NO_SUCH_MODEL — returned strings are caller-owned. */ @@ -118,6 +130,24 @@ char *aimux_error_model_type(const aimux_error_t *error); char *aimux_error_provider_id(const aimux_error_t *error); +/* + * AIMUX_E_RETRY — retrying stopped; the error keeps the per-attempt history. + * `aimux_error_message()` composes the summary ("Failed after N attempts…"). + */ + +/** Why retrying stopped: "maxRetriesExceeded" (every permitted attempt + * failed with a retryable error) or "errorNotRetryable" (a later attempt + * failed with a non-retryable error). Caller-owned string. */ +char *aimux_error_retry_reason(const aimux_error_t *error); +/** Number of recorded attempt errors; 0 under any other code or for NULL. */ +int32_t aimux_error_retry_count(const aimux_error_t *error); +/** The attempt error at `index` (0-based, oldest first; the last entry is + * the final attempt) as a NEW owned error — read it with these same getters + * and release it with aimux_error_free(), independently of the parent. + * NULL when `index` is out of range or under any other code. */ +aimux_error_t *aimux_error_retry_error_at(const aimux_error_t *error, + int32_t index); + #ifdef __cplusplus } #endif diff --git a/aimux-ffi/aimux-ffi.h b/aimux-ffi/aimux-ffi.h index 6918c62d..eaa2692e 100644 --- a/aimux-ffi/aimux-ffi.h +++ b/aimux-ffi/aimux-ffi.h @@ -15,7 +15,7 @@ * Getter strings are caller-owned (`aimux_free_string`). See aimux-error.h. * * Each prototype below identifies its expected high-level error: - * [AiMuxError] codes 1..13 + * [AiMuxError] codes 1..14 * [RecordingError] codes 100..105 * [C ABI] no expected high-level code * Every fallible call can additionally return a C ABI failure (200..206). @@ -48,7 +48,7 @@ * ## Concurrency * * All functions are synchronous (block until the operation completes). - * `opts_json.timeout` (`{"total_ms","first_chunk_ms","chunk_ms"}`, + * `opts_json.timeout` (`{"total_ms","step_ms","first_chunk_ms","chunk_ms"}`, * milliseconds; the latter two streaming-only) bounds a call from the callee * side; `aimux_abort_signal_*` plus the `*_with_abort` entry points cancel a * running call from another thread. diff --git a/aimux-ffi/src/lib.rs b/aimux-ffi/src/lib.rs index 3413eba5..0eb72557 100644 --- a/aimux-ffi/src/lib.rs +++ b/aimux-ffi/src/lib.rs @@ -12,7 +12,7 @@ //! value on failure (the out-parameter is left at its sentinel: handle 0, //! pointer NULL). Every non-NULL error has one code from [`aimux_error_code`] //! and one message from [`aimux_error_message`], and is released exactly once -//! with [`aimux_error_free`]. Codes 1..13 come from `AiMuxError`, 100..105 +//! with [`aimux_error_free`]. Codes 1..14 come from `AiMuxError`, 100..105 //! from `RecordingError`, and 200..206 identify failures detected while //! crossing the C ABI. //! @@ -53,6 +53,7 @@ use std::sync::{Arc, Mutex, OnceLock}; use serde::de::DeserializeOwned; +use aimux_core::AbortSignal; use aimux_core::AiMuxError; use aimux_core::generate::{ GenerateTextOptions, generate_object, generate_text, generate_text_as_openai, stream_text, @@ -63,7 +64,6 @@ use aimux_core::message::ModelPrompt; use aimux_core::openai_output::OpenAiStreamOptions; use aimux_core::provider::Provider; use aimux_core::recording::RecordingError; -use aimux_core::shared::AbortSignal; use aimux_core::trace::{RingTraceStore, TraceFilter, TraceLayer}; use aimux_providers::anthropic::{AnthropicConfig, AnthropicProvider}; use aimux_providers::anthropic_aws::{AnthropicAwsProvider, AnthropicAwsProviderConfig}; @@ -424,6 +424,9 @@ pub const AIMUX_E_NO_SUCH_PROVIDER: i32 = 10; pub const AIMUX_E_API_CALL: i32 = 11; pub const AIMUX_E_TIMEOUT: i32 = 12; pub const AIMUX_E_ABORTED: i32 = 13; +// `Retry` (newest variant) reclaims the slot the pre-unification `Other` +// vacated — the opaque-pointer ABI break means no old caller can misread it. +pub const AIMUX_E_RETRY: i32 = 14; // 100..105 preserve `RecordingError` as a separate high-level type while C // uses one code space for every returned error. @@ -452,6 +455,7 @@ pub const AIMUX_TRANSCRIPTION_NEXT_PART_TIMEOUT: i32 = 3; fn aimux_error_code_of(err: &AiMuxError) -> i32 { match err { AiMuxError::ApiCall { .. } => AIMUX_E_API_CALL, + AiMuxError::Retry(_) => AIMUX_E_RETRY, AiMuxError::JsonParse(_) => AIMUX_E_JSON_PARSE, AiMuxError::InvalidResponseData(_) => AIMUX_E_INVALID_RESPONSE_DATA, AiMuxError::Tool(_) => AIMUX_E_TOOL, @@ -462,7 +466,7 @@ fn aimux_error_code_of(err: &AiMuxError) -> i32 { AiMuxError::NoSuchModel { .. } => AIMUX_E_NO_SUCH_MODEL, AiMuxError::NoSuchProvider { .. } => AIMUX_E_NO_SUCH_PROVIDER, AiMuxError::Timeout(_) => AIMUX_E_TIMEOUT, - AiMuxError::Aborted => AIMUX_E_ABORTED, + AiMuxError::Aborted(_) => AIMUX_E_ABORTED, AiMuxError::Other(_) => AIMUX_E_OTHER, } } @@ -506,6 +510,13 @@ fn api_call(e: &AiMuxError) -> Option<&aimux_core::ApiCallError> { } } +fn retry(e: &AiMuxError) -> Option<&aimux_core::RetryError> { + match e { + AiMuxError::Retry(r) => Some(r), + _ => None, + } +} + /// Apply `f` to the `AiMuxError` stored in a returned error. /// /// The closure keeps the borrowed error scoped to this call instead of @@ -588,18 +599,108 @@ pub extern "C" fn aimux_error_provider_message(err: *const aimux_error_t) -> *mu ) } -/// `AIMUX_E_API_CALL`: provider request id. -#[unsafe(no_mangle)] -pub extern "C" fn aimux_error_request_id(err: *const aimux_error_t) -> *mut c_char { - opt_cstring(map_aimux_error(err, |e| api_call(e)?.request_id.clone()).flatten()) -} - /// `AIMUX_E_API_CALL`: raw response body. #[unsafe(no_mangle)] pub extern "C" fn aimux_error_response_body(err: *const aimux_error_t) -> *mut c_char { opt_cstring(map_aimux_error(err, |e| api_call(e)?.response_body.clone()).flatten()) } +/// `AIMUX_E_API_CALL`: sanitized request URL. +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_url(err: *const aimux_error_t) -> *mut c_char { + opt_cstring( + map_aimux_error(err, |e| Some(api_call(e)?.url.clone())) + .flatten() + .filter(|u| !u.is_empty()), + ) +} + +/// `AIMUX_E_API_CALL`: sanitized request body values as a JSON string; NULL +/// when the request carried none (JSON `null`). +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_request_body_values(err: *const aimux_error_t) -> *mut c_char { + opt_cstring( + map_aimux_error(err, |e| { + let values = &api_call(e)?.request_body_values; + if values.is_null() { + return None; + } + serde_json::to_string(values).ok() + }) + .flatten(), + ) +} + +/// `AIMUX_E_API_CALL`: sanitized response headers as one JSON object string +/// of name → value pairs, e.g. `{"retry-after-ms":"1500"}`. +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_response_headers(err: *const aimux_error_t) -> *mut c_char { + opt_cstring( + map_aimux_error(err, |e| { + serde_json::to_string(api_call(e)?.response_headers.as_ref()?).ok() + }) + .flatten(), + ) +} + +/// `AIMUX_E_API_CALL`: parsed provider error data as a JSON string. +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_provider_data(err: *const aimux_error_t) -> *mut c_char { + opt_cstring( + map_aimux_error(err, |e| { + serde_json::to_string(api_call(e)?.data.as_ref()?).ok() + }) + .flatten(), + ) +} + +/// `AIMUX_E_RETRY`: why retrying stopped — the wire name of +/// `RetryErrorReason`: "maxRetriesExceeded" or "errorNotRetryable". +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_retry_reason(err: *const aimux_error_t) -> *mut c_char { + let reason = map_aimux_error(err, |e| { + Some(match retry(e)?.reason { + aimux_core::RetryErrorReason::MaxRetriesExceeded => "maxRetriesExceeded", + aimux_core::RetryErrorReason::ErrorNotRetryable => "errorNotRetryable", + }) + }) + .flatten(); + opt_cstring(reason.map(str::to_string)) +} + +/// `AIMUX_E_RETRY`: number of recorded attempt errors; 0 under any other +/// code or for NULL. +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_retry_count(err: *const aimux_error_t) -> i32 { + map_aimux_error(err, |e| Some(retry(e)?.errors.len() as i32)) + .flatten() + .unwrap_or(0) +} + +/// `AIMUX_E_RETRY`: the attempt error at `index` (0-based, oldest first; the +/// last entry is the final attempt) as a NEW owned error the caller reads +/// with these same getters and releases with `aimux_error_free`, +/// independently of the parent. NULL when `index` is out of range or under +/// any other code. +#[unsafe(no_mangle)] +pub extern "C" fn aimux_error_retry_error_at( + err: *const aimux_error_t, + index: i32, +) -> *mut aimux_error_t { + if index < 0 { + return std::ptr::null_mut(); + } + map_aimux_error(err, |e| { + retry(e)?.errors.get(index as usize).map(|attempt| { + Box::into_raw(Box::new(aimux_error_t { + error: AiMuxFfiError::AiMux(attempt.clone()), + })) + }) + }) + .flatten() + .unwrap_or_else(std::ptr::null_mut) +} + /// `AIMUX_E_NO_SUCH_MODEL`: the model id that was asked for. #[unsafe(no_mangle)] pub extern "C" fn aimux_error_model_id(err: *const aimux_error_t) -> *mut c_char { @@ -1696,7 +1797,7 @@ fn stream_text_as_openai_with_signal( Some(signal) => { tokio::select! { biased; - _ = signal.cancelled() => Err(AiMuxError::Aborted), + _ = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), result = stream_text_as_openai(&*model, prompt, opts, stream_options) => result, } } @@ -1708,7 +1809,9 @@ fn stream_text_as_openai_with_signal( Some(signal) => { tokio::select! { biased; - _ = signal.cancelled() => return Err(AiMuxError::Aborted.into()), + _ = signal.cancelled() => { + return Err(AiMuxError::from_abort_signal(signal).into()); + } item = stream.next() => item, } } @@ -1719,7 +1822,16 @@ fn stream_text_as_openai_with_signal( }; // A chunk that cannot be serialized ends the stream with this // layer's ResultSerialization — never a silent "{}" placeholder. - let cstr = stream_part_cstring(&item?)?; + // Core yields a malformed frame as a recoverable Err and keeps the + // stream alive. Every consumer of this path types items as + // ChatCompletionChunk, so the error cannot ride this wire — skip + // it and keep pumping (full fidelity lives on the plain + // StreamPart path). + let cstr = match item { + Ok(chunk) => stream_part_cstring(&chunk)?, + Err(e) if e.is_recoverable_stream_error() => continue, + Err(e) => return Err(e.into()), + }; invoke_stream_callback("on_part", || { on_part(cstr.as_ptr(), stream_ctx as *mut c_void); })?; @@ -1758,7 +1870,7 @@ fn stream_text_with_signal( Some(signal) => { tokio::select! { biased; - _ = signal.cancelled() => Err(AiMuxError::Aborted), + _ = signal.cancelled() => Err(AiMuxError::from_abort_signal(signal)), result = stream_text(&*model, prompt, opts) => result, } } @@ -1770,7 +1882,9 @@ fn stream_text_with_signal( Some(signal) => { tokio::select! { biased; - _ = signal.cancelled() => return Err(AiMuxError::Aborted.into()), + _ = signal.cancelled() => { + return Err(AiMuxError::from_abort_signal(signal).into()); + } item = stream.next() => item, } } @@ -1781,7 +1895,16 @@ fn stream_text_with_signal( }; // A part that cannot be serialized ends the stream with this // layer's ResultSerialization — never a silent "{}" placeholder. - let cstr = stream_part_cstring(&item?)?; + // Core keeps the stream alive across a malformed frame; mirror it + // by delivering the error as a StreamPart::Error data frame (the + // documented non-terminal shape on this path) and keep pumping. + let cstr = match item { + Ok(part) => stream_part_cstring(&part)?, + Err(e) if e.is_recoverable_stream_error() => { + stream_part_cstring(&aimux_core::stream_part::StreamPart::Error { error: e })? + } + Err(e) => return Err(e.into()), + }; invoke_stream_callback("on_part", || { on_part(cstr.as_ptr(), stream_ctx as *mut c_void); })?; @@ -1950,7 +2073,7 @@ pub extern "C" fn aimux_embed( opts = serde_json::from_str(&s).map_err(|e| wire_err("opts_json", e))?; } opts.values = serde_json::from_str(&values_json).map_err(|e| wire_err("values_json", e))?; - run_json(async move { model.do_embed(&opts).await }) + run_json(async move { aimux_core::embedding_model::embed(model.as_ref(), opts).await }) }) } @@ -2001,7 +2124,9 @@ pub extern "C" fn aimux_speech_generate( }; let opts: aimux_core::speech_model::SpeechCallOptions = parse_json_arg(opts_json, "opts_json")?; - run_json(async move { model.do_generate(&opts).await }) + run_json( + async move { aimux_core::speech_model::generate_speech(model.as_ref(), opts).await }, + ) }) } @@ -2083,7 +2208,7 @@ pub extern "C" fn aimux_image_generate( }; let opts: aimux_core::image_model::ImageCallOptions = parse_json_arg(opts_json, "opts_json")?; - run_json(async move { model.do_generate(&opts).await }) + run_json(async move { aimux_core::image_model::generate_image(model.as_ref(), opts).await }) }) } @@ -2143,7 +2268,9 @@ pub extern "C" fn aimux_transcription_generate( aimux_core::transcription_model::AudioInput::Base64(audio_base64), media_type, ); - run_json(async move { model.do_generate(&opts).await }) + run_json( + async move { aimux_core::transcription_model::transcribe(model.as_ref(), opts).await }, + ) }) } @@ -2454,7 +2581,7 @@ pub extern "C" fn aimux_rerank( }; let opts: aimux_core::reranking_model::RerankingCallOptions = parse_json_arg(opts_json, "opts_json")?; - run_json(async move { model.do_rerank(&opts).await }) + run_json(async move { aimux_core::reranking_model::rerank(model.as_ref(), opts).await }) }) } @@ -2508,7 +2635,7 @@ pub extern "C" fn aimux_video_generate( }; let opts: aimux_core::video_model::VideoCallOptions = parse_json_arg(opts_json, "opts_json")?; - run_json(async move { model.do_generate(&opts).await }) + run_json(async move { aimux_core::video_model::generate_video(model.as_ref(), opts).await }) }) } @@ -2563,7 +2690,7 @@ pub extern "C" fn aimux_search( }; let opts: aimux_core::search_model::SearchCallOptions = parse_json_arg(opts_json, "opts_json")?; - run_json(async move { model.do_search(&opts).await }) + run_json(async move { aimux_core::search_model::search(model.as_ref(), opts).await }) }) } @@ -3105,7 +3232,7 @@ struct RouterFfiConfig { #[cfg(test)] mod tests { - use aimux_core::ApiCallError; + use aimux_core::{ApiCallError, RetryError, RetryErrorReason}; use std::time::Duration; use super::*; @@ -3137,7 +3264,7 @@ mod tests { fn expect_aimux_error(e: *mut aimux_error_t) -> (i32, String) { assert!(!e.is_null(), "expected a returned error"); let code = aimux_error_code(e); - if !(AIMUX_E_OTHER..=AIMUX_E_ABORTED).contains(&code) { + if !(AIMUX_E_OTHER..=AIMUX_E_RETRY).contains(&code) { panic!("expected an AiMuxError code, got {code}: {}", msg(e)); } let out = (code, take(aimux_error_message(e)).unwrap()); @@ -3166,18 +3293,26 @@ mod tests { /// NULL / -1 / 0 under every other; a NULL pointer answers as "no error". #[test] fn payload_getters_follow_the_code() { - let owner = boxed(AiMuxError::ApiCall(ApiCallError { + let owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError { status_code: Some(429), provider_code: Some("insufficient_quota".into()), - message: "quota".into(), response_body: Some("{}".into()), - request_id: Some("req_1".into()), - retry_after_ms: Some(1500), + response_headers: Some(std::collections::HashMap::from([( + "retry-after-ms".to_string(), + "1500".to_string(), + )])), + data: Some(serde_json::json!({"code": "insufficient_quota"})), is_retryable: true, - })); + ..ApiCallError::new( + "quota", + "https://api.example.test/v1", + serde_json::json!({"model": "m"}), + ) + }))); let h = owner; assert_eq!(aimux_error_code(h), AIMUX_E_API_CALL); assert_eq!(aimux_error_status(h), 429); + // The retry hint is derived from the response headers. assert_eq!(aimux_error_retry_ms(h), 1500); assert_eq!(aimux_error_retryable(h), 1); assert_eq!( @@ -3191,25 +3326,106 @@ mod tests { ); let m = take(aimux_error_message(h)).unwrap(); assert!(m.contains("quota") && m != "quota", "{m}"); - assert_eq!(take(aimux_error_request_id(h)).as_deref(), Some("req_1")); assert_eq!(take(aimux_error_response_body(h)).as_deref(), Some("{}")); + assert_eq!( + take(aimux_error_url(h)).as_deref(), + Some("https://api.example.test/v1") + ); + assert_eq!( + take(aimux_error_request_body_values(h)).as_deref(), + Some(r#"{"model":"m"}"#) + ); + assert_eq!( + take(aimux_error_response_headers(h)).as_deref(), + Some(r#"{"retry-after-ms":"1500"}"#) + ); + assert_eq!( + take(aimux_error_provider_data(h)).as_deref(), + Some(r#"{"code":"insufficient_quota"}"#) + ); assert!(aimux_error_model_id(h).is_null()); assert!(aimux_error_model_type(h).is_null()); assert!(aimux_error_provider_id(h).is_null()); + // Retry getters answer their sentinels under AIMUX_E_API_CALL. + assert!(aimux_error_retry_reason(h).is_null()); + assert_eq!(aimux_error_retry_count(h), 0); + assert!(aimux_error_retry_error_at(h, 0).is_null()); aimux_error_free(owner); - // Absent Option fields are NULL, not empty strings. - let owner = boxed(AiMuxError::ApiCall(ApiCallError { - message: "x".into(), - ..Default::default() - })); + // Absent Option fields are NULL, not empty strings; a JSON `null` + // request body is NULL too. + let owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError::new( + "x", + "https://api.example.test/v1", + serde_json::Value::Null, + )))); let h = owner; assert!(aimux_error_provider_code(h).is_null()); - assert!(aimux_error_request_id(h).is_null()); assert!(aimux_error_response_body(h).is_null()); + assert!(aimux_error_response_headers(h).is_null()); + assert!(aimux_error_provider_data(h).is_null()); + assert!(aimux_error_request_body_values(h).is_null()); assert_eq!(take(aimux_error_provider_message(h)).as_deref(), Some("x")); aimux_error_free(owner); + // Retry owns reason + attempt history; each history entry crosses the + // ABI as a new owned error read with the same getters. + let owner = boxed(AiMuxError::Retry(RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors: vec![ + AiMuxError::Other("first".into()), + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(500), + is_retryable: true, + ..ApiCallError::new( + "boom", + "https://api.example.test/v1", + serde_json::Value::Null, + ) + })), + ], + })); + let h = owner; + assert_eq!(aimux_error_code(h), AIMUX_E_RETRY); + assert_eq!( + take(aimux_error_retry_reason(h)).as_deref(), + Some("maxRetriesExceeded") + ); + assert_eq!(aimux_error_retry_count(h), 2); + // ApiCall-owned facts stay with the attempt that owns them. + assert_eq!(aimux_error_status(h), -1); + assert_eq!(aimux_error_retryable(h), 0); + assert!(aimux_error_provider_code(h).is_null()); + let m = take(aimux_error_message(h)).unwrap(); + assert!(m.contains("Failed after 2 attempts"), "{m}"); + let first = aimux_error_retry_error_at(h, 0); + assert_eq!(aimux_error_code(first), AIMUX_E_OTHER); + assert_eq!(take(aimux_error_message(first)).as_deref(), Some("first")); + aimux_error_free(first); + let last = aimux_error_retry_error_at(h, 1); + assert_eq!(aimux_error_code(last), AIMUX_E_API_CALL); + assert_eq!(aimux_error_status(last), 500); + assert_eq!(aimux_error_retryable(last), 1); + // The history entry outlives the parent: free order is unconstrained. + aimux_error_free(owner); + assert_eq!( + take(aimux_error_provider_message(last)).as_deref(), + Some("boom") + ); + aimux_error_free(last); + let owner = boxed(AiMuxError::Retry(RetryError { + reason: RetryErrorReason::ErrorNotRetryable, + errors: vec![AiMuxError::Other("bad".into())], + })); + let h = owner; + assert_eq!( + take(aimux_error_retry_reason(h)).as_deref(), + Some("errorNotRetryable") + ); + assert!(aimux_error_retry_error_at(h, 1).is_null(), "out of range"); + assert!(aimux_error_retry_error_at(h, -1).is_null()); + aimux_error_free(owner); + let owner = boxed(AiMuxError::NoSuchModel { model_id: "m".into(), model_type: "language".into(), @@ -3241,26 +3457,37 @@ mod tests { assert_eq!(aimux_error_retry_ms(n), -1); assert_eq!(aimux_error_retryable(n), 0); assert!(aimux_error_provider_code(n).is_null()); + assert!(aimux_error_url(n).is_null()); + assert!(aimux_error_request_body_values(n).is_null()); + assert!(aimux_error_response_headers(n).is_null()); + assert!(aimux_error_provider_data(n).is_null()); + assert!(aimux_error_retry_reason(n).is_null()); + assert_eq!(aimux_error_retry_count(n), 0); + assert!(aimux_error_retry_error_at(n, 0).is_null()); } /// The retry verdict crosses the ABI as its own getter. It is not /// derivable from `status`: both cases below report -1, and they disagree. #[test] fn retryable_crosses_the_abi_and_status_cannot_stand_in() { - let transport_owner = boxed(AiMuxError::ApiCall(ApiCallError { - message: "connection reset".into(), + let transport_owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError { is_retryable: true, - ..Default::default() - })); + ..ApiCallError::new( + "connection reset", + "https://api.example.test/v1", + serde_json::Value::Null, + ) + }))); let transport = transport_owner; assert_eq!(aimux_error_status(transport), -1); assert_eq!(aimux_error_retryable(transport), 1); aimux_error_free(transport_owner); - let no_key_owner = boxed(AiMuxError::ApiCall(ApiCallError { - message: "no api key".into(), - ..Default::default() - })); + let no_key_owner = boxed(AiMuxError::ApiCall(Box::new(ApiCallError::new( + "no api key", + "https://api.example.test/v1", + serde_json::Value::Null, + )))); let no_key = no_key_owner; assert_eq!( aimux_error_status(no_key), @@ -3281,9 +3508,17 @@ mod tests { #[test] fn unified_error_codes_cover_every_source() { // AiMuxError. - let e = boxed(AiMuxError::Aborted); + let e = boxed(AiMuxError::Aborted("request aborted".into())); assert_eq!(aimux_error_code(e), AIMUX_E_ABORTED); - assert_eq!(msg(e), AiMuxError::Aborted.to_string()); + assert_eq!(msg(e), "request aborted"); + + // Retry extends the contiguous run to 1..14. + let e = boxed(AiMuxError::Retry(RetryError { + reason: RetryErrorReason::ErrorNotRetryable, + errors: vec![AiMuxError::Other("bad".into())], + })); + assert_eq!(aimux_error_code(e), AIMUX_E_RETRY); + assert!(msg(e).contains("non-retryable")); // Recording. use RecordingError as R; @@ -3385,18 +3620,26 @@ mod tests { ); } - /// Pin the full 13-variant → code mapping. + /// Pin the full 14-variant → code mapping. #[test] fn error_code_mapping_covers_all_variants() { let s = |t: &str| t.to_string(); let cases: Vec<(AiMuxError, i32)> = vec![ ( - AiMuxError::ApiCall(ApiCallError { - message: s("x"), - ..Default::default() - }), + AiMuxError::ApiCall(Box::new(ApiCallError::new( + "x", + "https://example.test", + serde_json::json!({}), + ))), AIMUX_E_API_CALL, ), + ( + AiMuxError::Retry(RetryError { + reason: RetryErrorReason::MaxRetriesExceeded, + errors: vec![AiMuxError::Other(s("first")), AiMuxError::Other(s("last"))], + }), + AIMUX_E_RETRY, + ), (AiMuxError::JsonParse(s("x")), AIMUX_E_JSON_PARSE), ( AiMuxError::InvalidResponseData(s("x")), @@ -3427,7 +3670,10 @@ mod tests { AIMUX_E_NO_SUCH_PROVIDER, ), (AiMuxError::Timeout(s("x")), AIMUX_E_TIMEOUT), - (AiMuxError::Aborted, AIMUX_E_ABORTED), + ( + AiMuxError::Aborted("request aborted".into()), + AIMUX_E_ABORTED, + ), (AiMuxError::Other(s("x")), AIMUX_E_OTHER), ]; for (e, code) in cases { @@ -3621,6 +3867,157 @@ mod tests { })) } + /// Streams: delta, recoverable Err(JsonParse), delta, Finish. + struct RecoverableFrameModel; + #[async_trait::async_trait] + impl aimux_core::LanguageModel for RecoverableFrameModel { + fn provider(&self) -> &str { + "mock" + } + fn model_id(&self) -> &str { + "recoverable" + } + async fn do_generate( + &self, + _options: &aimux_core::options::CallOptions, + ) -> Result { + unimplemented!() + } + async fn do_stream( + &self, + _options: &aimux_core::options::CallOptions, + ) -> Result { + let delta = |text: &str| aimux_core::stream_part::StreamPart::TextDelta { + id: "1".into(), + delta: text.into(), + provider_metadata: None, + }; + Ok(aimux_core::result::StreamResult { + stream: Box::pin(futures::stream::iter([ + Ok(delta("a")), + Err(AiMuxError::JsonParse("bad frame".into())), + Ok(delta("b")), + Ok(aimux_core::stream_part::StreamPart::Finish { + finish_reason: aimux_core::types::FinishReason { + unified: aimux_core::types::FinishReasonUnified::Stop, + raw: None, + }, + usage: aimux_core::types::Usage::default(), + provider_metadata: None, + }), + ])), + request_body: None, + response_headers: None, + }) + } + } + + /// The pump must deliver a recoverable frame error as a StreamPart::Error + /// data part and keep going — parts after it still arrive and on_done + /// fires (the pre-fix pump returned the error and skipped on_done). + #[test] + fn recoverable_frame_error_does_not_end_the_ffi_stream() { + struct Collected { + parts: Vec, + done: bool, + } + extern "C-unwind" fn on_part(json: *const c_char, ctx: *mut c_void) { + let collected = unsafe { &mut *(ctx as *mut Collected) }; + let text = unsafe { std::ffi::CStr::from_ptr(json) } + .to_string_lossy() + .into_owned(); + collected.parts.push(text); + } + extern "C-unwind" fn on_done(ctx: *mut c_void) { + let collected = unsafe { &mut *(ctx as *mut Collected) }; + collected.done = true; + } + + let handle = intern_model(Arc::new(RecoverableFrameModel)); + let mut collected = Collected { + parts: Vec::new(), + done: false, + }; + let prompt = std::ffi::CString::new("\"hi\"").unwrap(); + let e = aimux_stream_text( + handle, + prompt.as_ptr(), + std::ptr::null(), + Some(on_part), + Some(on_done), + &mut collected as *mut Collected as *mut c_void, + ); + assert!(e.is_null(), "{}", msg(e)); + assert!( + collected.done, + "on_done must fire after a recoverable frame" + ); + let error_index = collected + .parts + .iter() + .position(|p| p.contains("\"Error\"")) + .expect("the recoverable frame arrives as a StreamPart::Error part"); + assert!( + collected.parts[error_index + 1..] + .iter() + .any(|p| p.contains("\"b\"")), + "parts after the recoverable frame still arrive: {:?}", + collected.parts + ); + } + + /// On the OpenAI-compatible path every item is a ChatCompletionChunk, so a + /// recoverable frame error is skipped — no error item, no truncation, and + /// on_done still fires. + #[test] + fn recoverable_frame_error_is_skipped_on_the_openai_stream() { + struct Collected { + parts: Vec, + done: bool, + } + extern "C-unwind" fn on_part(json: *const c_char, ctx: *mut c_void) { + let collected = unsafe { &mut *(ctx as *mut Collected) }; + let text = unsafe { std::ffi::CStr::from_ptr(json) } + .to_string_lossy() + .into_owned(); + collected.parts.push(text); + } + extern "C-unwind" fn on_done(ctx: *mut c_void) { + let collected = unsafe { &mut *(ctx as *mut Collected) }; + collected.done = true; + } + + let handle = intern_model(Arc::new(RecoverableFrameModel)); + let mut collected = Collected { + parts: Vec::new(), + done: false, + }; + let prompt = std::ffi::CString::new("\"hi\"").unwrap(); + let e = aimux_stream_text_as_openai( + handle, + prompt.as_ptr(), + std::ptr::null(), + Some(on_part), + Some(on_done), + &mut collected as *mut Collected as *mut c_void, + ); + assert!(e.is_null(), "{}", msg(e)); + assert!( + collected.done, + "on_done must fire after a recoverable frame" + ); + assert!( + collected.parts.iter().all(|p| !p.contains("\"error\"")), + "no error item may ride the chunk-typed wire: {:?}", + collected.parts + ); + assert!( + collected.parts.iter().any(|p| p.contains("\"b\"")), + "chunks after the recoverable frame still arrive: {:?}", + collected.parts + ); + } + #[test] fn router_new_builds_router_model() { let handles = [ @@ -3798,7 +4195,7 @@ mod tests { }; tokio::select! { _ = aborted => { - yield Err(AiMuxError::Aborted); + yield Err(AiMuxError::Aborted("request aborted".into())); break; } chunk = audio.next() => match chunk { diff --git a/aimux-ffi/src/transcription_session.rs b/aimux-ffi/src/transcription_session.rs index e51a9038..d6d71054 100644 --- a/aimux-ffi/src/transcription_session.rs +++ b/aimux-ffi/src/transcription_session.rs @@ -6,7 +6,7 @@ //! push / input-done / next-part operations. //! //! ```text -//! session_new ──► spawn tokio task ──► TranscriptionModel::do_stream(audio_rx) +//! session_new ──► spawn tokio task ──► stream_transcribe(model, audio_rx) //! │ parts flow out //! push_audio ──► bounded mpsc sender ▼ //! input_done ──► drop sender (= end of audio) @@ -29,8 +29,8 @@ use std::time::Duration; use futures::channel::mpsc; use futures::{SinkExt, StreamExt}; +use aimux_core::AbortSignal; use aimux_core::error::AiMuxError; -use aimux_core::shared::AbortSignal; use aimux_core::transcription_model::{ AudioChunk, InputAudioFormat, TranscriptionModel, TranscriptionStreamOptions, TranscriptionStreamPart, @@ -72,7 +72,7 @@ pub struct TranscriptionFfiSession { } impl TranscriptionFfiSession { - /// Spawn the driver for `model.do_stream` and return the session. + /// Spawn the live-transcription operation and return the session. /// /// `user_abort` (optional) and the internal drop token are both linked /// into the effective abort signal passed to the model: either firing @@ -111,7 +111,8 @@ impl TranscriptionFfiSession { include_raw_chunks: opts.include_raw_chunks, timeout: opts.timeout, }; - let result = model.do_stream(options).await; + let result = + aimux_core::transcription_model::stream_transcribe(model.as_ref(), options).await; match result { Ok(stream_result) => { let mut stream = stream_result.stream; diff --git a/aimux-ffi/tests/common/mod.rs b/aimux-ffi/tests/common/mod.rs index e69cb6e0..e0f76a45 100644 --- a/aimux-ffi/tests/common/mod.rs +++ b/aimux-ffi/tests/common/mod.rs @@ -6,8 +6,8 @@ use std::ffi::{CStr, CString}; use std::os::raw::c_char; use aimux_ffi::{ - AIMUX_E_ABORTED, AIMUX_E_FFI_CALLBACK_FAILURE, AIMUX_E_FFI_NULL_POINTER, AIMUX_E_OTHER, - AIMUX_E_RECORDING_INIT, AIMUX_E_RECORDING_WRITE, aimux_error_code, aimux_error_free, + AIMUX_E_FFI_CALLBACK_FAILURE, AIMUX_E_FFI_NULL_POINTER, AIMUX_E_OTHER, AIMUX_E_RECORDING_INIT, + AIMUX_E_RECORDING_WRITE, AIMUX_E_RETRY, aimux_error_code, aimux_error_free, aimux_error_message, aimux_error_t, aimux_free_string, }; @@ -48,7 +48,7 @@ pub fn ok(e: *mut aimux_error_t, name: &str) { pub fn expect_aimux_error(e: *mut aimux_error_t, name: &str) -> (i32, String) { assert!(!e.is_null(), "{name}: expected a returned error"); let code = aimux_error_code(e); - if !(AIMUX_E_OTHER..=AIMUX_E_ABORTED).contains(&code) { + if !(AIMUX_E_OTHER..=AIMUX_E_RETRY).contains(&code) { panic!( "{name}: expected an AiMuxError code, got {code}: {}", msg(e) diff --git a/aimux-ffi/tests/exports_smoke_test.rs b/aimux-ffi/tests/exports_smoke_test.rs index 3fbaa696..14786b25 100644 --- a/aimux-ffi/tests/exports_smoke_test.rs +++ b/aimux-ffi/tests/exports_smoke_test.rs @@ -52,31 +52,31 @@ use aimux_ffi::{ aimux_cohere_reranking_new_with_base, aimux_consume_stream_text, aimux_drop_handle, aimux_embed, aimux_error_code, aimux_error_free, aimux_error_message, aimux_error_model_id, aimux_error_model_type, aimux_error_provider_code, aimux_error_provider_id, - aimux_error_provider_message, aimux_error_request_id, aimux_error_response_body, - aimux_error_retry_ms, aimux_error_retryable, aimux_error_status, aimux_error_t, - aimux_file_upload, aimux_free_string, aimux_generate_object, aimux_generate_text, - aimux_generate_text_as_openai, aimux_get_model_specs, aimux_google_embedding_new, - aimux_google_embedding_new_with_base, aimux_google_image_new, aimux_google_image_new_with_base, - aimux_google_video_new, aimux_google_video_new_with_base, aimux_image_generate, - aimux_init_logging, aimux_init_proxy, aimux_init_recording, aimux_init_recording_ring, - aimux_init_recording_ring_default, aimux_list_sessions, aimux_mistral_new, - aimux_mistral_new_with_base, aimux_moa_new, aimux_mock_replay_new, aimux_openai_embedding_new, - aimux_openai_embedding_new_with_base, aimux_openai_files_new, aimux_openai_files_new_with_base, - aimux_openai_image_new, aimux_openai_image_new_with_base, aimux_openai_new, - aimux_openai_new_with_base, aimux_openai_speech_new, aimux_openai_speech_new_with_base, - aimux_openai_transcription_new, aimux_openai_transcription_new_with_base, - aimux_provider_from_env, aimux_provider_handle_new, aimux_provider_list_models, - aimux_provider_model, aimux_provider_new, aimux_recording_flush, aimux_recording_stop, - aimux_recording_try_flush, aimux_register_providers, aimux_rerank, aimux_router_new, - aimux_search, aimux_session_calls, aimux_session_infer_init, aimux_session_store_init, - aimux_speech_generate, aimux_stream_text, aimux_stream_text_as_openai, - aimux_stream_text_as_openai_with_abort, aimux_stream_text_with_abort, aimux_tavily_search_new, - aimux_tavily_search_new_with_base, aimux_trace_aggregate, aimux_trace_clear, - aimux_trace_export_jsonl, aimux_trace_new, aimux_trace_new_audited, aimux_trace_session_chain, - aimux_trace_session_trajectory, aimux_transcription_generate, aimux_transcription_input_done, - aimux_transcription_next_part, aimux_transcription_push_audio, - aimux_transcription_session_drop, aimux_transcription_session_new, aimux_vertex_new, - aimux_vertex_new_with_base, aimux_video_generate, aimux_xai_new, aimux_xai_new_with_base, + aimux_error_provider_message, aimux_error_response_body, aimux_error_retry_ms, + aimux_error_retryable, aimux_error_status, aimux_error_t, aimux_file_upload, aimux_free_string, + aimux_generate_object, aimux_generate_text, aimux_generate_text_as_openai, + aimux_get_model_specs, aimux_google_embedding_new, aimux_google_embedding_new_with_base, + aimux_google_image_new, aimux_google_image_new_with_base, aimux_google_video_new, + aimux_google_video_new_with_base, aimux_image_generate, aimux_init_logging, aimux_init_proxy, + aimux_init_recording, aimux_init_recording_ring, aimux_init_recording_ring_default, + aimux_list_sessions, aimux_mistral_new, aimux_mistral_new_with_base, aimux_moa_new, + aimux_mock_replay_new, aimux_openai_embedding_new, aimux_openai_embedding_new_with_base, + aimux_openai_files_new, aimux_openai_files_new_with_base, aimux_openai_image_new, + aimux_openai_image_new_with_base, aimux_openai_new, aimux_openai_new_with_base, + aimux_openai_speech_new, aimux_openai_speech_new_with_base, aimux_openai_transcription_new, + aimux_openai_transcription_new_with_base, aimux_provider_from_env, aimux_provider_handle_new, + aimux_provider_list_models, aimux_provider_model, aimux_provider_new, aimux_recording_flush, + aimux_recording_stop, aimux_recording_try_flush, aimux_register_providers, aimux_rerank, + aimux_router_new, aimux_search, aimux_session_calls, aimux_session_infer_init, + aimux_session_store_init, aimux_speech_generate, aimux_stream_text, + aimux_stream_text_as_openai, aimux_stream_text_as_openai_with_abort, + aimux_stream_text_with_abort, aimux_tavily_search_new, aimux_tavily_search_new_with_base, + aimux_trace_aggregate, aimux_trace_clear, aimux_trace_export_jsonl, aimux_trace_new, + aimux_trace_new_audited, aimux_trace_session_chain, aimux_trace_session_trajectory, + aimux_transcription_generate, aimux_transcription_input_done, aimux_transcription_next_part, + aimux_transcription_push_audio, aimux_transcription_session_drop, + aimux_transcription_session_new, aimux_vertex_new, aimux_vertex_new_with_base, + aimux_video_generate, aimux_xai_new, aimux_xai_new_with_base, }; use common::{c, expect_aimux_error, expect_failure, expect_ffi_error, msg, ok, take}; @@ -155,7 +155,7 @@ fn header_and_exports_agree() { exports.sort(); assert_eq!( exports.len(), - 109, + 115, "export count changed; update the headers" ); @@ -791,10 +791,19 @@ fn video_generate_fails_cleanly_on_unreachable_host() { &mut h, ); expect_handle(e, h, "video handle"); - let opts = c(r#"{"prompt":"waves","n":1,"provider_options":{}}"#); + // `n` and `provider_options` are omitted and `poll` is supplied: this + // pins the serde defaults and the poll wire keys at the ABI boundary — + // typed binding structs omit unset fields, so a strict parse here would + // regress every C-ABI language at once. + let opts = c(r#"{"prompt":"waves","poll":{"interval_ms":1,"timeout_ms":1}}"#); let mut out: *mut c_char = ptr::null_mut(); let e = aimux_video_generate(h, opts.as_ptr(), &mut out); - expect_ptr_aimux_failure(e, out, "video_generate"); + assert!(out.is_null(), "video_generate: expected NULL out-param"); + let (code, message) = expect_aimux_error(e, "video_generate"); + assert_ne!( + code, AIMUX_E_INVALID_ARGUMENT, + "options with defaults omitted must parse; got: {message}" + ); aimux_drop_handle(h); } @@ -978,7 +987,6 @@ fn utility_exports_return_clean_values() { for get in [ aimux_error_provider_code, aimux_error_provider_message, - aimux_error_request_id, aimux_error_response_body, aimux_error_model_id, aimux_error_model_type, diff --git a/bindings/c/example.c b/bindings/c/example.c index 60947e4c..f1957bcf 100644 --- a/bindings/c/example.c +++ b/bindings/c/example.c @@ -72,7 +72,6 @@ static int report(const char *what, aimux_error_t *e) { // Payload strings are NULL when the provider did not send them. print_owned("provider code", aimux_error_provider_code(e)); print_owned("provider message", aimux_error_provider_message(e)); - print_owned("request id", aimux_error_request_id(e)); break; } case AIMUX_E_TOKEN_EXPIRED: diff --git a/bindings/c/example.cpp b/bindings/c/example.cpp index 666ad4d3..be832c5b 100644 --- a/bindings/c/example.cpp +++ b/bindings/c/example.cpp @@ -32,7 +32,6 @@ class AimuxException : public std::runtime_error { retryable_(aimux_error_retryable(e) != 0), provider_code_(take(aimux_error_provider_code(e))), provider_message_(take(aimux_error_provider_message(e))), - request_id_(take(aimux_error_request_id(e))), response_body_(take(aimux_error_response_body(e))), model_id_(take(aimux_error_model_id(e))), model_type_(take(aimux_error_model_type(e))), provider_id_(take(aimux_error_provider_id(e))) { @@ -49,7 +48,6 @@ class AimuxException : public std::runtime_error { // Per-code payload; "" when the code does not carry the field. const std::string &providerCode() const { return provider_code_; } // API_CALL const std::string &providerMessage() const { return provider_message_; } // API_CALL - const std::string &requestId() const { return request_id_; } // API_CALL const std::string &responseBody() const { return response_body_; } // API_CALL const std::string &modelId() const { return model_id_; } // NO_SUCH_MODEL const std::string &modelType() const { return model_type_; } // NO_SUCH_MODEL @@ -60,16 +58,16 @@ class AimuxException : public std::runtime_error { int status_ = -1; int64_t retry_ms_ = -1; bool retryable_ = false; - std::string provider_code_, provider_message_, request_id_, response_body_, model_id_, + std::string provider_code_, provider_message_, response_body_, model_id_, model_type_, provider_id_; }; -// An AiMuxError code (1..13) becomes AimuxException. Codes 200..206 mean this +// An AiMuxError code (1..14) becomes AimuxException. Codes 200..206 mean this // program made a bad call (NULL argument, malformed JSON, dead handle). static void throw_if_failed(aimux_error_t *err, const std::string &what) { if (!err) return; int32_t code = aimux_error_code(err); - if (code >= AIMUX_E_OTHER && code <= AIMUX_E_ABORTED) { + if (code >= AIMUX_E_OTHER && code <= AIMUX_E_RETRY) { throw AimuxException(what, err); } std::string msg = take(aimux_error_message(err)); From d3d766fe477e55599b3473da64189621aa905a4b Mon Sep 17 00:00:00 2001 From: Card Cunningham Date: Mon, 24 Aug 2026 22:29:11 +0800 Subject: [PATCH 05/24] feat(bindings)!: expose retry, timeout, and video surfaces Node, Python, Go, Java, Kotlin, Swift, and Flutter carry the pipeline's typed error context (retryable, retry_ms, provider code/message, request id, response body) and the VideoPollOptions surface. The Node napi and Python pyo3 stream drivers forward recoverable frame errors as data and keep pumping (skipping them on the chunk-typed OpenAI path); terminal errors still end the stream. ts-rs generated types are regenerated and gated by scripts/gen_ts_types.py --check; web example fixtures pick up the recording step field. --- bindings/flutter/lib/aimux.dart | 6 +- bindings/flutter/lib/errors.dart | 149 +++++++++++-- bindings/flutter/lib/types.dart | 6 +- bindings/flutter/lib/types.g.dart | 2 + bindings/flutter/test/errors_test.dart | 74 ++++++- bindings/go/aimux.go | 45 +++- bindings/go/error.go | 67 ++++-- bindings/go/error_test.go | 38 +++- bindings/go/multimodal_test.go | 27 +++ bindings/go/multimodal_types.go | 109 ++++++---- bindings/go/types.go | 1 + bindings/go/wire_format_test.go | 8 +- .../ai/arcships/aimux/AimuxException.java | 174 ++++++++++++++-- .../main/java/ai/arcships/aimux/AimuxFFI.java | 24 ++- .../java/ai/arcships/aimux/AimuxResult.java | 7 +- .../main/java/ai/arcships/aimux/Model.java | 6 +- .../ai/arcships/aimux/MultimodalTypes.java | 196 +++++++++++++++--- .../main/java/ai/arcships/aimux/Types.java | 12 +- .../ai/arcships/aimux/AimuxExceptionTest.java | 25 ++- .../arcships/aimux/MultimodalTypesTest.java | 20 ++ .../main/kotlin/ai/arcships/aimux/Errors.kt | 133 ++++++++++-- .../main/kotlin/ai/arcships/aimux/Model.kt | 24 ++- .../kotlin/ai/arcships/aimux/Multimodal.kt | 2 +- .../ai/arcships/aimux/MultimodalTypes.kt | 22 ++ .../main/kotlin/ai/arcships/aimux/Types.kt | 2 + .../kotlin/ai/arcships/aimux/ContractTest.kt | 16 ++ .../kotlin/ai/arcships/aimux/ErrorsTest.kt | 51 ++++- bindings/node/Cargo.toml | 2 +- bindings/node/__test__/error.test.ts | 46 +++- bindings/node/__test__/index.test.ts | 7 + bindings/node/src/error.rs | 69 ++++-- bindings/node/src/error.ts | 35 +++- bindings/node/src/index.ts | 6 + bindings/node/src/lib.rs | 36 +++- bindings/node/src/multimodal.rs | 62 +++--- bindings/node/src/types.ts | 3 + bindings/node/src/types/AiMuxError.ts | 15 +- bindings/node/src/types/ApiCallError.ts | 50 +++-- bindings/node/src/types/CallOptions.ts | 42 ++-- .../node/src/types/EmbeddingCallOptions.ts | 15 +- bindings/node/src/types/HttpExchange.ts | 19 +- bindings/node/src/types/ImageCallOptions.ts | 27 ++- bindings/node/src/types/OutcomeRecord.ts | 7 +- .../node/src/types/RerankingCallOptions.ts | 19 +- bindings/node/src/types/RetryError.ts | 8 + bindings/node/src/types/RetryErrorReason.ts | 6 + bindings/node/src/types/SearchCallOptions.ts | 25 ++- bindings/node/src/types/SpeechCallOptions.ts | 25 ++- .../node/src/types/TimeoutConfiguration.ts | 15 +- .../src/types/TranscriptionCallOptions.ts | 9 + bindings/node/src/types/VideoCallOptions.ts | 43 ++-- bindings/node/src/types/VideoPollOptions.ts | 14 ++ bindings/node/src/types/VideoResult.ts | 10 +- bindings/python/python/aimux/__init__.py | 2 + bindings/python/python/aimux/wrapper.py | 1 + bindings/python/src/error.rs | 65 +++++- bindings/python/src/lib.rs | 20 ++ bindings/python/src/multimodal.rs | 57 +++-- bindings/python/tests/test_error.py | 66 ++++-- bindings/swift/Sources/Aimux/Aimux.swift | 122 +++++++++-- bindings/swift/Sources/Aimux/Types.swift | 10 +- contract-tests/fixtures/wire-format.json | 2 +- scripts/gen_ts_types.py | 12 +- tools/aimux-web/examples/gen_demo_fixtures.rs | 3 + tools/aimux-web/examples/gen_fixture.rs | 3 + tools/aimux-web/src/api/mod.rs | 2 +- 66 files changed, 1812 insertions(+), 414 deletions(-) create mode 100644 bindings/node/src/types/RetryError.ts create mode 100644 bindings/node/src/types/RetryErrorReason.ts create mode 100644 bindings/node/src/types/VideoPollOptions.ts diff --git a/bindings/flutter/lib/aimux.dart b/bindings/flutter/lib/aimux.dart index c7ee4c4e..16210981 100644 --- a/bindings/flutter/lib/aimux.dart +++ b/bindings/flutter/lib/aimux.dart @@ -592,8 +592,10 @@ class Model implements Finalizable { /// Stream text from the model. /// - /// Returns a Stream of StreamPart maps. Terminal failures throw - /// [AimuxException] via [Stream.addError] (and do not call on_done). + /// Returns a Stream of StreamPart maps. Recoverable frame errors arrive + /// as `Error` stream-part maps and the stream continues; only + /// transport/Core failures throw [AimuxException] via [Stream.addError] + /// (and do not call on_done). /// /// NOTE: the underlying `aimux_stream_text` call is synchronous — this /// method blocks the calling isolate until the stream completes, so all diff --git a/bindings/flutter/lib/errors.dart b/bindings/flutter/lib/errors.dart index 77b98eef..b034a25d 100644 --- a/bindings/flutter/lib/errors.dart +++ b/bindings/flutter/lib/errors.dart @@ -4,7 +4,7 @@ // Transport (aimux-error.h): every fallible C call returns // `aimux_error_t *` — NULL on success (the result is in the trailing // out-param), non-NULL on failure (out-param at its sentinel: 0 / NULL). The -// unified code selects AiMuxError (1..13), RecordingError (100..105), or a +// unified code selects AiMuxError (1..14), RecordingError (100..105), or a // C ABI failure (200..206). The last range maps to // StateError('aimux ffi: …'); Dart does not expose seven additional classes. // Every field is copied before the error is released with `aimux_error_free` @@ -22,8 +22,9 @@ import 'package:ffi/ffi.dart'; // ───────────────────────────────────────────────────────────────────────────── /// Machine-readable codes. Values match C `aimux_error_code_t` / Go `Code`. -/// 13 variant codes numbered consecutively 1–13 (1 is the catch-all). A code outside -/// that range is a header/library mismatch and fails with [StateError], not an +/// 14 live variant codes: 1–14 (1 is the catch-all, [retry] at 14 reclaims +/// the slot the pre-unification `Other` vacated). A code outside that +/// set is a header/library mismatch and fails with [StateError], not an /// error type. Every HTTP-shaped failure /// arrives as [apiCall], classified /// by [AimuxException.status] (401 auth, 404 model, 429 rate limit; @@ -44,6 +45,7 @@ abstract final class AimuxErrorCode { static const int timeout = 12; static const int aborted = 13; static const int other = 1; + static const int retry = 14; static const Map _text = { ok: 'OK', @@ -60,6 +62,7 @@ abstract final class AimuxErrorCode { timeout: 'Timeout', aborted: 'Aborted', other: 'Other', + retry: 'Retry', }; /// Core `error_type()` name. @@ -72,7 +75,7 @@ abstract final class AimuxErrorCode { /// Decode the `aimux_error_t *` [e] returned by a call that can fail in /// `AiMuxError` (`[AiMuxError]` in aimux-ffi.h). NULL → returns (success). -/// Codes 1..13 become [AimuxException]; 200..206 become [StateError]. +/// Codes 1..14 become [AimuxException]; 200..206 become [StateError]. /// The returned error is always freed. void expectAimuxError(Pointer e, String context) { if (e == nullptr) return; @@ -289,11 +292,22 @@ final _StrOfPtr _errorMessage = _strGetter('aimux_error_message'); final _StrOfPtr _errorProviderCode = _strGetter('aimux_error_provider_code'); final _StrOfPtr _errorProviderMessage = _strGetter('aimux_error_provider_message'); -final _StrOfPtr _errorRequestId = _strGetter('aimux_error_request_id'); final _StrOfPtr _errorResponseBody = _strGetter('aimux_error_response_body'); +final _StrOfPtr _errorUrl = _strGetter('aimux_error_url'); +final _StrOfPtr _errorRequestBodyValues = + _strGetter('aimux_error_request_body_values'); +final _StrOfPtr _errorResponseHeaders = + _strGetter('aimux_error_response_headers'); +final _StrOfPtr _errorProviderData = _strGetter('aimux_error_provider_data'); final _StrOfPtr _errorModelId = _strGetter('aimux_error_model_id'); final _StrOfPtr _errorModelType = _strGetter('aimux_error_model_type'); final _StrOfPtr _errorProviderId = _strGetter('aimux_error_provider_id'); +final _StrOfPtr _errorRetryReason = _strGetter('aimux_error_retry_reason'); +final _IntOfPtr _errorRetryCount = _i32Getter('aimux_error_retry_count'); +final Pointer Function(Pointer, int) _errorRetryErrorAt = + _errLib.lookupFunction Function(Pointer, Int32), + Pointer Function(Pointer, int)>( + 'aimux_error_retry_error_at'); /// Read an owned getter string for [error]; frees it; null when absent. String? _errStr(_StrOfPtr getter, Pointer error) { @@ -306,6 +320,42 @@ String? _errStr(_StrOfPtr getter, Pointer error) { } } +/// Decode a JSON-string getter value; null when absent or malformed. +Object? _jsonValue(String? json) { + if (json == null) return null; + try { + return jsonDecode(json); + } on FormatException { + return null; + } +} + +/// `aimux_error_response_headers` JSON object → string→string map. +Map? _headerMap(String? json) { + final decoded = _jsonValue(json); + return decoded is Map ? Map.from(decoded) : null; +} + +/// Copy the per-attempt history of an `AIMUX_E_RETRY` error. Each child from +/// `aimux_error_retry_error_at` is a NEW owned error (deep copy, oldest +/// first): decoded with the same getters and freed independently of the +/// parent. Padded so [RetryError.errors] is never empty even for a +/// deserialized value with no recorded attempts. +List _decodeRetryErrors(Pointer error, String context) { + final count = _errorRetryCount(error); + final errors = []; + for (var i = 0; i < count; i++) { + final child = _errorRetryErrorAt(error, i); + if (child == nullptr) continue; + try { + errors.add(AimuxException._decode(child, context)); + } finally { + _errorFree(child); + } + } + return errors.isEmpty ? [OtherError('unknown retry error')] : errors; +} + /// Run [fn] with a temporary native UTF-8 copy of [s]; always frees it. T withUtf8(String s, T Function(Pointer) fn) { final ptr = toCString(s); @@ -420,8 +470,10 @@ int construct2( /// - [message]: human-readable text /// /// Per-code payload lives on the carrying subclass only: [APICallError] -/// (`providerCode`/`providerMessage`/`requestId`/`responseBody`), [NoSuchModelError] -/// (`modelId`/`modelType`), [NoSuchProviderError] (`providerId`). +/// (`providerCode`/`providerMessage`/`url`/`requestBodyValues`/ +/// `responseHeaders`/`responseBody`/`data`), [RetryError] (`reason`/`errors`), +/// [NoSuchModelError] (`modelId`/`modelType`), [NoSuchProviderError] +/// (`providerId`). class AimuxException implements Exception { /// Human-readable failure text. final String message; @@ -452,7 +504,7 @@ class AimuxException implements Exception { /// Build the typed subclass from a returned `const aimux_error_t *` [error] /// via the `aimux_error_*` getters (payload getters only under the owning /// code; getter strings freed here). The caller ([expectAimuxError]) - /// frees it. A code outside 1..13 is a + /// frees it. A code outside the published set (1..14) is a /// contract violation and throws [StateError]. factory AimuxException._decode(Pointer error, String context) { final code = _errorCode(error); @@ -468,8 +520,17 @@ class AimuxException implements Exception { retryable: retryable, providerCode: _errStr(_errorProviderCode, error), providerMessage: _errStr(_errorProviderMessage, error), - requestId: _errStr(_errorRequestId, error), - responseBody: _errStr(_errorResponseBody, error)); + url: _errStr(_errorUrl, error), + requestBodyValues: + _jsonValue(_errStr(_errorRequestBodyValues, error)), + responseHeaders: _headerMap(_errStr(_errorResponseHeaders, error)), + responseBody: _errStr(_errorResponseBody, error), + data: _jsonValue(_errStr(_errorProviderData, error))); + case AimuxErrorCode.retry: + return RetryError(message, + reason: + RetryErrorReason.fromWire(_errStr(_errorRetryReason, error)), + errors: _decodeRetryErrors(error, context)); case AimuxErrorCode.noSuchModel: return NoSuchModelError(message, retryable: retryable, @@ -522,6 +583,14 @@ class AimuxException implements Exception { return NoSuchProviderError(message, status: status, retryMs: retryMs, retryable: retryable); case AimuxErrorCode.apiCall: return APICallError(message, status: status, retryMs: retryMs, retryable: retryable); + case AimuxErrorCode.retry: + // The real decode path ([_decode]) rebuilds the attempt history via + // the retry getters; a bare code carries only its message. + return RetryError( + message, + reason: RetryErrorReason.maxRetriesExceeded, + errors: [OtherError(message)], + ); case AimuxErrorCode.timeout: return AimuxTimeoutError(message, status: status, retryMs: retryMs, retryable: retryable); case AimuxErrorCode.aborted: @@ -623,23 +692,75 @@ class APICallError extends AimuxException { /// e.g. `slow down`. final String? providerMessage; - /// Provider request id, for support tickets. - final String? requestId; + /// Sanitized request URL. + final String? url; + + /// Sanitized request body values (decoded JSON; any JSON type). + final dynamic requestBodyValues; + + /// Sanitized response headers. (Request ids live here, not on a field: + /// the core no longer captures a distinct one.) + final Map? responseHeaders; /// Raw response body. final String? responseBody; + /// Parsed provider error data (the AI SDK's `APICallError.data`), + /// decoded JSON. + final dynamic data; + APICallError(super.message, {super.status, super.retryMs, super.retryable, this.providerCode, this.providerMessage, - this.requestId, - this.responseBody}) + this.url, + this.requestBodyValues, + this.responseHeaders, + this.responseBody, + this.data}) : super(code: AimuxErrorCode.apiCall); } +/// Why a [RetryError] stopped retrying (core serde camelCase wire names). +enum RetryErrorReason { + /// Every permitted attempt failed with a retryable error. + maxRetriesExceeded('maxRetriesExceeded'), + + /// A later attempt failed with a non-retryable error. + errorNotRetryable('errorNotRetryable'); + + const RetryErrorReason(this.wireValue); + final String wireValue; + + static RetryErrorReason fromWire(String? value) => + value == errorNotRetryable.wireValue + ? errorNotRetryable + : maxRetriesExceeded; +} + +/// Retrying stopped (`AIMUX_E_RETRY`): [reason] says why, [message] carries +/// the composed summary ("Failed after N attempts…"), and [errors] keeps the +/// per-attempt history, oldest first — each entry is itself a fully decoded +/// [AimuxException] (an [APICallError] keeps its whole detail set; a nested +/// [RetryError] recurses). +class RetryError extends AimuxException { + final RetryErrorReason reason; + + /// Per-attempt errors, oldest first; never empty. + final List errors; + + RetryError(super.message, + {required this.reason, required List errors}) + : assert(errors.isNotEmpty), + errors = List.unmodifiable(errors), + super(code: AimuxErrorCode.retry); + + /// The final attempt. + AimuxException get lastError => errors.last; +} + /// Request timed out. (Prefixed to avoid shadowing dart:async /// `TimeoutError`.) class AimuxTimeoutError extends AimuxException { diff --git a/bindings/flutter/lib/types.dart b/bindings/flutter/lib/types.dart index fd3b28ee..2fe7b7b3 100644 --- a/bindings/flutter/lib/types.dart +++ b/bindings/flutter/lib/types.dart @@ -736,21 +736,25 @@ class StreamTextResultAggregated { class TimeoutConfiguration { @JsonKey(name: 'total_ms') final int? totalMs; + @JsonKey(name: 'step_ms') + final int? stepMs; @JsonKey(name: 'first_chunk_ms') final int? firstChunkMs; @JsonKey(name: 'chunk_ms') final int? chunkMs; - TimeoutConfiguration({this.totalMs, this.firstChunkMs, this.chunkMs}); + TimeoutConfiguration({this.totalMs, this.stepMs, this.firstChunkMs, this.chunkMs}); factory TimeoutConfiguration.fromJson(Map json) => TimeoutConfiguration( totalMs: json['total_ms'] as int?, + stepMs: json['step_ms'] as int?, firstChunkMs: json['first_chunk_ms'] as int?, chunkMs: json['chunk_ms'] as int?, ); Map toJson() => { if (totalMs != null) 'total_ms': totalMs, + if (stepMs != null) 'step_ms': stepMs, if (firstChunkMs != null) 'first_chunk_ms': firstChunkMs, if (chunkMs != null) 'chunk_ms': chunkMs, }; diff --git a/bindings/flutter/lib/types.g.dart b/bindings/flutter/lib/types.g.dart index 374ea12f..21013b26 100644 --- a/bindings/flutter/lib/types.g.dart +++ b/bindings/flutter/lib/types.g.dart @@ -267,6 +267,7 @@ TimeoutConfiguration _$TimeoutConfigurationFromJson( Map json) => TimeoutConfiguration( totalMs: (json['total_ms'] as num?)?.toInt(), + stepMs: (json['step_ms'] as num?)?.toInt(), firstChunkMs: (json['first_chunk_ms'] as num?)?.toInt(), chunkMs: (json['chunk_ms'] as num?)?.toInt(), ); @@ -275,6 +276,7 @@ Map _$TimeoutConfigurationToJson( TimeoutConfiguration instance) => { 'total_ms': instance.totalMs, + 'step_ms': instance.stepMs, 'first_chunk_ms': instance.firstChunkMs, 'chunk_ms': instance.chunkMs, }; diff --git a/bindings/flutter/test/errors_test.dart b/bindings/flutter/test/errors_test.dart index 0c0d4272..3a2a4ba9 100644 --- a/bindings/flutter/test/errors_test.dart +++ b/bindings/flutter/test/errors_test.dart @@ -74,6 +74,7 @@ void main() { AimuxErrorCode.timeout: AimuxTimeoutError, AimuxErrorCode.aborted: RequestAbortedError, AimuxErrorCode.other: OtherError, + AimuxErrorCode.retry: RetryError, }; for (final entry in cases.entries) { final e = AimuxException.fromCode(entry.key, 'msg'); @@ -85,9 +86,73 @@ void main() { test('unknown code is rejected with StateError', () { // A code outside the published table is an ABI mismatch, not an error // kind. 1 is AIMUX_E_OTHER now because Other inherited the old UNKNOWN - // slot, so it resolves; 14 is the gap left behind. + // slot, so it resolves; 15 is the first unassigned value. expect(() => AimuxException.fromCode(999, 'future'), throwsStateError); - expect(() => AimuxException.fromCode(14, 'unused'), throwsStateError); + expect(() => AimuxException.fromCode(15, 'unused'), throwsStateError); + }); + + test('bare retry code synthesizes a single-attempt RetryError', () { + final e = AimuxException.fromCode(AimuxErrorCode.retry, 'gave up'); + expect(e, isA()); + final retry = e as RetryError; + expect(retry.reason, RetryErrorReason.maxRetriesExceeded); + expect(retry.errors, hasLength(1)); + expect(retry.lastError, isA()); + }); + }); + + group('RetryError', () { + test('keeps the per-attempt history, oldest first', () { + final e = RetryError( + 'Failed after 2 attempts. Last error: took too long', + reason: RetryErrorReason.maxRetriesExceeded, + errors: [ + APICallError('API call error: HTTP 429: slow down', + status: 429, retryMs: 1500, retryable: true), + AimuxTimeoutError('took too long'), + ], + ); + expect(e.code, AimuxErrorCode.retry); + expect(e.codeName, 'Retry'); + expect(e.errors, hasLength(2)); + expect(e.errors.first, isA()); + expect(e.lastError, isA()); + // The history is a snapshot, not a mutable list. + expect(() => e.errors.add(OtherError('x')), throwsUnsupportedError); + }); + + test('reason decodes from the core wire names', () { + expect(RetryErrorReason.fromWire('maxRetriesExceeded'), + RetryErrorReason.maxRetriesExceeded); + expect(RetryErrorReason.fromWire('errorNotRetryable'), + RetryErrorReason.errorNotRetryable); + // Defensive default for a missing/unknown wire value. + expect( + RetryErrorReason.fromWire(null), RetryErrorReason.maxRetriesExceeded); + }); + }); + + group('APICallError', () { + test('carries the enriched request/response context', () { + final e = APICallError( + 'API call error: HTTP 429: slow down', + status: 429, + retryMs: 1500, + retryable: true, + providerCode: 'rate_limit_exceeded', + providerMessage: 'slow down', + url: 'https://api.example', + requestBodyValues: {'model': 'm'}, + responseHeaders: {'retry-after-ms': '1500'}, + responseBody: '{"error":"slow down"}', + data: {'type': 'rate_limit'}, + ); + expect(e.retryable, isTrue); + expect(e.providerCode, 'rate_limit_exceeded'); + expect(e.url, 'https://api.example'); + expect(e.requestBodyValues, {'model': 'm'}); + expect(e.responseHeaders?['retry-after-ms'], '1500'); + expect(e.data, {'type': 'rate_limit'}); }); }); @@ -110,10 +175,11 @@ void main() { expect(AimuxErrorCode.noSuchProvider, 10); expect(AimuxErrorCode.name(AimuxErrorCode.noSuchProvider), 'NoSuchProvider'); - // The AIMUX_E_UNKNOWN catch-all is gone and Other took its slot, so the - // engine codes are contiguous 1–13. + // Engine codes are contiguous 1–14. expect(AimuxErrorCode.other, 1); expect(AimuxErrorCode.aborted, 13); + expect(AimuxErrorCode.retry, 14); + expect(AimuxErrorCode.name(AimuxErrorCode.retry), 'Retry'); }); }); diff --git a/bindings/go/aimux.go b/bindings/go/aimux.go index d53e7020..7d0d671b 100644 --- a/bindings/go/aimux.go +++ b/bindings/go/aimux.go @@ -53,6 +53,7 @@ import "C" import ( "context" + "encoding/json" "errors" "fmt" "io" @@ -1124,7 +1125,7 @@ func cstr(p *C.char) string { // ── Error decoding ─────────────────────────────────────────────────────────── // // Every fallible C call returns *C.aimux_error_t: nil = success, non-nil -// = failure. The unified code space distinguishes AiMuxError (1..13), +// = failure. The unified code space distinguishes AiMuxError (1..14), // RecordingError (100..105), and failures detected by the C ABI (200..206). // The latter collapse to a plain error in Go; no public Go error type is added // for those implementation failures. Every helper frees the pointer once. @@ -1149,8 +1150,9 @@ func expectFfiError(e *C.aimux_error_t) error { return ffiError(e) } -// expectAimuxError decodes an [AiMuxError] call: nil → nil; 1..13 → *Error; -// 200..206 → plain C ABI error. Any other code is an ABI contract violation. +// expectAimuxError decodes an [AiMuxError] call: nil → nil; 1..14 → +// *Error; 200..206 → plain C ABI error. Any other code is an ABI contract +// violation. func expectAimuxError(e *C.aimux_error_t) error { if e == nil { return nil @@ -1160,7 +1162,16 @@ func expectAimuxError(e *C.aimux_error_t) error { if ffiCodeFromC(codeValue) { return ffiError(e) } + return aimuxErrorFromC(e) +} + +// aimuxErrorFromC reads a non-nil AiMuxError-coded error into *Error via the +// per-variant getters. It does not free e — the caller owns the pointer — so +// it can recurse into the owned children aimux_error_retry_error_at returns. +// A non-AiMuxError code is an ABI contract violation and panics. +func aimuxErrorFromC(e *C.aimux_error_t) *Error { str := cstr + codeValue := int(C.aimux_error_code(e)) code, ok := codeFromC(codeValue) if !ok { panic(fmt.Sprintf("aimux: unknown aimux_error_code_t: %d", codeValue)) @@ -1182,13 +1193,39 @@ func expectAimuxError(e *C.aimux_error_t) error { err.RetryMs = int64(C.aimux_error_retry_ms(e)) err.ProviderCode = str(C.aimux_error_provider_code(e)) err.ProviderMessage = str(C.aimux_error_provider_message(e)) - err.RequestID = str(C.aimux_error_request_id(e)) err.ResponseBody = str(C.aimux_error_response_body(e)) + err.URL = str(C.aimux_error_url(e)) + // The JSON-string getters carry serde-produced JSON; an absent + // payload is NULL ("" here), never invalid JSON, so a failed + // Unmarshal just leaves the field at its zero value. + if s := str(C.aimux_error_request_body_values(e)); s != "" { + _ = json.Unmarshal([]byte(s), &err.RequestBodyValues) + } + if s := str(C.aimux_error_response_headers(e)); s != "" { + _ = json.Unmarshal([]byte(s), &err.ResponseHeaders) + } + if s := str(C.aimux_error_provider_data(e)); s != "" { + _ = json.Unmarshal([]byte(s), &err.Data) + } case CodeNoSuchModel: err.ModelID = str(C.aimux_error_model_id(e)) err.ModelType = str(C.aimux_error_model_type(e)) case CodeNoSuchProvider: err.ProviderID = str(C.aimux_error_provider_id(e)) + case CodeRetry: + err.Reason = RetryErrorReason(str(C.aimux_error_retry_reason(e))) + // Each attempt is a NEW owned error (deep copy) read with the same + // getters — recursion also covers a nested Retry — and freed + // independently of the parent, right after decoding. + count := int(C.aimux_error_retry_count(e)) + for i := 0; i < count; i++ { + child := C.aimux_error_retry_error_at(e, C.int32_t(i)) + if child == nil { + continue + } + err.Errors = append(err.Errors, aimuxErrorFromC(child)) + C.aimux_error_free(child) + } } // TokenExpired carries a 401 by contract even if C reports -1; every // other status is the observed one (ApiCall without a status = no HTTP diff --git a/bindings/go/error.go b/bindings/go/error.go index bd587f7c..43a7a858 100644 --- a/bindings/go/error.go +++ b/bindings/go/error.go @@ -10,9 +10,9 @@ import ( ) // Code is the machine-readable Aimux error code. Values match -// aimux-ffi aimux_error_code_t (1..13 = core AiMuxError variants; 1 is the -// catch-all Other). -// A code outside that range is a header/library mismatch and expectAimuxError +// aimux-ffi aimux_error_code_t (1..14 = core AiMuxError variants; +// 1 is the catch-all Other, 14 = Retry). +// A code outside that set is a header/library mismatch and expectAimuxError // panics rather than inventing an "unknown" variant. Recording failures are // a different type: see RecordingError. // @@ -37,6 +37,7 @@ const ( CodeAPICall Code = 11 CodeTimeout Code = 12 CodeAborted Code = 13 + CodeRetry Code = 14 ) // String returns the core error_type name (e.g. "ApiCall", "TokenExpired"). @@ -71,17 +72,20 @@ func (c Code) String() string { return "Aborted" case CodeOther: return "Other" + case CodeRetry: + return "Retry" default: return fmt.Sprintf("Code(%d)", int(c)) } } -// codeFromC maps a C aimux_error_code_t (1..13); false for any other value. +// codeFromC maps a C aimux_error_code_t (1..14); false for any other +// value. func codeFromC(code int) (Code, bool) { - if code < int(CodeOther) || code > int(CodeAborted) { - return 0, false + if code >= int(CodeOther) && code <= int(CodeRetry) { + return Code(code), true } - return Code(code), true + return 0, false } // Error is the structured aimux failure type for Go (openai-go style: one @@ -94,6 +98,7 @@ func codeFromC(code int) (Code, bool) { // if errors.As(err, &e) { // // e.Code, e.Status, e.RetryMs, e.Retryable, e.Message // // e.Code == aimux.CodeAPICall && e.Status == 429 → rate limited +// // e.Code == aimux.CodeRetry → e.Reason, e.Errors, e.LastError() // } // // Fields mirror the aimux-ffi aimux_error_* getters / core helpers: @@ -103,9 +108,12 @@ func codeFromC(code int) (Code, bool) { // - RetryMs: rate-limit hint, or -1 (0 = retry now) // - Retryable: the core's retry verdict, carried across the ABI // - Message: human-readable text -// - ProviderCode / ProviderMessage / RequestID / ResponseBody: CodeAPICall payload +// - ProviderCode / ProviderMessage / ResponseBody / URL / +// RequestBodyValues / ResponseHeaders / Data: CodeAPICall payload // - ModelID / ModelType: CodeNoSuchModel payload // - ProviderID: CodeNoSuchProvider payload +// - Reason / Errors: CodeRetry payload — why retrying stopped, and the +// per-attempt history type Error struct { Code Code Message string @@ -121,11 +129,46 @@ type Error struct { // field is set only under the Code named here and empty otherwise. ProviderCode string // CodeAPICall: provider's own error code, e.g. "insufficient_quota" ProviderMessage string // CodeAPICall: the failure's own text without the composed prefix Message carries, e.g. "slow down" - RequestID string // CodeAPICall: provider request id ResponseBody string // CodeAPICall: raw response body - ModelID string // CodeNoSuchModel: the model id asked for - ModelType string // CodeNoSuchModel: the model type it was asked for as - ProviderID string // CodeNoSuchProvider: the provider id asked for + URL string // CodeAPICall: sanitized request URL + // RequestBodyValues is the sanitized request body of the failed call, + // decoded from JSON (any JSON type). CodeAPICall only. + RequestBodyValues any + // ResponseHeaders is the sanitized response headers of the failed call. + // Provider request-id evidence, when present, is a header here — there + // is no separate request-id field. CodeAPICall only. + ResponseHeaders map[string]string + // Data is the provider's parsed error data, decoded from JSON. + // CodeAPICall only. + Data any + ModelID string // CodeNoSuchModel: the model id asked for + ModelType string // CodeNoSuchModel: the model type it was asked for as + ProviderID string // CodeNoSuchProvider: the provider id asked for + // Reason says why retrying stopped; Errors preserves each attempt as a + // concrete *Error (oldest first; an attempt can itself be any Code, + // with the full per-code payload). CodeRetry only. + Reason RetryErrorReason + Errors []*Error +} + +// RetryErrorReason explains why operation retry stopped. Values are the +// core's serde camelCase wire names, as aimux_error_retry_reason returns them. +type RetryErrorReason string + +const ( + // RetryMaxRetriesExceeded: every permitted attempt failed with a + // retryable error. + RetryMaxRetriesExceeded RetryErrorReason = "maxRetriesExceeded" + // RetryErrorNotRetryable: an attempt failed with a non-retryable error. + RetryErrorNotRetryable RetryErrorReason = "errorNotRetryable" +) + +// LastError returns the final attempt error, or nil for a non-retry error. +func (e *Error) LastError() *Error { + if e == nil || len(e.Errors) == 0 { + return nil + } + return e.Errors[len(e.Errors)-1] } // Error implements the error interface. diff --git a/bindings/go/error_test.go b/bindings/go/error_test.go index 8bf4e528..ddda2f58 100644 --- a/bindings/go/error_test.go +++ b/bindings/go/error_test.go @@ -74,9 +74,12 @@ func TestPayloadEngineFailure(t *testing.T) { if e.ProviderID != "no-such-provider" { t.Fatalf("ProviderID: got %q", e.ProviderID) } - if e.ProviderCode != "" || e.ProviderMessage != "" || e.RequestID != "" || e.ResponseBody != "" || e.ModelID != "" || e.ModelType != "" { + if e.ProviderCode != "" || e.ProviderMessage != "" || e.ResponseBody != "" || e.URL != "" || e.ModelID != "" || e.ModelType != "" { t.Fatalf("payload fields of other codes must be empty: %+v", e) } + if e.Reason != "" || len(e.Errors) != 0 || e.LastError() != nil { + t.Fatalf("retry payload of other codes must be empty: %+v", e) + } } // Use-after-close is the binding's own guard: an ErrClosed-wrapped error, @@ -468,7 +471,8 @@ func TestEmbedRejectsRawPassThroughOpts(t *testing.T) { } func TestCodeFromCRejectsOutOfRange(t *testing.T) { - for _, bad := range []int{0, 14, 15, 999} { + // 15 is the first unassigned value. + for _, bad := range []int{0, 15, 16, 999} { if _, ok := codeFromC(bad); ok { t.Fatalf("%d is not an AiMuxError variant", bad) } @@ -480,4 +484,34 @@ func TestCodeFromCRejectsOutOfRange(t *testing.T) { if c, ok := codeFromC(11); !ok || c != CodeAPICall { t.Fatalf("11 → %v, %v", c, ok) } + if c, ok := codeFromC(14); !ok || c != CodeRetry { + t.Fatalf("14 → %v, %v; want CodeRetry", c, ok) + } +} + +// A CodeRetry error keeps the per-attempt history in order; LastError is the +// final attempt. The reason values are the core's serde camelCase wire names. +func TestRetryErrorHistory(t *testing.T) { + first := &Error{Code: CodeAPICall, Message: "API call error: HTTP 429: slow down", Status: 429, Retryable: true} + last := &Error{Code: CodeAPICall, Message: "API call error: HTTP 401: bad key", Status: 401} + e := &Error{ + Code: CodeRetry, + Message: "Failed after 2 attempts with non-retryable error: 'bad key'", + Status: -1, + RetryMs: -1, + Reason: RetryErrorNotRetryable, + Errors: []*Error{first, last}, + } + if e.Code.String() != "Retry" { + t.Fatalf("Code.String: %s", e.Code.String()) + } + if e.LastError() != last { + t.Fatalf("LastError: got %v", e.LastError()) + } + if e.Errors[0] != first || !e.Errors[0].Retryable { + t.Fatalf("attempt order lost: %+v", e.Errors) + } + if e.Reason != "errorNotRetryable" || RetryMaxRetriesExceeded != "maxRetriesExceeded" { + t.Fatalf("reason wire names changed: %q", e.Reason) + } } diff --git a/bindings/go/multimodal_test.go b/bindings/go/multimodal_test.go index c4a52f81..b695ab57 100644 --- a/bindings/go/multimodal_test.go +++ b/bindings/go/multimodal_test.go @@ -244,6 +244,33 @@ func TestE2E_VideoResultParsing(t *testing.T) { } } +func TestVideoPollOptionsWireFormat(t *testing.T) { + intervalMS := uint64(1_000) + timeoutMS := uint64(120_000) + opts := VideoCallOptions{ + Poll: &VideoPollOptions{ + IntervalMS: &intervalMS, + TimeoutMS: &timeoutMS, + }, + } + + encoded, err := json.Marshal(opts) + if err != nil { + t.Fatalf("marshal video options: %v", err) + } + var wire map[string]any + if err := json.Unmarshal(encoded, &wire); err != nil { + t.Fatalf("decode video options: %v", err) + } + poll, ok := wire["poll"].(map[string]any) + if !ok { + t.Fatalf("poll missing from wire format: %s", encoded) + } + if poll["interval_ms"] != float64(intervalMS) || poll["timeout_ms"] != float64(timeoutMS) { + t.Fatalf("unexpected poll wire format: %s", encoded) + } +} + // ── E2E: Search result parsing ─────────────────────────────────────────────── func TestE2E_SearchResultParsing(t *testing.T) { diff --git a/bindings/go/multimodal_types.go b/bindings/go/multimodal_types.go index aec4cf7e..3eb2eabf 100644 --- a/bindings/go/multimodal_types.go +++ b/bindings/go/multimodal_types.go @@ -38,9 +38,11 @@ type EmbeddingResult struct { // EmbeddingCallOptions is the options for an embedding call. type EmbeddingCallOptions struct { - Values []string `json:"values,omitempty"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Values []string `json:"values,omitempty"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Speech (TTS) ────────────────────────────────────────────────────────────── @@ -78,14 +80,16 @@ type SpeechResult struct { // SpeechCallOptions is the options for speech generation. type SpeechCallOptions struct { - Text string `json:"text"` - Voice *string `json:"voice,omitempty"` - OutputFormat *string `json:"output_format,omitempty"` - Instructions *string `json:"instructions,omitempty"` - Speed *float64 `json:"speed,omitempty"` - Language *string `json:"language,omitempty"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Text string `json:"text"` + Voice *string `json:"voice,omitempty"` + OutputFormat *string `json:"output_format,omitempty"` + Instructions *string `json:"instructions,omitempty"` + Speed *float64 `json:"speed,omitempty"` + Language *string `json:"language,omitempty"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Image ───────────────────────────────────────────────────────────────────── @@ -121,15 +125,17 @@ type ImageResult struct { // ImageCallOptions is the options for image generation. type ImageCallOptions struct { - Prompt *string `json:"prompt,omitempty"` - N *int `json:"n,omitempty"` - Size *string `json:"size,omitempty"` - AspectRatio *string `json:"aspect_ratio,omitempty"` - Seed *uint64 `json:"seed,omitempty"` - Files []json.RawMessage `json:"files,omitempty"` - Mask json.RawMessage `json:"mask,omitempty"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Prompt *string `json:"prompt,omitempty"` + N *int `json:"n,omitempty"` + Size *string `json:"size,omitempty"` + AspectRatio *string `json:"aspect_ratio,omitempty"` + Seed *uint64 `json:"seed,omitempty"` + Files []json.RawMessage `json:"files,omitempty"` + Mask json.RawMessage `json:"mask,omitempty"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Transcription (STT) ─────────────────────────────────────────────────────── @@ -172,10 +178,12 @@ type TranscriptionResult struct { // TranscriptionCallOptions is the options for transcription. type TranscriptionCallOptions struct { - Audio json.RawMessage `json:"audio"` - MediaType string `json:"media_type"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Audio json.RawMessage `json:"audio"` + MediaType string `json:"media_type"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Reranking ──────────────────────────────────────────────────────────────── @@ -205,11 +213,13 @@ type RerankingResult struct { // RerankingCallOptions is the options for reranking. type RerankingCallOptions struct { - Documents json.RawMessage `json:"documents"` - Query string `json:"query"` - TopN *int `json:"top_n,omitempty"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Documents json.RawMessage `json:"documents"` + Query string `json:"query"` + TopN *int `json:"top_n,omitempty"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Video ─────────────────────────────────────────────────────────────────── @@ -255,15 +265,24 @@ type VideoResult struct { Response VideoResponse `json:"response"` } +// VideoPollOptions overrides Core's pacing for an asynchronous video job. +type VideoPollOptions struct { + IntervalMS *uint64 `json:"interval_ms,omitempty"` + TimeoutMS *uint64 `json:"timeout_ms,omitempty"` +} + // VideoCallOptions is the options for video generation. type VideoCallOptions struct { - Prompt *string `json:"prompt,omitempty"` - N *int `json:"n,omitempty"` - AspectRatio *string `json:"aspect_ratio,omitempty"` - Resolution *string `json:"resolution,omitempty"` - Seed *uint64 `json:"seed,omitempty"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Prompt *string `json:"prompt,omitempty"` + N *int `json:"n,omitempty"` + AspectRatio *string `json:"aspect_ratio,omitempty"` + Resolution *string `json:"resolution,omitempty"` + Seed *uint64 `json:"seed,omitempty"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Poll *VideoPollOptions `json:"poll,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Search ────────────────────────────────────────────────────────────────── @@ -294,14 +313,16 @@ type SearchResult struct { // SearchCallOptions is the options for a search call. type SearchCallOptions struct { - Query string `json:"query"` - MaxResults *int `json:"max_results,omitempty"` - IncludeRawContent *bool `json:"include_raw_content,omitempty"` - TimeRange *string `json:"time_range,omitempty"` - IncludeDomains []string `json:"include_domains,omitempty"` - ExcludeDomains []string `json:"exclude_domains,omitempty"` - ProviderOptions jsonObj `json:"provider_options"` - Headers map[string]string `json:"headers,omitempty"` + Query string `json:"query"` + MaxResults *int `json:"max_results,omitempty"` + IncludeRawContent *bool `json:"include_raw_content,omitempty"` + TimeRange *string `json:"time_range,omitempty"` + IncludeDomains []string `json:"include_domains,omitempty"` + ExcludeDomains []string `json:"exclude_domains,omitempty"` + MaxRetries *uint32 `json:"max_retries,omitempty"` + Timeout *TimeoutConfiguration `json:"timeout,omitempty"` + ProviderOptions jsonObj `json:"provider_options"` + Headers map[string]string `json:"headers,omitempty"` } // ── Files ─────────────────────────────────────────────────────────────────── diff --git a/bindings/go/types.go b/bindings/go/types.go index e2b6e7be..3c4bd62a 100644 --- a/bindings/go/types.go +++ b/bindings/go/types.go @@ -265,6 +265,7 @@ func MarshalMessages(msgs []ModelMessage) (string, error) { // TimeoutConfiguration sets per-call timeouts (RFC-0016 H3). type TimeoutConfiguration struct { TotalMs *uint64 `json:"total_ms,omitempty"` + StepMs *uint64 `json:"step_ms,omitempty"` FirstChunkMs *uint64 `json:"first_chunk_ms,omitempty"` ChunkMs *uint64 `json:"chunk_ms,omitempty"` } diff --git a/bindings/go/wire_format_test.go b/bindings/go/wire_format_test.go index f8d125aa..c54932e3 100644 --- a/bindings/go/wire_format_test.go +++ b/bindings/go/wire_format_test.go @@ -218,10 +218,12 @@ func TestWireFormatConsistency(t *testing.T) { if tc3.ChunkMs == nil || *tc3.ChunkMs != 500 { t.Error("expected ChunkMs=500") } - reencoded, _ := json.Marshal(tc3) - if string(reencoded) != wireJSON { - t.Errorf("round-trip mismatch: got %s, want %s", reencoded, wireJSON) + if tc3.StepMs != nil { + t.Error("expected nil StepMs") } + // No byte round-trip here: the fixture spells Rust's explicit + // nulls, while the Go struct is omitempty by convention (same + // reasoning as the GenerateTextOptions branch above). case "GenerateContent": // Go keeps result content parts as raw JSON by design diff --git a/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java index 93eaf815..db259467 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxException.java @@ -1,11 +1,17 @@ package ai.arcships.aimux; +import com.fasterxml.jackson.databind.JsonNode; import com.sun.jna.Pointer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; /** * AiMuxError hierarchy (OpenAI Java / Vercel AI SDK style). * - *

Raised when a fallible C ABI call returns an AiMuxError code (1–13). + *

Raised when a fallible C ABI call returns an AiMuxError code (1–14). * Recording failures use the * independent {@link RecordingException} type; C ABI failures (bad raw * wire JSON, use-after-close, re-entrant call) surface as plain @@ -23,7 +29,8 @@ * } * } * - *

Every instance carries {@link #getCode()} (C {@code aimux_error_code_t} 1–13), + *

Every instance carries {@link #getCode()} (C {@code aimux_error_code_t} + * 1–14, where {@link #AIMUX_E_RETRY} = 14), * {@link #getStatusCode()} (HTTP or {@code -1}), {@link #getRetryMs()} (hint * or {@code -1}; {@code 0} = retry now) and {@link #isRetryable()}. Message * text comes from the C layer. @@ -36,8 +43,9 @@ public class AimuxException extends RuntimeException { private static final long serialVersionUID = 1L; // ── aimux_error_code_t (aimux-error.h) ────────────────────────────────── - // 13 variant codes (1–13; 1 is the catch-all OTHER); every HTTP-shaped failure - // arrives as AIMUX_E_API_CALL. A code outside that range is a header/library + // 14 variant codes (1–14; 1 is the catch-all OTHER, 14 = RETRY reclaiming + // the slot the pre-unification Other vacated); every HTTP-shaped failure + // arrives as AIMUX_E_API_CALL. A code outside that set is a header/library // mismatch and fails with IllegalStateException, never an AimuxException. // Recording failures are a different type: see RecordingException. @@ -55,13 +63,14 @@ public class AimuxException extends RuntimeException { public static final int AIMUX_E_TIMEOUT = 12; public static final int AIMUX_E_ABORTED = 13; public static final int AIMUX_E_OTHER = 1; + public static final int AIMUX_E_RETRY = 14; private final int code; private final int status; private final long retryMs; // Set once by the fromC construction path; false for local / synthesized - // failures. Not a constructor param so the 13 subclass constructors keep + // failures. Not a constructor param so the subclass constructors keep // their public signatures. private boolean retryable; @@ -93,7 +102,7 @@ public AimuxException(String message, int code, int status, long retryMs, Throwa // ── Accessors ─────────────────────────────────────────────────────────── - /** C {@code aimux_error_code_t} value (1–13). */ + /** C {@code aimux_error_code_t} value (1–14). */ public int getCode() { return code; } @@ -128,8 +137,9 @@ public boolean isRetryable() { * prefixing {@code prefix} to the message. Reads the code, message, * retryable, status and the payload getters for that code only, freeing * every returned string. Does not own the pointer: the caller - * ({@link AimuxResult#expectAimuxError}) frees the returned error afterwards. - * A code outside 1–13 is a header/library mismatch → + * ({@link AimuxResult#expectAimuxError}) frees the returned error afterwards + * (retry attempt errors are new owned copies and are freed here). + * A code outside 1–14 is a header/library mismatch → * {@link IllegalStateException}. */ static AimuxException fromC(Pointer error, String prefix) { @@ -149,8 +159,16 @@ static AimuxException fromC(Pointer error, String prefix) { ex = new APICallError(msg, ffi.aimux_error_status(error), ffi.aimux_error_retry_ms(error), AimuxResult.takeString(ffi.aimux_error_provider_code(error)), AimuxResult.takeString(ffi.aimux_error_provider_message(error)), - AimuxResult.takeString(ffi.aimux_error_request_id(error)), - AimuxResult.takeString(ffi.aimux_error_response_body(error))); + AimuxResult.takeString(ffi.aimux_error_response_body(error)), + AimuxResult.takeString(ffi.aimux_error_url(error)), + parseJson(AimuxResult.takeString(ffi.aimux_error_request_body_values(error))), + headerMap(AimuxResult.takeString(ffi.aimux_error_response_headers(error))), + parseJson(AimuxResult.takeString(ffi.aimux_error_provider_data(error)))); + break; + case AIMUX_E_RETRY: + ex = new RetryError(msg, + RetryErrorReason.fromWire(AimuxResult.takeString(ffi.aimux_error_retry_reason(error))), + retryHistory(ffi, error, msg)); break; case AIMUX_E_NO_SUCH_MODEL: ex = new NoSuchModelError(msg, -1, -1L, @@ -168,6 +186,56 @@ static AimuxException fromC(Pointer error, String prefix) { return ex; } + /** + * Decode the per-attempt history of an {@link #AIMUX_E_RETRY} error. + * Each attempt is a new owned {@code aimux_error_t *} (index 0 = oldest) + * that can itself be any AiMuxError — including a nested Retry — and is + * freed here, independently of the parent. + */ + private static List retryHistory(AimuxFFI ffi, Pointer error, String fallbackMessage) { + int count = ffi.aimux_error_retry_count(error); + List errors = new ArrayList<>(); + for (int i = 0; i < count; i++) { + Pointer attempt = ffi.aimux_error_retry_error_at(error, i); + if (attempt == null) { + continue; + } + try { + errors.add(fromC(attempt, "")); + } finally { + ffi.aimux_error_free(attempt); + } + } + if (errors.isEmpty()) { + // A deserialized Retry may carry no attempts; keep RetryError total. + errors.add(new OtherError(fallbackMessage, -1, -1L)); + } + return errors; + } + + /** Parse a getter-returned JSON string; {@code null} (absent) stays {@code null}. */ + private static JsonNode parseJson(String json) { + if (json == null) { + return null; + } + try { + return Types.AimuxJson.MAPPER.readTree(json); + } catch (Exception e) { + return null; + } + } + + /** Response headers arrive as one JSON object string of string→string pairs. */ + private static Map headerMap(String json) { + JsonNode node = parseJson(json); + if (node == null || !node.isObject()) { + return null; + } + Map headers = new LinkedHashMap<>(); + node.fields().forEachRemaining(f -> headers.put(f.getKey(), f.getValue().asText())); + return headers; + } + /** * Local (non-FFI) failure: default status/retryMs = -1. */ @@ -205,6 +273,9 @@ private static AimuxException createByCode(int code, String message, int status, return new NoSuchProviderError(message, status, retryMs); case AIMUX_E_API_CALL: return new APICallError(message, status, retryMs); + case AIMUX_E_RETRY: + return new RetryError(message, RetryErrorReason.MAX_RETRIES_EXCEEDED, + Collections.singletonList(new OtherError(message, -1, -1L))); case AIMUX_E_TIMEOUT: return new TimeoutError(message, status, retryMs); case AIMUX_E_ABORTED: @@ -247,6 +318,8 @@ public static String codeName(int code) { return "Aborted"; case AIMUX_E_OTHER: return "Other"; + case AIMUX_E_RETRY: + return "Retry"; default: return "Code(" + code + ")"; } @@ -351,20 +424,28 @@ public String getProviderId() { public static class APICallError extends AimuxException { private final String providerCode; private final String providerMessage; - private final String requestId; private final String responseBody; + private final String url; + private final JsonNode requestBodyValues; + private final Map responseHeaders; + private final JsonNode data; public APICallError(String message, int status, long retryMs) { - this(message, status, retryMs, null, null, null, null); + this(message, status, retryMs, null, null, null, null, null, null, null); } public APICallError(String message, int status, long retryMs, - String providerCode, String providerMessage, String requestId, String responseBody) { + String providerCode, String providerMessage, String responseBody, + String url, JsonNode requestBodyValues, + Map responseHeaders, JsonNode data) { super(message, AIMUX_E_API_CALL, status, retryMs); this.providerCode = providerCode; this.providerMessage = providerMessage; - this.requestId = requestId; this.responseBody = responseBody; + this.url = url; + this.requestBodyValues = requestBodyValues; + this.responseHeaders = responseHeaders; + this.data = data; } /** The provider's own error code (e.g. {@code "insufficient_quota"}), or {@code null}. */ @@ -377,15 +458,68 @@ public String getProviderMessage() { return providerMessage; } - /** Provider request id (for support tickets), or {@code null}. */ - public String getRequestId() { - return requestId; - } - /** Raw response body, or {@code null}. */ public String getResponseBody() { return responseBody; } + + /** Sanitized request URL, or {@code null}. */ + public String getUrl() { + return url; + } + + /** Sanitized request body values (any JSON type), or {@code null}. */ + public JsonNode getRequestBodyValues() { + return requestBodyValues; + } + + /** Sanitized response headers (also carry provider request ids), or {@code null}. */ + public Map getResponseHeaders() { + return responseHeaders; + } + + /** Parsed provider error data (AI SDK {@code APICallError.data}), or {@code null}. */ + public JsonNode getData() { + return data; + } + } + + /** Why the retry loop gave up; wire names are the core's serde camelCase. */ + public enum RetryErrorReason { + /** Every permitted attempt failed with a retryable error. */ + MAX_RETRIES_EXCEEDED, + /** A later attempt failed with a non-retryable error. */ + ERROR_NOT_RETRYABLE; + + static RetryErrorReason fromWire(String value) { + return "errorNotRetryable".equals(value) + ? ERROR_NOT_RETRYABLE : MAX_RETRIES_EXCEEDED; + } + } + + /** + * The retry loop gave up (AI SDK {@code RetryError} analogue): + * {@link #getReason()} says why, {@link #getErrors()} is the per-attempt + * history (oldest first), {@link #getLastError()} the final attempt. + */ + public static class RetryError extends AimuxException { + private final RetryErrorReason reason; + private final List errors; + private final AimuxException lastError; + + public RetryError(String message, RetryErrorReason reason, List errors) { + super(message, AIMUX_E_RETRY, -1, -1L); + if (errors == null || errors.isEmpty()) { + throw new IllegalArgumentException("RetryError requires at least one error"); + } + this.reason = reason; + this.errors = Collections.unmodifiableList(new ArrayList<>(errors)); + this.lastError = this.errors.get(this.errors.size() - 1); + } + + public RetryErrorReason getReason() { return reason; } + public List getErrors() { return errors; } + public AimuxException getLastError() { return lastError; } } public static class TimeoutError extends AimuxException { @@ -394,7 +528,7 @@ public TimeoutError(String message, int status, long retryMs) { } } - /** Request aborted (not a DOM {@code AbortError}). */ + /** Request aborted (not a DOM {@code AbortError}); the message is the abort payload. */ public static class RequestAbortedError extends AimuxException { public RequestAbortedError(String message, int status, long retryMs) { super(message, AIMUX_E_ABORTED, status, retryMs); diff --git a/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java index ce5f710a..21c1a96a 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxFFI.java @@ -183,12 +183,30 @@ Pointer aimux_stream_text_as_openai(long handle, String promptJson, String optsJ /** AIMUX_E_API_CALL: owned string or NULL. */ Pointer aimux_error_provider_message(Pointer err); - /** AIMUX_E_API_CALL: owned string or NULL. */ - Pointer aimux_error_request_id(Pointer err); - /** AIMUX_E_API_CALL: owned string or NULL. */ Pointer aimux_error_response_body(Pointer err); + /** AIMUX_E_API_CALL: sanitized request URL; owned string or NULL. */ + Pointer aimux_error_url(Pointer err); + + /** AIMUX_E_API_CALL: sanitized request body values as a JSON string; owned or NULL. */ + Pointer aimux_error_request_body_values(Pointer err); + + /** AIMUX_E_API_CALL: response headers as one JSON object string of string→string pairs; owned or NULL. */ + Pointer aimux_error_response_headers(Pointer err); + + /** AIMUX_E_API_CALL: parsed provider error data as a JSON string; owned or NULL. */ + Pointer aimux_error_provider_data(Pointer err); + + /** AIMUX_E_RETRY: "maxRetriesExceeded" / "errorNotRetryable"; owned string or NULL. */ + Pointer aimux_error_retry_reason(Pointer err); + + /** AIMUX_E_RETRY: number of recorded attempt errors; 0 under any other code. */ + int aimux_error_retry_count(Pointer err); + + /** AIMUX_E_RETRY: attempt at index (0 = oldest) as a NEW OWNED error — release with {@link #aimux_error_free}; NULL when out of range or not Retry. */ + Pointer aimux_error_retry_error_at(Pointer err, int index); + /** AIMUX_E_NO_SUCH_MODEL: owned string or NULL. */ Pointer aimux_error_model_id(Pointer err); diff --git a/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java b/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java index 7a383962..f4eca39e 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/AimuxResult.java @@ -9,7 +9,7 @@ * *

Every fallible C call returns an {@code aimux_error_t *} ({@code null} * = success, result in the out-parameter). Its code identifies an AiMuxError - * (1–13), RecordingError (100–105), or a failure detected by the C ABI + * (1–14), RecordingError (100–105), or a failure detected by the C ABI * (200–206). The last range collapses to {@link IllegalStateException} * ({@code "aimux ffi: "} + message); Java does not expose seven additional * exception types. Each helper frees the pointer exactly once. User-triggerable @@ -64,7 +64,7 @@ private static String prefix(String context) { } /** - * Decode an error from a call that may return {@code AiMuxError}: 1–13 → + * Decode an error from a call that may return {@code AiMuxError}: 1–14 → * {@link AimuxException}; 200–206 → {@link IllegalStateException}. * Frees {@code e}. */ @@ -79,7 +79,8 @@ static RuntimeException expectAimuxError(Pointer e, String context) { if (isFfiCode(code)) { return ffiError(e, prefix); } - if (code < AimuxException.AIMUX_E_OTHER || code > AimuxException.AIMUX_E_ABORTED) { + if ((code < AimuxException.AIMUX_E_OTHER || code > AimuxException.AIMUX_E_ABORTED) + && code != AimuxException.AIMUX_E_RETRY) { return codeMismatch(code, prefix, "AiMuxError"); } return AimuxException.fromC(e, prefix); diff --git a/bindings/java/src/main/java/ai/arcships/aimux/Model.java b/bindings/java/src/main/java/ai/arcships/aimux/Model.java index 84b6fb5d..da205ada 100644 --- a/bindings/java/src/main/java/ai/arcships/aimux/Model.java +++ b/bindings/java/src/main/java/ai/arcships/aimux/Model.java @@ -629,8 +629,10 @@ public void invoke(Pointer streamCtx) { *

The FFI call starts on the first terminal operation of the returned * stream (mirror of Kotlin's {@code streamTextSequence}). Iteration pulls * parts from a {@link LinkedBlockingQueue} fed by the stream callbacks; - * the stream ends at the sentinel. Terminal stream failures throw - * {@link AimuxException} from the blocking FFI call. + * the stream ends at the sentinel. Recoverable frame errors arrive as + * {@code StreamPart.Error} data parts and the stream continues; only + * transport/Core failures throw {@link AimuxException} from the blocking + * FFI call. * * 

{@code
      * model.streamTextStream("\"Write a haiku\"").forEach(System.out::println);
diff --git a/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java b/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java
index 3ea1e0d3..d310ca38 100644
--- a/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java
+++ b/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java
@@ -254,6 +254,8 @@ public static class EmbeddingCallOptions {
         @JsonProperty("values") private List values = new ArrayList<>();
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         EmbeddingCallOptions() {}
@@ -271,6 +273,8 @@ public static EmbeddingCallOptions of(List values) {
         public List getValues() { return values; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -278,13 +282,20 @@ public static class Builder {
             private List values = new ArrayList<>();
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder values(List v) { this.values = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public EmbeddingCallOptions build() {
-                return new EmbeddingCallOptions(values, providerOptions, headers);
+                EmbeddingCallOptions result = new EmbeddingCallOptions(values, providerOptions, headers);
+                result.maxRetries = maxRetries;
+                result.timeout = timeout;
+                return result;
             }
         }
 
@@ -295,11 +306,13 @@ public boolean equals(Object o) {
             EmbeddingCallOptions that = (EmbeddingCallOptions) o;
             return Objects.equals(values, that.values)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
-        public int hashCode() { return Objects.hash(values, providerOptions, headers); }
+        public int hashCode() { return Objects.hash(values, providerOptions, headers, maxRetries, timeout); }
 
         @Override
         public String toString() { return "EmbeddingCallOptions(" + values + ")"; }
@@ -601,6 +614,8 @@ public static class SpeechCallOptions {
         @JsonProperty("language") private String language;
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         SpeechCallOptions() {}
@@ -630,6 +645,8 @@ public static SpeechCallOptions of(String text) {
         public String getLanguage() { return language; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -642,6 +659,8 @@ public static class Builder {
             private String language;
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder text(String v) { this.text = v; return this; }
             public Builder voice(String v) { this.voice = v; return this; }
@@ -651,10 +670,15 @@ public static class Builder {
             public Builder language(String v) { this.language = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public SpeechCallOptions build() {
-                return new SpeechCallOptions(text, voice, outputFormat, instructions,
+                SpeechCallOptions result = new SpeechCallOptions(text, voice, outputFormat, instructions,
                     speed, language, providerOptions, headers);
+                result.maxRetries = maxRetries;
+                result.timeout = timeout;
+                return result;
             }
         }
 
@@ -670,12 +694,14 @@ public boolean equals(Object o) {
                 && Objects.equals(speed, that.speed)
                 && Objects.equals(language, that.language)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(text, voice, outputFormat, instructions, speed, language, providerOptions, headers);
+            return Objects.hash(text, voice, outputFormat, instructions, speed, language, providerOptions, headers, maxRetries, timeout);
         }
 
         @Override
@@ -1003,6 +1029,8 @@ public static class ImageCallOptions {
         @JsonProperty("mask") private JsonNode mask;
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         ImageCallOptions() {}
@@ -1035,6 +1063,8 @@ public static ImageCallOptions of(String prompt) {
         public JsonNode getMask() { return mask; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -1048,6 +1078,8 @@ public static class Builder {
             private JsonNode mask;
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder prompt(String v) { this.prompt = v; return this; }
             public Builder n(Integer v) { this.n = v; return this; }
@@ -1058,10 +1090,15 @@ public static class Builder {
             public Builder mask(JsonNode v) { this.mask = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public ImageCallOptions build() {
-                return new ImageCallOptions(prompt, n, size, aspectRatio, seed, files, mask,
+                ImageCallOptions result = new ImageCallOptions(prompt, n, size, aspectRatio, seed, files, mask,
                     providerOptions, headers);
+                result.maxRetries = maxRetries;
+                result.timeout = timeout;
+                return result;
             }
         }
 
@@ -1078,12 +1115,14 @@ public boolean equals(Object o) {
                 && Objects.equals(files, that.files)
                 && Objects.equals(mask, that.mask)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(prompt, n, size, aspectRatio, seed, files, mask, providerOptions, headers);
+            return Objects.hash(prompt, n, size, aspectRatio, seed, files, mask, providerOptions, headers, maxRetries, timeout);
         }
 
         @Override
@@ -1355,6 +1394,8 @@ public static class TranscriptionCallOptions {
         @JsonProperty("media_type") private String mediaType = "";
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         TranscriptionCallOptions() {}
@@ -1375,6 +1416,8 @@ public static TranscriptionCallOptions of(JsonNode audio, String mediaType) {
         public String getMediaType() { return mediaType; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -1383,14 +1426,21 @@ public static class Builder {
             private String mediaType = "";
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder audio(JsonNode v) { this.audio = v; return this; }
             public Builder mediaType(String v) { this.mediaType = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public TranscriptionCallOptions build() {
-                return new TranscriptionCallOptions(audio, mediaType, providerOptions, headers);
+                TranscriptionCallOptions result = new TranscriptionCallOptions(audio, mediaType, providerOptions, headers);
+                result.maxRetries = maxRetries;
+                result.timeout = timeout;
+                return result;
             }
         }
 
@@ -1402,11 +1452,13 @@ public boolean equals(Object o) {
             return Objects.equals(audio, that.audio)
                 && Objects.equals(mediaType, that.mediaType)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
-        public int hashCode() { return Objects.hash(audio, mediaType, providerOptions, headers); }
+        public int hashCode() { return Objects.hash(audio, mediaType, providerOptions, headers, maxRetries, timeout); }
 
         @Override
         public String toString() { return "TranscriptionCallOptions(" + mediaType + ")"; }
@@ -1602,6 +1654,8 @@ public static class RerankingCallOptions {
         @JsonProperty("top_n") private Integer topN;
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         RerankingCallOptions() {}
@@ -1624,6 +1678,8 @@ public static RerankingCallOptions of(JsonNode documents, String query) {
         public Integer getTopN() { return topN; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -1633,15 +1689,22 @@ public static class Builder {
             private Integer topN;
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder documents(JsonNode v) { this.documents = v; return this; }
             public Builder query(String v) { this.query = v; return this; }
             public Builder topN(Integer v) { this.topN = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public RerankingCallOptions build() {
-                return new RerankingCallOptions(documents, query, topN, providerOptions, headers);
+                RerankingCallOptions result = new RerankingCallOptions(documents, query, topN, providerOptions, headers);
+                result.maxRetries = maxRetries;
+                result.timeout = timeout;
+                return result;
             }
         }
 
@@ -1654,11 +1717,13 @@ public boolean equals(Object o) {
                 && Objects.equals(query, that.query)
                 && Objects.equals(topN, that.topN)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
-        public int hashCode() { return Objects.hash(documents, query, topN, providerOptions, headers); }
+        public int hashCode() { return Objects.hash(documents, query, topN, providerOptions, headers, maxRetries, timeout); }
 
         @Override
         public String toString() { return "RerankingCallOptions(" + query + ")"; }
@@ -2421,6 +2486,54 @@ public boolean equals(Object o) {
         public String toString() { return "VideoFrameImage(" + image + ", " + frameType + ")"; }
     }
 
+    /** Per-call pacing overrides for the Core-owned video status poll loop. */
+    public static class VideoPollOptions {
+        @JsonProperty("interval_ms") private Long intervalMs;
+        @JsonProperty("timeout_ms") private Long timeoutMs;
+
+        @JsonCreator
+        VideoPollOptions() {}
+
+        private VideoPollOptions(Long intervalMs, Long timeoutMs) {
+            this.intervalMs = intervalMs;
+            this.timeoutMs = timeoutMs;
+        }
+
+        public Long getIntervalMs() { return intervalMs; }
+        public Long getTimeoutMs() { return timeoutMs; }
+
+        public static Builder builder() { return new Builder(); }
+
+        public static class Builder {
+            private Long intervalMs;
+            private Long timeoutMs;
+
+            public Builder intervalMs(Long v) { this.intervalMs = v; return this; }
+            public Builder timeoutMs(Long v) { this.timeoutMs = v; return this; }
+
+            public VideoPollOptions build() {
+                return new VideoPollOptions(intervalMs, timeoutMs);
+            }
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof VideoPollOptions)) return false;
+            VideoPollOptions that = (VideoPollOptions) o;
+            return Objects.equals(intervalMs, that.intervalMs)
+                && Objects.equals(timeoutMs, that.timeoutMs);
+        }
+
+        @Override
+        public int hashCode() { return Objects.hash(intervalMs, timeoutMs); }
+
+        @Override
+        public String toString() {
+            return "VideoPollOptions(" + intervalMs + ", " + timeoutMs + ")";
+        }
+    }
+
     /** Options for video generation. */
     public static class VideoCallOptions {
         @JsonProperty("prompt") private String prompt;
@@ -2436,6 +2549,9 @@ public static class VideoCallOptions {
         @JsonProperty("generate_audio") private Boolean generateAudio;
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("poll") private VideoPollOptions poll;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         VideoCallOptions() {}
@@ -2444,7 +2560,9 @@ private VideoCallOptions(String prompt, Integer n, String aspectRatio, String re
                                  Long duration, Double fps, Long seed,
                                  VideoFile image, List frameImages,
                                  List inputReferences, Boolean generateAudio,
-                                 JsonNode providerOptions, Map headers) {
+                                 JsonNode providerOptions, Map headers,
+                                 Integer maxRetries, VideoPollOptions poll,
+                                 Types.TimeoutConfiguration timeout) {
             this.prompt = prompt;
             this.n = n;
             this.aspectRatio = aspectRatio;
@@ -2458,11 +2576,14 @@ private VideoCallOptions(String prompt, Integer n, String aspectRatio, String re
             this.generateAudio = generateAudio;
             this.providerOptions = providerOptions;
             this.headers = headers;
+            this.maxRetries = maxRetries;
+            this.poll = poll;
+            this.timeout = timeout;
         }
 
         public static VideoCallOptions of(String prompt) {
             return new VideoCallOptions(prompt, null, null, null, null, null, null,
-                                        null, null, null, null, null, null);
+                                        null, null, null, null, null, null, null, null, null);
         }
 
         public String getPrompt() { return prompt; }
@@ -2478,6 +2599,9 @@ public static VideoCallOptions of(String prompt) {
         public Boolean getGenerateAudio() { return generateAudio; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public VideoPollOptions getPoll() { return poll; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -2495,6 +2619,9 @@ public static class Builder {
             private Boolean generateAudio;
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private VideoPollOptions poll;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder prompt(String v) { this.prompt = v; return this; }
             public Builder n(Integer v) { this.n = v; return this; }
@@ -2509,11 +2636,14 @@ public static class Builder {
             public Builder generateAudio(Boolean v) { this.generateAudio = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder poll(VideoPollOptions v) { this.poll = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public VideoCallOptions build() {
                 return new VideoCallOptions(prompt, n, aspectRatio, resolution, duration, fps, seed,
                                             image, frameImages, inputReferences, generateAudio,
-                                            providerOptions, headers);
+                                            providerOptions, headers, maxRetries, poll, timeout);
             }
         }
 
@@ -2534,14 +2664,17 @@ public boolean equals(Object o) {
                 && Objects.equals(inputReferences, that.inputReferences)
                 && Objects.equals(generateAudio, that.generateAudio)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(poll, that.poll)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
         public int hashCode() {
             return Objects.hash(prompt, n, aspectRatio, resolution, duration, fps, seed,
                                 image, frameImages, inputReferences, generateAudio,
-                                providerOptions, headers);
+                                providerOptions, headers, maxRetries, poll, timeout);
         }
 
         @Override
@@ -2748,6 +2881,8 @@ public static class SearchCallOptions {
         @JsonProperty("exclude_domains") private List excludeDomains = new ArrayList<>();
         @JsonProperty("provider_options") private JsonNode providerOptions;
         @JsonProperty("headers") private Map headers;
+        @JsonProperty("max_retries") private Integer maxRetries;
+        @JsonProperty("timeout") private Types.TimeoutConfiguration timeout;
 
         @JsonCreator
         SearchCallOptions() {}
@@ -2779,6 +2914,8 @@ public static SearchCallOptions of(String query) {
         public List getExcludeDomains() { return excludeDomains; }
         public JsonNode getProviderOptions() { return providerOptions; }
         public Map getHeaders() { return headers; }
+        public Integer getMaxRetries() { return maxRetries; }
+        public Types.TimeoutConfiguration getTimeout() { return timeout; }
 
         public static Builder builder() { return new Builder(); }
 
@@ -2791,6 +2928,8 @@ public static class Builder {
             private List excludeDomains = new ArrayList<>();
             private JsonNode providerOptions;
             private Map headers;
+            private Integer maxRetries;
+            private Types.TimeoutConfiguration timeout;
 
             public Builder query(String v) { this.query = v; return this; }
             public Builder maxResults(Integer v) { this.maxResults = v; return this; }
@@ -2800,10 +2939,15 @@ public static class Builder {
             public Builder excludeDomains(List v) { this.excludeDomains = v; return this; }
             public Builder providerOptions(JsonNode v) { this.providerOptions = v; return this; }
             public Builder headers(Map v) { this.headers = v; return this; }
+            public Builder maxRetries(Integer v) { this.maxRetries = v; return this; }
+            public Builder timeout(Types.TimeoutConfiguration v) { this.timeout = v; return this; }
 
             public SearchCallOptions build() {
-                return new SearchCallOptions(query, maxResults, includeRawContent, timeRange,
+                SearchCallOptions result = new SearchCallOptions(query, maxResults, includeRawContent, timeRange,
                     includeDomains, excludeDomains, providerOptions, headers);
+                result.maxRetries = maxRetries;
+                result.timeout = timeout;
+                return result;
             }
         }
 
@@ -2819,13 +2963,15 @@ public boolean equals(Object o) {
                 && Objects.equals(includeDomains, that.includeDomains)
                 && Objects.equals(excludeDomains, that.excludeDomains)
                 && Objects.equals(providerOptions, that.providerOptions)
-                && Objects.equals(headers, that.headers);
+                && Objects.equals(headers, that.headers)
+                && Objects.equals(maxRetries, that.maxRetries)
+                && Objects.equals(timeout, that.timeout);
         }
 
         @Override
         public int hashCode() {
             return Objects.hash(query, maxResults, includeRawContent, timeRange,
-                includeDomains, excludeDomains, providerOptions, headers);
+                includeDomains, excludeDomains, providerOptions, headers, maxRetries, timeout);
         }
 
         @Override
@@ -2985,7 +3131,3 @@ public boolean equals(Object o) {
 
 
 
-
-
-
-
diff --git a/bindings/java/src/main/java/ai/arcships/aimux/Types.java b/bindings/java/src/main/java/ai/arcships/aimux/Types.java
index db25fa24..534171d6 100644
--- a/bindings/java/src/main/java/ai/arcships/aimux/Types.java
+++ b/bindings/java/src/main/java/ai/arcships/aimux/Types.java
@@ -1605,19 +1605,22 @@ public int hashCode() {
      */
     public static class TimeoutConfiguration {
         @JsonProperty("total_ms") private Long totalMs;
+        @JsonProperty("step_ms") private Long stepMs;
         @JsonProperty("first_chunk_ms") private Long firstChunkMs;
         @JsonProperty("chunk_ms") private Long chunkMs;
 
         @JsonCreator
         TimeoutConfiguration() {}
 
-        private TimeoutConfiguration(Long totalMs, Long firstChunkMs, Long chunkMs) {
+        private TimeoutConfiguration(Long totalMs, Long stepMs, Long firstChunkMs, Long chunkMs) {
             this.totalMs = totalMs;
+            this.stepMs = stepMs;
             this.firstChunkMs = firstChunkMs;
             this.chunkMs = chunkMs;
         }
 
         public Long getTotalMs() { return totalMs; }
+        public Long getStepMs() { return stepMs; }
         public Long getFirstChunkMs() { return firstChunkMs; }
         public Long getChunkMs() { return chunkMs; }
 
@@ -1625,15 +1628,17 @@ private TimeoutConfiguration(Long totalMs, Long firstChunkMs, Long chunkMs) {
 
         public static class Builder {
             private Long totalMs;
+            private Long stepMs;
             private Long firstChunkMs;
             private Long chunkMs;
 
             public Builder totalMs(Long v) { this.totalMs = v; return this; }
+            public Builder stepMs(Long v) { this.stepMs = v; return this; }
             public Builder firstChunkMs(Long v) { this.firstChunkMs = v; return this; }
             public Builder chunkMs(Long v) { this.chunkMs = v; return this; }
 
             public TimeoutConfiguration build() {
-                return new TimeoutConfiguration(totalMs, firstChunkMs, chunkMs);
+                return new TimeoutConfiguration(totalMs, stepMs, firstChunkMs, chunkMs);
             }
         }
 
@@ -1643,13 +1648,14 @@ public boolean equals(Object o) {
             if (!(o instanceof TimeoutConfiguration)) return false;
             TimeoutConfiguration that = (TimeoutConfiguration) o;
             return Objects.equals(totalMs, that.totalMs)
+                && Objects.equals(stepMs, that.stepMs)
                 && Objects.equals(firstChunkMs, that.firstChunkMs)
                 && Objects.equals(chunkMs, that.chunkMs);
         }
 
         @Override
         public int hashCode() {
-            return Objects.hash(totalMs, firstChunkMs, chunkMs);
+            return Objects.hash(totalMs, stepMs, firstChunkMs, chunkMs);
         }
     }
 
diff --git a/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java b/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java
index fca9231e..0cb38965 100644
--- a/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java
+++ b/bindings/java/src/test/java/ai/arcships/aimux/AimuxExceptionTest.java
@@ -220,11 +220,13 @@ void ofMapsEveryVariantToExpectedClass() {
             .isInstanceOf(AimuxException.RequestAbortedError.class);
         assertThat(AimuxException.of(AimuxException.AIMUX_E_OTHER, "m"))
             .isInstanceOf(AimuxException.OtherError.class);
+        assertThat(AimuxException.of(AimuxException.AIMUX_E_RETRY, "m"))
+            .isInstanceOf(AimuxException.RetryError.class);
     }
 
     @Test
     void codesOutsideTheRustEnumAreRejected() {
-        // Out-of-range (the old recording slot 15) is a header/library mismatch.
+        // Out-of-range (15, the first unassigned value) is a header/library mismatch.
         assertThatThrownBy(() -> AimuxException.of(15, ""))
             .isInstanceOf(IllegalStateException.class);
         assertThatThrownBy(() -> AimuxException.of(999, "m"))
@@ -269,6 +271,7 @@ void extractHandleNullErrorPassesThrough() {
     void codeNameCoversKnownCodes() {
         assertThat(AimuxException.codeName(AimuxException.AIMUX_OK)).isEqualTo("OK");
         assertThat(AimuxException.codeName(AimuxException.AIMUX_E_API_CALL)).isEqualTo("ApiCall");
+        assertThat(AimuxException.codeName(AimuxException.AIMUX_E_RETRY)).isEqualTo("Retry");
         assertThat(AimuxException.codeName(999)).startsWith("Code(");
     }
 
@@ -286,8 +289,26 @@ void payloadIsNullForLocalFailures() {
         AimuxException.APICallError api = (AimuxException.APICallError)
             AimuxException.of(AimuxException.AIMUX_E_API_CALL, "x");
         assertThat(api.getProviderCode()).isNull();
-        assertThat(api.getRequestId()).isNull();
         assertThat(api.getResponseBody()).isNull();
+        assertThat(api.getUrl()).isNull();
+        assertThat(api.getRequestBodyValues()).isNull();
+        assertThat(api.getResponseHeaders()).isNull();
+        assertThat(api.getData()).isNull();
         assertThat(api.isRetryable()).isFalse();
     }
+
+    /** Code 14 locally: reason defaults, the single-attempt fallback keeps the type total. */
+    @Test
+    void retryErrorCarriesReasonAndHistory() {
+        AimuxException.RetryError e = (AimuxException.RetryError)
+            AimuxException.of(AimuxException.AIMUX_E_RETRY, "Failed after 2 attempts. Last error: x");
+        assertThat(e.getCode()).isEqualTo(AimuxException.AIMUX_E_RETRY);
+        assertThat(e.getReason()).isEqualTo(AimuxException.RetryErrorReason.MAX_RETRIES_EXCEEDED);
+        assertThat(e.getErrors()).hasSize(1);
+        assertThat(e.getLastError()).isSameAs(e.getErrors().get(0));
+        assertThat(e.getStatusCode()).isEqualTo(-1);
+        assertThatThrownBy(() -> new AimuxException.RetryError(
+                "m", AimuxException.RetryErrorReason.ERROR_NOT_RETRYABLE, java.util.Collections.emptyList()))
+            .isInstanceOf(IllegalArgumentException.class);
+    }
 }
diff --git a/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java b/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java
index 6fd11b8f..b34f50e0 100644
--- a/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java
+++ b/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java
@@ -233,6 +233,26 @@ void videoResultRoundTrip() throws Exception {
             + "\"media_type\":\"video/mp4\"}}],\"warnings\":[],\"response\":{}}"));
     }
 
+    @Test
+    void videoPollOptionsRoundTrip() throws Exception {
+        MultimodalTypes.VideoCallOptions options = MultimodalTypes.VideoCallOptions.builder()
+            .prompt("a cat")
+            .poll(MultimodalTypes.VideoPollOptions.builder()
+                .intervalMs(1_000L)
+                .timeoutMs(120_000L)
+                .build())
+            .build();
+
+        String json = M.writeValueAsString(options);
+        assertThat(M.readTree(json).path("poll").path("interval_ms").asLong()).isEqualTo(1_000L);
+        assertThat(M.readTree(json).path("poll").path("timeout_ms").asLong()).isEqualTo(120_000L);
+
+        MultimodalTypes.VideoCallOptions decoded =
+            M.readValue(json, MultimodalTypes.VideoCallOptions.class);
+        assertThat(decoded.getPoll()).isEqualTo(options.getPoll());
+        assertThat(decoded).isEqualTo(options);
+    }
+
     // ── Search ─────────────────────────────────────────────────────────────
 
     @Test
diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt
index e0fbf36d..85d6a500 100644
--- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt
+++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Errors.kt
@@ -1,12 +1,17 @@
 package ai.arcships.aimux
 
 import com.sun.jna.Pointer
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonElement
+import kotlinx.serialization.json.JsonObject
+import kotlinx.serialization.json.JsonPrimitive
 
 /**
  * Machine-readable codes matching aimux-ffi `aimux_error_code_t` (aimux-error.h).
- * 1..13 mirror the 13 core variants (1 is the catch-all `Other`). The
- * per-status codes (Provider, Http, RateLimited, Auth, ModelNotFound) are
- * gone, every HTTP-shaped failure arrives as [AIMUX_E_API_CALL].
+ * 1..14 mirror the 14 core variants (1 is the catch-all `Other`, 14 is `Retry`).
+ * The per-status codes (Provider, Http, RateLimited, Auth,
+ * ModelNotFound) are gone, every HTTP-shaped failure arrives as
+ * [AIMUX_E_API_CALL].
  */
 const val AIMUX_OK: Int = 0
 const val AIMUX_E_JSON_PARSE: Int = 2
@@ -22,6 +27,7 @@ const val AIMUX_E_API_CALL: Int = 11
 const val AIMUX_E_TIMEOUT: Int = 12
 const val AIMUX_E_ABORTED: Int = 13
 const val AIMUX_E_OTHER: Int = 1
+const val AIMUX_E_RETRY: Int = 14
 
 /**
  * Base exception for every core `AiMuxError`.
@@ -45,7 +51,7 @@ const val AIMUX_E_OTHER: Int = 1
  * }
  * ```
  *
- * Transport: Rust → C `aimux_error_t *` with code 1..13 → [fromC].
+ * Transport: Rust → C `aimux_error_t *` with code 1..14 → [fromC].
  * Primary path is not a JSON
  * error envelope.
  */
@@ -74,8 +80,9 @@ sealed class AimuxException(
          * Reads code, message and retryable for every code, the payload
          * getters only under their owning code, and frees every returned
          * string. Does not own the pointer: the caller ([expectAimuxError]) frees
-         * the returned error afterwards. Code [AIMUX_OK] or a code outside 1..13
-         * is a header/library mismatch and throws [IllegalStateException].
+         * the returned error afterwards (retry attempt errors are new owned
+         * copies and are freed here). Code [AIMUX_OK] or a code outside
+         * 1..14 is a header/library mismatch and throws [IllegalStateException].
          */
         @JvmStatic
         internal fun fromC(error: Pointer, prefix: String = ""): AimuxException {
@@ -93,8 +100,18 @@ sealed class AimuxException(
                     retryable = retryable,
                     providerCode = takeString(lib.aimux_error_provider_code(error)),
                     providerMessage = takeString(lib.aimux_error_provider_message(error)),
-                    requestId = takeString(lib.aimux_error_request_id(error)),
                     responseBody = takeString(lib.aimux_error_response_body(error)),
+                    url = takeString(lib.aimux_error_url(error)),
+                    requestBodyValues = parseJson(takeString(lib.aimux_error_request_body_values(error))),
+                    responseHeaders = headerMap(takeString(lib.aimux_error_response_headers(error))),
+                    data = parseJson(takeString(lib.aimux_error_provider_data(error))),
+                )
+                AIMUX_E_RETRY -> createByCode(
+                    code,
+                    msg,
+                    retryable = retryable,
+                    retryReason = RetryErrorReason.fromWire(takeString(lib.aimux_error_retry_reason(error))),
+                    retryErrors = retryHistory(error),
                 )
                 AIMUX_E_NO_SUCH_MODEL -> createByCode(
                     code,
@@ -114,7 +131,39 @@ sealed class AimuxException(
         }
 
         /**
-         * Build the subclass for a core / C error code (1..13).
+         * Decode the per-attempt history of an [AIMUX_E_RETRY] error. Each
+         * attempt is a new owned `aimux_error_t *` (index 0 = oldest) that can
+         * itself be any AiMuxError — including a nested Retry — and is freed
+         * here, independently of the parent.
+         */
+        private fun retryHistory(error: Pointer): List {
+            val lib = FFI.lib
+            return (0 until lib.aimux_error_retry_count(error)).mapNotNull { i ->
+                lib.aimux_error_retry_error_at(error, i)?.let { attempt ->
+                    try {
+                        fromC(attempt)
+                    } finally {
+                        lib.aimux_error_free(attempt)
+                    }
+                }
+            }
+        }
+
+        /** Parse a getter-returned JSON string; null (absent) stays null. */
+        private fun parseJson(json: String?): JsonElement? = try {
+            json?.let(Json::parseToJsonElement)
+        } catch (_: Exception) {
+            null
+        }
+
+        /** Response headers arrive as one JSON object string of string→string pairs. */
+        private fun headerMap(json: String?): Map? =
+            (parseJson(json) as? JsonObject)?.mapValues { (_, value) ->
+                (value as? JsonPrimitive)?.content ?: value.toString()
+            }
+
+        /**
+         * Build the subclass for a core / C error code (1..14).
          *
          * Any other code — [AIMUX_OK] on a failure path or a code this binding
          * does not know — is a header/library mismatch and throws
@@ -130,11 +179,16 @@ sealed class AimuxException(
             retryable: Boolean = false,
             providerCode: String? = null,
             providerMessage: String? = null,
-            requestId: String? = null,
             responseBody: String? = null,
+            url: String? = null,
+            requestBodyValues: JsonElement? = null,
+            responseHeaders: Map? = null,
+            data: JsonElement? = null,
             modelId: String = "",
             modelType: String = "",
             providerId: String = "",
+            retryReason: RetryErrorReason = RetryErrorReason.MAX_RETRIES_EXCEEDED,
+            retryErrors: List = emptyList(),
         ): AimuxException = when (code) {
             AIMUX_E_JSON_PARSE -> JSONParseError(message, status, retryMs, cause, retryable)
             AIMUX_E_INVALID_RESPONSE_DATA -> InvalidResponseDataError(message, status, retryMs, cause, retryable)
@@ -151,7 +205,18 @@ sealed class AimuxException(
             AIMUX_E_UNSUPPORTED_FUNCTIONALITY -> UnsupportedFunctionalityError(message, status, retryMs, cause, retryable)
             AIMUX_E_NO_SUCH_MODEL -> NoSuchModelError(message, status, retryMs, cause, retryable, modelId, modelType)
             AIMUX_E_NO_SUCH_PROVIDER -> NoSuchProviderError(message, status, retryMs, cause, retryable, providerId)
-            AIMUX_E_API_CALL -> APICallError(message, status, retryMs, cause, retryable, providerCode, providerMessage, requestId, responseBody)
+            AIMUX_E_API_CALL -> APICallError(
+                message, status, retryMs, cause, retryable,
+                providerCode, providerMessage, responseBody,
+                url, requestBodyValues, responseHeaders, data,
+            )
+            AIMUX_E_RETRY -> RetryError(
+                message,
+                retryReason,
+                // A deserialized Retry may carry no attempts; keep RetryError total.
+                retryErrors.ifEmpty { listOf(OtherError(message)) },
+                cause,
+            )
             AIMUX_E_TIMEOUT -> TimeoutError(message, status, retryMs, cause, retryable)
             AIMUX_E_ABORTED -> RequestAbortedError(message, status, retryMs, cause, retryable)
             AIMUX_E_OTHER -> OtherError(message, status, retryMs, cause, retryable)
@@ -175,6 +240,7 @@ sealed class AimuxException(
             AIMUX_E_TIMEOUT -> "Timeout"
             AIMUX_E_ABORTED -> "Aborted"
             AIMUX_E_OTHER -> "Other"
+            AIMUX_E_RETRY -> "Retry"
             else -> "Code($code)"
         }
     }
@@ -268,7 +334,8 @@ class NoSuchProviderError(
  * transport failure — and says nothing about whether a retry would help: read
  * [retryable] for that, never the status sentinel.
  * The provider's own detail, when the response carried it: [providerCode],
- * [providerMessage], [requestId], [responseBody] (null when absent or synthesized locally).
+ * [providerMessage], [responseBody], [url], [requestBodyValues],
+ * [responseHeaders], [data] (null when absent or synthesized locally).
  */
 class APICallError(
     message: String,
@@ -280,12 +347,50 @@ class APICallError(
     val providerCode: String? = null,
     /** The failure's own text without the composed prefix [message] carries, e.g. "slow down". */
     val providerMessage: String? = null,
-    /** Provider request id, for support tickets. */
-    val requestId: String? = null,
     /** Raw response body. */
     val responseBody: String? = null,
+    /** Sanitized request URL. */
+    val url: String? = null,
+    /** Sanitized request body values (any JSON type). */
+    val requestBodyValues: JsonElement? = null,
+    /** Sanitized response headers (also carry provider request ids). */
+    val responseHeaders: Map? = null,
+    /** Parsed provider error data (AI SDK `APICallError.data`). */
+    val data: JsonElement? = null,
 ) : AimuxException(message, AIMUX_E_API_CALL, status, retryMs, cause, retryable)
 
+/** Why the retry loop gave up ([RetryError.reason]); wire names are the core's serde camelCase. */
+enum class RetryErrorReason(val wireValue: String) {
+    /** Every permitted attempt failed with a retryable error. */
+    MAX_RETRIES_EXCEEDED("maxRetriesExceeded"),
+
+    /** A later attempt failed with a non-retryable error. */
+    ERROR_NOT_RETRYABLE("errorNotRetryable");
+
+    companion object {
+        fun fromWire(value: String?): RetryErrorReason =
+            if (value == ERROR_NOT_RETRYABLE.wireValue) ERROR_NOT_RETRYABLE else MAX_RETRIES_EXCEEDED
+    }
+}
+
+/**
+ * The retry loop gave up (AI SDK `RetryError` analogue): [reason] says why,
+ * [errors] is the per-attempt history (oldest first), [lastError] the final
+ * attempt.
+ */
+class RetryError(
+    message: String,
+    val reason: RetryErrorReason,
+    val errors: List,
+    cause: Throwable? = null,
+) : AimuxException(message, AIMUX_E_RETRY, -1, -1, cause) {
+    init {
+        require(errors.isNotEmpty()) { "RetryError requires at least one error" }
+    }
+
+    val lastError: AimuxException = errors.last()
+}
+
 class TimeoutError(
     message: String,
     status: Int = -1,
@@ -294,7 +399,7 @@ class TimeoutError(
     retryable: Boolean = false,
 ) : AimuxException(message, AIMUX_E_TIMEOUT, status, retryMs, cause, retryable)
 
-/** Request aborted (not a Java interruption). */
+/** Request aborted (not a Java interruption); the message is the abort payload. */
 class RequestAbortedError(
     message: String = "request aborted",
     status: Int = -1,
diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt
index 798c8850..d5b19f79 100644
--- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt
+++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Model.kt
@@ -8,7 +8,7 @@
  * Errors: every fallible C call returns an `aimux_error_t *` ([Pointer]?):
  * null = success, result written to the trailing out-parameter
  * ([LongByReference] for handles, [PointerByReference] for JSON strings);
- * non-null = failure. Its unified code identifies [AimuxException] (1..13),
+ * non-null = failure. Its unified code identifies [AimuxException] (1..14),
  * [RecordingException] (100..105), or a C ABI failure (200..206). The last
  * range maps to `IllegalStateException("aimux ffi: …")`. A decoder releases
  * the returned pointer with `aimux_error_free`.
@@ -117,8 +117,16 @@ internal interface AimuxFFI : Library {
     fun aimux_error_retry_ms(error: Pointer?): Long
     fun aimux_error_provider_code(error: Pointer?): Pointer?
     fun aimux_error_provider_message(error: Pointer?): Pointer?
-    fun aimux_error_request_id(error: Pointer?): Pointer?
     fun aimux_error_response_body(error: Pointer?): Pointer?
+    fun aimux_error_url(error: Pointer?): Pointer?
+    fun aimux_error_request_body_values(error: Pointer?): Pointer?
+    fun aimux_error_response_headers(error: Pointer?): Pointer?
+    fun aimux_error_provider_data(error: Pointer?): Pointer?
+    // AIMUX_E_RETRY payload: reason wire name, attempt count, and each attempt
+    // as a NEW OWNED error the caller frees with aimux_error_free.
+    fun aimux_error_retry_reason(error: Pointer?): Pointer?
+    fun aimux_error_retry_count(error: Pointer?): Int
+    fun aimux_error_retry_error_at(error: Pointer?, index: Int): Pointer?
     fun aimux_error_model_id(error: Pointer?): Pointer?
     fun aimux_error_model_type(error: Pointer?): Pointer?
     fun aimux_error_provider_id(error: Pointer?): Pointer?
@@ -251,7 +259,7 @@ private fun ffiError(e: Pointer, prefix: String): IllegalStateException {
 }
 
 /**
- * Decode an error from a call that may return `AiMuxError`: 1..13 →
+ * Decode an error from a call that may return `AiMuxError`: 1..14 →
  * [AimuxException]; 200..206 → [IllegalStateException]. Frees [e].
  */
 internal fun expectAimuxError(e: Pointer, context: String = ""): RuntimeException {
@@ -259,7 +267,7 @@ internal fun expectAimuxError(e: Pointer, context: String = ""): RuntimeExceptio
     try {
         val code = FFI.lib.aimux_error_code(e)
         if (isFfiCode(code)) return ffiError(e, prefix)
-        check(code in AIMUX_E_OTHER..AIMUX_E_ABORTED) {
+        check(code in AIMUX_E_OTHER..AIMUX_E_ABORTED || code == AIMUX_E_RETRY) {
             "${prefix}aimux ffi: expected AiMuxError code, got $code"
         }
         return AimuxException.fromC(e, prefix)
@@ -512,9 +520,11 @@ class Model internal constructor(handle: Long) : Closeable {
     /**
      * Stream text from the model.
      *
-     * Blocks the calling thread until the stream completes. Stream failures
-     * throw [AimuxException] (on_done is not invoked on failure). Mid-stream
-     * errors surface the same way after any parts already delivered to [onPart].
+     * Blocks the calling thread until the stream completes. Recoverable
+     * frame errors (a malformed SSE frame) arrive as `StreamPart::Error`
+     * data parts and the stream continues; only transport/Core failures
+     * throw [AimuxException] (on_done is not invoked on failure), after any
+     * parts already delivered to [onPart].
      *
      * @param promptJson JSON prompt string.
      * @param optsJson Optional JSON-serialized GenerateTextOptions.
diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt
index f117c34f..8cdeb6bc 100644
--- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt
+++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Multimodal.kt
@@ -8,7 +8,7 @@
  * JSON strings (base64 for binary), matching the C ABI wire format.
  *
  * Every fallible C call returns an `aimux_error_t *` (null = success) that
- * [expectAimuxError] decodes codes 1..13 as [AimuxException] and 200..206 as
+ * [expectAimuxError] decodes codes 1..14 as [AimuxException] and 200..206 as
  * [IllegalStateException] `"aimux ffi: …"` (malformed raw JSON is
  * caught before the C call by [requireJson] as [IllegalArgumentException]).
  * No JSON envelope on the primary path.
diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/MultimodalTypes.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/MultimodalTypes.kt
index 3f7955a5..3f547296 100644
--- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/MultimodalTypes.kt
+++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/MultimodalTypes.kt
@@ -78,6 +78,8 @@ data class EmbeddingResult(
 @Serializable
 data class EmbeddingCallOptions(
     val values: List = emptyList(),
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
@@ -134,6 +136,8 @@ data class SpeechCallOptions(
     val instructions: String? = null,
     val speed: Double? = null,
     val language: String? = null,
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
@@ -193,6 +197,8 @@ data class ImageCallOptions(
     val seed: Long? = null,
     val files: List = emptyList(),
     val mask: JsonElement? = null,
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
@@ -246,6 +252,8 @@ data class TranscriptionResult(
 data class TranscriptionCallOptions(
     val audio: JsonElement = JsonObject(emptyMap()),
     @SerialName("media_type") val mediaType: String = "",
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
@@ -286,6 +294,8 @@ data class RerankingCallOptions(
     val documents: JsonElement = JsonArray(emptyList()),
     val query: String = "",
     @SerialName("top_n") val topN: Int? = null,
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
@@ -350,6 +360,13 @@ data class VideoResult(
     val response: VideoResponse = VideoResponse(),
 )
 
+/** Per-call pacing overrides for the Core-owned video status poll loop. */
+@Serializable
+data class VideoPollOptions(
+    @SerialName("interval_ms") val intervalMs: Long? = null,
+    @SerialName("timeout_ms") val timeoutMs: Long? = null,
+)
+
 /** Options for video generation. */
 @Serializable
 data class VideoCallOptions(
@@ -358,6 +375,9 @@ data class VideoCallOptions(
     @SerialName("aspect_ratio") val aspectRatio: String? = null,
     val resolution: String? = null,
     val seed: Long? = null,
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val poll: VideoPollOptions? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
@@ -402,6 +422,8 @@ data class SearchCallOptions(
     @SerialName("time_range") val timeRange: String? = null,
     @SerialName("include_domains") val includeDomains: List = emptyList(),
     @SerialName("exclude_domains") val excludeDomains: List = emptyList(),
+    @SerialName("max_retries") val maxRetries: Long? = null,
+    val timeout: TimeoutConfiguration? = null,
     @SerialName("provider_options") val providerOptions: JsonElement? = null,
     val headers: Map? = null,
 )
diff --git a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt
index d9d44144..15980811 100644
--- a/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt
+++ b/bindings/kotlin/src/main/kotlin/ai/arcships/aimux/Types.kt
@@ -564,6 +564,8 @@ data class ModelMessage(
 data class TimeoutConfiguration(
     /** Overall timeout for the entire call (including retries and, for streaming, the whole stream), in milliseconds. */
     @SerialName("total_ms") val totalMs: Long? = null,
+    /** Timeout for one model operation attempt, in milliseconds. */
+    @SerialName("step_ms") val stepMs: Long? = null,
     /** Timeout waiting for the first stream chunk (streaming only). */
     @SerialName("first_chunk_ms") val firstChunkMs: Long? = null,
     /** Maximum idle time between consecutive stream chunks (streaming only). */
diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ContractTest.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ContractTest.kt
index eb0bd7dd..00dc0074 100644
--- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ContractTest.kt
+++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ContractTest.kt
@@ -147,6 +147,22 @@ class ContractTest {
         assertThat(opts.maxRetries).isEqualTo(3L)
     }
 
+    @Test
+    fun `video poll options use the core wire names`() {
+        val opts = VideoCallOptions(
+            prompt = "a cat",
+            poll = VideoPollOptions(intervalMs = 1_000L, timeoutMs = 120_000L),
+        )
+
+        val encoded = AimuxJson.encodeToString(VideoCallOptions.serializer(), opts)
+        assertThat(encoded).contains(
+            "\"poll\":{\"interval_ms\":1000,\"timeout_ms\":120000}",
+        )
+        val decoded = AimuxJson.decodeFromString(encoded)
+        assertThat(decoded.poll?.intervalMs).isEqualTo(1_000L)
+        assertThat(decoded.poll?.timeoutMs).isEqualTo(120_000L)
+    }
+
     /// RFC-0016 M10: `Usage.raw` with a vendor-specific field survives a
     /// Kotlin round-trip.
     @Test
diff --git a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt
index f1c934b6..04af9b5a 100644
--- a/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt
+++ b/bindings/kotlin/src/test/kotlin/ai/arcships/aimux/ErrorsTest.kt
@@ -2,6 +2,7 @@ package ai.arcships.aimux
 
 import com.sun.jna.ptr.LongByReference
 import com.sun.jna.ptr.PointerByReference
+import kotlinx.serialization.json.Json
 import org.assertj.core.api.Assertions.assertThat
 import org.assertj.core.api.Assertions.assertThatThrownBy
 import org.junit.jupiter.api.Test
@@ -125,15 +126,21 @@ class ErrorsTest {
         val apiEx = AimuxException.createByCode(
             AIMUX_E_API_CALL, "API call error: HTTP 429: slow down", 429, 1500,
             retryable = true, providerCode = "insufficient_quota", providerMessage = "slow down",
-            requestId = "req_123", responseBody = "{\"error\":{}}",
+            responseBody = "{\"error\":{}}", url = "https://api.example/v1/chat",
+            requestBodyValues = Json.parseToJsonElement("""{"model":"m"}"""),
+            responseHeaders = mapOf("retry-after-ms" to "1500"),
+            data = Json.parseToJsonElement("""{"code":"quota"}"""),
         ) as APICallError
         assertThat(apiEx.status).isEqualTo(429)
         assertThat(apiEx.retryMs).isEqualTo(1500L)
         assertThat(apiEx.retryable).isTrue()
         assertThat(apiEx.providerCode).isEqualTo("insufficient_quota")
         assertThat(apiEx.providerMessage).isEqualTo("slow down")
-        assertThat(apiEx.requestId).isEqualTo("req_123")
         assertThat(apiEx.responseBody).isEqualTo("{\"error\":{}}")
+        assertThat(apiEx.url).isEqualTo("https://api.example/v1/chat")
+        assertThat(apiEx.requestBodyValues.toString()).isEqualTo("""{"model":"m"}""")
+        assertThat(apiEx.responseHeaders).isEqualTo(mapOf("retry-after-ms" to "1500"))
+        assertThat(apiEx.data.toString()).isEqualTo("""{"code":"quota"}""")
 
         val modelEx = AimuxException.createByCode(
             AIMUX_E_NO_SUCH_MODEL, "no such model", modelId = "gpt-nope", modelType = "language",
@@ -142,11 +149,44 @@ class ErrorsTest {
         assertThat(modelEx.modelType).isEqualTo("language")
 
         // Absent: null on APICallError, "" on the id classes.
-        assertThat((AimuxException.createByCode(AIMUX_E_API_CALL, "x") as APICallError).requestId).isNull()
+        assertThat((AimuxException.createByCode(AIMUX_E_API_CALL, "x") as APICallError).url).isNull()
+        assertThat((AimuxException.createByCode(AIMUX_E_API_CALL, "x") as APICallError).responseHeaders).isNull()
         assertThat((AimuxException.createByCode(AIMUX_E_NO_SUCH_MODEL, "x") as NoSuchModelError).modelId).isEmpty()
         assertThat((AimuxException.createByCode(AIMUX_E_NO_SUCH_PROVIDER, "x") as NoSuchProviderError).providerId).isEmpty()
     }
 
+    /** Code 14: the retry loop's verdict plus the per-attempt history, oldest first. */
+    @Test
+    fun `createByCode builds RetryError with reason and attempt history`() {
+        val attempts = listOf(
+            AimuxException.createByCode(AIMUX_E_API_CALL, "HTTP 500", 500, retryable = true),
+            AimuxException.createByCode(AIMUX_E_API_CALL, "HTTP 401", 401),
+        )
+        val ex = AimuxException.createByCode(
+            AIMUX_E_RETRY, "Failed after 2 attempts with non-retryable error: 'HTTP 401'",
+            retryReason = RetryErrorReason.ERROR_NOT_RETRYABLE, retryErrors = attempts,
+        ) as RetryError
+        assertThat(ex.code).isEqualTo(AIMUX_E_RETRY)
+        assertThat(ex.reason).isEqualTo(RetryErrorReason.ERROR_NOT_RETRYABLE)
+        assertThat(ex.errors).hasSize(2)
+        assertThat(ex.lastError.status).isEqualTo(401)
+        assertThat(ex.status).isEqualTo(-1)
+
+        // Defensive: an empty history collapses to one OtherError carrying the message.
+        val fallback = AimuxException.createByCode(AIMUX_E_RETRY, "gave up") as RetryError
+        assertThat(fallback.errors).hasSize(1)
+        assertThat(fallback.lastError).isInstanceOf(OtherError::class.java)
+        assertThat(fallback.reason).isEqualTo(RetryErrorReason.MAX_RETRIES_EXCEEDED)
+    }
+
+    /** The wire names are the core's serde camelCase; unknown text falls back to the common case. */
+    @Test
+    fun `RetryErrorReason maps the wire names`() {
+        assertThat(RetryErrorReason.fromWire("errorNotRetryable")).isEqualTo(RetryErrorReason.ERROR_NOT_RETRYABLE)
+        assertThat(RetryErrorReason.fromWire("maxRetriesExceeded")).isEqualTo(RetryErrorReason.MAX_RETRIES_EXCEEDED)
+        assertThat(RetryErrorReason.fromWire(null)).isEqualTo(RetryErrorReason.MAX_RETRIES_EXCEEDED)
+    }
+
     /** 401 / 404 are the same class; only the status distinguishes them. */
     @Test
     fun `createByCode maps HTTP statuses onto APICallError`() {
@@ -175,11 +215,14 @@ class ErrorsTest {
         assertThat(noKey.retryable).isFalse()
     }
 
-    /** A code outside 1..13 is a header/library mismatch, not an error type. */
+    /** A code outside 1..14 is a header/library mismatch, not an error type. */
     @Test
     fun `createByCode rejects codes outside the enum`() {
         assertThatThrownBy { AimuxException.createByCode(999, "?") }
             .isInstanceOf(IllegalStateException::class.java)
+        // 15 is the first unassigned value and is rejected.
+        assertThatThrownBy { AimuxException.createByCode(15, "?") }
+            .isInstanceOf(IllegalStateException::class.java)
         assertThatThrownBy { AimuxException.createByCode(AIMUX_OK, "?") }
             .isInstanceOf(IllegalStateException::class.java)
     }
diff --git a/bindings/node/Cargo.toml b/bindings/node/Cargo.toml
index 3c0c3129..0c06edcb 100644
--- a/bindings/node/Cargo.toml
+++ b/bindings/node/Cargo.toml
@@ -16,7 +16,7 @@ serde = { version = "1", features = ["derive"] }
 serde_json = { version = "1", features = ["preserve_order"] }
 tokio = { version = "1", features = ["full"] }
 futures = "0.3"
-napi = { version = "3", features = ["napi8", "tokio_rt", "async"] }
+napi = { version = "3", features = ["napi8", "tokio_rt", "async", "serde-json"] }
 napi-derive = "3"
 
 # Keep this crate out of the main cargo workspace — napi-rs builds it independently.
diff --git a/bindings/node/__test__/error.test.ts b/bindings/node/__test__/error.test.ts
index e2aca314..5540c45e 100644
--- a/bindings/node/__test__/error.test.ts
+++ b/bindings/node/__test__/error.test.ts
@@ -12,6 +12,7 @@ import {
   streamText,
   AimuxError,
   APICallError,
+  RetryError,
   InvalidArgumentError,
   NoSuchProviderError,
   RecordingError,
@@ -95,7 +96,8 @@ test('payload fields live on the subclass the core fills them for', async (t) =>
 
   // ApiCall-only fields are absent here — the core fills them for no other
   // variant, so they must not exist as always-empty properties.
-  for (const key of ['status', 'retryMs', 'providerCode', 'responseBody', 'requestId']) {
+  const apiCallOnly = ['status', 'retryMs', 'providerCode', 'responseBody', 'url', 'requestBodyValues', 'responseHeaders', 'data']
+  for (const key of apiCallOnly) {
     t.false(key in err, `${key} must not exist on ${err.name}`)
   }
   t.false('retryable' in err)
@@ -126,8 +128,12 @@ test('APICallError carries every field the response produced', async (t) => {
     t.is(err.retryMs, 1500)
     t.true(err.retryable)
     t.is(err.providerCode, 'rate_limit_exceeded')
-    t.is(err.requestId, 'req_abc123')
     t.is(err.responseBody, body)
+    // Request context is carried on the error itself.
+    t.true(err.url?.startsWith(url), `url should start with ${url}: ${err.url}`)
+    t.is((err.requestBodyValues as { model?: string }).model, 'gpt-4o')
+    // Sanitized response headers, as sent (retryMs above is derived from them).
+    t.is(err.responseHeaders?.['retry-after-ms'], '1500')
     // The provider's text on its own; `message` is the composed form.
     t.is(err.providerMessage, 'slow down')
     t.true(err.message.includes('slow down'))
@@ -137,6 +143,42 @@ test('APICallError carries every field the response produced', async (t) => {
   }
 })
 
+test('exhausted retries throw RetryError with the per-attempt history', async (t) => {
+  const body = JSON.stringify({
+    error: { message: 'slow down', type: 'rate_limit_exceeded' },
+  })
+  const { server, url } = await startMockServer((_req, res) => {
+    // retry-after-ms 0 keeps the retry loop instant.
+    res.writeHead(429, { 'content-type': 'application/json', 'retry-after-ms': '0' })
+    res.end(body)
+  })
+  try {
+    const model = await native.openai('test-key', 'gpt-4o', { baseUrl: url, maxRetries: 1 })
+    const err = (await t.throwsAsync(() =>
+      model.generateText(JSON.stringify('hi')),
+    )) as RetryError
+    t.true(err instanceof RetryError)
+    t.true(err instanceof AimuxError)
+    t.false(err instanceof APICallError)
+    t.is(err.reason, 'maxRetriesExceeded')
+
+    // Complete per-attempt history, oldest first, each a full typed error.
+    t.is(err.errors.length, 2)
+    for (const attempt of err.errors) {
+      t.true(attempt instanceof APICallError)
+      t.is((attempt as APICallError).status, 429)
+      t.is((attempt as APICallError).providerCode, 'rate_limit_exceeded')
+    }
+    t.is(err.lastError, err.errors[1])
+    t.regex(err.message, /Failed after 2 attempts/)
+    // Retry exhaustion is not itself an API exchange: no ApiCall-only fields.
+    t.false('status' in err)
+    t.false('retryable' in err)
+  } finally {
+    await new Promise((resolve) => server.close(() => resolve()))
+  }
+})
+
 test('recorder init failures are RecordingError; flush is a no-op when none is installed', (t) => {
   // Nothing recording: nothing to flush is a success.
   recordingStop()
diff --git a/bindings/node/__test__/index.test.ts b/bindings/node/__test__/index.test.ts
index 010a5156..fedf568d 100644
--- a/bindings/node/__test__/index.test.ts
+++ b/bindings/node/__test__/index.test.ts
@@ -7,6 +7,7 @@ import {
   initRecordingRing,
   recordingStop,
 } from '../src/native.ts'
+import type { VideoCallOptions, VideoPollOptions } from '../src/index.ts'
 
 // These tests verify the native module loads and the API surface works.
 // They do NOT make real API calls — they test error handling for invalid keys.
@@ -18,6 +19,12 @@ test('native module loads and exports functions', (t) => {
   t.is(typeof provider, 'function')
 })
 
+test('video poll options are exported from the public type barrel', (t) => {
+  const poll: VideoPollOptions = { interval_ms: 1_000, timeout_ms: 120_000 }
+  const options: Pick = { poll }
+  t.deepEqual(options.poll, poll)
+})
+
 test('openai() creates a model instance with valid key format', async (t) => {
   // Even with a fake key, the provider should construct (the key is only
   // validated on the first API call, not at construction time).
diff --git a/bindings/node/src/error.rs b/bindings/node/src/error.rs
index 4ec399c1..da924f90 100644
--- a/bindings/node/src/error.rs
+++ b/bindings/node/src/error.rs
@@ -7,7 +7,7 @@
 
 use std::collections::HashMap;
 
-use aimux_core::{AiMuxError, recording::RecordingError};
+use aimux_core::{AiMuxError, RetryErrorReason, recording::RecordingError};
 use napi::bindgen_prelude::{
     Function, FunctionRef, JsValue, Object, Result as NapiResult, ToNapiValue, TypeName, ValueType,
 };
@@ -16,6 +16,7 @@ use napi_derive::napi;
 
 const ERROR_CLASS_NAMES: &[&str] = &[
     "APICallError",
+    "RetryError",
     "JSONParseError",
     "InvalidResponseDataError",
     "ToolError",
@@ -166,6 +167,7 @@ pub(crate) fn serialize_result(value: &T) -> MResult &'static str {
     match error {
         AiMuxError::ApiCall(_) => "APICallError",
+        AiMuxError::Retry(_) => "RetryError",
         AiMuxError::JsonParse(_) => "JSONParseError",
         AiMuxError::InvalidResponseData(_) => "InvalidResponseDataError",
         AiMuxError::Tool(_) => "ToolError",
@@ -176,7 +178,7 @@ fn aimux_error_class_name(error: &AiMuxError) -> &'static str {
         AiMuxError::NoSuchModel { .. } => "NoSuchModelError",
         AiMuxError::NoSuchProvider { .. } => "NoSuchProviderError",
         AiMuxError::Timeout(_) => "TimeoutError",
-        AiMuxError::Aborted => "RequestAbortedError",
+        AiMuxError::Aborted(_) => "RequestAbortedError",
         AiMuxError::Other(_) => "OtherError",
     }
 }
@@ -217,30 +219,57 @@ fn new_registered_error<'env>(
 
 /// Build the exact JS subclass registered for this core variant.
 fn create_aimux_throwable(env: &Env, error: &AiMuxError) -> NapiResult {
+    Ok(Error::from(create_aimux_error_object(env, error)?.to_unknown()))
+}
+
+/// Instantiate the registered subclass and fill its variant-owned fields.
+/// Recursive: a `Retry` error carries its attempt history as full instances.
+fn create_aimux_error_object<'env>(env: &'env Env, error: &AiMuxError) -> NapiResult> {
     let message = error.to_string();
     let mut obj = new_registered_error(env, aimux_error_class_name(error), &message)?;
     if let Some(status) = error.status_code() {
         obj.set("status", i32::from(status))?;
     }
-    if let AiMuxError::ApiCall(detail) = error {
-        obj.set("retryable", error.is_retryable())?;
-        if let Some(retry_ms) = error.retry_after_hint() {
-            obj.set("retryMs", retry_ms)?;
-        }
-        if let Some(provider_code) = &detail.provider_code {
-            obj.set("providerCode", provider_code.as_str())?;
-        }
-        if !detail.message.is_empty() {
-            obj.set("providerMessage", detail.message.as_str())?;
-        }
-        if let Some(request_id) = &detail.request_id {
-            obj.set("requestId", request_id.as_str())?;
+    match error {
+        AiMuxError::ApiCall(detail) => {
+            obj.set("retryable", detail.is_retryable)?;
+            if let Some(retry_ms) = error.retry_after_hint() {
+                obj.set("retryMs", retry_ms)?;
+            }
+            if !detail.url.is_empty() {
+                obj.set("url", detail.url.as_str())?;
+            }
+            if !detail.request_body_values.is_null() {
+                obj.set("requestBodyValues", detail.request_body_values.clone())?;
+            }
+            if let Some(provider_code) = &detail.provider_code {
+                obj.set("providerCode", provider_code.as_str())?;
+            }
+            if !detail.message.is_empty() {
+                obj.set("providerMessage", detail.message.as_str())?;
+            }
+            if let Some(response_body) = &detail.response_body {
+                obj.set("responseBody", response_body.as_str())?;
+            }
+            if let Some(headers) = &detail.response_headers {
+                obj.set("responseHeaders", headers.clone())?;
+            }
+            if let Some(data) = &detail.data {
+                obj.set("data", data.clone())?;
+            }
         }
-        if let Some(response_body) = &detail.response_body {
-            obj.set("responseBody", response_body.as_str())?;
+        AiMuxError::Retry(retry) => {
+            let reason = match retry.reason {
+                RetryErrorReason::MaxRetriesExceeded => "maxRetriesExceeded",
+                RetryErrorReason::ErrorNotRetryable => "errorNotRetryable",
+            };
+            obj.set("reason", reason)?;
+            let mut attempts = Vec::with_capacity(retry.errors.len());
+            for attempt in &retry.errors {
+                attempts.push(create_aimux_error_object(env, attempt)?);
+            }
+            obj.set("errors", attempts)?;
         }
-    }
-    match error {
         AiMuxError::NoSuchModel {
             model_id,
             model_type,
@@ -255,7 +284,7 @@ fn create_aimux_throwable(env: &Env, error: &AiMuxError) -> NapiResult {
         }
         _ => {}
     }
-    Ok(Error::from(obj.to_unknown()))
+    Ok(obj)
 }
 
 /// A [`BindingError`] as napi-rs's own plain `Error`: caller-side faults
diff --git a/bindings/node/src/error.ts b/bindings/node/src/error.ts
index 3ec248aa..936452a1 100644
--- a/bindings/node/src/error.ts
+++ b/bindings/node/src/error.ts
@@ -7,6 +7,8 @@
  * } catch (e) {
  *   if (e instanceof APICallError) {
  *     // e.status === 429 → rate limited (e.retryMs), 401 → auth, 404 → model
+ *   } else if (e instanceof RetryError) {
+ *     // every retry attempt failed — e.errors is the per-attempt history
  *   } else if (e instanceof AimuxError) {
  *     // any AiMuxError
  *   } else if (e instanceof Error && (e as { code?: string }).code === 'InvalidArg') {
@@ -38,15 +40,42 @@ export class APICallError extends AimuxError {
   declare readonly status?: number
   /** Rate-limit hint in ms; absent if none; `0` means retry immediately. */
   declare readonly retryMs?: number
+  /** Sanitized request URL; absent when unknown. */
+  declare readonly url?: string
+  /** Sanitized request body values sent to the provider. */
+  declare readonly requestBodyValues?: unknown
   /** Provider's machine-readable error code (e.g. `'rate_limit_exceeded'`). */
   declare readonly providerCode?: string
   /** Provider's failure text without Aimux's composed prefix. */
   declare readonly providerMessage?: string
   /** Raw error response body, verbatim. */
   declare readonly responseBody?: string
-  /** Provider-assigned request id (`x-request-id` / `request-id` header). */
-  declare readonly requestId?: string
+  /** Sanitized response headers. */
+  declare readonly responseHeaders?: Record
+  /** Parsed provider error data (AI SDK `APICallError.data`). */
+  declare readonly data?: unknown
 }
+
+/** Why a {@link RetryError} stopped retrying (core serde wire names). */
+export type RetryErrorReason = 'maxRetriesExceeded' | 'errorNotRetryable'
+
+/**
+ * A retried model operation gave up. `errors` is the complete per-attempt
+ * history, oldest first; each entry is itself a full member of this hierarchy
+ * (typically {@link APICallError} with its provider detail).
+ */
+export class RetryError extends AimuxError {
+  /** Why retrying stopped. */
+  declare readonly reason: RetryErrorReason
+  /** Per-attempt errors, oldest first. */
+  declare readonly errors: AimuxError[]
+
+  /** The final attempt's error (`undefined` only for a deserialized empty history). */
+  get lastError(): AimuxError | undefined {
+    return this.errors[this.errors.length - 1]
+  }
+}
+
 export class JSONParseError extends AimuxError {}
 export class InvalidResponseDataError extends AimuxError {}
 export class ToolError extends AimuxError {}
@@ -70,7 +99,7 @@ export class NoSuchProviderError extends AimuxError {
   declare readonly providerId: string
 }
 export class TimeoutError extends AimuxError {}
-/** Request aborted (not DOM `AbortError`). */
+/** Request aborted (not DOM `AbortError`); `message` is the abort payload. */
 export class RequestAbortedError extends AimuxError {}
 export class OtherError extends AimuxError {}
 
diff --git a/bindings/node/src/index.ts b/bindings/node/src/index.ts
index 54dd8998..34f1faf8 100644
--- a/bindings/node/src/index.ts
+++ b/bindings/node/src/index.ts
@@ -50,12 +50,16 @@ import type {
 
   ModelSpec,
   RuntimeModel,
+  VideoCallOptions,
+  VideoPollOptions,
 } from './types'
 
 // Error hierarchy (throw/catch). Wire payload type `AiMuxError` lives under StreamPart only.
 export {
   AimuxError,
   APICallError,
+  RetryError,
+  type RetryErrorReason,
   JSONParseError,
   InvalidResponseDataError,
   ToolError,
@@ -147,6 +151,8 @@ export type {
 
   ModelSpec,
   RuntimeModel,
+  VideoCallOptions,
+  VideoPollOptions,
 }
 
 /**
diff --git a/bindings/node/src/lib.rs b/bindings/node/src/lib.rs
index 5b44ab1c..10d8a941 100644
--- a/bindings/node/src/lib.rs
+++ b/bindings/node/src/lib.rs
@@ -56,14 +56,14 @@ pub struct Model {
 /// ```
 #[napi]
 pub struct AbortBridge {
-    signal: Arc,
+    signal: Arc,
 }
 
 #[napi]
 impl AbortBridge {
     #[napi(constructor)]
     pub fn new(signal: napi::bindgen_prelude::AbortSignal) -> Self {
-        let core = aimux_core::shared::AbortSignal::new();
+        let core = aimux_core::AbortSignal::new();
         let watcher = core.clone();
         signal.on_abort(move || watcher.abort());
         Self {
@@ -79,7 +79,7 @@ impl AbortBridge {
 }
 
 impl AbortBridge {
-    fn core_signal(&self) -> aimux_core::shared::AbortSignal {
+    fn core_signal(&self) -> aimux_core::AbortSignal {
         (*self.signal).clone()
     }
 }
@@ -356,6 +356,21 @@ impl Model {
                                             break; // receiver dropped (JS stopped iterating)
                                         }
                                     }
+                                    Err(e) if e.is_recoverable_stream_error() => {
+                                        // Core keeps the stream alive across a malformed
+                                        // frame; deliver it as a StreamPart::Error data item
+                                        // and keep pumping.
+                                        match serde_json::to_string(
+                                            &aimux_core::stream_part::StreamPart::Error { error: e },
+                                        ) {
+                                            Ok(json) => {
+                                                if tx.send(Ok(json)).await.is_err() {
+                                                    break;
+                                                }
+                                            }
+                                            Err(_) => break,
+                                        }
+                                    }
                                     Err(e) => {
                                         let _ = tx.send(Err(AiMuxBindingError::from(&e))).await;
                                         break;
@@ -487,6 +502,11 @@ impl Model {
                                             break;
                                         }
                                     }
+                                    Err(e) if e.is_recoverable_stream_error() => {
+                                        // Consumers type this path as ChatCompletionChunk;
+                                        // the error cannot ride it — skip and keep pumping
+                                        // (full fidelity lives on the StreamPart path).
+                                    }
                                     Err(e) => {
                                         let _ = tx.send(Err(AiMuxBindingError::from(&e))).await;
                                         break;
@@ -603,10 +623,7 @@ fn apply_provider_config_openai(
         config = config.with_project(proj);
     }
     if let Some(max) = cfg.max_retries {
-        config = config.with_retry_config(aimux_provider_utils::RetryConfig {
-            max_retries: max,
-            ..aimux_provider_utils::RetryConfig::default()
-        });
+        config.retry_config.max_retries = max;
     }
     if let Some(ref json_str) = cfg.body_overrides {
         let overrides: serde_json::Value = parse_wire_json("config.bodyOverrides", json_str)?;
@@ -971,10 +988,7 @@ pub async fn anthropic(
                         cfg = cfg.with_headers(h);
                     }
                     if let Some(max) = opts.max_retries {
-                        cfg = cfg.with_retry_config(aimux_provider_utils::RetryConfig {
-                            max_retries: max,
-                            ..aimux_provider_utils::RetryConfig::default()
-                        });
+                        cfg.retry_config.max_retries = max;
                     }
                     if let Some(ref json_str) = opts.body_overrides {
                         let overrides: serde_json::Value =
diff --git a/bindings/node/src/multimodal.rs b/bindings/node/src/multimodal.rs
index 8dadd640..665cf32d 100644
--- a/bindings/node/src/multimodal.rs
+++ b/bindings/node/src/multimodal.rs
@@ -51,9 +51,7 @@ impl EmbeddingModel {
                 let values: Vec = parse_wire_json("values_json", &values_json)?;
                 opts.values = values;
 
-                let result = self
-                    .inner
-                    .do_embed(&opts)
+                let result = aimux_core::embedding_model::embed(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -89,9 +87,7 @@ impl SpeechModel {
                 let mut opts: SpeechCallOptions = parse_wire_json("opts_json", &opts_json)?;
                 opts.abort_signal = bridge.map(|b| b.core_signal());
 
-                let result = self
-                    .inner
-                    .do_generate(&opts)
+                let result = aimux_core::speech_model::generate_speech(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -127,9 +123,7 @@ impl ImageModel {
                 let mut opts: ImageCallOptions = parse_wire_json("opts_json", &opts_json)?;
                 opts.abort_signal = bridge.map(|b| b.core_signal());
 
-                let result = self
-                    .inner
-                    .do_generate(&opts)
+                let result = aimux_core::image_model::generate_image(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -169,18 +163,19 @@ impl TranscriptionModel {
                     TranscriptionCallOptions::new(AudioInput::Base64(audio_base64), media_type);
                 if let Some(s) = opts_json.as_deref() {
                     if !s.trim().is_empty() && s.trim() != "null" {
-                        let parsed: TranscriptionCallOptions = parse_wire_json("opts_json", s)?;
-                        // Keep audio and media_type from our explicit args
-                        parsed
-                            .provider_options
-                            .inspect(|p| opts.provider_options = Some(p.clone()));
+                        let mut parsed: TranscriptionCallOptions =
+                            parse_wire_json("opts_json", s)?;
+                        // Required data comes from the explicit arguments; all
+                        // operation policy and optional fields come from JSON.
+                        parsed.audio = opts.audio;
+                        parsed.media_type = opts.media_type;
+                        opts = parsed;
                     }
                 }
                 opts.abort_signal = bridge.map(|b| b.core_signal());
 
-                let result = self
-                    .inner
-                    .do_generate(&opts)
+                let result =
+                    aimux_core::transcription_model::transcribe(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -222,16 +217,17 @@ impl RerankingModel {
                 let mut opts = RerankingCallOptions::new(query, docs);
                 if let Some(s) = opts_json.as_deref() {
                     if !s.trim().is_empty() && s.trim() != "null" {
-                        let parsed: RerankingCallOptions = parse_wire_json("opts_json", s)?;
-                        opts.provider_options = parsed.provider_options;
-                        opts.top_n = parsed.top_n;
+                        let mut parsed: RerankingCallOptions = parse_wire_json("opts_json", s)?;
+                        // Required data comes from the explicit arguments; all
+                        // operation policy and optional fields come from JSON.
+                        parsed.query = opts.query;
+                        parsed.documents = opts.documents;
+                        opts = parsed;
                     }
                 }
                 opts.abort_signal = bridge.map(|b| b.core_signal());
 
-                let result = self
-                    .inner
-                    .do_rerank(&opts)
+                let result = aimux_core::reranking_model::rerank(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -267,9 +263,7 @@ impl VideoModel {
                 let mut opts: VideoCallOptions = parse_wire_json("opts_json", &opts_json)?;
                 opts.abort_signal = bridge.map(|b| b.core_signal());
 
-                let result = self
-                    .inner
-                    .do_generate(&opts)
+                let result = aimux_core::video_model::generate_video(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -313,9 +307,7 @@ impl SearchModel {
                 }
                 opts.abort_signal = bridge.map(|b| b.core_signal());
 
-                let result = self
-                    .inner
-                    .do_search(&opts)
+                let result = aimux_core::search_model::search(self.inner.as_ref(), opts)
                     .await
                     .map_err(|e| AiMuxBindingError::from(&e))?;
                 serialize_result(&result)
@@ -662,7 +654,7 @@ pub struct TranscriptionSession {
     audio_tx: std::sync::Mutex>>,
     parts_rx:
         tokio::sync::Mutex>>,
-    token: aimux_core::shared::AbortSignal,
+    token: aimux_core::AbortSignal,
 }
 
 /// Start a streaming transcription session. `opts_json` (optional):
@@ -695,8 +687,8 @@ pub async fn start_transcription_session(
             };
 
             // Effective abort = user bridge OR close token (linked).
-            let token = aimux_core::shared::AbortSignal::new();
-            let effective = aimux_core::shared::AbortSignal::new();
+            let token = aimux_core::AbortSignal::new();
+            let effective = aimux_core::AbortSignal::new();
             let mut sources = vec![token.clone()];
             if let Some(b) = bridge {
                 sources.push(b.core_signal());
@@ -730,7 +722,11 @@ pub async fn start_transcription_session(
                     include_raw_chunks: opts.include_raw_chunks.unwrap_or(false),
                     timeout: opts.timeout,
                 };
-                let result = model.do_stream(options).await;
+                let result = aimux_core::transcription_model::stream_transcribe(
+                    model.as_ref(),
+                    options,
+                )
+                .await;
                 match result {
                     Ok(stream_result) => {
                         use futures::StreamExt;
diff --git a/bindings/node/src/types.ts b/bindings/node/src/types.ts
index c14882fa..a98add7f 100644
--- a/bindings/node/src/types.ts
+++ b/bindings/node/src/types.ts
@@ -88,6 +88,8 @@ export type * from './types/ResponseFormat'
 export type * from './types/ResponseInfo'
 export type * from './types/ResponseMetadata'
 export type * from './types/ResponseRecord'
+export type * from './types/RetryError'
+export type * from './types/RetryErrorReason'
 export type * from './types/Role'
 export type * from './types/RuntimeModel'
 export type * from './types/SearchCallOptions'
@@ -137,6 +139,7 @@ export type * from './types/VideoFile'
 export type * from './types/VideoFileData'
 export type * from './types/VideoFrameImage'
 export type * from './types/VideoFrameType'
+export type * from './types/VideoPollOptions'
 export type * from './types/VideoResponse'
 export type * from './types/VideoResult'
 export type * from './types/Warning'
diff --git a/bindings/node/src/types/AiMuxError.ts b/bindings/node/src/types/AiMuxError.ts
index bd81f0a3..84cd3574 100644
--- a/bindings/node/src/types/AiMuxError.ts
+++ b/bindings/node/src/types/AiMuxError.ts
@@ -1,22 +1,17 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
 import type { ApiCallError } from "./ApiCallError";
+import type { RetryError } from "./RetryError";
 
 /**
  * Unified error type for all aimux operations.
  *
  * Variants are cut by *what the caller does about it*, not by where the
- * failure came from. `ApiCall` carries the [`ApiCallError`] detail inline
- * (unboxed, like async-openai's `ApiError(ApiError)`); the size guard in
- * `error_value_golden_test` keeps the enum under clippy's
- * `result_large_err` threshold.
- *
- * Constructed as a named struct literal everywhere
- * (`ApiCallError { status_code: .., ..Default::default() }`), the same
- * shape as the AI SDK's named-options constructor.
+ * failure came from. `ApiCallError` is boxed only to keep the Rust enum
+ * compact; serde and every binding still observe the same object shape.
  */
-export type AiMuxError = { "ApiCall": ApiCallError } | { "JsonParse": string } | { "InvalidResponseData": string } | { "Tool": string } | { "InvalidArgument": string } | { "InvalidPrompt": string } | { "TokenExpired": string } | { "UnsupportedFunctionality": string } | { "NoSuchModel": { model_id: string, 
+export type AiMuxError = { "ApiCall": ApiCallError } | { "Retry": RetryError } | { "JsonParse": string } | { "InvalidResponseData": string } | { "Tool": string } | { "InvalidArgument": string } | { "InvalidPrompt": string } | { "TokenExpired": string } | { "UnsupportedFunctionality": string } | { "NoSuchModel": { model_id: string,
 /**
  * What kind of model was requested (`"languageModel"`,
  * `"imageModel"`, …), the AI SDK's `modelType`.
  */
-model_type: string, } } | { "NoSuchProvider": { provider_id: string, } } | { "Timeout": string } | "Aborted" | { "Other": string };
+model_type: string, } } | { "NoSuchProvider": { provider_id: string, } } | { "Timeout": string } | { "Aborted": string } | { "Other": string };
diff --git a/bindings/node/src/types/ApiCallError.ts b/bindings/node/src/types/ApiCallError.ts
index dfb24830..0deee628 100644
--- a/bindings/node/src/types/ApiCallError.ts
+++ b/bindings/node/src/types/ApiCallError.ts
@@ -1,4 +1,5 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
  * What a failed API call observed — the field set of the AI SDK's
@@ -12,51 +13,58 @@
  * body read) are `ApiCall` errors too — no response arrived, so
  * `status_code` is `None` and `is_retryable` is `true`, exactly as the AI
  * SDK's `handleFetchError` builds an `APICallError` with no `statusCode`
- * and `isRetryable: true`. `message` holds the provider's text **only** —
- * the status is a field, never baked into the string; `Display` composes
- * the human-readable form from the fields at print time, so nothing
- * downstream has to parse it back out.
+ * and `isRetryable: true`. `message` holds the provider's text verbatim when
+ * one is available; transport and parsing failures use a library message
+ * with the source detail appended. The HTTP status is always a field, never
+ * baked into the string; `Display` composes the human-readable form from the
+ * fields at print time, so nothing downstream has to parse it back out.
  *
- * `APICallError` fields not carried in this round: `url` /
- * `requestBodyValues` / `responseHeaders` (the request context is not
- * available at the error-construction sites today; all fields are
- * `#[serde(default)]`, so adding them later is not a breaking change).
+ * Every producer must provide the sanitized request URL and request values.
+ * This keeps transport and response failures self-contained and matches the
+ * AI SDK's `APICallError` contract.
  */
-export type ApiCallError = { 
+export type ApiCallError = {
+/**
+ * Sanitized request URL. Required for every API-derived failure.
+ */
+url: string,
+/**
+ * Sanitized request values used to create the API request.
+ */
+request_body_values: JsonValue,
 /**
  * HTTP status of the response, when it came from one
  * (`APICallError.statusCode`). Always the *observed* status: the HTTP
  * layer fills it for every response-derived error; errors built without
  * an HTTP exchange leave it `None`.
  */
-status_code: number | null, 
+status_code: number | null,
 /**
  * Provider's machine-readable code (OpenAI's `code`/`type`,
  * e.g. `"rate_limit_exceeded"` vs `"insufficient_quota"`). Never an HTTP
  * status. Our normalized take on `APICallError.data`.
  */
-provider_code: string | null, 
+provider_code: string | null,
 /**
- * The provider's message, verbatim. No status prefix.
+ * Human-readable failure text. Provider text is verbatim when available;
+ * locally detected transport/parse failures include their source detail.
+ * Never includes an HTTP status prefix.
  */
-message: string, 
+message: string,
 /**
  * The raw response body, verbatim (`APICallError.responseBody`) — the
  * lossless evidence when `message`/`provider_code` are extractions.
  * `None` when the error did not come from a response body.
  */
-response_body: string | null, 
+response_body: string | null,
 /**
- * Provider-assigned request id, from the `x-request-id` / `request-id`
- * response header when one was sent. Filled by the shared HTTP layer.
+ * Sanitized response headers.
  */
-request_id: string | null, 
+response_headers: { [key in string]: string } | null,
 /**
- * Retry hint in milliseconds, distilled from the `retry-after-ms` /
- * `retry-after` response headers on a 429. The classification lives in
- * `status_code`; this is the matching action input.
+ * Parsed provider error data.
  */
-retry_after_ms: number | null, 
+data: JsonValue | null,
 /**
  * Whether retrying can help (`APICallError.isRetryable`) — stored at
  * construction, exactly like the AI SDK: the response path computes it
diff --git a/bindings/node/src/types/CallOptions.ts b/bindings/node/src/types/CallOptions.ts
index 5e9056a6..9085e1f8 100644
--- a/bindings/node/src/types/CallOptions.ts
+++ b/bindings/node/src/types/CallOptions.ts
@@ -14,84 +14,84 @@ import type { JsonValue } from "./serde_json/JsonValue";
  * `GenerateTextOptions` (user-facing) which is converted to `CallOptions`
  * by the `generate_text` / `stream_text` functions.
  */
-export type CallOptions = { 
+export type CallOptions = {
 /**
  * The standardized prompt (message array). Required.
  */
-prompt: Array, 
+prompt: Array,
 /**
  * Maximum tokens to generate.
  */
-max_output_tokens: number | null, 
+max_output_tokens: number | null,
 /**
  * Sampling temperature.
  */
-temperature: number | null, 
+temperature: number | null,
 /**
  * Stop sequences.
  */
-stop_sequences: Array | null, 
+stop_sequences: Array | null,
 /**
  * Nucleus sampling `top_p`.
  */
-top_p: number | null, 
+top_p: number | null,
 /**
  * Top-k sampling.
  */
-top_k: number | null, 
+top_k: number | null,
 /**
  * Presence penalty.
  */
-presence_penalty: number | null, 
+presence_penalty: number | null,
 /**
  * Frequency penalty.
  */
-frequency_penalty: number | null, 
+frequency_penalty: number | null,
 /**
  * Response format (text or JSON).
  */
-response_format: ResponseFormat | null, 
+response_format: ResponseFormat | null,
 /**
  * Seed for reproducibility.
  */
-seed: number | null, 
+seed: number | null,
 /**
  * Tools available to the model (function tools and/or provider-defined tools).
  */
-tools: Array | null, 
+tools: Array | null,
 /**
  * How the model should choose tools.
  */
-tool_choice: ToolChoice, 
+tool_choice: ToolChoice,
 /**
  * Extra HTTP headers.
  */
-headers: { [key in string]: string } | null, 
+headers: { [key in string]: string } | null,
 /**
  * Provider-specific options (keyed by provider name).
  */
-provider_options: { [key in string]: JsonValue } | null, 
+provider_options: { [key in string]: JsonValue } | null,
 /**
  * Top-level reasoning effort. Maps to OpenAI `reasoning_effort` and
  * Anthropic `thinking` config.
  */
-reasoning: ReasoningEffort | null, 
+reasoning: ReasoningEffort | null,
 /**
  * Per-call request body overrides. Deep-merged into the provider-built
  * request body (after any built-in vendor override) before sending.
  * `null` values delete the corresponding key. See RFC-0017.
  */
-body_overrides: JsonValue | null, 
+body_overrides: JsonValue | null,
 /**
  * Per-call retry count override. `None` uses the provider's configured
- * `RetryConfig.max_retries`. `Some(0)` disables retries.
+ * Core operation retry. `Some(0)` disables retries.
  */
-max_retries: number | null, 
+max_retries: number | null,
 /**
  * Per-call timeout configuration (total / first-chunk / chunk idle).
  * `None` = no timeouts (provider defaults still apply at the HTTP layer).
  */
-timeout: TimeoutConfiguration | null, 
+timeout: TimeoutConfiguration | null,
 /**
  * Session identifier, for grouping consecutive calls into a session
  * (observability, see RFC-0024). Explicit values take precedence; when
@@ -100,7 +100,7 @@ timeout: TimeoutConfiguration | null,
  * local grouping, while session headers in `headers` are for upstream
  * routing — both may share an id value but travel different paths.
  */
-session_id: string | null, 
+session_id: string | null,
 /**
  * Emit raw provider stream chunks as `StreamPart::Raw` (debugging aid).
  * When `Some(true)`, streaming providers yield one `Raw` part per JSON
diff --git a/bindings/node/src/types/EmbeddingCallOptions.ts b/bindings/node/src/types/EmbeddingCallOptions.ts
index c78a3476..a3c424e8 100644
--- a/bindings/node/src/types/EmbeddingCallOptions.ts
+++ b/bindings/node/src/types/EmbeddingCallOptions.ts
@@ -1,4 +1,5 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
@@ -6,15 +7,23 @@ import type { JsonValue } from "./serde_json/JsonValue";
  *
  * Aligned with V4 `EmbeddingModelV4CallOptions`.
  */
-export type EmbeddingCallOptions = { 
+export type EmbeddingCallOptions = {
 /**
  * List of text values to generate embeddings for.
  */
-values: Array, 
+values: Array,
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null,
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null,
 /**
  * Additional provider-specific options, keyed by provider name.
  */
-provider_options: { [key in string]: JsonValue } | null, 
+provider_options: { [key in string]: JsonValue } | null,
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/HttpExchange.ts b/bindings/node/src/types/HttpExchange.ts
index 2d87f56c..112f3275 100644
--- a/bindings/node/src/types/HttpExchange.ts
+++ b/bindings/node/src/types/HttpExchange.ts
@@ -6,15 +6,26 @@ import type { TimingRecord } from "./TimingRecord";
 /**
  * ③ HTTP 侧:单次 attempt 的 wire 交换。
  */
-export type HttpExchange = { 
+export type HttpExchange = {
 /**
- * 第几次重试(0=首次);per-attempt 递增。
+ * Composite step this exchange belongs to (e.g. `router[0]:openai/gpt-4o`
+ * or `moa.ref[1]:...`). `None` for a plain, non-composite operation.
  */
-attempt: number, request: HttpRecord, 
+step?: string | null,
+/**
+ * Core operation attempt, starting at 1. Attempt numbers are unique
+ * across the whole call, including composite child steps, so
+ * `(attempt, exchange_index)` alone still identifies an exchange.
+ */
+attempt: number,
+/**
+ * HTTP exchange within the operation attempt, starting at 1.
+ */
+exchange_index: number, request: HttpRecord,
 /**
  * None = 请求失败未获响应。
  */
-response: ResponseRecord | null, timing: TimingRecord, error: string | null, 
+response: ResponseRecord | null, timing: TimingRecord, error: string | null,
 /**
  * 流式:该 exchange 是否已终结(收到 response 补全)。非流式恒 true。
  */
diff --git a/bindings/node/src/types/ImageCallOptions.ts b/bindings/node/src/types/ImageCallOptions.ts
index d32de995..6ccac7f4 100644
--- a/bindings/node/src/types/ImageCallOptions.ts
+++ b/bindings/node/src/types/ImageCallOptions.ts
@@ -2,6 +2,7 @@
 import type { AspectRatio } from "./AspectRatio";
 import type { ImageFile } from "./ImageFile";
 import type { Size } from "./Size";
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
@@ -9,42 +10,50 @@ import type { JsonValue } from "./serde_json/JsonValue";
  *
  * Aligned with V4 `ImageModelV4CallOptions`.
  */
-export type ImageCallOptions = { 
+export type ImageCallOptions = {
 /**
  * Prompt for the image generation. `None` for operations (e.g. upscaling)
  * that do not require a prompt.
  */
-prompt: string | null, 
+prompt: string | null,
 /**
  * Number of images to generate.
  */
-n: number, 
+n: number,
 /**
  * Size of the images, in `{width}x{height}` format.
  * `None` uses the provider's default size.
  */
-size: Size | null, 
+size: Size | null,
 /**
  * Aspect ratio of the images, in `{width}:{height}` format.
  * `None` uses the provider's default aspect ratio.
  */
-aspect_ratio: AspectRatio | null, 
+aspect_ratio: AspectRatio | null,
 /**
  * Seed for deterministic generation. `None` uses the provider's default.
  */
-seed: number | null, 
+seed: number | null,
 /**
  * Images for image editing or variation generation.
  */
-files: Array | null, 
+files: Array | null,
 /**
  * Mask image for inpainting operations.
  */
-mask: ImageFile | null, 
+mask: ImageFile | null,
 /**
  * Additional provider-specific options, keyed by provider name.
  */
-provider_options: { [key in string]: JsonValue }, 
+provider_options: { [key in string]: JsonValue },
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null,
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null,
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/OutcomeRecord.ts b/bindings/node/src/types/OutcomeRecord.ts
index f76f2ac9..b9fc7941 100644
--- a/bindings/node/src/types/OutcomeRecord.ts
+++ b/bindings/node/src/types/OutcomeRecord.ts
@@ -5,7 +5,12 @@ import type { JsonValue } from "./serde_json/JsonValue";
 /**
  * 最终结果摘要。
  */
-export type OutcomeRecord = { status: OutcomeStatus, finish_reason: string | null, error: string | null, 
+export type OutcomeRecord = { status: OutcomeStatus, finish_reason: string | null, error: string | null,
+/**
+ * Lossless structured domain error. In particular, `RetryError.errors`
+ * keeps the complete attempt history rather than only its display text.
+ */
+error_value: JsonValue | null,
 /**
  * 序列化的 Usage。
  */
diff --git a/bindings/node/src/types/RerankingCallOptions.ts b/bindings/node/src/types/RerankingCallOptions.ts
index 3de78c9a..2d4e93de 100644
--- a/bindings/node/src/types/RerankingCallOptions.ts
+++ b/bindings/node/src/types/RerankingCallOptions.ts
@@ -1,5 +1,6 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
 import type { RerankingDocuments } from "./RerankingDocuments";
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
@@ -7,23 +8,31 @@ import type { JsonValue } from "./serde_json/JsonValue";
  *
  * Aligned with V4 `RerankingModelV4CallOptions`.
  */
-export type RerankingCallOptions = { 
+export type RerankingCallOptions = {
 /**
  * Documents to rerank.
  */
-documents: RerankingDocuments, 
+documents: RerankingDocuments,
 /**
  * The query to rerank the documents against.
  */
-query: string, 
+query: string,
 /**
  * Optional limit: return only the top `n` documents.
  */
-top_n: number | null, 
+top_n: number | null,
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null,
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null,
 /**
  * Additional provider-specific options, keyed by provider name.
  */
-provider_options: { [key in string]: JsonValue } | null, 
+provider_options: { [key in string]: JsonValue } | null,
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/RetryError.ts b/bindings/node/src/types/RetryError.ts
new file mode 100644
index 00000000..0a76855d
--- /dev/null
+++ b/bindings/node/src/types/RetryError.ts
@@ -0,0 +1,8 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { AiMuxError } from "./AiMuxError";
+import type { RetryErrorReason } from "./RetryErrorReason";
+
+/**
+ * Complete error history for a retried model operation.
+ */
+export type RetryError = { reason: RetryErrorReason, errors: Array, };
diff --git a/bindings/node/src/types/RetryErrorReason.ts b/bindings/node/src/types/RetryErrorReason.ts
new file mode 100644
index 00000000..ce092dcf
--- /dev/null
+++ b/bindings/node/src/types/RetryErrorReason.ts
@@ -0,0 +1,6 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+/**
+ * Why the retry wrapper stopped retrying.
+ */
+export type RetryErrorReason = "maxRetriesExceeded" | "errorNotRetryable";
diff --git a/bindings/node/src/types/SearchCallOptions.ts b/bindings/node/src/types/SearchCallOptions.ts
index e11420fe..4d544c2c 100644
--- a/bindings/node/src/types/SearchCallOptions.ts
+++ b/bindings/node/src/types/SearchCallOptions.ts
@@ -1,42 +1,51 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
  * Options passed to [`SearchModel::do_search`].
  */
-export type SearchCallOptions = { 
+export type SearchCallOptions = {
 /**
  * The search query string.
  */
-query: string, 
+query: string,
 /**
  * Maximum number of results to return.
  */
-max_results: number | null, 
+max_results: number | null,
 /**
  * Whether to include raw page content in results.
  *
  * Provider support varies; providers that cannot honor this should
  * issue a warning rather than erroring.
  */
-include_raw_content: boolean | null, 
+include_raw_content: boolean | null,
 /**
  * Optional time range filter (e.g. `"day"`, `"week"`, `"month"`,
  * `"year"`). Provider support varies.
  */
-time_range: string | null, 
+time_range: string | null,
 /**
  * Optional list of domains to include in results.
  */
-include_domains: Array | null, 
+include_domains: Array | null,
 /**
  * Optional list of domains to exclude from results.
  */
-exclude_domains: Array | null, 
+exclude_domains: Array | null,
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null,
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null,
 /**
  * Additional provider-specific options, keyed by provider name.
  */
-provider_options: { [key in string]: JsonValue } | null, 
+provider_options: { [key in string]: JsonValue } | null,
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/SpeechCallOptions.ts b/bindings/node/src/types/SpeechCallOptions.ts
index 39cae2f4..c1131c41 100644
--- a/bindings/node/src/types/SpeechCallOptions.ts
+++ b/bindings/node/src/types/SpeechCallOptions.ts
@@ -1,4 +1,5 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
@@ -6,38 +7,46 @@ import type { JsonValue } from "./serde_json/JsonValue";
  *
  * Aligned with V4 `SpeechModelV4CallOptions`.
  */
-export type SpeechCallOptions = { 
+export type SpeechCallOptions = {
 /**
  * Text to convert to speech.
  */
-text: string, 
+text: string,
 /**
  * The voice to use (provider-specific voice ID, name, or other
  * identifier).
  */
-voice: string | null, 
+voice: string | null,
 /**
  * The desired output format for the audio, e.g. `"mp3"`, `"wav"`.
  */
-output_format: string | null, 
+output_format: string | null,
 /**
  * Instructions for the speech generation, e.g.
  * `"Speak in a slow and steady tone"`.
  */
-instructions: string | null, 
+instructions: string | null,
 /**
  * The speed of the speech generation.
  */
-speed: number | null, 
+speed: number | null,
 /**
  * The language for speech generation (ISO 639-1 code, e.g. `"en"`, or
  * `"auto"` for automatic detection). Provider support varies.
  */
-language: string | null, 
+language: string | null,
 /**
  * Additional provider-specific options, keyed by provider name.
  */
-provider_options: { [key in string]: JsonValue } | null, 
+provider_options: { [key in string]: JsonValue } | null,
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null,
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null,
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/TimeoutConfiguration.ts b/bindings/node/src/types/TimeoutConfiguration.ts
index 8e3a3444..58d90ce0 100644
--- a/bindings/node/src/types/TimeoutConfiguration.ts
+++ b/bindings/node/src/types/TimeoutConfiguration.ts
@@ -8,16 +8,25 @@
  * corresponding limit. A `total` timeout also covers retry backoff and the
  * whole streamed response.
  */
-export type TimeoutConfiguration = { 
+export type TimeoutConfiguration = {
 /**
  * Overall timeout for the entire call (including retries and, for
  * streaming, the whole stream), in milliseconds.
  */
-total_ms: number | null, 
+total_ms: number | null,
+/**
+ * Timeout for one generation step, including that step's attempts and
+ * retry backoff, in milliseconds. Aimux currently has one step.
+ */
+step_ms: number | null,
 /**
  * Timeout waiting for the first stream chunk (streaming only).
+ *
+ * Counted from operation start, so it also bounds stream establishment
+ * and any retries before the first semantic output: it is the
+ * user-perceived time-to-first-output budget, not a per-attempt timer.
  */
-first_chunk_ms: number | null, 
+first_chunk_ms: number | null,
 /**
  * Maximum idle time between consecutive stream chunks (streaming only).
  */
diff --git a/bindings/node/src/types/TranscriptionCallOptions.ts b/bindings/node/src/types/TranscriptionCallOptions.ts
index 79acc81e..d43f937a 100644
--- a/bindings/node/src/types/TranscriptionCallOptions.ts
+++ b/bindings/node/src/types/TranscriptionCallOptions.ts
@@ -1,5 +1,6 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
 import type { AudioInput } from "./AudioInput";
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
@@ -20,6 +21,14 @@ media_type: string,
  * Additional provider-specific options, keyed by provider name.
  */
 provider_options: { [key in string]: JsonValue } | null, 
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null, 
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null, 
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/VideoCallOptions.ts b/bindings/node/src/types/VideoCallOptions.ts
index ecdc9f26..6f8e9057 100644
--- a/bindings/node/src/types/VideoCallOptions.ts
+++ b/bindings/node/src/types/VideoCallOptions.ts
@@ -1,65 +1,80 @@
 // This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
 import type { AspectRatio } from "./AspectRatio";
 import type { Size } from "./Size";
+import type { TimeoutConfiguration } from "./TimeoutConfiguration";
 import type { VideoFile } from "./VideoFile";
 import type { VideoFrameImage } from "./VideoFrameImage";
+import type { VideoPollOptions } from "./VideoPollOptions";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
- * Options passed to [`VideoModel::do_generate`].
+ * Options passed to [`VideoModel::do_start`] and [`VideoModel::do_status`].
  *
  * Aligned with V4 `VideoModelV4CallOptions`.
  */
-export type VideoCallOptions = { 
+export type VideoCallOptions = {
 /**
  * Text prompt for the video generation. `None` when not required.
  */
-prompt: string | null, 
+prompt: string | null,
 /**
  * Number of videos to generate. Default `1`; most models only support
  * `n = 1` due to computational cost.
  */
-n: number, 
+n: number,
 /**
  * Aspect ratio, in `{width}:{height}` format (e.g. `"16:9"`).
  */
-aspect_ratio: AspectRatio | null, 
+aspect_ratio: AspectRatio | null,
 /**
  * Resolution, in `{width}x{height}` format (e.g. `"1280x720"`).
  */
-resolution: Size | null, 
+resolution: Size | null,
 /**
  * Duration of the video in seconds. Typically 3–10 seconds.
  */
-duration: number | null, 
+duration: number | null,
 /**
  * Frames per second. Common values: 24, 30, 60.
  */
-fps: number | null, 
+fps: number | null,
 /**
  * Seed for deterministic generation. `None` uses a random seed.
  */
-seed: number | null, 
+seed: number | null,
 /**
  * Input image for image-to-video generation (the starting frame).
  */
-image: VideoFile | null, 
+image: VideoFile | null,
 /**
  * Role-tagged image inputs for first-last-frame generation.
  */
-frame_images: Array | null, 
+frame_images: Array | null,
 /**
  * Reference inputs for reference-to-video generation (images or videos).
  */
-input_references: Array | null, 
+input_references: Array | null,
 /**
  * Whether the model should generate audio alongside the video.
  */
-generate_audio: boolean | null, 
+generate_audio: boolean | null,
 /**
  * Additional provider-specific options, keyed by provider name.
  */
-provider_options: { [key in string]: JsonValue }, 
+provider_options: { [key in string]: JsonValue },
+/**
+ * Per-call retry override. `None` uses the model default.
+ */
+max_retries: number | null,
+/**
+ * Per-call poll pacing override for the start/status flow. Unset fields
+ * fall back to the model's [`VideoModel::poll_config`].
+ */
+poll: VideoPollOptions | null,
+/**
+ * Per-call operation timeout.
+ */
+timeout: TimeoutConfiguration | null,
 /**
  * Additional HTTP headers to send with the request.
  */
diff --git a/bindings/node/src/types/VideoPollOptions.ts b/bindings/node/src/types/VideoPollOptions.ts
new file mode 100644
index 00000000..5950d03e
--- /dev/null
+++ b/bindings/node/src/types/VideoPollOptions.ts
@@ -0,0 +1,14 @@
+// This file was generated by [ts-rs](https://github.com/Aleph-Alpha/ts-rs). Do not edit this file manually.
+
+/**
+ * Per-call poll pacing for video generation (AI SDK `generateVideo` `poll`).
+ */
+export type VideoPollOptions = { 
+/**
+ * Delay between consecutive status checks, in milliseconds.
+ */
+interval_ms: number | null, 
+/**
+ * Maximum total time to wait for completion, in milliseconds.
+ */
+timeout_ms: number | null, };
diff --git a/bindings/node/src/types/VideoResult.ts b/bindings/node/src/types/VideoResult.ts
index 9bda831e..08e0c838 100644
--- a/bindings/node/src/types/VideoResult.ts
+++ b/bindings/node/src/types/VideoResult.ts
@@ -5,23 +5,23 @@ import type { Warning } from "./Warning";
 import type { JsonValue } from "./serde_json/JsonValue";
 
 /**
- * The result of [`VideoModel::do_generate`].
+ * The final result of a video generation operation.
  *
  * Aligned with V4 `VideoModelV4Result`.
  */
-export type VideoResult = { 
+export type VideoResult = {
 /**
  * Generated videos as URLs, base64 strings, or binary data.
  */
-videos: Array, 
+videos: Array,
 /**
  * Warnings for the call, e.g. unsupported features.
  */
-warnings: Array, 
+warnings: Array,
 /**
  * Additional provider-specific metadata, keyed by provider name.
  */
-provider_metadata: { [key in string]: JsonValue } | null, 
+provider_metadata: { [key in string]: JsonValue } | null,
 /**
  * Response information for telemetry and debugging.
  */
diff --git a/bindings/python/python/aimux/__init__.py b/bindings/python/python/aimux/__init__.py
index 6575b7ea..a0a676c8 100644
--- a/bindings/python/python/aimux/__init__.py
+++ b/bindings/python/python/aimux/__init__.py
@@ -10,6 +10,7 @@
 from .aimux import (
     AimuxError,
     APICallError,
+    RetryError,
     JSONParseError,
     InvalidResponseDataError,
     ToolError,
@@ -80,6 +81,7 @@
 __all__ = [
     "AimuxError",
     "APICallError",
+    "RetryError",
     "JSONParseError",
     "InvalidResponseDataError",
     "ToolError",
diff --git a/bindings/python/python/aimux/wrapper.py b/bindings/python/python/aimux/wrapper.py
index e0d1b816..6fbb09a1 100644
--- a/bindings/python/python/aimux/wrapper.py
+++ b/bindings/python/python/aimux/wrapper.py
@@ -817,6 +817,7 @@ class TimeoutConfiguration(BaseModel):
     """
 
     total_ms: Optional[int] = None
+    step_ms: Optional[int] = None
     first_chunk_ms: Optional[int] = None
     chunk_ms: Optional[int] = None
 
diff --git a/bindings/python/src/error.rs b/bindings/python/src/error.rs
index 0e256152..e2ae7ac2 100644
--- a/bindings/python/src/error.rs
+++ b/bindings/python/src/error.rs
@@ -23,13 +23,14 @@ use aimux_core::{AiMuxError, recording::RecordingError as CoreRecordingError};
 use pyo3::create_exception;
 use pyo3::exceptions::{PyException, PyRuntimeError, PyValueError};
 use pyo3::prelude::*;
-use pyo3::types::PyAnyMethods;
+use pyo3::types::{PyAnyMethods, PyList};
 
 // Base — catch-all for AiMuxError failures only. Recording and binding failures
 // have independent public exception types below.
 create_exception!(aimux, AimuxError, PyException, "AiMux failure");
 
 create_exception!(aimux, APICallError, AimuxError, "API call failure");
+create_exception!(aimux, RetryError, AimuxError, "Operation retry failure");
 create_exception!(
     aimux,
     JSONParseError,
@@ -203,15 +204,19 @@ fn recording_py_err(e: &CoreRecordingError) -> PyErr {
 }
 
 pub(crate) fn to_py_err(e: &AiMuxError) -> PyErr {
-    Python::with_gil(|py| match raise_variant(py, e) {
-        Ok(err) => err,
+    Python::with_gil(|py| match exception_instance(py, e) {
+        Ok(instance) => PyErr::from_value_bound(instance),
         Err(e) => e,
     })
 }
 
-fn raise_variant(py: Python<'_>, e: &AiMuxError) -> PyResult {
+fn exception_instance<'py>(
+    py: Python<'py>,
+    e: &AiMuxError,
+) -> PyResult> {
     let typ = match e {
         AiMuxError::ApiCall(_) => py.get_type_bound::(),
+        AiMuxError::Retry(_) => py.get_type_bound::(),
         AiMuxError::JsonParse(_) => py.get_type_bound::(),
         AiMuxError::InvalidResponseData(_) => py.get_type_bound::(),
         AiMuxError::Tool(_) => py.get_type_bound::(),
@@ -224,7 +229,7 @@ fn raise_variant(py: Python<'_>, e: &AiMuxError) -> PyResult {
         AiMuxError::NoSuchModel { .. } => py.get_type_bound::(),
         AiMuxError::NoSuchProvider { .. } => py.get_type_bound::(),
         AiMuxError::Timeout(_) => py.get_type_bound::(),
-        AiMuxError::Aborted => py.get_type_bound::(),
+        AiMuxError::Aborted(_) => py.get_type_bound::(),
         AiMuxError::Other(_) => py.get_type_bound::(),
     };
     let inst = typ.call1((e.to_string(),))?;
@@ -245,7 +250,42 @@ fn raise_variant(py: Python<'_>, e: &AiMuxError) -> PyResult {
                 (!d.message.is_empty()).then_some(d.message.as_str()),
             )?;
             inst.setattr("response_body", d.response_body.as_deref())?;
-            inst.setattr("request_id", d.request_id.as_deref())?;
+            inst.setattr("url", (!d.url.is_empty()).then_some(d.url.as_str()))?;
+            // JSON null projects to Python None, so no separate absent case.
+            inst.setattr(
+                "request_body_values",
+                json_value_to_python(py, &d.request_body_values)?,
+            )?;
+            inst.setattr("response_headers", d.response_headers.clone())?;
+            inst.setattr(
+                "data",
+                match &d.data {
+                    Some(data) => Some(json_value_to_python(py, data)?),
+                    None => None,
+                },
+            )?;
+        }
+        AiMuxError::Retry(retry) => {
+            // Attempt history, oldest first; each entry is itself a full
+            // exception instance (recursing through this same projection).
+            let errors = PyList::empty_bound(py);
+            for error in &retry.errors {
+                errors.append(exception_instance(py, error)?)?;
+            }
+            inst.setattr(
+                "reason",
+                match retry.reason {
+                    aimux_core::RetryErrorReason::MaxRetriesExceeded => "maxRetriesExceeded",
+                    aimux_core::RetryErrorReason::ErrorNotRetryable => "errorNotRetryable",
+                },
+            )?;
+            // A deserialized history may be empty; last_error is None then.
+            let last_error = match errors.len() {
+                0 => None,
+                n => Some(errors.get_item(n - 1)?),
+            };
+            inst.setattr("errors", &errors)?;
+            inst.setattr("last_error", last_error)?;
         }
         AiMuxError::TokenExpired(_) => {
             inst.setattr("status", 401)?;
@@ -265,7 +305,17 @@ fn raise_variant(py: Python<'_>, e: &AiMuxError) -> PyResult {
         }
         _ => {}
     }
-    Ok(PyErr::from_value_bound(inst))
+    Ok(inst)
+}
+
+fn json_value_to_python<'py>(
+    py: Python<'py>,
+    value: &serde_json::Value,
+) -> PyResult> {
+    py.import_bound("json")?.call_method1(
+        "loads",
+        (serde_json::to_string(value).expect("serde_json::Value is always serializable"),),
+    )
 }
 
 /// Register the exception hierarchy on the Python module.
@@ -273,6 +323,7 @@ pub(crate) fn register(m: &Bound<'_, PyModule>) -> PyResult<()> {
     let py = m.py();
     m.add("AimuxError", py.get_type_bound::())?;
     m.add("APICallError", py.get_type_bound::())?;
+    m.add("RetryError", py.get_type_bound::())?;
     m.add("JSONParseError", py.get_type_bound::())?;
     m.add(
         "InvalidResponseDataError",
diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs
index 49f3e539..b2f7af66 100644
--- a/bindings/python/src/lib.rs
+++ b/bindings/python/src/lib.rs
@@ -233,6 +233,21 @@ impl Model {
                                     break;
                                 }
                             }
+                            Err(e) if e.is_recoverable_stream_error() => {
+                                // Core keeps the stream alive across a malformed frame;
+                                // deliver it as a StreamPart::Error data item and keep
+                                // pumping.
+                                match serde_json::to_string(
+                                    &aimux_core::stream_part::StreamPart::Error { error: e },
+                                ) {
+                                    Ok(json) => {
+                                        if tx.send(Ok(json)).await.is_err() {
+                                            break;
+                                        }
+                                    }
+                                    Err(_) => break,
+                                }
+                            }
                             Err(e) => {
                                 let _ = tx.send(Err(e.into())).await;
                                 break;
@@ -337,6 +352,11 @@ impl Model {
                                     break;
                                 }
                             }
+                            Err(e) if e.is_recoverable_stream_error() => {
+                                // Consumers type this path as ChatCompletionChunk; the
+                                // error cannot ride it — skip and keep pumping (full
+                                // fidelity lives on the StreamPart path).
+                            }
                             Err(e) => {
                                 let _ = tx.send(Err(e.into())).await;
                                 break;
diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs
index 3e8df1a8..3ca843c8 100644
--- a/bindings/python/src/multimodal.rs
+++ b/bindings/python/src/multimodal.rs
@@ -54,7 +54,9 @@ impl EmbeddingModel {
         let values: Vec = wire_json("values_json", values_json)?;
         opts.values = values;
 
-        let result = crate::runtime().block_on(async move { self.inner.do_embed(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::embedding_model::embed(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -79,7 +81,9 @@ impl SpeechModel {
     pub fn generate(&self, opts_json: &str) -> PyResult {
         let opts: SpeechCallOptions = wire_json("opts_json", opts_json)?;
 
-        let result = crate::runtime().block_on(async move { self.inner.do_generate(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::speech_model::generate_speech(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -104,7 +108,9 @@ impl ImageModel {
     pub fn generate(&self, opts_json: &str) -> PyResult {
         let opts: ImageCallOptions = wire_json("opts_json", opts_json)?;
 
-        let result = crate::runtime().block_on(async move { self.inner.do_generate(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::image_model::generate_image(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -140,15 +146,18 @@ impl TranscriptionModel {
         );
         if let Some(s) = opts_json {
             if !s.trim().is_empty() && s.trim() != "null" {
-                let parsed: TranscriptionCallOptions = wire_json("opts_json", s)?;
-                // Keep audio and media_type from our explicit args.
-                if let Some(p) = &parsed.provider_options {
-                    opts.provider_options = Some(p.clone());
-                }
+                // Take every caller option; audio and media_type always come
+                // from the explicit args.
+                let mut parsed: TranscriptionCallOptions = wire_json("opts_json", s)?;
+                parsed.audio = opts.audio;
+                parsed.media_type = opts.media_type;
+                opts = parsed;
             }
         }
 
-        let result = crate::runtime().block_on(async move { self.inner.do_generate(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::transcription_model::transcribe(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -184,13 +193,18 @@ impl RerankingModel {
         let mut opts = RerankingCallOptions::new(query.to_string(), docs);
         if let Some(s) = opts_json {
             if !s.trim().is_empty() && s.trim() != "null" {
-                let parsed: RerankingCallOptions = wire_json("opts_json", s)?;
-                opts.provider_options = parsed.provider_options;
-                opts.top_n = parsed.top_n;
+                // Take every caller option; query and documents always come
+                // from the explicit args.
+                let mut parsed: RerankingCallOptions = wire_json("opts_json", s)?;
+                parsed.query = opts.query;
+                parsed.documents = opts.documents;
+                opts = parsed;
             }
         }
 
-        let result = crate::runtime().block_on(async move { self.inner.do_rerank(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::reranking_model::rerank(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -215,7 +229,9 @@ impl VideoModel {
     pub fn generate(&self, opts_json: &str) -> PyResult {
         let opts: VideoCallOptions = wire_json("opts_json", opts_json)?;
 
-        let result = crate::runtime().block_on(async move { self.inner.do_generate(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::video_model::generate_video(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -247,7 +263,9 @@ impl SearchModel {
             }
         }
 
-        let result = crate::runtime().block_on(async move { self.inner.do_search(&opts).await });
+        let result = crate::runtime().block_on(async move {
+            aimux_core::search_model::search(self.inner.as_ref(), opts).await
+        });
 
         match result {
             Ok(r) => serialize_result(&r),
@@ -528,7 +546,7 @@ pub struct TranscriptionSession {
     parts_rx: tokio::sync::Mutex<
         tokio::sync::mpsc::Receiver>,
     >,
-    token: aimux_core::shared::AbortSignal,
+    token: aimux_core::AbortSignal,
 }
 
 /// Start a streaming transcription session. `opts_json` (optional):
@@ -554,8 +572,8 @@ pub fn start_transcription_session(
         _ => SessionOpts::default(),
     };
 
-    let token = aimux_core::shared::AbortSignal::new();
-    let effective = aimux_core::shared::AbortSignal::new();
+    let token = aimux_core::AbortSignal::new();
+    let effective = aimux_core::AbortSignal::new();
     {
         let linked = effective.clone();
         let source = token.clone();
@@ -585,7 +603,8 @@ pub fn start_transcription_session(
             include_raw_chunks: opts.include_raw_chunks.unwrap_or(false),
             timeout: opts.timeout,
         };
-        let result = model.do_stream(options).await;
+        let result = aimux_core::transcription_model::stream_transcribe(model.as_ref(), options)
+            .await;
         match result {
             Ok(stream_result) => {
                 use futures::StreamExt;
diff --git a/bindings/python/tests/test_error.py b/bindings/python/tests/test_error.py
index 4b0a2c9a..cc6adc58 100644
--- a/bindings/python/tests/test_error.py
+++ b/bindings/python/tests/test_error.py
@@ -23,7 +23,10 @@
     "provider_code",
     "provider_message",
     "response_body",
-    "request_id",
+    "url",
+    "request_body_values",
+    "response_headers",
+    "data",
 )
 
 _ERROR_BODY = json.dumps(
@@ -37,7 +40,7 @@ def _free_port():
         return s.getsockname()[1]
 
 
-def _rate_limited_server(port):
+def _rate_limited_server(port, retry_after_ms="1500"):
     """Always answers 429 with a provider error body plus retry/request headers."""
 
     class Handler(BaseHTTPRequestHandler):
@@ -48,7 +51,7 @@ def do_POST(self):
             self.send_header("content-type", "application/json")
             self.send_header("content-length", str(len(resp)))
             self.send_header("x-request-id", "req_abc123")
-            self.send_header("retry-after-ms", "1500")
+            self.send_header("retry-after-ms", retry_after_ms)
             self.end_headers()
             self.wfile.write(resp)
 
@@ -58,6 +61,16 @@ def log_message(self, *args):
     HTTPServer(("127.0.0.1", port), Handler).serve_forever()
 
 
+def _wait_for_port(port):
+    for _ in range(50):
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.connect(("127.0.0.1", port))
+            break
+        except ConnectionRefusedError:
+            time.sleep(0.05)
+
+
 def test_no_such_provider_carries_only_its_payload_and_no_json():
     from aimux import AimuxError, NoSuchProviderError, provider
 
@@ -82,17 +95,12 @@ def test_api_call_error_carries_every_field_the_response_produced():
     proc = Process(target=_rate_limited_server, args=(port,))
     proc.start()
     try:
-        for _ in range(50):
-            try:
-                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-                    s.connect(("127.0.0.1", port))
-                break
-            except ConnectionRefusedError:
-                time.sleep(0.05)
+        _wait_for_port(port)
 
         model = openai("test-key", "gpt-4o", f"http://127.0.0.1:{port}")
+        # 429 is retryable; disable retries to observe the bare call failure.
         with pytest.raises(APICallError) as excinfo:
-            generate_text(model, "hi")
+            generate_text(model, "hi", {"max_retries": 0})
         err = excinfo.value
 
         assert type(err) is APICallError
@@ -100,7 +108,11 @@ def test_api_call_error_carries_every_field_the_response_produced():
         assert err.retry_ms == 1500
         assert err.retryable is True
         assert err.provider_code == "rate_limit_exceeded"
-        assert err.request_id == "req_abc123"
+        # No distinct request-id field: the raw header evidence carries it.
+        assert err.response_headers["x-request-id"] == "req_abc123"
+        assert err.response_headers["retry-after-ms"] == "1500"
+        assert f"127.0.0.1:{port}" in err.url
+        assert err.request_body_values["model"] == "gpt-4o"
         assert err.response_body == _ERROR_BODY
         # The provider's text on its own; str(e) is the composed form.
         assert err.provider_message == "slow down"
@@ -111,6 +123,36 @@ def test_api_call_error_carries_every_field_the_response_produced():
         proc.join(timeout=2)
 
 
+def test_retry_error_carries_reason_and_typed_attempt_history():
+    from aimux import APICallError, RetryError, generate_text, openai
+
+    port = _free_port()
+    # retry-after-ms: 1 keeps the honored retry delay negligible.
+    proc = Process(target=_rate_limited_server, args=(port, "1"))
+    proc.start()
+    try:
+        _wait_for_port(port)
+
+        model = openai("test-key", "gpt-4o", f"http://127.0.0.1:{port}")
+        with pytest.raises(RetryError) as excinfo:
+            generate_text(model, "hi", {"max_retries": 1})
+        err = excinfo.value
+
+        assert type(err) is RetryError
+        assert err.reason == "maxRetriesExceeded"
+        # Attempt history, oldest first; each entry is a full typed exception.
+        assert len(err.errors) == 2
+        assert err.last_error is err.errors[-1]
+        for attempt in err.errors:
+            assert type(attempt) is APICallError
+            assert attempt.status == 429
+            assert attempt.provider_code == "rate_limit_exceeded"
+        assert "Failed after 2 attempts" in str(err)
+    finally:
+        proc.terminate()
+        proc.join(timeout=2)
+
+
 def test_malformed_wire_json_is_a_value_error_not_a_core_error():
     from aimux import AimuxError, openai
 
diff --git a/bindings/swift/Sources/Aimux/Aimux.swift b/bindings/swift/Sources/Aimux/Aimux.swift
index e9adb03f..7f8d75e6 100644
--- a/bindings/swift/Sources/Aimux/Aimux.swift
+++ b/bindings/swift/Sources/Aimux/Aimux.swift
@@ -11,7 +11,7 @@ import Foundation
 //
 // Every fallible C function returns `aimux_error_t *` (`OpaquePointer?`):
 // NULL = success (result in the trailing out-param), non-NULL = failure. The
-// unified code is AiMuxError (1...13), RecordingError (100...105), or a C ABI
+// unified code is AiMuxError (1...14), RecordingError (100...105), or a C ABI
 // failure (200...206). The three `expect*` decoders copy the relevant fields, release
 // it with `aimux_error_free` (exactly once) and return the Swift error
 // to throw. Errors are not handles: never `aimux_drop_handle` one.
@@ -44,7 +44,7 @@ func expectFfiError(_ e: OpaquePointer, context: String) -> any Error {
     return invariant("aimux ffi: \(context): \(message)")
 }
 
-/// Decode a returned error from an `[AiMuxError]` call: 1...13 becomes
+/// Decode a returned error from an `[AiMuxError]` call: 1...14 become
 /// `AimuxError`; 200...206 is decoded by `expectFfiError`. Frees `e` once.
 func expectAimuxError(_ e: OpaquePointer, context: String) -> any Error {
     let code = aimux_error_code(e)
@@ -108,9 +108,18 @@ func ffiStringCall(
 // Errors
 // ─────────────────────────────────────────────────────────────────────────────
 
+/// Why a `.retry` failure stopped retrying. Raw values are the core's serde
+/// camelCase wire names (`aimux_error_retry_reason`).
+public enum RetryErrorReason: String, Equatable, Sendable {
+    /// Every permitted attempt failed with a retryable error.
+    case maxRetriesExceeded
+    /// A later attempt failed with an error not worth retrying.
+    case errorNotRetryable
+}
+
 /// Structured aimux failure type (Swift `Error`).
 ///
-/// Maps 1:1 from the 13 core `AiMuxError` variants. Every HTTP-shaped failure
+/// Maps 1:1 from the 14 core `AiMuxError` variants. Every HTTP-shaped failure
 /// is `.apiCall` (`AIMUX_E_API_CALL`). Only aimux-core produces these: a
 /// binding-local failure (raw JSON that does not parse, a typed value that
 /// fails to encode, library output that fails to decode) surfaces as the
@@ -146,8 +155,21 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl
     /// ever observed — a missing API key, an error built without a request, or
     /// a transport failure; read `retryable` to tell those apart, `status`
     /// cannot.
-    case apiCall(message: String, status: Int, retryMs: Int64, retryable: Bool, providerCode: String? = nil, providerMessage: String? = nil, requestId: String? = nil, responseBody: String? = nil)
+    case apiCall(
+        message: String, status: Int, retryMs: Int64, retryable: Bool,
+        providerCode: String? = nil, providerMessage: String? = nil,
+        responseBody: String? = nil, url: String? = nil,
+        requestBodyValues: JSONValue? = nil, responseHeaders: [String: String]? = nil,
+        providerData: JSONValue? = nil
+    )
+    /// Retrying stopped (`AIMUX_E_RETRY`): `reason` says why, `errors` is the
+    /// per-attempt history (oldest first; `errors.last` is the final attempt,
+    /// each element any `AimuxError`, `.apiCall` with the full detail set
+    /// included), `message` the composed summary ("Failed after N attempts…").
+    case retry(message: String, reason: RetryErrorReason, errors: [AimuxError])
     case timeout(message: String, status: Int, retryMs: Int64, retryable: Bool)
+    /// `message` is the abort payload verbatim ("request aborted" for signal
+    /// aborts).
     case aborted(message: String, status: Int, retryMs: Int64, retryable: Bool)
     case other(message: String, status: Int, retryMs: Int64, retryable: Bool)
 
@@ -165,11 +187,15 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl
              .unsupportedFunctionality(let m, let s, let r, let t),
              .noSuchModel(let m, let s, let r, let t, _, _),
              .noSuchProvider(let m, let s, let r, let t, _),
-             .apiCall(let m, let s, let r, let t, _, _, _, _),
+             .apiCall(let m, let s, let r, let t, _, _, _, _, _, _, _),
              .timeout(let m, let s, let r, let t),
              .aborted(let m, let s, let r, let t),
              .other(let m, let s, let r, let t):
             return (m, s, r, t)
+        case .retry(let m, _, _):
+            // Retrying already stopped: no single HTTP payload, and by
+            // definition not retryable.
+            return (m, -1, -1, false)
         }
     }
 
@@ -187,6 +213,7 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl
         case .noSuchModel: c = AIMUX_E_NO_SUCH_MODEL
         case .noSuchProvider: c = AIMUX_E_NO_SUCH_PROVIDER
         case .apiCall: c = AIMUX_E_API_CALL
+        case .retry: c = AIMUX_E_RETRY
         case .timeout: c = AIMUX_E_TIMEOUT
         case .aborted: c = AIMUX_E_ABORTED
         case .other: c = AIMUX_E_OTHER
@@ -216,28 +243,63 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl
 
     /// `.apiCall` only: the provider's own error code (e.g. `"insufficient_quota"`).
     public var providerCode: String? {
-        if case .apiCall(_, _, _, _, let v, _, _, _) = self { return v }
+        if case .apiCall(_, _, _, _, let v, _, _, _, _, _, _) = self { return v }
         return nil
     }
 
     /// `.apiCall` only: the failure's own text without the composed prefix `message` carries (e.g. `"slow down"`).
     public var providerMessage: String? {
-        if case .apiCall(_, _, _, _, _, let v, _, _) = self { return v }
+        if case .apiCall(_, _, _, _, _, let v, _, _, _, _, _) = self { return v }
         return nil
     }
 
-    /// `.apiCall` only: the provider request id, for support tickets.
-    public var requestId: String? {
-        if case .apiCall(_, _, _, _, _, _, let v, _) = self { return v }
+    /// `.apiCall` only: the raw response body.
+    public var responseBody: String? {
+        if case .apiCall(_, _, _, _, _, _, let v, _, _, _, _) = self { return v }
         return nil
     }
 
-    /// `.apiCall` only: the raw response body.
-    public var responseBody: String? {
-        if case .apiCall(_, _, _, _, _, _, _, let v) = self { return v }
+    /// `.apiCall` only: the sanitized request URL.
+    public var url: String? {
+        if case .apiCall(_, _, _, _, _, _, _, let v, _, _, _) = self { return v }
+        return nil
+    }
+
+    /// `.apiCall` only: the sanitized request body values (any JSON type).
+    public var requestBodyValues: JSONValue? {
+        if case .apiCall(_, _, _, _, _, _, _, _, let v, _, _) = self { return v }
         return nil
     }
 
+    /// `.apiCall` only: the sanitized response headers (includes the
+    /// retry-after evidence `retryMs` is derived from).
+    public var responseHeaders: [String: String]? {
+        if case .apiCall(_, _, _, _, _, _, _, _, _, let v, _) = self { return v }
+        return nil
+    }
+
+    /// `.apiCall` only: the provider's parsed error data (the AI SDK's
+    /// `APICallError.data`).
+    public var providerData: JSONValue? {
+        if case .apiCall(_, _, _, _, _, _, _, _, _, _, let v) = self { return v }
+        return nil
+    }
+
+    /// `.retry` only: why retrying stopped.
+    public var retryReason: RetryErrorReason? {
+        if case .retry(_, let v, _) = self { return v }
+        return nil
+    }
+
+    /// `.retry` only: the per-attempt history, oldest first.
+    public var retryErrors: [AimuxError]? {
+        if case .retry(_, _, let v) = self { return v }
+        return nil
+    }
+
+    /// `.retry` only: the final attempt's error (`retryErrors` last element).
+    public var lastError: AimuxError? { retryErrors?.last }
+
     /// `.noSuchModel` only: the model id that was asked for.
     public var modelId: String? {
         if case .noSuchModel(_, _, _, _, let v, _) = self { return v }
@@ -308,8 +370,24 @@ public enum AimuxError: Error, LocalizedError, CustomStringConvertible, Equatabl
             return .apiCall(message: message, status: status, retryMs: retryMs, retryable: retryable,
                             providerCode: takeCString(aimux_error_provider_code(h)),
                             providerMessage: takeCString(aimux_error_provider_message(h)),
-                            requestId: takeCString(aimux_error_request_id(h)),
-                            responseBody: takeCString(aimux_error_response_body(h)))
+                            responseBody: takeCString(aimux_error_response_body(h)),
+                            url: takeCString(aimux_error_url(h)),
+                            requestBodyValues: takeJSONValue(aimux_error_request_body_values(h)),
+                            responseHeaders: takeHeaderMap(aimux_error_response_headers(h)),
+                            providerData: takeJSONValue(aimux_error_provider_data(h)))
+        case AIMUX_E_RETRY:
+            let reason = takeCString(aimux_error_retry_reason(h))
+                .flatMap(RetryErrorReason.init(rawValue:)) ?? .maxRetriesExceeded
+            var errors: [AimuxError] = []
+            for i in 0..?) -> JSONValue? {
+        guard let s = takeCString(p) else { return nil }
+        return try? JSONDecoder().decode(JSONValue.self, from: Data(s.utf8))
+    }
+
+    /// Decode a getter-owned JSON object of string→string pairs (the
+    /// response-headers shape), freeing the C allocation.
+    private static func takeHeaderMap(_ p: UnsafeMutablePointer?) -> [String: String]? {
+        guard let s = takeCString(p) else { return nil }
+        return try? JSONDecoder().decode([String: String].self, from: Data(s.utf8))
+    }
 }
 
 // ─────────────────────────────────────────────────────────────────────────────
diff --git a/bindings/swift/Sources/Aimux/Types.swift b/bindings/swift/Sources/Aimux/Types.swift
index b4821116..5937c4f5 100644
--- a/bindings/swift/Sources/Aimux/Types.swift
+++ b/bindings/swift/Sources/Aimux/Types.swift
@@ -27,7 +27,7 @@ import Foundation
 /// On this toolchain `Bool` and `Double` decode are mutually exclusive (a JSON
 /// bool fails `Double` decode and a JSON number fails `Bool` decode), so the
 /// scalar ordering below is unambiguous.
-public enum JSONValue: Codable, Equatable {
+public enum JSONValue: Codable, Equatable, Sendable {
     case null
     case bool(Bool)
     case number(Double)
@@ -1076,17 +1076,21 @@ public struct StreamTextResultAggregated: Codable, Equatable {
 /// backoff and the whole streamed response.
 public struct TimeoutConfiguration: Codable, Equatable {
     public var totalMs: UInt64?
+    public var stepMs: UInt64?
     public var firstChunkMs: UInt64?
     public var chunkMs: UInt64?
 
     enum CodingKeys: String, CodingKey {
         case totalMs = "total_ms"
+        case stepMs = "step_ms"
         case firstChunkMs = "first_chunk_ms"
         case chunkMs = "chunk_ms"
     }
 
-    public init(totalMs: UInt64? = nil, firstChunkMs: UInt64? = nil, chunkMs: UInt64? = nil) {
-        self.totalMs = totalMs; self.firstChunkMs = firstChunkMs; self.chunkMs = chunkMs
+    public init(totalMs: UInt64? = nil, stepMs: UInt64? = nil,
+                firstChunkMs: UInt64? = nil, chunkMs: UInt64? = nil) {
+        self.totalMs = totalMs; self.stepMs = stepMs
+        self.firstChunkMs = firstChunkMs; self.chunkMs = chunkMs
     }
 }
 
diff --git a/contract-tests/fixtures/wire-format.json b/contract-tests/fixtures/wire-format.json
index 3746b12c..d0c8c32c 100644
--- a/contract-tests/fixtures/wire-format.json
+++ b/contract-tests/fixtures/wire-format.json
@@ -128,7 +128,7 @@
   {
     "name": "timeout_configuration_values",
     "type": "TimeoutConfiguration",
-    "json": "{\"total_ms\":5000,\"first_chunk_ms\":1000,\"chunk_ms\":500}",
+    "json": "{\"total_ms\":5000,\"step_ms\":null,\"first_chunk_ms\":1000,\"chunk_ms\":500}",
     "description": "TimeoutConfiguration with all three limits set (wire shape lock, RFC-0016 H3)"
   },
   {
diff --git a/scripts/gen_ts_types.py b/scripts/gen_ts_types.py
index febd19a1..26a0da90 100755
--- a/scripts/gen_ts_types.py
+++ b/scripts/gen_ts_types.py
@@ -34,6 +34,11 @@
 EXCLUDED = {"ProviderName.ts"}
 
 
+def normalize_generated(content: str) -> str:
+    """Ignore ts-rs's incidental end-of-line spaces in generated output."""
+    return "\n".join(line.rstrip() for line in content.splitlines()) + "\n"
+
+
 def regenerate(export_dir: Path) -> None:
     """Run the ts-rs export tests, writing the .ts files into export_dir.
 
@@ -54,7 +59,7 @@ def regenerate(export_dir: Path) -> None:
 def manifest(root: Path, exclude: set[str] = frozenset()) -> dict[str, str]:
     """Map POSIX-style relative path -> content for every .ts under root."""
     return {
-        p.relative_to(root).as_posix(): p.read_text(encoding="utf-8")
+        p.relative_to(root).as_posix(): normalize_generated(p.read_text(encoding="utf-8"))
         for p in root.rglob("*.ts")
         if p.name not in exclude
     }
@@ -80,7 +85,10 @@ def sync(expected: dict[str, str]) -> None:
     for rel, content in expected.items():
         path = TYPES_DIR / rel
         path.parent.mkdir(parents=True, exist_ok=True)
-        path.write_text(content, encoding="utf-8")
+        # `content` is canonicalized by manifest(); compare the raw file so a
+        # normal regeneration also removes ts-rs's incidental EOL spaces.
+        if not path.exists() or path.read_text(encoding="utf-8") != content:
+            path.write_text(content, encoding="utf-8")
 
 
 def main() -> int:
diff --git a/tools/aimux-web/examples/gen_demo_fixtures.rs b/tools/aimux-web/examples/gen_demo_fixtures.rs
index a76151e3..499fdfcf 100644
--- a/tools/aimux-web/examples/gen_demo_fixtures.rs
+++ b/tools/aimux-web/examples/gen_demo_fixtures.rs
@@ -67,7 +67,9 @@ fn recording(
             provider_options: None,
         },
         exchanges: vec![HttpExchange {
+            step: None,
             attempt: 0,
+            exchange_index: 0,
             request: HttpRecord {
                 method: "POST".into(),
                 url: "https://api.example.com/v1/chat/completions".to_string(),
@@ -96,6 +98,7 @@ fn recording(
             } else {
                 None
             },
+            error_value: None,
             usage: Some(usage),
         },
         complete: true,
diff --git a/tools/aimux-web/examples/gen_fixture.rs b/tools/aimux-web/examples/gen_fixture.rs
index 7feeb39d..533367c6 100644
--- a/tools/aimux-web/examples/gen_fixture.rs
+++ b/tools/aimux-web/examples/gen_fixture.rs
@@ -56,7 +56,9 @@ fn main() {
             provider_options: None,
         },
         exchanges: vec![HttpExchange {
+            step: None,
             attempt: 0,
+            exchange_index: 0,
             request: HttpRecord {
                 method: "POST".into(),
                 url: "https://api.openai.com/v1/chat/completions".into(),
@@ -81,6 +83,7 @@ fn main() {
             status: OutcomeStatus::Success,
             finish_reason: Some("stop".into()),
             error: None,
+            error_value: None,
             usage: Some(serde_json::json!({
                 "input_tokens": { "total": 12, "cache_read": 0 },
                 "output_tokens": { "total": 5 }
diff --git a/tools/aimux-web/src/api/mod.rs b/tools/aimux-web/src/api/mod.rs
index fb2683b7..b8beed57 100644
--- a/tools/aimux-web/src/api/mod.rs
+++ b/tools/aimux-web/src/api/mod.rs
@@ -62,7 +62,7 @@ pub fn err_response(e: AiMuxError) -> Response {
         | AiMuxError::NoSuchProvider { .. }
         | AiMuxError::NoSuchModel { .. } => StatusCode::BAD_REQUEST,
         AiMuxError::UnsupportedFunctionality(_) => StatusCode::NOT_IMPLEMENTED,
-        AiMuxError::Timeout(_) | AiMuxError::Aborted => StatusCode::GATEWAY_TIMEOUT,
+        AiMuxError::Timeout(_) | AiMuxError::Aborted(_) => StatusCode::GATEWAY_TIMEOUT,
         _ => StatusCode::INTERNAL_SERVER_ERROR,
     };
     (

From 5ddfd0f57f51985ea02d5eaa1d74892c3e28aa0a Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Mon, 24 Aug 2026 22:29:11 +0800
Subject: [PATCH 06/24] docs: ship the AI SDK request-pipeline spec
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The design spec lives at docs/ai-sdk-request-pipeline.md (moved out of
rfc/), with a difference table for every deliberate deviation from the
AI SDK baseline (63db193) and the §13 acceptance matrix. API docs across
all bindings pick up the retry/timeout error context, the video
start/status split with its idempotency gate, and the recoverable-frame
stream semantics; download/SSRF security is explicitly handed to the
PR that owns it.
---
 README.md                       |    4 +-
 docs/API.md                     |    4 +-
 docs/PROJECT-OVERVIEW.md        |    4 +-
 docs/README.md                  |    1 +
 docs/ai-sdk-request-pipeline.md | 1037 +++++++++++++++++++++++++++++++
 docs/api/c.md                   |   54 +-
 docs/api/flutter.md             |   12 +-
 docs/api/go.md                  |   17 +-
 docs/api/java.md                |   25 +-
 docs/api/kotlin.md              |   17 +-
 docs/api/node.md                |   19 +-
 docs/api/python.md              |   15 +-
 docs/api/reference.md           |    4 +-
 docs/api/rust.md                |   42 +-
 docs/api/swift.md               |    5 +-
 docs/error-model.md             |    6 +-
 rfc/0009-request-resilience.md  |    6 +
 rfc/0016-align-with-aisdk.md    |    7 +
 18 files changed, 1203 insertions(+), 76 deletions(-)
 create mode 100644 docs/ai-sdk-request-pipeline.md

diff --git a/README.md b/README.md
index c7bd386b..731a120f 100644
--- a/README.md
+++ b/README.md
@@ -98,7 +98,7 @@ aimux/
 ├── aimux-core            # Core abstractions: LanguageModel / Provider / Message / StreamPart
 ├── aimux-providers       # 329 provider implementations (251 registry-backed + native)
 ├── aimux-stream          # SSE / NDJSON stream parsing
-├── aimux-provider-utils  # HTTP utilities: retry, backoff, error parsing, API-key loading
+├── aimux-provider-utils  # One-exchange HTTP helpers, response handlers, API-key loading
 ├── aimux-ffi             # C ABI (opaque handles + JSON results + owned aimux_error_t *) for non-native bindings
 └── tools/                # aimux-cli (cache probe) · aimux-replay · aimux-web (console)
 ```
@@ -122,7 +122,7 @@ cargo add aimux-core aimux-providers
 | `aimux-core` | Core abstractions: `LanguageModel` / `Provider` / `Message` / `StreamPart` | [crates.io](https://crates.io/crates/aimux-core) |
 | `aimux-providers` | 325 provider implementations | [crates.io](https://crates.io/crates/aimux-providers) |
 | `aimux-stream` | SSE / NDJSON stream parsing | [crates.io](https://crates.io/crates/aimux-stream) |
-| `aimux-provider-utils` | HTTP utilities: retry, backoff, error parsing | [crates.io](https://crates.io/crates/aimux-provider-utils) |
+| `aimux-provider-utils` | One-exchange HTTP helpers and typed response handlers | [crates.io](https://crates.io/crates/aimux-provider-utils) |
 | `aimux-ffi` | C ABI for non-native bindings | [crates.io](https://crates.io/crates/aimux-ffi) |
 
 **Node.js**:
diff --git a/docs/API.md b/docs/API.md
index 5088b161..2d100395 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -111,7 +111,7 @@ Examples: [Node.js](api/node.md#text-generation) · [Python](api/python.md#text-
 | `instructions` | `string?` | System instructions |
 | `reasoning` | `ReasoningEffort?` | Reasoning effort |
 | `max_retries` | `number?` | Per-call retry override; `0` disables retries (`None` = provider default, 2) |
-| `timeout` | `TimeoutConfiguration?` | Per-call timeouts (total / first-chunk / chunk idle) — see [Timeouts](#timeouts) |
+| `timeout` | `TimeoutConfiguration?` | Per-call timeouts (total / step / first-chunk / chunk idle) — see [Timeouts](#timeouts) |
 | `body_overrides` | `object?` | Per-call request-body overrides, deep-merged; `null` values delete keys |
 | `headers` | `object?` | Extra HTTP headers |
 
@@ -213,7 +213,7 @@ pre-aborted signal fails fast without sending).
   never crosses the JSON boundary):
 
   ```rust
-  let signal = aimux_core::shared::AbortSignal::new();
+  let signal = aimux_core::AbortSignal::new();
   let opts = GenerateTextOptions { abort_signal: Some(signal.clone()), ..Default::default() };
   let task = tokio::spawn(generate_text(&model, "Explain Rust.", opts));
   signal.abort(); // cancels the call
diff --git a/docs/PROJECT-OVERVIEW.md b/docs/PROJECT-OVERVIEW.md
index 0b041d03..6f227c29 100644
--- a/docs/PROJECT-OVERVIEW.md
+++ b/docs/PROJECT-OVERVIEW.md
@@ -163,7 +163,7 @@ aimux/
 │   ├── 251 OpenAI compatible #   registry-backed: provider-registry.json + provider(name, ...) entry (RFC-0017 phase 4)
 │   └── modalities/search     #   voice / image / video / search implementations
 ├── aimux-stream             # SSE / NDJSON streaming parsing
-├── aimux-provider-utils     # HTTP utilities: retry, backoff, error parsing, API key loading
+├── aimux-provider-utils     # One-exchange HTTP helpers, response handlers, API-key loading
 ├── aimux-ffi                # C ABI (FFI infrastructure, shared by all bindings)
 └── bindings/                # 6 language bindings
     ├── node/                #   napi-rs v3 + typed TS wrapper
@@ -234,7 +234,7 @@ Covers scenarios such as tool calling, multi-turn dialogue, reasoning/thinking,
 
 - `shared_client()` connection pool sharing (providers no longer each establish their own connections)
 - TLS session reuse
-- Jitter backoff retry (referencing the catcher design)
+- Full Jitter backoff lives in `aimux-core::retry` as the `get_delay_ms` default (server retry hints are honored exactly, not jittered; RFC-0031)
 - Fixed timeout
 
 #### 6. Request-layer decoupling (RFC-0009 supplement)
diff --git a/docs/README.md b/docs/README.md
index 7ded561e..11df171f 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -15,6 +15,7 @@ Public documentation for aimux — a unified LLM access layer written in Rust.
 | [PROJECT-OVERVIEW.md](PROJECT-OVERVIEW.md) | Project overview, design decisions, and benchmark summary |
 | [PERF-RESULTS.md](PERF-RESULTS.md) | Performance benchmark results (aimux vs OpenAI SDK / Vercel AI SDK) |
 | [aimux-vs-aisdk-node.md](aimux-vs-aisdk-node.md) | Node.js developer-experience comparison vs Vercel AI SDK |
+| [ai-sdk-request-pipeline.md](ai-sdk-request-pipeline.md) | AI SDK-aligned request pipeline — operation retry, response handlers, timeout/abort design |
 
 For the project README, quickstart, and provider/binding tables, see the
 [top-level README](../README.md).
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
new file mode 100644
index 00000000..7311d726
--- /dev/null
+++ b/docs/ai-sdk-request-pipeline.md
@@ -0,0 +1,1037 @@
+# RFC-0031: 对齐 AI SDK 的请求管线
+
+> **Status**: IMPLEMENTED（2026-08-20；jitter 与首条 SSE error 两处有意差异见 §3）
+>
+> **Date**: 2026-08-19
+>
+> **Reference baseline**: Vercel AI SDK `63db193`
+>
+> **Scope**: operation retry、timeout/abort、POST/GET API helpers、response handlers、`ApiCallError` 与 `RetryError`
+
+---
+
+## 1. 摘要
+
+aimux 当前把 HTTP、retry、timeout、body 读取和错误解析集中在：
+
+```text
+send
+send_timed
+send_stream_timed
+send_with_retry_raw
+ErrorStructure
+parse_provider_error
+```
+
+结果是 retry 只包围 HTTP exchange，无法包围完整的 Provider operation；成功响应解析、
+Provider 业务错误和 semantic stream timeout 也处在错误的层级。
+
+本 RFC 不重新发明请求框架，直接采用 AI SDK 的现有分层和函数名称：
+
+```text
+Core user operation
+  └─ prepare_retries
+       └─ retry_with_exponential_backoff_respecting_retry_headers
+            └─ model.do_generate / do_stream / do_embed / ...
+                 └─ Provider
+                      └─ post_json_to_api / post_form_data_to_api /
+                         post_to_api / get_from_api
+                           ├─ failed_response_handler
+                           └─ successful_response_handler
+```
+
+核心不变量：
+
+> 一个用户级 model operation 只有一个通用 retry owner；一次 Provider Utils API helper
+> 调用只执行一次 fetch attempt，不做 retry。自动 redirect 属于这次 fetch attempt，
+> recording 仍把整条 redirect chain 记为一个 logical exchange。
+
+---
+
+## 2. 规范来源
+
+以下 AI SDK 文件是本 RFC 对应部分的规范来源：
+
+| 能力 | AI SDK 文件 |
+|---|---|
+| Core retry 包围 `model.doGenerate` | `packages/ai/src/generate-text/generate-text.ts` |
+| stream 建立阶段 retry 与 semantic timeout | `packages/ai/src/generate-text/stream-text.ts` |
+| `prepareRetries` | `packages/ai/src/util/prepare-retries.ts` |
+| APICall-aware retry | `packages/ai/src/util/retry-with-exponential-backoff.ts` |
+| 通用 retry primitive | `packages/provider-utils/src/retry-with-exponential-backoff.ts` |
+| `RetryError` | `packages/ai/src/util/retry-error.ts` |
+| POST helpers | `packages/provider-utils/src/post-to-api.ts` |
+| GET helper | `packages/provider-utils/src/get-from-api.ts` |
+| response handlers | `packages/provider-utils/src/response-handler.ts` |
+| fetch error normalization | `packages/provider-utils/src/handle-fetch-error.ts` |
+| `APICallError` | `packages/provider/src/errors/api-call-error.ts` |
+| timeout configuration | `packages/ai/src/prompt/request-options.ts` |
+| abort/timeout merge | `packages/ai/src/util/merge-abort-signals.ts`、`set-abort-timeout.ts` |
+
+除 §3 明列的 Rust/Aimux 差异外，实现应遵循这些文件的函数边界和行为，不得用新的
+`send_*`、`execute_api`、`call_to_api` 等公开抽象替代。
+
+---
+
+## 3. Aimux 仅有的适配差异
+
+| 差异 | 决策 |
+|---|---|
+| crate 依赖方向 | retry primitive 和 Core wrapper 都放进 `aimux-core::retry`，避免 `aimux-core` 反向依赖 `aimux-provider-utils`；函数边界和算法不变 |
+| Rust 错误类型 | AI SDK `RetryError.errors: unknown[]` 对应 `Vec`；这是同一错误历史的强类型表达 |
+| enum 尺寸 | `AiMuxError::ApiCall(Box)`；`Box` 只解决 Rust enum 尺寸，不改变序列化或 binding 语义 |
+| `cause` | 本轮不加入；AI SDK 用 `cause` 保存的底层错误追加到 message（`{upstream message}: {source}`），或保存在 data/response body，不得静默丢失 |
+| Gateway | 不增加 `GatewayError`；aimux 没有独立 AI Gateway 错误体系 |
+| 既有错误字段 | 保留 Aimux 的 `provider_code`；删除旧路径派生的 `request_id` / `retry_after_ms`，retry hint 的唯一事实源是 `response_headers` |
+| error context 安全 | AI SDK 传原始 request context；Aimux 写入 public error 前使用统一白名单/脱敏/大小限制，binary 与大型 data URL 只保留摘要 |
+| 错误 body 上限 | AI SDK 全量读取错误 body（仅受 2 GiB 防 OOM 上限，超限直接抛 `DownloadError` 替换原错误）；Aimux 的 `ApiCallError` 会跨 FFI 序列化并写入 recording，因此错误 body 采用 best-effort 截断读：public `response_body` 上限 64 KiB（lossy 解码后按字符边界执行，带 `…(truncated)` 标记），解析上限 1 MiB 保证超大但合法的错误 JSON 仍能进 provider mapper，读取中途连接死亡保留已收到的部分 |
+| provider 默认 retry | 保留既有 `RetryConfig` 源码与行为兼容；Core 读取 model config，per-call 只覆盖 `max_retries` |
+| jitter | AI SDK 默认不抖动但留 `getDelayInMs` 注入点；Aimux 用同一注入点默认注入 RFC-0009 Full Jitter，只作用于 exponential delay，server hint 精确遵守（§6.5） |
+| 第一条 SSE error | AI SDK OpenAI provider 会扫描到首个 semantic output；Aimux 保留 RFC-0016 的更窄 first-event peek，使立即到达的 error 成为 retry 边界内的 attempt 失败（§8.3） |
+| 默认超时 | 与 AI SDK 一致无默认 `total_ms`；Aimux 有意保留非流式 exchange 的 30s whole-response 上限，但从 shared-client 全局配置下移到单次 exchange，流式 exchange 豁免（§5.4） |
+| timeout 输入形状 | 保留跨语言已有的 object/struct 形式，不增加 AI SDK 的裸 number 简写；字段语义一致 |
+| `AbortSignal` | 保留既有的 `CancellationToken` 薄包装，只表示调用方取消；Rust 可直接 drop future，因此 timeout 由 Core deadline/select 表达，不照搬 JS 的 signal merge（§8.0） |
+| SSE framing | AI SDK 的 event-source handler 直接产出解析后的事件；Aimux 复用 `aimux_stream::SseStream` 做同一件事，framing 不留在 Provider |
+| tool timeout | aimux 目前不执行用户 tool，本 RFC 不增加 `tool_ms`/per-tool timeout |
+| observability | 每次 exchange 继续走 Aimux recording/tracing；这不改变 helper 的单次请求语义 |
+| FFI 运输 | 不由本 RFC 重新设计；C 错误 owner 与 nested error getter 由 RFC-0030 单独规定 |
+
+这些是适配，不是另一套架构。
+
+---
+
+## 4. 分层
+
+### 4.1 `aimux-core`
+
+负责：
+
+- 用户级 `generate_text` / `stream_text` / multimodal operations；
+- `prepare_retries` 和两层 exponential-backoff 函数；
+- `RetryError`；
+- 用户 abort、total/step/stream timeout；
+- semantic stream timeout。
+
+固定模块：
+
+```text
+aimux-core/src/abort_signal.rs   AbortSignal（实现从 shared.rs 迁出；旧 `shared::AbortSignal` 路径 re-export 保留）
+aimux-core/src/retry.rs
+aimux-core/src/timeout.rs
+```
+
+### 4.2 `aimux-provider-utils`
+
+负责：
+
+- shared reqwest client、pool、proxy；
+- 单次 POST/GET exchange；
+- `ResponseHandler` 与标准 handler factories；
+- body size limit、header extraction、fetch error normalization；
+- HTTP exchange recording。
+
+固定模块：
+
+```text
+aimux-provider-utils/src/post_to_api.rs
+aimux-provider-utils/src/get_from_api.rs
+aimux-provider-utils/src/response_handler.rs
+aimux-provider-utils/src/handle_fetch_error.rs
+aimux-provider-utils/src/extract_response_headers.rs
+aimux-provider-utils/src/read_response_with_size_limit.rs
+```
+
+`http.rs` 只保留 client/pool/proxy、HTTP value types 和私有的单次 reqwest primitive。
+
+### 4.3 Provider
+
+负责：
+
+- URL、headers、request body；
+- 为每个 HTTP operation 选择 successful/failed response handler；
+- Provider-specific schema、业务错误和结果转换；
+- SSE 或 custom protocol 到 `StreamPart` 的转换；
+- submit/poll/download 等有状态工作流；
+- 在现有 monolithic `do_generate` SPI 下，对已获得 job id 后的幂等 poll/download exchange 单独 retry。
+
+Provider 不对完整 model operation 执行通用 exponential retry；上述安全 exchange 是
+为避免重新 submit 的有状态例外。
+
+---
+
+## 5. Provider Utils API
+
+### 5.1 固定函数面
+
+| AI SDK | Aimux Rust |
+|---|---|
+| `postJsonToApi` | `post_json_to_api` |
+| `postFormDataToApi` | `post_form_data_to_api` |
+| `postToApi` | `post_to_api` |
+| `getFromApi` | `get_from_api` |
+| `ResponseHandler` | `ResponseHandler` |
+| `createJsonErrorResponseHandler` | `create_json_error_response_handler` |
+| `createJsonResponseHandler` | `create_json_response_handler` |
+| `createEventSourceResponseHandler` | `create_event_source_response_handler` |
+| `createBinaryResponseHandler` | `create_binary_response_handler` |
+| `createStatusCodeErrorResponseHandler` | `create_status_code_error_response_handler` |
+| `handleFetchError` | `handle_fetch_error` |
+| `extractResponseHeaders` | `extract_response_headers` |
+| `readResponseWithSizeLimit` | `read_response_with_size_limit` |
+
+### 5.2 `ResponseHandler`
+
+语义直接对应 AI SDK：
+
+```rust
+pub struct ResponseHandlerInput {
+    pub url: String,
+    pub request_body_values: serde_json::Value,
+    pub response: reqwest::Response,
+    /// Rust 适配：fetch 的 body 天然绑定 AbortSignal，reqwest 的不绑定。handler 里
+    /// 所有 body 读取（size-limited read、SSE 分帧）都必须 `select!` 这个信号。
+    pub abort_signal: Option,
+}
+
+pub struct ResponseHandlerOutput {
+    pub value: T,
+    pub raw_value: Option,
+    pub response_headers: Option>,
+}
+```
+
+具体实现可以使用泛型 async closure，不要求 `Arc>`。
+
+每个 API call 同时传入：
+
+```text
+successful_response_handler
+failed_response_handler
+```
+
+分发固定为：
+
+```text
+transport failure → handle_fetch_error
+2xx               → successful_response_handler
+non-2xx           → failed_response_handler
+```
+
+2xx 解析失败由 successful handler 产生 `ApiCallError`，不得再调用 failed handler。
+handler 返回的 ApiCall/Timeout/Aborted 原样透传；其他 handler failure 按 AI SDK 的
+`Failed to process successful/error response` 规则包装。由于本轮没有 `cause`，原始信息
+必须保存在 message、data 或受限长度的 response body 中。
+
+标准 handler factory 直接移植 AI SDK `response-handler.ts` 的同名行为。Provider-specific
+协议使用 custom `ResponseHandler`；例如 Bedrock binary event-stream 不是 SSE，不得套
+`create_event_source_response_handler`。
+
+各 factory 的固定签名（AI SDK 的 `schema` 参数在 Rust 里就是类型参数）：
+
+```rust
+pub fn create_json_response_handler()
+    -> ResponseHandler;                       // value: T, raw_value: Some(json)
+pub fn create_event_source_response_handler()
+    -> ResponseHandler>>;
+    // 用 aimux_stream::SseStream 分帧；每个 event 的 `data` 解析成 T；`[DONE]` 跳过；
+    // 单条解析失败 yield 一个 Err 项（见 §7.2 stream 行），流不终止，由 provider 决定怎么办
+pub fn create_binary_response_handler() -> ResponseHandler;
+pub fn create_json_error_response_handler(error_to_message: F) -> ResponseHandler
+where F: Fn(&serde_json::Value) -> ProviderErrorParts;  // { message, provider_code }
+pub fn create_status_code_error_response_handler() -> ResponseHandler;
+```
+
+aimux 所有 SSE provider 都按 `data` 里的 `type` 字段分派，没有人读 `event:` 名，所以
+T-typed handler 覆盖全部现有用法；需要同时接受正常 chunk 和 error chunk 的 provider 用
+`#[serde(untagged)] enum Chunk { Ok(..), Error(..) }` 作 T——这正是 AI SDK 里
+`z.union([chunkSchema, errorSchema])` 的 Rust 写法。
+
+`create_json_response_handler` 返回 `Bytes` 再让 Provider 自己 `serde_json::from_slice(..)?`
+是错误实现：schema 失败会变成没有 URL/status/headers 的 `JsonParse`，违反 §7.2。
+Provider 调用 JSON helper 之后不得再出现 `serde_json::from_slice` / `from_str` 解析响应
+主体；需要原始 `Value`（例如 `Usage.raw`）时读 `raw_value`。
+
+失败 handler 的归属与 AI SDK 相同：`@ai-sdk/provider-utils` 只提供 factory，
+`openaiFailedResponseHandler`、`anthropicFailedResponseHandler` 等常量定义在各 provider
+包里。Aimux 同样在 `aimux-providers` 的各 provider 模块内定义
+`_failed_response_handler()`（或同义 `static`），由该 endpoint 的 wire schema
+决定。`aimux-provider-utils` 不得提供 `create_openai_error_response_handler`、
+`create_json_path_error_response_handler` 这类“换个名字的 `ErrorStructure`”，也不得把
+OpenAI 形状的失败 handler 当作所有 Provider 的默认值——fal、ElevenLabs、Cohere 等的错误
+体不是 `{error:{message}}`，套 OpenAI handler 会丢 message。
+
+`ErrorStructure`、`DEFAULT_ERROR_STRUCTURE` 和 `parse_provider_error` 全部删除，不增加另一套
+全局默认错误 schema。handler 是否复用由具体 endpoint 的 wire schema 决定，不能由
+Provider 名称或调用点数量推断。
+
+### 5.3 `post_json_to_api` / `post_form_data_to_api`
+
+两者只负责准备 body/context 后调用 `post_to_api`：
+
+```text
+post_json_to_api:
+  content = JSON bytes
+  values  = structured request body
+
+post_form_data_to_api:
+  content = multipart body
+  values  = field summary; binary 仅保留类型与长度
+```
+
+body 类型由函数签名决定，不在运行时检查：`post_json_to_api` 接 `serde_json::Value`，
+`post_form_data_to_api` 接 `MultipartForm`，`post_to_api` 接 `HttpBody`，`get_from_api`
+没有 body。"`post_json_to_api` 收到非 JSON body 返回 `InvalidArgument`" 这类守卫是签名
+没设计好的补丁，不要出现。其余参数（url、headers、abort_signal、call_id、recording_context）
+沿用既有 `HttpRequest` 去掉 method/body 后的那部分。
+
+handler 接口仍接收当前 request context；当它被存入 `ApiCallError` 时必须通过统一
+redaction helper。不得依赖每个 Provider 自己记得脱敏。
+
+### 5.4 `post_to_api` / `get_from_api`
+
+它们遵循 AI SDK 的单次 fetch attempt 契约：
+
+- 不接受 `RetryConfig`；
+- 不接受 `TimeoutConfiguration`；
+- 不执行 backoff；
+- 不调用自身或另一个 API helper 进行 retry；
+- body read/parse 由选中的 response handler 完成；
+- abort signal 只负责中止本次 exchange。
+
+四个 helper 共用一个私有的单次 exchange primitive（`pub(crate)`，名字不进入公开 API）。
+该 primitive 是既有 `send_with_retry_raw` 中 **去掉 for 循环、backoff 与 `ErrorStructure`
+之后剩下的那部分**，而不是另起一份裸 `client.execute()`：
+
+- 共享 client/pool/proxy 选择；
+- abort-aware 的发送（`tokio::select!` on `abort_signal.cancelled()`）；
+- transport error → `handle_fetch_error`；
+- **recording / tracing**：`record_exchange` / `record_failed_exchange` /
+  `record_transport_closed` 和 `tracing::info!` 的 exchange 行原样保留。`HttpRequest.call_id`、
+  `recording_context` 不能变成未读字段。每次 helper 调用恰好记录一个 logical exchange；
+  redirect chain 仍属于该 exchange，retry 不再出现在这里。`attempt` 与 `exchange_index`
+  都由 §10.2 的 `RecordingContext` 提供。
+
+`get_from_api` 的 `request_body_values` 为 `{}`。provider 返回下载 URL 时所需的 SSRF
+防护（私网/loopback 拒绝、DNS pinning、redirect 逐跳校验和跨 origin credential 清理）
+由独立的 #163 交付；该 PR rebase 到本管线后负责接入对应 Provider call site。
+
+**默认超时的取舍**：Core 与 AI SDK 一样没有默认 `total_ms`，binding 也不得补 operation
+deadline；但 Aimux 恢复旧行为，给每个**非流式** helper exchange 保留固定 30s 的
+whole-response 上限（connect、headers、完整 body/handler parse）。它不放回 shared client，
+因为 shared client 同时服务长生命周期 stream；而是在单次 exchange primitive 周围执行。
+超时产生 no-status、retryable `ApiCallError`，因此仍由 Core 决定是否重试完整 operation，
+不会伪装成 non-retryable 的 Core `Timeout`。streaming handler 豁免这个 30s guard，只受
+connect timeout、caller abort、显式 Core total/step/first/chunk deadline 约束。这是相对 AI SDK
+fetch 无默认 exchange timeout 的有意差异，CHANGELOG 必须明确。
+
+### 5.5 `handle_fetch_error`
+
+与 AI SDK 一致：
+
+| 输入 | 输出 |
+|---|---|
+| Timeout/Aborted | 原样返回 |
+| 已是 `ApiCallError` | 原样返回 |
+| reqwest DNS/connect/TLS/socket/request transport failure | no-status、retryable `ApiCallError` |
+| 其他本地错误 | 原样返回 |
+
+它只标准化错误，不 retry。
+
+---
+
+## 6. Core retry
+
+### 6.1 固定函数面
+
+| AI SDK | Aimux Rust |
+|---|---|
+| `prepareRetries` | `prepare_retries` |
+| `retryWithExponentialBackoffRespectingRetryHeaders` | `retry_with_exponential_backoff_respecting_retry_headers` |
+| `retryWithExponentialBackoff` | `retry_with_exponential_backoff` |
+| `delay` | `delay` |
+| `mergeAbortSignals` | 不逐字移植；Core 的 `tokio::select!` 同时观察 caller signal、deadline 和 operation future（§8.0） |
+| `setAbortTimeout` | `timeout::OperationTimeout` 保存 deadline，驱动 future 直接 `sleep_until`（§8.0） |
+
+Core operation 把 per-call override、既有 model `RetryConfig` 和 caller abort 交给同名函数：
+
+```text
+prepare_retries(max_retries, retry_config, abort_signal)
+```
+
+`max_retries=None` 使用 `retry_config.max_retries`（默认 2）；per-call `Some(n)` 只覆盖
+该计数，`initial_delay` 和 `backoff_factor` 保持既有配置。`RetryConfig` 的 canonical
+定义移到 `aimux-core::retry`，`aimux-provider-utils::RetryConfig` 和
+`aimux_provider_utils::retry::RetryConfig` 都 re-export 同一类型。旧
+`retry_config` / `with_retry_config` 保留，不增加第二套 provider config 命名。
+
+两层函数的边界与 AI SDK 相同，**不得合并成一个 APICall-aware 的 primitive**：
+
+```rust
+// 通用 primitive：不认识 ApiCallError
+pub(crate) async fn retry_with_exponential_backoff(
+    op: F,
+    max_retries: u32,
+    initial_delay_ms: u64,
+    backoff_factor: u64,
+    abort_signal: Option<&AbortSignal>,
+    should_retry: impl FnMut(&AiMuxError) -> bool,
+    get_delay_ms: impl FnMut(&AiMuxError, u64 /* exponential */) -> u64,
+) -> Result;
+
+// APICall-aware 包装：只填两个 hook
+async fn retry_with_exponential_backoff_respecting_retry_headers<..>(
+    op, max_retries, initial_delay_ms, backoff_factor, abort_signal,
+) {
+    retry_with_exponential_backoff(
+        op, max_retries, initial_delay_ms, backoff_factor, abort_signal,
+        |e| matches!(e, AiMuxError::ApiCall(d) if d.is_retryable),
+        get_retry_delay_ms, // §6.5
+    )
+}
+```
+
+`prepare_retries(max_retries, retry_config, abort_signal)` 返回 `PreparedRetries { max_retries, retry }`
+的 Rust 形式是 `PreparedRetries::retry(&self, op)`，与 AI SDK 的 `{ maxRetries, retry }` 一致；
+不另起额外的 retry 类型。
+
+### 6.2 Retry 边界
+
+非流式：
+
+```text
+retry
+  └─ model.do_generate / do_embed / do_rerank / ...
+       ├─ HTTP exchange
+       ├─ body read
+       ├─ successful response parse
+       └─ Provider result conversion
+```
+
+流式：
+
+```text
+retry
+  └─ model.do_stream
+       ├─ HTTP exchange
+       ├─ peek 第一条 SSE 事件（RFC-0016 M3；是 error 就 Err 返回 → 本次 attempt 失败）
+       └─ 返回 parsed stream（被 peek 的事件重新接回流首，不丢）
+```
+
+stream 返回之后的 SSE error、parse error 或 transport failure 不自动 retry/reconnect。
+peek 在返回之前，因此它产生的错误按 §8.3 走正常 attempt 语义。
+
+### 6.3 `RetryError`
+
+```rust
+#[serde(rename_all = "camelCase")]
+pub enum RetryErrorReason {
+    MaxRetriesExceeded,
+    ErrorNotRetryable,
+}
+
+pub struct RetryError {
+    pub reason: RetryErrorReason,
+    pub errors: Vec,
+}
+
+impl RetryError {
+    pub fn last_error(&self) -> &AiMuxError {
+        self.errors.last().expect("RetryError always contains an error")
+    }
+}
+```
+
+AI SDK 使用 `errors: unknown[]`，因为 enclosing operation 的不同 attempt 可以失败为不同
+错误；Aimux 的等价强类型是 `Vec`，不是 `Vec`。
+
+`last_error` 从 `errors.last()` 派生，不重复序列化。AI SDK public reason union 中的 `abort`
+没有 retry primitive 生产路径；Aimux 同样让 Timeout/Aborted 原样返回，不创建 Abort reason。
+公开 reason 值固定为 `maxRetriesExceeded` / `errorNotRetryable`；各语言只做惯用命名映射，
+不得另造数值 code。
+
+### 6.4 精确行为
+
+默认值与 AI SDK 一致：
+
+```text
+max_retries = 2
+initial_delay = 2000ms
+backoff_factor = 2
+```
+
+| 情况 | 返回 |
+|---|---|
+| `max_retries=0` | 原错误 |
+| 第一次 non-retryable | 原错误 |
+| Timeout/Aborted | 原错误 |
+| retry 后成功 | 成功 |
+| retryable errors 耗尽 | `RetryError::MaxRetriesExceeded` |
+| retry 后遇到 non-retryable | `RetryError::ErrorNotRetryable`，保存全部 errors |
+
+`max_retries=2` 最多执行三次。`RetryError` 自身不可重试，也不得嵌套。
+判定顺序与 AI SDK 相同：先 `try_number > max_retries` → `MaxRetriesExceeded`（**不看最后一个
+错误是否 retryable**：`[retryable, retryable, non-retryable]` 在 `max_retries=2` 下是
+`MaxRetriesExceeded`），再 `should_retry` → delay 后重试，再 `try_number == 1` → 原错误，
+否则 `ErrorNotRetryable`。
+`should_retry` 固定为 `AiMuxError::ApiCall(e) && e.is_retryable`；其他 variant 不得自行
+加入通用 retry 白名单。
+
+固定 message 与 AI SDK 一致：
+
+```text
+Failed after {N} attempts. Last error: {last_error}
+Failed after {N} attempts with non-retryable error: '{last_error}'
+```
+
+### 6.5 Retry headers
+
+delay 依次读取：
+
+1. `retry-after-ms`；
+2. `retry-after` 数字秒；
+3. `retry-after` HTTP-date；
+4. exponential delay。
+
+server hint 只在 AI SDK 的 reasonable-delay 条件成立时使用：
+
+```text
+ms >= 0 && (ms < 60_000 || ms < exponential_delay)
+```
+
+事实源是 `ApiCallError.response_headers`，适用于任意 retryable status，不限于 429。
+不再复制成 `retry_after_ms`；`response_headers` 是唯一事实源。
+
+**Jitter**：AI SDK 的 primitive 通过 `getDelayInMs` 把 delay 策略交给调用方，默认
+`({exponentialBackoffDelay}) => exponentialBackoffDelay`，即默认不抖动但刻意留了注入点。
+Aimux 沿用这个注入点，默认策略保留 RFC-0009 的 Full Jitter——但只作用在 exponential delay
+上：`get_retry_delay_ms` 选中 server hint 时**原样返回 hint**，不再对 hint 取随机（这是
+旧 `get_retry_delay_ms_with_jitter` 把 `Retry-After` 提前的 bug）。
+
+```rust
+// 只回答“有没有可用的 server hint”，不回退到 exponential——
+// 用 `base != exponential` 反推 hint 是否胜出会在 hint 恰好等于 exponential 时误判。
+fn retry_after_hint_ms(error: &AiMuxError, exponential: u64) -> Option;
+
+match retry_after_hint_ms(error, exponential) {
+    Some(hint) => hint,              // 精确遵守，不抖
+    None       => rng.gen_range(0..=exponential),   // Full Jitter（RFC-0009）
+}
+```
+
+RFC-0009 在 jitter 上**不被** supersede；被 supersede 的只是 "retry 在 HTTP 层"。
+
+delay 必须 abort-aware；delay 期间 timeout/abort 直接返回对应错误，不进入 `RetryError`。
+
+---
+
+## 7. 错误模型
+
+### 7.1 `ApiCallError`
+
+字段直接对应 AI SDK，外加现有兼容字段：
+
+```rust
+pub struct ApiCallError {
+    pub message: String,
+    pub url: String,
+    pub request_body_values: serde_json::Value,
+
+    pub status_code: Option,
+    pub response_headers: Option>,
+    pub response_body: Option,
+    pub data: Option,
+    pub is_retryable: bool,
+
+    // Aimux extension:
+    pub provider_code: Option,
+}
+
+pub enum AiMuxError {
+    ApiCall(Box),
+    Retry(RetryError),
+    Timeout(String),
+    Aborted(String),
+    // existing variants...
+}
+```
+
+`url` 和 `request_body_values` 是 required，与 AI SDK 一致。所有 `ApiCallError` producer
+必须从当前 API operation 取得这两个值，不得以 `None` 掩盖缺失的 context。
+`ApiCallError` 不再实现 `Default`；invalid endpoint、client construction 等 pre-HTTP
+失败必须改成其真实的参数/配置/内部错误类型，不能伪造 API context。
+
+`Box` 是必要的 Rust 布局适配：新增 context 后不得删除现有
+`AiMuxError <= 128 bytes` guard；serde、TS 和 binding 的外部结构保持不变。
+`Retry(RetryError)` 同理：`Vec` 本身 24 字节可以不 box，但若 guard 失败则 box 它而不是放宽
+guard。`aimux-core/tests/error_value_golden_test.rs` 的 `error_size_is_pinned` 和
+`variant_set_is_exactly_thirteen`（改名为 fourteen 并加入 `Retry` 的 golden 行）必须在
+第一个提交里就通过——它们就是本节的验收测试，不是迁移尾声再修的东西。
+
+旧路径与本 RFC 同批删除，因此 `request_id` / `retry_after_ms` 不再属于 `ApiCallError`。
+failed handler **不得再从 headers 派生它们**：`response_headers` 是 retry hint 的唯一事实源。
+
+### 7.2 生产规则
+
+| 失败 | 错误 |
+|---|---|
+| DNS/connect/TLS/socket/body transport | no-status、retryable `ApiCallError` |
+| 408/409/429/5xx | 默认 retryable `ApiCallError` |
+| 其他 4xx | 默认 non-retryable `ApiCallError` |
+| 2xx（非流式）body/JSON/schema failure | observed-status、默认 non-retryable `ApiCallError`（AI SDK: `Invalid JSON response`） |
+| peek 到的第一条 SSE error（`do_stream` 返回前，§8.3） | 按其 status/code 构造 `ApiCallError`，retryability 照 HTTP 规则；作为 attempt 失败上抛 |
+| stream 内单条 chunk JSON/schema failure | `JsonParse` / `InvalidResponseData` 作为流的一个 `Err` 项（AI SDK: `JSONParseError`/`TypeValidationError` error part）；**不是** `ApiCallError`，也不终止流 |
+| stream 内 transport failure | no-status、retryable `ApiCallError` 作为 `Err` 项，随后流结束；Core 不重连（§8.3） |
+| 200 Provider business error | Provider 构造 `ApiCallError` 并显式决定 retryable |
+| 用户参数 / binding wire 错误 | 参数或 boundary error，不是 `ApiCallError` |
+
+`is_retryable` 的含义是“安全地重新执行 enclosing model operation”，不只是“该 HTTP
+错误暂时性”。Provider workflow 可以覆盖 HTTP status 的默认判断。
+
+`AiMuxError::JsonParse` / `InvalidResponseData` 两个 variant 保留（它们是 14-variant binding
+契约的一部分，stream chunk 行和 replay/tool-args 解析仍生产它们），但**非流式响应路径不再
+生产它们**：`From for AiMuxError` 仅供非响应路径使用，provider 代码里
+对响应体的 `?` 转换消失（§5.2）。
+
+### 7.3 Redaction
+
+进入错误前必须脱敏：
+
+- URL query；
+- `request_body_values`；
+- `response_headers`；
+- `data` 中的敏感字段。
+
+复用并整理现有 `aimux_core::recording::is_sensitive_key` 规则；logging、recording 和 error
+不得维护三套敏感键表。具体：`aimux-provider-utils/src/logging.rs` 里私有的
+`is_sensitive_key`（5 个子串）删除，`redact_value` / `redact_request_values` /
+`extract_response_headers` 全部调用 `aimux_core::recording::is_sensitive_key`。目前
+header 走 core 表、body 走 logging 表，正是本段禁止的状态。现有 `http.rs::redacted_response_headers` 提取成共享 helper，
+不重写第四份 header policy。`request_body_values` 和 `data` 还必须经过共享的深度/大小限制：binary、
+base64/data URL 和超长字符串只保留类型、长度或安全摘要。`response_body` 继续受现有
+size limit 和 UTF-8 安全截断约束，且不得被默认日志输出。
+
+### 7.4 Timeout 与 Abort
+
+保留简单 public variants，并保存实际 timeout 标签/时长：
+
+```rust
+AiMuxError::Timeout(String)
+AiMuxError::Aborted(String)
+```
+
+这对应 AI SDK 的 timeout/abort 错误；调用方取消使用稳定 message `request aborted`，timeout
+由触发的 deadline 直接生成 `{label} timeout of {ms}ms exceeded`。不新增只重复 message 的
+`TimeoutErrorData`。
+两者均不可 retry，且不得进入 `RetryError.errors`。
+
+---
+
+## 8. Timeout 与 stream
+
+### 8.0 `AbortSignal`
+
+既有 `AbortSignal` 就是 `tokio_util::sync::CancellationToken` 的薄包装，这已经是 Web
+`AbortSignal` 在 Rust 里的惯用对应。它只表达**调用方主动取消**；timeout 是 Core 的 deadline，
+不是写回 signal 的另一种 reason。完整定义（`aimux-core/src/abort_signal.rs`）：
+
+```rust
+#[derive(Debug, Clone, Default)]
+pub struct AbortSignal {
+    token: CancellationToken,
+}
+
+impl AbortSignal {
+    pub fn new() -> Self;
+    pub fn abort(&self);                                  // token.cancel()
+    pub fn is_aborted(&self) -> bool;                     // token.is_cancelled()
+    pub fn cancelled(&self) -> impl Future + Send + 'static; // token.clone().cancelled_owned()
+}
+```
+
+`aimux-core/src/timeout.rs` 在 operation 开始时记录 total/step deadlines，不 spawn timer task：
+
+```rust
+struct OperationTimeout {
+    total: Option,
+    step: Option,
+}
+
+async fn run(operation, caller: Option<&AbortSignal>, timeout: OperationTimeout) {
+    tokio::select! {
+        biased;
+        _ = caller.cancelled() => Err(AiMuxError::Aborted("request aborted".into())),
+        _ = sleep_until(timeout.deadline()) => Err(timeout.deadline().error()),
+        result = operation => result,
+    }
+}
+```
+
+AI SDK 的 `mergeAbortSignals` / `setAbortTimeout` 是围绕 `AbortController` 的 JS 写法；
+Rust future 可以被 drop，因此不需要先把 timeout 转写成另一棵 cancellation tree：外层
+`select!` 赢得 deadline/abort 分支时，operation future（包括当前 HTTP read 或 retry delay）立即
+被 drop。stream 返回后由 Core stream wrapper 用同一 operation deadline 继续 select。
+
+以下是被拒绝的实现，出现任何一条都算跑偏：
+
+- `AbortSignal` 内部持有 `sources: Arc<[AbortSignal]>`，`is_aborted()` / `cancelled()` 递归
+  遍历 sources；
+- `AbortSignal` 再加 `parent`、`OnceLock` 和 `abort_after()`：parent 先取消后，晚到的
+  child timeout 仍可写 child-local reason，使已观察到的取消原因从 Aborted 变成 Timeout；
+- `deadline` 只在 `is_aborted()` 被调用时惰性比较并触发 cancel；deadline 必须由正在驱动
+  operation/stream 的 `sleep_until` 观察；
+- `cancelled()` 返回 `Pin>` 并内部用 `FuturesUnordered` 聚合——既有签名
+  `impl Future + Send + 'static` 够用；
+- 继续放在 `shared.rs`：`shared.rs` 只放 `SharedHeaders` 等纯数据类型，`AbortSignal` 是
+  控制流原语，单独成文件。
+
+Timeout 到期后由 Core timeout/stream wrapper 直接构造 `AiMuxError::Timeout(msg)`；
+`AiMuxError::from_abort_signal` 只产生 `Aborted("request aborted")`。Provider Utils 只能观察 caller
+signal，不拥有或重建 total/step/semantic timeout。
+
+### 8.1 `TimeoutConfiguration`
+
+```rust
+pub struct TimeoutConfiguration {
+    pub total_ms: Option,
+    pub step_ms: Option,
+    pub first_chunk_ms: Option,
+    pub chunk_ms: Option,
+}
+```
+
+| 字段 | 范围 |
+|---|---|
+| `total_ms` | 整个用户 operation：attempts、backoff、parse 和 stream 生命周期 |
+| `step_ms` | 一个 generation step，包括该 step 内的 attempts/backoff |
+| `first_chunk_ms` | 每个 streaming step 从 `do_stream`/setup 开始到第一条 semantic output |
+| `chunk_ms` | 同一 streaming step 的 semantic outputs 之间 |
+
+Core 用同一个 `tokio::select!` 同时观察用户 abort、total/step deadline 与当前 operation
+future，不构造合并 signal，也不把 timeout 写进 `AbortSignal`。外部 abort 在第一次 attempt
+前触发时，不得产生 HTTP exchange。
+
+`first_chunk_ms` 在调用 `do_stream` **之前** arm，因此覆盖 request/setup 以及 §8.3 的
+first-event peek；慢握手、200 后无首帧、以及首条 semantic output 迟到都受同一 budget 约束。
+被 peek 的正常事件会重新接回流首，第一条 semantic output 对用户可见前清除 timer。aimux
+目前是单 step，`step_ms` 与 `total_ms` 作用域重合——照样接受并在同一处 arm，以便字段语义
+跨语言一致、将来多 step 时不改 contract。
+
+### 8.2 Semantic output
+
+分类直接跟随 AI SDK `stream-text.ts` 的 `isOutputChunkType`：
+
+| Aimux `StreamPart` | 重置 semantic timeout |
+|---|---|
+| 非空 `TextDelta` / `ReasoningDelta` / `ToolInputDelta` | 是 |
+| `ToolCall` / `File` | 是 |
+| start/end、metadata、source、raw、finish、error、空 delta | 否 |
+
+第一条 output 必须先清除 first-chunk timer，再对用户可见；随后启动/reset chunk timer。
+SSE keepalive 和原始 network bytes 不得重置 timer。
+
+### 8.3 Stream retry 边界
+
+`model.do_stream` 返回 parsed stream 前可以 retry；返回后：
+
+- 不自动 reconnect；
+- mid-stream transport/parse error 不 retry；
+- 即使尚未产生 token 也不重放。
+
+**第一条 SSE 事件是 error 时**：aimux 的 provider 现在会在 `do_stream` 返回前 peek 第一条
+事件，遇到 error 直接让 `do_stream` 返回 `Err`（RFC-0016 M3）。这发生在 retry 边界**之内**，
+所以它就是一次普通的 attempt 失败，按 `is_retryable` 走 Core retry——一个以 SSE error 形式
+送达的 429 和以 HTTP 429 送达的是同一件事。AI SDK 的 OpenAI provider 会扫描到首个
+semantic output，Aimux 则只检查第一条 event；Aimux 的检查窗口更窄，但两者都把窗口内的
+provider error 留在 operation retry 边界内。`do_stream` 返回 `Ok(stream)` 之后的 error 则按
+上面的规则不 retry。
+
+---
+
+## 9. Aimux-specific workflows
+
+### 9.1 Submit/poll/download
+
+`VideoModel` 已拆成 `do_start` / `do_status`（对齐 AI SDK generate-video 的
+start/status 流），Core 拥有 poll 循环并分别 retry 两个阶段：
+
+- `do_start` 由 Core retry；Core 为它铸造一次 `idempotency-key`（caller 提供的
+  优先），同一个 key 覆盖所有 replay，且不泄漏给 `do_status`（AI SDK
+  generate-video 同款行为）；
+- `do_status` 在同一个 operation reference 上由 Core 独立 retry——poll 失败
+  不会重新 submit；
+- `poll.timeout` 约束轮询节奏（与 AI SDK 相同，在两次 status 检查之间判定超
+  时）；挂起的 status GET 由 provider-utils 的单次交换 30s 响应上限兜底，
+  retry 次数有界；
+- Provider-specific polling delay 不是通用 exponential retry；
+- 耗尽产生的 `RetryError` 原样透传，外层不得再次 submit。
+
+### 9.2 Router / MoA
+
+Composite 外层不重放；重试放在已有语义边界内：
+
+- Router 当前 child 耗尽 retry 后才 fallback，routing 只执行一次；
+- Router stream 只 retry 选中 child 的 setup，返回 stream 后不 fallback/reconnect；
+- MoA 每个 reference 独立 retry，aggregator generate/stream setup 独立 retry；
+- `RouterModel` / `MoaModel` 的 model default 为 0，inner 耗尽产生的 `RetryError`
+  原样穿过外层，避免重跑 fallback/fanout。
+
+### 9.3 Realtime transcription
+
+- `stream_transcribe` 是 user operation，`do_stream` 是 Provider SPI；
+- live audio 不可重放，所以 session setup 只尝试一次，不使用 operation retry；
+- WebSocket 不强行套 POST/GET helper；handshake 仍必须提供等价的 URL/request context 和
+  abort/error normalization；
+- session 建立后不自动重建；
+- `next_part(timeout)` 是会话控制流，不是本 RFC 的 operation Timeout。
+
+### 9.4 Files
+
+`Files::upload_file` 按 AI SDK Provider SPI 本身就是公开 operation，没有对应的
+`do_upload_file`；本轮不为它虚构第二层。upload 只执行一次，Provider
+`RetryConfig` 不适用于 Files。这与旧 HTTP helper 会对 upload exchange 自动重试的
+行为不同；对非幂等 POST 保留该行为可能重复创建文件/上传会话，因此本轮
+明确不保留。Google Files 的 poll GET 如需独立 retry，应作为后续幂等
+exchange 改动，不得包住整个 upload。
+
+---
+
+## 10. User operations、recording 与 bindings
+
+### 10.1 所有 modality 使用 Core operation
+
+Core 必须提供：
+
+```text
+generate_text
+stream_text
+embed
+generate_image
+generate_video
+generate_speech
+transcribe
+stream_transcribe
+rerank
+search
+```
+
+所有 binding 调用这些 user operations，不得直接调用 provider-facing `do_*`。`do_*`
+继续作为 Provider SPI。
+
+为保留现有 model-level 配置，各 model trait 提供只读 `retry_config()`；只有
+Core user operation 读取它。Provider 的 `do_*` 不得重试整个 operation；§9.1 的
+已创建 job 安全 exchange 例外除外。
+
+### 10.2 Recording
+
+retry 上移后，记录必须区分：
+
+```text
+operation_attempt: 1, 2, 3...
+exchange_index: 1, 2, 3... within the attempt
+```
+
+一个 attempt 可包含 submit/poll/download 多个 exchanges。最终 outcome 保存完整
+`RetryError.errors`，不能只保存最后一个错误。
+
+机制：`RecordingContext` 持有一个跨 root/child 共享的 attempt allocator；Core 的 retry
+闭包每次调用前从 allocator 取一个全 call 唯一的 attempt id，并写入**该 context 本地**的
+current attempt。单次 exchange primitive 读取 local current attempt，再用同样属于该 context
+的 exchange counter 编号。child 共享 allocator，但不共享 current attempt / exchange counter，
+因此 MoA 并发或 Router 交错执行时，一个 sibling 的 `start_attempt` 不会重标另一个 sibling
+正在进行的 exchange。provider 不碰这些计数器。
+
+旧 `send_with_retry_raw` 在耗尽时打的 `tracing::error!`（RFC-0014 §4.2 failed 行）随
+retry 一起上移到 Core 的 retry 包装：每次 attempt 失败 `warn`，最终失败 `error`，
+`Aborted` 不算故障不打 error。单次 primitive 只打 exchange 级别的 `info`/`debug`。
+
+### 10.3 Bindings
+
+公开新增：
+
+```text
+RetryError extends AimuxError
+  reason
+  errors
+  lastError / last_error
+  message
+```
+
+每个 nested error 必须恢复真实具体类型。Node/Python 继续使用 napi-rs/PyO3 桥接；
+C-derived bindings 的运输由 RFC-0030 处理。本 RFC 不把 `RetryError.errors` 降级为 JSON
+envelope，也不重新引入 flat error struct。
+
+`ApiCallError` 的新增字段、`Aborted(String)` 和 `TimeoutConfiguration.step_ms` 必须同步到
+Node/Python/Go/Java/Kotlin/Swift/Flutter。
+
+---
+
+## 11. 迁移与改动面
+
+### 11.1 删除映射
+
+| 现有 Aimux | 目标 |
+|---|---|
+| `send_timed` | 删除；Core timeout + 对应 API helper |
+| `send_stream_timed` | 删除；Core semantic timeout + stream response handler |
+| public `send` / `send_stream` | 删除；改用 `post_*_to_api` / `get_from_api` |
+| `send_with_retry_raw` | 删除 |
+| Provider Utils retry execution | 删除；迁入 Core operation |
+| `ErrorStructure` / `DEFAULT_ERROR_STRUCTURE` | 删除 |
+| `parse_provider_error` | 删除；由 failed response handler 替代 |
+| `TimeoutBodyStream` | 删除；由 Core semantic stream timeout 替代 |
+| `get_retry_delay_ms_with_jitter` | 删除；Full Jitter 改为 `aimux-core::retry` 里 `get_delay_ms` hook 的默认实现，且不再抖动 server hint（§6.5） |
+| Provider `resolve_retry_config` | 删除；effective max retries 由 Core 解析 |
+| Provider-owned retry execution | 删除；旧 `RetryConfig` 类型/配置 API 保留，由 Core 执行 |
+| `aimux-provider-utils/src/retry.rs` | 仅保留 `RetryConfig` 兼容 re-export；retry 实现只在 `aimux-core::retry` |
+| `logging.rs::is_sensitive_key` | 删除；统一用 `aimux_core::recording::is_sensitive_key` |
+| `shared.rs::AbortSignal` | 实现迁到 `abort_signal.rs`，旧路径 re-export；timeout 不塞进 signal |
+
+旧 helper 不得保留为并行执行路径。若保留一版 deprecated forwarding wrapper，它只能转发，
+不得自行 retry/timeout。
+
+### 11.2 Provider 迁移规则
+
+每个 HTTP operation 按 method/body/protocol 选择 AI SDK 对应组合：
+
+| 请求/响应 | helper + successful handler |
+|---|---|
+| JSON POST → JSON | `post_json_to_api` + `create_json_response_handler` |
+| JSON POST → SSE | `post_json_to_api` + `create_event_source_response_handler` |
+| multipart POST → JSON | `post_form_data_to_api` + `create_json_response_handler` |
+| POST → binary | `post_to_api` + `create_binary_response_handler` |
+| GET | `get_from_api` + 对应 response handler |
+| custom protocol | 对应 helper + custom `ResponseHandler` |
+
+每个 operation 同时指定 failed handler。共享只以相同 wire schema 为依据，不预设“一个
+Provider 一个 handler”，也不预设“一个调用点一个 handler”。
+
+流式 provider 在拿到 event 流之后、`do_stream` 返回之前保留首条事件的 peek（§8.3）：
+handler 负责分帧和解析，peek 与"错误就 `Err` 返回"仍是 provider 的职责，不下沉到 handler
+（handler 不知道该 endpoint 的错误载荷长什么样）。
+
+### 11.3 Rust 和 binding 改动面
+
+| 区域 | 必要改动 |
+|---|---|
+| `aimux-core` | boxed `ApiCallError`、`RetryError`、Core retry/timeout wrappers、各 modality user operation |
+| `aimux-provider-utils` | 新 API helpers/handlers；删除旧 send/retry/error-structure 路径 |
+| `aimux-providers` | 每个 HTTP operation 迁移 helper 与两个 handlers |
+| recording/tracing | operation attempt 与 exchange index 分离 |
+| 7 个 bindings | 新错误字段、nested `RetryError`、`Aborted` reason、`step_ms`、改走 Core operations |
+| FFI | 仅按 RFC-0030 承载新增 domain error 数据；本 RFC 不改错误运输模型 |
+
+## 12. 实施顺序
+
+0. 先让 `error_value_golden_test`（size guard、variant 数、`Retry`/`Aborted(String)` golden）
+   和 §8.0 的 cancellation/deadline 测试通过——它们是后面每一步的回归网；
+1. `ApiCallError`/`RetryError`/redaction 与 Core retry 单元测试；
+2. Provider Utils API helpers 和标准 response handlers；
+3. OpenAI、Anthropic 作为迁移模板；
+4. 其余 language providers；
+5. multimodal、polling 和 custom-protocol providers；
+6. 所有 binding 改走 Core user operations；
+7. recording/tracing 迁移；
+8. 删除旧 send/retry/`ErrorStructure` 代码；
+9. 更新被 supersede 的 RFC 和 API docs。
+
+阶段 1–7 可以在同一分支上分多次提交，但**合并到 master 的那一刻**不得同时存在 HTTP retry
+与 Core retry 两套 owner。
+
+---
+
+## 13. 验收条件
+
+### Retry
+
+1. `max_retries=0` 返回原错误；
+2. 第一次 non-retryable 返回原错误；
+3. `max_retries=2` 的三次 retryable failure 产生长度 3 的 `MaxRetriesExceeded`；
+4. retryable 后 non-retryable 产生 `ErrorNotRetryable` 并保存完整历史；
+5. Timeout/Aborted 不进入 `RetryError`；
+6. `RetryError` 不嵌套且不可重试；
+7. `RetryError` message 与 AI SDK fixture 逐字一致；delay 的 **hint 分支**与 AI SDK 一致，
+   exponential 分支为 `rand(0..=exp)`（§6.5，测试注入固定 RNG 断言边界）；
+8. 503 的 `Retry-After` 生效且不被 jitter 提前；无 hint 时 delay ∈ [0, exponential]；
+8a. `[retryable, retryable, non-retryable]` @ `max_retries=2` → `MaxRetriesExceeded`；
+8b. primitive 本身不引用 `ApiCallError`（编译期可查：`retry_with_exponential_backoff` 的模块不 `use crate::error::ApiCallError`）。
+
+### HTTP/handlers
+
+9. 每次 API helper 调用只执行一个 fetch attempt 且不 retry；自动 redirect chain 记录为一个 logical exchange；
+10. successful/failed handler 只执行被 status 选中的一个；
+11. 2xx body/JSON/schema failure 带完整、已脱敏的 API context；
+12. 200 Provider business error 可以显式决定 retryability；
+13. request context、URL、response headers/data 使用统一 redaction；
+14. response body size limit 和 UTF-8 截断正确；
+15. custom binary/event-stream protocol 不被误当 SSE；
+15a. JSON helper 之后 provider 代码里不存在对响应体的 `serde_json::from_slice`/`from_str`；
+15b. `aimux-provider-utils` 不导出任何 provider 名字的 failed handler；
+15c. 每次 helper 调用恰好产生一条 exchange recording（含 transport failure）。
+
+### Timeout/stream/workflow
+
+16. total timeout 覆盖 attempts、backoff、parse 和 stream 生命周期；
+17. first/chunk timer 只由 semantic output 重置；
+18. first timer 在 `do_stream`/setup 前启动，并在第一条 output 可见前清除；
+19. stream 返回后的 error 不 reconnect/retry；`do_stream` 返回前 peek 到的 retryable error 走 Core retry；
+19a. stream 内 chunk 解析失败是 `JsonParse`/`InvalidResponseData` 的 `Err` 项且流继续；
+19b. peek 未消费掉数据：首条事件是正常 chunk 时它仍作为流的第一项被 yield（测试：单事件 SSE 响应，断言流里能拿到该 chunk）；
+19c. 首条 SSE error 为 429 时 `stream_text` 重试后成功；为 400 时立即返回原错误不重试；
+20. polling failure 在已有 job 上 retry，耗尽也不重复 submit；
+21. Router child 在 fallback 前耗尽 retry，MoA reference/aggregator 独立 retry，且不重跑 routing/fanout；
+21a. `AbortSignal` 符合 §8.0：只有 caller cancellation；total/step deadline 直接返回对应 Timeout message，不改变 signal；
+21b. 一个 operation 结束后没有残留 timer/forwarding task（实现不 spawn，测试以 paused time 验证 drop operation 后不再有可触发状态）。
+
+### Bindings/observability
+
+22. 所有 modality binding 调用 Core user operation；
+23. 七个 binding 无损恢复 `RetryError.errors` 的具体错误类型；
+24. `ApiCallError` 新字段和稳定的 `Aborted("request aborted")` 跨语言一致；
+25. recording 区分 operation attempt 与 exchange index。
+
+---
+
+## 14. 不做与拒绝方案
+
+本 RFC 不做：`cause`、`GatewayError`、mid-stream reconnect、stream 返回后的任何 retry、
+circuit breaker、tool execution timeout、per-child Router retry policy。
+
+拒绝：
+
+- **HTTP retry + Core retry**：会产生乘法 attempts、双重 backoff 和嵌套错误；
+- **只重命名 `send_with_retry_raw`**：retry 边界仍然错误；
+- **`RetryError.errors: Vec`**：不等价于 AI SDK 的 `unknown[]`，会丢失
+  operation 的异构错误历史；
+- **按 raw bytes 重置 stream timeout**：keepalive/framing 不是 semantic output；
+- **新的统一 `call_to_api`**：AI SDK 已有 POST/GET helpers 和 ResponseHandler 边界，
+  不再发明平行 API（四个 helper 背后的 `pub(crate)` 单次 primitive 不算，见 §5.4）；
+- **JSON handler 返回 `Bytes` 让 provider 再解析**：schema 失败会丢掉 API context；
+- **把 timeout 合并进 `AbortSignal`**：见 §8.0；JS 需要 AbortController 取消 fetch，Rust 直接
+  drop future，更简单也不会产生 parent/child reason 竞争；
+- **把第一条 SSE 事件的 peek 扩展为 AI SDK OpenAI provider 的完整扫描窗口**：baseline 会
+  扫描到首个 semantic output，Aimux 保留较窄的 first-event peek（§8.3）：同一个 429 不应
+  因为 provider 用 SSE 而非 HTTP status 送达就得不到重试，且 C-ABI binding 用一个 nullable
+  error 指针即可表达失败，无需在流里再解一层。代价是与 AI SDK 的 stream fixture 不能逐字
+  对照——接受。
+
+---
+
+## 15. 对既有 RFC 的影响
+
+| RFC | 影响 |
+|---|---|
+| RFC-0009 | 保留 shared client/pool/connect safety 与 Full Jitter（改挂在 Core retry 的 `get_delay_ms` 上）；supersede HTTP retry loop；30s client-wide timeout 改为仅包非流式 helper 的 per-exchange whole-response guard（§5.4） |
+| RFC-0016 | supersede “retry 在 HTTP 层”和 raw-byte stream timeout 的实现结论；**保留** M3 的首条 SSE 事件 peek，并明确它是 retry 边界内的 attempt 失败（§8.3） |
+| RFC-0017 | `max_retries` 仍可作 model default，但由 Core 解析和执行 |
+| RFC-0021 | composite 作为一个 attempt；per-child policy 另行设计 |
+| RFC-0023 | attempt 拆为 operation attempt + exchange index |
+| RFC-0028 | realtime session 控制流保持独立；仅 handshake 使用 operation policy |
+| RFC-0030 | 负责 nested errors 的 C ABI 运输；不改变本 RFC 的 domain error 模型 |
+
+---
+
+## 16. 最终契约
+
+> Core 用 AI SDK 同名 retry/timeout 函数包围完整 `model.do_*`；Provider 用 AI SDK
+> 同名 POST/GET helpers 和 successful/failed ResponseHandler 完成一次 fetch attempt；
+> stream 返回后不重放；旧 `send*`、HTTP retry 和 `ErrorStructure` 路径全部删除。
diff --git a/docs/api/c.md b/docs/api/c.md
index dbf9e508..d240ca42 100644
--- a/docs/api/c.md
+++ b/docs/api/c.md
@@ -71,7 +71,7 @@ The code range identifies the source:
 
 | Range | Meaning |
 |---|---|
-| `1..13` | `AiMuxError` |
+| `1..14` | `AiMuxError` (14 = `Retry`) |
 | `100..105` | `RecordingError` |
 | `200..206` | failure detected by the C ABI |
 
@@ -83,15 +83,18 @@ The code range identifies the source:
 ```c
 int32_t aimux_error_code(const aimux_error_t *err);        /* aimux_error_code_t; AIMUX_OK for NULL */
 char   *aimux_error_message(const aimux_error_t *err);     /* every code; owned → aimux_free_string */
-int32_t aimux_error_retryable(const aimux_error_t *err);   /* 1 = retrying may help; every code */
+int32_t aimux_error_retryable(const aimux_error_t *err);   /* 1 = retrying may help; only API_CALL can answer 1 */
 int32_t aimux_error_status(const aimux_error_t *err);      /* API_CALL: observed status or -1; TOKEN_EXPIRED: 401; else -1 */
 
 /* AIMUX_E_API_CALL — NULL / -1 under any other code */
 int64_t aimux_error_retry_ms(const aimux_error_t *err);         /* retry hint, or -1; 0 = retry now */
 char   *aimux_error_provider_code(const aimux_error_t *err);    /* provider's own code, e.g. "insufficient_quota" */
 char   *aimux_error_provider_message(const aimux_error_t *err); /* the failure's own text; message is the composed form */
-char   *aimux_error_request_id(const aimux_error_t *err);
-char   *aimux_error_response_body(const aimux_error_t *err);
+char   *aimux_error_response_body(const aimux_error_t *err);    /* provider response body; first 64 KiB, "…(truncated)" beyond */
+char   *aimux_error_url(const aimux_error_t *err);              /* sanitized request URL */
+char   *aimux_error_request_body_values(const aimux_error_t *err); /* sanitized request body values, as a JSON string */
+char   *aimux_error_response_headers(const aimux_error_t *err); /* sanitized headers, one JSON object string, e.g. {"retry-after-ms":"1500"} */
+char   *aimux_error_provider_data(const aimux_error_t *err);    /* parsed provider error data, as a JSON string */
 
 /* AIMUX_E_NO_SUCH_MODEL */
 char   *aimux_error_model_id(const aimux_error_t *err);
@@ -99,25 +102,48 @@ char   *aimux_error_model_type(const aimux_error_t *err);
 
 /* AIMUX_E_NO_SUCH_PROVIDER */
 char   *aimux_error_provider_id(const aimux_error_t *err);
+
+/* AIMUX_E_RETRY — retrying stopped; the error keeps the per-attempt history */
+char   *aimux_error_retry_reason(const aimux_error_t *err);     /* "maxRetriesExceeded" | "errorNotRetryable" */
+int32_t aimux_error_retry_count(const aimux_error_t *err);      /* recorded attempts; 0 under any other code */
+aimux_error_t *aimux_error_retry_error_at(const aimux_error_t *err, int32_t index); /* NEW owned error, or NULL */
 ```
 
 `code` and `message` answer for every non-`NULL` error. Every other getter
 belongs to an AiMuxError code and returns `NULL` / `-1` / `0` under any other
 code — read it inside the matching `case`. All returned `char *` are owned by the caller
 (`aimux_free_string`); a `NULL` payload string under its own code means the
-provider did not send it (`provider_code`, `request_id`, `response_body` are
-optional even under `AIMUX_E_API_CALL`). Every getter is `NULL`-safe.
+provider did not send it (`provider_code`, `response_body`, `url`,
+`response_headers`, `provider_data` are optional even under
+`AIMUX_E_API_CALL`). Provider request ids ride in `response_headers`.
+Every getter is `NULL`-safe.
 
 Branch on `retryable`, never on the `status` sentinel: a transport failure and
 a missing API key both report `status == -1` and disagree about whether a retry
-would help. `retry_ms` is a *hint* that rides along — it is `-1` whenever the
-provider advertised no delay (neither a `retry-after` / `retry-after-ms`
-response header nor a `retry_after_ms` / `retry_after` member in the JSON
-error payload), including on a retryable status, so fall back to your own
-exponential backoff when it is negative.
-
-The unified `aimux_error_code_t` keeps the existing AiMuxError values 1–13,
-adds RecordingError values 100–105, and assigns C ABI failures 200–206:
+would help. `retry_ms` is a *hint* that rides along — derived from the
+response headers (`retry-after-ms`, then `retry-after` in seconds or
+HTTP-date form); it is `-1` whenever the headers advertised no delay,
+including on a retryable status, so fall back to your own exponential
+backoff when it is negative.
+
+`AIMUX_E_RETRY` (14) means the core ran the operation's retry loop and gave
+up: `aimux_error_retry_reason` says why — `"maxRetriesExceeded"` (every
+permitted attempt failed with a retryable error) or `"errorNotRetryable"` (a
+later attempt failed with a non-retryable error) — and `aimux_error_message`
+composes the summary (`Failed after 2 attempts. Last error: …`). Walk the
+per-attempt history with `aimux_error_retry_count` /
+`aimux_error_retry_error_at` (index 0 = oldest; `count - 1` = the final
+attempt). Each attempt comes back as a **new owned** `aimux_error_t *`: read
+it with these same getters — it can itself be any AiMuxError code, including
+`AIMUX_E_API_CALL` with the full detail set — and release it with
+`aimux_error_free` independently of the parent (free order is
+unconstrained).
+
+The unified `aimux_error_code_t` extends the AiMuxError values to 1–14
+(`AIMUX_E_RETRY` = 14 reclaims the slot the pre-unification `Other` vacated —
+the opaque-pointer ABI break means no old caller can misread it), adds
+RecordingError values 100–105, and assigns C ABI
+failures 200–206:
 `NULL_POINTER`, `INVALID_UTF8`, `INVALID_WIRE_JSON`, `INVALID_HANDLE`,
 `REENTRANT_CALL`, `RESULT_SERIALIZATION`, and `CALLBACK_FAILURE`. Values are
 never renumbered or reused. A code outside the enum is a header/library
diff --git a/docs/api/flutter.md b/docs/api/flutter.md
index 87050af5..b6c49c28 100644
--- a/docs/api/flutter.md
+++ b/docs/api/flutter.md
@@ -76,7 +76,7 @@ shares a base with the other (both just `implements Exception`):
 
 | Source | Rust | Dart | C code |
 |---|---|---|---|
-| AiMux | `AiMuxError` | `AimuxException` hierarchy | 1..13 |
+| AiMux | `AiMuxError` | `AimuxException` hierarchy | 1..14 |
 | recorder | `RecordingError` | `RecordingException` | 100..105 |
 
 Every fallible C call returns an opaque `aimux_error_t *` (`NULL` =
@@ -84,7 +84,7 @@ success, result in the trailing out-parameter). The binding has one decoder
 (`errors.dart`): `expectAimuxError(e, context)` for model calls,
 `expectRecordingError(e, context)` for `initRecording` / `recordingTryFlush`,
 `expectFfiError(e, context)` for utilities that can only fail in the C ABI.
-One unified code selects 1..13, 100..105, or 200..206; each decoder copies the
+One unified code selects 1..14, 100..105, or 200..206; each decoder copies the
 relevant fields, releases the error with `aimux_error_free` exactly once, and
 throws the matching `AimuxException` subclass / `RecordingException`. Codes
 200..206 throw the native
@@ -103,6 +103,7 @@ Exception (implements)
       ├── UnsupportedFunctionalityError
       ├── NoSuchModelError / NoSuchProviderError
       ├── APICallError              // every HTTP-shaped failure; branch on status
+      ├── RetryError                // reason, errors, lastError
       ├── AimuxTimeoutError
       ├── RequestAbortedError
       └── OtherError
@@ -111,11 +112,16 @@ Exception (implements)
 Every instance has `message`, `code` (`AimuxErrorCode` constants matching C
 `aimux_error_code_t`), `status` (HTTP or `-1`), `retryMs` (hint or `-1`;
 `0` = retry now) and `retryable`. Per-code payload lives on the carrying
-subclass only: `APICallError.providerCode` / `.providerMessage` / `.requestId` / `.responseBody`
+subclass only: `APICallError.providerCode` / `.providerMessage` / `.responseBody`
 (`String?`), `NoSuchModelError.modelId` / `.modelType`, and
 `NoSuchProviderError.providerId`. A code outside the enum is a header/library
 mismatch and fails with `StateError`, not an error type.
 
+`APICallError` additionally exposes the sanitized URL/request values,
+response headers/body, parsed provider data/code, and retryability.
+`RetryError` preserves every concrete attempt error in `errors`, with
+`lastError` and `reason`.
+
 ```dart
 import 'package:aimux/aimux.dart'; // exports errors.dart
 
diff --git a/docs/api/go.md b/docs/api/go.md
index 5c99dee0..4e500ff0 100644
--- a/docs/api/go.md
+++ b/docs/api/go.md
@@ -56,11 +56,11 @@ if err != nil {
 | `Status` | HTTP status, or `-1` |
 | `RetryMs` | rate-limit hint, or `-1` (`0` = retry now) |
 | `Retryable` | the core's retry verdict; never derived from `Status` |
-| `ProviderCode`, `ProviderMessage`, `RequestID`, `ResponseBody` | `CodeAPICall` payload; empty under any other code |
+| `ProviderCode`, `ProviderMessage`, `ResponseBody`, `URL`, `RequestBodyValues`, `ResponseHeaders`, `Data` | `CodeAPICall` payload; empty under any other code (request-id evidence rides in `ResponseHeaders`) |
 | `ModelID`, `ModelType` | `CodeNoSuchModel` payload; empty under any other code |
 | `ProviderID` | `CodeNoSuchProvider` payload; empty under any other code |
 
-`Code` values 1..13 mirror aimux-core's `AiMuxError` variants. A code outside
+`Code` values 1..14 mirror aimux-core's `AiMuxError` variants. A code outside
 the enum is a header/library mismatch and fails with a `panic`, not an error
 type.
 
@@ -94,7 +94,7 @@ no Go type of their own; the binding maps them to native Go errors:
 
 Decoder: every fallible C call returns an opaque `aimux_error_t *` (NULL =
 success, result in the out-parameter). One `aimux_error_code()` distinguishes
-`AiMuxError` (1–13), `RecordingError` (100–105), and C ABI failures (200–206).
+`AiMuxError` (1–14), `RecordingError` (100–105), and C ABI failures (200–206).
 `expectAimuxError`, `expectRecordingError`, and `expectFfiError` enforce the
 range expected by each call; the first two restore `*Error` and
 `*RecordingError`, while 200–206 becomes a plain `error`. Every path frees the
@@ -111,13 +111,16 @@ design** — that one is opt-in, and every one of its five entry points has a
 |------------|--------------------------|
 | `aimux.go` `mustNew` — behind `OpenAI` / `OpenAIWithBase` / `Anthropic` / `AnthropicWithBase` / `DeepSeek` | **Yes, by design.** `regexp.MustCompile` convention: an `apiKey` / `modelID` / `baseURL` that is not valid UTF-8 or contains a NUL panics, as does any AiMuxError failure. Use `NewOpenAI` / `NewOpenAIWithBase` / `NewAnthropic` / `NewAnthropicWithBase` / `NewDeepSeek` for anything caller-supplied |
 | `aimux.go` `InitLogging` — `expectFfiError` returned an error | No. `level` is coerced first: empty, non-UTF-8, or NUL-bearing falls back to `"warn"`, which is what aimux-core does with an unparseable level anyway (`AIMUX_LOG` / `AIMUX_LOG_LEVEL` outrank it regardless). That leaves no documented failure for `aimux_init_logging`, so a non-nil error here is a header/library mismatch |
-| `aimux.go` `expectAimuxError` — `aimux_error_code_t` outside 1..13 | No. Header/library version mismatch |
+| `aimux.go` `expectAimuxError` — `aimux_error_code_t` outside 1..14 | No. Header/library version mismatch |
 | `aimux.go` `expectRecordingError` — `aimux_error_code_t` outside the enum | No. Header/library version mismatch |
 | `multimodal.go` `TranscriptionSession.NextPart` — unknown `aimux_transcription_next_part` state | No. Header/library version mismatch |
 
 The three mismatch panics are a contract violation the C header itself says to
 abort on, not an error to report.
 
+For `CodeRetry`, `Reason`, `Errors`, and `LastError()` preserve
+the concrete attempt history.
+
 ## Quick Start
 
 ```go
@@ -324,6 +327,8 @@ Video generation typically returns a URL (not binary).
 ```go
 prompt := "A cat playing piano"
 n := 1
+pollIntervalMs := uint64(1_000)
+pollTimeoutMs := uint64(120_000)
 
 videor, err := aimux.NewGoogleVideo("sk-...", "veo-3.0")
 if err != nil {
@@ -334,6 +339,10 @@ defer videor.Close()
 resultJSON, err := videor.Generate(&aimux.VideoCallOptions{
     Prompt: &prompt,
     N:      &n,
+    Poll: &aimux.VideoPollOptions{
+        IntervalMS: &pollIntervalMs,
+        TimeoutMS:  &pollTimeoutMs,
+    },
 })
 if err != nil {
     log.Fatal(err)
diff --git a/docs/api/java.md b/docs/api/java.md
index 6c7f11ec..c3e16b5f 100644
--- a/docs/api/java.md
+++ b/docs/api/java.md
@@ -75,6 +75,7 @@ RuntimeException
       ├── UnsupportedFunctionalityError
       ├── NoSuchModelError / NoSuchProviderError
       ├── APICallError               // every HTTP-shaped failure; classify on getStatusCode()
+      ├── RetryError                 // reason, errors, lastError
       ├── TimeoutError / RequestAbortedError
       └── OtherError
 ```
@@ -84,7 +85,7 @@ Every instance has:
 | Field | Meaning |
 |-------|---------|
 | `getMessage()` | human-readable text from C |
-| `getCode()` | `aimux_error_code_t` value 1–13 (matches `aimux-error.h`) |
+| `getCode()` | `aimux_error_code_t` value 1–14, where 14 = `Retry` (matches `aimux-error.h`) |
 | `getStatusCode()` | HTTP status, or `-1` |
 | `getRetryMs()` | rate-limit hint, or `-1` (`0` = retry now) |
 | `isRetryable()` | the `AiMuxError` retry verdict (not derivable from status) |
@@ -92,10 +93,15 @@ Every instance has:
 A code outside the enum is a header/library mismatch and fails with
 `IllegalStateException`, not an error type.
 
-Three subclasses carry the C payload of their variant (nullable `String`s,
-`null` when unavailable): `APICallError` —
-`getProviderCode()`, `getProviderMessage()`, `getRequestId()`, `getResponseBody()`;
-`NoSuchModelError` — `getModelId()`, `getModelType()`;
+Four subclasses carry the C payload of their variant (`null` when
+unavailable): `APICallError` — `getProviderCode()`, `getProviderMessage()`,
+`getResponseBody()`, `getUrl()`, `getRequestBodyValues()`,
+`getResponseHeaders()`, `getData()`; `RetryError` — `getReason()`
+(`RetryErrorReason.MAX_RETRIES_EXCEEDED` — every permitted attempt failed
+with a retryable error — or `ERROR_NOT_RETRYABLE` — a later attempt failed
+non-retryably), `getErrors()` (the per-attempt history, oldest first, each
+itself an `AimuxException` — typically `APICallError` with its full detail),
+`getLastError()`; `NoSuchModelError` — `getModelId()`, `getModelType()`;
 `NoSuchProviderError` — `getProviderId()`.
 
 Recording failures are a **separate type**, mirroring the two unrelated
@@ -134,11 +140,12 @@ types and share no base beyond `RuntimeException`.
 Transport: every fallible C call returns an opaque `aimux_error_t *`
 (JNA `Pointer`) — `null` on success with the result in a trailing out-parameter
 (`LongByReference` handle / `PointerByReference` JSON), non-null on failure.
-`AimuxResult` reads one unified code: 1–13 restores the matching
+`AimuxResult` reads one unified code: 1–14 restores the matching
 `AimuxException` subclass, 100–105 restores `RecordingException`, and 200–206
 becomes `IllegalStateException("aimux ffi: …")`. Payload getters are read only
-under their owning AiMuxError code. Every returned string is freed and the
-returned error is released with
+under their owning AiMuxError code; a `RetryError`'s attempt errors are new
+owned errors, decoded recursively and freed by the binding. Every returned
+string is freed and the returned error is released with
 `aimux_error_free` exactly once (errors are never handles). No JSON error
 envelopes on the main path. Subclasses are nested under `AimuxException` (e.g.
 `AimuxException.APICallError`).
@@ -360,7 +367,7 @@ in the wrapper-object wire form (e.g. `{"TextDelta":{...}}`).
 `EmbeddingCallOptions` / `EmbeddingResult`, `SpeechCallOptions` / `SpeechResult`,
 `ImageCallOptions` / `ImageResult`, `TranscriptionCallOptions` /
 `TranscriptionResult`, `RerankingCallOptions` / `RerankingResult`,
-`VideoCallOptions` / `VideoResult`, `SearchCallOptions` / `SearchResult`,
+`VideoCallOptions` / `VideoPollOptions` / `VideoResult`, `SearchCallOptions` / `SearchResult`,
 `UploadFileCallOptions` / `UploadFileResult`. Sealed unions (`AudioData`,
 `ImageOutputs`, `VideoData`) use custom Jackson serializers with the
 externally-tagged wire form. Response types (`EmbeddingResponse`, etc.) match
diff --git a/docs/api/kotlin.md b/docs/api/kotlin.md
index bd4f940f..f744f259 100644
--- a/docs/api/kotlin.md
+++ b/docs/api/kotlin.md
@@ -56,7 +56,7 @@ Two aimux exception types, each mirroring its own Rust type — **AiMux**
 (`AimuxException`) and **recorder** (`RecordingException`). They share no base
 beyond `RuntimeException`; catch each on its own. Every fallible C call returns
 an `aimux_error_t *` (null = success, result in the out-parameter). The
-binding reads one unified code: 1..13 restores an `AimuxException` subclass,
+binding reads one unified code: 1..14 restores an `AimuxException` subclass,
 100..105 restores `RecordingException`, and 200..206 becomes
 `IllegalStateException("aimux ffi: …")`. Payload getters are read only under
 their owning AiMuxError code.
@@ -80,7 +80,8 @@ RuntimeException
       ├── UnsupportedFunctionalityError
       ├── NoSuchModelError / NoSuchProviderError   // modelId + modelType / providerId
       ├── APICallError               // every HTTP-shaped failure; classify on status
-      │                              // + providerCode, providerMessage, requestId, responseBody (null when absent)
+      │                              // + providerCode, providerMessage, responseBody, url, requestBodyValues, responseHeaders, data (null when absent)
+      ├── RetryError                 // the retry loop gave up; reason, errors (oldest first), lastError
       ├── TimeoutError / RequestAbortedError
       └── OtherError
 ```
@@ -104,10 +105,16 @@ try {
 
 | Field | Meaning |
 |-------|---------|
-| `code` | `AIMUX_E_*` matching C `aimux_error_code_t` (1..13; 1 is the catch-all `Other`) |
+| `code` | `AIMUX_E_*` matching C `aimux_error_code_t` (1..14, where 14 = `Retry`; 1 is the catch-all `Other`) |
 | `status` | HTTP status when known; otherwise `-1` |
 | `retryMs` | Rate-limit hint in ms; `-1` if none; `0` = retry immediately |
 
+`RetryError` preserves the per-attempt history: `reason`
+(`RetryErrorReason.MAX_RETRIES_EXCEEDED` — every permitted attempt failed
+with a retryable error — or `ERROR_NOT_RETRYABLE` — a later attempt failed
+non-retryably), `errors` (oldest first, each itself an `AimuxException` —
+typically `APICallError` with its full detail), and `lastError`.
+
 Recording errors are a separate type, mirroring Rust's `recording::RecordingError`
 (C codes 100..105): `initRecording()` and `recordingTryFlush()` throw
 `RecordingException(code: RecordingErrorCode, message)` — a plain `RuntimeException`,
@@ -237,6 +244,10 @@ typed model surface: `Role`, `FinishReasonUnified`, `ReasoningEffort`,
 `FileBytes` / `FileData` (sealed), `GenerateContent` (sealed),
 `GenerateResult`, `GenerateTextResult`, `StreamPart` (sealed).
 
+`MultimodalTypes.kt` includes `VideoCallOptions.poll: VideoPollOptions?`;
+`intervalMs` and `timeoutMs` serialize as `interval_ms` / `timeout_ms` for the
+Core-owned video status loop.
+
 ## Coverage
 
 Full multimodal surface — text generation, streaming, embedding, TTS, STT
diff --git a/docs/api/node.md b/docs/api/node.md
index 0eb7307d..949a9ef7 100644
--- a/docs/api/node.md
+++ b/docs/api/node.md
@@ -89,6 +89,7 @@ checks); the recorder throws its own class:
 Error
  └── AimuxError
       ├── APICallError              // provider call/transport failure; status when observed
+      ├── RetryError                // the retry loop gave up; reason, errors (oldest first), lastError
       ├── JSONParseError / InvalidResponseDataError / ToolError
       ├── InvalidArgumentError / InvalidPromptError
       ├── TokenExpiredError
@@ -126,11 +127,15 @@ works directly for synchronous calls, promises, and stream/session errors;
 Every `AimuxError` instance has the ordinary `Error` fields. There is no aimux
 `code` discriminator and no JSON companion. Payload fields belong to the class
 that carries them: `APICallError` adds `retryable` and optional `status` /
-`retryMs` / `providerCode` /
-`providerMessage` / `responseBody` / `requestId`, `TokenExpiredError` carries
-`status: 401`, `NoSuchModelError` adds `modelId` / `modelType`, and
-`NoSuchProviderError` adds `providerId`. Missing HTTP status and retry hints are
-absent rather than represented by `-1`.
+`retryMs` / `url` / `requestBodyValues` / `responseHeaders` / `providerCode` /
+`providerMessage` / `responseBody` / `data`; `RetryError` adds `reason`
+(`'maxRetriesExceeded'` — every permitted attempt failed with a retryable
+error — or `'errorNotRetryable'` — a later attempt failed non-retryably),
+`errors` — the per-attempt history, oldest first, each itself an error from
+this hierarchy — and `lastError`; `TokenExpiredError` carries `status: 401`;
+`NoSuchModelError` adds `modelId` / `modelType`; and `NoSuchProviderError`
+adds `providerId`. Missing HTTP status and retry hints are absent rather than
+represented by `-1`.
 
 ```typescript
 import { generateText, AimuxError, APICallError } from '@arcships/aimux'
@@ -379,6 +384,7 @@ const videor = await googleVideo('sk-...', 'veo-3.0')
 const resultJson = await videor.generate(JSON.stringify({
   prompt: 'A cat playing piano',
   n: 1,
+  poll: { interval_ms: 1_000, timeout_ms: 120_000 },
   provider_options: {},
 }))
 const result = JSON.parse(resultJson)
@@ -389,6 +395,9 @@ if (result.videos[0].Url) {
 }
 ```
 
+The package root exports the generated `VideoCallOptions` and
+`VideoPollOptions` types; both poll fields are milliseconds.
+
 ## Reranking
 
 Reorders a document list by relevance.
diff --git a/docs/api/python.md b/docs/api/python.md
index de0f1e32..b354c4dd 100644
--- a/docs/api/python.md
+++ b/docs/api/python.md
@@ -336,6 +336,7 @@ OpenAI/Anthropic SDK style, same idea as Vercel AI SDK on JS):
 Exception
  └── AimuxError
       ├── APICallError              # provider call/transport failure; status when observed
+      ├── RetryError                # the retry loop gave up; reason, errors (oldest first), last_error
       ├── JSONParseError / InvalidResponseDataError / ToolError
       ├── InvalidArgumentError / InvalidPromptError
       ├── TokenExpiredError
@@ -368,11 +369,15 @@ pyo3 does, as Python builtins, never disguised as an `AimuxError`:
 JSON that parses but has a bad value is still `InvalidArgumentError`.
 
 Payload attributes belong to the class that carries them and are absent on the
-others. `APICallError` has `status` / `retryable` / `retry_ms` /
-`provider_code` / `provider_message` / `response_body` / `request_id`;
-optional values use Python's normal `None`. `TokenExpiredError` has
-`status == 401`, `NoSuchModelError` has `model_id` / `model_type`, and
-`NoSuchProviderError` has `provider_id`.
+others. `APICallError` has `status` / `retryable` / `retry_ms` / `url` /
+`request_body_values` / `response_headers` / `provider_code` /
+`provider_message` / `response_body` / `data`; optional values use Python's
+normal `None`. `RetryError` has `reason` (`"maxRetriesExceeded"` — every
+permitted attempt failed with a retryable error — or `"errorNotRetryable"` —
+a later attempt failed non-retryably), `errors` — the per-attempt history,
+oldest first, each itself an exception from this hierarchy — and
+`last_error`. `TokenExpiredError` has `status == 401`, `NoSuchModelError` has
+`model_id` / `model_type`, and `NoSuchProviderError` has `provider_id`.
 
 ```python
 from aimux import (
diff --git a/docs/api/reference.md b/docs/api/reference.md
index a3422512..ab1b5e71 100644
--- a/docs/api/reference.md
+++ b/docs/api/reference.md
@@ -11,10 +11,10 @@
 |------|-------------|----------------|-------|
 | `GenerateTextOptions` | [generate.rs](../../aimux-core/src/generate.rs) | [GenerateTextOptions.ts](../../bindings/node/src/types/GenerateTextOptions.ts) | per-call text options (temperature, tools, reasoning, timeout, …) |
 | `CallOptions` | [options.rs](../../aimux-core/src/options.rs) | [CallOptions.ts](../../bindings/node/src/types/CallOptions.ts) | base options for all modalities |
-| `TimeoutConfiguration` | [options.rs](../../aimux-core/src/options.rs) | [TimeoutConfiguration.ts](../../bindings/node/src/types/TimeoutConfiguration.ts) | `total_ms` / `first_chunk_ms` / `chunk_ms` |
+| `TimeoutConfiguration` | [options.rs](../../aimux-core/src/options.rs) | [TimeoutConfiguration.ts](../../bindings/node/src/types/TimeoutConfiguration.ts) | `total_ms` / `step_ms` / `first_chunk_ms` / `chunk_ms` |
 | `ResponseFormat` | [options.rs](../../aimux-core/src/options.rs) | [ResponseFormat.ts](../../bindings/node/src/types/ResponseFormat.ts) | text / json_object / json_schema |
 | `ReasoningEffort` | [types.rs](../../aimux-core/src/types.rs) | [ReasoningEffort.ts](../../bindings/node/src/types/ReasoningEffort.ts) | 7 levels, passed through verbatim |
-| `AbortSignal` | [shared.rs](../../aimux-core/src/shared.rs) | — (runtime handle) | Rust: options field; Node: `AbortBridge` + JS `AbortSignal` |
+| `AbortSignal` | [abort_signal.rs](../../aimux-core/src/abort_signal.rs) | — (runtime handle) | Rust: options field; Node: `AbortBridge` + JS `AbortSignal` |
 | `ProviderOptions` | [provider.rs](../../aimux-providers/src/provider.rs) | `ProviderConfig` (per binding) | `base_url` / `headers` / `organization` / `project` / `max_retries` / `body_overrides` |
 | `ProviderName` | [provider_name.rs](../../aimux-providers/src/provider_name.rs) | [ProviderName.ts](../../bindings/node/src/types/ProviderName.ts) | one constant per registry provider (251); also generated for Go / Python / Java / Kotlin / Swift / Dart (`scripts/gen_provider_names.py`) |
 | `GenerateTextResult` | [generate.rs](../../aimux-core/src/generate.rs) | [GenerateTextResult.ts](../../bindings/node/src/types/GenerateTextResult.ts) | `text`, `tool_calls`, `usage`, `warnings`, `raw` |
diff --git a/docs/api/rust.md b/docs/api/rust.md
index a806c0c6..fbdc04ba 100644
--- a/docs/api/rust.md
+++ b/docs/api/rust.md
@@ -83,7 +83,7 @@ Cancellation — set the runtime `abort_signal` handle (never crosses the JSON
 boundary; FFI bindings cannot set it):
 
 ```rust
-let signal = aimux_core::shared::AbortSignal::new();
+let signal = aimux_core::AbortSignal::new();
 let opts = GenerateTextOptions {
     abort_signal: Some(signal.clone()),
     ..Default::default()
@@ -115,9 +115,10 @@ while let Some(part) = stream.next().await {
 ```
 
 Streaming honors the same `timeout` / `abort_signal` options as
-[text generation](#text-generation); streamed timeouts surface as
-`StreamPart::Error { error: AiMuxError::Timeout(..) }` and aborting the
-signal ends the stream with `StreamPart::Error { error: AiMuxError::Aborted }`.
+[text generation](#text-generation); streamed timeouts surface as an
+`Err(AiMuxError::Timeout(..))` stream item, and aborting the signal ends the
+stream with `Err(AiMuxError::Aborted(..))`. Provider-reported error events
+remain `Ok(StreamPart::Error { .. })` data.
 
 > Stream part variants are documented in the [API overview](../API.md#streaming-generation).
 
@@ -174,11 +175,11 @@ let result = generate_text(&model, messages, opts).await?;
 Converts text into a vector representation.
 
 ```rust
-use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel};
+use aimux_core::embedding_model::{EmbeddingCallOptions, embed};
 
 let model = provider.embedding_model("text-embedding-3-small");
 let opts = EmbeddingCallOptions::new("hello");
-let result = model.do_embed(&opts).await?;
+let result = embed(&model, opts).await?;
 // result.embeddings[0] is Vec
 ```
 
@@ -187,11 +188,11 @@ let result = model.do_embed(&opts).await?;
 Converts text into speech audio.
 
 ```rust
-use aimux_core::speech_model::{SpeechCallOptions, SpeechModel};
+use aimux_core::speech_model::{SpeechCallOptions, generate_speech};
 
 let model = provider.speech("tts-1");
 let opts = SpeechCallOptions::new("Hello world!");
-let result = model.do_generate(&opts).await?;
+let result = generate_speech(&model, opts).await?;
 // result.audio is AudioData::Base64(String) or AudioData::Binary(Vec)
 ```
 
@@ -200,25 +201,25 @@ let result = model.do_generate(&opts).await?;
 Converts audio into text (non-streaming).
 
 ```rust
-use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, TranscriptionModel};
+use aimux_core::transcription_model::{AudioInput, TranscriptionCallOptions, transcribe};
 
 let model = provider.transcription("whisper-1");
 let opts = TranscriptionCallOptions::new(
     AudioInput::Base64(audio_base64),
     "audio/mp3",
 );
-let result = model.do_generate(&opts).await?;
+let result = transcribe(&model, opts).await?;
 // result.text, result.segments, result.language
 ```
 
 ## Image Generation
 
 ```rust
-use aimux_core::image_model::{ImageCallOptions, ImageModel};
+use aimux_core::image_model::{ImageCallOptions, generate_image};
 
 let model = provider.image("dall-e-3");
 let opts = ImageCallOptions { prompt: Some("A cute sea otter".into()), n: 1, .. };
-let result = model.do_generate(&opts).await?;
+let result = generate_image(&model, opts).await?;
 // result.images is ImageOutputs::Base64(Vec) or Binary(Vec>)
 ```
 
@@ -227,11 +228,11 @@ let result = model.do_generate(&opts).await?;
 Video generation typically returns a URL (not binary).
 
 ```rust
-use aimux_core::video_model::{VideoCallOptions, VideoModel};
+use aimux_core::video_model::{VideoCallOptions, generate_video};
 
 let model = provider.video("veo-3.0");
 let opts = VideoCallOptions { prompt: Some("A cat".into()), n: 1, .. };
-let result = model.do_generate(&opts).await?;
+let result = generate_video(&model, opts).await?;
 // result.videos[0] is VideoData::Url { url, media_type }
 ```
 
@@ -240,11 +241,11 @@ let result = model.do_generate(&opts).await?;
 Reorders a document list by relevance.
 
 ```rust
-use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, RerankingModel};
+use aimux_core::reranking_model::{RerankingCallOptions, RerankingDocuments, rerank};
 
 let model = provider.reranking_model("rerank-v3.0");
 let opts = RerankingCallOptions::new("What is Rust?", docs);
-let result = model.do_rerank(&opts).await?;
+let result = rerank(&model, opts).await?;
 // result.ranking sorted by score
 ```
 
@@ -253,11 +254,11 @@ let result = model.do_rerank(&opts).await?;
 Calls a search provider to obtain results.
 
 ```rust
-use aimux_core::search_model::{SearchCallOptions, SearchModel};
+use aimux_core::search_model::{SearchCallOptions, search};
 
 let model = provider.search_model("tavily-search");
 let opts = SearchCallOptions::new("What is Rust?");
-let result = model.do_search(&opts).await?;
+let result = search(&model, opts).await?;
 // result.results is Vec
 ```
 
@@ -291,7 +292,7 @@ The Rust core provides 10 traits/interfaces, implemented by each provider as nee
 | `TranscriptionModel` | `do_generate`, `do_stream` | Speech to text |
 | `ImageModel` | `do_generate` | Image generation |
 | `RerankingModel` | `do_rerank` | Reranking |
-| `VideoModel` | `do_generate` | Video generation |
+| `VideoModel` | `do_start`, `do_status` | Video generation (Core-driven start/poll flow) |
 | `SearchModel` | `do_search` | Search |
 | `Files` | `upload_file` | File upload |
 
@@ -321,7 +322,8 @@ Rust types are the canonical definitions — one module per feature in
 | `content` | `ContentPart` (Text / Image / File / Reasoning / ToolCall / ToolResult) |
 | `tool` | `Tool`, `FunctionTool`, `ProviderTool`, `ToolCall`, `ToolResult` |
 | `types` | `Usage`, `TokenUsage`, `FinishReason`, `ResponseMetadata`, `Warning` |
-| `shared` | `FileBytes`, `FileData`, `Size`, `AspectRatio`, `ResponseInfo`, `AbortSignal` |
+| `shared` | `FileBytes`, `FileData`, `Size`, `AspectRatio`, `ResponseInfo` |
+| `abort_signal` | `AbortSignal` |
 | `embedding_model` | `EmbeddingModel`, `EmbeddingCallOptions`, `EmbeddingResult` |
 | `speech_model` | `SpeechModel`, `SpeechCallOptions`, `SpeechResult` |
 | `transcription_model` | `TranscriptionModel`, `TranscriptionCallOptions`, `TranscriptionResult`, `AudioInput` |
diff --git a/docs/api/swift.md b/docs/api/swift.md
index 451851b3..2ef760fa 100644
--- a/docs/api/swift.md
+++ b/docs/api/swift.md
@@ -107,7 +107,7 @@ have no Swift type — see "C ABI failures" below.
 Every fallible C function returns an opaque `aimux_error_t *`
 (`OpaquePointer?`): `NULL` = success (the result is in the trailing
 out-parameter), non-`NULL` = failure. One unified code selects `AimuxError`
-(1...13), `RecordingError` (100...105), or a C ABI failure (200...206).
+(1...14), `RecordingError` (100...105), or a C ABI failure (200...206).
 The three decoders enforce the range expected by each call and restore the
 Swift error type; 200...206 collapses to `DecodingError.dataCorrupted`.
 Every path copies its strings (freed with
@@ -131,6 +131,7 @@ and yields the invariant `DecodingError.dataCorrupted("aimux ffi: :
 | `.noSuchModel` | `AIMUX_E_NO_SUCH_MODEL` (9) | Registry miss |
 | `.noSuchProvider` | `AIMUX_E_NO_SUCH_PROVIDER` (10) | Unknown provider id |
 | `.apiCall` | `AIMUX_E_API_CALL` (11) | Every HTTP-shaped failure; branch on `status` (401 auth, 404 model, 429 rate limit) |
+| `.retry` | `AIMUX_E_RETRY` (14) | Complete attempt history and stop reason |
 | `.timeout` | `AIMUX_E_TIMEOUT` (12) | Request timed out |
 | `.aborted` | `AIMUX_E_ABORTED` (13) | Request aborted |
 | `.other` | `AIMUX_E_OTHER` (1) | Unclassified core error |
@@ -145,7 +146,7 @@ with `DecodingError.dataCorrupted` from the decoder, not an error type.
 Every case carries `message`, `status` (`Int?` — `nil` when C reports no
 status), `retryMs` (`Int64?` — `nil` if none; `0` = retry now) and
 `retryable`. Three cases carry a
-typed payload as extra associated values: `.apiCall(providerCode:providerMessage:requestId:responseBody:)`
+typed payload as extra associated values: `.apiCall(providerCode:providerMessage:responseBody:url:requestBodyValues:responseHeaders:providerData:)`
 (all optional), `.noSuchModel(modelId:modelType:)` and
 `.noSuchProvider(providerId:)`; the same-named computed properties return
 `nil` on every other case. `e.code` returns the mapped `aimux_error_code_t`
diff --git a/docs/error-model.md b/docs/error-model.md
index 40c35576..ddb55459 100644
--- a/docs/error-model.md
+++ b/docs/error-model.md
@@ -20,10 +20,10 @@
 - C ABI 成功时返回 `NULL` 并写入出参；失败时返回一个由调用方释放一次的
   不透明 `aimux_error_t *`。
 - Go、Java、Kotlin、Swift 和 Flutter 将 C 错误还原为本语言错误。错误码
-  `1..13` 属于核心，`100..105` 属于录制，`200..206` 属于绑定边界。
+  `1..14` 属于核心，`100..105` 属于录制，`200..206` 属于绑定边界。
 
-结构化字段只放在真正拥有它的错误上。例如，HTTP 状态和请求 ID 属于
-`APICallError`，不放进通用基类。
+结构化字段只放在真正拥有它的错误上。例如，HTTP 状态和响应头属于
+`APICallError`，不放进通用基类（retry hint 与 request id 从 `response_headers` 读取）。
 
 ## 兼容性约定
 
diff --git a/rfc/0009-request-resilience.md b/rfc/0009-request-resilience.md
index ef3eb875..1617f7fa 100644
--- a/rfc/0009-request-resilience.md
+++ b/rfc/0009-request-resilience.md
@@ -4,6 +4,12 @@
 > **Date**: 2026-07-29
 > **Scope**: `aimux-provider-utils` references three specific design points from catcher (connection-pool config, jitter backoff, fixed timeouts) and implements request-layer optimization using reqwest natively + the existing retry.rs, without introducing a catcher-http dependency
 > **Related**: [RFC-0002](0002-provider-improvements.md) provider adapter layer improvements, [RFC-0003](0003-test-cassette.md) test cassette plan
+>
+> **Superseded in part by [RFC-0031](0031-ai-sdk-request-pipeline.md)**:
+> shared client/pool/connect-timeout remain, but HTTP-layer retry and fixed
+> whole-response timeout no longer describe the current architecture. Retry
+> and operation deadlines now belong to Core; provider-utils performs one
+> exchange per API-helper call.
 
 ## 1. Motivation
 
diff --git a/rfc/0016-align-with-aisdk.md b/rfc/0016-align-with-aisdk.md
index 5350dc68..2f72a554 100644
--- a/rfc/0016-align-with-aisdk.md
+++ b/rfc/0016-align-with-aisdk.md
@@ -4,6 +4,13 @@
 > **Date**: 2026-08-01
 > **Scope**: 系统对比 aimux 0.1.2 与 Vercel AI SDK (`@ai-sdk/openai` / `@ai-sdk/openai-compatible` / `@ai-sdk/provider` V4) 的接口与实现,识别能力缺口,并按优先级规划补齐路径
 > **Related**: [RFC-0009](0009-request-resilience.md) request resilience(retry/timeout,本 RFC 的 abortSignal/timeout 缺口与之相关),[RFC-0014](0014-logging.md) 统一日志体系(可观测性缺口依赖本 RFC 的 span 树)
+>
+> **Request-pipeline sections are superseded by
+> [RFC-0031](0031-ai-sdk-request-pipeline.md)**: retry/timeout ownership moved
+> to Core, provider-utils now exposes one-exchange API helpers plus response
+> handlers, and the old `send_timed` / `send_stream_timed` implementation was
+> removed. RFC-0016 M3's first-SSE-event peek is intentionally retained as
+> Aimux's documented deviation from the AI SDK.
 
 ---
 

From 1c5e73bb74c728203c0dc8f568347faf257ad37a Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Mon, 31 Aug 2026 16:14:53 +0800
Subject: [PATCH 07/24] fix: end the xAI stream on a terminal error; clamp the
 retry hint
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two defects this branch introduced, found by an adversarial review of the
diff (14 other candidate findings were refuted against the code).

xAI Responses yielded StreamPart::Error for a terminal in-stream error and
then continued reading. That was harmless on master, where both consumers
returned on the first Error part; this branch makes Error non-terminal so a
trailing Finish can still deliver usage and recording, which turns the
continue into a hang against a server that keeps the connection open after
the error. The sibling OpenAI Responses reducer was fixed this way in this
branch already, test included; xAI is its near-duplicate and was missed.

retry_after_hint() returns negative milliseconds for a Retry-After HTTP-date
in the past. The retry loop drops those via u64::try_from, but the C ABI
projects the raw value through aimux_error_retry_ms, whose header documents
-1 as 'absent' — so a hint one millisecond in the past reads as no hint, and
a larger negative reaches Go/Java as a sleep duration. Upstream's check is
'0 <= ms' (retry-with-exponential-backoff.ts); apply it at the source.
---
 aimux-core/src/error.rs                  |  4 +++
 aimux-providers/src/xai/responses/mod.rs | 39 ++++++++----------------
 2 files changed, 16 insertions(+), 27 deletions(-)

diff --git a/aimux-core/src/error.rs b/aimux-core/src/error.rs
index d52c7e1b..d8420941 100644
--- a/aimux-core/src/error.rs
+++ b/aimux-core/src/error.rs
@@ -315,6 +315,10 @@ impl AiMuxError {
                         })
                 })
             })
+            // Upstream's `0 <= ms` check: a hint in the past is no hint. It
+            // also keeps the value clear of the C ABI's -1 = absent sentinel
+            // (`aimux_error_retry_ms`).
+            .filter(|milliseconds| *milliseconds >= 0)
     }
 
     /// Returns the HTTP status code carried by this error, if any.
diff --git a/aimux-providers/src/xai/responses/mod.rs b/aimux-providers/src/xai/responses/mod.rs
index fa4426f0..70a85b07 100644
--- a/aimux-providers/src/xai/responses/mod.rs
+++ b/aimux-providers/src/xai/responses/mod.rs
@@ -109,25 +109,14 @@ impl LanguageModel for XaiResponsesModel {
         let provider_tool_names = request_result.provider_tool_names;
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), build_header_list(&headers), options),
             body.clone(),
             super::xai_successful_response_handler::(),
             super::xai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let _raw_value = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
@@ -377,25 +366,14 @@ impl LanguageModel for XaiResponsesModel {
         let endpoint = self.endpoint();
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             super::xai_event_source_response_handler::(),
             super::xai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
@@ -636,7 +614,14 @@ impl LanguageModel for XaiResponsesModel {
                                     stream_response_headers.clone(),
                                 ),
                             });
-                            continue;
+                            final_finish_reason = FinishReason {
+                                unified: FinishReasonUnified::Error,
+                                raw: Some("error".to_string()),
+                            };
+                            // A terminal error ends this stream; waiting for
+                            // more events can hang on a source that keeps the
+                            // connection open.
+                            break;
                         }
 
                         // ── custom_tool_call_input.delta / .done ──

From 03880e04c5e62b6944f5b200e8576157bdaf359d Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Mon, 31 Aug 2026 16:15:05 +0800
Subject: [PATCH 08/24] refactor: delete the boilerplate the pipeline migration
 created
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Unifying a mechanism should shrink the diff; this branch's provider layer
grew instead. Four repetitions carried no information, so they are gone:

- HttpRequest gains Default. 120 of the 131 construction sites listed the
  SSRF guard's four fields as all-defaults; only the eight sites that
  actually download a provider-supplied URL still spell them out.
- ResponseHandlerOutput.response_headers drops its Option. Every one of the
  twelve producers filled it with Some(..), so 63 files were unwrapping a
  value that could not be absent — google/files.rs even carried an
  .expect("header-only response handler always supplies headers").
- HttpRequest::new(url, headers, options) via a small ExchangeContext trait.
  111 sites copied the same three lines to inherit cancellation, call id and
  recording context from the caller's options; 72 collapse to one line. The
  rest keep the literal because they set a custom response_timeout or the
  download guard.
- Six providers whose error JSON is the shared { error: { message, type |
  code } } shape now call create_standard_json_error_response_handler.

Left alone deliberately: the 33 retry_config overrides return each
provider's configured value rather than the trait default, and 27 of the 32
error handlers parse genuinely different error shapes.
---
 aimux-provider-utils/src/http.rs              | 71 ++++++++++++++++++-
 aimux-provider-utils/src/lib.rs               |  7 +-
 aimux-provider-utils/src/post_to_api.rs       |  5 +-
 aimux-provider-utils/src/response_handler.rs  | 34 +++++++--
 .../tests/api_helpers_test.rs                 | 10 +--
 .../tests/http_recording_test.rs              |  5 +-
 aimux-provider-utils/tests/logging_test.rs    |  5 +-
 aimux-providers/src/anthropic/files.rs        | 13 +---
 aimux-providers/src/anthropic/mod.rs          |  5 +-
 aimux-providers/src/anthropic/stream.rs       |  7 +-
 aimux-providers/src/anthropic_aws/mod.rs      |  5 +-
 aimux-providers/src/assemblyai.rs             | 41 ++---------
 aimux-providers/src/aws_polly.rs              |  7 +-
 aimux-providers/src/azure/model.rs            |  5 +-
 aimux-providers/src/azure/responses.rs        | 30 ++------
 aimux-providers/src/bedrock/embedding.rs      |  7 +-
 aimux-providers/src/bedrock/image.rs          |  7 +-
 aimux-providers/src/bedrock/mod.rs            |  5 +-
 aimux-providers/src/bedrock/model.rs          | 16 ++---
 aimux-providers/src/bedrock/reranking.rs      |  7 +-
 aimux-providers/src/black_forest_labs.rs      | 13 +---
 aimux-providers/src/cartesia.rs               | 30 ++------
 aimux-providers/src/codex.rs                  | 26 ++-----
 aimux-providers/src/cohere/embedding.rs       | 15 +---
 aimux-providers/src/cohere/mod.rs             |  5 +-
 aimux-providers/src/cohere/model.rs           | 38 ++++------
 aimux-providers/src/cohere/reranking.rs       | 15 +---
 aimux-providers/src/dataforseo.rs             |  7 +-
 aimux-providers/src/deepgram.rs               |  7 +-
 aimux-providers/src/elevenlabs.rs             | 26 +++----
 aimux-providers/src/exa_ai.rs                 |  7 +-
 aimux-providers/src/fal.rs                    |  8 +--
 aimux-providers/src/firecrawl.rs              |  7 +-
 aimux-providers/src/gladia.rs                 | 36 +++-------
 aimux-providers/src/google/embedding.rs       | 14 ++--
 aimux-providers/src/google/files.rs           |  4 +-
 aimux-providers/src/google/image.rs           | 30 ++------
 aimux-providers/src/google/mod.rs             |  5 +-
 aimux-providers/src/google/model.rs           | 34 +++------
 aimux-providers/src/google/video.rs           | 22 ++----
 aimux-providers/src/google_pse.rs             |  7 +-
 aimux-providers/src/huggingface/responses.rs  | 32 ++-------
 aimux-providers/src/hume.rs                   |  7 +-
 aimux-providers/src/jina_ai.rs                |  7 +-
 aimux-providers/src/klingai.rs                | 30 ++------
 aimux-providers/src/linkup.rs                 | 27 +------
 aimux-providers/src/lmnt.rs                   | 15 +---
 aimux-providers/src/luma.rs                   | 32 ++-------
 aimux-providers/src/mistral/embedding.rs      | 15 +---
 aimux-providers/src/mistral/mod.rs            |  5 +-
 aimux-providers/src/mistral/model.rs          | 38 ++++------
 aimux-providers/src/open_responses.rs         | 40 ++++-------
 aimux-providers/src/openai/embedding.rs       | 15 +---
 aimux-providers/src/openai/files.rs           | 13 +---
 aimux-providers/src/openai/image.rs           | 38 +++-------
 aimux-providers/src/openai/model.rs           | 35 ++-------
 aimux-providers/src/openai/responses/mod.rs   | 30 ++------
 aimux-providers/src/openai/speech.rs          | 15 +---
 aimux-providers/src/openai/transcription.rs   | 15 +---
 aimux-providers/src/parallel_ai.rs            | 27 +------
 aimux-providers/src/prodia.rs                 | 49 +++----------
 aimux-providers/src/recraft.rs                |  2 +-
 aimux-providers/src/replicate.rs              |  6 +-
 aimux-providers/src/revai.rs                  | 53 +++++---------
 aimux-providers/src/runwayml.rs               | 30 ++------
 aimux-providers/src/searxng.rs                |  7 +-
 aimux-providers/src/serper.rs                 | 27 +------
 aimux-providers/src/stability.rs              | 17 +----
 aimux-providers/src/tavily.rs                 |  7 +-
 aimux-providers/src/tinyfish.rs               | 27 +------
 aimux-providers/src/vertex/anthropic_model.rs | 12 +---
 aimux-providers/src/vertex/embedding.rs       | 14 ++--
 aimux-providers/src/vertex/image.rs           | 30 ++------
 aimux-providers/src/vertex/mod.rs             |  5 +-
 aimux-providers/src/vertex/model.rs           | 14 ++--
 aimux-providers/src/vertex/transcription.rs   |  7 +-
 aimux-providers/src/vertex/video.rs           | 30 ++------
 aimux-providers/src/voyage/embedding.rs       | 16 +----
 aimux-providers/src/voyage/reranking.rs       | 15 +---
 aimux-providers/src/xai/mod.rs                |  4 +-
 aimux-providers/src/xai/model.rs              | 30 ++------
 aimux-providers/src/you_com.rs                | 27 +------
 82 files changed, 380 insertions(+), 1133 deletions(-)

diff --git a/aimux-provider-utils/src/http.rs b/aimux-provider-utils/src/http.rs
index 58b381ca..6b61d32a 100644
--- a/aimux-provider-utils/src/http.rs
+++ b/aimux-provider-utils/src/http.rs
@@ -204,7 +204,7 @@ impl HttpMethod {
 /// Metadata shared by `post_json_to_api`, `post_form_data_to_api`,
 /// `post_to_api`, and `get_from_api`. Method and body are fixed by the helper
 /// signature rather than supplied as runtime fields.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub struct HttpRequest {
     pub url: String,
     pub headers: Vec<(String, String)>,
@@ -235,6 +235,75 @@ pub struct HttpRequest {
     pub credentialed_origin: Option,
 }
 
+/// The per-operation context every model's call options carry into an
+/// exchange: cancellation, and (for language models) the call id and
+/// recording context the RFC-0031 pipeline threads through.
+pub trait ExchangeContext {
+    fn abort_signal(&self) -> Option;
+    /// Only language-model `CallOptions` participate in recording; the other
+    /// modalities have no call id to correlate on.
+    fn call_id(&self) -> Option {
+        None
+    }
+    fn recording_context(&self) -> Option {
+        None
+    }
+}
+
+impl ExchangeContext for aimux_core::options::CallOptions {
+    fn abort_signal(&self) -> Option {
+        self.abort_signal.clone()
+    }
+    fn call_id(&self) -> Option {
+        self.call_id.clone()
+    }
+    fn recording_context(&self) -> Option {
+        self.recording_context.clone()
+    }
+}
+
+macro_rules! exchange_context_abort_only {
+    ($($ty:path),* $(,)?) => {
+        $(impl ExchangeContext for $ty {
+            fn abort_signal(&self) -> Option {
+                self.abort_signal.clone()
+            }
+        })*
+    };
+}
+
+exchange_context_abort_only!(
+    aimux_core::search_model::SearchCallOptions,
+    aimux_core::speech_model::SpeechCallOptions,
+    aimux_core::image_model::ImageCallOptions,
+    aimux_core::video_model::VideoCallOptions,
+    aimux_core::transcription_model::TranscriptionCallOptions,
+    aimux_core::embedding_model::EmbeddingCallOptions,
+    aimux_core::reranking_model::RerankingCallOptions,
+    aimux_core::files_model::UploadFileCallOptions,
+);
+
+impl HttpRequest {
+    /// A plain exchange that inherits the operation's cancellation and
+    /// recording context. Downloads of provider-supplied URLs set the guard
+    /// fields explicitly instead.
+    #[must_use]
+    pub fn new(
+        url: impl Into,
+        headers: Vec<(String, String)>,
+        options: &impl ExchangeContext,
+    ) -> Self {
+        Self {
+            url: url.into(),
+            headers,
+            abort_signal: options.abort_signal(),
+            call_id: options.call_id(),
+            recording_context: options.recording_context(),
+            ..Self::default()
+        }
+    }
+}
+
 /// SSRF guard configuration carried by a validated exchange
 /// (`HttpRequest.validate_url`).
 #[derive(Debug, Clone)]
diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs
index ced24453..679f24be 100644
--- a/aimux-provider-utils/src/lib.rs
+++ b/aimux-provider-utils/src/lib.rs
@@ -29,7 +29,9 @@ pub use api_key::load_api_key;
 pub use download_guard::same_origin;
 pub use get_from_api::get_from_api;
 pub use headers::with_user_agent_suffix;
-pub use http::{HttpBody, HttpRequest, ProxyConfig, init_proxy, shared_client, sleep_or_abort};
+pub use http::{
+    ExchangeContext, HttpBody, HttpRequest, ProxyConfig, init_proxy, shared_client, sleep_or_abort,
+};
 pub use logging::{body_logging_enabled, init_logging, redact_body, redact_error_context};
 pub use multipart::{MultipartForm, media_type_to_extension};
 pub use post_to_api::{post_form_data_to_api, post_json_to_api, post_to_api};
@@ -37,7 +39,8 @@ pub use response_handler::{
     ProviderErrorParts, ResponseHandler, ResponseHandlerInput, ResponseHandlerOutput,
     create_binary_response_handler, create_event_source_response_handler,
     create_json_error_response_handler, create_json_response_handler,
-    create_status_code_error_response_handler, stream_error_api_call,
+    create_standard_json_error_response_handler, create_status_code_error_response_handler,
+    stream_error_api_call,
 };
 pub use retry::RetryConfig;
 pub use url::{validate_base_url, without_trailing_slash, without_trailing_slash_opt};
diff --git a/aimux-provider-utils/src/post_to_api.rs b/aimux-provider-utils/src/post_to_api.rs
index bc801813..1d9a3212 100644
--- a/aimux-provider-utils/src/post_to_api.rs
+++ b/aimux-provider-utils/src/post_to_api.rs
@@ -235,10 +235,7 @@ mod tests {
             abort_signal: None,
             call_id: None,
             recording_context: None,
-            response_timeout: None,
-            validate_url: false,
-            trusted_origin: None,
-            credentialed_origin: None,
+            ..Default::default()
         };
         let error = call_to_api_with_deadline(
             request.prepare(HttpMethod::Post, HttpBody::Json(serde_json::json!({}))),
diff --git a/aimux-provider-utils/src/response_handler.rs b/aimux-provider-utils/src/response_handler.rs
index 7c86ffc0..f0f8d4c9 100644
--- a/aimux-provider-utils/src/response_handler.rs
+++ b/aimux-provider-utils/src/response_handler.rs
@@ -27,7 +27,7 @@ pub struct ResponseHandlerInput {
 pub struct ResponseHandlerOutput {
     pub value: T,
     pub raw_value: Option,
-    pub response_headers: Option>,
+    pub response_headers: std::collections::HashMap,
 }
 
 type HandlerFuture =
@@ -133,7 +133,7 @@ where
         Ok(ResponseHandlerOutput {
             value,
             raw_value: Some(raw_value),
-            response_headers: Some(headers),
+            response_headers: headers,
         })
     })
 }
@@ -175,6 +175,28 @@ fn error_body_string(body: &[u8], read_truncated: bool) -> String {
     value
 }
 
+/// Handler for the `{ "error": { "message", "type" | "code" } }` error shape,
+/// which several providers share verbatim. Providers whose error JSON differs
+/// keep their own mapping via [`create_json_error_response_handler`].
+#[must_use]
+pub fn create_standard_json_error_response_handler() -> ResponseHandler {
+    create_json_error_response_handler(|data| {
+        let error = data.get("error").unwrap_or(data);
+        ProviderErrorParts {
+            message: error
+                .get("message")
+                .and_then(serde_json::Value::as_str)
+                .unwrap_or_default()
+                .to_owned(),
+            provider_code: error
+                .get("type")
+                .or_else(|| error.get("code"))
+                .and_then(serde_json::Value::as_str)
+                .map(str::to_owned),
+        }
+    })
+}
+
 /// Parse an error JSON and map its data to a provider message and code.
 #[must_use]
 pub fn create_json_error_response_handler(
@@ -226,7 +248,7 @@ where
         Ok(ResponseHandlerOutput {
             value: error,
             raw_value: None,
-            response_headers: Some(headers),
+            response_headers: headers,
         })
     })
 }
@@ -303,7 +325,7 @@ pub fn create_binary_response_handler() -> ResponseHandler {
         Ok(ResponseHandlerOutput {
             value,
             raw_value: None,
-            response_headers: Some(headers),
+            response_headers: headers,
         })
     })
 }
@@ -382,7 +404,7 @@ where
         Ok(ResponseHandlerOutput {
             value,
             raw_value: None,
-            response_headers: Some(output_headers),
+            response_headers: output_headers,
         })
     })
     .streaming()
@@ -419,7 +441,7 @@ pub fn create_status_code_error_response_handler() -> ResponseHandler HttpRequest {
         call_id: None,
         recording_context: None,
 
-        response_timeout: None,
-        validate_url: false,
-        trusted_origin: None,
-        credentialed_origin: None,
+        ..Default::default()
     }
 }
 
@@ -75,10 +72,7 @@ async fn post_json_dispatches_typed_success_handler() {
 
     assert_eq!(output.value, Reply { value: "ok".into() });
     assert_eq!(output.raw_value, Some(json!({"value": "ok"})));
-    assert_eq!(
-        output.response_headers.unwrap().get("x-test"),
-        Some(&"yes".into())
-    );
+    assert_eq!(output.response_headers.get("x-test"), Some(&"yes".into()));
 }
 
 #[tokio::test]
diff --git a/aimux-provider-utils/tests/http_recording_test.rs b/aimux-provider-utils/tests/http_recording_test.rs
index 076e48b7..f17933fc 100644
--- a/aimux-provider-utils/tests/http_recording_test.rs
+++ b/aimux-provider-utils/tests/http_recording_test.rs
@@ -53,10 +53,7 @@ fn request_with_signal(
             call_id, recorder,
         )),
 
-        response_timeout: None,
-        validate_url: false,
-        trusted_origin: None,
-        credentialed_origin: None,
+        ..Default::default()
     }
 }
 
diff --git a/aimux-provider-utils/tests/logging_test.rs b/aimux-provider-utils/tests/logging_test.rs
index 7d229644..2d7d0799 100644
--- a/aimux-provider-utils/tests/logging_test.rs
+++ b/aimux-provider-utils/tests/logging_test.rs
@@ -28,10 +28,7 @@ fn request(url: String) -> HttpRequest {
         call_id: None,
         recording_context: None,
 
-        response_timeout: None,
-        validate_url: false,
-        trusted_origin: None,
-        credentialed_origin: None,
+        ..Default::default()
     }
 }
 
diff --git a/aimux-providers/src/anthropic/files.rs b/aimux-providers/src/anthropic/files.rs
index e6bbb424..f7cd5e8b 100644
--- a/aimux-providers/src/anthropic/files.rs
+++ b/aimux-providers/src/anthropic/files.rs
@@ -131,18 +131,7 @@ impl Files for AnthropicFiles {
         // sets `Content-Type` from it, so it is intentionally not added to the
         // header list above.
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             HttpBody::Bytes(body, content_type),
             aimux_provider_utils::create_json_response_handler(),
             super::anthropic_failed_response_handler(),
diff --git a/aimux-providers/src/anthropic/mod.rs b/aimux-providers/src/anthropic/mod.rs
index 2e3ef928..5f2c5e39 100644
--- a/aimux-providers/src/anthropic/mod.rs
+++ b/aimux-providers/src/anthropic/mod.rs
@@ -341,10 +341,7 @@ impl Provider for AnthropicProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         anthropic_failed_response_handler(),
diff --git a/aimux-providers/src/anthropic/stream.rs b/aimux-providers/src/anthropic/stream.rs
index 7ac6affa..7a53b146 100644
--- a/aimux-providers/src/anthropic/stream.rs
+++ b/aimux-providers/src/anthropic/stream.rs
@@ -112,10 +112,7 @@ fn build_anthropic_request(
             abort_signal,
             call_id: recording_context.as_ref().map(|c| c.call_id.clone()),
             recording_context,
-            response_timeout: None,
-            validate_url: false,
-            trusted_origin: None,
-            credentialed_origin: None,
+            ..Default::default()
         },
         http_body,
     ))
@@ -830,7 +827,7 @@ pub(crate) async fn anthropic_stream_core(
     )
     .await?;
 
-    let response_headers = resp.response_headers.unwrap_or_default();
+    let response_headers = resp.response_headers;
     let mut sse_stream = resp.value;
     let first_event = match sse_stream.next().await {
         Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error),
diff --git a/aimux-providers/src/anthropic_aws/mod.rs b/aimux-providers/src/anthropic_aws/mod.rs
index 18619640..52699f9a 100644
--- a/aimux-providers/src/anthropic_aws/mod.rs
+++ b/aimux-providers/src/anthropic_aws/mod.rs
@@ -258,10 +258,7 @@ impl Provider for AnthropicAwsProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         crate::anthropic::anthropic_failed_response_handler(),
diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs
index e66ef007..6103ea4c 100644
--- a/aimux-providers/src/assemblyai.rs
+++ b/aimux-providers/src/assemblyai.rs
@@ -238,18 +238,7 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
 
         // Step 1: Upload audio.
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.upload_url(),
-                headers: header_list.clone(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.upload_url(), header_list.clone(), options),
             HttpBody::Bytes(audio_bytes, "application/octet-stream".to_string()),
             aimux_provider_utils::create_json_response_handler(),
             assemblyai_failed_response_handler(),
@@ -288,18 +277,7 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
         }
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.transcript_url(),
-                headers: header_list.clone(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.transcript_url(), header_list.clone(), options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             assemblyai_failed_response_handler(),
@@ -327,18 +305,7 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
             let resp = retries
                 .retry(|| {
                     aimux_provider_utils::get_from_api(
-                        HttpRequest {
-                            url: poll_url.clone(),
-                            headers: header_list.clone(),
-
-                            abort_signal: options.abort_signal.clone(),
-                            call_id: None,
-                            recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
-                        },
+                        HttpRequest::new(poll_url.clone(), header_list.clone(), options),
                         aimux_provider_utils::create_json_response_handler::<
                             AssemblyAITranscriptResponse,
                         >(),
@@ -347,7 +314,7 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
                 })
                 .await?;
 
-            response_headers = resp.response_headers.unwrap_or_default();
+            response_headers = resp.response_headers;
 
             raw_body = resp.raw_value.unwrap_or(Value::Null);
             let parsed = resp.value;
diff --git a/aimux-providers/src/aws_polly.rs b/aimux-providers/src/aws_polly.rs
index 6f42cea4..5216192f 100644
--- a/aimux-providers/src/aws_polly.rs
+++ b/aimux-providers/src/aws_polly.rs
@@ -286,10 +286,7 @@ impl SpeechModel for AwsPollySpeechModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()),
             aimux_provider_utils::create_binary_response_handler(),
@@ -297,7 +294,7 @@ impl SpeechModel for AwsPollySpeechModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let audio_bytes = resp.value.to_vec();
 
diff --git a/aimux-providers/src/azure/model.rs b/aimux-providers/src/azure/model.rs
index a680ade6..b7376abe 100644
--- a/aimux-providers/src/azure/model.rs
+++ b/aimux-providers/src/azure/model.rs
@@ -335,10 +335,7 @@ impl Provider for AzureProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         crate::openai::openai_failed_response_handler(),
diff --git a/aimux-providers/src/azure/responses.rs b/aimux-providers/src/azure/responses.rs
index 01fd62b5..17926a4f 100644
--- a/aimux-providers/src/azure/responses.rs
+++ b/aimux-providers/src/azure/responses.rs
@@ -272,25 +272,14 @@ impl LanguageModel for AzureResponsesModel {
         let provider_key = provider_key().to_string();
         let endpoint = self.endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             aimux_provider_utils::create_json_response_handler::(),
             crate::openai::openai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let raw_body = resp
             .raw_value
             .as_ref()
@@ -331,25 +320,14 @@ impl LanguageModel for AzureResponsesModel {
 
         let endpoint = self.endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             crate::openai::openai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
diff --git a/aimux-providers/src/bedrock/embedding.rs b/aimux-providers/src/bedrock/embedding.rs
index 2d3da0d4..56db97f2 100644
--- a/aimux-providers/src/bedrock/embedding.rs
+++ b/aimux-providers/src/bedrock/embedding.rs
@@ -207,10 +207,7 @@ impl EmbeddingModel for BedrockEmbeddingModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()),
             aimux_provider_utils::create_json_response_handler(),
@@ -219,7 +216,7 @@ impl EmbeddingModel for BedrockEmbeddingModel {
         .await?;
 
         // Capture response headers (needed for token count extraction).
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/bedrock/image.rs b/aimux-providers/src/bedrock/image.rs
index 21abca38..cbebdedb 100644
--- a/aimux-providers/src/bedrock/image.rs
+++ b/aimux-providers/src/bedrock/image.rs
@@ -290,17 +290,14 @@ impl ImageModel for BedrockImageModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()),
             aimux_provider_utils::create_json_response_handler::(),
             super::bedrock_failed_response_handler(),
         )
         .await?;
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let rb = resp.value;
 
diff --git a/aimux-providers/src/bedrock/mod.rs b/aimux-providers/src/bedrock/mod.rs
index 9696509f..df28bcfd 100644
--- a/aimux-providers/src/bedrock/mod.rs
+++ b/aimux-providers/src/bedrock/mod.rs
@@ -324,10 +324,7 @@ impl Provider for BedrockProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         bedrock_failed_response_handler(),
diff --git a/aimux-providers/src/bedrock/model.rs b/aimux-providers/src/bedrock/model.rs
index c62cfb8e..d7ad90f0 100644
--- a/aimux-providers/src/bedrock/model.rs
+++ b/aimux-providers/src/bedrock/model.rs
@@ -72,7 +72,7 @@ fn bedrock_event_stream_response_handler()
         Ok(aimux_provider_utils::ResponseHandlerOutput {
             value,
             raw_value: None,
-            response_headers: Some(output_headers),
+            response_headers: output_headers,
         })
     })
     .streaming()
@@ -205,10 +205,7 @@ impl LanguageModel for BedrockModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()),
             aimux_provider_utils::create_json_response_handler(),
@@ -216,7 +213,7 @@ impl LanguageModel for BedrockModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let data: BedrockConverseResponse = resp.value;
 
@@ -275,10 +272,7 @@ impl LanguageModel for BedrockModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()),
             bedrock_event_stream_response_handler(),
@@ -286,7 +280,7 @@ impl LanguageModel for BedrockModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         // Bedrock converse-stream returns binary AWS event stream format.
         // We read the full body and decode it, then emit stream parts.
diff --git a/aimux-providers/src/bedrock/reranking.rs b/aimux-providers/src/bedrock/reranking.rs
index 5396aef4..59658122 100644
--- a/aimux-providers/src/bedrock/reranking.rs
+++ b/aimux-providers/src/bedrock/reranking.rs
@@ -217,10 +217,7 @@ impl RerankingModel for BedrockRerankingModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(body_str.into_bytes(), "application/json".to_string()),
             aimux_provider_utils::create_json_response_handler::(),
@@ -229,7 +226,7 @@ impl RerankingModel for BedrockRerankingModel {
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs
index a351b3e4..ccaafba0 100644
--- a/aimux-providers/src/black_forest_labs.rs
+++ b/aimux-providers/src/black_forest_labs.rs
@@ -326,18 +326,7 @@ impl ImageModel for BlackForestLabsImageModel {
 
         // Submit
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.submit_endpoint(),
-                headers: header_list.clone(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.submit_endpoint(), header_list.clone(), options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             bfl_failed_response_handler(),
diff --git a/aimux-providers/src/cartesia.rs b/aimux-providers/src/cartesia.rs
index fc961113..c57d9613 100644
--- a/aimux-providers/src/cartesia.rs
+++ b/aimux-providers/src/cartesia.rs
@@ -200,25 +200,14 @@ impl SpeechModel for CartesiaSpeechModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers.into_iter().collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), headers.into_iter().collect(), options),
             Value::Object(body.clone()),
             aimux_provider_utils::create_binary_response_handler(),
             cartesia_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let audio_bytes = resp.value.to_vec();
 
@@ -779,25 +768,14 @@ impl TranscriptionModel for CartesiaTranscriptionModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers.into_iter().collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), headers.into_iter().collect(), options),
             HttpBody::Bytes(body_bytes, content_type),
             aimux_provider_utils::create_json_response_handler::(),
             cartesia_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let parsed = resp.value;
diff --git a/aimux-providers/src/codex.rs b/aimux-providers/src/codex.rs
index 7b026515..ccf99d3e 100644
--- a/aimux-providers/src/codex.rs
+++ b/aimux-providers/src/codex.rs
@@ -294,17 +294,7 @@ impl CodexModel {
 
         let endpoint = self.endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-                abort_signal: opts.abort_signal.clone(),
-                call_id: opts.call_id.clone(),
-                recording_context: opts.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), &opts),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             crate::openai::openai_failed_response_handler(),
@@ -312,7 +302,7 @@ impl CodexModel {
         .await
         .map_err(|e| self.map_subscription_401(e))?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let mut sse = resp.value;
         // (parsed response object, raw event payload) of the terminal event.
         let mut completed: Option<(Value, String)> = None;
@@ -407,10 +397,7 @@ impl CodexModel {
                 abort_signal: opts.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
@@ -419,7 +406,7 @@ impl CodexModel {
         .await
         .map_err(|e| self.map_subscription_401(e))?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
             Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error),
@@ -549,10 +536,7 @@ pub async fn codex_refresh_at(
             abort_signal: None,
             call_id: None,
             recording_context: None,
-            response_timeout: None,
-            validate_url: false,
-            trusted_origin: None,
-            credentialed_origin: None,
+            ..Default::default()
         },
         body,
         aimux_provider_utils::create_json_response_handler(),
diff --git a/aimux-providers/src/cohere/embedding.rs b/aimux-providers/src/cohere/embedding.rs
index 61577512..1c21a849 100644
--- a/aimux-providers/src/cohere/embedding.rs
+++ b/aimux-providers/src/cohere/embedding.rs
@@ -108,25 +108,14 @@ impl EmbeddingModel for CohereEmbeddingModel {
         header_list.push(("Content-Type".to_string(), "application/json".to_string()));
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             super::cohere_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/cohere/mod.rs b/aimux-providers/src/cohere/mod.rs
index eac6f97b..b4cc7531 100644
--- a/aimux-providers/src/cohere/mod.rs
+++ b/aimux-providers/src/cohere/mod.rs
@@ -161,10 +161,7 @@ impl Provider for CohereProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         cohere_failed_response_handler(),
diff --git a/aimux-providers/src/cohere/model.rs b/aimux-providers/src/cohere/model.rs
index 2539f141..3ea454fd 100644
--- a/aimux-providers/src/cohere/model.rs
+++ b/aimux-providers/src/cohere/model.rs
@@ -123,28 +123,21 @@ impl LanguageModel for CohereModel {
         let body = body_result.body.clone();
         let headers = self.build_headers(options.headers.as_ref());
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers
+            HttpRequest::new(
+                self.endpoint(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             body.clone(),
             aimux_provider_utils::create_json_response_handler(),
             super::cohere_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let data: ChatResponse = resp.value;
 
         // Build content array.
@@ -267,28 +260,21 @@ impl LanguageModel for CohereModel {
         let headers = self.build_headers(options.headers.as_ref());
         let endpoint = self.endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: headers
+            HttpRequest::new(
+                endpoint.clone(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             super::cohere_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
             Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error),
diff --git a/aimux-providers/src/cohere/reranking.rs b/aimux-providers/src/cohere/reranking.rs
index 146babf2..9a2c9181 100644
--- a/aimux-providers/src/cohere/reranking.rs
+++ b/aimux-providers/src/cohere/reranking.rs
@@ -151,18 +151,7 @@ impl RerankingModel for CohereRerankingModel {
         header_list.push(("Content-Type".to_string(), "application/json".to_string()));
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             body.clone(),
             aimux_provider_utils::create_json_response_handler::(),
             super::cohere_failed_response_handler(),
@@ -170,7 +159,7 @@ impl RerankingModel for CohereRerankingModel {
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/dataforseo.rs b/aimux-providers/src/dataforseo.rs
index 687057a3..b0016adb 100644
--- a/aimux-providers/src/dataforseo.rs
+++ b/aimux-providers/src/dataforseo.rs
@@ -279,10 +279,7 @@ impl SearchModel for DataforseoSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler::(),
@@ -291,7 +288,7 @@ impl SearchModel for DataforseoSearchModel {
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/deepgram.rs b/aimux-providers/src/deepgram.rs
index ff3cef25..2647b735 100644
--- a/aimux-providers/src/deepgram.rs
+++ b/aimux-providers/src/deepgram.rs
@@ -332,10 +332,7 @@ impl TranscriptionModel for DeepgramTranscriptionModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             HttpBody::Bytes(audio_bytes, options.media_type.clone()),
             aimux_provider_utils::create_json_response_handler::(),
@@ -343,7 +340,7 @@ impl TranscriptionModel for DeepgramTranscriptionModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let parsed = resp.value;
diff --git a/aimux-providers/src/elevenlabs.rs b/aimux-providers/src/elevenlabs.rs
index cb0fccdb..46f7176b 100644
--- a/aimux-providers/src/elevenlabs.rs
+++ b/aimux-providers/src/elevenlabs.rs
@@ -216,10 +216,7 @@ impl SpeechModel for ElevenLabsSpeechModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             Value::Object(body.clone()),
             aimux_provider_utils::create_binary_response_handler(),
@@ -227,7 +224,7 @@ impl SpeechModel for ElevenLabsSpeechModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let audio_bytes = resp.value.to_vec();
 
         let timestamp = chrono::Utc::now().to_rfc3339();
@@ -681,28 +678,21 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers
+            HttpRequest::new(
+                self.endpoint(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             HttpBody::Bytes(body_bytes, content_type),
             aimux_provider_utils::create_json_response_handler::(),
             elevenlabs_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let parsed = resp.value;
 
diff --git a/aimux-providers/src/exa_ai.rs b/aimux-providers/src/exa_ai.rs
index 61decc89..61b8d4ed 100644
--- a/aimux-providers/src/exa_ai.rs
+++ b/aimux-providers/src/exa_ai.rs
@@ -189,17 +189,14 @@ impl SearchModel for ExaAiSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler(),
             exa_failed_response_handler(),
         )
         .await?;
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value;
         let parsed: ExaResponse = resp.value;
 
diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs
index 37c5d549..eff9548e 100644
--- a/aimux-providers/src/fal.rs
+++ b/aimux-providers/src/fal.rs
@@ -341,7 +341,7 @@ impl TranscriptionModel for FalTranscriptionModel {
                 continue;
             };
 
-            response_headers = resp.response_headers.unwrap_or_default();
+            response_headers = resp.response_headers;
             raw_body = resp.raw_value.unwrap_or(Value::Null);
             parsed = resp.value;
             break;
@@ -586,7 +586,7 @@ impl ImageModel for FalImageModel {
         )
         .await?;
 
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let rb: Value = resp.value;
         let retries = retry::prepare_retries(
             options.max_retries,
@@ -848,7 +848,7 @@ impl VideoModel for FalVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -900,7 +900,7 @@ impl VideoModel for FalVideoModel {
             Err(error) => return Err(error),
         };
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let raw_body = resp.value;
 
         // Extract video URL from response.
diff --git a/aimux-providers/src/firecrawl.rs b/aimux-providers/src/firecrawl.rs
index 14386d59..f917f64f 100644
--- a/aimux-providers/src/firecrawl.rs
+++ b/aimux-providers/src/firecrawl.rs
@@ -196,17 +196,14 @@ impl SearchModel for FirecrawlSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler(),
             firecrawl_failed_response_handler(),
         )
         .await?;
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value;
         let parsed: FirecrawlResponse = resp.value;
 
diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs
index c4743f05..9bf05ffe 100644
--- a/aimux-providers/src/gladia.rs
+++ b/aimux-providers/src/gladia.rs
@@ -227,21 +227,14 @@ impl TranscriptionModel for GladiaTranscriptionModel {
         let (body_bytes, content_type) = form.finish();
 
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.upload_url(),
-                headers: headers
+            HttpRequest::new(
+                self.upload_url(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             HttpBody::Bytes(body_bytes, content_type),
             aimux_provider_utils::create_json_response_handler(),
             gladia_failed_response_handler(),
@@ -267,21 +260,14 @@ impl TranscriptionModel for GladiaTranscriptionModel {
         }
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.initiate_url(),
-                headers: headers
+            HttpRequest::new(
+                self.initiate_url(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             gladia_failed_response_handler(),
@@ -333,7 +319,7 @@ impl TranscriptionModel for GladiaTranscriptionModel {
                 })
                 .await?;
 
-            response_headers = resp.response_headers.unwrap_or_default();
+            response_headers = resp.response_headers;
             raw_body = resp.raw_value.unwrap_or(Value::Null);
             let parsed = resp.value;
 
diff --git a/aimux-providers/src/google/embedding.rs b/aimux-providers/src/google/embedding.rs
index 5710ac34..4eb061c4 100644
--- a/aimux-providers/src/google/embedding.rs
+++ b/aimux-providers/src/google/embedding.rs
@@ -128,10 +128,7 @@ impl EmbeddingModel for GoogleEmbeddingModel {
                     abort_signal: options.abort_signal.clone(),
                     call_id: None,
                     recording_context: None,
-                    response_timeout: None,
-                    validate_url: false,
-                    trusted_origin: None,
-                    credentialed_origin: None,
+                    ..Default::default()
                 },
                 Value::Object(body),
                 aimux_provider_utils::create_json_response_handler(),
@@ -139,7 +136,7 @@ impl EmbeddingModel for GoogleEmbeddingModel {
             )
             .await?;
 
-            let response_headers = resp.response_headers.unwrap_or_default();
+            let response_headers = resp.response_headers;
 
             let raw_value: Value = resp.value;
 
@@ -214,10 +211,7 @@ impl EmbeddingModel for GoogleEmbeddingModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
@@ -225,7 +219,7 @@ impl EmbeddingModel for GoogleEmbeddingModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs
index d3d818f3..ff44c790 100644
--- a/aimux-providers/src/google/files.rs
+++ b/aimux-providers/src/google/files.rs
@@ -211,7 +211,7 @@ impl Files for GoogleFiles {
                 Ok(aimux_provider_utils::ResponseHandlerOutput {
                     value: (),
                     raw_value: None,
-                    response_headers: Some(headers),
+                    response_headers: headers,
                 })
             }),
             super::google_failed_response_handler(),
@@ -227,8 +227,6 @@ impl Files for GoogleFiles {
 
         let upload_url = init_resp
             .response_headers
-            .as_ref()
-            .expect("header-only response handler always supplies headers")
             .get("x-goog-upload-url")
             .cloned()
             .ok_or_else(|| {
diff --git a/aimux-providers/src/google/image.rs b/aimux-providers/src/google/image.rs
index 85220385..75ff4da8 100644
--- a/aimux-providers/src/google/image.rs
+++ b/aimux-providers/src/google/image.rs
@@ -169,25 +169,14 @@ impl GoogleImageModel {
         let header_list = build_header_list(&headers);
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.predict_endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.predict_endpoint(), header_list, options),
             body,
             aimux_provider_utils::create_json_response_handler(),
             super::google_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let response_body: Value = resp.value;
 
@@ -324,25 +313,14 @@ impl GoogleImageModel {
         let header_list = build_header_list(&headers);
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.generate_content_endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.generate_content_endpoint(), header_list, options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             super::google_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let response_body: Value = resp.value;
 
diff --git a/aimux-providers/src/google/mod.rs b/aimux-providers/src/google/mod.rs
index 8bcc02d1..0a3fc576 100644
--- a/aimux-providers/src/google/mod.rs
+++ b/aimux-providers/src/google/mod.rs
@@ -237,10 +237,7 @@ impl Provider for GoogleProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         google_failed_response_handler(),
diff --git a/aimux-providers/src/google/model.rs b/aimux-providers/src/google/model.rs
index 393d4014..62a303f7 100644
--- a/aimux-providers/src/google/model.rs
+++ b/aimux-providers/src/google/model.rs
@@ -121,25 +121,18 @@ impl LanguageModel for GoogleModel {
         let (body, tool_warnings) = build_request_body_with_warnings(&self.model_id, options);
         let headers = self.build_headers(options.headers.as_ref());
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.generate_endpoint(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(
+                self.generate_endpoint(),
+                build_header_list(&headers),
+                options,
+            ),
             body.clone(),
             aimux_provider_utils::create_json_response_handler(),
             super::google_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let data: GenerateContentResponse = resp.value;
 
@@ -198,25 +191,14 @@ impl LanguageModel for GoogleModel {
         let headers = self.build_headers(options.headers.as_ref());
         let endpoint = self.stream_endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             super::google_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
diff --git a/aimux-providers/src/google/video.rs b/aimux-providers/src/google/video.rs
index 5ed56a6e..dce85c6b 100644
--- a/aimux-providers/src/google/video.rs
+++ b/aimux-providers/src/google/video.rs
@@ -125,10 +125,7 @@ impl VideoModel for GoogleVideoModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler(),
@@ -155,7 +152,7 @@ impl VideoModel for GoogleVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -182,24 +179,13 @@ impl VideoModel for GoogleVideoModel {
 
         let poll_url = self.operation_url(operation_name);
         let resp = aimux_provider_utils::get_from_api(
-            HttpRequest {
-                url: poll_url.clone(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(poll_url.clone(), header_list, options),
             aimux_provider_utils::create_json_response_handler::(),
             super::google_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let raw_body: Value = resp.value;
         // Check the in-band error first: a terminal response may carry both
diff --git a/aimux-providers/src/google_pse.rs b/aimux-providers/src/google_pse.rs
index dfcc14c9..7be6d7e6 100644
--- a/aimux-providers/src/google_pse.rs
+++ b/aimux-providers/src/google_pse.rs
@@ -246,17 +246,14 @@ impl SearchModel for GooglePseSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             aimux_provider_utils::create_json_response_handler::(),
             google_pse_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/huggingface/responses.rs b/aimux-providers/src/huggingface/responses.rs
index b06face2..b0ea9b2a 100644
--- a/aimux-providers/src/huggingface/responses.rs
+++ b/aimux-providers/src/huggingface/responses.rs
@@ -80,7 +80,7 @@ fn huggingface_successful_response_handler() -> aimux_provider_utils::ResponseHa
                     .and_then(Value::as_str)
                     .map(str::to_string),
                 response_body: Some(output.value.to_string()),
-                response_headers: output.response_headers.clone(),
+                response_headers: Some(output.response_headers.clone()),
                 ..ApiCallError::new(message, url, request_body_values)
             })));
         }
@@ -187,25 +187,14 @@ impl LanguageModel for HuggingFaceResponsesModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), build_header_list(&headers), options),
             body.clone(),
             huggingface_successful_response_handler(),
             huggingface_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let response: Value = resp.value;
 
@@ -250,25 +239,14 @@ impl LanguageModel for HuggingFaceResponsesModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), build_header_list(&headers), options),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             huggingface_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
             Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error),
diff --git a/aimux-providers/src/hume.rs b/aimux-providers/src/hume.rs
index c78e1a02..d71b4a0a 100644
--- a/aimux-providers/src/hume.rs
+++ b/aimux-providers/src/hume.rs
@@ -200,10 +200,7 @@ impl SpeechModel for HumeSpeechModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             Value::Object(body.clone()),
             aimux_provider_utils::create_binary_response_handler(),
@@ -211,7 +208,7 @@ impl SpeechModel for HumeSpeechModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let audio_bytes = resp.value.to_vec();
 
diff --git a/aimux-providers/src/jina_ai.rs b/aimux-providers/src/jina_ai.rs
index a820c128..cea88fbb 100644
--- a/aimux-providers/src/jina_ai.rs
+++ b/aimux-providers/src/jina_ai.rs
@@ -257,10 +257,7 @@ impl RerankingModel for JinaAiRerankingModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler::(),
@@ -269,7 +266,7 @@ impl RerankingModel for JinaAiRerankingModel {
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/klingai.rs b/aimux-providers/src/klingai.rs
index d8a11811..7bb6739e 100644
--- a/aimux-providers/src/klingai.rs
+++ b/aimux-providers/src/klingai.rs
@@ -293,18 +293,7 @@ impl VideoModel for KlingAIVideoModel {
         let submit_url = self.submit_endpoint(&mode);
         let request_body = Value::Object(body);
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: submit_url.clone(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(submit_url.clone(), header_list, options),
             request_body.clone(),
             aimux_provider_utils::create_json_response_handler::(),
             klingai_failed_response_handler(),
@@ -354,7 +343,7 @@ impl VideoModel for KlingAIVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -380,24 +369,13 @@ impl VideoModel for KlingAIVideoModel {
 
         let poll_url = self.poll_endpoint(mode, id, task_id);
         let resp = aimux_provider_utils::get_from_api(
-            HttpRequest {
-                url: poll_url.clone(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(poll_url.clone(), header_list, options),
             aimux_provider_utils::create_json_response_handler::(),
             klingai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let result: KlingAITaskResult = resp.value;
 
diff --git a/aimux-providers/src/linkup.rs b/aimux-providers/src/linkup.rs
index 3df59624..2cb7ab4a 100644
--- a/aimux-providers/src/linkup.rs
+++ b/aimux-providers/src/linkup.rs
@@ -21,24 +21,6 @@ use aimux_core::search_model::{
 };
 use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash};
 
-fn linkup_failed_response_handler() -> aimux_provider_utils::ResponseHandler {
-    aimux_provider_utils::create_json_error_response_handler(|data| {
-        let error = data.get("error").unwrap_or(data);
-        aimux_provider_utils::ProviderErrorParts {
-            message: error
-                .get("message")
-                .and_then(Value::as_str)
-                .unwrap_or_default()
-                .to_owned(),
-            provider_code: error
-                .get("type")
-                .or_else(|| error.get("code"))
-                .and_then(Value::as_str)
-                .map(str::to_owned),
-        }
-    })
-}
-
 /// Fixed model ID for the Linkup search model.
 const MODEL_ID: &str = "linkup-search";
 
@@ -247,19 +229,16 @@ impl SearchModel for LinkupSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler::(),
-            linkup_failed_response_handler(),
+            aimux_provider_utils::create_standard_json_error_response_handler(),
         )
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/lmnt.rs b/aimux-providers/src/lmnt.rs
index 8dd7e4c7..2c224d84 100644
--- a/aimux-providers/src/lmnt.rs
+++ b/aimux-providers/src/lmnt.rs
@@ -171,25 +171,14 @@ impl SpeechModel for LMNTSpeechModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers.into_iter().collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), headers.into_iter().collect(), options),
             Value::Object(body.clone()),
             aimux_provider_utils::create_binary_response_handler(),
             lmnt_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let audio_bytes = resp.value.to_vec();
 
diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs
index ee63671b..03120688 100644
--- a/aimux-providers/src/luma.rs
+++ b/aimux-providers/src/luma.rs
@@ -318,24 +318,13 @@ impl ImageModel for LumaImageModel {
 
         // Submit
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.generations_url(None),
-                headers: header_list.clone(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.generations_url(None), header_list.clone(), options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             luma_failed_response_handler(),
         )
         .await?;
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let submit_body: Value = resp.value;
 
         let generation_id = submit_body
@@ -357,18 +346,11 @@ impl ImageModel for LumaImageModel {
             let pr = retries
                 .retry(|| {
                     aimux_provider_utils::get_from_api(
-                        HttpRequest {
-                            url: self.generations_url(Some(&generation_id)),
-                            headers: header_list.clone(),
-
-                            abort_signal: options.abort_signal.clone(),
-                            call_id: None,
-                            recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
-                        },
+                        HttpRequest::new(
+                            self.generations_url(Some(&generation_id)),
+                            header_list.clone(),
+                            options,
+                        ),
                         aimux_provider_utils::create_json_response_handler::(),
                         luma_failed_response_handler(),
                     )
diff --git a/aimux-providers/src/mistral/embedding.rs b/aimux-providers/src/mistral/embedding.rs
index 76feea2f..6e1411de 100644
--- a/aimux-providers/src/mistral/embedding.rs
+++ b/aimux-providers/src/mistral/embedding.rs
@@ -102,25 +102,14 @@ impl EmbeddingModel for MistralEmbeddingModel {
         header_list.push(("Content-Type".to_string(), "application/json".to_string()));
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             super::mistral_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/mistral/mod.rs b/aimux-providers/src/mistral/mod.rs
index 6f8ca1c6..7b245b27 100644
--- a/aimux-providers/src/mistral/mod.rs
+++ b/aimux-providers/src/mistral/mod.rs
@@ -196,10 +196,7 @@ impl Provider for MistralProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         mistral_failed_response_handler(),
diff --git a/aimux-providers/src/mistral/model.rs b/aimux-providers/src/mistral/model.rs
index 4675a97e..547a877a 100644
--- a/aimux-providers/src/mistral/model.rs
+++ b/aimux-providers/src/mistral/model.rs
@@ -225,28 +225,21 @@ impl LanguageModel for MistralModel {
         let body = build_request_body(&self.model_id, options, false);
         let headers = self.build_headers(options.headers.as_ref());
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers
+            HttpRequest::new(
+                self.endpoint(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             body.clone(),
             aimux_provider_utils::create_json_response_handler(),
             super::mistral_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let data: ChatCompletionResponse = resp.value;
 
         let choice =
@@ -329,28 +322,21 @@ impl LanguageModel for MistralModel {
         let body = build_request_body(&self.model_id, options, true);
         let headers = self.build_headers(options.headers.as_ref());
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: headers
+            HttpRequest::new(
+                self.endpoint(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             super::mistral_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let mut sse_stream = resp.value;
 
         // Peek at the first SSE event to detect early errors.
diff --git a/aimux-providers/src/open_responses.rs b/aimux-providers/src/open_responses.rs
index 95eff627..1675befc 100644
--- a/aimux-providers/src/open_responses.rs
+++ b/aimux-providers/src/open_responses.rs
@@ -73,7 +73,7 @@ fn open_responses_successful_response_handler() -> aimux_provider_utils::Respons
                     .and_then(Value::as_str)
                     .map(str::to_string),
                 response_body: Some(output.value.to_string()),
-                response_headers: output.response_headers.clone(),
+                response_headers: Some(output.response_headers.clone()),
                 ..ApiCallError::new(message, url, request_body_values)
             })));
         }
@@ -315,25 +315,18 @@ impl LanguageModel for OpenResponsesModel {
             build_request_body(&self.model_id, options, &self.config.provider_options_name);
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.config.url.clone(),
-                headers: headers.into_iter().collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(
+                self.config.url.clone(),
+                headers.into_iter().collect(),
+                options,
+            ),
             body.clone(),
             open_responses_successful_response_handler(),
             open_responses_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw: Value = resp.value;
 
@@ -471,25 +464,18 @@ impl LanguageModel for OpenResponsesModel {
         };
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.config.url.clone(),
-                headers: headers.into_iter().collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(
+                self.config.url.clone(),
+                headers.into_iter().collect(),
+                options,
+            ),
             stream_body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             open_responses_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let mut sse_stream = resp.value;
 
diff --git a/aimux-providers/src/openai/embedding.rs b/aimux-providers/src/openai/embedding.rs
index 4ae9be4b..14def1ed 100644
--- a/aimux-providers/src/openai/embedding.rs
+++ b/aimux-providers/src/openai/embedding.rs
@@ -114,18 +114,7 @@ impl EmbeddingModel for OpenAIEmbeddingModel {
         header_list.push(("Content-Type".to_string(), "application/json".to_string()));
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             super::openai_failed_response_handler(),
@@ -134,7 +123,7 @@ impl EmbeddingModel for OpenAIEmbeddingModel {
 
         // `send` retries 408/409/429/5xx and returns an error for non-2xx, so an `Ok`
         // response here is guaranteed to be 2xx — no manual is_success() check.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs
index d1060b6a..e2aee9c5 100644
--- a/aimux-providers/src/openai/files.rs
+++ b/aimux-providers/src/openai/files.rs
@@ -168,18 +168,7 @@ impl Files for OpenAIFiles {
         // sets `Content-Type` from it, so it is intentionally not added to the
         // header list above.
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             HttpBody::Bytes(body, content_type),
             aimux_provider_utils::create_json_response_handler(),
             super::openai_failed_response_handler(),
diff --git a/aimux-providers/src/openai/image.rs b/aimux-providers/src/openai/image.rs
index 6ca993c6..ec157c4c 100644
--- a/aimux-providers/src/openai/image.rs
+++ b/aimux-providers/src/openai/image.rs
@@ -155,18 +155,7 @@ impl ImageModel for OpenAIImageModel {
                 build_edit_multipart(&self.model_id, options, &openai_options);
 
             let resp = aimux_provider_utils::post_to_api(
-                HttpRequest {
-                    url: self.edits_endpoint(),
-                    headers: build_header_list(&headers),
-
-                    abort_signal: options.abort_signal.clone(),
-                    call_id: None,
-                    recording_context: None,
-                    response_timeout: None,
-                    validate_url: false,
-                    trusted_origin: None,
-                    credentialed_origin: None,
-                },
+                HttpRequest::new(self.edits_endpoint(), build_header_list(&headers), options),
                 HttpBody::Bytes(form_body, content_type),
                 aimux_provider_utils::create_json_response_handler(),
                 super::openai_failed_response_handler(),
@@ -174,25 +163,18 @@ impl ImageModel for OpenAIImageModel {
             .await?;
 
             let val: Value = resp.value;
-            (val, resp.response_headers.unwrap_or_default(), None)
+            (val, resp.response_headers, None)
         } else {
             // ── Generation path: JSON body ──
             let openai_options = parse_generation_provider_options(&options.provider_options);
             let body = build_generation_body(&self.model_id, options, &openai_options);
 
             let resp = aimux_provider_utils::post_json_to_api(
-                HttpRequest {
-                    url: self.generations_endpoint(),
-                    headers: build_header_list(&headers),
-
-                    abort_signal: options.abort_signal.clone(),
-                    call_id: None,
-                    recording_context: None,
-                    response_timeout: None,
-                    validate_url: false,
-                    trusted_origin: None,
-                    credentialed_origin: None,
-                },
+                HttpRequest::new(
+                    self.generations_endpoint(),
+                    build_header_list(&headers),
+                    options,
+                ),
                 Value::Object(body.clone()),
                 aimux_provider_utils::create_json_response_handler(),
                 super::openai_failed_response_handler(),
@@ -200,11 +182,7 @@ impl ImageModel for OpenAIImageModel {
             .await?;
 
             let val: Value = resp.value;
-            (
-                val,
-                resp.response_headers.unwrap_or_default(),
-                Some(Value::Object(body)),
-            )
+            (val, resp.response_headers, Some(Value::Object(body)))
         };
 
         let images = extract_images(&body_value);
diff --git a/aimux-providers/src/openai/model.rs b/aimux-providers/src/openai/model.rs
index 82255536..59837557 100644
--- a/aimux-providers/src/openai/model.rs
+++ b/aimux-providers/src/openai/model.rs
@@ -279,25 +279,14 @@ pub async fn execute_generate(
     let body = request_result.body;
 
     let resp = aimux_provider_utils::post_json_to_api(
-        HttpRequest {
-            url: endpoint.to_string(),
-            headers: build_header_list(headers),
-
-            abort_signal: options.abort_signal.clone(),
-            call_id: options.call_id.clone(),
-            recording_context: options.recording_context.clone(),
-            response_timeout: None,
-            validate_url: false,
-            trusted_origin: None,
-            credentialed_origin: None,
-        },
+        HttpRequest::new(endpoint.to_string(), build_header_list(headers), options),
         body.clone(),
         aimux_provider_utils::create_json_response_handler::(),
         super::openai_failed_response_handler(),
     )
     .await?;
 
-    let response_headers = resp.response_headers.unwrap_or_default();
+    let response_headers = resp.response_headers;
 
     // Parse the raw body once: the `Value` keeps the provider's original
     // fields (incl. vendor-specific usage fields) for `Usage.raw` (M10).
@@ -445,25 +434,14 @@ pub async fn execute_stream(
     let RequestBodyResult { body, warnings } = request_result;
 
     let resp = aimux_provider_utils::post_json_to_api(
-        HttpRequest {
-            url: endpoint.to_string(),
-            headers: build_header_list(headers),
-
-            abort_signal: options.abort_signal.clone(),
-            call_id: options.call_id.clone(),
-            recording_context: options.recording_context.clone(),
-            response_timeout: None,
-            validate_url: false,
-            trusted_origin: None,
-            credentialed_origin: None,
-        },
+        HttpRequest::new(endpoint.to_string(), build_header_list(headers), options),
         body.clone(),
         aimux_provider_utils::create_event_source_response_handler::(),
         super::openai_failed_response_handler(),
     )
     .await?;
 
-    let response_headers = resp.response_headers.unwrap_or_default();
+    let response_headers = resp.response_headers;
 
     let mut sse_stream = resp.value;
 
@@ -928,10 +906,7 @@ pub async fn execute_list_models(
                     abort_signal: None,
                     call_id: None,
                     recording_context: None,
-                    response_timeout: None,
-                    validate_url: false,
-                    trusted_origin: None,
-                    credentialed_origin: None,
+                    ..Default::default()
                 },
                 aimux_provider_utils::create_json_response_handler(),
                 super::openai_failed_response_handler(),
diff --git a/aimux-providers/src/openai/responses/mod.rs b/aimux-providers/src/openai/responses/mod.rs
index 809964bf..91afa897 100644
--- a/aimux-providers/src/openai/responses/mod.rs
+++ b/aimux-providers/src/openai/responses/mod.rs
@@ -127,25 +127,14 @@ impl LanguageModel for OpenAIResponsesModel {
 
         let endpoint = self.endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             aimux_provider_utils::create_json_response_handler::(),
             super::openai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let raw_body = resp
             .raw_value
             .as_ref()
@@ -183,25 +172,14 @@ impl LanguageModel for OpenAIResponsesModel {
 
         let endpoint = self.endpoint();
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
             super::openai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
diff --git a/aimux-providers/src/openai/speech.rs b/aimux-providers/src/openai/speech.rs
index 87a073af..d1d30241 100644
--- a/aimux-providers/src/openai/speech.rs
+++ b/aimux-providers/src/openai/speech.rs
@@ -98,18 +98,7 @@ impl SpeechModel for OpenAISpeechModel {
             .collect();
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             Value::Object(body.clone()),
             aimux_provider_utils::create_binary_response_handler(),
             super::openai_failed_response_handler(),
@@ -118,7 +107,7 @@ impl SpeechModel for OpenAISpeechModel {
 
         // send() returns Ok only for 2xx responses; non-2xx (incl. 408/409/429/5xx
         // after exhausting retries) is mapped to an AiMuxError internally.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let audio_bytes = resp.value.to_vec();
 
diff --git a/aimux-providers/src/openai/transcription.rs b/aimux-providers/src/openai/transcription.rs
index 25258292..b0ca5777 100644
--- a/aimux-providers/src/openai/transcription.rs
+++ b/aimux-providers/src/openai/transcription.rs
@@ -344,25 +344,14 @@ impl TranscriptionModel for OpenAITranscriptionModel {
         let header_list: Vec<(String, String)> = headers.into_iter().collect();
 
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             HttpBody::Bytes(body_bytes, content_type),
             aimux_provider_utils::create_json_response_handler::(),
             super::openai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let parsed = resp.value;
diff --git a/aimux-providers/src/parallel_ai.rs b/aimux-providers/src/parallel_ai.rs
index 6376b762..41c1251e 100644
--- a/aimux-providers/src/parallel_ai.rs
+++ b/aimux-providers/src/parallel_ai.rs
@@ -21,24 +21,6 @@ use aimux_core::search_model::{
 };
 use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash};
 
-fn parallel_ai_failed_response_handler() -> aimux_provider_utils::ResponseHandler {
-    aimux_provider_utils::create_json_error_response_handler(|data| {
-        let error = data.get("error").unwrap_or(data);
-        aimux_provider_utils::ProviderErrorParts {
-            message: error
-                .get("message")
-                .and_then(Value::as_str)
-                .unwrap_or_default()
-                .to_owned(),
-            provider_code: error
-                .get("type")
-                .or_else(|| error.get("code"))
-                .and_then(Value::as_str)
-                .map(str::to_owned),
-        }
-    })
-}
-
 /// Fixed model ID for the Parallel AI search model.
 const MODEL_ID: &str = "parallel-search";
 
@@ -206,18 +188,15 @@ impl SearchModel for ParallelAiSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler::(),
-            parallel_ai_failed_response_handler(),
+            aimux_provider_utils::create_standard_json_error_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/prodia.rs b/aimux-providers/src/prodia.rs
index 216f01d8..fd0cba85 100644
--- a/aimux-providers/src/prodia.rs
+++ b/aimux-providers/src/prodia.rs
@@ -258,25 +258,14 @@ impl ImageModel for ProdiaImageModel {
         let header_list: Vec<(String, String)> = headers.into_iter().collect();
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             body,
             aimux_provider_utils::create_binary_response_handler(),
             prodia_failed_response_handler(),
         )
         .await?;
 
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let content_type = rh.get("content-type").cloned().unwrap_or_default();
 
         // Extract boundary
@@ -411,18 +400,11 @@ impl VideoModel for ProdiaVideoModel {
 
         // Submit job.
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: format!("{}/job", self.config.base_url),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(
+                format!("{}/job", self.config.base_url),
+                header_list,
+                options,
+            ),
             body,
             aimux_provider_utils::create_json_response_handler(),
             prodia_failed_response_handler(),
@@ -448,7 +430,7 @@ impl VideoModel for ProdiaVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -472,24 +454,13 @@ impl VideoModel for ProdiaVideoModel {
         let poll_url = format!("{}/job/{}", self.config.base_url, job_id);
 
         let resp = aimux_provider_utils::get_from_api(
-            HttpRequest {
-                url: poll_url.clone(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(poll_url.clone(), header_list, options),
             aimux_provider_utils::create_json_response_handler::(),
             prodia_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let raw_body = resp.value;
         let status_str = raw_body
diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs
index c6665781..d85a9b1d 100644
--- a/aimux-providers/src/recraft.rs
+++ b/aimux-providers/src/recraft.rs
@@ -337,7 +337,7 @@ impl ImageModel for RecraftImageModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let value: Value = resp.value;
         let retries = retry::prepare_retries(
             options.max_retries,
diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs
index 72f51635..1bcb1379 100644
--- a/aimux-providers/src/replicate.rs
+++ b/aimux-providers/src/replicate.rs
@@ -316,7 +316,7 @@ impl ImageModel for ReplicateImageModel {
         )
         .await?;
 
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let rb: Value = resp.value;
         let retries = retry::prepare_retries(
             options.max_retries,
@@ -509,7 +509,7 @@ impl VideoModel for ReplicateVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -550,7 +550,7 @@ impl VideoModel for ReplicateVideoModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let raw_body = resp.value;
         let status_str = raw_body
diff --git a/aimux-providers/src/revai.rs b/aimux-providers/src/revai.rs
index 9dba8327..8e2a3c14 100644
--- a/aimux-providers/src/revai.rs
+++ b/aimux-providers/src/revai.rs
@@ -229,21 +229,14 @@ impl TranscriptionModel for RevaiTranscriptionModel {
 
         // Submit job.
         let resp = aimux_provider_utils::post_form_data_to_api(
-            HttpRequest {
-                url: submit_url.clone(),
-                headers: headers
+            HttpRequest::new(
+                submit_url.clone(),
+                headers
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
                     .collect(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+                options,
+            ),
             form,
             aimux_provider_utils::create_json_response_handler::(),
             revai_failed_response_handler(),
@@ -291,21 +284,14 @@ impl TranscriptionModel for RevaiTranscriptionModel {
             let resp = retries
                 .retry(|| {
                     aimux_provider_utils::get_from_api(
-                        HttpRequest {
-                            url: poll_url.clone(),
-                            headers: headers
+                        HttpRequest::new(
+                            poll_url.clone(),
+                            headers
                                 .iter()
                                 .map(|(k, v)| (k.clone(), v.clone()))
                                 .collect(),
-
-                            abort_signal: options.abort_signal.clone(),
-                            call_id: None,
-                            recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
-                        },
+                            options,
+                        ),
                         aimux_provider_utils::create_json_response_handler::(),
                         revai_failed_response_handler(),
                     )
@@ -334,28 +320,21 @@ impl TranscriptionModel for RevaiTranscriptionModel {
         let resp = retries
             .retry(|| {
                 aimux_provider_utils::get_from_api(
-                    HttpRequest {
-                        url: self.transcript_url(&job_id),
-                        headers: headers
+                    HttpRequest::new(
+                        self.transcript_url(&job_id),
+                        headers
                             .iter()
                             .map(|(k, v)| (k.clone(), v.clone()))
                             .collect(),
-
-                        abort_signal: options.abort_signal.clone(),
-                        call_id: None,
-                        recording_context: None,
-                        response_timeout: None,
-                        validate_url: false,
-                        trusted_origin: None,
-                        credentialed_origin: None,
-                    },
+                        options,
+                    ),
                     aimux_provider_utils::create_json_response_handler::(),
                     revai_failed_response_handler(),
                 )
             })
             .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let parsed = resp.value;
 
diff --git a/aimux-providers/src/runwayml.rs b/aimux-providers/src/runwayml.rs
index 77bdcbb6..be5e8bc5 100644
--- a/aimux-providers/src/runwayml.rs
+++ b/aimux-providers/src/runwayml.rs
@@ -268,18 +268,7 @@ impl VideoModel for RunwaymlVideoModel {
 
         // Submit the task.
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: submit_url,
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(submit_url, header_list, options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             runwayml_failed_response_handler(),
@@ -296,7 +285,7 @@ impl VideoModel for RunwaymlVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -320,24 +309,13 @@ impl VideoModel for RunwaymlVideoModel {
         let poll_url = format!("{}/v1/tasks/{}", self.config.base_url, task_id);
 
         let resp = aimux_provider_utils::get_from_api(
-            HttpRequest {
-                url: poll_url.clone(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(poll_url.clone(), header_list, options),
             aimux_provider_utils::create_json_response_handler::(),
             runwayml_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.clone().unwrap_or_default();
+        let response_headers = resp.response_headers.clone();
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let task_details: RunwaymlTaskDetailsResponse = resp.value;
 
diff --git a/aimux-providers/src/searxng.rs b/aimux-providers/src/searxng.rs
index 5ab3605e..f4e58db8 100644
--- a/aimux-providers/src/searxng.rs
+++ b/aimux-providers/src/searxng.rs
@@ -176,16 +176,13 @@ impl SearchModel for SearxngSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             aimux_provider_utils::create_json_response_handler::(),
             aimux_provider_utils::create_status_code_error_response_handler(),
         )
         .await?;
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/serper.rs b/aimux-providers/src/serper.rs
index 0d7a3edb..5c87a912 100644
--- a/aimux-providers/src/serper.rs
+++ b/aimux-providers/src/serper.rs
@@ -17,24 +17,6 @@ use aimux_core::search_model::{
 };
 use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash};
 
-fn serper_failed_response_handler() -> aimux_provider_utils::ResponseHandler {
-    aimux_provider_utils::create_json_error_response_handler(|data| {
-        let error = data.get("error").unwrap_or(data);
-        aimux_provider_utils::ProviderErrorParts {
-            message: error
-                .get("message")
-                .and_then(Value::as_str)
-                .unwrap_or_default()
-                .to_owned(),
-            provider_code: error
-                .get("type")
-                .or_else(|| error.get("code"))
-                .and_then(Value::as_str)
-                .map(str::to_owned),
-        }
-    })
-}
-
 const MODEL_ID: &str = "serper-search";
 
 /// Configuration for the Serper provider.
@@ -184,17 +166,14 @@ impl SearchModel for SerperSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler(),
-            serper_failed_response_handler(),
+            aimux_provider_utils::create_standard_json_error_response_handler(),
         )
         .await?;
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value;
         let parsed: SerperResponse = resp.value;
 
diff --git a/aimux-providers/src/stability.rs b/aimux-providers/src/stability.rs
index 2a0e020b..d727cba9 100644
--- a/aimux-providers/src/stability.rs
+++ b/aimux-providers/src/stability.rs
@@ -90,7 +90,7 @@ fn stability_successful_response_handler() -> aimux_provider_utils::ResponseHand
         Ok(aimux_provider_utils::ResponseHandlerOutput {
             value,
             raw_value: None,
-            response_headers: Some(headers),
+            response_headers: headers,
         })
     })
 }
@@ -327,25 +327,14 @@ impl ImageModel for StabilityImageModel {
         let header_list: Vec<(String, String)> = headers.into_iter().collect();
 
         let resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             HttpBody::Bytes(body_bytes, content_type),
             stability_successful_response_handler(),
             stability_failed_response_handler(),
         )
         .await?;
 
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let image_bytes = match resp.value {
             StabilitySuccess::Json(v) => {
                 let b64 = v.get("image").and_then(|i| i.as_str()).ok_or_else(|| {
diff --git a/aimux-providers/src/tavily.rs b/aimux-providers/src/tavily.rs
index feefeebb..a01ab7ec 100644
--- a/aimux-providers/src/tavily.rs
+++ b/aimux-providers/src/tavily.rs
@@ -199,17 +199,14 @@ impl SearchModel for TavilySearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body,
             aimux_provider_utils::create_json_response_handler(),
             tavily_failed_response_handler(),
         )
         .await?;
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value;
         let parsed: TavilyResponse = resp.value;
 
diff --git a/aimux-providers/src/tinyfish.rs b/aimux-providers/src/tinyfish.rs
index bd8c721b..06bb368c 100644
--- a/aimux-providers/src/tinyfish.rs
+++ b/aimux-providers/src/tinyfish.rs
@@ -22,24 +22,6 @@ use aimux_core::shared::SharedHeaders;
 
 use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash};
 
-fn tinyfish_failed_response_handler() -> aimux_provider_utils::ResponseHandler {
-    aimux_provider_utils::create_json_error_response_handler(|data| {
-        let error = data.get("error").unwrap_or(data);
-        aimux_provider_utils::ProviderErrorParts {
-            message: error
-                .get("message")
-                .and_then(Value::as_str)
-                .unwrap_or_default()
-                .to_owned(),
-            provider_code: error
-                .get("type")
-                .or_else(|| error.get("code"))
-                .and_then(Value::as_str)
-                .map(str::to_owned),
-        }
-    })
-}
-
 /// Provider canonical name.
 const PROVIDER_NAME: &str = "tinyfish";
 
@@ -233,18 +215,15 @@ impl SearchModel for TinyfishSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             aimux_provider_utils::create_json_response_handler::(),
-            tinyfish_failed_response_handler(),
+            aimux_provider_utils::create_standard_json_error_response_handler(),
         )
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/vertex/anthropic_model.rs b/aimux-providers/src/vertex/anthropic_model.rs
index b49d099c..afcccf47 100644
--- a/aimux-providers/src/vertex/anthropic_model.rs
+++ b/aimux-providers/src/vertex/anthropic_model.rs
@@ -184,10 +184,7 @@ impl LanguageModel for VertexAnthropicModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body.clone(),
             aimux_provider_utils::create_json_response_handler(),
@@ -245,10 +242,7 @@ impl LanguageModel for VertexAnthropicModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
@@ -256,7 +250,7 @@ impl LanguageModel for VertexAnthropicModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let mut sse_stream = resp.value;
         let first_event = match sse_stream.next().await {
             Some(Err(error @ AiMuxError::ApiCall(_))) => return Err(error),
diff --git a/aimux-providers/src/vertex/embedding.rs b/aimux-providers/src/vertex/embedding.rs
index a69396ef..5ebd732b 100644
--- a/aimux-providers/src/vertex/embedding.rs
+++ b/aimux-providers/src/vertex/embedding.rs
@@ -149,10 +149,7 @@ impl EmbeddingModel for VertexEmbeddingModel {
                     abort_signal: options.abort_signal.clone(),
                     call_id: None,
                     recording_context: None,
-                    response_timeout: None,
-                    validate_url: false,
-                    trusted_origin: None,
-                    credentialed_origin: None,
+                    ..Default::default()
                 },
                 Value::Object(body),
                 aimux_provider_utils::create_json_response_handler(),
@@ -160,7 +157,7 @@ impl EmbeddingModel for VertexEmbeddingModel {
             )
             .await?;
 
-            let response_headers = resp.response_headers.unwrap_or_default();
+            let response_headers = resp.response_headers;
 
             let raw_value: Value = resp.value;
 
@@ -234,10 +231,7 @@ impl EmbeddingModel for VertexEmbeddingModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
@@ -245,7 +239,7 @@ impl EmbeddingModel for VertexEmbeddingModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/vertex/image.rs b/aimux-providers/src/vertex/image.rs
index 83180492..f93b5924 100644
--- a/aimux-providers/src/vertex/image.rs
+++ b/aimux-providers/src/vertex/image.rs
@@ -194,24 +194,13 @@ impl VertexImageModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.predict_endpoint(),
-                headers,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.predict_endpoint(), headers, options),
             body,
             aimux_provider_utils::create_json_response_handler(),
             crate::google::google_failed_response_handler(),
         )
         .await?;
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let rb: Value = resp.value;
 
         let images: Vec = rb
@@ -342,24 +331,13 @@ impl VertexImageModel {
         let headers = self.build_headers(options.headers.as_ref());
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.generate_content_endpoint(),
-                headers,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.generate_content_endpoint(), headers, options),
             body,
             aimux_provider_utils::create_json_response_handler(),
             google_failed_response_handler(),
         )
         .await?;
-        let rh = resp.response_headers.unwrap_or_default();
+        let rh = resp.response_headers;
         let rb: Value = resp.value;
 
         let mut images: Vec = Vec::new();
diff --git a/aimux-providers/src/vertex/mod.rs b/aimux-providers/src/vertex/mod.rs
index dcd457d3..9c828aa7 100644
--- a/aimux-providers/src/vertex/mod.rs
+++ b/aimux-providers/src/vertex/mod.rs
@@ -384,10 +384,7 @@ impl Provider for VertexProvider {
                             abort_signal: None,
                             call_id: None,
                             recording_context: None,
-                            response_timeout: None,
-                            validate_url: false,
-                            trusted_origin: None,
-                            credentialed_origin: None,
+                            ..Default::default()
                         },
                         aimux_provider_utils::create_json_response_handler(),
                         crate::google::google_failed_response_handler(),
diff --git a/aimux-providers/src/vertex/model.rs b/aimux-providers/src/vertex/model.rs
index a9263336..2b8b5d91 100644
--- a/aimux-providers/src/vertex/model.rs
+++ b/aimux-providers/src/vertex/model.rs
@@ -162,10 +162,7 @@ impl LanguageModel for VertexModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body.clone(),
             aimux_provider_utils::create_json_response_handler(),
@@ -173,7 +170,7 @@ impl LanguageModel for VertexModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let data: GenerateContentResponse = resp.value;
 
@@ -239,10 +236,7 @@ impl LanguageModel for VertexModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: options.call_id.clone(),
                 recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             body.clone(),
             aimux_provider_utils::create_event_source_response_handler::(),
@@ -250,7 +244,7 @@ impl LanguageModel for VertexModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         // Same source as the non-stream path: the response `Date` header.
         let response_timestamp = response_headers.get("date").cloned();
 
diff --git a/aimux-providers/src/vertex/transcription.rs b/aimux-providers/src/vertex/transcription.rs
index 7edfeb96..14f8e8a3 100644
--- a/aimux-providers/src/vertex/transcription.rs
+++ b/aimux-providers/src/vertex/transcription.rs
@@ -264,10 +264,7 @@ impl TranscriptionModel for VertexTranscriptionModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             request_body.clone(),
             aimux_provider_utils::create_json_response_handler::(),
@@ -275,7 +272,7 @@ impl TranscriptionModel for VertexTranscriptionModel {
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let parsed = resp.value;
diff --git a/aimux-providers/src/vertex/video.rs b/aimux-providers/src/vertex/video.rs
index c76c2acb..a5139f4c 100644
--- a/aimux-providers/src/vertex/video.rs
+++ b/aimux-providers/src/vertex/video.rs
@@ -158,18 +158,7 @@ impl VideoModel for VertexVideoModel {
             .collect();
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.predict_url(),
-                headers: header_list.clone(),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.predict_url(), header_list.clone(), options),
             body,
             aimux_provider_utils::create_json_response_handler(),
             crate::google::google_failed_response_handler(),
@@ -195,7 +184,7 @@ impl VideoModel for VertexVideoModel {
             response: VideoResponse {
                 timestamp: Some(chrono::Utc::now().to_rfc3339()),
                 model_id: Some(self.model_id.clone()),
-                headers: response_headers,
+                headers: Some(response_headers),
             },
         })
     }
@@ -222,24 +211,13 @@ impl VideoModel for VertexVideoModel {
 
         let poll_url = self.operation_url(operation_name);
         let resp = aimux_provider_utils::get_from_api(
-            HttpRequest {
-                url: poll_url.clone(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(poll_url.clone(), header_list, options),
             aimux_provider_utils::create_json_response_handler::(),
             crate::google::google_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
         let response_body = resp.raw_value.as_ref().map(ToString::to_string);
         let raw_body: Value = resp.value;
         // Check the in-band error first: a terminal response may carry both
diff --git a/aimux-providers/src/voyage/embedding.rs b/aimux-providers/src/voyage/embedding.rs
index a6cbc365..1e366693 100644
--- a/aimux-providers/src/voyage/embedding.rs
+++ b/aimux-providers/src/voyage/embedding.rs
@@ -105,26 +105,14 @@ impl EmbeddingModel for VoyageEmbeddingModel {
         // sets `Content-Type: application/json`, so it is intentionally not
         // added to the header list above.
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             Value::Object(body),
             aimux_provider_utils::create_json_response_handler(),
             voyage_failed_response_handler(),
         )
         .await?;
 
-        let response_headers: HashMap =
-            resp.response_headers.unwrap_or_default().clone();
+        let response_headers: HashMap = resp.response_headers.clone();
 
         let raw_value: Value = resp.value;
 
diff --git a/aimux-providers/src/voyage/reranking.rs b/aimux-providers/src/voyage/reranking.rs
index 8f2cb436..01369286 100644
--- a/aimux-providers/src/voyage/reranking.rs
+++ b/aimux-providers/src/voyage/reranking.rs
@@ -148,18 +148,7 @@ impl RerankingModel for VoyageRerankingModel {
         header_list.push(("Content-Type".to_string(), "application/json".to_string()));
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: header_list,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), header_list, options),
             body.clone(),
             aimux_provider_utils::create_json_response_handler::(),
             voyage_failed_response_handler(),
@@ -167,7 +156,7 @@ impl RerankingModel for VoyageRerankingModel {
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs
index fdd623c0..0b180ece 100644
--- a/aimux-providers/src/xai/mod.rs
+++ b/aimux-providers/src/xai/mod.rs
@@ -109,7 +109,7 @@ where
         let output = aimux_provider_utils::create_json_response_handler::()
             .handle(input)
             .await?;
-        let headers = output.response_headers.clone().unwrap_or_default();
+        let headers = output.response_headers.clone();
         let raw = output.value;
         if let Some(error) = raw.get("error").filter(|value| !value.is_null()) {
             let message = error
@@ -180,7 +180,7 @@ where
         let output = aimux_provider_utils::create_json_response_handler::()
             .handle(input)
             .await?;
-        let headers = output.response_headers.unwrap_or_default();
+        let headers = output.response_headers;
         let raw = output.value;
         let error = raw.get("error").filter(|value| !value.is_null());
         let message = error
diff --git a/aimux-providers/src/xai/model.rs b/aimux-providers/src/xai/model.rs
index cc0c2b4b..5799be56 100644
--- a/aimux-providers/src/xai/model.rs
+++ b/aimux-providers/src/xai/model.rs
@@ -113,18 +113,7 @@ impl LanguageModel for XaiModel {
         let body = request_result.body;
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.endpoint(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(self.endpoint(), build_header_list(&headers), options),
             body.clone(),
             super::xai_successful_response_handler::(),
             super::xai_failed_response_handler(),
@@ -132,7 +121,7 @@ impl LanguageModel for XaiModel {
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let _raw_value = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;
@@ -262,25 +251,14 @@ impl LanguageModel for XaiModel {
         let endpoint = self.endpoint();
 
         let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: endpoint.clone(),
-                headers: build_header_list(&headers),
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: options.call_id.clone(),
-                recording_context: options.recording_context.clone(),
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
+            HttpRequest::new(endpoint.clone(), build_header_list(&headers), options),
             body.clone(),
             super::xai_event_source_response_handler::(),
             super::xai_failed_response_handler(),
         )
         .await?;
 
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let mut sse_stream = resp.value;
 
diff --git a/aimux-providers/src/you_com.rs b/aimux-providers/src/you_com.rs
index 19dc8f06..dcab66c3 100644
--- a/aimux-providers/src/you_com.rs
+++ b/aimux-providers/src/you_com.rs
@@ -23,24 +23,6 @@ use aimux_core::shared::SharedHeaders;
 
 use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash};
 
-fn you_com_failed_response_handler() -> aimux_provider_utils::ResponseHandler {
-    aimux_provider_utils::create_json_error_response_handler(|data| {
-        let error = data.get("error").unwrap_or(data);
-        aimux_provider_utils::ProviderErrorParts {
-            message: error
-                .get("message")
-                .and_then(Value::as_str)
-                .unwrap_or_default()
-                .to_owned(),
-            provider_code: error
-                .get("type")
-                .or_else(|| error.get("code"))
-                .and_then(Value::as_str)
-                .map(str::to_owned),
-        }
-    })
-}
-
 /// Provider canonical name.
 const PROVIDER_NAME: &str = "you_com";
 
@@ -235,18 +217,15 @@ impl SearchModel for YouComSearchModel {
                 abort_signal: options.abort_signal.clone(),
                 call_id: None,
                 recording_context: None,
-                response_timeout: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
+                ..Default::default()
             },
             aimux_provider_utils::create_json_response_handler::(),
-            you_com_failed_response_handler(),
+            aimux_provider_utils::create_standard_json_error_response_handler(),
         )
         .await?;
 
         // Capture response headers.
-        let response_headers = resp.response_headers.unwrap_or_default();
+        let response_headers = resp.response_headers;
 
         let raw_body = resp.raw_value.unwrap_or(Value::Null);
         let data = resp.value;

From 01da456bedec0cc0fee336cf1f26de4ebce4aedc Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Thu, 3 Sep 2026 23:22:00 +0800
Subject: [PATCH 09/24] fix(core): deep-merge start/completion provider
 metadata for video generation

merge_provider_metadata() unions object fields within a shared provider key
instead of merged.entry(provider).or_insert(value), which silently dropped
every start-phase field whenever the completion phase reported metadata
under the same provider key. On a field collision the completion value wins;
this is documented on the helper.

Covered by merge_provider_metadata_unions_same_provider_key_across_phases in
aimux-core/src/video_model.rs.
---
 aimux-core/src/video_model.rs | 66 +++++++++++++++++++++++++++++++----
 1 file changed, 59 insertions(+), 7 deletions(-)

diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs
index 7a9c267a..342789c7 100644
--- a/aimux-core/src/video_model.rs
+++ b/aimux-core/src/video_model.rs
@@ -429,18 +429,47 @@ async fn start_and_poll(
                 let mut warnings = start.warnings;
                 warnings.append(&mut result.warnings);
                 result.warnings = warnings;
-                if let Some(start_meta) = start.provider_metadata {
-                    let merged = result
-                        .provider_metadata
-                        .get_or_insert_with(Default::default);
-                    for (provider, value) in start_meta {
-                        merged.entry(provider).or_insert(value);
+                result.provider_metadata =
+                    merge_provider_metadata(start.provider_metadata, result.provider_metadata);
+                return Ok(result);
+            }
+        }
+    }
+}
+
+/// Merge two phases' provider metadata one level deep.
+///
+/// A plain `HashMap::entry().or_insert()` at the provider-key level drops
+/// every field the other phase set once both phases report the *same*
+/// provider key — e.g. a start call's `job_id` disappearing once the
+/// completion call also reports `x_provider` metadata. Instead, when both
+/// sides have an object for the same provider key, their fields are unioned;
+/// on a field collision the later (`completion`/`b`) value wins. When a
+/// provider key appears on only one side, or either side's value for it is
+/// not a JSON object, that side's value is used as-is — there is nothing to
+/// union.
+fn merge_provider_metadata(
+    a: Option,
+    b: Option,
+) -> Option {
+    let Some(mut merged) = a else { return b };
+    let Some(b) = b else { return Some(merged) };
+    for (provider, b_value) in b {
+        match merged.get_mut(&provider) {
+            Some(a_value) => match (a_value.as_object_mut(), b_value.as_object()) {
+                (Some(a_obj), Some(b_obj)) => {
+                    for (field, value) in b_obj {
+                        a_obj.insert(field.clone(), value.clone());
                     }
                 }
-                return Ok(result);
+                _ => *a_value = b_value,
+            },
+            None => {
+                merged.insert(provider, b_value);
             }
         }
     }
+    Some(merged)
 }
 
 #[cfg(test)]
@@ -683,6 +712,29 @@ mod tests {
         assert!(messages[1].contains("from status"), "{messages:?}");
     }
 
+    #[test]
+    fn merge_provider_metadata_unions_same_provider_key_across_phases() {
+        let start = SharedProviderMetadata::from([(
+            "fal".to_string(),
+            serde_json::json!({ "job_id": "job-1", "region": "us-east" }),
+        )]);
+        let completion = SharedProviderMetadata::from([(
+            "fal".to_string(),
+            serde_json::json!({ "region": "eu-west", "seed": 42 }),
+        )]);
+
+        let merged = merge_provider_metadata(Some(start), Some(completion))
+            .expect("both phases reported metadata");
+
+        // `job_id` only came from the start phase and must survive the merge
+        // instead of being dropped by an `entry().or_insert()` collision.
+        assert_eq!(
+            merged["fal"],
+            serde_json::json!({ "job_id": "job-1", "region": "eu-west", "seed": 42 }),
+            "expected start-only fields preserved and completion to win on collision"
+        );
+    }
+
     #[tokio::test]
     async fn poll_timeout_fails_a_generation_that_never_completes() {
         let model = ScriptedVideoModel::new(vec![]);

From ea9a33d9082f5461f2ede431eda28a37599835ab Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Thu, 3 Sep 2026 23:39:20 +0800
Subject: [PATCH 10/24] fix(provider-utils): cap successful JSON bodies
 separately from binary downloads
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

create_json_response_handler() reused DEFAULT_MAX_DOWNLOAD_SIZE (2 GiB) for
successful JSON bodies, which are held simultaneously as raw bytes, a parsed
Value, and a deserialized struct — letting a single response balloon peak
memory well past the nominal cap. It now defaults to a separate 64 MiB
DEFAULT_MAX_JSON_RESPONSE_SIZE, configurable per request via the existing
HttpRequest struct (HttpRequest::max_json_response_bytes, threaded through
ResponseHandlerInput, matching how response_timeout is already overridden
per request). create_binary_response_handler is unaffected and keeps the
2 GiB bound.

Also drops the from_slice:: + from_value(value.clone()) double parse
in favor of deserializing T directly from the bytes; raw_value (used by a
few providers for Usage.raw) is now a best-effort second from_slice call
instead of a clone of the first parse's tree.

Covered by oversize_json_success_body_is_rejected in
aimux-provider-utils/tests/api_helpers_test.rs.
---
 aimux-provider-utils/src/http.rs              |  8 ++++
 aimux-provider-utils/src/post_to_api.rs       |  1 +
 .../src/read_response_with_size_limit.rs      | 11 ++++-
 aimux-provider-utils/src/response_handler.rs  | 41 ++++++++++-------
 .../tests/api_helpers_test.rs                 | 44 +++++++++++++++++++
 .../tests/download_guard_test.rs              |  1 +
 aimux-providers/src/black_forest_labs.rs      |  2 +
 aimux-providers/src/fal.rs                    |  6 +++
 aimux-providers/src/gladia.rs                 |  1 +
 aimux-providers/src/google/files.rs           |  3 ++
 aimux-providers/src/luma.rs                   |  1 +
 aimux-providers/src/recraft.rs                |  2 +
 aimux-providers/src/replicate.rs              |  4 ++
 docs/ai-sdk-request-pipeline.md               |  1 +
 14 files changed, 109 insertions(+), 17 deletions(-)

diff --git a/aimux-provider-utils/src/http.rs b/aimux-provider-utils/src/http.rs
index 6b61d32a..ac419d78 100644
--- a/aimux-provider-utils/src/http.rs
+++ b/aimux-provider-utils/src/http.rs
@@ -216,6 +216,12 @@ pub struct HttpRequest {
     /// (e.g. Replicate `prefer: wait`) declares its own bound here. Streaming
     /// exchanges are exempt regardless.
     pub response_timeout: Option,
+    /// Per-request override of the successful-JSON-body size cap (default
+    /// [`crate::read_response_with_size_limit::DEFAULT_MAX_JSON_RESPONSE_SIZE`]).
+    /// Only consulted by [`crate::response_handler::create_json_response_handler`];
+    /// binary downloads keep the separate, larger
+    /// `DEFAULT_MAX_DOWNLOAD_SIZE` bound regardless of this field.
+    pub max_json_response_bytes: Option,
     /// AI SDK `validateUrl`: set for URLs taken from provider responses
     /// (generated assets, polling and result URLs). The exchange then goes
     /// through the SSRF download guard — target and every DNS answer
@@ -322,6 +328,7 @@ pub(crate) struct PreparedRequest {
     call_id: Option,
     recording_context: Option,
     pub(crate) response_timeout: Option,
+    pub(crate) max_json_response_bytes: Option,
     pub(crate) validation: Option,
 }
 
@@ -345,6 +352,7 @@ impl HttpRequest {
             call_id: self.call_id,
             recording_context: self.recording_context,
             response_timeout: self.response_timeout,
+            max_json_response_bytes: self.max_json_response_bytes,
             validation: self.validate_url.then_some(DownloadValidation {
                 trusted_origin: self.trusted_origin,
                 credentialed_origin: self.credentialed_origin,
diff --git a/aimux-provider-utils/src/post_to_api.rs b/aimux-provider-utils/src/post_to_api.rs
index 1d9a3212..075d3361 100644
--- a/aimux-provider-utils/src/post_to_api.rs
+++ b/aimux-provider-utils/src/post_to_api.rs
@@ -160,6 +160,7 @@ async fn execute_exchange(
         request_body_values: request_body_values.clone(),
         response,
         abort_signal: request.abort_signal.clone(),
+        max_json_response_bytes: request.max_json_response_bytes,
     };
 
     if !(200..300).contains(&status) {
diff --git a/aimux-provider-utils/src/read_response_with_size_limit.rs b/aimux-provider-utils/src/read_response_with_size_limit.rs
index 74909c17..bd5d1934 100644
--- a/aimux-provider-utils/src/read_response_with_size_limit.rs
+++ b/aimux-provider-utils/src/read_response_with_size_limit.rs
@@ -5,9 +5,18 @@ use futures::StreamExt;
 
 use aimux_core::{AiMuxError, ApiCallError};
 
-/// Default maximum buffered response size (2 GiB, matching AI SDK).
+/// Default maximum buffered response size for binary downloads (2 GiB,
+/// matching AI SDK).
 pub const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 2 * 1024 * 1024 * 1024;
 
+/// Default maximum buffered response size for a successful JSON body (64
+/// MiB). AI SDK reuses its 2 GiB download bound for JSON bodies too, but a
+/// JSON success response is held simultaneously as raw bytes, a parsed
+/// `serde_json::Value`, and a deserialized struct — a 2 GiB cap lets a single
+/// response balloon to several times that in resident memory. This bound is
+/// per-request configurable via `HttpRequest::max_json_response_bytes`.
+pub const DEFAULT_MAX_JSON_RESPONSE_SIZE: usize = 64 * 1024 * 1024;
+
 /// Read a response incrementally and fail before unbounded allocation.
 ///
 /// # Errors
diff --git a/aimux-provider-utils/src/response_handler.rs b/aimux-provider-utils/src/response_handler.rs
index f0f8d4c9..db568ef8 100644
--- a/aimux-provider-utils/src/response_handler.rs
+++ b/aimux-provider-utils/src/response_handler.rs
@@ -11,7 +11,7 @@ use aimux_core::{AiMuxError, ApiCallError};
 
 use crate::extract_response_headers::extract_response_headers;
 use crate::read_response_with_size_limit::{
-    DEFAULT_MAX_DOWNLOAD_SIZE, read_response_with_size_limit,
+    DEFAULT_MAX_DOWNLOAD_SIZE, DEFAULT_MAX_JSON_RESPONSE_SIZE, read_response_with_size_limit,
 };
 
 /// Input supplied to one response handler.
@@ -20,6 +20,11 @@ pub struct ResponseHandlerInput {
     pub request_body_values: serde_json::Value,
     pub response: reqwest::Response,
     pub abort_signal: Option,
+    /// Per-request override of the successful-JSON-body size cap, forwarded
+    /// from `HttpRequest::max_json_response_bytes`. Only
+    /// [`create_json_response_handler`] consults this; other handlers ignore
+    /// it.
+    pub max_json_response_bytes: Option,
 }
 
 /// A parsed value plus response metadata.
@@ -90,6 +95,11 @@ impl ResponseHandler {
 }
 
 /// Parse a successful JSON response using the endpoint's response type.
+///
+/// The body is size-limited to [`DEFAULT_MAX_JSON_RESPONSE_SIZE`] (or
+/// `HttpRequest::max_json_response_bytes`, when the caller overrides it) —
+/// deliberately smaller than the binary-download bound, since a JSON success
+/// body is deserialized straight into `T` and held alongside the raw bytes.
 #[must_use]
 pub fn create_json_response_handler() -> ResponseHandler
 where
@@ -98,15 +108,25 @@ where
     ResponseHandler::new(|input| async move {
         let status = input.response.status().as_u16();
         let headers = extract_response_headers(input.response.headers());
+        let max_bytes = input
+            .max_json_response_bytes
+            .unwrap_or(DEFAULT_MAX_JSON_RESPONSE_SIZE);
         let body = read_response_with_size_limit(
             input.response,
             &input.url,
             &input.request_body_values,
-            DEFAULT_MAX_DOWNLOAD_SIZE,
+            max_bytes,
             input.abort_signal.as_ref(),
         )
         .await?;
-        let raw_value = serde_json::from_slice::(&body).map_err(|error| {
+        // Deserialize straight into `T` instead of parsing a
+        // `serde_json::Value` first and converting that — the previous
+        // two-step parse cloned the intermediate `Value` tree, holding
+        // bytes + `Value` + `T` at once. `raw_value` (needed by callers such
+        // as `Usage.raw`) is a best-effort second parse: `T`'s successful
+        // deserialization already proves `body` is valid JSON, so this
+        // cannot fail in a way that changes the outcome.
+        let value = serde_json::from_slice::(&body).map_err(|error| {
             AiMuxError::ApiCall(Box::new(ApiCallError {
                 status_code: Some(status),
                 response_body: Some(String::from_utf8_lossy(&body).into_owned()),
@@ -118,21 +138,10 @@ where
                 )
             }))
         })?;
-        let value = serde_json::from_value(raw_value.clone()).map_err(|error| {
-            AiMuxError::ApiCall(Box::new(ApiCallError {
-                status_code: Some(status),
-                response_body: Some(String::from_utf8_lossy(&body).into_owned()),
-                response_headers: Some(headers.clone()),
-                ..ApiCallError::new(
-                    format!("Invalid JSON response: {error}"),
-                    input.url,
-                    input.request_body_values,
-                )
-            }))
-        })?;
+        let raw_value = serde_json::from_slice::(&body).ok();
         Ok(ResponseHandlerOutput {
             value,
-            raw_value: Some(raw_value),
+            raw_value,
             response_headers: headers,
         })
     })
diff --git a/aimux-provider-utils/tests/api_helpers_test.rs b/aimux-provider-utils/tests/api_helpers_test.rs
index 308acd0e..6b365714 100644
--- a/aimux-provider-utils/tests/api_helpers_test.rs
+++ b/aimux-provider-utils/tests/api_helpers_test.rs
@@ -327,3 +327,47 @@ async fn oversize_error_body_is_truncated_not_replaced() {
     assert!(body.ends_with("…(truncated)"), "truncation marked");
     assert!(body.len() < 70 * 1024, "body actually truncated");
 }
+
+/// A JSON success body over the configured cap must fail with the
+/// size-limit `ApiCallError`, not be buffered in full. Uses
+/// `HttpRequest::max_json_response_bytes` (rather than waiting for a
+/// multi-hundred-MB fixture) to exercise the same code path a provider hits
+/// against the real 64 MiB default.
+#[tokio::test]
+async fn oversize_json_success_body_is_rejected() {
+    let server = MockServer::start().await;
+    let huge = format!("{{\"value\":\"{}\"}}", "x".repeat(2048));
+    Mock::given(method("POST"))
+        .and(path("/call"))
+        .respond_with(
+            ResponseTemplate::new(200)
+                .insert_header("content-type", "application/json")
+                .set_body_string(huge),
+        )
+        .expect(1)
+        .mount(&server)
+        .await;
+
+    let error = post_json_to_api::(
+        HttpRequest {
+            max_json_response_bytes: Some(1024),
+            ..request(format!("{}/call", server.uri()))
+        },
+        json!({"input": 1}),
+        create_json_response_handler(),
+        failed_response_handler(),
+    )
+    .await
+    .unwrap_err();
+
+    let AiMuxError::ApiCall(detail) = error else {
+        panic!("expected ApiCall");
+    };
+    assert!(
+        detail
+            .message
+            .contains("exceeded maximum size of 1024 bytes"),
+        "expected the size-limit error, got: {}",
+        detail.message
+    );
+}
diff --git a/aimux-provider-utils/tests/download_guard_test.rs b/aimux-provider-utils/tests/download_guard_test.rs
index 1b107ab2..52995b51 100644
--- a/aimux-provider-utils/tests/download_guard_test.rs
+++ b/aimux-provider-utils/tests/download_guard_test.rs
@@ -30,6 +30,7 @@ fn request(
         call_id: None,
         recording_context: None,
         response_timeout: None,
+        max_json_response_bytes: None,
         validate_url: true,
         trusted_origin: trusted_origin.map(str::to_owned),
         credentialed_origin: credentialed_origin.map(str::to_owned),
diff --git a/aimux-providers/src/black_forest_labs.rs b/aimux-providers/src/black_forest_labs.rs
index ccaafba0..ebfa09af 100644
--- a/aimux-providers/src/black_forest_labs.rs
+++ b/aimux-providers/src/black_forest_labs.rs
@@ -403,6 +403,7 @@ impl ImageModel for BlackForestLabsImageModel {
                             call_id: None,
                             recording_context: None,
                             response_timeout: None,
+                            max_json_response_bytes: None,
                             validate_url: true,
                             trusted_origin: Some(self.config.base_url.clone()),
                             // Headers are gated per URL by gated_headers above.
@@ -479,6 +480,7 @@ impl ImageModel for BlackForestLabsImageModel {
                         call_id: None,
                         recording_context: None,
                         response_timeout: None,
+                        max_json_response_bytes: None,
                         validate_url: true,
                         trusted_origin: Some(self.config.base_url.clone()),
                         // Headers are gated per URL by gated_headers above.
diff --git a/aimux-providers/src/fal.rs b/aimux-providers/src/fal.rs
index eff9548e..32460f8f 100644
--- a/aimux-providers/src/fal.rs
+++ b/aimux-providers/src/fal.rs
@@ -271,6 +271,7 @@ impl TranscriptionModel for FalTranscriptionModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
@@ -310,6 +311,7 @@ impl TranscriptionModel for FalTranscriptionModel {
                             call_id: None,
                             recording_context: None,
                             response_timeout: None,
+                            max_json_response_bytes: None,
                             validate_url: false,
                             trusted_origin: None,
                             credentialed_origin: None,
@@ -576,6 +578,7 @@ impl ImageModel for FalImageModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
@@ -619,6 +622,7 @@ impl ImageModel for FalImageModel {
                                 call_id: None,
                                 recording_context: None,
                                 response_timeout: None,
+                                max_json_response_bytes: None,
                                 validate_url: true,
                                 trusted_origin: Some(self.config.base_url.clone()),
                                 credentialed_origin: Some(self.config.base_url.clone()),
@@ -828,6 +832,7 @@ impl VideoModel for FalVideoModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
@@ -881,6 +886,7 @@ impl VideoModel for FalVideoModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs
index 9bf05ffe..bea76ab6 100644
--- a/aimux-providers/src/gladia.rs
+++ b/aimux-providers/src/gladia.rs
@@ -308,6 +308,7 @@ impl TranscriptionModel for GladiaTranscriptionModel {
                             call_id: None,
                             recording_context: None,
                             response_timeout: None,
+                            max_json_response_bytes: None,
                             validate_url: true,
                             trusted_origin: Some(self.config.base_url.clone()),
                             credentialed_origin: Some(self.config.base_url.clone()),
diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs
index ff44c790..d6099495 100644
--- a/aimux-providers/src/google/files.rs
+++ b/aimux-providers/src/google/files.rs
@@ -198,6 +198,7 @@ impl Files for GoogleFiles {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
@@ -257,6 +258,7 @@ impl Files for GoogleFiles {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: true,
                 trusted_origin: Some(self.config.base_url.clone()),
                 credentialed_origin: Some(self.config.base_url.clone()),
@@ -315,6 +317,7 @@ impl Files for GoogleFiles {
                     call_id: None,
                     recording_context: None,
                     response_timeout: None,
+                    max_json_response_bytes: None,
                     validate_url: false,
                     trusted_origin: None,
                     credentialed_origin: None,
diff --git a/aimux-providers/src/luma.rs b/aimux-providers/src/luma.rs
index 03120688..56d8d239 100644
--- a/aimux-providers/src/luma.rs
+++ b/aimux-providers/src/luma.rs
@@ -413,6 +413,7 @@ impl ImageModel for LumaImageModel {
                         call_id: None,
                         recording_context: None,
                         response_timeout: None,
+                        max_json_response_bytes: None,
                         validate_url: true,
                         trusted_origin: Some(self.config.base_url.clone()),
                         credentialed_origin: Some(self.config.base_url.clone()),
diff --git a/aimux-providers/src/recraft.rs b/aimux-providers/src/recraft.rs
index d85a9b1d..ab208cbe 100644
--- a/aimux-providers/src/recraft.rs
+++ b/aimux-providers/src/recraft.rs
@@ -271,6 +271,7 @@ async fn extract_images(
                         call_id: None,
                         recording_context: None,
                         response_timeout: None,
+                        max_json_response_bytes: None,
                         validate_url: true,
                         trusted_origin: Some(base_url.to_string()),
                         credentialed_origin: Some(base_url.to_string()),
@@ -327,6 +328,7 @@ impl ImageModel for RecraftImageModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
diff --git a/aimux-providers/src/replicate.rs b/aimux-providers/src/replicate.rs
index 1bcb1379..09eb8b21 100644
--- a/aimux-providers/src/replicate.rs
+++ b/aimux-providers/src/replicate.rs
@@ -306,6 +306,7 @@ impl ImageModel for ReplicateImageModel {
                 response_timeout: Some(Duration::from_secs(
                     max_wait.unwrap_or(DEFAULT_WAIT_SECONDS) + WAIT_TRANSPORT_MARGIN_SECONDS,
                 )),
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
@@ -349,6 +350,7 @@ impl ImageModel for ReplicateImageModel {
                             call_id: None,
                             recording_context: None,
                             response_timeout: None,
+                            max_json_response_bytes: None,
                             validate_url: true,
                             trusted_origin: Some(self.config.base_url.clone()),
                             credentialed_origin: Some(self.config.base_url.clone()),
@@ -482,6 +484,7 @@ impl VideoModel for ReplicateVideoModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
@@ -541,6 +544,7 @@ impl VideoModel for ReplicateVideoModel {
                 call_id: None,
                 recording_context: None,
                 response_timeout: None,
+                max_json_response_bytes: None,
                 validate_url: false,
                 trusted_origin: None,
                 credentialed_origin: None,
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index 7311d726..bf7e38d3 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -85,6 +85,7 @@ Core user operation
 | 既有错误字段 | 保留 Aimux 的 `provider_code`；删除旧路径派生的 `request_id` / `retry_after_ms`，retry hint 的唯一事实源是 `response_headers` |
 | error context 安全 | AI SDK 传原始 request context；Aimux 写入 public error 前使用统一白名单/脱敏/大小限制，binary 与大型 data URL 只保留摘要 |
 | 错误 body 上限 | AI SDK 全量读取错误 body（仅受 2 GiB 防 OOM 上限，超限直接抛 `DownloadError` 替换原错误）；Aimux 的 `ApiCallError` 会跨 FFI 序列化并写入 recording，因此错误 body 采用 best-effort 截断读：public `response_body` 上限 64 KiB（lossy 解码后按字符边界执行，带 `…(truncated)` 标记），解析上限 1 MiB 保证超大但合法的错误 JSON 仍能进 provider mapper，读取中途连接死亡保留已收到的部分 |
+| 成功 JSON body 上限 | AI SDK 对成功 body 复用同一个 2 GiB 下载上限；Aimux 的成功 JSON body 会同时以 bytes + `T` 的形式驻留（`raw_value` 是 best-effort 的第二次解析，不再额外持有 `Value` 克隆），2 GiB 上限会放大峰值内存，因此 `create_json_response_handler` 改用独立更小的默认值 `DEFAULT_MAX_JSON_RESPONSE_SIZE`（64 MiB），可通过 `HttpRequest::max_json_response_bytes` 按请求覆盖；`create_binary_response_handler` 不受影响，仍使用 `DEFAULT_MAX_DOWNLOAD_SIZE`（2 GiB） |
 | provider 默认 retry | 保留既有 `RetryConfig` 源码与行为兼容；Core 读取 model config，per-call 只覆盖 `max_retries` |
 | jitter | AI SDK 默认不抖动但留 `getDelayInMs` 注入点；Aimux 用同一注入点默认注入 RFC-0009 Full Jitter，只作用于 exponential delay，server hint 精确遵守（§6.5） |
 | 第一条 SSE error | AI SDK OpenAI provider 会扫描到首个 semantic output；Aimux 保留 RFC-0016 的更窄 first-event peek，使立即到达的 error 成为 retry 边界内的 attempt 失败（§8.3） |

From 0baf004acd93f426f2f5d3275e1c6ac1a576dfcf Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Thu, 3 Sep 2026 23:46:59 +0800
Subject: [PATCH 11/24] fix(providers): apply Core retry to file uploads
 (OpenAI, Anthropic, Google)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

upload_file() did a single post_to_api/post_json_to_api exchange with no
retry wrapper, so each provider's configured RetryConfig was silently
ignored for file operations even though the config surface stays exposed.

Wraps each exchange with aimux_core::retry::prepare_retries +
PreparedRetries::retry, the same primitive execute_list_models already uses
for a non-Core-operation exchange. Uploads are not billable and a failed
create-a-file request never returns an id to replay against, so a plain
retry of the whole exchange is safe — documented as the new §9.4 policy in
docs/ai-sdk-request-pipeline.md, superseding the prior "Files gets no
retry" decision.

Google's three-stage init/upload/poll flow retries each stage
independently rather than wrapping the whole upload, so an upload-stage
retry reuses the already-minted upload_url instead of re-running init, and
a poll-stage retry never re-uploads.

Covered by transient_failure_is_retried_and_succeeds in the OpenAI and
Anthropic files tests, and transient_upload_failure_is_retried_without_re_initiating
in the Google files test (also asserts the init stage runs exactly once).
---
 aimux-providers/src/anthropic/files.rs        |  29 ++-
 aimux-providers/src/google/files.rs           | 189 ++++++++++--------
 aimux-providers/src/openai/files.rs           |  29 ++-
 aimux-providers/tests/anthropic_files_test.rs |  40 ++++
 aimux-providers/tests/google_files_test.rs    |  80 +++++++-
 aimux-providers/tests/openai_files_test.rs    |  40 ++++
 docs/ai-sdk-request-pipeline.md               |  15 +-
 7 files changed, 322 insertions(+), 100 deletions(-)

diff --git a/aimux-providers/src/anthropic/files.rs b/aimux-providers/src/anthropic/files.rs
index f7cd5e8b..9d3cf309 100644
--- a/aimux-providers/src/anthropic/files.rs
+++ b/aimux-providers/src/anthropic/files.rs
@@ -130,13 +130,28 @@ impl Files for AnthropicFiles {
         // and its content-type are carried by `HttpBody::Bytes` — the HTTP layer
         // sets `Content-Type` from it, so it is intentionally not added to the
         // header list above.
-        let resp = aimux_provider_utils::post_to_api(
-            HttpRequest::new(self.endpoint(), header_list, options),
-            HttpBody::Bytes(body, content_type),
-            aimux_provider_utils::create_json_response_handler(),
-            super::anthropic_failed_response_handler(),
-        )
-        .await?;
+        //
+        // `upload_file` is a single-exchange Provider SPI operation (no Core
+        // `do_upload_file` wrapper, RFC-0031 §9.4), so nothing above this call
+        // retries it. Apply the Core retry primitive here, same as
+        // `execute_list_models`: the upload itself is not billable, and a
+        // failed create-a-file request never returns a file id to replay
+        // against, so a plain retry of the whole exchange is safe.
+        let retries = aimux_core::retry::prepare_retries(
+            None,
+            self.config.retry_config,
+            options.abort_signal.clone(),
+        );
+        let resp = retries
+            .retry(|| {
+                aimux_provider_utils::post_to_api(
+                    HttpRequest::new(self.endpoint(), header_list.clone(), options),
+                    HttpBody::Bytes(body.clone(), content_type.clone()),
+                    aimux_provider_utils::create_json_response_handler(),
+                    super::anthropic_failed_response_handler(),
+                )
+            })
+            .await?;
 
         let data: AnthropicFilesResponse = resp.value;
 
diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs
index d6099495..828bcc79 100644
--- a/aimux-providers/src/google/files.rs
+++ b/aimux-providers/src/google/files.rs
@@ -189,42 +189,59 @@ impl Files for GoogleFiles {
         ));
         init_headers.push(("Content-Type".to_string(), "application/json".to_string()));
 
-        let init_resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest {
-                url: self.init_endpoint(),
-                headers: init_headers,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                max_json_response_bytes: None,
-                validate_url: false,
-                trusted_origin: None,
-                credentialed_origin: None,
-            },
-            init_body_value,
-            aimux_provider_utils::ResponseHandler::new(|input| async move {
-                let headers =
-                    aimux_provider_utils::extract_response_headers::extract_response_headers(
-                        input.response.headers(),
-                    );
-                Ok(aimux_provider_utils::ResponseHandlerOutput {
-                    value: (),
-                    raw_value: None,
-                    response_headers: headers,
+        // `upload_file` is a single-exchange Provider SPI operation (no Core
+        // `do_upload_file` wrapper, RFC-0031 §9.4), so nothing above this call
+        // retries any of its three stages. Each stage is retried independently
+        // (same Core primitive `execute_list_models` uses) rather than
+        // wrapping the whole upload: a transient failure in the upload or
+        // poll stage must not replay a prior stage — the init exchange that
+        // minted `upload_url`, or the potentially large upload body itself.
+        let retries = aimux_core::retry::prepare_retries(
+            None,
+            self.config.retry_config,
+            options.abort_signal.clone(),
+        );
+        let init_resp = retries
+            .retry(|| async {
+                aimux_provider_utils::post_json_to_api(
+                    HttpRequest {
+                        url: self.init_endpoint(),
+                        headers: init_headers.clone(),
+
+                        abort_signal: options.abort_signal.clone(),
+                        call_id: None,
+                        recording_context: None,
+                        response_timeout: None,
+                        max_json_response_bytes: None,
+                        validate_url: false,
+                        trusted_origin: None,
+                        credentialed_origin: None,
+                    },
+                    init_body_value.clone(),
+                    aimux_provider_utils::ResponseHandler::new(|input| async move {
+                        let headers = aimux_provider_utils::extract_response_headers::extract_response_headers(
+                            input.response.headers(),
+                        );
+                        Ok(aimux_provider_utils::ResponseHandlerOutput {
+                            value: (),
+                            raw_value: None,
+                            response_headers: headers,
+                        })
+                    }),
+                    super::google_failed_response_handler(),
+                )
+                .await
+                // Applied per attempt (not just on final exhaustion) so a
+                // `RetryError`'s saved `errors` also carry this context.
+                .map_err(|e| match e {
+                    AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError {
+                        message: format!("Failed to initiate resumable upload: {}", d.message),
+                        ..*d
+                    })),
+                    e => e,
                 })
-            }),
-            super::google_failed_response_handler(),
-        )
-        .await
-        .map_err(|e| match e {
-            AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError {
-                message: format!("Failed to initiate resumable upload: {}", d.message),
-                ..*d
-            })),
-            e => e,
-        })?;
+            })
+            .await?;
 
         let upload_url = init_resp
             .response_headers
@@ -249,32 +266,40 @@ impl Files for GoogleFiles {
         // The upload URL comes from the init response's x-goog-upload-url
         // header and receives the user's file bytes; validate it. (AI SDK
         // fetches this URL unvalidated — kept stricter here deliberately.)
-        let upload_resp = aimux_provider_utils::post_to_api(
-            HttpRequest {
-                url: upload_url,
-                headers: upload_headers,
-
-                abort_signal: options.abort_signal.clone(),
-                call_id: None,
-                recording_context: None,
-                response_timeout: None,
-                max_json_response_bytes: None,
-                validate_url: true,
-                trusted_origin: Some(self.config.base_url.clone()),
-                credentialed_origin: Some(self.config.base_url.clone()),
-            },
-            HttpBody::Bytes(file_bytes, media_type.clone()),
-            aimux_provider_utils::create_json_response_handler::(),
-            super::google_failed_response_handler(),
-        )
-        .await
-        .map_err(|e| match e {
-            AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError {
-                message: format!("Failed to upload file data: {}", d.message),
-                ..*d
-            })),
-            e => e,
-        })?;
+        // Retried independently of the init stage above (§ note there): a
+        // retry here resends the same bytes to the same `upload_url`, never
+        // re-mints one.
+        let upload_resp = retries
+            .retry(|| async {
+                aimux_provider_utils::post_to_api(
+                    HttpRequest {
+                        url: upload_url.clone(),
+                        headers: upload_headers.clone(),
+
+                        abort_signal: options.abort_signal.clone(),
+                        call_id: None,
+                        recording_context: None,
+                        response_timeout: None,
+                        max_json_response_bytes: None,
+                        validate_url: true,
+                        trusted_origin: Some(self.config.base_url.clone()),
+                        credentialed_origin: Some(self.config.base_url.clone()),
+                    },
+                    HttpBody::Bytes(file_bytes.clone(), media_type.clone()),
+                    aimux_provider_utils::create_json_response_handler::(),
+                    super::google_failed_response_handler(),
+                )
+                .await
+                // Applied per attempt, same reasoning as the init stage above.
+                .map_err(|e| match e {
+                    AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError {
+                        message: format!("Failed to upload file data: {}", d.message),
+                        ..*d
+                    })),
+                    e => e,
+                })
+            })
+            .await?;
 
         let mut file = upload_resp.value.file;
 
@@ -308,24 +333,30 @@ impl Files for GoogleFiles {
 
             let poll_url = format!("{}/{}", self.config.base_url, file.name);
 
-            let poll_resp = aimux_provider_utils::get_from_api(
-                HttpRequest {
-                    url: poll_url.clone(),
-                    headers: poll_header_list.clone(),
-
-                    abort_signal: options.abort_signal.clone(),
-                    call_id: None,
-                    recording_context: None,
-                    response_timeout: None,
-                    max_json_response_bytes: None,
-                    validate_url: false,
-                    trusted_origin: None,
-                    credentialed_origin: None,
-                },
-                aimux_provider_utils::create_json_response_handler::(),
-                super::google_failed_response_handler(),
-            )
-            .await?;
+            // A status GET is idempotent, so it gets the same independent
+            // retry as the init/upload stages; a failure here never re-runs
+            // the upload.
+            let poll_resp = retries
+                .retry(|| {
+                    aimux_provider_utils::get_from_api(
+                        HttpRequest {
+                            url: poll_url.clone(),
+                            headers: poll_header_list.clone(),
+
+                            abort_signal: options.abort_signal.clone(),
+                            call_id: None,
+                            recording_context: None,
+                            response_timeout: None,
+                            max_json_response_bytes: None,
+                            validate_url: false,
+                            trusted_origin: None,
+                            credentialed_origin: None,
+                        },
+                        aimux_provider_utils::create_json_response_handler::(),
+                        super::google_failed_response_handler(),
+                    )
+                })
+                .await?;
 
             file = poll_resp.value;
             last_poll_body = poll_resp.raw_value.map(|value| value.to_string());
diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs
index e2aee9c5..00a39b43 100644
--- a/aimux-providers/src/openai/files.rs
+++ b/aimux-providers/src/openai/files.rs
@@ -167,13 +167,28 @@ impl Files for OpenAIFiles {
         // and its content-type are carried by `HttpBody::Bytes` — the HTTP layer
         // sets `Content-Type` from it, so it is intentionally not added to the
         // header list above.
-        let resp = aimux_provider_utils::post_to_api(
-            HttpRequest::new(self.endpoint(), header_list, options),
-            HttpBody::Bytes(body, content_type),
-            aimux_provider_utils::create_json_response_handler(),
-            super::openai_failed_response_handler(),
-        )
-        .await?;
+        //
+        // `upload_file` is a single-exchange Provider SPI operation (no Core
+        // `do_upload_file` wrapper, RFC-0031 §9.4), so nothing above this call
+        // retries it. Apply the Core retry primitive here, same as
+        // `execute_list_models`: the upload itself is not billable, and a
+        // failed create-a-file request never returns a file id to replay
+        // against, so a plain retry of the whole exchange is safe.
+        let retries = aimux_core::retry::prepare_retries(
+            None,
+            self.config.retry_config,
+            options.abort_signal.clone(),
+        );
+        let resp = retries
+            .retry(|| {
+                aimux_provider_utils::post_to_api(
+                    HttpRequest::new(self.endpoint(), header_list.clone(), options),
+                    HttpBody::Bytes(body.clone(), content_type.clone()),
+                    aimux_provider_utils::create_json_response_handler(),
+                    super::openai_failed_response_handler(),
+                )
+            })
+            .await?;
 
         let data: OpenAIFilesResponse = resp.value;
 
diff --git a/aimux-providers/tests/anthropic_files_test.rs b/aimux-providers/tests/anthropic_files_test.rs
index 4f6127f7..7153f09f 100644
--- a/aimux-providers/tests/anthropic_files_test.rs
+++ b/aimux-providers/tests/anthropic_files_test.rs
@@ -308,3 +308,43 @@ async fn has_correct_provider_name() {
 
     assert_eq!(files.provider(), "anthropic.files");
 }
+
+/// A transient 503 followed by 200 must succeed: `upload_file` retries the
+/// upload exchange using the provider's configured retry settings, the same
+/// way `list_models` already does.
+#[tokio::test]
+async fn transient_failure_is_retried_and_succeeds() {
+    let server = MockServer::start().await;
+    let attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let responder_attempts = std::sync::Arc::clone(&attempts);
+    Mock::given(method("POST"))
+        .and(path("/v1/files"))
+        .respond_with(move |_: &wiremock::Request| {
+            if responder_attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
+                ResponseTemplate::new(503)
+                    .insert_header("retry-after-ms", "0")
+                    .set_body_json(json!({"error": {"message": "try again"}}))
+            } else {
+                ResponseTemplate::new(200).set_body_json(file_response_body())
+            }
+        })
+        .mount(&server)
+        .await;
+
+    let config = AnthropicConfig::new("test-api-key")
+        .with_base_url(server.uri())
+        .with_retry_config(aimux_core::retry::RetryConfig {
+            max_retries: 1,
+            ..Default::default()
+        });
+    let provider = AnthropicProvider::new(config);
+    let files = provider.files();
+
+    let result = files.upload_file(&upload_options()).await.unwrap();
+
+    assert_eq!(
+        result.provider_reference.get("anthropic"),
+        Some(&"file-abc123".to_string())
+    );
+    assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
+}
diff --git a/aimux-providers/tests/google_files_test.rs b/aimux-providers/tests/google_files_test.rs
index 17dd6fcc..789f8199 100644
--- a/aimux-providers/tests/google_files_test.rs
+++ b/aimux-providers/tests/google_files_test.rs
@@ -401,7 +401,16 @@ async fn should_throw_when_initiation_request_fails() {
         .mount(&server)
         .await;
 
-    let provider = provider(&server);
+    // A persistent failure exhausts retries either way; pin max_retries to 0
+    // so this test (which only checks the message text, not retry
+    // exhaustion) stays fast.
+    let config = GoogleConfig::new("test-api-key")
+        .with_base_url(format!("{}/v1beta", server.uri()))
+        .with_retry_config(aimux_core::retry::RetryConfig {
+            max_retries: 0,
+            ..Default::default()
+        });
+    let provider = GoogleProvider::new(config);
     let files = provider.files();
 
     let result = files.upload_file(&upload_options()).await;
@@ -456,7 +465,16 @@ async fn should_throw_when_upload_request_fails() {
         .mount(&server)
         .await;
 
-    let provider = provider(&server);
+    // A persistent failure exhausts retries either way; pin max_retries to 0
+    // so this test (which only checks the message text, not retry
+    // exhaustion) stays fast.
+    let config = GoogleConfig::new("test-api-key")
+        .with_base_url(format!("{}/v1beta", server.uri()))
+        .with_retry_config(aimux_core::retry::RetryConfig {
+            max_retries: 0,
+            ..Default::default()
+        });
+    let provider = GoogleProvider::new(config);
     let files = provider.files();
 
     let result = files.upload_file(&upload_options()).await;
@@ -586,3 +604,61 @@ async fn should_omit_optional_fields_from_metadata_when_not_present() {
     assert!(google.get("expirationTime").is_none());
     assert!(google.get("sha256Hash").is_none());
 }
+
+/// A transient 503 on the upload stage followed by 200 must succeed, and
+/// must not re-run the init stage — a retried upload reuses the
+/// already-minted `upload_url` instead of requesting a new one.
+#[tokio::test]
+async fn transient_upload_failure_is_retried_without_re_initiating() {
+    let server = MockServer::start().await;
+    let upload_url = format!("{}/resume", server.uri());
+
+    let init_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let init_observed = std::sync::Arc::clone(&init_attempts);
+    Mock::given(method("POST"))
+        .and(path("/upload/v1beta/files"))
+        .respond_with(move |_: &wiremock::Request| {
+            init_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            ResponseTemplate::new(200).insert_header("x-goog-upload-url", upload_url.as_str())
+        })
+        .mount(&server)
+        .await;
+
+    let upload_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let upload_observed = std::sync::Arc::clone(&upload_attempts);
+    Mock::given(method("POST"))
+        .and(path("/resume"))
+        .respond_with(move |_: &wiremock::Request| {
+            if upload_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
+                ResponseTemplate::new(503)
+                    .insert_header("retry-after-ms", "0")
+                    .set_body_json(json!({"error": {"message": "try again"}}))
+            } else {
+                ResponseTemplate::new(200).set_body_json(json!({ "file": default_file_resource() }))
+            }
+        })
+        .mount(&server)
+        .await;
+
+    let config = GoogleConfig::new("test-api-key")
+        .with_base_url(format!("{}/v1beta", server.uri()))
+        .with_retry_config(aimux_core::retry::RetryConfig {
+            max_retries: 1,
+            ..Default::default()
+        });
+    let provider = GoogleProvider::new(config);
+    let files = provider.files();
+
+    let result = files.upload_file(&upload_options()).await.unwrap();
+
+    assert_eq!(
+        result.provider_reference.get("google"),
+        Some(&"https://generativelanguage.googleapis.com/v1beta/files/abc123".to_string())
+    );
+    assert_eq!(
+        init_attempts.load(std::sync::atomic::Ordering::SeqCst),
+        1,
+        "the init stage must not be replayed by an upload-stage retry"
+    );
+    assert_eq!(upload_attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
+}
diff --git a/aimux-providers/tests/openai_files_test.rs b/aimux-providers/tests/openai_files_test.rs
index 20d531de..ab503b3b 100644
--- a/aimux-providers/tests/openai_files_test.rs
+++ b/aimux-providers/tests/openai_files_test.rs
@@ -279,3 +279,43 @@ async fn should_set_specification_version_and_provider() {
     assert_eq!(files.specification_version(), "v4");
     assert_eq!(files.provider(), "openai.files");
 }
+
+/// A transient 503 followed by 200 must succeed: `upload_file` retries the
+/// upload exchange using the provider's configured retry settings, the same
+/// way `list_models` already does.
+#[tokio::test]
+async fn transient_failure_is_retried_and_succeeds() {
+    let server = MockServer::start().await;
+    let attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let responder_attempts = std::sync::Arc::clone(&attempts);
+    Mock::given(method("POST"))
+        .and(path("/files"))
+        .respond_with(move |_: &wiremock::Request| {
+            if responder_attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
+                ResponseTemplate::new(503)
+                    .insert_header("retry-after-ms", "0")
+                    .set_body_json(json!({"error": {"message": "try again"}}))
+            } else {
+                ResponseTemplate::new(200).set_body_json(file_response_body("file-retried"))
+            }
+        })
+        .mount(&server)
+        .await;
+
+    let config = OpenAIConfig::new("test-api-key")
+        .with_base_url(server.uri())
+        .with_retry_config(aimux_core::retry::RetryConfig {
+            max_retries: 1,
+            ..Default::default()
+        });
+    let provider = OpenAIProvider::new(config);
+    let files = provider.files();
+
+    let result = files.upload_file(&upload_options(None)).await.unwrap();
+
+    assert_eq!(
+        result.provider_reference.get("openai"),
+        Some(&"file-retried".to_string())
+    );
+    assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
+}
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index bf7e38d3..5b281884 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -788,11 +788,16 @@ Composite 外层不重放；重试放在已有语义边界内：
 ### 9.4 Files
 
 `Files::upload_file` 按 AI SDK Provider SPI 本身就是公开 operation，没有对应的
-`do_upload_file`；本轮不为它虚构第二层。upload 只执行一次，Provider
-`RetryConfig` 不适用于 Files。这与旧 HTTP helper 会对 upload exchange 自动重试的
-行为不同；对非幂等 POST 保留该行为可能重复创建文件/上传会话，因此本轮
-明确不保留。Google Files 的 poll GET 如需独立 retry，应作为后续幂等
-exchange 改动，不得包住整个 upload。
+`do_upload_file`；本轮不为它虚构第二层。**修订**（RFC-0031 review）：`upload_file`
+现在用 Provider 已配置的 `RetryConfig`（`aimux_core::retry::prepare_retries` +
+`PreparedRetries::retry`，与 `execute_list_models` 同一 primitive）包住单次 HTTP
+exchange——upload 不计费，失败的 create-file 请求不会返回可供下一次请求复用的
+file id，因此对整个 exchange 做 plain retry 是安全的；这与非幂等的“创建资源”类
+POST 通常不能自动重试的一般规则不同，此处的例外前提就是“不计费 + 无副作用可
+重放”。OpenAI 与 Anthropic 各是一次 exchange，直接整体重试。Google Files 是
+init/upload/poll 三阶段：三个 exchange **分别独立**重试，而不是包住整个
+`upload_file`——upload 阶段的重试只重发文件字节到已经拿到的 `upload_url`，不会
+重新调用 init 铸造新 URL；poll 阶段的重试只重发状态 GET，不会触发 upload。
 
 ---
 

From c38bfd4752a41963032f51d63f535a58cde54d9f Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Thu, 3 Sep 2026 23:50:45 +0800
Subject: [PATCH 12/24] fix(providers): retry assemblyai/gladia upload and
 submit stages independently
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

do_generate() ran upload -> submit -> poll as three plain exchanges inside
one Core-retried operation, only the poll loop had its own retry. A
transient failure of the submit exchange caused Core's outer retry to
replay the whole do_generate, re-uploading the (potentially large) audio
file for no reason.

Each stage now gets its own retry via the Core primitive, using the same
retry config as the existing poll-stage retry. When a stage's own retries
are exhausted it returns AiMuxError::Retry, which the outer Core retry
passes through unchanged instead of re-wrapping (see
retry_with_exponential_backoff's Err(AiMuxError::Retry(_)) arm in
aimux-core/src/retry.rs), so do_generate itself is never replayed either.

Checked the other transcription/single-modality providers for the same
shape: revai combines the upload and submit into one multipart exchange
(no separate re-uploadable stage), and fal/luma/black_forest_labs inline
their input (base64 data URL or a provider-hosted reference) into a single
submit request rather than uploading separately — none of them have the
two-exchange upload-then-submit shape this fix addresses.

Covered by transient_submit_failure_is_retried_without_re_uploading
(assemblyai) and transient_initiate_failure_is_retried_without_re_uploading
(gladia), both asserting the upload endpoint received exactly one request.
---
 aimux-providers/src/assemblyai.rs             | 55 +++++++-----
 aimux-providers/src/gladia.rs                 | 83 +++++++++++--------
 .../tests/assemblyai_transcription_test.rs    | 63 +++++++++++++-
 .../tests/gladia_transcription_test.rs        | 66 ++++++++++++++-
 4 files changed, 213 insertions(+), 54 deletions(-)

diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs
index 6103ea4c..dc5b6868 100644
--- a/aimux-providers/src/assemblyai.rs
+++ b/aimux-providers/src/assemblyai.rs
@@ -236,14 +236,32 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
         let headers = self.build_headers(options.headers.as_ref());
         let header_list: Vec<(String, String)> = headers.into_iter().collect();
 
+        // Core wraps the whole `do_generate` in one retry (RFC-0031 §6.2), so
+        // without per-stage retries here, a transient failure in a later
+        // stage would replay every earlier stage — re-uploading the audio to
+        // retry a submit or poll failure. Each stage gets its own retry
+        // against the provider's configured retry settings instead; an
+        // exhausted inner retry returns `AiMuxError::Retry`, which the outer
+        // Core retry passes through unchanged rather than re-wrapping
+        // (`retry_with_exponential_backoff`'s `Err(AiMuxError::Retry(_))`
+        // arm), so `do_generate` is never replayed either.
+        let retries = retry::prepare_retries(
+            options.max_retries,
+            self.retry_config(),
+            options.abort_signal.clone(),
+        );
+
         // Step 1: Upload audio.
-        let resp = aimux_provider_utils::post_to_api(
-            HttpRequest::new(self.upload_url(), header_list.clone(), options),
-            HttpBody::Bytes(audio_bytes, "application/octet-stream".to_string()),
-            aimux_provider_utils::create_json_response_handler(),
-            assemblyai_failed_response_handler(),
-        )
-        .await?;
+        let resp = retries
+            .retry(|| {
+                aimux_provider_utils::post_to_api(
+                    HttpRequest::new(self.upload_url(), header_list.clone(), options),
+                    HttpBody::Bytes(audio_bytes.clone(), "application/octet-stream".to_string()),
+                    aimux_provider_utils::create_json_response_handler(),
+                    assemblyai_failed_response_handler(),
+                )
+            })
+            .await?;
 
         let upload: AssemblyAIUploadResponse = resp.value;
 
@@ -276,20 +294,19 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
             }
         }
 
-        let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest::new(self.transcript_url(), header_list.clone(), options),
-            Value::Object(body),
-            aimux_provider_utils::create_json_response_handler(),
-            assemblyai_failed_response_handler(),
-        )
-        .await?;
+        let body = Value::Object(body);
+        let resp = retries
+            .retry(|| {
+                aimux_provider_utils::post_json_to_api(
+                    HttpRequest::new(self.transcript_url(), header_list.clone(), options),
+                    body.clone(),
+                    aimux_provider_utils::create_json_response_handler(),
+                    assemblyai_failed_response_handler(),
+                )
+            })
+            .await?;
 
         let submit: AssemblyAISubmitResponse = resp.value;
-        let retries = retry::prepare_retries(
-            options.max_retries,
-            self.retry_config(),
-            options.abort_signal.clone(),
-        );
 
         // Step 3: Poll for completion.
         let mut raw_body: Value;
diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs
index bea76ab6..c0865647 100644
--- a/aimux-providers/src/gladia.rs
+++ b/aimux-providers/src/gladia.rs
@@ -221,25 +221,43 @@ impl TranscriptionModel for GladiaTranscriptionModel {
 
         let headers = self.build_headers(options.headers.as_ref());
 
+        // Core wraps the whole `do_generate` in one retry (RFC-0031 §6.2), so
+        // without per-stage retries here, a transient failure in a later
+        // stage would replay every earlier stage — re-uploading the audio to
+        // retry an initiate or poll failure. Each stage gets its own retry
+        // against the provider's configured retry settings instead; an
+        // exhausted inner retry returns `AiMuxError::Retry`, which the outer
+        // Core retry passes through unchanged rather than re-wrapping
+        // (`retry_with_exponential_backoff`'s `Err(AiMuxError::Retry(_))`
+        // arm), so `do_generate` is never replayed either.
+        let retries = retry::prepare_retries(
+            options.max_retries,
+            self.retry_config(),
+            options.abort_signal.clone(),
+        );
+
         // Step 1: Upload audio.
         let mut form = MultipartForm::new();
         form.file("audio", &filename, &options.media_type, &audio_bytes)?;
         let (body_bytes, content_type) = form.finish();
 
-        let resp = aimux_provider_utils::post_to_api(
-            HttpRequest::new(
-                self.upload_url(),
-                headers
-                    .iter()
-                    .map(|(k, v)| (k.clone(), v.clone()))
-                    .collect(),
-                options,
-            ),
-            HttpBody::Bytes(body_bytes, content_type),
-            aimux_provider_utils::create_json_response_handler(),
-            gladia_failed_response_handler(),
-        )
-        .await?;
+        let resp = retries
+            .retry(|| {
+                aimux_provider_utils::post_to_api(
+                    HttpRequest::new(
+                        self.upload_url(),
+                        headers
+                            .iter()
+                            .map(|(k, v)| (k.clone(), v.clone()))
+                            .collect(),
+                        options,
+                    ),
+                    HttpBody::Bytes(body_bytes.clone(), content_type.clone()),
+                    aimux_provider_utils::create_json_response_handler(),
+                    gladia_failed_response_handler(),
+                )
+            })
+            .await?;
 
         let upload: GladiaUploadResponse = resp.value;
 
@@ -259,27 +277,26 @@ impl TranscriptionModel for GladiaTranscriptionModel {
             }
         }
 
-        let resp = aimux_provider_utils::post_json_to_api(
-            HttpRequest::new(
-                self.initiate_url(),
-                headers
-                    .iter()
-                    .map(|(k, v)| (k.clone(), v.clone()))
-                    .collect(),
-                options,
-            ),
-            Value::Object(body),
-            aimux_provider_utils::create_json_response_handler(),
-            gladia_failed_response_handler(),
-        )
-        .await?;
+        let body = Value::Object(body);
+        let resp = retries
+            .retry(|| {
+                aimux_provider_utils::post_json_to_api(
+                    HttpRequest::new(
+                        self.initiate_url(),
+                        headers
+                            .iter()
+                            .map(|(k, v)| (k.clone(), v.clone()))
+                            .collect(),
+                        options,
+                    ),
+                    body.clone(),
+                    aimux_provider_utils::create_json_response_handler(),
+                    gladia_failed_response_handler(),
+                )
+            })
+            .await?;
 
         let init: GladiaInitResponse = resp.value;
-        let retries = retry::prepare_retries(
-            options.max_retries,
-            self.retry_config(),
-            options.abort_signal.clone(),
-        );
 
         // Step 3: Poll for result.
         let mut raw_body: Value;
diff --git a/aimux-providers/tests/assemblyai_transcription_test.rs b/aimux-providers/tests/assemblyai_transcription_test.rs
index ba551e9a..afd5a796 100644
--- a/aimux-providers/tests/assemblyai_transcription_test.rs
+++ b/aimux-providers/tests/assemblyai_transcription_test.rs
@@ -1,4 +1,4 @@
-//! Rust translation of the AssemblyAI transcription model tests.
+//! Rust translation of the AssemblyAI transcription model tests.
 //! Source: `reference/ai/packages/assemblyai/src/assemblyai-transcription-model.test.ts`
 
 use std::collections::HashMap;
@@ -231,3 +231,64 @@ async fn should_include_provider_metadata_with_utterances() {
     let aai = md.get("assemblyai").unwrap();
     assert!(aai.get("utterances").is_some());
 }
+
+/// A transient 503 on the transcript-submit stage, followed by 200, must
+/// succeed without re-uploading the audio: the upload exchange is retried
+/// independently of the submit exchange, so Core's outer per-`do_generate`
+/// retry never needs to (and must not) replay the upload.
+#[tokio::test]
+async fn transient_submit_failure_is_retried_without_re_uploading() {
+    let server = MockServer::start().await;
+
+    let upload_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let upload_observed = std::sync::Arc::clone(&upload_attempts);
+    Mock::given(method("POST"))
+        .and(path("/v2/upload"))
+        .respond_with(move |_: &wiremock::Request| {
+            upload_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            ResponseTemplate::new(200)
+                .set_body_json(json!({"upload_url": "https://upload.assemblyai.com/test"}))
+        })
+        .mount(&server)
+        .await;
+
+    let submit_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let submit_observed = std::sync::Arc::clone(&submit_attempts);
+    Mock::given(method("POST"))
+        .and(path("/v2/transcript"))
+        .respond_with(move |_: &wiremock::Request| {
+            if submit_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
+                ResponseTemplate::new(503)
+                    .insert_header("retry-after-ms", "0")
+                    .set_body_json(json!({"error": "try again"}))
+            } else {
+                ResponseTemplate::new(200)
+                    .set_body_json(json!({"id": "test-id", "status": "queued"}))
+            }
+        })
+        .mount(&server)
+        .await;
+
+    Mock::given(method("GET"))
+        .and(path("/v2/transcript/test-id"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(fixture_transcript()))
+        .mount(&server)
+        .await;
+
+    let config = AssemblyAIConfig::new("test-api-key").with_base_url(server.uri());
+    let provider = AssemblyAIProvider::new(config);
+    let model = provider.transcription("universal-2");
+
+    let result = model
+        .do_generate(&options(mock_audio(), "audio/wav"))
+        .await
+        .unwrap();
+
+    assert_eq!(result.text, "Hello from AssemblyAI.");
+    assert_eq!(
+        upload_attempts.load(std::sync::atomic::Ordering::SeqCst),
+        1,
+        "the upload stage must not be replayed by a submit-stage retry"
+    );
+    assert_eq!(submit_attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
+}
diff --git a/aimux-providers/tests/gladia_transcription_test.rs b/aimux-providers/tests/gladia_transcription_test.rs
index 9a696f0c..ef46e4bd 100644
--- a/aimux-providers/tests/gladia_transcription_test.rs
+++ b/aimux-providers/tests/gladia_transcription_test.rs
@@ -1,4 +1,4 @@
-//! Rust translation of the Gladia transcription model tests.
+//! Rust translation of the Gladia transcription model tests.
 //! Source: `reference/ai/packages/gladia/src/gladia-transcription-model.test.ts`
 
 use std::collections::HashMap;
@@ -183,3 +183,67 @@ async fn should_include_provider_metadata() {
             .contains_key("gladia")
     );
 }
+
+/// A transient 503 on the initiate stage, followed by 200, must succeed
+/// without re-uploading the audio: the upload exchange is retried
+/// independently of the initiate exchange, so Core's outer per-`do_generate`
+/// retry never needs to (and must not) replay the upload.
+#[tokio::test]
+async fn transient_initiate_failure_is_retried_without_re_uploading() {
+    let server = MockServer::start().await;
+
+    let upload_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let upload_observed = std::sync::Arc::clone(&upload_attempts);
+    Mock::given(method("POST"))
+        .and(path("/v2/upload"))
+        .respond_with(move |_: &wiremock::Request| {
+            upload_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+            ResponseTemplate::new(200)
+                .set_body_json(json!({"audio_url": "https://upload.gladia.io/test"}))
+        })
+        .mount(&server)
+        .await;
+
+    let initiate_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+    let initiate_observed = std::sync::Arc::clone(&initiate_attempts);
+    let result_url = format!("{}/v2/result/test", server.uri());
+    Mock::given(method("POST"))
+        .and(path("/v2/pre-recorded"))
+        .respond_with(move |_: &wiremock::Request| {
+            if initiate_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
+                ResponseTemplate::new(503)
+                    .insert_header("retry-after-ms", "0")
+                    .set_body_json(json!({"error": "try again"}))
+            } else {
+                ResponseTemplate::new(200).set_body_json(json!({"result_url": result_url}))
+            }
+        })
+        .mount(&server)
+        .await;
+
+    Mock::given(method("GET"))
+        .and(path("/v2/result/test"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(fixture_result()))
+        .mount(&server)
+        .await;
+
+    let config = GladiaConfig::new("test-api-key").with_base_url(server.uri());
+    let provider = GladiaProvider::new(config);
+    let model = provider.transcription("default");
+
+    let result = model
+        .do_generate(&options(mock_audio(), "audio/wav"))
+        .await
+        .unwrap();
+
+    assert_eq!(result.text, "Hello from Gladia.");
+    assert_eq!(
+        upload_attempts.load(std::sync::atomic::Ordering::SeqCst),
+        1,
+        "the upload stage must not be replayed by an initiate-stage retry"
+    );
+    assert_eq!(
+        initiate_attempts.load(std::sync::atomic::Ordering::SeqCst),
+        2
+    );
+}

From 857f54c5cc7a8027192e0b78bdbfda9b6c0e23ae Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Thu, 3 Sep 2026 23:54:39 +0800
Subject: [PATCH 13/24] fix(core): batch video generation by
 max_videos_per_call and validate n
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

generate_video() performed exactly one do_start and never read options.n or
VideoModel::max_videos_per_call(), so n > 1 silently returned one video and
n = 0 still started a billed request.

start_and_poll() now validates n (0 is InvalidArgument, before any network
call), splits n into max_videos_per_call-sized batches (AI SDK
generateVideo/generateImage's algorithm: every batch but the last is
full-sized, the last is the remainder or a full batch), and runs every
batch's independent do_start/poll cycle concurrently via
futures::future::try_join_all, matching the AI SDK's Promise.all — not
sequentially. Each batch mints its own idempotency key, so a retried batch
never collides with another batch's replay. Batch results are flattened
back in order; provider metadata is aggregated across batches with the
same deep-merge policy as the start/completion merge (finding 5).

One Rust-specific note documented in code and in the RFC: on the first
batch failure, Rust drops the other in-flight batch futures (unlike
Promise.all, where sibling promises keep running); this doesn't change
correctness since nothing here reconnects to or cancels an already-started
provider job either way.

Covered by three new tests in aimux-core/src/video_model.rs:
n_zero_is_rejected_before_any_network_call,
n_above_max_per_call_splits_into_batches_and_aggregates (also asserts
per-batch idempotency keys and cross-batch metadata merge), and
batches_run_concurrently_not_sequentially.
---
 aimux-core/src/video_model.rs   | 272 ++++++++++++++++++++++++++++++--
 docs/ai-sdk-request-pipeline.md |  10 ++
 2 files changed, 266 insertions(+), 16 deletions(-)

diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs
index 342789c7..57f6472b 100644
--- a/aimux-core/src/video_model.rs
+++ b/aimux-core/src/video_model.rs
@@ -331,17 +331,31 @@ pub trait VideoModel: Send + Sync {
     ) -> Result;
 }
 
-/// User-facing video generation with Core-owned retry, polling, and timeout.
+/// User-facing video generation with Core-owned retry, polling, batching, and
+/// timeout.
 ///
-/// Orchestration (AI SDK `generate-video`): `do_start` is retried as one
-/// unit, then `do_status` is polled — each poll retried independently —
-/// against the same operation reference, so a transient poll failure never
-/// re-creates the billed task.
+/// Orchestration (AI SDK `generateVideo`, mirroring `generateImage`'s
+/// batching): `options.n` is split into `VideoModel::max_videos_per_call`
+/// sized batches, and every batch runs a full, independent `do_start`/poll
+/// cycle **concurrently** (the AI SDK batches images/videos via
+/// `Promise.all`, not sequentially) — each with its own idempotency key, so a
+/// retried batch never collides with another batch's replay. Within a batch,
+/// `do_start` is retried as one unit, then `do_status` is polled — each poll
+/// retried independently — against that batch's operation reference, so a
+/// transient poll failure never re-creates the billed task. Batch results are
+/// flattened back in batch order; provider metadata across batches is
+/// deep-merged the same way start/completion metadata is within a batch (see
+/// [`merge_provider_metadata`]).
 ///
 /// # Errors
 ///
-/// Returns the provider failure, retry exhaustion, poll or operation
-/// timeout, or caller abort.
+/// Returns `InvalidArgument` when `options.n == 0` (checked before any
+/// network call), or the first batch's provider failure, retry exhaustion,
+/// poll or operation timeout, or caller abort — Rust drops the other
+/// in-flight batches' futures on the first error, unlike `Promise.all`,
+/// which lets sibling settles run to completion; their partial network
+/// effects (e.g. an already-started but abandoned batch) are the same
+/// either way, since nothing here reconnects to or cancels a provider job.
 pub async fn generate_video(
     model: &dyn VideoModel,
     options: VideoCallOptions,
@@ -360,6 +374,12 @@ async fn start_and_poll(
     model: &dyn VideoModel,
     options: VideoCallOptions,
 ) -> Result {
+    if options.n == 0 {
+        return Err(AiMuxError::InvalidArgument(
+            "video generation `n` must be at least 1".to_string(),
+        ));
+    }
+
     let abort_signal = options.abort_signal.clone();
     let retries = retry::prepare_retries(
         options.max_retries,
@@ -367,6 +387,81 @@ async fn start_and_poll(
         abort_signal.clone(),
     );
 
+    let mut poll = model.poll_config();
+    if let Some(overrides) = options.poll {
+        if let Some(ms) = overrides.interval_ms {
+            poll.interval = Duration::from_millis(ms);
+        }
+        if let Some(ms) = overrides.timeout_ms {
+            poll.timeout = Duration::from_millis(ms);
+        }
+    }
+
+    let max_per_call = model.max_videos_per_call().unwrap_or(options.n).max(1);
+    let batch_counts = batch_video_counts(options.n, max_per_call);
+
+    let batches = batch_counts.into_iter().map(|n| {
+        let mut batch_options = options.clone();
+        batch_options.n = n;
+        start_and_poll_one_batch(model, batch_options, &retries, poll, abort_signal.clone())
+    });
+    let results = futures::future::try_join_all(batches).await?;
+
+    let mut videos = Vec::new();
+    let mut warnings = Vec::new();
+    let mut provider_metadata = None;
+    let mut response = None;
+    for result in results {
+        videos.extend(result.videos);
+        warnings.extend(result.warnings);
+        provider_metadata = merge_provider_metadata(provider_metadata, result.provider_metadata);
+        response.get_or_insert(result.response);
+    }
+
+    Ok(VideoResult {
+        videos,
+        warnings,
+        provider_metadata,
+        // `batch_counts` is never empty (`n >= 1` was checked above), so at
+        // least one batch ran and `response` is always `Some` here; the first
+        // batch's response represents the call for telemetry purposes.
+        response: response.unwrap_or_default(),
+    })
+}
+
+/// Split `n` into `max_per_call`-sized batches (AI SDK `generateVideo` /
+/// `generateImage`): every batch but the last is full-sized; the last batch
+/// takes the remainder, or a full batch if `n` divides evenly.
+fn batch_video_counts(n: u32, max_per_call: u32) -> Vec {
+    let max_per_call = max_per_call.max(1);
+    let call_count = n.div_ceil(max_per_call);
+    (0..call_count)
+        .map(|i| {
+            if i + 1 < call_count {
+                max_per_call
+            } else {
+                let remainder = n % max_per_call;
+                if remainder == 0 {
+                    max_per_call
+                } else {
+                    remainder
+                }
+            }
+        })
+        .collect()
+}
+
+/// Run one `do_start` + poll cycle for a single provider-sized batch.
+///
+/// Each batch mints its own idempotency key (AI SDK batching semantics): a
+/// batch, not the whole `n`-video request, is the unit of idempotent replay.
+async fn start_and_poll_one_batch(
+    model: &dyn VideoModel,
+    options: VideoCallOptions,
+    retries: &retry::PreparedRetries,
+    poll: VideoPollConfig,
+    abort_signal: Option,
+) -> Result {
     // `do_start` is billable: mint one idempotency key per logical start,
     // OUTSIDE the retry closure, so providers that honor this header can
     // deduplicate a replay where the first attempt succeeded but its response
@@ -385,15 +480,6 @@ async fn start_and_poll(
 
     let start = retries.retry(|| model.do_start(&start_options)).await?;
 
-    let mut poll = model.poll_config();
-    if let Some(overrides) = options.poll {
-        if let Some(ms) = overrides.interval_ms {
-            poll.interval = Duration::from_millis(ms);
-        }
-        if let Some(ms) = overrides.timeout_ms {
-            poll.timeout = Duration::from_millis(ms);
-        }
-    }
     let poll_started = Instant::now();
     loop {
         let elapsed = poll_started.elapsed();
@@ -800,4 +886,158 @@ mod tests {
                 .is_empty()
         );
     }
+
+    #[tokio::test]
+    async fn n_zero_is_rejected_before_any_network_call() {
+        let model = ScriptedVideoModel::new(vec![StatusStep::Complete]);
+        let mut options = fast_poll_options();
+        options.n = 0;
+
+        let error = generate_video(&model, options).await.unwrap_err();
+
+        assert!(
+            matches!(error, AiMuxError::InvalidArgument(_)),
+            "expected InvalidArgument, got {error:?}"
+        );
+        assert_eq!(
+            model.starts.load(Ordering::SeqCst),
+            0,
+            "n=0 must not start a billed request"
+        );
+    }
+
+    /// A model whose `do_start` is retriable per-batch (a small delay proves
+    /// batches overlap rather than running one after another) and whose
+    /// `do_status` completes immediately with one video + one provider
+    /// metadata field unique to the batch index.
+    struct BatchedVideoModel {
+        starts: AtomicU32,
+        start_idempotency_keys: Mutex>>,
+        start_delay: Duration,
+        max_per_call: Option,
+    }
+
+    #[async_trait]
+    impl VideoModel for BatchedVideoModel {
+        fn provider(&self) -> &str {
+            "test"
+        }
+        fn model_id(&self) -> &str {
+            "batched"
+        }
+        fn max_videos_per_call(&self) -> Option {
+            self.max_per_call
+        }
+        async fn do_start(
+            &self,
+            options: &VideoCallOptions,
+        ) -> Result {
+            let index = self.starts.fetch_add(1, Ordering::SeqCst);
+            self.start_idempotency_keys
+                .lock()
+                .expect("start key capture lock should not be poisoned")
+                .push(idempotency_key(options));
+            tokio::time::sleep(self.start_delay).await;
+            Ok(VideoOperationStart {
+                operation: serde_json::json!({ "task_id": format!("t-{index}") }),
+                warnings: vec![Warning::Other {
+                    message: format!("batch {index} start"),
+                }],
+                provider_metadata: Some(SharedProviderMetadata::from([(
+                    "test".to_string(),
+                    serde_json::json!({ format!("batch{index}"): index }),
+                )])),
+                response: VideoResponse::default(),
+            })
+        }
+        async fn do_status(
+            &self,
+            operation: &serde_json::Value,
+            _options: &VideoCallOptions,
+        ) -> Result {
+            let task_id = operation["task_id"].as_str().unwrap().to_string();
+            Ok(VideoOperationStatus::Completed(VideoResult {
+                videos: vec![VideoData::Url {
+                    url: format!("https://cdn/{task_id}.mp4"),
+                    media_type: "video/mp4".to_string(),
+                }],
+                warnings: Vec::new(),
+                provider_metadata: None,
+                response: VideoResponse::default(),
+            }))
+        }
+    }
+
+    #[tokio::test]
+    async fn n_above_max_per_call_splits_into_batches_and_aggregates() {
+        let model = BatchedVideoModel {
+            starts: AtomicU32::new(0),
+            start_idempotency_keys: Mutex::new(Vec::new()),
+            start_delay: Duration::from_millis(1),
+            max_per_call: Some(1),
+        };
+        let mut options = fast_poll_options();
+        options.n = 3;
+
+        let result = generate_video(&model, options).await.unwrap();
+
+        assert_eq!(
+            model.starts.load(Ordering::SeqCst),
+            3,
+            "n=3 with max_videos_per_call=1 must run three do_start batches"
+        );
+        assert_eq!(result.videos.len(), 3);
+        assert_eq!(result.warnings.len(), 3, "warnings from every batch kept");
+
+        let keys = model
+            .start_idempotency_keys
+            .lock()
+            .expect("start key capture lock should not be poisoned");
+        assert_eq!(keys.len(), 3);
+        assert!(keys.iter().all(Option::is_some));
+        let unique: std::collections::HashSet<_> = keys.iter().collect();
+        assert_eq!(
+            unique.len(),
+            3,
+            "each batch must mint its own idempotency key: {keys:?}"
+        );
+
+        // Provider metadata across batches is deep-merged (§ merge_provider_metadata),
+        // not `entry().or_insert()`-dropped: all three batch-specific fields survive
+        // under the shared "test" provider key.
+        let metadata = result
+            .provider_metadata
+            .expect("provider metadata aggregated");
+        let test_meta = metadata.get("test").expect("test provider metadata");
+        for i in 0..3 {
+            assert!(
+                test_meta.get(format!("batch{i}")).is_some(),
+                "batch {i}'s metadata field missing from aggregate: {test_meta:?}"
+            );
+        }
+    }
+
+    #[tokio::test]
+    async fn batches_run_concurrently_not_sequentially() {
+        let model = BatchedVideoModel {
+            starts: AtomicU32::new(0),
+            start_idempotency_keys: Mutex::new(Vec::new()),
+            start_delay: Duration::from_millis(80),
+            max_per_call: Some(1),
+        };
+        let mut options = fast_poll_options();
+        options.n = 3;
+
+        let started = Instant::now();
+        generate_video(&model, options).await.unwrap();
+        let elapsed = started.elapsed();
+
+        // Sequential batches would take >= 3 * 80ms = 240ms; concurrent
+        // batches take roughly one delay (~80ms) plus scheduling overhead.
+        assert!(
+            elapsed < Duration::from_millis(200),
+            "batches should overlap (AI SDK generateVideo/generateImage runs \
+             them concurrently via Promise.all), took {elapsed:?}"
+        );
+    }
 }
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index 5b281884..c46db21f 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -765,6 +765,16 @@ start/status 流），Core 拥有 poll 循环并分别 retry 两个阶段：
   retry 次数有界；
 - Provider-specific polling delay 不是通用 exponential retry；
 - 耗尽产生的 `RetryError` 原样透传，外层不得再次 submit。
+- **修订**（RFC-0031 review）：`options.n` 按 `VideoModel::max_videos_per_call`
+  切成多个 batch（AI SDK `generateVideo`/`generateImage` 同款算法：除最后一个
+  batch 外每个都是满批，最后一个是余数或整批），每个 batch **并发**跑一次
+  完整的独立 `do_start`/poll 流程（AI SDK 用 `Promise.all`；Rust 侧用
+  `futures::future::try_join_all`，第一个 batch 失败时其余尚在进行的 batch
+  future 被 drop，这点比 JS 更激进但不影响正确性——没有 batch 会重连或取消
+  已经发出的 provider 任务）——每个 batch 独立铸造 idempotency key。结果按
+  batch 顺序拼接 `videos`/`warnings`；`provider_metadata` 跨 batch 用与
+  start/completion 相同的 deep-merge 策略聚合。`n == 0` 在任何网络调用之前
+  返回 `InvalidArgument`。
 
 ### 9.2 Router / MoA
 

From 7a83301c1e0dce444984c11a22ebedf612f56107 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:01:31 +0800
Subject: [PATCH 14/24] fix(core): measure stream timeouts at the producer, not
 the consumer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The stream-consumption wrapper observed first_chunk_ms and chunk_ms inside
its own poll_next, so the deadlines kept elapsing while the caller was not
polling, and the biased select reported a timeout ahead of a chunk the
provider had already delivered on time. A consumer that polled late, or
spent time processing the previous item, was charged against the provider's
output budget.

The provider stream is now driven by a pump task spawned before stream_text
returns. It arms and resets the first-chunk and chunk-idle deadlines when
output arrives and forwards items through an unbounded channel; the returned
stream only forwards from that channel, watches caller abort, and stops at
the terminal item. The channel is unbounded so the pump never blocks on the
consumer (the AI SDK's streamText also consumes eagerly); the pump is aborted
when the returned stream is dropped, which drops the provider stream and
cancels the underlying HTTP exchange.

The spec's "nothing is spawned" claim is narrowed to the operation
deadlines; §3, §8.1 and acceptance item 21b/21c record the pump. The
first-chunk test that codified the consumer-delay behaviour now asserts the
opposite, and a paused-time matrix test covers chunk_ms with a slow consumer.
---
 aimux-core/Cargo.toml                         |   2 +-
 aimux-core/src/generate.rs                    | 107 +++++++++++++-----
 aimux-core/src/timeout.rs                     |   4 +-
 .../tests/stream_timeout_matrix_test.rs       |  35 +++++-
 docs/ai-sdk-request-pipeline.md               |  14 ++-
 5 files changed, 126 insertions(+), 36 deletions(-)

diff --git a/aimux-core/Cargo.toml b/aimux-core/Cargo.toml
index 5cc8a1e1..0ad5927d 100644
--- a/aimux-core/Cargo.toml
+++ b/aimux-core/Cargo.toml
@@ -22,7 +22,7 @@ ts-rs = { workspace = true }
 # declaration at the real API surface documents what Core actually asks
 # of the runtime (RFC-0016 had removed core's tokio dependency; RFC-0031
 # reintroduces it for operation deadlines and retry backoff).
-tokio = { version = "1", features = ["time", "macros"] }
+tokio = { version = "1", features = ["time", "macros", "rt", "sync"] }
 tokio-util = "0.7"
 httpdate = { workspace = true }
 rand = { workspace = true }
diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs
index 9059cecc..9db604db 100644
--- a/aimux-core/src/generate.rs
+++ b/aimux-core/src/generate.rs
@@ -1017,66 +1017,94 @@ pub async fn stream_text(
     // first-chunk budget armed at operation start continues until the first
     // output chunk.
     let operation_deadline = operation_timeout.deadline();
-    // Always wrapped, even with no abort signal or deadlines armed: this
-    // loop is also the stream's terminal fuse — `Finish` and non-recoverable
+    // The provider stream is driven by a pump task, not by the consumer's
+    // polls: first-chunk and chunk-idle deadlines measure when provider
+    // output *arrives*, so a caller that polls late still receives output the
+    // provider delivered on time. Pull-based streams cannot express this
+    // without a driver — a deadline observed only inside `poll_next` keeps
+    // elapsing while nobody polls. The channel is unbounded on purpose: the
+    // AI SDK's `streamText` also consumes eagerly, and a bounded channel
+    // would make the pump block on the consumer, reintroducing consumer
+    // latency into the provider timers. Buffering is bounded by the response
+    // itself; the pump stops at the terminal item and is aborted when the
+    // returned stream is dropped.
+    //
+    // Always wrapped, even with no abort signal or deadlines armed: the pump
+    // is also the stream's terminal fuse — `Finish` and non-recoverable
     // errors must end the stream under every configuration, or a provider
     // that keeps the connection open after `Finish` hangs a consumer
     // reading to end-of-stream.
     let stream: Pin> + Send>> = {
         let mut stream = stream;
-        let abort_signal = abort_signal;
         let chunk_ms = stream_timeout.chunk_ms;
-        Box::pin(async_stream::stream! {
+        let pump_abort = abort_signal.clone();
+        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+        let pump = tokio::spawn(async move {
             let mut chunk_deadline = first_chunk_deadline;
 
             loop {
                 let next = tokio::select! {
                     biased;
-                    () = timeout::wait_for_abort(abort_signal.as_ref()) => {
-                        yield Err(AiMuxError::from_abort_signal(
-                            abort_signal.as_ref().expect("abort future only resolves for a signal"),
-                        ));
-                        break;
-                    }
+                    // Dropping the provider stream cancels the in-flight HTTP
+                    // exchange; the consumer side reports the abort error.
+                    () = timeout::wait_for_abort(pump_abort.as_ref()) => break,
                     () = timeout::wait_for_deadline(operation_deadline) => {
-                        yield Err(operation_deadline
+                        let _ = tx.send(Err(operation_deadline
                             .expect("operation deadline future only resolves for a deadline")
-                            .error());
+                            .error()));
                         break;
                     }
                     () = timeout::wait_for_deadline(chunk_deadline) => {
-                        yield Err(chunk_deadline
+                        let _ = tx.send(Err(chunk_deadline
                             .expect("chunk deadline future only resolves for a deadline")
-                            .error());
+                            .error()));
                         break;
                     }
                     item = stream.next() => item,
                 };
 
                 let Some(item) = next else { break };
-                // A malformed individual frame does not prove that the
-                // transport is dead: SSE framing lets a later event still be
-                // valid. Transport/Core errors remain terminal so we do not
-                // keep polling a failed source indefinitely.
-                let terminal = matches!(&item, Ok(StreamPart::Finish { .. }))
-                    || matches!(&item, Err(error) if !error.is_recoverable_stream_error());
+                let terminal = is_terminal_item(&item);
                 if let Ok(part) = &item
                     && is_output_chunk(part)
                 {
                     chunk_deadline = match chunk_ms
-                        .map(|duration_ms| {
-                            timeout::TimeoutDeadline::from_now("Chunk", duration_ms)
-                        })
+                        .map(|duration_ms| timeout::TimeoutDeadline::from_now("Chunk", duration_ms))
                         .transpose()
                     {
                         Ok(deadline) => deadline,
                         Err(error) => {
-                            yield Err(error);
+                            let _ = tx.send(Err(error));
                             break;
                         }
                     };
                 }
 
+                // A closed receiver means the consumer dropped the stream.
+                if tx.send(item).is_err() || terminal {
+                    break;
+                }
+            }
+        });
+        let pump = AbortOnDrop(pump);
+        Box::pin(async_stream::stream! {
+            let _pump = pump;
+            loop {
+                let next = tokio::select! {
+                    biased;
+                    () = timeout::wait_for_abort(abort_signal.as_ref()) => {
+                        yield Err(AiMuxError::from_abort_signal(
+                            abort_signal.as_ref().expect("abort future only resolves for a signal"),
+                        ));
+                        break;
+                    }
+                    item = rx.recv() => item,
+                };
+                let Some(item) = next else { break };
+                // Mirror the pump's fuse here so the stream ends right after
+                // the terminal item, ahead of a later abort or the channel
+                // close racing it.
+                let terminal = is_terminal_item(&item);
                 yield item;
                 if terminal {
                     break;
@@ -1099,6 +1127,24 @@ pub async fn stream_text(
     })
 }
 
+/// A malformed individual frame does not prove that the transport is dead:
+/// SSE framing lets a later event still be valid. `Finish` and transport/Core
+/// errors are terminal so the stream is not polled past a dead source.
+fn is_terminal_item(item: &Result) -> bool {
+    matches!(item, Ok(StreamPart::Finish { .. }))
+        || matches!(item, Err(error) if !error.is_recoverable_stream_error())
+}
+
+/// Aborts the stream pump task when the returned stream is dropped, so an
+/// abandoned stream leaves no task driving the provider connection.
+struct AbortOnDrop(tokio::task::JoinHandle<()>);
+
+impl Drop for AbortOnDrop {
+    fn drop(&mut self) {
+        self.0.abort();
+    }
+}
+
 /// Run `do_generate` inside the RFC-0014 `generate` span and emit the
 /// `generate_end` event. A plain async fn (rather than an inline async block)
 /// so the `?` error type is pinned by the declared return type.
@@ -1418,7 +1464,11 @@ mod operation_retry_tests {
     }
 
     #[tokio::test]
-    async fn first_chunk_timeout_starts_before_the_returned_stream_is_polled() {
+    async fn late_consumer_still_receives_a_first_chunk_delivered_on_time() {
+        // The provider yields immediately; the first-chunk budget measures
+        // that arrival, not when the caller gets around to polling. Sleeping
+        // past the budget before the first poll must not turn delivered
+        // output into a timeout.
         let model = StreamModel::new(StreamBehavior::Normal);
         let options = GenerateTextOptions {
             max_retries: Some(0),
@@ -1430,12 +1480,13 @@ mod operation_retry_tests {
         };
         let mut result = stream_text(&model, "hello", options).await.unwrap();
 
-        tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        tokio::time::sleep(std::time::Duration::from_millis(20)).await;
 
         assert!(matches!(
-            result.stream.next().await.unwrap(),
-            Err(AiMuxError::Timeout(message)) if message == "First chunk timeout of 1ms exceeded"
+            result.stream.next().await.unwrap().unwrap(),
+            StreamPart::TextDelta { delta, .. } if delta == "hello"
         ));
+        assert!(result.stream.next().await.is_none());
     }
 
     #[tokio::test]
diff --git a/aimux-core/src/timeout.rs b/aimux-core/src/timeout.rs
index d46dc679..b3f9d5de 100644
--- a/aimux-core/src/timeout.rs
+++ b/aimux-core/src/timeout.rs
@@ -10,7 +10,9 @@ use crate::options::TimeoutConfiguration;
 
 /// Deadlines shared by the establishment and streaming phases of one model
 /// operation. No timer task is spawned: the future currently driving the
-/// operation observes the deadline directly.
+/// operation observes the deadline directly — for a returned stream that is
+/// the pump task in `generate::stream_text`, so deadlines measure provider
+/// output arrival rather than consumer polling.
 #[derive(Debug, Clone, Copy)]
 pub(crate) struct OperationTimeout {
     total: Option,
diff --git a/aimux-core/tests/stream_timeout_matrix_test.rs b/aimux-core/tests/stream_timeout_matrix_test.rs
index 9d839ab8..8121bf7d 100644
--- a/aimux-core/tests/stream_timeout_matrix_test.rs
+++ b/aimux-core/tests/stream_timeout_matrix_test.rs
@@ -1,4 +1,5 @@
-//! RFC-0031 §13 acceptance: chunk-timer semantics (item 17) and
+//! RFC-0031 §13 acceptance: chunk-timer semantics (item 17),
+//! producer-side timing (consumer delay never counts), and
 //! no-lingering-timer-state after drop (item 21b).
 //!
 //! Paused tokio time makes every assertion exact: auto-advance jumps to the
@@ -119,10 +120,38 @@ async fn semantic_output_resets_the_chunk_timer() {
     );
 }
 
+#[tokio::test(start_paused = true)]
+async fn consumer_delay_does_not_count_against_the_chunk_timer() {
+    // The second delta arrives at t=600, inside the 1000ms chunk budget
+    // measured from the first delta. A consumer that sleeps until t=2000
+    // before polling again must still receive it; only the *next* poll sees
+    // the timeout, which fired at t=1600 on the producer side.
+    let model = ProbeModel {
+        second_part_is_output: true,
+    };
+    let mut result = stream_text(&model, "hello", chunk_options()).await.unwrap();
+    assert!(result.stream.next().await.expect("first part").is_ok());
+
+    tokio::time::sleep(Duration::from_millis(2_000)).await;
+
+    match result.stream.next().await.expect("second part") {
+        Ok(StreamPart::TextDelta { delta, .. }) => assert_eq!(delta, "second"),
+        other => panic!("expected the buffered second delta, got {other:?}"),
+    }
+    match result.stream.next().await.expect("timeout") {
+        Err(AiMuxError::Timeout(message)) => {
+            assert_eq!(message, "Chunk timeout of 1000ms exceeded");
+        }
+        other => panic!("expected chunk timeout, got {other:?}"),
+    }
+    assert!(result.stream.next().await.is_none());
+}
+
 #[tokio::test(start_paused = true)]
 async fn dropped_operation_leaves_no_armed_timer_state() {
-    // §8.0/§13-21b: deadlines live inside the operation future, nothing is
-    // spawned. Dropping the stream mid-flight must leave no timer that could
+    // §8.0/§13-21b: deadlines live inside the pump task that drives the
+    // provider stream, and that task is aborted when the returned stream is
+    // dropped. Dropping the stream mid-flight must leave no timer that could
     // later fire and mutate the caller's signal (the rejected abort_after
     // design would trip this).
     let caller = AbortSignal::new();
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index c46db21f..3c19c58b 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -92,6 +92,7 @@ Core user operation
 | 默认超时 | 与 AI SDK 一致无默认 `total_ms`；Aimux 有意保留非流式 exchange 的 30s whole-response 上限，但从 shared-client 全局配置下移到单次 exchange，流式 exchange 豁免（§5.4） |
 | timeout 输入形状 | 保留跨语言已有的 object/struct 形式，不增加 AI SDK 的裸 number 简写；字段语义一致 |
 | `AbortSignal` | 保留既有的 `CancellationToken` 薄包装，只表示调用方取消；Rust 可直接 drop future，因此 timeout 由 Core deadline/select 表达，不照搬 JS 的 signal merge（§8.0） |
+| stream 驱动 | AI SDK 的 `streamText` 返回前即开始消费 provider stream（push 语义，timer 观察到的是 chunk 到达时刻）；Rust stream 是 pull 语义，只在 `poll_next` 里观察 deadline 会把消费方不 poll 的时间也算进 `first_chunk_ms`/`chunk_ms`。Aimux 在 `stream_text` 返回前 spawn 一个 pump task 驱动 provider stream 进 unbounded channel，deadline 在 pump 侧按到达时刻 arm/reset；返回的 stream drop 时 abort pump。缓冲无上界与 AI SDK 一致（受响应本身约束），pump 在 terminal item 处停止 |
 | SSE framing | AI SDK 的 event-source handler 直接产出解析后的事件；Aimux 复用 `aimux_stream::SseStream` 做同一件事，framing 不留在 Provider |
 | tool timeout | aimux 目前不执行用户 tool，本 RFC 不增加 `tool_ms`/per-tool timeout |
 | observability | 每次 exchange 继续走 Aimux recording/tracing；这不改变 helper 的单次请求语义 |
@@ -649,7 +650,7 @@ impl AbortSignal {
 }
 ```
 
-`aimux-core/src/timeout.rs` 在 operation 开始时记录 total/step deadlines，不 spawn timer task：
+`aimux-core/src/timeout.rs` 在 operation 开始时记录 total/step deadlines，不 spawn timer task（流式阶段的 deadline 由 §8.1 的 pump task 观察，它是唯一 spawn 的 task，随返回的 stream drop 而 abort）：
 
 ```rust
 struct OperationTimeout {
@@ -713,7 +714,13 @@ future，不构造合并 signal，也不把 timeout 写进 `AbortSignal`。外
 
 `first_chunk_ms` 在调用 `do_stream` **之前** arm，因此覆盖 request/setup 以及 §8.3 的
 first-event peek；慢握手、200 后无首帧、以及首条 semantic output 迟到都受同一 budget 约束。
-被 peek 的正常事件会重新接回流首，第一条 semantic output 对用户可见前清除 timer。aimux
+被 peek 的正常事件会重新接回流首，第一条 semantic output 对用户可见前清除 timer。
+
+`do_stream` 返回后，provider stream 由 Core spawn 的 pump task 驱动并写入 unbounded channel，
+`first_chunk_ms`/`chunk_ms`/total/step deadline 都在 pump 侧观察：timer 只度量 provider 输出
+的到达间隔，消费方晚 poll 或处理上一条 item 的时间不计入。已到达但尚未被消费的 output 在
+timeout 之前原样交付，timeout 作为其后的 terminal `Err` 项出现。返回的 stream 被 drop 时
+pump 被 abort，provider stream 随之 drop，取消底层 HTTP exchange。aimux
 目前是单 step，`step_ms` 与 `total_ms` 作用域重合——照样接受并在同一处 arm，以便字段语义
 跨语言一致、将来多 step 时不改 contract。
 
@@ -996,7 +1003,8 @@ handler 负责分帧和解析，peek 与"错误就 `Err` 返回"仍是 provider
 20. polling failure 在已有 job 上 retry，耗尽也不重复 submit；
 21. Router child 在 fallback 前耗尽 retry，MoA reference/aggregator 独立 retry，且不重跑 routing/fanout；
 21a. `AbortSignal` 符合 §8.0：只有 caller cancellation；total/step deadline 直接返回对应 Timeout message，不改变 signal；
-21b. 一个 operation 结束后没有残留 timer/forwarding task（实现不 spawn，测试以 paused time 验证 drop operation 后不再有可触发状态）。
+21b. 一个 operation 结束后没有残留 timer/forwarding task（唯一 spawn 的 stream pump 随返回的 stream drop 而 abort；测试以 paused time 验证 drop operation 后不再有可触发状态）；
+21c. `first_chunk_ms`/`chunk_ms` 只度量 provider 输出到达：消费方在 budget 之外才 poll 时仍收到已按时到达的 output（测试：producer 按时产出、consumer sleep 过 deadline 再 poll）。
 
 ### Bindings/observability
 

From 8fd89fea4b3fbef16982c3901d59048c980a7274 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:14:32 +0800
Subject: [PATCH 15/24] docs(core): link merge_provider_metadata as plain code,
 not an intra-doc link

`generate_video` is public and `merge_provider_metadata` is not, so the
intra-doc link failed `cargo doc` under `-D warnings`
(private_intra_doc_links).
---
 aimux-core/src/video_model.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs
index 57f6472b..4b1bd512 100644
--- a/aimux-core/src/video_model.rs
+++ b/aimux-core/src/video_model.rs
@@ -345,7 +345,7 @@ pub trait VideoModel: Send + Sync {
 /// transient poll failure never re-creates the billed task. Batch results are
 /// flattened back in batch order; provider metadata across batches is
 /// deep-merged the same way start/completion metadata is within a batch (see
-/// [`merge_provider_metadata`]).
+/// `merge_provider_metadata`).
 ///
 /// # Errors
 ///

From f315b861dc29a85d73eadae35697be86e2926bb5 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:19:07 +0800
Subject: [PATCH 16/24] refactor(core): shrink the video batching fix

Same behaviour, less of it:

- inline the batch split, dropping `batch_video_counts` and its `Vec`;
- fold the per-batch aggregation into the first batch's `VideoResult`
  instead of rebuilding one field at a time, which also removes the
  `response.unwrap_or_default()` fallback that could never be taken;
- collapse `merge_provider_metadata`'s nested match into one
  `(remove, b_value)` match, halving its branches;
- replace `BatchedVideoModel` and its two tests with one test on the
  existing `ScriptedVideoModel`. Concurrency is now asserted from a
  peak-in-flight high-water mark rather than a 200ms wall-clock
  threshold, so it is deterministic and does not need a timer.
---
 aimux-core/src/video_model.rs | 279 +++++++++-------------------------
 1 file changed, 75 insertions(+), 204 deletions(-)

diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs
index 4b1bd512..ecc3dbbf 100644
--- a/aimux-core/src/video_model.rs
+++ b/aimux-core/src/video_model.rs
@@ -397,58 +397,35 @@ async fn start_and_poll(
         }
     }
 
+    // Split `n` into `max_per_call`-sized batches (AI SDK `generateVideo` /
+    // `generateImage`): every batch but the last is full-sized, the last one
+    // takes the remainder (or a full batch when `n` divides evenly).
     let max_per_call = model.max_videos_per_call().unwrap_or(options.n).max(1);
-    let batch_counts = batch_video_counts(options.n, max_per_call);
-
-    let batches = batch_counts.into_iter().map(|n| {
+    let last = options.n.div_ceil(max_per_call) - 1;
+    let batches = (0..=last).map(|i| {
         let mut batch_options = options.clone();
-        batch_options.n = n;
+        batch_options.n = if i < last {
+            max_per_call
+        } else {
+            match options.n % max_per_call {
+                0 => max_per_call,
+                remainder => remainder,
+            }
+        };
         start_and_poll_one_batch(model, batch_options, &retries, poll, abort_signal.clone())
     });
-    let results = futures::future::try_join_all(batches).await?;
 
-    let mut videos = Vec::new();
-    let mut warnings = Vec::new();
-    let mut provider_metadata = None;
-    let mut response = None;
+    // `n >= 1` was checked above, so at least one batch ran; the first
+    // batch's `response` represents the call for telemetry purposes.
+    let mut results = futures::future::try_join_all(batches).await?.into_iter();
+    let mut merged = results.next().expect("at least one batch");
     for result in results {
-        videos.extend(result.videos);
-        warnings.extend(result.warnings);
-        provider_metadata = merge_provider_metadata(provider_metadata, result.provider_metadata);
-        response.get_or_insert(result.response);
+        merged.videos.extend(result.videos);
+        merged.warnings.extend(result.warnings);
+        merged.provider_metadata =
+            merge_provider_metadata(merged.provider_metadata, result.provider_metadata);
     }
-
-    Ok(VideoResult {
-        videos,
-        warnings,
-        provider_metadata,
-        // `batch_counts` is never empty (`n >= 1` was checked above), so at
-        // least one batch ran and `response` is always `Some` here; the first
-        // batch's response represents the call for telemetry purposes.
-        response: response.unwrap_or_default(),
-    })
-}
-
-/// Split `n` into `max_per_call`-sized batches (AI SDK `generateVideo` /
-/// `generateImage`): every batch but the last is full-sized; the last batch
-/// takes the remainder, or a full batch if `n` divides evenly.
-fn batch_video_counts(n: u32, max_per_call: u32) -> Vec {
-    let max_per_call = max_per_call.max(1);
-    let call_count = n.div_ceil(max_per_call);
-    (0..call_count)
-        .map(|i| {
-            if i + 1 < call_count {
-                max_per_call
-            } else {
-                let remainder = n % max_per_call;
-                if remainder == 0 {
-                    max_per_call
-                } else {
-                    remainder
-                }
-            }
-        })
-        .collect()
+    Ok(merged)
 }
 
 /// Run one `do_start` + poll cycle for a single provider-sized batch.
@@ -525,35 +502,26 @@ async fn start_and_poll_one_batch(
 
 /// Merge two phases' provider metadata one level deep.
 ///
-/// A plain `HashMap::entry().or_insert()` at the provider-key level drops
-/// every field the other phase set once both phases report the *same*
-/// provider key — e.g. a start call's `job_id` disappearing once the
-/// completion call also reports `x_provider` metadata. Instead, when both
-/// sides have an object for the same provider key, their fields are unioned;
-/// on a field collision the later (`completion`/`b`) value wins. When a
-/// provider key appears on only one side, or either side's value for it is
-/// not a JSON object, that side's value is used as-is — there is nothing to
-/// union.
+/// A plain `entry().or_insert()` at the provider-key level drops every field
+/// the other phase set once both phases report the *same* provider key — e.g.
+/// a start call's `job_id` disappearing once the completion call also reports
+/// metadata under `x_provider`. Objects under a shared key are unioned
+/// instead, with `b` (the later phase) winning a field collision. Anything
+/// that is not a pair of objects has nothing to union, so `b` replaces `a`.
 fn merge_provider_metadata(
     a: Option,
     b: Option,
 ) -> Option {
     let Some(mut merged) = a else { return b };
-    let Some(b) = b else { return Some(merged) };
-    for (provider, b_value) in b {
-        match merged.get_mut(&provider) {
-            Some(a_value) => match (a_value.as_object_mut(), b_value.as_object()) {
-                (Some(a_obj), Some(b_obj)) => {
-                    for (field, value) in b_obj {
-                        a_obj.insert(field.clone(), value.clone());
-                    }
-                }
-                _ => *a_value = b_value,
-            },
-            None => {
-                merged.insert(provider, b_value);
+    for (provider, b_value) in b.into_iter().flatten() {
+        let value = match (merged.remove(&provider), b_value) {
+            (Some(serde_json::Value::Object(mut a_obj)), serde_json::Value::Object(b_obj)) => {
+                a_obj.extend(b_obj);
+                serde_json::Value::Object(a_obj)
             }
-        }
+            (_, b_value) => b_value,
+        };
+        merged.insert(provider, value);
     }
     Some(merged)
 }
@@ -575,6 +543,10 @@ mod tests {
 
     struct ScriptedVideoModel {
         starts: AtomicU32,
+        /// High-water mark of `do_start` calls in flight at once, so a test
+        /// can assert batch concurrency without a wall-clock threshold.
+        peak_concurrent_starts: AtomicU32,
+        in_flight_starts: AtomicU32,
         statuses: AtomicU32,
         script: Vec,
         start_failures: u32,
@@ -586,6 +558,8 @@ mod tests {
         fn new(script: Vec) -> Self {
             Self {
                 starts: AtomicU32::new(0),
+                peak_concurrent_starts: AtomicU32::new(0),
+                in_flight_starts: AtomicU32::new(0),
                 statuses: AtomicU32::new(0),
                 script,
                 start_failures: 0,
@@ -635,6 +609,13 @@ mod tests {
             options: &VideoCallOptions,
         ) -> Result {
             let attempt = self.starts.fetch_add(1, Ordering::SeqCst);
+            let in_flight = self.in_flight_starts.fetch_add(1, Ordering::SeqCst) + 1;
+            self.peak_concurrent_starts
+                .fetch_max(in_flight, Ordering::SeqCst);
+            // Yield so concurrent batches actually overlap here rather than
+            // each running to completion on its first poll.
+            tokio::task::yield_now().await;
+            self.in_flight_starts.fetch_sub(1, Ordering::SeqCst);
             let key = idempotency_key(options);
             self.start_idempotency_keys
                 .lock()
@@ -800,24 +781,18 @@ mod tests {
 
     #[test]
     fn merge_provider_metadata_unions_same_provider_key_across_phases() {
-        let start = SharedProviderMetadata::from([(
-            "fal".to_string(),
-            serde_json::json!({ "job_id": "job-1", "region": "us-east" }),
-        )]);
-        let completion = SharedProviderMetadata::from([(
-            "fal".to_string(),
-            serde_json::json!({ "region": "eu-west", "seed": 42 }),
-        )]);
-
-        let merged = merge_provider_metadata(Some(start), Some(completion))
-            .expect("both phases reported metadata");
-
-        // `job_id` only came from the start phase and must survive the merge
-        // instead of being dropped by an `entry().or_insert()` collision.
+        let phase = |value| Some(SharedProviderMetadata::from([("fal".to_string(), value)]));
+        let merged = merge_provider_metadata(
+            phase(serde_json::json!({ "job_id": "job-1", "region": "us-east" })),
+            phase(serde_json::json!({ "region": "eu-west", "seed": 42 })),
+        )
+        .expect("both phases reported metadata");
+
+        // `job_id` is start-only and must survive instead of being dropped by
+        // an `entry().or_insert()` collision on the shared provider key.
         assert_eq!(
             merged["fal"],
-            serde_json::json!({ "job_id": "job-1", "region": "eu-west", "seed": 42 }),
-            "expected start-only fields preserved and completion to win on collision"
+            serde_json::json!({ "job_id": "job-1", "region": "eu-west", "seed": 42 })
         );
     }
 
@@ -906,138 +881,34 @@ mod tests {
         );
     }
 
-    /// A model whose `do_start` is retriable per-batch (a small delay proves
-    /// batches overlap rather than running one after another) and whose
-    /// `do_status` completes immediately with one video + one provider
-    /// metadata field unique to the batch index.
-    struct BatchedVideoModel {
-        starts: AtomicU32,
-        start_idempotency_keys: Mutex>>,
-        start_delay: Duration,
-        max_per_call: Option,
-    }
-
-    #[async_trait]
-    impl VideoModel for BatchedVideoModel {
-        fn provider(&self) -> &str {
-            "test"
-        }
-        fn model_id(&self) -> &str {
-            "batched"
-        }
-        fn max_videos_per_call(&self) -> Option {
-            self.max_per_call
-        }
-        async fn do_start(
-            &self,
-            options: &VideoCallOptions,
-        ) -> Result {
-            let index = self.starts.fetch_add(1, Ordering::SeqCst);
-            self.start_idempotency_keys
-                .lock()
-                .expect("start key capture lock should not be poisoned")
-                .push(idempotency_key(options));
-            tokio::time::sleep(self.start_delay).await;
-            Ok(VideoOperationStart {
-                operation: serde_json::json!({ "task_id": format!("t-{index}") }),
-                warnings: vec![Warning::Other {
-                    message: format!("batch {index} start"),
-                }],
-                provider_metadata: Some(SharedProviderMetadata::from([(
-                    "test".to_string(),
-                    serde_json::json!({ format!("batch{index}"): index }),
-                )])),
-                response: VideoResponse::default(),
-            })
-        }
-        async fn do_status(
-            &self,
-            operation: &serde_json::Value,
-            _options: &VideoCallOptions,
-        ) -> Result {
-            let task_id = operation["task_id"].as_str().unwrap().to_string();
-            Ok(VideoOperationStatus::Completed(VideoResult {
-                videos: vec![VideoData::Url {
-                    url: format!("https://cdn/{task_id}.mp4"),
-                    media_type: "video/mp4".to_string(),
-                }],
-                warnings: Vec::new(),
-                provider_metadata: None,
-                response: VideoResponse::default(),
-            }))
-        }
-    }
-
     #[tokio::test]
-    async fn n_above_max_per_call_splits_into_batches_and_aggregates() {
-        let model = BatchedVideoModel {
-            starts: AtomicU32::new(0),
-            start_idempotency_keys: Mutex::new(Vec::new()),
-            start_delay: Duration::from_millis(1),
-            max_per_call: Some(1),
-        };
+    async fn n_above_max_per_call_splits_into_concurrent_batches() {
+        // `ScriptedVideoModel::max_videos_per_call()` is 1, so n=3 must run
+        // three independent start/poll cycles, concurrently, each with its
+        // own idempotency key (a batch, not the whole request, is the unit of
+        // idempotent replay).
+        let model = ScriptedVideoModel::new(vec![]);
         let mut options = fast_poll_options();
         options.n = 3;
 
         let result = generate_video(&model, options).await.unwrap();
 
+        assert_eq!(model.starts.load(Ordering::SeqCst), 3);
+        assert_eq!(result.videos.len(), 3);
         assert_eq!(
-            model.starts.load(Ordering::SeqCst),
-            3,
-            "n=3 with max_videos_per_call=1 must run three do_start batches"
+            result.warnings.len(),
+            6,
+            "start + status warnings per batch"
+        );
+        assert!(
+            model.peak_concurrent_starts.load(Ordering::SeqCst) > 1,
+            "batches must overlap, not run one after another"
         );
-        assert_eq!(result.videos.len(), 3);
-        assert_eq!(result.warnings.len(), 3, "warnings from every batch kept");
-
         let keys = model
             .start_idempotency_keys
             .lock()
             .expect("start key capture lock should not be poisoned");
-        assert_eq!(keys.len(), 3);
-        assert!(keys.iter().all(Option::is_some));
         let unique: std::collections::HashSet<_> = keys.iter().collect();
-        assert_eq!(
-            unique.len(),
-            3,
-            "each batch must mint its own idempotency key: {keys:?}"
-        );
-
-        // Provider metadata across batches is deep-merged (§ merge_provider_metadata),
-        // not `entry().or_insert()`-dropped: all three batch-specific fields survive
-        // under the shared "test" provider key.
-        let metadata = result
-            .provider_metadata
-            .expect("provider metadata aggregated");
-        let test_meta = metadata.get("test").expect("test provider metadata");
-        for i in 0..3 {
-            assert!(
-                test_meta.get(format!("batch{i}")).is_some(),
-                "batch {i}'s metadata field missing from aggregate: {test_meta:?}"
-            );
-        }
-    }
-
-    #[tokio::test]
-    async fn batches_run_concurrently_not_sequentially() {
-        let model = BatchedVideoModel {
-            starts: AtomicU32::new(0),
-            start_idempotency_keys: Mutex::new(Vec::new()),
-            start_delay: Duration::from_millis(80),
-            max_per_call: Some(1),
-        };
-        let mut options = fast_poll_options();
-        options.n = 3;
-
-        let started = Instant::now();
-        generate_video(&model, options).await.unwrap();
-        let elapsed = started.elapsed();
-
-        // Sequential batches would take >= 3 * 80ms = 240ms; concurrent
-        // batches take roughly one delay (~80ms) plus scheduling overhead.
-        assert!(
-            elapsed < Duration::from_millis(200),
-            "batches should overlap (AI SDK generateVideo/generateImage runs \
-             them concurrently via Promise.all), took {elapsed:?}"
-        );
+        assert_eq!(unique.len(), 3, "one key per batch: {keys:?}");
     }
 }

From 6b187c94a829f5aa7aaf0424c3d4dd5df81496e5 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:19:07 +0800
Subject: [PATCH 17/24] refactor(providers,docs): de-duplicate the
 retry-placement rationale

The six review fixes repeated the same rationale verbatim in five
provider modules and restated the code in three doc sections. Keep one
short statement of the invariant per site and let the spec carry the
argument; the deviation table rows and acceptance items are untouched.

Test collapse, one regression test per finding:

- drop the Anthropic files retry test (same single-exchange shape and
  same finding as the OpenAI one, different fixture data only) and the
  Gladia per-stage test (same upload-then-submit shape as AssemblyAI's).
  Google Files keeps its own test: three-stage independent retry is a
  distinct contract.
- give `google_files_test` one `provider_with_retries` helper instead of
  three inline `GoogleConfig` builders.
---
 aimux-providers/src/anthropic/files.rs        | 10 ++-
 aimux-providers/src/assemblyai.rs             | 11 +---
 aimux-providers/src/gladia.rs                 | 11 +---
 aimux-providers/src/google/files.rs           | 21 ++----
 aimux-providers/src/openai/files.rs           | 10 ++-
 aimux-providers/tests/anthropic_files_test.rs | 40 ------------
 .../tests/gladia_transcription_test.rs        | 64 -------------------
 aimux-providers/tests/google_files_test.rs    | 45 +++++--------
 docs/ai-sdk-request-pipeline.md               | 34 ++++------
 9 files changed, 46 insertions(+), 200 deletions(-)

diff --git a/aimux-providers/src/anthropic/files.rs b/aimux-providers/src/anthropic/files.rs
index 9d3cf309..2f08101b 100644
--- a/aimux-providers/src/anthropic/files.rs
+++ b/aimux-providers/src/anthropic/files.rs
@@ -131,12 +131,10 @@ impl Files for AnthropicFiles {
         // sets `Content-Type` from it, so it is intentionally not added to the
         // header list above.
         //
-        // `upload_file` is a single-exchange Provider SPI operation (no Core
-        // `do_upload_file` wrapper, RFC-0031 §9.4), so nothing above this call
-        // retries it. Apply the Core retry primitive here, same as
-        // `execute_list_models`: the upload itself is not billable, and a
-        // failed create-a-file request never returns a file id to replay
-        // against, so a plain retry of the whole exchange is safe.
+        // Nothing above `upload_file` retries it (there is no Core
+        // `do_upload_file`), so the retry lives here — safe for the whole
+        // exchange because an upload is not billable and a failed
+        // create-a-file request returns no id to replay against (§9.4).
         let retries = aimux_core::retry::prepare_retries(
             None,
             self.config.retry_config,
diff --git a/aimux-providers/src/assemblyai.rs b/aimux-providers/src/assemblyai.rs
index dc5b6868..3bae0555 100644
--- a/aimux-providers/src/assemblyai.rs
+++ b/aimux-providers/src/assemblyai.rs
@@ -236,15 +236,10 @@ impl TranscriptionModel for AssemblyAITranscriptionModel {
         let headers = self.build_headers(options.headers.as_ref());
         let header_list: Vec<(String, String)> = headers.into_iter().collect();
 
-        // Core wraps the whole `do_generate` in one retry (RFC-0031 §6.2), so
-        // without per-stage retries here, a transient failure in a later
-        // stage would replay every earlier stage — re-uploading the audio to
-        // retry a submit or poll failure. Each stage gets its own retry
-        // against the provider's configured retry settings instead; an
+        // Per-stage retry, not one retry around `do_generate`: Core's outer
+        // retry would re-upload the audio to retry a later stage. An
         // exhausted inner retry returns `AiMuxError::Retry`, which the outer
-        // Core retry passes through unchanged rather than re-wrapping
-        // (`retry_with_exponential_backoff`'s `Err(AiMuxError::Retry(_))`
-        // arm), so `do_generate` is never replayed either.
+        // retry passes through, so `do_generate` is never replayed (§6.2).
         let retries = retry::prepare_retries(
             options.max_retries,
             self.retry_config(),
diff --git a/aimux-providers/src/gladia.rs b/aimux-providers/src/gladia.rs
index c0865647..fbaa7af6 100644
--- a/aimux-providers/src/gladia.rs
+++ b/aimux-providers/src/gladia.rs
@@ -221,15 +221,10 @@ impl TranscriptionModel for GladiaTranscriptionModel {
 
         let headers = self.build_headers(options.headers.as_ref());
 
-        // Core wraps the whole `do_generate` in one retry (RFC-0031 §6.2), so
-        // without per-stage retries here, a transient failure in a later
-        // stage would replay every earlier stage — re-uploading the audio to
-        // retry an initiate or poll failure. Each stage gets its own retry
-        // against the provider's configured retry settings instead; an
+        // Per-stage retry, not one retry around `do_generate`: Core's outer
+        // retry would re-upload the audio to retry a later stage. An
         // exhausted inner retry returns `AiMuxError::Retry`, which the outer
-        // Core retry passes through unchanged rather than re-wrapping
-        // (`retry_with_exponential_backoff`'s `Err(AiMuxError::Retry(_))`
-        // arm), so `do_generate` is never replayed either.
+        // retry passes through, so `do_generate` is never replayed (§6.2).
         let retries = retry::prepare_retries(
             options.max_retries,
             self.retry_config(),
diff --git a/aimux-providers/src/google/files.rs b/aimux-providers/src/google/files.rs
index 828bcc79..d0256391 100644
--- a/aimux-providers/src/google/files.rs
+++ b/aimux-providers/src/google/files.rs
@@ -189,13 +189,10 @@ impl Files for GoogleFiles {
         ));
         init_headers.push(("Content-Type".to_string(), "application/json".to_string()));
 
-        // `upload_file` is a single-exchange Provider SPI operation (no Core
-        // `do_upload_file` wrapper, RFC-0031 §9.4), so nothing above this call
-        // retries any of its three stages. Each stage is retried independently
-        // (same Core primitive `execute_list_models` uses) rather than
-        // wrapping the whole upload: a transient failure in the upload or
-        // poll stage must not replay a prior stage — the init exchange that
-        // minted `upload_url`, or the potentially large upload body itself.
+        // Nothing above `upload_file` retries it (§9.4), so the retry lives
+        // here — per stage, not around the whole upload: a failure in the
+        // upload or poll stage must not replay the init exchange that minted
+        // `upload_url`, nor resend the file body.
         let retries = aimux_core::retry::prepare_retries(
             None,
             self.config.retry_config,
@@ -231,8 +228,7 @@ impl Files for GoogleFiles {
                     super::google_failed_response_handler(),
                 )
                 .await
-                // Applied per attempt (not just on final exhaustion) so a
-                // `RetryError`'s saved `errors` also carry this context.
+                // Per attempt, so a `RetryError`'s saved `errors` carry it too.
                 .map_err(|e| match e {
                     AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError {
                         message: format!("Failed to initiate resumable upload: {}", d.message),
@@ -266,9 +262,6 @@ impl Files for GoogleFiles {
         // The upload URL comes from the init response's x-goog-upload-url
         // header and receives the user's file bytes; validate it. (AI SDK
         // fetches this URL unvalidated — kept stricter here deliberately.)
-        // Retried independently of the init stage above (§ note there): a
-        // retry here resends the same bytes to the same `upload_url`, never
-        // re-mints one.
         let upload_resp = retries
             .retry(|| async {
                 aimux_provider_utils::post_to_api(
@@ -290,7 +283,6 @@ impl Files for GoogleFiles {
                     super::google_failed_response_handler(),
                 )
                 .await
-                // Applied per attempt, same reasoning as the init stage above.
                 .map_err(|e| match e {
                     AiMuxError::ApiCall(d) => AiMuxError::ApiCall(Box::new(ApiCallError {
                         message: format!("Failed to upload file data: {}", d.message),
@@ -333,9 +325,6 @@ impl Files for GoogleFiles {
 
             let poll_url = format!("{}/{}", self.config.base_url, file.name);
 
-            // A status GET is idempotent, so it gets the same independent
-            // retry as the init/upload stages; a failure here never re-runs
-            // the upload.
             let poll_resp = retries
                 .retry(|| {
                     aimux_provider_utils::get_from_api(
diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs
index 00a39b43..2574cd76 100644
--- a/aimux-providers/src/openai/files.rs
+++ b/aimux-providers/src/openai/files.rs
@@ -168,12 +168,10 @@ impl Files for OpenAIFiles {
         // sets `Content-Type` from it, so it is intentionally not added to the
         // header list above.
         //
-        // `upload_file` is a single-exchange Provider SPI operation (no Core
-        // `do_upload_file` wrapper, RFC-0031 §9.4), so nothing above this call
-        // retries it. Apply the Core retry primitive here, same as
-        // `execute_list_models`: the upload itself is not billable, and a
-        // failed create-a-file request never returns a file id to replay
-        // against, so a plain retry of the whole exchange is safe.
+        // Nothing above `upload_file` retries it (there is no Core
+        // `do_upload_file`), so the retry lives here — safe for the whole
+        // exchange because an upload is not billable and a failed
+        // create-a-file request returns no id to replay against (§9.4).
         let retries = aimux_core::retry::prepare_retries(
             None,
             self.config.retry_config,
diff --git a/aimux-providers/tests/anthropic_files_test.rs b/aimux-providers/tests/anthropic_files_test.rs
index 7153f09f..4f6127f7 100644
--- a/aimux-providers/tests/anthropic_files_test.rs
+++ b/aimux-providers/tests/anthropic_files_test.rs
@@ -308,43 +308,3 @@ async fn has_correct_provider_name() {
 
     assert_eq!(files.provider(), "anthropic.files");
 }
-
-/// A transient 503 followed by 200 must succeed: `upload_file` retries the
-/// upload exchange using the provider's configured retry settings, the same
-/// way `list_models` already does.
-#[tokio::test]
-async fn transient_failure_is_retried_and_succeeds() {
-    let server = MockServer::start().await;
-    let attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
-    let responder_attempts = std::sync::Arc::clone(&attempts);
-    Mock::given(method("POST"))
-        .and(path("/v1/files"))
-        .respond_with(move |_: &wiremock::Request| {
-            if responder_attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
-                ResponseTemplate::new(503)
-                    .insert_header("retry-after-ms", "0")
-                    .set_body_json(json!({"error": {"message": "try again"}}))
-            } else {
-                ResponseTemplate::new(200).set_body_json(file_response_body())
-            }
-        })
-        .mount(&server)
-        .await;
-
-    let config = AnthropicConfig::new("test-api-key")
-        .with_base_url(server.uri())
-        .with_retry_config(aimux_core::retry::RetryConfig {
-            max_retries: 1,
-            ..Default::default()
-        });
-    let provider = AnthropicProvider::new(config);
-    let files = provider.files();
-
-    let result = files.upload_file(&upload_options()).await.unwrap();
-
-    assert_eq!(
-        result.provider_reference.get("anthropic"),
-        Some(&"file-abc123".to_string())
-    );
-    assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2);
-}
diff --git a/aimux-providers/tests/gladia_transcription_test.rs b/aimux-providers/tests/gladia_transcription_test.rs
index ef46e4bd..ab4ee57e 100644
--- a/aimux-providers/tests/gladia_transcription_test.rs
+++ b/aimux-providers/tests/gladia_transcription_test.rs
@@ -183,67 +183,3 @@ async fn should_include_provider_metadata() {
             .contains_key("gladia")
     );
 }
-
-/// A transient 503 on the initiate stage, followed by 200, must succeed
-/// without re-uploading the audio: the upload exchange is retried
-/// independently of the initiate exchange, so Core's outer per-`do_generate`
-/// retry never needs to (and must not) replay the upload.
-#[tokio::test]
-async fn transient_initiate_failure_is_retried_without_re_uploading() {
-    let server = MockServer::start().await;
-
-    let upload_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
-    let upload_observed = std::sync::Arc::clone(&upload_attempts);
-    Mock::given(method("POST"))
-        .and(path("/v2/upload"))
-        .respond_with(move |_: &wiremock::Request| {
-            upload_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
-            ResponseTemplate::new(200)
-                .set_body_json(json!({"audio_url": "https://upload.gladia.io/test"}))
-        })
-        .mount(&server)
-        .await;
-
-    let initiate_attempts = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
-    let initiate_observed = std::sync::Arc::clone(&initiate_attempts);
-    let result_url = format!("{}/v2/result/test", server.uri());
-    Mock::given(method("POST"))
-        .and(path("/v2/pre-recorded"))
-        .respond_with(move |_: &wiremock::Request| {
-            if initiate_observed.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0 {
-                ResponseTemplate::new(503)
-                    .insert_header("retry-after-ms", "0")
-                    .set_body_json(json!({"error": "try again"}))
-            } else {
-                ResponseTemplate::new(200).set_body_json(json!({"result_url": result_url}))
-            }
-        })
-        .mount(&server)
-        .await;
-
-    Mock::given(method("GET"))
-        .and(path("/v2/result/test"))
-        .respond_with(ResponseTemplate::new(200).set_body_json(fixture_result()))
-        .mount(&server)
-        .await;
-
-    let config = GladiaConfig::new("test-api-key").with_base_url(server.uri());
-    let provider = GladiaProvider::new(config);
-    let model = provider.transcription("default");
-
-    let result = model
-        .do_generate(&options(mock_audio(), "audio/wav"))
-        .await
-        .unwrap();
-
-    assert_eq!(result.text, "Hello from Gladia.");
-    assert_eq!(
-        upload_attempts.load(std::sync::atomic::Ordering::SeqCst),
-        1,
-        "the upload stage must not be replayed by an initiate-stage retry"
-    );
-    assert_eq!(
-        initiate_attempts.load(std::sync::atomic::Ordering::SeqCst),
-        2
-    );
-}
diff --git a/aimux-providers/tests/google_files_test.rs b/aimux-providers/tests/google_files_test.rs
index 789f8199..17f6a132 100644
--- a/aimux-providers/tests/google_files_test.rs
+++ b/aimux-providers/tests/google_files_test.rs
@@ -15,6 +15,7 @@ use wiremock::matchers::{method, path};
 use wiremock::{Mock, MockServer, ResponseTemplate};
 
 use aimux_core::files_model::{Files, UploadFileCallOptions, UploadFileData};
+use aimux_core::retry::RetryConfig;
 use aimux_core::shared::{FileBytes, SharedProviderOptions};
 use aimux_providers::{GoogleConfig, GoogleProvider};
 
@@ -36,8 +37,16 @@ fn default_file_resource() -> Value {
 }
 
 fn provider(server: &MockServer) -> GoogleProvider {
-    let config =
-        GoogleConfig::new("test-api-key").with_base_url(format!("{}/v1beta", server.uri()));
+    provider_with_retries(server, RetryConfig::default().max_retries)
+}
+
+fn provider_with_retries(server: &MockServer, max_retries: u32) -> GoogleProvider {
+    let config = GoogleConfig::new("test-api-key")
+        .with_base_url(format!("{}/v1beta", server.uri()))
+        .with_retry_config(RetryConfig {
+            max_retries,
+            ..Default::default()
+        });
     GoogleProvider::new(config)
 }
 
@@ -401,16 +410,8 @@ async fn should_throw_when_initiation_request_fails() {
         .mount(&server)
         .await;
 
-    // A persistent failure exhausts retries either way; pin max_retries to 0
-    // so this test (which only checks the message text, not retry
-    // exhaustion) stays fast.
-    let config = GoogleConfig::new("test-api-key")
-        .with_base_url(format!("{}/v1beta", server.uri()))
-        .with_retry_config(aimux_core::retry::RetryConfig {
-            max_retries: 0,
-            ..Default::default()
-        });
-    let provider = GoogleProvider::new(config);
+    // Only the message text matters here; skip the retries so it stays fast.
+    let provider = provider_with_retries(&server, 0);
     let files = provider.files();
 
     let result = files.upload_file(&upload_options()).await;
@@ -465,16 +466,8 @@ async fn should_throw_when_upload_request_fails() {
         .mount(&server)
         .await;
 
-    // A persistent failure exhausts retries either way; pin max_retries to 0
-    // so this test (which only checks the message text, not retry
-    // exhaustion) stays fast.
-    let config = GoogleConfig::new("test-api-key")
-        .with_base_url(format!("{}/v1beta", server.uri()))
-        .with_retry_config(aimux_core::retry::RetryConfig {
-            max_retries: 0,
-            ..Default::default()
-        });
-    let provider = GoogleProvider::new(config);
+    // Only the message text matters here; skip the retries so it stays fast.
+    let provider = provider_with_retries(&server, 0);
     let files = provider.files();
 
     let result = files.upload_file(&upload_options()).await;
@@ -640,13 +633,7 @@ async fn transient_upload_failure_is_retried_without_re_initiating() {
         .mount(&server)
         .await;
 
-    let config = GoogleConfig::new("test-api-key")
-        .with_base_url(format!("{}/v1beta", server.uri()))
-        .with_retry_config(aimux_core::retry::RetryConfig {
-            max_retries: 1,
-            ..Default::default()
-        });
-    let provider = GoogleProvider::new(config);
+    let provider = provider_with_retries(&server, 1);
     let files = provider.files();
 
     let result = files.upload_file(&upload_options()).await.unwrap();
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index 3c19c58b..4ec6f354 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -716,11 +716,9 @@ future，不构造合并 signal，也不把 timeout 写进 `AbortSignal`。外
 first-event peek；慢握手、200 后无首帧、以及首条 semantic output 迟到都受同一 budget 约束。
 被 peek 的正常事件会重新接回流首，第一条 semantic output 对用户可见前清除 timer。
 
-`do_stream` 返回后，provider stream 由 Core spawn 的 pump task 驱动并写入 unbounded channel，
-`first_chunk_ms`/`chunk_ms`/total/step deadline 都在 pump 侧观察：timer 只度量 provider 输出
-的到达间隔，消费方晚 poll 或处理上一条 item 的时间不计入。已到达但尚未被消费的 output 在
-timeout 之前原样交付，timeout 作为其后的 terminal `Err` 项出现。返回的 stream 被 drop 时
-pump 被 abort，provider stream 随之 drop，取消底层 HTTP exchange。aimux
+`do_stream` 返回后所有 deadline 都由 pump task 观察（见“stream 驱动”偏差行）：timer 只度量
+provider 输出的到达间隔，消费方晚 poll 的时间不计入；已到达但尚未被消费的 output 在 timeout
+之前原样交付，timeout 作为其后的 terminal `Err` 项出现。aimux
 目前是单 step，`step_ms` 与 `total_ms` 作用域重合——照样接受并在同一处 arm，以便字段语义
 跨语言一致、将来多 step 时不改 contract。
 
@@ -773,15 +771,10 @@ start/status 流），Core 拥有 poll 循环并分别 retry 两个阶段：
 - Provider-specific polling delay 不是通用 exponential retry；
 - 耗尽产生的 `RetryError` 原样透传，外层不得再次 submit。
 - **修订**（RFC-0031 review）：`options.n` 按 `VideoModel::max_videos_per_call`
-  切成多个 batch（AI SDK `generateVideo`/`generateImage` 同款算法：除最后一个
-  batch 外每个都是满批，最后一个是余数或整批），每个 batch **并发**跑一次
-  完整的独立 `do_start`/poll 流程（AI SDK 用 `Promise.all`；Rust 侧用
-  `futures::future::try_join_all`，第一个 batch 失败时其余尚在进行的 batch
-  future 被 drop，这点比 JS 更激进但不影响正确性——没有 batch 会重连或取消
-  已经发出的 provider 任务）——每个 batch 独立铸造 idempotency key。结果按
-  batch 顺序拼接 `videos`/`warnings`；`provider_metadata` 跨 batch 用与
-  start/completion 相同的 deep-merge 策略聚合。`n == 0` 在任何网络调用之前
-  返回 `InvalidArgument`。
+  切 batch，每个 batch 并发跑一次独立的 `do_start`/poll 流程并独立铸造
+  idempotency key（AI SDK 用 `Promise.all`，Rust 侧 `try_join_all` 会在首个
+  失败时 drop 其余 batch future——不影响正确性，两边都不会重连或取消已发出的
+  provider 任务）。`n == 0` 在任何网络调用之前返回 `InvalidArgument`。
 
 ### 9.2 Router / MoA
 
@@ -806,15 +799,10 @@ Composite 外层不重放；重试放在已有语义边界内：
 
 `Files::upload_file` 按 AI SDK Provider SPI 本身就是公开 operation，没有对应的
 `do_upload_file`；本轮不为它虚构第二层。**修订**（RFC-0031 review）：`upload_file`
-现在用 Provider 已配置的 `RetryConfig`（`aimux_core::retry::prepare_retries` +
-`PreparedRetries::retry`，与 `execute_list_models` 同一 primitive）包住单次 HTTP
-exchange——upload 不计费，失败的 create-file 请求不会返回可供下一次请求复用的
-file id，因此对整个 exchange 做 plain retry 是安全的；这与非幂等的“创建资源”类
-POST 通常不能自动重试的一般规则不同，此处的例外前提就是“不计费 + 无副作用可
-重放”。OpenAI 与 Anthropic 各是一次 exchange，直接整体重试。Google Files 是
-init/upload/poll 三阶段：三个 exchange **分别独立**重试，而不是包住整个
-`upload_file`——upload 阶段的重试只重发文件字节到已经拿到的 `upload_url`，不会
-重新调用 init 铸造新 URL；poll 阶段的重试只重发状态 GET，不会触发 upload。
+自己用 Provider 已配置的 `RetryConfig` 重试（与 `execute_list_models` 同一
+primitive）。这是“创建资源类 POST 不自动重试”一般规则的显式例外，前提是
+upload 不计费、失败的 create-file 请求不会返回可复用的 file id。多阶段的
+Google Files 按 exchange 分别重试，而不是包住整个 `upload_file`。
 
 ---
 

From edd98643cba3296c080bc5c672a16a12ffdf3d40 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:20:48 +0800
Subject: [PATCH 18/24] refactor(core,providers): trim the pump rationale; drop
 a whole-body clone
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The pump comment restated the loop it sits on; keep the two points the
code cannot show (why the channel is unbounded, why the pump is spawned
even with nothing armed).

xAI's successful-response handler parses the body into a `Value`, hands
it back as `raw_value`, and used `from_value(raw.clone())` to build `T` —
a second full tree resident at peak, the same amplification the JSON
size-cap fix removed from `create_json_response_handler`. Deserialize
from a borrow instead.
---
 aimux-core/src/generate.rs     | 27 +++++++++++----------------
 aimux-providers/src/xai/mod.rs |  5 ++++-
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/aimux-core/src/generate.rs b/aimux-core/src/generate.rs
index 9db604db..0b545711 100644
--- a/aimux-core/src/generate.rs
+++ b/aimux-core/src/generate.rs
@@ -1017,23 +1017,18 @@ pub async fn stream_text(
     // first-chunk budget armed at operation start continues until the first
     // output chunk.
     let operation_deadline = operation_timeout.deadline();
-    // The provider stream is driven by a pump task, not by the consumer's
-    // polls: first-chunk and chunk-idle deadlines measure when provider
-    // output *arrives*, so a caller that polls late still receives output the
-    // provider delivered on time. Pull-based streams cannot express this
-    // without a driver — a deadline observed only inside `poll_next` keeps
-    // elapsing while nobody polls. The channel is unbounded on purpose: the
-    // AI SDK's `streamText` also consumes eagerly, and a bounded channel
-    // would make the pump block on the consumer, reintroducing consumer
-    // latency into the provider timers. Buffering is bounded by the response
-    // itself; the pump stops at the terminal item and is aborted when the
-    // returned stream is dropped.
+    // A pump task drives the provider stream, so the deadlines below measure
+    // when provider output *arrives*: a deadline observed only inside
+    // `poll_next` keeps elapsing while nobody polls, charging the provider
+    // for the consumer's latency. The channel is unbounded so the pump never
+    // blocks on the consumer (which would put that latency back into the
+    // timers); the AI SDK's `streamText` consumes eagerly for the same
+    // reason, and buffering is bounded by the response itself.
     //
-    // Always wrapped, even with no abort signal or deadlines armed: the pump
-    // is also the stream's terminal fuse — `Finish` and non-recoverable
-    // errors must end the stream under every configuration, or a provider
-    // that keeps the connection open after `Finish` hangs a consumer
-    // reading to end-of-stream.
+    // Spawned even with no abort signal or deadlines armed: the pump is also
+    // the stream's terminal fuse — `Finish` and non-recoverable errors must
+    // end the stream under every configuration, or a provider that keeps the
+    // connection open after `Finish` hangs a consumer reading to end-of-stream.
     let stream: Pin> + Send>> = {
         let mut stream = stream;
         let chunk_ms = stream_timeout.chunk_ms;
diff --git a/aimux-providers/src/xai/mod.rs b/aimux-providers/src/xai/mod.rs
index 0b180ece..104ea039 100644
--- a/aimux-providers/src/xai/mod.rs
+++ b/aimux-providers/src/xai/mod.rs
@@ -133,7 +133,10 @@ where
                 ..aimux_core::ApiCallError::new(message, url, request_body_values)
             })));
         }
-        let value = serde_json::from_value(raw.clone()).map_err(|error| {
+        // Borrow `raw` instead of `from_value(raw.clone())`: the whole body is
+        // already a `Value` here and is handed back as `raw_value`, so the
+        // clone was a second full tree resident at peak.
+        let value = serde::Deserialize::deserialize(&raw).map_err(|error: serde_json::Error| {
             AiMuxError::ApiCall(Box::new(aimux_core::ApiCallError {
                 status_code: Some(status),
                 response_body: Some(raw.to_string()),

From 50876531abd7613043dbdc5677b172a9464345cd Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:26:57 +0800
Subject: [PATCH 19/24] fix(bindings): keep the explicit query in Node/Python
 search()

`transcribe` and `rerank` restore the fields their explicit arguments
own after deserializing `opts_json` over the options struct; `search`
replaced the whole struct instead, so `search("cats", {"query":"dogs"})`
silently searched for dogs.
---
 bindings/node/src/multimodal.rs   | 5 ++++-
 bindings/python/src/multimodal.rs | 5 ++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/bindings/node/src/multimodal.rs b/bindings/node/src/multimodal.rs
index 665cf32d..6428e1b2 100644
--- a/bindings/node/src/multimodal.rs
+++ b/bindings/node/src/multimodal.rs
@@ -301,7 +301,10 @@ impl SearchModel {
                 let mut opts = SearchCallOptions::new(query);
                 if let Some(s) = opts_json.as_deref() {
                     if !s.trim().is_empty() && s.trim() != "null" {
-                        let parsed: SearchCallOptions = parse_wire_json("opts_json", s)?;
+                        // Required data comes from the explicit arguments; all
+                        // operation policy and optional fields come from JSON.
+                        let mut parsed: SearchCallOptions = parse_wire_json("opts_json", s)?;
+                        parsed.query = opts.query;
                         opts = parsed;
                     }
                 }
diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs
index 3ca843c8..47bf948c 100644
--- a/bindings/python/src/multimodal.rs
+++ b/bindings/python/src/multimodal.rs
@@ -258,7 +258,10 @@ impl SearchModel {
         let mut opts = SearchCallOptions::new(query.to_string());
         if let Some(s) = opts_json {
             if !s.trim().is_empty() && s.trim() != "null" {
-                let parsed: SearchCallOptions = wire_json("opts_json", s)?;
+                // Take every caller option; the query always comes from the
+                // explicit arg.
+                let mut parsed: SearchCallOptions = wire_json("opts_json", s)?;
+                parsed.query = opts.query;
                 opts = parsed;
             }
         }

From 26ccdac3379eba1e6b1d29b447ba7b804591e94a Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:26:57 +0800
Subject: [PATCH 20/24] fix(java): type VideoCallOptions.fps as Long, not
 Double

`fps` is `Option` in Rust, and serde rejects a JSON float for it,
so Jackson's `24.0` failed to deserialize the *whole* `VideoCallOptions`
at the FFI boundary. Every other integer field in this struct
(`duration`, `seed`) is already `Long`.
---
 .../java/ai/arcships/aimux/MultimodalTypes.java   | 10 +++++-----
 .../ai/arcships/aimux/MultimodalTypesTest.java    | 15 +++++++++++++++
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java b/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java
index d310ca38..34fc9de7 100644
--- a/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java
+++ b/bindings/java/src/main/java/ai/arcships/aimux/MultimodalTypes.java
@@ -2541,7 +2541,7 @@ public static class VideoCallOptions {
         @JsonProperty("aspect_ratio") private String aspectRatio;
         @JsonProperty("resolution") private String resolution;
         @JsonProperty("duration") private Long duration;
-        @JsonProperty("fps") private Double fps;
+        @JsonProperty("fps") private Long fps;
         @JsonProperty("seed") private Long seed;
         @JsonProperty("image") private VideoFile image;
         @JsonProperty("frame_images") private List frameImages;
@@ -2557,7 +2557,7 @@ public static class VideoCallOptions {
         VideoCallOptions() {}
 
         private VideoCallOptions(String prompt, Integer n, String aspectRatio, String resolution,
-                                 Long duration, Double fps, Long seed,
+                                 Long duration, Long fps, Long seed,
                                  VideoFile image, List frameImages,
                                  List inputReferences, Boolean generateAudio,
                                  JsonNode providerOptions, Map headers,
@@ -2591,7 +2591,7 @@ public static VideoCallOptions of(String prompt) {
         public String getAspectRatio() { return aspectRatio; }
         public String getResolution() { return resolution; }
         public Long getDuration() { return duration; }
-        public Double getFps() { return fps; }
+        public Long getFps() { return fps; }
         public Long getSeed() { return seed; }
         public VideoFile getImage() { return image; }
         public List getFrameImages() { return frameImages; }
@@ -2611,7 +2611,7 @@ public static class Builder {
             private String aspectRatio;
             private String resolution;
             private Long duration;
-            private Double fps;
+            private Long fps;
             private Long seed;
             private VideoFile image;
             private List frameImages;
@@ -2628,7 +2628,7 @@ public static class Builder {
             public Builder aspectRatio(String v) { this.aspectRatio = v; return this; }
             public Builder resolution(String v) { this.resolution = v; return this; }
             public Builder duration(Long v) { this.duration = v; return this; }
-            public Builder fps(Double v) { this.fps = v; return this; }
+            public Builder fps(Long v) { this.fps = v; return this; }
             public Builder seed(Long v) { this.seed = v; return this; }
             public Builder image(VideoFile v) { this.image = v; return this; }
             public Builder frameImages(List v) { this.frameImages = v; return this; }
diff --git a/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java b/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java
index b34f50e0..99f709af 100644
--- a/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java
+++ b/bindings/java/src/test/java/ai/arcships/aimux/MultimodalTypesTest.java
@@ -253,6 +253,21 @@ void videoPollOptionsRoundTrip() throws Exception {
         assertThat(decoded).isEqualTo(options);
     }
 
+    @Test
+    void videoFpsIsEmittedAsAnIntegerForRustU32() throws Exception {
+        // `VideoCallOptions.fps` is `Option` in Rust; a JSON float makes
+        // serde reject the whole options object at the FFI boundary.
+        MultimodalTypes.VideoCallOptions options = MultimodalTypes.VideoCallOptions.builder()
+            .prompt("a cat")
+            .fps(24L)
+            .build();
+
+        String json = M.writeValueAsString(options);
+        assertThat(M.readTree(json).path("fps").isIntegralNumber())
+            .as("wire form: %s", json)
+            .isTrue();
+    }
+
     // ── Search ─────────────────────────────────────────────────────────────
 
     @Test

From 6a644cca7a2a934ad0343c54721d573980c24e08 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:26:57 +0800
Subject: [PATCH 21/24] fix(web): write fixtures with RECORDING_SCHEMA, not a
 hardcoded 1

`RECORDING_SCHEMA` is 2 since the exchange records gained
`step`/`attempt`/`exchange_index`, so both fixture generators were
stamping schema-2 payloads with schema 1.
---
 tools/aimux-web/examples/gen_demo_fixtures.rs | 2 +-
 tools/aimux-web/examples/gen_fixture.rs       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tools/aimux-web/examples/gen_demo_fixtures.rs b/tools/aimux-web/examples/gen_demo_fixtures.rs
index 499fdfcf..7a0792c9 100644
--- a/tools/aimux-web/examples/gen_demo_fixtures.rs
+++ b/tools/aimux-web/examples/gen_demo_fixtures.rs
@@ -51,7 +51,7 @@ fn recording(
         ..Default::default()
     };
     Recording {
-        schema: 1,
+        schema: aimux_core::recording::RECORDING_SCHEMA,
         call_id: call_id.into(),
         recorded_at: recorded_at.into(),
         input: InputRecord {
diff --git a/tools/aimux-web/examples/gen_fixture.rs b/tools/aimux-web/examples/gen_fixture.rs
index 533367c6..55e0bc55 100644
--- a/tools/aimux-web/examples/gen_fixture.rs
+++ b/tools/aimux-web/examples/gen_fixture.rs
@@ -40,7 +40,7 @@ fn main() {
     };
 
     let rec = Recording {
-        schema: 1,
+        schema: aimux_core::recording::RECORDING_SCHEMA,
         call_id: "call-fixture-1".into(),
         recorded_at: "2026-08-14T00:00:00.000Z".into(),
         input: InputRecord {

From 7ec05c79b68e45e3ac0889a038ec314d2be158d6 Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:40:22 +0800
Subject: [PATCH 22/24] fix(bindings): let Node/Python pass options to
 transcribe/rerank/search
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

These three methods take their required data as explicit arguments and
everything else as `opts_json`, but they deserialized the caller's
options object into the full options struct first — and `audio`,
`media_type`, `query` and `documents` are not optional there. Any
non-empty `opts_json` therefore failed with `missing field` before the
explicit arguments could be restored, so `max_retries` and `timeout`
were unreachable from Node and Python for all three modalities, which is
exactly what this branch set out to expose.

Fill the argument-owned fields in before deserializing. The values
inserted are shape-only placeholders that the existing `parsed.x =
opts.x` lines overwrite, so nothing large is copied; they are built from
the real enums (`AudioInput::Base64`, `RerankingDocuments::Text`) so a
variant rename fails to compile instead of at runtime.

`opts_with_args` is deliberately free of binding error types: it is the
only part with logic, and keeping it pure makes it unit-testable without
a Python interpreter or a Node runtime (the pyo3 test binary cannot link
libpython here). The wire-error classification is factored out of
`wire_json` / `parse_wire_json` so the wrapper keeps reporting malformed
text as `InvalidWireJson` and schema violations as `InvalidArgument`.

The C ABI is unaffected: it takes the whole `RerankingCallOptions` /
`SearchCallOptions` as one JSON blob, so its required fields are always
present. That is also why this is fixed at the binding sites rather than
by making the core fields `#[serde(default)]`.
---
 bindings/node/src/error.rs        |  10 ++-
 bindings/node/src/multimodal.rs   | 115 +++++++++++++++++++++++++++--
 bindings/python/src/error.rs      |  10 ++-
 bindings/python/src/multimodal.rs | 116 ++++++++++++++++++++++++++++--
 4 files changed, 239 insertions(+), 12 deletions(-)

diff --git a/bindings/node/src/error.rs b/bindings/node/src/error.rs
index da924f90..0b380639 100644
--- a/bindings/node/src/error.rs
+++ b/bindings/node/src/error.rs
@@ -142,7 +142,13 @@ pub(crate) fn parse_wire_json(
     argument: &'static str,
     json: &str,
 ) -> MResult {
-    serde_json::from_str(json).map_err(|e| match e.classify() {
+    serde_json::from_str(json).map_err(|e| wire_error(argument, &e))
+}
+
+/// The classification above, for callers that produce the `serde_json::Error`
+/// themselves rather than through [`parse_wire_json`].
+pub(crate) fn wire_error(argument: &'static str, e: &serde_json::Error) -> AiMuxBindingError {
+    match e.classify() {
         serde_json::error::Category::Data => AiMuxBindingError::from(AiMuxError::InvalidArgument(
             format!("invalid {argument}: {e}"),
         )),
@@ -151,7 +157,7 @@ pub(crate) fn parse_wire_json(
             message: e.to_string(),
         }
         .into(),
-    })
+    }
 }
 
 /// Serialize a result for the JS side; failure is a binding `ResultSerialization`.
diff --git a/bindings/node/src/multimodal.rs b/bindings/node/src/multimodal.rs
index 6428e1b2..1d609570 100644
--- a/bindings/node/src/multimodal.rs
+++ b/bindings/node/src/multimodal.rs
@@ -20,6 +20,53 @@ use aimux_core::transcription_model::{
 use aimux_core::video_model::{VideoCallOptions, VideoModel as VideoModelTrait};
 use napi_derive::napi;
 
+/// Fill in the fields the method's explicit arguments own, then deserialize.
+///
+/// Those fields are required on the Rust options struct, so a caller's
+/// options object on its own fails with `missing field` before the explicit
+/// arguments can be put back. The values inserted here are placeholders that
+/// satisfy the shape only; the caller overwrites each one from its own
+/// argument right after this returns. Pure JSON — no binding error types — so
+/// the parse contract is unit-testable without a Node runtime.
+fn opts_with_args(
+    json: &str,
+    owned_by_args: &[(&str, serde_json::Value)],
+) -> Result {
+    let mut value: serde_json::Value = serde_json::from_str(json)?;
+    // A non-object body needs no filling; `from_value` reports it as the same
+    // invalid-type data error it always did.
+    if let Some(object) = value.as_object_mut() {
+        for (field, placeholder) in owned_by_args {
+            object.insert((*field).to_string(), placeholder.clone());
+        }
+    }
+    serde_json::from_value(value)
+}
+
+/// [`opts_with_args`] under `parse_wire_json`'s error classification.
+fn parse_opts_json(
+    argument: &'static str,
+    json: &str,
+    owned_by_args: &[(&str, serde_json::Value)],
+) -> crate::error::MResult {
+    opts_with_args(json, owned_by_args).map_err(|e| crate::error::wire_error(argument, &e))
+}
+
+/// Shape-only stand-in for `TranscriptionCallOptions::audio`, built from the
+/// real enum so a variant rename is a compile error rather than a runtime one.
+fn audio_placeholder() -> serde_json::Value {
+    serde_json::to_value(AudioInput::Base64(String::new()))
+        .expect("AudioInput always serializes")
+}
+
+/// Shape-only stand-in for `RerankingCallOptions::documents`.
+fn documents_placeholder() -> serde_json::Value {
+    serde_json::to_value(aimux_core::reranking_model::RerankingDocuments::Text {
+        values: Vec::new(),
+    })
+    .expect("RerankingDocuments always serializes")
+}
+
 // ─────────────────────────────────────────────────────────────────────────────
 // EmbeddingModel
 // ─────────────────────────────────────────────────────────────────────────────
@@ -163,8 +210,14 @@ impl TranscriptionModel {
                     TranscriptionCallOptions::new(AudioInput::Base64(audio_base64), media_type);
                 if let Some(s) = opts_json.as_deref() {
                     if !s.trim().is_empty() && s.trim() != "null" {
-                        let mut parsed: TranscriptionCallOptions =
-                            parse_wire_json("opts_json", s)?;
+                        let mut parsed: TranscriptionCallOptions = parse_opts_json(
+                            "opts_json",
+                            s,
+                            &[
+                                ("audio", audio_placeholder()),
+                                ("media_type", serde_json::Value::from("")),
+                            ],
+                        )?;
                         // Required data comes from the explicit arguments; all
                         // operation policy and optional fields come from JSON.
                         parsed.audio = opts.audio;
@@ -217,7 +270,14 @@ impl RerankingModel {
                 let mut opts = RerankingCallOptions::new(query, docs);
                 if let Some(s) = opts_json.as_deref() {
                     if !s.trim().is_empty() && s.trim() != "null" {
-                        let mut parsed: RerankingCallOptions = parse_wire_json("opts_json", s)?;
+                        let mut parsed: RerankingCallOptions = parse_opts_json(
+                            "opts_json",
+                            s,
+                            &[
+                                ("query", serde_json::Value::from("")),
+                                ("documents", documents_placeholder()),
+                            ],
+                        )?;
                         // Required data comes from the explicit arguments; all
                         // operation policy and optional fields come from JSON.
                         parsed.query = opts.query;
@@ -303,7 +363,11 @@ impl SearchModel {
                     if !s.trim().is_empty() && s.trim() != "null" {
                         // Required data comes from the explicit arguments; all
                         // operation policy and optional fields come from JSON.
-                        let mut parsed: SearchCallOptions = parse_wire_json("opts_json", s)?;
+                        let mut parsed: SearchCallOptions = parse_opts_json(
+                            "opts_json",
+                            s,
+                            &[("query", serde_json::Value::from(""))],
+                        )?;
                         parsed.query = opts.query;
                         opts = parsed;
                     }
@@ -921,3 +985,46 @@ impl TranscriptionSession {
         self.token.abort();
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use aimux_core::search_model::SearchCallOptions;
+
+    /// A caller that passes any options at all must not trip `missing field`
+    /// on the data its explicit arguments already carry — that is what made
+    /// `maxRetries` / `timeout` unreachable from Node for these modalities.
+    #[test]
+    fn options_parse_without_the_fields_the_explicit_args_own() {
+        let transcription: TranscriptionCallOptions = opts_with_args(
+            r#"{"max_retries":3,"timeout":{"total_ms":1000}}"#,
+            &[
+                ("audio", audio_placeholder()),
+                ("media_type", serde_json::Value::from("")),
+            ],
+        )
+        .expect("transcription options with no audio/media_type");
+        assert_eq!(transcription.max_retries, Some(3));
+
+        let rerank: RerankingCallOptions = opts_with_args(
+            r#"{"top_n":2,"max_retries":1}"#,
+            &[
+                ("query", serde_json::Value::from("")),
+                ("documents", documents_placeholder()),
+            ],
+        )
+        .expect("rerank options with no query/documents");
+        assert_eq!(rerank.top_n, Some(2));
+
+        // The placeholder is overwritten by the caller, so it must lose to
+        // whatever the explicit argument holds — never to the JSON.
+        let mut search: SearchCallOptions = opts_with_args(
+            r#"{"query":"dogs","max_results":5}"#,
+            &[("query", serde_json::Value::from(""))],
+        )
+        .expect("search options with no query");
+        assert_eq!(search.max_results, Some(5));
+        search.query = "cats".to_string();
+        assert_eq!(search.query, "cats");
+    }
+}
diff --git a/bindings/python/src/error.rs b/bindings/python/src/error.rs
index e2ae7ac2..96768bc5 100644
--- a/bindings/python/src/error.rs
+++ b/bindings/python/src/error.rs
@@ -118,7 +118,13 @@ pub(crate) fn wire_json(
     argument: &'static str,
     s: &str,
 ) -> PyResult {
-    serde_json::from_str(s).map_err(|e| match e.classify() {
+    serde_json::from_str(s).map_err(|e| wire_error(argument, &e))
+}
+
+/// The classification above, for callers that produce the `serde_json::Error`
+/// themselves rather than through [`wire_json`].
+pub(crate) fn wire_error(argument: &'static str, e: &serde_json::Error) -> pyo3::PyErr {
+    match e.classify() {
         serde_json::error::Category::Data => to_py_err(&AiMuxError::InvalidArgument(format!(
             "invalid {argument}: {e}"
         ))),
@@ -126,7 +132,7 @@ pub(crate) fn wire_json(
             argument,
             message: e.to_string(),
         }),
-    })
+    }
 }
 
 /// Serialize a result for the wire; failure is the binding's, not an AiMuxError.
diff --git a/bindings/python/src/multimodal.rs b/bindings/python/src/multimodal.rs
index 47bf948c..3c728961 100644
--- a/bindings/python/src/multimodal.rs
+++ b/bindings/python/src/multimodal.rs
@@ -14,7 +14,8 @@
 use std::sync::Arc;
 
 use crate::error::{
-    BindingError, AiMuxBindingError, binding_py_err, serialize_result, to_py_err, wire_json,
+    BindingError, AiMuxBindingError, binding_py_err, serialize_result, to_py_err, wire_error,
+    wire_json,
 };
 use aimux_core::AiMuxError;
 use aimux_core::embedding_model::{EmbeddingCallOptions, EmbeddingModel as EmbeddingModelTrait};
@@ -30,6 +31,53 @@ use aimux_core::transcription_model::{
 use aimux_core::video_model::{VideoCallOptions, VideoModel as VideoModelTrait};
 use pyo3::prelude::*;
 
+/// Fill in the fields the method's explicit arguments own, then deserialize.
+///
+/// Those fields are required on the Rust options struct, so a caller's
+/// options object on its own fails with `missing field` before the explicit
+/// arguments can be put back. The values inserted here are placeholders that
+/// satisfy the shape only; the caller overwrites each one from its own
+/// argument right after this returns. Pure JSON — no binding error types — so
+/// the parse contract is unit-testable without a Python interpreter.
+fn opts_with_args(
+    json: &str,
+    owned_by_args: &[(&str, serde_json::Value)],
+) -> Result {
+    let mut value: serde_json::Value = serde_json::from_str(json)?;
+    // A non-object body needs no filling; `from_value` reports it as the same
+    // invalid-type data error it always did.
+    if let Some(object) = value.as_object_mut() {
+        for (field, placeholder) in owned_by_args {
+            object.insert((*field).to_string(), placeholder.clone());
+        }
+    }
+    serde_json::from_value(value)
+}
+
+/// [`opts_with_args`] under `wire_json`'s error classification.
+fn parse_opts_json(
+    argument: &'static str,
+    json: &str,
+    owned_by_args: &[(&str, serde_json::Value)],
+) -> PyResult {
+    opts_with_args(json, owned_by_args).map_err(|e| wire_error(argument, &e))
+}
+
+/// Shape-only stand-in for `TranscriptionCallOptions::audio`, built from the
+/// real enum so a variant rename is a compile error rather than a runtime one.
+fn audio_placeholder() -> serde_json::Value {
+    serde_json::to_value(AudioInput::Base64(String::new()))
+        .expect("AudioInput always serializes")
+}
+
+/// Shape-only stand-in for `RerankingCallOptions::documents`.
+fn documents_placeholder() -> serde_json::Value {
+    serde_json::to_value(aimux_core::reranking_model::RerankingDocuments::Text {
+        values: Vec::new(),
+    })
+    .expect("RerankingDocuments always serializes")
+}
+
 // ─────────────────────────────────────────────────────────────────────────────
 // EmbeddingModel
 // ─────────────────────────────────────────────────────────────────────────────
@@ -148,7 +196,14 @@ impl TranscriptionModel {
             if !s.trim().is_empty() && s.trim() != "null" {
                 // Take every caller option; audio and media_type always come
                 // from the explicit args.
-                let mut parsed: TranscriptionCallOptions = wire_json("opts_json", s)?;
+                let mut parsed: TranscriptionCallOptions = parse_opts_json(
+                    "opts_json",
+                    s,
+                    &[
+                        ("audio", audio_placeholder()),
+                        ("media_type", serde_json::Value::from("")),
+                    ],
+                )?;
                 parsed.audio = opts.audio;
                 parsed.media_type = opts.media_type;
                 opts = parsed;
@@ -195,7 +250,14 @@ impl RerankingModel {
             if !s.trim().is_empty() && s.trim() != "null" {
                 // Take every caller option; query and documents always come
                 // from the explicit args.
-                let mut parsed: RerankingCallOptions = wire_json("opts_json", s)?;
+                let mut parsed: RerankingCallOptions = parse_opts_json(
+                    "opts_json",
+                    s,
+                    &[
+                        ("query", serde_json::Value::from("")),
+                        ("documents", documents_placeholder()),
+                    ],
+                )?;
                 parsed.query = opts.query;
                 parsed.documents = opts.documents;
                 opts = parsed;
@@ -260,7 +322,11 @@ impl SearchModel {
             if !s.trim().is_empty() && s.trim() != "null" {
                 // Take every caller option; the query always comes from the
                 // explicit arg.
-                let mut parsed: SearchCallOptions = wire_json("opts_json", s)?;
+                let mut parsed: SearchCallOptions = parse_opts_json(
+                    "opts_json",
+                    s,
+                    &[("query", serde_json::Value::from(""))],
+                )?;
                 parsed.query = opts.query;
                 opts = parsed;
             }
@@ -794,3 +860,45 @@ impl TranscriptionSession {
         self.token.abort();
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A caller that passes any options at all must not trip `missing field`
+    /// on the data its explicit arguments already carry — that is what made
+    /// `max_retries` / `timeout` unreachable from Python for these modalities.
+    #[test]
+    fn options_parse_without_the_fields_the_explicit_args_own() {
+        let transcription: TranscriptionCallOptions = opts_with_args(
+            r#"{"max_retries":3,"timeout":{"total_ms":1000}}"#,
+            &[
+                ("audio", audio_placeholder()),
+                ("media_type", serde_json::Value::from("")),
+            ],
+        )
+        .expect("transcription options with no audio/media_type");
+        assert_eq!(transcription.max_retries, Some(3));
+
+        let rerank: RerankingCallOptions = opts_with_args(
+            r#"{"top_n":2,"max_retries":1}"#,
+            &[
+                ("query", serde_json::Value::from("")),
+                ("documents", documents_placeholder()),
+            ],
+        )
+        .expect("rerank options with no query/documents");
+        assert_eq!(rerank.top_n, Some(2));
+
+        // The placeholder is overwritten by the caller, so it must lose to
+        // whatever the explicit argument holds — never to the JSON.
+        let mut search: SearchCallOptions = opts_with_args(
+            r#"{"query":"dogs","max_results":5}"#,
+            &[("query", serde_json::Value::from(""))],
+        )
+        .expect("search options with no query");
+        assert_eq!(search.max_results, Some(5));
+        search.query = "cats".to_string();
+        assert_eq!(search.query, "cats");
+    }
+}

From 2d850db05639ce2721ed3e30de3dd08d1d15c51c Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Fri, 4 Sep 2026 00:40:22 +0800
Subject: [PATCH 23/24] docs: record that stream_text now requires a tokio
 runtime
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The pump task calls `tokio::spawn` unconditionally, where before a
runtime was only needed once a deadline was armed. Adds the missing §3
deviation row.
---
 docs/ai-sdk-request-pipeline.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index 4ec6f354..f6a34fd6 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -93,6 +93,7 @@ Core user operation
 | timeout 输入形状 | 保留跨语言已有的 object/struct 形式，不增加 AI SDK 的裸 number 简写；字段语义一致 |
 | `AbortSignal` | 保留既有的 `CancellationToken` 薄包装，只表示调用方取消；Rust 可直接 drop future，因此 timeout 由 Core deadline/select 表达，不照搬 JS 的 signal merge（§8.0） |
 | stream 驱动 | AI SDK 的 `streamText` 返回前即开始消费 provider stream（push 语义，timer 观察到的是 chunk 到达时刻）；Rust stream 是 pull 语义，只在 `poll_next` 里观察 deadline 会把消费方不 poll 的时间也算进 `first_chunk_ms`/`chunk_ms`。Aimux 在 `stream_text` 返回前 spawn 一个 pump task 驱动 provider stream 进 unbounded channel，deadline 在 pump 侧按到达时刻 arm/reset；返回的 stream drop 时 abort pump。缓冲无上界与 AI SDK 一致（受响应本身约束），pump 在 terminal item 处停止 |
+| runtime 依赖 | AI SDK 没有这层概念；Aimux 的 `stream_text` 因为 pump task 需要 `tokio::spawn`，现在无条件要求运行在 tokio runtime 上，此前只有 armed 的 deadline 才需要（`stream_text` 之外的 Core operation 不变） |
 | SSE framing | AI SDK 的 event-source handler 直接产出解析后的事件；Aimux 复用 `aimux_stream::SseStream` 做同一件事，framing 不留在 Provider |
 | tool timeout | aimux 目前不执行用户 tool，本 RFC 不增加 `tool_ms`/per-tool timeout |
 | observability | 每次 exchange 继续走 Aimux recording/tracing；这不改变 helper 的单次请求语义 |

From ebf0a22ed673315b4b6e8aebc603baa45d4563ba Mon Sep 17 00:00:00 2001
From: Card Cunningham 
Date: Mon, 7 Sep 2026 17:02:24 +0800
Subject: [PATCH 24/24] fix(core): scope caller video idempotency keys to each
 batch

---
 aimux-core/src/video_model.rs   | 45 +++++++++++++++++++++++++++++++++
 docs/ai-sdk-request-pipeline.md |  4 ++-
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/aimux-core/src/video_model.rs b/aimux-core/src/video_model.rs
index ecc3dbbf..8833c4ee 100644
--- a/aimux-core/src/video_model.rs
+++ b/aimux-core/src/video_model.rs
@@ -404,6 +404,17 @@ async fn start_and_poll(
     let last = options.n.div_ceil(max_per_call) - 1;
     let batches = (0..=last).map(|i| {
         let mut batch_options = options.clone();
+        if last > 0
+            && let Some(headers) = batch_options.headers.as_mut()
+        {
+            // Distinct batches must not deduplicate each other. Deriving the
+            // suffix from the batch index also keeps caller-key replays stable.
+            for (name, key) in headers.iter_mut() {
+                if name.eq_ignore_ascii_case("idempotency-key") {
+                    *key = format!("{key}:batch:{i}");
+                }
+            }
+        }
         batch_options.n = if i < last {
             max_per_call
         } else {
@@ -911,4 +922,38 @@ mod tests {
         let unique: std::collections::HashSet<_> = keys.iter().collect();
         assert_eq!(unique.len(), 3, "one key per batch: {keys:?}");
     }
+
+    #[tokio::test]
+    async fn caller_key_is_distinct_per_batch_and_stable_across_retries_and_replays() {
+        let model = ScriptedVideoModel::new(vec![]).with_start_failures(1);
+        let mut options = fast_poll_options();
+        options.n = 3;
+        options
+            .headers
+            .get_or_insert_with(SharedHeaders::new)
+            .insert("Idempotency-Key".into(), "caller-request-1".into());
+        generate_video(&model, options.clone()).await.unwrap();
+        let unique = {
+            let keys = model.start_idempotency_keys.lock().unwrap();
+            let unique: std::collections::HashSet<_> = keys.iter().cloned().collect();
+            assert_eq!(unique.len(), 3, "distinct batches share keys: {keys:?}");
+            assert_eq!(
+                keys.len(),
+                4,
+                "the failed batch must reuse its key on retry"
+            );
+            unique
+        };
+
+        let replay = ScriptedVideoModel::new(vec![]);
+        generate_video(&replay, options).await.unwrap();
+        let replay_keys = replay.start_idempotency_keys.lock().unwrap();
+        assert_eq!(
+            replay_keys
+                .iter()
+                .cloned()
+                .collect::>(),
+            unique
+        );
+    }
 }
diff --git a/docs/ai-sdk-request-pipeline.md b/docs/ai-sdk-request-pipeline.md
index f6a34fd6..3220e154 100644
--- a/docs/ai-sdk-request-pipeline.md
+++ b/docs/ai-sdk-request-pipeline.md
@@ -775,7 +775,9 @@ start/status 流），Core 拥有 poll 循环并分别 retry 两个阶段：
   切 batch，每个 batch 并发跑一次独立的 `do_start`/poll 流程并独立铸造
   idempotency key（AI SDK 用 `Promise.all`，Rust 侧 `try_join_all` 会在首个
   失败时 drop 其余 batch future——不影响正确性，两边都不会重连或取消已发出的
-  provider 任务）。`n == 0` 在任何网络调用之前返回 `InvalidArgument`。
+  provider 任务）。调用方提供 key 时，单批保持原键，多批按批次索引派生不同的
+  键；同批重试和同一调用方 key 的重放保持键稳定。`n == 0` 在任何网络调用
+  之前返回 `InvalidArgument`。
 
 ### 9.2 Router / MoA