mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: update gateway pool and usage flows
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
@@ -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]
|
||||||
|
|||||||
@@ -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,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
@@ -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)
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|||||||
@@ -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")
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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::{
|
||||||
|
|||||||
@@ -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
|
||||||
|
}
|
||||||
|
|||||||
@@ -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(
|
||||||
(
|
(
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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);
|
||||||
|
|||||||
@@ -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(
|
||||||
) {
|
) {
|
||||||
|
|||||||
@@ -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: '流中断',
|
||||||
|
|||||||
@@ -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 () => {
|
||||||
|
|||||||
@@ -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', () => {
|
||||||
|
|||||||
@@ -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 = (
|
||||||
|
|||||||
Reference in New Issue
Block a user