From 63ca0f9419e97cc494162f2a3ecf92bf56622646 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Thu, 10 Sep 2026 23:58:52 +0000 Subject: [PATCH 1/4] perf(attestation): write both chat signatures in one statement The provider (TEE) and gateway signature paths each stored the ecdsa and ed25519 signatures with two separate INSERT ... ON CONFLICT statements. On the streaming path this runs before the client receives [DONE], so on a replica far from the database every statement is a user-visible round trip (verification p50 is 38 ms on cpu01 and 377 ms on cpu02). AttestationRepository gains add_chat_signatures with a default one-by-one implementation, and the Postgres repository overrides it with a single multi-row upsert using the same conflict clause. Both callers collect the signatures first and write once. --- .../database/src/repositories/attestation.rs | 61 +++++++++++++++++++ .../src/attestation/chat_signatures.rs | 54 ++++++++-------- .../src/attestation/gateway_signatures.rs | 48 +++++++-------- crates/services/src/attestation/ports.rs | 14 +++++ 4 files changed, 124 insertions(+), 53 deletions(-) diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index e25197334..d3919fe0c 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -72,6 +72,67 @@ impl AttestationRepository for PgAttestationRepository { Ok(()) } + /// One round trip for all signatures of a chat: the streaming path stores + /// an ecdsa and an ed25519 signature before the client sees `[DONE]`, and + /// on a replica far from the database each statement is a full network + /// round trip. + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + if signatures.is_empty() { + return Ok(()); + } + let client = self + .pool + .get() + .await + .map_err(|e| AttestationError::RepositoryError(e.to_string()))?; + + let signature_kinds: Vec> = signatures + .iter() + .map(|signature| signature.signature_kind.map(|kind| kind.as_str())) + .collect(); + let mut sql = String::from( + "INSERT INTO chat_signatures (chat_id, text, signature, signing_address, signing_algo, signature_kind) VALUES ", + ); + let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); + for (index, (signature, signature_kind)) in + signatures.iter().zip(signature_kinds.iter()).enumerate() + { + if index > 0 { + sql.push_str(", "); + } + let base = index * 6; + sql.push_str(&format!( + "(${}, ${}, ${}, ${}, ${}, ${})", + base + 1, + base + 2, + base + 3, + base + 4, + base + 5, + base + 6 + )); + params.push(&chat_id); + params.push(&signature.text); + params.push(&signature.signature); + params.push(&signature.signing_address); + params.push(&signature.signing_algo); + params.push(signature_kind); + } + sql.push_str( + " ON CONFLICT (chat_id, signing_algo) DO UPDATE SET text = EXCLUDED.text, signature = EXCLUDED.signature, signing_address = EXCLUDED.signing_address, signature_kind = EXCLUDED.signature_kind, updated_at = NOW()", + ); + + client + .execute(sql.as_str(), ¶ms) + .await + .map_err(|e| AttestationError::RepositoryError(e.to_string()))?; + + Ok(()) + } + async fn get_chat_signature( &self, chat_id: &str, diff --git a/crates/services/src/attestation/chat_signatures.rs b/crates/services/src/attestation/chat_signatures.rs index 6fa936809..971f12961 100644 --- a/crates/services/src/attestation/chat_signatures.rs +++ b/crates/services/src/attestation/chat_signatures.rs @@ -76,6 +76,7 @@ impl AttestationService { let environment = get_environment(); let env_tag = format!("{TAG_ENVIRONMENT}:{environment}"); let result: Result<(), AttestationError> = async { + let mut signatures = Vec::with_capacity(2); for algo in ["ecdsa", "ed25519"] { let provider_signature = provider .get_signature(chat_id, Some(algo.to_string())) @@ -103,38 +104,39 @@ impl AttestationService { ); AttestationError::ProviderError(e.to_string()) })?; - let signature = ChatSignature { + signatures.push(ChatSignature { text: provider_signature.text, signature: provider_signature.signature, signing_address: provider_signature.signing_address, signing_algo: provider_signature.signing_algo, signature_kind: Some(SignatureKind::ProviderTee), - }; - - self.repository - .add_chat_signature(chat_id, signature) - .await - .map_err(|e| { - tracing::error!( - %chat_id, - error = %e, - "Failed to store chat signature in repository for algorithm: {}", - algo - ); - let duration = start_time.elapsed(); - self.metrics_service.record_count( - METRIC_VERIFICATION_FAILURE, - 1, - &[&format!("{TAG_REASON}:{REASON_REPOSITORY_ERROR}"), &env_tag], - ); - self.metrics_service.record_latency( - METRIC_VERIFICATION_DURATION, - duration, - &[&env_tag], - ); - AttestationError::RepositoryError(e.to_string()) - })?; + }); } + + // Both algorithms land in one statement: this runs before the + // client sees `[DONE]`, so every saved round trip is user-visible. + self.repository + .add_chat_signatures(chat_id, signatures) + .await + .map_err(|e| { + tracing::error!( + %chat_id, + error = %e, + "Failed to store chat signatures in repository" + ); + let duration = start_time.elapsed(); + self.metrics_service.record_count( + METRIC_VERIFICATION_FAILURE, + 1, + &[&format!("{TAG_REASON}:{REASON_REPOSITORY_ERROR}"), &env_tag], + ); + self.metrics_service.record_latency( + METRIC_VERIFICATION_DURATION, + duration, + &[&env_tag], + ); + AttestationError::RepositoryError(e.to_string()) + })?; Ok(()) } .await; diff --git a/crates/services/src/attestation/gateway_signatures.rs b/crates/services/src/attestation/gateway_signatures.rs index 630d7cd2e..c16479ceb 100644 --- a/crates/services/src/attestation/gateway_signatures.rs +++ b/crates/services/src/attestation/gateway_signatures.rs @@ -18,6 +18,7 @@ impl AttestationService { let env_tag = format!("{TAG_ENVIRONMENT}:{environment}"); let signature_text = format!("{request_hash}:{response_hash}"); + let mut signatures = Vec::with_capacity(2); for algo in ["ecdsa", "ed25519"] { let (signature_hex, signing_address) = match algo { "ed25519" => self.sign_ed25519_gateway_signature(&signature_text), @@ -26,35 +27,28 @@ impl AttestationService { "Unknown signing algorithm: {algo}" ))), }?; - - self.repository - .add_chat_signature( - signature_id, - ChatSignature { - text: signature_text.clone(), - signature: signature_hex, - signing_address, - signing_algo: algo.to_string(), - signature_kind: Some(SignatureKind::Gateway), - }, - ) - .await - .map_err(|e| { - tracing::error!( - "Failed to store {} signature in repository for algorithm: {}", - id_label, - algo - ); - AttestationError::RepositoryError(e.to_string()) - })?; - tracing::info!( - signature_kind = id_label, - signature_id = signature_id, - signing_algo = algo, - "Stored gateway signature" - ); + signatures.push(ChatSignature { + text: signature_text.clone(), + signature: signature_hex, + signing_address, + signing_algo: algo.to_string(), + signature_kind: Some(SignatureKind::Gateway), + }); } + self.repository + .add_chat_signatures(signature_id, signatures) + .await + .map_err(|e| { + tracing::error!("Failed to store {} signatures in repository", id_label); + AttestationError::RepositoryError(e.to_string()) + })?; + tracing::info!( + signature_kind = id_label, + signature_id = signature_id, + "Stored gateway signatures" + ); + let duration = start_time.elapsed(); self.metrics_service .record_count(METRIC_SIGNATURE_CREATION_SUCCESS, 1, &[&env_tag]); diff --git a/crates/services/src/attestation/ports.rs b/crates/services/src/attestation/ports.rs index a8fdac25a..e1a1599e3 100644 --- a/crates/services/src/attestation/ports.rs +++ b/crates/services/src/attestation/ports.rs @@ -101,6 +101,20 @@ pub trait AttestationRepository: Send + Sync { chat_id: &str, signature: ChatSignature, ) -> Result<(), AttestationError>; + + /// Store several signatures for one chat id. Backends that can write them + /// in a single statement should override this; the default stores them + /// one by one and stops at the first failure. + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + for signature in signatures { + self.add_chat_signature(chat_id, signature).await?; + } + Ok(()) + } async fn get_chat_signature( &self, chat_id: &str, From 975296066e15aeeb059e83128be24827fea2eb14 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:16:35 +0000 Subject: [PATCH 2/4] review: tolerate duplicate algorithms and keep partial signature writes - A multi-row ON CONFLICT DO UPDATE raises 21000 when two rows of the same statement share the arbiter key; signing_algo comes from the backend, so the repository now collapses duplicates and keeps the last one, matching the previous one-at-a-time overwrite. Unit-tested. - The provider path stores whatever was fetched before propagating a later fetch error, so a backend that serves ecdsa but fails ed25519 still leaves the chat verifiable by ecdsa, as it did before. --- .../database/src/repositories/attestation.rs | 63 +++++++++++++++++++ .../src/attestation/chat_signatures.rs | 26 +++++--- 2 files changed, 81 insertions(+), 8 deletions(-) diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index d3919fe0c..f27fa35bc 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -48,6 +48,27 @@ impl PgAttestationRepository { } } +/// Collapse signatures that share a `signing_algo`, keeping the last one. +/// +/// A multi-row `INSERT ... ON CONFLICT DO UPDATE` raises SQLSTATE 21000 +/// ("cannot affect row a second time") when two rows of the same statement +/// hit the same arbiter key, whereas writing them one at a time simply let +/// the later row overwrite the earlier. `signing_algo` is supplied by the +/// inference backend, so the batch must tolerate duplicates the same way. +fn last_signature_per_algo(signatures: Vec) -> Vec { + let mut deduped: Vec = Vec::with_capacity(signatures.len()); + for signature in signatures { + match deduped + .iter_mut() + .find(|existing| existing.signing_algo == signature.signing_algo) + { + Some(existing) => *existing = signature, + None => deduped.push(signature), + } + } + deduped +} + #[async_trait] impl AttestationRepository for PgAttestationRepository { async fn add_chat_signature( @@ -81,6 +102,7 @@ impl AttestationRepository for PgAttestationRepository { chat_id: &str, signatures: Vec, ) -> Result<(), AttestationError> { + let signatures = last_signature_per_algo(signatures); if signatures.is_empty() { return Ok(()); } @@ -164,3 +186,44 @@ impl AttestationRepository for PgAttestationRepository { self.row_to_chat_signature(row) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn signature(algo: &str, text: &str) -> ChatSignature { + ChatSignature { + text: text.to_string(), + signature: format!("sig-{text}"), + signing_address: "addr".to_string(), + signing_algo: algo.to_string(), + signature_kind: None, + } + } + + #[test] + fn distinct_algorithms_are_kept_in_order() { + let out = last_signature_per_algo(vec![signature("ecdsa", "a"), signature("ed25519", "b")]); + assert_eq!(out.len(), 2); + assert_eq!(out[0].signing_algo, "ecdsa"); + assert_eq!(out[1].signing_algo, "ed25519"); + } + + #[test] + fn duplicate_algorithm_keeps_the_last_signature() { + let out = last_signature_per_algo(vec![ + signature("ecdsa", "first"), + signature("ed25519", "b"), + signature("ecdsa", "second"), + ]); + assert_eq!(out.len(), 2); + assert_eq!(out[0].signing_algo, "ecdsa"); + assert_eq!(out[0].text, "second"); + assert_eq!(out[1].signing_algo, "ed25519"); + } + + #[test] + fn empty_input_stays_empty() { + assert!(last_signature_per_algo(Vec::new()).is_empty()); + } +} diff --git a/crates/services/src/attestation/chat_signatures.rs b/crates/services/src/attestation/chat_signatures.rs index 971f12961..0cfa7e28a 100644 --- a/crates/services/src/attestation/chat_signatures.rs +++ b/crates/services/src/attestation/chat_signatures.rs @@ -75,8 +75,8 @@ impl AttestationService { let environment = get_environment(); let env_tag = format!("{TAG_ENVIRONMENT}:{environment}"); - let result: Result<(), AttestationError> = async { - let mut signatures = Vec::with_capacity(2); + let mut signatures = Vec::with_capacity(2); + let fetched: Result<(), AttestationError> = async { for algo in ["ecdsa", "ed25519"] { let provider_signature = provider .get_signature(chat_id, Some(algo.to_string())) @@ -113,8 +113,19 @@ impl AttestationService { }); } - // Both algorithms land in one statement: this runs before the - // client sees `[DONE]`, so every saved round trip is user-visible. + Ok(()) + } + .await; + + // Store whatever was fetched even if a later algorithm failed: the + // one-at-a-time implementation persisted each signature as it came, + // so a backend that serves ecdsa but fails ed25519 still leaves the + // chat verifiable by ecdsa. Both algorithms land in one statement: + // this runs before the client sees `[DONE]`, so every saved round + // trip is user-visible. + let stored: Result<(), AttestationError> = if signatures.is_empty() { + Ok(()) + } else { self.repository .add_chat_signatures(chat_id, signatures) .await @@ -136,10 +147,9 @@ impl AttestationService { &[&env_tag], ); AttestationError::RepositoryError(e.to_string()) - })?; - Ok(()) - } - .await; + }) + }; + let result = fetched.and(stored); provider.unpin_chat_connection(chat_id); result?; From d8ad797bef61ca35a61483c10ec069502218bcd9 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:34:00 +0000 Subject: [PATCH 3/4] review: log signature id and error when a gateway signature write fails --- crates/services/src/attestation/gateway_signatures.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/services/src/attestation/gateway_signatures.rs b/crates/services/src/attestation/gateway_signatures.rs index c16479ceb..07c4e18bd 100644 --- a/crates/services/src/attestation/gateway_signatures.rs +++ b/crates/services/src/attestation/gateway_signatures.rs @@ -40,7 +40,12 @@ impl AttestationService { .add_chat_signatures(signature_id, signatures) .await .map_err(|e| { - tracing::error!("Failed to store {} signatures in repository", id_label); + tracing::error!( + signature_kind = id_label, + signature_id = signature_id, + error = %e, + "Failed to store gateway signatures in repository" + ); AttestationError::RepositoryError(e.to_string()) })?; tracing::info!( From 37b67d642918c5bdf277af9705dd9e3cb769cfa6 Mon Sep 17 00:00:00 2001 From: Pierre LE GUEN <26087574+PierreLeGuen@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:57:13 +0000 Subject: [PATCH 4/4] fix(attestation): preserve partial signatures when stream fetches time out --- .../chat_signature_lifecycle_tests.rs | 25 +- .../src/attestation/chat_signatures.rs | 77 ++-- crates/services/src/attestation/mod.rs | 3 + crates/services/src/attestation/ports.rs | 10 + .../attestation/provider_signature_tests.rs | 413 ++++++++++++++++++ .../services/src/attestation/service_trait.rs | 16 +- crates/services/src/completions/mod.rs | 8 +- 7 files changed, 514 insertions(+), 38 deletions(-) create mode 100644 crates/services/src/attestation/provider_signature_tests.rs diff --git a/crates/services/src/attestation/chat_signature_lifecycle_tests.rs b/crates/services/src/attestation/chat_signature_lifecycle_tests.rs index 27bce8d96..fd026083a 100644 --- a/crates/services/src/attestation/chat_signature_lifecycle_tests.rs +++ b/crates/services/src/attestation/chat_signature_lifecycle_tests.rs @@ -36,18 +36,35 @@ use crate::{ /// Repository that records stored signatures and succeeds. #[derive(Clone, Default)] -struct RecordingRepository { +pub(super) struct RecordingRepository { stored: Arc>>, + batch_sizes: Arc>>, } impl RecordingRepository { - fn stored(&self) -> Vec<(String, ChatSignature)> { + pub(super) fn stored(&self) -> Vec<(String, ChatSignature)> { self.stored.lock().map(|s| s.clone()).unwrap_or_default() } + + pub(super) fn batch_sizes(&self) -> Vec { + self.batch_sizes.lock().unwrap().clone() + } } #[async_trait] impl AttestationRepository for RecordingRepository { + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + self.batch_sizes.lock().unwrap().push(signatures.len()); + for signature in signatures { + self.add_chat_signature(chat_id, signature).await?; + } + Ok(()) + } + async fn add_chat_signature( &self, chat_id: &str, @@ -72,7 +89,7 @@ impl AttestationRepository for RecordingRepository { } /// Repository whose store always fails. -struct FailingRepository; +pub(super) struct FailingRepository; #[async_trait] impl AttestationRepository for FailingRepository { @@ -251,7 +268,7 @@ async fn pool_with_pinned_chat(chat_id: &str) -> (Arc, Ar (pool, provider) } -fn lifecycle_service( +pub(super) fn lifecycle_service( repository: Arc, pool: Arc, ) -> AttestationService { diff --git a/crates/services/src/attestation/chat_signatures.rs b/crates/services/src/attestation/chat_signatures.rs index 0cfa7e28a..41ed3ee05 100644 --- a/crates/services/src/attestation/chat_signatures.rs +++ b/crates/services/src/attestation/chat_signatures.rs @@ -7,11 +7,17 @@ use super::{ }; use crate::{metrics::consts::*, usage::StopReason}; -/// Upper bound on the gateway signature store at end-of-stream. The client is +/// Upper bound on signature finalization at end-of-stream. The client is /// still waiting for the held-back `[DONE]` while this runs, so it must be /// bounded; on timeout the routing pin is still released and the stream ends /// without a stored signature (logged by the caller). -pub(in crate::attestation) const STREAM_SIGNATURE_STORE_TIMEOUT: Duration = Duration::from_secs(5); +pub(crate) const STREAM_SIGNATURE_STORE_TIMEOUT: Duration = Duration::from_secs(5); + +/// Reserve the final second of the streaming deadline for persisting fetched +/// signatures. An outer timeout would otherwise drop an already-fetched ECDSA +/// signature if the ED25519 fetch stalls before the batch can be written. +pub(in crate::attestation) const PROVIDER_SIGNATURE_FETCH_TIMEOUT: Duration = + STREAM_SIGNATURE_STORE_TIMEOUT.saturating_sub(Duration::from_secs(1)); impl AttestationService { pub(in crate::attestation) async fn get_chat_signature_impl( @@ -54,6 +60,7 @@ impl AttestationService { pub(in crate::attestation) async fn store_chat_signature_from_provider_impl( &self, chat_id: &str, + fetch_deadline: Option, ) -> Result<(), AttestationError> { let start_time = std::time::Instant::now(); let provider = self @@ -78,32 +85,43 @@ impl AttestationService { let mut signatures = Vec::with_capacity(2); let fetched: Result<(), AttestationError> = async { for algo in ["ecdsa", "ed25519"] { - let provider_signature = provider - .get_signature(chat_id, Some(algo.to_string())) - .await - .map_err(|e| { - // The error string embeds the backend URL on connection - // failures — without it (and chat_id) these events are - // impossible to attribute to a model/backend. - tracing::error!( - %chat_id, - error = %e, - "Failed to get chat signature from provider for algorithm: {}", - algo - ); - let duration = start_time.elapsed(); - self.metrics_service.record_count( - METRIC_VERIFICATION_FAILURE, - 1, - &[&format!("{TAG_REASON}:{REASON_INFERENCE_ERROR}"), &env_tag], - ); - self.metrics_service.record_latency( - METRIC_VERIFICATION_DURATION, - duration, - &[&env_tag], - ); - AttestationError::ProviderError(e.to_string()) - })?; + let fetch = async { + provider + .get_signature(chat_id, Some(algo.to_string())) + .await + .map_err(|e| AttestationError::ProviderError(e.to_string())) + }; + let provider_signature = match fetch_deadline { + Some(deadline) => match tokio::time::timeout_at(deadline, fetch).await { + Ok(result) => result, + Err(_) => Err(AttestationError::ProviderError(format!( + "Timed out fetching chat signature for algorithm: {algo}" + ))), + }, + None => fetch.await, + } + .inspect_err(|e| { + // The error string embeds the backend URL on connection + // failures — without it (and chat_id) these events are + // impossible to attribute to a model/backend. + tracing::error!( + %chat_id, + error = %e, + "Failed to get chat signature from provider for algorithm: {}", + algo + ); + let duration = start_time.elapsed(); + self.metrics_service.record_count( + METRIC_VERIFICATION_FAILURE, + 1, + &[&format!("{TAG_REASON}:{REASON_INFERENCE_ERROR}"), &env_tag], + ); + self.metrics_service.record_latency( + METRIC_VERIFICATION_DURATION, + duration, + &[&env_tag], + ); + })?; signatures.push(ChatSignature { text: provider_signature.text, signature: provider_signature.signature, @@ -117,7 +135,8 @@ impl AttestationService { } .await; - // Store whatever was fetched even if a later algorithm failed: the + // Store whatever was fetched even if a later algorithm failed or hit + // the fetch deadline, leaving time before the outer stream timeout. The // one-at-a-time implementation persisted each signature as it came, // so a backend that serves ecdsa but fails ed25519 still leaves the // chat verifiable by ecdsa. Both algorithms land in one statement: diff --git a/crates/services/src/attestation/mod.rs b/crates/services/src/attestation/mod.rs index e54dfea42..0a2ef0812 100644 --- a/crates/services/src/attestation/mod.rs +++ b/crates/services/src/attestation/mod.rs @@ -11,6 +11,8 @@ mod lifecycle; pub mod measurement; pub mod models; pub mod ports; +#[cfg(test)] +mod provider_signature_tests; mod report; pub mod report_data; mod service_trait; @@ -22,6 +24,7 @@ use config::ItaAttestationConfig; use ed25519_dalek::{SigningKey, VerifyingKey}; use k256::ecdsa::{SigningKey as EcdsaSigningKey, VerifyingKey as EcdsaVerifyingKey}; +pub(crate) use chat_signatures::STREAM_SIGNATURE_STORE_TIMEOUT; pub use environment::{ compute_spki_hash, load_tls_cert_fingerprint, load_vpc_info, load_vpc_shared_secret, }; diff --git a/crates/services/src/attestation/ports.rs b/crates/services/src/attestation/ports.rs index e1a1599e3..85e8e46b9 100644 --- a/crates/services/src/attestation/ports.rs +++ b/crates/services/src/attestation/ports.rs @@ -23,6 +23,16 @@ pub trait AttestationServiceTrait: Send + Sync { chat_id: &str, ) -> Result<(), AttestationError>; + /// Fetch and store provider signatures before a stream sends `[DONE]`. + /// Implementations may reserve part of the caller's finalization deadline + /// to persist partial fetches; background callers use the method above. + async fn store_stream_chat_signature_from_provider( + &self, + chat_id: &str, + ) -> Result<(), AttestationError> { + self.store_chat_signature_from_provider(chat_id).await + } + /// Store a chat signature directly over gateway-emitted bytes. /// Creates a signature with text format "request_hash:response_hash" /// and stores both ECDSA and ED25519 signatures. diff --git a/crates/services/src/attestation/provider_signature_tests.rs b/crates/services/src/attestation/provider_signature_tests.rs new file mode 100644 index 000000000..a087cbd65 --- /dev/null +++ b/crates/services/src/attestation/provider_signature_tests.rs @@ -0,0 +1,413 @@ +//! Provider signature batching must retain successful fetches when the other +//! algorithm fails or times out before the streaming finalization deadline. +use std::{sync::Arc, time::Duration}; + +use async_trait::async_trait; +use config::ExternalProvidersConfig; +use inference_providers::{self as ip, mock::MockProvider, InferenceProvider}; + +use super::{ + chat_signature_lifecycle_tests::{lifecycle_service, FailingRepository, RecordingRepository}, + ports::{AttestationRepository, AttestationServiceTrait}, + AttestationError, AttestationService, ChatSignature, SignatureKind, + STREAM_SIGNATURE_STORE_TIMEOUT, +}; +use crate::{ + inference_provider_pool::InferenceProviderPool, + metrics::{ + capturing::{CapturingMetricsService, MetricValue}, + consts::*, + }, +}; + +#[derive(Clone, Copy)] +enum FetchBehavior { + Success, + FailSecond, + StallSecond, + StallFirst, + SlowSecond, +} + +struct ControlledProvider { + inner: MockProvider, + behavior: FetchBehavior, +} + +#[async_trait] +impl InferenceProvider for ControlledProvider { + async fn get_signature( + &self, + chat_id: &str, + signing_algo: Option, + ) -> Result { + match (signing_algo.as_deref(), self.behavior) { + (Some("ecdsa"), FetchBehavior::StallFirst) + | (Some("ed25519"), FetchBehavior::StallSecond) => std::future::pending().await, + (Some("ed25519"), FetchBehavior::FailSecond) => Err( + ip::CompletionError::CompletionError("synthetic fetch failure".to_string()), + ), + (Some("ed25519"), FetchBehavior::SlowSecond) => { + tokio::time::sleep(Duration::from_secs(6)).await; + self.inner.get_signature(chat_id, signing_algo).await + } + _ => { + // Spending part of the budget on ECDSA catches a timeout that + // is accidentally restarted for each algorithm. + if matches!(self.behavior, FetchBehavior::StallSecond) { + tokio::time::sleep(Duration::from_secs(2)).await; + } + self.inner.get_signature(chat_id, signing_algo).await + } + } + } + + fn unpin_chat_connection(&self, chat_id: &str) { + self.inner.unpin_chat_connection(chat_id); + } + async fn models(&self) -> Result { + self.inner.models().await + } + + async fn chat_completion_stream( + &self, + params: ip::ChatCompletionParams, + request_hash: String, + ) -> Result { + self.inner + .chat_completion_stream(params, request_hash) + .await + } + + async fn chat_completion( + &self, + params: ip::ChatCompletionParams, + request_hash: String, + ) -> Result { + self.inner.chat_completion(params, request_hash).await + } + + async fn text_completion_stream( + &self, + params: ip::CompletionParams, + ) -> Result { + self.inner.text_completion_stream(params).await + } + + async fn image_generation( + &self, + params: ip::ImageGenerationParams, + request_hash: String, + ) -> Result { + self.inner.image_generation(params, request_hash).await + } + + async fn image_edit( + &self, + params: Arc, + request_hash: String, + ) -> Result { + self.inner.image_edit(params, request_hash).await + } + + async fn score( + &self, + params: ip::ScoreParams, + request_hash: String, + ) -> Result { + self.inner.score(params, request_hash).await + } + + async fn rerank( + &self, + params: ip::RerankParams, + ) -> Result { + self.inner.rerank(params).await + } + + async fn embeddings_raw( + &self, + body: bytes::Bytes, + extra: std::collections::HashMap, + ) -> Result { + self.inner.embeddings_raw(body, extra).await + } + + async fn privacy_classify_raw( + &self, + body: bytes::Bytes, + extra: std::collections::HashMap, + ) -> Result { + self.inner.privacy_classify_raw(body, extra).await + } + + async fn get_attestation_report( + &self, + model: String, + signing_algo: Option, + nonce: Option, + signing_address: Option, + include_tls_fingerprint: bool, + ) -> Result, ip::models::AttestationError> { + self.inner + .get_attestation_report( + model, + signing_algo, + nonce, + signing_address, + include_tls_fingerprint, + ) + .await + } + + async fn audio_transcription( + &self, + params: ip::AudioTranscriptionParams, + request_hash: String, + ) -> Result { + self.inner.audio_transcription(params, request_hash).await + } +} + +async fn provider_service( + chat_id: &str, + behavior: FetchBehavior, + repository: Arc, +) -> ( + AttestationService, + Arc, + Arc, +) { + let pool = Arc::new(InferenceProviderPool::new( + None, + ExternalProvidersConfig::default(), + )); + let provider = Arc::new(ControlledProvider { + inner: MockProvider::new(), + behavior, + }); + pool.store_chat_id_mapping(chat_id.to_string(), provider.clone()) + .await; + let mut service = lifecycle_service(repository, pool); + let metrics = Arc::new(CapturingMetricsService::new()); + service.metrics_service = metrics.clone(); + (service, provider, metrics) +} + +fn assert_failure_metric(metrics: &CapturingMetricsService, reason: &str) { + let failures: Vec<_> = metrics + .get_metrics() + .into_iter() + .filter(|metric| metric.name == METRIC_VERIFICATION_FAILURE) + .collect(); + assert_eq!(failures.len(), 1); + assert!(matches!(failures[0].value, MetricValue::Count(1))); + assert!(failures[0].tags.contains(&format!("{TAG_REASON}:{reason}"))); + assert!(!metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_SUCCESS)); +} + +#[tokio::test] +async fn provider_success_stores_both_signatures_in_one_batch_and_unpins() { + let chat_id = "chatcmpl-provider-batch-success"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::Success, + Arc::new(repository.clone()), + ) + .await; + service + .store_stream_chat_signature_from_provider(chat_id) + .await + .unwrap(); + assert_eq!(repository.batch_sizes(), vec![2]); + let stored = repository.stored(); + assert_eq!( + stored + .iter() + .map(|(_, sig)| sig.signing_algo.as_str()) + .collect::>(), + vec!["ecdsa", "ed25519"] + ); + assert!(stored + .iter() + .all(|(id, sig)| id == chat_id && sig.signature_kind == Some(SignatureKind::ProviderTee))); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert!(metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_SUCCESS + && matches!(metric.value, MetricValue::Count(1)))); + assert!(!metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_FAILURE)); +} + +#[tokio::test] +async fn provider_second_fetch_error_preserves_first_signature_and_unpins() { + let chat_id = "chatcmpl-provider-partial-error"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::FailSecond, + Arc::new(repository.clone()), + ) + .await; + let result = service + .store_stream_chat_signature_from_provider(chat_id) + .await; + assert!(matches!(result, Err(AttestationError::ProviderError(_)))); + assert_eq!(repository.batch_sizes(), vec![1]); + assert_eq!(repository.stored()[0].1.signing_algo, "ecdsa"); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_INFERENCE_ERROR); +} + +/// Make the flush consume part of its reserved budget, so a passing test +/// requires the write to finish before the actual five-second outer timeout. +struct DelayedRepository(RecordingRepository); + +#[async_trait] +impl AttestationRepository for DelayedRepository { + async fn add_chat_signature( + &self, + chat_id: &str, + signature: ChatSignature, + ) -> Result<(), AttestationError> { + self.0.add_chat_signature(chat_id, signature).await + } + async fn add_chat_signatures( + &self, + chat_id: &str, + signatures: Vec, + ) -> Result<(), AttestationError> { + tokio::time::sleep(Duration::from_millis(500)).await; + self.0.add_chat_signatures(chat_id, signatures).await + } + async fn get_chat_signature( + &self, + chat_id: &str, + signing_algo: &str, + ) -> Result { + self.0.get_chat_signature(chat_id, signing_algo).await + } +} + +#[tokio::test(start_paused = true)] +async fn provider_stalled_second_fetch_flushes_first_before_stream_timeout() { + let chat_id = "chatcmpl-provider-partial-timeout"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::StallSecond, + Arc::new(DelayedRepository(repository.clone())), + ) + .await; + let started = tokio::time::Instant::now(); + // The same hard deadline used by InterceptStream must not cancel the + // partial write, even after ECDSA has used two seconds of the fetch budget. + let result = tokio::time::timeout( + STREAM_SIGNATURE_STORE_TIMEOUT, + service.store_stream_chat_signature_from_provider(chat_id), + ) + .await + .expect("partial signature must be flushed before stream finalization cancels it"); + assert!( + matches!(result, Err(AttestationError::ProviderError(message)) if message.contains("Timed out") && message.contains("ed25519")) + ); + assert_eq!(repository.batch_sizes(), vec![1]); + assert_eq!(repository.stored()[0].1.signing_algo, "ecdsa"); + assert!(started.elapsed() >= Duration::from_millis(4500)); + assert!(started.elapsed() < STREAM_SIGNATURE_STORE_TIMEOUT); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_INFERENCE_ERROR); +} + +#[tokio::test(start_paused = true)] +async fn provider_stalled_first_fetch_unpins_without_an_empty_write() { + let chat_id = "chatcmpl-provider-first-timeout"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::StallFirst, + Arc::new(repository.clone()), + ) + .await; + let result = tokio::time::timeout( + STREAM_SIGNATURE_STORE_TIMEOUT, + service.store_stream_chat_signature_from_provider(chat_id), + ) + .await + .unwrap(); + assert!( + matches!(result, Err(AttestationError::ProviderError(message)) if message.contains("Timed out")) + ); + assert!(repository.batch_sizes().is_empty()); + assert!(repository.stored().is_empty()); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_INFERENCE_ERROR); +} + +#[tokio::test] +async fn provider_repository_error_is_metriced_and_unpins() { + let chat_id = "chatcmpl-provider-repository-error"; + let (service, provider, metrics) = + provider_service(chat_id, FetchBehavior::Success, Arc::new(FailingRepository)).await; + let result = service + .store_stream_chat_signature_from_provider(chat_id) + .await; + assert!(matches!(result, Err(AttestationError::RepositoryError(_)))); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert_failure_metric(&metrics, REASON_REPOSITORY_ERROR); +} + +#[tokio::test(start_paused = true)] +async fn background_provider_fetch_keeps_its_existing_timeout() { + let chat_id = "chatcmpl-provider-background-slow"; + let repository = RecordingRepository::default(); + let (service, provider, metrics) = provider_service( + chat_id, + FetchBehavior::SlowSecond, + Arc::new(repository.clone()), + ) + .await; + let started = tokio::time::Instant::now(); + service + .store_chat_signature_from_provider(chat_id) + .await + .unwrap(); + assert!(started.elapsed() >= Duration::from_secs(6)); + assert_eq!(repository.batch_sizes(), vec![2]); + assert_eq!(repository.stored().len(), 2); + assert_eq!( + provider.inner.unpinned_chat_ids(), + vec![chat_id.to_string()] + ); + assert!(metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_SUCCESS)); + assert!(!metrics + .get_metrics() + .iter() + .any(|metric| metric.name == METRIC_VERIFICATION_FAILURE)); +} diff --git a/crates/services/src/attestation/service_trait.rs b/crates/services/src/attestation/service_trait.rs index e85855d7a..a0dedd58a 100644 --- a/crates/services/src/attestation/service_trait.rs +++ b/crates/services/src/attestation/service_trait.rs @@ -1,6 +1,7 @@ use async_trait::async_trait; use super::{ + chat_signatures::PROVIDER_SIGNATURE_FETCH_TIMEOUT, ita::{ItaTokenQuery, ItaTokenResponse}, models::AttestationReport, ports, AttestationError, AttestationService, SignatureLookupResult, @@ -21,7 +22,20 @@ impl ports::AttestationServiceTrait for AttestationService { &self, chat_id: &str, ) -> Result<(), AttestationError> { - self.store_chat_signature_from_provider_impl(chat_id).await + self.store_chat_signature_from_provider_impl(chat_id, None) + .await + } + + async fn store_stream_chat_signature_from_provider( + &self, + chat_id: &str, + ) -> Result<(), AttestationError> { + // Start the budget before provider lookup and share it across both + // algorithms. Non-streaming/background callers retain their existing + // provider timeouts, without this streaming-specific deadline. + let deadline = tokio::time::Instant::now() + PROVIDER_SIGNATURE_FETCH_TIMEOUT; + self.store_chat_signature_from_provider_impl(chat_id, Some(deadline)) + .await } async fn store_chat_signature( diff --git a/crates/services/src/completions/mod.rs b/crates/services/src/completions/mod.rs index bb7ddef87..f63bc8d8e 100644 --- a/crates/services/src/completions/mod.rs +++ b/crates/services/src/completions/mod.rs @@ -1,6 +1,7 @@ pub mod ports; use crate::attestation::ports::AttestationServiceTrait; +use crate::attestation::STREAM_SIGNATURE_STORE_TIMEOUT; use crate::inference_provider_pool::InferenceProviderPool; use crate::models::ModelsRepository; use crate::responses::models::ResponseId; @@ -22,7 +23,6 @@ use std::task::{Context, Poll}; use std::time::{Duration, Instant}; use tracing::Instrument; -const FINALIZE_TIMEOUT_SECS: u64 = 5; /// A raw provider terminal marker is only expected to contain `data: [DONE]` /// plus its SSE line ending. Bound it before accepting it as a signed terminal /// event so malformed padding cannot be retained or signed. @@ -213,8 +213,8 @@ where Box::pin(async move { match tokio::time::timeout( - Duration::from_secs(FINALIZE_TIMEOUT_SECS), - attestation_service.store_chat_signature_from_provider(&chat_id), + STREAM_SIGNATURE_STORE_TIMEOUT, + attestation_service.store_stream_chat_signature_from_provider(&chat_id), ) .await { @@ -227,7 +227,7 @@ where %organization_id, %model_id, "Timeout storing chat signature after {}s", - FINALIZE_TIMEOUT_SECS + STREAM_SIGNATURE_STORE_TIMEOUT.as_secs() ); // The provider-store implementation normally unpins after // it completes. A timeout cancels that future before its