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

fix: preserve provider auth metadata

Browse Source
This commit is contained in:
elky
2026-06-20 22:11:42 +08:00
parent 1123095bb7
commit f1c30439ff
51 changed files with 2098 additions and 420 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/observability/monitoring/cache_payloads.rs
+10 -1
View File
@@ -1,4 +1,5 @@
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_auth_config_uses_header_authorization;
use aether_crypto::decrypt_python_fernet_ciphertext;
#[cfg(test)]
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
@@ -28,7 +29,15 @@ pub(super) fn admin_monitoring_masked_provider_key_prefix(
) -> Option<String> {
match key.auth_type.trim() {
"service_account" | "vertex_ai" => Some("[Service Account]".to_string()),
"oauth" => Some("[OAuth Token]".to_string()),
"oauth" => {
if provider_key_auth_config_uses_header_authorization(
state.parse_catalog_auth_config_json(key).as_ref(),
) {
Some("[OAuth Header]".to_string())
} else {
Some("[OAuth Token]".to_string())
}
}
_ => {
let full_key = key
.encrypted_api_key
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
+16 -2
View File
@@ -1,6 +1,7 @@
use super::super::helpers::admin_provider_oauth_key_name_from_auth_config;
use super::super::token_import::{
build_provider_access_token_import_auth_config, provider_type_supports_access_token_import,
build_provider_access_token_import_auth_config, decode_access_token_expires_at,
provider_oauth_import_authorization_bearer_token, provider_type_supports_access_token_import,
};
use super::kiro_import::execute_admin_provider_oauth_kiro_batch_import;
use super::parse::{
@@ -394,9 +395,22 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
let AdminProviderOAuthResolvedBatchImport {
access_token,
mut auth_config,
expires_at,
mut expires_at,
} = resolved_import;
apply_admin_provider_oauth_batch_import_hints(provider_type, entry, &mut auth_config);
if let Some(header_access_token) =
provider_oauth_import_authorization_bearer_token(entry.request_headers.as_ref())
{
if let Some(header_expires_at) =
decode_access_token_expires_at(&header_access_token).or(entry.expires_at)
{
auth_config.insert("expires_at".to_string(), json!(header_expires_at));
expires_at = Some(header_expires_at);
} else {
auth_config.remove("expires_at");
expires_at = None;
}
}
let duplicate =
match find_duplicate_provider_oauth_key(state, provider_id, &auth_config, None).await {
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/parse.rs
+138 -2
View File
@@ -1,4 +1,8 @@
use super::super::token_import::{import_tokens_from_raw_token, normalize_provider_import_tokens};
use super::super::token_import::{
import_tokens_from_raw_token, normalize_provider_import_tokens,
normalize_provider_oauth_import_headers_from_object,
provider_oauth_import_authorization_bearer_token,
};
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
use crate::handlers::admin::provider::oauth::state::{current_unix_secs, json_u64_value};
use axum::{
@@ -9,6 +13,7 @@ use axum::{
};
use serde::Deserialize;
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone, Deserialize)]
@@ -22,6 +27,7 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
pub parse_error: Option<String>,
pub refresh_token: Option<String>,
pub access_token: Option<String>,
pub export_access_token: Option<String>,
pub raw_credentials: Option<serde_json::Value>,
pub expires_at: Option<u64>,
pub account_id: Option<String>,
@@ -37,6 +43,7 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
pub sso_rw_token: Option<String>,
pub cf_cookies: Option<String>,
pub cf_clearance: Option<String>,
pub request_headers: Option<BTreeMap<String, String>>,
pub user_agent: Option<String>,
pub browser_profile: Option<String>,
}
@@ -183,6 +190,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
parse_error: None,
refresh_token,
access_token,
export_access_token: None,
raw_credentials: None,
expires_at: None,
account_id: None,
@@ -198,6 +206,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
sso_rw_token: grok_cookie_value(raw_token, "sso-rw"),
cf_cookies: grok_cookie_profile(raw_token),
cf_clearance: grok_cookie_value(raw_token, "cf_clearance"),
request_headers: None,
user_agent: None,
browser_profile: None,
})
@@ -237,10 +246,24 @@ fn extract_admin_provider_oauth_batch_import_entry(
.as_deref()
.and_then(|cookie| grok_cookie_value(cookie, "sso"))
});
let request_headers = normalize_provider_oauth_import_headers_from_object(object);
let header_bearer_token =
provider_oauth_import_authorization_bearer_token(request_headers.as_ref());
let export_access_token = if refresh_token.is_none() {
access_token
.as_ref()
.filter(|value| header_bearer_token.as_deref() != Some(value.as_str()))
.cloned()
} else {
None
};
let (refresh_token, access_token) = normalize_provider_import_tokens(
provider_type,
refresh_token.as_deref(),
access_token.as_deref().or(session_token.as_deref()),
access_token
.as_deref()
.or(session_token.as_deref())
.or(header_bearer_token.as_deref()),
);
let windsurf_api_key = is_windsurf
.then(|| {
@@ -386,6 +409,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
parse_error: None,
refresh_token,
access_token,
export_access_token,
raw_credentials,
expires_at,
account_id,
@@ -401,6 +425,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
sso_rw_token,
cf_cookies,
cf_clearance,
request_headers,
user_agent,
browser_profile,
})
@@ -479,6 +504,7 @@ fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
parse_error: Some(error),
refresh_token: None,
access_token: None,
export_access_token: None,
raw_credentials: None,
expires_at: None,
account_id: None,
@@ -494,6 +520,7 @@ fn parse_error_entry(error: String) -> AdminProviderOAuthBatchImportEntry {
sso_rw_token: None,
cf_cookies: None,
cf_clearance: None,
request_headers: None,
user_agent: None,
browser_profile: None,
}
@@ -569,6 +596,11 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
.entry("pool_tier".to_string())
.or_insert_with(|| json!(pool_tier));
}
if let Some(access_token) = entry.export_access_token.as_ref() {
auth_config
.entry("access_token".to_string())
.or_insert_with(|| json!(access_token));
}
if let Some(user_id) = entry.user_id.as_ref() {
auth_config
.entry("user_id".to_string())
@@ -599,6 +631,11 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
.entry("cf_clearance".to_string())
.or_insert_with(|| json!(cf_clearance));
}
if let Some(request_headers) = entry.request_headers.as_ref() {
auth_config
.entry("headers".to_string())
.or_insert_with(|| json!(request_headers));
}
if let Some(user_agent) = entry.user_agent.as_ref() {
auth_config
.entry("user_agent".to_string())
@@ -721,11 +758,110 @@ mod tests {
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].refresh_token, None);
assert_eq!(entries[0].access_token.as_deref(), Some("at_1"));
assert_eq!(entries[0].export_access_token.as_deref(), Some("at_1"));
assert_eq!(entries[0].expires_at, Some(2_100_000_000));
assert_eq!(entries[0].account_id.as_deref(), Some("acc-1"));
assert_eq!(entries[0].email.as_deref(), Some("[email protected]"));
}
#[test]
fn parses_codex_import_header_overrides() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"codex",
r#"[{"access_token":"jwt_access","headers":{"authorization":"Bearer session-1","chatgpt-account-id":"acc-1","host":"blocked.example"}}]"#,
);
assert_eq!(entries.len(), 1);
let headers = entries[0]
.request_headers
.as_ref()
.expect("headers should parse");
assert_eq!(
headers.get("authorization"),
Some(&"Bearer session-1".to_string())
);
assert_eq!(
headers.get("chatgpt-account-id"),
Some(&"acc-1".to_string())
);
assert!(!headers.contains_key("host"));
}
#[test]
fn treats_import_authorization_header_as_session_access_token() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"codex",
r#"[{"access_token":"jwt_access","headers":{"authorization":"Bearer session-1"}}]"#,
);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].access_token.as_deref(), Some("jwt_access"));
assert_eq!(
entries[0].export_access_token.as_deref(),
Some("jwt_access")
);
let entries = parse_admin_provider_oauth_batch_import_entries(
"codex",
r#"[{"headers":{"authorization":"Bearer session-1"}}]"#,
);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].access_token.as_deref(), Some("session-1"));
assert!(entries[0].export_access_token.is_none());
}
#[test]
fn applies_codex_import_header_overrides_to_auth_config() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"codex",
r#"{"accessToken":"at_1","headers":{"Authorization":"Bearer session-1"}}"#,
);
let mut auth_config = serde_json::Map::new();
apply_admin_provider_oauth_batch_import_hints("codex", &entries[0], &mut auth_config);
assert_eq!(
auth_config.get("headers"),
Some(&json!({"authorization":"Bearer session-1"}))
);
assert_eq!(auth_config.get("access_token"), Some(&json!("at_1")));
}
#[test]
fn does_not_export_access_token_when_it_is_the_header_bearer() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"codex",
r#"{"accessToken":"session-1","headers":{"Authorization":"Bearer session-1"}}"#,
);
let mut auth_config = serde_json::Map::new();
apply_admin_provider_oauth_batch_import_hints("codex", &entries[0], &mut auth_config);
assert!(entries[0].export_access_token.is_none());
assert!(auth_config.get("access_token").is_none());
assert_eq!(
auth_config.get("headers"),
Some(&json!({"authorization":"Bearer session-1"}))
);
}
#[test]
fn does_not_export_access_token_from_refresh_token_payload() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"codex",
r#"{"refreshToken":"refresh-1","accessToken":"refreshed-access"}"#,
);
let mut auth_config = serde_json::Map::new();
apply_admin_provider_oauth_batch_import_hints("codex", &entries[0], &mut auth_config);
assert_eq!(entries[0].refresh_token.as_deref(), Some("refresh-1"));
assert_eq!(entries[0].access_token.as_deref(), Some("refreshed-access"));
assert!(entries[0].export_access_token.is_none());
assert!(auth_config.get("access_token").is_none());
}
#[test]
fn parses_common_chatgpt_web_json_aliases() {
let entries = parse_admin_provider_oauth_batch_import_entries(
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
+108 -3
View File
@@ -16,7 +16,9 @@ use super::super::state::{
};
use super::helpers::admin_provider_oauth_key_name_from_auth_config;
use super::token_import::{
build_provider_access_token_import_auth_config, normalize_provider_import_tokens,
build_provider_access_token_import_auth_config, decode_access_token_expires_at,
normalize_provider_import_tokens, normalize_provider_oauth_import_headers_from_object,
provider_oauth_import_authorization_bearer_token_from_object,
provider_type_supports_access_token_import,
};
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_import_provider_id;
@@ -151,6 +153,20 @@ fn import_payload_u64_any(
})
}
fn import_payload_export_access_token(
payload: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
if import_payload_string(payload, "refresh_token", "refreshToken").is_some() {
return None;
}
let access_token = import_payload_string(payload, "access_token", "accessToken")?;
let header_bearer = provider_oauth_import_authorization_bearer_token_from_object(payload);
if header_bearer.as_deref() == Some(access_token.as_str()) {
return None;
}
Some(access_token)
}
fn apply_single_import_hints(
provider_type: &str,
payload: &serde_json::Map<String, serde_json::Value>,
@@ -267,6 +283,18 @@ fn apply_single_import_hints(
}
});
}
if let Some(access_token) = import_payload_export_access_token(payload) {
auth_config
.entry("access_token".to_string())
.or_insert_with(|| json!(access_token));
}
if let Some(request_headers) = normalize_provider_oauth_import_headers_from_object(payload) {
auth_config
.entry("headers".to_string())
.or_insert_with(|| json!(request_headers));
}
}
async fn resolve_admin_provider_oauth_single_import_tokens(
@@ -497,7 +525,8 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
"session_token",
"sessionToken",
],
);
)
.or_else(|| provider_oauth_import_authorization_bearer_token_from_object(&raw_payload));
let imported_expires_at =
import_payload_u64_any(&raw_payload, &["expires_at", "expiresAt", "expired"]);
let name = raw_payload
@@ -614,9 +643,22 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
let AdminProviderOAuthSingleImportTokens {
access_token,
mut auth_config,
expires_at,
mut expires_at,
} = resolved_import;
apply_single_import_hints(&provider_type, &raw_payload, &mut auth_config);
if let Some(header_access_token) =
provider_oauth_import_authorization_bearer_token_from_object(&raw_payload)
{
if let Some(header_expires_at) =
decode_access_token_expires_at(&header_access_token).or(imported_expires_at)
{
auth_config.insert("expires_at".to_string(), json!(header_expires_at));
expires_at = Some(header_expires_at);
} else {
auth_config.remove("expires_at");
expires_at = None;
}
}
let has_refresh_token = auth_config
.get("refresh_token")
.and_then(serde_json::Value::as_str)
@@ -776,6 +818,69 @@ mod tests {
assert_eq!(auth_config.get("user_agent"), Some(&json!("antigravity")));
}
#[test]
fn single_import_preserves_distinct_top_level_access_token_for_export() {
let payload = json!({
"access_token": "jwt-access-token",
"headers": {
"authorization": "Bearer imported-session-token"
},
})
.as_object()
.cloned()
.expect("payload should be an object");
let mut auth_config = serde_json::Map::new();
apply_single_import_hints("codex", &payload, &mut auth_config);
assert_eq!(
auth_config.get("access_token"),
Some(&json!("jwt-access-token"))
);
assert_eq!(
auth_config.get("headers"),
Some(&json!({"authorization":"Bearer imported-session-token"}))
);
}
#[test]
fn single_import_does_not_preserve_header_bearer_as_export_access_token() {
let payload = json!({
"access_token": "imported-session-token",
"headers": {
"authorization": "Bearer imported-session-token"
},
})
.as_object()
.cloned()
.expect("payload should be an object");
let mut auth_config = serde_json::Map::new();
apply_single_import_hints("codex", &payload, &mut auth_config);
assert!(auth_config.get("access_token").is_none());
assert_eq!(
auth_config.get("headers"),
Some(&json!({"authorization":"Bearer imported-session-token"}))
);
}
#[test]
fn single_import_does_not_preserve_access_token_from_refresh_token_payload_for_export() {
let payload = json!({
"refresh_token": "refresh-token",
"access_token": "refreshed-access-token",
})
.as_object()
.cloned()
.expect("payload should be an object");
let mut auth_config = serde_json::Map::new();
apply_single_import_hints("codex", &payload, &mut auth_config);
assert!(auth_config.get("access_token").is_none());
}
#[test]
fn windsurf_import_error_redacts_http_body() {
let error = OAuthError::HttpStatus {
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/token_import.rs
+125 -1
View File
@@ -5,6 +5,7 @@ use base64::{
Engine as _,
};
use serde_json::{json, Map, Value};
use std::collections::BTreeMap;
fn decode_base64_url_part(value: &str) -> Option<Vec<u8>> {
URL_SAFE_NO_PAD
@@ -117,6 +118,95 @@ pub(super) fn decode_access_token_expires_at(access_token: &str) -> Option<u64>
json_u64_value(claims.get("exp"))
}
fn normalize_import_header_name(raw: &str) -> Option<String> {
let value = raw.trim().to_ascii_lowercase();
if value.is_empty()
|| matches!(
value.as_str(),
"host"
| "content-length"
| "connection"
| "keep-alive"
| "proxy-authenticate"
| "proxy-authorization"
| "proxy-connection"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
)
{
return None;
}
http::header::HeaderName::from_bytes(value.as_bytes())
.ok()
.map(|name| name.as_str().to_string())
}
fn normalize_import_header_value(value: &Value) -> Option<String> {
let value = match value {
Value::String(raw) => raw.trim().to_string(),
Value::Number(raw) => raw.to_string(),
Value::Bool(raw) => raw.to_string(),
_ => return None,
};
if value.is_empty() || http::header::HeaderValue::from_str(&value).is_err() {
return None;
}
Some(value)
}
pub(super) fn normalize_provider_oauth_import_headers(
value: Option<&Value>,
) -> Option<BTreeMap<String, String>> {
let object = value?.as_object()?;
let mut headers = BTreeMap::new();
for (raw_key, raw_value) in object {
let Some(key) = normalize_import_header_name(raw_key) else {
continue;
};
let Some(value) = normalize_import_header_value(raw_value) else {
continue;
};
headers.insert(key, value);
}
(!headers.is_empty()).then_some(headers)
}
pub(super) fn normalize_provider_oauth_import_headers_from_object(
object: &Map<String, Value>,
) -> Option<BTreeMap<String, String>> {
normalize_provider_oauth_import_headers(
object
.get("headers")
.or_else(|| object.get("request_headers"))
.or_else(|| object.get("requestHeaders"))
.or_else(|| object.get("header_overrides"))
.or_else(|| object.get("headerOverrides"))
.or_else(|| object.get("extra_headers"))
.or_else(|| object.get("extraHeaders")),
)
}
pub(super) fn provider_oauth_import_authorization_bearer_token(
headers: Option<&BTreeMap<String, String>>,
) -> Option<String> {
let value = headers?.get("authorization")?.trim();
let (scheme, token) = value.split_once(' ')?;
if !scheme.eq_ignore_ascii_case("bearer") {
return None;
}
let token = token.trim();
(!token.is_empty()).then(|| token.to_string())
}
pub(super) fn provider_oauth_import_authorization_bearer_token_from_object(
object: &Map<String, Value>,
) -> Option<String> {
let headers = normalize_provider_oauth_import_headers_from_object(object)?;
provider_oauth_import_authorization_bearer_token(Some(&headers))
}
pub(super) fn provider_type_supports_access_token_import(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
@@ -176,7 +266,9 @@ pub(super) fn build_provider_access_token_import_auth_config(
mod tests {
use super::{
build_provider_access_token_import_auth_config, decode_access_token_expires_at,
looks_like_access_token, normalize_provider_import_tokens, normalize_single_import_tokens,
looks_like_access_token, normalize_provider_import_tokens,
normalize_provider_oauth_import_headers, normalize_single_import_tokens,
provider_oauth_import_authorization_bearer_token,
};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use serde_json::json;
@@ -247,6 +339,38 @@ mod tests {
);
}
#[test]
fn normalizes_import_header_overrides_for_auth_headers() {
let headers = normalize_provider_oauth_import_headers(Some(&json!({
"Authorization": " Bearer session-token ",
"X-Feature": true,
"Host": "evil.example",
"X-Bad": "line\nbreak",
})))
.expect("headers should normalize");
assert_eq!(
headers.get("authorization"),
Some(&"Bearer session-token".to_string())
);
assert_eq!(headers.get("x-feature"), Some(&"true".to_string()));
assert!(!headers.contains_key("host"));
assert!(!headers.contains_key("x-bad"));
}
#[test]
fn extracts_import_authorization_header_bearer_token() {
let headers = normalize_provider_oauth_import_headers(Some(&json!({
"Authorization": "Bearer at-session-token",
})))
.expect("headers should normalize");
assert_eq!(
provider_oauth_import_authorization_bearer_token(Some(&headers)).as_deref(),
Some("at-session-token")
);
}
#[test]
fn builds_chatgpt_web_temporary_auth_config_from_access_token() {
let token = unsigned_jwt(json!({
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs
+5 -1
View File
@@ -76,8 +76,12 @@ async fn execute_chatgpt_web_quota_plan(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let spec =
let mut spec =
build_chatgpt_web_pool_quota_request(&transport.key.id, &endpoint.base_url, authorization);
crate::provider_transport::apply_local_auth_config_header_overrides(
&mut spec.headers,
transport.key.decrypted_auth_config.as_deref(),
);
let resolved_transport_profile = state.resolve_transport_profile(transport);
let plan = super::shared::build_provider_quota_execution_plan(
transport,
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs
+7 -2
View File
@@ -16,12 +16,17 @@ pub(super) fn build_codex_quota_request_spec(
.decrypted_auth_config
.as_deref()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok());
build_codex_pool_quota_request(
let mut request = build_codex_pool_quota_request(
&transport.key.id,
resolved_oauth_auth,
Some(transport.key.decrypted_api_key.as_str()),
auth_config.as_ref(),
)
)?;
crate::provider_transport::apply_local_auth_config_header_overrides(
&mut request.headers,
transport.key.decrypted_auth_config.as_deref(),
);
Ok(request)
}
pub(super) async fn execute_codex_quota_plan(
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
-7
View File
@@ -405,7 +405,6 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
cost_limit_per_key_tokens: None,
rate_limit_cooldown_seconds: 300,
overload_cooldown_seconds: 30,
health_policy_enabled: true,
probing_enabled: false,
probing_target_percent: None,
probing_target_count: None,
@@ -466,10 +465,6 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.and_then(json_u64)
.filter(|value| *value > 0)
.unwrap_or(30),
health_policy_enabled: pool_advanced
.get("health_policy_enabled")
.and_then(Value::as_bool)
.unwrap_or(true),
probing_enabled: pool_advanced
.get("probing_enabled")
.and_then(Value::as_bool)
@@ -583,7 +578,6 @@ mod tests {
"cost_limit_per_key_tokens": 12000,
"rate_limit_cooldown_seconds": 420,
"overload_cooldown_seconds": 45,
"health_policy_enabled": false,
"probing_enabled": true,
"probing_target_percent": 25,
"probing_target_count": 3,
@@ -624,7 +618,6 @@ mod tests {
assert_eq!(config.cost_limit_per_key_tokens, Some(12_000));
assert_eq!(config.rate_limit_cooldown_seconds, 420);
assert_eq!(config.overload_cooldown_seconds, 45);
assert!(!config.health_policy_enabled);
assert!(config.probing_enabled);
assert_eq!(config.probing_target_percent, Some(25.0));
assert_eq!(config.probing_target_count, Some(3));
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
+1 -6
View File
@@ -437,10 +437,6 @@ pub(crate) async fn record_admin_provider_pool_error(
error_body: Option<&str>,
response_headers: Option<&BTreeMap<String, String>>,
) {
if !pool_config.health_policy_enabled {
return;
}
let error_message = extract_error_message(error_body).to_ascii_lowercase();
if status_code == 401 {
@@ -553,7 +549,7 @@ pub(crate) async fn record_admin_provider_pool_stream_timeout(
key_id: &str,
pool_config: &AdminProviderPoolConfig,
) {
if !pool_config.health_policy_enabled || pool_config.stream_timeout_threshold == 0 {
if pool_config.stream_timeout_threshold == 0 {
return;
}
@@ -637,7 +633,6 @@ mod tests {
cost_limit_per_key_tokens: Some(10_000),
rate_limit_cooldown_seconds: 300,
overload_cooldown_seconds: 30,
health_policy_enabled: true,
probing_enabled: false,
probing_target_percent: None,
probing_target_count: None,
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
+5 -1
View File
@@ -4,7 +4,8 @@ use crate::handlers::admin::provider::shared::support::{
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
use crate::provider_key_auth::{
provider_key_auth_semantics, provider_key_can_refresh_oauth, provider_key_effective_api_formats,
provider_key_auth_config_uses_header_authorization, provider_key_auth_semantics,
provider_key_can_refresh_oauth, provider_key_effective_api_formats,
};
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_admin::provider::quota as admin_provider_quota_pure;
@@ -1114,6 +1115,8 @@ pub(super) fn build_admin_pool_key_payload(
.and_then(|config| config.get("access_token_import_temporary"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let oauth_header_auth = auth_semantics.oauth_managed()
&& provider_key_auth_config_uses_header_authorization(auth_config.as_ref());
let account_status_code = admin_pool_trimmed_string_from_map(account_snapshot, "code");
let account_status_label =
admin_pool_trimmed_string(account_snapshot.and_then(|item| item.get("label")));
@@ -1193,6 +1196,7 @@ pub(super) fn build_admin_pool_key_payload(
"can_edit_oauth".to_string(),
json!(auth_semantics.can_edit_oauth()),
);
payload.insert("oauth_header_auth".to_string(), json!(oauth_header_auth));
payload.insert("oauth_expires_at".to_string(), json!(oauth_expires_at));
payload.insert("oauth_invalid_at".to_string(), json!(oauth_invalid_at));
payload.insert(
apps/aether-gateway/src/handlers/admin/provider/pool_admin/read_routes/resolve_selection.rs
+6 -1
View File
@@ -3,7 +3,10 @@ use super::{
AdminPoolResolveSelectionRequest, ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::provider_key_auth::{provider_key_auth_semantics, provider_key_can_refresh_oauth};
use crate::provider_key_auth::{
provider_key_auth_config_uses_header_authorization, provider_key_auth_semantics,
provider_key_can_refresh_oauth,
};
use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure;
use axum::{
@@ -106,6 +109,8 @@ pub(super) async fn build_admin_pool_resolve_selection_response(
"can_refresh_oauth": provider_key_can_refresh_oauth(auth_semantics, auth_config.as_ref()),
"can_export_oauth": auth_semantics.can_export_oauth(),
"can_edit_oauth": auth_semantics.can_edit_oauth(),
"oauth_header_auth": auth_semantics.oauth_managed()
&& provider_key_auth_config_uses_header_authorization(auth_config.as_ref()),
})
})
.collect::<Vec<_>>();
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test.rs
+1
View File
@@ -2180,6 +2180,7 @@ async fn provider_query_execute_openai_image_test_candidate(
} else {
crate::provider_transport::build_openai_image_headers(
crate::provider_transport::ProviderOpenAiImageHeadersInput {
transport: &transport,
headers: &parts.headers,
auth_header: &auth_header,
auth_value: &auth_value,
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
-1
View File
@@ -43,7 +43,6 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) cost_limit_per_key_tokens: Option<u64>,
pub(crate) rate_limit_cooldown_seconds: u64,
pub(crate) overload_cooldown_seconds: u64,
pub(crate) health_policy_enabled: bool,
pub(crate) probing_enabled: bool,
pub(crate) probing_target_percent: Option<f64>,
pub(crate) probing_target_count: Option<u64>,
apps/aether-gateway/src/handlers/admin/provider/write/reveal.rs
+119
View File
@@ -97,6 +97,7 @@ fn provider_oauth_export_payload(
if let Some(access_token) = fallback_access_token
.map(str::trim)
.filter(|value| !value.is_empty() && *value != "__placeholder__")
.filter(|value| !oauth_export_fallback_matches_authorization_header(&payload, value))
{
payload.insert("access_token".to_string(), json!(access_token));
}
@@ -116,6 +117,43 @@ fn provider_oauth_export_payload(
payload
}
fn oauth_export_fallback_matches_authorization_header(
payload: &serde_json::Map<String, serde_json::Value>,
fallback_access_token: &str,
) -> bool {
let fallback_access_token = fallback_access_token.trim();
if fallback_access_token.is_empty() {
return false;
}
let Some(authorization) = payload
.get("headers")
.and_then(serde_json::Value::as_object)
.and_then(|headers| {
headers
.iter()
.find(|(key, _)| key.trim().eq_ignore_ascii_case("authorization"))
.and_then(|(_, value)| value.as_str())
})
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return false;
};
if authorization == fallback_access_token {
return true;
}
let mut parts = authorization.splitn(2, char::is_whitespace);
let Some(scheme) = parts.next() else {
return false;
};
let Some(token) = parts.next() else {
return false;
};
scheme.eq_ignore_ascii_case("bearer") && token.trim() == fallback_access_token
}
fn json_map_has_non_empty_string(
map: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
@@ -186,3 +224,84 @@ pub(crate) async fn build_admin_export_key_payload(
);
Ok(serde_json::Value::Object(payload))
}
#[cfg(test)]
mod tests {
use super::provider_oauth_export_payload;
use serde_json::json;
#[test]
fn oauth_export_preserves_imported_request_headers() {
let auth_config = json!({
"provider_type": "codex",
"email": "[email protected]",
"headers": {
"authorization": "Bearer imported-session",
"chatgpt-account-id": "acct-1"
}
})
.as_object()
.cloned()
.expect("auth_config should be an object");
let payload = provider_oauth_export_payload("codex", &auth_config, None, Some("fallback"));
assert_eq!(
payload.get("headers"),
Some(&json!({
"authorization": "Bearer imported-session",
"chatgpt-account-id": "acct-1"
}))
);
assert_eq!(payload.get("access_token"), Some(&json!("fallback")));
}
#[test]
fn oauth_export_does_not_promote_imported_header_bearer_to_access_token() {
let auth_config = json!({
"provider_type": "codex",
"email": "[email protected]",
"headers": {
"authorization": "Bearer imported-session"
}
})
.as_object()
.cloned()
.expect("auth_config should be an object");
let payload =
provider_oauth_export_payload("codex", &auth_config, None, Some("imported-session"));
assert_eq!(
payload.get("headers"),
Some(&json!({"authorization": "Bearer imported-session"}))
);
assert!(payload.get("access_token").is_none());
}
#[test]
fn oauth_export_keeps_explicit_access_token_even_with_header_bearer() {
let auth_config = json!({
"provider_type": "codex",
"access_token": "jwt-access-token",
"headers": {
"authorization": "Bearer imported-session"
}
})
.as_object()
.cloned()
.expect("auth_config should be an object");
let payload =
provider_oauth_export_payload("codex", &auth_config, None, Some("imported-session"));
assert_eq!(
payload.get("access_token"),
Some(&json!("jwt-access-token"))
);
assert_eq!(
payload.get("headers"),
Some(&json!({"authorization": "Bearer imported-session"}))
);
}
}
apps/aether-gateway/src/handlers/public/system_modules_helpers/keys_grouped.rs
+17 -6
View File
@@ -1,16 +1,27 @@
use super::enabled_key_capability_short_names;
use crate::handlers::shared::unix_secs_to_rfc3339;
use crate::provider_key_auth::provider_key_effective_api_formats;
use crate::handlers::shared::{parse_catalog_auth_config_json, unix_secs_to_rfc3339};
use crate::provider_key_auth::{
provider_key_auth_config_uses_header_authorization, provider_key_effective_api_formats,
};
use crate::AppState;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use serde_json::json;
use std::collections::{BTreeMap, HashMap};
use std::time::{SystemTime, UNIX_EPOCH};
fn grouped_key_masked_label(auth_type: &str) -> &'static str {
match auth_type.trim() {
fn grouped_key_masked_label(state: &AppState, key: &StoredProviderCatalogKey) -> &'static str {
match key.auth_type.trim() {
"service_account" | "vertex_ai" => "[Service Account]",
"oauth" => "[OAuth Token]",
"oauth" => {
if provider_key_auth_config_uses_header_authorization(
parse_catalog_auth_config_json(state, key).as_ref(),
) {
"[OAuth Header]"
} else {
"[OAuth Token]"
}
}
_ => "[API Key]",
}
}
@@ -150,7 +161,7 @@ pub(crate) async fn build_admin_keys_grouped_by_format_payload(
"provider_id": key.provider_id,
"name": key.name,
"auth_type": key.auth_type,
"api_key_masked": grouped_key_masked_label(&key.auth_type),
"api_key_masked": grouped_key_masked_label(state, &key),
"internal_priority": key.internal_priority,
"global_priority_by_format": key.global_priority_by_format,
"rate_multipliers": key.rate_multipliers,
apps/aether-gateway/src/handlers/shared/catalog.rs
+15 -3
View File
@@ -1,7 +1,8 @@
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
use crate::provider_key_auth::{
provider_key_auth_semantics, provider_key_can_refresh_oauth,
provider_key_configured_api_formats, provider_key_inherits_provider_api_formats,
provider_key_auth_config_uses_header_authorization, provider_key_auth_semantics,
provider_key_can_refresh_oauth, provider_key_configured_api_formats,
provider_key_inherits_provider_api_formats,
};
use crate::AppState;
use aether_admin::provider::quota as admin_provider_quota_pure;
@@ -132,7 +133,15 @@ pub(crate) fn take_secret_suffix(value: &str, suffix_chars: usize) -> &str {
pub(crate) fn masked_catalog_api_key(state: &AppState, key: &StoredProviderCatalogKey) -> String {
match key.auth_type.trim() {
"service_account" | "vertex_ai" => "[Service Account]".to_string(),
"oauth" => "[OAuth Token]".to_string(),
"oauth" => {
if provider_key_auth_config_uses_header_authorization(
parse_catalog_auth_config_json(state, key).as_ref(),
) {
"[OAuth Header]".to_string()
} else {
"[OAuth Token]".to_string()
}
}
_ => {
let Some(ciphertext) = key
.encrypted_api_key
@@ -2207,6 +2216,8 @@ pub(crate) fn build_admin_provider_key_response(
.and_then(|config| config.get("access_token_import_temporary"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let oauth_header_auth = auth_semantics.oauth_managed()
&& provider_key_auth_config_uses_header_authorization(auth_config.as_ref());
let oauth_plan_type = derive_catalog_oauth_plan_type(key, provider_type, auth_config.as_ref());
let (
health_score,
@@ -2279,6 +2290,7 @@ pub(crate) fn build_admin_provider_key_response(
"can_edit_oauth".to_string(),
json!(auth_semantics.can_edit_oauth()),
);
payload.insert("oauth_header_auth".to_string(), json!(oauth_header_auth));
payload.insert("name".to_string(), json!(key.name));
payload.insert("rate_multipliers".to_string(), json!(key.rate_multipliers));
payload.insert(