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

feat(billing): 引入 model_id 精确计费查找路径，传播模型 ID 至用量事件与 report context

Browse Source 
- UsageEventData 新增 model_id / global_model_id 字段，write.rs seed 结构体同步补充
- report_context 新增 model_id / global_model_id / global_model_name，各 payload 构建时从 candidate 传入
- event_enrichment 优先按 model_id 精确查找计费上下文，回退为按名称多轮查找（保留 NoRule 结果降级逻辑）
- BillingReadRepository 新增 find_model_context_by_model_id，memory/sql 分别实现；SQL 查询重构支持按 provider_model_name 和 mappings 匹配并按优先级排序
- pricing.rs 修复：model_tiered_pricing 为空 tiers 时回退到 default_tiered_pricing
- request_metadata 允许字段列表补充 model_id / global_model_id / global_model_name
- admin usage 路由信息输出及脱敏字段列表同步新增三个 model 相关字段
This commit is contained in:
fawney19
2026-04-24 14:40:28 +08:00
parent f3c9835759
commit 780f09c1a2
20 changed files with 665 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/payload.rs
View File

@@ -91,6 +91,9 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),

18
apps/aether-gateway/src/ai_pipeline/planner/report_context.rs
View File

@@ -16,6 +16,9 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) endpoint_id: &'a str,
pub(crate) key_id: &'a str,
pub(crate) key_name: Option<&'a str>,
pub(crate) model_id: Option<&'a str>,
pub(crate) global_model_id: Option<&'a str>,
pub(crate) global_model_name: Option<&'a str>,
pub(crate) provider_api_format: &'a str,
pub(crate) client_api_format: &'a str,
pub(crate) mapped_model: Option<&'a str>,
@@ -136,6 +139,21 @@ pub(crate) fn build_local_execution_report_context(
if let Some(key_name) = parts.key_name {
object.insert("key_name".to_string(), Value::String(key_name.to_string()));
}
if let Some(model_id) = parts.model_id {
object.insert("model_id".to_string(), Value::String(model_id.to_string()));
}
if let Some(global_model_id) = parts.global_model_id {
object.insert(
"global_model_id".to_string(),
Value::String(global_model_id.to_string()),
);
}
if let Some(global_model_name) = parts.global_model_name {
object.insert(
"global_model_name".to_string(),
Value::String(global_model_name.to_string()),
);
}
if let Some(mapped_model) = parts.mapped_model {
object.insert(
"mapped_model".to_string(),

3
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/decision.rs
View File

@@ -76,6 +76,9 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
mapped_model: None,

3
apps/aether-gateway/src/ai_pipeline/planner/specialized/image/decision.rs
View File

@@ -68,6 +68,9 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),

3
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/decision.rs
View File

@@ -59,6 +59,9 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: None,
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),

3
apps/aether-gateway/src/ai_pipeline/planner/standard/family/payload.rs
View File

@@ -65,6 +65,9 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),

3
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/payload.rs
View File

@@ -82,6 +82,9 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: "openai:chat",
mapped_model: Some(&resolved.mapped_model),

3
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/payload.rs
View File

@@ -86,6 +86,9 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
endpoint_id: &candidate.endpoint_id,
key_id: &candidate.key_id,
key_name: Some(&candidate.key_name),
model_id: Some(&candidate.model_id),
global_model_id: Some(&candidate.global_model_id),
global_model_name: Some(&candidate.global_model_name),
provider_api_format: &resolved.provider_api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),

15
apps/aether-gateway/src/data/state/integrations.rs
View File

