feat: update gateway pool and usage flows

This commit is contained in:
fawney19
2026-05-12 21:05:11 +08:00
parent 38012c62ff
commit d1a47c068e
18 changed files with 582 additions and 328 deletions

View File

@@ -3,6 +3,7 @@ pub(crate) use crate::handlers::admin::{
build_internal_control_error_response, create_provider_oauth_catalog_key, build_internal_control_error_response, create_provider_oauth_catalog_key,
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response, find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
maybe_build_local_admin_response, provider_oauth_runtime_endpoint_for_provider, maybe_build_local_admin_response, provider_oauth_runtime_endpoint_for_provider,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally, refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally,
refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally, refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally,
refresh_provider_oauth_account_state_after_update, update_existing_provider_oauth_catalog_key, refresh_provider_oauth_account_state_after_update, update_existing_provider_oauth_catalog_key,

View File

@@ -32,7 +32,8 @@ use crate::handlers::shared::provider_pool::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value, admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
}; };
use crate::handlers::shared::provider_pool::{ use crate::handlers::shared::provider_pool::{
try_claim_admin_provider_pool_key, AdminProviderPoolConfig, AdminProviderPoolRuntimeState, read_admin_provider_pool_key_cooldown_reason, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState,
}; };
use crate::handlers::shared::{ use crate::handlers::shared::{
parse_catalog_auth_config_json, provider_key_health_summary, parse_catalog_auth_config_json, provider_key_health_summary,
@@ -537,7 +538,7 @@ impl<'a> PoolKeyCursor<'a> {
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> { async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
while let Some(candidate) = self.queued_candidates.pop_front() { while let Some(candidate) = self.queued_candidates.pop_front() {
let mut candidate = candidate; let mut candidate = candidate;
if !self.attach_pool_key_lease(&mut candidate).await { if self.skip_candidate_if_runtime_cooldown(&candidate).await {
continue; continue;
} }
candidate.orchestration.pool_key_index = Some(self.next_pool_key_index); candidate.orchestration.pool_key_index = Some(self.next_pool_key_index);
@@ -548,43 +549,42 @@ impl<'a> PoolKeyCursor<'a> {
None None
} }
async fn attach_pool_key_lease( async fn skip_candidate_if_runtime_cooldown(
&mut self, &mut self,
candidate: &mut EligibleLocalExecutionCandidate, candidate: &EligibleLocalExecutionCandidate,
) -> bool { ) -> bool {
if !pool_key_order_requires_exclusive_lease(&self.pool_key_order) { match read_admin_provider_pool_key_cooldown_reason(
return true;
}
let owner = self.state.app().tunnel.local_instance_id();
match try_claim_admin_provider_pool_key(
self.state.app().runtime_state.as_ref(), self.state.app().runtime_state.as_ref(),
candidate.candidate.provider_id.as_str(), candidate.candidate.provider_id.as_str(),
candidate.candidate.key_id.as_str(), candidate.candidate.key_id.as_str(),
owner,
) )
.await .await
{ {
Ok(Some(lease)) => { Ok(Some(_)) => {
candidate.orchestration.pool_key_lease = Some(lease); self.record_skip_reason("pool_cooldown");
self.skipped_candidates
.push(SkippedLocalExecutionCandidate {
candidate: candidate.candidate.clone(),
skip_reason: "pool_cooldown",
transport: Some(candidate.transport.clone()),
ranking: candidate.ranking.clone(),
extra_data: None,
});
true true
} }
Ok(None) => { Ok(None) => false,
self.record_skip_reason("pool_key_lease_busy");
false
}
Err(err) => { Err(err) => {
warn!( warn!(
event_name = "pool_key_lease_claim_failed", event_name = "pool_key_cooldown_check_failed",
log_type = "event", log_type = "event",
provider_id = %candidate.candidate.provider_id, provider_id = %candidate.candidate.provider_id,
endpoint_id = %candidate.candidate.endpoint_id, endpoint_id = %candidate.candidate.endpoint_id,
model_id = %candidate.candidate.model_id, model_id = %candidate.candidate.model_id,
key_id = %candidate.candidate.key_id, key_id = %candidate.candidate.key_id,
error = ?err, error = ?err,
"gateway pool scheduler failed to claim pool key lease; scheduling key without lease" "gateway pool scheduler failed to read pool key cooldown; scheduling key"
); );
true false
} }
} }
} }
@@ -707,13 +707,6 @@ impl<'a> PoolKeyCursor<'a> {
} }
} }
fn pool_key_order_requires_exclusive_lease(order: &StoredPoolKeyCandidateOrder) -> bool {
!matches!(
order,
StoredPoolKeyCandidateOrder::CacheAffinity | StoredPoolKeyCandidateOrder::SingleAccount
)
}
fn pool_candidate_transport_policy_facts( fn pool_candidate_transport_policy_facts(
candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate, candidate: &aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate,
) -> CandidateTransportPolicyFacts<'_> { ) -> CandidateTransportPolicyFacts<'_> {
@@ -1137,8 +1130,7 @@ fn apply_ai_pool_orchestration(
mod tests { mod tests {
use super::{ use super::{
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context, apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
pool_config_for_candidate, pool_key_order_requires_exclusive_lease, PoolCatalogKeyContext, pool_config_for_candidate, PoolCatalogKeyContext, PoolKeyCursor,
PoolKeyCursor,
}; };
use crate::ai_serving::{ use crate::ai_serving::{
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate, apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
@@ -1146,8 +1138,7 @@ mod tests {
}; };
use crate::data::GatewayDataState; use crate::data::GatewayDataState;
use crate::handlers::shared::provider_pool::{ use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, release_admin_provider_pool_key_lease, record_admin_provider_pool_error, AdminProviderPoolRuntimeState,
try_claim_admin_provider_pool_key, AdminProviderPoolRuntimeState,
}; };
use crate::orchestration::LocalExecutionCandidateMetadata; use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::{AppState, LocalExecutionRuntimeMissDiagnostic}; use crate::{AppState, LocalExecutionRuntimeMissDiagnostic};
@@ -1901,24 +1892,6 @@ mod tests {
assert_eq!(lru_cursor.pool_key_order, StoredPoolKeyCandidateOrder::Lru); assert_eq!(lru_cursor.pool_key_order, StoredPoolKeyCandidateOrder::Lru);
} }
#[test]
fn pool_key_lease_exclusivity_follows_distribution_mode() {
assert!(!pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::CacheAffinity
));
assert!(!pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::SingleAccount
));
assert!(pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::Lru
));
assert!(pool_key_order_requires_exclusive_lease(
&StoredPoolKeyCandidateOrder::LoadBalance {
seed: "seed".to_string()
}
));
}
#[test] #[test]
fn pool_key_cursor_records_runtime_miss_when_exhausted_without_returning_key() { fn pool_key_cursor_records_runtime_miss_when_exhausted_without_returning_key() {
let app = AppState::new().expect("state should build"); let app = AppState::new().expect("state should build");
@@ -1941,8 +1914,8 @@ mod tests {
); );
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None) let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None)
.with_runtime_miss_diagnostic(trace_id, true); .with_runtime_miss_diagnostic(trace_id, true);
cursor.record_skip_reason("pool_key_lease_busy"); cursor.record_skip_reason("pool_cooldown");
cursor.record_skip_reason("pool_key_lease_busy"); cursor.record_skip_reason("pool_cooldown");
cursor.record_skip_reason("transport_snapshot_missing"); cursor.record_skip_reason("transport_snapshot_missing");
cursor.log_exhausted(); cursor.log_exhausted();
@@ -1953,11 +1926,11 @@ mod tests {
.expect("runtime miss diagnostic should exist"); .expect("runtime miss diagnostic should exist");
assert_eq!(diagnostic.reason, "all_candidates_skipped"); assert_eq!(diagnostic.reason, "all_candidates_skipped");
assert_eq!(diagnostic.skipped_candidate_count, Some(1)); assert_eq!(diagnostic.skipped_candidate_count, Some(1));
assert_eq!(diagnostic.skip_reasons.get("pool_key_lease_busy"), Some(&1)); assert_eq!(diagnostic.skip_reasons.get("pool_cooldown"), Some(&1));
} }
#[tokio::test] #[tokio::test]
async fn pool_key_cursor_freezes_queued_candidates_after_window_refill() { async fn pool_key_cursor_rechecks_cooldown_for_frozen_window_candidates() {
let app = AppState::new().expect("state should build"); let app = AppState::new().expect("state should build");
let provider_config = Some(json!({ let provider_config = Some(json!({
"pool_advanced": { "pool_advanced": {
@@ -2009,33 +1982,24 @@ mod tests {
let second = cursor let second = cursor
.next_key() .next_key()
.await .await
.expect("second key should keep the frozen queue order"); .expect("second key should skip the cooled-down frozen key");
assert_eq!(second.candidate.key_id, "key-b"); assert_eq!(second.candidate.key_id, "key-c");
assert_eq!(second.orchestration.pool_key_index, Some(1)); assert_eq!(second.orchestration.pool_key_index, Some(1));
let third = cursor
.next_key()
.await
.expect("third key should keep the frozen queue order");
assert_eq!(third.candidate.key_id, "key-c");
assert_eq!(third.orchestration.pool_key_index, Some(2));
let skipped = cursor.take_skipped_candidates(); let skipped = cursor.take_skipped_candidates();
assert!(skipped.is_empty()); assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-b", "pool_cooldown")]
);
} }
#[tokio::test] #[tokio::test]
async fn pool_key_cursor_skips_busy_lease_and_claims_returned_key() { async fn pool_key_cursor_allows_parallel_requests_to_use_same_healthy_key() {
let app = AppState::new().expect("state should build"); let app = AppState::new().expect("state should build");
let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } })); let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } }));
let held_lease = try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
"key-a",
"test-owner",
)
.await
.expect("lease claim should not fail")
.expect("first key should be claimed by test");
let group = sample_eligible_candidate( let group = sample_eligible_candidate(
"provider-pool", "provider-pool",
"endpoint-1", "endpoint-1",
@@ -2052,87 +2016,19 @@ mod tests {
10, 10,
provider_config.clone(), provider_config.clone(),
), ),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config), sample_eligible_candidate(
]);
let candidate = cursor
.next_key()
.await
.expect("cursor should skip busy key and return next key");
assert_eq!(candidate.candidate.key_id, "key-b");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert_eq!(
cursor.skip_reason_counts.get("pool_key_lease_busy"),
Some(&1)
);
let returned_lease = candidate
.orchestration
.pool_key_lease
.clone()
.expect("returned key should carry its lease");
assert!(
try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool", "provider-pool",
"endpoint-1",
"key-b", "key-b",
"other-owner",
)
.await
.expect("second claim should not fail")
.is_none(),
"returned key should remain leased until execution releases it"
);
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &held_lease)
.await
.expect("held lease release should not fail");
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &returned_lease)
.await
.expect("returned lease release should not fail");
}
#[tokio::test]
async fn pool_key_cursor_allows_busy_lease_for_cache_affinity_mode() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({
"pool_advanced": {
"scheduling_presets": [
{"preset": "cache_affinity", "enabled": true}
]
}
}));
let held_lease = try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
"key-a",
"test-owner",
)
.await
.expect("lease claim should not fail")
.expect("first key should be claimed by test");
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config.clone(),
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10, 10,
provider_config.clone(), provider_config.clone(),
), ),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
]); ]);
let candidate = cursor let candidate = cursor
.next_key() .next_key()
.await .await
.expect("cache-affinity cursor should allow repeated busy key"); .expect("cursor should return the first healthy key");
assert_eq!(candidate.candidate.key_id, "key-a"); assert_eq!(candidate.candidate.key_id, "key-a");
assert_eq!(candidate.orchestration.pool_key_index, Some(0)); assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert!(candidate.orchestration.pool_key_lease.is_none()); assert!(candidate.orchestration.pool_key_lease.is_none());
@@ -2140,9 +2036,82 @@ mod tests {
.skip_reason_counts .skip_reason_counts
.contains_key("pool_key_lease_busy")); .contains_key("pool_key_lease_busy"));
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &held_lease) let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config,
);
let mut second_cursor =
PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
second_cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
),
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-b",
10,
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
),
]);
let second_candidate = second_cursor
.next_key()
.await .await
.expect("held lease release should not fail"); .expect("second request should also be allowed to pick the same healthy key");
assert_eq!(second_candidate.candidate.key_id, "key-a");
assert!(second_candidate.orchestration.pool_key_lease.is_none());
}
#[tokio::test]
async fn pool_key_cursor_skips_key_after_account_cooldown_is_recorded() {
let app = AppState::new().expect("state should build");
let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } }));
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config.clone(),
);
let pool_config = pool_config_for_candidate(&group).expect("pool config should parse");
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
cursor.queued_candidates = VecDeque::from([
sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
provider_config.clone(),
),
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
]);
record_admin_provider_pool_error(
app.runtime_state.as_ref(),
"provider-pool",
"key-a",
&pool_config,
429,
None,
None,
)
.await;
let candidate = cursor
.next_key()
.await
.expect("cursor should skip cooled-down key and return next key");
assert_eq!(candidate.candidate.key_id, "key-b");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert!(candidate.orchestration.pool_key_lease.is_none());
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1));
} }
#[tokio::test] #[tokio::test]
@@ -2177,21 +2146,7 @@ mod tests {
provider_config.clone(), provider_config.clone(),
); );
let pool_config = pool_config_for_candidate(&group).expect("pool config should parse"); let pool_config = pool_config_for_candidate(&group).expect("pool config should parse");
let mut held_leases = Vec::new();
for key_id in ["key-00000", "key-00001"] { for key_id in ["key-00000", "key-00001"] {
held_leases.push(
try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
key_id,
"large-pool-test",
)
.await
.expect("lease claim should not fail")
.expect("test key should claim"),
);
}
for key_id in ["key-00002", "key-00003"] {
record_admin_provider_pool_error( record_admin_provider_pool_error(
app.runtime_state.as_ref(), app.runtime_state.as_ref(),
"provider-pool", "provider-pool",
@@ -2219,22 +2174,16 @@ mod tests {
); );
let mut returned_ids = Vec::new(); let mut returned_ids = Vec::new();
let mut returned_leases = Vec::new();
for _ in 0..10 { for _ in 0..10 {
let candidate = cursor let candidate = cursor
.next_key() .next_key()
.await .await
.expect("large pool should return first page candidates"); .expect("large pool should return first page candidates");
returned_ids.push(candidate.candidate.key_id.clone()); returned_ids.push(candidate.candidate.key_id.clone());
returned_leases.push( assert!(candidate.orchestration.pool_key_lease.is_none());
candidate
.orchestration
.pool_key_lease
.expect("lru pool candidate should carry exclusive lease"),
);
} }
assert_eq!(returned_ids.first().map(String::as_str), Some("key-00004")); assert_eq!(returned_ids.first().map(String::as_str), Some("key-00002"));
assert_eq!(returned_ids.last().map(String::as_str), Some("key-00013")); assert_eq!(returned_ids.last().map(String::as_str), Some("key-00011"));
assert_eq!(cursor.scanned_keys, 64); assert_eq!(cursor.scanned_keys, 64);
assert!( assert!(
cursor.queued_candidates.len() <= cursor.window_size as usize, cursor.queued_candidates.len() <= cursor.window_size as usize,
@@ -2251,39 +2200,18 @@ mod tests {
None, None,
) )
.await; .await;
held_leases.push(
try_claim_admin_provider_pool_key(
app.runtime_state.as_ref(),
"provider-pool",
"key-00015",
"large-pool-test",
)
.await
.expect("dynamic lease claim should not fail")
.expect("dynamic busy key should claim"),
);
let candidate = cursor let candidate = cursor
.next_key() .next_key()
.await .await
.expect("cursor should keep the frozen window despite later runtime changes"); .expect("cursor should keep the frozen window despite later runtime changes");
assert_eq!(candidate.candidate.key_id, "key-00014"); assert_eq!(candidate.candidate.key_id, "key-00012");
returned_ids.push(candidate.candidate.key_id.clone()); returned_ids.push(candidate.candidate.key_id.clone());
returned_leases.push( assert!(candidate.orchestration.pool_key_lease.is_none());
candidate
.orchestration
.pool_key_lease
.expect("lru pool candidate should carry exclusive lease"),
);
while let Some(candidate) = cursor.next_key().await { while let Some(candidate) = cursor.next_key().await {
returned_ids.push(candidate.candidate.key_id.clone()); returned_ids.push(candidate.candidate.key_id.clone());
returned_leases.push( assert!(candidate.orchestration.pool_key_lease.is_none());
candidate
.orchestration
.pool_key_lease
.expect("lru pool candidate should carry exclusive lease"),
);
} }
let max_returned_windows = aether_dispatch_core::DEFAULT_POOL_MAX_SCAN let max_returned_windows = aether_dispatch_core::DEFAULT_POOL_MAX_SCAN
@@ -2293,34 +2221,23 @@ mod tests {
<= (max_returned_windows * aether_dispatch_core::DEFAULT_POOL_WINDOW_SIZE) as usize, <= (max_returned_windows * aether_dispatch_core::DEFAULT_POOL_WINDOW_SIZE) as usize,
"cursor should only return bounded frozen windows per request" "cursor should only return bounded frozen windows per request"
); );
assert_eq!(returned_ids.len(), 125); assert_eq!(returned_ids.len(), 127);
assert_eq!(returned_ids.last().map(String::as_str), Some("key-00463")); assert_eq!(returned_ids.last().map(String::as_str), Some("key-00463"));
assert_eq!(cursor.scanned_keys, 512); assert_eq!(cursor.scanned_keys, 512);
assert_eq!( assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&3));
cursor.skip_reason_counts.get("pool_key_lease_busy"), assert!(!cursor
Some(&3) .skip_reason_counts
); .contains_key("pool_key_lease_busy"));
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&2)); for skipped in ["key-00000", "key-00001", "key-00014"] {
for skipped in ["key-00000", "key-00001", "key-00002", "key-00003"] {
assert!( assert!(
!returned_ids.iter().any(|key_id| key_id == skipped), !returned_ids.iter().any(|key_id| key_id == skipped),
"{skipped} should have been skipped" "{skipped} should have been skipped"
); );
} }
assert!( assert!(
returned_ids.iter().any(|key_id| key_id == "key-00014"), returned_ids.iter().any(|key_id| key_id == "key-00015"),
"key-00014 was already frozen in the current window" "key-00015 should not be blocked by request-scoped leases"
); );
assert!(
!returned_ids.iter().any(|key_id| key_id == "key-00015"),
"key-00015 should still be skipped because exclusive lease is checked at dispatch time"
);
for lease in held_leases.into_iter().chain(returned_leases) {
release_admin_provider_pool_key_lease(app.runtime_state.as_ref(), &lease)
.await
.expect("lease release should not fail");
}
} }
#[test] #[test]

