mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
Merge branch 'fawney19:main' into main
This commit is contained in:
@@ -709,31 +709,39 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
}
|
||||
|
||||
async fn refill_queued_candidates(&mut self) -> bool {
|
||||
let mut candidates = Vec::new();
|
||||
let refill_target = self.window_size.max(1) as usize;
|
||||
// Keep pool expansion bounded; the cursor freezes one small window at a time.
|
||||
while candidates.len() < refill_target {
|
||||
let Some(mut page_candidates) = self.next_page_candidates().await else {
|
||||
break;
|
||||
};
|
||||
candidates.append(&mut page_candidates);
|
||||
}
|
||||
|
||||
if candidates.is_empty() {
|
||||
return false;
|
||||
}
|
||||
loop {
|
||||
let mut candidates = Vec::new();
|
||||
// Keep pool expansion bounded; the cursor freezes one small window at a time.
|
||||
while candidates.len() < refill_target {
|
||||
let Some(mut page_candidates) = self.next_page_candidates().await else {
|
||||
break;
|
||||
};
|
||||
candidates.append(&mut page_candidates);
|
||||
}
|
||||
|
||||
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
self.sticky_session_token.as_deref(),
|
||||
)
|
||||
.await;
|
||||
scheduled.truncate(refill_target);
|
||||
self.record_skipped_candidates(&skipped);
|
||||
self.queued_candidates.extend(scheduled.drain(..));
|
||||
self.skipped_candidates.append(&mut skipped);
|
||||
!self.queued_candidates.is_empty()
|
||||
if candidates.is_empty() {
|
||||
return false;
|
||||
}
|
||||
|
||||
let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
|
||||
self.state,
|
||||
candidates,
|
||||
self.sticky_session_token.as_deref(),
|
||||
)
|
||||
.await;
|
||||
self.record_skipped_candidates(&skipped);
|
||||
self.skipped_candidates.append(&mut skipped);
|
||||
|
||||
if scheduled.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
scheduled.truncate(refill_target);
|
||||
self.queued_candidates.extend(scheduled.drain(..));
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
|
||||
@@ -2798,6 +2806,78 @@ mod tests {
|
||||
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_continues_after_exhausted_window() {
|
||||
let provider_config = Some(json!({
|
||||
"pool_advanced": {
|
||||
"skip_exhausted_accounts": true
|
||||
}
|
||||
}));
|
||||
let (provider, endpoint, mut keys, rows) = large_pool_fixture(3, provider_config.clone());
|
||||
for key in keys.iter_mut().take(2) {
|
||||
key.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"provider_type": "openai",
|
||||
"exhausted": true,
|
||||
"usage_ratio": 1.0,
|
||||
"windows": [
|
||||
{
|
||||
"code": "daily",
|
||||
"used_ratio": 1.0,
|
||||
"remaining_ratio": 0.0
|
||||
}
|
||||
]
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
let data_state =
|
||||
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
|
||||
Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
keys,
|
||||
)),
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
|
||||
)
|
||||
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
|
||||
let app = AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data_state);
|
||||
let group = sample_eligible_candidate(
|
||||
"provider-pool",
|
||||
"endpoint-1",
|
||||
"pool-group",
|
||||
10,
|
||||
provider_config,
|
||||
);
|
||||
|
||||
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
|
||||
cursor.window_size = 2;
|
||||
cursor.page_size = 2;
|
||||
cursor.max_scanned_keys = 4;
|
||||
|
||||
let candidate = cursor
|
||||
.next_key()
|
||||
.await
|
||||
.expect("cursor should scan past an exhausted window");
|
||||
assert_eq!(candidate.candidate.key_id, "key-00002");
|
||||
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
|
||||
assert!(candidate.orchestration.pool_key_lease.is_none());
|
||||
assert_eq!(
|
||||
cursor
|
||||
.skip_reason_counts
|
||||
.get(aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON),
|
||||
Some(&2)
|
||||
);
|
||||
|
||||
let skipped = cursor.take_skipped_candidates();
|
||||
assert_eq!(skipped.len(), 2);
|
||||
assert!(skipped.iter().all(|candidate| {
|
||||
candidate.skip_reason == aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
|
||||
}));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
|
||||
const KEY_COUNT: usize = 2048;
|
||||
|
||||
@@ -53,6 +53,15 @@ fn admin_usage_aggregation_by_provider_json(
|
||||
json!(rows
|
||||
.iter()
|
||||
.map(|row| {
|
||||
let identity_source = match row.secondary_name.as_deref() {
|
||||
Some("legacy_name") => "legacy_name",
|
||||
_ => "provider_id",
|
||||
};
|
||||
let provider_id = if identity_source == "provider_id" {
|
||||
json!(row.group_key)
|
||||
} else {
|
||||
serde_json::Value::Null
|
||||
};
|
||||
let success_count = row.success_count.unwrap_or_default();
|
||||
let error_count = row.request_count.saturating_sub(success_count);
|
||||
let success_rate = if row.request_count == 0 {
|
||||
@@ -67,7 +76,9 @@ fn admin_usage_aggregation_by_provider_json(
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(row.group_key.as_str());
|
||||
json!({
|
||||
"provider_id": row.group_key,
|
||||
"provider_id": provider_id,
|
||||
"provider_key": row.group_key,
|
||||
"provider_identity_source": identity_source,
|
||||
"provider": provider_name,
|
||||
"request_count": row.request_count,
|
||||
"total_tokens": row.total_tokens,
|
||||
|
||||
@@ -2,7 +2,9 @@ use super::state::{
|
||||
decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string,
|
||||
json_u64_value,
|
||||
};
|
||||
use crate::handlers::admin::admin_provider_pool_config;
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
|
||||
use crate::provider_key_auth::provider_active_api_formats;
|
||||
use crate::GatewayError;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
@@ -168,6 +170,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
||||
.app()
|
||||
.invalidate_local_oauth_refresh_entry(&key.id)
|
||||
.await;
|
||||
seed_provider_oauth_pool_score(state, provider_id, key, now_unix_secs).await;
|
||||
}
|
||||
Ok(created)
|
||||
}
|
||||
@@ -221,10 +224,75 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
.app()
|
||||
.invalidate_local_oauth_refresh_entry(&key.id)
|
||||
.await;
|
||||
seed_provider_oauth_pool_score(state, &existing_key.provider_id, key, now_unix_secs).await;
|
||||
}
|
||||
Ok(persisted)
|
||||
}
|
||||
|
||||
async fn seed_provider_oauth_pool_score(
|
||||
state: &AdminAppState<'_>,
|
||||
provider_id: &str,
|
||||
key: &StoredProviderCatalogKey,
|
||||
now_unix_secs: u64,
|
||||
) {
|
||||
let provider_id = provider_id.to_string();
|
||||
let provider = match state
|
||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
{
|
||||
Ok(mut providers) => providers.pop(),
|
||||
Err(err) => {
|
||||
tracing::debug!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %key.id,
|
||||
error = ?err,
|
||||
"gateway provider oauth provisioning: failed to read provider for pool score seed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let Some(provider) = provider else {
|
||||
return;
|
||||
};
|
||||
let Some(pool_config) = admin_provider_pool_config(&provider) else {
|
||||
return;
|
||||
};
|
||||
let endpoints = match state
|
||||
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
|
||||
.await
|
||||
{
|
||||
Ok(endpoints) => endpoints,
|
||||
Err(err) => {
|
||||
tracing::debug!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %key.id,
|
||||
error = ?err,
|
||||
"gateway provider oauth provisioning: failed to read endpoints for pool score seed"
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
|
||||
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
|
||||
state.as_ref(),
|
||||
&provider,
|
||||
&pool_config,
|
||||
&endpoints,
|
||||
std::slice::from_ref(key),
|
||||
now_unix_secs,
|
||||
score_ensure_budget,
|
||||
)
|
||||
.await
|
||||
{
|
||||
tracing::debug!(
|
||||
provider_id = %provider_id,
|
||||
key_id = %key.id,
|
||||
error = ?err,
|
||||
"gateway provider oauth provisioning: failed to seed pool score row"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
fn provider_oauth_catalog_key_api_formats(
|
||||
provider_type: &str,
|
||||
api_formats: &[String],
|
||||
|
||||
@@ -115,6 +115,30 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
|
||||
admin_menu_group: Some("system"),
|
||||
admin_menu_order: 60,
|
||||
},
|
||||
AdminModuleDefinition {
|
||||
name: "payment_gateways",
|
||||
display_name: "支付配置",
|
||||
description: "配置易支付、支付宝官方、微信支付官方和 Stripe 等支付网关",
|
||||
category: "integration",
|
||||
env_key: "PAYMENT_GATEWAYS_AVAILABLE",
|
||||
default_available: true,
|
||||
admin_route: Some("/admin/payment-gateways"),
|
||||
admin_menu_icon: Some("CreditCard"),
|
||||
admin_menu_group: None,
|
||||
admin_menu_order: 70,
|
||||
},
|
||||
AdminModuleDefinition {
|
||||
name: "referral",
|
||||
display_name: "邀请返利",
|
||||
description: "管理用户邀请关系与返利记录,支持比例返利和人头返利",
|
||||
category: "integration",
|
||||
env_key: "REFERRAL_AVAILABLE",
|
||||
default_available: true,
|
||||
admin_route: Some("/admin/referrals"),
|
||||
admin_menu_icon: Some("Gift"),
|
||||
admin_menu_group: Some("management"),
|
||||
admin_menu_order: 75,
|
||||
},
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
|
||||
@@ -46,7 +46,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
|
||||
|
||||
let (endpoints_result, keys_result) = tokio::join!(
|
||||
state.list_provider_catalog_endpoints_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_key_summaries_by_provider_ids(&provider_ids),
|
||||
state.list_provider_catalog_keys_by_provider_ids(&provider_ids),
|
||||
);
|
||||
|
||||
let active_endpoints = endpoints_result
|
||||
|
||||
@@ -5,7 +5,12 @@ use aether_crypto::{
|
||||
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
|
||||
};
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, StoredProviderCatalogEndpoint,
|
||||
StoredProviderCatalogKey, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats,
|
||||
StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use axum::body::Body;
|
||||
use axum::routing::any;
|
||||
use axum::{extract::Request, Json, Router};
|
||||
@@ -22,6 +27,103 @@ use crate::constants::{
|
||||
};
|
||||
use crate::data::GatewayDataState;
|
||||
|
||||
struct SummaryNullingProviderCatalogReadRepository {
|
||||
inner: InMemoryProviderCatalogReadRepository,
|
||||
}
|
||||
|
||||
impl SummaryNullingProviderCatalogReadRepository {
|
||||
fn seed(
|
||||
providers: Vec<StoredProviderCatalogProvider>,
|
||||
endpoints: Vec<StoredProviderCatalogEndpoint>,
|
||||
keys: Vec<StoredProviderCatalogKey>,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: InMemoryProviderCatalogReadRepository::seed(providers, endpoints, keys),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl ProviderCatalogReadRepository for SummaryNullingProviderCatalogReadRepository {
|
||||
async fn list_providers(
|
||||
&self,
|
||||
active_only: bool,
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
self.inner.list_providers(active_only).await
|
||||
}
|
||||
|
||||
async fn list_providers_by_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
|
||||
self.inner.list_providers_by_ids(provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_ids(
|
||||
&self,
|
||||
endpoint_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
self.inner.list_endpoints_by_ids(endpoint_ids).await
|
||||
}
|
||||
|
||||
async fn list_endpoints_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
|
||||
self.inner
|
||||
.list_endpoints_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_keys_by_ids(
|
||||
&self,
|
||||
key_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
self.inner.list_keys_by_ids(key_ids).await
|
||||
}
|
||||
|
||||
async fn list_keys_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
self.inner.list_keys_by_provider_ids(provider_ids).await
|
||||
}
|
||||
|
||||
async fn list_key_summaries_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
|
||||
let mut keys = self.inner.list_keys_by_provider_ids(provider_ids).await?;
|
||||
for key in &mut keys {
|
||||
key.internal_priority = 50;
|
||||
key.global_priority_by_format = None;
|
||||
key.rate_multipliers = None;
|
||||
key.request_count = None;
|
||||
key.success_count = None;
|
||||
key.error_count = None;
|
||||
key.total_response_time_ms = None;
|
||||
key.circuit_breaker_by_format = None;
|
||||
}
|
||||
Ok(keys)
|
||||
}
|
||||
|
||||
async fn list_keys_page(
|
||||
&self,
|
||||
query: &ProviderCatalogKeyListQuery,
|
||||
) -> Result<StoredProviderCatalogKeyPage, DataLayerError> {
|
||||
self.inner.list_keys_page(query).await
|
||||
}
|
||||
|
||||
async fn list_key_stats_by_provider_ids(
|
||||
&self,
|
||||
provider_ids: &[String],
|
||||
) -> Result<Vec<StoredProviderCatalogKeyStats>, DataLayerError> {
|
||||
self.inner
|
||||
.list_key_stats_by_provider_ids(provider_ids)
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_provider_keys_locally_with_trusted_admin_principal() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
@@ -2195,7 +2297,7 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
|
||||
key_b.created_at_unix_ms = Some(1_711_100_000);
|
||||
key_b.updated_at_unix_secs = Some(1_711_100_100);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
let provider_catalog_repository = Arc::new(SummaryNullingProviderCatalogReadRepository::seed(
|
||||
vec![
|
||||
sample_provider("provider-openai", "openai", 10),
|
||||
sample_provider("provider-claude", "claude", 20)
|
||||
@@ -2252,6 +2354,12 @@ async fn gateway_handles_admin_keys_grouped_by_format_locally_with_trusted_admin
|
||||
"https://api.openai.example"
|
||||
);
|
||||
assert_eq!(payload["openai:chat"][0]["capabilities"], json!(["1h缓存"]));
|
||||
assert_eq!(payload["openai:chat"][0]["format_priority"], 3);
|
||||
assert_eq!(
|
||||
payload["openai:chat"][0]["global_priority_by_format"]["openai:chat"],
|
||||
3
|
||||
);
|
||||
assert_eq!(payload["openai:chat"][0]["internal_priority"], 10);
|
||||
assert_eq!(payload["claude:messages"][0]["provider_active"], false);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
|
||||
@@ -12,8 +12,12 @@ use aether_data::repository::management_tokens::{
|
||||
use aether_data::repository::oauth_providers::{
|
||||
InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
|
||||
};
|
||||
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
|
||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
|
||||
use aether_data_contracts::repository::pool_scores::{
|
||||
GetPoolMemberScoresByIdsQuery, PoolMemberHardState, PoolMemberIdentity, PoolScoreReadRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
|
||||
};
|
||||
@@ -32,6 +36,7 @@ use super::super::{
|
||||
use crate::admin_api::{
|
||||
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
|
||||
};
|
||||
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
|
||||
use crate::audit::AdminAuditEvent;
|
||||
use crate::constants::{
|
||||
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
|
||||
@@ -1930,6 +1935,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
|
||||
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
|
||||
provider.provider_type = "chatgpt_web".to_string();
|
||||
provider.config = Some(json!({"pool_advanced": {}}));
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-chatgpt-web-image",
|
||||
"provider-chatgpt-web",
|
||||
@@ -1941,6 +1947,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
vec![endpoint],
|
||||
vec![],
|
||||
));
|
||||
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
|
||||
|
||||
let (token_url, token_handle) = start_server(token_server).await;
|
||||
let gateway = build_router_with_state(
|
||||
@@ -1950,6 +1957,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository))
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_provider_oauth_token_url_for_tests(
|
||||
@@ -2030,6 +2038,20 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
|
||||
assert_eq!(auth_config["plan_type"], "plus");
|
||||
assert_eq!(auth_config["user_id"], "user-pool-image");
|
||||
|
||||
let score_scope = provider_key_pool_score_scope();
|
||||
let score_identity =
|
||||
PoolMemberIdentity::provider_api_key("provider-chatgpt-web", persisted.id.clone());
|
||||
let scores = pool_score_repository
|
||||
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
|
||||
ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)],
|
||||
})
|
||||
.await
|
||||
.expect("pool score should load");
|
||||
assert_eq!(scores.len(), 1);
|
||||
assert_eq!(scores[0].member_id, persisted.id);
|
||||
assert_eq!(scores[0].hard_state, PoolMemberHardState::Unknown);
|
||||
assert!(scores[0].score > 0.0);
|
||||
|
||||
gateway_handle.abort();
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
@@ -665,8 +665,14 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
|
||||
let provider_items = provider_payload.as_array().expect("array response");
|
||||
assert_eq!(provider_items.len(), 2);
|
||||
assert_eq!(provider_items[0]["provider"], "OpenAI");
|
||||
assert_eq!(provider_items[0]["provider_id"], "provider-openai");
|
||||
assert_eq!(provider_items[0]["provider_key"], "provider-openai");
|
||||
assert_eq!(provider_items[0]["provider_identity_source"], "provider_id");
|
||||
assert_eq!(provider_items[0]["output_tokens"], 40);
|
||||
assert_eq!(provider_items[1]["provider"], "Anthropic");
|
||||
assert_eq!(provider_items[1]["provider_id"], "provider-anthropic");
|
||||
assert_eq!(provider_items[1]["provider_key"], "provider-anthropic");
|
||||
assert_eq!(provider_items[1]["provider_identity_source"], "provider_id");
|
||||
assert_eq!(provider_items[1]["output_tokens"], 20);
|
||||
|
||||
let api_format_response = admin_request(reqwest::Client::new().get(format!(
|
||||
@@ -693,6 +699,59 @@ async fn gateway_handles_admin_usage_aggregation_stats_locally_with_trusted_admi
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_admin_usage_aggregation_stats_for_legacy_provider_name_rows() {
|
||||
let legacy_usage = {
|
||||
let mut usage = sample_usage_row(
|
||||
"usage-legacy",
|
||||
"req-legacy",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"Legacy Provider",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
);
|
||||
usage.provider_id = None;
|
||||
usage.total_tokens = usage.input_tokens;
|
||||
usage
|
||||
};
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![legacy_usage]));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(
|
||||
usage_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/aggregation/stats?group_by=provider&limit=10&start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0"
|
||||
)))
|
||||
.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");
|
||||
let items = payload.as_array().expect("array response");
|
||||
assert_eq!(items.len(), 1);
|
||||
assert_eq!(items[0]["provider"], "Legacy Provider");
|
||||
assert_eq!(items[0]["provider_id"], serde_json::Value::Null);
|
||||
assert_eq!(items[0]["provider_key"], "Legacy Provider");
|
||||
assert_eq!(items[0]["provider_identity_source"], "legacy_name");
|
||||
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_service_unavailable_for_admin_usage_replay_without_provider_catalog_reader(
|
||||
) {
|
||||
|
||||
@@ -375,6 +375,51 @@ fn ensure_codex_chat_reasoning_defaults(
|
||||
.or_insert_with(|| json!(CODEX_DEFAULT_REASONING_SUMMARY));
|
||||
}
|
||||
|
||||
fn codex_tool_type_rejects_top_level_name(tool_type: &str) -> bool {
|
||||
let normalized = tool_type.trim().to_ascii_lowercase();
|
||||
!normalized.is_empty()
|
||||
&& normalized != "function"
|
||||
&& normalized != "custom"
|
||||
&& normalized != "namespace"
|
||||
}
|
||||
|
||||
fn strip_codex_hosted_tool_names_for_backend(body_object: &mut serde_json::Map<String, Value>) {
|
||||
let Some(tools) = body_object.get_mut("tools").and_then(Value::as_array_mut) else {
|
||||
return;
|
||||
};
|
||||
|
||||
for tool in tools {
|
||||
let Some(tool_object) = tool.as_object_mut() else {
|
||||
continue;
|
||||
};
|
||||
if tool_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(codex_tool_type_rejects_top_level_name)
|
||||
{
|
||||
tool_object.remove("name");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn strip_codex_hosted_tool_choice_name_for_backend(
|
||||
body_object: &mut serde_json::Map<String, Value>,
|
||||
) {
|
||||
let Some(tool_choice_object) = body_object
|
||||
.get_mut("tool_choice")
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
if tool_choice_object
|
||||
.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(codex_tool_type_rejects_top_level_name)
|
||||
{
|
||||
tool_choice_object.remove("name");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn apply_codex_openai_responses_special_body_edits(
|
||||
provider_request_body: &mut Value,
|
||||
provider_type: &str,
|
||||
@@ -420,6 +465,8 @@ pub fn apply_codex_openai_responses_special_body_edits(
|
||||
{
|
||||
body_object.insert("instructions".to_string(), json!(""));
|
||||
}
|
||||
strip_codex_hosted_tool_names_for_backend(body_object);
|
||||
strip_codex_hosted_tool_choice_name_for_backend(body_object);
|
||||
if is_openai_image_request(provider_api_format)
|
||||
|| codex_openai_responses_tool_choice_references_image_generation(body_object)
|
||||
{
|
||||
@@ -620,6 +667,87 @@ mod tests {
|
||||
assert_eq!(provider_request_body["parallel_tool_calls"], json!(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_responses_body_edits_preserve_function_tools_for_codex_backend() {
|
||||
let mut provider_request_body = json!({
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"tools": [{
|
||||
"type": "function",
|
||||
"name": "lookup_account",
|
||||
"description": "Lookup an account by id.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"account_id": {
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": ["account_id"],
|
||||
"additionalProperties": false
|
||||
},
|
||||
"strict": true
|
||||
}],
|
||||
"tool_choice": {
|
||||
"type": "function",
|
||||
"name": "lookup_account"
|
||||
}
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
provider_request_body["tools"][0]["name"],
|
||||
json!("lookup_account")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_request_body["tools"][0]["parameters"]["properties"]["account_id"]["type"],
|
||||
json!("string")
|
||||
);
|
||||
assert_eq!(
|
||||
provider_request_body["tool_choice"]["name"],
|
||||
json!("lookup_account")
|
||||
);
|
||||
assert!(provider_request_body["tools"][0].get("function").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_responses_body_edits_strip_name_from_hosted_web_search_tool() {
|
||||
let mut provider_request_body = json!({
|
||||
"input": [],
|
||||
"model": "gpt-5.4",
|
||||
"tools": [{
|
||||
"type": "web_search",
|
||||
"name": "web_search"
|
||||
}],
|
||||
"tool_choice": {
|
||||
"type": "web_search",
|
||||
"name": "web_search"
|
||||
}
|
||||
});
|
||||
|
||||
apply_codex_openai_responses_special_body_edits(
|
||||
&mut provider_request_body,
|
||||
"codex",
|
||||
"openai:responses",
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert!(provider_request_body["tools"][0].get("name").is_none());
|
||||
assert!(provider_request_body["tool_choice"].get("name").is_none());
|
||||
assert_eq!(
|
||||
provider_request_body["tool_choice"]["type"],
|
||||
json!("web_search")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compact_body_edits_strip_include_store_and_stream() {
|
||||
let mut provider_request_body = json!({
|
||||
|
||||
@@ -0,0 +1,10 @@
|
||||
-- Usage is a historical fact table. Backfill nullable provider_id snapshots
|
||||
-- from the unique provider name where the catalog row still exists.
|
||||
|
||||
UPDATE `usage` AS usage_rows
|
||||
JOIN providers
|
||||
ON providers.name = TRIM(usage_rows.provider_name)
|
||||
SET usage_rows.provider_id = providers.id
|
||||
WHERE usage_rows.provider_id IS NULL
|
||||
AND TRIM(COALESCE(usage_rows.provider_name, '')) <> ''
|
||||
AND LOWER(TRIM(COALESCE(usage_rows.provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending');
|
||||
@@ -4427,21 +4427,6 @@ END $mig$;
|
||||
|
||||
|
||||
|
||||
--
|
||||
-- Name: usage usage_provider_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||
--
|
||||
|
||||
DO $mig$ BEGIN
|
||||
ALTER TABLE ONLY public.usage
|
||||
ADD CONSTRAINT usage_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES public.providers(id) ON DELETE SET NULL;
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
WHEN duplicate_table THEN NULL;
|
||||
WHEN invalid_table_definition THEN NULL;
|
||||
END $mig$;
|
||||
|
||||
|
||||
|
||||
--
|
||||
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||
--
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
-- Usage is a historical fact table. Keep the provider_id snapshot even if the
|
||||
-- provider catalog row is deleted, and backfill rows that can still be matched
|
||||
-- by the unique provider name.
|
||||
|
||||
UPDATE public.usage AS usage_rows
|
||||
SET provider_id = providers.id
|
||||
FROM public.providers AS providers
|
||||
WHERE usage_rows.provider_id IS NULL
|
||||
AND BTRIM(COALESCE(usage_rows.provider_name, '')) <> ''
|
||||
AND lower(BTRIM(COALESCE(usage_rows.provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
AND providers.name = BTRIM(usage_rows.provider_name);
|
||||
|
||||
ALTER TABLE ONLY public.usage
|
||||
DROP CONSTRAINT IF EXISTS usage_provider_id_fkey;
|
||||
@@ -0,0 +1,18 @@
|
||||
-- Usage is a historical fact table. Backfill nullable provider_id snapshots
|
||||
-- from the unique provider name where the catalog row still exists.
|
||||
|
||||
UPDATE "usage"
|
||||
SET provider_id = (
|
||||
SELECT providers.id
|
||||
FROM providers
|
||||
WHERE providers.name = TRIM("usage".provider_name)
|
||||
LIMIT 1
|
||||
)
|
||||
WHERE provider_id IS NULL
|
||||
AND TRIM(COALESCE(provider_name, '')) <> ''
|
||||
AND LOWER(TRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM providers
|
||||
WHERE providers.name = TRIM("usage".provider_name)
|
||||
);
|
||||
@@ -627,21 +627,6 @@ END $mig$;
|
||||
|
||||
|
||||
|
||||
--
|
||||
-- Name: usage usage_provider_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||
--
|
||||
|
||||
DO $mig$ BEGIN
|
||||
ALTER TABLE ONLY public.usage
|
||||
ADD CONSTRAINT usage_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES public.providers(id) ON DELETE SET NULL;
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
WHEN duplicate_table THEN NULL;
|
||||
WHEN invalid_table_definition THEN NULL;
|
||||
END $mig$;
|
||||
|
||||
|
||||
|
||||
--
|
||||
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||
--
|
||||
|
||||
@@ -552,21 +552,6 @@ END $mig$;
|
||||
|
||||
|
||||
|
||||
--
|
||||
-- Name: usage usage_provider_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||
--
|
||||
|
||||
DO $mig$ BEGIN
|
||||
ALTER TABLE ONLY public.usage
|
||||
ADD CONSTRAINT usage_provider_id_fkey FOREIGN KEY (provider_id) REFERENCES public.providers(id) ON DELETE SET NULL;
|
||||
EXCEPTION
|
||||
WHEN duplicate_object THEN NULL;
|
||||
WHEN duplicate_table THEN NULL;
|
||||
WHEN invalid_table_definition THEN NULL;
|
||||
END $mig$;
|
||||
|
||||
|
||||
|
||||
--
|
||||
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||
--
|
||||
|
||||
@@ -7,7 +7,7 @@ use tracing::info;
|
||||
// Generated by build.rs from schema/bootstrap/postgres.
|
||||
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
||||
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520000000;
|
||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520010000;
|
||||
|
||||
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
||||
SELECT COUNT(*)::BIGINT
|
||||
|
||||
@@ -311,6 +311,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
||||
20260519120000,
|
||||
20260519130000,
|
||||
20260520000000,
|
||||
20260520010000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -631,6 +632,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
20260519120000,
|
||||
20260519130000,
|
||||
20260520000000,
|
||||
20260520010000,
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -653,6 +655,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
20260519120000,
|
||||
20260519130000,
|
||||
20260520000000,
|
||||
20260520010000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -1175,6 +1178,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
20260519120000,
|
||||
20260519130000,
|
||||
20260520000000,
|
||||
20260520010000,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -983,6 +983,30 @@ fn usage_provider_display_name(item: &StoredRequestUsageAudit) -> Option<String>
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_provider_id(item: &StoredRequestUsageAudit) -> Option<String> {
|
||||
let provider_id = item.provider_id.as_deref()?.trim();
|
||||
if provider_id.is_empty() || usage_reserved_provider_label(provider_id) {
|
||||
None
|
||||
} else {
|
||||
Some(provider_id.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_provider_aggregation_identity(
|
||||
item: &StoredRequestUsageAudit,
|
||||
) -> Option<(String, Option<String>, String)> {
|
||||
let display_name = usage_provider_display_name(item);
|
||||
if let Some(provider_id) = usage_provider_id(item) {
|
||||
return Some((provider_id, display_name, "provider_id".to_string()));
|
||||
}
|
||||
let display_name = display_name?;
|
||||
Some((
|
||||
display_name.clone(),
|
||||
Some(display_name),
|
||||
"legacy_name".to_string(),
|
||||
))
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
async fn find_by_id(
|
||||
@@ -1167,14 +1191,14 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
|| item.created_at_unix_ms >= query.created_until_unix_secs
|
||||
|| matches!(item.status.as_str(), "pending" | "streaming")
|
||||
|| (query.exclude_reserved_provider_labels
|
||||
&& usage_provider_display_name(item).is_none())
|
||||
&& usage_provider_aggregation_identity(item).is_none())
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
let provider_display_name =
|
||||
let provider_identity =
|
||||
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
|
||||
match usage_provider_display_name(item) {
|
||||
match usage_provider_aggregation_identity(item) {
|
||||
Some(value) => Some(value),
|
||||
None => continue,
|
||||
}
|
||||
@@ -1184,19 +1208,11 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
|
||||
let group_key = match query.group_by {
|
||||
UsageAuditAggregationGroupBy::Model => item.model.clone(),
|
||||
UsageAuditAggregationGroupBy::Provider => {
|
||||
let display_name = provider_display_name
|
||||
.as_deref()
|
||||
.expect("provider display name is set for provider aggregation");
|
||||
item.provider_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|provider_id| {
|
||||
!provider_id.is_empty() && !usage_reserved_provider_label(provider_id)
|
||||
})
|
||||
.unwrap_or(display_name)
|
||||
.to_string()
|
||||
}
|
||||
UsageAuditAggregationGroupBy::Provider => provider_identity
|
||||
.as_ref()
|
||||
.expect("provider identity is set for provider aggregation")
|
||||
.0
|
||||
.clone(),
|
||||
UsageAuditAggregationGroupBy::ApiFormat => item
|
||||
.api_format
|
||||
.clone()
|
||||
@@ -1211,7 +1227,17 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
&& (bucket.display_name.is_none()
|
||||
|| bucket.display_name.as_deref() == Some("Unknown"))
|
||||
{
|
||||
bucket.display_name = provider_display_name;
|
||||
bucket.display_name = provider_identity
|
||||
.as_ref()
|
||||
.and_then(|(_, display_name, _)| display_name.clone());
|
||||
}
|
||||
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
|
||||
&& (bucket.secondary_name.is_none()
|
||||
|| bucket.secondary_name.as_deref() == Some("legacy_name"))
|
||||
{
|
||||
bucket.secondary_name = provider_identity
|
||||
.as_ref()
|
||||
.map(|(_, _, identity_source)| identity_source.clone());
|
||||
}
|
||||
bucket.request_count = bucket.request_count.saturating_add(1);
|
||||
bucket.total_tokens = bucket.total_tokens.saturating_add(item.total_tokens);
|
||||
@@ -3149,6 +3175,12 @@ mod tests {
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_aggregation_skips_unknown_provider_labels() {
|
||||
let valid_provider = sample_usage("req-valid-provider", 300);
|
||||
|
||||
let mut legacy_provider = sample_usage("req-legacy-provider", 250);
|
||||
legacy_provider.provider_id = None;
|
||||
legacy_provider.provider_name = "Legacy Provider".to_string();
|
||||
|
||||
let mut unknown = sample_usage("req-unknown-provider", 100);
|
||||
unknown.provider_id = None;
|
||||
unknown.provider_name = "unknown".to_string();
|
||||
@@ -3158,7 +3190,8 @@ mod tests {
|
||||
typo_unknown.provider_name = "unknow".to_string();
|
||||
|
||||
let repository = InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage("req-valid-provider", 300),
|
||||
valid_provider,
|
||||
legacy_provider,
|
||||
unknown,
|
||||
typo_unknown,
|
||||
]);
|
||||
@@ -3174,9 +3207,29 @@ mod tests {
|
||||
.await
|
||||
.expect("aggregation should succeed");
|
||||
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].group_key, "provider-1");
|
||||
assert_eq!(rows[0].display_name.as_deref(), Some("OpenAI"));
|
||||
assert_eq!(rows.len(), 2);
|
||||
let provider_id_row = rows
|
||||
.iter()
|
||||
.find(|row| row.group_key == "provider-1")
|
||||
.expect("provider_id row should be present");
|
||||
assert_eq!(provider_id_row.display_name.as_deref(), Some("OpenAI"));
|
||||
assert_eq!(
|
||||
provider_id_row.secondary_name.as_deref(),
|
||||
Some("provider_id")
|
||||
);
|
||||
|
||||
let legacy_name_row = rows
|
||||
.iter()
|
||||
.find(|row| row.group_key == "Legacy Provider")
|
||||
.expect("legacy provider name row should be present");
|
||||
assert_eq!(
|
||||
legacy_name_row.display_name.as_deref(),
|
||||
Some("Legacy Provider")
|
||||
);
|
||||
assert_eq!(
|
||||
legacy_name_row.secondary_name.as_deref(),
|
||||
Some("legacy_name")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -3193,11 +3246,15 @@ mod tests {
|
||||
pending_provider.provider_id = None;
|
||||
pending_provider.provider_name = "pending".to_string();
|
||||
|
||||
let mut id_only_provider = sample_usage("req-id-only-provider", 350);
|
||||
id_only_provider.provider_name = "unknown".to_string();
|
||||
|
||||
let repository = InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage("req-valid-provider", 400),
|
||||
unknown,
|
||||
typo_unknown,
|
||||
pending_provider,
|
||||
id_only_provider,
|
||||
]);
|
||||
|
||||
let model_rows = repository
|
||||
@@ -3212,7 +3269,7 @@ mod tests {
|
||||
.expect("model aggregation should succeed");
|
||||
assert_eq!(model_rows.len(), 1);
|
||||
assert_eq!(model_rows[0].group_key, "gpt-4.1");
|
||||
assert_eq!(model_rows[0].request_count, 1);
|
||||
assert_eq!(model_rows[0].request_count, 2);
|
||||
|
||||
let api_format_rows = repository
|
||||
.aggregate_usage_audits(&UsageAuditAggregationQuery {
|
||||
@@ -3226,7 +3283,7 @@ mod tests {
|
||||
.expect("api format aggregation should succeed");
|
||||
assert_eq!(api_format_rows.len(), 1);
|
||||
assert_eq!(api_format_rows[0].group_key, "openai:chat");
|
||||
assert_eq!(api_format_rows[0].request_count, 1);
|
||||
assert_eq!(api_format_rows[0].request_count, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
||||
@@ -1488,8 +1488,16 @@ const REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL: &str =
|
||||
include_str!("queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql");
|
||||
|
||||
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
|
||||
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
|
||||
const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_id, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')";
|
||||
const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = r#" AND (
|
||||
(
|
||||
BTRIM(COALESCE("usage".provider_id, '')) <> ''
|
||||
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
)
|
||||
OR (
|
||||
BTRIM(COALESCE("usage".provider_name, '')) <> ''
|
||||
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
)
|
||||
)"#;
|
||||
const USAGE_RAW_PROVIDER_GROUP_KEY_SQL: &str = r#"CASE
|
||||
WHEN BTRIM(COALESCE("usage".provider_id, '')) = ''
|
||||
OR lower(BTRIM(COALESCE("usage".provider_id, ''))) IN ('unknown', 'unknow', 'pending')
|
||||
@@ -1502,13 +1510,33 @@ const USAGE_RAW_PROVIDER_DISPLAY_NAME_SQL: &str = r#"CASE
|
||||
THEN NULL
|
||||
ELSE BTRIM("usage".provider_name)
|
||||
END"#;
|
||||
const USAGE_PROVIDER_IDENTITY_SOURCE_SQL: &str = r#"CASE
|
||||
WHEN BTRIM(COALESCE("usage".provider_id, '')) <> ''
|
||||
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
THEN 'provider_id'
|
||||
WHEN BTRIM(COALESCE("usage".provider_name, '')) <> ''
|
||||
AND lower(BTRIM(COALESCE("usage".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
THEN 'legacy_name'
|
||||
ELSE NULL
|
||||
END"#;
|
||||
const USAGE_PROVIDER_IDENTITY_JOIN_SQL: &str = r#" LEFT JOIN providers AS provider_by_id
|
||||
ON BTRIM(COALESCE("usage".provider_id, '')) <> ''
|
||||
AND lower(BTRIM(COALESCE("usage".provider_id, ''))) NOT IN ('unknown', 'unknow', 'pending')
|
||||
AND provider_by_id.id = BTRIM("usage".provider_id)"#;
|
||||
const USAGE_RESOLVED_PROVIDER_GROUP_KEY_SQL: &str = r#"COALESCE(
|
||||
provider_by_id.id,
|
||||
BTRIM("usage".provider_id)
|
||||
CASE
|
||||
WHEN BTRIM(COALESCE("usage".provider_id, '')) = ''
|
||||
OR lower(BTRIM(COALESCE("usage".provider_id, ''))) IN ('unknown', 'unknow', 'pending')
|
||||
THEN NULL
|
||||
ELSE BTRIM("usage".provider_id)
|
||||
END,
|
||||
CASE
|
||||
WHEN BTRIM(COALESCE("usage".provider_name, '')) = ''
|
||||
OR lower(BTRIM(COALESCE("usage".provider_name, ''))) IN ('unknown', 'unknow', 'pending')
|
||||
THEN NULL
|
||||
ELSE BTRIM("usage".provider_name)
|
||||
END
|
||||
)"#;
|
||||
const USAGE_RESOLVED_PROVIDER_DISPLAY_NAME_SQL: &str = r#"COALESCE(
|
||||
provider_by_id.name,
|
||||
@@ -1558,9 +1586,9 @@ fn usage_audit_aggregation_sql_fragments(
|
||||
filtered_extra_where: "",
|
||||
group_key_expr: "provider_group_key",
|
||||
display_name_expr: "provider_display_name",
|
||||
secondary_name_expr: "NULL::varchar",
|
||||
secondary_name_expr: "provider_identity_source",
|
||||
aggregate_display_name_expr: "MAX(display_name)",
|
||||
aggregate_secondary_name_expr: "NULL::varchar",
|
||||
aggregate_secondary_name_expr: "CASE WHEN COUNT(*) FILTER (WHERE secondary_name = 'provider_id') > 0 THEN 'provider_id' WHEN COUNT(*) FILTER (WHERE secondary_name = 'legacy_name') > 0 THEN 'legacy_name' ELSE NULL END",
|
||||
avg_response_time_expr: "AVG(response_time_ms::DOUBLE PRECISION)",
|
||||
success_count_expr: "COALESCE(SUM(success_flag), 0)::BIGINT",
|
||||
},
|
||||
@@ -6687,10 +6715,10 @@ ORDER BY request_count DESC, group_key ASC
|
||||
) -> Result<Vec<StoredUsageAuditAggregation>, DataLayerError> {
|
||||
let fragments = usage_audit_aggregation_sql_fragments(query.group_by);
|
||||
let provider_extra_where =
|
||||
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
|
||||
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
|
||||
|| query.exclude_reserved_provider_labels
|
||||
{
|
||||
USAGE_PROVIDER_IDENTITY_FILTER_SQL
|
||||
} else if query.exclude_reserved_provider_labels {
|
||||
USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL
|
||||
} else {
|
||||
""
|
||||
};
|
||||
@@ -6702,6 +6730,7 @@ WITH filtered_usage AS (
|
||||
"usage".user_id AS user_id,
|
||||
{provider_group_key_expr} AS provider_group_key,
|
||||
{provider_display_name_expr} AS provider_display_name,
|
||||
{provider_identity_source_expr} AS provider_identity_source,
|
||||
COALESCE("usage".api_format, 'unknown') AS api_format_group_key,
|
||||
GREATEST(COALESCE("usage".input_tokens, 0), 0) AS input_tokens,
|
||||
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
|
||||
@@ -6835,6 +6864,7 @@ LIMIT $3
|
||||
provider_identity_join = fragments.provider_identity_join,
|
||||
provider_group_key_expr = fragments.provider_group_key_expr,
|
||||
provider_display_name_expr = fragments.provider_display_name_expr,
|
||||
provider_identity_source_expr = USAGE_PROVIDER_IDENTITY_SOURCE_SQL,
|
||||
group_key_expr = fragments.group_key_expr,
|
||||
display_name_expr = fragments.display_name_expr,
|
||||
secondary_name_expr = fragments.secondary_name_expr,
|
||||
|
||||
@@ -464,16 +464,24 @@ fn usage_sql_aggregate_usage_audits_supports_daily_model_and_provider_aggregates
|
||||
#[test]
|
||||
fn usage_sql_provider_aggregation_excludes_unknown_provider_labels() {
|
||||
let source = include_str!("mod.rs");
|
||||
assert!(source.contains(
|
||||
r#"const USAGE_PROVIDER_IDENTITY_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_id, '')) <> ''"#
|
||||
));
|
||||
assert!(source.contains("const USAGE_PROVIDER_IDENTITY_FILTER_SQL"));
|
||||
assert!(source.contains("const USAGE_PROVIDER_IDENTITY_SOURCE_SQL"));
|
||||
assert!(source.contains(r#"BTRIM(COALESCE("usage".provider_id, '')) <> ''"#));
|
||||
assert!(source.contains(r#"BTRIM(COALESCE("usage".provider_name, '')) <> ''"#));
|
||||
assert!(source.contains("LEFT JOIN providers AS provider_by_id"));
|
||||
assert!(source.contains("provider_by_id.id = BTRIM(\"usage\".provider_id)"));
|
||||
assert!(source.contains("COALESCE(\n provider_by_id.id,\n CASE"));
|
||||
assert!(
|
||||
source.contains("COALESCE(\n provider_by_id.id,\n BTRIM(\"usage\".provider_id)")
|
||||
source.contains(
|
||||
"ELSE BTRIM(\"usage\".provider_id)\n END,\n CASE\n WHEN BTRIM(COALESCE(\"usage\".provider_name, ''))"
|
||||
)
|
||||
);
|
||||
assert!(source.contains("COALESCE(\n provider_by_id.name,"));
|
||||
assert!(!source.contains("provider_by_name.name = BTRIM(\"usage\".provider_name)"));
|
||||
assert!(source.contains("{provider_identity_source_expr} AS provider_identity_source"));
|
||||
assert!(source.contains(r#"secondary_name_expr: "provider_identity_source""#));
|
||||
assert!(source
|
||||
.contains("COUNT(*) FILTER (WHERE secondary_name = 'provider_id') > 0 THEN 'provider_id'"));
|
||||
assert!(source.contains(
|
||||
"if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {\n return self.aggregate_usage_audits_raw(query).await;"
|
||||
));
|
||||
|
||||
@@ -223,10 +223,19 @@ WHERE request_id = ?
|
||||
AND status IN ('streaming', 'success')
|
||||
"#;
|
||||
|
||||
const SQLITE_PROVIDER_NAME_IS_NOT_RESERVED: &str = r#"
|
||||
provider_name IS NOT NULL
|
||||
AND TRIM(provider_name) <> ''
|
||||
AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending')
|
||||
const SQLITE_PROVIDER_IDENTITY_IS_NOT_RESERVED: &str = r#"
|
||||
(
|
||||
(
|
||||
provider_id IS NOT NULL
|
||||
AND TRIM(provider_id) <> ''
|
||||
AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending')
|
||||
)
|
||||
OR (
|
||||
provider_name IS NOT NULL
|
||||
AND TRIM(provider_name) <> ''
|
||||
AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending')
|
||||
)
|
||||
)
|
||||
"#;
|
||||
|
||||
const SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR: &str = r#"
|
||||
@@ -617,6 +626,23 @@ THEN TRIM(provider_id) ELSE TRIM(provider_name) END"
|
||||
}
|
||||
}
|
||||
|
||||
fn sqlite_usage_aggregation_secondary_expr(group_by: UsageAuditAggregationGroupBy) -> &'static str {
|
||||
match group_by {
|
||||
UsageAuditAggregationGroupBy::Provider => {
|
||||
"CASE WHEN SUM(CASE WHEN provider_id IS NOT NULL \
|
||||
AND TRIM(provider_id) <> '' \
|
||||
AND LOWER(TRIM(provider_id)) NOT IN ('unknown', 'unknow', 'pending') \
|
||||
THEN 1 ELSE 0 END) > 0 THEN 'provider_id' \
|
||||
WHEN SUM(CASE WHEN provider_name IS NOT NULL \
|
||||
AND TRIM(provider_name) <> '' \
|
||||
AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') \
|
||||
THEN 1 ELSE 0 END) > 0 THEN 'legacy_name' \
|
||||
ELSE NULL END"
|
||||
}
|
||||
_ => "NULL",
|
||||
}
|
||||
}
|
||||
|
||||
fn sqlite_aggregate_u64(row: &SqliteRow, field: &str) -> Result<u64, DataLayerError> {
|
||||
Ok(row.try_get::<i64, _>(field).map_sql_err()?.max(0) as u64)
|
||||
}
|
||||
@@ -682,7 +708,7 @@ fn decode_sqlite_usage_aggregation_row(
|
||||
Ok(StoredUsageAuditAggregation {
|
||||
group_key: row.try_get::<String, _>("group_key").map_sql_err()?,
|
||||
display_name: row.try_get("display_name").map_sql_err()?,
|
||||
secondary_name: None,
|
||||
secondary_name: row.try_get("secondary_name").map_sql_err()?,
|
||||
request_count: sqlite_aggregate_u64(row, "request_count")?,
|
||||
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
|
||||
output_tokens: sqlite_aggregate_u64(row, "output_tokens")?,
|
||||
@@ -1323,8 +1349,9 @@ FROM "usage"
|
||||
}
|
||||
|
||||
let group_expr = sqlite_usage_aggregation_group_expr(query.group_by);
|
||||
let secondary_expr = sqlite_usage_aggregation_secondary_expr(query.group_by);
|
||||
let display_expr = if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
|
||||
"NULLIF(TRIM(provider_name), '')"
|
||||
"CASE WHEN provider_name IS NOT NULL AND TRIM(provider_name) <> '' AND LOWER(TRIM(provider_name)) NOT IN ('unknown', 'unknow', 'pending') THEN TRIM(provider_name) ELSE NULL END"
|
||||
} else {
|
||||
"NULL"
|
||||
};
|
||||
@@ -1348,6 +1375,7 @@ FROM "usage"
|
||||
SELECT
|
||||
{group_expr} AS group_key,
|
||||
{display_expr} AS display_name,
|
||||
{secondary_expr} AS secondary_name,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens,
|
||||
COALESCE(SUM(MAX(COALESCE(output_tokens, 0), 0)), 0) AS output_tokens,
|
||||
@@ -1368,6 +1396,7 @@ FROM "usage"
|
||||
"#,
|
||||
effective_input_expr = SQLITE_USAGE_EFFECTIVE_INPUT_TOKENS_EXPR,
|
||||
total_input_context_expr = SQLITE_USAGE_TOTAL_INPUT_CONTEXT_EXPR,
|
||||
secondary_expr = secondary_expr,
|
||||
cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR
|
||||
));
|
||||
let mut has_where = false;
|
||||
@@ -1383,7 +1412,7 @@ FROM "usage"
|
||||
builder.push("status NOT IN ('pending', 'streaming')");
|
||||
if query.exclude_reserved_provider_labels {
|
||||
push_sqlite_usage_where(&mut builder, &mut has_where);
|
||||
builder.push(SQLITE_PROVIDER_NAME_IS_NOT_RESERVED);
|
||||
builder.push(SQLITE_PROVIDER_IDENTITY_IS_NOT_RESERVED);
|
||||
}
|
||||
if matches!(query.group_by, UsageAuditAggregationGroupBy::User) {
|
||||
push_sqlite_usage_where(&mut builder, &mut has_where);
|
||||
|
||||
@@ -380,6 +380,81 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_quota_exhaustion_snapshot_expires_after_reset_at() {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system time should be after unix epoch")
|
||||
.as_secs();
|
||||
|
||||
let mut expired = sample_key(None);
|
||||
expired.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"updated_at": now.saturating_sub(600),
|
||||
"windows": [{
|
||||
"code": "5h",
|
||||
"used_ratio": 1.0,
|
||||
"reset_at": now.saturating_sub(60),
|
||||
"is_exhausted": true
|
||||
}]
|
||||
}
|
||||
}));
|
||||
assert!(!provider_pool_key_account_quota_exhausted(
|
||||
&expired, "codex"
|
||||
));
|
||||
|
||||
let mut active = sample_key(None);
|
||||
active.status_snapshot = Some(json!({
|
||||
"quota": {
|
||||
"version": 2,
|
||||
"provider_type": "codex",
|
||||
"code": "exhausted",
|
||||
"exhausted": true,
|
||||
"updated_at": now,
|
||||
"windows": [{
|
||||
"code": "5h",
|
||||
"used_ratio": 1.0,
|
||||
"reset_at": now.saturating_add(3600),
|
||||
"is_exhausted": true
|
||||
}]
|
||||
}
|
||||
}));
|
||||
assert!(provider_pool_key_account_quota_exhausted(&active, "codex"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_quota_exhaustion_metadata_expires_after_reset_at() {
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.expect("system time should be after unix epoch")
|
||||
.as_secs();
|
||||
|
||||
assert!(!provider_pool_key_account_quota_exhausted(
|
||||
&sample_key(Some(json!({
|
||||
"codex": {
|
||||
"updated_at": now.saturating_sub(600),
|
||||
"primary_used_percent": 100.0,
|
||||
"primary_reset_at": now.saturating_sub(60)
|
||||
}
|
||||
}))),
|
||||
"codex",
|
||||
));
|
||||
assert!(provider_pool_key_account_quota_exhausted(
|
||||
&sample_key(Some(json!({
|
||||
"codex": {
|
||||
"updated_at": now,
|
||||
"primary_used_percent": 100.0,
|
||||
"primary_reset_at": now.saturating_add(3600)
|
||||
}
|
||||
}))),
|
||||
"codex",
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn grok_quota_tier_boundaries_match_pool_modes() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -10,8 +10,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
@@ -255,6 +256,19 @@ fn infer_chatgpt_web_image_quota_limit(
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_current_unix_secs().is_some_and(|now| {
|
||||
let mut image_quota = Map::new();
|
||||
if let Some(value) = bucket.get("image_quota_reset_at") {
|
||||
image_quota.insert("reset_at".to_string(), value.clone());
|
||||
}
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
&image_quota,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now,
|
||||
)
|
||||
}) {
|
||||
return false;
|
||||
}
|
||||
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -10,8 +10,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
@@ -121,6 +122,37 @@ pub fn build_codex_pool_quota_request(
|
||||
})
|
||||
}
|
||||
|
||||
fn codex_window_reset_elapsed(bucket: &Map<String, Value>, prefix: &str) -> bool {
|
||||
let Some(now_unix_secs) = provider_pool_current_unix_secs() else {
|
||||
return false;
|
||||
};
|
||||
let mut window = Map::new();
|
||||
for (target, source) in [
|
||||
("reset_at", format!("{prefix}_reset_at")),
|
||||
("next_reset_at", format!("{prefix}_next_reset_at")),
|
||||
("reset_seconds", format!("{prefix}_reset_seconds")),
|
||||
(
|
||||
"reset_after_seconds",
|
||||
format!("{prefix}_reset_after_seconds"),
|
||||
),
|
||||
] {
|
||||
if let Some(value) = bucket.get(source.as_str()) {
|
||||
window.insert(target.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
&window,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now_unix_secs,
|
||||
)
|
||||
}
|
||||
|
||||
fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str) -> bool {
|
||||
let used_percent_key = format!("{prefix}_used_percent");
|
||||
provider_pool_json_f64(bucket.get(used_percent_key.as_str()))
|
||||
.is_some_and(|value| value >= 100.0 && !codex_window_reset_elapsed(bucket, prefix))
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
|
||||
return false;
|
||||
@@ -130,7 +162,6 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
|
||||
return true;
|
||||
}
|
||||
provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0)
|
||||
|| provider_pool_json_f64(bucket.get("secondary_used_percent"))
|
||||
.is_some_and(|value| value >= 100.0)
|
||||
codex_window_used_percent_exhausted(bucket, "primary")
|
||||
|| codex_window_used_percent_exhausted(bucket, "secondary")
|
||||
}
|
||||
|
||||
@@ -7,8 +7,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
|
||||
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
|
||||
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
|
||||
@@ -191,6 +192,8 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
|
||||
let mut model_count = 0usize;
|
||||
let mut exhausted_count = 0usize;
|
||||
let now_unix_secs = provider_pool_current_unix_secs();
|
||||
let bucket_updated_at = provider_pool_timestamp_unix_secs(bucket.get("updated_at"));
|
||||
for (model_key, item) in models.iter() {
|
||||
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
|
||||
continue;
|
||||
@@ -207,12 +210,14 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
continue;
|
||||
}
|
||||
model_count += 1;
|
||||
if provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|
||||
let quota_exhausted = provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|
||||
|| provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
|
||||
|| provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
|
||||
|| provider_pool_json_f64(item.get("remaining_fraction"))
|
||||
.is_some_and(|value| value <= 0.0)
|
||||
{
|
||||
.is_some_and(|value| value <= 0.0);
|
||||
let reset_elapsed = now_unix_secs
|
||||
.is_some_and(|now| provider_pool_reset_deadline_elapsed(item, bucket_updated_at, now));
|
||||
if quota_exhausted && !reset_elapsed {
|
||||
exhausted_count += 1;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -11,8 +11,9 @@ use crate::provider::{
|
||||
ProviderPoolMemberInput,
|
||||
};
|
||||
use crate::quota::{
|
||||
provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision,
|
||||
provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
|
||||
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
|
||||
provider_pool_timestamp_unix_secs,
|
||||
};
|
||||
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
|
||||
|
||||
@@ -151,6 +152,15 @@ fn normalize_kiro_version(value: &str) -> &str {
|
||||
}
|
||||
|
||||
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
|
||||
if provider_pool_current_unix_secs().is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(
|
||||
bucket,
|
||||
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
|
||||
now,
|
||||
)
|
||||
}) {
|
||||
return false;
|
||||
}
|
||||
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
@@ -135,7 +137,7 @@ pub(crate) fn provider_pool_json_f64(value: Option<&Value>) -> Option<f64> {
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
||||
pub(crate) fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
||||
let mut timestamp = provider_pool_json_f64(value)?;
|
||||
if timestamp <= 0.0 {
|
||||
return None;
|
||||
@@ -146,6 +148,49 @@ fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
|
||||
Some(timestamp as u64)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_pool_current_unix_secs() -> Option<u64> {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
}
|
||||
|
||||
fn provider_pool_reset_deadline_unix_secs(
|
||||
item: &Map<String, Value>,
|
||||
fallback_observed_at: Option<u64>,
|
||||
) -> Option<u64> {
|
||||
provider_pool_timestamp_unix_secs(item.get("reset_at"))
|
||||
.or_else(|| provider_pool_timestamp_unix_secs(item.get("next_reset_at")))
|
||||
.or_else(|| {
|
||||
let reset_seconds = provider_pool_json_f64(item.get("reset_seconds"))
|
||||
.or_else(|| provider_pool_json_f64(item.get("reset_after_seconds")))?;
|
||||
if reset_seconds < 0.0 {
|
||||
return None;
|
||||
}
|
||||
let base = provider_pool_timestamp_unix_secs(item.get("observed_at"))
|
||||
.or_else(|| provider_pool_timestamp_unix_secs(item.get("updated_at")))
|
||||
.or(fallback_observed_at)?;
|
||||
Some(base.saturating_add(reset_seconds.ceil() as u64))
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn provider_pool_reset_deadline_elapsed(
|
||||
item: &Map<String, Value>,
|
||||
fallback_observed_at: Option<u64>,
|
||||
now_unix_secs: u64,
|
||||
) -> bool {
|
||||
provider_pool_reset_deadline_unix_secs(item, fallback_observed_at)
|
||||
.is_some_and(|reset_at| reset_at <= now_unix_secs)
|
||||
}
|
||||
|
||||
fn provider_pool_quota_window_is_exhausted(window: &Map<String, Value>) -> bool {
|
||||
provider_pool_json_bool(window.get("is_exhausted"))
|
||||
.or_else(|| {
|
||||
provider_pool_json_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn provider_pool_quota_snapshot_matches_provider(
|
||||
quota_snapshot: &Map<String, Value>,
|
||||
provider_type: &str,
|
||||
@@ -203,19 +248,45 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
|
||||
}
|
||||
let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
|
||||
if exhausted {
|
||||
let windows_max_ratio = quota_snapshot
|
||||
let now_unix_secs = provider_pool_current_unix_secs();
|
||||
let snapshot_observed_at =
|
||||
provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at"))
|
||||
.or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at")));
|
||||
|
||||
if let Some(windows) = quota_snapshot
|
||||
.get("windows")
|
||||
.and_then(Value::as_array)
|
||||
.filter(|w| !w.is_empty())
|
||||
.and_then(|windows| {
|
||||
windows
|
||||
.iter()
|
||||
.filter_map(Value::as_object)
|
||||
.filter_map(|w| w.get("used_ratio"))
|
||||
.filter_map(Value::as_f64)
|
||||
.max_by(f64::total_cmp)
|
||||
});
|
||||
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
|
||||
.filter(|windows| !windows.is_empty())
|
||||
{
|
||||
let mut saw_exhausted_window = false;
|
||||
let mut saw_active_exhausted_window = false;
|
||||
let mut windows_max_ratio = None::<f64>;
|
||||
|
||||
for window in windows.iter().filter_map(Value::as_object) {
|
||||
if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) {
|
||||
windows_max_ratio =
|
||||
Some(windows_max_ratio.map_or(ratio, |current| current.max(ratio)));
|
||||
}
|
||||
if provider_pool_quota_window_is_exhausted(window) {
|
||||
saw_exhausted_window = true;
|
||||
let reset_elapsed = now_unix_secs.is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
|
||||
});
|
||||
if !reset_elapsed {
|
||||
saw_active_exhausted_window = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if saw_exhausted_window {
|
||||
return Some(saw_active_exhausted_window);
|
||||
}
|
||||
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
|
||||
return Some(false);
|
||||
}
|
||||
} else if now_unix_secs.is_some_and(|now| {
|
||||
provider_pool_reset_deadline_elapsed(quota_snapshot, snapshot_observed_at, now)
|
||||
}) {
|
||||
return Some(false);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -488,6 +488,7 @@ pub fn build_local_request_candidate_status_record(
|
||||
demoted_by: metadata.demoted_by.clone(),
|
||||
routing_trace: metadata.routing_trace.clone(),
|
||||
});
|
||||
let extra_data = mark_request_candidate_stream_completed_if_success(status, extra_data);
|
||||
let created_at_unix_ms = started_at_unix_ms.or(finished_at_unix_ms);
|
||||
|
||||
Some(UpsertRequestCandidateRecord {
|
||||
@@ -568,7 +569,7 @@ pub fn build_report_request_candidate_status_record(
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests: None,
|
||||
extra_data: slot.extra_data,
|
||||
extra_data: mark_request_candidate_stream_completed_if_success(status, slot.extra_data),
|
||||
required_capabilities: None,
|
||||
created_at_unix_ms: Some(created_at_unix_ms),
|
||||
started_at_unix_ms,
|
||||
@@ -829,6 +830,23 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
|
||||
(!extra_data.is_empty()).then_some(Value::Object(extra_data))
|
||||
}
|
||||
|
||||
fn mark_request_candidate_stream_completed_if_success(
|
||||
status: RequestCandidateStatus,
|
||||
extra_data: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
if status != RequestCandidateStatus::Success {
|
||||
return extra_data;
|
||||
}
|
||||
|
||||
let mut object = match extra_data {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(other) => return Some(other),
|
||||
None => Map::new(),
|
||||
};
|
||||
object.insert("stream_completed".to_string(), Value::Bool(true));
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
fn merge_request_candidate_extra_data(
|
||||
existing: Option<Value>,
|
||||
overlay: Option<Value>,
|
||||
@@ -1295,6 +1313,53 @@ mod tests {
|
||||
assert_eq!(record.finished_at_unix_ms, Some(123));
|
||||
assert_eq!(record.created_at_unix_ms, Some(123));
|
||||
assert_eq!(record.status, RequestCandidateStatus::Success);
|
||||
assert_eq!(
|
||||
record
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("stream_completed")),
|
||||
Some(&json!(true))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_success_status_marks_stream_completed_for_pending_cleanup_recovery() {
|
||||
let mut plan = sample_plan();
|
||||
plan.candidate_id = Some("cand-1".to_string());
|
||||
let report_context = json!({
|
||||
"request_id": "req-1",
|
||||
"candidate_id": "cand-1",
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"user_id": "user-1",
|
||||
"api_key_id": "api-key-1",
|
||||
"client_api_format": "openai:responses",
|
||||
"provider_api_format": "openai:responses",
|
||||
});
|
||||
|
||||
let record =
|
||||
build_local_request_candidate_status_record(LocalRequestCandidateStatusRecordInput {
|
||||
plan: &plan,
|
||||
report_context: Some(&report_context),
|
||||
status_update: SchedulerRequestCandidateStatusUpdate {
|
||||
status: RequestCandidateStatus::Success,
|
||||
status_code: Some(200),
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: Some(25),
|
||||
started_at_unix_ms: Some(1_000),
|
||||
finished_at_unix_ms: Some(1_025),
|
||||
},
|
||||
})
|
||||
.expect("success status record should build");
|
||||
|
||||
assert_eq!(
|
||||
record
|
||||
.extra_data
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("stream_completed")),
|
||||
Some(&json!(true))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
import apiClient from './client'
|
||||
|
||||
const MODULE_MANAGEMENT_ORDER_CONFIG_KEY = 'module_management.extension_order'
|
||||
|
||||
export interface ModuleStatus {
|
||||
name: string
|
||||
available: boolean
|
||||
@@ -75,6 +77,20 @@ const CHAT_PII_REDACTION_DEFAULT_CONFIG: ChatPiiRedactionConfig = {
|
||||
placeholder_prefix: 'AETHER',
|
||||
}
|
||||
|
||||
export function normalizeModuleManagementOrder(value: unknown): string[] {
|
||||
if (!Array.isArray(value)) return []
|
||||
const seen = new Set<string>()
|
||||
const order: string[] = []
|
||||
for (const item of value) {
|
||||
if (typeof item !== 'string') continue
|
||||
const name = item.trim()
|
||||
if (!name || seen.has(name)) continue
|
||||
seen.add(name)
|
||||
order.push(name)
|
||||
}
|
||||
return order
|
||||
}
|
||||
|
||||
function cloneDefaultChatPiiRedactionRules(): ChatPiiRedactionRule[] {
|
||||
return CHAT_PII_REDACTION_DEFAULT_RULES.map(rule => ({ ...rule }))
|
||||
}
|
||||
@@ -189,6 +205,31 @@ export const modulesApi = {
|
||||
return response.data
|
||||
},
|
||||
|
||||
async getModuleManagementOrder(): Promise<string[]> {
|
||||
try {
|
||||
const response = await apiClient.get<{ key: string; value: unknown }>(
|
||||
`/api/admin/system/configs/${MODULE_MANAGEMENT_ORDER_CONFIG_KEY}`
|
||||
)
|
||||
return normalizeModuleManagementOrder(response.data.value)
|
||||
} catch (err) {
|
||||
const status = (err as { response?: { status?: number } }).response?.status
|
||||
if (status === 404) return []
|
||||
throw err
|
||||
}
|
||||
},
|
||||
|
||||
async updateModuleManagementOrder(order: string[]): Promise<string[]> {
|
||||
const normalized = normalizeModuleManagementOrder(order)
|
||||
const response = await apiClient.put<{ key: string; value: unknown }>(
|
||||
`/api/admin/system/configs/${MODULE_MANAGEMENT_ORDER_CONFIG_KEY}`,
|
||||
{
|
||||
value: normalized,
|
||||
description: '模块管理扩展模块展示顺序',
|
||||
},
|
||||
)
|
||||
return normalizeModuleManagementOrder(response.data.value)
|
||||
},
|
||||
|
||||
async getChatPiiRedactionConfig(): Promise<ChatPiiRedactionConfig> {
|
||||
const [enabled, rules, cacheTtlSeconds, placeholderPrefix] = await Promise.all([
|
||||
getSystemConfigValue(CHAT_PII_REDACTION_CONFIG_KEYS.enabled),
|
||||
|
||||
@@ -71,7 +71,9 @@ export interface UsageByUser {
|
||||
}
|
||||
|
||||
export interface UsageByProvider {
|
||||
provider_id: string
|
||||
provider_id?: string | null
|
||||
provider_key?: string
|
||||
provider_identity_source?: 'provider_id' | 'legacy_name'
|
||||
provider: string
|
||||
request_count: number
|
||||
total_tokens: number
|
||||
|
||||
@@ -46,7 +46,7 @@
|
||||
</TableRow>
|
||||
<TableRow
|
||||
v-for="provider in data"
|
||||
:key="provider.provider"
|
||||
:key="provider.providerKey ?? provider.providerId ?? provider.provider"
|
||||
>
|
||||
<TableCell class="font-medium py-2 px-2">
|
||||
{{ provider.provider }}
|
||||
|
||||
@@ -129,6 +129,9 @@ describe('useUsageData', () => {
|
||||
])
|
||||
getUsageByProviderMock.mockResolvedValueOnce([
|
||||
{
|
||||
provider_id: 'provider-openai',
|
||||
provider_key: 'provider-openai',
|
||||
provider_identity_source: 'provider_id',
|
||||
provider: 'OpenAI',
|
||||
request_count: 3,
|
||||
total_tokens: 300,
|
||||
@@ -159,6 +162,11 @@ describe('useUsageData', () => {
|
||||
})
|
||||
expect(modelStats.value).toHaveLength(1)
|
||||
expect(providerStats.value).toHaveLength(1)
|
||||
expect(providerStats.value[0]).toMatchObject({
|
||||
providerId: 'provider-openai',
|
||||
providerKey: 'provider-openai',
|
||||
providerIdentitySource: 'provider_id',
|
||||
})
|
||||
expect(apiFormatStats.value).toHaveLength(1)
|
||||
expect(availableModels.value).toEqual(['gpt-5'])
|
||||
expect(availableProviders.value).toEqual(['OpenAI'])
|
||||
|
||||
@@ -173,6 +173,9 @@ export function useUsageData(options: UseUsageDataOptions) {
|
||||
|
||||
const visibleProviderData = providerData.filter(item => isUsageProviderVisible(item.provider))
|
||||
providerStats.value = visibleProviderData.map(item => ({
|
||||
providerId: item.provider_id,
|
||||
providerKey: item.provider_key,
|
||||
providerIdentitySource: item.provider_identity_source,
|
||||
provider: item.provider,
|
||||
requests: item.request_count,
|
||||
totalTokens: item.total_tokens || 0,
|
||||
|
||||
@@ -41,6 +41,9 @@ export interface EnhancedModelStatsItem extends ModelStatsItem {
|
||||
|
||||
// 提供商统计
|
||||
export interface ProviderStatsItem {
|
||||
providerId?: string | null
|
||||
providerKey?: string
|
||||
providerIdentitySource?: 'provider_id' | 'legacy_name'
|
||||
provider: string
|
||||
requests: number
|
||||
totalTokens: number
|
||||
|
||||
@@ -727,7 +727,7 @@ const navigation = computed(() => {
|
||||
items: [
|
||||
{ name: '钱包中心', href: '/dashboard/wallet', icon: Wallet },
|
||||
{ name: '套餐中心', href: '/dashboard/billing', icon: Package },
|
||||
{ name: '我的邀请', href: '/dashboard/referral', icon: Gift },
|
||||
...(moduleStore.isActive('referral') ? [{ name: '我的邀请', href: '/dashboard/referral', icon: Gift }] : []),
|
||||
{ name: '使用统计', href: '/dashboard/usage', icon: BarChart3 },
|
||||
]
|
||||
}
|
||||
@@ -749,19 +749,21 @@ const navigation = computed(() => {
|
||||
Puzzle,
|
||||
Server,
|
||||
SlidersHorizontal,
|
||||
CreditCard,
|
||||
Gift,
|
||||
}
|
||||
|
||||
// 添加模块菜单项(按 admin_menu_order 排序,只显示已激活的)
|
||||
const moduleMenuItems = Object.values(moduleStore.modules)
|
||||
.filter(m => m.active && m.admin_route && m.admin_menu_group === 'system')
|
||||
.sort((a, b) => a.admin_menu_order - b.admin_menu_order)
|
||||
.map(m => ({
|
||||
name: m.display_name,
|
||||
href: m.admin_route ?? '',
|
||||
icon: iconMap[m.admin_menu_icon || ''] || Puzzle
|
||||
}))
|
||||
const activeModuleItems = (group: string) =>
|
||||
Object.values(moduleStore.modules)
|
||||
.filter(m => m.active && m.admin_route && m.admin_menu_group === group)
|
||||
.sort((a, b) => a.admin_menu_order - b.admin_menu_order)
|
||||
.map(m => ({
|
||||
name: m.display_name,
|
||||
href: m.admin_route ?? '',
|
||||
icon: iconMap[m.admin_menu_icon || ''] || Puzzle
|
||||
}))
|
||||
|
||||
systemItems.push(...moduleMenuItems)
|
||||
systemItems.push(...activeModuleItems('system'))
|
||||
|
||||
// 模块管理和系统设置放在最后
|
||||
systemItems.push({ name: '模块管理', href: '/admin/modules', icon: Puzzle })
|
||||
@@ -788,9 +790,8 @@ const navigation = computed(() => {
|
||||
{ name: '号池管理', href: '/admin/pool', icon: Database },
|
||||
{ name: '独立密钥', href: '/admin/keys', icon: Key },
|
||||
{ name: '钱包管理', href: '/admin/wallets', icon: Wallet },
|
||||
{ name: '支付配置', href: '/admin/payment-gateways', icon: CreditCard },
|
||||
{ name: '套餐管理', href: '/admin/billing-plans', icon: Package },
|
||||
{ name: '邀请返利', href: '/admin/referrals', icon: Gift },
|
||||
...activeModuleItems('management'),
|
||||
{ name: '异步任务', href: '/admin/async-tasks', icon: Zap },
|
||||
{ name: '使用记录', href: '/admin/usage', icon: BarChart3 },
|
||||
]
|
||||
|
||||
@@ -9,6 +9,7 @@ import type { User as AdminUser } from '@/api/users'
|
||||
import type { AdminApiKeysResponse } from '@/api/admin'
|
||||
import type { Profile, UsageResponse } from '@/api/me'
|
||||
import type { ProviderWithEndpointsSummary, GlobalModelResponse } from '@/api/endpoints/types'
|
||||
import type { ModuleStatus } from '@/api/modules'
|
||||
|
||||
// ========== 用户数据 ==========
|
||||
|
||||
@@ -902,7 +903,7 @@ export const MOCK_USAGE_RESPONSE: UsageResponse = {
|
||||
|
||||
// ========== 系统配置 ==========
|
||||
|
||||
export const MOCK_SYSTEM_CONFIGS = [
|
||||
export const MOCK_SYSTEM_CONFIGS: Array<{ key: string; value: unknown; description?: string }> = [
|
||||
{ key: 'rate_limit_enabled', value: true, description: '是否启用速率限制' },
|
||||
{ key: 'default_rate_limit', value: 60, description: '默认速率限制(请求/分钟)' },
|
||||
{ key: 'cache_enabled', value: true, description: '是否启用缓存' },
|
||||
@@ -914,6 +915,161 @@ export const MOCK_SYSTEM_CONFIGS = [
|
||||
{ key: 'proxy_node_metrics_cleanup_batch_size', value: 5000, description: '代理节点指标每批次清理条数' }
|
||||
]
|
||||
|
||||
const MOCK_MODULE_DEFINITIONS: Array<Omit<ModuleStatus, 'active' | 'health'> & { health?: ModuleStatus['health'] }> = [
|
||||
{
|
||||
name: 'management_tokens',
|
||||
display_name: '访问令牌',
|
||||
description: '管理 API 访问令牌,支持细粒度权限控制和 IP 限制',
|
||||
category: 'security',
|
||||
available: true,
|
||||
enabled: true,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/management-tokens',
|
||||
admin_menu_icon: null,
|
||||
admin_menu_group: null,
|
||||
admin_menu_order: 0,
|
||||
},
|
||||
{
|
||||
name: 'ldap',
|
||||
display_name: 'LDAP 认证',
|
||||
description: '支持通过 LDAP/Active Directory 进行用户认证',
|
||||
category: 'auth',
|
||||
available: true,
|
||||
enabled: false,
|
||||
config_validated: false,
|
||||
config_error: '请先配置 LDAP 连接信息',
|
||||
admin_route: '/admin/ldap',
|
||||
admin_menu_icon: 'Users',
|
||||
admin_menu_group: 'system',
|
||||
admin_menu_order: 50,
|
||||
},
|
||||
{
|
||||
name: 'oauth',
|
||||
display_name: 'OAuth 登录',
|
||||
description: '支持通过第三方 OAuth Provider 登录/绑定账号',
|
||||
category: 'auth',
|
||||
available: true,
|
||||
enabled: true,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/oauth',
|
||||
admin_menu_icon: 'Key',
|
||||
admin_menu_group: null,
|
||||
admin_menu_order: 55,
|
||||
},
|
||||
{
|
||||
name: 'notification_email',
|
||||
display_name: '异常通知',
|
||||
description: '为 5xx 异常发送邮件通知,可在模块管理中启用或禁用',
|
||||
category: 'integration',
|
||||
available: true,
|
||||
enabled: false,
|
||||
config_validated: false,
|
||||
config_error: '请先完成邮件配置(SMTP)',
|
||||
admin_route: null,
|
||||
admin_menu_icon: 'Mail',
|
||||
admin_menu_group: 'system',
|
||||
admin_menu_order: 58,
|
||||
},
|
||||
{
|
||||
name: 'chat_pii_redaction',
|
||||
display_name: '敏感信息保护',
|
||||
description: '发送给供应商前将聊天消息中的敏感信息替换为占位符,返回客户端前自动还原。',
|
||||
category: 'security',
|
||||
available: true,
|
||||
enabled: false,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/modules/chat-pii-redaction',
|
||||
admin_menu_icon: 'ShieldCheck',
|
||||
admin_menu_group: 'system',
|
||||
admin_menu_order: 59,
|
||||
},
|
||||
{
|
||||
name: 'model_directives',
|
||||
display_name: '模型后缀参数',
|
||||
description: '允许通过模型名后缀覆盖推理参数',
|
||||
category: 'integration',
|
||||
available: true,
|
||||
enabled: true,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/model-directives',
|
||||
admin_menu_icon: 'SlidersHorizontal',
|
||||
admin_menu_group: null,
|
||||
admin_menu_order: 59,
|
||||
},
|
||||
{
|
||||
name: 'gemini_files',
|
||||
display_name: '文件缓存',
|
||||
description: '管理 Gemini Files API 上传的文件,支持文件上传、查看和删除',
|
||||
category: 'integration',
|
||||
available: true,
|
||||
enabled: false,
|
||||
config_validated: false,
|
||||
config_error: '至少启用一个具有「Gemini 文件 API」能力的 Key',
|
||||
admin_route: '/admin/gemini-files',
|
||||
admin_menu_icon: 'FileUp',
|
||||
admin_menu_group: 'system',
|
||||
admin_menu_order: 60,
|
||||
health: 'degraded',
|
||||
},
|
||||
{
|
||||
name: 'proxy_nodes',
|
||||
display_name: '代理节点',
|
||||
description: '添加Http/Socket代理节点, 或使用Aether-Proxy自动连接代理节点.',
|
||||
category: 'integration',
|
||||
available: true,
|
||||
enabled: true,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/proxy-nodes',
|
||||
admin_menu_icon: 'Server',
|
||||
admin_menu_group: 'system',
|
||||
admin_menu_order: 60,
|
||||
},
|
||||
{
|
||||
name: 'payment_gateways',
|
||||
display_name: '支付配置',
|
||||
description: '配置易支付、支付宝官方、微信支付官方和 Stripe 等支付网关',
|
||||
category: 'integration',
|
||||
available: true,
|
||||
enabled: false,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/payment-gateways',
|
||||
admin_menu_icon: 'CreditCard',
|
||||
admin_menu_group: null,
|
||||
admin_menu_order: 70,
|
||||
},
|
||||
{
|
||||
name: 'referral',
|
||||
display_name: '邀请返利',
|
||||
description: '管理用户邀请关系与返利记录,支持比例返利和人头返利',
|
||||
category: 'integration',
|
||||
available: true,
|
||||
enabled: false,
|
||||
config_validated: true,
|
||||
config_error: null,
|
||||
admin_route: '/admin/referrals',
|
||||
admin_menu_icon: 'Gift',
|
||||
admin_menu_group: 'management',
|
||||
admin_menu_order: 75,
|
||||
},
|
||||
]
|
||||
|
||||
export const MOCK_MODULE_STATUSES: Record<string, ModuleStatus> = Object.fromEntries(
|
||||
MOCK_MODULE_DEFINITIONS.map(module => [
|
||||
module.name,
|
||||
{
|
||||
...module,
|
||||
active: module.available && module.enabled && module.config_validated,
|
||||
health: module.health ?? 'healthy',
|
||||
},
|
||||
])
|
||||
) as Record<string, ModuleStatus>
|
||||
|
||||
// ========== API 格式 ==========
|
||||
|
||||
export const MOCK_API_FORMATS = {
|
||||
|
||||
@@ -22,6 +22,7 @@ import {
|
||||
MOCK_PROVIDERS,
|
||||
MOCK_GLOBAL_MODELS,
|
||||
MOCK_SYSTEM_CONFIGS,
|
||||
MOCK_MODULE_STATUSES,
|
||||
MOCK_API_FORMATS
|
||||
} from './data'
|
||||
|
||||
@@ -1290,6 +1291,13 @@ const mockHandlers: Record<string, (config: AxiosRequestConfig) => Promise<Axios
|
||||
return createMockResponse({ requests: [] })
|
||||
},
|
||||
|
||||
// ========== Admin: Modules ==========
|
||||
'GET /api/admin/modules/status': async () => {
|
||||
await delay()
|
||||
requireAdmin()
|
||||
return createMockResponse(MOCK_MODULE_STATUSES)
|
||||
},
|
||||
|
||||
// ========== Admin: System ==========
|
||||
'GET /api/admin/system/configs': async () => {
|
||||
await delay()
|
||||
@@ -1788,6 +1796,71 @@ function generateMockModelsForProvider(providerId: string) {
|
||||
|
||||
// ========== 注册动态路由 ==========
|
||||
|
||||
// 系统配置详情
|
||||
registerDynamicRoute('GET', '/api/admin/system/configs/:configKey', async (_config, params) => {
|
||||
await delay()
|
||||
requireAdmin()
|
||||
const key = decodeURIComponent(params.configKey)
|
||||
const entry = MOCK_SYSTEM_CONFIGS.find(item => item.key === key)
|
||||
if (!entry) {
|
||||
throw { response: createMockResponse({ detail: `配置项 '${key}' 不存在` }, 404) }
|
||||
}
|
||||
return createMockResponse({ key: entry.key, value: entry.value, description: entry.description })
|
||||
})
|
||||
|
||||
// 系统配置更新
|
||||
registerDynamicRoute('PUT', '/api/admin/system/configs/:configKey', async (config, params) => {
|
||||
await delay()
|
||||
requireAdmin()
|
||||
const key = decodeURIComponent(params.configKey)
|
||||
const body = JSON.parse(config.data || '{}') as { value?: unknown; description?: string }
|
||||
const index = MOCK_SYSTEM_CONFIGS.findIndex(item => item.key === key)
|
||||
const entry = {
|
||||
key,
|
||||
value: body.value ?? null,
|
||||
description: body.description,
|
||||
}
|
||||
if (index === -1) {
|
||||
MOCK_SYSTEM_CONFIGS.push(entry)
|
||||
} else {
|
||||
MOCK_SYSTEM_CONFIGS[index] = {
|
||||
...MOCK_SYSTEM_CONFIGS[index],
|
||||
...entry,
|
||||
}
|
||||
}
|
||||
return createMockResponse(entry)
|
||||
})
|
||||
|
||||
// 模块状态详情
|
||||
registerDynamicRoute('GET', '/api/admin/modules/status/:moduleName', async (_config, params) => {
|
||||
await delay()
|
||||
requireAdmin()
|
||||
const moduleStatus = MOCK_MODULE_STATUSES[params.moduleName]
|
||||
if (!moduleStatus) {
|
||||
throw { response: createMockResponse({ detail: '模块不存在' }, 404) }
|
||||
}
|
||||
return createMockResponse(moduleStatus)
|
||||
})
|
||||
|
||||
// 模块启用状态更新
|
||||
registerDynamicRoute('PUT', '/api/admin/modules/status/:moduleName/enabled', async (config, params) => {
|
||||
await delay()
|
||||
requireAdmin()
|
||||
const moduleStatus = MOCK_MODULE_STATUSES[params.moduleName]
|
||||
if (!moduleStatus) {
|
||||
throw { response: createMockResponse({ detail: '模块不存在' }, 404) }
|
||||
}
|
||||
const body = JSON.parse(config.data || '{}') as { enabled?: boolean }
|
||||
const enabled = body.enabled === true
|
||||
const updated = {
|
||||
...moduleStatus,
|
||||
enabled,
|
||||
active: moduleStatus.available && enabled && moduleStatus.config_validated,
|
||||
}
|
||||
MOCK_MODULE_STATUSES[params.moduleName] = updated
|
||||
return createMockResponse(updated)
|
||||
})
|
||||
|
||||
// Provider 详情
|
||||
registerDynamicRoute('GET', '/api/admin/providers/:providerId/summary', async (_config, params) => {
|
||||
await delay()
|
||||
|
||||
@@ -79,9 +79,25 @@
|
||||
</div>
|
||||
|
||||
<!-- 扩展模块 -->
|
||||
<h3 class="text-sm font-semibold text-muted-foreground uppercase tracking-wider mb-4">
|
||||
扩展模块
|
||||
</h3>
|
||||
<div class="mb-4 flex items-center justify-between gap-3">
|
||||
<h3 class="text-sm font-semibold text-muted-foreground uppercase tracking-wider">
|
||||
扩展模块
|
||||
</h3>
|
||||
<Button
|
||||
v-if="hasCustomModuleOrder"
|
||||
variant="outline"
|
||||
size="sm"
|
||||
class="gap-1.5"
|
||||
:disabled="loading || orderSaving"
|
||||
@click="resetModuleOrder"
|
||||
>
|
||||
<RotateCcw
|
||||
class="w-3.5 h-3.5"
|
||||
:class="{ 'animate-spin': orderSaving }"
|
||||
/>
|
||||
恢复默认
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
<!-- 模块卡片网格 -->
|
||||
<div class="grid grid-cols-1 md:grid-cols-2 xl:grid-cols-3 gap-5">
|
||||
@@ -89,11 +105,23 @@
|
||||
v-for="module in filteredModules"
|
||||
:key="module.name"
|
||||
class="group relative border rounded-2xl p-6 transition-all duration-200 hover:shadow-lg"
|
||||
:class="{
|
||||
'bg-muted/40 border-muted': !module.available,
|
||||
'border-primary/40 bg-gradient-to-br from-primary/5 to-primary/10 shadow-sm': module.active,
|
||||
'border-border bg-card hover:border-primary/20': !module.active && module.available
|
||||
}"
|
||||
:class="[
|
||||
{
|
||||
'bg-muted/40 border-muted': !module.available,
|
||||
'border-primary/40 bg-gradient-to-br from-primary/5 to-primary/10 shadow-sm': module.active,
|
||||
'border-border bg-card hover:border-primary/20': !module.active && module.available
|
||||
},
|
||||
draggedModuleName === module.name ? 'opacity-70 ring-2 ring-primary/30' : '',
|
||||
dragOverModuleName === module.name ? 'ring-2 ring-primary/40 border-primary/50' : '',
|
||||
canReorderModules ? 'cursor-grab active:cursor-grabbing' : ''
|
||||
]"
|
||||
:draggable="canReorderModules"
|
||||
:title="orderSaving ? '正在保存排序' : '拖拽卡片调整顺序'"
|
||||
@dragstart="handleModuleDragStart(module.name, $event)"
|
||||
@dragend="handleModuleDragEnd"
|
||||
@dragover.prevent="handleModuleDragOver(module.name)"
|
||||
@dragleave="handleModuleDragLeave(module.name)"
|
||||
@drop.prevent="handleModuleDrop(module.name)"
|
||||
>
|
||||
<!-- 状态指示器 -->
|
||||
<div class="absolute top-5 right-5">
|
||||
@@ -121,7 +149,7 @@
|
||||
class="w-5 h-5"
|
||||
/>
|
||||
</div>
|
||||
<div class="flex-1 min-w-0 pt-1">
|
||||
<div class="flex-1 min-w-0 pt-1 pr-8">
|
||||
<h4 class="font-semibold text-base truncate">
|
||||
{{ module.display_name }}
|
||||
</h4>
|
||||
@@ -214,7 +242,17 @@
|
||||
<script setup lang="ts">
|
||||
import { ref, computed, onMounted } from 'vue'
|
||||
import { useRouter } from 'vue-router'
|
||||
import { RefreshCw, Puzzle, Users, Shield, Gauge, Link, Search, Settings } from 'lucide-vue-next'
|
||||
import {
|
||||
RefreshCw,
|
||||
Puzzle,
|
||||
Users,
|
||||
Shield,
|
||||
Gauge,
|
||||
Link,
|
||||
Search,
|
||||
Settings,
|
||||
RotateCcw,
|
||||
} from 'lucide-vue-next'
|
||||
import Button from '@/components/ui/button.vue'
|
||||
import Switch from '@/components/ui/switch.vue'
|
||||
import Input from '@/components/ui/input.vue'
|
||||
@@ -224,6 +262,7 @@ import { useModuleStore } from '@/stores/modules'
|
||||
import { BUILTIN_TOOLS } from '@/config/builtin-tools'
|
||||
import { log } from '@/utils/logger'
|
||||
import { getErrorMessage } from '@/types/api-error'
|
||||
import { modulesApi, type ModuleStatus } from '@/api/modules'
|
||||
|
||||
const router = useRouter()
|
||||
const { success, error } = useToast()
|
||||
@@ -232,6 +271,10 @@ const moduleStore = useModuleStore()
|
||||
const loading = ref(false)
|
||||
const toggling = ref<Record<string, boolean>>({})
|
||||
const searchQuery = ref('')
|
||||
const moduleOrder = ref<string[]>([])
|
||||
const orderSaving = ref(false)
|
||||
const draggedModuleName = ref<string | null>(null)
|
||||
const dragOverModuleName = ref<string | null>(null)
|
||||
|
||||
// 过滤后的内置工具
|
||||
const filteredBuiltinTools = computed(() => {
|
||||
@@ -262,12 +305,63 @@ function getModuleStatusCopy(module: { name: string; enabled: boolean; active: b
|
||||
return '已开启'
|
||||
}
|
||||
|
||||
// 所有模块列表(按 admin_menu_order 排序)
|
||||
const allModules = computed(() => {
|
||||
return Object.values(moduleStore.modules)
|
||||
.sort((a, b) => a.admin_menu_order - b.admin_menu_order)
|
||||
function compareModuleDefaultOrder(a: ModuleStatus, b: ModuleStatus) {
|
||||
return a.admin_menu_order - b.admin_menu_order ||
|
||||
a.display_name.localeCompare(b.display_name, 'zh-Hans') ||
|
||||
a.name.localeCompare(b.name)
|
||||
}
|
||||
|
||||
function applySavedModuleOrder(modules: ModuleStatus[], order: string[]) {
|
||||
if (order.length === 0) return modules
|
||||
const modulesByName = new Map(modules.map(module => [module.name, module]))
|
||||
const seen = new Set<string>()
|
||||
const ordered: ModuleStatus[] = []
|
||||
|
||||
for (const moduleName of order) {
|
||||
const module = modulesByName.get(moduleName)
|
||||
if (!module || seen.has(moduleName)) continue
|
||||
seen.add(moduleName)
|
||||
ordered.push(module)
|
||||
}
|
||||
|
||||
for (const module of modules) {
|
||||
if (!seen.has(module.name)) {
|
||||
ordered.push(module)
|
||||
}
|
||||
}
|
||||
|
||||
return ordered
|
||||
}
|
||||
|
||||
function normalizeOrderForCurrentModules(order: string[]) {
|
||||
const availableNames = new Set(defaultOrderedModules.value.map(module => module.name))
|
||||
return order.filter(moduleName => availableNames.has(moduleName))
|
||||
}
|
||||
|
||||
function moveNameToTargetIndex(names: string[], draggedName: string, targetName: string) {
|
||||
const fromIndex = names.indexOf(draggedName)
|
||||
const targetIndex = names.indexOf(targetName)
|
||||
if (fromIndex === -1 || targetIndex === -1 || fromIndex === targetIndex) return names
|
||||
|
||||
const next = [...names]
|
||||
const [dragged] = next.splice(fromIndex, 1)
|
||||
next.splice(targetIndex, 0, dragged)
|
||||
return next
|
||||
}
|
||||
|
||||
// 后端默认顺序
|
||||
const defaultOrderedModules = computed(() => {
|
||||
return Object.values(moduleStore.modules).sort(compareModuleDefaultOrder)
|
||||
})
|
||||
|
||||
// 所有模块列表(应用自定义展示顺序)
|
||||
const allModules = computed(() => {
|
||||
return applySavedModuleOrder(defaultOrderedModules.value, moduleOrder.value)
|
||||
})
|
||||
|
||||
const hasCustomModuleOrder = computed(() => moduleOrder.value.length > 0)
|
||||
const canReorderModules = computed(() => !orderSaving.value && allModules.value.length > 1)
|
||||
|
||||
// 过滤后的模块列表
|
||||
const filteredModules = computed(() => {
|
||||
if (!searchQuery.value.trim()) {
|
||||
@@ -286,7 +380,11 @@ const filteredModules = computed(() => {
|
||||
async function fetchModules() {
|
||||
loading.value = true
|
||||
try {
|
||||
await moduleStore.fetchModules()
|
||||
const [, savedOrder] = await Promise.all([
|
||||
moduleStore.fetchModules(),
|
||||
modulesApi.getModuleManagementOrder(),
|
||||
])
|
||||
moduleOrder.value = normalizeOrderForCurrentModules(savedOrder)
|
||||
} catch (err) {
|
||||
error('获取模块列表失败')
|
||||
log.error('获取模块列表失败:', err)
|
||||
@@ -309,6 +407,76 @@ async function toggleModule(moduleName: string, enabled: boolean) {
|
||||
}
|
||||
}
|
||||
|
||||
async function saveModuleOrder(nextOrder: string[]) {
|
||||
if (orderSaving.value) return
|
||||
const previousOrder = [...moduleOrder.value]
|
||||
moduleOrder.value = normalizeOrderForCurrentModules(nextOrder)
|
||||
orderSaving.value = true
|
||||
try {
|
||||
await modulesApi.updateModuleManagementOrder(moduleOrder.value)
|
||||
success('模块顺序已保存')
|
||||
} catch (err) {
|
||||
moduleOrder.value = previousOrder
|
||||
error(getErrorMessage(err, '保存模块顺序失败'))
|
||||
log.error('保存模块顺序失败:', err)
|
||||
} finally {
|
||||
orderSaving.value = false
|
||||
}
|
||||
}
|
||||
|
||||
function resetModuleOrder() {
|
||||
saveModuleOrder([])
|
||||
}
|
||||
|
||||
function isInteractiveDragTarget(target: EventTarget | null) {
|
||||
return target instanceof HTMLElement &&
|
||||
target.closest('button, a, input, textarea, select, [role="switch"]') !== null
|
||||
}
|
||||
|
||||
function handleModuleDragStart(moduleName: string, event: DragEvent) {
|
||||
if (!canReorderModules.value || isInteractiveDragTarget(event.target)) {
|
||||
event.preventDefault()
|
||||
return
|
||||
}
|
||||
draggedModuleName.value = moduleName
|
||||
if (event.dataTransfer) {
|
||||
event.dataTransfer.effectAllowed = 'move'
|
||||
event.dataTransfer.setData('text/plain', moduleName)
|
||||
}
|
||||
}
|
||||
|
||||
function handleModuleDragEnd() {
|
||||
draggedModuleName.value = null
|
||||
dragOverModuleName.value = null
|
||||
}
|
||||
|
||||
function handleModuleDragOver(moduleName: string) {
|
||||
if (!canReorderModules.value || !draggedModuleName.value || draggedModuleName.value === moduleName) {
|
||||
dragOverModuleName.value = null
|
||||
return
|
||||
}
|
||||
dragOverModuleName.value = moduleName
|
||||
}
|
||||
|
||||
function handleModuleDragLeave(moduleName: string) {
|
||||
if (dragOverModuleName.value === moduleName) {
|
||||
dragOverModuleName.value = null
|
||||
}
|
||||
}
|
||||
|
||||
function handleModuleDrop(targetModuleName: string) {
|
||||
const draggedName = draggedModuleName.value
|
||||
handleModuleDragEnd()
|
||||
if (!canReorderModules.value || !draggedName || draggedName === targetModuleName) return
|
||||
|
||||
const nextOrder = moveNameToTargetIndex(
|
||||
allModules.value.map(module => module.name),
|
||||
draggedName,
|
||||
targetModuleName,
|
||||
)
|
||||
saveModuleOrder(nextOrder)
|
||||
}
|
||||
|
||||
onMounted(() => {
|
||||
fetchModules()
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user