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

fix provider pool quota status handling

Browse Source
This commit is contained in:
fawney19
2026-05-07 11:28:44 +08:00
parent 01c9f9be49
commit fd44906bb7
37 changed files with 1844 additions and 500 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
crates/aether-admin/src/provider/pool.rs
View File

@@ -395,17 +395,38 @@ pub fn admin_pool_is_oauth_invalid(key: &StoredProviderCatalogKey, now_unix_secs
if key.auth_type.trim() != "oauth" {
return false;
}
if key
.oauth_invalid_reason
.as_deref()
.is_some_and(|value| !value.trim().is_empty())
{
return true;
if let Some(reason) = key.oauth_invalid_reason.as_deref().map(str::trim) {
let account_state = provider_status::resolve_pool_account_state(
None,
key.upstream_metadata.as_ref(),
Some(reason),
);
if account_state.blocked && !account_state.recoverable {
return true;
}
if admin_pool_reason_has_tag(reason, "[REFRESH_FAILED]") {
return key
.expires_at_unix_secs
.is_none_or(|value| value == 0 || value <= now_unix_secs);
}
if admin_pool_reason_has_tag(reason, "[REQUEST_FAILED]") {
return false;
}
if !reason.is_empty() {
return true;
}
}
key.expires_at_unix_secs
.is_some_and(|value| value > 0 && value <= now_unix_secs)
}
fn admin_pool_reason_has_tag(reason: &str, tag: &str) -> bool {
reason
.lines()
.map(str::trim)
.any(|line| line.starts_with(tag))
}
pub fn admin_pool_matches_quick_selector(
key: &StoredProviderCatalogKey,
selector: &str,

76
crates/aether-admin/src/provider/quota.rs
View File

@@ -396,7 +396,8 @@ fn codex_merge_invalid_reason(current: &str, candidate_reason: &str) -> String {
return current.to_string();
}
if current.starts_with(OAUTH_EXPIRED_PREFIX)
&& candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
&& (candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
|| candidate_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
{
return current.to_string();
}
@@ -872,9 +873,11 @@ pub fn parse_chatgpt_web_conversation_init_response(
#[cfg(test)]
mod tests {
use super::{
codex_runtime_invalid_reason, parse_chatgpt_web_conversation_init_response,
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
codex_build_invalid_state, codex_runtime_invalid_reason,
parse_chatgpt_web_conversation_init_response, OAUTH_ACCOUNT_BLOCK_PREFIX,
OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
#[test]
@@ -900,6 +903,73 @@ mod tests {
assert_eq!(codex_runtime_invalid_reason(403, Some("forbidden")), None);
}
#[test]
fn codex_invalid_state_keeps_oauth_expired_over_refresh_failure() {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.oauth_invalid_at_unix_secs = Some(100);
key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
assert_eq!(
codex_build_invalid_state(
&key,
format!("{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"),
200,
),
(
Some(100),
Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
)
);
assert_eq!(
codex_build_invalid_state(
&key,
format!("{OAUTH_REQUEST_FAILED_PREFIX}账号状态检查失败"),
200,
),
(
Some(100),
Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
)
);
}
#[test]
fn codex_invalid_state_allows_account_block_to_override_oauth_expired() {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key-1".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.oauth_invalid_at_unix_secs = Some(100);
key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
assert_eq!(
codex_build_invalid_state(
&key,
format!("{OAUTH_ACCOUNT_BLOCK_PREFIX}account has been deactivated"),
200,
),
(
Some(200),
Some(format!(
"{OAUTH_ACCOUNT_BLOCK_PREFIX}account has been deactivated"
))
)
);
}
#[test]
fn parses_chatgpt_web_image_quota_from_conversation_init() {
let parsed = parse_chatgpt_web_conversation_init_response(

88
crates/aether-admin/src/provider/status.rs
View File

@@ -1,11 +1,6 @@
use serde_json::Value;
use std::collections::BTreeSet;
const OAUTH_ACCOUNT_BLOCK_PREFIX: &str = "[ACCOUNT_BLOCK] ";
const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] ";
const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] ";
const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] ";
const ACCOUNT_BLOCK_REASON_KEYWORDS: &[&str] = &[
"suspended",
"banned",
@@ -326,55 +321,55 @@ fn resolve_from_metadata(
fn resolve_from_oauth_invalid_reason(reason: Option<&str>) -> Option<PoolAccountState> {
let text = clean_text(reason)?;
if text.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX) {
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_ACCOUNT_BLOCK_PREFIX)))
.unwrap_or_else(|| "账号异常".to_string());
let (code, label) = classify_block_reason(&cleaned);
if let Some(cleaned) = tagged_reason(&text, "ACCOUNT_BLOCK") {
let reason = if cleaned.is_empty() {
"账号异常".to_string()
} else {
cleaned
};
let (code, label) = classify_block_reason(&reason);
return Some(PoolAccountState {
blocked: true,
code: Some(code.to_string()),
label: Some(label.to_string()),
reason: Some(cleaned),
reason: Some(reason),
source: Some("oauth_invalid".to_string()),
recoverable: false,
});
}
if text.starts_with(OAUTH_EXPIRED_PREFIX) {
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_EXPIRED_PREFIX)))
.unwrap_or_else(|| "OAuth Token 已过期且无法续期".to_string());
if let Some(cleaned) = tagged_reason(&text, "OAUTH_EXPIRED") {
let reason = if cleaned.is_empty() {
"OAuth Token 已过期且无法续期".to_string()
} else {
cleaned
};
return Some(PoolAccountState {
blocked: true,
code: Some("oauth_token_invalid".to_string()),
label: Some("Token 失效".to_string()),
reason: Some(cleaned),
reason: Some(reason),
source: Some("oauth_invalid".to_string()),
recoverable: false,
});
}
if text.starts_with(OAUTH_REFRESH_FAILED_PREFIX) {
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_REFRESH_FAILED_PREFIX)))
.unwrap_or_else(|| "OAuth Token 续期失败".to_string());
return Some(PoolAccountState {
blocked: true,
code: Some("oauth_token_invalid".to_string()),
label: Some("Token 失效".to_string()),
reason: Some(cleaned),
source: Some("oauth_refresh".to_string()),
recoverable: false,
});
}
if text.starts_with(OAUTH_REQUEST_FAILED_PREFIX) {
let cleaned = clean_text(Some(text.trim_start_matches(OAUTH_REQUEST_FAILED_PREFIX)))
.unwrap_or_else(|| "账号状态检查失败".to_string());
if let Some(cleaned) = tagged_reason(&text, "REQUEST_FAILED") {
let reason = if cleaned.is_empty() {
"账号状态检查失败".to_string()
} else {
cleaned
};
return Some(PoolAccountState {
blocked: false,
code: Some("oauth_request_failed".to_string()),
label: Some("请求失败".to_string()),
reason: Some(cleaned),
reason: Some(reason),
source: Some("oauth_request".to_string()),
recoverable: true,
});
}
if tagged_reason(&text, "REFRESH_FAILED").is_some() {
return None;
}
if text.starts_with('[') {
return None;
}
@@ -462,20 +457,8 @@ pub fn resolve_account_status_snapshot(
};
}
if let Some(cleaned) = tagged_reason(&text, "REFRESH_FAILED") {
let reason = if cleaned.is_empty() {
"OAuth Token 续期失败".to_string()
} else {
cleaned
};
return AccountStatusSnapshot {
code: "oauth_token_invalid".to_string(),
label: Some("Token 失效".to_string()),
reason: Some(reason),
blocked: true,
source: Some("oauth_refresh".to_string()),
recoverable: false,
};
if tagged_reason(&text, "REFRESH_FAILED").is_some() {
return AccountStatusSnapshot::default();
}
if text.starts_with('[') {
@@ -578,31 +561,30 @@ mod tests {
}
#[test]
fn resolves_refresh_failed_as_token_invalid_pool_state() {
fn ignores_refresh_failed_as_pool_account_block() {
let state = resolve_pool_account_state(
Some("codex"),
None,
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已失效"),
);
assert!(state.blocked);
assert!(!state.recoverable);
assert_eq!(state.code.as_deref(), Some("oauth_token_invalid"));
assert_eq!(state.label.as_deref(), Some("Token 失效"));
assert!(!state.blocked);
assert_eq!(state.code.as_deref(), None);
assert_eq!(state.label.as_deref(), None);
assert!(!should_auto_remove_account_state(&state));
}
#[test]
fn account_snapshot_marks_refresh_failed_as_token_invalid() {
fn account_snapshot_ignores_refresh_failed_as_account_block() {
let snapshot = resolve_account_status_snapshot(
Some("codex"),
None,
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 已失效"),
);
assert_eq!(snapshot.code, "oauth_token_invalid");
assert_eq!(snapshot.label.as_deref(), Some("Token 失效"));
assert!(snapshot.blocked);
assert_eq!(snapshot.code, "ok");
assert_eq!(snapshot.label.as_deref(), None);
assert!(!snapshot.blocked);
assert!(!snapshot.recoverable);
}

13
crates/aether-data/migrations/postgres/20260505130000_project_upstream_stream_in_usage_billing_facts.sql
View File

@@ -1,19 +1,6 @@
ALTER TABLE IF EXISTS public.usage
ADD COLUMN IF NOT EXISTS upstream_is_stream boolean;
UPDATE public.usage
SET upstream_is_stream = COALESCE(
CASE
WHEN (request_metadata->>'upstream_is_stream') IN ('true', 'false')
THEN (request_metadata->>'upstream_is_stream')::boolean
ELSE NULL
END,
COALESCE(is_stream, FALSE)
)
WHERE upstream_is_stream IS NULL;
ANALYZE public.usage;
COMMENT ON COLUMN public.usage.upstream_is_stream IS
'Resolved upstream stream mode from request_metadata.upstream_is_stream, falling back to is_stream for legacy rows.';

222
crates/aether-data/src/repository/provider_catalog/memory.rs
View File

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use serde_json::{json, Map, Value};
use super::{
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogReadRepository,
@@ -88,6 +89,8 @@ impl InMemoryProviderCatalogReadRepository {
{
key.last_used_at_unix_secs = recomputed_last_used_at_unix_secs;
}
apply_codex_window_usage_stats_delta(&mut key.status_snapshot, delta);
}
pub(crate) fn rebuild_usage_stats(
@@ -156,6 +159,156 @@ fn apply_f64_delta(current: f64, delta: f64) -> f64 {
}
}
fn json_u64(value: Option<&Value>) -> Option<u64> {
value.and_then(|value| {
value.as_u64().or_else(|| {
value
.as_str()
.and_then(|text| text.trim().parse::<u64>().ok())
})
})
}
fn json_i64(value: Option<&Value>) -> Option<i64> {
value.and_then(|value| {
value.as_i64().or_else(|| {
value
.as_str()
.and_then(|text| text.trim().parse::<i64>().ok())
})
})
}
fn json_f64(value: Option<&Value>) -> Option<f64> {
value.and_then(|value| {
value.as_f64().or_else(|| {
value
.as_str()
.and_then(|text| text.trim().parse::<f64>().ok())
})
})
}
fn apply_i64_delta_to_json_u64(current: u64, delta: i64) -> u64 {
if delta >= 0 {
current.saturating_add(delta as u64)
} else {
current.saturating_sub(delta.unsigned_abs())
}
}
fn apply_f64_delta_to_json_cost(current: f64, delta: f64) -> f64 {
let current = if current.is_finite() { current } else { 0.0 };
let delta = if delta.is_finite() { delta } else { 0.0 };
(current + delta).max(0.0)
}
fn codex_window_matches_usage_time(window: &Map<String, Value>, usage_created_at: u64) -> bool {
let code = window
.get("code")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if !code.eq_ignore_ascii_case("5h") && !code.eq_ignore_ascii_case("weekly") {
return false;
}
let Some(reset_at) = json_u64(window.get("reset_at")) else {
return false;
};
let Some(window_minutes) = json_u64(window.get("window_minutes")) else {
return false;
};
let Some(window_seconds) = window_minutes.checked_mul(60) else {
return false;
};
let Some(window_start) = reset_at.checked_sub(window_seconds) else {
return false;
};
let usage_reset_at = json_u64(window.get("usage_reset_at")).unwrap_or(0);
let start = window_start.max(usage_reset_at);
usage_created_at >= start && usage_created_at < reset_at
}
fn apply_codex_window_usage_stats_delta(
status_snapshot: &mut Option<Value>,
delta: &ProviderApiKeyUsageDelta,
) {
let Some(usage_created_at) = delta.usage_created_at_unix_secs else {
return;
};
if delta.request_count == 0 && delta.total_tokens == 0 && delta.total_cost_usd == 0.0 {
return;
}
let Some(quota) = status_snapshot
.as_mut()
.and_then(Value::as_object_mut)
.and_then(|snapshot| snapshot.get_mut("quota"))
.and_then(Value::as_object_mut)
else {
return;
};
let quota_provider_type = quota
.get("provider_type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if !quota_provider_type.eq_ignore_ascii_case("codex") {
return;
}
let Some(windows) = quota.get_mut("windows").and_then(Value::as_array_mut) else {
return;
};
for window in windows.iter_mut().filter_map(Value::as_object_mut) {
if !codex_window_matches_usage_time(window, usage_created_at) {
continue;
}
let usage = window
.entry("usage".to_string())
.or_insert_with(|| json!({}))
.as_object_mut();
let Some(usage) = usage else {
window.insert("usage".to_string(), json!({}));
let Some(usage) = window.get_mut("usage").and_then(Value::as_object_mut) else {
continue;
};
let request_count = apply_i64_delta_to_json_u64(0, delta.request_count);
let total_tokens = apply_i64_delta_to_json_u64(0, delta.total_tokens);
let total_cost_usd = apply_f64_delta_to_json_cost(0.0, delta.total_cost_usd);
usage.insert("request_count".to_string(), json!(request_count));
usage.insert("total_tokens".to_string(), json!(total_tokens));
usage.insert(
"total_cost_usd".to_string(),
json!(format!("{total_cost_usd:.8}")),
);
continue;
};
let request_count = apply_i64_delta_to_json_u64(
json_i64(usage.get("request_count")).unwrap_or(0).max(0) as u64,
delta.request_count,
);
let total_tokens = apply_i64_delta_to_json_u64(
json_i64(usage.get("total_tokens")).unwrap_or(0).max(0) as u64,
delta.total_tokens,
);
let total_cost_usd = apply_f64_delta_to_json_cost(
json_f64(usage.get("total_cost_usd")).unwrap_or(0.0),
delta.total_cost_usd,
);
usage.insert("request_count".to_string(), json!(request_count));
usage.insert("total_tokens".to_string(), json!(total_tokens));
usage.insert(
"total_cost_usd".to_string(),
json!(format!("{total_cost_usd:.8}")),
);
}
}
#[async_trait]
impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
async fn list_providers(
@@ -573,6 +726,8 @@ mod tests {
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use crate::repository::usage::ProviderApiKeyUsageDelta;
use serde_json::json;
fn sample_provider(id: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
@@ -716,6 +871,73 @@ mod tests {
assert_eq!(stored[0].expires_at_unix_secs, Some(4_102_444_800));
}
#[tokio::test]
async fn materializes_codex_window_usage_stats_delta_in_memory() {
let mut key = sample_key("key-1", "provider-1");
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "codex",
"windows": [
{
"code": "5h",
"reset_at": 120_000u64,
"window_minutes": 300u64,
"usage": {
"request_count": 1,
"total_tokens": 10,
"total_cost_usd": "0.10000000"
}
},
{
"code": "weekly",
"reset_at": 700_000u64,
"window_minutes": 10_080u64
}
]
}
}));
let repository = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1")],
vec![],
vec![key],
);
repository.apply_usage_stats_delta(
"key-1",
&ProviderApiKeyUsageDelta {
request_count: 2,
total_tokens: 25,
total_cost_usd: 0.25,
usage_created_at_unix_secs: Some(110_000),
..ProviderApiKeyUsageDelta::default()
},
None,
);
let stored = repository
.list_keys_by_ids(&["key-1".to_string()])
.await
.expect("keys should read");
let windows = stored[0].status_snapshot.as_ref().expect("snapshot")["quota"]["windows"]
.as_array()
.expect("windows");
let five_h = windows
.iter()
.find(|window| window["code"] == json!("5h"))
.expect("5h window should exist");
let weekly = windows
.iter()
.find(|window| window["code"] == json!("weekly"))
.expect("weekly window should exist");
assert_eq!(five_h["usage"]["request_count"], json!(3));
assert_eq!(five_h["usage"]["total_tokens"], json!(35));
assert_eq!(five_h["usage"]["total_cost_usd"], json!("0.35000000"));
assert_eq!(weekly["usage"]["request_count"], json!(2));
assert_eq!(weekly["usage"]["total_tokens"], json!(25));
assert_eq!(weekly["usage"]["total_cost_usd"], json!("0.25000000"));
}
#[tokio::test]
async fn paginates_provider_keys_with_search_and_active_filter() {
let mut alpha = sample_key("key-1", "provider-1");

6
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -501,6 +501,7 @@ pub(crate) struct ProviderApiKeyUsageContribution {
pub total_cost_usd: f64,
pub total_response_time_ms: i64,
pub last_used_at_unix_secs: Option<u64>,
pub usage_created_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default)]
@@ -513,6 +514,7 @@ pub(crate) struct ProviderApiKeyUsageDelta {
pub total_response_time_ms: i64,
pub candidate_last_used_at_unix_secs: Option<u64>,
pub removed_last_used_at_unix_secs: Option<u64>,
pub usage_created_at_unix_secs: Option<u64>,
}
impl ProviderApiKeyUsageDelta {
@@ -529,6 +531,7 @@ impl ProviderApiKeyUsageDelta {
total_response_time_ms: after.total_response_time_ms - before.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
removed_last_used_at_unix_secs: None,
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
}
}
@@ -542,6 +545,7 @@ impl ProviderApiKeyUsageDelta {
total_response_time_ms: after.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
removed_last_used_at_unix_secs: None,
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
}
}
@@ -555,6 +559,7 @@ impl ProviderApiKeyUsageDelta {
total_response_time_ms: -before.total_response_time_ms,
candidate_last_used_at_unix_secs: None,
removed_last_used_at_unix_secs: before.last_used_at_unix_secs,
usage_created_at_unix_secs: before.usage_created_at_unix_secs,
}
}
@@ -664,6 +669,7 @@ pub(crate) fn provider_api_key_usage_contribution(
0
},
last_used_at_unix_secs: Some(usage.created_at_unix_ms),
usage_created_at_unix_secs: Some(usage.created_at_unix_ms),
})
}

