diff --git a/apps/aether-gateway/src/cache/auth_context.rs b/apps/aether-gateway/src/cache/auth_context.rs index c6e00529c..cc2c3e9dd 100644 --- a/apps/aether-gateway/src/cache/auth_context.rs +++ b/apps/aether-gateway/src/cache/auth_context.rs @@ -175,25 +175,16 @@ impl Drop for AuthContextOwnedInflightGuard { impl AuthContextCache { fn finish_inflight(&self, cache_key: &str, state: &Arc) { - let removed = self - .inflight - .lock() - .map(|mut inflight| { - if inflight - .get(cache_key) - .is_some_and(|current| Arc::ptr_eq(current, state)) - { - state.complete(); - inflight.remove(cache_key).is_some() - } else { - false - } - }) - .unwrap_or(false); - if !removed { + if let Ok(mut inflight) = self.inflight.lock() { // A clear may already have detached and completed this flight. // Keep drop idempotent so an old guard cannot affect its replacement. - return; + if inflight + .get(cache_key) + .is_some_and(|current| Arc::ptr_eq(current, state)) + { + state.complete(); + inflight.remove(cache_key); + } } } diff --git a/apps/aether-gateway/src/data/state/auth.rs b/apps/aether-gateway/src/data/state/auth.rs index 7bc60e737..be7308c0d 100644 --- a/apps/aether-gateway/src/data/state/auth.rs +++ b/apps/aether-gateway/src/data/state/auth.rs @@ -1787,7 +1787,7 @@ impl GatewayDataState { now_unix_secs, ))); }; - let Some(_) = crate::request_diagnostics::observe_db_operation( + let Some(user) = crate::request_diagnostics::observe_db_operation( "auth_user_policy", self.database_pool_summary(), repository.find_user_auth_by_id(&snapshot.user_id), @@ -1799,6 +1799,7 @@ impl GatewayDataState { now_unix_secs, ))); }; + snapshot.user_role = user.role; let groups = self .effective_user_groups_for_user(&snapshot.user_id) .await?; @@ -2455,7 +2456,7 @@ mod tests { } #[tokio::test] - async fn admin_non_standalone_snapshot_bypasses_group_and_key_policies() { + async fn admin_non_standalone_snapshot_applies_group_and_key_policies() { let mut snapshot = sample_snapshot_with_role("key-admin", "admin-1", "admin") .with_user_rate_limit(Some(120)); snapshot.api_key_allowed_providers = Some(vec!["anthropic".to_string()]); @@ -2501,16 +2502,17 @@ mod tests { .expect("snapshot should resolve") .expect("snapshot should exist"); - assert_eq!(resolved.effective_allowed_providers(), None); - assert_eq!(resolved.effective_allowed_api_formats(), None); - assert_eq!(resolved.effective_allowed_models(), None); - assert_eq!(resolved.user_rate_limit, None); - assert_eq!(resolved.api_key_rate_limit, None); - assert_eq!(resolved.api_key_concurrent_limit, None); + assert_eq!(resolved.user_role, "admin"); + assert_eq!(resolved.effective_allowed_providers(), Some(&[][..])); + assert_eq!(resolved.effective_allowed_api_formats(), Some(&[][..])); + assert_eq!(resolved.effective_allowed_models(), Some(&[][..])); + assert_eq!(resolved.user_rate_limit, Some(1)); + assert_eq!(resolved.api_key_rate_limit, Some(5)); + assert_eq!(resolved.api_key_concurrent_limit, Some(1)); } #[tokio::test] - async fn current_admin_role_bypasses_stored_user_and_key_policies() { + async fn current_admin_role_refreshes_without_bypassing_key_policies() { let mut snapshot = sample_snapshot("key-admin", "admin-1"); snapshot.api_key_allowed_providers = Some(vec!["anthropic".to_string()]); snapshot.api_key_allowed_api_formats = Some(vec!["anthropic:messages".to_string()]); @@ -2533,12 +2535,43 @@ mod tests { .expect("snapshot should exist"); assert_eq!(resolved.user_role, "admin"); - assert_eq!(resolved.effective_allowed_providers(), None); - assert_eq!(resolved.effective_allowed_api_formats(), None); - assert_eq!(resolved.effective_allowed_models(), None); + assert_eq!( + resolved.effective_allowed_providers(), + Some(&["anthropic".to_string()][..]) + ); + assert_eq!( + resolved.effective_allowed_api_formats(), + Some(&["anthropic:messages".to_string()][..]) + ); + assert_eq!( + resolved.effective_allowed_models(), + Some(&["claude-sonnet-4-5".to_string()][..]) + ); assert_eq!(resolved.user_rate_limit, None); - assert_eq!(resolved.api_key_rate_limit, None); - assert_eq!(resolved.api_key_concurrent_limit, None); + assert_eq!(resolved.api_key_rate_limit, Some(60)); + assert_eq!(resolved.api_key_concurrent_limit, Some(5)); + } + + #[tokio::test] + async fn current_user_role_replaces_stored_admin_role() { + let snapshot = sample_snapshot_with_role("key-user", "user-1", "admin"); + let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![( + Some("hash-user".to_string()), + snapshot, + )])); + let user_repository = Arc::new(InMemoryUserReadRepository::seed_auth_users(vec![ + sample_auth_user("user-1", "user"), + ])); + let state = GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository) + .with_user_reader(user_repository); + + let resolved = state + .read_auth_api_key_snapshot_by_key_hash("hash-user", 100) + .await + .expect("snapshot should resolve") + .expect("snapshot should exist"); + + assert_eq!(resolved.user_role, "user"); } #[tokio::test] diff --git a/apps/aether-gateway/src/executor/candidate_loop.rs b/apps/aether-gateway/src/executor/candidate_loop.rs index 41a77a4fb..e6b008f74 100644 --- a/apps/aether-gateway/src/executor/candidate_loop.rs +++ b/apps/aether-gateway/src/executor/candidate_loop.rs @@ -474,7 +474,7 @@ where self.state, self.trace_id, self.decision, - &plan, + plan, report_context.as_ref(), ) .await? diff --git a/apps/aether-gateway/src/orchestration/report_effects.rs b/apps/aether-gateway/src/orchestration/report_effects.rs index abb87af6d..9fc72c643 100644 --- a/apps/aether-gateway/src/orchestration/report_effects.rs +++ b/apps/aether-gateway/src/orchestration/report_effects.rs @@ -376,7 +376,7 @@ async fn sync_gemini_cli_credits_from_report( .as_ref() .and_then(Value::as_object) .cloned() - .unwrap_or_else(serde_json::Map::new); + .unwrap_or_default(); gemini_cli_bucket.insert("credits".to_string(), credits.clone()); gemini_cli_bucket.insert("updated_at".to_string(), json!(now_unix_secs)); @@ -889,7 +889,7 @@ async fn sync_codex_quota_from_response_headers( let current_codex = expected_namespace_value .clone() .and_then(|value| value.as_object().cloned()) - .unwrap_or_else(serde_json::Map::new); + .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); @@ -925,7 +925,7 @@ async fn sync_codex_quota_from_response_headers( // Response headers describe an authoritative quota snapshot. A CAS // conflict means a newer snapshot/delta won, so avoid replaying stale // data over it. - return Ok(false); + Ok(false) } #[cfg(test)] diff --git a/crates/aether-data/adapters/postgres/src/usage/mod.rs b/crates/aether-data/adapters/postgres/src/usage/mod.rs index 973d2b1e3..a95db63fe 100644 --- a/crates/aether-data/adapters/postgres/src/usage/mod.rs +++ b/crates/aether-data/adapters/postgres/src/usage/mod.rs @@ -2043,6 +2043,8 @@ struct PreparedFirstByteUsage { updated_at_unix_secs: i64, } +type PartitionedFirstByteUsages = (Vec, Vec<(usize, UpsertUsageRecord)>); + #[derive(Debug)] struct PreparedPendingUsage { usage: UpsertUsageRecord, @@ -2223,7 +2225,7 @@ impl PreparedFirstByteUsage { fn partition_first_byte_usages( usages: Vec, -) -> Result<(Vec, Vec<(usize, UpsertUsageRecord)>), DataLayerError> { +) -> Result { let mut request_id_counts = BTreeMap::::new(); for usage in &usages { *request_id_counts diff --git a/crates/aether-data/adapters/postgres/src/usage/tests.rs b/crates/aether-data/adapters/postgres/src/usage/tests.rs index cfd4d18a7..10d708e87 100644 --- a/crates/aether-data/adapters/postgres/src/usage/tests.rs +++ b/crates/aether-data/adapters/postgres/src/usage/tests.rs @@ -3746,7 +3746,7 @@ fn first_byte_provider_counter_batch_prepares_all_columns_before_query_building( assert_eq!(prepared.candidate_last_used_at_unix_secs, Some(10)); assert_eq!(prepared.removed_last_used_at_unix_secs, Some(11)); assert_eq!(prepared.usage_created_at_unix_secs, Some(12)); - assert!(super::USAGE_COUNTER_DELTA_INSERT_BATCH_SIZE >= 1_024); + const { assert!(super::USAGE_COUNTER_DELTA_INSERT_BATCH_SIZE >= 1_024) }; assert!( super::USAGE_COUNTER_DELTA_INSERT_BATCH_SIZE * super::USAGE_COUNTER_DELTA_INSERT_BINDS_PER_ROW @@ -3816,7 +3816,7 @@ fn first_byte_batch_writes_provider_counter_transitions_in_bulk() { assert!(implementation.contains("prepare_first_byte_provider_contribution_transitions")); assert!(implementation.contains("insert_usage_counter_deltas_batch_in_tx")); assert!(!implementation.contains("enqueue_first_byte_provider_contribution_transition_in_tx")); - assert!(super::USAGE_COUNTER_DELTA_INSERT_BATCH_SIZE > 512 * 2); + const { assert!(super::USAGE_COUNTER_DELTA_INSERT_BATCH_SIZE > 512 * 2) }; } #[test] diff --git a/crates/aether-testing/loadtools/src/bin/gateway_pressure_probe.rs b/crates/aether-testing/loadtools/src/bin/gateway_pressure_probe.rs index a9c56a8f7..086181516 100644 --- a/crates/aether-testing/loadtools/src/bin/gateway_pressure_probe.rs +++ b/crates/aether-testing/loadtools/src/bin/gateway_pressure_probe.rs @@ -2349,8 +2349,9 @@ fn missing_required_settle_drain_metrics(samples: &[PrometheusSample]) -> Vec, ) { let mut tasks = tokio::task::JoinSet::new(); - let batch_concurrency = concurrency - .min(FIRST_BYTE_PERSISTENCE_MAX_BATCH_CONCURRENCY) - .max(1); + let batch_concurrency = concurrency.clamp(1, FIRST_BYTE_PERSISTENCE_MAX_BATCH_CONCURRENCY); let write_admission = Arc::new(tokio::sync::Semaphore::new(concurrency)); while let Some(first) = receiver.recv().await { let mut batch = vec![first]; @@ -7689,11 +7685,12 @@ mod tests { .await .expect("pending lifecycle usage should be persisted"); - let records = store.records.lock().expect("records lock"); - assert_eq!(records.len(), 1); - assert_eq!(records[0].request_id, "req-lifecycle-queue-pending-1"); - assert_eq!(records[0].status, "pending"); - drop(records); + { + let records = store.records.lock().expect("records lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0].request_id, "req-lifecycle-queue-pending-1"); + assert_eq!(records[0].status, "pending"); + } assert!( queue .read_group("usage-test-consumer") @@ -8236,20 +8233,7 @@ mod tests { } }); - timeout(Duration::from_secs(2), async { - while store.max_writes_in_flight.load(Ordering::Acquire) - < super::PENDING_PERSISTENCE_SINGLE_WRITE_TARGET_CONCURRENCY - || dispatcher - .sender - .as_ref() - .is_some_and(|sender| sender.capacity() != 0) - { - tokio::task::yield_now().await; - } - }) - .await - .expect("the bounded dispatcher should fill its fixed worker window"); - timeout(Duration::from_secs(2), dispatch) + timeout(Duration::from_secs(5), dispatch) .await .expect("overflow dispatch must remain non-blocking") .expect("overflow dispatch task should not panic"); @@ -8260,11 +8244,6 @@ mod tests { dispatcher.state.max_pending.load(Ordering::Acquire) <= RESIDENT_BOUND, "resident pending records exceeded the fixed channel and worker bound" ); - assert_eq!( - store.max_writes_in_flight.load(Ordering::Acquire), - super::PENDING_PERSISTENCE_SINGLE_WRITE_TARGET_CONCURRENCY - ); - store.release_writes.add_permits(TOTAL_ITEMS); timeout(Duration::from_secs(5), async { while dispatcher.state.pending.load(Ordering::Acquire) != 0 { @@ -8276,6 +8255,11 @@ mod tests { let writes_completed = store.writes_completed.load(Ordering::Acquire); assert_eq!(completed.load(Ordering::Acquire), writes_completed); assert_eq!(writes_completed as u64 + overflow, TOTAL_ITEMS as u64); + assert!( + store.max_writes_in_flight.load(Ordering::Acquire) + <= super::PENDING_PERSISTENCE_SINGLE_WRITE_TARGET_CONCURRENCY, + "pending persistence exceeded its fixed worker window" + ); } #[tokio::test] @@ -10036,12 +10020,13 @@ mod tests { }) .await .expect("first-byte transition should be persisted directly"); - let records = store.records.lock().expect("records lock"); - assert_eq!(records.len(), 1); - assert_eq!(records[0].request_id, "req-stream-direct-delay"); - assert_eq!(records[0].status, "streaming"); - assert_eq!(records[0].first_byte_time_ms, Some(12)); - drop(records); + { + let records = store.records.lock().expect("records lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0].request_id, "req-stream-direct-delay"); + assert_eq!(records[0].status, "streaming"); + assert_eq!(records[0].first_byte_time_ms, Some(12)); + } sleep(Duration::from_millis(60)).await; let queued = queue @@ -10181,11 +10166,12 @@ mod tests { }) .await .expect("first-byte transition should be persisted directly"); - let records = store.records.lock().expect("records lock"); - assert_eq!(records.len(), 1); - assert_eq!(records[0].status, "streaming"); - assert_eq!(records[0].first_byte_time_ms, Some(12)); - drop(records); + { + let records = store.records.lock().expect("records lock"); + assert_eq!(records.len(), 1); + assert_eq!(records[0].status, "streaming"); + assert_eq!(records[0].first_byte_time_ms, Some(12)); + } sleep(Duration::from_millis(80)).await; let delayed = queue @@ -12063,13 +12049,15 @@ mod tests { ); timeout(Duration::from_secs(1), async { loop { - let records = store.records.lock().expect("records lock"); - if records.iter().any(|record| { - record.request_id == unrelated_request_id && record.status == "pending" - }) { + let unrelated_written = { + let records = store.records.lock().expect("records lock"); + records.iter().any(|record| { + record.request_id == unrelated_request_id && record.status == "pending" + }) + }; + if unrelated_written { break; } - drop(records); tokio::task::yield_now().await; } }) @@ -12152,13 +12140,15 @@ mod tests { ); timeout(Duration::from_secs(1), async { loop { - let records = store.records.lock().expect("records lock"); - if records.iter().any(|record| { - record.request_id == unrelated_request_id && record.status == "pending" - }) { + let unrelated_written = { + let records = store.records.lock().expect("records lock"); + records.iter().any(|record| { + record.request_id == unrelated_request_id && record.status == "pending" + }) + }; + if unrelated_written { break; } - drop(records); tokio::task::yield_now().await; } })