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

refactor gateway orchestration and failover effects

Browse Source
This commit is contained in:
fawney19
2026-04-18 11:35:11 +08:00
parent 3321bb3ccc
commit 569242d72f
42 changed files with 4365 additions and 1486 deletions
Download Patch File Download Diff File Expand all files Collapse all files

625
apps/aether-gateway/src/usage/reporting/codex_realtime_quota.rs
View File

@@ -1,625 +0,0 @@
use std::collections::{BTreeMap, HashMap};
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
use crate::clock::current_unix_secs;
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::{AppState, GatewayError};
use aether_admin::provider::quota as admin_provider_quota_pure;
use serde_json::Value;
const CACHE_TTL_SECONDS: u64 = 30;
const CACHE_MAX_ENTRIES: usize = 4096;
type HeaderFingerprintCache = Mutex<HashMap<String, (String, Instant)>>;
static HEADER_FINGERPRINT_CACHE: OnceLock<HeaderFingerprintCache> = OnceLock::new();
fn header_fingerprint_cache() -> &'static HeaderFingerprintCache {
HEADER_FINGERPRINT_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn report_context_key_id(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|context| context.get("key_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn is_volatile_compare_field(key: &str) -> bool {
key == "updated_at" || key.ends_with("_reset_seconds") || key.ends_with("_reset_after_seconds")
}
fn canonicalize_value(value: &Value) -> Value {
match value {
Value::Array(items) => Value::Array(items.iter().map(canonicalize_value).collect()),
Value::Object(object) => {
let mut entries = object.iter().collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.cmp(right.0));
let mut normalized = serde_json::Map::new();
for (key, value) in entries {
normalized.insert(key.clone(), canonicalize_value(value));
}
Value::Object(normalized)
}
_ => value.clone(),
}
}
fn fingerprint_codex_payload(value: &Value) -> Option<String> {
let object = value.as_object()?;
let mut entries = object
.iter()
.filter(|(key, _)| !is_volatile_compare_field(key))
.collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.cmp(right.0));
let mut normalized = serde_json::Map::new();
for (key, value) in entries {
normalized.insert(key.clone(), canonicalize_value(value));
}
serde_json::to_string(&Value::Object(normalized)).ok()
}
fn get_cached_fingerprint(key_id: &str, now: Instant) -> Option<String> {
let mut cache = header_fingerprint_cache()
.lock()
.expect("codex realtime quota cache should lock");
match cache.get(key_id) {
Some((fingerprint, expires_at)) if *expires_at > now => Some(fingerprint.clone()),
Some(_) => {
cache.remove(key_id);
None
}
None => None,
}
}
fn set_cached_fingerprint(key_id: &str, fingerprint: String, now: Instant) {
let mut cache = header_fingerprint_cache()
.lock()
.expect("codex realtime quota cache should lock");
cache.insert(
key_id.to_string(),
(
fingerprint,
now.checked_add(Duration::from_secs(CACHE_TTL_SECONDS))
.unwrap_or(now),
),
);
cache.retain(|_, (_, expires_at)| *expires_at > now);
if cache.len() <= CACHE_MAX_ENTRIES {
return;
}
let mut entries = cache
.iter()
.map(|(key, (_, expires_at))| (key.clone(), *expires_at))
.collect::<Vec<_>>();
entries.sort_by_key(|entry| entry.1);
for (key, _) in entries.into_iter().take(cache.len() - CACHE_MAX_ENTRIES) {
cache.remove(&key);
}
}
fn merge_metadata_object(
current: Option<&Value>,
section_key: &str,
section_value: Value,
) -> Option<Value> {
let mut merged = current
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
merged.insert(section_key.to_string(), section_value);
Some(Value::Object(merged))
}
pub(super) async fn sync_codex_quota_from_response_headers(
state: &AppState,
report_context: Option<&Value>,
headers: &BTreeMap<String, String>,
) -> Result<bool, GatewayError> {
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
let now_unix_secs = current_unix_secs();
let Some(parsed) = admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
else {
return Ok(false);
};
let Some(incoming_fingerprint) = fingerprint_codex_payload(&parsed) else {
return Ok(false);
};
let now = Instant::now();
if get_cached_fingerprint(&key_id, now).as_deref() == Some(incoming_fingerprint.as_str()) {
return Ok(false);
}
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
else {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
if !provider.provider_type.trim().eq_ignore_ascii_case("codex") {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
}
let current_codex = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.cloned()
.unwrap_or_else(serde_json::Map::new);
let current_codex = Value::Object(current_codex);
let Some(current_fingerprint) = fingerprint_codex_payload(&current_codex) else {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
if current_fingerprint == incoming_fingerprint {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
}
let updated_upstream_metadata =
merge_metadata_object(key.upstream_metadata.as_ref(), "codex", parsed);
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"response_headers",
);
let mut updated_key = key;
updated_key.upstream_metadata = updated_upstream_metadata;
updated_key.status_snapshot = updated_status_snapshot;
updated_key.updated_at_unix_secs = Some(now_unix_secs);
let updated = state
.update_provider_catalog_key(&updated_key)
.await?
.is_some();
if updated {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
}
Ok(updated)
}
#[cfg(test)]
pub(super) fn clear_codex_quota_fingerprint_cache() {
if let Some(cache) = HEADER_FINGERPRINT_CACHE.get() {
cache
.lock()
.expect("codex realtime quota cache should lock")
.clear();
}
}
#[cfg(test)]
mod tests {
use super::{clear_codex_quota_fingerprint_cache, sync_codex_quota_from_response_headers};
use crate::data::GatewayDataState;
use crate::AppState;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::{json, Value};
use std::collections::BTreeMap;
use std::sync::Arc;
fn build_state(repository: Arc<InMemoryProviderCatalogReadRepository>) -> AppState {
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
)
}
fn sample_provider(provider_id: &str, provider_type: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
provider_id.to_string(),
provider_type.to_string(),
None,
provider_type.to_string(),
)
.expect("provider should build")
}
fn sample_key(
key_id: &str,
provider_id: &str,
upstream_metadata: Option<Value>,
) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
key_id.to_string(),
provider_id.to_string(),
"default".to_string(),
"bearer".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:cli"])),
"sk-codex-test".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
key.upstream_metadata = upstream_metadata;
key
}
fn quota_snapshot<'a>(key: &'a StoredProviderCatalogKey) -> &'a serde_json::Map<String, Value> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.expect("quota snapshot should exist")
}
fn paid_headers(
primary_used_percent: &str,
secondary_used_percent: &str,
primary_reset_after_seconds: &str,
secondary_reset_after_seconds: &str,
) -> BTreeMap<String, String> {
BTreeMap::from([
("x-codex-plan-type".to_string(), "team".to_string()),
(
"x-codex-primary-used-percent".to_string(),
primary_used_percent.to_string(),
),
(
"x-codex-secondary-used-percent".to_string(),
secondary_used_percent.to_string(),
),
(
"x-codex-primary-window-minutes".to_string(),
"300".to_string(),
),
(
"x-codex-secondary-window-minutes".to_string(),
"10080".to_string(),
),
(
"x-codex-primary-reset-after-seconds".to_string(),
primary_reset_after_seconds.to_string(),
),
(
"x-codex-secondary-reset-after-seconds".to_string(),
secondary_reset_after_seconds.to_string(),
),
(
"x-codex-primary-reset-at".to_string(),
"1776148929".to_string(),
),
(
"x-codex-secondary-reset-at".to_string(),
"1776657828".to_string(),
),
(
"x-codex-credits-has-credits".to_string(),
"False".to_string(),
),
("x-codex-credits-balance".to_string(), "".to_string()),
("x-codex-credits-unlimited".to_string(), "False".to_string()),
])
}
#[tokio::test]
async fn sync_codex_quota_replaces_existing_codex_fields_and_preserves_other_sections() {
clear_codex_quota_fingerprint_cache();
let mut key = sample_key(
"key-codex-1",
"provider-codex",
Some(json!({
"codex": {
"legacy_marker": "drop-me",
"secondary_used_percent": 2.0,
"credits_balance": 42.0,
"account_disabled": true,
"reason": "deactivated_workspace"
},
"other": {
"value": true
}
})),
);
key.status_snapshot = Some(json!({
"oauth": {
"code": "valid",
"label": "有效",
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"blocked": false,
"recoverable": false
},
"quota": {
"code": "unknown",
"exhausted": false
}
}));
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![key],
));
let state = build_state(Arc::clone(&repository));
let updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-1",
"key_id": "key-codex-1"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("codex realtime sync should succeed");
assert!(updated);
let reloaded = repository
.list_keys_by_ids(&["key-codex-1".to_string()])
.await
.expect("keys should list");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.get("plan_type"), Some(&json!("team")));
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
assert_eq!(codex.get("has_credits"), Some(&json!(false)));
assert_eq!(codex.get("credits_unlimited"), Some(&json!(false)));
assert!(codex.get("legacy_marker").is_none());
assert!(codex.get("credits_balance").is_none());
assert!(codex.get("account_disabled").is_none());
assert!(codex.get("reason").is_none());
assert!(codex.get("updated_at").and_then(Value::as_u64).is_some());
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("other")),
Some(&json!({"value": true}))
);
let quota = quota_snapshot(&reloaded[0]);
assert_eq!(quota.get("version"), Some(&json!(2)));
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
assert_eq!(quota.get("source"), Some(&json!("response_headers")));
assert_eq!(quota.get("code"), Some(&json!("exhausted")));
assert_eq!(quota.get("exhausted"), Some(&json!(true)));
assert_eq!(quota.get("plan_type"), Some(&json!("team")));
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
let windows = quota
.get("windows")
.and_then(Value::as_array)
.expect("windows should be array");
assert_eq!(windows.len(), 2);
assert_eq!(windows[0].get("code"), Some(&json!("weekly")));
assert_eq!(windows[1].get("code"), Some(&json!("5h")));
let oauth = reloaded[0]
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("oauth"))
.and_then(Value::as_object)
.expect("oauth snapshot should exist");
assert_eq!(oauth.get("code"), Some(&json!("valid")));
let account = reloaded[0]
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object)
.expect("account snapshot should exist");
assert_eq!(account.get("code"), Some(&json!("ok")));
}
#[tokio::test]
async fn sync_codex_quota_skips_non_codex_provider() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai")],
Vec::new(),
vec![sample_key("key-openai-1", "provider-openai", None)],
));
let state = build_state(Arc::clone(&repository));
let updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-openai-realtime-1",
"key_id": "key-openai-1"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("non-codex realtime sync should not fail");
assert!(!updated);
let reloaded = repository
.list_keys_by_ids(&["key-openai-1".to_string()])
.await
.expect("keys should list");
assert_eq!(reloaded[0].upstream_metadata, None);
}
#[tokio::test]
async fn sync_codex_quota_skips_when_headers_do_not_contain_codex_metadata() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key("key-codex-2", "provider-codex", None)],
));
let state = build_state(Arc::clone(&repository));
let updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-2",
"key_id": "key-codex-2"
})),
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
)
.await
.expect("empty codex headers should not fail");
assert!(!updated);
let reloaded = repository
.list_keys_by_ids(&["key-codex-2".to_string()])
.await
.expect("keys should list");
assert_eq!(reloaded[0].upstream_metadata, None);
}
#[tokio::test]
async fn sync_codex_quota_cache_hit_skips_when_only_reset_countdown_changes() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key("key-codex-3", "provider-codex", None)],
));
let state = build_state(Arc::clone(&repository));
let key_ids = ["key-codex-3".to_string()];
let first_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-3",
"key_id": "key-codex-3"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("first realtime sync should succeed");
assert!(first_updated);
let first_snapshot = repository
.list_keys_by_ids(&key_ids)
.await
.expect("keys should list")
.into_iter()
.next()
.expect("key should exist");
let second_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-3",
"key_id": "key-codex-3"
})),
&paid_headers("100", "31", "42", "84"),
)
.await
.expect("second realtime sync should succeed");
assert!(!second_updated);
let second_snapshot = repository
.list_keys_by_ids(&key_ids)
.await
.expect("keys should list")
.into_iter()
.next()
.expect("key should exist");
assert_eq!(
first_snapshot.upstream_metadata,
second_snapshot.upstream_metadata
);
assert_eq!(
first_snapshot.updated_at_unix_secs,
second_snapshot.updated_at_unix_secs
);
}
#[tokio::test]
async fn sync_codex_quota_updates_when_usage_percent_changes() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key("key-codex-4", "provider-codex", None)],
));
let state = build_state(Arc::clone(&repository));
let first_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-4",
"key_id": "key-codex-4"
})),
&paid_headers("98", "31", "15160", "524059"),
)
.await
.expect("first realtime sync should succeed");
assert!(first_updated);
let second_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-4",
"key_id": "key-codex-4"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("second realtime sync should succeed");
assert!(second_updated);
let reloaded = repository
.list_keys_by_ids(&["key-codex-4".to_string()])
.await
.expect("keys should list");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
}
}

