refactor: 拆分 gateway 单体为独立 crate,新增 systemd 部署方案

将 gateway 内部的 model-fetch、provider-transport、scheduler-core、
usage-runtime、video-tasks-core 模块提取为独立 crate;重构 gateway
内部模块结构(state/router/cache/data/query 等);移除大量遗留模块
文件;新增 systemd 二进制部署骨架及相关文档;更新前端 usage 相关
API 和组件。
This commit is contained in:
fawney19
2026-04-05 20:23:16 +08:00
parent cbc811f6ce
commit 763ff03a7b
777 changed files with 42659 additions and 21469 deletions

View File

@@ -7,6 +7,10 @@ repository.workspace = true
description = "Shared billing domain core for Aether Rust migration"
[dependencies]
aether-data.workspace = true
aether-usage-runtime.workspace = true
async-trait.workspace = true
serde.workspace = true
serde_json.workspace = true
thiserror.workspace = true
tokio.workspace = true

View File

@@ -0,0 +1,223 @@
use aether_data::repository::billing::StoredBillingModelContext;
use aether_data::DataLayerError;
use aether_usage_runtime::{UsageEvent, UsageEventType};
use async_trait::async_trait;
use serde_json::{Map, Value};
use crate::{BillingModelPricingSnapshot, BillingService, BillingUsageInput};
#[async_trait]
pub trait BillingModelContextLookup: Send + Sync {
async fn find_billing_model_context(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError>;
}
pub async fn enrich_usage_event_with_billing(
data: &dyn BillingModelContextLookup,
event: &mut UsageEvent,
) -> Result<(), DataLayerError> {
if !matches!(event.event_type, UsageEventType::Completed) {
event.data.total_cost_usd = Some(0.0);
event.data.actual_total_cost_usd = Some(0.0);
return Ok(());
}
let Some(provider_id) = event
.data
.provider_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(());
};
let model_name = event.data.model.trim();
if model_name.is_empty() {
return Ok(());
}
let Some(context) = data
.find_billing_model_context(
provider_id,
event.data.provider_api_key_id.as_deref(),
model_name,
)
.await?
else {
return Ok(());
};
let pricing = map_pricing_context(context);
let input = BillingUsageInput {
task_type: event
.data
.request_type
.clone()
.unwrap_or_else(|| "chat".to_string()),
api_format: event
.data
.endpoint_api_format
.clone()
.or_else(|| event.data.api_format.clone()),
request_count: if event.data.status_code.unwrap_or_default() >= 400
|| event.data.error_message.is_some()
{
0
} else {
1
},
input_tokens: event.data.input_tokens.unwrap_or_default() as i64,
output_tokens: event.data.output_tokens.unwrap_or_default() as i64,
cache_creation_tokens: event.data.cache_creation_input_tokens.unwrap_or_default() as i64,
cache_read_tokens: event.data.cache_read_input_tokens.unwrap_or_default() as i64,
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
};
let computation = BillingService::new()
.calculate(&pricing, &input)
.map_err(|err| {
DataLayerError::UnexpectedValue(format!("billing calculation failed: {err}"))
})?;
event.data.total_cost_usd = Some(computation.cost_result.cost);
event.data.actual_total_cost_usd = Some(computation.actual_total_cost);
merge_billing_snapshot_metadata(
&mut event.data.request_metadata,
&computation.cost_result.snapshot,
computation.rate_multiplier,
computation.is_free_tier,
)?;
Ok(())
}
fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot {
BillingModelPricingSnapshot {
provider_id: context.provider_id,
provider_billing_type: context.provider_billing_type,
provider_api_key_id: context.provider_api_key_id,
provider_api_key_rate_multipliers: context.provider_api_key_rate_multipliers,
provider_api_key_cache_ttl_minutes: context.provider_api_key_cache_ttl_minutes,
global_model_id: context.global_model_id,
global_model_name: context.global_model_name,
global_model_config: context.global_model_config,
default_price_per_request: context.default_price_per_request,
default_tiered_pricing: context.default_tiered_pricing,
model_id: context.model_id,
model_provider_model_name: context.model_provider_model_name,
model_config: context.model_config,
model_price_per_request: context.model_price_per_request,
model_tiered_pricing: context.model_tiered_pricing,
}
}
fn merge_billing_snapshot_metadata(
request_metadata: &mut Option<Value>,
snapshot: &crate::BillingSnapshot,
rate_multiplier: f64,
is_free_tier: bool,
) -> Result<(), DataLayerError> {
let snapshot = serde_json::to_value(snapshot).map_err(|err| {
DataLayerError::UnexpectedValue(format!("failed to serialize billing snapshot: {err}"))
})?;
let mut metadata = match request_metadata.take() {
Some(Value::Object(object)) => object,
_ => Map::new(),
};
metadata.insert("billing_snapshot".to_string(), snapshot);
metadata.insert("rate_multiplier".to_string(), Value::from(rate_multiplier));
metadata.insert("is_free_tier".to_string(), Value::from(is_free_tier));
*request_metadata = Some(Value::Object(metadata));
Ok(())
}
#[cfg(test)]
mod tests {
use aether_data::repository::billing::StoredBillingModelContext;
use aether_usage_runtime::{UsageEvent, UsageEventData, UsageEventType};
use async_trait::async_trait;
use serde_json::json;
use serde_json::Value;
use super::{enrich_usage_event_with_billing, BillingModelContextLookup};
struct TestLookup {
context: Option<StoredBillingModelContext>,
}
#[async_trait]
impl BillingModelContextLookup for TestLookup {
async fn find_billing_model_context(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, aether_data::DataLayerError> {
Ok(self.context.clone())
}
}
#[tokio::test]
async fn enriches_completed_usage_event_with_billing_snapshot() {
let lookup = TestLookup {
context: Some(
StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
Some(json!({"openai:chat": 0.5})),
Some(60),
"global-model-1".to_string(),
"gpt-5".to_string(),
None,
Some(0.02),
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0,"cache_creation_price_per_1m":3.75,"cache_read_price_per_1m":0.30}]})),
Some("model-1".to_string()),
Some("gpt-5-upstream".to_string()),
None,
None,
None,
)
.expect("billing context should build"),
),
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-billing-1",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
provider_id: Some("provider-1".to_string()),
provider_api_key_id: Some("key-1".to_string()),
request_type: Some("chat".to_string()),
api_format: Some("openai:chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
input_tokens: Some(1_000),
output_tokens: Some(500),
cache_read_input_tokens: Some(100),
status_code: Some(200),
..UsageEventData::default()
},
);
enrich_usage_event_with_billing(&lookup, &mut event)
.await
.expect("billing should succeed");
assert!(event.data.total_cost_usd.unwrap_or_default() > 0.0);
assert!(event.data.actual_total_cost_usd.unwrap_or_default() > 0.0);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_snapshot"))
.and_then(|value| value.get("status"))
.and_then(Value::as_str),
Some("complete")
);
}
}

View File

@@ -1,4 +1,5 @@
mod default_rule;
mod event_enrichment;
mod formula_engine;
mod models;
mod precision;
@@ -6,14 +7,17 @@ mod pricing;
mod schema;
mod service;
mod token_normalization;
mod usage_mapper;
pub use aether_usage_runtime::{
map_usage, map_usage_from_response, StandardizedUsage, UsageMapper,
};
pub use default_rule::{normalize_task_type, DefaultBillingRuleGenerator, VirtualBillingRule};
pub use event_enrichment::{enrich_usage_event_with_billing, BillingModelContextLookup};
pub use formula_engine::{
extract_variable_names, BillingIncompleteError, ExpressionEvaluationError, FormulaEngine,
FormulaEvaluationResult, FormulaEvaluationStatus, UnsafeExpressionError,
};
pub use models::{BillingDimension, BillingUnit, CostBreakdown, StandardizedUsage};
pub use models::{BillingDimension, BillingUnit, CostBreakdown};
pub use precision::{
quantize_cost, quantize_display, quantize_value, BILLING_DISPLAY_PRECISION,
BILLING_STORAGE_PRECISION,
@@ -24,4 +28,3 @@ pub use schema::{
};
pub use service::BillingService;
pub use token_normalization::normalize_input_tokens_for_billing;
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};

View File

@@ -33,63 +33,6 @@ impl BillingDimension {
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
pub struct StandardizedUsage {
pub input_tokens: i64,
pub output_tokens: i64,
pub cache_creation_tokens: i64,
pub cache_read_tokens: i64,
pub reasoning_tokens: i64,
pub cache_storage_token_hours: f64,
pub request_count: i64,
pub dimensions: BTreeMap<String, serde_json::Value>,
}
impl StandardizedUsage {
pub fn new() -> Self {
Self {
request_count: 1,
..Self::default()
}
}
pub fn get(&self, field_name: &str) -> Option<serde_json::Value> {
match field_name {
"input_tokens" => Some(serde_json::json!(self.input_tokens)),
"output_tokens" => Some(serde_json::json!(self.output_tokens)),
"cache_creation_tokens" => Some(serde_json::json!(self.cache_creation_tokens)),
"cache_read_tokens" => Some(serde_json::json!(self.cache_read_tokens)),
"reasoning_tokens" => Some(serde_json::json!(self.reasoning_tokens)),
"cache_storage_token_hours" => Some(serde_json::json!(self.cache_storage_token_hours)),
"request_count" => Some(serde_json::json!(self.request_count)),
"extra" | "dimensions" => Some(serde_json::json!(self.dimensions)),
_ => self.dimensions.get(field_name).cloned(),
}
}
pub fn set(&mut self, field_name: &str, value: impl Into<serde_json::Value>) {
let value = value.into();
match field_name {
"input_tokens" => self.input_tokens = as_i64(&value, 0),
"output_tokens" => self.output_tokens = as_i64(&value, 0),
"cache_creation_tokens" => self.cache_creation_tokens = as_i64(&value, 0),
"cache_read_tokens" => self.cache_read_tokens = as_i64(&value, 0),
"reasoning_tokens" => self.reasoning_tokens = as_i64(&value, 0),
"cache_storage_token_hours" => self.cache_storage_token_hours = as_f64(&value, 0.0),
"request_count" => self.request_count = as_i64(&value, 0),
"extra" | "dimensions" => {
self.dimensions = match value {
serde_json::Value::Object(map) => map.into_iter().collect(),
_ => BTreeMap::new(),
}
}
_ => {
self.dimensions.insert(field_name.to_string(), value);
}
}
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
pub struct CostBreakdown {
pub costs: BTreeMap<String, f64>,
@@ -98,20 +41,9 @@ pub struct CostBreakdown {
pub effective_prices: BTreeMap<String, f64>,
}
fn as_i64(value: &serde_json::Value, default: i64) -> i64 {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
.unwrap_or(default)
}
fn as_f64(value: &serde_json::Value, default: f64) -> f64 {
value.as_f64().unwrap_or(default)
}
#[cfg(test)]
mod tests {
use super::{BillingDimension, BillingUnit, StandardizedUsage};
use super::{BillingDimension, BillingUnit};
#[test]
fn dimension_calculates_per_million_tokens() {
@@ -124,17 +56,4 @@ mod tests {
};
assert_eq!(dimension.calculate(500_000.0, 2.0), 1.0);
}
#[test]
fn standardized_usage_reads_and_writes_known_and_extra_fields() {
let mut usage = StandardizedUsage::new();
usage.set("input_tokens", 10);
usage.set("custom_dimension", "value");
assert_eq!(usage.get("input_tokens"), Some(serde_json::json!(10)));
assert_eq!(
usage.get("custom_dimension"),
Some(serde_json::json!("value"))
);
}
}

View File

@@ -15,6 +15,7 @@ futures-util.workspace = true
redis.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
sqlx = { workspace = true, features = ["migrate", "macros"] }
thiserror.workspace = true
tokio.workspace = true

View File

@@ -145,6 +145,7 @@ mod tests {
assert!(backends.transactions().postgres().is_none());
assert!(backends.workers().redis().is_none());
assert!(backends.write().shadow_results().is_none());
assert!(backends.write().settlement().is_none());
assert!(backends.write().usage().is_none());
}
@@ -191,7 +192,9 @@ mod tests {
assert!(backends.write().management_tokens().is_some());
assert!(backends.write().oauth_providers().is_some());
assert!(backends.write().proxy_nodes().is_some());
assert!(backends.write().provider_catalog().is_some());
assert!(backends.write().provider_quotas().is_some());
assert!(backends.write().settlement().is_some());
assert!(backends.write().usage().is_some());
assert!(backends.write().wallets().is_some());
assert!(backends.config().postgres.is_some());
@@ -220,6 +223,7 @@ mod tests {
assert!(backends.read().oauth_providers().is_none());
assert!(backends.transactions().postgres().is_none());
assert!(backends.write().shadow_results().is_none());
assert!(backends.write().settlement().is_none());
assert!(backends.write().usage().is_none());
assert!(backends.config().redis.is_some());
}

View File

@@ -45,9 +45,11 @@ use crate::repository::proxy_nodes::{
use crate::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqlxProviderQuotaRepository,
};
use crate::repository::settlement::{SettlementWriteRepository, SqlxSettlementRepository};
use crate::repository::shadow_results::{
ShadowResultReadRepository, ShadowResultWriteRepository, SqlxShadowResultRepository,
};
use crate::repository::system::{AdminSystemStats, StoredSystemConfigEntry};
use crate::repository::usage::{
SqlxUsageReadRepository, UsageReadRepository, UsageWriteRepository,
};
@@ -260,6 +262,10 @@ impl PostgresBackend {
Arc::new(SqlxWalletRepository::new(self.pool_clone()))
}
pub fn settlement_write_repository(&self) -> Arc<dyn SettlementWriteRepository> {
Arc::new(SqlxSettlementRepository::new(self.pool_clone()))
}
pub fn video_task_read_repository(&self) -> Arc<dyn VideoTaskReadRepository> {
Arc::new(SqlxVideoTaskReadRepository::new(self.pool_clone()))
}
@@ -322,19 +328,20 @@ impl PostgresBackend {
pub async fn list_system_config_entries(
&self,
) -> Result<Vec<(String, serde_json::Value, Option<String>, Option<u64>)>, DataLayerError> {
) -> Result<Vec<StoredSystemConfigEntry>, DataLayerError> {
let rows = sqlx::query(LIST_SYSTEM_CONFIG_ENTRIES_SQL)
.fetch_all(&self.pool)
.await?;
rows.into_iter()
.map(|row| {
Ok((
row.try_get("key")?,
row.try_get("value")?,
row.try_get("description")?,
row.try_get::<Option<i64>, _>("updated_at_unix_secs")?
Ok(StoredSystemConfigEntry {
key: row.try_get("key")?,
value: row.try_get("value")?,
description: row.try_get("description")?,
updated_at_unix_secs: row
.try_get::<Option<i64>, _>("updated_at_unix_secs")?
.map(|value| value.max(0) as u64),
))
})
})
.collect()
}
@@ -344,7 +351,7 @@ impl PostgresBackend {
key: &str,
value: &serde_json::Value,
description: Option<&str>,
) -> Result<(String, serde_json::Value, Option<String>, Option<u64>), DataLayerError> {
) -> Result<StoredSystemConfigEntry, DataLayerError> {
let row = sqlx::query(UPSERT_SYSTEM_CONFIG_ENTRY_SQL)
.bind(uuid::Uuid::new_v4().to_string())
.bind(key)
@@ -352,13 +359,14 @@ impl PostgresBackend {
.bind(description)
.fetch_one(&self.pool)
.await?;
Ok((
row.try_get("key")?,
row.try_get("value")?,
row.try_get("description")?,
row.try_get::<Option<i64>, _>("updated_at_unix_secs")?
Ok(StoredSystemConfigEntry {
key: row.try_get("key")?,
value: row.try_get("value")?,
description: row.try_get("description")?,
updated_at_unix_secs: row
.try_get::<Option<i64>, _>("updated_at_unix_secs")?
.map(|value| value.max(0) as u64),
))
})
}
pub async fn delete_system_config_value(&self, key: &str) -> Result<bool, DataLayerError> {
@@ -369,16 +377,16 @@ impl PostgresBackend {
Ok(result.rows_affected() > 0)
}
pub async fn read_admin_system_stats(&self) -> Result<(u64, u64, u64, u64), DataLayerError> {
pub async fn read_admin_system_stats(&self) -> Result<AdminSystemStats, DataLayerError> {
let row = sqlx::query(READ_ADMIN_SYSTEM_STATS_SQL)
.fetch_one(&self.pool)
.await?;
Ok((
row.try_get::<i64, _>("total_users")?.max(0) as u64,
row.try_get::<i64, _>("active_users")?.max(0) as u64,
row.try_get::<i64, _>("total_api_keys")?.max(0) as u64,
row.try_get::<i64, _>("total_requests")?.max(0) as u64,
))
Ok(AdminSystemStats {
total_users: row.try_get::<i64, _>("total_users")?.max(0) as u64,
active_users: row.try_get::<i64, _>("active_users")?.max(0) as u64,
total_api_keys: row.try_get::<i64, _>("total_api_keys")?.max(0) as u64,
total_requests: row.try_get::<i64, _>("total_requests")?.max(0) as u64,
})
}
}
@@ -431,6 +439,7 @@ mod tests {
let _usage_writer = backend.usage_write_repository();
let _wallet_reader = backend.wallet_read_repository();
let _wallet_writer = backend.wallet_write_repository();
let _settlement_writer = backend.settlement_write_repository();
let _video_task_reader = backend.video_task_read_repository();
let _video_task_writer = backend.video_task_write_repository();
let _transaction_runner = backend.transaction_runner();

View File

@@ -13,6 +13,7 @@ use crate::repository::oauth_providers::OAuthProviderWriteRepository;
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
use crate::repository::settlement::SettlementWriteRepository;
use crate::repository::shadow_results::ShadowResultWriteRepository;
use crate::repository::usage::UsageWriteRepository;
use crate::repository::video_tasks::VideoTaskWriteRepository;
@@ -32,6 +33,7 @@ pub struct DataWriteRepositories {
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
settlement: Option<Arc<dyn SettlementWriteRepository>>,
usage: Option<Arc<dyn UsageWriteRepository>>,
video_tasks: Option<Arc<dyn VideoTaskWriteRepository>>,
wallets: Option<Arc<dyn WalletWriteRepository>>,
@@ -55,6 +57,7 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
.field("has_settlement", &self.settlement.is_some())
.field("has_usage", &self.usage.is_some())
.field("has_video_tasks", &self.video_tasks.is_some())
.field("has_wallets", &self.wallets.is_some())
@@ -78,6 +81,7 @@ impl DataWriteRepositories {
proxy_nodes: postgres.map(PostgresBackend::proxy_node_write_repository),
provider_catalog: postgres.map(PostgresBackend::provider_catalog_write_repository),
provider_quotas: postgres.map(PostgresBackend::provider_quota_write_repository),
settlement: postgres.map(PostgresBackend::settlement_write_repository),
usage: postgres.map(PostgresBackend::usage_write_repository),
video_tasks: postgres.map(PostgresBackend::video_task_write_repository),
wallets: postgres.map(PostgresBackend::wallet_write_repository),
@@ -136,6 +140,10 @@ impl DataWriteRepositories {
self.provider_catalog.clone()
}
pub fn settlement(&self) -> Option<Arc<dyn SettlementWriteRepository>> {
self.settlement.clone()
}
pub fn video_tasks(&self) -> Option<Arc<dyn VideoTaskWriteRepository>> {
self.video_tasks.clone()
}
@@ -157,6 +165,7 @@ impl DataWriteRepositories {
|| self.proxy_nodes.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()
|| self.settlement.is_some()
|| self.usage.is_some()
|| self.video_tasks.is_some()
|| self.wallets.is_some()
@@ -198,6 +207,7 @@ mod tests {
assert!(write.proxy_nodes().is_some());
assert!(write.provider_catalog().is_some());
assert!(write.provider_quotas().is_some());
assert!(write.settlement().is_some());
assert!(write.usage().is_some());
assert!(write.video_tasks().is_some());
assert!(write.wallets().is_some());

View File

@@ -0,0 +1,301 @@
use async_trait::async_trait;
use crate::repository::auth::ResolvedAuthApiKeySnapshot;
use crate::repository::candidates::DecisionTrace;
use crate::repository::usage::StoredRequestUsageAudit;
use crate::DataLayerError;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct RequestAuditBundle {
pub request_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub usage: Option<StoredRequestUsageAudit>,
#[serde(skip_serializing_if = "Option::is_none")]
pub decision_trace: Option<DecisionTrace>,
#[serde(skip_serializing_if = "Option::is_none")]
pub auth_snapshot: Option<ResolvedAuthApiKeySnapshot>,
}
#[async_trait]
pub trait RequestAuditReader {
async fn find_request_usage_audit_by_request_id(
&self,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError>;
async fn read_request_decision_trace(
&self,
request_id: &str,
attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError>;
async fn read_resolved_auth_api_key_snapshot(
&self,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<ResolvedAuthApiKeySnapshot>, DataLayerError>;
}
pub async fn read_request_audit_bundle(
state: &impl RequestAuditReader,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
let usage = state
.find_request_usage_audit_by_request_id(request_id)
.await?;
let decision_trace = state
.read_request_decision_trace(request_id, attempted_only)
.await?;
let auth_snapshot = if let Some(usage) = usage.as_ref() {
match (usage.user_id.as_deref(), usage.api_key_id.as_deref()) {
(Some(user_id), Some(api_key_id)) => {
state
.read_resolved_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await?
}
_ => None,
}
} else {
None
};
if usage.is_none() && decision_trace.is_none() && auth_snapshot.is_none() {
return Ok(None);
}
Ok(Some(RequestAuditBundle {
request_id: request_id.to_string(),
usage,
decision_trace,
auth_snapshot,
}))
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use async_trait::async_trait;
use super::{read_request_audit_bundle, RequestAuditReader};
use crate::repository::auth::{ResolvedAuthApiKeySnapshot, StoredAuthApiKeySnapshot};
use crate::repository::candidates::{
DecisionTrace, DecisionTraceCandidate, RequestCandidateFinalStatus, RequestCandidateStatus,
StoredRequestCandidate,
};
use crate::repository::usage::StoredRequestUsageAudit;
use crate::DataLayerError;
#[derive(Default)]
struct FakeRequestAuditReader {
usage: Option<StoredRequestUsageAudit>,
decision_trace: Option<DecisionTrace>,
auth_snapshot: Option<ResolvedAuthApiKeySnapshot>,
auth_snapshot_reads: AtomicUsize,
}
#[async_trait]
impl RequestAuditReader for FakeRequestAuditReader {
async fn find_request_usage_audit_by_request_id(
&self,
_request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
Ok(self.usage.clone())
}
async fn read_request_decision_trace(
&self,
_request_id: &str,
_attempted_only: bool,
) -> Result<Option<DecisionTrace>, DataLayerError> {
Ok(self.decision_trace.clone())
}
async fn read_resolved_auth_api_key_snapshot(
&self,
_user_id: &str,
_api_key_id: &str,
_now_unix_secs: u64,
) -> Result<Option<ResolvedAuthApiKeySnapshot>, DataLayerError> {
self.auth_snapshot_reads.fetch_add(1, Ordering::Relaxed);
Ok(self.auth_snapshot.clone())
}
}
#[tokio::test]
async fn read_request_audit_bundle_resolves_usage_trace_and_auth_snapshot() {
let state = FakeRequestAuditReader {
usage: Some(sample_usage("req-audit-1")),
decision_trace: Some(sample_decision_trace("req-audit-1")),
auth_snapshot: Some(sample_resolved_auth_snapshot("user-1", "api-key-1")),
auth_snapshot_reads: AtomicUsize::new(0),
};
let bundle = read_request_audit_bundle(&state, "req-audit-1", true, 123)
.await
.expect("bundle should read")
.expect("bundle should exist");
assert_eq!(bundle.request_id, "req-audit-1");
assert_eq!(
bundle
.usage
.as_ref()
.map(|usage| usage.provider_name.as_str()),
Some("OpenAI")
);
assert_eq!(
bundle
.decision_trace
.as_ref()
.map(|trace| trace.total_candidates),
Some(1)
);
assert_eq!(
bundle
.auth_snapshot
.as_ref()
.map(|snapshot| snapshot.api_key_id.as_str()),
Some("api-key-1")
);
assert_eq!(state.auth_snapshot_reads.load(Ordering::Relaxed), 1);
}
#[tokio::test]
async fn read_request_audit_bundle_returns_none_when_all_sources_are_empty() {
let state = FakeRequestAuditReader::default();
let bundle = read_request_audit_bundle(&state, "req-audit-empty", false, 123)
.await
.expect("bundle should read");
assert!(bundle.is_none());
assert_eq!(state.auth_snapshot_reads.load(Ordering::Relaxed), 0);
}
fn sample_usage(request_id: &str) -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
"OpenAI".to_string(),
"gpt-4.1".to_string(),
None,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
false,
false,
120,
40,
160,
0.24,
0.36,
Some(200),
None,
None,
Some(450),
Some(120),
"completed".to_string(),
"settled".to_string(),
100,
101,
Some(102),
)
.expect("usage should build")
}
fn sample_decision_trace(request_id: &str) -> DecisionTrace {
let candidate = StoredRequestCandidate::new(
"cand-1".to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(37),
None,
None,
None,
100,
Some(101),
Some(102),
)
.expect("candidate should build");
DecisionTrace {
request_id: request_id.to_string(),
total_candidates: 1,
final_status: RequestCandidateFinalStatus::Success,
total_latency_ms: 37,
candidates: vec![DecisionTraceCandidate {
candidate,
provider_name: Some("OpenAI".to_string()),
provider_website: None,
provider_type: Some("custom".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
provider_key_name: Some("prod".to_string()),
provider_key_auth_type: Some("api_key".to_string()),
provider_key_capabilities: None,
provider_key_is_active: Some(true),
}],
}
}
fn sample_resolved_auth_snapshot(
user_id: &str,
api_key_id: &str,
) -> ResolvedAuthApiKeySnapshot {
let stored = StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(4_102_444_800),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("auth snapshot should build");
ResolvedAuthApiKeySnapshot::from_stored(stored, 123)
}
}

View File

@@ -5,8 +5,11 @@ mod types;
pub use memory::InMemoryAuthApiKeySnapshotRepository;
pub use sql::SqlxAuthApiKeySnapshotReadRepository;
pub use types::{
AuthApiKeyExportSummary, AuthApiKeyLookupKey, AuthApiKeyReadRepository,
AuthApiKeyWriteRepository, AuthRepository, CreateStandaloneApiKeyRecord,
CreateUserApiKeyRecord, StandaloneApiKeyExportListQuery, StoredAuthApiKeyExportRecord,
StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord, UpdateUserApiKeyBasicRecord,
read_resolved_auth_api_key_snapshot, read_resolved_auth_api_key_snapshot_by_key_hash,
read_resolved_auth_api_key_snapshot_by_user_api_key_ids, AuthApiKeyExportSummary,
AuthApiKeyLookupKey, AuthApiKeyReadRepository, AuthApiKeyWriteRepository, AuthRepository,
CreateStandaloneApiKeyRecord, CreateUserApiKeyRecord, ResolvedAuthApiKeySnapshot,
ResolvedAuthApiKeySnapshotReader, StandaloneApiKeyExportListQuery,
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot, UpdateStandaloneApiKeyBasicRecord,
UpdateUserApiKeyBasicRecord,
};

View File

@@ -124,6 +124,131 @@ impl StoredAuthApiKeySnapshot {
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ResolvedAuthApiKeySnapshot {
pub user_id: String,
pub username: String,
pub email: Option<String>,
pub user_role: String,
pub user_auth_source: String,
pub user_is_active: bool,
pub user_is_deleted: bool,
pub user_rate_limit: Option<i32>,
pub user_allowed_providers: Option<Vec<String>>,
pub user_allowed_api_formats: Option<Vec<String>>,
pub user_allowed_models: Option<Vec<String>>,
pub api_key_id: String,
pub api_key_name: Option<String>,
pub api_key_is_active: bool,
pub api_key_is_locked: bool,
pub api_key_is_standalone: bool,
pub api_key_rate_limit: Option<i32>,
pub api_key_concurrent_limit: Option<i32>,
pub api_key_expires_at_unix_secs: Option<u64>,
pub api_key_allowed_providers: Option<Vec<String>>,
pub api_key_allowed_api_formats: Option<Vec<String>>,
pub api_key_allowed_models: Option<Vec<String>>,
pub currently_usable: bool,
}
impl ResolvedAuthApiKeySnapshot {
pub fn from_stored(snapshot: StoredAuthApiKeySnapshot, now_unix_secs: u64) -> Self {
let currently_usable = snapshot.is_currently_usable(now_unix_secs);
Self {
user_id: snapshot.user_id,
username: snapshot.username,
email: snapshot.email,
user_role: snapshot.user_role,
user_auth_source: snapshot.user_auth_source,
user_is_active: snapshot.user_is_active,
user_is_deleted: snapshot.user_is_deleted,
user_rate_limit: snapshot.user_rate_limit,
user_allowed_providers: snapshot.user_allowed_providers,
user_allowed_api_formats: snapshot.user_allowed_api_formats,
user_allowed_models: snapshot.user_allowed_models,
api_key_id: snapshot.api_key_id,
api_key_name: snapshot.api_key_name,
api_key_is_active: snapshot.api_key_is_active,
api_key_is_locked: snapshot.api_key_is_locked,
api_key_is_standalone: snapshot.api_key_is_standalone,
api_key_rate_limit: snapshot.api_key_rate_limit,
api_key_concurrent_limit: snapshot.api_key_concurrent_limit,
api_key_expires_at_unix_secs: snapshot.api_key_expires_at_unix_secs,
api_key_allowed_providers: snapshot.api_key_allowed_providers,
api_key_allowed_api_formats: snapshot.api_key_allowed_api_formats,
api_key_allowed_models: snapshot.api_key_allowed_models,
currently_usable,
}
}
pub fn effective_allowed_providers(&self) -> Option<&[String]> {
self.api_key_allowed_providers
.as_deref()
.or(self.user_allowed_providers.as_deref())
}
pub fn effective_allowed_api_formats(&self) -> Option<&[String]> {
self.api_key_allowed_api_formats
.as_deref()
.or(self.user_allowed_api_formats.as_deref())
}
pub fn effective_allowed_models(&self) -> Option<&[String]> {
self.api_key_allowed_models
.as_deref()
.or(self.user_allowed_models.as_deref())
}
}
#[async_trait]
pub trait ResolvedAuthApiKeySnapshotReader: Send + Sync {
async fn find_stored_auth_api_key_snapshot(
&self,
key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError>;
}
pub async fn read_resolved_auth_api_key_snapshot(
reader: &impl ResolvedAuthApiKeySnapshotReader,
key: AuthApiKeyLookupKey<'_>,
now_unix_secs: u64,
) -> Result<Option<ResolvedAuthApiKeySnapshot>, crate::DataLayerError> {
Ok(reader
.find_stored_auth_api_key_snapshot(key)
.await?
.map(|snapshot| ResolvedAuthApiKeySnapshot::from_stored(snapshot, now_unix_secs)))
}
pub async fn read_resolved_auth_api_key_snapshot_by_key_hash(
reader: &impl ResolvedAuthApiKeySnapshotReader,
key_hash: &str,
now_unix_secs: u64,
) -> Result<Option<ResolvedAuthApiKeySnapshot>, crate::DataLayerError> {
read_resolved_auth_api_key_snapshot(
reader,
AuthApiKeyLookupKey::KeyHash(key_hash),
now_unix_secs,
)
.await
}
pub async fn read_resolved_auth_api_key_snapshot_by_user_api_key_ids(
reader: &impl ResolvedAuthApiKeySnapshotReader,
user_id: &str,
api_key_id: &str,
now_unix_secs: u64,
) -> Result<Option<ResolvedAuthApiKeySnapshot>, crate::DataLayerError> {
read_resolved_auth_api_key_snapshot(
reader,
AuthApiKeyLookupKey::UserApiKeyIds {
user_id,
api_key_id,
},
now_unix_secs,
)
.await
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAuthApiKeyExportRecord {
pub user_id: String,
@@ -485,7 +610,14 @@ fn parse_u64_i64(value: i64, field_name: &str) -> Result<u64, crate::DataLayerEr
#[cfg(test)]
mod tests {
use super::{StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot};
use async_trait::async_trait;
use super::{
read_resolved_auth_api_key_snapshot_by_key_hash,
read_resolved_auth_api_key_snapshot_by_user_api_key_ids, AuthApiKeyLookupKey,
ResolvedAuthApiKeySnapshot, ResolvedAuthApiKeySnapshotReader, StoredAuthApiKeyExportRecord,
StoredAuthApiKeySnapshot,
};
#[test]
fn rejects_non_array_allowed_providers() {
@@ -611,6 +743,50 @@ mod tests {
assert!(!snapshot.is_currently_usable(101));
}
#[test]
fn resolved_snapshot_prefers_api_key_lists_over_user_lists() {
let snapshot = StoredAuthApiKeySnapshot::new(
"user-1".to_string(),
"alice".to_string(),
None,
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
"key-1".to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["anthropic"])),
None,
None,
)
.expect("snapshot should build");
let resolved = ResolvedAuthApiKeySnapshot::from_stored(snapshot, 150);
assert!(resolved.currently_usable);
assert_eq!(
resolved.effective_allowed_providers(),
Some(&["anthropic".to_string()][..])
);
assert_eq!(
resolved.effective_allowed_api_formats(),
Some(&["openai:chat".to_string()][..])
);
assert_eq!(
resolved.effective_allowed_models(),
Some(&["gpt-4.1".to_string()][..])
);
}
#[test]
fn export_record_rejects_negative_totals() {
assert!(StoredAuthApiKeyExportRecord::new(
@@ -665,4 +841,89 @@ mod tests {
assert_eq!(record.total_requests, 12);
assert_eq!(record.total_cost_usd, 1.25);
}
#[derive(Default)]
struct FakeResolvedAuthApiKeySnapshotReader {
stored: Option<StoredAuthApiKeySnapshot>,
}
#[async_trait]
impl ResolvedAuthApiKeySnapshotReader for FakeResolvedAuthApiKeySnapshotReader {
async fn find_stored_auth_api_key_snapshot(
&self,
_key: AuthApiKeyLookupKey<'_>,
) -> Result<Option<StoredAuthApiKeySnapshot>, crate::DataLayerError> {
Ok(self.stored.clone())
}
}
#[tokio::test]
async fn reads_resolved_auth_snapshot_by_user_api_key_ids() {
let reader = FakeResolvedAuthApiKeySnapshotReader {
stored: Some(sample_auth_snapshot("key-1", "user-1")),
};
let snapshot = read_resolved_auth_api_key_snapshot_by_user_api_key_ids(
&reader, "user-1", "key-1", 150,
)
.await
.expect("snapshot should read")
.expect("snapshot should exist");
assert_eq!(snapshot.user_id, "user-1");
assert_eq!(snapshot.api_key_id, "key-1");
assert!(snapshot.currently_usable);
}
#[tokio::test]
async fn reads_resolved_auth_snapshot_by_key_hash() {
let reader = FakeResolvedAuthApiKeySnapshotReader {
stored: Some(sample_auth_snapshot("key-1", "user-1")),
};
let snapshot = read_resolved_auth_api_key_snapshot_by_key_hash(&reader, "hash-lookup", 150)
.await
.expect("snapshot should read")
.expect("snapshot should exist");
assert_eq!(
snapshot.effective_allowed_providers(),
Some(&["openai".to_string()][..])
);
assert_eq!(
snapshot.effective_allowed_api_formats(),
Some(&["openai:chat".to_string()][..])
);
assert_eq!(
snapshot.effective_allowed_models(),
Some(&["gpt-4.1".to_string()][..])
);
}
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {
StoredAuthApiKeySnapshot::new(
user_id.to_string(),
"alice".to_string(),
Some("alice@example.com".to_string()),
"user".to_string(),
"local".to_string(),
true,
false,
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
api_key_id.to_string(),
Some("default".to_string()),
true,
false,
false,
Some(60),
Some(5),
Some(200),
Some(serde_json::json!(["openai"])),
Some(serde_json::json!(["openai:chat"])),
Some(serde_json::json!(["gpt-4.1"])),
)
.expect("snapshot should build")
}
}

View File

@@ -4,4 +4,8 @@ mod types;
pub use memory::InMemoryBillingReadRepository;
pub use sql::SqlxBillingReadRepository;
pub use types::{BillingReadRepository, StoredBillingModelContext};
pub use types::{
AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingPresetApplyResult,
AdminBillingRuleRecord, AdminBillingRuleWriteInput, BillingReadRepository,
StoredBillingModelContext,
};

View File

@@ -74,6 +74,74 @@ impl StoredBillingModelContext {
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminBillingRuleRecord {
pub id: String,
pub name: String,
pub task_type: String,
pub global_model_id: Option<String>,
pub model_id: Option<String>,
pub expression: String,
pub variables: Value,
pub dimension_mappings: Value,
pub is_enabled: bool,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AdminBillingRuleWriteInput {
pub name: String,
pub task_type: String,
pub global_model_id: Option<String>,
pub model_id: Option<String>,
pub expression: String,
pub variables: Value,
pub dimension_mappings: Value,
pub is_enabled: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminBillingCollectorRecord {
pub id: String,
pub api_format: String,
pub task_type: String,
pub dimension_name: String,
pub source_type: String,
pub source_path: Option<String>,
pub value_type: String,
pub transform_expression: Option<String>,
pub default_value: Option<String>,
pub priority: i32,
pub is_enabled: bool,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct AdminBillingCollectorWriteInput {
pub api_format: String,
pub task_type: String,
pub dimension_name: String,
pub source_type: String,
pub source_path: Option<String>,
pub value_type: String,
pub transform_expression: Option<String>,
pub default_value: Option<String>,
pub priority: i32,
pub is_enabled: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminBillingPresetApplyResult {
pub preset: String,
pub mode: String,
pub created: u64,
pub updated: u64,
pub skipped: u64,
pub errors: Vec<String>,
}
#[async_trait]
pub trait BillingReadRepository: Send + Sync {
async fn find_model_context(

View File

@@ -5,7 +5,9 @@ mod types;
pub use memory::InMemoryRequestCandidateRepository;
pub use sql::SqlxRequestCandidateReadRepository;
pub use types::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
RequestCandidateRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
build_decision_trace, derive_request_candidate_final_status, DecisionTrace,
DecisionTraceCandidate, PublicHealthStatusCount, PublicHealthTimelineBucket,
RequestCandidateFinalStatus, RequestCandidateReadRepository, RequestCandidateRepository,
RequestCandidateStatus, RequestCandidateTrace, RequestCandidateWriteRepository,
StoredRequestCandidate, UpsertRequestCandidateRecord,
};

View File

@@ -1,5 +1,11 @@
use std::collections::BTreeMap;
use async_trait::async_trait;
use crate::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RequestCandidateStatus {
@@ -185,6 +191,210 @@ impl StoredRequestCandidate {
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RequestCandidateFinalStatus {
Success,
Failed,
Cancelled,
Streaming,
Pending,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct RequestCandidateTrace {
pub request_id: String,
pub total_candidates: usize,
pub final_status: RequestCandidateFinalStatus,
pub total_latency_ms: u64,
pub candidates: Vec<StoredRequestCandidate>,
}
impl RequestCandidateTrace {
pub fn from_candidates(
request_id: impl Into<String>,
all_candidates: Vec<StoredRequestCandidate>,
attempted_only: bool,
) -> Option<Self> {
if all_candidates.is_empty() {
return None;
}
let candidates = if attempted_only {
all_candidates
.iter()
.filter(|candidate| {
candidate
.status
.is_attempted(candidate.started_at_unix_secs)
})
.cloned()
.collect::<Vec<_>>()
} else {
all_candidates.clone()
};
let total_latency_ms = candidates
.iter()
.filter(|candidate| {
matches!(
candidate.status,
RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
) && candidate.latency_ms.is_some()
})
.map(|candidate| candidate.latency_ms.unwrap_or(0))
.sum();
let final_status_source = if attempted_only && candidates.is_empty() {
&all_candidates
} else {
&candidates
};
Some(Self {
request_id: request_id.into(),
total_candidates: candidates.len(),
final_status: derive_request_candidate_final_status(final_status_source),
total_latency_ms,
candidates,
})
}
}
pub fn derive_request_candidate_final_status(
candidates: &[StoredRequestCandidate],
) -> RequestCandidateFinalStatus {
let has_success = candidates.iter().any(|candidate| {
candidate.status == RequestCandidateStatus::Success
|| matches!(candidate.status_code, Some(status_code) if (200..300).contains(&status_code))
});
if has_success {
return RequestCandidateFinalStatus::Success;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Streaming)
{
return RequestCandidateFinalStatus::Streaming;
}
if candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Pending)
{
return RequestCandidateFinalStatus::Pending;
}
let has_cancelled = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled);
let has_failed = candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Failed);
if has_cancelled && !has_failed {
return RequestCandidateFinalStatus::Cancelled;
}
RequestCandidateFinalStatus::Failed
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DecisionTraceCandidate {
#[serde(flatten)]
pub candidate: StoredRequestCandidate,
pub provider_name: Option<String>,
pub provider_website: Option<String>,
pub provider_type: Option<String>,
pub endpoint_api_format: Option<String>,
pub endpoint_api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub provider_key_name: Option<String>,
pub provider_key_auth_type: Option<String>,
pub provider_key_capabilities: Option<serde_json::Value>,
pub provider_key_is_active: Option<bool>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct DecisionTrace {
pub request_id: String,
pub total_candidates: usize,
pub final_status: RequestCandidateFinalStatus,
pub total_latency_ms: u64,
pub candidates: Vec<DecisionTraceCandidate>,
}
pub fn build_decision_trace(
trace: RequestCandidateTrace,
providers: Vec<StoredProviderCatalogProvider>,
endpoints: Vec<StoredProviderCatalogEndpoint>,
keys: Vec<StoredProviderCatalogKey>,
) -> DecisionTrace {
let provider_map = providers
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let endpoint_map = endpoints
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
let key_map = keys
.into_iter()
.map(|item| (item.id.clone(), item))
.collect::<BTreeMap<_, _>>();
DecisionTrace {
request_id: trace.request_id,
total_candidates: trace.total_candidates,
final_status: trace.final_status,
total_latency_ms: trace.total_latency_ms,
candidates: trace
.candidates
.into_iter()
.map(|candidate| {
enrich_decision_trace_candidate(candidate, &provider_map, &endpoint_map, &key_map)
})
.collect(),
}
}
fn enrich_decision_trace_candidate(
candidate: StoredRequestCandidate,
provider_map: &BTreeMap<String, StoredProviderCatalogProvider>,
endpoint_map: &BTreeMap<String, StoredProviderCatalogEndpoint>,
key_map: &BTreeMap<String, StoredProviderCatalogKey>,
) -> DecisionTraceCandidate {
let provider = candidate
.provider_id
.as_ref()
.and_then(|provider_id| provider_map.get(provider_id));
let endpoint = candidate
.endpoint_id
.as_ref()
.and_then(|endpoint_id| endpoint_map.get(endpoint_id));
let provider_key = candidate
.key_id
.as_ref()
.and_then(|key_id| key_map.get(key_id));
DecisionTraceCandidate {
provider_name: provider.map(|item| item.name.clone()),
provider_website: provider.and_then(|item| item.website.clone()),
provider_type: provider.map(|item| item.provider_type.clone()),
endpoint_api_format: endpoint.map(|item| item.api_format.clone()),
endpoint_api_family: endpoint.and_then(|item| item.api_family.clone()),
endpoint_kind: endpoint.and_then(|item| item.endpoint_kind.clone()),
provider_key_name: provider_key
.map(|item| item.name.clone())
.or_else(|| candidate.api_key_name.clone()),
provider_key_auth_type: provider_key.map(|item| item.auth_type.clone()),
provider_key_capabilities: provider_key.and_then(|item| item.capabilities.clone()),
provider_key_is_active: provider_key.map(|item| item.is_active),
candidate,
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PublicHealthStatusCount {
pub endpoint_id: String,
@@ -313,7 +523,14 @@ impl<T> RequestCandidateRepository for T where
#[cfg(test)]
mod tests {
use super::{RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord};
use super::{
build_decision_trace, derive_request_candidate_final_status, DecisionTrace,
DecisionTraceCandidate, RequestCandidateFinalStatus, RequestCandidateStatus,
RequestCandidateTrace, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use crate::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
#[test]
fn parses_status_from_database_text() {
@@ -359,6 +576,184 @@ mod tests {
.is_err());
}
fn sample_candidate(
id: &str,
request_id: &str,
candidate_index: i32,
status: RequestCandidateStatus,
started_at_unix_secs: Option<i64>,
latency_ms: Option<i32>,
status_code: Option<i32>,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
request_id.to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
Some("alice".to_string()),
Some("default".to_string()),
candidate_index,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("provider-key-1".to_string()),
status,
None,
false,
status_code,
None,
None,
latency_ms,
Some(1),
None,
None,
100 + i64::from(candidate_index),
started_at_unix_secs,
started_at_unix_secs.map(|value| value + 1),
)
.expect("candidate should build")
}
#[test]
fn derives_request_candidate_final_status_preferring_success() {
let candidates = vec![sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Success,
Some(100),
Some(25),
Some(200),
)];
assert_eq!(
derive_request_candidate_final_status(&candidates),
RequestCandidateFinalStatus::Success
);
}
#[test]
fn request_candidate_trace_filters_attempted_rows() {
let trace = RequestCandidateTrace::from_candidates(
"req-1",
vec![
sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Pending,
None,
None,
None,
),
sample_candidate(
"cand-2",
"req-1",
1,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(502),
),
],
true,
)
.expect("trace should exist");
assert_eq!(trace.total_candidates, 1);
assert_eq!(trace.candidates[0].id, "cand-2");
assert_eq!(trace.final_status, RequestCandidateFinalStatus::Failed);
assert_eq!(trace.total_latency_ms, 33);
}
fn sample_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
Some("https://openai.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
}
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
}
fn sample_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"provider-key-1".to_string(),
"provider-1".to_string(),
"prod-key".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
}
#[test]
fn build_decision_trace_enriches_candidate_with_provider_catalog_metadata() {
let trace = RequestCandidateTrace::from_candidates(
"req-1",
vec![sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Failed,
Some(101),
Some(37),
Some(502),
)],
true,
)
.expect("trace should exist");
assert_eq!(
build_decision_trace(
trace,
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
),
DecisionTrace {
request_id: "req-1".to_string(),
total_candidates: 1,
final_status: RequestCandidateFinalStatus::Failed,
total_latency_ms: 37,
candidates: vec![DecisionTraceCandidate {
candidate: sample_candidate(
"cand-1",
"req-1",
0,
RequestCandidateStatus::Failed,
Some(101),
Some(37),
Some(502),
),
provider_name: Some("OpenAI".to_string()),
provider_website: Some("https://openai.com".to_string()),
provider_type: Some("custom".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
endpoint_api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
provider_key_name: Some("prod-key".to_string()),
provider_key_auth_type: Some("api_key".to_string()),
provider_key_capabilities: Some(serde_json::json!({"cache_1h": true})),
provider_key_is_active: Some(true),
}],
}
);
}
#[test]
fn rejects_negative_created_at() {
assert!(StoredRequestCandidate::new(

View File

@@ -1,4 +1,5 @@
pub mod announcements;
pub mod audit;
pub mod auth;
pub mod auth_modules;
pub mod billing;
@@ -9,9 +10,12 @@ pub mod global_models;
pub mod management_tokens;
pub mod oauth_providers;
pub mod provider_catalog;
pub mod provider_oauth;
pub mod proxy_nodes;
pub mod quota;
pub mod settlement;
pub mod shadow_results;
pub mod system;
pub mod usage;
pub mod users;
pub mod video_tasks;

View File

@@ -0,0 +1,191 @@
const KIRO_DEVICE_AUTH_SESSION_PREFIX: &str = "device_auth_session:";
const PROVIDER_OAUTH_BATCH_TASK_PREFIX: &str = "provider_oauth_batch_task:";
const PROVIDER_OAUTH_STATE_PREFIX: &str = "provider_oauth_state:";
pub const KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS: u64 = 60;
pub const PROVIDER_OAUTH_BATCH_TASK_TTL_SECS: u64 = 24 * 60 * 60;
pub const PROVIDER_OAUTH_STATE_TTL_SECS: u64 = 600;
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminProviderOAuthDeviceSession {
pub provider_id: String,
pub region: String,
pub client_id: String,
pub client_secret: String,
pub device_code: String,
pub interval: u64,
pub expires_at_unix_secs: u64,
pub status: String,
pub proxy_node_id: Option<String>,
pub created_at_unix_secs: u64,
pub key_id: Option<String>,
pub email: Option<String>,
pub replaced: bool,
pub error_msg: Option<String>,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminProviderOAuthState {
pub key_id: String,
pub provider_id: String,
pub provider_type: String,
pub pkce_verifier: Option<String>,
}
pub fn provider_oauth_device_session_storage_key(session_id: &str) -> String {
format!("{KIRO_DEVICE_AUTH_SESSION_PREFIX}{session_id}")
}
pub fn provider_oauth_state_storage_key(nonce: &str) -> String {
format!("{PROVIDER_OAUTH_STATE_PREFIX}{nonce}")
}
pub fn provider_oauth_batch_task_storage_key(task_id: &str) -> String {
format!("{PROVIDER_OAUTH_BATCH_TASK_PREFIX}{task_id}")
}
pub fn build_provider_oauth_batch_task_status_payload(
provider_id: &str,
state: &serde_json::Map<String, serde_json::Value>,
) -> serde_json::Value {
let now_unix_secs = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or(0);
let raw_status = state
.get("status")
.and_then(serde_json::Value::as_str)
.unwrap_or("failed");
let normalized_status = match raw_status {
"submitted" | "processing" | "completed" | "failed" => raw_status,
_ => "failed",
};
let error_samples = state
.get("error_samples")
.and_then(serde_json::Value::as_array)
.map(|items| {
items
.iter()
.filter(|item| item.is_object())
.cloned()
.collect::<Vec<_>>()
})
.unwrap_or_default();
serde_json::json!({
"task_id": state
.get("task_id")
.and_then(serde_json::Value::as_str)
.unwrap_or_default(),
"provider_id": provider_id,
"provider_type": state
.get("provider_type")
.and_then(serde_json::Value::as_str)
.unwrap_or_default(),
"status": normalized_status,
"total": state.get("total").and_then(serde_json::Value::as_i64).unwrap_or(0),
"processed": state.get("processed").and_then(serde_json::Value::as_i64).unwrap_or(0),
"success": state.get("success").and_then(serde_json::Value::as_i64).unwrap_or(0),
"failed": state.get("failed").and_then(serde_json::Value::as_i64).unwrap_or(0),
"progress_percent": state
.get("progress_percent")
.and_then(serde_json::Value::as_i64)
.unwrap_or(0)
.clamp(0, 100),
"message": state.get("message").cloned().unwrap_or(serde_json::Value::Null),
"error": state.get("error").cloned().unwrap_or(serde_json::Value::Null),
"error_samples": error_samples,
"created_at": state
.get("created_at")
.and_then(serde_json::Value::as_u64)
.unwrap_or(now_unix_secs),
"started_at": state.get("started_at").cloned().unwrap_or(serde_json::Value::Null),
"finished_at": state
.get("finished_at")
.cloned()
.unwrap_or(serde_json::Value::Null),
"updated_at": state
.get("updated_at")
.and_then(serde_json::Value::as_u64)
.unwrap_or(now_unix_secs),
})
}
#[cfg(test)]
mod tests {
use super::{
build_provider_oauth_batch_task_status_payload, provider_oauth_batch_task_storage_key,
provider_oauth_device_session_storage_key, provider_oauth_state_storage_key,
KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS, PROVIDER_OAUTH_BATCH_TASK_TTL_SECS,
PROVIDER_OAUTH_STATE_TTL_SECS,
};
use serde_json::json;
#[test]
fn builds_provider_oauth_storage_keys_with_expected_prefixes() {
assert_eq!(
provider_oauth_device_session_storage_key("session-123"),
"device_auth_session:session-123"
);
assert_eq!(
provider_oauth_state_storage_key("nonce-123"),
"provider_oauth_state:nonce-123"
);
assert_eq!(
provider_oauth_batch_task_storage_key("task-123"),
"provider_oauth_batch_task:task-123"
);
}
#[test]
fn provider_oauth_storage_ttls_match_gateway_expectations() {
assert_eq!(KIRO_DEVICE_AUTH_SESSION_TTL_BUFFER_SECS, 60);
assert_eq!(PROVIDER_OAUTH_BATCH_TASK_TTL_SECS, 24 * 60 * 60);
assert_eq!(PROVIDER_OAUTH_STATE_TTL_SECS, 600);
}
#[test]
fn batch_task_status_payload_normalizes_status_and_clamps_progress() {
let input = json!({
"task_id": "task-123",
"provider_type": "codex",
"status": "weird",
"total": 4,
"processed": 2,
"success": 1,
"failed": 1,
"progress_percent": 999,
"error_samples": [
{"detail": "x"},
"skip-me"
],
"created_at": 1u64,
"updated_at": 2u64
});
let payload = build_provider_oauth_batch_task_status_payload(
"provider-123",
input.as_object().expect("input should be object"),
);
assert_eq!(
payload.get("provider_id").and_then(|v| v.as_str()),
Some("provider-123")
);
assert_eq!(
payload.get("status").and_then(|v| v.as_str()),
Some("failed")
);
assert_eq!(
payload.get("progress_percent").and_then(|v| v.as_i64()),
Some(100)
);
assert_eq!(
payload
.get("error_samples")
.and_then(|v| v.as_array())
.map(Vec::len),
Some(1)
);
}
}

View File

@@ -0,0 +1,228 @@
use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
use async_trait::async_trait;
use super::types::{SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput};
use crate::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
use crate::DataLayerError;
#[derive(Debug)]
enum InMemorySettlementWalletStore {
Owned(RwLock<BTreeMap<String, StoredWalletSnapshot>>),
Shared(Arc<InMemoryWalletRepository>),
}
impl Default for InMemorySettlementWalletStore {
fn default() -> Self {
Self::Owned(RwLock::new(BTreeMap::new()))
}
}
impl InMemorySettlementWalletStore {
fn seeded<I>(items: I) -> Self
where
I: IntoIterator<Item = StoredWalletSnapshot>,
{
let mut wallets_by_id = BTreeMap::new();
for item in items {
wallets_by_id.insert(item.id.clone(), item);
}
Self::Owned(RwLock::new(wallets_by_id))
}
fn with_mut<R>(&self, f: impl FnOnce(&mut BTreeMap<String, StoredWalletSnapshot>) -> R) -> R {
match self {
Self::Owned(wallets_by_id) => {
let mut wallets = wallets_by_id.write().expect("settlement repo lock");
f(&mut wallets)
}
Self::Shared(repository) => repository.with_wallets_mut(f),
}
}
}
#[derive(Debug, Default)]
pub struct InMemorySettlementRepository {
wallets: InMemorySettlementWalletStore,
provider_monthly_used: RwLock<BTreeMap<String, f64>>,
}
impl InMemorySettlementRepository {
pub fn seed<I>(items: I) -> Self
where
I: IntoIterator<Item = StoredWalletSnapshot>,
{
Self {
wallets: InMemorySettlementWalletStore::seeded(items),
provider_monthly_used: RwLock::new(BTreeMap::new()),
}
}
pub fn from_wallet_repository(wallet_repository: Arc<InMemoryWalletRepository>) -> Self {
Self {
wallets: InMemorySettlementWalletStore::Shared(wallet_repository),
provider_monthly_used: RwLock::new(BTreeMap::new()),
}
}
}
#[async_trait]
impl SettlementWriteRepository for InMemorySettlementRepository {
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
input.validate()?;
if input.billing_status != "pending" {
return Ok(Some(StoredUsageSettlement {
request_id: input.request_id,
wallet_id: None,
billing_status: input.billing_status,
wallet_balance_before: None,
wallet_balance_after: None,
wallet_recharge_balance_before: None,
wallet_recharge_balance_after: None,
wallet_gift_balance_before: None,
wallet_gift_balance_after: None,
provider_monthly_used_usd: None,
finalized_at_unix_secs: input.finalized_at_unix_secs,
}));
}
let final_billing_status = if input.status == "completed" {
"settled"
} else {
"void"
};
let mut settlement = self.wallets.with_mut(|wallets| {
let wallet_id = input
.api_key_id
.as_deref()
.and_then(|api_key_id| {
wallets
.values()
.find(|wallet| wallet.api_key_id.as_deref() == Some(api_key_id))
.map(|wallet| wallet.id.clone())
})
.or_else(|| {
input.user_id.as_deref().and_then(|user_id| {
wallets
.values()
.find(|wallet| wallet.user_id.as_deref() == Some(user_id))
.map(|wallet| wallet.id.clone())
})
});
let wallet = wallet_id
.as_deref()
.and_then(|wallet_id| wallets.get_mut(wallet_id));
let mut settlement = StoredUsageSettlement {
request_id: input.request_id.clone(),
wallet_id: None,
billing_status: final_billing_status.to_string(),
wallet_balance_before: None,
wallet_balance_after: None,
wallet_recharge_balance_before: None,
wallet_recharge_balance_after: None,
wallet_gift_balance_before: None,
wallet_gift_balance_after: None,
provider_monthly_used_usd: None,
finalized_at_unix_secs: input.finalized_at_unix_secs,
};
if let Some(wallet) = wallet {
let before_recharge = wallet.balance;
let before_gift = wallet.gift_balance;
let before_total = before_recharge + before_gift;
settlement.wallet_id = Some(wallet.id.clone());
settlement.wallet_balance_before = Some(before_total);
settlement.wallet_recharge_balance_before = Some(before_recharge);
settlement.wallet_gift_balance_before = Some(before_gift);
if final_billing_status == "settled" {
if wallet.limit_mode.eq_ignore_ascii_case("unlimited") {
wallet.total_consumed += input.total_cost_usd;
} else {
let gift_deduction = before_gift.max(0.0).min(input.total_cost_usd);
let recharge_deduction = input.total_cost_usd - gift_deduction;
wallet.gift_balance = before_gift - gift_deduction;
wallet.balance = before_recharge - recharge_deduction;
wallet.total_consumed += input.total_cost_usd;
}
}
settlement.wallet_recharge_balance_after = Some(wallet.balance);
settlement.wallet_gift_balance_after = Some(wallet.gift_balance);
settlement.wallet_balance_after = Some(wallet.balance + wallet.gift_balance);
}
settlement
});
if final_billing_status == "settled" {
if let Some(provider_id) = input.provider_id {
let mut quotas = self
.provider_monthly_used
.write()
.expect("provider quota lock");
let value = quotas.entry(provider_id).or_insert(0.0);
*value += input.actual_total_cost_usd;
settlement.provider_monthly_used_usd = Some(*value);
}
}
Ok(Some(settlement))
}
}
#[cfg(test)]
mod tests {
use super::InMemorySettlementRepository;
use crate::repository::settlement::{SettlementWriteRepository, UsageSettlementInput};
use crate::repository::wallet::StoredWalletSnapshot;
fn sample_wallet() -> StoredWalletSnapshot {
StoredWalletSnapshot::new(
"wallet-1".to_string(),
Some("user-1".to_string()),
Some("key-1".to_string()),
10.0,
2.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
0.0,
0.0,
0.0,
0.0,
100,
)
.expect("wallet should build")
}
#[tokio::test]
async fn settles_usage_against_wallet_and_provider_quota() {
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
provider_id: Some("provider-1".to_string()),
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 1.5,
finalized_at_unix_secs: Some(200),
})
.await
.expect("settlement should succeed")
.expect("settlement should exist");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(9.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
}
}

View File

@@ -0,0 +1,9 @@
mod memory;
mod sql;
mod types;
pub use memory::InMemorySettlementRepository;
pub use sql::SqlxSettlementRepository;
pub use types::{
SettlementRepository, SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
};

View File

@@ -0,0 +1,279 @@
use async_trait::async_trait;
use sqlx::{PgPool, Row};
use super::types::{SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput};
use crate::postgres::PostgresTransactionRunner;
use crate::DataLayerError;
const FINALIZE_USAGE_BILLING_SQL: &str = r#"
UPDATE "usage"
SET
billing_status = $2,
finalized_at = COALESCE(finalized_at, to_timestamp($3))
WHERE request_id = $1
"#;
#[derive(Debug, Clone)]
pub struct SqlxSettlementRepository {
tx_runner: PostgresTransactionRunner,
}
impl SqlxSettlementRepository {
pub fn new(pool: PgPool) -> Self {
let tx_runner = PostgresTransactionRunner::new(pool);
Self { tx_runner }
}
}
#[async_trait]
impl SettlementWriteRepository for SqlxSettlementRepository {
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
input.validate()?;
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
let row = sqlx::query(
r#"
SELECT
request_id,
wallet_id,
billing_status,
CAST(wallet_balance_before AS DOUBLE PRECISION) AS wallet_balance_before,
CAST(wallet_balance_after AS DOUBLE PRECISION) AS wallet_balance_after,
CAST(wallet_recharge_balance_before AS DOUBLE PRECISION) AS wallet_recharge_balance_before,
CAST(wallet_recharge_balance_after AS DOUBLE PRECISION) AS wallet_recharge_balance_after,
CAST(wallet_gift_balance_before AS DOUBLE PRECISION) AS wallet_gift_balance_before,
CAST(wallet_gift_balance_after AS DOUBLE PRECISION) AS wallet_gift_balance_after,
provider_id,
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
FROM "usage"
WHERE request_id = $1
FOR UPDATE
"#,
)
.bind(&input.request_id)
.fetch_optional(&mut **tx)
.await?;
let Some(usage_row) = row else {
return Ok(None);
};
let current_billing_status: String = usage_row.try_get("billing_status")?;
if current_billing_status == "settled" || current_billing_status == "void" {
return Ok(Some(StoredUsageSettlement {
request_id: usage_row.try_get("request_id")?,
wallet_id: usage_row.try_get("wallet_id")?,
billing_status: current_billing_status,
wallet_balance_before: usage_row.try_get("wallet_balance_before")?,
wallet_balance_after: usage_row.try_get("wallet_balance_after")?,
wallet_recharge_balance_before: usage_row
.try_get("wallet_recharge_balance_before")?,
wallet_recharge_balance_after: usage_row
.try_get("wallet_recharge_balance_after")?,
wallet_gift_balance_before: usage_row
.try_get("wallet_gift_balance_before")?,
wallet_gift_balance_after: usage_row
.try_get("wallet_gift_balance_after")?,
provider_monthly_used_usd: None,
finalized_at_unix_secs: usage_row
.try_get::<Option<i64>, _>("finalized_at_unix_secs")?
.map(|value| value as u64),
}));
}
let final_billing_status = if input.status == "completed" {
"settled"
} else {
"void"
};
let finalized_at =
i64::try_from(input.finalized_at_unix_secs.unwrap_or_else(|| {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}))
.map_err(|_| {
DataLayerError::InvalidInput("finalized_at overflow".to_string())
})?;
let mut settlement = StoredUsageSettlement {
request_id: input.request_id.clone(),
wallet_id: None,
billing_status: final_billing_status.to_string(),
wallet_balance_before: None,
wallet_balance_after: None,
wallet_recharge_balance_before: None,
wallet_recharge_balance_after: None,
wallet_gift_balance_before: None,
wallet_gift_balance_after: None,
provider_monthly_used_usd: None,
finalized_at_unix_secs: Some(finalized_at as u64),
};
if final_billing_status == "settled" {
let wallet_row = if let Some(api_key_id) = input
.api_key_id
.as_deref()
.filter(|value| !value.is_empty())
{
sqlx::query(
r#"
SELECT
id,
CAST(balance AS DOUBLE PRECISION) AS balance,
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
limit_mode
FROM wallets
WHERE api_key_id = $1
FOR UPDATE
LIMIT 1
"#,
)
.bind(api_key_id)
.fetch_optional(&mut **tx)
.await?
} else {
None
};
let wallet_row = if wallet_row.is_some() {
wallet_row
} else if let Some(user_id) =
input.user_id.as_deref().filter(|value| !value.is_empty())
{
sqlx::query(
r#"
SELECT
id,
CAST(balance AS DOUBLE PRECISION) AS balance,
CAST(gift_balance AS DOUBLE PRECISION) AS gift_balance,
limit_mode
FROM wallets
WHERE user_id = $1
FOR UPDATE
LIMIT 1
"#,
)
.bind(user_id)
.fetch_optional(&mut **tx)
.await?
} else {
None
};
if let Some(wallet_row) = wallet_row {
let wallet_id: String = wallet_row.try_get("id")?;
let before_recharge: f64 = wallet_row.try_get("balance")?;
let before_gift: f64 = wallet_row.try_get("gift_balance")?;
let limit_mode: String = wallet_row.try_get("limit_mode")?;
let before_total = before_recharge + before_gift;
let mut after_recharge = before_recharge;
let mut after_gift = before_gift;
if !limit_mode.eq_ignore_ascii_case("unlimited") {
let gift_deduction = before_gift.max(0.0).min(input.total_cost_usd);
let recharge_deduction = input.total_cost_usd - gift_deduction;
after_gift = before_gift - gift_deduction;
after_recharge = before_recharge - recharge_deduction;
}
sqlx::query(
r#"
UPDATE wallets
SET
balance = $2,
gift_balance = $3,
total_consumed = CAST(total_consumed AS DOUBLE PRECISION) + $4,
updated_at = NOW()
WHERE id = $1
"#,
)
.bind(&wallet_id)
.bind(after_recharge)
.bind(after_gift)
.bind(input.total_cost_usd)
.execute(&mut **tx)
.await?;
settlement.wallet_id = Some(wallet_id.clone());
settlement.wallet_balance_before = Some(before_total);
settlement.wallet_balance_after = Some(after_recharge + after_gift);
settlement.wallet_recharge_balance_before = Some(before_recharge);
settlement.wallet_recharge_balance_after = Some(after_recharge);
settlement.wallet_gift_balance_before = Some(before_gift);
settlement.wallet_gift_balance_after = Some(after_gift);
sqlx::query(
r#"
UPDATE "usage"
SET
wallet_id = $2,
wallet_balance_before = $3,
wallet_balance_after = $4,
wallet_recharge_balance_before = $5,
wallet_recharge_balance_after = $6,
wallet_gift_balance_before = $7,
wallet_gift_balance_after = $8
WHERE request_id = $1
"#,
)
.bind(&input.request_id)
.bind(&wallet_id)
.bind(before_total)
.bind(after_recharge + after_gift)
.bind(before_recharge)
.bind(after_recharge)
.bind(before_gift)
.bind(after_gift)
.execute(&mut **tx)
.await?;
}
if let Some(provider_id) = input
.provider_id
.as_deref()
.filter(|value| !value.is_empty())
{
let quota_row = sqlx::query(
r#"
UPDATE providers
SET
monthly_used_usd = COALESCE(monthly_used_usd, 0) + $2,
updated_at = NOW()
WHERE id = $1
RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
"#,
)
.bind(provider_id)
.bind(input.actual_total_cost_usd)
.fetch_optional(&mut **tx)
.await?;
settlement.provider_monthly_used_usd =
quota_row.and_then(|row| row.try_get("monthly_used_usd").ok());
}
}
sqlx::query(FINALIZE_USAGE_BILLING_SQL)
.bind(&input.request_id)
.bind(final_billing_status)
.bind(finalized_at)
.execute(&mut **tx)
.await?;
Ok(Some(settlement))
})
})
.await
}
}
#[cfg(test)]
mod tests {
#[test]
fn finalize_usage_billing_sql_does_not_require_usage_updated_at_column() {
assert!(!super::FINALIZE_USAGE_BILLING_SQL.contains("updated_at"));
}
}

View File

@@ -0,0 +1,83 @@
use async_trait::async_trait;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UsageSettlementInput {
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub provider_id: Option<String>,
pub status: String,
pub billing_status: String,
pub total_cost_usd: f64,
pub actual_total_cost_usd: f64,
pub finalized_at_unix_secs: Option<u64>,
}
impl UsageSettlementInput {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
if self.request_id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"settlement request_id cannot be empty".to_string(),
));
}
if self.status.trim().is_empty() || self.billing_status.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"settlement status cannot be empty".to_string(),
));
}
if !self.total_cost_usd.is_finite() || !self.actual_total_cost_usd.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"settlement cost must be finite".to_string(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredUsageSettlement {
pub request_id: String,
pub wallet_id: Option<String>,
pub billing_status: String,
pub wallet_balance_before: Option<f64>,
pub wallet_balance_after: Option<f64>,
pub wallet_recharge_balance_before: Option<f64>,
pub wallet_recharge_balance_after: Option<f64>,
pub wallet_gift_balance_before: Option<f64>,
pub wallet_gift_balance_after: Option<f64>,
pub provider_monthly_used_usd: Option<f64>,
pub finalized_at_unix_secs: Option<u64>,
}
#[async_trait]
pub trait SettlementWriteRepository: Send + Sync {
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, crate::DataLayerError>;
}
pub trait SettlementRepository: SettlementWriteRepository + Send + Sync {}
impl<T> SettlementRepository for T where T: SettlementWriteRepository + Send + Sync {}
#[cfg(test)]
mod tests {
use super::UsageSettlementInput;
#[test]
fn rejects_invalid_settlement_input() {
let input = UsageSettlementInput {
request_id: "".to_string(),
user_id: None,
api_key_id: None,
provider_id: None,
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 0.1,
actual_total_cost_usd: 0.1,
finalized_at_unix_secs: None,
};
assert!(input.validate().is_err());
}
}

View File

@@ -0,0 +1,22 @@
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredSystemConfigEntry {
pub key: String,
pub value: serde_json::Value,
pub description: Option<String>,
pub updated_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminSecurityBlacklistEntry {
pub ip_address: String,
pub reason: String,
pub ttl_seconds: Option<i64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize, Default)]
pub struct AdminSystemStats {
pub total_users: u64,
pub active_users: u64,
pub total_api_keys: u64,
pub total_requests: u64,
}

View File

@@ -5,6 +5,6 @@ mod types;
pub use memory::InMemoryUserReadRepository;
pub use sql::SqlxUserReadRepository;
pub use types::{
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, UserExportListQuery,
UserExportSummary, UserReadRepository,
StoredUserAuthRecord, StoredUserExportRow, StoredUserPreferenceRecord, StoredUserSessionRecord,
StoredUserSummary, UserExportListQuery, UserExportSummary, UserReadRepository,
};

View File

@@ -213,6 +213,165 @@ impl StoredUserExportRow {
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredUserSessionRecord {
pub id: String,
pub user_id: String,
pub client_device_id: String,
pub device_label: Option<String>,
pub refresh_token_hash: String,
pub prev_refresh_token_hash: Option<String>,
pub rotated_at: Option<DateTime<Utc>>,
pub last_seen_at: Option<DateTime<Utc>>,
pub expires_at: Option<DateTime<Utc>>,
pub revoked_at: Option<DateTime<Utc>>,
pub revoke_reason: Option<String>,
pub ip_address: Option<String>,
pub user_agent: Option<String>,
pub created_at: Option<DateTime<Utc>>,
pub updated_at: Option<DateTime<Utc>>,
}
impl StoredUserSessionRecord {
pub const REFRESH_GRACE_SECONDS: i64 = 10;
pub const TOUCH_INTERVAL_SECONDS: i64 = 300;
#[allow(clippy::too_many_arguments)]
pub fn new(
id: String,
user_id: String,
client_device_id: String,
device_label: Option<String>,
refresh_token_hash: String,
prev_refresh_token_hash: Option<String>,
rotated_at: Option<DateTime<Utc>>,
last_seen_at: Option<DateTime<Utc>>,
expires_at: Option<DateTime<Utc>>,
revoked_at: Option<DateTime<Utc>>,
revoke_reason: Option<String>,
ip_address: Option<String>,
user_agent: Option<String>,
created_at: Option<DateTime<Utc>>,
updated_at: Option<DateTime<Utc>>,
) -> Result<Self, crate::DataLayerError> {
if id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"user_sessions.id is empty".to_string(),
));
}
if user_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"user_sessions.user_id is empty".to_string(),
));
}
if client_device_id.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"user_sessions.client_device_id is empty".to_string(),
));
}
if refresh_token_hash.trim().is_empty() {
return Err(crate::DataLayerError::UnexpectedValue(
"user_sessions.refresh_token_hash is empty".to_string(),
));
}
Ok(Self {
id,
user_id,
client_device_id,
device_label,
refresh_token_hash,
prev_refresh_token_hash,
rotated_at,
last_seen_at,
expires_at,
revoked_at,
revoke_reason,
ip_address,
user_agent,
created_at,
updated_at,
})
}
pub fn hash_refresh_token(token: &str) -> String {
use sha2::Digest;
let mut hasher = sha2::Sha256::new();
hasher.update(token.as_bytes());
format!("{:x}", hasher.finalize())
}
pub fn verify_refresh_token(&self, token: &str, now: DateTime<Utc>) -> (bool, bool) {
let token_hash = Self::hash_refresh_token(token);
if self.refresh_token_hash == token_hash {
return (true, false);
}
let Some(prev_hash) = self.prev_refresh_token_hash.as_ref() else {
return (false, false);
};
let Some(rotated_at) = self.rotated_at else {
return (false, false);
};
if prev_hash == &token_hash
&& now.signed_duration_since(rotated_at).num_seconds() <= Self::REFRESH_GRACE_SECONDS
{
return (true, true);
}
(false, false)
}
pub fn is_revoked(&self) -> bool {
self.revoked_at.is_some()
}
pub fn is_expired(&self, now: DateTime<Utc>) -> bool {
self.expires_at.is_none_or(|expires_at| expires_at <= now)
}
pub fn should_touch(&self, now: DateTime<Utc>) -> bool {
self.last_seen_at
.map(|last_seen_at| {
now.signed_duration_since(last_seen_at).num_seconds()
>= Self::TOUCH_INTERVAL_SECONDS
})
.unwrap_or(true)
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredUserPreferenceRecord {
pub user_id: String,
pub avatar_url: Option<String>,
pub bio: Option<String>,
pub default_provider_id: Option<String>,
pub default_provider_name: Option<String>,
pub theme: String,
pub language: String,
pub timezone: String,
pub email_notifications: bool,
pub usage_alerts: bool,
pub announcement_notifications: bool,
}
impl StoredUserPreferenceRecord {
pub fn default_for_user(user_id: impl Into<String>) -> Self {
Self {
user_id: user_id.into(),
avatar_url: None,
bio: None,
default_provider_id: None,
default_provider_name: None,
theme: "light".to_string(),
language: "zh-CN".to_string(),
timezone: "Asia/Shanghai".to_string(),
email_notifications: true,
usage_alerts: true,
announcement_notifications: true,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct UserExportListQuery {
pub skip: usize,
@@ -336,9 +495,13 @@ fn parse_string_list_array(
#[cfg(test)]
mod tests {
use chrono::{Duration, Utc};
use serde_json::Value;
use super::{StoredUserAuthRecord, StoredUserExportRow};
use super::{
StoredUserAuthRecord, StoredUserExportRow, StoredUserPreferenceRecord,
StoredUserSessionRecord,
};
#[test]
fn builds_user_export_row_with_allowed_lists() {
@@ -447,4 +610,49 @@ mod tests {
);
assert_eq!(row.allowed_models, Some(vec!["gpt-4.1".to_string()]));
}
#[test]
fn user_session_previous_refresh_token_has_grace_window() {
let now = Utc::now();
let session = StoredUserSessionRecord::new(
"session-1".to_string(),
"user-1".to_string(),
"device-1".to_string(),
None,
StoredUserSessionRecord::hash_refresh_token("current-token"),
Some(StoredUserSessionRecord::hash_refresh_token("prev-token")),
Some(now - Duration::seconds(StoredUserSessionRecord::REFRESH_GRACE_SECONDS - 1)),
None,
None,
None,
None,
None,
None,
None,
None,
)
.expect("session should build");
assert_eq!(
session.verify_refresh_token("prev-token", now),
(true, true)
);
assert_eq!(
session.verify_refresh_token("current-token", now),
(true, false)
);
}
#[test]
fn user_preference_defaults_match_gateway_expectations() {
let record = StoredUserPreferenceRecord::default_for_user("user-1");
assert_eq!(record.user_id, "user-1");
assert_eq!(record.theme, "light");
assert_eq!(record.language, "zh-CN");
assert_eq!(record.timezone, "Asia/Shanghai");
assert!(record.email_notifications);
assert!(record.usage_alerts);
assert!(record.announcement_notifications);
}
}

View File

@@ -4,7 +4,17 @@ use std::sync::RwLock;
use async_trait::async_trait;
use super::types::{
StoredUsageSettlement, StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey,
AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminWalletLedgerQuery,
AdminWalletListQuery, AdminWalletRefundRequestListQuery, CompleteAdminWalletRefundInput,
CreateManualWalletRechargeInput, CreateWalletRechargeOrderInput,
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, FailAdminWalletRefundInput,
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage,
StoredAdminWalletLedgerPage, StoredAdminWalletListItem, StoredAdminWalletListPage,
StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage,
StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome,
WalletReadRepository, WalletWriteRepository,
};
use crate::DataLayerError;
@@ -12,7 +22,6 @@ use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryWalletRepository {
wallets_by_id: RwLock<BTreeMap<String, StoredWalletSnapshot>>,
provider_monthly_used: RwLock<BTreeMap<String, f64>>,
}
impl InMemoryWalletRepository {
@@ -26,9 +35,16 @@ impl InMemoryWalletRepository {
}
Self {
wallets_by_id: RwLock::new(wallets_by_id),
provider_monthly_used: RwLock::new(BTreeMap::new()),
}
}
pub(crate) fn with_wallets_mut<R>(
&self,
f: impl FnOnce(&mut BTreeMap<String, StoredWalletSnapshot>) -> R,
) -> R {
let mut wallets = self.wallets_by_id.write().expect("wallet repo lock");
f(&mut wallets)
}
}
#[async_trait]
@@ -96,112 +112,329 @@ impl WalletReadRepository for InMemoryWalletRepository {
.cloned()
.collect())
}
async fn list_admin_wallets(
&self,
query: &AdminWalletListQuery,
) -> Result<StoredAdminWalletListPage, DataLayerError> {
let wallets = self.wallets_by_id.read().expect("wallet repo lock");
let mut items = wallets
.values()
.filter(|wallet| {
query
.status
.as_deref()
.is_none_or(|expected| wallet.status == expected)
})
.filter(|wallet| match query.owner_type.as_deref() {
Some("user") => wallet.user_id.is_some(),
Some("api_key") => wallet.api_key_id.is_some(),
_ => true,
})
.map(|wallet| StoredAdminWalletListItem {
id: wallet.id.clone(),
user_id: wallet.user_id.clone(),
api_key_id: wallet.api_key_id.clone(),
balance: wallet.balance,
gift_balance: wallet.gift_balance,
limit_mode: wallet.limit_mode.clone(),
currency: wallet.currency.clone(),
status: wallet.status.clone(),
total_recharged: wallet.total_recharged,
total_consumed: wallet.total_consumed,
total_refunded: wallet.total_refunded,
total_adjusted: wallet.total_adjusted,
user_name: None,
api_key_name: None,
created_at_unix_secs: None,
updated_at_unix_secs: Some(wallet.updated_at_unix_secs),
})
.collect::<Vec<_>>();
items.sort_by(|left, right| {
right
.updated_at_unix_secs
.cmp(&left.updated_at_unix_secs)
.then_with(|| right.id.cmp(&left.id))
});
let total = items.len() as u64;
let items = items
.into_iter()
.skip(query.offset)
.take(query.limit)
.collect::<Vec<_>>();
Ok(StoredAdminWalletListPage { items, total })
}
async fn list_admin_wallet_ledger(
&self,
_query: &AdminWalletLedgerQuery,
) -> Result<StoredAdminWalletLedgerPage, DataLayerError> {
Ok(StoredAdminWalletLedgerPage::default())
}
async fn list_admin_wallet_refund_requests(
&self,
_query: &AdminWalletRefundRequestListQuery,
) -> Result<StoredAdminWalletRefundRequestPage, DataLayerError> {
Ok(StoredAdminWalletRefundRequestPage::default())
}
async fn list_admin_wallet_transactions(
&self,
_wallet_id: &str,
_limit: usize,
_offset: usize,
) -> Result<StoredAdminWalletTransactionPage, DataLayerError> {
Ok(StoredAdminWalletTransactionPage::default())
}
async fn find_wallet_today_usage(
&self,
_wallet_id: &str,
_billing_timezone: &str,
) -> Result<Option<StoredWalletDailyUsageLedger>, DataLayerError> {
Ok(None)
}
async fn list_wallet_daily_usage_history(
&self,
_wallet_id: &str,
_billing_timezone: &str,
_limit: usize,
) -> Result<StoredWalletDailyUsageLedgerPage, DataLayerError> {
Ok(StoredWalletDailyUsageLedgerPage::default())
}
async fn list_admin_wallet_refunds(
&self,
_wallet_id: &str,
_limit: usize,
_offset: usize,
) -> Result<StoredAdminWalletRefundPage, DataLayerError> {
Ok(StoredAdminWalletRefundPage::default())
}
async fn list_admin_payment_orders(
&self,
_query: &AdminPaymentOrderListQuery,
) -> Result<StoredAdminPaymentOrderPage, DataLayerError> {
Ok(StoredAdminPaymentOrderPage::default())
}
async fn find_admin_payment_order(
&self,
_order_id: &str,
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
Ok(None)
}
async fn list_wallet_payment_orders_by_user_id(
&self,
_user_id: &str,
_limit: usize,
_offset: usize,
) -> Result<StoredAdminPaymentOrderPage, DataLayerError> {
Ok(StoredAdminPaymentOrderPage::default())
}
async fn find_wallet_payment_order_by_user_id(
&self,
_user_id: &str,
_order_id: &str,
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
Ok(None)
}
async fn find_wallet_refund(
&self,
_wallet_id: &str,
_refund_id: &str,
) -> Result<Option<super::types::StoredAdminWalletRefund>, DataLayerError> {
Ok(None)
}
async fn list_admin_payment_callbacks(
&self,
_payment_method: Option<&str>,
_limit: usize,
_offset: usize,
) -> Result<StoredAdminPaymentCallbackPage, DataLayerError> {
Ok(StoredAdminPaymentCallbackPage::default())
}
}
#[async_trait]
impl WalletWriteRepository for InMemoryWalletRepository {
async fn settle_usage(
async fn create_wallet_recharge_order(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
input.validate()?;
if input.billing_status != "pending" {
return Ok(Some(StoredUsageSettlement {
request_id: input.request_id,
wallet_id: None,
billing_status: input.billing_status,
wallet_balance_before: None,
wallet_balance_after: None,
wallet_recharge_balance_before: None,
wallet_recharge_balance_after: None,
wallet_gift_balance_before: None,
wallet_gift_balance_after: None,
provider_monthly_used_usd: None,
finalized_at_unix_secs: input.finalized_at_unix_secs,
}));
}
input: CreateWalletRechargeOrderInput,
) -> Result<CreateWalletRechargeOrderOutcome, DataLayerError> {
let mut wallets = self.wallets_by_id.write().expect("wallet repo lock");
let wallet_id = input
.api_key_id
.as_deref()
.and_then(|api_key_id| {
wallets
.values()
.find(|wallet| wallet.api_key_id.as_deref() == Some(api_key_id))
.map(|wallet| wallet.id.clone())
})
.or_else(|| {
input.user_id.as_deref().and_then(|user_id| {
wallets
.values()
.find(|wallet| wallet.user_id.as_deref() == Some(user_id))
.map(|wallet| wallet.id.clone())
})
});
let wallet = wallet_id
.as_deref()
.and_then(|wallet_id| wallets.get_mut(wallet_id));
let final_billing_status = if input.status == "completed" {
"settled"
} else {
"void"
};
let mut settlement = StoredUsageSettlement {
request_id: input.request_id,
wallet_id: None,
billing_status: final_billing_status.to_string(),
wallet_balance_before: None,
wallet_balance_after: None,
wallet_recharge_balance_before: None,
wallet_recharge_balance_after: None,
wallet_gift_balance_before: None,
wallet_gift_balance_after: None,
provider_monthly_used_usd: None,
finalized_at_unix_secs: input.finalized_at_unix_secs,
};
if let Some(wallet) = wallet {
let before_recharge = wallet.balance;
let before_gift = wallet.gift_balance;
let before_total = before_recharge + before_gift;
settlement.wallet_id = Some(wallet.id.clone());
settlement.wallet_balance_before = Some(before_total);
settlement.wallet_recharge_balance_before = Some(before_recharge);
settlement.wallet_gift_balance_before = Some(before_gift);
if final_billing_status == "settled" {
if wallet.limit_mode.eq_ignore_ascii_case("unlimited") {
wallet.total_consumed += input.total_cost_usd;
} else {
let gift_deduction = before_gift.max(0.0).min(input.total_cost_usd);
let recharge_deduction = input.total_cost_usd - gift_deduction;
wallet.gift_balance = before_gift - gift_deduction;
wallet.balance = before_recharge - recharge_deduction;
wallet.total_consumed += input.total_cost_usd;
}
}
settlement.wallet_recharge_balance_after = Some(wallet.balance);
settlement.wallet_gift_balance_after = Some(wallet.gift_balance);
settlement.wallet_balance_after = Some(wallet.balance + wallet.gift_balance);
let wallet = wallets
.values_mut()
.find(|wallet| wallet.user_id.as_deref() == Some(input.user_id.as_str()));
if wallet
.as_ref()
.is_some_and(|wallet| wallet.status != "active")
{
return Ok(CreateWalletRechargeOrderOutcome::WalletInactive);
}
let wallet_id = wallet
.map(|wallet| wallet.id.clone())
.or(input.preferred_wallet_id)
.unwrap_or_else(|| "wallet-memory".to_string());
Ok(CreateWalletRechargeOrderOutcome::Created(
StoredAdminPaymentOrder {
id: "payment-order-memory".to_string(),
order_no: input.order_no,
wallet_id,
user_id: Some(input.user_id),
amount_usd: input.amount_usd,
pay_amount: input.pay_amount,
pay_currency: input.pay_currency,
exchange_rate: input.exchange_rate,
refunded_amount_usd: 0.0,
refundable_amount_usd: 0.0,
payment_method: input.payment_method,
gateway_order_id: Some(input.gateway_order_id),
gateway_response: Some(input.gateway_response),
status: "pending".to_string(),
created_at_unix_secs: 0,
paid_at_unix_secs: None,
credited_at_unix_secs: None,
expires_at_unix_secs: Some(input.expires_at_unix_secs),
},
))
}
if final_billing_status == "settled" {
if let Some(provider_id) = input.provider_id {
let mut quotas = self
.provider_monthly_used
.write()
.expect("provider quota lock");
let value = quotas.entry(provider_id).or_insert(0.0);
*value += input.actual_total_cost_usd;
settlement.provider_monthly_used_usd = Some(*value);
}
async fn create_wallet_refund_request(
&self,
input: CreateWalletRefundRequestInput,
) -> Result<CreateWalletRefundRequestOutcome, DataLayerError> {
let wallets = self.wallets_by_id.read().expect("wallet repo lock");
let Some(wallet) = wallets.get(&input.wallet_id) else {
return Ok(CreateWalletRefundRequestOutcome::WalletMissing);
};
if input.amount_usd > wallet.balance {
return Ok(CreateWalletRefundRequestOutcome::RefundAmountExceedsAvailableBalance);
}
Ok(CreateWalletRefundRequestOutcome::Created(
super::types::StoredAdminWalletRefund {
id: "refund-memory".to_string(),
refund_no: input.refund_no,
wallet_id: input.wallet_id,
user_id: Some(input.user_id),
payment_order_id: input.payment_order_id,
source_type: input
.source_type
.unwrap_or_else(|| "wallet_balance".to_string()),
source_id: input.source_id,
refund_mode: input
.refund_mode
.unwrap_or_else(|| "offline_payout".to_string()),
amount_usd: input.amount_usd,
status: "pending_approval".to_string(),
reason: input.reason,
failure_reason: None,
gateway_refund_id: None,
payout_method: None,
payout_reference: None,
payout_proof: None,
requested_by: None,
approved_by: None,
processed_by: None,
created_at_unix_secs: 0,
updated_at_unix_secs: 0,
processed_at_unix_secs: None,
completed_at_unix_secs: None,
},
))
}
Ok(Some(settlement))
async fn process_payment_callback(
&self,
_input: ProcessPaymentCallbackInput,
) -> Result<ProcessPaymentCallbackOutcome, DataLayerError> {
Ok(ProcessPaymentCallbackOutcome::Failed {
duplicate: false,
error: "payment callback is not supported in memory wallet repository".to_string(),
})
}
async fn adjust_wallet_balance(
&self,
_input: AdjustWalletBalanceInput,
) -> Result<
Option<(
StoredWalletSnapshot,
super::types::StoredAdminWalletTransaction,
)>,
DataLayerError,
> {
Ok(None)
}
async fn create_manual_wallet_recharge(
&self,
_input: CreateManualWalletRechargeInput,
) -> Result<Option<(StoredWalletSnapshot, StoredAdminPaymentOrder)>, DataLayerError> {
Ok(None)
}
async fn process_admin_wallet_refund(
&self,
_input: ProcessAdminWalletRefundInput,
) -> Result<
WalletMutationOutcome<(
StoredWalletSnapshot,
super::types::StoredAdminWalletRefund,
super::types::StoredAdminWalletTransaction,
)>,
DataLayerError,
> {
Ok(WalletMutationOutcome::NotFound)
}
async fn complete_admin_wallet_refund(
&self,
_input: CompleteAdminWalletRefundInput,
) -> Result<WalletMutationOutcome<super::types::StoredAdminWalletRefund>, DataLayerError> {
Ok(WalletMutationOutcome::NotFound)
}
async fn fail_admin_wallet_refund(
&self,
_input: FailAdminWalletRefundInput,
) -> Result<
WalletMutationOutcome<(
StoredWalletSnapshot,
super::types::StoredAdminWalletRefund,
Option<super::types::StoredAdminWalletTransaction>,
)>,
DataLayerError,
> {
Ok(WalletMutationOutcome::NotFound)
}
async fn expire_admin_payment_order(
&self,
_order_id: &str,
) -> Result<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>, DataLayerError> {
Ok(WalletMutationOutcome::NotFound)
}
async fn fail_admin_payment_order(
&self,
_order_id: &str,
) -> Result<WalletMutationOutcome<StoredAdminPaymentOrder>, DataLayerError> {
Ok(WalletMutationOutcome::NotFound)
}
async fn credit_admin_payment_order(
&self,
_input: CreditAdminPaymentOrderInput,
) -> Result<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>, DataLayerError> {
Ok(WalletMutationOutcome::NotFound)
}
}
@@ -209,8 +442,7 @@ impl WalletWriteRepository for InMemoryWalletRepository {
mod tests {
use super::InMemoryWalletRepository;
use crate::repository::wallet::{
StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey, WalletReadRepository,
WalletWriteRepository,
AdminWalletListQuery, StoredWalletSnapshot, WalletLookupKey, WalletReadRepository,
};
fn sample_wallet() -> StoredWalletSnapshot {
@@ -244,27 +476,73 @@ mod tests {
}
#[tokio::test]
async fn settles_usage_against_wallet_and_provider_quota() {
let repository = InMemoryWalletRepository::seed(vec![sample_wallet()]);
let settlement = repository
.settle_usage(UsageSettlementInput {
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
provider_id: Some("provider-1".to_string()),
status: "completed".to_string(),
billing_status: "pending".to_string(),
total_cost_usd: 3.0,
actual_total_cost_usd: 1.5,
finalized_at_unix_secs: Some(200),
async fn lists_admin_wallets_with_filters_and_pagination() {
let repository = InMemoryWalletRepository::seed(vec![
sample_wallet(),
StoredWalletSnapshot::new(
"wallet-2".to_string(),
Some("user-2".to_string()),
None,
3.0,
1.0,
"finite".to_string(),
"USD".to_string(),
"inactive".to_string(),
0.0,
0.0,
0.0,
0.0,
90,
)
.expect("wallet should build"),
StoredWalletSnapshot::new(
"wallet-3".to_string(),
None,
Some("key-3".to_string()),
5.0,
0.0,
"unlimited".to_string(),
"USD".to_string(),
"active".to_string(),
0.0,
0.0,
0.0,
0.0,
110,
)
.expect("wallet should build"),
]);
let page = repository
.list_admin_wallets(&AdminWalletListQuery {
status: Some("active".to_string()),
owner_type: Some("api_key".to_string()),
limit: 1,
offset: 0,
})
.await
.expect("settlement should succeed")
.expect("settlement should exist");
.expect("list should succeed");
assert_eq!(settlement.billing_status, "settled");
assert_eq!(settlement.wallet_balance_before, Some(12.0));
assert_eq!(settlement.wallet_balance_after, Some(9.0));
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
assert_eq!(page.total, 2);
assert_eq!(page.items.len(), 1);
assert_eq!(page.items[0].id, "wallet-3");
assert_eq!(page.items[0].updated_at_unix_secs, Some(110));
}
#[tokio::test]
async fn daily_usage_queries_default_to_empty_in_memory() {
let repository = InMemoryWalletRepository::seed(vec![sample_wallet()]);
let today = repository
.find_wallet_today_usage("wallet-1", "Asia/Shanghai")
.await
.expect("lookup should succeed");
let history = repository
.list_wallet_daily_usage_history("wallet-1", "Asia/Shanghai", 20)
.await
.expect("history should succeed");
assert!(today.is_none());
assert_eq!(history.total, 0);
assert!(history.items.is_empty());
}
}

View File

@@ -5,6 +5,19 @@ mod types;
pub use memory::InMemoryWalletRepository;
pub use sql::SqlxWalletRepository;
pub use types::{
StoredUsageSettlement, StoredWalletSnapshot, UsageSettlementInput, WalletLookupKey,
AdjustWalletBalanceInput, AdminPaymentCallbackRecord, AdminPaymentOrderListQuery,
AdminWalletLedgerQuery, AdminWalletListQuery, AdminWalletPaymentOrderRecord,
AdminWalletRefundRecord, AdminWalletRefundRequestListQuery, AdminWalletTransactionRecord,
CompleteAdminWalletRefundInput, CreateManualWalletRechargeInput,
CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
FailAdminWalletRefundInput, ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput,
ProcessPaymentCallbackOutcome, StoredAdminPaymentCallback, StoredAdminPaymentCallbackPage,
StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, StoredAdminWalletLedgerItem,
StoredAdminWalletLedgerPage, StoredAdminWalletListItem, StoredAdminWalletListPage,
StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestItem,
StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome,
WalletReadRepository, WalletRepository, WalletWriteRepository,
};

File diff suppressed because it is too large Load Diff

View File

@@ -94,53 +94,500 @@ impl StoredWalletSnapshot {
}
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct AdminWalletListQuery {
pub status: Option<String>,
pub owner_type: Option<String>,
pub limit: usize,
pub offset: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UsageSettlementInput {
pub request_id: String,
pub struct StoredAdminWalletListItem {
pub id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub provider_id: Option<String>,
pub balance: f64,
pub gift_balance: f64,
pub limit_mode: String,
pub currency: String,
pub status: String,
pub billing_status: String,
pub total_cost_usd: f64,
pub actual_total_cost_usd: f64,
pub finalized_at_unix_secs: Option<u64>,
pub total_recharged: f64,
pub total_consumed: f64,
pub total_refunded: f64,
pub total_adjusted: f64,
pub user_name: Option<String>,
pub api_key_name: Option<String>,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
}
impl UsageSettlementInput {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
if self.request_id.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"wallet settlement request_id cannot be empty".to_string(),
));
}
if self.status.trim().is_empty() || self.billing_status.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(
"wallet settlement status cannot be empty".to_string(),
));
}
if !self.total_cost_usd.is_finite() || !self.actual_total_cost_usd.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"wallet settlement cost must be finite".to_string(),
));
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletListPage {
pub items: Vec<StoredAdminWalletListItem>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct AdminWalletLedgerQuery {
pub category: Option<String>,
pub reason_code: Option<String>,
pub owner_type: Option<String>,
pub limit: usize,
pub offset: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredUsageSettlement {
pub request_id: String,
pub wallet_id: Option<String>,
pub billing_status: String,
pub wallet_balance_before: Option<f64>,
pub wallet_balance_after: Option<f64>,
pub wallet_recharge_balance_before: Option<f64>,
pub wallet_recharge_balance_after: Option<f64>,
pub wallet_gift_balance_before: Option<f64>,
pub wallet_gift_balance_after: Option<f64>,
pub provider_monthly_used_usd: Option<f64>,
pub finalized_at_unix_secs: Option<u64>,
pub struct StoredAdminWalletLedgerItem {
pub id: String,
pub wallet_id: String,
pub category: String,
pub reason_code: String,
pub amount: f64,
pub balance_before: f64,
pub balance_after: f64,
pub recharge_balance_before: f64,
pub recharge_balance_after: f64,
pub gift_balance_before: f64,
pub gift_balance_after: f64,
pub link_type: Option<String>,
pub link_id: Option<String>,
pub operator_id: Option<String>,
pub operator_name: Option<String>,
pub operator_email: Option<String>,
pub description: Option<String>,
pub wallet_user_id: Option<String>,
pub wallet_user_name: Option<String>,
pub wallet_api_key_id: Option<String>,
pub api_key_name: Option<String>,
pub wallet_status: String,
pub created_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletLedgerPage {
pub items: Vec<StoredAdminWalletLedgerItem>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct AdminWalletRefundRequestListQuery {
pub status: Option<String>,
pub limit: usize,
pub offset: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletRefundRequestItem {
pub id: String,
pub refund_no: String,
pub wallet_id: String,
pub user_id: Option<String>,
pub payment_order_id: Option<String>,
pub source_type: String,
pub source_id: Option<String>,
pub refund_mode: String,
pub amount_usd: f64,
pub status: String,
pub reason: Option<String>,
pub failure_reason: Option<String>,
pub gateway_refund_id: Option<String>,
pub payout_method: Option<String>,
pub payout_reference: Option<String>,
pub payout_proof: Option<serde_json::Value>,
pub requested_by: Option<String>,
pub approved_by: Option<String>,
pub processed_by: Option<String>,
pub wallet_user_id: Option<String>,
pub wallet_user_name: Option<String>,
pub wallet_api_key_id: Option<String>,
pub api_key_name: Option<String>,
pub wallet_status: String,
pub created_at_unix_secs: Option<u64>,
pub updated_at_unix_secs: Option<u64>,
pub processed_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletRefundRequestPage {
pub items: Vec<StoredAdminWalletRefundRequestItem>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletTransaction {
pub id: String,
pub wallet_id: String,
pub category: String,
pub reason_code: String,
pub amount: f64,
pub balance_before: f64,
pub balance_after: f64,
pub recharge_balance_before: f64,
pub recharge_balance_after: f64,
pub gift_balance_before: f64,
pub gift_balance_after: f64,
pub link_type: Option<String>,
pub link_id: Option<String>,
pub operator_id: Option<String>,
pub operator_name: Option<String>,
pub operator_email: Option<String>,
pub description: Option<String>,
pub created_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminWalletTransactionRecord {
pub id: String,
pub wallet_id: String,
pub category: String,
pub reason_code: String,
pub amount: f64,
pub balance_before: f64,
pub balance_after: f64,
pub recharge_balance_before: f64,
pub recharge_balance_after: f64,
pub gift_balance_before: f64,
pub gift_balance_after: f64,
pub link_type: Option<String>,
pub link_id: Option<String>,
pub operator_id: Option<String>,
pub description: Option<String>,
pub created_at_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletTransactionPage {
pub items: Vec<StoredAdminWalletTransaction>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredWalletDailyUsageLedger {
pub id: Option<String>,
pub billing_date: String,
pub billing_timezone: String,
pub total_cost_usd: f64,
pub total_requests: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub cache_creation_tokens: u64,
pub cache_read_tokens: u64,
pub first_finalized_at_unix_secs: Option<u64>,
pub last_finalized_at_unix_secs: Option<u64>,
pub aggregated_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredWalletDailyUsageLedgerPage {
pub items: Vec<StoredWalletDailyUsageLedger>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletRefund {
pub id: String,
pub refund_no: String,
pub wallet_id: String,
pub user_id: Option<String>,
pub payment_order_id: Option<String>,
pub source_type: String,
pub source_id: Option<String>,
pub refund_mode: String,
pub amount_usd: f64,
pub status: String,
pub reason: Option<String>,
pub failure_reason: Option<String>,
pub gateway_refund_id: Option<String>,
pub payout_method: Option<String>,
pub payout_reference: Option<String>,
pub payout_proof: Option<serde_json::Value>,
pub requested_by: Option<String>,
pub approved_by: Option<String>,
pub processed_by: Option<String>,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
pub processed_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminWalletRefundRecord {
pub id: String,
pub refund_no: String,
pub wallet_id: String,
pub user_id: Option<String>,
pub payment_order_id: Option<String>,
pub source_type: String,
pub source_id: Option<String>,
pub refund_mode: String,
pub amount_usd: f64,
pub status: String,
pub reason: Option<String>,
pub failure_reason: Option<String>,
pub gateway_refund_id: Option<String>,
pub payout_method: Option<String>,
pub payout_reference: Option<String>,
pub payout_proof: Option<serde_json::Value>,
pub requested_by: Option<String>,
pub approved_by: Option<String>,
pub processed_by: Option<String>,
pub created_at_unix_secs: u64,
pub updated_at_unix_secs: u64,
pub processed_at_unix_secs: Option<u64>,
pub completed_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminWalletRefundPage {
pub items: Vec<StoredAdminWalletRefund>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct AdminPaymentOrderListQuery {
pub status: Option<String>,
pub payment_method: Option<String>,
pub limit: usize,
pub offset: usize,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminPaymentOrder {
pub id: String,
pub order_no: String,
pub wallet_id: String,
pub user_id: Option<String>,
pub amount_usd: f64,
pub pay_amount: Option<f64>,
pub pay_currency: Option<String>,
pub exchange_rate: Option<f64>,
pub refunded_amount_usd: f64,
pub refundable_amount_usd: f64,
pub payment_method: String,
pub gateway_order_id: Option<String>,
pub gateway_response: Option<serde_json::Value>,
pub status: String,
pub created_at_unix_secs: u64,
pub paid_at_unix_secs: Option<u64>,
pub credited_at_unix_secs: Option<u64>,
pub expires_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminWalletPaymentOrderRecord {
pub id: String,
pub order_no: String,
pub wallet_id: String,
pub user_id: Option<String>,
pub amount_usd: f64,
pub pay_amount: Option<f64>,
pub pay_currency: Option<String>,
pub exchange_rate: Option<f64>,
pub refunded_amount_usd: f64,
pub refundable_amount_usd: f64,
pub payment_method: String,
pub gateway_order_id: Option<String>,
pub status: String,
pub gateway_response: Option<serde_json::Value>,
pub created_at_unix_secs: u64,
pub paid_at_unix_secs: Option<u64>,
pub credited_at_unix_secs: Option<u64>,
pub expires_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminPaymentOrderPage {
pub items: Vec<StoredAdminPaymentOrder>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminPaymentCallback {
pub id: String,
pub payment_order_id: Option<String>,
pub payment_method: String,
pub callback_key: String,
pub order_no: Option<String>,
pub gateway_order_id: Option<String>,
pub payload_hash: Option<String>,
pub signature_valid: bool,
pub status: String,
pub payload: Option<serde_json::Value>,
pub error_message: Option<String>,
pub created_at_unix_secs: u64,
pub processed_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminPaymentCallbackRecord {
pub id: String,
pub payment_order_id: Option<String>,
pub payment_method: String,
pub callback_key: String,
pub order_no: Option<String>,
pub gateway_order_id: Option<String>,
pub payload_hash: Option<String>,
pub signature_valid: bool,
pub status: String,
pub payload: Option<serde_json::Value>,
pub error_message: Option<String>,
pub created_at_unix_secs: u64,
pub processed_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminPaymentCallbackPage {
pub items: Vec<StoredAdminPaymentCallback>,
pub total: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CreateWalletRechargeOrderInput {
pub preferred_wallet_id: Option<String>,
pub user_id: String,
pub amount_usd: f64,
pub pay_amount: Option<f64>,
pub pay_currency: Option<String>,
pub exchange_rate: Option<f64>,
pub payment_method: String,
pub gateway_order_id: String,
pub gateway_response: serde_json::Value,
pub order_no: String,
pub expires_at_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum CreateWalletRechargeOrderOutcome {
Created(StoredAdminPaymentOrder),
WalletInactive,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CreateWalletRefundRequestInput {
pub wallet_id: String,
pub user_id: String,
pub amount_usd: f64,
pub payment_order_id: Option<String>,
pub source_type: Option<String>,
pub source_id: Option<String>,
pub refund_mode: Option<String>,
pub reason: Option<String>,
pub idempotency_key: Option<String>,
pub refund_no: String,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum CreateWalletRefundRequestOutcome {
Created(StoredAdminWalletRefund),
Duplicate(StoredAdminWalletRefund),
WalletMissing,
RefundAmountExceedsAvailableBalance,
PaymentOrderNotFound,
PaymentOrderNotRefundable,
RefundAmountExceedsAvailableOrderAmount,
DuplicateRejected,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProcessPaymentCallbackInput {
pub payment_method: String,
pub callback_key: String,
pub order_no: Option<String>,
pub gateway_order_id: Option<String>,
pub amount_usd: f64,
pub pay_amount: Option<f64>,
pub pay_currency: Option<String>,
pub exchange_rate: Option<f64>,
pub payload_hash: String,
pub payload: serde_json::Value,
pub signature_valid: bool,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum ProcessPaymentCallbackOutcome {
DuplicateProcessed {
order_id: Option<String>,
},
Failed {
duplicate: bool,
error: String,
},
AlreadyCredited {
duplicate: bool,
order_id: String,
order_no: String,
wallet_id: String,
},
Applied {
duplicate: bool,
order_id: String,
order_no: String,
wallet_id: String,
order: StoredAdminPaymentOrder,
},
}
#[derive(Debug, Clone, PartialEq)]
pub enum WalletMutationOutcome<T> {
Applied(T),
NotFound,
Invalid(String),
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdjustWalletBalanceInput {
pub wallet_id: String,
pub amount_usd: f64,
pub balance_type: String,
pub operator_id: Option<String>,
pub description: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CreateManualWalletRechargeInput {
pub wallet_id: String,
pub amount_usd: f64,
pub payment_method: String,
pub operator_id: Option<String>,
pub description: Option<String>,
pub order_no: String,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProcessAdminWalletRefundInput {
pub wallet_id: String,
pub refund_id: String,
pub operator_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CompleteAdminWalletRefundInput {
pub wallet_id: String,
pub refund_id: String,
pub gateway_refund_id: Option<String>,
pub payout_reference: Option<String>,
pub payout_proof: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct FailAdminWalletRefundInput {
pub wallet_id: String,
pub refund_id: String,
pub reason: String,
pub operator_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CreditAdminPaymentOrderInput {
pub order_id: String,
pub gateway_order_id: Option<String>,
pub pay_amount: Option<f64>,
pub pay_currency: Option<String>,
pub exchange_rate: Option<f64>,
pub gateway_response_patch: Option<serde_json::Value>,
pub operator_id: Option<String>,
}
#[async_trait]
@@ -159,14 +606,156 @@ pub trait WalletReadRepository: Send + Sync {
&self,
api_key_ids: &[String],
) -> Result<Vec<StoredWalletSnapshot>, crate::DataLayerError>;
async fn list_admin_wallets(
&self,
query: &AdminWalletListQuery,
) -> Result<StoredAdminWalletListPage, crate::DataLayerError>;
async fn list_admin_wallet_ledger(
&self,
query: &AdminWalletLedgerQuery,
) -> Result<StoredAdminWalletLedgerPage, crate::DataLayerError>;
async fn list_admin_wallet_refund_requests(
&self,
query: &AdminWalletRefundRequestListQuery,
) -> Result<StoredAdminWalletRefundRequestPage, crate::DataLayerError>;
async fn list_admin_wallet_transactions(
&self,
wallet_id: &str,
limit: usize,
offset: usize,
) -> Result<StoredAdminWalletTransactionPage, crate::DataLayerError>;
async fn find_wallet_today_usage(
&self,
wallet_id: &str,
billing_timezone: &str,
) -> Result<Option<StoredWalletDailyUsageLedger>, crate::DataLayerError>;
async fn list_wallet_daily_usage_history(
&self,
wallet_id: &str,
billing_timezone: &str,
limit: usize,
) -> Result<StoredWalletDailyUsageLedgerPage, crate::DataLayerError>;
async fn list_admin_wallet_refunds(
&self,
wallet_id: &str,
limit: usize,
offset: usize,
) -> Result<StoredAdminWalletRefundPage, crate::DataLayerError>;
async fn list_admin_payment_orders(
&self,
query: &AdminPaymentOrderListQuery,
) -> Result<StoredAdminPaymentOrderPage, crate::DataLayerError>;
async fn find_admin_payment_order(
&self,
order_id: &str,
) -> Result<Option<StoredAdminPaymentOrder>, crate::DataLayerError>;
async fn list_wallet_payment_orders_by_user_id(
&self,
user_id: &str,
limit: usize,
offset: usize,
) -> Result<StoredAdminPaymentOrderPage, crate::DataLayerError>;
async fn find_wallet_payment_order_by_user_id(
&self,
user_id: &str,
order_id: &str,
) -> Result<Option<StoredAdminPaymentOrder>, crate::DataLayerError>;
async fn find_wallet_refund(
&self,
wallet_id: &str,
refund_id: &str,
) -> Result<Option<StoredAdminWalletRefund>, crate::DataLayerError>;
async fn list_admin_payment_callbacks(
&self,
payment_method: Option<&str>,
limit: usize,
offset: usize,
) -> Result<StoredAdminPaymentCallbackPage, crate::DataLayerError>;
}
#[async_trait]
pub trait WalletWriteRepository: Send + Sync {
async fn settle_usage(
async fn create_wallet_recharge_order(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, crate::DataLayerError>;
input: CreateWalletRechargeOrderInput,
) -> Result<CreateWalletRechargeOrderOutcome, crate::DataLayerError>;
async fn create_wallet_refund_request(
&self,
input: CreateWalletRefundRequestInput,
) -> Result<CreateWalletRefundRequestOutcome, crate::DataLayerError>;
async fn process_payment_callback(
&self,
input: ProcessPaymentCallbackInput,
) -> Result<ProcessPaymentCallbackOutcome, crate::DataLayerError>;
async fn adjust_wallet_balance(
&self,
input: AdjustWalletBalanceInput,
) -> Result<Option<(StoredWalletSnapshot, StoredAdminWalletTransaction)>, crate::DataLayerError>;
async fn create_manual_wallet_recharge(
&self,
input: CreateManualWalletRechargeInput,
) -> Result<Option<(StoredWalletSnapshot, StoredAdminPaymentOrder)>, crate::DataLayerError>;
async fn process_admin_wallet_refund(
&self,
input: ProcessAdminWalletRefundInput,
) -> Result<
WalletMutationOutcome<(
StoredWalletSnapshot,
StoredAdminWalletRefund,
StoredAdminWalletTransaction,
)>,
crate::DataLayerError,
>;
async fn complete_admin_wallet_refund(
&self,
input: CompleteAdminWalletRefundInput,
) -> Result<WalletMutationOutcome<StoredAdminWalletRefund>, crate::DataLayerError>;
async fn fail_admin_wallet_refund(
&self,
input: FailAdminWalletRefundInput,
) -> Result<
WalletMutationOutcome<(
StoredWalletSnapshot,
StoredAdminWalletRefund,
Option<StoredAdminWalletTransaction>,
)>,
crate::DataLayerError,
>;
async fn expire_admin_payment_order(
&self,
order_id: &str,
) -> Result<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>, crate::DataLayerError>;
async fn fail_admin_payment_order(
&self,
order_id: &str,
) -> Result<WalletMutationOutcome<StoredAdminPaymentOrder>, crate::DataLayerError>;
async fn credit_admin_payment_order(
&self,
input: CreditAdminPaymentOrderInput,
) -> Result<WalletMutationOutcome<(StoredAdminPaymentOrder, bool)>, crate::DataLayerError>;
}
pub trait WalletRepository: WalletReadRepository + WalletWriteRepository + Send + Sync {}
@@ -175,7 +764,8 @@ impl<T> WalletRepository for T where T: WalletReadRepository + WalletWriteReposi
#[cfg(test)]
mod tests {
use super::{StoredWalletSnapshot, UsageSettlementInput};
use super::StoredWalletSnapshot;
use crate::repository::settlement::UsageSettlementInput;
#[test]
fn rejects_invalid_wallet_snapshot() {

View File

@@ -0,0 +1,19 @@
[package]
name = "aether-model-fetch"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
aether-contracts.workspace = true
aether-data.workspace = true
aether-provider-transport.workspace = true
aether-scheduler-core.workspace = true
async-trait.workspace = true
regex.workspace = true
serde_json.workspace = true
uuid.workspace = true
[dev-dependencies]
tokio.workspace = true

View File

@@ -0,0 +1,226 @@
use std::collections::BTreeSet;
use aether_data::repository::global_models::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, StoredAdminGlobalModelPage,
StoredAdminProviderModel, UpsertAdminProviderModelRecord,
};
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::matches_model_mapping;
use async_trait::async_trait;
use serde_json::Value;
use uuid::Uuid;
use crate::json_string_list;
#[async_trait]
pub trait ModelFetchAssociationStore {
type Error: Send;
fn has_global_model_reader(&self) -> bool;
fn has_global_model_writer(&self) -> bool;
fn model_fetch_internal_error(&self, message: String) -> Self::Error;
async fn list_admin_provider_models(
&self,
query: &AdminProviderModelListQuery,
) -> Result<Vec<StoredAdminProviderModel>, Self::Error>;
async fn list_admin_global_models(
&self,
query: &AdminGlobalModelListQuery,
) -> Result<StoredAdminGlobalModelPage, Self::Error>;
async fn create_admin_provider_model(
&self,
record: &UpsertAdminProviderModelRecord,
) -> Result<Option<StoredAdminProviderModel>, Self::Error>;
async fn list_provider_catalog_keys_by_provider_ids(
&self,
provider_ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, Self::Error>;
async fn delete_admin_provider_model(
&self,
provider_id: &str,
model_id: &str,
) -> Result<bool, Self::Error>;
}
pub async fn sync_provider_model_whitelist_associations<S>(
state: &S,
provider_id: &str,
current_allowed_models: &[String],
) -> Result<(), S::Error>
where
S: ModelFetchAssociationStore + Sync + ?Sized,
{
if !state.has_global_model_reader() || !state.has_global_model_writer() {
return Ok(());
}
auto_associate_provider_by_key_whitelist(state, provider_id, current_allowed_models).await?;
auto_disassociate_provider_by_key_whitelist(state, provider_id).await?;
Ok(())
}
async fn auto_associate_provider_by_key_whitelist<S>(
state: &S,
provider_id: &str,
allowed_models: &[String],
) -> Result<(), S::Error>
where
S: ModelFetchAssociationStore + Sync + ?Sized,
{
if allowed_models.is_empty() {
return Ok(());
}
let provider_models = state
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: provider_id.to_string(),
is_active: None,
offset: 0,
limit: 10_000,
})
.await?;
let linked_global_model_ids = provider_models
.iter()
.map(|model| model.global_model_id.clone())
.collect::<BTreeSet<_>>();
let existing_provider_model_names = provider_models
.iter()
.map(|model| model.provider_model_name.clone())
.collect::<BTreeSet<_>>();
let global_models = state
.list_admin_global_models(&AdminGlobalModelListQuery {
offset: 0,
limit: 10_000,
is_active: Some(true),
search: None,
})
.await?
.items;
for global_model in global_models {
if linked_global_model_ids.contains(&global_model.id)
|| existing_provider_model_names.contains(&global_model.name)
{
continue;
}
let mappings = global_model_mapping_patterns(global_model.config.as_ref());
if mappings.is_empty() {
continue;
}
if !allowed_models.iter().any(|allowed_model| {
mappings
.iter()
.any(|pattern| matches_model_mapping(pattern, allowed_model))
}) {
continue;
}
let record = UpsertAdminProviderModelRecord::new(
Uuid::new_v4().to_string(),
provider_id.to_string(),
global_model.id.clone(),
global_model.name.clone(),
None,
None,
None,
None,
None,
None,
None,
None,
true,
true,
None,
)
.map_err(|err| state.model_fetch_internal_error(err.to_string()))?;
state.create_admin_provider_model(&record).await?;
}
Ok(())
}
async fn auto_disassociate_provider_by_key_whitelist<S>(
state: &S,
provider_id: &str,
) -> Result<(), S::Error>
where
S: ModelFetchAssociationStore + Sync + ?Sized,
{
let keys = state
.list_provider_catalog_keys_by_provider_ids(&[provider_id.to_string()])
.await?;
let active_non_oauth_keys = keys
.into_iter()
.filter(|key| key.is_active)
.filter(|key| !is_oauth_auth_type(&key.auth_type))
.collect::<Vec<_>>();
if active_non_oauth_keys.is_empty() {
return Ok(());
}
if active_non_oauth_keys
.iter()
.any(|key| key.allowed_models.is_none())
{
return Ok(());
}
let all_allowed_models = active_non_oauth_keys
.iter()
.flat_map(|key| json_string_list(key.allowed_models.as_ref()))
.collect::<BTreeSet<_>>();
let provider_models = state
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: provider_id.to_string(),
is_active: None,
offset: 0,
limit: 10_000,
})
.await?;
for model in provider_models {
let mappings = global_model_mapping_patterns(model.global_model_config.as_ref());
if mappings.is_empty() {
continue;
}
let matched = all_allowed_models.iter().any(|allowed_model| {
mappings
.iter()
.any(|pattern| matches_model_mapping(pattern, allowed_model))
});
if matched {
continue;
}
state
.delete_admin_provider_model(provider_id, &model.id)
.await?;
}
Ok(())
}
fn global_model_mapping_patterns(config: Option<&Value>) -> Vec<String> {
config
.and_then(Value::as_object)
.and_then(|object| object.get("model_mappings"))
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default()
}
fn is_oauth_auth_type(value: &str) -> bool {
matches!(value.trim().to_ascii_lowercase().as_str(), "oauth" | "kiro")
}

View File

@@ -0,0 +1,77 @@
const MODEL_FETCH_INTERVAL_MINUTES_DEFAULT: u64 = 1440;
const MODEL_FETCH_INTERVAL_MINUTES_MIN: u64 = 60;
const MODEL_FETCH_INTERVAL_MINUTES_MAX: u64 = 10080;
const MODEL_FETCH_STARTUP_DELAY_SECONDS_DEFAULT: u64 = 10;
pub fn model_fetch_interval_minutes() -> u64 {
std::env::var("MODEL_FETCH_INTERVAL_MINUTES")
.ok()
.and_then(|value| value.parse::<u64>().ok())
.map(|value| {
value.clamp(
MODEL_FETCH_INTERVAL_MINUTES_MIN,
MODEL_FETCH_INTERVAL_MINUTES_MAX,
)
})
.unwrap_or(MODEL_FETCH_INTERVAL_MINUTES_DEFAULT)
}
pub fn model_fetch_startup_enabled() -> bool {
std::env::var("MODEL_FETCH_STARTUP_ENABLED")
.ok()
.map(|value| value.trim().eq_ignore_ascii_case("true"))
.unwrap_or(true)
}
pub fn model_fetch_startup_delay_seconds() -> u64 {
std::env::var("MODEL_FETCH_STARTUP_DELAY_SECONDS")
.ok()
.and_then(|value| value.parse::<u64>().ok())
.unwrap_or(MODEL_FETCH_STARTUP_DELAY_SECONDS_DEFAULT)
}
#[cfg(test)]
mod tests {
use super::{
model_fetch_interval_minutes, model_fetch_startup_delay_seconds,
model_fetch_startup_enabled,
};
struct TestEnvVarGuard {
key: &'static str,
previous: Option<String>,
}
impl Drop for TestEnvVarGuard {
fn drop(&mut self) {
if let Some(previous) = self.previous.as_deref() {
std::env::set_var(self.key, previous);
} else {
std::env::remove_var(self.key);
}
}
}
fn set_test_env_var(key: &'static str, value: &str) -> TestEnvVarGuard {
let previous = std::env::var(key).ok();
std::env::set_var(key, value);
TestEnvVarGuard { key, previous }
}
#[test]
fn interval_minutes_clamps_to_supported_bounds() {
let _interval = set_test_env_var("MODEL_FETCH_INTERVAL_MINUTES", "5");
assert_eq!(model_fetch_interval_minutes(), 60);
let _interval = set_test_env_var("MODEL_FETCH_INTERVAL_MINUTES", "20000");
assert_eq!(model_fetch_interval_minutes(), 10080);
}
#[test]
fn startup_flags_read_from_environment() {
let _enabled = set_test_env_var("MODEL_FETCH_STARTUP_ENABLED", "false");
let _delay = set_test_env_var("MODEL_FETCH_STARTUP_DELAY_SECONDS", "3");
assert!(!model_fetch_startup_enabled());
assert_eq!(model_fetch_startup_delay_seconds(), 3);
}
}

View File

@@ -0,0 +1,17 @@
mod association_sync;
mod config;
mod logic;
mod transport;
pub use association_sync::{
sync_provider_model_whitelist_associations, ModelFetchAssociationStore,
};
pub use config::{
model_fetch_interval_minutes, model_fetch_startup_delay_seconds, model_fetch_startup_enabled,
};
pub use logic::{
aggregate_models_for_cache, apply_model_filters, build_models_fetch_url,
endpoint_supports_rust_models_fetch, extract_error_message, json_string_list,
parse_models_response, select_models_fetch_endpoint, ModelFetchRunSummary, ModelsFetchSuccess,
};
pub use transport::{build_models_fetch_execution_plan, ModelFetchTransportRuntime};

View File

@@ -0,0 +1,510 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_provider_transport::provider_types::provider_type_supports_model_fetch;
use regex::Regex;
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelFetchRunSummary {
pub attempted: usize,
pub succeeded: usize,
pub failed: usize,
pub skipped: usize,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ModelsFetchSuccess {
pub fetched_model_ids: Vec<String>,
pub cached_models: Vec<Value>,
}
pub fn extract_error_message(value: &Value) -> Option<String> {
value
.get("error")
.and_then(Value::as_object)
.and_then(|error| error.get("message"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| {
value
.get("message")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
pub fn build_models_fetch_url(
provider_type: &str,
endpoint_api_format: &str,
base_url: &str,
) -> Option<(String, String)> {
let api_format = normalize_api_format(endpoint_api_format);
if !provider_type_supports_model_fetch(provider_type) {
return None;
}
let url = if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
build_v1_models_url(base_url)
} else if api_format.starts_with("gemini:") {
build_gemini_models_url(base_url)
} else {
return None;
}?;
Some((url, api_format))
}
pub fn parse_models_response(
endpoint_api_format: &str,
body: &Value,
) -> Result<ModelsFetchSuccess, String> {
let api_format = normalize_api_format(endpoint_api_format);
let mut cached_models = Vec::new();
let mut fetched_model_ids = Vec::new();
let mut seen = BTreeSet::new();
if api_format.starts_with("openai:") || api_format.starts_with("claude:") {
let items = if let Some(items) = body.get("data").and_then(Value::as_array) {
items
} else if let Some(items) = body.as_array() {
items
} else {
return Err("models response is missing data array".to_string());
};
for item in items {
let Some(model_id) = item
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
if !seen.insert(model_id.to_string()) {
continue;
}
fetched_model_ids.push(model_id.to_string());
cached_models.push(normalize_cached_model(item, model_id, &api_format));
}
} else if api_format.starts_with("gemini:") {
let items = body
.get("models")
.and_then(Value::as_array)
.ok_or_else(|| "gemini models response is missing models array".to_string())?;
for item in items {
let Some(name) = item
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let model_id = name.strip_prefix("models/").unwrap_or(name).trim();
if model_id.is_empty() || !seen.insert(model_id.to_string()) {
continue;
}
fetched_model_ids.push(model_id.to_string());
cached_models.push(normalize_cached_model(item, model_id, &api_format));
}
} else {
return Err("models response parser does not support this provider format".to_string());
}
Ok(ModelsFetchSuccess {
fetched_model_ids,
cached_models,
})
}
pub fn select_models_fetch_endpoint(
endpoints: &[StoredProviderCatalogEndpoint],
key: &StoredProviderCatalogKey,
) -> Option<StoredProviderCatalogEndpoint> {
let key_formats = json_string_list(key.api_formats.as_ref())
.into_iter()
.map(|value| normalize_api_format(&value))
.collect::<BTreeSet<_>>();
endpoints
.iter()
.filter(|endpoint| endpoint.is_active)
.find(|endpoint| {
let api_format = normalize_api_format(&endpoint.api_format);
(key_formats.is_empty() || key_formats.contains(&api_format))
&& endpoint_supports_rust_models_fetch(&endpoint.api_format)
})
.cloned()
}
pub fn endpoint_supports_rust_models_fetch(api_format: &str) -> bool {
let api_format = normalize_api_format(api_format);
matches!(
api_format.as_str(),
"openai:chat"
| "openai:cli"
| "openai:responses"
| "openai:compact"
| "claude:chat"
| "claude:cli"
| "gemini:chat"
| "gemini:cli"
)
}
pub fn apply_model_filters(
fetched_model_ids: &[String],
locked_models: Vec<String>,
include_patterns: Vec<String>,
exclude_patterns: Vec<String>,
) -> Vec<String> {
let mut filtered = BTreeSet::new();
for model_id in fetched_model_ids {
if model_id.trim().is_empty() {
continue;
}
let included = if include_patterns.is_empty() {
true
} else {
include_patterns
.iter()
.any(|pattern| wildcard_matches(pattern, model_id))
};
if !included {
continue;
}
let excluded = exclude_patterns
.iter()
.any(|pattern| wildcard_matches(pattern, model_id));
if !excluded {
filtered.insert(model_id.trim().to_string());
}
}
for model in locked_models {
let trimmed = model.trim();
if !trimmed.is_empty() {
filtered.insert(trimmed.to_string());
}
}
filtered.into_iter().collect()
}
pub fn json_string_list(value: Option<&Value>) -> Vec<String> {
value
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
})
.unwrap_or_default()
}
pub fn aggregate_models_for_cache(models: &[Value]) -> Vec<Value> {
let mut aggregated = BTreeMap::<String, serde_json::Map<String, Value>>::new();
let mut order = Vec::<String>::new();
for model in models {
let Some(object) = model.as_object() else {
continue;
};
let Some(model_id) = object
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let entry = aggregated.entry(model_id.to_string()).or_insert_with(|| {
order.push(model_id.to_string());
let mut cloned = object.clone();
cloned.remove("api_format");
cloned
});
let api_formats = object
.get("api_formats")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
let existing_formats = entry
.get("api_formats")
.and_then(Value::as_array)
.map(|items| {
items
.iter()
.filter_map(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
let merged_formats = existing_formats
.union(&api_formats)
.cloned()
.map(Value::String)
.collect::<Vec<_>>();
entry.insert("api_formats".to_string(), Value::Array(merged_formats));
for (key, value) in object {
if key == "api_format" || entry.contains_key(key) {
continue;
}
entry.insert(key.clone(), value.clone());
}
}
order
.into_iter()
.filter_map(|model_id| aggregated.remove(&model_id))
.map(Value::Object)
.collect()
}
fn build_v1_models_url(base_url: &str) -> Option<String> {
let (trimmed_base_url, query) = split_url_query(base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
if trimmed_base_url.is_empty() {
return None;
}
let mut url = if trimmed_base_url.ends_with("/v1") {
format!("{trimmed_base_url}/models")
} else {
format!("{trimmed_base_url}/v1/models")
};
if let Some(query) = query.filter(|value| !value.trim().is_empty()) {
url.push('?');
url.push_str(query);
}
Some(url)
}
fn build_gemini_models_url(base_url: &str) -> Option<String> {
let (trimmed_base_url, base_query) = split_url_query(base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
if trimmed_base_url.is_empty() {
return None;
}
let mut url = if trimmed_base_url.ends_with("/v1beta") {
format!("{trimmed_base_url}/models")
} else if trimmed_base_url.contains("/v1beta/models") {
trimmed_base_url.to_string()
} else {
format!("{trimmed_base_url}/v1beta/models")
};
if let Some(query) = base_query.filter(|value| !value.trim().is_empty()) {
url.push('?');
url.push_str(query);
}
Some(url)
}
fn split_url_query(base_url: &str) -> (&str, Option<&str>) {
let trimmed = base_url.trim();
trimmed
.split_once('?')
.map(|(base, query)| (base, Some(query)))
.unwrap_or((trimmed, None))
}
fn normalize_cached_model(item: &Value, model_id: &str, api_format: &str) -> Value {
let mut object = item.as_object().cloned().unwrap_or_default();
object.insert("id".to_string(), Value::String(model_id.to_string()));
object.insert(
"api_formats".to_string(),
Value::Array(vec![Value::String(api_format.to_string())]),
);
object.remove("api_format");
Value::Object(object)
}
fn wildcard_matches(pattern: &str, model_id: &str) -> bool {
let mut regex = String::from("^");
for ch in pattern.chars() {
match ch {
'*' => regex.push_str(".*"),
'?' => regex.push('.'),
other => regex.push_str(&regex::escape(&other.to_string())),
}
}
regex.push('$');
Regex::new(&regex)
.ok()
.is_some_and(|compiled| compiled.is_match(model_id))
}
fn normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}
#[cfg(test)]
mod tests {
use aether_data::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use serde_json::json;
use super::{
aggregate_models_for_cache, apply_model_filters, build_gemini_models_url,
build_models_fetch_url, parse_models_response, select_models_fetch_endpoint,
};
fn sample_endpoint(
provider_id: &str,
endpoint_id: &str,
api_format: &str,
base_url: &str,
) -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
endpoint_id.to_string(),
provider_id.to_string(),
api_format.to_string(),
None,
None,
true,
)
.expect("endpoint should build")
.with_transport_fields(
base_url.to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_key(provider_id: &str, key_id: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
key_id.to_string(),
provider_id.to_string(),
"primary".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:chat"])),
"encrypted".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build")
}
#[test]
fn apply_model_filters_respects_include_exclude_and_locked_models() {
let filtered = apply_model_filters(
&[
"gpt-5".to_string(),
"gpt-beta".to_string(),
"claude-4".to_string(),
],
vec!["locked-model".to_string()],
vec!["gpt-*".to_string()],
vec!["gpt-beta".to_string()],
);
assert_eq!(
filtered,
vec!["gpt-5".to_string(), "locked-model".to_string()]
);
}
#[test]
fn aggregate_models_for_cache_merges_api_formats_by_model_id() {
let aggregated = aggregate_models_for_cache(&[
json!({"id":"gpt-5","api_formats":["openai:chat"]}),
json!({"id":"gpt-5","api_formats":["openai:cli"]}),
]);
assert_eq!(aggregated.len(), 1);
assert_eq!(
aggregated[0]["api_formats"],
json!(["openai:chat", "openai:cli"])
);
}
#[test]
fn build_gemini_models_url_preserves_base_query() {
let url =
build_gemini_models_url("https://generativelanguage.googleapis.com/v1beta?key=abc")
.expect("gemini models url should build");
assert_eq!(
url,
"https://generativelanguage.googleapis.com/v1beta/models?key=abc"
);
}
#[test]
fn build_models_fetch_url_rejects_provider_types_without_fetch_support() {
assert_eq!(
build_models_fetch_url("vertex_ai", "gemini:chat", "https://example.com"),
None
);
}
#[test]
fn parse_models_response_normalizes_openai_payload() {
let parsed = parse_models_response(
"openai:chat",
&json!({"data": [{"id": "gpt-5"}, {"id": "gpt-5"}]}),
)
.expect("response should parse");
assert_eq!(parsed.fetched_model_ids, vec!["gpt-5".to_string()]);
assert_eq!(
parsed.cached_models[0]["api_formats"],
json!(["openai:chat"])
);
}
#[test]
fn select_models_fetch_endpoint_respects_key_api_formats() {
let key = sample_key("provider-1", "key-1");
let endpoints = vec![
sample_endpoint(
"provider-1",
"endpoint-cli",
"openai:cli",
"https://example.com",
),
sample_endpoint(
"provider-1",
"endpoint-chat",
"openai:chat",
"https://example.com",
),
];
let selected =
select_models_fetch_endpoint(&endpoints, &key).expect("endpoint should be selected");
assert_eq!(selected.id, "endpoint-chat");
}
}

View File

@@ -0,0 +1,267 @@
use std::collections::BTreeMap;
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use aether_provider_transport::auth::{
resolve_local_gemini_auth, resolve_local_openai_chat_auth, resolve_local_standard_auth,
};
use aether_provider_transport::url::build_passthrough_path_url;
use aether_provider_transport::vertex::resolve_local_vertex_api_key_query_auth;
use aether_provider_transport::{
apply_local_header_rules, ensure_upstream_auth_header, resolve_transport_execution_timeouts,
resolve_transport_tls_profile, GatewayProviderTransportSnapshot, LocalResolvedOAuthRequestAuth,
};
use async_trait::async_trait;
use serde_json::json;
use crate::build_models_fetch_url;
#[async_trait]
pub trait ModelFetchTransportRuntime: Send + Sync {
async fn resolve_local_oauth_request_auth(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<LocalResolvedOAuthRequestAuth>, String>;
async fn resolve_model_fetch_proxy(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot>;
}
pub async fn build_models_fetch_execution_plan(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<ExecutionPlan, String> {
let (upstream_url, provider_api_format) = build_models_fetch_url(
&transport.provider.provider_type,
&transport.endpoint.api_format,
&transport.endpoint.base_url,
)
.ok_or_else(|| "Rust models fetch does not support this provider format yet".to_string())?;
let (auth_header_name, auth_header_value) = resolve_models_fetch_auth(runtime, transport)
.await?
.ok_or_else(|| {
"Rust models fetch auth resolution is not supported for this key".to_string()
})?;
let mut headers = BTreeMap::from([(auth_header_name.clone(), auth_header_value.clone())]);
if !apply_local_header_rules(
&mut headers,
transport.endpoint.header_rules.as_ref(),
&[auth_header_name.as_str()],
&json!({}),
None,
) {
return Err("Endpoint header_rules application failed".to_string());
}
ensure_upstream_auth_header(&mut headers, &auth_header_name, &auth_header_value);
Ok(ExecutionPlan {
request_id: format!("req-model-fetch-{}", transport.key.id),
candidate_id: None,
provider_name: Some(transport.provider.name.clone()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: "GET".to_string(),
url: upstream_url,
headers,
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: false,
client_api_format: provider_api_format.clone(),
provider_api_format,
model_name: None,
proxy: runtime.resolve_model_fetch_proxy(transport).await,
tls_profile: resolve_transport_tls_profile(transport),
timeouts: resolve_transport_execution_timeouts(transport),
})
}
async fn resolve_models_fetch_auth(
runtime: &(impl ModelFetchTransportRuntime + ?Sized),
transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<(String, String)>, String> {
if transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
|| transport.key.auth_type.trim().eq_ignore_ascii_case("kiro")
{
return match runtime.resolve_local_oauth_request_auth(transport).await {
Ok(Some(LocalResolvedOAuthRequestAuth::Header { name, value })) => {
Ok(Some((name, value)))
}
Ok(Some(LocalResolvedOAuthRequestAuth::Kiro(_))) => Ok(None),
Ok(None) => Ok(None),
Err(err) => Err(err),
};
}
if let Some(auth) = resolve_local_openai_chat_auth(transport) {
return Ok(Some(auth));
}
if let Some(auth) = resolve_local_standard_auth(transport) {
return Ok(Some(auth));
}
if let Some(auth) = resolve_local_gemini_auth(transport) {
return Ok(Some(auth));
}
if let Some(query_auth) = resolve_local_vertex_api_key_query_auth(transport) {
let url = build_passthrough_path_url(
&transport.endpoint.base_url,
"/v1/publishers/google/models",
Some(&format!("{}={}", query_auth.name, query_auth.value)),
&[],
);
if url.is_some() {
return Ok(None);
}
}
Ok(None)
}
#[cfg(test)]
mod tests {
use aether_contracts::ProxySnapshot;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use async_trait::async_trait;
use super::{build_models_fetch_execution_plan, ModelFetchTransportRuntime};
struct TestRuntime {
oauth_auth: Option<aether_provider_transport::LocalResolvedOAuthRequestAuth>,
proxy: Option<ProxySnapshot>,
}
#[async_trait]
impl ModelFetchTransportRuntime for TestRuntime {
async fn resolve_local_oauth_request_auth(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Result<Option<aether_provider_transport::LocalResolvedOAuthRequestAuth>, String>
{
Ok(self.oauth_auth.clone())
}
async fn resolve_model_fetch_proxy(
&self,
_transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot> {
self.proxy.clone()
}
}
fn sample_transport(api_format: &str, auth_type: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider One".to_string(),
provider_type: "openai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: Some(30.0),
stream_first_byte_timeout_secs: Some(5.0),
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: api_format.to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://example.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: auth_type.to_string(),
is_active: true,
api_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[tokio::test]
async fn builds_openai_models_fetch_plan_from_transport_snapshot() {
let runtime = TestRuntime {
oauth_auth: None,
proxy: None,
};
let plan = build_models_fetch_execution_plan(
&runtime,
&sample_transport("openai:chat", "api_key"),
)
.await
.expect("plan");
assert_eq!(plan.method, "GET");
assert_eq!(plan.url, "https://example.com/v1/models");
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer secret")
);
assert_eq!(plan.provider_api_format, "openai:chat");
}
#[tokio::test]
async fn builds_oauth_models_fetch_plan_from_runtime_auth() {
let runtime = TestRuntime {
oauth_auth: Some(
aether_provider_transport::LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer oauth-token".to_string(),
},
),
proxy: Some(ProxySnapshot {
enabled: Some(true),
mode: Some("fixed".to_string()),
node_id: None,
label: None,
url: Some("http://proxy.internal".to_string()),
extra: None,
}),
};
let plan =
build_models_fetch_execution_plan(&runtime, &sample_transport("openai:chat", "oauth"))
.await
.expect("plan");
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer oauth-token")
);
assert_eq!(
plan.proxy.as_ref().and_then(|proxy| proxy.url.as_deref()),
Some("http://proxy.internal")
);
}
}

View File

@@ -0,0 +1,28 @@
[package]
name = "aether-provider-transport"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Provider transport core extracted from aether-gateway"
[dependencies]
aether-contracts.workspace = true
aether-crypto.workspace = true
aether-data.workspace = true
aether-video-tasks-core.workspace = true
async-trait.workspace = true
http.workspace = true
regex.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true
url.workspace = true
uuid.workspace = true
[dev-dependencies]
axum = { version = "0.8", features = ["ws"] }

View File

@@ -0,0 +1,23 @@
mod auth;
mod policy;
mod request;
mod url;
pub use auth::{
build_antigravity_static_identity_headers, resolve_local_antigravity_request_auth,
AntigravityRequestAuth, AntigravityRequestAuthSupport, AntigravityRequestAuthUnsupportedReason,
ANTIGRAVITY_PROVIDER_TYPE, ANTIGRAVITY_REQUEST_USER_AGENT,
};
pub use policy::{
classify_local_antigravity_request_support, AntigravityRequestSideSpec,
AntigravityRequestSideSupport, AntigravityRequestSideUnsupportedReason,
};
pub use request::{
build_antigravity_safe_v1internal_request, classify_antigravity_safe_request_body,
AntigravityEnvelopeRequestType, AntigravityRequestEnvelopeSupport,
AntigravityRequestEnvelopeUnsupportedReason,
};
pub use url::{
build_antigravity_v1internal_url, AntigravityRequestUrlAction,
ANTIGRAVITY_V1INTERNAL_PATH_TEMPLATE,
};

View File

@@ -0,0 +1,213 @@
use std::collections::BTreeMap;
use serde_json::Value;
use super::super::snapshot::GatewayProviderTransportSnapshot;
pub const ANTIGRAVITY_PROVIDER_TYPE: &str = "antigravity";
pub const ANTIGRAVITY_REQUEST_USER_AGENT: &str = "antigravity";
const ANTIGRAVITY_CLIENT_NAME: &str = "antigravity";
const ANTIGRAVITY_GOOG_API_CLIENT: &str = "gl-node/18.18.2 fire/0.8.6 grpc/1.10.x";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AntigravityRequestAuth {
pub project_id: String,
pub client_version: Option<String>,
pub session_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AntigravityRequestAuthSupport {
Supported(AntigravityRequestAuth),
Unsupported(AntigravityRequestAuthUnsupportedReason),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AntigravityRequestAuthUnsupportedReason {
WrongProviderType,
MissingAuthConfig,
InvalidAuthConfigJson,
ComplexDynamicAuthConfig,
MissingProjectId,
}
pub fn resolve_local_antigravity_request_auth(
transport: &GatewayProviderTransportSnapshot,
) -> AntigravityRequestAuthSupport {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(ANTIGRAVITY_PROVIDER_TYPE)
{
return AntigravityRequestAuthSupport::Unsupported(
AntigravityRequestAuthUnsupportedReason::WrongProviderType,
);
}
let Some(raw_auth_config) = transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return AntigravityRequestAuthSupport::Unsupported(
AntigravityRequestAuthUnsupportedReason::MissingAuthConfig,
);
};
let Ok(auth_config) = serde_json::from_str::<Value>(raw_auth_config) else {
return AntigravityRequestAuthSupport::Unsupported(
AntigravityRequestAuthUnsupportedReason::InvalidAuthConfigJson,
);
};
if contains_blocked_auth_fields(&auth_config) {
return AntigravityRequestAuthSupport::Unsupported(
AntigravityRequestAuthUnsupportedReason::ComplexDynamicAuthConfig,
);
}
let Some(project_id) = find_string_by_paths(
&auth_config,
&[
&["project_id"],
&["projectId"],
&["project", "id"],
&["project", "project_id"],
&["project", "projectId"],
&["antigravity", "project_id"],
&["antigravity", "projectId"],
&["metadata", "project_id"],
&["metadata", "projectId"],
],
) else {
return AntigravityRequestAuthSupport::Unsupported(
AntigravityRequestAuthUnsupportedReason::MissingProjectId,
);
};
let client_version = find_string_by_paths(
&auth_config,
&[
&["client_version"],
&["clientVersion"],
&["antigravity", "client_version"],
&["antigravity", "clientVersion"],
&["metadata", "client_version"],
&["metadata", "clientVersion"],
],
);
let session_id = find_string_by_paths(
&auth_config,
&[
&["session_id"],
&["sessionId"],
&["antigravity", "session_id"],
&["antigravity", "sessionId"],
&["metadata", "session_id"],
&["metadata", "sessionId"],
],
);
AntigravityRequestAuthSupport::Supported(AntigravityRequestAuth {
project_id,
client_version,
session_id,
})
}
pub fn build_antigravity_static_identity_headers(
auth: &AntigravityRequestAuth,
) -> BTreeMap<String, String> {
let mut headers = BTreeMap::from([
(
String::from("x-client-name"),
String::from(ANTIGRAVITY_CLIENT_NAME),
),
(
String::from("x-goog-api-client"),
String::from(ANTIGRAVITY_GOOG_API_CLIENT),
),
]);
if let Some(client_version) = auth
.client_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
headers.insert(String::from("x-client-version"), client_version.to_string());
}
if let Some(session_id) = auth
.session_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
headers.insert(String::from("x-vscode-sessionid"), session_id.to_string());
}
headers
}
fn find_string_by_paths(value: &Value, paths: &[&[&str]]) -> Option<String> {
for path in paths {
let mut current = value;
let mut matched = true;
for segment in *path {
let Some(next) = current.get(*segment) else {
matched = false;
break;
};
current = next;
}
if !matched {
continue;
}
if let Some(string) = current
.as_str()
.map(str::trim)
.filter(|item| !item.is_empty())
{
return Some(string.to_string());
}
}
None
}
fn contains_blocked_auth_fields(value: &Value) -> bool {
match value {
Value::Object(map) => map.iter().any(|(key, inner)| {
is_blocked_auth_key(key.as_str()) || contains_blocked_auth_fields(inner)
}),
Value::Array(items) => items.iter().any(contains_blocked_auth_fields),
_ => false,
}
}
fn is_blocked_auth_key(key: &str) -> bool {
matches!(
key.trim().to_ascii_lowercase().as_str(),
"private_key"
| "privateKey"
| "private_key_id"
| "privateKeyId"
| "service_account"
| "serviceAccount"
| "service_account_json"
| "serviceAccountJson"
| "service_account_key"
| "serviceAccountKey"
| "credential_source"
| "credentialSource"
| "token_url"
| "tokenUrl"
| "auth_uri"
| "authUri"
| "subject"
| "audience"
)
}

View File

@@ -0,0 +1,113 @@
use serde_json::Value;
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::auth::{
resolve_local_antigravity_request_auth, AntigravityRequestAuth, AntigravityRequestAuthSupport,
AntigravityRequestAuthUnsupportedReason, ANTIGRAVITY_PROVIDER_TYPE,
};
use super::request::{
classify_antigravity_safe_request_body, AntigravityEnvelopeRequestType,
AntigravityRequestEnvelopeUnsupportedReason,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AntigravityRequestSideSpec {
pub auth: AntigravityRequestAuth,
pub request_type: AntigravityEnvelopeRequestType,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AntigravityRequestSideSupport {
Supported(AntigravityRequestSideSpec),
Unsupported(AntigravityRequestSideUnsupportedReason),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AntigravityRequestSideUnsupportedReason {
InactiveTransport,
WrongProviderType,
UnsupportedApiFormat,
UnsupportedCustomPath,
UnsupportedHeaderRules,
UnsupportedBodyRules,
UnsupportedNetworkConfig,
UnsupportedAuth(AntigravityRequestAuthUnsupportedReason),
UnsupportedEnvelope(AntigravityRequestEnvelopeUnsupportedReason),
}
pub fn classify_local_antigravity_request_support(
transport: &GatewayProviderTransportSnapshot,
request_body: &Value,
request_type: AntigravityEnvelopeRequestType,
) -> AntigravityRequestSideSupport {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::InactiveTransport,
);
}
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(ANTIGRAVITY_PROVIDER_TYPE)
{
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::WrongProviderType,
);
}
let endpoint_format = transport.endpoint.api_format.trim();
if !endpoint_format.eq_ignore_ascii_case("gemini:chat")
&& !endpoint_format.eq_ignore_ascii_case("gemini:cli")
{
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedApiFormat,
);
}
if transport
.endpoint
.custom_path
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
{
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedCustomPath,
);
}
if transport.endpoint.header_rules.is_some() {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedHeaderRules,
);
}
if transport.endpoint.body_rules.is_some() {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedBodyRules,
);
}
if transport.provider.proxy.is_some()
|| transport.endpoint.proxy.is_some()
|| transport.key.proxy.is_some()
|| transport.key.fingerprint.is_some()
{
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedNetworkConfig,
);
}
let auth = match resolve_local_antigravity_request_auth(transport) {
AntigravityRequestAuthSupport::Supported(auth) => auth,
AntigravityRequestAuthSupport::Unsupported(reason) => {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedAuth(reason),
);
}
};
if let Err(reason) = classify_antigravity_safe_request_body(request_body) {
return AntigravityRequestSideSupport::Unsupported(
AntigravityRequestSideUnsupportedReason::UnsupportedEnvelope(reason),
);
}
AntigravityRequestSideSupport::Supported(AntigravityRequestSideSpec { auth, request_type })
}

View File

@@ -0,0 +1,120 @@
use serde_json::{Map, Value};
use super::auth::{AntigravityRequestAuth, ANTIGRAVITY_REQUEST_USER_AGENT};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AntigravityEnvelopeRequestType {
Agent,
EndpointTest,
}
impl AntigravityEnvelopeRequestType {
fn as_str(self) -> &'static str {
match self {
Self::Agent => "agent",
Self::EndpointTest => "endpoint_test",
}
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum AntigravityRequestEnvelopeSupport {
Supported(Value),
Unsupported(AntigravityRequestEnvelopeUnsupportedReason),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum AntigravityRequestEnvelopeUnsupportedReason {
NonObjectBody,
MissingContents,
MissingRequestId,
MissingModel,
ComplexEnvelopeTransform,
}
pub fn classify_antigravity_safe_request_body(
request_body: &Value,
) -> Result<(), AntigravityRequestEnvelopeUnsupportedReason> {
let Value::Object(map) = request_body else {
return Err(AntigravityRequestEnvelopeUnsupportedReason::NonObjectBody);
};
if !map.contains_key("contents") {
return Err(AntigravityRequestEnvelopeUnsupportedReason::MissingContents);
}
if contains_blocked_request_features(request_body) {
return Err(AntigravityRequestEnvelopeUnsupportedReason::ComplexEnvelopeTransform);
}
Ok(())
}
pub fn build_antigravity_safe_v1internal_request(
auth: &AntigravityRequestAuth,
request_id: &str,
model: &str,
request_body: &Value,
request_type: AntigravityEnvelopeRequestType,
) -> AntigravityRequestEnvelopeSupport {
if request_id.trim().is_empty() {
return AntigravityRequestEnvelopeSupport::Unsupported(
AntigravityRequestEnvelopeUnsupportedReason::MissingRequestId,
);
}
if model.trim().is_empty() {
return AntigravityRequestEnvelopeSupport::Unsupported(
AntigravityRequestEnvelopeUnsupportedReason::MissingModel,
);
}
if let Err(reason) = classify_antigravity_safe_request_body(request_body) {
return AntigravityRequestEnvelopeSupport::Unsupported(reason);
}
let Value::Object(source) = request_body else {
return AntigravityRequestEnvelopeSupport::Unsupported(
AntigravityRequestEnvelopeUnsupportedReason::NonObjectBody,
);
};
let mut inner_request: Map<String, Value> = source.clone();
inner_request.remove("model");
inner_request.remove("safetySettings");
inner_request.remove("safety_settings");
AntigravityRequestEnvelopeSupport::Supported(serde_json::json!({
"project": auth.project_id,
"requestId": request_id,
"request": Value::Object(inner_request),
"model": model,
"userAgent": ANTIGRAVITY_REQUEST_USER_AGENT,
"requestType": request_type.as_str(),
}))
}
fn contains_blocked_request_features(value: &Value) -> bool {
match value {
Value::Object(map) => map.iter().any(|(key, inner)| {
is_blocked_request_key(key.as_str()) || contains_blocked_request_features(inner)
}),
Value::Array(items) => items.iter().any(contains_blocked_request_features),
_ => false,
}
}
fn is_blocked_request_key(key: &str) -> bool {
matches!(
key.trim(),
"systemInstruction"
| "system_instruction"
| "tools"
| "toolConfig"
| "tool_config"
| "thinkingConfig"
| "thinking_config"
| "imageConfig"
| "image_config"
| "functionCall"
| "function_call"
| "functionResponse"
| "function_response"
)
}

View File

@@ -0,0 +1,69 @@
use std::collections::BTreeMap;
use url::form_urlencoded;
pub const ANTIGRAVITY_V1INTERNAL_PATH_TEMPLATE: &str = "/v1internal:{action}";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AntigravityRequestUrlAction {
GenerateContent,
StreamGenerateContent,
}
impl AntigravityRequestUrlAction {
fn as_str(self) -> &'static str {
match self {
Self::GenerateContent => "generateContent",
Self::StreamGenerateContent => "streamGenerateContent",
}
}
fn is_stream(self) -> bool {
matches!(self, Self::StreamGenerateContent)
}
}
pub fn build_antigravity_v1internal_url(
base_url: &str,
action: AntigravityRequestUrlAction,
query: Option<&BTreeMap<String, String>>,
) -> Option<String> {
let trimmed_base = base_url.trim();
if trimmed_base.is_empty() {
return None;
}
let path = ANTIGRAVITY_V1INTERNAL_PATH_TEMPLATE.replace("{action}", action.as_str());
let mut url = format!("{}{}", trimmed_base.trim_end_matches('/'), path);
let mut params = BTreeMap::new();
if let Some(query) = query {
for (key, value) in query {
let key = key.trim();
let value = value.trim();
if key.is_empty() || value.is_empty() || key.eq_ignore_ascii_case("beta") {
continue;
}
params.insert(key.to_string(), value.to_string());
}
}
if action.is_stream() {
params
.entry(String::from("alt"))
.or_insert_with(|| String::from("sse"));
}
if !params.is_empty() {
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in params {
serializer.append_pair(key.as_str(), value.as_str());
}
let query_string = serializer.finish();
if !query_string.is_empty() {
url.push('?');
url.push_str(&query_string);
}
}
Some(url)
}

View File

@@ -0,0 +1,144 @@
use std::collections::BTreeMap;
use super::headers::should_skip_upstream_passthrough_header;
use super::snapshot::GatewayProviderTransportSnapshot;
fn collect_passthrough_headers(
headers: &http::HeaderMap,
extra_headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {
let mut out = BTreeMap::new();
for (name, value) in headers.iter() {
let Ok(value) = value.to_str() else {
continue;
};
let key = name.as_str().to_ascii_lowercase();
if should_skip_upstream_passthrough_header(&key) {
continue;
}
let value = value.trim();
if value.is_empty() {
continue;
}
out.insert(key, value.to_string());
}
for (key, value) in extra_headers {
let normalized_key = key.to_ascii_lowercase();
let value = value.trim();
if value.is_empty() {
continue;
}
out.insert(normalized_key, value.to_string());
}
out
}
pub fn build_passthrough_headers(
headers: &http::HeaderMap,
extra_headers: &BTreeMap<String, String>,
content_type: Option<&str>,
) -> BTreeMap<String, String> {
let mut out = collect_passthrough_headers(headers, extra_headers);
out.entry("content-type".to_string()).or_insert_with(|| {
content_type
.filter(|value| !value.trim().is_empty())
.unwrap_or("application/json")
.trim()
.to_string()
});
out.remove("content-length");
out
}
pub fn build_openai_passthrough_headers(
headers: &http::HeaderMap,
auth_header: &str,
auth_value: &str,
extra_headers: &BTreeMap<String, String>,
content_type: Option<&str>,
) -> BTreeMap<String, String> {
let mut out = build_passthrough_headers(headers, extra_headers, content_type);
ensure_upstream_auth_header(&mut out, auth_header, auth_value);
out
}
pub fn build_passthrough_headers_with_auth(
headers: &http::HeaderMap,
auth_header: &str,
auth_value: &str,
extra_headers: &BTreeMap<String, String>,
) -> BTreeMap<String, String> {
let mut out = collect_passthrough_headers(headers, extra_headers);
ensure_upstream_auth_header(&mut out, auth_header, auth_value);
out.remove("content-length");
out
}
pub fn ensure_upstream_auth_header(
headers: &mut BTreeMap<String, String>,
auth_header: &str,
auth_value: &str,
) {
let header_name = auth_header.trim().to_ascii_lowercase();
let header_value = auth_value.trim();
if header_name.is_empty() || header_value.is_empty() {
return;
}
if headers
.get(&header_name)
.map(|value| value.trim().is_empty())
.unwrap_or(true)
{
headers.insert(header_name, header_value.to_string());
}
}
pub fn resolve_local_openai_chat_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
if !matches!(auth_type.as_str(), "api_key" | "bearer") {
return None;
}
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
Some(("authorization".to_string(), format!("Bearer {secret}")))
}
pub fn resolve_local_standard_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
match auth_type.as_str() {
"api_key" => Some(("x-api-key".to_string(), secret.to_string())),
"bearer" => Some(("authorization".to_string(), format!("Bearer {secret}"))),
_ => None,
}
}
pub fn resolve_local_gemini_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<(String, String)> {
let auth_type = transport.key.auth_type.trim().to_ascii_lowercase();
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
match auth_type.as_str() {
"api_key" => Some(("x-goog-api-key".to_string(), secret.to_string())),
"bearer" => Some(("authorization".to_string(), format!("Bearer {secret}"))),
_ => None,
}
}

View File

@@ -0,0 +1,600 @@
use std::collections::BTreeMap;
use serde_json::Value;
use url::form_urlencoded;
const UNSAFE_AUTH_CONFIG_HEADER_NAMES: &[&str] = &[
"api-key",
"authorization",
"content-length",
"content-type",
"cookie",
"host",
"proxy-authorization",
"x-api-key",
"x-goog-api-key",
];
const UNSAFE_AUTH_CONFIG_QUERY_NAMES: &[&str] = &[
"access_token",
"api_key",
"apikey",
"authorization",
"key",
"token",
];
const SENSITIVE_AUTH_CONFIG_KEYS: &[&str] = &[
"access_token",
"api_key",
"apikey",
"authorization",
"client_email",
"client_id",
"client_secret",
"expires_at",
"id_token",
"key",
"private_key",
"refresh_token",
"service_account",
"token",
"token_uri",
];
const IGNORABLE_AUTH_CONFIG_METADATA_KEYS: &[&str] = &[
"account_id",
"account_name",
"account_user_id",
"auth_method",
"email",
"model_regions",
"organizations",
"plan_type",
"project_id",
"provider_type",
"region",
"tier",
"user_id",
"workspace_id",
"workspace_name",
];
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LocalAuthConfigSafeSubset {
pub headers: BTreeMap<String, String>,
pub query: BTreeMap<String, String>,
pub path: Option<String>,
}
impl LocalAuthConfigSafeSubset {
fn is_empty(&self) -> bool {
self.headers.is_empty() && self.query.is_empty() && self.path.is_none()
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LocalAuthConfigAbsorption {
Missing,
Unsupported,
Absorbed {
base_url: String,
header_rules: Option<Value>,
custom_path: Option<String>,
},
}
pub fn absorb_local_auth_config_safe_subset(
base_url: &str,
header_rules: Option<Value>,
custom_path: Option<String>,
raw_auth_config: Option<&str>,
) -> LocalAuthConfigAbsorption {
let Some(raw_auth_config) = raw_auth_config
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return LocalAuthConfigAbsorption::Missing;
};
let subset = match parse_local_auth_config_safe_subset(raw_auth_config) {
Ok(subset) => subset,
Err(()) => return LocalAuthConfigAbsorption::Unsupported,
};
if subset.is_empty() {
return LocalAuthConfigAbsorption::Unsupported;
}
let header_rules = match merge_auth_config_header_rules(header_rules, &subset.headers) {
Some(rules) => rules,
None => return LocalAuthConfigAbsorption::Unsupported,
};
let base_url = match merge_auth_config_base_url(base_url, &subset.query) {
Some(value) => value,
None => return LocalAuthConfigAbsorption::Unsupported,
};
let custom_path = match merge_auth_config_custom_path(custom_path, subset.path) {
Some(path) => path,
None => return LocalAuthConfigAbsorption::Unsupported,
};
LocalAuthConfigAbsorption::Absorbed {
base_url,
header_rules,
custom_path,
}
}
fn parse_local_auth_config_safe_subset(raw: &str) -> Result<LocalAuthConfigSafeSubset, ()> {
let parsed: Value = serde_json::from_str(raw).map_err(|_| ())?;
let object = parsed.as_object().ok_or(())?;
let mut headers = BTreeMap::new();
let mut query = BTreeMap::new();
let mut path = None;
parse_local_auth_config_object(object, &mut headers, &mut query, &mut path, true)?;
Ok(LocalAuthConfigSafeSubset {
headers,
query,
path,
})
}
fn parse_local_auth_config_object(
object: &serde_json::Map<String, Value>,
headers: &mut BTreeMap<String, String>,
query: &mut BTreeMap<String, String>,
path: &mut Option<String>,
allow_metadata: bool,
) -> Result<(), ()> {
for (key, value) in object {
let normalized = key.trim().to_ascii_lowercase();
match normalized.as_str() {
"headers" | "extra_headers" | "extraheaders" => {
merge_string_map(headers, value, normalize_auth_config_header_name)?
}
"query" | "query_params" | "queryparams" => {
merge_string_map(query, value, normalize_auth_config_query_key)?
}
"path" | "custom_path" => {
let value = value.as_str().ok_or(())?;
let normalized = normalize_auth_config_path(value).ok_or(())?;
*path = Some(normalized);
}
"custompath" => {
let value = value.as_str().ok_or(())?;
let normalized = normalize_auth_config_path(value).ok_or(())?;
*path = Some(normalized);
}
"transport" | "request" => {
let nested = value.as_object().ok_or(())?;
parse_local_auth_config_object(nested, headers, query, path, false)?;
}
_ if allow_metadata && is_ignorable_auth_config_metadata_key(&normalized) => {}
_ if allow_metadata && is_sensitive_auth_config_key(&normalized) => return Err(()),
_ => return Err(()),
}
}
Ok(())
}
fn merge_string_map(
out: &mut BTreeMap<String, String>,
value: &Value,
normalize_key: fn(&str) -> Option<String>,
) -> Result<(), ()> {
let object = value.as_object().ok_or(())?;
for (raw_key, raw_value) in object {
let key = normalize_key(raw_key).ok_or(())?;
let value = parse_static_auth_config_value(raw_value).ok_or(())?;
out.insert(key, value);
}
Ok(())
}
fn parse_static_auth_config_value(value: &Value) -> Option<String> {
match value {
Value::String(raw) => {
let normalized = raw.trim();
if normalized.is_empty() {
None
} else {
Some(normalized.to_string())
}
}
Value::Number(raw) => Some(raw.to_string()),
Value::Bool(raw) => Some(raw.to_string()),
_ => None,
}
}
fn normalize_auth_config_header_name(raw: &str) -> Option<String> {
let value = raw.trim().to_ascii_lowercase();
if value.is_empty()
|| value.chars().any(|char| char.is_ascii_control())
|| UNSAFE_AUTH_CONFIG_HEADER_NAMES.contains(&value.as_str())
{
return None;
}
http::header::HeaderName::from_bytes(value.as_bytes())
.ok()
.map(|name| name.as_str().to_string())
}
fn normalize_auth_config_query_key(raw: &str) -> Option<String> {
let value = raw.trim();
if value.is_empty()
|| value.chars().any(|char| matches!(char, '&' | '=' | '#'))
|| value.chars().any(|char| char.is_ascii_control())
|| UNSAFE_AUTH_CONFIG_QUERY_NAMES
.iter()
.any(|blocked| value.eq_ignore_ascii_case(blocked))
{
return None;
}
Some(value.to_string())
}
fn is_sensitive_auth_config_key(key: &str) -> bool {
SENSITIVE_AUTH_CONFIG_KEYS
.iter()
.any(|blocked| key.eq_ignore_ascii_case(blocked))
}
fn is_ignorable_auth_config_metadata_key(key: &str) -> bool {
IGNORABLE_AUTH_CONFIG_METADATA_KEYS
.iter()
.any(|allowed| key.eq_ignore_ascii_case(allowed))
}
fn normalize_auth_config_path(raw: &str) -> Option<String> {
let value = raw.trim();
if value.is_empty()
|| !value.starts_with('/')
|| value.contains("://")
|| value
.chars()
.any(|char| matches!(char, '{' | '}' | '$' | '#'))
|| value.chars().any(|char| char.is_ascii_control())
{
return None;
}
Some(value.to_string())
}
fn merge_auth_config_header_rules(
existing_rules: Option<Value>,
headers: &BTreeMap<String, String>,
) -> Option<Option<Value>> {
if headers.is_empty() {
return Some(existing_rules);
}
let mut merged = match existing_rules {
Some(Value::Array(items)) => items,
Some(_) => return None,
None => Vec::new(),
};
for (key, value) in headers {
merged.push(serde_json::json!({
"action": "set",
"key": key,
"value": value,
}));
}
Some(Some(Value::Array(merged)))
}
fn merge_auth_config_custom_path(
existing_custom_path: Option<String>,
path_override: Option<String>,
) -> Option<Option<String>> {
let base_path = path_override.or(existing_custom_path);
let Some(base_path) = base_path else {
return Some(None);
};
let (path_only, query) = split_path_and_query(&base_path)?;
if query.is_empty() {
return Some(Some(path_only));
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in query {
serializer.append_pair(&key, &value);
}
Some(Some(format!("{path_only}?{}", serializer.finish())))
}
fn merge_auth_config_base_url(base_url: &str, query: &BTreeMap<String, String>) -> Option<String> {
if query.is_empty() {
return Some(base_url.to_string());
}
let raw_base_url = base_url.trim();
let had_implicit_root = raw_base_url
.split_once("://")
.map(|(_, rest)| {
let authority = rest.split_once('?').map(|(head, _)| head).unwrap_or(rest);
!authority.contains('/')
})
.unwrap_or(false);
let mut url = url::Url::parse(raw_base_url).ok()?;
let mut merged = BTreeMap::new();
for (key, value) in url.query_pairs() {
let value = value.trim();
if value.is_empty() {
return None;
}
merged.insert(key.into_owned(), value.to_string());
}
for (key, value) in query {
merged.insert(key.clone(), value.clone());
}
if merged.is_empty() {
url.set_query(None);
return Some(url.to_string());
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in merged {
serializer.append_pair(&key, &value);
}
url.set_query(Some(&serializer.finish()));
let mut normalized = url.to_string();
if had_implicit_root {
normalized = normalized.replacen("/?", "?", 1);
}
Some(normalized)
}
fn split_path_and_query(path: &str) -> Option<(String, BTreeMap<String, String>)> {
let normalized = normalize_auth_config_path(path)?;
let (path_only, query_part) = if let Some((path, query)) = normalized.split_once('?') {
(path.to_string(), Some(query))
} else {
(normalized, None)
};
if path_only.is_empty() {
return None;
}
let mut query = BTreeMap::new();
if let Some(query_part) = query_part.filter(|value| !value.trim().is_empty()) {
for (key, value) in form_urlencoded::parse(query_part.as_bytes()) {
let key = normalize_auth_config_query_key(key.as_ref())?;
let value = value.trim();
if value.is_empty() {
return None;
}
query.insert(key, value.to_string());
}
}
Some((path_only, query))
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{absorb_local_auth_config_safe_subset, LocalAuthConfigAbsorption};
#[test]
fn absorbs_static_headers_and_query_into_existing_transport_fields() {
let result = absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
Some(json!([{"action":"set","key":"x-base","value":"1"}])),
None,
Some(
r#"{
"headers": {"x-account-id": "acc-1"},
"query": {"tenant": "demo"}
}"#,
),
);
let LocalAuthConfigAbsorption::Absorbed {
base_url,
header_rules,
custom_path,
} = result
else {
panic!("auth_config should be absorbed");
};
assert_eq!(base_url, "https://api.openai.example/v1?tenant=demo");
assert_eq!(
header_rules,
Some(json!([
{"action":"set","key":"x-base","value":"1"},
{"action":"set","key":"x-account-id","value":"acc-1"}
]))
);
assert_eq!(custom_path, None);
}
#[test]
fn absorbs_path_override_and_query_aliases() {
let result = absorb_local_auth_config_safe_subset(
"https://generativelanguage.googleapis.com/v1beta",
None,
Some("/v1beta/models/original:generateContent".to_string()),
Some(
r#"{
"extra_headers": {"x-tenant": "demo"},
"query_params": {"alt": "sse"},
"custom_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent"
}"#,
),
);
let LocalAuthConfigAbsorption::Absorbed {
base_url,
header_rules,
custom_path,
} = result
else {
panic!("auth_config should be absorbed");
};
assert_eq!(
base_url,
"https://generativelanguage.googleapis.com/v1beta?alt=sse"
);
assert_eq!(
header_rules,
Some(json!([{"action":"set","key":"x-tenant","value":"demo"}]))
);
assert_eq!(
custom_path.as_deref(),
Some("/v1beta/models/gemini-2.5-pro:streamGenerateContent")
);
}
#[test]
fn rejects_unknown_keys_and_reserved_headers() {
assert_eq!(
absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(r#"{"provider_type":"custom"}"#),
),
LocalAuthConfigAbsorption::Unsupported
);
assert_eq!(
absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(r#"{"headers":{"authorization":"Bearer x"}}"#),
),
LocalAuthConfigAbsorption::Unsupported
);
assert_eq!(
absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(r#"{"query":{"key":"secret"}}"#),
),
LocalAuthConfigAbsorption::Unsupported
);
}
#[test]
fn absorbs_query_only_configs_into_base_url_for_dynamic_path_formats() {
let result = absorb_local_auth_config_safe_subset(
"https://generativelanguage.googleapis.com/v1beta",
None,
None,
Some(r#"{"query":{"alt":"sse"}}"#),
);
let LocalAuthConfigAbsorption::Absorbed {
base_url,
header_rules,
custom_path,
} = result
else {
panic!("query-only auth_config should be absorbed");
};
assert_eq!(
base_url,
"https://generativelanguage.googleapis.com/v1beta?alt=sse"
);
assert_eq!(header_rules, None);
assert_eq!(custom_path, None);
}
#[test]
fn absorbs_camel_case_transport_keys_with_ignorable_metadata() {
let result = absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(
r#"{
"email": "user@example.com",
"plan_type": "plus",
"request": {
"extraHeaders": {"x-org-id": "org-1"},
"queryParams": {"tenant": "demo", "retry": 2, "stream": true},
"customPath": "/v1/responses"
}
}"#,
),
);
let LocalAuthConfigAbsorption::Absorbed {
base_url,
header_rules,
custom_path,
} = result
else {
panic!("camelCase auth_config should be absorbed");
};
assert_eq!(
base_url,
"https://api.openai.example/v1?retry=2&stream=true&tenant=demo"
);
assert_eq!(
header_rules,
Some(json!([{"action":"set","key":"x-org-id","value":"org-1"}]))
);
assert_eq!(custom_path.as_deref(), Some("/v1/responses"));
}
#[test]
fn rejects_sensitive_oauth_fields_even_with_transport_subset() {
assert_eq!(
absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(
r#"{
"headers": {"x-org-id": "org-1"},
"refresh_token": "rt-1"
}"#,
),
),
LocalAuthConfigAbsorption::Unsupported
);
assert_eq!(
absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(
r#"{
"query": {"tenant": "demo"},
"access_token": "at-1"
}"#,
),
),
LocalAuthConfigAbsorption::Unsupported
);
}
#[test]
fn rejects_metadata_only_auth_config_without_transport_subset() {
assert_eq!(
absorb_local_auth_config_safe_subset(
"https://api.openai.example/v1",
None,
None,
Some(
r#"{
"email": "user@example.com",
"plan_type": "plus",
"workspace_name": "demo"
}"#,
),
),
LocalAuthConfigAbsorption::Unsupported
);
}
}

View File

@@ -0,0 +1,129 @@
use super::snapshot::GatewayProviderTransportSnapshot;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct ProviderTransportSnapshotCacheKey {
provider_id: String,
endpoint_id: String,
key_id: String,
}
impl ProviderTransportSnapshotCacheKey {
pub fn new(provider_id: &str, endpoint_id: &str, key_id: &str) -> Option<Self> {
let provider_id = provider_id.trim();
let endpoint_id = endpoint_id.trim();
let key_id = key_id.trim();
if provider_id.is_empty() || endpoint_id.is_empty() || key_id.is_empty() {
return None;
}
Some(Self {
provider_id: provider_id.to_string(),
endpoint_id: endpoint_id.to_string(),
key_id: key_id.to_string(),
})
}
}
pub fn provider_transport_snapshot_looks_refreshed(
current: &GatewayProviderTransportSnapshot,
refreshed: &GatewayProviderTransportSnapshot,
) -> bool {
current.key.decrypted_api_key != refreshed.key.decrypted_api_key
|| current.key.decrypted_auth_config != refreshed.key.decrypted_auth_config
|| current.key.expires_at_unix_secs != refreshed.key.expires_at_unix_secs
}
#[cfg(test)]
mod tests {
use super::{provider_transport_snapshot_looks_refreshed, ProviderTransportSnapshotCacheKey};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
fn sample_snapshot() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider".to_string(),
provider_type: "openai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai".to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://example.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "Key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: Some(1),
proxy: None,
fingerprint: None,
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: Some("{\"token\":\"x\"}".to_string()),
},
}
}
#[test]
fn cache_key_requires_non_empty_segments() {
assert!(ProviderTransportSnapshotCacheKey::new("provider", "endpoint", "key").is_some());
assert!(ProviderTransportSnapshotCacheKey::new("", "endpoint", "key").is_none());
assert!(ProviderTransportSnapshotCacheKey::new("provider", " ", "key").is_none());
assert!(ProviderTransportSnapshotCacheKey::new("provider", "endpoint", "").is_none());
}
#[test]
fn refresh_detection_tracks_key_material_and_expiry() {
let current = sample_snapshot();
let mut refreshed = current.clone();
assert!(!provider_transport_snapshot_looks_refreshed(
&current, &refreshed
));
refreshed.key.decrypted_api_key = "sk-updated".to_string();
assert!(provider_transport_snapshot_looks_refreshed(
&current, &refreshed
));
let mut refreshed = current.clone();
refreshed.key.decrypted_auth_config = Some("{\"token\":\"y\"}".to_string());
assert!(provider_transport_snapshot_looks_refreshed(
&current, &refreshed
));
let mut refreshed = current.clone();
refreshed.key.expires_at_unix_secs = Some(2);
assert!(provider_transport_snapshot_looks_refreshed(
&current, &refreshed
));
}
}

View File

@@ -0,0 +1,9 @@
mod auth;
mod policy;
mod request;
mod url;
pub use auth::supports_local_claude_code_auth;
pub use policy::supports_local_claude_code_transport_with_network;
pub use request::{build_claude_code_passthrough_headers, sanitize_claude_code_request_body};
pub use url::build_claude_code_messages_url;

View File

@@ -0,0 +1,8 @@
use super::super::auth::resolve_local_standard_auth;
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::super::supports_local_oauth_request_auth_resolution;
pub fn supports_local_claude_code_auth(transport: &GatewayProviderTransportSnapshot) -> bool {
resolve_local_standard_auth(transport).is_some()
|| supports_local_oauth_request_auth_resolution(transport)
}

View File

@@ -0,0 +1,53 @@
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::super::{
body_rules_are_locally_supported, header_rules_are_locally_supported,
resolve_transport_tls_profile, supports_local_oauth_request_auth_resolution,
transport_proxy_is_locally_supported,
};
use super::auth::supports_local_claude_code_auth;
pub fn supports_local_claude_code_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
) -> bool {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return false;
}
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code")
{
return false;
}
if !transport
.endpoint
.api_format
.trim()
.eq_ignore_ascii_case(api_format.trim())
{
return false;
}
if !header_rules_are_locally_supported(transport.endpoint.header_rules.as_ref())
|| !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref())
{
return false;
}
if !supports_local_claude_code_auth(transport) {
return false;
}
if transport.key.decrypted_auth_config.is_some()
&& !supports_local_oauth_request_auth_resolution(transport)
{
return false;
}
if !transport_proxy_is_locally_supported(transport) {
return false;
}
if transport.key.fingerprint.is_some() && resolve_transport_tls_profile(transport).is_none() {
return false;
}
true
}

View File

@@ -0,0 +1,330 @@
use std::collections::{BTreeMap, BTreeSet};
use serde_json::{Map, Value};
use super::super::auth::build_openai_passthrough_headers;
const DEFAULT_ANTHROPIC_VERSION: &str = "2023-06-01";
const DEFAULT_ACCEPT: &str = "application/json";
const STREAM_HELPER_METHOD: &str = "stream";
const DUMMY_THINKING_SIGNATURE: &str = "skip_thought_signature_validator";
const REQUIRED_BETA_TOKENS: &[&str] = &[
"claude-code-20250219",
"oauth-2025-04-20",
"interleaved-thinking-2025-05-14",
];
const EXCLUDED_BETA_TOKENS: &[&str] = &["context-1m-2025-08-07"];
pub fn build_claude_code_passthrough_headers(
headers: &http::HeaderMap,
auth_header: &str,
auth_value: &str,
extra_headers: &BTreeMap<String, String>,
stream: bool,
fingerprint: Option<&Value>,
) -> BTreeMap<String, String> {
let mut out = build_openai_passthrough_headers(
headers,
auth_header,
auth_value,
extra_headers,
Some("application/json"),
);
out.insert("accept".to_string(), DEFAULT_ACCEPT.to_string());
out.insert(
"anthropic-version".to_string(),
DEFAULT_ANTHROPIC_VERSION.to_string(),
);
out.insert(
"anthropic-beta".to_string(),
merge_anthropic_beta_tokens(out.get("anthropic-beta").map(String::as_str)),
);
out.insert("x-stainless-lang".to_string(), "js".to_string());
out.insert(
"x-stainless-package-version".to_string(),
"0.70.0".to_string(),
);
out.insert("x-stainless-os".to_string(), "Linux".to_string());
out.insert("x-stainless-arch".to_string(), "arm64".to_string());
out.insert("x-stainless-runtime".to_string(), "node".to_string());
out.insert(
"x-stainless-runtime-version".to_string(),
"v24.13.0".to_string(),
);
out.insert("x-stainless-retry-count".to_string(), "0".to_string());
out.insert("x-stainless-timeout".to_string(), "600".to_string());
out.insert("x-app".to_string(), "cli".to_string());
out.insert(
"anthropic-dangerous-direct-browser-access".to_string(),
"true".to_string(),
);
if stream {
out.insert(
"x-stainless-helper-method".to_string(),
STREAM_HELPER_METHOD.to_string(),
);
} else {
out.remove("x-stainless-helper-method");
}
if let Some(fingerprint) = fingerprint.and_then(Value::as_object) {
override_header_from_fingerprint(
&mut out,
fingerprint,
"stainless_package_version",
"x-stainless-package-version",
);
override_header_from_fingerprint(&mut out, fingerprint, "stainless_os", "x-stainless-os");
override_header_from_fingerprint(
&mut out,
fingerprint,
"stainless_arch",
"x-stainless-arch",
);
override_header_from_fingerprint(
&mut out,
fingerprint,
"stainless_runtime_version",
"x-stainless-runtime-version",
);
override_header_from_fingerprint(
&mut out,
fingerprint,
"stainless_timeout",
"x-stainless-timeout",
);
override_header_from_fingerprint(&mut out, fingerprint, "user_agent", "user-agent");
}
out
}
pub fn sanitize_claude_code_request_body(body: &mut Value) {
let Some(body_object) = body.as_object_mut() else {
return;
};
let thinking_enabled = body_object
.get("thinking")
.and_then(Value::as_object)
.and_then(|thinking| thinking.get("type"))
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| matches!(value.to_ascii_lowercase().as_str(), "enabled" | "adaptive"));
let Some(messages) = body_object
.get_mut("messages")
.and_then(Value::as_array_mut)
else {
return;
};
for message in messages {
let Some(message_object) = message.as_object_mut() else {
continue;
};
let role = message_object
.get("role")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_string();
let Some(content) = message_object
.get_mut("content")
.and_then(Value::as_array_mut)
else {
continue;
};
let mut filtered = Vec::with_capacity(content.len());
for block in std::mem::take(content) {
let Value::Object(block_object) = block else {
filtered.push(block);
continue;
};
if keep_claude_code_block(&block_object, &role, thinking_enabled) {
filtered.push(Value::Object(block_object));
}
}
*content = filtered;
}
}
fn keep_claude_code_block(
block_object: &Map<String, Value>,
role: &str,
thinking_enabled: bool,
) -> bool {
let block_type = block_object
.get("type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if matches!(block_type, "thinking" | "redacted_thinking") {
let signature = block_object
.get("signature")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
return thinking_enabled
&& role.eq_ignore_ascii_case("assistant")
&& !signature.is_empty()
&& signature != DUMMY_THINKING_SIGNATURE;
}
if block_type.is_empty() && block_object.contains_key("thinking") {
return false;
}
true
}
fn merge_anthropic_beta_tokens(incoming: Option<&str>) -> String {
let mut seen = BTreeSet::new();
let mut merged = Vec::new();
for token in REQUIRED_BETA_TOKENS {
append_beta_token(&mut seen, &mut merged, token);
}
for token in incoming.unwrap_or_default().split(',') {
let token = token.trim();
if EXCLUDED_BETA_TOKENS
.iter()
.any(|excluded| token.eq_ignore_ascii_case(excluded))
{
continue;
}
append_beta_token(&mut seen, &mut merged, token);
}
merged.join(",")
}
fn append_beta_token(seen: &mut BTreeSet<String>, merged: &mut Vec<String>, token: &str) {
let normalized = token.trim();
if normalized.is_empty() {
return;
}
let key = normalized.to_ascii_lowercase();
if seen.insert(key) {
merged.push(normalized.to_string());
}
}
fn override_header_from_fingerprint(
headers: &mut BTreeMap<String, String>,
fingerprint: &Map<String, Value>,
fingerprint_key: &str,
header_key: &str,
) {
let Some(value) = fingerprint
.get(fingerprint_key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return;
};
headers.insert(header_key.to_string(), value.to_string());
}
#[cfg(test)]
mod tests {
use super::{build_claude_code_passthrough_headers, sanitize_claude_code_request_body};
use serde_json::json;
use std::collections::BTreeMap;
#[test]
fn claude_code_headers_merge_required_betas_and_stream_helper() {
let mut headers = http::HeaderMap::new();
headers.insert(
"anthropic-beta",
http::HeaderValue::from_static("context-1m-2025-08-07,custom-beta"),
);
headers.insert(
"user-agent",
http::HeaderValue::from_static("Claude-Code/Test"),
);
let built = build_claude_code_passthrough_headers(
&headers,
"authorization",
"Bearer upstream-token",
&BTreeMap::new(),
true,
Some(&json!({
"user_agent":"Claude-Code/9.9",
"stainless_package_version":"1.0.5",
"stainless_runtime_version":"v22.12.0",
"stainless_timeout":"900"
})),
);
assert_eq!(
built.get("anthropic-beta").map(String::as_str),
Some(
"claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,custom-beta"
)
);
assert_eq!(
built.get("anthropic-version").map(String::as_str),
Some("2023-06-01")
);
assert_eq!(
built.get("accept").map(String::as_str),
Some("application/json")
);
assert_eq!(
built.get("x-stainless-helper-method").map(String::as_str),
Some("stream")
);
assert_eq!(built.get("x-app").map(String::as_str), Some("cli"));
assert_eq!(
built.get("x-stainless-package-version").map(String::as_str),
Some("1.0.5")
);
assert_eq!(
built.get("x-stainless-runtime-version").map(String::as_str),
Some("v22.12.0")
);
assert_eq!(
built.get("x-stainless-timeout").map(String::as_str),
Some("900")
);
assert_eq!(
built.get("user-agent").map(String::as_str),
Some("Claude-Code/9.9")
);
assert_eq!(
built.get("authorization").map(String::as_str),
Some("Bearer upstream-token")
);
}
#[test]
fn claude_code_body_sanitizer_drops_invalid_thinking_blocks() {
let mut body = json!({
"thinking": {"type":"enabled"},
"messages": [{
"role":"assistant",
"content":[
{"type":"thinking","thinking":"keep","signature":"sig_valid"},
{"type":"thinking","thinking":"drop-empty","signature":""},
{"type":"redacted_thinking","data":"keep-redacted","signature":"sig_redacted"},
{"type":"redacted_thinking","data":"drop-no-signature"},
{"thinking":"drop-no-type"},
{"type":"text","text":"ok"}
]
}]
});
sanitize_claude_code_request_body(&mut body);
assert_eq!(
body["messages"][0]["content"],
json!([
{"type":"thinking","thinking":"keep","signature":"sig_valid"},
{"type":"redacted_thinking","data":"keep-redacted","signature":"sig_redacted"},
{"type":"text","text":"ok"}
])
);
}
}

View File

@@ -0,0 +1,81 @@
use std::collections::BTreeMap;
use url::form_urlencoded;
pub fn build_claude_code_messages_url(upstream_base_url: &str, query: Option<&str>) -> String {
let (trimmed_base_url, base_query) = split_query(upstream_base_url.trim());
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
let mut url =
if trimmed_base_url.ends_with("/v1/messages") || trimmed_base_url.ends_with("/messages") {
trimmed_base_url.to_string()
} else if trimmed_base_url.ends_with("/v1") {
format!("{trimmed_base_url}/messages")
} else {
format!("{trimmed_base_url}/v1/messages")
};
append_merged_query(&mut url, base_query, query);
url
}
fn split_query(value: &str) -> (&str, Option<&str>) {
value
.split_once('?')
.map(|(base, query)| (base, Some(query)))
.unwrap_or((value, None))
}
fn append_merged_query(url: &mut String, base_query: Option<&str>, request_query: Option<&str>) {
let Some(query) = merge_query_layers(base_query, request_query) else {
return;
};
if url.contains('?') {
url.push('&');
} else {
url.push('?');
}
url.push_str(&query);
}
fn merge_query_layers(base_query: Option<&str>, request_query: Option<&str>) -> Option<String> {
let mut merged = BTreeMap::new();
for source in [base_query, request_query] {
let Some(source) = source.map(str::trim).filter(|value| !value.is_empty()) else {
continue;
};
for (key, value) in form_urlencoded::parse(source.as_bytes()) {
merged.insert(key.into_owned(), value.into_owned());
}
}
if merged.is_empty() {
return None;
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in merged {
serializer.append_pair(&key, &value);
}
Some(serializer.finish())
}
#[cfg(test)]
mod tests {
use super::build_claude_code_messages_url;
#[test]
fn keeps_existing_messages_suffix_without_duplication() {
assert_eq!(
build_claude_code_messages_url("https://api.anthropic.com/v1/messages", None),
"https://api.anthropic.com/v1/messages"
);
}
#[test]
fn appends_messages_and_merges_query() {
assert_eq!(
build_claude_code_messages_url(
"https://api.anthropic.com/v1?beta=true",
Some("foo=bar"),
),
"https://api.anthropic.com/v1/messages?beta=true&foo=bar"
);
}
}

View File

@@ -0,0 +1,444 @@
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use serde_json::{json, Value};
use super::oauth_refresh::{
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshError,
LocalResolvedOAuthRequestAuth,
};
use super::snapshot::GatewayProviderTransportSnapshot;
const AUTH_HEADER_NAME: &str = "authorization";
const OAUTH_REFRESH_SKEW_SECS: u64 = 120;
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct GenericOAuthTemplate {
provider_type: &'static str,
token_url: &'static str,
client_id: &'static str,
client_secret: &'static str,
scopes: &'static [&'static str],
uses_json_payload: bool,
}
const GENERIC_OAUTH_TEMPLATES: &[GenericOAuthTemplate] = &[
GenericOAuthTemplate {
provider_type: "claude_code",
token_url: "https://console.anthropic.com/v1/oauth/token",
client_id: "9d1c250a-e61b-44d9-88ed-5944d1962f5e",
client_secret: "",
scopes: &["org:create_api_key", "user:profile", "user:inference"],
uses_json_payload: true,
},
GenericOAuthTemplate {
provider_type: "codex",
token_url: "https://auth.openai.com/oauth/token",
client_id: "app_EMoamEEZ73f0CkXaXp7hrann",
client_secret: "",
scopes: &["openid", "email", "profile", "offline_access"],
uses_json_payload: false,
},
GenericOAuthTemplate {
provider_type: "gemini_cli",
token_url: "https://oauth2.googleapis.com/token",
client_id: "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com",
client_secret: "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",
scopes: &[
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
"https://www.googleapis.com/auth/userinfo.profile",
],
uses_json_payload: false,
},
GenericOAuthTemplate {
provider_type: "antigravity",
token_url: "https://oauth2.googleapis.com/token",
client_id: "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com",
client_secret: "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",
scopes: &[
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
"https://www.googleapis.com/auth/userinfo.profile",
"https://www.googleapis.com/auth/cclog",
"https://www.googleapis.com/auth/experimentsandconfigs",
],
uses_json_payload: false,
},
];
pub fn supports_local_generic_oauth_request_auth_resolution(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
transport.key.auth_type.trim().eq_ignore_ascii_case("oauth")
&& template_for_provider_type(transport.provider.provider_type.as_str()).is_some()
}
#[derive(Debug, Clone, Default)]
pub struct GenericOAuthRefreshAdapter {
token_url_overrides: BTreeMap<String, String>,
}
impl GenericOAuthRefreshAdapter {
pub fn with_token_url_for_tests(
mut self,
provider_type: &str,
token_url: impl Into<String>,
) -> Self {
self.token_url_overrides
.insert(provider_type.trim().to_ascii_lowercase(), token_url.into());
self
}
fn token_url_for_template(&self, template: GenericOAuthTemplate) -> String {
self.token_url_overrides
.get(template.provider_type)
.cloned()
.unwrap_or_else(|| template.token_url.to_string())
}
fn auth_config_from_transport(transport: &GatewayProviderTransportSnapshot) -> Option<Value> {
transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| serde_json::from_str::<Value>(value).ok())
}
fn auth_config_from_entry(
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<Value> {
entry
.metadata
.as_ref()
.filter(|_| {
entry
.provider_type
.eq_ignore_ascii_case(transport.provider.provider_type.as_str())
})
.cloned()
}
fn base_auth_config(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Option<Value> {
entry
.and_then(|cached| Self::auth_config_from_entry(transport, cached))
.or_else(|| Self::auth_config_from_transport(transport))
}
fn resolve_direct_header(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
if !supports_local_generic_oauth_request_auth_resolution(transport) {
return None;
}
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() || secret == PLACEHOLDER_API_KEY {
return None;
}
let auth_config = Self::auth_config_from_transport(transport);
let refreshable = auth_config
.as_ref()
.and_then(refresh_token_from_auth_config)
.is_some();
if refreshable && auth_config_expires_soon(auth_config.as_ref()) {
return None;
}
Some(LocalResolvedOAuthRequestAuth::Header {
name: AUTH_HEADER_NAME.to_string(),
value: format!("Bearer {secret}"),
})
}
fn build_cached_entry(
&self,
template: GenericOAuthTemplate,
access_token: &str,
metadata: Value,
expires_at_unix_secs: Option<u64>,
) -> CachedOAuthEntry {
CachedOAuthEntry {
provider_type: template.provider_type.to_string(),
auth_header_name: AUTH_HEADER_NAME.to_string(),
auth_header_value: format!("Bearer {access_token}"),
expires_at_unix_secs,
metadata: Some(metadata),
}
}
}
#[async_trait]
impl LocalOAuthRefreshAdapter for GenericOAuthRefreshAdapter {
fn provider_type(&self) -> &'static str {
"generic_oauth"
}
fn supports(&self, transport: &GatewayProviderTransportSnapshot) -> bool {
supports_local_generic_oauth_request_auth_resolution(transport)
}
fn resolve_cached(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
if !entry
.provider_type
.eq_ignore_ascii_case(transport.provider.provider_type.as_str())
{
return None;
}
if expires_at_requires_refresh(entry.expires_at_unix_secs) {
return None;
}
let name = entry.auth_header_name.trim();
let value = entry.auth_header_value.trim();
if name.is_empty() || value.is_empty() {
return None;
}
Some(LocalResolvedOAuthRequestAuth::Header {
name: name.to_ascii_lowercase(),
value: value.to_string(),
})
}
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
self.resolve_direct_header(transport)
}
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool {
if !supports_local_generic_oauth_request_auth_resolution(transport) {
return false;
}
if entry
.and_then(|cached| self.resolve_cached(transport, cached))
.is_some()
|| self.resolve_direct_header(transport).is_some()
{
return false;
}
self.base_auth_config(transport, entry)
.as_ref()
.and_then(refresh_token_from_auth_config)
.is_some()
}
async fn refresh(
&self,
client: &reqwest::Client,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
let Some(template) = template_for_provider_type(transport.provider.provider_type.as_str())
else {
return Ok(None);
};
let mut metadata = self
.base_auth_config(transport, entry)
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
let Some(refresh_token) = metadata.get("refresh_token").and_then(non_empty_string) else {
return Ok(None);
};
let token_url = self.token_url_for_template(template);
let scope = (!template.scopes.is_empty()).then(|| template.scopes.join(" "));
let request = client.post(token_url);
let response = if template.uses_json_payload {
let mut body = serde_json::Map::from_iter([
(
"grant_type".to_string(),
Value::String("refresh_token".to_string()),
),
(
"client_id".to_string(),
Value::String(template.client_id.to_string()),
),
(
"refresh_token".to_string(),
Value::String(refresh_token.clone()),
),
]);
if let Some(scope) = scope.as_ref() {
body.insert("scope".to_string(), Value::String(scope.clone()));
}
request
.header("Content-Type", "application/json")
.header("Accept", "application/json")
.json(&Value::Object(body))
.send()
.await
} else {
let mut form = vec![
("grant_type", "refresh_token".to_string()),
("client_id", template.client_id.to_string()),
("refresh_token", refresh_token.clone()),
];
if let Some(scope) = scope.as_ref() {
form.push(("scope", scope.clone()));
}
if !template.client_secret.trim().is_empty() {
form.push(("client_secret", template.client_secret.to_string()));
}
request
.header("Content-Type", "application/x-www-form-urlencoded")
.header("Accept", "application/json")
.form(&form)
.send()
.await
}
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: template.provider_type,
source,
})?;
let status = response.status();
let body = response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: template.provider_type,
source,
})?;
if !status.is_success() {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: template.provider_type,
status_code: status.as_u16(),
body_excerpt: truncate_body(&body),
});
}
let payload: Value =
serde_json::from_str(&body).map_err(|_| LocalOAuthRefreshError::InvalidResponse {
provider_type: template.provider_type,
message: "generic oauth refresh returned non-json body".to_string(),
})?;
let Some(access_token) = payload.get("access_token").and_then(non_empty_string) else {
return Err(LocalOAuthRefreshError::InvalidResponse {
provider_type: template.provider_type,
message: "generic oauth refresh returned empty access_token".to_string(),
});
};
let expires_at_unix_secs = resolve_expires_at(payload.get("expires_in"));
metadata.insert(
"provider_type".to_string(),
Value::String(template.provider_type.to_string()),
);
metadata.insert("updated_at".to_string(), json!(current_unix_secs()));
if let Some(refresh_token) = payload.get("refresh_token").and_then(non_empty_string) {
metadata.insert("refresh_token".to_string(), Value::String(refresh_token));
}
if let Some(token_type) = payload.get("token_type").and_then(non_empty_string) {
metadata.insert("token_type".to_string(), Value::String(token_type));
}
if let Some(scope) = payload.get("scope").and_then(non_empty_string) {
metadata.insert("scope".to_string(), Value::String(scope));
}
match expires_at_unix_secs {
Some(expires_at_unix_secs) => {
metadata.insert("expires_at".to_string(), json!(expires_at_unix_secs));
}
None => {
metadata.remove("expires_at");
}
}
Ok(Some(self.build_cached_entry(
template,
access_token.as_str(),
Value::Object(metadata),
expires_at_unix_secs,
)))
}
}
fn template_for_provider_type(provider_type: &str) -> Option<GenericOAuthTemplate> {
let normalized = provider_type.trim();
GENERIC_OAUTH_TEMPLATES
.iter()
.find(|template| normalized.eq_ignore_ascii_case(template.provider_type))
.copied()
}
fn refresh_token_from_auth_config(auth_config: &Value) -> Option<String> {
auth_config
.as_object()
.and_then(|object| object.get("refresh_token"))
.and_then(non_empty_string)
}
fn auth_config_expires_soon(auth_config: Option<&Value>) -> bool {
expires_at_requires_refresh(
auth_config
.and_then(|value| value.as_object())
.and_then(|object| object.get("expires_at"))
.and_then(|value| parse_u64_value(Some(value))),
)
}
fn expires_at_requires_refresh(expires_at_unix_secs: Option<u64>) -> bool {
expires_at_unix_secs
.map(|expires_at_unix_secs| {
current_unix_secs() >= expires_at_unix_secs.saturating_sub(OAUTH_REFRESH_SKEW_SECS)
})
.unwrap_or(false)
}
fn resolve_expires_at(expires_in: Option<&Value>) -> Option<u64> {
parse_u64_value(expires_in).map(|expires_in| current_unix_secs().saturating_add(expires_in))
}
fn parse_u64_value(value: Option<&Value>) -> Option<u64> {
match value? {
Value::Number(number) => number.as_u64(),
Value::String(string) => string.trim().parse::<u64>().ok(),
_ => None,
}
}
fn non_empty_string(value: &Value) -> Option<String> {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default()
}
fn truncate_body(body: &str) -> String {
let body = body.trim();
if body.is_empty() {
return String::from("-");
}
body.chars().take(500).collect()
}

View File

@@ -0,0 +1,41 @@
pub fn should_skip_request_header(name: &str) -> bool {
let normalized = name.to_ascii_lowercase();
matches!(
normalized.as_str(),
"connection"
| "keep-alive"
| "proxy-authenticate"
| "proxy-authorization"
| "proxy-connection"
| "te"
| "trailer"
| "transfer-encoding"
| "upgrade"
| "x-aether-execution-path"
| "x-aether-dependency-reason"
| "x-aether-control-execute-fallback"
| "x-aether-rate-limit-preflight"
)
}
pub fn should_skip_upstream_passthrough_header(name: &str) -> bool {
matches!(
name.to_ascii_lowercase().as_str(),
"authorization"
| "x-api-key"
| "x-goog-api-key"
| "host"
| "content-length"
| "transfer-encoding"
| "connection"
| "accept-encoding"
| "content-encoding"
| "x-real-ip"
| "x-real-proto"
| "x-forwarded-for"
| "x-forwarded-proto"
| "x-forwarded-scheme"
| "x-forwarded-host"
| "x-forwarded-port"
) || should_skip_request_header(name)
}

View File

@@ -0,0 +1,31 @@
mod auth;
mod converter;
mod credentials;
mod headers;
mod policy;
mod refresh;
mod request;
mod url;
pub use auth::{
build_kiro_request_auth_from_config, resolve_local_kiro_bearer_auth,
resolve_local_kiro_request_auth, supports_local_kiro_auth_prerequisites,
supports_local_kiro_request_auth_resolution, KiroBearerAuth, KiroRequestAuth, KIRO_AUTH_HEADER,
PROVIDER_TYPE,
};
pub use converter::convert_claude_messages_to_conversation_state;
pub use credentials::{generate_machine_id, normalize_machine_id, KiroAuthConfig};
pub use headers::{build_generate_assistant_headers, AWS_EVENTSTREAM_CONTENT_TYPE};
pub use policy::{
supports_local_kiro_request_transport, supports_local_kiro_request_transport_with_network,
};
pub use refresh::KiroOAuthRefreshAdapter;
pub use request::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
build_kiro_provider_headers, build_kiro_provider_request_body,
header_rules_are_locally_supported, supports_local_kiro_request_shape,
};
pub use url::{
build_kiro_generate_assistant_response_url, resolve_kiro_base_url,
GENERATE_ASSISTANT_RESPONSE_PATH, KIRO_ENVELOPE_NAME,
};

View File

@@ -0,0 +1,311 @@
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::credentials::{generate_machine_id, KiroAuthConfig};
pub const PROVIDER_TYPE: &str = "kiro";
pub const KIRO_AUTH_HEADER: &str = "authorization";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KiroBearerAuth {
pub name: &'static str,
pub value: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KiroRequestAuth {
pub name: &'static str,
pub value: String,
pub auth_config: KiroAuthConfig,
pub machine_id: String,
}
pub fn build_kiro_request_auth_from_config(
auth_config: KiroAuthConfig,
fallback_secret: Option<&str>,
) -> Option<KiroRequestAuth> {
let fallback_secret = fallback_secret
.map(str::trim)
.filter(|value| !value.is_empty() && *value != "__placeholder__");
let token = auth_config
.cached_access_token()
.filter(|_| !auth_config.cached_access_token_requires_refresh(120))
.or(fallback_secret)?;
let machine_id = generate_machine_id(&auth_config, Some(token))?;
Some(KiroRequestAuth {
name: KIRO_AUTH_HEADER,
value: format!("Bearer {token}"),
auth_config,
machine_id,
})
}
pub fn resolve_local_kiro_bearer_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<KiroBearerAuth> {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(PROVIDER_TYPE)
{
return None;
}
if transport.key.decrypted_auth_config.is_some() {
return None;
}
if !transport
.key
.auth_type
.trim()
.eq_ignore_ascii_case("bearer")
{
return None;
}
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
Some(KiroBearerAuth {
name: KIRO_AUTH_HEADER,
value: format!("Bearer {secret}"),
})
}
pub fn supports_local_kiro_auth_prerequisites(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
resolve_local_kiro_bearer_auth(transport).is_some()
}
pub fn resolve_local_kiro_request_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<KiroRequestAuth> {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(PROVIDER_TYPE)
{
return None;
}
if !transport
.key
.auth_type
.trim()
.eq_ignore_ascii_case("bearer")
{
return None;
}
let auth_config = KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref())
.unwrap_or(KiroAuthConfig {
auth_method: None,
refresh_token: None,
expires_at: None,
profile_arn: None,
region: None,
auth_region: None,
api_region: None,
client_id: None,
client_secret: None,
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: None,
});
let fallback_secret = transport
.key
.decrypted_api_key
.trim()
.strip_prefix("__placeholder__")
.map(|_| "")
.unwrap_or(transport.key.decrypted_api_key.trim());
build_kiro_request_auth_from_config(auth_config, Some(fallback_secret))
}
pub fn supports_local_kiro_request_auth_resolution(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
resolve_local_kiro_request_auth(transport).is_some()
|| KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref())
.is_some_and(|auth_config| {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(PROVIDER_TYPE)
&& transport
.key
.auth_type
.trim()
.eq_ignore_ascii_case("bearer")
&& auth_config.can_refresh_access_token()
})
}
#[cfg(test)]
mod tests {
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
resolve_local_kiro_bearer_auth, resolve_local_kiro_request_auth,
supports_local_kiro_auth_prerequisites, supports_local_kiro_request_auth_resolution,
KIRO_AUTH_HEADER,
};
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Kiro".to_string(),
provider_type: "kiro".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "claude:cli".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("cli".to_string()),
is_active: true,
base_url: "https://kiro.example".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:cli".to_string()]),
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "upstream-key".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn resolves_bearer_auth_for_known_kiro_subset() {
let auth = resolve_local_kiro_bearer_auth(&sample_transport())
.expect("kiro bearer auth should resolve");
assert_eq!(auth.name, KIRO_AUTH_HEADER);
assert_eq!(auth.value, "Bearer upstream-key");
assert!(supports_local_kiro_auth_prerequisites(&sample_transport()));
}
#[test]
fn rejects_auth_config_subset() {
let mut transport = sample_transport();
transport.key.decrypted_auth_config = Some("{\"mode\":\"custom\"}".to_string());
assert!(resolve_local_kiro_bearer_auth(&transport).is_none());
assert!(!supports_local_kiro_auth_prerequisites(&transport));
}
#[test]
fn rejects_non_bearer_subset() {
let mut transport = sample_transport();
transport.key.auth_type = "api_key".to_string();
assert!(resolve_local_kiro_bearer_auth(&transport).is_none());
}
#[test]
fn resolves_request_auth_from_cached_access_token() {
let mut transport = sample_transport();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"access_token":"cached-token",
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"api_region":"us-west-2"
}"#
.to_string(),
);
let auth = resolve_local_kiro_request_auth(&transport)
.expect("request auth should resolve from cached token");
assert_eq!(auth.name, KIRO_AUTH_HEADER);
assert_eq!(auth.value, "Bearer cached-token");
assert_eq!(auth.auth_config.effective_api_region(), "us-west-2");
assert_eq!(
auth.machine_id,
"123e4567e89b12d3a456426614174000123e4567e89b12d3a456426614174000"
);
}
#[test]
fn skips_expired_cached_access_token_without_fallback_secret() {
let mut transport = sample_transport();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"access_token":"expired-token",
"expires_at": 1,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
}"#
.to_string(),
);
assert!(resolve_local_kiro_request_auth(&transport).is_none());
}
#[test]
fn falls_back_to_decrypted_api_key_when_cached_access_token_is_expired() {
let mut transport = sample_transport();
transport.key.decrypted_api_key = "live-upstream-token".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"access_token":"expired-token",
"expires_at": 1,
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
}"#
.to_string(),
);
let auth = resolve_local_kiro_request_auth(&transport)
.expect("request auth should fall back to decrypted api key");
assert_eq!(auth.value, "Bearer live-upstream-token");
}
#[test]
fn supports_refresh_only_resolution_without_cached_access_token() {
let mut transport = sample_transport();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
}"#
.to_string(),
);
assert!(resolve_local_kiro_request_auth(&transport).is_none());
assert!(supports_local_kiro_request_auth_resolution(&transport));
}
}

View File

@@ -0,0 +1,716 @@
use std::collections::BTreeSet;
use serde_json::{json, Map, Value};
use tracing::warn;
use uuid::Uuid;
const SYSTEM_CHUNKED_POLICY: &str = "When the Write or Edit tool has content size limits, always comply silently. Never suggest bypassing these limits via alternative tools. Never ask the user whether to switch approaches. Complete all chunked operations without commentary.";
const WRITE_TOOL_DESCRIPTION_SUFFIX: &str = "- IMPORTANT: If the content to write exceeds 150 lines, you MUST only write the first 50 lines using this tool, then use `Edit` tool to append the remaining content in chunks of no more than 50 lines each. If needed, leave a unique placeholder to help append content. Do NOT attempt to write all content at once.";
const EDIT_TOOL_DESCRIPTION_SUFFIX: &str = "- IMPORTANT: If the `new_string` content exceeds 50 lines, you MUST split it into multiple Edit calls, each replacing no more than 50 lines at a time. If used to append content, leave a unique placeholder to help append content. On the final chunk, do NOT include the placeholder.";
pub fn convert_claude_messages_to_conversation_state(
request_body: &Value,
model: &str,
) -> Option<Value> {
let model_id = model.trim();
if model_id.is_empty() {
return None;
}
let messages = request_body.get("messages")?.as_array()?;
if messages.is_empty() {
return None;
}
let conversation_id = request_body
.get("metadata")
.and_then(Value::as_object)
.and_then(|metadata| {
metadata
.get("user_id")
.or_else(|| metadata.get("userId"))
.and_then(Value::as_str)
})
.and_then(extract_session_id)
.unwrap_or_else(|| Uuid::new_v4().to_string());
let agent_continuation_id = Uuid::new_v4().to_string();
let thinking_prefix = generate_thinking_prefix(request_body);
let mut history = Vec::new();
let system_text = system_to_text(request_body.get("system"));
if !system_text.is_empty() {
history.push(json!({
"userInputMessage": {
"content": format!("{system_text}\n{SYSTEM_CHUNKED_POLICY}"),
"modelId": model_id,
"origin": "AI_EDITOR"
}
}));
history.push(json!({
"assistantResponseMessage": {
"content": "I will follow these instructions."
}
}));
}
let last_is_assistant = messages
.last()
.and_then(Value::as_object)
.and_then(|message| message.get("role"))
.and_then(Value::as_str)
.is_some_and(|role| role == "assistant");
let history_end_index = if last_is_assistant {
messages.len()
} else {
messages.len().saturating_sub(1)
};
let mut user_buffer = Vec::new();
for message in &messages[..history_end_index] {
let Some(message) = message.as_object() else {
continue;
};
match message.get("role").and_then(Value::as_str) {
Some("user") => user_buffer.push(message),
Some("assistant") => {
if let Some(user_item) = flush_user_buffer(&mut user_buffer, model_id) {
history.push(user_item);
} else if history.is_empty()
|| history
.last()
.and_then(Value::as_object)
.is_some_and(|item| item.contains_key("assistantResponseMessage"))
{
history.push(json!({
"userInputMessage": {
"content": "Continue.",
"modelId": model_id,
"origin": "AI_EDITOR"
}
}));
}
if let Some(assistant_item) = convert_assistant_message(message) {
history.push(json!({"assistantResponseMessage": assistant_item}));
}
}
_ => {}
}
}
if let Some(tail_user) = flush_user_buffer(&mut user_buffer, model_id) {
history.push(tail_user);
history.push(json!({"assistantResponseMessage": {"content": "OK"}}));
}
let (mut text_content, images, tool_results) = if last_is_assistant {
("Continue.".to_string(), Vec::new(), Vec::new())
} else {
let last = messages.last()?.as_object()?;
if last.get("role").and_then(Value::as_str) != Some("user") {
return None;
}
process_message_content(last.get("content"))
};
let mut tools = convert_tools(request_body.get("tools"));
let mut history_tool_names = BTreeSet::new();
let mut history_tool_result_ids = BTreeSet::new();
let mut history_tool_use_ids = BTreeSet::new();
for item in &history {
let Some(item) = item.as_object() else {
continue;
};
if let Some(user_input) = item.get("userInputMessage").and_then(Value::as_object) {
if let Some(results) = user_input
.get("userInputMessageContext")
.and_then(Value::as_object)
.and_then(|ctx| ctx.get("toolResults"))
.and_then(Value::as_array)
{
for result in results {
if let Some(tool_use_id) = result
.as_object()
.and_then(|result| result.get("toolUseId"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
history_tool_result_ids.insert(tool_use_id.to_string());
}
}
}
}
if let Some(assistant) = item
.get("assistantResponseMessage")
.and_then(Value::as_object)
{
if let Some(tool_uses) = assistant.get("toolUses").and_then(Value::as_array) {
for tool_use in tool_uses {
let Some(tool_use) = tool_use.as_object() else {
continue;
};
if let Some(name) = tool_use
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
history_tool_names.insert(name.to_string());
}
if let Some(tool_use_id) = tool_use
.get("toolUseId")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
history_tool_use_ids.insert(tool_use_id.to_string());
}
}
}
}
}
let existing_tool_names = tools
.iter()
.filter_map(|tool| {
tool.get("toolSpecification")
.and_then(Value::as_object)
.and_then(|spec| spec.get("name"))
.and_then(Value::as_str)
.map(|name| name.to_ascii_lowercase())
})
.collect::<BTreeSet<_>>();
for tool_name in history_tool_names {
if !existing_tool_names.contains(&tool_name.to_ascii_lowercase()) {
tools.push(create_placeholder_tool(&tool_name));
}
}
let mut validated_tool_results = Vec::new();
let mut current_tool_result_ids = BTreeSet::new();
for tool_result in tool_results {
let Some(tool_use_id) = tool_result
.get("toolUseId")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
if !history_tool_use_ids.contains(tool_use_id)
|| history_tool_result_ids.contains(tool_use_id)
{
continue;
}
current_tool_result_ids.insert(tool_use_id.to_string());
validated_tool_results.push(tool_result);
}
let orphaned_tool_use_ids = history_tool_use_ids
.difference(&history_tool_result_ids)
.filter(|tool_use_id| !current_tool_result_ids.contains(*tool_use_id))
.cloned()
.collect::<BTreeSet<_>>();
if !orphaned_tool_use_ids.is_empty() {
warn!(
"kiro: removing {} orphaned tool_use(s) from history",
orphaned_tool_use_ids.len()
);
for item in &mut history {
let Some(item) = item.as_object_mut() else {
continue;
};
let Some(assistant) = item
.get_mut("assistantResponseMessage")
.and_then(Value::as_object_mut)
else {
continue;
};
let Some(tool_uses) = assistant.get_mut("toolUses").and_then(Value::as_array_mut)
else {
continue;
};
tool_uses.retain(|tool_use| {
!tool_use
.get("toolUseId")
.and_then(Value::as_str)
.is_some_and(|tool_use_id| orphaned_tool_use_ids.contains(tool_use_id))
});
if tool_uses.is_empty() {
assistant.remove("toolUses");
}
}
}
let mut user_context = Map::new();
if !tools.is_empty() {
user_context.insert("tools".to_string(), Value::Array(tools));
}
if !validated_tool_results.is_empty() {
user_context.insert(
"toolResults".to_string(),
Value::Array(validated_tool_results),
);
}
if let Some(thinking_prefix) = thinking_prefix.as_deref() {
if !has_thinking_tags(&text_content) {
text_content = format!("{thinking_prefix}\n{text_content}");
}
}
let mut user_input = Map::new();
user_input.insert(
"userInputMessageContext".to_string(),
Value::Object(user_context),
);
user_input.insert("content".to_string(), Value::String(text_content));
user_input.insert("modelId".to_string(), Value::String(model_id.to_string()));
user_input.insert("origin".to_string(), Value::String("AI_EDITOR".to_string()));
if !images.is_empty() {
user_input.insert("images".to_string(), Value::Array(images));
}
Some(json!({
"agentContinuationId": agent_continuation_id,
"agentTaskType": "vibe",
"chatTriggerType": "MANUAL",
"currentMessage": {
"userInputMessage": Value::Object(user_input)
},
"conversationId": conversation_id,
"history": history,
}))
}
fn extract_session_id(user_id: &str) -> Option<String> {
let position = user_id.find("session_")?;
let candidate = user_id.get(position + "session_".len()..position + "session_".len() + 36)?;
(candidate.matches('-').count() == 4).then(|| candidate.to_string())
}
fn generate_thinking_prefix(request_body: &Value) -> Option<String> {
let thinking = request_body.get("thinking")?.as_object()?;
match thinking.get("type").and_then(Value::as_str).map(str::trim) {
Some("enabled") => {
let budget_tokens = thinking
.get("budget_tokens")
.and_then(Value::as_i64)
.unwrap_or_default();
Some(format!(
"<thinking_mode>enabled</thinking_mode><max_thinking_length>{budget_tokens}</max_thinking_length>"
))
}
Some("adaptive") => {
let effort = request_body
.get("output_config")
.and_then(Value::as_object)
.and_then(|cfg| cfg.get("effort"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("high");
Some(format!(
"<thinking_mode>adaptive</thinking_mode><thinking_effort>{effort}</thinking_effort>"
))
}
_ => None,
}
}
fn has_thinking_tags(content: &str) -> bool {
content.contains("<thinking_mode>") || content.contains("<max_thinking_length>")
}
fn system_to_text(system: Option<&Value>) -> String {
match system {
Some(Value::String(text)) => text.clone(),
Some(Value::Array(items)) => items
.iter()
.filter_map(|item| {
item.as_object()
.and_then(|item| item.get("text"))
.and_then(Value::as_str)
.map(ToOwned::to_owned)
})
.collect::<Vec<_>>()
.join("\n"),
_ => String::new(),
}
}
fn flush_user_buffer(user_buffer: &mut Vec<&Map<String, Value>>, model_id: &str) -> Option<Value> {
if user_buffer.is_empty() {
return None;
}
let mut parts = Vec::new();
let mut images = Vec::new();
let mut tool_results = Vec::new();
for message in user_buffer.drain(..) {
let (text, mut message_images, mut message_tool_results) =
process_message_content(message.get("content"));
if !text.is_empty() {
parts.push(text);
}
images.append(&mut message_images);
tool_results.append(&mut message_tool_results);
}
let mut payload = Map::new();
payload.insert("content".to_string(), Value::String(parts.join("\n")));
payload.insert("modelId".to_string(), Value::String(model_id.to_string()));
payload.insert("origin".to_string(), Value::String("AI_EDITOR".to_string()));
if !images.is_empty() {
payload.insert("images".to_string(), Value::Array(images));
}
if !tool_results.is_empty() {
payload.insert(
"userInputMessageContext".to_string(),
json!({"toolResults": tool_results}),
);
}
Some(json!({"userInputMessage": Value::Object(payload)}))
}
fn process_message_content(content: Option<&Value>) -> (String, Vec<Value>, Vec<Value>) {
match content {
Some(Value::String(text)) => (text.clone(), Vec::new(), Vec::new()),
Some(Value::Array(blocks)) => {
let mut text_parts = Vec::new();
let mut images = Vec::new();
let mut tool_results = Vec::new();
for block in blocks {
let Some(block) = block.as_object() else {
continue;
};
match block
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"text" => {
if let Some(text) = block.get("text").and_then(Value::as_str) {
text_parts.push(text.to_string());
}
}
"image" => {
let Some(source) = block.get("source").and_then(Value::as_object) else {
continue;
};
let Some(format) = source
.get("media_type")
.or_else(|| source.get("mediaType"))
.and_then(Value::as_str)
.and_then(image_format)
else {
continue;
};
let Some(bytes) = source.get("data").and_then(Value::as_str) else {
continue;
};
images.push(json!({
"format": format,
"source": {"bytes": bytes}
}));
}
"tool_result" => {
let Some(tool_use_id) = block
.get("tool_use_id")
.or_else(|| block.get("toolUseId"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let text = match block.get("content") {
Some(Value::String(text)) => text.clone(),
Some(Value::Array(items)) => items
.iter()
.filter_map(|item| {
item.as_object()
.filter(|item| {
item.get("type").and_then(Value::as_str) == Some("text")
})
.and_then(|item| item.get("text"))
.and_then(Value::as_str)
.map(ToOwned::to_owned)
})
.collect::<Vec<_>>()
.join("\n"),
Some(other) => {
serde_json::to_string(other).unwrap_or_else(|_| other.to_string())
}
None => String::new(),
};
let is_error = block
.get("is_error")
.or_else(|| block.get("isError"))
.and_then(Value::as_bool)
.unwrap_or(false);
tool_results.push(json!({
"toolUseId": tool_use_id,
"content": [{"text": text}],
"status": if is_error { "error" } else { "success" },
"isError": is_error,
}));
}
_ => {}
}
}
(text_parts.join(""), images, tool_results)
}
_ => (String::new(), Vec::new(), Vec::new()),
}
}
fn image_format(media_type: &str) -> Option<&'static str> {
let (prefix, suffix) = media_type.split_once('/')?;
if prefix != "image" {
return None;
}
match suffix.trim().to_ascii_lowercase().as_str() {
"jpeg" => Some("jpeg"),
"png" => Some("png"),
"gif" => Some("gif"),
"webp" => Some("webp"),
"jpg" => Some("jpeg"),
_ => None,
}
}
fn clean_tool_schema(value: &Value) -> Value {
match value {
Value::Object(object) => {
let mut out = Map::new();
for (key, inner) in object {
if key == "additionalProperties" {
continue;
}
if key == "required" && inner.as_array().is_some_and(|items| items.is_empty()) {
continue;
}
out.insert(key.clone(), clean_tool_schema(inner));
}
Value::Object(out)
}
Value::Array(items) => Value::Array(items.iter().map(clean_tool_schema).collect()),
_ => value.clone(),
}
}
fn convert_tools(tools: Option<&Value>) -> Vec<Value> {
let Some(tools) = tools.and_then(Value::as_array) else {
return Vec::new();
};
tools
.iter()
.filter_map(|tool| {
let tool = tool.as_object()?;
let name = tool.get("name")?.as_str()?.trim();
if name.is_empty() {
return None;
}
let mut description = tool
.get("description")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_string();
let suffix = match name {
"Write" => Some(WRITE_TOOL_DESCRIPTION_SUFFIX),
"Edit" => Some(EDIT_TOOL_DESCRIPTION_SUFFIX),
_ => None,
};
if let Some(suffix) = suffix {
description = if description.is_empty() {
suffix.to_string()
} else {
format!("{description}\n{suffix}")
};
}
if description.len() > 10_000 {
description.truncate(10_000);
}
let input_schema = tool
.get("input_schema")
.or_else(|| tool.get("inputSchema"))
.filter(|value| value.is_object())
.map(clean_tool_schema)
.unwrap_or_else(|| json!({}));
Some(json!({
"toolSpecification": {
"name": name,
"description": description,
"inputSchema": {
"json": input_schema
}
}
}))
})
.collect()
}
fn create_placeholder_tool(name: &str) -> Value {
json!({
"toolSpecification": {
"name": name,
"description": "Tool used in conversation history",
"inputSchema": {
"json": {
"type": "object",
"properties": {}
}
}
}
})
}
fn convert_assistant_message(message: &Map<String, Value>) -> Option<Value> {
let content = message.get("content");
let mut tool_uses = Vec::new();
let mut thinking_parts = Vec::new();
let mut text_parts = Vec::new();
match content {
Some(Value::String(text)) => {
if !text.is_empty() {
text_parts.push(text.clone());
}
}
Some(Value::Array(blocks)) => {
for block in blocks {
let Some(block) = block.as_object() else {
continue;
};
match block
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"thinking" => {
if let Some(thinking) = block.get("thinking").and_then(Value::as_str) {
if !thinking.is_empty() {
thinking_parts.push(thinking.to_string());
}
}
}
"text" => {
if let Some(text) = block.get("text").and_then(Value::as_str) {
if !text.is_empty() {
text_parts.push(text.to_string());
}
}
}
"tool_use" => {
let Some(tool_use_id) = block
.get("id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let Some(name) = block
.get("name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
let input = block
.get("input")
.filter(|value| value.is_object())
.cloned()
.unwrap_or_else(|| json!({}));
tool_uses.push(json!({
"toolUseId": tool_use_id,
"name": name,
"input": input
}));
}
_ => {}
}
}
}
_ => {}
}
let thinking_str = thinking_parts.join("");
let text_str = text_parts.join("");
let mut content_str = if thinking_str.is_empty() {
text_str
} else if text_str.is_empty() {
format!("<thinking>{thinking_str}</thinking>")
} else {
format!("<thinking>{thinking_str}</thinking>\n\n{text_str}")
};
if content_str.is_empty() && !tool_uses.is_empty() {
content_str = " ".to_string();
}
if content_str.is_empty() && tool_uses.is_empty() {
return None;
}
let mut out = Map::new();
out.insert("content".to_string(), Value::String(content_str));
if !tool_uses.is_empty() {
out.insert("toolUses".to_string(), Value::Array(tool_uses));
}
Some(Value::Object(out))
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::convert_claude_messages_to_conversation_state;
#[test]
fn converts_simple_claude_request_into_conversation_state() {
let conversation_state = convert_claude_messages_to_conversation_state(
&json!({
"messages": [
{"role":"user","content":"hello"}
],
"thinking": {"type": "enabled", "budget_tokens": 128},
"tools": [
{"name":"Write","description":"write file","input_schema":{"type":"object","properties":{},"required":[]}}
]
}),
"claude-sonnet-4-upstream",
)
.expect("conversation state should build");
assert_eq!(
conversation_state
.get("currentMessage")
.and_then(|value| value.get("userInputMessage"))
.and_then(|value| value.get("content"))
.and_then(|value| value.as_str()),
Some(
"<thinking_mode>enabled</thinking_mode><max_thinking_length>128</max_thinking_length>\nhello"
)
);
assert_eq!(
conversation_state
.get("currentMessage")
.and_then(|value| value.get("userInputMessage"))
.and_then(|value| value.get("userInputMessageContext"))
.and_then(|value| value.get("tools"))
.and_then(|value| value.as_array())
.map(Vec::len),
Some(1)
);
}
}

View File

@@ -0,0 +1,436 @@
use serde_json::Value;
use sha2::{Digest, Sha256};
use std::time::{SystemTime, UNIX_EPOCH};
pub const DEFAULT_REGION: &str = "us-east-1";
pub const DEFAULT_KIRO_VERSION: &str = "0.8.0";
pub const DEFAULT_NODE_VERSION: &str = "22.21.1";
pub const DEFAULT_SYSTEM_VERSION: &str = "other#unknown";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KiroAuthConfig {
pub auth_method: Option<String>,
pub refresh_token: Option<String>,
pub expires_at: Option<u64>,
pub profile_arn: Option<String>,
pub region: Option<String>,
pub auth_region: Option<String>,
pub api_region: Option<String>,
pub client_id: Option<String>,
pub client_secret: Option<String>,
pub machine_id: Option<String>,
pub kiro_version: Option<String>,
pub system_version: Option<String>,
pub node_version: Option<String>,
pub access_token: Option<String>,
}
impl KiroAuthConfig {
pub fn from_raw_json(raw: Option<&str>) -> Option<Self> {
let raw = raw?.trim();
if raw.is_empty() {
return None;
}
let parsed: Value = serde_json::from_str(raw).ok()?;
Self::from_json_value(&parsed)
}
pub fn from_json_value(raw: &Value) -> Option<Self> {
let object = raw.as_object()?;
Some(Self {
auth_method: get_nonempty_string(
object,
&["auth_method", "authMethod", "auth_type", "authType"],
)
.map(|value| normalize_auth_method(&value)),
refresh_token: get_nonempty_string(object, &["refresh_token", "refreshToken"]),
expires_at: get_epoch_seconds(object.get("expires_at"))
.or_else(|| get_epoch_seconds(object.get("expiresAt"))),
profile_arn: get_nonempty_string(object, &["profile_arn", "profileArn"]),
region: get_nonempty_string(object, &["region"]),
auth_region: get_nonempty_string(object, &["auth_region", "authRegion"]),
api_region: get_nonempty_string(object, &["api_region", "apiRegion"]),
client_id: get_nonempty_string(object, &["client_id", "clientId"]),
client_secret: get_nonempty_string(object, &["client_secret", "clientSecret"]),
machine_id: get_nonempty_string(object, &["machine_id", "machineId"]),
kiro_version: get_nonempty_string(object, &["kiro_version", "kiroVersion"]),
system_version: get_nonempty_string(object, &["system_version", "systemVersion"]),
node_version: get_nonempty_string(object, &["node_version", "nodeVersion"]),
access_token: get_nonempty_string(object, &["access_token", "accessToken"]),
})
}
pub fn to_json_value(&self) -> Value {
let mut object = serde_json::Map::new();
insert_optional_string(&mut object, "auth_method", self.auth_method.as_deref());
insert_optional_string(&mut object, "refresh_token", self.refresh_token.as_deref());
if let Some(expires_at) = self.expires_at {
object.insert("expires_at".to_string(), Value::from(expires_at));
}
insert_optional_string(&mut object, "profile_arn", self.profile_arn.as_deref());
insert_optional_string(&mut object, "region", self.region.as_deref());
insert_optional_string(&mut object, "auth_region", self.auth_region.as_deref());
insert_optional_string(&mut object, "api_region", self.api_region.as_deref());
insert_optional_string(&mut object, "client_id", self.client_id.as_deref());
insert_optional_string(&mut object, "client_secret", self.client_secret.as_deref());
insert_optional_string(&mut object, "machine_id", self.machine_id.as_deref());
insert_optional_string(&mut object, "kiro_version", self.kiro_version.as_deref());
insert_optional_string(
&mut object,
"system_version",
self.system_version.as_deref(),
);
insert_optional_string(&mut object, "node_version", self.node_version.as_deref());
insert_optional_string(&mut object, "access_token", self.access_token.as_deref());
Value::Object(object)
}
pub fn effective_api_region(&self) -> &str {
self.api_region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_auth_region(&self) -> &str {
self.auth_region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
self.region
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
})
.unwrap_or(DEFAULT_REGION)
}
pub fn effective_kiro_version(&self) -> &str {
self.kiro_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_KIRO_VERSION)
}
pub fn effective_system_version(&self) -> &str {
self.system_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_SYSTEM_VERSION)
}
pub fn effective_node_version(&self) -> &str {
self.node_version
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_NODE_VERSION)
}
pub fn cached_access_token(&self) -> Option<&str> {
self.access_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn cached_access_token_requires_refresh(&self, skew_seconds: u64) -> bool {
let Some(expires_at) = self.expires_at else {
return false;
};
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default();
now >= expires_at.saturating_sub(skew_seconds)
}
pub fn is_idc_auth(&self) -> bool {
let explicit_method = self
.auth_method
.as_deref()
.map(normalize_auth_method)
.unwrap_or_else(|| "social".to_string());
if explicit_method != "social" {
return explicit_method == "idc";
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
pub fn profile_arn_for_payload(&self) -> Option<&str> {
if self.is_idc_auth() {
return None;
}
self.profile_arn
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
}
pub fn can_refresh_access_token(&self) -> bool {
let refresh_token = self
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| value.len() >= 100 && !value.contains("..."));
if refresh_token.is_none() {
return false;
}
if !self.is_idc_auth() {
return true;
}
self.client_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
&& self
.client_secret
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.is_some()
}
}
pub fn normalize_machine_id(raw: &str) -> Option<String> {
let raw = raw.trim();
if raw.is_empty() {
return None;
}
if raw.len() == 64 && raw.bytes().all(|byte| byte.is_ascii_hexdigit()) {
return Some(raw.to_ascii_lowercase());
}
if raw.len() == 36
&& raw.chars().enumerate().all(|(idx, ch)| match idx {
8 | 13 | 18 | 23 => ch == '-',
_ => ch.is_ascii_hexdigit(),
})
{
let normalized = raw.replace('-', "").to_ascii_lowercase();
return Some(format!("{normalized}{normalized}"));
}
None
}
pub fn generate_machine_id(
auth_config: &KiroAuthConfig,
fallback_secret: Option<&str>,
) -> Option<String> {
if let Some(machine_id) = auth_config
.machine_id
.as_deref()
.and_then(normalize_machine_id)
{
return Some(machine_id);
}
let seed = auth_config
.refresh_token
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| {
fallback_secret
.map(str::trim)
.filter(|value| !value.is_empty())
})?;
let mut hasher = Sha256::new();
hasher.update(b"KotlinNativeAPI/");
hasher.update(seed.as_bytes());
Some(format!("{:x}", hasher.finalize()))
}
fn get_nonempty_string(object: &serde_json::Map<String, Value>, keys: &[&str]) -> Option<String> {
keys.iter()
.find_map(|key| object.get(*key))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn insert_optional_string(
object: &mut serde_json::Map<String, Value>,
key: &str,
value: Option<&str>,
) {
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
object.insert(key.to_string(), Value::String(value.to_string()));
}
fn get_epoch_seconds(value: Option<&Value>) -> Option<u64> {
match value? {
Value::Number(number) => number.as_u64().or_else(|| {
number
.as_i64()
.and_then(|value| (value >= 0).then_some(value as u64))
}),
Value::String(text) => text.trim().parse::<u64>().ok(),
_ => None,
}
}
fn normalize_auth_method(raw: &str) -> String {
let value = raw.trim().to_ascii_lowercase();
match value.as_str() {
"" => "social".to_string(),
"builder-id"
| "builder_id"
| "builderid"
| "device"
| "device-auth"
| "device_authorization"
| "iam"
| "identity-center"
| "identity_center"
| "identitycenter"
| "idc" => "idc".to_string(),
_ => value,
}
}
#[cfg(test)]
mod tests {
use super::{generate_machine_id, normalize_machine_id, KiroAuthConfig, DEFAULT_REGION};
#[test]
fn normalizes_uuid_machine_id() {
assert_eq!(
normalize_machine_id("123e4567-e89b-12d3-a456-426614174000").as_deref(),
Some("123e4567e89b12d3a456426614174000123e4567e89b12d3a456426614174000")
);
}
#[test]
fn hashes_refresh_token_into_machine_id() {
let auth_config = KiroAuthConfig {
auth_method: None,
refresh_token: Some("r".repeat(128)),
expires_at: None,
profile_arn: None,
region: None,
auth_region: None,
api_region: None,
client_id: None,
client_secret: None,
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: None,
};
let machine_id = generate_machine_id(&auth_config, None).expect("machine id should exist");
assert_eq!(machine_id.len(), 64);
}
#[test]
fn parses_auth_config_aliases() {
let auth_config = KiroAuthConfig::from_raw_json(Some(
r#"{
"authMethod":"identity_center",
"refreshToken":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"expires_at": 4102444800,
"profileArn":"arn:aws:bedrock:demo",
"apiRegion":"us-west-2",
"clientId":"cid",
"clientSecret":"secret",
"machineId":"123e4567-e89b-12d3-a456-426614174000",
"kiroVersion":"1.2.3",
"systemVersion":"darwin#24.6.0",
"nodeVersion":"22.21.1",
"accessToken":"cached-token"
}"#,
))
.expect("auth config should parse");
assert_eq!(auth_config.auth_method.as_deref(), Some("idc"));
assert_eq!(
auth_config.refresh_token.as_deref(),
Some(
"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
)
);
assert_eq!(auth_config.expires_at, Some(4_102_444_800));
assert_eq!(
auth_config.profile_arn.as_deref(),
Some("arn:aws:bedrock:demo")
);
assert_eq!(auth_config.client_id.as_deref(), Some("cid"));
assert_eq!(auth_config.client_secret.as_deref(), Some("secret"));
assert_eq!(auth_config.effective_api_region(), "us-west-2");
assert_eq!(auth_config.effective_kiro_version(), "1.2.3");
assert_eq!(auth_config.effective_system_version(), "darwin#24.6.0");
assert_eq!(auth_config.effective_node_version(), "22.21.1");
assert_eq!(auth_config.access_token.as_deref(), Some("cached-token"));
assert!(auth_config.is_idc_auth());
assert!(auth_config.profile_arn_for_payload().is_none());
assert_eq!(auth_config.effective_auth_region(), "us-east-1");
assert!(auth_config.can_refresh_access_token());
assert_eq!(DEFAULT_REGION, "us-east-1");
}
#[test]
fn infers_idc_when_client_credentials_exist() {
let auth_config = KiroAuthConfig::from_raw_json(Some(
r#"{
"refreshToken":"rt-1",
"clientId":"cid",
"clientSecret":"secret",
"profileArn":"arn:aws:bedrock:demo"
}"#,
))
.expect("auth config should parse");
assert!(auth_config.is_idc_auth());
assert!(auth_config.profile_arn_for_payload().is_none());
}
#[test]
fn round_trips_json_value() {
let auth_config = KiroAuthConfig::from_raw_json(Some(
r#"{
"auth_method":"social",
"refreshToken":"rt-1....................................................................................................",
"expires_at": 4102444800,
"profileArn":"arn:aws:bedrock:demo",
"region":"eu-north-1",
"apiRegion":"us-west-2",
"machineId":"123e4567-e89b-12d3-a456-426614174000",
"kiroVersion":"1.2.3",
"systemVersion":"darwin#24.6.0",
"nodeVersion":"22.21.1",
"accessToken":"cached-token"
}"#,
))
.expect("auth config should parse");
let value = auth_config.to_json_value();
let reparsed = KiroAuthConfig::from_json_value(&value).expect("auth config should reparse");
assert_eq!(reparsed, auth_config);
}
}

View File

@@ -0,0 +1,122 @@
use std::collections::BTreeMap;
use uuid::Uuid;
use super::credentials::KiroAuthConfig;
pub const AWS_EVENTSTREAM_CONTENT_TYPE: &str = "application/vnd.amazon.eventstream";
const AWS_SDK_JS_MAIN_VERSION: &str = "1.0.27";
const CODEWHISPERER_OPTOUT: &str = "true";
const KIRO_AGENT_MODE: &str = "vibe";
fn build_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
if machine_id.trim().is_empty() {
format!("KiroIDE-{kiro_version}")
} else {
format!("KiroIDE-{kiro_version}-{machine_id}")
}
}
fn build_x_amz_user_agent_main(kiro_version: &str, machine_id: &str) -> String {
format!(
"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} {}",
build_kiro_ide_tag(kiro_version, machine_id)
)
}
fn build_user_agent_main(
system_version: &str,
node_version: &str,
kiro_version: &str,
machine_id: &str,
) -> String {
format!(
"aws-sdk-js/{AWS_SDK_JS_MAIN_VERSION} ua/2.1 os/{system_version} lang/js md/nodejs#{node_version} api/codewhispererstreaming#{AWS_SDK_JS_MAIN_VERSION} m/E {}",
build_kiro_ide_tag(kiro_version, machine_id)
)
}
pub fn build_generate_assistant_headers(
auth_config: &KiroAuthConfig,
machine_id: &str,
) -> BTreeMap<String, String> {
let kiro_version = auth_config.effective_kiro_version();
let system_version = auth_config.effective_system_version();
let node_version = auth_config.effective_node_version();
let region = auth_config.effective_api_region();
let host = format!("q.{region}.amazonaws.com");
BTreeMap::from([
(
"accept".to_string(),
AWS_EVENTSTREAM_CONTENT_TYPE.to_string(),
),
(
"amz-sdk-invocation-id".to_string(),
Uuid::new_v4().to_string(),
),
(
"amz-sdk-request".to_string(),
"attempt=1; max=3".to_string(),
),
("connection".to_string(), "close".to_string()),
("content-type".to_string(), "application/json".to_string()),
("host".to_string(), host),
(
"user-agent".to_string(),
build_user_agent_main(system_version, node_version, kiro_version, machine_id),
),
(
"x-amz-user-agent".to_string(),
build_x_amz_user_agent_main(kiro_version, machine_id),
),
(
"x-amzn-codewhisperer-optout".to_string(),
CODEWHISPERER_OPTOUT.to_string(),
),
(
"x-amzn-kiro-agent-mode".to_string(),
KIRO_AGENT_MODE.to_string(),
),
])
}
#[cfg(test)]
mod tests {
use super::super::credentials::KiroAuthConfig;
use super::{build_generate_assistant_headers, AWS_EVENTSTREAM_CONTENT_TYPE};
#[test]
fn builds_generate_assistant_headers_for_region() {
let auth_config = KiroAuthConfig {
auth_method: None,
refresh_token: None,
expires_at: None,
profile_arn: None,
region: None,
auth_region: None,
api_region: Some("us-west-2".to_string()),
client_id: None,
client_secret: None,
machine_id: None,
kiro_version: Some("1.2.3".to_string()),
system_version: Some("darwin#24.6.0".to_string()),
node_version: Some("22.21.1".to_string()),
access_token: None,
};
let headers = build_generate_assistant_headers(&auth_config, "machine-123");
assert_eq!(
headers.get("accept").map(String::as_str),
Some(AWS_EVENTSTREAM_CONTENT_TYPE)
);
assert_eq!(
headers.get("host").map(String::as_str),
Some("q.us-west-2.amazonaws.com")
);
assert_eq!(
headers.get("x-amzn-kiro-agent-mode").map(String::as_str),
Some("vibe")
);
}
}

View File

@@ -0,0 +1,133 @@
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::super::{resolve_transport_tls_profile, transport_proxy_is_locally_supported};
use super::{supports_local_kiro_request_auth_resolution, supports_local_kiro_request_shape};
pub fn supports_local_kiro_request_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return false;
}
if !transport
.endpoint
.api_format
.trim()
.eq_ignore_ascii_case("claude:cli")
{
return false;
}
if !supports_local_kiro_request_auth_resolution(transport) {
return false;
}
if !supports_local_kiro_request_shape(
transport.endpoint.header_rules.as_ref(),
transport.endpoint.body_rules.as_ref(),
) {
return false;
}
true
}
pub fn supports_local_kiro_request_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_kiro_request_transport(transport)
&& transport_proxy_is_locally_supported(transport)
&& (transport.key.fingerprint.is_none()
|| resolve_transport_tls_profile(transport).is_some())
}
#[cfg(test)]
mod tests {
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
supports_local_kiro_request_transport, supports_local_kiro_request_transport_with_network,
};
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Kiro".to_string(),
provider_type: "kiro".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "claude:cli".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("cli".to_string()),
is_active: true,
base_url: "https://kiro.example".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:cli".to_string()]),
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some(
r#"{
"access_token":"cached-token",
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000"
}"#
.to_string(),
),
},
}
}
#[test]
fn supports_kiro_request_transport_when_cached_access_token_exists() {
assert!(supports_local_kiro_request_transport(&sample_transport()));
assert!(supports_local_kiro_request_transport_with_network(
&sample_transport()
));
}
#[test]
fn supports_kiro_request_transport_when_refresh_only_auth_exists() {
let mut transport = sample_transport();
transport.key.decrypted_api_key = "__placeholder__".to_string();
transport.key.decrypted_auth_config = Some(
r#"{
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr"
}"#
.to_string(),
);
assert!(supports_local_kiro_request_transport(&transport));
assert!(supports_local_kiro_request_transport_with_network(
&transport
));
}
}

View File

@@ -0,0 +1,702 @@
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use serde_json::{json, Value};
use super::super::oauth_refresh::{
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshError,
LocalResolvedOAuthRequestAuth,
};
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::auth::{
build_kiro_request_auth_from_config, resolve_local_kiro_request_auth, PROVIDER_TYPE,
};
use super::credentials::{generate_machine_id, KiroAuthConfig};
const IDC_AMZ_USER_AGENT: &str = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
#[derive(Debug, Clone, Default)]
pub struct KiroOAuthRefreshAdapter {
social_refresh_base_url: Option<String>,
idc_refresh_base_url: Option<String>,
}
impl KiroOAuthRefreshAdapter {
pub fn with_refresh_base_urls(
mut self,
social_refresh_base_url: Option<String>,
idc_refresh_base_url: Option<String>,
) -> Self {
self.social_refresh_base_url = social_refresh_base_url;
self.idc_refresh_base_url = idc_refresh_base_url;
self
}
pub async fn refresh_auth_config(
&self,
client: &reqwest::Client,
auth_config: &KiroAuthConfig,
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
if auth_config.is_idc_auth() {
self.refresh_idc_token(client, auth_config).await
} else {
self.refresh_social_token(client, auth_config).await
}
}
fn social_refresh_url(&self, auth_config: &KiroAuthConfig) -> String {
if let Some(base_url) = self
.social_refresh_base_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
return format!("{}/refreshToken", base_url.trim_end_matches('/'));
}
let region = auth_config.effective_auth_region();
format!("https://prod.{region}.auth.desktop.kiro.dev/refreshToken")
}
fn idc_refresh_url(&self, auth_config: &KiroAuthConfig) -> String {
if let Some(base_url) = self
.idc_refresh_base_url
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
return format!("{}/token", base_url.trim_end_matches('/'));
}
let region = auth_config.effective_auth_region();
format!("https://oidc.{region}.amazonaws.com/token")
}
fn auth_config_from_entry(entry: &CachedOAuthEntry) -> Option<KiroAuthConfig> {
entry
.metadata
.as_ref()
.filter(|_| entry.provider_type.eq_ignore_ascii_case(PROVIDER_TYPE))
.and_then(KiroAuthConfig::from_json_value)
}
fn base_auth_config(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Option<KiroAuthConfig> {
entry.and_then(Self::auth_config_from_entry).or_else(|| {
KiroAuthConfig::from_raw_json(transport.key.decrypted_auth_config.as_deref())
})
}
fn build_cached_entry(auth_config: &KiroAuthConfig) -> Option<CachedOAuthEntry> {
let request_auth = build_kiro_request_auth_from_config(auth_config.clone(), None)?;
Some(CachedOAuthEntry {
provider_type: PROVIDER_TYPE.to_string(),
auth_header_name: request_auth.name.to_string(),
auth_header_value: request_auth.value,
expires_at_unix_secs: auth_config.expires_at,
metadata: Some(auth_config.to_json_value()),
})
}
async fn refresh_social_token(
&self,
client: &reqwest::Client,
auth_config: &KiroAuthConfig,
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
let url = self.social_refresh_url(auth_config);
let host = reqwest::Url::parse(&url)
.ok()
.and_then(|value| value.host_str().map(ToOwned::to_owned))
.unwrap_or_else(|| {
format!(
"prod.{}.auth.desktop.kiro.dev",
auth_config.effective_auth_region()
)
});
let machine_id = generate_machine_id(auth_config, None).ok_or_else(|| {
LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "missing machine_id seed for social refresh".to_string(),
}
})?;
let kiro_version = auth_config.effective_kiro_version();
let user_agent = build_kiro_ide_tag(kiro_version, &machine_id);
let response = client
.post(url)
.header("User-Agent", user_agent)
.header("Host", host)
.header("Accept", "application/json, text/plain, */*")
.header("Content-Type", "application/json")
.header("Connection", "close")
.header("Accept-Encoding", "gzip, compress, deflate, br")
.json(&json!({
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default()
}))
.send()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
let status = response.status();
let body = response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
if !status.is_success() {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: PROVIDER_TYPE,
status_code: status.as_u16(),
body_excerpt: truncate_body(&body),
});
}
let payload: Value =
serde_json::from_str(&body).map_err(|_| LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "social refresh returned non-json body".to_string(),
})?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "social refresh returned empty accessToken".to_string(),
})?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(resolve_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id = Some(machine_id);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
if let Some(profile_arn) = payload
.get("profileArn")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.profile_arn = Some(profile_arn.to_string());
}
Ok(refreshed)
}
async fn refresh_idc_token(
&self,
client: &reqwest::Client,
auth_config: &KiroAuthConfig,
) -> Result<KiroAuthConfig, LocalOAuthRefreshError> {
let url = self.idc_refresh_url(auth_config);
let host = reqwest::Url::parse(&url)
.ok()
.and_then(|value| value.host_str().map(ToOwned::to_owned))
.unwrap_or_else(|| {
format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region())
});
let response = client
.post(url)
.header("Content-Type", "application/json")
.header("Host", host)
.header("x-amz-user-agent", IDC_AMZ_USER_AGENT)
.header("User-Agent", "node")
.header("Accept", "*/*")
.json(&json!({
"clientId": auth_config
.client_id
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"clientSecret": auth_config
.client_secret
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"grantType": "refresh_token"
}))
.send()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
let status = response.status();
let body = response
.text()
.await
.map_err(|source| LocalOAuthRefreshError::Transport {
provider_type: PROVIDER_TYPE,
source,
})?;
if !status.is_success() {
return Err(LocalOAuthRefreshError::HttpStatus {
provider_type: PROVIDER_TYPE,
status_code: status.as_u16(),
body_excerpt: truncate_body(&body),
});
}
let payload: Value =
serde_json::from_str(&body).map_err(|_| LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "idc refresh returned non-json body".to_string(),
})?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| LocalOAuthRefreshError::InvalidResponse {
provider_type: PROVIDER_TYPE,
message: "idc refresh returned empty accessToken".to_string(),
})?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(resolve_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id = generate_machine_id(auth_config, None);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
Ok(refreshed)
}
fn refreshable_auth_config(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Option<KiroAuthConfig> {
let auth_config = self.base_auth_config(transport, entry)?;
auth_config
.can_refresh_access_token()
.then_some(auth_config)
}
}
#[async_trait]
impl LocalOAuthRefreshAdapter for KiroOAuthRefreshAdapter {
fn provider_type(&self) -> &'static str {
PROVIDER_TYPE
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
let auth_config = Self::auth_config_from_entry(entry)?;
let request_auth = build_kiro_request_auth_from_config(auth_config, None)?;
Some(LocalResolvedOAuthRequestAuth::Kiro(request_auth))
}
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
resolve_local_kiro_request_auth(transport).map(LocalResolvedOAuthRequestAuth::Kiro)
}
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool {
entry
.and_then(|cached| self.resolve_cached(transport, cached))
.is_none()
&& self.resolve_without_refresh(transport).is_none()
&& self.refreshable_auth_config(transport, entry).is_some()
}
async fn refresh(
&self,
client: &reqwest::Client,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
let Some(auth_config) = self.refreshable_auth_config(transport, entry) else {
return Ok(None);
};
let refreshed = if auth_config.is_idc_auth() {
self.refresh_idc_token(client, &auth_config).await?
} else {
self.refresh_social_token(client, &auth_config).await?
};
Ok(Self::build_cached_entry(&refreshed))
}
}
fn build_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
if machine_id.trim().is_empty() {
format!("KiroIDE-{kiro_version}")
} else {
format!("KiroIDE-{kiro_version}-{machine_id}")
}
}
fn resolve_expires_at(payload: &Value) -> u64 {
let expires_in = payload
.get("expiresIn")
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_str()?.parse::<u64>().ok())
})
.unwrap_or(3600);
current_unix_secs().saturating_add(expires_in)
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default()
}
fn truncate_body(body: &str) -> String {
let body = body.trim();
if body.is_empty() {
return String::from("-");
}
body.chars().take(500).collect()
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use super::super::super::oauth_refresh::{
LocalOAuthRefreshAdapter, LocalResolvedOAuthRequestAuth,
};
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{KiroOAuthRefreshAdapter, IDC_AMZ_USER_AGENT};
use axum::body::to_bytes;
use axum::extract::Request;
use axum::response::IntoResponse;
use axum::routing::any;
use axum::{Json, Router};
use http::StatusCode;
use serde_json::{json, Value};
use tokio::task::JoinHandle;
#[derive(Debug, Clone)]
struct SeenRefreshRequest {
body: Value,
authorization: String,
host: String,
user_agent: String,
x_amz_user_agent: String,
}
fn sample_transport(raw_auth_config: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Kiro".to_string(),
provider_type: "kiro".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "claude:cli".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("cli".to_string()),
is_active: true,
base_url: "https://kiro.example".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: Some(vec!["claude:cli".to_string()]),
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some(raw_auth_config.to_string()),
},
}
}
async fn start_server(app: Router) -> (String, JoinHandle<()>) {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let addr = listener
.local_addr()
.expect("listener should expose local addr");
let handle = tokio::spawn(async move {
axum::serve(listener, app).await.expect("server should run");
});
(format!("http://{addr}"), handle)
}
#[tokio::test]
async fn refreshes_social_token_via_adapter() {
let seen_request = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
let seen_request_clone = Arc::clone(&seen_request);
let server = Router::new().route(
"/refreshToken",
any(move |request: Request| {
let seen_request_inner = Arc::clone(&seen_request_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let body: Value =
serde_json::from_slice(&raw_body).expect("body should parse as json");
*seen_request_inner.lock().expect("mutex should lock") =
Some(SeenRefreshRequest {
body,
authorization: parts
.headers
.get("authorization")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
host: parts
.headers
.get("host")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
user_agent: parts
.headers
.get("user-agent")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
x_amz_user_agent: parts
.headers
.get("x-amz-user-agent")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
(
StatusCode::OK,
Json(json!({
"accessToken": "cached-kiro-access-token",
"refreshToken": "s".repeat(120),
"expiresIn": 3600,
"profileArn": "arn:aws:bedrock:demo"
})),
)
.into_response()
}
}),
);
let (server_url, server_handle) = start_server(server).await;
let adapter =
KiroOAuthRefreshAdapter::default().with_refresh_base_urls(Some(server_url), None);
let transport = sample_transport(
r#"{
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"machine_id":"123e4567-e89b-12d3-a456-426614174000",
"kiro_version":"1.2.3"
}"#,
);
let entry = adapter
.refresh(&reqwest::Client::new(), &transport, None)
.await
.expect("refresh should succeed")
.expect("cached entry should exist");
let resolved = adapter
.resolve_cached(&transport, &entry)
.expect("cached entry should resolve");
let seen_request = seen_request
.lock()
.expect("mutex should lock")
.clone()
.expect("refresh request should be captured");
assert_eq!(seen_request.body["refreshToken"], json!("r".repeat(120)));
assert_eq!(seen_request.authorization, "");
assert!(!seen_request.user_agent.is_empty());
assert_eq!(seen_request.x_amz_user_agent, "");
assert!(!seen_request.host.trim().is_empty());
match resolved {
LocalResolvedOAuthRequestAuth::Kiro(auth) => {
assert_eq!(auth.value, "Bearer cached-kiro-access-token");
assert_eq!(
auth.auth_config.profile_arn.as_deref(),
Some("arn:aws:bedrock:demo")
);
assert!(auth.auth_config.expires_at.is_some());
}
other => panic!("unexpected resolved auth: {other:?}"),
}
server_handle.abort();
}
#[tokio::test]
async fn refreshes_idc_token_via_adapter() {
let seen_request = Arc::new(Mutex::new(None::<SeenRefreshRequest>));
let seen_request_clone = Arc::clone(&seen_request);
let server = Router::new().route(
"/token",
any(move |request: Request| {
let seen_request_inner = Arc::clone(&seen_request_clone);
async move {
let (parts, body) = request.into_parts();
let raw_body = to_bytes(body, usize::MAX).await.expect("body should read");
let body: Value =
serde_json::from_slice(&raw_body).expect("body should parse as json");
*seen_request_inner.lock().expect("mutex should lock") =
Some(SeenRefreshRequest {
body,
authorization: parts
.headers
.get("authorization")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
host: parts
.headers
.get("host")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
user_agent: parts
.headers
.get("user-agent")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
x_amz_user_agent: parts
.headers
.get("x-amz-user-agent")
.and_then(|value| value.to_str().ok())
.unwrap_or_default()
.to_string(),
});
(
StatusCode::OK,
Json(json!({
"accessToken": "cached-idc-access-token",
"refreshToken": "i".repeat(120),
"expiresIn": 1800
})),
)
.into_response()
}
}),
);
let (server_url, server_handle) = start_server(server).await;
let adapter =
KiroOAuthRefreshAdapter::default().with_refresh_base_urls(None, Some(server_url));
let transport = sample_transport(
r#"{
"auth_method":"identity_center",
"refresh_token":"rrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrrr",
"client_id":"cid",
"client_secret":"secret",
"profile_arn":"arn:aws:bedrock:demo"
}"#,
);
let entry = adapter
.refresh(&reqwest::Client::new(), &transport, None)
.await
.expect("refresh should succeed")
.expect("cached entry should exist");
let resolved = adapter
.resolve_cached(&transport, &entry)
.expect("cached entry should resolve");
let seen_request = seen_request
.lock()
.expect("mutex should lock")
.clone()
.expect("refresh request should be captured");
assert_eq!(
seen_request.body["grantType"].as_str(),
Some("refresh_token")
);
assert_eq!(seen_request.body["clientId"].as_str(), Some("cid"));
assert_eq!(seen_request.user_agent, "node");
assert_eq!(seen_request.x_amz_user_agent, IDC_AMZ_USER_AGENT);
assert!(!seen_request.host.trim().is_empty());
match resolved {
LocalResolvedOAuthRequestAuth::Kiro(auth) => {
assert_eq!(auth.value, "Bearer cached-idc-access-token");
assert!(auth.auth_config.profile_arn_for_payload().is_none());
assert!(auth.auth_config.expires_at.is_some());
}
other => panic!("unexpected resolved auth: {other:?}"),
}
server_handle.abort();
}
}

View File

@@ -0,0 +1,284 @@
use std::collections::BTreeMap;
use serde_json::{json, Value};
pub use super::super::rules::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
header_rules_are_locally_supported,
};
use super::super::should_skip_upstream_passthrough_header;
use super::converter::convert_claude_messages_to_conversation_state;
use super::credentials::KiroAuthConfig;
use super::headers::build_generate_assistant_headers;
pub fn supports_local_kiro_request_shape(
header_rules: Option<&Value>,
body_rules: Option<&Value>,
) -> bool {
header_rules_are_locally_supported(header_rules) && body_rules_are_locally_supported(body_rules)
}
pub fn build_kiro_provider_request_body(
body_json: &Value,
mapped_model: &str,
auth_config: &KiroAuthConfig,
body_rules: Option<&Value>,
) -> Option<Value> {
let conversation_state =
convert_claude_messages_to_conversation_state(body_json, mapped_model)?;
let mut provider_request_body = json!({
"conversationState": conversation_state
});
let mut inference_config = serde_json::Map::new();
if let Some(max_tokens) = body_json
.get("max_tokens")
.and_then(|value| {
value
.as_i64()
.or_else(|| value.as_u64().map(|value| value as i64))
})
.filter(|value| *value > 0)
{
inference_config.insert("maxTokens".to_string(), Value::from(max_tokens));
}
if let Some(temperature) = body_json
.get("temperature")
.and_then(Value::as_f64)
.filter(|value| *value >= 0.0)
{
inference_config.insert("temperature".to_string(), Value::from(temperature));
}
if let Some(top_p) = body_json
.get("top_p")
.and_then(Value::as_f64)
.filter(|value| *value > 0.0)
{
inference_config.insert("topP".to_string(), Value::from(top_p));
}
if !inference_config.is_empty() {
provider_request_body.as_object_mut()?.insert(
"inferenceConfig".to_string(),
Value::Object(inference_config),
);
}
if let Some(profile_arn) = auth_config.profile_arn_for_payload() {
provider_request_body.as_object_mut()?.insert(
"profileArn".to_string(),
Value::String(profile_arn.to_string()),
);
}
if !apply_local_body_rules(&mut provider_request_body, body_rules, Some(body_json)) {
return None;
}
Some(provider_request_body)
}
pub fn build_kiro_provider_headers(
headers: &http::HeaderMap,
provider_request_body: &Value,
original_request_body: &Value,
header_rules: Option<&Value>,
auth_header: &str,
auth_value: &str,
auth_config: &KiroAuthConfig,
machine_id: &str,
) -> Option<BTreeMap<String, String>> {
let mut out = BTreeMap::new();
for (name, value) in headers {
let Ok(value) = value.to_str() else {
continue;
};
let key = name.as_str().to_ascii_lowercase();
if should_skip_upstream_passthrough_header(&key) {
continue;
}
let value = value.trim();
if value.is_empty() {
continue;
}
out.insert(key, value.to_string());
}
if !apply_local_header_rules(
&mut out,
header_rules,
&[auth_header, "content-type"],
provider_request_body,
Some(original_request_body),
) {
return None;
}
for (key, value) in build_generate_assistant_headers(auth_config, machine_id) {
out.insert(key, value);
}
out.insert(
auth_header.trim().to_ascii_lowercase(),
auth_value.trim().to_string(),
);
out.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
out.remove("content-length");
Some(out)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::super::credentials::KiroAuthConfig;
use super::{
build_kiro_provider_headers, build_kiro_provider_request_body,
supports_local_kiro_request_shape,
};
#[test]
fn supports_empty_local_request_shape() {
assert!(supports_local_kiro_request_shape(None, None));
}
#[test]
fn rejects_unsupported_rule_shape() {
assert!(!supports_local_kiro_request_shape(
Some(&json!({"action":"set"})),
None
));
}
#[test]
fn supports_simple_header_and_body_rules() {
assert!(supports_local_kiro_request_shape(
Some(&json!([{"action":"set","key":"x-provider-extra","value":"1"}])),
Some(&json!([{"action":"set","path":"debugTag","value":true}]))
));
}
#[test]
fn wraps_claude_request_into_kiro_payload_before_body_rules() {
let auth_config = KiroAuthConfig {
auth_method: None,
refresh_token: Some("r".repeat(128)),
expires_at: None,
profile_arn: Some("arn:aws:bedrock:demo".to_string()),
region: None,
auth_region: None,
api_region: Some("us-east-1".to_string()),
client_id: None,
client_secret: None,
machine_id: Some("123e4567-e89b-12d3-a456-426614174000".to_string()),
kiro_version: None,
system_version: None,
node_version: None,
access_token: Some("cached-token".to_string()),
};
let payload = build_kiro_provider_request_body(
&json!({
"messages": [{"role":"user","content":"hello"}],
"max_tokens": 64
}),
"claude-sonnet-4-upstream",
&auth_config,
Some(&json!([
{"action":"set","path":"debugTag","value":"kiro-local"}
])),
)
.expect("payload should build");
assert!(payload.get("conversationState").is_some());
assert_eq!(
payload
.get("inferenceConfig")
.and_then(|value| value.get("maxTokens")),
Some(&json!(64))
);
assert_eq!(
payload.get("profileArn"),
Some(&json!("arn:aws:bedrock:demo"))
);
assert_eq!(payload.get("debugTag"), Some(&json!("kiro-local")));
}
#[test]
fn applies_header_rules_before_kiro_extra_headers() {
let auth_config = KiroAuthConfig {
auth_method: None,
refresh_token: Some("r".repeat(128)),
expires_at: None,
profile_arn: None,
region: None,
auth_region: None,
api_region: Some("us-east-1".to_string()),
client_id: None,
client_secret: None,
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: Some("cached-token".to_string()),
};
let headers = build_kiro_provider_headers(
&http::HeaderMap::new(),
&json!({"conversationState": {}}),
&json!({"messages": []}),
Some(&json!([
{"action":"set","key":"accept","value":"text/plain"},
{"action":"set","key":"x-endpoint-tag","value":"kiro-local"}
])),
"authorization",
"Bearer cached-token",
&auth_config,
"machine-123",
)
.expect("headers should build");
assert_eq!(
headers.get("accept").map(String::as_str),
Some("application/vnd.amazon.eventstream")
);
assert_eq!(
headers.get("authorization").map(String::as_str),
Some("Bearer cached-token")
);
assert_eq!(
headers.get("x-endpoint-tag").map(String::as_str),
Some("kiro-local")
);
}
#[test]
fn omits_profile_arn_for_idc_auth() {
let auth_config = KiroAuthConfig {
auth_method: Some("identity_center".to_string()),
refresh_token: Some("r".repeat(128)),
expires_at: None,
profile_arn: Some("arn:aws:bedrock:demo".to_string()),
region: None,
auth_region: None,
api_region: Some("us-east-1".to_string()),
client_id: Some("cid".to_string()),
client_secret: Some("secret".to_string()),
machine_id: None,
kiro_version: None,
system_version: None,
node_version: None,
access_token: Some("cached-token".to_string()),
};
let payload = build_kiro_provider_request_body(
&json!({
"messages": [{"role":"user","content":"hello"}]
}),
"claude-sonnet-4-upstream",
&auth_config,
None,
)
.expect("payload should build");
assert!(payload.get("profileArn").is_none());
}
}

View File

@@ -0,0 +1,71 @@
use super::super::url::build_passthrough_path_url;
use super::credentials::DEFAULT_REGION;
pub const GENERATE_ASSISTANT_RESPONSE_PATH: &str = "/generateAssistantResponse";
pub const KIRO_ENVELOPE_NAME: &str = "kiro:generateAssistantResponse";
pub fn resolve_kiro_base_url(upstream_base_url: &str, api_region: Option<&str>) -> String {
let region = api_region
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(DEFAULT_REGION);
upstream_base_url
.trim()
.replace("{region}", region)
.trim_end_matches('/')
.to_string()
}
pub fn build_kiro_generate_assistant_response_url(
upstream_base_url: &str,
query: Option<&str>,
api_region: Option<&str>,
) -> Option<String> {
let upstream_base_url = resolve_kiro_base_url(upstream_base_url, api_region);
build_passthrough_path_url(
upstream_base_url.as_str(),
GENERATE_ASSISTANT_RESPONSE_PATH,
query,
&[],
)
}
#[cfg(test)]
mod tests {
use super::{
build_kiro_generate_assistant_response_url, resolve_kiro_base_url,
GENERATE_ASSISTANT_RESPONSE_PATH, KIRO_ENVELOPE_NAME,
};
#[test]
fn exposes_kiro_request_constants() {
assert_eq!(
GENERATE_ASSISTANT_RESPONSE_PATH,
"/generateAssistantResponse"
);
assert_eq!(KIRO_ENVELOPE_NAME, "kiro:generateAssistantResponse");
}
#[test]
fn builds_generate_assistant_response_url() {
assert_eq!(
build_kiro_generate_assistant_response_url(
"https://kiro.{region}.example?tenant=demo",
Some("stream=true"),
Some("us-west-2")
)
.as_deref(),
Some(
"https://kiro.us-west-2.example/generateAssistantResponse?stream=true&tenant=demo"
)
);
}
#[test]
fn resolves_region_placeholder_in_base_url() {
assert_eq!(
resolve_kiro_base_url("https://kiro.{region}.example/", Some("us-west-2")),
"https://kiro.us-west-2.example"
);
}
}

View File

@@ -0,0 +1,50 @@
pub mod antigravity;
pub mod auth;
mod auth_config;
mod cache;
pub mod claude_code;
mod generic_oauth;
mod headers;
pub mod kiro;
mod network;
pub mod oauth_refresh;
pub mod policy;
pub mod provider_types;
pub mod rules;
pub mod snapshot;
pub mod url;
pub mod vertex;
mod video;
pub use auth::{build_passthrough_headers, ensure_upstream_auth_header};
pub use cache::{provider_transport_snapshot_looks_refreshed, ProviderTransportSnapshotCacheKey};
pub use generic_oauth::{
supports_local_generic_oauth_request_auth_resolution, GenericOAuthRefreshAdapter,
};
pub use headers::{should_skip_request_header, should_skip_upstream_passthrough_header};
pub use network::{
resolve_transport_execution_timeouts, resolve_transport_proxy_snapshot,
resolve_transport_proxy_snapshot_with_tunnel_affinity, resolve_transport_tls_profile,
transport_proxy_is_locally_supported, TransportTunnelAffinityLookup,
TransportTunnelAttachmentOwner,
};
pub use oauth_refresh::{
supports_local_oauth_request_auth_resolution, CachedOAuthEntry, LocalOAuthRefreshCoordinator,
LocalOAuthRefreshError, LocalResolvedOAuthRequestAuth,
};
pub use policy::{
supports_local_gemini_transport, supports_local_gemini_transport_with_network,
supports_local_standard_transport,
};
pub use rules::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
header_rules_are_locally_supported,
};
pub use snapshot::{
read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
ProviderTransportSnapshotSource,
};
pub use video::{
reconstruct_local_video_task_snapshot, resolve_local_video_task_transport,
VideoTaskTransportSnapshotLookup,
};

View File

@@ -0,0 +1,396 @@
use aether_contracts::{ExecutionTimeouts, ProxySnapshot};
use async_trait::async_trait;
use serde_json::{json, Map, Value};
use tracing::warn;
use super::snapshot::GatewayProviderTransportSnapshot;
const TUNNEL_BASE_URL_EXTRA_KEY: &str = "tunnel_base_url";
const TUNNEL_OWNER_INSTANCE_ID_EXTRA_KEY: &str = "tunnel_owner_instance_id";
const TUNNEL_OWNER_OBSERVED_AT_EXTRA_KEY: &str = "tunnel_owner_observed_at_unix_secs";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TransportTunnelAttachmentOwner {
pub gateway_instance_id: String,
pub relay_base_url: String,
pub observed_at_unix_secs: u64,
}
#[async_trait]
pub trait TransportTunnelAffinityLookup: Send + Sync {
async fn lookup_tunnel_attachment_owner(
&self,
node_id: &str,
) -> Result<Option<TransportTunnelAttachmentOwner>, String>;
}
pub fn resolve_transport_execution_timeouts(
transport: &GatewayProviderTransportSnapshot,
) -> Option<ExecutionTimeouts> {
let total_ms = transport
.provider
.request_timeout_secs
.filter(|value| value.is_finite() && *value > 0.0)
.map(|value| (value * 1000.0).round() as u64);
let first_byte_ms = transport
.provider
.stream_first_byte_timeout_secs
.filter(|value| value.is_finite() && *value > 0.0)
.map(|value| (value * 1000.0).round() as u64);
if total_ms.is_none() && first_byte_ms.is_none() {
return None;
}
Some(ExecutionTimeouts {
total_ms,
first_byte_ms,
..ExecutionTimeouts::default()
})
}
pub fn resolve_transport_proxy_snapshot(
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot> {
let raw = effective_proxy_config(transport)?;
proxy_snapshot_from_value(raw)
}
pub async fn resolve_transport_proxy_snapshot_with_tunnel_affinity(
lookup: &dyn TransportTunnelAffinityLookup,
transport: &GatewayProviderTransportSnapshot,
) -> Option<ProxySnapshot> {
let mut snapshot = resolve_transport_proxy_snapshot(transport)?;
let Some(node_id) = snapshot
.node_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Some(snapshot);
};
let owner = match lookup.lookup_tunnel_attachment_owner(node_id).await {
Ok(owner) => owner,
Err(error) => {
warn!(error = %error, node_id = node_id, "failed to load tunnel attachment owner");
None
}
};
let Some(owner) = owner else {
return Some(snapshot);
};
let mut extra = snapshot
.extra
.take()
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
let configured_tunnel_base_url = extra
.get(TUNNEL_BASE_URL_EXTRA_KEY)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
if configured_tunnel_base_url.is_none() {
extra.insert(
TUNNEL_BASE_URL_EXTRA_KEY.to_string(),
Value::String(owner.relay_base_url.clone()),
);
}
extra.insert(
TUNNEL_OWNER_INSTANCE_ID_EXTRA_KEY.to_string(),
Value::String(owner.gateway_instance_id),
);
extra.insert(
TUNNEL_OWNER_OBSERVED_AT_EXTRA_KEY.to_string(),
json!(owner.observed_at_unix_secs),
);
snapshot.extra = Some(Value::Object(extra));
Some(snapshot)
}
pub fn transport_proxy_is_locally_supported(transport: &GatewayProviderTransportSnapshot) -> bool {
let has_configured_proxy = transport.provider.proxy.is_some()
|| transport.endpoint.proxy.is_some()
|| transport.key.proxy.is_some();
if !has_configured_proxy {
return true;
}
let Some(snapshot) = resolve_transport_proxy_snapshot(transport) else {
return false;
};
if snapshot.enabled == Some(false) {
return true;
}
snapshot
.url
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())
|| snapshot
.node_id
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
pub fn resolve_transport_tls_profile(
transport: &GatewayProviderTransportSnapshot,
) -> Option<String> {
transport
.key
.fingerprint
.as_ref()
.and_then(|value| value.get("tls_profile"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn effective_proxy_config(transport: &GatewayProviderTransportSnapshot) -> Option<&Value> {
for candidate in [
transport.key.proxy.as_ref(),
transport.endpoint.proxy.as_ref(),
transport.provider.proxy.as_ref(),
]
.into_iter()
.flatten()
{
if proxy_enabled(candidate) {
return Some(candidate);
}
}
None
}
fn proxy_enabled(value: &Value) -> bool {
value
.as_object()
.and_then(|object| object.get("enabled"))
.and_then(Value::as_bool)
.unwrap_or(true)
}
fn proxy_snapshot_from_value(value: &Value) -> Option<ProxySnapshot> {
let object = value.as_object()?;
let enabled = object.get("enabled").and_then(Value::as_bool);
let mode = json_string_field(object, "mode");
let node_id = json_string_field(object, "node_id");
let label = json_string_field(object, "label");
let url = json_string_field(object, "url").or_else(|| json_string_field(object, "proxy_url"));
let mut extra = Map::new();
for (key, value) in object {
if matches!(
key.as_str(),
"enabled" | "mode" | "node_id" | "label" | "url" | "proxy_url"
) {
continue;
}
extra.insert(key.clone(), value.clone());
}
Some(ProxySnapshot {
enabled,
mode,
node_id,
label,
url,
extra: if extra.is_empty() {
None
} else {
Some(Value::Object(extra))
},
})
}
fn json_string_field(object: &Map<String, Value>, key: &str) -> Option<String> {
object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use async_trait::async_trait;
use serde_json::{json, Value};
use super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
resolve_transport_proxy_snapshot, resolve_transport_proxy_snapshot_with_tunnel_affinity,
resolve_transport_tls_profile, transport_proxy_is_locally_supported,
TransportTunnelAffinityLookup, TransportTunnelAttachmentOwner,
};
#[derive(Default)]
struct TestTunnelAffinityLookup {
owners: BTreeMap<String, TransportTunnelAttachmentOwner>,
}
#[async_trait]
impl TransportTunnelAffinityLookup for TestTunnelAffinityLookup {
async fn lookup_tunnel_attachment_owner(
&self,
node_id: &str,
) -> Result<Option<TransportTunnelAttachmentOwner>, String> {
Ok(self.owners.get(node_id).cloned())
}
}
fn sample_lookup() -> TestTunnelAffinityLookup {
let mut owners = BTreeMap::new();
owners.insert(
"proxy-node-1".to_string(),
TransportTunnelAttachmentOwner {
gateway_instance_id: "gateway-b".to_string(),
relay_base_url: "http://gateway-b.internal".to_string(),
observed_at_unix_secs: 4_102_444_800u64,
},
);
TestTunnelAffinityLookup { owners }
}
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "provider".to_string(),
provider_type: "custom".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: Some(json!({"url":"http://provider-proxy:8080"})),
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:chat".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://api.openai.example".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: Some(json!({"enabled":false,"url":"http://endpoint-proxy:8080"})),
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: Some(json!({"node_id":"proxy-node-1","kind":"manual"})),
fingerprint: Some(json!({"tls_profile":"chrome_136"})),
decrypted_api_key: "sk-test".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn resolves_transport_proxy_with_key_precedence() {
let snapshot = resolve_transport_proxy_snapshot(&sample_transport())
.expect("proxy snapshot should resolve");
assert_eq!(snapshot.node_id.as_deref(), Some("proxy-node-1"));
assert_eq!(snapshot.url, None);
assert_eq!(snapshot.extra, Some(json!({"kind":"manual"})));
}
#[tokio::test]
async fn enriches_transport_proxy_snapshot_with_tunnel_owner_hint() {
let state = sample_lookup();
let snapshot =
resolve_transport_proxy_snapshot_with_tunnel_affinity(&state, &sample_transport())
.await
.expect("proxy snapshot should resolve");
assert_eq!(snapshot.node_id.as_deref(), Some("proxy-node-1"));
assert_eq!(
snapshot
.extra
.as_ref()
.and_then(|value| value.get("tunnel_base_url"))
.and_then(Value::as_str),
Some("http://gateway-b.internal")
);
assert_eq!(
snapshot
.extra
.as_ref()
.and_then(|value| value.get("tunnel_owner_instance_id"))
.and_then(Value::as_str),
Some("gateway-b")
);
}
#[tokio::test]
async fn preserves_explicit_tunnel_base_url_when_owner_hint_exists() {
let mut transport = sample_transport();
transport.key.proxy = Some(json!({
"node_id": "proxy-node-1",
"kind": "manual",
"tunnel_base_url": "http://configured-gateway.internal",
}));
let state = sample_lookup();
let snapshot = resolve_transport_proxy_snapshot_with_tunnel_affinity(&state, &transport)
.await
.expect("proxy snapshot should resolve");
assert_eq!(
snapshot
.extra
.as_ref()
.and_then(|value| value.get("tunnel_base_url"))
.and_then(Value::as_str),
Some("http://configured-gateway.internal")
);
assert_eq!(
snapshot
.extra
.as_ref()
.and_then(|value| value.get("tunnel_owner_instance_id"))
.and_then(Value::as_str),
Some("gateway-b")
);
}
#[test]
fn resolves_transport_tls_profile_from_key_fingerprint() {
assert_eq!(
resolve_transport_tls_profile(&sample_transport()).as_deref(),
Some("chrome_136")
);
assert!(transport_proxy_is_locally_supported(&sample_transport()));
}
}

View File

@@ -0,0 +1,470 @@
use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use aether_data::redis::{RedisLockKey, RedisLockRunner};
use async_trait::async_trait;
use serde_json::Value;
use thiserror::Error;
use tokio::sync::Mutex;
use super::generic_oauth::supports_local_generic_oauth_request_auth_resolution;
pub use super::generic_oauth::GenericOAuthRefreshAdapter;
use super::kiro::{
supports_local_kiro_request_auth_resolution, KiroOAuthRefreshAdapter, KiroRequestAuth,
};
use super::snapshot::GatewayProviderTransportSnapshot;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum LocalResolvedOAuthRequestAuth {
#[allow(dead_code)]
Header {
name: String,
value: String,
},
Kiro(KiroRequestAuth),
}
#[derive(Debug, Clone, PartialEq)]
pub struct LocalOAuthResolution {
pub auth: Option<LocalResolvedOAuthRequestAuth>,
pub refreshed_entry: Option<CachedOAuthEntry>,
pub refresh_in_flight: bool,
}
#[derive(Debug, Clone, PartialEq)]
pub struct CachedOAuthEntry {
pub provider_type: String,
pub auth_header_name: String,
pub auth_header_value: String,
pub expires_at_unix_secs: Option<u64>,
pub metadata: Option<Value>,
}
#[derive(Debug, Error)]
pub enum LocalOAuthRefreshError {
#[error("{provider_type} oauth refresh request failed: {source}")]
Transport {
provider_type: &'static str,
#[source]
source: reqwest::Error,
},
#[error("{provider_type} oauth refresh returned HTTP {status_code}: {body_excerpt}")]
HttpStatus {
provider_type: &'static str,
status_code: u16,
body_excerpt: String,
},
#[error("{provider_type} oauth refresh returned invalid response: {message}")]
InvalidResponse {
provider_type: &'static str,
message: String,
},
}
#[async_trait]
pub trait LocalOAuthRefreshAdapter: Send + Sync {
fn provider_type(&self) -> &'static str;
fn supports(&self, transport: &GatewayProviderTransportSnapshot) -> bool {
transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(self.provider_type())
}
fn resolve_cached(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth>;
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth>;
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool;
async fn refresh(
&self,
client: &reqwest::Client,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError>;
}
pub struct LocalOAuthRefreshCoordinator {
adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>,
cache: Mutex<BTreeMap<String, CachedOAuthEntry>>,
key_locks: Mutex<BTreeMap<String, Arc<Mutex<()>>>>,
}
impl fmt::Debug for LocalOAuthRefreshCoordinator {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("LocalOAuthRefreshCoordinator")
.field("adapter_count", &self.adapters.len())
.finish()
}
}
impl Default for LocalOAuthRefreshCoordinator {
fn default() -> Self {
Self::new()
}
}
impl LocalOAuthRefreshCoordinator {
const DISTRIBUTED_REFRESH_LOCK_TTL_MS: u64 = 30_000;
pub fn new() -> Self {
Self {
adapters: vec![
Arc::new(KiroOAuthRefreshAdapter::default()),
Arc::new(GenericOAuthRefreshAdapter::default()),
],
cache: Mutex::new(BTreeMap::new()),
key_locks: Mutex::new(BTreeMap::new()),
}
}
async fn lock_for_key(&self, key_id: &str) -> Arc<Mutex<()>> {
let mut key_locks = self.key_locks.lock().await;
key_locks
.entry(key_id.to_string())
.or_insert_with(|| Arc::new(Mutex::new(())))
.clone()
}
async fn cached_entry(&self, key_id: &str) -> Option<CachedOAuthEntry> {
self.cache.lock().await.get(key_id).cloned()
}
async fn insert_cached_entry(&self, key_id: &str, entry: CachedOAuthEntry) {
self.cache.lock().await.insert(key_id.to_string(), entry);
}
pub async fn resolve_with_result(
&self,
client: &reqwest::Client,
transport: &GatewayProviderTransportSnapshot,
distributed_lock: Option<&RedisLockRunner>,
distributed_owner: Option<&str>,
) -> Result<Option<LocalOAuthResolution>, LocalOAuthRefreshError> {
let Some(adapter) = self
.adapters
.iter()
.find(|adapter| adapter.supports(transport))
else {
return Ok(None);
};
let key_id = transport.key.id.trim();
let cached_entry = if key_id.is_empty() {
None
} else {
self.cached_entry(key_id).await
};
if let Some(auth) = cached_entry
.as_ref()
.and_then(|entry| adapter.resolve_cached(transport, entry))
{
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if let Some(auth) = adapter.resolve_without_refresh(transport) {
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if !adapter.should_refresh(transport, cached_entry.as_ref()) {
return Ok(None);
}
if key_id.is_empty() {
return Ok(None);
}
let key_lock = self.lock_for_key(key_id).await;
let _key_guard = key_lock.lock().await;
let cached_entry = self.cached_entry(key_id).await;
if let Some(auth) = cached_entry
.as_ref()
.and_then(|entry| adapter.resolve_cached(transport, entry))
{
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if let Some(auth) = adapter.resolve_without_refresh(transport) {
return Ok(Some(LocalOAuthResolution::resolved(auth, None)));
}
if !adapter.should_refresh(transport, cached_entry.as_ref()) {
return Ok(None);
}
let distributed_lease = match (distributed_lock, distributed_owner) {
(Some(lock), Some(owner)) if !owner.trim().is_empty() => {
let lock_key = RedisLockKey(format!("provider_oauth_refresh_lock:{key_id}"));
match lock
.try_acquire(
&lock_key,
owner,
Some(Self::DISTRIBUTED_REFRESH_LOCK_TTL_MS),
)
.await
{
Ok(Some(lease)) => Some(lease),
Ok(None) => return Ok(Some(LocalOAuthResolution::refresh_in_flight())),
Err(err) => {
tracing::warn!(
key_id = %key_id,
provider_type = adapter.provider_type(),
error = ?err,
"gateway local oauth refresh distributed lock unavailable"
);
None
}
}
}
_ => None,
};
let refresh_result = adapter
.refresh(client, transport, cached_entry.as_ref())
.await;
if let (Some(lock), Some(lease)) = (distributed_lock, distributed_lease.as_ref()) {
if let Err(err) = lock.release(lease).await {
tracing::warn!(
key_id = %key_id,
provider_type = adapter.provider_type(),
error = ?err,
"gateway local oauth refresh distributed lock release failed"
);
}
}
let Some(refreshed_entry) = refresh_result? else {
return Ok(None);
};
self.insert_cached_entry(key_id, refreshed_entry.clone())
.await;
Ok(adapter
.resolve_cached(transport, &refreshed_entry)
.map(|auth| LocalOAuthResolution::resolved(auth, Some(refreshed_entry))))
}
pub fn with_adapters_for_tests(adapters: Vec<Arc<dyn LocalOAuthRefreshAdapter>>) -> Self {
Self {
adapters,
cache: Mutex::new(BTreeMap::new()),
key_locks: Mutex::new(BTreeMap::new()),
}
}
}
impl LocalOAuthResolution {
fn resolved(
auth: LocalResolvedOAuthRequestAuth,
refreshed_entry: Option<CachedOAuthEntry>,
) -> Self {
Self {
auth: Some(auth),
refreshed_entry,
refresh_in_flight: false,
}
}
fn refresh_in_flight() -> Self {
Self {
auth: None,
refreshed_entry: None,
refresh_in_flight: true,
}
}
}
pub fn supports_local_oauth_request_auth_resolution(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_kiro_request_auth_resolution(transport)
|| supports_local_generic_oauth_request_auth_resolution(transport)
}
#[cfg(test)]
mod tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
CachedOAuthEntry, LocalOAuthRefreshAdapter, LocalOAuthRefreshCoordinator,
LocalOAuthRefreshError, LocalOAuthResolution, LocalResolvedOAuthRequestAuth,
};
use async_trait::async_trait;
use std::sync::Arc;
#[derive(Debug)]
struct TestAdapter {
refresh_hits: Arc<AtomicUsize>,
}
#[async_trait]
impl LocalOAuthRefreshAdapter for TestAdapter {
fn provider_type(&self) -> &'static str {
"test-oauth"
}
fn resolve_cached(
&self,
_transport: &GatewayProviderTransportSnapshot,
entry: &CachedOAuthEntry,
) -> Option<LocalResolvedOAuthRequestAuth> {
(entry.provider_type == "test-oauth").then(|| LocalResolvedOAuthRequestAuth::Header {
name: entry.auth_header_name.clone(),
value: entry.auth_header_value.clone(),
})
}
fn resolve_without_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
) -> Option<LocalResolvedOAuthRequestAuth> {
let secret = transport.key.decrypted_api_key.trim();
(!secret.is_empty() && secret != "__placeholder__").then(|| {
LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: format!("Bearer {secret}"),
}
})
}
fn should_refresh(
&self,
transport: &GatewayProviderTransportSnapshot,
entry: Option<&CachedOAuthEntry>,
) -> bool {
entry.is_none() && transport.key.decrypted_api_key.trim() == "__placeholder__"
}
async fn refresh(
&self,
_client: &reqwest::Client,
_transport: &GatewayProviderTransportSnapshot,
_entry: Option<&CachedOAuthEntry>,
) -> Result<Option<CachedOAuthEntry>, LocalOAuthRefreshError> {
self.refresh_hits.fetch_add(1, Ordering::SeqCst);
Ok(Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "Bearer refreshed-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
}))
}
}
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "test".to_string(),
provider_type: "test-oauth".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "claude:cli".to_string(),
api_family: Some("claude".to_string()),
endpoint_kind: Some("cli".to_string()),
is_active: true,
base_url: "https://example.test".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "__placeholder__".to_string(),
decrypted_auth_config: Some("{\"refresh_token\":\"rt-1\"}".to_string()),
},
}
}
#[tokio::test]
async fn coordinator_reuses_runtime_cached_refresh_result() {
let refresh_hits = Arc::new(AtomicUsize::new(0));
let coordinator =
LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![Arc::new(TestAdapter {
refresh_hits: Arc::clone(&refresh_hits),
})]);
let transport = sample_transport();
let client = reqwest::Client::new();
let first = coordinator
.resolve_with_result(&client, &transport, None, None)
.await
.expect("first resolve should succeed");
let second = coordinator
.resolve_with_result(&client, &transport, None, None)
.await
.expect("second resolve should succeed");
assert_eq!(refresh_hits.load(Ordering::SeqCst), 1);
assert_eq!(
first,
Some(LocalOAuthResolution {
auth: Some(LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer refreshed-token".to_string(),
}),
refreshed_entry: Some(CachedOAuthEntry {
provider_type: "test-oauth".to_string(),
auth_header_name: "authorization".to_string(),
auth_header_value: "Bearer refreshed-token".to_string(),
expires_at_unix_secs: Some(4_102_444_800),
metadata: None,
}),
refresh_in_flight: false,
})
);
assert_eq!(
second,
Some(LocalOAuthResolution {
auth: Some(LocalResolvedOAuthRequestAuth::Header {
name: "authorization".to_string(),
value: "Bearer refreshed-token".to_string(),
}),
refreshed_entry: None,
refresh_in_flight: false,
})
);
}
}

View File

@@ -0,0 +1,137 @@
use super::provider_types::{
provider_type_supports_local_openai_chat_transport,
provider_type_supports_local_same_format_transport,
};
use super::snapshot::GatewayProviderTransportSnapshot;
use super::{
body_rules_are_locally_supported, header_rules_are_locally_supported,
resolve_transport_tls_profile, supports_local_oauth_request_auth_resolution,
transport_proxy_is_locally_supported,
};
pub fn supports_local_openai_chat_transport(transport: &GatewayProviderTransportSnapshot) -> bool {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return false;
}
if !transport
.endpoint
.api_format
.trim()
.eq_ignore_ascii_case("openai:chat")
{
return false;
}
if !header_rules_are_locally_supported(transport.endpoint.header_rules.as_ref())
|| !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref())
{
return false;
}
if transport.key.decrypted_auth_config.is_some()
&& !supports_local_oauth_request_auth_resolution(transport)
{
return false;
}
if !transport_proxy_is_locally_supported(transport) {
return false;
}
if transport.key.fingerprint.is_some() && resolve_transport_tls_profile(transport).is_none() {
return false;
}
if !provider_type_supports_local_openai_chat_transport(&transport.provider.provider_type) {
return false;
}
true
}
pub fn supports_local_standard_transport(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
) -> bool {
supports_local_same_format_transport(transport, api_format, false)
}
pub fn supports_local_gemini_transport(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
) -> bool {
supports_local_same_format_transport(transport, api_format, false)
}
pub fn supports_local_standard_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
) -> bool {
supports_local_same_format_transport(transport, api_format, true)
}
pub fn supports_local_gemini_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
) -> bool {
supports_local_same_format_transport(transport, api_format, true)
}
fn supports_local_same_format_transport(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
allow_network_passthrough: bool,
) -> bool {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return false;
}
if !transport
.endpoint
.api_format
.trim()
.eq_ignore_ascii_case(api_format.trim())
{
return false;
}
if !header_rules_are_locally_supported(transport.endpoint.header_rules.as_ref())
|| !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref())
{
return false;
}
if transport.key.decrypted_auth_config.is_some()
&& !supports_local_oauth_request_auth_resolution(transport)
{
return false;
}
let has_custom_path = transport
.endpoint
.custom_path
.as_deref()
.is_some_and(|value| !value.trim().is_empty());
if has_custom_path && !allow_network_passthrough {
return false;
}
if allow_network_passthrough {
if !transport_proxy_is_locally_supported(transport) {
return false;
}
if transport.key.fingerprint.is_some() && resolve_transport_tls_profile(transport).is_none()
{
return false;
}
} else if transport.provider.proxy.is_some()
|| transport.endpoint.proxy.is_some()
|| transport.key.proxy.is_some()
|| transport
.key
.fingerprint
.as_ref()
.and_then(|value| value.get("tls_profile"))
.and_then(|value| value.as_str())
.is_some_and(|value| !value.trim().is_empty())
{
return false;
}
if !provider_type_supports_local_same_format_transport(&transport.provider.provider_type) {
return false;
}
true
}

View File

@@ -0,0 +1,139 @@
#[derive(Debug, Clone, Copy)]
pub struct ProviderOAuthTemplate {
pub provider_type: &'static str,
pub display_name: &'static str,
pub authorize_url: &'static str,
pub token_url: &'static str,
pub client_id: &'static str,
pub client_secret: &'static str,
pub scopes: &'static [&'static str],
pub redirect_uri: &'static str,
pub use_pkce: bool,
}
pub fn provider_type_is_fixed(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"claude_code" | "kiro" | "codex" | "gemini_cli" | "antigravity" | "vertex_ai"
)
}
pub fn provider_type_enables_format_conversion_by_default(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"claude_code" | "kiro" | "codex" | "antigravity" | "vertex_ai"
)
}
pub fn fixed_provider_template(
provider_type: &str,
) -> Option<(&'static str, &'static [&'static str])> {
match provider_type.trim().to_ascii_lowercase().as_str() {
"claude_code" => Some(("https://api.anthropic.com", &["claude:cli"])),
"codex" => Some((
"https://chatgpt.com/backend-api/codex",
&["openai:cli", "openai:compact"],
)),
"kiro" => Some(("https://q.{region}.amazonaws.com", &["claude:cli"])),
"gemini_cli" => Some(("https://cloudcode-pa.googleapis.com", &["gemini:cli"])),
"vertex_ai" => Some((
"https://aiplatform.googleapis.com",
&["gemini:chat", "claude:chat"],
)),
"antigravity" => Some(("https://cloudcode-pa.googleapis.com", &["gemini:chat"])),
_ => None,
}
}
pub fn provider_type_supports_model_fetch(provider_type: &str) -> bool {
!matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"vertex_ai" | "antigravity" | "codex" | "kiro" | "claude_code"
)
}
pub fn provider_type_supports_local_openai_chat_transport(provider_type: &str) -> bool {
!matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"antigravity" | "claude_code" | "codex" | "gemini_cli" | "kiro" | "vertex_ai"
)
}
pub fn provider_type_supports_local_same_format_transport(provider_type: &str) -> bool {
!matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"antigravity" | "claude_code" | "kiro" | "vertex_ai"
)
}
pub fn is_codex_cli_backend_url(url: &str) -> bool {
let url = url.trim().to_ascii_lowercase();
url.contains("/codex") && (url.contains("/backend-api/") || url.contains("/backendapi/"))
}
pub fn provider_type_is_fixed_for_admin_oauth(provider_type: &str) -> bool {
provider_type_is_fixed(provider_type)
}
pub fn provider_type_admin_oauth_template(provider_type: &str) -> Option<ProviderOAuthTemplate> {
match provider_type.trim().to_ascii_lowercase().as_str() {
"claude_code" => Some(ProviderOAuthTemplate {
provider_type: "claude_code",
display_name: "ClaudeCode",
authorize_url: "https://claude.ai/oauth/authorize",
token_url: "https://console.anthropic.com/v1/oauth/token",
client_id: "9d1c250a-e61b-44d9-88ed-5944d1962f5e",
client_secret: "",
scopes: &["org:create_api_key", "user:profile", "user:inference"],
redirect_uri: "http://localhost:54545/callback",
use_pkce: true,
}),
"codex" => Some(ProviderOAuthTemplate {
provider_type: "codex",
display_name: "Codex",
authorize_url: "https://auth.openai.com/oauth/authorize",
token_url: "https://auth.openai.com/oauth/token",
client_id: "app_EMoamEEZ73f0CkXaXp7hrann",
client_secret: "",
scopes: &["openid", "email", "profile", "offline_access"],
redirect_uri: "http://localhost:1455/auth/callback",
use_pkce: true,
}),
"gemini_cli" => Some(ProviderOAuthTemplate {
provider_type: "gemini_cli",
display_name: "GeminiCli",
authorize_url: "https://accounts.google.com/o/oauth2/v2/auth",
token_url: "https://oauth2.googleapis.com/token",
client_id: "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com",
client_secret: "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",
scopes: &[
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
"https://www.googleapis.com/auth/userinfo.profile",
],
redirect_uri: "http://localhost:8085/oauth2callback",
use_pkce: false,
}),
"antigravity" => Some(ProviderOAuthTemplate {
provider_type: "antigravity",
display_name: "Antigravity",
authorize_url: "https://accounts.google.com/o/oauth2/v2/auth",
token_url: "https://oauth2.googleapis.com/token",
client_id: "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com",
client_secret: "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",
scopes: &[
"https://www.googleapis.com/auth/cloud-platform",
"https://www.googleapis.com/auth/userinfo.email",
"https://www.googleapis.com/auth/userinfo.profile",
"https://www.googleapis.com/auth/cclog",
"https://www.googleapis.com/auth/experimentsandconfigs",
],
redirect_uri: "http://localhost:51121/oauth2callback",
use_pkce: true,
}),
_ => None,
}
}
pub const ADMIN_PROVIDER_OAUTH_TEMPLATE_TYPES: &[&str] =
&["claude_code", "codex", "gemini_cli", "antigravity"];

View File

@@ -0,0 +1,820 @@
use std::collections::{BTreeMap, HashSet};
use regex::Regex;
use serde_json::{Map, Value};
const ORIGINAL_PLACEHOLDER: &str = "{{$original}}";
const CONDITION_SOURCES: &[&str] = &["current", "original"];
const CONDITION_TYPE_VALUES: &[&str] = &["string", "number", "boolean", "array", "object", "null"];
#[derive(Debug, Clone, PartialEq, Eq)]
enum BodyPathSegment {
Key(String),
Index(isize),
}
pub fn header_rules_are_locally_supported(rules: Option<&Value>) -> bool {
let Some(rules) = rules else {
return true;
};
let Some(rules) = rules.as_array() else {
return false;
};
rules.iter().all(|rule| {
let Some(rule) = rule.as_object() else {
return false;
};
if rule
.get("condition")
.is_some_and(|value| !value.is_null() && !condition_is_locally_supported(value))
{
return false;
}
match rule
.get("action")
.and_then(Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.as_deref()
{
Some("set") => {
rule.get("key")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty())
&& rule.get("value").is_some_and(Value::is_string)
}
Some("drop") => rule
.get("key")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty()),
Some("rename") => {
rule.get("from")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty())
&& rule
.get("to")
.and_then(Value::as_str)
.is_some_and(|value| !value.trim().is_empty())
}
_ => false,
}
})
}
pub fn apply_local_header_rules(
headers: &mut BTreeMap<String, String>,
rules: Option<&Value>,
protected_keys: &[&str],
body: &Value,
original_body: Option<&Value>,
) -> bool {
let Some(rules) = rules else {
return true;
};
let Some(rules) = rules.as_array() else {
return false;
};
let protected_keys: HashSet<String> = protected_keys
.iter()
.map(|value| value.trim().to_ascii_lowercase())
.collect();
for rule in rules {
let Some(rule) = rule.as_object() else {
return false;
};
if let Some(condition) = rule.get("condition").filter(|value| !value.is_null()) {
if !condition_is_locally_supported(condition) {
return false;
}
if !evaluate_local_condition(body, condition, original_body) {
continue;
}
}
match rule
.get("action")
.and_then(Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.as_deref()
{
Some("set") => {
let Some(key) = rule.get("key").and_then(Value::as_str).map(str::trim) else {
return false;
};
let Some(value) = rule.get("value").and_then(Value::as_str) else {
return false;
};
let key = key.to_ascii_lowercase();
if !protected_keys.contains(&key) {
headers.insert(key, value.to_string());
}
}
Some("drop") => {
let Some(key) = rule.get("key").and_then(Value::as_str).map(str::trim) else {
return false;
};
let key = key.to_ascii_lowercase();
if !protected_keys.contains(&key) {
headers.remove(&key);
}
}
Some("rename") => {
let Some(from) = rule.get("from").and_then(Value::as_str).map(str::trim) else {
return false;
};
let Some(to) = rule.get("to").and_then(Value::as_str).map(str::trim) else {
return false;
};
let from = from.to_ascii_lowercase();
let to = to.to_ascii_lowercase();
if protected_keys.contains(&from) || protected_keys.contains(&to) {
continue;
}
if let Some(value) = headers.remove(&from) {
headers.insert(to, value);
}
}
_ => return false,
}
}
true
}
pub fn body_rules_are_locally_supported(rules: Option<&Value>) -> bool {
let Some(rules) = rules else {
return true;
};
let Some(rules) = rules.as_array() else {
return false;
};
rules.iter().all(|rule| {
let Some(rule) = rule.as_object() else {
return false;
};
if rule
.get("condition")
.is_some_and(|value| !value.is_null() && !condition_is_locally_supported(value))
{
return false;
}
match rule
.get("action")
.and_then(Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.as_deref()
{
Some("set") => {
rule.get("path")
.and_then(Value::as_str)
.and_then(parse_body_path)
.is_some()
&& !rule.get("value").is_some_and(contains_original_placeholder)
}
Some("drop") => rule
.get("path")
.and_then(Value::as_str)
.and_then(parse_body_path)
.is_some(),
Some("rename") => {
rule.get("from")
.and_then(Value::as_str)
.and_then(parse_body_path)
.is_some()
&& rule
.get("to")
.and_then(Value::as_str)
.and_then(parse_body_path)
.is_some()
}
_ => false,
}
})
}
pub fn apply_local_body_rules(
body: &mut Value,
rules: Option<&Value>,
original_body: Option<&Value>,
) -> bool {
let Some(rules) = rules else {
return true;
};
let Some(rules) = rules.as_array() else {
return false;
};
for rule in rules {
let Some(rule) = rule.as_object() else {
return false;
};
if let Some(condition) = rule.get("condition").filter(|value| !value.is_null()) {
if !condition_is_locally_supported(condition) {
return false;
}
if !evaluate_local_condition(body, condition, original_body) {
continue;
}
}
match rule
.get("action")
.and_then(Value::as_str)
.map(str::trim)
.map(str::to_ascii_lowercase)
.as_deref()
{
Some("set") => {
let Some(path) = rule
.get("path")
.and_then(Value::as_str)
.and_then(parse_body_path)
else {
return false;
};
let value = rule.get("value").cloned().unwrap_or(Value::Null);
if contains_original_placeholder(&value) {
return false;
}
let _ = set_nested_value(body, &path, value);
}
Some("drop") => {
let Some(path) = rule
.get("path")
.and_then(Value::as_str)
.and_then(parse_body_path)
else {
return false;
};
let _ = delete_nested_value(body, &path);
}
Some("rename") => {
let Some(from) = rule
.get("from")
.and_then(Value::as_str)
.and_then(parse_body_path)
else {
return false;
};
let Some(to) = rule
.get("to")
.and_then(Value::as_str)
.and_then(parse_body_path)
else {
return false;
};
let _ = rename_nested_value(body, &from, &to);
}
_ => return false,
}
}
true
}
fn condition_is_locally_supported(condition: &Value) -> bool {
let Some(condition) = condition.as_object() else {
return false;
};
if let Some(children) = condition.get("all").and_then(Value::as_array) {
return !children.is_empty() && children.iter().all(condition_is_locally_supported);
}
if let Some(children) = condition.get("any").and_then(Value::as_array) {
return !children.is_empty() && children.iter().all(condition_is_locally_supported);
}
let source = condition
.get("source")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or("current");
if !CONDITION_SOURCES.contains(&source) {
return false;
}
let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else {
return false;
};
let Some(path) = condition
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.and_then(parse_body_path)
else {
return false;
};
if path.is_empty() {
return false;
}
match op {
"exists" | "not_exists" | "eq" | "neq" => true,
"gt" | "lt" | "gte" | "lte" => condition
.get("value")
.is_some_and(|value| value.as_f64().is_some() && !value.is_boolean()),
"starts_with" | "ends_with" | "matches" => condition
.get("value")
.and_then(Value::as_str)
.is_some_and(|value| {
if op == "matches" {
Regex::new(value).is_ok()
} else {
true
}
}),
"contains" => condition.get("value").is_some(),
"in" => condition.get("value").is_some_and(Value::is_array),
"type_is" => condition
.get("value")
.and_then(Value::as_str)
.is_some_and(|value| CONDITION_TYPE_VALUES.contains(&value)),
_ => false,
}
}
fn evaluate_local_condition(
body: &Value,
condition: &Value,
original_body: Option<&Value>,
) -> bool {
let Some(condition) = condition.as_object() else {
return false;
};
if let Some(children) = condition.get("all").and_then(Value::as_array) {
return !children.is_empty()
&& children
.iter()
.all(|child| evaluate_local_condition(body, child, original_body));
}
if let Some(children) = condition.get("any").and_then(Value::as_array) {
return !children.is_empty()
&& children
.iter()
.any(|child| evaluate_local_condition(body, child, original_body));
}
let source = condition
.get("source")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or("current");
let target = if source.eq_ignore_ascii_case("original") {
original_body.unwrap_or(body)
} else {
body
};
let Some(op) = condition.get("op").and_then(Value::as_str).map(str::trim) else {
return false;
};
let Some(path) = condition
.get("path")
.and_then(Value::as_str)
.map(str::trim)
.and_then(parse_body_path)
else {
return false;
};
let current_value = get_nested_value(target, &path);
if op == "exists" {
return current_value.is_some();
}
if op == "not_exists" {
return current_value.is_none();
}
let Some(current_value) = current_value else {
return false;
};
let expected = condition.get("value");
match op {
"eq" => expected == Some(&current_value),
"neq" => expected != Some(&current_value),
"gt" | "lt" | "gte" | "lte" => {
let Some(current) = json_number(&current_value) else {
return false;
};
let Some(expected) = expected.and_then(json_number) else {
return false;
};
match op {
"gt" => current > expected,
"lt" => current < expected,
"gte" => current >= expected,
"lte" => current <= expected,
_ => false,
}
}
"starts_with" => current_value
.as_str()
.zip(expected.and_then(Value::as_str))
.is_some_and(|(current, expected)| current.starts_with(expected)),
"ends_with" => current_value
.as_str()
.zip(expected.and_then(Value::as_str))
.is_some_and(|(current, expected)| current.ends_with(expected)),
"contains" => match (current_value, expected) {
(Value::String(current), Some(Value::String(expected))) => current.contains(expected),
(Value::Array(current), Some(expected)) => {
current.iter().any(|value| value == expected)
}
_ => false,
},
"matches" => current_value
.as_str()
.zip(expected.and_then(Value::as_str))
.is_some_and(|(current, expected)| {
Regex::new(expected)
.map(|pattern| pattern.is_match(current))
.unwrap_or(false)
}),
"in" => expected
.and_then(Value::as_array)
.is_some_and(|values| values.iter().any(|value| value == &current_value)),
"type_is" => expected
.and_then(Value::as_str)
.is_some_and(|expected| match expected {
"string" => current_value.is_string(),
"number" => current_value.as_f64().is_some() && !current_value.is_boolean(),
"boolean" => current_value.is_boolean(),
"array" => current_value.is_array(),
"object" => current_value.is_object(),
"null" => current_value.is_null(),
_ => false,
}),
_ => false,
}
}
fn json_number(value: &Value) -> Option<f64> {
value.as_f64().filter(|_| !value.is_boolean())
}
fn parse_body_path(path: &str) -> Option<Vec<BodyPathSegment>> {
let raw = path.trim();
if raw.is_empty() {
return None;
}
let chars: Vec<char> = raw.chars().collect();
let mut parts = Vec::new();
let mut current = String::new();
let mut expect_key = true;
let mut index = 0usize;
while index < chars.len() {
let ch = chars[index];
if ch == '\\' && chars.get(index + 1).copied() == Some('.') {
current.push('.');
expect_key = false;
index += 2;
continue;
}
if ch == '.' {
if !current.is_empty() {
parts.push(BodyPathSegment::Key(std::mem::take(&mut current)));
} else if expect_key {
return None;
}
expect_key = true;
index += 1;
continue;
}
if ch == '[' {
if !current.is_empty() {
parts.push(BodyPathSegment::Key(std::mem::take(&mut current)));
}
let mut close_index = index + 1;
while close_index < chars.len() && chars[close_index] != ']' {
close_index += 1;
}
if close_index >= chars.len() {
return None;
}
let inner = chars[index + 1..close_index]
.iter()
.collect::<String>()
.trim()
.to_string();
if inner.is_empty() || inner == "*" {
return None;
}
let Ok(index_value) = inner.parse::<isize>() else {
return None;
};
parts.push(BodyPathSegment::Index(index_value));
expect_key = false;
index = close_index + 1;
continue;
}
current.push(ch);
expect_key = false;
index += 1;
}
if !current.is_empty() {
parts.push(BodyPathSegment::Key(current));
} else if expect_key {
return None;
}
(!parts.is_empty()).then_some(parts)
}
fn contains_original_placeholder(value: &Value) -> bool {
match value {
Value::String(value) => value.contains(ORIGINAL_PLACEHOLDER),
Value::Array(items) => items.iter().any(contains_original_placeholder),
Value::Object(items) => items.values().any(contains_original_placeholder),
_ => false,
}
}
fn resolve_index(len: usize, index: isize) -> Option<usize> {
if index >= 0 {
((index as usize) < len).then_some(index as usize)
} else {
let resolved = len as isize + index;
(resolved >= 0).then_some(resolved as usize)
}
}
fn get_nested_value(value: &Value, path: &[BodyPathSegment]) -> Option<Value> {
let mut current = value;
for segment in path {
match segment {
BodyPathSegment::Key(key) => {
current = current.as_object()?.get(key)?;
}
BodyPathSegment::Index(index) => {
let values = current.as_array()?;
let resolved = resolve_index(values.len(), *index)?;
current = values.get(resolved)?;
}
}
}
Some(current.clone())
}
fn get_existing_child_mut<'a>(
current: &'a mut Value,
segment: &BodyPathSegment,
) -> Option<&'a mut Value> {
match segment {
BodyPathSegment::Key(key) => current.as_object_mut()?.get_mut(key),
BodyPathSegment::Index(index) => {
let values = current.as_array_mut()?;
let resolved = resolve_index(values.len(), *index)?;
values.get_mut(resolved)
}
}
}
fn set_nested_value(current: &mut Value, path: &[BodyPathSegment], value: Value) -> bool {
let Some((last, parents)) = path.split_last() else {
return false;
};
let mut current = current;
for (offset, segment) in parents.iter().enumerate() {
let next = &path[offset + 1];
match segment {
BodyPathSegment::Key(key) => {
let Some(object) = current.as_object_mut() else {
return false;
};
match next {
BodyPathSegment::Key(_) => {
let child = object
.entry(key.clone())
.or_insert_with(|| Value::Object(Map::new()));
if !child.is_object() {
*child = Value::Object(Map::new());
}
current = child;
}
BodyPathSegment::Index(_) => {
let Some(child) = object.get_mut(key) else {
return false;
};
if !child.is_array() {
return false;
}
current = child;
}
}
}
BodyPathSegment::Index(index) => {
let Some(values) = current.as_array_mut() else {
return false;
};
let Some(resolved) = resolve_index(values.len(), *index) else {
return false;
};
current = &mut values[resolved];
}
}
}
match last {
BodyPathSegment::Key(key) => {
let Some(object) = current.as_object_mut() else {
return false;
};
object.insert(key.clone(), value);
true
}
BodyPathSegment::Index(index) => {
let Some(values) = current.as_array_mut() else {
return false;
};
let Some(resolved) = resolve_index(values.len(), *index) else {
return false;
};
values[resolved] = value;
true
}
}
}
fn delete_nested_value(current: &mut Value, path: &[BodyPathSegment]) -> bool {
let Some((last, parents)) = path.split_last() else {
return false;
};
let mut current = current;
for segment in parents {
let Some(child) = get_existing_child_mut(current, segment) else {
return false;
};
current = child;
}
match last {
BodyPathSegment::Key(key) => current
.as_object_mut()
.and_then(|object| object.remove(key))
.is_some(),
BodyPathSegment::Index(index) => {
let Some(values) = current.as_array_mut() else {
return false;
};
let Some(resolved) = resolve_index(values.len(), *index) else {
return false;
};
values.remove(resolved);
true
}
}
}
fn rename_nested_value(
current: &mut Value,
from: &[BodyPathSegment],
to: &[BodyPathSegment],
) -> bool {
if from == to {
return get_nested_value(current, from).is_some();
}
let Some(value) = get_nested_value(current, from) else {
return false;
};
if !set_nested_value(current, to, value) {
return false;
}
delete_nested_value(current, from)
}
#[cfg(test)]
mod tests {
use super::{
apply_local_body_rules, apply_local_header_rules, body_rules_are_locally_supported,
header_rules_are_locally_supported,
};
#[test]
fn header_rules_allow_simple_set_drop_and_rename() {
let rules = serde_json::json!([
{"action":"set","key":"x-added","value":"1"},
{"action":"drop","key":"x-drop"},
{"action":"rename","from":"x-old","to":"x-new"}
]);
assert!(header_rules_are_locally_supported(Some(&rules)));
let mut headers = std::collections::BTreeMap::from([
("x-drop".to_string(), "drop-me".to_string()),
("x-old".to_string(), "old-value".to_string()),
("authorization".to_string(), "Bearer keep".to_string()),
]);
assert!(apply_local_header_rules(
&mut headers,
Some(&rules),
&["authorization", "content-type"],
&serde_json::json!({}),
None,
));
assert_eq!(headers.get("x-added").map(String::as_str), Some("1"));
assert!(!headers.contains_key("x-drop"));
assert_eq!(headers.get("x-new").map(String::as_str), Some("old-value"));
assert_eq!(
headers.get("authorization").map(String::as_str),
Some("Bearer keep")
);
}
#[test]
fn header_rules_allow_simple_conditions() {
let rules = serde_json::json!([
{"action":"set","key":"x-added","value":"1","condition":{"path":"metadata.mode","op":"eq","value":"safe"}},
{"action":"set","key":"x-from-original","value":"1","condition":{"path":"metadata.client","op":"exists","source":"original"}}
]);
assert!(header_rules_are_locally_supported(Some(&rules)));
let mut headers = std::collections::BTreeMap::new();
assert!(apply_local_header_rules(
&mut headers,
Some(&rules),
&[],
&serde_json::json!({"metadata":{"mode":"safe"}}),
Some(&serde_json::json!({"metadata":{"client":"desktop"}})),
));
assert_eq!(headers.get("x-added").map(String::as_str), Some("1"));
assert_eq!(
headers.get("x-from-original").map(String::as_str),
Some("1")
);
}
#[test]
fn body_rules_allow_simple_nested_set_drop_and_rename() {
let rules = serde_json::json!([
{"action":"set","path":"metadata.mode","value":"safe"},
{"action":"drop","path":"tools[1]"},
{"action":"rename","from":"messages[0].content","to":"messages[0].text"}
]);
assert!(body_rules_are_locally_supported(Some(&rules)));
let mut body = serde_json::json!({
"messages": [{"content":"hello"}],
"tools": [{"name":"a"},{"name":"b"}]
});
assert!(apply_local_body_rules(&mut body, Some(&rules), None));
assert_eq!(body["metadata"]["mode"], "safe");
assert_eq!(body["tools"], serde_json::json!([{"name":"a"}]));
assert_eq!(body["messages"][0]["text"], "hello");
assert!(body["messages"][0].get("content").is_none());
}
#[test]
fn body_rules_allow_simple_conditions() {
let rules = serde_json::json!([
{"action":"set","path":"instructions","value":"You are GPT-5.","condition":{"path":"instructions","op":"not_exists"}},
{"action":"set","path":"metadata.origin","value":"desktop","condition":{"path":"metadata.client","op":"exists","source":"original"}}
]);
assert!(body_rules_are_locally_supported(Some(&rules)));
let original = serde_json::json!({
"metadata": {
"client": "desktop"
}
});
let mut body = serde_json::json!({
"metadata": {
"mode": "safe"
}
});
assert!(apply_local_body_rules(
&mut body,
Some(&rules),
Some(&original)
));
assert_eq!(body["instructions"], "You are GPT-5.");
assert_eq!(body["metadata"]["origin"], "desktop");
}
#[test]
fn body_rules_reject_placeholder_and_wildcard_paths() {
let placeholder_rules =
serde_json::json!([{"action":"set","path":"model","value":"{{$original}}"}]);
assert!(!body_rules_are_locally_supported(Some(&placeholder_rules)));
let wildcard_rules = serde_json::json!([{"action":"drop","path":"items[*].value"}]);
assert!(!body_rules_are_locally_supported(Some(&wildcard_rules)));
}
}

View File

@@ -0,0 +1,985 @@
use aether_data::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::DataLayerError;
use async_trait::async_trait;
use super::auth_config::{absorb_local_auth_config_safe_subset, LocalAuthConfigAbsorption};
#[path = "snapshot_mapping.rs"]
mod snapshot_mapping;
use self::snapshot_mapping::{fallback_encryption_keys, map_endpoint, map_key, map_provider};
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct GatewayProviderTransportSnapshot {
pub provider: GatewayProviderTransportProvider,
pub endpoint: GatewayProviderTransportEndpoint,
pub key: GatewayProviderTransportKey,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct GatewayProviderTransportProvider {
pub id: String,
pub name: String,
pub provider_type: String,
pub website: Option<String>,
pub is_active: bool,
pub keep_priority_on_conversion: bool,
pub enable_format_conversion: bool,
pub concurrent_limit: Option<i32>,
pub max_retries: Option<i32>,
pub proxy: Option<serde_json::Value>,
pub request_timeout_secs: Option<f64>,
pub stream_first_byte_timeout_secs: Option<f64>,
pub config: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct GatewayProviderTransportEndpoint {
pub id: String,
pub provider_id: String,
pub api_format: String,
pub api_family: Option<String>,
pub endpoint_kind: Option<String>,
pub is_active: bool,
pub base_url: String,
pub header_rules: Option<serde_json::Value>,
pub body_rules: Option<serde_json::Value>,
pub max_retries: Option<i32>,
pub custom_path: Option<String>,
pub config: Option<serde_json::Value>,
pub format_acceptance_config: Option<serde_json::Value>,
pub proxy: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct GatewayProviderTransportKey {
pub id: String,
pub provider_id: String,
pub name: String,
pub auth_type: String,
pub is_active: bool,
pub api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
pub capabilities: Option<serde_json::Value>,
pub rate_multipliers: Option<serde_json::Value>,
pub global_priority_by_format: Option<serde_json::Value>,
pub expires_at_unix_secs: Option<u64>,
pub proxy: Option<serde_json::Value>,
pub fingerprint: Option<serde_json::Value>,
pub decrypted_api_key: String,
pub decrypted_auth_config: Option<String>,
}
#[async_trait]
pub trait ProviderTransportSnapshotSource: Send + Sync {
fn encryption_key(&self) -> Option<&str>;
async fn list_provider_catalog_providers_by_ids(
&self,
ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError>;
async fn list_provider_catalog_endpoints_by_ids(
&self,
ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError>;
async fn list_provider_catalog_keys_by_ids(
&self,
ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError>;
}
pub async fn read_provider_transport_snapshot(
state: &dyn ProviderTransportSnapshotSource,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
) -> Result<Option<GatewayProviderTransportSnapshot>, DataLayerError> {
let Some(encryption_key) = state.encryption_key() else {
return Ok(None);
};
let fallback_encryption_keys = fallback_encryption_keys(encryption_key);
let providers = state
.list_provider_catalog_providers_by_ids(&[provider_id.to_string()])
.await?;
let endpoints = state
.list_provider_catalog_endpoints_by_ids(&[endpoint_id.to_string()])
.await?;
let keys = state
.list_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await?;
let Some(provider) = providers.into_iter().next() else {
return Ok(None);
};
let Some(endpoint) = endpoints.into_iter().next() else {
return Ok(None);
};
let Some(key) = keys.into_iter().next() else {
return Ok(None);
};
if endpoint.provider_id != provider.id {
return Err(DataLayerError::UnexpectedValue(format!(
"provider_endpoints.provider_id mismatch: expected {}, got {}",
provider.id, endpoint.provider_id
)));
}
if key.provider_id != provider.id {
return Err(DataLayerError::UnexpectedValue(format!(
"provider_api_keys.provider_id mismatch: expected {}, got {}",
provider.id, key.provider_id
)));
}
let provider = map_provider(provider);
let mut endpoint = map_endpoint(endpoint);
let mut key = map_key(key, encryption_key, &fallback_encryption_keys)?;
if let LocalAuthConfigAbsorption::Absorbed {
base_url,
header_rules,
custom_path,
} = absorb_local_auth_config_safe_subset(
&endpoint.base_url,
endpoint.header_rules.clone(),
endpoint.custom_path.clone(),
key.decrypted_auth_config.as_deref(),
) {
endpoint.base_url = base_url;
endpoint.header_rules = header_rules;
endpoint.custom_path = custom_path;
key.decrypted_auth_config = None;
}
Ok(Some(GatewayProviderTransportSnapshot {
provider,
endpoint,
key,
}))
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::DataLayerError;
use async_trait::async_trait;
use super::super::policy::{
supports_local_openai_chat_transport, supports_local_standard_transport_with_network,
};
use super::{
map_key, read_provider_transport_snapshot, GatewayProviderTransportSnapshot,
ProviderTransportSnapshotSource,
};
struct TestSnapshotSource {
repository: Arc<InMemoryProviderCatalogReadRepository>,
encryption_key: Option<String>,
}
impl TestSnapshotSource {
fn new(
repository: Arc<InMemoryProviderCatalogReadRepository>,
encryption_key: impl Into<Option<String>>,
) -> Self {
Self {
repository,
encryption_key: encryption_key.into(),
}
}
}
#[async_trait]
impl ProviderTransportSnapshotSource for TestSnapshotSource {
fn encryption_key(&self) -> Option<&str> {
self.encryption_key.as_deref()
}
async fn list_provider_catalog_providers_by_ids(
&self,
ids: &[String],
) -> Result<Vec<StoredProviderCatalogProvider>, DataLayerError> {
self.repository.list_providers_by_ids(ids).await
}
async fn list_provider_catalog_endpoints_by_ids(
&self,
ids: &[String],
) -> Result<Vec<StoredProviderCatalogEndpoint>, DataLayerError> {
self.repository.list_endpoints_by_ids(ids).await
}
async fn list_provider_catalog_keys_by_ids(
&self,
ids: &[String],
) -> Result<Vec<StoredProviderCatalogKey>, DataLayerError> {
self.repository.list_keys_by_ids(ids).await
}
}
fn sample_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
Some("https://openai.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
true,
Some(32),
Some(3),
Some(serde_json::json!({"url":"http://provider-proxy"})),
Some(20.0),
Some(8.0),
Some(serde_json::json!({"region":"global"})),
)
}
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
Some(serde_json::json!([{"action":"set","key":"x-test","value":"1"}])),
Some(serde_json::json!([{"action":"drop","path":"stream"}])),
Some(2),
Some("/v1/chat/completions".to_string()),
Some(serde_json::json!({"api_version":"v1"})),
Some(serde_json::json!({"allow":["openai:chat"]})),
Some(serde_json::json!({"url":"http://endpoint-proxy"})),
)
.expect("endpoint transport fields should build")
}
fn sample_key() -> StoredProviderCatalogKey {
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}",
)
.expect("auth config ciphertext should build");
StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"prod-key".to_string(),
"api_key".to_string(),
Some(serde_json::json!({"cache_1h": true})),
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat", "openai:cli"])),
encrypted_api_key,
Some(encrypted_auth_config),
Some(serde_json::json!({"openai:chat": 0.8})),
Some(serde_json::json!({"openai:chat": 1})),
Some(serde_json::json!(["gpt-4.1", "gpt-4.1-mini"])),
Some(1_800_000_000),
Some(serde_json::json!({"node_id":"proxy-node-1"})),
Some(serde_json::json!({"tls_profile":"chrome_136"})),
)
.expect("key transport fields should build")
}
fn read_state() -> TestSnapshotSource {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()))
}
#[tokio::test]
async fn reads_decrypted_provider_transport_snapshot() {
let state = read_state();
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-1", "key-1")
.await
.expect("snapshot should read")
.expect("snapshot should exist");
assert_eq!(
snapshot,
GatewayProviderTransportSnapshot {
provider: super::GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "OpenAI".to_string(),
provider_type: "custom".to_string(),
website: Some("https://openai.com".to_string()),
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: Some(32),
max_retries: Some(3),
proxy: Some(serde_json::json!({"url":"http://provider-proxy"})),
request_timeout_secs: Some(20.0),
stream_first_byte_timeout_secs: Some(8.0),
config: Some(serde_json::json!({"region":"global"})),
},
endpoint: super::GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:chat".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://api.openai.com".to_string(),
header_rules: Some(
serde_json::json!([{"action":"set","key":"x-test","value":"1"}]),
),
body_rules: Some(serde_json::json!([{"action":"drop","path":"stream"}])),
max_retries: Some(2),
custom_path: Some("/v1/chat/completions".to_string()),
config: Some(serde_json::json!({"api_version":"v1"})),
format_acceptance_config: Some(serde_json::json!({"allow":["openai:chat"]}),),
proxy: Some(serde_json::json!({"url":"http://endpoint-proxy"})),
},
key: super::GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "prod-key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["openai:chat".to_string(), "openai:cli".to_string(),]),
allowed_models: Some(vec!["gpt-4.1".to_string(), "gpt-4.1-mini".to_string(),]),
capabilities: Some(serde_json::json!({"cache_1h": true})),
rate_multipliers: Some(serde_json::json!({"openai:chat": 0.8})),
global_priority_by_format: Some(serde_json::json!({"openai:chat": 1})),
expires_at_unix_secs: Some(1_800_000_000),
proxy: Some(serde_json::json!({"node_id":"proxy-node-1"})),
fingerprint: Some(serde_json::json!({"tls_profile":"chrome_136"})),
decrypted_api_key: "sk-live-openai".to_string(),
decrypted_auth_config: Some(
"{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}".to_string(),
),
},
}
);
}
#[tokio::test]
async fn returns_none_when_encryption_key_is_not_configured() {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = TestSnapshotSource::new(repository, None);
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-1", "key-1")
.await
.expect("snapshot read should not error");
assert!(snapshot.is_none());
}
#[tokio::test]
async fn absorbs_safe_auth_config_into_local_transport_fields() {
let provider = sample_provider();
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-safe-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
Some(serde_json::json!([{"action":"set","key":"x-test","value":"1"}])),
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport fields should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"headers":{"x-account-id":"acc-1"},"query":{"tenant":"demo"}}"#,
)
.expect("auth config ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-safe-1".to_string(),
"provider-1".to_string(),
"safe-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport fields should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state =
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()));
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-safe-1", "key-safe-1")
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.key.decrypted_auth_config, None);
assert_eq!(
snapshot.endpoint.base_url,
"https://api.openai.com?tenant=demo"
);
assert_eq!(snapshot.endpoint.custom_path.as_deref(), None);
assert_eq!(
snapshot.endpoint.header_rules,
Some(serde_json::json!([
{"action":"set","key":"x-test","value":"1"},
{"action":"set","key":"x-account-id","value":"acc-1"}
]))
);
assert!(supports_local_openai_chat_transport(&snapshot));
}
#[tokio::test]
async fn accepts_plaintext_legacy_key_material() {
let provider = sample_provider();
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-legacy-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
Some(serde_json::json!([{"action":"set","key":"x-test","value":"1"}])),
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport fields should build");
let key = StoredProviderCatalogKey::new(
"key-legacy-1".to_string(),
"provider-1".to_string(),
"legacy-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
"sk-plaintext-openai".to_string(),
Some(r#"{"headers":{"x-account-id":"acc-legacy"}}"#.to_string()),
None,
None,
None,
None,
None,
None,
)
.expect("key transport fields should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state =
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()));
let snapshot = read_provider_transport_snapshot(
&state,
"provider-1",
"endpoint-legacy-1",
"key-legacy-1",
)
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.key.decrypted_api_key, "sk-plaintext-openai");
assert_eq!(snapshot.key.decrypted_auth_config, None);
assert_eq!(
snapshot.endpoint.header_rules,
Some(serde_json::json!([
{"action":"set","key":"x-test","value":"1"},
{"action":"set","key":"x-account-id","value":"acc-legacy"}
]))
);
}
#[tokio::test]
async fn rejects_fernet_shaped_key_material_when_encryption_key_is_wrong() {
let key = sample_key();
let error = map_key(key, "wrong-encryption-key", &[])
.expect_err("snapshot read should fail for Fernet-shaped data with wrong key");
assert!(matches!(error, DataLayerError::UnexpectedValue(message)
if message.contains("failed to decrypt provider_api_keys.api_key")));
}
#[test]
fn decrypts_fernet_shaped_key_material_with_fallback_encryption_key() {
let key = sample_key();
let mapped = map_key(
key,
"wrong-encryption-key",
&[DEVELOPMENT_ENCRYPTION_KEY.to_string()],
)
.expect("fallback key should decrypt");
assert_eq!(mapped.decrypted_api_key, "sk-live-openai");
assert_eq!(
mapped.decrypted_auth_config.as_deref(),
Some("{\"refresh_token\":\"rt-1\",\"project\":\"demo\"}")
);
}
#[test]
fn accepts_stringified_allowed_models_in_transport_key() {
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-compat-1".to_string(),
"provider-1".to_string(),
"compat-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:chat"])),
encrypted_api_key,
None,
None,
None,
Some(serde_json::json!("[\"gpt-5.2\", \"gpt-5\"]")),
None,
None,
None,
)
.expect("key transport fields should build");
let mapped =
map_key(key, DEVELOPMENT_ENCRYPTION_KEY, &[]).expect("stringified list should parse");
assert_eq!(
mapped.allowed_models,
Some(vec!["gpt-5.2".to_string(), "gpt-5".to_string()])
);
}
#[test]
fn accepts_single_string_api_format_in_transport_key() {
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-compat-2".to_string(),
"provider-1".to_string(),
"compat-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!("openai:chat")),
encrypted_api_key,
None,
None,
None,
Some(serde_json::json!("gpt-5.2")),
None,
None,
None,
)
.expect("key transport fields should build");
let mapped =
map_key(key, DEVELOPMENT_ENCRYPTION_KEY, &[]).expect("single string should parse");
assert_eq!(mapped.api_formats, Some(vec!["openai:chat".to_string()]));
assert_eq!(mapped.allowed_models, Some(vec!["gpt-5.2".to_string()]));
}
#[tokio::test]
async fn keeps_unsupported_auth_config_blocking_local_transport() {
let provider = sample_provider();
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-safe-2".to_string(),
"provider-1".to_string(),
"openai:cli".to_string(),
Some("openai".to_string()),
Some("cli".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport fields should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"refresh_token":"rt-1","project":"demo"}"#,
)
.expect("auth config ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-safe-2".to_string(),
"provider-1".to_string(),
"unsafe-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:cli"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport fields should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state =
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()));
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-safe-2", "key-safe-2")
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(
snapshot.key.decrypted_auth_config.as_deref(),
Some(r#"{"refresh_token":"rt-1","project":"demo"}"#)
);
assert!(!supports_local_standard_transport_with_network(
&snapshot,
"openai:cli"
));
}
#[tokio::test]
async fn absorbs_query_only_auth_config_into_gemini_base_url() {
let provider = sample_provider();
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-safe-3".to_string(),
"provider-1".to_string(),
"gemini:chat".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com/v1beta".to_string(),
None,
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport fields should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"query":{"alt":"sse"}}"#,
)
.expect("auth config ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-safe-3".to_string(),
"provider-1".to_string(),
"safe-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:chat"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport fields should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state =
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()));
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-safe-3", "key-safe-3")
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(
snapshot.endpoint.base_url,
"https://generativelanguage.googleapis.com/v1beta?alt=sse"
);
assert_eq!(snapshot.endpoint.custom_path, None);
assert_eq!(snapshot.key.decrypted_auth_config, None);
}
#[tokio::test]
async fn absorbs_transport_subset_when_metadata_is_present() {
let provider = sample_provider();
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-safe-4".to_string(),
"provider-1".to_string(),
"openai:cli".to_string(),
Some("openai".to_string()),
Some("cli".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com/v1".to_string(),
Some(serde_json::json!([{"action":"set","key":"x-base","value":"1"}])),
None,
Some(2),
None,
None,
None,
None,
)
.expect("endpoint transport fields should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let encrypted_auth_config = encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{
"email":"user@example.com",
"plan_type":"plus",
"transport":{
"extraHeaders":{"x-org-id":"org-1"},
"queryParams":{"tenant":"demo","retry":2}
}
}"#,
)
.expect("auth config ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-safe-4".to_string(),
"provider-1".to_string(),
"safe-key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["openai:cli"])),
encrypted_api_key,
Some(encrypted_auth_config),
None,
None,
None,
None,
None,
None,
)
.expect("key transport fields should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state =
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()));
let snapshot =
read_provider_transport_snapshot(&state, "provider-1", "endpoint-safe-4", "key-safe-4")
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.key.decrypted_auth_config, None);
assert_eq!(
snapshot.endpoint.base_url,
"https://api.openai.com/v1?retry=2&tenant=demo"
);
assert_eq!(
snapshot.endpoint.header_rules,
Some(serde_json::json!([
{"action":"set","key":"x-base","value":"1"},
{"action":"set","key":"x-org-id","value":"org-1"}
]))
);
assert!(supports_local_standard_transport_with_network(
&snapshot,
"openai:cli"
));
}
#[tokio::test]
async fn normalizes_json_null_transport_fields_before_local_support_checks() {
let provider = sample_provider().with_transport_fields(
true,
false,
false,
None,
Some(2),
Some(serde_json::Value::Null),
Some(20.0),
Some(8.0),
Some(serde_json::Value::Null),
);
let endpoint = StoredProviderCatalogEndpoint::new(
"endpoint-null-json".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.openai.com".to_string(),
Some(serde_json::Value::Null),
Some(serde_json::Value::Null),
Some(2),
None,
Some(serde_json::Value::Null),
Some(serde_json::Value::Null),
Some(serde_json::Value::Null),
)
.expect("endpoint transport fields should build");
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-live-openai")
.expect("api key ciphertext should build");
let key = StoredProviderCatalogKey::new(
"key-null-json".to_string(),
"provider-1".to_string(),
"safe-key".to_string(),
"api_key".to_string(),
Some(serde_json::Value::Null),
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::Value::Null),
encrypted_api_key,
None,
Some(serde_json::Value::Null),
Some(serde_json::Value::Null),
Some(serde_json::Value::Null),
None,
Some(serde_json::Value::Null),
Some(serde_json::Value::Null),
)
.expect("key transport fields should build");
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![key],
));
let state =
TestSnapshotSource::new(repository, Some(DEVELOPMENT_ENCRYPTION_KEY.to_string()));
let snapshot = read_provider_transport_snapshot(
&state,
"provider-1",
"endpoint-null-json",
"key-null-json",
)
.await
.expect("snapshot read should succeed")
.expect("snapshot should exist");
assert_eq!(snapshot.provider.proxy, None);
assert_eq!(snapshot.provider.config, None);
assert_eq!(snapshot.endpoint.header_rules, None);
assert_eq!(snapshot.endpoint.body_rules, None);
assert_eq!(snapshot.endpoint.config, None);
assert_eq!(snapshot.endpoint.format_acceptance_config, None);
assert_eq!(snapshot.endpoint.proxy, None);
assert_eq!(snapshot.key.api_formats, None);
assert_eq!(snapshot.key.allowed_models, None);
assert_eq!(snapshot.key.capabilities, None);
assert_eq!(snapshot.key.rate_multipliers, None);
assert_eq!(snapshot.key.global_priority_by_format, None);
assert_eq!(snapshot.key.proxy, None);
assert_eq!(snapshot.key.fingerprint, None);
assert!(supports_local_openai_chat_transport(&snapshot));
}
}

View File

@@ -0,0 +1,229 @@
use aether_crypto::{decrypt_python_fernet_ciphertext, looks_like_python_fernet_ciphertext};
use aether_data::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data::DataLayerError;
use super::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey, GatewayProviderTransportProvider,
};
pub(super) fn map_provider(
provider: StoredProviderCatalogProvider,
) -> GatewayProviderTransportProvider {
GatewayProviderTransportProvider {
id: provider.id,
name: provider.name,
provider_type: provider.provider_type,
website: provider.website,
is_active: provider.is_active,
keep_priority_on_conversion: provider.keep_priority_on_conversion,
enable_format_conversion: provider.enable_format_conversion,
concurrent_limit: provider.concurrent_limit,
max_retries: provider.max_retries,
proxy: normalize_optional_json(provider.proxy),
request_timeout_secs: provider.request_timeout_secs,
stream_first_byte_timeout_secs: provider.stream_first_byte_timeout_secs,
config: normalize_optional_json(provider.config),
}
}
pub(super) fn map_endpoint(
endpoint: StoredProviderCatalogEndpoint,
) -> GatewayProviderTransportEndpoint {
GatewayProviderTransportEndpoint {
id: endpoint.id,
provider_id: endpoint.provider_id,
api_format: endpoint.api_format,
api_family: endpoint.api_family,
endpoint_kind: endpoint.endpoint_kind,
is_active: endpoint.is_active,
base_url: endpoint.base_url,
header_rules: normalize_optional_json(endpoint.header_rules),
body_rules: normalize_optional_json(endpoint.body_rules),
max_retries: endpoint.max_retries,
custom_path: endpoint.custom_path,
config: normalize_optional_json(endpoint.config),
format_acceptance_config: normalize_optional_json(endpoint.format_acceptance_config),
proxy: normalize_optional_json(endpoint.proxy),
}
}
pub(super) fn map_key(
key: StoredProviderCatalogKey,
encryption_key: &str,
fallback_encryption_keys: &[String],
) -> Result<GatewayProviderTransportKey, DataLayerError> {
let decrypted_api_key = decrypt_secret(
encryption_key,
fallback_encryption_keys,
&key.encrypted_api_key,
"provider_api_keys.api_key",
)?;
let decrypted_auth_config = key
.encrypted_auth_config
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(|ciphertext| {
decrypt_secret(
encryption_key,
fallback_encryption_keys,
ciphertext,
"provider_api_keys.auth_config",
)
})
.transpose()?;
Ok(GatewayProviderTransportKey {
id: key.id,
provider_id: key.provider_id,
name: key.name,
auth_type: key.auth_type,
is_active: key.is_active,
api_formats: normalize_string_list(
normalize_optional_json(key.api_formats),
"provider_api_keys.api_formats",
)?,
allowed_models: normalize_string_list(
normalize_optional_json(key.allowed_models),
"provider_api_keys.allowed_models",
)?,
capabilities: normalize_optional_json(key.capabilities),
rate_multipliers: normalize_optional_json(key.rate_multipliers),
global_priority_by_format: normalize_optional_json(key.global_priority_by_format),
expires_at_unix_secs: key.expires_at_unix_secs,
proxy: normalize_optional_json(key.proxy),
fingerprint: normalize_optional_json(key.fingerprint),
decrypted_api_key,
decrypted_auth_config,
})
}
fn normalize_optional_json(value: Option<serde_json::Value>) -> Option<serde_json::Value> {
match value {
Some(serde_json::Value::Null) | None => None,
Some(value) => Some(value),
}
}
fn decrypt_secret(
encryption_key: &str,
fallback_encryption_keys: &[String],
ciphertext: &str,
field_name: &str,
) -> Result<String, DataLayerError> {
match decrypt_python_fernet_ciphertext(encryption_key, ciphertext) {
Ok(value) => Ok(value),
Err(_error) if should_use_plaintext_secret(ciphertext, field_name) => {
Ok(ciphertext.trim().to_string())
}
Err(error) => {
for fallback_encryption_key in fallback_encryption_keys {
if let Ok(value) =
decrypt_python_fernet_ciphertext(fallback_encryption_key, ciphertext)
{
return Ok(value);
}
}
Err(DataLayerError::UnexpectedValue(format!(
"failed to decrypt {field_name}: {error}"
)))
}
}
}
pub(super) fn fallback_encryption_keys(primary_encryption_key: &str) -> Vec<String> {
let mut keys = Vec::new();
for env_key in ["AETHER_GATEWAY_DATA_ENCRYPTION_KEY", "ENCRYPTION_KEY"] {
let Ok(value) = std::env::var(env_key) else {
continue;
};
let value = value.trim();
if value.is_empty()
|| value == primary_encryption_key
|| keys.iter().any(|existing| existing == value)
{
continue;
}
keys.push(value.to_string());
}
keys
}
fn should_use_plaintext_secret(ciphertext: &str, field_name: &str) -> bool {
let ciphertext = ciphertext.trim();
if ciphertext.is_empty() {
return false;
}
if looks_like_python_fernet_ciphertext(ciphertext) {
return false;
}
match field_name {
"provider_api_keys.api_key" => !ciphertext.starts_with('{') && !ciphertext.starts_with('['),
"provider_api_keys.auth_config" => {
ciphertext.starts_with('{') || ciphertext.starts_with('[')
}
_ => false,
}
}
fn normalize_string_list(
raw: Option<serde_json::Value>,
field_name: &str,
) -> Result<Option<Vec<String>>, DataLayerError> {
let Some(raw) = raw else {
return Ok(None);
};
normalize_string_list_value(&raw, field_name)
}
fn normalize_string_list_value(
raw: &serde_json::Value,
field_name: &str,
) -> Result<Option<Vec<String>>, DataLayerError> {
match raw {
serde_json::Value::Null => Ok(None),
serde_json::Value::Array(items) => normalize_string_list_array(items, field_name).map(Some),
serde_json::Value::String(raw) => normalize_embedded_string_list(raw, field_name),
_ => Err(DataLayerError::UnexpectedValue(format!(
"{field_name} is not a JSON array"
))),
}
}
fn normalize_embedded_string_list(
raw: &str,
field_name: &str,
) -> Result<Option<Vec<String>>, DataLayerError> {
let raw = raw.trim();
if raw.is_empty() || raw.eq_ignore_ascii_case("null") {
return Ok(None);
}
if let Ok(decoded) = serde_json::from_str::<serde_json::Value>(raw) {
return normalize_string_list_value(&decoded, field_name);
}
Ok(Some(vec![raw.to_string()]))
}
fn normalize_string_list_array(
items: &[serde_json::Value],
field_name: &str,
) -> Result<Vec<String>, DataLayerError> {
let mut values = Vec::with_capacity(items.len());
for item in items {
let Some(value) = item.as_str() else {
return Err(DataLayerError::UnexpectedValue(format!(
"{field_name} contains a non-string item"
)));
};
let value = value.trim();
if !value.is_empty() {
values.push(value.to_string());
}
}
Ok(values)
}

View File

@@ -0,0 +1,321 @@
use std::collections::BTreeMap;
use super::provider_types::is_codex_cli_backend_url;
use url::form_urlencoded;
pub fn build_openai_chat_url(upstream_base_url: &str, query: Option<&str>) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
let mut url = if trimmed.ends_with("/v1") {
format!("{trimmed}/chat/completions")
} else {
format!("{trimmed}/v1/chat/completions")
};
append_merged_query(&mut url, base_query, None, query, &[]);
url
}
pub fn build_openai_cli_url(upstream_base_url: &str, query: Option<&str>, compact: bool) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
let suffix = if compact {
"/responses/compact"
} else {
"/responses"
};
let mut url = if is_codex_cli_backend_url(trimmed)
|| trimmed.ends_with("/codex")
|| trimmed.ends_with("/v1")
{
format!("{trimmed}{suffix}")
} else {
format!("{trimmed}/v1{suffix}")
};
append_merged_query(&mut url, base_query, None, query, &[]);
url
}
pub fn build_claude_messages_url(upstream_base_url: &str, query: Option<&str>) -> String {
let (trimmed, base_query) = split_base_url_query(upstream_base_url);
let trimmed = trimmed.trim_end_matches('/');
let mut url = if trimmed.ends_with("/v1") {
format!("{trimmed}/messages")
} else {
format!("{trimmed}/v1/messages")
};
append_merged_query(&mut url, base_query, None, query, &[]);
url
}
pub fn build_gemini_content_url(
upstream_base_url: &str,
model: &str,
stream: bool,
query: Option<&str>,
) -> Option<String> {
let (trimmed_base_url, base_query) = split_base_url_query(upstream_base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
let trimmed_model = model.trim();
if trimmed_base_url.is_empty() || trimmed_model.is_empty() {
return None;
}
let operation = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
let mut url = if trimmed_base_url.ends_with("/v1beta") {
format!("{trimmed_base_url}/models/{trimmed_model}:{operation}")
} else if trimmed_base_url.contains("/v1beta/models/") {
format!("{trimmed_base_url}:{operation}")
} else {
format!("{trimmed_base_url}/v1beta/models/{trimmed_model}:{operation}")
};
append_merged_query(&mut url, base_query, None, query, &["key"]);
Some(url)
}
pub fn build_gemini_video_predict_long_running_url(
upstream_base_url: &str,
model: &str,
query: Option<&str>,
) -> Option<String> {
let (trimmed_base_url, base_query) = split_base_url_query(upstream_base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
let trimmed_model = model.trim();
if trimmed_base_url.is_empty() || trimmed_model.is_empty() {
return None;
}
let mut url = if trimmed_base_url.ends_with("/v1beta") {
format!("{trimmed_base_url}/models/{trimmed_model}:predictLongRunning")
} else if trimmed_base_url.contains("/v1beta/models/") {
format!("{trimmed_base_url}:predictLongRunning")
} else {
format!("{trimmed_base_url}/v1beta/models/{trimmed_model}:predictLongRunning")
};
append_merged_query(&mut url, base_query, None, query, &["key"]);
Some(url)
}
pub fn build_passthrough_path_url(
upstream_base_url: &str,
path: &str,
query: Option<&str>,
blocked_keys: &[&str],
) -> Option<String> {
let (trimmed_base_url, base_query) = split_base_url_query(upstream_base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
let trimmed_path = path.trim();
if trimmed_base_url.is_empty() || trimmed_path.is_empty() {
return None;
}
let (trimmed_path, path_query) = split_path_query(trimmed_path);
let normalized_base_url =
if trimmed_base_url.ends_with("/v1beta") && trimmed_path.starts_with("/v1beta") {
trimmed_base_url.trim_end_matches("/v1beta")
} else {
trimmed_base_url
};
let mut url = format!("{normalized_base_url}{trimmed_path}");
append_merged_query(&mut url, base_query, path_query, query, blocked_keys);
Some(url)
}
pub fn build_gemini_files_passthrough_url(
upstream_base_url: &str,
path: &str,
query: Option<&str>,
) -> Option<String> {
let (trimmed_base_url, base_query) = split_base_url_query(upstream_base_url);
let trimmed_base_url = trimmed_base_url.trim_end_matches('/');
let trimmed_path = path.trim();
if trimmed_base_url.is_empty() || trimmed_path.is_empty() {
return None;
}
let (trimmed_path, path_query) = split_path_query(trimmed_path);
let normalized_base_url = if trimmed_base_url.ends_with("/v1beta")
&& (trimmed_path.starts_with("/v1beta/") || trimmed_path.starts_with("/upload/v1beta/"))
{
trimmed_base_url.trim_end_matches("/v1beta")
} else {
trimmed_base_url
};
let mut url = format!("{normalized_base_url}{trimmed_path}");
append_merged_query(&mut url, base_query, path_query, query, &["key"]);
Some(url)
}
fn split_base_url_query(base_url: &str) -> (&str, Option<&str>) {
let trimmed = base_url.trim();
trimmed
.split_once('?')
.map(|(base, query)| (base, Some(query)))
.unwrap_or((trimmed, None))
}
fn split_path_query(path: &str) -> (&str, Option<&str>) {
path.split_once('?')
.map(|(path, query)| (path, Some(query)))
.unwrap_or((path, None))
}
fn append_merged_query(
url: &mut String,
base_query: Option<&str>,
path_query: Option<&str>,
request_query: Option<&str>,
blocked_keys: &[&str],
) {
let Some(query) = merge_query_layers(base_query, path_query, request_query, blocked_keys)
else {
return;
};
if url.contains('?') {
url.push('&');
} else {
url.push('?');
}
url.push_str(&query);
}
fn merge_query_layers(
base_query: Option<&str>,
path_query: Option<&str>,
request_query: Option<&str>,
blocked_keys: &[&str],
) -> Option<String> {
if blocked_keys.is_empty()
&& path_query.is_none()
&& base_query.is_none()
&& request_query
.map(str::trim)
.is_some_and(|value| !value.is_empty())
{
return request_query
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
}
let mut merged = BTreeMap::new();
for source in [base_query, path_query, request_query] {
merge_query_string(&mut merged, source, blocked_keys);
}
if merged.is_empty() {
return None;
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in merged {
serializer.append_pair(&key, &value);
}
Some(serializer.finish())
}
fn merge_query_string(
out: &mut BTreeMap<String, String>,
query: Option<&str>,
blocked_keys: &[&str],
) {
let Some(query) = query.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
for (key, value) in form_urlencoded::parse(query.as_bytes()) {
if blocked_keys
.iter()
.any(|blocked| key.as_ref().eq_ignore_ascii_case(blocked))
{
continue;
}
out.insert(key.into_owned(), value.into_owned());
}
}
#[cfg(test)]
mod tests {
use super::{
build_gemini_content_url, build_gemini_files_passthrough_url,
build_gemini_video_predict_long_running_url, build_openai_chat_url,
build_passthrough_path_url,
};
#[test]
fn merges_base_url_query_for_same_format_urls() {
assert_eq!(
build_openai_chat_url(
"https://api.openai.example/v1?tenant=demo",
Some("mode=fast&tenant=override")
),
"https://api.openai.example/v1/chat/completions?mode=fast&tenant=override"
);
}
#[test]
fn merges_base_url_query_for_dynamic_gemini_content_urls() {
assert_eq!(
build_gemini_content_url(
"https://generativelanguage.googleapis.com/v1beta?alt=sse",
"gemini-2.5-pro",
true,
Some("foo=bar&key=secret")
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse&foo=bar"
)
);
}
#[test]
fn merges_base_path_and_request_query_for_passthrough_paths() {
assert_eq!(
build_passthrough_path_url(
"https://api.openai.example/v1?tenant=demo",
"/videos/generations?variant=video",
Some("size=1024"),
&[]
)
.as_deref(),
Some(
"https://api.openai.example/v1/videos/generations?size=1024&tenant=demo&variant=video"
)
);
}
#[test]
fn merges_base_url_query_for_gemini_files_passthrough_urls() {
assert_eq!(
build_gemini_files_passthrough_url(
"https://generativelanguage.googleapis.com/v1beta?alt=media",
"/upload/v1beta/files?uploadType=resumable",
Some("key=secret&pageSize=10")
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/upload/v1beta/files?alt=media&pageSize=10&uploadType=resumable"
)
);
}
#[test]
fn merges_base_url_query_for_gemini_video_urls() {
assert_eq!(
build_gemini_video_predict_long_running_url(
"https://generativelanguage.googleapis.com/v1beta?alt=sse",
"veo-3.0-generate-preview",
Some("foo=bar&key=secret")
)
.as_deref(),
Some(
"https://generativelanguage.googleapis.com/v1beta/models/veo-3.0-generate-preview:predictLongRunning?alt=sse&foo=bar"
)
);
}
}

View File

@@ -0,0 +1,19 @@
mod auth;
mod policy;
mod url;
pub use auth::{
resolve_local_vertex_api_key_query_auth, VertexApiKeyQueryAuth, VERTEX_API_KEY_QUERY_PARAM,
};
pub use policy::{
supports_local_vertex_api_key_gemini_transport,
supports_local_vertex_api_key_gemini_transport_with_network,
supports_local_vertex_api_key_imagen_transport,
supports_local_vertex_api_key_imagen_transport_with_network,
};
pub use url::{
build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url,
VERTEX_API_KEY_BASE_URL,
};
pub const PROVIDER_TYPE: &str = "vertex_ai";

View File

@@ -0,0 +1,130 @@
use super::super::snapshot::GatewayProviderTransportSnapshot;
pub const VERTEX_API_KEY_QUERY_PARAM: &str = "key";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VertexApiKeyQueryAuth {
pub name: &'static str,
pub value: String,
}
pub fn resolve_local_vertex_api_key_query_auth(
transport: &GatewayProviderTransportSnapshot,
) -> Option<VertexApiKeyQueryAuth> {
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
{
return None;
}
if transport.key.decrypted_auth_config.is_some() {
return None;
}
if !transport
.key
.auth_type
.trim()
.eq_ignore_ascii_case("api_key")
{
return None;
}
let secret = transport.key.decrypted_api_key.trim();
if secret.is_empty() {
return None;
}
Some(VertexApiKeyQueryAuth {
name: VERTEX_API_KEY_QUERY_PARAM,
value: secret.to_string(),
})
}
#[cfg(test)]
mod tests {
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{resolve_local_vertex_api_key_query_auth, VERTEX_API_KEY_QUERY_PARAM};
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Vertex".to_string(),
provider_type: "vertex_ai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "gemini:chat".to_string(),
api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://aiplatform.googleapis.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:chat".to_string()]),
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn resolves_query_auth_for_vertex_api_key_subset() {
let auth = resolve_local_vertex_api_key_query_auth(&sample_transport())
.expect("vertex api key query auth should resolve");
assert_eq!(auth.name, VERTEX_API_KEY_QUERY_PARAM);
assert_eq!(auth.value, "vertex-secret");
}
#[test]
fn rejects_non_api_key_transport() {
let mut transport = sample_transport();
transport.key.auth_type = "service_account".to_string();
assert!(resolve_local_vertex_api_key_query_auth(&transport).is_none());
}
#[test]
fn rejects_vertex_auth_config_transport() {
let mut transport = sample_transport();
transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string());
assert!(resolve_local_vertex_api_key_query_auth(&transport).is_none());
}
}

View File

@@ -0,0 +1,204 @@
use super::super::snapshot::GatewayProviderTransportSnapshot;
use super::super::{
body_rules_are_locally_supported, header_rules_are_locally_supported,
resolve_transport_tls_profile, transport_proxy_is_locally_supported,
};
use super::auth::resolve_local_vertex_api_key_query_auth;
pub fn supports_local_vertex_api_key_gemini_transport(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_vertex_api_key_same_format_transport(
transport,
&["gemini:chat", "gemini:cli"],
false,
)
}
pub fn supports_local_vertex_api_key_gemini_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_vertex_api_key_same_format_transport(
transport,
&["gemini:chat", "gemini:cli"],
true,
)
}
pub fn supports_local_vertex_api_key_imagen_transport(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_vertex_api_key_same_format_transport(transport, &["gemini:chat"], false)
}
pub fn supports_local_vertex_api_key_imagen_transport_with_network(
transport: &GatewayProviderTransportSnapshot,
) -> bool {
supports_local_vertex_api_key_same_format_transport(transport, &["gemini:chat"], true)
}
fn supports_local_vertex_api_key_same_format_transport(
transport: &GatewayProviderTransportSnapshot,
api_formats: &[&str],
allow_network_passthrough: bool,
) -> bool {
if !transport.provider.is_active || !transport.endpoint.is_active || !transport.key.is_active {
return false;
}
if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case(super::PROVIDER_TYPE)
{
return false;
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if !api_formats
.iter()
.any(|api_format| endpoint_api_format.eq_ignore_ascii_case(api_format))
{
return false;
}
if !header_rules_are_locally_supported(transport.endpoint.header_rules.as_ref())
|| !body_rules_are_locally_supported(transport.endpoint.body_rules.as_ref())
{
return false;
}
if resolve_local_vertex_api_key_query_auth(transport).is_none() {
return false;
}
let has_custom_path = transport
.endpoint
.custom_path
.as_deref()
.is_some_and(|value: &str| !value.trim().is_empty());
let has_tls_profile = resolve_transport_tls_profile(transport)
.as_deref()
.is_some_and(|value: &str| !value.trim().is_empty());
if has_custom_path && !allow_network_passthrough {
return false;
}
if allow_network_passthrough {
if !transport_proxy_is_locally_supported(transport) {
return false;
}
if transport.key.fingerprint.is_some() && resolve_transport_tls_profile(transport).is_none()
{
return false;
}
} else if transport.provider.proxy.is_some()
|| transport.endpoint.proxy.is_some()
|| transport.key.proxy.is_some()
|| has_tls_profile
{
return false;
}
true
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::super::super::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use super::{
supports_local_vertex_api_key_gemini_transport,
supports_local_vertex_api_key_gemini_transport_with_network,
};
fn sample_transport() -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Vertex".to_string(),
provider_type: "vertex_ai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "gemini:chat".to_string(),
api_family: Some("gemini".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://aiplatform.googleapis.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["gemini:chat".to_string()]),
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "vertex-secret".to_string(),
decrypted_auth_config: None,
},
}
}
#[test]
fn supports_vertex_api_key_same_format_subset() {
assert!(supports_local_vertex_api_key_gemini_transport(
&sample_transport()
));
}
#[test]
fn supports_vertex_api_key_gemini_cli_subset() {
let mut transport = sample_transport();
transport.endpoint.api_format = "gemini:cli".to_string();
assert!(supports_local_vertex_api_key_gemini_transport(&transport));
}
#[test]
fn rejects_vertex_service_account_subset() {
let mut transport = sample_transport();
transport.key.auth_type = "service_account".to_string();
transport.key.decrypted_auth_config = Some("{\"project_id\":\"demo-project\"}".to_string());
assert!(!supports_local_vertex_api_key_gemini_transport(&transport));
}
#[test]
fn allows_network_passthrough_for_custom_path_with_local_proxy_support() {
let mut transport = sample_transport();
transport.endpoint.custom_path =
Some("/v1/publishers/google/models/gemini-2.5-pro:generateContent".to_string());
transport.key.proxy = Some(json!({"url":"http://proxy.example:8080"}));
transport.key.fingerprint = Some(json!({"tls_profile":"chrome_136"}));
assert!(!supports_local_vertex_api_key_gemini_transport(&transport));
assert!(supports_local_vertex_api_key_gemini_transport_with_network(
&transport
));
}
}

View File

@@ -0,0 +1,121 @@
use std::collections::BTreeMap;
use url::form_urlencoded;
use super::super::url::build_passthrough_path_url;
pub const VERTEX_API_KEY_BASE_URL: &str = "https://aiplatform.googleapis.com";
pub fn build_vertex_api_key_gemini_content_url(
model: &str,
stream: bool,
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
}
pub fn build_vertex_api_key_imagen_content_url(
model: &str,
stream: bool,
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
build_vertex_api_key_google_model_url(model, stream, api_key, request_query)
}
fn build_vertex_api_key_google_model_url(
model: &str,
stream: bool,
api_key: &str,
request_query: Option<&str>,
) -> Option<String> {
let trimmed_model = model.trim();
let trimmed_api_key = api_key.trim();
if trimmed_model.is_empty() || trimmed_api_key.is_empty() {
return None;
}
let action = if stream {
"streamGenerateContent"
} else {
"generateContent"
};
let path = format!("/v1/publishers/google/models/{trimmed_model}:{action}");
let merged_query = build_vertex_api_key_query(trimmed_api_key, request_query, stream);
build_passthrough_path_url(VERTEX_API_KEY_BASE_URL, &path, merged_query.as_deref(), &[])
}
fn build_vertex_api_key_query(
api_key: &str,
request_query: Option<&str>,
stream: bool,
) -> Option<String> {
let mut merged = BTreeMap::new();
merge_query_string(&mut merged, request_query);
merged.remove("beta");
merged.insert("key".to_string(), api_key.to_string());
if stream {
merged
.entry("alt".to_string())
.or_insert_with(|| "sse".to_string());
}
let mut serializer = form_urlencoded::Serializer::new(String::new());
for (key, value) in merged {
serializer.append_pair(&key, &value);
}
let query = serializer.finish();
if query.is_empty() {
None
} else {
Some(query)
}
}
fn merge_query_string(out: &mut BTreeMap<String, String>, query: Option<&str>) {
let Some(query) = query.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
for (key, value) in form_urlencoded::parse(query.as_bytes()) {
out.insert(key.into_owned(), value.into_owned());
}
}
#[cfg(test)]
mod tests {
use super::{build_vertex_api_key_gemini_content_url, build_vertex_api_key_imagen_content_url};
#[test]
fn builds_vertex_gemini_api_key_stream_url() {
assert_eq!(
build_vertex_api_key_gemini_content_url(
"gemini-2.5-pro",
true,
"vertex-secret",
Some("foo=bar&beta=v1")
)
.as_deref(),
Some(
"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.5-pro:streamGenerateContent?alt=sse&foo=bar&key=vertex-secret"
)
);
}
#[test]
fn builds_vertex_imagen_api_key_sync_url() {
assert_eq!(
build_vertex_api_key_imagen_content_url(
"imagen-3.0-generate-001",
false,
"vertex-secret",
Some("view=full")
)
.as_deref(),
Some(
"https://aiplatform.googleapis.com/v1/publishers/google/models/imagen-3.0-generate-001:generateContent?key=vertex-secret&view=full"
)
);
}
}

View File

@@ -0,0 +1,285 @@
use aether_data::repository::video_tasks::StoredVideoTask;
use aether_video_tasks_core::{
LocalVideoTaskSnapshot, LocalVideoTaskTransport, LocalVideoTaskTransportBridgeInput,
};
use async_trait::async_trait;
use super::auth::{resolve_local_gemini_auth, resolve_local_standard_auth};
use super::network::resolve_transport_execution_timeouts;
use super::policy::{supports_local_gemini_transport, supports_local_standard_transport};
use super::snapshot::GatewayProviderTransportSnapshot;
#[async_trait]
pub trait VideoTaskTransportSnapshotLookup: Send + Sync {
async fn read_video_task_provider_transport_snapshot(
&self,
provider_id: &str,
endpoint_id: &str,
key_id: &str,
) -> Result<Option<GatewayProviderTransportSnapshot>, String>;
}
pub fn resolve_local_video_task_transport(
transport: &GatewayProviderTransportSnapshot,
api_format: &str,
model_name: Option<String>,
) -> Option<LocalVideoTaskTransport> {
let api_format = api_format.trim();
let (auth_header, auth_value) = match api_format {
"openai:video" => {
if !supports_local_standard_transport(transport, api_format) {
return None;
}
resolve_local_standard_auth(transport)?
}
"gemini:video" => {
if !supports_local_gemini_transport(transport, api_format) {
return None;
}
resolve_local_gemini_auth(transport)?
}
_ => return None,
};
Some(LocalVideoTaskTransport::from_bridge_input(
LocalVideoTaskTransportBridgeInput {
upstream_base_url: transport.endpoint.base_url.clone(),
provider_name: Some(transport.provider.name.clone()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
auth_header,
auth_value,
content_type: Some("application/json".to_string()),
model_name,
proxy: None,
tls_profile: None,
timeouts: resolve_transport_execution_timeouts(transport),
},
))
}
pub async fn reconstruct_local_video_task_snapshot(
lookup: &dyn VideoTaskTransportSnapshotLookup,
task: &StoredVideoTask,
) -> Result<Option<LocalVideoTaskSnapshot>, String> {
let provider_api_format = task
.provider_api_format
.as_deref()
.unwrap_or_default()
.trim();
if !matches!(provider_api_format, "openai:video" | "gemini:video") {
return Ok(None);
}
let Some(provider_id) = task.provider_id.as_deref() else {
return Ok(None);
};
let Some(endpoint_id) = task.endpoint_id.as_deref() else {
return Ok(None);
};
let Some(key_id) = task.key_id.as_deref() else {
return Ok(None);
};
let Some(transport) = lookup
.read_video_task_provider_transport_snapshot(provider_id, endpoint_id, key_id)
.await?
else {
return Ok(None);
};
let Some(local_transport) =
resolve_local_video_task_transport(&transport, provider_api_format, task.model.clone())
else {
return Ok(None);
};
Ok(LocalVideoTaskSnapshot::from_stored_task_with_transport(
task,
local_transport,
))
}
#[cfg(test)]
mod tests {
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
use aether_video_tasks_core::LocalVideoTaskSnapshot;
use async_trait::async_trait;
use serde_json::json;
use super::{
reconstruct_local_video_task_snapshot, resolve_local_video_task_transport,
VideoTaskTransportSnapshotLookup,
};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
fn sample_transport(api_format: &str, auth_type: &str) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider One".to_string(),
provider_type: "openai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: Some(30.0),
stream_first_byte_timeout_secs: Some(5.0),
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: api_format.to_string(),
api_family: None,
endpoint_kind: None,
is_active: true,
base_url: "https://example.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "key".to_string(),
auth_type: auth_type.to_string(),
is_active: true,
api_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}
}
fn sample_stored_video_task() -> StoredVideoTask {
StoredVideoTask {
id: "task-1".to_string(),
short_id: Some("short-1".to_string()),
request_id: "request-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: Some("user".to_string()),
api_key_name: Some("key".to_string()),
external_task_id: Some("upstream-task-1".to_string()),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
client_api_format: Some("openai:video".to_string()),
provider_api_format: Some("openai:video".to_string()),
format_converted: false,
model: Some("sora".to_string()),
prompt: Some("generate".to_string()),
original_request_body: Some(json!({"prompt": "generate"})),
duration_seconds: None,
resolution: None,
aspect_ratio: None,
size: Some("1024x1024".to_string()),
status: VideoTaskStatus::Submitted,
progress_percent: 0,
progress_message: None,
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: None,
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: 1,
submitted_at_unix_secs: Some(1),
completed_at_unix_secs: None,
updated_at_unix_secs: 1,
error_code: None,
error_message: None,
video_url: None,
request_metadata: None,
}
}
struct TestLookup(Option<GatewayProviderTransportSnapshot>);
#[async_trait]
impl VideoTaskTransportSnapshotLookup for TestLookup {
async fn read_video_task_provider_transport_snapshot(
&self,
_provider_id: &str,
_endpoint_id: &str,
_key_id: &str,
) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
Ok(self.0.clone())
}
}
#[test]
fn resolves_openai_video_transport() {
let transport = resolve_local_video_task_transport(
&sample_transport("openai:video", "bearer"),
"openai:video",
Some("sora".to_string()),
)
.expect("transport");
assert_eq!(
transport.headers.get("authorization").map(String::as_str),
Some("Bearer secret")
);
assert_eq!(transport.model_name.as_deref(), Some("sora"));
assert_eq!(transport.provider_id, "provider-1");
}
#[test]
fn resolves_gemini_video_transport() {
let transport = resolve_local_video_task_transport(
&sample_transport("gemini:video", "api_key"),
"gemini:video",
Some("veo".to_string()),
)
.expect("transport");
assert_eq!(
transport.headers.get("x-goog-api-key").map(String::as_str),
Some("secret")
);
assert_eq!(transport.model_name.as_deref(), Some("veo"));
assert_eq!(transport.endpoint_id, "endpoint-1");
}
#[test]
fn rejects_mismatched_video_transport_format() {
let transport = sample_transport("openai:chat", "bearer");
assert!(resolve_local_video_task_transport(&transport, "openai:video", None).is_none());
}
#[tokio::test]
async fn reconstructs_openai_video_snapshot_via_lookup_trait() {
let lookup = TestLookup(Some(sample_transport("openai:video", "bearer")));
let snapshot = reconstruct_local_video_task_snapshot(&lookup, &sample_stored_video_task())
.await
.expect("lookup should succeed")
.expect("snapshot");
match snapshot {
LocalVideoTaskSnapshot::OpenAi(seed) => {
assert_eq!(seed.transport.provider_id, "provider-1");
assert_eq!(seed.transport.model_name.as_deref(), Some("sora"));
}
LocalVideoTaskSnapshot::Gemini(_) => panic!("expected openai snapshot"),
}
}
}

View File

@@ -9,8 +9,11 @@ description = "Shared runtime/bootstrap helpers for Aether Rust services"
[dependencies]
async-stream.workspace = true
axum = { version = "0.8" }
chrono.workspace = true
futures-util.workspace = true
redis.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true

View File

@@ -2,7 +2,7 @@ use crate::config::ServiceRuntimeConfig;
use crate::error::RuntimeBootstrapError;
pub fn init_service_runtime(config: ServiceRuntimeConfig) -> Result<(), RuntimeBootstrapError> {
crate::tracing::init_tracing(config)?;
crate::tracing::init_tracing(config.clone())?;
crate::metrics::init_metrics(config);
Ok(())
}

View File

@@ -1,6 +1,6 @@
use crate::observability::ServiceObservabilityConfig;
use crate::observability::{FileLoggingConfig, LogDestination, ServiceObservabilityConfig};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServiceRuntimeConfig {
pub service_name: &'static str,
pub default_log_filter: &'static str,
@@ -21,6 +21,26 @@ impl ServiceRuntimeConfig {
self
}
pub const fn with_log_destination(mut self, log_destination: LogDestination) -> Self {
self.observability.log_destination = log_destination;
self
}
pub fn with_file_logging(mut self, file_logging: FileLoggingConfig) -> Self {
self.observability.file_logging = Some(file_logging);
self
}
pub fn with_node_role(mut self, node_role: impl Into<String>) -> Self {
self.observability.node_role = Some(node_role.into());
self
}
pub fn with_instance_id(mut self, instance_id: impl Into<String>) -> Self {
self.observability.instance_id = Some(instance_id.into());
self
}
pub const fn with_metrics_namespace(mut self, metrics_namespace: &'static str) -> Self {
self.observability.metrics_namespace = metrics_namespace;
self

View File

@@ -7,6 +7,7 @@ mod error;
pub mod metrics;
mod observability;
pub mod queue;
pub mod redaction;
pub mod shutdown;
pub mod task;
mod tracing;
@@ -23,9 +24,14 @@ pub use distributed::{
};
pub use error::RuntimeBootstrapError;
pub use metrics::{prometheus_response, service_up_sample, MetricKind, MetricLabel, MetricSample};
pub use observability::ServiceObservabilityConfig;
pub use observability::{
FileLoggingConfig, LogDestination, LogRotation, ServiceObservabilityConfig,
};
pub use queue::{
bounded_queue, BoundedQueueReceiver, BoundedQueueSender, QueueSendError, QueueSnapshot,
};
pub use redaction::{summarize_text_payload, TextPayloadSummary};
pub use shutdown::wait_for_shutdown_signal;
pub use tracing::{init_reloadable_tracing, LogFormat, LogReloader};
pub use tracing::{
init_reloadable_service_tracing, init_reloadable_tracing, LogFormat, LogReloader,
};

View File

@@ -1,9 +1,57 @@
use crate::tracing::LogFormat;
use std::path::PathBuf;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogDestination {
Stdout,
File,
Both,
}
impl LogDestination {
pub const fn needs_file_sink(self) -> bool {
matches!(self, Self::File | Self::Both)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum LogRotation {
Hourly,
Daily,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FileLoggingConfig {
pub dir: PathBuf,
pub rotation: LogRotation,
pub retention_days: u64,
pub max_files: usize,
}
impl FileLoggingConfig {
pub fn new(
dir: impl Into<PathBuf>,
rotation: LogRotation,
retention_days: u64,
max_files: usize,
) -> Self {
Self {
dir: dir.into(),
rotation,
retention_days,
max_files,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ServiceObservabilityConfig {
pub log_format: LogFormat,
pub metrics_namespace: &'static str,
pub log_destination: LogDestination,
pub file_logging: Option<FileLoggingConfig>,
pub node_role: Option<String>,
pub instance_id: Option<String>,
}
impl ServiceObservabilityConfig {
@@ -11,6 +59,30 @@ impl ServiceObservabilityConfig {
Self {
log_format,
metrics_namespace,
log_destination: LogDestination::Stdout,
file_logging: None,
node_role: None,
instance_id: None,
}
}
pub const fn with_log_destination(mut self, log_destination: LogDestination) -> Self {
self.log_destination = log_destination;
self
}
pub fn with_file_logging(mut self, file_logging: FileLoggingConfig) -> Self {
self.file_logging = Some(file_logging);
self
}
pub fn with_node_role(mut self, node_role: impl Into<String>) -> Self {
self.node_role = Some(node_role.into());
self
}
pub fn with_instance_id(mut self, instance_id: impl Into<String>) -> Self {
self.instance_id = Some(instance_id.into());
self
}
}

View File

@@ -0,0 +1,36 @@
use std::fmt::Write as _;
use sha2::{Digest, Sha256};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TextPayloadSummary {
pub bytes: usize,
pub sha256: String,
}
pub fn summarize_text_payload(text: &str) -> TextPayloadSummary {
let digest = Sha256::digest(text.as_bytes());
let mut sha256 = String::with_capacity(digest.len() * 2);
for byte in digest {
write!(&mut sha256, "{byte:02x}").expect("writing to string should not fail");
}
TextPayloadSummary {
bytes: text.len(),
sha256,
}
}
#[cfg(test)]
mod tests {
use super::summarize_text_payload;
#[test]
fn summarizes_text_payload_without_exposing_content() {
let summary = summarize_text_payload("secret-body");
assert_eq!(summary.bytes, 11);
assert_eq!(
summary.sha256,
"7c3029502007b2beae470d090221f0a8f7708be361be4662f4b649426e5767b3"
);
}
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,16 @@
[package]
name = "aether-scheduler-core"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Pure scheduler health and quota logic extracted from aether-gateway"
[dependencies]
aether-contracts.workspace = true
aether-data.workspace = true
aether-wallet.workspace = true
regex.workspace = true
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true

View File

@@ -0,0 +1,173 @@
use sha2::{Digest, Sha256};
use crate::SchedulerMinimalCandidateSelectionCandidate;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SchedulerAffinityTarget {
pub provider_id: String,
pub endpoint_id: String,
pub key_id: String,
}
pub fn build_scheduler_affinity_cache_key_for_api_key_id(
api_key_id: &str,
api_format: &str,
global_model_name: &str,
) -> Option<String> {
let api_key_id = api_key_id.trim();
if api_key_id.is_empty() {
return None;
}
let api_format = crate::normalize_api_format(api_format);
let global_model_name = global_model_name.trim();
if api_format.is_empty() || global_model_name.is_empty() {
return None;
}
Some(format!(
"scheduler_affinity:{api_key_id}:{api_format}:{global_model_name}"
))
}
pub fn compare_affinity_order(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
affinity_key: Option<&str>,
) -> std::cmp::Ordering {
let Some(affinity_key) = affinity_key else {
return std::cmp::Ordering::Equal;
};
candidate_affinity_hash(affinity_key, left).cmp(&candidate_affinity_hash(affinity_key, right))
}
pub fn candidate_affinity_hash(
affinity_key: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> u64 {
let mut hasher = Sha256::new();
hasher.update(affinity_key.as_bytes());
hasher.update(b":");
hasher.update(candidate.provider_id.as_bytes());
hasher.update(b":");
hasher.update(candidate.endpoint_id.as_bytes());
hasher.update(b":");
hasher.update(candidate.key_id.as_bytes());
let digest = hasher.finalize();
u64::from_be_bytes([
digest[0], digest[1], digest[2], digest[3], digest[4], digest[5], digest[6], digest[7],
])
}
pub fn matches_affinity_target(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
target: &SchedulerAffinityTarget,
) -> bool {
candidate.provider_id == target.provider_id
&& candidate.endpoint_id == target.endpoint_id
&& candidate.key_id == target.key_id
}
pub fn candidate_key(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> (String, String, String) {
(
candidate.provider_id.clone(),
candidate.endpoint_id.clone(),
candidate.key_id.clone(),
)
}
#[cfg(test)]
mod tests {
use super::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
compare_affinity_order, matches_affinity_target, SchedulerAffinityTarget,
};
use crate::SchedulerMinimalCandidateSelectionCandidate;
fn sample_candidate(id: &str) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "custom".to_string(),
provider_priority: 1,
endpoint_id: format!("endpoint-{id}"),
endpoint_api_format: "openai:chat".to_string(),
key_id: format!("key-{id}"),
key_name: format!("Key {id}"),
key_auth_type: "api_key".to_string(),
key_internal_priority: 1,
key_global_priority_for_format: Some(1),
key_capabilities: None,
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
mapping_matched_model: None,
}
}
#[test]
fn builds_normalized_scheduler_affinity_cache_key() {
assert_eq!(
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "OPENAI:CHAT", "gpt-5"),
Some("scheduler_affinity:api-key-1:openai:chat:gpt-5".to_string())
);
}
#[test]
fn rejects_blank_affinity_key_components() {
assert_eq!(
build_scheduler_affinity_cache_key_for_api_key_id("", "openai:chat", "gpt-5"),
None
);
assert_eq!(
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "", "gpt-5"),
None
);
assert_eq!(
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", ""),
None
);
}
#[test]
fn affinity_hash_and_order_are_candidate_specific() {
let left = sample_candidate("1");
let right = sample_candidate("2");
assert_ne!(
candidate_affinity_hash("api-key-1", &left),
candidate_affinity_hash("api-key-1", &right)
);
assert_ne!(
compare_affinity_order(&left, &right, Some("api-key-1")),
std::cmp::Ordering::Equal
);
assert_eq!(
compare_affinity_order(&left, &right, None),
std::cmp::Ordering::Equal
);
}
#[test]
fn affinity_target_and_candidate_key_reuse_candidate_identity() {
let candidate = sample_candidate("1");
let target = SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
endpoint_id: candidate.endpoint_id.clone(),
key_id: candidate.key_id.clone(),
};
assert!(matches_affinity_target(&candidate, &target));
assert_eq!(
candidate_key(&candidate),
(
"provider-1".to_string(),
"endpoint-1".to_string(),
"key-1".to_string()
)
);
}
}

View File

@@ -0,0 +1,108 @@
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct SchedulerAuthConstraints {
pub allowed_providers: Option<Vec<String>>,
pub allowed_api_formats: Option<Vec<String>>,
pub allowed_models: Option<Vec<String>>,
}
pub fn auth_constraints_allow_provider(
constraints: Option<&SchedulerAuthConstraints>,
provider_id: &str,
provider_name: &str,
) -> bool {
let Some(allowed) =
constraints.and_then(|constraints| constraints.allowed_providers.as_deref())
else {
return true;
};
allowed.iter().any(|value| {
value.trim().eq_ignore_ascii_case(provider_id.trim())
|| value.trim().eq_ignore_ascii_case(provider_name.trim())
})
}
pub fn auth_constraints_allow_api_format(
constraints: Option<&SchedulerAuthConstraints>,
api_format: &str,
) -> bool {
let Some(allowed) =
constraints.and_then(|constraints| constraints.allowed_api_formats.as_deref())
else {
return true;
};
allowed
.iter()
.any(|value| crate::normalize_api_format(value) == api_format)
}
pub fn auth_constraints_allow_model(
constraints: Option<&SchedulerAuthConstraints>,
requested_model_name: &str,
resolved_global_model_name: &str,
) -> bool {
let Some(allowed) = constraints.and_then(|constraints| constraints.allowed_models.as_deref())
else {
return true;
};
allowed
.iter()
.any(|value| value == requested_model_name || value == resolved_global_model_name)
}
#[cfg(test)]
mod tests {
use super::{
auth_constraints_allow_api_format, auth_constraints_allow_model,
auth_constraints_allow_provider, SchedulerAuthConstraints,
};
fn sample_constraints() -> SchedulerAuthConstraints {
SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string(), "OpenAI".to_string()]),
allowed_api_formats: Some(vec!["OPENAI:CHAT".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
}
}
#[test]
fn constraints_allow_matching_provider_identifier_or_name() {
let constraints = sample_constraints();
assert!(auth_constraints_allow_provider(
Some(&constraints),
"provider-1",
"other"
));
assert!(auth_constraints_allow_provider(
Some(&constraints),
"other",
"openai"
));
assert!(!auth_constraints_allow_provider(
Some(&constraints),
"other",
"other"
));
}
#[test]
fn constraints_normalize_api_formats_and_models() {
let constraints = sample_constraints();
assert!(auth_constraints_allow_api_format(
Some(&constraints),
"openai:chat"
));
assert!(auth_constraints_allow_model(
Some(&constraints),
"gpt-5",
"gpt-5"
));
assert!(!auth_constraints_allow_model(
Some(&constraints),
"gpt-4.1",
"gpt-4.1"
));
}
}

View File

@@ -0,0 +1,708 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use aether_data::DataLayerError;
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub struct SchedulerMinimalCandidateSelectionCandidate {
pub provider_id: String,
pub provider_name: String,
pub provider_type: String,
pub provider_priority: i32,
pub endpoint_id: String,
pub endpoint_api_format: String,
pub key_id: String,
pub key_name: String,
pub key_auth_type: String,
pub key_internal_priority: i32,
pub key_global_priority_for_format: Option<i32>,
pub key_capabilities: Option<serde_json::Value>,
pub model_id: String,
pub global_model_id: String,
pub global_model_name: String,
pub selected_provider_model_name: String,
pub mapping_matched_model: Option<String>,
}
pub fn candidate_supports_required_capability(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
required_capability: &str,
) -> bool {
let required_capability = required_capability.trim();
if required_capability.is_empty() {
return true;
}
let Some(capabilities) = candidate.key_capabilities.as_ref() else {
return false;
};
if let Some(object) = capabilities.as_object() {
return object.iter().any(|(key, value)| {
key.eq_ignore_ascii_case(required_capability)
&& match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => {
value.as_i64().is_some_and(|value| value > 0)
}
_ => false,
}
});
}
if let Some(items) = capabilities.as_array() {
return items.iter().any(|value| {
value
.as_str()
.is_some_and(|value| value.eq_ignore_ascii_case(required_capability))
});
}
false
}
pub fn auth_api_key_concurrency_limit_reached(
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
api_key_id: &str,
concurrent_limit: usize,
) -> bool {
if api_key_id.trim().is_empty() || concurrent_limit == 0 {
return false;
}
crate::count_recent_active_requests_for_api_key(recent_candidates, api_key_id, now_unix_secs)
>= concurrent_limit
}
pub fn build_minimal_candidate_selection(
rows: Vec<StoredMinimalCandidateSelectionRow>,
normalized_api_format: &str,
requested_model_name: &str,
resolved_global_model_name: &str,
require_streaming: bool,
auth_constraints: Option<&crate::SchedulerAuthConstraints>,
affinity_key: Option<&str>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Ok(Vec::new());
}
if !crate::auth_constraints_allow_model(
auth_constraints,
requested_model_name,
resolved_global_model_name,
) {
return Ok(Vec::new());
}
let mut candidates = Vec::new();
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
&row.provider_id,
&row.provider_name,
) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
let Some((selected_provider_model_name, mapping_matched_model)) =
crate::resolve_provider_model_name(&row, requested_model_name, normalized_api_format)
else {
continue;
};
candidates.push(SchedulerMinimalCandidateSelectionCandidate {
provider_id: row.provider_id,
provider_name: row.provider_name,
provider_type: row.provider_type,
provider_priority: row.provider_priority,
endpoint_id: row.endpoint_id,
endpoint_api_format: row.endpoint_api_format,
key_id: row.key_id,
key_name: row.key_name,
key_auth_type: row.key_auth_type,
key_internal_priority: row.key_internal_priority,
key_global_priority_for_format: crate::extract_global_priority_for_format(
row.key_global_priority_by_format.as_ref(),
normalized_api_format,
)?,
key_capabilities: row.key_capabilities,
model_id: row.model_id,
global_model_id: row.global_model_id,
global_model_name: row.global_model_name,
selected_provider_model_name,
mapping_matched_model,
});
}
candidates.sort_by(|left, right| {
left.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
});
Ok(candidates)
}
pub fn collect_global_model_names_for_required_capability(
rows: Vec<StoredMinimalCandidateSelectionRow>,
normalized_api_format: &str,
required_capability: &str,
require_streaming: bool,
auth_constraints: Option<&crate::SchedulerAuthConstraints>,
) -> Vec<String> {
if normalized_api_format.is_empty() || required_capability.trim().is_empty() {
return Vec::new();
}
if !crate::auth_constraints_allow_api_format(auth_constraints, normalized_api_format) {
return Vec::new();
}
let mut model_names = BTreeSet::new();
for row in rows {
if !crate::auth_constraints_allow_provider(
auth_constraints,
&row.provider_id,
&row.provider_name,
) {
continue;
}
if !crate::row_supports_required_capability(&row, required_capability) {
continue;
}
if require_streaming && !row.supports_streaming() {
continue;
}
if !crate::auth_constraints_allow_model(
auth_constraints,
&row.global_model_name,
&row.global_model_name,
) {
continue;
}
model_names.insert(row.global_model_name);
}
model_names.into_iter().collect()
}
pub fn collect_selectable_candidates_from_keys(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
selectable_keys: &BTreeSet<(String, String, String)>,
cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
) -> Vec<SchedulerMinimalCandidateSelectionCandidate> {
let mut selected = Vec::new();
let mut emitted_keys = BTreeSet::new();
if let Some(target) = cached_affinity_target {
if let Some(candidate) = candidates
.iter()
.find(|candidate| crate::matches_affinity_target(candidate, target))
.cloned()
{
let key = crate::candidate_key(&candidate);
if selectable_keys.contains(&key) && emitted_keys.insert(key) {
selected.push(candidate);
}
}
}
for candidate in candidates {
let key = crate::candidate_key(&candidate);
if !selectable_keys.contains(&key) || !emitted_keys.insert(key) {
continue;
}
selected.push(candidate);
}
selected
}
pub fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
affinity_key: Option<&str>,
) {
candidates.sort_by(|left, right| {
left.key_global_priority_for_format
.unwrap_or(i32::MAX)
.cmp(&right.key_global_priority_for_format.unwrap_or(i32::MAX))
.then_with(|| compare_provider_key_health_order(left, right, provider_key_rpm_states))
.then_with(|| crate::compare_affinity_order(left, right, affinity_key))
.then(left.provider_priority.cmp(&right.provider_priority))
.then(left.key_internal_priority.cmp(&right.key_internal_priority))
.then(left.provider_id.cmp(&right.provider_id))
.then(left.endpoint_id.cmp(&right.endpoint_id))
.then(left.key_id.cmp(&right.key_id))
.then(
left.selected_provider_model_name
.cmp(&right.selected_provider_model_name),
)
});
}
pub fn candidate_is_selectable_with_runtime_state(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
recent_candidates: &[StoredRequestCandidate],
provider_concurrent_limits: &BTreeMap<String, usize>,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
now_unix_secs: u64,
cached_affinity_target: Option<&crate::SchedulerAffinityTarget>,
provider_quota_blocks_requests: bool,
rpm_reset_at: Option<u64>,
) -> bool {
if provider_quota_blocks_requests {
return false;
}
if crate::is_candidate_in_recent_failure_cooldown(
recent_candidates,
candidate.provider_id.as_str(),
candidate.endpoint_id.as_str(),
candidate.key_id.as_str(),
now_unix_secs,
) {
return false;
}
if provider_concurrent_limits
.get(&candidate.provider_id)
.is_some_and(|limit| {
crate::count_recent_active_requests_for_provider(
recent_candidates,
candidate.provider_id.as_str(),
now_unix_secs,
) >= *limit
})
{
return false;
}
let is_cached_user = cached_affinity_target
.is_some_and(|target| crate::matches_affinity_target(candidate, target));
if let Some(provider_key) = provider_key_rpm_states.get(&candidate.key_id) {
if crate::is_provider_key_circuit_open(provider_key, candidate.endpoint_api_format.as_str())
{
return false;
}
if crate::provider_key_health_score(provider_key, candidate.endpoint_api_format.as_str())
.is_some_and(|score| score <= 0.0)
{
return false;
}
if !crate::provider_key_rpm_allows_request_since(
provider_key,
recent_candidates,
now_unix_secs,
is_cached_user,
rpm_reset_at,
) {
return false;
}
}
true
}
fn compare_provider_key_health_order(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> std::cmp::Ordering {
let left_bucket = candidate_provider_key_health_bucket(left, provider_key_rpm_states);
let right_bucket = candidate_provider_key_health_bucket(right, provider_key_rpm_states);
right_bucket.cmp(&left_bucket).then_with(|| {
let left_score = candidate_provider_key_health_score(left, provider_key_rpm_states);
let right_score = candidate_provider_key_health_score(right, provider_key_rpm_states);
right_score
.partial_cmp(&left_score)
.unwrap_or(std::cmp::Ordering::Equal)
})
}
fn candidate_provider_key_health_bucket(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> Option<crate::ProviderKeyHealthBucket> {
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|key| {
crate::provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
})
}
fn candidate_provider_key_health_score(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
) -> f64 {
provider_key_rpm_states
.get(&candidate.key_id)
.and_then(|key| {
crate::effective_provider_key_health_score(key, candidate.endpoint_api_format.as_str())
})
.unwrap_or(1.0)
}
#[cfg(test)]
mod tests {
use std::collections::{BTreeMap, BTreeSet};
use aether_data::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use super::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability,
collect_selectable_candidates_from_keys, reorder_candidates_by_scheduler_health,
SchedulerMinimalCandidateSelectionCandidate,
};
use crate::SchedulerAuthConstraints;
fn sample_row(id: &str) -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "custom".to_string(),
provider_priority: 10,
provider_is_active: true,
endpoint_id: format!("endpoint-{id}"),
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: format!("key-{id}"),
key_name: format!("prod-{id}"),
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: Some(serde_json::json!({"cache_1h": true})),
key_internal_priority: 50,
key_global_priority_by_format: Some(serde_json::json!({"openai:chat": 2})),
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
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: format!("gpt-5-upstream-{id}"),
model_provider_model_mappings: Some(vec![StoredProviderModelMapping {
name: format!("gpt-5-canary-{id}"),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]),
model_supports_streaming: None,
model_is_active: true,
model_is_available: true,
}
}
fn sample_candidate(
id: &str,
capabilities: Option<serde_json::Value>,
) -> SchedulerMinimalCandidateSelectionCandidate {
SchedulerMinimalCandidateSelectionCandidate {
provider_id: format!("provider-{id}"),
provider_name: format!("Provider {id}"),
provider_type: "openai".to_string(),
provider_priority: 0,
endpoint_id: format!("endpoint-{id}"),
endpoint_api_format: "openai:chat".to_string(),
key_id: format!("key-{id}"),
key_name: format!("key-{id}"),
key_auth_type: "bearer".to_string(),
key_internal_priority: 0,
key_global_priority_for_format: None,
key_capabilities: capabilities,
model_id: format!("model-{id}"),
global_model_id: format!("global-model-{id}"),
global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
mapping_matched_model: None,
}
}
fn sample_key(id: &str, health_score: f64) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
format!("key-{id}"),
format!("provider-{id}"),
format!("key-{id}"),
"api_key".to_string(),
None,
true,
)
.expect("provider key should build");
key.health_by_format = Some(serde_json::json!({
"openai:chat": {
"health_score": health_score
}
}));
key
}
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
created_at_unix_secs: i64,
) -> StoredRequestCandidate {
let finished_at_unix_secs = match status {
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming => None,
_ => Some(created_at_unix_secs),
};
StoredRequestCandidate::new(
id.to_string(),
format!("req-{id}"),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
finished_at_unix_secs,
)
.expect("candidate should build")
}
#[test]
fn reads_required_capability_from_object_and_array_forms() {
assert!(candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!({"vision": true}))),
"vision"
));
assert!(candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!(["vision", "tools"]))),
"tools"
));
assert!(!candidate_supports_required_capability(
&sample_candidate("1", Some(serde_json::json!({"vision": false}))),
"vision"
));
}
#[test]
fn builds_minimal_candidate_selection_with_auth_constraints() {
let mut disallowed = sample_row("2");
disallowed.provider_id = "provider-blocked".to_string();
disallowed.provider_name = "Blocked".to_string();
let constraints = SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string()]),
allowed_api_formats: Some(vec!["OPENAI:CHAT".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let candidates = build_minimal_candidate_selection(
vec![sample_row("1"), disallowed],
"openai:chat",
"gpt-5",
"gpt-5",
false,
Some(&constraints),
None,
)
.expect("candidate selection should build");
assert_eq!(candidates.len(), 1);
assert_eq!(candidates[0].provider_id, "provider-1");
assert_eq!(candidates[0].selected_provider_model_name, "gpt-5-canary-1");
}
#[test]
fn collects_global_model_names_for_required_capability_with_auth_constraints() {
let mut disallowed = sample_row("2");
disallowed.global_model_name = "gpt-4.1".to_string();
disallowed.provider_id = "provider-blocked".to_string();
disallowed.provider_name = "Blocked".to_string();
let constraints = SchedulerAuthConstraints {
allowed_providers: Some(vec!["provider-1".to_string()]),
allowed_api_formats: Some(vec!["openai:chat".to_string()]),
allowed_models: Some(vec!["gpt-5".to_string()]),
};
let model_names = collect_global_model_names_for_required_capability(
vec![sample_row("1"), disallowed],
"openai:chat",
"cache_1h",
false,
Some(&constraints),
);
assert_eq!(model_names, vec!["gpt-5".to_string()]);
}
#[test]
fn reorders_candidates_by_health_before_affinity_tiebreak() {
let mut candidates = vec![
sample_candidate("1", None),
sample_candidate("2", None),
sample_candidate("3", None),
];
let provider_key_rpm_states = BTreeMap::from([
("key-1".to_string(), sample_key("1", 0.95)),
("key-2".to_string(), sample_key("2", 0.40)),
("key-3".to_string(), sample_key("3", 0.95)),
]);
reorder_candidates_by_scheduler_health(
&mut candidates,
&provider_key_rpm_states,
Some("api-key-1"),
);
assert_ne!(candidates[0].key_id, "key-2");
assert_ne!(candidates[1].key_id, "key-2");
assert_eq!(candidates[2].key_id, "key-2");
}
#[test]
fn collects_selectable_candidates_with_affinity_priority_and_dedup() {
let candidates = vec![
sample_candidate("1", None),
sample_candidate("2", None),
sample_candidate("1", None),
];
let selectable_keys = BTreeSet::from([
(
"provider-1".to_string(),
"endpoint-1".to_string(),
"key-1".to_string(),
),
(
"provider-2".to_string(),
"endpoint-2".to_string(),
"key-2".to_string(),
),
]);
let selected = collect_selectable_candidates_from_keys(
candidates,
&selectable_keys,
Some(&crate::SchedulerAffinityTarget {
provider_id: "provider-2".to_string(),
endpoint_id: "endpoint-2".to_string(),
key_id: "key-2".to_string(),
}),
);
assert_eq!(selected.len(), 2);
assert_eq!(selected[0].key_id, "key-2");
assert_eq!(selected[1].key_id, "key-1");
}
#[test]
fn candidate_selectability_respects_provider_concurrency_limit() {
let recent_candidates = vec![stored_candidate("one", RequestCandidateStatus::Pending, 95)];
let provider_concurrent_limits = BTreeMap::from([("provider-1".to_string(), 1)]);
assert!(!candidate_is_selectable_with_runtime_state(
&sample_candidate("1", None),
&recent_candidates,
&provider_concurrent_limits,
&BTreeMap::new(),
100,
None,
false,
None,
));
}
#[test]
fn candidate_selectability_rejects_quota_or_zero_health() {
let provider_key_rpm_states = BTreeMap::from([("key-1".to_string(), sample_key("1", 0.0))]);
assert!(!candidate_is_selectable_with_runtime_state(
&sample_candidate("1", None),
&[],
&BTreeMap::new(),
&provider_key_rpm_states,
100,
None,
false,
None,
));
assert!(!candidate_is_selectable_with_runtime_state(
&sample_candidate("1", None),
&[],
&BTreeMap::new(),
&BTreeMap::new(),
100,
None,
true,
None,
));
}
#[test]
fn detects_auth_api_key_concurrency_limit_from_recent_active_requests() {
let recent_candidates = vec![StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build")];
assert!(auth_api_key_concurrency_limit_reached(
&recent_candidates,
100,
"api-key-1",
1,
));
assert!(!auth_api_key_concurrency_limit_reached(
&recent_candidates,
100,
"api-key-1",
2,
));
}
}

View File

@@ -0,0 +1,948 @@
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
const FAILURE_COOLDOWN_WINDOW_SECS: u64 = 60;
const FAILURE_COOLDOWN_THRESHOLD: usize = 2;
const ACTIVE_REQUEST_WINDOW_SECS: u64 = 300;
pub const PROVIDER_KEY_RPM_WINDOW_SECS: u64 = 60;
const PROBE_PHASE_REQUESTS: u32 = 100;
const PROBE_RESERVATION_RATIO: f64 = 0.1;
const STABLE_MIN_RESERVATION_RATIO: f64 = 0.1;
const STABLE_MAX_RESERVATION_RATIO: f64 = 0.35;
const SUCCESS_COUNT_FOR_FULL_CONFIDENCE: u32 = 50;
const COOLDOWN_HOURS_FOR_FULL_CONFIDENCE: f64 = 24.0;
const LOW_LOAD_THRESHOLD: f64 = 0.5;
const HIGH_LOAD_THRESHOLD: f64 = 0.8;
const ENFORCEMENT_CONFIDENCE_THRESHOLD: f64 = 0.6;
const HEALTH_DEGRADED_THRESHOLD: f64 = 0.8;
const HEALTH_LOW_THRESHOLD: f64 = 0.5;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ProviderKeyHealthBucket {
Low,
Degraded,
Healthy,
}
impl ProviderKeyHealthBucket {
fn from_score(score: f64) -> Self {
let score = score.clamp(0.0, 1.0);
if score < HEALTH_LOW_THRESHOLD {
return Self::Low;
}
if score < HEALTH_DEGRADED_THRESHOLD {
return Self::Degraded;
}
Self::Healthy
}
}
pub fn is_candidate_in_recent_failure_cooldown(
recent_candidates: &[StoredRequestCandidate],
provider_id: &str,
endpoint_id: &str,
key_id: &str,
now_unix_secs: u64,
) -> bool {
let mut recent_failures = 0usize;
for candidate in recent_candidates {
if candidate.provider_id.as_deref() != Some(provider_id)
|| candidate.endpoint_id.as_deref() != Some(endpoint_id)
|| candidate.key_id.as_deref() != Some(key_id)
{
continue;
}
let observed_at_unix_secs = candidate
.finished_at_unix_secs
.or(candidate.started_at_unix_secs)
.unwrap_or(candidate.created_at_unix_secs);
if now_unix_secs.saturating_sub(observed_at_unix_secs) > FAILURE_COOLDOWN_WINDOW_SECS {
continue;
}
match candidate.status {
RequestCandidateStatus::Success => return false,
RequestCandidateStatus::Failed | RequestCandidateStatus::Cancelled => {
recent_failures += 1;
if recent_failures >= FAILURE_COOLDOWN_THRESHOLD {
return true;
}
}
RequestCandidateStatus::Available
| RequestCandidateStatus::Unused
| RequestCandidateStatus::Pending
| RequestCandidateStatus::Streaming
| RequestCandidateStatus::Skipped => {}
}
}
false
}
pub fn count_recent_active_requests_for_provider(
recent_candidates: &[StoredRequestCandidate],
provider_id: &str,
now_unix_secs: u64,
) -> usize {
recent_candidates
.iter()
.filter(|candidate| candidate.provider_id.as_deref() == Some(provider_id))
.filter(|candidate| is_recently_active(candidate, now_unix_secs))
.count()
}
pub fn count_recent_active_requests_for_api_key(
recent_candidates: &[StoredRequestCandidate],
api_key_id: &str,
now_unix_secs: u64,
) -> usize {
recent_candidates
.iter()
.filter(|candidate| candidate.api_key_id.as_deref() == Some(api_key_id))
.filter(|candidate| is_recently_active(candidate, now_unix_secs))
.count()
}
pub fn effective_provider_key_rpm_limit(
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) -> Option<usize> {
if let Some(limit) = key.rpm_limit.filter(|limit| *limit > 0) {
return usize::try_from(limit).ok();
}
let learned_limit = key
.learned_rpm_limit
.filter(|limit| *limit > 0)
.and_then(|limit| usize::try_from(limit).ok())?;
if provider_key_reservation_confidence(key, now_unix_secs) < ENFORCEMENT_CONFIDENCE_THRESHOLD {
return None;
}
Some(learned_limit)
}
pub fn count_recent_rpm_requests_for_provider_key(
recent_candidates: &[StoredRequestCandidate],
key_id: &str,
now_unix_secs: u64,
) -> usize {
count_recent_rpm_requests_for_provider_key_since(recent_candidates, key_id, now_unix_secs, None)
}
pub fn count_recent_rpm_requests_for_provider_key_since(
recent_candidates: &[StoredRequestCandidate],
key_id: &str,
now_unix_secs: u64,
reset_after_unix_secs: Option<u64>,
) -> usize {
let mut attempted_count = 0usize;
let mut max_observed = 0usize;
for candidate in recent_candidates {
if candidate.key_id.as_deref() != Some(key_id) {
continue;
}
if !is_recent_rpm_observation(candidate, now_unix_secs) {
continue;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
if reset_after_unix_secs.is_some_and(|reset_after| observed_at_unix_secs <= reset_after) {
continue;
}
attempted_count += 1;
max_observed = max_observed.max(candidate.concurrent_requests.unwrap_or_default() as usize);
}
max_observed.max(attempted_count)
}
pub fn provider_key_rpm_allows_request(
key: &StoredProviderCatalogKey,
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
is_cached_user: bool,
) -> bool {
provider_key_rpm_allows_request_since(
key,
recent_candidates,
now_unix_secs,
is_cached_user,
None,
)
}
pub fn provider_key_rpm_allows_request_since(
key: &StoredProviderCatalogKey,
recent_candidates: &[StoredRequestCandidate],
now_unix_secs: u64,
is_cached_user: bool,
reset_after_unix_secs: Option<u64>,
) -> bool {
let Some(effective_limit) = effective_provider_key_rpm_limit(key, now_unix_secs) else {
return true;
};
if effective_limit == 0 {
return false;
}
let current_usage = count_recent_rpm_requests_for_provider_key_since(
recent_candidates,
key.id.as_str(),
now_unix_secs,
reset_after_unix_secs,
);
if is_cached_user {
return current_usage < effective_limit;
}
let available_for_new = available_provider_key_rpm_slots_for_new_user(
key,
current_usage,
effective_limit,
now_unix_secs,
);
current_usage < available_for_new
}
pub fn provider_key_health_score(key: &StoredProviderCatalogKey, api_format: &str) -> Option<f64> {
let score = key
.health_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
.and_then(|payload| payload.get("health_score"))
.and_then(json_value_as_f64)?;
Some(score.clamp(0.0, 1.0))
}
pub fn aggregate_provider_key_health_score(key: &StoredProviderCatalogKey) -> Option<f64> {
let health_by_format = key.health_by_format.as_ref()?.as_object()?;
let mut scores = Vec::new();
for payload in health_by_format.values() {
let Some(score) = payload
.as_object()
.and_then(|payload| payload.get("health_score"))
.and_then(json_value_as_f64)
else {
continue;
};
scores.push(score.clamp(0.0, 1.0));
}
scores.into_iter().reduce(f64::min)
}
pub fn effective_provider_key_health_score(
key: &StoredProviderCatalogKey,
api_format: &str,
) -> Option<f64> {
provider_key_health_score(key, api_format).or_else(|| aggregate_provider_key_health_score(key))
}
pub fn provider_key_health_bucket(
key: &StoredProviderCatalogKey,
api_format: &str,
) -> Option<ProviderKeyHealthBucket> {
effective_provider_key_health_score(key, api_format).map(ProviderKeyHealthBucket::from_score)
}
pub fn is_provider_key_circuit_open(key: &StoredProviderCatalogKey, api_format: &str) -> bool {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|values| values.get(api_format))
.and_then(serde_json::Value::as_object)
.and_then(|payload| payload.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
fn available_provider_key_rpm_slots_for_new_user(
key: &StoredProviderCatalogKey,
current_usage: usize,
effective_limit: usize,
now_unix_secs: u64,
) -> usize {
let reservation_ratio =
provider_key_dynamic_reservation_ratio(key, current_usage, effective_limit, now_unix_secs);
usize::max(
1,
(effective_limit as f64 * (1.0 - reservation_ratio)).floor() as usize,
)
}
fn provider_key_dynamic_reservation_ratio(
key: &StoredProviderCatalogKey,
current_usage: usize,
effective_limit: usize,
now_unix_secs: u64,
) -> f64 {
let total_requests = provider_key_total_requests(key);
if total_requests < PROBE_PHASE_REQUESTS {
return PROBE_RESERVATION_RATIO;
}
let confidence = provider_key_reservation_confidence(key, now_unix_secs);
let load_ratio = provider_key_load_ratio(current_usage, effective_limit);
if load_ratio < LOW_LOAD_THRESHOLD {
return STABLE_MIN_RESERVATION_RATIO;
}
if load_ratio < HIGH_LOAD_THRESHOLD {
let load_factor =
(load_ratio - LOW_LOAD_THRESHOLD) / (HIGH_LOAD_THRESHOLD - LOW_LOAD_THRESHOLD);
return STABLE_MIN_RESERVATION_RATIO
+ confidence
* load_factor
* (STABLE_MAX_RESERVATION_RATIO - STABLE_MIN_RESERVATION_RATIO);
}
STABLE_MIN_RESERVATION_RATIO
+ confidence * (STABLE_MAX_RESERVATION_RATIO - STABLE_MIN_RESERVATION_RATIO)
}
fn is_recently_active(candidate: &StoredRequestCandidate, now_unix_secs: u64) -> bool {
if candidate.finished_at_unix_secs.is_some() {
return false;
}
if !matches!(
candidate.status,
RequestCandidateStatus::Pending | RequestCandidateStatus::Streaming
) {
return false;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
now_unix_secs.saturating_sub(observed_at_unix_secs) <= ACTIVE_REQUEST_WINDOW_SECS
}
fn is_recent_rpm_observation(candidate: &StoredRequestCandidate, now_unix_secs: u64) -> bool {
if !candidate
.status
.is_attempted(candidate.started_at_unix_secs)
{
return false;
}
let observed_at_unix_secs = candidate
.started_at_unix_secs
.unwrap_or(candidate.created_at_unix_secs);
now_unix_secs.saturating_sub(observed_at_unix_secs) <= PROVIDER_KEY_RPM_WINDOW_SECS
}
fn provider_key_total_requests(key: &StoredProviderCatalogKey) -> u32 {
let request_count = key.request_count.unwrap_or_default();
if request_count > 0 {
return request_count;
}
let history_count = key
.adjustment_history
.as_ref()
.and_then(serde_json::Value::as_array)
.map(|values| values.len() as u32 * 10)
.unwrap_or_default();
key.concurrent_429_count.unwrap_or_default()
+ key.rpm_429_count.unwrap_or_default()
+ key.success_count.unwrap_or_default()
+ history_count
}
fn provider_key_load_ratio(current_usage: usize, effective_limit: usize) -> f64 {
if effective_limit == 0 {
return 0.0;
}
(current_usage as f64 / effective_limit as f64).min(1.0)
}
fn provider_key_reservation_confidence(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> f64 {
let request_count = key.request_count.unwrap_or_default() as f64;
let success_count = key.success_count.unwrap_or_default() as f64;
let success_score = if request_count >= SUCCESS_COUNT_FOR_FULL_CONFIDENCE as f64 {
let success_rate = if request_count > 0.0 {
success_count / request_count
} else {
0.0
};
success_rate * 0.4
} else if request_count > 0.0 {
let success_rate = success_count / request_count;
let progress_ratio = request_count / SUCCESS_COUNT_FOR_FULL_CONFIDENCE as f64;
success_rate * progress_ratio * 0.4
} else {
0.0
};
let cooldown_score = match key.last_429_at_unix_secs {
Some(last_429_at_unix_secs) => {
let hours_since_429 =
now_unix_secs.saturating_sub(last_429_at_unix_secs) as f64 / 3600.0;
(hours_since_429 / COOLDOWN_HOURS_FOR_FULL_CONFIDENCE).min(1.0) * 0.3
}
None => 0.3,
};
let stability_score = provider_key_stability_score(key);
(success_score + cooldown_score + stability_score).min(1.0)
}
fn provider_key_stability_score(key: &StoredProviderCatalogKey) -> f64 {
let Some(history) = key
.adjustment_history
.as_ref()
.and_then(serde_json::Value::as_array)
else {
return 0.15;
};
if history.len() < 3 {
return 0.15;
}
let recent = if history.len() > 5 {
&history[history.len() - 5..]
} else {
history.as_slice()
};
let limits = recent
.iter()
.filter_map(|entry| entry.get("new_limit"))
.filter_map(json_value_as_f64)
.collect::<Vec<_>>();
if limits.len() < 2 {
return 0.15;
}
let mean = limits.iter().sum::<f64>() / limits.len() as f64;
let variance = limits
.iter()
.map(|limit| {
let delta = *limit - mean;
delta * delta
})
.sum::<f64>()
/ (limits.len() as f64 - 1.0);
let stability_ratio = (1.0 - variance / 10.0).max(0.0);
stability_ratio * 0.3
}
fn json_value_as_f64(value: &serde_json::Value) -> Option<f64> {
value
.as_f64()
.or_else(|| value.as_i64().map(|raw| raw as f64))
.or_else(|| value.as_u64().map(|raw| raw as f64))
}
#[cfg(test)]
mod tests {
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::repository::provider_catalog::StoredProviderCatalogKey;
use super::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
count_recent_active_requests_for_provider, count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket,
};
fn stored_candidate(
id: &str,
status: RequestCandidateStatus,
created_at_unix_secs: i64,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
format!("req-{id}"),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
status,
None,
false,
None,
None,
None,
None,
None,
None,
None,
created_at_unix_secs,
Some(created_at_unix_secs),
Some(created_at_unix_secs),
)
.expect("candidate should build")
}
fn provider_catalog_key(id: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
id.to_string(),
"provider-a".to_string(),
"primary".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("provider key should build")
}
#[test]
fn cooldown_triggers_after_two_recent_failures() {
let recent_candidates = vec![
stored_candidate("one", RequestCandidateStatus::Failed, 95),
stored_candidate("two", RequestCandidateStatus::Cancelled, 99),
];
assert!(is_candidate_in_recent_failure_cooldown(
&recent_candidates,
"provider-a",
"endpoint-a",
"key-a",
100,
));
}
#[test]
fn recent_success_clears_cooldown() {
let recent_candidates = vec![
stored_candidate("one", RequestCandidateStatus::Failed, 95),
stored_candidate("two", RequestCandidateStatus::Success, 99),
stored_candidate("three", RequestCandidateStatus::Cancelled, 98),
];
assert!(!is_candidate_in_recent_failure_cooldown(
&recent_candidates,
"provider-a",
"endpoint-a",
"key-a",
100,
));
}
#[test]
fn counts_only_recently_active_provider_requests() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
95,
Some(95),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"two".to_string(),
"req-two".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Streaming,
None,
false,
None,
None,
None,
None,
None,
None,
None,
96,
Some(96),
None,
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"three".to_string(),
"req-three".to_string(),
None,
Some("api-key-1".to_string()),
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
None,
None,
None,
None,
None,
None,
None,
97,
Some(97),
Some(98),
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_active_requests_for_provider(&recent_candidates, "provider-a", 100),
2
);
assert_eq!(
count_recent_active_requests_for_api_key(&recent_candidates, "api-key-1", 100),
2
);
}
#[test]
fn fixed_provider_key_rpm_limit_takes_precedence() {
let key = provider_catalog_key("key-a").with_rate_limit_fields(
Some(120),
Some(80),
None,
None,
None,
None,
Some(10),
Some(10),
);
assert_eq!(effective_provider_key_rpm_limit(&key, 100), Some(120));
}
#[test]
fn learned_provider_key_rpm_limit_requires_confidence() {
let low_confidence = provider_catalog_key("key-a").with_rate_limit_fields(
None,
Some(80),
Some(0),
Some(0),
Some(99),
None,
Some(5),
Some(1),
);
assert_eq!(effective_provider_key_rpm_limit(&low_confidence, 100), None);
let high_confidence = provider_catalog_key("key-a").with_rate_limit_fields(
None,
Some(80),
Some(0),
Some(0),
None,
Some(serde_json::json!([
{"new_limit": 80},
{"new_limit": 81},
{"new_limit": 80},
])),
Some(120),
Some(118),
);
assert_eq!(
effective_provider_key_rpm_limit(&high_confidence, 100),
Some(80)
);
}
#[test]
fn counts_recent_provider_key_rpm_from_snapshot_or_recent_attempts() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(7),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"two".to_string(),
"req-two".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(502),
None,
None,
Some(10),
None,
None,
None,
98,
Some(98),
Some(99),
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_rpm_requests_for_provider_key(&recent_candidates, "key-a", 100),
7
);
}
#[test]
fn ignores_rpm_observations_before_reset_watermark() {
let recent_candidates = vec![
StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(7),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build"),
StoredRequestCandidate::new(
"two".to_string(),
"req-two".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(2),
None,
None,
99,
Some(99),
Some(100),
)
.expect("candidate should build"),
];
assert_eq!(
count_recent_rpm_requests_for_provider_key_since(
&recent_candidates,
"key-a",
100,
Some(98),
),
2
);
}
#[test]
fn provider_key_rpm_reserves_capacity_for_new_users() {
let key = provider_catalog_key("key-a").with_rate_limit_fields(
Some(10),
None,
None,
None,
None,
None,
Some(5),
Some(5),
);
let recent_candidates = vec![StoredRequestCandidate::new(
"one".to_string(),
"req-one".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-a".to_string()),
Some("endpoint-a".to_string()),
Some("key-a".to_string()),
RequestCandidateStatus::Success,
None,
false,
Some(200),
None,
None,
Some(10),
Some(9),
None,
None,
95,
Some(95),
Some(96),
)
.expect("candidate should build")];
assert!(!provider_key_rpm_allows_request(
&key,
&recent_candidates,
100,
false,
));
assert!(provider_key_rpm_allows_request(
&key,
&recent_candidates,
100,
true,
));
assert!(provider_key_rpm_allows_request_since(
&key,
&recent_candidates,
100,
false,
Some(97),
));
}
#[test]
fn reads_provider_key_health_and_circuit_status_for_api_format() {
let key = provider_catalog_key("key-a").with_health_fields(
Some(serde_json::json!({
"openai:chat": {"health_score": 0.25},
"openai:responses": {"health_score": 0.75}
})),
Some(serde_json::json!({
"openai:chat": {"open": true},
"openai:responses": {"open": false}
})),
);
assert_eq!(provider_key_health_score(&key, "openai:chat"), Some(0.25));
assert_eq!(
provider_key_health_score(&key, "openai:responses"),
Some(0.75)
);
assert!(is_provider_key_circuit_open(&key, "openai:chat"));
assert!(!is_provider_key_circuit_open(&key, "openai:responses"));
}
#[test]
fn aggregates_provider_key_health_score_with_lower_bound_strategy() {
let key = provider_catalog_key("key-a").with_health_fields(
Some(serde_json::json!({
"openai:chat": {"health_score": 0.85},
"openai:responses": {"health_score": 0.45},
"claude:chat": {"health_score": 0.70}
})),
None,
);
assert_eq!(aggregate_provider_key_health_score(&key), Some(0.45));
assert_eq!(
effective_provider_key_health_score(&key, "gemini:chat"),
Some(0.45)
);
}
#[test]
fn classifies_provider_key_health_bucket_from_effective_score() {
let low = provider_catalog_key("key-low").with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.30}})),
None,
);
let degraded = provider_catalog_key("key-degraded").with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.65}})),
None,
);
let healthy = provider_catalog_key("key-healthy").with_health_fields(
Some(serde_json::json!({"openai:chat": {"health_score": 0.92}})),
None,
);
assert_eq!(
provider_key_health_bucket(&low, "openai:chat"),
Some(ProviderKeyHealthBucket::Low)
);
assert_eq!(
provider_key_health_bucket(&degraded, "openai:chat"),
Some(ProviderKeyHealthBucket::Degraded)
);
assert_eq!(
provider_key_health_bucket(&healthy, "openai:chat"),
Some(ProviderKeyHealthBucket::Healthy)
);
}
}

View File

@@ -0,0 +1,45 @@
mod affinity;
mod auth;
mod candidate;
mod health;
mod model;
mod provider;
mod request_candidate;
pub use affinity::{
build_scheduler_affinity_cache_key_for_api_key_id, candidate_affinity_hash, candidate_key,
compare_affinity_order, matches_affinity_target, SchedulerAffinityTarget,
};
pub use auth::{
auth_constraints_allow_api_format, auth_constraints_allow_model,
auth_constraints_allow_provider, SchedulerAuthConstraints,
};
pub use candidate::{
auth_api_key_concurrency_limit_reached, build_minimal_candidate_selection,
candidate_is_selectable_with_runtime_state, candidate_supports_required_capability,
collect_global_model_names_for_required_capability, collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health, SchedulerMinimalCandidateSelectionCandidate,
};
pub use health::{
aggregate_provider_key_health_score, count_recent_active_requests_for_api_key,
count_recent_active_requests_for_provider, count_recent_rpm_requests_for_provider_key,
count_recent_rpm_requests_for_provider_key_since, effective_provider_key_health_score,
effective_provider_key_rpm_limit, is_candidate_in_recent_failure_cooldown,
is_provider_key_circuit_open, provider_key_health_bucket, provider_key_health_score,
provider_key_rpm_allows_request, provider_key_rpm_allows_request_since,
ProviderKeyHealthBucket, PROVIDER_KEY_RPM_WINDOW_SECS,
};
pub use model::{
candidate_model_names, extract_global_priority_for_format, matches_model_mapping,
normalize_api_format, resolve_provider_model_name, resolve_requested_global_model_name,
row_supports_required_capability, select_provider_model_name,
};
pub use provider::{build_provider_concurrent_limit_map, should_skip_provider_quota};
pub use request_candidate::{
build_execution_request_candidate_seed, build_local_request_candidate_status_record,
build_report_request_candidate_status_record, execution_error_details,
finalize_execution_request_candidate_report_context, is_terminal_candidate_status,
parse_request_candidate_report_context, resolve_report_request_candidate_slot,
SchedulerExecutionRequestCandidateSeed, SchedulerRequestCandidateReportContext,
SchedulerResolvedReportRequestCandidateSlot,
};

View File

@@ -0,0 +1,257 @@
use std::collections::BTreeSet;
use aether_data::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data::DataLayerError;
use regex::Regex;
pub fn resolve_requested_global_model_name(
rows: &[StoredMinimalCandidateSelectionRow],
requested_model_name: &str,
api_format: &str,
) -> Option<String> {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_name == requested_model_name
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.model_provider_model_mappings
.as_ref()
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping_scope_matches(mapping, api_format)
&& mapping.name == requested_model_name
})
})
})
})
.or_else(|| {
resolve_global_model_name_by(rows, |row| {
row.global_model_mappings.as_ref().is_some_and(|patterns| {
patterns
.iter()
.any(|pattern| matches_model_mapping(pattern, requested_model_name))
})
})
})
}
fn resolve_global_model_name_by<F>(
rows: &[StoredMinimalCandidateSelectionRow],
matches: F,
) -> Option<String>
where
F: Fn(&StoredMinimalCandidateSelectionRow) -> bool,
{
let mut matches = rows
.iter()
.filter(|row| matches(row))
.map(|row| row.global_model_name.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<BTreeSet<_>>()
.into_iter();
matches.next()
}
pub fn resolve_provider_model_name(
row: &StoredMinimalCandidateSelectionRow,
requested_model_name: &str,
api_format: &str,
) -> Option<(String, Option<String>)> {
let selected_provider_model_name = select_provider_model_name(row, api_format);
let Some(key_allowed_models) = row.key_allowed_models.as_ref() else {
return Some((selected_provider_model_name, None));
};
if key_allowed_models.is_empty() {
return None;
}
if key_allowed_models
.iter()
.any(|value| value == requested_model_name)
{
return Some((selected_provider_model_name, None));
}
let candidate_models = candidate_model_names(row, api_format);
let mut sorted_allowed_models = key_allowed_models
.iter()
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.collect::<Vec<_>>();
sorted_allowed_models.sort();
for allowed_model in &sorted_allowed_models {
if candidate_models.contains(allowed_model.as_str()) {
return Some((allowed_model.clone(), Some(allowed_model.clone())));
}
}
let Some(global_model_mappings) = row.global_model_mappings.as_ref() else {
return None;
};
for allowed_model in sorted_allowed_models {
for pattern in global_model_mappings {
if matches_model_mapping(pattern, &allowed_model) {
return Some((allowed_model.clone(), Some(allowed_model)));
}
}
}
None
}
pub fn select_provider_model_name(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> String {
let Some(mappings) = row.model_provider_model_mappings.as_ref() else {
return row.model_provider_model_name.clone();
};
let mut scoped = mappings
.iter()
.filter(|mapping| mapping_scope_matches(mapping, api_format))
.collect::<Vec<_>>();
if scoped.is_empty() {
return row.model_provider_model_name.clone();
}
scoped.sort_by(|left, right| {
left.priority
.cmp(&right.priority)
.then(left.name.cmp(&right.name))
});
let top_priority = scoped[0].priority;
scoped
.into_iter()
.find(|mapping| mapping.priority == top_priority)
.map(|mapping| mapping.name.clone())
.unwrap_or_else(|| row.model_provider_model_name.clone())
}
pub fn candidate_model_names(
row: &StoredMinimalCandidateSelectionRow,
api_format: &str,
) -> BTreeSet<String> {
let mut names = BTreeSet::from([row.model_provider_model_name.clone()]);
if let Some(mappings) = row.model_provider_model_mappings.as_ref() {
for mapping in mappings {
if mapping_scope_matches(mapping, api_format) {
names.insert(mapping.name.clone());
}
}
}
names
}
fn mapping_scope_matches(mapping: &StoredProviderModelMapping, api_format: &str) -> bool {
let Some(api_formats) = mapping.api_formats.as_ref() else {
return true;
};
api_formats
.iter()
.any(|value| normalize_api_format(value) == api_format)
}
pub fn row_supports_required_capability(
row: &StoredMinimalCandidateSelectionRow,
required_capability: &str,
) -> bool {
capabilities_support_required_capability(row.key_capabilities.as_ref(), required_capability)
}
fn capabilities_support_required_capability(
capabilities: Option<&serde_json::Value>,
required_capability: &str,
) -> bool {
let required_capability = required_capability.trim();
if required_capability.is_empty() {
return true;
}
let Some(capabilities) = capabilities else {
return false;
};
if let Some(object) = capabilities.as_object() {
return object.iter().any(|(key, value)| {
key.eq_ignore_ascii_case(required_capability)
&& match value {
serde_json::Value::Bool(value) => *value,
serde_json::Value::String(value) => value.eq_ignore_ascii_case("true"),
serde_json::Value::Number(value) => {
value.as_i64().is_some_and(|value| value > 0)
}
_ => false,
}
});
}
if let Some(items) = capabilities.as_array() {
return items.iter().any(|value| {
value
.as_str()
.is_some_and(|value| value.eq_ignore_ascii_case(required_capability))
});
}
false
}
pub fn matches_model_mapping(pattern: &str, model_name: &str) -> bool {
let Ok(compiled) = Regex::new(&format!("^(?:{pattern})$")) else {
return false;
};
compiled.is_match(model_name)
}
pub fn extract_global_priority_for_format(
raw: Option<&serde_json::Value>,
api_format: &str,
) -> Result<Option<i32>, DataLayerError> {
let Some(raw) = raw else {
return Ok(None);
};
let Some(object) = raw.as_object() else {
return Err(DataLayerError::UnexpectedValue(
"provider_api_keys.global_priority_by_format is not a JSON object".to_string(),
));
};
let Some(value) = object
.iter()
.find(|(key, _)| normalize_api_format(key) == api_format)
.map(|(_, value)| value)
else {
return Ok(None);
};
if let Some(value) = value.as_i64() {
return i32::try_from(value).map(Some).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.global_priority_by_format value: {value}"
))
});
}
if let Some(value) = value.as_str() {
let value = value.trim().parse::<i32>().map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"invalid provider_api_keys.global_priority_by_format value: {value}"
))
})?;
return Ok(Some(value));
}
Err(DataLayerError::UnexpectedValue(
"provider_api_keys.global_priority_by_format contains a non-integer value".to_string(),
))
}
pub fn normalize_api_format(value: &str) -> String {
value.trim().to_ascii_lowercase()
}

View File

@@ -0,0 +1,127 @@
use std::collections::BTreeMap;
use aether_data::repository::provider_catalog::StoredProviderCatalogProvider;
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
use aether_wallet::{ProviderBillingType, ProviderQuotaSnapshot};
pub fn should_skip_provider_quota(quota: &StoredProviderQuotaSnapshot, now_unix_secs: u64) -> bool {
let snapshot = ProviderQuotaSnapshot {
provider_id: quota.provider_id.clone(),
billing_type: ProviderBillingType::parse(&quota.billing_type),
monthly_quota_usd: quota.monthly_quota_usd,
monthly_used_usd: quota.monthly_used_usd,
quota_reset_day: quota.quota_reset_day,
quota_last_reset_at_unix_secs: quota.quota_last_reset_at_unix_secs,
quota_expires_at_unix_secs: quota.quota_expires_at_unix_secs,
is_active: quota.is_active,
};
if !snapshot.is_active || snapshot.is_expired(now_unix_secs) {
return true;
}
match snapshot.billing_type {
ProviderBillingType::MonthlyQuota | ProviderBillingType::FreeTier => snapshot
.remaining_quota_usd()
.is_some_and(|remaining| remaining <= 0.0),
ProviderBillingType::PayAsYouGo | ProviderBillingType::Unknown => false,
}
}
pub fn build_provider_concurrent_limit_map(
providers: Vec<StoredProviderCatalogProvider>,
) -> BTreeMap<String, usize> {
providers
.into_iter()
.filter_map(|provider| {
provider
.concurrent_limit
.and_then(|limit| usize::try_from(limit).ok())
.filter(|limit| *limit > 0)
.map(|limit| (provider.id, limit))
})
.collect()
}
#[cfg(test)]
mod tests {
use super::{build_provider_concurrent_limit_map, should_skip_provider_quota};
use aether_data::repository::provider_catalog::StoredProviderCatalogProvider;
use aether_data::repository::quota::StoredProviderQuotaSnapshot;
fn sample_provider(id: &str, concurrent_limit: Option<i32>) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
id.to_string(),
format!("provider-{id}"),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
concurrent_limit,
None,
None,
None,
None,
None,
)
}
#[test]
fn skips_inactive_or_exhausted_monthly_quota_provider() {
let inactive = StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"monthly_quota".to_string(),
Some(10.0),
1.0,
Some(30),
Some(1_000),
None,
false,
)
.expect("quota should build");
assert!(should_skip_provider_quota(&inactive, 2_000));
let exhausted = StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"monthly_quota".to_string(),
Some(10.0),
10.0,
Some(30),
Some(1_000),
None,
true,
)
.expect("quota should build");
assert!(should_skip_provider_quota(&exhausted, 2_000));
let payg = StoredProviderQuotaSnapshot::new(
"provider-1".to_string(),
"pay_as_you_go".to_string(),
None,
10.0,
None,
None,
None,
true,
)
.expect("quota should build");
assert!(!should_skip_provider_quota(&payg, 2_000));
}
#[test]
fn builds_provider_concurrent_limit_map_for_positive_limits_only() {
let limits = build_provider_concurrent_limit_map(vec![
sample_provider("provider-a", Some(10)),
sample_provider("provider-b", Some(0)),
sample_provider("provider-c", None),
]);
assert_eq!(limits.get("provider-a"), Some(&10));
assert!(!limits.contains_key("provider-b"));
assert!(!limits.contains_key("provider-c"));
}
}

View File

@@ -0,0 +1,664 @@
use aether_contracts::{ExecutionError, ExecutionPlan};
use aether_data::repository::candidates::{
RequestCandidateStatus, StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use serde_json::Value;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SchedulerRequestCandidateReportContext {
pub request_id: Option<String>,
pub candidate_id: Option<String>,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub candidate_index: Option<u32>,
pub retry_index: u32,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub client_api_format: Option<String>,
pub provider_api_format: Option<String>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct SchedulerResolvedReportRequestCandidateSlot {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
pub candidate_index: u32,
pub retry_index: u32,
pub provider_id: Option<String>,
pub endpoint_id: Option<String>,
pub key_id: Option<String>,
pub extra_data: Option<Value>,
pub created_at_unix_secs: u64,
pub started_at_unix_secs: Option<u64>,
pub finished_at_unix_secs: Option<u64>,
}
pub struct SchedulerExecutionRequestCandidateSeed {
pub upsert_record: UpsertRequestCandidateRecord,
pub report_context: Value,
}
pub fn execution_error_details(
error: Option<&ExecutionError>,
body_json: Option<&Value>,
) -> (Option<String>, Option<String>) {
match error {
Some(error) => (
Some(format!("{:?}", error.kind)),
Some(error.message.trim().to_string()).filter(|value| !value.is_empty()),
),
None => (
None,
body_json
.and_then(extract_error_message)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
),
}
}
pub fn parse_request_candidate_report_context(
report_context: Option<&Value>,
) -> Option<SchedulerRequestCandidateReportContext> {
let report_context = report_context?;
let retry_index = report_context
.get("retry_index")
.and_then(Value::as_u64)
.unwrap_or_default();
Some(SchedulerRequestCandidateReportContext {
request_id: string_field(report_context, "request_id"),
candidate_id: string_field(report_context, "candidate_id"),
user_id: string_field(report_context, "user_id"),
api_key_id: string_field(report_context, "api_key_id"),
candidate_index: report_context
.get("candidate_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
retry_index: u32::try_from(retry_index).unwrap_or(u32::MAX),
provider_id: string_field(report_context, "provider_id"),
endpoint_id: string_field(report_context, "endpoint_id"),
key_id: string_field(report_context, "key_id"),
client_api_format: string_field(report_context, "client_api_format"),
provider_api_format: string_field(report_context, "provider_api_format"),
})
}
pub fn resolve_report_request_candidate_slot(
existing_candidates: &[StoredRequestCandidate],
metadata: SchedulerRequestCandidateReportContext,
now_unix_secs: u64,
generated_candidate_id: String,
) -> Option<SchedulerResolvedReportRequestCandidateSlot> {
let request_id = metadata.request_id.clone()?;
let matched_candidate = match_existing_report_candidate(existing_candidates, &metadata);
let synthesized_extra_data = build_report_candidate_extra_data(&metadata);
let created_at_unix_secs = matched_candidate
.as_ref()
.map(|candidate| candidate.created_at_unix_secs)
.unwrap_or(now_unix_secs);
let candidate_index = matched_candidate
.as_ref()
.map(|candidate| candidate.candidate_index)
.or(metadata.candidate_index)
.unwrap_or_else(|| next_candidate_index(existing_candidates));
let retry_index = matched_candidate
.as_ref()
.map(|candidate| candidate.retry_index)
.unwrap_or(metadata.retry_index);
Some(SchedulerResolvedReportRequestCandidateSlot {
id: matched_candidate
.as_ref()
.map(|candidate| candidate.id.clone())
.or(metadata.candidate_id)
.unwrap_or(generated_candidate_id),
request_id,
user_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.user_id.clone())
.or(metadata.user_id),
api_key_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.api_key_id.clone())
.or(metadata.api_key_id),
candidate_index,
retry_index,
provider_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.provider_id.clone())
.or(metadata.provider_id),
endpoint_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.endpoint_id.clone())
.or(metadata.endpoint_id),
key_id: matched_candidate
.as_ref()
.and_then(|candidate| candidate.key_id.clone())
.or(metadata.key_id),
extra_data: matched_candidate
.as_ref()
.and_then(|candidate| candidate.extra_data.clone())
.or(synthesized_extra_data),
created_at_unix_secs,
started_at_unix_secs: matched_candidate
.as_ref()
.and_then(|candidate| candidate.started_at_unix_secs),
finished_at_unix_secs: matched_candidate
.as_ref()
.and_then(|candidate| candidate.finished_at_unix_secs),
})
}
pub fn build_execution_request_candidate_seed(
plan: &ExecutionPlan,
report_context: Option<&Value>,
started_at_unix_secs: u64,
generated_candidate_id: String,
) -> SchedulerExecutionRequestCandidateSeed {
let mut context = report_context
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let request_id = string_field(&Value::Object(context.clone()), "request_id")
.unwrap_or_else(|| plan.request_id.clone());
let candidate_index = context
.get("candidate_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(0);
let retry_index = context
.get("retry_index")
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.unwrap_or(0);
let candidate_id = string_field(&Value::Object(context.clone()), "candidate_id")
.unwrap_or(generated_candidate_id);
let user_id = string_field(&Value::Object(context.clone()), "user_id");
let api_key_id = string_field(&Value::Object(context.clone()), "api_key_id");
context.insert("request_id".to_string(), Value::String(request_id.clone()));
context.insert(
"candidate_id".to_string(),
Value::String(candidate_id.clone()),
);
context.insert(
"candidate_index".to_string(),
Value::Number(candidate_index.into()),
);
context.insert(
"provider_id".to_string(),
Value::String(plan.provider_id.clone()),
);
context.insert(
"endpoint_id".to_string(),
Value::String(plan.endpoint_id.clone()),
);
context.insert("key_id".to_string(), Value::String(plan.key_id.clone()));
SchedulerExecutionRequestCandidateSeed {
upsert_record: UpsertRequestCandidateRecord {
id: candidate_id,
request_id,
user_id,
api_key_id,
username: None,
api_key_name: None,
candidate_index,
retry_index,
provider_id: Some(plan.provider_id.clone()),
endpoint_id: Some(plan.endpoint_id.clone()),
key_id: Some(plan.key_id.clone()),
status: RequestCandidateStatus::Pending,
skip_reason: None,
is_cached: Some(false),
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: Some(started_at_unix_secs),
started_at_unix_secs: Some(started_at_unix_secs),
finished_at_unix_secs: None,
},
report_context: Value::Object(context),
}
}
pub fn build_local_request_candidate_status_record(
plan: &ExecutionPlan,
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
) -> Option<UpsertRequestCandidateRecord> {
let candidate_id = plan
.candidate_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?;
let metadata = parse_request_candidate_report_context(report_context)?;
let candidate_index = metadata.candidate_index?;
Some(UpsertRequestCandidateRecord {
id: candidate_id.to_string(),
request_id: plan.request_id.clone(),
user_id: metadata.user_id,
api_key_id: metadata.api_key_id,
username: None,
api_key_name: None,
candidate_index,
retry_index: metadata.retry_index,
provider_id: Some(plan.provider_id.clone()),
endpoint_id: Some(plan.endpoint_id.clone()),
key_id: Some(plan.key_id.clone()),
status,
skip_reason: None,
is_cached: None,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: None,
started_at_unix_secs,
finished_at_unix_secs,
})
}
pub fn build_report_request_candidate_status_record(
slot: SchedulerResolvedReportRequestCandidateSlot,
status: RequestCandidateStatus,
status_code: Option<u16>,
error_type: Option<String>,
error_message: Option<String>,
latency_ms: Option<u64>,
started_at_unix_secs: Option<u64>,
finished_at_unix_secs: Option<u64>,
now_unix_secs: u64,
) -> UpsertRequestCandidateRecord {
let terminal_unix_secs = finished_at_unix_secs.unwrap_or(now_unix_secs);
let started_at_unix_secs = started_at_unix_secs
.or(slot.started_at_unix_secs)
.or_else(|| status.is_attempted(None).then_some(terminal_unix_secs));
let finished_at_unix_secs = finished_at_unix_secs
.or(slot.finished_at_unix_secs)
.or_else(|| is_terminal_candidate_status(status).then_some(terminal_unix_secs));
UpsertRequestCandidateRecord {
id: slot.id,
request_id: slot.request_id,
user_id: slot.user_id,
api_key_id: slot.api_key_id,
username: None,
api_key_name: None,
candidate_index: slot.candidate_index,
retry_index: slot.retry_index,
provider_id: slot.provider_id,
endpoint_id: slot.endpoint_id,
key_id: slot.key_id,
status,
skip_reason: None,
is_cached: None,
status_code,
error_type,
error_message,
latency_ms,
concurrent_requests: None,
extra_data: slot.extra_data,
required_capabilities: None,
created_at_unix_secs: Some(slot.created_at_unix_secs),
started_at_unix_secs,
finished_at_unix_secs,
}
}
pub fn finalize_execution_request_candidate_report_context(
report_context: Value,
candidate_id: &str,
) -> Value {
let mut context = report_context.as_object().cloned().unwrap_or_default();
let candidate_id = candidate_id.trim();
if !candidate_id.is_empty() {
context.insert(
"candidate_id".to_string(),
Value::String(candidate_id.to_string()),
);
}
Value::Object(context)
}
pub fn is_terminal_candidate_status(status: RequestCandidateStatus) -> bool {
matches!(
status,
RequestCandidateStatus::Unused
| RequestCandidateStatus::Success
| RequestCandidateStatus::Failed
| RequestCandidateStatus::Cancelled
| RequestCandidateStatus::Skipped
)
}
fn extract_error_message(body_json: &Value) -> Option<&str> {
body_json
.get("error")
.and_then(|error| {
error
.get("message")
.and_then(Value::as_str)
.or_else(|| error.as_str())
})
.or_else(|| body_json.get("message").and_then(Value::as_str))
}
fn string_field(value: &Value, key: &str) -> Option<String> {
value
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn match_existing_report_candidate<'a>(
candidates: &'a [StoredRequestCandidate],
metadata: &SchedulerRequestCandidateReportContext,
) -> Option<&'a StoredRequestCandidate> {
if let Some(candidate_id) = metadata.candidate_id.as_deref() {
if let Some(candidate) = candidates
.iter()
.find(|candidate| candidate.id == candidate_id)
{
return Some(candidate);
}
}
if let Some(candidate_index) = metadata.candidate_index {
if let Some(candidate) = candidates.iter().find(|candidate| {
candidate.candidate_index == candidate_index
&& candidate.retry_index == metadata.retry_index
}) {
return Some(candidate);
}
}
candidates
.iter()
.filter(|candidate| {
candidate.provider_id.as_deref() == metadata.provider_id.as_deref()
&& candidate.endpoint_id.as_deref() == metadata.endpoint_id.as_deref()
&& candidate.key_id.as_deref() == metadata.key_id.as_deref()
})
.max_by_key(|candidate| {
(
candidate.retry_index,
candidate.candidate_index,
candidate.created_at_unix_secs,
)
})
}
fn next_candidate_index(candidates: &[StoredRequestCandidate]) -> u32 {
candidates
.iter()
.map(|candidate| candidate.candidate_index)
.max()
.map(|value| value.saturating_add(1))
.unwrap_or_default()
}
fn build_report_candidate_extra_data(
metadata: &SchedulerRequestCandidateReportContext,
) -> Option<Value> {
let mut extra_data = serde_json::Map::new();
extra_data.insert("gateway_execution_runtime".to_string(), Value::Bool(true));
extra_data.insert("phase".to_string(), Value::String("3c_trial".to_string()));
if let Some(client_api_format) = metadata.client_api_format.clone() {
extra_data.insert(
"client_api_format".to_string(),
Value::String(client_api_format),
);
}
if let Some(provider_api_format) = metadata.provider_api_format.clone() {
extra_data.insert(
"provider_api_format".to_string(),
Value::String(provider_api_format),
);
}
(!extra_data.is_empty()).then_some(Value::Object(extra_data))
}
#[cfg(test)]
mod tests {
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, RequestBody,
};
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use serde_json::{json, Value};
use super::{
build_execution_request_candidate_seed, build_local_request_candidate_status_record,
build_report_request_candidate_status_record, execution_error_details,
finalize_execution_request_candidate_report_context,
parse_request_candidate_report_context, resolve_report_request_candidate_slot,
SchedulerResolvedReportRequestCandidateSlot,
};
fn sample_candidate(
id: &str,
candidate_index: u32,
retry_index: u32,
) -> StoredRequestCandidate {
StoredRequestCandidate::new(
id.to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("key-1".to_string()),
None,
None,
candidate_index as i32,
retry_index as i32,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("catalog-key-1".to_string()),
RequestCandidateStatus::Pending,
None,
false,
None,
None,
None,
None,
None,
None,
None,
100,
Some(110),
None,
)
.expect("candidate should build")
}
fn sample_plan() -> ExecutionPlan {
ExecutionPlan {
request_id: "req-1".to_string(),
candidate_id: None,
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://example.com/v1/chat/completions".to_string(),
headers: Default::default(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({"model": "gpt-5"})),
stream: false,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".to_string()),
proxy: None,
tls_profile: None,
timeouts: None,
}
}
#[test]
fn parses_report_context_and_resolves_existing_candidate_slot() {
let metadata = parse_request_candidate_report_context(Some(&json!({
"request_id": "req-1",
"candidate_index": 1,
"retry_index": 2,
"provider_id": "provider-1",
"endpoint_id": "endpoint-1",
"key_id": "catalog-key-1",
"client_api_format": "openai:chat"
})))
.expect("metadata");
let slot = resolve_report_request_candidate_slot(
&[sample_candidate("cand-1", 1, 2)],
metadata,
123,
"generated-1".to_string(),
)
.expect("slot");
assert_eq!(slot.id, "cand-1");
assert_eq!(slot.candidate_index, 1);
assert_eq!(slot.retry_index, 2);
assert_eq!(slot.request_id, "req-1");
}
#[test]
fn resolves_error_details_from_execution_error_or_body_json() {
let error = ExecutionError {
kind: ExecutionErrorKind::Upstream5xx,
phase: ExecutionPhase::FirstByte,
message: " upstream failed ".to_string(),
upstream_status: Some(502),
retryable: true,
failover_recommended: true,
};
assert_eq!(
execution_error_details(Some(&error), None),
(
Some("Upstream5xx".to_string()),
Some("upstream failed".to_string())
)
);
assert_eq!(
execution_error_details(None, Some(&json!({"error": {"message": "bad request"}}))),
(None, Some("bad request".to_string()))
);
}
#[test]
fn builds_execution_request_candidate_seed_and_finalizes_report_context() {
let seed = build_execution_request_candidate_seed(
&sample_plan(),
Some(&json!({
"request_id": "req-override",
"candidate_index": 3,
"retry_index": 2,
"user_id": "user-1",
"api_key_id": "api-key-1",
"client_api_format": "openai:chat"
})),
123,
"generated-1".to_string(),
);
assert_eq!(seed.upsert_record.id, "generated-1");
assert_eq!(seed.upsert_record.request_id, "req-override");
assert_eq!(seed.upsert_record.candidate_index, 3);
assert_eq!(seed.upsert_record.retry_index, 2);
assert_eq!(seed.upsert_record.user_id.as_deref(), Some("user-1"));
assert_eq!(
seed.report_context
.get("provider_id")
.and_then(Value::as_str),
Some("provider-1")
);
let finalized =
finalize_execution_request_candidate_report_context(seed.report_context, "cand-final");
assert_eq!(
finalized.get("candidate_id").and_then(Value::as_str),
Some("cand-final")
);
}
#[test]
fn builds_local_request_candidate_status_record() {
let mut plan = sample_plan();
plan.candidate_id = Some("cand-1".to_string());
let record = build_local_request_candidate_status_record(
&plan,
Some(&json!({
"candidate_index": 1,
"retry_index": 2,
"user_id": "user-1",
"api_key_id": "api-key-1"
})),
RequestCandidateStatus::Failed,
Some(500),
Some("Upstream5xx".to_string()),
Some("boom".to_string()),
Some(42),
Some(100),
Some(101),
)
.expect("record should build");
assert_eq!(record.id, "cand-1");
assert_eq!(record.candidate_index, 1);
assert_eq!(record.retry_index, 2);
assert_eq!(record.user_id.as_deref(), Some("user-1"));
assert_eq!(record.status, RequestCandidateStatus::Failed);
}
#[test]
fn builds_report_request_candidate_status_record_with_terminal_timestamps() {
let record = build_report_request_candidate_status_record(
SchedulerResolvedReportRequestCandidateSlot {
id: "cand-1".to_string(),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
candidate_index: 1,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
extra_data: None,
created_at_unix_secs: 10,
started_at_unix_secs: None,
finished_at_unix_secs: None,
},
RequestCandidateStatus::Success,
Some(200),
None,
None,
Some(12),
None,
None,
123,
);
assert_eq!(record.started_at_unix_secs, Some(123));
assert_eq!(record.finished_at_unix_secs, Some(123));
assert_eq!(record.created_at_unix_secs, Some(10));
assert_eq!(record.status, RequestCandidateStatus::Success);
}
}

View File

@@ -0,0 +1,17 @@
[package]
name = "aether-usage-runtime"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
description = "Usage runtime shared types and queue primitives extracted from aether-gateway"
[dependencies]
aether-contracts.workspace = true
aether-data.workspace = true
async-trait.workspace = true
base64.workspace = true
serde.workspace = true
serde_json.workspace = true
tokio.workspace = true
tracing.workspace = true

View File

@@ -0,0 +1,112 @@
use aether_data::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageRuntimeConfig {
pub enabled: bool,
pub stream_key: String,
pub consumer_group: String,
pub dlq_stream_key: String,
pub stream_maxlen: usize,
pub consumer_batch_size: usize,
pub consumer_block_ms: u64,
pub reclaim_idle_ms: u64,
pub reclaim_count: usize,
pub reclaim_interval_ms: u64,
}
impl Default for UsageRuntimeConfig {
fn default() -> Self {
Self {
enabled: false,
stream_key: "usage:events".to_string(),
consumer_group: "usage_consumers".to_string(),
dlq_stream_key: "usage:events:dlq".to_string(),
stream_maxlen: 2_000,
consumer_batch_size: 200,
consumer_block_ms: 500,
reclaim_idle_ms: 30_000,
reclaim_count: 200,
reclaim_interval_ms: 5_000,
}
}
}
impl UsageRuntimeConfig {
pub fn disabled() -> Self {
Self::default()
}
pub fn validate(&self) -> Result<(), DataLayerError> {
if !self.enabled {
return Ok(());
}
if self.stream_key.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime stream_key cannot be empty".to_string(),
));
}
if self.consumer_group.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_group cannot be empty".to_string(),
));
}
if self.dlq_stream_key.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime dlq_stream_key cannot be empty".to_string(),
));
}
if self.stream_maxlen == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime stream_maxlen must be positive".to_string(),
));
}
if self.consumer_batch_size == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_batch_size must be positive".to_string(),
));
}
if self.consumer_block_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_block_ms must be positive".to_string(),
));
}
if self.reclaim_idle_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime reclaim_idle_ms must be positive".to_string(),
));
}
if self.reclaim_count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime reclaim_count must be positive".to_string(),
));
}
if self.reclaim_interval_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime reclaim_interval_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::UsageRuntimeConfig;
#[test]
fn disabled_config_is_valid() {
assert!(UsageRuntimeConfig::disabled().validate().is_ok());
}
#[test]
fn enabled_config_rejects_empty_stream_key() {
let config = UsageRuntimeConfig {
enabled: true,
stream_key: String::new(),
..UsageRuntimeConfig::default()
};
assert!(config.validate().is_err());
}
}

View File

@@ -0,0 +1,216 @@
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data::DataLayerError;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const USAGE_EVENT_VERSION: u8 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageEventType {
Pending,
Streaming,
Completed,
Failed,
Cancelled,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct UsageEventData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub user_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub api_key_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub username: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub api_key_name: Option<String>,
pub provider_name: String,
pub model: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub target_model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_endpoint_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_api_key_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub api_format: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub api_family: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub endpoint_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub endpoint_api_format: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_api_family: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_endpoint_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub has_format_conversion: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub is_stream: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub total_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_creation_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub cache_read_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub output_price_per_1m: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub total_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub actual_total_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub status_code: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error_message: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub error_category: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_time_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub first_byte_time_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_metadata: Option<Value>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct UsageEvent {
pub event_type: UsageEventType,
pub request_id: String,
pub timestamp_ms: u64,
pub data: UsageEventData,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
struct UsageEventEnvelope {
v: u8,
#[serde(rename = "type")]
event_type: UsageEventType,
request_id: String,
timestamp_ms: u64,
data: UsageEventData,
}
impl UsageEvent {
pub fn new(
event_type: UsageEventType,
request_id: impl Into<String>,
data: UsageEventData,
) -> Self {
Self {
event_type,
request_id: request_id.into(),
timestamp_ms: now_ms(),
data,
}
}
pub fn to_stream_fields(&self) -> Result<BTreeMap<String, String>, DataLayerError> {
let payload = UsageEventEnvelope {
v: USAGE_EVENT_VERSION,
event_type: self.event_type,
request_id: self.request_id.clone(),
timestamp_ms: self.timestamp_ms,
data: self.data.clone(),
};
let payload = serde_json::to_string(&payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to serialize usage event payload: {err}"
))
})?;
Ok(BTreeMap::from([("payload".to_string(), payload)]))
}
pub fn from_stream_fields(fields: &BTreeMap<String, String>) -> Result<Self, DataLayerError> {
let payload = fields.get("payload").ok_or_else(|| {
DataLayerError::UnexpectedValue(
"usage event stream entry missing payload field".to_string(),
)
})?;
let envelope: UsageEventEnvelope = serde_json::from_str(payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to deserialize usage event payload: {err}"
))
})?;
if envelope.v != USAGE_EVENT_VERSION {
return Err(DataLayerError::UnexpectedValue(format!(
"unsupported usage event version: {}",
envelope.v
)));
}
Ok(Self {
event_type: envelope.event_type,
request_id: envelope.request_id,
timestamp_ms: envelope.timestamp_ms,
data: envelope.data,
})
}
}
pub fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
#[cfg(test)]
mod tests {
use super::{UsageEvent, UsageEventData, UsageEventType};
#[test]
fn usage_event_round_trips_through_stream_fields() {
let event = UsageEvent::new(
UsageEventType::Completed,
"req-1",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
input_tokens: Some(10),
output_tokens: Some(20),
..UsageEventData::default()
},
);
let fields = event.to_stream_fields().expect("event should serialize");
let parsed = UsageEvent::from_stream_fields(&fields).expect("event should parse");
assert_eq!(parsed.request_id, "req-1");
assert_eq!(parsed.event_type, UsageEventType::Completed);
assert_eq!(parsed.data.total_tokens, None);
assert_eq!(parsed.data.output_tokens, Some(20));
}
}

View File

@@ -0,0 +1,44 @@
pub mod config;
pub mod event;
pub mod queue;
pub mod record;
pub mod report;
pub mod report_context;
pub mod runtime;
pub mod settlement;
pub mod standardized_usage;
pub mod usage_mapper;
pub mod worker;
pub mod write;
pub use config::UsageRuntimeConfig;
pub use event::{now_ms, UsageEvent, UsageEventData, UsageEventType, USAGE_EVENT_VERSION};
pub use queue::UsageQueue;
pub use record::build_upsert_usage_record_from_event;
pub use report::{
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
infer_internal_finalize_signature, is_local_ai_stream_report_kind,
is_local_ai_sync_report_kind, normalize_gemini_file_name, report_request_id,
resolve_internal_finalize_route, should_handle_local_stream_report,
should_handle_local_sync_report, sync_report_represents_failure, GatewayStreamReportRequest,
GatewaySyncReportRequest, GeminiFileMappingEntry, InternalFinalizeRoute,
GEMINI_FILE_MAPPING_TTL_SECONDS,
};
pub use report_context::{
build_locally_actionable_report_context_from_request_candidate,
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
};
pub use runtime::{UsageBillingEventEnricher, UsageRuntime, UsageRuntimeAccess};
pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
pub use standardized_usage::StandardizedUsage;
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
pub use worker::{
build_usage_queue_worker, write_event_record, UsageDataEventRecorder, UsageEventRecorder,
UsageQueueWorker, UsageRecordWriter,
};
pub use write::{
build_pending_usage_record, build_stream_terminal_usage_event,
build_stream_terminal_usage_outcome, build_streaming_usage_record,
build_sync_terminal_usage_event, build_sync_terminal_usage_outcome,
build_terminal_usage_event_from_outcome, TerminalUsageOutcome, UsageTerminalState,
};

View File

@@ -0,0 +1,162 @@
use serde_json::json;
use aether_data::redis::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data::DataLayerError;
use super::config::UsageRuntimeConfig;
use super::event::UsageEvent;
#[derive(Debug, Clone)]
pub struct UsageQueue {
runner: RedisStreamRunner,
config: UsageRuntimeConfig,
stream: RedisStreamName,
group: RedisConsumerGroup,
dlq_stream: RedisStreamName,
}
impl UsageQueue {
pub fn new(
runner: RedisStreamRunner,
config: UsageRuntimeConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
let tuned_runner = RedisStreamRunner::new(
runner.client().clone(),
runner.keyspace().clone(),
usage_stream_runner_config(&config),
)?;
Ok(Self {
runner: tuned_runner,
stream: RedisStreamName(config.stream_key.clone()),
group: RedisConsumerGroup(config.consumer_group.clone()),
dlq_stream: RedisStreamName(config.dlq_stream_key.clone()),
config,
})
}
pub async fn ensure_consumer_group(&self) -> Result<(), DataLayerError> {
self.runner
.ensure_consumer_group(&self.stream, &self.group, "0-0")
.await
}
pub async fn enqueue(&self, event: &UsageEvent) -> Result<String, DataLayerError> {
let fields = event.to_stream_fields()?;
self.runner
.append_fields_with_maxlen(&self.stream, &fields, Some(self.config.stream_maxlen))
.await
}
pub async fn read_group(
&self,
consumer: &RedisConsumerName,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
self.runner
.read_group(&self.stream, &self.group, consumer)
.await
}
pub async fn claim_stale(
&self,
consumer: &RedisConsumerName,
start_id: &str,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
Ok(self
.runner
.claim_stale(
&self.stream,
&self.group,
consumer,
start_id,
RedisStreamReclaimConfig {
min_idle_ms: self.config.reclaim_idle_ms,
count: self.config.reclaim_count,
},
)
.await?
.entries)
}
pub async fn ack_and_delete(&self, ids: &[String]) -> Result<(), DataLayerError> {
self.runner.ack(&self.stream, &self.group, ids).await?;
self.runner.delete(&self.stream, ids).await?;
Ok(())
}
pub async fn push_dead_letter(
&self,
entry: &RedisStreamEntry,
error: &str,
) -> Result<String, DataLayerError> {
self.runner
.append_json(
&self.dlq_stream,
"payload",
&json!({
"entry_id": entry.id,
"fields": entry.fields,
"error": error,
}),
)
.await
}
}
fn usage_stream_runner_config(config: &UsageRuntimeConfig) -> RedisStreamRunnerConfig {
let read_block_ms = config.consumer_block_ms.max(1);
let command_timeout_ms = read_block_ms.saturating_add(2_000).max(5_000);
RedisStreamRunnerConfig {
command_timeout_ms: Some(command_timeout_ms),
read_block_ms: Some(read_block_ms),
read_count: config.consumer_batch_size.max(1),
}
}
#[cfg(test)]
mod tests {
use super::{usage_stream_runner_config, UsageQueue};
use crate::UsageRuntimeConfig;
use aether_data::redis::{RedisClientConfig, RedisClientFactory, RedisStreamRunner};
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
aether_data::redis::RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn usage_queue_applies_runtime_block_and_batch_settings() {
let config = UsageRuntimeConfig {
enabled: true,
consumer_block_ms: 750,
consumer_batch_size: 123,
..UsageRuntimeConfig::default()
};
let queue = UsageQueue::new(sample_runner(), config)
.expect("usage queue should build from runtime config");
assert_eq!(
queue.runner.config(),
usage_stream_runner_config(&queue.config)
);
assert_eq!(queue.runner.config().read_block_ms, Some(750));
assert_eq!(queue.runner.config().read_count, 123);
assert_eq!(queue.runner.config().command_timeout_ms, Some(5_000));
}
}

View File

@@ -0,0 +1,116 @@
use aether_data::repository::usage::UpsertUsageRecord;
use aether_data::DataLayerError;
use crate::{UsageEvent, UsageEventType};
pub fn build_upsert_usage_record_from_event(
event: &UsageEvent,
) -> Result<UpsertUsageRecord, DataLayerError> {
let (status, billing_status) = lifecycle_status_and_billing(event.event_type);
let data = event.data.clone();
let now_unix_secs = event.timestamp_ms / 1_000;
Ok(UpsertUsageRecord {
request_id: event.request_id.clone(),
user_id: data.user_id,
api_key_id: data.api_key_id,
username: data.username,
api_key_name: data.api_key_name,
provider_name: data.provider_name,
model: data.model,
target_model: data.target_model,
provider_id: empty_to_none(data.provider_id),
provider_endpoint_id: empty_to_none(data.provider_endpoint_id),
provider_api_key_id: empty_to_none(data.provider_api_key_id),
request_type: data.request_type,
api_format: data.api_format,
api_family: data.api_family,
endpoint_kind: data.endpoint_kind,
endpoint_api_format: data.endpoint_api_format,
provider_api_family: data.provider_api_family,
provider_endpoint_kind: data.provider_endpoint_kind,
has_format_conversion: data.has_format_conversion,
is_stream: data.is_stream,
input_tokens: data.input_tokens,
output_tokens: data.output_tokens,
total_tokens: data.total_tokens,
cache_creation_input_tokens: data.cache_creation_input_tokens,
cache_read_input_tokens: data.cache_read_input_tokens,
cache_creation_cost_usd: data.cache_creation_cost_usd,
cache_read_cost_usd: data.cache_read_cost_usd,
output_price_per_1m: data.output_price_per_1m,
total_cost_usd: data.total_cost_usd,
actual_total_cost_usd: data.actual_total_cost_usd,
status_code: data.status_code,
error_message: data.error_message,
error_category: data.error_category,
response_time_ms: data.response_time_ms,
first_byte_time_ms: data.first_byte_time_ms,
status: status.to_string(),
billing_status: billing_status.to_string(),
request_headers: data.request_headers,
request_body: data.request_body,
provider_request_headers: data.provider_request_headers,
provider_request_body: data.provider_request_body,
response_headers: data.response_headers,
response_body: data.response_body,
client_response_headers: data.client_response_headers,
client_response_body: data.client_response_body,
request_metadata: data.request_metadata,
finalized_at_unix_secs: Some(now_unix_secs),
created_at_unix_secs: Some(now_unix_secs),
updated_at_unix_secs: now_unix_secs,
})
}
fn lifecycle_status_and_billing(event_type: UsageEventType) -> (&'static str, &'static str) {
match event_type {
UsageEventType::Pending => ("pending", "pending"),
UsageEventType::Streaming => ("streaming", "pending"),
UsageEventType::Completed => ("completed", "pending"),
UsageEventType::Failed => ("failed", "void"),
UsageEventType::Cancelled => ("cancelled", "void"),
}
}
fn empty_to_none(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
#[cfg(test)]
mod tests {
use crate::{UsageEvent, UsageEventData, UsageEventType};
use super::build_upsert_usage_record_from_event;
#[test]
fn builds_upsert_record_from_terminal_event() {
let record = build_upsert_usage_record_from_event(&UsageEvent {
event_type: UsageEventType::Completed,
request_id: "req-1".to_string(),
timestamp_ms: 1_700_000_000_000,
data: UsageEventData {
user_id: Some("user-1".to_string()),
api_key_id: Some("key-1".to_string()),
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
api_format: Some("openai:chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
input_tokens: Some(10),
output_tokens: Some(20),
total_tokens: Some(30),
status_code: Some(200),
..UsageEventData::default()
},
})
.expect("record should build");
assert_eq!(record.request_id, "req-1");
assert_eq!(record.status, "completed");
assert_eq!(record.billing_status, "pending");
assert_eq!(record.total_tokens, Some(30));
assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000));
}
}

View File

@@ -0,0 +1,643 @@
use std::collections::BTreeMap;
use aether_contracts::ExecutionTelemetry;
use base64::Engine as _;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const GEMINI_FILE_MAPPING_TTL_SECONDS: u64 = 60 * 60 * 48;
const GEMINI_FILE_MAPPING_CACHE_PREFIX: &str = "gemini_files:key";
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct GatewaySyncReportRequest {
pub trace_id: String,
pub report_kind: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub report_context: Option<serde_json::Value>,
pub status_code: u16,
pub headers: BTreeMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body_json: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_body_json: Option<serde_json::Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub body_base64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub telemetry: Option<ExecutionTelemetry>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct GatewayStreamReportRequest {
pub trace_id: String,
pub report_kind: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub report_context: Option<serde_json::Value>,
pub status_code: u16,
pub headers: BTreeMap<String, String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_body_base64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_body_base64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub telemetry: Option<ExecutionTelemetry>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InternalFinalizeRoute {
pub public_path: &'static str,
pub route_family: &'static str,
pub route_kind: &'static str,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GeminiFileMappingEntry {
pub file_name: String,
pub display_name: Option<String>,
pub mime_type: Option<String>,
}
pub fn infer_internal_finalize_signature(payload: &GatewaySyncReportRequest) -> Option<String> {
let report_context = payload.report_context.as_ref()?;
let from_context = report_context
.get("client_api_format")
.and_then(serde_json::Value::as_str)
.or_else(|| {
report_context
.get("provider_api_format")
.and_then(serde_json::Value::as_str)
})
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
if from_context.is_some() {
return from_context;
}
let report_kind = payload.report_kind.trim().to_ascii_lowercase();
if report_kind.starts_with("openai_chat_") {
return Some("openai:chat".to_string());
}
if report_kind.starts_with("openai_compact_") {
return Some("openai:compact".to_string());
}
if report_kind.starts_with("openai_cli_") {
return Some("openai:cli".to_string());
}
if report_kind.starts_with("openai_video_") {
return Some("openai:video".to_string());
}
if report_kind.starts_with("claude_chat_") {
return Some("claude:chat".to_string());
}
if report_kind.starts_with("claude_cli_") {
return Some("claude:cli".to_string());
}
if report_kind.starts_with("gemini_chat_") {
return Some("gemini:chat".to_string());
}
if report_kind.starts_with("gemini_cli_") {
return Some("gemini:cli".to_string());
}
if report_kind.starts_with("gemini_video_") {
return Some("gemini:video".to_string());
}
None
}
pub fn resolve_internal_finalize_route(signature: &str) -> Option<InternalFinalizeRoute> {
match signature {
"openai:chat" => Some(InternalFinalizeRoute {
public_path: "/v1/chat/completions",
route_family: "openai",
route_kind: "chat",
}),
"openai:cli" => Some(InternalFinalizeRoute {
public_path: "/v1/responses",
route_family: "openai",
route_kind: "cli",
}),
"openai:compact" => Some(InternalFinalizeRoute {
public_path: "/v1/responses/compact",
route_family: "openai",
route_kind: "compact",
}),
"openai:video" => Some(InternalFinalizeRoute {
public_path: "/v1/videos",
route_family: "openai",
route_kind: "video",
}),
"claude:chat" => Some(InternalFinalizeRoute {
public_path: "/v1/messages",
route_family: "claude",
route_kind: "chat",
}),
"claude:cli" => Some(InternalFinalizeRoute {
public_path: "/v1/messages",
route_family: "claude",
route_kind: "cli",
}),
"gemini:chat" => Some(InternalFinalizeRoute {
public_path: "/v1beta/models",
route_family: "gemini",
route_kind: "chat",
}),
"gemini:cli" => Some(InternalFinalizeRoute {
public_path: "/v1beta/models",
route_family: "gemini",
route_kind: "cli",
}),
"gemini:video" => Some(InternalFinalizeRoute {
public_path: "/v1beta/models",
route_family: "gemini",
route_kind: "video",
}),
_ => None,
}
}
pub fn normalize_gemini_file_name(file_name: &str) -> Option<String> {
let file_name = file_name.trim();
if file_name.is_empty() {
return None;
}
if file_name.starts_with("files/") {
Some(file_name.to_string())
} else {
Some(format!("files/{file_name}"))
}
}
pub fn gemini_file_mapping_cache_key(file_name: &str) -> String {
format!("{GEMINI_FILE_MAPPING_CACHE_PREFIX}:{file_name}")
}
pub fn extract_gemini_file_mapping_entries(
payload: &GatewaySyncReportRequest,
) -> Vec<GeminiFileMappingEntry> {
let Some(body) = extract_sync_report_body_json(payload) else {
return Vec::new();
};
let Some(object) = body.as_object() else {
return Vec::new();
};
let mut entries = Vec::new();
maybe_push_gemini_file_mapping_entry(&mut entries, object);
if let Some(file_object) = object.get("file").and_then(Value::as_object) {
maybe_push_gemini_file_mapping_entry(&mut entries, file_object);
}
if let Some(files) = object.get("files").and_then(Value::as_array) {
for item in files {
if let Some(file_object) = item.as_object() {
maybe_push_gemini_file_mapping_entry(&mut entries, file_object);
}
}
}
entries
}
pub fn report_request_id(report_context: Option<&serde_json::Value>) -> &str {
report_context
.and_then(|context| context.get("request_id"))
.and_then(Value::as_str)
.filter(|value| !value.trim().is_empty())
.unwrap_or("-")
}
pub fn is_local_ai_sync_report_kind(report_kind: &str) -> bool {
matches!(
report_kind,
"openai_chat_sync_success"
| "claude_chat_sync_success"
| "gemini_chat_sync_success"
| "openai_chat_sync_error"
| "claude_chat_sync_error"
| "gemini_chat_sync_error"
| "openai_cli_sync_success"
| "claude_cli_sync_success"
| "gemini_cli_sync_success"
| "openai_cli_sync_error"
| "openai_compact_sync_error"
| "claude_cli_sync_error"
| "gemini_cli_sync_error"
| "openai_video_create_sync_success"
| "openai_video_remix_sync_success"
| "gemini_video_create_sync_success"
| "openai_video_delete_sync_success"
| "openai_video_cancel_sync_success"
| "gemini_video_cancel_sync_success"
| "openai_video_create_sync_error"
| "openai_video_remix_sync_error"
| "gemini_video_create_sync_error"
| "openai_video_delete_sync_error"
| "openai_video_cancel_sync_error"
| "gemini_video_cancel_sync_error"
| "gemini_files_store_mapping"
| "gemini_files_delete_mapping"
)
}
pub fn is_local_ai_stream_report_kind(report_kind: &str) -> bool {
matches!(
report_kind,
"openai_chat_stream_success"
| "claude_chat_stream_success"
| "gemini_chat_stream_success"
| "openai_cli_stream_success"
| "claude_cli_stream_success"
| "gemini_cli_stream_success"
)
}
pub fn sync_report_represents_failure(
payload: &GatewaySyncReportRequest,
error_type: Option<&str>,
) -> bool {
if payload.report_kind == "openai_video_delete_sync_success" && payload.status_code == 404 {
return false;
}
payload.status_code >= 400
|| payload.report_kind.contains("error")
|| error_type.is_some()
|| payload
.body_json
.as_ref()
.and_then(|body| body.get("error"))
.is_some()
}
pub fn should_handle_local_sync_report(
report_context: Option<&serde_json::Value>,
report_kind: &str,
) -> bool {
crate::report_context::report_context_is_locally_actionable(report_context)
&& is_local_ai_sync_report_kind(report_kind)
}
pub fn should_handle_local_stream_report(
report_context: Option<&serde_json::Value>,
report_kind: &str,
) -> bool {
crate::report_context::report_context_is_locally_actionable(report_context)
&& is_local_ai_stream_report_kind(report_kind)
}
fn maybe_push_gemini_file_mapping_entry(
entries: &mut Vec<GeminiFileMappingEntry>,
object: &serde_json::Map<String, Value>,
) {
let file_name = object
.get("name")
.and_then(Value::as_str)
.and_then(normalize_gemini_file_name);
let Some(file_name) = file_name else {
return;
};
if entries.iter().any(|entry| entry.file_name == file_name) {
return;
}
entries.push(GeminiFileMappingEntry {
file_name,
display_name: object
.get("displayName")
.or_else(|| object.get("display_name"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
mime_type: object
.get("mimeType")
.or_else(|| object.get("mime_type"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
});
}
fn extract_sync_report_body_json(payload: &GatewaySyncReportRequest) -> Option<Value> {
if let Some(body_json) = payload.body_json.as_ref() {
return Some(body_json.clone());
}
if let Some(client_body_json) = payload.client_body_json.as_ref() {
return Some(client_body_json.clone());
}
if !content_type_starts_with(&payload.headers, "application/json") {
return None;
}
let body_base64 = payload.body_base64.as_deref()?;
let bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.ok()?;
serde_json::from_slice(&bytes).ok()
}
fn content_type_starts_with(headers: &BTreeMap<String, String>, expected_prefix: &str) -> bool {
headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case("content-type"))
.map(|(_, value)| value.trim().to_ascii_lowercase())
.is_some_and(|value| value.starts_with(expected_prefix))
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use base64::Engine as _;
use serde_json::json;
use super::{
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
infer_internal_finalize_signature, is_local_ai_stream_report_kind,
is_local_ai_sync_report_kind, normalize_gemini_file_name, report_request_id,
resolve_internal_finalize_route, should_handle_local_stream_report,
should_handle_local_sync_report, sync_report_represents_failure, GatewaySyncReportRequest,
GeminiFileMappingEntry, InternalFinalizeRoute,
};
fn sample_sync_report(report_kind: &str, status_code: u16) -> GatewaySyncReportRequest {
GatewaySyncReportRequest {
trace_id: "trace-123".to_string(),
report_kind: report_kind.to_string(),
report_context: None,
status_code,
headers: BTreeMap::new(),
body_json: None,
client_body_json: None,
body_base64: None,
telemetry: None,
}
}
fn sample_sync_report_with_context(
report_kind: &str,
report_context: serde_json::Value,
) -> GatewaySyncReportRequest {
GatewaySyncReportRequest {
trace_id: "trace-123".to_string(),
report_kind: report_kind.to_string(),
report_context: Some(report_context),
status_code: 200,
headers: BTreeMap::new(),
body_json: None,
client_body_json: None,
body_base64: None,
telemetry: None,
}
}
#[test]
fn classifies_local_ai_sync_report_kinds() {
assert!(is_local_ai_sync_report_kind(
"openai_video_create_sync_success"
));
assert!(is_local_ai_sync_report_kind("gemini_files_delete_mapping"));
assert!(!is_local_ai_sync_report_kind("unknown_sync_kind"));
}
#[test]
fn classifies_local_ai_stream_report_kinds() {
assert!(is_local_ai_stream_report_kind("openai_chat_stream_success"));
assert!(!is_local_ai_stream_report_kind("openai_chat_stream_error"));
}
#[test]
fn treats_openai_video_delete_404_success_as_non_failure() {
let payload = sample_sync_report("openai_video_delete_sync_success", 404);
assert!(!sync_report_represents_failure(&payload, None));
}
#[test]
fn detects_sync_report_failure_from_status_kind_error_type_or_body() {
let status_payload = sample_sync_report("openai_chat_sync_success", 500);
assert!(sync_report_represents_failure(&status_payload, None));
let kind_payload = sample_sync_report("openai_chat_sync_error", 200);
assert!(sync_report_represents_failure(&kind_payload, None));
let error_type_payload = sample_sync_report("openai_chat_sync_success", 200);
assert!(sync_report_represents_failure(
&error_type_payload,
Some("authentication_error")
));
let mut error_body_payload = sample_sync_report("openai_chat_sync_success", 200);
error_body_payload.body_json = Some(json!({"error": {"message": "bad request"}}));
assert!(sync_report_represents_failure(&error_body_payload, None));
let success_payload = sample_sync_report("openai_chat_sync_success", 200);
assert!(!sync_report_represents_failure(&success_payload, None));
}
#[test]
fn infers_internal_finalize_signature_from_context_or_report_kind() {
let from_context = sample_sync_report_with_context(
"unknown_sync_finalize",
json!({"client_api_format": "gemini:video"}),
);
assert_eq!(
infer_internal_finalize_signature(&from_context),
Some("gemini:video".to_string())
);
let from_report_kind =
sample_sync_report_with_context("openai_video_create_sync_finalize", json!({}));
assert_eq!(
infer_internal_finalize_signature(&from_report_kind),
Some("openai:video".to_string())
);
let unknown = sample_sync_report("unknown_sync_finalize", 200);
assert_eq!(infer_internal_finalize_signature(&unknown), None);
}
#[test]
fn resolves_internal_finalize_route_for_supported_signatures() {
assert_eq!(
resolve_internal_finalize_route("openai:compact"),
Some(InternalFinalizeRoute {
public_path: "/v1/responses/compact",
route_family: "openai",
route_kind: "compact",
})
);
assert_eq!(
resolve_internal_finalize_route("gemini:video"),
Some(InternalFinalizeRoute {
public_path: "/v1beta/models",
route_family: "gemini",
route_kind: "video",
})
);
assert_eq!(resolve_internal_finalize_route("unknown:kind"), None);
}
#[test]
fn normalizes_gemini_file_names() {
assert_eq!(
normalize_gemini_file_name("abc123"),
Some("files/abc123".to_string())
);
assert_eq!(
normalize_gemini_file_name("files/abc123"),
Some("files/abc123".to_string())
);
assert_eq!(normalize_gemini_file_name(" "), None);
}
#[test]
fn builds_gemini_file_mapping_cache_keys() {
assert_eq!(
gemini_file_mapping_cache_key("files/abc123"),
"gemini_files:key:files/abc123"
);
}
#[test]
fn extracts_gemini_file_mapping_entries_from_supported_shapes() {
let payload = GatewaySyncReportRequest {
trace_id: "trace-123".to_string(),
report_kind: "gemini_files_store_mapping".to_string(),
report_context: None,
status_code: 200,
headers: BTreeMap::new(),
body_json: Some(json!({
"name": "abc123",
"displayName": "root-name",
"file": {
"name": "files/def456",
"mimeType": "image/png"
},
"files": [
{
"name": "abc123",
"display_name": "deduped"
},
{
"name": "ghi789",
"display_name": "third"
}
]
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let entries = extract_gemini_file_mapping_entries(&payload);
assert_eq!(
entries,
vec![
GeminiFileMappingEntry {
file_name: "files/abc123".to_string(),
display_name: Some("root-name".to_string()),
mime_type: None,
},
GeminiFileMappingEntry {
file_name: "files/def456".to_string(),
display_name: None,
mime_type: Some("image/png".to_string()),
},
GeminiFileMappingEntry {
file_name: "files/ghi789".to_string(),
display_name: Some("third".to_string()),
mime_type: None,
}
]
);
}
#[test]
fn extracts_gemini_file_mapping_entries_from_base64_json_body() {
let encoded_body = base64::engine::general_purpose::STANDARD.encode(
serde_json::to_vec(&json!({
"name": "base64-file"
}))
.expect("json should encode"),
);
let payload = GatewaySyncReportRequest {
trace_id: "trace-123".to_string(),
report_kind: "gemini_files_store_mapping".to_string(),
report_context: None,
status_code: 200,
headers: BTreeMap::from([(
"content-type".to_string(),
"application/json; charset=utf-8".to_string(),
)]),
body_json: None,
client_body_json: None,
body_base64: Some(encoded_body),
telemetry: None,
};
let entries = extract_gemini_file_mapping_entries(&payload);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].file_name, "files/base64-file");
}
#[test]
fn reads_report_request_id_from_context() {
assert_eq!(
report_request_id(Some(&json!({"request_id": "req-123"}))),
"req-123"
);
assert_eq!(report_request_id(Some(&json!({"request_id": " "}))), "-");
assert_eq!(report_request_id(None), "-");
}
#[test]
fn decides_when_local_sync_report_should_be_handled() {
assert!(should_handle_local_sync_report(
Some(&json!({
"request_id": "req-123",
"provider_id": "provider-123",
"endpoint_id": "endpoint-123",
"key_id": "key-123"
})),
"openai_chat_sync_success"
));
assert!(!should_handle_local_sync_report(
Some(&json!({"request_id": "req-123"})),
"openai_chat_sync_success"
));
assert!(!should_handle_local_sync_report(
Some(&json!({
"request_id": "req-123",
"provider_id": "provider-123",
"endpoint_id": "endpoint-123",
"key_id": "key-123"
})),
"unknown_sync_kind"
));
}
#[test]
fn decides_when_local_stream_report_should_be_handled() {
assert!(should_handle_local_stream_report(
Some(&json!({
"request_id": "req-123",
"provider_id": "provider-123",
"endpoint_id": "endpoint-123",
"key_id": "key-123"
})),
"openai_chat_stream_success"
));
assert!(!should_handle_local_stream_report(
Some(&json!({
"request_id": "req-123",
"provider_id": "provider-123",
"endpoint_id": "endpoint-123",
"key_id": "key-123"
})),
"openai_chat_stream_error"
));
}
}

View File

@@ -0,0 +1,235 @@
use aether_data::repository::candidates::StoredRequestCandidate;
use aether_data::repository::video_tasks::StoredVideoTask;
use serde_json::{Map, Value};
pub fn report_context_is_locally_actionable(report_context: Option<&Value>) -> bool {
let Some(context) = report_context else {
return false;
};
has_non_empty_str(context, "request_id")
&& (has_non_empty_str(context, "candidate_id")
|| has_u64(context, "candidate_index")
|| has_non_empty_str(context, "provider_id")
|| has_non_empty_str(context, "endpoint_id")
|| has_non_empty_str(context, "key_id"))
}
pub fn build_locally_actionable_report_context_from_request_candidate(
context: &Value,
candidate: &StoredRequestCandidate,
) -> Option<Value> {
let mut object = context.as_object()?.clone();
insert_missing_string_value(&mut object, "candidate_id", Some(candidate.id.as_str()));
if !object.contains_key("candidate_index") {
object.insert(
"candidate_index".to_string(),
Value::Number(candidate.candidate_index.into()),
);
}
insert_missing_optional_string_value(
&mut object,
"provider_id",
candidate.provider_id.as_deref(),
);
insert_missing_optional_string_value(
&mut object,
"endpoint_id",
candidate.endpoint_id.as_deref(),
);
insert_missing_optional_string_value(&mut object, "key_id", candidate.key_id.as_deref());
insert_missing_optional_string_value(&mut object, "user_id", candidate.user_id.as_deref());
insert_missing_optional_string_value(
&mut object,
"api_key_id",
candidate.api_key_id.as_deref(),
);
let resolved = Value::Object(object);
report_context_is_locally_actionable(Some(&resolved)).then_some(resolved)
}
pub fn build_locally_actionable_report_context_from_video_task(
context: &Value,
task: &StoredVideoTask,
) -> Option<Value> {
let mut object = context.as_object()?.clone();
insert_missing_string_value(&mut object, "request_id", Some(task.request_id.as_str()));
insert_missing_optional_string_value(&mut object, "provider_id", task.provider_id.as_deref());
insert_missing_optional_string_value(&mut object, "endpoint_id", task.endpoint_id.as_deref());
insert_missing_optional_string_value(&mut object, "key_id", task.key_id.as_deref());
insert_missing_optional_string_value(&mut object, "user_id", task.user_id.as_deref());
insert_missing_optional_string_value(&mut object, "api_key_id", task.api_key_id.as_deref());
insert_missing_optional_string_value(
&mut object,
"client_api_format",
task.client_api_format.as_deref(),
);
insert_missing_optional_string_value(
&mut object,
"provider_api_format",
task.provider_api_format.as_deref(),
);
Some(Value::Object(object))
}
fn insert_missing_string_value(object: &mut Map<String, Value>, key: &str, value: Option<&str>) {
if object.contains_key(key) {
return;
}
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
object.insert(key.to_string(), Value::String(value.to_string()));
}
fn insert_missing_optional_string_value(
object: &mut Map<String, Value>,
key: &str,
value: Option<&str>,
) {
insert_missing_string_value(object, key, value);
}
fn has_non_empty_str(value: &Value, key: &str) -> bool {
value
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
fn has_u64(value: &Value, key: &str) -> bool {
value.get(key).and_then(Value::as_u64).is_some()
}
#[cfg(test)]
mod tests {
use aether_data::repository::candidates::{RequestCandidateStatus, StoredRequestCandidate};
use aether_data::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
use serde_json::{json, Value};
use super::{
build_locally_actionable_report_context_from_request_candidate,
build_locally_actionable_report_context_from_video_task,
report_context_is_locally_actionable,
};
fn sample_candidate() -> StoredRequestCandidate {
StoredRequestCandidate {
id: "cand-1".to_string(),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: None,
api_key_name: None,
candidate_index: 0,
retry_index: 0,
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
status: RequestCandidateStatus::Pending,
skip_reason: None,
is_cached: false,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_secs: 1,
started_at_unix_secs: None,
finished_at_unix_secs: None,
}
}
fn sample_video_task() -> StoredVideoTask {
StoredVideoTask {
id: "task-1".to_string(),
short_id: Some("short-1".to_string()),
request_id: "req-1".to_string(),
user_id: Some("user-1".to_string()),
api_key_id: Some("api-key-1".to_string()),
username: None,
api_key_name: None,
external_task_id: Some("ext-1".to_string()),
provider_id: Some("provider-1".to_string()),
endpoint_id: Some("endpoint-1".to_string()),
key_id: Some("key-1".to_string()),
client_api_format: Some("openai:video".to_string()),
provider_api_format: Some("openai:video".to_string()),
format_converted: false,
model: Some("sora".to_string()),
prompt: None,
original_request_body: None,
duration_seconds: None,
resolution: None,
aspect_ratio: None,
size: None,
status: VideoTaskStatus::Submitted,
progress_percent: 0,
progress_message: None,
retry_count: 0,
poll_interval_seconds: 10,
next_poll_at_unix_secs: None,
poll_count: 0,
max_poll_count: 360,
created_at_unix_secs: 1,
submitted_at_unix_secs: Some(1),
completed_at_unix_secs: None,
updated_at_unix_secs: 1,
error_code: None,
error_message: None,
video_url: None,
request_metadata: None,
}
}
#[test]
fn detects_locally_actionable_report_context() {
assert!(report_context_is_locally_actionable(Some(&json!({
"request_id": "req-1",
"provider_id": "provider-1"
}))));
assert!(!report_context_is_locally_actionable(Some(&json!({
"request_id": "req-1"
}))));
}
#[test]
fn patches_locally_actionable_report_context_from_candidate() {
let resolved = build_locally_actionable_report_context_from_request_candidate(
&json!({"request_id": "req-1"}),
&sample_candidate(),
)
.expect("candidate context should resolve");
assert_eq!(
resolved.get("candidate_id").and_then(Value::as_str),
Some("cand-1")
);
assert_eq!(
resolved.get("provider_id").and_then(Value::as_str),
Some("provider-1")
);
}
#[test]
fn patches_locally_actionable_report_context_from_video_task() {
let resolved = build_locally_actionable_report_context_from_video_task(
&json!({"local_task_id": "task-1"}),
&sample_video_task(),
)
.expect("video task context should resolve");
assert_eq!(
resolved.get("request_id").and_then(Value::as_str),
Some("req-1")
);
assert_eq!(
resolved.get("provider_api_format").and_then(Value::as_str),
Some("openai:video")
);
}
}

View File

@@ -0,0 +1,315 @@
use std::sync::Arc;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_data::redis::RedisStreamRunner;
use aether_data::DataLayerError;
use async_trait::async_trait;
use tracing::warn;
use crate::{
build_pending_usage_record, build_stream_terminal_usage_outcome, build_streaming_usage_record,
build_sync_terminal_usage_outcome, build_terminal_usage_event_from_outcome,
build_upsert_usage_record_from_event, build_usage_queue_worker, settle_usage_if_needed,
GatewayStreamReportRequest, GatewaySyncReportRequest, UsageEvent, UsageQueue,
UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter, UsageTerminalState,
};
#[async_trait]
pub trait UsageBillingEventEnricher: Send + Sync {
async fn enrich_usage_event(&self, event: &mut UsageEvent) -> Result<(), DataLayerError>;
}
pub trait UsageRuntimeAccess:
UsageRecordWriter + UsageSettlementWriter + UsageBillingEventEnricher + Send + Sync
{
fn has_usage_writer(&self) -> bool;
fn has_usage_worker_runner(&self) -> bool;
fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
}
#[derive(Debug, Clone)]
pub struct UsageRuntime {
config: UsageRuntimeConfig,
}
impl Default for UsageRuntime {
fn default() -> Self {
Self::disabled()
}
}
impl UsageRuntime {
pub fn disabled() -> Self {
Self {
config: UsageRuntimeConfig::disabled(),
}
}
pub fn new(config: UsageRuntimeConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub fn is_enabled(&self) -> bool {
self.config.enabled
}
pub fn can_spawn_worker<T>(&self, data: &T) -> bool
where
T: UsageRuntimeAccess,
{
self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_runner()
}
pub fn spawn_worker<T>(&self, data: Arc<T>) -> Option<tokio::task::JoinHandle<()>>
where
T: UsageRuntimeAccess + 'static,
{
if !self.can_spawn_worker(data.as_ref()) {
return None;
}
let runner = data.usage_worker_runner()?;
let worker = build_usage_queue_worker(runner, data, self.config.clone()).ok()?;
Some(worker.spawn())
}
pub async fn record_pending<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) where
T: UsageRuntimeAccess,
{
if !self.is_enabled() {
return;
}
let now_unix_secs = now_unix_secs();
match build_pending_usage_record(plan, report_context, now_unix_secs) {
Ok(record) => {
if let Err(err) = data.upsert_usage_record(record).await {
warn!(
event_name = "usage_pending_record_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to record sync pending usage"
);
}
}
Err(err) => {
warn!(
event_name = "usage_pending_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build sync pending usage"
)
}
}
}
pub async fn record_stream_started<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
headers: &std::collections::BTreeMap<String, String>,
telemetry: Option<&ExecutionTelemetry>,
) where
T: UsageRuntimeAccess,
{
if !self.is_enabled() {
return;
}
let now_unix_secs = now_unix_secs();
match build_streaming_usage_record(
plan,
report_context,
status_code,
headers,
telemetry,
now_unix_secs,
) {
Ok(record) => {
if let Err(err) = data.upsert_usage_record(record).await {
warn!(
event_name = "usage_stream_record_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to record stream usage"
);
}
}
Err(err) => {
warn!(
event_name = "usage_stream_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build stream usage"
)
}
}
}
pub async fn record_sync_terminal<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
) where
T: UsageRuntimeAccess,
{
if !self.is_enabled() {
return;
}
match build_terminal_usage_event_from_outcome(build_sync_terminal_usage_outcome(
plan,
report_context,
payload,
)) {
Ok(mut event) => {
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_sync_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to enrich sync usage event with billing"
);
}
self.enqueue_or_write_terminal(data, event).await
}
Err(err) => {
warn!(
event_name = "usage_sync_terminal_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build sync terminal usage event"
)
}
}
}
pub async fn record_stream_terminal<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewayStreamReportRequest,
cancelled: bool,
) where
T: UsageRuntimeAccess,
{
if !self.is_enabled() {
return;
}
let mut outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
if cancelled {
outcome.terminal_state = UsageTerminalState::Cancelled;
}
match build_terminal_usage_event_from_outcome(outcome) {
Ok(mut event) => {
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_stream_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to enrich stream usage event with billing"
);
}
self.enqueue_or_write_terminal(data, event).await
}
Err(err) => {
warn!(
event_name = "usage_stream_terminal_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build stream terminal usage event"
)
}
}
}
async fn enqueue_or_write_terminal<T>(&self, data: &T, event: UsageEvent)
where
T: UsageRuntimeAccess,
{
if let Some(runner) = data.usage_worker_runner() {
match UsageQueue::new(runner, self.config.clone()) {
Ok(queue) => match queue.enqueue(&event).await {
Ok(_) => return,
Err(err) => {
warn!(
event_name = "usage_terminal_enqueue_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "direct_write",
error = %err,
"usage runtime failed to enqueue terminal usage event; falling back to direct write"
)
}
},
Err(err) => {
warn!(
event_name = "usage_terminal_queue_init_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "direct_write",
error = %err,
"usage runtime failed to build queue; falling back to direct write"
)
}
}
}
match build_upsert_usage_record_from_event(&event) {
Ok(record) => match data.upsert_usage_record(record).await {
Ok(Some(stored)) => {
if let Err(err) = settle_usage_if_needed(data, &stored).await {
warn!(
event_name = "usage_terminal_settlement_failed",
log_type = "event",
request_id = %event.request_id,
error = %err,
"usage runtime failed to settle terminal usage directly"
);
}
}
Ok(None) => {}
Err(err) => {
warn!(
event_name = "usage_terminal_upsert_failed",
log_type = "event",
request_id = %event.request_id,
error = %err,
"usage runtime failed to upsert terminal usage directly"
);
}
},
Err(err) => {
warn!(
event_name = "usage_terminal_upsert_build_failed",
log_type = "event",
request_id = %event.request_id,
error = %err,
"usage runtime failed to build terminal usage upsert"
)
}
}
}
}
fn now_unix_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}

View File

@@ -0,0 +1,190 @@
use aether_data::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data::repository::usage::StoredRequestUsageAudit;
use aether_data::{DataLayerError, DataLayerError::InvalidInput};
use async_trait::async_trait;
#[async_trait]
pub trait UsageSettlementWriter: Send + Sync {
fn has_usage_settlement_writer(&self) -> bool;
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError>;
}
pub async fn settle_usage_if_needed(
writer: &dyn UsageSettlementWriter,
usage: &StoredRequestUsageAudit,
) -> Result<(), DataLayerError> {
if !writer.has_usage_settlement_writer() || usage.billing_status != "pending" {
return Ok(());
}
if !matches!(usage.status.as_str(), "completed" | "failed" | "cancelled") {
return Ok(());
}
let finalized_at_unix_secs = usage
.finalized_at_unix_secs
.or(Some(usage.updated_at_unix_secs));
let input = UsageSettlementInput {
request_id: usage.request_id.clone(),
user_id: usage.user_id.clone(),
api_key_id: usage.api_key_id.clone(),
provider_id: usage.provider_id.clone(),
status: usage.status.clone(),
billing_status: usage.billing_status.clone(),
total_cost_usd: finite_cost(usage.total_cost_usd)?,
actual_total_cost_usd: finite_cost(usage.actual_total_cost_usd)?,
finalized_at_unix_secs,
};
let _ = writer.settle_usage(input).await?;
Ok(())
}
fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
if value.is_finite() {
Ok(value)
} else {
Err(InvalidInput(
"wallet settlement cost must be finite".to_string(),
))
}
}
#[cfg(test)]
mod tests {
use std::sync::Mutex;
use super::{settle_usage_if_needed, UsageSettlementWriter};
use aether_data::repository::settlement::UsageSettlementInput;
use aether_data::repository::usage::StoredRequestUsageAudit;
use async_trait::async_trait;
#[derive(Default)]
struct TestSettlementWriter {
has_writer: bool,
inputs: Mutex<Vec<UsageSettlementInput>>,
}
#[async_trait]
impl UsageSettlementWriter for TestSettlementWriter {
fn has_usage_settlement_writer(&self) -> bool {
self.has_writer
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<
Option<aether_data::repository::settlement::StoredUsageSettlement>,
aether_data::DataLayerError,
> {
self.inputs
.lock()
.expect("settlement inputs lock")
.push(input);
Ok(None)
}
}
fn sample_usage() -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
"usage-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
Some("key-1".to_string()),
None,
None,
"openai".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
None,
None,
None,
None,
None,
None,
None,
None,
None,
false,
false,
0,
0,
0,
1.25,
0.75,
Some(200),
None,
None,
None,
None,
"completed".to_string(),
"pending".to_string(),
100,
200,
None,
)
.expect("usage should build")
}
#[tokio::test]
async fn settles_pending_terminal_usage() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let usage = sample_usage();
settle_usage_if_needed(&writer, &usage)
.await
.expect("settlement should succeed");
let inputs = writer.inputs.lock().expect("settlement inputs lock");
assert_eq!(inputs.len(), 1);
assert_eq!(inputs[0].request_id, "req-1");
assert_eq!(inputs[0].status, "completed");
assert_eq!(inputs[0].billing_status, "pending");
assert_eq!(inputs[0].finalized_at_unix_secs, Some(200));
assert_eq!(inputs[0].total_cost_usd, 1.25);
assert_eq!(inputs[0].actual_total_cost_usd, 0.75);
}
#[tokio::test]
async fn skips_when_usage_is_not_pending_or_terminal() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.billing_status = "settled".to_string();
usage.status = "streaming".to_string();
settle_usage_if_needed(&writer, &usage)
.await
.expect("skipped settlement should succeed");
let inputs = writer.inputs.lock().expect("settlement inputs lock");
assert!(inputs.is_empty());
}
#[tokio::test]
async fn rejects_non_finite_costs_before_writing() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.total_cost_usd = f64::NAN;
let err = settle_usage_if_needed(&writer, &usage)
.await
.expect_err("non-finite costs should be rejected");
assert!(matches!(err, aether_data::DataLayerError::InvalidInput(_)));
let inputs = writer.inputs.lock().expect("settlement inputs lock");
assert!(inputs.is_empty());
}
}

View File

@@ -0,0 +1,87 @@
use std::collections::BTreeMap;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
pub struct StandardizedUsage {
pub input_tokens: i64,
pub output_tokens: i64,
pub cache_creation_tokens: i64,
pub cache_read_tokens: i64,
pub reasoning_tokens: i64,
pub cache_storage_token_hours: f64,
pub request_count: i64,
pub dimensions: BTreeMap<String, serde_json::Value>,
}
impl StandardizedUsage {
pub fn new() -> Self {
Self {
request_count: 1,
..Self::default()
}
}
pub fn get(&self, field_name: &str) -> Option<serde_json::Value> {
match field_name {
"input_tokens" => Some(serde_json::json!(self.input_tokens)),
"output_tokens" => Some(serde_json::json!(self.output_tokens)),
"cache_creation_tokens" => Some(serde_json::json!(self.cache_creation_tokens)),
"cache_read_tokens" => Some(serde_json::json!(self.cache_read_tokens)),
"reasoning_tokens" => Some(serde_json::json!(self.reasoning_tokens)),
"cache_storage_token_hours" => Some(serde_json::json!(self.cache_storage_token_hours)),
"request_count" => Some(serde_json::json!(self.request_count)),
"extra" | "dimensions" => Some(serde_json::json!(self.dimensions)),
_ => self.dimensions.get(field_name).cloned(),
}
}
pub fn set(&mut self, field_name: &str, value: impl Into<serde_json::Value>) {
let value = value.into();
match field_name {
"input_tokens" => self.input_tokens = as_i64(&value, 0),
"output_tokens" => self.output_tokens = as_i64(&value, 0),
"cache_creation_tokens" => self.cache_creation_tokens = as_i64(&value, 0),
"cache_read_tokens" => self.cache_read_tokens = as_i64(&value, 0),
"reasoning_tokens" => self.reasoning_tokens = as_i64(&value, 0),
"cache_storage_token_hours" => self.cache_storage_token_hours = as_f64(&value, 0.0),
"request_count" => self.request_count = as_i64(&value, 0),
"extra" | "dimensions" => {
self.dimensions = match value {
serde_json::Value::Object(map) => map.into_iter().collect(),
_ => BTreeMap::new(),
}
}
_ => {
self.dimensions.insert(field_name.to_string(), value);
}
}
}
}
fn as_i64(value: &serde_json::Value, default: i64) -> i64 {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
.unwrap_or(default)
}
fn as_f64(value: &serde_json::Value, default: f64) -> f64 {
value.as_f64().unwrap_or(default)
}
#[cfg(test)]
mod tests {
use super::StandardizedUsage;
#[test]
fn standardized_usage_reads_and_writes_known_and_extra_fields() {
let mut usage = StandardizedUsage::new();
usage.set("input_tokens", 10);
usage.set("custom_dimension", "value");
assert_eq!(usage.get("input_tokens"), Some(serde_json::json!(10)));
assert_eq!(
usage.get("custom_dimension"),
Some(serde_json::json!("value"))
);
}
}

Some files were not shown because too many files have changed in this diff Show More