@@ -157,6 +157,21 @@ impl MinimalCandidateSelectionRowSource for GatewayDataState {
#[async_trait]
impl BillingModelContextLookup for GatewayDataState {
async fn find_billing_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
GatewayDataState::find_billing_model_context_by_model_id(
self,
provider_id,
provider_api_key_id,
model_id,
)
.await
}
async fn find_billing_model_context(
&self,
provider_id: &str,

16
apps/aether-gateway/src/data/state/runtime.rs
View File

@@ -1241,6 +1241,22 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_billing_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.find_model_context_by_model_id(provider_id, provider_api_key_id, model_id)
.await
}
None => Ok(None),
}
}
pub(crate) async fn read_request_candidate_trace(
&self,
request_id: &str,

25
crates/aether-admin/src/observability/usage.rs
View File

@@ -285,6 +285,9 @@ fn admin_usage_strip_routing_metadata(metadata: &mut serde_json::Map<String, Val
metadata.remove("candidate_id");
metadata.remove("candidate_index");
metadata.remove("key_name");
metadata.remove("model_id");
metadata.remove("global_model_id");
metadata.remove("global_model_name");
metadata.remove("planner_kind");
metadata.remove("route_family");
metadata.remove("route_kind");
@@ -795,6 +798,17 @@ fn admin_usage_routing_json(
"key_name",
provider_key_name.or_else(|| item.routing_key_name()),
);
maybe_insert_string_field(&mut routing, "model_id", item.routing_model_id());
maybe_insert_string_field(
&mut routing,
"global_model_id",
item.routing_global_model_id(),
);
maybe_insert_string_field(
&mut routing,
"global_model_name",
item.routing_global_model_name(),
);
maybe_insert_string_field(&mut routing, "planner_kind", item.routing_planner_kind());
maybe_insert_string_field(&mut routing, "route_family", item.routing_route_family());
maybe_insert_string_field(&mut routing, "route_kind", item.routing_route_kind());
@@ -2395,7 +2409,10 @@ mod tests {
fn detail_payload_exposes_typed_routing_section() {
let item = StoredRequestUsageAudit {
request_metadata: Some(json!({
"trace_id": "trace-routing-detail"
"trace_id": "trace-routing-detail",
"model_id": "model-1",
"global_model_id": "global-model-1",
"global_model_name": "gpt-5"
})),
candidate_id: Some("cand-1".to_string()),
candidate_index: Some(2),
@@ -2423,6 +2440,9 @@ mod tests {
assert_eq!(payload["routing"]["candidate_id"], "cand-1");
assert_eq!(payload["routing"]["candidate_index"], 2);
assert_eq!(payload["routing"]["key_name"], "resolved-primary");
assert_eq!(payload["routing"]["model_id"], "model-1");
assert_eq!(payload["routing"]["global_model_id"], "global-model-1");
assert_eq!(payload["routing"]["global_model_name"], "gpt-5");
assert_eq!(payload["routing"]["planner_kind"], "claude_cli_sync");
assert_eq!(payload["routing"]["route_family"], "claude");
assert_eq!(payload["routing"]["route_kind"], "cli");
@@ -2437,6 +2457,9 @@ mod tests {
assert!(payload["metadata"]["candidate_id"].is_null());
assert!(payload["metadata"]["candidate_index"].is_null());
assert!(payload["metadata"]["key_name"].is_null());
assert!(payload["metadata"]["model_id"].is_null());
assert!(payload["metadata"]["global_model_id"].is_null());
assert!(payload["metadata"]["global_model_name"].is_null());
assert!(payload["metadata"]["planner_kind"].is_null());
assert!(payload["metadata"]["trace_id"].is_null());
assert!(payload["metadata"]["route_family"].is_null());

207
crates/aether-billing/src/event_enrichment.rs
View File

@@ -4,10 +4,23 @@ use aether_usage_runtime::{UsageEvent, UsageEventType};
use async_trait::async_trait;
use serde_json::{Map, Value};
use crate::{BillingModelPricingSnapshot, BillingService, BillingUsageInput};
use crate::{
BillingComputation, BillingModelPricingSnapshot, BillingService, BillingSnapshotStatus,
BillingUsageInput,
};
#[async_trait]
pub trait BillingModelContextLookup: Send + Sync {
async fn find_billing_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
let _ = (provider_id, provider_api_key_id, model_id);
Ok(None)
}
async fn find_billing_model_context(
&self,
provider_id: &str,
@@ -35,23 +48,78 @@ pub async fn enrich_usage_event_with_billing(
else {
return Ok(());
};
let model_name = event.data.model.trim();
if model_name.is_empty() {
if let Some(model_id) = event
.data
.model_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
if let Some(context) = data
.find_billing_model_context_by_model_id(
provider_id,
event.data.provider_api_key_id.as_deref(),
model_id,
)
.await?
{
let pricing = map_pricing_context(context);
let computation = calculate_billing_computation(&pricing, event)?;
apply_billing_computation(event, computation)?;
return Ok(());
}
}
let mut first_no_rule = None;
for lookup_name in billing_model_lookup_names(&event.data) {
let Some(context) = data
.find_billing_model_context(
provider_id,
event.data.provider_api_key_id.as_deref(),
lookup_name,
)
.await?
else {
continue;
};
let pricing = map_pricing_context(context);
let computation = calculate_billing_computation(&pricing, event)?;
if matches!(
computation.cost_result.status,
BillingSnapshotStatus::NoRule
) {
first_no_rule.get_or_insert(computation);
continue;
}
apply_billing_computation(event, computation)?;
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(());
};
if let Some(computation) = first_no_rule {
apply_billing_computation(event, computation)?;
}
Ok(())
}
let pricing = map_pricing_context(context);
fn billing_model_lookup_names(data: &aether_usage_runtime::UsageEventData) -> Vec<&str> {
let mut names = Vec::new();
for value in [data.target_model.as_deref(), Some(data.model.as_str())]
.into_iter()
.flatten()
{
let value = value.trim();
if !value.is_empty() && !names.contains(&value) {
names.push(value);
}
}
names
}
fn calculate_billing_computation(
pricing: &BillingModelPricingSnapshot,
event: &UsageEvent,
) -> Result<BillingComputation, DataLayerError> {
let input = BillingUsageInput {
task_type: event
.data
@@ -85,11 +153,17 @@ pub async fn enrich_usage_event_with_billing(
cache_ttl_minutes: pricing.provider_api_key_cache_ttl_minutes,
};
let computation = BillingService::new()
.calculate(&pricing, &input)
BillingService::new()
.calculate(pricing, &input)
.map_err(|err| {
DataLayerError::UnexpectedValue(format!("billing calculation failed: {err}"))
})?;
})
}
fn apply_billing_computation(
event: &mut UsageEvent,
computation: BillingComputation,
) -> Result<(), DataLayerError> {
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(
@@ -97,8 +171,7 @@ pub async fn enrich_usage_event_with_billing(
&computation.cost_result.snapshot,
computation.rate_multiplier,
computation.is_free_tier,
)?;
Ok(())
)
}
fn map_pricing_context(context: StoredBillingModelContext) -> BillingModelPricingSnapshot {
@@ -153,11 +226,22 @@ mod tests {
use super::{enrich_usage_event_with_billing, BillingModelContextLookup};
struct TestLookup {
context: Option<StoredBillingModelContext>,
name_context: Option<StoredBillingModelContext>,
model_id_context: Option<StoredBillingModelContext>,
}
#[async_trait]
impl BillingModelContextLookup for TestLookup {
async fn find_billing_model_context_by_model_id(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_model_id: &str,
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
{
Ok(self.model_id_context.clone())
}
async fn find_billing_model_context(
&self,
_provider_id: &str,
@@ -165,14 +249,14 @@ mod tests {
_global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, aether_data_contracts::DataLayerError>
{
Ok(self.context.clone())
Ok(self.name_context.clone())
}
}
#[tokio::test]
async fn enriches_completed_usage_event_with_billing_snapshot() {
let lookup = TestLookup {
context: Some(
name_context: Some(
StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
@@ -192,6 +276,7 @@ mod tests {
)
.expect("billing context should build"),
),
model_id_context: None,
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
@@ -229,4 +314,82 @@ mod tests {
Some("complete")
);
}
#[tokio::test]
async fn enriches_by_provider_model_id_before_name_fallback() {
let blank_name_context = StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
Some(60),
"global-model-blank".to_string(),
"claude-sonnet-4-6".to_string(),
None,
None,
None,
Some("model-blank".to_string()),
Some("claude-sonnet-4-6".to_string()),
None,
None,
None,
)
.expect("blank billing context should build");
let priced_model_context = StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
Some(60),
"global-model-priced".to_string(),
"claude-sonnet-4-6".to_string(),
None,
None,
None,
Some("model-priced".to_string()),
Some("claude-sonnet-4-6".to_string()),
None,
None,
Some(
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}),
),
)
.expect("priced billing context should build");
let lookup = TestLookup {
name_context: Some(blank_name_context),
model_id_context: Some(priced_model_context),
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-billing-model-id-1",
UsageEventData {
provider_name: "NekoCode".to_string(),
model: "claude-sonnet-4-6".to_string(),
model_id: Some("model-priced".to_string()),
provider_id: Some("provider-1".to_string()),
provider_api_key_id: Some("key-1".to_string()),
request_type: Some("chat".to_string()),
input_tokens: Some(1_000),
output_tokens: Some(500),
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_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")
);
}
}

62
crates/aether-billing/src/pricing.rs
View File

@@ -24,6 +24,7 @@ impl BillingModelPricingSnapshot {
pub fn effective_tiered_pricing(&self) -> Option<&Value> {
self.model_tiered_pricing
.as_ref()
.filter(|value| has_tiered_pricing_tiers(value))
.or(self.default_tiered_pricing.as_ref())
}
@@ -58,6 +59,67 @@ impl BillingModelPricingSnapshot {
}
}
fn has_tiered_pricing_tiers(value: &Value) -> bool {
value
.get("tiers")
.and_then(Value::as_array)
.is_some_and(|tiers| !tiers.is_empty())
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::BillingModelPricingSnapshot;
fn snapshot(
model_tiered_pricing: Option<serde_json::Value>,
default_tiered_pricing: Option<serde_json::Value>,
) -> BillingModelPricingSnapshot {
BillingModelPricingSnapshot {
provider_id: "provider-1".to_string(),
provider_billing_type: None,
provider_api_key_id: None,
provider_api_key_rate_multipliers: None,
provider_api_key_cache_ttl_minutes: None,
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(),
global_model_config: None,
default_price_per_request: None,
default_tiered_pricing,
model_id: Some("model-1".to_string()),
model_provider_model_name: Some("gpt-5-upstream".to_string()),
model_config: None,
model_price_per_request: None,
model_tiered_pricing,
}
}
#[test]
fn empty_provider_tiered_pricing_inherits_global_default() {
let default_pricing =
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
let pricing = snapshot(Some(json!({})), Some(default_pricing.clone()));
assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing));
let pricing = snapshot(Some(json!({"tiers": []})), Some(default_pricing.clone()));
assert_eq!(pricing.effective_tiered_pricing(), Some(&default_pricing));
}
#[test]
fn populated_provider_tiered_pricing_overrides_global_default() {
let provider_pricing =
json!({"tiers":[{"up_to":null,"input_price_per_1m":1.0,"output_price_per_1m":2.0}]});
let default_pricing =
json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]});
let pricing = snapshot(Some(provider_pricing.clone()), Some(default_pricing));
assert_eq!(pricing.effective_tiered_pricing(), Some(&provider_pricing));
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BillingUsageInput {
pub task_type: String,

10
crates/aether-data-contracts/src/repository/billing/types.rs
View File

@@ -150,4 +150,14 @@ pub trait BillingReadRepository: Send + Sync {
provider_api_key_id: Option<&str>,
global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError>;
async fn find_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, crate::DataLayerError> {
let _ = (provider_id, provider_api_key_id, model_id);
Ok(None)
}
}

12
crates/aether-data-contracts/src/repository/usage/types.rs
View File

@@ -500,6 +500,18 @@ impl StoredRequestUsageAudit {
.or_else(|| self.request_metadata_string("key_name"))
}
pub fn routing_model_id(&self) -> Option<&str> {
self.request_metadata_string("model_id")
}
pub fn routing_global_model_id(&self) -> Option<&str> {
self.request_metadata_string("global_model_id")
}
pub fn routing_global_model_name(&self) -> Option<&str> {
self.request_metadata_string("global_model_name")
}
pub fn routing_planner_kind(&self) -> Option<&str> {
self.planner_kind
.as_deref()

166
crates/aether-data/src/repository/billing/memory.rs
View File

@@ -44,12 +44,21 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
provider_api_key_id: Option<&str>,
global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
let by_key = self.by_key.read().expect("billing repository lock");
if let Some(value) = find_context_by_provider_model_name(
&by_key,
provider_id,
provider_api_key_id,
global_model_name,
) {
return Ok(Some(value));
}
let key = (
provider_id.to_string(),
global_model_name.to_string(),
provider_api_key_id.map(ToOwned::to_owned),
);
let by_key = self.by_key.read().expect("billing repository lock");
if let Some(value) = by_key.get(&key) {
return Ok(Some(value.clone()));
}
@@ -68,6 +77,88 @@ impl BillingReadRepository for InMemoryBillingReadRepository {
})
.map(|(_, value)| value.clone()))
}
async fn find_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
let by_key = self.by_key.read().expect("billing repository lock");
if let Some(value) =
find_context_by_model_id_and_key(&by_key, provider_id, provider_api_key_id, model_id)
{
return Ok(Some(value));
}
if let Some(value) = find_context_by_model_id_and_key(&by_key, provider_id, None, model_id)
{
return Ok(Some(value));
}
Ok(by_key
.iter()
.find(|((stored_provider_id, _, _), value)| {
stored_provider_id == provider_id && value.model_id.as_deref() == Some(model_id)
})
.map(|(_, value)| value.clone()))
}
}
fn find_context_by_provider_model_name(
by_key: &BillingContextMap,
provider_id: &str,
provider_api_key_id: Option<&str>,
provider_model_name: &str,
) -> Option<StoredBillingModelContext> {
find_context_by_provider_model_name_and_key(
by_key,
provider_id,
provider_api_key_id,
provider_model_name,
)
.or_else(|| {
find_context_by_provider_model_name_and_key(by_key, provider_id, None, provider_model_name)
})
.or_else(|| {
by_key
.iter()
.find(|((stored_provider_id, _, _), value)| {
stored_provider_id == provider_id
&& value.model_provider_model_name.as_deref() == Some(provider_model_name)
})
.map(|(_, value)| value.clone())
})
}
fn find_context_by_provider_model_name_and_key(
by_key: &BillingContextMap,
provider_id: &str,
provider_api_key_id: Option<&str>,
provider_model_name: &str,
) -> Option<StoredBillingModelContext> {
by_key
.iter()
.find(|((stored_provider_id, _, stored_key_id), value)| {
stored_provider_id == provider_id
&& stored_key_id.as_deref() == provider_api_key_id
&& value.model_provider_model_name.as_deref() == Some(provider_model_name)
})
.map(|(_, value)| value.clone())
}
fn find_context_by_model_id_and_key(
by_key: &BillingContextMap,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Option<StoredBillingModelContext> {
by_key
.iter()
.find(|((stored_provider_id, _, stored_key_id), value)| {
stored_provider_id == provider_id
&& stored_key_id.as_deref() == provider_api_key_id
&& value.model_id.as_deref() == Some(model_id)
})
.map(|(_, value)| value.clone())
}
#[cfg(test)]
@@ -110,4 +201,77 @@ mod tests {
assert_eq!(stored.provider_id, "provider-1");
assert_eq!(stored.global_model_name, "gpt-5");
}
#[tokio::test]
async fn resolves_by_provider_model_name_before_global_name_collision() {
let global_named_context = StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
Some(60),
"global-model-blank".to_string(),
"claude-sonnet-4-6".to_string(),
None,
None,
None,
Some("model-blank".to_string()),
Some("blank-upstream".to_string()),
None,
None,
None,
)
.expect("blank billing context should build");
let provider_priced_context = StoredBillingModelContext::new(
"provider-1".to_string(),
Some("pay_as_you_go".to_string()),
Some("key-1".to_string()),
None,
Some(60),
"global-model-priced".to_string(),
"claude-opus-4-6".to_string(),
None,
None,
Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]})),
Some("model-priced".to_string()),
Some("claude-sonnet-4-6".to_string()),
None,
None,
None,
)
.expect("priced billing context should build");
let repository = InMemoryBillingReadRepository::seed(vec![
global_named_context,
provider_priced_context,
]);
let stored = repository
.find_model_context("provider-1", Some("key-1"), "claude-sonnet-4-6")
.await
.expect("lookup should succeed")
.expect("context should exist");
assert_eq!(stored.global_model_name, "claude-opus-4-6");
assert_eq!(
stored.model_provider_model_name.as_deref(),
Some("claude-sonnet-4-6")
);
assert!(stored.default_tiered_pricing.is_some());
}
#[tokio::test]
async fn resolves_by_model_id() {
let repository = InMemoryBillingReadRepository::seed(vec![sample_context()]);
let stored = repository
.find_model_context_by_model_id("provider-1", Some("key-1"), "model-1")
.await
.expect("lookup should succeed")
.expect("context should exist");
assert_eq!(stored.global_model_name, "gpt-5");
assert_eq!(
stored.model_provider_model_name.as_deref(),
Some("gpt-5-upstream")
);
}
}

