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

merge(main): sync latest main into security branch

Browse Source
This commit is contained in:
elky
2026-09-05 00:30:16 +08:00
parent 1c89b5f9ab 27b0381a9a
commit 0e3bd7eff4
40 changed files with 2513 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+62
View File
@@ -5,6 +5,7 @@ use super::shared::{
quota_key_auto_removed, quota_refresh_success_invalid_state,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::AdminImportProviderModelsRequest;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::{
@@ -22,6 +23,63 @@ use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
fn antigravity_discovered_model_ids(metadata_update: Option<&serde_json::Value>) -> Vec<String> {
metadata_update
.and_then(|value| value.pointer("/antigravity/quota_by_model"))
.and_then(serde_json::Value::as_object)
.into_iter()
.flat_map(|models| models.keys())
.map(String::as_str)
.filter(|model_id| aether_model_fetch::antigravity_model_id_is_routable(model_id))
.map(ToOwned::to_owned)
.collect()
}
async fn sync_antigravity_discovered_models(
state: &AdminAppState<'_>,
provider_id: &str,
metadata_update: Option<&serde_json::Value>,
) {
if !state.has_global_model_data_reader() || !state.has_global_model_data_writer() {
return;
}
let model_ids = antigravity_discovered_model_ids(metadata_update);
if model_ids.is_empty() {
return;
}
let result = state
.build_admin_import_provider_models_payload(
provider_id,
AdminImportProviderModelsRequest {
model_ids,
tiered_pricing: None,
price_per_request: None,
},
)
.await;
match result {
Ok(payload) => {
let errors = payload
.get("errors")
.and_then(serde_json::Value::as_array)
.map(Vec::len)
.unwrap_or(0);
if errors > 0 {
warn!(
provider_id,
errors, "Antigravity discovered-model catalog sync completed with item errors"
);
}
}
Err(error) => warn!(
provider_id,
error = %error,
"Antigravity discovered-model catalog sync failed"
),
}
}
async fn execute_antigravity_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
@@ -321,6 +379,10 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
continue;
}
if status == "success" {
sync_antigravity_discovered_models(state, &provider.id, metadata_update.as_ref()).await;
}
if status == "success" {
success_count += 1;
} else {
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+73 -5
View File
@@ -606,6 +606,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,
@@ -671,6 +715,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);
@@ -1972,7 +2019,19 @@ mod tests {
.expect("key should build");
key.encrypted_auth_config = Some(encrypted_auth_config.clone());
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_fence = ProviderTransportCredentialFence {
encrypted_auth_config,
@@ -2070,6 +2129,13 @@ mod tests {
let key_id = "key-codex-reset-credential-generation";
let (app, repository, credential) = codex_reset_state_machine_test_state(key_id);
let admin_state = AdminAppState::new(&app);
let original_metadata = repository
.list_keys_by_ids(&[key_id.to_string()])
.await
.expect("key should load before reservation")
.pop()
.expect("key should exist before reservation")
.upstream_metadata;
let result = reserve_codex_account_reset(
&admin_state,
@@ -2093,10 +2159,7 @@ mod tests {
.expect("key should reload")
.pop()
.expect("key should exist");
assert_eq!(
stored.upstream_metadata.unwrap()["codex"],
json!({"credential_generation":"credential-v1"})
);
assert_eq!(stored.upstream_metadata, original_metadata);
}
#[tokio::test]
@@ -2231,6 +2294,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}])
);
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/state/exchange.rs
+18 -4
View File
@@ -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,
@@ -39,7 +40,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)))
@@ -71,7 +79,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
@@ -99,7 +110,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 {