i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-04 16:37:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'zhefox/main' into zhefox-main

Browse Source 
# Conflicts:
#	crates/aether-admin/src/provider/quota.rs
#	crates/aether-ai/formats/src/formats/openai/chat/stream.rs
#	crates/aether-ai/formats/src/formats/openai/responses/mod.rs
#	crates/aether-provider/pool/src/provider.rs
#	crates/aether-provider/pool/src/quota.rs
This commit is contained in:
zhefox
2026-09-02 15:25:27 +08:00
parent dbbe7b22ab 144a28f544
commit a0369cf49a
229 changed files with 51071 additions and 896 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
+8
View File
@@ -167,6 +167,14 @@ fn parse_pool_score_rules(pool_advanced: &Map<String, Value>) -> PoolMemberScore
fn normalize_pool_preset_mode(preset: &str, raw_mode: Option<&Value>) -> Option<String> {
match preset {
"cache_affinity" => Some(
raw_mode
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| matches!(*value, "single_account" | "lru"))
.unwrap_or("single_account")
.to_string(),
),
"free_team_first" | "free_first" | "team_first" | "plus_first" | "pro_first" => {
let default_mode = match preset {
"free_team_first" => "both",
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
+1
View File
@@ -1342,6 +1342,7 @@ pub(super) fn build_admin_pool_key_payload(
json!(key.internal_priority),
);
payload.insert("rpm_limit".to_string(), json!(key.rpm_limit));
payload.insert("concurrent_limit".to_string(), json!(key.concurrent_limit));
payload.insert(
"cache_ttl_minutes".to_string(),
json!(key.cache_ttl_minutes),
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+2
View File
@@ -3177,6 +3177,7 @@ async fn provider_query_execute_standard_test_candidate(
crate::ai_serving::openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
request_model,
),
)
else {
@@ -3238,6 +3239,7 @@ async fn provider_query_execute_standard_test_candidate(
crate::ai_serving::openai_responses_reasoning_replay_policy(
transport.provider.provider_type.as_str(),
transport.endpoint.base_url.as_str(),
request_model,
),
)
.is_err()
apps/aether-gateway/src/handlers/admin/users/billing.rs
+70
View File
@@ -34,6 +34,22 @@ fn admin_user_id_from_billing_path(request_path: &str, suffix: &str) -> Option<S
}
}
fn admin_user_entitlement_ids_from_path(request_path: &str) -> Option<(String, String)> {
let rest = request_path
.trim_end_matches('/')
.strip_prefix("/api/admin/users/")?;
let mut parts = rest.split('/');
let user_id = parts.next()?.trim();
if parts.next()? != "billing" || parts.next()? != "entitlements" {
return None;
}
let entitlement_id = parts.next()?.trim();
if user_id.is_empty() || entitlement_id.is_empty() || parts.next().is_some() {
return None;
}
Some((user_id.to_string(), entitlement_id.to_string()))
}
fn admin_user_billing_operator_id(request_context: &AdminRequestContext<'_>) -> Option<String> {
request_context
.decision()
@@ -202,6 +218,60 @@ pub(in super::super) async fn build_admin_list_user_billing_entitlements_respons
}
}
pub(in super::super) async fn build_admin_revoke_user_billing_entitlement_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let Some((user_id, entitlement_id)) =
admin_user_entitlement_ids_from_path(request_context.path())
else {
return Ok(build_admin_users_bad_request_response("缺少套餐权益 ID"));
};
if state.find_user_auth_by_id(&user_id).await?.is_none() {
return Ok((
http::StatusCode::NOT_FOUND,
Json(json!({ "detail": "用户不存在" })),
)
.into_response());
}
match state
.app()
.revoke_user_plan_entitlement(&user_id, &entitlement_id)
.await?
{
crate::LocalMutationOutcome::Applied(()) => {}
crate::LocalMutationOutcome::NotFound => {
return Ok((
http::StatusCode::NOT_FOUND,
Json(json!({ "detail": "套餐权益不存在或已失效" })),
)
.into_response());
}
crate::LocalMutationOutcome::Invalid(detail) => {
return Ok(build_admin_users_bad_request_response(detail));
}
crate::LocalMutationOutcome::Unavailable => {
return Ok(build_admin_users_data_unavailable_response());
}
}
let entitlements = match load_admin_user_entitlements_payload(state, &user_id).await? {
Some(value) => value,
None => return Ok(build_admin_users_data_unavailable_response()),
};
Ok(attach_admin_audit_response(
Json(json!({
"items": entitlements["items"].clone(),
"entitlements": entitlements["items"].clone(),
"total": entitlements["total"].clone(),
}))
.into_response(),
"admin_user_plan_revoked",
"revoke_user_billing_entitlement",
"user_plan_entitlement",
&entitlement_id,
))
}
pub(in super::super) async fn build_admin_grant_user_billing_plan_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
apps/aether-gateway/src/handlers/admin/users/mod.rs
+1
View File
@@ -28,6 +28,7 @@ use self::batch::{
use self::billing::{
build_admin_grant_user_billing_plan_response,
build_admin_list_user_billing_entitlements_response,
build_admin_revoke_user_billing_entitlement_response,
};
use self::groups::{
build_admin_create_user_group_response, build_admin_delete_user_group_response,
apps/aether-gateway/src/handlers/admin/users/routes.rs
+12 -4
View File
@@ -8,10 +8,11 @@ use super::{
build_admin_list_user_group_members_response, build_admin_list_user_groups_response,
build_admin_list_user_sessions_response, build_admin_list_users_response,
build_admin_replace_user_group_members_response, build_admin_resolve_user_selection_response,
build_admin_reveal_user_api_key_response, build_admin_set_default_user_group_response,
build_admin_toggle_user_api_key_lock_response, build_admin_update_user_api_key_response,
build_admin_update_user_group_response, build_admin_update_user_response,
build_admin_user_batch_action_response, build_admin_users_data_unavailable_response,
build_admin_reveal_user_api_key_response, build_admin_revoke_user_billing_entitlement_response,
build_admin_set_default_user_group_response, build_admin_toggle_user_api_key_lock_response,
build_admin_update_user_api_key_response, build_admin_update_user_group_response,
build_admin_update_user_response, build_admin_user_batch_action_response,
build_admin_users_data_unavailable_response,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::GatewayError;
@@ -58,6 +59,10 @@ fn is_admin_users_route(request_context: &AdminRequestContext<'_>) -> bool {
&& path.starts_with("/api/admin/users/")
&& path.ends_with("/billing/grant-plan")
&& path.matches('/').count() == 6)
|| (request_context.method() == http::Method::DELETE
&& path.starts_with("/api/admin/users/")
&& path.contains("/billing/entitlements/")
&& path.matches('/').count() == 7)
|| ((request_context.method() == http::Method::GET
|| request_context.method() == http::Method::PUT
|| request_context.method() == http::Method::DELETE)
@@ -155,6 +160,9 @@ pub(super) async fn maybe_build_local_admin_users_routes_response(
build_admin_grant_user_billing_plan_response(state, request_context, request_body)
.await?,
)),
Some("revoke_user_billing_entitlement") => Ok(Some(
build_admin_revoke_user_billing_entitlement_response(state, request_context).await?,
)),
Some("get_user") => Ok(Some(
build_admin_get_user_response(state, request_context).await?,
)),
apps/aether-gateway/src/handlers/proxy/mod.rs
+94 -11
View File
@@ -105,6 +105,12 @@ const LOCAL_EXECUTION_LOOP_DETECTED_DETAIL: &str =
"Gateway detected an execution runtime request loop back into the local frontdoor";
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
"当前调用方 API Key 并发请求数已达上限，请稍后重试";
const PROVIDER_KEY_CAPACITY_LIMIT_REACHED_DETAIL: &str =
"所有可用上游账号当前均已达到并发或 RPM 上限，请稍后重试";
const PROVIDER_KEY_CAPACITY_LIMIT_SKIP_REASONS: &[&str] = &[
"provider_key_concurrency_limit_reached",
"key_rpm_exhausted",
];
const LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL: &str =
"当前 AI 请求在本地执行规划阶段超时，请稍后重试";
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
@@ -1095,6 +1101,7 @@ async fn proxy_request_inner(
),
}
let (mut parts, body) = request.into_parts();
crate::ai_serving::codex_context::install_codex_fingerprint_context_slot(&mut parts);
let redaction_slot = crate::privacy::RedactionSessionSlot::default();
parts.extensions.insert(redaction_slot.clone());
parts
@@ -1350,6 +1357,7 @@ async fn proxy_request_inner(
.extensions
.get::<crate::middleware::CfConnectingIp>()
.map(|value| value.0.as_str()),
client_ip,
local_proxy_body.as_ref(),
)
.await
@@ -1913,12 +1921,23 @@ async fn proxy_request_inner(
.all_candidates_skipped_for_reason(AUTH_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
|| local_execution_runtime_miss_context
.all_candidates_skipped_for_reason(LEGACY_API_KEY_CONCURRENCY_LIMIT_SKIP_REASON);
let local_execution_runtime_miss_detail = (!auth_api_key_concurrency_limited)
.then(|| {
let provider_key_capacity_limited = local_execution_runtime_miss_diagnostic
.as_ref()
.map(|diagnostic| diagnostic_is_provider_key_capacity_limited(Some(diagnostic)))
.unwrap_or_else(|| {
local_execution_runtime_miss_context
.all_provider_request_body_build_failures_detail()
.all_candidates_skipped_for_reasons(PROVIDER_KEY_CAPACITY_LIMIT_SKIP_REASONS)
});
let local_execution_runtime_miss_detail = provider_key_capacity_limited
.then_some(PROVIDER_KEY_CAPACITY_LIMIT_REACHED_DETAIL.to_string())
.or_else(|| {
(!auth_api_key_concurrency_limited)
.then(|| {
local_execution_runtime_miss_context
.all_provider_request_body_build_failures_detail()
})
.flatten()
})
.flatten()
.or_else(|| {
local_execution_runtime_miss_detail(
control_decision,
@@ -2034,7 +2053,7 @@ async fn proxy_request_inner(
let mut response = build_local_http_error_response(
&trace_id,
control_decision,
http::StatusCode::SERVICE_UNAVAILABLE,
local_execution_runtime_miss_status(provider_key_capacity_limited),
local_execution_runtime_miss_client_message(
local_execution_runtime_miss_detail.as_str(),
)
@@ -2360,6 +2379,30 @@ fn diagnostic_is_auth_api_key_concurrency_limited(
}))
}
fn diagnostic_is_provider_key_capacity_limited(
diagnostic: Option<&LocalExecutionRuntimeMissDiagnostic>,
) -> bool {
let Some(diagnostic) = diagnostic else {
return false;
};
PROVIDER_KEY_CAPACITY_LIMIT_SKIP_REASONS.contains(&diagnostic.reason.as_str())
|| (diagnostic.candidate_count.is_some_and(|candidate_count| {
candidate_count > 0
&& diagnostic.skipped_candidate_count.unwrap_or(0) >= candidate_count
}) && !diagnostic.skip_reasons.is_empty()
&& diagnostic.skip_reasons.iter().all(|(reason, count)| {
PROVIDER_KEY_CAPACITY_LIMIT_SKIP_REASONS.contains(&reason.as_str()) && *count > 0
}))
}
fn local_execution_runtime_miss_status(provider_key_capacity_limited: bool) -> http::StatusCode {
if provider_key_capacity_limited {
http::StatusCode::TOO_MANY_REQUESTS
} else {
http::StatusCode::SERVICE_UNAVAILABLE
}
}
fn local_execution_runtime_miss_route_detail(
decision: Option<&GatewayControlDecision>,
) -> Option<&'static str> {
@@ -2398,14 +2441,15 @@ mod tests {
use super::{
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
owner_forward_request_is_stream, restore_redacted_stream_execution_response,
restore_redacted_sync_execution_response, routing_overlay_allows_affinity_target,
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
RequestBodyBufferPolicy,
diagnostic_is_auth_api_key_concurrency_limited,
diagnostic_is_provider_key_capacity_limited, local_execution_runtime_miss_detail,
local_execution_runtime_miss_status, owner_forward_request_is_stream,
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
routing_overlay_allows_affinity_target, GatewayControlDecision,
LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError, RequestBodyBufferPolicy,
};
use axum::body::{to_bytes, Body, Bytes};
use axum::http::{header, HeaderMap, HeaderValue, Method, Response};
use axum::http::{header, HeaderMap, HeaderValue, Method, Response, StatusCode};
use serde_json::json;
use tokio::sync::Semaphore;
@@ -2885,6 +2929,45 @@ mod tests {
Some("当前调用方 API Key 并发请求数已达上限，请稍后重试")
);
}
#[test]
fn provider_key_capacity_requires_every_skip_reason_to_be_capacity_related() {
let capacity_limited = LocalExecutionRuntimeMissDiagnostic {
reason: "candidate_evaluation_incomplete".to_string(),
candidate_count: Some(2),
skipped_candidate_count: Some(2),
skip_reasons: std::collections::BTreeMap::from([
("provider_key_concurrency_limit_reached".to_string(), 1),
("key_rpm_exhausted".to_string(), 1),
]),
..LocalExecutionRuntimeMissDiagnostic::default()
};
let mixed_failure = LocalExecutionRuntimeMissDiagnostic {
reason: "all_candidates_skipped".to_string(),
candidate_count: Some(2),
skipped_candidate_count: Some(2),
skip_reasons: std::collections::BTreeMap::from([
("provider_key_concurrency_limit_reached".to_string(), 1),
("account_quota_exhausted".to_string(), 1),
]),
..LocalExecutionRuntimeMissDiagnostic::default()
};
assert!(diagnostic_is_provider_key_capacity_limited(Some(
&capacity_limited
)));
assert!(!diagnostic_is_provider_key_capacity_limited(Some(
&mixed_failure
)));
assert_eq!(
local_execution_runtime_miss_status(true),
StatusCode::TOO_MANY_REQUESTS
);
assert_eq!(
local_execution_runtime_miss_status(false),
StatusCode::SERVICE_UNAVAILABLE
);
}
}
#[path = "finalize.rs"]
apps/aether-gateway/src/handlers/proxy/websocket/live/http.rs
+16 -1
View File
@@ -1062,13 +1062,28 @@ mod tests {
#[tokio::test]
async fn codex_realtime_call_remains_on_the_live_handler() {
let decision = GatewayControlDecision::synthetic(
let mut decision = GatewayControlDecision::synthetic(
"/v1/realtime/calls",
Some("ai_public".to_string()),
Some("codex".to_string()),
Some("live".to_string()),
Some("codex:live".to_string()),
);
decision.auth_context = Some(GatewayControlAuthContext {
user_id: "user-codex-realtime".to_string(),
api_key_id: "key-codex-realtime".to_string(),
username: Some("codex-realtime".to_string()),
api_key_name: Some("codex-realtime".to_string()),
balance_remaining: None,
access_allowed: true,
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: true,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: None,
ip_rules: None,
});
let request_context = GatewayPublicRequestContext {
trace_id: "trace-codex-realtime-call".to_string(),
request_method: http::Method::POST,
apps/aether-gateway/src/handlers/proxy/websocket/live/planner.rs
+16 -2
View File
@@ -49,6 +49,8 @@ pub(super) enum LiveAuthMode {
pub(super) struct PlannedLiveCandidate {
pub(super) execution: AiExecutionDecision,
pub(super) pinned_candidate: ResponsesWebSocketPinnedCandidate,
pub(super) codex_fingerprint_context:
aether_provider_transport::CodexFingerprintConvergenceContext,
pub(super) client_model: String,
pub(super) provider_model: String,
pub(super) auth_mode: LiveAuthMode,
@@ -228,8 +230,9 @@ async fn plan_live_candidate_inner(
if validate_model(client_model).is_err() || client_model.len() > MAX_LIVE_MODEL_BYTES {
return Ok(None);
}
let parts = build_live_planning_parts(headers, remote_addr);
let mut parts = build_live_planning_parts(headers, remote_addr);
let body = json!({"model": client_model, "input": []});
crate::ai_serving::codex_context::install_codex_fingerprint_context_slot(&mut parts);
let execution = maybe_build_pinned_stream_local_same_format_provider_decision_payload(
state,
&parts,
@@ -338,6 +341,8 @@ async fn plan_live_candidate_inner(
Ok(Some(PlannedLiveCandidate {
execution,
pinned_candidate,
codex_fingerprint_context:
crate::ai_serving::codex_context::resolve_codex_fingerprint_context(&parts, &body),
client_model: client_model.to_string(),
provider_model,
auth_mode,
@@ -560,7 +565,11 @@ pub(super) fn build_live_stream_admission_attempt(
remote_addr: &SocketAddr,
upstream_url: String,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
let parts = build_live_planning_parts(headers, remote_addr);
let mut parts = build_live_planning_parts(headers, remote_addr);
crate::ai_serving::codex_context::restore_codex_logical_turn_context(
&mut parts,
&candidate.codex_fingerprint_context,
);
let body = json!({"model": candidate.client_model.as_str(), "input": []});
let mut execution = candidate.execution.clone();
execution.upstream_url = Some(upstream_url);
@@ -922,6 +931,11 @@ mod tests {
"key-1",
)
.unwrap(),
codex_fingerprint_context:
aether_provider_transport::CodexFingerprintConvergenceContext::new(
"test-live-turn",
1,
),
client_model: "global-model".to_string(),
provider_model: "provider-model".to_string(),
auth_mode,
apps/aether-gateway/src/handlers/proxy/websocket/live/registry.rs
+5
View File
@@ -718,6 +718,11 @@ mod tests {
PlannedLiveCandidate {
execution,
pinned_candidate: binding.pinned_candidate.clone(),
codex_fingerprint_context:
aether_provider_transport::CodexFingerprintConvergenceContext::new(
"test-live-turn",
1,
),
client_model: binding.client_model.clone(),
provider_model: binding.provider_model.clone(),
auth_mode: binding.auth_mode,
apps/aether-gateway/src/handlers/proxy/websocket/responses/adapters/codex.rs
+117 -10
View File
@@ -13,8 +13,9 @@ use crate::ai_serving::AiExecutionDecision;
use crate::clock::current_unix_secs;
use crate::handlers::proxy::websocket::transport::UpstreamWebSocketErrorCodes;
use crate::orchestration::{
codex_account_id_from_headers, codex_quota_exhaustion_reset_at,
sync_codex_websocket_quota_metadata, ResponsesWebSocketAdapter,
codex_account_id_from_headers, codex_model_quota_exhaustion_reset_at,
codex_quota_exhaustion_reset_at, sync_codex_websocket_quota_metadata,
ResponsesWebSocketAdapter,
};
use crate::AppState;
@@ -114,16 +115,37 @@ impl ResponsesWebSocketProtocolAdapter for CodexResponsesWebSocketAdapter {
event: &Value,
) -> Option<ResponsesWebSocketAdapterObservation> {
let rate_limits = parse_codex_rate_limits(event)?;
let exhausted =
let account_exhausted =
aether_admin::provider::quota::codex_rate_limit_metadata_exhausted(&rate_limits);
let retry_exclusion_until_unix_secs =
codex_quota_exhaustion_reset_at(&rate_limits, current_unix_secs());
let active_limit_exhausted =
aether_admin::provider::quota::codex_websocket_response_has_usage_limit_error(event);
let now_unix_secs = current_unix_secs();
let scoped_reset_at = if account_exhausted {
codex_quota_exhaustion_reset_at(&rate_limits, now_unix_secs)
} else {
codex_model_quota_exhaustion_reset_at(&rate_limits, now_unix_secs)
};
let retry_exclusion_until_unix_secs = active_limit_exhausted
.then(|| {
aether_admin::provider::quota::codex_websocket_usage_limit_reset_at(
event,
now_unix_secs,
)
})
.flatten()
.or(scoped_reset_at);
Some(ResponsesWebSocketAdapterObservation {
drain: exhausted.then_some(ResponsesWebSocketDrainDirective {
error_code: "codex_account_quota_exhausted",
retry_current_turn: true,
retry_exclusion_until_unix_secs,
}),
drain: (account_exhausted || active_limit_exhausted).then_some(
ResponsesWebSocketDrainDirective {
error_code: if account_exhausted {
"codex_account_quota_exhausted"
} else {
"codex_active_limit_exhausted"
},
retry_current_turn: true,
retry_exclusion_until_unix_secs,
},
),
quota_metadata: Some(rate_limits),
})
}
@@ -337,6 +359,91 @@ mod tests {
);
}
#[test]
fn model_scoped_usage_limit_error_drains_without_account_exhaustion() {
let adapter = CodexResponsesWebSocketAdapter;
let event = json!({
"type": "error",
"error": {
"type": "usage_limit_reached",
"plan_type": "pro",
},
"status_code": 429,
"headers": {
"X-Codex-Plan-Type": "pro",
"X-Codex-Active-Limit": "codex_bengalfox",
"X-Codex-Primary-Used-Percent": "100",
"X-Codex-Primary-Window-Minutes": "300",
"X-Codex-Primary-Reset-At": "4000000000",
"X-Codex-Bengalfox-Limit-Name": "GPT-5.3-Codex-Spark",
"X-Codex-Bengalfox-Primary-Used-Percent": "100",
"X-Codex-Bengalfox-Primary-Window-Minutes": "300",
"X-Codex-Bengalfox-Primary-Reset-At": "4000000000",
},
});
let observation = adapter
.observe_upstream_event(&event)
.expect("model-scoped quota error should be observed");
let drain = observation
.drain
.expect("model-scoped quota error should retry the current turn");
let quota = observation
.quota_metadata
.expect("model-scoped quota metadata should be retained");
assert_eq!(drain.error_code, "codex_active_limit_exhausted");
assert!(drain.retry_current_turn);
assert_eq!(
drain.retry_exclusion_until_unix_secs,
Some(4_000_000_000u64)
);
assert_eq!(quota["spark_primary_used_percent"], json!(100.0));
assert!(quota.get("allowed").is_none());
assert!(quota.get("limit_reached").is_none());
assert!(!aether_admin::provider::quota::codex_rate_limit_metadata_exhausted(&quota));
}
#[test]
fn account_scoped_usage_limit_error_keeps_account_drain_semantics() {
let adapter = CodexResponsesWebSocketAdapter;
let event = json!({
"type": "error",
"error": {
"type": "usage_limit_reached",
"plan_type": "free",
"resets_at": 4_000_000_000u64,
},
"status_code": 429,
"headers": {
"X-Codex-Plan-Type": "free",
"X-Codex-Primary-Used-Percent": "100",
"X-Codex-Primary-Window-Minutes": "43200",
"X-Codex-Primary-Reset-At": "4000000000",
},
});
let observation = adapter
.observe_upstream_event(&event)
.expect("account quota error should be observed");
let drain = observation
.drain
.expect("account quota error should retry the current turn");
let quota = observation
.quota_metadata
.expect("account quota metadata should be retained");
assert_eq!(drain.error_code, "codex_account_quota_exhausted");
assert!(drain.retry_current_turn);
assert_eq!(
drain.retry_exclusion_until_unix_secs,
Some(4_000_000_000u64)
);
assert_eq!(quota["allowed"], json!(false));
assert_eq!(quota["limit_reached"], json!(true));
assert!(aether_admin::provider::quota::codex_rate_limit_metadata_exhausted(&quota));
}
#[test]
fn only_known_codex_pre_response_signals_are_safe_to_rebind() {
let adapter = CodexResponsesWebSocketAdapter;
apps/aether-gateway/src/handlers/proxy/websocket/responses/admission.rs
+12 -9
View File
@@ -11,7 +11,7 @@ use aether_contracts::ExecutionPlan;
use crate::execution_runtime::acquire_upstream_execution_gate;
use crate::provider_pool_demand::{
acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
acquire_provider_pool_execution_guard, ProviderPoolInFlightAdmission, ProviderPoolInFlightGuard,
};
use crate::upstream_admission::UpstreamTargetAdmissionPermit;
use crate::{AppState, GatewayError};
@@ -41,14 +41,17 @@ impl ResponsesWebSocketTurnAdmission {
return Err(error);
}
};
let provider_pool = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
&plan.request_id,
plan.candidate_id.as_deref(),
&plan.key_id,
)
.await;
let provider_pool = match acquire_provider_pool_execution_guard(state, plan).await? {
ProviderPoolInFlightAdmission::Acquired(guard) => guard,
ProviderPoolInFlightAdmission::Saturated { limit } => {
drop(upstream_target);
drop(upstream_execution);
return Err(GatewayError::Client {
status: http::StatusCode::TOO_MANY_REQUESTS,
message: format!("上游账号并发已达上限 ({limit})"),
});
}
};
Ok(Self {
upstream_execution,
apps/aether-gateway/src/handlers/proxy/websocket/responses/client.rs
+20 -3
View File
@@ -1,5 +1,6 @@
//! Client-side Responses WebSocket event forwarding and follow-up planning.
use aether_provider_transport::CodexFingerprintConvergenceContext;
use axum::extract::ws::{Message as AxumWsMessage, WebSocket};
use futures_util::SinkExt;
use serde_json::Value;
@@ -248,7 +249,14 @@ pub(super) async fn forward_client_message(
// derive one strong live control snapshot that every stage below
// shares. The connection's Upgrade-time decision is only the
// immutable identity seed.
let planning_parts = build_planning_parts(context);
let logical_turn_id = Uuid::now_v7().to_string();
let mut planning_parts = build_planning_parts(context);
let codex_fingerprint_context =
crate::ai_serving::codex_context::attach_codex_logical_turn_context(
&mut planning_parts,
&client_event,
&logical_turn_id,
);
let turn_control = match resolve_responses_websocket_turn_control(
state,
context,
@@ -453,6 +461,8 @@ pub(super) async fn forward_client_message(
context,
planning_parts,
client_event,
logical_turn_id,
codex_fingerprint_context,
turn_control,
turn_redaction_session,
)
@@ -466,6 +476,8 @@ pub(super) async fn forward_client_message(
planning_parts,
client_event,
requested_model,
logical_turn_id,
codex_fingerprint_context,
turn_control,
raw_responses_lite_static_config
.expect("independent turns always retain their raw static config"),
@@ -513,6 +525,8 @@ async fn forward_pinned_continuation(
context: &WebSocketRequestContext,
planning_parts: http::request::Parts,
client_event: Value,
logical_turn_id: String,
codex_fingerprint_context: CodexFingerprintConvergenceContext,
turn_control: ResponsesWebSocketTurnControl,
turn_redaction_session: Option<RedactionSession>,
) -> RelayDisposition {
@@ -556,7 +570,6 @@ async fn forward_pinned_continuation(
};
let turn_request_id = Uuid::new_v4().to_string();
let logical_turn_id = Uuid::new_v4().to_string();
let planned = match await_owned_responses_websocket_plan(spawn_owned_responses_websocket_plan(
state.clone(),
planning_parts,
@@ -766,6 +779,7 @@ async fn forward_pinned_continuation(
bound.body_normalization = normalization;
bound.turn_state.begin(
LogicalTurn::new(client_event, turn_index, logical_turn_id)
.with_codex_fingerprint_context(codex_fingerprint_context)
.with_provider_store(provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
turn,
@@ -800,12 +814,13 @@ async fn forward_replanned_response_create(
planning_parts: http::request::Parts,
client_event: Value,
requested_model: String,
logical_turn_id: String,
codex_fingerprint_context: CodexFingerprintConvergenceContext,
turn_control: ResponsesWebSocketTurnControl,
raw_responses_lite_static_config: ResponsesLiteStaticConfig,
turn_redaction_session: Option<RedactionSession>,
) -> RelayDisposition {
let turn_request_id = Uuid::new_v4().to_string();
let logical_turn_id = Uuid::new_v4().to_string();
let now_unix_secs = current_unix_secs();
let excluded_key_ids = bound.exhausted_exclusions.key_ids(now_unix_secs);
let excluded_codex_account_ids = bound.exhausted_exclusions.codex_account_ids(now_unix_secs);
@@ -992,6 +1007,7 @@ async fn forward_replanned_response_create(
bound.body_normalization = normalization;
bound.turn_state.begin(
LogicalTurn::new(client_event.clone(), turn_index, logical_turn_id.clone())
.with_codex_fingerprint_context(codex_fingerprint_context.clone())
.with_provider_store(provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
turn,
@@ -1076,6 +1092,7 @@ async fn forward_replanned_response_create(
bound.binding_identity = replacement.binding_identity;
bound.turn_state.begin(
LogicalTurn::new(client_event, turn_index, logical_turn_id)
.with_codex_fingerprint_context(codex_fingerprint_context)
.with_provider_store(provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
turn,
apps/aether-gateway/src/handlers/proxy/websocket/responses/quota.rs
+8 -1
View File
@@ -146,6 +146,7 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
};
let turn_index = active.turn_index;
let logical_turn_id = active.logical_turn_id.clone();
let codex_fingerprint_context = active.codex_fingerprint_context.clone();
let turn_attempt = active.turn_attempt;
let retry_exclusion_until_unix_secs = bound
@@ -154,7 +155,13 @@ pub(super) async fn retry_active_turn_after_quota_exhaustion(
let exhausted_key = record_exhausted_bound_key(bound, retry_exclusion_until_unix_secs);
let exhausted_key_id = exhausted_key.as_ref().map(|(key_id, _)| key_id.clone());
let planning_parts = build_planning_parts(context);
let mut planning_parts = build_planning_parts(context);
if let Some(codex_fingerprint_context) = codex_fingerprint_context.as_ref() {
crate::ai_serving::codex_context::restore_codex_logical_turn_context(
&mut planning_parts,
codex_fingerprint_context,
);
}
let turn_request_id = Uuid::new_v4().to_string();
let now_unix_secs = current_unix_secs();
let excluded_key_ids = bound.exhausted_exclusions.key_ids(now_unix_secs);
apps/aether-gateway/src/handlers/proxy/websocket/responses/session.rs
+9 -2
View File
@@ -444,7 +444,14 @@ async fn bootstrap_responses_websocket(
let raw_responses_lite_static_config =
ResponsesLiteStaticConfig::from_response_create(&first_event);
let planning_parts = build_planning_parts(context);
let first_logical_turn_id = Uuid::now_v7().to_string();
let mut planning_parts = build_planning_parts(context);
let first_codex_fingerprint_context =
crate::ai_serving::codex_context::attach_codex_logical_turn_context(
&mut planning_parts,
&first_event,
&first_logical_turn_id,
);
let turn_control = match resolve_responses_websocket_turn_control(
&state,
context,
@@ -860,7 +867,6 @@ async fn bootstrap_responses_websocket(
return None;
}
};
let first_logical_turn_id = Uuid::new_v4().to_string();
let first_turn_decision = prepare_responses_websocket_turn_decision(
&decision,
context.trace_id.clone(),
@@ -954,6 +960,7 @@ async fn bootstrap_responses_websocket(
}
bound.turn_state.begin(
LogicalTurn::new(first_event, 1, first_logical_turn_id)
.with_codex_fingerprint_context(first_codex_fingerprint_context)
.with_provider_store(first_provider_event.get("store") == Some(&Value::Bool(true)))
.with_turn_control(turn_control),
first_turn,
apps/aether-gateway/src/handlers/proxy/websocket/responses/turn_state.rs
+13
View File
@@ -5,6 +5,7 @@
//! 非法组合只能靠调用点的 if 和「记得同时改另外两个字段」来避免。这里把它收敛成
//! 一个枚举：合法组合由类型保证，转换只能走受控 API。
use aether_provider_transport::CodexFingerprintConvergenceContext;
use serde_json::Value;
use super::control::ResponsesWebSocketTurnControl;
@@ -26,6 +27,9 @@ pub(super) struct LogicalTurn {
pub(super) provider_store: bool,
pub(super) turn_index: u64,
pub(super) logical_turn_id: String,
/// Immutable Codex client identity for every provider attempt belonging to
/// this logical turn. A transparent re-plan must never mint a new turn.
pub(super) codex_fingerprint_context: Option<CodexFingerprintConvergenceContext>,
pub(super) turn_attempt: u32,
pub(super) retry_attempted: bool,
pub(super) retry_unsafe_reason: Option<&'static str>,
@@ -42,6 +46,7 @@ impl LogicalTurn {
provider_store: false,
turn_index,
logical_turn_id,
codex_fingerprint_context: None,
turn_attempt: 1,
retry_attempted: false,
retry_unsafe_reason: None,
@@ -54,6 +59,14 @@ impl LogicalTurn {
self
}
pub(super) fn with_codex_fingerprint_context(
mut self,
context: CodexFingerprintConvergenceContext,
) -> Self {
self.codex_fingerprint_context = Some(context);
self
}
pub(super) fn with_provider_store(mut self, provider_store: bool) -> Self {
self.provider_store = provider_store;
self
apps/aether-gateway/src/handlers/public/mod.rs
+1 -1
View File
@@ -28,5 +28,5 @@ pub(crate) use self::support::{
build_api_key_install_session_response, build_proxy_node_install_session_response,
build_unhandled_public_support_response, matches_model_mapping_for_models,
maybe_build_local_admin_announcements_response, maybe_build_local_public_support_response,
CreateApiKeyInstallSessionRequest,
vscodex_ws_proxy, CreateApiKeyInstallSessionRequest,
};
apps/aether-gateway/src/handlers/public/support.rs
+10
View File
@@ -48,6 +48,8 @@ mod support_payment;
mod support_test_connection;
#[path = "support/user_me.rs"]
mod support_user_me;
#[path = "support/user_me_vscodex.rs"]
mod support_vscodex;
#[path = "support/wallet.rs"]
mod support_wallet;
@@ -89,6 +91,8 @@ use self::support_oauth::maybe_build_local_oauth_response;
use self::support_payment::maybe_build_local_payment_callback_response;
use self::support_test_connection::maybe_build_local_test_connection_response;
use self::support_user_me::maybe_build_local_users_me_response;
pub(crate) use self::support_vscodex::vscodex_ws_proxy;
use self::support_vscodex::{handle_users_me_vscodex_request, maybe_build_local_vscodex_response};
use self::support_wallet::{
build_wallet_balance_payload_for_auth_scope, build_wallet_balance_payload_for_user,
build_wallet_live_today_usage_payload_for_api_key,
@@ -121,6 +125,7 @@ pub(crate) async fn maybe_build_local_public_support_response(
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
cf_connecting_ip: Option<&str>,
client_ip: std::net::IpAddr,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
@@ -192,6 +197,11 @@ pub(crate) async fn maybe_build_local_public_support_response(
.await;
}
if decision.route_family.as_deref() == Some("vscodex") {
return maybe_build_local_vscodex_response(state, request_context, client_ip, request_body)
.await;
}
if decision.route_family.as_deref() == Some("install") {
return maybe_build_local_install_response(state, request_context).await;
}
apps/aether-gateway/src/handlers/public/support/user_me.rs
+3 -2
View File
@@ -2,8 +2,9 @@ use super::{
auth_password_policy_level, base_url_from_request, build_auth_error_response,
build_auth_wallet_summary_payload, decrypt_catalog_secret_with_fallbacks,
encrypt_catalog_secret_with_fallbacks, handle_auth_me,
handle_users_me_api_key_install_session_create, query_param_optional_bool, query_param_value,
resolve_authenticated_local_user, sanitize_public_model_config_for_user, unix_secs_to_rfc3339,
handle_users_me_api_key_install_session_create, handle_users_me_vscodex_request,
query_param_optional_bool, query_param_value, resolve_authenticated_local_user,
sanitize_public_model_config_for_user, unix_secs_to_rfc3339,
users_me_api_key_install_sessions_path_matches, validate_auth_register_password, AppState,
AuthenticatedLocalUserContext, GatewayPublicRequestContext, PUBLIC_CAPABILITY_DEFINITIONS,
};
apps/aether-gateway/src/handlers/public/support/user_me_routes.rs
+12 -3
View File
@@ -18,9 +18,10 @@ use super::{
handle_users_me_preferences_put, handle_users_me_providers_get, handle_users_me_referral_get,
handle_users_me_sessions_get, handle_users_me_update_session, handle_users_me_usage_active_get,
handle_users_me_usage_get, handle_users_me_usage_heatmap_get,
handle_users_me_usage_interval_timeline_get, users_me_api_key_capabilities_path_matches,
users_me_api_key_detail_path_matches, users_me_api_key_install_sessions_path_matches,
users_me_api_key_providers_path_matches, users_me_management_token_detail_path_matches,
handle_users_me_usage_interval_timeline_get, handle_users_me_vscodex_request,
users_me_api_key_capabilities_path_matches, users_me_api_key_detail_path_matches,
users_me_api_key_install_sessions_path_matches, users_me_api_key_providers_path_matches,
users_me_management_token_detail_path_matches,
users_me_management_token_regenerate_path_matches,
users_me_management_token_toggle_path_matches, users_me_management_tokens_root,
users_me_session_detail_path_matches, AppState, GatewayPublicRequestContext,
@@ -55,6 +56,14 @@ pub(crate) async fn maybe_build_local_users_me_response(
{
Some(handle_users_me_delete_other_sessions(state, request_context, headers).await)
}
Some(
"vscodex_devices_list"
| "vscodex_pairing_create"
| "vscodex_device_delete"
| "vscodex_ws_ticket_create",
) => Some(
handle_users_me_vscodex_request(state, request_context, headers, request_body).await,
),
Some("session_delete")
if users_me_session_detail_path_matches(&request_context.request_path) =>
{
apps/aether-gateway/src/handlers/public/support/user_me_vscodex.rs
+931
View File
@@ -0,0 +1,931 @@
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::{Arc, LazyLock, Mutex};
use std::time::Duration;
use axum::body::{Body, Bytes};
use axum::extract::{
ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage, WebSocket, WebSocketUpgrade},
ConnectInfo, State,
};
use axum::http::{self, header};
use axum::response::{IntoResponse, Response};
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use serde_json::{json, Map, Value};
use tokio::sync::Semaphore;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::tungstenite::protocol::{
CloseFrame as TungsteniteCloseFrame, WebSocketConfig,
};
use tokio_tungstenite::tungstenite::Message as TungsteniteMessage;
use tracing::warn;
use super::{
build_auth_error_response, build_auth_json_response, module_available_from_env,
resolve_authenticated_local_user, AppState, GatewayPublicRequestContext,
};
const VSCODEX_ENABLED_ENV: &str = "AETHER_VSCODEX_ENABLED";
const VSCODEX_INTERNAL_URL_ENV: &str = "AETHER_VSCODEX_INTERNAL_URL";
const VSCODEX_INTERNAL_TOKEN_ENV: &str = "AETHER_VSCODEX_INTERNAL_TOKEN";
const VSCODEX_REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
const VSCODEX_MAX_RESPONSE_BYTES: usize = 1024 * 1024;
const VSCODEX_WS_MAX_MESSAGE_BYTES: usize = 16 * 1024 * 1024;
const VSCODEX_WS_MAX_CONNECTIONS: usize = 256;
const VSCODEX_WS_MAX_CONNECTIONS_PER_IP: usize = 16;
const VSCODEX_DEVICE_PATH_PREFIX: &str = "/api/users/me/vscodex/devices/";
const VSCODEX_CLIENT_IP_HEADER: &str = "x-aether-client-ip";
static VSCODEX_HTTP_CLIENT: LazyLock<Result<reqwest::Client, reqwest::Error>> =
LazyLock::new(|| {
reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.build()
});
static VSCODEX_WS_CONNECTIONS: LazyLock<Arc<Semaphore>> =
LazyLock::new(|| Arc::new(Semaphore::new(VSCODEX_WS_MAX_CONNECTIONS)));
static VSCODEX_WS_CONNECTIONS_BY_IP: LazyLock<Arc<VscodexWsIpConnectionLimiter>> =
LazyLock::new(|| {
Arc::new(VscodexWsIpConnectionLimiter::new(
VSCODEX_WS_MAX_CONNECTIONS_PER_IP,
))
});
#[derive(Debug)]
struct VscodexWsIpConnectionLimiter {
max_connections: usize,
active: Mutex<HashMap<IpAddr, usize>>,
}
impl VscodexWsIpConnectionLimiter {
fn new(max_connections: usize) -> Self {
Self {
max_connections: max_connections.max(1),
active: Mutex::new(HashMap::new()),
}
}
fn try_acquire(self: &Arc<Self>, client_ip: IpAddr) -> Option<VscodexWsIpConnectionPermit> {
let mut active = self
.active
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let current = active.get(&client_ip).copied().unwrap_or_default();
if current >= self.max_connections {
return None;
}
active.insert(client_ip, current.saturating_add(1));
Some(VscodexWsIpConnectionPermit {
limiter: Arc::clone(self),
client_ip,
})
}
fn release(&self, client_ip: IpAddr) {
let mut active = self
.active
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let Some(current) = active.get_mut(&client_ip) else {
return;
};
if *current <= 1 {
active.remove(&client_ip);
} else {
*current -= 1;
}
}
#[cfg(test)]
fn active_ip_count(&self) -> usize {
self.active
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.len()
}
}
#[derive(Debug)]
struct VscodexWsIpConnectionPermit {
limiter: Arc<VscodexWsIpConnectionLimiter>,
client_ip: IpAddr,
}
impl Drop for VscodexWsIpConnectionPermit {
fn drop(&mut self) {
self.limiter.release(self.client_ip);
}
}
#[derive(Debug)]
struct VscodexSidecarConfig {
base_url: reqwest::Url,
authorization: reqwest::header::HeaderValue,
http_client: reqwest::Client,
}
#[derive(Debug, Default, Deserialize)]
struct CreatePairingRequest {
name: Option<String>,
}
#[derive(Debug, Deserialize)]
struct CreateWsTicketRequest {
device_id: String,
}
#[derive(Debug, Deserialize)]
struct ExchangePairingRequest {
code: String,
name: Option<String>,
}
pub(crate) async fn vscodex_ws_proxy(
State(state): State<AppState>,
ConnectInfo(remote_addr): ConnectInfo<std::net::SocketAddr>,
ws: WebSocketUpgrade,
headers: http::HeaderMap,
) -> Response<Body> {
let request_permit = match state.try_acquire_request_permit().await {
Ok(value) => value,
Err(err) => {
warn!(error = ?err, "VS Codex WebSocket request admission rejected");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"服务繁忙，请稍后重试",
false,
);
}
};
let client_ip = crate::headers::effective_client_ip(&headers, &remote_addr);
match state.admin_security_ip_blacklisted(client_ip).await {
Ok(true) => {
return build_auth_error_response(
http::StatusCode::FORBIDDEN,
"当前 IP 已被禁止访问",
false,
)
}
Ok(false) => {}
Err(err) => warn!(
client_ip = %client_ip,
error = ?err,
"VS Codex WebSocket IP blacklist check failed open"
),
}
let connection_permit = match Arc::clone(&VSCODEX_WS_CONNECTIONS).try_acquire_owned() {
Ok(value) => value,
Err(_) => {
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 连接数已达上限",
false,
)
}
};
// Only active connections have entries, and each already owns one of the 256 global slots.
let ip_connection_permit = match VSCODEX_WS_CONNECTIONS_BY_IP.try_acquire(client_ip) {
Some(value) => value,
None => {
warn!(
client_ip = %client_ip,
limit = VSCODEX_WS_MAX_CONNECTIONS_PER_IP,
"VS Codex per-IP WebSocket connection limit reached"
);
let mut response = build_auth_error_response(
http::StatusCode::TOO_MANY_REQUESTS,
"当前 IP 的 VS Codex 连接数已达上限",
false,
);
response
.headers_mut()
.insert(header::RETRY_AFTER, http::HeaderValue::from_static("1"));
return response;
}
};
let config = match load_vscodex_sidecar_config() {
Ok(Some(value)) => value,
Ok(None) => {
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务未启用",
false,
)
}
Err(detail) => {
warn!(error = %detail, "VS Codex WebSocket sidecar configuration is invalid");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
let sidecar_url = match build_vscodex_websocket_url(&config.base_url) {
Ok(value) => value,
Err(detail) => {
warn!(error = %detail, "could not build VS Codex sidecar WebSocket URL");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
let mut sidecar_request = match sidecar_url.as_str().into_client_request() {
Ok(value) => value,
Err(err) => {
warn!(error = %err, "could not build VS Codex sidecar WebSocket request");
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
for header_name in [header::ORIGIN, header::SEC_WEBSOCKET_PROTOCOL] {
if let Some(value) = headers.get(&header_name) {
sidecar_request
.headers_mut()
.insert(header_name, value.clone());
}
}
let mut sidecar_config = WebSocketConfig::default();
sidecar_config.max_message_size = Some(VSCODEX_WS_MAX_MESSAGE_BYTES);
sidecar_config.max_frame_size = Some(VSCODEX_WS_MAX_MESSAGE_BYTES);
let (sidecar_socket, sidecar_response) = match tokio::time::timeout(
VSCODEX_REQUEST_TIMEOUT,
tokio_tungstenite::connect_async_with_config(sidecar_request, Some(sidecar_config), true),
)
.await
{
Ok(Ok(value)) => value,
Ok(Err(err)) => {
warn!(error = %err, "VS Codex sidecar WebSocket handshake failed");
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务暂时不可用",
false,
);
}
Err(_) => {
warn!("VS Codex sidecar WebSocket handshake timed out");
return build_auth_error_response(
http::StatusCode::GATEWAY_TIMEOUT,
"VS Codex 服务请求超时",
false,
);
}
};
let selected_protocol = sidecar_response
.headers()
.get(header::SEC_WEBSOCKET_PROTOCOL)
.and_then(|value| value.to_str().ok())
.map(str::to_string);
let ws = ws
.max_message_size(VSCODEX_WS_MAX_MESSAGE_BYTES)
.max_frame_size(VSCODEX_WS_MAX_MESSAGE_BYTES);
let ws = match selected_protocol {
Some(protocol) => ws.protocols([protocol]),
None => ws,
};
drop(request_permit);
ws.on_upgrade(move |browser_socket| async move {
let _connection_permit = connection_permit;
bridge_vscodex_websockets(browser_socket, sidecar_socket, ip_connection_permit).await;
})
}
pub(super) async fn maybe_build_local_vscodex_response(
_state: &AppState,
request_context: &GatewayPublicRequestContext,
client_ip: std::net::IpAddr,
request_body: Option<&Bytes>,
) -> Option<Response<Body>> {
let decision = request_context.control_decision.as_ref()?;
if decision.route_family.as_deref() != Some("vscodex") {
return None;
}
if decision.route_kind.as_deref() != Some("pairing_exchange")
|| !matches!(
request_context.request_path.as_str(),
"/api/vscodex/pair" | "/api/vscodex/pair/"
)
{
return Some(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"VS Codex 接口不存在",
false,
));
}
let config = match load_vscodex_sidecar_config() {
Ok(Some(value)) => value,
Ok(None) => {
return Some(build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务未启用",
false,
))
}
Err(detail) => {
warn!(
error = %detail,
"VS Codex sidecar configuration is invalid"
);
return Some(build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
));
}
};
let payload = match parse_pairing_exchange_request(request_body) {
Ok(value) => value,
Err(response) => return Some(response),
};
let url = match append_vscodex_sidecar_path(&config.base_url, &["v1", "pairings", "exchange"]) {
Ok(value) => value,
Err(detail) => {
warn!(error = %detail, "could not build VS Codex pairing exchange URL");
return Some(build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
));
}
};
let request =
build_authenticated_sidecar_request(&config, reqwest::Method::POST, url, Some(payload))
.header(VSCODEX_CLIENT_IP_HEADER, client_ip.to_string());
Some(send_vscodex_sidecar_request(request, "public", "pairing_exchange").await)
}
pub(super) async fn handle_users_me_vscodex_request(
state: &AppState,
request_context: &GatewayPublicRequestContext,
headers: &http::HeaderMap,
request_body: Option<&Bytes>,
) -> Response<Body> {
let auth = match resolve_authenticated_local_user(state, request_context, headers).await {
Ok(value) => value,
Err(response) => return response,
};
let config = match load_vscodex_sidecar_config() {
Ok(Some(value)) => value,
Ok(None) => {
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务未启用",
false,
)
}
Err(detail) => {
warn!(
user_id = %auth.user.id,
error = %detail,
"VS Codex sidecar configuration is invalid"
);
return build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
);
}
};
let Some(route_kind) = request_context
.control_decision
.as_ref()
.and_then(|decision| decision.route_kind.as_deref())
else {
return build_auth_error_response(
http::StatusCode::NOT_FOUND,
"VS Codex 接口不存在",
false,
);
};
let request = match build_vscodex_sidecar_request(
&config,
&auth.user.id,
route_kind,
&request_context.request_path,
request_body,
) {
Ok(value) => value,
Err(response) => return response,
};
send_vscodex_sidecar_request(request, &auth.user.id, route_kind).await
}
fn load_vscodex_sidecar_config() -> Result<Option<VscodexSidecarConfig>, String> {
if !module_available_from_env(VSCODEX_ENABLED_ENV, false) {
return Ok(None);
}
let raw_url = required_env(VSCODEX_INTERNAL_URL_ENV)?;
let base_url = reqwest::Url::parse(&raw_url)
.map_err(|err| format!("{VSCODEX_INTERNAL_URL_ENV} is invalid: {err}"))?;
if !matches!(base_url.scheme(), "http" | "https")
|| !base_url.has_host()
|| !base_url.username().is_empty()
|| base_url.password().is_some()
|| base_url.query().is_some()
|| base_url.fragment().is_some()
|| base_url.cannot_be_a_base()
{
return Err(format!(
"{VSCODEX_INTERNAL_URL_ENV} must be an HTTP(S) base URL without credentials, query, or fragment"
));
}
let token = required_env(VSCODEX_INTERNAL_TOKEN_ENV)?;
let authorization = reqwest::header::HeaderValue::from_str(&format!("Bearer {token}"))
.map_err(|_| format!("{VSCODEX_INTERNAL_TOKEN_ENV} is not a valid HTTP credential"))?;
let http_client = VSCODEX_HTTP_CLIENT
.as_ref()
.map_err(|err| format!("could not initialize VS Codex HTTP client: {err}"))?
.clone();
Ok(Some(VscodexSidecarConfig {
base_url,
authorization,
http_client,
}))
}
fn required_env(key: &str) -> Result<String, String> {
std::env::var(key)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.ok_or_else(|| format!("{key} is required"))
}
fn build_vscodex_sidecar_request(
config: &VscodexSidecarConfig,
user_id: &str,
route_kind: &str,
request_path: &str,
request_body: Option<&Bytes>,
) -> Result<reqwest::RequestBuilder, Response<Body>> {
let (method, suffix, payload) = match route_kind {
"vscodex_devices_list" => (reqwest::Method::GET, vec!["devices"], None),
"vscodex_pairing_create" => (
reqwest::Method::POST,
vec!["pairings"],
Some(parse_pairing_request(request_body)?),
),
"vscodex_device_delete" => {
let Some(device_id) = vscodex_device_id_from_path(request_path) else {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"设备标识无效",
false,
));
};
(reqwest::Method::DELETE, vec!["devices", device_id], None)
}
"vscodex_ws_ticket_create" => (
reqwest::Method::POST,
vec!["ws-tickets"],
Some(parse_ws_ticket_request(request_body)?),
),
_ => {
return Err(build_auth_error_response(
http::StatusCode::NOT_FOUND,
"VS Codex 接口不存在",
false,
))
}
};
let url = build_vscodex_sidecar_url(&config.base_url, user_id, &suffix).map_err(|detail| {
warn!(user_id = %user_id, error = %detail, "could not build VS Codex sidecar URL");
build_auth_error_response(
http::StatusCode::SERVICE_UNAVAILABLE,
"VS Codex 服务配置不完整",
false,
)
})?;
Ok(build_authenticated_sidecar_request(
config, method, url, payload,
))
}
fn build_authenticated_sidecar_request(
config: &VscodexSidecarConfig,
method: reqwest::Method,
url: reqwest::Url,
payload: Option<Value>,
) -> reqwest::RequestBuilder {
let mut request = config
.http_client
.request(method, url)
.header(header::AUTHORIZATION, config.authorization.clone())
.header(header::ACCEPT, "application/json")
.timeout(VSCODEX_REQUEST_TIMEOUT);
if let Some(payload) = payload {
request = request.json(&payload);
}
request
}
fn build_vscodex_sidecar_url(
base_url: &reqwest::Url,
user_id: &str,
suffix: &[&str],
) -> Result<reqwest::Url, String> {
let mut segments = vec!["internal", "v1", "users", user_id];
segments.extend(suffix.iter().copied());
append_vscodex_sidecar_path(base_url, &segments)
}
fn append_vscodex_sidecar_path(
base_url: &reqwest::Url,
suffix: &[&str],
) -> Result<reqwest::Url, String> {
let mut url = base_url.clone();
let mut path_segments = url
.path_segments_mut()
.map_err(|_| "VS Codex sidecar URL cannot contain path segments".to_string())?;
path_segments.pop_if_empty();
path_segments.extend(suffix.iter().copied());
drop(path_segments);
Ok(url)
}
fn build_vscodex_websocket_url(base_url: &reqwest::Url) -> Result<reqwest::Url, String> {
let mut url = append_vscodex_sidecar_path(base_url, &["api", "vscodex", "ws"])?;
let scheme = match url.scheme() {
"http" => "ws",
"https" => "wss",
_ => return Err("VS Codex sidecar URL must use HTTP(S)".to_string()),
};
url.set_scheme(scheme)
.map_err(|_| "could not convert VS Codex sidecar URL to WebSocket".to_string())?;
Ok(url)
}
fn parse_pairing_request(request_body: Option<&Bytes>) -> Result<Value, Response<Body>> {
let payload = parse_json_request::<CreatePairingRequest>(request_body, true)?;
let mut object = Map::new();
if let Some(name) = payload.name {
object.insert("name".to_string(), Value::String(name));
}
Ok(Value::Object(object))
}
fn parse_ws_ticket_request(request_body: Option<&Bytes>) -> Result<Value, Response<Body>> {
let payload = parse_json_request::<CreateWsTicketRequest>(request_body, false)?;
let device_id = payload.device_id.trim();
if !valid_vscodex_device_id(device_id) {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"设备标识无效",
false,
));
}
Ok(json!({ "device_id": device_id }))
}
fn parse_pairing_exchange_request(request_body: Option<&Bytes>) -> Result<Value, Response<Body>> {
let payload = parse_json_request::<ExchangePairingRequest>(request_body, false)?;
let code = payload.code.trim();
if code.is_empty() || code.len() > 256 || code.chars().any(char::is_control) {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"配对码无效",
false,
));
}
let mut object = Map::from_iter([("code".to_string(), Value::String(code.to_string()))]);
if let Some(name) = payload.name {
object.insert("name".to_string(), Value::String(name));
}
Ok(Value::Object(object))
}
fn parse_json_request<T>(
request_body: Option<&Bytes>,
empty_object_allowed: bool,
) -> Result<T, Response<Body>>
where
T: serde::de::DeserializeOwned,
{
let body = request_body.filter(|body| !body.is_empty());
let result = match body {
Some(body) => serde_json::from_slice(body),
None if empty_object_allowed => serde_json::from_slice(b"{}"),
None => {
return Err(build_auth_error_response(
http::StatusCode::BAD_REQUEST,
"缺少请求体",
false,
))
}
};
result.map_err(|_| {
build_auth_error_response(http::StatusCode::BAD_REQUEST, "请求数据验证失败", false)
})
}
fn vscodex_device_id_from_path(path: &str) -> Option<&str> {
let trimmed = path.trim_end_matches('/');
let device_id = trimmed.strip_prefix(VSCODEX_DEVICE_PATH_PREFIX)?;
if device_id.contains('/') || !valid_vscodex_device_id(device_id) {
return None;
}
Some(device_id)
}
fn valid_vscodex_device_id(value: &str) -> bool {
!value.is_empty() && value.len() <= 128 && !value.chars().any(char::is_control)
}
async fn send_vscodex_sidecar_request(
request: reqwest::RequestBuilder,
request_scope: &str,
operation: &str,
) -> Response<Body> {
let mut upstream = match request.send().await {
Ok(value) => value,
Err(err) => {
warn!(
request_scope = %request_scope,
operation = %operation,
error = %err,
"VS Codex sidecar request failed"
);
let (status, detail) = if err.is_timeout() {
(http::StatusCode::GATEWAY_TIMEOUT, "VS Codex 服务请求超时")
} else {
(http::StatusCode::BAD_GATEWAY, "VS Codex 服务暂时不可用")
};
return build_auth_error_response(status, detail, false);
}
};
let status = http::StatusCode::from_u16(upstream.status().as_u16())
.unwrap_or(http::StatusCode::BAD_GATEWAY);
if matches!(
status,
http::StatusCode::UNAUTHORIZED | http::StatusCode::FORBIDDEN
) {
warn!(
request_scope = %request_scope,
operation = %operation,
upstream_status = status.as_u16(),
"VS Codex sidecar rejected gateway credentials"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务鉴权失败",
false,
);
}
if status.is_redirection() {
warn!(
request_scope = %request_scope,
operation = %operation,
upstream_status = status.as_u16(),
"VS Codex sidecar returned an unexpected redirect"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
if status == http::StatusCode::NO_CONTENT {
return vscodex_no_store_response(status.into_response(), None);
}
let mut response_body = Vec::new();
while let Some(chunk) = match upstream.chunk().await {
Ok(value) => value,
Err(err) => {
warn!(
request_scope = %request_scope,
operation = %operation,
error = %err,
"could not read VS Codex sidecar response"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
} {
if response_body.len().saturating_add(chunk.len()) > VSCODEX_MAX_RESPONSE_BYTES {
warn!(
request_scope = %request_scope,
operation = %operation,
"VS Codex sidecar response exceeded the size limit"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
response_body.extend_from_slice(&chunk);
}
if response_body.is_empty() {
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
let payload = match serde_json::from_slice(&response_body) {
Ok(value) => value,
Err(err) => {
warn!(
request_scope = %request_scope,
operation = %operation,
upstream_status = status.as_u16(),
error = %err,
"VS Codex sidecar returned non-JSON data"
);
return build_auth_error_response(
http::StatusCode::BAD_GATEWAY,
"VS Codex 服务返回无效响应",
false,
);
}
};
let retry_after = upstream.headers().get(header::RETRY_AFTER).cloned();
vscodex_no_store_response(build_auth_json_response(status, payload, None), retry_after)
}
fn vscodex_no_store_response(
mut response: Response<Body>,
retry_after: Option<http::HeaderValue>,
) -> Response<Body> {
response.headers_mut().insert(
header::CACHE_CONTROL,
http::HeaderValue::from_static("no-store"),
);
if let Some(retry_after) = retry_after {
response
.headers_mut()
.insert(header::RETRY_AFTER, retry_after);
}
response
}
async fn bridge_vscodex_websockets<S>(
browser_socket: WebSocket,
sidecar_socket: S,
ip_connection_permit: VscodexWsIpConnectionPermit,
) where
S: futures_util::Stream<
Item = Result<TungsteniteMessage, tokio_tungstenite::tungstenite::Error>,
> + futures_util::Sink<TungsteniteMessage, Error = tokio_tungstenite::tungstenite::Error>
+ Unpin
+ Send
+ 'static,
{
let (mut browser_tx, mut browser_rx) = browser_socket.split();
let (mut sidecar_tx, mut sidecar_rx) = sidecar_socket.split();
let mut ip_connection_permit = Some(ip_connection_permit);
loop {
tokio::select! {
browser_message = browser_rx.next() => {
match browser_message {
Some(Ok(message)) => {
let close = matches!(message, AxumMessage::Close(_));
if let Err(err) = sidecar_tx.send(axum_to_tungstenite_message(message)).await {
warn!(error = %err, "could not forward VS Codex browser WebSocket frame");
break;
}
if close {
break;
}
}
Some(Err(err)) => {
warn!(error = %err, "VS Codex browser WebSocket read failed");
break;
}
None => break,
}
}
sidecar_message = sidecar_rx.next() => {
match sidecar_message {
Some(Ok(TungsteniteMessage::Frame(_))) => continue,
Some(Ok(message)) => {
if ip_connection_permit.is_some() && vscodex_ws_authentication_succeeded(&message) {
ip_connection_permit.take();
}
let close = matches!(message, TungsteniteMessage::Close(_));
if let Err(err) = browser_tx.send(tungstenite_to_axum_message(message)).await {
warn!(error = %err, "could not forward VS Codex sidecar WebSocket frame");
break;
}
if close {
break;
}
}
Some(Err(err)) => {
warn!(error = %err, "VS Codex sidecar WebSocket read failed");
break;
}
None => break,
}
}
}
}
let _ = sidecar_tx.close().await;
let _ = browser_tx.close().await;
}
fn vscodex_ws_authentication_succeeded(message: &TungsteniteMessage) -> bool {
let TungsteniteMessage::Text(text) = message else {
return false;
};
serde_json::from_str::<Value>(text.as_ref())
.ok()
.and_then(|payload| {
payload
.get("type")
.and_then(Value::as_str)
.map(str::to_string)
})
.as_deref()
== Some("auth.ok")
}
fn axum_to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
match message {
AxumMessage::Text(text) => TungsteniteMessage::Text(text.to_string().into()),
AxumMessage::Binary(bytes) => TungsteniteMessage::Binary(bytes),
AxumMessage::Ping(bytes) => TungsteniteMessage::Ping(bytes),
AxumMessage::Pong(bytes) => TungsteniteMessage::Pong(bytes),
AxumMessage::Close(frame) => {
TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
code: frame.code.into(),
reason: frame.reason.to_string().into(),
}))
}
}
}
fn tungstenite_to_axum_message(message: TungsteniteMessage) -> AxumMessage {
match message {
TungsteniteMessage::Text(text) => AxumMessage::Text(text.to_string().into()),
TungsteniteMessage::Binary(bytes) => AxumMessage::Binary(bytes),
TungsteniteMessage::Ping(bytes) => AxumMessage::Ping(bytes),
TungsteniteMessage::Pong(bytes) => AxumMessage::Pong(bytes),
TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
code: frame.code.into(),
reason: frame.reason.to_string().into(),
})),
TungsteniteMessage::Frame(_) => AxumMessage::Close(None),
}
}
#[cfg(test)]
mod tests {
use super::{vscodex_ws_authentication_succeeded, VscodexWsIpConnectionLimiter};
use std::sync::Arc;
use tokio_tungstenite::tungstenite::Message;
#[test]
fn vscodex_ws_ip_limiter_releases_and_removes_inactive_ips() {
let limiter = Arc::new(VscodexWsIpConnectionLimiter::new(1));
let client_ip = "198.51.100.10".parse().expect("IP should parse");
let permit = limiter
.try_acquire(client_ip)
.expect("first connection should acquire");
assert_eq!(limiter.active_ip_count(), 1);
assert!(limiter.try_acquire(client_ip).is_none());
drop(permit);
assert_eq!(limiter.active_ip_count(), 0);
assert!(limiter.try_acquire(client_ip).is_some());
}
#[test]
fn vscodex_ws_ip_limiter_releases_only_after_sidecar_auth_success() {
assert!(vscodex_ws_authentication_succeeded(&Message::Text(
r#"{"type":"auth.ok","role":"operator"}"#.into()
)));
assert!(!vscodex_ws_authentication_succeeded(&Message::Text(
r#"{"type":"auth","token":"client-controlled"}"#.into()
)));
assert!(!vscodex_ws_authentication_succeeded(&Message::Binary(
br#"{"type":"auth.ok"}"#.to_vec().into()
)));
}
}
apps/aether-gateway/src/handlers/shared/catalog.rs
+28 -3
View File
@@ -443,7 +443,18 @@ fn provider_quota_metadata_bucket<'a>(
fn provider_quota_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
let mut parsed = match value {
Some(Value::Number(number)) => number.as_f64(),
Some(Value::String(text)) => text.trim().parse::<f64>().ok(),
Some(Value::String(text)) => {
let text = text.trim();
if let Ok(timestamp) = text.parse::<f64>() {
Some(timestamp)
} else {
return chrono::DateTime::parse_from_rfc3339(text)
.ok()?
.timestamp()
.try_into()
.ok();
}
}
_ => None,
}?;
if !parsed.is_finite() || parsed <= 0.0 {
@@ -560,7 +571,10 @@ fn model_quota_window_snapshot(
.map(|value| value.clamp(0.0, 1.0))
.or_else(|| used_ratio.map(|value| (1.0 - value).max(0.0)));
let reset_at = provider_quota_timestamp_unix_secs(
item.get("reset_at").or_else(|| item.get("next_reset_at")),
item.get("reset_at")
.or_else(|| item.get("next_reset_at"))
.or_else(|| item.get("reset_time"))
.or_else(|| item.get("next_reset_time")),
);
let reset_seconds = quota_window_reset_seconds(observed_at_unix_secs, reset_at);
let is_exhausted = item
@@ -4094,7 +4108,8 @@ mod tests {
},
"claude-sonnet-4-6": {
"display_name": "Claude Sonnet 4.6 (Thinking)",
"remaining_fraction": 0.3
"remaining_fraction": 0.3,
"reset_time": "2026-04-07T12:34:56Z"
}
}
}
@@ -4139,6 +4154,16 @@ mod tests {
label_for_model("claude-sonnet-4-6"),
Some(json!("Claude Sonnet 4.6 (Thinking)"))
);
let claude_window = windows
.iter()
.filter_map(Value::as_object)
.find(|window| window.get("model") == Some(&json!("claude-sonnet-4-6")))
.expect("Claude quota window should exist");
assert_eq!(
claude_window.get("reset_at"),
Some(&json!(1_775_565_296u64))
);
assert_eq!(claude_window.get("reset_seconds"), Some(&json!(12_011u64)));
}
#[test]
apps/aether-gateway/src/handlers/shared/request_utils.rs
+7 -1
View File
@@ -495,8 +495,14 @@ pub(crate) fn public_support_local_requires_buffered_body(
Some(
"api_keys_create"
| "api_key_install_session_create"
| "management_tokens_create",
| "management_tokens_create"
| "vscodex_pairing_create"
| "vscodex_ws_ticket_create",
),
) | (
Some("vscodex"),
http::Method::POST,
Some("pairing_exchange"),
) | (
Some("wallet"),
http::Method::POST,