From c856e4306f7ebd6fb4f8769379fef0238530bb6d Mon Sep 17 00:00:00 2001 From: Zyy Date: Mon, 27 Jul 2026 14:55:30 +0800 Subject: [PATCH] fix(xai): mark rate limits as quota exhausted Classify xAI 429 rate-limit responses as quota exhaustion so the scheduler removes the account from active rotation until the normal recovery flow makes it available again. Cover the provider replay path and transport classification with tests. Co-authored-by: Codex --- backend/crates/providers/xai/src/provider.rs | 1 + .../providers/xai/src/transport/network.rs | 39 +++++++++++++------ .../providers/xai/tests/provider/contract.rs | 17 ++++++++ .../providers/xai/tests/transport/network.rs | 31 +++++++++++++++ 4 files changed, 76 insertions(+), 12 deletions(-) diff --git a/backend/crates/providers/xai/src/provider.rs b/backend/crates/providers/xai/src/provider.rs index a09266839..b74d6dc5a 100644 --- a/backend/crates/providers/xai/src/provider.rs +++ b/backend/crates/providers/xai/src/provider.rs @@ -1714,6 +1714,7 @@ fn explicit_rejection_is_replay_safe(kind: ProviderErrorKind, status: u16) -> bo (ProviderErrorKind::Unauthorized, 401) | (ProviderErrorKind::QuotaExhausted, 402) | (ProviderErrorKind::QuotaExhausted, 403) + | (ProviderErrorKind::QuotaExhausted, 429) | (ProviderErrorKind::RateLimited, 429) ) } diff --git a/backend/crates/providers/xai/src/transport/network.rs b/backend/crates/providers/xai/src/transport/network.rs index a91cf2b4d..8ef0f25f1 100644 --- a/backend/crates/providers/xai/src/transport/network.rs +++ b/backend/crates/providers/xai/src/transport/network.rs @@ -1003,7 +1003,13 @@ async fn classify_inference_status(response: Response) -> GrokInferenceTransport StatusCode::REQUEST_TIMEOUT | StatusCode::GATEWAY_TIMEOUT => { GrokInferenceTransportErrorKind::Timeout } - StatusCode::TOO_MANY_REQUESTS => GrokInferenceTransportErrorKind::RateLimited, + StatusCode::TOO_MANY_REQUESTS => { + if quota_exhausted(&metadata) { + GrokInferenceTransportErrorKind::QuotaExhausted + } else { + GrokInferenceTransportErrorKind::RateLimited + } + } status if status.is_server_error() => GrokInferenceTransportErrorKind::Unavailable, _ => GrokInferenceTransportErrorKind::Protocol, }; @@ -1070,17 +1076,10 @@ fn first_string(object: &serde_json::Map, fields: &[&str]) -> Opt } fn classify_forbidden(metadata: &InferenceErrorMetadata) -> GrokInferenceTransportErrorKind { - let text = forbidden_metadata_text(metadata); + let text = inference_metadata_text(metadata); if credential_rejected(&text) { GrokInferenceTransportErrorKind::Unauthorized - } else if contains_any( - &text, - &[ - "subscription:free-usage-exhausted", - "used all the included free usage for model", - "personal-team-blocked:spending-limit", - ], - ) { + } else if quota_exhausted(metadata) { GrokInferenceTransportErrorKind::QuotaExhausted } else { GrokInferenceTransportErrorKind::PermissionDenied @@ -1088,7 +1087,7 @@ fn classify_forbidden(metadata: &InferenceErrorMetadata) -> GrokInferenceTranspo } fn forbidden_requires_recovery(metadata: &InferenceErrorMetadata) -> bool { - let text = forbidden_metadata_text(metadata); + let text = inference_metadata_text(metadata); credential_rejected(&text) || text.contains("access to the chat endpoint is denied") || text.trim_matches([' ', '.', '!', '\t', '\r', '\n']) == "access denied" @@ -1101,7 +1100,23 @@ fn reasoning_decode_failed(metadata: &InferenceErrorMetadata) -> bool { }) } -fn forbidden_metadata_text(metadata: &InferenceErrorMetadata) -> String { +fn quota_exhausted(metadata: &InferenceErrorMetadata) -> bool { + contains_any( + &inference_metadata_text(metadata), + &[ + "subscription:free-usage-exhausted", + "subscription_free_usage_exhausted", + "free-usage-exhausted", + "free_usage_exhausted", + "used all the included free usage", + "used all your free usage", + "personal-team-blocked:spending-limit", + "personal_team_blocked_spending_limit", + ], + ) +} + +fn inference_metadata_text(metadata: &InferenceErrorMetadata) -> String { [ metadata.code.as_deref().unwrap_or_default(), metadata.error_type.as_deref().unwrap_or_default(), diff --git a/backend/crates/providers/xai/tests/provider/contract.rs b/backend/crates/providers/xai/tests/provider/contract.rs index 18df6126b..521bc0ee3 100644 --- a/backend/crates/providers/xai/tests/provider/contract.rs +++ b/backend/crates/providers/xai/tests/provider/contract.rs @@ -1208,6 +1208,23 @@ async fn explicit_http_429_marks_provider_error_replay_safe() { assert!(error.replay_is_safe()); } +#[tokio::test] +async fn explicit_quota_http_429_marks_provider_error_replay_safe() { + let error = mapped_transport_error( + GrokInferenceTransportError::new( + GrokInferenceTransportErrorKind::QuotaExhausted, + UpstreamSendState::Sent, + ) + .with_status(429), + false, + ) + .await; + + assert_eq!(error.kind(), ProviderErrorKind::QuotaExhausted); + assert_eq!(error.upstream_status(), Some(429)); + assert!(error.replay_is_safe()); +} + #[tokio::test] async fn explicit_http_408_does_not_mark_provider_error_replay_safe() { let error = mapped_transport_error( diff --git a/backend/crates/providers/xai/tests/transport/network.rs b/backend/crates/providers/xai/tests/transport/network.rs index a0cabbc53..a3403db69 100644 --- a/backend/crates/providers/xai/tests/transport/network.rs +++ b/backend/crates/providers/xai/tests/transport/network.rs @@ -358,6 +358,37 @@ async fn inference_transport_should_scope_forbidden_failures_from_safe_metadata( } } +#[tokio::test] +async fn inference_transport_should_classify_free_usage_429_as_quota_exhausted() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(ResponseTemplate::new(429).set_body_json(json!({ + "error": { + "code": "subscription:free-usage-exhausted", + "message": "You have used all your free usage" + } + }))) + .expect(1) + .mount(&server) + .await; + let origin = Url::parse(&server.uri()).expect("wiremock origin"); + let error = inference_transport(&origin) + .execute(inference_request(&origin)) + .await + .expect_err("free usage exhaustion must fail"); + + assert_eq!( + error.kind(), + GrokInferenceTransportErrorKind::QuotaExhausted + ); + assert_eq!(error.status(), Some(429)); + assert_eq!( + error.upstream_code().map(|code| code.as_str()), + Some("subscription_free_usage_exhausted") + ); +} + #[tokio::test] async fn inference_transport_should_bound_retry_after_to_the_safe_window() { for (header, expected) in [