From c06806b1f560a2cb73e0de914063ec7563eb62e7 Mon Sep 17 00:00:00 2001 From: Coffee <95295094+hanakannzashi@users.noreply.github.com> Date: Tue, 1 Sep 2026 13:58:43 +0800 Subject: [PATCH 1/3] fix: gateway-sign alias and legacy completion responses --- crates/api/src/routes/completions.rs | 253 +++++++++++++----- .../tests/e2e_all/model_alias_transparency.rs | 129 ++++++++- .../tests/e2e_all/signature_verification.rs | 158 +++++++++++ 3 files changed, 469 insertions(+), 71 deletions(-) diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index 1a87122f5..a42c66508 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -1519,18 +1519,20 @@ async fn chat_completions_inner( service_request.original_request = None; } - // Auto-redact re-serializes response bytes after restoring the original - // values. A provider signature covers the pre-redaction bytes, so an - // attested response must use a Gateway signature over the bytes returned - // to the client instead. + // Auto-redact and alias handling can change the bytes returned to the + // client. A provider signature covers the upstream request and response, + // so it cannot verify those public bytes. + let alias_requires_gateway_signature = + alias_canonical.is_some() && model_attestation_supported.unwrap_or(false); let gateway_signature_enabled = usage_mode.gateway_signature_enabled - || auto_redact_requires_gateway_signature(auto_redact_enabled, model_attestation_supported); - // Never publish a provider signature over bytes that auto-redact changes. - // If metadata is unavailable, we cannot safely create a Gateway signature - // either, but omitting a signature is still better than returning one that - // cannot verify. + || auto_redact_requires_gateway_signature(auto_redact_enabled, model_attestation_supported) + || alias_requires_gateway_signature; + let public_response_rewritten = auto_redact_enabled || alias_canonical.is_some(); + // If model metadata is unavailable, we cannot safely mint a Gateway + // signature. Do not fall back to a provider signature over different + // public bytes. service_request.skip_provider_chat_signature = - usage_mode.gateway_signature_enabled || auto_redact_enabled; + usage_mode.gateway_signature_enabled || public_response_rewritten; let redaction_map = Arc::new(redaction_map); // Check if streaming is requested @@ -1709,7 +1711,8 @@ async fn chat_completions_inner( // store. let Some(mut chunk) = event.chunk else { if event.is_done_marker() { - if auto_redact_enabled + if gateway_signature_enabled + || auto_redact_enabled || rewrite_public_stream_usage || strip_intermediate_usage { @@ -1748,7 +1751,7 @@ async fn chat_completions_inner( )); } } - if gateway_signature_enabled || auto_redact_enabled { + if gateway_signature_enabled || public_response_rewritten { let mut chat_id = public_signature_chat_id.lock().await; if chat_id.is_none() { *chat_id = Some(chat.id.clone()); @@ -2019,7 +2022,7 @@ async fn chat_completions_inner( .release_chat_signature_pin(&chat_id) .await; } - } else if auto_redact_enabled { + } else if public_response_rewritten { if let Some(chat_id) = public_signature_chat_id_for_chain.lock().await.clone() { @@ -2175,7 +2178,7 @@ async fn chat_completions_inner( _ => body_bytes, }; - if auto_redact_enabled { + if public_response_rewritten { if gateway_signature_enabled { let response_hash = hex::encode(Sha256::digest(&body_bytes)); if let Err(error) = app_state @@ -2190,7 +2193,7 @@ async fn chat_completions_inner( tracing::error!( chat_id = %response_with_bytes.response.id, error = %error, - "Failed to store auto-redacted chat completion signature" + "Failed to store public chat completion signature" ); } } else { @@ -2341,6 +2344,8 @@ async fn completions_inner( request: CompletionRequest, request_id: Uuid, ) -> axum::response::Response { + let request_hash = body_hash.hash.clone(); + // Reject E2E encryption: validate for parity (an invalid version still 400s // the same way chat does), then refuse if any encryption header is present. let encryption_headers = match crate::routes::common::validate_encryption_headers(&headers) { @@ -2430,7 +2435,24 @@ async fn completions_inner( .resolve_alias_cached(&request.model) .await; - let service_request = convert_text_request_to_service( + let resolved_model_name = alias_canonical.as_deref().unwrap_or(&request.model); + let legacy_gateway_signature_enabled = + match app_state.models_service.get_models_with_pricing().await { + Ok(models) => models + .iter() + .find(|model| model.model_name == resolved_model_name) + .is_some_and(|model| model.attestation_supported), + Err(error) => { + tracing::warn!( + model = %request.model, + error = %error, + "Failed to read cached model metadata for legacy completion signing" + ); + false + } + }; + + let mut service_request = convert_text_request_to_service( &request, prompt, api_key.api_key.created_by_user_id.0, @@ -2441,6 +2463,10 @@ async fn completions_inner( body_hash, request_id, ); + // This endpoint always converts the provider's chat-completion payload + // into the legacy completion format, so a provider signature cannot + // verify the bytes returned to the client. + service_request.skip_provider_chat_signature = true; if request.stream == Some(true) { match app_state @@ -2501,6 +2527,11 @@ async fn completions_inner( let organization_id = api_key.organization.id.0; let model_for_err = request.model.clone(); + let public_signature_hasher = Arc::new(tokio::sync::Mutex::new(Sha256::new())); + let public_signature_chat_id = + Arc::new(tokio::sync::Mutex::new(stream_chat_id.clone())); + let stream_error_count = Arc::new(std::sync::atomic::AtomicU32::new(0)); + let attestation_service = app_state.attestation_service.clone(); // Warning to inject into the first streamed chunk of an // alias-served response (issue #573). @@ -2509,61 +2540,138 @@ async fn completions_inner( |canonical| alias_warning_message(&request.model, canonical), ))); let pending_warning = alias_warning_pending.clone(); + let public_signature_hasher_for_chunks = public_signature_hasher.clone(); + let public_signature_chat_id_for_chunks = public_signature_chat_id.clone(); + let stream_error_count_for_chunks = stream_error_count.clone(); let byte_stream = peekable_stream .filter_map(move |result| { let model_for_err = model_for_err.clone(); let pending_warning = pending_warning.clone(); - std::future::ready(match result { - // Control lines (blank/comment/[DONE]) carry no - // parsed payload — skip; the gateway appends its - // own [DONE] terminator below. This route reshapes - // chat chunks into text-completion format, so it - // always re-serializes (no byte passthrough). - Ok(event) => event.chunk.map(|chunk| { - let text_chunk = chat_chunk_to_text_chunk(chunk); - // The first chunk of an alias-served response - // gets a top-level "warning" (issue #573). - let alias_warning = - pending_warning.lock().ok().and_then(|mut g| g.take()); - let json_data = match alias_warning { - Some(warning) => { - serde_json::to_value(&text_chunk).map(|mut v| { - if let Some(obj) = v.as_object_mut() { - obj.insert( - "warning".to_string(), - serde_json::Value::String(warning), - ); - } - v.to_string() - }) + let public_signature_hasher = public_signature_hasher_for_chunks.clone(); + let public_signature_chat_id = public_signature_chat_id_for_chunks.clone(); + let stream_error_count = stream_error_count_for_chunks.clone(); + async move { + match result { + // Control lines (blank/comment/[DONE]) carry no + // parsed payload — skip; the gateway appends its + // own [DONE] terminator below. This route reshapes + // chat chunks into text-completion format, so it + // always re-serializes (no byte passthrough). + Ok(event) => { + let chunk = event.chunk?; + let chat_id = match &chunk { + inference_providers::StreamChunk::Chat(chunk) => { + chunk.id.clone() + } + inference_providers::StreamChunk::Text(chunk) => { + chunk.id.clone() + } + }; + let mut stored_chat_id = public_signature_chat_id.lock().await; + if stored_chat_id.is_none() { + *stored_chat_id = Some(chat_id); } - None => serde_json::to_string(&text_chunk), + drop(stored_chat_id); + + let text_chunk = chat_chunk_to_text_chunk(chunk); + // The first chunk of an alias-served response + // gets a top-level "warning" (issue #573). + let alias_warning = + pending_warning.lock().ok().and_then(|mut g| g.take()); + let json_data = match alias_warning { + Some(warning) => { + serde_json::to_value(&text_chunk).map(|mut v| { + if let Some(obj) = v.as_object_mut() { + obj.insert( + "warning".to_string(), + serde_json::Value::String(warning), + ); + } + v.to_string() + }) + } + None => serde_json::to_string(&text_chunk), + } + .unwrap_or_else(|e| { + tracing::error!( + %organization_id, + "Failed to serialize text completion chunk: {e}" + ); + "{}".to_string() + }); + let bytes = Bytes::from(format!("data: {json_data}\n\n")); + if legacy_gateway_signature_enabled { + public_signature_hasher.lock().await.update(&bytes); + } + Some(Ok::(bytes)) } - .unwrap_or_else(|e| { + Err(e) => { + stream_error_count + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); tracing::error!( %organization_id, - "Failed to serialize text completion chunk: {e}" + model = %model_for_err, + error_type = %completion_stream_error_category(&e), + "Text completion stream error" ); - "{}".to_string() - }); - Ok::(Bytes::from(format!( - "data: {json_data}\n\n" - ))) - }), - Err(e) => { - tracing::error!( - %organization_id, - model = %model_for_err, - error_type = %completion_stream_error_category(&e), - "Text completion stream error" - ); - Some(Ok::(sse_error_frame(&e))) + Some(Ok::(sse_error_frame(&e))) + } } - }) + } }) - .chain(futures::stream::once(async move { - Ok::(Bytes::from_static(b"data: [DONE]\n\n")) + .chain(futures::stream::once({ + let public_signature_hasher = public_signature_hasher.clone(); + let public_signature_chat_id = public_signature_chat_id.clone(); + let stream_error_count = stream_error_count.clone(); + let request_hash = request_hash.clone(); + let model_name = request.model.clone(); + async move { + let done = Bytes::from_static(b"data: [DONE]\n\n"); + let chat_id = public_signature_chat_id.lock().await.clone(); + let stream_errored = stream_error_count + .load(std::sync::atomic::Ordering::Relaxed) + > 0; + + if legacy_gateway_signature_enabled && !stream_errored { + let response_hash = { + let mut hasher = public_signature_hasher.lock().await; + hasher.update(&done); + hex::encode(hasher.clone().finalize()) + }; + if let Some(chat_id) = chat_id { + if let Err(error) = attestation_service + .store_chat_signature_and_unpin( + &chat_id, + request_hash, + response_hash, + ) + .await + { + tracing::error!( + %organization_id, + model = %model_name, + error = %error, + "Failed to store legacy completion stream signature" + ); + } + } else { + tracing::warn!( + %organization_id, + model = %model_name, + "Cannot store legacy completion stream signature: no chat_id observed" + ); + } + } else if let Some(chat_id) = chat_id { + // The service did not fetch a provider signature for this + // transformed response, so release its routing pin here. + attestation_service + .release_chat_signature_pin(&chat_id) + .await; + } + + Ok::(done) + } })); let mut response_builder = Response::builder() @@ -2619,13 +2727,18 @@ async fn completions_inner( .await { Ok(response_with_bytes) => { - let inference_id = hash_inference_id_to_uuid(&response_with_bytes.response.id); + let chat_id = response_with_bytes.response.id.clone(); + let inference_id = hash_inference_id_to_uuid(&chat_id); let completion = chat_response_to_text_response(response_with_bytes.response); let body_bytes = match serde_json::to_vec(&completion) { Ok(b) => b, Err(e) => { tracing::error!(error = %e, "failed to serialize text completion response"); + app_state + .attestation_service + .release_chat_signature_pin(&chat_id) + .await; return ( StatusCode::INTERNAL_SERVER_ERROR, ResponseJson(ErrorResponse::new( @@ -2649,6 +2762,26 @@ async fn completions_inner( None => body_bytes, }; + if legacy_gateway_signature_enabled { + let response_hash = hex::encode(Sha256::digest(&body_bytes)); + if let Err(error) = app_state + .attestation_service + .store_chat_signature_and_unpin(&chat_id, request_hash, response_hash) + .await + { + tracing::error!( + chat_id = %chat_id, + error = %error, + "Failed to store legacy completion signature" + ); + } + } else { + app_state + .attestation_service + .release_chat_signature_pin(&chat_id) + .await; + } + let mut response_builder = Response::builder() .status(StatusCode::OK) .header(header::CONTENT_TYPE, "application/json") diff --git a/crates/api/tests/e2e_all/model_alias_transparency.rs b/crates/api/tests/e2e_all/model_alias_transparency.rs index 946f074c4..eb8297f6f 100644 --- a/crates/api/tests/e2e_all/model_alias_transparency.rs +++ b/crates/api/tests/e2e_all/model_alias_transparency.rs @@ -6,6 +6,7 @@ use crate::common::*; use api::models::BatchUpdateModelApiRequest; +use bytes::Bytes; /// Create a synthetic model and deprecate it in favor of the e2e Qwen mock /// model, returning the deprecated (alias) name. This reproduces the exact @@ -67,10 +68,13 @@ async fn test_aliased_request_warns_non_streaming() { let org = setup_org_with_credits(&server, 10_000_000_000i64).await; let api_key = get_api_key_for_org(&server, org.id).await; + let request_body = chat_body(&alias, false); + let request_json = serde_json::to_string(&request_body).expect("request should serialize"); let response = server .post("/v1/chat/completions") .add_header("Authorization", format!("Bearer {api_key}")) - .json(&chat_body(&alias, false)) + .content_type("application/json") + .bytes(Bytes::from(request_json.clone())) .await; assert_eq!(response.status_code(), 200, "{}", response.text()); @@ -85,7 +89,9 @@ async fn test_aliased_request_warns_non_streaming() { assert_eq!(header, format!("{alias} -> {E2E_QWEN_MODEL_NAME}")); // Body carries the canonical model and a top-level warning - let body: serde_json::Value = response.json(); + let response_text = response.text(); + let body: serde_json::Value = + serde_json::from_str(&response_text).expect("chat response should be JSON"); assert_eq!(body["model"], E2E_QWEN_MODEL_NAME); let warning = body["warning"] .as_str() @@ -94,21 +100,55 @@ async fn test_aliased_request_warns_non_streaming() { warning.contains(&alias) && warning.contains(E2E_QWEN_MODEL_NAME), "warning should name both alias and canonical model: {warning}" ); + + let chat_id = body["id"].as_str().expect("response should have an id"); + let signature_response = server + .get(format!("/v1/signature/{chat_id}?signing_algo=ecdsa").as_str()) + .add_header("Authorization", format!("Bearer {api_key}")) + .await; + assert_eq!( + signature_response.status_code(), + 200, + "gateway signature should be available: {}", + signature_response.text() + ); + let signature = signature_response.json::(); + assert_eq!(signature["signature_kind"], "gateway"); + assert_eq!( + signature["text"], + format!( + "{}:{}", + compute_sha256(&request_json), + compute_sha256(&response_text) + ) + ); } #[tokio::test] async fn test_aliased_request_warns_streaming_first_chunk() { - let server = setup_test_server().await; + use http_body_util::BodyExt; + use tower::ServiceExt; + + let (server, router) = setup_test_server_and_router().await; let alias = setup_deprecated_alias(&server).await; let org = setup_org_with_credits(&server, 10_000_000_000i64).await; let api_key = get_api_key_for_org(&server, org.id).await; - let response = server - .post("/v1/chat/completions") - .add_header("Authorization", format!("Bearer {api_key}")) - .json(&chat_body(&alias, true)) - .await; - assert_eq!(response.status_code(), 200, "{}", response.text()); + let mut request_body = chat_body(&alias, true); + request_body["stream_options"] = serde_json::json!({ + "continuous_usage_stats": true + }); + let request_json = serde_json::to_string(&request_body).expect("request should serialize"); + let request = axum::http::Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("Authorization", format!("Bearer {api_key}")) + .header("Content-Type", "application/json") + .body(axum::body::Body::from(request_json.clone())) + .expect("request should build"); + let response = router.clone().oneshot(request).await; + let response = response.expect("router should serve the streaming request"); + assert_eq!(response.status(), axum::http::StatusCode::OK); let header = response .headers() @@ -119,8 +159,27 @@ async fn test_aliased_request_warns_streaming_first_chunk() { .to_string(); assert_eq!(header, format!("{alias} -> {E2E_QWEN_MODEL_NAME}")); - // Only the FIRST data chunk carries the warning - let text = response.text(); + // Stop as soon as [DONE] is observed. The signature must already be + // available at this point; polling further frames would hide a race in + // which the route stores the signature after exposing the terminator. + let mut body = response.into_body(); + let mut received = Vec::new(); + let mut saw_done = false; + while let Some(frame) = body.frame().await { + let frame = frame.expect("stream frame should not error"); + let Some(data) = frame.data_ref() else { + continue; + }; + received.extend_from_slice(data); + if String::from_utf8_lossy(&received).contains("data: [DONE]") { + saw_done = true; + break; + } + } + let text = String::from_utf8(received).expect("SSE body should be UTF-8"); + assert!(saw_done, "stream should end with [DONE]: {text}"); + + // Only the FIRST data chunk carries the warning. let mut data_chunks = text .lines() .filter_map(|l| l.strip_prefix("data: ")) @@ -128,6 +187,10 @@ async fn test_aliased_request_warns_streaming_first_chunk() { .map(|d| serde_json::from_str::(d).expect("chunk should parse")); let first = data_chunks.next().expect("stream should have chunks"); + let chat_id = first["id"] + .as_str() + .expect("first stream chunk should have an id") + .to_string(); let warning = first["warning"] .as_str() .expect("first chunk of aliased stream must carry a warning"); @@ -143,6 +206,50 @@ async fn test_aliased_request_warns_streaming_first_chunk() { "only the first chunk should carry the warning, got: {chunk}" ); } + + let signature_request = axum::http::Request::builder() + .method("GET") + .uri(format!("/v1/signature/{chat_id}?signing_algo=ecdsa")) + .header("Authorization", format!("Bearer {api_key}")) + .body(axum::body::Body::empty()) + .expect("signature request should build"); + let signature_response = router.clone().oneshot(signature_request).await; + let signature_response = signature_response.expect("router should serve signature request"); + let signature_status = signature_response.status(); + let signature_bytes = signature_response + .into_body() + .collect() + .await + .expect("signature body should collect") + .to_bytes(); + assert_eq!( + signature_status, + axum::http::StatusCode::OK, + "gateway signature should be available: {}", + String::from_utf8_lossy(&signature_bytes) + ); + let signature: serde_json::Value = + serde_json::from_slice(&signature_bytes).expect("signature response should be JSON"); + assert_eq!(signature["signature_kind"], "gateway"); + assert_eq!( + signature["text"], + format!( + "{}:{}", + compute_sha256(&request_json), + compute_sha256(&text) + ) + ); + + while let Some(frame) = body.frame().await { + let frame = frame.expect("trailing frame should not error"); + if let Some(data) = frame.data_ref() { + assert!( + data.is_empty(), + "no bytes may follow [DONE]: {:?}", + String::from_utf8_lossy(data) + ); + } + } } #[tokio::test] diff --git a/crates/api/tests/e2e_all/signature_verification.rs b/crates/api/tests/e2e_all/signature_verification.rs index 8f914905c..cb6c06609 100644 --- a/crates/api/tests/e2e_all/signature_verification.rs +++ b/crates/api/tests/e2e_all/signature_verification.rs @@ -2,12 +2,170 @@ use crate::common::*; +use bytes::Bytes; use inference_providers::StreamChunk; // ============================================ // Streaming Signature Verification Tests // ============================================ +#[tokio::test] +async fn test_legacy_completion_gateway_signature_hashes_public_json() { + let server = setup_test_server().await; + setup_qwen_model(&server).await; + let org = setup_org_with_credits(&server, 10_000_000_000i64).await; + let api_key = get_api_key_for_org(&server, org.id).await; + + let request_body = serde_json::json!({ + "model": E2E_QWEN_MODEL_NAME, + "prompt": "Respond with only two words.", + "max_tokens": 16 + }); + let request_json = serde_json::to_string(&request_body).expect("request should serialize"); + let response = server + .post("/v1/completions") + .add_header("Authorization", format!("Bearer {api_key}")) + .content_type("application/json") + .bytes(Bytes::from(request_json.clone())) + .await; + assert_eq!(response.status_code(), 200, "{}", response.text()); + + let response_text = response.text(); + let completion: serde_json::Value = + serde_json::from_str(&response_text).expect("legacy response should be JSON"); + let chat_id = completion["id"] + .as_str() + .expect("legacy response should include an id"); + + let signature_response = server + .get(format!("/v1/signature/{chat_id}?signing_algo=ecdsa").as_str()) + .add_header("Authorization", format!("Bearer {api_key}")) + .await; + assert_eq!( + signature_response.status_code(), + 200, + "gateway signature should be available: {}", + signature_response.text() + ); + let signature = signature_response.json::(); + assert_eq!(signature["signature_kind"], "gateway"); + assert_eq!( + signature["text"], + format!( + "{}:{}", + compute_sha256(&request_json), + compute_sha256(&response_text) + ) + ); +} + +#[tokio::test] +async fn test_legacy_stream_gateway_signature_is_ready_at_done() { + use http_body_util::BodyExt; + use tower::ServiceExt; + + let (server, router) = setup_test_server_and_router().await; + setup_qwen_model(&server).await; + let org = setup_org_with_credits(&server, 10_000_000_000i64).await; + let api_key = get_api_key_for_org(&server, org.id).await; + + let request_body = serde_json::json!({ + "model": E2E_QWEN_MODEL_NAME, + "prompt": "Respond with only two words.", + "max_tokens": 16, + "stream": true + }); + let request_json = serde_json::to_string(&request_body).expect("request should serialize"); + let request = axum::http::Request::builder() + .method("POST") + .uri("/v1/completions") + .header("Authorization", format!("Bearer {api_key}")) + .header("Content-Type", "application/json") + .body(axum::body::Body::from(request_json.clone())) + .expect("request should build"); + let response = router + .clone() + .oneshot(request) + .await + .expect("router should serve the streaming request"); + assert_eq!(response.status(), axum::http::StatusCode::OK); + + let mut body = response.into_body(); + let mut received = Vec::new(); + let mut saw_done = false; + while let Some(frame) = body.frame().await { + let frame = frame.expect("stream frame should not error"); + let Some(data) = frame.data_ref() else { + continue; + }; + received.extend_from_slice(data); + if String::from_utf8_lossy(&received).contains("data: [DONE]") { + saw_done = true; + break; + } + } + let response_text = String::from_utf8(received).expect("SSE body should be UTF-8"); + assert!( + saw_done, + "legacy stream should end with [DONE]: {response_text}" + ); + + let chat_id = response_text + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|data| data.trim() != "[DONE]") + .find_map(|data| serde_json::from_str::(data).ok()) + .and_then(|chunk| chunk["id"].as_str().map(ToOwned::to_owned)) + .expect("legacy stream should include an id"); + + let signature_request = axum::http::Request::builder() + .method("GET") + .uri(format!("/v1/signature/{chat_id}?signing_algo=ecdsa")) + .header("Authorization", format!("Bearer {api_key}")) + .body(axum::body::Body::empty()) + .expect("signature request should build"); + let signature_response = router + .clone() + .oneshot(signature_request) + .await + .expect("router should serve signature request"); + let signature_status = signature_response.status(); + let signature_bytes = signature_response + .into_body() + .collect() + .await + .expect("signature body should collect") + .to_bytes(); + assert_eq!( + signature_status, + axum::http::StatusCode::OK, + "gateway signature must be available the instant [DONE] is decoded: {}", + String::from_utf8_lossy(&signature_bytes) + ); + let signature: serde_json::Value = + serde_json::from_slice(&signature_bytes).expect("signature response should be JSON"); + assert_eq!(signature["signature_kind"], "gateway"); + assert_eq!( + signature["text"], + format!( + "{}:{}", + compute_sha256(&request_json), + compute_sha256(&response_text) + ) + ); + + while let Some(frame) = body.frame().await { + let frame = frame.expect("trailing frame should not error"); + if let Some(data) = frame.data_ref() { + assert!( + data.is_empty(), + "no bytes may follow [DONE]: {:?}", + String::from_utf8_lossy(data) + ); + } + } +} + #[tokio::test] async fn test_streaming_chat_completion_signature_verification() { let server = setup_test_server().await; From eb19842f394385b934b8415d7c1f46f5d27420e6 Mon Sep 17 00:00:00 2001 From: Coffee <95295094+hanakannzashi@users.noreply.github.com> Date: Thu, 3 Sep 2026 16:12:00 +0800 Subject: [PATCH 2/3] fix: release pins when rewrite streams are cancelled --- crates/api/src/routes/completions.rs | 82 ++++++++++- crates/api/tests/common/mod.rs | 21 +++ .../tests/e2e_all/signature_verification.rs | 133 ++++++++++++++++++ 3 files changed, 233 insertions(+), 3 deletions(-) diff --git a/crates/api/src/routes/completions.rs b/crates/api/src/routes/completions.rs index a42c66508..240cbef2e 100644 --- a/crates/api/src/routes/completions.rs +++ b/crates/api/src/routes/completions.rs @@ -18,6 +18,7 @@ use axum::{ response::{IntoResponse, Json as ResponseJson, Response}, }; use futures::stream::StreamExt; +use services::attestation::ports::AttestationServiceTrait; use services::auto_redact::{self, AutoRedactError, RedactionMap, StreamUnredact}; use services::common::encryption_headers as service_encryption_headers; use services::completions::{ @@ -43,6 +44,40 @@ const USAGE_RECORDING_TIMEOUT_SECS: u64 = 5; // the remaining stream (including the buffered control events) flow through. const MAX_LEADING_CONTROL_EVENTS: usize = 32; +/// Releases the provider-routing signature pin if a response body is dropped +/// before its normal end-of-stream finalization runs. +struct StreamSignaturePinGuard { + attestation_service: Arc, + chat_id: Option, + armed: Arc, +} + +impl Drop for StreamSignaturePinGuard { + fn drop(&mut self) { + let Some(chat_id) = self.chat_id.clone() else { + return; + }; + if !self.armed.swap(false, std::sync::atomic::Ordering::AcqRel) { + return; + } + + let attestation_service = self.attestation_service.clone(); + let handle = match tokio::runtime::Handle::try_current() { + Ok(handle) => handle, + Err(_) => { + tracing::error!(%chat_id, "Cannot release signature routing pin: no Tokio runtime available"); + return; + } + }; + std::mem::drop(handle.spawn(async move { + attestation_service + .release_chat_signature_pin(&chat_id) + .await; + tracing::debug!(%chat_id, "Released signature routing pin after stream cancellation"); + })); + } +} + /// Insert validated E2EE headers into a provider `extra` HashMap. fn insert_encryption_headers( encryption_headers: &crate::routes::common::EncryptionHeaders, @@ -1641,10 +1676,19 @@ async fn chat_completions_inner( )); let final_stream_usage_for_chain = final_stream_usage.clone(); let public_signature_hasher = Arc::new(tokio::sync::Mutex::new(Sha256::new())); - let public_signature_chat_id = Arc::new(tokio::sync::Mutex::new(None::)); + let public_signature_chat_id = + Arc::new(tokio::sync::Mutex::new(stream_chat_id.clone())); let public_signature_hasher_for_chain = public_signature_hasher.clone(); let public_signature_chat_id_for_chain = public_signature_chat_id.clone(); let attestation_service_for_chain = app_state.attestation_service.clone(); + let stream_signature_pin_armed = + Arc::new(std::sync::atomic::AtomicBool::new(stream_chat_id.is_some())); + let stream_signature_pin_armed_for_chain = stream_signature_pin_armed.clone(); + let stream_signature_pin_guard = StreamSignaturePinGuard { + attestation_service: app_state.attestation_service.clone(), + chat_id: stream_chat_id.clone(), + armed: stream_signature_pin_armed, + }; // Re-attach any stashed leading control events, then convert // to a raw bytes stream. @@ -1897,6 +1941,7 @@ async fn chat_completions_inner( let organization_id = api_key.organization.id.0; let model_name = request.model.clone(); let request_hash = request_hash.clone(); + let stream_signature_pin_armed = stream_signature_pin_armed_for_chain; async move { let mut combined: Vec = Vec::new(); let error_count_final = @@ -2032,6 +2077,12 @@ async fn chat_completions_inner( } } + // The tail owns all normal signature finalization and + // explicit-release paths. If the body is dropped before + // this point, the response-body guard releases the pin. + stream_signature_pin_armed + .store(false, std::sync::atomic::Ordering::Release); + if combined.is_empty() { // Avoid emitting an empty body frame. None @@ -2041,7 +2092,12 @@ async fn chat_completions_inner( } }) .filter_map(std::future::ready), - ); + ) + .map(move |item| { + // Keep cleanup alive while Axum owns the response body. + let _ = &stream_signature_pin_guard; + item + }); // Look up which trust tier served this stream. The pool stores a // chat_id → provider mapping when the first chunk arrives; we read @@ -2540,6 +2596,14 @@ async fn completions_inner( |canonical| alias_warning_message(&request.model, canonical), ))); let pending_warning = alias_warning_pending.clone(); + let stream_signature_pin_armed = + Arc::new(std::sync::atomic::AtomicBool::new(stream_chat_id.is_some())); + let stream_signature_pin_armed_for_chain = stream_signature_pin_armed.clone(); + let stream_signature_pin_guard = StreamSignaturePinGuard { + attestation_service: app_state.attestation_service.clone(), + chat_id: stream_chat_id.clone(), + armed: stream_signature_pin_armed, + }; let public_signature_hasher_for_chunks = public_signature_hasher.clone(); let public_signature_chat_id_for_chunks = public_signature_chat_id.clone(); let stream_error_count_for_chunks = stream_error_count.clone(); @@ -2626,6 +2690,7 @@ async fn completions_inner( let stream_error_count = stream_error_count.clone(); let request_hash = request_hash.clone(); let model_name = request.model.clone(); + let stream_signature_pin_armed = stream_signature_pin_armed_for_chain; async move { let done = Bytes::from_static(b"data: [DONE]\n\n"); let chat_id = public_signature_chat_id.lock().await.clone(); @@ -2670,9 +2735,20 @@ async fn completions_inner( .await; } + // The legacy tail owns normal finalization. A dropped + // response body before it is polled is cleaned up by + // the response-body guard. + stream_signature_pin_armed + .store(false, std::sync::atomic::Ordering::Release); + Ok::(done) } - })); + })) + .map(move |item| { + // Keep cleanup alive while Axum owns the response body. + let _ = &stream_signature_pin_guard; + item + }); let mut response_builder = Response::builder() .status(StatusCode::OK) diff --git a/crates/api/tests/common/mod.rs b/crates/api/tests/common/mod.rs index 561a52b64..dbeb1da70 100644 --- a/crates/api/tests/common/mod.rs +++ b/crates/api/tests/common/mod.rs @@ -478,6 +478,27 @@ pub async fn setup_test_server_and_router() -> (axum_test::TestServer, axum::Rou (server, router) } +/// Like [`setup_test_server_with_pool`], with the underlying router for tests +/// that need both mock-provider control and frame-by-frame response polling. +pub async fn setup_test_server_with_pool_and_router() -> ( + axum_test::TestServer, + axum::Router, + std::sync::Arc, + std::sync::Arc, + Arc, +) { + let infra = setup_test_infrastructure().await; + let (server, inference_provider_pool, mock_provider, router) = + build_test_server_components(infra.database.clone(), infra.config).await; + ( + server, + router, + inference_provider_pool, + mock_provider, + infra.database, + ) +} + pub async fn setup_test_server_real_providers() -> ( axum_test::TestServer, Arc, diff --git a/crates/api/tests/e2e_all/signature_verification.rs b/crates/api/tests/e2e_all/signature_verification.rs index cb6c06609..9fe9e65af 100644 --- a/crates/api/tests/e2e_all/signature_verification.rs +++ b/crates/api/tests/e2e_all/signature_verification.rs @@ -2,9 +2,23 @@ use crate::common::*; +use api::models::BatchUpdateModelApiRequest; use bytes::Bytes; use inference_providers::StreamChunk; +fn first_stream_chat_id(response_text: &str) -> String { + response_text + .lines() + .filter_map(|line| line.strip_prefix("data: ")) + .filter(|data| data.trim() != "[DONE]") + .find_map(|data| match serde_json::from_str::(data) { + Ok(StreamChunk::Chat(chunk)) => Some(chunk.id), + Ok(StreamChunk::Text(chunk)) => Some(chunk.id), + _ => None, + }) + .expect("stream should include a chat completion id") +} + // ============================================ // Streaming Signature Verification Tests // ============================================ @@ -166,6 +180,125 @@ async fn test_legacy_stream_gateway_signature_is_ready_at_done() { } } +#[tokio::test] +async fn test_dropping_alias_stream_releases_signature_routing_pin() { + use http_body_util::BodyExt; + use tower::ServiceExt; + + let (server, router, _pool, mock, _database) = setup_test_server_with_pool_and_router().await; + setup_qwen_model(&server).await; + let alias = format!("test-signature-alias-{}", uuid::Uuid::new_v4()); + let mut batch = BatchUpdateModelApiRequest::new(); + batch.insert( + E2E_QWEN_MODEL_NAME.to_string(), + serde_json::from_value(serde_json::json!({ "aliases": [alias] })) + .expect("alias update should deserialize"), + ); + admin_batch_upsert_models(&server, batch, get_session_id()).await; + let org = setup_org_with_credits(&server, 10_000_000_000i64).await; + let api_key = get_api_key_for_org(&server, org.id).await; + + let request = axum::http::Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("Authorization", format!("Bearer {api_key}")) + .header("Content-Type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "model": alias, + "messages": [{ "role": "user", "content": "Respond with two words." }], + "stream": true, + "stream_options": { "continuous_usage_stats": true }, + "nonce": 905 + }) + .to_string(), + )) + .expect("request should build"); + let response = router + .oneshot(request) + .await + .expect("router should serve the streaming request"); + assert_eq!(response.status(), axum::http::StatusCode::OK); + + let mut body = response.into_body(); + let frame = body + .frame() + .await + .expect("stream should yield a first frame") + .expect("first frame should not error"); + let first_bytes = frame.data_ref().expect("first frame should contain data"); + let first_response = String::from_utf8(first_bytes.to_vec()).expect("SSE should be UTF-8"); + let chat_id = first_stream_chat_id(&first_response); + + drop(body); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if mock.unpinned_chat_ids() == vec![chat_id.clone()] { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("cancelling an alias stream should release its signature routing pin"); + assert_eq!(mock.unpinned_chat_ids(), vec![chat_id]); +} + +#[tokio::test] +async fn test_dropping_legacy_stream_releases_signature_routing_pin() { + use http_body_util::BodyExt; + use tower::ServiceExt; + + let (server, router, _pool, mock, _database) = setup_test_server_with_pool_and_router().await; + setup_qwen_model(&server).await; + let org = setup_org_with_credits(&server, 10_000_000_000i64).await; + let api_key = get_api_key_for_org(&server, org.id).await; + + let request = axum::http::Request::builder() + .method("POST") + .uri("/v1/completions") + .header("Authorization", format!("Bearer {api_key}")) + .header("Content-Type", "application/json") + .body(axum::body::Body::from( + serde_json::json!({ + "model": E2E_QWEN_MODEL_NAME, + "prompt": "Respond with two words.", + "stream": true, + "nonce": 906 + }) + .to_string(), + )) + .expect("request should build"); + let response = router + .oneshot(request) + .await + .expect("router should serve the streaming request"); + assert_eq!(response.status(), axum::http::StatusCode::OK); + + let mut body = response.into_body(); + let frame = body + .frame() + .await + .expect("stream should yield a first frame") + .expect("first frame should not error"); + let first_bytes = frame.data_ref().expect("first frame should contain data"); + let first_response = String::from_utf8(first_bytes.to_vec()).expect("SSE should be UTF-8"); + let chat_id = first_stream_chat_id(&first_response); + + drop(body); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + if mock.unpinned_chat_ids() == vec![chat_id.clone()] { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("cancelling a legacy stream should release its signature routing pin"); + assert_eq!(mock.unpinned_chat_ids(), vec![chat_id]); +} + #[tokio::test] async fn test_streaming_chat_completion_signature_verification() { let server = setup_test_server().await; From 5da28816a72ef894e7dcbba5aefeb9a13270cf8a Mon Sep 17 00:00:00 2001 From: Coffee <95295094+hanakannzashi@users.noreply.github.com> Date: Fri, 4 Sep 2026 09:52:01 +0800 Subject: [PATCH 3/3] chore: rerun checks after retarget