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

feat(gateway): Codex/OpenAI Responses WebSocket 代理模式

Browse Source 
在 /v1/responses 上支持 WebSocket 升级，把客户端帧中继到上游 Codex /
OpenAI Responses WebSocket 端点，同时保持既有的路由、鉴权、配额与用量
语义：

- 路由与准入：control/route/ai.rs 识别 WebSocket 升级请求；
  websocket/ingress.rs 复用 API Key 鉴权、IP 规则与并发许可，并引入
  独立的 WebSocket 连接许可
- 中继：websocket/responses/* 按 connection / session / turn 分层，
  帧解析归一化、socket 写入有界、continuation 保持调度亲和性
- 配额：orchestration/codex_quota_breaker.rs 在账号配额耗尽时熔断并
  自动恢复，不再直接断开客户端连接
- 用量：每个 turn 的终态用量落库，request_metadata 记录
  websocket_mode / websocket_transport，管理端与 usage 视图暴露
  is_websocket
- 管理端：provider 可配置 Responses WebSocket 开关
This commit is contained in:
AAEE86
2026-08-17 14:50:33 +08:00
committed by ZheFox
parent 9a0d346ff3
commit 71b54070e8
72 changed files with 10441 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-provider/pool/src/lib.rs
+117 -3
View File
@@ -34,9 +34,10 @@ pub use providers::{
WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, WINDSURF_USER_STATUS_PATH,
};
pub use quota::{
provider_pool_key_account_quota_exhausted, provider_pool_key_scheduling_label,
provider_pool_member_quota_snapshot, provider_pool_quota_metadata_provider_type,
provider_pool_quota_metadata_updated_at, provider_pool_quota_snapshot_updated_at,
provider_pool_key_account_quota_exhausted, provider_pool_key_quota_hard_blocked,
provider_pool_key_scheduling_label, provider_pool_member_quota_snapshot,
provider_pool_quota_metadata_provider_type, provider_pool_quota_metadata_updated_at,
provider_pool_quota_snapshot_updated_at,
};
pub use quota_refresh::ProviderPoolQuotaRequestSpec;
pub use service::ProviderPoolService;
@@ -693,6 +694,15 @@ mod tests {
#[test]
fn provider_quota_exhaustion_is_adapter_owned() {
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"allowed": false,
"limit_reached": true
}
}))),
"codex",
));
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
@@ -758,6 +768,35 @@ mod tests {
}))),
"codex",
));
assert!(!provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"allowed": true,
"primary_used_percent": 100.0
}
}))),
"codex",
));
let mut explicit_codex_limit = sample_key(None);
explicit_codex_limit.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"allowed": false,
"limit_reached": true,
"usage_ratio": 0.91,
"windows": [{
"code": "weekly",
"used_ratio": 0.91
}]
}
}));
assert!(provider_pool_key_account_quota_exhausted(
&explicit_codex_limit,
"codex",
));
}
#[test]
@@ -817,6 +856,8 @@ mod tests {
&sample_key(Some(json!({
"codex": {
"updated_at": now.saturating_sub(600),
"allowed": false,
"limit_reached": true,
"primary_used_percent": 100.0,
"primary_reset_at": now.saturating_sub(60)
}
@@ -835,6 +876,79 @@ mod tests {
));
}
#[test]
fn codex_explicit_quota_block_is_hard_until_reset() {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time should be after unix epoch")
.as_secs();
assert!(provider_pool_key_quota_hard_blocked(
&sample_key(Some(json!({
"codex": {
"updated_at": now,
"allowed": false,
"limit_reached": true,
"primary_reset_at": now.saturating_add(3600)
}
}))),
"codex",
));
assert!(!provider_pool_key_quota_hard_blocked(
&sample_key(Some(json!({
"codex": {
"updated_at": now,
"primary_used_percent": 100.0,
"primary_reset_at": now.saturating_add(3600)
}
}))),
"codex",
));
assert!(!provider_pool_key_quota_hard_blocked(
&sample_key(Some(json!({
"codex": {
"updated_at": now.saturating_sub(600),
"allowed": false,
"limit_reached": true,
"primary_reset_at": now.saturating_sub(60)
}
}))),
"codex",
));
let mut snapshot_blocked = sample_key(None);
snapshot_blocked.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"allowed": false,
"limit_reached": true,
"usage_ratio": 0.91,
"reset_at": now.saturating_add(3600)
}
}));
assert!(provider_pool_key_quota_hard_blocked(
&snapshot_blocked,
"codex",
));
snapshot_blocked.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"exhausted": true,
"allowed": false,
"limit_reached": true,
"usage_ratio": 0.91,
"reset_at": now.saturating_sub(60)
}
}));
assert!(!provider_pool_key_quota_hard_blocked(
&snapshot_blocked,
"codex",
));
}
#[test]
fn grok_quota_tier_boundaries_match_pool_modes() {
assert_eq!(
crates/aether-provider/pool/src/provider.rs
+5
View File
@@ -64,6 +64,7 @@ pub trait ProviderPoolAdapter: Send + Sync {
quota_reset_seconds: provider_pool_quota_reset_seconds(input.key),
account_blocked: provider_pool_account_blocked(input.key),
quota_exhausted: self.quota_exhausted(input),
quota_hard_blocked: self.quota_hard_blocked(input),
..PoolMemberSignals::default()
}
}
@@ -72,6 +73,10 @@ pub trait ProviderPoolAdapter: Send + Sync {
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
.unwrap_or(false)
}
fn quota_hard_blocked(&self, _input: &ProviderPoolMemberInput<'_>) -> bool {
false
}
}
pub(crate) fn provider_pool_matching_endpoint<F>(
crates/aether-provider/pool/src/providers/codex.rs
+66 -2
View File
@@ -11,8 +11,9 @@ use crate::provider::{
};
use crate::quota::{
provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
provider_pool_member_quota_snapshot, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
};
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -48,6 +49,23 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
}
fn quota_exhausted(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
if let Some(quota_snapshot) =
provider_pool_member_quota_snapshot(input.key, input.provider_type)
{
let explicitly_exhausted = provider_pool_json_bool(quota_snapshot.get("allowed"))
== Some(false)
|| provider_pool_json_bool(quota_snapshot.get("limit_reached")) == Some(true);
if explicitly_exhausted {
let observed_at = provider_pool_timestamp_unix_secs(
quota_snapshot
.get("observed_at")
.or_else(|| quota_snapshot.get("updated_at")),
);
return !provider_pool_current_unix_secs().is_some_and(|now_unix_secs| {
provider_pool_reset_deadline_elapsed(quota_snapshot, observed_at, now_unix_secs)
});
}
}
if let Some(exhausted) =
provider_pool_quota_snapshot_exhausted_decision(input.key, input.provider_type)
{
@@ -57,6 +75,10 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
.is_some_and(quota_exhausted_from_bucket)
}
fn quota_hard_blocked(&self, input: &ProviderPoolMemberInput<'_>) -> bool {
codex_explicit_quota_block_active(input.key, input.provider_type)
}
fn quota_refresh_endpoint(
&self,
endpoints: &[StoredProviderCatalogEndpoint],
@@ -72,6 +94,35 @@ impl ProviderPoolAdapter for CodexProviderPoolAdapter {
}
}
fn codex_explicit_quota_block_active(
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
provider_type: &str,
) -> bool {
let Some(quota_snapshot) = provider_pool_member_quota_snapshot(key, provider_type) else {
return provider_pool_metadata_bucket(key.upstream_metadata.as_ref(), provider_type)
.is_some_and(|bucket| {
(provider_pool_json_bool(bucket.get("allowed")) == Some(false)
|| provider_pool_json_bool(bucket.get("limit_reached")) == Some(true))
&& !["primary", "secondary"]
.into_iter()
.any(|prefix| codex_window_reset_elapsed(bucket, prefix))
});
};
let explicitly_blocked = provider_pool_json_bool(quota_snapshot.get("allowed")) == Some(false)
|| provider_pool_json_bool(quota_snapshot.get("limit_reached")) == Some(true);
if !explicitly_blocked {
return false;
}
let observed_at = provider_pool_timestamp_unix_secs(
quota_snapshot
.get("observed_at")
.or_else(|| quota_snapshot.get("updated_at")),
);
!provider_pool_current_unix_secs().is_some_and(|now_unix_secs| {
provider_pool_reset_deadline_elapsed(quota_snapshot, observed_at, now_unix_secs)
})
}
fn build_codex_wham_headers(
resolved_oauth_auth: Option<(String, String)>,
decrypted_api_key: Option<&str>,
@@ -248,6 +299,19 @@ fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
let allowed = provider_pool_json_bool(bucket.get("allowed"));
let limit_reached = provider_pool_json_bool(bucket.get("limit_reached"));
if allowed == Some(false) || limit_reached == Some(true) {
let reset_elapsed = ["primary", "secondary"]
.into_iter()
.any(|prefix| codex_window_reset_elapsed(bucket, prefix));
if !reset_elapsed {
return true;
}
}
if allowed == Some(true) || limit_reached == Some(false) {
return false;
}
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
return false;
}
crates/aether-provider/pool/src/quota.rs
+12
View File
@@ -18,6 +18,18 @@ pub fn provider_pool_key_account_quota_exhausted(
})
}
pub fn provider_pool_key_quota_hard_blocked(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> bool {
let adapter = ProviderPoolService::with_builtin_adapters().adapter(provider_type);
adapter.quota_hard_blocked(&ProviderPoolMemberInput {
provider_type,
key,
auth_config: None,
})
}
pub fn provider_pool_member_quota_snapshot<'a>(
key: &'a StoredProviderCatalogKey,
provider_type: &str,