106
crates/aether-data/src/repository/usage/postgres/mod.rs
View File

@@ -1333,15 +1333,20 @@ const REBUILD_API_KEY_USAGE_STATS_SQL: &str =
const APPLY_PROVIDER_API_KEY_USAGE_DELTA_SQL: &str =
include_str!("queries/apply_provider_api_key_usage_delta_sql.sql");
const APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL: &str =
include_str!("queries/apply_provider_api_key_codex_window_usage_delta_sql.sql");
const RESET_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
include_str!("queries/reset_provider_api_key_usage_stats_sql.sql");
const REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
include_str!("queries/rebuild_provider_api_key_usage_stats_sql.sql");
const REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL: &str =
include_str!("queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql");
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str =
" AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str = " AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
struct UsageAuditAggregationSqlFragments {
filtered_extra_where: &'static str,
@@ -6395,33 +6400,38 @@ WHERE stats_daily_api_key.date >=
return Ok(Vec::new());
}
let (table_name, group_column, display_name_expr, avg_response_time_expr, success_count_expr) =
match group_by {
UsageAuditAggregationGroupBy::Model => (
"stats_user_daily_model",
"model",
"NULL::varchar",
"NULL::DOUBLE PRECISION",
"NULL::BIGINT",
),
UsageAuditAggregationGroupBy::Provider => (
"stats_user_daily_provider",
"provider_name",
"provider_name",
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
"COALESCE(SUM(success_requests), 0)::BIGINT",
),
UsageAuditAggregationGroupBy::ApiFormat => (
"stats_user_daily_api_format",
"api_format",
"NULL::varchar",
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
"NULL::BIGINT",
),
UsageAuditAggregationGroupBy::User => {
return Ok(Vec::new());
}
};
let (
table_name,
group_column,
display_name_expr,
avg_response_time_expr,
success_count_expr,
) = match group_by {
UsageAuditAggregationGroupBy::Model => (
"stats_user_daily_model",
"model",
"NULL::varchar",
"NULL::DOUBLE PRECISION",
"NULL::BIGINT",
),
UsageAuditAggregationGroupBy::Provider => (
"stats_user_daily_provider",
"provider_name",
"provider_name",
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
"COALESCE(SUM(success_requests), 0)::BIGINT",
),
UsageAuditAggregationGroupBy::ApiFormat => (
"stats_user_daily_api_format",
"api_format",
"NULL::varchar",
"CASE WHEN COALESCE(SUM(response_time_samples), 0) > 0 THEN COALESCE(SUM(response_time_sum_ms), 0) / COALESCE(SUM(response_time_samples), 0) ELSE NULL END",
"NULL::BIGINT",
),
UsageAuditAggregationGroupBy::User => {
return Ok(Vec::new());
}
};
let provider_extra_where = if matches!(group_by, UsageAuditAggregationGroupBy::Provider) {
" AND BTRIM(COALESCE(provider_name, '')) <> '' AND lower(BTRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')"
@@ -7925,6 +7935,10 @@ ORDER BY "usage".user_id ASC
.await
.map_postgres_err()?
.rows_affected();
sqlx::query(REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL)
.execute(&mut **tx)
.await
.map_postgres_err()?;
Ok(rows_affected)
}) as BoxFuture<'_, Result<u64, DataLayerError>>
})
@@ -8371,6 +8385,40 @@ async fn apply_provider_api_key_usage_delta_in_tx(
.execute(&mut **tx)
.await
.map_postgres_err()?;
apply_provider_api_key_codex_window_usage_delta_in_tx(tx, key_id, delta).await?;
Ok(())
}
async fn apply_provider_api_key_codex_window_usage_delta_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
key_id: &str,
delta: &ProviderApiKeyUsageDelta,
) -> Result<(), DataLayerError> {
let Some(usage_created_at_unix_secs) = delta.usage_created_at_unix_secs else {
return Ok(());
};
if delta.request_count == 0 && delta.total_tokens == 0 && delta.total_cost_usd == 0.0 {
return Ok(());
}
let total_cost_usd_delta = if delta.total_cost_usd.is_finite() {
delta.total_cost_usd
} else {
0.0
};
sqlx::query(APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL)
.bind(key_id)
.bind(i64::try_from(usage_created_at_unix_secs).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"provider api key window usage timestamp exceeds i64: {usage_created_at_unix_secs}"
))
})?)
.bind(delta.request_count)
.bind(delta.total_tokens)
.bind(total_cost_usd_delta)
.execute(&mut **tx)
.await
.map_postgres_err()?;
Ok(())
}

