mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
Merge remote-tracking branch 'upstream/main'
# Conflicts: # apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/payloads.rs # apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/reads.rs # apps/aether-gateway/src/handlers/admin/provider/endpoints_admin/update.rs # apps/aether-gateway/src/tests/control/admin/endpoints/routes.rs # frontend/src/features/models/components/GlobalModelFormDialog.vue # frontend/src/features/providers/components/ProviderModelFormDialog.vue # frontend/src/features/providers/components/provider-tabs/__tests__/model-test-request.spec.ts # frontend/src/features/providers/components/provider-tabs/model-test-request.ts
This commit is contained in:
@@ -2,14 +2,20 @@ use std::env;
|
||||
use std::process::Command;
|
||||
|
||||
fn main() {
|
||||
println!("cargo:rerun-if-env-changed=AETHER_BUILD_VERSION");
|
||||
println!("cargo:rerun-if-env-changed=AETHER_VERSION");
|
||||
println!("cargo:rerun-if-env-changed=GITHUB_REF_NAME");
|
||||
println!("cargo:rerun-if-changed=../../.git/HEAD");
|
||||
|
||||
let package_version = env::var("CARGO_PKG_VERSION").unwrap_or_else(|_| "unknown".to_string());
|
||||
let version = env::var("AETHER_VERSION")
|
||||
let version = env::var("AETHER_BUILD_VERSION")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| {
|
||||
env::var("AETHER_VERSION")
|
||||
.ok()
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
})
|
||||
.or_else(|| {
|
||||
env::var("GITHUB_REF_NAME")
|
||||
.ok()
|
||||
|
||||
@@ -49,10 +49,10 @@ pub(crate) use aether_ai_formats::api::{
|
||||
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
|
||||
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
|
||||
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
|
||||
AiControlPlanRequest, ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind,
|
||||
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
|
||||
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
|
||||
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
sanitize_request_path_and_query, AiControlPlanRequest, ExecutionRuntimeAuthContext,
|
||||
LocalCoreSyncErrorKind, LocalOpenAiImageSpec, LocalSameFormatProviderFamily,
|
||||
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
|
||||
LocalStandardSpec, StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
|
||||
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
|
||||
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
|
||||
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,
|
||||
|
||||
@@ -157,7 +157,9 @@ pub(crate) fn resolve_local_decision_execution_runtime_auth_context(
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<ExecutionRuntimeAuthContext> {
|
||||
resolve_decision_execution_runtime_auth_context(decision).filter(|auth_context| {
|
||||
!auth_context.user_id.trim().is_empty() && !auth_context.api_key_id.trim().is_empty()
|
||||
auth_context.access_allowed
|
||||
&& !auth_context.user_id.trim().is_empty()
|
||||
&& !auth_context.api_key_id.trim().is_empty()
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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};
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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::{
|
||||
|
||||
@@ -502,7 +502,9 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
|
||||
.summarize_usage_audits(&UsageAuditSummaryQuery {
|
||||
created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
..Default::default()
|
||||
user_id: query_param_value(query, "user_id"),
|
||||
provider_name: query_param_value(query, "provider"),
|
||||
model: query_param_value(query, "model"),
|
||||
})
|
||||
.await?;
|
||||
return Ok(Some(build_admin_usage_summary_stats_response_from_summary(
|
||||
|
||||
@@ -1,11 +1,9 @@
|
||||
use crate::handlers::admin::shared::AdminTypedObjectPatch;
|
||||
use crate::provider_key_auth::provider_key_effective_api_formats;
|
||||
use aether_admin::provider::endpoints as admin_provider_endpoints_pure;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
|
||||
pub(super) fn key_api_formats_without_entry(
|
||||
key: &StoredProviderCatalogKey,
|
||||
@@ -15,50 +13,14 @@ pub(super) fn key_api_formats_without_entry(
|
||||
}
|
||||
|
||||
pub(super) fn endpoint_key_counts_by_format(
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
keys: &[StoredProviderCatalogKey],
|
||||
) -> (
|
||||
std::collections::BTreeMap<String, usize>,
|
||||
std::collections::BTreeMap<String, usize>,
|
||||
) {
|
||||
let mut active_endpoint_formats = BTreeSet::new();
|
||||
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
|
||||
active_endpoint_formats.insert(endpoint.api_format.clone());
|
||||
}
|
||||
|
||||
let mut total_by_format = BTreeMap::<String, BTreeSet<String>>::new();
|
||||
let mut active_by_format = BTreeMap::<String, BTreeSet<String>>::new();
|
||||
for key in keys {
|
||||
for api_format in
|
||||
provider_key_effective_api_formats(key, &provider.provider_type, endpoints)
|
||||
{
|
||||
if !active_endpoint_formats.contains(&api_format) {
|
||||
continue;
|
||||
}
|
||||
total_by_format
|
||||
.entry(api_format.clone())
|
||||
.or_default()
|
||||
.insert(key.id.clone());
|
||||
if key.is_active {
|
||||
active_by_format
|
||||
.entry(api_format)
|
||||
.or_default()
|
||||
.insert(key.id.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
(
|
||||
total_by_format
|
||||
.into_iter()
|
||||
.map(|(api_format, keys)| (api_format, keys.len()))
|
||||
.collect(),
|
||||
active_by_format
|
||||
.into_iter()
|
||||
.map(|(api_format, keys)| (api_format, keys.len()))
|
||||
.collect(),
|
||||
)
|
||||
admin_provider_endpoints_pure::endpoint_key_counts_by_format(provider_type, endpoints, keys)
|
||||
}
|
||||
|
||||
pub(super) fn build_admin_provider_endpoint_response(
|
||||
|
||||
@@ -39,7 +39,7 @@ pub(crate) async fn build_admin_provider_endpoints_payload(
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
let (total_keys_by_format, active_keys_by_format) =
|
||||
endpoint_key_counts_by_format(&provider, &endpoints, &keys);
|
||||
endpoint_key_counts_by_format(&provider.provider_type, &endpoints, &keys);
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
@@ -52,15 +52,17 @@ pub(crate) async fn build_admin_provider_endpoints_payload(
|
||||
.skip(skip)
|
||||
.take(limit)
|
||||
.map(|endpoint| {
|
||||
let endpoint_api_format =
|
||||
aether_ai_formats::normalize_api_format_alias(&endpoint.api_format);
|
||||
build_admin_provider_endpoint_response(
|
||||
&endpoint,
|
||||
&provider.name,
|
||||
total_keys_by_format
|
||||
.get(endpoint.api_format.as_str())
|
||||
.get(endpoint_api_format.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
active_keys_by_format
|
||||
.get(endpoint.api_format.as_str())
|
||||
.get(endpoint_api_format.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
now_unix_secs,
|
||||
@@ -93,23 +95,27 @@ pub(crate) async fn build_admin_endpoint_payload(
|
||||
.await
|
||||
.ok()
|
||||
.unwrap_or_default();
|
||||
let (total_keys_by_format, active_keys_by_format) =
|
||||
endpoint_key_counts_by_format(&provider, std::slice::from_ref(&endpoint), &keys);
|
||||
let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format(
|
||||
&provider.provider_type,
|
||||
std::slice::from_ref(&endpoint),
|
||||
&keys,
|
||||
);
|
||||
let now_unix_secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let endpoint_api_format = aether_ai_formats::normalize_api_format_alias(&endpoint.api_format);
|
||||
|
||||
Some(build_admin_provider_endpoint_response(
|
||||
&endpoint,
|
||||
&provider.name,
|
||||
total_keys_by_format
|
||||
.get(endpoint.api_format.as_str())
|
||||
.get(endpoint_api_format.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
active_keys_by_format
|
||||
.get(endpoint.api_format.as_str())
|
||||
.get(endpoint_api_format.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
now_unix_secs,
|
||||
|
||||
@@ -147,19 +147,23 @@ pub(super) async fn maybe_handle(
|
||||
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
let (total_keys_by_format, active_keys_by_format) =
|
||||
endpoint_key_counts_by_format(&provider, std::slice::from_ref(&updated), &keys);
|
||||
let (total_keys_by_format, active_keys_by_format) = endpoint_key_counts_by_format(
|
||||
&provider.provider_type,
|
||||
std::slice::from_ref(&updated),
|
||||
&keys,
|
||||
);
|
||||
let updated_api_format = aether_ai_formats::normalize_api_format_alias(&updated.api_format);
|
||||
|
||||
Ok(Some(
|
||||
Json(build_admin_provider_endpoint_response(
|
||||
&updated,
|
||||
&provider.name,
|
||||
total_keys_by_format
|
||||
.get(updated.api_format.as_str())
|
||||
.get(updated_api_format.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
active_keys_by_format
|
||||
.get(updated.api_format.as_str())
|
||||
.get(updated_api_format.as_str())
|
||||
.copied()
|
||||
.unwrap_or(0),
|
||||
now_unix_secs,
|
||||
|
||||
@@ -30,6 +30,25 @@ use axum::response::IntoResponse;
|
||||
use axum::Json;
|
||||
use serde_json::json;
|
||||
|
||||
async fn apply_supplied_auth_context(
|
||||
state: &AppState,
|
||||
decision: &mut GatewayControlDecision,
|
||||
auth_context: Option<crate::control::GatewayControlAuthContext>,
|
||||
) -> Result<bool, GatewayError> {
|
||||
let Some(auth_context) = auth_context else {
|
||||
return Ok(false);
|
||||
};
|
||||
let refreshed = crate::control::refresh_execution_runtime_auth_context(
|
||||
state,
|
||||
auth_context,
|
||||
decision.auth_endpoint_signature.as_deref(),
|
||||
)
|
||||
.await?;
|
||||
decision.local_auth_rejection = refreshed.local_rejection.clone();
|
||||
decision.auth_context = Some(refreshed);
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
state: &AppState,
|
||||
request_context: &GatewayPublicRequestContext,
|
||||
@@ -206,11 +225,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
Json(build_internal_gateway_fallback_plan_payload(None)).into_response(),
|
||||
));
|
||||
};
|
||||
let provided_auth_context = payload.auth_context.is_some();
|
||||
if let Some(auth_context) = payload.auth_context {
|
||||
resolved.auth_context = Some(auth_context);
|
||||
resolved.local_auth_rejection = None;
|
||||
}
|
||||
let provided_auth_context =
|
||||
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||
let auth_context = resolved.auth_context.as_ref();
|
||||
if auth_context
|
||||
.map(|value| !value.access_allowed)
|
||||
@@ -308,11 +324,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
Json(build_internal_gateway_fallback_plan_payload(None)).into_response(),
|
||||
));
|
||||
};
|
||||
let provided_auth_context = payload.auth_context.is_some();
|
||||
if let Some(auth_context) = payload.auth_context {
|
||||
resolved.auth_context = Some(auth_context);
|
||||
resolved.local_auth_rejection = None;
|
||||
}
|
||||
let provided_auth_context =
|
||||
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||
let auth_context = resolved.auth_context.as_ref();
|
||||
if auth_context
|
||||
.map(|value| !value.access_allowed)
|
||||
@@ -407,11 +420,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
else {
|
||||
return Ok(Some(build_internal_gateway_proxy_public_response()));
|
||||
};
|
||||
let provided_auth_context = payload.auth_context.is_some();
|
||||
if let Some(auth_context) = payload.auth_context {
|
||||
resolved.auth_context = Some(auth_context);
|
||||
resolved.local_auth_rejection = None;
|
||||
}
|
||||
let provided_auth_context =
|
||||
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||
if let Some(mut planned) = api::maybe_build_sync_plan_payload(
|
||||
state,
|
||||
&parts,
|
||||
@@ -474,11 +484,8 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
else {
|
||||
return Ok(Some(build_internal_gateway_proxy_public_response()));
|
||||
};
|
||||
let provided_auth_context = payload.auth_context.is_some();
|
||||
if let Some(auth_context) = payload.auth_context {
|
||||
resolved.auth_context = Some(auth_context);
|
||||
resolved.local_auth_rejection = None;
|
||||
}
|
||||
let provided_auth_context =
|
||||
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||
if let Some(mut planned) = api::maybe_build_stream_plan_payload(
|
||||
state,
|
||||
&parts,
|
||||
@@ -546,10 +553,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(auth_context) = payload.auth_context {
|
||||
resolved.auth_context = Some(auth_context);
|
||||
resolved.local_auth_rejection = None;
|
||||
}
|
||||
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||
if let Some(plan_payload) = api::maybe_build_sync_plan_payload(
|
||||
state,
|
||||
&parts,
|
||||
@@ -630,10 +634,7 @@ pub(crate) async fn maybe_build_local_internal_proxy_response_impl(
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(auth_context) = payload.auth_context {
|
||||
resolved.auth_context = Some(auth_context);
|
||||
resolved.local_auth_rejection = None;
|
||||
}
|
||||
apply_supplied_auth_context(state, &mut resolved, payload.auth_context).await?;
|
||||
if let Some(plan_payload) = api::maybe_build_stream_plan_payload(
|
||||
state,
|
||||
&parts,
|
||||
|
||||
@@ -7,7 +7,7 @@ use crate::constants::{
|
||||
};
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::control::GatewayPublicRequestContext;
|
||||
use crate::middleware::{should_downgrade_access_log, RequestLogEmitted};
|
||||
use crate::middleware::{sanitize_access_log_path, should_downgrade_access_log, RequestLogEmitted};
|
||||
use crate::AppState;
|
||||
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
|
||||
use axum::body::{Body, Bytes};
|
||||
@@ -95,11 +95,12 @@ pub(super) fn finalize_gateway_response(
|
||||
.map(|auth_context| auth_context.api_key_id.as_str())
|
||||
.unwrap_or("-");
|
||||
let status_code = response.status().as_u16();
|
||||
let sanitized_path_and_query = sanitize_access_log_path(path_and_query);
|
||||
emit_admin_audit(
|
||||
&mut response,
|
||||
trace_id,
|
||||
method,
|
||||
path_and_query,
|
||||
sanitized_path_and_query.as_str(),
|
||||
control_decision,
|
||||
);
|
||||
if response.status().is_server_error() {
|
||||
@@ -112,7 +113,7 @@ pub(super) fn finalize_gateway_response(
|
||||
request_id,
|
||||
remote_addr = %remote_addr,
|
||||
method = %method,
|
||||
path = %path_and_query,
|
||||
path = %sanitized_path_and_query,
|
||||
user_id,
|
||||
api_key_id,
|
||||
route_class,
|
||||
@@ -122,7 +123,7 @@ pub(super) fn finalize_gateway_response(
|
||||
elapsed_ms,
|
||||
"gateway request failed"
|
||||
);
|
||||
} else if should_downgrade_access_log(method, path_and_query) {
|
||||
} else if should_downgrade_access_log(method, sanitized_path_and_query.as_str()) {
|
||||
trace!(
|
||||
event_name = "http_request_completed",
|
||||
log_type = "access",
|
||||
@@ -132,7 +133,7 @@ pub(super) fn finalize_gateway_response(
|
||||
request_id,
|
||||
remote_addr = %remote_addr,
|
||||
method = %method,
|
||||
path = %path_and_query,
|
||||
path = %sanitized_path_and_query,
|
||||
user_id,
|
||||
api_key_id,
|
||||
route_class,
|
||||
@@ -152,7 +153,7 @@ pub(super) fn finalize_gateway_response(
|
||||
request_id,
|
||||
remote_addr = %remote_addr,
|
||||
method = %method,
|
||||
path = %path_and_query,
|
||||
path = %sanitized_path_and_query,
|
||||
user_id,
|
||||
api_key_id,
|
||||
route_class,
|
||||
@@ -254,3 +255,106 @@ pub(super) fn finalize_gateway_response_with_context(
|
||||
request_permit,
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::finalize_gateway_response;
|
||||
use crate::control::GatewayControlDecision;
|
||||
use crate::AppState;
|
||||
use axum::body::Body;
|
||||
use axum::http::{Method, Response, StatusCode};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
use tracing_subscriber::filter::LevelFilter;
|
||||
use tracing_subscriber::prelude::*;
|
||||
|
||||
#[derive(Clone, Default)]
|
||||
struct SharedBuffer(Arc<Mutex<Vec<u8>>>);
|
||||
|
||||
struct SharedBufferWriter(Arc<Mutex<Vec<u8>>>);
|
||||
|
||||
impl SharedBuffer {
|
||||
fn lines(&self) -> Vec<serde_json::Value> {
|
||||
String::from_utf8(self.0.lock().expect("buffer should lock").clone())
|
||||
.expect("buffer should contain valid utf-8")
|
||||
.lines()
|
||||
.filter(|line| !line.trim().is_empty())
|
||||
.map(|line| serde_json::from_str(line).expect("json log line should parse"))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
impl std::io::Write for SharedBufferWriter {
|
||||
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
||||
self.0
|
||||
.lock()
|
||||
.expect("buffer should lock")
|
||||
.extend_from_slice(buf);
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> std::io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<'a> tracing_subscriber::fmt::writer::MakeWriter<'a> for SharedBuffer {
|
||||
type Writer = SharedBufferWriter;
|
||||
|
||||
fn make_writer(&'a self) -> Self::Writer {
|
||||
SharedBufferWriter(Arc::clone(&self.0))
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finalize_gateway_response_logs_sanitized_path_and_query() {
|
||||
let state = AppState::new().expect("gateway state should build");
|
||||
let writer = SharedBuffer::default();
|
||||
let subscriber = tracing_subscriber::registry().with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.json()
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
let remote_addr = "127.0.0.1:8080"
|
||||
.parse()
|
||||
.expect("remote address should parse");
|
||||
let control_decision = GatewayControlDecision::synthetic(
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent",
|
||||
Some("ai_public".to_string()),
|
||||
Some("gemini".to_string()),
|
||||
Some("generate_content".to_string()),
|
||||
Some("gemini:generate_content".to_string()),
|
||||
);
|
||||
|
||||
let response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.body(Body::empty())
|
||||
.expect("response should build");
|
||||
|
||||
let _response = finalize_gateway_response(
|
||||
&state,
|
||||
response,
|
||||
"trace-finalize",
|
||||
&remote_addr,
|
||||
&Method::GET,
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse",
|
||||
Some(&control_decision),
|
||||
"execution_runtime_sync",
|
||||
&Instant::now(),
|
||||
None,
|
||||
);
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(
|
||||
logs[0]["path"],
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -30,10 +30,10 @@ mod refunds;
|
||||
use self::flow::handle_wallet_flow;
|
||||
pub(in crate::handlers::public::support) use self::reads::build_wallet_balance_payload_for_user;
|
||||
use self::reads::{
|
||||
build_wallet_daily_usage_payload, build_wallet_payload, build_wallet_zero_today_entry,
|
||||
handle_wallet_balance, handle_wallet_today_cost, handle_wallet_transactions,
|
||||
parse_wallet_limit, parse_wallet_offset, wallet_fixed_offset, wallet_today_billing_date_string,
|
||||
wallet_transaction_payload_from_record,
|
||||
build_wallet_daily_usage_payload, build_wallet_live_today_usage_payload_for_user,
|
||||
build_wallet_payload, build_wallet_zero_today_entry, handle_wallet_balance,
|
||||
handle_wallet_today_cost, handle_wallet_transactions, parse_wallet_limit, parse_wallet_offset,
|
||||
wallet_fixed_offset, wallet_transaction_payload_from_record,
|
||||
};
|
||||
pub(crate) use self::recharge::sanitize_wallet_gateway_response;
|
||||
use self::recharge::{
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
use super::{
|
||||
build_auth_error_response, build_auth_json_response, build_wallet_daily_usage_payload,
|
||||
build_wallet_payload, build_wallet_zero_today_entry, http, parse_wallet_limit,
|
||||
parse_wallet_offset, resolve_authenticated_local_user, unix_secs_to_rfc3339,
|
||||
wallet_fixed_offset, wallet_today_billing_date_string, wallet_transaction_payload_from_record,
|
||||
AppState, Body, GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
|
||||
build_wallet_live_today_usage_payload_for_user, build_wallet_payload,
|
||||
build_wallet_zero_today_entry, http, parse_wallet_limit, parse_wallet_offset,
|
||||
resolve_authenticated_local_user, unix_secs_to_rfc3339, wallet_fixed_offset,
|
||||
wallet_transaction_payload_from_record, AppState, Body, GatewayPublicRequestContext, Response,
|
||||
WALLET_LEGACY_TIMEZONE,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
@@ -98,32 +99,43 @@ pub(super) async fn handle_wallet_flow(
|
||||
return build_auth_json_response(http::StatusCode::OK, payload, None);
|
||||
};
|
||||
|
||||
let mut today_entry = build_wallet_zero_today_entry();
|
||||
if let Ok(Some(today_usage)) = state
|
||||
.find_wallet_today_usage(&wallet.id, WALLET_LEGACY_TIMEZONE)
|
||||
.await
|
||||
let mut today_entry =
|
||||
match build_wallet_live_today_usage_payload_for_user(state, &auth.user.id).await {
|
||||
Ok(Some(today_usage)) => today_usage,
|
||||
_ => build_wallet_zero_today_entry(),
|
||||
};
|
||||
if today_entry
|
||||
.get("total_requests")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap_or_default()
|
||||
== 0
|
||||
{
|
||||
today_entry = build_wallet_daily_usage_payload(
|
||||
today_usage.id,
|
||||
today_usage.billing_date,
|
||||
today_usage.billing_timezone,
|
||||
today_usage.total_cost_usd,
|
||||
today_usage.total_requests,
|
||||
today_usage.input_tokens,
|
||||
today_usage.output_tokens,
|
||||
today_usage.cache_creation_tokens,
|
||||
today_usage.cache_read_tokens,
|
||||
today_usage
|
||||
.first_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
today_usage
|
||||
.last_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
today_usage
|
||||
.aggregated_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
true,
|
||||
);
|
||||
if let Ok(Some(today_usage)) = state
|
||||
.find_wallet_today_usage(&wallet.id, WALLET_LEGACY_TIMEZONE)
|
||||
.await
|
||||
{
|
||||
today_entry = build_wallet_daily_usage_payload(
|
||||
today_usage.id,
|
||||
today_usage.billing_date,
|
||||
today_usage.billing_timezone,
|
||||
today_usage.total_cost_usd,
|
||||
today_usage.total_requests,
|
||||
today_usage.input_tokens,
|
||||
today_usage.output_tokens,
|
||||
today_usage.cache_creation_tokens,
|
||||
today_usage.cache_read_tokens,
|
||||
today_usage
|
||||
.first_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
today_usage
|
||||
.last_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
today_usage
|
||||
.aggregated_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
true,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let fetch_size = offset.saturating_add(limit).min(5200);
|
||||
|
||||
@@ -1,11 +1,11 @@
|
||||
use super::{
|
||||
build_auth_error_response, build_auth_json_response, build_auth_wallet_summary_payload, http,
|
||||
query_param_value, resolve_authenticated_local_user, unix_secs_to_rfc3339, AppState, Body,
|
||||
GatewayError, GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
|
||||
GatewayPublicRequestContext, Response, WALLET_LEGACY_TIMEZONE,
|
||||
};
|
||||
use crate::handlers::shared::round_to;
|
||||
use aether_data_contracts::repository::usage::UsageSettledCostSummaryQuery;
|
||||
use chrono::Utc;
|
||||
use chrono::{TimeZone, Utc};
|
||||
use serde_json::json;
|
||||
|
||||
const WALLET_TODAY_COST_UNAVAILABLE_DETAIL: &str = "钱包今日费用数据暂不可用";
|
||||
@@ -143,6 +143,26 @@ pub(super) fn wallet_today_billing_date_string() -> String {
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn wallet_today_usage_window() -> Result<(String, String, u64, u64), String> {
|
||||
let offset = wallet_fixed_offset();
|
||||
let today = Utc::now().with_timezone(&offset).date_naive();
|
||||
let Some(local_start_naive) = today.and_hms_opt(0, 0, 0) else {
|
||||
return Err("wallet today start is invalid".to_string());
|
||||
};
|
||||
let Some(local_start) = offset.from_local_datetime(&local_start_naive).single() else {
|
||||
return Err("wallet today local start is ambiguous".to_string());
|
||||
};
|
||||
let local_end = local_start + chrono::Duration::days(1);
|
||||
let start_unix_secs = local_start.timestamp().max(0) as u64;
|
||||
let end_unix_secs = local_end.timestamp().max(0) as u64;
|
||||
Ok((
|
||||
today.to_string(),
|
||||
WALLET_LEGACY_TIMEZONE.to_string(),
|
||||
start_unix_secs,
|
||||
end_unix_secs,
|
||||
))
|
||||
}
|
||||
|
||||
pub(super) fn build_wallet_daily_usage_payload(
|
||||
id: Option<String>,
|
||||
date: String,
|
||||
@@ -193,6 +213,43 @@ pub(super) fn build_wallet_zero_today_entry() -> serde_json::Value {
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) async fn build_wallet_live_today_usage_payload_for_user(
|
||||
state: &AppState,
|
||||
user_id: &str,
|
||||
) -> Result<Option<serde_json::Value>, String> {
|
||||
if !state.has_usage_data_reader() {
|
||||
return Ok(None);
|
||||
}
|
||||
let (date, timezone, start_unix_secs, end_unix_secs) = wallet_today_usage_window()?;
|
||||
let summary = state
|
||||
.summarize_usage_settled_cost(&UsageSettledCostSummaryQuery {
|
||||
created_from_unix_secs: start_unix_secs,
|
||||
created_until_unix_secs: end_unix_secs,
|
||||
user_id: Some(user_id.to_string()),
|
||||
})
|
||||
.await
|
||||
.map_err(|err| format!("wallet today cost lookup failed: {err:?}"))?;
|
||||
Ok(Some(build_wallet_daily_usage_payload(
|
||||
None,
|
||||
date,
|
||||
timezone,
|
||||
summary.total_cost_usd,
|
||||
summary.total_requests,
|
||||
summary.input_tokens,
|
||||
summary.output_tokens,
|
||||
summary.cache_creation_tokens,
|
||||
summary.cache_read_tokens,
|
||||
summary
|
||||
.first_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
summary
|
||||
.last_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339),
|
||||
Some(Utc::now().to_rfc3339()),
|
||||
true,
|
||||
)))
|
||||
}
|
||||
|
||||
pub(super) fn wallet_transaction_payload_from_record(
|
||||
record: &aether_data::repository::wallet::StoredAdminWalletTransaction,
|
||||
) -> serde_json::Value {
|
||||
@@ -253,65 +310,17 @@ pub(super) async fn handle_wallet_today_cost(
|
||||
Ok(value) => value,
|
||||
Err(response) => return response,
|
||||
};
|
||||
let today = Utc::now().date_naive();
|
||||
let Some(start_of_day) = today.and_hms_opt(0, 0, 0) else {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
"wallet today start is invalid",
|
||||
match build_wallet_live_today_usage_payload_for_user(state, &auth.user.id).await {
|
||||
Ok(Some(payload)) => build_auth_json_response(http::StatusCode::OK, payload, None),
|
||||
Ok(None) => build_auth_error_response(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
WALLET_TODAY_COST_UNAVAILABLE_DETAIL,
|
||||
false,
|
||||
);
|
||||
};
|
||||
let start_unix_secs = u64::try_from(
|
||||
chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(start_of_day, chrono::Utc)
|
||||
.timestamp(),
|
||||
)
|
||||
.unwrap_or_default();
|
||||
let end_unix_secs = start_unix_secs.saturating_add(24 * 3600);
|
||||
|
||||
let summary = match state
|
||||
.summarize_usage_settled_cost(&UsageSettledCostSummaryQuery {
|
||||
created_from_unix_secs: start_unix_secs,
|
||||
created_until_unix_secs: end_unix_secs,
|
||||
user_id: Some(auth.user.id.clone()),
|
||||
})
|
||||
.await
|
||||
{
|
||||
Ok(value) => value,
|
||||
Err(err) => {
|
||||
return build_auth_error_response(
|
||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
format!("wallet today cost lookup failed: {err:?}"),
|
||||
false,
|
||||
)
|
||||
),
|
||||
Err(detail) => {
|
||||
build_auth_error_response(http::StatusCode::INTERNAL_SERVER_ERROR, detail, false)
|
||||
}
|
||||
};
|
||||
|
||||
let first_finalized_at = summary
|
||||
.first_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
let last_finalized_at = summary
|
||||
.last_finalized_at_unix_secs
|
||||
.and_then(unix_secs_to_rfc3339);
|
||||
|
||||
build_auth_json_response(
|
||||
http::StatusCode::OK,
|
||||
json!({
|
||||
"id": serde_json::Value::Null,
|
||||
"date": today.to_string(),
|
||||
"timezone": "UTC",
|
||||
"total_cost": round_to(summary.total_cost_usd, 6),
|
||||
"total_requests": summary.total_requests,
|
||||
"input_tokens": summary.input_tokens,
|
||||
"output_tokens": summary.output_tokens,
|
||||
"cache_creation_tokens": summary.cache_creation_tokens,
|
||||
"cache_read_tokens": summary.cache_read_tokens,
|
||||
"first_finalized_at": first_finalized_at,
|
||||
"last_finalized_at": last_finalized_at,
|
||||
"aggregated_at": Utc::now().to_rfc3339(),
|
||||
"is_today": true,
|
||||
}),
|
||||
None,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) async fn handle_wallet_transactions(
|
||||
|
||||
@@ -9,7 +9,10 @@ use clap::{Args as ClapArgs, Parser, Subcommand, ValueEnum};
|
||||
use tracing::{debug, info, warn};
|
||||
|
||||
use aether_crypto::warm_python_fernet_secret;
|
||||
use aether_data::lifecycle::export::{export_database_jsonl, import_database_jsonl, ExportDomain};
|
||||
use aether_data::lifecycle::export::{
|
||||
copy_database_records, export_database_jsonl, import_database_jsonl, DataCopyOptions,
|
||||
ExportDomain,
|
||||
};
|
||||
use aether_data::{DatabaseDriver, SqlDatabaseConfig, SqlPoolConfig, DEFAULT_SQLITE_DATABASE_URL};
|
||||
use aether_gateway::{
|
||||
attach_static_frontend, build_router_with_state, set_gateway_frontdoor_app_port, AppState,
|
||||
@@ -579,6 +582,8 @@ enum DataCommand {
|
||||
Export(DataExportArgs),
|
||||
/// Import database-neutral JSONL into the selected SQL database.
|
||||
Import(DataImportArgs),
|
||||
/// Copy persistent SQL data directly between two databases without a JSONL file.
|
||||
Copy(DataCopyArgs),
|
||||
}
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
@@ -602,6 +607,27 @@ struct DataImportArgs {
|
||||
input: PathBuf,
|
||||
}
|
||||
|
||||
#[derive(ClapArgs, Debug, Clone)]
|
||||
struct DataCopyArgs {
|
||||
#[arg(long, value_enum)]
|
||||
source_driver: DatabaseDriverArg,
|
||||
|
||||
#[arg(long)]
|
||||
source_url: String,
|
||||
|
||||
#[arg(long, value_enum)]
|
||||
target_driver: DatabaseDriverArg,
|
||||
|
||||
#[arg(long)]
|
||||
target_url: String,
|
||||
|
||||
#[arg(long, value_enum, value_delimiter = ',')]
|
||||
domains: Vec<ExportDomainArg>,
|
||||
|
||||
#[arg(long)]
|
||||
omit_request_body_details: bool,
|
||||
}
|
||||
|
||||
impl GatewayLoggingArgs {
|
||||
fn apply_to_runtime_config(
|
||||
&self,
|
||||
@@ -1251,6 +1277,7 @@ async fn run_data_command(command: &DataCommand) -> Result<(), Box<dyn std::erro
|
||||
match command {
|
||||
DataCommand::Export(args) => run_data_export(args).await,
|
||||
DataCommand::Import(args) => run_data_import(args).await,
|
||||
DataCommand::Copy(args) => run_data_copy(args).await,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1267,11 +1294,11 @@ fn required_sql_database_config(
|
||||
}
|
||||
|
||||
fn requested_export_domains(args: &DataExportArgs) -> Vec<ExportDomain> {
|
||||
args.domains
|
||||
.iter()
|
||||
.copied()
|
||||
.map(Into::into)
|
||||
.collect::<Vec<_>>()
|
||||
requested_domains(&args.domains)
|
||||
}
|
||||
|
||||
fn requested_domains(domains: &[ExportDomainArg]) -> Vec<ExportDomain> {
|
||||
domains.iter().copied().map(Into::into).collect::<Vec<_>>()
|
||||
}
|
||||
|
||||
fn current_unix_secs() -> Result<u64, std::time::SystemTimeError> {
|
||||
@@ -1324,6 +1351,61 @@ async fn run_data_import(args: &DataImportArgs) -> Result<(), Box<dyn std::error
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn copy_database_config(
|
||||
driver: DatabaseDriverArg,
|
||||
url: &str,
|
||||
label: &str,
|
||||
) -> Result<SqlDatabaseConfig, Box<dyn std::error::Error>> {
|
||||
let url = url.trim();
|
||||
if url.is_empty() {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidInput,
|
||||
format!("{label} database URL must not be empty"),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
let driver = DatabaseDriver::from(driver);
|
||||
Ok(SqlDatabaseConfig::new(
|
||||
driver,
|
||||
url,
|
||||
SqlPoolConfig {
|
||||
require_ssl: false,
|
||||
..SqlPoolConfig::default()
|
||||
},
|
||||
)?)
|
||||
}
|
||||
|
||||
async fn run_data_copy(args: &DataCopyArgs) -> Result<(), Box<dyn std::error::Error>> {
|
||||
let source = copy_database_config(args.source_driver, &args.source_url, "source")?;
|
||||
let target = copy_database_config(args.target_driver, &args.target_url, "target")?;
|
||||
let source_driver = source.driver;
|
||||
let target_driver = target.driver;
|
||||
let domains = requested_domains(&args.domains);
|
||||
let created_at_unix_secs = current_unix_secs()?;
|
||||
let imported = copy_database_records(
|
||||
source,
|
||||
target,
|
||||
domains,
|
||||
created_at_unix_secs,
|
||||
DataCopyOptions {
|
||||
omit_request_body_details: args.omit_request_body_details,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
||||
info!(
|
||||
source_driver = %source_driver,
|
||||
target_driver = %target_driver,
|
||||
imported,
|
||||
"database copy complete"
|
||||
);
|
||||
println!(
|
||||
"copied {} records from {} to {} without a JSONL file",
|
||||
imported, source_driver, target_driver
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn run_explicit_migrations(args: &Args) -> Result<(), Box<dyn std::error::Error>> {
|
||||
if args.data.effective_sql_database_config().is_none() {
|
||||
return Err(std::io::Error::new(
|
||||
|
||||
@@ -8,6 +8,7 @@ use axum::middleware::Next;
|
||||
use axum::response::Response;
|
||||
use tracing::{info, trace, warn};
|
||||
|
||||
use crate::ai_serving::api::sanitize_request_path_and_query;
|
||||
use crate::constants::{
|
||||
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
|
||||
};
|
||||
@@ -52,14 +53,19 @@ pub(crate) fn should_downgrade_access_log(method: &Method, path: &str) -> bool {
|
||||
|| normalized_path.starts_with("/api/admin/monitoring/trace/")
|
||||
}
|
||||
|
||||
pub(crate) fn sanitize_access_log_path(path: &str) -> String {
|
||||
sanitize_request_path_and_query(path, None).unwrap_or_else(|| "/".to_string())
|
||||
}
|
||||
|
||||
pub(crate) async fn access_log_middleware(mut request: Request<Body>, next: Next) -> Response {
|
||||
let started_at = Instant::now();
|
||||
let method = request.method().clone();
|
||||
let path = request
|
||||
let raw_path = request
|
||||
.uri()
|
||||
.path_and_query()
|
||||
.map(|value| value.as_str().to_string())
|
||||
.unwrap_or_else(|| "/".to_string());
|
||||
let path = sanitize_access_log_path(&raw_path);
|
||||
let trace_id = extract_or_generate_trace_id(request.headers());
|
||||
if !request.headers().contains_key(TRACE_ID_HEADER) {
|
||||
request.headers_mut().insert(
|
||||
@@ -158,7 +164,7 @@ pub(crate) async fn access_log_middleware(mut request: Request<Body>, next: Next
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{access_log_middleware, should_downgrade_access_log};
|
||||
use super::{access_log_middleware, sanitize_access_log_path, should_downgrade_access_log};
|
||||
use crate::constants::{
|
||||
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER,
|
||||
TRACE_ID_HEADER,
|
||||
@@ -212,6 +218,56 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn access_log_path_redacts_credential_query_values() {
|
||||
assert_eq!(
|
||||
sanitize_access_log_path(
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse&pageSize=10&token=hidden"
|
||||
),
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse&pageSize=10"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn access_log_emits_sanitized_path() {
|
||||
let writer = SharedBuffer::default();
|
||||
let subscriber = tracing_subscriber::registry().with(
|
||||
tracing_subscriber::fmt::layer()
|
||||
.json()
|
||||
.flatten_event(true)
|
||||
.with_current_span(false)
|
||||
.with_span_list(false)
|
||||
.with_writer(writer.clone())
|
||||
.with_filter(LevelFilter::INFO),
|
||||
);
|
||||
let dispatch = tracing::Dispatch::new(subscriber);
|
||||
let _guard = tracing::dispatcher::set_default(&dispatch);
|
||||
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent",
|
||||
get(|| async { Response::new(Body::empty()) }),
|
||||
)
|
||||
.layer(axum::middleware::from_fn(access_log_middleware));
|
||||
|
||||
let _response = app
|
||||
.oneshot(
|
||||
Request::builder()
|
||||
.uri("/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse")
|
||||
.body(Body::empty())
|
||||
.expect("request should build"),
|
||||
)
|
||||
.await
|
||||
.expect("request should succeed");
|
||||
|
||||
let logs = writer.lines();
|
||||
assert_eq!(logs.len(), 1);
|
||||
assert_eq!(
|
||||
logs[0]["path"],
|
||||
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "current_thread")]
|
||||
async fn access_log_emits_completed_events_by_default() {
|
||||
let writer = SharedBuffer::default();
|
||||
|
||||
@@ -3,7 +3,7 @@ mod frontdoor_cors;
|
||||
mod strip_cf_headers;
|
||||
|
||||
pub(crate) use access_log::{
|
||||
access_log_middleware, should_downgrade_access_log, RequestLogEmitted,
|
||||
access_log_middleware, sanitize_access_log_path, should_downgrade_access_log, RequestLogEmitted,
|
||||
};
|
||||
pub(crate) use frontdoor_cors::frontdoor_cors_middleware;
|
||||
pub use strip_cf_headers::strip_cf_headers_middleware;
|
||||
|
||||
@@ -1,5 +1,31 @@
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn admin_system_build_version_contract_uses_explicit_local_build_arg() {
|
||||
let build_rs = read_workspace_file("apps/aether-gateway/build.rs");
|
||||
for pattern in [
|
||||
"cargo:rerun-if-env-changed=AETHER_BUILD_VERSION",
|
||||
"env::var(\"AETHER_BUILD_VERSION\")",
|
||||
] {
|
||||
assert!(
|
||||
build_rs.contains(pattern),
|
||||
"apps/aether-gateway/build.rs should consume explicit build version pattern {pattern}"
|
||||
);
|
||||
}
|
||||
|
||||
let dockerfile = read_workspace_file("Dockerfile.app.local");
|
||||
for pattern in [
|
||||
"ARG AETHER_BUILD_VERSION",
|
||||
"ENV AETHER_BUILD_VERSION=${AETHER_BUILD_VERSION}",
|
||||
"AETHER_VERSION=${AETHER_BUILD_VERSION}",
|
||||
] {
|
||||
assert!(
|
||||
dockerfile.contains(pattern),
|
||||
"Dockerfile.app.local should pass explicit build version pattern {pattern}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_system_and_endpoint_roots_stay_thin() {
|
||||
let system_mod = read_workspace_file("apps/aether-gateway/src/handlers/admin/system/mod.rs");
|
||||
|
||||
@@ -262,6 +262,91 @@ async fn gateway_counts_fixed_provider_oauth_keys_for_inherited_endpoint_formats
|
||||
gateway_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_counts_keys_with_null_api_formats_for_each_fixed_provider_endpoint() {
|
||||
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/endpoints",
|
||||
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 inherited_key = sample_key(
|
||||
"key-codex-oauth",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"codex-token",
|
||||
);
|
||||
inherited_key.auth_type = "oauth".to_string();
|
||||
inherited_key.api_formats = None;
|
||||
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![
|
||||
sample_endpoint(
|
||||
"endpoint-codex-responses",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
)
|
||||
.with_timestamps(Some(1_711_000_000), Some(1_711_000_100)),
|
||||
sample_endpoint(
|
||||
"endpoint-codex-image",
|
||||
"provider-codex",
|
||||
"openai:image",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
)
|
||||
.with_timestamps(Some(1_710_000_000), Some(1_710_000_100)),
|
||||
],
|
||||
vec![inherited_key],
|
||||
));
|
||||
|
||||
let (_, 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,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!(
|
||||
"{gateway_url}/api/admin/endpoints/providers/provider-codex/endpoints?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(), 2);
|
||||
assert_eq!(items[0]["api_format"], "openai:responses");
|
||||
assert_eq!(items[0]["total_keys"], 1);
|
||||
assert_eq!(items[0]["active_keys"], 1);
|
||||
assert_eq!(items[1]["api_format"], "openai:image");
|
||||
assert_eq!(items[1]["total_keys"], 1);
|
||||
assert_eq!(items[1]["active_keys"], 1);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_service_unavailable_for_admin_provider_endpoint_create_when_catalog_writer_unavailable(
|
||||
) {
|
||||
|
||||
@@ -387,6 +387,71 @@ async fn gateway_handles_admin_usage_stats_locally_with_trusted_admin_principal(
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_filters_admin_usage_stats_by_query_fields_locally() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
start_usage_upstream("/api/admin/usage/stats").await;
|
||||
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_usage_row(
|
||||
"usage-1",
|
||||
"req-1",
|
||||
Some("user-1"),
|
||||
Some("key-1"),
|
||||
Some("primary"),
|
||||
"OpenAI",
|
||||
"gpt-5",
|
||||
"completed",
|
||||
120,
|
||||
30,
|
||||
0.3,
|
||||
0.36,
|
||||
DAY_1_UNIX_SECS,
|
||||
),
|
||||
sample_usage_row(
|
||||
"usage-2",
|
||||
"req-2",
|
||||
Some("user-2"),
|
||||
Some("key-2"),
|
||||
Some("secondary"),
|
||||
"Anthropic",
|
||||
"claude-3-7",
|
||||
"failed",
|
||||
40,
|
||||
10,
|
||||
0.1,
|
||||
0.12,
|
||||
DAY_2_UNIX_SECS,
|
||||
),
|
||||
]));
|
||||
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(GatewayDataState::with_usage_reader_for_tests(
|
||||
usage_repository,
|
||||
)),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = admin_request(reqwest::Client::new().get(format!(
|
||||
"{gateway_url}/api/admin/usage/stats?start_date=2024-03-21&end_date=2024-03-22&tz_offset_minutes=0&user_id=user-2&provider=Anthropic&model=claude-3-7"
|
||||
)))
|
||||
.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["total_requests"], 1);
|
||||
assert_eq!(payload["total_tokens"], 70);
|
||||
assert_eq!(payload["error_count"], 1);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_defaults_admin_usage_stats_to_bounded_recent_window_when_query_missing() {
|
||||
let (upstream_url, upstream_hits, upstream_handle) =
|
||||
|
||||
@@ -11,6 +11,9 @@ use crate::tests::{
|
||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||
EXECUTION_PATH_HEADER,
|
||||
};
|
||||
use aether_data::repository::billing::InMemoryBillingReadRepository;
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use base64::Engine as _;
|
||||
|
||||
#[tokio::test]
|
||||
@@ -1154,6 +1157,118 @@ async fn gateway_handles_internal_gateway_decision_sync_locally_with_supplied_au
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_internal_decision_sync_revalidates_supplied_auth_context_wallet() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
let upstream_hits_clone = Arc::clone(&upstream_hits);
|
||||
let upstream = Router::new().route(
|
||||
"/{*path}",
|
||||
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("proxied"))
|
||||
}
|
||||
}),
|
||||
);
|
||||
|
||||
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
|
||||
None,
|
||||
unrestricted_models_snapshot("api-key-empty-wallet", "user-empty-wallet"),
|
||||
)]));
|
||||
let candidate_repository =
|
||||
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
|
||||
sample_models_candidate_row("provider-1", "openai", "openai:chat", "gpt-5", 10),
|
||||
]));
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider("provider-1", "openai", 10)],
|
||||
vec![sample_endpoint(
|
||||
"endpoint-provider-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"https://api.openai.example",
|
||||
)],
|
||||
vec![sample_key(
|
||||
"key-provider-1",
|
||||
"provider-1",
|
||||
"openai:chat",
|
||||
"sk-upstream-openai",
|
||||
)],
|
||||
));
|
||||
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let billing_repository = Arc::new(InMemoryBillingReadRepository::seed(Vec::new()));
|
||||
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 (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(
|
||||
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_usage_billing_and_wallet_for_tests(
|
||||
auth_repository,
|
||||
candidate_repository,
|
||||
provider_catalog_repository,
|
||||
request_candidate_repository,
|
||||
usage_repository,
|
||||
billing_repository,
|
||||
wallet_repository,
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
),
|
||||
),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{gateway_url}/api/internal/gateway/decision-sync"))
|
||||
.json(&json!({
|
||||
"trace_id": "trace-internal-decision-sync-empty-wallet",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {
|
||||
"content-type": "application/json",
|
||||
},
|
||||
"body_json": {
|
||||
"model": "gpt-5",
|
||||
"messages": [],
|
||||
},
|
||||
"auth_context": {
|
||||
"user_id": "user-empty-wallet",
|
||||
"api_key_id": "api-key-empty-wallet",
|
||||
"access_allowed": true,
|
||||
}
|
||||
}))
|
||||
.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["action"], "fallback_plan");
|
||||
assert_eq!(payload["auth_context"], serde_json::Value::Null);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_auth_context() {
|
||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||
|
||||
@@ -4133,13 +4133,7 @@ async fn gateway_handles_wallet_balance_locally_without_proxying_upstream() {
|
||||
#[tokio::test]
|
||||
async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
||||
let auth_now = Utc::now();
|
||||
let usage_now = chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
|
||||
auth_now
|
||||
.date_naive()
|
||||
.and_hms_opt(12, 0, 0)
|
||||
.expect("midday should be valid"),
|
||||
chrono::Utc,
|
||||
);
|
||||
let usage_now = auth_now - chrono::Duration::minutes(30);
|
||||
let user = sample_auth_user(auth_now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
@@ -4211,6 +4205,72 @@ async fn gateway_handles_wallet_today_cost_locally_without_proxying_upstream() {
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_wallet_flow_today_entry_uses_live_settled_usage() {
|
||||
let auth_now = Utc::now();
|
||||
let user = sample_auth_user(auth_now);
|
||||
let access_token = build_test_auth_token(
|
||||
"access",
|
||||
serde_json::Map::from_iter([
|
||||
("user_id".to_string(), json!(user.id)),
|
||||
("role".to_string(), json!(user.role)),
|
||||
(
|
||||
"created_at".to_string(),
|
||||
json!(user.created_at.map(|value| value.to_rfc3339())),
|
||||
),
|
||||
("session_id".to_string(), json!("session-wallet-flow-today")),
|
||||
]),
|
||||
auth_now + chrono::Duration::hours(1),
|
||||
);
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||
sample_user_usage_audit(
|
||||
"usage-wallet-flow-today",
|
||||
"req-wallet-flow-today",
|
||||
"user-auth-1",
|
||||
"gpt-4.1",
|
||||
"OpenAI",
|
||||
"completed",
|
||||
auth_now - chrono::Duration::minutes(5),
|
||||
),
|
||||
]));
|
||||
let (gateway_url, upstream_hits, gateway_handle, upstream_handle) =
|
||||
start_auth_gateway_with_usage_state(
|
||||
user,
|
||||
sample_auth_wallet("user-auth-1", auth_now),
|
||||
[sample_auth_session(
|
||||
"user-auth-1",
|
||||
"session-wallet-flow-today",
|
||||
"device-wallet-flow-today",
|
||||
"refresh-token-placeholder",
|
||||
auth_now,
|
||||
)],
|
||||
usage_repository,
|
||||
)
|
||||
.await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{gateway_url}/api/wallet/flow?limit=20&offset=0"))
|
||||
.header("authorization", format!("Bearer {access_token}"))
|
||||
.header("x-client-device-id", "device-wallet-flow-today")
|
||||
.header("user-agent", "AetherTest/1.0")
|
||||
.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["today_entry"]["is_today"], true);
|
||||
assert_eq!(payload["today_entry"]["timezone"], "Asia/Shanghai");
|
||||
assert_eq!(payload["today_entry"]["total_requests"], 1);
|
||||
assert_eq!(payload["today_entry"]["input_tokens"], 120);
|
||||
assert_eq!(payload["today_entry"]["cache_read_tokens"], 15);
|
||||
assert_eq!(payload["today_entry"]["total_cost"], 1.25);
|
||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||
|
||||
gateway_handle.abort();
|
||||
upstream_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_returns_service_unavailable_for_wallet_today_cost_without_usage_reader() {
|
||||
let now = Utc::now();
|
||||
|
||||
@@ -24,14 +24,9 @@ pub(crate) async fn resolve_wallet_auth_gate(
|
||||
auth_snapshot.api_key_is_standalone,
|
||||
)
|
||||
.await?;
|
||||
let is_admin = wallet_auth_allows_admin_bypass(
|
||||
&auth_snapshot.user_role,
|
||||
auth_snapshot.api_key_is_standalone,
|
||||
);
|
||||
|
||||
let decision = match wallet.as_ref() {
|
||||
Some(wallet) => map_wallet_snapshot(wallet).access_decision(is_admin),
|
||||
None if is_admin => WalletAccessDecision::allowed(None),
|
||||
Some(wallet) => map_wallet_snapshot(wallet).access_decision(false),
|
||||
None => WalletAccessDecision::wallet_unavailable(None),
|
||||
};
|
||||
if !auth_snapshot.api_key_is_standalone {
|
||||
@@ -83,19 +78,50 @@ fn map_wallet_snapshot(snapshot: &StoredWalletSnapshot) -> WalletSnapshot {
|
||||
}
|
||||
}
|
||||
|
||||
fn wallet_auth_allows_admin_bypass(user_role: &str, api_key_is_standalone: bool) -> bool {
|
||||
user_role.eq_ignore_ascii_case("admin") && !api_key_is_standalone
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_data::repository::wallet::StoredWalletSnapshot;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_data::repository::usage::InMemoryUsageReadRepository;
|
||||
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
|
||||
use aether_data_contracts::repository::billing::{
|
||||
BillingReadRepository, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
|
||||
};
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use aether_wallet::{WalletAccessFailure, WalletLimitMode, WalletSnapshot, WalletStatus};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
local_rejection_from_wallet_access, map_wallet_snapshot, wallet_auth_allows_admin_bypass,
|
||||
local_rejection_from_wallet_access, map_wallet_snapshot, resolve_wallet_auth_gate,
|
||||
};
|
||||
use crate::control::GatewayLocalAuthRejection;
|
||||
use crate::data::auth::GatewayAuthApiKeySnapshot;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::AppState;
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FixedQuotaBillingReadRepository {
|
||||
quota: Option<UserDailyQuotaAvailabilityRecord>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl BillingReadRepository for FixedQuotaBillingReadRepository {
|
||||
async fn find_model_context(
|
||||
&self,
|
||||
_provider_id: &str,
|
||||
_provider_api_key_id: Option<&str>,
|
||||
_global_model_name: &str,
|
||||
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
async fn find_user_daily_quota_availability(
|
||||
&self,
|
||||
_user_id: &str,
|
||||
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, DataLayerError> {
|
||||
Ok(self.quota.clone())
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_wallet_snapshot_and_derives_balance_denied() {
|
||||
@@ -145,8 +171,143 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn standalone_key_never_uses_admin_wallet_bypass() {
|
||||
assert!(wallet_auth_allows_admin_bypass("admin", false));
|
||||
assert!(!wallet_auth_allows_admin_bypass("admin", true));
|
||||
fn admin_user_with_empty_finite_wallet_is_balance_denied() {
|
||||
let stored = StoredWalletSnapshot::new(
|
||||
"wallet-1".to_string(),
|
||||
Some("admin-1".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 decision = map_wallet_snapshot(&stored).access_decision(true);
|
||||
|
||||
assert_eq!(decision.failure, Some(WalletAccessFailure::BalanceDenied));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ordinary_user_key_without_quota_denies_empty_wallet() {
|
||||
let state = state_with_wallet_and_quota(empty_user_wallet(), None);
|
||||
let auth_snapshot = ordinary_user_api_key_snapshot();
|
||||
|
||||
let decision = resolve_wallet_auth_gate(&state, &auth_snapshot)
|
||||
.await
|
||||
.expect("wallet gate should resolve")
|
||||
.expect("wallet gate should return a decision");
|
||||
|
||||
assert!(!decision.allowed);
|
||||
assert_eq!(decision.failure, Some(WalletAccessFailure::BalanceDenied));
|
||||
assert_eq!(
|
||||
local_rejection_from_wallet_access(&decision),
|
||||
Some(GatewayLocalAuthRejection::BalanceDenied {
|
||||
remaining: Some(0.0),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ordinary_user_key_with_remaining_quota_allows_empty_wallet() {
|
||||
let state = state_with_wallet_and_quota(
|
||||
empty_user_wallet(),
|
||||
Some(quota_availability(10.0, 4.0, false)),
|
||||
);
|
||||
let auth_snapshot = ordinary_user_api_key_snapshot();
|
||||
|
||||
let decision = resolve_wallet_auth_gate(&state, &auth_snapshot)
|
||||
.await
|
||||
.expect("wallet gate should resolve")
|
||||
.expect("wallet gate should return a decision");
|
||||
|
||||
assert!(decision.allowed);
|
||||
assert_eq!(decision.failure, None);
|
||||
assert_eq!(decision.remaining, Some(4.0));
|
||||
}
|
||||
|
||||
fn state_with_wallet_and_quota(
|
||||
wallet: StoredWalletSnapshot,
|
||||
quota: Option<UserDailyQuotaAvailabilityRecord>,
|
||||
) -> AppState {
|
||||
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
|
||||
let billing_repository: Arc<dyn BillingReadRepository> =
|
||||
Arc::new(FixedQuotaBillingReadRepository { quota });
|
||||
let wallet_repository = Arc::new(InMemoryWalletRepository::seed(vec![wallet]));
|
||||
let data = GatewayDataState::with_usage_billing_and_wallet_for_tests(
|
||||
usage_repository,
|
||||
billing_repository,
|
||||
wallet_repository,
|
||||
);
|
||||
AppState::new()
|
||||
.expect("state should build")
|
||||
.with_data_state_for_tests(data)
|
||||
}
|
||||
|
||||
fn empty_user_wallet() -> StoredWalletSnapshot {
|
||||
StoredWalletSnapshot::new(
|
||||
"wallet-user-1".to_string(),
|
||||
Some("user-1".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")
|
||||
}
|
||||
|
||||
fn quota_availability(
|
||||
total_quota_usd: f64,
|
||||
remaining_usd: f64,
|
||||
allow_wallet_overage: bool,
|
||||
) -> UserDailyQuotaAvailabilityRecord {
|
||||
UserDailyQuotaAvailabilityRecord {
|
||||
has_active_daily_quota: true,
|
||||
total_quota_usd,
|
||||
used_usd: total_quota_usd - remaining_usd,
|
||||
remaining_usd,
|
||||
allow_wallet_overage,
|
||||
}
|
||||
}
|
||||
|
||||
fn ordinary_user_api_key_snapshot() -> GatewayAuthApiKeySnapshot {
|
||||
GatewayAuthApiKeySnapshot {
|
||||
user_id: "user-1".to_string(),
|
||||
username: "ordinary-user".to_string(),
|
||||
email: Some("ordinary@example.com".to_string()),
|
||||
user_role: "user".to_string(),
|
||||
user_auth_source: "local".to_string(),
|
||||
user_is_active: true,
|
||||
user_is_deleted: false,
|
||||
user_rate_limit: None,
|
||||
user_allowed_providers: None,
|
||||
user_allowed_api_formats: None,
|
||||
user_allowed_models: None,
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
api_key_name: Some("admin-created-key".to_string()),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: None,
|
||||
api_key_concurrent_limit: None,
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user