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/aimux-core/Cargo.toml b/aimux-core/Cargo.toml index afda2ec1..0ad5927d 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", "rt", "sync"] } 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..d8420941 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,48 @@ 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), + } + }) + }) + }) + }) + // 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. /// - /// 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 +346,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 +388,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 +415,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..0b545711 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,127 @@ 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(); + // 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. + // + // 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; + 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; + // 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) => { + let _ = tx.send(Err(operation_deadline + .expect("operation deadline future only resolves for a deadline") + .error())); + break; + } + () = timeout::wait_for_deadline(chunk_deadline) => { + let _ = tx.send(Err(chunk_deadline + .expect("chunk deadline future only resolves for a deadline") + .error())); + break; + } + item = stream.next() => item, + }; + + let Some(item) = next else { break }; + 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)) + .transpose() + { + Ok(deadline) => deadline, + 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; + } + } + }) + }; let stream = crate::recording::RecordingOutcomeStream::new( stream, recorder.clone(), @@ -899,6 +1122,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. @@ -1036,3 +1277,456 @@ 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 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), + 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(20)).await; + + assert!(matches!( + result.stream.next().await.unwrap().unwrap(), + StreamPart::TextDelta { delta, .. } if delta == "hello" + )); + assert!(result.stream.next().await.is_none()); + } + + #[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..b3f9d5de --- /dev/null +++ b/aimux-core/src/timeout.rs @@ -0,0 +1,265 @@ +//! 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 — 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, + 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/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..8833c4ee 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,662 @@ 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, batching, and +/// timeout. +/// +/// 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 `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, +) -> 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 { + 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, + model.retry_config(), + 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); + } + } + + // 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 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 { + 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()) + }); + + // `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 { + merged.videos.extend(result.videos); + merged.warnings.extend(result.warnings); + merged.provider_metadata = + merge_provider_metadata(merged.provider_metadata, result.provider_metadata); + } + Ok(merged) +} + +/// 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 + // 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 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; + 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 `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 }; + 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) +} + +#[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, + /// 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, + start_idempotency_keys: Mutex>>, + status_idempotency_keys: Mutex>>, + } + + impl ScriptedVideoModel { + 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, + 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 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() + .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:?}"); + } + + #[test] + fn merge_provider_metadata_unions_same_provider_key_across_phases() { + 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 }) + ); + } + + #[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() + ); + } + + #[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" + ); + } + + #[tokio::test] + 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!( + 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" + ); + let keys = model + .start_idempotency_keys + .lock() + .expect("start key capture lock should not be poisoned"); + 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/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..8121bf7d --- /dev/null +++ b/aimux-core/tests/stream_timeout_matrix_test.rs @@ -0,0 +1,179 @@ +//! 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 +//! 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 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 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(); + 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()); +} 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/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 37de35a3..ac419d78 100644 --- a/aimux-provider-utils/src/http.rs +++ b/aimux-provider-utils/src/http.rs @@ -1,280 +1,127 @@ -//! 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(), - RedirectMode::Automatic, - ) - }) - .as_ref() - .map_err(|e| client_init_error(e)) -} - -/// 校验下载共享 Client（带 30s 整体超时，禁自动重定向——逐跳手动校验）。 -static DOWNLOAD: OnceLock> = OnceLock::new(); - -/// Shared client for validated downloads whose hop is not pinned (trusted -/// origin, or a proxy owns resolution). Redirects are followed manually by -/// [`send_validated_redirects`] so each hop is re-validated. -fn download_client() -> Result<&'static Client, AiMuxError> { - DOWNLOAD - .get_or_init(|| { - build_client( - PoolConfig::default(), - TimeoutConfig::default(), - global_proxy(), - RedirectMode::Manual, - ) - }) - .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(), - RedirectMode::Automatic, - ) - }) - .as_ref() - .map_err(|e| client_init_error(e)) +/// 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()) } -/// 用给定配置构建一个 reqwest Client。构建失败返回错误字符串（reqwest 仅 -/// 提供 Display），由调用方决定映射——不再 `expect`（issue #115：受限环境 -/// 下 TLS/资源初始化失败不应在 panic=abort 产物中直接终止宿主进程）。 -/// One redirect ceiling for both modes: reqwest's automatic following and -/// the manual validated-download loop. +/// Redirect ceiling for the manual validated-download loop. const MAX_REDIRECTS: usize = 10; -/// How a client treats redirects. `Manual` is for the validated-download -/// path, which follows redirects itself so every hop can be re-validated -/// and re-pinned. -#[derive(Clone, Copy)] -enum RedirectMode { - Automatic, - Manual, -} - -impl RedirectMode { - fn policy(self) -> reqwest::redirect::Policy { - match self { - Self::Automatic => reqwest::redirect::Policy::limited(MAX_REDIRECTS), - Self::Manual => reqwest::redirect::Policy::none(), - } - } -} - -fn build_client( - pool: PoolConfig, - timeout: TimeoutConfig, - proxy: ProxyConfig, - redirects: RedirectMode, -) -> Result { - apply_proxy(client_builder(&pool, &timeout, redirects), &proxy) +/// Non-redirecting client for a validated hop on the trusted origin. +/// Redirects are followed manually so every hop is re-validated; built per +/// exchange (downloads are rare and each hop is one request, so there is no +/// pool to share — and no stale-runtime pool to reuse). +fn download_client() -> Result { + let builder = Client::builder() + .connect_timeout(Duration::from_millis(10_000)) + .redirect(reqwest::redirect::Policy::none()); + apply_proxy(builder, &global_proxy()) .build() - .map_err(|e| e.to_string()) -} - -fn client_builder( - pool: &PoolConfig, - timeout: &TimeoutConfig, - redirects: RedirectMode, -) -> reqwest::ClientBuilder { - let mut b = Client::builder() - .redirect(redirects.policy()) - .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 + .map_err(|e| AiMuxError::Other(format!("download client initialization failed: {e}"))) } -/// Build a one-off client whose resolver only ever answers with the -/// pre-validated addresses. Used for downloads whose DNS results passed the -/// SSRF guard; prevents rebinding between validation and connection. -/// -/// `reqwest` matches `resolve` entries by exact host, so the URL's host (and -/// port, which `reqwest` reuses for the socket) must be supplied explicitly; -/// port 0 would make it dial port 0. +/// Client for one validated hop off the trusted origin: the connection is +/// pinned to exactly the DNS answers that passed the guard (resolve +/// overrides), defeating TTL-0 rebinding. The proxy configuration is applied +/// so reqwest makes the per-URL routing decision itself: when a proxy +/// carries the request the proxy resolves the target (a trusted transport, +/// the override below is unused), and any request the proxy rules send +/// DIRECT still connects only through the validated, pinned addresses. fn pinned_client(url: &str, addresses: &[std::net::IpAddr]) -> Result { if addresses.is_empty() { return Err(AiMuxError::Other( @@ -288,1584 +135,848 @@ fn pinned_client(url: &str, addresses: &[std::net::IpAddr]) -> Result = addresses .iter() .map(|address| std::net::SocketAddr::new(*address, port)) .collect(); - b = b.resolve_to_addrs(&host, &socket_addresses); - b.build().map_err(|e| { - AiMuxError::Other(format!("pinned download client initialization failed: {e}")) - }) + let builder = Client::builder() + .connect_timeout(Duration::from_millis(10_000)) + .redirect(reqwest::redirect::Policy::none()); + apply_proxy(builder, &global_proxy()) + .resolve_to_addrs(&host, &socket_addresses) + .build() + .map_err(|e| { + AiMuxError::Other(format!("pinned download client initialization failed: {e}")) + }) } -/// 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). +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 不参与重建。 -#[derive(Debug, Clone)] +/// 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, Default)] 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, + /// Per-exchange whole-response hang guard. `None` uses the 30s default; + /// a provider whose endpoint legitimately holds the connection longer + /// (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 + /// validated and pinned, redirects followed manually with each hop + /// re-validated. + pub validate_url: bool, + /// AI SDK `trustedOrigin`: a developer-configured origin (normally the + /// provider's `base_url`) whose same-origin URLs are exempt from the + /// address blocklist, so self-hosted deployments serving assets from + /// their own (possibly private) origin keep working. Reachability only — + /// never about headers, and never derived from response data. + pub trusted_origin: Option, + /// AI SDK `credentialedOrigin`: caller headers (which may carry the + /// provider API key) are sent only while the target is same-origin with + /// this value — from the first request and on every redirect hop. `None` + /// means the caller gates its own headers (e.g. BFL's host allowlist). + 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 + } } -/// HTTP 响应（非流式，纯数据）。 -/// -/// `body` 为完整响应体字节；provider 用 `serde_json::from_slice` 等解析。 -#[derive(Debug)] -pub struct HttpResponse { - pub status: u16, - pub headers: HashMap, - pub body: Bytes, +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() + } } -/// HTTP 流式响应。 -/// -/// `body` 为字节流（`BoxStream>`，不依赖 reqwest）。 -/// 传给 `SseStream::new` 即可解析 SSE——`SseStream` 是泛型的，能直接接收。 -pub struct HttpStreamResponse { - pub status: u16, - pub headers: HashMap, - pub body: BoxStream<'static, Result>, +macro_rules! exchange_context_abort_only { + ($($ty:path),* $(,)?) => { + $(impl ExchangeContext for $ty { + fn abort_signal(&self) -> Option { + self.abort_signal.clone() + } + })* + }; } -// ════════════════════════════════════════════════════════════════════════════ -// 执行（含 retry，内部逻辑） -// ════════════════════════════════════════════════════════════════════════════ +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() + } + } +} -/// 发送非流式请求，带 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()?; - send_via( - RequestTransport::Direct(client), - request, - retry_config, - error_structure, - ) - .await +/// SSRF guard configuration carried by a validated exchange +/// (`HttpRequest.validate_url`). +#[derive(Debug, Clone)] +pub(crate) struct DownloadValidation { + pub(crate) trusted_origin: Option, + pub(crate) credentialed_origin: Option, } -/// Fetch a **provider-supplied URL** (a generated asset, polling URL, result -/// URL, or upload URL taken from a response body or header) with SSRF -/// protection: the URL is validated against AI SDK's address blocklists, -/// every DNS answer is checked and the connection pinned to the validated -/// addresses (defeating TTL-0 rebinding), every redirect hop is re-validated, -/// and headers are sanitized (hop-by-hop, forwarding, and metadata-service -/// headers dropped). -/// -/// Both origins mirror AI SDK's options and must come from developer -/// configuration or a provider's own allowlist — never from a response: -/// -/// - `trusted_origin` (normally the configured `base_url`) exempts -/// same-origin URLs from the address blocklist so self-hosted deployments -/// keep working. It is only about reachability, never about headers. -/// - `credentialed_origin` — AI SDK's `credentialedOrigin` — confines caller -/// headers (which may carry the provider API key) to that origin, from the -/// first request and on every redirect hop; once stripped they are never -/// restored. `None` means the caller gates its own headers (e.g. BFL's -/// host allowlist); headers then still strip on any redirect leaving the -/// request's origin. -/// -/// # Errors -/// -/// Returns [`AiMuxError::InvalidArgument`] when the URL or a DNS answer -/// fails validation, plus the same transport/response errors as [`send`]. -pub async fn send_validated( - request: HttpRequest, - trusted_origin: Option<&str>, - credentialed_origin: Option<&str>, - retry_config: RetryConfig, - error_structure: &ErrorStructure, -) -> Result { - auto_init_from_env(); - send_via( - RequestTransport::Validated { - trusted_origin, - credentialed_origin, - }, - request, - retry_config, - error_structure, - ) - .await +#[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, + pub(crate) response_timeout: Option, + pub(crate) max_json_response_bytes: Option, + pub(crate) validation: Option, +} + +impl HttpRequest { + pub(crate) fn prepare(self, method: HttpMethod, body: HttpBody) -> PreparedRequest { + // AI SDK gates credentialed headers independently of validateUrl: + // when the target is off the credentialed origin they never leave, + // whichever transport the exchange takes. + let mut headers = self.headers; + if let Some(origin) = self.credentialed_origin.as_deref() + && !crate::download_guard::same_origin(&self.url, origin) + { + crate::download_guard::retain_user_agent(&mut headers); + } + PreparedRequest { + method, + body, + url: self.url, + headers, + abort_signal: self.abort_signal, + 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, + }), + } + } } -async fn send_via( - transport: RequestTransport<'_>, - request: HttpRequest, - retry_config: RetryConfig, - error_structure: &ErrorStructure, -) -> Result { - let started = Instant::now(); - let (resp, attempt) = - send_with_retry_raw(&transport, &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() - }) - })?, - }; - - if body_logging_enabled() { - tracing::trace!( - target: "aimux_provider_utils::http", - body = %redact_body(&String::from_utf8_lossy(&body)), - "response_body" - ); + record_transport_closed(self.request); + self.armed = false; + } } - // RFC-0023: 成功 attempt 的 exchange(body 补全;attempt 来自重试循环)。 - // 门控提前:录制上下文不存在时零成本(S5 性能——不构造 exchange)。 - if request.recording_context.is_none() { - return Ok(HttpResponse { - status, - headers, - body, - }); + fn disarm(&mut self) { + self.armed = false; } - 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, - }, - ); - // 非流式:exchange 即终点 → 传输封闭(与层 A 幂等)。 - record_transport_closed(&request); +} - Ok(HttpResponse { - status, - headers, - body, - }) +impl Drop for ExchangeGuard<'_> { + fn drop(&mut self) { + self.fail("exchange future cancelled"); + } } -/// 发送流式请求，返回字节流。内部用 `shared_streaming_client()`（无整体超时）。 -/// -/// retry 仅覆盖**建连阶段**——`.send()` 返回 200 后立即返回字节流，流中途 -/// 出错不重试（已吐 token 后重试会重复内容，RFC-0009 §5）。 -/// -/// # 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()?; +/// 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 (resp, attempt) = send_with_retry_raw( - &RequestTransport::Direct(client), - &request, - retry_config, - error_structure, - ) - .await?; - - let status = resp.status().as_u16(); - let headers = collect_headers(resp.headers()); + 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); tracing::debug!( target: "aimux_provider_utils::http", - status = status, - "stream_connected" + 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" ); - - // 流式录制上下文:仅调用点带 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 等流终结补全)。 - record_exchange( - &request, - aimux_core::recording::HttpExchange { - 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, - }, - error: None, - finalized: false, - }, + 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" ); - Some(StreamRec { - request: request.clone(), - attempt, - status, - headers: resp_headers, - ttfb_ms: None, - body: ByteAccumulator::new(RECORD_BODY_CAP), - chunks: 0usize, - }) + } + + let sent = if request.validation.is_some() { + send_validated_redirects(request).await } else { - None + match shared_client() { + Ok(client) => send_one_request(&client, request).await, + Err(error) => Err(error), + } + }; + let response = match sent { + 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" + ); + Ok(observe_response( + response, + request, + started, + latency_ms, + attempt, + exchange_index, + )) +} - // 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()) - }); - 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() - }) - }) +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()) + })) }) - .boxed(), - started: false, - start: started, - chunks: 0u64, - 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, - chunks: usize, +fn is_redirect_status(status: reqwest::StatusCode) -> bool { + matches!( + status, + reqwest::StatusCode::MOVED_PERMANENTLY + | reqwest::StatusCode::FOUND + | reqwest::StatusCode::SEE_OTHER + | reqwest::StatusCode::TEMPORARY_REDIRECT + | reqwest::StatusCode::PERMANENT_REDIRECT + ) } -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); - } +fn redirect_error( + request: &PreparedRequest, + message: impl Into, + status: reqwest::StatusCode, +) -> AiMuxError { + AiMuxError::ApiCall(Box::new(ApiCallError { + status_code: Some(status.as_u16()), + ..api_call_error(request, message) + })) } -/// 流式响应体观测包装：发出「首字节」与「流结束」两个 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, -} - -impl Stream for ObservedByteStream { - 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); +/// Send one validated-download exchange, following redirects manually so +/// every hop is re-validated and its connection pinned to the DNS answers +/// that passed the guard (see `download_guard`). The whole chain is one +/// exchange to the caller, matching how an auto-following client behaves. +async fn send_validated_redirects( + request: &PreparedRequest, +) -> Result { + let validation = request + .validation + .clone() + .expect("send_validated_redirects requires a validation config"); + let trusted_origin = validation.trusted_origin.as_deref(); + let credentialed_origin = validation.credentialed_origin.as_deref(); + let mut current = request.clone(); + // First-request credential gating already ran in `prepare`; this strips + // hop-by-hop, forwarding, and metadata-service headers for every hop. + crate::download_guard::sanitize_download_headers(&mut current.headers); + // Headers never travel past this origin on a redirect; stripping is + // one-way, so a hop back onto it cannot restore them. + let credential_anchor = credentialed_origin.unwrap_or(&request.url); + let mut pinned = + crate::download_guard::validate_download_target(¤t.url, trusted_origin).await?; + for redirect_count in 0..=MAX_REDIRECTS { + // Empty pins mean the hop is on the trusted origin. + let client = if pinned.is_empty() { + download_client()? + } else { + pinned_client(¤t.url, &pinned)? + }; + let response = send_one_request(&client, ¤t).await?; + let status = response.status(); + if !is_redirect_status(status) { + return Ok(response); } - let result = this.inner.as_mut().poll_next(cx); - 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; - } - } - 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(None) => { - 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(None); - } - } - Poll::Pending => {} + let Some(location) = response.headers().get(reqwest::header::LOCATION) else { + return Ok(response); + }; + if redirect_count == MAX_REDIRECTS { + return Err(redirect_error(request, "too many redirects", status)); } - result - } -} - -impl Drop for ObservedByteStream { - fn drop(&mut self) { - // 消费方提前放弃流:仍补全已累积的 body,避免骨架悬空。 - if self.record.is_some() - && !self.done - && let Some(rec) = &mut self.record + let location = location + .to_str() + .map(str::to_owned) + .map_err(|_| redirect_error(request, "redirect location is not valid UTF-8", status))?; + // Dropping the unconsumed 3xx body releases its connection before a + // potentially slow DNS check for the next hop. + drop(response); + let base = url::Url::parse(¤t.url) + .map_err(|e| AiMuxError::InvalidArgument(format!("invalid request URL: {e}")))?; + let next = base + .join(&location) + .map_err(|e| redirect_error(request, format!("invalid redirect URL: {e}"), status))?; + // Fetch treats a redirect to a non-HTTP(S) scheme (data:, file:, ...) + // as a network error; following one would let the redirecting server + // fabricate a response outside the transport. + if !matches!(next.scheme(), "http" | "https") { + return Err(redirect_error( + request, + format!("redirect to non-HTTP scheme: {}", next.scheme()), + status, + )); + } + let next = next.to_string(); + let hop_trusted = crate::download_guard::hop_trusted_origin(trusted_origin, ¤t.url); + pinned = crate::download_guard::validate_download_target(&next, hop_trusted).await?; + // Credentials are scoped to the credential anchor, not the previous + // hop: once a redirect leaves it, headers cannot come back. + if !crate::download_guard::same_origin(&next, credential_anchor) { + crate::download_guard::retain_user_agent(&mut current.headers); + } + if status == reqwest::StatusCode::SEE_OTHER + || (matches!( + status, + reqwest::StatusCode::MOVED_PERMANENTLY | reqwest::StatusCode::FOUND + ) && matches!(current.method, HttpMethod::Post)) { - rec.finalize(Some("abandoned".to_string())); + current.method = HttpMethod::Get; + current.body = HttpBody::Empty; } + current.url = next; } + unreachable!("redirect loop returns on response or error") } -// ════════════════════════════════════════════════════════════════════════════ -// 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)?; - - 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, - } -} -/// 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?, +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), }; - // 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); + 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(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(), + 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, }) } -/// 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. +/// Abort-aware provider polling delay. This is not the operation retry delay; +/// Core owns exponential backoff. /// -/// 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). +/// # Errors /// -/// 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, +/// 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(()) + } + } } -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)) - }) +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), }; - 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)) - } + record_exchange( + request, + HttpExchange { + step: request + .recording_context + .as_ref() + .and_then(RecordingContext::step), + attempt, + exchange_index, + request: to_http_record(request), + response: Some(response), + timing: TimingRecord { + latency_ms, + ttfb_ms: Some(latency_ms), + }, + error: None, + finalized: false, + }, + ); + ResponseRecording { + context: context.clone(), + attempt, + exchange_index, + status: status.as_u16(), + headers: redacted_headers, + ttfb_ms: Some(latency_ms), + body: ByteAccumulator::new(RECORD_BODY_CAP), + chunks: 0, + finalized: false, } - } + }); + let observed = ObservedBody { + inner: response.bytes_stream().boxed(), + started, + chunks: 0, + done: false, + 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, + done: bool, + recording: Option, } -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(); +impl Stream for ObservedBody { + type Item = Result; - if this.done { + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.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; + let result = self.inner.as_mut().poll_next(context); + match &result { + Poll::Ready(Some(Ok(bytes))) => { + self.chunks += 1; + let chunks = self.chunks; + if let Some(recording) = &mut self.recording { + recording.body.push(bytes); + recording.chunks = chunks; } - 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(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; - this.sleep = None; - Poll::Ready(None) + } + Poll::Ready(None) => { + self.done = true; + tracing::debug!( + target: "aimux_provider_utils::http", + chunks = self.chunks, + duration_ms = self.started.elapsed().as_millis() as u64, + "response_body_end" + ); + if let Some(recording) = &mut self.recording { + recording.finalize(None); } - Poll::Pending => Poll::Pending, } + Poll::Pending => {} } + result } } -// ════════════════════════════════════════════════════════════════════════════ -// 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(()) +impl Drop for ObservedBody { + fn drop(&mut self) { + if !self.done + && let Some(recording) = &mut self.recording + { + recording.finalize(Some("response body abandoned".into())); } } } -/// 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; +struct ResponseRecording { + context: RecordingContext, + attempt: u32, + exchange_index: u32, + status: u16, + headers: Vec<(String, String)>, + ttfb_ms: Option, + body: ByteAccumulator, + chunks: usize, + finalized: bool, +} -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), - } +impl ResponseRecording { + fn finalize(&mut self, error: Option) { + if self.finalized { + return; } - None => Ok(read.await), + 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); } } -// ── 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] -} +const RECORD_BODY_CAP: usize = 1 << 20; -/// 有界字节累积器:流式录制按原始字节累积,流末统一解码。 -/// -/// 解决每个网络 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); +fn record_exchange(request: &PreparedRequest, exchange: HttpExchange) { + if let Some(context) = &request.recording_context { + context + .recorder + .record_exchange(&context.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); +fn record_transport_closed(request: &PreparedRequest) { + if let Some(context) = &request.recording_context { + context.recorder.record_transport_closed(&context.call_id); } } -/// 声明传输层封闭(层 B;P2 语义:不再有 exchange)。 -fn record_transport_closed(request: &HttpRequest) { - if let Some(ctx) = &request.recording_context { - ctx.recorder.record_transport_closed(&ctx.call_id); - } -} - -/// 收集响应头并脱敏(录制 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, - } -} - -/// 录制失败 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 是纯数据，可重复读）。 -/// How one attempt of the retry loop reaches the network: directly on a -/// client, or through the validated-download redirect loop. -enum RequestTransport<'a> { - Direct(&'a Client), - Validated { - trusted_origin: Option<&'a str>, - credentialed_origin: Option<&'a str>, - }, -} - -impl RequestTransport<'_> { - async fn send(&self, request: &HttpRequest) -> Result { - match self { - Self::Direct(client) => send_request(client, request).await, - Self::Validated { - trusted_origin, - credentialed_origin, - } => send_validated_redirects(request, *trusted_origin, *credentialed_origin).await, - } - } -} - -/// The redirect statuses fetch follows; 300/304 are responses, not hops. -fn is_redirect_status(status: reqwest::StatusCode) -> bool { - matches!( - status, - reqwest::StatusCode::MOVED_PERMANENTLY - | reqwest::StatusCode::FOUND - | reqwest::StatusCode::SEE_OTHER - | reqwest::StatusCode::TEMPORARY_REDIRECT - | reqwest::StatusCode::PERMANENT_REDIRECT - ) -} - -fn redirect_error(message: impl Into, status: reqwest::StatusCode) -> AiMuxError { - AiMuxError::ApiCall(ApiCallError { - message: message.into(), - status_code: Some(status.as_u16()), - is_retryable: false, - ..Default::default() - }) -} - -/// Send one attempt of a validated download, following redirects manually so -/// every hop is validated and its connection pinned to the DNS answers that -/// passed the guard. The whole chain counts as one attempt to the retry -/// layer, matching how the auto-following shared client behaves. -async fn send_validated_redirects( - request: &HttpRequest, - trusted_origin: Option<&str>, - credentialed_origin: Option<&str>, -) -> Result { - let mut current = request.clone(); - crate::download_guard::sanitize_download_headers(&mut current.headers); - // AI SDK's credentialedOrigin: when set, caller headers (which may carry - // the provider API key) are confined to that exact origin from the very - // first request. When unset the caller has already gated its own headers - // (e.g. BFL's host allowlist); redirects still strip below either way. - if let Some(origin) = credentialed_origin - && !crate::download_guard::same_origin(¤t.url, origin) - { - crate::download_guard::retain_user_agent(&mut current.headers); - } - // Headers never travel past this origin on a redirect; stripping is - // one-way, so a hop back onto it cannot restore them. - let credential_anchor = credentialed_origin.unwrap_or(&request.url); - let mut pinned = - crate::download_guard::validate_download_target(¤t.url, trusted_origin).await?; - for redirect_count in 0..=MAX_REDIRECTS { - let hop_client; - let client: &Client = if pinned.is_empty() { - // Empty pins mean the hop is on the trusted origin. - download_client()? - } else { - // DNS answers were validated; pin them so a direct connection - // can only reach the addresses that passed the guard. When a - // configured proxy carries the request instead, the proxy - // resolves the target and the pin is deliberately unused. - hop_client = pinned_client(¤t.url, &pinned)?; - &hop_client - }; - let response = send_request(client, ¤t).await?; - let status = response.status(); - if !is_redirect_status(status) { - return Ok(response); - } - let Some(location) = response.headers().get(reqwest::header::LOCATION) else { - return Ok(response); - }; - if redirect_count == MAX_REDIRECTS { - return Err(redirect_error("too many redirects", status)); - } - let location = location - .to_str() - .map(str::to_owned) - .map_err(|_| redirect_error("redirect location is not valid UTF-8", status))?; - // Dropping the unconsumed 3xx body releases its connection before a - // potentially slow DNS check for the next hop. - drop(response); - let base = url::Url::parse(¤t.url) - .map_err(|e| AiMuxError::InvalidArgument(format!("invalid request URL: {e}")))?; - let next = base - .join(&location) - .map_err(|e| redirect_error(format!("invalid redirect URL: {e}"), status))?; - // Fetch treats a redirect to a non-HTTP(S) scheme (data:, file:, ...) - // as a network error; following one would let the redirecting server - // fabricate a response outside the transport. - if !matches!(next.scheme(), "http" | "https") { - return Err(redirect_error( - format!("redirect to non-HTTP scheme: {}", next.scheme()), - status, - )); - } - let next = next.to_string(); - let hop_trusted = crate::download_guard::hop_trusted_origin(trusted_origin, ¤t.url); - pinned = crate::download_guard::validate_download_target(&next, hop_trusted).await?; - // Credentials are scoped to the credential anchor, not the previous - // hop: once a redirect leaves it, headers cannot come back. - if !crate::download_guard::same_origin(&next, credential_anchor) { - crate::download_guard::retain_user_agent(&mut current.headers); - } - if status == reqwest::StatusCode::SEE_OTHER - || (matches!( - status, - reqwest::StatusCode::MOVED_PERMANENTLY | reqwest::StatusCode::FOUND - ) && matches!(current.method, HttpMethod::Post)) - { - current.method = HttpMethod::Get; - current.body = HttpBody::Empty; - } - current.url = next; - } - unreachable!("redirect loop returns on response or error") -} - -async fn send_with_retry_raw( - transport: &RequestTransport<'_>, - 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 = transport.send(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); @@ -1874,547 +985,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; - } - - #[tokio::test] - async fn pinned_client_resolves_only_through_validated_addresses() { - use tokio::io::{AsyncReadExt, AsyncWriteExt}; - - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let port = listener.local_addr().unwrap().port(); - let server = tokio::spawn(async move { - let (mut stream, _) = listener.accept().await.unwrap(); - let mut request = [0; 1024]; - let _ = stream.read(&mut request).await.unwrap(); - stream - .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok") - .await - .unwrap(); - }); - // .example never resolves in real DNS, so success proves the resolve - // override was used rather than the system resolver. - let url = format!("http://download.example:{port}/file"); - let client = - pinned_client(&url, &["127.0.0.1".parse().unwrap()]).expect("pinned client builds"); - let response = tokio::time::timeout(Duration::from_secs(2), client.get(url).send()) - .await - .expect("request must not hang") - .expect("the validated address must be used"); - assert_eq!(response.status(), reqwest::StatusCode::OK); - tokio::time::timeout(Duration::from_secs(2), server) - .await - .expect("test server must finish") - .unwrap(); - } -} diff --git a/aimux-provider-utils/src/lib.rs b/aimux-provider-utils/src/lib.rs index 8d89e6cf..679f24be 100644 --- a/aimux-provider-utils/src/lib.rs +++ b/aimux-provider-utils/src/lib.rs @@ -2,16 +2,22 @@ //! //! 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; mod download_guard; +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 @@ -21,20 +27,20 @@ pub mod ws; 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, HttpMethod, HttpRequest, HttpResponse, HttpStreamResponse, PoolConfig, ProxyConfig, - RequestTimeout, TimeoutConfig, init_proxy, send, send_stream, send_stream_timed, send_timed, - send_validated, shared_client, shared_streaming_client, sleep_or_abort, + ExchangeContext, HttpBody, HttpRequest, ProxyConfig, init_proxy, shared_client, sleep_or_abort, }; -pub use logging::{body_logging_enabled, init_logging, redact_body}; +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_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/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..075d3361 --- /dev/null +++ b/aimux-provider-utils/src/post_to_api.rs @@ -0,0 +1,268 @@ +//! 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> { + // A provider whose endpoint legitimately holds the connection longer + // (e.g. Replicate `prefer: wait`) declares its own bound on the request. + let response_timeout = request + .response_timeout + .unwrap_or(EXCHANGE_RESPONSE_TIMEOUT); + call_to_api_with_deadline( + request, + request_body_values, + successful_response_handler, + failed_response_handler, + 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(), + max_json_response_bytes: request.max_json_response_bytes, + }; + + 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, + ..Default::default() + }; + 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..bd5d1934 --- /dev/null +++ b/aimux-provider-utils/src/read_response_with_size_limit.rs @@ -0,0 +1,143 @@ +//! Bounded response reads. + +use bytes::Bytes; +use futures::StreamExt; + +use aimux_core::{AiMuxError, ApiCallError}; + +/// 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 +/// +/// 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..db568ef8 --- /dev/null +++ b/aimux-provider-utils/src/response_handler.rs @@ -0,0 +1,456 @@ +//! 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, DEFAULT_MAX_JSON_RESPONSE_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, + /// 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. +#[derive(Debug)] +pub struct ResponseHandlerOutput { + pub value: T, + pub raw_value: Option, + pub response_headers: std::collections::HashMap, +} + +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. +/// +/// 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 + 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 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, + max_bytes, + input.abort_signal.as_ref(), + ) + .await?; + // 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()), + response_headers: Some(headers.clone()), + ..ApiCallError::new( + format!("Invalid JSON response: {error}"), + input.url.clone(), + input.request_body_values.clone(), + ) + })) + })?; + let raw_value = serde_json::from_slice::(&body).ok(); + Ok(ResponseHandlerOutput { + value, + raw_value, + response_headers: 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 +} + +/// 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( + 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: 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: 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: 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: 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 6c834550..c45d20ea 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 { @@ -156,21 +168,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. @@ -200,12 +234,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}"))) } } } @@ -220,12 +254,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}"))) } } } @@ -244,7 +278,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()))); } @@ -275,12 +309,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}")))); } }, } @@ -328,10 +362,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..6b365714 --- /dev/null +++ b/aimux-provider-utils/tests/api_helpers_test.rs @@ -0,0 +1,373 @@ +//! 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, + + ..Default::default() + } +} + +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.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"); +} + +/// 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 a6478bf7..52995b51 100644 --- a/aimux-provider-utils/tests/download_guard_test.rs +++ b/aimux-provider-utils/tests/download_guard_test.rs @@ -1,4 +1,4 @@ -//! SSRF guard wiring tests for `send_validated`. +//! SSRF guard wiring tests for `get_from_api` with `validate_url`. //! //! The guard must reject provider-supplied URLs that point at //! private/loopback/link-local space (before any connection is attempted), @@ -13,21 +13,44 @@ use wiremock::matchers::{header, method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; use aimux_core::error::AiMuxError; -use aimux_provider_utils::response::DEFAULT_ERROR_STRUCTURE; -use aimux_provider_utils::{HttpBody, HttpMethod, HttpRequest, RetryConfig, send_validated}; - -fn request(url: String) -> HttpRequest { +use aimux_provider_utils::{ + HttpRequest, create_binary_response_handler, create_json_response_handler, + create_status_code_error_response_handler, get_from_api, +}; + +fn request( + url: String, + trusted_origin: Option<&str>, + credentialed_origin: Option<&str>, +) -> HttpRequest { HttpRequest { - method: HttpMethod::Get, url, headers: vec![("authorization".into(), "Bearer test".into())], - body: HttpBody::Empty, abort_signal: None, 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), } } +async fn download( + url: String, + trusted_origin: Option<&str>, + credentialed_origin: Option<&str>, +) -> Result { + get_from_api( + request(url, trusted_origin, credentialed_origin), + create_binary_response_handler(), + create_status_code_error_response_handler(), + ) + .await + .map(|output| output.value) +} + #[tokio::test] async fn rejects_non_public_literal_urls_before_connecting() { for url in [ @@ -36,15 +59,9 @@ async fn rejects_non_public_literal_urls_before_connecting() { "http://[::ffff:169.254.169.254]/meta", "http://10.1.2.3/file", ] { - let error = send_validated( - request(url.into()), - None, - None, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await - .expect_err("non-public literal must be rejected"); + let error = download(url.into(), None, None) + .await + .expect_err("non-public literal must be rejected"); assert!( matches!(error, AiMuxError::InvalidArgument(ref m) if m.contains("non-public")), "unexpected error for {url}: {error:?}" @@ -63,20 +80,18 @@ async fn trusted_origin_download_succeeds_and_keeps_auth_headers() { .mount(&server) .await; - let response = send_validated( - request(format!("{}/generated/file", server.uri())), - Some(&server.uri()), - Some(&server.uri()), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + let output = get_from_api( + request( + format!("{}/generated/file", server.uri()), + Some(&server.uri()), + Some(&server.uri()), + ), + create_json_response_handler::(), + create_status_code_error_response_handler(), ) .await .unwrap(); - assert_eq!(response.status, 200); - assert_eq!( - serde_json::from_slice::(&response.body).unwrap(), - json!({"value": "ok"}) - ); + assert_eq!(output.value, json!({"value": "ok"})); } #[tokio::test] @@ -95,17 +110,14 @@ async fn follows_a_relative_redirect_within_the_trusted_origin() { .mount(&server) .await; - let response = send_validated( - request(format!("{}/start", server.uri())), + let body = download( + format!("{}/start", server.uri()), Some(&server.uri()), Some(&server.uri()), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, ) .await .unwrap(); - assert_eq!(response.status, 200); - assert_eq!(response.body.as_ref(), b"done"); + assert_eq!(body.as_ref(), b"done"); } #[tokio::test] @@ -120,12 +132,10 @@ async fn rejects_a_redirect_to_a_non_public_target() { .mount(&server) .await; - let error = send_validated( - request(format!("{}/start", server.uri())), + let error = download( + format!("{}/start", server.uri()), Some(&server.uri()), Some(&server.uri()), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, ) .await .expect_err("redirect to metadata IP must be rejected"); @@ -151,12 +161,10 @@ async fn rejects_a_redirect_onto_a_foreign_loopback_origin() { .mount(&server) .await; - let error = send_validated( - request(format!("{}/start", server.uri())), + let error = download( + format!("{}/start", server.uri()), Some(&server.uri()), Some(&server.uri()), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, ) .await .expect_err("foreign loopback origin must be rejected"); @@ -180,12 +188,10 @@ async fn rejects_a_redirect_to_a_data_url() { .mount(&server) .await; - let error = send_validated( - request(format!("{}/start", server.uri())), + let error = download( + format!("{}/start", server.uri()), Some(&server.uri()), Some(&server.uri()), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, ) .await .expect_err("redirect to a data: URL must be rejected"); @@ -205,16 +211,18 @@ async fn sanitizes_metadata_and_cookie_headers_from_download_requests() { .mount(&server) .await; - let mut req = request(format!("{}/file", server.uri())); + let mut req = request( + format!("{}/file", server.uri()), + Some(&server.uri()), + Some(&server.uri()), + ); req.headers.push(("Cookie".into(), "session=secret".into())); req.headers .push(("Metadata-Flavor".into(), "Google".into())); - send_validated( + get_from_api( req, - Some(&server.uri()), - Some(&server.uri()), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + create_binary_response_handler(), + create_status_code_error_response_handler(), ) .await .unwrap(); diff --git a/aimux-provider-utils/tests/http_recording_test.rs b/aimux-provider-utils/tests/http_recording_test.rs index 145433ce..f17933fc 100644 --- a/aimux-provider-utils/tests/http_recording_test.rs +++ b/aimux-provider-utils/tests/http_recording_test.rs @@ -1,540 +1,284 @@ -//! 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, + )), + + ..Default::default() } } -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); -} - -#[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); + recorder.record_transport_closed(call_id); + recorder.flush(); } -#[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..2d7d0799 100644 --- a/aimux-provider-utils/tests/logging_test.rs +++ b/aimux-provider-utils/tests/logging_test.rs @@ -1,227 +1,106 @@ -//! 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, + + ..Default::default() } } -/// 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)); -} 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..2f08101b 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,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 = send( - 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, - }, + // + // 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, - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; + 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 = - 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..5f2c5e39 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + anthropic_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { @@ -345,7 +360,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..7a53b146 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,17 @@ 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, + ..Default::default() + }, + http_body, + )) } /// Read a string field, dropping absent / non-string values. @@ -681,17 +711,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 +727,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 +803,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 +819,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; + 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 +861,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 +879,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 +974,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 +1075,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 +1133,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 +1154,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 +1175,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..52699f9a 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + crate::anthropic::anthropic_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { @@ -278,7 +277,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..3bae0555 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)] @@ -206,24 +236,29 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); + // 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 + // retry passes through, so `do_generate` is never replayed (§6.2). + let retries = retry::prepare_retries( + options.max_retries, + self.retry_config(), + options.abort_signal.clone(), + ); + // Step 1: Upload audio. - let resp = send( - 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, - ) - .await?; - - let upload: AssemblyAIUploadResponse = serde_json::from_slice(&resp.body)?; + 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; // Step 2: Submit transcript request. let mut body = Map::new(); @@ -254,23 +289,19 @@ impl TranscriptionModel for AssemblyAITranscriptionModel { } } - let resp = send( - 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, - ) - .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 = serde_json::from_slice(&resp.body)?; + let submit: AssemblyAISubmitResponse = resp.value; // Step 3: Poll for completion. let mut raw_body: Value; @@ -282,26 +313,23 @@ 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?; + let poll_url = self.transcript_status_url(&submit.id); + let resp = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest::new(poll_url.clone(), header_list.clone(), options), + aimux_provider_utils::create_json_response_handler::< + AssemblyAITranscriptResponse, + >(), + assemblyai_failed_response_handler(), + ) + }) + .await?; - response_headers = resp.headers; + response_headers = resp.response_headers; - raw_body = serde_json::from_slice(&resp.body)?; - let parsed: AssemblyAITranscriptResponse = serde_json::from_value(raw_body.clone())?; + 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 +436,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..5216192f 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,25 @@ 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, + ..Default::default() }, - 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; - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); @@ -470,48 +482,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 +521,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..b7376abe 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + crate::openai::openai_failed_response_handler(), + ) + }) + .await?; // Azure response: { data: [{ id, model, ... }] } #[derive(serde::Deserialize)] @@ -357,7 +357,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 +479,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 +506,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 +513,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 +526,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..17926a4f 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,29 @@ 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( - 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(), - }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + 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 status = resp.status; - let response_headers = resp.headers; - - let data: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + 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 +318,31 @@ 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( - 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(), - }, - retry_config, - &DEFAULT_ERROR_STRUCTURE, - options.timeout.map(Into::into), + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + 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 status = resp.status; - let response_headers = resp.headers; + let response_headers = resp.response_headers; - 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..56db97f2 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,26 @@ 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, + ..Default::default() }, - 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; - 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..cbebdedb 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,24 @@ 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, + ..Default::default() }, - 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; + 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 +317,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..df28bcfd 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + bedrock_failed_response_handler(), + ) + }) + .await?; // AWS response: { modelSummaries: [{ modelId, modelName, ... }] } #[derive(serde::Deserialize)] @@ -321,7 +345,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..d7ad90f0 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: 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,25 @@ 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(), + ..Default::default() }, - 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; - 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 +264,30 @@ 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(), + ..Default::default() }, - 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; // 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 +689,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,27 @@ 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, + ..Default::default() }, - 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; - 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 04019dec..ebfa09af 100644 --- a/aimux-providers/src/black_forest_labs.rs +++ b/aimux-providers/src/black_forest_labs.rs @@ -17,12 +17,50 @@ 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, send_validated, - sleep_or_abort, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, sleep_or_abort, without_trailing_slash}; + +/// AI SDK's `isTrustedUrl` (black-forest-labs-api.ts): credentials may go to +/// the configured origin, or over HTTPS to `bfl.ai` and its subdomains — BFL +/// serves polling and asset URLs from regional clusters. Allowlisted hosts +/// still get full URL/DNS validation; this gates only the headers. +fn bfl_trusted_url(url: &str, base_url: &str) -> bool { + if aimux_provider_utils::same_origin(url, base_url) { + return true; + } + let Ok(parsed) = url::Url::parse(url) else { + return false; + }; + parsed.scheme() == "https" + && parsed + .host_str() + .is_some_and(|host| host == "bfl.ai" || host.ends_with(".bfl.ai")) +} + +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; @@ -140,23 +178,6 @@ fn gcd(a: u32, b: u32) -> u32 { if b == 0 { a } else { gcd(b, a % b) } } -/// AI SDK's `isTrustedUrl` (black-forest-labs-api.ts): credentials may go to -/// the configured origin, or over HTTPS to `bfl.ai` and its subdomains — BFL -/// serves polling and asset URLs from regional clusters. Allowlisted hosts -/// still get full URL/DNS validation; this gates only the headers. -fn bfl_trusted_url(url: &str, base_url: &str) -> bool { - if aimux_provider_utils::same_origin(url, base_url) { - return true; - } - let Ok(parsed) = url::Url::parse(url) else { - return false; - }; - parsed.scheme() == "https" - && parsed - .host_str() - .is_some_and(|host| host == "bfl.ai" || host.ends_with(".bfl.ai")) -} - #[async_trait] impl ImageModel for BlackForestLabsImageModel { fn provider(&self) -> &str { @@ -302,33 +323,16 @@ impl ImageModel for BlackForestLabsImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list: Vec<(String, String)> = headers.into_iter().collect(); - // AI SDK gates headers per URL via isTrustedUrl; response-supplied - // targets outside the BFL allowlist get none. - let gated_headers = |url: &str| -> Vec<(String, String)> { - if bfl_trusted_url(url, &self.config.base_url) { - header_list.clone() - } else { - vec![] - } - }; // Submit - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.submit_endpoint(), header_list.clone(), options), + 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") @@ -342,6 +346,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 @@ -368,29 +377,45 @@ impl ImageModel for BlackForestLabsImageModel { let mut result_end_time = None; let mut result_duration = None; + // AI SDK gates headers per URL via isTrustedUrl; response-supplied + // targets outside the BFL allowlist get none. + let gated_headers = |url: &str| -> Vec<(String, String)> { + if bfl_trusted_url(url, &self.config.base_url) { + header_list.clone() + } else { + vec![] + } + }; + for _ in 0..max_attempts { // AI SDK polls polling_url with validateUrl: true and gates the // headers itself via isTrustedUrl (base_url origin or HTTPS // *.bfl.ai), so credentialed_origin is None here. let poll_url = poll_url_with_id.to_string(); - let pr = send_validated( - HttpRequest { - method: HttpMethod::Get, - headers: gated_headers(&poll_url), - url: poll_url, - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - Some(&self.config.base_url), - 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.clone(), + headers: gated_headers(&poll_url), + + 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()), + // Headers are gated per URL by gated_headers above. + credentialed_origin: 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") @@ -416,13 +441,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), @@ -438,26 +467,31 @@ impl ImageModel for BlackForestLabsImageModel { })?; // Download image; result.sample is a URL from the poll response body. - // AI SDK sends its headers to trusted BFL hosts on the download too - // (isTrustedUrl-gated), so mirror the poll's header policy. - let ir = send_validated( - HttpRequest { - method: HttpMethod::Get, - headers: gated_headers(&image_url), - url: image_url, - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - Some(&self.config.base_url), - None, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - let image_bytes = ir.body.to_vec(); + // AI SDK sends its headers to trusted BFL hosts on the download too, + // gated by the same allowlist. + let ir = retries + .retry(|| { + aimux_provider_utils::get_from_api( + HttpRequest { + url: image_url.clone(), + headers: gated_headers(&image_url), + + 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()), + // Headers are gated per URL by gated_headers above. + credentialed_origin: 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 @@ -501,22 +535,3 @@ impl ImageModel for BlackForestLabsImageModel { }) } } - -#[cfg(test)] -mod tests { - use super::bfl_trusted_url; - - #[test] - fn credentials_go_to_the_base_url_and_bfl_hosts_only() { - let base = "https://api.bfl.ai"; - assert!(bfl_trusted_url("https://api.bfl.ai/v1/get_result", base)); - assert!(bfl_trusted_url( - "https://api.us1.bfl.ai/v1/get_result", - base - )); - assert!(bfl_trusted_url("https://bfl.ai/x", base)); - assert!(!bfl_trusted_url("http://api.us1.bfl.ai/x", base)); - assert!(!bfl_trusted_url("https://evil-bfl.ai/x", base)); - assert!(!bfl_trusted_url("https://attacker.example/x", base)); - } -} diff --git a/aimux-providers/src/cartesia.rs b/aimux-providers/src/cartesia.rs index 619baea2..c57d9613 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,17 @@ impl SpeechModel for CartesiaSpeechModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - let audio_bytes = resp.body.to_vec(); + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); @@ -756,27 +767,18 @@ impl TranscriptionModel for CartesiaTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( - 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, + let resp = aimux_provider_utils::post_to_api( + 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.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); + let response_headers = resp.response_headers; - 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..ccf99d3e 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,25 @@ impl CodexModel { let headers = self.inner().build_headers(opts.headers.as_ref()); let (body, warnings) = self.body(&opts, true)?; - let resp = send_stream_timed( - HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), - 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), + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + 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(), ) .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; + 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 +319,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 +332,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 +368,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 +389,38 @@ 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(), + ..Default::default() }, - 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; + 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 +441,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 +529,22 @@ 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, + ..Default::default() }, - 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..1c21a849 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,17 @@ impl EmbeddingModel for CohereEmbeddingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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..b4cc7531 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + cohere_failed_response_handler(), + ) + }) + .await?; // Cohere v2: { models: [{ name, endpoints, ... }] } #[derive(serde::Deserialize)] struct Resp { @@ -165,7 +180,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..3ea454fd 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,23 @@ 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( - HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), - headers: headers + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new( + self.endpoint(), + 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), + options, + ), + 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; + let data: ChatResponse = resp.value; // Build content array. let mut content = Vec::new(); @@ -272,35 +258,43 @@ 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( - HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), - headers: headers + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new( + endpoint.clone(), + 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), + options, + ), + 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; + 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 +304,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 +314,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 +515,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 +572,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..9a2c9181 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,19 @@ impl RerankingModel for CohereRerankingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), header_list, options), + 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; - 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..b0016adb 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,27 @@ 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, + ..Default::default() }, - 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; - 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..2647b735 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,26 @@ 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, + ..Default::default() }, - 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; - 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..46f7176b 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,27 @@ 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, + ..Default::default() }, - 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; + let audio_bytes = resp.value.to_vec(); let timestamp = chrono::Utc::now().to_rfc3339(); @@ -640,29 +677,24 @@ impl TranscriptionModel for ElevenLabsTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( - HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), - headers: headers + let resp = aimux_provider_utils::post_to_api( + HttpRequest::new( + self.endpoint(), + 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, + options, + ), + 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; + 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..61b8d4ed 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,24 @@ 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, + ..Default::default() }, - 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; + let response_body = resp.raw_value; + let parsed: ExaResponse = resp.value; Ok(SearchResult { results: map_results(parsed.results), @@ -193,7 +207,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 3bf570b4..32460f8f 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, send_validated, - 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,96 @@ 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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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; + 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 +566,36 @@ 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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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; + 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,25 +609,30 @@ 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()) { - // images[].url comes from the queue result response body. - let ir = send_validated( - 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, - }, - Some(&self.config.base_url), - Some(&self.config.base_url), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - downloaded.push(ir.body.to_vec()); + // images[].url comes from the queue result response body, so + // it goes through the SSRF download guard. + 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, + 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()), + }, + aimux_provider_utils::create_binary_response_handler(), + fal_failed_response_handler(), + ) + }) + .await?; + downloaded.push(ir.value.to_vec()); } } @@ -636,7 +699,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`. @@ -644,7 +708,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, @@ -726,7 +791,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(); @@ -752,76 +820,94 @@ 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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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: Some(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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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; + let raw_body = resp.value; // Extract video URL from response. let videos: Vec = if let Some(video) = raw_body @@ -844,15 +930,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..f917f64f 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,24 @@ 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, + ..Default::default() }, - 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; + let response_body = resp.raw_value; + let parsed: FirecrawlResponse = resp.value; Ok(SearchResult { results: map_results(parsed.data.web), @@ -209,7 +214,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 7f339c96..fbaa7af6 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, send_validated, 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)] @@ -201,31 +221,40 @@ impl TranscriptionModel for GladiaTranscriptionModel { let headers = self.build_headers(options.headers.as_ref()); + // 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 + // retry passes through, so `do_generate` is never replayed (§6.2). + 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 = send( - 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, - ) - .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 = serde_json::from_slice(&resp.body)?; + let upload: GladiaUploadResponse = resp.value; // Step 2: Initiate transcription. let mut body = Map::new(); @@ -243,26 +272,26 @@ impl TranscriptionModel for GladiaTranscriptionModel { } } - let resp = send( - 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, - ) - .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 = serde_json::from_slice(&resp.body)?; + let init: GladiaInitResponse = resp.value; // Step 3: Poll for result. let mut raw_body: Value; @@ -277,30 +306,35 @@ impl TranscriptionModel for GladiaTranscriptionModel { // AI SDK polls result_url with validateUrl: true and // credentialedOrigin = the API origin: the target is validated // and headers survive only while it stays on base_url's origin. - let resp = send_validated( - 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, - }, - Some(&self.config.base_url), - Some(&self.config.base_url), - 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, + 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()), + }, + 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; + raw_body = resp.raw_value.unwrap_or(Value::Null); + let parsed = resp.value; if parsed.status == "done" { let result = parsed.result.ok_or_else(|| { @@ -348,13 +382,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..4eb061c4 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,25 @@ 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, + ..Default::default() }, - 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; - 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 +203,25 @@ 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, + ..Default::default() }, - 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; - 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 d0e420c5..d0256391 100644 --- a/aimux-providers/src/google/files.rs +++ b/aimux-providers/src/google/files.rs @@ -25,19 +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, send_validated, 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 { @@ -117,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, @@ -197,31 +189,58 @@ impl Files for GoogleFiles { )); init_headers.push(("Content-Type".to_string(), "application/json".to_string())); - let init_resp = send( - 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, - ) - .await - .map_err(|e| match e { - AiMuxError::ApiCall(d) => AiMuxError::ApiCall(ApiCallError { - message: format!("Failed to initiate resumable upload: {}", d.message), - ..d - }), - e => e, - })?; + // 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, + 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 + // 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), + ..*d + })), + e => e, + }) + }) + .await?; let upload_url = init_resp - .headers + .response_headers .get("x-goog-upload-url") .cloned() .ok_or_else(|| { @@ -239,38 +258,42 @@ impl Files for GoogleFiles { ), ]; + let upload_request_url = upload_url.clone(); // 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 = send_validated( - 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, - }, - Some(&self.config.base_url), - Some(&self.config.base_url), - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, - ) - .await - .map_err(|e| match e { - AiMuxError::ApiCall(d) => AiMuxError::ApiCall(ApiCallError { - message: format!("Failed to upload file data: {}", d.message), - ..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 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 + .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; // Step 3: Poll if file is PROCESSING. let poll_interval_ms = google_options.poll_interval_ms.unwrap_or(2000); @@ -283,9 +306,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!( @@ -302,37 +325,47 @@ impl Files for GoogleFiles { let poll_url = format!("{}/{}", self.config.base_url, file.name); - let poll_resp = send( - HttpRequest { - method: HttpMethod::Get, - url: poll_url, - 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, - ) - .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); + 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()); + 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..75ff4da8 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,17 @@ impl GoogleImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list = build_header_list(&headers); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +312,17 @@ impl GoogleImageModel { let headers = self.build_headers(options.headers.as_ref()); let header_list = build_header_list(&headers); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +350,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..0a3fc576 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + google_failed_response_handler(), + ) + }) + .await?; // Gemini response: { models: [{ name: "models/gemini-...", displayName, ... }] } #[derive(serde::Deserialize)] @@ -211,7 +257,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..62a303f7 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,21 @@ 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( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +189,38 @@ 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( - HttpRequest { - method: HttpMethod::Post, - url: self.stream_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), + let endpoint = self.stream_endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +254,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 +587,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 +887,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..dce85c6b 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,24 @@ 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, + ..Default::default() }, - 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 +145,68 @@ 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: Some(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::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; + 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 +235,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..7be6d7e6 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,25 @@ 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, + ..Default::default() }, - 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; - 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..b0ea9b2a 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: Some(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,19 @@ impl LanguageModel for HuggingFaceResponsesModel { let body = request.body; let headers = self.build_headers(options.headers.as_ref()); - let resp = send_timed( - 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), + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), build_header_list(&headers), options), + 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; - // 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 +238,33 @@ impl LanguageModel for HuggingFaceResponsesModel { let body = request.body; let headers = self.build_headers(options.headers.as_ref()); - let resp = send_stream_timed( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; - let sse_stream = SseStream::new(resp.body); + 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), + 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 +277,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 +533,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 +1070,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 +1215,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..d71b4a0a 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,25 @@ 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, + ..Default::default() }, - 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; - 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..cea88fbb 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,27 @@ 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, + ..Default::default() }, - 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; - 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..7bb6739e 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,35 @@ impl VideoModel for KlingAIVideoModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit task. - let resp = send( - HttpRequest { - method: HttpMethod::Post, - url: self.submit_endpoint(&mode), - 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, + let submit_url = self.submit_endpoint(&mode); + let request_body = Value::Object(body); + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(submit_url.clone(), header_list, options), + 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 +334,114 @@ 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: Some(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 headers = self.build_headers(options.headers.as_ref()); + let header_list: Vec<(String, String)> = headers.into_iter().collect(); - let result: KlingAITaskResult = serde_json::from_slice(&resp.body)?; + let poll_url = self.poll_endpoint(mode, id, task_id); + let resp = aimux_provider_utils::get_from_api( + HttpRequest::new(poll_url.clone(), header_list, options), + aimux_provider_utils::create_json_response_handler::(), + klingai_failed_response_handler(), + ) + .await?; - if result.code != 0 { - return Err(AiMuxError::ApiCall(ApiCallError { - status_code: Some(resp.status), - provider_code: Some(result.code.to_string()), - message: result + let response_headers = resp.response_headers; + 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..2cb7ab4a 100644 --- a/aimux-providers/src/linkup.rs +++ b/aimux-providers/src/linkup.rs @@ -19,10 +19,7 @@ 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}; /// Fixed model ID for the Linkup search model. const MODEL_ID: &str = "linkup-search"; @@ -224,28 +221,27 @@ 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, + ..Default::default() }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler::(), + aimux_provider_utils::create_standard_json_error_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; - 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..2c224d84 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,17 @@ impl SpeechModel for LMNTSpeechModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 2ba2403a..56d8d239 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, send_validated, - 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,15 @@ impl ImageModel for LumaImageModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; - let submit_body: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers; + let submit_body: Value = resp.value; let generation_id = submit_body .get("id") @@ -328,26 +334,30 @@ 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::new( + self.generations_url(Some(&generation_id)), + header_list.clone(), + options, + ), + 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 +374,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), @@ -386,25 +400,30 @@ impl ImageModel for LumaImageModel { )) })?; - // Download image; assets.image is a URL from the poll response body. - let ir = send_validated( - HttpRequest { - method: HttpMethod::Get, - url: image_url, - headers: vec![], - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - Some(&self.config.base_url), - Some(&self.config.base_url), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - let image_bytes = ir.body.to_vec(); + // Download image; assets.image is a URL from the poll response body, + // so it goes through the SSRF download guard. + 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, + 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()), + }, + 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..6e1411de 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,17 @@ impl EmbeddingModel for MistralEmbeddingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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..7b245b27 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + mistral_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { #[serde(default)] @@ -157,7 +214,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..547a877a 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,23 @@ 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( - HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), - headers: headers + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new( + self.endpoint(), + 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), + options, + ), + 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; + let data: ChatCompletionResponse = resp.value; let choice = data.choices.into_iter().next().ok_or_else(|| { @@ -336,43 +321,44 @@ 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( - HttpRequest { - method: HttpMethod::Post, - url: self.endpoint(), - headers: headers + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new( + self.endpoint(), + 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), + options, + ), + 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; + 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 +381,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 +399,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 +558,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..1675befc 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: Some(output.response_headers.clone()), + ..ApiCallError::new(message, url, request_body_values) + }))); + } + Ok(output) + }) +} // == Config == @@ -267,47 +314,21 @@ impl LanguageModel for OpenResponsesModel { let (body, warnings) = build_request_body(&self.model_id, options, &self.config.provider_options_name); - let resp = send_timed( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +463,54 @@ impl LanguageModel for OpenResponsesModel { b }; - let resp = send_stream_timed( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +533,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 +797,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 e043e65c..14def1ed 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,19 @@ impl EmbeddingModel for OpenAIEmbeddingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), header_list, options), + 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; - 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 diff --git a/aimux-providers/src/openai/files.rs b/aimux-providers/src/openai/files.rs index 99cb8873..2574cd76 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,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 = send( - 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, - }, + // + // 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, - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - - let data: OpenAIFilesResponse = serde_json::from_slice::(&resp.body)?; + 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; // 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..ec157c4c 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,35 @@ impl ImageModel for OpenAIImageModel { let (form_body, content_type) = build_edit_multipart(&self.model_id, options, &openai_options); - let resp = send( - 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, + let resp = aimux_provider_utils::post_to_api( + 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(), ) .await?; - let val: Value = serde_json::from_slice(&resp.body)?; - (val, resp.headers, None) + let val: Value = resp.value; + (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 = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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(), ) .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, 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..59837557 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,25 @@ 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( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + 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). - 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 +426,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 +433,36 @@ 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( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; + let response_headers = resp.response_headers; - 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 +474,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 +508,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 +522,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 +552,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 +762,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 +886,35 @@ 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, + ..Default::default() + }, + 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..91afa897 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,29 @@ impl LanguageModel for OpenAIResponsesModel { let body = request_result.body; let provider_key = self.provider_options_name().to_string(); - let resp = send_timed( - 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), + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + 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 status = resp.status; - let response_headers = resp.headers; - - let data: Value = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + 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 +170,31 @@ impl LanguageModel for OpenAIResponsesModel { .and_then(serde_json::Value::as_bool) == Some(true); - let resp = send_stream_timed( - 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), + let endpoint = self.endpoint(); + let resp = aimux_provider_utils::post_json_to_api( + 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 status = resp.status; - let response_headers = resp.headers; + let response_headers = resp.response_headers; - 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..d1d30241 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,19 @@ impl SpeechModel for OpenAISpeechModel { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), header_list, options), + 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; - 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..b0ca5777 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,18 @@ 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( - 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, + let resp = aimux_provider_utils::post_to_api( + 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.headers; - - let raw_body: Value = serde_json::from_slice(&resp.body).unwrap_or(Value::Null); + let response_headers = resp.response_headers; - 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 +487,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 +506,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 +574,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 +640,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..41c1251e 100644 --- a/aimux-providers/src/parallel_ai.rs +++ b/aimux-providers/src/parallel_ai.rs @@ -19,10 +19,7 @@ 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}; /// Fixed model ID for the Parallel AI search model. const MODEL_ID: &str = "parallel-search"; @@ -183,27 +180,26 @@ 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, + ..Default::default() }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler::(), + aimux_provider_utils::create_standard_json_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; - 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..fd0cba85 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,15 @@ 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( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), header_list, options), + body, + aimux_provider_utils::create_binary_response_handler(), + prodia_failed_response_handler(), ) .await?; - let rh = resp.headers; + let rh = resp.response_headers; let content_type = rh.get("content-type").cloned().unwrap_or_default(); // Extract boundary @@ -264,7 +281,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 +331,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 +379,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 +399,20 @@ impl VideoModel for ProdiaVideoModel { let header_list: Vec<(String, String)> = headers.into_iter().collect(); // Submit job. - let resp = send( - HttpRequest { - method: HttpMethod::Post, - url: format!("{}/job", self.config.base_url), - headers: header_list.clone(), - body: HttpBody::Json(body), - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new( + format!("{}/job", self.config.base_url), + header_list, + options, + ), + 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 +423,68 @@ 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: Some(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::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; + 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 +500,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 e9f52e8e..ab208cbe 100644 --- a/aimux-providers/src/recraft.rs +++ b/aimux-providers/src/recraft.rs @@ -19,12 +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, send_validated, - 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"; @@ -211,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, base_url: &str, ) -> Result { let items = response.get("data").and_then(|d| d.as_array()); @@ -242,25 +258,30 @@ async fn extract_images( let mut binaries = Vec::with_capacity(urls.len()); for url in &urls { - // data[].url is a generated-image URL from the response body. - let resp = send_validated( - HttpRequest { - method: HttpMethod::Get, - url: url.clone(), - headers: vec![], - body: HttpBody::Empty, - - abort_signal: abort_signal.clone(), - call_id: None, - recording_context: None, - }, - Some(base_url), - Some(base_url), - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, - ) - .await?; - binaries.push(resp.body.to_vec()); + // data[].url is a generated-image URL from the response body, so it + // goes through the SSRF download guard. + 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, + 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()), + }, + aimux_provider_utils::create_binary_response_handler(), + recraft_failed_response_handler(), + ) + }) + .await?; + binaries.push(resp.value.to_vec()); } Ok(ImageOutputs::Binary(binaries)) } @@ -298,27 +319,41 @@ 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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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; + 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(), &self.config.base_url).await?; + let images = extract_images( + &value, + &retries, + options.abort_signal.clone(), + &self.config.base_url, + ) + .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 f6887869..09eb8b21 100644 --- a/aimux-providers/src/replicate.rs +++ b/aimux-providers/src/replicate.rs @@ -4,6 +4,7 @@ //! (`reference/ai/packages/replicate/src/replicate-image-model.ts`). use std::collections::HashMap; +use std::time::Duration; use async_trait::async_trait; use serde_json::{Map, Value, json}; @@ -13,12 +14,31 @@ 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, send_validated, - sleep_or_abort, without_trailing_slash, -}; +use aimux_provider_utils::{HttpRequest, load_api_key, without_trailing_slash}; + +/// Replicate's `prefer: wait` holds the connection for at most this long when +/// no explicit duration is given (their documented default and maximum). +const DEFAULT_WAIT_SECONDS: u64 = 60; + +/// Headroom added on top of the wait window for connect/TLS/body transport +/// when sizing the create exchange's `response_timeout`. +const WAIT_TRANSPORT_MARGIN_SECONDS: u64 = 5; + +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 +291,39 @@ 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, + // This exchange legitimately holds the connection for the + // whole wait window; widen the hang guard past it so a + // slow-but-alive generation is not misread as a dead + // transport (retryable → re-create → double billing). + 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, }, - 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; + 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"] { @@ -300,28 +335,32 @@ impl ImageModel for ReplicateImageModel { _ => Vec::new(), }; - // Download images + // Download images; output URLs come from the prediction response + // body, so they go through the SSRF download guard. let mut downloaded: Vec> = Vec::new(); for url in &urls { - // output URLs come from the prediction response body. - let ir = send_validated( - HttpRequest { - method: HttpMethod::Get, - url: url.clone(), - headers: vec![], - body: HttpBody::Empty, - - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - Some(&self.config.base_url), - Some(&self.config.base_url), - 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, + 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()), + }, + aimux_provider_utils::create_binary_response_handler(), + replicate_failed_response_handler(), + ) + }) + .await?; + downloaded.push(ir.value.to_vec()); } Ok(ImageResult { @@ -343,7 +382,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`. @@ -351,7 +391,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, @@ -408,7 +449,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(); @@ -431,23 +475,28 @@ 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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: 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()) @@ -456,51 +505,80 @@ 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: Some(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, + response_timeout: None, + max_json_response_bytes: None, + validate_url: false, + trusted_origin: None, + credentialed_origin: None, + }, + aimux_provider_utils::create_json_response_handler::(), + replicate_failed_response_handler(), + ) + .await?; + + 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 + .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. @@ -528,15 +606,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..8e2a3c14 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,39 @@ 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( - HttpRequest { - method: HttpMethod::Post, - url: self.jobs_url(), - headers: headers + let resp = aimux_provider_utils::post_form_data_to_api( + HttpRequest::new( + submit_url.clone(), + 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, + options, + ), + 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 +265,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 +280,63 @@ 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::new( + poll_url.clone(), + headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + options, + ), + 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::new( + self.transcript_url(&job_id), + headers + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(), + options, + ), + 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; + 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..be5e8bc5 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,104 @@ impl VideoModel for RunwaymlVideoModel { let submit_url = format!("{}{submit_path}", self.config.base_url); // Submit the task. - let resp = send( - HttpRequest { - method: HttpMethod::Post, - url: submit_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(), - &RUNWAYML_ERROR_STRUCTURE, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(submit_url, header_list, options), + 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 response_headers = resp.response_headers; + let task: RunwaymlTaskCreationResponse = resp.value; - let timestamp = chrono::Utc::now().to_rfc3339(); - - 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), }, }) } + + 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::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(); + 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..f4e58db8 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,30 @@ 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, + ..Default::default() }, - 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; - 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..5c87a912 100644 --- a/aimux-providers/src/serper.rs +++ b/aimux-providers/src/serper.rs @@ -15,10 +15,7 @@ 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}; const MODEL_ID: &str = "serper-search"; @@ -161,24 +158,24 @@ 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, + ..Default::default() }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + body, + aimux_provider_utils::create_json_response_handler(), + aimux_provider_utils::create_standard_json_error_response_handler(), ) .await?; - let response_headers = resp.headers; - - let parsed: SerperResponse = serde_json::from_slice(&resp.body)?; + let response_headers = resp.response_headers; + let response_body = resp.raw_value; + let parsed: SerperResponse = resp.value; Ok(SearchResult { results: map_results(parsed.organic), @@ -187,7 +184,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..d727cba9 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: headers, + }) + }) +} /// Map a Stability model ID to its generate-endpoint sub-path. /// @@ -267,37 +326,27 @@ 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( - 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, + let resp = aimux_provider_utils::post_to_api( + 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.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; + 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..a01ab7ec 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,24 @@ 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, + ..Default::default() }, - 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; + let response_body = resp.raw_value; + let parsed: TavilyResponse = resp.value; Ok(SearchResult { results: map_results(parsed.results), @@ -211,7 +217,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..06bb368c 100644 --- a/aimux-providers/src/tinyfish.rs +++ b/aimux-providers/src/tinyfish.rs @@ -13,17 +13,14 @@ 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}; /// Provider canonical name. const PROVIDER_NAME: &str = "tinyfish"; @@ -204,38 +201,32 @@ 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, + ..Default::default() }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + aimux_provider_utils::create_standard_json_error_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; - 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..afcccf47 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,23 @@ 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(), + ..Default::default() }, - 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 +234,44 @@ 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(), + ..Default::default() }, - 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; + 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 +280,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 +298,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 +372,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 +439,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 +483,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 +500,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..5ebd732b 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,25 @@ 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, + ..Default::default() }, - 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; - 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 +223,25 @@ 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, + ..Default::default() }, - 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; - 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..f93b5924 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,15 @@ impl VertexImageModel { let headers = self.build_headers(options.headers.as_ref()); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers; - let rb: Value = serde_json::from_slice(&resp.body)?; + let rh = resp.response_headers; + let rb: Value = resp.value; let images: Vec = rb .get("predictions") @@ -342,23 +330,15 @@ 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( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.generate_content_endpoint(), headers, options), + 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; + 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 +408,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..9c828aa7 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,24 @@ 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, + ..Default::default() + }, + aimux_provider_utils::create_json_response_handler(), + crate::google::google_failed_response_handler(), + ) + }) + .await?; #[derive(serde::Deserialize)] struct Resp { @@ -402,7 +403,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..2b8b5d91 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,25 @@ 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(), + ..Default::default() }, - 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; - 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 +227,48 @@ 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(), + ..Default::default() }, - 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; // 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 +302,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 +533,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 +650,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,26 @@ 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, + ..Default::default() }, - 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; - 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..a5139f4c 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,16 @@ impl VideoModel for VertexVideoModel { .map(|(k, v)| (k.clone(), v.clone())) .collect(); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.predict_url(), header_list.clone(), options), + 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 +177,68 @@ 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: Some(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.vertex operation reference is missing operation_name".to_string(), + ) + })?; - let resp = send( - HttpRequest { - method: HttpMethod::Get, - url: self.operation_url(&operation_name), - headers: header_list.clone(), - body: HttpBody::Empty, + 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(); - abort_signal: options.abort_signal.clone(), - call_id: None, - recording_context: None, - }, - RetryConfig::default(), - &GOOGLE_ERROR_STRUCTURE, - ) - .await?; + let poll_url = self.operation_url(operation_name); + let resp = aimux_provider_utils::get_from_api( + HttpRequest::new(poll_url.clone(), header_list, options), + 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; + 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 +266,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..1e366693 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,17 @@ 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( - 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, + let resp = aimux_provider_utils::post_json_to_api( + 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.headers.clone(); + let response_headers: HashMap = resp.response_headers.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..01369286 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,19 @@ impl RerankingModel for VoyageRerankingModel { .collect(); header_list.push(("Content-Type".to_string(), "application/json".to_string())); - let resp = send( - 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, + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), header_list, options), + 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; - 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..104ea039 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,191 @@ 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(); + 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) + }))); + } + // 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()), + 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; + 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 +261,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 +280,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 +320,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..5799be56 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,19 @@ 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( - 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), + let resp = aimux_provider_utils::post_json_to_api( + HttpRequest::new(self.endpoint(), build_header_list(&headers), options), + 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; - 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 +248,34 @@ 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( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.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; - 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 +283,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 +309,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 +326,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 +522,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..70a85b07 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,18 @@ impl LanguageModel for XaiResponsesModel { let body = request_result.body; let provider_tool_names = request_result.provider_tool_names; - let resp = send_timed( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.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; - 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 +363,36 @@ 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( - 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), + let resp = aimux_provider_utils::post_json_to_api( + 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.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; + + 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 +414,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,16 +606,22 @@ 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; + 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 ── @@ -942,11 +889,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..dcab66c3 100644 --- a/aimux-providers/src/you_com.rs +++ b/aimux-providers/src/you_com.rs @@ -14,17 +14,14 @@ 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}; /// Provider canonical name. const PROVIDER_NAME: &str = "you_com"; @@ -206,38 +203,32 @@ 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, + ..Default::default() }, - RetryConfig::default(), - &DEFAULT_ERROR_STRUCTURE, + aimux_provider_utils::create_json_response_handler::(), + aimux_provider_utils::create_standard_json_error_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; - 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/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/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/gladia_transcription_test.rs b/aimux-providers/tests/gladia_transcription_test.rs index 9a696f0c..ab4ee57e 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; 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_files_test.rs b/aimux-providers/tests/google_files_test.rs index 17dd6fcc..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,7 +410,8 @@ async fn should_throw_when_initiation_request_fails() { .mount(&server) .await; - let provider = provider(&server); + // 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; @@ -456,7 +466,8 @@ async fn should_throw_when_upload_request_fails() { .mount(&server) .await; - let provider = provider(&server); + // 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; @@ -586,3 +597,55 @@ 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 provider = provider_with_retries(&server, 1); + 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/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_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/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, } } 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)); 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..34fc9de7 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;
@@ -2428,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;
@@ -2436,15 +2549,20 @@ 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() {}
 
         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) {
+                                 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; }
@@ -2470,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; }
@@ -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(); }
 