94
crates/aether-data/src/repository/billing/sql.rs
View File

@@ -23,8 +23,7 @@ SELECT
m.tiered_pricing AS model_tiered_pricing
FROM providers p
INNER JOIN global_models gm
ON gm.name = $2
AND gm.is_active = TRUE
ON gm.is_active = TRUE
LEFT JOIN models m
ON m.global_model_id = gm.id
AND m.provider_id = p.id
@@ -33,7 +32,71 @@ LEFT JOIN provider_api_keys pak
ON pak.id = $3
AND pak.provider_id = p.id
WHERE p.id = $1
ORDER BY COALESCE(m.is_available, FALSE) DESC, m.created_at ASC
AND (
gm.name = $2
OR m.provider_model_name = $2
OR (
m.provider_model_mappings IS NOT NULL
AND (
m.provider_model_mappings @> jsonb_build_array(jsonb_build_object('name', $2::TEXT))
OR m.provider_model_mappings @> jsonb_build_array(to_jsonb($2::TEXT))
OR m.provider_model_mappings @> jsonb_build_object('name', $2::TEXT)
OR m.provider_model_mappings = to_jsonb($2::TEXT)
)
)
)
ORDER BY
CASE
WHEN m.provider_model_name = $2 THEN 0
WHEN m.provider_model_mappings IS NOT NULL
AND (
m.provider_model_mappings @> jsonb_build_array(jsonb_build_object('name', $2::TEXT))
OR m.provider_model_mappings @> jsonb_build_array(to_jsonb($2::TEXT))
OR m.provider_model_mappings @> jsonb_build_object('name', $2::TEXT)
OR m.provider_model_mappings = to_jsonb($2::TEXT)
) THEN 1
WHEN gm.name = $2 THEN 2
ELSE 3
END ASC,
COALESCE(m.is_available, FALSE) DESC,
CASE
WHEN m.tiered_pricing IS NOT NULL OR m.price_per_request IS NOT NULL THEN 0
WHEN gm.default_tiered_pricing IS NOT NULL OR gm.default_price_per_request IS NOT NULL THEN 1
ELSE 2
END ASC,
m.created_at ASC
LIMIT 1
"#;
const FIND_MODEL_CONTEXT_BY_MODEL_ID_SQL: &str = r#"
SELECT
p.id AS provider_id,
CAST(p.billing_type AS TEXT) AS provider_billing_type,
pak.id AS provider_api_key_id,
pak.rate_multipliers AS provider_api_key_rate_multipliers,
pak.cache_ttl_minutes AS provider_api_key_cache_ttl_minutes,
gm.id AS global_model_id,
gm.name AS global_model_name,
gm.config AS global_model_config,
CAST(gm.default_price_per_request AS DOUBLE PRECISION) AS default_price_per_request,
gm.default_tiered_pricing AS default_tiered_pricing,
m.id AS model_id,
m.provider_model_name AS model_provider_model_name,
m.config AS model_config,
CAST(m.price_per_request AS DOUBLE PRECISION) AS model_price_per_request,
m.tiered_pricing AS model_tiered_pricing
FROM providers p
INNER JOIN models m
ON m.id = $2
AND m.provider_id = p.id
AND m.is_active = TRUE
INNER JOIN global_models gm
ON gm.id = m.global_model_id
AND gm.is_active = TRUE
LEFT JOIN provider_api_keys pak
ON pak.id = $3
AND pak.provider_id = p.id
WHERE p.id = $1
LIMIT 1
"#;
@@ -62,6 +125,22 @@ impl SqlxBillingReadRepository {
.map_postgres_err()?;
row.as_ref().map(map_row).transpose()
}
pub async fn find_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
let row = sqlx::query(FIND_MODEL_CONTEXT_BY_MODEL_ID_SQL)
.bind(provider_id)
.bind(model_id)
.bind(provider_api_key_id)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
row.as_ref().map(map_row).transpose()
}
}
#[async_trait]
@@ -74,6 +153,15 @@ impl BillingReadRepository for SqlxBillingReadRepository {
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
Self::find_model_context(self, provider_id, provider_api_key_id, global_model_name).await
}
async fn find_model_context_by_model_id(
&self,
provider_id: &str,
provider_api_key_id: Option<&str>,
model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
Self::find_model_context_by_model_id(self, provider_id, provider_api_key_id, model_id).await
}
}
fn map_row(row: &sqlx::postgres::PgRow) -> Result<StoredBillingModelContext, DataLayerError> {

4
crates/aether-usage-runtime/src/event.rs
View File

@@ -33,6 +33,10 @@ pub struct UsageEventData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub target_model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub global_model_id: 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>,

18
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -75,6 +75,9 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_non_null_value(source, target, "billing_snapshot");
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
copy_non_empty_string(source, target, "billing_snapshot_status");
copy_non_empty_string(source, target, "model_id");
copy_non_empty_string(source, target, "global_model_id");
copy_non_empty_string(source, target, "global_model_name");
copy_non_null_value(source, target, "dimensions");
copy_non_null_value(source, target, "billing_rule_snapshot");
copy_non_null_value(source, target, "scheduling_audit");
@@ -97,6 +100,9 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
remove_non_null_value(&mut source, target, "billing_snapshot");
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
remove_non_empty_string(&mut source, target, "model_id");
remove_non_empty_string(&mut source, target, "global_model_id");
remove_non_empty_string(&mut source, target, "global_model_name");
remove_non_null_value(&mut source, target, "dimensions");
remove_non_null_value(&mut source, target, "billing_rule_snapshot");
remove_non_null_value(&mut source, target, "scheduling_audit");
@@ -378,6 +384,9 @@ mod tests {
"billing_snapshot": {"status": "complete"},
"billing_snapshot_schema_version": "2.0",
"billing_snapshot_status": "complete",
"model_id": "model-1",
"global_model_id": "global-model-1",
"global_model_name": "gpt-5",
"dimensions": {"total_input_context": 10},
"rate_multiplier": 1.25,
"is_free_tier": false,
@@ -405,6 +414,9 @@ mod tests {
"billing_snapshot": {"status": "complete"},
"billing_snapshot_schema_version": "2.0",
"billing_snapshot_status": "complete",
"model_id": "model-1",
"global_model_id": "global-model-1",
"global_model_name": "gpt-5",
"dimensions": {"total_input_context": 10},
"rate_multiplier": 1.25,
"is_free_tier": false,
@@ -455,6 +467,9 @@ mod tests {
"client_requested_stream": false,
"upstream_is_stream": true,
"provider_id": "provider-1",
"model_id": "model-1",
"global_model_id": "global-model-1",
"global_model_name": "gpt-5",
"billing_snapshot": {"status": "complete"}
})
.as_object()
@@ -468,6 +483,9 @@ mod tests {
json!({
"client_requested_stream": false,
"upstream_is_stream": true,
"model_id": "model-1",
"global_model_id": "global-model-1",
"global_model_name": "gpt-5",
"billing_snapshot": {"status": "complete"}
})
);

24
crates/aether-usage-runtime/src/write.rs
View File

@@ -74,6 +74,8 @@ pub struct LifecycleUsageSeed {
pub provider_name: String,
pub model: String,
pub target_model: Option<String>,
pub model_id: Option<String>,
pub global_model_id: Option<String>,
pub provider_id: Option<String>,
pub provider_endpoint_id: Option<String>,
pub provider_api_key_id: Option<String>,
@@ -110,6 +112,8 @@ pub struct TerminalUsageContextSeed {
pub provider_name: String,
pub model: String,
pub target_model: Option<String>,
pub model_id: Option<String>,
pub global_model_id: Option<String>,
pub provider_id: Option<String>,
pub provider_endpoint_id: Option<String>,
pub provider_api_key_id: Option<String>,
@@ -168,6 +172,8 @@ pub struct TerminalUsageSeed {
pub provider_name: String,
pub model: String,
pub target_model: Option<String>,
pub model_id: Option<String>,
pub global_model_id: Option<String>,
pub provider_id: Option<String>,
pub provider_endpoint_id: Option<String>,
pub provider_api_key_id: Option<String>,
@@ -262,6 +268,8 @@ pub fn build_lifecycle_usage_seed(
provider_name,
model,
target_model: context_string(context, "mapped_model"),
model_id: context_string(context, "model_id"),
global_model_id: context_string(context, "global_model_id"),
provider_id: empty_to_none(
context_string(context, "provider_id")
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
@@ -486,6 +494,8 @@ fn build_terminal_usage_event_from_seed_impl(
provider_name,
model,
target_model,
model_id,
global_model_id,
provider_id,
provider_endpoint_id,
provider_api_key_id,
@@ -548,6 +558,8 @@ fn build_terminal_usage_event_from_seed_impl(
provider_name,
model,
target_model,
model_id,
global_model_id,
provider_id,
provider_endpoint_id,
provider_api_key_id,
@@ -657,6 +669,8 @@ pub fn build_terminal_usage_context_seed(
.or_else(|| non_empty_str(plan.model_name.as_deref()))
.unwrap_or_else(|| "unknown".to_string()),
target_model: context_string(context, "mapped_model"),
model_id: context_string(context, "model_id"),
global_model_id: context_string(context, "global_model_id"),
provider_id: context_string(context, "provider_id")
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
provider_endpoint_id: context_string(context, "endpoint_id")
@@ -813,6 +827,8 @@ pub fn build_sync_terminal_usage_seed(
provider_name: context_seed.provider_name,
model: context_seed.model,
target_model: context_seed.target_model,
model_id: context_seed.model_id,
global_model_id: context_seed.global_model_id,
provider_id: context_seed.provider_id,
provider_endpoint_id: context_seed.provider_endpoint_id,
provider_api_key_id: context_seed.provider_api_key_id,
@@ -881,6 +897,8 @@ pub fn build_stream_terminal_usage_seed(
provider_name: context_seed.provider_name,
model: context_seed.model,
target_model: context_seed.target_model,
model_id: context_seed.model_id,
global_model_id: context_seed.global_model_id,
provider_id: context_seed.provider_id,
provider_endpoint_id: context_seed.provider_endpoint_id,
provider_api_key_id: context_seed.provider_api_key_id,
@@ -1242,6 +1260,8 @@ fn build_usage_event_data_seed_with_detail(
provider_name,
model,
target_model: context_string(context, "mapped_model"),
model_id: context_string(context, "model_id"),
global_model_id: context_string(context, "global_model_id"),
provider_id: context_string(context, "provider_id")
.or_else(|| non_empty_str(Some(plan.provider_id.as_str()))),
provider_endpoint_id: context_string(context, "endpoint_id")
@@ -3523,6 +3543,8 @@ mod tests {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
target_model: None,
model_id: None,
global_model_id: None,
provider_id: Some("provider-1".to_string()),
provider_endpoint_id: Some("endpoint-1".to_string()),
provider_api_key_id: Some("upstream-key-1".to_string()),
@@ -3646,6 +3668,8 @@ mod tests {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
target_model: None,
model_id: None,
global_model_id: None,
provider_id: Some("provider-1".to_string()),
provider_endpoint_id: Some("endpoint-1".to_string()),
provider_api_key_id: Some("upstream-key-1".to_string()),