i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-08 18:37: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
apps/aether-gateway/src/orchestration/codex_quota_breaker.rs
+374
View File
@@ -0,0 +1,374 @@
//! Short-lived runtime circuit breaker for exhausted Codex accounts.
//!
//! Persisted provider-key quota snapshots remain the durable source of truth.
//! This module closes the interval between receiving a definitive WebSocket
//! `usage_limit_reached` event and every scheduler/cache replica observing the
//! persisted snapshot. The account-scoped entry also protects pools that
//! contain more than one catalog key for the same ChatGPT account.
use std::collections::BTreeMap;
use serde_json::{json, Map, Value};
use sha2::{Digest, Sha256};
use tracing::{info, warn};
use crate::clock::current_unix_secs;
use crate::{AppState, GatewayError};
const CODEX_QUOTA_BREAKER_KEY_PREFIX: &str = "aether:codex:quota-breaker:v1";
const CODEX_QUOTA_BREAKER_FALLBACK_TTL_SECONDS: u64 = 300;
const CODEX_QUOTA_BREAKER_MAX_TTL_SECONDS: u64 = 31 * 24 * 60 * 60;
/// Installs immediate runtime exclusions for a definitive Codex quota
/// exhaustion signal. When RuntimeState is backed by Redis the exclusions are
/// shared by every gateway node; the in-memory backend still protects the
/// current node.
pub(crate) async fn install_codex_quota_exhaustion_breaker(
state: &AppState,
report_context: Option<&Value>,
quota_metadata: &Value,
source: &str,
) -> Result<bool, GatewayError> {
if !aether_admin::provider::quota::codex_rate_limit_metadata_exhausted(quota_metadata) {
return Ok(false);
}
let keys = codex_quota_breaker_keys_from_report_context(report_context);
if keys.is_empty() {
return Ok(false);
}
let now_unix_secs = current_unix_secs();
let (ttl_seconds, reset_at_unix_secs) = codex_quota_breaker_ttl(quota_metadata, now_unix_secs);
let value = json!({
"version": 1,
"observed_at": now_unix_secs,
"reset_at": reset_at_unix_secs,
"source": source,
})
.to_string();
for key in &keys {
state.runtime_kv_setex(key, &value, ttl_seconds).await?;
}
info!(
event_name = "codex_account_quota_breaker_installed",
log_type = "event",
scope_count = keys.len(),
ttl_seconds,
reset_at_unix_secs = ?reset_at_unix_secs,
source,
"gateway installed immediate Codex quota exhaustion exclusions"
);
Ok(true)
}
/// Returns whether a planned Codex request is temporarily blocked by a
/// definitive account quota signal that has not yet been observed in the
/// durable provider catalog.
pub(crate) async fn codex_quota_breaker_blocks_candidate(
state: &AppState,
provider_type: Option<&str>,
key_id: Option<&str>,
provider_request_headers: &BTreeMap<String, String>,
) -> Result<bool, GatewayError> {
if !provider_type.is_some_and(|value| value.trim().eq_ignore_ascii_case("codex")) {
return Ok(false);
}
for key in codex_quota_breaker_keys(
key_id,
codex_account_id_from_headers(provider_request_headers),
) {
if state.runtime_kv_exists(&key).await? {
return Ok(true);
}
}
Ok(false)
}
fn codex_quota_breaker_keys_from_report_context(report_context: Option<&Value>) -> Vec<String> {
let key_id = report_context
.and_then(|context| context.get("key_id"))
.and_then(Value::as_str);
let account_id = report_context
.and_then(|context| context.get("provider_request_headers"))
.and_then(Value::as_object)
.and_then(account_id_from_header_object);
codex_quota_breaker_keys(key_id, account_id)
}
fn codex_quota_breaker_keys(key_id: Option<&str>, account_id: Option<&str>) -> Vec<String> {
let mut keys = Vec::with_capacity(2);
if let Some(account_id) = normalize_identifier(account_id) {
keys.push(codex_quota_breaker_runtime_key("account", account_id));
}
if let Some(key_id) = normalize_identifier(key_id) {
keys.push(codex_quota_breaker_runtime_key("key", key_id));
}
keys
}
fn codex_quota_breaker_runtime_key(scope: &str, identifier: &str) -> String {
format!(
"{CODEX_QUOTA_BREAKER_KEY_PREFIX}:{scope}:{}",
opaque_identifier(identifier)
)
}
fn opaque_identifier(identifier: &str) -> String {
let digest = Sha256::digest(identifier.as_bytes());
let mut encoded = String::with_capacity(digest.len().saturating_mul(2));
for byte in digest {
use std::fmt::Write as _;
let _ = write!(encoded, "{byte:02x}");
}
encoded
}
fn normalize_identifier(value: Option<&str>) -> Option<&str> {
value.map(str::trim).filter(|value| !value.is_empty())
}
pub(crate) fn codex_account_id_from_headers(headers: &BTreeMap<String, String>) -> Option<&str> {
headers.iter().find_map(|(name, value)| {
name.trim()
.eq_ignore_ascii_case("chatgpt-account-id")
.then_some(value.as_str())
.and_then(|value| normalize_identifier(Some(value)))
})
}
fn account_id_from_header_object(headers: &Map<String, Value>) -> Option<&str> {
headers.iter().find_map(|(name, value)| {
name.trim()
.eq_ignore_ascii_case("chatgpt-account-id")
.then(|| value.as_str())
.flatten()
.and_then(|value| normalize_identifier(Some(value)))
})
}
fn codex_quota_breaker_ttl(quota_metadata: &Value, now_unix_secs: u64) -> (u64, Option<u64>) {
let reset_at = codex_quota_exhaustion_reset_at(quota_metadata, now_unix_secs);
let ttl_seconds = reset_at
.and_then(|reset_at| reset_at.checked_sub(now_unix_secs))
.filter(|ttl| *ttl > 0)
.unwrap_or(CODEX_QUOTA_BREAKER_FALLBACK_TTL_SECONDS)
.clamp(1, CODEX_QUOTA_BREAKER_MAX_TTL_SECONDS);
(ttl_seconds, reset_at)
}
/// Returns the latest reset deadline required for a currently exhausted Codex
/// quota window. It is shared by the distributed breaker and the per-socket
/// retry exclusion so both stop excluding the account at the same time.
pub(crate) fn codex_quota_exhaustion_reset_at(
quota_metadata: &Value,
now_unix_secs: u64,
) -> Option<u64> {
let Some(metadata) = quota_metadata.as_object() else {
return None;
};
let exhausted_windows = ["primary", "secondary"]
.into_iter()
.filter(|prefix| codex_window_is_exhausted(metadata, prefix))
.collect::<Vec<_>>();
let prefixes = if exhausted_windows.is_empty() {
vec!["primary", "secondary"]
} else {
exhausted_windows
};
prefixes
.iter()
.filter_map(|prefix| codex_window_reset_at(metadata, prefix, now_unix_secs))
.filter(|reset_at| *reset_at > now_unix_secs)
.max()
}
fn codex_window_is_exhausted(metadata: &Map<String, Value>, prefix: &str) -> bool {
let used_percent = metadata
.get(&format!("{prefix}_used_percent"))
.and_then(aether_admin::provider::quota::coerce_json_f64);
used_percent.is_some_and(|used_percent| used_percent >= 100.0 - 1e-6)
}
fn codex_window_reset_at(
metadata: &Map<String, Value>,
prefix: &str,
observed_at_unix_secs: u64,
) -> Option<u64> {
metadata
.get(&format!("{prefix}_reset_at"))
.and_then(aether_admin::provider::quota::coerce_json_u64)
.filter(|reset_at| *reset_at > observed_at_unix_secs)
.or_else(|| {
metadata
.get(&format!("{prefix}_reset_after_seconds"))
.and_then(aether_admin::provider::quota::coerce_json_u64)
.and_then(|seconds| observed_at_unix_secs.checked_add(seconds))
})
}
pub(crate) fn log_codex_quota_breaker_install_failure(error: &GatewayError) {
warn!(
event_name = "codex_account_quota_breaker_install_failed",
log_type = "ops",
error = ?error,
"gateway could not install the immediate Codex quota exhaustion breaker"
);
}
pub(crate) fn log_codex_quota_breaker_check_failure(error: &GatewayError) {
warn!(
event_name = "codex_account_quota_breaker_check_failed",
log_type = "ops",
transport = "websocket",
websocket = true,
error = ?error,
"gateway could not check the immediate Codex quota exhaustion breaker; allowing candidate selection"
);
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use super::{
codex_quota_breaker_blocks_candidate, codex_quota_breaker_keys,
codex_quota_breaker_keys_from_report_context, codex_quota_breaker_ttl,
codex_quota_exhaustion_reset_at, install_codex_quota_exhaustion_breaker,
};
use crate::AppState;
#[test]
fn account_scope_is_shared_across_catalog_keys() {
let first = codex_quota_breaker_keys(Some("key-first"), Some("account-123"));
let second = codex_quota_breaker_keys(Some("key-second"), Some("account-123"));
assert_eq!(first.first(), second.first());
assert_ne!(first.last(), second.last());
assert!(!first.first().is_some_and(|key| key.contains("account-123")));
}
#[test]
fn report_context_and_planned_headers_use_the_same_account_scope() {
let report_context = json!({
"key_id": "key-first",
"provider_request_headers": {
"ChatGPT-Account-ID": "account-123"
}
});
let context_keys = codex_quota_breaker_keys_from_report_context(Some(&report_context));
let planned_keys = codex_quota_breaker_keys(Some("key-second"), Some("account-123"));
assert_eq!(context_keys.first(), planned_keys.first());
}
#[test]
fn ttl_uses_the_exhausted_window_reset_deadline() {
let (ttl, reset_at) = codex_quota_breaker_ttl(
&json!({
"primary_used_percent": 100,
"primary_reset_at": 1_000,
"secondary_used_percent": 10,
"secondary_reset_at": 2_000,
}),
500,
);
assert_eq!(ttl, 500);
assert_eq!(reset_at, Some(1_000));
}
#[test]
fn ttl_falls_back_when_an_error_has_no_reset_metadata() {
let (ttl, reset_at) = codex_quota_breaker_ttl(&json!({"allowed": false}), 500);
assert_eq!(ttl, 300);
assert_eq!(reset_at, None);
}
#[test]
fn reset_deadline_uses_relative_metadata_when_absolute_reset_is_stale() {
assert_eq!(
codex_quota_exhaustion_reset_at(
&json!({
"primary_used_percent": 100,
"primary_reset_at": 999,
"primary_reset_after_seconds": 120,
}),
1_000,
),
Some(1_120)
);
}
#[test]
fn account_header_matching_is_case_insensitive() {
let headers =
BTreeMap::from([("CHATGPT-ACCOUNT-ID".to_string(), "account-123".to_string())]);
let keys = codex_quota_breaker_keys(
Some("key-first"),
headers.iter().find_map(|(name, value)| {
name.eq_ignore_ascii_case("chatgpt-account-id")
.then_some(value.as_str())
}),
);
assert_eq!(keys.len(), 2);
}
#[tokio::test]
async fn exhausted_account_immediately_blocks_a_different_catalog_key() {
let state = AppState::new().expect("gateway state should build");
let report_context = json!({
"key_id": "key-first",
"provider_request_headers": {
"ChatGPT-Account-ID": "account-123"
}
});
let quota_metadata = json!({
"allowed": false,
"limit_reached": true,
"primary_used_percent": 100,
"primary_reset_after_seconds": 60,
});
assert!(install_codex_quota_exhaustion_breaker(
&state,
Some(&report_context),
&quota_metadata,
"test",
)
.await
.expect("breaker installation should succeed"));
let same_account_other_key =
BTreeMap::from([("chatgpt-account-id".to_string(), "account-123".to_string())]);
assert!(codex_quota_breaker_blocks_candidate(
&state,
Some("codex"),
Some("key-second"),
&same_account_other_key,
)
.await
.expect("breaker lookup should succeed"));
let other_account =
BTreeMap::from([("chatgpt-account-id".to_string(), "account-456".to_string())]);
assert!(!codex_quota_breaker_blocks_candidate(
&state,
Some("codex"),
Some("key-second"),
&other_account,
)
.await
.expect("breaker lookup should succeed"));
}
}
apps/aether-gateway/src/orchestration/effects.rs
+260 -12
View File
@@ -29,7 +29,8 @@ use tracing::warn;
use super::{
classify_failure_disposition, local_failover_error_message, project_local_adaptive_rate_limit,
project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed,
project_local_key_circuit_failure, project_local_success_health, FailureScope,
project_local_key_circuit_failure, project_local_success_health,
resolve_local_failover_analysis_for_attempt, FailureScope, LocalFailoverAnalysis,
LocalFailoverClassification,
};
use crate::ai_serving::extract_pool_sticky_session_token;
@@ -325,6 +326,41 @@ pub(crate) fn spawn_local_oauth_success_effect(
});
}
/// Inputs for the terminal effects of a failed streaming attempt.
///
/// The status/body are deliberately supplied by the transport-specific caller:
/// a WebSocket terminal event may carry its own status and error body, while a
/// normal stream failure gets them from the HTTP response. Keeping this type at
/// the orchestration boundary prevents each transport from rebuilding the
/// health, adaptive, OAuth, and pool effect sequence independently.
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalStreamFailureEffect<'a> {
pub(crate) status_code: u16,
pub(crate) headers: &'a BTreeMap<String, String>,
pub(crate) response_text: Option<&'a str>,
pub(crate) stream_timeout: bool,
}
impl<'a> LocalStreamFailureEffect<'a> {
pub(crate) const fn new(
status_code: u16,
headers: &'a BTreeMap<String, String>,
response_text: Option<&'a str>,
) -> Self {
Self {
status_code,
headers,
response_text,
stream_timeout: false,
}
}
pub(crate) const fn with_stream_timeout(mut self) -> Self {
self.stream_timeout = true;
self
}
}
struct PoolFeedbackContext {
pool_config: AdminProviderPoolConfig,
sticky_session_token: Option<String>,
@@ -377,36 +413,157 @@ pub(crate) async fn apply_local_execution_effect(
}
LocalExecutionEffect::PoolSuccessSync { payload } => {
record_sync_pool_success_effect(state, context, payload).await;
release_pool_key_lease_effect(state, context).await;
release_local_pool_key_lease(state, context).await;
}
LocalExecutionEffect::PoolSuccessStream { payload } => {
record_stream_pool_success_effect(state, context, payload).await;
release_pool_key_lease_effect(state, context).await;
release_local_pool_key_lease(state, context).await;
}
LocalExecutionEffect::PoolError(effect) => {
record_pool_error_effect(state, context, effect).await;
release_pool_key_lease_effect(state, context).await;
release_local_pool_key_lease(state, context).await;
}
LocalExecutionEffect::PoolStreamTimeout => {
record_pool_stream_timeout_effect(state, context).await;
release_pool_key_lease_effect(state, context).await;
release_local_pool_key_lease(state, context).await;
}
}
}
async fn release_pool_key_lease_effect(state: &AppState, context: LocalExecutionEffectContext<'_>) {
/// Apply the provider/key effects shared by every successful streaming
/// transport. Usage persistence and request-candidate terminal status remain
/// owned by the report layer; this helper only projects execution health,
/// adaptive state, and pool feedback.
pub(crate) async fn apply_local_stream_success_effects(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
payload: &GatewayStreamReportRequest,
) {
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
)
.await;
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::PoolSuccessStream { payload },
)
.await;
}
/// Apply the provider/key effects shared by every failed streaming attempt.
/// The returned analysis is the same failover classification used by the
/// normal stream runtime, allowing the caller to make a transport-specific
/// retry/close decision without re-running policy evaluation.
pub(crate) async fn apply_local_stream_failure_effects(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
effect: LocalStreamFailureEffect<'_>,
) -> LocalFailoverAnalysis {
let analysis = resolve_local_failover_analysis_for_attempt(
state,
context.plan,
context.report_context,
effect.status_code,
effect.response_text,
)
.await;
if effect.stream_timeout {
apply_local_execution_effect(state, context, LocalExecutionEffect::PoolStreamTimeout).await;
}
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: effect.status_code,
classification: analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
status_code: effect.status_code,
classification: analysis.classification,
headers: Some(effect.headers),
}),
)
.await;
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
status_code: effect.status_code,
classification: analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code: effect.status_code,
response_text: effect.response_text,
}),
)
.await;
apply_local_execution_effect(
state,
context,
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
status_code: effect.status_code,
classification: analysis.classification,
headers: effect.headers,
error_body: effect.response_text,
}),
)
.await;
analysis
}
pub(crate) async fn release_local_pool_key_lease(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
) {
let metadata = local_execution_candidate_metadata_from_report_context(context.report_context);
let Some(lease) = metadata.pool_key_lease else {
return;
};
release_pool_key_lease(state, &lease).await;
}
/// Releases a lease carried by a planned-but-not-started report context. This
/// path has no execution plan yet, so it intentionally omits candidate health
/// logging and only performs the distributed lock cleanup.
pub(crate) async fn release_pool_key_lease_from_report_context(
state: &AppState,
report_context: Option<&Value>,
) {
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
let Some(lease) = metadata.pool_key_lease else {
return;
};
release_pool_key_lease(state, &lease).await;
}
async fn release_pool_key_lease(state: &AppState, lease: &aether_runtime_state::RuntimeLockLease) {
if let Err(err) =
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), &lease).await
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), lease).await
{
warn!(
error = ?err,
provider_id = %context.plan.provider_id,
key_id = %context.plan.key_id,
"gateway orchestration effects: failed to release pool key lease"
"gateway orchestration effects: failed to release a planned pool key lease"
);
}
}
@@ -1967,20 +2124,22 @@ mod tests {
use serde_json::{json, Value};
use super::{
apply_local_execution_effect, execution_plan_bearer_matches_transport,
apply_local_execution_effect, apply_local_stream_failure_effects,
apply_local_stream_success_effects, execution_plan_bearer_matches_transport,
local_candidate_failure_should_apply_key_effects,
local_candidate_failure_should_record_pool_error, pool_score_feedback_gate_allows,
pool_score_hard_state_for_status, resolve_pool_feedback_context,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect,
LocalPoolErrorEffect, ProviderKeyEffectLockPool,
LocalPoolErrorEffect, LocalStreamFailureEffect, ProviderKeyEffectLockPool,
};
use crate::data::{GatewayDataConfig, GatewayDataState};
use crate::orchestration::{
apply_local_report_effect, LocalFailoverClassification, LocalReportEffect,
};
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::usage::GatewayStreamReportRequest;
use crate::AppState;
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id,
@@ -2031,6 +2190,22 @@ mod tests {
plan
}
fn sample_stream_report() -> GatewayStreamReportRequest {
GatewayStreamReportRequest {
trace_id: "trace-stream-effects".to_string(),
report_kind: "openai_chat_stream_success".to_string(),
report_context: None,
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: None,
telemetry: None,
}
}
#[test]
fn pool_score_feedback_gate_suppresses_repeated_success_writes() {
super::POOL_SCORE_FEEDBACK_GATE.clear();
@@ -2828,6 +3003,79 @@ mod tests {
.is_some());
}
#[tokio::test]
async fn stream_success_effect_helper_projects_health_and_scheduler_affinity() {
let state = AppState::new().expect("gateway state should build");
let plan = sample_plan();
let report_context = json!({
"api_key_id": "api-key-1",
"client_api_format": "openai:chat",
"model": "gpt-5",
});
let cache_key =
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
let payload = sample_stream_report();
apply_local_stream_success_effects(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: Some(&report_context),
},
&payload,
)
.await;
assert_eq!(
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
Some(SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
key_id: "key-1".to_string(),
})
);
}
#[tokio::test]
async fn stream_failure_effect_helper_returns_analysis_and_projects_health() {
let state = health_state();
let plan = sample_plan();
let headers = BTreeMap::new();
let analysis = apply_local_stream_failure_effects(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: None,
},
LocalStreamFailureEffect::new(503, &headers, Some("upstream unavailable"))
.with_stream_timeout(),
)
.await;
assert_eq!(
analysis.classification,
LocalFailoverClassification::UseDefault
);
assert_eq!(analysis.decision.as_str(), "use_default");
let stored_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
.await
.expect("provider catalog keys should load")
.into_iter()
.next()
.expect("stored key should exist");
assert_eq!(
stored_key
.health_by_format
.as_ref()
.and_then(|value| value.get("openai:chat"))
.and_then(|value| value.get("consecutive_failures"))
.and_then(Value::as_u64),
Some(1)
);
}
#[tokio::test]
async fn configured_stop_pattern_keeps_scheduler_affinity_cache() {
let state = AppState::new().expect("gateway state should build");
apps/aether-gateway/src/orchestration/mod.rs
+15 -3
View File
@@ -7,6 +7,7 @@ use crate::AppState;
mod adaptive;
mod attempt;
mod classifier;
mod codex_quota_breaker;
mod effects;
mod health;
mod oauth_error;
@@ -32,11 +33,20 @@ pub(crate) use self::classifier::{
FailureTokenAction, LocalFailoverClassification, LocalFailoverInput,
LocalTransportFailoverClassification,
};
pub(crate) use self::codex_quota_breaker::{
codex_account_id_from_headers, codex_quota_breaker_blocks_candidate,
codex_quota_exhaustion_reset_at, install_codex_quota_exhaustion_breaker,
log_codex_quota_breaker_check_failure, log_codex_quota_breaker_install_failure,
};
pub(crate) use self::effects::{
apply_local_execution_effect, spawn_local_oauth_success_effect, LocalAdaptiveRateLimitEffect,
LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect,
apply_local_stream_failure_effects, apply_local_stream_success_effects,
release_local_pool_key_lease,
release_pool_key_lease_from_report_context, LocalAdaptiveRateLimitEffect,
LocalStreamFailureEffect,
};
pub(crate) use self::health::{
project_local_failure_health, project_local_key_circuit_closed,
@@ -48,8 +58,9 @@ pub(crate) use self::oauth_error::{
pub(crate) use self::policy::{
append_local_failover_policy_to_value, codex_cyber_flag_passthrough_enabled,
cyber_continue_failover_enabled, local_failover_policy_from_report_context,
local_failover_policy_from_transport, resolve_local_failover_policy, LocalFailoverPolicy,
LocalFailoverRegexRule, CYBER_CONTINUE_FAILOVER_CONFIG_KEY,
local_failover_policy_from_transport, resolve_local_failover_policy,
responses_websocket_adapter, LocalFailoverPolicy, LocalFailoverRegexRule,
ResponsesWebSocketAdapter, CYBER_CONTINUE_FAILOVER_CONFIG_KEY, RESPONSES_WEBSOCKET_CONFIG_KEY,
};
pub(crate) use self::recovery::{
analyze_local_failover, analyze_local_transport_error, apply_provider_failure_disposition,
@@ -59,7 +70,8 @@ pub(crate) use self::recovery::{
#[cfg(test)]
pub(crate) use self::report_effects::clear_local_report_effect_caches_for_tests;
pub(crate) use self::report_effects::{
apply_local_report_effect, store_local_gemini_file_mapping, LocalReportEffect,
apply_local_report_effect, store_local_gemini_file_mapping,
sync_codex_websocket_quota_metadata, LocalReportEffect,
};
pub(crate) async fn resolve_local_failover_analysis_for_attempt(
apps/aether-gateway/src/orchestration/policy.rs
+98 -1
View File
@@ -8,6 +8,7 @@ use crate::provider_transport::GatewayProviderTransportSnapshot;
use crate::AppState;
pub(crate) const CYBER_CONTINUE_FAILOVER_CONFIG_KEY: &str = "cyber_continue_failover";
pub(crate) const RESPONSES_WEBSOCKET_CONFIG_KEY: &str = "responses_websocket";
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct LocalFailoverPolicy {
@@ -297,6 +298,63 @@ pub(crate) fn codex_cyber_flag_passthrough_enabled(
.unwrap_or(true)
}
/// Selects the protocol adapter responsible for one eligible Responses
/// WebSocket upstream. Provider-scoped feature switches remain the source of
/// truth; this enum only identifies provider-specific extensions around the
/// otherwise standard Responses WebSocket protocol.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ResponsesWebSocketAdapter {
/// A provider that speaks the standard OpenAI Responses WebSocket protocol.
Standard,
/// Standard protocol plus Codex account and quota extensions.
Codex,
}
impl ResponsesWebSocketAdapter {
pub(crate) fn supports_provider_type(self, provider_type: &str) -> bool {
match self {
Self::Standard => {
!provider_type.trim().is_empty()
&& !provider_type.trim().eq_ignore_ascii_case("codex")
}
Self::Codex => provider_type.trim().eq_ignore_ascii_case("codex"),
}
}
}
/// Whether a provider explicitly enables the standard Responses WebSocket
/// bridge. The setting is provider-scoped so rollout remains opt-in per
/// verified upstream.
pub(crate) fn responses_websocket_enabled(provider_config: Option<&Value>) -> bool {
provider_config
.and_then(|config| config.get(RESPONSES_WEBSOCKET_CONFIG_KEY))
.and_then(Value::as_object)
.and_then(|responses| responses.get("enabled"))
.and_then(Value::as_bool)
.unwrap_or(false)
}
/// Returns the enabled Responses WebSocket adapter for a provider. The shared
/// protocol bridge remains opt-in, while this resolver isolates provider-only
/// extensions from candidate planning and the session engine.
pub(crate) fn responses_websocket_adapter(
provider_type: &str,
provider_config: Option<&Value>,
) -> Option<ResponsesWebSocketAdapter> {
let provider_type = provider_type.trim();
if provider_type.is_empty() {
return None;
}
if !responses_websocket_enabled(provider_config) {
return None;
}
Some(if provider_type.eq_ignore_ascii_case("codex") {
ResponsesWebSocketAdapter::Codex
} else {
ResponsesWebSocketAdapter::Standard
})
}
fn local_failover_regex_rule_to_value(rule: &LocalFailoverRegexRule) -> Value {
json!({
"pattern": rule.pattern,
@@ -368,7 +426,9 @@ mod tests {
use super::{
append_local_failover_policy_to_value, local_failover_policy_from_report_context,
local_failover_policy_from_transport, LocalFailoverPolicy, LocalFailoverRegexRule,
local_failover_policy_from_transport, responses_websocket_adapter,
responses_websocket_enabled, LocalFailoverPolicy, LocalFailoverRegexRule,
ResponsesWebSocketAdapter,
};
use crate::provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -593,4 +653,41 @@ mod tests {
transport.provider.provider_type = "llm".to_string();
assert!(local_failover_policy_from_transport(&transport).stop_cyber_policy_errors);
}
#[test]
fn responses_websocket_requires_an_explicit_provider_switch() {
assert!(!responses_websocket_enabled(None));
assert!(!responses_websocket_enabled(Some(&json!({
"responses_websocket": {"enabled": false}
}))));
assert!(responses_websocket_enabled(Some(&json!({
"responses_websocket": {"enabled": true}
}))));
assert_eq!(
responses_websocket_adapter(
"custom",
Some(&json!({"responses_websocket": {"enabled": false}})),
),
None
);
assert_eq!(
responses_websocket_adapter(
"custom",
Some(&json!({"responses_websocket": {"enabled": true}})),
),
Some(ResponsesWebSocketAdapter::Standard)
);
assert_eq!(
responses_websocket_adapter(
"codex",
Some(&json!({"responses_websocket": {"enabled": true}})),
),
Some(ResponsesWebSocketAdapter::Codex)
);
assert!(ResponsesWebSocketAdapter::Codex.supports_provider_type("CODEX"));
assert!(!ResponsesWebSocketAdapter::Codex.supports_provider_type("openai"));
assert!(ResponsesWebSocketAdapter::Standard.supports_provider_type("custom"));
assert!(!ResponsesWebSocketAdapter::Standard.supports_provider_type("codex"));
}
}
apps/aether-gateway/src/orchestration/report_effects.rs
+222 -7
View File
@@ -16,12 +16,16 @@ use serde_json::{json, Value};
use tracing::warn;
use uuid::Uuid;
use super::codex_quota_breaker::{
install_codex_quota_exhaustion_breaker, log_codex_quota_breaker_install_failure,
};
use crate::clock::current_unix_secs;
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::log_ids::short_request_id;
use crate::{AppState, GatewayError};
const RUNTIME_METADATA_CAS_MAX_ATTEMPTS: usize = 16;
const CODEX_WEBSOCKET_RATE_LIMITS_REPORT_CONTEXT_FIELD: &str = "codex_websocket_rate_limits";
static GROK_CHINESE_WAIT_DURATION_RE: OnceLock<Regex> = OnceLock::new();
static GROK_ENGLISH_WAIT_DURATION_RE: OnceLock<Regex> = OnceLock::new();
@@ -95,6 +99,12 @@ fn report_context_provider_response_headers(
(!out.is_empty()).then_some(out)
}
fn codex_websocket_quota_from_report_context(report_context: Option<&Value>) -> Option<Value> {
report_context
.and_then(|context| context.get(CODEX_WEBSOCKET_RATE_LIMITS_REPORT_CONTEXT_FIELD))
.filter(|value| value.as_object().is_some_and(|object| !object.is_empty()))
.cloned()
}
fn merge_metadata_object(
current: Option<&Value>,
section_key: &str,
@@ -257,6 +267,35 @@ fn gemini_cli_credits_from_stream_payload(
latest
}
fn codex_websocket_quota_from_stream_payload(
payload: &GatewayStreamReportRequest,
now_unix_secs: u64,
) -> Option<Value> {
let body_base64 = payload.provider_body_base64.as_deref()?;
let body = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.ok()?;
let text = std::str::from_utf8(&body).ok()?;
let mut latest = None::<Value>;
for raw_line in text.lines() {
let line = raw_line.trim_matches('\r').trim();
let data = line.strip_prefix("data:").map(str::trim).unwrap_or(line);
if data.is_empty() || data == "[DONE]" || data.starts_with(':') {
continue;
}
let Ok(value) = serde_json::from_str::<Value>(data) else {
continue;
};
if let Some(quota) = admin_provider_quota_pure::parse_codex_websocket_rate_limits_response(
&value,
now_unix_secs,
) {
latest = Some(quota);
}
}
latest
}
async fn sync_gemini_cli_credits_from_report(
state: &AppState,
report_context: Option<&Value>,
@@ -572,7 +611,24 @@ async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncR
}
async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStreamReportRequest) {
if (200..300).contains(&payload.status_code) {
let websocket_quota_seen = match sync_codex_websocket_quota_from_stream_payload(state, payload)
.await
{
Ok(Some(_)) => true,
Ok(None) => false,
Err(err) => {
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist Codex realtime quota from WebSocket response body"
);
false
}
};
if !websocket_quota_seen && (200..300).contains(&payload.status_code) {
if let Err(err) = sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
@@ -762,11 +818,6 @@ async fn sync_codex_quota_from_response_headers(
report_context: Option<&Value>,
headers: &BTreeMap<String, String>,
) -> Result<bool, GatewayError> {
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
let now_unix_secs = current_unix_secs();
let observed_at_unix_secs = report_context_u64(
report_context,
@@ -791,6 +842,97 @@ async fn sync_codex_quota_from_response_headers(
}) else {
return Ok(false);
};
sync_codex_quota_metadata_with_observation(
state,
report_context,
parsed,
"response_headers",
observed_at_unix_secs,
request_started_at_unix_ms,
request_order_id,
observed_reset_generation,
observed_credential_generation,
)
.await
}
async fn sync_codex_websocket_quota_from_stream_payload(
state: &AppState,
payload: &GatewayStreamReportRequest,
) -> Result<Option<bool>, GatewayError> {
let now_unix_secs = current_unix_secs();
let parsed = codex_websocket_quota_from_report_context(payload.report_context.as_ref())
.or_else(|| codex_websocket_quota_from_stream_payload(payload, now_unix_secs));
let Some(parsed) = parsed else {
return Ok(None);
};
Ok(Some(
sync_codex_quota_metadata(
state,
payload.report_context.as_ref(),
parsed,
"websocket_response_body",
)
.await?,
))
}
async fn sync_codex_quota_metadata(
state: &AppState,
report_context: Option<&Value>,
parsed: Value,
source: &'static str,
) -> Result<bool, GatewayError> {
let observed_at_unix_secs = parsed
.get("updated_at")
.and_then(admin_provider_quota_pure::coerce_json_u64)
.filter(|value| *value > 0)
.unwrap_or_else(current_unix_secs);
let request_started_at_unix_ms =
report_context_u64(report_context, "provider_request_started_at_unix_ms");
let request_order_id = report_context_string(report_context, "provider_request_order_id");
let observed_reset_generation =
report_context_u64(report_context, "codex_quota_reset_generation");
let observed_credential_generation =
report_context_string(report_context, "codex_credential_generation");
sync_codex_quota_metadata_with_observation(
state,
report_context,
parsed,
source,
observed_at_unix_secs,
request_started_at_unix_ms,
request_order_id,
observed_reset_generation,
observed_credential_generation,
)
.await
}
#[allow(clippy::too_many_arguments)]
async fn sync_codex_quota_metadata_with_observation(
state: &AppState,
report_context: Option<&Value>,
parsed: Value,
source: &'static str,
observed_at_unix_secs: u64,
request_started_at_unix_ms: Option<u64>,
request_order_id: Option<&str>,
observed_reset_generation: Option<u64>,
observed_credential_generation: Option<&str>,
) -> Result<bool, GatewayError> {
if admin_provider_quota_pure::codex_rate_limit_metadata_exhausted(&parsed) {
if let Err(error) =
install_codex_quota_exhaustion_breaker(state, report_context, &parsed, source).await
{
log_codex_quota_breaker_install_failure(&error);
}
}
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
// Runtime headers can be partial (for example only the primary window),
// so absence never authoritatively removes another stored window.
let coverage = admin_provider_quota_pure::CodexQuotaWindowCoverage::Patch;
@@ -845,7 +987,7 @@ async fn sync_codex_quota_from_response_headers(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"response_headers",
source,
);
let updated = state
.update_provider_catalog_key_runtime_metadata(
@@ -872,6 +1014,14 @@ async fn sync_codex_quota_from_response_headers(
Ok(false)
}
pub(crate) async fn sync_codex_websocket_quota_metadata(
state: &AppState,
report_context: Option<&Value>,
parsed: Value,
) -> Result<bool, GatewayError> {
sync_codex_quota_metadata(state, report_context, parsed, "websocket_response_body").await
}
#[cfg(test)]
pub(crate) fn clear_local_report_effect_caches_for_tests() {}
@@ -1104,6 +1254,33 @@ mod tests {
assert_eq!(status["quota"]["provider_type"], json!("gemini_cli"));
}
#[test]
fn codex_quota_snapshot_match_requires_explicit_signal_projection() {
let parsed = json!({
"allowed": false,
"limit_reached": true,
});
assert!(!codex_quota_snapshot_matches_metadata(
Some(&json!({
"quota": {
"exhausted": true
}
})),
&parsed,
));
assert!(codex_quota_snapshot_matches_metadata(
Some(&json!({
"quota": {
"exhausted": true,
"allowed": false,
"limit_reached": true
}
})),
&parsed,
));
}
#[test]
fn grok_quota_feedback_decrements_the_matching_window() {
let mut bucket = json!({
@@ -1290,4 +1467,42 @@ mod tests {
None
);
}
#[test]
fn codex_quota_snapshot_does_not_roll_back_exhaustion() {
let current = json!({
"allowed": false,
"limit_reached": true,
"primary_used_percent": 100.0,
"primary_reset_at": 2_000,
"updated_at": 100
});
let delayed = json!({
"allowed": true,
"limit_reached": false,
"primary_used_percent": 99.0,
"primary_reset_at": 2_000,
"updated_at": 101
});
assert!(codex_snapshot_regresses(&current, &delayed));
}
#[test]
fn codex_quota_snapshot_allows_a_new_reset_window() {
let current = json!({
"allowed": false,
"limit_reached": true,
"primary_used_percent": 100.0,
"primary_reset_at": 2_000,
"updated_at": 100
});
let refreshed = json!({
"allowed": true,
"limit_reached": false,
"primary_used_percent": 1.0,
"primary_reset_at": 1_000,
"updated_at": 101
});
assert!(!codex_snapshot_regresses(&current, &refreshed));
}
}