feat(gateway): add Codex Live and OpenAI Realtime

Implement preflighted Live/Realtime WebSocket transports, protocol-aware authentication, usage auditing, UI filtering, and legacy Codex permission migration.
This commit is contained in:
ZheFox
2026-08-21 04:27:34 +08:00
parent fe38dcd294
commit 2c89202001
105 changed files with 7553 additions and 947 deletions
@@ -310,6 +310,11 @@ fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditLi
return false;
}
}
if let Some(is_websocket) = query.is_websocket {
if item.is_websocket() != is_websocket {
return false;
}
}
if query.error_only
&& item.status != "failed"
&& item.status_code.unwrap_or_default() < 400
@@ -391,6 +396,11 @@ fn usage_matches_keyword_search_query(
return false;
}
}
if let Some(is_websocket) = query.is_websocket {
if item.is_websocket() != is_websocket {
return false;
}
}
if query.error_only
&& item.status != "failed"
&& item.status_code.unwrap_or_default() < 400
@@ -2396,7 +2406,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
continue;
}
let entry = totals.entry(api_key_id.to_string()).or_insert(0);
*entry = (*entry).saturating_add(item.total_tokens);
if item.usage_available() {
*entry = (*entry).saturating_add(item.total_tokens);
}
}
Ok(totals)
}
@@ -2436,7 +2448,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
..Default::default()
});
entry.request_count = entry.request_count.saturating_add(1);
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
if item.usage_available() {
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
}
}
Ok(totals.into_values().collect())
}
@@ -2469,8 +2483,10 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
..StoredProviderApiKeyUsageSummary::default()
});
entry.request_count = entry.request_count.saturating_add(1);
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
entry.total_cost_usd += item.total_cost_usd;
if item.usage_available() {
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
entry.total_cost_usd += item.total_cost_usd;
}
entry.last_used_at_unix_secs = Some(
entry
.last_used_at_unix_secs
@@ -2524,8 +2540,10 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
}
summary.request_count = summary.request_count.saturating_add(1);
summary.total_tokens = summary.total_tokens.saturating_add(item.total_tokens);
summary.total_cost_usd += item.total_cost_usd;
if item.usage_available() {
summary.total_tokens = summary.total_tokens.saturating_add(item.total_tokens);
summary.total_cost_usd += item.total_cost_usd;
}
}
summaries.push(summary);
@@ -2598,6 +2616,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
let date_key = dt.date_naive().to_string();
let entry = daily.entry(date_key).or_insert((0, 0, 0.0, 0.0));
entry.0 += 1;
if !item.usage_available() {
continue;
}
let cache_creation = if item.cache_creation_input_tokens == 0
&& (item.cache_creation_ephemeral_5m_input_tokens
+ item.cache_creation_ephemeral_1h_input_tokens)
@@ -15,7 +15,8 @@ use crate::repository::usage::{
};
use aether_data_contracts::repository::usage::{
usage_body_ref, ProviderApiKeyWindowUsageRequest, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery,
UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery,
UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery,
UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery,
UsageTimeSeriesGranularity,
};
@@ -332,6 +333,79 @@ async fn provider_aggregation_skips_unknown_provider_labels() {
);
}
#[tokio::test]
async fn unmetered_session_audit_counts_lifecycle_without_token_or_cost_contribution() {
let metered = sample_usage("req-metered", 100);
let mut live = sample_usage("req-live", 200);
live.request_metadata = Some(json!({
"usage_available": false,
"websocket_mode": true,
"websocket_transport": "codex_live_direct",
}));
live.billing_status = "void".to_string();
live.input_tokens = 0;
live.output_tokens = 0;
live.total_tokens = 0;
live.cache_creation_input_tokens = 0;
live.cache_creation_ephemeral_5m_input_tokens = 0;
live.cache_creation_ephemeral_1h_input_tokens = 0;
live.cache_read_input_tokens = 0;
live.total_cost_usd = 0.0;
live.actual_total_cost_usd = 0.0;
let repository = InMemoryUsageReadRepository::seed(vec![metered, live]);
let listed = repository
.list_usage_audits(&UsageAuditListQuery {
created_from_unix_secs: Some(0),
created_until_unix_secs: Some(1_000),
newest_first: true,
..UsageAuditListQuery::default()
})
.await
.expect("audit list should succeed");
assert_eq!(listed.len(), 2);
assert!(!listed
.iter()
.find(|item| item.request_id == "req-live")
.expect("Live row should remain visible")
.usage_available());
let aggregate = repository
.aggregate_usage_audits(&UsageAuditAggregationQuery {
created_from_unix_secs: 0,
created_until_unix_secs: 1_000,
group_by: UsageAuditAggregationGroupBy::Model,
limit: 10,
exclude_reserved_provider_labels: false,
})
.await
.expect("aggregate should succeed");
assert_eq!(aggregate.len(), 1);
assert_eq!(aggregate[0].request_count, 2);
assert_eq!(aggregate[0].total_tokens, 150);
let summary = repository
.summarize_usage_audits(&UsageAuditSummaryQuery {
created_from_unix_secs: 0,
created_until_unix_secs: 1_000,
..UsageAuditSummaryQuery::default()
})
.await
.expect("summary should succeed");
assert_eq!(summary.total_requests, 2);
assert_eq!(summary.recorded_total_tokens, 150);
let provider_key_summaries = repository
.summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()])
.await
.expect("provider key lifecycle summary should succeed");
let provider_key_summary = provider_key_summaries
.get("provider-key-1")
.expect("provider key summary");
assert_eq!(provider_key_summary.request_count, 2);
assert_eq!(provider_key_summary.total_tokens, 150);
}
#[tokio::test]
async fn aggregation_can_skip_unknown_provider_records_for_model_and_api_format() {
let mut unknown = sample_usage("req-unknown-provider", 100);
@@ -1910,6 +1984,54 @@ async fn list_usage_audits_applies_second_based_time_filters() {
assert_eq!(items[0].request_id, "req-2");
}
#[tokio::test]
async fn usage_audit_websocket_filter_applies_to_list_count_and_keyword_search() {
let mut websocket = sample_usage("req-ws", 2);
websocket.request_metadata = Some(json!({
"websocket_mode": true,
"websocket_transport": "codex_live_direct",
}));
let repository =
InMemoryUsageReadRepository::seed(vec![sample_usage("req-http", 1), websocket]);
let list_query = crate::repository::usage::UsageAuditListQuery {
is_websocket: Some(true),
..Default::default()
};
let listed = repository
.list_usage_audits(&list_query)
.await
.expect("WebSocket list should succeed");
assert_eq!(listed.len(), 1);
assert_eq!(listed[0].request_id, "req-ws");
assert_eq!(
repository
.count_usage_audits(&list_query)
.await
.expect("WebSocket count should succeed"),
1
);
let keyword_query = UsageAuditKeywordSearchQuery {
is_websocket: Some(true),
keywords: vec!["gpt-4.1".to_string()],
..Default::default()
};
let keyword_matches = repository
.list_usage_audits_by_keyword_search(&keyword_query)
.await
.expect("WebSocket keyword list should succeed");
assert_eq!(keyword_matches.len(), 1);
assert_eq!(keyword_matches[0].request_id, "req-ws");
assert_eq!(
repository
.count_usage_audits_by_keyword_search(&keyword_query)
.await
.expect("WebSocket keyword count should succeed"),
1
);
}
#[tokio::test]
async fn dashboard_and_leaderboard_total_tokens_use_effective_cache_aware_tokens() {
let mut item = sample_usage("req-cache-aware-total", 1_711_000_000);