mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix: restore provider key circuit breaker backoff
This commit is contained in:
@@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRe
|
||||
use crate::handlers::admin::provider::write::normalize::{
|
||||
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
|
||||
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
|
||||
validate_vertex_api_formats,
|
||||
normalize_max_probe_interval_minutes, validate_vertex_api_formats,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::{
|
||||
@@ -181,7 +181,8 @@ pub(crate) async fn build_admin_create_provider_key_record(
|
||||
key.rpm_limit = payload.rpm_limit;
|
||||
key.concurrent_limit = normalize_optional_api_key_concurrent_limit(payload.concurrent_limit)?;
|
||||
key.cache_ttl_minutes = payload.cache_ttl_minutes.unwrap_or(5);
|
||||
key.max_probe_interval_minutes = payload.max_probe_interval_minutes.unwrap_or(32);
|
||||
key.max_probe_interval_minutes =
|
||||
normalize_max_probe_interval_minutes(payload.max_probe_interval_minutes.unwrap_or(32))?;
|
||||
key.request_count = Some(0);
|
||||
key.success_count = Some(0);
|
||||
key.error_count = Some(0);
|
||||
|
||||
@@ -41,8 +41,9 @@ async fn build_admin_provider_key_items_payload(
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let items = key_page
|
||||
.items
|
||||
let keys = key_page.items;
|
||||
|
||||
let items = keys
|
||||
.into_iter()
|
||||
.map(|key| {
|
||||
let api_formats =
|
||||
|
||||
@@ -2,7 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePa
|
||||
use crate::handlers::admin::provider::write::normalize::{
|
||||
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
|
||||
normalize_api_format_list, normalize_auth_type, normalize_auth_type_by_format,
|
||||
validate_vertex_api_formats,
|
||||
normalize_max_probe_interval_minutes, validate_vertex_api_formats,
|
||||
};
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::handlers::admin::shared::{
|
||||
@@ -293,7 +293,8 @@ pub(crate) async fn build_admin_update_provider_key_record(
|
||||
updated.cache_ttl_minutes = cache_ttl_minutes;
|
||||
}
|
||||
if let Some(max_probe_interval_minutes) = payload.max_probe_interval_minutes {
|
||||
updated.max_probe_interval_minutes = max_probe_interval_minutes;
|
||||
updated.max_probe_interval_minutes =
|
||||
normalize_max_probe_interval_minutes(max_probe_interval_minutes)?;
|
||||
}
|
||||
if let Some(is_active) = payload.is_active {
|
||||
updated.is_active = is_active;
|
||||
|
||||
@@ -115,6 +115,14 @@ pub(crate) fn normalize_auth_type(value: Option<&str>) -> Result<String, String>
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_max_probe_interval_minutes(value: i32) -> Result<i32, String> {
|
||||
if (0..=32).contains(&value) {
|
||||
Ok(value)
|
||||
} else {
|
||||
Err("max_probe_interval_minutes 必须在 0 到 32 之间".to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_pool_advanced_config(
|
||||
value: Option<serde_json::Value>,
|
||||
) -> Result<Option<serde_json::Value>, String> {
|
||||
|
||||
@@ -19,7 +19,8 @@ use tracing::warn;
|
||||
use super::{
|
||||
local_failover_error_message, project_local_adaptive_rate_limit,
|
||||
project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed,
|
||||
project_local_key_circuit_open, project_local_success_health, LocalFailoverClassification,
|
||||
project_local_key_circuit_failure, project_local_key_circuit_open,
|
||||
project_local_success_health, LocalFailoverClassification,
|
||||
};
|
||||
use crate::ai_serving::extract_pool_sticky_session_token;
|
||||
use crate::client_session_affinity::{
|
||||
@@ -521,21 +522,38 @@ async fn record_health_failure_effect(
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let observed_at_unix_secs = current_unix_secs();
|
||||
let Some(health_by_format) = project_local_failure_health(
|
||||
current_key.health_by_format.as_ref(),
|
||||
api_format,
|
||||
effect.classification,
|
||||
effect.status_code,
|
||||
current_unix_secs(),
|
||||
observed_at_unix_secs,
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
let consecutive_failures = health_by_format
|
||||
.get(api_format)
|
||||
.and_then(|value| value.get("consecutive_failures"))
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0);
|
||||
let circuit_breaker_by_format = project_local_key_circuit_failure(
|
||||
current_key.circuit_breaker_by_format.as_ref(),
|
||||
api_format,
|
||||
observed_at_unix_secs,
|
||||
consecutive_failures,
|
||||
current_key.max_probe_interval_minutes,
|
||||
);
|
||||
let circuit_breaker_update = circuit_breaker_by_format
|
||||
.as_ref()
|
||||
.or(current_key.circuit_breaker_by_format.as_ref());
|
||||
|
||||
if let Err(err) = state
|
||||
.update_provider_catalog_key_format_health(
|
||||
.update_provider_catalog_key_health_state(
|
||||
&context.plan.key_id,
|
||||
api_format,
|
||||
&health_by_format,
|
||||
current_key.is_active,
|
||||
Some(&health_by_format),
|
||||
circuit_breaker_update,
|
||||
)
|
||||
.await
|
||||
{
|
||||
@@ -705,6 +723,7 @@ async fn open_pool_key_circuit_breaker(
|
||||
api_format,
|
||||
reason,
|
||||
current_unix_secs(),
|
||||
current_key.max_probe_interval_minutes,
|
||||
) else {
|
||||
return;
|
||||
};
|
||||
@@ -1831,6 +1850,51 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_failure_opens_circuit_after_eight_consecutive_failures() {
|
||||
let state = health_state();
|
||||
let plan = sample_plan();
|
||||
|
||||
for _ in 0..8 {
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: None,
|
||||
},
|
||||
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||
status_code: 503,
|
||||
classification: LocalFailoverClassification::RetryUpstreamFailure,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let stored_key = state
|
||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
|
||||
.await
|
||||
.expect("provider catalog keys should load")
|
||||
.into_iter()
|
||||
.next()
|
||||
.expect("stored key should exist");
|
||||
let circuit = stored_key
|
||||
.circuit_breaker_by_format
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("openai:chat"))
|
||||
.expect("format circuit should be stored");
|
||||
assert_eq!(circuit["open"], json!(true));
|
||||
assert_eq!(circuit["reason"], json!("consecutive_failures_8"));
|
||||
assert_eq!(circuit["probe_interval_minutes"], json!(1));
|
||||
assert!(circuit["next_probe_at_unix_secs"].as_u64().is_some());
|
||||
assert_eq!(
|
||||
circuit["request_results_window"]
|
||||
.as_array()
|
||||
.map(Vec::len)
|
||||
.unwrap_or_default(),
|
||||
8
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_success_projection_resets_key_health_for_format() {
|
||||
let state = health_state();
|
||||
|
||||
@@ -4,7 +4,8 @@ use super::LocalFailoverClassification;
|
||||
use crate::handlers::shared::unix_secs_to_rfc3339;
|
||||
|
||||
const LOCAL_HEALTH_SCORE_FLOOR: f64 = 0.2;
|
||||
const LOCAL_KEY_CIRCUIT_PROBE_DELAY_SECS: u64 = 32 * 60;
|
||||
pub(crate) const LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD: u64 = 8;
|
||||
const LOCAL_KEY_CIRCUIT_MAX_PROBE_INTERVAL_MINUTES: u64 = 32;
|
||||
|
||||
pub(crate) fn project_local_failure_health(
|
||||
current_health_by_format: Option<&Value>,
|
||||
@@ -79,6 +80,7 @@ pub(crate) fn project_local_key_circuit_open(
|
||||
api_format: &str,
|
||||
reason: &str,
|
||||
observed_at_unix_secs: u64,
|
||||
max_probe_interval_minutes: i32,
|
||||
) -> Option<Value> {
|
||||
let api_format = api_format.trim();
|
||||
let reason = reason.trim();
|
||||
@@ -86,23 +88,143 @@ pub(crate) fn project_local_key_circuit_open(
|
||||
return None;
|
||||
}
|
||||
|
||||
let next_probe_at_unix_secs =
|
||||
observed_at_unix_secs.saturating_add(LOCAL_KEY_CIRCUIT_PROBE_DELAY_SECS);
|
||||
let mut circuit_by_format = current_circuit_by_format
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let current = circuit_by_format
|
||||
.get(api_format)
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let max_probe_interval_minutes =
|
||||
normalize_max_probe_interval_minutes(max_probe_interval_minutes);
|
||||
let probe_interval_minutes =
|
||||
next_circuit_probe_interval_minutes(¤t, max_probe_interval_minutes);
|
||||
let next_probe_at_unix_secs =
|
||||
next_probe_at_unix_secs(observed_at_unix_secs, probe_interval_minutes);
|
||||
let open_at = current
|
||||
.get("open_at")
|
||||
.filter(|_| current_bool(¤t, "open"))
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!(unix_secs_to_rfc3339(observed_at_unix_secs)));
|
||||
let half_open_failures = current
|
||||
.get("half_open_failures")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0)
|
||||
.saturating_add(u64::from(current_bool(¤t, "open")));
|
||||
let request_results_window =
|
||||
append_request_result_window(¤t, observed_at_unix_secs, false);
|
||||
circuit_by_format.insert(
|
||||
api_format.to_string(),
|
||||
json!({
|
||||
"open": true,
|
||||
"open_at": unix_secs_to_rfc3339(observed_at_unix_secs),
|
||||
"open_at": open_at,
|
||||
"reason": reason,
|
||||
"next_probe_at": unix_secs_to_rfc3339(next_probe_at_unix_secs),
|
||||
"next_probe_at_unix_secs": next_probe_at_unix_secs,
|
||||
"probe_interval_minutes": probe_interval_minutes,
|
||||
"max_probe_interval_minutes": max_probe_interval_minutes,
|
||||
"last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
|
||||
"last_probe_failure_at": if half_open_failures > 0 {
|
||||
json!(unix_secs_to_rfc3339(observed_at_unix_secs))
|
||||
} else {
|
||||
Value::Null
|
||||
},
|
||||
"half_open_until": Value::Null,
|
||||
"half_open_successes": 0,
|
||||
"half_open_failures": 0,
|
||||
"half_open_failures": half_open_failures,
|
||||
"request_results_window": request_results_window,
|
||||
}),
|
||||
);
|
||||
Some(Value::Object(circuit_by_format))
|
||||
}
|
||||
|
||||
pub(crate) fn project_local_key_circuit_failure(
|
||||
current_circuit_by_format: Option<&Value>,
|
||||
api_format: &str,
|
||||
observed_at_unix_secs: u64,
|
||||
consecutive_failures: u64,
|
||||
max_probe_interval_minutes: i32,
|
||||
) -> Option<Value> {
|
||||
let api_format = api_format.trim();
|
||||
if api_format.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut circuit_by_format = current_circuit_by_format
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let current = circuit_by_format
|
||||
.get(api_format)
|
||||
.and_then(Value::as_object)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let request_results_window =
|
||||
append_request_result_window(¤t, observed_at_unix_secs, false);
|
||||
let already_open = current_bool(¤t, "open");
|
||||
if !already_open && consecutive_failures < LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD {
|
||||
circuit_by_format.insert(
|
||||
api_format.to_string(),
|
||||
json!({
|
||||
"open": false,
|
||||
"open_at": Value::Null,
|
||||
"reason": Value::Null,
|
||||
"next_probe_at": Value::Null,
|
||||
"next_probe_at_unix_secs": Value::Null,
|
||||
"probe_interval_minutes": 0,
|
||||
"max_probe_interval_minutes": normalize_max_probe_interval_minutes(max_probe_interval_minutes),
|
||||
"failure_count": consecutive_failures,
|
||||
"last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
|
||||
"last_probe_failure_at": Value::Null,
|
||||
"half_open_until": Value::Null,
|
||||
"half_open_successes": 0,
|
||||
"half_open_failures": 0,
|
||||
"request_results_window": request_results_window,
|
||||
}),
|
||||
);
|
||||
return Some(Value::Object(circuit_by_format));
|
||||
}
|
||||
|
||||
let max_probe_interval_minutes =
|
||||
normalize_max_probe_interval_minutes(max_probe_interval_minutes);
|
||||
let probe_interval_minutes =
|
||||
next_circuit_probe_interval_minutes(¤t, max_probe_interval_minutes);
|
||||
let next_probe_at_unix_secs =
|
||||
next_probe_at_unix_secs(observed_at_unix_secs, probe_interval_minutes);
|
||||
let open_at = current
|
||||
.get("open_at")
|
||||
.filter(|_| already_open)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!(unix_secs_to_rfc3339(observed_at_unix_secs)));
|
||||
let half_open_failures = current
|
||||
.get("half_open_failures")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(0)
|
||||
.saturating_add(u64::from(already_open));
|
||||
|
||||
circuit_by_format.insert(
|
||||
api_format.to_string(),
|
||||
json!({
|
||||
"open": true,
|
||||
"open_at": open_at,
|
||||
"reason": format!("consecutive_failures_{LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD}"),
|
||||
"next_probe_at": unix_secs_to_rfc3339(next_probe_at_unix_secs),
|
||||
"next_probe_at_unix_secs": next_probe_at_unix_secs,
|
||||
"probe_interval_minutes": probe_interval_minutes,
|
||||
"max_probe_interval_minutes": max_probe_interval_minutes,
|
||||
"failure_count": consecutive_failures,
|
||||
"last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
|
||||
"last_probe_failure_at": if already_open {
|
||||
json!(unix_secs_to_rfc3339(observed_at_unix_secs))
|
||||
} else {
|
||||
Value::Null
|
||||
},
|
||||
"half_open_until": Value::Null,
|
||||
"half_open_successes": 0,
|
||||
"half_open_failures": half_open_failures,
|
||||
"request_results_window": request_results_window,
|
||||
}),
|
||||
);
|
||||
Some(Value::Object(circuit_by_format))
|
||||
@@ -137,6 +259,60 @@ pub(crate) fn project_local_key_circuit_closed(
|
||||
Some(Value::Object(circuit_by_format))
|
||||
}
|
||||
|
||||
fn current_bool(current: &serde_json::Map<String, Value>, field: &str) -> bool {
|
||||
current.get(field).and_then(Value::as_bool).unwrap_or(false)
|
||||
}
|
||||
|
||||
fn normalize_max_probe_interval_minutes(value: i32) -> u64 {
|
||||
value.clamp(0, LOCAL_KEY_CIRCUIT_MAX_PROBE_INTERVAL_MINUTES as i32) as u64
|
||||
}
|
||||
|
||||
fn next_circuit_probe_interval_minutes(
|
||||
current: &serde_json::Map<String, Value>,
|
||||
max_probe_interval_minutes: u64,
|
||||
) -> u64 {
|
||||
if max_probe_interval_minutes == 0 {
|
||||
return 0;
|
||||
}
|
||||
if !current_bool(current, "open") {
|
||||
return 1.min(max_probe_interval_minutes);
|
||||
}
|
||||
current
|
||||
.get("probe_interval_minutes")
|
||||
.and_then(Value::as_u64)
|
||||
.unwrap_or(1)
|
||||
.max(1)
|
||||
.saturating_mul(2)
|
||||
.min(max_probe_interval_minutes)
|
||||
}
|
||||
|
||||
fn next_probe_at_unix_secs(observed_at_unix_secs: u64, interval_minutes: u64) -> u64 {
|
||||
observed_at_unix_secs.saturating_add(interval_minutes.saturating_mul(60))
|
||||
}
|
||||
|
||||
fn append_request_result_window(
|
||||
current: &serde_json::Map<String, Value>,
|
||||
observed_at_unix_secs: u64,
|
||||
ok: bool,
|
||||
) -> Value {
|
||||
let mut window = current
|
||||
.get("request_results_window")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
window.push(json!({
|
||||
"ts": observed_at_unix_secs,
|
||||
"ok": ok,
|
||||
}));
|
||||
let keep = usize::try_from(LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD)
|
||||
.unwrap_or(8)
|
||||
.max(1);
|
||||
if window.len() > keep {
|
||||
window = window.split_off(window.len() - keep);
|
||||
}
|
||||
Value::Array(window)
|
||||
}
|
||||
|
||||
fn local_candidate_failure_should_project_health(
|
||||
classification: LocalFailoverClassification,
|
||||
status_code: u16,
|
||||
@@ -182,7 +358,8 @@ mod tests {
|
||||
|
||||
use super::{
|
||||
project_local_failure_health, project_local_key_circuit_closed,
|
||||
project_local_key_circuit_open, project_local_success_health,
|
||||
project_local_key_circuit_failure, project_local_key_circuit_open,
|
||||
project_local_success_health,
|
||||
};
|
||||
use crate::orchestration::LocalFailoverClassification;
|
||||
|
||||
@@ -269,6 +446,7 @@ mod tests {
|
||||
"openai:chat",
|
||||
"account_deactivated_401",
|
||||
1_760_000_000,
|
||||
32,
|
||||
)
|
||||
.expect("projection should exist");
|
||||
|
||||
@@ -279,7 +457,47 @@ mod tests {
|
||||
);
|
||||
assert_eq!(
|
||||
projected["openai:chat"]["next_probe_at_unix_secs"],
|
||||
json!(1_760_001_920u64)
|
||||
json!(1_760_000_060u64)
|
||||
);
|
||||
assert_eq!(projected["openai:chat"]["probe_interval_minutes"], json!(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn consecutive_failure_circuit_opens_after_threshold_and_backs_off() {
|
||||
let before_threshold =
|
||||
project_local_key_circuit_failure(None, "openai:chat", 1_760_000_000, 7, 32)
|
||||
.expect("projection should exist");
|
||||
assert_eq!(before_threshold["openai:chat"]["open"], json!(false));
|
||||
|
||||
let opened = project_local_key_circuit_failure(
|
||||
Some(&before_threshold),
|
||||
"openai:chat",
|
||||
1_760_000_060,
|
||||
8,
|
||||
32,
|
||||
)
|
||||
.expect("projection should exist");
|
||||
assert_eq!(opened["openai:chat"]["open"], json!(true));
|
||||
assert_eq!(
|
||||
opened["openai:chat"]["reason"],
|
||||
json!("consecutive_failures_8")
|
||||
);
|
||||
assert_eq!(opened["openai:chat"]["probe_interval_minutes"], json!(1));
|
||||
assert_eq!(
|
||||
opened["openai:chat"]["next_probe_at_unix_secs"],
|
||||
json!(1_760_000_120u64)
|
||||
);
|
||||
|
||||
let backed_off =
|
||||
project_local_key_circuit_failure(Some(&opened), "openai:chat", 1_760_000_120, 9, 32)
|
||||
.expect("projection should exist");
|
||||
assert_eq!(
|
||||
backed_off["openai:chat"]["probe_interval_minutes"],
|
||||
json!(2)
|
||||
);
|
||||
assert_eq!(
|
||||
backed_off["openai:chat"]["next_probe_at_unix_secs"],
|
||||
json!(1_760_000_240u64)
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -34,7 +34,8 @@ pub(crate) use self::effects::{
|
||||
LocalPoolErrorEffect,
|
||||
};
|
||||
pub(crate) use self::health::{
|
||||
project_local_failure_health, project_local_key_circuit_closed, project_local_key_circuit_open,
|
||||
project_local_failure_health, project_local_key_circuit_closed,
|
||||
project_local_key_circuit_failure, project_local_key_circuit_open,
|
||||
project_local_success_health,
|
||||
};
|
||||
pub(crate) use self::policy::{
|
||||
|
||||
@@ -406,7 +406,6 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
|
||||
"fn candidate_provider_key_health_bucket(",
|
||||
"fn candidate_provider_key_health_score(",
|
||||
"count_recent_active_requests_for_provider(",
|
||||
"is_candidate_in_recent_failure_cooldown(",
|
||||
"provider_key_health_score(",
|
||||
"provider_key_rpm_allows_request_since(",
|
||||
"read_recent_request_candidates(128)",
|
||||
|
||||
@@ -118,6 +118,107 @@ async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principa
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_provider_keys_expose_circuit_breaker_and_recover_clears_it() {
|
||||
let key = sample_key("key-1", "provider-1", "openai:chat", "sk-test-a").with_health_fields(
|
||||
Some(json!({"openai:chat": {
|
||||
"health_score": 0.2,
|
||||
"consecutive_failures": 8,
|
||||
"last_failure_at": "2026-03-26T12:00:00+00:00"
|
||||
}})),
|
||||
Some(json!({"openai:chat": {
|
||||
"open": true,
|
||||
"open_at": "2026-03-26T12:00:00+00:00",
|
||||
"reason": "consecutive_failures_8",
|
||||
"next_probe_at": "2026-03-26T12:01:00+00:00",
|
||||
"next_probe_at_unix_secs": 1774526460u64,
|
||||
"probe_interval_minutes": 1,
|
||||
"max_probe_interval_minutes": 32,
|
||||
"half_open_until": null,
|
||||
"half_open_successes": 0,
|
||||
"half_open_failures": 0
|
||||
}})),
|
||||
);
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"https://example.com/v1",
|
||||
)],
|
||||
vec![key],
|
||||
));
|
||||
let gateway_state = AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
|
||||
&provider_catalog_repository,
|
||||
)),
|
||||
);
|
||||
let gateway = build_router_with_state(gateway_state);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-1/keys?skip=0&limit=50"
|
||||
))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload[0]["circuit_breaker_open"], true);
|
||||
assert_eq!(
|
||||
payload[0]["circuit_breaker_by_format"]["openai:chat"]["reason"],
|
||||
"consecutive_failures_8"
|
||||
);
|
||||
assert_eq!(
|
||||
payload[0]["circuit_breaker_by_format"]["openai:chat"]["probe_interval_minutes"],
|
||||
1
|
||||
);
|
||||
assert!(
|
||||
payload[0]["circuit_breaker_by_format"]["openai:chat"]["next_probe_at_unix_secs"]
|
||||
.as_u64()
|
||||
.is_some()
|
||||
);
|
||||
|
||||
let recover_response = client
|
||||
.patch(format!(
|
||||
"{gateway_url}/api/admin/endpoints/health/keys/key-1"
|
||||
))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("recover request should succeed");
|
||||
assert_eq!(recover_response.status(), StatusCode::OK);
|
||||
|
||||
let response = client
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-1/keys?skip=0&limit=50"
|
||||
))
|
||||
.header(GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload[0]["circuit_breaker_open"], false);
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_keys_page_locally_with_total() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
Reference in New Issue
Block a user