@@ -2487,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;
@@ -2495,13 +2619,16 @@ 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; }
             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; }
@@ -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..99f709af 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,41 @@ 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);
+    }
+
+    @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
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..0b380639 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",
@@ -141,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}"),
         )),
@@ -150,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`.
@@ -166,6 +173,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 +184,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 +225,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 +290,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..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
 // ─────────────────────────────────────────────────────────────────────────────
@@ -51,9 +98,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 +134,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 +170,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 +210,25 @@ 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_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;
+                        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 +270,24 @@ 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_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;
+                        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 +323,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)
@@ -307,15 +361,20 @@ 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_opts_json(
+                            "opts_json",
+                            s,
+                            &[("query", serde_json::Value::from(""))],
+                        )?;
+                        parsed.query = opts.query;
                         opts = parsed;
                     }
                 }
                 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 +721,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 +754,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 +789,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;
@@ -922,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/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..8bef0979 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..fa26fb2d 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,17 +13,25 @@
  * 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 = { 
+/**
+ * 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
@@ -37,7 +46,9 @@ status_code: number | 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, 
 /**
@@ -47,16 +58,13 @@ message: string,
  */
 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..53c07237 100644
--- a/bindings/node/src/types/CallOptions.ts
+++ b/bindings/node/src/types/CallOptions.ts
@@ -84,7 +84,7 @@ reasoning: ReasoningEffort | 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, 
 /**
diff --git a/bindings/node/src/types/EmbeddingCallOptions.ts b/bindings/node/src/types/EmbeddingCallOptions.ts
index c78a3476..996b0591 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";
 
 /**
@@ -11,6 +12,14 @@ export type EmbeddingCallOptions = {
  * List of text values to generate embeddings for.
  */
 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.
  */
diff --git a/bindings/node/src/types/HttpExchange.ts b/bindings/node/src/types/HttpExchange.ts
index 2d87f56c..336f2176 100644
--- a/bindings/node/src/types/HttpExchange.ts
+++ b/bindings/node/src/types/HttpExchange.ts
@@ -8,9 +8,20 @@ import type { TimingRecord } from "./TimingRecord";
  */
 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 = 请求失败未获响应。
  */
diff --git a/bindings/node/src/types/ImageCallOptions.ts b/bindings/node/src/types/ImageCallOptions.ts
index d32de995..93cd3890 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";
 
 /**
@@ -45,6 +46,14 @@ mask: ImageFile | null,
  * Additional provider-specific options, keyed by provider name.
  */
 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..c843d810 100644
--- a/bindings/node/src/types/OutcomeRecord.ts
+++ b/bindings/node/src/types/OutcomeRecord.ts
@@ -6,6 +6,11 @@ import type { JsonValue } from "./serde_json/JsonValue";
  * 最终结果摘要。
  */
 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..8c41f7f5 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";
 
 /**
@@ -20,6 +21,14 @@ query: string,
  * Optional limit: return only the top `n` documents.
  */
 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.
  */
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..0a7152b3 100644
--- a/bindings/node/src/types/SearchCallOptions.ts
+++ b/bindings/node/src/types/SearchCallOptions.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";
 
 /**
@@ -33,6 +34,14 @@ include_domains: Array | null,
  * Optional list of domains to exclude from results.
  */
 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.
  */
diff --git a/bindings/node/src/types/SpeechCallOptions.ts b/bindings/node/src/types/SpeechCallOptions.ts
index 39cae2f4..0872b66b 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";
 
 /**
@@ -38,6 +39,14 @@ language: string | null,
  * 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/TimeoutConfiguration.ts b/bindings/node/src/types/TimeoutConfiguration.ts
index 8e3a3444..d547f0a0 100644
--- a/bindings/node/src/types/TimeoutConfiguration.ts
+++ b/bindings/node/src/types/TimeoutConfiguration.ts
@@ -14,8 +14,17 @@ export type TimeoutConfiguration = {
  * streaming, the whole stream), in milliseconds.
  */
 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, 
 /**
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..22bf461a 100644
--- a/bindings/node/src/types/VideoCallOptions.ts
+++ b/bindings/node/src/types/VideoCallOptions.ts
@@ -1,12 +1,14 @@
 // 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`.
  */
