admin/openhuman
1
0
Fork 0
mirror of https://github.com/tinyhumansai/openhuman.git synced 2026-07-27 21:08:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(embeddings): add 429 backoff to stop 31K Sentry event flood (#2898) (#2904)

Browse Source
This commit is contained in:
YOMXXX
2026-05-29 19:46:35 +05:30
committed by GitHub
parent c52f75d258
commit e9882f63c7
6 changed files with 796 additions and 138 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/core/observability.rs
+44
View File
@@ -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
src/openhuman/embeddings/cohere.rs
+221 -47
View File
@@ -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<String>) -> 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<Vec<Vec<f32>>> {
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 ...): <body>" —
// 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"
);
}
}
src/openhuman/embeddings/mod.rs
+1
View File
@@ -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;
src/openhuman/embeddings/openai.rs
+146 -91
View File
@@ -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<Vec<Vec<f32>>> {
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 ...): <body>" —
// that format preserves the "(429 " substring required by the
// TransientUpstreamHttp classifier in core::observability.
unreachable!("embed retry loop must exit via return or bail")
}
}
src/openhuman/embeddings/openai_tests.rs
+249
View File
@@ -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
);
}
src/openhuman/embeddings/retry_after.rs
+135
View File
@@ -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<u64> {
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);
}
}