test(runtime): drain observers around chat SQL budget capture

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<S: tracing::Subscriber> tracing_subscriber::Layer<S> 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,20 +854,23 @@ 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;
         }
         let baseline = budgets
             .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"
         );
     }
 }

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