diff --git a/.config/nextest.toml b/.config/nextest.toml index 41628550e..5051225a4 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -32,7 +32,7 @@ max-threads = 4 # These tests intentionally touch database-wide state (DDL, broad encryption # scans, or ACCESS EXCLUSIVE locks), so they must not overlap any other test. [[profile.default.overrides]] -filter = 'package(api) & binary(e2e_all) & (test(/^database_encryption::/) | test(/^reporting_usage::usage_export::usage_export_database_timeout_returns_504_and_stops_query$/) | test(/^reporting_usage::usage_summary::usage_summary_database_timeout_returns_504$/))' +filter = 'package(api) & binary(e2e_all) & (test(/^database_encryption::/) | test(/^prepared_statement_cache::/) | test(/^reporting_usage::usage_export::usage_export_database_timeout_returns_504_and_stops_query$/) | test(/^reporting_usage::usage_summary::usage_summary_database_timeout_returns_504$/))' test-group = "e2e-db" threads-required = "num-test-threads" diff --git a/crates/api/tests/e2e_all/main.rs b/crates/api/tests/e2e_all/main.rs index 7fe4d3771..a77661513 100644 --- a/crates/api/tests/e2e_all/main.rs +++ b/crates/api/tests/e2e_all/main.rs @@ -67,6 +67,7 @@ mod org_system_prompt; mod organization_deletion; mod pagination_validation; mod patroni_failover; +mod prepared_statement_cache; mod privacy_classify; mod privacy_redact; mod provider_errors; diff --git a/crates/api/tests/e2e_all/prepared_statement_cache.rs b/crates/api/tests/e2e_all/prepared_statement_cache.rs new file mode 100644 index 000000000..ab8757f49 --- /dev/null +++ b/crates/api/tests/e2e_all/prepared_statement_cache.rs @@ -0,0 +1,397 @@ +//! Hot-path queries run through the connection's prepared-statement cache. +//! During a rolling deploy a replica on the previous release keeps serving +//! after the new release has run its migrations; if one of those migrations +//! adds a column to a table read with `SELECT *`, Postgres refuses the old +//! plan with "cached plan must not change result type". The repositories must +//! recover from that transparently, not surface a 500. + +use crate::common::*; +use anyhow::{ensure, Context, Result}; +use database::{models::RecordUsageRequest, repositories::OrganizationUsageRepository}; +use futures::future::join_all; +use uuid::Uuid; + +#[tokio::test] +async fn cached_statements_recover_from_a_column_added_under_them() { + let (server, database) = setup_test_server_with_database().await; + let (session_id, _email) = setup_unique_test_session(&database).await; + let org = create_org_with_session(&server, &session_id).await; + let api_key = get_api_key_for_org_with_session(&server, org.id.clone(), &session_id).await; + + let probe = || async { + server + .get("/v1/files?limit=1") + .add_header("Authorization", format!("Bearer {api_key}")) + .add_header("User-Agent", MOCK_USER_AGENT) + .await + }; + + // Warm the cache: API-key auth resolves the workspace and organization + // with a `SELECT w.*` on every request. Probes run concurrently so the + // statement is cached on every connection of the test pool, not just + // one; that is the shape of a warm production replica. + for response in join_all((0..8).map(|_| probe())).await { + assert_eq!(response.status_code(), 200, "{}", response.text()); + } + + // A migration on the next release adds a column to `workspaces` while + // this process still holds statements prepared against the old shape. + let column = format!( + "rolling_deploy_{}", + &uuid::Uuid::new_v4().simple().to_string()[..8] + ); + { + let client = database + .pool() + .get() + .await + .expect("failed to get database connection"); + client + .execute( + &format!("ALTER TABLE workspaces ADD COLUMN {column} TEXT"), + &[], + ) + .await + .expect("failed to add column"); + } + + // Every request must still succeed: the stale plan is dropped from the + // cache and the statement re-prepared inside the repository retry. + // Collect instead of asserting so the column is dropped even on failure; + // a leftover column would fail the database-encryption classification + // scans that share this database. + // Concurrent again: every warm connection holds a stale plan, and each + // request must recover on whichever connection it lands on. + let outcomes: Vec<_> = join_all((0..8).map(|_| probe())) + .await + .into_iter() + .map(|response| (response.status_code(), response.text())) + .collect(); + + // Best-effort cleanup that never panics: a leftover column breaks the + // database-encryption classification scans that share this database, so + // the drop must run even when the pool or the statement misbehaves, and + // the assertions below must still report the real outcome. + let cleanup = async { + let client = database.pool().get().await?; + client + .execute( + &format!("ALTER TABLE workspaces DROP COLUMN IF EXISTS {column}"), + &[], + ) + .await?; + Ok::<(), Box>(()) + }; + let cleanup = cleanup.await; + + for (status, body) in outcomes { + assert_eq!( + status, 200, + "request after a schema change must recover, got: {body}" + ); + } + + cleanup.expect("failed to drop the temporary column"); +} + +enum UsageStatementInvalidation { + AddedColumn, + DeallocatedBalance, +} + +async fn usage_fixture() -> (database::pool::DbPool, RecordUsageRequest) { + let pool = db_setup::create_test_pool().await; + let client = pool.get().await.expect("fixture connection"); + let suffix = Uuid::new_v4().to_string(); + let user_id = Uuid::parse_str(MOCK_USER_ID).unwrap(); + let row = client + .query_one( + "WITH org AS ( + INSERT INTO organizations (name) VALUES ($1) RETURNING id + ), workspace AS ( + INSERT INTO workspaces (name, organization_id, created_by_user_id) + SELECT $1, org.id, $2 FROM org RETURNING id + ), api_key AS ( + INSERT INTO api_keys (key_hash, key_prefix, name, workspace_id, created_by_user_id) + SELECT lpad(replace(workspace.id::text, '-', ''), 64, '0'), 'sk-test', $1, workspace.id, $2 + FROM workspace RETURNING id + ) + SELECT org.id AS organization_id, workspace.id AS workspace_id, + api_key.id AS api_key_id, models.id AS model_id + FROM org, workspace, api_key, models WHERE models.model_name = $3", + &[&suffix, &user_id, &E2E_QWEN_MODEL_NAME], + ) + .await + .expect("test-owned usage fixture"); + let request = RecordUsageRequest { + organization_id: row.get("organization_id"), + workspace_id: row.get("workspace_id"), + api_key_id: row.get("api_key_id"), + model_id: row.get("model_id"), + model_name: E2E_QWEN_MODEL_NAME.to_string(), + input_tokens: 12, + output_tokens: 8, + input_cost: 12_000_000, + output_cost: 16_000_000, + total_cost: 28_000_000, + inference_type: services::usage::InferenceType::ChatCompletion + .as_str() + .to_string(), + ttft_ms: None, + avg_itl_ms: None, + inference_id: Some(Uuid::new_v4()), + provider_request_id: None, + stop_reason: None, + response_id: None, + image_count: None, + cache_read_tokens: 0, + cache_write_tokens: 0, + billing_details: None, + service_tier: None, + context_band: None, + served_provider_tier: None, + served_provider_type: None, + served_via_fallback: false, + }; + drop(client); + (pool, request) +} + +async fn assert_usage_recovers_without_double_charging(invalidation: UsageStatementInvalidation) { + let (pool, request) = usage_fixture().await; + let repository = OrganizationUsageRepository::new(pool.clone()); + let pool_size = pool.status().expect("active pool").max_size; + // More warm connections than retry_db!'s three attempts proves that + // recovery cannot rely on eventually acquiring a cold connection. + assert!(pool_size > 3); + let mut connections = Vec::new(); + for _ in 0..pool_size { + connections.push(pool.get().await.expect("reserve pool connection")); + } + for _ in 0..pool_size { + // Only this connection is available to the real repository. Retain + // it after warming, then release the next unvisited connection. + drop(connections.remove(0)); + let mut warm_request = request.clone(); + warm_request.inference_id = Some(Uuid::new_v4()); + assert!( + repository + .record_usage(warm_request) + .await + .unwrap() + .was_inserted + ); + let client = pool.get().await.expect("reacquire warmed connection"); + let cached_insert: bool = client + .query_one( + "SELECT EXISTS ( + SELECT 1 FROM pg_prepared_statements + WHERE statement LIKE '%INSERT INTO organization_usage_log%' + AND statement NOT LIKE '%pg_prepared_statements%' + )", + &[], + ) + .await + .expect("inspect warmed session") + .get(0); + assert!( + cached_insert, + "every session must hold the real usage INSERT" + ); + connections.push(client); + } + + let column = format!("usage_cache_{}", Uuid::new_v4().simple()); + match invalidation { + UsageStatementInvalidation::AddedColumn => { + connections[0] + .batch_execute(&format!( + "ALTER TABLE organization_usage_log ADD COLUMN {column} TEXT" + )) + .await + .expect("invalidate the cached INSERT RETURNING result type"); + } + UsageStatementInvalidation::DeallocatedBalance => { + for client in &connections { + // Leave the usage INSERT valid, so it succeeds before the + // balance UPSERT fails. Recovery must roll back that usage + // row and restart the entire transaction to charge once. + let balance_statement: String = client + .query_one( + "SELECT name FROM pg_prepared_statements WHERE statement LIKE $1", + &[&"%INSERT INTO organization_balance%"], + ) + .await + .expect("every warmed session must have the balance UPSERT") + .get(0); + let quoted_name = balance_statement.replace('"', "\"\""); + client + .batch_execute(&format!("DEALLOCATE \"{quoted_name}\"")) + .await + .expect("invalidate only the server's balance statement"); + } + } + } + drop(connections); + + // Return errors instead of panicking until after DDL cleanup, so a failed + // regression does not leave an unclassified column for encryption tests. + let outcome = async { + let recovered = repository + .record_usage(request.clone()) + .await + .context("record usage after invalidating every warm connection")?; + ensure!(recovered.was_inserted, "recovered usage must be inserted"); + let mut reserved = Vec::new(); + for _ in 1..pool_size { + reserved.push(pool.get().await?); + } + let mut duplicate_request = request.clone(); + duplicate_request.total_cost *= 2; + let duplicate = repository.record_usage(duplicate_request).await?; + ensure!(duplicate.id == recovered.id && !duplicate.was_inserted); + ensure!(duplicate.total_cost == request.total_cost); + + // Only one connection is free: a new write must succeed on the + // session whose duplicate transaction just rolled back. + let mut after_rollback = request.clone(); + after_rollback.inference_id = Some(Uuid::new_v4()); + ensure!(repository.record_usage(after_rollback).await?.was_inserted); + + let expected_requests = pool_size as i64 + 2; + let expected_spend = expected_requests * request.total_cost; + let expected_tokens = + expected_requests * i64::from(request.input_tokens + request.output_tokens); + let client = pool.get().await?; + let usage = client + .query_one( + "SELECT COUNT(*)::BIGINT, SUM(total_cost)::BIGINT, SUM(total_tokens)::BIGINT, + COUNT(*) FILTER (WHERE inference_id = $2)::BIGINT + FROM organization_usage_log WHERE organization_id = $1", + &[&request.organization_id, &request.inference_id], + ) + .await?; + ensure!( + usage.get::<_, i64>(0) == expected_requests, + "usage row count" + ); + ensure!(usage.get::<_, i64>(1) == expected_spend, "usage cost"); + ensure!(usage.get::<_, i64>(2) == expected_tokens, "usage tokens"); + ensure!( + usage.get::<_, i64>(3) == 1, + "recovered inference appears once" + ); + drop(client); + let balance = repository + .get_balance(request.organization_id) + .await? + .context("organization balance must exist")?; + ensure!( + balance.total_requests == expected_requests, + "balance requests" + ); + ensure!( + balance.total_spent == expected_spend, + "balance charged exactly once" + ); + ensure!(balance.total_tokens == expected_tokens, "balance tokens"); + Ok::<(), anyhow::Error>(()) + } + .await; + + let cleanup: Result<()> = async { + if matches!(invalidation, UsageStatementInvalidation::AddedColumn) { + pool.get() + .await? + .batch_execute(&format!( + "ALTER TABLE organization_usage_log DROP COLUMN IF EXISTS {column}" + )) + .await?; + } + Ok(()) + } + .await; + outcome.expect("usage transaction must recover without duplicate billing"); + cleanup.expect("remove temporary usage column"); +} + +#[tokio::test] +async fn usage_transaction_recovers_from_changed_result_type_without_double_charging() { + assert_usage_recovers_without_double_charging(UsageStatementInvalidation::AddedColumn).await; +} + +#[tokio::test] +async fn usage_transaction_recovers_from_missing_balance_statement_without_double_charging() { + assert_usage_recovers_without_double_charging(UsageStatementInvalidation::DeallocatedBalance) + .await; +} + +#[tokio::test] +async fn batched_signatures_reuse_a_prepared_statement() { + use services::attestation::{ports::AttestationRepository, ChatSignature, SignatureKind}; + + let (_server, database) = setup_test_server_with_database().await; + let pool = database.pool(); + // Keep one connection available so both repository calls and the session + // inspection below use the same PostgreSQL prepared-statement cache. + let mut held_connections = Vec::new(); + for _ in 1..pool.status().expect("database pool").max_size { + held_connections.push(pool.get().await.expect("hold another connection")); + } + + let chat_id = format!("batch-cache-{}", uuid::Uuid::new_v4()); + let signatures = |text: &str| { + ["ecdsa", "ed25519"] + .into_iter() + .map(|algo| ChatSignature { + text: text.to_string(), + signature: format!("synthetic-{algo}"), + signing_address: "synthetic-address".to_string(), + signing_algo: algo.to_string(), + signature_kind: Some(SignatureKind::ProviderTee), + }) + .collect() + }; + let cached_names = || async { + pool.get() + .await + .expect("inspect the warmed connection") + .query( + "SELECT name FROM pg_prepared_statements WHERE statement LIKE $1 ORDER BY name", + &[&"INSERT INTO chat_signatures%"], + ) + .await + .expect("inspect prepared statements") + .into_iter() + .map(|row| row.get::<_, String>("name")) + .collect::>() + }; + + database + .attestation + .add_chat_signatures(&chat_id, signatures("first")) + .await + .expect("store both signatures"); + let first_names = cached_names().await; + assert_eq!(first_names.len(), 1, "the batch must be cached"); + + database + .attestation + .add_chat_signatures(&chat_id, signatures("updated")) + .await + .expect("update both signatures"); + assert_eq!( + cached_names().await, + first_names, + "reuse the same statement" + ); + for algo in ["ecdsa", "ed25519"] { + let stored = database + .attestation + .get_chat_signature(&chat_id, algo) + .await + .expect("read the updated signature"); + assert_eq!(stored.text, "updated"); + assert_eq!(stored.signature_kind, Some(SignatureKind::ProviderTee)); + } +} diff --git a/crates/database/src/repositories/api_key.rs b/crates/database/src/repositories/api_key.rs index d587eede9..e6b885087 100644 --- a/crates/database/src/repositories/api_key.rs +++ b/crates/database/src/repositories/api_key.rs @@ -1,5 +1,6 @@ use crate::models::ApiKey; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -149,7 +150,7 @@ impl ApiKeyRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT ak.* FROM api_keys ak @@ -192,7 +193,7 @@ impl ApiKeyRepository { .map_err(RepositoryError::PoolError)?; client - .execute( + .cached_execute( "UPDATE api_keys SET last_used_at = NOW() WHERE id = $1", &[&id], ) diff --git a/crates/database/src/repositories/attestation.rs b/crates/database/src/repositories/attestation.rs index e25197334..f975154d0 100644 --- a/crates/database/src/repositories/attestation.rs +++ b/crates/database/src/repositories/attestation.rs @@ -1,3 +1,4 @@ +use crate::repositories::statement_cache::CachedStatements; use async_trait::async_trait; use services::attestation::{ ports::AttestationRepository, AttestationError, ChatSignature, SignatureKind, @@ -48,6 +49,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( @@ -62,7 +84,7 @@ impl AttestationRepository for PgAttestationRepository { .map_err(|e| AttestationError::RepositoryError(e.to_string()))?; let signature_kind = signature.signature_kind.map(|kind| kind.as_str()); client - .execute( + .cached_execute( "INSERT INTO chat_signatures (chat_id, text, signature, signing_address, signing_algo, signature_kind) VALUES ($1, $2, $3, $4, $5, $6) 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()", &[&chat_id, &signature.text, &signature.signature, &signature.signing_address, &signature.signing_algo, &signature_kind], ) @@ -72,6 +94,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 + .cached_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 +187,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/database/src/repositories/mod.rs b/crates/database/src/repositories/mod.rs index f5411510a..37db640ad 100644 --- a/crates/database/src/repositories/mod.rs +++ b/crates/database/src/repositories/mod.rs @@ -33,6 +33,7 @@ pub mod retry; pub mod service; pub mod service_usage_repository_impl; pub mod session; +pub mod statement_cache; pub mod usage_repository_impl; pub mod user; pub mod utils; diff --git a/crates/database/src/repositories/model.rs b/crates/database/src/repositories/model.rs index 3b41eb40d..be6428cae 100644 --- a/crates/database/src/repositories/model.rs +++ b/crates/database/src/repositories/model.rs @@ -1,6 +1,7 @@ use crate::constants::DEFAULT_MODEL_OWNED_BY; use crate::models::{Model, ModelHistory, UpdateModelPricingRequest}; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -269,7 +270,7 @@ impl ModelRepository { .map_err(RepositoryError::PoolError)?; client - .query( + .cached_query( r#" SELECT id, model_name, model_display_name, model_description, model_icon, @@ -1301,7 +1302,7 @@ impl ModelRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT m.id, diff --git a/crates/database/src/repositories/organization_limits.rs b/crates/database/src/repositories/organization_limits.rs index f56c52e94..3bf05ed60 100644 --- a/crates/database/src/repositories/organization_limits.rs +++ b/crates/database/src/repositories/organization_limits.rs @@ -1,5 +1,6 @@ use crate::models::{OrganizationLimitsHistory, UpdateOrganizationLimitsDbRequest}; use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -122,7 +123,7 @@ impl OrganizationLimitsRepository { .map_err(RepositoryError::PoolError)?; client - .query( + .cached_query( r#" SELECT id, organization_id, spend_limit, credit_type, source, currency, effective_from, effective_until, diff --git a/crates/database/src/repositories/organization_staking_farm_sources.rs b/crates/database/src/repositories/organization_staking_farm_sources.rs index edace0482..817d7e333 100644 --- a/crates/database/src/repositories/organization_staking_farm_sources.rs +++ b/crates/database/src/repositories/organization_staking_farm_sources.rs @@ -1,4 +1,5 @@ use crate::pool::DbPool; +use crate::repositories::statement_cache::CachedStatements; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -165,7 +166,7 @@ impl StakingFarmRepository for OrganizationStakingFarmSourcesRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT id, organization_id, near_account_id, network_id, contract_id, farm_product_id, farm_price_id, credit_nano_usd_per_reward_unit, diff --git a/crates/database/src/repositories/organization_usage.rs b/crates/database/src/repositories/organization_usage.rs index ae3cba21f..1e0001297 100644 --- a/crates/database/src/repositories/organization_usage.rs +++ b/crates/database/src/repositories/organization_usage.rs @@ -3,6 +3,9 @@ use crate::models::{ ServedProviderType, StopReason, }; use crate::pool::DbPool; +use crate::repositories::statement_cache::{ + invalidate_pool_statement_caches, is_stale_statement, CachedStatements, +}; use crate::repositories::utils::map_db_error; use crate::retry_db; use anyhow::{Context, Result}; @@ -47,7 +50,7 @@ impl OrganizationUsageRepository { .map_err(RepositoryError::PoolError)?; client - .query_one( + .cached_query_one( r#" SELECT COALESCE(SUM(total_cost), 0)::BIGINT as total_spend FROM organization_usage_log @@ -77,6 +80,16 @@ impl OrganizationUsageRepository { .context("Failed to get database connection") .map_err(RepositoryError::PoolError)?; + // A stale statement inside the transaction aborts it; evict the + // statement from every pooled connection so the retry does not + // trip over another stale copy, then let `retry_db!` restart. + let pool = deadpool_postgres::Client::pool(&client); + let map_tx_error = |err: tokio_postgres::Error| { + if is_stale_statement(&err) { + invalidate_pool_statement_caches(pool.as_ref()); + } + map_db_error(err) + }; let transaction = client.transaction().await.map_err(map_db_error)?; let id = Uuid::new_v4(); @@ -93,7 +106,7 @@ impl OrganizationUsageRepository { .served_provider_type .map(|provider| provider.as_str()); let maybe_row = transaction - .query_opt( + .cached_query_opt( r#" INSERT INTO organization_usage_log ( id, organization_id, workspace_id, api_key_id, @@ -140,13 +153,13 @@ impl OrganizationUsageRepository { ], ) .await - .map_err(map_db_error)?; + .map_err(map_tx_error)?; let (row, was_inserted) = match maybe_row { Some(row) => { // New insert succeeded — update organization balance transaction - .execute( + .cached_execute( r#" INSERT INTO organization_balance ( organization_id, @@ -172,7 +185,7 @@ impl OrganizationUsageRepository { ], ) .await - .map_err(map_db_error)?; + .map_err(map_tx_error)?; transaction.commit().await.map_err(map_db_error)?; (row, true) @@ -188,7 +201,7 @@ impl OrganizationUsageRepository { ); let existing = client - .query_one( + .cached_query_one( r#" SELECT * FROM organization_usage_log @@ -220,7 +233,7 @@ impl OrganizationUsageRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT organization_id, total_spent, last_usage_at, total_requests, total_tokens, updated_at diff --git a/crates/database/src/repositories/statement_cache.rs b/crates/database/src/repositories/statement_cache.rs new file mode 100644 index 000000000..4612e63d5 --- /dev/null +++ b/crates/database/src/repositories/statement_cache.rs @@ -0,0 +1,187 @@ +//! Per-connection prepared-statement reuse for hot-path queries. +//! +//! `tokio_postgres` prepares a raw SQL string on every call, so each query is +//! two network round trips (Parse/Describe, then Bind/Execute). deadpool keeps +//! a per-connection statement cache; going through it makes a repeated query +//! one round trip. On a replica far from the database this is the difference +//! between ~60 ms and ~120 ms per query. +//! +//! A cached statement can outlive the server-side object it points at: +//! +//! - after a migration adds a column to a table read with `SELECT *`, Postgres +//! refuses the old plan with SQLSTATE 0A000 "cached plan must not change +//! result type" (a replica still on the previous release during a rolling +//! deploy hits this on every warm connection); +//! - a statement can be missing on the server altogether, SQLSTATE 26000, +//! if the connection was reset underneath the cache. +//! +//! Both are handled the same way: the stale entry is evicted from every +//! connection's cache in the pool (a schema change invalidates all of them +//! at once), the statement is re-prepared and re-run on the same connection +//! when no transaction is involved, and [`map_db_error`](super::utils::map_db_error) +//! classifies the error as retryable so `retry_db!` re-runs transactional +//! blocks with a fresh statement. + +use async_trait::async_trait; +use deadpool_postgres::{Client, Pool, Transaction}; +use tokio_postgres::{error::SqlState, types::ToSql, Error, Row}; + +/// PostgreSQL routine that raises the stale-plan error; stable across server +/// locales, unlike the message text. +const STALE_PLAN_ROUTINE: &str = "RevalidateCachedQuery"; +const STALE_PLAN_MESSAGE: &str = "cached plan must not change result type"; + +/// True when `err` is Postgres refusing a prepared statement whose result +/// columns changed since it was prepared. +pub fn is_stale_cached_plan(err: &Error) -> bool { + err.as_db_error().is_some_and(|db_err| { + db_err.code() == &SqlState::FEATURE_NOT_SUPPORTED + && (db_err.routine() == Some(STALE_PLAN_ROUTINE) + || db_err.message().contains(STALE_PLAN_MESSAGE)) + }) +} + +/// True when the server no longer has the prepared statement the cache +/// handed out. +pub fn is_missing_prepared_statement(err: &Error) -> bool { + err.as_db_error() + .is_some_and(|db_err| db_err.code() == &SqlState::INVALID_SQL_STATEMENT_NAME) +} + +/// True for either way a cached statement handle can go bad. +pub fn is_stale_statement(err: &Error) -> bool { + is_stale_cached_plan(err) || is_missing_prepared_statement(err) +} + +/// Evict every cached statement on every connection of `pool`. A schema +/// change invalidates the same statement on all warm connections, so a +/// per-connection eviction would leave the next retry to trip over another +/// stale copy. +pub fn invalidate_pool_statement_caches(pool: Option<&Pool>) { + if let Some(pool) = pool { + pool.manager().statement_caches.clear(); + } +} + +/// Run queries through the connection's prepared-statement cache. +#[async_trait] +pub trait CachedStatements { + async fn cached_query( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error>; + + async fn cached_query_opt( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error>; + + async fn cached_query_one( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result; + + async fn cached_execute(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) + -> Result; +} + +/// Runs a cached statement, and on a stale handle evicts it pool-wide and +/// re-runs once on the same connection (outside a transaction only). +macro_rules! run_cached { + ($self:ident, $method:ident, $sql:expr, $params:expr) => {{ + let statement = $self.prepare_cached($sql).await?; + match $self.$method(&statement, $params).await { + Err(err) if is_stale_statement(&err) => { + $self.evict_stale($sql); + if !$self.can_retry_in_place() { + return Err(err); + } + let statement = $self.prepare_cached($sql).await?; + $self.$method(&statement, $params).await + } + result => result, + } + }}; +} + +macro_rules! impl_cached_statements { + ($ty:ty) => { + #[async_trait] + impl CachedStatements for $ty { + async fn cached_query( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error> { + run_cached!(self, query, sql, params) + } + + async fn cached_query_opt( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result, Error> { + run_cached!(self, query_opt, sql, params) + } + + async fn cached_query_one( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result { + run_cached!(self, query_one, sql, params) + } + + async fn cached_execute( + &self, + sql: &str, + params: &[&(dyn ToSql + Sync)], + ) -> Result { + run_cached!(self, execute, sql, params) + } + } + }; +} + +trait StaleStatementCleanup { + /// Drop stale statements after the server rejected `sql`. + fn evict_stale(&self, sql: &str); + /// Whether the statement can be re-prepared and re-run on this + /// connection right away. + fn can_retry_in_place(&self) -> bool; +} + +impl StaleStatementCleanup for Client { + fn evict_stale(&self, _sql: &str) { + tracing::warn!("Dropping cached prepared statements after a schema change; retrying"); + match Client::pool(self) { + Some(pool) => invalidate_pool_statement_caches(Some(&pool)), + None => self.statement_cache.clear(), + } + } + + fn can_retry_in_place(&self) -> bool { + true + } +} + +impl StaleStatementCleanup for Transaction<'_> { + fn evict_stale(&self, _sql: &str) { + tracing::warn!("Dropping cached prepared statements after a schema change; retrying"); + // The pool is not reachable from a transaction; callers that hold + // the client evict pool-wide through `invalidate_pool_statement_caches`. + self.statement_cache.clear(); + } + + fn can_retry_in_place(&self) -> bool { + // The transaction is aborted after the error; the surrounding + // `retry_db!` restarts it with a fresh statement. + false + } +} + +impl_cached_statements!(Client); +impl_cached_statements!(Transaction<'_>); diff --git a/crates/database/src/repositories/utils.rs b/crates/database/src/repositories/utils.rs index 6bf50cd0b..20763849f 100644 --- a/crates/database/src/repositories/utils.rs +++ b/crates/database/src/repositories/utils.rs @@ -41,6 +41,17 @@ pub fn map_db_error(err: tokio_postgres::Error) -> RepositoryError { &SqlState::QUERY_CANCELED => RepositoryError::QueryTimeout, + // A cached prepared statement outlived the server-side object it + // points at (schema change during a rolling deploy, or a reset + // connection). The caches have already been evicted (see + // `statement_cache`); classify as retryable so `retry_db!` + // re-prepares and re-runs it. + _ if super::statement_cache::is_stale_statement(&err) => { + RepositoryError::ConnectionFailed( + "cached prepared statement is no longer valid".to_string(), + ) + } + // Default case - wrap in generic database error _ => RepositoryError::DatabaseError(anyhow::anyhow!( "Database error ({}): {}", diff --git a/crates/database/src/repositories/workspace.rs b/crates/database/src/repositories/workspace.rs index 92f24b1ab..6302b2e76 100644 --- a/crates/database/src/repositories/workspace.rs +++ b/crates/database/src/repositories/workspace.rs @@ -1,3 +1,4 @@ +use crate::repositories::statement_cache::CachedStatements; use crate::{ models::{CreateWorkspaceRequest, UpdateWorkspaceRequest, Workspace}, pool::DbPool, @@ -380,7 +381,7 @@ impl WorkspaceRepository { .map_err(RepositoryError::PoolError)?; client - .query_opt( + .cached_query_opt( r#" SELECT w.*, 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