From e9882f63c7a4b7d8102ec1598c2b27d822e7cb8c Mon Sep 17 00:00:00 2001 From: YOMXXX Date: Fri, 29 May 2026 22:16:35 +0800 Subject: [PATCH] fix(embeddings): add 429 backoff to stop 31K Sentry event flood (#2898) (#2904) --- src/core/observability.rs | 44 ++++ src/openhuman/embeddings/cohere.rs | 268 +++++++++++++++++++---- src/openhuman/embeddings/mod.rs | 1 + src/openhuman/embeddings/openai.rs | 237 ++++++++++++-------- src/openhuman/embeddings/openai_tests.rs | 249 +++++++++++++++++++++ src/openhuman/embeddings/retry_after.rs | 135 ++++++++++++ 6 files changed, 796 insertions(+), 138 deletions(-) create mode 100644 src/openhuman/embeddings/retry_after.rs diff --git a/src/core/observability.rs b/src/core/observability.rs index edd02e4d1..66ab428b9 100644 --- a/src/core/observability.rs +++ b/src/core/observability.rs @@ -2202,6 +2202,50 @@ mod tests { ); } + /// Task B (issue #2898): prove the canonical 429 error message produced by + /// the embedding clients is already classified as `TransientUpstreamHttp` + /// so Sentry events are suppressed even without backoff. + /// + /// The `is_transient_upstream_http_message` matcher checks for + /// `"api error (429 "` (case-insensitive), which is present in both the + /// OpenAI and Cohere canonical error shapes. + #[test] + fn embedding_429_classifies_as_transient_upstream_http() { + // OpenAI/Voyage canonical shape (openai.rs emit site). + let msg = "Embedding API error (429 Too Many Requests): Rate limit exceeded."; + assert_eq!( + expected_error_kind(msg), + Some(ExpectedErrorKind::TransientUpstreamHttp), + "OpenAI 429 must classify as TransientUpstreamHttp: {msg}" + ); + + // Cohere canonical shape (cohere.rs emit site). + let cohere_msg = "Cohere embed API error (429 Too Many Requests): rate limit exceeded."; + assert_eq!( + expected_error_kind(cohere_msg), + Some(ExpectedErrorKind::TransientUpstreamHttp), + "Cohere 429 must classify as TransientUpstreamHttp: {cohere_msg}" + ); + + // After-cap bail shape from the retry loop (openai.rs). + let cap_msg = + "Embedding API error (429 Too Many Requests): rate limit exceeded after 3 retries"; + assert_eq!( + expected_error_kind(cap_msg), + Some(ExpectedErrorKind::TransientUpstreamHttp), + "retry-cap bail message must classify as TransientUpstreamHttp: {cap_msg}" + ); + + // After-cap bail shape from the retry loop (cohere.rs). + let cohere_cap_msg = + "Cohere embed API error (429 Too Many Requests): rate limit exceeded after 3 retries"; + assert_eq!( + expected_error_kind(cohere_cap_msg), + Some(ExpectedErrorKind::TransientUpstreamHttp), + "Cohere retry-cap bail message must classify as TransientUpstreamHttp: {cohere_cap_msg}" + ); + } + #[test] fn classifies_backend_env_api_key_not_configured() { // TAURI-RUST-2H5 (~5 K events): backend embedding endpoint returns a diff --git a/src/openhuman/embeddings/cohere.rs b/src/openhuman/embeddings/cohere.rs index 9621046be..fe5601b0f 100644 --- a/src/openhuman/embeddings/cohere.rs +++ b/src/openhuman/embeddings/cohere.rs @@ -8,6 +8,7 @@ use async_trait::async_trait; +use super::retry_after::{backoff_ms_for_attempt, MAX_429_RETRIES}; use super::EmbeddingProvider; pub const COHERE_API_BASE: &str = "https://api.cohere.com"; @@ -18,6 +19,7 @@ pub struct CohereEmbedding { api_key: String, model: String, dims: usize, + base_url: String, } impl CohereEmbedding { @@ -33,9 +35,20 @@ impl CohereEmbedding { api_key: api_key.to_string(), model, dims, + base_url: COHERE_API_BASE.to_string(), } } + /// Test-only base URL override. OpenAI's base URL is constructor-injected + /// since its `new()` already takes one; Cohere historically hardcoded + /// `COHERE_API_BASE`, so this builder fills the gap for the 429 backoff + /// tests. + #[cfg(test)] + pub(crate) fn with_base_url(mut self, base: impl Into) -> Self { + self.base_url = base.into(); + self + } + fn http_client(&self) -> reqwest::Client { crate::openhuman::config::build_runtime_proxy_client("embeddings.cohere") } @@ -65,14 +78,21 @@ impl EmbeddingProvider for CohereEmbedding { self.dims } + /// Sends a POST request to the Cohere embed API. + /// + /// On 429 (Too Many Requests) or 503 (Service Unavailable) the call is + /// retried up to `MAX_429_RETRIES` times with exponential backoff. When + /// the server supplies a `Retry-After` header its value (delta-seconds) is + /// preferred over the computed backoff. After all retries are exhausted the + /// canonical error message is returned so the `TransientUpstreamHttp` + /// classifier in `core::observability` demotes it to a warning breadcrumb + /// instead of a Sentry error event. async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { if texts.is_empty() { return Ok(Vec::new()); } - super::rate_limit::acquire_embedding_slot(COHERE_API_BASE).await; - - let url = format!("{COHERE_API_BASE}/v2/embed"); + let url = format!("{}/v2/embed", self.base_url); tracing::debug!( target: "embeddings.cohere", @@ -86,61 +106,114 @@ impl EmbeddingProvider for CohereEmbedding { "embedding_types": ["float"], }); - let resp = self - .http_client() - .post(&url) - .header("Content-Type", "application/json") - .header("Authorization", format!("Bearer {}", self.api_key)) - .json(&body) - .send() - .await?; + // Retry loop: handles 429 Too Many Requests and 503 Service Unavailable + // with Retry-After–aware exponential backoff. + for attempt in 0..=MAX_429_RETRIES { + // Proactively gate every outbound attempt (initial + retries) against + // the per-endpoint rate budget. The chokepoint must sit inside the + // loop: a single pre-loop acquire would let retried 429/503 attempts + // bypass token consumption and let concurrent callers blow past the + // cap, ironically triggering more 429s. Token consumption tracks the + // number of HTTP attempts (1 + retries actually executed). Loopback + // endpoints are exempt (see `rate_limit`). + super::rate_limit::acquire_embedding_slot(&self.base_url).await; + + let resp = self + .http_client() + .post(&url) + .header("Content-Type", "application/json") + .header("Authorization", format!("Bearer {}", self.api_key)) + .json(&body) + .send() + .await?; - if !resp.status().is_success() { let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); - let message = format!("Cohere embed API error ({status}): {text}"); - crate::core::observability::report_error_or_expected( - &message, - "embeddings", - "cohere_embed", - &[("model", self.model.as_str()), ("failure", "non_2xx")], - ); - anyhow::bail!(message); - } - let payload: CohereEmbedResponse = resp - .json() - .await - .map_err(|e| anyhow::anyhow!("Cohere embed response parse failed: {e}"))?; + // Retry on 429 and 503 — both can carry a Retry-After header. + let is_retryable = status.as_u16() == 429 || status.as_u16() == 503; - let embeddings = payload.embeddings.float; + if is_retryable && attempt < MAX_429_RETRIES { + // Read Retry-After before consuming the body. + let retry_after_val = resp + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_owned()); - if embeddings.len() != texts.len() { - anyhow::bail!( - "Cohere embed count mismatch: sent {} texts, got {} embeddings", - texts.len(), - embeddings.len() - ); - } + let body_text = resp.text().await.unwrap_or_default(); + tracing::debug!( + target: "embeddings.cohere", + "[embeddings] cohere {} body on retry: {body_text}", + status.as_u16() + ); - for (i, vec) in embeddings.iter().enumerate() { - if self.dims > 0 && vec.len() != self.dims { + let delay_ms = backoff_ms_for_attempt(attempt, retry_after_val.as_deref()); + + tracing::debug!( + target: "embeddings.cohere", + "[embeddings] cohere {}, retrying in {}ms (attempt {}/{})", + status.as_u16(), delay_ms, attempt + 1, MAX_429_RETRIES + ); + + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + continue; + } + + if !status.is_success() { + let text = resp.text().await.unwrap_or_default(); + let message = format!("Cohere embed API error ({status}): {text}"); + crate::core::observability::report_error_or_expected( + &message, + "embeddings", + "cohere_embed", + &[("model", self.model.as_str()), ("failure", "non_2xx")], + ); + anyhow::bail!(message); + } + + let payload: CohereEmbedResponse = resp + .json() + .await + .map_err(|e| anyhow::anyhow!("Cohere embed response parse failed: {e}"))?; + + let embeddings = payload.embeddings.float; + + if embeddings.len() != texts.len() { anyhow::bail!( - "Cohere embed dimension mismatch at index {i}: expected {}, got {}", - self.dims, - vec.len() + "Cohere embed count mismatch: sent {} texts, got {} embeddings", + texts.len(), + embeddings.len() ); } + + for (i, vec) in embeddings.iter().enumerate() { + if self.dims > 0 && vec.len() != self.dims { + anyhow::bail!( + "Cohere embed dimension mismatch at index {i}: expected {}, got {}", + self.dims, + vec.len() + ); + } + } + + tracing::debug!( + target: "embeddings.cohere", + "[cohere] embed success: model={}, count={}, dims={}", + self.model, embeddings.len(), + embeddings.first().map(|v| v.len()).unwrap_or(0) + ); + + return Ok(embeddings); } - tracing::debug!( - target: "embeddings.cohere", - "[cohere] embed success: model={}, count={}, dims={}", - self.model, embeddings.len(), - embeddings.first().map(|v| v.len()).unwrap_or(0) - ); - - Ok(embeddings) + // The loop always exits via `return Ok(...)`, `bail!(...)`, or + // `continue`; this point is structurally unreachable. On the final + // attempt (`attempt == MAX_429_RETRIES`) the retryable guard is false + // and execution falls into the non-2xx branch above, which bails with + // the body-bearing format "Cohere embed API error (429 ...): " — + // that format preserves the "(429 " substring required by the + // TransientUpstreamHttp classifier in core::observability. + unreachable!("cohere embed retry loop must exit via return or bail") } } @@ -176,4 +249,105 @@ mod tests { let p = CohereEmbedding::new("k", "", 0); assert!(p.embed(&[]).await.unwrap().is_empty()); } + + // ── 429 backoff tests ────────────────────────────────────── + + use axum::{http::StatusCode, routing::post, Router}; + use std::{ + net::SocketAddr, + sync::{Arc, Mutex}, + }; + + async fn start_mock(app: Router) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr: SocketAddr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + format!("http://127.0.0.1:{}", addr.port()) + } + + /// Cohere 429 then success — verifies retry recovers. + /// + /// The mock returns 429 (with Retry-After: 0 for zero real-wall-clock delay) + /// twice, then 200 on the third call. The real `CohereEmbedding::embed` is + /// driven via `with_base_url` pointing at the axum mock server. + #[tokio::test] + async fn cohere_embed_429_then_success() { + let counter = Arc::new(Mutex::new(0u32)); + let counter_clone = counter.clone(); + + let app = Router::new().route( + "/v2/embed", + post(move || { + let counter = counter_clone.clone(); + async move { + let mut n = counter.lock().unwrap(); + *n += 1; + if *n <= 2 { + axum::response::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .header("Retry-After", "0") + .body(axum::body::Body::from(r#"{"message":"rate limited"}"#)) + .unwrap() + } else { + axum::response::Response::builder() + .status(StatusCode::OK) + .header("Content-Type", "application/json") + .body(axum::body::Body::from( + r#"{"embeddings":{"float":[[0.1,0.2]]}}"#, + )) + .unwrap() + } + } + }), + ); + + let base_url = start_mock(app).await; + let p = CohereEmbedding::new("test-key", "embed-english-v3.0", 2).with_base_url(&base_url); + + let result = p.embed(&["hello"]).await.unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(*counter.lock().unwrap(), 3, "should have taken 3 requests"); + } + + /// Cohere 429 indefinitely — verify bail with canonical message after retry + /// cap, and that exactly `MAX_429_RETRIES + 1` requests were made. + #[tokio::test] + async fn cohere_embed_429_indefinite_bails_after_cap() { + let counter = Arc::new(Mutex::new(0u32)); + let counter_clone = counter.clone(); + + let app = Router::new().route( + "/v2/embed", + post(move || { + let counter = counter_clone.clone(); + async move { + let mut n = counter.lock().unwrap(); + *n += 1; + axum::response::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .header("Retry-After", "0") + .body(axum::body::Body::from(r#"{"message":"always limited"}"#)) + .unwrap() + } + }), + ); + + let base_url = start_mock(app).await; + let p = CohereEmbedding::new("test-key", "embed-english-v3.0", 2).with_base_url(&base_url); + + let err = p.embed(&["hello"]).await.unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("429"), + "should contain 429 in error message: {msg}" + ); + // MAX_429_RETRIES retries + 1 initial = MAX_429_RETRIES + 1 total requests + assert_eq!( + *counter.lock().unwrap(), + MAX_429_RETRIES + 1, + "should make exactly MAX_429_RETRIES+1 requests" + ); + } } diff --git a/src/openhuman/embeddings/mod.rs b/src/openhuman/embeddings/mod.rs index 1f3e5cd89..b990490dd 100644 --- a/src/openhuman/embeddings/mod.rs +++ b/src/openhuman/embeddings/mod.rs @@ -21,6 +21,7 @@ pub mod ollama; pub mod openai; mod provider_trait; pub mod rate_limit; +pub mod retry_after; mod rpc; mod schemas; pub mod voyage; diff --git a/src/openhuman/embeddings/openai.rs b/src/openhuman/embeddings/openai.rs index 64a98d7e7..0c85acff4 100644 --- a/src/openhuman/embeddings/openai.rs +++ b/src/openhuman/embeddings/openai.rs @@ -5,6 +5,7 @@ use async_trait::async_trait; +use super::retry_after::{backoff_ms_for_attempt, MAX_429_RETRIES}; use super::EmbeddingProvider; /// Embedding provider for OpenAI and compatible APIs (e.g., LocalAI, Ollama). @@ -89,19 +90,19 @@ impl EmbeddingProvider for OpenAiEmbedding { } /// Sends a POST request to the embedding API. + /// + /// On 429 (Too Many Requests) or 503 (Service Unavailable) the call is + /// retried up to `MAX_429_RETRIES` times with exponential backoff. When + /// the server supplies a `Retry-After` header its value (delta-seconds) is + /// preferred over the computed backoff. After all retries are exhausted the + /// canonical error message is returned so the `TransientUpstreamHttp` + /// classifier in `core::observability` demotes it to a warning breadcrumb + /// instead of a Sentry error event. async fn embed(&self, texts: &[&str]) -> anyhow::Result>> { if texts.is_empty() { return Ok(Vec::new()); } - // Proactively gate the outbound request against the per-endpoint rate - // budget so cloud backends (OpenHuman/Voyage, OpenAI, custom remote - // endpoints) stay under their account quota instead of tripping 429s. - // This is the single chokepoint every cloud embed funnels through — - // the `cloud` provider delegates here, and `openai`/`custom:` use it - // directly. Loopback endpoints are exempt (see `rate_limit`). - super::rate_limit::acquire_embedding_slot(&self.base_url).await; - let url = self.embeddings_url(); tracing::debug!( @@ -115,98 +116,152 @@ impl EmbeddingProvider for OpenAiEmbedding { "input": texts, }); - let mut req = self - .http_client() - .post(&url) - .header("Content-Type", "application/json") - .json(&body); + // Retry loop: handles 429 Too Many Requests and 503 Service Unavailable + // with Retry-After–aware exponential backoff. + for attempt in 0..=MAX_429_RETRIES { + let mut req = self + .http_client() + .post(&url) + .header("Content-Type", "application/json") + .json(&body); - // Only set Authorization header when an API key is configured. - if !self.api_key.is_empty() { - req = req.header("Authorization", format!("Bearer {}", self.api_key)); - } - - let resp = req.send().await?; - - if !resp.status().is_success() { - let status = resp.status(); - let status_str = status.as_u16().to_string(); - let text = resp.text().await.unwrap_or_default(); - tracing::debug!( - target: "openai::embed", - "[openai] embed error: status={status}, body={text}" - ); - let message = format!("Embedding API error ({status}): {text}"); - // Use `report_error_or_expected` so transient upstream HTTP failures - // (e.g. 429 Too Many Requests, which the memory_tree job runner - // already retries with backoff) log a warning breadcrumb instead of - // firing a Sentry error event per attempt. - crate::core::observability::report_error_or_expected( - message.as_str(), - "embeddings", - "openai_embed", - &[ - ("model", self.model.as_str()), - ("status", status_str.as_str()), - ("failure", "non_2xx"), - ], - ); - anyhow::bail!(message); - } - - let json: serde_json::Value = resp.json().await?; - let data = json - .get("data") - .and_then(|d| d.as_array()) - .ok_or_else(|| anyhow::anyhow!("Invalid embedding response: missing 'data'"))?; - - // Validate that the response count matches the input count. - if data.len() != texts.len() { - anyhow::bail!( - "openai embed count mismatch: sent {} texts, got {} items in 'data'", - texts.len(), - data.len() - ); - } - - let mut embeddings = Vec::with_capacity(data.len()); - for (i, item) in data.iter().enumerate() { - let embedding = item - .get("embedding") - .and_then(|e| e.as_array()) - .ok_or_else(|| { - anyhow::anyhow!("Invalid embedding item at index {i}: missing 'embedding'") - })?; - - let mut vec = Vec::with_capacity(embedding.len()); - for (j, v) in embedding.iter().enumerate() { - #[allow(clippy::cast_possible_truncation)] - let f = v.as_f64().ok_or_else(|| { - anyhow::anyhow!("non-numeric value at data[{i}].embedding[{j}]: {v}") - })? as f32; - vec.push(f); + // Only set Authorization header when an API key is configured. + if !self.api_key.is_empty() { + req = req.header("Authorization", format!("Bearer {}", self.api_key)); } - // Validate dimensions. - if self.dims > 0 && vec.len() != self.dims { + // Proactively gate every outbound attempt (initial + retries) against + // the per-endpoint rate budget so cloud backends (OpenHuman/Voyage, + // OpenAI, custom remote endpoints) stay under their account quota + // instead of tripping 429s. The chokepoint must sit inside the loop: + // a single pre-loop acquire would let retried 429/503 attempts bypass + // token consumption and let concurrent callers blow past the cap, + // ironically triggering more 429s. Token consumption tracks the number + // of HTTP attempts (1 + retries actually executed). Loopback endpoints + // are exempt (see `rate_limit`). + super::rate_limit::acquire_embedding_slot(&self.base_url).await; + + let resp = req.send().await?; + + let status = resp.status(); + + // Retry on 429 and 503 — both can carry a Retry-After header. + let is_retryable = status.as_u16() == 429 || status.as_u16() == 503; + + if is_retryable && attempt < MAX_429_RETRIES { + // Read Retry-After before consuming the body. + let retry_after_val = resp + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_owned()); + + let body_text = resp.text().await.unwrap_or_default(); + tracing::debug!( + target: "openai::embed", + "[embeddings] openai {} body on retry: {body_text}", + status.as_u16() + ); + + let delay_ms = backoff_ms_for_attempt(attempt, retry_after_val.as_deref()); + + tracing::debug!( + target: "openai::embed", + "[embeddings] openai {}, retrying in {}ms (attempt {}/{})", + status.as_u16(), delay_ms, attempt + 1, MAX_429_RETRIES + ); + + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + continue; + } + + if !status.is_success() { + let status_str = status.as_u16().to_string(); + let text = resp.text().await.unwrap_or_default(); + tracing::debug!( + target: "openai::embed", + "[openai] embed error: status={status}, body={text}" + ); + let message = format!("Embedding API error ({status}): {text}"); + // Use `report_error_or_expected` so transient upstream HTTP + // failures (e.g. 429 Too Many Requests after retry cap) log a + // warning breadcrumb instead of firing a Sentry error event. + crate::core::observability::report_error_or_expected( + message.as_str(), + "embeddings", + "openai_embed", + &[ + ("model", self.model.as_str()), + ("status", status_str.as_str()), + ("failure", "non_2xx"), + ], + ); + anyhow::bail!(message); + } + + let json: serde_json::Value = resp.json().await?; + let data = json + .get("data") + .and_then(|d| d.as_array()) + .ok_or_else(|| anyhow::anyhow!("Invalid embedding response: missing 'data'"))?; + + // Validate that the response count matches the input count. + if data.len() != texts.len() { anyhow::bail!( - "openai embed dimension mismatch at index {i}: expected {}, got {}", - self.dims, - vec.len() + "openai embed count mismatch: sent {} texts, got {} items in 'data'", + texts.len(), + data.len() ); } - embeddings.push(vec); + let mut embeddings = Vec::with_capacity(data.len()); + for (i, item) in data.iter().enumerate() { + let embedding = item + .get("embedding") + .and_then(|e| e.as_array()) + .ok_or_else(|| { + anyhow::anyhow!("Invalid embedding item at index {i}: missing 'embedding'") + })?; + + let mut vec = Vec::with_capacity(embedding.len()); + for (j, v) in embedding.iter().enumerate() { + #[allow(clippy::cast_possible_truncation)] + let f = v.as_f64().ok_or_else(|| { + anyhow::anyhow!("non-numeric value at data[{i}].embedding[{j}]: {v}") + })? as f32; + vec.push(f); + } + + // Validate dimensions. + if self.dims > 0 && vec.len() != self.dims { + anyhow::bail!( + "openai embed dimension mismatch at index {i}: expected {}, got {}", + self.dims, + vec.len() + ); + } + + embeddings.push(vec); + } + + tracing::debug!( + target: "openai::embed", + "[openai] embed success: model={}, count={}, dims={}", + self.model, embeddings.len(), + embeddings.first().map(|v| v.len()).unwrap_or(0) + ); + + return Ok(embeddings); } - tracing::debug!( - target: "openai::embed", - "[openai] embed success: model={}, count={}, dims={}", - self.model, embeddings.len(), - embeddings.first().map(|v| v.len()).unwrap_or(0) - ); - - Ok(embeddings) + // The loop always exits via `return Ok(...)`, `bail!(...)`, or + // `continue`; this point is structurally unreachable. On the final + // attempt (`attempt == MAX_429_RETRIES`) the retryable guard is false + // and execution falls into the non-2xx branch above, which bails with + // the body-bearing format "Embedding API error (429 ...): " — + // that format preserves the "(429 " substring required by the + // TransientUpstreamHttp classifier in core::observability. + unreachable!("embed retry loop must exit via return or bail") } } diff --git a/src/openhuman/embeddings/openai_tests.rs b/src/openhuman/embeddings/openai_tests.rs index 41d633a79..f77a28eb9 100644 --- a/src/openhuman/embeddings/openai_tests.rs +++ b/src/openhuman/embeddings/openai_tests.rs @@ -439,3 +439,252 @@ async fn embed_with_explicit_api_path() { let result = p.embed(&["test"]).await.unwrap(); assert_eq!(result.len(), 1); } + +// ── 429 backoff / Retry-After tests ────────────────────── + +/// Mock returns 429 twice (with Retry-After: 0) then 200 — verify embed +/// succeeds and that exactly 3 requests were made (initial + 2 retries). +#[tokio::test] +async fn embed_429_retries_with_retry_after_then_succeeds() { + use std::sync::{Arc, Mutex}; + let counter = Arc::new(Mutex::new(0u32)); + let counter_clone = counter.clone(); + + let app = Router::new().route( + "/v1/embeddings", + post(move || { + let counter = counter_clone.clone(); + async move { + let mut n = counter.lock().unwrap(); + *n += 1; + if *n <= 2 { + // Return 429 with Retry-After: 0 (zero delay) for fast tests. + axum::response::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .header("Retry-After", "0") + .body(axum::body::Body::from( + r#"{"error":{"message":"rate limited"}}"#, + )) + .unwrap() + } else { + axum::response::Response::builder() + .status(StatusCode::OK) + .header("Content-Type", "application/json") + .body(axum::body::Body::from( + r#"{"data":[{"embedding":[1.0,2.0]}]}"#, + )) + .unwrap() + } + } + }), + ); + let url = start_mock(app).await; + let p = OpenAiEmbedding::new(&url, "k", "m", 2); + + let result = p.embed(&["hello"]).await.unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0], vec![1.0_f32, 2.0]); + assert_eq!( + *counter.lock().unwrap(), + 3, + "expected 3 total requests (1 initial + 2 retries)" + ); +} + +/// Mock returns 429 indefinitely — verify bail with canonical message and +/// that exactly MAX_429_RETRIES + 1 requests were made. +#[tokio::test] +async fn embed_429_indefinite_bails_after_retry_cap() { + use crate::openhuman::embeddings::retry_after::MAX_429_RETRIES; + use std::sync::{Arc, Mutex}; + let counter = Arc::new(Mutex::new(0u32)); + let counter_clone = counter.clone(); + + let app = Router::new().route( + "/v1/embeddings", + post(move || { + let counter = counter_clone.clone(); + async move { + let mut n = counter.lock().unwrap(); + *n += 1; + axum::response::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .header("Retry-After", "0") + .body(axum::body::Body::from( + r#"{"error":{"message":"always rate limited"}}"#, + )) + .unwrap() + } + }), + ); + let url = start_mock(app).await; + let p = OpenAiEmbedding::new(&url, "k", "m", 1); + + let err = p.embed(&["hi"]).await.unwrap_err(); + let msg = err.to_string(); + + // The error message must contain "429" so the classifier still suppresses it. + assert!( + msg.contains("429"), + "bail message must contain 429 for Sentry classifier: {msg}" + ); + // Must match is_transient_upstream_http_message via "(429 " substring. + assert!( + crate::core::observability::is_transient_message_failure(&msg), + "bail message must classify as transient: {msg}" + ); + + let requests = *counter.lock().unwrap(); + assert_eq!( + requests, + MAX_429_RETRIES + 1, + "expected exactly MAX_429_RETRIES+1={} requests, got {}", + MAX_429_RETRIES + 1, + requests + ); +} + +/// Mock returns 429 without Retry-After header — verify the no-header code +/// path is taken, that a retry is attempted, and that the request succeeds. +/// Uses `Retry-After: 0` on the *second* attempt to confirm header parsing +/// is independent from the no-header first attempt. +#[tokio::test] +async fn embed_429_without_retry_after_uses_exponential_backoff() { + use crate::openhuman::embeddings::retry_after::{backoff_ms_for_attempt, BASE_BACKOFF_MS}; + use std::sync::{Arc, Mutex}; + + // Confirm the helper returns the exponential base when header is absent. + assert_eq!( + backoff_ms_for_attempt(0, None), + BASE_BACKOFF_MS, + "attempt 0 without header should use BASE_BACKOFF_MS" + ); + assert_eq!( + backoff_ms_for_attempt(1, None), + BASE_BACKOFF_MS * 2, + "attempt 1 without header should double" + ); + + // Confirm header path: Retry-After: 0 overrides the exponential base. + assert_eq!( + backoff_ms_for_attempt(0, Some("0")), + 0, + "Retry-After: 0 should yield zero-ms delay" + ); + + // End-to-end: first request returns 429 (no Retry-After), second succeeds. + // Use Retry-After: 0 on the 429 to avoid real-wall-clock delay in CI. + let counter = Arc::new(Mutex::new(0u32)); + let counter_clone = counter.clone(); + + let app = Router::new().route( + "/v1/embeddings", + post(move || { + let counter = counter_clone.clone(); + async move { + let mut n = counter.lock().unwrap(); + *n += 1; + if *n == 1 { + // Mock still sets Retry-After: 0 here to keep this test fast; + // the no-header exponential branch is exercised end-to-end by + // `embed_429_no_retry_after_header_falls_back_to_exponential` + // below and by unit tests in `retry_after.rs`. + axum::response::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .header("Retry-After", "0") + .body(axum::body::Body::from( + r#"{"error":{"message":"rate limited"}}"#, + )) + .unwrap() + } else { + axum::response::Response::builder() + .status(StatusCode::OK) + .header("Content-Type", "application/json") + .body(axum::body::Body::from(r#"{"data":[{"embedding":[9.9]}]}"#)) + .unwrap() + } + } + }), + ); + let url = start_mock(app).await; + let p = OpenAiEmbedding::new(&url, "k", "m", 1); + + let result = p.embed(&["hi"]).await; + assert!(result.is_ok(), "should succeed after retry: {:?}", result); + assert_eq!( + *counter.lock().unwrap(), + 2, + "expected 2 total requests (1 initial + 1 retry)" + ); +} + +/// End-to-end coverage for the no-`Retry-After` exponential-backoff branch. +/// +/// The previous "no Retry-After" test still set `Retry-After: 0` on the mock, +/// so the embed loop never actually called `backoff_ms_for_attempt(.., None)` — +/// the exponential fallback was only exercised by unit tests in `retry_after.rs`. +/// This test omits the header entirely so the loop must use the `BASE_BACKOFF_MS` +/// path. We tolerate the ~1 s wait (one retry @ `BASE_BACKOFF_MS = 1000 ms`) to +/// keep the assertion meaningful: the real backoff schedule is what we care about. +#[tokio::test] +async fn embed_429_no_retry_after_header_falls_back_to_exponential() { + use crate::openhuman::embeddings::retry_after::BASE_BACKOFF_MS; + use std::sync::{Arc, Mutex}; + use tokio::time::Instant; + + let counter = Arc::new(Mutex::new(0u32)); + let counter_clone = counter.clone(); + + let app = Router::new().route( + "/v1/embeddings", + post(move || { + let counter = counter_clone.clone(); + async move { + let mut n = counter.lock().unwrap(); + *n += 1; + if *n == 1 { + // Crucial: omit the Retry-After header entirely so the embed + // loop hits `backoff_ms_for_attempt(0, None)` = + // `BASE_BACKOFF_MS` and sleeps. If a future refactor breaks + // the fallback (e.g. defaulting to 0 ms) the elapsed-time + // assertion below will catch it. + axum::response::Response::builder() + .status(StatusCode::TOO_MANY_REQUESTS) + .body(axum::body::Body::from( + r#"{"error":{"message":"rate limited"}}"#, + )) + .unwrap() + } else { + axum::response::Response::builder() + .status(StatusCode::OK) + .header("Content-Type", "application/json") + .body(axum::body::Body::from(r#"{"data":[{"embedding":[4.2]}]}"#)) + .unwrap() + } + } + }), + ); + let url = start_mock(app).await; + let p = OpenAiEmbedding::new(&url, "k", "m", 1); + + let start = Instant::now(); + let result = p.embed(&["hi"]).await; + let elapsed = start.elapsed(); + + assert!(result.is_ok(), "should succeed after retry: {:?}", result); + assert_eq!( + *counter.lock().unwrap(), + 2, + "expected 2 total requests (1 initial + 1 retry)" + ); + // The fallback must actually wait the exponential base — within a 250 ms + // jitter window for slow CI runners. Without this we couldn't tell whether + // the no-header branch was taken or silently short-circuited. + let min_wait = std::time::Duration::from_millis(BASE_BACKOFF_MS.saturating_sub(250)); + assert!( + elapsed >= min_wait, + "expected elapsed >= ~{}ms (BASE_BACKOFF_MS minus jitter), got {:?}", + min_wait.as_millis(), + elapsed + ); +} diff --git a/src/openhuman/embeddings/retry_after.rs b/src/openhuman/embeddings/retry_after.rs new file mode 100644 index 000000000..b799f77a6 --- /dev/null +++ b/src/openhuman/embeddings/retry_after.rs @@ -0,0 +1,135 @@ +//! Retry-After header parsing and 429/503 backoff logic for embedding clients. +//! +//! The HTTP `Retry-After` header may arrive as either: +//! - A non-negative integer: delta-seconds from now (e.g. `Retry-After: 30`) +//! - An HTTP-date string: absolute point in time (e.g. `Retry-After: Wed, 21 Oct 2015 07:28:00 GMT`) +//! +//! This module prefers the delta-seconds form and falls back to exponential +//! backoff when the header is absent or unparseable. See RFC 9110 §10.2.4. + +/// Maximum number of 429/503 retries before giving up. +/// +/// Three retries means the client makes up to four total attempts: the +/// original request plus three retries. This caps the per-call delay at +/// roughly `BASE_BACKOFF_MS * 2^2 = 4 s` in the no-`Retry-After` path (or +/// the server-directed duration in the header path). +pub const MAX_429_RETRIES: u32 = 3; + +/// Base exponential-backoff delay in milliseconds when `Retry-After` is absent. +/// +/// Sequence (per attempt): 1 s, 2 s, 4 s — capped at 30 s. +pub const BASE_BACKOFF_MS: u64 = 1_000; + +/// Maximum backoff delay in milliseconds regardless of `Retry-After` or +/// computed exponent. Prevents a misbehaving server from parking the caller +/// for an unreasonably long time. +pub const MAX_BACKOFF_MS: u64 = 30_000; + +/// Parse the `Retry-After` header value into a delay in milliseconds. +/// +/// Accepts the delta-seconds form only (e.g. `"30"`, `"0"`). HTTP-date form +/// is not parsed — fall back to exponential backoff when the value is not a +/// non-negative integer. +/// +/// Returns `None` when the header is absent, empty, or not a valid +/// non-negative integer. +pub fn parse_retry_after_ms(header_value: Option<&str>) -> Option { + let s = header_value?.trim(); + // Only accept the delta-seconds form: a non-negative integer. + let secs: u64 = s.parse().ok()?; + Some(secs.saturating_mul(1_000).min(MAX_BACKOFF_MS)) +} + +/// Compute the delay for attempt `n` (0-indexed) with optional `Retry-After` +/// override. Uses exponential backoff as fallback: +/// `BASE_BACKOFF_MS * 2^n`, capped at `MAX_BACKOFF_MS`. +pub fn backoff_ms_for_attempt(attempt: u32, retry_after_header: Option<&str>) -> u64 { + if let Some(ms) = parse_retry_after_ms(retry_after_header) { + return ms; + } + BASE_BACKOFF_MS + .saturating_mul(2u64.saturating_pow(attempt)) + .min(MAX_BACKOFF_MS) +} + +#[cfg(test)] +mod tests { + use super::*; + + // ── parse_retry_after_ms ──────────────────────────────────── + + #[test] + fn parses_integer_seconds() { + assert_eq!(parse_retry_after_ms(Some("30")), Some(30_000)); + } + + #[test] + fn parses_zero() { + assert_eq!(parse_retry_after_ms(Some("0")), Some(0)); + } + + #[test] + fn parses_whitespace_padded() { + assert_eq!(parse_retry_after_ms(Some(" 5 ")), Some(5_000)); + } + + #[test] + fn caps_at_max_backoff() { + // 9999 s * 1000 ms/s > MAX_BACKOFF_MS + assert_eq!(parse_retry_after_ms(Some("9999")), Some(MAX_BACKOFF_MS)); + } + + #[test] + fn returns_none_for_http_date() { + // HTTP-date form — not parsed; fall through to exponential backoff. + assert_eq!( + parse_retry_after_ms(Some("Wed, 21 Oct 2015 07:28:00 GMT")), + None + ); + } + + #[test] + fn returns_none_for_none_input() { + assert_eq!(parse_retry_after_ms(None), None); + } + + #[test] + fn returns_none_for_empty_string() { + assert_eq!(parse_retry_after_ms(Some("")), None); + } + + #[test] + fn returns_none_for_negative() { + // Negative integers are not valid delta-seconds per RFC 9110. + assert_eq!(parse_retry_after_ms(Some("-1")), None); + } + + // ── backoff_ms_for_attempt ───────────────────────────────── + + #[test] + fn uses_retry_after_when_present() { + assert_eq!(backoff_ms_for_attempt(0, Some("5")), 5_000); + assert_eq!(backoff_ms_for_attempt(2, Some("5")), 5_000); + } + + #[test] + fn falls_back_to_exponential_when_header_absent() { + assert_eq!(backoff_ms_for_attempt(0, None), BASE_BACKOFF_MS); + assert_eq!(backoff_ms_for_attempt(1, None), BASE_BACKOFF_MS * 2); + assert_eq!(backoff_ms_for_attempt(2, None), BASE_BACKOFF_MS * 4); + } + + #[test] + fn falls_back_to_exponential_when_header_unparseable() { + assert_eq!( + backoff_ms_for_attempt(0, Some("not-a-number")), + BASE_BACKOFF_MS + ); + } + + #[test] + fn exponential_caps_at_max() { + // 2^10 = 1024 — well past the cap + assert_eq!(backoff_ms_for_attempt(10, None), MAX_BACKOFF_MS); + } +}