i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 18:59:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(proxy): 重构 Proxy 节点管理与隧道系统

Browse Source 
- 重构 proxy_nodes 管理端，支持节点注册、心跳、隧道生命周期管理
- 增强 tunnel 嵌入式 hub 和隧道协议
- 重构 aether-proxy 配置、隧道客户端、心跳和调度机制
- 调整 admin OAuth/配额/导入等处理器的参数传递
- 扩展数据迁移模块
- 补充 proxy nodes、OAuth、配额、系统导入等测试
- 更新前端 proxy nodes 视图和 API
This commit is contained in:
fawney19
2026-04-14 22:51:02 +08:00
parent fb31928e44
commit a4e7ac1df6
55 changed files with 3477 additions and 892 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
+19 -6
View File
@@ -8,9 +8,11 @@ use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider
use crate::handlers::admin::provider::oauth::provisioning::build_provider_oauth_auth_config_from_token_payload;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update;
use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
};
@@ -116,7 +118,18 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
.list_provider_catalog_endpoints_by_provider_ids(&[provider_id.to_string()])
.await?;
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id);
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
proxy_node_id,
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let mut results = Vec::with_capacity(entries.len());
let mut success = 0usize;
let mut failed = 0usize;
@@ -126,7 +139,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
state,
template,
entry.refresh_token.as_str(),
proxy_node_id,
request_proxy.clone(),
)
.await
{
@@ -189,7 +202,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
expires_at,
)
.await?
@@ -232,7 +245,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
expires_at,
)
.await?
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs
+23 -9
View File
@@ -2,9 +2,11 @@ use super::parse::{AdminProviderOAuthBatchImportEntry, AdminProviderOAuthBatchIm
use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider_oauth_key;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update;
use crate::handlers::admin::provider::oauth::state::decode_jwt_claims;
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
@@ -14,6 +16,7 @@ use aether_admin::provider::oauth::{
build_kiro_batch_import_key_name, coerce_admin_provider_oauth_import_str,
parse_admin_provider_oauth_kiro_batch_import_entries,
};
use aether_contracts::ProxySnapshot;
use serde_json::{json, Map, Value};
use std::collections::BTreeSet;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -102,7 +105,7 @@ fn admin_provider_oauth_kiro_refresh_error_detail(
async fn refresh_admin_provider_oauth_kiro_auth_config(
state: &AdminAppState<'_>,
auth_config: &AdminKiroAuthConfig,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
social_refresh_base_url: Option<&str>,
idc_refresh_base_url: Option<&str>,
) -> Result<AdminKiroAuthConfig, String> {
@@ -164,7 +167,7 @@ async fn refresh_admin_provider_oauth_kiro_auth_config(
"grantType": "refresh_token",
})),
None,
proxy_node_id,
proxy.clone(),
)
.await
.map_err(|err| format!("IDC refresh 请求失败: {err}"))?;
@@ -268,7 +271,7 @@ async fn refresh_admin_provider_oauth_kiro_auth_config(
.unwrap_or_default(),
})),
None,
proxy_node_id,
proxy,
)
.await
.map_err(|err| format!("social refresh 请求失败: {err}"))?;
@@ -352,7 +355,18 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(&[provider_id.to_string()])
.await?;
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id);
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
proxy_node_id,
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let social_refresh_base_url =
admin_provider_oauth_kiro_refresh_base_url_override(state, "kiro_social_refresh");
let idc_refresh_base_url =
@@ -392,7 +406,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
refreshed_auth_config = match refresh_admin_provider_oauth_kiro_auth_config(
state,
&refreshed_auth_config,
proxy_node_id,
request_proxy.clone(),
social_refresh_base_url.as_deref(),
idc_refresh_base_url.as_deref(),
)
@@ -477,7 +491,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
refreshed_auth_config.expires_at,
)
.await?
@@ -511,7 +525,7 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
&access_token,
&auth_config,
&provider_oauth_active_api_formats(&endpoints),
key_proxy.clone(),
None,
refreshed_auth_config.expires_at,
)
.await?
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
+18 -1
View File
@@ -1,5 +1,6 @@
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::quota::codex::refresh_codex_provider_quota_locally;
use super::super::super::runtime::provider_oauth_runtime_endpoint_for_provider;
use super::super::super::state::{
admin_provider_oauth_template, enrich_admin_provider_oauth_auth_config,
is_fixed_provider_type_for_provider_oauth, json_non_empty_string, json_u64_value,
@@ -123,6 +124,22 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
"该 Provider 不支持 OAuth 授权",
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
payload.proxy_node_id.as_deref(),
&[
key.proxy.as_ref(),
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_payload = match state
.exchange_admin_provider_oauth_code(
@@ -130,7 +147,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
payload.proxy_node_id.as_deref(),
request_proxy,
)
.await
{
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/provider.rs
+22 -10
View File
@@ -2,10 +2,11 @@ use super::super::super::duplicates::find_duplicate_provider_oauth_key;
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::provisioning::{
build_provider_oauth_auth_config_from_token_payload, create_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use super::super::super::runtime::refresh_provider_oauth_account_state_after_update;
use super::super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
is_fixed_provider_type_for_provider_oauth,
@@ -114,6 +115,21 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
"该 Provider 不支持 OAuth 授权",
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
payload.proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_payload = match state
.exchange_admin_provider_oauth_code(
@@ -121,7 +137,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
payload.proxy_node_id.as_deref(),
request_proxy,
)
.await
{
@@ -138,11 +154,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(payload.proxy_node_id.as_deref());
let duplicate = match state
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
.await
@@ -163,7 +175,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
expires_at,
)
.await?
@@ -204,7 +216,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
expires_at,
)
.await?
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/authorize.rs
+18 -6
View File
@@ -1,5 +1,6 @@
use super::session::AdminProviderOAuthDeviceAuthorizePayload;
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::runtime::provider_oauth_runtime_endpoint_for_provider;
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response, current_unix_secs,
default_kiro_device_start_url, generate_provider_oauth_nonce, json_non_empty_string,
@@ -69,6 +70,21 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
"设备授权仅支持 Kiro provider",
));
}
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
payload.proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let region = normalize_kiro_device_region(Some(payload.region.as_str())).ok_or_else(|| {
build_internal_control_error_response(http::StatusCode::BAD_REQUEST, "region 格式无效")
@@ -85,11 +101,7 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
};
let client_registration = match state
.register_admin_kiro_device_oidc_client(
&region,
&start_url,
payload.proxy_node_id.as_deref(),
)
.register_admin_kiro_device_oidc_client(&region, &start_url, request_proxy.clone())
.await
{
Ok(payload) => payload,
@@ -114,7 +126,7 @@ pub(super) async fn handle_admin_provider_oauth_device_authorize(
&client_id,
&client_secret,
&start_url,
payload.proxy_node_id.as_deref(),
request_proxy,
)
.await
{
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/poll.rs
+23 -11
View File
@@ -5,9 +5,11 @@ use crate::handlers::admin::provider::oauth::duplicates::find_duplicate_provider
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, provider_oauth_active_api_formats,
provider_oauth_key_proxy_value, update_existing_provider_oauth_catalog_key,
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update;
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response, build_kiro_device_key_name,
current_unix_secs, decode_jwt_claims, json_non_empty_string, json_u64_value,
@@ -125,6 +127,21 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
"Provider 不存在",
));
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider("kiro", &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
session.proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_result = match state
.poll_admin_kiro_device_token(
@@ -132,7 +149,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&session.client_id,
&session.client_secret,
&session.device_code,
session.proxy_node_id.as_deref(),
request_proxy,
)
.await
{
@@ -252,12 +269,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
}
};
let key_proxy = provider_oauth_key_proxy_value(session.proxy_node_id.as_deref());
let api_formats = provider_oauth_active_api_formats(
&state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?,
);
let api_formats = provider_oauth_active_api_formats(&endpoints);
let mut replaced = false;
let persisted_key = if let Some(existing_key) = duplicate {
replaced = true;
@@ -266,7 +278,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
Some(expires_at),
)
.await?
@@ -288,7 +300,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
Some(expires_at),
)
.await?
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
+22 -14
View File
@@ -2,10 +2,11 @@ use super::super::duplicates::find_duplicate_provider_oauth_key;
use super::super::errors::build_internal_control_error_response;
use super::super::provisioning::{
build_provider_oauth_auth_config_from_token_payload, create_provider_oauth_catalog_key,
provider_oauth_active_api_formats, provider_oauth_key_proxy_value,
update_existing_provider_oauth_catalog_key,
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
};
use super::super::runtime::refresh_provider_oauth_account_state_after_update;
use super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
exchange_admin_provider_oauth_refresh_token, is_fixed_provider_type_for_provider_oauth,
@@ -96,13 +97,24 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
let Some(template) = admin_provider_oauth_template(&provider_type) else {
return Ok(build_admin_provider_oauth_backend_unavailable_response());
};
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let runtime_endpoint = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints);
let request_proxy = state
.resolve_admin_provider_oauth_operation_proxy_snapshot(
proxy_node_id.as_deref(),
&[
runtime_endpoint
.as_ref()
.and_then(|endpoint| endpoint.proxy.as_ref()),
provider.proxy.as_ref(),
],
)
.await;
let token_payload = match state
.exchange_admin_provider_oauth_refresh_token(
template,
refresh_token_input,
proxy_node_id.as_deref(),
)
.exchange_admin_provider_oauth_refresh_token(template, refresh_token_input, request_proxy)
.await
{
Ok(payload) => payload,
@@ -124,11 +136,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
auth_config.insert("refresh_token".to_string(), json!(refresh_token));
}
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let api_formats = provider_oauth_active_api_formats(&endpoints);
let key_proxy = provider_oauth_key_proxy_value(proxy_node_id.as_deref());
let duplicate = match state
.find_duplicate_provider_oauth_key(&provider_id, &auth_config, None)
.await
@@ -149,7 +157,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
&existing_key,
&access_token,
&auth_config,
key_proxy.clone(),
None,
expires_at,
)
.await?
@@ -189,7 +197,7 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
&access_token,
&auth_config,
&api_formats,
key_proxy.clone(),
None,
expires_at,
)
.await?
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/request.rs
+1 -1
View File
@@ -79,7 +79,7 @@ pub(super) async fn parse_admin_provider_oauth_refresh_request(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, endpoints)
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints)
else {
return Ok(RefreshDispatch::Respond(response::control_error_response(
http::StatusCode::BAD_REQUEST,
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+15 -17
View File
@@ -1,13 +1,14 @@
use super::shared::{
coerce_json_f64, coerce_json_string, execute_provider_quota_plan,
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
coerce_json_f64, coerce_json_string, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_antigravity_usage_response;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -31,6 +32,14 @@ async fn execute_antigravity_quota_plan(
.or_insert_with(|| "antigravity".to_string());
let body = json!({ "project": project_id });
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("antigravity-quota:{}", transport.key.id),
candidate_id: None,
@@ -56,20 +65,9 @@ async fn execute_antigravity_quota_plan(
client_api_format: "gemini:chat".to_string(),
provider_api_format: "antigravity:fetch_available_models".to_string(),
model_name: Some("fetchAvailableModels".to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
proxy,
tls_profile: state.resolve_transport_tls_profile(transport),
timeouts: state
.resolve_transport_execution_timeouts(transport)
.or(Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(30_000),
write_ms: Some(30_000),
pool_ms: Some(30_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
})),
timeouts,
};
execute_provider_quota_plan(state, transport, plan, "antigravity").await
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs
+15 -15
View File
@@ -1,9 +1,12 @@
use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};
use super::super::shared::{
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
ProviderQuotaExecutionOutcome,
};
use super::parse::normalize_codex_plan_type;
use crate::handlers::admin::provider::shared::payloads::CODEX_WHAM_USAGE_URL;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_contracts::{ExecutionPlan, RequestBody};
use std::collections::BTreeMap;
pub(super) fn build_codex_refresh_headers(
@@ -58,6 +61,14 @@ pub(super) async fn execute_codex_quota_plan(
transport: &AdminGatewayProviderTransportSnapshot,
headers: BTreeMap<String, String>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("codex-quota:{}", transport.key.id),
candidate_id: None,
@@ -79,20 +90,9 @@ pub(super) async fn execute_codex_quota_plan(
client_api_format: "openai:cli".to_string(),
provider_api_format: "openai:cli".to_string(),
model_name: Some("codex-wham-usage".to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
proxy,
tls_profile: state.resolve_transport_tls_profile(transport),
timeouts: state
.resolve_transport_execution_timeouts(transport)
.or(Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(30_000),
write_ms: Some(30_000),
pool_ms: Some(30_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
})),
timeouts,
};
execute_provider_quota_plan(state, transport, plan, "codex").await
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/plan.rs
+12 -14
View File
@@ -1,3 +1,4 @@
use super::super::shared::default_provider_quota_execution_timeouts;
use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};
use crate::handlers::admin::provider::shared::payloads::{
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
@@ -6,7 +7,7 @@ use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_contracts::{ExecutionPlan, RequestBody};
use std::collections::BTreeMap;
use url::form_urlencoded;
use uuid::Uuid;
@@ -66,6 +67,14 @@ pub(super) async fn execute_kiro_quota_plan(
transport: &AdminGatewayProviderTransportSnapshot,
auth: &AdminKiroRequestAuth,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("kiro-quota:{}", transport.key.id),
candidate_id: None,
@@ -87,20 +96,9 @@ pub(super) async fn execute_kiro_quota_plan(
client_api_format: "claude:cli".to_string(),
provider_api_format: "kiro:usage".to_string(),
model_name: Some("kiro-usage-limits".to_string()),
proxy: state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await,
proxy,
tls_profile: state.resolve_transport_tls_profile(transport),
timeouts: state
.resolve_transport_execution_timeouts(transport)
.or(Some(ExecutionTimeouts {
connect_ms: Some(30_000),
read_ms: Some(30_000),
write_ms: Some(30_000),
pool_ms: Some(30_000),
total_ms: Some(30_000),
..ExecutionTimeouts::default()
})),
timeouts,
};
execute_provider_quota_plan(state, transport, plan, "kiro").await
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+26 -1
View File
@@ -4,16 +4,37 @@ use crate::handlers::admin::provider::shared::payloads::{
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionResult};
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
const PROVIDER_QUOTA_DEFAULT_TIMEOUT_MS: u64 = 30_000;
const PROVIDER_QUOTA_PROXY_TIMEOUT_MS: u64 = 60_000;
pub(super) enum ProviderQuotaExecutionOutcome {
Response(ExecutionResult),
Failure(String),
}
pub(super) fn default_provider_quota_execution_timeouts(
proxy: Option<&ProxySnapshot>,
) -> ExecutionTimeouts {
let timeout_ms = if proxy.is_some() {
PROVIDER_QUOTA_PROXY_TIMEOUT_MS
} else {
PROVIDER_QUOTA_DEFAULT_TIMEOUT_MS
};
ExecutionTimeouts {
connect_ms: Some(timeout_ms),
read_ms: Some(timeout_ms),
write_ms: Some(timeout_ms),
pool_ms: Some(timeout_ms),
total_ms: Some(timeout_ms),
..ExecutionTimeouts::default()
}
}
pub(super) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
}
@@ -127,6 +148,9 @@ pub(super) async fn execute_provider_quota_plan(
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let proxy_source = state
.resolve_transport_proxy_source_with_tunnel_affinity(transport)
.await;
let proxy_url_present = plan
.proxy
.as_ref()
@@ -138,6 +162,7 @@ pub(super) async fn execute_provider_quota_plan(
endpoint_id = %transport.endpoint.id,
url = %plan.url,
tls_profile = ?plan.tls_profile.as_deref(),
proxy_source = ?proxy_source,
proxy_node_id = ?proxy_node_id,
proxy_url_present,
error = %error,
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
+34 -20
View File
@@ -9,28 +9,34 @@ use aether_data_contracts::repository::provider_catalog::{
pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
provider_type: &str,
endpoints: Vec<StoredProviderCatalogEndpoint>,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
let provider_type = provider_type.trim().to_ascii_lowercase();
match provider_type.as_str() {
"codex" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:cli")
}),
"antigravity" => endpoints.into_iter().find(|endpoint| {
endpoint.is_active
&& (endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:chat")
|| endpoint
"codex" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:cli"))
}),
.eq_ignore_ascii_case("openai:cli")
})
.cloned(),
"antigravity" => endpoints
.iter()
.find(|endpoint| {
endpoint.is_active
&& (endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:chat")
|| endpoint
.api_format
.trim()
.eq_ignore_ascii_case("gemini:cli"))
})
.cloned(),
"kiro" => endpoints
.iter()
.find(|endpoint| {
@@ -41,8 +47,16 @@ pub(crate) fn provider_oauth_runtime_endpoint_for_provider(
.eq_ignore_ascii_case("claude:cli")
})
.cloned()
.or_else(|| endpoints.into_iter().find(|endpoint| endpoint.is_active)),
_ => endpoints.into_iter().find(|endpoint| endpoint.is_active),
.or_else(|| {
endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned()
}),
_ => endpoints
.iter()
.find(|endpoint| endpoint.is_active)
.cloned(),
}
}
@@ -59,7 +73,7 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
let endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, endpoints)
let Some(endpoint) = provider_oauth_runtime_endpoint_for_provider(&provider_type, &endpoints)
else {
return Ok((false, None));
};
apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs
+25 -10
View File
@@ -3,16 +3,25 @@ use super::super::errors::{
};
use super::json_non_empty_string;
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
use aether_contracts::ProxySnapshot;
use axum::{body::Body, http, response::Response};
use url::form_urlencoded;
fn provider_oauth_transport_error_detail(prefix: &str, error: &str) -> String {
let error = error.trim();
if error.is_empty() {
return prefix.to_string();
}
format!("{prefix}: {error}")
}
pub(crate) async fn exchange_admin_provider_oauth_code(
state: &AdminAppState<'_>,
template: AdminProviderOAuthTemplate,
code: &str,
state_nonce: &str,
pkce_verifier: Option<&str>,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url);
let response = if template.provider_type == "claude_code" {
@@ -63,7 +72,7 @@ pub(crate) async fn exchange_admin_provider_oauth_code(
Some("application/json"),
Some(serde_json::Value::Object(body)),
None,
proxy_node_id,
proxy.clone(),
)
.await
} else {
@@ -100,12 +109,15 @@ pub(crate) async fn exchange_admin_provider_oauth_code(
Some("application/x-www-form-urlencoded"),
None,
Some(form_body),
proxy_node_id,
proxy.clone(),
)
.await
}
.map_err(|_| {
build_internal_control_error_response(http::StatusCode::BAD_REQUEST, "token exchange 失败")
.map_err(|error| {
build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
provider_oauth_transport_error_detail("token exchange 失败", &error),
)
})?;
if !response.status.is_success() {
@@ -134,7 +146,7 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
state: &AdminAppState<'_>,
template: AdminProviderOAuthTemplate,
refresh_token: &str,
proxy_node_id: Option<&str>,
proxy: Option<ProxySnapshot>,
) -> Result<serde_json::Value, Response<Body>> {
let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url);
let scope = template.scopes.join(" ");
@@ -175,7 +187,7 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
Some("application/json"),
Some(serde_json::Value::Object(body)),
None,
proxy_node_id,
proxy.clone(),
)
.await
} else {
@@ -211,14 +223,17 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
Some("application/x-www-form-urlencoded"),
None,
Some(form_body),
proxy_node_id,
proxy.clone(),
)
.await
}
.map_err(|_| {
.map_err(|error| {
build_internal_control_error_response(
http::StatusCode::BAD_REQUEST,
"Refresh Token 验证失败: token exchange 失败",
provider_oauth_transport_error_detail(
"Refresh Token 验证失败: token exchange 失败",
&error,
),
)
})?;
apps/aether-gateway/src/handlers/admin/provider/write/normalize.rs
+5 -8
View File
@@ -28,9 +28,6 @@ pub(crate) fn normalize_pool_advanced_config(
serde_json::Value::Null => Ok(None),
// `pool_advanced: {}` still means "enable pool mode with defaults".
serde_json::Value::Object(map) => Ok(Some(serde_json::Value::Object(map))),
// Backward compatibility for older boolean payloads.
serde_json::Value::Bool(true) => Ok(Some(serde_json::json!({}))),
serde_json::Value::Bool(false) => Ok(None),
_ => Err("pool_advanced 必须是 JSON 对象".to_string()),
}
}
@@ -78,14 +75,14 @@ mod tests {
}
#[test]
fn normalize_pool_advanced_accepts_legacy_booleans() {
fn normalize_pool_advanced_rejects_legacy_booleans() {
assert_eq!(
normalize_pool_advanced_config(Some(json!(true))).expect("true should normalize"),
Some(json!({}))
normalize_pool_advanced_config(Some(json!(true))).unwrap_err(),
"pool_advanced 必须是 JSON 对象"
);
assert_eq!(
normalize_pool_advanced_config(Some(json!(false))).expect("false should normalize"),
None
normalize_pool_advanced_config(Some(json!(false))).unwrap_err(),
"pool_advanced 必须是 JSON 对象"
);
}
}