diff --git a/crates/runtime/tests/chat_sql_budget.rs b/crates/runtime/tests/chat_sql_budget.rs index b578e905ff..4fd37c8030 100644 --- a/crates/runtime/tests/chat_sql_budget.rs +++ b/crates/runtime/tests/chat_sql_budget.rs @@ -12,6 +12,65 @@ use sha2::{Digest, Sha256}; use std::sync::{Arc, Mutex}; use tracing_subscriber::prelude::*; +/// The run drain does not own the detached observer. Its completion counter is +/// published only after the production worker has finished all of its SQL. +async fn wait_for_chat_observers(state: &astra_runtime::AppState, expected_observers: u64) { + let registry = state.metrics_registry(); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + loop { + let metrics = registry.render_prometheus(); + let mut completed = 0; + for line in metrics.lines().filter(|line| { + line.starts_with("astra_turn_observer_runs_total{") + || line.starts_with("astra_turn_observer_dispatches_total{") + }) { + let (series, value) = line.rsplit_once(' ').unwrap(); + let value: u64 = value.parse().unwrap(); + if value == 0 { + continue; + } + assert!( + !series.contains("outcome=\"error\"") + && !series.contains("outcome=\"dropped_full\""), + "observer must finish without errors or dropped work: {line}" + ); + if series.starts_with("astra_turn_observer_runs_total{") { + completed += value; + } + } + assert!(completed <= expected_observers, "unexpected extra observer"); + if completed == expected_observers { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("each chat must finish its detached observer before ending SQL capture"); +} + +/// Measure request SQL on an initialized pool, not timing-dependent SQLx +/// connection setup. Every statement after this boundary still counts. +async fn prime_chat_pool(ctx: &MatrixE2eCtx) { + let pool = ctx.shared_pool.get(); + // This serial probe needs only a small warm reserve, not the deployment's + // full connection limit. Additional connections opened during a request + // remain visible to the raw SQL budget. + let reserve = pool.options().get_max_connections().min(4); + tokio::time::timeout(std::time::Duration::from_secs(10), async { + let mut connections = Vec::new(); + for _ in 0..reserve { + connections.push(pool.acquire().await.unwrap()); + } + drop(connections); + while pool.num_idle() < reserve as usize { + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + }) + .await + .expect("initialize and return the chat-pool reserve before SQL capture"); +} + #[derive(Default)] struct Statement(String, f64); @@ -105,6 +164,23 @@ impl tracing_subscriber::Layer for SqlCapture { const CONTROL_OBSERVATION_SQL: &str = "SELECT requested.ordinal, run_id, session_id, status, waiting_for, parent_run_id, ancestor_path, run_generation, last_event_idx, CAST(cancellation_requested_at IS NOT NULL AS SIGNED) AS cancellation_requested, CAST(owner_pod_id IS NOT NULL AND owner_lease_expires_at IS NOT NULL AND owner_lease_expires_at >= NOW(6) AS SIGNED) AS owner_lease_live FROM (VALUES ROW(0,?,?)) requested(ordinal, requested_user_id, requested_run_id) LEFT JOIN agent_runs ON user_id = requested_user_id AND run_id = requested_run_id AND run_id IN (?)"; +// The account/session snapshot is cached for 30 seconds. A slow fixture can +// cross that boundary independently of readonly tool rounds. Attribute only +// this exact query, at most once per request, while retaining the raw total. +const AUTHENTICATION_SNAPSHOT_SQL: &str = "SELECT u.user_id, u.username, u.email, u.display_name, u.is_active, CAST(UNIX_TIMESTAMP(s.session_expires_at) AS SIGNED) AS session_expires_at_unix, m.provider_id AS memoria_provider_id, m.external_subject AS memoria_subject FROM (SELECT MAX(expires_at) AS session_expires_at FROM auth_refresh_tokens WHERE user_id = ? AND session_id = ? AND is_revoked = 0 AND expires_at > NOW()) s LEFT JOIN auth_users u ON u.user_id = ? LEFT JOIN (SELECT astra_user_id, provider_id, external_subject FROM auth_external_identities WHERE astra_user_id = ? AND provider_id LIKE 'memoria:%' ORDER BY CASE WHEN provider_id = ? THEN 0 ELSE 1 END, updated_at DESC, provider_id ASC LIMIT 1) m ON m.astra_user_id = u.user_id"; + +fn authentication_snapshot_count(statements: &[String]) -> usize { + let count = statements + .iter() + .filter(|sql| sql.as_str() == AUTHENTICATION_SNAPSHOT_SQL) + .count(); + assert!( + count <= 1, + "a request must not repeat its authentication snapshot" + ); + count +} + fn answer(message: Value, reason: &str) -> ProviderResponse { ProviderResponse::OpenAi(json!({ "choices": [{"index": 0, "message": message, "finish_reason": reason}], @@ -422,7 +498,11 @@ async fn ordinary_chat_and_readonly_rounds_have_bounded_total_sql() { "{body}" ); assert!(body.contains("Warm context ready."), "{body}"); + // Shutdown drain would reject the measured turn. The complete SSE + // response settles the warm-up turn; separately await its observer. + wait_for_chat_observers(&b.ctx.app_state, 1).await; } + prime_chat_pool(&b.ctx).await; let message = if executes_agent { "List agent status." } else if readonly_rounds == 0 { @@ -515,6 +595,7 @@ async fn ordinary_chat_and_readonly_rounds_have_bounded_total_sql() { .drain_background_runs(std::time::Duration::from_secs(5)) .await ); + wait_for_chat_observers(&b.ctx.app_state, 1 + u64::from(warm_context)).await; let samples = capture.0.lock().unwrap().take().unwrap(); let execution_ms = execution_started.elapsed().as_secs_f64() * 1000.0; // Explicit cold-replay verification follows the ordinary request's @@ -773,6 +854,7 @@ async fn ordinary_chat_and_readonly_rounds_have_bounded_total_sql() { statements.len(), statements.join("\n") ); + let authentication_reads = authentication_snapshot_count(statements); if *scenario == "ordinary" { continue; } @@ -780,13 +862,15 @@ async fn ordinary_chat_and_readonly_rounds_have_bounded_total_sql() { .iter() .find(|(name, _, warm, _, _)| *name == "ordinary" && warm == warm_context) .unwrap(); - // Explain only the exact observed periodic reads; the independent - // raw total budget above never deducts a query. - let expected = baseline.4.len() as i64 + *polls as i64 - baseline.3 as i64; + // Explain exact periodic reads and at most one request-auth cache + // miss. The independent raw total budget above never deducts a query. + let expected = baseline.4.len() as i64 + *polls as i64 - baseline.3 as i64 + + authentication_reads as i64 + - authentication_snapshot_count(&baseline.4) as i64; assert_eq!( statements.len() as i64, expected, - "{scenario} {rounds} readonly rounds (warm={warm_context}, polls={polls}) must add no SQL beyond the exact observed periodic controls" + "{scenario} {rounds} readonly rounds (warm={warm_context}, polls={polls}) must add no SQL beyond exact periodic controls and one request-auth cache miss" ); } } diff --git a/docs/guides/testing.md b/docs/guides/testing.md index d6a343f487..3aae13d488 100644 --- a/docs/guides/testing.md +++ b/docs/guides/testing.md @@ -493,6 +493,33 @@ Optionally set **`ASTRA_AUTO_CREATE_DATABASE=1`** so the first `ensure_core_schema` (server or online tests) runs `CREATE DATABASE IF NOT EXISTS` for that effective name (bootstrap catalog defaults to `mysql`). +### Ordinary-chat SQL budget + +`chat_sql_budget::ordinary_chat_and_readonly_rounds_have_bounded_total_sql` +uses the same isolated MatrixOne setup and scripted loopback provider as the +system HTTP journeys. Run this test by exact name in its own process because +it owns the global SQL tracing subscriber: + +```bash +cargo test -p astra-runtime --features e2e-hooks --test chat_sql_budget \ + ordinary_chat_and_readonly_rounds_have_bounded_total_sql -- --ignored --exact --nocapture +``` + +The 20-statement limit measures each request on an initialized server pool, +including transactions, periodic control reads and the real observer's SQL. +Pool initialization happens before capture; no statements are filtered out +inside the measured interval. Warm-context setup consumes its complete SSE +response and waits for observer completion before the next request starts; +it must not use shutdown drain, which closes admission for that next request. +Capture ends only after +both the measured run and its observer finish: the detached observer is not +owned by `drain_background_runs`. Observer errors, dropped work or a missing +completion fail the probe instead of silently undercounting SQL. +The readonly-round comparison accounts for exact periodic control reads and +at most one account/session authentication snapshot when its 30-second cache +expires. Those statements still count toward the raw 20-statement limit; +repeated authentication snapshots or any other extra readonly-round SQL fail. + ### Durable provider concurrency and cancellation The ignored runtime test