i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-09 18:59:50 +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/ai_serving/planner/specialized/image/request.rs
+1
View File
@@ -210,6 +210,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
})
} else {
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
transport,
headers: effective_headers,
auth_header: &auth_header,
auth_value: &auth_value,
apps/aether-gateway/src/ai_serving/planner/standard/family/request.rs
+1
View File
@@ -1227,6 +1227,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
let effective_headers = input.effective_headers(&parts.headers);
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
transport,
headers: effective_headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+1
View File
@@ -1053,6 +1053,7 @@ async fn resolve_openai_chat_to_openai_image_payload_parts(
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
transport,
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/decision/request.rs
+19 -6
View File
@@ -46,12 +46,12 @@ use crate::ai_serving::transport::kiro::{
KiroRequestAuth, KIRO_ENVELOPE_NAME,
};
use crate::ai_serving::transport::{
build_grok_browser_headers, build_grok_upstream_url, build_kiro_cross_format_upstream_url,
build_openai_image_headers, build_openai_image_upstream_url,
build_standard_provider_request_headers, build_windsurf_cascade_headers,
build_windsurf_cascade_request_body, build_windsurf_cascade_upstream_url,
is_gemini_cli_provider_transport, is_windsurf_provider_transport,
local_standard_transport_unsupported_reason_with_network,
apply_local_auth_config_header_overrides, build_grok_browser_headers, build_grok_upstream_url,
build_kiro_cross_format_upstream_url, build_openai_image_headers,
build_openai_image_upstream_url, build_standard_provider_request_headers,
build_windsurf_cascade_headers, build_windsurf_cascade_request_body,
build_windsurf_cascade_upstream_url, is_gemini_cli_provider_transport,
is_windsurf_provider_transport, local_standard_transport_unsupported_reason_with_network,
local_windsurf_request_transport_unsupported_reason_with_network,
openai_image_transport_unsupported_reason, resolve_openai_image_auth, GrokHeaderInput,
ProviderOpenAiImageHeadersInput, StandardProviderRequestHeadersInput,
@@ -681,6 +681,10 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
apply_local_auth_config_header_overrides(
&mut provider_request_headers,
transport.key.decrypted_auth_config.as_deref(),
);
}
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
@@ -857,6 +861,10 @@ async fn build_gemini_cli_openai_responses_payload_parts(
Some(trace_id),
resolved.transport.key.decrypted_auth_config.as_deref(),
);
apply_local_auth_config_header_overrides(
&mut provider_request_headers,
resolved.transport.key.decrypted_auth_config.as_deref(),
);
request_identity_response_encoding_when_redacted(
&mut provider_request_headers,
request_redacted,
@@ -1143,6 +1151,7 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
};
let Some(mut provider_request_headers) =
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
transport,
headers: &parts.headers,
auth_header: &prepared_candidate.auth_header,
auth_value: &prepared_candidate.auth_value,
@@ -1181,6 +1190,10 @@ async fn resolve_openai_responses_to_openai_image_payload_parts(
Some(trace_id),
transport.key.decrypted_auth_config.as_deref(),
);
apply_local_auth_config_header_overrides(
&mut provider_request_headers,
transport.key.decrypted_auth_config.as_deref(),
);
}
let (execution_strategy, conversion_mode) =
apps/aether-gateway/src/ai_serving/transport.rs
+2 -2
View File
@@ -59,8 +59,8 @@ pub(crate) mod windsurf {
}
pub(crate) use aether_provider_transport::{
append_transport_diagnostics_to_value, apply_local_body_rules,
apply_local_body_rules_with_request_headers, apply_local_header_rules,
append_transport_diagnostics_to_value, apply_local_auth_config_header_overrides,
apply_local_body_rules, apply_local_body_rules_with_request_headers, apply_local_header_rules,
apply_local_header_rules_with_request_headers, apply_standard_provider_request_body_rules,
apply_standard_provider_request_body_rules_with_request_headers,
apply_transport_request_body_semantics, body_rules_are_locally_supported,
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(
apps/aether-gateway/src/provider_key_auth.rs
+39 -1
View File
@@ -95,6 +95,23 @@ pub(crate) fn provider_key_can_refresh_oauth(
.is_some_and(|value| !value.is_empty())
}
pub(crate) fn provider_key_auth_config_uses_header_authorization(
auth_config: Option<&Map<String, Value>>,
) -> bool {
auth_config
.and_then(|config| config.get("headers"))
.and_then(Value::as_object)
.is_some_and(|headers| {
headers.iter().any(|(key, value)| {
key.trim().eq_ignore_ascii_case("authorization")
&& value
.as_str()
.map(str::trim)
.is_some_and(|value| !value.is_empty())
})
})
}
fn normalized_auth_type(key: &StoredProviderCatalogKey) -> String {
key.auth_type.trim().to_ascii_lowercase()
}
@@ -255,7 +272,8 @@ pub(crate) fn provider_key_effective_api_formats(
#[cfg(test)]
mod tests {
use super::{
provider_active_api_formats, provider_key_auth_semantics, provider_key_can_refresh_oauth,
provider_active_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_effective_api_formats,
provider_key_inherits_provider_api_formats, ProviderKeyCredentialKind,
ProviderKeyRuntimeAuthKind,
@@ -377,6 +395,26 @@ mod tests {
));
}
#[test]
fn detects_oauth_header_authorization_in_auth_config() {
assert!(provider_key_auth_config_uses_header_authorization(
json!({
"headers": {
"Authorization": "Bearer imported-session"
}
})
.as_object()
));
assert!(!provider_key_auth_config_uses_header_authorization(
json!({
"headers": {
"authorization": " "
}
})
.as_object()
));
}
#[test]
fn recognizes_legacy_kiro_bearer_key_with_auth_config_as_oauth_managed() {
let mut key = sample_key("bearer");
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
+226
View File
@@ -498,6 +498,81 @@ async fn gateway_admin_provider_keys_prefers_upstream_plan_type_over_auth_config
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_admin_provider_keys_marks_oauth_header_auth() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-codex/keys",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth-header",
"provider-codex",
"openai:responses",
"imported-session-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","headers":{"authorization":"Bearer imported-session-token"}}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![key],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex/keys?skip=0&limit=50"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
let items = payload.as_array().expect("payload should be an array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["api_key_masked"], "[OAuth Header]");
assert_eq!(items[0]["oauth_header_auth"], true);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_creates_admin_provider_key_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -1123,6 +1198,157 @@ async fn gateway_exports_admin_provider_key_access_token_when_refresh_token_is_m
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_export_does_not_emit_access_token_from_imported_authorization_header() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-codex-a/export",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut key = sample_key(
"key-codex-a",
"provider-codex",
"openai:responses",
"imported-session-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","headers":{"authorization":"Bearer imported-session-token"}}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex", 10)],
vec![],
vec![key],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/keys/key-codex-a/export"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(
payload["headers"]["authorization"],
"Bearer imported-session-token"
);
assert!(payload.get("access_token").is_none());
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_export_preserves_distinct_imported_access_token_with_authorization_header() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/endpoints/keys/key-codex-a/export",
any(move |_request: Request| {
let upstream_hits_inner = Arc::clone(&upstream_hits_clone);
async move {
*upstream_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected upstream hit"))
}
}),
);
let mut key = sample_key(
"key-codex-a",
"provider-codex",
"openai:responses",
"jwt-access-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","email":"[email protected]","access_token":"jwt-access-token","headers":{"authorization":"Bearer imported-session-token"}}"#,
)
.expect("auth config ciphertext should build"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex", 10)],
vec![],
vec![key],
));
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/endpoints/keys/key-codex-a/export"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["provider_type"], "codex");
assert_eq!(payload["access_token"], "jwt-access-token");
assert_eq!(
payload["headers"]["authorization"],
"Bearer imported-session-token"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_clears_admin_provider_key_oauth_invalid_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
apps/aether-gateway/src/tests/control/admin/oauth.rs
+89
View File
@@ -3417,6 +3417,95 @@ async fn gateway_imports_codex_access_token_without_refresh_token_as_temporary_a
token_handle.abort();
}
#[tokio::test]
async fn gateway_imports_codex_header_authorization_without_overwriting_payload_access_token() {
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
"endpoint-codex-chat",
"provider-codex",
"openai:chat",
"https://chatgpt.com/backend-api/codex",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let access_token =
sample_codex_access_token_with_profile_email("[email protected]", "acct-profile-123");
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-codex/import-refresh-token"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.json(&json!({
"access_token": access_token,
"headers": {
"authorization": "Bearer imported-session-token",
"chatgpt-account-id": "acct-header"
},
"name": "temporary-codex-header-auth",
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
let reloaded = provider_catalog_repository
.list_keys_by_provider_ids(&["provider-codex".to_string()])
.await
.expect("keys should load");
let persisted = reloaded.first().expect("persisted key should exist");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should be present"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, access_token);
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["access_token"], access_token);
assert_eq!(
auth_config["headers"]["authorization"],
"Bearer imported-session-token"
);
assert_eq!(auth_config["headers"]["chatgpt-account-id"], "acct-header");
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_imports_chatgpt_web_access_token_without_refresh_token_as_temporary_account() {
let token_hits = Arc::new(Mutex::new(0usize));
apps/aether-gateway/src/tests/control/admin/pool.rs
+130
View File
@@ -3040,6 +3040,136 @@ async fn gateway_pool_plan_free_selector_prefers_upstream_plan_type() {
);
}
#[tokio::test]
async fn gateway_pool_keys_mark_oauth_header_auth() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth-header",
"provider-codex",
"openai:responses",
"imported-session-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","headers":{"authorization":"Bearer imported-session-token"}}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-codex/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys.len(), 1);
assert_eq!(keys[0]["oauth_header_auth"], true);
}
#[tokio::test]
async fn gateway_pool_resolve_selection_marks_oauth_header_auth() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-oauth-header",
"provider-codex",
"openai:responses",
"imported-session-token",
);
key.auth_type = "oauth".to_string();
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","headers":{"authorization":"Bearer imported-session-token"}}"#,
)
.expect("auth config should encrypt"),
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_reader_for_tests(provider_catalog_repository)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
);
let response = local_admin_pool_response(
&state,
http::Method::POST,
"/api/admin/pool/provider-codex/keys/resolve-selection",
Some(json!({})),
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let items = payload["items"].as_array().expect("items should be array");
assert_eq!(items.len(), 1);
assert_eq!(items[0]["oauth_header_auth"], true);
}
#[tokio::test]
async fn gateway_handles_admin_pool_resolve_selection_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));