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

Merge remote-tracking branch 'origin/pr-484'

Browse Source
This commit is contained in:
fawney19
2026-05-18 18:01:57 +08:00
parent be939f7e63 1b24e1c22a
commit 7f76eff827
17 changed files with 925 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
apps/aether-gateway/src/control/auth/gate.rs
View File

@@ -106,7 +106,7 @@ async fn balance_capacity_rejection(
requested_model: Option<&str>,
body: &Bytes,
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
if auth_context.api_key_is_standalone || auth_context.admin_bypass_limits {
if auth_context.api_key_is_standalone {
return Ok(None);
}
if auth_context.local_rejection.is_some() {
@@ -816,6 +816,43 @@ mod tests {
}
}
#[tokio::test]
async fn admin_bypass_limits_does_not_skip_exhausted_daily_quota_capacity() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}]
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(0.0, false), context);
let mut decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
if let Some(auth_context) = decision.auth_context.as_mut() {
auth_context.admin_bypass_limits = true;
}
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(
br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"stream":true}"#,
);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("quota rejection should resolve");
assert_eq!(
rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(0.0),
})
);
}
#[tokio::test]
async fn positive_balance_still_denies_known_cost_above_available_capacity() {
let context = billing_context_with_pricing(

5
apps/aether-gateway/src/control/auth/mod.rs
View File

@@ -9,7 +9,8 @@ pub(crate) use gate::{
request_model_local_rejection, should_buffer_request_for_local_auth,
trusted_auth_local_rejection, GatewayLocalAuthRejection,
};
pub(super) use resolution::{resolve_control_decision_auth, ControlDecisionAuthResolution};
pub(crate) use resolution::{
resolve_execution_runtime_auth_context, GatewayAdminPrincipalContext, GatewayControlAuthContext,
refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
GatewayAdminPrincipalContext, GatewayControlAuthContext,
};
pub(super) use resolution::{resolve_control_decision_auth, ControlDecisionAuthResolution};

229
apps/aether-gateway/src/control/auth/resolution.rs
View File

@@ -433,7 +433,14 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
let _ = trace_id;
if let Some(auth_context) = decision.auth_context.clone() {
return Ok(Some(auth_context));
return Ok(Some(
refresh_execution_runtime_auth_context(
state,
auth_context,
decision.auth_endpoint_signature.as_deref(),
)
.await?,
));
}
let Some(auth_endpoint_signature) = decision.auth_endpoint_signature.as_deref() else {
@@ -445,7 +452,14 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
};
if let Some(auth_context) = get_cached_auth_context(state, &cache_key) {
return Ok(Some(auth_context));
let refreshed = refresh_execution_runtime_auth_context(
state,
auth_context,
Some(auth_endpoint_signature),
)
.await?;
put_cached_auth_context(state, cache_key, refreshed.clone());
return Ok(Some(refreshed));
}
if let Some(auth_context) =
@@ -461,6 +475,56 @@ pub(crate) async fn resolve_execution_runtime_auth_context(
Ok(None)
}
pub(crate) async fn refresh_execution_runtime_auth_context(
state: &AppState,
auth_context: GatewayControlAuthContext,
auth_endpoint_signature: Option<&str>,
) -> Result<GatewayControlAuthContext, GatewayError> {
if auth_context.local_rejection.is_some() || !auth_context.access_allowed {
return Ok(auth_context);
}
let Some(auth_endpoint_signature) = auth_endpoint_signature
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(auth_context);
};
if !state.has_auth_api_key_reader()
|| auth_context.user_id.trim().is_empty()
|| auth_context.api_key_id.trim().is_empty()
{
return Ok(auth_context);
}
let snapshot = state
.data
.read_auth_api_key_snapshot(
&auth_context.user_id,
&auth_context.api_key_id,
current_unix_secs(),
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let Some(snapshot) = snapshot else {
let mut denied = auth_context;
denied.access_allowed = false;
denied.local_rejection = Some(GatewayLocalAuthRejection::InvalidApiKey);
denied.balance_remaining = None;
return Ok(denied);
};
let wallet_access = resolve_wallet_auth_gate(state, &snapshot).await?;
Ok(build_data_backed_auth_context(
state,
snapshot,
auth_endpoint_signature,
Some(true),
auth_context.balance_remaining,
wallet_access,
)
.await)
}
fn put_cached_auth_context(
state: &AppState,
cache_key: String,
@@ -609,7 +673,7 @@ async fn build_data_backed_auth_context(
.api_key_expires_at_unix_secs
.is_some_and(|expires_at| expires_at < current_unix_secs());
let locked_api_key = snapshot.api_key_is_locked && !snapshot.api_key_is_standalone;
let access_allowed = header_access_allowed
let key_access_allowed = header_access_allowed
.map(|value| value && snapshot.currently_usable)
.unwrap_or(snapshot.currently_usable);
let wallet_remaining = wallet_access
@@ -656,7 +720,7 @@ async fn build_data_backed_auth_context(
user_id: snapshot.user_id,
api_key_id: snapshot.api_key_id,
balance_remaining: wallet_remaining.or(balance_remaining),
access_allowed,
access_allowed: key_access_allowed && local_rejection.is_none(),
user_rate_limit: snapshot.user_rate_limit,
api_key_rate_limit: snapshot.api_key_rate_limit,
api_key_is_standalone: snapshot.api_key_is_standalone,
@@ -835,13 +899,20 @@ mod tests {
InMemoryAuthApiKeySnapshotRepository, StoredAuthApiKeySnapshot,
};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::wallet::{
InMemoryWalletRepository, StoredWalletSnapshot, WalletReadRepository,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
use axum::http::{HeaderMap, Uri};
use super::{resolve_data_backed_auth_context, GatewayLocalAuthRejection};
use super::{
resolve_data_backed_auth_context, resolve_execution_runtime_auth_context,
GatewayLocalAuthRejection,
};
use crate::control::auth::credentials::hash_api_key;
use crate::control::GatewayControlDecision;
use crate::data::GatewayDataState;
use crate::AppState;
@@ -946,6 +1017,154 @@ mod tests {
assert_eq!(repository.touch_count("key-1"), 1);
}
#[tokio::test]
async fn data_backed_auth_context_marks_wallet_denial_as_not_allowed() {
let api_key = "sk-test-empty-wallet";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(api_key)),
sample_snapshot("key-empty-wallet", "user-empty-wallet"),
)]));
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
StoredWalletSnapshot::new(
"wallet-empty".to_string(),
Some("user-empty-wallet".to_string()),
None,
0.0,
0.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
0.0,
0.0,
0.0,
0.0,
100,
)
.expect("wallet should build"),
]));
let data =
GatewayDataState::with_auth_and_wallet_for_tests(auth_repository, wallet_repository);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data);
let mut headers = HeaderMap::new();
headers.insert(
http::header::AUTHORIZATION,
format!("Bearer {api_key}").parse().unwrap(),
);
let auth_context = resolve_data_backed_auth_context(
&state,
&headers,
&uri("/v1/chat/completions"),
Some("openai:chat"),
)
.await
.expect("resolution should succeed")
.expect("auth context should exist");
assert_eq!(
auth_context.local_rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(0.0),
})
);
assert!(!auth_context.access_allowed);
}
#[tokio::test]
async fn execution_runtime_auth_context_revalidates_cached_wallet_state() {
let api_key = "sk-test-runtime-wallet-cache";
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(api_key)),
sample_snapshot("key-runtime-wallet-cache", "user-runtime-wallet-cache"),
)]));
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![
StoredWalletSnapshot::new(
"wallet-runtime-cache".to_string(),
Some("user-runtime-wallet-cache".to_string()),
None,
10.0,
0.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
10.0,
0.0,
0.0,
0.0,
100,
)
.expect("wallet should build"),
]));
let data = GatewayDataState::with_auth_and_wallet_for_tests(
auth_repository,
Arc::clone(&wallet_repository),
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data);
let decision = GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
);
let mut headers = HeaderMap::new();
headers.insert("x-api-key", api_key.parse().unwrap());
let first = resolve_execution_runtime_auth_context(
&state,
&decision,
&headers,
&uri("/v1/chat/completions"),
"trace-runtime-wallet-cache",
)
.await
.expect("resolution should succeed")
.expect("auth context should exist");
assert!(first.access_allowed);
wallet_repository
.update_auth_user_wallet_snapshot(
"user-runtime-wallet-cache",
0.0,
0.0,
"finite",
"USD",
"active",
10.0,
10.0,
0.0,
0.0,
Some(101),
)
.await
.expect("wallet update should succeed")
.expect("wallet should exist");
let second = resolve_execution_runtime_auth_context(
&state,
&decision,
&headers,
&uri("/v1/chat/completions"),
"trace-runtime-wallet-cache",
)
.await
.expect("resolution should succeed")
.expect("auth context should exist");
assert_eq!(
second.local_rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(0.0),
})
);
assert!(!second.access_allowed);
}
#[tokio::test]
async fn data_backed_auth_context_allows_provider_id_for_matching_provider_type() {
let api_key = "sk-test-provider-id";

7
apps/aether-gateway/src/control/mod.rs
View File

@@ -8,9 +8,10 @@ mod public;
mod route;
pub(crate) use auth::{
extract_requested_model, request_model_local_rejection, resolve_execution_runtime_auth_context,
should_buffer_request_for_local_auth, trusted_auth_local_rejection,
GatewayAdminPrincipalContext, GatewayControlAuthContext, GatewayLocalAuthRejection,
extract_requested_model, refresh_execution_runtime_auth_context, request_model_local_rejection,
resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth,
trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext,
GatewayLocalAuthRejection,
};
pub(crate) use execute::{allows_control_execute_emergency, maybe_execute_via_control};
pub(crate) use management_token_permissions::{