diff --git a/apps/aether-gateway/src/data/state/catalog.rs b/apps/aether-gateway/src/data/state/catalog.rs index be7d08775..a31fa0583 100644 --- a/apps/aether-gateway/src/data/state/catalog.rs +++ b/apps/aether-gateway/src/data/state/catalog.rs @@ -1,14 +1,14 @@ use super::{ ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery, GeminiFileMappingStats, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - PublicHealthStatusCount, PublicHealthTimelineBucket, StoredGeminiFileMapping, - StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, - StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, - StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate, - UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, + ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogKeyStatusSnapshotUpdate, PublicHealthStatusCount, PublicHealthTimelineBucket, + StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, + StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, + StoredRequestCandidate, UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord, }; impl GatewayDataState { @@ -561,6 +561,20 @@ impl GatewayDataState { Ok(updated) } + pub(crate) async fn compare_and_update_provider_catalog_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + let updated = match &self.provider_catalog_writer { + Some(repository) => repository.compare_and_update_key_admin_state(update).await, + None => Ok(false), + }?; + // Clear on both success and conflict so a retry cannot reuse the stale + // credential snapshot that lost the CAS. + self.clear_provider_catalog_cache(); + Ok(updated) + } + pub(crate) async fn update_provider_catalog_keys( &self, keys: &[StoredProviderCatalogKey], diff --git a/apps/aether-gateway/src/data/state/mod.rs b/apps/aether-gateway/src/data/state/mod.rs index 685639d4f..fcffb2a1e 100644 --- a/apps/aether-gateway/src/data/state/mod.rs +++ b/apps/aether-gateway/src/data/state/mod.rs @@ -121,13 +121,13 @@ use aether_data_contracts::repository::pool_scores::{ UpsertPoolMemberScore, }; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, - ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, - ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, - ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, - StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, - StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, + ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate, + ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery, + ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, + ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, + ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, + StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; use aether_data_contracts::repository::quota::{ ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot, diff --git a/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs b/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs index 1f13401b6..ac336ba7b 100644 --- a/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs +++ b/apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs @@ -2573,6 +2573,7 @@ fn json_execution_result( candidate_id: plan.candidate_id.clone(), status_code, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(body), body_bytes_b64: None, @@ -2615,6 +2616,7 @@ fn bytes_execution_result( candidate_id: plan.candidate_id.clone(), status_code, headers, + response_observation: None, body: Some(ResponseBody { json_body: None, body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)), @@ -2637,6 +2639,7 @@ fn execution_result_frame_stream( payload: StreamFramePayload::Headers { status_code: result.status_code, headers: result.headers.clone(), + response_observation: result.response_observation.clone(), }, }, StreamFrame { diff --git a/apps/aether-gateway/src/execution_runtime/fallback.rs b/apps/aether-gateway/src/execution_runtime/fallback.rs index e475850e4..3c7287b6d 100644 --- a/apps/aether-gateway/src/execution_runtime/fallback.rs +++ b/apps/aether-gateway/src/execution_runtime/fallback.rs @@ -480,6 +480,7 @@ mod tests { candidate_id: None, status_code: 502, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -519,6 +520,7 @@ mod tests { candidate_id: None, status_code: 502, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -583,6 +585,7 @@ mod tests { candidate_id: None, status_code: 429, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -614,6 +617,7 @@ mod tests { candidate_id: None, status_code: 401, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -645,6 +649,7 @@ mod tests { candidate_id: None, status_code: 502, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: Some(ExecutionError { @@ -708,6 +713,7 @@ mod tests { candidate_id: None, status_code: 404, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -900,6 +906,7 @@ mod tests { candidate_id: None, status_code: 200, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -1022,6 +1029,7 @@ mod tests { candidate_id: None, status_code: 429, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -1068,6 +1076,7 @@ mod tests { candidate_id: None, status_code: 429, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -1174,6 +1183,7 @@ mod tests { candidate_id: None, status_code: 200, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -1211,6 +1221,7 @@ mod tests { candidate_id: None, status_code: 400, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, @@ -1259,6 +1270,7 @@ mod tests { candidate_id: None, status_code: 429, headers: Default::default(), + response_observation: None, body: None, telemetry: None, error: None, diff --git a/apps/aether-gateway/src/execution_runtime/grok.rs b/apps/aether-gateway/src/execution_runtime/grok.rs index a6f37cabf..af27d5750 100644 --- a/apps/aether-gateway/src/execution_runtime/grok.rs +++ b/apps/aether-gateway/src/execution_runtime/grok.rs @@ -841,6 +841,7 @@ fn encode_grok_headers_frame( payload: StreamFramePayload::Headers { status_code, headers, + response_observation: None, }, }) } @@ -2157,6 +2158,7 @@ fn grok_execution_result( candidate_id: plan.candidate_id.clone(), status_code, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(body_json), body_bytes_b64: None, @@ -2220,6 +2222,7 @@ fn grok_collected_frame_stream( "application/json".to_string() }, )]), + response_observation: None, }, }, StreamFrame { diff --git a/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs b/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs index 1ffa21955..1afbf8ab0 100644 --- a/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs +++ b/apps/aether-gateway/src/execution_runtime/kiro_web_search.rs @@ -279,6 +279,7 @@ fn raw_response_frame_stream( payload: StreamFramePayload::Headers { status_code, headers, + response_observation: None, }, }, StreamFrame { @@ -1449,6 +1450,7 @@ mod tests { candidate_id: None, status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "jsonrpc": "2.0", diff --git a/apps/aether-gateway/src/execution_runtime/oauth_retry.rs b/apps/aether-gateway/src/execution_runtime/oauth_retry.rs index 7ede8cf7b..293193ebd 100644 --- a/apps/aether-gateway/src/execution_runtime/oauth_retry.rs +++ b/apps/aether-gateway/src/execution_runtime/oauth_retry.rs @@ -2,10 +2,10 @@ use aether_contracts::ExecutionPlan; use tracing::warn; use crate::orchestration::{ - oauth_status_may_be_invalid as status_may_be_oauth_invalid, + local_failover_error_message, oauth_status_may_be_invalid as status_may_be_oauth_invalid, oauth_status_proves_access_token_invalid as status_proves_access_token_invalid, }; -use crate::state::AgentIdentityAuthConfigFence; +use crate::state::{AgentIdentityAuthConfigFence, CodexRuntimeOAuthObservation}; use crate::{provider_transport::LocalOAuthRefreshError, AppState}; pub(crate) async fn refresh_oauth_plan_auth_for_retry( @@ -14,6 +14,9 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry( status_code: u16, response_text: Option<&str>, trace_id: &str, + report_context: Option<&serde_json::Value>, + request_started_at_unix_ms: Option, + request_order_id: Option<&str>, ) -> bool { if !status_may_be_oauth_invalid(status_code, response_text) { return false; @@ -109,15 +112,49 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry( body_excerpt, .. }) if matches!(refresh_status_code, 400 | 401 | 403) => { - if let Err(err) = state - .persist_local_oauth_refresh_failure_state( - &transport, - refresh_status_code, - body_excerpt.as_str(), - access_token_invalid_proven, - ) - .await - { + let observed_credential_generation = + report_context_string(report_context, "codex_credential_generation"); + let runtime_invalid_message = local_failover_error_message(response_text); + let runtime_invalid_reason = + aether_admin::provider::quota::codex_runtime_invalid_reason( + status_code, + runtime_invalid_message.as_deref(), + ); + let persist_result = match (request_started_at_unix_ms, request_order_id) { + (Some(request_started_at_unix_ms), Some(request_order_id)) + if transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") => + { + state + .persist_local_oauth_refresh_failure_state_observed( + &transport, + refresh_status_code, + body_excerpt.as_str(), + access_token_invalid_proven, + CodexRuntimeOAuthObservation { + request_started_at_unix_ms, + request_order_id, + observed_credential_generation, + runtime_invalid_reason: runtime_invalid_reason.as_deref(), + }, + ) + .await + } + _ => { + state + .persist_local_oauth_refresh_failure_state( + &transport, + refresh_status_code, + body_excerpt.as_str(), + access_token_invalid_proven, + ) + .await + } + }; + if let Err(err) = persist_result { warn!( event_name = "local_oauth_retry_refresh_failure_persist_failed", log_type = "ops", @@ -161,6 +198,17 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry( } } +fn report_context_string<'a>( + report_context: Option<&'a serde_json::Value>, + field: &str, +) -> Option<&'a str> { + report_context + .and_then(|context| context.get(field)) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> { plan.headers .iter() @@ -209,6 +257,7 @@ mod tests { use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; @@ -312,7 +361,7 @@ mod tests { } #[tokio::test] - async fn auto_removes_codex_key_after_request_proven_terminal_refresh_failure() { + async fn retains_codex_key_after_request_proven_terminal_refresh_failure() { let token_hits = Arc::new(Mutex::new(0usize)); let token_hits_clone = Arc::clone(&token_hits); let token_server = Router::new().route( @@ -458,16 +507,34 @@ mod tests { 401, Some(r#"{"error":"oauth_token_invalid"}"#), "trace-oauth-retry", + None, + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), ) .await; assert!(!retried); assert_eq!(*token_hits.lock().expect("mutex should lock"), 1); - let keys = provider_catalog_repository + let stored_key = provider_catalog_repository .list_keys_by_ids(&["key-codex-oauth-retry".to_string()]) .await - .expect("keys should read"); - assert!(keys.is_empty()); + .expect("keys should read") + .into_iter() + .next() + .expect("request-scoped refresh failure should retain the key"); + let invalid_reason = stored_key + .oauth_invalid_reason + .as_deref() + .expect("combined invalid reason should persist"); + assert!(invalid_reason.contains("[OAUTH_EXPIRED]")); + assert!(invalid_reason.contains("[REFRESH_FAILED]")); + assert_eq!( + stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_id")), + Some(&json!("01900000-0000-7000-8000-000000000010")) + ); token_handle.abort(); } @@ -619,6 +686,9 @@ mod tests { 401, Some(r#"{"error":"invalid_token"}"#), "trace-claude-oauth-fence-first", + None, + None, + None, ) .await ); @@ -647,6 +717,9 @@ mod tests { 401, Some(r#"{"error":"invalid_token"}"#), "trace-claude-oauth-fence-stale", + None, + None, + None, ) .await ); @@ -665,6 +738,7 @@ mod tests { .expect("Claude key should load") .pop() .expect("Claude key should exist"); + let expected_admin_replacement = admin_replacement.clone(); admin_replacement.encrypted_api_key = Some( encrypt_python_fernet_plaintext( DEVELOPMENT_ENCRYPTION_KEY, @@ -673,10 +747,23 @@ mod tests { .expect("admin access token should encrypt"), ); admin_replacement.expires_at_unix_secs = Some(4_102_444_800); - provider_catalog_repository - .update_key(&admin_replacement) + assert!(provider_catalog_repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: expected_admin_replacement + .encrypted_auth_config + .clone(), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: expected_admin_replacement.encrypted_api_key.clone(), + auth_type: expected_admin_replacement.auth_type.clone(), + provider_id: expected_admin_replacement.provider_id.clone(), + provider_type: "claude_code".to_string(), + }, + key: admin_replacement, + codex_rotation: None, + reset_oauth_runtime: true, + }) .await - .expect("admin replacement should persist"); + .expect("admin replacement CAS should run")); let admin_result = state .force_local_oauth_refresh_entry(&stale_transport) diff --git a/apps/aether-gateway/src/execution_runtime/response_header_rules.rs b/apps/aether-gateway/src/execution_runtime/response_header_rules.rs index c5d7f6118..16eff1bf3 100644 --- a/apps/aether-gateway/src/execution_runtime/response_header_rules.rs +++ b/apps/aether-gateway/src/execution_runtime/response_header_rules.rs @@ -10,6 +10,10 @@ use crate::{AppState, GatewayError}; const RESPONSE_HEADER_RULES_KEY: &str = "response_header_rules"; const RESPONSE_HEADER_RULES_CAMEL_KEY: &str = "responseHeaderRules"; const PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY: &str = "provider_response_headers"; +const PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY: &str = "provider_request_started_at_unix_ms"; +const PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY: &str = "provider_request_order_id"; +const PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY: &str = + "provider_response_headers_observed_at_unix_ms"; const RESPONSE_HEADER_RULE_PROTECTED_KEYS: &[&str] = &["content-length"]; const RESPONSE_HEADER_RULES_CACHE_TTL: Duration = Duration::from_secs(5); @@ -98,6 +102,9 @@ pub(crate) async fn apply_endpoint_response_header_rules( pub(crate) fn attach_provider_response_headers_to_report_context( report_context: Option, provider_headers: &BTreeMap, + provider_request_started_at_unix_ms: u64, + provider_response_headers_observed_at_unix_ms: u64, + provider_request_order_id: &str, ) -> Option { let provider_headers = serde_json::to_value(provider_headers).ok()?; let mut object = match report_context { @@ -105,9 +112,99 @@ pub(crate) fn attach_provider_response_headers_to_report_context( Some(other) => Map::from_iter([("seed".to_string(), other)]), None => Map::new(), }; - object.insert( - PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(), - provider_headers, - ); + let observation_is_absent = !object.contains_key(PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY) + && !object.contains_key(PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY) + && !object.contains_key(PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY) + && !object.contains_key(PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY); + if observation_is_absent { + object.insert( + PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(), + provider_headers, + ); + object.insert( + PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY.to_string(), + Value::from(provider_request_started_at_unix_ms), + ); + object.insert( + PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY.to_string(), + Value::from(provider_response_headers_observed_at_unix_ms), + ); + object.insert( + PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY.to_string(), + Value::from(provider_request_order_id), + ); + } Some(Value::Object(object)) } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn provider_response_observation_is_first_write_wins() { + let first_headers = + BTreeMap::from([("x-codex-primary-used-percent".to_string(), "10".to_string())]); + let second_headers = + BTreeMap::from([("x-codex-primary-used-percent".to_string(), "20".to_string())]); + + let report_context = attach_provider_response_headers_to_report_context( + Some(json!("seed-value")), + &first_headers, + 100, + 200, + "observation-1", + ); + let report_context = attach_provider_response_headers_to_report_context( + report_context, + &second_headers, + 300, + 400, + "observation-2", + ) + .expect("report context should exist"); + + assert_eq!(report_context["seed"], json!("seed-value")); + assert_eq!( + report_context["provider_response_headers"]["x-codex-primary-used-percent"], + json!("10") + ); + assert_eq!( + report_context["provider_request_started_at_unix_ms"], + json!(100) + ); + assert_eq!( + report_context["provider_response_headers_observed_at_unix_ms"], + json!(200) + ); + assert_eq!( + report_context["provider_request_order_id"], + json!("observation-1") + ); + } + + #[test] + fn provider_response_observation_does_not_complete_a_partial_triplet() { + let report_context = attach_provider_response_headers_to_report_context( + Some(json!({"provider_response_headers": {"x-existing": "1"}})), + &BTreeMap::from([("x-new".to_string(), "2".to_string())]), + 300, + 400, + "observation-2", + ) + .expect("report context should exist"); + + assert_eq!( + report_context["provider_response_headers"]["x-existing"], + json!("1") + ); + assert!(report_context + .get("provider_request_started_at_unix_ms") + .is_none()); + assert!(report_context + .get("provider_response_headers_observed_at_unix_ms") + .is_none()); + assert!(report_context.get("provider_request_order_id").is_none()); + } +} diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index 13b0c70ca..31b3f059c 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -11,8 +11,8 @@ use std::time::{Duration, Instant}; use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope}; use aether_contracts::{ - ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage, - StreamFrame, StreamFramePayload, + ExecutionPlan, ExecutionResponseObservation, ExecutionStreamTerminalSummary, + ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload, }; use aether_data_contracts::repository::candidates::{ RequestCandidateStatus, UpsertRequestCandidateRecord, @@ -112,12 +112,13 @@ use crate::execution_runtime::{ use crate::log_ids::short_request_id; use crate::orchestration::{ apply_local_execution_effect, build_local_error_flow_metadata, classify_failure_disposition, - cyber_continue_failover_enabled, trace_upstream_response_body, with_error_flow_report_context, + cyber_continue_failover_enabled, spawn_local_oauth_success_effect, + trace_upstream_response_body, with_error_flow_report_context, with_upstream_response_report_context, FailureDisposition, FailureTokenAction, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis, LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, - LocalPoolErrorEffect, + LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; use crate::provider_pool_demand::{ acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard, @@ -1249,6 +1250,9 @@ async fn execute_in_process_stream_with_oauth_retry( retry_status_code, response_text.as_deref(), trace_id, + report_context, + Some(execution.response_observation.request_started_at_unix_ms), + Some(&execution.response_observation.request_order_id), ) .await { @@ -2818,6 +2822,7 @@ async fn execute_stream_from_direct_passthrough( stream_precommit_committed: _, response, started_at: upstream_started_at, + response_observation, stream_first_byte_timeout, upstream_target_permit, } = execution; @@ -2834,8 +2839,23 @@ async fn execute_stream_from_direct_passthrough( let request_id = plan.request_id.clone(); let candidate_id = plan.candidate_id.clone(); let request_id_for_log = short_request_id(request_id.as_str()); - let mut report_context = - attach_provider_response_headers_to_report_context(report_context, &headers); + let mut report_context = attach_provider_response_headers_to_report_context( + report_context, + &headers, + response_observation.request_started_at_unix_ms, + response_observation.response_headers_observed_at_unix_ms, + &response_observation.request_order_id, + ); + spawn_local_oauth_success_effect( + state.clone(), + &plan, + report_context.as_ref(), + LocalOAuthSuccessEffect { + status_code, + request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms), + request_order_id: Some(&response_observation.request_order_id), + }, + ); if status_code == 200 { seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { @@ -3819,6 +3839,7 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), + None, ) .await; } @@ -3891,6 +3912,7 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), + None, ) .await; } @@ -3963,6 +3985,7 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), + None, ) .await; } @@ -4035,6 +4058,7 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), + None, ) .await; } @@ -4192,6 +4216,15 @@ async fn execute_execution_runtime_stream_inner( record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } + let report_context = attach_provider_response_headers_to_report_context( + report_context, + &execution.headers, + execution.response_observation.request_started_at_unix_ms, + execution + .response_observation + .response_headers_observed_at_unix_ms, + &execution.response_observation.request_order_id, + ); let stream_precommit_committed = execution.stream_precommit_committed; let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream_with_retry_scope( @@ -4211,6 +4244,7 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out, retry_fallback_out, + None, ) .await; } @@ -4327,6 +4361,15 @@ async fn execute_execution_runtime_stream_inner( record_stream_pending_lifecycle(state, seed, &mut stage_trace).await; lifecycle_pending_recorded = true; } + let report_context = attach_provider_response_headers_to_report_context( + report_context, + &execution.headers, + execution.response_observation.request_started_at_unix_ms, + execution + .response_observation + .response_headers_observed_at_unix_ms, + &execution.response_observation.request_order_id, + ); let stream_precommit_committed = execution.stream_precommit_committed; let frame_stream = build_direct_execution_frame_stream(execution).boxed(); return execute_stream_from_frame_stream_with_retry_scope( @@ -4346,10 +4389,13 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), + None, ) .await; } + let remote_request_started_at_unix_ms = current_request_candidate_unix_ms(); + let remote_request_order_id = uuid::Uuid::now_v7().to_string(); let response = match post_stream_plan_to_remote_execution_runtime( state, remote_execution_runtime_base_url, @@ -4431,6 +4477,12 @@ async fn execute_execution_runtime_stream_inner( )?)); } + let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms(); + let remote_fallback_observation = ExecutionResponseObservation { + request_started_at_unix_ms: remote_request_started_at_unix_ms, + response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms, + request_order_id: remote_request_order_id, + }; let frame_stream = response .bytes_stream() .map_err(|err| IoError::other(err.to_string())) @@ -4452,6 +4504,7 @@ async fn execute_execution_runtime_stream_inner( provider_pool_in_flight_guard.take(), retry_scope_out.as_deref_mut(), retry_fallback_out.as_deref_mut(), + Some(remote_fallback_observation), ) .await; } @@ -5481,6 +5534,7 @@ async fn execute_stream_from_frame_stream( in_flight_guard, None, None, + None, ) .await } @@ -5503,6 +5557,7 @@ async fn execute_stream_from_frame_stream_with_retry_scope( in_flight_guard: Option, mut retry_scope_out: Option<&mut AiAttemptRetryScope>, mut retry_fallback_out: Option<&mut Option>>, + fallback_response_observation: Option, ) -> Result>, GatewayError> { let request_id = plan.request_id.as_str(); let request_id_for_log = short_request_id(request_id); @@ -5535,14 +5590,37 @@ async fn execute_stream_from_frame_stream_with_retry_scope( let StreamFramePayload::Headers { status_code, mut headers, + response_observation, } = first_frame.payload else { return Err(GatewayError::Internal( "execution runtime stream must start with headers frame".to_string(), )); }; - let mut report_context = - attach_provider_response_headers_to_report_context(report_context, &headers); + let response_observation = response_observation + .or(fallback_response_observation) + .unwrap_or(ExecutionResponseObservation { + request_started_at_unix_ms: candidate_started_unix_secs, + response_headers_observed_at_unix_ms: current_request_candidate_unix_ms(), + request_order_id: uuid::Uuid::now_v7().to_string(), + }); + let mut report_context = attach_provider_response_headers_to_report_context( + report_context, + &headers, + response_observation.request_started_at_unix_ms, + response_observation.response_headers_observed_at_unix_ms, + &response_observation.request_order_id, + ); + spawn_local_oauth_success_effect( + state.clone(), + &plan, + report_context.as_ref(), + LocalOAuthSuccessEffect { + status_code, + request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms), + request_order_id: Some(&response_observation.request_order_id), + }, + ); if status_code == 200 { seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await; if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { @@ -8310,6 +8388,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { @@ -8389,6 +8468,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { @@ -8431,6 +8511,7 @@ mod tests { None, Some(&mut retry_scope), None, + None, ) .await .expect("prefetch transport execution should resolve"); @@ -8480,6 +8561,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { @@ -8522,6 +8604,7 @@ mod tests { None, Some(&mut retry_scope), None, + None, ) .await .expect("prefetch HTTP status execution should resolve"); @@ -8680,6 +8763,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); for chunk in chunks { @@ -8725,6 +8809,7 @@ mod tests { None, Some(&mut retry_scope), Some(&mut fallback_response), + None, ) .await .expect("native Anthropic stream execution should succeed"); @@ -9364,6 +9449,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { @@ -9413,6 +9499,7 @@ mod tests { None, None, None, + None, ), ) .await @@ -9850,6 +9937,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); } @@ -11532,6 +11620,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { @@ -11660,6 +11749,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { @@ -12384,6 +12474,7 @@ mod tests { "content-type".to_string(), "text/event-stream".to_string(), )]), + response_observation: None, }, })); yield Ok::(ndjson_frame(StreamFrame { diff --git a/apps/aether-gateway/src/execution_runtime/stream_pump.rs b/apps/aether-gateway/src/execution_runtime/stream_pump.rs index ce7515ffd..dbe07ec8f 100644 --- a/apps/aether-gateway/src/execution_runtime/stream_pump.rs +++ b/apps/aether-gateway/src/execution_runtime/stream_pump.rs @@ -4,8 +4,9 @@ use std::io::Error as IoError; use std::time::{Duration, Instant}; use aether_contracts::{ - ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionStreamTerminalSummary, - ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType, + ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionResponseObservation, + ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame, StreamFramePayload, + StreamFrameType, }; use async_stream::stream; use axum::body::Bytes; @@ -44,6 +45,7 @@ pub(crate) fn build_direct_execution_frame_stream( stream_precommit_committed: _, response, started_at, + response_observation, stream_first_byte_timeout, upstream_target_permit, } = execution; @@ -108,7 +110,11 @@ pub(crate) fn build_direct_execution_frame_stream( } } - match encode_headers_frame(status_code, response_headers) { + match encode_headers_frame( + status_code, + response_headers, + &response_observation, + ) { Ok(frame) => yield Ok(frame), Err(err) => { yield Err(err); @@ -153,7 +159,11 @@ pub(crate) fn build_direct_execution_frame_stream( upstream_bytes, first_byte_timeout, }) => { - match encode_headers_frame(status_code, original_headers) { + match encode_headers_frame( + status_code, + original_headers, + &response_observation, + ) { Ok(frame) => yield Ok(frame), Err(err) => { yield Err(err); @@ -192,7 +202,11 @@ pub(crate) fn build_direct_execution_frame_stream( return; } - match encode_headers_frame(status_code, headers) { + match encode_headers_frame( + status_code, + headers, + &response_observation, + ) { Ok(frame) => yield Ok(frame), Err(err) => { yield Err(err); @@ -611,12 +625,14 @@ pub(crate) fn build_direct_execution_frame_stream( fn encode_headers_frame( status_code: u16, headers: BTreeMap, + response_observation: &ExecutionResponseObservation, ) -> Result { encode_stream_frame_ndjson(&StreamFrame { frame_type: StreamFrameType::Headers, payload: StreamFramePayload::Headers { status_code, headers, + response_observation: Some(response_observation.clone()), }, }) } @@ -1606,43 +1622,47 @@ mod tests { .expect("listener should bind"); let addr = listener.local_addr().expect("local addr should resolve"); let server = tokio::spawn(async move { - let app = Router::new().route( - "/responses", - post(|| async { - let body = serde_json::json!({ - "id": "resp_sync_bridge_123", - "object": "response", - "model": "gpt-5.4", - "status": "completed", - "output": [{ - "type": "message", - "id": "msg_sync_bridge_123", - "role": "assistant", - "content": [{ - "type": "output_text", - "text": "Hello from buffered JSON stream", - "annotations": [] - }] - }], - "usage": { - "input_tokens": 1, - "output_tokens": 2, - "total_tokens": 3 - } - }); - let mut response = axum::http::Response::new(Body::from( - serde_json::to_vec(&body).expect("json should encode"), - )); - response.headers_mut().insert( - header::CONTENT_TYPE, - HeaderValue::from_static("application/json"), - ); - response - }), - ); - axum::serve(listener, app) + let (mut socket, _) = listener.accept().await.expect("client should connect"); + let mut request = [0_u8; 4096]; + let _ = socket + .read(&mut request) .await - .expect("server should start"); + .expect("request should read"); + let body = serde_json::to_vec(&serde_json::json!({ + "id": "resp_sync_bridge_123", + "object": "response", + "model": "gpt-5.4", + "status": "completed", + "output": [{ + "type": "message", + "id": "msg_sync_bridge_123", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "Hello from buffered JSON stream", + "annotations": [] + }] + }], + "usage": { + "input_tokens": 1, + "output_tokens": 2, + "total_tokens": 3 + } + })) + .expect("json should encode"); + socket + .write_all( + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\n\r\n", + body.len() + ) + .as_bytes(), + ) + .await + .expect("headers should write"); + socket.flush().await.expect("headers should flush"); + tokio::time::sleep(Duration::from_millis(75)).await; + socket.write_all(&body).await.expect("body should write"); }); let runtime = DirectSyncExecutionRuntime::new(); @@ -1678,6 +1698,12 @@ mod tests { }) .await .expect("stream execution should succeed"); + let expected_observation = execution.response_observation.clone(); + assert!( + expected_observation.response_headers_observed_at_unix_ms + >= expected_observation.request_started_at_unix_ms + ); + assert!(!expected_observation.request_order_id.is_empty()); let frames = build_direct_execution_frame_stream(execution) .map(|item| item.expect("frame should encode")) @@ -1691,6 +1717,10 @@ mod tests { let header_frame: Value = serde_json::from_str(&frames[0]).expect("headers frame should parse"); + let encoded_observation: aether_contracts::ExecutionResponseObservation = + serde_json::from_value(header_frame["payload"]["response_observation"].clone()) + .expect("headers frame should retain the response observation"); + assert_eq!(encoded_observation, expected_observation); assert_eq!( header_frame .get("payload") diff --git a/apps/aether-gateway/src/execution_runtime/sync/execution.rs b/apps/aether-gateway/src/execution_runtime/sync/execution.rs index 5cdc6bc0d..fe30181f7 100644 --- a/apps/aether-gateway/src/execution_runtime/sync/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/sync/execution.rs @@ -5,8 +5,8 @@ use std::time::{Duration, Instant}; use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY}; use aether_contracts::{ - ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult, - ExecutionTelemetry, + ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, + ExecutionResponseObservation, ExecutionResult, ExecutionTelemetry, }; use aether_data_contracts::repository::candidates::RequestCandidateStatus; use aether_scheduler_core::{ @@ -70,11 +70,12 @@ use crate::execution_runtime::{ }; use crate::log_ids::short_request_id; use crate::orchestration::{ - apply_local_execution_effect, build_local_error_flow_metadata, trace_upstream_response_body, - with_error_flow_report_context, with_upstream_response_report_context, - LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, - LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, - LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, + apply_local_execution_effect, build_local_error_flow_metadata, + spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context, + with_upstream_response_report_context, LocalAdaptiveRateLimitEffect, + LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, + LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect, + LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard; use crate::request_candidate_runtime::{ @@ -1379,7 +1380,19 @@ async fn execute_direct_sync_runtime_candidate( candidate_started_unix_ms, event.status_code, event.ttfb_ms, - ) + ); + spawn_local_oauth_success_effect( + state_for_response_started.clone(), + plan, + report_context, + LocalOAuthSuccessEffect { + status_code: event.status_code, + request_started_at_unix_ms: Some( + event.response_observation.request_started_at_unix_ms, + ), + request_order_id: Some(&event.response_observation.request_order_id), + }, + ); }) .await .map_err(SyncExecutionFailure::from_transport); @@ -1483,12 +1496,25 @@ async fn execute_openai_image_sync_upstream_sse_candidate( OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot); progress.record_connecting().await; + let request_started_at_unix_ms = current_request_candidate_unix_ms(); + let request_order_id = uuid::Uuid::now_v7().to_string(); let response = send_request(plan, request_body) .await .map_err(SyncExecutionFailure::from_transport)?; let ttfb_ms = started_at.elapsed().as_millis() as u64; + let response_headers_observed_at_unix_ms = current_request_candidate_unix_ms(); let status_code = response.status_code(); let headers = response.headers(); + spawn_local_oauth_success_effect( + state.clone(), + plan, + report_context, + LocalOAuthSuccessEffect { + status_code, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: Some(&request_order_id), + }, + ); progress.record_response_started(status_code, ttfb_ms).await; let mut body_bytes = Vec::new(); @@ -1569,6 +1595,11 @@ async fn execute_openai_image_sync_upstream_sse_candidate( candidate_id: plan.candidate_id.clone(), status_code, headers, + response_observation: Some(ExecutionResponseObservation { + request_started_at_unix_ms, + response_headers_observed_at_unix_ms, + request_order_id, + }), body, telemetry: Some(ExecutionTelemetry { ttfb_ms: Some(ttfb_ms), @@ -2461,6 +2492,16 @@ async fn execute_execution_runtime_sync_impl( }; let mut candidate_first_byte_elapsed_ms = calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result); + let initial_response_observed_at_unix_ms = current_request_candidate_unix_ms(); + let mut provider_response_observation = + result + .response_observation + .clone() + .unwrap_or(ExecutionResponseObservation { + request_started_at_unix_ms: candidate_started_unix_secs, + response_headers_observed_at_unix_ms: initial_response_observed_at_unix_ms, + request_order_id: uuid::Uuid::now_v7().to_string(), + }); let mut oauth_retry_attempted = false; let ( result_error_type, @@ -2473,6 +2514,18 @@ async fn execute_execution_runtime_sync_impl( local_failover_response_text, local_failover_analysis, ) = loop { + spawn_local_oauth_success_effect( + state.clone(), + &plan, + report_context.as_ref(), + LocalOAuthSuccessEffect { + status_code: result.status_code, + request_started_at_unix_ms: Some( + provider_response_observation.request_started_at_unix_ms, + ), + request_order_id: Some(&provider_response_observation.request_order_id), + }, + ); let result_latency_ms = result .telemetry .as_ref() @@ -2534,10 +2587,15 @@ async fn execute_execution_runtime_sync_impl( result.status_code, local_failover_response_text.as_deref(), trace_id, + report_context.as_ref(), + Some(provider_response_observation.request_started_at_unix_ms), + Some(&provider_response_observation.request_order_id), ) .await { oauth_retry_attempted = true; + let retry_started_at_unix_ms = current_request_candidate_unix_ms(); + let retry_request_order_id = uuid::Uuid::now_v7().to_string(); match crate::execution_runtime::execute_execution_runtime_sync_plan( state, Some(trace_id), @@ -2546,6 +2604,16 @@ async fn execute_execution_runtime_sync_impl( .await { Ok(retry_result) => { + let retry_response_observed_at_unix_ms = current_request_candidate_unix_ms(); + provider_response_observation = retry_result + .response_observation + .clone() + .unwrap_or(ExecutionResponseObservation { + request_started_at_unix_ms: retry_started_at_unix_ms, + response_headers_observed_at_unix_ms: + retry_response_observed_at_unix_ms, + request_order_id: retry_request_order_id, + }); candidate_first_byte_elapsed_ms = calibrated_sync_candidate_first_byte_elapsed_ms( candidate_started_at, @@ -2594,6 +2662,13 @@ async fn execute_execution_runtime_sync_impl( local_failover_analysis, ); }; + let mut report_context = attach_provider_response_headers_to_report_context( + report_context, + &headers, + provider_response_observation.request_started_at_unix_ms, + provider_response_observation.response_headers_observed_at_unix_ms, + &provider_response_observation.request_order_id, + ); if result.status_code >= 400 { apply_local_execution_effect( state, @@ -2739,8 +2814,6 @@ async fn execute_execution_runtime_sync_impl( } let status_code = result.status_code; let has_body_bytes = body_base64.is_some(); - let mut report_context = - attach_provider_response_headers_to_report_context(report_context, &headers); if (200..300).contains(&status_code) { seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await; if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) { @@ -3231,6 +3304,8 @@ async fn execute_sync_via_remote_execution_runtime( candidate_started_unix_secs: u64, candidate_started_at: Instant, ) -> Result { + let remote_request_started_at_unix_ms = current_request_candidate_unix_ms(); + let remote_request_order_id = uuid::Uuid::now_v7().to_string(); let response = match post_sync_plan_to_remote_execution_runtime( state, remote_execution_runtime_base_url, @@ -3299,11 +3374,19 @@ async fn execute_sync_via_remote_execution_runtime( )); } - response - .json() + let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms(); + let mut result = response + .json::() .await - .map(RemoteSyncFallbackOutcome::Executed) - .map_err(|err| GatewayError::Internal(err.to_string())) + .map_err(|err| GatewayError::Internal(err.to_string()))?; + result + .response_observation + .get_or_insert(ExecutionResponseObservation { + request_started_at_unix_ms: remote_request_started_at_unix_ms, + response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms, + request_order_id: remote_request_order_id, + }); + Ok(RemoteSyncFallbackOutcome::Executed(result)) } #[cfg(test)] diff --git a/apps/aether-gateway/src/execution_runtime/transport.rs b/apps/aether-gateway/src/execution_runtime/transport.rs index 2a3786592..e17b394e3 100644 --- a/apps/aether-gateway/src/execution_runtime/transport.rs +++ b/apps/aether-gateway/src/execution_runtime/transport.rs @@ -9,12 +9,12 @@ use std::sync::{Arc, LazyLock, Mutex as StdMutex, OnceLock, RwLock as StdRwLock} use std::time::{Duration, Instant}; use aether_contracts::{ - ExecutionPlan, ExecutionResponseBodyMode, ExecutionResult, ExecutionTelemetry, ProxySnapshot, - ResolvedTransportProfile, ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, - EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER, - EXECUTION_RESPONSE_BODY_MODE_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ, - TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, - TRANSPORT_HTTP_MODE_HTTP1_ONLY, + ExecutionPlan, ExecutionResponseBodyMode, ExecutionResponseObservation, ExecutionResult, + ExecutionTelemetry, ProxySnapshot, ResolvedTransportProfile, ResponseBody, + EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, + EXECUTION_REQUEST_HTTP1_ONLY_HEADER, EXECUTION_RESPONSE_BODY_MODE_HEADER, + TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS, + TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, }; use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation; use aether_http::{apply_http_client_config, HttpClientConfig}; @@ -691,6 +691,7 @@ pub(crate) struct DirectUpstreamStreamExecution { pub(crate) stream_precommit_committed: bool, pub(crate) response: DirectUpstreamResponse, pub(crate) started_at: Instant, + pub(crate) response_observation: ExecutionResponseObservation, pub(crate) stream_first_byte_timeout: Option, pub(crate) upstream_target_permit: Option, } @@ -699,6 +700,7 @@ pub(crate) struct DirectUpstreamStreamExecution { pub(crate) struct DirectSyncResponseStarted { pub(crate) status_code: u16, pub(crate) ttfb_ms: u64, + pub(crate) response_observation: ExecutionResponseObservation, } impl DirectSyncExecutionRuntime { @@ -724,14 +726,23 @@ impl DirectSyncExecutionRuntime { let body_bytes = build_request_body(plan)?; let started_at = Instant::now(); + let request_started_at_unix_ms = crate::clock::current_unix_ms(); + let request_order_id = uuid::Uuid::now_v7().to_string(); with_non_stream_total_timeout(plan, async move { let response = send_request_inner(plan, body_bytes, false).await?; let ttfb_ms = started_at.elapsed().as_millis() as u64; + let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms(); let status_code = response.status_code(); let headers = response.headers(); + let response_observation = ExecutionResponseObservation { + request_started_at_unix_ms, + response_headers_observed_at_unix_ms, + request_order_id, + }; on_response_started(DirectSyncResponseStarted { status_code, ttfb_ms, + response_observation: response_observation.clone(), }); let (body_bytes, stream_ttfb_ms) = response.bytes_with_stream_timeout(plan, started_at).await?; @@ -752,6 +763,7 @@ impl DirectSyncExecutionRuntime { candidate_id: plan.candidate_id.clone(), status_code, headers, + response_observation: Some(response_observation), body, telemetry: Some(ExecutionTelemetry { ttfb_ms: stream_ttfb_ms.or(Some(ttfb_ms)), @@ -776,6 +788,8 @@ impl DirectSyncExecutionRuntime { ); let started_at = Instant::now(); + let request_started_at_unix_ms = crate::clock::current_unix_ms(); + let request_order_id = uuid::Uuid::now_v7().to_string(); let response = send_request(plan, body_bytes).await?; observe_gateway_stage_ms( "direct_send_headers", @@ -783,6 +797,7 @@ impl DirectSyncExecutionRuntime { ); let status_code = response.status_code(); let headers = response.headers(); + let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms(); let stream_summary_report_context = build_stream_summary_report_context(plan); @@ -797,6 +812,11 @@ impl DirectSyncExecutionRuntime { stream_precommit_committed: false, response: response.into_direct_upstream_response(), started_at, + response_observation: ExecutionResponseObservation { + request_started_at_unix_ms, + response_headers_observed_at_unix_ms, + request_order_id, + }, stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan), upstream_target_permit: None, }) @@ -834,7 +854,7 @@ pub(crate) async fn execute_sync_plan_with_report_context( } if resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).is_some() { - return execute_sync_plan_via_local_tunnel(state, plan) + return execute_sync_plan_via_local_tunnel(state, plan, report_context) .await .map_err(|err| GatewayError::Internal(err.to_string())); } @@ -857,7 +877,24 @@ pub(crate) async fn execute_sync_plan_with_report_context( Ok(None) => {} Err(err) => return Err(GatewayError::Internal(err.to_string())), } - match DirectSyncExecutionRuntime::new().execute_sync(plan).await { + let state_for_response_started = state.clone(); + match DirectSyncExecutionRuntime::new() + .execute_sync_with_response_started(plan, move |event| { + crate::orchestration::spawn_local_oauth_success_effect( + state_for_response_started, + plan, + report_context, + crate::orchestration::LocalOAuthSuccessEffect { + status_code: event.status_code, + request_started_at_unix_ms: Some( + event.response_observation.request_started_at_unix_ms, + ), + request_order_id: Some(&event.response_observation.request_order_id), + }, + ); + }) + .await + { Ok(result) => { record_manual_proxy_request_outcome(state, plan, result.status_code).await; Ok(result) @@ -889,6 +926,8 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel( plan.body.body_bytes_b64.is_some(), )?; let started_at = Instant::now(); + let request_started_at_unix_ms = crate::clock::current_unix_ms(); + let request_order_id = uuid::Uuid::now_v7().to_string(); let response = state .tunnel .open_direct_relay_stream( @@ -900,6 +939,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel( .map_err(ExecutionRuntimeTransportError::RelayError)?; let status_code = response.status(); let headers = collect_tunnel_response_headers(response.headers()); + let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms(); Ok(Some(DirectUpstreamStreamExecution { request_id: plan.request_id.clone(), @@ -912,6 +952,11 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel( stream_precommit_committed: false, response: DirectUpstreamResponse::LocalTunnel(response), started_at, + response_observation: ExecutionResponseObservation { + request_started_at_unix_ms, + response_headers_observed_at_unix_ms, + request_order_id, + }, stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan), upstream_target_permit: None, })) @@ -991,13 +1036,19 @@ fn manual_proxy_node_id(proxy: Option<&ProxySnapshot>) -> Option { async fn execute_sync_plan_via_local_tunnel( state: &AppState, plan: &ExecutionPlan, + report_context: Option<&serde_json::Value>, ) -> Result { - with_non_stream_total_timeout(plan, execute_sync_plan_via_local_tunnel_inner(state, plan)).await + with_non_stream_total_timeout( + plan, + execute_sync_plan_via_local_tunnel_inner(state, plan, report_context), + ) + .await } async fn execute_sync_plan_via_local_tunnel_inner( state: &AppState, plan: &ExecutionPlan, + report_context: Option<&serde_json::Value>, ) -> Result { let node_id = resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).ok_or_else(|| { ExecutionRuntimeTransportError::RelayError("local tunnel node unavailable".to_string()) @@ -1030,6 +1081,8 @@ async fn execute_sync_plan_via_local_tunnel_inner( "gateway execution runtime local tunnel request prepared" ); let started_at = Instant::now(); + let request_started_at_unix_ms = crate::clock::current_unix_ms(); + let request_order_id = uuid::Uuid::now_v7().to_string(); let mut response = state .tunnel .open_direct_relay_stream( @@ -1040,8 +1093,24 @@ async fn execute_sync_plan_via_local_tunnel_inner( .await .map_err(ExecutionRuntimeTransportError::RelayError)?; let ttfb_ms = started_at.elapsed().as_millis() as u64; + let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms(); let status_code = response.status(); let headers = collect_tunnel_response_headers(response.headers()); + let response_observation = ExecutionResponseObservation { + request_started_at_unix_ms, + response_headers_observed_at_unix_ms, + request_order_id, + }; + crate::orchestration::spawn_local_oauth_success_effect( + state.clone(), + plan, + report_context, + crate::orchestration::LocalOAuthSuccessEffect { + status_code, + request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms), + request_order_id: Some(&response_observation.request_order_id), + }, + ); let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-"); let (body_bytes, stream_ttfb_ms) = collect_local_tunnel_response_body(response, plan, started_at).await?; @@ -1095,6 +1164,7 @@ async fn execute_sync_plan_via_local_tunnel_inner( candidate_id: plan.candidate_id.clone(), status_code, headers, + response_observation: Some(response_observation), body, telemetry: Some(ExecutionTelemetry { ttfb_ms: stream_ttfb_ms.or(Some(ttfb_ms)), @@ -5605,6 +5675,8 @@ mod tests { ) .await .expect("headers should write"); + socket.flush().await.expect("headers should flush"); + tokio::time::sleep(std::time::Duration::from_millis(40)).await; socket .write_all(b"b\r\ndata: one\n\n\r\n") .await @@ -5634,12 +5706,34 @@ mod tests { let body = result .body + .clone() .and_then(|body| body.body_bytes_b64) .and_then(|body| base64::engine::general_purpose::STANDARD.decode(body).ok()) .expect("stream body should be captured as bytes"); let body = String::from_utf8(body).expect("stream body should be utf8"); assert!(body.contains("data: one")); assert!(body.contains("data: two")); + let observation = result + .response_observation + .expect("stream sync execution should preserve header observation"); + let telemetry = result + .telemetry + .expect("stream sync execution should include telemetry"); + let ttfb_ms = telemetry + .ttfb_ms + .expect("stream sync execution should measure the first body byte"); + assert!( + observation.response_headers_observed_at_unix_ms + >= observation.request_started_at_unix_ms + ); + assert!( + observation + .response_headers_observed_at_unix_ms + .saturating_sub(observation.request_started_at_unix_ms) + < ttfb_ms, + "header observation must not be derived from body-byte ttfb" + ); + assert!(!observation.request_order_id.is_empty()); } #[tokio::test] diff --git a/apps/aether-gateway/src/execution_runtime/windsurf.rs b/apps/aether-gateway/src/execution_runtime/windsurf.rs index 27aa3936a..98b695337 100644 --- a/apps/aether-gateway/src/execution_runtime/windsurf.rs +++ b/apps/aether-gateway/src/execution_runtime/windsurf.rs @@ -266,6 +266,7 @@ pub(crate) async fn maybe_execute_windsurf_sync( candidate_id: prepared.candidate_id, status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(body_json), body_bytes_b64: None, @@ -527,6 +528,7 @@ fn build_windsurf_stream_frame_stream( ("cache-control".to_string(), "no-cache".to_string()), ("content-type".to_string(), "text/event-stream".to_string()), ]), + response_observation: None, }, }); diff --git a/apps/aether-gateway/src/executor/orchestration.rs b/apps/aether-gateway/src/executor/orchestration.rs index eb7e59f03..366690283 100644 --- a/apps/aether-gateway/src/executor/orchestration.rs +++ b/apps/aether-gateway/src/executor/orchestration.rs @@ -1687,6 +1687,7 @@ mod tests { CONTENT_TYPE.as_str().to_string(), "application/json".to_string(), )]), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(body_json), body_bytes_b64: None, diff --git a/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/codex_reset_credit.rs b/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/codex_reset_credit.rs index 4910ea4c2..234b20753 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/codex_reset_credit.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/codex_reset_credit.rs @@ -55,6 +55,33 @@ pub(super) async fn maybe_handle( if idempotency_key.is_empty() { return Ok(Some(bad_request_response("idempotency_key 不能为空"))); } + if idempotency_key.len() > 256 { + return Ok(Some(bad_request_response( + "idempotency_key 不能超过 256 个字节", + ))); + } + let expected_credential_generation = match payload.expected_credential_generation { + serde_json::Value::Null => None, + serde_json::Value::String(value) => { + let value = value.trim().to_string(); + if value.is_empty() { + return Ok(Some(bad_request_response( + "expected_credential_generation 不能为空字符串", + ))); + } + if value.len() > 256 { + return Ok(Some(bad_request_response( + "expected_credential_generation 不能超过 256 个字节", + ))); + } + Some(value) + } + _ => { + return Ok(Some(bad_request_response( + "expected_credential_generation 必须是字符串或 null", + ))); + } + }; let Some(key) = state .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id)) @@ -93,9 +120,15 @@ pub(super) async fn maybe_handle( ))); }; - let (status, payload) = - consume_codex_reset_credit_locally(state, &provider, &endpoint, key, &idempotency_key) - .await?; + let (status, payload) = consume_codex_reset_credit_locally( + state, + &provider, + &endpoint, + key, + &idempotency_key, + expected_credential_generation.as_deref(), + ) + .await?; Ok(Some((status, Json(payload)).into_response())) } diff --git a/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs b/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs index c2200221d..66af2c12c 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs @@ -1,7 +1,10 @@ use crate::handlers::admin::admin_provider_pool_config; use crate::handlers::admin::provider::shared::paths::admin_update_key_id; use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch; -use crate::handlers::admin::provider::write::keys::admin_provider_key_update_requires_immediate_model_fetch; +use crate::handlers::admin::provider::write::keys::{ + admin_provider_key_update_requires_immediate_model_fetch, + build_provider_catalog_key_admin_cas_update, +}; use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::maintenance::ensure_provider_key_pool_scores_for_keys; use crate::provider_key_auth::provider_key_effective_api_formats; @@ -82,7 +85,25 @@ pub(super) async fn maybe_handle( Ok(record) => record, Err(detail) => return Ok(Some(bad_request_response(detail))), }; - let Some(mut updated) = state.update_provider_catalog_key(&updated_record).await? else { + let admin_update = build_provider_catalog_key_admin_cas_update( + &existing_key, + updated_record.clone(), + &provider.provider_type, + ); + if !state + .compare_and_update_provider_catalog_key_admin_state(&admin_update) + .await? + { + return Ok(Some(conflict_response( + "Key 凭据或配置已被其他请求更新，请刷新后重试", + ))); + } + let Some(mut updated) = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id)) + .await? + .into_iter() + .next() + else { return Ok(None); }; if updated_record.learned_rpm_limit != existing_key.learned_rpm_limit { @@ -183,3 +204,11 @@ fn not_found_response(detail: impl Into) -> Response { ) .into_response() } + +fn conflict_response(detail: impl Into) -> Response { + ( + http::StatusCode::CONFLICT, + Json(json!({ "detail": detail.into() })), + ) + .into_response() +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs index cc440a89d..dd3a73d56 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs @@ -22,16 +22,60 @@ use crate::handlers::admin::request::{AdminAppState, AdminRequestContext}; use crate::handlers::shared::sync_provider_key_oauth_status_snapshot; use crate::provider_key_auth::provider_key_is_oauth_managed; use crate::GatewayError; -use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthRuntimeStateCasUpdate; +use aether_data_contracts::repository::provider_catalog::{ + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, + ProviderCatalogUpstreamMetadataNamespaceExpectation, +}; use axum::{ body::{Body, Bytes}, http, response::{IntoResponse, Response}, Json, }; -use serde_json::json; +use serde_json::{json, Value}; use std::time::{SystemTime, UNIX_EPOCH}; +const CODEX_OAUTH_COMPLETE_NAMESPACE_CAS_MAX_RETRIES: usize = 3; +const CODEX_CREDENTIAL_GENERATION_KEY: &str = "credential_generation"; + +#[derive(Debug, PartialEq)] +enum CodexOAuthCompleteCasMissAction { + AlreadyCompleted, + RetryNamespace(Option), + Conflict, +} + +fn codex_oauth_complete_cas_miss_action( + latest_encrypted_auth_config: Option<&str>, + latest_upstream_metadata: Option<&Value>, + latest_status_snapshot: Option<&Value>, + expected_encrypted_auth_config: Option<&str>, + persisted_encrypted_auth_config: &str, + expected_codex_metadata_value: Option<&Value>, + replacement_codex_metadata_value: &Value, +) -> CodexOAuthCompleteCasMissAction { + let latest_codex_metadata_value = latest_upstream_metadata + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + let quota_is_cleared = latest_status_snapshot + .and_then(Value::as_object) + .and_then(|snapshot| snapshot.get("quota")) + == Some(&Value::Null); + if latest_encrypted_auth_config == Some(persisted_encrypted_auth_config) + && latest_codex_metadata_value.as_ref() == Some(replacement_codex_metadata_value) + && quota_is_cleared + { + return CodexOAuthCompleteCasMissAction::AlreadyCompleted; + } + if latest_encrypted_auth_config != expected_encrypted_auth_config + || latest_codex_metadata_value.as_ref() == expected_codex_metadata_value + { + return CodexOAuthCompleteCasMissAction::Conflict; + } + CodexOAuthCompleteCasMissAction::RetryNamespace(latest_codex_metadata_value) +} + pub(super) async fn handle_admin_provider_oauth_complete_key( state: &AdminAppState<'_>, request_context: &AdminRequestContext<'_>, @@ -277,29 +321,98 @@ pub(super) async fn handle_admin_provider_oauth_complete_key( .and_then(|snapshot| snapshot.get("oauth")) .cloned() .unwrap_or(serde_json::Value::Null); + let mut expected_codex_metadata_value = key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + let mut status_snapshot_patch = + serde_json::Map::from_iter([("oauth".to_string(), oauth_status)]); + if provider_type == "codex" { + status_snapshot_patch.insert("quota".to_string(), serde_json::Value::Null); + } let persisted_encrypted_auth_config = recovered_key .encrypted_auth_config .clone() .expect("recovered auth config should be present"); - let updated_result = state - .app() - .compare_and_update_provider_catalog_key_oauth_runtime_state( - &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - key_id: key_id.clone(), - expected_encrypted_auth_config: state_data.expected_encrypted_auth_config, - expected_credential: None, - encrypted_auth_config: persisted_encrypted_auth_config.clone(), - encrypted_api_key_update: Some(encrypted_api_key), - expires_at_unix_secs_update: Some(expires_at), - oauth_invalid_at_unix_secs: None, - oauth_invalid_reason: None, - reset_error_count: true, - upstream_metadata_patch: None, - status_snapshot_patch: json!({ "oauth": oauth_status }), - updated_at_unix_secs: Some(now_unix_secs), - }, - ) - .await; + let replacement_codex_metadata_value = json!({ + CODEX_CREDENTIAL_GENERATION_KEY: uuid::Uuid::now_v7().to_string() + }); + let expected_encrypted_auth_config = state_data.expected_encrypted_auth_config.clone(); + let updated_result: Result = async { + let max_namespace_retries = if provider_type == "codex" { + CODEX_OAUTH_COMPLETE_NAMESPACE_CAS_MAX_RETRIES + } else { + 0 + }; + for retry in 0..=max_namespace_retries { + let updated = state + .app() + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.clone(), + expected_encrypted_auth_config: expected_encrypted_auth_config.clone(), + expected_credential: None, + expected_upstream_metadata_namespace: (provider_type == "codex").then( + || ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex_metadata_value.clone(), + }, + ), + encrypted_auth_config: persisted_encrypted_auth_config.clone(), + encrypted_api_key_update: Some(encrypted_api_key.clone()), + expires_at_unix_secs_update: Some(expires_at), + oauth_invalid_at_unix_secs: None, + oauth_invalid_reason: None, + reset_error_count: true, + upstream_metadata_patch: (provider_type == "codex") + .then(|| json!({"codex": replacement_codex_metadata_value.clone()})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: serde_json::Value::Object( + status_snapshot_patch.clone(), + ), + updated_at_unix_secs: Some(now_unix_secs), + }, + ) + .await?; + if updated { + return Ok(true); + } + if provider_type != "codex" { + return Ok(false); + } + + let Some(latest_key) = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id)) + .await? + .into_iter() + .next() + else { + return Ok(false); + }; + match codex_oauth_complete_cas_miss_action( + latest_key.encrypted_auth_config.as_deref(), + latest_key.upstream_metadata.as_ref(), + latest_key.status_snapshot.as_ref(), + expected_encrypted_auth_config.as_deref(), + &persisted_encrypted_auth_config, + expected_codex_metadata_value.as_ref(), + &replacement_codex_metadata_value, + ) { + CodexOAuthCompleteCasMissAction::AlreadyCompleted => return Ok(true), + CodexOAuthCompleteCasMissAction::Conflict => return Ok(false), + CodexOAuthCompleteCasMissAction::RetryNamespace(latest_codex_metadata_value) => { + if retry == max_namespace_retries { + return Ok(false); + } + expected_codex_metadata_value = latest_codex_metadata_value; + } + } + } + Ok(false) + } + .await; let _ = state .app() .invalidate_local_oauth_refresh_entry(&key_id) @@ -397,3 +510,78 @@ pub(super) async fn handle_admin_provider_oauth_complete_key( })) .into_response()) } + +#[cfg(test)] +mod tests { + use super::{codex_oauth_complete_cas_miss_action, CodexOAuthCompleteCasMissAction}; + use serde_json::json; + + #[test] + fn codex_oauth_complete_retries_only_when_namespace_changed() { + let expected_codex = json!({"request_id": "old"}); + let replacement_codex = json!({"credential_generation": "generation-new"}); + let latest_metadata = json!({ + "codex": {"request_id": "new"}, + "unrelated": {"preserved": true} + }); + + assert_eq!( + codex_oauth_complete_cas_miss_action( + Some("old-auth"), + Some(&latest_metadata), + None, + Some("old-auth"), + "new-auth", + Some(&expected_codex), + &replacement_codex, + ), + CodexOAuthCompleteCasMissAction::RetryNamespace(Some(json!({ + "request_id": "new" + }))) + ); + assert_eq!( + codex_oauth_complete_cas_miss_action( + Some("old-auth"), + Some(&json!({"codex": expected_codex.clone()})), + None, + Some("old-auth"), + "new-auth", + Some(&expected_codex), + &replacement_codex, + ), + CodexOAuthCompleteCasMissAction::Conflict + ); + } + + #[test] + fn codex_oauth_complete_accepts_an_ambiguous_success_but_rejects_auth_rotation() { + let replacement_codex = json!({"credential_generation": "generation-new"}); + assert_eq!( + codex_oauth_complete_cas_miss_action( + Some("new-auth"), + Some(&json!({ + "codex": replacement_codex.clone(), + "unrelated": {"preserved": true} + })), + Some(&json!({"quota": null})), + Some("old-auth"), + "new-auth", + Some(&json!({"request_id": "old"})), + &replacement_codex, + ), + CodexOAuthCompleteCasMissAction::AlreadyCompleted + ); + assert_eq!( + codex_oauth_complete_cas_miss_action( + Some("other-auth"), + Some(&json!({"codex": {"request_id": "new"}})), + Some(&json!({"quota": null})), + Some("old-auth"), + "new-auth", + Some(&json!({"request_id": "old"})), + &replacement_codex, + ), + CodexOAuthCompleteCasMissAction::Conflict + ); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs index 5c3eaafa2..c27160e9b 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs @@ -12,6 +12,7 @@ use crate::ai_serving::{ build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope, }; use crate::handlers::admin::admin_provider_pool_config; +use crate::handlers::admin::provider::write::keys::build_provider_catalog_key_admin_cas_update; use crate::handlers::admin::request::AdminAppState; use crate::provider_key_auth::provider_active_api_formats; use crate::GatewayError; @@ -286,6 +287,61 @@ fn grok_oauth_catalog_key_fingerprint( grok_browser_transport_fingerprint_from_auth_config(auth_config) } +pub(crate) fn rotate_codex_credential_generation( + key: &mut StoredProviderCatalogKey, + provider_type: &str, +) { + if !provider_type.trim().eq_ignore_ascii_case("codex") { + return; + } + + let mut upstream_metadata = key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + upstream_metadata.insert( + "codex".to_string(), + json!({ + aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY: + Uuid::now_v7().to_string(), + }), + ); + key.upstream_metadata = Some(Value::Object(upstream_metadata)); + + if let Some(mut status_snapshot) = key + .status_snapshot + .as_ref() + .and_then(Value::as_object) + .cloned() + { + status_snapshot.insert("quota".to_string(), Value::Null); + key.status_snapshot = Some(Value::Object(status_snapshot)); + } +} + +pub(crate) fn ensure_codex_credential_generation_rotated( + key: &mut StoredProviderCatalogKey, + provider_type: &str, + previous_generation: Option<&str>, +) { + if !provider_type.trim().eq_ignore_ascii_case("codex") { + return; + } + + let current_generation = key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .and_then(|codex| aether_admin::provider::quota::codex_credential_generation(Some(codex))); + let already_rotated = current_generation.is_some() && current_generation != previous_generation; + if !already_rotated { + rotate_codex_credential_generation(key, provider_type); + } +} + pub(crate) async fn create_provider_oauth_catalog_key( state: &AdminAppState<'_>, provider_id: &str, @@ -344,6 +400,7 @@ pub(crate) async fn create_provider_oauth_catalog_key( record.circuit_breaker_by_format = Some(json!({})); record.created_at_unix_ms = Some(now_unix_secs); record.updated_at_unix_secs = Some(now_unix_secs); + rotate_codex_credential_generation(&mut record, provider_type); let created = state.create_provider_catalog_key(&record).await?; if let Some(key) = created.as_ref() { let _ = state @@ -395,17 +452,23 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key( updated.proxy = Some(proxy); } updated.updated_at_unix_secs = Some(now_unix_secs); - if state.update_provider_catalog_key(&updated).await?.is_none() { - return Ok(None); - } + rotate_codex_credential_generation(&mut updated, provider_type); + let admin_update = + build_provider_catalog_key_admin_cas_update(existing_key, updated.clone(), provider_type); if !state - .clear_provider_catalog_key_oauth_invalid_marker(&updated.id) + .compare_and_update_provider_catalog_key_admin_state(&admin_update) .await? { return Ok(None); } let persisted = state - .reset_provider_catalog_key_recovery_state(&updated.id) + .reset_provider_catalog_key_recovery_state_fenced( + &updated.id, + updated + .encrypted_auth_config + .as_deref() + .expect("OAuth update always supplies encrypted auth_config"), + ) .await?; if let Some(key) = persisted.as_ref() { let _ = state @@ -502,10 +565,12 @@ fn provider_oauth_catalog_key_api_formats( #[cfg(test)] mod tests { use super::{ - grok_oauth_catalog_key_fingerprint, provider_oauth_token_payload_expires_at_unix_secs, + ensure_codex_credential_generation_rotated, grok_oauth_catalog_key_fingerprint, + provider_oauth_token_payload_expires_at_unix_secs, rotate_codex_credential_generation, }; + use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; - use serde_json::json; + use serde_json::{json, Value}; fn sample_unsigned_jwt(payload: serde_json::Value) -> String { let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#); @@ -609,4 +674,101 @@ mod tests { assert!(grok_oauth_catalog_key_fingerprint("openai", auth_config).is_none()); } + + #[test] + fn codex_credential_rotation_replaces_quota_namespace_and_preserves_unrelated_state() { + let mut key = StoredProviderCatalogKey::new( + "key".to_string(), + "provider".to_string(), + "Codex".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "old-generation", + "primary_used_percent": 75.0, + }, + "unrelated": {"preserved": true}, + })); + key.status_snapshot = Some(json!({ + "oauth": {"status": "valid"}, + "quota": {"used_ratio": 0.75}, + })); + + rotate_codex_credential_generation(&mut key, "codex"); + + let codex = key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(Value::as_object) + .expect("codex namespace should exist"); + assert_eq!(codex.len(), 1); + assert_ne!( + codex + .get(aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY) + .and_then(Value::as_str), + Some("old-generation") + ); + assert_eq!( + key.upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/unrelated/preserved")), + Some(&json!(true)) + ); + assert_eq!( + key.status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.get("quota")), + Some(&Value::Null) + ); + assert_eq!( + key.status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.pointer("/oauth/status")), + Some(&json!("valid")) + ); + } + + #[test] + fn codex_credential_rotation_ensure_does_not_rotate_twice_in_one_write() { + let mut key = StoredProviderCatalogKey::new( + "key".to_string(), + "provider".to_string(), + "Codex".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.upstream_metadata = Some(json!({ + "codex": {"credential_generation": "generation-before-write"} + })); + + rotate_codex_credential_generation(&mut key, "codex"); + let builder_generation = key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/credential_generation")) + .and_then(Value::as_str) + .expect("builder should rotate the generation") + .to_string(); + + ensure_codex_credential_generation_rotated( + &mut key, + "codex", + Some("generation-before-write"), + ); + + assert_eq!( + key.upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/credential_generation")) + .and_then(Value::as_str), + Some(builder_generation.as_str()) + ); + } } diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs index 8714372a2..868ffc20f 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs @@ -470,6 +470,7 @@ mod tests { candidate_id: None, status_code: 403, headers: BTreeMap::new(), + response_observation: None, body: Some(ResponseBody { json_body: None, body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)), diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs index db00ae0b4..ea0f33418 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs @@ -18,24 +18,229 @@ use self::plan::{ execute_codex_reset_credit_plan, }; use super::shared::{ - build_quota_snapshot_payload, extract_execution_error_message, - oauth_refresh_auto_removed_result, persist_fenced_provider_quota_refresh_state, - persist_provider_quota_refresh_state, provider_auto_remove_banned_keys, + build_quota_snapshot_payload, complete_codex_account_reset, extract_execution_error_message, + oauth_refresh_auto_removed_result, persist_codex_provider_quota_refresh_state, + persist_fenced_provider_quota_refresh_state, provider_auto_remove_banned_keys, provider_auto_remove_quota_exhausted_keys, quota_key_auto_removed, - quota_refresh_success_invalid_state, should_auto_remove_oauth_invalid_key, - ProviderQuotaExecutionOutcome, + quota_refresh_success_invalid_state, reserve_codex_account_reset, + should_auto_remove_oauth_invalid_key, CodexAccountResetCompleteResult, + CodexAccountResetReserveResult, CodexAccountResetTerminal, ProviderQuotaExecutionOutcome, }; -use crate::handlers::admin::request::AdminAppState; +use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot}; use crate::provider_key_auth::provider_key_is_oauth_managed; +use crate::state::ProviderTransportCredentialFence; use crate::GatewayError; use aether_contracts::ProxySnapshot; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogKeyOAuthCredentialCasDelete, + ProviderCatalogUpstreamMetadataNamespaceExpectation, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use axum::http::StatusCode; use serde_json::{json, Map, Value}; -use std::time::{SystemTime, UNIX_EPOCH}; + +const CODEX_OAUTH_CREDENTIAL_STABILIZATION_ATTEMPTS: usize = 3; +const CODEX_RESET_QUOTA_RECONCILIATION_DELAYS_MS: [u64; 4] = [1_000, 2_000, 4_000, 8_000]; + +enum CodexOAuthRequestPreparation { + Ready { + transport: AdminGatewayProviderTransportSnapshot, + auth: (String, String), + credential_fence: ProviderTransportCredentialFence, + }, + MissingAuth, + Conflict, +} + +async fn prepare_codex_oauth_request( + state: &AdminAppState<'_>, + initial_transport: &AdminGatewayProviderTransportSnapshot, +) -> Result { + for _ in 0..CODEX_OAUTH_CREDENTIAL_STABILIZATION_ATTEMPTS { + let Some(transport) = state + .read_provider_transport_snapshot_uncached( + &initial_transport.provider.id, + &initial_transport.endpoint.id, + &initial_transport.key.id, + ) + .await? + else { + return Ok(CodexOAuthRequestPreparation::Conflict); + }; + if !crate::state::provider_transport_context_allows_credential_rotation( + initial_transport, + &transport, + ) { + return Ok(CodexOAuthRequestPreparation::Conflict); + } + let Some(before_fence) = state + .app() + .capture_provider_transport_credential_fence(&transport) + .await? + else { + continue; + }; + + let resolved_auth = state.resolve_local_oauth_header_auth(&transport).await?; + let Some(current_transport) = state + .read_provider_transport_snapshot_uncached( + &initial_transport.provider.id, + &initial_transport.endpoint.id, + &initial_transport.key.id, + ) + .await? + else { + return Ok(CodexOAuthRequestPreparation::Conflict); + }; + if !crate::state::provider_transport_context_allows_credential_rotation( + initial_transport, + ¤t_transport, + ) { + return Ok(CodexOAuthRequestPreparation::Conflict); + } + let Some(after_fence) = state + .app() + .capture_provider_transport_credential_fence(¤t_transport) + .await? + else { + continue; + }; + if before_fence != after_fence { + continue; + } + + return Ok(match resolved_auth { + Some(auth) => CodexOAuthRequestPreparation::Ready { + transport: current_transport, + auth, + credential_fence: after_fence, + }, + None => CodexOAuthRequestPreparation::MissingAuth, + }); + } + + Ok(CodexOAuthRequestPreparation::Conflict) +} + +fn codex_reset_refresh_succeeded(payload: Option<&Value>, key_id: &str) -> bool { + payload + .and_then(|payload| payload.get("results")) + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(Value::as_object) + .find(|item| item.get("key_id").and_then(Value::as_str) == Some(key_id)) + .and_then(|item| item.get("status")) + .and_then(Value::as_str) + .is_some_and(|status| status.eq_ignore_ascii_case("success")) +} + +async fn codex_reset_fence_is_still_pending( + state: &AdminAppState<'_>, + key_id: &str, + expected_credential: &ProviderTransportCredentialFence, + reset_fence: &super::shared::CodexAccountResetFence, +) -> Result { + let Some(key) = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await? + .into_iter() + .next() + else { + return Ok(false); + }; + if key.encrypted_auth_config.as_deref() + != Some(expected_credential.encrypted_auth_config.as_str()) + || key.encrypted_api_key != expected_credential.credential.encrypted_api_key + || key.auth_type != expected_credential.credential.auth_type + || key.provider_id != expected_credential.credential.provider_id + { + return Ok(false); + } + let provider_type_matches = state + .read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id)) + .await? + .into_iter() + .next() + .is_some_and(|provider| { + provider.provider_type == expected_credential.credential.provider_type + }); + if !provider_type_matches { + return Ok(false); + } + + let codex = key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .and_then(Value::as_object); + Ok(codex.is_some_and(|codex| { + codex + .get(aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY) + .and_then(Value::as_str) + == Some(reset_fence.id.as_str()) + && codex + .get(aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(aether_admin::provider::quota::coerce_json_u64) + == Some(reset_fence.generation) + && codex + .get( + aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_PENDING_GENERATION_KEY, + ) + .and_then(aether_admin::provider::quota::coerce_json_u64) + == Some(reset_fence.generation) + && codex + .get(aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY) + .and_then(Value::as_bool) + == Some(true) + })) +} + +async fn refresh_codex_quota_after_reset_until_settled( + state: &AdminAppState<'_>, + provider: &StoredProviderCatalogProvider, + endpoint: &StoredProviderCatalogEndpoint, + key: &StoredProviderCatalogKey, + reset_fence: &super::shared::CodexAccountResetFence, + expected_credential: &ProviderTransportCredentialFence, +) -> Result, GatewayError> { + let mut latest_payload = None; + for attempt in 0..=CODEX_RESET_QUOTA_RECONCILIATION_DELAYS_MS.len() { + if attempt > 0 { + tokio::time::sleep(std::time::Duration::from_millis( + CODEX_RESET_QUOTA_RECONCILIATION_DELAYS_MS[attempt - 1], + )) + .await; + } + if !codex_reset_fence_is_still_pending(state, &key.id, expected_credential, reset_fence) + .await? + { + break; + } + + let payload = refresh_codex_provider_quota_locally_with_reset_fence( + state, + provider, + endpoint, + vec![key.clone()], + None, + Some(reset_fence.id.as_str()), + Some(reset_fence.generation), + Some(expected_credential), + ) + .await?; + let refresh_succeeded = codex_reset_refresh_succeeded(payload.as_ref(), &key.id); + latest_payload = payload; + if !refresh_succeeded + || !codex_reset_fence_is_still_pending(state, &key.id, expected_credential, reset_fence) + .await? + { + break; + } + } + Ok(latest_payload) +} fn merge_codex_quota_metadata( header_metadata: Option<&serde_json::Value>, @@ -53,6 +258,26 @@ fn merge_codex_quota_metadata( serde_json::Value::Object(merged) } +fn codex_quota_window_coverage( + body_json: Option<&Value>, +) -> aether_admin::provider::quota::CodexQuotaWindowCoverage { + let body = body_json.and_then(Value::as_object); + let has_account_snapshot = body + .and_then(|body| body.get("rate_limit")) + .and_then(Value::as_object) + .is_some(); + let has_spark_snapshot = body + .and_then(|body| body.get("additional_rate_limits")) + .and_then(Value::as_array) + .is_some(); + + match (has_account_snapshot, has_spark_snapshot) { + (true, true) => aether_admin::provider::quota::CodexQuotaWindowCoverage::FullSnapshot, + (true, false) => aether_admin::provider::quota::CodexQuotaWindowCoverage::AccountSnapshot, + _ => aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch, + } +} + fn truncate_codex_reset_credit_detail_error(message: impl Into) -> String { let message = message.into(); let mut sanitized = message.replace('\n', " "); @@ -198,6 +423,10 @@ fn codex_consume_success_status(outcome: &str) -> &'static str { } } +fn codex_reset_credit_outcome_allows_usage_drop(outcome: &str) -> bool { + matches!(outcome, "reset" | "already_redeemed") +} + fn codex_extract_refresh_result_fields( refresh_payload: Option<&Value>, key_id: &str, @@ -247,12 +476,74 @@ fn codex_extract_refresh_result_fields( ) } +async fn finish_codex_reset_replay( + state: &AdminAppState<'_>, + provider: &StoredProviderCatalogProvider, + endpoint: &StoredProviderCatalogEndpoint, + key: &StoredProviderCatalogKey, + credential: &ProviderTransportCredentialFence, + terminal: CodexAccountResetTerminal, +) -> Result<(StatusCode, Value), GatewayError> { + let mut refresh_status = "skipped".to_string(); + let mut refresh_error = None; + let mut metadata = None; + let mut quota_snapshot = None; + if codex_reset_credit_outcome_allows_usage_drop(&terminal.outcome) { + let fence = super::shared::CodexAccountResetFence { + unix_ms: crate::clock::current_unix_ms(), + id: format!("reset:{}", terminal.idempotency_key), + generation: terminal.generation, + }; + if codex_reset_fence_is_still_pending(state, &key.id, credential, &fence).await? { + match refresh_codex_quota_after_reset_until_settled( + state, provider, endpoint, key, &fence, credential, + ) + .await + { + Ok(payload) => { + (refresh_status, refresh_error, metadata, quota_snapshot) = + codex_extract_refresh_result_fields(payload.as_ref(), &key.id); + } + Err(err) => { + refresh_status = "failed".to_string(); + refresh_error = + Some(truncate_codex_reset_credit_detail_error(err.into_message())); + } + } + } + } + let mut payload = Map::new(); + payload.insert("key_id".to_string(), json!(key.id)); + payload.insert( + "status".to_string(), + json!(codex_consume_success_status(&terminal.outcome)), + ); + payload.insert("outcome".to_string(), json!(terminal.outcome)); + payload.insert( + "idempotency_key".to_string(), + json!(terminal.idempotency_key), + ); + payload.insert("replay".to_string(), json!(true)); + payload.insert("refresh_status".to_string(), json!(refresh_status)); + if let Some(refresh_error) = refresh_error { + payload.insert("refresh_error".to_string(), json!(refresh_error)); + } + if let Some(metadata) = metadata { + payload.insert("metadata".to_string(), metadata); + } + if let Some(quota_snapshot) = quota_snapshot { + payload.insert("quota_snapshot".to_string(), quota_snapshot); + } + Ok((StatusCode::OK, Value::Object(payload))) +} + pub(crate) async fn consume_codex_reset_credit_locally( state: &AdminAppState<'_>, provider: &StoredProviderCatalogProvider, endpoint: &StoredProviderCatalogEndpoint, key: StoredProviderCatalogKey, idempotency_key: &str, + expected_credential_generation: Option<&str>, ) -> Result<(StatusCode, Value), GatewayError> { let transport = match state .read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id) @@ -273,22 +564,48 @@ pub(crate) async fn consume_codex_reset_credit_locally( }; let is_oauth_managed = provider_key_is_oauth_managed(&key, provider.provider_type.as_str()); - let resolved_oauth_auth = if is_oauth_managed { - state.resolve_local_oauth_header_auth(&transport).await? - } else { - None - }; - if is_oauth_managed && resolved_oauth_auth.is_none() { + if !is_oauth_managed { return Ok(( StatusCode::BAD_REQUEST, json!({ "key_id": key.id, "status": "error", "outcome": "error", - "message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token", + "message": "Codex reset credit 仅支持 OAuth 托管账号", }), )); } + let (transport, resolved_oauth_auth, reset_credential_fence) = + match prepare_codex_oauth_request(state, &transport).await? { + CodexOAuthRequestPreparation::Ready { + transport, + auth, + credential_fence, + } => (transport, Some(auth), credential_fence), + CodexOAuthRequestPreparation::MissingAuth => { + return Ok(( + StatusCode::BAD_REQUEST, + json!({ + "key_id": key.id, + "status": "error", + "outcome": "error", + "message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token", + }), + )); + } + CodexOAuthRequestPreparation::Conflict => { + return Ok(( + StatusCode::CONFLICT, + json!({ + "key_id": key.id, + "status": "error", + "outcome": "error", + "idempotency_key": idempotency_key, + "message": "Codex credential changed before reset credit could be consumed", + }), + )); + } + }; let request_spec = match build_codex_reset_credit_consume_request_spec( &transport, @@ -309,6 +626,79 @@ pub(crate) async fn consume_codex_reset_credit_locally( } }; + let reservation = match reserve_codex_account_reset( + state, + &key.id, + reset_credential_fence.encrypted_auth_config.as_str(), + &reset_credential_fence.credential, + expected_credential_generation, + idempotency_key, + ) + .await? + { + Some(CodexAccountResetReserveResult::Reserved(reservation)) => reservation, + Some(CodexAccountResetReserveResult::Replay(terminal)) => { + return finish_codex_reset_replay( + state, + provider, + endpoint, + &key, + &reset_credential_fence, + terminal, + ) + .await; + } + Some(CodexAccountResetReserveResult::LegacyReplay) => { + return Ok(( + StatusCode::OK, + json!({ + "key_id": key.id, + "status": "success", + "outcome": "historical_replay", + "idempotency_key": idempotency_key, + "refresh_status": "skipped", + }), + )); + } + Some(CodexAccountResetReserveResult::Busy(active)) => { + return Ok(( + StatusCode::CONFLICT, + json!({ + "key_id": key.id, + "status": "error", + "outcome": "busy", + "idempotency_key": idempotency_key, + "active_idempotency_key": active.idempotency_key, + "message": "Another Codex reset credit operation is unresolved", + }), + )); + } + Some(CodexAccountResetReserveResult::CredentialGenerationMismatch) => { + return Ok(( + StatusCode::CONFLICT, + json!({ + "key_id": key.id, + "status": "error", + "outcome": "credential_changed", + "idempotency_key": idempotency_key, + "message": "Codex credential changed since this reset request was prepared", + }), + )); + } + None => { + return Ok(( + StatusCode::CONFLICT, + json!({ + "key_id": key.id, + "status": "error", + "outcome": "error", + "idempotency_key": idempotency_key, + "message": "Codex reset reservation could not be persisted", + }), + )); + } + }; + let result = match execute_codex_reset_credit_plan(state, &transport, request_spec, None).await? { ProviderQuotaExecutionOutcome::Response(result) => result, @@ -331,11 +721,11 @@ pub(crate) async fn consume_codex_reset_credit_locally( .and_then(|body| body.json_body.as_ref()); let outcome = normalize_codex_reset_credit_consume_outcome(body_json) .unwrap_or_else(|| "unknown".to_string()); - let known_non_error_outcome = matches!( + let known_terminal_outcome = matches!( outcome.as_str(), "reset" | "already_redeemed" | "nothing_to_reset" | "no_credit" ); - if result.status_code >= 400 && !known_non_error_outcome { + if !known_terminal_outcome { let detail = extract_execution_error_message(&result) .unwrap_or_else(|| format!("HTTP {}", result.status_code)); return Ok(( @@ -345,19 +735,88 @@ pub(crate) async fn consume_codex_reset_credit_locally( "status": "error", "outcome": "error", "idempotency_key": idempotency_key, - "message": format!("reset credit consume 返回状态码 {}: {detail}", result.status_code), + "message": format!("reset credit consume outcome is ambiguous: {detail}"), "status_code": result.status_code, }), )); } - let (refresh_status, refresh_error, metadata, quota_snapshot) = - match refresh_codex_provider_quota_locally( + let fence_unix_ms = result + .response_observation + .as_ref() + .map(|observation| observation.response_headers_observed_at_unix_ms) + .unwrap_or_else(crate::clock::current_unix_ms); + let Some(completed) = complete_codex_account_reset( + state, + &key.id, + reset_credential_fence.encrypted_auth_config.as_str(), + &reset_credential_fence.credential, + &reservation, + &outcome, + fence_unix_ms, + ) + .await? + else { + return Ok(( + StatusCode::CONFLICT, + json!({ + "key_id": key.id, + "status": "error", + "outcome": "error", + "idempotency_key": idempotency_key, + "message": "Codex reset completion could not be persisted", + }), + )); + }; + + let (effective_outcome, reset_fence) = match completed { + CodexAccountResetCompleteResult::Activated(fence) => (outcome.clone(), Some(fence)), + CodexAccountResetCompleteResult::Noop(terminal) + | CodexAccountResetCompleteResult::Replay(terminal) => { + let fence = + codex_reset_credit_outcome_allows_usage_drop(&terminal.outcome).then(|| { + super::shared::CodexAccountResetFence { + unix_ms: fence_unix_ms, + id: format!("reset:{}", terminal.idempotency_key), + generation: terminal.generation, + } + }); + (terminal.outcome, fence) + } + }; + + let (refresh_status, refresh_error, metadata, quota_snapshot) = match reset_fence.as_ref() { + Some(reset_fence) => { + match refresh_codex_quota_after_reset_until_settled( + state, + provider, + endpoint, + &key, + reset_fence, + &reset_credential_fence, + ) + .await + { + Ok(refresh_payload) => { + codex_extract_refresh_result_fields(refresh_payload.as_ref(), &key.id) + } + Err(err) => ( + "failed".to_string(), + Some(truncate_codex_reset_credit_detail_error(err.into_message())), + None, + None, + ), + } + } + None => match refresh_codex_provider_quota_locally_with_reset_fence( state, provider, endpoint, vec![key.clone()], None, + None, + None, + Some(&reset_credential_fence), ) .await { @@ -370,15 +829,16 @@ pub(crate) async fn consume_codex_reset_credit_locally( None, None, ), - }; + }, + }; let mut payload = Map::new(); payload.insert("key_id".to_string(), json!(key.id)); payload.insert( "status".to_string(), - json!(codex_consume_success_status(&outcome)), + json!(codex_consume_success_status(&effective_outcome)), ); - payload.insert("outcome".to_string(), json!(outcome)); + payload.insert("outcome".to_string(), json!(effective_outcome)); payload.insert("idempotency_key".to_string(), json!(idempotency_key)); payload.insert("refresh_status".to_string(), json!(refresh_status)); if let Some(refresh_error) = refresh_error { @@ -400,6 +860,29 @@ pub(crate) async fn refresh_codex_provider_quota_locally( endpoint: &StoredProviderCatalogEndpoint, keys: Vec, proxy_override: Option, +) -> Result, GatewayError> { + refresh_codex_provider_quota_locally_with_reset_fence( + state, + provider, + endpoint, + keys, + proxy_override, + None, + None, + None, + ) + .await +} + +async fn refresh_codex_provider_quota_locally_with_reset_fence( + state: &AdminAppState<'_>, + provider: &StoredProviderCatalogProvider, + endpoint: &StoredProviderCatalogEndpoint, + keys: Vec, + proxy_override: Option, + account_reset_fence_id: Option<&str>, + authoritative_reset_generation: Option, + expected_reset_credential: Option<&crate::state::ProviderTransportCredentialFence>, ) -> Result, GatewayError> { let mut results = Vec::new(); let mut success_count = 0usize; @@ -412,7 +895,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally( for key in keys { let had_oauth_refresh_issue = codex_oauth_refresh_issue_reason(key.oauth_invalid_reason.as_deref()); - let transport = match state + let initial_transport = match state .read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id) .await? { @@ -429,48 +912,70 @@ pub(crate) async fn refresh_codex_provider_quota_locally( } }; let is_oauth_managed = provider_key_is_oauth_managed(&key, provider.provider_type.as_str()); - let quota_auth_config_fence = if is_oauth_managed { - match state - .app() - .capture_provider_transport_auth_config_fence(&transport) - .await? - { - Some(ciphertext) => Some(ciphertext), - None => { + let (transport, resolved_oauth_auth, quota_credential_fence) = if is_oauth_managed { + match prepare_codex_oauth_request(state, &initial_transport).await? { + CodexOAuthRequestPreparation::Ready { + transport, + auth, + credential_fence, + } => (transport, Some(auth), Some(credential_fence)), + CodexOAuthRequestPreparation::MissingAuth => { failed_count += 1; results.push(json!({ "key_id": key.id, "key_name": key.name, "status": "error", - "message": "OAuth credential changed before quota refresh", + "message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token", })); continue; } + CodexOAuthRequestPreparation::Conflict => { + if quota_key_auto_removed(state, &key.id).await? { + auto_removed_count += 1; + results.push(oauth_refresh_auto_removed_result(&key)); + } else { + failed_count += 1; + results.push(json!({ + "key_id": key.id, + "key_name": key.name, + "status": "error", + "message": "OAuth credential changed before quota refresh", + })); + } + continue; + } } } else { - None + (initial_transport, None, None) }; - - let resolved_oauth_auth = if is_oauth_managed { - state.resolve_local_oauth_header_auth(&transport).await? - } else { - None - }; - if is_oauth_managed && quota_key_auto_removed(state, &key.id).await? { - auto_removed_count += 1; - results.push(oauth_refresh_auto_removed_result(&key)); - continue; - } - if is_oauth_managed && resolved_oauth_auth.is_none() { - failed_count += 1; - results.push(json!({ - "key_id": key.id, - "key_name": key.name, - "status": "error", - "message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token", - })); - continue; + if let Some(expected_reset_credential) = expected_reset_credential { + if quota_credential_fence.as_ref() != Some(expected_reset_credential) { + failed_count += 1; + results.push(json!({ + "key_id": key.id, + "key_name": key.name, + "status": "error", + "message": "Codex credential changed after reset credit was consumed", + })); + continue; + } } + let transport_codex_metadata = transport + .key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")); + let observed_reset_generation = authoritative_reset_generation.or_else(|| { + Some( + aether_admin::provider::quota::codex_quota_account_reset_generation( + transport_codex_metadata, + ), + ) + }); + let observed_credential_generation = + aether_admin::provider::quota::codex_credential_generation(transport_codex_metadata) + .map(ToOwned::to_owned); let request_spec = match build_codex_quota_request_spec(&transport, resolved_oauth_auth.clone()) { @@ -487,6 +992,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally( } }; + let quota_request_fallback_started_at_unix_ms = crate::clock::current_unix_ms(); + let quota_request_fallback_order_id = uuid::Uuid::now_v7().to_string(); let result = match execute_codex_quota_plan( state, &transport, @@ -508,16 +1015,25 @@ pub(crate) async fn refresh_codex_provider_quota_locally( continue; } }; - let now_unix_secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .ok() - .map(|duration| duration.as_secs()) - .unwrap_or(0); + let quota_response_fallback_observed_at_unix_ms = crate::clock::current_unix_ms(); + let quota_response_observation = result.response_observation.as_ref(); + let quota_request_started_at_unix_ms = quota_response_observation + .map(|observation| observation.request_started_at_unix_ms) + .unwrap_or(quota_request_fallback_started_at_unix_ms); + let quota_response_observed_at_unix_ms = quota_response_observation + .map(|observation| observation.response_headers_observed_at_unix_ms) + .unwrap_or(quota_response_fallback_observed_at_unix_ms); + let quota_request_order_id = quota_response_observation + .map(|observation| observation.request_order_id.as_str()) + .unwrap_or(quota_request_fallback_order_id.as_str()); + let now_unix_secs = quota_response_observed_at_unix_ms / 1_000; let header_metadata = parse_codex_usage_headers(&result.headers, now_unix_secs); let mut metadata_update = header_metadata .as_ref() .map(|metadata| json!({ "codex": metadata })); + let mut quota_window_coverage = + aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch; let (mut oauth_invalid_at_unix_secs, mut oauth_invalid_reason) = (None, None); let mut status = "error".to_string(); let mut message = None::; @@ -544,6 +1060,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally( now_unix_secs, ) .await?; + quota_window_coverage = codex_quota_window_coverage(Some(body_json)); metadata_update = Some(json!({ "codex": codex_metadata })); @@ -586,6 +1103,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally( } 402 => { if codex_looks_like_workspace_deactivated(err_msg.as_deref()) { + quota_window_coverage = + aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch; let mut codex_meta = metadata_update .as_ref() .and_then(|value| value.get("codex")) @@ -627,6 +1146,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally( oauth_invalid_reason = reason; status = "workspace_deactivated".to_string(); } else { + quota_window_coverage = + aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch; let plan_type = transport .key .decrypted_auth_config @@ -667,24 +1188,45 @@ pub(crate) async fn refresh_codex_provider_quota_locally( } } - let persisted = if let Some(expected_auth_config) = quota_auth_config_fence.as_deref() { + let persisted = if let Some(expected_credential) = quota_credential_fence.as_ref() { persist_fenced_provider_quota_refresh_state( state, &key.id, - expected_auth_config, + expected_credential.encrypted_auth_config.as_str(), metadata_update.as_ref(), oauth_invalid_at_unix_secs, oauth_invalid_reason.clone(), + aether_admin::provider::quota::CodexQuotaMergeContext { + observed_at_unix_secs: now_unix_secs, + request_started_at_unix_ms: Some(quota_request_started_at_unix_ms), + request_order_id: Some(quota_request_order_id), + observed_reset_generation, + authoritative_reset_generation, + observed_credential_generation: observed_credential_generation.as_deref(), + account_reset_fence_id, + coverage: quota_window_coverage, + }, + Some(&expected_credential.credential), ) .await? } else { - persist_provider_quota_refresh_state( + persist_codex_provider_quota_refresh_state( state, &key.id, metadata_update.as_ref(), oauth_invalid_at_unix_secs, oauth_invalid_reason.clone(), None, + aether_admin::provider::quota::CodexQuotaMergeContext { + observed_at_unix_secs: now_unix_secs, + request_started_at_unix_ms: Some(quota_request_started_at_unix_ms), + request_order_id: Some(quota_request_order_id), + observed_reset_generation, + authoritative_reset_generation, + observed_credential_generation: observed_credential_generation.as_deref(), + account_reset_fence_id, + coverage: quota_window_coverage, + }, ) .await? }; @@ -698,26 +1240,57 @@ pub(crate) async fn refresh_codex_provider_quota_locally( })); continue; } - let credential_cas_delete = quota_auth_config_fence.as_ref().map(|auth_config| { + let persisted_key = state + .read_provider_catalog_keys_by_ids(&[key.id.clone()]) + .await? + .into_iter() + .next(); + let persisted_codex_metadata = persisted_key + .as_ref() + .and_then(|key| key.upstream_metadata.as_ref()) + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + if let Some(codex_metadata) = persisted_codex_metadata.as_ref() { + metadata_update = Some(json!({"codex": codex_metadata})); + } + let persisted_codex_object = persisted_codex_metadata + .as_ref() + .and_then(serde_json::Value::as_object); + let request_owns_persisted_oauth_state = quota_credential_fence.is_none() + || (persisted_codex_object + .and_then(|codex| codex.get("oauth_state_request_started_at_unix_ms")) + .and_then(aether_admin::provider::quota::coerce_json_u64) + == Some(quota_request_started_at_unix_ms) + && persisted_codex_object + .and_then(|codex| codex.get("oauth_state_request_id")) + .and_then(serde_json::Value::as_str) + == Some(quota_request_order_id)); + let credential_cas_delete = quota_credential_fence.as_ref().map(|credential_fence| { ProviderCatalogKeyOAuthCredentialCasDelete { key_id: key.id.clone(), - expected_encrypted_auth_config: Some(auth_config.clone()), - expected_credential: ProviderCatalogKeyOAuthCredentialFence { - encrypted_api_key: key.encrypted_api_key.clone(), - auth_type: key.auth_type.clone(), - provider_id: key.provider_id.clone(), - provider_type: provider.provider_type.clone(), - }, + expected_encrypted_auth_config: Some( + credential_fence.encrypted_auth_config.clone(), + ), + expected_credential: credential_fence.credential.clone(), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: persisted_codex_metadata.clone(), + }, + ), } }); - let should_auto_remove_hard_banned = - provider_auto_remove_banned_keys(provider.config.as_ref()) - && should_auto_remove_oauth_invalid_key( - &key, - oauth_invalid_reason.as_deref(), - matches!(status_code, Some(401 | 403)), - now_unix_secs, - ); + let should_auto_remove_hard_banned = request_owns_persisted_oauth_state + && provider_auto_remove_banned_keys(provider.config.as_ref()) + && should_auto_remove_oauth_invalid_key( + persisted_key.as_ref().unwrap_or(&key), + persisted_key + .as_ref() + .and_then(|key| key.oauth_invalid_reason.as_deref()), + matches!(status_code, Some(401 | 403)), + now_unix_secs, + ); let auto_removed_hard_banned = if should_auto_remove_hard_banned { match credential_cas_delete.as_ref() { Some(delete) => { @@ -735,6 +1308,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally( auto_removed_hard_banned_count += 1; } let auto_removed_quota_exhausted = if !auto_removed_hard_banned + && request_owns_persisted_oauth_state && status == "quota_exhausted" && provider_auto_remove_quota_exhausted_keys(provider.config.as_ref()) { @@ -798,7 +1372,9 @@ pub(crate) async fn refresh_codex_provider_quota_locally( } if let Some(quota_snapshot) = build_quota_snapshot_payload( "codex", - key.status_snapshot.as_ref(), + persisted_key + .as_ref() + .and_then(|key| key.status_snapshot.as_ref()), metadata_update.as_ref(), ) { payload.insert("quota_snapshot".to_string(), quota_snapshot); @@ -866,6 +1442,19 @@ mod tests { ); } + #[test] + fn codex_reset_credit_only_allows_usage_drop_after_confirmed_redemption() { + assert!(codex_reset_credit_outcome_allows_usage_drop("reset")); + assert!(codex_reset_credit_outcome_allows_usage_drop( + "already_redeemed" + )); + assert!(!codex_reset_credit_outcome_allows_usage_drop( + "nothing_to_reset" + )); + assert!(!codex_reset_credit_outcome_allows_usage_drop("no_credit")); + assert!(!codex_reset_credit_outcome_allows_usage_drop("unknown")); + } + #[test] fn codex_reset_credit_detail_failure_records_attempt_time() { let mut metadata = Map::new(); @@ -880,4 +1469,29 @@ mod tests { Some(&json!(1_777_000_000u64)) ); } + + #[test] + fn codex_quota_coverage_only_replaces_observed_window_families() { + assert_eq!( + codex_quota_window_coverage(Some(&json!({"credits":{"balance":5}}))), + aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch + ); + assert_eq!( + codex_quota_window_coverage(None), + aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch + ); + assert_eq!( + codex_quota_window_coverage(Some(&json!({ + "rate_limit":{"primary_window":{}} + }))), + aether_admin::provider::quota::CodexQuotaWindowCoverage::AccountSnapshot + ); + assert_eq!( + codex_quota_window_coverage(Some(&json!({ + "rate_limit":{"primary_window":{}}, + "additional_rate_limits":[] + }))), + aether_admin::provider::quota::CodexQuotaWindowCoverage::FullSnapshot + ); + } } diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/grok.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/grok.rs index 0604412be..2a9b707c3 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/grok.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/grok.rs @@ -796,6 +796,7 @@ mod tests { candidate_id: None, status_code: 403, headers: BTreeMap::new(), + response_observation: None, body: Some(ResponseBody { json_body: None, body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)), diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs index 822e86c9f..6c5e55937 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs @@ -14,8 +14,9 @@ use aether_contracts::{ ResolvedTransportProfile, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, }; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, - ProviderCatalogKeyStatusSnapshotUpdate, StoredProviderCatalogEndpoint, + ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, + ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, + ProviderCatalogUpstreamMetadataNamespaceExpectation, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, }; use aether_provider_pool::{ProviderPoolQuotaRequestSpec, ProviderPoolService}; @@ -24,6 +25,54 @@ use tracing::warn; const PROVIDER_QUOTA_DEFAULT_TIMEOUT_MS: u64 = 30_000; const PROVIDER_QUOTA_PROXY_TIMEOUT_MS: u64 = 60_000; +const CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS: usize = 16; +const CODEX_RESET_HISTORY_LIMIT: usize = 64; +const CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY: &str = + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY; +const CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY: &str = + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CodexAccountResetFence { + pub unix_ms: u64, + pub id: String, + pub generation: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CodexAccountResetFenceInstall { + Owned(CodexAccountResetFence), + Superseded, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CodexAccountResetReservation { + pub idempotency_key: String, + pub generation: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct CodexAccountResetTerminal { + pub idempotency_key: String, + pub generation: u64, + pub outcome: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CodexAccountResetReserveResult { + Reserved(CodexAccountResetReservation), + Replay(CodexAccountResetTerminal), + LegacyReplay, + Busy(CodexAccountResetReservation), + CredentialGenerationMismatch, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum CodexAccountResetCompleteResult { + Activated(CodexAccountResetFence), + Noop(CodexAccountResetTerminal), + Replay(CodexAccountResetTerminal), +} pub(super) enum ProviderQuotaExecutionOutcome { Response(ExecutionResult), @@ -211,6 +260,602 @@ pub(super) fn quota_refresh_success_invalid_state( admin_provider_quota_pure::quota_refresh_success_invalid_state(key) } +fn merge_codex_oauth_response_state( + latest_key: &StoredProviderCatalogKey, + incoming_invalid_at_unix_secs: Option, + incoming_invalid_reason: Option<&str>, + observed_at_unix_secs: u64, +) -> (Option, Option) { + match incoming_invalid_reason { + Some(reason) => admin_provider_quota_pure::codex_build_invalid_state( + latest_key, + reason.to_string(), + incoming_invalid_at_unix_secs.unwrap_or(observed_at_unix_secs), + ), + None => admin_provider_quota_pure::quota_refresh_success_invalid_state(latest_key), + } +} + +fn codex_reset_credential_matches( + key: &StoredProviderCatalogKey, + expected_encrypted_auth_config: &str, + expected_credential: &ProviderCatalogKeyOAuthCredentialFence, +) -> bool { + key.encrypted_auth_config.as_deref() == Some(expected_encrypted_auth_config) + && key.encrypted_api_key == expected_credential.encrypted_api_key + && key.auth_type == expected_credential.auth_type + && key.provider_id == expected_credential.provider_id +} + +fn codex_reset_reservation_from_object( + codex: &serde_json::Map, +) -> Option { + let reservation = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY)? + .as_object()?; + let idempotency_key = reservation.get("idempotency_key")?.as_str()?.trim(); + let generation = reservation + .get("generation") + .and_then(admin_provider_quota_pure::coerce_json_u64) + .filter(|generation| *generation > 0)?; + (!idempotency_key.is_empty()).then(|| CodexAccountResetReservation { + idempotency_key: idempotency_key.to_string(), + generation, + }) +} + +fn codex_reset_history_from_object( + codex: &serde_json::Map, +) -> Vec { + codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_HISTORY_KEY) + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(|value| { + let object = value.as_object()?; + let idempotency_key = object.get("idempotency_key")?.as_str()?.trim(); + let generation = object + .get("generation") + .and_then(admin_provider_quota_pure::coerce_json_u64)?; + let outcome = object.get("outcome")?.as_str()?.trim(); + (!idempotency_key.is_empty() && !outcome.is_empty()).then(|| { + CodexAccountResetTerminal { + idempotency_key: idempotency_key.to_string(), + generation, + outcome: outcome.to_string(), + } + }) + }) + .collect() +} + +fn codex_reset_write_bounded_history( + codex: &mut serde_json::Map, + terminal: &CodexAccountResetTerminal, +) { + let mut history = codex_reset_history_from_object(codex); + history.retain(|entry| entry.idempotency_key != terminal.idempotency_key); + history.push(terminal.clone()); + if history.len() > CODEX_RESET_HISTORY_LIMIT { + history.drain(..history.len() - CODEX_RESET_HISTORY_LIMIT); + } + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_HISTORY_KEY.to_string(), + serde_json::Value::Array( + history + .into_iter() + .map(|entry| { + serde_json::json!({ + "idempotency_key": entry.idempotency_key, + "generation": entry.generation, + "outcome": entry.outcome, + }) + }) + .collect(), + ), + ); + + codex_reset_write_processed_id(codex, &terminal.idempotency_key); +} + +fn codex_reset_write_processed_id( + codex: &mut serde_json::Map, + idempotency_key: &str, +) -> bool { + let mut processed_ids = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_PROCESSED_IDS_KEY) + .and_then(serde_json::Value::as_array) + .into_iter() + .flatten() + .filter_map(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) + .collect::>(); + let already_processed = processed_ids.iter().any(|value| value == idempotency_key); + processed_ids.retain(|value| value != idempotency_key); + processed_ids.push(idempotency_key.to_string()); + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_PROCESSED_IDS_KEY.to_string(), + serde_json::Value::Array( + processed_ids + .into_iter() + .map(serde_json::Value::String) + .collect(), + ), + ); + already_processed +} + +async fn persist_codex_reset_namespace( + state: &AdminAppState<'_>, + key: &StoredProviderCatalogKey, + expected_codex: Option, + next_codex: serde_json::Value, + expected_encrypted_auth_config: &str, + expected_credential: &ProviderCatalogKeyOAuthCredentialFence, + updated_at_unix_secs: u64, +) -> Result { + state + .app() + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key.id.clone(), + expected_encrypted_auth_config: Some(expected_encrypted_auth_config.to_string()), + expected_credential: Some(expected_credential.clone()), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex, + }, + ), + encrypted_auth_config: expected_encrypted_auth_config.to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: key.oauth_invalid_reason.clone(), + upstream_metadata_patch: Some(serde_json::json!({"codex": next_codex})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: serde_json::json!({}), + reset_error_count: false, + updated_at_unix_secs: Some(updated_at_unix_secs), + }, + ) + .await +} + +pub(crate) async fn reserve_codex_account_reset( + state: &AdminAppState<'_>, + key_id: &str, + expected_encrypted_auth_config: &str, + expected_credential: &ProviderCatalogKeyOAuthCredentialFence, + expected_credential_generation: Option<&str>, + idempotency_key: &str, +) -> Result, GatewayError> { + let idempotency_key = idempotency_key.trim(); + if idempotency_key.is_empty() { + return Ok(None); + } + for attempt in 0..CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let Some(key) = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await? + .into_iter() + .next() + else { + return Ok(None); + }; + if !codex_reset_credential_matches( + &key, + expected_encrypted_auth_config, + expected_credential, + ) { + return Ok(None); + } + let expected_codex = key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + if !admin_provider_quota_pure::codex_credential_generation_matches( + expected_codex.as_ref(), + expected_credential_generation, + ) { + return Ok(Some( + CodexAccountResetReserveResult::CredentialGenerationMismatch, + )); + } + let mut codex = expected_codex + .as_ref() + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + if let Some(terminal) = codex_reset_history_from_object(&codex) + .into_iter() + .find(|entry| entry.idempotency_key == idempotency_key) + { + return Ok(Some(CodexAccountResetReserveResult::Replay(terminal))); + } + if codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_PROCESSED_IDS_KEY) + .and_then(serde_json::Value::as_array) + .is_some_and(|ids| { + ids.iter() + .any(|value| value.as_str() == Some(idempotency_key)) + }) + { + let active_generation = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .unwrap_or(0); + let active_fence_matches = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY) + .and_then(serde_json::Value::as_str) + == Some(format!("reset:{idempotency_key}").as_str()); + if active_generation > 0 && active_fence_matches { + return Ok(Some(CodexAccountResetReserveResult::Replay( + CodexAccountResetTerminal { + idempotency_key: idempotency_key.to_string(), + generation: active_generation, + outcome: "already_redeemed".to_string(), + }, + ))); + } + return Ok(Some(CodexAccountResetReserveResult::LegacyReplay)); + } + if let Some(reservation) = codex_reset_reservation_from_object(&codex) { + return Ok(Some(if reservation.idempotency_key == idempotency_key { + CodexAccountResetReserveResult::Reserved(reservation) + } else { + CodexAccountResetReserveResult::Busy(reservation) + })); + } + let active_generation = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .unwrap_or(0); + let sequence = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_SEQUENCE_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .unwrap_or(active_generation) + .max(active_generation); + let Some(generation) = sequence.checked_add(1) else { + return Err(GatewayError::Internal( + "Codex reset generation exhausted".to_string(), + )); + }; + let reservation = CodexAccountResetReservation { + idempotency_key: idempotency_key.to_string(), + generation, + }; + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_SEQUENCE_KEY.to_string(), + serde_json::json!(generation), + ); + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY.to_string(), + serde_json::json!({ + "idempotency_key": reservation.idempotency_key, + "generation": reservation.generation, + }), + ); + if persist_codex_reset_namespace( + state, + &key, + expected_codex, + serde_json::Value::Object(codex), + expected_encrypted_auth_config, + expected_credential, + crate::clock::current_unix_secs(), + ) + .await? + { + return Ok(Some(CodexAccountResetReserveResult::Reserved(reservation))); + } + if attempt + 1 < CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + tokio::task::yield_now().await; + } + } + Ok(None) +} + +pub(crate) async fn complete_codex_account_reset( + state: &AdminAppState<'_>, + key_id: &str, + expected_encrypted_auth_config: &str, + expected_credential: &ProviderCatalogKeyOAuthCredentialFence, + reservation: &CodexAccountResetReservation, + outcome: &str, + fence_unix_ms: u64, +) -> Result, GatewayError> { + let outcome = outcome.trim(); + let activates = matches!(outcome, "reset" | "already_redeemed"); + let noop = matches!(outcome, "nothing_to_reset" | "no_credit"); + if (!activates && !noop) || fence_unix_ms == 0 { + return Ok(None); + } + for attempt in 0..CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let Some(key) = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await? + .into_iter() + .next() + else { + return Ok(None); + }; + if !codex_reset_credential_matches( + &key, + expected_encrypted_auth_config, + expected_credential, + ) { + return Ok(None); + } + let expected_codex = key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + let mut codex = expected_codex + .as_ref() + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + let existing_terminal = codex_reset_history_from_object(&codex) + .into_iter() + .find(|entry| entry.idempotency_key == reservation.idempotency_key); + let upgrades_noop = existing_terminal.as_ref().is_some_and(|terminal| { + terminal.generation == reservation.generation + && !matches!(terminal.outcome.as_str(), "reset" | "already_redeemed") + && activates + }); + if let Some(terminal) = existing_terminal.as_ref().filter(|_| !upgrades_noop) { + return Ok(Some(CodexAccountResetCompleteResult::Replay( + terminal.clone(), + ))); + } + if !upgrades_noop + && codex_reset_reservation_from_object(&codex).as_ref() != Some(reservation) + { + return Ok(None); + } + let terminal = CodexAccountResetTerminal { + idempotency_key: reservation.idempotency_key.clone(), + generation: reservation.generation, + outcome: outcome.to_string(), + }; + codex_reset_write_bounded_history(&mut codex, &terminal); + if codex_reset_reservation_from_object(&codex).as_ref() == Some(reservation) { + codex.remove(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY); + } + let active_generation = codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .unwrap_or(0); + let completed = if activates && active_generation <= reservation.generation { + let fence = CodexAccountResetFence { + unix_ms: fence_unix_ms, + id: format!("reset:{}", reservation.idempotency_key), + generation: reservation.generation, + }; + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY.to_string(), + serde_json::json!(reservation.generation), + ); + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_PENDING_GENERATION_KEY + .to_string(), + serde_json::json!(reservation.generation), + ); + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY.to_string(), + serde_json::json!(true), + ); + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY.to_string(), + serde_json::json!(fence_unix_ms), + ); + codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY.to_string(), + serde_json::json!(fence.id), + ); + CodexAccountResetCompleteResult::Activated(fence) + } else if activates { + CodexAccountResetCompleteResult::Replay(terminal) + } else { + CodexAccountResetCompleteResult::Noop(terminal) + }; + if persist_codex_reset_namespace( + state, + &key, + expected_codex, + serde_json::Value::Object(codex), + expected_encrypted_auth_config, + expected_credential, + fence_unix_ms / 1_000, + ) + .await? + { + return Ok(Some(completed)); + } + if attempt + 1 < CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + tokio::task::yield_now().await; + } + } + Ok(None) +} + +pub(crate) async fn persist_codex_account_reset_fence( + state: &AdminAppState<'_>, + key_id: &str, + expected_encrypted_auth_config: Option<&str>, + expected_credential: Option<&ProviderCatalogKeyOAuthCredentialFence>, + fence_unix_ms: u64, + fence_id: &str, + idempotency_key: &str, +) -> Result, GatewayError> { + let fence_id = fence_id.trim(); + let idempotency_key = idempotency_key.trim(); + if fence_unix_ms == 0 || fence_id.is_empty() || idempotency_key.is_empty() { + return Ok(None); + } + + for attempt in 0..CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let Some(latest_key) = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await? + .into_iter() + .next() + else { + return Ok(None); + }; + let expected_codex = latest_key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + let mut next_codex = expected_codex + .as_ref() + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + let already_processed = codex_reset_write_processed_id(&mut next_codex, idempotency_key); + let stored_fence = next_codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .zip( + next_codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned), + ); + let owns_stored_fence = already_processed + && stored_fence + .as_ref() + .is_some_and(|(_, stored_id)| stored_id == fence_id); + let installs_fence = !already_processed + && stored_fence + .as_ref() + .is_none_or(|(stored_unix_ms, stored_id)| { + (fence_unix_ms, fence_id) > (*stored_unix_ms, stored_id.as_str()) + }); + if installs_fence { + next_codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY.to_string(), + serde_json::json!(fence_unix_ms), + ); + next_codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY.to_string(), + serde_json::json!(fence_id), + ); + next_codex.insert( + admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY.to_string(), + serde_json::json!(true), + ); + } + let effective_fence = CodexAccountResetFence { + unix_ms: next_codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .unwrap_or(fence_unix_ms), + id: next_codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY) + .and_then(serde_json::Value::as_str) + .unwrap_or(fence_id) + .to_string(), + generation: next_codex + .get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(admin_provider_quota_pure::coerce_json_u64) + .unwrap_or(0), + }; + let install = if installs_fence || owns_stored_fence { + CodexAccountResetFenceInstall::Owned(effective_fence) + } else { + CodexAccountResetFenceInstall::Superseded + }; + if already_processed { + let credential_matches = match (expected_encrypted_auth_config, expected_credential) { + (Some(expected_auth), Some(expected_credential)) => { + latest_key.encrypted_auth_config.as_deref() == Some(expected_auth) + && latest_key.encrypted_api_key == expected_credential.encrypted_api_key + && latest_key.auth_type == expected_credential.auth_type + && latest_key.provider_id == expected_credential.provider_id + } + (Some(expected_auth), None) => { + latest_key.encrypted_auth_config.as_deref() == Some(expected_auth) + } + (None, _) => true, + }; + return Ok(credential_matches.then_some(install)); + } + let next_codex = serde_json::Value::Object(next_codex); + + let persisted = if let Some(expected_encrypted_auth_config) = expected_encrypted_auth_config + { + if latest_key.encrypted_auth_config.as_deref() != Some(expected_encrypted_auth_config) { + return Ok(None); + } + state + .app() + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.to_string(), + expected_encrypted_auth_config: Some( + expected_encrypted_auth_config.to_string(), + ), + expected_credential: expected_credential.cloned(), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex, + }, + ), + encrypted_auth_config: expected_encrypted_auth_config.to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), + upstream_metadata_patch: Some(serde_json::json!({ + "codex": next_codex + })), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: serde_json::json!({}), + reset_error_count: false, + updated_at_unix_secs: Some(fence_unix_ms / 1_000), + }, + ) + .await? + } else { + state + .app() + .update_provider_catalog_key_runtime_metadata( + &ProviderCatalogKeyRuntimeMetadataUpdate { + key_id: key_id.to_string(), + namespace: "codex".to_string(), + expected_upstream_metadata_value: expected_codex, + upstream_metadata_value: next_codex, + status_snapshot_patch: serde_json::json!({}), + updated_at_unix_secs: Some(fence_unix_ms / 1_000), + }, + ) + .await? + }; + if persisted { + return Ok(Some(install)); + } + if attempt + 1 < CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let backoff_us = 50_u64.saturating_mul((attempt + 1) as u64).min(1_000); + tokio::time::sleep(std::time::Duration::from_micros(backoff_us)).await; + } + } + Ok(None) +} + pub(super) fn coerce_json_string(value: Option<&serde_json::Value>) -> Option { admin_provider_quota_pure::coerce_json_string(value) } @@ -286,6 +931,28 @@ pub(super) fn build_provider_quota_execution_plan( } } +fn codex_reset_refresh_is_superseded( + current: Option<&serde_json::Value>, + context: admin_provider_quota_pure::CodexQuotaMergeContext<'_>, +) -> bool { + let Some(incoming_fence_id) = context + .account_reset_fence_id + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return false; + }; + current + .and_then(serde_json::Value::as_object) + .and_then(|codex| { + codex.get(admin_provider_quota_pure::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY) + }) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .is_some_and(|stored_fence_id| stored_fence_id != incoming_fence_id) +} + pub(crate) async fn persist_provider_quota_refresh_state( state: &AdminAppState<'_>, key_id: &str, @@ -306,6 +973,154 @@ pub(crate) async fn persist_provider_quota_refresh_state( .await } +pub(crate) async fn persist_codex_provider_quota_refresh_state( + state: &AdminAppState<'_>, + key_id: &str, + metadata_update: Option<&serde_json::Value>, + oauth_invalid_at_unix_secs: Option, + oauth_invalid_reason: Option, + encrypted_auth_config: Option, + merge_context: admin_provider_quota_pure::CodexQuotaMergeContext<'_>, +) -> Result { + let Some(incoming_codex) = metadata_update.and_then(|value| value.get("codex")) else { + return persist_provider_quota_refresh_state( + state, + key_id, + metadata_update, + oauth_invalid_at_unix_secs, + oauth_invalid_reason, + encrypted_auth_config, + ) + .await; + }; + + for attempt in 0..CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let Some(mut latest_key) = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await? + .into_iter() + .next() + else { + return Ok(false); + }; + let expected_codex = latest_key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + if !admin_provider_quota_pure::codex_credential_generation_matches( + expected_codex.as_ref(), + merge_context.observed_credential_generation, + ) { + return Ok(true); + } + if codex_reset_refresh_is_superseded(expected_codex.as_ref(), merge_context) { + return Ok(true); + } + let Some(outcome) = admin_provider_quota_pure::merge_codex_quota_metadata_snapshot( + expected_codex.as_ref(), + incoming_codex, + merge_context, + ) else { + return Ok(false); + }; + let merged_update = serde_json::json!({"codex": outcome.metadata.clone()}); + latest_key.upstream_metadata = Some(merge_upstream_metadata( + latest_key.upstream_metadata.as_ref(), + &merged_update, + )); + let current_encrypted_auth_config = latest_key.encrypted_auth_config.clone(); + if let Some(encrypted_auth_config) = encrypted_auth_config.as_ref() { + latest_key.encrypted_auth_config = Some(encrypted_auth_config.clone()); + } + if encrypted_auth_config.is_some() { + ( + latest_key.oauth_invalid_at_unix_secs, + latest_key.oauth_invalid_reason, + ) = merge_codex_oauth_response_state( + &latest_key, + oauth_invalid_at_unix_secs, + oauth_invalid_reason.as_deref(), + merge_context.observed_at_unix_secs, + ); + } + latest_key.status_snapshot = sync_provider_key_quota_status_snapshot( + latest_key.status_snapshot.as_ref(), + "codex", + latest_key.upstream_metadata.as_ref(), + "refresh_api", + ); + if encrypted_auth_config.is_some() { + latest_key.status_snapshot = sync_provider_key_oauth_status_snapshot( + latest_key.status_snapshot.as_ref(), + &latest_key, + ); + } + latest_key.updated_at_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()); + + let persisted = if let Some(encrypted_auth_config) = encrypted_auth_config.as_ref() { + state + .app() + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.to_string(), + expected_encrypted_auth_config: current_encrypted_auth_config, + expected_credential: None, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex, + }, + ), + encrypted_auth_config: encrypted_auth_config.clone(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), + upstream_metadata_patch: Some(serde_json::json!({ + "codex": outcome.metadata + })), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: provider_quota_refresh_status_patch( + latest_key.status_snapshot.as_ref(), + ), + reset_error_count: false, + updated_at_unix_secs: latest_key.updated_at_unix_secs, + }, + ) + .await? + } else { + state + .app() + .update_provider_catalog_key_runtime_metadata( + &ProviderCatalogKeyRuntimeMetadataUpdate { + key_id: key_id.to_string(), + namespace: "codex".to_string(), + expected_upstream_metadata_value: expected_codex, + upstream_metadata_value: outcome.metadata, + status_snapshot_patch: provider_quota_refresh_status_patch( + latest_key.status_snapshot.as_ref(), + ), + updated_at_unix_secs: latest_key.updated_at_unix_secs, + }, + ) + .await? + }; + if persisted { + return Ok(true); + } + if attempt + 1 < CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let backoff_us = 50_u64.saturating_mul((attempt + 1) as u64).min(1_000); + tokio::time::sleep(std::time::Duration::from_micros(backoff_us)).await; + } + } + Ok(false) +} + /// Persist a Codex Agent Identity quota response only when the exact encrypted /// auth_config used for the request is still installed. Metadata, OAuth state, /// and their status projection share one repository CAS so a replacement cannot @@ -317,6 +1132,8 @@ pub(crate) async fn persist_fenced_provider_quota_refresh_state( metadata_update: Option<&serde_json::Value>, oauth_invalid_at_unix_secs: Option, oauth_invalid_reason: Option, + merge_context: admin_provider_quota_pure::CodexQuotaMergeContext<'_>, + expected_credential: Option<&ProviderCatalogKeyOAuthCredentialFence>, ) -> Result { let expected_encrypted_auth_config = expected_encrypted_auth_config.trim(); if expected_encrypted_auth_config.is_empty() { @@ -327,64 +1144,153 @@ pub(crate) async fn persist_fenced_provider_quota_refresh_state( "fenced quota metadata update must be an object".to_string(), )); } - let Some(mut latest_key) = state - .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) - .await? - .into_iter() - .next() - else { - return Ok(false); - }; - if latest_key.encrypted_auth_config.as_deref() != Some(expected_encrypted_auth_config) { - return Ok(false); - } + for attempt in 0..CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let Some(mut latest_key) = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await? + .into_iter() + .next() + else { + return Ok(false); + }; + if latest_key.encrypted_auth_config.as_deref() != Some(expected_encrypted_auth_config) { + return Ok(false); + } - let quota_snapshot_provider_type = - metadata_update.and_then(aether_provider_pool::provider_pool_quota_metadata_provider_type); - if let Some(metadata_update) = metadata_update { + let expected_codex = latest_key + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned(); + if !admin_provider_quota_pure::codex_credential_generation_matches( + expected_codex.as_ref(), + merge_context.observed_credential_generation, + ) { + return Ok(true); + } + if codex_reset_refresh_is_superseded(expected_codex.as_ref(), merge_context) { + return Ok(true); + } + let merged_codex = match metadata_update.and_then(|value| value.get("codex")) { + Some(incoming_codex) => { + let Some(outcome) = admin_provider_quota_pure::merge_codex_quota_metadata_snapshot( + expected_codex.as_ref(), + incoming_codex, + merge_context, + ) else { + return Ok(false); + }; + outcome.metadata + } + None => expected_codex + .clone() + .unwrap_or_else(|| serde_json::json!({})), + }; + let expected_codex_object = expected_codex + .as_ref() + .and_then(serde_json::Value::as_object); + let stale_oauth_state = admin_provider_quota_pure::codex_oauth_state_request_order_is_stale( + expected_codex_object, + merge_context.request_started_at_unix_ms, + merge_context.request_order_id, + ); + let mut merged_codex = merged_codex.as_object().cloned().unwrap_or_default(); + if !stale_oauth_state { + if let Some(request_started_at_unix_ms) = merge_context.request_started_at_unix_ms { + merged_codex.insert( + CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY.to_string(), + serde_json::json!(request_started_at_unix_ms), + ); + if let Some(request_order_id) = merge_context + .request_order_id + .map(str::trim) + .filter(|value| !value.is_empty()) + { + merged_codex.insert( + CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY.to_string(), + serde_json::json!(request_order_id), + ); + } else { + merged_codex.remove(CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY); + } + } + ( + latest_key.oauth_invalid_at_unix_secs, + latest_key.oauth_invalid_reason, + ) = merge_codex_oauth_response_state( + &latest_key, + oauth_invalid_at_unix_secs, + oauth_invalid_reason.as_deref(), + merge_context.observed_at_unix_secs, + ); + } + let merged_metadata_update = serde_json::json!({ + "codex": serde_json::Value::Object(merged_codex) + }); latest_key.upstream_metadata = Some(merge_upstream_metadata( latest_key.upstream_metadata.as_ref(), - metadata_update, + &merged_metadata_update, )); - } - latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs; - latest_key.oauth_invalid_reason = oauth_invalid_reason; - if let Some(provider_type) = quota_snapshot_provider_type.as_deref() { - latest_key.status_snapshot = sync_provider_key_quota_status_snapshot( + if metadata_update + .and_then(|value| value.get("codex")) + .is_some() + { + latest_key.status_snapshot = sync_provider_key_quota_status_snapshot( + latest_key.status_snapshot.as_ref(), + "codex", + latest_key.upstream_metadata.as_ref(), + "refresh_api", + ); + } + latest_key.status_snapshot = sync_provider_key_oauth_status_snapshot( latest_key.status_snapshot.as_ref(), - provider_type, - latest_key.upstream_metadata.as_ref(), - "refresh_api", + &latest_key, ); - } - latest_key.status_snapshot = - sync_provider_key_oauth_status_snapshot(latest_key.status_snapshot.as_ref(), &latest_key); - latest_key.updated_at_unix_secs = SystemTime::now() - .duration_since(UNIX_EPOCH) - .ok() - .map(|duration| duration.as_secs()); + latest_key.updated_at_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()); - state - .app() - .compare_and_update_provider_catalog_key_oauth_runtime_state( - &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - key_id: key_id.to_string(), - expected_encrypted_auth_config: Some(expected_encrypted_auth_config.to_string()), - expected_credential: None, - encrypted_auth_config: expected_encrypted_auth_config.to_string(), - encrypted_api_key_update: None, - expires_at_unix_secs_update: None, - oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, - oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), - upstream_metadata_patch: metadata_update.cloned(), - status_snapshot_patch: provider_quota_refresh_status_patch( - latest_key.status_snapshot.as_ref(), - ), - reset_error_count: false, - updated_at_unix_secs: latest_key.updated_at_unix_secs, - }, - ) - .await + let updated = state + .app() + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.to_string(), + expected_encrypted_auth_config: Some( + expected_encrypted_auth_config.to_string(), + ), + expected_credential: expected_credential.cloned(), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex, + }, + ), + encrypted_auth_config: expected_encrypted_auth_config.to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), + upstream_metadata_patch: Some(merged_metadata_update), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: provider_quota_refresh_status_patch( + latest_key.status_snapshot.as_ref(), + ), + reset_error_count: false, + updated_at_unix_secs: latest_key.updated_at_unix_secs, + }, + ) + .await?; + if updated { + return Ok(true); + } + if attempt + 1 < CODEX_QUOTA_PERSIST_CAS_MAX_ATTEMPTS { + let backoff_us = 50_u64.saturating_mul((attempt + 1) as u64).min(1_000); + tokio::time::sleep(std::time::Duration::from_micros(backoff_us)).await; + } + } + Ok(false) } async fn persist_provider_quota_refresh_state_after_read( @@ -639,10 +1545,1190 @@ mod tests { use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogKey, + StoredProviderCatalogProvider, }; use serde_json::json; use std::sync::Arc; + fn codex_merge_context( + request_started_at_unix_ms: u64, + ) -> admin_provider_quota_pure::CodexQuotaMergeContext<'static> { + codex_merge_context_with_id(request_started_at_unix_ms, None) + } + + fn codex_merge_context_with_id( + request_started_at_unix_ms: u64, + request_order_id: Option<&'static str>, + ) -> admin_provider_quota_pure::CodexQuotaMergeContext<'static> { + admin_provider_quota_pure::CodexQuotaMergeContext { + observed_at_unix_secs: request_started_at_unix_ms / 1_000, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id, + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: None, + coverage: admin_provider_quota_pure::CodexQuotaWindowCoverage::AccountSnapshot, + } + } + + fn codex_reset_merge_context( + request_started_at_unix_ms: u64, + fence_id: &'static str, + ) -> admin_provider_quota_pure::CodexQuotaMergeContext<'static> { + admin_provider_quota_pure::CodexQuotaMergeContext { + observed_at_unix_secs: request_started_at_unix_ms / 1_000, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: Some("reset-refresh"), + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: Some(fence_id), + coverage: admin_provider_quota_pure::CodexQuotaWindowCoverage::AccountSnapshot, + } + } + + fn codex_refresh_test_state( + key_id: &str, + encrypted_auth_config: Option<&str>, + ) -> (AppState, Arc) { + let mut key = StoredProviderCatalogKey::new( + key_id.to_string(), + "provider-codex-refresh".to_string(), + "Codex Refresh".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.encrypted_auth_config = encrypted_auth_config.map(ToOwned::to_owned); + key.upstream_metadata = Some(json!({ + "codex": { + "plan_type": "plus", + "primary_used_percent": 60.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 200_000u64, + "updated_at": 200u64 + } + })); + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![], + vec![], + vec![key], + )); + let app = AppState::new() + .expect("app should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone( + &repository, + )), + ); + (app, repository) + } + + fn codex_reset_state_machine_test_state( + key_id: &str, + ) -> ( + AppState, + Arc, + ProviderCatalogKeyOAuthCredentialFence, + ) { + let provider = StoredProviderCatalogProvider::new( + "provider-codex-reset-state".to_string(), + "Codex Reset State".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"); + let mut key = StoredProviderCatalogKey::new( + key_id.to_string(), + provider.id.clone(), + "Codex Reset".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.encrypted_auth_config = Some("auth-v1".to_string()); + key.upstream_metadata = Some(json!({ + "codex": {"credential_generation": "credential-v1"} + })); + let credential = ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: None, + auth_type: key.auth_type.clone(), + provider_id: provider.id.clone(), + provider_type: provider.provider_type.clone(), + }; + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![], + vec![key], + )); + let app = AppState::new() + .expect("app should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone( + &repository, + )), + ); + (app, repository, credential) + } + + #[tokio::test] + async fn codex_reset_reservation_serializes_ids_and_reuses_same_generation() { + let key_id = "key-codex-reset-reservation"; + let (app, repository, credential) = codex_reset_state_machine_test_state(key_id); + let admin_state = AdminAppState::new(&app); + + let first = reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-a", + ) + .await + .expect("reservation should complete") + .expect("reservation should exist"); + let first = match first { + CodexAccountResetReserveResult::Reserved(value) => value, + other => panic!("unexpected first reservation: {other:?}"), + }; + let same = reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-a", + ) + .await + .expect("same reservation should complete") + .expect("same reservation should exist"); + assert_eq!( + same, + CodexAccountResetReserveResult::Reserved(first.clone()) + ); + let other = reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-b", + ) + .await + .expect("busy check should complete") + .expect("busy reservation should exist"); + assert_eq!(other, CodexAccountResetReserveResult::Busy(first.clone())); + + // An ambiguous upstream result does not call complete; the durable + // reservation continues to block another id while the same id resumes. + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.unwrap()["codex"]; + assert_eq!(codex["account_quota_reset_sequence"], json!(1u64)); + assert_eq!( + codex["account_quota_reset_reservation"]["idempotency_key"], + json!("reset-a") + ); + assert!(codex.get("account_quota_reset_generation").is_none()); + } + + #[tokio::test] + async fn codex_reset_reservation_rejects_replaced_credential_generation() { + let key_id = "key-codex-reset-credential-generation"; + let (app, repository, credential) = codex_reset_state_machine_test_state(key_id); + let admin_state = AdminAppState::new(&app); + + let result = reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-before-rebind"), + "reset-from-old-account", + ) + .await + .expect("generation fence should complete") + .expect("generation mismatch should be explicit"); + assert_eq!( + result, + CodexAccountResetReserveResult::CredentialGenerationMismatch + ); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!( + stored.upstream_metadata.unwrap()["codex"], + json!({"credential_generation":"credential-v1"}) + ); + } + + #[tokio::test] + async fn codex_reset_noop_does_not_activate_but_later_id_gets_new_generation() { + let key_id = "key-codex-reset-noop"; + let (app, repository, credential) = codex_reset_state_machine_test_state(key_id); + let admin_state = AdminAppState::new(&app); + let reservation = match reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-noop", + ) + .await + .unwrap() + .unwrap() + { + CodexAccountResetReserveResult::Reserved(value) => value, + other => panic!("unexpected reservation: {other:?}"), + }; + assert!(matches!( + complete_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + &reservation, + "nothing_to_reset", + 200_000, + ) + .await + .unwrap(), + Some(CodexAccountResetCompleteResult::Noop(_)) + )); + let next = reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-next", + ) + .await + .unwrap() + .unwrap(); + assert!(matches!( + next, + CodexAccountResetReserveResult::Reserved(CodexAccountResetReservation { + generation: 2, + .. + }) + )); + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .unwrap() + .pop() + .unwrap(); + let codex = &stored.upstream_metadata.unwrap()["codex"]; + assert!(codex.get("account_quota_reset_generation").is_none()); + } + + async fn complete_codex_reset_in_order( + first_outcome: &str, + second_outcome: &str, + ) -> serde_json::Value { + let key_id = format!("key-codex-reset-order-{first_outcome}"); + let (app, repository, credential) = codex_reset_state_machine_test_state(&key_id); + let admin_state = AdminAppState::new(&app); + let reservation = match reserve_codex_account_reset( + &admin_state, + &key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "same-id", + ) + .await + .unwrap() + .unwrap() + { + CodexAccountResetReserveResult::Reserved(value) => value, + other => panic!("unexpected reservation: {other:?}"), + }; + complete_codex_account_reset( + &admin_state, + &key_id, + "auth-v1", + &credential, + &reservation, + first_outcome, + 200_000, + ) + .await + .unwrap() + .expect("first completion should persist"); + complete_codex_account_reset( + &admin_state, + &key_id, + "auth-v1", + &credential, + &reservation, + second_outcome, + 210_000, + ) + .await + .unwrap() + .expect("second completion should converge"); + repository + .list_keys_by_ids(&[key_id]) + .await + .unwrap() + .pop() + .unwrap() + .upstream_metadata + .unwrap()["codex"] + .clone() + } + + #[tokio::test] + async fn codex_reset_activation_wins_over_noop_in_both_completion_orders() { + for codex in [ + complete_codex_reset_in_order("nothing_to_reset", "reset").await, + complete_codex_reset_in_order("reset", "nothing_to_reset").await, + ] { + assert_eq!(codex["account_quota_reset_generation"], json!(1u64)); + assert_eq!(codex["account_quota_reset_pending_generation"], json!(1u64)); + assert_eq!(codex["account_quota_reset_pending"], json!(true)); + assert_eq!( + codex["account_quota_reset_history"][0]["outcome"], + json!("reset") + ); + } + } + + #[tokio::test] + async fn delayed_reset_upgrade_preserves_the_next_generation_reservation() { + let key_id = "key-codex-reset-upgrade-next-generation"; + let (app, repository, credential) = codex_reset_state_machine_test_state(key_id); + let admin_state = AdminAppState::new(&app); + let first = match reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-first", + ) + .await + .unwrap() + .unwrap() + { + CodexAccountResetReserveResult::Reserved(value) => value, + other => panic!("unexpected first reservation: {other:?}"), + }; + complete_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + &first, + "nothing_to_reset", + 200_000, + ) + .await + .unwrap() + .expect("first noop should complete"); + let second = match reserve_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + Some("credential-v1"), + "reset-second", + ) + .await + .unwrap() + .unwrap() + { + CodexAccountResetReserveResult::Reserved(value) => value, + other => panic!("unexpected second reservation: {other:?}"), + }; + assert_eq!(second.generation, 2); + + assert!(matches!( + complete_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + &first, + "reset", + 210_000, + ) + .await + .unwrap(), + Some(CodexAccountResetCompleteResult::Activated( + CodexAccountResetFence { generation: 1, .. } + )) + )); + let after_upgrade = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .unwrap() + .pop() + .unwrap(); + assert_eq!( + after_upgrade.upstream_metadata.as_ref().unwrap()["codex"] + ["account_quota_reset_reservation"], + json!({ + "idempotency_key": "reset-second", + "generation": 2, + }) + ); + + assert!(matches!( + complete_codex_account_reset( + &admin_state, + key_id, + "auth-v1", + &credential, + &second, + "reset", + 220_000, + ) + .await + .unwrap(), + Some(CodexAccountResetCompleteResult::Activated( + CodexAccountResetFence { generation: 2, .. } + )) + )); + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .unwrap() + .pop() + .unwrap(); + let codex = &stored.upstream_metadata.unwrap()["codex"]; + assert_eq!(codex["account_quota_reset_generation"], json!(2u64)); + assert!(codex.get("account_quota_reset_reservation").is_none()); + } + + #[tokio::test] + async fn stale_codex_refresh_cannot_lower_realtime_usage() { + let key_id = "key-codex-refresh-monotonic"; + let (app, repository) = codex_refresh_test_state(key_id, None); + let admin_state = AdminAppState::new(&app); + let stale_refresh = json!({"codex": { + "plan_type": "plus", + "primary_used_percent": 50.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64, + "updated_at": 100u64 + }}); + + assert!(persist_codex_provider_quota_refresh_state( + &admin_state, + key_id, + Some(&stale_refresh), + None, + None, + None, + codex_merge_context(100_000), + ) + .await + .expect("refresh persistence should complete")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!(codex["primary_used_percent"], json!(60.0)); + assert_eq!( + codex["account_quota_request_started_at_unix_ms"], + json!(200_000u64) + ); + } + + #[tokio::test] + async fn codex_reset_fence_is_idempotent_and_rejects_pre_reset_response() { + let key_id = "key-codex-reset-fence"; + let (app, repository) = codex_refresh_test_state(key_id, None); + let admin_state = AdminAppState::new(&app); + + let initial_fence = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 250_000, + "fence-a", + "redeem-once", + ) + .await + .expect("reset fence should persist") + .expect("reset fence should be returned"); + let initial_fence = match initial_fence { + CodexAccountResetFenceInstall::Owned(fence) => fence, + CodexAccountResetFenceInstall::Superseded => { + panic!("initial reset should own its fence") + } + }; + let duplicate_fence = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 300_000, + "fence-a", + "redeem-once", + ) + .await + .expect("duplicate reset fence should be idempotent") + .expect("duplicate should return the installed fence"); + let duplicate_fence = match duplicate_fence { + CodexAccountResetFenceInstall::Owned(fence) => fence, + CodexAccountResetFenceInstall::Superseded => { + panic!("duplicate active reset should retain ownership") + } + }; + assert_eq!(duplicate_fence, initial_fence); + + let stale = json!({"codex": { + "primary_used_percent": 100.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + assert!(persist_codex_provider_quota_refresh_state( + &admin_state, + key_id, + Some(&stale), + None, + None, + None, + codex_merge_context(200_000), + ) + .await + .expect("stale response should be harmlessly acknowledged")); + + let baseline = json!({"codex": { + "primary_used_percent": 0.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + assert!(persist_codex_provider_quota_refresh_state( + &admin_state, + key_id, + Some(&baseline), + None, + None, + None, + codex_reset_merge_context(260_000, "fence-a"), + ) + .await + .expect("reset baseline should persist")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!(codex["primary_used_percent"], json!(0.0)); + assert_eq!(codex["account_quota_reset_fence_id"], json!("fence-a")); + assert_eq!(codex["account_quota_reset_pending"], json!(false)); + } + + #[tokio::test] + async fn codex_reset_fence_barrier_never_moves_backward_and_remembers_processed_ids() { + let key_id = "key-codex-reset-fence-order"; + let (app, repository) = codex_refresh_test_state(key_id, None); + let admin_state = AdminAppState::new(&app); + + let newer = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 300_000, + "fence-newer", + "redeem-newer", + ) + .await + .expect("newer reset fence should persist") + .expect("newer reset fence should be returned"); + assert!(matches!( + newer, + CodexAccountResetFenceInstall::Owned(CodexAccountResetFence { + unix_ms: 300_000, + ref id, + .. + }) if id == "fence-newer" + )); + let delayed_older = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 250_000, + "fence-older", + "redeem-older", + ) + .await + .expect("older reset should be recorded") + .expect("install result should be returned"); + assert_eq!(delayed_older, CodexAccountResetFenceInstall::Superseded); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!( + codex["account_quota_reset_fence_unix_ms"], + json!(300_000u64) + ); + assert_eq!(codex["account_quota_reset_fence_id"], json!("fence-newer")); + let processed_ids = codex["account_quota_reset_processed_ids"] + .as_array() + .expect("processed reset ids should be an array"); + assert_eq!(processed_ids.len(), 2); + assert!(processed_ids.contains(&json!("redeem-older"))); + assert!(processed_ids.contains(&json!("redeem-newer"))); + } + + #[tokio::test] + async fn concurrent_codex_reset_fences_converge_on_newest_barrier() { + let key_id = "key-codex-reset-fence-concurrent"; + let (app, repository) = codex_refresh_test_state(key_id, None); + let admin_state = AdminAppState::new(&app); + + let (older, newer) = tokio::join!( + persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 250_000, + "fence-older", + "redeem-older", + ), + persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 300_000, + "fence-newer", + "redeem-newer", + ), + ); + let older = older + .expect("older reset should complete") + .expect("older reset should return an install result"); + let newer = newer + .expect("newer reset should complete") + .expect("newer reset should return an install result"); + assert!(matches!( + (older, newer), + ( + CodexAccountResetFenceInstall::Owned(_), + CodexAccountResetFenceInstall::Owned(_) + ) | ( + CodexAccountResetFenceInstall::Superseded, + CodexAccountResetFenceInstall::Owned(_) + ) + )); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!( + codex["account_quota_reset_fence_unix_ms"], + json!(300_000u64) + ); + assert_eq!(codex["account_quota_reset_fence_id"], json!("fence-newer")); + assert_eq!(codex["account_quota_reset_pending"], json!(true)); + let processed_ids = codex["account_quota_reset_processed_ids"] + .as_array() + .expect("processed reset ids should be an array"); + assert_eq!(processed_ids.len(), 2); + assert!(processed_ids.contains(&json!("redeem-older"))); + assert!(processed_ids.contains(&json!("redeem-newer"))); + } + + #[tokio::test] + async fn superseded_codex_reset_refresh_cannot_confirm_newer_fence() { + let key_id = "key-codex-reset-fence-stale-refresh"; + let (app, repository) = codex_refresh_test_state(key_id, None); + let admin_state = AdminAppState::new(&app); + + for (fence_unix_ms, fence_id, redeem_id) in [ + (250_000, "fence-older", "redeem-older"), + (300_000, "fence-newer", "redeem-newer"), + ] { + let install = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + fence_unix_ms, + fence_id, + redeem_id, + ) + .await + .expect("reset fence should persist") + .expect("reset fence should return an install result"); + assert!(matches!(install, CodexAccountResetFenceInstall::Owned(_))); + } + + let stale_baseline = json!({"codex": { + "primary_used_percent": 0.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + assert!(persist_codex_provider_quota_refresh_state( + &admin_state, + key_id, + Some(&stale_baseline), + None, + None, + None, + admin_provider_quota_pure::CodexQuotaMergeContext { + observed_at_unix_secs: 310, + request_started_at_unix_ms: Some(310_000), + request_order_id: Some("older-reset-late-refresh"), + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: Some("fence-older"), + coverage: admin_provider_quota_pure::CodexQuotaWindowCoverage::AccountSnapshot, + }, + ) + .await + .expect("superseded reset refresh should be harmlessly acknowledged")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!(codex["primary_used_percent"], json!(60.0)); + assert_eq!(codex["account_quota_reset_fence_id"], json!("fence-newer")); + assert_eq!(codex["account_quota_reset_pending"], json!(true)); + } + + #[tokio::test] + async fn replaying_historical_codex_reset_does_not_reopen_pending() { + let key_id = "key-codex-reset-fence-replay"; + let (app, repository) = codex_refresh_test_state(key_id, None); + let admin_state = AdminAppState::new(&app); + + for (fence_unix_ms, fence_id, redeem_id, request_started_at_unix_ms, usage) in [ + (250_000, "fence-a", "redeem-a", 260_000, 20.0), + (300_000, "fence-b", "redeem-b", 310_000, 0.0), + ] { + let install = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + fence_unix_ms, + fence_id, + redeem_id, + ) + .await + .expect("reset fence should persist") + .expect("reset fence should be returned"); + assert!(matches!(install, CodexAccountResetFenceInstall::Owned(_))); + let baseline = json!({"codex": { + "primary_used_percent": usage, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + assert!(persist_codex_provider_quota_refresh_state( + &admin_state, + key_id, + Some(&baseline), + None, + None, + None, + admin_provider_quota_pure::CodexQuotaMergeContext { + observed_at_unix_secs: request_started_at_unix_ms / 1_000, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: Some("reset-refresh"), + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: Some(fence_id), + coverage: admin_provider_quota_pure::CodexQuotaWindowCoverage::AccountSnapshot, + }, + ) + .await + .expect("reset baseline should persist")); + } + + let replay = persist_codex_account_reset_fence( + &admin_state, + key_id, + None, + None, + 350_000, + "fence-a-replay", + "redeem-a", + ) + .await + .expect("historical replay should be idempotent") + .expect("historical replay should return an install result"); + assert_eq!(replay, CodexAccountResetFenceInstall::Superseded); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!(codex["primary_used_percent"], json!(0.0)); + assert_eq!(codex["account_quota_reset_fence_id"], json!("fence-b")); + assert_eq!(codex["account_quota_reset_pending"], json!(false)); + assert_eq!( + codex["account_quota_reset_processed_ids"], + json!(["redeem-a", "redeem-b"]) + ); + } + + #[tokio::test] + async fn fenced_stale_codex_refresh_keeps_usage_and_oauth_state() { + let key_id = "key-codex-fenced-refresh-monotonic"; + let (app, repository) = codex_refresh_test_state(key_id, Some("auth-v1")); + let admin_state = AdminAppState::new(&app); + let stale_refresh = json!({"codex": { + "primary_used_percent": 50.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some(&stale_refresh), + Some(300), + Some("refresh-state".to_string()), + codex_merge_context(100_000), + None, + ) + .await + .expect("fenced refresh persistence should complete")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["codex"]["primary_used_percent"], + json!(60.0) + ); + assert_eq!(stored.oauth_invalid_at_unix_secs, None); + assert_eq!(stored.oauth_invalid_reason, None); + } + + #[tokio::test] + async fn fenced_older_refresh_cannot_overwrite_newer_oauth_state() { + let key_id = "key-codex-fenced-oauth-watermark"; + let (app, repository) = codex_refresh_test_state(key_id, Some("auth-v1")); + let admin_state = AdminAppState::new(&app); + let quota = |used_percent| { + json!({"codex": { + "primary_used_percent": used_percent, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}) + }; + + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a(70.0)), + None, + None, + codex_merge_context(300_000), + None, + ) + .await + .expect("newer refresh should persist")); + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a(65.0)), + Some(250), + Some("stale-invalid".to_string()), + codex_merge_context(250_000), + None, + ) + .await + .expect("older refresh should merge without replacing OAuth state")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.oauth_invalid_at_unix_secs, None); + assert_eq!(stored.oauth_invalid_reason, None); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["codex"] + [CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY], + json!(300_000u64) + ); + } + + #[tokio::test] + async fn fenced_same_millisecond_refresh_uses_request_id_for_oauth_order() { + let key_id = "key-codex-fenced-oauth-id-watermark"; + let (app, repository) = codex_refresh_test_state(key_id, Some("auth-v1")); + let admin_state = AdminAppState::new(&app); + let quota = json!({"codex": { + "primary_used_percent": 70.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + None, + None, + codex_merge_context_with_id(300_000, Some("request-b")), + None, + ) + .await + .expect("newer same-millisecond refresh should persist")); + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + Some(300), + Some("stale-invalid".to_string()), + codex_merge_context_with_id(300_000, Some("request-a")), + None, + ) + .await + .expect("older same-millisecond refresh should merge without replacing OAuth state")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.oauth_invalid_at_unix_secs, None); + assert_eq!(stored.oauth_invalid_reason, None); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY], + json!(300_000u64) + ); + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY], + json!("request-b") + ); + } + + #[tokio::test] + async fn fenced_same_millisecond_newer_request_id_can_replace_oauth_state() { + let key_id = "key-codex-fenced-oauth-id-watermark-newer-invalid"; + let (app, repository) = codex_refresh_test_state(key_id, Some("auth-v1")); + let admin_state = AdminAppState::new(&app); + let quota = json!({"codex": { + "primary_used_percent": 70.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + None, + None, + codex_merge_context_with_id(300_000, Some("request-a")), + None, + ) + .await + .expect("older same-millisecond refresh should persist")); + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + Some(300), + Some("newer-invalid".to_string()), + codex_merge_context_with_id(300_000, Some("request-b")), + None, + ) + .await + .expect("newer same-millisecond refresh should replace OAuth state")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.oauth_invalid_at_unix_secs, Some(300)); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("newer-invalid") + ); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY], + json!(300_000u64) + ); + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY], + json!("request-b") + ); + } + + #[tokio::test] + async fn fenced_older_success_cannot_clear_newer_oauth_invalid_state() { + let key_id = "key-codex-newer-invalid-older-success"; + let (app, repository) = codex_refresh_test_state(key_id, Some("auth-v1")); + let admin_state = AdminAppState::new(&app); + let quota = json!({"codex": { + "primary_used_percent": 70.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + Some(300), + Some("newer-invalid".to_string()), + codex_merge_context_with_id(300_000, Some("request-newer")), + None, + ) + .await + .expect("newer invalid response should persist")); + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + None, + None, + codex_merge_context_with_id(250_000, Some("request-older")), + None, + ) + .await + .expect("older success should be harmlessly acknowledged")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.oauth_invalid_at_unix_secs, Some(300)); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("newer-invalid") + ); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY], + json!(300_000u64) + ); + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY], + json!("request-newer") + ); + } + + #[tokio::test] + async fn fenced_same_millisecond_older_success_cannot_clear_newer_invalid_state() { + let key_id = "key-codex-same-ms-newer-invalid-older-success"; + let (app, repository) = codex_refresh_test_state(key_id, Some("auth-v1")); + let admin_state = AdminAppState::new(&app); + let quota = json!({"codex": { + "primary_used_percent": 70.0, + "primary_reset_at": 2_000_000_000u64, + "primary_window_minutes": 300u64 + }}); + + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + Some(300), + Some("newer-invalid".to_string()), + codex_merge_context_with_id(300_000, Some("request-b")), + None, + ) + .await + .expect("newer same-millisecond invalid response should persist")); + assert!(persist_fenced_provider_quota_refresh_state( + &admin_state, + key_id, + "auth-v1", + Some("a), + None, + None, + codex_merge_context_with_id(300_000, Some("request-a")), + None, + ) + .await + .expect("older same-millisecond success should be acknowledged")); + + let stored = repository + .list_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.oauth_invalid_at_unix_secs, Some(300)); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("newer-invalid") + ); + let codex = &stored.upstream_metadata.as_ref().unwrap()["codex"]; + assert_eq!( + codex[CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY], + json!("request-b") + ); + } + #[tokio::test] async fn metadata_cas_conflict_does_not_persist_stale_oauth_runtime_state() { let mut key = StoredProviderCatalogKey::new( diff --git a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs index eaeefbb3e..2130651ea 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs @@ -178,6 +178,7 @@ fn provider_query_execution_json_body_decodes_stream_encoded_json_response() { "content-type".to_string(), "application/json".to_string(), )]), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: None, body_bytes_b64: Some(encoded_body), @@ -436,6 +437,7 @@ fn provider_query_standard_test_aggregates_responses_stream_body() { candidate_id: Some("candidate-0".to_string()), status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: None, body_bytes_b64: Some( @@ -468,6 +470,7 @@ fn provider_query_standard_test_aggregates_responses_image_generation_call() { candidate_id: Some("candidate-0".to_string()), status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: None, body_bytes_b64: Some( @@ -601,6 +604,7 @@ fn provider_query_search_success_requires_non_empty_output() { candidate_id: Some("candidate-0".to_string()), status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(body), body_bytes_b64: None, @@ -751,6 +755,7 @@ fn provider_query_standard_test_rejects_gemini_success_without_visible_output() candidate_id: Some("candidate-0".to_string()), status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "candidates": [{ diff --git a/apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs b/apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs index ed7e747a9..83853d399 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs @@ -122,6 +122,7 @@ pub(crate) struct AdminProviderQuotaRefreshRequest { #[derive(Debug, Deserialize)] pub(crate) struct AdminCodexResetCreditConsumeRequest { pub(crate) idempotency_key: String, + pub(crate) expected_credential_generation: serde_json::Value, } #[derive(Debug, Deserialize)] @@ -339,3 +340,37 @@ pub(crate) struct AdminImportProviderModelsRequest { )] pub(crate) price_per_request: Option, } + +#[cfg(test)] +mod tests { + use super::AdminCodexResetCreditConsumeRequest; + + #[test] + fn codex_reset_credit_consume_requires_an_explicit_credential_generation() { + assert!( + serde_json::from_value::( + serde_json::json!({"idempotency_key":"reset-old-client"}), + ) + .is_err() + ); + + let legacy_account = + serde_json::from_value::(serde_json::json!({ + "idempotency_key":"reset-legacy-account", + "expected_credential_generation":null, + })) + .expect("explicit null should fence an account without a generation"); + assert!(legacy_account.expected_credential_generation.is_null()); + + let generated_account = + serde_json::from_value::(serde_json::json!({ + "idempotency_key":"reset-generated-account", + "expected_credential_generation":"credential-v2", + })) + .expect("string generation should deserialize"); + assert_eq!( + generated_account.expected_credential_generation, + serde_json::json!("credential-v2") + ); + } +} diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs b/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs index e165abd9f..535b09f39 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs @@ -1,3 +1,4 @@ +use crate::handlers::admin::provider::oauth::provisioning::rotate_codex_credential_generation; use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest; use crate::handlers::admin::provider::write::normalize::{ normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, @@ -216,6 +217,7 @@ pub(crate) async fn build_admin_create_provider_key_record( )?; key.created_at_unix_ms = Some(now_unix_secs); key.updated_at_unix_secs = Some(now_unix_secs); + rotate_codex_credential_generation(&mut key, &provider.provider_type); Ok(key) } diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/keys/mod.rs b/apps/aether-gateway/src/handlers/admin/provider/write/keys/mod.rs index e75dfa4f0..b7b5175d7 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/keys/mod.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/keys/mod.rs @@ -6,6 +6,7 @@ pub(crate) use self::update::build_admin_update_provider_key_record; pub(crate) use self::update::{ admin_provider_key_update_requires_immediate_model_fetch, build_admin_update_provider_key_record_with_existing_keys, + build_provider_catalog_key_admin_cas_update, }; mod batch; diff --git a/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs b/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs index 58032d69f..5c114e2b6 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs @@ -1,3 +1,4 @@ +use crate::handlers::admin::provider::oauth::provisioning::rotate_codex_credential_generation; use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch; use crate::handlers::admin::provider::write::normalize::{ normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys, @@ -13,6 +14,7 @@ use crate::handlers::admin::shared::{ use crate::handlers::shared::normalize_optional_api_key_concurrent_limit; use crate::provider_key_auth::provider_key_is_oauth_managed; use aether_data_contracts::repository::provider_catalog::{ + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyOAuthCredentialFence, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use aether_provider_transport::provider_types::provider_type_is_fixed; @@ -368,6 +370,12 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys( .duration_since(UNIX_EPOCH) .ok() .map(|duration| duration.as_secs()); + let credential_identity_changed = !updated.auth_type.eq_ignore_ascii_case(&existing.auth_type) + || updated.encrypted_api_key != existing.encrypted_api_key + || updated.encrypted_auth_config != existing.encrypted_auth_config; + if credential_identity_changed { + rotate_codex_credential_generation(&mut updated, &provider.provider_type); + } Ok(updated) } @@ -382,6 +390,49 @@ pub(crate) fn admin_provider_key_update_requires_immediate_model_fetch( && (!existing.auto_fetch_models || filters_changed || locked_models_changed) } +pub(crate) fn build_provider_catalog_key_admin_cas_update( + existing: &StoredProviderCatalogKey, + updated: StoredProviderCatalogKey, + provider_type: &str, +) -> ProviderCatalogKeyAdminCasUpdate { + let previous_generation = existing + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/credential_generation")) + .and_then(serde_json::Value::as_str); + let next_generation = updated + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/credential_generation")) + .and_then(serde_json::Value::as_str); + let credential_changed = existing.auth_type != updated.auth_type + || existing.encrypted_api_key != updated.encrypted_api_key + || existing.encrypted_auth_config != updated.encrypted_auth_config; + let codex_rotation = provider_type + .trim() + .eq_ignore_ascii_case("codex") + .then(|| next_generation.filter(|next| Some(*next) != previous_generation)) + .flatten() + .map(|generation| { + json!({ + aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY: generation, + }) + }); + + ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: existing.encrypted_auth_config.clone(), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: existing.encrypted_api_key.clone(), + auth_type: existing.auth_type.clone(), + provider_id: existing.provider_id.clone(), + provider_type: provider_type.to_string(), + }, + key: updated, + codex_rotation, + reset_oauth_runtime: credential_changed, + } +} + fn raw_secret_auth_type(value: &str) -> bool { matches!( value.trim().to_ascii_lowercase().as_str(), diff --git a/apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs b/apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs index e1ef6a081..cd57b5f37 100644 --- a/apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs +++ b/apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs @@ -186,6 +186,15 @@ impl<'a> AdminAppState<'a> { self.app.update_provider_catalog_key(key).await } + pub(crate) async fn compare_and_update_provider_catalog_key_admin_state( + &self, + update: &aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + self.app + .compare_and_update_provider_catalog_key_admin_state(update) + .await + } + pub(crate) async fn compare_and_update_provider_catalog_key_adaptive_state( &self, update: &aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyAdaptiveStateUpdate, diff --git a/apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs b/apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs index 0e2995f1d..4b3b7dab7 100644 --- a/apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs +++ b/apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs @@ -513,6 +513,7 @@ impl<'a> AdminAppState<'a> { provider_id: key.provider_id.clone(), provider_type: provider.provider_type.clone(), }, + expected_upstream_metadata_namespace: None, }, ) .await diff --git a/apps/aether-gateway/src/handlers/admin/request/system/import.rs b/apps/aether-gateway/src/handlers/admin/request/system/import.rs index d02406e1e..312d41ab1 100644 --- a/apps/aether-gateway/src/handlers/admin/request/system/import.rs +++ b/apps/aether-gateway/src/handlers/admin/request/system/import.rs @@ -4,10 +4,12 @@ use crate::api::ai::admin_endpoint_signature_parts; use crate::handlers::admin::admin_provider_pool_config; use crate::handlers::admin::model::ADMIN_EXTERNAL_MODELS_PROXY_NODE_CONFIG_KEY; use crate::handlers::admin::provider::endpoints_admin::payloads::AdminProviderEndpointUpdatePatch; +use crate::handlers::admin::provider::oauth::provisioning::ensure_codex_credential_generation_rotated; use crate::handlers::admin::provider::shared::payloads::{ AdminProviderCreateRequest, AdminProviderKeyCreateRequest, AdminProviderKeyUpdatePatch, AdminProviderUpdatePatch, }; +use crate::handlers::admin::provider::write::keys::build_provider_catalog_key_admin_cas_update; use crate::handlers::admin::shared::{ normalize_json_array, normalize_json_object, normalize_string_list, }; @@ -377,10 +379,14 @@ fn normalize_import_key_raw_payload( fn apply_imported_oauth_key_credentials( state: &AdminAppState<'_>, + provider_type: &str, + previous_codex_credential_generation: Option<&str>, raw_key: &Map, normalized_auth_config: Option<&Value>, record: &mut aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey, ) -> Result { + let previous_encrypted_api_key = record.encrypted_api_key.clone(); + let previous_encrypted_auth_config = record.encrypted_auth_config.clone(); let mut credentials_supplied = false; let mut api_key_supplied = false; if let Some(api_key_value) = raw_key.get("api_key") { @@ -424,10 +430,19 @@ fn apply_imported_oauth_key_credentials( api_key_supplied, ); + let credential_material_changed = record.encrypted_api_key != previous_encrypted_api_key + || record.encrypted_auth_config != previous_encrypted_auth_config; if credentials_supplied { record.oauth_invalid_at_unix_secs = None; record.oauth_invalid_reason = None; } + if credential_material_changed { + ensure_codex_credential_generation_rotated( + record, + provider_type, + previous_codex_credential_generation, + ); + } Ok(credentials_supplied) } @@ -1861,6 +1876,15 @@ impl<'a> AdminAppState<'a> { if let Some(existing_index) = existing_key_index { let existing_key = existing_keys[existing_index].clone(); + let previous_codex_credential_generation = existing_key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .and_then(|codex| { + aether_admin::provider::quota::codex_credential_generation(Some(codex)) + }) + .map(ToOwned::to_owned); match merge_mode { AdminImportMergeMode::Skip => { stats.keys.skipped += 1; @@ -1890,6 +1914,8 @@ impl<'a> AdminAppState<'a> { let oauth_credentials_supplied = if auth_type == "oauth" { invalid!(apply_imported_oauth_key_credentials( self, + &provider.provider_type, + previous_codex_credential_generation.as_deref(), &raw_key, normalized_auth_config.as_ref(), &mut updated, @@ -1903,8 +1929,31 @@ impl<'a> AdminAppState<'a> { imported_key.fingerprint.clone(), "fingerprint", )); - let Some(mut persisted) = - self.update_provider_catalog_key(&updated).await? + let admin_update = build_provider_catalog_key_admin_cas_update( + &existing_key, + updated.clone(), + &provider.provider_type, + ); + if !self + .compare_and_update_provider_catalog_key_admin_state(&admin_update) + .await? + { + return Ok(Err(( + http::StatusCode::CONFLICT, + json!({ + "detail": format!( + "Provider '{provider_name}' 的 Key 已被其他请求更新，请重试" + ) + }), + ))); + } + let Some(mut persisted) = self + .read_provider_catalog_keys_by_ids(std::slice::from_ref( + &updated.id, + )) + .await? + .into_iter() + .next() else { return Ok(Err(invalid_request(format!( "更新 Provider '{provider_name}' 的 Key 失败" @@ -1926,16 +1975,15 @@ impl<'a> AdminAppState<'a> { persisted = reloaded; } if oauth_credentials_supplied { - if !self - .clear_provider_catalog_key_oauth_invalid_marker(&updated.id) - .await? - { - return Ok(Err(invalid_request(format!( - "更新 Provider '{provider_name}' 的 Key 失败" - )))); - } let Some(reloaded) = self - .reset_provider_catalog_key_recovery_state(&updated.id) + .reset_provider_catalog_key_recovery_state_fenced( + &updated.id, + updated.encrypted_auth_config.as_deref().ok_or_else(|| { + GatewayError::Internal(format!( + "OAuth Provider '{provider_name}' imported without auth_config" + )) + })?, + ) .await? else { return Ok(Err(invalid_request(format!( @@ -1975,6 +2023,8 @@ impl<'a> AdminAppState<'a> { let oauth_credentials_supplied = if auth_type == "oauth" { invalid!(apply_imported_oauth_key_credentials( self, + &provider.provider_type, + None, &raw_key, normalized_auth_config.as_ref(), &mut record, diff --git a/apps/aether-gateway/src/model_fetch/runtime.rs b/apps/aether-gateway/src/model_fetch/runtime.rs index 5502fb3f0..748381e92 100644 --- a/apps/aether-gateway/src/model_fetch/runtime.rs +++ b/apps/aether-gateway/src/model_fetch/runtime.rs @@ -867,6 +867,7 @@ mod tests { candidate_id: None, status_code: 200, headers: Default::default(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(body), body_bytes_b64: None, diff --git a/apps/aether-gateway/src/orchestration/effects.rs b/apps/aether-gateway/src/orchestration/effects.rs index 8e4f99d27..30f848998 100644 --- a/apps/aether-gateway/src/orchestration/effects.rs +++ b/apps/aether-gateway/src/orchestration/effects.rs @@ -237,6 +237,13 @@ pub(crate) struct LocalOAuthInvalidationEffect<'a> { pub(crate) response_text: Option<&'a str>, } +#[derive(Debug, Clone, Copy)] +pub(crate) struct LocalOAuthSuccessEffect<'a> { + pub(crate) status_code: u16, + pub(crate) request_started_at_unix_ms: Option, + pub(crate) request_order_id: Option<&'a str>, +} + #[derive(Debug, Clone, Copy)] pub(crate) enum LocalExecutionEffect<'a> { AttemptFailure(LocalAttemptFailureEffect), @@ -245,6 +252,7 @@ pub(crate) enum LocalExecutionEffect<'a> { HealthSuccess(LocalHealthSuccessEffect), AdaptiveSuccess(LocalAdaptiveSuccessEffect), OauthInvalidation(LocalOAuthInvalidationEffect<'a>), + OauthSuccess(LocalOAuthSuccessEffect<'a>), PoolSuccessSync { payload: &'a GatewaySyncReportRequest, }, @@ -255,6 +263,68 @@ pub(crate) enum LocalExecutionEffect<'a> { PoolStreamTimeout, } +#[derive(Debug)] +struct OwnedLocalOAuthSuccessEffect { + status_code: u16, + provider_id: String, + endpoint_id: String, + key_id: String, + authorization: String, + request_started_at_unix_ms: u64, + request_order_id: String, + observed_credential_generation: Option, +} + +fn owned_local_oauth_success_effect( + plan: &ExecutionPlan, + report_context: Option<&Value>, + effect: LocalOAuthSuccessEffect<'_>, +) -> Option { + if !(200..300).contains(&effect.status_code) { + return None; + } + let request_started_at_unix_ms = effect.request_started_at_unix_ms?; + let request_order_id = effect + .request_order_id + .map(str::trim) + .filter(|value| !value.is_empty())?; + let authorization = execution_plan_authorization(plan)?.trim().to_string(); + if authorization.is_empty() || bearer_access_token(&authorization).is_none() { + return None; + } + Some(OwnedLocalOAuthSuccessEffect { + status_code: effect.status_code, + provider_id: plan.provider_id.clone(), + endpoint_id: plan.endpoint_id.clone(), + key_id: plan.key_id.clone(), + authorization, + request_started_at_unix_ms, + request_order_id: request_order_id.to_string(), + observed_credential_generation: report_context_string_field( + report_context, + "codex_credential_generation", + ) + .map(ToOwned::to_owned), + }) +} + +/// Schedule a fenced Codex OAuth-success observation after provider headers are available. +/// Only the small identity/observation tuple is moved into the task; request bodies and plans +/// remain owned by the caller. +pub(crate) fn spawn_local_oauth_success_effect( + state: AppState, + plan: &ExecutionPlan, + report_context: Option<&Value>, + effect: LocalOAuthSuccessEffect<'_>, +) { + let Some(effect) = owned_local_oauth_success_effect(plan, report_context, effect) else { + return; + }; + tokio::spawn(async move { + record_oauth_success_effect_owned(&state, effect).await; + }); +} + struct PoolFeedbackContext { pool_config: AdminProviderPoolConfig, sticky_session_token: Option, @@ -302,6 +372,9 @@ pub(crate) async fn apply_local_execution_effect( LocalExecutionEffect::OauthInvalidation(effect) => { record_oauth_invalidation_effect(state, context, effect).await; } + LocalExecutionEffect::OauthSuccess(effect) => { + record_oauth_success_effect(state, context, effect).await; + } LocalExecutionEffect::PoolSuccessSync { payload } => { record_sync_pool_success_effect(state, context, payload).await; release_pool_key_lease_effect(state, context).await; @@ -349,6 +422,16 @@ fn report_context_string_field<'a>( .filter(|value| !value.is_empty()) } +fn report_context_u64_field(report_context: Option<&Value>, field: &str) -> Option { + report_context + .and_then(|context| context.get(field)) + .and_then(|value| { + value + .as_u64() + .or_else(|| value.as_str().and_then(|value| value.parse::().ok())) + }) +} + fn local_scheduler_affinity_cache_key(report_context: Option<&Value>) -> Option { let client_session_affinity = local_client_session_affinity(report_context); let policy_context = scheduler_affinity_policy_context_from_report_context(report_context); @@ -1463,9 +1546,23 @@ async fn record_oauth_invalidation_effect( ) else { return; }; + let request_started_at_unix_ms = report_context_u64_field( + context.report_context, + "provider_request_started_at_unix_ms", + ); + let request_order_id = + report_context_string_field(context.report_context, "provider_request_order_id"); + let observed_credential_generation = + report_context_string_field(context.report_context, "codex_credential_generation"); if let Err(err) = state - .mark_provider_transport_oauth_invalid_fenced(&transport, invalid_reason.as_str()) + .mark_provider_transport_oauth_invalid_fenced( + &transport, + invalid_reason.as_str(), + request_started_at_unix_ms, + request_order_id, + observed_credential_generation, + ) .await { warn!( @@ -1475,6 +1572,70 @@ async fn record_oauth_invalidation_effect( } } +async fn record_oauth_success_effect( + state: &AppState, + context: LocalExecutionEffectContext<'_>, + effect: LocalOAuthSuccessEffect<'_>, +) { + let Some(effect) = + owned_local_oauth_success_effect(context.plan, context.report_context, effect) + else { + return; + }; + record_oauth_success_effect_owned(state, effect).await; +} + +async fn record_oauth_success_effect_owned(state: &AppState, effect: OwnedLocalOAuthSuccessEffect) { + if !(200..300).contains(&effect.status_code) { + return; + } + let transport = match state + .read_provider_transport_snapshot(&effect.provider_id, &effect.endpoint_id, &effect.key_id) + .await + { + Ok(Some(transport)) => transport, + Ok(None) => return, + Err(err) => { + warn!( + "gateway orchestration effects: failed to read transport snapshot for oauth success provider {} endpoint {} key {}: {:?}", + effect.provider_id, effect.endpoint_id, effect.key_id, err + ); + return; + } + }; + if !transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + || crate::provider_transport::is_codex_agent_identity_transport(&transport) + || !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") + || crate::provider_transport::resolve_local_generic_oauth_transport_authorization( + &transport, + ) + .as_deref() + .and_then(bearer_access_token) + != bearer_access_token(effect.authorization.as_str()) + { + return; + } + + if let Err(err) = state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(effect.request_started_at_unix_ms), + Some(effect.request_order_id.as_str()), + effect.observed_credential_generation.as_deref(), + ) + .await + { + warn!( + "gateway orchestration effects: failed to persist oauth success for provider {} endpoint {} key {}: {:?}", + effect.provider_id, effect.endpoint_id, effect.key_id, err + ); + } +} + fn execution_plan_uses_codex_agent_identity(plan: &ExecutionPlan) -> bool { execution_plan_authorization(plan) .is_some_and(crate::provider_transport::is_codex_agent_identity_authorization) @@ -1802,6 +1963,7 @@ mod tests { StoredProviderCatalogProvider, }; use aether_test_support::ManagedRedisServer; + use aether_usage_runtime::GatewaySyncReportRequest; use serde_json::{json, Value}; use super::{ @@ -1811,11 +1973,13 @@ mod tests { pool_score_hard_state_for_status, resolve_pool_feedback_context, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, - LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, - ProviderKeyEffectLockPool, + LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, + LocalPoolErrorEffect, ProviderKeyEffectLockPool, }; use crate::data::{GatewayDataConfig, GatewayDataState}; - use crate::orchestration::LocalFailoverClassification; + use crate::orchestration::{ + apply_local_report_effect, LocalFailoverClassification, LocalReportEffect, + }; use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL; use crate::AppState; use aether_scheduler_core::{ @@ -3496,6 +3660,277 @@ mod tests { ); } + #[tokio::test] + async fn oauth_success_clears_recoverable_codex_invalid_state() { + let mut key = sample_codex_key(); + key.oauth_invalid_at_unix_secs = Some(100); + key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string()); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current", + "oauth_state_request_started_at_unix_ms": 100_000u64, + "oauth_state_request_id": "00000001-86a0-7000-8000-000000000001" + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let plan = sample_codex_plan(); + let report_context = json!({ + "codex_credential_generation": "credential-generation-current", + "provider_request_started_at_unix_ms": 200_000u64, + "provider_request_order_id": "00000003-0d40-7000-8000-000000000001" + }); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }, + LocalExecutionEffect::OauthSuccess(LocalOAuthSuccessEffect { + status_code: 200, + request_started_at_unix_ms: Some(200_000), + request_order_id: Some("00000003-0d40-7000-8000-000000000001"), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!(stored_key.oauth_invalid_at_unix_secs, None); + assert_eq!(stored_key.oauth_invalid_reason, None); + assert_eq!( + stored_key.upstream_metadata.as_ref().and_then( + |metadata| metadata.pointer("/codex/oauth_state_request_started_at_unix_ms") + ), + Some(&json!(200_000u64)) + ); + assert_eq!( + stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_id")), + Some(&json!("00000003-0d40-7000-8000-000000000001")) + ); + } + + #[tokio::test] + async fn same_millisecond_older_oauth_success_does_not_clear_newer_codex_invalidation() { + let mut key = sample_codex_key(); + key.oauth_invalid_at_unix_secs = Some(300); + key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string()); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current", + "oauth_state_request_started_at_unix_ms": 300_000u64, + "oauth_state_request_id": "00000004-93e0-7000-8000-000000000002" + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let plan = sample_codex_plan(); + let report_context = json!({ + "codex_credential_generation": "credential-generation-current", + "provider_request_started_at_unix_ms": 300_000u64, + "provider_request_order_id": "00000004-93e0-7000-8000-000000000001" + }); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }, + LocalExecutionEffect::OauthSuccess(LocalOAuthSuccessEffect { + status_code: 200, + request_started_at_unix_ms: Some(300_000), + request_order_id: Some("00000004-93e0-7000-8000-000000000001"), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!(stored_key.oauth_invalid_at_unix_secs, Some(300)); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] session expired") + ); + assert_eq!( + stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_id")), + Some(&json!("00000004-93e0-7000-8000-000000000002")) + ); + } + + #[tokio::test] + async fn oauth_success_preserves_codex_account_block() { + let mut key = sample_codex_key(); + key.oauth_invalid_at_unix_secs = Some(100); + key.oauth_invalid_reason = Some("[ACCOUNT_BLOCK] account deactivated".to_string()); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current" + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let plan = sample_codex_plan(); + let report_context = json!({ + "codex_credential_generation": "credential-generation-current", + "provider_request_started_at_unix_ms": 200_000u64, + "provider_request_order_id": "00000003-0d40-7000-8000-000000000001" + }); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }, + LocalExecutionEffect::OauthSuccess(LocalOAuthSuccessEffect { + status_code: 200, + request_started_at_unix_ms: Some(200_000), + request_order_id: Some("00000003-0d40-7000-8000-000000000001"), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!(stored_key.oauth_invalid_at_unix_secs, Some(100)); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[ACCOUNT_BLOCK] account deactivated") + ); + assert_eq!( + stored_key.upstream_metadata.as_ref().and_then( + |metadata| metadata.pointer("/codex/oauth_state_request_started_at_unix_ms") + ), + Some(&json!(200_000u64)) + ); + } + + #[tokio::test] + async fn oauth_success_does_not_clear_codex_invalid_state_for_replaced_bearer() { + let mut key = sample_codex_key(); + key.oauth_invalid_at_unix_secs = Some(100); + key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string()); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current" + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let mut plan = sample_codex_plan(); + plan.headers.insert( + "authorization".to_string(), + "Bearer replaced-access-token".to_string(), + ); + let report_context = json!({ + "codex_credential_generation": "credential-generation-current", + "provider_request_started_at_unix_ms": 200_000u64, + "provider_request_order_id": "00000003-0d40-7000-8000-000000000001" + }); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }, + LocalExecutionEffect::OauthSuccess(LocalOAuthSuccessEffect { + status_code: 200, + request_started_at_unix_ms: Some(200_000), + request_order_id: Some("00000003-0d40-7000-8000-000000000001"), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!(stored_key.oauth_invalid_at_unix_secs, Some(100)); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] session expired") + ); + assert!(stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_started_at_unix_ms")) + .is_none()); + } + + #[tokio::test] + async fn oauth_success_effect_ignores_non_success_status() { + let mut key = sample_codex_key(); + key.oauth_invalid_at_unix_secs = Some(100); + key.oauth_invalid_reason = Some("[OAUTH_EXPIRED] session expired".to_string()); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current" + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let plan = sample_codex_plan(); + let report_context = json!({ + "codex_credential_generation": "credential-generation-current", + "provider_request_started_at_unix_ms": 200_000u64, + "provider_request_order_id": "00000003-0d40-7000-8000-000000000001" + }); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }, + LocalExecutionEffect::OauthSuccess(LocalOAuthSuccessEffect { + status_code: 304, + request_started_at_unix_ms: Some(200_000), + request_order_id: Some("00000003-0d40-7000-8000-000000000001"), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!(stored_key.oauth_invalid_at_unix_secs, Some(100)); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] session expired") + ); + assert!(stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_started_at_unix_ms")) + .is_none()); + } + #[tokio::test] async fn oauth_invalidation_marks_generic_codex_403_as_token_invalid() { let state = codex_state(); @@ -3581,7 +4016,7 @@ mod tests { } #[tokio::test] - async fn oauth_invalidation_auto_removes_inactive_pat_owner() { + async fn oauth_invalidation_auto_remove_keeps_inactive_pat_owner() { let state = codex_state_with_auto_remove(); let plan = sample_codex_plan(); @@ -3600,6 +4035,39 @@ mod tests { ) .await; + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("recoverable token invalidation should retain the key"); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] Personal access token owner is inactive.") + ); + } + + #[tokio::test] + async fn oauth_invalidation_auto_removes_account_block() { + let state = codex_state_with_auto_remove(); + let plan = sample_codex_plan(); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: None, + }, + LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { + status_code: 403, + response_text: Some( + r#"{"error":{"message":"account has been deactivated"},"status":403}"#, + ), + }), + ) + .await; + let stored_keys = state .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) .await @@ -3709,6 +4177,155 @@ mod tests { assert_eq!(stored_key.oauth_invalid_reason, None); } + #[tokio::test] + async fn oauth_invalidation_rejects_stale_codex_credential_generation() { + let mut key = sample_codex_key(); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current" + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let plan = sample_codex_plan(); + let report_context = json!({ + "codex_credential_generation": "credential-generation-stale", + "provider_request_started_at_unix_ms": 100_000u64, + "provider_request_order_id": "00000001-86a0-7000-8000-000000000001" + }); + + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&report_context), + }, + LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { + status_code: 401, + response_text: Some(r#"{"error":{"message":"session expired"}}"#), + }), + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!(stored_key.oauth_invalid_at_unix_secs, None); + assert_eq!(stored_key.oauth_invalid_reason, None); + assert!(stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_started_at_unix_ms")) + .is_none()); + } + + #[tokio::test] + async fn delayed_codex_quota_report_does_not_clear_newer_oauth_invalidation() { + let mut key = sample_codex_key(); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-generation-current", + "account_quota_reset_generation": 0 + } + })); + let state = codex_state_with_provider_and_key(sample_codex_provider(), key); + let plan = sample_codex_plan(); + let older_request_id = "00000001-86a0-7000-8000-000000000001"; + let newer_request_id = "00000001-86a0-7000-8000-000000000002"; + let older_uuid = uuid::Uuid::parse_str(older_request_id).expect("older id should parse"); + let newer_uuid = uuid::Uuid::parse_str(newer_request_id).expect("newer id should parse"); + assert_eq!(older_uuid.get_version_num(), 7); + assert_eq!(newer_uuid.get_version_num(), 7); + assert!(older_uuid < newer_uuid); + + let invalidation_context = json!({ + "codex_credential_generation": "credential-generation-current", + "provider_request_started_at_unix_ms": 100_000u64, + "provider_request_order_id": newer_request_id + }); + apply_local_execution_effect( + &state, + LocalExecutionEffectContext { + plan: &plan, + report_context: Some(&invalidation_context), + }, + LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect { + status_code: 401, + response_text: Some(r#"{"error":{"message":"session expired"}}"#), + }), + ) + .await; + + let delayed_report_context = json!({ + "key_id": plan.key_id, + "codex_credential_generation": "credential-generation-current", + "codex_quota_reset_generation": 0, + "provider_request_started_at_unix_ms": 100_000u64, + "provider_response_headers_observed_at_unix_ms": 110_000u64, + "provider_request_order_id": older_request_id + }); + let delayed_headers = BTreeMap::from([ + ("x-codex-plan-type".to_string(), "free".to_string()), + ("x-codex-primary-used-percent".to_string(), "25".to_string()), + ( + "x-codex-primary-reset-at".to_string(), + "2000000000".to_string(), + ), + ( + "x-codex-primary-window-minutes".to_string(), + "300".to_string(), + ), + ]); + let delayed_report = GatewaySyncReportRequest { + trace_id: "trace-delayed-codex-quota".to_string(), + report_kind: "openai_responses_sync_success".to_string(), + report_context: Some(delayed_report_context), + status_code: 200, + headers: delayed_headers, + body_json: None, + client_body_json: None, + body_base64: None, + telemetry: None, + }; + apply_local_report_effect( + &state, + LocalReportEffect::Sync { + payload: &delayed_report, + }, + ) + .await; + + let stored_key = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id)) + .await + .expect("provider catalog keys should load") + .into_iter() + .next() + .expect("stored key should exist"); + assert_eq!( + stored_key.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] session expired") + ); + assert!(stored_key.oauth_invalid_at_unix_secs.is_some()); + assert_eq!( + stored_key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/oauth_state_request_id")), + Some(&json!(newer_request_id)) + ); + assert_eq!( + stored_key + .status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.pointer("/oauth/code")), + Some(&json!("expired")) + ); + } + #[tokio::test] async fn health_failure_updates_codex_key_for_current_bearer_request() { let state = codex_state(); diff --git a/apps/aether-gateway/src/orchestration/mod.rs b/apps/aether-gateway/src/orchestration/mod.rs index f79b25e8e..b3f842093 100644 --- a/apps/aether-gateway/src/orchestration/mod.rs +++ b/apps/aether-gateway/src/orchestration/mod.rs @@ -33,10 +33,10 @@ pub(crate) use self::classifier::{ LocalTransportFailoverClassification, }; pub(crate) use self::effects::{ - apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, - LocalAttemptFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, - LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, - LocalPoolErrorEffect, + apply_local_execution_effect, spawn_local_oauth_success_effect, LocalAdaptiveRateLimitEffect, + LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect, + LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect, + LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect, }; pub(crate) use self::health::{ project_local_failure_health, project_local_key_circuit_closed, diff --git a/apps/aether-gateway/src/orchestration/policy.rs b/apps/aether-gateway/src/orchestration/policy.rs index 8a5b9f166..2515f133c 100644 --- a/apps/aether-gateway/src/orchestration/policy.rs +++ b/apps/aether-gateway/src/orchestration/policy.rs @@ -227,6 +227,30 @@ pub(crate) fn append_local_failover_policy_to_value( "local_failover_policy".to_string(), local_failover_policy_to_value(&local_failover_policy_from_transport(transport)), ); + if transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + { + let codex = transport + .key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")); + object.insert( + "codex_quota_reset_generation".to_string(), + Value::from(aether_admin::provider::quota::codex_quota_account_reset_generation(codex)), + ); + if let Some(generation) = aether_admin::provider::quota::codex_credential_generation(codex) + { + object.insert( + "codex_credential_generation".to_string(), + Value::String(generation.to_string()), + ); + } + } Value::Object(object) } @@ -459,6 +483,26 @@ mod tests { ); } + #[test] + fn codex_report_context_captures_quota_and_credential_generations() { + let mut transport = sample_transport(None, None, None); + transport.provider.provider_type = "codex".to_string(); + transport.key.upstream_metadata = Some(json!({ + "codex": { + "account_quota_reset_generation": 7, + "credential_generation": "credential-generation-7" + } + })); + + let report_context = append_local_failover_policy_to_value(json!({}), &transport); + + assert_eq!(report_context["codex_quota_reset_generation"], json!(7u64)); + assert_eq!( + report_context["codex_credential_generation"], + json!("credential-generation-7") + ); + } + #[test] fn transport_error_failover_defaults_to_continue_and_accepts_explicit_stop() { let default_policy = diff --git a/apps/aether-gateway/src/orchestration/report_effects.rs b/apps/aether-gateway/src/orchestration/report_effects.rs index 9fc72c643..0ff62d451 100644 --- a/apps/aether-gateway/src/orchestration/report_effects.rs +++ b/apps/aether-gateway/src/orchestration/report_effects.rs @@ -1,6 +1,6 @@ -use std::collections::{BTreeMap, HashMap}; -use std::sync::{Mutex, OnceLock}; -use std::time::{Duration, Instant}; +use std::collections::BTreeMap; +use std::sync::OnceLock; +use std::time::Duration; use aether_admin::provider::quota as admin_provider_quota_pure; use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyRuntimeMetadataUpdate; @@ -21,13 +21,7 @@ use crate::handlers::shared::sync_provider_key_quota_status_snapshot; use crate::log_ids::short_request_id; use crate::{AppState, GatewayError}; -const CODEX_QUOTA_CACHE_TTL_SECONDS: u64 = 30; -const CODEX_QUOTA_CACHE_MAX_ENTRIES: usize = 4096; const RUNTIME_METADATA_CAS_MAX_ATTEMPTS: usize = 16; - -type HeaderFingerprintCache = Mutex>; - -static CODEX_QUOTA_HEADER_FINGERPRINT_CACHE: OnceLock = OnceLock::new(); static GROK_CHINESE_WAIT_DURATION_RE: OnceLock = OnceLock::new(); static GROK_ENGLISH_WAIT_DURATION_RE: OnceLock = OnceLock::new(); @@ -62,10 +56,6 @@ pub(crate) async fn apply_local_report_effect(state: &AppState, effect: LocalRep } } -fn codex_quota_header_fingerprint_cache() -> &'static HeaderFingerprintCache { - CODEX_QUOTA_HEADER_FINGERPRINT_CACHE.get_or_init(|| Mutex::new(HashMap::new())) -} - fn report_context_key_id(report_context: Option<&Value>) -> Option { report_context .and_then(|context| context.get("key_id")) @@ -75,6 +65,20 @@ fn report_context_key_id(report_context: Option<&Value>) -> Option { .map(ToOwned::to_owned) } +fn report_context_u64(report_context: Option<&Value>, key: &str) -> Option { + report_context + .and_then(|context| context.get(key)) + .and_then(admin_provider_quota_pure::coerce_json_u64) +} + +fn report_context_string<'a>(report_context: Option<&'a Value>, key: &str) -> Option<&'a str> { + report_context + .and_then(|context| context.get(key)) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + fn report_context_provider_response_headers( report_context: Option<&Value>, ) -> Option> { @@ -91,86 +95,6 @@ fn report_context_provider_response_headers( (!out.is_empty()).then_some(out) } -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::>(); - 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 { - let object = value.as_object()?; - let mut entries = object - .iter() - .filter(|(key, _)| !is_volatile_compare_field(key)) - .collect::>(); - 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_codex_quota_fingerprint(key_id: &str, now: Instant) -> Option { - let mut cache = codex_quota_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_codex_quota_fingerprint(key_id: &str, fingerprint: String, now: Instant) { - let mut cache = codex_quota_header_fingerprint_cache() - .lock() - .expect("codex realtime quota cache should lock"); - cache.insert( - key_id.to_string(), - ( - fingerprint, - now.checked_add(Duration::from_secs(CODEX_QUOTA_CACHE_TTL_SECONDS)) - .unwrap_or(now), - ), - ); - - cache.retain(|_, (_, expires_at)| *expires_at > now); - if cache.len() <= CODEX_QUOTA_CACHE_MAX_ENTRIES { - return; - } - - let mut entries = cache - .iter() - .map(|(key, (_, expires_at))| (key.clone(), *expires_at)) - .collect::>(); - entries.sort_by_key(|entry| entry.1); - for (key, _) in entries - .into_iter() - .take(cache.len() - CODEX_QUOTA_CACHE_MAX_ENTRIES) - { - cache.remove(&key); - } -} - fn merge_metadata_object( current: Option<&Value>, section_key: &str, @@ -593,21 +517,23 @@ async fn sync_grok_quota_from_report_context( async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncReportRequest) { apply_local_gemini_file_mapping_report_effect(state, payload).await; - if let Err(err) = 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" - ); + if (200..300).contains(&payload.status_code) { + if let Err(err) = 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" + ); + } } if let Err(err) = sync_grok_quota_from_report_context( state, @@ -646,21 +572,23 @@ async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncR } async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStreamReportRequest) { - if let Err(err) = 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" - ); + if (200..300).contains(&payload.status_code) { + if let Err(err) = 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" + ); + } } if let Err(err) = sync_grok_quota_from_report_context( state, @@ -840,103 +768,112 @@ async fn sync_codex_quota_from_response_headers( }; let now_unix_secs = current_unix_secs(); + let observed_at_unix_secs = report_context_u64( + report_context, + "provider_response_headers_observed_at_unix_ms", + ) + .map(|value| value / 1_000) + .filter(|value| *value > 0) + .unwrap_or(now_unix_secs); + let request_started_at_unix_ms = + report_context_u64(report_context, "provider_request_started_at_unix_ms"); + let request_order_id = report_context_string(report_context, "provider_request_order_id"); + let observed_reset_generation = + report_context_u64(report_context, "codex_quota_reset_generation"); + let observed_credential_generation = + report_context_string(report_context, "codex_credential_generation"); let provider_headers = report_context_provider_response_headers(report_context); let parsed_from_provider_headers = provider_headers.as_ref().and_then(|headers| { - admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs) + admin_provider_quota_pure::parse_codex_usage_headers(headers, observed_at_unix_secs) }); - let Some(parsed) = parsed_from_provider_headers - .or_else(|| 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 { + let Some(parsed) = parsed_from_provider_headers.or_else(|| { + admin_provider_quota_pure::parse_codex_usage_headers(headers, observed_at_unix_secs) + }) else { return Ok(false); }; + // Runtime headers can be partial (for example only the primary window), + // so absence never authoritatively removes another stored window. + let coverage = admin_provider_quota_pure::CodexQuotaWindowCoverage::Patch; - let now = Instant::now(); - if get_cached_codex_quota_fingerprint(&key_id, now).as_deref() - == Some(incoming_fingerprint.as_str()) - { - return Ok(false); + for attempt in 0..RUNTIME_METADATA_CAS_MAX_ATTEMPTS { + let Some(key) = state + .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id)) + .await? + .into_iter() + .next() + else { + 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 { + return Ok(false); + }; + if !provider.provider_type.trim().eq_ignore_ascii_case("codex") { + return Ok(false); + } + + let expected_namespace_value = + upstream_metadata_namespace_value(key.upstream_metadata.as_ref(), "codex"); + let Some(outcome) = admin_provider_quota_pure::merge_codex_quota_metadata_snapshot( + expected_namespace_value.as_ref(), + &parsed, + admin_provider_quota_pure::CodexQuotaMergeContext { + observed_at_unix_secs, + request_started_at_unix_ms, + request_order_id, + observed_reset_generation, + authoritative_reset_generation: None, + observed_credential_generation, + account_reset_fence_id: None, + coverage, + }, + ) else { + return Ok(false); + }; + if !outcome.changed { + return Ok(false); + } + let next_codex = outcome.metadata; + let updated_upstream_metadata = + merge_metadata_object(key.upstream_metadata.as_ref(), "codex", next_codex.clone()); + 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 updated = state + .update_provider_catalog_key_runtime_metadata( + &ProviderCatalogKeyRuntimeMetadataUpdate { + key_id: key_id.clone(), + namespace: "codex".to_string(), + expected_upstream_metadata_value: expected_namespace_value, + upstream_metadata_value: next_codex, + status_snapshot_patch: quota_status_snapshot_patch( + updated_status_snapshot.as_ref(), + ), + updated_at_unix_secs: Some(observed_at_unix_secs), + }, + ) + .await?; + if updated { + return Ok(true); + } + if attempt + 1 < RUNTIME_METADATA_CAS_MAX_ATTEMPTS { + let backoff_us = 50_u64.saturating_mul((attempt + 1) as u64).min(1_000); + tokio::time::sleep(Duration::from_micros(backoff_us)).await; + } } - - let Some(key) = state - .read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id)) - .await? - .into_iter() - .next() - else { - set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint.clone(), 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_codex_quota_fingerprint(&key_id, incoming_fingerprint.clone(), now); - return Ok(false); - }; - if !provider.provider_type.trim().eq_ignore_ascii_case("codex") { - set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint.clone(), now); - return Ok(false); - } - - let expected_namespace_value = - upstream_metadata_namespace_value(key.upstream_metadata.as_ref(), "codex"); - let current_codex = expected_namespace_value - .clone() - .and_then(|value| value.as_object().cloned()) - .unwrap_or_default(); - let current_codex = Value::Object(current_codex); - let Some(current_fingerprint) = fingerprint_codex_payload(¤t_codex) else { - set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint.clone(), now); - return Ok(false); - }; - if current_fingerprint == incoming_fingerprint { - set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint.clone(), now); - return Ok(false); - } - - let updated_upstream_metadata = - merge_metadata_object(key.upstream_metadata.as_ref(), "codex", parsed.clone()); - 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 updated = state - .update_provider_catalog_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate { - key_id: key_id.clone(), - namespace: "codex".to_string(), - expected_upstream_metadata_value: expected_namespace_value, - upstream_metadata_value: parsed.clone(), - status_snapshot_patch: quota_status_snapshot_patch(updated_status_snapshot.as_ref()), - updated_at_unix_secs: Some(now_unix_secs), - }) - .await?; - if updated { - set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint.clone(), now); - return Ok(true); - } - // Response headers describe an authoritative quota snapshot. A CAS - // conflict means a newer snapshot/delta won, so avoid replaying stale - // data over it. Ok(false) } #[cfg(test)] -pub(crate) fn clear_local_report_effect_caches_for_tests() { - if let Some(cache) = CODEX_QUOTA_HEADER_FINGERPRINT_CACHE.get() { - cache - .lock() - .expect("codex realtime quota cache should lock") - .clear(); - } -} +pub(crate) fn clear_local_report_effect_caches_for_tests() {} #[cfg(test)] mod tests { @@ -950,6 +887,151 @@ mod tests { use crate::data::GatewayDataState; + fn codex_headers(used_percent: f64, reset_at: u64) -> BTreeMap { + BTreeMap::from([ + ("x-codex-plan-type".to_string(), "free".to_string()), + ( + "x-codex-primary-used-percent".to_string(), + used_percent.to_string(), + ), + ("x-codex-primary-reset-at".to_string(), reset_at.to_string()), + ( + "x-codex-primary-window-minutes".to_string(), + "300".to_string(), + ), + ]) + } + + fn codex_test_state(key_id: &str) -> AppState { + let provider = StoredProviderCatalogProvider::new( + "codex-provider".to_string(), + "Codex".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"); + let key = StoredProviderCatalogKey::new( + key_id.to_string(), + provider.id.clone(), + "Codex Key".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![], + vec![key], + )); + AppState::new() + .expect("gateway state should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(repository), + ) + } + + async fn stored_codex_metadata(state: &AppState, key_id: &str) -> Value { + state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .and_then(|key| key.upstream_metadata) + .and_then(|metadata| metadata.get("codex").cloned()) + .expect("codex metadata should exist") + } + + #[tokio::test] + async fn out_of_order_codex_quota_reports_keep_max_usage_and_allow_next_reset() { + clear_local_report_effect_caches_for_tests(); + let key_id = "codex-out-of-order-quota-key"; + let state = codex_test_state(key_id); + let report_context = json!({"key_id": key_id}); + let reset_at = 2_000_000_000u64; + let higher = codex_headers(60.0, reset_at); + let lower = codex_headers(50.0, reset_at); + + let (higher_result, lower_result) = tokio::join!( + sync_codex_quota_from_response_headers(&state, Some(&report_context), &higher), + sync_codex_quota_from_response_headers(&state, Some(&report_context), &lower), + ); + higher_result.expect("higher quota report should complete"); + lower_result.expect("lower quota report should complete"); + + let stored = stored_codex_metadata(&state, key_id).await; + assert_eq!(stored["primary_used_percent"], json!(60.0)); + assert_eq!(stored["primary_reset_at"], json!(reset_at)); + + let next_reset_at = reset_at + 18_000; + assert!(sync_codex_quota_from_response_headers( + &state, + Some(&report_context), + &codex_headers(2.0, next_reset_at), + ) + .await + .expect("new reset quota report should complete")); + let reset = stored_codex_metadata(&state, key_id).await; + assert_eq!(reset["primary_used_percent"], json!(2.0)); + assert_eq!(reset["primary_reset_at"], json!(next_reset_at)); + } + + #[tokio::test] + async fn non_success_sync_report_does_not_persist_codex_quota_headers() { + let key_id = "codex-non-success-sync-report"; + let state = codex_test_state(key_id); + let payload = GatewaySyncReportRequest { + trace_id: "trace-non-success-sync".to_string(), + report_kind: "openai_responses_sync_error".to_string(), + report_context: Some(json!({"key_id": key_id})), + status_code: 401, + headers: codex_headers(85.0, 2_000_000_000), + body_json: None, + client_body_json: None, + body_base64: None, + telemetry: None, + }; + + apply_local_report_effect(&state, LocalReportEffect::Sync { payload: &payload }).await; + + let stored = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert!(stored.upstream_metadata.is_none()); + } + + #[tokio::test] + async fn non_success_stream_report_does_not_persist_codex_quota_headers() { + let key_id = "codex-non-success-stream-report"; + let state = codex_test_state(key_id); + let payload = GatewayStreamReportRequest { + trace_id: "trace-non-success-stream".to_string(), + report_kind: "openai_responses_stream_error".to_string(), + report_context: Some(json!({"key_id": key_id})), + status_code: 429, + headers: codex_headers(95.0, 2_000_000_000), + provider_body_base64: None, + provider_body_state: None, + client_body_base64: None, + client_body_state: None, + terminal_summary: None, + telemetry: None, + }; + + apply_local_report_effect(&state, LocalReportEffect::Stream { payload: &payload }).await; + + let stored = state + .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert!(stored.upstream_metadata.is_none()); + } + #[tokio::test] async fn gemini_report_metadata_write_preserves_adaptive_and_other_provider_state() { let provider = StoredProviderCatalogProvider::new( diff --git a/apps/aether-gateway/src/state/catalog.rs b/apps/aether-gateway/src/state/catalog.rs index a809aed38..3aeed16c9 100644 --- a/apps/aether-gateway/src/state/catalog.rs +++ b/apps/aether-gateway/src/state/catalog.rs @@ -662,6 +662,21 @@ impl AppState { Ok(updated) } + pub(crate) async fn compare_and_update_provider_catalog_key_admin_state( + &self, + update: &provider_catalog::ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + let updated = self + .data + .compare_and_update_provider_catalog_key_admin_state(update) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))?; + // A conflict means another instance changed credentials. Invalidate on + // both outcomes before the caller reloads or reports the conflict. + self.invalidate_provider_routing_caches(); + Ok(updated) + } + pub(crate) async fn update_provider_catalog_keys( &self, keys: &[provider_catalog::StoredProviderCatalogKey], diff --git a/apps/aether-gateway/src/state/mod.rs b/apps/aether-gateway/src/state/mod.rs index f8b19ab19..c342a7c52 100644 --- a/apps/aether-gateway/src/state/mod.rs +++ b/apps/aether-gateway/src/state/mod.rs @@ -38,7 +38,10 @@ pub(crate) use self::cache::{ PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL, }; pub use self::cors::FrontdoorCorsConfig; -pub(crate) use self::oauth::AgentIdentityAuthConfigFence; +pub(crate) use self::oauth::{ + provider_transport_context_allows_credential_rotation, AgentIdentityAuthConfigFence, + CodexRuntimeOAuthObservation, ProviderTransportCredentialFence, +}; pub(crate) use self::types::{ AdminWalletMutationOutcome, GatewayAdminPaymentCallbackView, GatewayUserPreferenceView, GatewayUserSessionView, LocalExecutionRuntimeMissDiagnostic, LocalMutationOutcome, diff --git a/apps/aether-gateway/src/state/oauth.rs b/apps/aether-gateway/src/state/oauth.rs index c5b6507e8..1f53e3a48 100644 --- a/apps/aether-gateway/src/state/oauth.rs +++ b/apps/aether-gateway/src/state/oauth.rs @@ -19,7 +19,7 @@ use aether_contracts::{ use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - StoredProviderCatalogKey, + ProviderCatalogUpstreamMetadataNamespaceExpectation, StoredProviderCatalogKey, }; use aether_runtime_state::RuntimeLockLease; use base64::{engine::general_purpose::STANDARD, Engine as _}; @@ -42,11 +42,20 @@ const OAUTH_ACCOUNT_BLOCK_PREFIX: &str = "[ACCOUNT_BLOCK] "; const OAUTH_EXPIRED_PREFIX: &str = "[OAUTH_EXPIRED] "; const OAUTH_REFRESH_FAILED_PREFIX: &str = "[REFRESH_FAILED] "; const OAUTH_REQUEST_FAILED_PREFIX: &str = "[REQUEST_FAILED] "; +const CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS: usize = 16; #[derive(Debug, Clone, PartialEq)] -struct ProviderTransportCredentialFence { - encrypted_auth_config: String, - credential: ProviderCatalogKeyOAuthCredentialFence, +pub(crate) struct ProviderTransportCredentialFence { + pub(crate) encrypted_auth_config: String, + pub(crate) credential: ProviderCatalogKeyOAuthCredentialFence, +} + +#[derive(Debug, Clone, Copy)] +pub(crate) struct CodexRuntimeOAuthObservation<'a> { + pub(crate) request_started_at_unix_ms: u64, + pub(crate) request_order_id: &'a str, + pub(crate) observed_credential_generation: Option<&'a str>, + pub(crate) runtime_invalid_reason: Option<&'a str>, } struct GatewayLocalOAuthHttpExecutor<'a> { @@ -183,6 +192,18 @@ fn local_oauth_transport_context_allows_reload( true } +pub(crate) fn provider_transport_context_allows_credential_rotation( + initial: &provider_transport::GatewayProviderTransportSnapshot, + current: &provider_transport::GatewayProviderTransportSnapshot, +) -> bool { + initial.provider.id == current.provider.id + && initial.endpoint.id == current.endpoint.id + && initial.key.id == current.key.id + && initial.key.auth_type == current.key.auth_type + && initial.provider.provider_type == current.provider.provider_type + && local_oauth_transport_context_allows_reload(initial, current) +} + fn discard_failed_local_oauth_refresh_resolution( resolution: &mut Option, ) { @@ -1757,7 +1778,7 @@ impl AppState { .map(|fence| fence.encrypted_auth_config)) } - async fn capture_provider_transport_credential_fence( + pub(crate) async fn capture_provider_transport_credential_fence( &self, transport: &provider_transport::GatewayProviderTransportSnapshot, ) -> Result, GatewayError> { @@ -1924,6 +1945,9 @@ impl AppState { &self, transport: &provider_transport::GatewayProviderTransportSnapshot, invalid_reason: &str, + request_started_at_unix_ms: Option, + request_order_id: Option<&str>, + observed_credential_generation: Option<&str>, ) -> Result { let invalid_reason = invalid_reason.trim(); let key_id = transport.key.id.trim(); @@ -1937,108 +1961,348 @@ impl AppState { else { return Ok(false); }; - - let Some(mut latest_key) = self - .data - .list_provider_catalog_keys_by_ids(&[key_id.to_string()]) - .await - .map_err(|err| GatewayError::Internal(err.to_string()))? - .into_iter() - .next() - else { - return Ok(false); - }; - if latest_key.encrypted_auth_config.as_deref() - != Some(expected_credential_fence.encrypted_auth_config.as_str()) - || latest_key.encrypted_api_key - != expected_credential_fence.credential.encrypted_api_key - || latest_key.auth_type != expected_credential_fence.credential.auth_type - || latest_key.provider_id != expected_credential_fence.credential.provider_id - || !provider_key_is_oauth_managed(&latest_key, provider_type) - { - return Ok(false); - } - + let request_order_id = request_order_id + .map(str::trim) + .filter(|value| !value.is_empty()); + let is_codex = provider_type.trim().eq_ignore_ascii_case("codex"); let now_unix_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .ok() .map(|duration| duration.as_secs()) .unwrap_or(0); - let (oauth_invalid_at_unix_secs, oauth_invalid_reason) = merge_runtime_oauth_invalid_state( - provider_type, - &latest_key, - invalid_reason, - now_unix_secs, - ); - if oauth_invalid_at_unix_secs == latest_key.oauth_invalid_at_unix_secs - && oauth_invalid_reason == latest_key.oauth_invalid_reason + for attempt in 0..CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS { + let Some(mut latest_key) = self + .data + .list_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))? + .into_iter() + .next() + else { + return Ok(false); + }; + if latest_key.encrypted_auth_config.as_deref() + != Some(expected_credential_fence.encrypted_auth_config.as_str()) + || latest_key.encrypted_api_key + != expected_credential_fence.credential.encrypted_api_key + || latest_key.auth_type != expected_credential_fence.credential.auth_type + || latest_key.provider_id != expected_credential_fence.credential.provider_id + || !provider_key_is_oauth_managed(&latest_key, provider_type) + { + return Ok(false); + } + + let expected_codex = is_codex + .then(|| { + latest_key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .cloned() + }) + .flatten(); + let expected_codex_object = expected_codex.as_ref().and_then(Value::as_object); + if is_codex + && !admin_provider_quota_pure::codex_credential_generation_matches( + expected_codex.as_ref(), + observed_credential_generation, + ) + { + return Ok(false); + } + if is_codex + && request_started_at_unix_ms.is_some() + && admin_provider_quota_pure::codex_oauth_state_request_order_is_stale( + expected_codex_object, + request_started_at_unix_ms, + request_order_id, + ) + { + return Ok(false); + } + + let (oauth_invalid_at_unix_secs, oauth_invalid_reason) = + merge_runtime_oauth_invalid_state( + provider_type, + &latest_key, + invalid_reason, + now_unix_secs, + ); + let invalid_state_changed = oauth_invalid_at_unix_secs + != latest_key.oauth_invalid_at_unix_secs + || oauth_invalid_reason != latest_key.oauth_invalid_reason; + let next_codex = (is_codex && request_started_at_unix_ms.is_some()).then(|| { + let mut codex = expected_codex_object.cloned().unwrap_or_default(); + codex.insert( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY.to_string(), + json!(request_started_at_unix_ms.unwrap_or_default()), + ); + if let Some(request_order_id) = request_order_id { + codex.insert( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY + .to_string(), + json!(request_order_id), + ); + } else { + codex.remove( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY, + ); + } + Value::Object(codex) + }); + if !invalid_state_changed && next_codex.is_none() { + return Ok(false); + } + + latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs; + latest_key.oauth_invalid_reason = oauth_invalid_reason; + latest_key.updated_at_unix_secs = Some(now_unix_secs); + let current_status_snapshot = latest_key.status_snapshot.take(); + latest_key.status_snapshot = + sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key); + + let updated = self + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.to_string(), + expected_encrypted_auth_config: Some( + expected_credential_fence.encrypted_auth_config.clone(), + ), + expected_credential: Some(expected_credential_fence.credential.clone()), + expected_upstream_metadata_namespace: is_codex.then(|| { + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex.clone(), + } + }), + encrypted_auth_config: expected_credential_fence + .encrypted_auth_config + .clone(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), + upstream_metadata_patch: next_codex + .as_ref() + .map(|codex| json!({"codex": codex})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: provider_key_oauth_status_snapshot_update( + &latest_key, + ) + .status_snapshot_patch, + reset_error_count: false, + updated_at_unix_secs: latest_key.updated_at_unix_secs, + }, + ) + .await?; + if !updated { + if attempt + 1 < CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS { + tokio::task::yield_now().await; + } + continue; + } + + let delete_namespace_expectation = + is_codex.then(|| ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: next_codex.or(expected_codex), + }); + let auto_removed = if admin_provider_quota_pure::provider_auto_remove_banned_keys( + transport.provider.config.as_ref(), + ) && (!is_codex + || oauth_invalid_reason_is_account_block( + latest_key.oauth_invalid_reason.as_deref(), + )) + && admin_provider_quota_pure::should_auto_remove_oauth_invalid_key( + &latest_key, + None, + true, + now_unix_secs, + ) { + self.delete_provider_transport_oauth_credential_fenced( + key_id, + &expected_credential_fence, + delete_namespace_expectation, + ) + .await? + } else { + false + }; + let _ = self.invalidate_local_oauth_refresh_entry(key_id).await; + tracing::info!( + key_id, + provider_id = %transport.provider.id, + provider_type, + updated = true, + auto_removed, + "gateway fenced OAuth invalidation persisted" + ); + return Ok(true); + } + Ok(false) + } + + pub(crate) async fn mark_provider_transport_oauth_success_fenced( + &self, + transport: &provider_transport::GatewayProviderTransportSnapshot, + request_started_at_unix_ms: Option, + request_order_id: Option<&str>, + observed_credential_generation: Option<&str>, + ) -> Result { + let key_id = transport.key.id.trim(); + if key_id.is_empty() + || !transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + || !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") + || provider_transport::is_codex_agent_identity_transport(transport) { return Ok(false); } - latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs; - latest_key.oauth_invalid_reason = oauth_invalid_reason; - latest_key.updated_at_unix_secs = Some(now_unix_secs); - let current_status_snapshot = latest_key.status_snapshot.take(); - latest_key.status_snapshot = - sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key); - - let updated = self - .compare_and_update_provider_catalog_key_oauth_runtime_state( - &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { - key_id: key_id.to_string(), - expected_encrypted_auth_config: Some( - expected_credential_fence.encrypted_auth_config.clone(), - ), - expected_credential: Some(expected_credential_fence.credential.clone()), - encrypted_auth_config: expected_credential_fence.encrypted_auth_config.clone(), - encrypted_api_key_update: None, - expires_at_unix_secs_update: None, - oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, - oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), - upstream_metadata_patch: None, - status_snapshot_patch: provider_key_oauth_status_snapshot_update(&latest_key) - .status_snapshot_patch, - reset_error_count: false, - updated_at_unix_secs: latest_key.updated_at_unix_secs, - }, - ) - .await?; - let auto_removed = if updated - && admin_provider_quota_pure::provider_auto_remove_banned_keys( - transport.provider.config.as_ref(), - ) - && admin_provider_quota_pure::should_auto_remove_oauth_invalid_key( - &latest_key, - None, - true, - now_unix_secs, - ) { - self.delete_provider_transport_oauth_credential_fenced( - key_id, - &expected_credential_fence, - ) - .await? - } else { - false + let Some(request_started_at_unix_ms) = request_started_at_unix_ms else { + // Without an order fence, an older 2xx could clear a newer 401. + return Ok(false); }; - if updated { - let _ = self.invalidate_local_oauth_refresh_entry(key_id).await; + let Some(request_order_id) = request_order_id + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + return Ok(false); + }; + let Some(expected_credential_fence) = self + .capture_provider_transport_credential_fence(transport) + .await? + else { + return Ok(false); + }; + let now_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()) + .unwrap_or(0); + + for attempt in 0..CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS { + let Some(mut latest_key) = self + .data + .list_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))? + .into_iter() + .next() + else { + return Ok(false); + }; + if latest_key.encrypted_auth_config.as_deref() + != Some(expected_credential_fence.encrypted_auth_config.as_str()) + || latest_key.encrypted_api_key + != expected_credential_fence.credential.encrypted_api_key + || latest_key.auth_type != expected_credential_fence.credential.auth_type + || latest_key.provider_id != expected_credential_fence.credential.provider_id + || !provider_key_is_oauth_managed(&latest_key, "codex") + { + return Ok(false); + } + + let expected_codex = match latest_key.upstream_metadata.as_ref() { + Some(Value::Object(metadata)) => metadata.get("codex").cloned(), + Some(_) => return Ok(false), + None => None, + }; + let expected_codex_object = match expected_codex.as_ref() { + Some(Value::Object(codex)) => Some(codex), + Some(_) => return Ok(false), + None => None, + }; + if !admin_provider_quota_pure::codex_credential_generation_matches( + expected_codex.as_ref(), + observed_credential_generation, + ) || admin_provider_quota_pure::codex_oauth_success_request_order_is_stale( + expected_codex_object, + Some(request_started_at_unix_ms), + Some(request_order_id), + ) { + return Ok(false); + } + + let mut next_codex = expected_codex_object.cloned().unwrap_or_default(); + next_codex.insert( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY.to_string(), + json!(request_started_at_unix_ms), + ); + next_codex.insert( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY.to_string(), + json!(request_order_id), + ); + let next_codex = Value::Object(next_codex); + + ( + latest_key.oauth_invalid_at_unix_secs, + latest_key.oauth_invalid_reason, + ) = local_oauth_refresh_success_invalid_state(&latest_key); + latest_key.updated_at_unix_secs = Some(now_unix_secs); + let current_status_snapshot = latest_key.status_snapshot.take(); + latest_key.status_snapshot = + sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key); + + let updated = self + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.to_string(), + expected_encrypted_auth_config: Some( + expected_credential_fence.encrypted_auth_config.clone(), + ), + expected_credential: Some(expected_credential_fence.credential.clone()), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex, + }, + ), + encrypted_auth_config: expected_credential_fence + .encrypted_auth_config + .clone(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), + upstream_metadata_patch: Some(json!({"codex": next_codex})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: provider_key_oauth_status_snapshot_update( + &latest_key, + ) + .status_snapshot_patch, + reset_error_count: false, + updated_at_unix_secs: latest_key.updated_at_unix_secs, + }, + ) + .await?; + if updated { + tracing::info!( + key_id, + provider_id = %transport.provider.id, + request_started_at_unix_ms, + request_order_id, + "gateway fenced OAuth success persisted" + ); + return Ok(true); + } + if attempt + 1 < CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS { + tokio::task::yield_now().await; + } } - tracing::info!( - key_id, - provider_id = %transport.provider.id, - provider_type, - updated, - auto_removed, - "gateway fenced OAuth invalidation persisted" - ); - Ok(updated) + Ok(false) } async fn delete_provider_transport_oauth_credential_fenced( &self, key_id: &str, expected: &ProviderTransportCredentialFence, + expected_upstream_metadata_namespace: Option< + ProviderCatalogUpstreamMetadataNamespaceExpectation, + >, ) -> Result { let deleted = self .compare_and_delete_provider_catalog_key_oauth_credential( @@ -2046,6 +2310,7 @@ impl AppState { key_id: key_id.to_string(), expected_encrypted_auth_config: Some(expected.encrypted_auth_config.clone()), expected_credential: expected.credential.clone(), + expected_upstream_metadata_namespace, }, ) .await?; @@ -2215,12 +2480,14 @@ impl AppState { key_id: key_id.to_string(), expected_encrypted_auth_config: Some(expected_encrypted_auth_config), expected_credential: Some(expected_credential_fence.credential.clone()), + expected_upstream_metadata_namespace: None, encrypted_auth_config, encrypted_api_key_update: None, expires_at_unix_secs_update: None, oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), upstream_metadata_patch: None, + upstream_metadata_namespace_to_remove: None, status_snapshot_patch: provider_key_oauth_status_snapshot_update( &latest_key, ) @@ -2329,12 +2596,14 @@ impl AppState { expected_credential_fence.encrypted_auth_config.clone(), ), expected_credential: Some(expected_credential_fence.credential.clone()), + expected_upstream_metadata_namespace: None, encrypted_auth_config: encrypted_auth_config.to_string(), encrypted_api_key_update: Some(encrypted_api_key.clone()), expires_at_unix_secs_update: Some(entry.expires_at_unix_secs), oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), upstream_metadata_patch: None, + upstream_metadata_namespace_to_remove: None, status_snapshot_patch: provider_key_oauth_status_snapshot_update( &latest_key, ) @@ -2404,6 +2673,205 @@ impl AppState { Ok(()) } + pub(crate) async fn persist_local_oauth_refresh_failure_state_observed( + &self, + transport: &provider_transport::GatewayProviderTransportSnapshot, + status_code: u16, + body_excerpt: &str, + access_token_invalid_proven: bool, + observation: CodexRuntimeOAuthObservation<'_>, + ) -> Result { + let key_id = transport.key.id.trim(); + let request_order_id = observation.request_order_id.trim(); + if key_id.is_empty() + || request_order_id.is_empty() + || !transport + .provider + .provider_type + .trim() + .eq_ignore_ascii_case("codex") + || !transport.key.auth_type.trim().eq_ignore_ascii_case("oauth") + || provider_transport::is_codex_agent_identity_transport(transport) + { + return Ok(false); + } + let Some(expected_credential_fence) = self + .capture_provider_transport_credential_fence(transport) + .await? + else { + return Ok(false); + }; + let refresh_reason = format!( + "{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 ({status_code}): {}", + normalize_local_oauth_refresh_error_message(Some(status_code), Some(body_excerpt)) + ); + let now_unix_secs = SystemTime::now() + .duration_since(UNIX_EPOCH) + .ok() + .map(|duration| duration.as_secs()) + .unwrap_or(0); + + for attempt in 0..CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS { + let Some(mut latest_key) = self + .data + .list_provider_catalog_keys_by_ids(&[key_id.to_string()]) + .await + .map_err(|err| GatewayError::Internal(err.to_string()))? + .into_iter() + .next() + else { + return Ok(false); + }; + if latest_key.encrypted_auth_config.as_deref() + != Some(expected_credential_fence.encrypted_auth_config.as_str()) + || latest_key.encrypted_api_key + != expected_credential_fence.credential.encrypted_api_key + || latest_key.auth_type != expected_credential_fence.credential.auth_type + || latest_key.provider_id != expected_credential_fence.credential.provider_id + || !provider_key_is_oauth_managed(&latest_key, "codex") + { + return Ok(false); + } + + let expected_codex = match latest_key.upstream_metadata.as_ref() { + Some(Value::Object(metadata)) => metadata.get("codex").cloned(), + Some(_) => return Ok(false), + None => None, + }; + let expected_codex_object = match expected_codex.as_ref() { + Some(Value::Object(codex)) => Some(codex), + Some(_) => return Ok(false), + None => None, + }; + if !admin_provider_quota_pure::codex_credential_generation_matches( + expected_codex.as_ref(), + observation.observed_credential_generation, + ) || admin_provider_quota_pure::codex_oauth_state_request_order_is_stale( + expected_codex_object, + Some(observation.request_started_at_unix_ms), + Some(request_order_id), + ) { + return Ok(false); + } + + let mut next_codex = expected_codex_object.cloned().unwrap_or_default(); + next_codex.insert( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY.to_string(), + json!(observation.request_started_at_unix_ms), + ); + next_codex.insert( + admin_provider_quota_pure::CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY.to_string(), + json!(request_order_id), + ); + let next_codex = Value::Object(next_codex); + + if let Some(runtime_invalid_reason) = observation + .runtime_invalid_reason + .map(str::trim) + .filter(|value| !value.is_empty()) + { + ( + latest_key.oauth_invalid_at_unix_secs, + latest_key.oauth_invalid_reason, + ) = merge_runtime_oauth_invalid_state( + "codex", + &latest_key, + runtime_invalid_reason, + now_unix_secs, + ); + } + if let Some(merged_reason) = merge_local_oauth_refresh_failure_reason( + latest_key.oauth_invalid_reason.as_deref(), + &refresh_reason, + ) { + latest_key.oauth_invalid_at_unix_secs = latest_key + .oauth_invalid_at_unix_secs + .or(Some(now_unix_secs)); + latest_key.oauth_invalid_reason = Some(merged_reason); + } + latest_key.updated_at_unix_secs = Some(now_unix_secs); + let current_status_snapshot = latest_key.status_snapshot.take(); + latest_key.status_snapshot = + sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key); + + let updated = self + .compare_and_update_provider_catalog_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.to_string(), + expected_encrypted_auth_config: Some( + expected_credential_fence.encrypted_auth_config.clone(), + ), + expected_credential: Some(expected_credential_fence.credential.clone()), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: expected_codex, + }, + ), + encrypted_auth_config: expected_credential_fence + .encrypted_auth_config + .clone(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, + oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), + upstream_metadata_patch: Some(json!({"codex": next_codex.clone()})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: provider_key_oauth_status_snapshot_update( + &latest_key, + ) + .status_snapshot_patch, + reset_error_count: false, + updated_at_unix_secs: latest_key.updated_at_unix_secs, + }, + ) + .await?; + if !updated { + if attempt + 1 < CODEX_OAUTH_INVALIDATION_CAS_MAX_ATTEMPTS { + tokio::task::yield_now().await; + } + continue; + } + + let auto_removed = if admin_provider_quota_pure::provider_auto_remove_banned_keys( + transport.provider.config.as_ref(), + ) && oauth_invalid_reason_is_account_block( + latest_key.oauth_invalid_reason.as_deref(), + ) + && admin_provider_quota_pure::should_auto_remove_oauth_invalid_key( + &latest_key, + None, + access_token_invalid_proven, + now_unix_secs, + ) { + self.delete_provider_transport_oauth_credential_fenced( + key_id, + &expected_credential_fence, + Some(ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(next_codex), + }), + ) + .await? + } else { + false + }; + let _ = self.invalidate_local_oauth_refresh_entry(key_id).await; + tracing::info!( + key_id, + provider_id = %transport.provider.id, + status_code, + request_started_at_unix_ms = observation.request_started_at_unix_ms, + request_order_id, + updated = true, + auto_removed, + "gateway observed local OAuth refresh failure state persisted" + ); + return Ok(auto_removed); + } + Ok(false) + } + pub(crate) async fn persist_local_oauth_refresh_failure_state( &self, transport: &provider_transport::GatewayProviderTransportSnapshot, @@ -2495,6 +2963,7 @@ impl AppState { expected_credential_fence.encrypted_auth_config.clone(), ), expected_credential: Some(expected_credential_fence.credential.clone()), + expected_upstream_metadata_namespace: None, encrypted_auth_config: expected_credential_fence .encrypted_auth_config .clone(), @@ -2503,6 +2972,7 @@ impl AppState { oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs, oauth_invalid_reason: latest_key.oauth_invalid_reason.clone(), upstream_metadata_patch: None, + upstream_metadata_namespace_to_remove: None, status_snapshot_patch: provider_key_oauth_status_snapshot_update( &latest_key, ) @@ -2554,6 +3024,7 @@ impl AppState { self.delete_provider_transport_oauth_credential_fenced( key_id, &expected_credential_fence, + None, ) .await? } else if !transport @@ -2927,8 +3398,9 @@ mod tests { use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository, - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyListQuery, + ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogReadRepository, + ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, }; @@ -2938,10 +3410,10 @@ mod tests { use tokio::sync::Notify; use super::{ - AgentIdentityAuthConfigFence, AppState, ProviderTransportSnapshotCacheKey, - ProviderTransportSnapshotFlight, ProviderTransportSnapshotFlightResult, - ProviderTransportSnapshotInflightRegistration, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, - PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL, + AgentIdentityAuthConfigFence, AppState, CodexRuntimeOAuthObservation, + ProviderTransportSnapshotCacheKey, ProviderTransportSnapshotFlight, + ProviderTransportSnapshotFlightResult, ProviderTransportSnapshotInflightRegistration, + PROVIDER_TRANSPORT_SNAPSHOT_CACHE_STALE_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_TTL, }; use crate::data::GatewayDataState; @@ -3104,6 +3576,40 @@ mod tests { (state, repository, encrypted_auth_config) } + async fn replace_key_credential_for_test( + repository: &InMemoryProviderCatalogReadRepository, + replacement: StoredProviderCatalogKey, + provider_type: &str, + ) { + let current = repository + .list_keys_by_ids(std::slice::from_ref(&replacement.id)) + .await + .expect("current key should load before replacement") + .into_iter() + .next() + .expect("current key should exist before replacement"); + let codex_rotation = provider_type.eq_ignore_ascii_case("codex").then(|| { + json!({ + "credential_generation": "test-admin-replacement-generation" + }) + }); + assert!(repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: current.encrypted_auth_config.clone(), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: current.encrypted_api_key.clone(), + auth_type: current.auth_type.clone(), + provider_id: current.provider_id.clone(), + provider_type: provider_type.to_string(), + }, + key: replacement, + codex_rotation, + reset_oauth_runtime: true, + }) + .await + .expect("credential replacement CAS should run")); + } + fn state_with_global_format_conversion(enabled: bool) -> AppState { let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( vec![sample_provider()], @@ -4006,10 +4512,7 @@ mod tests { encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "admin-replacement") .expect("replacement credential should encrypt"), ); - repository - .update_key(&replacement) - .await - .expect("replacement credential should persist"); + replace_key_credential_for_test(repository.as_ref(), replacement, "vertex_ai").await; assert!( state @@ -4093,10 +4596,7 @@ mod tests { ) .expect("replacement config should encrypt"), ); - repository - .update_key(&replaced) - .await - .expect("replacement should persist"); + replace_key_credential_for_test(repository.as_ref(), replaced, "codex").await; assert!(matches!( state @@ -4152,10 +4652,7 @@ mod tests { replaced.encrypted_auth_config = Some(replacement_auth_config.clone()); replaced.encrypted_api_key = Some(replacement_api_key.clone()); replaced.expires_at_unix_secs = None; - repository - .update_key(&replaced) - .await - .expect("replacement should persist"); + replace_key_credential_for_test(repository.as_ref(), replaced, "codex").await; let refreshed_entry = crate::provider_transport::CachedOAuthEntry { provider_type: "codex".to_string(), @@ -4220,10 +4717,7 @@ mod tests { .pop() .expect("key should exist"); replaced.encrypted_api_key = Some(replacement_api_key.clone()); - repository - .update_key(&replaced) - .await - .expect("access token replacement should persist"); + replace_key_credential_for_test(repository.as_ref(), replaced, "codex").await; assert!(!state .persist_local_oauth_refresh_failure_state( @@ -4273,15 +4767,15 @@ mod tests { .pop() .expect("key should exist"); replaced.encrypted_api_key = Some(replacement_api_key.clone()); - repository - .update_key(&replaced) - .await - .expect("access token replacement should persist"); + replace_key_credential_for_test(repository.as_ref(), replaced, "codex").await; assert!(!state .mark_provider_transport_oauth_invalid_fenced( &stale_transport, "[OAUTH_EXPIRED] stale request", + None, + None, + None, ) .await .expect("stale invalidation should be ignored")); @@ -4299,4 +4793,807 @@ mod tests { assert!(stored.oauth_invalid_at_unix_secs.is_none()); assert!(stored.oauth_invalid_reason.is_none()); } + + #[tokio::test] + async fn newer_runtime_oauth_success_clears_older_invalid_and_advances_watermark() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_invalid_fenced( + &transport, + "[OAUTH_EXPIRED] stale request", + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("older invalidation should persist")); + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_001), + Some("01900000-0000-7000-8000-000000000011"), + None, + ) + .await + .expect("newer success should persist")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(stored.oauth_invalid_at_unix_secs.is_none()); + assert!(stored.oauth_invalid_reason.is_none()); + let codex = stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .expect("Codex metadata should exist"); + assert_eq!( + codex["oauth_state_request_started_at_unix_ms"], + json!(1_001_u64) + ); + assert_eq!( + codex["oauth_state_request_id"], + json!("01900000-0000-7000-8000-000000000011") + ); + assert_eq!( + stored + .status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.get("oauth")) + .and_then(|oauth| oauth.get("requires_reauth")), + Some(&json!(false)) + ); + } + + #[tokio::test] + async fn runtime_oauth_success_allows_equal_quota_watermark() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + repository + .update_key_oauth_runtime_state( + "key-1", + Some(1), + Some("[OAUTH_EXPIRED] response raced with quota persistence"), + None, + Some(1), + ) + .await + .expect("invalid state should persist"); + repository + .upsert_key_upstream_metadata_namespace( + "key-1", + "codex", + &json!({ + "account_quota_request_started_at_unix_ms": 1_000_u64, + "account_quota_request_id": "01900000-0000-7000-8000-000000000010" + }), + Some(1), + ) + .await + .expect("quota watermark should persist"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("the same response should be idempotently accepted")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(stored.oauth_invalid_at_unix_secs.is_none()); + assert!(stored.oauth_invalid_reason.is_none()); + let codex = stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .expect("Codex metadata should exist"); + assert_eq!( + codex["oauth_state_request_id"], + json!("01900000-0000-7000-8000-000000000010") + ); + assert_eq!( + codex["account_quota_request_id"], + json!("01900000-0000-7000-8000-000000000010") + ); + } + + #[tokio::test] + async fn older_runtime_oauth_success_cannot_clear_newer_invalid() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_invalid_fenced( + &transport, + "[OAUTH_EXPIRED] newer request", + Some(1_001), + Some("01900000-0000-7000-8000-000000000011"), + None, + ) + .await + .expect("newer invalidation should persist")); + assert!(!state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("older success should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] newer request") + ); + } + + #[tokio::test] + async fn runtime_oauth_success_uses_request_id_for_same_millisecond_ordering() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_invalid_fenced( + &transport, + "[OAUTH_EXPIRED] same millisecond request", + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("invalidation should persist")); + assert!(!state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000009"), + None, + ) + .await + .expect("lower request id should be ignored")); + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000011"), + None, + ) + .await + .expect("higher request id should persist")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(stored.oauth_invalid_reason.is_none()); + } + + #[tokio::test] + async fn runtime_oauth_success_preserves_account_block() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + repository + .update_key_oauth_runtime_state( + "key-1", + Some(1), + Some("[ACCOUNT_BLOCK] account deactivated"), + None, + Some(1), + ) + .await + .expect("account block should persist"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("success watermark should persist")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("[ACCOUNT_BLOCK] account deactivated") + ); + assert_eq!(stored.oauth_invalid_at_unix_secs, Some(1)); + } + + #[tokio::test] + async fn runtime_oauth_success_rejects_credential_generation_mismatch() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + repository + .update_key_oauth_runtime_state( + "key-1", + Some(1), + Some("[OAUTH_EXPIRED] current generation invalid"), + None, + Some(1), + ) + .await + .expect("invalid state should persist"); + repository + .upsert_key_upstream_metadata_namespace( + "key-1", + "codex", + &json!({"credential_generation": "credential-v2"}), + Some(1), + ) + .await + .expect("credential generation should persist"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(!state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + Some("credential-v1"), + ) + .await + .expect("stale generation should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] current generation invalid") + ); + } + + #[tokio::test] + async fn runtime_oauth_success_rejects_replaced_transport_credential() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-old"); + repository + .update_key_oauth_runtime_state( + "key-1", + Some(1), + Some("[OAUTH_EXPIRED] replacement generation invalid"), + None, + Some(1), + ) + .await + .expect("invalid state should persist"); + let stale_transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + let replacement_api_key = + encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "access-admin") + .expect("replacement api key should encrypt"); + let mut replaced = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should load") + .pop() + .expect("key should exist"); + replaced.encrypted_api_key = Some(replacement_api_key); + replace_key_credential_for_test(repository.as_ref(), replaced, "codex").await; + repository + .update_key_oauth_runtime_state( + "key-1", + Some(1), + Some("[OAUTH_EXPIRED] replacement generation invalid"), + None, + Some(1), + ) + .await + .expect("replacement invalid state should persist"); + + assert!(!state + .mark_provider_transport_oauth_success_fenced( + &stale_transport, + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("stale transport should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] replacement generation invalid") + ); + } + + #[tokio::test] + async fn runtime_oauth_success_requires_complete_request_order() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + repository + .update_key_oauth_runtime_state( + "key-1", + Some(1), + Some("[OAUTH_EXPIRED] unordered response"), + None, + Some(1), + ) + .await + .expect("invalid state should persist"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(!state + .mark_provider_transport_oauth_success_fenced( + &transport, + None, + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("missing timestamp should be ignored")); + assert!(!state + .mark_provider_transport_oauth_success_fenced(&transport, Some(1_000), None, None,) + .await + .expect("missing request id should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert_eq!( + stored.oauth_invalid_reason.as_deref(), + Some("[OAUTH_EXPIRED] unordered response") + ); + } + + #[tokio::test] + async fn older_observed_refresh_failure_cannot_override_newer_success() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_001), + Some("01900000-0000-7000-8000-000000000011"), + None, + ) + .await + .expect("newer success should persist")); + assert!(!state + .persist_local_oauth_refresh_failure_state_observed( + &transport, + 401, + r#"{"error":"invalid_grant"}"#, + true, + CodexRuntimeOAuthObservation { + request_started_at_unix_ms: 1_000, + request_order_id: "01900000-0000-7000-8000-000000000010", + observed_credential_generation: None, + runtime_invalid_reason: None, + }, + ) + .await + .expect("older refresh failure should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(stored.oauth_invalid_reason.is_none()); + assert_eq!( + stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(|codex| codex.get("oauth_state_request_id")), + Some(&json!("01900000-0000-7000-8000-000000000011")) + ); + } + + #[tokio::test] + async fn newer_success_clears_observed_refresh_failure() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(!state + .persist_local_oauth_refresh_failure_state_observed( + &transport, + 401, + r#"{"error":"invalid_grant"}"#, + true, + CodexRuntimeOAuthObservation { + request_started_at_unix_ms: 1_000, + request_order_id: "01900000-0000-7000-8000-000000000010", + observed_credential_generation: None, + runtime_invalid_reason: None, + }, + ) + .await + .expect("refresh failure should persist without auto-removal")); + let failed = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(failed + .oauth_invalid_reason + .as_deref() + .is_some_and(|reason| reason.starts_with("[REFRESH_FAILED]"))); + + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_001), + Some("01900000-0000-7000-8000-000000000011"), + None, + ) + .await + .expect("newer success should persist")); + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(stored.oauth_invalid_at_unix_secs.is_none()); + assert!(stored.oauth_invalid_reason.is_none()); + } + + #[tokio::test] + async fn observed_refresh_failure_rejects_credential_generation_mismatch() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + repository + .upsert_key_upstream_metadata_namespace( + "key-1", + "codex", + &json!({"credential_generation": "credential-v2"}), + Some(1), + ) + .await + .expect("credential generation should persist"); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(!state + .persist_local_oauth_refresh_failure_state_observed( + &transport, + 401, + r#"{"error":"invalid_grant"}"#, + true, + CodexRuntimeOAuthObservation { + request_started_at_unix_ms: 1_000, + request_order_id: "01900000-0000-7000-8000-000000000010", + observed_credential_generation: Some("credential-v1"), + runtime_invalid_reason: None, + }, + ) + .await + .expect("stale generation should be ignored")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should remain"); + assert!(stored.oauth_invalid_reason.is_none()); + assert!(stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(|codex| codex.get("oauth_state_request_id")) + .is_none()); + } + + #[tokio::test] + async fn codex_runtime_invalid_is_not_auto_removed_before_newer_success() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let mut provider = repository + .list_providers_by_ids(&["provider-1".to_string()]) + .await + .expect("provider should load") + .pop() + .expect("provider should exist"); + provider.config = Some(json!({ + "pool_advanced": {"auto_remove_banned_keys": true} + })); + repository + .update_provider(&provider) + .await + .expect("provider config should persist"); + state.clear_provider_transport_snapshot_cache(); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_invalid_fenced( + &transport, + "oauth_token_invalid", + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("runtime invalid should persist")); + let invalid = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key lookup should succeed") + .into_iter() + .next() + .expect("recoverable runtime invalid must retain the key"); + assert_eq!( + invalid.oauth_invalid_reason.as_deref(), + Some("oauth_token_invalid") + ); + assert!(invalid.oauth_invalid_at_unix_secs.is_some()); + assert_eq!( + invalid + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(|codex| codex.get("oauth_state_request_id")), + Some(&json!("01900000-0000-7000-8000-000000000010")) + ); + assert!(state + .mark_provider_transport_oauth_success_fenced( + &transport, + Some(1_001), + Some("01900000-0000-7000-8000-000000000011"), + None, + ) + .await + .expect("newer success should clear invalid")); + } + + #[tokio::test] + async fn codex_observed_terminal_refresh_failure_is_persisted_without_auto_removal() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let mut provider = repository + .list_providers_by_ids(&["provider-1".to_string()]) + .await + .expect("provider should load") + .pop() + .expect("provider should exist"); + provider.config = Some(json!({ + "pool_advanced": {"auto_remove_banned_keys": true} + })); + repository + .update_provider(&provider) + .await + .expect("provider config should persist"); + state.clear_provider_transport_snapshot_cache(); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(!state + .persist_local_oauth_refresh_failure_state_observed( + &transport, + 401, + r#"{"error":{"code":"invalid_grant","message":"refresh token invalid"}}"#, + true, + CodexRuntimeOAuthObservation { + request_started_at_unix_ms: 1_000, + request_order_id: "01900000-0000-7000-8000-000000000010", + observed_credential_generation: None, + runtime_invalid_reason: Some("[OAUTH_EXPIRED] session expired"), + }, + ) + .await + .expect("terminal refresh failure should persist without auto-removal")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key lookup should succeed") + .into_iter() + .next() + .expect("request-scoped terminal refresh failure must retain the key"); + let reason = stored + .oauth_invalid_reason + .as_deref() + .expect("combined invalid reason should persist"); + assert!(reason.starts_with("[OAUTH_EXPIRED] session expired")); + assert!(reason.contains("[REFRESH_FAILED]")); + assert!(stored.oauth_invalid_at_unix_secs.is_some()); + assert_eq!( + stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(|codex| codex.get("oauth_state_request_id")), + Some(&json!("01900000-0000-7000-8000-000000000010")) + ); + } + + #[tokio::test] + async fn codex_runtime_account_block_is_still_auto_removed() { + let initial_config = json!({ + "provider_type": "codex", + "refresh_token": "refresh-stable", + "expires_at": 4102444800_u64 + }); + let (state, repository, _) = codex_oauth_state(&initial_config, "access-current"); + let mut provider = repository + .list_providers_by_ids(&["provider-1".to_string()]) + .await + .expect("provider should load") + .pop() + .expect("provider should exist"); + provider.config = Some(json!({ + "pool_advanced": {"auto_remove_banned_keys": true} + })); + repository + .update_provider(&provider) + .await + .expect("provider config should persist"); + state.clear_provider_transport_snapshot_cache(); + let transport = state + .read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1") + .await + .expect("transport should load") + .expect("transport should exist"); + + assert!(state + .mark_provider_transport_oauth_invalid_fenced( + &transport, + "[ACCOUNT_BLOCK] account deactivated", + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + None, + ) + .await + .expect("account block should persist and auto-remove")); + assert!(repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key lookup should succeed") + .is_empty()); + } } diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs index 651542973..44c706b45 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs @@ -774,6 +774,105 @@ async fn generic_key_routes_reject_agent_identity_credential_writes() { gateway_handle.abort(); } +#[tokio::test] +async fn generic_codex_key_credential_switch_rotates_generation_and_clears_quota() { + let mut existing_key = sample_key( + "key-codex-existing", + "provider-codex", + "openai:responses", + "old-oauth-access-token", + ); + existing_key.auth_type = "oauth".to_string(); + existing_key.encrypted_auth_config = Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + r#"{"provider_type":"codex","refresh_token":"old-refresh-token"}"#, + ) + .expect("old auth config should encrypt"), + ); + existing_key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "generation-before-switch", + "primary_used_percent": 80.0, + }, + "unrelated": {"preserved": true}, + })); + existing_key.status_snapshot = Some(json!({ + "oauth": {"status": "valid"}, + "quota": {"used_ratio": 0.8}, + })); + let mut provider = sample_provider("provider-codex", "codex", 10); + provider.provider_type = "codex".to_string(); + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![provider], + vec![], + vec![existing_key], + )); + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests( + provider_catalog_repository.clone(), + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .put(format!( + "{gateway_url}/api/admin/endpoints/keys/key-codex-existing" + )) + .header(crate::constants::GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") + .json(&json!({ + "auth_type": "api_key", + "api_key": "new-codex-api-key" + })) + .send() + .await + .expect("credential switch should complete"); + assert_eq!(response.status(), StatusCode::OK); + + let reloaded = provider_catalog_repository + .list_keys_by_ids(&["key-codex-existing".to_string()]) + .await + .expect("key should reload"); + assert_eq!(reloaded.len(), 1); + let key = &reloaded[0]; + assert_eq!(key.auth_type, "api_key"); + let codex = key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(serde_json::Value::as_object) + .expect("codex metadata should exist"); + assert_eq!(codex.len(), 1, "unexpected Codex metadata: {codex:?}"); + assert_ne!( + codex + .get(aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY) + .and_then(serde_json::Value::as_str), + Some("generation-before-switch") + ); + assert_eq!( + key.upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/unrelated/preserved")), + Some(&json!(true)) + ); + assert_eq!( + key.status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.get("quota")), + Some(&serde_json::Value::Null) + ); + + gateway_handle.abort(); +} + #[tokio::test] async fn provider_key_concurrent_limit_create_and_list_responses() { let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs index 7a3dd059b..0c6bc508b 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -1,7 +1,9 @@ use std::collections::BTreeMap; use std::sync::{Arc, Mutex}; -use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY}; +use aether_crypto::{ + decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY, +}; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository; use aether_data_contracts::repository::provider_catalog::{ @@ -121,6 +123,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a "1900500000".to_string(), ), ]), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "plan_type": "plus", @@ -295,6 +298,467 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a upstream_handle.abort(); } +#[test] +fn gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh() { + run_provider_quota_test( + "gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh", + gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh_impl, + ); +} + +async fn gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh_impl() { + let token_hits = Arc::new(Mutex::new(0usize)); + let token_server = Router::new().route( + "/oauth/token", + post({ + let token_hits = Arc::clone(&token_hits); + move || { + let token_hits = Arc::clone(&token_hits); + async move { + *token_hits.lock().expect("mutex should lock") += 1; + Json(json!({ + "access_token": "refreshed-codex-access-token", + "refresh_token": "rotated-codex-refresh-token", + "token_type": "Bearer", + "expires_in": 3_600, + "account_id": "acct-quota-refresh", + "plan_type": "plus" + })) + } + } + }), + ); + + let seen_requests = Arc::new(Mutex::new(Vec::<(String, String)>::new())); + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any({ + let seen_requests = Arc::clone(&seen_requests); + move |request: Request| { + let seen_requests = Arc::clone(&seen_requests); + async move { + let plan: aether_contracts::ExecutionPlan = serde_json::from_slice( + &to_bytes(request.into_body(), usize::MAX) + .await + .expect("body should read"), + ) + .expect("plan should parse"); + seen_requests.lock().expect("mutex should lock").push(( + plan.url.clone(), + plan.headers + .get("authorization") + .cloned() + .unwrap_or_default(), + )); + let body_json = match plan.url.as_str() { + "https://chatgpt.com/backend-api/wham/usage" => json!({ + "plan_type": "plus", + "rate_limit": { + "primary_window": { + "used_percent": 23.0, + "reset_at": 1_900_000_000u64, + "window_minutes": 300 + } + } + }), + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits" => { + json!({"available_count": 0, "credits": []}) + } + url => panic!("unexpected execution runtime URL: {url}"), + }; + let result = aether_contracts::ExecutionResult { + request_id: plan.request_id, + candidate_id: None, + status_code: 200, + headers: BTreeMap::new(), + response_observation: None, + body: Some(aether_contracts::ResponseBody { + json_body: Some(body_json), + body_bytes_b64: None, + }), + telemetry: None, + error: None, + }; + (StatusCode::OK, Json(result)) + } + } + }), + ); + + let mut key = sample_key( + "key-codex-expired-quota", + "provider-codex-expired-quota", + "openai:responses", + "expired-codex-access-token", + ); + key.auth_type = "oauth".to_string(); + key.expires_at_unix_secs = Some(1); + key.encrypted_auth_config = Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + &json!({ + "provider_type": "codex", + "refresh_token": "expired-codex-refresh-token", + "expires_at": 1, + "account_id": "acct-quota-refresh", + "plan_type": "plus" + }) + .to_string(), + ) + .expect("auth config should encrypt"), + ); + key.upstream_metadata = Some(json!({ + "codex": {"credential_generation": "credential-quota-refresh"} + })); + + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![StoredProviderCatalogProvider::new( + "provider-codex-expired-quota".to_string(), + "codex".to_string(), + Some("https://example.com".to_string()), + "codex".to_string(), + ) + .expect("provider should build")], + vec![sample_endpoint( + "endpoint-codex-expired-quota", + "provider-codex-expired-quota", + "openai:responses", + "https://chatgpt.com/backend-api", + )], + vec![key], + )); + + let (token_url, token_handle) = start_server(token_server).await; + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let oauth_refresh = + crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![ + Arc::new( + crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default() + .with_token_url_for_tests("codex", format!("{token_url}/oauth/token")), + ), + ]); + let gateway = build_router_with_state( + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests( + provider_catalog_repository.clone(), + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ) + .with_oauth_refresh_coordinator_for_tests(oauth_refresh), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!( + "{gateway_url}/api/admin/endpoints/providers/provider-codex-expired-quota/refresh-quota" + )) + .header(GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("json body should parse"); + assert_eq!(payload["success"], 1, "payload={payload}"); + assert_eq!(payload["failed"], 0, "payload={payload}"); + assert_eq!(payload["results"][0]["status"], "success"); + assert_eq!(*token_hits.lock().expect("mutex should lock"), 1); + assert_eq!( + seen_requests.lock().expect("mutex should lock").as_slice(), + [ + ( + "https://chatgpt.com/backend-api/wham/usage".to_string(), + "Bearer refreshed-codex-access-token".to_string(), + ), + ( + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits".to_string(), + "Bearer refreshed-codex-access-token".to_string(), + ), + ] + ); + + let reloaded = provider_catalog_repository + .list_keys_by_ids(&["key-codex-expired-quota".to_string()]) + .await + .expect("key should reload"); + let persisted = reloaded.first().expect("key should remain installed"); + let decrypted_api_key = decrypt_python_fernet_ciphertext( + DEVELOPMENT_ENCRYPTION_KEY, + persisted + .encrypted_api_key + .as_deref() + .expect("api key should persist"), + ) + .expect("api key should decrypt"); + assert_eq!(decrypted_api_key, "refreshed-codex-access-token"); + let decrypted_auth_config = decrypt_python_fernet_ciphertext( + DEVELOPMENT_ENCRYPTION_KEY, + persisted + .encrypted_auth_config + .as_deref() + .expect("auth config should persist"), + ) + .expect("auth config should decrypt"); + let auth_config: serde_json::Value = + serde_json::from_str(&decrypted_auth_config).expect("auth config should parse"); + assert_eq!(auth_config["refresh_token"], "rotated-codex-refresh-token"); + assert_eq!( + persisted + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/credential_generation")), + Some(&json!("credential-quota-refresh")) + ); + assert_eq!( + persisted + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/primary_used_percent")), + Some(&json!(23.0)) + ); + + gateway_handle.abort(); + execution_runtime_handle.abort(); + token_handle.abort(); +} + +#[test] +fn gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative() { + run_provider_quota_test( + "gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative", + gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative_impl, + ); +} + +async fn gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative_impl() { + const RESET_FENCE_UNIX_MS: u64 = 1_800_000_000_000; + const RESET_AT_UNIX_SECS: u64 = 2_000_000_000; + + let usage_hits = Arc::new(Mutex::new(0usize)); + let detail_hits = Arc::new(Mutex::new(0usize)); + let seen_urls = Arc::new(Mutex::new(Vec::::new())); + let execution_runtime = Router::new().route( + "/v1/execute/sync", + any({ + let usage_hits = Arc::clone(&usage_hits); + let detail_hits = Arc::clone(&detail_hits); + let seen_urls = Arc::clone(&seen_urls); + move |request: Request| { + let usage_hits = Arc::clone(&usage_hits); + let detail_hits = Arc::clone(&detail_hits); + let seen_urls = Arc::clone(&seen_urls); + async move { + let plan: aether_contracts::ExecutionPlan = serde_json::from_slice( + &to_bytes(request.into_body(), usize::MAX) + .await + .expect("body should read"), + ) + .expect("plan should parse"); + seen_urls + .lock() + .expect("mutex should lock") + .push(plan.url.clone()); + assert_eq!( + plan.headers.get("authorization").map(String::as_str), + Some("Bearer codex-reset-access-token") + ); + + let (body_json, response_observation) = match plan.url.as_str() { + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume" => { + assert_eq!(plan.method, "POST"); + assert_eq!( + plan.body.json_body, + Some(json!({"redeem_request_id": "reset-e2e"})) + ); + ( + json!({"outcome": "reset"}), + Some(aether_contracts::ExecutionResponseObservation { + request_started_at_unix_ms: RESET_FENCE_UNIX_MS - 100, + response_headers_observed_at_unix_ms: RESET_FENCE_UNIX_MS, + request_order_id: "consume-reset-e2e".to_string(), + }), + ) + } + "https://chatgpt.com/backend-api/wham/usage" => { + let hit = { + let mut hits = usage_hits.lock().expect("mutex should lock"); + *hits += 1; + *hits + }; + let used_percent = match hit { + 1 => 100.0, + 2 => 0.0, + _ => panic!("unexpected wham/usage request #{hit}"), + }; + ( + json!({ + "plan_type": "plus", + "rate_limit": { + "primary_window": { + "used_percent": used_percent, + "reset_at": RESET_AT_UNIX_SECS, + "window_minutes": 300 + } + } + }), + Some(aether_contracts::ExecutionResponseObservation { + request_started_at_unix_ms: RESET_FENCE_UNIX_MS + + (hit as u64 * 1_000), + response_headers_observed_at_unix_ms: RESET_FENCE_UNIX_MS + + (hit as u64 * 1_000) + + 100, + request_order_id: format!("usage-reset-e2e-{hit}"), + }), + ) + } + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits" => { + *detail_hits.lock().expect("mutex should lock") += 1; + (json!({"available_count": 0, "credits": []}), None) + } + url => panic!("unexpected execution runtime URL: {url}"), + }; + let result = aether_contracts::ExecutionResult { + request_id: plan.request_id, + candidate_id: None, + status_code: 200, + headers: BTreeMap::new(), + response_observation, + body: Some(aether_contracts::ResponseBody { + json_body: Some(body_json), + body_bytes_b64: None, + }), + telemetry: None, + error: None, + }; + (StatusCode::OK, Json(result)) + } + } + }), + ); + + let mut key = sample_key( + "key-codex-reset", + "provider-codex-reset", + "openai:responses", + "codex-reset-access-token", + ); + key.auth_type = "oauth".to_string(); + key.expires_at_unix_secs = Some(4_102_444_800); + key.encrypted_auth_config = Some( + encrypt_python_fernet_plaintext( + DEVELOPMENT_ENCRYPTION_KEY, + &json!({ + "provider_type": "codex", + "refresh_token": "codex-reset-refresh-token", + "expires_at": 4_102_444_800u64, + "account_id": "acct-reset-e2e", + "plan_type": "plus" + }) + .to_string(), + ) + .expect("auth config should encrypt"), + ); + key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "credential-reset-e2e", + "plan_type": "plus", + "primary_used_percent": 100.0, + "primary_reset_at": RESET_AT_UNIX_SECS, + "primary_window_minutes": 300, + "updated_at": (RESET_FENCE_UNIX_MS / 1_000) - 1, + "account_quota_request_started_at_unix_ms": RESET_FENCE_UNIX_MS - 1_000, + "account_quota_request_id": "usage-before-reset" + } + })); + + let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed( + vec![StoredProviderCatalogProvider::new( + "provider-codex-reset".to_string(), + "codex".to_string(), + Some("https://example.com".to_string()), + "codex".to_string(), + ) + .expect("provider should build")], + vec![sample_endpoint( + "endpoint-codex-reset", + "provider-codex-reset", + "openai:responses", + "https://chatgpt.com/backend-api", + )], + vec![key], + )); + + let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await; + let gateway = build_router_with_state( + build_state_with_execution_runtime_override(execution_runtime_url) + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests( + provider_catalog_repository.clone(), + ) + .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; + + let response = reqwest::Client::new() + .post(format!( + "{gateway_url}/api/admin/endpoints/keys/key-codex-reset/codex-reset-credit/consume" + )) + .header(GATEWAY_HEADER, "rust-phase3b") + .header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123") + .header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin") + .header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123") + .json(&json!({ + "idempotency_key": "reset-e2e", + "expected_credential_generation": "credential-reset-e2e" + })) + .send() + .await + .expect("request should succeed"); + + assert_eq!(response.status(), StatusCode::OK); + let payload: serde_json::Value = response.json().await.expect("json body should parse"); + assert_eq!(payload["status"], "success", "payload={payload}"); + assert_eq!(payload["outcome"], "reset"); + assert_eq!(payload["refresh_status"], "success"); + assert_eq!(*usage_hits.lock().expect("mutex should lock"), 2); + assert_eq!(*detail_hits.lock().expect("mutex should lock"), 2); + assert_eq!( + seen_urls.lock().expect("mutex should lock").as_slice(), + [ + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume", + "https://chatgpt.com/backend-api/wham/usage", + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits", + "https://chatgpt.com/backend-api/wham/usage", + "https://chatgpt.com/backend-api/wham/rate-limit-reset-credits", + ] + ); + + let reloaded = provider_catalog_repository + .list_keys_by_ids(&["key-codex-reset".to_string()]) + .await + .expect("key should reload"); + let codex = reloaded[0] + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .expect("codex metadata should persist"); + assert_eq!(codex["primary_used_percent"], json!(0.0)); + assert_eq!(codex["primary_reset_at"], json!(RESET_AT_UNIX_SECS)); + assert_eq!(codex["account_quota_reset_pending"], json!(false)); + assert_eq!( + codex["account_quota_reset_processed_ids"], + json!(["reset-e2e"]) + ); + + gateway_handle.abort(); + execution_runtime_handle.abort(); +} + #[tokio::test] async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_required() { let upstream = Router::new().route( @@ -318,6 +782,7 @@ async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_req candidate_id: None, status_code: 402, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "error": { @@ -439,6 +904,7 @@ async fn gateway_auto_removes_codex_key_when_quota_proves_oauth_invalid() { candidate_id: None, status_code: 401, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "error": { @@ -594,6 +1060,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_requested_codex_keys "1900000000".to_string(), ), ]), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "plan_type": "plus", @@ -740,6 +1207,7 @@ async fn gateway_refreshes_admin_provider_quota_for_codex_proxy_with_extended_ti candidate_id: None, status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "plan_type": "plus", @@ -897,6 +1365,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_ad candidate_id: None, status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "subscriptionInfo": { @@ -1280,6 +1749,7 @@ async fn gateway_refresh_kiro_quota_reconciles_missing_fixed_endpoint_before_ref candidate_id: None, status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "subscriptionInfo": { @@ -1443,6 +1913,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_gemini_cli_with_trus candidate_id: None, status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "buckets": [ @@ -1840,6 +2311,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru candidate_id: None, status_code: 200, headers: BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "models": { diff --git a/apps/aether-gateway/src/tests/control/admin/oauth.rs b/apps/aether-gateway/src/tests/control/admin/oauth.rs index ba6676988..804d868eb 100644 --- a/apps/aether-gateway/src/tests/control/admin/oauth.rs +++ b/apps/aether-gateway/src/tests/control/admin/oauth.rs @@ -3440,11 +3440,16 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p let token_hits_clone = Arc::clone(&token_hits); let seen_token = Arc::new(Mutex::new(None::)); let seen_token_clone = Arc::clone(&seen_token); + let namespace_race_repository = Arc::new(Mutex::new( + None::>, + )); + let namespace_race_repository_clone = Arc::clone(&namespace_race_repository); let token_server = Router::new().route( "/oauth/token", any(move |request: Request| { let token_hits_inner = Arc::clone(&token_hits_clone); let seen_token_inner = Arc::clone(&seen_token_clone); + let namespace_race_repository_inner = Arc::clone(&namespace_race_repository_clone); async move { *token_hits_inner.lock().expect("mutex should lock") += 1; let (parts, body) = request.into_parts(); @@ -3459,6 +3464,23 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p body: String::from_utf8(raw_body.to_vec()) .expect("token request body should be utf8"), }); + let repository = namespace_race_repository_inner + .lock() + .expect("mutex should lock") + .clone() + .expect("provider catalog repository should be installed"); + repository + .upsert_key_upstream_metadata_namespace( + "key-codex-oauth", + "codex", + &json!({ + "primary_used_percent": 80.0, + "account_quota_request_id": "concurrent-refresh" + }), + Some(1_800_000_001), + ) + .await + .expect("concurrent Codex namespace update should succeed"); Json(json!({ "access_token": "new-codex-access-token", "refresh_token": "new-codex-refresh-token", @@ -3492,6 +3514,33 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p key.circuit_breaker_by_format = Some(json!({ "openai:chat": {"state": "open"} })); + key.upstream_metadata = Some(json!({ + "codex": { + "primary_used_percent": 100.0, + "primary_reset_at": 4_200_000_000u64, + "account_quota_reset_fence_unix_ms": 1_800_000_000_000u64, + "account_quota_reset_fence_id": "old-account-fence", + "account_quota_reset_processed_ids": ["old-account-reset"], + "account_quota_reset_pending": true + }, + "unrelated_runtime": { + "preserved": true + } + })); + key.status_snapshot = Some(json!({ + "oauth": { + "code": "invalid" + }, + "quota": { + "provider_type": "codex", + "usage_ratio": 1.0, + "windows": [{ + "kind": "primary", + "usage": 100.0, + "reset_at": 4_200_000_000u64 + }] + } + })); let score_identity = PoolMemberIdentity::provider_api_key("provider-codex", "key-codex-oauth"); let score_scope = provider_key_pool_score_scope(); @@ -3510,6 +3559,8 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p vec![], vec![key], )); + *namespace_race_repository.lock().expect("mutex should lock") = + Some(Arc::clone(&provider_catalog_repository)); let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![invalid_score])); @@ -3593,6 +3644,37 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p assert_eq!(persisted.error_count, Some(0)); assert_eq!(persisted.health_by_format, Some(json!({}))); assert_eq!(persisted.circuit_breaker_by_format, Some(json!({}))); + assert_eq!( + persisted + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("codex")) + .and_then(Value::as_object) + .map(|codex| codex.keys().cloned().collect::>()), + Some(vec!["credential_generation".to_string()]), + "explicit Codex reauthorization must not carry quota/reset state across accounts" + ); + assert!(persisted + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/codex/credential_generation")) + .and_then(Value::as_str) + .is_some_and(|generation| !generation.is_empty())); + assert_eq!( + persisted + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("unrelated_runtime")), + Some(&json!({"preserved": true})) + ); + assert_eq!( + persisted + .status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.get("quota")), + Some(&Value::Null) + ); let scores = pool_score_repository .get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery { ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)], @@ -6164,6 +6246,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p candidate_id: None, status_code: 401, headers: std::collections::BTreeMap::new(), + response_observation: None, body: None, telemetry: None, error: None, @@ -6523,6 +6606,7 @@ async fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_im candidate_id: None, status_code: 200, headers: std::collections::BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "plan_type": "plus", @@ -6757,6 +6841,7 @@ async fn run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test( candidate_id: None, status_code: 200, headers: std::collections::BTreeMap::new(), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "subscriptionInfo": { diff --git a/apps/aether-gateway/src/tests/control/admin/system_import.rs b/apps/aether-gateway/src/tests/control/admin/system_import.rs index 174712420..1a436ce44 100644 --- a/apps/aether-gateway/src/tests/control/admin/system_import.rs +++ b/apps/aether-gateway/src/tests/control/admin/system_import.rs @@ -2317,6 +2317,17 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp ) .expect("auth config should encrypt"), ); + existing_key.upstream_metadata = Some(json!({ + "codex": { + "credential_generation": "generation-before-import", + "primary_used_percent": 90.0, + }, + "unrelated": {"preserved": true}, + })); + existing_key.status_snapshot = Some(json!({ + "oauth": {"status": "invalid"}, + "quota": {"used_ratio": 0.9}, + })); let score_identity = PoolMemberIdentity::provider_api_key("provider-codex-existing", "key-codex-existing"); @@ -2417,6 +2428,31 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp assert_eq!(key.error_count, Some(0)); assert_eq!(key.health_by_format, Some(json!({}))); assert_eq!(key.circuit_breaker_by_format, Some(json!({}))); + let codex = key + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")) + .and_then(Value::as_object) + .expect("Codex metadata should exist"); + assert_eq!(codex.len(), 1); + assert_ne!( + codex + .get(aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY) + .and_then(Value::as_str), + Some("generation-before-import") + ); + assert_eq!( + key.upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/unrelated/preserved")), + Some(&json!(true)) + ); + assert_eq!( + key.status_snapshot + .as_ref() + .and_then(|snapshot| snapshot.get("quota")), + Some(&Value::Null) + ); assert_eq!( decrypt_python_fernet_ciphertext( DEVELOPMENT_ENCRYPTION_KEY, diff --git a/apps/aether-gateway/src/tests/control/proxy/embeddings.rs b/apps/aether-gateway/src/tests/control/proxy/embeddings.rs index 5ba243a6a..dbe2e8edd 100644 --- a/apps/aether-gateway/src/tests/control/proxy/embeddings.rs +++ b/apps/aether-gateway/src/tests/control/proxy/embeddings.rs @@ -657,6 +657,7 @@ fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult { candidate_id: plan.candidate_id.clone(), status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "object": "list", @@ -679,6 +680,7 @@ fn gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult { candidate_id: plan.candidate_id.clone(), status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "model": "gemini-embedding-2-preview", @@ -703,6 +705,7 @@ fn vertex_gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionRe candidate_id: plan.candidate_id.clone(), status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "predictions": [ @@ -727,6 +730,7 @@ fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionRes candidate_id: plan.candidate_id.clone(), status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "model": "gemini-embedding-2-preview", @@ -752,6 +756,7 @@ fn aliyun_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult { candidate_id: plan.candidate_id.clone(), status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "output": { diff --git a/apps/aether-gateway/src/tests/control/proxy/rerank.rs b/apps/aether-gateway/src/tests/control/proxy/rerank.rs index f4d3c341f..81c0fadf9 100644 --- a/apps/aether-gateway/src/tests/control/proxy/rerank.rs +++ b/apps/aether-gateway/src/tests/control/proxy/rerank.rs @@ -151,6 +151,7 @@ fn rerank_execution_result(plan: &ExecutionPlan) -> ExecutionResult { candidate_id: plan.candidate_id.clone(), status_code: 200, headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "model": "upstream-rerank", diff --git a/apps/aether-gateway/src/tests/files/stream.rs b/apps/aether-gateway/src/tests/files/stream.rs index 527998ed9..895f304dd 100644 --- a/apps/aether-gateway/src/tests/files/stream.rs +++ b/apps/aether-gateway/src/tests/files/stream.rs @@ -148,6 +148,7 @@ async fn gateway_executes_gemini_files_download_via_local_decision_gate_with_loc "content-type".to_string(), "application/octet-stream".to_string(), )]), + response_observation: None, }, }, StreamFrame { diff --git a/apps/aether-gateway/src/tests/usage/local.rs b/apps/aether-gateway/src/tests/usage/local.rs index bdb4c12ba..d59d43a95 100644 --- a/apps/aether-gateway/src/tests/usage/local.rs +++ b/apps/aether-gateway/src/tests/usage/local.rs @@ -774,6 +774,7 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha "content-type".to_string(), "application/json".to_string(), )]), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "error": { @@ -1056,6 +1057,7 @@ async fn sync_transport_error_policy_stops_or_retries_candidates_end_to_end_impl "content-type".to_string(), "application/json".to_string(), )]), + response_observation: None, body: Some(aether_contracts::ResponseBody { json_body: Some(json!({ "id": "chatcmpl-transport-policy", diff --git a/apps/aether-gateway/src/tests/video/stream.rs b/apps/aether-gateway/src/tests/video/stream.rs index 56eec5c91..f41b009fd 100644 --- a/apps/aether-gateway/src/tests/video/stream.rs +++ b/apps/aether-gateway/src/tests/video/stream.rs @@ -212,6 +212,7 @@ async fn gateway_executes_openai_video_content_from_reconstructed_data_task_with "content-type".to_string(), "video/mp4".to_string(), )]), + response_observation: None, }, }, StreamFrame { diff --git a/crates/aether-admin/src/provider/quota.rs b/crates/aether-admin/src/provider/quota.rs index b4f755471..cbfc261e5 100644 --- a/crates/aether-admin/src/provider/quota.rs +++ b/crates/aether-admin/src/provider/quota.rs @@ -616,6 +616,1191 @@ pub fn build_codex_quota_exhausted_fallback_metadata( serde_json::Value::Object(object) } +const CODEX_QUOTA_WINDOW_SUFFIXES: &[&str] = &[ + "used_percent", + "reset_seconds", + "reset_after_seconds", + "reset_at", + "next_reset_at", + "window_minutes", +]; +const CODEX_QUOTA_RESET_DEADLINE_TOLERANCE_SECONDS: u64 = 30; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CodexQuotaWindowCoverage { + Patch, + AccountSnapshot, + FullSnapshot, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct CodexQuotaMergeContext<'a> { + pub observed_at_unix_secs: u64, + pub request_started_at_unix_ms: Option, + pub request_order_id: Option<&'a str>, + /// Reset generation captured before the upstream request was sent. + /// Once a key has entered generation 1, an absent or different value is + /// treated as a pre-reset observation for account quota and metadata. + pub observed_reset_generation: Option, + /// Generation owned by an explicit reset reconciliation request. Normal + /// quota observations leave this unset, but a complete account snapshot + /// started after the reset fence may still reconcile a delayed reset. + pub authoritative_reset_generation: Option, + /// Non-secret credential generation captured with the transport snapshot. + pub observed_credential_generation: Option<&'a str>, + /// Identifies the reset-credit fence that authorized this observation. + /// Retained for generation-0 rolling-upgrade compatibility. + pub account_reset_fence_id: Option<&'a str>, + pub coverage: CodexQuotaWindowCoverage, +} + +impl<'a> CodexQuotaMergeContext<'a> { + fn request_order(self) -> Option> { + codex_quota_request_order(self.request_started_at_unix_ms, self.request_order_id) + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct CodexQuotaMergeOutcome { + pub metadata: serde_json::Value, + pub changed: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CodexQuotaWindowFamily { + Account, + Spark, +} + +impl CodexQuotaWindowFamily { + fn watermark_key(self) -> &'static str { + match self { + Self::Account => CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_KEY, + Self::Spark => CODEX_QUOTA_SPARK_REQUEST_WATERMARK_KEY, + } + } + + fn watermark_id_key(self) -> &'static str { + match self { + Self::Account => CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_ID_KEY, + Self::Spark => CODEX_QUOTA_SPARK_REQUEST_WATERMARK_ID_KEY, + } + } +} + +pub const CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_KEY: &str = + "account_quota_request_started_at_unix_ms"; +pub const CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_ID_KEY: &str = "account_quota_request_id"; +pub const CODEX_QUOTA_SPARK_REQUEST_WATERMARK_KEY: &str = "spark_quota_request_started_at_unix_ms"; +pub const CODEX_QUOTA_SPARK_REQUEST_WATERMARK_ID_KEY: &str = "spark_quota_request_id"; +pub const CODEX_QUOTA_METADATA_REQUEST_WATERMARK_KEY: &str = + "quota_metadata_request_started_at_unix_ms"; +pub const CODEX_QUOTA_METADATA_REQUEST_WATERMARK_ID_KEY: &str = "quota_metadata_request_id"; +pub const CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY: &str = "oauth_state_request_started_at_unix_ms"; +pub const CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY: &str = "oauth_state_request_id"; +pub const CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY: &str = "account_quota_reset_fence_unix_ms"; +pub const CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY: &str = "account_quota_reset_fence_id"; +pub const CODEX_QUOTA_ACCOUNT_RESET_PROCESSED_IDS_KEY: &str = "account_quota_reset_processed_ids"; +pub const CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY: &str = "account_quota_reset_pending"; +pub const CODEX_QUOTA_ACCOUNT_RESET_SEQUENCE_KEY: &str = "account_quota_reset_sequence"; +pub const CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY: &str = "account_quota_reset_generation"; +pub const CODEX_QUOTA_ACCOUNT_RESET_PENDING_GENERATION_KEY: &str = + "account_quota_reset_pending_generation"; +pub const CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY: &str = "account_quota_reset_reservation"; +pub const CODEX_QUOTA_ACCOUNT_RESET_HISTORY_KEY: &str = "account_quota_reset_history"; +pub const CODEX_CREDENTIAL_GENERATION_KEY: &str = "credential_generation"; + +pub fn codex_quota_account_reset_generation(codex: Option<&serde_json::Value>) -> u64 { + codex + .and_then(serde_json::Value::as_object) + .and_then(|codex| codex.get(CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY)) + .and_then(coerce_json_u64) + .unwrap_or(0) +} + +pub fn codex_credential_generation(codex: Option<&serde_json::Value>) -> Option<&str> { + codex + .and_then(serde_json::Value::as_object) + .and_then(|codex| codex.get(CODEX_CREDENTIAL_GENERATION_KEY)) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) +} + +pub fn codex_credential_generation_matches( + codex: Option<&serde_json::Value>, + observed: Option<&str>, +) -> bool { + let observed = observed.map(str::trim).filter(|value| !value.is_empty()); + codex_credential_generation(codex) == observed +} + +fn codex_quota_observation_matches_reset_generation( + object: &serde_json::Map, + context: CodexQuotaMergeContext<'_>, +) -> bool { + let active_generation = object + .get(CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(coerce_json_u64) + .unwrap_or(0); + if active_generation == 0 { + context.observed_reset_generation.unwrap_or(0) == 0 + } else { + context.observed_reset_generation == Some(active_generation) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +struct CodexQuotaRequestOrder<'a> { + started_at_unix_ms: u64, + request_id: Option<&'a str>, +} + +fn codex_quota_request_order<'a>( + started_at_unix_ms: Option, + request_id: Option<&'a str>, +) -> Option> { + started_at_unix_ms.map(|started_at_unix_ms| CodexQuotaRequestOrder { + started_at_unix_ms, + request_id: request_id + .map(str::trim) + .filter(|request_id| !request_id.is_empty()), + }) +} + +fn codex_quota_read_request_order<'a>( + object: &'a serde_json::Map, + watermark_key: &str, + watermark_id_key: &str, +) -> Option> { + object + .get(watermark_key) + .and_then(coerce_json_u64) + .map(|started_at_unix_ms| CodexQuotaRequestOrder { + started_at_unix_ms, + request_id: object + .get(watermark_id_key) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|request_id| !request_id.is_empty()), + }) +} + +fn codex_quota_request_order_is_stale( + incoming: Option>, + current: Option>, +) -> bool { + match (incoming, current) { + (Some(incoming), Some(current)) => incoming <= current, + (None, Some(_)) => true, + _ => false, + } +} + +/// Returns whether an incoming Codex observation is older than or identical +/// to the stored request order. Request ids break ties within one millisecond; +/// a legacy watermark without an id sorts before one that has an id. +pub fn codex_request_order_is_stale( + incoming_started_at_unix_ms: Option, + incoming_request_id: Option<&str>, + stored_started_at_unix_ms: Option, + stored_request_id: Option<&str>, +) -> bool { + codex_quota_request_order_is_stale( + codex_quota_request_order(incoming_started_at_unix_ms, incoming_request_id), + codex_quota_request_order(stored_started_at_unix_ms, stored_request_id), + ) +} + +/// Compares an OAuth-state observation against every persisted Codex response +/// watermark. Quota-only responses also prove request ordering, so an older +/// runtime authentication failure cannot override them. +pub fn codex_oauth_state_request_order_is_stale( + codex: Option<&serde_json::Map>, + incoming_started_at_unix_ms: Option, + incoming_request_id: Option<&str>, +) -> bool { + let stored = codex.and_then(|codex| { + [ + ( + CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY, + CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY, + ), + ( + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_ID_KEY, + ), + ( + CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_ID_KEY, + ), + ( + CODEX_QUOTA_SPARK_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_SPARK_REQUEST_WATERMARK_ID_KEY, + ), + ] + .into_iter() + .filter_map(|(watermark_key, watermark_id_key)| { + codex_quota_read_request_order(codex, watermark_key, watermark_id_key) + }) + .max() + }); + codex_quota_request_order_is_stale( + codex_quota_request_order(incoming_started_at_unix_ms, incoming_request_id), + stored, + ) +} + +/// Compares a successful OAuth observation against every persisted Codex +/// response watermark. Equality is allowed because quota persistence and the +/// success effect may independently process the same upstream response. +pub fn codex_oauth_success_request_order_is_stale( + codex: Option<&serde_json::Map>, + incoming_started_at_unix_ms: Option, + incoming_request_id: Option<&str>, +) -> bool { + let stored = codex.and_then(|codex| { + [ + ( + CODEX_OAUTH_STATE_REQUEST_WATERMARK_KEY, + CODEX_OAUTH_STATE_REQUEST_WATERMARK_ID_KEY, + ), + ( + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_ID_KEY, + ), + ( + CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_ACCOUNT_REQUEST_WATERMARK_ID_KEY, + ), + ( + CODEX_QUOTA_SPARK_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_SPARK_REQUEST_WATERMARK_ID_KEY, + ), + ] + .into_iter() + .filter_map(|(watermark_key, watermark_id_key)| { + codex_quota_read_request_order(codex, watermark_key, watermark_id_key) + }) + .max() + }); + match ( + codex_quota_request_order(incoming_started_at_unix_ms, incoming_request_id), + stored, + ) { + (Some(incoming), Some(stored)) => incoming < stored, + (None, Some(_)) => true, + _ => false, + } +} + +fn codex_quota_request_order_is_newer( + incoming: CodexQuotaRequestOrder<'_>, + current: Option>, +) -> bool { + current.is_none_or(|current| incoming > current) +} + +fn codex_quota_write_request_order( + object: &mut serde_json::Map, + watermark_key: &str, + watermark_id_key: &str, + order: CodexQuotaRequestOrder<'_>, +) { + object.insert(watermark_key.to_string(), json!(order.started_at_unix_ms)); + if let Some(request_id) = order.request_id { + object.insert(watermark_id_key.to_string(), json!(request_id)); + } else { + object.remove(watermark_id_key); + } +} + +fn codex_quota_is_request_order_key(key: &str) -> bool { + key == CODEX_QUOTA_METADATA_REQUEST_WATERMARK_KEY + || key == CODEX_QUOTA_METADATA_REQUEST_WATERMARK_ID_KEY + || [ + CodexQuotaWindowFamily::Account, + CodexQuotaWindowFamily::Spark, + ] + .into_iter() + .any(|family| key == family.watermark_key() || key == family.watermark_id_key()) +} + +fn codex_quota_is_reset_fence_key(key: &str) -> bool { + matches!( + key, + CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY + | CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY + | CODEX_QUOTA_ACCOUNT_RESET_PROCESSED_IDS_KEY + | CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY + | CODEX_QUOTA_ACCOUNT_RESET_SEQUENCE_KEY + | CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY + | CODEX_QUOTA_ACCOUNT_RESET_PENDING_GENERATION_KEY + | CODEX_QUOTA_ACCOUNT_RESET_RESERVATION_KEY + | CODEX_QUOTA_ACCOUNT_RESET_HISTORY_KEY + | CODEX_CREDENTIAL_GENERATION_KEY + ) +} + +#[derive(Debug, Clone, Copy)] +struct CodexQuotaAccountResetFence<'a> { + unix_ms: u64, + id: &'a str, + pending: bool, +} + +fn codex_quota_account_reset_fence( + object: &serde_json::Map, +) -> Option> { + let unix_ms = object + .get(CODEX_QUOTA_ACCOUNT_RESET_FENCE_UNIX_MS_KEY) + .and_then(coerce_json_u64) + .filter(|value| *value > 0)?; + let id = object + .get(CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty())?; + let pending = object + .get(CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + Some(CodexQuotaAccountResetFence { + unix_ms, + id, + pending, + }) +} + +fn codex_quota_reset_fence_authorizes( + fence: CodexQuotaAccountResetFence<'_>, + context: CodexQuotaMergeContext<'_>, +) -> bool { + context + .account_reset_fence_id + .map(str::trim) + .filter(|value| !value.is_empty()) + == Some(fence.id) +} + +fn codex_quota_request_started_after_reset_fence( + fence: CodexQuotaAccountResetFence<'_>, + context: CodexQuotaMergeContext<'_>, +) -> bool { + context + .request_started_at_unix_ms + .is_some_and(|started_at| started_at > fence.unix_ms) +} + +fn codex_quota_reset_fence_blocks( + fence: CodexQuotaAccountResetFence<'_>, + context: CodexQuotaMergeContext<'_>, +) -> bool { + !codex_quota_reset_fence_authorizes(fence, context) + && !codex_quota_request_started_after_reset_fence(fence, context) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum CodexQuotaWindowSlot { + Primary, + Secondary, + SparkPrimary, + SparkSecondary, +} + +impl CodexQuotaWindowSlot { + const ALL: [Self; 4] = [ + Self::Primary, + Self::Secondary, + Self::SparkPrimary, + Self::SparkSecondary, + ]; + + fn prefix(self) -> &'static str { + match self { + Self::Primary => "primary", + Self::Secondary => "secondary", + Self::SparkPrimary => "spark_primary", + Self::SparkSecondary => "spark_secondary", + } + } + + fn family(self) -> CodexQuotaWindowFamily { + match self { + Self::Primary | Self::Secondary => CodexQuotaWindowFamily::Account, + Self::SparkPrimary | Self::SparkSecondary => CodexQuotaWindowFamily::Spark, + } + } +} + +#[derive(Debug, Clone)] +struct CodexQuotaWindowObservation { + slot: CodexQuotaWindowSlot, + fields: serde_json::Map, + window_minutes: Option, + deadline: Option, + disabled: bool, +} + +impl CodexQuotaWindowObservation { + fn active(&self) -> bool { + !self.disabled && !self.fields.is_empty() + } + + fn used_percent(&self) -> Option { + self.fields + .get("used_percent") + .and_then(coerce_json_f64) + .filter(|value| value.is_finite()) + } + + fn persist_deadline(&mut self) { + if let Some(deadline) = self.deadline { + self.fields.remove("next_reset_at"); + self.fields.insert("reset_at".to_string(), json!(deadline)); + } + } +} + +fn codex_quota_window_key(slot: CodexQuotaWindowSlot, suffix: &str) -> String { + format!("{}_{suffix}", slot.prefix()) +} + +fn codex_quota_is_window_key(key: &str) -> bool { + CodexQuotaWindowSlot::ALL.iter().any(|slot| { + CODEX_QUOTA_WINDOW_SUFFIXES + .iter() + .any(|suffix| key == codex_quota_window_key(*slot, suffix)) + }) +} + +fn codex_quota_read_window( + object: &serde_json::Map, + slot: CodexQuotaWindowSlot, + observed_at_unix_secs: Option, +) -> Option { + let mut fields = serde_json::Map::new(); + for suffix in CODEX_QUOTA_WINDOW_SUFFIXES { + let key = codex_quota_window_key(slot, suffix); + let Some(raw) = object.get(&key) else { + continue; + }; + let normalized = match *suffix { + "used_percent" => coerce_json_f64(raw) + .filter(|value| value.is_finite()) + .map(|value| json!(value)), + _ => coerce_json_u64(raw).map(|value| json!(value)), + }; + if let Some(value) = normalized { + fields.insert((*suffix).to_string(), value); + } + } + if fields.is_empty() { + return None; + } + + let window_minutes = fields.get("window_minutes").and_then(coerce_json_u64); + let disabled = window_minutes == Some(0); + let explicit_deadline = fields + .get("reset_at") + .and_then(coerce_json_u64) + .filter(|value| *value > 0) + .or_else(|| { + fields + .get("next_reset_at") + .and_then(coerce_json_u64) + .filter(|value| *value > 0) + }); + let reset_after_seconds = fields + .get("reset_after_seconds") + .and_then(coerce_json_u64) + .or_else(|| fields.get("reset_seconds").and_then(coerce_json_u64)); + let deadline = explicit_deadline.or_else(|| { + observed_at_unix_secs + .zip(reset_after_seconds) + .map(|(observed_at, reset_after)| observed_at.saturating_add(reset_after)) + }); + + Some(CodexQuotaWindowObservation { + slot, + fields, + window_minutes: window_minutes.filter(|value| *value > 0), + deadline, + disabled, + }) +} + +fn codex_quota_read_family_windows( + object: &serde_json::Map, + family: CodexQuotaWindowFamily, + observed_at_unix_secs: Option, +) -> Vec { + CodexQuotaWindowSlot::ALL + .iter() + .copied() + .filter(|slot| slot.family() == family) + .filter_map(|slot| codex_quota_read_window(object, slot, observed_at_unix_secs)) + .collect() +} + +pub fn codex_quota_metadata_has_account_windows(value: &serde_json::Value) -> bool { + value.as_object().is_some_and(|object| { + !codex_quota_read_family_windows(object, CodexQuotaWindowFamily::Account, None).is_empty() + }) +} + +pub fn codex_quota_metadata_has_spark_windows(value: &serde_json::Value) -> bool { + value.as_object().is_some_and(|object| { + !codex_quota_read_family_windows(object, CodexQuotaWindowFamily::Spark, None).is_empty() + }) +} + +fn codex_quota_same_window_identity( + current: &CodexQuotaWindowObservation, + incoming: &CodexQuotaWindowObservation, +) -> bool { + match (current.window_minutes, incoming.window_minutes) { + (Some(current), Some(incoming)) => current == incoming, + // Old metadata did not always carry a duration. Only fall back to its + // storage slot when at least one side has that legacy shape. + _ => current.slot == incoming.slot, + } +} + +fn codex_quota_merge_same_window( + current: &CodexQuotaWindowObservation, + incoming: &CodexQuotaWindowObservation, +) -> CodexQuotaWindowObservation { + // A deadline must be present on both sides to prove a natural generation + // change. Legacy metadata without one stays monotonic until a later pair of + // observations establishes the deadline. + if current + .deadline + .zip(incoming.deadline) + .is_some_and(|(current, incoming)| { + incoming.saturating_add(CODEX_QUOTA_RESET_DEADLINE_TOLERANCE_SECONDS) < current + }) + { + return current.clone(); + } + + if current + .deadline + .zip(incoming.deadline) + .is_some_and(|(current, incoming)| { + incoming > current.saturating_add(CODEX_QUOTA_RESET_DEADLINE_TOLERANCE_SECONDS) + }) + { + let mut next = incoming.clone(); + next.persist_deadline(); + return next; + } + + let mut merged = current.clone(); + for (suffix, value) in &incoming.fields { + if suffix != "used_percent" { + merged.fields.insert(suffix.clone(), value.clone()); + } + } + merged.window_minutes = incoming.window_minutes.or(current.window_minutes); + // Keep the established deadline when observations differ only by normal + // countdown/clock jitter so repeated responses cannot inch it forward. + merged.deadline = current.deadline.or(incoming.deadline); + merged.persist_deadline(); + + let used_percent = match (current.used_percent(), incoming.used_percent()) { + (Some(current), Some(incoming)) => Some(current.max(incoming)), + (Some(current), None) => Some(current), + (None, Some(incoming)) => Some(incoming), + (None, None) => None, + }; + if let Some(used_percent) = used_percent { + merged + .fields + .insert("used_percent".to_string(), json!(used_percent)); + } else { + merged.fields.remove("used_percent"); + } + merged +} + +fn codex_quota_merge_stale_same_window_usage( + current: &CodexQuotaWindowObservation, + incoming: &CodexQuotaWindowObservation, +) -> CodexQuotaWindowObservation { + let same_generation = match (current.deadline, incoming.deadline) { + (Some(current), Some(incoming)) => { + current.abs_diff(incoming) <= CODEX_QUOTA_RESET_DEADLINE_TOLERANCE_SECONDS + } + (None, None) => true, + _ => false, + }; + if !same_generation { + return current.clone(); + } + + let Some(incoming_used_percent) = incoming.used_percent() else { + return current.clone(); + }; + if current + .used_percent() + .is_some_and(|current_used_percent| current_used_percent >= incoming_used_percent) + { + return current.clone(); + } + + let mut merged = current.clone(); + merged + .fields + .insert("used_percent".to_string(), json!(incoming_used_percent)); + merged +} + +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +struct CodexQuotaWindowAssignmentScore { + unmatched: usize, + worst_deadline_rank: u8, + deadline_rank_sum: usize, + deadline_distance_sum: u128, + slot_mismatches: usize, + assignment_key: Vec, +} + +fn codex_quota_window_deadline_match_score( + current: &CodexQuotaWindowObservation, + incoming: &CodexQuotaWindowObservation, +) -> (u8, u64) { + match (current.deadline, incoming.deadline) { + (Some(current), Some(incoming)) => { + let distance = current.abs_diff(incoming); + ( + u8::from(distance > CODEX_QUOTA_RESET_DEADLINE_TOLERANCE_SECONDS), + distance, + ) + } + (None, None) => (0, 0), + _ => (2, 0), + } +} + +fn codex_quota_window_assignment_score( + current: &[CodexQuotaWindowObservation], + incoming: &[CodexQuotaWindowObservation], + assignment: &[Option], +) -> CodexQuotaWindowAssignmentScore { + let mut unmatched = 0; + let mut worst_deadline_rank = 0; + let mut deadline_rank_sum = 0; + let mut deadline_distance_sum = 0u128; + let mut slot_mismatches = 0; + + for (incoming_index, current_index) in assignment.iter().copied().enumerate() { + if !incoming[incoming_index].active() { + continue; + } + let Some(current_index) = current_index else { + unmatched += 1; + continue; + }; + let current_window = ¤t[current_index]; + let incoming_window = &incoming[incoming_index]; + let (deadline_rank, deadline_distance) = + codex_quota_window_deadline_match_score(current_window, incoming_window); + worst_deadline_rank = worst_deadline_rank.max(deadline_rank); + deadline_rank_sum += usize::from(deadline_rank); + deadline_distance_sum += u128::from(deadline_distance); + slot_mismatches += usize::from(current_window.slot != incoming_window.slot); + } + + CodexQuotaWindowAssignmentScore { + unmatched, + worst_deadline_rank, + deadline_rank_sum, + deadline_distance_sum, + slot_mismatches, + assignment_key: assignment + .iter() + .map(|index| index.unwrap_or(usize::MAX)) + .collect(), + } +} + +fn codex_quota_search_window_assignments( + current: &[CodexQuotaWindowObservation], + incoming: &[CodexQuotaWindowObservation], + incoming_index: usize, + used_current: &mut [bool], + assignment: &mut [Option], + best: &mut Option<(CodexQuotaWindowAssignmentScore, Vec>)>, +) { + if incoming_index == incoming.len() { + let score = codex_quota_window_assignment_score(current, incoming, assignment); + if best + .as_ref() + .is_none_or(|(best_score, _)| score < *best_score) + { + *best = Some((score, assignment.to_vec())); + } + return; + } + + assignment[incoming_index] = None; + codex_quota_search_window_assignments( + current, + incoming, + incoming_index + 1, + used_current, + assignment, + best, + ); + if !incoming[incoming_index].active() { + return; + } + + for (current_index, current_window) in current.iter().enumerate() { + if used_current[current_index] + || !current_window.active() + || !codex_quota_same_window_identity(current_window, &incoming[incoming_index]) + { + continue; + } + used_current[current_index] = true; + assignment[incoming_index] = Some(current_index); + codex_quota_search_window_assignments( + current, + incoming, + incoming_index + 1, + used_current, + assignment, + best, + ); + used_current[current_index] = false; + } + assignment[incoming_index] = None; +} + +fn codex_quota_match_windows( + current: &[CodexQuotaWindowObservation], + incoming: &[CodexQuotaWindowObservation], +) -> Vec> { + let mut best = None; + codex_quota_search_window_assignments( + current, + incoming, + 0, + &mut vec![false; current.len()], + &mut vec![None; incoming.len()], + &mut best, + ); + best.map(|(_, assignment)| assignment) + .unwrap_or_else(|| vec![None; incoming.len()]) +} + +fn codex_quota_reset_observation_proves_new_baseline( + current: &[CodexQuotaWindowObservation], + incoming: &[CodexQuotaWindowObservation], +) -> bool { + let current = current + .iter() + .filter(|window| window.active()) + .cloned() + .collect::>(); + let incoming = incoming + .iter() + .filter(|window| window.active()) + .cloned() + .collect::>(); + if current.is_empty() { + return !incoming.is_empty(); + } + + let incoming_is_confirmed_zero_baseline = !incoming.is_empty() + && incoming + .iter() + .all(|window| window.used_percent() == Some(0.0)); + if incoming_is_confirmed_zero_baseline + && current + .iter() + .all(|window| window.used_percent() == Some(0.0)) + { + return true; + } + + let window_matches = codex_quota_match_windows(¤t, &incoming); + incoming + .iter() + .enumerate() + .any(|(incoming_index, incoming_window)| { + let Some(current_index) = window_matches[incoming_index] else { + return true; + }; + let current_window = ¤t[current_index]; + if current_window + .deadline + .zip(incoming_window.deadline) + .is_some_and(|(current, incoming)| { + incoming > current.saturating_add(CODEX_QUOTA_RESET_DEADLINE_TOLERANCE_SECONDS) + }) + { + return true; + } + current_window + .used_percent() + .zip(incoming_window.used_percent()) + .is_some_and(|(current, incoming)| incoming < current) + }) +} + +fn codex_quota_write_family_windows( + object: &mut serde_json::Map, + family: CodexQuotaWindowFamily, + windows: &[CodexQuotaWindowObservation], +) { + for slot in CodexQuotaWindowSlot::ALL + .iter() + .copied() + .filter(|slot| slot.family() == family) + { + for suffix in CODEX_QUOTA_WINDOW_SUFFIXES { + object.remove(&codex_quota_window_key(slot, suffix)); + } + } + for window in windows.iter().filter(|window| window.active()) { + for (suffix, value) in &window.fields { + object.insert(codex_quota_window_key(window.slot, suffix), value.clone()); + } + } +} + +fn codex_quota_stabilize_legacy_deadlines( + current: &serde_json::Map, + merged: &mut serde_json::Map, +) { + let observed_at = current.get("updated_at").and_then(coerce_json_u64); + for family in [ + CodexQuotaWindowFamily::Account, + CodexQuotaWindowFamily::Spark, + ] { + let mut windows = codex_quota_read_family_windows(current, family, observed_at); + if windows.iter().any(|window| window.deadline.is_some()) { + for window in &mut windows { + window.persist_deadline(); + } + codex_quota_write_family_windows(merged, family, &windows); + } + } +} + +fn codex_quota_family_authoritative( + coverage: CodexQuotaWindowCoverage, + family: CodexQuotaWindowFamily, +) -> bool { + matches!(coverage, CodexQuotaWindowCoverage::FullSnapshot) + || matches!( + (coverage, family), + ( + CodexQuotaWindowCoverage::AccountSnapshot, + CodexQuotaWindowFamily::Account + ) + ) +} + +fn codex_quota_processes_family( + coverage: CodexQuotaWindowCoverage, + family: CodexQuotaWindowFamily, + incoming: &[CodexQuotaWindowObservation], +) -> bool { + match coverage { + CodexQuotaWindowCoverage::Patch => !incoming.is_empty(), + CodexQuotaWindowCoverage::AccountSnapshot => family == CodexQuotaWindowFamily::Account, + CodexQuotaWindowCoverage::FullSnapshot => true, + } +} + +fn codex_quota_apply_family( + current_object: &serde_json::Map, + incoming_object: &serde_json::Map, + merged: &mut serde_json::Map, + family: CodexQuotaWindowFamily, + context: CodexQuotaMergeContext<'_>, +) { + let active_reset_generation = current_object + .get(CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY) + .and_then(coerce_json_u64) + .unwrap_or(0); + let generation_matches = + codex_quota_observation_matches_reset_generation(current_object, context); + if family == CodexQuotaWindowFamily::Account && !generation_matches { + return; + } + let current_observed_at = current_object.get("updated_at").and_then(coerce_json_u64); + let current = codex_quota_read_family_windows(current_object, family, current_observed_at); + let mut incoming = codex_quota_read_family_windows( + incoming_object, + family, + Some(context.observed_at_unix_secs), + ); + if context.coverage == CodexQuotaWindowCoverage::Patch + && current.iter().filter(|window| window.active()).count() > 1 + { + // A partial header reports upstream's primary/secondary name, while + // paid accounts may store those windows in the opposite slots. Without + // a duration there is no stable identity, so leave both windows alone. + incoming.retain(|window| !window.active() || window.window_minutes.is_some()); + } + if !codex_quota_processes_family(context.coverage, family, &incoming) { + return; + } + + let account_reset_fence = (family == CodexQuotaWindowFamily::Account) + .then(|| codex_quota_account_reset_fence(current_object)) + .flatten(); + if active_reset_generation == 0 + && account_reset_fence.is_some_and(|fence| codex_quota_reset_fence_blocks(fence, context)) + { + return; + } + let authoritative = codex_quota_family_authoritative(context.coverage, family); + let pending_generation_matches = active_reset_generation > 0 + && current_object + .get(CODEX_QUOTA_ACCOUNT_RESET_PENDING_GENERATION_KEY) + .and_then(coerce_json_u64) + == Some(active_reset_generation); + let generation_authorizes_reset = pending_generation_matches + && (context.authoritative_reset_generation == Some(active_reset_generation) + || account_reset_fence.is_some_and(|fence| { + codex_quota_request_started_after_reset_fence(fence, context) + })); + let legacy_fence_authorizes_reset = active_reset_generation == 0 + && account_reset_fence.is_some_and(|fence| { + codex_quota_reset_fence_authorizes(fence, context) + || codex_quota_request_started_after_reset_fence(fence, context) + }); + let reset_baseline = account_reset_fence.is_some_and(|fence| { + fence.pending + && authoritative + && incoming.iter().any(CodexQuotaWindowObservation::active) + && codex_quota_reset_observation_proves_new_baseline(¤t, &incoming) + && (generation_authorizes_reset || legacy_fence_authorizes_reset) + }); + if account_reset_fence.is_some_and(|fence| fence.pending) && !reset_baseline { + return; + } + let stored_watermark = codex_quota_read_request_order( + current_object, + family.watermark_key(), + family.watermark_id_key(), + ); + let stale_family = !reset_baseline + && codex_quota_request_order_is_stale(context.request_order(), stored_watermark); + let window_matches = (!reset_baseline).then(|| codex_quota_match_windows(¤t, &incoming)); + let mut next = if reset_baseline || (authoritative && !stale_family) { + Vec::new() + } else { + current + .iter() + .filter(|window| window.active()) + .cloned() + .collect::>() + }; + + for (incoming_index, incoming_window) in incoming + .iter() + .enumerate() + .filter(|(_, window)| window.active()) + { + let current_index = window_matches + .as_ref() + .and_then(|matches| matches[incoming_index]); + let Some(current_index) = current_index else { + if !stale_family { + next.retain(|window| window.slot != incoming_window.slot); + let mut accepted = incoming_window.clone(); + accepted.persist_deadline(); + next.push(accepted); + } + continue; + }; + let mut merged_window = if stale_family { + codex_quota_merge_stale_same_window_usage(¤t[current_index], incoming_window) + } else { + codex_quota_merge_same_window(¤t[current_index], incoming_window) + }; + let target_slot = if authoritative && !stale_family { + incoming_window.slot + } else { + current[current_index].slot + }; + merged_window.slot = target_slot; + next.retain(|window| window.slot != target_slot); + next.push(merged_window); + } + + if !authoritative && !stale_family { + for incoming_window in incoming.iter().filter(|window| window.disabled) { + next.retain(|window| window.slot != incoming_window.slot); + } + } + next.sort_by_key(|window| window.slot); + + codex_quota_write_family_windows(merged, family, &next); + if reset_baseline { + merged.insert( + CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY.to_string(), + json!(false), + ); + } + if !stale_family { + if let Some(incoming_order) = context + .request_order() + .filter(|incoming| codex_quota_request_order_is_newer(*incoming, stored_watermark)) + { + codex_quota_write_request_order( + merged, + family.watermark_key(), + family.watermark_id_key(), + incoming_order, + ); + } + } +} + +fn codex_quota_semantic_metadata( + object: &serde_json::Map, +) -> serde_json::Map { + object + .iter() + .filter(|(key, _)| { + key.as_str() != "updated_at" + && !key.ends_with("_reset_seconds") + && !key.ends_with("_reset_after_seconds") + }) + .map(|(key, value)| (key.clone(), value.clone())) + .collect() +} + +/// Merge a parsed Codex quota observation into the stored flat metadata. +/// +/// Positive `window_minutes` values identify windows independently of the +/// primary/secondary storage slot. Within one reset deadline usage is +/// monotonic; advancing the deadline starts a new generation. Snapshot modes +/// replace the covered window families, while patch mode leaves absent windows +/// alone. A request-start/id watermark prevents a delayed request from +/// restoring a superseded window shape. +pub fn merge_codex_quota_metadata_snapshot( + current: Option<&serde_json::Value>, + incoming: &serde_json::Value, + context: CodexQuotaMergeContext<'_>, +) -> Option { + let incoming_object = incoming.as_object()?; + let current_object = current + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + if !codex_credential_generation_matches(current, context.observed_credential_generation) { + return Some(CodexQuotaMergeOutcome { + metadata: serde_json::Value::Object(current_object), + changed: false, + }); + } + let mut merged = current_object.clone(); + codex_quota_stabilize_legacy_deadlines(¤t_object, &mut merged); + + let stored_metadata_watermark = std::iter::once(codex_quota_read_request_order( + ¤t_object, + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_ID_KEY, + )) + .chain( + [ + CodexQuotaWindowFamily::Account, + CodexQuotaWindowFamily::Spark, + ] + .into_iter() + .map(|family| { + codex_quota_read_request_order( + ¤t_object, + family.watermark_key(), + family.watermark_id_key(), + ) + }), + ) + .flatten() + .max(); + let has_incoming_metadata = incoming_object.keys().any(|key| { + key != "updated_at" + && !codex_quota_is_request_order_key(key) + && !codex_quota_is_reset_fence_key(key) + && !codex_quota_is_window_key(key) + }); + let stale_metadata = has_incoming_metadata + && (codex_quota_request_order_is_stale(context.request_order(), stored_metadata_watermark) + || !codex_quota_observation_matches_reset_generation(¤t_object, context) + || (codex_quota_account_reset_generation(Some(&serde_json::Value::Object( + current_object.clone(), + ))) == 0 + && codex_quota_account_reset_fence(¤t_object) + .is_some_and(|fence| codex_quota_reset_fence_blocks(fence, context)))); + + if has_incoming_metadata && !stale_metadata { + for (key, value) in incoming_object { + if key == "updated_at" + || codex_quota_is_request_order_key(key) + || codex_quota_is_reset_fence_key(key) + || codex_quota_is_window_key(key) + { + continue; + } + merged.insert(key.clone(), value.clone()); + } + if let Some(incoming_order) = context.request_order().filter(|incoming| { + codex_quota_request_order_is_newer(*incoming, stored_metadata_watermark) + }) { + codex_quota_write_request_order( + &mut merged, + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_KEY, + CODEX_QUOTA_METADATA_REQUEST_WATERMARK_ID_KEY, + incoming_order, + ); + } + } + + codex_quota_apply_family( + ¤t_object, + incoming_object, + &mut merged, + CodexQuotaWindowFamily::Account, + context, + ); + codex_quota_apply_family( + ¤t_object, + incoming_object, + &mut merged, + CodexQuotaWindowFamily::Spark, + context, + ); + + let changed = + codex_quota_semantic_metadata(¤t_object) != codex_quota_semantic_metadata(&merged); + if changed { + let updated_at_unix_secs = current_object + .get("updated_at") + .and_then(coerce_json_u64) + .unwrap_or_default() + .max(context.observed_at_unix_secs); + merged.insert("updated_at".to_string(), json!(updated_at_unix_secs)); + Some(CodexQuotaMergeOutcome { + metadata: serde_json::Value::Object(merged), + changed: true, + }) + } else { + Some(CodexQuotaMergeOutcome { + metadata: serde_json::Value::Object(current_object), + changed: false, + }) + } +} + fn codex_write_window( target: &mut serde_json::Map, source: &serde_json::Map, @@ -661,6 +1846,37 @@ fn codex_window_has_active_limit(source: &serde_json::Map 0.0) } +fn codex_window_is_explicitly_disabled( + source: &serde_json::Map, +) -> bool { + let used_percent_is_zero = source + .get("used_percent") + .and_then(coerce_json_f64) + .is_some_and(|value| value == 0.0); + let reset_after_is_zero = source + .get("reset_after_seconds") + .and_then(coerce_json_u64) + .is_some_and(|value| value == 0); + let reset_at_is_empty = source.get("reset_at").is_some_and(|value| { + value.is_null() + || value.as_str().is_some_and(|value| value.trim().is_empty()) + || coerce_json_u64(value).is_some_and(|value| value == 0) + }); + let duration_is_zero = ["window_minutes", "limit_window_seconds"] + .iter() + .find_map(|key| source.get(*key).and_then(coerce_json_u64)) + .is_some_and(|value| value == 0); + + used_percent_is_zero && reset_after_is_zero && reset_at_is_empty && duration_is_zero +} + +fn codex_write_disabled_window( + target: &mut serde_json::Map, + target_prefix: &str, +) { + target.insert(format!("{target_prefix}_window_minutes"), json!(0u64)); +} + fn codex_find_spark_rate_limit( root: &serde_json::Map, ) -> Option<&serde_json::Map> { @@ -757,6 +1973,9 @@ pub fn parse_codex_wham_usage_response( codex_write_window(&mut result, &primary_window, "secondary"); } else { codex_write_window(&mut result, &primary_window, "primary"); + if codex_window_is_explicitly_disabled(&secondary_window) { + codex_write_disabled_window(&mut result, "secondary"); + } } if let Some(spark_rate_limit) = codex_find_spark_rate_limit(root) { @@ -1193,6 +2412,13 @@ pub fn parse_codex_usage_headers( .and_then(|value| value.parse::().ok()) { object.insert("reset_at".to_string(), json!(value)); + } else if normalized + .get(&reset_at_key) + .is_some_and(|value| value.is_empty()) + { + // Preserve an explicitly empty reset-at long enough to recognize + // the complete all-zero secondary-window disabled marker. + object.insert("reset_at".to_string(), serde_json::Value::Null); } if let Some(value) = normalized .get(&window_minutes_key) @@ -1212,6 +2438,9 @@ pub fn parse_codex_usage_headers( codex_write_window(&mut result, &primary_window, "secondary"); } else { codex_write_window(&mut result, &primary_window, "primary"); + if codex_window_is_explicitly_disabled(&secondary_window) { + codex_write_disabled_window(&mut result, "secondary"); + } } if let Some(value) = normalized @@ -1284,8 +2513,30 @@ fn codex_merge_invalid_reason(current: &str, candidate_reason: &str) -> String { } return format!("{current}\n{candidate_reason}"); } - if current.starts_with(OAUTH_EXPIRED_PREFIX) - && candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX) + if candidate_reason.starts_with(OAUTH_EXPIRED_PREFIX) { + let candidate_lines = candidate_reason + .lines() + .map(str::trim) + .filter(|line| !line.is_empty()) + .collect::>(); + let missing_refresh_failures = current + .lines() + .map(str::trim) + .filter(|line| line.starts_with(OAUTH_REFRESH_FAILED_PREFIX)) + .filter(|line| !candidate_lines.contains(line)) + .collect::>(); + if missing_refresh_failures.is_empty() { + return candidate_reason.to_string(); + } + return format!( + "{candidate_reason}\n{}", + missing_refresh_failures.join("\n") + ); + } + if candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX) + && current.lines().map(str::trim).any(|line| { + line.starts_with(OAUTH_EXPIRED_PREFIX) || line.starts_with(OAUTH_REFRESH_FAILED_PREFIX) + }) { return current.to_string(); } @@ -2122,22 +3373,1351 @@ pub fn parse_chatgpt_web_conversation_init_response( #[cfg(test)] mod tests { use super::{ - codex_build_invalid_state, codex_runtime_invalid_reason, extract_execution_error_detail, - normalize_codex_reset_credit_consume_outcome, parse_antigravity_usage_response, - parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response, - parse_codex_usage_headers, parse_codex_wham_reset_credits_detail_response, - parse_codex_wham_usage_response, parse_gemini_cli_retrieve_user_quota_response, + codex_build_invalid_state, codex_oauth_success_request_order_is_stale, + codex_runtime_invalid_reason, extract_execution_error_detail, + merge_codex_quota_metadata_snapshot, normalize_codex_reset_credit_consume_outcome, + parse_antigravity_usage_response, parse_chatgpt_web_conversation_init_response, + parse_codex_backend_me_response, parse_codex_usage_headers, + parse_codex_wham_reset_credits_detail_response, parse_codex_wham_usage_response, + parse_gemini_cli_retrieve_user_quota_response, parse_gemini_cli_v1internal_credits_response, parse_windsurf_model_configs_response, parse_windsurf_rate_limit_response, parse_windsurf_user_status_response, provider_auto_remove_quota_exhausted_keys, quota_refresh_success_invalid_state, - should_auto_remove_structured_reason, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, - OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX, + should_auto_remove_structured_reason, CodexQuotaMergeContext, CodexQuotaWindowCoverage, + OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, + OAUTH_REQUEST_FAILED_PREFIX, }; use aether_contracts::{ExecutionResult, ResponseBody}; use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey; use serde_json::json; use std::collections::BTreeMap; + fn merge_codex_quota( + current: Option<&serde_json::Value>, + incoming: &serde_json::Value, + observed_at_unix_secs: u64, + request_started_at_unix_ms: u64, + coverage: CodexQuotaWindowCoverage, + ) -> super::CodexQuotaMergeOutcome { + merge_codex_quota_metadata_snapshot( + current, + incoming, + CodexQuotaMergeContext { + observed_at_unix_secs, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: None, + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: None, + coverage, + }, + ) + .expect("quota metadata should merge") + } + + fn merge_codex_quota_ordered( + current: Option<&serde_json::Value>, + incoming: &serde_json::Value, + observed_at_unix_secs: u64, + request_started_at_unix_ms: u64, + request_order_id: &str, + coverage: CodexQuotaWindowCoverage, + ) -> super::CodexQuotaMergeOutcome { + merge_codex_quota_metadata_snapshot( + current, + incoming, + CodexQuotaMergeContext { + observed_at_unix_secs, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: Some(request_order_id), + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: None, + coverage, + }, + ) + .expect("quota metadata should merge") + } + + fn merge_codex_quota_after_explicit_reset( + current: Option<&serde_json::Value>, + incoming: &serde_json::Value, + observed_at_unix_secs: u64, + request_started_at_unix_ms: u64, + ) -> super::CodexQuotaMergeOutcome { + merge_codex_quota_metadata_snapshot( + current, + incoming, + CodexQuotaMergeContext { + observed_at_unix_secs, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: Some("reset-refresh"), + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: Some("reset-fence"), + coverage: CodexQuotaWindowCoverage::AccountSnapshot, + }, + ) + .expect("reset quota metadata should merge") + } + + #[test] + fn codex_oauth_success_allows_equal_quota_and_oauth_watermarks() { + let quota_watermark = json!({ + "account_quota_request_started_at_unix_ms": 1_000_u64, + "account_quota_request_id": "01900000-0000-7000-8000-000000000010" + }); + let oauth_watermark = json!({ + "oauth_state_request_started_at_unix_ms": 1_000_u64, + "oauth_state_request_id": "01900000-0000-7000-8000-000000000010" + }); + + for current in ["a_watermark, &oauth_watermark] { + assert!(!codex_oauth_success_request_order_is_stale( + current.as_object(), + Some(1_000), + Some("01900000-0000-7000-8000-000000000010"), + )); + } + } + + #[test] + fn codex_oauth_success_rejects_older_order_and_same_millisecond_lower_id() { + let current = json!({ + "quota_metadata_request_started_at_unix_ms": 1_000_u64, + "quota_metadata_request_id": "01900000-0000-7000-8000-000000000010" + }); + + assert!(codex_oauth_success_request_order_is_stale( + current.as_object(), + Some(999), + Some("01900000-0000-7000-8000-000000000099"), + )); + assert!(codex_oauth_success_request_order_is_stale( + current.as_object(), + Some(1_000), + Some("01900000-0000-7000-8000-000000000009"), + )); + } + + #[test] + fn codex_quota_merge_keeps_usage_monotonic_within_one_generation() { + let current = json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 90_000u64, + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 50.0, + "primary_reset_at": 2_012u64, + "primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &incoming, + 101, + 90_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + } + + #[test] + fn codex_quota_same_value_advances_request_watermark_once() { + let current = json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 90_000u64, + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &incoming, + 101, + 100_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!( + outcome.metadata["account_quota_request_started_at_unix_ms"], + json!(100_000u64) + ); + assert_eq!(outcome.metadata["updated_at"], json!(101u64)); + + let duplicate = merge_codex_quota( + Some(&outcome.metadata), + &incoming, + 102, + 100_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + assert!(!duplicate.changed); + assert_eq!(duplicate.metadata, outcome.metadata); + } + + #[test] + fn codex_quota_request_id_breaks_same_millisecond_ties() { + let first = merge_codex_quota_ordered( + None, + &json!({ + "plan_type": "free", + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }), + 100, + 100_000, + "request-a", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + assert_eq!( + first.metadata["account_quota_request_id"], + json!("request-a") + ); + assert_eq!( + first.metadata["quota_metadata_request_id"], + json!("request-a") + ); + + let newer = merge_codex_quota_ordered( + Some(&first.metadata), + &json!({ + "plan_type": "team", + "primary_used_percent": 5.0, + "primary_reset_at": 3_000_000u64, + "primary_window_minutes": 43_800u64 + }), + 101, + 100_000, + "request-z", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + assert!(newer.changed); + assert_eq!(newer.metadata["plan_type"], json!("team")); + assert_eq!(newer.metadata["primary_window_minutes"], json!(43_800u64)); + assert_eq!( + newer.metadata["account_quota_request_id"], + json!("request-z") + ); + assert_eq!( + newer.metadata["quota_metadata_request_id"], + json!("request-z") + ); + + let delayed = merge_codex_quota_ordered( + Some(&newer.metadata), + &json!({ + "plan_type": "plus", + "primary_used_percent": 80.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }), + 102, + 100_000, + "request-m", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + assert!(!delayed.changed); + assert_eq!(delayed.metadata, newer.metadata); + } + + #[test] + fn codex_quota_stale_same_millisecond_request_cannot_advance_generation() { + let newer = merge_codex_quota_ordered( + None, + &json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 1_001u64, + "primary_window_minutes": 300u64 + }), + 1, + 100_000, + "request-b", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + assert_eq!(newer.metadata["primary_reset_at"], json!(1_001u64)); + assert_eq!(newer.metadata["primary_used_percent"], json!(60.0)); + + let delayed = merge_codex_quota_ordered( + Some(&newer.metadata), + &json!({ + "primary_used_percent": 5.0, + "primary_reset_after_seconds": 60u64, + "primary_window_minutes": 300u64 + }), + 1_000, + 100_000, + "request-a", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!delayed.changed); + assert_eq!(delayed.metadata, newer.metadata); + assert_eq!(delayed.metadata["primary_reset_at"], json!(1_001u64)); + assert_eq!(delayed.metadata["primary_used_percent"], json!(60.0)); + assert_eq!( + delayed.metadata["account_quota_request_id"], + json!("request-b") + ); + } + + #[test] + fn codex_quota_stale_request_can_raise_usage_within_same_generation() { + let current = json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 1_001u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "account_quota_request_id": "request-b", + "updated_at": 1u64 + }); + let delayed = merge_codex_quota_ordered( + Some(¤t), + &json!({ + "primary_used_percent": 80.0, + "primary_reset_at": 1_010u64, + "primary_window_minutes": 300u64 + }), + 2, + 100_000, + "request-a", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(delayed.changed); + assert_eq!(delayed.metadata["primary_used_percent"], json!(80.0)); + assert_eq!(delayed.metadata["primary_reset_at"], json!(1_001u64)); + assert_eq!( + delayed.metadata["account_quota_request_id"], + json!("request-b") + ); + } + + #[test] + fn codex_quota_request_id_supersedes_legacy_same_millisecond_watermark() { + let legacy = json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + + let outcome = merge_codex_quota_ordered( + Some(&legacy), + &json!({ + "primary_used_percent": 5.0, + "primary_reset_at": 3_000_000u64, + "primary_window_minutes": 43_800u64 + }), + 101, + 100_000, + "request-a", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_window_minutes"], json!(43_800u64)); + assert_eq!( + outcome.metadata["account_quota_request_id"], + json!("request-a") + ); + } + + #[test] + fn codex_quota_merge_allows_usage_drop_after_deadline_advances() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 2.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &incoming, + 101, + 100_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_used_percent"], json!(2.0)); + assert_eq!(outcome.metadata["primary_reset_at"], json!(20_000u64)); + assert_eq!(outcome.metadata["updated_at"], json!(101u64)); + } + + #[test] + fn codex_quota_deadline_tolerance_has_inclusive_thirty_second_boundary() { + let current = json!({ + "primary_used_percent": 90.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "updated_at": 100u64 + }); + let cases = [ + (2_030u64, 10.0, 90.0, 2_000u64), + (2_031u64, 10.0, 10.0, 2_031u64), + (1_970u64, 95.0, 95.0, 2_000u64), + (1_969u64, 95.0, 90.0, 2_000u64), + ]; + + for (incoming_deadline, incoming_usage, expected_usage, expected_deadline) in cases { + let outcome = merge_codex_quota( + Some(¤t), + &json!({ + "primary_used_percent": incoming_usage, + "primary_reset_at": incoming_deadline, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert_eq!( + outcome.metadata["primary_used_percent"], + json!(expected_usage), + "incoming deadline {incoming_deadline}" + ); + assert_eq!( + outcome.metadata["primary_reset_at"], + json!(expected_deadline), + "incoming deadline {incoming_deadline}" + ); + } + } + + #[test] + fn codex_quota_legacy_window_without_deadline_stays_conservative() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_window_minutes": 300u64, + "updated_at": 100u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &json!({ + "primary_used_percent": 1.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert_eq!(outcome.metadata["primary_used_percent"], json!(100.0)); + assert_eq!(outcome.metadata["primary_reset_at"], json!(20_000u64)); + } + + #[test] + fn codex_quota_explicit_reset_allows_usage_drop_with_same_deadline() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "account_quota_request_id": "before-reset", + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_processed_ids": ["redeem-once"], + "account_quota_reset_pending": true, + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }); + + let outcome = + merge_codex_quota_after_explicit_reset(Some(¤t), &incoming, 110, 110_000); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_used_percent"], json!(0.0)); + assert_eq!(outcome.metadata["primary_reset_at"], json!(20_000u64)); + assert_eq!( + outcome.metadata["account_quota_reset_pending"], + json!(false) + ); + } + + #[test] + fn codex_quota_reset_fence_rejects_pre_reset_response_after_baseline() { + let baseline = merge_codex_quota_after_explicit_reset( + Some(&json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "account_quota_request_id": "before-reset", + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": true, + "updated_at": 100u64 + })), + &json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + ); + + let delayed = merge_codex_quota_ordered( + Some(&baseline.metadata), + &json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 120, + 100_000, + "old-in-flight", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!delayed.changed); + assert_eq!(delayed.metadata, baseline.metadata); + assert_eq!(delayed.metadata["primary_used_percent"], json!(0.0)); + } + + #[test] + fn codex_quota_reset_fence_rejects_pre_reset_account_metadata() { + let current = json!({ + "plan_type": "plus", + "reset_credits": { + "available_count": 0, + "updated_at": 110, + }, + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 110_000u64, + "account_quota_request_id": "after-reset", + "quota_metadata_request_started_at_unix_ms": 110_000u64, + "quota_metadata_request_id": "after-reset", + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": false, + "updated_at": 110u64 + }); + let delayed = merge_codex_quota_ordered( + Some(¤t), + &json!({ + "plan_type": "team", + "reset_credits": { + "available_count": 1, + "updated_at": 100, + }, + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 120, + 100_000, + "before-reset", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!delayed.changed); + assert_eq!(delayed.metadata, current); + assert_eq!(delayed.metadata["plan_type"], json!("plus")); + assert_eq!( + delayed.metadata["reset_credits"]["available_count"], + json!(0) + ); + } + + #[test] + fn codex_quota_pending_reset_ignores_runtime_patch_but_not_spark() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "spark_primary_used_percent": 20.0, + "spark_primary_reset_at": 30_000u64, + "spark_primary_window_minutes": 300u64, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": true, + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "spark_primary_used_percent": 30.0, + "spark_primary_reset_at": 30_000u64, + "spark_primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota_ordered( + Some(¤t), + &incoming, + 110, + 110_000, + "after-reset-runtime", + CodexQuotaWindowCoverage::Patch, + ); + + assert_eq!(outcome.metadata["primary_used_percent"], json!(100.0)); + assert_eq!(outcome.metadata["spark_primary_used_percent"], json!(30.0)); + assert_eq!(outcome.metadata["account_quota_reset_pending"], json!(true)); + } + + #[test] + fn codex_quota_reset_fence_treats_same_millisecond_request_as_pre_reset() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": true, + "updated_at": 100u64 + }); + + let outcome = merge_codex_quota_ordered( + Some(¤t), + &json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 110, + 105_000, + "uuid-that-sorts-after-fence", + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + } + + #[test] + fn codex_quota_reset_pending_waits_for_usage_to_drop() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": true, + "updated_at": 100u64 + }); + let unchanged = merge_codex_quota_after_explicit_reset( + Some(¤t), + &json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + ); + assert!(!unchanged.changed); + assert_eq!( + unchanged.metadata["account_quota_reset_pending"], + json!(true) + ); + + let reset = merge_codex_quota_after_explicit_reset( + Some(&unchanged.metadata), + &json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 111, + 111_000, + ); + assert!(reset.changed); + assert_eq!(reset.metadata["primary_used_percent"], json!(0.0)); + assert_eq!(reset.metadata["account_quota_reset_pending"], json!(false)); + } + + #[test] + fn codex_quota_reset_pending_disambiguates_equal_duration_slot_swap() { + let current = json!({ + "primary_used_percent": 90.0, + "primary_reset_at": 1_000u64, + "primary_window_minutes": 300u64, + "secondary_used_percent": 10.0, + "secondary_reset_at": 2_000u64, + "secondary_window_minutes": 300u64, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": true, + "updated_at": 100u64 + }); + let swapped_but_increased = json!({ + "primary_used_percent": 20.0, + "primary_reset_at": 2_005u64, + "primary_window_minutes": 300u64, + "secondary_used_percent": 95.0, + "secondary_reset_at": 1_005u64, + "secondary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota_after_explicit_reset( + Some(¤t), + &swapped_but_increased, + 110, + 110_000, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + assert_eq!(outcome.metadata["account_quota_reset_pending"], json!(true)); + } + + #[test] + fn codex_quota_reset_pending_accepts_authoritative_zero_after_zero_baseline() { + let current = json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset-fence", + "account_quota_reset_pending": true, + "updated_at": 100u64 + }); + let confirmed = merge_codex_quota_after_explicit_reset( + Some(¤t), + &json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + ); + + assert!(confirmed.changed); + assert_eq!(confirmed.metadata["primary_used_percent"], json!(0.0)); + assert_eq!( + confirmed.metadata["account_quota_reset_pending"], + json!(false) + ); + } + + #[test] + fn codex_quota_pre_reset_generation_cannot_touch_account_but_spark_still_merges() { + let current = json!({ + "plan_type": "plus", + "primary_used_percent": 20.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "spark_primary_used_percent": 10.0, + "spark_primary_reset_at": 3_000u64, + "spark_primary_window_minutes": 300u64, + "account_quota_reset_generation": 2u64, + "updated_at": 100u64 + }); + let incoming = json!({ + "plan_type": "team", + "primary_used_percent": 90.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "spark_primary_used_percent": 30.0, + "spark_primary_reset_at": 3_000u64, + "spark_primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota_metadata_snapshot( + Some(¤t), + &incoming, + CodexQuotaMergeContext { + observed_at_unix_secs: 110, + request_started_at_unix_ms: Some(110_000), + request_order_id: Some("pre-reset"), + observed_reset_generation: Some(1), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: None, + coverage: CodexQuotaWindowCoverage::FullSnapshot, + }, + ) + .expect("quota metadata should merge"); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["plan_type"], json!("plus")); + assert_eq!(outcome.metadata["primary_used_percent"], json!(20.0)); + assert_eq!(outcome.metadata["spark_primary_used_percent"], json!(30.0)); + } + + #[test] + fn codex_quota_credential_generation_mismatch_rejects_every_field() { + let current = json!({ + "credential_generation": "credential-new", + "plan_type": "plus", + "primary_used_percent": 20.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "spark_primary_used_percent": 10.0, + "spark_primary_reset_at": 3_000u64, + "spark_primary_window_minutes": 300u64, + "updated_at": 100u64 + }); + let incoming = json!({ + "plan_type": "team", + "primary_used_percent": 90.0, + "spark_primary_used_percent": 30.0 + }); + + let outcome = merge_codex_quota_metadata_snapshot( + Some(¤t), + &incoming, + CodexQuotaMergeContext { + observed_at_unix_secs: 110, + request_started_at_unix_ms: Some(110_000), + request_order_id: Some("old-credential"), + observed_reset_generation: Some(0), + authoritative_reset_generation: None, + observed_credential_generation: Some("credential-old"), + account_reset_fence_id: None, + coverage: CodexQuotaWindowCoverage::FullSnapshot, + }, + ) + .expect("quota metadata should be acknowledged"); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + } + + #[test] + fn codex_quota_only_latest_reset_generation_can_close_pending_zero_baseline() { + let current = json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_reset_generation": 2u64, + "account_quota_reset_pending_generation": 2u64, + "account_quota_reset_pending": true, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset:b", + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 0.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }); + let merge = |generation| { + merge_codex_quota_metadata_snapshot( + Some(¤t), + &incoming, + CodexQuotaMergeContext { + observed_at_unix_secs: 110, + request_started_at_unix_ms: Some(110_000), + request_order_id: Some("reset-refresh"), + observed_reset_generation: Some(generation), + authoritative_reset_generation: Some(generation), + observed_credential_generation: None, + account_reset_fence_id: Some("reset:b"), + coverage: CodexQuotaWindowCoverage::AccountSnapshot, + }, + ) + .expect("quota metadata should merge") + }; + + let stale = merge(1); + assert!(!stale.changed); + assert_eq!(stale.metadata["account_quota_reset_pending"], json!(true)); + + let current_generation = merge(2); + assert!(current_generation.changed); + assert_eq!( + current_generation.metadata["account_quota_reset_pending"], + json!(false) + ); + } + + #[test] + fn codex_quota_pending_reset_recovers_from_later_current_generation_snapshot() { + let current = json!({ + "primary_used_percent": 100.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_reset_generation": 2u64, + "account_quota_reset_pending_generation": 2u64, + "account_quota_reset_pending": true, + "account_quota_reset_fence_unix_ms": 105_000u64, + "account_quota_reset_fence_id": "reset:b", + "updated_at": 100u64 + }); + let incoming = json!({ + "primary_used_percent": 20.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }); + let merge = |request_started_at_unix_ms, coverage| { + merge_codex_quota_metadata_snapshot( + Some(¤t), + &incoming, + CodexQuotaMergeContext { + observed_at_unix_secs: 120, + request_started_at_unix_ms: Some(request_started_at_unix_ms), + request_order_id: Some("later-refresh"), + observed_reset_generation: Some(2), + authoritative_reset_generation: None, + observed_credential_generation: None, + account_reset_fence_id: None, + coverage, + }, + ) + .expect("quota metadata should merge") + }; + + let pre_reset_request = merge(100_000, CodexQuotaWindowCoverage::AccountSnapshot); + assert!(!pre_reset_request.changed); + assert_eq!( + pre_reset_request.metadata["account_quota_reset_pending"], + json!(true) + ); + + let partial_headers = merge(120_000, CodexQuotaWindowCoverage::Patch); + assert!(!partial_headers.changed); + assert_eq!( + partial_headers.metadata["account_quota_reset_pending"], + json!(true) + ); + + let settled = merge(120_000, CodexQuotaWindowCoverage::AccountSnapshot); + assert!(settled.changed); + assert_eq!(settled.metadata["primary_used_percent"], json!(20.0)); + assert_eq!( + settled.metadata["account_quota_reset_pending"], + json!(false) + ); + } + + #[test] + fn codex_quota_merge_derives_stable_deadline_from_observation_time() { + let first = merge_codex_quota( + None, + &json!({ + "primary_used_percent": 60.0, + "primary_reset_after_seconds": 900u64, + "primary_window_minutes": 300u64 + }), + 1_000, + 900_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + assert_eq!(first.metadata["primary_reset_at"], json!(1_900u64)); + + let delayed = merge_codex_quota( + Some(&first.metadata), + &json!({ + "primary_used_percent": 50.0, + "primary_reset_after_seconds": 890u64, + "primary_window_minutes": 300u64 + }), + 1_010, + 800_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!delayed.changed); + assert_eq!(delayed.metadata["primary_used_percent"], json!(60.0)); + assert_eq!(delayed.metadata["primary_reset_at"], json!(1_900u64)); + } + + #[test] + fn codex_quota_merge_replaces_paid_windows_with_monthly_shape() { + let current = json!({ + "primary_used_percent": 70.0, + "primary_reset_at": 10_000u64, + "primary_window_minutes": 10_080u64, + "secondary_used_percent": 20.0, + "secondary_reset_at": 2_000u64, + "secondary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + let monthly = json!({ + "primary_used_percent": 5.0, + "primary_reset_at": 3_000_000u64, + "primary_window_minutes": 43_800u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &monthly, + 200, + 150_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_window_minutes"], json!(43_800u64)); + assert_eq!(outcome.metadata["primary_used_percent"], json!(5.0)); + assert!(outcome.metadata.get("secondary_used_percent").is_none()); + assert!(outcome.metadata.get("secondary_window_minutes").is_none()); + } + + #[test] + fn codex_quota_merge_stale_request_cannot_restore_old_window_shape() { + let monthly = json!({ + "primary_used_percent": 5.0, + "primary_reset_at": 3_000_000u64, + "primary_window_minutes": 43_800u64, + "account_quota_request_started_at_unix_ms": 150_000u64, + "updated_at": 200u64 + }); + let delayed_paid = json!({ + "primary_used_percent": 70.0, + "primary_reset_at": 10_000u64, + "primary_window_minutes": 10_080u64, + "secondary_used_percent": 20.0, + "secondary_reset_at": 2_000u64, + "secondary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(&monthly), + &delayed_paid, + 210, + 100_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, monthly); + } + + #[test] + fn codex_quota_merge_matches_duration_when_windows_move_slots() { + let free = json!({ + "primary_used_percent": 40.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + let paid = json!({ + "primary_used_percent": 10.0, + "primary_reset_at": 10_000u64, + "primary_window_minutes": 10_080u64, + "secondary_used_percent": 30.0, + "secondary_reset_at": 2_005u64, + "secondary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(&free), + &paid, + 110, + 110_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_window_minutes"], json!(10_080u64)); + assert_eq!(outcome.metadata["secondary_window_minutes"], json!(300u64)); + assert_eq!(outcome.metadata["secondary_used_percent"], json!(40.0)); + } + + #[test] + fn codex_quota_merge_uses_deadline_to_disambiguate_equal_duration_slot_swap() { + let current = json!({ + "primary_used_percent": 90.0, + "primary_reset_at": 1_000u64, + "primary_window_minutes": 300u64, + "secondary_used_percent": 10.0, + "secondary_reset_at": 2_000u64, + "secondary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + let swapped = json!({ + "primary_used_percent": 20.0, + "primary_reset_at": 2_005u64, + "primary_window_minutes": 300u64, + "secondary_used_percent": 95.0, + "secondary_reset_at": 1_005u64, + "secondary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &swapped, + 110, + 110_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_used_percent"], json!(20.0)); + assert_eq!(outcome.metadata["primary_reset_at"], json!(2_000u64)); + assert_eq!(outcome.metadata["secondary_used_percent"], json!(95.0)); + assert_eq!(outcome.metadata["secondary_reset_at"], json!(1_000u64)); + } + + #[test] + fn codex_quota_patch_ignores_primary_only_header_for_paid_windows() { + let current = json!({ + "primary_used_percent": 20.0, + "primary_reset_at": 10_000u64, + "primary_window_minutes": 10_080u64, + "secondary_used_percent": 40.0, + "secondary_reset_at": 2_000u64, + "secondary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + let headers = + BTreeMap::from([("x-codex-primary-used-percent".to_string(), "80".to_string())]); + let partial = parse_codex_usage_headers(&headers, 110) + .expect("partial Codex usage headers should parse"); + + let outcome = merge_codex_quota( + Some(¤t), + &partial, + 110, + 110_000, + CodexQuotaWindowCoverage::Patch, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + } + + #[test] + fn codex_quota_patch_without_duration_uses_slot_for_single_window() { + let current = json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &json!({ "primary_used_percent": 70.0 }), + 110, + 110_000, + CodexQuotaWindowCoverage::Patch, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_used_percent"], json!(70.0)); + assert_eq!(outcome.metadata["primary_window_minutes"], json!(300u64)); + } + + #[test] + fn codex_quota_patch_with_duration_matches_across_paid_slots() { + let current = json!({ + "primary_used_percent": 20.0, + "primary_reset_at": 10_000u64, + "primary_window_minutes": 10_080u64, + "secondary_used_percent": 40.0, + "secondary_reset_at": 2_000u64, + "secondary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 100u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &json!({ + "primary_used_percent": 70.0, + "primary_reset_at": 2_005u64, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + CodexQuotaWindowCoverage::Patch, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_used_percent"], json!(20.0)); + assert_eq!(outcome.metadata["secondary_used_percent"], json!(70.0)); + assert_eq!(outcome.metadata["secondary_window_minutes"], json!(300u64)); + } + + #[test] + fn codex_quota_account_snapshot_preserves_spark_windows() { + let current = json!({ + "primary_used_percent": 50.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "spark_primary_used_percent": 25.0, + "spark_primary_reset_at": 3_000u64, + "spark_primary_window_minutes": 300u64, + "updated_at": 100u64 + }); + let outcome = merge_codex_quota( + Some(¤t), + &json!({ + "primary_used_percent": 60.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }), + 110, + 110_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["spark_primary_used_percent"], json!(25.0)); + assert_eq!(outcome.metadata["spark_primary_reset_at"], json!(3_000u64)); + } + + #[test] + fn codex_quota_full_snapshot_removes_absent_spark_and_zero_windows() { + let current = json!({ + "primary_used_percent": 50.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "secondary_used_percent": 20.0, + "secondary_window_minutes": 10_080u64, + "spark_primary_used_percent": 25.0, + "spark_primary_window_minutes": 300u64, + "updated_at": 100u64 + }); + let outcome = merge_codex_quota( + Some(¤t), + &json!({ + "primary_used_percent": 55.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64, + "secondary_window_minutes": 0u64 + }), + 110, + 110_000, + CodexQuotaWindowCoverage::FullSnapshot, + ); + + assert!(outcome.changed); + assert!(outcome.metadata.get("secondary_window_minutes").is_none()); + assert!(outcome.metadata.get("spark_primary_used_percent").is_none()); + assert!(outcome + .metadata + .get("spark_primary_window_minutes") + .is_none()); + } + + #[test] + fn codex_quota_merge_legacy_slot_without_identity_is_monotonic() { + let current = json!({ + "primary_used_percent": 60.0, + "account_quota_request_started_at_unix_ms": 110_000u64, + "updated_at": 100u64 + }); + let outcome = merge_codex_quota( + Some(¤t), + &json!({ "primary_used_percent": 50.0 }), + 110, + 110_000, + CodexQuotaWindowCoverage::Patch, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + } + + #[test] + fn codex_quota_stale_request_cannot_overwrite_account_metadata() { + let current = json!({ + "plan_type": "free", + "credits_balance": 20.0, + "primary_used_percent": 10.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 43_800u64, + "quota_metadata_request_started_at_unix_ms": 200_000u64, + "account_quota_request_started_at_unix_ms": 200_000u64, + "updated_at": 200u64 + }); + let stale = json!({ + "plan_type": "plus", + "credits_balance": 2.0, + "primary_used_percent": 80.0, + "primary_reset_at": 2_000u64, + "primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &stale, + 210, + 100_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(!outcome.changed); + assert_eq!(outcome.metadata, current); + } + + #[test] + fn codex_quota_metadata_uses_the_newest_family_watermark() { + let current = json!({ + "plan_type": "free", + "credits_balance": 20.0, + "primary_used_percent": 40.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64, + "quota_metadata_request_started_at_unix_ms": 100_000u64, + "account_quota_request_started_at_unix_ms": 200_000u64, + "updated_at": 200u64 + }); + let delayed = json!({ + "plan_type": "plus", + "credits_balance": 2.0, + "primary_used_percent": 50.0, + "primary_reset_at": 20_000u64, + "primary_window_minutes": 300u64 + }); + + let outcome = merge_codex_quota( + Some(¤t), + &delayed, + 210, + 150_000, + CodexQuotaWindowCoverage::AccountSnapshot, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["plan_type"], json!("free")); + assert_eq!(outcome.metadata["credits_balance"], json!(20.0)); + assert_eq!(outcome.metadata["primary_used_percent"], json!(50.0)); + assert_eq!( + outcome.metadata["quota_metadata_request_started_at_unix_ms"], + json!(100_000u64) + ); + assert_eq!( + outcome.metadata["account_quota_request_started_at_unix_ms"], + json!(200_000u64) + ); + } + #[test] fn execution_error_detail_preserves_structured_code_and_message() { let result = ExecutionResult { @@ -2145,6 +4725,7 @@ mod tests { candidate_id: None, status_code: 401, headers: BTreeMap::new(), + response_observation: None, body: Some(ResponseBody { json_body: Some(json!({ "error": { @@ -2281,6 +4862,56 @@ mod tests { ); } + #[test] + fn codex_invalid_state_preserves_refresh_failure_when_oauth_expired_arrives_later() { + 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_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + )); + let expected = format!( + "{OAUTH_EXPIRED_PREFIX}session expired\n{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + ); + + assert_eq!( + codex_build_invalid_state(&key, format!("{OAUTH_EXPIRED_PREFIX}session expired"), 200,), + (Some(200), Some(expected.clone())) + ); + + key.oauth_invalid_reason = Some(format!( + "{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + )); + assert_eq!( + codex_build_invalid_state(&key, expected.clone(), 200), + (Some(200), Some(expected)) + ); + + key.oauth_invalid_reason = Some(format!( + "{OAUTH_EXPIRED_PREFIX}old session expired\n{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + )); + assert_eq!( + codex_build_invalid_state( + &key, + format!("{OAUTH_EXPIRED_PREFIX}new session expired"), + 300, + ), + ( + Some(300), + Some(format!( + "{OAUTH_EXPIRED_PREFIX}new session expired\n{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + )) + ) + ); + } + #[test] fn codex_invalid_state_keeps_oauth_expired_over_request_failure() { let mut key = StoredProviderCatalogKey::new( @@ -2308,6 +4939,37 @@ mod tests { ); } + #[test] + fn codex_invalid_state_keeps_refresh_failure_over_request_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_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + )); + + assert_eq!( + codex_build_invalid_state( + &key, + format!("{OAUTH_REQUEST_FAILED_PREFIX}账号状态检查失败"), + 200, + ), + ( + Some(100), + Some(format!( + "{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败 (401): refresh_token 无效" + )) + ) + ); + } + #[test] fn codex_invalid_state_allows_account_block_to_override_oauth_expired() { let mut key = StoredProviderCatalogKey::new( @@ -2585,7 +5247,7 @@ mod tests { } #[test] - fn parses_codex_monthly_header_without_zero_secondary_placeholder() { + fn codex_quota_parses_monthly_header_with_zero_secondary_tombstone() { let headers = BTreeMap::from([ ("x-codex-plan-type".to_string(), "team".to_string()), ("x-codex-primary-used-percent".to_string(), "14".to_string()), @@ -2625,6 +5287,111 @@ mod tests { Some(&json!(43_800u64)) ); assert!(parsed.get("secondary_used_percent").is_none()); + assert_eq!(parsed.get("secondary_window_minutes"), Some(&json!(0u64))); + } + + #[test] + fn codex_quota_monthly_header_patch_removes_previous_five_hour_window() { + let current = json!({ + "primary_used_percent": 30.0, + "primary_reset_at": 1_790_000_000u64, + "primary_window_minutes": 10_080u64, + "secondary_used_percent": 70.0, + "secondary_reset_at": 1_785_000_000u64, + "secondary_window_minutes": 300u64, + "account_quota_request_started_at_unix_ms": 100_000u64, + "updated_at": 1_784_000_000u64 + }); + let headers = BTreeMap::from([ + ("x-codex-plan-type".to_string(), "team".to_string()), + ("x-codex-primary-used-percent".to_string(), "14".to_string()), + ( + "x-codex-primary-reset-after-seconds".to_string(), + "2627672".to_string(), + ), + ( + "x-codex-primary-reset-at".to_string(), + "1786915122".to_string(), + ), + ( + "x-codex-primary-window-minutes".to_string(), + "43800".to_string(), + ), + ( + "x-codex-secondary-used-percent".to_string(), + "0".to_string(), + ), + ( + "x-codex-secondary-reset-after-seconds".to_string(), + "0".to_string(), + ), + ("x-codex-secondary-reset-at".to_string(), "".to_string()), + ( + "x-codex-secondary-window-minutes".to_string(), + "0".to_string(), + ), + ]); + let monthly = parse_codex_usage_headers(&headers, 1_784_287_450) + .expect("monthly Codex headers should parse"); + + let outcome = merge_codex_quota( + Some(¤t), + &monthly, + 1_784_287_450, + 110_000, + CodexQuotaWindowCoverage::Patch, + ); + + assert!(outcome.changed); + assert_eq!(outcome.metadata["primary_used_percent"], json!(14.0)); + assert_eq!(outcome.metadata["primary_window_minutes"], json!(43_800u64)); + assert!(outcome.metadata.get("secondary_used_percent").is_none()); + assert!(outcome.metadata.get("secondary_window_minutes").is_none()); + } + + #[test] + fn codex_quota_parses_monthly_body_with_complete_zero_secondary_tombstone() { + let parsed = parse_codex_wham_usage_response( + &json!({ + "plan_type": "team", + "rate_limit": { + "primary_window": { + "used_percent": 14.0, + "reset_after_seconds": 2_627_672u64, + "reset_at": 1_786_915_122u64, + "window_minutes": 43_800u64 + }, + "secondary_window": { + "used_percent": 0.0, + "reset_after_seconds": 0u64, + "reset_at": null, + "window_minutes": 0u64 + } + } + }), + 1_784_287_450, + ) + .expect("monthly Codex body should parse"); + + assert_eq!(parsed.get("secondary_window_minutes"), Some(&json!(0u64))); + } + + #[test] + fn codex_quota_partial_zero_secondary_header_does_not_emit_tombstone() { + let headers = BTreeMap::from([ + ("x-codex-primary-used-percent".to_string(), "14".to_string()), + ( + "x-codex-primary-window-minutes".to_string(), + "43800".to_string(), + ), + ( + "x-codex-secondary-window-minutes".to_string(), + "0".to_string(), + ), + ]); + let parsed = parse_codex_usage_headers(&headers, 1_784_287_450) + .expect("partial Codex headers should parse"); + assert!(parsed.get("secondary_window_minutes").is_none()); } diff --git a/crates/aether-contracts/src/frame.rs b/crates/aether-contracts/src/frame.rs index 7eed808ed..584d6b8a6 100644 --- a/crates/aether-contracts/src/frame.rs +++ b/crates/aether-contracts/src/frame.rs @@ -2,7 +2,10 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; -use crate::{ExecutionError, ExecutionStreamTerminalSummary, ExecutionTelemetry}; +use crate::{ + ExecutionError, ExecutionResponseObservation, ExecutionStreamTerminalSummary, + ExecutionTelemetry, +}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] @@ -21,6 +24,8 @@ pub enum StreamFramePayload { status_code: u16, #[serde(default)] headers: BTreeMap, + #[serde(default, skip_serializing_if = "Option::is_none")] + response_observation: Option, }, Data { #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/crates/aether-contracts/src/lib.rs b/crates/aether-contracts/src/lib.rs index 8dd71ebbe..f9a9233b1 100644 --- a/crates/aether-contracts/src/lib.rs +++ b/crates/aether-contracts/src/lib.rs @@ -19,7 +19,7 @@ pub use plan::{ TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY, TRANSPORT_POOL_SCOPE_KEY, }; -pub use result::{ExecutionResult, ExecutionTelemetry, ResponseBody}; +pub use result::{ExecutionResponseObservation, ExecutionResult, ExecutionTelemetry, ResponseBody}; pub use usage::{ ExecutionStreamTerminalSummary, StandardizedUsage, USAGE_SERVER_NOW_UNIX_MS_HEADER, }; diff --git a/crates/aether-contracts/src/result.rs b/crates/aether-contracts/src/result.rs index 995a4725d..3d4802dd9 100644 --- a/crates/aether-contracts/src/result.rs +++ b/crates/aether-contracts/src/result.rs @@ -5,6 +5,13 @@ use serde_json::Value; use crate::ExecutionError; +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ExecutionResponseObservation { + pub request_started_at_unix_ms: u64, + pub response_headers_observed_at_unix_ms: u64, + pub request_order_id: String, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct ExecutionTelemetry { #[serde(default, skip_serializing_if = "Option::is_none")] @@ -32,6 +39,8 @@ pub struct ExecutionResult { #[serde(default)] pub headers: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] + pub response_observation: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub body: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub telemetry: Option, diff --git a/crates/aether-data/adapters/mysql/src/provider_catalog.rs b/crates/aether-data/adapters/mysql/src/provider_catalog.rs index d7baba44d..208d57156 100644 --- a/crates/aether-data/adapters/mysql/src/provider_catalog.rs +++ b/crates/aether-data/adapters/mysql/src/provider_catalog.rs @@ -8,8 +8,8 @@ use sqlx::{ }; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate, + ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, @@ -936,6 +936,84 @@ WHERE id = ? self.reload_key(&key.id, "updated").await } + pub async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + validate_admin_key_cas_update(update)?; + let key = &update.key; + let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; + let rotation_json = update + .codex_rotation + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys Codex rotation is not serializable: {err}" + )) + })?; + let mut builder = QueryBuilder::::new("UPDATE provider_api_keys SET provider_id = "); + push_admin_key_assignments(&mut builder, key, updated_at)?; + if update.reset_oauth_runtime { + builder.push(", oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0"); + } + if let Some(rotation_json) = rotation_json.as_deref() { + builder + .push(", upstream_metadata = JSON_SET(COALESCE(upstream_metadata, JSON_OBJECT()), '$.codex', CAST(") + .push_bind(rotation_json) + .push(" AS JSON))"); + } + if update.codex_rotation.is_some() || update.reset_oauth_runtime { + builder.push(", status_snapshot = "); + match (update.codex_rotation.is_some(), update.reset_oauth_runtime) { + (true, true) => builder.push("JSON_SET(COALESCE(status_snapshot, JSON_OBJECT()), '$.quota', CAST('null' AS JSON), '$.oauth', CAST('null' AS JSON))"), + (true, false) => builder.push("JSON_SET(COALESCE(status_snapshot, JSON_OBJECT()), '$.quota', CAST('null' AS JSON))"), + (false, true) => builder.push("JSON_SET(COALESCE(status_snapshot, JSON_OBJECT()), '$.oauth', CAST('null' AS JSON))"), + (false, false) => unreachable!(), + }; + } + builder + .push(" WHERE id = ") + .push_bind(&key.id) + .push(" AND BINARY api_key <=> BINARY ") + .push_bind(update.expected_credential.encrypted_api_key.as_deref()) + .push(" AND BINARY auth_config <=> BINARY ") + .push_bind(update.expected_encrypted_auth_config.as_deref()) + .push(" AND BINARY auth_type = BINARY ") + .push_bind(&update.expected_credential.auth_type) + .push(" AND BINARY provider_id = BINARY ") + .push_bind(&update.expected_credential.provider_id) + .push( + " AND EXISTS (SELECT 1 FROM providers WHERE BINARY providers.id = BINARY provider_api_keys.provider_id AND BINARY providers.provider_type = BINARY ", + ) + .push_bind(&update.expected_credential.provider_type) + .push(")"); + if update.codex_rotation.is_some() { + builder + .push(" AND JSON_TYPE(COALESCE(upstream_metadata, JSON_OBJECT())) = 'OBJECT'") + .push(" AND NOT (BINARY api_key <=> BINARY ") + .push_bind(key.encrypted_api_key.as_deref()) + .push(" AND BINARY auth_config <=> BINARY ") + .push_bind(key.encrypted_auth_config.as_deref()) + .push(" AND BINARY auth_type = BINARY ") + .push_bind(&key.auth_type) + .push(" AND BINARY provider_id = BINARY ") + .push_bind(&key.provider_id) + .push(")"); + } + if update.codex_rotation.is_some() || update.reset_oauth_runtime { + builder.push(" AND JSON_TYPE(COALESCE(status_snapshot, JSON_OBJECT())) = 'OBJECT'"); + } + let rows_affected = builder + .build() + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + pub async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -1007,6 +1085,10 @@ WHERE id = ? || expected.auth_type.trim().is_empty() || expected.provider_id.trim().is_empty() || expected.provider_type.trim().is_empty() + || delete + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) { return Err(DataLayerError::InvalidInput( "provider catalog OAuth credential CAS delete contains empty fields".to_string(), @@ -1030,6 +1112,33 @@ WHERE id = ? ) .push_bind(&expected.provider_type) .push(")"); + if let Some(expected) = delete.expected_upstream_metadata_namespace.as_ref() { + let namespace_path = format!( + "$.{}", + serde_json::to_string(&expected.namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + let expected_value = expected + .expected_value + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata expected value is not serializable: {err}" + )) + })?; + builder + .push(" AND JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'") + .push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path) + .push(") <=> CAST(") + .push_bind(expected_value) + .push(" AS JSON)"); + } let rows_affected = builder .build() .execute(&self.pool) @@ -1339,6 +1448,25 @@ WHERE id = ? "provider catalog OAuth credential fence must not contain empty fields".to_string(), )); } + if update + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth runtime metadata namespace must not be empty".to_string(), + )); + } + if update + .upstream_metadata_namespace_to_remove + .as_deref() + .is_some_and(|namespace| namespace.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth runtime metadata namespace to remove must not be empty" + .to_string(), + )); + } if !update.status_snapshot_patch.is_object() { return Err(DataLayerError::InvalidInput( "provider catalog status snapshot patch must be an object".to_string(), @@ -1353,6 +1481,22 @@ WHERE id = ? "provider catalog upstream metadata patch must be an object".to_string(), )); } + if update + .upstream_metadata_namespace_to_remove + .as_ref() + .is_some_and(|namespace| { + update + .upstream_metadata_patch + .as_ref() + .and_then(serde_json::Value::as_object) + .is_some_and(|patch| patch.contains_key(namespace)) + }) + { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth runtime metadata namespace cannot be patched and removed in the same update" + .to_string(), + )); + } let mut builder = QueryBuilder::::new("UPDATE provider_api_keys SET oauth_invalid_at = "); builder @@ -1375,9 +1519,29 @@ WHERE id = ? "provider_api_keys.expires_at", )?); } - if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() { + if update.upstream_metadata_patch.is_some() + || update.upstream_metadata_namespace_to_remove.is_some() + { builder.push(", upstream_metadata = "); - push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?; + if update.upstream_metadata_namespace_to_remove.is_some() { + builder.push("JSON_REMOVE("); + } + if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() { + push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?; + } else { + builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')"); + } + if let Some(namespace) = update.upstream_metadata_namespace_to_remove.as_ref() { + let namespace_path = format!( + "$.{}", + serde_json::to_string(namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + builder.push(", ").push_bind(namespace_path).push(")"); + } } builder.push(", status_snapshot = "); push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; @@ -1411,6 +1575,39 @@ WHERE id = ? .push_bind(&expected.provider_type) .push(")"); } + if update.expected_upstream_metadata_namespace.is_some() + || update.upstream_metadata_patch.is_some() + || update.upstream_metadata_namespace_to_remove.is_some() + { + builder + .push(" AND JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'"); + } + if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() { + let namespace_path = format!( + "$.{}", + serde_json::to_string(&expected.namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + let expected_value = expected + .expected_value + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata expected value is not serializable: {err}" + )) + })?; + builder + .push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path) + .push(") <=> CAST(") + .push_bind(expected_value) + .push(" AS JSON)"); + } let rows_affected = builder .build() .execute(&self.pool) @@ -1618,6 +1815,7 @@ WHERE id = ? ) .push(" WHERE id = ") .push_bind(&update.key_id) + .push(" AND JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'") .push(" AND JSON_EXTRACT(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") .push_bind(namespace_path) .push(") <=> CAST(") @@ -1897,6 +2095,13 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository { Self::update_key(self, key).await } + async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + Self::compare_and_update_key_admin_state(self, update).await + } + async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -2426,6 +2631,10 @@ SET model_exclude_patterns = ?, updated_at = ? WHERE id = ? + AND BINARY provider_id = BINARY ? + AND BINARY auth_type = BINARY ? + AND BINARY api_key <=> BINARY ? + AND BINARY auth_config <=> BINARY ? "# } @@ -2500,7 +2709,152 @@ fn key_update_query( "provider_api_keys.model_exclude_patterns", )?) .bind(updated_at) - .bind(&key.id)) + .bind(&key.id) + .bind(&key.provider_id) + .bind(&key.auth_type) + .bind(&key.encrypted_api_key) + .bind(&key.encrypted_auth_config)) +} + +fn push_admin_key_assignments<'args>( + builder: &mut QueryBuilder<'args, MySql>, + key: &'args StoredProviderCatalogKey, + updated_at: i64, +) -> Result<(), DataLayerError> { + builder + .push_bind(&key.provider_id) + .push(", name = ") + .push_bind(&key.name) + .push(", api_key = ") + .push_bind(&key.encrypted_api_key) + .push(", auth_type = ") + .push_bind(&key.auth_type) + .push(", capabilities = ") + .push_bind(optional_json_to_string( + &key.capabilities, + "provider_api_keys.capabilities", + )?) + .push(", is_active = ") + .push_bind(key.is_active) + .push(", api_formats = ") + .push_bind(optional_json_to_string( + &key.api_formats, + "provider_api_keys.api_formats", + )?) + .push(", auth_type_by_format = ") + .push_bind(optional_json_to_string( + &key.auth_type_by_format, + "provider_api_keys.auth_type_by_format", + )?) + .push(", allow_auth_channel_mismatch_formats = ") + .push_bind(optional_json_to_string( + &key.allow_auth_channel_mismatch_formats, + "provider_api_keys.allow_auth_channel_mismatch_formats", + )?) + .push(", auth_config = ") + .push_bind(&key.encrypted_auth_config) + .push(", note = ") + .push_bind(&key.note) + .push(", internal_priority = ") + .push_bind(key.internal_priority) + .push(", rate_multipliers = ") + .push_bind(optional_json_to_string( + &key.rate_multipliers, + "provider_api_keys.rate_multipliers", + )?) + .push(", global_priority_by_format = ") + .push_bind(optional_json_to_string( + &key.global_priority_by_format, + "provider_api_keys.global_priority_by_format", + )?) + .push(", allowed_models = ") + .push_bind(optional_json_to_string( + &key.allowed_models, + "provider_api_keys.allowed_models", + )?) + .push(", expires_at = ") + .push_bind(optional_i64_from_u64( + key.expires_at_unix_secs, + "provider_api_keys.expires_at", + )?) + .push(", cache_ttl_minutes = ") + .push_bind(key.cache_ttl_minutes) + .push(", max_probe_interval_minutes = ") + .push_bind(key.max_probe_interval_minutes) + .push(", proxy = ") + .push_bind(optional_json_to_string( + &key.proxy, + "provider_api_keys.proxy", + )?) + .push(", fingerprint = ") + .push_bind(optional_json_to_string( + &key.fingerprint, + "provider_api_keys.fingerprint", + )?) + .push(", rpm_limit = ") + .push_bind(optional_i64_from_u32(key.rpm_limit)) + .push(", concurrent_limit = ") + .push_bind(key.concurrent_limit) + .push(", auto_fetch_models = ") + .push_bind(key.auto_fetch_models) + .push(", locked_models = ") + .push_bind(optional_json_to_string( + &key.locked_models, + "provider_api_keys.locked_models", + )?) + .push(", model_include_patterns = ") + .push_bind(optional_json_to_string( + &key.model_include_patterns, + "provider_api_keys.model_include_patterns", + )?) + .push(", model_exclude_patterns = ") + .push_bind(optional_json_to_string( + &key.model_exclude_patterns, + "provider_api_keys.model_exclude_patterns", + )?) + .push(", updated_at = ") + .push_bind(updated_at); + Ok(()) +} + +fn validate_admin_key_cas_update( + update: &ProviderCatalogKeyAdminCasUpdate, +) -> Result<(), DataLayerError> { + validate_key(&update.key)?; + let expected = &update.expected_credential; + if expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + || expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || update + .expected_encrypted_auth_config + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog admin credential fence contains empty fields".to_string(), + )); + } + let Some(rotation) = update.codex_rotation.as_ref() else { + return Ok(()); + }; + let valid_rotation = expected.provider_type.eq_ignore_ascii_case("codex") + && rotation.as_object().is_some_and(|object| { + object.len() == 1 + && object + .get("credential_generation") + .and_then(serde_json::Value::as_str) + .is_some_and(|generation| !generation.trim().is_empty()) + }); + if !valid_rotation { + return Err(DataLayerError::InvalidInput( + "provider catalog Codex rotation must contain only credential_generation".to_string(), + )); + } + Ok(()) } fn optional_json_from_string( @@ -2880,6 +3234,34 @@ mod tests { ] { assert!(!sql.contains(runtime_assignment)); } + assert!(sql.contains("binary api_key <=> binary ?")); + assert!(sql.contains("binary auth_config <=> binary ?")); + } + + #[test] + fn admin_credential_cas_has_atomic_rotation_guards() { + let source = include_str!("provider_catalog.rs"); + for predicate in [ + "BINARY api_key <=> BINARY ", + "BINARY auth_config <=> BINARY ", + "JSON_TYPE(COALESCE(upstream_metadata, JSON_OBJECT())) = 'OBJECT'", + "JSON_TYPE(COALESCE(status_snapshot, JSON_OBJECT())) = 'OBJECT'", + "JSON_SET(COALESCE(upstream_metadata, JSON_OBJECT()), '$.codex'", + "oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0", + ] { + assert!( + source.contains(predicate), + "missing admin CAS guard: {predicate}" + ); + } + } + + #[test] + fn runtime_metadata_cas_requires_an_object_metadata_root() { + let source = include_str!("provider_catalog.rs"); + assert!( + source.contains("JSON_TYPE(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'OBJECT'") + ); } #[tokio::test] async fn empty_id_lists_do_not_connect_to_lazy_pool() { diff --git a/crates/aether-data/adapters/postgres/src/provider_catalog.rs b/crates/aether-data/adapters/postgres/src/provider_catalog.rs index 13a3efbf4..5cd541702 100644 --- a/crates/aether-data/adapters/postgres/src/provider_catalog.rs +++ b/crates/aether-data/adapters/postgres/src/provider_catalog.rs @@ -9,8 +9,8 @@ use sqlx::{ }; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate, + ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, @@ -380,6 +380,10 @@ SET auth_type_by_format = $27, allow_auth_channel_mismatch_formats = $28 WHERE id = $1 + AND provider_id = $2 + AND auth_type = $4 + AND api_key IS NOT DISTINCT FROM $5 + AND auth_config IS NOT DISTINCT FROM $6 "#; const KEY_RUNTIME_HEALTH_CAS_SQL: &str = r#" @@ -405,6 +409,7 @@ SET ELSE TO_TIMESTAMP($5::double precision) END WHERE id = $1 + AND jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' AND (COALESCE(upstream_metadata, '{}'::jsonb) -> $2) IS NOT DISTINCT FROM $6::jsonb "#; @@ -500,6 +505,112 @@ fn key_update_query(key: &StoredProviderCatalogKey) -> Query<'_, Postgres, PgArg .bind(&key.allow_auth_channel_mismatch_formats) } +fn push_admin_key_assignments<'args>( + builder: &mut QueryBuilder<'args, Postgres>, + key: &'args StoredProviderCatalogKey, +) { + builder + .push_bind(&key.provider_id) + .push(", api_formats = ") + .push_bind(&key.api_formats) + .push(", auth_type = ") + .push_bind(&key.auth_type) + .push(", api_key = ") + .push_bind(&key.encrypted_api_key) + .push(", auth_config = ") + .push_bind(&key.encrypted_auth_config) + .push(", name = ") + .push_bind(&key.name) + .push(", note = ") + .push_bind(&key.note) + .push(", rate_multipliers = ") + .push_bind(&key.rate_multipliers) + .push(", internal_priority = ") + .push_bind(key.internal_priority) + .push(", global_priority_by_format = ") + .push_bind(&key.global_priority_by_format) + .push(", rpm_limit = ") + .push_bind(key.rpm_limit.map(|value| value as i32)) + .push(", concurrent_limit = ") + .push_bind(key.concurrent_limit) + .push(", allowed_models = ") + .push_bind(&key.allowed_models) + .push(", capabilities = ") + .push_bind(&key.capabilities) + .push(", cache_ttl_minutes = ") + .push_bind(key.cache_ttl_minutes) + .push(", max_probe_interval_minutes = ") + .push_bind(key.max_probe_interval_minutes) + .push(", auto_fetch_models = ") + .push_bind(key.auto_fetch_models) + .push(", locked_models = ") + .push_bind(&key.locked_models) + .push(", model_include_patterns = ") + .push_bind(&key.model_include_patterns) + .push(", model_exclude_patterns = ") + .push_bind(&key.model_exclude_patterns) + .push(", proxy = ") + .push_bind(&key.proxy) + .push(", fingerprint = ") + .push_bind(&key.fingerprint) + .push(", expires_at = CASE WHEN ") + .push_bind(key.expires_at_unix_secs.map(|value| value as f64)) + .push("::double precision IS NULL THEN NULL ELSE TO_TIMESTAMP(") + .push_bind(key.expires_at_unix_secs.map(|value| value as f64)) + .push("::double precision) END, is_active = ") + .push_bind(key.is_active) + .push(", updated_at = CASE WHEN ") + .push_bind(key.updated_at_unix_secs.map(|value| value as f64)) + .push("::double precision IS NULL THEN NOW() ELSE TO_TIMESTAMP(") + .push_bind(key.updated_at_unix_secs.map(|value| value as f64)) + .push("::double precision) END, auth_type_by_format = ") + .push_bind(&key.auth_type_by_format) + .push(", allow_auth_channel_mismatch_formats = ") + .push_bind(&key.allow_auth_channel_mismatch_formats); +} + +fn validate_admin_key_cas_update( + update: &ProviderCatalogKeyAdminCasUpdate, +) -> Result<(), DataLayerError> { + validate_key_for_update(&update.key)?; + let expected = &update.expected_credential; + if update.key.name.trim().is_empty() + || update.key.auth_type.trim().is_empty() + || expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + || expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || update + .expected_encrypted_auth_config + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog admin credential fence contains empty fields".to_string(), + )); + } + let Some(rotation) = update.codex_rotation.as_ref() else { + return Ok(()); + }; + let valid_rotation = expected.provider_type.eq_ignore_ascii_case("codex") + && rotation.as_object().is_some_and(|object| { + object.len() == 1 + && object + .get("credential_generation") + .and_then(serde_json::Value::as_str) + .is_some_and(|generation| !generation.trim().is_empty()) + }); + if !valid_rotation { + return Err(DataLayerError::InvalidInput( + "provider catalog Codex rotation must contain only credential_generation".to_string(), + )); + } + Ok(()) +} + #[derive(Debug, Clone)] pub struct SqlxProviderCatalogReadRepository { pool: PgPool, @@ -901,11 +1012,29 @@ WHERE id = $1 || expected.provider_id.trim().is_empty() || expected.provider_type.trim().is_empty() }) + || update + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) + || update + .upstream_metadata_namespace_to_remove + .as_deref() + .is_some_and(|namespace| namespace.trim().is_empty()) || !update.status_snapshot_patch.is_object() || update .upstream_metadata_patch .as_ref() .is_some_and(|patch| !patch.is_object()) + || update + .upstream_metadata_namespace_to_remove + .as_ref() + .is_some_and(|namespace| { + update + .upstream_metadata_patch + .as_ref() + .and_then(serde_json::Value::as_object) + .is_some_and(|patch| patch.contains_key(namespace)) + }) { return Err(DataLayerError::InvalidInput( "provider catalog OAuth runtime CAS requires key_id, auth_config, and object status patch" @@ -932,27 +1061,40 @@ SET ELSE TO_TIMESTAMP($7::double precision) END, upstream_metadata = CASE - WHEN $8::jsonb IS NULL THEN upstream_metadata - ELSE COALESCE(upstream_metadata, '{}'::jsonb) || $8::jsonb + WHEN $8::jsonb IS NULL AND $9::text IS NULL THEN upstream_metadata + WHEN $9::text IS NULL THEN COALESCE(upstream_metadata, '{}'::jsonb) || $8::jsonb + ELSE (COALESCE(upstream_metadata, '{}'::jsonb) || COALESCE($8::jsonb, '{}'::jsonb)) - $9 END, - status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $9::jsonb)::json, - error_count = CASE WHEN $10::boolean THEN 0 ELSE error_count END, + status_snapshot = (COALESCE(status_snapshot::jsonb, '{}'::jsonb) || $10::jsonb)::json, + error_count = CASE WHEN $11::boolean THEN 0 ELSE error_count END, updated_at = CASE - WHEN $11::double precision IS NULL THEN NOW() - ELSE TO_TIMESTAMP($11::double precision) + WHEN $12::double precision IS NULL THEN NOW() + ELSE TO_TIMESTAMP($12::double precision) END WHERE id = $1 - AND auth_config IS NOT DISTINCT FROM $12 - AND ($13::boolean IS FALSE OR api_key IS NOT DISTINCT FROM $14) - AND ($15::text IS NULL OR auth_type = $15) - AND ($16::text IS NULL OR provider_id = $16) + AND auth_config IS NOT DISTINCT FROM $13 AND ( - $17::text IS NULL + ($8::jsonb IS NULL AND $9::text IS NULL) + OR jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' + ) + AND ($14::boolean IS FALSE OR api_key IS NOT DISTINCT FROM $15) + AND ($16::text IS NULL OR auth_type = $16) + AND ($17::text IS NULL OR provider_id = $17) + AND ( + $18::text IS NULL OR EXISTS ( SELECT 1 FROM providers WHERE providers.id = provider_api_keys.provider_id - AND providers.provider_type = $17 + AND providers.provider_type = $18 + ) + ) + AND ( + $19::boolean IS FALSE + OR ( + jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' + AND (COALESCE(upstream_metadata, '{}'::jsonb) -> $20) + IS NOT DISTINCT FROM $21::jsonb ) ) "#, @@ -970,6 +1112,7 @@ WHERE id = $1 .map(|value| value as f64), ) .bind(update.upstream_metadata_patch.as_ref()) + .bind(update.upstream_metadata_namespace_to_remove.as_deref()) .bind(&update.status_snapshot_patch) .bind(update.reset_error_count) .bind(update.updated_at_unix_secs.map(|value| value as f64)) @@ -999,6 +1142,19 @@ WHERE id = $1 .as_ref() .map(|expected| expected.provider_type.as_str()), ) + .bind(update.expected_upstream_metadata_namespace.is_some()) + .bind( + update + .expected_upstream_metadata_namespace + .as_ref() + .map(|expected| expected.namespace.as_str()), + ) + .bind( + update + .expected_upstream_metadata_namespace + .as_ref() + .and_then(|expected| expected.expected_value.as_ref()), + ) .execute(&self.pool) .await .map_postgres_err()? @@ -2008,6 +2164,76 @@ WHERE id = $1 }) } + pub async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + validate_admin_key_cas_update(update)?; + let key = &update.key; + let mut builder = + QueryBuilder::::new("UPDATE provider_api_keys SET provider_id = "); + push_admin_key_assignments(&mut builder, key); + if update.reset_oauth_runtime { + builder.push(", oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0"); + } + if let Some(rotation) = update.codex_rotation.as_ref() { + builder + .push(", upstream_metadata = jsonb_set(COALESCE(upstream_metadata, '{}'::jsonb), '{codex}', ") + .push_bind(rotation) + .push("::jsonb, true)"); + } + if update.codex_rotation.is_some() || update.reset_oauth_runtime { + builder.push(", status_snapshot = "); + match (update.codex_rotation.is_some(), update.reset_oauth_runtime) { + (true, true) => builder.push("jsonb_set(jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{quota}', 'null'::jsonb, true), '{oauth}', 'null'::jsonb, true)::json"), + (true, false) => builder.push("jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{quota}', 'null'::jsonb, true)::json"), + (false, true) => builder.push("jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{oauth}', 'null'::jsonb, true)::json"), + (false, false) => unreachable!(), + }; + } + builder + .push(" WHERE id = ") + .push_bind(&key.id) + .push(" AND api_key IS NOT DISTINCT FROM ") + .push_bind(update.expected_credential.encrypted_api_key.as_deref()) + .push(" AND auth_config IS NOT DISTINCT FROM ") + .push_bind(update.expected_encrypted_auth_config.as_deref()) + .push(" AND auth_type = ") + .push_bind(&update.expected_credential.auth_type) + .push(" AND provider_id = ") + .push_bind(&update.expected_credential.provider_id) + .push( + " AND EXISTS (SELECT 1 FROM providers WHERE providers.id = provider_api_keys.provider_id AND providers.provider_type = ", + ) + .push_bind(&update.expected_credential.provider_type) + .push(")"); + if update.codex_rotation.is_some() { + builder + .push(" AND jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'") + .push(" AND NOT (api_key IS NOT DISTINCT FROM ") + .push_bind(key.encrypted_api_key.as_deref()) + .push(" AND auth_config IS NOT DISTINCT FROM ") + .push_bind(key.encrypted_auth_config.as_deref()) + .push(" AND auth_type = ") + .push_bind(&key.auth_type) + .push(" AND provider_id = ") + .push_bind(&key.provider_id) + .push(")"); + } + if update.codex_rotation.is_some() || update.reset_oauth_runtime { + builder.push( + " AND jsonb_typeof(COALESCE(status_snapshot::jsonb, '{}'::jsonb)) = 'object'", + ); + } + let rows_affected = builder + .build() + .execute(&self.pool) + .await + .map_postgres_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + pub async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -2088,6 +2314,10 @@ WHERE id = $1 || expected.auth_type.trim().is_empty() || expected.provider_id.trim().is_empty() || expected.provider_type.trim().is_empty() + || delete + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) { return Err(DataLayerError::InvalidInput( "provider catalog OAuth credential CAS delete contains empty fields".to_string(), @@ -2107,6 +2337,14 @@ WHERE id = $1 WHERE providers.id = provider_api_keys.provider_id AND providers.provider_type = $6 ) + AND ( + $7::boolean IS FALSE + OR ( + jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object' + AND (COALESCE(upstream_metadata, '{}'::jsonb) -> $8) + IS NOT DISTINCT FROM $9::jsonb + ) + ) "#, ) .bind(&delete.key_id) @@ -2115,6 +2353,19 @@ WHERE id = $1 .bind(&expected.auth_type) .bind(&expected.provider_id) .bind(&expected.provider_type) + .bind(delete.expected_upstream_metadata_namespace.is_some()) + .bind( + delete + .expected_upstream_metadata_namespace + .as_ref() + .map(|expected| expected.namespace.as_str()), + ) + .bind( + delete + .expected_upstream_metadata_namespace + .as_ref() + .and_then(|expected| expected.expected_value.as_ref()), + ) .execute(&self.pool) .await .map_postgres_err()? @@ -2663,6 +2914,13 @@ impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository { Self::update_key(self, key).await } + async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + Self::compare_and_update_key_admin_state(self, update).await + } + async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -3407,6 +3665,7 @@ mod tests { fn runtime_metadata_cas_compares_only_the_requested_namespace() { let sql = super::KEY_RUNTIME_METADATA_CAS_SQL.to_ascii_lowercase(); assert!(sql.contains("upstream_metadata, '{}'::jsonb) -> $2")); + assert!(sql.contains("jsonb_typeof(coalesce(upstream_metadata, '{}'::jsonb)) = 'object'")); assert!(sql.contains("is not distinct from $6::jsonb")); assert!(sql.contains("status_snapshot::jsonb")); assert!(!sql.contains("is_active")); @@ -3436,5 +3695,26 @@ mod tests { } assert!(sql.contains("is_active = $25")); assert!(sql.contains("rpm_limit = $12")); + assert!(sql.contains("api_key is not distinct from $5")); + assert!(sql.contains("auth_config is not distinct from $6")); + } + + #[test] + fn admin_credential_cas_has_atomic_rotation_guards() { + let source = include_str!("provider_catalog.rs"); + for predicate in [ + "api_key IS NOT DISTINCT FROM ", + "auth_config IS NOT DISTINCT FROM ", + "jsonb_typeof(COALESCE(upstream_metadata, '{}'::jsonb)) = 'object'", + "jsonb_typeof(COALESCE(status_snapshot::jsonb, '{}'::jsonb)) = 'object'", + "jsonb_set(COALESCE(upstream_metadata, '{}'::jsonb), '{codex}'", + "jsonb_set(COALESCE(status_snapshot::jsonb, '{}'::jsonb), '{quota}', 'null'::jsonb, true)::json", + "oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0", + ] { + assert!( + source.contains(predicate), + "missing admin CAS guard: {predicate}" + ); + } } } diff --git a/crates/aether-data/adapters/sqlite/src/provider_catalog.rs b/crates/aether-data/adapters/sqlite/src/provider_catalog.rs index 4f730875e..567f28362 100644 --- a/crates/aether-data/adapters/sqlite/src/provider_catalog.rs +++ b/crates/aether-data/adapters/sqlite/src/provider_catalog.rs @@ -8,8 +8,8 @@ use sqlx::{ }; use aether_data_contracts::repository::provider_catalog::{ - ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate, - ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate, + ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate, @@ -1117,6 +1117,85 @@ WHERE id = ? self.reload_key(&key.id, "updated").await } + pub async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + validate_admin_key_cas_update(update)?; + let key = &update.key; + let updated_at = key.updated_at_unix_secs.unwrap_or_else(current_unix_secs) as i64; + let rotation_json = update + .codex_rotation + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys Codex rotation is not serializable: {err}" + )) + })?; + let mut builder = + QueryBuilder::::new("UPDATE provider_api_keys SET provider_id = "); + push_admin_key_assignments(&mut builder, key, updated_at)?; + if update.reset_oauth_runtime { + builder.push(", oauth_invalid_at = NULL, oauth_invalid_reason = NULL, error_count = 0"); + } + if let Some(rotation_json) = rotation_json.as_deref() { + builder + .push(", upstream_metadata = json_set(COALESCE(upstream_metadata, '{}'), '$.codex', json(") + .push_bind(rotation_json) + .push("))"); + } + if update.codex_rotation.is_some() || update.reset_oauth_runtime { + builder.push(", status_snapshot = "); + match (update.codex_rotation.is_some(), update.reset_oauth_runtime) { + (true, true) => builder.push("json_set(COALESCE(status_snapshot, '{}'), '$.quota', json('null'), '$.oauth', json('null'))"), + (true, false) => builder.push("json_set(COALESCE(status_snapshot, '{}'), '$.quota', json('null'))"), + (false, true) => builder.push("json_set(COALESCE(status_snapshot, '{}'), '$.oauth', json('null'))"), + (false, false) => unreachable!(), + }; + } + builder + .push(" WHERE id = ") + .push_bind(&key.id) + .push(" AND api_key IS ") + .push_bind(update.expected_credential.encrypted_api_key.as_deref()) + .push(" AND auth_config IS ") + .push_bind(update.expected_encrypted_auth_config.as_deref()) + .push(" AND auth_type = ") + .push_bind(&update.expected_credential.auth_type) + .push(" AND provider_id = ") + .push_bind(&update.expected_credential.provider_id) + .push( + " AND EXISTS (SELECT 1 FROM providers WHERE providers.id = provider_api_keys.provider_id AND providers.provider_type = ", + ) + .push_bind(&update.expected_credential.provider_type) + .push(")"); + if update.codex_rotation.is_some() { + builder + .push(" AND CASE WHEN upstream_metadata IS NULL THEN 1 WHEN json_valid(upstream_metadata) THEN json_type(upstream_metadata) = 'object' ELSE 0 END") + .push(" AND NOT (api_key IS ") + .push_bind(key.encrypted_api_key.as_deref()) + .push(" AND auth_config IS ") + .push_bind(key.encrypted_auth_config.as_deref()) + .push(" AND auth_type = ") + .push_bind(&key.auth_type) + .push(" AND provider_id = ") + .push_bind(&key.provider_id) + .push(")"); + } + if update.codex_rotation.is_some() || update.reset_oauth_runtime { + builder.push(" AND CASE WHEN status_snapshot IS NULL THEN 1 WHEN json_valid(status_snapshot) THEN json_type(status_snapshot) = 'object' ELSE 0 END"); + } + let rows_affected = builder + .build() + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); + Ok(rows_affected > 0) + } + pub async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -1188,37 +1267,73 @@ WHERE id = ? || expected.auth_type.trim().is_empty() || expected.provider_id.trim().is_empty() || expected.provider_type.trim().is_empty() + || delete + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) { return Err(DataLayerError::InvalidInput( "provider catalog OAuth credential CAS delete contains empty fields".to_string(), )); } - let rows_affected = sqlx::query( - r#" -DELETE FROM provider_api_keys -WHERE id = ? - AND auth_config IS ? - AND api_key IS ? - AND auth_type = ? - AND provider_id = ? - AND EXISTS ( - SELECT 1 - FROM providers - WHERE providers.id = provider_api_keys.provider_id - AND providers.provider_type = ? - ) -"#, - ) - .bind(&delete.key_id) - .bind(delete.expected_encrypted_auth_config.as_deref()) - .bind(expected.encrypted_api_key.as_deref()) - .bind(&expected.auth_type) - .bind(&expected.provider_id) - .bind(&expected.provider_type) - .execute(&self.pool) - .await - .map_sql_err()? - .rows_affected(); + let mut builder = QueryBuilder::::new("DELETE FROM provider_api_keys WHERE id = "); + builder + .push_bind(&delete.key_id) + .push(" AND auth_config IS ") + .push_bind(delete.expected_encrypted_auth_config.as_deref()) + .push(" AND api_key IS ") + .push_bind(expected.encrypted_api_key.as_deref()) + .push(" AND auth_type = ") + .push_bind(&expected.auth_type) + .push(" AND provider_id = ") + .push_bind(&expected.provider_id) + .push( + " AND EXISTS (SELECT 1 FROM providers WHERE \ + providers.id = provider_api_keys.provider_id \ + AND providers.provider_type = ", + ) + .push_bind(&expected.provider_type) + .push(")"); + if let Some(expected) = delete.expected_upstream_metadata_namespace.as_ref() { + let namespace_path = format!( + "$.{}", + serde_json::to_string(&expected.namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + builder + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'object'"); + if let Some(expected_value) = expected.expected_value.as_ref() { + let expected_value = serde_json::to_string(expected_value).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata expected value is not serializable: {err}" + )) + })?; + builder + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path.clone()) + .push(") = json_type(json(") + .push_bind(expected_value.clone()) + .push(")) AND json_extract(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path) + .push(") IS json_extract(json_object('value', json(") + .push_bind(expected_value) + .push(")), '$.value')"); + } else { + builder + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path) + .push(") IS NULL"); + } + } + let rows_affected = builder + .build() + .execute(&self.pool) + .await + .map_sql_err()? + .rows_affected(); Ok(rows_affected > 0) } @@ -1522,6 +1637,25 @@ WHERE id = ? "provider catalog OAuth credential fence must not contain empty fields".to_string(), )); } + if update + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth runtime metadata namespace must not be empty".to_string(), + )); + } + if update + .upstream_metadata_namespace_to_remove + .as_deref() + .is_some_and(|namespace| namespace.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth runtime metadata namespace to remove must not be empty" + .to_string(), + )); + } if !update.status_snapshot_patch.is_object() { return Err(DataLayerError::InvalidInput( "provider catalog status snapshot patch must be an object".to_string(), @@ -1536,6 +1670,22 @@ WHERE id = ? "provider catalog upstream metadata patch must be an object".to_string(), )); } + if update + .upstream_metadata_namespace_to_remove + .as_ref() + .is_some_and(|namespace| { + update + .upstream_metadata_patch + .as_ref() + .and_then(serde_json::Value::as_object) + .is_some_and(|patch| patch.contains_key(namespace)) + }) + { + return Err(DataLayerError::InvalidInput( + "provider catalog OAuth runtime metadata namespace cannot be patched and removed in the same update" + .to_string(), + )); + } let mut builder = QueryBuilder::::new("UPDATE provider_api_keys SET oauth_invalid_at = "); builder @@ -1558,9 +1708,29 @@ WHERE id = ? "provider_api_keys.expires_at", )?); } - if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() { + if update.upstream_metadata_patch.is_some() + || update.upstream_metadata_namespace_to_remove.is_some() + { builder.push(", upstream_metadata = "); - push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?; + if update.upstream_metadata_namespace_to_remove.is_some() { + builder.push("json_remove("); + } + if let Some(metadata_patch) = update.upstream_metadata_patch.as_ref() { + push_upstream_metadata_shallow_patch(&mut builder, metadata_patch)?; + } else { + builder.push("COALESCE(NULLIF(upstream_metadata, ''), '{}')"); + } + if let Some(namespace) = update.upstream_metadata_namespace_to_remove.as_ref() { + let namespace_path = format!( + "$.{}", + serde_json::to_string(namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + builder.push(", ").push_bind(namespace_path).push(")"); + } } builder.push(", status_snapshot = "); push_status_snapshot_shallow_patch(&mut builder, &update.status_snapshot_patch)?; @@ -1594,6 +1764,45 @@ WHERE id = ? .push_bind(&expected.provider_type) .push(")"); } + if update.expected_upstream_metadata_namespace.is_some() + || update.upstream_metadata_patch.is_some() + || update.upstream_metadata_namespace_to_remove.is_some() + { + builder + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'object'"); + } + if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() { + let namespace_path = format!( + "$.{}", + serde_json::to_string(&expected.namespace).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata namespace is not serializable: {err}" + )) + })? + ); + if let Some(expected_value) = expected.expected_value.as_ref() { + let expected_value = serde_json::to_string(expected_value).map_err(|err| { + DataLayerError::UnexpectedValue(format!( + "provider_api_keys.upstream_metadata expected value is not serializable: {err}" + )) + })?; + builder + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path.clone()) + .push(") = json_type(json(") + .push_bind(expected_value.clone()) + .push(")) AND json_extract(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path) + .push(") IS json_extract(json_object('value', json(") + .push_bind(expected_value) + .push(")), '$.value')"); + } else { + builder + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") + .push_bind(namespace_path) + .push(") IS NULL"); + } + } let rows_affected = builder .build() .execute(&self.pool) @@ -1800,7 +2009,8 @@ WHERE id = ? .unwrap_or_else(current_unix_secs) as i64, ) .push(" WHERE id = ") - .push_bind(&update.key_id); + .push_bind(&update.key_id) + .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}')) = 'object'"); if let Some(expected_metadata_value) = expected_metadata_value { builder .push(" AND json_type(COALESCE(NULLIF(upstream_metadata, ''), '{}'), ") @@ -2092,6 +2302,13 @@ impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository { Self::update_key(self, key).await } + async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + Self::compare_and_update_key_admin_state(self, update).await + } + async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -2593,6 +2810,10 @@ SET model_exclude_patterns = ?, updated_at = ? WHERE id = ? + AND provider_id = ? + AND auth_type = ? + AND api_key IS ? + AND auth_config IS ? "# } @@ -2667,7 +2888,152 @@ fn key_update_query( "provider_api_keys.model_exclude_patterns", )?) .bind(updated_at) - .bind(&key.id)) + .bind(&key.id) + .bind(&key.provider_id) + .bind(&key.auth_type) + .bind(&key.encrypted_api_key) + .bind(&key.encrypted_auth_config)) +} + +fn push_admin_key_assignments<'args>( + builder: &mut QueryBuilder<'args, Sqlite>, + key: &'args StoredProviderCatalogKey, + updated_at: i64, +) -> Result<(), DataLayerError> { + builder + .push_bind(&key.provider_id) + .push(", name = ") + .push_bind(&key.name) + .push(", api_key = ") + .push_bind(&key.encrypted_api_key) + .push(", auth_type = ") + .push_bind(&key.auth_type) + .push(", capabilities = ") + .push_bind(optional_json_to_string( + &key.capabilities, + "provider_api_keys.capabilities", + )?) + .push(", is_active = ") + .push_bind(key.is_active) + .push(", api_formats = ") + .push_bind(optional_json_to_string( + &key.api_formats, + "provider_api_keys.api_formats", + )?) + .push(", auth_type_by_format = ") + .push_bind(optional_json_to_string( + &key.auth_type_by_format, + "provider_api_keys.auth_type_by_format", + )?) + .push(", allow_auth_channel_mismatch_formats = ") + .push_bind(optional_json_to_string( + &key.allow_auth_channel_mismatch_formats, + "provider_api_keys.allow_auth_channel_mismatch_formats", + )?) + .push(", auth_config = ") + .push_bind(&key.encrypted_auth_config) + .push(", note = ") + .push_bind(&key.note) + .push(", internal_priority = ") + .push_bind(key.internal_priority) + .push(", rate_multipliers = ") + .push_bind(optional_json_to_string( + &key.rate_multipliers, + "provider_api_keys.rate_multipliers", + )?) + .push(", global_priority_by_format = ") + .push_bind(optional_json_to_string( + &key.global_priority_by_format, + "provider_api_keys.global_priority_by_format", + )?) + .push(", allowed_models = ") + .push_bind(optional_json_to_string( + &key.allowed_models, + "provider_api_keys.allowed_models", + )?) + .push(", expires_at = ") + .push_bind(optional_i64_from_u64( + key.expires_at_unix_secs, + "provider_api_keys.expires_at", + )?) + .push(", cache_ttl_minutes = ") + .push_bind(key.cache_ttl_minutes) + .push(", max_probe_interval_minutes = ") + .push_bind(key.max_probe_interval_minutes) + .push(", proxy = ") + .push_bind(optional_json_to_string( + &key.proxy, + "provider_api_keys.proxy", + )?) + .push(", fingerprint = ") + .push_bind(optional_json_to_string( + &key.fingerprint, + "provider_api_keys.fingerprint", + )?) + .push(", rpm_limit = ") + .push_bind(optional_i64_from_u32(key.rpm_limit)) + .push(", concurrent_limit = ") + .push_bind(key.concurrent_limit) + .push(", auto_fetch_models = ") + .push_bind(key.auto_fetch_models) + .push(", locked_models = ") + .push_bind(optional_json_to_string( + &key.locked_models, + "provider_api_keys.locked_models", + )?) + .push(", model_include_patterns = ") + .push_bind(optional_json_to_string( + &key.model_include_patterns, + "provider_api_keys.model_include_patterns", + )?) + .push(", model_exclude_patterns = ") + .push_bind(optional_json_to_string( + &key.model_exclude_patterns, + "provider_api_keys.model_exclude_patterns", + )?) + .push(", updated_at = ") + .push_bind(updated_at); + Ok(()) +} + +fn validate_admin_key_cas_update( + update: &ProviderCatalogKeyAdminCasUpdate, +) -> Result<(), DataLayerError> { + validate_key(&update.key)?; + let expected = &update.expected_credential; + if expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + || expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || update + .expected_encrypted_auth_config + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + { + return Err(DataLayerError::InvalidInput( + "provider catalog admin credential fence contains empty fields".to_string(), + )); + } + let Some(rotation) = update.codex_rotation.as_ref() else { + return Ok(()); + }; + let valid_rotation = expected.provider_type.eq_ignore_ascii_case("codex") + && rotation.as_object().is_some_and(|object| { + object.len() == 1 + && object + .get("credential_generation") + .and_then(serde_json::Value::as_str) + .is_some_and(|generation| !generation.trim().is_empty()) + }); + if !valid_rotation { + return Err(DataLayerError::InvalidInput( + "provider catalog Codex rotation must contain only credential_generation".to_string(), + )); + } + Ok(()) } fn optional_json_from_string( @@ -3020,14 +3386,265 @@ mod tests { use crate::run_migrations; use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, - ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, - ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogUpstreamMetadataNamespaceUpdate, - StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, + ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogUpstreamMetadataNamespaceExpectation, + ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint, + StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use serde_json::json; + #[tokio::test] + async fn sqlite_admin_credential_cas_rotates_codex_namespace_atomically() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let repository = SqliteProviderCatalogReadRepository::new(pool); + repository + .create_provider( + &StoredProviderCatalogProvider::new( + "admin-cas-provider".to_string(), + "Admin CAS Provider".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"), + None, + ) + .await + .expect("provider should create"); + + let mut original = StoredProviderCatalogKey::new( + "admin-cas-key".to_string(), + "admin-cas-provider".to_string(), + "Before".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + original.encrypted_api_key = Some("api-v1".to_string()); + original.encrypted_auth_config = Some("auth-v1".to_string()); + original.learned_rpm_limit = Some(17); + original.error_count = Some(9); + original.oauth_invalid_at_unix_secs = Some(123); + original.oauth_invalid_reason = Some("old credential invalid".to_string()); + original.upstream_metadata = Some(json!({ + "codex": {"credential_generation":"generation-v1","used_percent":90}, + "other": {"preserved":true} + })); + original.status_snapshot = Some(json!({ + "quota":{"used_ratio":0.9}, + "oauth":{"status":"valid"} + })); + repository + .create_key(&original) + .await + .expect("key should create"); + + let mut replacement = original.clone(); + replacement.name = "After".to_string(); + replacement.auth_type = "api_key".to_string(); + replacement.encrypted_api_key = Some("api-v2".to_string()); + replacement.encrypted_auth_config = None; + replacement.updated_at_unix_secs = Some(200); + let update = ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-v1".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-v1".to_string()), + auth_type: "oauth".to_string(), + provider_id: "admin-cas-provider".to_string(), + provider_type: "codex".to_string(), + }, + key: replacement, + codex_rotation: Some(json!({"credential_generation":"generation-v2"})), + reset_oauth_runtime: true, + }; + assert!(repository + .compare_and_update_key_admin_state(&update) + .await + .expect("credential rotation should run")); + assert!(!repository + .compare_and_update_key_admin_state(&update) + .await + .expect("stale credential fence should be a CAS miss")); + + let stored = repository + .list_keys_by_ids(&["admin-cas-key".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.name, "After"); + assert_eq!(stored.auth_type, "api_key"); + assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-v2")); + assert_eq!(stored.encrypted_auth_config, None); + assert_eq!(stored.learned_rpm_limit, Some(17)); + assert_eq!(stored.error_count, Some(0)); + assert_eq!(stored.oauth_invalid_at_unix_secs, None); + assert_eq!(stored.oauth_invalid_reason, None); + assert_eq!( + stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.get("codex")), + Some(&json!({"credential_generation":"generation-v2"})) + ); + assert_eq!( + stored + .upstream_metadata + .as_ref() + .and_then(|metadata| metadata.pointer("/other/preserved")), + Some(&json!(true)) + ); + assert_eq!( + stored + .status_snapshot + .as_ref() + .and_then(|status| status.get("quota")), + Some(&serde_json::Value::Null) + ); + assert_eq!( + stored + .status_snapshot + .as_ref() + .and_then(|status| status.get("oauth")), + Some(&serde_json::Value::Null) + ); + } + + #[tokio::test] + async fn sqlite_admin_credential_rotation_rejects_non_object_runtime_roots() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let repository = SqliteProviderCatalogReadRepository::new(pool); + repository + .create_provider( + &StoredProviderCatalogProvider::new( + "admin-invalid-provider".to_string(), + "Admin Invalid Provider".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"), + None, + ) + .await + .expect("provider should create"); + let mut original = StoredProviderCatalogKey::new( + "admin-invalid-key".to_string(), + "admin-invalid-provider".to_string(), + "Before".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + original.encrypted_api_key = Some("api-v1".to_string()); + original.encrypted_auth_config = Some("auth-v1".to_string()); + original.upstream_metadata = Some(json!(null)); + original.status_snapshot = Some(json!({"quota":{"used_ratio":0.9}})); + repository + .create_key(&original) + .await + .expect("key should create"); + let mut replacement = original.clone(); + replacement.encrypted_api_key = Some("api-v2".to_string()); + let update = ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-v1".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-v1".to_string()), + auth_type: "oauth".to_string(), + provider_id: "admin-invalid-provider".to_string(), + provider_type: "codex".to_string(), + }, + key: replacement, + codex_rotation: Some(json!({"credential_generation":"generation-v2"})), + reset_oauth_runtime: true, + }; + assert!(!repository + .compare_and_update_key_admin_state(&update) + .await + .expect("invalid metadata root should be a CAS miss")); + let stored = repository + .list_keys_by_ids(&["admin-invalid-key".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-v1")); + assert_eq!(stored.upstream_metadata, Some(json!(null))); + } + + #[tokio::test] + async fn sqlite_ordinary_admin_update_cannot_replace_credentials() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let repository = SqliteProviderCatalogReadRepository::new(pool); + repository + .create_provider( + &StoredProviderCatalogProvider::new( + "ordinary-cas-provider".to_string(), + "Ordinary CAS Provider".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"), + None, + ) + .await + .expect("provider should create"); + let mut original = StoredProviderCatalogKey::new( + "ordinary-cas-key".to_string(), + "ordinary-cas-provider".to_string(), + "Before".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + original.encrypted_api_key = Some("api-v1".to_string()); + original.encrypted_auth_config = Some("auth-v1".to_string()); + repository + .create_key(&original) + .await + .expect("key should create"); + + let mut forbidden = original.clone(); + forbidden.name = "Must Not Persist".to_string(); + forbidden.encrypted_api_key = Some("api-v2".to_string()); + assert!(repository.update_key(&forbidden).await.is_err()); + + let stored = repository + .list_keys_by_ids(&["ordinary-cas-key".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.name, "Before"); + assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-v1")); + assert_eq!(stored.encrypted_auth_config.as_deref(), Some("auth-v1")); + } + #[tokio::test] async fn sqlite_repository_reads_provider_catalog_contract_views() { let pool = sqlx::sqlite::SqlitePoolOptions::new() @@ -3269,6 +3886,128 @@ mod tests { assert_eq!(status["known_boundary"], serde_json::Value::Null); } + #[tokio::test] + async fn sqlite_runtime_metadata_cas_rejects_non_object_metadata_roots() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .expect("sqlite pool should connect"); + run_migrations(&pool) + .await + .expect("sqlite migrations should run"); + let repository = SqliteProviderCatalogReadRepository::new(pool); + repository + .create_provider( + &StoredProviderCatalogProvider::new( + "invalid-root-provider".to_string(), + "Invalid Root Provider".to_string(), + None, + "codex".to_string(), + ) + .expect("provider should build"), + None, + ) + .await + .expect("provider should create"); + + for (index, invalid_root) in [json!(null), json!([]), json!("invalid"), json!(1)] + .into_iter() + .enumerate() + { + let key_id = format!("invalid-root-{index}"); + let mut key = StoredProviderCatalogKey::new( + key_id.clone(), + "invalid-root-provider".to_string(), + format!("Invalid Root {index}"), + "oauth".to_string(), + None, + true, + ) + .expect("key should build"); + key.upstream_metadata = Some(invalid_root.clone()); + key.status_snapshot = Some(json!({"quota":{"remaining":9}})); + key.updated_at_unix_secs = Some(10); + repository + .create_key(&key) + .await + .expect("key should create"); + + assert!(!repository + .update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate { + key_id: key_id.clone(), + namespace: "codex".to_string(), + expected_upstream_metadata_value: None, + upstream_metadata_value: json!({"remaining":1}), + status_snapshot_patch: json!({"quota":{"remaining":1}}), + updated_at_unix_secs: Some(20), + }) + .await + .expect("non-object root should be a CAS miss")); + assert!(!repository + .compare_and_update_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: key_id.clone(), + expected_encrypted_auth_config: None, + expected_credential: None, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: None, + }, + ), + encrypted_auth_config: "next-auth".to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: None, + oauth_invalid_reason: None, + upstream_metadata_patch: Some(json!({"codex":{"remaining":1}})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: json!({"quota":{"remaining":1}}), + reset_error_count: false, + updated_at_unix_secs: Some(20), + }, + ) + .await + .expect("OAuth metadata update should be a CAS miss")); + assert!(!repository + .compare_and_delete_key_oauth_credential( + &ProviderCatalogKeyOAuthCredentialCasDelete { + key_id: key_id.clone(), + expected_encrypted_auth_config: None, + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: None, + auth_type: "oauth".to_string(), + provider_id: "invalid-root-provider".to_string(), + provider_type: "codex".to_string(), + }, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: None, + }, + ), + }, + ) + .await + .expect("OAuth credential delete should be a CAS miss")); + + let stored = repository + .list_keys_by_ids(&[key_id]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.upstream_metadata, Some(invalid_root)); + assert_eq!( + stored.status_snapshot, + Some(json!({"quota":{"remaining":9}})) + ); + assert_eq!(stored.encrypted_auth_config, None); + assert_eq!(stored.updated_at_unix_secs, Some(10)); + } + } + #[tokio::test] async fn sqlite_oauth_runtime_cas_fences_auth_config_and_preserves_admin_fields() { let pool = sqlx::sqlite::SqlitePoolOptions::new() @@ -3322,6 +4061,10 @@ mod tests { "quota": {"remaining": 7}, "admin": {"label": "keep"} })); + key.upstream_metadata = Some(json!({ + "codex": {"remaining": 5}, + "admin": {"keep": true} + })); repository .create_key(&key) .await @@ -3336,12 +4079,19 @@ mod tests { provider_id: "oauth-cas-provider".to_string(), provider_type: "codex".to_string(), }), + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(json!({"remaining": 5})), + }, + ), encrypted_auth_config: "encrypted-auth-v2".to_string(), encrypted_api_key_update: Some("encrypted-api-v2".to_string()), expires_at_unix_secs_update: Some(Some(4_102_555_900)), oauth_invalid_at_unix_secs: None, oauth_invalid_reason: None, upstream_metadata_patch: Some(json!({"codex": {"remaining": 3}})), + upstream_metadata_namespace_to_remove: None, status_snapshot_patch: json!({ "oauth": {"invalid": false, "task_id": "task-v2"}, "runtime": {"generation": 2} @@ -3394,6 +4144,160 @@ mod tests { assert_eq!(status["admin"], json!({"label": "keep"})); assert_eq!(status["runtime"], json!({"generation": 2})); + let stale_metadata_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + expected_encrypted_auth_config: Some("encrypted-auth-v2".to_string()), + expected_credential: None, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(json!({"remaining": 5})), + }, + ), + encrypted_auth_config: "encrypted-auth-v3".to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + upstream_metadata_patch: Some(json!({"codex": {"remaining": 0}})), + status_snapshot_patch: json!({"quota": {"remaining": 0}}), + updated_at_unix_secs: Some(201), + ..update.clone() + }; + assert!(!repository + .compare_and_update_key_oauth_runtime_state(&stale_metadata_update) + .await + .expect("stale metadata namespace should conflict")); + let stored_after_metadata_conflict = repository + .list_keys_by_ids(&[key.id.clone()]) + .await + .expect("key should reload after metadata conflict") + .pop() + .expect("key should exist"); + assert_eq!( + stored_after_metadata_conflict + .encrypted_auth_config + .as_deref(), + Some("encrypted-auth-v2") + ); + assert_eq!( + stored_after_metadata_conflict + .upstream_metadata + .as_ref() + .unwrap()["codex"], + json!({"remaining": 3}) + ); + assert_eq!( + stored_after_metadata_conflict + .status_snapshot + .as_ref() + .unwrap()["quota"], + json!({"remaining": 7}) + ); + + let mut absent_key = StoredProviderCatalogKey::new( + "oauth-cas-absent-key".to_string(), + "oauth-cas-provider".to_string(), + "Absent Namespace".to_string(), + "oauth".to_string(), + None, + true, + ) + .expect("absent namespace key should build") + .with_transport_fields( + None, + Some("encrypted-api-absent".to_string()), + Some("encrypted-auth-absent".to_string()), + None, + None, + None, + None, + None, + None, + ) + .expect("absent namespace key transport should build"); + absent_key.upstream_metadata = Some(json!({"admin": {"keep": true}})); + repository + .create_key(&absent_key) + .await + .expect("absent namespace key should create"); + let absent_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: absent_key.id.clone(), + expected_encrypted_auth_config: Some("encrypted-auth-absent".to_string()), + expected_credential: None, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: None, + }, + ), + encrypted_auth_config: "encrypted-auth-absent".to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: None, + oauth_invalid_reason: None, + upstream_metadata_patch: Some(json!({"codex": {"remaining": 1}})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: json!({}), + reset_error_count: false, + updated_at_unix_secs: Some(202), + }; + assert!(repository + .compare_and_update_key_oauth_runtime_state(&absent_update) + .await + .expect("absent metadata namespace should match")); + let stored_absent = repository + .list_keys_by_ids(&[absent_key.id]) + .await + .expect("absent namespace key should reload") + .pop() + .expect("absent namespace key should exist"); + assert_eq!( + stored_absent.upstream_metadata.as_ref().unwrap()["codex"], + json!({"remaining": 1}) + ); + assert_eq!( + stored_absent.upstream_metadata.as_ref().unwrap()["admin"], + json!({"keep": true}) + ); + + let remove_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(json!({"remaining": 1})), + }, + ), + upstream_metadata_patch: Some(json!({"runtime": {"generation": 2}})), + upstream_metadata_namespace_to_remove: Some("codex".to_string()), + status_snapshot_patch: json!({"quota": null}), + updated_at_unix_secs: Some(203), + ..absent_update + }; + assert!(repository + .compare_and_update_key_oauth_runtime_state(&remove_update) + .await + .expect("matching metadata namespace removal should succeed")); + assert!(!repository + .compare_and_update_key_oauth_runtime_state(&remove_update) + .await + .expect("stale metadata namespace removal should be a CAS miss")); + let stored_after_remove = repository + .list_keys_by_ids(&[stored_absent.id]) + .await + .expect("removed namespace key should reload") + .pop() + .expect("removed namespace key should exist"); + let metadata = stored_after_remove + .upstream_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .expect("metadata should remain an object"); + assert!(!metadata.contains_key("codex")); + assert_eq!(metadata["admin"], json!({"keep": true})); + assert_eq!(metadata["runtime"], json!({"generation": 2})); + assert_eq!( + stored_after_remove.status_snapshot.as_ref().unwrap()["quota"], + serde_json::Value::Null + ); + let stale_api_key_update = ProviderCatalogKeyOAuthRuntimeStateCasUpdate { expected_encrypted_auth_config: Some("encrypted-auth-v2".to_string()), expected_credential: Some(ProviderCatalogKeyOAuthCredentialFence { @@ -3455,13 +4359,14 @@ mod tests { provider_id: "oauth-cas-provider".to_string(), provider_type: "codex".to_string(), }, + expected_upstream_metadata_namespace: None, }; assert!(!repository .compare_and_delete_key_oauth_credential(&stale_delete) .await .expect("stale credential delete should conflict")); - let current_delete = ProviderCatalogKeyOAuthCredentialCasDelete { + let stale_namespace_delete = ProviderCatalogKeyOAuthCredentialCasDelete { key_id: stored_after_stale.id.clone(), expected_encrypted_auth_config: stored_after_stale.encrypted_auth_config.clone(), expected_credential: ProviderCatalogKeyOAuthCredentialFence { @@ -3470,6 +4375,35 @@ mod tests { provider_id: stored_after_stale.provider_id.clone(), provider_type: "codex".to_string(), }, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(json!({"remaining": 3})), + }, + ), + }; + let current_codex = json!({"remaining": 2, "request": "newer"}); + assert!(repository + .upsert_key_upstream_metadata_namespace( + &stored_after_stale.id, + "codex", + ¤t_codex, + Some(202), + ) + .await + .expect("newer codex namespace should persist")); + assert!(!repository + .compare_and_delete_key_oauth_credential(&stale_namespace_delete) + .await + .expect("stale namespace delete should conflict")); + let current_delete = ProviderCatalogKeyOAuthCredentialCasDelete { + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(current_codex), + }, + ), + ..stale_namespace_delete }; assert!(repository .compare_and_delete_key_oauth_credential(¤t_delete) diff --git a/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs b/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs index 8ce89d06b..7d35e5e05 100644 --- a/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data/contracts/src/repository/provider_catalog/mod.rs @@ -4,10 +4,12 @@ mod types; pub use snapshot::ProviderCatalogSnapshot; pub use types::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, + ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, + ProviderCatalogUpstreamMetadataNamespaceExpectation, ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, diff --git a/crates/aether-data/contracts/src/repository/provider_catalog/types.rs b/crates/aether-data/contracts/src/repository/provider_catalog/types.rs index c6a28f932..7039a5878 100644 --- a/crates/aether-data/contracts/src/repository/provider_catalog/types.rs +++ b/crates/aether-data/contracts/src/repository/provider_catalog/types.rs @@ -79,13 +79,50 @@ pub struct ProviderCatalogKeyOAuthCredentialFence { pub provider_type: String, } +/// Administrator-owned key replacement fenced by the exact credential state +/// observed while the edit was prepared. This prevents an older admin request +/// from restoring credentials that a concurrent request already replaced. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ProviderCatalogKeyAdminCasUpdate { + pub expected_encrypted_auth_config: Option, + pub expected_credential: ProviderCatalogKeyOAuthCredentialFence, + /// Full requested key. Repositories merge its administrator-owned fields + /// while preserving the currently stored runtime-owned fields. + pub key: StoredProviderCatalogKey, + /// Optional replacement for the complete Codex metadata namespace. When + /// present it must be an object containing only a non-empty string + /// `credential_generation`; repositories also clear the quota snapshot in + /// the same atomic write. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub codex_rotation: Option, + /// Clear OAuth invalid markers in the same credential replacement write. + /// This must be true whenever new credential material supersedes the old + /// credential, so a post-CAS unfenced cleanup is never required. + #[serde(default)] + pub reset_oauth_runtime: bool, +} + /// Atomic key deletion fenced by the exact OAuth credential generation that -/// produced the terminal failure. +/// produced the terminal failure and, when supplied, one metadata namespace. #[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] pub struct ProviderCatalogKeyOAuthCredentialCasDelete { pub key_id: String, pub expected_encrypted_auth_config: Option, pub expected_credential: ProviderCatalogKeyOAuthCredentialFence, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expected_upstream_metadata_namespace: + Option, +} + +/// Optional single-namespace metadata fence for an OAuth runtime CAS. +/// +/// The outer option on the owning update controls whether the namespace is +/// compared. Within an expectation, `None` requires the namespace to be absent. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct ProviderCatalogUpstreamMetadataNamespaceExpectation { + pub namespace: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expected_value: Option, } /// Agent/runtime-owned OAuth state update fenced by the exact encrypted @@ -98,6 +135,11 @@ pub struct ProviderCatalogKeyOAuthRuntimeStateCasUpdate { pub expected_encrypted_auth_config: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub expected_credential: Option, + /// Optional metadata namespace value compared atomically with the OAuth + /// credential fence. `expected_value: None` means the namespace must be absent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expected_upstream_metadata_namespace: + Option, pub encrypted_auth_config: String, /// Optional access-token ciphertext replacement owned by refresh success. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -110,6 +152,10 @@ pub struct ProviderCatalogKeyOAuthRuntimeStateCasUpdate { /// Top-level runtime metadata namespaces to merge in the same fenced write. #[serde(default, skip_serializing_if = "Option::is_none")] pub upstream_metadata_patch: Option, + /// Optional top-level runtime metadata namespace to remove in the same + /// fenced write. It must not also appear in `upstream_metadata_patch`. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub upstream_metadata_namespace_to_remove: Option, pub status_snapshot_patch: serde_json::Value, #[serde(default)] pub reset_error_count: bool, @@ -822,6 +868,18 @@ pub trait ProviderCatalogWriteRepository: Send + Sync { key: &StoredProviderCatalogKey, ) -> Result; + /// Compare-and-swap administrator-owned key configuration. Credential + /// rotation, Codex namespace replacement, and quota invalidation must be + /// committed atomically with the configuration update. + async fn compare_and_update_key_admin_state( + &self, + _update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + Err(crate::DataLayerError::InvalidConfiguration( + "provider catalog admin CAS updates are not supported by this repository".to_string(), + )) + } + async fn update_keys( &self, keys: &[StoredProviderCatalogKey], diff --git a/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs b/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs index c23a7fb6a..f1d91ef7e 100644 --- a/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs +++ b/crates/aether-data/runtime/src/repository/provider_catalog/memory.rs @@ -7,10 +7,10 @@ use serde_json::{json, Map, Value}; use super::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery, - ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, - ProviderCatalogReadRepository, ProviderCatalogSnapshot, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, + ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot, ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, @@ -538,10 +538,106 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { key.id ))); }; + if !admin_credential_snapshot_matches(stored, key) { + return Err(DataLayerError::UnexpectedValue(format!( + "provider catalog key {} credential state changed", + key.id + ))); + } *stored = merge_admin_key_update(stored, key); Ok(stored.clone()) } + async fn compare_and_update_key_admin_state( + &self, + update: &ProviderCatalogKeyAdminCasUpdate, + ) -> Result { + validate_admin_cas_update(update)?; + + let mut index = self + .index + .write() + .expect("provider catalog repository lock"); + let Some(stored) = index.keys.get(&update.key.id) else { + return Ok(false); + }; + let expected = &update.expected_credential; + let provider_type_matches = index + .providers + .get(&stored.provider_id) + .is_some_and(|provider| provider.provider_type == expected.provider_type); + if stored.encrypted_auth_config != update.expected_encrypted_auth_config + || stored.encrypted_api_key != expected.encrypted_api_key + || stored.auth_type != expected.auth_type + || stored.provider_id != expected.provider_id + || !provider_type_matches + { + return Ok(false); + } + + if (update.codex_rotation.is_some() + && stored + .upstream_metadata + .as_ref() + .is_some_and(|metadata| !metadata.is_object())) + || ((update.codex_rotation.is_some() || update.reset_oauth_runtime) + && stored + .status_snapshot + .as_ref() + .is_some_and(|snapshot| !snapshot.is_object())) + { + return Ok(false); + } + + // A rotation marker is meaningful only when the requested key really + // changes the fenced credential. Rejecting a no-op rotation prevents a + // caller from clearing quota or OAuth state without replacing secrets. + if update.codex_rotation.is_some() + && expected.encrypted_api_key == update.key.encrypted_api_key + && update.expected_encrypted_auth_config == update.key.encrypted_auth_config + && expected.auth_type == update.key.auth_type + && expected.provider_id == update.key.provider_id + { + return Ok(false); + } + + let mut merged = merge_admin_key_update(stored, &update.key); + if let Some(codex_rotation) = update.codex_rotation.as_ref() { + let mut upstream_metadata = stored + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + upstream_metadata.insert("codex".to_string(), codex_rotation.clone()); + merged.upstream_metadata = Some(Value::Object(upstream_metadata)); + + let mut status_snapshot = stored + .status_snapshot + .as_ref() + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + status_snapshot.insert("quota".to_string(), Value::Null); + merged.status_snapshot = Some(Value::Object(status_snapshot)); + } + if update.reset_oauth_runtime { + merged.error_count = Some(0); + merged.oauth_invalid_at_unix_secs = None; + merged.oauth_invalid_reason = None; + if let Some(status_snapshot) = merged + .status_snapshot + .as_mut() + .and_then(Value::as_object_mut) + { + status_snapshot.insert("oauth".to_string(), Value::Null); + } + } + + index.keys.insert(update.key.id.clone(), merged); + Ok(true) + } + async fn update_keys( &self, keys: &[StoredProviderCatalogKey], @@ -551,11 +647,17 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { .write() .expect("provider catalog repository lock"); for key in keys { - if !index.keys.contains_key(&key.id) { + let Some(stored) = index.keys.get(&key.id) else { return Err(DataLayerError::UnexpectedValue(format!( "provider catalog key {} not found", key.id ))); + }; + if !admin_credential_snapshot_matches(stored, key) { + return Err(DataLayerError::UnexpectedValue(format!( + "provider catalog key {} credential state changed", + key.id + ))); } } for key in keys { @@ -716,6 +818,10 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { || expected.auth_type.trim().is_empty() || expected.provider_id.trim().is_empty() || expected.provider_type.trim().is_empty() + || delete + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) { return Err(DataLayerError::InvalidInput( "provider catalog OAuth credential CAS delete contains empty fields".to_string(), @@ -728,6 +834,24 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { let Some(key) = index.keys.get(&delete.key_id) else { return Ok(false); }; + if let Some(expected) = delete.expected_upstream_metadata_namespace.as_ref() { + if key + .upstream_metadata + .as_ref() + .is_some_and(|metadata| !metadata.is_object()) + { + return Ok(false); + } + let current_namespace = key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(&expected.namespace)) + .cloned(); + if current_namespace != expected.expected_value { + return Ok(false); + } + } let provider_type_matches = index .providers .get(&key.provider_id) @@ -829,6 +953,24 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { || expected.provider_id.trim().is_empty() || expected.provider_type.trim().is_empty() }) + || update + .expected_upstream_metadata_namespace + .as_ref() + .is_some_and(|expected| expected.namespace.trim().is_empty()) + || update + .upstream_metadata_namespace_to_remove + .as_deref() + .is_some_and(|namespace| namespace.trim().is_empty()) + || update + .upstream_metadata_namespace_to_remove + .as_ref() + .is_some_and(|namespace| { + update + .upstream_metadata_patch + .as_ref() + .and_then(Value::as_object) + .is_some_and(|patch| patch.contains_key(namespace)) + }) || !update.status_snapshot_patch.is_object() || update .upstream_metadata_patch @@ -856,6 +998,27 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { { return Ok(false); } + if (update.expected_upstream_metadata_namespace.is_some() + || update.upstream_metadata_patch.is_some() + || update.upstream_metadata_namespace_to_remove.is_some()) + && key + .upstream_metadata + .as_ref() + .is_some_and(|metadata| !metadata.is_object()) + { + return Ok(false); + } + if let Some(expected) = update.expected_upstream_metadata_namespace.as_ref() { + let current_namespace = key + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .and_then(|metadata| metadata.get(&expected.namespace)) + .cloned(); + if current_namespace != expected.expected_value { + return Ok(false); + } + } if let Some(expected) = update.expected_credential.as_ref() { let provider_type_matches = index .providers @@ -869,6 +1032,32 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { return Ok(false); } } + let upstream_metadata_update = if update.upstream_metadata_patch.is_some() + || update.upstream_metadata_namespace_to_remove.is_some() + { + let upstream_metadata = json_object_for_merge( + key.upstream_metadata.as_ref(), + "provider catalog upstream metadata", + )?; + let metadata_patch = update + .upstream_metadata_patch + .as_ref() + .and_then(Value::as_object) + .cloned() + .unwrap_or_default(); + let mut upstream_metadata = merge_json_objects(upstream_metadata, metadata_patch); + if let Some(namespace) = update.upstream_metadata_namespace_to_remove.as_ref() { + upstream_metadata.remove(namespace); + } + Some(Value::Object(upstream_metadata)) + } else { + None + }; + let status_snapshot = json_object_for_merge( + key.status_snapshot.as_ref(), + "provider catalog status snapshot", + )?; + let status_snapshot = Value::Object(merge_json_objects(status_snapshot, patch)); let key = index .keys .get_mut(&update.key_id) @@ -885,26 +1074,10 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { if update.reset_error_count { key.error_count = Some(0); } - if let Some(metadata_patch) = update - .upstream_metadata_patch - .as_ref() - .and_then(Value::as_object) - .cloned() - { - let upstream_metadata = json_object_for_merge( - key.upstream_metadata.as_ref(), - "provider catalog upstream metadata", - )?; - key.upstream_metadata = Some(Value::Object(merge_json_objects( - upstream_metadata, - metadata_patch, - ))); + if let Some(upstream_metadata) = upstream_metadata_update { + key.upstream_metadata = Some(upstream_metadata); } - let status_snapshot = json_object_for_merge( - key.status_snapshot.as_ref(), - "provider catalog status snapshot", - )?; - key.status_snapshot = Some(Value::Object(merge_json_objects(status_snapshot, patch))); + key.status_snapshot = Some(status_snapshot); key.updated_at_unix_secs = Some( update .updated_at_unix_secs @@ -1024,6 +1197,13 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository { let Some(key) = index.keys.get_mut(&update.key_id) else { return Ok(false); }; + if key + .upstream_metadata + .as_ref() + .is_some_and(|metadata| !metadata.is_object()) + { + return Ok(false); + } let current_namespace = key .upstream_metadata .as_ref() @@ -1185,6 +1365,55 @@ fn merge_admin_key_update( merged } +fn admin_credential_snapshot_matches( + stored: &StoredProviderCatalogKey, + requested: &StoredProviderCatalogKey, +) -> bool { + stored.provider_id == requested.provider_id + && stored.auth_type == requested.auth_type + && stored.encrypted_api_key == requested.encrypted_api_key + && stored.encrypted_auth_config == requested.encrypted_auth_config +} + +fn validate_admin_cas_update( + update: &ProviderCatalogKeyAdminCasUpdate, +) -> Result<(), DataLayerError> { + let expected = &update.expected_credential; + let invalid_credential = update.key.id.trim().is_empty() + || update.key.provider_id.trim().is_empty() + || update.key.name.trim().is_empty() + || update.key.auth_type.trim().is_empty() + || expected.auth_type.trim().is_empty() + || expected.provider_id.trim().is_empty() + || expected.provider_type.trim().is_empty() + || expected + .encrypted_api_key + .as_deref() + .is_some_and(|value| value.trim().is_empty()) + || update + .expected_encrypted_auth_config + .as_deref() + .is_some_and(|value| value.trim().is_empty()); + let invalid_rotation = update.codex_rotation.as_ref().is_some_and(|rotation| { + let Some(rotation) = rotation.as_object() else { + return true; + }; + !expected.provider_type.eq_ignore_ascii_case("codex") + || rotation.len() != 1 + || rotation + .get("credential_generation") + .and_then(Value::as_str) + .is_none_or(|generation| generation.trim().is_empty()) + }); + if invalid_credential || invalid_rotation { + return Err(DataLayerError::InvalidInput( + "provider catalog admin CAS requires an exact credential fence and a valid Codex rotation marker" + .to_string(), + )); + } + Ok(()) +} + fn adaptive_status_snapshot_patch(patch: &Value) -> Result, DataLayerError> { const OWNED_FIELDS: [&str; 6] = [ "observation_count", @@ -1236,14 +1465,16 @@ mod tests { use super::InMemoryProviderCatalogReadRepository; use crate::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, - ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, - ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, - ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogReadRepository, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, + ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, + ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, + ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceExpectation, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider, }; use crate::repository::usage::ProviderApiKeyUsageDelta; + use crate::DataLayerError; use serde_json::{json, Value}; use std::sync::Arc; use tokio::sync::Barrier; @@ -1421,12 +1652,14 @@ mod tests { provider_id: "provider-1".to_string(), provider_type: "custom".to_string(), }), + expected_upstream_metadata_namespace: None, encrypted_auth_config: "ciphertext-auth-2".to_string(), encrypted_api_key_update: Some("ciphertext-api-2".to_string()), expires_at_unix_secs_update: Some(Some(456)), oauth_invalid_at_unix_secs: None, oauth_invalid_reason: None, upstream_metadata_patch: Some(json!({"codex":{"remaining":3}})), + upstream_metadata_namespace_to_remove: None, status_snapshot_patch: json!({"oauth":{"code":"none"}}), reset_error_count: false, updated_at_unix_secs: Some(123), @@ -1465,6 +1698,172 @@ mod tests { ); } + #[tokio::test] + async fn oauth_runtime_cas_fences_metadata_namespace_value_and_absence() { + let repository_with_metadata = |metadata: Value| { + let mut key = sample_key("key-1", "provider-1") + .with_transport_fields( + None, + "ciphertext-api-1".to_string(), + Some("ciphertext-auth-1".to_string()), + None, + None, + None, + None, + None, + None, + ) + .expect("key transport should build"); + key.upstream_metadata = Some(metadata); + key.status_snapshot = Some(json!({"oauth":{"generation":1}})); + InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider("provider-1")], + vec![], + vec![key], + ) + }; + let update = |expected_value: Option, remaining: u64| { + ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: "key-1".to_string(), + expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: None, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value, + }, + ), + encrypted_auth_config: "ciphertext-auth-1".to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: None, + oauth_invalid_reason: None, + upstream_metadata_patch: Some(json!({"codex":{"remaining":remaining}})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: json!({"oauth":{"generation":remaining}}), + reset_error_count: false, + updated_at_unix_secs: Some(123), + } + }; + + let exact_repository = repository_with_metadata(json!({ + "codex":{"remaining":5}, + "other":{"keep":true} + })); + assert!(exact_repository + .compare_and_update_key_oauth_runtime_state(&update(Some(json!({"remaining":5})), 4,)) + .await + .expect("matching metadata namespace CAS should succeed")); + assert!(!exact_repository + .compare_and_update_key_oauth_runtime_state(&update(Some(json!({"remaining":5})), 0,)) + .await + .expect("stale metadata namespace CAS should conflict")); + let stored = exact_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["codex"], + json!({"remaining":4}) + ); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["other"], + json!({"keep":true}) + ); + assert_eq!( + stored.status_snapshot.as_ref().unwrap()["oauth"], + json!({"generation":4}) + ); + + let absent_repository = repository_with_metadata(json!({"other":{"keep":true}})); + assert!(absent_repository + .compare_and_update_key_oauth_runtime_state(&update(None, 1)) + .await + .expect("absent metadata namespace CAS should succeed")); + let stored = absent_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["codex"], + json!({"remaining":1}) + ); + + let null_repository = repository_with_metadata(json!({"codex":null})); + assert!(!null_repository + .compare_and_update_key_oauth_runtime_state(&update(None, 1)) + .await + .expect("JSON null must not compare equal to an absent namespace")); + let stored = null_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!( + stored.upstream_metadata.as_ref().unwrap()["codex"], + Value::Null + ); + assert_eq!( + stored.status_snapshot.as_ref().unwrap()["oauth"], + json!({"generation":1}) + ); + + let removal_repository = repository_with_metadata(json!({ + "codex":{"remaining":5}, + "admin":{"keep":true} + })); + let removal_update = |expected_remaining: u64| { + let mut update = update( + Some(json!({"remaining":expected_remaining})), + expected_remaining, + ); + update.upstream_metadata_patch = Some(json!({"runtime":{"generation":2}})); + update.upstream_metadata_namespace_to_remove = Some("codex".to_string()); + update.status_snapshot_patch = json!({"quota":null}); + update + }; + assert!(!removal_repository + .compare_and_update_key_oauth_runtime_state(&removal_update(4)) + .await + .expect("stale removal fence should be a CAS miss")); + assert!(removal_repository + .compare_and_update_key_oauth_runtime_state(&removal_update(5)) + .await + .expect("matching removal fence should succeed")); + let stored = removal_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload after namespace removal") + .pop() + .expect("key should exist"); + let metadata = stored + .upstream_metadata + .as_ref() + .and_then(Value::as_object) + .expect("metadata should remain an object"); + assert!(!metadata.contains_key("codex")); + assert_eq!(metadata["admin"], json!({"keep":true})); + assert_eq!(metadata["runtime"], json!({"generation":2})); + assert_eq!( + stored.status_snapshot.as_ref().unwrap()["quota"], + Value::Null + ); + + let mut ambiguous_update = removal_update(5); + ambiguous_update.upstream_metadata_patch = Some(json!({"codex":{"remaining":0}})); + assert!(matches!( + removal_repository + .compare_and_update_key_oauth_runtime_state(&ambiguous_update) + .await, + Err(DataLayerError::InvalidInput(_)) + )); + } + #[tokio::test] async fn oauth_runtime_cas_rejects_changed_credential_context() { let repository = || { @@ -1495,12 +1894,14 @@ mod tests { provider_id: "provider-1".to_string(), provider_type: "custom".to_string(), }), + expected_upstream_metadata_namespace: None, encrypted_auth_config: "ciphertext-auth-2".to_string(), encrypted_api_key_update: Some("ciphertext-api-2".to_string()), expires_at_unix_secs_update: None, oauth_invalid_at_unix_secs: None, oauth_invalid_reason: None, upstream_metadata_patch: None, + upstream_metadata_namespace_to_remove: None, status_snapshot_patch: json!({}), reset_error_count: false, updated_at_unix_secs: Some(123), @@ -1514,10 +1915,21 @@ mod tests { .pop() .expect("key should exist"); key.encrypted_api_key = Some("ciphertext-admin".to_string()); - api_key_repository - .update_key(&key) + assert!(api_key_repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("ciphertext-api-1".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }, + key, + codex_rotation: None, + reset_oauth_runtime: true, + }) .await - .expect("api key replacement should persist"); + .expect("api key replacement should persist")); assert!(!api_key_repository .compare_and_update_key_oauth_runtime_state(&update()) .await @@ -1531,10 +1943,21 @@ mod tests { .pop() .expect("key should exist"); key.auth_type = "oauth".to_string(); - auth_type_repository - .update_key(&key) + assert!(auth_type_repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("ciphertext-api-1".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }, + key, + codex_rotation: None, + reset_oauth_runtime: true, + }) .await - .expect("auth type replacement should persist"); + .expect("auth type replacement should persist")); assert!(!auth_type_repository .compare_and_update_key_oauth_runtime_state(&update()) .await @@ -1586,6 +2009,7 @@ mod tests { provider_id: "provider-1".to_string(), provider_type: "custom".to_string(), }, + expected_upstream_metadata_namespace: None, }; let mut replacement = repository @@ -1597,10 +2021,21 @@ mod tests { replacement.encrypted_api_key = Some("ciphertext-api-2".to_string()); replacement.encrypted_auth_config = Some("ciphertext-auth-2".to_string()); replacement.auth_type = "oauth".to_string(); - repository - .update_key(&replacement) + assert!(repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("ciphertext-api-1".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }, + key: replacement, + codex_rotation: None, + reset_oauth_runtime: true, + }) .await - .expect("replacement should persist"); + .expect("replacement should persist")); assert!(!repository .compare_and_delete_key_oauth_credential(&stale_delete) @@ -1625,6 +2060,7 @@ mod tests { provider_id: "provider-1".to_string(), provider_type: "custom".to_string(), }, + expected_upstream_metadata_namespace: None, }; assert!(repository .compare_and_delete_key_oauth_credential(¤t_delete) @@ -1637,6 +2073,74 @@ mod tests { .is_empty()); } + #[tokio::test] + async fn oauth_credential_cas_delete_rejects_changed_metadata_namespace() { + let mut key = sample_key("key-1", "provider-1") + .with_transport_fields( + None, + "ciphertext-api-1".to_string(), + Some("ciphertext-auth-1".to_string()), + None, + None, + None, + None, + None, + None, + ) + .expect("key transport should build"); + key.upstream_metadata = Some(json!({ + "codex": {"oauth_state_request_started_at_unix_ms": 100}, + "admin": {"keep": true} + })); + let repository = InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider("provider-1")], + vec![], + vec![key], + ); + let delete = |expected_value: Value| ProviderCatalogKeyOAuthCredentialCasDelete { + key_id: "key-1".to_string(), + expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("ciphertext-api-1".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: Some(expected_value), + }, + ), + }; + let stale_delete = delete(json!({"oauth_state_request_started_at_unix_ms": 100})); + let current_codex = json!({"oauth_state_request_started_at_unix_ms": 200}); + assert!(repository + .upsert_key_upstream_metadata_namespace("key-1", "codex", ¤t_codex, Some(200)) + .await + .expect("newer codex namespace should persist")); + + assert!(!repository + .compare_and_delete_key_oauth_credential(&stale_delete) + .await + .expect("stale namespace delete should be a CAS miss")); + assert_eq!( + repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should remain after stale delete")[0] + .upstream_metadata + .as_ref() + .unwrap()["codex"], + current_codex + ); + + assert!(repository + .compare_and_delete_key_oauth_credential(&delete(current_codex)) + .await + .expect("current namespace delete should succeed")); + } + #[tokio::test] async fn materializes_codex_window_usage_stats_delta_in_memory() { let mut key = sample_key("key-1", "provider-1"); @@ -2278,6 +2782,94 @@ mod tests { .expect("matching namespace CAS should succeed")); } + #[tokio::test] + async fn runtime_metadata_namespace_cas_rejects_non_object_metadata_roots() { + for invalid_root in [json!(null), json!([]), json!("invalid"), json!(1)] { + let mut key = sample_key("key-1", "provider-1"); + key.upstream_metadata = Some(invalid_root.clone()); + key.status_snapshot = Some(json!({"quota":{"remaining":9}})); + key.updated_at_unix_secs = Some(10); + let repository = InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider("provider-1")], + vec![], + vec![key], + ); + + assert!(!repository + .update_key_runtime_metadata(&ProviderCatalogKeyRuntimeMetadataUpdate { + key_id: "key-1".to_string(), + namespace: "codex".to_string(), + expected_upstream_metadata_value: None, + upstream_metadata_value: json!({"remaining":1}), + status_snapshot_patch: json!({"quota":{"remaining":1}}), + updated_at_unix_secs: Some(20), + }) + .await + .expect("non-object root should be a CAS miss")); + assert!(!repository + .compare_and_update_key_oauth_runtime_state( + &ProviderCatalogKeyOAuthRuntimeStateCasUpdate { + key_id: "key-1".to_string(), + expected_encrypted_auth_config: None, + expected_credential: None, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: None, + }, + ), + encrypted_auth_config: "next-auth".to_string(), + encrypted_api_key_update: None, + expires_at_unix_secs_update: None, + oauth_invalid_at_unix_secs: None, + oauth_invalid_reason: None, + upstream_metadata_patch: Some(json!({"codex":{"remaining":1}})), + upstream_metadata_namespace_to_remove: None, + status_snapshot_patch: json!({"quota":{"remaining":1}}), + reset_error_count: false, + updated_at_unix_secs: Some(20), + }, + ) + .await + .expect("OAuth metadata update should be a CAS miss")); + assert!(!repository + .compare_and_delete_key_oauth_credential( + &ProviderCatalogKeyOAuthCredentialCasDelete { + key_id: "key-1".to_string(), + expected_encrypted_auth_config: None, + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: None, + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }, + expected_upstream_metadata_namespace: Some( + ProviderCatalogUpstreamMetadataNamespaceExpectation { + namespace: "codex".to_string(), + expected_value: None, + }, + ), + }, + ) + .await + .expect("OAuth credential delete should be a CAS miss")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.upstream_metadata, Some(invalid_root)); + assert_eq!( + stored.status_snapshot, + Some(json!({"quota":{"remaining":9}})) + ); + assert_eq!(stored.encrypted_auth_config, None); + assert_eq!(stored.updated_at_unix_secs, Some(10)); + } + } + #[tokio::test] async fn stale_admin_update_preserves_concurrent_runtime_owned_fields() { let mut key = sample_key("key-1", "provider-1"); @@ -2365,6 +2957,273 @@ mod tests { ); } + #[tokio::test] + async fn admin_cas_rotates_codex_namespace_and_rejects_stale_credentials() { + let mut key = sample_key("key-1", "provider-1") + .with_transport_fields( + None, + "api-old".to_string(), + Some("auth-old".to_string()), + None, + None, + None, + None, + None, + None, + ) + .expect("key transport should build"); + key.auth_type = "oauth".to_string(); + key.upstream_metadata = Some(json!({ + "codex": {"credential_generation":"generation-old","used_percent":80}, + "runtime": {"keep":true} + })); + key.status_snapshot = Some(json!({ + "quota":{"used_ratio":0.8}, + "oauth":{"invalid":true} + })); + key.oauth_invalid_at_unix_secs = Some(123); + key.oauth_invalid_reason = Some("expired".to_string()); + key.error_count = Some(9); + key.health_by_format = Some(json!({"openai:chat":{"failures":2}})); + let mut provider = sample_provider("provider-1"); + provider.provider_type = "codex".to_string(); + let repository = + InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![], vec![key.clone()]); + + let mut requested = key.clone(); + requested.name = "rotated".to_string(); + requested.encrypted_api_key = Some("api-new".to_string()); + requested.encrypted_auth_config = Some("auth-new".to_string()); + requested.upstream_metadata = Some(json!({"caller":"must-not-replace-runtime"})); + requested.status_snapshot = Some(json!({"caller":"must-not-replace-runtime"})); + let update = ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-old".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-old".to_string()), + auth_type: "oauth".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "codex".to_string(), + }, + key: requested, + codex_rotation: Some(json!({"credential_generation":"generation-new"})), + reset_oauth_runtime: true, + }; + assert!(repository + .compare_and_update_key_admin_state(&update) + .await + .expect("matching admin CAS should succeed")); + assert!(!repository + .compare_and_update_key_admin_state(&update) + .await + .expect("stale admin CAS should miss")); + + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.name, "rotated"); + assert_eq!( + stored.upstream_metadata, + Some(json!({ + "codex":{"credential_generation":"generation-new"}, + "runtime":{"keep":true} + })) + ); + assert_eq!( + stored.status_snapshot, + Some(json!({"quota":null,"oauth":null})) + ); + assert_eq!(stored.oauth_invalid_at_unix_secs, None); + assert_eq!(stored.oauth_invalid_reason, None); + assert_eq!(stored.error_count, Some(0)); + assert_eq!( + stored.health_by_format, + Some(json!({"openai:chat":{"failures":2}})) + ); + } + + #[tokio::test] + async fn admin_rotation_fails_closed_for_malformed_runtime_roots() { + for malformed_field in ["upstream_metadata", "status_snapshot"] { + let mut key = sample_key("key-1", "provider-1"); + key.auth_type = "oauth".to_string(); + key.encrypted_api_key = Some("api-old".to_string()); + key.encrypted_auth_config = Some("auth-old".to_string()); + key.upstream_metadata = Some(json!({"runtime":{"keep":true}})); + key.status_snapshot = Some(json!({"oauth":{"keep":true}})); + if malformed_field == "upstream_metadata" { + key.upstream_metadata = Some(json!([])); + } else { + key.status_snapshot = Some(json!(null)); + } + let original = key.clone(); + let mut requested = key.clone(); + requested.name = "must-not-persist".to_string(); + requested.encrypted_api_key = Some("api-new".to_string()); + requested.encrypted_auth_config = Some("auth-new".to_string()); + let mut provider = sample_provider("provider-1"); + provider.provider_type = "codex".to_string(); + let repository = + InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![], vec![key]); + + assert!(!repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-old".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-old".to_string()), + auth_type: "oauth".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "codex".to_string(), + }, + key: requested, + codex_rotation: Some(json!({ + "credential_generation":"generation-new" + })), + reset_oauth_runtime: true, + }) + .await + .expect("malformed runtime root should be a CAS miss")); + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored, original); + } + } + + #[tokio::test] + async fn ordinary_admin_update_rejects_stale_credential_snapshot() { + let mut key = sample_key("key-1", "provider-1"); + key.encrypted_api_key = Some("api-old".to_string()); + key.encrypted_auth_config = Some("auth-old".to_string()); + let mut stale = key.clone(); + stale.name = "stale-name".to_string(); + let repository = InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider("provider-1")], + vec![], + vec![key], + ); + + let mut replacement = stale.clone(); + replacement.encrypted_api_key = Some("api-new".to_string()); + replacement.encrypted_auth_config = Some("auth-new".to_string()); + assert!(repository + .compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-old".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-old".to_string()), + auth_type: "api_key".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "custom".to_string(), + }, + key: replacement, + codex_rotation: None, + reset_oauth_runtime: true, + }) + .await + .expect("credential replacement should succeed")); + + assert!(matches!( + repository.update_key(&stale).await, + Err(DataLayerError::UnexpectedValue(_)) + )); + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.name, "stale-name"); + assert_eq!(stored.encrypted_api_key.as_deref(), Some("api-new")); + assert_eq!(stored.encrypted_auth_config.as_deref(), Some("auth-new")); + } + + #[tokio::test] + async fn admin_rotation_rejects_noop_credential_and_reset_rejects_bad_status_root() { + let mut key = sample_key("key-1", "provider-1"); + key.auth_type = "oauth".to_string(); + key.encrypted_api_key = Some("api-old".to_string()); + key.encrypted_auth_config = Some("auth-old".to_string()); + key.error_count = Some(7); + key.upstream_metadata = Some(json!({"codex":{"credential_generation":"old"}})); + key.status_snapshot = Some(json!({"quota":{"used_ratio":0.4}})); + let mut provider = sample_provider("provider-1"); + provider.provider_type = "codex".to_string(); + let repository = + InMemoryProviderCatalogReadRepository::seed(vec![provider], vec![], vec![key.clone()]); + let no_op = ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-old".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-old".to_string()), + auth_type: "oauth".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "codex".to_string(), + }, + key: key.clone(), + codex_rotation: Some(json!({"credential_generation":"new"})), + reset_oauth_runtime: true, + }; + assert!(!repository + .compare_and_update_key_admin_state(&no_op) + .await + .expect("no-op rotation should be a CAS miss")); + + let reset_valid = ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-old".to_string()), + expected_credential: no_op.expected_credential.clone(), + key: key.clone(), + codex_rotation: None, + reset_oauth_runtime: true, + }; + assert!(repository + .compare_and_update_key_admin_state(&reset_valid) + .await + .expect("valid OAuth runtime reset should succeed")); + let stored = repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.error_count, Some(0)); + + let mut malformed = key.clone(); + malformed.status_snapshot = Some(json!("invalid")); + let reset_only = ProviderCatalogKeyAdminCasUpdate { + expected_encrypted_auth_config: Some("auth-old".to_string()), + expected_credential: ProviderCatalogKeyOAuthCredentialFence { + encrypted_api_key: Some("api-old".to_string()), + auth_type: "oauth".to_string(), + provider_id: "provider-1".to_string(), + provider_type: "codex".to_string(), + }, + key: malformed.clone(), + codex_rotation: None, + reset_oauth_runtime: true, + }; + let malformed_repository = InMemoryProviderCatalogReadRepository::seed( + vec![sample_provider("provider-1")], + vec![], + vec![malformed], + ); + assert!(!malformed_repository + .compare_and_update_key_admin_state(&reset_only) + .await + .expect("reset with malformed status root should be a CAS miss")); + + let stored = malformed_repository + .list_keys_by_ids(&["key-1".to_string()]) + .await + .expect("key should reload") + .pop() + .expect("key should exist"); + assert_eq!(stored.error_count, Some(7)); + } + #[tokio::test] async fn updates_endpoint() { let repository = InMemoryProviderCatalogReadRepository::seed( diff --git a/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs b/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs index 2fbfb7fa7..d89d2c3cb 100644 --- a/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs +++ b/crates/aether-data/runtime/src/repository/provider_catalog/mod.rs @@ -3,10 +3,12 @@ mod memory; #[allow(unused_imports)] pub(crate) use aether_data_contracts::repository::provider_catalog::{ ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate, - ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, + ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate, + ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot, + ProviderCatalogUpstreamMetadataNamespaceExpectation, ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage, diff --git a/crates/aether-gateway/execution/src/stream/ndjson.rs b/crates/aether-gateway/execution/src/stream/ndjson.rs index dad3c8be5..990eb2411 100644 --- a/crates/aether-gateway/execution/src/stream/ndjson.rs +++ b/crates/aether-gateway/execution/src/stream/ndjson.rs @@ -28,6 +28,11 @@ mod tests { payload: StreamFramePayload::Headers { status_code: 200, headers: BTreeMap::from([("content-type".into(), "text/event-stream".into())]), + response_observation: Some(aether_contracts::ExecutionResponseObservation { + request_started_at_unix_ms: 100, + response_headers_observed_at_unix_ms: 125, + request_order_id: "0198-order-1".to_string(), + }), }, }; diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index 71e25a644..bc00a0850 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -1475,6 +1475,7 @@ mod tests { candidate_id: plan.candidate_id.clone(), status_code: self.status_code, headers: BTreeMap::new(), + response_observation: None, body: Some(ResponseBody { json_body: Some(self.response_body.clone()), body_bytes_b64: None, @@ -1526,6 +1527,7 @@ mod tests { candidate_id: plan.candidate_id.clone(), status_code, headers: BTreeMap::new(), + response_observation: None, body: Some(ResponseBody { json_body: Some(response_body), body_bytes_b64: None, @@ -1582,6 +1584,7 @@ mod tests { candidate_id: plan.candidate_id.clone(), status_code, headers: BTreeMap::new(), + response_observation: None, body: Some(ResponseBody { json_body: Some(response_body), body_bytes_b64: None, diff --git a/frontend/src/api/endpoints/keys.ts b/frontend/src/api/endpoints/keys.ts index 08006df60..796ca190a 100644 --- a/frontend/src/api/endpoints/keys.ts +++ b/frontend/src/api/endpoints/keys.ts @@ -321,6 +321,7 @@ export async function refreshProviderQuota( export interface ConsumeCodexResetCreditPayload { idempotency_key: string + expected_credential_generation: string | null } export interface ConsumeCodexResetCreditResult { @@ -331,6 +332,8 @@ export interface ConsumeCodexResetCreditResult { | 'already_redeemed' | 'nothing_to_reset' | 'no_credit' + | 'historical_replay' + | 'credential_changed' | 'unknown' | 'error' | string diff --git a/frontend/src/api/endpoints/types/provider.ts b/frontend/src/api/endpoints/types/provider.ts index 351f87be1..1aadc704a 100644 --- a/frontend/src/api/endpoints/types/provider.ts +++ b/frontend/src/api/endpoints/types/provider.ts @@ -323,6 +323,7 @@ export interface EndpointAPIKey { // Codex 上游元数据类型 export interface CodexUpstreamMetadata { + credential_generation?: string updated_at?: number // 更新时间（Unix 时间戳） plan_type?: string // 套餐类型 primary_used_percent?: number // 周限额窗口使用百分比 @@ -348,6 +349,10 @@ export interface CodexUpstreamMetadata { has_credits?: boolean // 是否有积分 credits_balance?: number // 积分余额 reset_credits?: QuotaResetCreditsSnapshot | null // Codex earned rate-limit reset credits + account_quota_reset_reservation?: { + idempotency_key?: string | null + generation?: number | null + } | null } export interface AntigravityModelQuota { diff --git a/frontend/src/features/providers/components/ProviderDetailDrawer.vue b/frontend/src/features/providers/components/ProviderDetailDrawer.vue index c77ccb2bf..d44736fb9 100644 --- a/frontend/src/features/providers/components/ProviderDetailDrawer.vue +++ b/frontend/src/features/providers/components/ProviderDetailDrawer.vue @@ -292,7 +292,7 @@
@@ -303,7 +303,9 @@ :disabled="consumingCodexResetCreditKeyId === key.id" @click="handleConsumeCodexResetCredit(key)" > - {{ consumingCodexResetCreditKeyId === key.id ? legacyT('重置中...') : legacyT('点击以进行重置') }} + {{ consumingCodexResetCreditKeyId === key.id + ? legacyT('重置中...') + : legacyT(hasPendingCodexResetCredit(key) ? '继续确认重置' : '点击以进行重置') }} ).active_idempotency_key + return typeof activeKey === 'string' && activeKey.trim() ? activeKey.trim() : null +} + +function codexResetCreditCredentialChangedFromError(error: unknown): boolean { + if (typeof error !== 'object' || error === null || !('response' in error)) return false + const response = (error as { response?: { data?: unknown } }).response + if (typeof response?.data !== 'object' || response.data === null) return false + return (response.data as Record).outcome === 'credential_changed' +} + async function handleConsumeCodexResetCredit(key: EndpointAPIKey) { if (!canConsumeCodexResetCredit(key)) return + const credentialGeneration = getCodexCredentialGeneration(key) + if (credentialGeneration === undefined) return + const pendingIdempotencyKey = getPendingCodexResetCreditIdempotencyKey(key) const earliest = getVisibleCodexResetCreditItems(key)[0] const detailMessage = earliest ? `\n当前最早过期项：${earliest.displayKey}，${formatCodexResetCreditExpiresAt(earliest.expiresAt)} 过期。` : '' const confirmed = await confirm({ title: legacyT('确认使用 Codex 重置机会'), - message: `${legacyT('将消耗 1 次 Codex 重置机会。操作完成后会重新刷新账号配额状态。')}${detailMessage}`, + message: pendingIdempotencyKey + ? legacyT('将继续确认上次尚未完成的 Codex 重置请求。') + : `${legacyT('将消耗 1 次 Codex 重置机会。操作完成后会重新刷新账号配额状态。')}${detailMessage}`, confirmText: legacyT('确认重置'), cancelText: legacyT('取消'), variant: 'warning', @@ -1759,10 +1785,15 @@ async function handleConsumeCodexResetCredit(key: EndpointAPIKey) { consumingCodexResetCreditKeyId.value = key.id try { - const idempotencyKey = createCodexResetCreditIdempotencyKey() + const idempotencyKey = pendingIdempotencyKey + || readPendingCodexResetCreditIdempotencyKey(key.id, credentialGeneration) + || createCodexResetCreditIdempotencyKey() + rememberPendingCodexResetCreditIdempotencyKey(key.id, idempotencyKey, credentialGeneration) const result = await consumeCodexResetCredit(key.id, { idempotency_key: idempotencyKey, + expected_credential_generation: credentialGeneration, }) + clearPendingCodexResetCreditIdempotencyKeyForOutcome(key.id, result.outcome) applyQuotaResults([{ key_id: result.key_id, status: result.refresh_status === 'success' ? 'success' : result.status, @@ -1781,6 +1812,16 @@ async function handleConsumeCodexResetCredit(key: EndpointAPIKey) { } emit('refresh') } catch (err: unknown) { + const activeIdempotencyKey = codexResetCreditActiveIdempotencyKeyFromError(err) + if (codexResetCreditCredentialChangedFromError(err)) { + clearPendingCodexResetCreditIdempotencyKey(key.id) + } else if (activeIdempotencyKey) { + rememberPendingCodexResetCreditIdempotencyKey( + key.id, + activeIdempotencyKey, + credentialGeneration, + ) + } showError(localizedApiError(err, 'Codex 重置机会使用失败'), legacyT('错误')) await Promise.all([loadProvider(), loadEndpoints()]) emit('refresh') @@ -1909,6 +1950,9 @@ function getCodexQuotaDisplayFromMetadata(metadata: CodexUpstreamMetadata | null if (!metadata) return null const display: CodexUpstreamMetadata = {} + if (metadata.credential_generation?.trim()) { + display.credential_generation = metadata.credential_generation.trim() + } if (metadata.plan_type) display.plan_type = metadata.plan_type const numberFields: (keyof CodexUpstreamMetadata)[] = [ @@ -1938,6 +1982,9 @@ function getCodexQuotaDisplayFromMetadata(metadata: CodexUpstreamMetadata | null numberFields.forEach(field => copyCodexNumberField(display, metadata, field)) if (metadata.has_credits !== undefined) display.has_credits = metadata.has_credits if (metadata.reset_credits) display.reset_credits = metadata.reset_credits + if (metadata.account_quota_reset_reservation) { + display.account_quota_reset_reservation = metadata.account_quota_reset_reservation + } return Object.keys(display).length > 0 ? display : null } @@ -2025,6 +2072,7 @@ function hasCodexQuotaDisplayData(key: EndpointAPIKey): boolean { || codex.spark_primary_used_percent !== undefined || codex.spark_secondary_used_percent !== undefined || codexDisplayHasResetCredits(codex) + || hasPendingCodexResetCredit(key) ) } @@ -2052,10 +2100,29 @@ function getVisibleCodexResetCreditItems(key: EndpointAPIKey) { return getVisibleCodexResetCreditItemsFromSnapshot(getCodexResetCreditsDisplay(key)) } +function getPendingCodexResetCreditIdempotencyKey(key: EndpointAPIKey): string | null { + const serverReservation = getCodexResetCreditReservationIdempotencyKey(getCodexQuotaDisplay(key)) + if (serverReservation) return serverReservation + const credentialGeneration = getCodexCredentialGeneration(key) + return credentialGeneration === undefined + ? null + : readPendingCodexResetCreditIdempotencyKey(key.id, credentialGeneration) +} + +function getCodexCredentialGeneration(key: EndpointAPIKey): string | null | undefined { + const codex = key.upstream_metadata?.codex + if (!codex || typeof codex !== 'object') return undefined + return codex.credential_generation?.trim() || null +} + +function hasPendingCodexResetCredit(key: EndpointAPIKey): boolean { + return getPendingCodexResetCreditIdempotencyKey(key) !== null +} + function canConsumeCodexResetCredit(key: EndpointAPIKey): boolean { return provider.value?.provider_type === 'codex' - && getCodexResetCreditAvailableCount(key) !== null - && (getCodexResetCreditAvailableCount(key) ?? 0) > 0 + && getCodexCredentialGeneration(key) !== undefined + && (hasPendingCodexResetCredit(key) || (getCodexResetCreditAvailableCount(key) ?? 0) > 0) && !consumingCodexResetCreditKeyId.value } diff --git a/frontend/src/features/providers/components/__tests__/codex-reset-credit-display.spec.ts b/frontend/src/features/providers/components/__tests__/codex-reset-credit-display.spec.ts index 7f33318f2..a26ab0a13 100644 --- a/frontend/src/features/providers/components/__tests__/codex-reset-credit-display.spec.ts +++ b/frontend/src/features/providers/components/__tests__/codex-reset-credit-display.spec.ts @@ -1,12 +1,18 @@ import { describe, expect, it } from 'vitest' import { + clearPendingCodexResetCreditIdempotencyKey, + clearPendingCodexResetCreditIdempotencyKeyForOutcome, createCodexResetCreditIdempotencyKey, formatCodexResetCreditCount, formatCodexResetCreditExpiresAt, getCodexResetCreditAvailableCount, + getCodexResetCreditReservationIdempotencyKey, getVisibleCodexResetCreditItems, + isCodexResetCreditTerminalOutcome, mergeCodexQuotaDisplays, + readPendingCodexResetCreditIdempotencyKey, + rememberPendingCodexResetCreditIdempotencyKey, } from '@/features/providers/components/codex-reset-credit-display' import type { QuotaResetCreditsSnapshot } from '@/api/endpoints/types' @@ -59,6 +65,19 @@ describe('codex reset credit display helpers', () => { expect(getVisibleCodexResetCreditItems(snapshot, 1_700_000_000)).toEqual([]) }) + it('recovers the active idempotency key from a persisted server reservation', () => { + expect(getCodexResetCreditReservationIdempotencyKey({ + reset_credits: { available_count: 0 }, + account_quota_reset_reservation: { + idempotency_key: ' server-active-attempt ', + generation: 3, + }, + })).toBe('server-active-attempt') + expect(getCodexResetCreditReservationIdempotencyKey({ + account_quota_reset_reservation: { idempotency_key: ' ' }, + })).toBeNull() + }) + it('sorts available detail items by remaining time and labels visible items with short ordinal keys', () => { const snapshot: QuotaResetCreditsSnapshot = { available_count: 7, @@ -139,4 +158,164 @@ describe('codex reset credit display helpers', () => { getRandomValues: array => array, })).toBe('existing-random-uuid') }) + + it('keeps an unresolved reset idempotency key until a terminal response clears it', () => { + const values = new Map() + const storage = { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => values.set(key, value), + removeItem: (key: string) => values.delete(key), + } + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-1', + 'reset-attempt-1', + 'credential-v1', + storage, + ) + expect(readPendingCodexResetCreditIdempotencyKey('key-1', 'credential-v1', storage)) + .toBe('reset-attempt-1') + expect(readPendingCodexResetCreditIdempotencyKey('key-2', 'credential-v1', storage)).toBeNull() + + clearPendingCodexResetCreditIdempotencyKey('key-1', storage) + expect(readPendingCodexResetCreditIdempotencyKey('key-1', 'credential-v1', storage)).toBeNull() + }) + + it('clears pending idempotency keys for compatibility history replays', () => { + const values = new Map() + const storage = { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => values.set(key, value), + removeItem: (key: string) => values.delete(key), + } + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-1', + 'reset-attempt-1', + 'credential-v1', + storage, + ) + expect(isCodexResetCreditTerminalOutcome('historical_replay')).toBe(true) + expect(clearPendingCodexResetCreditIdempotencyKeyForOutcome( + 'key-1', + 'historical_replay', + storage, + )).toBe(true) + expect(readPendingCodexResetCreditIdempotencyKey('key-1', 'credential-v1', storage)).toBeNull() + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-1', + 'reset-attempt-2', + 'credential-v1', + storage, + ) + expect(clearPendingCodexResetCreditIdempotencyKeyForOutcome('key-1', 'unknown', storage)) + .toBe(false) + expect(readPendingCodexResetCreditIdempotencyKey('key-1', 'credential-v1', storage)) + .toBe('reset-attempt-2') + expect(isCodexResetCreditTerminalOutcome('unknown')).toBe(false) + expect(isCodexResetCreditTerminalOutcome('error')).toBe(false) + }) + + it('keeps the active idempotency key in memory when session storage is unavailable', () => { + const unavailableStorage = { + getItem: () => { throw new Error('storage unavailable') }, + setItem: () => { throw new Error('storage unavailable') }, + removeItem: () => { throw new Error('storage unavailable') }, + } + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-without-storage', + 'reset-attempt-original', + 'credential-v1', + unavailableStorage, + ) + expect(readPendingCodexResetCreditIdempotencyKey( + 'key-without-storage', + 'credential-v1', + unavailableStorage, + )) + .toBe('reset-attempt-original') + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-without-storage', + 'reset-attempt-from-conflict', + 'credential-v1', + unavailableStorage, + ) + expect(readPendingCodexResetCreditIdempotencyKey( + 'key-without-storage', + 'credential-v1', + unavailableStorage, + )) + .toBe('reset-attempt-from-conflict') + + clearPendingCodexResetCreditIdempotencyKey('key-without-storage', unavailableStorage) + expect(readPendingCodexResetCreditIdempotencyKey( + 'key-without-storage', + 'credential-v1', + unavailableStorage, + )) + .toBeNull() + }) + + it('does not resurrect a cleared key when session storage removal fails', () => { + const values = new Map() + const partiallyUnavailableStorage = { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => values.set(key, value), + removeItem: () => { throw new Error('storage removal unavailable') }, + } + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-with-stale-storage', + 'terminal-attempt', + 'credential-v1', + partiallyUnavailableStorage, + ) + clearPendingCodexResetCreditIdempotencyKey( + 'key-with-stale-storage', + partiallyUnavailableStorage, + ) + + expect(readPendingCodexResetCreditIdempotencyKey( + 'key-with-stale-storage', + 'credential-v1', + partiallyUnavailableStorage, + )).toBeNull() + }) + + it('drops a pending attempt when the credential generation changes', () => { + const values = new Map() + const storage = { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => values.set(key, value), + removeItem: (key: string) => values.delete(key), + } + + rememberPendingCodexResetCreditIdempotencyKey( + 'key-rotated', + 'attempt-for-account-a', + 'credential-a', + storage, + ) + expect(readPendingCodexResetCreditIdempotencyKey('key-rotated', 'credential-b', storage)) + .toBeNull() + expect([...values.values()]).toEqual([]) + }) + + it('never replays a legacy v1 pending value without a credential generation', () => { + const values = new Map([ + ['aether:codex-reset-credit-pending:v1:key-legacy', 'legacy-attempt'], + ]) + const storage = { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => values.set(key, value), + removeItem: (key: string) => values.delete(key), + } + + expect(readPendingCodexResetCreditIdempotencyKey('key-legacy', 'credential-current', storage)) + .toBeNull() + expect([...values.values()]).toEqual([]) + }) }) diff --git a/frontend/src/features/providers/components/codex-reset-credit-display.ts b/frontend/src/features/providers/components/codex-reset-credit-display.ts index 44c274072..2b6f4c95c 100644 --- a/frontend/src/features/providers/components/codex-reset-credit-display.ts +++ b/frontend/src/features/providers/components/codex-reset-credit-display.ts @@ -66,6 +66,84 @@ interface CodexResetCreditCrypto { getRandomValues: (array: Uint8Array) => Uint8Array } +interface CodexResetCreditPendingStorage { + getItem(key: string): string | null + setItem(key: string, value: string): void + removeItem(key: string): void +} + +const CODEX_RESET_CREDIT_PENDING_STORAGE_PREFIX = 'aether:codex-reset-credit-pending:v2:' +const CODEX_RESET_CREDIT_LEGACY_PENDING_STORAGE_PREFIX = 'aether:codex-reset-credit-pending:v1:' + +interface CodexResetCreditPendingAttempt { + idempotencyKey: string + credentialGeneration: string | null +} + +const pendingCodexResetCreditAttempts = new Map() + +const CODEX_RESET_CREDIT_TERMINAL_OUTCOMES = new Set([ + 'reset', + 'already_redeemed', + 'nothing_to_reset', + 'no_credit', + 'historical_replay', +]) + +function codexResetCreditPendingStorageKey(keyId: string): string { + return `${CODEX_RESET_CREDIT_PENDING_STORAGE_PREFIX}${keyId.trim()}` +} + +function codexResetCreditLegacyPendingStorageKey(keyId: string): string { + return `${CODEX_RESET_CREDIT_LEGACY_PENDING_STORAGE_PREFIX}${keyId.trim()}` +} + +function normalizeCodexCredentialGeneration(value: string | null | undefined): string | null { + const normalized = value?.trim() + return normalized ? normalized : null +} + +function parseCodexResetCreditPendingAttempt(value: string): CodexResetCreditPendingAttempt | null { + try { + const parsed = JSON.parse(value) as unknown + if (typeof parsed !== 'object' || parsed === null) return null + const record = parsed as Record + const idempotencyKey = typeof record.idempotencyKey === 'string' + ? record.idempotencyKey.trim() + : '' + const generation = record.credentialGeneration + if (!idempotencyKey || idempotencyKey.length > 256) return null + if (generation !== null && typeof generation !== 'string') return null + const credentialGeneration = normalizeCodexCredentialGeneration(generation) + if (typeof generation === 'string' && credentialGeneration === null) return null + return { idempotencyKey, credentialGeneration } + } catch { + return null + } +} + +function resolveCodexResetCreditPendingStorage( + storage: CodexResetCreditPendingStorage | undefined, +): CodexResetCreditPendingStorage | undefined { + if (storage) return storage + try { + return globalThis.sessionStorage + } catch { + return undefined + } +} + +export function isCodexResetCreditTerminalOutcome(outcome: string): boolean { + return CODEX_RESET_CREDIT_TERMINAL_OUTCOMES.has(outcome) +} + +export function getCodexResetCreditReservationIdempotencyKey( + display: CodexUpstreamMetadata | null | undefined, +): string | null { + const value = display?.account_quota_reset_reservation?.idempotency_key?.trim() + return value && value.length <= 256 ? value : null +} + export function createCodexResetCreditIdempotencyKey( cryptoSource: CodexResetCreditCrypto | undefined = globalThis.crypto, ): string { @@ -82,6 +160,91 @@ export function createCodexResetCreditIdempotencyKey( return `${hex.slice(0, 8)}-${hex.slice(8, 12)}-${hex.slice(12, 16)}-${hex.slice(16, 20)}-${hex.slice(20)}` } +export function readPendingCodexResetCreditIdempotencyKey( + keyId: string, + expectedCredentialGeneration: string | null, + storage?: CodexResetCreditPendingStorage, +): string | null { + const storageKey = codexResetCreditPendingStorageKey(keyId) + const expectedGeneration = normalizeCodexCredentialGeneration(expectedCredentialGeneration) + const resolvedStorage = resolveCodexResetCreditPendingStorage(storage) + try { + resolvedStorage?.removeItem(codexResetCreditLegacyPendingStorageKey(keyId)) + } catch { + // Legacy values are never read, so failed cleanup cannot replay them. + } + if (pendingCodexResetCreditAttempts.has(storageKey)) { + const attempt = pendingCodexResetCreditAttempts.get(storageKey) + if (attempt?.credentialGeneration === expectedGeneration) return attempt.idempotencyKey + pendingCodexResetCreditAttempts.set(storageKey, null) + try { + resolvedStorage?.removeItem(storageKey) + } catch { + // The in-memory tombstone still prevents a stale generation from replaying. + } + return null + } + try { + const value = resolvedStorage?.getItem(storageKey)?.trim() + const attempt = value ? parseCodexResetCreditPendingAttempt(value) : null + if (attempt?.credentialGeneration === expectedGeneration) { + pendingCodexResetCreditAttempts.set(storageKey, attempt) + return attempt.idempotencyKey + } + pendingCodexResetCreditAttempts.set(storageKey, null) + if (value) resolvedStorage?.removeItem(storageKey) + } catch { + // Fall through to the in-memory copy when browser storage is unavailable. + } + return null +} + +export function rememberPendingCodexResetCreditIdempotencyKey( + keyId: string, + idempotencyKey: string, + credentialGeneration: string | null, + storage?: CodexResetCreditPendingStorage, +): void { + const normalized = idempotencyKey.trim() + if (!normalized || normalized.length > 256) return + const storageKey = codexResetCreditPendingStorageKey(keyId) + const attempt: CodexResetCreditPendingAttempt = { + idempotencyKey: normalized, + credentialGeneration: normalizeCodexCredentialGeneration(credentialGeneration), + } + pendingCodexResetCreditAttempts.set(storageKey, attempt) + try { + resolveCodexResetCreditPendingStorage(storage)?.setItem(storageKey, JSON.stringify(attempt)) + } catch { + // The in-memory copy still permits retrying the same request in this page session. + } +} + +export function clearPendingCodexResetCreditIdempotencyKey( + keyId: string, + storage?: CodexResetCreditPendingStorage, +): void { + const storageKey = codexResetCreditPendingStorageKey(keyId) + pendingCodexResetCreditAttempts.set(storageKey, null) + try { + const resolvedStorage = resolveCodexResetCreditPendingStorage(storage) + resolvedStorage?.removeItem(storageKey) + resolvedStorage?.removeItem(codexResetCreditLegacyPendingStorageKey(keyId)) + } catch { + // Ignore unavailable browser storage after a terminal server response. + } +} + +export function clearPendingCodexResetCreditIdempotencyKeyForOutcome( + keyId: string, + outcome: string, + storage?: CodexResetCreditPendingStorage, +): boolean { + if (!isCodexResetCreditTerminalOutcome(outcome)) return false + clearPendingCodexResetCreditIdempotencyKey(keyId, storage) + return true +} + function codexResetCreditRemainingSeconds( item: QuotaResetCreditSnapshot, snapshot: QuotaResetCreditsSnapshot,