Skip to main content

firstpass_proxy/
proxy.rs

1//! Axum wiring: routes, request/response shapes, and observe-mode trace construction
2//! (SPEC §7.1, §7.1a — forward unchanged, record asynchronously, zero added latency).
3
4use std::sync::Arc;
5use std::time::Instant;
6
7use axum::body::Body;
8use axum::extract::{Request, State};
9use axum::http::HeaderMap;
10use axum::middleware::Next;
11use axum::response::{IntoResponse, Response};
12use axum::routing::{get, post};
13use axum::{Extension, Json, Router};
14use bytes::Bytes;
15use firstpass_core::features::{hour_bucket, token_bucket};
16use firstpass_core::hashchain::sha256_hex;
17use firstpass_core::{
18    Attempt, DeferredVerdict, FEATURE_VERSION, Features, FinalOutcome, GENESIS_HASH, Mode,
19    PolicyRef, RequestInfo, Score, ServedFrom, TaskKind, Trace, Verdict,
20};
21use serde::Deserialize;
22use serde_json::Value;
23use tokio::sync::mpsc::error::TrySendError;
24use uuid::Uuid;
25
26use crate::config::ProxyConfig;
27use crate::error::ProxyError;
28use crate::gate::{GateHealthRegistry, resolve_gates};
29use crate::provider::{Auth, ChatMessage, ModelRequest, ModelResponse, ProviderRegistry};
30use crate::router::{EnforceCtx, EngineOutcome, route_enforce};
31use crate::store;
32use crate::tenant_auth::{TenantId, auth_middleware};
33use crate::upstream::{forward_anthropic, forward_anthropic_streaming};
34use firstpass_core::Route;
35
36/// Shared state handed to every request handler. Cheap to clone: an `Arc`ed config, a
37/// pooled HTTP client, and a bounded channel sender.
38#[derive(Clone)]
39pub struct AppState {
40    /// Static proxy configuration.
41    pub config: Arc<ProxyConfig>,
42    /// Shared, connection-pooled HTTP client used to call upstream (observe passthrough).
43    pub http: reqwest::Client,
44    /// Multi-provider registry used by the enforce-mode escalation engine.
45    pub providers: ProviderRegistry,
46    /// Per-gate error budgets (auto-disable), shared across requests.
47    pub gate_health: Arc<GateHealthRegistry>,
48    /// Fire-and-forget sender to the background trace writer.
49    pub traces: store::TraceSender,
50    /// Optional online/adaptive conformal serve threshold (Gibbs-Candès ACI). `None` = fixed
51    /// `serve_threshold` from config (default). When present, `/v1/feedback` nudges it live and the
52    /// enforce path reads its current value per request — the reactive, self-tuning loop.
53    pub adaptive: Option<Arc<std::sync::Mutex<firstpass_core::conformal::AdaptiveConformal>>>,
54    /// Per-tenant request rate limiter (ADR 0004 §D6). `None` (the default) disables rate
55    /// limiting entirely — set via [`build_tenant_rate_limiter`] from
56    /// [`ProxyConfig::tenant_rate_per_sec`].
57    pub tenant_rate_limiter: Option<Arc<governor::DefaultKeyedRateLimiter<String>>>,
58}
59
60/// Build the per-tenant keyed rate limiter from config (ADR 0004 §D6). Returns `None` when
61/// `FIRSTPASS_TENANT_RATE_PER_SEC` is unset (the default) — single-operator and existing
62/// deployments see no limiter and no behavior change.
63#[must_use]
64pub fn build_tenant_rate_limiter(
65    config: &ProxyConfig,
66) -> Option<Arc<governor::DefaultKeyedRateLimiter<String>>> {
67    let per_sec = config.tenant_rate_per_sec?;
68    Some(Arc::new(governor::RateLimiter::keyed(
69        governor::Quota::per_second(per_sec),
70    )))
71}
72
73/// Axum middleware (ADR 0004 §D6): enforce the per-tenant request rate limit. Must run AFTER
74/// [`auth_middleware`] so the resolved [`TenantId`] is already in request extensions. A no-op
75/// (never returns 429) when [`AppState::tenant_rate_limiter`] is `None`.
76pub async fn tenant_rate_limit_middleware(
77    State(state): State<AppState>,
78    Extension(tenant): Extension<TenantId>,
79    req: Request,
80    next: Next,
81) -> Response {
82    if let Some(limiter) = &state.tenant_rate_limiter
83        && limiter.check_key(&tenant.0).is_err()
84    {
85        return ProxyError::RateLimited.into_response();
86    }
87    next.run(req).await
88}
89
90impl std::fmt::Debug for AppState {
91    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92        f.debug_struct("AppState")
93            .field("config", &self.config)
94            .finish_non_exhaustive()
95    }
96}
97
98/// Fire-and-forget a trace at the background writer: non-blocking, and bounded. If the writer has
99/// fallen behind enough to fill the buffer, or is gone, the trace is dropped with a warning rather
100/// than blocking the hot path or growing memory without limit (the audit chain over persisted
101/// traces stays valid; a dropped trace is simply absent).
102fn offer_trace(traces: &store::TraceSender, trace: Trace) {
103    record_trace_metrics(&trace);
104    match traces.try_send(trace) {
105        Ok(()) => {}
106        Err(TrySendError::Full(_)) => {
107            tracing::warn!("trace channel full; dropping trace (writer behind under load)");
108            metrics::counter!("firstpass_traces_dropped_total").increment(1);
109        }
110        Err(TrySendError::Closed(_)) => {
111            tracing::warn!("trace writer is gone; dropping trace");
112        }
113    }
114}
115
116/// Record the real signals every trace carries: enforce-mode latency/escalations (observe mode
117/// forwards unchanged, so its wall-clock time isn't a routing-decision latency), and what got
118/// served — regardless of mode, since an upstream failure is worth counting either way.
119fn record_trace_metrics(trace: &Trace) {
120    if trace.mode == Mode::Enforce {
121        metrics::histogram!("firstpass_enforce_latency_ms")
122            .record(trace.final_.total_latency_ms as f64);
123        if trace.final_.escalations > 0 {
124            metrics::counter!("firstpass_escalations_total")
125                .increment(u64::from(trace.final_.escalations));
126        }
127    }
128    let served_from = match trace.final_.served_from {
129        ServedFrom::Attempt => "attempt",
130        ServedFrom::BestAttempt => "best_attempt",
131        ServedFrom::Error => "error",
132    };
133    metrics::counter!("firstpass_served_total", "served_from" => served_from).increment(1);
134    if trace.final_.served_from == ServedFrom::Error {
135        metrics::counter!("firstpass_upstream_failures_total").increment(1);
136    }
137}
138
139/// Max accepted request body. Explicit (not axum's ~2 MB default) so it's an intentional ceiling:
140/// generous enough to pass through large multimodal/long-context requests, bounded so a single
141/// oversized body can't exhaust memory.
142const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
143
144/// Build the axum router: `POST /v1/messages`, `GET /v1/capabilities`, `GET /healthz`,
145/// `GET /metrics`.
146///
147/// # Errors
148/// [`ProxyError::Internal`] if the Prometheus recorder fails to install (see
149/// [`crate::metrics::install`]).
150pub fn app(state: AppState) -> Result<Router, ProxyError> {
151    crate::metrics::install()?;
152    let max_concurrency = state.config.max_concurrency;
153
154    // Tenant-facing business routes: every one runs the auth middleware, which injects the resolved
155    // `TenantId` into request extensions (the authenticated tenant when `require_auth` is on, the
156    // static default when off — ADR 0004 §D1/§D2). Operator routes (`/healthz`, `/metrics`) are
157    // NOT tenant-facing and stay outside the auth layer.
158    // Per-tenant rate limit (ADR 0004 §D6) runs INSIDE (after) the auth layer below — axum layers
159    // wrap outward-in, so a layer added earlier in the chain executes later on the request path —
160    // so the resolved `TenantId` is already in extensions when this middleware checks it. A no-op
161    // when `FIRSTPASS_TENANT_RATE_PER_SEC` is unset.
162    let business = Router::new()
163        .route("/v1/messages", post(messages))
164        .route("/v1/feedback", post(feedback))
165        .route("/v1/capabilities", get(capabilities))
166        .layer(axum::middleware::from_fn_with_state(
167            state.clone(),
168            tenant_rate_limit_middleware,
169        ))
170        .layer(axum::middleware::from_fn_with_state(
171            state.clone(),
172            auth_middleware,
173        ));
174
175    Ok(Router::new()
176        .merge(business)
177        .route("/healthz", get(healthz))
178        .route("/metrics", get(crate::metrics::handler))
179        // Explicit body-size ceiling (DoS/OOM guard) across every route.
180        .layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES))
181        // Concurrency load-shed: cap in-flight requests under the cap rather than falling over.
182        // Deliberately NOT a request timeout — that would sever in-flight SSE streams.
183        .layer(tower::limit::GlobalConcurrencyLimitLayer::new(
184            max_concurrency,
185        ))
186        .with_state(state))
187}
188
189/// `GET /healthz` — liveness probe.
190async fn healthz() -> impl IntoResponse {
191    Json(serde_json::json!({ "status": "ok" }))
192}
193
194/// `GET /v1/capabilities` — agent-first discovery (SPEC §0.2, §7.4): what this proxy speaks,
195/// which modes are live, the first enforce route's ladder/gates, and how to turn it off.
196async fn capabilities(State(state): State<AppState>) -> impl IntoResponse {
197    // Report the first enforce route's ladder + gates, so an agent can discover what it's routed
198    // through. Empty when no routing config is loaded (pure observe deployment).
199    let (ladder, gates) = state
200        .config
201        .routing
202        .as_ref()
203        .and_then(|c| c.routes.iter().find(|r| r.mode == Mode::Enforce))
204        .map(|r| (r.ladder.clone(), r.gates.clone()))
205        .unwrap_or_default();
206    Json(serde_json::json!({
207        "service": "firstpass",
208        "version": env!("CARGO_PKG_VERSION"),
209        "feature_version": FEATURE_VERSION,
210        "modes": ["observe", "enforce"],
211        "wire_apis": ["anthropic.messages"],
212        "ladder": ladder,
213        "gates": gates,
214        "feedback_api": "POST /v1/feedback",
215        "offboarding": "unset ANTHROPIC_BASE_URL",
216    }))
217}
218
219/// Body of `POST /v1/feedback`: a downstream outcome reported for a past decision.
220#[derive(Debug, Deserialize)]
221struct FeedbackRequest {
222    /// The `trace_id` of the decision this outcome is about.
223    trace_id: String,
224    /// The gate/source id, e.g. `"tests"` or `"feedback:ci"`.
225    gate_id: String,
226    /// `"pass"` | `"fail"` | `"abstain"`.
227    verdict: String,
228    /// Optional confidence in `[0, 1]`.
229    #[serde(default)]
230    score: Option<f64>,
231    /// Who reported it (a CI system, a human reviewer, a deferred gate).
232    reporter: String,
233}
234
235/// `POST /v1/feedback` — attach a downstream outcome (deferred verdict) to a past trace, closing
236/// the outcome-feedback loop (SPEC §8.3.4). The verdict is stored in a **separate** table keyed
237/// by `trace_id`; the sealed, hashed trace is never mutated, so the audit chain stays verifiable.
238/// Returns `202 Accepted`. This is the signal that later calibrates the gates.
239async fn feedback(
240    State(state): State<AppState>,
241    Extension(TenantId(tenant)): Extension<TenantId>,
242    body: Bytes,
243) -> Response {
244    let req: FeedbackRequest = match serde_json::from_slice(&body) {
245        Ok(r) => r,
246        Err(e) => {
247            return ProxyError::BadRequest(format!("invalid feedback body: {e}")).into_response();
248        }
249    };
250    let verdict = match req.verdict.as_str() {
251        "pass" => Verdict::Pass,
252        "fail" => Verdict::Fail,
253        "abstain" => Verdict::Abstain,
254        other => {
255            return ProxyError::BadRequest(format!("unknown verdict {other:?}")).into_response();
256        }
257    };
258    let score = match req.score {
259        Some(s) => match Score::new(s) {
260            Ok(sc) => Some(sc),
261            Err(_) => {
262                return ProxyError::BadRequest(format!("score {s} out of range [0,1]"))
263                    .into_response();
264            }
265        },
266        None => None,
267    };
268
269    let db = state.config.db_path.clone();
270
271    // Reject feedback for an unknown trace, so orphan outcomes can't accumulate — AND deny
272    // cross-tenant feedback (IDOR, ADR 0004 §D4). `trace_exists` is scoped to the caller's tenant,
273    // so a trace owned by another tenant is indistinguishable from a missing one: both return `404`
274    // (never `403`, which would be an existence oracle).
275    let (db_check, tenant_check, tid_check) = (db.clone(), tenant.clone(), req.trace_id.clone());
276    match tokio::task::spawn_blocking(move || {
277        store::trace_exists(&db_check, &tenant_check, &tid_check)
278    })
279    .await
280    {
281        Ok(Ok(true)) => {}
282        Ok(Ok(false)) => {
283            return ProxyError::NotFound(format!("unknown trace_id {:?}", req.trace_id))
284                .into_response();
285        }
286        Ok(Err(e)) => {
287            tracing::error!(%e, "feedback: trace_exists check failed");
288            return ProxyError::Internal(e.to_string()).into_response();
289        }
290        Err(e) => {
291            tracing::error!(%e, "feedback: trace_exists task panicked");
292            return ProxyError::Internal(e.to_string()).into_response();
293        }
294    }
295
296    // Correctness signal for the online adaptive loop — only a clear Pass/Fail nudges the threshold.
297    let feedback_signal = match verdict {
298        Verdict::Pass => Some(true),
299        Verdict::Fail => Some(false),
300        Verdict::Abstain => None,
301    };
302    let dv = DeferredVerdict {
303        gate_id: req.gate_id,
304        verdict,
305        score,
306        reported_at: jiff::Timestamp::now(),
307        reporter: req.reporter,
308    };
309    let trace_id = req.trace_id.clone();
310    match tokio::task::spawn_blocking(move || store::append_deferred(&db, &req.trace_id, &dv)).await
311    {
312        Ok(Ok(())) => {
313            // Close the reactive loop: nudge the live serve threshold toward the target.
314            if let (Some(a), Some(correct)) = (state.adaptive.as_ref(), feedback_signal)
315                && let Ok(mut g) = a.lock()
316            {
317                g.observe_served(correct);
318            }
319            (
320                axum::http::StatusCode::ACCEPTED,
321                Json(serde_json::json!({ "status": "recorded", "trace_id": trace_id })),
322            )
323                .into_response()
324        }
325        Ok(Err(e)) => {
326            tracing::error!(%e, "feedback: append_deferred failed");
327            ProxyError::Internal(e.to_string()).into_response()
328        }
329        Err(e) => {
330            tracing::error!(%e, "feedback: append_deferred task panicked");
331            ProxyError::Internal(e.to_string()).into_response()
332        }
333    }
334}
335
336/// The header a caller may set to group requests into a session for the audit trail. When
337/// absent, each request is its own session (keyed by its own trace id).
338const SESSION_HEADER: &str = "x-firstpass-session";
339
340/// Header carrying the calling agent identity (feature/routing signal).
341const AGENT_HEADER: &str = "x-firstpass-agent";
342/// Header carrying the calling subagent identity.
343const SUBAGENT_HEADER: &str = "x-firstpass-subagent";
344
345/// `POST /v1/messages` — dispatch on the matched route's mode. **Enforce** routes run the
346/// escalation engine (gate + escalate + failover); everything else is an **observe**
347/// passthrough (forward unchanged, trace asynchronously). Either way the trace is recorded
348/// off the response path.
349async fn messages(
350    State(state): State<AppState>,
351    Extension(TenantId(tenant)): Extension<TenantId>,
352    headers: HeaderMap,
353    body: Bytes,
354) -> Response {
355    let session_header = header_str(&headers, SESSION_HEADER);
356
357    // Only parse the request for routing when a routing config is loaded — an observe-only
358    // deployment does zero on-path parsing and keeps its zero-added-latency guarantee.
359    if let Some(routing) = state.config.routing.as_ref() {
360        let features = extract_features(&headers, &body);
361        if let Some(route) = routing
362            .route_for(&features)
363            .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
364        {
365            // Clone the matched route so no borrow of `state.config` is held across the await;
366            // routes are tiny (a handful of strings).
367            let route = route.clone();
368            if enforce_can_handle(&features, &body, routing.escalation.enforce_structured) {
369                return handle_enforce(
370                    &state,
371                    &headers,
372                    &body,
373                    features,
374                    &route,
375                    session_header,
376                    tenant,
377                )
378                .await;
379            }
380            // With `enforce_structured` off (default), tool/image/tool-block requests fall through
381            // to transparent observe passthrough (correct, un-gated) rather than being routed —
382            // byte-identical to before ADR 0005. (With it on, `enforce_can_handle` returns true and
383            // we never reach here.)
384            tracing::info!(
385                "enforce route matched but request carries tools/images; serving via observe passthrough"
386            );
387        }
388    }
389    observe_passthrough(state, headers, body, session_header, tenant).await
390}
391
392/// Whether the enforce path can faithfully handle this request.
393///
394/// Request content is carried verbatim (ADR 0005 P1), so tool_use/tool_result/image blocks survive
395/// the round trip, and a streaming client is served the gated result as SSE (P3). Whether such
396/// requests are actually *routed* is the operator's call:
397/// - `enforce_structured == false` (default): tools/images/tool-blocks fall back to observe —
398///   byte-identical to before the ADR (invariant I1).
399/// - `enforce_structured == true` (opt-in, live-verified): tools, images, and streaming all route
400///   through enforce.
401fn enforce_can_handle(features: &Features, body: &[u8], enforce_structured: bool) -> bool {
402    if enforce_structured {
403        return true;
404    }
405    features.tool_count == 0 && !features.has_images && !messages_have_tool_blocks(body)
406}
407
408/// Whether any message carries a `tool_use` or `tool_result` content block (a multi-turn tool
409/// conversation), which the text-only enforce normalization would drop.
410fn messages_have_tool_blocks(body: &[u8]) -> bool {
411    serde_json::from_slice::<Value>(body)
412        .ok()
413        .and_then(|json| {
414            json.get("messages")
415                .and_then(Value::as_array)
416                .map(|messages| messages.iter().any(message_has_tool_block))
417        })
418        .unwrap_or(false)
419}
420
421/// Whether a single message's content contains a `tool_use` or `tool_result` block.
422fn message_has_tool_block(message: &Value) -> bool {
423    message
424        .get("content")
425        .and_then(Value::as_array)
426        .is_some_and(|blocks| {
427            blocks.iter().any(|block| {
428                matches!(
429                    block.get("type").and_then(Value::as_str),
430                    Some("tool_use" | "tool_result")
431                )
432            })
433        })
434}
435
436/// Whether the request opts into server-sent-events streaming (`"stream": true`).
437fn is_stream_request(body: &[u8]) -> bool {
438    serde_json::from_slice::<Value>(body)
439        .ok()
440        .and_then(|json| json.get("stream").and_then(Value::as_bool))
441        .unwrap_or(false)
442}
443
444/// Read a header as an owned `String`, if present and valid UTF-8.
445fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
446    headers
447        .get(name)
448        .and_then(|v| v.to_str().ok())
449        .map(str::to_owned)
450}
451
452/// Build the routing/telemetry feature vector from request headers + body (best-effort;
453/// malformed fields fall back to safe defaults — this must never fail a request).
454fn extract_features(headers: &HeaderMap, body: &[u8]) -> Features {
455    let (_model, tool_count, has_images) = request_features(body);
456    let mut f = Features::new(TaskKind::Other);
457    f.agent = header_str(headers, AGENT_HEADER);
458    f.subagent = header_str(headers, SUBAGENT_HEADER);
459    f.tool_count = tool_count;
460    f.has_images = has_images;
461    // Pre-call we don't know the token count, so bucket by request byte size — a coarse,
462    // monotonic proxy that never exposes the exact prompt (matches the privacy contract).
463    f.prompt_token_bucket = token_bucket(body.len() as u64);
464    f.hour_bucket = hour_bucket(jiff::Timestamp::now());
465    f
466}
467
468/// Enforce mode (SPEC §7.1): run the escalation engine and serve the first output that clears
469/// the route's gates, escalating on failure with cross-provider failover.
470async fn handle_enforce(
471    state: &AppState,
472    headers: &HeaderMap,
473    body: &Bytes,
474    features: Features,
475    route: &Route,
476    session_header: Option<String>,
477    tenant: String,
478) -> Response {
479    let Some(base_request) = parse_model_request(body) else {
480        return ProxyError::BadRequest(
481            "request body is not a valid Anthropic Messages request".to_owned(),
482        )
483        .into_response();
484    };
485    let auth = Auth::from_headers(headers);
486    let gate_defs = state
487        .config
488        .routing
489        .as_ref()
490        .map_or(&[][..], |cfg| &cfg.gate_defs);
491    let gates = resolve_gates(&route.gates, gate_defs, &state.providers, &auth);
492    let session_id = session_header.unwrap_or_else(|| Uuid::now_v7().to_string());
493    let (budget, max_rungs, speculation, serve_threshold) = match state.config.routing.as_ref() {
494        Some(cfg) => (
495            cfg.budget.per_request_usd,
496            cfg.escalation.max_rungs_per_request,
497            cfg.escalation.speculation,
498            cfg.escalation.serve_threshold,
499        ),
500        None => (None, 3, 0, None),
501    };
502    // Online adaptive conformal: serve against the LIVE-tracked threshold (updated by /v1/feedback).
503    // Falls back to the fixed config threshold when adaptive is off or its lock is poisoned.
504    let serve_threshold = state
505        .adaptive
506        .as_ref()
507        .and_then(|a| a.lock().ok().map(|g| g.threshold()))
508        .or(serve_threshold);
509
510    let ctx = EnforceCtx {
511        ladder: &route.ladder,
512        gates: &gates,
513        health: &state.gate_health,
514        base_request: &base_request,
515        providers: &state.providers,
516        auth: &auth,
517        prices: &state.config.prices,
518        budget_per_request_usd: budget,
519        max_rungs,
520        speculation,
521        serve_threshold,
522        features,
523        // The tenant stamped on the enforce trace is the resolved identity from the auth layer
524        // (authenticated key, or the static default when auth is off) — never the request body.
525        tenant_id: tenant,
526        session_id,
527        prompt_hash: prompt_hash(&state.config.prompt_salt, body),
528        api: "anthropic.messages".to_owned(),
529        policy_id: "static-ladder@v0".to_owned(),
530    };
531
532    let (outcome, trace) = route_enforce(ctx).await;
533    // The trace is already built; enqueue it off-path (non-blocking `try_send`, so no spawn needed).
534    offer_trace(&state.traces, trace);
535
536    match outcome {
537        EngineOutcome::Served(resp) => {
538            let message = anthropic_response_json(&resp);
539            // A streaming client gets the gated result re-emitted as SSE (ADR 0005 P3): the gate
540            // needs the whole candidate, so enforce can't stream token-by-token from the model — it
541            // buffers to gate, then streams the served blocks out. tool_use blocks are preserved.
542            if is_stream_request(body) {
543                (
544                    axum::http::StatusCode::OK,
545                    [(
546                        axum::http::header::CONTENT_TYPE,
547                        "text/event-stream; charset=utf-8",
548                    )],
549                    anthropic_sse_from_message(&message),
550                )
551                    .into_response()
552            } else {
553                (axum::http::StatusCode::OK, Json(message)).into_response()
554            }
555        }
556        EngineOutcome::Failed(msg) => ProxyError::Engine(msg).into_response(),
557    }
558}
559
560/// Parse an Anthropic Messages request body into the normalized [`ModelRequest`]. Returns
561/// `None` if the body isn't valid JSON or lacks a `messages` array.
562///
563// Message content is preserved **verbatim** (string or array of blocks) — a plain-string content
564// serializes byte-identical on the wire, and tool_use/tool_result/image blocks survive the round
565// trip (ADR 0005, invariant I2). Gates operate on `ChatMessage::text_view()`, not the raw content,
566// so gate behavior is unchanged. Which requests actually enter enforce is still governed by
567// `enforce_can_handle`; this function only guarantees no fidelity is lost once they do.
568fn parse_model_request(body: &[u8]) -> Option<ModelRequest> {
569    let json: Value = serde_json::from_slice(body).ok()?;
570    let messages_json = json.get("messages")?.as_array()?;
571    let messages = messages_json
572        .iter()
573        .map(|m| ChatMessage {
574            role: m
575                .get("role")
576                .and_then(Value::as_str)
577                .unwrap_or("user")
578                .to_owned(),
579            content: m
580                .get("content")
581                .cloned()
582                .unwrap_or_else(|| Value::String(String::new())),
583        })
584        .collect();
585    let system = json
586        .get("system")
587        .and_then(Value::as_str)
588        .map(str::to_owned);
589    let max_tokens = json
590        .get("max_tokens")
591        .and_then(Value::as_u64)
592        .and_then(|n| u32::try_from(n).ok())
593        .unwrap_or(1024);
594    let tools = json.get("tools").cloned().unwrap_or(Value::Null);
595    Some(ModelRequest {
596        model: json
597            .get("model")
598            .and_then(Value::as_str)
599            .unwrap_or_default()
600            .to_owned(),
601        system,
602        messages,
603        max_tokens,
604        tools,
605    })
606}
607
608/// Render a served [`ModelResponse`] back into an Anthropic Messages response envelope, so the
609/// caller sees the same wire shape regardless of which provider actually answered.
610///
611/// The `content` blocks come **verbatim** from the upstream response (`resp.raw`) when it is an
612/// Anthropic message — so `tool_use` / `thinking` / multiple text blocks reach the caller intact
613/// (ADR 0005 I2). Only when `raw` has no Anthropic `content` array (a synthetic response, or the
614/// OpenAI adapter, which has `choices` instead) do we fall back to a single reconstructed text
615/// block. The envelope (`id`, `model`, `usage`) is always normalized so the served model id is the
616/// prefixed ladder id, not the bare wire id.
617fn anthropic_response_json(resp: &ModelResponse) -> Value {
618    let content = resp
619        .raw
620        .get("content")
621        .filter(|c| c.is_array())
622        .cloned()
623        .unwrap_or_else(|| serde_json::json!([{ "type": "text", "text": resp.text }]));
624    serde_json::json!({
625        "id": format!("msg_{}", Uuid::now_v7()),
626        "type": "message",
627        "role": "assistant",
628        "model": resp.model,
629        "content": content,
630        "usage": { "input_tokens": resp.in_tokens, "output_tokens": resp.out_tokens },
631    })
632}
633
634/// Append one `event: <type>\ndata: <json>\n\n` SSE frame.
635fn sse_event(out: &mut String, event: &str, data: &Value) {
636    out.push_str("event: ");
637    out.push_str(event);
638    out.push_str("\ndata: ");
639    out.push_str(&data.to_string());
640    out.push_str("\n\n");
641}
642
643/// Re-emit a served Anthropic message envelope (from [`anthropic_response_json`]) as an SSE stream
644/// body, so a `stream: true` client is served even though enforce buffered the response to gate it
645/// (ADR 0005 P3). The gate needs the full candidate, so this is not token-by-token streaming from
646/// the model — each content block is emitted as a single delta. `tool_use` blocks are preserved:
647/// their `input` is streamed as one `input_json_delta` (invariant I2), so the caller reconstructs
648/// the exact tool call.
649fn anthropic_sse_from_message(message: &Value) -> String {
650    let mut out = String::new();
651
652    // message_start carries the envelope with content emptied — the blocks stream next.
653    let mut start_msg = message.clone();
654    start_msg["content"] = Value::Array(Vec::new());
655    sse_event(
656        &mut out,
657        "message_start",
658        &serde_json::json!({ "type": "message_start", "message": start_msg }),
659    );
660
661    let empty = Vec::new();
662    let blocks = message
663        .get("content")
664        .and_then(Value::as_array)
665        .unwrap_or(&empty);
666    for (i, block) in blocks.iter().enumerate() {
667        match block.get("type").and_then(Value::as_str) {
668            Some("tool_use") => {
669                // Start with an empty input object, then stream the real input as one JSON delta.
670                let mut shell = block.clone();
671                shell["input"] = serde_json::json!({});
672                sse_event(
673                    &mut out,
674                    "content_block_start",
675                    &serde_json::json!({ "type": "content_block_start", "index": i, "content_block": shell }),
676                );
677                let input_json = block
678                    .get("input")
679                    .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
680                sse_event(
681                    &mut out,
682                    "content_block_delta",
683                    &serde_json::json!({ "type": "content_block_delta", "index": i,
684                        "delta": { "type": "input_json_delta", "partial_json": input_json } }),
685                );
686            }
687            _ => {
688                // text (and any other text-bearing block): start empty, stream the text as one delta.
689                let text = block.get("text").and_then(Value::as_str).unwrap_or("");
690                sse_event(
691                    &mut out,
692                    "content_block_start",
693                    &serde_json::json!({ "type": "content_block_start", "index": i,
694                        "content_block": { "type": "text", "text": "" } }),
695                );
696                sse_event(
697                    &mut out,
698                    "content_block_delta",
699                    &serde_json::json!({ "type": "content_block_delta", "index": i,
700                        "delta": { "type": "text_delta", "text": text } }),
701                );
702            }
703        }
704        sse_event(
705            &mut out,
706            "content_block_stop",
707            &serde_json::json!({ "type": "content_block_stop", "index": i }),
708        );
709    }
710
711    let out_tokens = message
712        .pointer("/usage/output_tokens")
713        .cloned()
714        .unwrap_or_else(|| Value::from(0));
715    sse_event(
716        &mut out,
717        "message_delta",
718        &serde_json::json!({ "type": "message_delta", "delta": { "stop_reason": "end_turn" },
719            "usage": { "output_tokens": out_tokens } }),
720    );
721    sse_event(
722        &mut out,
723        "message_stop",
724        &serde_json::json!({ "type": "message_stop" }),
725    );
726    out
727}
728
729/// Observe mode (SPEC §7.1a): forward unchanged, return unchanged, trace asynchronously.
730async fn observe_passthrough(
731    state: AppState,
732    headers: HeaderMap,
733    body: Bytes,
734    session_header: Option<String>,
735    tenant: String,
736) -> Response {
737    // Streaming requests are relayed chunk-by-chunk rather than buffered (SPEC §7.4).
738    if is_stream_request(&body) {
739        return observe_stream(state, headers, body, session_header, tenant).await;
740    }
741    let start = Instant::now();
742    let result = forward_anthropic(
743        &state.http,
744        &state.config.upstream_anthropic,
745        &headers,
746        body.clone(),
747    )
748    .await;
749    let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
750
751    match result {
752        Ok((status, resp_headers, resp_body)) => {
753            // Build + record the trace on a detached task so neither JSON parsing nor the
754            // channel send touches the response path: observe mode adds zero latency to what
755            // the caller sees (SPEC §7.1a). `Bytes` clones are cheap (refcounted).
756            spawn_trace(
757                &state,
758                body,
759                Some(resp_body.clone()),
760                latency_ms,
761                session_header,
762                tenant,
763            );
764            (status, resp_headers, resp_body).into_response()
765        }
766        Err(err) => {
767            spawn_trace(&state, body, None, latency_ms, session_header, tenant);
768            err.into_response()
769        }
770    }
771}
772
773/// Observe mode for a streaming request (`stream: true`): relay the upstream SSE response
774/// chunk-by-chunk instead of buffering, so streaming is preserved to the caller and
775/// time-to-first-byte stays low. `latency_ms` is time-to-response-headers (the added-latency
776/// figure that matters), recorded off the response path.
777///
778// ponytail: streamed-response token usage lives in the SSE `message_start`/`message_delta` events
779// we don't buffer, so the trace records request-side features + latency now; parsing usage from a
780// teed SSE stream is the follow-on.
781async fn observe_stream(
782    state: AppState,
783    headers: HeaderMap,
784    body: Bytes,
785    session_header: Option<String>,
786    tenant: String,
787) -> Response {
788    let start = Instant::now();
789    let result = forward_anthropic_streaming(
790        &state.http,
791        &state.config.upstream_anthropic,
792        &headers,
793        body.clone(),
794    )
795    .await;
796    let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
797
798    match result {
799        Ok((status, resp_headers, response)) => {
800            spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
801            let stream_body = Body::from_stream(response.bytes_stream());
802            (status, resp_headers, stream_body).into_response()
803        }
804        Err(err) => {
805            spawn_trace(&state, body, None, latency_ms, session_header, tenant);
806            err.into_response()
807        }
808    }
809}
810
811/// Enqueue a request-side trace for a streamed observe response, off the response path.
812fn spawn_stream_trace(
813    state: &AppState,
814    req_body: Bytes,
815    latency_ms: u64,
816    session_header: Option<String>,
817    tenant: String,
818) {
819    let config = state.config.clone();
820    let traces = state.traces.clone();
821    tokio::spawn(async move {
822        let mut trace =
823            build_stream_trace(&config, &req_body, latency_ms, session_header.as_deref());
824        // Stamp the resolved tenant identity — never the config default nor anything request-borne.
825        trace.tenant_id = tenant;
826        offer_trace(&traces, trace);
827    });
828}
829
830/// Construct the trace and enqueue it for the background writer, entirely off the response
831/// path. Fire-and-forget: if the writer has shut down we log rather than propagate — recording
832/// must never affect what the caller sees. `resp_body` is `Some` for a forwarded response and
833/// `None` when the upstream call failed outright.
834fn spawn_trace(
835    state: &AppState,
836    req_body: Bytes,
837    resp_body: Option<Bytes>,
838    latency_ms: u64,
839    session_header: Option<String>,
840    tenant: String,
841) {
842    let config = state.config.clone();
843    let traces = state.traces.clone();
844    tokio::spawn(async move {
845        let mut trace = match resp_body {
846            Some(resp) => build_trace(
847                &config,
848                &req_body,
849                &resp,
850                latency_ms,
851                session_header.as_deref(),
852            ),
853            None => build_error_trace(&config, &req_body, latency_ms, session_header.as_deref()),
854        };
855        // Stamp the resolved tenant identity — never the config default nor anything request-borne.
856        trace.tenant_id = tenant;
857        offer_trace(&traces, trace);
858    });
859}
860
861/// Session id for the trace: the caller-supplied header, or the trace's own id when absent.
862fn session_id(session_header: Option<&str>, trace_id: Uuid) -> String {
863    session_header
864        .map(str::to_owned)
865        .unwrap_or_else(|| trace_id.to_string())
866}
867
868/// Salted hash of the raw request body — the only trace of the prompt that ever touches
869/// storage (SPEC: never log or persist raw prompt text).
870fn prompt_hash(salt: &str, body: &[u8]) -> String {
871    let mut salted = Vec::with_capacity(salt.len() + body.len());
872    salted.extend_from_slice(salt.as_bytes());
873    salted.extend_from_slice(body);
874    sha256_hex(&salted)
875}
876
877/// Best-effort request-side feature extraction: model name, tool count, and whether any
878/// message carries image content. Malformed/absent fields fall back to safe defaults rather
879/// than failing the request — this is telemetry, not the served response.
880fn request_features(body: &[u8]) -> (Option<String>, u32, bool) {
881    let Ok(json) = serde_json::from_slice::<Value>(body) else {
882        return (None, 0, false);
883    };
884    let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
885    let tool_count = json
886        .get("tools")
887        .and_then(Value::as_array)
888        .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
889    let has_images = json
890        .get("messages")
891        .and_then(Value::as_array)
892        .is_some_and(|messages| messages.iter().any(message_has_image));
893    (model, tool_count, has_images)
894}
895
896/// Whether a single message's content contains an image block (`{"type": "image", ...}`).
897fn message_has_image(message: &Value) -> bool {
898    message
899        .get("content")
900        .and_then(Value::as_array)
901        .is_some_and(|blocks| {
902            blocks
903                .iter()
904                .any(|block| block.get("type").and_then(Value::as_str) == Some("image"))
905        })
906}
907
908/// Response-side usage extraction: model name and token counts, defaulting to `0` when the
909/// upstream response doesn't carry them (e.g. an error body).
910fn response_usage(body: &[u8]) -> (Option<String>, u64, u64) {
911    let Ok(json) = serde_json::from_slice::<Value>(body) else {
912        return (None, 0, 0);
913    };
914    let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
915    let in_tokens = json
916        .pointer("/usage/input_tokens")
917        .and_then(Value::as_u64)
918        .unwrap_or(0);
919    let out_tokens = json
920        .pointer("/usage/output_tokens")
921        .and_then(Value::as_u64)
922        .unwrap_or(0);
923    (model, in_tokens, out_tokens)
924}
925
926/// Build the observe-mode trace for a request that was successfully forwarded and answered.
927fn build_trace(
928    config: &ProxyConfig,
929    req_body: &Bytes,
930    resp_body: &Bytes,
931    latency_ms: u64,
932    session_header: Option<&str>,
933) -> Trace {
934    let (req_model, tool_count, has_images) = request_features(req_body);
935    let (resp_model, in_tokens, out_tokens) = response_usage(resp_body);
936    let model = resp_model
937        .or(req_model)
938        .unwrap_or_else(|| "unknown".to_owned());
939
940    let cost_usd = config
941        .prices
942        .cost_usd(&format!("anthropic/{model}"), in_tokens, out_tokens)
943        .unwrap_or(0.0);
944
945    let attempt = Attempt {
946        rung: 0,
947        model,
948        provider: "anthropic".to_owned(),
949        in_tokens,
950        out_tokens,
951        cost_usd,
952        latency_ms,
953        gates: Vec::new(),
954        verdict: Verdict::Pass,
955    };
956
957    let mut trace = base_trace(config, req_body, latency_ms, session_header);
958    trace.request.features.prompt_token_bucket = token_bucket(in_tokens);
959    trace.request.features.tool_count = tool_count;
960    trace.request.features.has_images = has_images;
961    trace.attempts.push(attempt);
962    trace.final_ = FinalOutcome {
963        served_rung: Some(0),
964        served_from: ServedFrom::Attempt,
965        total_cost_usd: cost_usd,
966        gate_cost_usd: 0.0,
967        total_latency_ms: latency_ms,
968        escalations: 0,
969        counterfactual_baseline_usd: cost_usd,
970        savings_usd: 0.0,
971    };
972    trace.recompute_savings();
973    trace
974}
975
976/// Build the observe-mode trace for a **streamed** response: we relayed real bytes to the caller,
977/// but the token usage lives in the SSE events we didn't buffer, so it's recorded as served with
978/// unknown (zero) usage — honest about what we served without inventing token counts.
979fn build_stream_trace(
980    config: &ProxyConfig,
981    req_body: &Bytes,
982    latency_ms: u64,
983    session_header: Option<&str>,
984) -> Trace {
985    let (req_model, tool_count, has_images) = request_features(req_body);
986    let model = req_model.unwrap_or_else(|| "unknown".to_owned());
987
988    let attempt = Attempt {
989        rung: 0,
990        model,
991        provider: "anthropic".to_owned(),
992        in_tokens: 0,
993        out_tokens: 0,
994        cost_usd: 0.0,
995        latency_ms,
996        gates: Vec::new(),
997        verdict: Verdict::Pass,
998    };
999
1000    let mut trace = base_trace(config, req_body, latency_ms, session_header);
1001    trace.request.features.tool_count = tool_count;
1002    trace.request.features.has_images = has_images;
1003    trace.attempts.push(attempt);
1004    trace.final_ = FinalOutcome {
1005        served_rung: Some(0),
1006        served_from: ServedFrom::Attempt,
1007        total_cost_usd: 0.0,
1008        gate_cost_usd: 0.0,
1009        total_latency_ms: latency_ms,
1010        escalations: 0,
1011        counterfactual_baseline_usd: 0.0,
1012        savings_usd: 0.0,
1013    };
1014    trace.recompute_savings();
1015    trace
1016}
1017
1018/// Build the observe-mode trace for a request whose upstream call failed outright (no
1019/// response to report usage from). Recorded with `served_from: Error` and no attempts —
1020/// keep the audit trail honest that nothing was served.
1021fn build_error_trace(
1022    config: &ProxyConfig,
1023    req_body: &Bytes,
1024    latency_ms: u64,
1025    session_header: Option<&str>,
1026) -> Trace {
1027    let (_, tool_count, has_images) = request_features(req_body);
1028    let mut trace = base_trace(config, req_body, latency_ms, session_header);
1029    trace.request.features.tool_count = tool_count;
1030    trace.request.features.has_images = has_images;
1031    trace.final_ = FinalOutcome {
1032        served_rung: None,
1033        served_from: ServedFrom::Error,
1034        total_cost_usd: 0.0,
1035        gate_cost_usd: 0.0,
1036        total_latency_ms: latency_ms,
1037        escalations: 0,
1038        counterfactual_baseline_usd: 0.0,
1039        savings_usd: 0.0,
1040    };
1041    trace.recompute_savings();
1042    trace
1043}
1044
1045/// The parts of a trace that don't depend on whether the call succeeded: identity, policy,
1046/// and the request-side feature vector minus token bucket (which needs response usage).
1047fn base_trace(
1048    config: &ProxyConfig,
1049    req_body: &Bytes,
1050    latency_ms: u64,
1051    session_header: Option<&str>,
1052) -> Trace {
1053    let trace_id = Uuid::now_v7();
1054    let mut features = Features::new(TaskKind::Other);
1055    features.hour_bucket = hour_bucket(jiff::Timestamp::now());
1056
1057    Trace {
1058        trace_id,
1059        prev_hash: GENESIS_HASH.to_owned(),
1060        tenant_id: config.tenant_id.clone(),
1061        session_id: session_id(session_header, trace_id),
1062        ts: jiff::Timestamp::now(),
1063        mode: Mode::Observe,
1064        policy: PolicyRef {
1065            id: "observe-passthrough@v0".to_owned(),
1066            explore: false,
1067        },
1068        request: RequestInfo {
1069            api: "anthropic.messages".to_owned(),
1070            prompt_hash: prompt_hash(&config.prompt_salt, req_body),
1071            features,
1072        },
1073        attempts: Vec::new(),
1074        deferred: Vec::new(),
1075        final_: FinalOutcome {
1076            served_rung: None,
1077            served_from: ServedFrom::Error,
1078            total_cost_usd: 0.0,
1079            gate_cost_usd: 0.0,
1080            total_latency_ms: latency_ms,
1081            escalations: 0,
1082            counterfactual_baseline_usd: 0.0,
1083            savings_usd: 0.0,
1084        },
1085    }
1086}
1087
1088#[cfg(test)]
1089mod tests {
1090    use bytes::Bytes;
1091
1092    use super::*;
1093
1094    fn test_config() -> ProxyConfig {
1095        ProxyConfig::from_lookup(|_| None).unwrap()
1096    }
1097
1098    #[test]
1099    fn build_trace_maps_request_and_response_fields() {
1100        let config = test_config();
1101        let req = Bytes::from_static(
1102            br#"{"model":"claude-haiku-4-5","tools":[{"name":"a"}],"messages":[]}"#,
1103        );
1104        let resp = Bytes::from_static(
1105            br#"{"model":"claude-haiku-4-5","usage":{"input_tokens":1200,"output_tokens":300}}"#,
1106        );
1107
1108        let trace = build_trace(&config, &req, &resp, 42, Some("sess-1"));
1109
1110        assert_eq!(trace.request.api, "anthropic.messages");
1111        assert_eq!(trace.session_id, "sess-1");
1112        assert_eq!(trace.attempts.len(), 1);
1113        let attempt = &trace.attempts[0];
1114        assert_eq!(attempt.model, "claude-haiku-4-5");
1115        assert_eq!(attempt.provider, "anthropic");
1116        assert_eq!(attempt.in_tokens, 1200);
1117        assert_eq!(attempt.out_tokens, 300);
1118        assert!(attempt.cost_usd > 0.0);
1119        assert_eq!(trace.request.features.tool_count, 1);
1120        assert!(!trace.request.features.has_images);
1121        assert_eq!(trace.final_.served_rung, Some(0));
1122    }
1123
1124    #[test]
1125    fn build_trace_falls_back_to_trace_id_session_when_header_absent() {
1126        let config = test_config();
1127        let req = Bytes::from_static(b"{}");
1128        let resp = Bytes::from_static(b"{}");
1129
1130        let trace = build_trace(&config, &req, &resp, 1, None);
1131
1132        assert_eq!(trace.session_id, trace.trace_id.to_string());
1133    }
1134
1135    #[test]
1136    fn build_error_trace_has_no_attempts_and_served_from_error() {
1137        let config = test_config();
1138        let req = Bytes::from_static(br#"{"model":"claude-haiku-4-5"}"#);
1139
1140        let trace = build_error_trace(&config, &req, 7, None);
1141
1142        assert!(trace.attempts.is_empty());
1143        assert_eq!(trace.final_.served_from, ServedFrom::Error);
1144        assert_eq!(trace.final_.served_rung, None);
1145    }
1146
1147    #[test]
1148    fn message_with_image_block_sets_has_images() {
1149        let req = br#"{"messages":[{"role":"user","content":[{"type":"image"}]}]}"#;
1150        let (_, _, has_images) = request_features(req);
1151        assert!(has_images);
1152    }
1153
1154    #[test]
1155    fn prompt_hash_never_contains_raw_prompt_text() {
1156        let hash = prompt_hash("salt", b"super secret prompt");
1157        assert!(!hash.contains("secret"));
1158        assert_eq!(hash.len(), 64);
1159    }
1160
1161    #[test]
1162    fn parse_model_request_preserves_content_verbatim_and_projects_text() {
1163        let body = br#"{"model":"m","system":"sys","max_tokens":50,
1164            "messages":[{"role":"user","content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]},
1165                        {"role":"assistant","content":"c"}]}"#;
1166        let req = parse_model_request(body).unwrap();
1167        assert_eq!(req.system.as_deref(), Some("sys"));
1168        assert_eq!(req.max_tokens, 50);
1169        assert_eq!(req.messages.len(), 2);
1170        // I2: the block array is carried verbatim, not flattened away...
1171        assert_eq!(
1172            req.messages[0].content,
1173            serde_json::json!([{"type":"text","text":"a"},{"type":"text","text":"b"}])
1174        );
1175        // ...and a plain string stays a plain string (I1: byte-identical on the wire).
1176        assert_eq!(req.messages[1].content, Value::String("c".to_owned()));
1177        // Gates see the same text they always did.
1178        assert_eq!(req.messages[0].text_view(), "a\nb");
1179        assert_eq!(req.messages[1].text_view(), "c");
1180    }
1181
1182    #[test]
1183    fn tool_and_image_blocks_survive_the_request_round_trip() {
1184        // ADR 0005 I2: tool_use / tool_result / image blocks are never dropped on the request side.
1185        let body = br#"{"model":"m","max_tokens":50,"messages":[
1186            {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
1187            {"role":"user","content":[
1188                {"type":"tool_result","tool_use_id":"t1","content":"2"},
1189                {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
1190            ]}]}"#;
1191        let req = parse_model_request(body).unwrap();
1192        let round_tripped = serde_json::to_value(&req.messages).unwrap();
1193        assert_eq!(
1194            round_tripped,
1195            serde_json::json!([
1196                {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
1197                {"role":"user","content":[
1198                    {"type":"tool_result","tool_use_id":"t1","content":"2"},
1199                    {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
1200                ]}
1201            ])
1202        );
1203    }
1204
1205    #[test]
1206    fn text_message_serializes_byte_identical_to_a_plain_string() {
1207        // I1: a string-content message must not gain array wrapping on the wire.
1208        let m = ChatMessage::text("user", "hello");
1209        assert_eq!(
1210            serde_json::to_string(&m).unwrap(),
1211            r#"{"role":"user","content":"hello"}"#
1212        );
1213    }
1214
1215    #[test]
1216    fn parse_model_request_rejects_non_message_bodies() {
1217        assert!(parse_model_request(b"not json").is_none());
1218        assert!(parse_model_request(br#"{"no":"messages"}"#).is_none());
1219    }
1220
1221    // --- Enforce-path handler tests (drive `messages` end-to-end with mock providers) ---
1222
1223    use crate::provider::{MockProvider, ModelResponse, Provider, ProviderError, ProviderRegistry};
1224    use axum::extract::State;
1225    use std::collections::HashMap;
1226    use std::sync::Arc;
1227    use tokio::sync::mpsc;
1228
1229    fn model_resp(model: &str, text: &str) -> ModelResponse {
1230        ModelResponse {
1231            model: model.to_owned(),
1232            text: text.to_owned(),
1233            in_tokens: 1000,
1234            out_tokens: 400,
1235            raw: serde_json::Value::Null,
1236        }
1237    }
1238
1239    /// Build an `AppState` whose anthropic provider answers the given per-model outcomes, with an
1240    /// enforce route over `ladder`/`gates`. Returns the state and the trace receiver.
1241    fn enforce_state(
1242        ladder: &[&str],
1243        gates: &[&str],
1244        outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
1245    ) -> (AppState, mpsc::Receiver<Trace>) {
1246        let toml = format!(
1247            "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [{}]\ngates = [{}]\n",
1248            ladder
1249                .iter()
1250                .map(|m| format!("\"{m}\""))
1251                .collect::<Vec<_>>()
1252                .join(", "),
1253            gates
1254                .iter()
1255                .map(|g| format!("\"{g}\""))
1256                .collect::<Vec<_>>()
1257                .join(", "),
1258        );
1259        let config = ProxyConfig::from_lookup(|k| match k {
1260            "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
1261            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
1262            _ => None,
1263        })
1264        .unwrap();
1265
1266        let mut outs = HashMap::new();
1267        for (model, out) in outcomes {
1268            outs.insert(model.to_owned(), out);
1269        }
1270        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
1271        map.insert(
1272            "anthropic".to_owned(),
1273            Arc::new(MockProvider::new("anthropic", outs)),
1274        );
1275        let providers = ProviderRegistry::from_map(map);
1276
1277        let (traces, rx) = mpsc::channel(64);
1278        let state = AppState {
1279            config: Arc::new(config),
1280            http: reqwest::Client::new(),
1281            providers,
1282            gate_health: Arc::new(GateHealthRegistry::new()),
1283            traces,
1284            adaptive: None,
1285            tenant_rate_limiter: None,
1286        };
1287        (state, rx)
1288    }
1289
1290    fn user_body() -> Bytes {
1291        Bytes::from_static(
1292            br#"{"model":"ignored","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}"#,
1293        )
1294    }
1295
1296    async fn body_json(resp: Response) -> Value {
1297        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
1298            .await
1299            .unwrap();
1300        serde_json::from_slice(&bytes).unwrap()
1301    }
1302
1303    #[tokio::test]
1304    async fn enforce_serves_first_pass_and_returns_anthropic_shape() {
1305        let (state, mut rx) = enforce_state(
1306            &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
1307            &["non-empty"],
1308            vec![(
1309                "anthropic/claude-haiku-4-5",
1310                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
1311            )],
1312        );
1313        let resp = messages(
1314            State(state),
1315            Extension(TenantId("default".to_owned())),
1316            HeaderMap::new(),
1317            user_body(),
1318        )
1319        .await;
1320        assert_eq!(resp.status(), axum::http::StatusCode::OK);
1321        let json = body_json(resp).await;
1322        assert_eq!(json["type"], "message");
1323        assert_eq!(json["content"][0]["text"], "hello");
1324        assert_eq!(json["model"], "anthropic/claude-haiku-4-5");
1325
1326        let trace = rx.try_recv().expect("a trace was enqueued");
1327        assert_eq!(trace.mode, Mode::Enforce);
1328        assert_eq!(trace.final_.served_rung, Some(0));
1329        assert_eq!(trace.attempts.len(), 1);
1330    }
1331
1332    #[tokio::test]
1333    async fn enforce_escalates_then_serves_and_traces_two_attempts() {
1334        let (state, mut rx) = enforce_state(
1335            &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
1336            &["non-empty"],
1337            vec![
1338                (
1339                    "anthropic/claude-haiku-4-5",
1340                    Ok(model_resp("anthropic/claude-haiku-4-5", "   ")),
1341                ), // fails
1342                (
1343                    "anthropic/claude-sonnet-5",
1344                    Ok(model_resp("anthropic/claude-sonnet-5", "answer")),
1345                ),
1346            ],
1347        );
1348        let resp = messages(
1349            State(state),
1350            Extension(TenantId("default".to_owned())),
1351            HeaderMap::new(),
1352            user_body(),
1353        )
1354        .await;
1355        let json = body_json(resp).await;
1356        assert_eq!(json["content"][0]["text"], "answer");
1357
1358        let trace = rx.try_recv().expect("trace enqueued");
1359        assert_eq!(trace.attempts.len(), 2);
1360        assert_eq!(trace.final_.escalations, 1);
1361        assert_eq!(trace.final_.served_rung, Some(1));
1362    }
1363
1364    #[tokio::test]
1365    async fn enforce_all_rungs_error_returns_502() {
1366        let (state, mut rx) = enforce_state(
1367            &["anthropic/claude-haiku-4-5"],
1368            &["non-empty"],
1369            vec![(
1370                "anthropic/claude-haiku-4-5",
1371                Err(ProviderError::Transport("down".into())),
1372            )],
1373        );
1374        let resp = messages(
1375            State(state),
1376            Extension(TenantId("default".to_owned())),
1377            HeaderMap::new(),
1378            user_body(),
1379        )
1380        .await;
1381        assert_eq!(resp.status(), axum::http::StatusCode::BAD_GATEWAY);
1382        // A trace is still recorded for the failed decision.
1383        assert!(rx.try_recv().is_ok());
1384    }
1385
1386    #[tokio::test]
1387    async fn no_routing_config_falls_through_to_observe_not_enforce() {
1388        // config with no routing => enforce path never runs; observe attempts a real upstream
1389        // call which fails fast against an unroutable host. We only assert it did NOT take the
1390        // enforce branch (which would have used the mock and returned 200 with our text).
1391        let config = ProxyConfig::from_lookup(|k| match k {
1392            "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
1393            _ => None,
1394        })
1395        .unwrap();
1396        let (traces, _rx) = mpsc::channel(64);
1397        let state = AppState {
1398            config: Arc::new(config),
1399            http: reqwest::Client::new(),
1400            providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
1401            gate_health: Arc::new(GateHealthRegistry::new()),
1402            traces,
1403            adaptive: None,
1404            tenant_rate_limiter: None,
1405        };
1406        let resp = messages(
1407            State(state),
1408            Extension(TenantId("default".to_owned())),
1409            HeaderMap::new(),
1410            user_body(),
1411        )
1412        .await;
1413        // Observe path forwards upstream; the bogus host yields a gateway error, not our 200.
1414        assert_ne!(resp.status(), axum::http::StatusCode::OK);
1415    }
1416
1417    #[test]
1418    fn detects_stream_requests() {
1419        assert!(is_stream_request(br#"{"stream": true}"#));
1420        assert!(!is_stream_request(br#"{"stream": false}"#));
1421        assert!(!is_stream_request(br#"{"model":"m"}"#));
1422        assert!(!is_stream_request(b"not json"));
1423    }
1424
1425    #[test]
1426    fn detects_tool_blocks_in_messages() {
1427        let with =
1428            br#"{"messages":[{"role":"user","content":[{"type":"tool_result","content":"42"}]}]}"#;
1429        let without = br#"{"messages":[{"role":"user","content":"hi"}]}"#;
1430        assert!(messages_have_tool_blocks(with));
1431        assert!(!messages_have_tool_blocks(without));
1432    }
1433
1434    #[test]
1435    fn enforce_only_handles_plain_text() {
1436        let plain =
1437            Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
1438        let tools = Bytes::from_static(
1439            br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
1440        );
1441        let f_plain = extract_features(&HeaderMap::new(), &plain);
1442        let f_tools = extract_features(&HeaderMap::new(), &tools);
1443        // Default (enforce_structured = false): plain text routes, tools fall back to observe.
1444        assert!(enforce_can_handle(&f_plain, &plain, false));
1445        assert!(!enforce_can_handle(&f_tools, &tools, false));
1446    }
1447
1448    #[test]
1449    fn structured_enforce_routes_tools_and_streaming() {
1450        // ADR 0005 P2+P3: with the opt-in flag on, tool and streaming requests both route through
1451        // enforce (streaming is served as the gated result re-emitted as SSE).
1452        let tools = Bytes::from_static(
1453            br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
1454        );
1455        let streaming_tools = Bytes::from_static(
1456            br#"{"model":"m","stream":true,"tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
1457        );
1458        let f = extract_features(&HeaderMap::new(), &tools);
1459        assert!(enforce_can_handle(&f, &tools, true));
1460        assert!(enforce_can_handle(&f, &streaming_tools, true));
1461    }
1462
1463    #[test]
1464    fn enforce_sse_reemission_preserves_text_and_tool_use() {
1465        // ADR 0005 P3 + I2: a served response with a text block AND a tool_use block round-trips
1466        // through the SSE re-emitter — the tool call's input survives as an input_json_delta.
1467        let resp = ModelResponse {
1468            model: "anthropic/claude-haiku-4-5".to_owned(),
1469            text: "let me check".to_owned(),
1470            in_tokens: 5,
1471            out_tokens: 7,
1472            raw: serde_json::json!({
1473                "content": [
1474                    { "type": "text", "text": "let me check" },
1475                    { "type": "tool_use", "id": "tu_1", "name": "get_weather", "input": { "city": "Paris" } }
1476                ]
1477            }),
1478        };
1479        let sse = anthropic_sse_from_message(&anthropic_response_json(&resp));
1480
1481        // Parse every data frame structurally (key order is not part of the contract).
1482        let frames: Vec<Value> = sse
1483            .lines()
1484            .filter_map(|l| l.strip_prefix("data: "))
1485            .map(|d| serde_json::from_str::<Value>(d).expect("each SSE data frame is valid JSON"))
1486            .collect();
1487
1488        // Full lifecycle, in order.
1489        assert_eq!(frames.first().unwrap()["type"], "message_start");
1490        assert_eq!(frames.last().unwrap()["type"], "message_stop");
1491        // The text block streams its text as a text_delta.
1492        assert!(frames.iter().any(|f| f["delta"]["type"] == "text_delta"
1493            && f["delta"]["text"] == "let me check"));
1494        // The tool_use block is present with its id/name, and its input streams as one JSON delta —
1495        // not dropped (ADR 0005 I2).
1496        assert!(
1497            frames
1498                .iter()
1499                .any(|f| f["content_block"]["type"] == "tool_use"
1500                    && f["content_block"]["name"] == "get_weather"
1501                    && f["content_block"]["id"] == "tu_1")
1502        );
1503        assert!(
1504            frames
1505                .iter()
1506                .any(|f| f["delta"]["type"] == "input_json_delta"
1507                    && f["delta"]["partial_json"] == r#"{"city":"Paris"}"#)
1508        );
1509    }
1510
1511    /// B2: an enforce route serves plain text (200 from the mock) but falls back to transparent
1512    /// observe passthrough for tool/image requests rather than dropping blocks — proven by the
1513    /// tool request hitting the (bogus) upstream instead of the enforcing mock.
1514    #[tokio::test]
1515    async fn enforce_falls_back_to_observe_for_tool_requests() {
1516        let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
1517        let config = ProxyConfig::from_lookup(|k| match k {
1518            "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
1519            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
1520            "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
1521            _ => None,
1522        })
1523        .unwrap();
1524        let mut outs = HashMap::new();
1525        outs.insert(
1526            "anthropic/m".to_owned(),
1527            Ok(model_resp("anthropic/m", "hello")),
1528        );
1529        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
1530        map.insert(
1531            "anthropic".to_owned(),
1532            Arc::new(MockProvider::new("anthropic", outs)),
1533        );
1534        let (traces, _rx) = mpsc::channel(64);
1535        let state = AppState {
1536            config: Arc::new(config),
1537            http: reqwest::Client::new(),
1538            providers: ProviderRegistry::from_map(map),
1539            gate_health: Arc::new(GateHealthRegistry::new()),
1540            traces,
1541            adaptive: None,
1542            tenant_rate_limiter: None,
1543        };
1544
1545        // Plain text enforces: the mock serves 200.
1546        let plain =
1547            Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
1548        let resp = messages(
1549            State(state.clone()),
1550            Extension(TenantId("default".to_owned())),
1551            HeaderMap::new(),
1552            plain,
1553        )
1554        .await;
1555        assert_eq!(
1556            resp.status(),
1557            axum::http::StatusCode::OK,
1558            "plain text should enforce"
1559        );
1560
1561        // Declares tools => cannot enforce faithfully => observe fallback => bogus upstream => not 200.
1562        let tools = Bytes::from_static(
1563            br#"{"model":"m","tools":[{"name":"get_weather"}],"messages":[{"role":"user","content":"hi"}]}"#,
1564        );
1565        let resp = messages(
1566            State(state.clone()),
1567            Extension(TenantId("default".to_owned())),
1568            HeaderMap::new(),
1569            tools,
1570        )
1571        .await;
1572        assert_ne!(
1573            resp.status(),
1574            axum::http::StatusCode::OK,
1575            "tool request must fall back to observe, not enforce"
1576        );
1577
1578        // tool_result block in a message => same fallback.
1579        let toolres = Bytes::from_static(
1580            br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"x","content":"42"}]}]}"#,
1581        );
1582        let resp = messages(
1583            State(state),
1584            Extension(TenantId("default".to_owned())),
1585            HeaderMap::new(),
1586            toolres,
1587        )
1588        .await;
1589        assert_ne!(
1590            resp.status(),
1591            axum::http::StatusCode::OK,
1592            "tool_result request must fall back to observe, not enforce"
1593        );
1594    }
1595
1596    // --- Feedback API tests (drive `feedback` against a real temp trace store) ---
1597
1598    /// Persist one trace to a fresh temp DB and return (state, db_path, trace_id).
1599    async fn feedback_state() -> (AppState, std::path::PathBuf, String) {
1600        let db = std::env::temp_dir().join(format!("firstpass-feedback-{}.db", Uuid::now_v7()));
1601        let (tx, handle) = crate::store::open(&db).unwrap();
1602
1603        let mut trace = build_error_trace(
1604            &ProxyConfig::from_lookup(|_| None).unwrap(),
1605            &Bytes::from_static(b"{}"),
1606            5,
1607            Some("sess-fb"),
1608        );
1609        trace.attempts.push(Attempt {
1610            rung: 0,
1611            model: "anthropic/claude-haiku-4-5".into(),
1612            provider: "anthropic".into(),
1613            in_tokens: 10,
1614            out_tokens: 5,
1615            cost_usd: 0.001,
1616            latency_ms: 5,
1617            gates: vec![],
1618            verdict: Verdict::Pass,
1619        });
1620        let trace_id = trace.trace_id.to_string();
1621        tx.try_send(trace).unwrap();
1622        drop(tx);
1623        handle.await.unwrap();
1624
1625        let db_str = db.to_string_lossy().into_owned();
1626        let config = ProxyConfig::from_lookup(move |k| match k {
1627            "FIRSTPASS_DB" => Some(db_str.clone()),
1628            _ => None,
1629        })
1630        .unwrap();
1631        let (traces, _rx) = mpsc::channel(64);
1632        let state = AppState {
1633            config: Arc::new(config),
1634            http: reqwest::Client::new(),
1635            providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
1636            gate_health: Arc::new(GateHealthRegistry::new()),
1637            traces,
1638            adaptive: None,
1639            tenant_rate_limiter: None,
1640        };
1641        (state, db, trace_id)
1642    }
1643
1644    #[tokio::test]
1645    async fn feedback_nudges_the_adaptive_threshold() {
1646        use firstpass_core::conformal::AdaptiveConformal;
1647        let (mut state, _db, trace_id) = feedback_state().await;
1648        let aci = Arc::new(std::sync::Mutex::new(AdaptiveConformal::new(0.1, 0.2, 0.5)));
1649        state.adaptive = Some(aci.clone());
1650        let before = aci.lock().unwrap().threshold();
1651
1652        // A served FAILURE raises the threshold (serve more conservatively).
1653        let fail = Bytes::from(
1654            serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "fail", "reporter": "ci" })
1655                .to_string(),
1656        );
1657        assert_eq!(
1658            feedback(
1659                State(state.clone()),
1660                Extension(TenantId("default".to_owned())),
1661                fail
1662            )
1663            .await
1664            .status(),
1665            axum::http::StatusCode::ACCEPTED
1666        );
1667        let after_fail = aci.lock().unwrap().threshold();
1668        assert!(
1669            after_fail > before,
1670            "served fail should raise the live threshold: {before} -> {after_fail}"
1671        );
1672
1673        // A served PASS nudges it back down — the loop is reactive both ways.
1674        let pass = Bytes::from(
1675            serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "pass", "reporter": "ci" })
1676                .to_string(),
1677        );
1678        let _ = feedback(
1679            State(state),
1680            Extension(TenantId("default".to_owned())),
1681            pass,
1682        )
1683        .await;
1684        assert!(aci.lock().unwrap().threshold() < after_fail);
1685    }
1686
1687    #[tokio::test]
1688    async fn feedback_records_a_deferred_verdict_without_breaking_the_chain() {
1689        let (state, db, trace_id) = feedback_state().await;
1690        let body = Bytes::from(
1691            serde_json::json!({
1692                "trace_id": trace_id,
1693                "gate_id": "tests",
1694                "verdict": "pass",
1695                "score": 1.0,
1696                "reporter": "ci",
1697            })
1698            .to_string(),
1699        );
1700        let resp = feedback(
1701            State(state),
1702            Extension(TenantId("default".to_owned())),
1703            body,
1704        )
1705        .await;
1706        assert_eq!(resp.status(), axum::http::StatusCode::ACCEPTED);
1707
1708        // The deferred verdict is visible on the trace view...
1709        let view = crate::store::load_trace_view(&db, "default", &trace_id)
1710            .unwrap()
1711            .unwrap();
1712        assert_eq!(view.deferred.len(), 1);
1713        assert_eq!(view.deferred[0].gate_id, "tests");
1714        // ...and the sealed chain still verifies (the outcome didn't mutate the trace).
1715        let traces = crate::store::load_all_traces(&db).unwrap();
1716        firstpass_core::verify_chain(&traces, GENESIS_HASH).unwrap();
1717
1718        let _ = std::fs::remove_file(&db);
1719    }
1720
1721    #[tokio::test]
1722    async fn feedback_for_unknown_trace_is_404() {
1723        let (state, db, _trace_id) = feedback_state().await;
1724        let body = Bytes::from(
1725            serde_json::json!({
1726                "trace_id": "does-not-exist",
1727                "gate_id": "tests",
1728                "verdict": "pass",
1729                "reporter": "ci",
1730            })
1731            .to_string(),
1732        );
1733        let resp = feedback(
1734            State(state),
1735            Extension(TenantId("default".to_owned())),
1736            body,
1737        )
1738        .await;
1739        assert_eq!(resp.status(), axum::http::StatusCode::NOT_FOUND);
1740        let _ = std::fs::remove_file(&db);
1741    }
1742
1743    /// D4 IDOR: a real trace owned by "default" cannot receive feedback from another tenant. The
1744    /// caller gets a `404` (not `403`), so there is no existence oracle across the boundary.
1745    #[tokio::test]
1746    async fn feedback_across_tenants_is_404_not_403() {
1747        let (state, db, trace_id) = feedback_state().await;
1748        let body = Bytes::from(
1749            serde_json::json!({
1750                "trace_id": trace_id,
1751                "gate_id": "tests",
1752                "verdict": "pass",
1753                "score": 1.0,
1754                "reporter": "attacker",
1755            })
1756            .to_string(),
1757        );
1758        // Caller authenticated as a *different* tenant than the trace's owner.
1759        let resp = feedback(
1760            State(state),
1761            Extension(TenantId("tenant-b".to_owned())),
1762            body,
1763        )
1764        .await;
1765        assert_eq!(
1766            resp.status(),
1767            axum::http::StatusCode::NOT_FOUND,
1768            "cross-tenant feedback must look exactly like a missing trace"
1769        );
1770        let _ = std::fs::remove_file(&db);
1771    }
1772
1773    #[tokio::test]
1774    async fn feedback_rejects_bad_verdict_and_score() {
1775        let (state, db, trace_id) = feedback_state().await;
1776        let bad_verdict = Bytes::from(
1777            serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "maybe", "reporter": "x" })
1778                .to_string(),
1779        );
1780        assert_eq!(
1781            feedback(
1782                State(state.clone()),
1783                Extension(TenantId("default".to_owned())),
1784                bad_verdict
1785            )
1786            .await
1787            .status(),
1788            axum::http::StatusCode::BAD_REQUEST
1789        );
1790        let bad_score = Bytes::from(
1791            serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "pass", "score": 9.0, "reporter": "x" })
1792                .to_string(),
1793        );
1794        assert_eq!(
1795            feedback(
1796                State(state),
1797                Extension(TenantId("default".to_owned())),
1798                bad_score
1799            )
1800            .await
1801            .status(),
1802            axum::http::StatusCode::BAD_REQUEST
1803        );
1804        let _ = std::fs::remove_file(&db);
1805    }
1806
1807    #[tokio::test]
1808    async fn metrics_endpoint_renders_after_a_real_request() {
1809        use tower::ServiceExt;
1810
1811        let (state, mut rx) = enforce_state(
1812            &["anthropic/claude-haiku-4-5"],
1813            &["non-empty"],
1814            vec![(
1815                "anthropic/claude-haiku-4-5",
1816                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
1817            )],
1818        );
1819        let router = app(state).expect("prometheus recorder installs");
1820
1821        let req = axum::http::Request::builder()
1822            .method("POST")
1823            .uri("/v1/messages")
1824            .header("content-type", "application/json")
1825            .body(Body::from(user_body()))
1826            .unwrap();
1827        let resp = router.clone().oneshot(req).await.unwrap();
1828        assert_eq!(resp.status(), axum::http::StatusCode::OK);
1829        rx.try_recv().expect("a trace was enqueued");
1830
1831        let metrics_req = axum::http::Request::builder()
1832            .method("GET")
1833            .uri("/metrics")
1834            .body(Body::empty())
1835            .unwrap();
1836        let metrics_resp = router.oneshot(metrics_req).await.unwrap();
1837        assert_eq!(metrics_resp.status(), axum::http::StatusCode::OK);
1838        let bytes = axum::body::to_bytes(metrics_resp.into_body(), 1 << 20)
1839            .await
1840            .unwrap();
1841        let body = String::from_utf8(bytes.to_vec()).unwrap();
1842        assert!(
1843            body.contains("firstpass_enforce_latency_ms"),
1844            "metrics body missing enforce latency histogram: {body}"
1845        );
1846        assert!(
1847            body.contains("firstpass_served_total"),
1848            "metrics body missing served counter: {body}"
1849        );
1850    }
1851
1852    // --- Multi-tenant auth (ADR 0004 §D1) integration tests, driven through the real router ---
1853
1854    /// Build an `AppState` whose config toggles auth and (optionally) carries a tenant-keys JSON.
1855    fn auth_state(require_auth: bool, keys_json: Option<String>) -> AppState {
1856        auth_state_rated(require_auth, keys_json, None)
1857    }
1858
1859    /// Like [`auth_state`], but also wires `FIRSTPASS_TENANT_RATE_PER_SEC` (ADR 0004 §D6) when
1860    /// `rate_per_sec` is `Some`.
1861    fn auth_state_rated(
1862        require_auth: bool,
1863        keys_json: Option<String>,
1864        rate_per_sec: Option<u32>,
1865    ) -> AppState {
1866        let config = ProxyConfig::from_lookup(|k| match k {
1867            "FIRSTPASS_REQUIRE_AUTH" => require_auth.then(|| "true".to_owned()),
1868            "FIRSTPASS_TENANT_KEYS_JSON" => keys_json.clone(),
1869            "FIRSTPASS_TENANT_RATE_PER_SEC" => rate_per_sec.map(|n| n.to_string()),
1870            _ => None,
1871        })
1872        .unwrap();
1873        let (traces, _rx) = mpsc::channel(64);
1874        // Deliberately leak the receiver for the test's lifetime so the sender never reports the
1875        // channel closed (the auth tests exercise `/v1/capabilities`, which enqueues no trace).
1876        std::mem::forget(_rx);
1877        let providers: HashMap<String, Arc<dyn Provider>> = HashMap::new();
1878        let tenant_rate_limiter = build_tenant_rate_limiter(&config);
1879        AppState {
1880            config: Arc::new(config),
1881            http: reqwest::Client::new(),
1882            providers: ProviderRegistry::from_map(providers),
1883            gate_health: Arc::new(GateHealthRegistry::new()),
1884            traces,
1885            adaptive: None,
1886            tenant_rate_limiter,
1887        }
1888    }
1889
1890    fn cap_request(auth_header: Option<&str>) -> axum::http::Request<Body> {
1891        let mut b = axum::http::Request::builder()
1892            .method("GET")
1893            .uri("/v1/capabilities");
1894        if let Some(h) = auth_header {
1895            b = b.header("authorization", h);
1896        }
1897        b.body(Body::empty()).unwrap()
1898    }
1899
1900    #[tokio::test]
1901    async fn auth_off_allows_unauthenticated_request() {
1902        use tower::ServiceExt;
1903        let router = app(auth_state(false, None)).expect("router");
1904        let resp = router.oneshot(cap_request(None)).await.unwrap();
1905        // Default-off: no key required, request proceeds to the handler.
1906        assert_eq!(resp.status(), axum::http::StatusCode::OK);
1907    }
1908
1909    #[tokio::test]
1910    async fn auth_on_missing_key_is_401_opaque() {
1911        use tower::ServiceExt;
1912        let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1913        let keys = format!("{{\"tenant-a\": {hash:?}}}");
1914        let router = app(auth_state(true, Some(keys))).expect("router");
1915
1916        let resp = router.oneshot(cap_request(None)).await.unwrap();
1917        assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
1918        let json = body_json(resp).await;
1919        assert_eq!(json["error"]["type"], "unauthorized");
1920        // Opaque: the body must not name tenants or hint which key would work.
1921        let msg = json["error"]["message"].as_str().unwrap();
1922        assert!(!msg.contains("tenant"), "no tenant oracle in body: {msg}");
1923    }
1924
1925    #[tokio::test]
1926    async fn auth_on_invalid_key_is_401() {
1927        use tower::ServiceExt;
1928        let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1929        let keys = format!("{{\"tenant-a\": {hash:?}}}");
1930        let router = app(auth_state(true, Some(keys))).expect("router");
1931
1932        let resp = router
1933            .oneshot(cap_request(Some("Bearer wrong-key")))
1934            .await
1935            .unwrap();
1936        assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
1937    }
1938
1939    #[tokio::test]
1940    async fn auth_on_valid_key_proceeds() {
1941        use tower::ServiceExt;
1942        let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1943        let keys = format!("{{\"tenant-a\": {hash:?}}}");
1944        let router = app(auth_state(true, Some(keys))).expect("router");
1945
1946        // Keyed format: `<tenant_id>.<secret>`.
1947        let resp = router
1948            .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
1949            .await
1950            .unwrap();
1951        // A valid key clears the middleware and reaches the handler.
1952        assert_eq!(resp.status(), axum::http::StatusCode::OK);
1953    }
1954
1955    /// Two tenants, keyed so requests carry a real, distinct `TenantId` (ADR 0004 §D6).
1956    fn two_tenant_state(rate_per_sec: Option<u32>) -> AppState {
1957        let hash_a = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1958        let hash_b = crate::tenant_auth::TenantKeys::hash_key("key-b").unwrap();
1959        let keys = format!("{{\"tenant-a\": {hash_a:?}, \"tenant-b\": {hash_b:?}}}");
1960        auth_state_rated(true, Some(keys), rate_per_sec)
1961    }
1962
1963    #[tokio::test]
1964    async fn tenant_exceeding_rate_limit_gets_429_opaque() {
1965        use tower::ServiceExt;
1966        // Burst capacity == rate for `Quota::per_second`, so 1 req/sec allows exactly 1 request
1967        // before the next is rejected. Only 2 requests (not 3+) to stay well clear of the 1s
1968        // replenish window even under real per-request Argon2-verify latency in the auth layer.
1969        let router = app(two_tenant_state(Some(1))).expect("router");
1970
1971        let r1 = router
1972            .clone()
1973            .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
1974            .await
1975            .unwrap();
1976        assert_eq!(r1.status(), axum::http::StatusCode::OK);
1977
1978        let r2 = router
1979            .clone()
1980            .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
1981            .await
1982            .unwrap();
1983        assert_eq!(r2.status(), axum::http::StatusCode::TOO_MANY_REQUESTS);
1984        let json = body_json(r2).await;
1985        assert_eq!(json["error"]["type"], "rate_limited");
1986        // Opaque: no bucket state or limit value leaked to the caller.
1987        let msg = json["error"]["message"].as_str().unwrap();
1988        assert!(!msg.contains('1'), "no limit value in body: {msg}");
1989    }
1990
1991    #[tokio::test]
1992    async fn rate_limit_buckets_are_independent_per_tenant() {
1993        use tower::ServiceExt;
1994        let router = app(two_tenant_state(Some(1))).expect("router");
1995
1996        // Tenant A exhausts its 1 req/sec budget...
1997        let a1 = router
1998            .clone()
1999            .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
2000            .await
2001            .unwrap();
2002        assert_eq!(a1.status(), axum::http::StatusCode::OK);
2003        let a2 = router
2004            .clone()
2005            .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
2006            .await
2007            .unwrap();
2008        assert_eq!(a2.status(), axum::http::StatusCode::TOO_MANY_REQUESTS);
2009
2010        // ...but tenant B, on the same gate/route, is unaffected (independent bucket).
2011        let b1 = router
2012            .clone()
2013            .oneshot(cap_request(Some("Bearer tenant-b.key-b")))
2014            .await
2015            .unwrap();
2016        assert_eq!(b1.status(), axum::http::StatusCode::OK);
2017    }
2018
2019    #[tokio::test]
2020    async fn rate_limit_unset_never_429s() {
2021        use tower::ServiceExt;
2022        // Backward-compat: with no FIRSTPASS_TENANT_RATE_PER_SEC, drive many requests through and
2023        // confirm none are ever rate-limited (default-off).
2024        let router = app(two_tenant_state(None)).expect("router");
2025        for _ in 0..20 {
2026            let resp = router
2027                .clone()
2028                .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
2029                .await
2030                .unwrap();
2031            assert_eq!(resp.status(), axum::http::StatusCode::OK);
2032        }
2033    }
2034}