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

fix: use codex quota refresh for account checks

Browse Source
This commit is contained in:
fawney19
2026-05-16 20:16:23 +08:00
parent 56994d4c29
commit 3a5922d4ee
16 changed files with 150 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
apps/aether-gateway/src/admin_api.rs
View File

@@ -3,10 +3,9 @@ pub(crate) use crate::handlers::admin::{
build_internal_control_error_response, create_provider_oauth_catalog_key,
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
maybe_build_local_admin_response, persist_provider_quota_refresh_state,
provider_account_self_check_endpoint_for_provider,
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_account_self_check,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
reconcile_admin_fixed_provider_template_endpoints,
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
update_existing_provider_oauth_catalog_key, AdminAppState,
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,

3
apps/aether-gateway/src/handlers/admin/mod.rs
View File

@@ -27,8 +27,7 @@ pub(crate) use self::provider::oauth::provisioning::{
};
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::{
persist_provider_quota_refresh_state, provider_account_self_check_endpoint_for_provider,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_account_self_check,
persist_provider_quota_refresh_state, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh,
};
pub(crate) use self::provider::oauth::runtime::{

45
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
View File

@@ -8,7 +8,8 @@ use self::invalid::{
codex_structured_invalid_reason,
};
use self::parse::{
parse_codex_backend_me_response, parse_codex_usage_headers, parse_codex_wham_usage_response,
build_codex_quota_exhausted_fallback_metadata, parse_codex_usage_headers,
parse_codex_wham_usage_response,
};
use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
use super::shared::{
@@ -110,7 +111,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("backend-api/me 请求执行失败: {detail}"),
"message": format!("wham/usage 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
@@ -137,9 +138,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
.as_ref()
.and_then(|body| body.json_body.as_ref())
{
if let Some(parsed) = parse_codex_backend_me_response(body_json, now_unix_secs)
.or_else(|| parse_codex_wham_usage_response(body_json, now_unix_secs))
{
if let Some(parsed) = parse_codex_wham_usage_response(body_json, now_unix_secs) {
metadata_update = Some(json!({
"codex": merge_codex_quota_metadata(header_metadata.as_ref(), &parsed)
}));
@@ -152,21 +151,21 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
status = "success".to_string();
} else {
status = "no_metadata".to_string();
message = Some("backend-api/me 响应中未包含账号信息".to_string());
message = Some("响应中未包含限额信息".to_string());
}
} else {
message = Some("无法解析 backend-api/me API 响应".to_string());
message = Some("无法解析 wham/usage API 响应".to_string());
}
} else {
let err_msg = extract_execution_error_message(&result);
message = Some(match err_msg.as_deref() {
Some(detail) if !detail.is_empty() => {
format!(
"backend-api/me API 返回状态码 {}: {}",
"wham/usage API 返回状态码 {}: {}",
result.status_code, detail
)
}
_ => format!("backend-api/me API 返回状态码 {}", result.status_code),
_ => format!("wham/usage API 返回状态码 {}", result.status_code),
});
match result.status_code {
@@ -223,14 +222,26 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
oauth_invalid_reason = reason;
status = "workspace_deactivated".to_string();
} else {
let (at, reason) = codex_build_invalid_state(
&key,
codex_structured_invalid_reason(402, err_msg.as_deref()),
now_unix_secs,
);
oauth_invalid_at_unix_secs = at;
oauth_invalid_reason = reason;
status = "payment_required".to_string();
let plan_type = transport
.key
.decrypted_auth_config
.as_deref()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok())
.and_then(|value| {
value
.get("plan_type")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned)
});
metadata_update = Some(json!({
"codex": build_codex_quota_exhausted_fallback_metadata(
plan_type.as_deref(),
now_unix_secs,
)
}));
(oauth_invalid_at_unix_secs, oauth_invalid_reason) =
quota_refresh_success_invalid_state(&key);
status = "quota_exhausted".to_string();
}
}
403 => {

7
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/parse.rs
View File

@@ -22,13 +22,6 @@ pub(super) fn parse_codex_wham_usage_response(
admin_provider_quota_pure::parse_codex_wham_usage_response(value, updated_at_unix_secs)
}
pub(super) fn parse_codex_backend_me_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
admin_provider_quota_pure::parse_codex_backend_me_response(value, updated_at_unix_secs)
}
pub(super) fn parse_codex_usage_headers(
headers: &BTreeMap<String, String>,
updated_at_unix_secs: u64,

16
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
View File

@@ -84,22 +84,6 @@ pub(crate) fn provider_quota_refresh_endpoint_for_provider(
)
}
pub(crate) fn provider_type_supports_account_self_check(provider_type: &str) -> bool {
ProviderPoolService::with_builtin_adapters().supports_account_self_check(provider_type)
}
pub(crate) fn provider_account_self_check_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
ProviderPoolService::with_builtin_adapters().account_self_check_endpoint_for_provider(
provider_type,
endpoints,
include_inactive,
)
}
pub(crate) fn provider_quota_refresh_missing_endpoint_message(provider_type: &str) -> String {
ProviderPoolService::with_builtin_adapters()
.quota_refresh_missing_endpoint_message(provider_type)

2
apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs
View File

@@ -226,7 +226,7 @@ pub(crate) struct AdminProviderUpdateRequest {
pub(crate) type AdminProviderUpdatePatch = AdminTypedObjectPatch<AdminProviderUpdateRequest>;
pub(crate) const CODEX_WHAM_USAGE_URL: &str = "https://chatgpt.com/backend-api/me";
pub(crate) const CODEX_WHAM_USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";
pub(crate) const KIRO_USAGE_LIMITS_PATH: &str = "/getUsageLimits";
pub(crate) const KIRO_USAGE_SDK_VERSION: &str = "1.0.0";
pub(crate) const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";

8
apps/aether-gateway/src/maintenance/runtime/account_self_check.rs
View File

@@ -14,8 +14,8 @@ use serde_json::{json, Value};
use tracing::{debug, info, warn};
use crate::admin_api::{
admin_provider_pool_config, provider_account_self_check_endpoint_for_provider,
provider_type_supports_account_self_check, refresh_provider_pool_quota_locally, AdminAppState,
admin_provider_pool_config, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh, refresh_provider_pool_quota_locally, AdminAppState,
};
use crate::{AppState, GatewayError};
@@ -547,7 +547,7 @@ fn endpoint_for_self_check(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
provider_account_self_check_endpoint_for_provider(provider_type, endpoints, true)
provider_quota_refresh_endpoint_for_provider(provider_type, endpoints, true)
}
fn gateway_error_message(err: GatewayError) -> String {
@@ -637,7 +637,7 @@ pub(crate) async fn perform_account_self_check_once_with_config(
summary.providers_skipped = summary.providers_skipped.saturating_add(1);
continue;
};
if !provider_type_supports_account_self_check(&provider_type) {
if !provider_type_supports_quota_refresh(&provider_type) {
summary.providers_skipped = summary.providers_skipped.saturating_add(1);
continue;
}

123
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
View File

@@ -29,7 +29,6 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
accept: String,
provider_api_format: String,
total_ms: Option<u64>,
}
@@ -69,7 +68,6 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.get("authorization")
.cloned()
.unwrap_or_default(),
accept: plan.headers.get("accept").cloned().unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
total_ms: plan
.timeouts
@@ -80,22 +78,41 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
headers: BTreeMap::from([
(
"x-codex-primary-reset-after-seconds".to_string(),
"18000".to_string(),
),
(
"x-codex-primary-reset-at".to_string(),
"1900000000".to_string(),
),
(
"x-codex-secondary-reset-after-seconds".to_string(),
"604800".to_string(),
),
(
"x-codex-secondary-reset-at".to_string(),
"1900500000".to_string(),
),
]),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"user": {
"id": "user-codex-123",
"email": "codex@example.com",
"name": "Codex User"
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": 12.5,
"window_minutes": 300
},
"secondary_window": {
"used_percent": 55.0,
"window_minutes": 10080
}
},
"account": {
"id": "acct-codex-123",
"name": "Personal",
"plan_type": "plus"
},
"plan": {
"type": "Plus",
"title": "ChatGPT Plus"
"credits": {
"has_credits": true,
"balance": 42.0,
"unlimited": false
}
})),
body_bytes_b64: None,
@@ -167,22 +184,18 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
);
assert_eq!(payload["results"][0]["quota_snapshot"]["plan_type"], "plus");
assert_eq!(
payload["results"][0]["quota_snapshot"]["exhausted"],
json!(false)
payload["results"][0]["quota_snapshot"]["reset_at"],
1_900_000_000u64
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["credits"]["balance"],
json!(42.0)
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["windows"]
.as_array()
.map(Vec::len),
Some(0usize)
);
assert_eq!(
payload["results"][0]["metadata"]["email"],
"codex@example.com"
);
assert_eq!(
payload["results"][0]["metadata"]["account_id"],
"acct-codex-123"
Some(2usize)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -193,13 +206,12 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/me"
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer sk-codex-123"
);
assert_eq!(seen_execution_runtime_request.accept, "application/json");
assert_eq!(
seen_execution_runtime_request.provider_api_format,
"openai:responses"
@@ -225,23 +237,33 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("email")),
Some(&json!("codex@example.com"))
.and_then(|value| value.get("primary_used_percent")),
Some(&json!(55.0))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("account_id")),
Some(&json!("acct-codex-123"))
.and_then(|value| value.get("primary_reset_at")),
Some(&json!(1_900_500_000u64))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("secondary_used_percent")),
Some(&json!(12.5))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("secondary_reset_at")),
Some(&json!(1_900_000_000u64))
);
assert!(reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("primary_used_percent"))
.is_none());
gateway_handle.abort();
execution_runtime_handle.abort();
@@ -249,7 +271,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
}
#[tokio::test]
async fn gateway_marks_codex_key_invalid_when_backend_me_returns_payment_required() {
async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_required() {
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-codex/refresh-quota",
any(move |_request: Request| async move {
@@ -337,18 +359,31 @@ async fn gateway_marks_codex_key_invalid_when_backend_me_returns_payment_require
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 0);
assert_eq!(payload["failed"], 1);
assert_eq!(payload["results"][0]["status"], "payment_required");
assert_eq!(payload["results"][0]["status"], "quota_exhausted");
assert_eq!(payload["results"][0]["status_code"], 402);
assert_eq!(
payload["results"][0]["quota_snapshot"]["provider_type"],
"codex"
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["exhausted"],
json!(true)
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert!(reloaded[0].oauth_invalid_at_unix_secs.is_some());
assert_eq!(reloaded[0].oauth_invalid_at_unix_secs, None);
assert_eq!(reloaded[0].oauth_invalid_reason, None);
assert_eq!(
reloaded[0].oauth_invalid_reason.as_deref(),
Some("[ACCOUNT_BLOCK] payment required")
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("primary_used_percent")),
Some(&json!(100.0))
);
gateway_handle.abort();
@@ -1345,7 +1380,7 @@ async fn gateway_reports_codex_quota_runtime_failures_locally_without_falling_ba
assert!(payload["results"][0]["message"]
.as_str()
.expect("message should be string")
.contains("backend-api/me 请求执行失败: execution runtime returned HTTP 500"));
.contains("wham/usage 请求执行失败: execution runtime returned HTTP 500"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository

14
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -2454,9 +2454,8 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
.as_str()
.expect("account_state_recheck_error should be string when recheck is attempted");
assert!(
account_state_recheck_error == "backend-api/me API 返回状态码 401"
|| account_state_recheck_error == "backend-api/me API 返回状态码 403"
|| account_state_recheck_error.starts_with("backend-api/me 请求执行失败:"),
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -4926,9 +4925,8 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.as_str()
.expect("account_state_recheck_error should be string when attempted");
assert!(
account_state_recheck_error == "backend-api/me API 返回状态码 401"
|| account_state_recheck_error == "backend-api/me API 返回状态码 403"
|| account_state_recheck_error.starts_with("backend-api/me 请求执行失败:"),
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
} else {
@@ -4946,7 +4944,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/me"
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
@@ -4970,7 +4968,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.expect("refreshed api key should decrypt");
assert_eq!(decrypted_api_key, "refreshed-codex-access-token");
if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "backend-api/me API 返回状态码 401"
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 401"
{
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert_eq!(