mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 08:27:46 +08:00
fix oauth identity and codex reset credits
This commit is contained in:
@@ -570,6 +570,50 @@ pub(crate) async fn reserve_codex_account_reset(
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn record_locally_consumed_codex_reset_credit(
|
||||
codex: &mut serde_json::Map<String, serde_json::Value>,
|
||||
observed_at_unix_secs: u64,
|
||||
) {
|
||||
let Some(reset_credits) = codex
|
||||
.get_mut("reset_credits")
|
||||
.and_then(serde_json::Value::as_object_mut)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
let Some(available_count) = reset_credits
|
||||
.get("available_count")
|
||||
.and_then(admin_provider_quota_pure::coerce_json_u64)
|
||||
else {
|
||||
return;
|
||||
};
|
||||
|
||||
reset_credits.insert(
|
||||
"available_count".to_string(),
|
||||
serde_json::json!(available_count.saturating_sub(1)),
|
||||
);
|
||||
reset_credits.insert(
|
||||
"updated_at".to_string(),
|
||||
serde_json::json!(observed_at_unix_secs),
|
||||
);
|
||||
reset_credits.insert(
|
||||
"detail_source".to_string(),
|
||||
serde_json::json!("local_consume"),
|
||||
);
|
||||
reset_credits.insert(
|
||||
"detail_status".to_string(),
|
||||
serde_json::json!("pending_refresh"),
|
||||
);
|
||||
reset_credits.remove("detail_error");
|
||||
if let Some(credits) = reset_credits
|
||||
.get_mut("credits")
|
||||
.and_then(serde_json::Value::as_array_mut)
|
||||
{
|
||||
if !credits.is_empty() {
|
||||
credits.remove(0);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn complete_codex_account_reset(
|
||||
state: &AdminAppState<'_>,
|
||||
key_id: &str,
|
||||
@@ -635,6 +679,9 @@ pub(crate) async fn complete_codex_account_reset(
|
||||
generation: reservation.generation,
|
||||
outcome: outcome.to_string(),
|
||||
};
|
||||
if outcome == "reset" {
|
||||
record_locally_consumed_codex_reset_credit(&mut codex, fence_unix_ms / 1_000);
|
||||
}
|
||||
codex_reset_write_bounded_history(&mut codex, &terminal);
|
||||
if codex_reset_reservation_from_object(&codex).as_ref() == Some(reservation) {
|
||||
codex.remove(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY);
|
||||
@@ -1671,7 +1718,19 @@ mod tests {
|
||||
.expect("key should build");
|
||||
key.encrypted_auth_config = Some("auth-v1".to_string());
|
||||
key.upstream_metadata = Some(json!({
|
||||
"codex": {"credential_generation": "credential-v1"}
|
||||
"codex": {
|
||||
"credential_generation": "credential-v1",
|
||||
"reset_credits": {
|
||||
"available_count": 2,
|
||||
"updated_at": 100u64,
|
||||
"detail_source": "wham_readonly",
|
||||
"detail_status": "available",
|
||||
"credits": [
|
||||
{"id": "credit-1", "expires_at": 20_000u64},
|
||||
{"id": "credit-2", "expires_at": 30_000u64}
|
||||
]
|
||||
}
|
||||
}
|
||||
}));
|
||||
let credential = ProviderCatalogKeyOAuthCredentialFence {
|
||||
encrypted_api_key: None,
|
||||
@@ -1926,6 +1985,11 @@ mod tests {
|
||||
codex["account_quota_reset_history"][0]["outcome"],
|
||||
json!("reset")
|
||||
);
|
||||
assert_eq!(codex["reset_credits"]["available_count"], json!(1u64));
|
||||
assert_eq!(
|
||||
codex["reset_credits"]["credits"],
|
||||
json!([{"id": "credit-2", "expires_at": 30_000u64}])
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -4,8 +4,9 @@ use super::super::errors::{
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
|
||||
use aether_contracts::ProxySnapshot;
|
||||
use aether_oauth::provider::providers::{
|
||||
ClaudeCodeProviderOAuthAdapter, GenericProviderOAuthAdapter, CLAUDE_CODE_PROVIDER_TYPE,
|
||||
CLAUDE_CODE_TOKEN_URL, CLAUDE_CODE_WEB_BASE_URL,
|
||||
AntigravityProviderOAuthAdapter, ClaudeCodeProviderOAuthAdapter, GenericProviderOAuthAdapter,
|
||||
ANTIGRAVITY_USER_INFO_URL, CLAUDE_CODE_PROVIDER_TYPE, CLAUDE_CODE_TOKEN_URL,
|
||||
CLAUDE_CODE_WEB_BASE_URL,
|
||||
};
|
||||
use aether_oauth::provider::{
|
||||
ProviderOAuthCookieAuthorizationInput, ProviderOAuthService, ProviderOAuthTransportContext,
|
||||
@@ -43,7 +44,14 @@ fn provider_oauth_exchange_context(
|
||||
fn provider_oauth_service_for_template(
|
||||
template: AdminProviderOAuthTemplate,
|
||||
token_url: String,
|
||||
antigravity_user_info_url: String,
|
||||
) -> Result<ProviderOAuthService, Response<Body>> {
|
||||
if template.provider_type.eq_ignore_ascii_case("antigravity") {
|
||||
let adapter = AntigravityProviderOAuthAdapter::default()
|
||||
.with_token_url_override(token_url)
|
||||
.with_user_info_url_override(antigravity_user_info_url);
|
||||
return Ok(ProviderOAuthService::new().with_adapter(Arc::new(adapter)));
|
||||
}
|
||||
GenericProviderOAuthAdapter::for_provider_type(template.provider_type)
|
||||
.map(|adapter| adapter.with_token_url_override(token_url))
|
||||
.map(|adapter| ProviderOAuthService::new().with_adapter(Arc::new(adapter)))
|
||||
@@ -75,7 +83,10 @@ pub(crate) async fn exchange_admin_provider_oauth_code(
|
||||
proxy: Option<ProxySnapshot>,
|
||||
) -> Result<serde_json::Value, Response<Body>> {
|
||||
let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url);
|
||||
let service = provider_oauth_service_for_template(template, token_url)?;
|
||||
let antigravity_user_info_url =
|
||||
state.provider_oauth_token_url("antigravity_user_info", ANTIGRAVITY_USER_INFO_URL);
|
||||
let service =
|
||||
provider_oauth_service_for_template(template, token_url, antigravity_user_info_url)?;
|
||||
let ctx = provider_oauth_exchange_context(template.provider_type, proxy);
|
||||
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
|
||||
let result = service
|
||||
@@ -103,7 +114,10 @@ pub(crate) async fn exchange_admin_provider_oauth_refresh_token(
|
||||
proxy: Option<ProxySnapshot>,
|
||||
) -> Result<serde_json::Value, Response<Body>> {
|
||||
let token_url = state.provider_oauth_token_url(template.provider_type, template.token_url);
|
||||
let service = provider_oauth_service_for_template(template, token_url)?;
|
||||
let antigravity_user_info_url =
|
||||
state.provider_oauth_token_url("antigravity_user_info", ANTIGRAVITY_USER_INFO_URL);
|
||||
let service =
|
||||
provider_oauth_service_for_template(template, token_url, antigravity_user_info_url)?;
|
||||
let ctx = provider_oauth_exchange_context(template.provider_type, proxy);
|
||||
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
|
||||
let input = aether_oauth::provider::ProviderOAuthImportInput {
|
||||
|
||||
Reference in New Issue
Block a user