View File

@@ -29,12 +29,14 @@ pub(crate) use self::provider::oauth::quota::antigravity::refresh_antigravity_pr
pub(crate) use self::provider::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally; pub(crate) use self::provider::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::codex::refresh_codex_provider_quota_locally; pub(crate) use self::provider::oauth::quota::codex::refresh_codex_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::kiro::refresh_kiro_provider_quota_locally; pub(crate) use self::provider::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh;
pub(crate) use self::provider::oauth::runtime::{ pub(crate) use self::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update, provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
}; };
pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response; pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response;
pub(crate) use self::provider::pool::config::admin_provider_pool_config; pub(crate) use self::provider::pool::config::admin_provider_pool_config;
pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response; pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response;
pub(crate) use self::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
pub(crate) use self::provider::{ pub(crate) use self::provider::{
maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response, maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response,
}; };

View File

@@ -18,6 +18,24 @@ use super::super::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_
use super::super::oauth::quota::codex::refresh_codex_provider_quota_locally; use super::super::oauth::quota::codex::refresh_codex_provider_quota_locally;
use super::super::oauth::quota::kiro::refresh_kiro_provider_quota_locally; use super::super::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
use super::super::oauth::quota::shared::normalize_string_id_list; use super::super::oauth::quota::shared::normalize_string_id_list;
use super::super::oauth::quota::shared::{
provider_type_supports_quota_refresh, unsupported_provider_quota_refresh_message,
};
use super::super::oauth::runtime::provider_oauth_runtime_endpoint_for_provider;
use super::super::write::provider::reconcile_admin_fixed_provider_template_endpoints;
fn unsupported_provider_quota_refresh_response(provider_type: &str) -> Response<Body> {
let message = unsupported_provider_quota_refresh_message(provider_type);
Json(json!({
"success": 0,
"failed": 0,
"total": 0,
"results": [],
"message": message,
"auto_removed": 0,
}))
.into_response()
}
pub(super) async fn maybe_handle( pub(super) async fn maybe_handle(
state: &AdminAppState<'_>, state: &AdminAppState<'_>,
@@ -85,41 +103,48 @@ pub(super) async fn maybe_handle(
let raw_key_ids = payload.key_ids; let raw_key_ids = payload.key_ids;
let selected_key_ids = normalize_string_id_list(raw_key_ids.clone()); let selected_key_ids = normalize_string_id_list(raw_key_ids.clone());
let explicit_key_ids_requested = raw_key_ids.is_some(); let explicit_key_ids_requested = raw_key_ids.is_some();
let endpoints = state
let is_fixed_provider = state
.fixed_provider_template(&provider.provider_type)
.is_some();
if !is_fixed_provider && !provider_type_supports_quota_refresh(&normalized_provider_type) {
return Ok(None);
}
let mut endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id)) .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?; .await?;
let endpoint = match normalized_provider_type.as_str() { let mut endpoint =
"codex" => endpoints.into_iter().find(|endpoint| { provider_oauth_runtime_endpoint_for_provider(&normalized_provider_type, &endpoints);
endpoint.is_active
&& crate::ai_serving::is_openai_responses_format(&endpoint.api_format) if endpoint.is_none() && is_fixed_provider {
}), if !state.has_provider_catalog_data_writer() {
"antigravity" => endpoints.into_iter().find(|endpoint| { if !provider_type_supports_quota_refresh(&normalized_provider_type) {
endpoint.is_active return Ok(Some(unsupported_provider_quota_refresh_response(
&& endpoint &normalized_provider_type,
.api_format )));
.trim() }
.eq_ignore_ascii_case("gemini:generate_content") return Ok(Some(
}), (
"kiro" => endpoints http::StatusCode::BAD_REQUEST,
.iter() Json(json!({ "detail": "固定 Provider 端点缺失,且 provider catalog writer 不可用,无法自动补全端点" })),
.find(|endpoint| { )
endpoint.is_active .into_response(),
&& endpoint ));
.api_format }
.trim() reconcile_admin_fixed_provider_template_endpoints(state, &provider).await?;
.eq_ignore_ascii_case("claude:messages") endpoints = state
}) .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.cloned() .await?;
.or_else(|| endpoints.into_iter().find(|endpoint| endpoint.is_active)), endpoint =
"chatgpt_web" => endpoints.into_iter().find(|endpoint| { provider_oauth_runtime_endpoint_for_provider(&normalized_provider_type, &endpoints);
endpoint.is_active }
&& endpoint
.api_format if !provider_type_supports_quota_refresh(&normalized_provider_type) {
.trim() return Ok(Some(unsupported_provider_quota_refresh_response(
.eq_ignore_ascii_case("openai:image") &normalized_provider_type,
}), )));
_ => return Ok(None), }
};
let Some(endpoint) = endpoint else { let Some(endpoint) = endpoint else {
let detail = match normalized_provider_type.as_str() { let detail = match normalized_provider_type.as_str() {
@@ -127,6 +152,8 @@ pub(super) async fn maybe_handle(
"antigravity" => "找不到有效的 gemini:generate_content 端点", "antigravity" => "找不到有效的 gemini:generate_content 端点",
"kiro" => "找不到有效的 Kiro 端点", "kiro" => "找不到有效的 Kiro 端点",
"chatgpt_web" => "找不到有效的 openai:image 端点", "chatgpt_web" => "找不到有效的 openai:image 端点",
"claude_code" => "找不到有效的 claude:messages 端点",
"gemini_cli" | "vertex_ai" => "找不到有效的 gemini:generate_content 端点",
_ => "找不到有效端点", _ => "找不到有效端点",
}; };
return Ok(Some( return Ok(Some(

View File

@@ -50,6 +50,25 @@ pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Ve
admin_provider_quota_pure::normalize_string_id_list(values) admin_provider_quota_pure::normalize_string_id_list(values)
} }
pub(crate) fn provider_type_supports_quota_refresh(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"codex" | "kiro" | "antigravity" | "chatgpt_web"
)
}
pub(crate) fn unsupported_provider_quota_refresh_message(provider_type: &str) -> String {
match provider_type.trim().to_ascii_lowercase().as_str() {
"claude_code" => "Claude Code 暂不支持自动刷新额度:上游没有稳定可用的账号额度查询接口",
"gemini_cli" => {
"Gemini CLI 暂不支持自动刷新额度:当前只能通过模型同步/缓存快照展示已知配额信息"
}
"vertex_ai" => "Vertex AI 暂不支持自动刷新额度:额度属于 Google Cloud 项目/区域配额",
_ => "该 Provider 暂不支持自动刷新额度",
}
.to_string()
}
pub(super) fn coerce_json_u64(value: &serde_json::Value) -> Option<u64> { pub(super) fn coerce_json_u64(value: &serde_json::Value) -> Option<u64> {
admin_provider_quota_pure::coerce_json_u64(value) admin_provider_quota_pure::coerce_json_u64(value)
} }

View File

@@ -60,6 +60,48 @@ pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
.find(|endpoint| endpoint.is_active) .find(|endpoint| endpoint.is_active)
.cloned() .cloned()
}), }),
"claude_code" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned(),
"gemini_cli" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned(),
"vertex_ai" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:generate_content")
})
.cloned()
.or_else(|| {
endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:messages")
})
.cloned()
}),
_ => endpoints _ => endpoints
.iter() .iter()
.find(|endpoint| endpoint.is_active) .find(|endpoint| endpoint.is_active)