166
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_codex_window_usage_delta_sql.sql Normal file
View File

@@ -0,0 +1,166 @@
WITH target_key AS (
SELECT
id,
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
FROM provider_api_keys
WHERE id = $1
AND jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
FOR UPDATE
),
window_items AS (
SELECT
target_key.id,
window_rows.window_item,
window_rows.window_ordinality
FROM target_key
CROSS JOIN LATERAL jsonb_array_elements(target_key.snapshot -> 'quota' -> 'windows')
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
),
parsed_windows AS (
SELECT
window_items.id,
window_items.window_item,
window_items.window_ordinality,
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
CASE
WHEN text_values.reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.reset_at_text) < 19
OR (
length(text_values.reset_at_text) = 19
AND text_values.reset_at_text <= '9223372036854775807'
)
)
THEN text_values.reset_at_text::BIGINT
ELSE NULL
END AS reset_at,
CASE
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
AND (
length(text_values.window_minutes_text) < 19
OR (
length(text_values.window_minutes_text) = 19
AND text_values.window_minutes_text <= '9223372036854775807'
)
)
THEN text_values.window_minutes_text::BIGINT
ELSE NULL
END AS window_minutes,
CASE
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.usage_reset_at_text) < 19
OR (
length(text_values.usage_reset_at_text) = 19
AND text_values.usage_reset_at_text <= '9223372036854775807'
)
)
THEN text_values.usage_reset_at_text::BIGINT
ELSE NULL
END AS usage_reset_at,
CASE
WHEN text_values.request_count_text ~ '^[0-9]+$'
AND (
length(text_values.request_count_text) < 19
OR (
length(text_values.request_count_text) = 19
AND text_values.request_count_text <= '9223372036854775807'
)
)
THEN text_values.request_count_text::BIGINT
ELSE 0
END AS current_request_count,
CASE
WHEN text_values.total_tokens_text ~ '^[0-9]+$'
AND (
length(text_values.total_tokens_text) < 19
OR (
length(text_values.total_tokens_text) = 19
AND text_values.total_tokens_text <= '9223372036854775807'
)
)
THEN text_values.total_tokens_text::BIGINT
ELSE 0
END AS current_total_tokens,
CASE
WHEN text_values.total_cost_usd_text ~ '^[-+]?[0-9]+([.][0-9]+)?$'
THEN text_values.total_cost_usd_text::DOUBLE PRECISION
ELSE 0
END AS current_total_cost_usd
FROM window_items
CROSS JOIN LATERAL (
SELECT
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'request_count', '')) AS request_count_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_tokens', '')) AS total_tokens_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_cost_usd', '')) AS total_cost_usd_text
) AS text_values
),
window_usage AS (
SELECT
parsed_windows.*,
CASE
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
THEN parsed_windows.window_minutes * 60
ELSE NULL
END AS window_seconds
FROM parsed_windows
),
updated_windows AS (
SELECT
window_usage.id,
jsonb_agg(
CASE
WHEN window_usage.window_code IN ('5h', 'weekly')
AND window_usage.reset_at IS NOT NULL
AND window_usage.window_seconds IS NOT NULL
AND window_usage.reset_at >= window_usage.window_seconds
AND window_usage.reset_at > $2
AND GREATEST(
window_usage.reset_at - window_usage.window_seconds,
COALESCE(window_usage.usage_reset_at, 0)
) <= $2
THEN jsonb_set(
window_usage.window_item,
'{usage}',
jsonb_build_object(
'request_count',
LEAST(
GREATEST(window_usage.current_request_count::NUMERIC + $3::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_tokens',
LEAST(
GREATEST(window_usage.current_total_tokens::NUMERIC + $4::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_cost_usd',
to_char(
GREATEST(COALESCE(window_usage.current_total_cost_usd, 0) + $5, 0),
'FM999999999999999990.00000000'
)
),
true
)
ELSE window_usage.window_item
END
ORDER BY window_usage.window_ordinality
) AS windows
FROM window_usage
GROUP BY window_usage.id
)
UPDATE provider_api_keys AS keys
SET
status_snapshot = jsonb_set(
target_key.snapshot,
'{quota,windows}',
updated_windows.windows,
true
)::json,
updated_at = NOW()
FROM target_key
JOIN updated_windows ON updated_windows.id = target_key.id
WHERE keys.id = target_key.id

164
crates/aether-data/src/repository/usage/postgres/queries/rebuild_provider_api_key_codex_window_usage_stats_sql.sql Normal file
View File

@@ -0,0 +1,164 @@
WITH target_keys AS (
SELECT
id,
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
FROM provider_api_keys
WHERE jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
),
window_items AS (
SELECT
target_keys.id,
window_rows.window_item,
window_rows.window_ordinality
FROM target_keys
CROSS JOIN LATERAL jsonb_array_elements(target_keys.snapshot -> 'quota' -> 'windows')
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
),
parsed_windows AS (
SELECT
window_items.id,
window_items.window_item,
window_items.window_ordinality,
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
CASE
WHEN text_values.reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.reset_at_text) < 19
OR (
length(text_values.reset_at_text) = 19
AND text_values.reset_at_text <= '9223372036854775807'
)
)
THEN text_values.reset_at_text::BIGINT
ELSE NULL
END AS reset_at,
CASE
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
AND (
length(text_values.window_minutes_text) < 19
OR (
length(text_values.window_minutes_text) = 19
AND text_values.window_minutes_text <= '9223372036854775807'
)
)
THEN text_values.window_minutes_text::BIGINT
ELSE NULL
END AS window_minutes,
CASE
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.usage_reset_at_text) < 19
OR (
length(text_values.usage_reset_at_text) = 19
AND text_values.usage_reset_at_text <= '9223372036854775807'
)
)
THEN text_values.usage_reset_at_text::BIGINT
ELSE NULL
END AS usage_reset_at
FROM window_items
CROSS JOIN LATERAL (
SELECT
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text
) AS text_values
),
window_usage AS (
SELECT
parsed_windows.*,
CASE
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
THEN parsed_windows.window_minutes * 60
ELSE NULL
END AS window_seconds
FROM parsed_windows
),
window_bounds AS (
SELECT
window_usage.*,
CASE
WHEN window_usage.window_code IN ('5h', 'weekly')
AND window_usage.reset_at IS NOT NULL
AND window_usage.window_seconds IS NOT NULL
AND window_usage.reset_at >= window_usage.window_seconds
THEN GREATEST(
window_usage.reset_at - window_usage.window_seconds,
COALESCE(window_usage.usage_reset_at, 0)
)
ELSE NULL
END AS window_start,
CASE
WHEN window_usage.window_code IN ('5h', 'weekly')
AND window_usage.reset_at IS NOT NULL
AND window_usage.window_seconds IS NOT NULL
AND window_usage.reset_at >= window_usage.window_seconds
THEN window_usage.reset_at
ELSE NULL
END AS window_end
FROM window_usage
),
aggregated AS (
SELECT
window_bounds.id,
window_bounds.window_ordinality,
COUNT("usage".id)::BIGINT AS request_count,
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)::BIGINT), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(COALESCE("usage".total_cost_usd, 0)), 0) AS DOUBLE PRECISION) AS total_cost_usd
FROM window_bounds
LEFT JOIN usage_billing_facts AS "usage"
ON window_bounds.window_start IS NOT NULL
AND window_bounds.window_end IS NOT NULL
AND "usage".provider_api_key_id = window_bounds.id
AND "usage".created_at >= to_timestamp(window_bounds.window_start::DOUBLE PRECISION)
AND "usage".created_at < to_timestamp(window_bounds.window_end::DOUBLE PRECISION)
GROUP BY
window_bounds.id,
window_bounds.window_ordinality
),
updated_windows AS (
SELECT
window_bounds.id,
jsonb_agg(
CASE
WHEN window_bounds.window_start IS NOT NULL
AND window_bounds.window_end IS NOT NULL
THEN jsonb_set(
window_bounds.window_item,
'{usage}',
jsonb_build_object(
'request_count',
COALESCE(aggregated.request_count, 0),
'total_tokens',
COALESCE(aggregated.total_tokens, 0),
'total_cost_usd',
to_char(
GREATEST(COALESCE(aggregated.total_cost_usd, 0), 0),
'FM999999999999999990.00000000'
)
),
true
)
ELSE window_bounds.window_item
END
ORDER BY window_bounds.window_ordinality
) AS windows
FROM window_bounds
LEFT JOIN aggregated
ON aggregated.id = window_bounds.id
AND aggregated.window_ordinality = window_bounds.window_ordinality
GROUP BY window_bounds.id
)
UPDATE provider_api_keys AS keys
SET
status_snapshot = jsonb_set(
target_keys.snapshot,
'{quota,windows}',
updated_windows.windows,
true
)::json,
updated_at = NOW()
FROM target_keys
JOIN updated_windows ON updated_windows.id = target_keys.id
WHERE keys.id = target_keys.id

