From d8902ea612df75a31a0abdb944e3d766892f1c66 Mon Sep 17 00:00:00 2001 From: MMEXA Date: Thu, 30 Jul 2026 03:15:51 +0800 Subject: [PATCH 1/5] fix(gateway): buffer admin pool batch update bodies --- .../src/control/tests/admin_pool.rs | 12 +++++ .../src/handlers/shared/request_utils.rs | 1 + .../src/tests/control/admin/pool.rs | 47 ++++++++++--------- 3 files changed, 39 insertions(+), 21 deletions(-) diff --git a/apps/aether-gateway/src/control/tests/admin_pool.rs b/apps/aether-gateway/src/control/tests/admin_pool.rs index dd177c268..373d818d7 100644 --- a/apps/aether-gateway/src/control/tests/admin_pool.rs +++ b/apps/aether-gateway/src/control/tests/admin_pool.rs @@ -1,6 +1,8 @@ use http::Uri; use super::{classify_control_route, headers}; +use crate::control::GatewayPublicRequestContext; +use crate::handlers::shared::local_proxy_route_requires_buffered_body; #[test] fn classifies_admin_pool_overview_as_admin_proxy_route() { @@ -92,6 +94,16 @@ fn classifies_admin_pool_provider_key_routes_as_admin_proxy_route() { batch_update.route_kind.as_deref(), Some("batch_update_keys") ); + let batch_update_context = GatewayPublicRequestContext::from_request_parts( + "trace-admin-pool-batch-update", + &http::Method::PATCH, + &batch_update_uri, + &headers, + Some(batch_update), + ); + assert!(local_proxy_route_requires_buffered_body( + &batch_update_context + )); let resolve_selection_uri: Uri = "/api/admin/pool/provider-1/keys/resolve-selection" .parse() diff --git a/apps/aether-gateway/src/handlers/shared/request_utils.rs b/apps/aether-gateway/src/handlers/shared/request_utils.rs index 8267b1382..a9a7895fb 100644 --- a/apps/aether-gateway/src/handlers/shared/request_utils.rs +++ b/apps/aether-gateway/src/handlers/shared/request_utils.rs @@ -377,6 +377,7 @@ pub(crate) fn admin_proxy_local_requires_buffered_body( | (Some("users_manage"), http::Method::PATCH, Some("lock_user_api_key")) | (Some("pool_manage"), http::Method::POST, Some("batch_import_keys")) | (Some("pool_manage"), http::Method::POST, Some("batch_action_keys")) + | (Some("pool_manage"), http::Method::PATCH, Some("batch_update_keys")) | (Some("pool_manage"), http::Method::POST, Some("resolve_selection")) | (Some("usage_manage"), http::Method::POST, Some("replay")) | (Some("wallets_manage"), http::Method::POST, Some("adjust_balance")) diff --git a/apps/aether-gateway/src/tests/control/admin/pool.rs b/apps/aether-gateway/src/tests/control/admin/pool.rs index a2c98945a..c201ce0e8 100644 --- a/apps/aether-gateway/src/tests/control/admin/pool.rs +++ b/apps/aether-gateway/src/tests/control/admin/pool.rs @@ -3615,19 +3615,26 @@ async fn gateway_batch_updates_shared_pool_key_configuration() { Vec::new(), vec![first_key, second_key], )); - let state = AppState::new() - .expect("gateway should build") - .with_data_state_for_tests( - GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone( - &provider_catalog_repository, - )), - ); + let gateway = build_router_with_state( + AppState::new() + .expect("gateway should build") + .with_data_state_for_tests( + GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone( + &provider_catalog_repository, + )), + ), + ); + let (gateway_url, gateway_handle) = start_server(gateway).await; - let response = local_admin_pool_response( - &state, - http::Method::PATCH, - "/api/admin/pool/provider-openai/keys/batch-update", - Some(json!({ + let response = reqwest::Client::new() + .patch(format!( + "{gateway_url}/api/admin/pool/provider-openai/keys/batch-update" + )) + .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!({ "key_ids": ["key-openai-b", "key-openai-a", "key-openai-a"], "patch": { "api_formats": ["openai:responses"], @@ -3638,17 +3645,13 @@ async fn gateway_batch_updates_shared_pool_key_configuration() { "locked_models": [], "note": null } - })), - ) - .await; + })) + .send() + .await + .expect("request should succeed"); assert_eq!(response.status(), StatusCode::OK); - let payload: serde_json::Value = serde_json::from_slice( - &to_bytes(response.into_body(), usize::MAX) - .await - .expect("body should read"), - ) - .expect("json body should parse"); + let payload: serde_json::Value = response.json().await.expect("json body should parse"); assert_eq!(payload["affected"], json!(2)); assert_eq!(payload["model_sync"], serde_json::Value::Null); @@ -3670,6 +3673,8 @@ async fn gateway_batch_updates_shared_pool_key_configuration() { assert_eq!(key.locked_models, None); assert_eq!(key.note, None); } + + gateway_handle.abort(); } #[tokio::test] From e55793c76594625ec67f01ab400cdf4de3c20d53 Mon Sep 17 00:00:00 2001 From: MMEXA Date: Thu, 30 Jul 2026 04:09:33 +0800 Subject: [PATCH 2/5] fix(ci): satisfy gateway clippy on upstream baseline --- .../src/execution_runtime/stream/execution_failures.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs index 117bdb71d..17af0c8e2 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs @@ -148,7 +148,7 @@ pub(super) fn build_stream_failure_from_execution_error( error: &ExecutionError, ) -> StreamFailureReport { let transport_error = execution_error_is_transport(error); - let status_code = error.upstream_status.unwrap_or_else(|| { + let status_code = error.upstream_status.unwrap_or({ if matches!( error.kind, ExecutionErrorKind::ConnectTimeout From 8cf9af79da3be1ff93f47f31552f02bf37f3e56d Mon Sep 17 00:00:00 2001 From: MMEXA Date: Thu, 30 Jul 2026 05:45:38 +0800 Subject: [PATCH 3/5] fix(ci): remove redundant usage policy update --- crates/aether-usage/runtime/src/runtime.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/aether-usage/runtime/src/runtime.rs b/crates/aether-usage/runtime/src/runtime.rs index 60d59217a..926ea1185 100644 --- a/crates/aether-usage/runtime/src/runtime.rs +++ b/crates/aether-usage/runtime/src/runtime.rs @@ -12692,7 +12692,6 @@ mod tests { apply_usage_body_capture_policy_to_event( UsageBodyCapturePolicy { record_level: UsageRequestRecordLevel::Basic, - ..UsageBodyCapturePolicy::default() }, &mut event, ); From d7d8db45bab8908cbb65f7720cdb6924d8ad17d6 Mon Sep 17 00:00:00 2001 From: MMEXA Date: Thu, 30 Jul 2026 06:44:59 +0800 Subject: [PATCH 4/5] test(gateway): align tunnel error fixture with failover policy --- .../src/execution_runtime/stream/execution.rs | 41 ++++++++++++------- 1 file changed, 27 insertions(+), 14 deletions(-) diff --git a/apps/aether-gateway/src/execution_runtime/stream/execution.rs b/apps/aether-gateway/src/execution_runtime/stream/execution.rs index d7fd67e77..811ec68ea 100644 --- a/apps/aether-gateway/src/execution_runtime/stream/execution.rs +++ b/apps/aether-gateway/src/execution_runtime/stream/execution.rs @@ -13980,20 +13980,6 @@ mod tests { #[tokio::test] async fn execute_execution_runtime_stream_returns_client_error_with_local_tunnel_message_before_first_data( ) { - let state = AppState::new().expect("app state should build"); - let tunnel_app = state.tunnel.app_state(); - let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8); - let (proxy_close_tx, _) = watch::channel(false); - tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new( - 901, - "node-1".to_string(), - "Node 1".to_string(), - proxy_tx, - proxy_close_tx, - 16, - 2, - ))); - let plan = ExecutionPlan { request_id: "req-client-stream-error-1".into(), candidate_id: Some("cand-client-stream-error-1".into()), @@ -14019,6 +14005,33 @@ mod tests { ..ExecutionTimeouts::default() }), }; + let provider_catalog = provider_catalog_for_plan( + &plan, + Some(json!({ + "failover_rules": { + "stop_on_transport_errors": true, + }, + })), + ); + let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests( + Arc::new(provider_catalog), + "development-key", + ); + let state = AppState::new() + .expect("app state should build") + .with_data_state_for_tests(data_state); + let tunnel_app = state.tunnel.app_state(); + let (proxy_tx, mut proxy_rx) = aether_runtime::bounded_queue(8); + let (proxy_close_tx, _) = watch::channel(false); + tunnel_app.hub.register_proxy(Arc::new(TunnelProxyConn::new( + 901, + "node-1".to_string(), + "Node 1".to_string(), + proxy_tx, + proxy_close_tx, + 16, + 2, + ))); let decision = test_decision(); let state_for_task = state.clone(); From 6c733f75902f819e009d3eac90feb3fb2d5d7984 Mon Sep 17 00:00:00 2001 From: MMEXA Date: Thu, 30 Jul 2026 06:44:59 +0800 Subject: [PATCH 5/5] fix(usage): preserve request diagnostics in event seeds --- crates/aether-usage/runtime/src/write.rs | 21 ++++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/crates/aether-usage/runtime/src/write.rs b/crates/aether-usage/runtime/src/write.rs index 2875dd098..f3a7f2071 100644 --- a/crates/aether-usage/runtime/src/write.rs +++ b/crates/aether-usage/runtime/src/write.rs @@ -2093,6 +2093,11 @@ fn build_runtime_request_metadata_seed_from_parts( if let Some(user_agent) = context_string(context, "user_agent") { metadata.insert("user_agent".to_string(), Value::String(user_agent)); } + for key in ["end_to_end_time_ms", "end_to_end_first_byte_time_ms"] { + if let Some(value) = context_u64(context, key) { + metadata.insert(key.to_string(), Value::from(value)); + } + } if let Some(client_requested_stream) = context_bool(context, "client_requested_stream") { metadata.insert( "client_requested_stream".to_string(), @@ -6710,7 +6715,9 @@ mod tests { }, "original_request_body": { "payload": "x".repeat(MAX_USAGE_CAPTURE_BYTES + 1) - } + }, + "end_to_end_time_ms": 125, + "end_to_end_first_byte_time_ms": 47 })), ); @@ -6727,6 +6734,18 @@ mod tests { "payload": "x".repeat(MAX_USAGE_CAPTURE_BYTES + 1) })) ); + assert_eq!( + data.request_metadata + .as_ref() + .and_then(|metadata| metadata.get("end_to_end_time_ms")), + Some(&json!(125)) + ); + assert_eq!( + data.request_metadata + .as_ref() + .and_then(|metadata| metadata.get("end_to_end_first_byte_time_ms")), + Some(&json!(47)) + ); } #[test]