From 41bb94533b7edfddd58a76f761544f72fd46426d Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Wed, 13 May 2026 16:22:27 +0800 Subject: [PATCH 1/4] feat(proxy, admin): add runtime model status, cooldown, background health check, and /admin/v1/models/status - Add ModelRuntimeStatusTracker with cooldown/unhealthy/healthy states - Set cooldown on request-path retryable failures (30s TTL) - Add background_model_check config (direct-model only) for periodic health probes - Mark unhealthy on background check failure, ignore configured transient statuses (408/429) - Filter routing targets by runtime status (skip cooldown/unhealthy unless all would be excluded) - Add GET /admin/v1/models/status returning per-model runtime status with stale awareness - Add BackgroundModelCheck schema (direct-model only, routing models rejected) - Wire background checker in aisix-server with configurable interval --- crates/aisix-admin/src/lib.rs | 91 ++++++ .../aisix-admin/src/models_status_handler.rs | 72 +++++ crates/aisix-admin/src/openapi.rs | 72 ++++- crates/aisix-admin/src/state.rs | 12 +- crates/aisix-core/src/models/mod.rs | 2 +- crates/aisix-core/src/models/model.rs | 42 +++ crates/aisix-core/src/models/schema.rs | 86 +++++- crates/aisix-proxy/src/background.rs | 177 +++++++++++ crates/aisix-proxy/src/chat.rs | 79 ++++- crates/aisix-proxy/src/health.rs | 252 ++++++++++++++++ crates/aisix-proxy/src/lib.rs | 274 ++++++++++++++++-- crates/aisix-proxy/src/routing.rs | 55 ++-- crates/aisix-proxy/src/state.rs | 9 +- crates/aisix-server/src/main.rs | 51 +++- tests/e2e/src/harness/admin.ts | 4 + 15 files changed, 1198 insertions(+), 80 deletions(-) create mode 100644 crates/aisix-admin/src/models_status_handler.rs create mode 100644 crates/aisix-proxy/src/background.rs diff --git a/crates/aisix-admin/src/lib.rs b/crates/aisix-admin/src/lib.rs index a7f1f714..50e7f98f 100644 --- a/crates/aisix-admin/src/lib.rs +++ b/crates/aisix-admin/src/lib.rs @@ -39,6 +39,7 @@ pub mod etcd_store; mod guardrails_handlers; mod health_handler; mod models_handlers; +mod models_status_handler; mod observability_exporters_handlers; mod openapi; mod playground_handler; @@ -73,6 +74,10 @@ pub fn build_router(state: AdminState) -> Router { .put(models_handlers::update_model) .delete(models_handlers::delete_model), ) + .route( + "/admin/v1/models/status", + get(models_status_handler::get_models_status), + ) .route( "/admin/v1/apikeys", get(apikeys_handlers::list_apikeys).post(apikeys_handlers::create_apikey), @@ -1089,4 +1094,90 @@ mod tests { // Empty snapshot → empty model list, but endpoint is operational. assert!(v["models"].is_array()); } + + #[tokio::test] + async fn models_status_returns_direct_and_routing_rows() { + use aisix_core::resource::ResourceEntry; + use aisix_core::Model; + use aisix_proxy::ModelRuntimeStatusTracker; + + let handle = SnapshotHandle::new(AisixSnapshot::new()); + let store = InMemoryStore::new() as Arc; + let runtime = Arc::new(ModelRuntimeStatusTracker::new()); + + let direct: Model = serde_json::from_value(model_payload("gpt4")).unwrap(); + store + .put_model(ResourceEntry { + id: "direct-1".into(), + value: direct, + revision: 1, + }) + .await + .unwrap(); + + let routing: Model = serde_json::from_value(json!({ + "display_name": "router", + "routing": { + "targets": [{"model": "gpt4"}] + } + })) + .unwrap(); + store + .put_model(ResourceEntry { + id: "routing-1".into(), + value: routing, + revision: 1, + }) + .await + .unwrap(); + + runtime.record_ignored_check("direct-1", 429, "ignored_transient_error"); + + let state = AdminState::new(handle, store, &cfg()).with_runtime_status_tracker(runtime); + let app = build_router(state); + let resp = run(app, auth_req("GET", "/admin/v1/models/status", None)).await; + assert_eq!(resp.status(), StatusCode::OK); + let rows = body_json(resp).await; + let rows = rows.as_array().unwrap(); + assert_eq!(rows.len(), 2); + + let direct = rows.iter().find(|row| row["id"] == "direct-1").unwrap(); + assert_eq!(direct["kind"], "direct"); + assert_eq!(direct["status"], "healthy"); + assert_eq!(direct["last_check_status"], 429); + assert_eq!(direct["status_reason"], "ignored_transient_error"); + + let routing = rows.iter().find(|row| row["id"] == "routing-1").unwrap(); + assert_eq!(routing["kind"], "routing"); + assert_eq!(routing["status"], "not_applicable"); + } + + #[tokio::test] + async fn create_model_accepts_background_model_check() { + let app = build_router(build_state()); + let resp = run( + app, + auth_req( + "POST", + "/admin/v1/models", + Some(json!({ + "display_name": "bg-model", + "provider": "openai", + "model_name": "gpt-4o-mini", + "provider_key_id": "11111111-1111-1111-1111-111111111111", + "background_model_check": { + "enabled": true, + "interval_seconds": 1, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "ignore_statuses": [408, 429], + "stale_after_seconds": 90 + } + })), + ), + ) + .await; + assert_eq!(resp.status(), StatusCode::OK); + } } diff --git a/crates/aisix-admin/src/models_status_handler.rs b/crates/aisix-admin/src/models_status_handler.rs new file mode 100644 index 00000000..cd63f5a5 --- /dev/null +++ b/crates/aisix-admin/src/models_status_handler.rs @@ -0,0 +1,72 @@ +//! `GET /admin/v1/models/status` — runtime per-model status. + +use axum::extract::State; +use axum::Json; +use serde::Serialize; +use std::time::Duration; + +use aisix_proxy::{RuntimeStatus, RuntimeStatusSnapshot}; + +use crate::auth::AdminAuth; +use crate::error::AdminError; +use crate::state::AdminState; + +#[derive(Debug, Clone, Copy, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ModelKind { + Direct, + Routing, +} + +#[derive(Debug, Serialize)] +pub struct ModelStatusView { + pub id: String, + pub display_name: String, + pub kind: ModelKind, + #[serde(flatten)] + pub details: RuntimeStatusSnapshot, +} + +pub async fn get_models_status( + _auth: AdminAuth, + State(state): State, +) -> Result>, AdminError> { + let all_models = state.store.list_models().await?; + let tracker = state.runtime_status_tracker.as_ref(); + + let models = all_models + .into_iter() + .map(|entry| { + if entry.value.is_routing() { + ModelStatusView { + id: entry.id, + display_name: entry.value.display_name, + kind: ModelKind::Routing, + details: RuntimeStatusSnapshot { + status: RuntimeStatus::NotApplicable, + ..RuntimeStatusSnapshot::default() + }, + } + } else { + let details = tracker + .map(|t| { + let stale_after = entry + .value + .background_model_check + .as_ref() + .map(|cfg| Duration::from_secs(cfg.stale_after_seconds)); + t.status_with_stale(&entry.id, stale_after) + }) + .unwrap_or_default(); + ModelStatusView { + id: entry.id, + display_name: entry.value.display_name, + kind: ModelKind::Direct, + details, + } + } + }) + .collect(); + + Ok(Json(models)) +} diff --git a/crates/aisix-admin/src/openapi.rs b/crates/aisix-admin/src/openapi.rs index 87f69bcd..866740ac 100644 --- a/crates/aisix-admin/src/openapi.rs +++ b/crates/aisix-admin/src/openapi.rs @@ -115,6 +115,18 @@ const OPENAPI_JSON: &str = r##"{ }, "delete": { "summary": "delete model", "responses": {"200": {"description": "OK"}, "404": {"description": "not found"}} } }, + "/admin/v1/models/status": { + "get": { + "summary": "list per-model runtime status", + "description": "Returns runtime routing/exclusion state for every Model. Direct models surface live runtime state keyed by resolved direct-model id; routing models return `not_applicable`. Request-path retryable failures surface as `cooldown`. Background checks can surface `unhealthy` or a healthy row with `status_reason=ignored_transient_error`.", + "responses": { + "200": { + "description": "OK", + "content": {"application/json": {"schema": {"type": "array", "items": {"$ref": "#/components/schemas/ModelStatusView"}}}} + } + } + } + }, "/admin/v1/apikeys": { "get": { "summary": "list api keys", "responses": {"200": {"description": "OK"}} }, "post": { @@ -265,9 +277,10 @@ const OPENAPI_JSON: &str = r##"{ "timeout": {"type": "integer", "minimum": 0, "description": "Request timeout in milliseconds. Absent or 0 = no timeout."}, "rate_limit": {"$ref": "#/components/schemas/RateLimit"}, "routing": {"$ref": "#/components/schemas/Routing"}, - "cost": {"$ref": "#/components/schemas/ModelCost"} + "cost": {"$ref": "#/components/schemas/ModelCost"}, + "background_model_check": {"$ref": "#/components/schemas/BackgroundModelCheck"} }, - "description": "A direct model ships `provider` + `model_name` + `provider_key_id`; a routing model ships `routing` and omits the upstream triple." + "description": "A direct model ships `provider` + `model_name` + `provider_key_id`; a routing model ships `routing` and omits the upstream triple. `background_model_check` is direct-model-only and rejected on routing models." }, "ModelEntry": { "type": "object", @@ -278,6 +291,55 @@ const OPENAPI_JSON: &str = r##"{ "revision": {"type": "integer"} } }, + "BackgroundModelCheck": { + "type": "object", + "required": ["enabled", "interval_seconds", "timeout_seconds", "prompt", "max_tokens", "stale_after_seconds"], + "properties": { + "enabled": {"type": "boolean", "description": "Turns the periodic direct-model probe on or off."}, + "interval_seconds": {"type": "integer", "minimum": 1, "description": "Probe interval in seconds."}, + "timeout_seconds": {"type": "integer", "minimum": 1, "description": "Per-probe timeout in seconds."}, + "prompt": {"type": "string", "minLength": 1, "description": "Minimal prompt used by the background probe request."}, + "max_tokens": {"type": "integer", "minimum": 1, "description": "Max completion tokens used by the probe request."}, + "ignore_statuses": { + "type": "array", + "description": "Upstream HTTP statuses that should be recorded without marking the model unhealthy. Typical values are 408 and 429.", + "items": {"type": "integer", "minimum": 100, "maximum": 599} + }, + "stale_after_seconds": {"type": "integer", "minimum": 1, "description": "Age threshold after which an unhealthy background-check result is treated as stale and stops excluding the model."} + }, + "description": "Periodic direct-model health-check configuration. Rejected on routing models." + }, + "ModelStatusView": { + "type": "object", + "required": ["id", "display_name", "kind", "status"], + "properties": { + "id": {"type": "string", "description": "Resolved model id. Direct-model runtime status is keyed by this id."}, + "display_name": {"type": "string"}, + "kind": {"$ref": "#/components/schemas/ModelKind"}, + "status": {"$ref": "#/components/schemas/RuntimeStatus"}, + "cooldown_until": {"$ref": "#/components/schemas/SystemTime"}, + "last_checked_at": {"$ref": "#/components/schemas/SystemTime"}, + "last_check_status": {"type": "integer", "minimum": 100, "maximum": 599}, + "status_reason": {"type": "string", "description": "Machine-readable explanation such as `retryable_failure`, `background_check_failed`, or `ignored_transient_error`."} + }, + "description": "Per-model runtime routing status. Routing rows always return `kind=routing` and `status=not_applicable`." + }, + "ModelKind": { + "type": "string", + "enum": ["direct", "routing"] + }, + "RuntimeStatus": { + "type": "string", + "enum": ["healthy", "unhealthy", "cooldown", "not_applicable"] + }, + "SystemTime": { + "type": "object", + "required": ["secs_since_epoch", "nanos_since_epoch"], + "properties": { + "secs_since_epoch": {"type": "integer", "minimum": 0}, + "nanos_since_epoch": {"type": "integer", "minimum": 0, "maximum": 999999999} + } + }, "ApiKey": { "type": "object", "required": ["key_hash", "allowed_models"], @@ -436,6 +498,7 @@ mod tests { "/admin/openapi-scalar", "/admin/v1/models", "/admin/v1/models/{id}", + "/admin/v1/models/status", "/admin/v1/apikeys", "/admin/v1/apikeys/{id}", "/admin/v1/apikeys/{id}/rotate", @@ -459,6 +522,11 @@ mod tests { for schema in [ "Model", "ModelEntry", + "BackgroundModelCheck", + "ModelStatusView", + "ModelKind", + "RuntimeStatus", + "SystemTime", "ApiKey", "ApiKeyEntry", "ProviderKey", diff --git a/crates/aisix-admin/src/state.rs b/crates/aisix-admin/src/state.rs index b68537dc..c24d7ef6 100644 --- a/crates/aisix-admin/src/state.rs +++ b/crates/aisix-admin/src/state.rs @@ -15,7 +15,7 @@ use aisix_core::snapshot::SnapshotHandle; use aisix_core::{AdminConfig, AisixSnapshot}; use aisix_etcd::WatchStatus; use aisix_obs::Metrics; -use aisix_proxy::{HealthTracker, LivezState}; +use aisix_proxy::{HealthTracker, LivezState, ModelRuntimeStatusTracker}; use axum::Router; use std::sync::Arc; @@ -30,6 +30,9 @@ pub struct AdminState { /// Shared in-process health tracker from the proxy. Used by the /// `/admin/v1/health` endpoint to report per-model health status. pub health_tracker: Option>, + /// Shared in-process runtime status tracker from the proxy. Used by + /// `/admin/v1/models/status` to report direct-model runtime state. + pub runtime_status_tracker: Option>, /// Watch supervisor's freshness state. When wired, the /// `/admin/v1/health` endpoint includes etcd revision + /// snapshot age so operators can detect a frozen / wedged config @@ -57,6 +60,7 @@ impl AdminState { store, metrics: None, health_tracker: None, + runtime_status_tracker: None, watch_status: None, livez_state: Arc::new(LivezState::new()), proxy_router: None, @@ -93,6 +97,12 @@ impl AdminState { self } + /// Attach the in-process runtime status tracker from the proxy. + pub fn with_runtime_status_tracker(mut self, tracker: Arc) -> Self { + self.runtime_status_tracker = Some(tracker); + self + } + /// Wire the proxy router so the playground endpoint can forward /// requests to it in-process via `tower::ServiceExt::oneshot`. pub fn with_proxy_router(mut self, router: Router) -> Self { diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index 14b7db61..ea2c58cd 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -32,7 +32,7 @@ pub use guardrail::{ BedrockAWSCredentials, BedrockConfig, BedrockLatencyMode, Guardrail, GuardrailHookPoint, GuardrailKind, KeywordConfig, KeywordPattern, }; -pub use model::{Model, Provider}; +pub use model::{BackgroundModelCheck, Model, Provider}; pub use observability_exporter::{ExporterKind, ObservabilityExporter, OtlpHttpConfig}; pub use provider_key::ProviderKey; pub use rate_limit::RateLimit; diff --git a/crates/aisix-core/src/models/model.rs b/crates/aisix-core/src/models/model.rs index 282b9361..eb3df854 100644 --- a/crates/aisix-core/src/models/model.rs +++ b/crates/aisix-core/src/models/model.rs @@ -81,6 +81,19 @@ impl ModelCost { } } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct BackgroundModelCheck { + pub enabled: bool, + pub interval_seconds: u64, + pub timeout_seconds: u64, + pub prompt: String, + pub max_tokens: u32, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub ignore_statuses: Vec, + pub stale_after_seconds: u64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct Model { @@ -124,6 +137,10 @@ pub struct Model { #[serde(default, skip_serializing_if = "Option::is_none")] pub cost: Option, + /// Optional direct-model-only background health-check configuration. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub background_model_check: Option, + /// Non-schema runtime id. Not part of the JSON payload — filled in by /// the snapshot loader from the etcd key path. Kept here so `Resource` /// can return a `&str` id. @@ -227,6 +244,31 @@ mod tests { assert_eq!(m.name(), "my-gpt4"); } + #[test] + fn direct_model_can_deserialize_background_check() { + let m: Model = serde_json::from_str( + r#"{ + "display_name": "my-gpt4", + "provider": "openai", + "model_name": "gpt-4o", + "provider_key_id": "11111111-1111-1111-1111-111111111111", + "background_model_check": { + "enabled": true, + "interval_seconds": 30, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "ignore_statuses": [408, 429], + "stale_after_seconds": 90 + } + }"#, + ) + .unwrap(); + let bg = m.background_model_check.unwrap(); + assert!(bg.enabled); + assert_eq!(bg.ignore_statuses, vec![408, 429]); + } + #[test] fn provider_default_urls_are_stable() { assert_eq!( diff --git a/crates/aisix-core/src/models/schema.rs b/crates/aisix-core/src/models/schema.rs index 7ae46d48..9a57bd81 100644 --- a/crates/aisix-core/src/models/schema.rs +++ b/crates/aisix-core/src/models/schema.rs @@ -149,6 +149,30 @@ fn model_schema() -> Value { "input_per_1k": { "type": "number", "minimum": 0 }, "output_per_1k": { "type": "number", "minimum": 0 } } + }, + "background_model_check": { + "type": "object", + "required": [ + "enabled", + "interval_seconds", + "timeout_seconds", + "prompt", + "max_tokens", + "stale_after_seconds" + ], + "additionalProperties": false, + "properties": { + "enabled": { "type": "boolean" }, + "interval_seconds": { "type": "integer", "minimum": 1 }, + "timeout_seconds": { "type": "integer", "minimum": 1 }, + "prompt": { "type": "string", "minLength": 1 }, + "max_tokens": { "type": "integer", "minimum": 1 }, + "ignore_statuses": { + "type": "array", + "items": { "type": "integer", "minimum": 100, "maximum": 599 } + }, + "stale_after_seconds": { "type": "integer", "minimum": 1 } + } } }, // Direct vs routing model: a model EITHER ships a `routing` @@ -161,7 +185,8 @@ fn model_schema() -> Value { "not": { "anyOf": [ { "required": ["provider"] }, { "required": ["model_name"] }, - { "required": ["provider_key_id"] } + { "required": ["provider_key_id"] }, + { "required": ["background_model_check"] } ]} }, { @@ -524,6 +549,65 @@ mod tests { assert!(validate_model(&v).is_err()); } + #[test] + fn direct_model_background_check_passes() { + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "background_model_check": { + "enabled": true, + "interval_seconds": 30, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "ignore_statuses": [408, 429], + "stale_after_seconds": 90 + } + }); + validate_model(&v).unwrap(); + } + + #[test] + fn routing_model_background_check_fails() { + let v = json!({ + "display_name": "router-1", + "routing": { + "targets": [{"model": "my-gpt4"}] + }, + "background_model_check": { + "enabled": true, + "interval_seconds": 30, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "stale_after_seconds": 90 + } + }); + assert!(validate_model(&v).is_err()); + } + + #[test] + fn background_check_rejects_invalid_ignore_status() { + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "background_model_check": { + "enabled": true, + "interval_seconds": 30, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "ignore_statuses": [99], + "stale_after_seconds": 90 + } + }); + assert!(validate_model(&v).is_err()); + } + #[test] fn schemas_initialise_once() { let a = Arc::as_ptr(&*SCHEMAS); diff --git a/crates/aisix-proxy/src/background.rs b/crates/aisix-proxy/src/background.rs new file mode 100644 index 00000000..3c389b14 --- /dev/null +++ b/crates/aisix-proxy/src/background.rs @@ -0,0 +1,177 @@ +use std::sync::Arc; +use std::time::Duration; + +use aisix_core::models::BackgroundModelCheck; +use aisix_core::{AisixSnapshot, Model}; +use aisix_gateway::{BridgeContext, BridgeError, ChatFormat, ChatMessage, Hub}; + +use crate::dispatch; +use crate::health::ModelRuntimeStatusTracker; + +pub async fn run_background_model_check_once( + snapshot: Arc, + hub: Arc, + tracker: Arc, + request_id: &str, +) { + for entry in snapshot.models.entries() { + let model = &entry.value; + if model.is_routing() { + continue; + } + let Some(cfg) = model.background_model_check.as_ref() else { + continue; + }; + if !cfg.enabled { + continue; + } + let outcome = check_direct_model(&snapshot, &hub, &entry.id, model, cfg, request_id).await; + match outcome { + Ok(()) => tracker.clear_unhealthy(&entry.id), + Err(BridgeError::UpstreamStatus { status, .. }) if cfg.ignore_statuses.contains(&status) => { + tracker.record_ignored_check(&entry.id, status, "ignored_transient_error") + } + Err(BridgeError::Timeout { .. }) if cfg.ignore_statuses.contains(&408) => { + tracker.record_ignored_check(&entry.id, 408, "ignored_transient_error") + } + Err(err) => tracker.mark_unhealthy( + &entry.id, + background_status_code(&err), + "background_check_failed", + ), + } + } +} + +async fn check_direct_model( + snapshot: &AisixSnapshot, + hub: &Hub, + model_id: &str, + model: &Model, + cfg: &BackgroundModelCheck, + request_id: &str, +) -> Result<(), BridgeError> { + let provider = dispatch::require_provider(model).map_err(|e| BridgeError::Config(e.to_string()))?; + let pk_entry = dispatch::resolve_provider_key(snapshot, model) + .map_err(|e| BridgeError::Config(e.to_string()))?; + let bridge = hub + .get(provider) + .ok_or_else(|| BridgeError::Config("no bridge registered for provider".into()))?; + + let req = ChatFormat { + model: model.display_name.clone(), + messages: vec![ChatMessage::user(cfg.prompt.clone())], + max_tokens: Some(cfg.max_tokens), + ..ChatFormat::new(model.display_name.clone(), vec![]) + }; + let model_arc = Arc::new(model.clone()); + let pk_arc = Arc::new(pk_entry.value.clone()); + let ctx = BridgeContext::new(request_id, model_arc, pk_arc).with_deadline(timeout(cfg)); + + let _ = bridge.chat(&req, &ctx).await?; + let _ = model_id; + Ok(()) +} + +fn timeout(cfg: &BackgroundModelCheck) -> Duration { + Duration::from_secs(cfg.timeout_seconds) +} + +fn background_status_code(err: &BridgeError) -> Option { + match err { + BridgeError::UpstreamStatus { status, .. } => Some(*status), + BridgeError::Timeout { .. } => Some(408), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use aisix_core::models::Provider; + use aisix_core::resource::ResourceEntry; + use aisix_provider_openai::OpenAiBridge; + use reqwest::Client; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + fn openai_test_bridge() -> OpenAiBridge { + let client = Client::builder().user_agent("aisix-test/0.1").no_proxy().build().unwrap(); + OpenAiBridge::with_client(client) + } + + fn provider_key_entry(id: &str, api_base: &str) -> ResourceEntry { + let cfg = format!( + r#"{{"display_name":"pk-{id}","secret":"sk-upstream","api_base":"{api_base}"}}"# + ); + let pk: aisix_core::ProviderKey = serde_json::from_str(&cfg).unwrap(); + ResourceEntry::new(id, pk, 1) + } + + fn direct_model_entry(id: &str, name: &str, pk_id: &str, enabled: bool, ignore: &[u16]) -> ResourceEntry { + let cfg = serde_json::json!({ + "display_name": name, + "provider": "openai", + "model_name": "gpt-4o-mini", + "provider_key_id": pk_id, + "background_model_check": { + "enabled": enabled, + "interval_seconds": 30, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "ignore_statuses": ignore, + "stale_after_seconds": 90 + } + }); + let model: Model = serde_json::from_value(cfg).unwrap(); + ResourceEntry::new(id, model, 1) + } + + #[tokio::test] + async fn background_check_marks_unhealthy_on_failure() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(503).set_body_string("down")) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(openai_test_bridge())); + let snapshot = Arc::new(AisixSnapshot::new()); + snapshot.provider_keys.insert(provider_key_entry("pk-1", &upstream.uri())); + snapshot.models.insert(direct_model_entry("m-1", "bg-model", "pk-1", true, &[408, 429])); + let tracker = Arc::new(ModelRuntimeStatusTracker::new()); + + run_background_model_check_once(snapshot, hub, tracker.clone(), "bg-check-1").await; + + let status = tracker.status("m-1"); + assert_eq!(status.status, crate::RuntimeStatus::Unhealthy); + assert_eq!(status.last_check_status, Some(503)); + } + + #[tokio::test] + async fn background_check_ignores_configured_transient_status() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(429).set_body_string("slow down")) + .mount(&upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(openai_test_bridge())); + let snapshot = Arc::new(AisixSnapshot::new()); + snapshot.provider_keys.insert(provider_key_entry("pk-1", &upstream.uri())); + snapshot.models.insert(direct_model_entry("m-1", "bg-model", "pk-1", true, &[408, 429])); + let tracker = Arc::new(ModelRuntimeStatusTracker::new()); + + run_background_model_check_once(snapshot, hub, tracker.clone(), "bg-check-1").await; + + let status = tracker.status("m-1"); + assert_eq!(status.status, crate::RuntimeStatus::Healthy); + assert_eq!(status.last_check_status, Some(429)); + assert_eq!(status.status_reason.as_deref(), Some("ignored_transient_error")); + } +} diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index d05faaf3..16d89e66 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -40,6 +40,13 @@ use crate::state::ProxyState; /// Header set on every non-streaming response indicating whether the /// response came from the cache (`hit`) or the upstream (`miss`). pub const CACHE_HEADER: &str = "x-aisix-cache"; +const RETRYABLE_FAILURE_COOLDOWN: Duration = Duration::from_secs(30); + +#[derive(Clone)] +struct AttemptModel { + id: String, + model: aisix_core::Model, +} pub async fn chat_completions( State(state): State, @@ -426,7 +433,7 @@ async fn dispatch( // Resolve the attempt-list of underlying Model entries. For a // routing model we walk targets per the configured strategy; for a // single-provider Model we just dispatch to it directly. - let attempt_models: Vec = + let attempt_models: Vec = if let Some(routing) = virtual_entry.value.routing.as_ref() { let names = state.routing.pick_targets(&req.model, routing); if names.is_empty() { @@ -441,11 +448,17 @@ async fn dispatch( "routing target {name:?} does not resolve to a Model" ))) })?; - resolved.push(target_entry.value.clone()); + resolved.push(AttemptModel { + id: target_entry.id.clone(), + model: target_entry.value.clone(), + }); } - resolved + filter_attempt_models(&state.runtime_status, resolved) } else { - vec![virtual_entry.value.clone()] + vec![AttemptModel { + id: virtual_entry.id.clone(), + model: virtual_entry.value.clone(), + }] }; // For non-routing requests, surface a misconfigured bridge as a @@ -453,7 +466,7 @@ async fn dispatch( // Routing requests rely on the loop's `is_retryable` path so a // single bad provider doesn't take down the whole request. if attempt_models.len() == 1 { - let only = &attempt_models[0]; + let only = &attempt_models[0].model; let provider = crate::dispatch::require_provider(only).map_err(with_model)?; if state.hub.get(provider).is_none() { return Err(with_model(ProxyError::ProviderUnavailable)); @@ -473,7 +486,7 @@ async fn dispatch( // is genuinely hard (we'd have to buffer the stream to detect // failure mid-flight) and not worth the complexity for V1. if req.is_streaming() { - let model = &attempt_models[0]; + let model = &attempt_models[0].model; let provider = crate::dispatch::require_provider(model).map_err(with_model)?; let pk_entry = crate::dispatch::resolve_provider_key(&snapshot, model).map_err(with_model)?; @@ -687,6 +700,7 @@ async fn dispatch( // cache hit we don't know (or care) which target ran the // original call; the fingerprint identified the answer. let provider_label = attempt_models[0] + .model .provider .map(|p| format!("{p:?}").to_lowercase()) .unwrap_or_else(|| "unknown".into()); @@ -754,7 +768,8 @@ async fn dispatch( .map(|routing| routing.retry_on_429_or_default()) .unwrap_or(false); - for model in &attempt_models { + for attempt in &attempt_models { + let model = &attempt.model; let Some(provider) = model.provider else { last_err = Some(BridgeError::Config("model has no provider".into())); continue; @@ -782,6 +797,7 @@ async fn dispatch( match bridge.chat(req, &ctx).await { Ok(resp) => { state.health.record_success(&model.display_name); + state.runtime_status.mark_healthy(&attempt.id); chosen_provider = Some(format!("{provider:?}").to_lowercase()); upstream = Some(resp); break; @@ -797,6 +813,11 @@ async fn dispatch( ); if retryable { state.health.record_failure(&model.display_name); + state.runtime_status.mark_cooldown( + &attempt.id, + RETRYABLE_FAILURE_COOLDOWN, + "retryable_failure", + ); } last_err = Some(err); if !retryable { @@ -960,6 +981,50 @@ async fn dispatch( }) } +fn filter_attempt_models( + runtime_status: &crate::ModelRuntimeStatusTracker, + attempts: Vec, +) -> Vec { + let mut healthy = Vec::new(); + let mut cooldown_only = Vec::new(); + let mut unhealthy_count = 0usize; + + for attempt in attempts.iter().cloned() { + let stale_after = attempt + .model + .background_model_check + .as_ref() + .map(|cfg| Duration::from_secs(cfg.stale_after_seconds)); + let status = runtime_status.status_with_stale(&attempt.id, stale_after); + match status.status { + crate::RuntimeStatus::Unhealthy => unhealthy_count += 1, + crate::RuntimeStatus::Cooldown => cooldown_only.push(attempt), + crate::RuntimeStatus::Healthy | crate::RuntimeStatus::NotApplicable => { + healthy.push(attempt) + } + } + } + + if !healthy.is_empty() { + return healthy; + } + if unhealthy_count < attempts.len() && !cooldown_only.is_empty() { + return attempts + .into_iter() + .filter(|attempt| { + let stale_after = attempt + .model + .background_model_check + .as_ref() + .map(|cfg| Duration::from_secs(cfg.stale_after_seconds)); + runtime_status.should_skip_for_routing(&attempt.id, stale_after) + != crate::RuntimeStatus::Unhealthy + }) + .collect(); + } + attempts +} + /// Wire-shape label for `FinishReason`. cp-api stores this verbatim /// in `dpmgr_usage_events.finish_reason`; the dashboard reads it back /// to distinguish normal stops from truncation / content_filter. diff --git a/crates/aisix-proxy/src/health.rs b/crates/aisix-proxy/src/health.rs index e8375eb6..1dd10560 100644 --- a/crates/aisix-proxy/src/health.rs +++ b/crates/aisix-proxy/src/health.rs @@ -16,6 +16,7 @@ use dashmap::DashMap; use std::sync::atomic::AtomicBool; use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::{Duration, SystemTime}; use axum::http::header::{HeaderName, HeaderValue, CONTENT_TYPE}; use axum::http::StatusCode; @@ -99,6 +100,95 @@ pub enum HealthLevel { Down, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "snake_case")] +pub enum RuntimeStatus { + Healthy, + Unhealthy, + Cooldown, + NotApplicable, +} + +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize)] +pub struct RuntimeStatusSnapshot { + pub status: RuntimeStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub cooldown_until: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub last_checked_at: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub last_check_status: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub status_reason: Option, +} + +impl Default for RuntimeStatusSnapshot { + fn default() -> Self { + Self { + status: RuntimeStatus::Healthy, + cooldown_until: None, + last_checked_at: None, + last_check_status: None, + status_reason: None, + } + } +} + +#[derive(Debug, Clone)] +struct RuntimeEntry { + unhealthy: bool, + cooldown_until: Option, + last_checked_at: Option, + last_check_status: Option, + status_reason: Option, +} + +impl Default for RuntimeEntry { + fn default() -> Self { + Self { + unhealthy: false, + cooldown_until: None, + last_checked_at: None, + last_check_status: None, + status_reason: None, + } + } +} + +impl RuntimeEntry { + fn snapshot(&self, now: SystemTime, stale_after: Option) -> RuntimeStatusSnapshot { + let cooldown_until = self.cooldown_until.filter(|until| *until > now); + let unhealthy = self.unhealthy && !self.is_stale(now, stale_after); + let status = if cooldown_until.is_some() { + RuntimeStatus::Cooldown + } else if unhealthy { + RuntimeStatus::Unhealthy + } else { + RuntimeStatus::Healthy + }; + RuntimeStatusSnapshot { + status, + cooldown_until, + last_checked_at: self.last_checked_at, + last_check_status: self.last_check_status, + status_reason: self.status_reason.clone(), + } + } + + fn is_stale(&self, now: SystemTime, stale_after: Option) -> bool { + let Some(stale_after) = stale_after else { + return false; + }; + let Some(last_checked_at) = self.last_checked_at else { + return false; + }; + match now.duration_since(last_checked_at) { + Ok(elapsed) => elapsed > stale_after, + Err(_) => false, + } + } +} + impl From for u8 { fn from(h: HealthLevel) -> u8 { match h { @@ -170,6 +260,11 @@ pub struct HealthTracker { entries: DashMap, } +#[derive(Default, Debug)] +pub struct ModelRuntimeStatusTracker { + entries: DashMap, +} + impl HealthTracker { pub fn new() -> Self { Self::default() @@ -211,10 +306,114 @@ impl HealthTracker { } } +impl ModelRuntimeStatusTracker { + pub fn new() -> Self { + Self::default() + } + + pub fn mark_cooldown(&self, model_id: &str, ttl: Duration, reason: impl Into) { + let now = SystemTime::now(); + let until = now + ttl; + let reason = reason.into(); + self.entries + .entry(model_id.to_string()) + .and_modify(|entry| { + entry.cooldown_until = Some(until); + entry.status_reason = Some(reason.clone()); + }) + .or_insert_with(|| RuntimeEntry { + cooldown_until: Some(until), + status_reason: Some(reason), + ..RuntimeEntry::default() + }); + } + + pub fn mark_healthy(&self, model_id: &str) { + if let Some(mut entry) = self.entries.get_mut(model_id) { + entry.unhealthy = false; + entry.cooldown_until = None; + entry.status_reason = None; + } + } + + pub fn clear_unhealthy(&self, model_id: &str) { + if let Some(mut entry) = self.entries.get_mut(model_id) { + entry.unhealthy = false; + if entry.status_reason.as_deref() == Some("background_check_failed") { + entry.status_reason = None; + } + } + } + + pub fn mark_unhealthy(&self, model_id: &str, status: Option, reason: impl Into) { + let now = SystemTime::now(); + let reason = reason.into(); + self.entries + .entry(model_id.to_string()) + .and_modify(|entry| { + entry.unhealthy = true; + entry.last_checked_at = Some(now); + entry.last_check_status = status; + entry.status_reason = Some(reason.clone()); + }) + .or_insert_with(|| RuntimeEntry { + unhealthy: true, + last_checked_at: Some(now), + last_check_status: status, + status_reason: Some(reason), + ..RuntimeEntry::default() + }); + } + + pub fn record_ignored_check( + &self, + model_id: &str, + status: u16, + reason: impl Into, + ) { + let now = SystemTime::now(); + let reason = reason.into(); + self.entries + .entry(model_id.to_string()) + .and_modify(|entry| { + entry.last_checked_at = Some(now); + entry.last_check_status = Some(status); + entry.status_reason = Some(reason.clone()); + }) + .or_insert_with(|| RuntimeEntry { + last_checked_at: Some(now), + last_check_status: Some(status), + status_reason: Some(reason), + ..RuntimeEntry::default() + }); + } + + pub fn status(&self, model_id: &str) -> RuntimeStatusSnapshot { + self.status_with_stale(model_id, None) + } + + pub fn status_with_stale( + &self, + model_id: &str, + stale_after: Option, + ) -> RuntimeStatusSnapshot { + let now = SystemTime::now(); + self.entries + .get(model_id) + .map(|entry| entry.snapshot(now, stale_after)) + .unwrap_or_default() + } + + pub fn should_skip_for_routing(&self, model_id: &str, stale_after: Option) -> RuntimeStatus { + self.status_with_stale(model_id, stale_after).status + } +} + #[cfg(test)] mod tests { use super::*; use axum::body::to_bytes; + use std::thread; #[test] fn new_model_is_healthy() { @@ -302,4 +501,57 @@ mod tests { assert!(text.contains("[-]shutdown failed: reason withheld")); assert!(text.contains("livez check failed")); } + + #[test] + fn runtime_tracker_defaults_to_healthy() { + let t = ModelRuntimeStatusTracker::new(); + let s = t.status("m-1"); + assert_eq!(s.status, RuntimeStatus::Healthy); + assert!(s.cooldown_until.is_none()); + } + + #[test] + fn runtime_tracker_cooldown_expires() { + let t = ModelRuntimeStatusTracker::new(); + t.mark_cooldown("m-1", Duration::from_millis(5), "retryable_failure"); + assert_eq!(t.status("m-1").status, RuntimeStatus::Cooldown); + thread::sleep(Duration::from_millis(10)); + assert_eq!(t.status("m-1").status, RuntimeStatus::Healthy); + } + + #[test] + fn runtime_tracker_unhealthy_then_healthy() { + let t = ModelRuntimeStatusTracker::new(); + t.mark_unhealthy("m-1", Some(500), "background_check_failed"); + let unhealthy = t.status("m-1"); + assert_eq!(unhealthy.status, RuntimeStatus::Unhealthy); + assert_eq!(unhealthy.last_check_status, Some(500)); + t.mark_healthy("m-1"); + assert_eq!(t.status("m-1").status, RuntimeStatus::Healthy); + } + + #[test] + fn runtime_tracker_ignored_status_does_not_mark_unhealthy() { + let t = ModelRuntimeStatusTracker::new(); + t.record_ignored_check("m-1", 429, "ignored_transient_error"); + let s = t.status("m-1"); + assert_eq!(s.status, RuntimeStatus::Healthy); + assert_eq!(s.last_check_status, Some(429)); + assert_eq!(s.status_reason.as_deref(), Some("ignored_transient_error")); + } + + #[test] + fn runtime_tracker_unhealthy_becomes_healthy_after_stale_window() { + let t = ModelRuntimeStatusTracker::new(); + t.mark_unhealthy("m-1", Some(503), "background_check_failed"); + assert_eq!( + t.status_with_stale("m-1", Some(Duration::from_secs(60))).status, + RuntimeStatus::Unhealthy + ); + std::thread::sleep(Duration::from_millis(15)); + assert_eq!( + t.status_with_stale("m-1", Some(Duration::from_millis(1))).status, + RuntimeStatus::Healthy + ); + } } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 73ce02a0..4a8c9d24 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -26,6 +26,7 @@ mod audio; mod auth; +pub mod background; pub mod budget; mod chat; mod completions; @@ -47,7 +48,9 @@ mod state; pub use auth::AuthenticatedKey; pub use error::{ErrorEnvelope, ProxyError}; -pub use health::{HealthTracker, LivezState}; +pub use health::{ + HealthTracker, LivezState, ModelRuntimeStatusTracker, RuntimeStatus, RuntimeStatusSnapshot, +}; pub use state::ProxyState; use axum::extract::State; @@ -2015,12 +2018,7 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!( - status, - StatusCode::OK, - "{}", - String::from_utf8_lossy(&bytes) - ); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "fallback worked"); } @@ -2085,12 +2083,7 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!( - status, - StatusCode::BAD_REQUEST, - "{}", - String::from_utf8_lossy(&bytes) - ); + assert_eq!(status, StatusCode::BAD_REQUEST, "{}", String::from_utf8_lossy(&bytes)); } #[tokio::test] @@ -2158,12 +2151,7 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!( - status, - StatusCode::OK, - "{}", - String::from_utf8_lossy(&bytes) - ); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "after retries"); } @@ -2233,33 +2221,261 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!( - status, - StatusCode::OK, - "{}", - String::from_utf8_lossy(&bytes) - ); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "429 fallback worked"); } #[tokio::test] - async fn routing_to_missing_target_returns_400() { - // Routing references a Model that isn't in the snapshot — this - // is a misconfiguration and should surface as a clean 400. + async fn routing_skips_target_in_runtime_cooldown() { + let cooled_upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-cooled", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "should not be called"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .expect(0) + .mount(&cooled_upstream) + .await; + + let good_upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-good", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "cooldown skipped"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .expect(1) + .mount(&good_upstream) + .await; + let hub = Arc::new(Hub::new()); hub.register(Provider::Openai, Arc::new(openai_test_bridge())); let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(pk_entry_with_id("pk-cooled", &cooled_upstream.uri())); + snap.provider_keys + .insert(pk_entry_with_id("pk-good", &good_upstream.uri())); + snap.models + .insert(model_entry_with_id("m-cooled", "primary", "pk-cooled")); + snap.models + .insert(model_entry_with_id("m-good", "secondary", "pk-good")); snap.models.insert(routing_entry( "smart", "failover", - &["nonexistent"], + &["primary", "secondary"], + None, None, None, + )); + snap.apikeys.insert(apikey_entry("sk-caller", &["smart"])); + + let state = build_state(snap, hub); + state.runtime_status.mark_cooldown( + "m-cooled", + std::time::Duration::from_secs(30), + "retryable_failure", + ); + let app = build_router(state); + let body = serde_json::json!({ + "model": "smart", + "messages": [{"role": "user", "content": "hi"}] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + + let resp = run(app, req).await; + let status = resp.status(); + let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["choices"][0]["message"]["content"], "cooldown skipped"); + } + + #[tokio::test] + async fn routing_ignores_cooldown_when_it_would_empty_all_candidates() { + let primary_upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-primary", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "cooldown fallback"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .expect(1) + .mount(&primary_upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(openai_test_bridge())); + + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(pk_entry_with_id("pk-primary", &primary_upstream.uri())); + snap.models + .insert(model_entry_with_id("m-primary", "primary", "pk-primary")); + snap.models.insert(routing_entry( + "smart", + "failover", + &["primary"], + None, + None, + None, + )); + snap.apikeys.insert(apikey_entry("sk-caller", &["smart"])); + + let state = build_state(snap, hub); + state.runtime_status.mark_cooldown( + "m-primary", + std::time::Duration::from_secs(30), + "retryable_failure", + ); + let app = build_router(state); + let body = serde_json::json!({ + "model": "smart", + "messages": [{"role": "user", "content": "hi"}] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + + let resp = run(app, req).await; + let status = resp.status(); + let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["choices"][0]["message"]["content"], "cooldown fallback"); + } + + #[tokio::test] + async fn routing_retryable_failure_puts_target_into_cooldown_for_next_request() { + let flaky_upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(502).set_body_string("temporary upstream failure"), + ) + .expect(1) + .mount(&flaky_upstream) + .await; + + let stable_upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "cmpl-stable", + "model": "gpt-4o", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "stable fallback"}, + "finish_reason": "stop" + }], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2} + }))) + .expect(2) + .mount(&stable_upstream) + .await; + + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(openai_test_bridge())); + + let snap = AisixSnapshot::new(); + snap.provider_keys + .insert(pk_entry_with_id("pk-flaky", &flaky_upstream.uri())); + snap.provider_keys + .insert(pk_entry_with_id("pk-stable", &stable_upstream.uri())); + snap.models + .insert(model_entry_with_id("m-flaky", "primary", "pk-flaky")); + snap.models + .insert(model_entry_with_id("m-stable", "secondary", "pk-stable")); + snap.models.insert(routing_entry( + "smart", + "failover", + &["primary", "secondary"], + Some(0), + Some(1), None, )); snap.apikeys.insert(apikey_entry("sk-caller", &["smart"])); + + let state = build_state(snap, hub); + let app = build_router(state.clone()); + let body = serde_json::json!({ + "model": "smart", + "messages": [{"role": "user", "content": "first"}] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + + let resp = run(app, req).await; + assert_eq!(resp.status(), StatusCode::OK); + assert_eq!(state.runtime_status.status("m-flaky").status, RuntimeStatus::Cooldown); + + let app = build_router(state); + let body = serde_json::json!({ + "model": "smart", + "messages": [{"role": "user", "content": "second"}] + }); + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer sk-caller") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + + let resp = run(app, req).await; + let status = resp.status(); + let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["choices"][0]["message"]["content"], "stable fallback"); + } + + #[tokio::test] + async fn routing_to_missing_target_returns_400() { + // Routing references a Model that isn't in the snapshot — this + // is a misconfiguration and should surface as a clean 400. + let hub = Arc::new(Hub::new()); + hub.register(Provider::Openai, Arc::new(openai_test_bridge())); + + let snap = AisixSnapshot::new(); + snap.models + .insert(routing_entry("smart", "failover", &["nonexistent"], None, None, None)); + snap.apikeys.insert(apikey_entry("sk-caller", &["smart"])); // No upstream provider_key needed — the routing target itself // is missing so dispatch fails before any provider lookup. diff --git a/crates/aisix-proxy/src/routing.rs b/crates/aisix-proxy/src/routing.rs index 2a07968a..942a3eb2 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -60,11 +60,7 @@ impl RoutingRegistry { return Vec::new(); } let start = self.starting_index(virtual_name, routing); - attempt_order( - &routing.targets, - start, - routing.max_fallbacks_or_default() + 1, - ) + attempt_order(&routing.targets, start, routing.max_fallbacks_or_default() + 1) } fn starting_index(&self, virtual_name: &str, routing: &Routing) -> usize { @@ -411,40 +407,25 @@ mod tests { #[test] fn is_retryable_distinguishes_4xx_from_other_failures() { - assert!(!is_retryable( - &BridgeError::UpstreamStatus { - status: 400, - message: "bad request".into(), - }, - false - )); - assert!(!is_retryable( - &BridgeError::UpstreamStatus { - status: 429, - message: "rate limited".into(), - }, - false - )); - assert!(is_retryable( - &BridgeError::UpstreamStatus { - status: 429, - message: "rate limited".into(), - }, - true - )); - assert!(is_retryable( - &BridgeError::UpstreamStatus { - status: 502, - message: "bad gateway".into(), - }, - false - )); + assert!(!is_retryable(&BridgeError::UpstreamStatus { + status: 400, + message: "bad request".into(), + }, false)); + assert!(!is_retryable(&BridgeError::UpstreamStatus { + status: 429, + message: "rate limited".into(), + }, false)); + assert!(is_retryable(&BridgeError::UpstreamStatus { + status: 429, + message: "rate limited".into(), + }, true)); + assert!(is_retryable(&BridgeError::UpstreamStatus { + status: 502, + message: "bad gateway".into(), + }, false)); assert!(is_retryable(&BridgeError::Timeout { elapsed_ms: 1 }, false)); assert!(is_retryable(&BridgeError::Transport("conn".into()), false)); - assert!(is_retryable( - &BridgeError::UpstreamDecode("x".into()), - false - )); + assert!(is_retryable(&BridgeError::UpstreamDecode("x".into()), false)); assert!(is_retryable(&BridgeError::Config("bad key".into()), false)); assert!(is_retryable(&BridgeError::StreamAborted, false)); } diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index 000d6603..7deffc3d 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -24,7 +24,7 @@ use aisix_ratelimit::Limiter; use std::sync::Arc; use crate::budget::BudgetClient; -use crate::health::{HealthTracker, LivezState}; +use crate::health::{HealthTracker, LivezState, ModelRuntimeStatusTracker}; use crate::routing::RoutingRegistry; #[derive(Clone)] @@ -46,6 +46,10 @@ pub struct ProxyState { pub health: Arc, /// Public liveness state served on `GET /livez`. pub livez: Arc, + /// Runtime model-status tracker keyed by resolved direct-model id. + /// Used for request-path cooldown/background health exclusion and + /// surfaced by `GET /admin/v1/models/status`. + pub runtime_status: Arc, /// CP-side usage telemetry sink. Backed by an mpsc channel into the /// sender worker spawned in aisix-server (see `telemetry::spawn`). /// Defaults to a no-op sink when running outside managed mode so @@ -73,6 +77,7 @@ impl ProxyState { budgets: Arc::new(BudgetClient::disabled()), health: Arc::new(HealthTracker::new()), livez: Arc::new(LivezState::new()), + runtime_status: Arc::new(ModelRuntimeStatusTracker::new()), usage_sink: UsageSink::disabled(), otlp_fan_out: OtlpHttpFanOut::new(), request_body_limit_bytes: cfg.request_body_limit_bytes, @@ -98,6 +103,7 @@ impl ProxyState { budgets: Arc::new(BudgetClient::disabled()), health: Arc::new(HealthTracker::new()), livez: Arc::new(LivezState::new()), + runtime_status: Arc::new(ModelRuntimeStatusTracker::new()), usage_sink: UsageSink::disabled(), otlp_fan_out: OtlpHttpFanOut::new(), request_body_limit_bytes: cfg.request_body_limit_bytes, @@ -126,6 +132,7 @@ impl ProxyState { budgets: Arc::new(BudgetClient::disabled()), health: Arc::new(HealthTracker::new()), livez: Arc::new(LivezState::new()), + runtime_status: Arc::new(ModelRuntimeStatusTracker::new()), usage_sink: UsageSink::disabled(), otlp_fan_out: OtlpHttpFanOut::new(), request_body_limit_bytes: cfg.request_body_limit_bytes, diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 7895c823..1062dfd4 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -32,6 +32,7 @@ use aisix_provider_anthropic::AnthropicBridge; use aisix_provider_deepseek::deepseek_bridge; use aisix_provider_gemini::gemini_bridge; use aisix_provider_openai::OpenAiBridge; +use aisix_proxy::background::run_background_model_check_once; use aisix_proxy::budget::BudgetClient; use aisix_proxy::ProxyState; use aisix_ratelimit::Limiter; @@ -481,8 +482,39 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { // Clone shared trackers before consuming proxy_state in build_router. let health_tracker = proxy_state.health.clone(); let livez_state = proxy_state.livez.clone(); + let runtime_status_tracker = proxy_state.runtime_status.clone(); + let background_snapshot = snapshot_handle.clone(); + let background_hub = hub.clone(); + let background_runtime_status_tracker = runtime_status_tracker.clone(); + let background_cancel_rx = cancel_rx.clone(); let proxy_router = aisix_proxy::build_router(proxy_state); + let background_check_task = tokio::spawn(async move { + let mut cancel = background_cancel_rx; + loop { + if *cancel.borrow() { + break; + } + let snapshot = background_snapshot.load(); + run_background_model_check_once( + snapshot.clone(), + background_hub.clone(), + background_runtime_status_tracker.clone(), + "background-model-check", + ) + .await; + let sleep_for = background_check_interval(snapshot.as_ref()); + tokio::select! { + changed = cancel.changed() => { + if changed.is_err() || *cancel.borrow() { + break; + } + } + _ = tokio::time::sleep(sleep_for) => {} + } + } + }); + // Admin router + listener are only built in standalone mode. // In managed mode (`cfg.managed.enabled = true`) the DP reads // configuration exclusively from etcd; exposing admin writes or @@ -496,6 +528,9 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { // per-model upstream failure counts. .with_health_tracker(health_tracker) .with_livez_state(livez_state.clone()) + // Share runtime status so /admin/v1/models/status exposes + // direct-model cooldown/background-health state. + .with_runtime_status_tracker(runtime_status_tracker) // Share the supervisor's freshness state so /admin/v1/health // exposes etcd watch staleness — without this, a wedged // watch lets the gateway serve stale config indefinitely @@ -516,7 +551,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { // Drop unused shared components so the compiler can see they // don't escape managed mode. The health tracker exists on // proxy_state and keeps working regardless. - let _ = (&health_tracker, &livez_state); + let _ = (&health_tracker, &livez_state, &runtime_status_tracker); tracing::info!("managed mode enabled — admin surface not bound"); None }; @@ -553,6 +588,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { if let Some(task) = telemetry_task { let _ = task.await; } + let _ = background_check_task.await; tracing::info!("aisix shut down cleanly"); Ok(()) } @@ -783,6 +819,19 @@ fn build_hub() -> Hub { hub } +fn background_check_interval(snapshot: &aisix_core::AisixSnapshot) -> std::time::Duration { + let min_interval = snapshot + .models + .entries() + .into_iter() + .filter_map(|entry| entry.value.background_model_check.clone()) + .filter(|cfg| cfg.enabled) + .map(|cfg| cfg.interval_seconds) + .min() + .unwrap_or(1); + std::time::Duration::from_secs(min_interval.max(1)) +} + /// Completes when the process receives SIGINT or SIGTERM (best-effort on /// Windows — Ctrl+C only) OR when another part of the system has already /// flipped the cancel channel. diff --git a/tests/e2e/src/harness/admin.ts b/tests/e2e/src/harness/admin.ts index 9580d1e5..bae457e2 100644 --- a/tests/e2e/src/harness/admin.ts +++ b/tests/e2e/src/harness/admin.ts @@ -36,6 +36,10 @@ export class AdminClient { return (res.items ?? []).map((entry) => entry.value); } + async listModelStatuses(): Promise>> { + return this.json>>("GET", "/admin/v1/models/status"); + } + async json>( method: string, path: string, From 6b109ed591477121b0fe164166ab138efbf2b5e4 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Thu, 14 May 2026 08:08:36 +0800 Subject: [PATCH 2/4] fix(cooldown): decouple from is_retryable, honor Retry-After, fail fast on all-filtered MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses HIGH/MEDIUM findings from issue #264 audit on PR #268. HIGH: - H1: cooldown trigger is now a configurable status list (default [401, 408, 429, 500, 502, 503, 504]) independent of is_retryable. 401 / 404 / 408 cool down even though they are non-retryable: the same target will keep failing on subsequent requests, so the routing filter should skip it. - H2: upstream Retry-After header (seconds form) is parsed by the OpenAI / Anthropic bridges and threaded through BridgeError::UpstreamStatus.retry_after. The cooldown layer honors it (clamped by max_seconds, default 600s). - H3: filter_attempt_models returns an AllUnhealthy outcome when every candidate is excluded; the dispatch loop converts it into a new ProxyError::AllCandidatesUnavailable that maps to 503 + Retry-After. Routing models can opt into the legacy "send to known-bad" behavior via `on_all_filtered: original_order`. MEDIUM: - M1: 429 always cools down regardless of retry_on_429. Retry (same request) and cooldown (next request) are independent layers. - M3: background checks run under a 4-permit semaphore so 100 direct models cannot fan out 100 concurrent provider probes; schema enforces interval_seconds >= 5. New per-direct-model `cooldown` config block (all fields optional): enabled / default_seconds / max_seconds / honor_retry_after / trigger_statuses / trigger_on_timeout / trigger_on_transport. New per-routing-model `on_all_filtered` policy: `fail` (default) or `original_order`. E2E: - 4 baseline files (runtime-status, background-health, runtime-mixed-filtering, retry-on-429-vs-background-ignore) that were claimed in the original PR description but never committed. Now committed and updated for the new contract (status_reason names, interval_seconds=5 floor). - 1 new file `cooldown-contract-e2e.test.ts` with 5 contract tests: H1 401 cooldown, M1 429-cools-down-without-retry, H2 Retry-After-drives-TTL, H3 fail-fast 503, H3 original_order escape hatch. - 2 pre-existing tests (fallback-e2e, routing-strategies-e2e) fixed to avoid warming cooldown in their readiness probes — use bare waitConfigPropagation() so per-target hit counts measure the test request, not the probe. Rust unit coverage: 22 new chat.rs tests for decide_cooldown + filter_attempt_models (covers default policy, override, Retry-After clamping, all-filtered branching, on_all_filtered policy enforcement). cargo test --workspace: 686 pass / 0 fail cargo clippy --workspace --all-targets -- -D warnings: clean cargo fmt --all -- --check: clean e2e: 92 pass / 0 fail across 47 files --- crates/aisix-admin/src/lib.rs | 4 +- crates/aisix-core/src/lib.rs | 7 +- crates/aisix-core/src/models/mod.rs | 6 +- crates/aisix-core/src/models/model.rs | 167 +++++ crates/aisix-core/src/models/routing.rs | 56 ++ crates/aisix-core/src/models/schema.rs | 158 +++- crates/aisix-gateway/src/bridge.rs | 112 ++- crates/aisix-gateway/src/lib.rs | 2 +- crates/aisix-provider-anthropic/src/bridge.rs | 14 +- crates/aisix-provider-openai/src/bridge.rs | 73 +- crates/aisix-proxy/src/audio.rs | 20 +- crates/aisix-proxy/src/background.rs | 125 +++- crates/aisix-proxy/src/chat.rs | 429 ++++++++++- crates/aisix-proxy/src/error.rs | 35 +- crates/aisix-proxy/src/health.rs | 33 +- crates/aisix-proxy/src/lib.rs | 68 +- crates/aisix-proxy/src/messages.rs | 19 +- crates/aisix-proxy/src/rerank.rs | 10 +- crates/aisix-proxy/src/responses.rs | 10 +- crates/aisix-proxy/src/routing.rs | 44 +- .../src/cases/background-health-e2e.test.ts | 264 +++++++ .../src/cases/cooldown-contract-e2e.test.ts | 706 ++++++++++++++++++ tests/e2e/src/cases/fallback-e2e.test.ts | 7 + ...ry-on-429-vs-background-ignore-e2e.test.ts | 157 ++++ .../src/cases/routing-strategies-e2e.test.ts | 49 +- .../cases/runtime-mixed-filtering-e2e.test.ts | 187 +++++ .../e2e/src/cases/runtime-status-e2e.test.ts | 186 +++++ tests/e2e/src/harness/upstream-openai.ts | 13 + 28 files changed, 2775 insertions(+), 186 deletions(-) create mode 100644 tests/e2e/src/cases/background-health-e2e.test.ts create mode 100644 tests/e2e/src/cases/cooldown-contract-e2e.test.ts create mode 100644 tests/e2e/src/cases/retry-on-429-vs-background-ignore-e2e.test.ts create mode 100644 tests/e2e/src/cases/runtime-mixed-filtering-e2e.test.ts create mode 100644 tests/e2e/src/cases/runtime-status-e2e.test.ts diff --git a/crates/aisix-admin/src/lib.rs b/crates/aisix-admin/src/lib.rs index 50e7f98f..40871446 100644 --- a/crates/aisix-admin/src/lib.rs +++ b/crates/aisix-admin/src/lib.rs @@ -1167,7 +1167,9 @@ mod tests { "provider_key_id": "11111111-1111-1111-1111-111111111111", "background_model_check": { "enabled": true, - "interval_seconds": 1, + // Minimum interval is 5s in schema; using 30 to + // mirror a realistic operator config. + "interval_seconds": 30, "timeout_seconds": 10, "prompt": "Respond with OK", "max_tokens": 8, diff --git a/crates/aisix-core/src/lib.rs b/crates/aisix-core/src/lib.rs index a03be886..e75bbde1 100644 --- a/crates/aisix-core/src/lib.rs +++ b/crates/aisix-core/src/lib.rs @@ -32,9 +32,10 @@ pub use error::{ pub use models::{ validate_apikey, validate_cache_policy, validate_guardrail, validate_model, validate_observability_exporter, validate_provider_key, AisixSnapshot, ApiKey, CachePolicy, - ExporterKind, Guardrail, GuardrailHookPoint, GuardrailKind, KeywordConfig, KeywordPattern, - Model, ObservabilityExporter, Provider, ProviderKey, RateLimit, Routing, RoutingStrategy, - RoutingTarget, SchemaError, + CooldownConfig, ExporterKind, Guardrail, GuardrailHookPoint, GuardrailKind, KeywordConfig, + KeywordPattern, Model, ObservabilityExporter, OnAllFilteredPolicy, Provider, ProviderKey, + RateLimit, Routing, RoutingStrategy, RoutingTarget, SchemaError, + DEFAULT_COOLDOWN_TRIGGER_STATUSES, }; pub use resource::{Resource, ResourceEntry}; pub use snapshot::{ResourceTable, SnapshotHandle}; diff --git a/crates/aisix-core/src/models/mod.rs b/crates/aisix-core/src/models/mod.rs index ea2c58cd..6d494456 100644 --- a/crates/aisix-core/src/models/mod.rs +++ b/crates/aisix-core/src/models/mod.rs @@ -32,11 +32,13 @@ pub use guardrail::{ BedrockAWSCredentials, BedrockConfig, BedrockLatencyMode, Guardrail, GuardrailHookPoint, GuardrailKind, KeywordConfig, KeywordPattern, }; -pub use model::{BackgroundModelCheck, Model, Provider}; +pub use model::{ + BackgroundModelCheck, CooldownConfig, Model, Provider, DEFAULT_COOLDOWN_TRIGGER_STATUSES, +}; pub use observability_exporter::{ExporterKind, ObservabilityExporter, OtlpHttpConfig}; pub use provider_key::ProviderKey; pub use rate_limit::RateLimit; -pub use routing::{Routing, RoutingStrategy, RoutingTarget}; +pub use routing::{OnAllFilteredPolicy, Routing, RoutingStrategy, RoutingTarget}; pub use schema::{ validate_apikey, validate_cache_policy, validate_guardrail, validate_model, validate_observability_exporter, validate_provider_key, SchemaError, diff --git a/crates/aisix-core/src/models/model.rs b/crates/aisix-core/src/models/model.rs index eb3df854..e1e7793c 100644 --- a/crates/aisix-core/src/models/model.rs +++ b/crates/aisix-core/src/models/model.rs @@ -94,6 +94,98 @@ pub struct BackgroundModelCheck { pub stale_after_seconds: u64, } +/// Request-path cooldown configuration for a direct model. Controls +/// which upstream failures temporarily exclude this model from routing +/// candidate selection, and for how long. +/// +/// Cooldown is **independent** of request retry semantics — i.e. +/// `Routing.retry_on_429` governs whether a 429 is retried within the +/// current request, but `CooldownConfig.trigger_statuses` governs +/// whether 429 takes the model out of rotation for subsequent +/// requests. The two layers serve different purposes: +/// - retry: short-window in-request recovery +/// - cooldown: medium-window cross-request backpressure +/// +/// All fields are optional; defaults preserve a safe behavior for any +/// direct model that doesn't ship a `cooldown` block. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)] +#[serde(deny_unknown_fields)] +pub struct CooldownConfig { + /// Whether cooldown is active for this model. Default: true. + /// Set to `false` to disable cooldown entirely (the model stays in + /// rotation regardless of upstream failures). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub enabled: Option, + /// Cooldown TTL in seconds when the upstream did not supply a + /// `Retry-After` header (or `honor_retry_after=false`). Default: 30. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub default_seconds: Option, + /// Upper bound on cooldown TTL. Caps a misbehaving upstream that + /// returns an unreasonable `Retry-After` value. Default: 600 (10 min). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_seconds: Option, + /// Whether to use the upstream's `Retry-After` header (seconds form) + /// as the cooldown TTL when present. Default: true. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub honor_retry_after: Option, + /// Status codes that trigger cooldown. Default: + /// `[401, 408, 429, 500, 502, 503, 504]` — auth failures and rate + /// limits + transient server errors. `400/403/422` etc. are caller + /// mistakes and intentionally excluded. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trigger_statuses: Option>, + /// Whether request-path timeouts trigger cooldown. Default: true. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trigger_on_timeout: Option, + /// Whether transport / decode / stream-abort errors trigger + /// cooldown. Default: true. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub trigger_on_transport: Option, +} + +/// Default cooldown trigger statuses applied when the operator does +/// not override `trigger_statuses` on a direct model. +pub const DEFAULT_COOLDOWN_TRIGGER_STATUSES: &[u16] = &[401, 408, 429, 500, 502, 503, 504]; + +const DEFAULT_COOLDOWN_SECONDS: u64 = 30; +const DEFAULT_COOLDOWN_MAX_SECONDS: u64 = 600; + +impl CooldownConfig { + pub fn enabled_or_default(&self) -> bool { + self.enabled.unwrap_or(true) + } + + pub fn default_seconds_or_default(&self) -> u64 { + self.default_seconds.unwrap_or(DEFAULT_COOLDOWN_SECONDS) + } + + pub fn max_seconds_or_default(&self) -> u64 { + self.max_seconds.unwrap_or(DEFAULT_COOLDOWN_MAX_SECONDS) + } + + pub fn honor_retry_after_or_default(&self) -> bool { + self.honor_retry_after.unwrap_or(true) + } + + /// Effective trigger-status list — operator override OR built-in + /// default. Returned as `Cow` so callers can avoid copies on the + /// default path. + pub fn effective_trigger_statuses(&self) -> std::borrow::Cow<'_, [u16]> { + match &self.trigger_statuses { + Some(list) => std::borrow::Cow::Borrowed(list.as_slice()), + None => std::borrow::Cow::Borrowed(DEFAULT_COOLDOWN_TRIGGER_STATUSES), + } + } + + pub fn trigger_on_timeout_or_default(&self) -> bool { + self.trigger_on_timeout.unwrap_or(true) + } + + pub fn trigger_on_transport_or_default(&self) -> bool { + self.trigger_on_transport.unwrap_or(true) + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct Model { @@ -141,6 +233,12 @@ pub struct Model { #[serde(default, skip_serializing_if = "Option::is_none")] pub background_model_check: Option, + /// Optional direct-model-only request-path cooldown configuration. + /// When absent, default cooldown semantics apply (see + /// [`CooldownConfig`] field docs for defaults). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cooldown: Option, + /// Non-schema runtime id. Not part of the JSON payload — filled in by /// the snapshot loader from the etcd key path. Kept here so `Resource` /// can return a `&str` id. @@ -244,6 +342,75 @@ mod tests { assert_eq!(m.name(), "my-gpt4"); } + #[test] + fn cooldown_config_defaults_via_helpers() { + let cfg = CooldownConfig::default(); + assert!(cfg.enabled_or_default()); + assert_eq!(cfg.default_seconds_or_default(), 30); + assert_eq!(cfg.max_seconds_or_default(), 600); + assert!(cfg.honor_retry_after_or_default()); + assert_eq!( + cfg.effective_trigger_statuses().as_ref(), + DEFAULT_COOLDOWN_TRIGGER_STATUSES, + ); + assert!(cfg.trigger_on_timeout_or_default()); + assert!(cfg.trigger_on_transport_or_default()); + } + + #[test] + fn cooldown_default_trigger_statuses_match_advertised_set() { + // Lock the documented default so a future change has to update + // both the constant and the test, surfaced as one diff. + assert_eq!( + DEFAULT_COOLDOWN_TRIGGER_STATUSES, + &[401, 408, 429, 500, 502, 503, 504] + ); + } + + #[test] + fn cooldown_config_partial_override_keeps_other_defaults() { + let cfg: CooldownConfig = serde_json::from_str(r#"{"default_seconds": 90}"#).unwrap(); + assert_eq!(cfg.default_seconds_or_default(), 90); + // Other fields fall back to defaults. + assert!(cfg.enabled_or_default()); + assert_eq!(cfg.max_seconds_or_default(), 600); + assert!(cfg.honor_retry_after_or_default()); + } + + #[test] + fn cooldown_config_disable_via_enabled_false() { + let cfg: CooldownConfig = serde_json::from_str(r#"{"enabled": false}"#).unwrap(); + assert!(!cfg.enabled_or_default()); + } + + #[test] + fn cooldown_config_override_trigger_statuses() { + let cfg: CooldownConfig = serde_json::from_str(r#"{"trigger_statuses": [503]}"#).unwrap(); + assert_eq!(cfg.effective_trigger_statuses().as_ref(), &[503]); + } + + #[test] + fn direct_model_can_deserialize_cooldown_config() { + let m: Model = serde_json::from_str( + r#"{ + "display_name": "my-gpt4", + "provider": "openai", + "model_name": "gpt-4o", + "provider_key_id": "11111111-1111-1111-1111-111111111111", + "cooldown": { + "enabled": true, + "default_seconds": 45, + "trigger_statuses": [429, 503] + } + }"#, + ) + .unwrap(); + let cooldown = m.cooldown.unwrap(); + assert!(cooldown.enabled_or_default()); + assert_eq!(cooldown.default_seconds_or_default(), 45); + assert_eq!(cooldown.effective_trigger_statuses().as_ref(), &[429, 503]); + } + #[test] fn direct_model_can_deserialize_background_check() { let m: Model = serde_json::from_str( diff --git a/crates/aisix-core/src/models/routing.rs b/crates/aisix-core/src/models/routing.rs index d4c39459..42b91960 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -56,6 +56,27 @@ impl RoutingTarget { } } +/// Behavior when every candidate target is filtered out by the +/// runtime status layer (all in cooldown or background-unhealthy). +/// +/// `Fail` is the default because sending traffic to a target we know +/// is currently bad — just because every other target is also bad — +/// amplifies cascading outages. Operators that prefer the legacy +/// behavior (try every candidate regardless of known state) can opt +/// into `OriginalOrder` per routing model. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum OnAllFilteredPolicy { + /// Return 503 with a Retry-After hint derived from the next + /// cooldown expiry (or a fixed fallback if every candidate is + /// background-unhealthy with no cooldown timer). Default. + #[default] + Fail, + /// Send to the original candidate list anyway, in declaration + /// order. Preserves availability over caller-facing correctness. + OriginalOrder, +} + #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct Routing { @@ -72,6 +93,10 @@ pub struct Routing { /// Whether upstream 429 participates in retries and failover. #[serde(default, skip_serializing_if = "Option::is_none")] pub retry_on_429: Option, + /// Policy for the case where every candidate is filtered out by + /// runtime status. See [`OnAllFilteredPolicy`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub on_all_filtered: Option, } impl Routing { @@ -91,6 +116,10 @@ impl Routing { self.retry_on_429.unwrap_or(false) } + pub fn on_all_filtered_or_default(&self) -> OnAllFilteredPolicy { + self.on_all_filtered.unwrap_or_default() + } + pub fn is_empty(&self) -> bool { self.targets.is_empty() } @@ -140,6 +169,7 @@ mod tests { retries: Some(0), max_fallbacks: Some(0), retry_on_429: None, + on_all_filtered: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } @@ -152,10 +182,36 @@ mod tests { retries: None, max_fallbacks: Some(99), retry_on_429: None, + on_all_filtered: None, }; assert_eq!(r.max_fallbacks_or_default(), 0); } + #[test] + fn on_all_filtered_defaults_to_fail() { + let r: Routing = serde_json::from_str(r#"{"targets":[{"model":"a"}]}"#).unwrap(); + assert_eq!(r.on_all_filtered_or_default(), OnAllFilteredPolicy::Fail); + } + + #[test] + fn on_all_filtered_parses_original_order() { + let r: Routing = serde_json::from_str( + r#"{"targets":[{"model":"a"}],"on_all_filtered":"original_order"}"#, + ) + .unwrap(); + assert_eq!( + r.on_all_filtered_or_default(), + OnAllFilteredPolicy::OriginalOrder + ); + } + + #[test] + fn on_all_filtered_rejects_unknown_value() { + let r: Result = + serde_json::from_str(r#"{"targets":[{"model":"a"}],"on_all_filtered":"explode"}"#); + assert!(r.is_err()); + } + #[test] fn missing_weight_defaults_to_one() { let t = RoutingTarget::new("x"); diff --git a/crates/aisix-core/src/models/schema.rs b/crates/aisix-core/src/models/schema.rs index 9a57bd81..eb89c20e 100644 --- a/crates/aisix-core/src/models/schema.rs +++ b/crates/aisix-core/src/models/schema.rs @@ -138,7 +138,11 @@ fn model_schema() -> Value { }, "retries": { "type": "integer", "minimum": 0 }, "max_fallbacks": { "type": "integer", "minimum": 0 }, - "retry_on_429": { "type": "boolean" } + "retry_on_429": { "type": "boolean" }, + "on_all_filtered": { + "type": "string", + "enum": ["fail", "original_order"] + } } }, "cost": { @@ -163,7 +167,10 @@ fn model_schema() -> Value { "additionalProperties": false, "properties": { "enabled": { "type": "boolean" }, - "interval_seconds": { "type": "integer", "minimum": 1 }, + // Minimum 5s guards against misconfiguration. Setting + // interval_seconds=1 with multiple direct models would + // burn provider quota and money very quickly. + "interval_seconds": { "type": "integer", "minimum": 5 }, "timeout_seconds": { "type": "integer", "minimum": 1 }, "prompt": { "type": "string", "minLength": 1 }, "max_tokens": { "type": "integer", "minimum": 1 }, @@ -173,6 +180,22 @@ fn model_schema() -> Value { }, "stale_after_seconds": { "type": "integer", "minimum": 1 } } + }, + "cooldown": { + "type": "object", + "additionalProperties": false, + "properties": { + "enabled": { "type": "boolean" }, + "default_seconds": { "type": "integer", "minimum": 0 }, + "max_seconds": { "type": "integer", "minimum": 1 }, + "honor_retry_after": { "type": "boolean" }, + "trigger_statuses": { + "type": "array", + "items": { "type": "integer", "minimum": 100, "maximum": 599 } + }, + "trigger_on_timeout": { "type": "boolean" }, + "trigger_on_transport": { "type": "boolean" } + } } }, // Direct vs routing model: a model EITHER ships a `routing` @@ -186,7 +209,8 @@ fn model_schema() -> Value { { "required": ["provider"] }, { "required": ["model_name"] }, { "required": ["provider_key_id"] }, - { "required": ["background_model_check"] } + { "required": ["background_model_check"] }, + { "required": ["cooldown"] } ]} }, { @@ -588,6 +612,134 @@ mod tests { assert!(validate_model(&v).is_err()); } + #[test] + fn direct_model_cooldown_block_passes() { + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "cooldown": { + "enabled": true, + "default_seconds": 30, + "max_seconds": 600, + "honor_retry_after": true, + "trigger_statuses": [401, 408, 429, 500, 502, 503, 504], + "trigger_on_timeout": true, + "trigger_on_transport": true + } + }); + validate_model(&v).unwrap(); + } + + #[test] + fn cooldown_block_partial_override_passes() { + // Only set one field — defaults fill the rest at runtime. + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "cooldown": { + "default_seconds": 90 + } + }); + validate_model(&v).unwrap(); + } + + #[test] + fn routing_model_cooldown_block_fails() { + // Cooldown is direct-model-only — routing models project to + // their underlying targets and have no upstream of their own. + let v = json!({ + "display_name": "router-1", + "routing": { "targets": [{"model": "x"}] }, + "cooldown": { "default_seconds": 30 } + }); + assert!(validate_model(&v).is_err()); + } + + #[test] + fn cooldown_rejects_invalid_status_code() { + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "cooldown": { "trigger_statuses": [99] } + }); + assert!(validate_model(&v).is_err()); + } + + #[test] + fn cooldown_max_seconds_must_be_positive() { + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "cooldown": { "max_seconds": 0 } + }); + assert!(validate_model(&v).is_err()); + } + + #[test] + fn routing_on_all_filtered_fail_passes() { + let v = json!({ + "display_name": "router-1", + "routing": { + "targets": [{"model": "a"}], + "on_all_filtered": "fail" + } + }); + validate_model(&v).unwrap(); + } + + #[test] + fn routing_on_all_filtered_original_order_passes() { + let v = json!({ + "display_name": "router-1", + "routing": { + "targets": [{"model": "a"}], + "on_all_filtered": "original_order" + } + }); + validate_model(&v).unwrap(); + } + + #[test] + fn routing_on_all_filtered_rejects_unknown_value() { + let v = json!({ + "display_name": "router-1", + "routing": { + "targets": [{"model": "a"}], + "on_all_filtered": "yolo" + } + }); + assert!(validate_model(&v).is_err()); + } + + #[test] + fn background_check_interval_below_min_fails() { + // Minimum interval is 5s — guards misconfiguration from + // burning provider quota on a 1s loop. + let v = json!({ + "display_name": "x", + "provider": "openai", + "model_name": "g", + "provider_key_id": "pk-1", + "background_model_check": { + "enabled": true, + "interval_seconds": 1, + "timeout_seconds": 10, + "prompt": "Respond with OK", + "max_tokens": 8, + "stale_after_seconds": 90 + } + }); + assert!(validate_model(&v).is_err()); + } + #[test] fn background_check_rejects_invalid_ignore_status() { let v = json!({ diff --git a/crates/aisix-gateway/src/bridge.rs b/crates/aisix-gateway/src/bridge.rs index 61011cd4..db5996c7 100644 --- a/crates/aisix-gateway/src/bridge.rs +++ b/crates/aisix-gateway/src/bridge.rs @@ -71,8 +71,18 @@ impl BridgeContext { pub enum BridgeError { #[error("upstream request timed out after {elapsed_ms}ms")] Timeout { elapsed_ms: u64 }, + /// Upstream returned a non-2xx HTTP status. `retry_after` carries + /// the upstream's `Retry-After` header parsed to a Duration when + /// present — used by the cooldown layer to honor provider-supplied + /// backoff hints. Bridges that cannot parse the header (or where + /// the header is absent) leave this `None`; the cooldown layer + /// falls back to its configured default in that case. #[error("upstream returned HTTP {status}: {message}")] - UpstreamStatus { status: u16, message: String }, + UpstreamStatus { + status: u16, + message: String, + retry_after: Option, + }, #[error("upstream returned an unparseable body: {0}")] UpstreamDecode(String), #[error("bridge is misconfigured: {0}")] @@ -83,6 +93,52 @@ pub enum BridgeError { StreamAborted, } +impl BridgeError { + /// Convenience constructor for upstream status errors when no + /// `Retry-After` is available. Keeps existing call sites readable. + pub fn upstream_status(status: u16, message: impl Into) -> Self { + Self::UpstreamStatus { + status, + message: message.into(), + retry_after: None, + } + } + + /// Convenience constructor for upstream status errors that carry + /// a parsed `Retry-After` hint. + pub fn upstream_status_with_retry_after( + status: u16, + message: impl Into, + retry_after: Option, + ) -> Self { + Self::UpstreamStatus { + status, + message: message.into(), + retry_after, + } + } +} + +/// Parse the `Retry-After` response header into a Duration. +/// +/// Per RFC 9110 §10.2.3, `Retry-After` may be either: +/// - a non-negative integer number of seconds, or +/// - an HTTP-date. +/// +/// We accept the seconds form (which is what OpenAI / Anthropic / +/// DeepSeek / Gemini all return on 429). The HTTP-date form is rare +/// for AI providers and parsing it pulls in `httpdate`; skip for V1 +/// — callers fall back to the configured default cooldown TTL. +/// +/// Returns `None` when the header is absent, unparseable, or the +/// seconds value is unreasonable (the cooldown layer applies a +/// `max_seconds` clamp regardless). +pub fn parse_retry_after(headers: &http::HeaderMap) -> Option { + let raw = headers.get(http::header::RETRY_AFTER)?.to_str().ok()?; + let seconds: u64 = raw.trim().parse().ok()?; + Some(Duration::from_secs(seconds)) +} + impl BridgeError { /// Stable HTTP status mapping. The proxy layer uses this to build /// its OpenAI-compatible `{error:{message,type,...}}` envelope. @@ -209,26 +265,56 @@ mod tests { #[test] fn upstream_4xx_passes_through_5xx_collapses_to_502() { - let e400 = BridgeError::UpstreamStatus { - status: 429, - message: "rate limit".into(), - }; + let e400 = BridgeError::upstream_status(429, "rate limit"); assert_eq!(e400.http_status(), 429); - let e500 = BridgeError::UpstreamStatus { - status: 503, - message: "busy".into(), - }; + let e500 = BridgeError::upstream_status(503, "busy"); assert_eq!(e500.http_status(), 502); - let e3xx = BridgeError::UpstreamStatus { - status: 301, - message: "redirect".into(), - }; + let e3xx = BridgeError::upstream_status(301, "redirect"); // Non-4xx collapses too — redirects we don't follow are 502-worthy. assert_eq!(e3xx.http_status(), 502); } + #[test] + fn upstream_status_carries_retry_after_when_provided() { + let e = BridgeError::upstream_status_with_retry_after( + 429, + "slow down", + Some(Duration::from_secs(60)), + ); + match e { + BridgeError::UpstreamStatus { retry_after, .. } => { + assert_eq!(retry_after, Some(Duration::from_secs(60))); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[test] + fn parse_retry_after_handles_seconds_form() { + let mut h = http::HeaderMap::new(); + h.insert(http::header::RETRY_AFTER, "30".parse().unwrap()); + assert_eq!(parse_retry_after(&h), Some(Duration::from_secs(30))); + } + + #[test] + fn parse_retry_after_returns_none_for_http_date_form() { + let mut h = http::HeaderMap::new(); + h.insert( + http::header::RETRY_AFTER, + "Wed, 21 Oct 2026 07:28:00 GMT".parse().unwrap(), + ); + // V1: HTTP-date form is intentionally not parsed. + assert_eq!(parse_retry_after(&h), None); + } + + #[test] + fn parse_retry_after_returns_none_when_absent() { + let h = http::HeaderMap::new(); + assert_eq!(parse_retry_after(&h), None); + } + #[test] fn transport_and_decode_errors_collapse_to_502() { assert_eq!( diff --git a/crates/aisix-gateway/src/lib.rs b/crates/aisix-gateway/src/lib.rs index c147952f..7e19232e 100644 --- a/crates/aisix-gateway/src/lib.rs +++ b/crates/aisix-gateway/src/lib.rs @@ -25,7 +25,7 @@ pub mod chat; pub mod hub; pub mod sse; -pub use bridge::{Bridge, BridgeContext, BridgeError, ChatChunkStream}; +pub use bridge::{parse_retry_after, Bridge, BridgeContext, BridgeError, ChatChunkStream}; pub use chat::{ ChatChunk, ChatDelta, ChatFormat, ChatMessage, ChatResponse, EmbeddingObject, EmbeddingRequest, EmbeddingResponse, EmbeddingUsage, FinishReason, Role, UsageStats, diff --git a/crates/aisix-provider-anthropic/src/bridge.rs b/crates/aisix-provider-anthropic/src/bridge.rs index 53e55e83..44882aed 100644 --- a/crates/aisix-provider-anthropic/src/bridge.rs +++ b/crates/aisix-provider-anthropic/src/bridge.rs @@ -105,11 +105,13 @@ fn upstream_model(ctx: &BridgeContext) -> Result<&str, BridgeError> { } async fn map_http_error(status: StatusCode, resp: reqwest::Response) -> BridgeError { + let retry_after = aisix_gateway::parse_retry_after(resp.headers()); let message = resp.text().await.unwrap_or_default(); - BridgeError::UpstreamStatus { - status: status.as_u16(), - message: truncate(&message, 1024), - } + BridgeError::upstream_status_with_retry_after( + status.as_u16(), + truncate(&message, 1024), + retry_after, + ) } fn truncate(s: &str, n: usize) -> String { @@ -354,7 +356,9 @@ mod tests { let ctx = sample_ctx(&server.uri()); let err = bridge.chat(&req(), &ctx).await.unwrap_err(); match err { - BridgeError::UpstreamStatus { status, message } => { + BridgeError::UpstreamStatus { + status, message, .. + } => { assert_eq!(status, 400); assert!(message.contains("invalid_request")); } diff --git a/crates/aisix-provider-openai/src/bridge.rs b/crates/aisix-provider-openai/src/bridge.rs index 73ab2a24..545d9746 100644 --- a/crates/aisix-provider-openai/src/bridge.rs +++ b/crates/aisix-provider-openai/src/bridge.rs @@ -102,11 +102,13 @@ fn upstream_model(ctx: &BridgeContext) -> Result<&str, BridgeError> { } async fn map_http_error(status: StatusCode, resp: reqwest::Response) -> BridgeError { + let retry_after = aisix_gateway::parse_retry_after(resp.headers()); let message = resp.text().await.unwrap_or_default(); - BridgeError::UpstreamStatus { - status: status.as_u16(), - message: truncate(&message, 1024), - } + BridgeError::upstream_status_with_retry_after( + status.as_u16(), + truncate(&message, 1024), + retry_after, + ) } fn truncate(s: &str, n: usize) -> String { @@ -468,7 +470,9 @@ mod tests { let ctx = sample_ctx(&server.uri()); let err = bridge.chat(&req(), &ctx).await.unwrap_err(); match err { - BridgeError::UpstreamStatus { status, message } => { + BridgeError::UpstreamStatus { + status, message, .. + } => { assert_eq!(status, 429); assert!(message.contains("slow down")); } @@ -476,6 +480,65 @@ mod tests { } } + #[tokio::test] + async fn non_streaming_429_surfaces_retry_after_header() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "42") + .set_body_string("slow down"), + ) + .mount(&server) + .await; + + let bridge = OpenAiBridge::new(); + let ctx = sample_ctx(&server.uri()); + let err = bridge.chat(&req(), &ctx).await.unwrap_err(); + match err { + BridgeError::UpstreamStatus { + status, + retry_after, + .. + } => { + assert_eq!(status, 429); + assert_eq!(retry_after, Some(Duration::from_secs(42))); + } + other => panic!("unexpected: {other:?}"), + } + } + + #[tokio::test] + async fn non_streaming_503_with_garbled_retry_after_is_none() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(503) + .insert_header("retry-after", "Wed, 21 Oct 2026 07:28:00 GMT") + .set_body_string("down"), + ) + .mount(&server) + .await; + + let bridge = OpenAiBridge::new(); + let ctx = sample_ctx(&server.uri()); + let err = bridge.chat(&req(), &ctx).await.unwrap_err(); + match err { + BridgeError::UpstreamStatus { + status, + retry_after, + .. + } => { + assert_eq!(status, 503); + // HTTP-date form is intentionally not parsed in V1. + assert!(retry_after.is_none()); + } + other => panic!("unexpected: {other:?}"), + } + } + #[tokio::test] async fn non_streaming_decode_error_on_malformed_body() { let server = MockServer::start().await; diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 9c0ea0d3..f00add8b 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -341,12 +341,14 @@ async fn multipart_dispatch( let status = resp.status(); if !status.is_success() { let s = status.as_u16(); + let retry_after = aisix_gateway::parse_retry_after(resp.headers()); let msg = resp.text().await.unwrap_or_default(); return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::UpstreamStatus { - status: s, - message: msg.chars().take(1024).collect(), - }, + aisix_gateway::BridgeError::upstream_status_with_retry_after( + s, + msg.chars().take(1024).collect::(), + retry_after, + ), )); } @@ -421,12 +423,14 @@ async fn speech_dispatch( let status = resp.status(); if !status.is_success() { let s = status.as_u16(); + let retry_after = aisix_gateway::parse_retry_after(resp.headers()); let msg = resp.text().await.unwrap_or_default(); return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::UpstreamStatus { - status: s, - message: msg.chars().take(1024).collect(), - }, + aisix_gateway::BridgeError::upstream_status_with_retry_after( + s, + msg.chars().take(1024).collect::(), + retry_after, + ), )); } diff --git a/crates/aisix-proxy/src/background.rs b/crates/aisix-proxy/src/background.rs index 3c389b14..d3dee250 100644 --- a/crates/aisix-proxy/src/background.rs +++ b/crates/aisix-proxy/src/background.rs @@ -4,16 +4,35 @@ use std::time::Duration; use aisix_core::models::BackgroundModelCheck; use aisix_core::{AisixSnapshot, Model}; use aisix_gateway::{BridgeContext, BridgeError, ChatFormat, ChatMessage, Hub}; +use tokio::sync::Semaphore; use crate::dispatch; use crate::health::ModelRuntimeStatusTracker; +/// Cap on the number of background model checks that may run +/// concurrently across all configured direct models. Each check +/// issues a real chat completion against the upstream provider — +/// burning the operator's quota and dollars — so we serialize them +/// to keep the cost bounded regardless of how many direct models +/// the operator has registered. +/// +/// Rationale: a deployment with 100 direct models all configured +/// with the same `interval_seconds` would otherwise fan out 100 +/// concurrent requests to upstream providers every interval. The +/// semaphore turns that into a slow trickle of ≤4 in-flight checks +/// at any time. The total cost-per-interval is unchanged; the +/// burstiness (and the chance of self-induced 429 on small +/// accounts) is dampened. +const MAX_CONCURRENT_BACKGROUND_CHECKS: usize = 4; + pub async fn run_background_model_check_once( snapshot: Arc, hub: Arc, tracker: Arc, request_id: &str, ) { + let semaphore = Arc::new(Semaphore::new(MAX_CONCURRENT_BACKGROUND_CHECKS)); + let mut tasks = Vec::new(); for entry in snapshot.models.entries() { let model = &entry.value; if model.is_routing() { @@ -25,20 +44,52 @@ pub async fn run_background_model_check_once( if !cfg.enabled { continue; } - let outcome = check_direct_model(&snapshot, &hub, &entry.id, model, cfg, request_id).await; - match outcome { - Ok(()) => tracker.clear_unhealthy(&entry.id), - Err(BridgeError::UpstreamStatus { status, .. }) if cfg.ignore_statuses.contains(&status) => { - tracker.record_ignored_check(&entry.id, status, "ignored_transient_error") - } - Err(BridgeError::Timeout { .. }) if cfg.ignore_statuses.contains(&408) => { - tracker.record_ignored_check(&entry.id, 408, "ignored_transient_error") - } - Err(err) => tracker.mark_unhealthy( - &entry.id, - background_status_code(&err), - "background_check_failed", - ), + let id = entry.id.clone(); + let model = model.clone(); + let cfg = cfg.clone(); + let snapshot = Arc::clone(&snapshot); + let hub = Arc::clone(&hub); + let tracker = Arc::clone(&tracker); + let request_id = request_id.to_string(); + let permits = Arc::clone(&semaphore); + tasks.push(tokio::spawn(async move { + // Acquire concurrency permit. If the semaphore is closed + // (only happens during shutdown), the check is skipped + // and the tracker stays at last-known state — fine. + let _permit = match permits.acquire_owned().await { + Ok(p) => p, + Err(_) => return, + }; + run_one(snapshot, hub, tracker, &id, &model, &cfg, &request_id).await; + })); + } + for task in tasks { + let _ = task.await; + } +} + +async fn run_one( + snapshot: Arc, + hub: Arc, + tracker: Arc, + id: &str, + model: &Model, + cfg: &BackgroundModelCheck, + request_id: &str, +) { + let outcome = check_direct_model(&snapshot, &hub, id, model, cfg, request_id).await; + match outcome { + Ok(()) => tracker.clear_unhealthy(id), + Err(BridgeError::UpstreamStatus { status, .. }) + if cfg.ignore_statuses.contains(&status) => + { + tracker.record_ignored_check(id, status, "ignored_transient_error") + } + Err(BridgeError::Timeout { .. }) if cfg.ignore_statuses.contains(&408) => { + tracker.record_ignored_check(id, 408, "ignored_transient_error") + } + Err(err) => { + tracker.mark_unhealthy(id, background_status_code(&err), "background_check_failed") } } } @@ -51,7 +102,8 @@ async fn check_direct_model( cfg: &BackgroundModelCheck, request_id: &str, ) -> Result<(), BridgeError> { - let provider = dispatch::require_provider(model).map_err(|e| BridgeError::Config(e.to_string()))?; + let provider = + dispatch::require_provider(model).map_err(|e| BridgeError::Config(e.to_string()))?; let pk_entry = dispatch::resolve_provider_key(snapshot, model) .map_err(|e| BridgeError::Config(e.to_string()))?; let bridge = hub @@ -96,7 +148,11 @@ mod tests { use wiremock::{Mock, MockServer, ResponseTemplate}; fn openai_test_bridge() -> OpenAiBridge { - let client = Client::builder().user_agent("aisix-test/0.1").no_proxy().build().unwrap(); + let client = Client::builder() + .user_agent("aisix-test/0.1") + .no_proxy() + .build() + .unwrap(); OpenAiBridge::with_client(client) } @@ -108,7 +164,13 @@ mod tests { ResourceEntry::new(id, pk, 1) } - fn direct_model_entry(id: &str, name: &str, pk_id: &str, enabled: bool, ignore: &[u16]) -> ResourceEntry { + fn direct_model_entry( + id: &str, + name: &str, + pk_id: &str, + enabled: bool, + ignore: &[u16], + ) -> ResourceEntry { let cfg = serde_json::json!({ "display_name": name, "provider": "openai", @@ -140,8 +202,16 @@ mod tests { let hub = Arc::new(Hub::new()); hub.register(Provider::Openai, Arc::new(openai_test_bridge())); let snapshot = Arc::new(AisixSnapshot::new()); - snapshot.provider_keys.insert(provider_key_entry("pk-1", &upstream.uri())); - snapshot.models.insert(direct_model_entry("m-1", "bg-model", "pk-1", true, &[408, 429])); + snapshot + .provider_keys + .insert(provider_key_entry("pk-1", &upstream.uri())); + snapshot.models.insert(direct_model_entry( + "m-1", + "bg-model", + "pk-1", + true, + &[408, 429], + )); let tracker = Arc::new(ModelRuntimeStatusTracker::new()); run_background_model_check_once(snapshot, hub, tracker.clone(), "bg-check-1").await; @@ -163,8 +233,16 @@ mod tests { let hub = Arc::new(Hub::new()); hub.register(Provider::Openai, Arc::new(openai_test_bridge())); let snapshot = Arc::new(AisixSnapshot::new()); - snapshot.provider_keys.insert(provider_key_entry("pk-1", &upstream.uri())); - snapshot.models.insert(direct_model_entry("m-1", "bg-model", "pk-1", true, &[408, 429])); + snapshot + .provider_keys + .insert(provider_key_entry("pk-1", &upstream.uri())); + snapshot.models.insert(direct_model_entry( + "m-1", + "bg-model", + "pk-1", + true, + &[408, 429], + )); let tracker = Arc::new(ModelRuntimeStatusTracker::new()); run_background_model_check_once(snapshot, hub, tracker.clone(), "bg-check-1").await; @@ -172,6 +250,9 @@ mod tests { let status = tracker.status("m-1"); assert_eq!(status.status, crate::RuntimeStatus::Healthy); assert_eq!(status.last_check_status, Some(429)); - assert_eq!(status.status_reason.as_deref(), Some("ignored_transient_error")); + assert_eq!( + status.status_reason.as_deref(), + Some("ignored_transient_error") + ); } } diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 16d89e66..98613c7b 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -40,7 +40,13 @@ use crate::state::ProxyState; /// Header set on every non-streaming response indicating whether the /// response came from the cache (`hit`) or the upstream (`miss`). pub const CACHE_HEADER: &str = "x-aisix-cache"; -const RETRYABLE_FAILURE_COOLDOWN: Duration = Duration::from_secs(30); + +/// Default Retry-After (in seconds) returned to the client when every +/// candidate is background-unhealthy and no cooldown timer is available +/// to derive a more precise hint. Operators tune per-model cooldown +/// TTLs via `cooldown.default_seconds`; this is only the all-unhealthy +/// fallback for the `on_all_filtered: fail` path. +const FALLBACK_ALL_UNHEALTHY_RETRY_AFTER: Duration = Duration::from_secs(30); #[derive(Clone)] struct AttemptModel { @@ -48,6 +54,88 @@ struct AttemptModel { model: aisix_core::Model, } +/// Decide whether a bridge error should trigger cooldown on the +/// failing direct model, and for how long. +/// +/// Cooldown is **independent** of `is_retryable`. A 401 (auth failure) +/// is non-retryable — retrying the same target in the current request +/// is pointless — but it should still cooldown because every +/// subsequent request that lands on the same target will see the same +/// 401. Conversely, a transient timeout may be retryable AND should +/// also cooldown. +/// +/// `Retry-After` from upstream is honored when `honor_retry_after` +/// is set (default true), clamped to `max_seconds`. +fn decide_cooldown( + err: &BridgeError, + cfg: Option<&aisix_core::CooldownConfig>, +) -> Option<(Duration, &'static str)> { + use aisix_core::CooldownConfig; + let default_cfg = CooldownConfig::default(); + let cfg = cfg.unwrap_or(&default_cfg); + if !cfg.enabled_or_default() { + return None; + } + let max = Duration::from_secs(cfg.max_seconds_or_default()); + let default_ttl = Duration::from_secs(cfg.default_seconds_or_default()); + let clamp = |d: Duration| -> Duration { + let bounded = d.min(max); + // 0 means "disabled" via field; if we got here, default_ttl + // should never be 0 — fall back to max if someone set + // default_seconds=0 by mistake but didn't disable. + if bounded.is_zero() { + max + } else { + bounded + } + }; + + match err { + BridgeError::UpstreamStatus { + status, + retry_after, + .. + } => { + let triggers = cfg.effective_trigger_statuses(); + if !triggers.contains(status) { + return None; + } + let ttl = if cfg.honor_retry_after_or_default() { + retry_after.map(clamp).unwrap_or(default_ttl) + } else { + default_ttl + }; + Some((ttl, cooldown_reason_for_status(*status))) + } + BridgeError::Timeout { .. } if cfg.trigger_on_timeout_or_default() => { + Some((default_ttl, "request_timeout")) + } + BridgeError::Transport(_) | BridgeError::StreamAborted + if cfg.trigger_on_transport_or_default() => + { + Some((default_ttl, "transport_error")) + } + BridgeError::UpstreamDecode(_) if cfg.trigger_on_transport_or_default() => { + Some((default_ttl, "upstream_decode_error")) + } + // Config errors mean WE are misconfigured (missing provider, bad + // bridge registration). Cooling down doesn't help; let it + // surface and operator fixes the snapshot. + BridgeError::Config(_) => None, + _ => None, + } +} + +fn cooldown_reason_for_status(status: u16) -> &'static str { + match status { + 401 => "upstream_auth_failure", + 408 => "upstream_request_timeout", + 429 => "upstream_rate_limited", + 500..=599 => "upstream_server_error", + _ => "upstream_status_failure", + } +} + pub async fn chat_completions( State(state): State, auth: AuthenticatedKey, @@ -453,7 +541,23 @@ async fn dispatch( model: target_entry.value.clone(), }); } - filter_attempt_models(&state.runtime_status, resolved) + match filter_attempt_models( + &state.runtime_status, + resolved, + routing.on_all_filtered_or_default(), + ) { + FilterOutcome::Selected(list) => list, + FilterOutcome::AllUnhealthy { retry_after_secs } => { + tracing::warn!( + virtual_model = %req.model, + retry_after_secs, + "all routing candidates are unavailable; failing fast", + ); + return Err(with_model(ProxyError::AllCandidatesUnavailable { + retry_after_secs, + })); + } + } } else { vec![AttemptModel { id: virtual_entry.id.clone(), @@ -813,11 +917,16 @@ async fn dispatch( ); if retryable { state.health.record_failure(&model.display_name); - state.runtime_status.mark_cooldown( - &attempt.id, - RETRYABLE_FAILURE_COOLDOWN, - "retryable_failure", - ); + } + // Cooldown decision is independent of retry — a + // non-retryable 401 still cools down because the + // same key will keep failing for upcoming + // requests; a retryable 502 also cools down so + // the next request prefers a different target. + if let Some((ttl, reason)) = + decide_cooldown(&err, attempt.model.cooldown.as_ref()) + { + state.runtime_status.mark_cooldown(&attempt.id, ttl, reason); } last_err = Some(err); if !retryable { @@ -981,10 +1090,29 @@ async fn dispatch( }) } +/// Outcome of routing-candidate filtering. Lifts the "all candidates +/// excluded" case out into a typed result so the dispatch loop can +/// short-circuit to a 503 + Retry-After instead of sending traffic to +/// a target we just confirmed is bad. +enum FilterOutcome { + /// At least one candidate survived the filter. The returned vector + /// is the filtered attempt list, in the original strategy order + /// minus the excluded entries. + Selected(Vec), + /// Every candidate is currently background-unhealthy and the + /// routing model is configured with `on_all_filtered: fail`. The + /// caller should surface a 503 with the supplied Retry-After hint + /// (in seconds), if any. + AllUnhealthy { retry_after_secs: Option }, +} + fn filter_attempt_models( runtime_status: &crate::ModelRuntimeStatusTracker, attempts: Vec, -) -> Vec { + policy: aisix_core::OnAllFilteredPolicy, +) -> FilterOutcome { + use aisix_core::OnAllFilteredPolicy; + let mut healthy = Vec::new(); let mut cooldown_only = Vec::new(); let mut unhealthy_count = 0usize; @@ -995,8 +1123,8 @@ fn filter_attempt_models( .background_model_check .as_ref() .map(|cfg| Duration::from_secs(cfg.stale_after_seconds)); - let status = runtime_status.status_with_stale(&attempt.id, stale_after); - match status.status { + let snapshot = runtime_status.status_with_stale(&attempt.id, stale_after); + match snapshot.status { crate::RuntimeStatus::Unhealthy => unhealthy_count += 1, crate::RuntimeStatus::Cooldown => cooldown_only.push(attempt), crate::RuntimeStatus::Healthy | crate::RuntimeStatus::NotApplicable => { @@ -1006,10 +1134,14 @@ fn filter_attempt_models( } if !healthy.is_empty() { - return healthy; + return FilterOutcome::Selected(healthy); } + // No healthy candidates — prefer cooldown over unhealthy when + // some non-unhealthy candidates exist. Sending to a target whose + // cooldown timer hasn't expired is still better than sending to + // a target that an active probe just confirmed is broken. if unhealthy_count < attempts.len() && !cooldown_only.is_empty() { - return attempts + let filtered: Vec = attempts .into_iter() .filter(|attempt| { let stale_after = attempt @@ -1021,8 +1153,23 @@ fn filter_attempt_models( != crate::RuntimeStatus::Unhealthy }) .collect(); + return FilterOutcome::Selected(filtered); + } + // All candidates are excluded. Policy decides. + // + // Retry-After for the fail path is a coarse fallback (30s by + // default — see FALLBACK_ALL_UNHEALTHY_RETRY_AFTER). We could + // try to derive it from per-candidate cooldown timers, but the + // categorisation above routes cooldown candidates into + // `cooldown_only` (returned via the Selected branch above), so + // by construction every candidate that reaches here is in the + // background-unhealthy state and has no cooldown timer to read. + match policy { + OnAllFilteredPolicy::Fail => FilterOutcome::AllUnhealthy { + retry_after_secs: Some(FALLBACK_ALL_UNHEALTHY_RETRY_AFTER.as_secs()), + }, + OnAllFilteredPolicy::OriginalOrder => FilterOutcome::Selected(attempts), } - attempts } /// Wire-shape label for `FinishReason`. cp-api stores this verbatim @@ -1539,3 +1686,259 @@ fn error_frame_payload(error_type: &str, message: &str) -> String { r#"{"error":{"message":"error","type":"internal_error"}}"#.into() }) } + +#[cfg(test)] +mod cooldown_tests { + use super::*; + use aisix_core::CooldownConfig; + use std::time::Duration as StdDuration; + + fn upstream(status: u16) -> BridgeError { + BridgeError::upstream_status(status, "boom") + } + + fn upstream_with_retry_after(status: u16, secs: u64) -> BridgeError { + BridgeError::upstream_status_with_retry_after( + status, + "rate limited", + Some(StdDuration::from_secs(secs)), + ) + } + + #[test] + fn default_config_cooldowns_429() { + let (ttl, reason) = decide_cooldown(&upstream(429), None).unwrap(); + assert_eq!(ttl, StdDuration::from_secs(30)); + assert_eq!(reason, "upstream_rate_limited"); + } + + #[test] + fn default_config_cooldowns_401_even_though_non_retryable() { + // H1 contract: 401 cools down. Non-retryable upstream errors + // (auth failure) should still take the target out of rotation, + // because the same key will keep failing on subsequent + // requests. The retry-vs-cooldown split is the whole point. + let (ttl, reason) = decide_cooldown(&upstream(401), None).unwrap(); + assert_eq!(ttl, StdDuration::from_secs(30)); + assert_eq!(reason, "upstream_auth_failure"); + } + + #[test] + fn default_config_cooldowns_408() { + let (_, reason) = decide_cooldown(&upstream(408), None).unwrap(); + assert_eq!(reason, "upstream_request_timeout"); + } + + #[test] + fn default_config_cooldowns_5xx() { + for status in [500, 502, 503, 504] { + let (_, reason) = decide_cooldown(&upstream(status), None).unwrap(); + assert_eq!(reason, "upstream_server_error", "status={status}"); + } + } + + #[test] + fn default_config_skips_400_and_other_4xx() { + // Caller bugs (400, 403, 422) are not cooldown signals — the + // model didn't fail, the request did. + assert!(decide_cooldown(&upstream(400), None).is_none()); + assert!(decide_cooldown(&upstream(403), None).is_none()); + assert!(decide_cooldown(&upstream(422), None).is_none()); + } + + #[test] + fn default_config_cooldowns_timeout_and_transport_errors() { + assert!(decide_cooldown(&BridgeError::Timeout { elapsed_ms: 30_000 }, None).is_some()); + assert!(decide_cooldown(&BridgeError::Transport("conn refused".into()), None).is_some()); + assert!(decide_cooldown(&BridgeError::StreamAborted, None).is_some()); + assert!(decide_cooldown(&BridgeError::UpstreamDecode("bad json".into()), None).is_some()); + } + + #[test] + fn config_disabled_skips_cooldown() { + let cfg = CooldownConfig { + enabled: Some(false), + ..Default::default() + }; + assert!(decide_cooldown(&upstream(429), Some(&cfg)).is_none()); + assert!(decide_cooldown(&upstream(500), Some(&cfg)).is_none()); + } + + #[test] + fn config_override_trigger_statuses_excludes_429() { + // Operator wants 429 to NOT cool down (e.g. heavy retry policy + // already handles burst). 500s still cool down. + let cfg = CooldownConfig { + trigger_statuses: Some(vec![500, 502, 503]), + ..Default::default() + }; + assert!(decide_cooldown(&upstream(429), Some(&cfg)).is_none()); + assert!(decide_cooldown(&upstream(503), Some(&cfg)).is_some()); + } + + #[test] + fn honor_retry_after_uses_upstream_hint() { + let (ttl, _) = decide_cooldown(&upstream_with_retry_after(429, 75), None).unwrap(); + assert_eq!(ttl, StdDuration::from_secs(75)); + } + + #[test] + fn honor_retry_after_clamps_to_max_seconds() { + // Upstream is misbehaving — Retry-After: 100000. Clamp to + // configured max so we don't lose the target for hours. + let cfg = CooldownConfig { + max_seconds: Some(60), + ..Default::default() + }; + let (ttl, _) = + decide_cooldown(&upstream_with_retry_after(429, 100_000), Some(&cfg)).unwrap(); + assert_eq!(ttl, StdDuration::from_secs(60)); + } + + #[test] + fn honor_retry_after_disabled_falls_back_to_default() { + let cfg = CooldownConfig { + honor_retry_after: Some(false), + default_seconds: Some(45), + ..Default::default() + }; + let (ttl, _) = decide_cooldown(&upstream_with_retry_after(429, 5), Some(&cfg)).unwrap(); + assert_eq!(ttl, StdDuration::from_secs(45)); + } + + #[test] + fn trigger_on_timeout_false_disables_timeout_cooldown() { + let cfg = CooldownConfig { + trigger_on_timeout: Some(false), + ..Default::default() + }; + assert!(decide_cooldown(&BridgeError::Timeout { elapsed_ms: 1 }, Some(&cfg)).is_none()); + } + + #[test] + fn config_error_never_cools_down() { + // Misconfig = WE are wrong; cooling down doesn't help. + assert!(decide_cooldown(&BridgeError::Config("bad key".into()), None).is_none()); + } +} + +#[cfg(test)] +mod filter_tests { + use super::*; + use aisix_core::{Model, OnAllFilteredPolicy}; + use std::time::Duration as StdDuration; + + fn am(id: &str) -> AttemptModel { + let model: Model = serde_json::from_str(&format!( + r#"{{ + "display_name": "{id}", + "provider": "openai", + "model_name": "gpt-4o-mini", + "provider_key_id": "pk-{id}" + }}"# + )) + .unwrap(); + AttemptModel { + id: id.to_string(), + model, + } + } + + #[test] + fn healthy_only_returns_all_healthy() { + let t = crate::ModelRuntimeStatusTracker::new(); + let attempts = vec![am("a"), am("b")]; + match filter_attempt_models(&t, attempts, OnAllFilteredPolicy::Fail) { + FilterOutcome::Selected(list) => { + assert_eq!(list.len(), 2); + } + other => panic!( + "expected Selected, got {:?}", + std::mem::discriminant(&other) + ), + } + } + + #[test] + fn cooldown_skipped_when_healthy_present() { + let t = crate::ModelRuntimeStatusTracker::new(); + t.mark_cooldown("a", StdDuration::from_secs(30), "retryable_failure"); + let attempts = vec![am("a"), am("b")]; + match filter_attempt_models(&t, attempts, OnAllFilteredPolicy::Fail) { + FilterOutcome::Selected(list) => { + assert_eq!(list.len(), 1); + assert_eq!(list[0].id, "b"); + } + _ => panic!("expected Selected"), + } + } + + #[test] + fn all_unhealthy_fail_policy_returns_retry_after_hint() { + // H3 contract: every candidate background-unhealthy, no + // cooldown timer → return 503 + fallback Retry-After (30s + // default). The dispatch loop converts this to a + // ProxyError::AllCandidatesUnavailable. + let t = crate::ModelRuntimeStatusTracker::new(); + t.mark_unhealthy("a", Some(503), "background_check_failed"); + t.mark_unhealthy("b", Some(503), "background_check_failed"); + let attempts = vec![am("a"), am("b")]; + match filter_attempt_models(&t, attempts, OnAllFilteredPolicy::Fail) { + FilterOutcome::AllUnhealthy { retry_after_secs } => { + assert_eq!(retry_after_secs, Some(30)); + } + _ => panic!("expected AllUnhealthy"), + } + } + + #[test] + fn one_cooldown_with_all_else_unhealthy_keeps_the_cooldown_candidate() { + // Mixed scenario: candidates a/b are background-unhealthy, c + // is in cooldown. The filter should pick c (cooldown beats + // unhealthy), not fail. + let t = crate::ModelRuntimeStatusTracker::new(); + t.mark_unhealthy("a", Some(503), "background_check_failed"); + t.mark_unhealthy("b", Some(503), "background_check_failed"); + t.mark_cooldown("c", StdDuration::from_secs(30), "x"); + let attempts = vec![am("a"), am("b"), am("c")]; + match filter_attempt_models(&t, attempts, OnAllFilteredPolicy::Fail) { + FilterOutcome::Selected(list) => { + assert_eq!(list.len(), 1); + assert_eq!(list[0].id, "c"); + } + _ => panic!("expected Selected with cooldown candidate"), + } + } + + #[test] + fn all_unhealthy_original_order_policy_returns_full_list() { + // Legacy opt-in: send to all candidates regardless. + let t = crate::ModelRuntimeStatusTracker::new(); + t.mark_unhealthy("a", Some(503), "background_check_failed"); + t.mark_unhealthy("b", Some(503), "background_check_failed"); + let attempts = vec![am("a"), am("b")]; + match filter_attempt_models(&t, attempts, OnAllFilteredPolicy::OriginalOrder) { + FilterOutcome::Selected(list) => { + assert_eq!(list.len(), 2); + } + _ => panic!("expected Selected under OriginalOrder policy"), + } + } + + #[test] + fn cooldown_no_unhealthy_returns_cooldown_candidates() { + // No healthy, no unhealthy — all candidates have a cooldown + // timer set. Routing should still pick from them (better than + // erroring out when we don't have evidence anyone is *broken*). + let t = crate::ModelRuntimeStatusTracker::new(); + t.mark_cooldown("a", StdDuration::from_secs(30), "x"); + t.mark_cooldown("b", StdDuration::from_secs(30), "x"); + let attempts = vec![am("a"), am("b")]; + match filter_attempt_models(&t, attempts, OnAllFilteredPolicy::Fail) { + FilterOutcome::Selected(list) => { + assert_eq!(list.len(), 2); + } + _ => panic!("expected Selected for cooldown-only"), + } + } +} diff --git a/crates/aisix-proxy/src/error.rs b/crates/aisix-proxy/src/error.rs index ad0f6ed5..4e30fd34 100644 --- a/crates/aisix-proxy/src/error.rs +++ b/crates/aisix-proxy/src/error.rs @@ -72,6 +72,13 @@ pub enum ProxyError { InvalidRequest(String), #[error("no bridge registered for provider")] ProviderUnavailable, + /// Every routing candidate was excluded by the runtime status layer + /// (all in cooldown or background-unhealthy) and the routing model + /// is configured with `on_all_filtered: fail`. Caller-visible as + /// 503 with a Retry-After hint derived from the nearest cooldown + /// expiry. See [`aisix_core::OnAllFilteredPolicy`]. + #[error("all routing candidates are unavailable")] + AllCandidatesUnavailable { retry_after_secs: Option }, /// Caller-visible message MUST NOT carry the matched-pattern detail. /// Per #153, leaking the matched literal back to the caller defeats /// the point of an output guardrail (the whole purpose is to keep the @@ -109,6 +116,7 @@ impl ProxyError { ProxyError::ModelNotFound(_) => StatusCode::NOT_FOUND, ProxyError::InvalidRequest(_) => StatusCode::BAD_REQUEST, ProxyError::ProviderUnavailable => StatusCode::SERVICE_UNAVAILABLE, + ProxyError::AllCandidatesUnavailable { .. } => StatusCode::SERVICE_UNAVAILABLE, ProxyError::ContentFiltered(_) => StatusCode::UNPROCESSABLE_ENTITY, ProxyError::BudgetExceeded(_) => StatusCode::TOO_MANY_REQUESTS, ProxyError::RequestTooLarge { .. } => StatusCode::PAYLOAD_TOO_LARGE, @@ -127,6 +135,7 @@ impl ProxyError { ProxyError::InvalidRequest(_) => "invalid_request_error", ProxyError::RequestTooLarge { .. } => "invalid_request_error", ProxyError::ProviderUnavailable => "provider_unavailable", + ProxyError::AllCandidatesUnavailable { .. } => "all_candidates_unavailable", ProxyError::ContentFiltered(_) => "content_filter", ProxyError::BudgetExceeded(_) => "billing_error", ProxyError::RateLimit(_) => "rate_limit_exceeded", @@ -140,6 +149,7 @@ impl ProxyError { pub fn retry_after_secs(&self) -> Option { match self { ProxyError::RateLimit(e) => e.retry_after_secs(), + ProxyError::AllCandidatesUnavailable { retry_after_secs } => *retry_after_secs, _ => None, } } @@ -188,10 +198,7 @@ mod tests { #[test] fn bridge_error_inherits_status_and_type() { - let bridge_err = BridgeError::UpstreamStatus { - status: 429, - message: "rate limited".into(), - }; + let bridge_err = BridgeError::upstream_status(429, "rate limited"); let wrapped = ProxyError::Bridge(bridge_err); assert_eq!(wrapped.status(), StatusCode::TOO_MANY_REQUESTS); assert_eq!(wrapped.kind(), "upstream_error"); @@ -199,14 +206,26 @@ mod tests { #[test] fn bridge_5xx_collapses_via_bridge_error_mapping() { - let bridge_err = BridgeError::UpstreamStatus { - status: 503, - message: "busy".into(), - }; + let bridge_err = BridgeError::upstream_status(503, "busy"); let wrapped = ProxyError::Bridge(bridge_err); assert_eq!(wrapped.status(), StatusCode::BAD_GATEWAY); } + #[test] + fn all_candidates_unavailable_is_503_with_optional_retry_after() { + let with_hint = ProxyError::AllCandidatesUnavailable { + retry_after_secs: Some(42), + }; + assert_eq!(with_hint.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(with_hint.kind(), "all_candidates_unavailable"); + assert_eq!(with_hint.retry_after_secs(), Some(42)); + + let no_hint = ProxyError::AllCandidatesUnavailable { + retry_after_secs: None, + }; + assert_eq!(no_hint.retry_after_secs(), None); + } + #[test] fn envelope_omits_null_param_and_code_on_wire() { let env = ProxyError::ModelNotFound("x".into()).envelope(); diff --git a/crates/aisix-proxy/src/health.rs b/crates/aisix-proxy/src/health.rs index 1dd10560..96b9126b 100644 --- a/crates/aisix-proxy/src/health.rs +++ b/crates/aisix-proxy/src/health.rs @@ -134,7 +134,7 @@ impl Default for RuntimeStatusSnapshot { } } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] struct RuntimeEntry { unhealthy: bool, cooldown_until: Option, @@ -143,18 +143,6 @@ struct RuntimeEntry { status_reason: Option, } -impl Default for RuntimeEntry { - fn default() -> Self { - Self { - unhealthy: false, - cooldown_until: None, - last_checked_at: None, - last_check_status: None, - status_reason: None, - } - } -} - impl RuntimeEntry { fn snapshot(&self, now: SystemTime, stale_after: Option) -> RuntimeStatusSnapshot { let cooldown_until = self.cooldown_until.filter(|until| *until > now); @@ -365,12 +353,7 @@ impl ModelRuntimeStatusTracker { }); } - pub fn record_ignored_check( - &self, - model_id: &str, - status: u16, - reason: impl Into, - ) { + pub fn record_ignored_check(&self, model_id: &str, status: u16, reason: impl Into) { let now = SystemTime::now(); let reason = reason.into(); self.entries @@ -404,7 +387,11 @@ impl ModelRuntimeStatusTracker { .unwrap_or_default() } - pub fn should_skip_for_routing(&self, model_id: &str, stale_after: Option) -> RuntimeStatus { + pub fn should_skip_for_routing( + &self, + model_id: &str, + stale_after: Option, + ) -> RuntimeStatus { self.status_with_stale(model_id, stale_after).status } } @@ -545,12 +532,14 @@ mod tests { let t = ModelRuntimeStatusTracker::new(); t.mark_unhealthy("m-1", Some(503), "background_check_failed"); assert_eq!( - t.status_with_stale("m-1", Some(Duration::from_secs(60))).status, + t.status_with_stale("m-1", Some(Duration::from_secs(60))) + .status, RuntimeStatus::Unhealthy ); std::thread::sleep(Duration::from_millis(15)); assert_eq!( - t.status_with_stale("m-1", Some(Duration::from_millis(1))).status, + t.status_with_stale("m-1", Some(Duration::from_millis(1))) + .status, RuntimeStatus::Healthy ); } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 4a8c9d24..60fa062c 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -2018,7 +2018,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::OK, + "{}", + String::from_utf8_lossy(&bytes) + ); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "fallback worked"); } @@ -2083,7 +2088,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::BAD_REQUEST, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::BAD_REQUEST, + "{}", + String::from_utf8_lossy(&bytes) + ); } #[tokio::test] @@ -2151,7 +2161,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::OK, + "{}", + String::from_utf8_lossy(&bytes) + ); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "after retries"); } @@ -2221,7 +2236,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::OK, + "{}", + String::from_utf8_lossy(&bytes) + ); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "429 fallback worked"); } @@ -2306,7 +2326,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::OK, + "{}", + String::from_utf8_lossy(&bytes) + ); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "cooldown skipped"); } @@ -2370,7 +2395,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::OK, + "{}", + String::from_utf8_lossy(&bytes) + ); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "cooldown fallback"); } @@ -2380,9 +2410,7 @@ data: [DONE]\n\n"; let flaky_upstream = MockServer::start().await; Mock::given(method("POST")) .and(path("/chat/completions")) - .respond_with( - ResponseTemplate::new(502).set_body_string("temporary upstream failure"), - ) + .respond_with(ResponseTemplate::new(502).set_body_string("temporary upstream failure")) .expect(1) .mount(&flaky_upstream) .await; @@ -2442,7 +2470,10 @@ data: [DONE]\n\n"; let resp = run(app, req).await; assert_eq!(resp.status(), StatusCode::OK); - assert_eq!(state.runtime_status.status("m-flaky").status, RuntimeStatus::Cooldown); + assert_eq!( + state.runtime_status.status("m-flaky").status, + RuntimeStatus::Cooldown + ); let app = build_router(state); let body = serde_json::json!({ @@ -2460,7 +2491,12 @@ data: [DONE]\n\n"; let resp = run(app, req).await; let status = resp.status(); let bytes = to_bytes(resp.into_body(), 65536).await.unwrap(); - assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&bytes)); + assert_eq!( + status, + StatusCode::OK, + "{}", + String::from_utf8_lossy(&bytes) + ); let v: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); assert_eq!(v["choices"][0]["message"]["content"], "stable fallback"); } @@ -2473,8 +2509,14 @@ data: [DONE]\n\n"; hub.register(Provider::Openai, Arc::new(openai_test_bridge())); let snap = AisixSnapshot::new(); - snap.models - .insert(routing_entry("smart", "failover", &["nonexistent"], None, None, None)); + snap.models.insert(routing_entry( + "smart", + "failover", + &["nonexistent"], + None, + None, + None, + )); snap.apikeys.insert(apikey_entry("sk-caller", &["smart"])); // No upstream provider_key needed — the routing target itself // is missing so dispatch fails before any provider lookup. diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 24e07046..d71c15d3 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -231,16 +231,19 @@ async fn dispatch( if !status.is_success() { let status_u16 = status.as_u16(); + let retry_after = aisix_gateway::parse_retry_after(upstream_resp.headers()); let message = upstream_resp.text().await.unwrap_or_default(); + let truncated = if message.len() > 1024 { + format!("{}…", &message[..1024]) + } else { + message + }; return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::UpstreamStatus { - status: status_u16, - message: if message.len() > 1024 { - format!("{}…", &message[..1024]) - } else { - message - }, - }, + aisix_gateway::BridgeError::upstream_status_with_retry_after( + status_u16, + truncated, + retry_after, + ), )); } diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 002e2100..41c4fadb 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -187,12 +187,14 @@ async fn dispatch( if !status.is_success() { let status_u16 = status.as_u16(); + let retry_after = aisix_gateway::parse_retry_after(upstream_resp.headers()); let message = upstream_resp.text().await.unwrap_or_default(); return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::UpstreamStatus { - status: status_u16, - message: message.chars().take(1024).collect(), - }, + aisix_gateway::BridgeError::upstream_status_with_retry_after( + status_u16, + message.chars().take(1024).collect::(), + retry_after, + ), )); } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 6accf3b0..5fdd1e99 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -159,12 +159,14 @@ async fn dispatch( if !status.is_success() { let status_u16 = status.as_u16(); + let retry_after = aisix_gateway::parse_retry_after(upstream_resp.headers()); let message = upstream_resp.text().await.unwrap_or_default(); return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::UpstreamStatus { - status: status_u16, - message: message.chars().take(1024).collect(), - }, + aisix_gateway::BridgeError::upstream_status_with_retry_after( + status_u16, + message.chars().take(1024).collect::(), + retry_after, + ), )); } diff --git a/crates/aisix-proxy/src/routing.rs b/crates/aisix-proxy/src/routing.rs index 942a3eb2..b6aefa94 100644 --- a/crates/aisix-proxy/src/routing.rs +++ b/crates/aisix-proxy/src/routing.rs @@ -60,7 +60,11 @@ impl RoutingRegistry { return Vec::new(); } let start = self.starting_index(virtual_name, routing); - attempt_order(&routing.targets, start, routing.max_fallbacks_or_default() + 1) + attempt_order( + &routing.targets, + start, + routing.max_fallbacks_or_default() + 1, + ) } fn starting_index(&self, virtual_name: &str, routing: &Routing) -> usize { @@ -150,6 +154,7 @@ mod tests { retries: None, max_fallbacks, retry_on_429: None, + on_all_filtered: None, } } @@ -407,25 +412,28 @@ mod tests { #[test] fn is_retryable_distinguishes_4xx_from_other_failures() { - assert!(!is_retryable(&BridgeError::UpstreamStatus { - status: 400, - message: "bad request".into(), - }, false)); - assert!(!is_retryable(&BridgeError::UpstreamStatus { - status: 429, - message: "rate limited".into(), - }, false)); - assert!(is_retryable(&BridgeError::UpstreamStatus { - status: 429, - message: "rate limited".into(), - }, true)); - assert!(is_retryable(&BridgeError::UpstreamStatus { - status: 502, - message: "bad gateway".into(), - }, false)); + assert!(!is_retryable( + &BridgeError::upstream_status(400, "bad request"), + false + )); + assert!(!is_retryable( + &BridgeError::upstream_status(429, "rate limited"), + false + )); + assert!(is_retryable( + &BridgeError::upstream_status(429, "rate limited"), + true + )); + assert!(is_retryable( + &BridgeError::upstream_status(502, "bad gateway"), + false + )); assert!(is_retryable(&BridgeError::Timeout { elapsed_ms: 1 }, false)); assert!(is_retryable(&BridgeError::Transport("conn".into()), false)); - assert!(is_retryable(&BridgeError::UpstreamDecode("x".into()), false)); + assert!(is_retryable( + &BridgeError::UpstreamDecode("x".into()), + false + )); assert!(is_retryable(&BridgeError::Config("bad key".into()), false)); assert!(is_retryable(&BridgeError::StreamAborted, false)); } diff --git a/tests/e2e/src/cases/background-health-e2e.test.ts b/tests/e2e/src/cases/background-health-e2e.test.ts new file mode 100644 index 00000000..b00f9c9c --- /dev/null +++ b/tests/e2e/src/cases/background-health-e2e.test.ts @@ -0,0 +1,264 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-background-health-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +describe("background health e2e", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let unhealthyUpstream: OpenAiUpstream | undefined; + let ignoredUpstream: OpenAiUpstream | undefined; + let stableUpstream: OpenAiUpstream | undefined; + let unhealthyModelID = ""; + let ignoredModelID = ""; + let stableModelID = ""; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + unhealthyUpstream = await startOpenAiUpstream({ + status: 503, + errorBody: { error: { message: "background unhealthy", type: "server_error" } }, + }); + ignoredUpstream = await startOpenAiUpstream({ + status: 429, + errorBody: { error: { message: "background ignore", type: "rate_limit_error" } }, + }); + stableUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-stable-bg", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "healthy target" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const unhealthyPk = await admin.createProviderKey({ + display_name: "bg-unhealthy-pk", + secret: "sk-mock", + api_base: `${unhealthyUpstream.baseUrl}/v1`, + }); + const ignoredPk = await admin.createProviderKey({ + display_name: "bg-ignored-pk", + secret: "sk-mock", + api_base: `${ignoredUpstream.baseUrl}/v1`, + }); + const stablePk = await admin.createProviderKey({ + display_name: "bg-stable-pk", + secret: "sk-mock", + api_base: `${stableUpstream.baseUrl}/v1`, + }); + + unhealthyModelID = ( + await admin.createModel({ + display_name: "bg-unhealthy", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: unhealthyPk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 90, + }, + }) + ).id; + + ignoredModelID = ( + await admin.createModel({ + display_name: "bg-ignored", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: ignoredPk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 90, + }, + }) + ).id; + + stableModelID = ( + await admin.createModel({ + display_name: "bg-stable", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: stablePk.id, + }) + ).id; + + await admin.createModel({ + display_name: "bg-router", + routing: { + strategy: "failover", + targets: [ + { model: "bg-unhealthy" }, + { model: "bg-stable" }, + ], + max_fallbacks: 1, + }, + }); + + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["bg-router", "bg-stable"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await unhealthyUpstream?.close(); + await ignoredUpstream?.close(); + await stableUpstream?.close(); + }); + + test("background unhealthy is surfaced and ignored 429 stays visible but healthy", async (ctx) => { + if (!etcdReachable || !app || !admin || !unhealthyUpstream || !ignoredUpstream || !stableUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + let lastStatuses: Array> = []; + try { + await waitConfigPropagation(async () => { + lastStatuses = await admin!.listModelStatuses(); + const unhealthy = lastStatuses.find((row) => row.id === unhealthyModelID); + const ignored = lastStatuses.find((row) => row.id === ignoredModelID); + return unhealthy?.status === "unhealthy" && ignored?.last_check_status === 429; + }); + } catch (err) { + throw new Error( + `${(err as Error).message}\nlast model statuses: ${JSON.stringify(lastStatuses, null, 2)}`, + ); + } + + const statuses = await admin.listModelStatuses(); + const unhealthy = statuses.find((row) => row.id === unhealthyModelID)!; + const ignored = statuses.find((row) => row.id === ignoredModelID)!; + const stable = statuses.find((row) => row.id === stableModelID)!; + + expect(unhealthy.status).toBe("unhealthy"); + expect(unhealthy.last_check_status).toBe(503); + expect(unhealthy.status_reason).toBe("background_check_failed"); + + expect(ignored.status).toBe("healthy"); + expect(ignored.last_check_status).toBe(429); + expect(ignored.status_reason).toBe("ignored_transient_error"); + + expect(stable.status).toBe("healthy"); + + await waitConfigPropagation(async () => { + try { + const probe = await client.chat.completions.create({ + model: "bg-router", + messages: [{ role: "user", content: "ready-bg-router" }], + }); + return probe.choices[0]?.message.content === "healthy target"; + } catch { + return false; + } + }); + + const unhealthyBaseline = unhealthyUpstream.receivedRequests.length; + const stableBaseline = stableUpstream.receivedRequests.length; + + const completion = await client.chat.completions.create({ + model: "bg-router", + messages: [{ role: "user", content: "skip background unhealthy" }], + }); + expect(completion.choices[0]?.message.content).toBe("healthy target"); + expect(unhealthyUpstream.receivedRequests.length - unhealthyBaseline).toBe(0); + expect(stableUpstream.receivedRequests.length - stableBaseline).toBe(1); + }); + + test("active background checks keep unhealthy state fresh even with a short stale window", async (ctx) => { + if (!etcdReachable || !app || !admin || !unhealthyUpstream) { + ctx.skip(); + return; + } + + const shortStalePk = await admin.createProviderKey({ + display_name: "bg-stale-short-pk", + secret: "sk-mock", + api_base: `${unhealthyUpstream.baseUrl}/v1`, + }); + await admin.createModel({ + display_name: "bg-stale-short", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: shortStalePk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [], + // Stale window must be > probe interval so a successful + // probe at t≈interval refreshes last_checked_at into the + // future — keeping the unhealthy entry "fresh" across the + // wait below. If stale_after < interval, the entry could + // expire between probes and the test would race. + stale_after_seconds: 8, + }, + }); + + const statuses = await admin.listModelStatuses(); + const staleModel = statuses.find((row) => row.display_name === "bg-stale-short"); + expect(staleModel).toBeTruthy(); + + await waitConfigPropagation(async () => { + const rows = await admin!.listModelStatuses(); + const row = rows.find((item) => item.display_name === "bg-stale-short"); + return row?.status === "unhealthy"; + }); + + // Wait longer than one probe interval but well within the + // stale window — an active probe must fire and re-mark the + // entry unhealthy, keeping the state fresh. + await new Promise((r) => setTimeout(r, 6500)); + + const rows = await admin.listModelStatuses(); + const row = rows.find((item) => item.display_name === "bg-stale-short")!; + expect(row.status).toBe("unhealthy"); + }); +}); diff --git a/tests/e2e/src/cases/cooldown-contract-e2e.test.ts b/tests/e2e/src/cases/cooldown-contract-e2e.test.ts new file mode 100644 index 00000000..13ffbd38 --- /dev/null +++ b/tests/e2e/src/cases/cooldown-contract-e2e.test.ts @@ -0,0 +1,706 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +/** + * E2E contract tests for the post-#268 cooldown / filter behavior. + * + * These tests cover the audit findings from issue #264 that the + * initial PR landed: + * - H1: 401 (auth failure) cools down even though it is non-retryable. + * - H2: Retry-After header from upstream drives the cooldown TTL. + * - H3: When every routing candidate is filtered out, the proxy + * fails fast with 503 + Retry-After (default policy). + * - H3 escape hatch: `on_all_filtered: original_order` preserves the + * legacy "send to known-bad" behavior for operators that + * explicitly opt in. + * - M1: 429 cools down regardless of `retry_on_429` — cooldown and + * retry are independent layers. + * + * Every test exercises the real backend with mock upstreams (no + * stubbing of the bridge or network layer) and asserts both client- + * visible outcome and per-upstream request evidence. + */ + +const CALLER_PLAINTEXT = "sk-cooldown-contract-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +describe("cooldown contract (H1) — 401 cools down despite being non-retryable", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let authFailUpstream: OpenAiUpstream | undefined; + let stableUpstream: OpenAiUpstream | undefined; + let authFailModelID = ""; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + authFailUpstream = await startOpenAiUpstream({ + status: 401, + errorBody: { + error: { message: "Incorrect API key provided", type: "invalid_request_error" }, + }, + }); + stableUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-h1-stable", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "h1 stable fallback" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const failPk = await admin.createProviderKey({ + display_name: "h1-401-pk", + secret: "sk-mock", + api_base: `${authFailUpstream.baseUrl}/v1`, + }); + const stablePk = await admin.createProviderKey({ + display_name: "h1-stable-pk", + secret: "sk-mock", + api_base: `${stableUpstream.baseUrl}/v1`, + }); + + authFailModelID = ( + await admin.createModel({ + display_name: "h1-401-model", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: failPk.id, + }) + ).id; + await admin.createModel({ + display_name: "h1-stable-model", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: stablePk.id, + }); + await admin.createModel({ + display_name: "h1-router", + routing: { + strategy: "failover", + targets: [{ model: "h1-401-model" }, { model: "h1-stable-model" }], + max_fallbacks: 1, + }, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["h1-router", "h1-stable-model"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await authFailUpstream?.close(); + await stableUpstream?.close(); + }); + + test("a 401 surfaces as 401 to the caller (non-retryable) AND cools down the target; next request skips the cooled target", async (ctx) => { + if (!etcdReachable || !app || !admin || !authFailUpstream || !stableUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + try { + const probe = await client.chat.completions.create({ + model: "h1-stable-model", + messages: [{ role: "user", content: "ready-h1-stable" }], + }); + return probe.choices[0]?.message.content === "h1 stable fallback"; + } catch { + return false; + } + }); + + // First request: 401 is non-retryable, so the dispatch loop + // does NOT failover within this request — the 401 propagates + // to the caller. This matches PR #263's retry semantics where + // non-retryable 4xx stops the loop entirely. The new behavior + // is purely that COOLDOWN is recorded anyway. + let firstError: unknown; + try { + await client.chat.completions.create({ + model: "h1-router", + messages: [{ role: "user", content: "trip 401 cooldown" }], + }); + } catch (err) { + firstError = err; + } + expect(firstError).toBeTruthy(); + expect((firstError as { status?: number }).status).toBe(401); + expect(authFailUpstream.receivedRequests.length).toBeGreaterThanOrEqual(1); + + // The H1 contract: even though 401 is non-retryable, cooldown + // was set on the failing target. + const statuses = await admin.listModelStatuses(); + const failed = statuses.find((row) => row.id === authFailModelID)!; + expect(failed.status).toBe("cooldown"); + expect(failed.status_reason).toBe("upstream_auth_failure"); + expect(failed.cooldown_until).toBeTruthy(); + + // Second request: routing filter sees the 401 target is in + // cooldown, picks the stable fallback as the first attempt. + const failBaseline = authFailUpstream.receivedRequests.length; + const stableBaseline = stableUpstream.receivedRequests.length; + + const second = await client.chat.completions.create({ + model: "h1-router", + messages: [{ role: "user", content: "skip cooled-down 401 target" }], + }); + expect(second.choices[0]?.message.content).toBe("h1 stable fallback"); + expect(authFailUpstream.receivedRequests.length - failBaseline).toBe(0); + expect(stableUpstream.receivedRequests.length - stableBaseline).toBe(1); + }); +}); + +describe("cooldown contract (M1) — 429 cools down even when retry_on_429=false", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let rateLimitedUpstream: OpenAiUpstream | undefined; + let stableUpstream: OpenAiUpstream | undefined; + let rateLimitedModelID = ""; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + rateLimitedUpstream = await startOpenAiUpstream({ + status: 429, + errorBody: { + error: { message: "Rate limit reached", type: "rate_limit_error" }, + }, + }); + stableUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-m1-stable", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "m1 stable fallback" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const rateLimitedPk = await admin.createProviderKey({ + display_name: "m1-429-pk", + secret: "sk-mock", + api_base: `${rateLimitedUpstream.baseUrl}/v1`, + }); + const stablePk = await admin.createProviderKey({ + display_name: "m1-stable-pk", + secret: "sk-mock", + api_base: `${stableUpstream.baseUrl}/v1`, + }); + + rateLimitedModelID = ( + await admin.createModel({ + display_name: "m1-429-model", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: rateLimitedPk.id, + }) + ).id; + await admin.createModel({ + display_name: "m1-stable-model", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: stablePk.id, + }); + // Note: retry_on_429 defaults to false here. The pre-#268-fix + // behavior would NOT cooldown a 429 in this configuration; the + // M1 fix decouples cooldown from retry so the 429 must still + // mark the target for backoff. + await admin.createModel({ + display_name: "m1-router", + routing: { + strategy: "failover", + targets: [{ model: "m1-429-model" }, { model: "m1-stable-model" }], + max_fallbacks: 1, + }, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["m1-router", "m1-stable-model"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await rateLimitedUpstream?.close(); + await stableUpstream?.close(); + }); + + test("429 surfaces to caller (retry_on_429=false → non-retryable) AND cools down; next request skips the cooled target", async (ctx) => { + if (!etcdReachable || !app || !admin || !rateLimitedUpstream || !stableUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + try { + const probe = await client.chat.completions.create({ + model: "m1-stable-model", + messages: [{ role: "user", content: "ready-m1-stable" }], + }); + return probe.choices[0]?.message.content === "m1 stable fallback"; + } catch { + return false; + } + }); + + // First request: with retry_on_429 unset (default false), the + // 429 is non-retryable so the dispatch loop does not failover + // — the 429 propagates to the caller. This is the same + // retry-vs-cooldown decoupling the M1 contract pins. + let firstError: unknown; + try { + await client.chat.completions.create({ + model: "m1-router", + messages: [{ role: "user", content: "trip 429 cooldown" }], + }); + } catch (err) { + firstError = err; + } + expect(firstError).toBeTruthy(); + expect((firstError as { status?: number }).status).toBe(429); + expect(rateLimitedUpstream.receivedRequests.length).toBeGreaterThanOrEqual(1); + + // The M1 contract: 429 cools down regardless of retry_on_429. + const statuses = await admin.listModelStatuses(); + const rateLimited = statuses.find((row) => row.id === rateLimitedModelID)!; + expect(rateLimited.status).toBe("cooldown"); + expect(rateLimited.status_reason).toBe("upstream_rate_limited"); + + // Second request: routing filter skips the cooled target. + const rateLimitedBaseline = rateLimitedUpstream.receivedRequests.length; + const stableBaseline = stableUpstream.receivedRequests.length; + + const second = await client.chat.completions.create({ + model: "m1-router", + messages: [{ role: "user", content: "skip 429 target after cooldown" }], + }); + expect(second.choices[0]?.message.content).toBe("m1 stable fallback"); + expect(rateLimitedUpstream.receivedRequests.length - rateLimitedBaseline).toBe(0); + expect(stableUpstream.receivedRequests.length - stableBaseline).toBe(1); + }); +}); + +describe("cooldown contract (H2) — Retry-After header from upstream drives TTL", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let upstream: OpenAiUpstream | undefined; + let stableUpstream: OpenAiUpstream | undefined; + let rateLimitedModelID = ""; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + upstream = await startOpenAiUpstream({ + status: 429, + errorBody: { error: { message: "slow down", type: "rate_limit_error" } }, + // Force the upstream to advertise a long Retry-After. The + // gateway must honor it instead of using the default 30s. + // Test asserts cooldown_until lands well beyond the default. + responseHeaders: { "retry-after": "180" }, + }); + stableUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-h2-stable", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "h2 stable fallback" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const upstreamPk = await admin.createProviderKey({ + display_name: "h2-upstream-pk", + secret: "sk-mock", + api_base: `${upstream.baseUrl}/v1`, + }); + const stablePk = await admin.createProviderKey({ + display_name: "h2-stable-pk", + secret: "sk-mock", + api_base: `${stableUpstream.baseUrl}/v1`, + }); + + rateLimitedModelID = ( + await admin.createModel({ + display_name: "h2-429-with-retry-after", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: upstreamPk.id, + }) + ).id; + await admin.createModel({ + display_name: "h2-stable-model", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: stablePk.id, + }); + await admin.createModel({ + display_name: "h2-router", + routing: { + strategy: "failover", + targets: [ + { model: "h2-429-with-retry-after" }, + { model: "h2-stable-model" }, + ], + max_fallbacks: 1, + }, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["h2-router", "h2-stable-model"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await upstream?.close(); + await stableUpstream?.close(); + }); + + test("Retry-After: 180 from upstream produces a cooldown_until far beyond the 30s default", async (ctx) => { + if (!etcdReachable || !app || !admin || !upstream || !stableUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + try { + const probe = await client.chat.completions.create({ + model: "h2-stable-model", + messages: [{ role: "user", content: "ready-h2-stable" }], + }); + return probe.choices[0]?.message.content === "h2 stable fallback"; + } catch { + return false; + } + }); + + const before = Date.now(); + // 429 with retry_on_429 unset (default false) → non-retryable + // → propagates to caller. We don't care about the response + // path here, only about the cooldown TTL the gateway derived + // from the upstream's Retry-After header. + let firstError: unknown; + try { + await client.chat.completions.create({ + model: "h2-router", + messages: [{ role: "user", content: "trip 429 retry-after" }], + }); + } catch (err) { + firstError = err; + } + expect((firstError as { status?: number }).status).toBe(429); + + const statuses = await admin.listModelStatuses(); + const row = statuses.find((r) => r.id === rateLimitedModelID)!; + expect(row.status).toBe("cooldown"); + expect(row.cooldown_until).toBeTruthy(); + + // The reported cooldown_until should be ~180s in the future. + // 30s default is the wrong answer — that would prove H2 isn't + // wired. `cooldown_until` is serialized via SystemTime's default + // serde shape: `{secs_since_epoch, nanos_since_epoch}`. + const cooldownUntil = row.cooldown_until as { secs_since_epoch: number }; + const cooldownUntilMs = cooldownUntil.secs_since_epoch * 1000; + const horizonMs = cooldownUntilMs - before; + expect(horizonMs).toBeGreaterThan(120_000); + expect(horizonMs).toBeLessThan(240_000); + }); +}); + +describe("filter contract (H3) — all candidates unhealthy returns 503", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let downAUpstream: OpenAiUpstream | undefined; + let downBUpstream: OpenAiUpstream | undefined; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + downAUpstream = await startOpenAiUpstream({ + status: 503, + errorBody: { error: { message: "all-down A", type: "server_error" } }, + }); + downBUpstream = await startOpenAiUpstream({ + status: 503, + errorBody: { error: { message: "all-down B", type: "server_error" } }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const aPk = await admin.createProviderKey({ + display_name: "h3-down-a-pk", + secret: "sk-mock", + api_base: `${downAUpstream.baseUrl}/v1`, + }); + const bPk = await admin.createProviderKey({ + display_name: "h3-down-b-pk", + secret: "sk-mock", + api_base: `${downBUpstream.baseUrl}/v1`, + }); + + await admin.createModel({ + display_name: "h3-down-a", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: aPk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 120, + }, + }); + await admin.createModel({ + display_name: "h3-down-b", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: bPk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 120, + }, + }); + // Default on_all_filtered policy is "fail" — explicit here for + // clarity even though it's the default. + await admin.createModel({ + display_name: "h3-router-fail", + routing: { + strategy: "failover", + targets: [{ model: "h3-down-a" }, { model: "h3-down-b" }], + max_fallbacks: 1, + on_all_filtered: "fail", + }, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["h3-router-fail"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await downAUpstream?.close(); + await downBUpstream?.close(); + }); + + test("when both candidates are background-unhealthy, the proxy fails fast with 503 + Retry-After", async (ctx) => { + if (!etcdReachable || !app || !admin || !downAUpstream || !downBUpstream) { + ctx.skip(); + return; + } + + // Wait for the background probe to mark both candidates unhealthy. + await waitConfigPropagation(async () => { + const statuses = await admin!.listModelStatuses(); + const a = statuses.find((row) => row.display_name === "h3-down-a"); + const b = statuses.find((row) => row.display_name === "h3-down-b"); + return a?.status === "unhealthy" && b?.status === "unhealthy"; + }); + + const aBaseline = downAUpstream.receivedRequests.length; + const bBaseline = downBUpstream.receivedRequests.length; + + // Raw fetch instead of the OpenAI SDK so we can inspect the + // 503 response shape and Retry-After header directly. The SDK + // would convert 503 to an APIError and obscure the headers. + const resp = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${CALLER_PLAINTEXT}`, + }, + body: JSON.stringify({ + model: "h3-router-fail", + messages: [{ role: "user", content: "should fail fast" }], + }), + }); + + expect(resp.status).toBe(503); + expect(resp.headers.get("retry-after")).toBeTruthy(); + const body = (await resp.json()) as { error?: { type?: string } }; + expect(body.error?.type).toBe("all_candidates_unavailable"); + + // Crucially: the proxy must NOT have sent a request to either + // upstream during the fail-fast path. This is the whole point + // of H3 — don't pile on known-bad targets. + expect(downAUpstream.receivedRequests.length - aBaseline).toBe(0); + expect(downBUpstream.receivedRequests.length - bBaseline).toBe(0); + }); +}); + +describe("filter contract (H3 escape hatch) — original_order sends to known-bad", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let downUpstream: OpenAiUpstream | undefined; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + downUpstream = await startOpenAiUpstream({ + status: 503, + errorBody: { error: { message: "still down", type: "server_error" } }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const pk = await admin.createProviderKey({ + display_name: "h3-escape-pk", + secret: "sk-mock", + api_base: `${downUpstream.baseUrl}/v1`, + }); + + await admin.createModel({ + display_name: "h3-escape-down", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 120, + }, + }); + await admin.createModel({ + display_name: "h3-router-escape", + routing: { + strategy: "failover", + targets: [{ model: "h3-escape-down" }], + max_fallbacks: 0, + on_all_filtered: "original_order", + }, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["h3-router-escape"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await downUpstream?.close(); + }); + + test("with on_all_filtered=original_order, the proxy still tries the unhealthy target (legacy opt-in)", async (ctx) => { + if (!etcdReachable || !app || !admin || !downUpstream) { + ctx.skip(); + return; + } + + await waitConfigPropagation(async () => { + const statuses = await admin!.listModelStatuses(); + const row = statuses.find((r) => r.display_name === "h3-escape-down"); + return row?.status === "unhealthy"; + }); + + const baseline = downUpstream.receivedRequests.length; + + // With original_order, the request goes out to the unhealthy + // target. The upstream still returns 503 and the SDK throws, + // but the key point is that a request *was sent*. + const resp = await fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + "content-type": "application/json", + authorization: `Bearer ${CALLER_PLAINTEXT}`, + }, + body: JSON.stringify({ + model: "h3-router-escape", + messages: [{ role: "user", content: "send anyway please" }], + }), + }); + + // 503 surfaces from the upstream (collapsed to 502 by the bridge + // mapping in BridgeError::http_status, but the upstream request + // is what we're verifying here). + expect([502, 503]).toContain(resp.status); + expect(downUpstream.receivedRequests.length - baseline).toBe(1); + }); +}); diff --git a/tests/e2e/src/cases/fallback-e2e.test.ts b/tests/e2e/src/cases/fallback-e2e.test.ts index 5630ac04..a0407149 100644 --- a/tests/e2e/src/cases/fallback-e2e.test.ts +++ b/tests/e2e/src/cases/fallback-e2e.test.ts @@ -76,6 +76,13 @@ describe("fallback e2e: virtual routing fails over from 5xx to next target", () provider: "openai", model_name: "gpt-4o-mini", provider_key_id: badPk.id, + // This test is specifically about retry-time failover (bad→good + // within one request). Cooldown (post-PR #268) would mark fb-bad + // after the readiness probe exercises the failover path, and + // the subsequent test request would then skip fb-bad — defeating + // the very contract the test is pinning. Disable cooldown here + // so the test exercises only its target behavior. + cooldown: { enabled: false }, }); await admin.createModel({ display_name: "fb-good", diff --git a/tests/e2e/src/cases/retry-on-429-vs-background-ignore-e2e.test.ts b/tests/e2e/src/cases/retry-on-429-vs-background-ignore-e2e.test.ts new file mode 100644 index 00000000..03affbe3 --- /dev/null +++ b/tests/e2e/src/cases/retry-on-429-vs-background-ignore-e2e.test.ts @@ -0,0 +1,157 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-retry-429-vs-bg-ignore"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +describe("retry_on_429 vs background ignore e2e", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let rateLimitedUpstream: OpenAiUpstream | undefined; + let fallbackUpstream: OpenAiUpstream | undefined; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + rateLimitedUpstream = await startOpenAiUpstream({ + scriptedResponses: [ + { + status: 429, + errorBody: { error: { message: "background ignore", type: "rate_limit_error" } }, + }, + { + status: 429, + errorBody: { error: { message: "request path 429", type: "rate_limit_error" } }, + }, + { + status: 429, + errorBody: { error: { message: "request path 429 retry", type: "rate_limit_error" } }, + }, + ], + }); + fallbackUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-429-fallback", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "request-path 429 fallback" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const primaryPk = await admin.createProviderKey({ + display_name: "retry-429-primary-pk", + secret: "sk-mock", + api_base: `${rateLimitedUpstream.baseUrl}/v1`, + }); + const fallbackPk = await admin.createProviderKey({ + display_name: "retry-429-fallback-pk", + secret: "sk-mock", + api_base: `${fallbackUpstream.baseUrl}/v1`, + }); + + await admin.createModel({ + display_name: "retry-429-primary", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: primaryPk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 90, + }, + }); + await admin.createModel({ + display_name: "retry-429-fallback", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: fallbackPk.id, + }); + await admin.createModel({ + display_name: "retry-429-router", + routing: { + strategy: "failover", + targets: [ + { model: "retry-429-primary" }, + { model: "retry-429-fallback" }, + ], + retries: 1, + max_fallbacks: 1, + retry_on_429: true, + }, + }); + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["retry-429-router", "retry-429-fallback"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await rateLimitedUpstream?.close(); + await fallbackUpstream?.close(); + }); + + test("background 429 stays healthy while request-path 429 still retries and fails over", async (ctx) => { + if (!etcdReachable || !app || !admin || !rateLimitedUpstream || !fallbackUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + const statuses = await admin!.listModelStatuses(); + const row = statuses.find((item) => item.display_name === "retry-429-primary"); + return row?.status === "healthy" && row?.last_check_status === 429; + }); + + const completion = await client.chat.completions.create({ + model: "retry-429-router", + messages: [{ role: "user", content: "429 request path should still retry" }], + }); + + expect(completion.choices[0]?.message.content).toBe("request-path 429 fallback"); + + const statuses = await admin.listModelStatuses(); + const row = statuses.find((item) => item.display_name === "retry-429-primary")!; + expect(row.status).toBe("cooldown"); + // After PR #268 H1/M1 contract: cooldown reason reflects the + // upstream HTTP semantic, not the gateway-internal retry flag. + // 429 maps to upstream_rate_limited regardless of whether + // retry_on_429 was true (here) or false (see m1_decoupling test). + expect(row.status_reason).toBe("upstream_rate_limited"); + }); +}); diff --git a/tests/e2e/src/cases/routing-strategies-e2e.test.ts b/tests/e2e/src/cases/routing-strategies-e2e.test.ts index 83864337..1c5355db 100644 --- a/tests/e2e/src/cases/routing-strategies-e2e.test.ts +++ b/tests/e2e/src/cases/routing-strategies-e2e.test.ts @@ -129,17 +129,12 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - await waitConfigPropagation(async () => { - try { - const probe = await client.chat.completions.create({ - model: "routing-retry-virtual", - messages: [{ role: "user", content: "ready-routing-retry" }], - }); - return probe.choices[0]?.message.content === "after retries"; - } catch { - return false; - } - }); + // Bare propagation wait — probing the virtual would warm the + // primary's cooldown (post-PR #268 contract: every retryable + // upstream failure cools down the failing direct target) and + // throw off the per-target hit counts below. The secondary's + // direct readiness was already established above. + await waitConfigPropagation(); const primaryBaseline = primary.receivedRequests.length; const secondaryBaseline = secondary.receivedRequests.length; @@ -205,17 +200,11 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - await waitConfigPropagation(async () => { - try { - const probe = await client.chat.completions.create({ - model: "routing-429-virtual", - messages: [{ role: "user", content: "ready-routing-429" }], - }); - return probe.choices[0]?.message.content === "429 fallback worked"; - } catch { - return false; - } - }); + // Bare propagation wait — probing the virtual would warm the + // primary's 429 cooldown (post-PR #268: 429 cools down even + // when retry_on_429=true, since cooldown is independent of + // retry) and zero out the per-target counts. + await waitConfigPropagation(); const primaryBaseline = primary.receivedRequests.length; const secondaryBaseline = secondary.receivedRequests.length; @@ -388,17 +377,11 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - await waitConfigPropagation(async () => { - try { - const probe = await client.chat.completions.create({ - model: "routing-weighted-virtual", - messages: [{ role: "user", content: "ready-routing-weighted" }], - }); - return probe.choices[0]?.message.content === "weighted fallback worked"; - } catch { - return false; - } - }); + // Bare propagation wait — probing the virtual would warm the + // primary's 502 cooldown (post-PR #268 contract) and skew the + // per-target hit counts. Both direct models' readiness was + // already established above. + await waitConfigPropagation(); const beforeBaseline = zeroWeightBefore.receivedRequests.length; const primaryBaseline = weightedPrimary.receivedRequests.length; diff --git a/tests/e2e/src/cases/runtime-mixed-filtering-e2e.test.ts b/tests/e2e/src/cases/runtime-mixed-filtering-e2e.test.ts new file mode 100644 index 00000000..e62d372a --- /dev/null +++ b/tests/e2e/src/cases/runtime-mixed-filtering-e2e.test.ts @@ -0,0 +1,187 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-runtime-mixed-filtering-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +describe("runtime mixed filtering e2e", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let unhealthyUpstream: OpenAiUpstream | undefined; + let cooldownUpstream: OpenAiUpstream | undefined; + let healthyUpstream: OpenAiUpstream | undefined; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + unhealthyUpstream = await startOpenAiUpstream({ + status: 503, + errorBody: { error: { message: "unhealthy target", type: "server_error" } }, + }); + cooldownUpstream = await startOpenAiUpstream({ + scriptedResponses: [ + { + status: 502, + errorBody: { error: { message: "cooldown target failed", type: "server_error" } }, + }, + { + nonStreamBody: { + id: "cmpl-cooldown-recovered", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "should not be selected second" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }, + ], + }); + healthyUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-healthy-mixed", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "healthy candidate won" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const unhealthyPk = await admin.createProviderKey({ + display_name: "mixed-unhealthy-pk", + secret: "sk-mock", + api_base: `${unhealthyUpstream.baseUrl}/v1`, + }); + const cooldownPk = await admin.createProviderKey({ + display_name: "mixed-cooldown-pk", + secret: "sk-mock", + api_base: `${cooldownUpstream.baseUrl}/v1`, + }); + const healthyPk = await admin.createProviderKey({ + display_name: "mixed-healthy-pk", + secret: "sk-mock", + api_base: `${healthyUpstream.baseUrl}/v1`, + }); + + await admin.createModel({ + display_name: "mixed-unhealthy", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: unhealthyPk.id, + background_model_check: { + enabled: true, + interval_seconds: 5, + timeout_seconds: 10, + prompt: "Respond with OK", + max_tokens: 8, + ignore_statuses: [408, 429], + stale_after_seconds: 90, + }, + }); + await admin.createModel({ + display_name: "mixed-cooldown", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: cooldownPk.id, + }); + await admin.createModel({ + display_name: "mixed-healthy", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: healthyPk.id, + }); + + await admin.createModel({ + display_name: "mixed-router", + routing: { + strategy: "failover", + targets: [ + { model: "mixed-unhealthy" }, + { model: "mixed-cooldown" }, + { model: "mixed-healthy" }, + ], + max_fallbacks: 2, + }, + }); + + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["mixed-router", "mixed-cooldown", "mixed-healthy"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await unhealthyUpstream?.close(); + await cooldownUpstream?.close(); + await healthyUpstream?.close(); + }); + + test("routing skips unhealthy first, then cooldown, and lands on healthy candidate", async (ctx) => { + if (!etcdReachable || !app || !admin || !unhealthyUpstream || !cooldownUpstream || !healthyUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + const statuses = await admin!.listModelStatuses(); + const unhealthy = statuses.find((row) => row.display_name === "mixed-unhealthy"); + return unhealthy?.status === "unhealthy"; + }); + + const first = await client.chat.completions.create({ + model: "mixed-router", + messages: [{ role: "user", content: "trip cooldown in middle candidate" }], + }); + expect(first.choices[0]?.message.content).toBe("healthy candidate won"); + + const unhealthyBaseline = unhealthyUpstream.receivedRequests.length; + const cooldownBaseline = cooldownUpstream.receivedRequests.length; + const healthyBaseline = healthyUpstream.receivedRequests.length; + + const second = await client.chat.completions.create({ + model: "mixed-router", + messages: [{ role: "user", content: "mixed filtering second pass" }], + }); + expect(second.choices[0]?.message.content).toBe("healthy candidate won"); + + expect(unhealthyUpstream.receivedRequests.length - unhealthyBaseline).toBe(0); + expect(cooldownUpstream.receivedRequests.length - cooldownBaseline).toBe(0); + expect(healthyUpstream.receivedRequests.length - healthyBaseline).toBe(1); + }); +}); diff --git a/tests/e2e/src/cases/runtime-status-e2e.test.ts b/tests/e2e/src/cases/runtime-status-e2e.test.ts new file mode 100644 index 00000000..4524b947 --- /dev/null +++ b/tests/e2e/src/cases/runtime-status-e2e.test.ts @@ -0,0 +1,186 @@ +import { createHash } from "node:crypto"; +import OpenAI from "openai"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + AdminClient, + EtcdClient, + spawnApp, + startOpenAiUpstream, + waitConfigPropagation, + type OpenAiUpstream, + type SpawnedApp, +} from "../harness/index.js"; + +const CALLER_PLAINTEXT = "sk-runtime-status-caller"; +const CALLER_KEY_HASH = createHash("sha256") + .update(CALLER_PLAINTEXT) + .digest("hex"); + +describe("runtime status e2e", () => { + let app: SpawnedApp | undefined; + let admin: AdminClient | undefined; + let etcdReachable = false; + let flakyUpstream: OpenAiUpstream | undefined; + let stableUpstream: OpenAiUpstream | undefined; + let flakyModelID = ""; + let stableModelID = ""; + let routerModelID = ""; + + beforeAll(async () => { + etcdReachable = await new EtcdClient().ping(); + if (!etcdReachable) return; + + flakyUpstream = await startOpenAiUpstream({ + scriptedResponses: [ + { + status: 502, + errorBody: { error: { message: "temporary upstream failure", type: "server_error" } }, + }, + { + nonStreamBody: { + id: "cmpl-flaky-recovered", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "flaky recovered" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }, + ], + }); + stableUpstream = await startOpenAiUpstream({ + nonStreamBody: { + id: "cmpl-stable", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "stable fallback" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + }); + + app = await spawnApp(); + admin = new AdminClient(app.adminUrl, app.adminKey); + + const flakyPk = await admin.createProviderKey({ + display_name: "runtime-flaky-pk", + secret: "sk-mock", + api_base: `${flakyUpstream.baseUrl}/v1`, + }); + const stablePk = await admin.createProviderKey({ + display_name: "runtime-stable-pk", + secret: "sk-mock", + api_base: `${stableUpstream.baseUrl}/v1`, + }); + const flakyModel = await admin.createModel({ + display_name: "runtime-flaky", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: flakyPk.id, + }); + flakyModelID = flakyModel.id; + const stableModel = await admin.createModel({ + display_name: "runtime-stable", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: stablePk.id, + }); + stableModelID = stableModel.id; + const routerModel = await admin.createModel({ + display_name: "runtime-router", + routing: { + strategy: "failover", + targets: [{ model: "runtime-flaky" }, { model: "runtime-stable" }], + retries: 0, + max_fallbacks: 1, + }, + }); + routerModelID = routerModel.id; + await admin.createApiKey({ + key_hash: CALLER_KEY_HASH, + allowed_models: ["runtime-router", "runtime-stable", "runtime-flaky"], + }); + }); + + afterAll(async () => { + await app?.exit(); + await flakyUpstream?.close(); + await stableUpstream?.close(); + }); + + test("retryable failure cools down the direct model, routing skips it, and admin surfaces runtime status", async (ctx) => { + if (!etcdReachable || !app || !admin || !flakyUpstream || !stableUpstream) { + ctx.skip(); + return; + } + + const client = new OpenAI({ + apiKey: CALLER_PLAINTEXT, + baseURL: `${app.proxyUrl}/v1`, + maxRetries: 0, + }); + + await waitConfigPropagation(async () => { + try { + const probe = await client.chat.completions.create({ + model: "runtime-stable", + messages: [{ role: "user", content: "ready-runtime-stable" }], + }); + return probe.choices[0]?.message.content === "stable fallback"; + } catch { + return false; + } + }); + + await waitConfigPropagation(); + + const first = await client.chat.completions.create({ + model: "runtime-router", + messages: [{ role: "user", content: "trip cooldown" }], + }); + expect(first.choices[0]?.message.content).toBe("stable fallback"); + expect(flakyUpstream.receivedRequests.length).toBeGreaterThanOrEqual(1); + expect(stableUpstream.receivedRequests.length).toBeGreaterThanOrEqual(1); + + const statusesAfterFirst = await admin.listModelStatuses(); + const flakyAfterFirst = statusesAfterFirst.find((row) => row.id === flakyModelID)!; + expect(flakyAfterFirst.status).toBe("cooldown"); + // After PR #268 H1 contract: cooldown reason is per-error-category. + // A 502 from the upstream maps to `upstream_server_error`. + expect(flakyAfterFirst.status_reason).toBe("upstream_server_error"); + + const flakyBaseline = flakyUpstream.receivedRequests.length; + const stableBaseline = stableUpstream.receivedRequests.length; + + const second = await client.chat.completions.create({ + model: "runtime-router", + messages: [{ role: "user", content: "skip cooled target" }], + }); + expect(second.choices[0]?.message.content).toBe("stable fallback"); + expect(flakyUpstream.receivedRequests.length - flakyBaseline).toBe(0); + expect(stableUpstream.receivedRequests.length - stableBaseline).toBe(1); + + const statuses = await admin.listModelStatuses(); + const flaky = statuses.find((row) => row.id === flakyModelID)!; + const stable = statuses.find((row) => row.id === stableModelID)!; + const router = statuses.find((row) => row.id === routerModelID)!; + + expect(flaky.status).toBe("cooldown"); + expect(flaky.status_reason).toBe("upstream_server_error"); + expect(flaky.cooldown_until).toBeTruthy(); + expect(stable.status).toBe("healthy"); + expect(router.status).toBe("not_applicable"); + }); +}); diff --git a/tests/e2e/src/harness/upstream-openai.ts b/tests/e2e/src/harness/upstream-openai.ts index 3a5cfd6b..262432f2 100644 --- a/tests/e2e/src/harness/upstream-openai.ts +++ b/tests/e2e/src/harness/upstream-openai.ts @@ -18,6 +18,12 @@ export interface OpenAiUpstreamOptions { disconnectAfterEvents?: number; /** Per-request response script; used in order before static opts. */ scriptedResponses?: OpenAiUpstreamStep[]; + /** + * Extra response headers to set on every reply. Used by the cooldown + * contract tests to assert that the gateway honors `Retry-After` + * from the upstream when computing the cooldown TTL. + */ + responseHeaders?: Record; } export interface OpenAiUpstreamStep { @@ -28,6 +34,8 @@ export interface OpenAiUpstreamStep { status?: number; errorBody?: unknown; disconnectAfterEvents?: number; + /** Extra response headers, same semantics as on the top-level options. */ + responseHeaders?: Record; } export interface OpenAiUpstream { @@ -72,6 +80,11 @@ export async function startOpenAiUpstream( if (step.responseDelayMs) await sleep(step.responseDelayMs); + const extraHeaders = { ...(opts.responseHeaders ?? {}), ...(step.responseHeaders ?? {}) }; + for (const [k, v] of Object.entries(extraHeaders)) { + res.setHeader(k, v); + } + const status = step.status ?? 200; if (status >= 400) { res.statusCode = status; From 2ea8a172304971e4c0b57883288788a75e6825d1 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Thu, 14 May 2026 08:26:57 +0800 Subject: [PATCH 3/4] fix(cooldown): address audit findings on prior cooldown commit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent audit on commit 6b109ed flagged these blockers: H-1: cooldown was wired only into chat.rs. Anthropic /v1/messages, OpenAI /v1/responses, audio/speech, audio/transcriptions, and /v1/rerank all silently dropped the cooldown signal. A 401 from Anthropic via /v1/messages would never take that target out of rotation. - Extract decide_cooldown + reason_for_status into crate::cooldown - Call from messages.rs (both passthrough + cross_provider_dispatch paths), responses.rs, audio.rs (speech + transcribe), rerank.rs H-2: routing-strategies-e2e readiness was weakened to bare waitConfigPropagation(), losing the verification that the virtual record had reached the DP. Restored by gating on admin.listModels() finding the virtual entry — proves propagation without sending dispatcher traffic that would warm cooldown. M-1: H2 e2e relied on undocumented serde-default SystemTime shape. Added wire-shape assertion so a future serialization change fails with a clear message instead of NaN. M-2: cooldown.default_seconds: 0 was schema-accepted but at runtime silently re-derived as max_seconds (default 600s) — the opposite of the operator's likely intent. Now treated as "disable cooldown on this model" with new unit test pinning the contract. M-5: OnAllFilteredPolicy::Fail rustdoc claimed Retry-After is "derived from the next cooldown expiry"; impl returns a fixed 30s fallback. Doc corrected to match the actual constant-fallback behavior with rationale. Drive-by: tests/e2e/src/harness/admin.ts listModels was decoding the response as `{items: [...]}` but the admin endpoint returns a bare array of ResourceEntry. The helper had no callers before now, so the bug was latent. cargo test --workspace: 689 pass / 0 fail (+3 for the new cooldown.rs unit tests). cargo clippy --workspace --all-targets -- -D warnings: clean. cargo fmt --all -- --check: clean. e2e: 92 pass / 0 fail across 47 files against real backend. Audit-flagged LOW items deferred: HTTP-date Retry-After parsing and CooldownConfig::default() docs. M-3 (BridgeError field addition) and M-4 (interval_seconds floor 1→5) are doc-only — call them out as breaking-on-config-reload in the PR body for upgrade notes. --- crates/aisix-core/src/models/routing.rs | 15 +- crates/aisix-proxy/src/audio.rs | 40 +++-- crates/aisix-proxy/src/chat.rs | 85 +-------- crates/aisix-proxy/src/cooldown.rs | 165 ++++++++++++++++++ crates/aisix-proxy/src/lib.rs | 1 + crates/aisix-proxy/src/messages.rs | 48 +++-- crates/aisix-proxy/src/rerank.rs | 20 ++- crates/aisix-proxy/src/responses.rs | 20 ++- .../src/cases/cooldown-contract-e2e.test.ts | 13 +- .../src/cases/routing-strategies-e2e.test.ts | 55 ++++-- tests/e2e/src/harness/admin.ts | 8 +- 11 files changed, 325 insertions(+), 145 deletions(-) create mode 100644 crates/aisix-proxy/src/cooldown.rs diff --git a/crates/aisix-core/src/models/routing.rs b/crates/aisix-core/src/models/routing.rs index 42b91960..53c37370 100644 --- a/crates/aisix-core/src/models/routing.rs +++ b/crates/aisix-core/src/models/routing.rs @@ -67,13 +67,22 @@ impl RoutingTarget { #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum OnAllFilteredPolicy { - /// Return 503 with a Retry-After hint derived from the next - /// cooldown expiry (or a fixed fallback if every candidate is - /// background-unhealthy with no cooldown timer). Default. + /// Return 503 with a fixed Retry-After hint (currently 30 seconds — + /// see `FALLBACK_ALL_UNHEALTHY_RETRY_AFTER` in + /// `crates/aisix-proxy/src/chat.rs`). Default. + /// + /// The hint is intentionally coarse: by the time the filter + /// reaches the all-filtered branch, every candidate is + /// background-unhealthy with no live cooldown timer (cooldown + /// candidates are returned via the Selected branch one tier up). + /// A future version may derive the hint from probe metadata; the + /// current contract is a flat fallback. #[default] Fail, /// Send to the original candidate list anyway, in declaration /// order. Preserves availability over caller-facing correctness. + /// Use only when the operator explicitly accepts the risk of + /// sending traffic to a target the gateway just probed as broken. OriginalOrder, } diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index f00add8b..55b6fcfc 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -343,16 +343,22 @@ async fn multipart_dispatch( let s = status.as_u16(); let retry_after = aisix_gateway::parse_retry_after(resp.headers()); let msg = resp.text().await.unwrap_or_default(); - return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::upstream_status_with_retry_after( - s, - msg.chars().take(1024).collect::(), - retry_after, - ), - )); + let err = aisix_gateway::BridgeError::upstream_status_with_retry_after( + s, + msg.chars().take(1024).collect::(), + retry_after, + ); + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state + .runtime_status + .mark_cooldown(&model_entry.id, ttl, reason); + } + return Err(ProxyError::Bridge(err)); } state.health.record_success(&model_name); + state.runtime_status.mark_healthy(&model_entry.id); // Relay response headers that matter for the client. let upstream_headers = resp.headers().clone(); @@ -425,16 +431,22 @@ async fn speech_dispatch( let s = status.as_u16(); let retry_after = aisix_gateway::parse_retry_after(resp.headers()); let msg = resp.text().await.unwrap_or_default(); - return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::upstream_status_with_retry_after( - s, - msg.chars().take(1024).collect::(), - retry_after, - ), - )); + let err = aisix_gateway::BridgeError::upstream_status_with_retry_after( + s, + msg.chars().take(1024).collect::(), + retry_after, + ); + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state + .runtime_status + .mark_cooldown(&model_entry.id, ttl, reason); + } + return Err(ProxyError::Bridge(err)); } state.health.record_success(&model_name); + state.runtime_status.mark_healthy(&model_entry.id); let upstream_headers = resp.headers().clone(); let body_bytes = resp diff --git a/crates/aisix-proxy/src/chat.rs b/crates/aisix-proxy/src/chat.rs index 98613c7b..df4dc8c3 100644 --- a/crates/aisix-proxy/src/chat.rs +++ b/crates/aisix-proxy/src/chat.rs @@ -54,87 +54,10 @@ struct AttemptModel { model: aisix_core::Model, } -/// Decide whether a bridge error should trigger cooldown on the -/// failing direct model, and for how long. -/// -/// Cooldown is **independent** of `is_retryable`. A 401 (auth failure) -/// is non-retryable — retrying the same target in the current request -/// is pointless — but it should still cooldown because every -/// subsequent request that lands on the same target will see the same -/// 401. Conversely, a transient timeout may be retryable AND should -/// also cooldown. -/// -/// `Retry-After` from upstream is honored when `honor_retry_after` -/// is set (default true), clamped to `max_seconds`. -fn decide_cooldown( - err: &BridgeError, - cfg: Option<&aisix_core::CooldownConfig>, -) -> Option<(Duration, &'static str)> { - use aisix_core::CooldownConfig; - let default_cfg = CooldownConfig::default(); - let cfg = cfg.unwrap_or(&default_cfg); - if !cfg.enabled_or_default() { - return None; - } - let max = Duration::from_secs(cfg.max_seconds_or_default()); - let default_ttl = Duration::from_secs(cfg.default_seconds_or_default()); - let clamp = |d: Duration| -> Duration { - let bounded = d.min(max); - // 0 means "disabled" via field; if we got here, default_ttl - // should never be 0 — fall back to max if someone set - // default_seconds=0 by mistake but didn't disable. - if bounded.is_zero() { - max - } else { - bounded - } - }; - - match err { - BridgeError::UpstreamStatus { - status, - retry_after, - .. - } => { - let triggers = cfg.effective_trigger_statuses(); - if !triggers.contains(status) { - return None; - } - let ttl = if cfg.honor_retry_after_or_default() { - retry_after.map(clamp).unwrap_or(default_ttl) - } else { - default_ttl - }; - Some((ttl, cooldown_reason_for_status(*status))) - } - BridgeError::Timeout { .. } if cfg.trigger_on_timeout_or_default() => { - Some((default_ttl, "request_timeout")) - } - BridgeError::Transport(_) | BridgeError::StreamAborted - if cfg.trigger_on_transport_or_default() => - { - Some((default_ttl, "transport_error")) - } - BridgeError::UpstreamDecode(_) if cfg.trigger_on_transport_or_default() => { - Some((default_ttl, "upstream_decode_error")) - } - // Config errors mean WE are misconfigured (missing provider, bad - // bridge registration). Cooling down doesn't help; let it - // surface and operator fixes the snapshot. - BridgeError::Config(_) => None, - _ => None, - } -} - -fn cooldown_reason_for_status(status: u16) -> &'static str { - match status { - 401 => "upstream_auth_failure", - 408 => "upstream_request_timeout", - 429 => "upstream_rate_limited", - 500..=599 => "upstream_server_error", - _ => "upstream_status_failure", - } -} +// Per-attempt cooldown decision lives in `crate::cooldown` so every +// dispatch path (chat, messages, responses, audio, rerank) shares the +// same logic. See cooldown.rs for the audit context (#264 H-1). +use crate::cooldown::decide_cooldown; pub async fn chat_completions( State(state): State, diff --git a/crates/aisix-proxy/src/cooldown.rs b/crates/aisix-proxy/src/cooldown.rs new file mode 100644 index 00000000..9b3c761d --- /dev/null +++ b/crates/aisix-proxy/src/cooldown.rs @@ -0,0 +1,165 @@ +//! Cooldown decision helper shared across every dispatch path. +//! +//! The proxy exposes more than one upstream endpoint family: +//! - `/v1/chat/completions` (chat.rs) +//! - `/v1/messages` (messages.rs — Anthropic-shape) +//! - `/v1/responses` (responses.rs — OpenAI Responses API passthrough) +//! - `/v1/audio/{speech,transcriptions,translations}` (audio.rs) +//! - `/v1/rerank` (rerank.rs) +//! +//! Every one of those paths surfaces upstream failures as +//! [`BridgeError`]. The runtime-status layer (see [`crate::health`]) +//! and the cross-request cooldown contract pinned by issue #264 apply +//! to **all** of them: a 401 from Anthropic via `/v1/messages` should +//! take that direct model out of rotation for the next request just +//! as a 401 via `/v1/chat/completions` does. +//! +//! This module owns the per-attempt cooldown decision. Each dispatch +//! path calls [`decide_cooldown`] after building a `BridgeError` and +//! before returning the error to the client. If a cooldown is +//! warranted, the caller hands the result to +//! [`crate::health::ModelRuntimeStatusTracker::mark_cooldown`]. +//! +//! Keeping this logic in one place (rather than per-dispatch) is what +//! prevents the H-1 audit class of bug: a new dispatch path silently +//! forgets to cool down because the routing-loop in chat.rs is where +//! cooldown lived historically. + +use std::time::Duration; + +use aisix_core::CooldownConfig; +use aisix_gateway::BridgeError; + +/// Decide whether a bridge error should trigger cooldown on the +/// failing direct model, and for how long. +/// +/// Cooldown is **independent** of `is_retryable`. A 401 (auth failure) +/// is non-retryable — retrying the same target in the current request +/// is pointless — but it should still cooldown because every +/// subsequent request that lands on the same target will see the same +/// 401. Conversely, a transient timeout may be retryable AND should +/// also cooldown. +/// +/// `Retry-After` from upstream is honored when `honor_retry_after` +/// is set (default true), clamped to `max_seconds`. A configured +/// `default_seconds: 0` is treated as "do not cool down on this +/// category" — matches the operator's likely intent of "disable +/// cooldown TTL" (M-2 audit on PR #268). +pub fn decide_cooldown( + err: &BridgeError, + cfg: Option<&CooldownConfig>, +) -> Option<(Duration, &'static str)> { + let default_cfg = CooldownConfig::default(); + let cfg = cfg.unwrap_or(&default_cfg); + if !cfg.enabled_or_default() { + return None; + } + + let default_secs = cfg.default_seconds_or_default(); + // A configured `default_seconds: 0` is a per-category disable. + // The schema allows `minimum: 0` for this reason — operators that + // want NO cooldown on any failure of this model set it to 0. + if default_secs == 0 { + return None; + } + + let max = Duration::from_secs(cfg.max_seconds_or_default()); + let default_ttl = Duration::from_secs(default_secs); + let clamp = |d: Duration| -> Duration { d.min(max) }; + + match err { + BridgeError::UpstreamStatus { + status, + retry_after, + .. + } => { + let triggers = cfg.effective_trigger_statuses(); + if !triggers.contains(status) { + return None; + } + let ttl = if cfg.honor_retry_after_or_default() { + retry_after.map(clamp).unwrap_or(default_ttl) + } else { + default_ttl + }; + Some((ttl, reason_for_status(*status))) + } + BridgeError::Timeout { .. } if cfg.trigger_on_timeout_or_default() => { + Some((default_ttl, "request_timeout")) + } + BridgeError::Transport(_) | BridgeError::StreamAborted + if cfg.trigger_on_transport_or_default() => + { + Some((default_ttl, "transport_error")) + } + BridgeError::UpstreamDecode(_) if cfg.trigger_on_transport_or_default() => { + Some((default_ttl, "upstream_decode_error")) + } + // Config errors mean WE are misconfigured (missing provider, bad + // bridge registration). Cooling down doesn't help; let it + // surface and operator fixes the snapshot. + BridgeError::Config(_) => None, + _ => None, + } +} + +/// Map an HTTP status to a stable `status_reason` token surfaced on +/// `/admin/v1/models/status`. Kept narrow and operator-friendly — +/// callers should not synthesize their own reason strings. +pub fn reason_for_status(status: u16) -> &'static str { + match status { + 401 => "upstream_auth_failure", + 408 => "upstream_request_timeout", + 429 => "upstream_rate_limited", + 500..=599 => "upstream_server_error", + _ => "upstream_status_failure", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + fn upstream(status: u16) -> BridgeError { + BridgeError::upstream_status(status, "boom") + } + + #[test] + fn default_seconds_zero_disables_cooldown_entirely() { + // M-2 contract: 0 = off for any error category. Previously + // (pre-audit fix) 0 silently re-derived as the max_seconds + // cap, which surprised operators who expected disable. + let cfg = CooldownConfig { + default_seconds: Some(0), + ..Default::default() + }; + assert!(decide_cooldown(&upstream(429), Some(&cfg)).is_none()); + assert!(decide_cooldown(&upstream(503), Some(&cfg)).is_none()); + assert!(decide_cooldown(&BridgeError::Timeout { elapsed_ms: 1 }, Some(&cfg)).is_none()); + } + + #[test] + fn config_disabled_skips_cooldown() { + let cfg = CooldownConfig { + enabled: Some(false), + ..Default::default() + }; + assert!(decide_cooldown(&upstream(429), Some(&cfg)).is_none()); + } + + #[test] + fn honor_retry_after_clamps_to_max_seconds() { + let cfg = CooldownConfig { + max_seconds: Some(60), + ..Default::default() + }; + let err = BridgeError::upstream_status_with_retry_after( + 429, + "rl", + Some(Duration::from_secs(100_000)), + ); + let (ttl, _) = decide_cooldown(&err, Some(&cfg)).unwrap(); + assert_eq!(ttl, Duration::from_secs(60)); + } +} diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index 60fa062c..df5f0e5f 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -30,6 +30,7 @@ pub mod background; pub mod budget; mod chat; mod completions; +pub(crate) mod cooldown; mod dispatch; mod embeddings; mod error; diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index d71c15d3..8656eb44 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -183,6 +183,7 @@ async fn dispatch( state, body, model, + &model_entry.id, &pk_entry.value, &model_name, request_id, @@ -238,13 +239,23 @@ async fn dispatch( } else { message }; - return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::upstream_status_with_retry_after( - status_u16, - truncated, - retry_after, - ), - )); + let err = aisix_gateway::BridgeError::upstream_status_with_retry_after( + status_u16, + truncated, + retry_after, + ); + // Apply the cross-request cooldown contract to the + // Anthropic-passthrough path too — without this, a 401 / 429 / + // 5xx via /v1/messages would never mark the direct model and + // subsequent requests would keep hitting the same broken + // upstream. See `crate::cooldown` for the shared decision. + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state + .runtime_status + .mark_cooldown(&model_entry.id, ttl, reason); + } + return Err(ProxyError::Bridge(err)); } // Update health tracker on success. @@ -377,6 +388,7 @@ async fn cross_provider_dispatch( state: &ProxyState, body: &Value, model: &aisix_core::Model, + model_id: &str, provider_key: &aisix_core::ProviderKey, model_name: &str, request_id: &str, @@ -428,11 +440,16 @@ async fn cross_provider_dispatch( let provider_label = format!("{provider:?}").to_lowercase(); if is_stream { - let upstream = bridge - .chat_stream(&chat, &ctx) - .await - .map_err(ProxyError::Bridge)?; + let upstream = bridge.chat_stream(&chat, &ctx).await.map_err(|err| { + if let Some((ttl, reason)) = + crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state.runtime_status.mark_cooldown(model_id, ttl, reason); + } + ProxyError::Bridge(err) + })?; state.health.record_success(model_name); + state.runtime_status.mark_healthy(model_id); let message_id = format!("msg_{}", Uuid::new_v4().simple()); let encoder = AnthropicSseEncoder::new(message_id, model_name, 0); @@ -466,8 +483,15 @@ async fn cross_provider_dispatch( } // Non-streaming. - let resp = bridge.chat(&chat, &ctx).await.map_err(ProxyError::Bridge)?; + let resp = bridge.chat(&chat, &ctx).await.map_err(|err| { + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state.runtime_status.mark_cooldown(model_id, ttl, reason); + } + ProxyError::Bridge(err) + })?; state.health.record_success(model_name); + state.runtime_status.mark_healthy(model_id); let metrics = AnthropicUsageMetrics { prompt_tokens: resp.usage.prompt_tokens, diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 41c4fadb..bf71ab70 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -189,16 +189,22 @@ async fn dispatch( let status_u16 = status.as_u16(); let retry_after = aisix_gateway::parse_retry_after(upstream_resp.headers()); let message = upstream_resp.text().await.unwrap_or_default(); - return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::upstream_status_with_retry_after( - status_u16, - message.chars().take(1024).collect::(), - retry_after, - ), - )); + let err = aisix_gateway::BridgeError::upstream_status_with_retry_after( + status_u16, + message.chars().take(1024).collect::(), + retry_after, + ); + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state + .runtime_status + .mark_cooldown(&model_entry.id, ttl, reason); + } + return Err(ProxyError::Bridge(err)); } state.health.record_success(&model_name); + state.runtime_status.mark_healthy(&model_entry.id); let upstream_headers = upstream_resp.headers().clone(); let body_bytes = upstream_resp diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index 5fdd1e99..e01c5c51 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -161,16 +161,22 @@ async fn dispatch( let status_u16 = status.as_u16(); let retry_after = aisix_gateway::parse_retry_after(upstream_resp.headers()); let message = upstream_resp.text().await.unwrap_or_default(); - return Err(ProxyError::Bridge( - aisix_gateway::BridgeError::upstream_status_with_retry_after( - status_u16, - message.chars().take(1024).collect::(), - retry_after, - ), - )); + let err = aisix_gateway::BridgeError::upstream_status_with_retry_after( + status_u16, + message.chars().take(1024).collect::(), + retry_after, + ); + if let Some((ttl, reason)) = crate::cooldown::decide_cooldown(&err, model.cooldown.as_ref()) + { + state + .runtime_status + .mark_cooldown(&model_entry.id, ttl, reason); + } + return Err(ProxyError::Bridge(err)); } state.health.record_success(&model_name); + state.runtime_status.mark_healthy(&model_entry.id); let provider_label = "openai".to_string(); diff --git a/tests/e2e/src/cases/cooldown-contract-e2e.test.ts b/tests/e2e/src/cases/cooldown-contract-e2e.test.ts index 13ffbd38..d4565574 100644 --- a/tests/e2e/src/cases/cooldown-contract-e2e.test.ts +++ b/tests/e2e/src/cases/cooldown-contract-e2e.test.ts @@ -468,9 +468,16 @@ describe("cooldown contract (H2) — Retry-After header from upstream drives TTL // The reported cooldown_until should be ~180s in the future. // 30s default is the wrong answer — that would prove H2 isn't // wired. `cooldown_until` is serialized via SystemTime's default - // serde shape: `{secs_since_epoch, nanos_since_epoch}`. - const cooldownUntil = row.cooldown_until as { secs_since_epoch: number }; - const cooldownUntilMs = cooldownUntil.secs_since_epoch * 1000; + // serde shape: `{secs_since_epoch, nanos_since_epoch}`. Assert + // the wire shape first so a future change to the admin + // serialization (e.g. RFC3339 strings) fails this test with a + // clear message rather than a silent NaN parse. + const cooldownUntil = row.cooldown_until as { secs_since_epoch?: unknown }; + expect( + typeof cooldownUntil.secs_since_epoch, + "cooldown_until shape changed — update this test to match the new admin wire format", + ).toBe("number"); + const cooldownUntilMs = (cooldownUntil.secs_since_epoch as number) * 1000; const horizonMs = cooldownUntilMs - before; expect(horizonMs).toBeGreaterThan(120_000); expect(horizonMs).toBeLessThan(240_000); diff --git a/tests/e2e/src/cases/routing-strategies-e2e.test.ts b/tests/e2e/src/cases/routing-strategies-e2e.test.ts index 1c5355db..379c6bbd 100644 --- a/tests/e2e/src/cases/routing-strategies-e2e.test.ts +++ b/tests/e2e/src/cases/routing-strategies-e2e.test.ts @@ -129,12 +129,20 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - // Bare propagation wait — probing the virtual would warm the - // primary's cooldown (post-PR #268 contract: every retryable - // upstream failure cools down the failing direct target) and - // throw off the per-target hit counts below. The secondary's - // direct readiness was already established above. - await waitConfigPropagation(); + // Probing the virtual would warm the primary's cooldown + // (post-PR #268: every retryable upstream failure cools down the + // failing direct target) and zero out the per-target hit counts + // below. Instead, gate on the admin snapshot containing the + // virtual record — that proves the routing config has propagated + // to the DP without sending any traffic through the dispatcher. + await waitConfigPropagation(async () => { + try { + const models = await admin!.listModels(); + return models.some((m) => m.display_name === "routing-retry-virtual"); + } catch { + return false; + } + }); const primaryBaseline = primary.receivedRequests.length; const secondaryBaseline = secondary.receivedRequests.length; @@ -200,11 +208,18 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - // Bare propagation wait — probing the virtual would warm the - // primary's 429 cooldown (post-PR #268: 429 cools down even - // when retry_on_429=true, since cooldown is independent of - // retry) and zero out the per-target counts. - await waitConfigPropagation(); + // Gate on admin-snapshot presence rather than probing the + // virtual — probe would warm the primary's 429 cooldown and + // zero out per-target counts (post-PR #268: 429 cools down + // regardless of retry_on_429). + await waitConfigPropagation(async () => { + try { + const models = await admin!.listModels(); + return models.some((m) => m.display_name === "routing-429-virtual"); + } catch { + return false; + } + }); const primaryBaseline = primary.receivedRequests.length; const secondaryBaseline = secondary.receivedRequests.length; @@ -377,11 +392,19 @@ describe("routing strategies and retry behavior e2e", () => { maxRetries: 0, }); - // Bare propagation wait — probing the virtual would warm the - // primary's 502 cooldown (post-PR #268 contract) and skew the - // per-target hit counts. Both direct models' readiness was - // already established above. - await waitConfigPropagation(); + // Gate on admin-snapshot presence rather than probing the + // virtual — probe would warm the weighted primary's 502 cooldown + // and skew per-target hit counts. Both direct models' readiness + // was already established above; this confirms the routing record + // has reached the DP snapshot. + await waitConfigPropagation(async () => { + try { + const models = await admin!.listModels(); + return models.some((m) => m.display_name === "routing-weighted-virtual"); + } catch { + return false; + } + }); const beforeBaseline = zeroWeightBefore.receivedRequests.length; const primaryBaseline = weightedPrimary.receivedRequests.length; diff --git a/tests/e2e/src/harness/admin.ts b/tests/e2e/src/harness/admin.ts index bae457e2..0de44c78 100644 --- a/tests/e2e/src/harness/admin.ts +++ b/tests/e2e/src/harness/admin.ts @@ -29,11 +29,15 @@ export class AdminClient { } async listModels(): Promise>> { - const res = await this.json<{ items?: Array<{ value: Record }> }>( + // GET /admin/v1/models returns a bare JSON array of + // ResourceEntry objects (`{id, value, revision}`). + // Callers downstream usually only care about the inner value + // (which carries `display_name`, `provider`, etc.), so unwrap it. + const entries = await this.json }>>( "GET", "/admin/v1/models", ); - return (res.items ?? []).map((entry) => entry.value); + return entries.map((entry) => entry.value); } async listModelStatuses(): Promise>> { From 7c4dfb237eb95793d5ce9bc357f57ad4e7df5e38 Mon Sep 17 00:00:00 2001 From: Ming Wen Date: Thu, 14 May 2026 08:42:46 +0800 Subject: [PATCH 4/4] fix(cooldown): cool down on transport/decode errors + Anthropic passthrough mark_healthy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-2 independent audit on commit 2ea8a17 caught that the cooldown wiring on the five non-chat dispatchers was incomplete: only the `!status.is_success()` branch ran `decide_cooldown`. The `?`-escape paths on `.send().await` (transport / TCP reset / DNS) and on `.json()` / `.bytes()` decode skipped cooldown entirely. So a hard upstream outage on /v1/messages — exactly the H-1 failure mode the prior commit claimed to fix — would still leave the target healthy in routing. Round-2 HIGH: - Add `crate::cooldown::note_failure(tracker, model_id, cfg, err)` that runs the cooldown decision and marks the tracker, returning the BridgeError unchanged so call sites can keep using `?`. - Wire it into every `.map_err(|e| BridgeError::Transport(...))` and `.map_err(|e| BridgeError::UpstreamDecode(...))` on messages.rs (2 sites), responses.rs (2 sites), rerank.rs (2 sites), audio.rs (4 sites across multipart + speech). Round-2 MEDIUM: - messages.rs Anthropic-passthrough success path now also calls `state.runtime_status.mark_healthy(&model_entry.id)` so a target that recovers via /v1/messages exits `cooldown` cleanly on /admin/v1/models/status. Cross-provider sibling paths already did this; passthrough was the outlier. 3 new unit tests pin the contract: - `note_failure_marks_cooldown_for_transport_errors` - `note_failure_marks_cooldown_for_decode_errors` - `note_failure_no_op_when_cooldown_disabled` cargo test --workspace: 692 pass / 0 fail. cargo clippy --workspace --all-targets -- -D warnings: clean. cargo fmt --all -- --check: clean. e2e (real backend + mock upstreams): 92 pass / 0 fail across 47 files. Audit LOW (admin-snapshot readiness race window) deferred: practical gap is sub-100ms with 50ms poll cadence, no observed flakes. --- crates/aisix-proxy/src/audio.rs | 36 ++++++++++++-- crates/aisix-proxy/src/cooldown.rs | 73 +++++++++++++++++++++++++++++ crates/aisix-proxy/src/messages.rs | 31 ++++++++++-- crates/aisix-proxy/src/rerank.rs | 18 ++++++- crates/aisix-proxy/src/responses.rs | 18 ++++++- 5 files changed, 164 insertions(+), 12 deletions(-) diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index 55b6fcfc..ff918c39 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -335,7 +335,14 @@ async fn multipart_dispatch( .multipart(form) .send() .await - .map_err(|e| aisix_gateway::BridgeError::Transport(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::Transport(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let status = resp.status(); @@ -365,7 +372,14 @@ async fn multipart_dispatch( let body_bytes = resp .bytes() .await - .map_err(|e| aisix_gateway::BridgeError::UpstreamDecode(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::UpstreamDecode(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes)); @@ -423,7 +437,14 @@ async fn speech_dispatch( .json(&body) .send() .await - .map_err(|e| aisix_gateway::BridgeError::Transport(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::Transport(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let status = resp.status(); @@ -452,7 +473,14 @@ async fn speech_dispatch( let body_bytes = resp .bytes() .await - .map_err(|e| aisix_gateway::BridgeError::UpstreamDecode(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::UpstreamDecode(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let mut out = axum::response::Response::new(axum::body::Body::from(body_bytes)); diff --git a/crates/aisix-proxy/src/cooldown.rs b/crates/aisix-proxy/src/cooldown.rs index 9b3c761d..f5d7d22f 100644 --- a/crates/aisix-proxy/src/cooldown.rs +++ b/crates/aisix-proxy/src/cooldown.rs @@ -30,6 +30,8 @@ use std::time::Duration; use aisix_core::CooldownConfig; use aisix_gateway::BridgeError; +use crate::health::ModelRuntimeStatusTracker; + /// Decide whether a bridge error should trigger cooldown on the /// failing direct model, and for how long. /// @@ -103,6 +105,29 @@ pub fn decide_cooldown( } } +/// Record a failed dispatch attempt against `model_id`: run the +/// cooldown decision and, if it fires, mark the runtime tracker. +/// Returns the error unchanged so call sites can keep using `?`. +/// +/// This is the right entry point for every `.map_err` on the proxy's +/// upstream call paths — including the early `.send().await` and +/// body-decode failures that bypass the `!status.is_success()` branch. +/// The audit on PR #268 (round 2) caught exactly that gap: +/// transport / decode errors were threading through `?` without ever +/// hitting `mark_cooldown`, so a TCP-reset against Anthropic via +/// `/v1/messages` would leave the target healthy in routing. +pub fn note_failure( + tracker: &ModelRuntimeStatusTracker, + model_id: &str, + cfg: Option<&CooldownConfig>, + err: BridgeError, +) -> BridgeError { + if let Some((ttl, reason)) = decide_cooldown(&err, cfg) { + tracker.mark_cooldown(model_id, ttl, reason); + } + err +} + /// Map an HTTP status to a stable `status_reason` token surfaced on /// `/admin/v1/models/status`. Kept narrow and operator-friendly — /// callers should not synthesize their own reason strings. @@ -148,6 +173,54 @@ mod tests { assert!(decide_cooldown(&upstream(429), Some(&cfg)).is_none()); } + #[test] + fn note_failure_marks_cooldown_for_transport_errors() { + // Round-2 audit contract: a Transport error (TCP reset, DNS + // failure, …) must mark cooldown when trigger_on_transport + // is on (default). The non-status `?` paths in messages.rs / + // responses.rs / audio.rs / rerank.rs all route through here. + let tracker = ModelRuntimeStatusTracker::new(); + let err = BridgeError::Transport("connection refused".into()); + let returned = note_failure(&tracker, "m-1", None, err); + // Error returned unchanged. + assert!(matches!(returned, BridgeError::Transport(_))); + // Tracker now reports cooldown for this target. + assert_eq!( + tracker.status("m-1").status, + crate::health::RuntimeStatus::Cooldown + ); + assert_eq!( + tracker.status("m-1").status_reason.as_deref(), + Some("transport_error") + ); + } + + #[test] + fn note_failure_marks_cooldown_for_decode_errors() { + let tracker = ModelRuntimeStatusTracker::new(); + let err = BridgeError::UpstreamDecode("bad json".into()); + let _ = note_failure(&tracker, "m-1", None, err); + assert_eq!( + tracker.status("m-1").status, + crate::health::RuntimeStatus::Cooldown + ); + } + + #[test] + fn note_failure_no_op_when_cooldown_disabled() { + let tracker = ModelRuntimeStatusTracker::new(); + let cfg = CooldownConfig { + enabled: Some(false), + ..Default::default() + }; + let err = BridgeError::Transport("nope".into()); + let _ = note_failure(&tracker, "m-1", Some(&cfg), err); + assert_eq!( + tracker.status("m-1").status, + crate::health::RuntimeStatus::Healthy + ); + } + #[test] fn honor_retry_after_clamps_to_max_seconds() { let cfg = CooldownConfig { diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index 8656eb44..4027bcbc 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -225,7 +225,14 @@ async fn dispatch( let upstream_resp = req_builder .send() .await - .map_err(|e| aisix_gateway::BridgeError::Transport(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::Transport(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let status = upstream_resp.status(); @@ -258,8 +265,14 @@ async fn dispatch( return Err(ProxyError::Bridge(err)); } - // Update health tracker on success. + // Update health trackers on success — both the display-name-keyed + // observational signal AND the id-keyed runtime status that + // routing filters consult. Without `mark_healthy` here, a target + // that recovered via the Anthropic passthrough would stay in + // `cooldown` on /admin/v1/models/status until its TTL naturally + // expired (round-2 audit MEDIUM on PR #268). state.health.record_success(&model_name); + state.runtime_status.mark_healthy(&model_entry.id); let provider_label = "anthropic".to_string(); @@ -305,11 +318,21 @@ async fn dispatch( metrics: AnthropicUsageMetrics::default(), }) } else { - // Non-streaming: deserialise and re-serialise as JSON. + // Non-streaming: deserialise and re-serialise as JSON. Decode + // failures cool down the target — a body the bridge can't + // parse is a real upstream problem worth taking out of + // rotation, not a caller bug. let json_body: Value = upstream_resp .json() .await - .map_err(|e| aisix_gateway::BridgeError::UpstreamDecode(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::UpstreamDecode(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let metrics = anthropic_metrics_from_response_json(&json_body); diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index bf71ab70..8576ee49 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -180,7 +180,14 @@ async fn dispatch( .json(body) .send() .await - .map_err(|e| aisix_gateway::BridgeError::Transport(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::Transport(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let status = upstream_resp.status(); @@ -210,7 +217,14 @@ async fn dispatch( let body_bytes = upstream_resp .bytes() .await - .map_err(|e| aisix_gateway::BridgeError::UpstreamDecode(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::UpstreamDecode(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let mut resp = axum::response::Response::new(axum::body::Body::from(body_bytes)); diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index e01c5c51..7df76411 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -152,7 +152,14 @@ async fn dispatch( .json(body) .send() .await - .map_err(|e| aisix_gateway::BridgeError::Transport(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::Transport(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; let status = upstream_resp.status(); @@ -206,7 +213,14 @@ async fn dispatch( let json_body: Value = upstream_resp .json() .await - .map_err(|e| aisix_gateway::BridgeError::UpstreamDecode(e.to_string())) + .map_err(|e| { + crate::cooldown::note_failure( + &state.runtime_status, + &model_entry.id, + model.cooldown.as_ref(), + aisix_gateway::BridgeError::UpstreamDecode(e.to_string()), + ) + }) .map_err(ProxyError::Bridge)?; Ok((Json(json_body).into_response(), provider_label))