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 }, 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 { 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, 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, 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, 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, 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 { 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, 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, 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, requested_processing_tier: Option<&str>, cache_ttl_minutes: Option, ) -> Result, 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::, 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::, 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, 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, requested_processing_tier: Option<&str>, cache_ttl_minutes: Option, ) -> 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 { ["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 { 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::>(); 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 { 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, 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) -> 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) -> 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, ) -> Option { 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, model_tiered_pricing: Option, rate_multipliers: Option, 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, model_context_by_model_id_calls: Arc, } 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, model_context_by_model_id_calls: Arc, ) -> 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, 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, 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, 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("a_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("a_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 )); } }