@@ -60,6 +62,19 @@ generate_audio: boolean | null,
  * Additional provider-specific options, keyed by provider name.
  */
 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..8a8a96bd 100644
--- a/bindings/node/src/types/VideoResult.ts
+++ b/bindings/node/src/types/VideoResult.ts
@@ -5,7 +5,7 @@ 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`.
  */
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..96768bc5 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,
@@ -117,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}"
         ))),
@@ -125,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.
@@ -203,15 +210,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 +235,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 +256,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 +311,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 +329,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..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
 // ─────────────────────────────────────────────────────────────────────────────
@@ -54,7 +102,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 +129,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 +156,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 +194,25 @@ 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 = 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;
             }
         }
 
-        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 +248,25 @@ 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 = parse_opts_json(
+                    "opts_json",
+                    s,
+                    &[
+                        ("query", serde_json::Value::from("")),
+                        ("documents", documents_placeholder()),
+                    ],
+                )?;
+                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 +291,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),
@@ -242,12 +320,21 @@ 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 = parse_opts_json(
+                    "opts_json",
+                    s,
+                    &[("query", serde_json::Value::from(""))],
+                )?;
+                parsed.query = opts.query;
                 opts = parsed;
             }
         }
 
-        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 +615,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 +641,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 +672,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;
@@ -772,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");
+    }
+}
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/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..3220e154
--- /dev/null
+++ b/docs/ai-sdk-request-pipeline.md
@@ -0,0 +1,1052 @@
+# 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，读取中途连接死亡保留已收到的部分 |
+| 成功 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） |
+| 默认超时 | 与 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 处停止 |
+| 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 的单次请求语义 |
+| 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（流式阶段的 deadline 由 §8.1 的 pump task 观察，它是唯一 spawn 的 task，随返回的 stream drop 而 abort）：
+
+```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。
+
+`do_stream` 返回后所有 deadline 都由 pump task 观察（见“stream 驱动”偏差行）：timer 只度量
+provider 输出的到达间隔，消费方晚 poll 的时间不计入；已到达但尚未被消费的 output 在 timeout
+之前原样交付，timeout 作为其后的 terminal `Err` 项出现。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。
+- **修订**（RFC-0031 review）：`options.n` 按 `VideoModel::max_videos_per_call`
+  切 batch，每个 batch 并发跑一次独立的 `do_start`/poll 流程并独立铸造
+  idempotency key（AI SDK 用 `Promise.all`，Rust 侧 `try_join_all` 会在首个
+  失败时 drop 其余 batch future——不影响正确性，两边都不会重连或取消已发出的
+  provider 任务）。调用方提供 key 时，单批保持原键，多批按批次索引派生不同的
+  键；同批重试和同一调用方 key 的重放保持键稳定。`n == 0` 在任何网络调用
+  之前返回 `InvalidArgument`。
+
+### 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`；本轮不为它虚构第二层。**修订**（RFC-0031 review）：`upload_file`
+自己用 Provider 已配置的 `RetryConfig` 重试（与 `execute_list_models` 同一
+primitive）。这是“创建资源类 POST 不自动重试”一般规则的显式例外，前提是
+upload 不计费、失败的 create-file 请求不会返回可复用的 file id。多阶段的
+Google Files 按 exchange 分别重试，而不是包住整个 `upload_file`。
+
+---
+
+## 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 的 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
+
+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.
 
 ---
 
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..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 {
@@ -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..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 {
@@ -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,
     };
     (