perf(attestation): write both chat signatures in one statement

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<ChatSignature>) -> Vec<ChatSignature> {
+    let mut deduped: Vec<ChatSignature> = 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<ChatSignature>,
+    ) -> 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<Option<&str>> = 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(), &params)
+            .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());
+    }
+}

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<Mutex<Vec<(String, ChatSignature)>>>,
+    batch_sizes: Arc<Mutex<Vec<usize>>>,
 }
 
 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<usize> {
+        self.batch_sizes.lock().unwrap().clone()
+    }
 }
 
 #[async_trait]
 impl AttestationRepository for RecordingRepository {
+    async fn add_chat_signatures(
+        &self,
+        chat_id: &str,
+        signatures: Vec<ChatSignature>,
+    ) -> 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<InferenceProviderPool>, Ar
     (pool, provider)
 }
 
-fn lifecycle_service(
+pub(super) fn lifecycle_service(
     repository: Arc<dyn AttestationRepository + Send + Sync>,
     pool: Arc<InferenceProviderPool>,
 ) -> AttestationService {

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<tokio::time::Instant>,
     ) -> 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?;