From 713010fa0a95bbaaa0173ecb1ac7f3e5f2383e03 Mon Sep 17 00:00:00 2001 From: elky Date: Wed, 22 Jul 2026 11:23:53 +0800 Subject: [PATCH] fix(gateway): restore auth role refresh and Rust checks Refresh the resolved user role without bypassing owner group and key policies. Resolve Rust 1.95 Clippy failures and make the pending persistence bound test scheduler-independent. --- apps/aether-gateway/src/cache/auth_context.rs | 25 ++--- apps/aether-gateway/src/data/state/auth.rs | 61 +++++++++--- .../src/executor/candidate_loop.rs | 2 +- .../src/orchestration/report_effects.rs | 6 +- .../adapters/postgres/src/usage/mod.rs | 4 +- .../adapters/postgres/src/usage/tests.rs | 4 +- .../src/bin/gateway_pressure_probe.rs | 5 +- crates/aether-usage/runtime/src/runtime.rs | 94 +++++++++---------- 8 files changed, 109 insertions(+), 92 deletions(-) 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; } })