36
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -235,18 +235,31 @@ fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
}
#[test]
fn usage_sql_summarizes_provider_key_window_usage_from_billing_facts() {
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("UNNEST"));
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
.contains("LEFT JOIN usage_billing_facts AS \"usage\""));
fn usage_sql_materializes_provider_key_window_usage_in_status_snapshot() {
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("UPDATE provider_api_keys AS keys"));
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("jsonb_set"));
assert!(
super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("created_at >= to_timestamp")
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'{quota,windows}'")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'usage'"));
assert!(
super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL.contains("created_at < to_timestamp")
!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("usage_billing_facts")
);
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
.contains("COUNT(\"usage\".id)::BIGINT AS request_count"));
}
#[test]
fn usage_sql_rebuilds_provider_key_window_usage_into_status_snapshot() {
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
.contains("UPDATE provider_api_keys AS keys"));
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
.contains("usage_billing_facts"));
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
.contains("provider_type', ''))) = 'codex'"));
assert!(
super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL.contains("'{quota,windows}'")
);
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL.contains("'{usage}'"));
}
#[test]
@@ -470,7 +483,7 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(super::SUMMARIZE_USAGE_TOTALS_BY_USER_IDS_SQL
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(super::SUMMARIZE_PROVIDER_API_KEY_WINDOW_USAGE_SQL
assert!(!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("usage_billing_facts AS \"usage\""));
}
@@ -498,7 +511,10 @@ fn usage_billing_facts_projects_upstream_stream_mode() {
assert!(migration.contains("COALESCE(usage_rows.upstream_is_stream"));
assert!(migration.contains("COALESCE(usage_rows.is_stream, FALSE)"));
assert!(migration.contains("ADD COLUMN IF NOT EXISTS upstream_is_stream boolean"));
assert!(migration.contains("request_metadata->>'upstream_is_stream'"));
assert!(
!migration.contains("request_metadata->>'upstream_is_stream'"),
"migration should avoid backfilling historical usage rows from request metadata"
);
}
#[test]

54
crates/aether-scheduler-core/src/request_candidate.rs
View File

@@ -24,6 +24,8 @@ pub struct SchedulerRequestCandidateReportContext {
pub body_rules: Option<Value>,
pub proxy: Option<Value>,
pub error_flow: Option<Value>,
pub candidate_group_id: Option<String>,
pub pool_key_index: Option<u32>,
pub ranking_mode: Option<String>,
pub priority_mode: Option<String>,
pub ranking_index: Option<u32>,
@@ -65,6 +67,8 @@ struct ReportCandidateExtraDataInput {
body_rules: Option<Value>,
proxy: Option<Value>,
error_flow: Option<Value>,
candidate_group_id: Option<String>,
pool_key_index: Option<u32>,
ranking_mode: Option<String>,
priority_mode: Option<String>,
ranking_index: Option<u32>,
@@ -153,6 +157,8 @@ pub fn parse_request_candidate_report_context(
.get("error_flow")
.cloned()
.filter(|value| !value.is_null()),
candidate_group_id: string_field(report_context, "candidate_group_id"),
pool_key_index: u32_field(report_context, "pool_key_index"),
ranking_mode: string_field(report_context, "ranking_mode"),
priority_mode: string_field(report_context, "priority_mode"),
ranking_index: u32_field(report_context, "ranking_index"),
@@ -188,6 +194,8 @@ pub fn resolve_report_request_candidate_slot(
body_rules,
proxy,
error_flow,
candidate_group_id,
pool_key_index,
ranking_mode,
priority_mode,
ranking_index,
@@ -206,6 +214,8 @@ pub fn resolve_report_request_candidate_slot(
body_rules,
proxy,
error_flow,
candidate_group_id,
pool_key_index,
ranking_mode,
priority_mode,
ranking_index,
@@ -323,6 +333,8 @@ pub fn build_execution_request_candidate_seed(
body_rules: metadata.body_rules,
proxy: metadata.proxy,
error_flow: metadata.error_flow,
candidate_group_id: metadata.candidate_group_id,
pool_key_index: metadata.pool_key_index,
ranking_mode: metadata.ranking_mode,
priority_mode: metadata.priority_mode,
ranking_index: metadata.ranking_index,
@@ -426,6 +438,8 @@ pub fn build_local_request_candidate_status_record(
body_rules: metadata.body_rules.clone(),
proxy: metadata.proxy.clone(),
error_flow: metadata.error_flow.clone(),
candidate_group_id: metadata.candidate_group_id.clone(),
pool_key_index: metadata.pool_key_index,
ranking_mode: metadata.ranking_mode.clone(),
priority_mode: metadata.priority_mode.clone(),
ranking_index: metadata.ranking_index,
@@ -660,6 +674,8 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
body_rules,
proxy,
error_flow,
candidate_group_id,
pool_key_index,
ranking_mode,
priority_mode,
ranking_index,
@@ -703,6 +719,22 @@ fn build_report_candidate_extra_data(input: ReportCandidateExtraDataInput) -> Op
if let Some(error_flow) = error_flow {
extra_data.insert("error_flow".to_string(), error_flow);
}
if let Some(candidate_group_id) = candidate_group_id {
extra_data.insert(
"candidate_group_id".to_string(),
Value::String(candidate_group_id.clone()),
);
extra_data.insert(
"pool_group_id".to_string(),
Value::String(candidate_group_id),
);
}
if let Some(pool_key_index) = pool_key_index {
extra_data.insert(
"pool_key_index".to_string(),
Value::Number(pool_key_index.into()),
);
}
if let Some(ranking_mode) = ranking_mode {
extra_data.insert("ranking_mode".to_string(), Value::String(ranking_mode));
}
@@ -880,7 +912,9 @@ mod tests {
"classification": "retry_upstream_failure",
"decision": "retry_next_candidate",
"propagation": "suppressed"
}
},
"candidate_group_id": "pool-group-1",
"pool_key_index": 2
})))
.expect("metadata");
@@ -935,6 +969,24 @@ mod tests {
.and_then(|value| value.get("propagation")),
Some(&json!("suppressed"))
);
assert_eq!(
slot.extra_data
.as_ref()
.and_then(|value| value.get("candidate_group_id")),
Some(&json!("pool-group-1"))
);
assert_eq!(
slot.extra_data
.as_ref()
.and_then(|value| value.get("pool_group_id")),
Some(&json!("pool-group-1"))
);
assert_eq!(
slot.extra_data
.as_ref()
.and_then(|value| value.get("pool_key_index")),
Some(&json!(2))
);
}
#[test]

149
crates/aether-usage-runtime/src/runtime.rs
View File

@@ -328,6 +328,26 @@ impl UsageRuntime {
self.enqueue_or_write_terminal(data, event).await;
}
pub async fn record_terminal_event_direct<T>(&self, data: &T, mut event: UsageEvent)
where
T: UsageRuntimeAccess,
{
if !self.is_enabled() {
return;
}
apply_body_capture_policy_from_data(data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %event.request_id,
error = %err,
"usage runtime failed to enrich terminal usage event with billing"
);
}
self.write_terminal_direct(data, &event).await;
}
async fn enqueue_or_write_terminal<T>(&self, data: &T, event: UsageEvent)
where
T: UsageRuntimeAccess,
@@ -360,7 +380,14 @@ impl UsageRuntime {
}
}
match build_upsert_usage_record_from_event(&event) {
self.write_terminal_direct(data, &event).await;
}
async fn write_terminal_direct<T>(&self, data: &T, event: &UsageEvent)
where
T: UsageRuntimeAccess,
{
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 {
@@ -518,7 +545,9 @@ fn now_unix_secs() -> u64 {
mod tests {
use std::sync::Mutex;
use aether_data::driver::redis::RedisStreamRunner;
use aether_data::driver::redis::{
RedisClientConfig, RedisClientFactory, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
};
@@ -542,6 +571,29 @@ mod tests {
records: Mutex<Vec<UpsertUsageRecord>>,
}
struct RedisConfiguredUsageStore {
inner: NoRedisUsageStore,
runner: 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(),
RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
#[async_trait]
impl UsageRecordWriter for NoRedisUsageStore {
async fn upsert_usage_record(
@@ -601,6 +653,64 @@ mod tests {
}
}
#[async_trait]
impl UsageRecordWriter for RedisConfiguredUsageStore {
async fn upsert_usage_record(
&self,
record: UpsertUsageRecord,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
self.inner.upsert_usage_record(record).await
}
}
#[async_trait]
impl UsageSettlementWriter for RedisConfiguredUsageStore {
fn has_usage_settlement_writer(&self) -> bool {
false
}
async fn settle_usage(
&self,
_input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
Ok(None)
}
}
#[async_trait]
impl UsageBillingEventEnricher for RedisConfiguredUsageStore {
async fn enrich_usage_event(&self, _event: &mut UsageEvent) -> Result<(), DataLayerError> {
Ok(())
}
}
#[async_trait]
impl ManualProxyNodeCounter for RedisConfiguredUsageStore {
async fn increment_manual_proxy_node_requests(
&self,
_node_id: &str,
_total_delta: i64,
_failed_delta: i64,
_latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
Ok(())
}
}
impl UsageRuntimeAccess for RedisConfiguredUsageStore {
fn has_usage_writer(&self) -> bool {
true
}
fn has_usage_worker_runner(&self) -> bool {
true
}
fn usage_worker_runner(&self) -> Option<RedisStreamRunner> {
Some(self.runner.clone())
}
}
#[tokio::test]
async fn terminal_usage_without_redis_writes_directly_to_usage_repository() {
let runtime = UsageRuntime::new(UsageRuntimeConfig {
@@ -633,6 +743,41 @@ mod tests {
assert_eq!(records[0].total_tokens, Some(12));
}
#[tokio::test]
async fn direct_terminal_usage_bypasses_redis_queue_and_writes_repository() {
let runtime = UsageRuntime::new(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})
.expect("usage runtime should build");
let runner = sample_runner();
let store = RedisConfiguredUsageStore {
inner: NoRedisUsageStore::default(),
runner,
};
let event = UsageEvent::new(
UsageEventType::Failed,
"req-direct-terminal-1",
UsageEventData {
user_id: Some("user-direct-terminal-1".to_string()),
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
status_code: Some(503),
error_message: Some("upstream failed".to_string()),
..UsageEventData::default()
},
);
runtime.record_terminal_event_direct(&store, event).await;
let records = store.inner.records.lock().expect("records lock");
assert_eq!(records.len(), 1);
assert_eq!(records[0].request_id, "req-direct-terminal-1");
assert_eq!(records[0].status, "failed");
assert_eq!(records[0].billing_status, "void");
assert_eq!(records[0].status_code, Some(503));
}
#[test]
fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() {
let mut event = UsageEvent::new(