View File

@@ -14,10 +14,6 @@ pub(super) fn pool_cooldown_key(provider_id: &str, key_id: &str) -> String {
format!("ap:{provider_id}:cooldown:{key_id}") format!("ap:{provider_id}:cooldown:{key_id}")
} }
pub(super) fn pool_lease_key(provider_id: &str, key_id: &str) -> String {
format!("ap:{provider_id}:lease:{key_id}")
}
pub(super) fn pool_cooldown_index_key(provider_id: &str) -> String { pub(super) fn pool_cooldown_index_key(provider_id: &str) -> String {
format!("ap:{provider_id}:cooldown_idx") format!("ap:{provider_id}:cooldown_idx")
} }

View File

@@ -1,23 +1,4 @@
use super::keys::pool_lease_key;
use aether_runtime_state::{DataLayerError, RuntimeLockLease, RuntimeState}; use aether_runtime_state::{DataLayerError, RuntimeLockLease, RuntimeState};
use std::time::Duration;
pub(crate) const ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS: u64 = 15 * 60 * 1000;
pub(crate) async fn try_claim_admin_provider_pool_key(
runtime: &RuntimeState,
provider_id: &str,
key_id: &str,
owner: &str,
) -> Result<Option<RuntimeLockLease>, DataLayerError> {
runtime
.lock_try_acquire(
&pool_lease_key(provider_id, key_id),
owner,
Duration::from_millis(ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS),
)
.await
}
pub(crate) async fn release_admin_provider_pool_key_lease( pub(crate) async fn release_admin_provider_pool_key_lease(
runtime: &RuntimeState, runtime: &RuntimeState,

View File

@@ -5,16 +5,14 @@ mod reads;
mod status; mod status;
mod writes; mod writes;
pub(crate) use self::leases::{ pub(crate) use self::leases::release_admin_provider_pool_key_lease;
release_admin_provider_pool_key_lease, try_claim_admin_provider_pool_key,
ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS,
};
pub(crate) use self::mutations::{ pub(crate) use self::mutations::{
clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost, clear_admin_provider_pool_cooldown, reset_admin_provider_pool_cost,
}; };
pub(crate) use self::reads::{ pub(crate) use self::reads::{
read_admin_provider_pool_cooldown_count, read_admin_provider_pool_cooldown_counts, read_admin_provider_pool_cooldown_count, read_admin_provider_pool_cooldown_counts,
read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_runtime_state, read_admin_provider_pool_cooldown_key_ids, read_admin_provider_pool_key_cooldown_reason,
read_admin_provider_pool_runtime_state,
}; };
pub(crate) use self::status::build_admin_provider_pool_status_payload; pub(crate) use self::status::build_admin_provider_pool_status_payload;
pub(crate) use self::writes::{ pub(crate) use self::writes::{

View File

@@ -7,7 +7,7 @@ use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_af
use crate::handlers::admin::provider::shared::support::{ use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolConfig, AdminProviderPoolRuntimeState,
}; };
use aether_runtime_state::RuntimeState; use aether_runtime_state::{DataLayerError, RuntimeState};
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn; use tracing::warn;
@@ -202,3 +202,13 @@ pub(crate) async fn read_admin_provider_pool_cooldown_key_ids(
.await .await
.unwrap_or_default() .unwrap_or_default()
} }
pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
runtime: &RuntimeState,
provider_id: &str,
key_id: &str,
) -> Result<Option<String>, DataLayerError> {
runtime
.kv_get(&pool_cooldown_key(provider_id, key_id))
.await
}

View File

@@ -328,10 +328,7 @@ pub(crate) async fn maybe_build_local_admin_management_tokens_response(
.admin_principal .admin_principal
.as_ref() .as_ref()
.and_then(|principal| principal.management_token_permissions.as_deref()) .and_then(|principal| principal.management_token_permissions.as_deref())
.map_or( .is_none_or(management_token_permissions_cover_all_assignable_permissions);
true,
management_token_permissions_cover_all_assignable_permissions,
);
if is_management_token && !management_token_is_full { if is_management_token && !management_token_is_full {
return Ok(Some( return Ok(Some(
( (

View File

@@ -2,10 +2,10 @@ pub(crate) use super::super::admin::provider::pool::config::{
admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value, admin_provider_pool_cache_affinity_enabled, admin_provider_pool_config_from_config_value,
}; };
pub(crate) use super::super::admin::provider::pool::runtime::{ pub(crate) use super::super::admin::provider::pool::runtime::{
admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_runtime_state, admin_provider_pool_key_circuit_breaker_reason, read_admin_provider_pool_key_cooldown_reason,
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout, read_admin_provider_pool_runtime_state, record_admin_provider_pool_error,
record_admin_provider_pool_success, release_admin_provider_pool_key_lease, record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
try_claim_admin_provider_pool_key, ADMIN_PROVIDER_POOL_KEY_LEASE_TTL_MS, release_admin_provider_pool_key_lease,
}; };
pub(crate) use super::super::admin::provider::shared::support::{ pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset, AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,

View File

@@ -16,6 +16,7 @@ use tracing::{debug, info, warn};
use crate::admin_api::{ use crate::admin_api::{
admin_provider_pool_config, provider_oauth_runtime_endpoint_for_provider, admin_provider_pool_config, provider_oauth_runtime_endpoint_for_provider,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally, refresh_antigravity_provider_quota_locally, refresh_chatgpt_web_provider_quota_locally,
refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally, AdminAppState, refresh_codex_provider_quota_locally, refresh_kiro_provider_quota_locally, AdminAppState,
}; };
@@ -108,10 +109,7 @@ fn now_unix_secs() -> u64 {
} }
fn provider_supports_quota_probe(provider_type: &str) -> bool { fn provider_supports_quota_probe(provider_type: &str) -> bool {
matches!( provider_type_supports_quota_refresh(provider_type)
provider_type.trim().to_ascii_lowercase().as_str(),
"codex" | "kiro" | "antigravity" | "chatgpt_web"
)
} }
fn json_number(value: Option<&Value>) -> Option<f64> { fn json_number(value: Option<&Value>) -> Option<f64> {
@@ -426,6 +424,37 @@ fn endpoint_for_probe(
provider_oauth_runtime_endpoint_for_provider(provider_type, endpoints) provider_oauth_runtime_endpoint_for_provider(provider_type, endpoints)
} }
async fn endpoint_for_probe_with_reconcile(
state: &AppState,
admin_state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
provider_type: &str,
endpoints_by_provider: &mut BTreeMap<String, Vec<StoredProviderCatalogEndpoint>>,
) -> Result<Option<StoredProviderCatalogEndpoint>, GatewayError> {
let endpoints = endpoints_by_provider
.get(&provider.id)
.map(Vec::as_slice)
.unwrap_or(&[]);
if let Some(endpoint) = endpoint_for_probe(provider_type, endpoints) {
return Ok(Some(endpoint));
}
if admin_state
.fixed_provider_template(&provider.provider_type)
.is_none()
{
return Ok(None);
}
reconcile_admin_fixed_provider_template_endpoints(admin_state, provider).await?;
let refreshed = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let endpoint = endpoint_for_probe(provider_type, &refreshed);
endpoints_by_provider.insert(provider.id.clone(), refreshed);
Ok(endpoint)
}
async fn refresh_provider_probe_keys( async fn refresh_provider_probe_keys(
admin_state: &AdminAppState<'_>, admin_state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider, provider: &StoredProviderCatalogProvider,
@@ -688,10 +717,15 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
}; };
for (provider, provider_type, pool_config) in providers { for (provider, provider_type, pool_config) in providers {
let endpoints = endpoints_by_provider let Some(endpoint) = endpoint_for_probe_with_reconcile(
.remove(&provider.id) state,
.unwrap_or_default(); &admin_state,
let Some(endpoint) = endpoint_for_probe(&provider_type, &endpoints) else { &provider,
&provider_type,
&mut endpoints_by_provider,
)
.await?
else {
summary.providers_skipped += 1; summary.providers_skipped += 1;
debug!( debug!(
provider_id = %provider.id, provider_id = %provider.id,
@@ -700,6 +734,9 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
); );
continue; continue;
}; };
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_else(|| vec![endpoint.clone()]);
let interval_minutes = pool_config.probing_interval_minutes; let interval_minutes = pool_config.probing_interval_minutes;
let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60); let interval_seconds = interval_minutes.clamp(1, 1440).saturating_mul(60);

View File

@@ -829,6 +829,233 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_ad
upstream_handle.abort(); upstream_handle.abort();
} }
#[tokio::test]
async fn gateway_refresh_kiro_quota_reconciles_missing_fixed_endpoint_before_refresh() {
let seen_endpoint_id = Arc::new(Mutex::new(None::<String>));
let seen_endpoint_id_clone = Arc::clone(&seen_endpoint_id);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_endpoint_id_inner = Arc::clone(&seen_endpoint_id_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_endpoint_id_inner.lock().expect("mutex should lock") =
Some(plan.endpoint_id.clone());
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
"subscriptionTitle": "KIRO PRO"
},
"usageBreakdownList": [{
"currentUsageWithPrecision": 1.0,
"usageLimitWithPrecision": 10.0,
"nextDateReset": 1_900_000_000u64
}]
})),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"access_token":"kiro-access-token","api_region":"us-west-2"}"#,
)
.expect("auth config ciphertext should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "__placeholder__")
.expect("api key ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-kiro-reconcile".to_string(),
"provider-kiro-reconcile".to_string(),
"default".to_string(),
"bearer".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["claude:messages"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-kiro-reconcile".to_string(),
"kiro".to_string(),
Some("https://example.com".to_string()),
"kiro".to_string(),
)
.expect("provider should build")],
Vec::new(),
vec![key],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-kiro-reconcile/refresh-quota"
))
.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["success"], 1);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["results"][0]["status"], "success");
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&["provider-kiro-reconcile".to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), 1);
assert_eq!(endpoints[0].api_format, "claude:messages");
assert_eq!(endpoints[0].base_url, "https://q.{region}.amazonaws.com");
assert_eq!(
*seen_endpoint_id.lock().expect("mutex should lock"),
Some(endpoints[0].id.clone())
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_before_clear_message(
) {
let cases = [
(
"provider-claude-code-reconcile",
"claude_code",
1usize,
"claude:messages",
"https://api.anthropic.com",
"Claude Code 暂不支持自动刷新额度",
),
(
"provider-gemini-cli-reconcile",
"gemini_cli",
1usize,
"gemini:generate_content",
"https://cloudcode-pa.googleapis.com",
"Gemini CLI 暂不支持自动刷新额度",
),
(
"provider-vertex-ai-reconcile",
"vertex_ai",
2usize,
"gemini:generate_content",
"https://aiplatform.googleapis.com",
"Vertex AI 暂不支持自动刷新额度",
),
];
let providers = cases
.iter()
.map(|(provider_id, provider_type, _, _, _, _)| {
StoredProviderCatalogProvider::new(
(*provider_id).to_string(),
(*provider_type).to_string(),
Some("https://example.com".to_string()),
(*provider_type).to_string(),
)
.expect("provider should build")
})
.collect::<Vec<_>>();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
providers,
Vec::new(),
Vec::new(),
));
let gateway = build_router_with_state(
build_state_with_execution_runtime_override("http://127.0.0.1:1")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let client = reqwest::Client::new();
for (provider_id, _, endpoint_count, api_format, base_url, message_prefix) in cases {
let response = client
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/{provider_id}/refresh-quota"
))
.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["success"], 0);
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 0);
assert!(payload["message"]
.as_str()
.expect("message should be string")
.starts_with(message_prefix));
let endpoints = provider_catalog_repository
.list_endpoints_by_provider_ids(&[provider_id.to_string()])
.await
.expect("endpoints should read");
assert_eq!(endpoints.len(), endpoint_count);
assert!(endpoints
.iter()
.any(|endpoint| endpoint.api_format == api_format && endpoint.base_url == base_url));
}
gateway_handle.abort();
}
#[tokio::test] #[tokio::test]
async fn gateway_reports_codex_quota_runtime_failures_locally_without_falling_back_to_admin_passthrough( async fn gateway_reports_codex_quota_runtime_failures_locally_without_falling_back_to_admin_passthrough(
) { ) {

View File

@@ -662,19 +662,6 @@ const formatNumber = (num: number): string => {
return num.toLocaleString('zh-CN') return num.toLocaleString('zh-CN')
} }
// 计算最终状态:优先检查进行中状态,再使用外部状态码
const computedFinalStatus = computed(() => {
const hasPending = trace.value?.candidates?.some(
c => c.status === 'pending' || c.status === 'streaming'
)
return resolveTimelineFinalStatus({
hasPendingCandidates: hasPending,
statusCode: props.overrideStatusCode,
requestStatus: props.requestStatus ?? usageData.value?.status,
traceFinalStatus: trace.value?.final_status,
})
})
// 获取最终状态标签 // 获取最终状态标签
const getFinalStatusLabel = (status: string) => { const getFinalStatusLabel = (status: string) => {
const labels: Record<string, string> = { const labels: Record<string, string> = {
@@ -806,11 +793,25 @@ const STATUS_PRIORITY: Record<string, number> = {
} }
const isParticipatedCandidate = (candidate: CandidateRecord): boolean => { const isParticipatedCandidate = (candidate: CandidateRecord): boolean => {
if (candidate.status === 'available' || candidate.status === 'unused') return false return TIMELINE_STATUS.includes(candidate.status)
if (candidate.status === 'pending' && !candidate.started_at) return false
return true
} }
const isLiveCandidate = (candidate: CandidateRecord): boolean => {
if (candidate.status === 'streaming') return true
return candidate.status === 'pending' && Boolean(candidate.started_at)
}
// 计算最终状态:优先检查真正已启动的进行中状态,再使用外部状态码
const computedFinalStatus = computed(() => {
const hasPending = trace.value?.candidates?.some(isLiveCandidate)
return resolveTimelineFinalStatus({
hasPendingCandidates: hasPending,
statusCode: props.overrideStatusCode,
requestStatus: props.requestStatus ?? usageData.value?.status,
traceFinalStatus: trace.value?.final_status,
})
})
const compareBySchedulingOrder = (a: CandidateRecord, b: CandidateRecord): number => { const compareBySchedulingOrder = (a: CandidateRecord, b: CandidateRecord): number => {
if (a.candidate_index !== b.candidate_index) { if (a.candidate_index !== b.candidate_index) {
return a.candidate_index - b.candidate_index return a.candidate_index - b.candidate_index
@@ -1695,10 +1696,7 @@ const propsRequestIsActive = computed(() => {
}) })
const traceHasActiveCandidate = computed(() => { const traceHasActiveCandidate = computed(() => {
return rawTimeline.value.some((candidate) => { return rawTimeline.value.some(isLiveCandidate)
const status = getDisplayStatus(candidate)
return status === 'pending' || status === 'streaming'
})
}) })
const traceFinalIsTerminal = computed(() => { const traceFinalIsTerminal = computed(() => {
@@ -1751,15 +1749,12 @@ watch(groupedTimeline, (newGroups) => {
} }
// 查找正在进行的组 // 查找正在进行的组
const activeIdx = newGroups.findIndex(g => g.primaryStatus === 'pending' || g.primaryStatus === 'streaming') const activeIdx = newGroups.findIndex(g => g.allAttempts.some(isLiveCandidate))
if (activeIdx >= 0) { if (activeIdx >= 0) {
selectedGroupIndex.value = activeIdx selectedGroupIndex.value = activeIdx
// 选中正在进行的尝试,而非最后一个 // 选中正在进行的尝试,而非最后一个
const group = newGroups[activeIdx] const group = newGroups[activeIdx]
const attemptIdx = group.allAttempts.findIndex(a => { const attemptIdx = group.allAttempts.findIndex(isLiveCandidate)
const status = getDisplayStatus(a)
return status === 'pending' || status === 'streaming'
})
selectedAttemptIndex.value = attemptIdx >= 0 ? attemptIdx : group.allAttempts.length - 1 selectedAttemptIndex.value = attemptIdx >= 0 ? attemptIdx : group.allAttempts.length - 1
return return
} }
@@ -1853,8 +1848,8 @@ const formatDuration = (startStr: string, endStr: string): string => {
// 获取状态标签 // 获取状态标签
const getStatusLabel = (status: string) => { const getStatusLabel = (status: string) => {
const labels: Record<string, string> = { const labels: Record<string, string> = {
available: '未执行', available: '可用未尝试',
unused: '未执行', unused: '未使用',
pending: '进行中', pending: '进行中',
streaming: '传输中', streaming: '传输中',
stream_interrupted: '流中断', stream_interrupted: '流中断',

View File

@@ -180,7 +180,7 @@ describe('HorizontalRequestTimeline', () => {
]) ])
}) })
it('orders visible candidates by scheduling index and hides unstarted lazy candidates', async () => { it('orders visible candidates by scheduling index and includes unattempted candidates', async () => {
const trace = buildTrace([ const trace = buildTrace([
buildCandidate({ buildCandidate({
id: 'cand-success', id: 'cand-success',
@@ -244,7 +244,17 @@ describe('HorizontalRequestTimeline', () => {
const labels = [...root.querySelectorAll<HTMLElement>('.node-label')] const labels = [...root.querySelectorAll<HTMLElement>('.node-label')]
.map(label => label.textContent?.trim()) .map(label => label.textContent?.trim())
expect(labels).toEqual(['Provider Skipped', 'Provider Failed', 'Provider Success']) expect(labels).toEqual([
'Provider Available',
'Provider Skipped',
'Provider Pending',
'Provider Failed',
'Provider Success',
])
const nodeDots = [...root.querySelectorAll<HTMLElement>('.node-dot')]
expect(nodeDots[0].classList.contains('status-available')).toBe(true)
expect(nodeDots[2].classList.contains('status-pending')).toBe(true)
}) })
it('uses candidate terminal status for node colors instead of overriding with HTTP code', async () => { it('uses candidate terminal status for node colors instead of overriding with HTTP code', async () => {

View File

@@ -24,7 +24,7 @@ function buildCandidate(
} }
describe('poolTrace', () => { describe('poolTrace', () => {
it('keeps only pool nodes that actually participated in scheduling audit fallback', () => { it('keeps pool audit nodes even when they were not attempted', () => {
const attempts = buildPoolAttemptCandidatesFromAudit([], [ const attempts = buildPoolAttemptCandidatesFromAudit([], [
{ {
candidate_index: 0, candidate_index: 0,
@@ -68,11 +68,13 @@ describe('poolTrace', () => {
}, },
], 'req-1') ], 'req-1')
expect(attempts).toHaveLength(2) expect(attempts).toHaveLength(3)
expect(attempts[0].key_id).toBe('key-success') expect(attempts[0].key_id).toBe('key-success')
expect(attempts[0].status).toBe('success') expect(attempts[0].status).toBe('success')
expect(attempts[1].key_id).toBe('key-skipped') expect(attempts[1].key_id).toBe('key-skipped')
expect(attempts[1].status).toBe('skipped') expect(attempts[1].status).toBe('skipped')
expect(attempts[2].key_id).toBe('key-available')
expect(attempts[2].status).toBe('available')
}) })
it('preserves real trace attempts even when audit status is non-standard', () => { it('preserves real trace attempts even when audit status is non-standard', () => {

View File

@@ -12,11 +12,6 @@ export const TIMELINE_STATUS: CandidateRecord['status'][] = [
'stream_interrupted', 'stream_interrupted',
] ]
const POOL_HIDDEN_STATUS = new Set<CandidateRecord['status']>([
'available',
'unused',
])
const PROVIDER_TYPE_LIKE_NAMES = new Set<string>([ const PROVIDER_TYPE_LIKE_NAMES = new Set<string>([
'codex', 'codex',
'kiro', 'kiro',
@@ -40,9 +35,7 @@ export const makeAttemptKey = (candidateIndex: number, retryIndex: number): stri
} }
export const isPoolParticipatedCandidate = (candidate: CandidateRecord): boolean => { export const isPoolParticipatedCandidate = (candidate: CandidateRecord): boolean => {
if (POOL_HIDDEN_STATUS.has(candidate.status)) return false return TIMELINE_STATUS.includes(candidate.status)
if (candidate.status === 'pending' && !candidate.started_at) return false
return true
} }
export const isAttemptedCandidate = ( export const isAttemptedCandidate = (