diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index e25197334..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( @@ -72,6 +93,68 @@ 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> { + let signatures = last_signature_per_algo(signatures); + 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, @@ -103,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_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 6fa936809..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 @@ -75,70 +82,94 @@ 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 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 signature = ChatSignature { + 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, 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()) - })?; + }); } + Ok(()) } .await; + // 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: + // 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 + .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()) + }) + }; + let result = fetched.and(stored); + provider.unpin_chat_connection(chat_id); result?; diff --git a/crates/services/src/attestation/gateway_signatures.rs b/crates/services/src/attestation/gateway_signatures.rs index 630d7cd2e..07c4e18bd 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,33 @@ 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!( + signature_kind = id_label, + signature_id = signature_id, + error = %e, + "Failed to store gateway signatures in repository" + ); + 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/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 a8fdac25a..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. @@ -101,6 +111,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, 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 f1b03dfc6..1bb776d70 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