Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions crates/api/tests/e2e_all/signature_verification.rs
Original file line number Diff line number Diff line change
Expand Up @@ -249,6 +249,62 @@ async fn test_raw_stream_with_upstream_done_retains_provider_signature() {
);
}

#[tokio::test]
async fn test_raw_stream_error_does_not_store_a_signature() {
let (server, _pool, mock, database) = setup_test_server_with_pool().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;

mock.set_default_response(
inference_providers::mock::ResponseTemplate::new("partial output").with_stream_error_after(
1,
inference_providers::CompletionError::HttpError {
status_code: 503,
message: "upstream stream failed".to_string(),
is_external: false,
},
),
)
.await;

let request_body = serde_json::json!({
"model": E2E_QWEN_MODEL_NAME,
"messages": [{ "role": "user", "content": "Respond with two words." }],
"stream": true,
"stream_options": { "continuous_usage_stats": true },
"nonce": 903
});
let response = server
.post("/v1/chat/completions")
.add_header("Authorization", format!("Bearer {api_key}"))
.json(&request_body)
.await;
assert_eq!(response.status_code(), 200, "{}", response.text());

let response_text = response.text();
assert!(response_text.contains("error"));
let chat_id = first_stream_chat_id(&response_text);
let client = database
.pool()
.get()
.await
.expect("database should connect");
let row = client
.query_one(
"SELECT COUNT(*) FROM chat_signatures WHERE chat_id = $1",
&[&chat_id],
)
.await
.expect("signature count query should succeed");
let signature_count: i64 = row.get(0);
assert_eq!(
signature_count, 0,
"error streams must not store a signature"
);
assert_eq!(mock.unpinned_chat_ids(), vec![chat_id]);
}

#[tokio::test]
async fn test_legacy_completion_gateway_signature_hashes_public_json() {
let server = setup_test_server().await;
Expand Down
12 changes: 6 additions & 6 deletions crates/services/src/completions/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,12 +160,12 @@ where

let attestation_service = self.attestation_service.clone();

// For a clean upstream EOF, a provider signature only verifies the
// exact upstream stream. If the route will synthesize a terminator,
// do not publish it as a signature for the client response. The route
// may create a Gateway signature for that synthesized terminator;
// either way this releases the routing pin.
if self.last_error.is_none() && !self.saw_upstream_done_marker {
// A provider signature only verifies the exact upstream stream. Do
// not publish one if the stream reported an error or if the route
// will synthesize its terminator. The route may create a Gateway
// signature for a clean synthesized terminator; either way this
// releases the routing pin.
if self.last_error.is_some() || !self.saw_upstream_done_marker {
Comment thread
hanakannzashi marked this conversation as resolved.
return Box::pin(async move {
attestation_service
.release_chat_signature_pin(&chat_id)
Expand Down
Loading