fix(gateway): buffer admin pool batch update bodies

This commit is contained in:
MMEXA
2026-07-30 05:14:34 +08:00
parent a04673a90d
commit d8902ea612
3 changed files with 39 additions and 21 deletions
@@ -1,6 +1,8 @@
use http::Uri; use http::Uri;
use super::{classify_control_route, headers}; use super::{classify_control_route, headers};
use crate::control::GatewayPublicRequestContext;
use crate::handlers::shared::local_proxy_route_requires_buffered_body;
#[test] #[test]
fn classifies_admin_pool_overview_as_admin_proxy_route() { 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(), batch_update.route_kind.as_deref(),
Some("batch_update_keys") 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" let resolve_selection_uri: Uri = "/api/admin/pool/provider-1/keys/resolve-selection"
.parse() .parse()
@@ -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("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_import_keys"))
| (Some("pool_manage"), http::Method::POST, Some("batch_action_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("pool_manage"), http::Method::POST, Some("resolve_selection"))
| (Some("usage_manage"), http::Method::POST, Some("replay")) | (Some("usage_manage"), http::Method::POST, Some("replay"))
| (Some("wallets_manage"), http::Method::POST, Some("adjust_balance")) | (Some("wallets_manage"), http::Method::POST, Some("adjust_balance"))
@@ -3615,19 +3615,26 @@ async fn gateway_batch_updates_shared_pool_key_configuration() {
Vec::new(), Vec::new(),
vec![first_key, second_key], vec![first_key, second_key],
)); ));
let state = AppState::new() let gateway = build_router_with_state(
.expect("gateway should build") AppState::new()
.with_data_state_for_tests( .expect("gateway should build")
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone( .with_data_state_for_tests(
&provider_catalog_repository, 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( let response = reqwest::Client::new()
&state, .patch(format!(
http::Method::PATCH, "{gateway_url}/api/admin/pool/provider-openai/keys/batch-update"
"/api/admin/pool/provider-openai/keys/batch-update", ))
Some(json!({ .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"], "key_ids": ["key-openai-b", "key-openai-a", "key-openai-a"],
"patch": { "patch": {
"api_formats": ["openai:responses"], "api_formats": ["openai:responses"],
@@ -3638,17 +3645,13 @@ async fn gateway_batch_updates_shared_pool_key_configuration() {
"locked_models": [], "locked_models": [],
"note": null "note": null
} }
})), }))
) .send()
.await; .await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK); assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice( let payload: serde_json::Value = response.json().await.expect("json body should parse");
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
assert_eq!(payload["affected"], json!(2)); assert_eq!(payload["affected"], json!(2));
assert_eq!(payload["model_sync"], serde_json::Value::Null); 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.locked_models, None);
assert_eq!(key.note, None); assert_eq!(key.note, None);
} }
gateway_handle.abort();
} }
#[tokio::test] #[tokio::test]