183
apps/aether-gateway/src/usage/reporting/mod.rs
View File

@@ -4,22 +4,20 @@ use aether_contracts::ExecutionError;
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{execution_error_details, SchedulerRequestCandidateStatusUpdate};
use tracing::{debug, warn};
use uuid::Uuid;
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::clock::current_unix_ms;
use crate::log_ids::short_request_id;
use crate::orchestration::{apply_local_report_effect, LocalReportEffect};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::{AppState, GatewayError};
mod codex_realtime_quota;
mod context;
use context::{report_context_is_locally_actionable, resolve_locally_actionable_report_context};
use aether_usage_runtime::{
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, normalize_gemini_file_name,
report_request_id, should_handle_local_stream_report, should_handle_local_sync_report,
sync_report_represents_failure, GEMINI_FILE_MAPPING_TTL_SECONDS,
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, report_request_id,
should_handle_local_stream_report, should_handle_local_sync_report,
sync_report_represents_failure,
};
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
@@ -188,7 +186,6 @@ pub(crate) async fn submit_stream_report(
}
async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportRequest) {
apply_local_gemini_file_mapping_side_effect(state, payload).await;
let terminal_unix_ms = current_unix_ms();
let (error_type, error_message) =
execution_error_details(None::<&ExecutionError>, payload.body_json.as_ref());
@@ -215,22 +212,7 @@ async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportR
},
)
.await;
if let Err(err) = codex_realtime_quota::sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist codex realtime quota from sync response headers"
);
}
apply_local_report_effect(state, LocalReportEffect::Sync { payload }).await;
}
async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamReportRequest) {
@@ -253,154 +235,7 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
},
)
.await;
if let Err(err) = codex_realtime_quota::sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist codex realtime quota from stream response headers"
);
}
}
async fn apply_local_gemini_file_mapping_side_effect(
state: &AppState,
payload: &GatewaySyncReportRequest,
) {
match payload.report_kind.as_str() {
"gemini_files_store_mapping" => {
if payload.status_code >= 300 {
return;
}
let key_id = payload
.report_context
.as_ref()
.and_then(|context| context.get("file_key_id"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let user_id = payload
.report_context
.as_ref()
.and_then(|context| context.get("user_id"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let Some(key_id) = key_id else {
return;
};
for entry in extract_gemini_file_mapping_entries(payload) {
if let Err(err) = store_local_gemini_file_mapping(
state,
entry.file_name.as_str(),
key_id,
user_id,
entry.display_name.as_deref(),
entry.mime_type.as_deref(),
)
.await
{
warn!(
event_name = "gemini_file_mapping_store_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
file_name = %entry.file_name,
error = ?err,
"gateway failed to persist gemini file mapping locally"
);
}
}
}
"gemini_files_delete_mapping" if payload.status_code < 300 => {
let file_name = payload
.report_context
.as_ref()
.and_then(|context| context.get("file_name"))
.and_then(serde_json::Value::as_str)
.and_then(normalize_gemini_file_name);
let Some(file_name) = file_name else {
return;
};
if let Err(err) = delete_local_gemini_file_mapping(state, file_name.as_str()).await {
warn!(
event_name = "gemini_file_mapping_delete_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
file_name = %file_name,
error = ?err,
"gateway failed to delete gemini file mapping locally"
);
}
}
_ => {}
}
}
pub(crate) async fn store_local_gemini_file_mapping(
state: &AppState,
file_name: &str,
key_id: &str,
user_id: Option<&str>,
display_name: Option<&str>,
mime_type: Option<&str>,
) -> Result<(), GatewayError> {
let Some(file_name) = normalize_gemini_file_name(file_name) else {
return Ok(());
};
let expires_at_unix_secs = current_unix_secs().saturating_add(GEMINI_FILE_MAPPING_TTL_SECONDS);
let _stored = state
.upsert_gemini_file_mapping(
aether_data::repository::gemini_file_mappings::UpsertGeminiFileMappingRecord {
id: Uuid::new_v4().to_string(),
file_name: file_name.clone(),
key_id: key_id.to_string(),
user_id: user_id.map(ToOwned::to_owned),
display_name: display_name.map(ToOwned::to_owned),
mime_type: mime_type.map(ToOwned::to_owned),
source_hash: None,
expires_at_unix_secs,
},
)
.await?;
state
.cache_set_string_with_ttl(
gemini_file_mapping_cache_key(file_name.as_str()).as_str(),
key_id,
GEMINI_FILE_MAPPING_TTL_SECONDS,
)
.await?;
Ok(())
}
async fn delete_local_gemini_file_mapping(
state: &AppState,
file_name: &str,
) -> Result<(), GatewayError> {
let Some(file_name) = normalize_gemini_file_name(file_name) else {
return Ok(());
};
let _deleted = state
.delete_gemini_file_mapping_by_file_name(file_name.as_str())
.await?;
state
.cache_delete_key(gemini_file_mapping_cache_key(file_name.as_str()).as_str())
.await?;
Ok(())
apply_local_report_effect(state, LocalReportEffect::Stream { payload }).await;
}
#[cfg(test)]
@@ -824,7 +659,7 @@ mod tests {
#[tokio::test]
async fn submit_sync_report_updates_codex_quota_from_response_headers() {
super::codex_realtime_quota::clear_codex_quota_fingerprint_cache();
crate::orchestration::clear_local_report_effect_caches_for_tests();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider(
@@ -888,7 +723,7 @@ mod tests {
#[tokio::test]
async fn submit_stream_report_updates_codex_quota_from_response_headers() {
super::codex_realtime_quota::clear_codex_quota_fingerprint_cache();
crate::orchestration::clear_local_report_effect_caches_for_tests();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider(