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

feat(gateway): add Codex Live and OpenAI Realtime

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/runtime/src/repository/candidate_selection/memory.rs
+76 -10
View File
@@ -1,5 +1,6 @@
use std::sync::RwLock;
use aether_data_contracts::repository::candidate_selection::provider_model_mapping_api_format_covers;
use async_trait::async_trait;
use super::{
@@ -242,9 +243,13 @@ fn row_matches_requested_model(
.is_some_and(|mappings| {
mappings.iter().any(|mapping| {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_scope_covers(value, api_format))
formats.iter().any(|value| {
provider_model_mapping_api_format_covers(
&row.provider_type,
value,
api_format,
)
})
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
@@ -298,9 +303,9 @@ fn mapping_scope_matches(
api_format: &str,
) -> bool {
mapping.api_formats.as_ref().is_none_or(|formats| {
formats
.iter()
.any(|value| api_format_scope_covers(value, api_format))
formats.iter().any(|value| {
provider_model_mapping_api_format_covers(&row.provider_type, value, api_format)
})
}) && mapping.endpoint_ids.as_ref().is_none_or(|endpoint_ids| {
endpoint_ids
.iter()
@@ -308,10 +313,6 @@ fn mapping_scope_matches(
})
}
fn api_format_scope_covers(allowed: &str, requested: &str) -> bool {
aether_ai_formats::api_format_permission_covers(allowed, requested)
}
fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format: &str) -> bool {
let provider_type = row.provider_type.trim().to_ascii_lowercase();
let auth_type = row.key_auth_type.trim().to_ascii_lowercase();
@@ -325,6 +326,7 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
| "openai:responses:compact"
| "openai:search"
| "openai:image"
| "codex:live"
)
}
"chatgpt_web" => {
@@ -496,6 +498,70 @@ mod tests {
);
}
#[tokio::test]
async fn codex_live_accepts_oauth_keys_without_accepting_api_keys() {
let mut oauth = sample_row("codex-live-oauth", "codex:live", "gpt-live", 10);
oauth.provider_type = "codex".to_string();
oauth.key_auth_type = "oauth".to_string();
let mut api_key = oauth.clone();
api_key.provider_id = "codex-live-api-key".to_string();
api_key.endpoint_id = "endpoint-codex-live-api-key".to_string();
api_key.key_id = "key-codex-live-api-key".to_string();
api_key.model_id = "model-codex-live-api-key".to_string();
api_key.key_auth_type = "api_key".to_string();
api_key.provider_priority = 20;
let repository =
InMemoryMinimalCandidateSelectionReadRepository::seed(vec![api_key, oauth]);
let rows = repository
.list_for_exact_api_format_and_requested_model("codex:live", "gpt-live")
.await
.expect("Codex Live candidate should load");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].provider_id, "codex-live-oauth");
assert_eq!(rows[0].key_auth_type, "oauth");
}
#[tokio::test]
async fn codex_live_reuses_only_codex_legacy_responses_model_mappings() {
let legacy_row = |provider_id: &str, provider_type: &str, auth_type: &str| {
let mut row = sample_row(
provider_id,
"codex:live",
&format!("global-{provider_id}"),
10,
);
row.provider_type = provider_type.to_string();
row.key_auth_type = auth_type.to_string();
row.key_api_formats = Some(vec!["codex:live".to_string()]);
row.model_provider_model_name = format!("upstream-{provider_id}");
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "legacy-live-alias".to_string(),
priority: 1,
api_formats: Some(vec!["openai:responses".to_string()]),
endpoint_ids: None,
operations: None,
}]);
row
};
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
legacy_row("codex-provider", "codex", "oauth"),
legacy_row("openai-provider", "openai", "api_key"),
legacy_row("custom-provider", "custom", "api_key"),
]);
let rows = repository
.list_for_exact_api_format_and_requested_model("codex:live", "legacy-live-alias")
.await
.expect("legacy Codex Live mapping should resolve");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].provider_id, "codex-provider");
assert_eq!(rows[0].provider_type, "codex");
}
#[tokio::test]
async fn includes_grok_oauth_rows_for_chat_models() {
let mut row = sample_row(
crates/aether-data/runtime/src/repository/usage/memory.rs
+27 -6
View File
@@ -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)
crates/aether-data/runtime/src/repository/usage/memory/tests.rs
+123 -1
View File
@@ -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);