i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-09 04:30:20 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
e8d9877b79ceed831801e8984232a6101af076a4
Aether/apps/aether-gateway/src/control/auth/gate.rs
T

2174 lines
74 KiB
Rust
Raw Blame History

use axum::body::Bytes;
use axum::http::Uri;
use super::super::GatewayControlDecision;
use super::credentials::{contains_string, extract_requested_model};
use super::GatewayControlAuthContext;
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AppState, GatewayError};
const DAILY_QUOTA_EPSILON_USD: f64 = 0.000_000_01;
const AUTH_PRICING_VALIDATION_CACHE_TTL: std::time::Duration = std::time::Duration::from_secs(5);
// Billing mutations clear this cache locally. The bounded stale window leaves
// room for cross-node propagation without synchronously reloading at every
// short TTL boundary.
const AUTH_CAPACITY_CACHE_STALE_TTL: std::time::Duration = std::time::Duration::from_secs(60);
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum GatewayLocalAuthRejection {
InvalidApiKey,
LockedApiKey,
WalletUnavailable,
BalanceDenied { remaining: Option<f64> },
ProviderNotAllowed { provider: String },
ApiFormatNotAllowed { api_format: String },
ModelNotAllowed { model: String },
IpNotAllowed { remote_ip: String },
}
pub(crate) fn trusted_auth_local_rejection(
decision: Option<&GatewayControlDecision>,
_headers: &http::HeaderMap,
) -> Option<GatewayLocalAuthRejection> {
let decision = decision?;
if decision.route_class.as_deref() != Some("ai_public") {
return None;
}
decision
.local_auth_rejection
.clone()
.or_else(|| decision.auth_context.as_ref()?.local_rejection.clone())
}
pub(crate) fn should_buffer_request_for_local_auth(
decision: Option<&GatewayControlDecision>,
headers: &http::HeaderMap,
) -> bool {
let Some(decision) = decision else {
return false;
};
decision.route_class.as_deref() == Some("ai_public")
&& decision.route_kind.as_deref() != Some("files")
&& crate::headers::is_json_request(headers)
}
pub(crate) async fn request_model_local_rejection(
state: &AppState,
decision: Option<&GatewayControlDecision>,
uri: &Uri,
headers: &http::HeaderMap,
body: &Bytes,
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
let Some(decision) = decision else {
return Ok(None);
};
if decision.route_class.as_deref() != Some("ai_public") {
return Ok(None);
}
let Some(auth_context) = decision.auth_context.as_ref() else {
return Ok(None);
};
let requested_model = extract_requested_model(decision, uri, headers, body);
if let (Some(allowed_models), Some(requested_model)) = (
auth_context.allowed_models.as_deref(),
requested_model.as_deref(),
) {
if !contains_string(allowed_models, requested_model)
&& !model_directive_base_model_is_allowed_for_request(
decision,
requested_model,
allowed_models,
)
&& !request_model_resolves_to_allowed_model(
state,
decision,
requested_model,
allowed_models,
)
.await?
{
return Ok(Some(GatewayLocalAuthRejection::ModelNotAllowed {
model: requested_model.to_string(),
}));
}
}
Ok(None)
}
pub(crate) async fn execution_plan_balance_capacity_rejection(
state: &AppState,
decision: &GatewayControlDecision,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
let started_at = std::time::Instant::now();
let result =
execution_plan_balance_capacity_rejection_inner(state, decision, plan, report_context)
.await;
observe_gateway_stage_ms(
"auth_capacity_total",
started_at.elapsed().as_millis() as u64,
);
result
}
async fn execution_plan_balance_capacity_rejection_inner(
state: &AppState,
decision: &GatewayControlDecision,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<Option<GatewayLocalAuthRejection>, GatewayError> {
let Some(auth_context) = decision.auth_context.as_ref() else {
return Ok(None);
};
if auth_context.local_rejection.is_some() {
return Ok(None);
}
if auth_context.api_key_is_standalone {
validate_execution_plan_pricing_configuration_for_plan(state, plan, report_context).await?;
return Ok(None);
}
let Some(available_usd) = available_balance_capacity_usd(state, auth_context).await? else {
validate_execution_plan_pricing_configuration_for_plan(state, plan, report_context).await?;
return Ok(None);
};
match estimate_execution_plan_cost_upper_bound_usd(state, plan, report_context).await? {
Some(estimated_cost_usd)
if estimated_cost_usd <= available_usd + DAILY_QUOTA_EPSILON_USD =>
{
Ok(None)
}
Some(_) | None if available_usd <= DAILY_QUOTA_EPSILON_USD => {
Ok(Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(0.0),
}))
}
Some(_) => Ok(Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(available_usd),
})),
None => Ok(None),
}
}
async fn validate_execution_plan_pricing_configuration_for_plan(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<(), GatewayError> {
let model_id = report_context_string_field(report_context, "model_id");
let global_model_name = report_context_string_field(report_context, "global_model_name");
let requested_processing_tier =
aether_data_contracts::repository::usage::extract_provider_service_tier_from_body(
plan.body.json_body.as_ref(),
);
validate_execution_plan_pricing_for_unavailable_estimate(
state,
plan,
model_id,
global_model_name,
requested_processing_tier.as_deref(),
)
.await
}
async fn available_balance_capacity_usd(
state: &AppState,
auth_context: &GatewayControlAuthContext,
) -> Result<Option<f64>, GatewayError> {
let quota_started_at = std::time::Instant::now();
let quota_result = state
.find_user_daily_quota_availability_for_auth(&auth_context.user_id)
.await;
observe_gateway_stage_ms(
"auth_capacity_quota",
quota_started_at.elapsed().as_millis() as u64,
);
let quota = quota_result?.filter(|quota| quota.has_active_daily_quota);
let wallet_started_at = std::time::Instant::now();
let wallet_result = state
.read_wallet_snapshot_for_auth(
&auth_context.user_id,
&auth_context.api_key_id,
auth_context.api_key_is_standalone,
)
.await;
observe_gateway_stage_ms(
"auth_capacity_wallet",
wallet_started_at.elapsed().as_millis() as u64,
);
let wallet = wallet_result?;
let wallet_available_usd = wallet.as_ref().and_then(wallet_finite_available_usd);
let wallet_is_unlimited = wallet
.as_ref()
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
Ok(match quota.as_ref() {
Some(_) if wallet_is_unlimited => None,
Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)),
Some(quota) => Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)),
None if wallet_is_unlimited => None,
None => wallet_available_usd,
})
}
fn wallet_finite_available_usd(
wallet: &aether_data::repository::wallet::StoredWalletSnapshot,
) -> Option<f64> {
if !wallet.status.eq_ignore_ascii_case("active")
|| wallet.limit_mode.eq_ignore_ascii_case("unlimited")
{
return None;
}
Some(wallet.balance.max(0.0) + wallet.gift_balance.max(0.0))
}
async fn estimate_execution_plan_cost_upper_bound_usd(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<Option<f64>, GatewayError> {
let started_at = std::time::Instant::now();
let result =
estimate_execution_plan_cost_upper_bound_usd_inner(state, plan, report_context).await;
observe_gateway_stage_ms(
"auth_capacity_cost_estimate",
started_at.elapsed().as_millis() as u64,
);
result
}
async fn estimate_execution_plan_cost_upper_bound_usd_inner(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<Option<f64>, GatewayError> {
let api_format = crate::ai_serving::normalize_api_format_alias(&plan.provider_api_format);
let body_json = plan.body.json_body.as_ref();
let requested_processing_tier =
aether_data_contracts::repository::usage::extract_provider_service_tier_from_body(
body_json,
);
let model_id = report_context_string_field(report_context, "model_id");
let global_model_name = report_context_string_field(report_context, "global_model_name");
let Some(task_type) = authorization_task_type(&api_format, report_context) else {
validate_execution_plan_pricing_for_unavailable_estimate(
state,
plan,
model_id,
global_model_name,
requested_processing_tier.as_deref(),
)
.await?;
return Ok(None);
};
let Some(body_json) = body_json else {
validate_execution_plan_pricing_for_unavailable_estimate(
state,
plan,
model_id,
global_model_name,
requested_processing_tier.as_deref(),
)
.await?;
return Ok(None);
};
if !openai_request_input_is_self_contained(&api_format, body_json) {
validate_execution_plan_pricing_for_unavailable_estimate(
state,
plan,
model_id,
global_model_name,
requested_processing_tier.as_deref(),
)
.await?;
return Ok(None);
}
let input_tokens = json_token_count_upper_bound(body_json);
let Ok(input_tokens) = i64::try_from(input_tokens) else {
validate_execution_plan_pricing_for_unavailable_estimate(
state,
plan,
model_id,
global_model_name,
requested_processing_tier.as_deref(),
)
.await?;
return Ok(None);
};
let max_output_tokens = max_output_tokens_from_request(body_json)
.map(|value| value.saturating_mul(output_choice_count_upper_bound(&api_format, body_json)))
.and_then(|value| i64::try_from(value).ok());
let cache_ttl_minutes =
aether_data_contracts::repository::usage::resolve_provider_cache_ttl_minutes(
Some(&api_format),
plan.model_name.as_deref(),
global_model_name,
Some(body_json),
);
if model_id.is_none() && global_model_name.is_none() {
return Ok(None);
}
let cache_key = execution_plan_cost_upper_bound_cache_key(
plan,
model_id,
global_model_name,
&api_format,
input_tokens,
max_output_tokens,
requested_processing_tier.as_deref(),
cache_ttl_minutes,
);
let ttl = state.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if ttl.is_zero() {
let _permit = state.acquire_auth_snapshot_load_gate().await?;
return calculate_execution_plan_cost_upper_bound(
state,
plan,
model_id,
global_model_name,
&api_format,
task_type,
input_tokens,
max_output_tokens,
requested_processing_tier.as_deref(),
cache_ttl_minutes,
)
.await;
}
state
.auth_request_cost_upper_bound_cache
.get_or_load(cache_key, ttl, || async {
let _permit = state.acquire_auth_snapshot_load_gate().await?;
calculate_execution_plan_cost_upper_bound(
state,
plan,
model_id,
global_model_name,
&api_format,
task_type,
input_tokens,
max_output_tokens,
requested_processing_tier.as_deref(),
cache_ttl_minutes,
)
.await
})
.await
}
#[allow(clippy::too_many_arguments)]
async fn calculate_execution_plan_cost_upper_bound(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
api_format: &str,
task_type: &str,
input_tokens: i64,
max_output_tokens: Option<i64>,
requested_processing_tier: Option<&str>,
cache_ttl_minutes: Option<i64>,
) -> Result<Option<f64>, GatewayError> {
let context =
load_execution_plan_billing_context(state, plan, model_id, global_model_name).await?;
let Some(context) = context else {
return Ok(None);
};
let mut estimate =
aether_billing::BillingAuthorizationEstimateInput::new(task_type, input_tokens);
estimate.api_format = Some(api_format.to_string());
estimate.requested_processing_tier = requested_processing_tier.map(ToOwned::to_owned);
estimate.cache_ttl_minutes = cache_ttl_minutes;
estimate.max_output_tokens = max_output_tokens;
aether_billing::BillingService::new()
.estimate_authorization_cost_upper_bound(
&aether_billing::BillingModelPricingSnapshot::from(context),
&estimate,
)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
async fn validate_execution_plan_pricing_for_unavailable_estimate(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
requested_processing_tier: Option<&str>,
) -> Result<(), GatewayError> {
let started_at = std::time::Instant::now();
let result = validate_execution_plan_pricing_for_unavailable_estimate_inner(
state,
plan,
model_id,
global_model_name,
requested_processing_tier,
)
.await;
observe_gateway_stage_ms(
"auth_capacity_pricing_validation",
started_at.elapsed().as_millis() as u64,
);
result
}
async fn validate_execution_plan_pricing_for_unavailable_estimate_inner(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
requested_processing_tier: Option<&str>,
) -> Result<(), GatewayError> {
if model_id.is_none() && global_model_name.is_none() {
return Ok(());
}
let capacity_ttl = state.frontdoor_runtime_guards.auth_capacity_cache_ttl;
if capacity_ttl.is_zero() {
return validate_execution_plan_pricing_uncached(
state,
plan,
model_id,
global_model_name,
requested_processing_tier,
)
.await;
}
let ttl = capacity_ttl.max(AUTH_PRICING_VALIDATION_CACHE_TTL);
let cache_key = execution_plan_pricing_validation_cache_key(
plan,
model_id,
global_model_name,
requested_processing_tier,
);
let cache = state.auth_request_cost_upper_bound_cache.clone();
cache
.get_or_load_once_stale_while_revalidating(
cache_key,
ttl,
AUTH_CAPACITY_CACHE_STALE_TTL,
|| async {
validate_execution_plan_pricing_uncached(
state,
plan,
model_id,
global_model_name,
requested_processing_tier,
)
.await?;
Ok::<Option<f64>, GatewayError>(Some(0.0))
},
|| {
let state = state.clone();
let plan = plan.clone();
let model_id = model_id.map(ToOwned::to_owned);
let global_model_name = global_model_name.map(ToOwned::to_owned);
let requested_processing_tier = requested_processing_tier.map(ToOwned::to_owned);
async move {
validate_execution_plan_pricing_uncached(
&state,
&plan,
model_id.as_deref(),
global_model_name.as_deref(),
requested_processing_tier.as_deref(),
)
.await?;
Ok::<Option<f64>, GatewayError>(Some(0.0))
}
},
crate::cache::CacheLoadObserver::default(),
)
.await?;
Ok(())
}
async fn validate_execution_plan_pricing_uncached(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
requested_processing_tier: Option<&str>,
) -> Result<(), GatewayError> {
let _permit = state.acquire_auth_snapshot_load_gate().await?;
let Some(context) =
load_execution_plan_billing_context(state, plan, model_id, global_model_name).await?
else {
return Ok(());
};
aether_billing::BillingModelPricingSnapshot::from(context)
.validate_authorization_pricing_configuration(requested_processing_tier)
.map_err(|err| GatewayError::Internal(err.to_string()))
}
async fn load_execution_plan_billing_context(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
) -> Result<
Option<aether_data_contracts::repository::billing::StoredBillingModelContext>,
GatewayError,
> {
let context = match model_id {
Some(model_id) => state
.data
.find_billing_model_context_by_model_id(&plan.provider_id, Some(&plan.key_id), model_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?,
None => state
.data
.find_billing_model_context(
&plan.provider_id,
Some(&plan.key_id),
global_model_name.expect("global model name should exist"),
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?,
};
Ok(context)
}
fn execution_plan_cost_upper_bound_cache_key(
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
api_format: &str,
input_tokens: i64,
max_output_tokens: Option<i64>,
requested_processing_tier: Option<&str>,
cache_ttl_minutes: Option<i64>,
) -> String {
format!(
"{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}",
plan.provider_id,
plan.key_id,
model_id.unwrap_or(""),
global_model_name.unwrap_or(""),
api_format,
input_tokens,
max_output_tokens
.map(|value| value.to_string())
.unwrap_or_else(|| "none".to_string()),
requested_processing_tier.unwrap_or("standard"),
cache_ttl_minutes
.map(|value| value.to_string())
.unwrap_or_else(|| "none".to_string()),
)
}
fn execution_plan_pricing_validation_cache_key(
plan: &aether_contracts::ExecutionPlan,
model_id: Option<&str>,
global_model_name: Option<&str>,
requested_processing_tier: Option<&str>,
) -> String {
format!(
"pricing-validation\x1f{}\x1f{}\x1f{}\x1f{}\x1f{}",
plan.provider_id,
plan.key_id,
model_id.unwrap_or(""),
global_model_name.unwrap_or(""),
requested_processing_tier.unwrap_or("standard"),
)
}
fn authorization_task_type<'a>(
api_format: &str,
report_context: Option<&'a serde_json::Value>,
) -> Option<&'a str> {
if report_context
.and_then(|context| context.get("image_request"))
.is_some()
|| api_format == "openai:image"
{
return None;
}
if api_format.ends_with(":embedding") {
return Some("embedding");
}
if api_format.ends_with(":rerank") {
return Some("rerank");
}
Some("chat")
}
fn report_context_string_field<'a>(
report_context: Option<&'a serde_json::Value>,
key: &str,
) -> Option<&'a str> {
report_context
.and_then(|context| context.get(key))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn max_output_tokens_from_request(value: &serde_json::Value) -> Option<u64> {
["max_tokens", "max_completion_tokens", "max_output_tokens"]
.iter()
.filter_map(|field| value.get(*field).and_then(serde_json::Value::as_u64))
.filter(|value| *value > 0)
.max()
}
fn output_choice_count_upper_bound(api_format: &str, value: &serde_json::Value) -> u64 {
if api_format != "openai:chat" {
return 1;
}
value
.get("n")
.and_then(serde_json::Value::as_u64)
.filter(|value| *value > 0)
.unwrap_or(1)
}
fn openai_request_input_is_self_contained(api_format: &str, value: &serde_json::Value) -> bool {
if !api_format.starts_with("openai:") {
return false;
}
let Some(object) = value.as_object() else {
return true;
};
if ["previous_response_id", "conversation"]
.iter()
.any(|key| object.get(*key).is_some_and(has_reference_value))
{
return false;
}
if object
.get("prompt")
.and_then(serde_json::Value::as_object)
.and_then(|prompt| prompt.get("id"))
.is_some_and(has_reference_value)
{
return false;
}
!contains_indirect_request_input(value)
}
fn contains_indirect_request_input(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::Array(items) => items.iter().any(contains_indirect_request_input),
serde_json::Value::Object(object) => {
let item_type = object
.get("type")
.and_then(serde_json::Value::as_str)
.map(|value| value.trim().to_ascii_lowercase());
if item_type.as_deref().is_some_and(|item_type| {
matches!(
item_type,
"url"
| "item_reference"
| "input_file"
| "input_image"
| "input_audio"
| "image_url"
| "file_search"
| "web_search"
| "web_search_preview"
| "computer_use"
| "computer_use_preview"
| "code_interpreter"
| "mcp"
| "image_generation"
)
}) {
return true;
}
if ["file_id", "file_uri", "fileUri", "vector_store_ids"]
.iter()
.any(|key| object.get(*key).is_some_and(has_reference_value))
{
return true;
}
object.values().any(contains_indirect_request_input)
}
_ => false,
}
}
fn has_reference_value(value: &serde_json::Value) -> bool {
match value {
serde_json::Value::Null => false,
serde_json::Value::String(value) => !value.trim().is_empty(),
serde_json::Value::Array(values) => !values.is_empty(),
serde_json::Value::Object(values) => !values.is_empty(),
_ => true,
}
}
fn json_token_count_upper_bound(value: &serde_json::Value) -> u64 {
serde_json::to_vec(value)
.map(|bytes| u64::try_from(bytes.len()).unwrap_or(u64::MAX))
.unwrap_or(u64::MAX)
}
fn model_directive_base_model_is_allowed_for_request(
decision: &GatewayControlDecision,
requested_model: &str,
allowed_models: &[String],
) -> bool {
let Some(client_api_format) = decision
.auth_endpoint_signature
.as_deref()
.map(crate::ai_serving::normalize_api_format_alias)
.filter(|value| !value.trim().is_empty())
else {
return false;
};
for api_format in candidate_api_formats_for_model_resolution(&client_api_format) {
let resolution = decision
.model_directive_policy
.resolve_reasoning(&api_format, Some(requested_model));
if resolution
.base_model()
.is_some_and(|base_model| contains_string(allowed_models, base_model))
{
return true;
}
}
false
}
async fn request_model_resolves_to_allowed_model(
state: &AppState,
decision: &GatewayControlDecision,
requested_model: &str,
allowed_models: &[String],
) -> Result<bool, GatewayError> {
let Some(client_api_format) = decision
.auth_endpoint_signature
.as_deref()
.map(crate::ai_serving::normalize_api_format_alias)
.filter(|value| !value.trim().is_empty())
else {
return Ok(false);
};
for api_format in candidate_api_formats_for_model_resolution(&client_api_format) {
let resolution = decision
.model_directive_policy
.resolve_reasoning(&api_format, Some(requested_model));
let routing_model = resolution.base_model().unwrap_or(requested_model);
let rows = {
// Model alias authorization runs before candidate planning, so its database read must
// participate in the same foreground DB admission budget as the rest of auth. Keep
// the permit scoped to this one read; callers do not hold this gate, and releasing it
// here avoids carrying a DB permit through pure filtering or subsequent formats.
let _permit = state.acquire_auth_snapshot_load_gate().await?;
state
.list_minimal_candidate_selection_rows_for_api_format(&api_format)
.await?
};
let matching_rows = rows
.into_iter()
.filter(|row| {
aether_scheduler_core::row_supports_requested_model_with_model_directives(
row,
routing_model,
&api_format,
false,
)
})
.collect::<Vec<_>>();
let Some(resolved_global_model) =
aether_scheduler_core::resolve_requested_global_model_name_with_model_directives(
&matching_rows,
routing_model,
&api_format,
false,
)
else {
continue;
};
if contains_string(allowed_models, &resolved_global_model) {
return Ok(true);
}
}
Ok(false)
}
fn candidate_api_formats_for_model_resolution(client_api_format: &str) -> Vec<String> {
let mut api_formats = Vec::new();
push_unique_api_format(&mut api_formats, client_api_format);
for api_format in crate::ai_serving::request_candidate_api_formats(client_api_format, false) {
push_unique_api_format(&mut api_formats, api_format);
}
api_formats
}
fn push_unique_api_format(api_formats: &mut Vec<String>, api_format: &str) {
let api_format = crate::ai_serving::normalize_api_format_alias(api_format);
if api_format.is_empty() || api_formats.iter().any(|value| value == &api_format) {
return;
}
api_formats.push(api_format);
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::wallet::StoredWalletSnapshot;
use aether_data_contracts::repository::billing::{
BillingReadRepository, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::DataLayerError;
use aether_runtime::ConcurrencyGate;
use async_trait::async_trait;
use axum::body::Bytes;
use axum::http::{HeaderMap, Uri};
use serde_json::json;
use super::{
available_balance_capacity_usd, execution_plan_balance_capacity_rejection,
execution_plan_cost_upper_bound_cache_key, max_output_tokens_from_request,
openai_request_input_is_self_contained, output_choice_count_upper_bound,
request_model_local_rejection, GatewayLocalAuthRejection,
};
use crate::control::{GatewayControlAuthContext, GatewayControlDecision};
use crate::data::GatewayDataState;
use crate::AppState;
fn sample_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-1".to_string(),
provider_name: "Provider 1".to_string(),
provider_type: "openai".to_string(),
provider_priority: 0,
provider_is_active: true,
endpoint_id: "endpoint-1".to_string(),
endpoint_api_format: "openai:chat".to_string(),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
endpoint_is_active: true,
key_id: "key-1".to_string(),
key_name: "key".to_string(),
key_auth_type: "api_key".to_string(),
key_is_active: true,
key_api_formats: Some(vec!["openai:chat".to_string()]),
key_allowed_models: None,
key_capabilities: None,
key_internal_priority: 0,
key_global_priority_by_format: None,
model_id: "model-1".to_string(),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_mappings: Some(vec!["gpt-5(?:\\.\\d+)?".to_string()]),
global_model_supports_streaming: Some(true),
model_provider_model_name: "gpt-5-upstream".to_string(),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: "gpt-5-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
endpoint_ids: None,
operations: None,
}]),
model_supports_streaming: Some(true),
model_is_active: true,
model_is_available: true,
}
}
fn sample_row_for_api_format(api_format: &str) -> StoredMinimalCandidateSelectionRow {
let mut row = sample_row();
let api_family = api_format
.split_once(':')
.map(|(family, _)| family)
.unwrap_or(api_format);
row.provider_id = format!("provider-{api_family}");
row.provider_name = format!("Provider {api_family}");
row.provider_type = api_family.to_string();
row.endpoint_id = format!("endpoint-{api_family}");
row.endpoint_api_format = api_format.to_string();
row.endpoint_api_family = Some(api_family.to_string());
row.key_id = format!("key-{api_family}");
row.key_api_formats = Some(vec![api_format.to_string()]);
if let Some(mappings) = row.model_provider_model_mappings.as_mut() {
for mapping in mappings {
mapping.api_formats = Some(vec![api_format.to_string()]);
}
}
row
}
fn decision_with_allowed_models(allowed_models: Vec<String>) -> GatewayControlDecision {
let mut decision = GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
);
decision.auth_context = Some(GatewayControlAuthContext {
user_id: "user-1".to_string(),
api_key_id: "api-key-1".to_string(),
username: None,
api_key_name: None,
balance_remaining: None,
access_allowed: true,
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: Some(allowed_models),
ip_rules: None,
});
decision
}
fn state_with_rows(rows: Vec<StoredMinimalCandidateSelectionRow>) -> AppState {
let repository = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows));
let data = GatewayDataState::with_minimal_candidate_selection_reader_for_tests(repository);
AppState::new()
.expect("state should build")
.with_data_state_for_tests(data)
}
fn state_with_quota_and_wallet(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
) -> AppState {
state_with_quota_context_and_wallet(quota, context, sample_wallet("user-1", 30.0))
}
fn state_with_quota_context_and_wallet(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
wallet: StoredWalletSnapshot,
) -> AppState {
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_row(),
]));
let billing_repository = Arc::new(FixedBillingReadRepository::new(quota, context));
let data = GatewayDataState::with_minimal_candidate_selection_and_billing_for_tests(
candidate_repository,
billing_repository,
);
AppState::new()
.expect("state should build")
.with_data_state_for_tests(data)
.with_auth_wallets_for_tests(vec![wallet])
}
fn state_with_model_mapping() -> AppState {
state_with_rows(vec![sample_row()])
}
fn execution_plan(body: serde_json::Value, api_format: &str) -> ExecutionPlan {
ExecutionPlan {
request_id: "request-1".to_string(),
candidate_id: Some("candidate-1".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://api.openai.com/v1/responses".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(body),
stream: false,
client_api_format: api_format.to_string(),
provider_api_format: api_format.to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
}
}
fn billing_report_context() -> serde_json::Value {
json!({
"model_id": "model-1",
"global_model_name": "gpt-5"
})
}
fn estimate_from_billing_context(
context: &StoredBillingModelContext,
api_format: &str,
input_tokens: i64,
max_output_tokens: Option<i64>,
) -> Option<f64> {
let mut estimate =
aether_billing::BillingAuthorizationEstimateInput::new("chat", input_tokens);
estimate.api_format = Some(api_format.to_string());
estimate.max_output_tokens = max_output_tokens;
aether_billing::BillingService::new()
.estimate_authorization_cost_upper_bound(
&aether_billing::BillingModelPricingSnapshot::from(context),
&estimate,
)
.expect("estimate should calculate")
}
fn json_headers() -> HeaderMap {
let mut headers = HeaderMap::new();
headers.insert(
axum::http::header::CONTENT_TYPE,
"application/json"
.parse()
.expect("content type should parse"),
);
headers
}
fn billing_context_with_pricing(
default_tiered_pricing: Option<serde_json::Value>,
model_tiered_pricing: Option<serde_json::Value>,
rate_multipliers: Option<serde_json::Value>,
billing_type: Option<&str>,
) -> StoredBillingModelContext {
StoredBillingModelContext::new(
"provider-1".to_string(),
billing_type.map(ToOwned::to_owned),
Some("key-1".to_string()),
rate_multipliers,
Some(60),
"global-model-1".to_string(),
"gpt-5".to_string(),
None,
None,
default_tiered_pricing,
Some("model-1".to_string()),
Some("gpt-5-upstream".to_string()),
None,
None,
model_tiered_pricing,
)
.expect("billing context should build")
}
fn sample_wallet(user_id: &str, balance: f64) -> StoredWalletSnapshot {
StoredWalletSnapshot::new(
format!("wallet-{user_id}"),
Some(user_id.to_string()),
None,
balance,
0.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
balance,
0.0,
0.0,
0.0,
100,
)
.expect("wallet should build")
}
fn quota_availability(
remaining_usd: f64,
allow_wallet_overage: bool,
) -> UserDailyQuotaAvailabilityRecord {
UserDailyQuotaAvailabilityRecord {
has_active_daily_quota: true,
total_quota_usd: remaining_usd,
used_usd: 0.0,
remaining_usd,
allow_wallet_overage,
}
}
#[derive(Debug)]
struct FixedBillingReadRepository {
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
quota_calls: Arc<AtomicUsize>,
model_context_by_model_id_calls: Arc<AtomicUsize>,
}
impl FixedBillingReadRepository {
fn new(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
) -> Self {
Self {
quota,
context,
quota_calls: Arc::new(AtomicUsize::new(0)),
model_context_by_model_id_calls: Arc::new(AtomicUsize::new(0)),
}
}
fn with_counters(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
quota_calls: Arc<AtomicUsize>,
model_context_by_model_id_calls: Arc<AtomicUsize>,
) -> Self {
Self {
quota,
context,
quota_calls,
model_context_by_model_id_calls,
}
}
}
#[async_trait]
impl BillingReadRepository for FixedBillingReadRepository {
async fn find_model_context(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
Ok(Some(self.context.clone()))
}
async fn find_model_context_by_model_id(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
self.model_context_by_model_id_calls
.fetch_add(1, Ordering::AcqRel);
Ok(Some(self.context.clone()))
}
async fn find_user_daily_quota_availability(
&self,
_user_id: &str,
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, DataLayerError> {
self.quota_calls.fetch_add(1, Ordering::AcqRel);
Ok(Some(self.quota.clone()))
}
}
#[tokio::test]
async fn model_rejection_allows_requested_model_that_resolves_to_allowed_global_model() {
let state = state_with_model_mapping();
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(br#"{"model":"gpt-5.2","messages":[]}"#);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_alias_resolution_waits_for_auth_database_gate() {
let mut state = state_with_model_mapping();
state.auth_snapshot_load_gate = Some(Arc::new(ConcurrencyGate::new(
"test_auth_model_resolution",
1,
)));
let held = state
.acquire_auth_snapshot_load_gate()
.await
.expect("auth gate acquisition should succeed")
.expect("auth gate should be configured");
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let headers = json_headers();
let body = Bytes::from_static(br#"{"model":"gpt-5.2","messages":[]}"#);
let blocked = tokio::time::timeout(
Duration::from_millis(25),
request_model_local_rejection(&state, Some(&decision), &uri, &headers, &body),
)
.await;
assert!(
blocked.is_err(),
"model alias candidate reads must wait for the auth DB gate"
);
drop(held);
let rejection = tokio::time::timeout(
Duration::from_secs(1),
request_model_local_rejection(&state, Some(&decision), &uri, &headers, &body),
)
.await
.expect("model alias resolution should resume after releasing the auth gate")
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_rejection_allows_cross_format_provider_mapping_to_allowed_global_model() {
let mut row = sample_row_for_api_format("gemini:generate_content");
row.model_provider_model_name = "gemini-2.5-pro-upstream".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gemini-2.5-pro-alias".to_string(),
priority: 1,
api_formats: Some(vec!["gemini:generate_content".to_string()]),
endpoint_ids: None,
operations: None,
}]);
let state = state_with_rows(vec![row]);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(br#"{"model":"gemini-2.5-pro-alias","messages":[]}"#);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_rejection_allows_cross_format_regex_mapping_to_allowed_global_model() {
let state = state_with_rows(vec![sample_row_for_api_format("claude:messages")]);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(br#"{"model":"gpt-5.2","messages":[]}"#);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_rejection_denies_requested_model_outside_allowed_global_models() {
let state = state_with_model_mapping();
let decision = decision_with_allowed_models(vec!["gpt-4.1".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(br#"{"model":"gpt-5.2","messages":[]}"#);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("model rejection should resolve");
assert_eq!(
rejection,
Some(GatewayLocalAuthRejection::ModelNotAllowed {
model: "gpt-5.2".to_string(),
})
);
}
#[tokio::test]
async fn model_rejection_reuses_request_policy_snapshot_for_directive_base_model() {
let state = state_with_rows(Vec::new());
let mut decision = decision_with_allowed_models(vec!["gpt-5.6-sol".to_string()]);
decision.model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::from_config_values(
Some(&json!(true)),
None,
);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(br#"{"model":"gpt-5.6-sol-high","messages":[]}"#);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn model_rejection_uses_custom_policy_suffix_for_base_model_authorization() {
let state = state_with_rows(Vec::new());
let mut decision = decision_with_allowed_models(vec!["deployment-alias".to_string()]);
decision.model_directive_policy =
crate::system_features::ModelDirectivePolicySnapshot::from_config_values(
Some(&json!(true)),
Some(&json!({
"reasoning_effort": {
"api_formats": {
"openai:chat": {
"suffixes": ["VendorFuture"],
"mappings": {
"VendorFuture": { "reasoning_effort": "high" }
}
}
}
}
})),
);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body =
Bytes::from_static(br#"{"model":"deployment-alias-VendorFuture","messages":[]}"#);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("model rejection should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn positive_balance_allows_unbounded_output_request_without_cost_estimate() {
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,
);
for allow_wallet_overage in [false, true] {
let state = state_with_quota_and_wallet(
quota_availability(50.0, allow_wallet_overage),
context.clone(),
);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"stream": true
}),
"openai:chat",
);
let rejection = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect("quota rejection should resolve");
assert_eq!(rejection, None);
}
}
#[tokio::test]
async fn unlimited_wallet_capacity_ignores_exhausted_non_overage_quota() {
let context = billing_context_with_pricing(None, None, None, None);
let mut wallet = sample_wallet("user-1", 0.0);
wallet.limit_mode = "unlimited".to_string();
let state =
state_with_quota_context_and_wallet(quota_availability(0.0, false), context, wallet);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let auth_context = decision
.auth_context
.as_ref()
.expect("decision should include auth context");
let capacity = available_balance_capacity_usd(&state, auth_context)
.await
.expect("capacity should resolve");
assert_eq!(capacity, None);
}
#[tokio::test]
async fn positive_balance_does_not_allow_historical_invalid_processing_pricing() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}],
"processing_tiers": {
"priority": {
"tiers": [{}],
"price_multiplier": -1.0
}
}
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(50.0, true), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"service_tier": "priority",
"max_completion_tokens": 1
}),
"openai:chat",
);
let error = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect_err("invalid configured processing pricing must stop authorization");
assert!(error
.into_message()
.contains("explicit catalog contains malformed or unrecognized prices"));
}
#[tokio::test]
async fn standalone_key_does_not_bypass_invalid_processing_pricing_validation() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{"up_to": null, "input_price_per_1m": 1.0}],
"processing_tiers": {
"priority": {"tiers": [{}], "price_multiplier": 2.0}
}
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(50.0, true), context);
let mut decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
decision
.auth_context
.as_mut()
.expect("auth context should exist")
.api_key_is_standalone = true;
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"service_tier": "priority",
"max_completion_tokens": 1
}),
"openai:chat",
);
let error = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect_err("standalone keys must still validate configured pricing");
assert!(error
.into_message()
.contains("explicit catalog contains malformed or unrecognized prices"));
}
#[tokio::test]
async fn unlimited_wallet_does_not_bypass_invalid_processing_pricing_validation() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{"up_to": null, "input_price_per_1m": 1.0}],
"processing_tiers": {
"priority": {"tiers": [{}], "price_multiplier": 2.0}
}
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(0.0, true), context);
{
let store = state
.auth_wallet_store
.as_ref()
.expect("test wallet store should exist");
let mut wallets = store.lock().expect("wallet store should lock");
wallets
.get_mut("wallet-user-1")
.expect("test wallet should exist")
.limit_mode = "unlimited".to_string();
}
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"service_tier": "priority",
"max_completion_tokens": 1
}),
"openai:chat",
);
let error = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect_err("unlimited wallets must still validate configured pricing");
assert!(error
.into_message()
.contains("explicit catalog contains malformed or unrecognized prices"));
}
#[tokio::test]
async fn standalone_and_unlimited_paths_keep_allowing_valid_pricing() {
let valid_context = billing_context_with_pricing(
Some(json!({
"tiers": [{"up_to": null, "input_price_per_1m": 1.0}]
})),
None,
None,
None,
);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"max_completion_tokens": 1
}),
"openai:chat",
);
let report_context = billing_report_context();
let standalone_state =
state_with_quota_and_wallet(quota_availability(50.0, true), valid_context.clone());
let mut standalone_decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
standalone_decision
.auth_context
.as_mut()
.expect("auth context should exist")
.api_key_is_standalone = true;
assert_eq!(
execution_plan_balance_capacity_rejection(
&standalone_state,
&standalone_decision,
&plan,
Some(&report_context),
)
.await
.expect("valid standalone pricing should resolve"),
None
);
let unlimited_state =
state_with_quota_and_wallet(quota_availability(0.0, true), valid_context);
{
let store = unlimited_state
.auth_wallet_store
.as_ref()
.expect("test wallet store should exist");
let mut wallets = store.lock().expect("wallet store should lock");
wallets
.get_mut("wallet-user-1")
.expect("test wallet should exist")
.limit_mode = "unlimited".to_string();
}
assert_eq!(
execution_plan_balance_capacity_rejection(
&unlimited_state,
&decision_with_allowed_models(vec!["gpt-5".to_string()]),
&plan,
Some(&report_context),
)
.await
.expect("valid unlimited-wallet pricing should resolve"),
None
);
}
#[tokio::test]
async fn auth_capacity_reuses_quota_wallet_and_cost_estimate_within_ttl() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 0.0,
"output_price_per_1m": 60.0
}]
})),
None,
None,
None,
);
let quota_calls = Arc::new(AtomicUsize::new(0));
let model_context_calls = Arc::new(AtomicUsize::new(0));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_row(),
]));
let billing_repository = Arc::new(FixedBillingReadRepository::with_counters(
quota_availability(1.0, true),
context,
Arc::clone(&quota_calls),
Arc::clone(&model_context_calls),
));
let data = GatewayDataState::with_minimal_candidate_selection_and_billing_for_tests(
candidate_repository,
billing_repository,
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data)
.with_auth_wallets_for_tests(vec![sample_wallet("user-1", 30.0)]);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 100_000
}),
"openai:chat",
);
let report_context = billing_report_context();
let first = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("first auth capacity check should resolve");
assert_eq!(first, None);
assert_eq!(quota_calls.load(Ordering::Acquire), 1);
assert_eq!(model_context_calls.load(Ordering::Acquire), 1);
let store = state
.auth_wallet_store
.as_ref()
.expect("test wallet store should exist");
{
let mut wallets = store.lock().expect("wallet store should lock");
let wallet = wallets
.get_mut("wallet-user-1")
.expect("test wallet should exist");
wallet.balance = 0.0;
wallet.gift_balance = 0.0;
}
let cached = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("cached auth capacity check should resolve");
assert_eq!(cached, None);
assert_eq!(quota_calls.load(Ordering::Acquire), 1);
assert_eq!(model_context_calls.load(Ordering::Acquire), 1);
state.invalidate_auth_context_cache();
let refreshed = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("refreshed auth capacity check should resolve");
assert_eq!(
refreshed,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(1.0),
})
);
assert_eq!(quota_calls.load(Ordering::Acquire), 2);
}
#[tokio::test]
async fn standalone_auth_capacity_reuses_pricing_validation_within_ttl() {
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 quota_calls = Arc::new(AtomicUsize::new(0));
let model_context_calls = Arc::new(AtomicUsize::new(0));
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_row(),
]));
let billing_repository = Arc::new(FixedBillingReadRepository::with_counters(
quota_availability(1.0, true),
context,
Arc::clone(&quota_calls),
Arc::clone(&model_context_calls),
));
let data = GatewayDataState::with_minimal_candidate_selection_and_billing_for_tests(
candidate_repository,
billing_repository,
);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data);
let mut decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
decision
.auth_context
.as_mut()
.expect("auth context should exist")
.api_key_is_standalone = true;
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}]
}),
"openai:chat",
);
let report_context = billing_report_context();
for _ in 0..2 {
assert_eq!(
execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("standalone pricing validation should resolve"),
None
);
}
assert_eq!(quota_calls.load(Ordering::Acquire), 0);
assert_eq!(model_context_calls.load(Ordering::Acquire), 1);
let cache_key = super::execution_plan_pricing_validation_cache_key(
&plan,
Some("model-1"),
Some("gpt-5"),
None,
);
assert_eq!(
state.auth_request_cost_upper_bound_cache.get(
&cache_key,
state.frontdoor_runtime_guards.auth_capacity_cache_ttl,
),
Some(Some(0.0))
);
state.invalidate_provider_routing_caches();
assert_eq!(
state.auth_request_cost_upper_bound_cache.get(
&cache_key,
state.frontdoor_runtime_guards.auth_capacity_cache_ttl,
),
None
);
}
#[tokio::test]
async fn admin_bypass_limits_does_not_skip_unbounded_zero_balance_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 plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"stream": true
}),
"openai:chat",
);
let report_context = billing_report_context();
let rejection = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("quota rejection should resolve");
assert_eq!(
rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(0.0),
})
);
}
#[tokio::test]
async fn zero_balance_allows_a_proven_free_tier_execution_plan() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 100.0,
"output_price_per_1m": 100.0
}]
})),
None,
None,
Some("free_tier"),
);
let state = state_with_quota_and_wallet(quota_availability(0.0, false), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"max_completion_tokens": 1_000_000
}),
"openai:chat",
);
let report_context = billing_report_context();
let rejection = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("free tier capacity check should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn finalized_chat_output_fields_and_choice_count_bound_capacity() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 0.0,
"output_price_per_1m": 20.0
}]
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(50.0, false), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 1,
"max_completion_tokens": 1_000_000,
"n": 3
}),
"openai:chat",
);
let rejection = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect("quota rejection should resolve");
assert_eq!(
rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(50.0),
})
);
}
#[tokio::test]
async fn stateful_responses_request_skips_unprovable_capacity_rejection() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 100.0,
"output_price_per_1m": 100.0
}]
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(0.01, false), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"input": "continue",
"previous_response_id": "resp_123",
"max_output_tokens": 1_000_000
}),
"openai:responses",
);
let report_context = billing_report_context();
let rejection = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&report_context),
)
.await
.expect("stateful request capacity check should resolve");
assert_eq!(rejection, None);
}
#[tokio::test]
async fn stateful_unavailable_estimate_still_rejects_invalid_processing_pricing() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}],
"processing_tiers": {
"priority": {
"tiers": [{}],
"price_multiplier": 2.0
}
}
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(1.0, false), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"input": "continue",
"previous_response_id": "resp_123",
"service_tier": "priority",
"max_output_tokens": 1
}),
"openai:responses",
);
let error = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect_err("unavailable estimates must still validate configured processing pricing");
assert!(error
.into_message()
.contains("explicit catalog contains malformed or unrecognized prices"));
}
#[tokio::test]
async fn wallet_overage_policy_extends_known_cost_capacity_when_enabled() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 0.0,
"output_price_per_1m": 70.0
}]
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(50.0, true), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let plan = execution_plan(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": "hi"}],
"max_tokens": 1_000_000
}),
"openai:chat",
);
let rejection = execution_plan_balance_capacity_rejection(
&state,
&decision,
&plan,
Some(&billing_report_context()),
)
.await
.expect("quota rejection should resolve");
assert_eq!(rejection, None);
}
#[test]
fn daily_quota_estimate_falls_back_to_default_tiers_when_model_tiers_empty() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 3.0,
"output_price_per_1m": 15.0
}]
})),
Some(json!({})),
None,
None,
);
let estimate =
estimate_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000))
.expect("estimate should be bounded");
assert_eq!(estimate, 18.75);
}
#[test]
fn daily_quota_estimate_applies_provider_key_rate_multiplier() {
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,
Some(json!({ "openai:chat": 2.0 })),
None,
);
let estimate =
estimate_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000))
.expect("estimate should be bounded");
assert_eq!(estimate, 6.5);
}
#[test]
fn daily_quota_estimate_treats_free_tier_as_zero_cost() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 3.0,
"output_price_per_1m": 15.0
}]
})),
None,
Some(json!({ "openai:chat": 10.0 })),
Some("free_tier"),
);
let estimate =
estimate_from_billing_context(&context, "openai:chat", 1_000_000, Some(1_000_000))
.expect("estimate should be bounded");
assert_eq!(estimate, 0.0);
}
#[test]
fn output_bound_uses_largest_supported_field_and_chat_choice_count() {
let body = json!({
"max_tokens": 1,
"max_completion_tokens": 100_000,
"max_output_tokens": 50_000,
"n": 3
});
assert_eq!(max_output_tokens_from_request(&body), Some(100_000));
assert_eq!(output_choice_count_upper_bound("openai:chat", &body), 3);
assert_eq!(
output_choice_count_upper_bound("openai:responses", &body),
1
);
}
#[test]
fn authorization_cache_key_includes_effective_cache_ttl() {
let plan = execution_plan(json!({"model": "gpt-5.6-sol"}), "openai:responses");
let without_ttl = execution_plan_cost_upper_bound_cache_key(
&plan,
Some("model-1"),
Some("gpt-5.6-sol"),
"openai:responses",
100,
Some(10),
Some("priority"),
None,
);
let with_ttl = execution_plan_cost_upper_bound_cache_key(
&plan,
Some("model-1"),
Some("gpt-5.6-sol"),
"openai:responses",
100,
Some(10),
Some("priority"),
Some(30),
);
assert_ne!(without_ttl, with_ttl);
}
#[test]
fn indirect_request_inputs_are_not_treated_as_body_bounded() {
let self_contained = json!({
"input": [{
"role": "user",
"content": [{"type": "input_text", "text": "hello"}]
}],
"tools": [{"type": "function", "name": "lookup", "parameters": {}}]
});
assert!(openai_request_input_is_self_contained(
"openai:responses",
&self_contained
));
for indirect in [
json!({"input": "continue", "previous_response_id": "resp_123"}),
json!({"input": "hello", "conversation": "conv_123"}),
json!({"prompt": {"id": "pmpt_123", "variables": {}}}),
json!({"input": [{"type": "item_reference", "id": "item_123"}]}),
json!({"input": [{"type": "input_file", "file_id": "file_123"}]}),
json!({"input": [{"type": "input_image", "image_url": "https://example.test/a.png"}]}),
json!({"input": [{"type": "url", "url": "https://example.test/document"}]}),
json!({"input": [{"file_uri": "https://example.test/file"}]}),
json!({"input": "search", "tools": [{"type": "file_search", "vector_store_ids": ["vs_123"]}]}),
] {
assert!(!openai_request_input_is_self_contained(
"openai:responses",
&indirect
));
}
assert!(!openai_request_input_is_self_contained(
"claude:messages",
&self_contained
));
}
}
Reference in New Issue View Git Blame Copy Permalink