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, Dialect, FEATURE_VERSION, Features, FinalOutcome, GENESIS_HASH, Mode,
19    ModelRef, PolicyRef, ProbeRegime, ProbeSignal, RequestInfo, RoutingMode, Score, ServedFrom,
20    TaskKind, Trace, Verdict,
21};
22use serde::Deserialize;
23use serde_json::Value;
24use std::future::Future;
25use std::time::Duration;
26use tokio::sync::mpsc::error::TrySendError;
27use uuid::Uuid;
28
29use crate::config::ProxyConfig;
30use crate::error::ProxyError;
31use crate::gate::{GateHealthRegistry, aggregate_with_policy, resolve_gates};
32use crate::provider::{Auth, ChatMessage, ModelRequest, ModelResponse, ProviderRegistry};
33use crate::router::{EnforceCtx, EngineOutcome, route_enforce};
34use crate::store;
35use crate::tenant_auth::{TenantId, auth_middleware};
36use crate::upstream::{
37    forward_anthropic, forward_anthropic_streaming, forward_openai, forward_openai_streaming,
38};
39use firstpass_core::Route;
40use firstpass_core::trace::ShadowSignal;
41
42/// Shared state handed to every request handler. Cheap to clone: an `Arc`ed config, a
43/// pooled HTTP client, and a bounded channel sender.
44#[derive(Clone)]
45pub struct AppState {
46    /// Static proxy configuration.
47    pub config: Arc<ProxyConfig>,
48    /// Shared, connection-pooled HTTP client used to call upstream (observe passthrough).
49    pub http: reqwest::Client,
50    /// Multi-provider registry used by the enforce-mode escalation engine.
51    pub providers: ProviderRegistry,
52    /// Per-gate error budgets (auto-disable), shared across requests.
53    pub gate_health: Arc<GateHealthRegistry>,
54    /// Per-(tenant, route) shadow spend for the current UTC day (ADR 0009 D2). Shadow makes real
55    /// model calls, so the daily ceiling is enforced through this rather than trusted.
56    pub shadow_ledger: Arc<crate::shadow::ShadowLedger>,
57    /// Per-(tenant, route) guardrail state (ADR 0009 D3): the trailing window of resolved
58    /// outcomes and whether a route is currently demoted.
59    pub guardrails: Arc<crate::guard::GuardrailRegistry>,
60    /// Fire-and-forget sender to the background trace writer.
61    pub traces: store::TraceSender,
62    /// Optional online/adaptive conformal serve threshold (Gibbs-Candès ACI). `None` = fixed
63    /// `serve_threshold` from config (default). When present, `/v1/feedback` nudges it live and the
64    /// enforce path reads its current value per request — the reactive, self-tuning loop.
65    pub adaptive: Option<Arc<std::sync::Mutex<firstpass_core::conformal::AdaptiveConformal>>>,
66    /// Optional UCB1 start-rung bandit (predict-to-start, verify-to-serve). `None` (default) =
67    /// start every request at rung 0, byte-identical to today. When present, `handle_enforce`
68    /// queries it for a predicted start rung per request and feeds back gate verdicts for online
69    /// learning — all in-memory, per-process.
70    pub bandit: Option<Arc<std::sync::Mutex<crate::bandit::StartRungBandit>>>,
71    /// Optional per-query gate-pass predictor (ADR 0008 Phase 2). `None` (default) = no
72    /// prediction, byte-identical to today. When `Some`, `handle_enforce` records its
73    /// `P(gate-pass)` for the start rung on the receipt in **shadow** (never acted on) and
74    /// feeds this request's attempts back for online learning — in-memory, per-process, warm-
75    /// started from receipts on boot.
76    pub predictor: Option<Arc<std::sync::Mutex<firstpass_core::PassPredictor>>>,
77    /// Per-tenant request rate limiter (ADR 0004 §D6). `None` (the default) disables rate
78    /// limiting entirely — set via [`build_tenant_rate_limiter`] from
79    /// [`ProxyConfig::tenant_rate_per_sec`].
80    pub tenant_rate_limiter: Option<Arc<governor::DefaultKeyedRateLimiter<String>>>,
81    /// Durable-receipts spill handle (`FIRSTPASS_RECEIPTS=durable`). `None` in best-effort mode
82    /// (the default) — behavior is byte-identical to before. When `Some`, `offer_trace` appends
83    /// to `<db_path>.spill.jsonl` on channel-full instead of dropping.
84    pub spill: Option<store::SpillHandle>,
85}
86
87/// Build the per-tenant keyed rate limiter from config (ADR 0004 §D6). Returns `None` when
88/// `FIRSTPASS_TENANT_RATE_PER_SEC` is unset (the default) — single-operator and existing
89/// deployments see no limiter and no behavior change.
90#[must_use]
91pub fn build_tenant_rate_limiter(
92    config: &ProxyConfig,
93) -> Option<Arc<governor::DefaultKeyedRateLimiter<String>>> {
94    let per_sec = config.tenant_rate_per_sec?;
95    Some(Arc::new(governor::RateLimiter::keyed(
96        governor::Quota::per_second(per_sec),
97    )))
98}
99
100/// Axum middleware (ADR 0004 §D6): enforce the per-tenant request rate limit. Must run AFTER
101/// [`auth_middleware`] so the resolved [`TenantId`] is already in request extensions. A no-op
102/// (never returns 429) when [`AppState::tenant_rate_limiter`] is `None`.
103pub async fn tenant_rate_limit_middleware(
104    State(state): State<AppState>,
105    Extension(tenant): Extension<TenantId>,
106    req: Request,
107    next: Next,
108) -> Response {
109    if let Some(limiter) = &state.tenant_rate_limiter
110        && limiter.check_key(&tenant.0).is_err()
111    {
112        return ProxyError::RateLimited.into_response();
113    }
114    next.run(req).await
115}
116
117impl std::fmt::Debug for AppState {
118    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
119        f.debug_struct("AppState")
120            .field("config", &self.config)
121            .finish_non_exhaustive()
122    }
123}
124
125/// Fire-and-forget a trace at the background writer: non-blocking, and bounded.
126///
127/// **Best-effort mode** (`spill` is `None`): if the writer has fallen behind enough to fill the
128/// buffer, the trace is dropped with a warning rather than blocking the hot path or growing memory
129/// without limit (the audit chain over persisted traces stays valid; a dropped trace is simply
130/// absent).
131///
132/// **Durable mode** (`spill` is `Some`): on `TrySendError::Full` the trace is serialised as a
133/// JSON line and appended (with `sync_data`) to the spill file. This blocks the calling task on a
134/// disk write — the deliberate tradeoff of durable mode; it only fires under sustained
135/// backpressure. The writer drains the spill file at startup and on channel-empty so the chain
136/// stays valid.
137///
138/// ponytail: the spill write holds `Mutex<File>` across a `sync_data` call on the calling tokio
139/// task — fine for the slow backpressure path; use `spawn_blocking` if disk latency at p99 under
140/// sustained overload is measurable.
141fn offer_trace(traces: &store::TraceSender, spill: Option<&store::SpillHandle>, trace: Trace) {
142    record_trace_metrics(&trace);
143    match traces.try_send(trace) {
144        Ok(()) => {}
145        Err(TrySendError::Full(t)) => {
146            if let Some(handle) = spill {
147                match store::append_to_spill(handle, &t) {
148                    Ok(()) => {
149                        metrics::counter!("firstpass_receipts_spilled_total").increment(1);
150                    }
151                    Err(e) => {
152                        tracing::error!(%e, "durable mode: spill write failed; trace lost");
153                        metrics::counter!("firstpass_traces_dropped_total").increment(1);
154                    }
155                }
156            } else {
157                tracing::warn!("trace channel full; dropping trace (writer behind under load)");
158                metrics::counter!("firstpass_traces_dropped_total").increment(1);
159            }
160        }
161        Err(TrySendError::Closed(_)) => {
162            tracing::warn!("trace writer is gone; dropping trace");
163        }
164    }
165}
166
167/// Record the real signals every trace carries: enforce-mode latency/escalations (observe mode
168/// forwards unchanged, so its wall-clock time isn't a routing-decision latency), and what got
169/// served — regardless of mode, since an upstream failure is worth counting either way.
170fn record_trace_metrics(trace: &Trace) {
171    if trace.mode == Mode::Enforce {
172        metrics::histogram!("firstpass_enforce_latency_ms")
173            .record(trace.final_.total_latency_ms as f64);
174        if trace.final_.escalations > 0 {
175            metrics::counter!("firstpass_escalations_total")
176                .increment(u64::from(trace.final_.escalations));
177        }
178    }
179    let served_from = match trace.final_.served_from {
180        ServedFrom::Attempt => "attempt",
181        ServedFrom::BestAttempt => "best_attempt",
182        ServedFrom::Error => "error",
183    };
184    metrics::counter!("firstpass_served_total", "served_from" => served_from).increment(1);
185    if trace.final_.served_from == ServedFrom::Error {
186        metrics::counter!("firstpass_upstream_failures_total").increment(1);
187    }
188    // The value signals: what was spent, what proof cost, and what routing saved vs always-top
189    // (§9.1 counterfactual). Monotonic gauges because `metrics` counters are integer-only and
190    // these are USD floats; scrape-side `rate()`/`increase()` work the same.
191    metrics::gauge!("firstpass_cost_usd_total").increment(trace.final_.total_cost_usd);
192    metrics::gauge!("firstpass_gate_cost_usd_total").increment(trace.final_.gate_cost_usd);
193    metrics::gauge!("firstpass_baseline_usd_total")
194        .increment(trace.final_.counterfactual_baseline_usd);
195    metrics::gauge!("firstpass_savings_usd_total").increment(trace.final_.savings_usd);
196    // Which rung actually served, labeled by model — the shape of the ladder in production.
197    if let Some(rung) = trace.final_.served_rung {
198        let model = trace
199            .attempts
200            .iter()
201            .find(|a| a.rung == rung)
202            .map(|a| a.model.clone())
203            .unwrap_or_else(|| "unknown".to_owned());
204        metrics::counter!(
205            "firstpass_served_rung_total",
206            "rung" => rung.to_string(),
207            "model" => model
208        )
209        .increment(1);
210    }
211}
212
213/// Max accepted request body. Explicit (not axum's ~2 MB default) so it's an intentional ceiling:
214/// generous enough to pass through large multimodal/long-context requests, bounded so a single
215/// oversized body can't exhaust memory.
216const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
217
218// ── Epsilon-greedy helpers ────────────────────────────────────────────────────
219
220/// Map a `u128` seed to a uniform float in `[0, 1)` via two SplitMix64 finalizer rounds.
221///
222/// Used to derive the per-request epsilon-greedy draw from `Uuid::now_v7().as_u128()` — no
223/// new dependencies needed. The two 64-bit halves are finalised separately then XOR-folded
224/// to a single u64 to mix time and random UUID bits.
225///
226/// ponytail: not a general-purpose RNG; replace with `rand` if more draws per request
227/// are ever needed.
228pub(crate) fn u01(seed: u128) -> f64 {
229    let lo = splitmix64_finalise(seed as u64);
230    let hi = splitmix64_finalise((seed >> 64) as u64);
231    // 53-bit mantissa of f64 → uniform on [0, 1).
232    ((lo ^ hi) >> 11) as f64 * (1.0_f64 / (1u64 << 53) as f64)
233}
234
235fn splitmix64_finalise(mut z: u64) -> u64 {
236    z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
237    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
238    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
239    z ^ (z >> 31)
240}
241
242/// Propensity of the logging policy choosing `chosen` under epsilon-greedy over `k` rungs
243/// where `greedy` is the deterministic choice.
244///
245/// `p = (1 − ε) · 𝟙[chosen == greedy] + ε / K`
246///
247/// Both terms apply when the epsilon branch fires and coincidentally lands on the greedy rung.
248#[must_use]
249pub(crate) fn epsilon_propensity(chosen: u32, greedy: u32, epsilon: f64, k: usize) -> f64 {
250    let greedy_term = if chosen == greedy { 1.0 - epsilon } else { 0.0 };
251    greedy_term + epsilon / k as f64
252}
253
254/// Build the axum router: `POST /v1/messages`, `GET /v1/capabilities`, `GET /healthz`,
255/// `GET /metrics`.
256///
257/// # Errors
258/// [`ProxyError::Internal`] if the Prometheus recorder fails to install (see
259/// [`crate::metrics::install`]).
260pub fn app(state: AppState) -> Result<Router, ProxyError> {
261    crate::metrics::install()?;
262    let max_concurrency = state.config.max_concurrency;
263
264    // Tenant-facing business routes: every one runs the auth middleware, which injects the resolved
265    // `TenantId` into request extensions (the authenticated tenant when `require_auth` is on, the
266    // static default when off — ADR 0004 §D1/§D2). Operator routes (`/healthz`, `/metrics`) are
267    // NOT tenant-facing and stay outside the auth layer.
268    // Per-tenant rate limit (ADR 0004 §D6) runs INSIDE (after) the auth layer below — axum layers
269    // wrap outward-in, so a layer added earlier in the chain executes later on the request path —
270    // so the resolved `TenantId` is already in extensions when this middleware checks it. A no-op
271    // when `FIRSTPASS_TENANT_RATE_PER_SEC` is unset.
272    let business = Router::new()
273        .route("/v1/messages", post(messages))
274        .route("/v1/chat/completions", post(chat_completions))
275        .route("/v1/feedback", post(feedback))
276        .route("/v1/capabilities", get(capabilities))
277        .layer(axum::middleware::from_fn_with_state(
278            state.clone(),
279            tenant_rate_limit_middleware,
280        ))
281        .layer(axum::middleware::from_fn_with_state(
282            state.clone(),
283            auth_middleware,
284        ));
285
286    Ok(Router::new()
287        .merge(business)
288        .route("/healthz", get(healthz))
289        .route("/metrics", get(crate::metrics::handler))
290        // Explicit body-size ceiling (DoS/OOM guard) across every route.
291        .layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES))
292        // Concurrency load-shed: cap in-flight requests under the cap rather than falling over.
293        // Deliberately NOT a request timeout — that would sever in-flight SSE streams.
294        .layer(tower::limit::GlobalConcurrencyLimitLayer::new(
295            max_concurrency,
296        ))
297        .with_state(state))
298}
299
300/// `GET /healthz` — liveness probe.
301async fn healthz() -> impl IntoResponse {
302    Json(serde_json::json!({ "status": "ok" }))
303}
304
305/// `GET /v1/capabilities` — agent-first discovery (SPEC §0.2, §7.4): what this proxy speaks,
306/// which modes are live, the first enforce route's ladder/gates, and how to turn it off.
307async fn capabilities(State(state): State<AppState>) -> impl IntoResponse {
308    // Report the first enforce route's ladder + gates, so an agent can discover what it's routed
309    // through. Empty when no routing config is loaded (pure observe deployment).
310    let (ladder, gates) = state
311        .config
312        .routing
313        .as_ref()
314        .and_then(|c| c.routes.iter().find(|r| r.mode == Mode::Enforce))
315        .map(|r| (r.ladder.clone(), r.gates.clone()))
316        .unwrap_or_default();
317    let routing_modes: Vec<serde_json::Value> = RoutingMode::ALL
318        .iter()
319        .map(|m| {
320            let p = m.preset();
321            serde_json::json!({
322                "name": m.as_str(),
323                "description": p.description,
324                "tradeoff": p.tradeoff,
325            })
326        })
327        .collect();
328    Json(serde_json::json!({
329        "service": "firstpass",
330        "version": env!("CARGO_PKG_VERSION"),
331        "feature_version": FEATURE_VERSION,
332        "modes": ["observe", "enforce"],
333        "routing_modes": routing_modes,
334        "wire_apis": ["anthropic.messages", "openai.chat_completions"],
335        "ladder": ladder,
336        "gates": gates,
337        "feedback_api": "POST /v1/feedback",
338        "offboarding": "unset ANTHROPIC_BASE_URL (or OPENAI_BASE_URL for OpenAI clients)",
339    }))
340}
341
342/// Body of `POST /v1/feedback`: a downstream outcome reported for a past decision.
343#[derive(Debug, Deserialize)]
344struct FeedbackRequest {
345    /// The `trace_id` of the decision this outcome is about.
346    trace_id: String,
347    /// The gate/source id, e.g. `"tests"` or `"feedback:ci"`.
348    gate_id: String,
349    /// `"pass"` | `"fail"` | `"abstain"`.
350    verdict: String,
351    /// Optional confidence in `[0, 1]`.
352    #[serde(default)]
353    score: Option<f64>,
354    /// Who reported it (a CI system, a human reviewer, a deferred gate).
355    reporter: String,
356}
357
358/// `POST /v1/feedback` — attach a downstream outcome (deferred verdict) to a past trace, closing
359/// the outcome-feedback loop (SPEC §8.3.4). The verdict is stored in a **separate** table keyed
360/// by `trace_id`; the sealed, hashed trace is never mutated, so the audit chain stays verifiable.
361/// Returns `202 Accepted`. This is the signal that later calibrates the gates.
362async fn feedback(
363    State(state): State<AppState>,
364    Extension(TenantId(tenant)): Extension<TenantId>,
365    body: Bytes,
366) -> Response {
367    let req: FeedbackRequest = match serde_json::from_slice(&body) {
368        Ok(r) => r,
369        Err(e) => {
370            return ProxyError::BadRequest(format!("invalid feedback body: {e}")).into_response();
371        }
372    };
373    let verdict = match req.verdict.as_str() {
374        "pass" => Verdict::Pass,
375        "fail" => Verdict::Fail,
376        "abstain" => Verdict::Abstain,
377        other => {
378            return ProxyError::BadRequest(format!("unknown verdict {other:?}")).into_response();
379        }
380    };
381    let score = match req.score {
382        Some(s) => match Score::new(s) {
383            Ok(sc) => Some(sc),
384            Err(_) => {
385                return ProxyError::BadRequest(format!("score {s} out of range [0,1]"))
386                    .into_response();
387            }
388        },
389        None => None,
390    };
391
392    let db = state.config.db_path.clone();
393
394    // Reject feedback for an unknown trace, so orphan outcomes can't accumulate — AND deny
395    // cross-tenant feedback (IDOR, ADR 0004 §D4). `trace_exists` is scoped to the caller's tenant,
396    // so a trace owned by another tenant is indistinguishable from a missing one: both return `404`
397    // (never `403`, which would be an existence oracle).
398    let (db_check, tenant_check, tid_check) = (db.clone(), tenant.clone(), req.trace_id.clone());
399    match tokio::task::spawn_blocking(move || {
400        store::trace_exists(&db_check, &tenant_check, &tid_check)
401    })
402    .await
403    {
404        Ok(Ok(true)) => {}
405        Ok(Ok(false)) => {
406            return ProxyError::NotFound(format!("unknown trace_id {:?}", req.trace_id))
407                .into_response();
408        }
409        Ok(Err(e)) => {
410            tracing::error!(%e, "feedback: trace_exists check failed");
411            return ProxyError::Internal(e.to_string()).into_response();
412        }
413        Err(e) => {
414            tracing::error!(%e, "feedback: trace_exists task panicked");
415            return ProxyError::Internal(e.to_string()).into_response();
416        }
417    }
418
419    // Correctness signal for the online adaptive loop — only a clear Pass/Fail nudges the threshold.
420    let feedback_signal = match verdict {
421        Verdict::Pass => Some(true),
422        Verdict::Fail => Some(false),
423        Verdict::Abstain => None,
424    };
425    let dv = DeferredVerdict {
426        gate_id: req.gate_id,
427        verdict,
428        score,
429        reported_at: jiff::Timestamp::now(),
430        reporter: req.reporter,
431    };
432    let trace_id = req.trace_id.clone();
433    match tokio::task::spawn_blocking(move || store::append_deferred(&db, &req.trace_id, &dv)).await
434    {
435        Ok(Ok(())) => {
436            // Close the reactive loop: nudge the live serve threshold toward the target.
437            if let (Some(a), Some(correct)) = (state.adaptive.as_ref(), feedback_signal)
438                && let Ok(mut g) = a.lock()
439            {
440                g.observe_served(correct);
441                metrics::gauge!("firstpass_serve_threshold").set(g.threshold());
442                metrics::gauge!("firstpass_realized_served_failure")
443                    .set(g.realized_served_failure());
444            }
445
446            // Which route produced the decision this outcome is about. Read from the trace
447            // rather than assumed, and tenant-scoped by the same ownership check `trace_exists`
448            // makes. Traces written before route recording fall back to 0 — the previous
449            // behaviour — so old logs keep working rather than being silently dropped.
450            let attributed_route = tokio::task::spawn_blocking({
451                let (db, ten, tid) = (
452                    state.config.db_path.clone(),
453                    tenant.clone(),
454                    trace_id.clone(),
455                );
456                move || store::trace_route(&db, &ten, &tid)
457            })
458            .await
459            .ok()
460            .and_then(Result::ok)
461            .flatten()
462            .flatten()
463            .unwrap_or(0) as usize;
464
465            // Feed the guardrail (ADR 0009 D3). This is deliberately the ONLY thing that feeds
466            // it: a resolved downstream outcome, not a gate verdict. Gate verdicts are
467            // Firstpass's own opinion, and a guardrail fed on them would be grading its own
468            // homework — it would keep enforcing precisely when its own judgement had drifted.
469            if let (Some(cfg), Some(correct)) = (
470                state.config.routing.as_ref().and_then(|r| r.guardrail),
471                feedback_signal,
472            ) {
473                let cooldown = state
474                    .config
475                    .routing
476                    .as_ref()
477                    .map_or(3_600, |r| r.guardrail_cooldown_secs);
478                let reaction = state.guardrails.record(
479                    &tenant,
480                    attributed_route,
481                    &cfg,
482                    correct,
483                    jiff::Timestamp::now().as_second(),
484                    cooldown,
485                );
486                match &reaction {
487                    crate::guard::Reaction::Demoted(v) => {
488                        metrics::counter!("firstpass_guardrail_demotions_total").increment(1);
489                        tracing::error!(
490                            tenant = %tenant,
491                            n = v.n,
492                            rate = v.rate,
493                            bound = v.bound,
494                            alpha = cfg.alpha,
495                            "GUARDRAIL: served-failure bound exceeded target — route demoted to                              observe; traffic now serves as it would without Firstpass"
496                        );
497                    }
498                    crate::guard::Reaction::Alarmed(v) => {
499                        metrics::counter!("firstpass_guardrail_breaches_total").increment(1);
500                        tracing::error!(
501                            tenant = %tenant,
502                            n = v.n,
503                            rate = v.rate,
504                            bound = v.bound,
505                            alpha = cfg.alpha,
506                            "GUARDRAIL: served-failure bound exceeded target (alarm only —                              routing unchanged)"
507                        );
508                    }
509                    crate::guard::Reaction::None => {}
510                }
511            }
512            (
513                axum::http::StatusCode::ACCEPTED,
514                Json(serde_json::json!({ "status": "recorded", "trace_id": trace_id })),
515            )
516                .into_response()
517        }
518        Ok(Err(e)) => {
519            tracing::error!(%e, "feedback: append_deferred failed");
520            ProxyError::Internal(e.to_string()).into_response()
521        }
522        Err(e) => {
523            tracing::error!(%e, "feedback: append_deferred task panicked");
524            ProxyError::Internal(e.to_string()).into_response()
525        }
526    }
527}
528
529/// The header a caller may set to group requests into a session for the audit trail. When
530/// absent, each request is its own session (keyed by its own trace id).
531const SESSION_HEADER: &str = "x-firstpass-session";
532
533/// Header carrying the calling agent identity (feature/routing signal).
534const AGENT_HEADER: &str = "x-firstpass-agent";
535/// Header carrying the calling subagent identity.
536const SUBAGENT_HEADER: &str = "x-firstpass-subagent";
537/// Kick off a shadow evaluation for an observed request, detached (ADR 0009 D2).
538///
539/// Called only once the caller's response already exists, so nothing here can change what was
540/// served, how long it took, or whether it succeeded. The join handle is dropped deliberately —
541/// nothing waits on this, because a measurement must never be able to delay a served answer.
542#[expect(
543    clippy::too_many_arguments,
544    reason = "mirrors the enforce context it evaluates; a wrapper struct would only move the list"
545)]
546fn spawn_shadow(
547    state: &AppState,
548    route: &firstpass_core::Route,
549    route_ix: usize,
550    body: Bytes,
551    auth: Auth,
552    features: Features,
553    tenant: String,
554    session_id: String,
555    api: &str,
556) {
557    let Some(shadow) = route.shadow else {
558        return;
559    };
560    // Keyed on the session so a conversation is consistently in or out of the sample, and under
561    // its own hash tag so the sample does not correlate with the rollout arm.
562    if !shadow.sampled(&state.config.prompt_salt, &session_id) {
563        return;
564    }
565    let (state, route, api) = (state.clone(), route.clone(), api.to_owned());
566    tokio::spawn(async move {
567        let signal = evaluate_shadow(
568            &state, &route, shadow, route_ix, &body, auth, &features, tenant, session_id, &api, 0.0,
569        )
570        .await;
571        // A shadow failure is data, not an incident: recorded, never surfaced to the caller.
572        tracing::debug!(
573            would_pass = signal.would_pass,
574            projected_usd = signal.projected_cost_usd,
575            skipped = ?signal.skipped,
576            "shadow evaluation complete"
577        );
578    });
579}
580
581/// Run one shadow evaluation and return the counterfactual signal (ADR 0009 D2).
582///
583/// Called from a detached task *after* the observed response has already been handed to the
584/// caller, so nothing here can affect what was served, its timing, or its bytes. Every failure
585/// path returns a signal describing the failure rather than propagating: a measurement must never
586/// be able to take down a request path that already succeeded.
587#[expect(
588    clippy::too_many_arguments,
589    reason = "mirrors the enforce context; grouping them into a struct would only move the list"
590)]
591async fn evaluate_shadow(
592    state: &AppState,
593    route: &firstpass_core::Route,
594    shadow: firstpass_core::rollout::Shadow,
595    route_ix: usize,
596    body: &Bytes,
597    auth: Auth,
598    features: &Features,
599    tenant: String,
600    session_id: String,
601    api: &str,
602    actual_cost_usd: f64,
603) -> ShadowSignal {
604    let skipped = |why: &str| ShadowSignal {
605        would_serve_rung: None,
606        would_pass: false,
607        projected_cost_usd: 0.0,
608        actual_cost_usd,
609        skipped: Some(why.to_owned()),
610    };
611
612    let now = jiff::Timestamp::now();
613    if !state
614        .shadow_ledger
615        .may_spend(&tenant, route_ix, shadow.max_usd_per_day, now)
616    {
617        // Recorded, not silent: an operator whose projection quietly stopped tracking would keep
618        // trusting a number that no longer describes their traffic.
619        return skipped("budget_exhausted");
620    }
621
622    let Some(base_request) = parse_model_request(body) else {
623        return skipped("unparseable_request");
624    };
625    let gate_defs = state
626        .config
627        .routing
628        .as_ref()
629        .map_or(&[][..], |cfg| &cfg.gate_defs);
630    let gates = resolve_gates(
631        &route.gates,
632        gate_defs,
633        &state.providers,
634        &auth,
635        &state.config.prices,
636    );
637    let (budget, max_rungs) = state
638        .config
639        .routing
640        .as_ref()
641        .map_or((None, u32::MAX), |cfg| {
642            (
643                cfg.budget.per_request_usd,
644                cfg.escalation.max_rungs_per_request,
645            )
646        });
647
648    let ctx = EnforceCtx {
649        ladder: &route.ladder,
650        gates: &gates,
651        health: &state.gate_health,
652        base_request: &base_request,
653        providers: &state.providers,
654        auth: &auth,
655        prices: &state.config.prices,
656        budget_per_request_usd: budget,
657        max_rungs,
658        // Shadow is a measurement, not a latency-sensitive serve: no speculation (it would spend
659        // more to save time nobody is waiting on) and no elastic verification.
660        speculation: 0,
661        serve_threshold: None,
662        elastic: None,
663        features: features.clone(),
664        start_rung: 0,
665        tenant_id: tenant.clone(),
666        session_id,
667        prompt_hash: prompt_hash(&state.config.prompt_salt, body),
668        api: api.to_owned(),
669        policy_id: "shadow".to_owned(),
670    };
671
672    let (outcome, trace) = route_enforce(ctx).await;
673    let spent = trace.final_.total_cost_usd;
674    state.shadow_ledger.debit(&tenant, route_ix, spent, now);
675
676    // The served rung is recorded on the trace's final outcome; read it there rather than
677    // re-deriving it from the attempts, so shadow and enforce agree by construction.
678    let would_pass = matches!(outcome, EngineOutcome::Served(_));
679    let would_serve_rung = if would_pass {
680        trace.final_.served_rung
681    } else {
682        None
683    };
684    ShadowSignal {
685        would_serve_rung,
686        would_pass,
687        projected_cost_usd: spent,
688        actual_cost_usd,
689        skipped: None,
690    }
691}
692
693/// Per-request routing-mode override. Case-insensitive; unknown values are logged and ignored
694/// (fall through to route-level / global-default). Valid values: observe|cost|balanced|quality|latency|max.
695const MODE_PROFILE_HEADER: &str = "x-firstpass-mode";
696
697/// Resolve the effective [`RoutingMode`] for this request.
698///
699/// Precedence (highest first):
700/// 1. `x-firstpass-mode` request header (case-insensitive; unknown values → warn + fall through)
701/// 2. `route.routing_mode` (per-route config)
702/// 3. `config.default_routing_mode` (global `FIRSTPASS_MODE_PROFILE` env var, default `Balanced`)
703///
704/// When nothing is set, returns `Balanced` — a strict no-op over existing config.
705fn resolve_mode(headers: &HeaderMap, route: &Route, config: &ProxyConfig) -> RoutingMode {
706    // (a) per-request header wins over everything
707    if let Some(val) = header_str(headers, MODE_PROFILE_HEADER) {
708        match val.trim().to_ascii_lowercase().as_str() {
709            "observe" => return RoutingMode::Observe,
710            "cost" => return RoutingMode::Cost,
711            "balanced" => return RoutingMode::Balanced,
712            "quality" => return RoutingMode::Quality,
713            "latency" => return RoutingMode::Latency,
714            "max" => return RoutingMode::Max,
715            other => {
716                tracing::warn!(
717                    value = other,
718                    "unknown x-firstpass-mode value; ignoring \
719                     (valid: observe|cost|balanced|quality|latency|max)"
720                );
721            }
722        }
723    }
724    // (b) per-route config
725    if let Some(m) = route.routing_mode {
726        return m;
727    }
728    // (c) global default (env FIRSTPASS_MODE_PROFILE, default Balanced)
729    config.default_routing_mode
730}
731
732/// `POST /v1/messages` — dispatch on the matched route's mode. **Enforce** routes run the
733/// escalation engine (gate + escalate + failover); everything else is an **observe**
734/// passthrough (forward unchanged, trace asynchronously). Either way the trace is recorded
735/// off the response path.
736async fn messages(
737    State(state): State<AppState>,
738    Extension(TenantId(tenant)): Extension<TenantId>,
739    headers: HeaderMap,
740    body: Bytes,
741) -> Response {
742    let session_header = header_str(&headers, SESSION_HEADER);
743
744    // Only parse the request for routing when a routing config is loaded — an observe-only
745    // deployment does zero on-path parsing and keeps its zero-added-latency guarantee.
746    if let Some(routing) = state.config.routing.as_ref() {
747        let features = extract_features(&headers, &body);
748        if let Some(route) = routing
749            .route_for(&features)
750            .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
751        {
752            // Index of the matched route, so shadow spend is budgeted per route rather than
753            // pooled across a config that may define several.
754            let route_ix = routing
755                .routes
756                .iter()
757                .position(|r| std::ptr::eq(r, route))
758                .unwrap_or(0);
759            // Clone the matched route so no borrow of `state.config` is held across the await;
760            // routes are tiny (a handful of strings).
761            let route = route.clone();
762            // Resolve routing-mode preset (header > route > global default).
763            let routing_mode = resolve_mode(&headers, &route, &state.config);
764            // Observe mode forces the observe passthrough path — no gating, no escalation.
765            if routing_mode == RoutingMode::Observe {
766                return observe_passthrough(state, headers, body, session_header, tenant).await;
767            }
768            // Guardrail demotion (ADR 0009 D3). A route whose served-failure bound has breached
769            // serves exactly as it would without Firstpass until the cooldown lapses. This is
770            // checked BEFORE rollout so a demotion cannot be partially overridden by an arm.
771            if state
772                .guardrails
773                .is_demoted(&tenant, route_ix, jiff::Timestamp::now().as_second())
774            {
775                return observe_passthrough(state, headers, body, session_header, tenant).await;
776            }
777            // Percentage rollout (ADR 0009 D1). Bucketing is a pure function of a STABLE key, so
778            // a conversation stays in one arm for its whole life — a per-request draw would flip
779            // a multi-turn agent mid-thread and, worse, would make the served population an
780            // unstable sample, which is exactly the population the published bound is over.
781            if let Some(rollout) = route.rollout.as_ref() {
782                let key_value = match rollout.key {
783                    firstpass_core::RolloutKey::Session => session_header
784                        .clone()
785                        .unwrap_or_else(|| firstpass_core::rollout::request_identity(&body)),
786                    // No session to hold constant, so each request is bucketed independently.
787                    // NOTE: unlike the session and tenant keys — whose values are recorded on the
788                    // trace — this bucket is recorded but NOT recomputable by an auditor, because
789                    // the input it hashes is the request body and raw prompts are deliberately
790                    // never stored. The recorded `bucket`/`enforced` remain authoritative.
791                    firstpass_core::RolloutKey::Request => {
792                        firstpass_core::rollout::request_identity(&body)
793                    }
794                    firstpass_core::RolloutKey::Tenant => tenant.to_string(),
795                };
796                let decision =
797                    firstpass_core::rollout::decide(&state.config.prompt_salt, rollout, &key_value);
798                if !decision.enforced {
799                    // The control arm. Served exactly as it would be without Firstpass.
800                    let shadow_ctx = route.shadow.map(|_| {
801                        (
802                            body.clone(),
803                            Auth::from_headers(&headers),
804                            features.clone(),
805                            tenant.clone(),
806                            session_header
807                                .clone()
808                                .unwrap_or_else(|| Uuid::now_v7().to_string()),
809                        )
810                    });
811                    let resp =
812                        observe_passthrough(state.clone(), headers, body, session_header, tenant)
813                            .await;
814                    // Strictly after the response object exists: shadow cannot affect what was
815                    // served, only what we later know about what we would have served.
816                    if let Some((b, auth, feats, ten, sess)) = shadow_ctx {
817                        spawn_shadow(
818                            &state,
819                            &route,
820                            route_ix,
821                            b,
822                            auth,
823                            feats,
824                            ten,
825                            sess,
826                            "anthropic",
827                        );
828                    }
829                    return resp;
830                }
831            }
832            if enforce_can_handle(
833                &features,
834                &body,
835                routing.escalation.enforce_structured,
836                &route.ladder,
837                &state.providers,
838                Dialect::Anthropic,
839            ) {
840                return handle_enforce(
841                    &state,
842                    &headers,
843                    &body,
844                    features,
845                    &route,
846                    route_ix,
847                    session_header,
848                    tenant,
849                    routing_mode,
850                )
851                .await;
852            }
853            // Structured request that can't be routed faithfully (flag off, or a ladder rung's
854            // dialect doesn't carry structured content verbatim yet): transparent observe
855            // passthrough — correct and un-gated beats routed and corrupted.
856            tracing::info!(
857                "enforce route matched but structured request can't be routed faithfully (flag/ladder); serving via observe passthrough"
858            );
859        }
860    }
861    observe_passthrough(state, headers, body, session_header, tenant).await
862}
863
864/// Whether the enforce path can faithfully handle this request.
865///
866/// **Verbatim-carry path** (ADR 0005 P4): when all ladder rungs carry the inbound dialect
867/// verbatim ([`crate::provider::Provider::carries_structured_verbatim`]), the original request
868/// body is forwarded byte-for-byte with only the model swapped, so every caller field survives.
869///
870/// **Translation path** (OpenAI-inbound → Anthropic ladder): for `Dialect::Openai` inbound
871/// requests hitting an all-Anthropic ladder, we translate the body to Anthropic shape — covers
872/// text, tools, tool_calls, and tool_result messages. `image_url` with http(s) URLs is not
873/// translatable (we can't relay them to Anthropic's vision API without fetching) → fallback.
874///
875/// `enforce_structured == false` restores the pre-ADR-0005 behavior: structured requests always
876/// fall back to transparent observe passthrough.
877fn enforce_can_handle(
878    features: &Features,
879    body: &[u8],
880    enforce_structured: bool,
881    ladder: &[String],
882    providers: &crate::provider::ProviderRegistry,
883    inbound: Dialect,
884) -> bool {
885    let structured = features.tool_count > 0
886        || features.has_images
887        || match inbound {
888            Dialect::Anthropic => messages_have_tool_blocks(body),
889            Dialect::Openai => openai_messages_have_tool_calls(body),
890            Dialect::Gemini => false,
891        };
892    if !structured {
893        return true;
894    }
895    if !enforce_structured {
896        return false;
897    }
898    // Path 1: verbatim carry — every rung speaks the inbound dialect natively.
899    let all_verbatim = ladder.iter().all(|rung| {
900        let provider_id = rung.split('/').next().unwrap_or_default();
901        providers
902            .get(provider_id)
903            .is_some_and(|p| p.carries_structured_verbatim(inbound))
904    });
905    if all_verbatim {
906        return true;
907    }
908    // Path 2: translation — OpenAI inbound → all-Anthropic ladder, when content is translatable.
909    // text/tools/tool_calls/tool_results are covered; http(s) image_url is not (can't relay to
910    // Anthropic's vision API without fetching). Conservative: any http(s) image → observe.
911    if inbound == Dialect::Openai && !openai_has_http_images(body) {
912        let all_anthropic = ladder.iter().all(|rung| {
913            let pid = rung.split('/').next().unwrap_or_default();
914            providers
915                .get(pid)
916                .is_some_and(|p| p.carries_structured_verbatim(Dialect::Anthropic))
917        });
918        if all_anthropic {
919            return true;
920        }
921    }
922    false
923}
924
925/// Whether any message carries a `tool_use` or `tool_result` content block (a multi-turn tool
926/// conversation), which the text-only enforce normalization would drop.
927fn messages_have_tool_blocks(body: &[u8]) -> bool {
928    serde_json::from_slice::<Value>(body)
929        .ok()
930        .and_then(|json| {
931            json.get("messages")
932                .and_then(Value::as_array)
933                .map(|messages| messages.iter().any(message_has_tool_block))
934        })
935        .unwrap_or(false)
936}
937
938/// Whether a single message's content contains a `tool_use` or `tool_result` block.
939fn message_has_tool_block(message: &Value) -> bool {
940    message
941        .get("content")
942        .and_then(Value::as_array)
943        .is_some_and(|blocks| {
944            blocks.iter().any(|block| {
945                matches!(
946                    block.get("type").and_then(Value::as_str),
947                    Some("tool_use" | "tool_result")
948                )
949            })
950        })
951}
952
953/// Whether the request opts into server-sent-events streaming (`"stream": true`).
954fn is_stream_request(body: &[u8]) -> bool {
955    serde_json::from_slice::<Value>(body)
956        .ok()
957        .and_then(|json| json.get("stream").and_then(Value::as_bool))
958        .unwrap_or(false)
959}
960
961/// Read a header as an owned `String`, if present and valid UTF-8.
962fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
963    headers
964        .get(name)
965        .and_then(|v| v.to_str().ok())
966        .map(str::to_owned)
967}
968
969/// Build the routing/telemetry feature vector from request headers + body (best-effort;
970/// malformed fields fall back to safe defaults — this must never fail a request).
971fn extract_features(headers: &HeaderMap, body: &[u8]) -> Features {
972    let (_model, tool_count, has_images) = request_features(body);
973    let mut f = Features::new(TaskKind::Other);
974    f.agent = header_str(headers, AGENT_HEADER);
975    f.subagent = header_str(headers, SUBAGENT_HEADER);
976    f.tool_count = tool_count;
977    f.has_images = has_images;
978    // Pre-call we don't know the token count, so bucket by request byte size — a coarse,
979    // monotonic proxy that never exposes the exact prompt (matches the privacy contract).
980    f.prompt_token_bucket = token_bucket(body.len() as u64);
981    f.hour_bucket = hour_bucket(jiff::Timestamp::now());
982    f
983}
984
985/// Enforce mode (SPEC §7.1): run the escalation engine and serve the first output that clears
986/// the route's gates, escalating on failure with cross-provider failover.
987#[allow(clippy::too_many_arguments)]
988async fn handle_enforce(
989    state: &AppState,
990    headers: &HeaderMap,
991    body: &Bytes,
992    features: Features,
993    route: &Route,
994    // Index of the matched route, stamped on the trace so a downstream outcome arriving
995    // later can be attributed back to the route that produced it (ADR 0009 D3).
996    route_ix: usize,
997    session_header: Option<String>,
998    tenant: String,
999    routing_mode: RoutingMode,
1000) -> Response {
1001    // A streaming client gets its SSE connection opened IMMEDIATELY: the routing pipeline
1002    // (model call + gates + possible escalation) runs in a spawned task while the response body
1003    // emits standards-compliant SSE comment keepalives (`: firstpass routing`) every few seconds,
1004    // so no client or proxy idle-timeout fires during a long escalation. When the pipeline
1005    // resolves, the gated result streams out as the usual Anthropic event sequence (ADR 0005 P3);
1006    // a pipeline error becomes an SSE `error` event (status is already 200 by then — the SSE
1007    // error frame is the in-band channel the protocol defines for exactly this).
1008    if is_stream_request(body) {
1009        let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
1010        let (state_c, headers_c, body_c, route_c) =
1011            (state.clone(), headers.clone(), body.clone(), route.clone());
1012        tokio::spawn(async move {
1013            let out = enforce_pipeline(
1014                &state_c,
1015                &headers_c,
1016                &body_c,
1017                features,
1018                &route_c,
1019                route_ix,
1020                session_header,
1021                tenant,
1022                routing_mode,
1023            )
1024            .await;
1025            let _ = tx.send(out);
1026        });
1027        return sse_keepalive_response(rx, anthropic_sse_from_message);
1028    }
1029    match enforce_pipeline(
1030        state,
1031        headers,
1032        body,
1033        features,
1034        route,
1035        route_ix,
1036        session_header,
1037        tenant,
1038        routing_mode,
1039    )
1040    .await
1041    {
1042        Ok(message) => (axum::http::StatusCode::OK, Json(message)).into_response(),
1043        Err(e) => e.into_response(),
1044    }
1045}
1046
1047/// Inner pipeline: resolve gates → route the ladder → bookkeeping (bandit, trace) → the served
1048/// [`ModelResponse`]. Shared by both Anthropic and OpenAI enforce paths; callers handle dialect-
1049/// specific parsing before and response rendering after.
1050#[allow(clippy::too_many_arguments)] // 10 params: all are genuinely distinct, not groupable
1051async fn enforce_pipeline_inner(
1052    state: &AppState,
1053    body: &Bytes,
1054    base_request: ModelRequest,
1055    auth: Auth,
1056    features: Features,
1057    route: &Route,
1058    session_header: Option<String>,
1059    tenant: String,
1060    api: &str,
1061    routing_mode: RoutingMode,
1062    route_ix: usize,
1063) -> Result<ModelResponse, ProxyError> {
1064    let gate_defs = state
1065        .config
1066        .routing
1067        .as_ref()
1068        .map_or(&[][..], |cfg| &cfg.gate_defs);
1069    let gates = resolve_gates(
1070        &route.gates,
1071        gate_defs,
1072        &state.providers,
1073        &auth,
1074        &state.config.prices,
1075    );
1076    let session_id = session_header.unwrap_or_else(|| Uuid::now_v7().to_string());
1077    let (budget, max_rungs, speculation, serve_threshold, elastic) =
1078        match state.config.routing.as_ref() {
1079            Some(cfg) => (
1080                cfg.budget.per_request_usd,
1081                cfg.escalation.max_rungs_per_request,
1082                cfg.escalation.speculation,
1083                cfg.escalation.serve_threshold,
1084                cfg.escalation.elastic.as_ref(),
1085            ),
1086            None => (None, 3, 0, None, None),
1087        };
1088    // Online adaptive conformal: serve against the LIVE-tracked threshold (updated by /v1/feedback).
1089    // Falls back to the fixed config threshold when adaptive is off or its lock is poisoned.
1090    let serve_threshold = state
1091        .adaptive
1092        .as_ref()
1093        .and_then(|a| a.lock().ok().map(|g| g.threshold()))
1094        .or(serve_threshold);
1095
1096    // Apply routing-mode preset overrides on top of config values.
1097    // Balanced preset has all None/false — byte-identical to existing behaviour.
1098    let preset = routing_mode.preset();
1099    let max_rungs = if let Some(delta) = preset.max_rungs_delta {
1100        (max_rungs as i32 + delta).max(1) as u32
1101    } else {
1102        max_rungs
1103    };
1104    let speculation = preset.speculation.unwrap_or(speculation);
1105
1106    // Predict-to-start (bandit): choose where the ladder starts for this context.
1107    // The gate still verifies the chosen rung's output before serving — prediction errors cost
1108    // money/latency but can never cause a wrong answer to be served.
1109    let bandit_ctx = crate::bandit::ContextBucket::from_features(&features);
1110
1111    // Step 1: greedy start — bandit prediction or rung 0 (cold-start / no bandit).
1112    // Thompson sampling returns its own Monte-Carlo selection propensity (the policy is
1113    // stochastic by nature); UCB1 returns None and relies on the epsilon overlay as before.
1114    let (greedy_rung, base_policy_id, ts_propensity) = {
1115        let (chosen, ts_p) = state
1116            .bandit
1117            .as_ref()
1118            .and_then(|b| b.lock().ok())
1119            .map(|mut b| {
1120                b.choose_start_with_propensity(&bandit_ctx, &route.ladder, &state.config.prices)
1121            })
1122            .unwrap_or((0, None));
1123        let policy = if ts_p.is_some() {
1124            "bandit@v2-ts".to_owned()
1125        } else if chosen > 0 {
1126            "bandit@v1".to_owned()
1127        } else {
1128            "static-ladder@v0".to_owned()
1129        };
1130        (chosen, policy, ts_p)
1131    };
1132
1133    // Step 2: epsilon-greedy overlay — randomise a fraction of start-rung choices so the
1134    // logging policy is stochastic and IPS/SNIPS off-policy estimates are valid
1135    // (Horvitz-Thompson 1952). Propensity p = (1−ε)·𝟙[chosen==greedy] + ε/K is recorded on
1136    // every trace; the bandit still observes all gate verdicts (learning is uninterrupted).
1137    let exploration_epsilon = state
1138        .config
1139        .routing
1140        .as_ref()
1141        .and_then(|cfg| cfg.escalation.exploration.as_ref())
1142        .map(|e| e.epsilon);
1143
1144    let (start_rung, policy_id, explore_flag, propensity) = if ts_propensity.is_some() {
1145        // Thompson IS the stochastic logging policy — its MC propensity is logged directly and
1146        // the epsilon overlay is redundant (warn once if both are configured).
1147        if exploration_epsilon.is_some() {
1148            tracing::warn!(
1149                "bandit.algorithm = thompson already logs propensities; \
1150                 [escalation.exploration] epsilon is ignored"
1151            );
1152        }
1153        (greedy_rung, base_policy_id, false, ts_propensity)
1154    } else if let Some(epsilon) = exploration_epsilon {
1155        let k = route.ladder.len().max(1);
1156        // Derive a per-request uniform draw from a fresh UUIDv7's bits — no new deps.
1157        let u = u01(Uuid::now_v7().as_u128());
1158        let (chosen, eps_branch) = if u < epsilon {
1159            // Epsilon branch: uniform over 0..k
1160            let idx = ((u / epsilon) * k as f64) as u32;
1161            (idx.min(k as u32 - 1), true)
1162        } else {
1163            (greedy_rung, false)
1164        };
1165        let p = epsilon_propensity(chosen, greedy_rung, epsilon, k);
1166        (chosen, format!("{base_policy_id}+eps"), eps_branch, Some(p))
1167    } else {
1168        (greedy_rung, base_policy_id, false, None)
1169    };
1170
1171    // Apply start_at_top mode override: Max mode skips bandit/epsilon and jumps to top rung.
1172    // ponytail: if ladder is empty start_rung stays 0 (saturating_sub handles it).
1173    let start_rung = if preset.start_at_top {
1174        route.ladder.len().saturating_sub(1) as u32
1175    } else {
1176        start_rung
1177    };
1178
1179    // Speculative-deferral band: prefetch only when the bandit's gate-pass estimate for the
1180    // chosen start rung is in the configured marginal zone — where the next rung is *probably
1181    // but not certainly* needed, the only place parallel spend reliably buys latency
1182    // (speculative cascades). Confident-pass or confident-fail contexts run serial and keep
1183    // the speculative tokens. No band / no bandit / cold context ⇒ configured behavior.
1184    let speculation = match state
1185        .config
1186        .routing
1187        .as_ref()
1188        .and_then(|cfg| cfg.escalation.speculation_band)
1189    {
1190        Some([lo, hi]) if speculation > 0 => {
1191            let estimate = state
1192                .bandit
1193                .as_ref()
1194                .and_then(|b| b.lock().ok())
1195                .and_then(|b| b.pass_estimate(&bandit_ctx, start_rung));
1196            match estimate {
1197                Some(p) if p < lo || p > hi => {
1198                    metrics::counter!("firstpass_speculation_skipped_total").increment(1);
1199                    0
1200                }
1201                _ => speculation,
1202            }
1203        }
1204        _ => speculation,
1205    };
1206
1207    // Emit metric whenever the bandit is configured (includes cold-start rung-0 choices).
1208    if state.bandit.is_some() {
1209        metrics::counter!(
1210            "firstpass_bandit_start_rung",
1211            "rung" => start_rung.to_string()
1212        )
1213        .increment(1);
1214    }
1215
1216    let ctx = EnforceCtx {
1217        ladder: &route.ladder,
1218        gates: &gates,
1219        health: &state.gate_health,
1220        base_request: &base_request,
1221        providers: &state.providers,
1222        auth: &auth,
1223        prices: &state.config.prices,
1224        budget_per_request_usd: budget,
1225        max_rungs,
1226        speculation,
1227        serve_threshold,
1228        elastic,
1229        features,
1230        start_rung,
1231        // The tenant stamped on the enforce trace is the resolved identity from the auth layer
1232        // (authenticated key, or the static default when auth is off) — never the request body.
1233        tenant_id: tenant,
1234        session_id,
1235        prompt_hash: prompt_hash(&state.config.prompt_salt, body),
1236        api: api.to_owned(),
1237        policy_id,
1238    };
1239
1240    let (outcome, mut trace) = route_enforce(ctx).await;
1241
1242    // Stamp which route produced this decision, so a downstream outcome arriving minutes later
1243    // can be attributed back to it (ADR 0009 D3). Without it the guardrail pools every outcome
1244    // onto one route, and in a multi-route config a failing route hides behind healthy siblings.
1245    trace.route_ix = Some(u32::try_from(route_ix).unwrap_or(0));
1246    // Patch explore/propensity onto the trace now that we know whether the epsilon branch fired.
1247    // route_enforce leaves these at (false, None); we own the trace before it's hashed+stored.
1248    trace.policy.explore = explore_flag;
1249    trace.policy.propensity = propensity;
1250    // Stamp the resolved mode profile when it's not Balanced (the default).
1251    // None → absent from JSON → byte-identical for existing traces.
1252    if routing_mode != RoutingMode::Balanced {
1253        trace.policy.mode_profile = Some(routing_mode.as_str().to_owned());
1254    }
1255
1256    // Online bandit learning: feed back every gate verdict from this request so the bandit
1257    // refines its start-rung estimates. Cheap in-memory update; done before offer_trace so the
1258    // trace borrow is still live (we read attempts, then pass trace to offer_trace by value).
1259    if let Some(bandit) = state.bandit.as_ref()
1260        && let Ok(mut b) = bandit.lock()
1261    {
1262        for attempt in &trace.attempts {
1263            b.observe(&bandit_ctx, attempt.rung, attempt.verdict);
1264        }
1265    }
1266
1267    // ── Per-query gate-pass predictor (ADR 0008 Phase 2) ────────────────────────────────────
1268    // Record the predicted P(gate-pass) for the start rung on the receipt in SHADOW (never acted
1269    // on), then learn online from this request's attempts. Default-off (predictor = None):
1270    // trace.predicted_pass stays None → byte-identical to today. The predictor never touches
1271    // serving; it only writes a receipt field, updates in-memory weights, and emits a metric.
1272    if let Some(predictor) = state.predictor.as_ref()
1273        && let Ok(mut p) = predictor.lock()
1274    {
1275        // Read the routed features from the trace (the owned `features` was moved into the ctx).
1276        let predicted = p.predict(&trace.request.features, start_rung);
1277        for attempt in &trace.attempts {
1278            match attempt.verdict {
1279                Verdict::Pass => p.update(&trace.request.features, attempt.rung, true),
1280                Verdict::Fail => p.update(&trace.request.features, attempt.rung, false),
1281                Verdict::Abstain => {} // no clear label — don't train on it
1282            }
1283        }
1284        trace.predicted_pass = Some(predicted);
1285        metrics::histogram!("firstpass_predictor_pass_prob").record(predicted);
1286    }
1287
1288    // ── Shadow probe (ADR 0008 Phase 1) ─────────────────────────────────────────────────────
1289    // Measure the k-sample gate-pass-count signal on a sampled fraction of requests.
1290    // Default-off (probe = None): zero extra provider calls, trace.probe stays None — byte-identical.
1291    // When on: k model calls at the start_rung model run concurrently; gate evals are read-only.
1292    // INVARIANT: gate_health.record() is NEVER called from the probe path — shadow must not
1293    //            trip error budgets or alter any mutable registry state.
1294    if let Some(probe_cfg) = state
1295        .config
1296        .routing
1297        .as_ref()
1298        .and_then(|c| c.escalation.probe)
1299        && u01(Uuid::now_v7().as_u128()) < probe_cfg.sample_rate
1300    {
1301        // Clamp start_rung to the ladder bounds (same as run_serial/run_speculative).
1302        let probe_rung = (start_rung as usize).min(route.ladder.len().saturating_sub(1));
1303        if let Some(probe_model_str) = route.ladder.get(probe_rung).cloned() {
1304            let probe_provider = ModelRef::parse(&probe_model_str)
1305                .ok()
1306                .and_then(|m| state.providers.get(&m.provider));
1307
1308            if let Some(probe_provider) = probe_provider {
1309                // Spawn k model calls concurrently.
1310                // ponytail: gate evals run sequentially after all calls complete — simple and correct.
1311                let mut join_set = tokio::task::JoinSet::new();
1312                for _ in 0..probe_cfg.k {
1313                    let mut probe_req = base_request.clone();
1314                    probe_req.model = probe_model_str.clone();
1315                    let probe_auth = auth.clone();
1316                    let prov = probe_provider.clone();
1317                    join_set.spawn(async move { prov.complete(&probe_req, &probe_auth).await });
1318                }
1319
1320                // Build the fail-closed id set — mirrors router::run_serial exactly.
1321                // ponytail: owned strings avoid lifetime/async issues; update if serve rule changes.
1322                let fail_closed_owned: std::collections::HashSet<String> = gates
1323                    .iter()
1324                    .filter(|g| g.abstain_fails_closed())
1325                    .map(|g| g.id().to_owned())
1326                    .collect();
1327
1328                let mut gate_pass_count = 0u32;
1329                let mut probe_cost_usd = 0.0f64;
1330
1331                while let Some(task_result) = join_set.join_next().await {
1332                    let Ok(Ok(probe_resp)) = task_result else {
1333                        continue; // provider error on a sample = not-passed; count honestly
1334                    };
1335                    probe_cost_usd += state
1336                        .config
1337                        .prices
1338                        .cost_usd(
1339                            &probe_model_str,
1340                            probe_resp.in_tokens,
1341                            probe_resp.out_tokens,
1342                        )
1343                        .unwrap_or(0.0);
1344
1345                    let mut probe_gate_req = base_request.clone();
1346                    probe_gate_req.model = probe_model_str.clone();
1347
1348                    // Run gates — READ-ONLY: deliberately no gate_health.record() calls so the
1349                    // shadow probe never trips error budgets or mutates registry state.
1350                    let mut probe_gate_results = Vec::with_capacity(gates.len());
1351                    for g in &gates {
1352                        // Respect disabled status (read; no write) so a sick gate isn't re-probed.
1353                        if !state.gate_health.enabled(&trace.tenant_id, g.id()) {
1354                            continue;
1355                        }
1356                        let r = g.evaluate(&probe_gate_req, &probe_resp).await;
1357                        // NOTE: gate_health.record() intentionally NOT called here.
1358                        probe_gate_results.push(r);
1359                    }
1360
1361                    let fail_closed_refs: std::collections::HashSet<&str> =
1362                        fail_closed_owned.iter().map(|s| s.as_str()).collect();
1363                    let verdict = aggregate_with_policy(&probe_gate_results, &fail_closed_refs);
1364                    // Mirror should_serve from router.rs exactly (private there; replicated here).
1365                    // ponytail: if the serve rule in router.rs changes, update this too.
1366                    let passes = match serve_threshold {
1367                        None => verdict == Verdict::Pass,
1368                        Some(t) => crate::calibrate::gate_score(&probe_gate_results, verdict) >= t,
1369                    };
1370                    if passes {
1371                        gate_pass_count += 1;
1372                    }
1373                }
1374
1375                let regime = ProbeRegime::classify(gate_pass_count, probe_cfg.k);
1376                let regime_label = match regime {
1377                    ProbeRegime::ConfidentPass => "confident_pass",
1378                    ProbeRegime::ConfidentFail => "confident_fail",
1379                    ProbeRegime::Ambiguous => "ambiguous",
1380                };
1381                metrics::counter!(
1382                    "firstpass_probe_regime_total",
1383                    "regime" => regime_label
1384                )
1385                .increment(1);
1386                metrics::gauge!("firstpass_probe_cost_usd_total").increment(probe_cost_usd);
1387                trace.probe = Some(ProbeSignal {
1388                    k: probe_cfg.k,
1389                    gate_pass_count,
1390                    regime,
1391                    probe_cost_usd,
1392                });
1393            }
1394        }
1395    }
1396    // ── end shadow probe ─────────────────────────────────────────────────────────────────────
1397
1398    // The trace is already built; enqueue it off-path (non-blocking `try_send`, so no spawn needed).
1399    offer_trace(&state.traces, state.spill.as_ref(), trace);
1400
1401    match outcome {
1402        EngineOutcome::Served(resp) => Ok(resp),
1403        EngineOutcome::Failed(msg) => Err(ProxyError::Engine(msg)),
1404    }
1405}
1406
1407/// Anthropic enforce pipeline: parse → inner pipeline → Anthropic-shaped response JSON.
1408/// Shared verbatim by the buffered (non-streaming) and keepalive-streaming paths.
1409#[allow(clippy::too_many_arguments)]
1410async fn enforce_pipeline(
1411    state: &AppState,
1412    headers: &HeaderMap,
1413    body: &Bytes,
1414    features: Features,
1415    route: &Route,
1416    // Route that produced this decision; threaded through for outcome attribution.
1417    route_ix: usize,
1418    session_header: Option<String>,
1419    tenant: String,
1420    routing_mode: RoutingMode,
1421) -> Result<Value, ProxyError> {
1422    let Some(base_request) = parse_model_request(body) else {
1423        return Err(ProxyError::BadRequest(
1424            "request body is not a valid Anthropic Messages request".to_owned(),
1425        ));
1426    };
1427    let auth = Auth::from_headers(headers);
1428    let resp = enforce_pipeline_inner(
1429        state,
1430        body,
1431        base_request,
1432        auth,
1433        features,
1434        route,
1435        session_header,
1436        tenant,
1437        "anthropic.messages",
1438        routing_mode,
1439        route_ix,
1440    )
1441    .await?;
1442    Ok(anthropic_response_json(&resp))
1443}
1444
1445/// OpenAI enforce pipeline: parse (with raw-carry for all-OpenAI ladders, else translation) →
1446/// inner pipeline → OpenAI `chat.completion` JSON.
1447#[allow(clippy::too_many_arguments)]
1448async fn enforce_pipeline_openai(
1449    state: &AppState,
1450    headers: &HeaderMap,
1451    body: &Bytes,
1452    features: Features,
1453    route: &Route,
1454    // Route that produced this decision; threaded through for outcome attribution.
1455    route_ix: usize,
1456    session_header: Option<String>,
1457    tenant: String,
1458    routing_mode: RoutingMode,
1459) -> Result<Value, ProxyError> {
1460    // Decide between verbatim raw-carry (all-OpenAI ladder) and translation (Anthropic ladder).
1461    let providers = &state.providers;
1462    let all_openai = route.ladder.iter().all(|rung| {
1463        let pid = rung.split('/').next().unwrap_or_default();
1464        providers
1465            .get(pid)
1466            .is_some_and(|p| p.carries_structured_verbatim(Dialect::Openai))
1467    });
1468    let Some(base_request) = parse_openai_request(body, all_openai) else {
1469        return Err(ProxyError::BadRequest(
1470            "request body is not a valid OpenAI Chat Completions request".to_owned(),
1471        ));
1472    };
1473    let auth = Auth::from_headers(headers);
1474    let resp = enforce_pipeline_inner(
1475        state,
1476        body,
1477        base_request,
1478        auth,
1479        features,
1480        route,
1481        session_header,
1482        tenant,
1483        "openai.chat_completions",
1484        routing_mode,
1485        route_ix,
1486    )
1487    .await?;
1488    Ok(openai_response_json(&resp))
1489}
1490
1491/// Interval between SSE comment keepalives while the enforce pipeline is still routing.
1492const SSE_KEEPALIVE_EVERY: Duration = Duration::from_secs(5);
1493
1494/// A 200 `text/event-stream` response whose body emits comment keepalives until `rx` resolves,
1495/// then the gated result formatted by `format_message` (or an SSE `error` event on failure).
1496/// SSE comment lines (leading `:`) are defined by the EventSource spec to be ignored by every
1497/// conforming parser — they keep the connection alive without confusing any client.
1498///
1499/// `format_message` converts the gated result `Value` to the dialect-appropriate SSE frame
1500/// string: pass [`anthropic_sse_from_message`] for Anthropic clients,
1501/// [`openai_sse_from_message`] for OpenAI clients.
1502fn sse_keepalive_response(
1503    rx: tokio::sync::oneshot::Receiver<Result<Value, ProxyError>>,
1504    format_message: fn(&Value) -> String,
1505) -> Response {
1506    let mut ticks = tokio::time::interval(SSE_KEEPALIVE_EVERY);
1507    ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1508    ticks.reset(); // skip the immediate first tick — the first keepalive fires after one period
1509    let stream = KeepaliveStream {
1510        rx: Some(rx),
1511        ticks,
1512        format_message,
1513    };
1514    (
1515        axum::http::StatusCode::OK,
1516        [(
1517            axum::http::header::CONTENT_TYPE,
1518            "text/event-stream; charset=utf-8",
1519        )],
1520        axum::body::Body::from_stream(stream),
1521    )
1522        .into_response()
1523}
1524
1525/// Hand-rolled [`futures_core::Stream`]-shaped body (no new dependency): comment keepalives
1526/// while the pipeline runs, then the final SSE frame, then end-of-stream.
1527struct KeepaliveStream {
1528    /// `Some` until the pipeline resolves and the final frame has been emitted.
1529    rx: Option<tokio::sync::oneshot::Receiver<Result<Value, ProxyError>>>,
1530    ticks: tokio::time::Interval,
1531    /// Converts the served result `Value` to the caller's dialect SSE frames.
1532    format_message: fn(&Value) -> String,
1533}
1534
1535impl futures_core::Stream for KeepaliveStream {
1536    type Item = Result<Bytes, std::convert::Infallible>;
1537
1538    fn poll_next(
1539        mut self: std::pin::Pin<&mut Self>,
1540        cx: &mut std::task::Context<'_>,
1541    ) -> std::task::Poll<Option<Self::Item>> {
1542        use std::task::Poll;
1543        let Some(rx) = self.rx.as_mut() else {
1544            return Poll::Ready(None); // final frame already emitted
1545        };
1546        if let Poll::Ready(out) = std::pin::Pin::new(rx).poll(cx) {
1547            let fmt = self.format_message;
1548            let frame = match out {
1549                Ok(Ok(message)) => fmt(&message),
1550                Ok(Err(e)) => sse_error_event(&e),
1551                Err(_) => sse_error_event(&ProxyError::Internal(
1552                    "enforce pipeline task dropped".to_owned(),
1553                )),
1554            };
1555            self.rx = None;
1556            return Poll::Ready(Some(Ok(Bytes::from(frame))));
1557        }
1558        if self.ticks.poll_tick(cx).is_ready() {
1559            return Poll::Ready(Some(Ok(Bytes::from_static(b": firstpass routing\n\n"))));
1560        }
1561        Poll::Pending
1562    }
1563}
1564
1565/// Render a pipeline error as the Anthropic SSE `error` event (client-safe message only —
1566/// internal detail is logged by the error type, never sent).
1567fn sse_error_event(e: &ProxyError) -> String {
1568    let mut out = String::new();
1569    sse_event(
1570        &mut out,
1571        "error",
1572        &serde_json::json!({
1573            "type": "error",
1574            "error": { "type": "api_error", "message": e.client_message() }
1575        }),
1576    );
1577    out
1578}
1579
1580/// Parse an Anthropic Messages request body into the normalized [`ModelRequest`]. Returns
1581/// `None` if the body isn't valid JSON or lacks a `messages` array.
1582///
1583// Message content is preserved **verbatim** (string or array of blocks) — a plain-string content
1584// serializes byte-identical on the wire, and tool_use/tool_result/image blocks survive the round
1585// trip (ADR 0005, invariant I2). Gates operate on `ChatMessage::text_view()`, not the raw content,
1586// so gate behavior is unchanged. Which requests actually enter enforce is still governed by
1587// `enforce_can_handle`; this function only guarantees no fidelity is lost once they do.
1588fn parse_model_request(body: &[u8]) -> Option<ModelRequest> {
1589    let json: Value = serde_json::from_slice(body).ok()?;
1590    let raw = json.clone();
1591    let messages_json = json.get("messages")?.as_array()?;
1592    let messages = messages_json
1593        .iter()
1594        .map(|m| ChatMessage {
1595            role: m
1596                .get("role")
1597                .and_then(Value::as_str)
1598                .unwrap_or("user")
1599                .to_owned(),
1600            content: m
1601                .get("content")
1602                .cloned()
1603                .unwrap_or_else(|| Value::String(String::new())),
1604        })
1605        .collect();
1606    let system = json
1607        .get("system")
1608        .and_then(Value::as_str)
1609        .map(str::to_owned);
1610    let max_tokens = json
1611        .get("max_tokens")
1612        .and_then(Value::as_u64)
1613        .and_then(|n| u32::try_from(n).ok())
1614        .unwrap_or(1024);
1615    let tools = json.get("tools").cloned().unwrap_or(Value::Null);
1616    Some(ModelRequest {
1617        model: json
1618            .get("model")
1619            .and_then(Value::as_str)
1620            .unwrap_or_default()
1621            .to_owned(),
1622        system,
1623        messages,
1624        max_tokens,
1625        tools,
1626        raw,
1627    })
1628}
1629
1630/// Render a served [`ModelResponse`] back into an Anthropic Messages response envelope, so the
1631/// caller sees the same wire shape regardless of which provider actually answered.
1632///
1633/// The `content` blocks come **verbatim** from the upstream response (`resp.raw`) when it is an
1634/// Anthropic message — so `tool_use` / `thinking` / multiple text blocks reach the caller intact
1635/// (ADR 0005 I2). Only when `raw` has no Anthropic `content` array (a synthetic response, or the
1636/// OpenAI adapter, which has `choices` instead) do we fall back to a single reconstructed text
1637/// block. The envelope (`id`, `model`, `usage`) is always normalized so the served model id is the
1638/// prefixed ladder id, not the bare wire id.
1639fn anthropic_response_json(resp: &ModelResponse) -> Value {
1640    let content = resp
1641        .raw
1642        .get("content")
1643        .filter(|c| c.is_array())
1644        .cloned()
1645        .unwrap_or_else(|| serde_json::json!([{ "type": "text", "text": resp.text }]));
1646    serde_json::json!({
1647        "id": format!("msg_{}", Uuid::now_v7()),
1648        "type": "message",
1649        "role": "assistant",
1650        "model": resp.model,
1651        "content": content,
1652        "usage": { "input_tokens": resp.in_tokens, "output_tokens": resp.out_tokens },
1653    })
1654}
1655
1656/// Append one `event: <type>\ndata: <json>\n\n` SSE frame.
1657fn sse_event(out: &mut String, event: &str, data: &Value) {
1658    out.push_str("event: ");
1659    out.push_str(event);
1660    out.push_str("\ndata: ");
1661    out.push_str(&data.to_string());
1662    out.push_str("\n\n");
1663}
1664
1665/// Re-emit a served Anthropic message envelope (from [`anthropic_response_json`]) as an SSE stream
1666/// body, so a `stream: true` client is served even though enforce buffered the response to gate it
1667/// (ADR 0005 P3). The gate needs the full candidate, so this is not token-by-token streaming from
1668/// the model — each content block is emitted as a single delta. `tool_use` blocks are preserved:
1669/// their `input` is streamed as one `input_json_delta` (invariant I2), so the caller reconstructs
1670/// the exact tool call.
1671fn anthropic_sse_from_message(message: &Value) -> String {
1672    let mut out = String::new();
1673
1674    // message_start carries the envelope with content emptied — the blocks stream next.
1675    let mut start_msg = message.clone();
1676    start_msg["content"] = Value::Array(Vec::new());
1677    sse_event(
1678        &mut out,
1679        "message_start",
1680        &serde_json::json!({ "type": "message_start", "message": start_msg }),
1681    );
1682
1683    let empty = Vec::new();
1684    let blocks = message
1685        .get("content")
1686        .and_then(Value::as_array)
1687        .unwrap_or(&empty);
1688    for (i, block) in blocks.iter().enumerate() {
1689        match block.get("type").and_then(Value::as_str) {
1690            Some("tool_use") => {
1691                // Start with an empty input object, then stream the real input as one JSON delta.
1692                let mut shell = block.clone();
1693                shell["input"] = serde_json::json!({});
1694                sse_event(
1695                    &mut out,
1696                    "content_block_start",
1697                    &serde_json::json!({ "type": "content_block_start", "index": i, "content_block": shell }),
1698                );
1699                let input_json = block
1700                    .get("input")
1701                    .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
1702                sse_event(
1703                    &mut out,
1704                    "content_block_delta",
1705                    &serde_json::json!({ "type": "content_block_delta", "index": i,
1706                        "delta": { "type": "input_json_delta", "partial_json": input_json } }),
1707                );
1708            }
1709            _ => {
1710                // text (and any other text-bearing block): start empty, stream the text as one delta.
1711                let text = block.get("text").and_then(Value::as_str).unwrap_or("");
1712                sse_event(
1713                    &mut out,
1714                    "content_block_start",
1715                    &serde_json::json!({ "type": "content_block_start", "index": i,
1716                        "content_block": { "type": "text", "text": "" } }),
1717                );
1718                sse_event(
1719                    &mut out,
1720                    "content_block_delta",
1721                    &serde_json::json!({ "type": "content_block_delta", "index": i,
1722                        "delta": { "type": "text_delta", "text": text } }),
1723                );
1724            }
1725        }
1726        sse_event(
1727            &mut out,
1728            "content_block_stop",
1729            &serde_json::json!({ "type": "content_block_stop", "index": i }),
1730        );
1731    }
1732
1733    let out_tokens = message
1734        .pointer("/usage/output_tokens")
1735        .cloned()
1736        .unwrap_or_else(|| Value::from(0));
1737    sse_event(
1738        &mut out,
1739        "message_delta",
1740        &serde_json::json!({ "type": "message_delta", "delta": { "stop_reason": "end_turn" },
1741            "usage": { "output_tokens": out_tokens } }),
1742    );
1743    sse_event(
1744        &mut out,
1745        "message_stop",
1746        &serde_json::json!({ "type": "message_stop" }),
1747    );
1748    out
1749}
1750
1751/// Observe mode (SPEC §7.1a): forward unchanged, return unchanged, trace asynchronously.
1752async fn observe_passthrough(
1753    state: AppState,
1754    headers: HeaderMap,
1755    body: Bytes,
1756    session_header: Option<String>,
1757    tenant: String,
1758) -> Response {
1759    // Streaming requests are relayed chunk-by-chunk rather than buffered (SPEC §7.4).
1760    if is_stream_request(&body) {
1761        return observe_stream(state, headers, body, session_header, tenant).await;
1762    }
1763    let start = Instant::now();
1764    let result = forward_anthropic(
1765        &state.http,
1766        &state.config.upstream_anthropic,
1767        &headers,
1768        body.clone(),
1769    )
1770    .await;
1771    let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
1772
1773    match result {
1774        Ok((status, resp_headers, resp_body)) => {
1775            // Build + record the trace on a detached task so neither JSON parsing nor the
1776            // channel send touches the response path: observe mode adds zero latency to what
1777            // the caller sees (SPEC §7.1a). `Bytes` clones are cheap (refcounted).
1778            spawn_trace(
1779                &state,
1780                body,
1781                Some(resp_body.clone()),
1782                latency_ms,
1783                session_header,
1784                tenant,
1785            );
1786            (status, resp_headers, resp_body).into_response()
1787        }
1788        Err(err) => {
1789            spawn_trace(&state, body, None, latency_ms, session_header, tenant);
1790            err.into_response()
1791        }
1792    }
1793}
1794
1795/// Observe mode for a streaming request (`stream: true`): relay the upstream SSE response
1796/// chunk-by-chunk instead of buffering, so streaming is preserved to the caller and
1797/// time-to-first-byte stays low. `latency_ms` is time-to-response-headers (the added-latency
1798/// figure that matters), recorded off the response path.
1799///
1800// ponytail: streamed-response token usage lives in the SSE `message_start`/`message_delta` events
1801// we don't buffer, so the trace records request-side features + latency now; parsing usage from a
1802// teed SSE stream is the follow-on.
1803async fn observe_stream(
1804    state: AppState,
1805    headers: HeaderMap,
1806    body: Bytes,
1807    session_header: Option<String>,
1808    tenant: String,
1809) -> Response {
1810    let start = Instant::now();
1811    let result = forward_anthropic_streaming(
1812        &state.http,
1813        &state.config.upstream_anthropic,
1814        &headers,
1815        body.clone(),
1816    )
1817    .await;
1818    let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
1819
1820    match result {
1821        Ok((status, resp_headers, response)) => {
1822            spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
1823            let stream_body = Body::from_stream(response.bytes_stream());
1824            (status, resp_headers, stream_body).into_response()
1825        }
1826        Err(err) => {
1827            spawn_trace(&state, body, None, latency_ms, session_header, tenant);
1828            err.into_response()
1829        }
1830    }
1831}
1832
1833/// Enqueue a request-side trace for a streamed observe response, off the response path.
1834fn spawn_stream_trace(
1835    state: &AppState,
1836    req_body: Bytes,
1837    latency_ms: u64,
1838    session_header: Option<String>,
1839    tenant: String,
1840) {
1841    let config = state.config.clone();
1842    let traces = state.traces.clone();
1843    let spill = state.spill.clone();
1844    tokio::spawn(async move {
1845        let mut trace =
1846            build_stream_trace(&config, &req_body, latency_ms, session_header.as_deref());
1847        // Stamp the resolved tenant identity — never the config default nor anything request-borne.
1848        trace.tenant_id = tenant;
1849        offer_trace(&traces, spill.as_ref(), trace);
1850    });
1851}
1852
1853/// Construct the trace and enqueue it for the background writer, entirely off the response
1854/// path. Fire-and-forget: if the writer has shut down we log rather than propagate — recording
1855/// must never affect what the caller sees. `resp_body` is `Some` for a forwarded response and
1856/// `None` when the upstream call failed outright.
1857fn spawn_trace(
1858    state: &AppState,
1859    req_body: Bytes,
1860    resp_body: Option<Bytes>,
1861    latency_ms: u64,
1862    session_header: Option<String>,
1863    tenant: String,
1864) {
1865    let config = state.config.clone();
1866    let traces = state.traces.clone();
1867    let spill = state.spill.clone();
1868    tokio::spawn(async move {
1869        let mut trace = match resp_body {
1870            Some(resp) => build_trace(
1871                &config,
1872                &req_body,
1873                &resp,
1874                latency_ms,
1875                session_header.as_deref(),
1876            ),
1877            None => build_error_trace(&config, &req_body, latency_ms, session_header.as_deref()),
1878        };
1879        // Stamp the resolved tenant identity — never the config default nor anything request-borne.
1880        trace.tenant_id = tenant;
1881        offer_trace(&traces, spill.as_ref(), trace);
1882    });
1883}
1884
1885/// Session id for the trace: the caller-supplied header, or the trace's own id when absent.
1886fn session_id(session_header: Option<&str>, trace_id: Uuid) -> String {
1887    session_header
1888        .map(str::to_owned)
1889        .unwrap_or_else(|| trace_id.to_string())
1890}
1891
1892/// Salted hash of the raw request body — the only trace of the prompt that ever touches
1893/// storage (SPEC: never log or persist raw prompt text).
1894fn prompt_hash(salt: &str, body: &[u8]) -> String {
1895    let mut salted = Vec::with_capacity(salt.len() + body.len());
1896    salted.extend_from_slice(salt.as_bytes());
1897    salted.extend_from_slice(body);
1898    sha256_hex(&salted)
1899}
1900
1901/// Best-effort request-side feature extraction: model name, tool count, and whether any
1902/// message carries image content. Malformed/absent fields fall back to safe defaults rather
1903/// than failing the request — this is telemetry, not the served response.
1904fn request_features(body: &[u8]) -> (Option<String>, u32, bool) {
1905    let Ok(json) = serde_json::from_slice::<Value>(body) else {
1906        return (None, 0, false);
1907    };
1908    let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
1909    let tool_count = json
1910        .get("tools")
1911        .and_then(Value::as_array)
1912        .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
1913    let has_images = json
1914        .get("messages")
1915        .and_then(Value::as_array)
1916        .is_some_and(|messages| messages.iter().any(message_has_image));
1917    (model, tool_count, has_images)
1918}
1919
1920/// Whether a single message's content contains an image block (`{"type": "image", ...}`).
1921fn message_has_image(message: &Value) -> bool {
1922    message
1923        .get("content")
1924        .and_then(Value::as_array)
1925        .is_some_and(|blocks| {
1926            blocks
1927                .iter()
1928                .any(|block| block.get("type").and_then(Value::as_str) == Some("image"))
1929        })
1930}
1931
1932/// Response-side usage extraction: model name and token counts, defaulting to `0` when the
1933/// upstream response doesn't carry them (e.g. an error body).
1934fn response_usage(body: &[u8]) -> (Option<String>, u64, u64) {
1935    let Ok(json) = serde_json::from_slice::<Value>(body) else {
1936        return (None, 0, 0);
1937    };
1938    let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
1939    let in_tokens = json
1940        .pointer("/usage/input_tokens")
1941        .and_then(Value::as_u64)
1942        .unwrap_or(0);
1943    let out_tokens = json
1944        .pointer("/usage/output_tokens")
1945        .and_then(Value::as_u64)
1946        .unwrap_or(0);
1947    (model, in_tokens, out_tokens)
1948}
1949
1950/// Build the observe-mode trace for a request that was successfully forwarded and answered.
1951fn build_trace(
1952    config: &ProxyConfig,
1953    req_body: &Bytes,
1954    resp_body: &Bytes,
1955    latency_ms: u64,
1956    session_header: Option<&str>,
1957) -> Trace {
1958    let (req_model, tool_count, has_images) = request_features(req_body);
1959    let (resp_model, in_tokens, out_tokens) = response_usage(resp_body);
1960    let model = resp_model
1961        .or(req_model)
1962        .unwrap_or_else(|| "unknown".to_owned());
1963
1964    let cost_usd = config
1965        .prices
1966        .cost_usd(&format!("anthropic/{model}"), in_tokens, out_tokens)
1967        .unwrap_or(0.0);
1968
1969    let attempt = Attempt {
1970        rung: 0,
1971        model,
1972        provider: "anthropic".to_owned(),
1973        in_tokens,
1974        out_tokens,
1975        cost_usd,
1976        latency_ms,
1977        gates: Vec::new(),
1978        verdict: Verdict::Pass,
1979    };
1980
1981    let mut trace = base_trace(config, req_body, latency_ms, session_header);
1982    trace.request.features.prompt_token_bucket = token_bucket(in_tokens);
1983    trace.request.features.tool_count = tool_count;
1984    trace.request.features.has_images = has_images;
1985    trace.attempts.push(attempt);
1986    trace.final_ = FinalOutcome {
1987        served_rung: Some(0),
1988        served_from: ServedFrom::Attempt,
1989        total_cost_usd: cost_usd,
1990        gate_cost_usd: 0.0,
1991        total_latency_ms: latency_ms,
1992        escalations: 0,
1993        counterfactual_baseline_usd: cost_usd,
1994        savings_usd: 0.0,
1995    };
1996    trace.recompute_savings();
1997    trace
1998}
1999
2000/// Build the observe-mode trace for a **streamed** response: we relayed real bytes to the caller,
2001/// but the token usage lives in the SSE events we didn't buffer, so it's recorded as served with
2002/// unknown (zero) usage — honest about what we served without inventing token counts.
2003fn build_stream_trace(
2004    config: &ProxyConfig,
2005    req_body: &Bytes,
2006    latency_ms: u64,
2007    session_header: Option<&str>,
2008) -> Trace {
2009    let (req_model, tool_count, has_images) = request_features(req_body);
2010    let model = req_model.unwrap_or_else(|| "unknown".to_owned());
2011
2012    let attempt = Attempt {
2013        rung: 0,
2014        model,
2015        provider: "anthropic".to_owned(),
2016        in_tokens: 0,
2017        out_tokens: 0,
2018        cost_usd: 0.0,
2019        latency_ms,
2020        gates: Vec::new(),
2021        verdict: Verdict::Pass,
2022    };
2023
2024    let mut trace = base_trace(config, req_body, latency_ms, session_header);
2025    trace.request.features.tool_count = tool_count;
2026    trace.request.features.has_images = has_images;
2027    trace.attempts.push(attempt);
2028    trace.final_ = FinalOutcome {
2029        served_rung: Some(0),
2030        served_from: ServedFrom::Attempt,
2031        total_cost_usd: 0.0,
2032        gate_cost_usd: 0.0,
2033        total_latency_ms: latency_ms,
2034        escalations: 0,
2035        counterfactual_baseline_usd: 0.0,
2036        savings_usd: 0.0,
2037    };
2038    trace.recompute_savings();
2039    trace
2040}
2041
2042/// Build the observe-mode trace for a request whose upstream call failed outright (no
2043/// response to report usage from). Recorded with `served_from: Error` and no attempts —
2044/// keep the audit trail honest that nothing was served.
2045fn build_error_trace(
2046    config: &ProxyConfig,
2047    req_body: &Bytes,
2048    latency_ms: u64,
2049    session_header: Option<&str>,
2050) -> Trace {
2051    let (_, tool_count, has_images) = request_features(req_body);
2052    let mut trace = base_trace(config, req_body, latency_ms, session_header);
2053    trace.request.features.tool_count = tool_count;
2054    trace.request.features.has_images = has_images;
2055    trace.final_ = FinalOutcome {
2056        served_rung: None,
2057        served_from: ServedFrom::Error,
2058        total_cost_usd: 0.0,
2059        gate_cost_usd: 0.0,
2060        total_latency_ms: latency_ms,
2061        escalations: 0,
2062        counterfactual_baseline_usd: 0.0,
2063        savings_usd: 0.0,
2064    };
2065    trace.recompute_savings();
2066    trace
2067}
2068
2069/// The parts of a trace that don't depend on whether the call succeeded: identity, policy,
2070/// and the request-side feature vector minus token bucket (which needs response usage).
2071fn base_trace(
2072    config: &ProxyConfig,
2073    req_body: &Bytes,
2074    latency_ms: u64,
2075    session_header: Option<&str>,
2076) -> Trace {
2077    let trace_id = Uuid::now_v7();
2078    let mut features = Features::new(TaskKind::Other);
2079    features.hour_bucket = hour_bucket(jiff::Timestamp::now());
2080
2081    Trace {
2082        trace_id,
2083        prev_hash: GENESIS_HASH.to_owned(),
2084        tenant_id: config.tenant_id.clone(),
2085        session_id: session_id(session_header, trace_id),
2086        ts: jiff::Timestamp::now(),
2087        mode: Mode::Observe,
2088        policy: PolicyRef {
2089            id: "observe-passthrough@v0".to_owned(),
2090            explore: false,
2091            propensity: None,
2092            mode_profile: None,
2093        },
2094        request: RequestInfo {
2095            api: "anthropic.messages".to_owned(),
2096            prompt_hash: prompt_hash(&config.prompt_salt, req_body),
2097            features,
2098        },
2099        attempts: Vec::new(),
2100        deferred: Vec::new(),
2101        final_: FinalOutcome {
2102            served_rung: None,
2103            served_from: ServedFrom::Error,
2104            total_cost_usd: 0.0,
2105            gate_cost_usd: 0.0,
2106            total_latency_ms: latency_ms,
2107            escalations: 0,
2108            counterfactual_baseline_usd: 0.0,
2109            savings_usd: 0.0,
2110        },
2111        probe: None,
2112        rollout: None,
2113        shadow: None,
2114        route_ix: None,
2115        predicted_pass: None,
2116        elastic: None,
2117    }
2118}
2119
2120// ── OpenAI-inbound detection helpers ─────────────────────────────────────────
2121
2122/// Whether any OpenAI-format message has `tool_calls` on an assistant turn, or a `role:"tool"`
2123/// message (a multi-turn tool conversation that would need translation).
2124fn openai_messages_have_tool_calls(body: &[u8]) -> bool {
2125    serde_json::from_slice::<Value>(body)
2126        .ok()
2127        .and_then(|json| {
2128            json.get("messages").and_then(Value::as_array).map(|msgs| {
2129                msgs.iter().any(|m| {
2130                    m.get("tool_calls").is_some()
2131                        || m.get("role").and_then(Value::as_str) == Some("tool")
2132                })
2133            })
2134        })
2135        .unwrap_or(false)
2136}
2137
2138/// Whether any OpenAI-format message has an `image_url` content part whose URL is an
2139/// http(s) URL (not a data: URI). These cannot be forwarded to Anthropic's vision API without
2140/// fetching, so they are treated as non-translatable → observe fallback.
2141fn openai_has_http_images(body: &[u8]) -> bool {
2142    serde_json::from_slice::<Value>(body)
2143        .ok()
2144        .and_then(|json| {
2145            json.get("messages").and_then(Value::as_array).map(|msgs| {
2146                msgs.iter().any(|m| {
2147                    m.get("content")
2148                        .and_then(Value::as_array)
2149                        .is_some_and(|parts| {
2150                            parts.iter().any(|p| {
2151                                p.get("type").and_then(Value::as_str) == Some("image_url")
2152                                    && p.pointer("/image_url/url")
2153                                        .and_then(Value::as_str)
2154                                        .is_some_and(|u| {
2155                                            u.starts_with("http://") || u.starts_with("https://")
2156                                        })
2157                            })
2158                        })
2159                })
2160            })
2161        })
2162        .unwrap_or(false)
2163}
2164
2165/// Whether any OpenAI-format message has an `image_url` content part (data: or http(s)).
2166/// Used by `extract_openai_features` to set `has_images`.
2167fn openai_messages_have_images(body: &[u8]) -> bool {
2168    serde_json::from_slice::<Value>(body)
2169        .ok()
2170        .and_then(|json| {
2171            json.get("messages").and_then(Value::as_array).map(|msgs| {
2172                msgs.iter().any(|m| {
2173                    m.get("content")
2174                        .and_then(Value::as_array)
2175                        .is_some_and(|parts| {
2176                            parts
2177                                .iter()
2178                                .any(|p| p.get("type").and_then(Value::as_str) == Some("image_url"))
2179                        })
2180                })
2181            })
2182        })
2183        .unwrap_or(false)
2184}
2185
2186/// Build the routing/telemetry feature vector from an OpenAI Chat Completions request body.
2187/// Parallel to [`extract_features`] but understands OpenAI format (image_url vs image blocks).
2188fn extract_openai_features(headers: &HeaderMap, body: &[u8]) -> Features {
2189    let Ok(json) = serde_json::from_slice::<Value>(body) else {
2190        let mut f = Features::new(TaskKind::Other);
2191        f.hour_bucket = hour_bucket(jiff::Timestamp::now());
2192        return f;
2193    };
2194    let tool_count = json
2195        .get("tools")
2196        .and_then(Value::as_array)
2197        .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
2198    let has_images = openai_messages_have_images(body);
2199    let mut f = Features::new(TaskKind::Other);
2200    f.agent = header_str(headers, AGENT_HEADER);
2201    f.subagent = header_str(headers, SUBAGENT_HEADER);
2202    f.tool_count = tool_count;
2203    f.has_images = has_images;
2204    f.prompt_token_bucket = token_bucket(body.len() as u64);
2205    f.hour_bucket = hour_bucket(jiff::Timestamp::now());
2206    f
2207}
2208
2209// ── OpenAI → internal translation ────────────────────────────────────────────
2210
2211/// Parse a `data:image/<type>;base64,<data>` URL into `(media_type, base64_data)`.
2212fn parse_data_url(url: &str) -> Option<(&str, &str)> {
2213    let rest = url.strip_prefix("data:")?;
2214    let (meta, data) = rest.split_once(',')?;
2215    let media_type = meta.strip_suffix(";base64")?;
2216    Some((media_type, data))
2217}
2218
2219/// Translate an OpenAI user content value to Anthropic content blocks.
2220/// Returns `None` for any `image_url` part with an http(s) URL (non-translatable).
2221fn translate_openai_user_content(content: &Value) -> Option<Value> {
2222    match content {
2223        // Plain string → keep as-is (most common path)
2224        Value::String(_) => Some(content.clone()),
2225        Value::Array(parts) => {
2226            let mut blocks: Vec<Value> = Vec::with_capacity(parts.len());
2227            for part in parts {
2228                match part.get("type").and_then(Value::as_str) {
2229                    Some("text") => {
2230                        let text = part.get("text").and_then(Value::as_str).unwrap_or("");
2231                        blocks.push(serde_json::json!({ "type": "text", "text": text }));
2232                    }
2233                    Some("image_url") => {
2234                        let url = part.pointer("/image_url/url").and_then(Value::as_str)?;
2235                        if url.starts_with("http://") || url.starts_with("https://") {
2236                            return None; // not translatable — caller falls back to observe
2237                        }
2238                        // data: URI → Anthropic base64 image block
2239                        let (media_type, data) = parse_data_url(url)?;
2240                        blocks.push(serde_json::json!({
2241                            "type": "image",
2242                            "source": { "type": "base64", "media_type": media_type, "data": data }
2243                        }));
2244                    }
2245                    _ => {} // skip unknown content part types conservatively
2246                }
2247            }
2248            Some(Value::Array(blocks))
2249        }
2250        _ => Some(Value::String(String::new())),
2251    }
2252}
2253
2254/// Translate OpenAI `tools` array to Anthropic tools format.
2255/// OpenAI: `[{"type":"function","function":{"name":"...","description":"...","parameters":{...}}}]`
2256/// Anthropic: `[{"name":"...","description":"...","input_schema":{...}}]`
2257fn translate_openai_tools(tools: &Value) -> Value {
2258    let Some(arr) = tools.as_array() else {
2259        return Value::Null;
2260    };
2261    let converted: Vec<Value> = arr
2262        .iter()
2263        .map(|tool| {
2264            let func = tool.get("function").unwrap_or(&Value::Null);
2265            let mut out = serde_json::json!({
2266                "name": func.get("name").cloned().unwrap_or(Value::String(String::new())),
2267                "input_schema": func.get("parameters").cloned()
2268                    .unwrap_or_else(|| serde_json::json!({ "type": "object" })),
2269            });
2270            if let Some(desc) = func.get("description") {
2271                out["description"] = desc.clone();
2272            }
2273            out
2274        })
2275        .collect();
2276    Value::Array(converted)
2277}
2278
2279/// Translate OpenAI `tool_choice` to Anthropic `tool_choice`. Best-effort.
2280fn translate_openai_tool_choice(tc: &Value) -> Value {
2281    match tc {
2282        Value::String(s) => match s.as_str() {
2283            "auto" => serde_json::json!({ "type": "auto" }),
2284            "required" => serde_json::json!({ "type": "any" }),
2285            // ponytail: "none" has no direct Anthropic equivalent; omit = no constraint
2286            _ => serde_json::json!({ "type": "auto" }),
2287        },
2288        Value::Object(_) => {
2289            // {"type":"function","function":{"name":"foo"}} → {"type":"tool","name":"foo"}
2290            if tc.get("type").and_then(Value::as_str) == Some("function") {
2291                let name = tc.pointer("/function/name").cloned().unwrap_or(Value::Null);
2292                serde_json::json!({ "type": "tool", "name": name })
2293            } else {
2294                serde_json::json!({ "type": "auto" })
2295            }
2296        }
2297        _ => serde_json::json!({ "type": "auto" }),
2298    }
2299}
2300
2301/// Parse an OpenAI Chat Completions request body into the normalized [`ModelRequest`].
2302///
2303/// `carry_raw`: when `true` (all-OpenAI-dialect ladder), the original JSON is stored in
2304/// `raw` for verbatim carry — only the model is swapped, every other field survives intact.
2305/// When `false` (translation path to Anthropic ladder), `raw` is `Null` so
2306/// `anthropic_wire_body` reconstructs from the translated normalized fields.
2307///
2308/// Returns `None` if:
2309/// - the body isn't valid JSON or lacks a `messages` array, OR
2310/// - a user message contains an `image_url` with an http(s) URL (non-translatable; caller
2311///   should have already fallen back via `enforce_can_handle` but this is a defense-in-depth
2312///   guard — `None` → `BadRequest` rather than silently dropping the image).
2313pub fn parse_openai_request(body: &[u8], carry_raw: bool) -> Option<ModelRequest> {
2314    let json: Value = serde_json::from_slice(body).ok()?;
2315    let raw = if carry_raw { json.clone() } else { Value::Null };
2316
2317    let messages_json = json.get("messages")?.as_array()?;
2318
2319    let mut system: Option<String> = None;
2320    let mut messages: Vec<ChatMessage> = Vec::with_capacity(messages_json.len());
2321    let mut tools = Value::Null;
2322    let mut tool_choice_override: Option<Value> = None;
2323
2324    for msg in messages_json {
2325        let role = msg.get("role").and_then(Value::as_str).unwrap_or("user");
2326        match role {
2327            "system" => {
2328                // First system message wins; subsequent ones are appended as user blocks.
2329                // ponytail: Anthropic doesn't support multiple system messages inline;
2330                // we take the last one here. A proper multi-system-message implementation
2331                // would concatenate them, but that's rare in practice.
2332                if let Some(s) = msg.get("content").and_then(Value::as_str) {
2333                    system = Some(s.to_owned());
2334                }
2335            }
2336            "user" => {
2337                let content_val = msg.get("content").unwrap_or(&Value::Null);
2338                let translated = translate_openai_user_content(content_val)?;
2339                messages.push(ChatMessage {
2340                    role: "user".to_owned(),
2341                    content: translated,
2342                });
2343            }
2344            "assistant" => {
2345                if let Some(tc_arr) = msg.get("tool_calls").and_then(Value::as_array) {
2346                    // Tool-call turn: translate tool_calls to Anthropic tool_use blocks.
2347                    let mut blocks: Vec<Value> = Vec::new();
2348                    // Text before tool calls (may be null or absent)
2349                    if let Some(text) = msg.get("content").and_then(Value::as_str)
2350                        && !text.is_empty()
2351                    {
2352                        blocks.push(serde_json::json!({ "type": "text", "text": text }));
2353                    }
2354                    for tc in tc_arr {
2355                        let id = tc.get("id").and_then(Value::as_str).unwrap_or("");
2356                        let func = tc.get("function").unwrap_or(&Value::Null);
2357                        let name = func.get("name").and_then(Value::as_str).unwrap_or("");
2358                        let args_str = func
2359                            .get("arguments")
2360                            .and_then(Value::as_str)
2361                            .unwrap_or("{}");
2362                        let input: Value = serde_json::from_str(args_str)
2363                            .unwrap_or_else(|_| serde_json::json!({}));
2364                        blocks.push(serde_json::json!({
2365                            "type": "tool_use",
2366                            "id": id,
2367                            "name": name,
2368                            "input": input,
2369                        }));
2370                    }
2371                    messages.push(ChatMessage {
2372                        role: "assistant".to_owned(),
2373                        content: Value::Array(blocks),
2374                    });
2375                } else {
2376                    // Regular text assistant message
2377                    let content = msg
2378                        .get("content")
2379                        .cloned()
2380                        .unwrap_or_else(|| Value::String(String::new()));
2381                    messages.push(ChatMessage {
2382                        role: "assistant".to_owned(),
2383                        content,
2384                    });
2385                }
2386            }
2387            "tool" => {
2388                // role:"tool" → Anthropic tool_result block (wrapped in user turn)
2389                let tool_call_id = msg
2390                    .get("tool_call_id")
2391                    .and_then(Value::as_str)
2392                    .unwrap_or("");
2393                let content = msg
2394                    .get("content")
2395                    .cloned()
2396                    .unwrap_or_else(|| Value::String(String::new()));
2397                let result_block = serde_json::json!({
2398                    "type": "tool_result",
2399                    "tool_use_id": tool_call_id,
2400                    "content": content,
2401                });
2402                messages.push(ChatMessage {
2403                    role: "user".to_owned(),
2404                    content: Value::Array(vec![result_block]),
2405                });
2406            }
2407            _ => {} // skip unknown roles
2408        }
2409    }
2410
2411    // Translate tools and tool_choice (only when NOT raw-carry; raw-carry forwards them as-is).
2412    if !carry_raw {
2413        if let Some(t) = json.get("tools") {
2414            tools = translate_openai_tools(t);
2415        }
2416        if let Some(tc) = json.get("tool_choice") {
2417            tool_choice_override = Some(translate_openai_tool_choice(tc));
2418        }
2419    } else {
2420        tools = json.get("tools").cloned().unwrap_or(Value::Null);
2421    }
2422
2423    let max_tokens = json
2424        .get("max_tokens")
2425        .and_then(Value::as_u64)
2426        .and_then(|n| u32::try_from(n).ok())
2427        .unwrap_or(1024);
2428
2429    let model = json
2430        .get("model")
2431        .and_then(Value::as_str)
2432        .unwrap_or_default()
2433        .to_owned();
2434
2435    // For translation path, embed tool_choice into tools value so it's available downstream.
2436    // ponytail: this stuffs tool_choice into the Anthropic body via the raw=Null path in
2437    // anthropic_wire_body, which rebuilds from normalized fields. Tool_choice isn't a
2438    // ModelRequest field, so we carry it via a synthetic tools wrapper... actually we don't
2439    // need this — anthropic_wire_body rebuilds from normalized fields that include `tools`
2440    // but not `tool_choice`. The translation path loses tool_choice for non-raw-carry. This
2441    // is the known ceiling; full fidelity on mixed ladders requires adding tool_choice to
2442    // ModelRequest or always using raw carry.
2443    let _ = tool_choice_override; // accepted limitation on translation path
2444
2445    Some(ModelRequest {
2446        model,
2447        system,
2448        messages,
2449        max_tokens,
2450        tools,
2451        raw,
2452    })
2453}
2454
2455// ── Internal → OpenAI response rendering ─────────────────────────────────────
2456
2457/// Extract `(content_text, tool_calls)` from a served [`ModelResponse`]'s raw value.
2458///
2459/// Handles both Anthropic-format raw (has `content` array → translate to OpenAI shape)
2460/// and OpenAI-format raw (has `choices` → pass through content/tool_calls from the wire).
2461fn extract_openai_content_and_tools(raw: &Value, text: &str) -> (Value, Option<Value>) {
2462    // Anthropic-format: content array with text and/or tool_use blocks
2463    if let Some(blocks) = raw.get("content").and_then(Value::as_array) {
2464        let mut text_parts: Vec<&str> = Vec::new();
2465        let mut tool_calls: Vec<Value> = Vec::new();
2466        for block in blocks {
2467            match block.get("type").and_then(Value::as_str) {
2468                Some("text") => {
2469                    if let Some(t) = block.get("text").and_then(Value::as_str) {
2470                        text_parts.push(t);
2471                    }
2472                }
2473                Some("tool_use") => {
2474                    let id = block.get("id").and_then(Value::as_str).unwrap_or("");
2475                    let name = block.get("name").and_then(Value::as_str).unwrap_or("");
2476                    let input_str = block
2477                        .get("input")
2478                        .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
2479                    tool_calls.push(serde_json::json!({
2480                        "id": id,
2481                        "type": "function",
2482                        "function": { "name": name, "arguments": input_str },
2483                    }));
2484                }
2485                _ => {}
2486            }
2487        }
2488        let content_text = if tool_calls.is_empty() || !text_parts.is_empty() {
2489            Value::String(text_parts.join(""))
2490        } else {
2491            Value::Null // tool-only response: null content per OpenAI spec
2492        };
2493        let tc = if tool_calls.is_empty() {
2494            None
2495        } else {
2496            Some(Value::Array(tool_calls))
2497        };
2498        return (content_text, tc);
2499    }
2500
2501    // OpenAI-format raw (all-OpenAI-ladder path): extract from choices
2502    if let Some(msg) = raw.pointer("/choices/0/message") {
2503        let content = msg
2504            .get("content")
2505            .cloned()
2506            .unwrap_or(Value::String(text.to_owned()));
2507        let tc = msg.get("tool_calls").cloned();
2508        return (content, tc);
2509    }
2510
2511    // Fallback: use the text projection
2512    (Value::String(text.to_owned()), None)
2513}
2514
2515/// Render a served [`ModelResponse`] back as an OpenAI `chat.completion` JSON envelope,
2516/// so an OpenAI-client caller sees the standard wire shape regardless of which rung answered.
2517fn openai_response_json(resp: &ModelResponse) -> Value {
2518    let (content_text, tool_calls) = extract_openai_content_and_tools(&resp.raw, &resp.text);
2519    let finish_reason = if tool_calls.is_some() {
2520        "tool_calls"
2521    } else {
2522        "stop"
2523    };
2524    let mut message = serde_json::json!({
2525        "role": "assistant",
2526        "content": content_text,
2527    });
2528    if let Some(tc) = tool_calls {
2529        message["tool_calls"] = tc;
2530    }
2531    serde_json::json!({
2532        "id": format!("chatcmpl-{}", Uuid::now_v7()),
2533        "object": "chat.completion",
2534        "created": jiff::Timestamp::now().as_second(),
2535        "model": resp.model,
2536        "choices": [{
2537            "index": 0,
2538            "message": message,
2539            "finish_reason": finish_reason,
2540        }],
2541        "usage": {
2542            "prompt_tokens": resp.in_tokens,
2543            "completion_tokens": resp.out_tokens,
2544            "total_tokens": resp.in_tokens + resp.out_tokens,
2545        }
2546    })
2547}
2548
2549/// Re-emit a served OpenAI `chat.completion` envelope as an SSE stream body
2550/// (`data: chat.completion.chunk` frames ending with `data: [DONE]`), so a `stream: true`
2551/// OpenAI client is served even though enforce buffered the full response to gate it.
2552fn openai_sse_from_message(message: &Value) -> String {
2553    let id = message
2554        .get("id")
2555        .and_then(Value::as_str)
2556        .unwrap_or("chatcmpl-unknown")
2557        .to_owned();
2558    let created = message.get("created").cloned().unwrap_or(Value::from(0));
2559    let model = message
2560        .get("model")
2561        .and_then(Value::as_str)
2562        .unwrap_or("unknown");
2563
2564    let choices = message.get("choices").and_then(Value::as_array);
2565    let msg = choices
2566        .and_then(|c| c.first())
2567        .and_then(|c| c.get("message"));
2568    let content = msg
2569        .and_then(|m| m.get("content"))
2570        .cloned()
2571        .unwrap_or(Value::Null);
2572    let tool_calls = msg.and_then(|m| m.get("tool_calls")).cloned();
2573    let finish_reason = choices
2574        .and_then(|c| c.first())
2575        .and_then(|c| c.get("finish_reason"))
2576        .cloned()
2577        .unwrap_or_else(|| Value::String("stop".to_owned()));
2578
2579    let mut out = String::new();
2580    let chunk = |delta: Value| {
2581        serde_json::json!({
2582            "id": id,
2583            "object": "chat.completion.chunk",
2584            "created": created,
2585            "model": model,
2586            "choices": [{ "index": 0, "delta": delta, "finish_reason": Value::Null }]
2587        })
2588    };
2589
2590    // Role delta
2591    let role_chunk = chunk(serde_json::json!({ "role": "assistant", "content": "" }));
2592    out.push_str("data: ");
2593    out.push_str(&role_chunk.to_string());
2594    out.push_str("\n\n");
2595
2596    // Content delta (if any)
2597    if let Value::String(text) = &content
2598        && !text.is_empty()
2599    {
2600        let content_chunk = chunk(serde_json::json!({ "content": text }));
2601        out.push_str("data: ");
2602        out.push_str(&content_chunk.to_string());
2603        out.push_str("\n\n");
2604    }
2605
2606    // Tool calls delta (if any)
2607    if let Some(tc) = tool_calls {
2608        let tc_chunk = chunk(serde_json::json!({ "tool_calls": tc }));
2609        out.push_str("data: ");
2610        out.push_str(&tc_chunk.to_string());
2611        out.push_str("\n\n");
2612    }
2613
2614    // Finish chunk
2615    let finish_chunk = serde_json::json!({
2616        "id": id,
2617        "object": "chat.completion.chunk",
2618        "created": created,
2619        "model": model,
2620        "choices": [{ "index": 0, "delta": {}, "finish_reason": finish_reason }]
2621    });
2622    out.push_str("data: ");
2623    out.push_str(&finish_chunk.to_string());
2624    out.push_str("\n\n");
2625
2626    out.push_str("data: [DONE]\n\n");
2627    out
2628}
2629
2630// ── OpenAI handler path ───────────────────────────────────────────────────────
2631
2632/// Enforce mode for an OpenAI-inbound request: run the escalation engine and serve the
2633/// first output that clears the route's gates, rendered as an OpenAI `chat.completion`.
2634#[allow(clippy::too_many_arguments)]
2635async fn handle_enforce_openai(
2636    state: &AppState,
2637    headers: &HeaderMap,
2638    body: &Bytes,
2639    features: Features,
2640    route: &Route,
2641    // Index of the matched route, stamped on the trace so a downstream outcome arriving
2642    // later can be attributed back to the route that produced it (ADR 0009 D3).
2643    route_ix: usize,
2644    session_header: Option<String>,
2645    tenant: String,
2646    routing_mode: RoutingMode,
2647) -> Response {
2648    if is_stream_request(body) {
2649        let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
2650        let (state_c, headers_c, body_c, route_c) =
2651            (state.clone(), headers.clone(), body.clone(), route.clone());
2652        tokio::spawn(async move {
2653            let out = enforce_pipeline_openai(
2654                &state_c,
2655                &headers_c,
2656                &body_c,
2657                features,
2658                &route_c,
2659                route_ix,
2660                session_header,
2661                tenant,
2662                routing_mode,
2663            )
2664            .await;
2665            let _ = tx.send(out);
2666        });
2667        return sse_keepalive_response(rx, openai_sse_from_message);
2668    }
2669    match enforce_pipeline_openai(
2670        state,
2671        headers,
2672        body,
2673        features,
2674        route,
2675        route_ix,
2676        session_header,
2677        tenant,
2678        routing_mode,
2679    )
2680    .await
2681    {
2682        Ok(message) => (axum::http::StatusCode::OK, Json(message)).into_response(),
2683        Err(e) => e.into_response(),
2684    }
2685}
2686
2687/// Observe mode for `POST /v1/chat/completions`: forward unchanged to the OpenAI upstream,
2688/// return unchanged, trace asynchronously.
2689///
2690// ponytail: observe trace stamps api="anthropic.messages" (base_trace default); update
2691// base_trace to accept an api param if the distinction matters for audit consumers.
2692async fn observe_passthrough_openai(
2693    state: AppState,
2694    headers: HeaderMap,
2695    body: Bytes,
2696    session_header: Option<String>,
2697    tenant: String,
2698) -> Response {
2699    if is_stream_request(&body) {
2700        return observe_stream_openai(state, headers, body, session_header, tenant).await;
2701    }
2702    let start = Instant::now();
2703    let result = forward_openai(
2704        &state.http,
2705        &state.config.upstream_openai,
2706        &headers,
2707        body.clone(),
2708    )
2709    .await;
2710    let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
2711    match result {
2712        Ok((status, resp_headers, resp_body)) => {
2713            spawn_trace(
2714                &state,
2715                body,
2716                Some(resp_body.clone()),
2717                latency_ms,
2718                session_header,
2719                tenant,
2720            );
2721            (status, resp_headers, resp_body).into_response()
2722        }
2723        Err(err) => {
2724            spawn_trace(&state, body, None, latency_ms, session_header, tenant);
2725            err.into_response()
2726        }
2727    }
2728}
2729
2730/// Observe streaming for `POST /v1/chat/completions`.
2731async fn observe_stream_openai(
2732    state: AppState,
2733    headers: HeaderMap,
2734    body: Bytes,
2735    session_header: Option<String>,
2736    tenant: String,
2737) -> Response {
2738    let start = Instant::now();
2739    let result = forward_openai_streaming(
2740        &state.http,
2741        &state.config.upstream_openai,
2742        &headers,
2743        body.clone(),
2744    )
2745    .await;
2746    let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
2747    match result {
2748        Ok((status, resp_headers, response)) => {
2749            spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
2750            let stream_body = Body::from_stream(response.bytes_stream());
2751            (status, resp_headers, stream_body).into_response()
2752        }
2753        Err(err) => {
2754            spawn_trace(&state, body, None, latency_ms, session_header, tenant);
2755            err.into_response()
2756        }
2757    }
2758}
2759
2760/// `POST /v1/chat/completions` — OpenAI-compatible inbound endpoint (SPEC §M1).
2761/// Observe mode: transparent passthrough to the OpenAI upstream base URL.
2762/// Enforce mode: translate to the internal `ModelRequest`, run the escalation engine,
2763/// render the result as an OpenAI `chat.completion` (or `chat.completion.chunk` SSE stream).
2764async fn chat_completions(
2765    State(state): State<AppState>,
2766    Extension(TenantId(tenant)): Extension<TenantId>,
2767    headers: HeaderMap,
2768    body: Bytes,
2769) -> Response {
2770    let session_header = header_str(&headers, SESSION_HEADER);
2771
2772    if let Some(routing) = state.config.routing.as_ref() {
2773        let features = extract_openai_features(&headers, &body);
2774        if let Some(route) = routing
2775            .route_for(&features)
2776            .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
2777        {
2778            // Index of the matched route, for outcome attribution (ADR 0009 D3).
2779            let route_ix = routing
2780                .routes
2781                .iter()
2782                .position(|r| std::ptr::eq(r, route))
2783                .unwrap_or(0);
2784            let route = route.clone();
2785            // Resolve routing-mode preset (header > route > global default).
2786            let routing_mode = resolve_mode(&headers, &route, &state.config);
2787            // Observe mode forces the observe passthrough path — no gating, no escalation.
2788            if routing_mode == RoutingMode::Observe {
2789                return observe_passthrough_openai(state, headers, body, session_header, tenant)
2790                    .await;
2791            }
2792            if enforce_can_handle(
2793                &features,
2794                &body,
2795                routing.escalation.enforce_structured,
2796                &route.ladder,
2797                &state.providers,
2798                Dialect::Openai,
2799            ) {
2800                return handle_enforce_openai(
2801                    &state,
2802                    &headers,
2803                    &body,
2804                    features,
2805                    &route,
2806                    route_ix,
2807                    session_header,
2808                    tenant,
2809                    routing_mode,
2810                )
2811                .await;
2812            }
2813            tracing::info!(
2814                "enforce route matched but OpenAI structured request can't be routed faithfully (flag/ladder); serving via observe passthrough"
2815            );
2816        }
2817    }
2818    observe_passthrough_openai(state, headers, body, session_header, tenant).await
2819}
2820
2821#[cfg(test)]
2822mod tests {
2823    use bytes::Bytes;
2824
2825    use super::*;
2826
2827    fn test_config() -> ProxyConfig {
2828        ProxyConfig::from_lookup(|_| None).unwrap()
2829    }
2830
2831    #[test]
2832    fn build_trace_maps_request_and_response_fields() {
2833        let config = test_config();
2834        let req = Bytes::from_static(
2835            br#"{"model":"claude-haiku-4-5","tools":[{"name":"a"}],"messages":[]}"#,
2836        );
2837        let resp = Bytes::from_static(
2838            br#"{"model":"claude-haiku-4-5","usage":{"input_tokens":1200,"output_tokens":300}}"#,
2839        );
2840
2841        let trace = build_trace(&config, &req, &resp, 42, Some("sess-1"));
2842
2843        assert_eq!(trace.request.api, "anthropic.messages");
2844        assert_eq!(trace.session_id, "sess-1");
2845        assert_eq!(trace.attempts.len(), 1);
2846        let attempt = &trace.attempts[0];
2847        assert_eq!(attempt.model, "claude-haiku-4-5");
2848        assert_eq!(attempt.provider, "anthropic");
2849        assert_eq!(attempt.in_tokens, 1200);
2850        assert_eq!(attempt.out_tokens, 300);
2851        assert!(attempt.cost_usd > 0.0);
2852        assert_eq!(trace.request.features.tool_count, 1);
2853        assert!(!trace.request.features.has_images);
2854        assert_eq!(trace.final_.served_rung, Some(0));
2855    }
2856
2857    #[test]
2858    fn build_trace_falls_back_to_trace_id_session_when_header_absent() {
2859        let config = test_config();
2860        let req = Bytes::from_static(b"{}");
2861        let resp = Bytes::from_static(b"{}");
2862
2863        let trace = build_trace(&config, &req, &resp, 1, None);
2864
2865        assert_eq!(trace.session_id, trace.trace_id.to_string());
2866    }
2867
2868    #[test]
2869    fn build_error_trace_has_no_attempts_and_served_from_error() {
2870        let config = test_config();
2871        let req = Bytes::from_static(br#"{"model":"claude-haiku-4-5"}"#);
2872
2873        let trace = build_error_trace(&config, &req, 7, None);
2874
2875        assert!(trace.attempts.is_empty());
2876        assert_eq!(trace.final_.served_from, ServedFrom::Error);
2877        assert_eq!(trace.final_.served_rung, None);
2878    }
2879
2880    #[test]
2881    fn message_with_image_block_sets_has_images() {
2882        let req = br#"{"messages":[{"role":"user","content":[{"type":"image"}]}]}"#;
2883        let (_, _, has_images) = request_features(req);
2884        assert!(has_images);
2885    }
2886
2887    #[test]
2888    fn prompt_hash_never_contains_raw_prompt_text() {
2889        let hash = prompt_hash("salt", b"super secret prompt");
2890        assert!(!hash.contains("secret"));
2891        assert_eq!(hash.len(), 64);
2892    }
2893
2894    #[test]
2895    fn parse_model_request_preserves_content_verbatim_and_projects_text() {
2896        let body = br#"{"model":"m","system":"sys","max_tokens":50,
2897            "messages":[{"role":"user","content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]},
2898                        {"role":"assistant","content":"c"}]}"#;
2899        let req = parse_model_request(body).unwrap();
2900        assert_eq!(req.system.as_deref(), Some("sys"));
2901        assert_eq!(req.max_tokens, 50);
2902        assert_eq!(req.messages.len(), 2);
2903        // I2: the block array is carried verbatim, not flattened away...
2904        assert_eq!(
2905            req.messages[0].content,
2906            serde_json::json!([{"type":"text","text":"a"},{"type":"text","text":"b"}])
2907        );
2908        // ...and a plain string stays a plain string (I1: byte-identical on the wire).
2909        assert_eq!(req.messages[1].content, Value::String("c".to_owned()));
2910        // Gates see the same text they always did.
2911        assert_eq!(req.messages[0].text_view(), "a\nb");
2912        assert_eq!(req.messages[1].text_view(), "c");
2913    }
2914
2915    #[test]
2916    fn tool_and_image_blocks_survive_the_request_round_trip() {
2917        // ADR 0005 I2: tool_use / tool_result / image blocks are never dropped on the request side.
2918        let body = br#"{"model":"m","max_tokens":50,"messages":[
2919            {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
2920            {"role":"user","content":[
2921                {"type":"tool_result","tool_use_id":"t1","content":"2"},
2922                {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
2923            ]}]}"#;
2924        let req = parse_model_request(body).unwrap();
2925        let round_tripped = serde_json::to_value(&req.messages).unwrap();
2926        assert_eq!(
2927            round_tripped,
2928            serde_json::json!([
2929                {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
2930                {"role":"user","content":[
2931                    {"type":"tool_result","tool_use_id":"t1","content":"2"},
2932                    {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
2933                ]}
2934            ])
2935        );
2936    }
2937
2938    #[test]
2939    fn text_message_serializes_byte_identical_to_a_plain_string() {
2940        // I1: a string-content message must not gain array wrapping on the wire.
2941        let m = ChatMessage::text("user", "hello");
2942        assert_eq!(
2943            serde_json::to_string(&m).unwrap(),
2944            r#"{"role":"user","content":"hello"}"#
2945        );
2946    }
2947
2948    #[test]
2949    fn parse_model_request_rejects_non_message_bodies() {
2950        assert!(parse_model_request(b"not json").is_none());
2951        assert!(parse_model_request(br#"{"no":"messages"}"#).is_none());
2952    }
2953
2954    // --- Enforce-path handler tests (drive `messages` end-to-end with mock providers) ---
2955
2956    use crate::provider::{MockProvider, ModelResponse, Provider, ProviderError, ProviderRegistry};
2957    use axum::extract::State;
2958    use std::collections::HashMap;
2959    use std::sync::Arc;
2960    use tokio::sync::mpsc;
2961
2962    fn model_resp(model: &str, text: &str) -> ModelResponse {
2963        ModelResponse {
2964            model: model.to_owned(),
2965            text: text.to_owned(),
2966            in_tokens: 1000,
2967            out_tokens: 400,
2968            raw: serde_json::Value::Null,
2969        }
2970    }
2971
2972    /// Build an `AppState` whose anthropic provider answers the given per-model outcomes, with an
2973    /// enforce route over `ladder`/`gates`. Returns the state and the trace receiver.
2974    fn enforce_state(
2975        ladder: &[&str],
2976        gates: &[&str],
2977        outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
2978    ) -> (AppState, mpsc::Receiver<Trace>) {
2979        let toml = format!(
2980            "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [{}]\ngates = [{}]\n",
2981            ladder
2982                .iter()
2983                .map(|m| format!("\"{m}\""))
2984                .collect::<Vec<_>>()
2985                .join(", "),
2986            gates
2987                .iter()
2988                .map(|g| format!("\"{g}\""))
2989                .collect::<Vec<_>>()
2990                .join(", "),
2991        );
2992        let config = ProxyConfig::from_lookup(|k| match k {
2993            "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
2994            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
2995            _ => None,
2996        })
2997        .unwrap();
2998
2999        let mut outs = HashMap::new();
3000        for (model, out) in outcomes {
3001            outs.insert(model.to_owned(), out);
3002        }
3003        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3004        map.insert(
3005            "anthropic".to_owned(),
3006            Arc::new(MockProvider::new("anthropic", outs)),
3007        );
3008        let providers = ProviderRegistry::from_map(map);
3009
3010        let (traces, rx) = mpsc::channel(64);
3011        let state = AppState {
3012            config: Arc::new(config),
3013            http: reqwest::Client::new(),
3014            providers,
3015            gate_health: Arc::new(GateHealthRegistry::new()),
3016            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3017            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3018            traces,
3019            adaptive: None,
3020            bandit: None,
3021            predictor: None,
3022            tenant_rate_limiter: None,
3023            spill: None,
3024        };
3025        (state, rx)
3026    }
3027
3028    fn user_body() -> Bytes {
3029        Bytes::from_static(
3030            br#"{"model":"ignored","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}"#,
3031        )
3032    }
3033
3034    async fn body_json(resp: Response) -> Value {
3035        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
3036            .await
3037            .unwrap();
3038        serde_json::from_slice(&bytes).unwrap()
3039    }
3040
3041    #[tokio::test]
3042    async fn mode_profile_stamped_on_trace_when_non_balanced() {
3043        let (state, mut rx) = enforce_state(
3044            &["anthropic/claude-haiku-4-5"],
3045            &["non-empty"],
3046            vec![(
3047                "anthropic/claude-haiku-4-5",
3048                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3049            )],
3050        );
3051        let mut headers = HeaderMap::new();
3052        headers.insert("x-firstpass-mode", "quality".parse().unwrap());
3053        let resp = messages(
3054            State(state),
3055            Extension(TenantId("default".to_owned())),
3056            headers,
3057            user_body(),
3058        )
3059        .await;
3060        assert_eq!(resp.status(), axum::http::StatusCode::OK);
3061        let trace = rx.try_recv().expect("trace enqueued");
3062        assert_eq!(
3063            trace.policy.mode_profile.as_deref(),
3064            Some("quality"),
3065            "mode_profile must be stamped when quality mode is active"
3066        );
3067    }
3068
3069    #[tokio::test]
3070    async fn mode_profile_absent_from_trace_when_balanced() {
3071        let (state, mut rx) = enforce_state(
3072            &["anthropic/claude-haiku-4-5"],
3073            &["non-empty"],
3074            vec![(
3075                "anthropic/claude-haiku-4-5",
3076                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3077            )],
3078        );
3079        // No mode header, no route routing_mode → Balanced by default → None.
3080        let resp = messages(
3081            State(state),
3082            Extension(TenantId("default".to_owned())),
3083            HeaderMap::new(),
3084            user_body(),
3085        )
3086        .await;
3087        assert_eq!(resp.status(), axum::http::StatusCode::OK);
3088        let trace = rx.try_recv().expect("trace enqueued");
3089        assert!(
3090            trace.policy.mode_profile.is_none(),
3091            "mode_profile must be absent when Balanced (byte-identical invariant)"
3092        );
3093    }
3094
3095    #[tokio::test]
3096    async fn enforce_serves_first_pass_and_returns_anthropic_shape() {
3097        let (state, mut rx) = enforce_state(
3098            &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
3099            &["non-empty"],
3100            vec![(
3101                "anthropic/claude-haiku-4-5",
3102                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3103            )],
3104        );
3105        let resp = messages(
3106            State(state),
3107            Extension(TenantId("default".to_owned())),
3108            HeaderMap::new(),
3109            user_body(),
3110        )
3111        .await;
3112        assert_eq!(resp.status(), axum::http::StatusCode::OK);
3113        let json = body_json(resp).await;
3114        assert_eq!(json["type"], "message");
3115        assert_eq!(json["content"][0]["text"], "hello");
3116        assert_eq!(json["model"], "anthropic/claude-haiku-4-5");
3117
3118        let trace = rx.try_recv().expect("a trace was enqueued");
3119        assert_eq!(trace.mode, Mode::Enforce);
3120        assert_eq!(trace.final_.served_rung, Some(0));
3121        assert_eq!(trace.attempts.len(), 1);
3122    }
3123
3124    #[tokio::test]
3125    async fn enforce_escalates_then_serves_and_traces_two_attempts() {
3126        let (state, mut rx) = enforce_state(
3127            &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
3128            &["non-empty"],
3129            vec![
3130                (
3131                    "anthropic/claude-haiku-4-5",
3132                    Ok(model_resp("anthropic/claude-haiku-4-5", "   ")),
3133                ), // fails
3134                (
3135                    "anthropic/claude-sonnet-5",
3136                    Ok(model_resp("anthropic/claude-sonnet-5", "answer")),
3137                ),
3138            ],
3139        );
3140        let resp = messages(
3141            State(state),
3142            Extension(TenantId("default".to_owned())),
3143            HeaderMap::new(),
3144            user_body(),
3145        )
3146        .await;
3147        let json = body_json(resp).await;
3148        assert_eq!(json["content"][0]["text"], "answer");
3149
3150        let trace = rx.try_recv().expect("trace enqueued");
3151        assert_eq!(trace.attempts.len(), 2);
3152        assert_eq!(trace.final_.escalations, 1);
3153        assert_eq!(trace.final_.served_rung, Some(1));
3154    }
3155
3156    #[tokio::test]
3157    async fn enforce_all_rungs_error_returns_502() {
3158        let (state, mut rx) = enforce_state(
3159            &["anthropic/claude-haiku-4-5"],
3160            &["non-empty"],
3161            vec![(
3162                "anthropic/claude-haiku-4-5",
3163                Err(ProviderError::Transport("down".into())),
3164            )],
3165        );
3166        let resp = messages(
3167            State(state),
3168            Extension(TenantId("default".to_owned())),
3169            HeaderMap::new(),
3170            user_body(),
3171        )
3172        .await;
3173        assert_eq!(resp.status(), axum::http::StatusCode::BAD_GATEWAY);
3174        // A trace is still recorded for the failed decision.
3175        assert!(rx.try_recv().is_ok());
3176    }
3177
3178    #[tokio::test]
3179    async fn no_routing_config_falls_through_to_observe_not_enforce() {
3180        // config with no routing => enforce path never runs; observe attempts a real upstream
3181        // call which fails fast against an unroutable host. We only assert it did NOT take the
3182        // enforce branch (which would have used the mock and returned 200 with our text).
3183        let config = ProxyConfig::from_lookup(|k| match k {
3184            "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3185            _ => None,
3186        })
3187        .unwrap();
3188        let (traces, _rx) = mpsc::channel(64);
3189        let state = AppState {
3190            config: Arc::new(config),
3191            http: reqwest::Client::new(),
3192            providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
3193            gate_health: Arc::new(GateHealthRegistry::new()),
3194            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3195            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3196            traces,
3197            adaptive: None,
3198            bandit: None,
3199            predictor: None,
3200            tenant_rate_limiter: None,
3201            spill: None,
3202        };
3203        let resp = messages(
3204            State(state),
3205            Extension(TenantId("default".to_owned())),
3206            HeaderMap::new(),
3207            user_body(),
3208        )
3209        .await;
3210        // Observe path forwards upstream; the bogus host yields a gateway error, not our 200.
3211        assert_ne!(resp.status(), axum::http::StatusCode::OK);
3212    }
3213
3214    // ── resolve_mode tests ────────────────────────────────────────────────────
3215
3216    fn bare_enforce_route() -> Route {
3217        use firstpass_core::config::{Match, Mode};
3218        Route {
3219            match_: Match::default(),
3220            mode: Mode::Enforce,
3221            ladder: vec!["anthropic/claude-haiku-4-5".to_owned()],
3222            gates: vec![],
3223            deferred_gates: vec![],
3224            routing_mode: None,
3225            rollout: None,
3226            shadow: None,
3227        }
3228    }
3229
3230    #[test]
3231    fn balanced_preset_application_is_noop() {
3232        // The core invariant: applying the Balanced preset to any base values returns them unchanged.
3233        let base_max_rungs = 3u32;
3234        let base_speculation = 2u32;
3235        let preset = RoutingMode::Balanced.preset();
3236        let max_rungs = if let Some(d) = preset.max_rungs_delta {
3237            (base_max_rungs as i32 + d).max(1) as u32
3238        } else {
3239            base_max_rungs
3240        };
3241        let speculation = preset.speculation.unwrap_or(base_speculation);
3242        let start_at_top = preset.start_at_top;
3243        assert_eq!(max_rungs, 3, "Balanced must not change max_rungs");
3244        assert_eq!(speculation, 2, "Balanced must not change speculation");
3245        assert!(!start_at_top, "Balanced must not set start_at_top");
3246    }
3247
3248    #[test]
3249    fn resolve_mode_header_wins_over_route_and_global() {
3250        let mut headers = HeaderMap::new();
3251        headers.insert("x-firstpass-mode", "cost".parse().unwrap());
3252        let mut route = bare_enforce_route();
3253        route.routing_mode = Some(RoutingMode::Quality); // lower priority
3254        let mut config = test_config();
3255        config.default_routing_mode = RoutingMode::Max; // lowest priority
3256        assert_eq!(
3257            resolve_mode(&headers, &route, &config),
3258            RoutingMode::Cost,
3259            "header must win"
3260        );
3261    }
3262
3263    #[test]
3264    fn resolve_mode_route_wins_over_global_when_no_header() {
3265        let mut route = bare_enforce_route();
3266        route.routing_mode = Some(RoutingMode::Latency);
3267        let mut config = test_config();
3268        config.default_routing_mode = RoutingMode::Max;
3269        assert_eq!(
3270            resolve_mode(&HeaderMap::new(), &route, &config),
3271            RoutingMode::Latency,
3272            "route must beat global default"
3273        );
3274    }
3275
3276    #[test]
3277    fn resolve_mode_global_when_header_and_route_absent() {
3278        let mut config = test_config();
3279        config.default_routing_mode = RoutingMode::Quality;
3280        assert_eq!(
3281            resolve_mode(&HeaderMap::new(), &bare_enforce_route(), &config),
3282            RoutingMode::Quality
3283        );
3284    }
3285
3286    #[test]
3287    fn resolve_mode_unknown_header_falls_through_to_route() {
3288        let mut headers = HeaderMap::new();
3289        // An unrecognised value must be ignored (warn + fall through).
3290        headers.insert("x-firstpass-mode", "turbo-mode".parse().unwrap());
3291        let mut route = bare_enforce_route();
3292        route.routing_mode = Some(RoutingMode::Cost);
3293        let config = test_config();
3294        assert_eq!(
3295            resolve_mode(&headers, &route, &config),
3296            RoutingMode::Cost,
3297            "unknown header value must fall through to route"
3298        );
3299    }
3300
3301    #[test]
3302    fn resolve_mode_header_case_insensitive() {
3303        let mut headers = HeaderMap::new();
3304        headers.insert("x-firstpass-mode", "QUALITY".parse().unwrap());
3305        assert_eq!(
3306            resolve_mode(&headers, &bare_enforce_route(), &test_config()),
3307            RoutingMode::Quality
3308        );
3309    }
3310
3311    #[test]
3312    fn resolve_mode_no_mode_set_returns_balanced() {
3313        // Default config has Balanced; route has None → must return Balanced.
3314        assert_eq!(
3315            resolve_mode(&HeaderMap::new(), &bare_enforce_route(), &test_config()),
3316            RoutingMode::Balanced
3317        );
3318    }
3319
3320    #[test]
3321    fn capabilities_json_includes_routing_modes() {
3322        let modes: Vec<&'static str> = RoutingMode::ALL.iter().map(|m| m.as_str()).collect();
3323        assert!(modes.contains(&"balanced"));
3324        assert!(modes.contains(&"cost"));
3325        assert!(modes.contains(&"quality"));
3326        assert!(modes.contains(&"latency"));
3327        assert!(modes.contains(&"max"));
3328        assert!(modes.contains(&"observe"));
3329    }
3330
3331    #[test]
3332    fn detects_stream_requests() {
3333        assert!(is_stream_request(br#"{"stream": true}"#));
3334        assert!(!is_stream_request(br#"{"stream": false}"#));
3335        assert!(!is_stream_request(br#"{"model":"m"}"#));
3336        assert!(!is_stream_request(b"not json"));
3337    }
3338
3339    #[test]
3340    fn detects_tool_blocks_in_messages() {
3341        let with =
3342            br#"{"messages":[{"role":"user","content":[{"type":"tool_result","content":"42"}]}]}"#;
3343        let without = br#"{"messages":[{"role":"user","content":"hi"}]}"#;
3344        assert!(messages_have_tool_blocks(with));
3345        assert!(!messages_have_tool_blocks(without));
3346    }
3347
3348    #[test]
3349    fn enforce_only_handles_plain_text() {
3350        let plain =
3351            Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
3352        let tools = Bytes::from_static(
3353            br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3354        );
3355        let f_plain = extract_features(&HeaderMap::new(), &plain);
3356        let f_tools = extract_features(&HeaderMap::new(), &tools);
3357        let anthropic_ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
3358        let providers = test_registry();
3359        // Opted out (enforce_structured = false): plain text routes, tools fall back to observe.
3360        assert!(enforce_can_handle(
3361            &f_plain,
3362            &plain,
3363            false,
3364            &anthropic_ladder,
3365            &providers,
3366            Dialect::Anthropic,
3367        ));
3368        assert!(!enforce_can_handle(
3369            &f_tools,
3370            &tools,
3371            false,
3372            &anthropic_ladder,
3373            &providers,
3374            Dialect::Anthropic,
3375        ));
3376    }
3377
3378    #[test]
3379    fn structured_enforce_routes_tools_and_streaming() {
3380        // ADR 0005 P2+P3: with the opt-in flag on, tool and streaming requests both route through
3381        // enforce (streaming is served as the gated result re-emitted as SSE).
3382        let tools = Bytes::from_static(
3383            br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3384        );
3385        let streaming_tools = Bytes::from_static(
3386            br#"{"model":"m","stream":true,"tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3387        );
3388        let f = extract_features(&HeaderMap::new(), &tools);
3389        let anthropic_ladder = vec![
3390            "anthropic/claude-haiku-4-5".to_owned(),
3391            "anthropic/claude-sonnet-5".to_owned(),
3392        ];
3393        let providers = test_registry();
3394        assert!(enforce_can_handle(
3395            &f,
3396            &tools,
3397            true,
3398            &anthropic_ladder,
3399            &providers,
3400            Dialect::Anthropic,
3401        ));
3402        assert!(enforce_can_handle(
3403            &f,
3404            &streaming_tools,
3405            true,
3406            &anthropic_ladder,
3407            &providers,
3408            Dialect::Anthropic,
3409        ));
3410    }
3411
3412    /// Registry with the built-in `anthropic` (verbatim carrier) + `openai` (not yet) providers.
3413    fn test_registry() -> crate::provider::ProviderRegistry {
3414        crate::provider::ProviderRegistry::new("http://localhost", "http://localhost")
3415    }
3416
3417    #[test]
3418    fn fidelity_guard_blocks_structured_on_non_verbatim_ladder() {
3419        // The default-on guard: a tool request routes through an all-Anthropic ladder, but a
3420        // ladder containing an OpenAI-dialect rung (structured translation not built) falls back
3421        // to observe — un-gated is safe, corrupted is not. Plain text routes on any ladder.
3422        let tools = Bytes::from_static(
3423            br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3424        );
3425        let plain =
3426            Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
3427        let f_tools = extract_features(&HeaderMap::new(), &tools);
3428        let f_plain = extract_features(&HeaderMap::new(), &plain);
3429        let providers = test_registry();
3430        let mixed_ladder = vec![
3431            "openai/gpt-4.1-mini".to_owned(),
3432            "anthropic/claude-sonnet-5".to_owned(),
3433        ];
3434        assert!(!enforce_can_handle(
3435            &f_tools,
3436            &tools,
3437            true,
3438            &mixed_ladder,
3439            &providers,
3440            Dialect::Anthropic,
3441        ));
3442        assert!(enforce_can_handle(
3443            &f_plain,
3444            &plain,
3445            true,
3446            &mixed_ladder,
3447            &providers,
3448            Dialect::Anthropic,
3449        ));
3450    }
3451
3452    #[test]
3453    fn enforce_sse_reemission_preserves_text_and_tool_use() {
3454        // ADR 0005 P3 + I2: a served response with a text block AND a tool_use block round-trips
3455        // through the SSE re-emitter — the tool call's input survives as an input_json_delta.
3456        let resp = ModelResponse {
3457            model: "anthropic/claude-haiku-4-5".to_owned(),
3458            text: "let me check".to_owned(),
3459            in_tokens: 5,
3460            out_tokens: 7,
3461            raw: serde_json::json!({
3462                "content": [
3463                    { "type": "text", "text": "let me check" },
3464                    { "type": "tool_use", "id": "tu_1", "name": "get_weather", "input": { "city": "Paris" } }
3465                ]
3466            }),
3467        };
3468        let sse = anthropic_sse_from_message(&anthropic_response_json(&resp));
3469
3470        // Parse every data frame structurally (key order is not part of the contract).
3471        let frames: Vec<Value> = sse
3472            .lines()
3473            .filter_map(|l| l.strip_prefix("data: "))
3474            .map(|d| serde_json::from_str::<Value>(d).expect("each SSE data frame is valid JSON"))
3475            .collect();
3476
3477        // Full lifecycle, in order.
3478        assert_eq!(frames.first().unwrap()["type"], "message_start");
3479        assert_eq!(frames.last().unwrap()["type"], "message_stop");
3480        // The text block streams its text as a text_delta.
3481        assert!(frames.iter().any(|f| f["delta"]["type"] == "text_delta"
3482            && f["delta"]["text"] == "let me check"));
3483        // The tool_use block is present with its id/name, and its input streams as one JSON delta —
3484        // not dropped (ADR 0005 I2).
3485        assert!(
3486            frames
3487                .iter()
3488                .any(|f| f["content_block"]["type"] == "tool_use"
3489                    && f["content_block"]["name"] == "get_weather"
3490                    && f["content_block"]["id"] == "tu_1")
3491        );
3492        assert!(
3493            frames
3494                .iter()
3495                .any(|f| f["delta"]["type"] == "input_json_delta"
3496                    && f["delta"]["partial_json"] == r#"{"city":"Paris"}"#)
3497        );
3498    }
3499
3500    /// ADR 0005, default-on: an enforce route now serves BOTH plain text and tool requests (the
3501    /// mock ladder carries structured content verbatim). Setting `enforce_structured = false`
3502    /// restores the old behavior: tool requests fall back to transparent observe passthrough —
3503    /// proven by the tool request hitting the (bogus) upstream instead of the enforcing mock.
3504    #[tokio::test]
3505    async fn enforce_falls_back_to_observe_for_tool_requests() {
3506        let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
3507        let config = ProxyConfig::from_lookup(|k| match k {
3508            "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
3509            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
3510            "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3511            _ => None,
3512        })
3513        .unwrap();
3514        let mut outs = HashMap::new();
3515        outs.insert(
3516            "anthropic/m".to_owned(),
3517            Ok(model_resp("anthropic/m", "hello")),
3518        );
3519        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3520        map.insert(
3521            "anthropic".to_owned(),
3522            Arc::new(MockProvider::new("anthropic", outs)),
3523        );
3524        let (traces, _rx) = mpsc::channel(64);
3525        let state = AppState {
3526            config: Arc::new(config),
3527            http: reqwest::Client::new(),
3528            providers: ProviderRegistry::from_map(map),
3529            gate_health: Arc::new(GateHealthRegistry::new()),
3530            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3531            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3532            traces,
3533            adaptive: None,
3534            bandit: None,
3535            predictor: None,
3536            tenant_rate_limiter: None,
3537            spill: None,
3538        };
3539
3540        // Plain text enforces: the mock serves 200.
3541        let plain =
3542            Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
3543        let resp = messages(
3544            State(state.clone()),
3545            Extension(TenantId("default".to_owned())),
3546            HeaderMap::new(),
3547            plain,
3548        )
3549        .await;
3550        assert_eq!(
3551            resp.status(),
3552            axum::http::StatusCode::OK,
3553            "plain text should enforce"
3554        );
3555
3556        // Default-on (enforce_structured = true) + verbatim-carrying ladder: tools now ENFORCE —
3557        // the mock serves 200 and the tool request never touches the bogus upstream.
3558        let tools = Bytes::from_static(
3559            br#"{"model":"m","tools":[{"name":"get_weather"}],"messages":[{"role":"user","content":"hi"}]}"#,
3560        );
3561        let resp = messages(
3562            State(state.clone()),
3563            Extension(TenantId("default".to_owned())),
3564            HeaderMap::new(),
3565            tools.clone(),
3566        )
3567        .await;
3568        assert_eq!(
3569            resp.status(),
3570            axum::http::StatusCode::OK,
3571            "tool request must route through enforce by default (ADR 0005 default-on)"
3572        );
3573
3574        // Opt-out (enforce_structured = false): the same tool request falls back to observe —
3575        // it hits the bogus upstream and is not 200.
3576        let toml_off = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n[escalation]\nenforce_structured = false\n";
3577        let config_off = ProxyConfig::from_lookup(|k| match k {
3578            "FIRSTPASS_CONFIG_TOML" => Some(toml_off.to_owned()),
3579            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
3580            "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3581            _ => None,
3582        })
3583        .unwrap();
3584        let state_off = AppState {
3585            config: Arc::new(config_off),
3586            ..state.clone()
3587        };
3588        let resp = messages(
3589            State(state_off),
3590            Extension(TenantId("default".to_owned())),
3591            HeaderMap::new(),
3592            tools,
3593        )
3594        .await;
3595        assert_ne!(
3596            resp.status(),
3597            axum::http::StatusCode::OK,
3598            "with enforce_structured = false a tool request must fall back to observe"
3599        );
3600
3601        // tool_result block in a message => same fallback.
3602        let toolres = Bytes::from_static(
3603            br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"x","content":"42"}]}]}"#,
3604        );
3605        let resp = messages(
3606            State(state),
3607            Extension(TenantId("default".to_owned())),
3608            HeaderMap::new(),
3609            toolres,
3610        )
3611        .await;
3612        assert_eq!(
3613            resp.status(),
3614            axum::http::StatusCode::OK,
3615            "tool_result blocks route through enforce by default too (verbatim carry)"
3616        );
3617    }
3618
3619    // --- Feedback API tests (drive `feedback` against a real temp trace store) ---
3620
3621    /// Persist one trace to a fresh temp DB and return (state, db_path, trace_id).
3622    async fn feedback_state() -> (AppState, std::path::PathBuf, String) {
3623        let db = std::env::temp_dir().join(format!("firstpass-feedback-{}.db", Uuid::now_v7()));
3624        let (tx, handle) = crate::store::open(&db).unwrap();
3625
3626        let mut trace = build_error_trace(
3627            &ProxyConfig::from_lookup(|_| None).unwrap(),
3628            &Bytes::from_static(b"{}"),
3629            5,
3630            Some("sess-fb"),
3631        );
3632        trace.attempts.push(Attempt {
3633            rung: 0,
3634            model: "anthropic/claude-haiku-4-5".into(),
3635            provider: "anthropic".into(),
3636            in_tokens: 10,
3637            out_tokens: 5,
3638            cost_usd: 0.001,
3639            latency_ms: 5,
3640            gates: vec![],
3641            verdict: Verdict::Pass,
3642        });
3643        let trace_id = trace.trace_id.to_string();
3644        tx.try_send(trace).unwrap();
3645        drop(tx);
3646        handle.await.unwrap();
3647
3648        let db_str = db.to_string_lossy().into_owned();
3649        let config = ProxyConfig::from_lookup(move |k| match k {
3650            "FIRSTPASS_DB" => Some(db_str.clone()),
3651            _ => None,
3652        })
3653        .unwrap();
3654        let (traces, _rx) = mpsc::channel(64);
3655        let state = AppState {
3656            config: Arc::new(config),
3657            http: reqwest::Client::new(),
3658            providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
3659            gate_health: Arc::new(GateHealthRegistry::new()),
3660            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3661            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3662            traces,
3663            adaptive: None,
3664            bandit: None,
3665            predictor: None,
3666            tenant_rate_limiter: None,
3667            spill: None,
3668        };
3669        (state, db, trace_id)
3670    }
3671
3672    #[tokio::test]
3673    async fn feedback_nudges_the_adaptive_threshold() {
3674        use firstpass_core::conformal::AdaptiveConformal;
3675        let (mut state, _db, trace_id) = feedback_state().await;
3676        let aci = Arc::new(std::sync::Mutex::new(AdaptiveConformal::new(0.1, 0.2, 0.5)));
3677        state.adaptive = Some(aci.clone());
3678        let before = aci.lock().unwrap().threshold();
3679
3680        // A served FAILURE raises the threshold (serve more conservatively).
3681        let fail = Bytes::from(
3682            serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "fail", "reporter": "ci" })
3683                .to_string(),
3684        );
3685        assert_eq!(
3686            feedback(
3687                State(state.clone()),
3688                Extension(TenantId("default".to_owned())),
3689                fail
3690            )
3691            .await
3692            .status(),
3693            axum::http::StatusCode::ACCEPTED
3694        );
3695        let after_fail = aci.lock().unwrap().threshold();
3696        assert!(
3697            after_fail > before,
3698            "served fail should raise the live threshold: {before} -> {after_fail}"
3699        );
3700
3701        // A served PASS nudges it back down — the loop is reactive both ways.
3702        let pass = Bytes::from(
3703            serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "pass", "reporter": "ci" })
3704                .to_string(),
3705        );
3706        let _ = feedback(
3707            State(state),
3708            Extension(TenantId("default".to_owned())),
3709            pass,
3710        )
3711        .await;
3712        assert!(aci.lock().unwrap().threshold() < after_fail);
3713    }
3714
3715    #[tokio::test]
3716    async fn feedback_records_a_deferred_verdict_without_breaking_the_chain() {
3717        let (state, db, trace_id) = feedback_state().await;
3718        let body = Bytes::from(
3719            serde_json::json!({
3720                "trace_id": trace_id,
3721                "gate_id": "tests",
3722                "verdict": "pass",
3723                "score": 1.0,
3724                "reporter": "ci",
3725            })
3726            .to_string(),
3727        );
3728        let resp = feedback(
3729            State(state),
3730            Extension(TenantId("default".to_owned())),
3731            body,
3732        )
3733        .await;
3734        assert_eq!(resp.status(), axum::http::StatusCode::ACCEPTED);
3735
3736        // The deferred verdict is visible on the trace view...
3737        let view = crate::store::load_trace_view(&db, "default", &trace_id)
3738            .unwrap()
3739            .unwrap();
3740        assert_eq!(view.deferred.len(), 1);
3741        assert_eq!(view.deferred[0].gate_id, "tests");
3742        // ...and the sealed chain still verifies (the outcome didn't mutate the trace).
3743        let traces = crate::store::load_all_traces(&db).unwrap();
3744        firstpass_core::verify_chain(&traces, GENESIS_HASH).unwrap();
3745
3746        let _ = std::fs::remove_file(&db);
3747    }
3748
3749    #[tokio::test]
3750    async fn feedback_for_unknown_trace_is_404() {
3751        let (state, db, _trace_id) = feedback_state().await;
3752        let body = Bytes::from(
3753            serde_json::json!({
3754                "trace_id": "does-not-exist",
3755                "gate_id": "tests",
3756                "verdict": "pass",
3757                "reporter": "ci",
3758            })
3759            .to_string(),
3760        );
3761        let resp = feedback(
3762            State(state),
3763            Extension(TenantId("default".to_owned())),
3764            body,
3765        )
3766        .await;
3767        assert_eq!(resp.status(), axum::http::StatusCode::NOT_FOUND);
3768        let _ = std::fs::remove_file(&db);
3769    }
3770
3771    /// D4 IDOR: a real trace owned by "default" cannot receive feedback from another tenant. The
3772    /// caller gets a `404` (not `403`), so there is no existence oracle across the boundary.
3773    #[tokio::test]
3774    async fn feedback_across_tenants_is_404_not_403() {
3775        let (state, db, trace_id) = feedback_state().await;
3776        let body = Bytes::from(
3777            serde_json::json!({
3778                "trace_id": trace_id,
3779                "gate_id": "tests",
3780                "verdict": "pass",
3781                "score": 1.0,
3782                "reporter": "attacker",
3783            })
3784            .to_string(),
3785        );
3786        // Caller authenticated as a *different* tenant than the trace's owner.
3787        let resp = feedback(
3788            State(state),
3789            Extension(TenantId("tenant-b".to_owned())),
3790            body,
3791        )
3792        .await;
3793        assert_eq!(
3794            resp.status(),
3795            axum::http::StatusCode::NOT_FOUND,
3796            "cross-tenant feedback must look exactly like a missing trace"
3797        );
3798        let _ = std::fs::remove_file(&db);
3799    }
3800
3801    #[tokio::test]
3802    async fn feedback_rejects_bad_verdict_and_score() {
3803        let (state, db, trace_id) = feedback_state().await;
3804        let bad_verdict = Bytes::from(
3805            serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "maybe", "reporter": "x" })
3806                .to_string(),
3807        );
3808        assert_eq!(
3809            feedback(
3810                State(state.clone()),
3811                Extension(TenantId("default".to_owned())),
3812                bad_verdict
3813            )
3814            .await
3815            .status(),
3816            axum::http::StatusCode::BAD_REQUEST
3817        );
3818        let bad_score = Bytes::from(
3819            serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "pass", "score": 9.0, "reporter": "x" })
3820                .to_string(),
3821        );
3822        assert_eq!(
3823            feedback(
3824                State(state),
3825                Extension(TenantId("default".to_owned())),
3826                bad_score
3827            )
3828            .await
3829            .status(),
3830            axum::http::StatusCode::BAD_REQUEST
3831        );
3832        let _ = std::fs::remove_file(&db);
3833    }
3834
3835    #[tokio::test]
3836    async fn metrics_endpoint_renders_after_a_real_request() {
3837        use tower::ServiceExt;
3838
3839        let (state, mut rx) = enforce_state(
3840            &["anthropic/claude-haiku-4-5"],
3841            &["non-empty"],
3842            vec![(
3843                "anthropic/claude-haiku-4-5",
3844                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3845            )],
3846        );
3847        let router = app(state).expect("prometheus recorder installs");
3848
3849        let req = axum::http::Request::builder()
3850            .method("POST")
3851            .uri("/v1/messages")
3852            .header("content-type", "application/json")
3853            .body(Body::from(user_body()))
3854            .unwrap();
3855        let resp = router.clone().oneshot(req).await.unwrap();
3856        assert_eq!(resp.status(), axum::http::StatusCode::OK);
3857        rx.try_recv().expect("a trace was enqueued");
3858
3859        let metrics_req = axum::http::Request::builder()
3860            .method("GET")
3861            .uri("/metrics")
3862            .body(Body::empty())
3863            .unwrap();
3864        let metrics_resp = router.oneshot(metrics_req).await.unwrap();
3865        assert_eq!(metrics_resp.status(), axum::http::StatusCode::OK);
3866        let bytes = axum::body::to_bytes(metrics_resp.into_body(), 1 << 20)
3867            .await
3868            .unwrap();
3869        let body = String::from_utf8(bytes.to_vec()).unwrap();
3870        assert!(
3871            body.contains("firstpass_enforce_latency_ms"),
3872            "metrics body missing enforce latency histogram: {body}"
3873        );
3874        assert!(
3875            body.contains("firstpass_served_total"),
3876            "metrics body missing served counter: {body}"
3877        );
3878    }
3879
3880    // --- Multi-tenant auth (ADR 0004 §D1) integration tests, driven through the real router ---
3881
3882    /// Build an `AppState` whose config toggles auth and (optionally) carries a tenant-keys JSON.
3883    fn auth_state(require_auth: bool, keys_json: Option<String>) -> AppState {
3884        auth_state_rated(require_auth, keys_json, None)
3885    }
3886
3887    /// Like [`auth_state`], but also wires `FIRSTPASS_TENANT_RATE_PER_SEC` (ADR 0004 §D6) when
3888    /// `rate_per_sec` is `Some`.
3889    fn auth_state_rated(
3890        require_auth: bool,
3891        keys_json: Option<String>,
3892        rate_per_sec: Option<u32>,
3893    ) -> AppState {
3894        let config = ProxyConfig::from_lookup(|k| match k {
3895            "FIRSTPASS_REQUIRE_AUTH" => require_auth.then(|| "true".to_owned()),
3896            "FIRSTPASS_TENANT_KEYS_JSON" => keys_json.clone(),
3897            "FIRSTPASS_TENANT_RATE_PER_SEC" => rate_per_sec.map(|n| n.to_string()),
3898            _ => None,
3899        })
3900        .unwrap();
3901        let (traces, _rx) = mpsc::channel(64);
3902        // Deliberately leak the receiver for the test's lifetime so the sender never reports the
3903        // channel closed (the auth tests exercise `/v1/capabilities`, which enqueues no trace).
3904        std::mem::forget(_rx);
3905        let providers: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3906        let tenant_rate_limiter = build_tenant_rate_limiter(&config);
3907        AppState {
3908            config: Arc::new(config),
3909            http: reqwest::Client::new(),
3910            providers: ProviderRegistry::from_map(providers),
3911            gate_health: Arc::new(GateHealthRegistry::new()),
3912            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3913            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3914            traces,
3915            adaptive: None,
3916            bandit: None,
3917            predictor: None,
3918            tenant_rate_limiter,
3919            spill: None,
3920        }
3921    }
3922
3923    fn cap_request(auth_header: Option<&str>) -> axum::http::Request<Body> {
3924        let mut b = axum::http::Request::builder()
3925            .method("GET")
3926            .uri("/v1/capabilities");
3927        if let Some(h) = auth_header {
3928            b = b.header("authorization", h);
3929        }
3930        b.body(Body::empty()).unwrap()
3931    }
3932
3933    #[tokio::test]
3934    async fn auth_off_allows_unauthenticated_request() {
3935        use tower::ServiceExt;
3936        let router = app(auth_state(false, None)).expect("router");
3937        let resp = router.oneshot(cap_request(None)).await.unwrap();
3938        // Default-off: no key required, request proceeds to the handler.
3939        assert_eq!(resp.status(), axum::http::StatusCode::OK);
3940    }
3941
3942    #[tokio::test]
3943    async fn auth_on_missing_key_is_401_opaque() {
3944        use tower::ServiceExt;
3945        let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3946        let keys = format!("{{\"tenant-a\": {hash:?}}}");
3947        let router = app(auth_state(true, Some(keys))).expect("router");
3948
3949        let resp = router.oneshot(cap_request(None)).await.unwrap();
3950        assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
3951        let json = body_json(resp).await;
3952        assert_eq!(json["error"]["type"], "unauthorized");
3953        // Opaque: the body must not name tenants or hint which key would work.
3954        let msg = json["error"]["message"].as_str().unwrap();
3955        assert!(!msg.contains("tenant"), "no tenant oracle in body: {msg}");
3956    }
3957
3958    #[tokio::test]
3959    async fn auth_on_invalid_key_is_401() {
3960        use tower::ServiceExt;
3961        let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3962        let keys = format!("{{\"tenant-a\": {hash:?}}}");
3963        let router = app(auth_state(true, Some(keys))).expect("router");
3964
3965        let resp = router
3966            .oneshot(cap_request(Some("Bearer wrong-key")))
3967            .await
3968            .unwrap();
3969        assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
3970    }
3971
3972    #[tokio::test]
3973    async fn auth_on_valid_key_proceeds() {
3974        use tower::ServiceExt;
3975        let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3976        let keys = format!("{{\"tenant-a\": {hash:?}}}");
3977        let router = app(auth_state(true, Some(keys))).expect("router");
3978
3979        // Keyed format: `<tenant_id>.<secret>`.
3980        let resp = router
3981            .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
3982            .await
3983            .unwrap();
3984        // A valid key clears the middleware and reaches the handler.
3985        assert_eq!(resp.status(), axum::http::StatusCode::OK);
3986    }
3987
3988    /// Two tenants, keyed so requests carry a real, distinct `TenantId` (ADR 0004 §D6).
3989    fn two_tenant_state(rate_per_sec: Option<u32>) -> AppState {
3990        let hash_a = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3991        let hash_b = crate::tenant_auth::TenantKeys::hash_key("key-b").unwrap();
3992        let keys = format!("{{\"tenant-a\": {hash_a:?}, \"tenant-b\": {hash_b:?}}}");
3993        auth_state_rated(true, Some(keys), rate_per_sec)
3994    }
3995
3996    #[tokio::test]
3997    async fn tenant_exceeding_rate_limit_gets_429_opaque() {
3998        use tower::ServiceExt;
3999        // Burst capacity == rate for `Quota::per_second`, so 1 req/sec allows one request through
4000        // and rejects the rest of a burst. The requests are fired CONCURRENTLY so they reach the
4001        // limiter in one tight window — a serial sequence is wall-clock sensitive (each request
4002        // pays a deliberately slow Argon2 verify, and on a loaded CI runner >1s can elapse between
4003        // limiter checks, refilling the bucket and turning the expected 429 into a legit 200).
4004        let router = app(two_tenant_state(Some(1))).expect("router");
4005
4006        let (r1, r2, r3, r4) = tokio::join!(
4007            router
4008                .clone()
4009                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4010            router
4011                .clone()
4012                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4013            router
4014                .clone()
4015                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4016            router
4017                .clone()
4018                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4019        );
4020        let responses = [r1.unwrap(), r2.unwrap(), r3.unwrap(), r4.unwrap()];
4021        let ok = responses
4022            .iter()
4023            .filter(|r| r.status() == axum::http::StatusCode::OK)
4024            .count();
4025        assert!(ok >= 1, "the burst's first request must pass");
4026        let limited: Vec<_> = responses
4027            .into_iter()
4028            .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
4029            .collect();
4030        assert!(
4031            !limited.is_empty(),
4032            "a 4-request burst against 1 req/sec must trip the limiter"
4033        );
4034
4035        let json = body_json(limited.into_iter().next().unwrap()).await;
4036        assert_eq!(json["error"]["type"], "rate_limited");
4037        // Opaque: no bucket state or limit value leaked to the caller.
4038        let msg = json["error"]["message"].as_str().unwrap();
4039        assert!(!msg.contains('1'), "no limit value in body: {msg}");
4040    }
4041
4042    #[tokio::test]
4043    async fn rate_limit_buckets_are_independent_per_tenant() {
4044        use tower::ServiceExt;
4045        let router = app(two_tenant_state(Some(1))).expect("router");
4046
4047        // Tenant A bursts past its 1 req/sec budget (concurrent — see the sibling test for why a
4048        // serial sequence would be wall-clock flaky)...
4049        let (a1, a2, a3) = tokio::join!(
4050            router
4051                .clone()
4052                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4053            router
4054                .clone()
4055                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4056            router
4057                .clone()
4058                .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4059        );
4060        let a_limited = [a1.unwrap(), a2.unwrap(), a3.unwrap()]
4061            .iter()
4062            .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
4063            .count();
4064        assert!(a_limited >= 1, "tenant A's burst must trip its limiter");
4065
4066        // ...but tenant B, on the same gate/route, is unaffected (independent bucket).
4067        let b1 = router
4068            .clone()
4069            .oneshot(cap_request(Some("Bearer tenant-b.key-b")))
4070            .await
4071            .unwrap();
4072        assert_eq!(b1.status(), axum::http::StatusCode::OK);
4073    }
4074
4075    #[tokio::test]
4076    async fn rate_limit_unset_never_429s() {
4077        use tower::ServiceExt;
4078        // Backward-compat: with no FIRSTPASS_TENANT_RATE_PER_SEC, drive many requests through and
4079        // confirm none are ever rate-limited (default-off).
4080        let router = app(two_tenant_state(None)).expect("router");
4081        for _ in 0..20 {
4082            let resp = router
4083                .clone()
4084                .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
4085                .await
4086                .unwrap();
4087            assert_eq!(resp.status(), axum::http::StatusCode::OK);
4088        }
4089    }
4090
4091    // ── epsilon-greedy helper unit tests ──────────────────────────────────────
4092
4093    #[test]
4094    fn u01_is_deterministic_and_in_range() {
4095        let s1 = u01(0xDEAD_BEEF_CAFE_1234_u128);
4096        let s2 = u01(0xDEAD_BEEF_CAFE_1234_u128);
4097        assert_eq!(s1, s2, "u01 must be deterministic for the same seed");
4098        assert!((0.0..1.0).contains(&s1), "u01 must return [0, 1), got {s1}");
4099
4100        // Different seeds produce different values (highly likely with SM64).
4101        let s3 = u01(0x1234_5678_9ABC_DEF0_u128);
4102        assert_ne!(s1, s3, "different seeds should give different values");
4103
4104        // Verify range across a spread of seeds.
4105        for i in 0u64..256 {
4106            let v = u01(i as u128);
4107            assert!((0.0..1.0).contains(&v), "seed {i}: u01={v} out of [0,1)");
4108        }
4109    }
4110
4111    #[test]
4112    fn epsilon_propensity_formula() {
4113        let epsilon = 0.2_f64;
4114        let k = 3_usize;
4115        let greedy = 1_u32;
4116
4117        // Greedy rung chosen: both (1-ε) and ε/K terms apply.
4118        let p_greedy = epsilon_propensity(greedy, greedy, epsilon, k);
4119        let expected_greedy = (1.0 - epsilon) + epsilon / k as f64;
4120        assert!(
4121            (p_greedy - expected_greedy).abs() < 1e-12,
4122            "{p_greedy} != {expected_greedy}"
4123        );
4124
4125        // Non-greedy rung: only ε/K term.
4126        let p_other = epsilon_propensity(0, greedy, epsilon, k);
4127        let expected_other = epsilon / k as f64;
4128        assert!(
4129            (p_other - expected_other).abs() < 1e-12,
4130            "{p_other} != {expected_other}"
4131        );
4132
4133        // All propensities are in (0, 1].
4134        for chosen in 0..k as u32 {
4135            let p = epsilon_propensity(chosen, greedy, epsilon, k);
4136            assert!(
4137                p > 0.0 && p <= 1.0,
4138                "propensity {p} out of (0,1] for chosen={chosen}"
4139            );
4140        }
4141    }
4142
4143    #[test]
4144    fn epsilon_branch_and_greedy_branch_both_occur_over_many_seeds() {
4145        // With epsilon=0.3 over 200 sequential seeds, both branches must fire.
4146        let epsilon = 0.3_f64;
4147        let mut saw_explore = false;
4148        let mut saw_greedy = false;
4149        for i in 0u64..200 {
4150            let u = u01(i as u128);
4151            if u < epsilon {
4152                saw_explore = true;
4153            } else {
4154                saw_greedy = true;
4155            }
4156            if saw_explore && saw_greedy {
4157                break;
4158            }
4159        }
4160        assert!(
4161            saw_explore,
4162            "epsilon branch must fire with epsilon=0.3 over 200 seeds"
4163        );
4164        assert!(
4165            saw_greedy,
4166            "greedy branch must occur with epsilon=0.3 over 200 seeds"
4167        );
4168    }
4169
4170    #[test]
4171    fn epsilon_propensity_sums_to_one_over_all_rungs() {
4172        // Sum of propensities over all K rungs equals 1 (it is a valid probability distribution).
4173        let epsilon = 0.15_f64;
4174        let k = 4_usize;
4175        let greedy = 2_u32;
4176        let total: f64 = (0..k as u32)
4177            .map(|r| epsilon_propensity(r, greedy, epsilon, k))
4178            .sum();
4179        assert!(
4180            (total - 1.0).abs() < 1e-12,
4181            "propensities must sum to 1, got {total}"
4182        );
4183    }
4184
4185    /// Poll the hand-rolled keepalive stream directly under paused tokio time: one keepalive
4186    /// per idle interval while the pipeline runs, then the final SSE frame, then end-of-stream.
4187    #[tokio::test(start_paused = true)]
4188    async fn keepalive_stream_ticks_then_emits_final_frame() {
4189        let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
4190        let mut ticks = tokio::time::interval(SSE_KEEPALIVE_EVERY);
4191        ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
4192        ticks.reset();
4193        let mut stream = KeepaliveStream {
4194            rx: Some(rx),
4195            ticks,
4196            format_message: anthropic_sse_from_message,
4197        };
4198        /// One poll of the stream: `Some(item)` if ready, `None` if pending right now.
4199        async fn next(
4200            stream: &mut KeepaliveStream,
4201        ) -> Option<Option<Result<Bytes, std::convert::Infallible>>> {
4202            std::future::poll_fn(|cx| {
4203                std::task::Poll::Ready(
4204                    match futures_core::Stream::poll_next(std::pin::Pin::new(&mut *stream), cx) {
4205                        std::task::Poll::Ready(item) => Some(item),
4206                        std::task::Poll::Pending => None,
4207                    },
4208                )
4209            })
4210            .await
4211        }
4212
4213        // Pipeline still running, no interval elapsed: nothing to emit yet.
4214        assert!(
4215            next(&mut stream).await.is_none(),
4216            "no frame before an interval"
4217        );
4218
4219        // Advance past one keepalive interval: a comment frame is emitted.
4220        tokio::time::advance(SSE_KEEPALIVE_EVERY + Duration::from_millis(1)).await;
4221        let frame = next(&mut stream)
4222            .await
4223            .expect("keepalive due")
4224            .unwrap()
4225            .unwrap();
4226        assert!(
4227            frame.starts_with(b": "),
4228            "keepalive must be an SSE comment (ignored by every conforming parser)"
4229        );
4230
4231        // Pipeline resolves: the final frame is the full Anthropic event sequence, then EOS.
4232        let message = serde_json::json!({
4233            "id": "msg_1", "type": "message", "role": "assistant", "model": "m",
4234            "content": [{ "type": "text", "text": "done" }],
4235            "usage": { "input_tokens": 1, "output_tokens": 1 }
4236        });
4237        tx.send(Ok(message)).unwrap();
4238        let frame = next(&mut stream)
4239            .await
4240            .expect("final frame")
4241            .unwrap()
4242            .unwrap();
4243        let text = String::from_utf8(frame.to_vec()).unwrap();
4244        assert!(text.contains("event: message_start"));
4245        assert!(text.contains("event: message_stop"));
4246        let eos = next(&mut stream)
4247            .await
4248            .expect("stream must end after the final frame");
4249        assert!(eos.is_none(), "end-of-stream after the final frame");
4250    }
4251
4252    /// E2E: a `stream: true` enforce request (default-on structured) is answered 200
4253    /// `text/event-stream` whose body carries the full gated event sequence.
4254    #[tokio::test]
4255    async fn streaming_enforce_serves_full_sse_sequence() {
4256        let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
4257        let config = ProxyConfig::from_lookup(|k| match k {
4258            "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
4259            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
4260            "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
4261            _ => None,
4262        })
4263        .unwrap();
4264        let mut outs = HashMap::new();
4265        outs.insert(
4266            "anthropic/m".to_owned(),
4267            Ok(model_resp("anthropic/m", "gated answer")),
4268        );
4269        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
4270        map.insert(
4271            "anthropic".to_owned(),
4272            Arc::new(MockProvider::new("anthropic", outs)),
4273        );
4274        let (traces, _rx) = mpsc::channel(64);
4275        let state = AppState {
4276            config: Arc::new(config),
4277            http: reqwest::Client::new(),
4278            providers: ProviderRegistry::from_map(map),
4279            gate_health: Arc::new(GateHealthRegistry::new()),
4280            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
4281            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
4282            traces,
4283            adaptive: None,
4284            bandit: None,
4285            predictor: None,
4286            tenant_rate_limiter: None,
4287            spill: None,
4288        };
4289        let body = Bytes::from_static(
4290            br#"{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}"#,
4291        );
4292        let resp = messages(
4293            State(state),
4294            Extension(TenantId("default".to_owned())),
4295            HeaderMap::new(),
4296            body,
4297        )
4298        .await;
4299        assert_eq!(resp.status(), axum::http::StatusCode::OK);
4300        assert!(
4301            resp.headers()
4302                .get(axum::http::header::CONTENT_TYPE)
4303                .and_then(|v| v.to_str().ok())
4304                .is_some_and(|ct| ct.starts_with("text/event-stream")),
4305            "streaming client must get SSE"
4306        );
4307        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
4308            .await
4309            .unwrap();
4310        let text = String::from_utf8(bytes.to_vec()).unwrap();
4311        assert!(text.contains("event: message_start"));
4312        assert!(text.contains("gated answer"));
4313        assert!(text.contains("event: message_stop"));
4314    }
4315
4316    // ── OpenAI inbound (SPEC §M1) ──────────────────────────────────────────────
4317
4318    // --- Golden translation tests ---
4319
4320    #[test]
4321    fn parse_openai_request_plain_text() {
4322        // Simple user message → normalized ModelRequest (translation path, carry_raw=false)
4323        let body = br#"{"model":"gpt-4o","max_tokens":256,"messages":[{"role":"user","content":"hello"}]}"#;
4324        let req = parse_openai_request(body, false).expect("must parse");
4325        assert_eq!(req.model, "gpt-4o");
4326        assert_eq!(req.max_tokens, 256);
4327        assert_eq!(req.messages.len(), 1);
4328        assert_eq!(req.messages[0].role, "user");
4329        assert_eq!(req.messages[0].content, Value::String("hello".to_owned()));
4330        assert!(req.system.is_none());
4331        // Translation path: raw must be Null so anthropic_wire_body rebuilds from fields
4332        assert_eq!(req.raw, Value::Null);
4333    }
4334
4335    #[test]
4336    fn parse_openai_request_system_message() {
4337        let body = br#"{"model":"gpt-4o","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hi"}]}"#;
4338        let req = parse_openai_request(body, false).expect("must parse");
4339        assert_eq!(req.system.as_deref(), Some("be concise"));
4340        assert_eq!(req.messages.len(), 1);
4341        assert_eq!(req.messages[0].role, "user");
4342    }
4343
4344    #[test]
4345    fn parse_openai_request_tool_calls_translate_to_tool_use() {
4346        let body = br#"{
4347            "model":"gpt-4o",
4348            "messages":[
4349                {"role":"user","content":"what's the weather?"},
4350                {"role":"assistant","content":null,"tool_calls":[
4351                    {"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Paris\"}"}}
4352                ]},
4353                {"role":"tool","tool_call_id":"call_1","content":"15C, cloudy"}
4354            ]
4355        }"#;
4356        let req = parse_openai_request(body, false).expect("must parse");
4357        // 3 messages → user + assistant (tool_use blocks) + user (tool_result)
4358        assert_eq!(req.messages.len(), 3);
4359        // Assistant turn: Anthropic tool_use block
4360        let asst = &req.messages[1];
4361        assert_eq!(asst.role, "assistant");
4362        let blocks = asst.content.as_array().expect("content array");
4363        assert_eq!(blocks[0]["type"], "tool_use");
4364        assert_eq!(blocks[0]["name"], "get_weather");
4365        assert_eq!(blocks[0]["id"], "call_1");
4366        assert_eq!(blocks[0]["input"]["city"], "Paris");
4367        // Tool result turn: role becomes "user", tool_result block
4368        let tool_msg = &req.messages[2];
4369        assert_eq!(tool_msg.role, "user");
4370        let result_blocks = tool_msg.content.as_array().expect("result blocks");
4371        assert_eq!(result_blocks[0]["type"], "tool_result");
4372        assert_eq!(result_blocks[0]["tool_use_id"], "call_1");
4373    }
4374
4375    #[test]
4376    fn parse_openai_request_tools_translate_to_anthropic_format() {
4377        let body = br#"{
4378            "model":"gpt-4o",
4379            "messages":[{"role":"user","content":"use a tool"}],
4380            "tools":[{"type":"function","function":{"name":"search","description":"web search","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}}]
4381        }"#;
4382        let req = parse_openai_request(body, false).expect("must parse");
4383        let tools = req.tools.as_array().expect("tools array");
4384        assert_eq!(tools.len(), 1);
4385        assert_eq!(tools[0]["name"], "search");
4386        assert_eq!(tools[0]["description"], "web search");
4387        assert_eq!(tools[0]["input_schema"]["type"], "object");
4388    }
4389
4390    #[test]
4391    fn parse_openai_request_raw_carry_preserves_full_body() {
4392        // All-OpenAI ladder path: raw = original JSON, no translation
4393        let body =
4394            br#"{"model":"gpt-4o","max_tokens":100,"messages":[{"role":"user","content":"hi"}]}"#;
4395        let req = parse_openai_request(body, true).expect("must parse");
4396        assert!(req.raw.is_object(), "raw must be the full JSON object");
4397        assert_eq!(req.raw["model"], "gpt-4o");
4398        assert_eq!(req.raw["max_tokens"], 100);
4399        // Tools remain in OpenAI shape (not translated) on the raw-carry path
4400        assert!(req.tools.is_null(), "no tools in this request");
4401    }
4402
4403    #[test]
4404    fn parse_openai_request_http_image_returns_none() {
4405        // Non-translatable: http image URL → caller must use observe fallback
4406        let body = br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
4407            {"type":"text","text":"describe this"},
4408            {"type":"image_url","image_url":{"url":"https://example.com/cat.png"}}
4409        ]}]}"#;
4410        let result = parse_openai_request(body, false);
4411        assert!(result.is_none(), "http image URL must fail translation");
4412    }
4413
4414    #[test]
4415    fn parse_openai_request_data_url_image_translates_to_anthropic_base64() {
4416        let body = br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
4417            {"type":"text","text":"describe"},
4418            {"type":"image_url","image_url":{"url":"data:image/png;base64,iVBORw0KGgo="}}
4419        ]}]}"#;
4420        let req = parse_openai_request(body, false).expect("data URL must parse");
4421        let blocks = req.messages[0].content.as_array().expect("blocks");
4422        let img = blocks
4423            .iter()
4424            .find(|b| b["type"] == "image")
4425            .expect("image block");
4426        assert_eq!(img["source"]["type"], "base64");
4427        assert_eq!(img["source"]["media_type"], "image/png");
4428        assert_eq!(img["source"]["data"], "iVBORw0KGgo=");
4429    }
4430
4431    // --- Response rendering tests ---
4432
4433    #[test]
4434    fn openai_response_json_renders_text_response() {
4435        let resp = ModelResponse {
4436            model: "gpt-4o".to_owned(),
4437            text: "Hello!".to_owned(),
4438            in_tokens: 10,
4439            out_tokens: 5,
4440            raw: serde_json::json!({
4441                "content": [{ "type": "text", "text": "Hello!" }]
4442            }),
4443        };
4444        let json = openai_response_json(&resp);
4445        assert_eq!(json["object"], "chat.completion");
4446        assert_eq!(json["model"], "gpt-4o");
4447        assert_eq!(json["choices"][0]["message"]["role"], "assistant");
4448        assert_eq!(json["choices"][0]["message"]["content"], "Hello!");
4449        assert_eq!(json["choices"][0]["finish_reason"], "stop");
4450        assert_eq!(json["usage"]["prompt_tokens"], 10);
4451        assert_eq!(json["usage"]["completion_tokens"], 5);
4452    }
4453
4454    #[test]
4455    fn openai_response_json_renders_tool_call() {
4456        let resp = ModelResponse {
4457            model: "gpt-4o".to_owned(),
4458            text: String::new(),
4459            in_tokens: 20,
4460            out_tokens: 15,
4461            raw: serde_json::json!({
4462                "content": [{
4463                    "type": "tool_use",
4464                    "id": "call_abc",
4465                    "name": "search",
4466                    "input": {"q": "Rust async"}
4467                }]
4468            }),
4469        };
4470        let json = openai_response_json(&resp);
4471        assert_eq!(json["choices"][0]["finish_reason"], "tool_calls");
4472        let tc = &json["choices"][0]["message"]["tool_calls"][0];
4473        assert_eq!(tc["id"], "call_abc");
4474        assert_eq!(tc["type"], "function");
4475        assert_eq!(tc["function"]["name"], "search");
4476        // content is null for tool-only responses (OpenAI spec)
4477        assert_eq!(json["choices"][0]["message"]["content"], Value::Null);
4478    }
4479
4480    #[test]
4481    fn openai_sse_from_message_plain_text_has_role_then_content_then_stop() {
4482        let resp = ModelResponse {
4483            model: "gpt-4o".to_owned(),
4484            text: "Hi there!".to_owned(),
4485            in_tokens: 5,
4486            out_tokens: 3,
4487            raw: serde_json::json!({
4488                "content": [{ "type": "text", "text": "Hi there!" }]
4489            }),
4490        };
4491        let sse = openai_sse_from_message(&openai_response_json(&resp));
4492        // Every line is either blank or starts with "data: "
4493        for line in sse.lines() {
4494            assert!(
4495                line.is_empty() || line.starts_with("data: "),
4496                "bad SSE line: {line:?}"
4497            );
4498        }
4499        let frames: Vec<&str> = sse
4500            .lines()
4501            .filter_map(|l| l.strip_prefix("data: "))
4502            .collect();
4503        // Last frame must be [DONE]
4504        assert_eq!(*frames.last().unwrap(), "[DONE]");
4505        // Role delta must appear
4506        let role_frame: Value = serde_json::from_str(frames[0]).unwrap();
4507        assert_eq!(role_frame["choices"][0]["delta"]["role"], "assistant");
4508        // Content delta must appear somewhere
4509        assert!(frames.iter().any(|f| {
4510            if *f == "[DONE]" {
4511                return false;
4512            }
4513            serde_json::from_str::<Value>(f)
4514                .ok()
4515                .is_some_and(|v| v["choices"][0]["delta"]["content"] == "Hi there!")
4516        }));
4517        // Finish reason must appear in a non-[DONE] frame
4518        assert!(frames.iter().any(|f| {
4519            if *f == "[DONE]" {
4520                return false;
4521            }
4522            serde_json::from_str::<Value>(f)
4523                .ok()
4524                .is_some_and(|v| v["choices"][0]["finish_reason"] == "stop")
4525        }));
4526    }
4527
4528    // --- Detection helper tests ---
4529
4530    #[test]
4531    fn detects_openai_tool_calls() {
4532        let with_tool_calls = Bytes::from_static(br#"{"messages":[
4533            {"role":"user","content":"hi"},
4534            {"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"f","arguments":"{}"}}]}
4535        ]}"#);
4536        let without = Bytes::from_static(br#"{"messages":[{"role":"user","content":"hi"}]}"#);
4537        let with_tool_msg = Bytes::from_static(
4538            br#"{"messages":[
4539            {"role":"tool","tool_call_id":"c1","content":"result"}
4540        ]}"#,
4541        );
4542        assert!(openai_messages_have_tool_calls(&with_tool_calls));
4543        assert!(!openai_messages_have_tool_calls(&without));
4544        assert!(openai_messages_have_tool_calls(&with_tool_msg));
4545    }
4546
4547    #[test]
4548    fn detects_openai_http_images() {
4549        let http_img = Bytes::from_static(
4550            br#"{"messages":[{"role":"user","content":[
4551            {"type":"image_url","image_url":{"url":"https://example.com/img.png"}}
4552        ]}]}"#,
4553        );
4554        let data_img = Bytes::from_static(
4555            br#"{"messages":[{"role":"user","content":[
4556            {"type":"image_url","image_url":{"url":"data:image/png;base64,abc"}}
4557        ]}]}"#,
4558        );
4559        let no_img = Bytes::from_static(br#"{"messages":[{"role":"user","content":"hi"}]}"#);
4560        assert!(openai_has_http_images(&http_img));
4561        assert!(!openai_has_http_images(&data_img));
4562        assert!(!openai_has_http_images(&no_img));
4563    }
4564
4565    #[test]
4566    fn enforce_can_handle_openai_inbound_all_openai_ladder() {
4567        // All-OpenAI ladder: verbatim carry, no translation needed → enforce allowed.
4568        let tools_body = Bytes::from_static(br#"{"model":"gpt-4o","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"c","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#);
4569        let f = extract_openai_features(&HeaderMap::new(), &tools_body);
4570        let ladder = vec!["openai/gpt-4o-mini".to_owned(), "openai/gpt-4o".to_owned()];
4571        let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
4572        assert!(enforce_can_handle(
4573            &f,
4574            &tools_body,
4575            true,
4576            &ladder,
4577            &providers,
4578            Dialect::Openai
4579        ));
4580    }
4581
4582    #[test]
4583    fn enforce_can_handle_openai_inbound_all_anthropic_ladder_no_http_image() {
4584        // Translation path: OpenAI inbound + all-Anthropic ladder + no http images → allowed.
4585        let tools_body = Bytes::from_static(br#"{"model":"gpt-4o","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"c","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#);
4586        let f = extract_openai_features(&HeaderMap::new(), &tools_body);
4587        let ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
4588        let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
4589        assert!(enforce_can_handle(
4590            &f,
4591            &tools_body,
4592            true,
4593            &ladder,
4594            &providers,
4595            Dialect::Openai,
4596        ));
4597    }
4598
4599    #[test]
4600    fn enforce_can_handle_openai_inbound_http_image_falls_back() {
4601        // Non-translatable: http image URL → enforce not possible, observe fallback.
4602        let img_body = Bytes::from_static(
4603            br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
4604            {"type":"image_url","image_url":{"url":"https://example.com/img.png"}}
4605        ]}]}"#,
4606        );
4607        let f = extract_openai_features(&HeaderMap::new(), &img_body);
4608        let ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
4609        let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
4610        assert!(!enforce_can_handle(
4611            &f,
4612            &img_body,
4613            true,
4614            &ladder,
4615            &providers,
4616            Dialect::Openai,
4617        ));
4618    }
4619
4620    // --- E2E handler tests ---
4621
4622    /// Build a minimal enforce AppState backed by MockProvider for OpenAI-inbound tests.
4623    fn openai_enforce_state(mock_resp: ModelResponse) -> AppState {
4624        let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"mock/m\"]\ngates = [\"non-empty\"]\n";
4625        let config = ProxyConfig::from_lookup(|k| match k {
4626            "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
4627            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
4628            _ => None,
4629        })
4630        .unwrap();
4631        let mut outs = HashMap::new();
4632        outs.insert("mock/m".to_owned(), Ok(mock_resp));
4633        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
4634        map.insert("mock".to_owned(), Arc::new(MockProvider::new("mock", outs)));
4635        let (traces, _rx) = mpsc::channel(64);
4636        std::mem::forget(_rx);
4637        let tenant_rate_limiter = build_tenant_rate_limiter(&config);
4638        AppState {
4639            config: Arc::new(config),
4640            http: reqwest::Client::new(),
4641            providers: ProviderRegistry::from_map(map),
4642            gate_health: Arc::new(GateHealthRegistry::new()),
4643            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
4644            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
4645            traces,
4646            adaptive: None,
4647            bandit: None,
4648            predictor: None,
4649            tenant_rate_limiter,
4650            spill: None,
4651        }
4652    }
4653
4654    #[tokio::test]
4655    async fn chat_completions_plain_text_enforce_returns_openai_shape() {
4656        let mock = model_resp("mock/m", "gated answer");
4657        let state = openai_enforce_state(mock);
4658        let body = Bytes::from_static(
4659            br#"{"model":"gpt-4o","messages":[{"role":"user","content":"hello"}]}"#,
4660        );
4661        let resp = chat_completions(
4662            State(state),
4663            Extension(TenantId("default".to_owned())),
4664            HeaderMap::new(),
4665            body,
4666        )
4667        .await;
4668        assert_eq!(resp.status(), axum::http::StatusCode::OK);
4669        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
4670            .await
4671            .unwrap();
4672        let json: Value = serde_json::from_slice(&bytes).expect("must be JSON");
4673        assert_eq!(json["object"], "chat.completion");
4674        assert_eq!(json["choices"][0]["message"]["role"], "assistant");
4675        assert_eq!(json["choices"][0]["message"]["content"], "gated answer");
4676        assert_eq!(json["choices"][0]["finish_reason"], "stop");
4677        assert!(
4678            json["id"]
4679                .as_str()
4680                .is_some_and(|id| id.starts_with("chatcmpl-"))
4681        );
4682    }
4683
4684    #[tokio::test]
4685    async fn chat_completions_stream_true_returns_sse_with_openai_chunks() {
4686        let mock = model_resp("mock/m", "gated answer");
4687        let state = openai_enforce_state(mock);
4688        let body = Bytes::from_static(
4689            br#"{"model":"gpt-4o","stream":true,"messages":[{"role":"user","content":"hello"}]}"#,
4690        );
4691        let resp = chat_completions(
4692            State(state),
4693            Extension(TenantId("default".to_owned())),
4694            HeaderMap::new(),
4695            body,
4696        )
4697        .await;
4698        assert_eq!(resp.status(), axum::http::StatusCode::OK);
4699        assert!(
4700            resp.headers()
4701                .get(axum::http::header::CONTENT_TYPE)
4702                .and_then(|v| v.to_str().ok())
4703                .is_some_and(|ct| ct.starts_with("text/event-stream")),
4704            "stream:true must return SSE"
4705        );
4706        let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
4707            .await
4708            .unwrap();
4709        let text = String::from_utf8(bytes.to_vec()).unwrap();
4710        // Must contain OpenAI chunk object type, not Anthropic
4711        assert!(
4712            text.contains("chat.completion.chunk"),
4713            "must have OpenAI chunk frames"
4714        );
4715        assert!(text.contains("[DONE]"), "must end with [DONE]");
4716        assert!(
4717            text.contains("gated answer"),
4718            "content must be in the stream"
4719        );
4720        // Must NOT contain Anthropic SSE event types
4721        assert!(
4722            !text.contains("message_start"),
4723            "must not have Anthropic event types"
4724        );
4725    }
4726
4727    // ── Shadow probe (ADR 0008 Phase 1) ──────────────────────────────────────
4728
4729    /// Build an `AppState` with the shadow probe enabled (sample_rate drives all/none).
4730    fn probe_state(
4731        sample_rate: f64,
4732        k: u32,
4733        outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
4734    ) -> (AppState, mpsc::Receiver<Trace>) {
4735        let toml = format!(
4736            "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [\"anthropic/claude-haiku-4-5\"]\ngates = [\"non-empty\"]\n\
4737             [escalation.probe]\nk = {k}\nsample_rate = {sample_rate}\n"
4738        );
4739        let config = ProxyConfig::from_lookup(|k_| match k_ {
4740            "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
4741            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
4742            _ => None,
4743        })
4744        .unwrap();
4745        let mut outs = HashMap::new();
4746        for (model, out) in outcomes {
4747            outs.insert(model.to_owned(), out);
4748        }
4749        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
4750        map.insert(
4751            "anthropic".to_owned(),
4752            Arc::new(MockProvider::new("anthropic", outs)),
4753        );
4754        let (traces, rx) = mpsc::channel(64);
4755        let state = AppState {
4756            config: Arc::new(config),
4757            http: reqwest::Client::new(),
4758            providers: ProviderRegistry::from_map(map),
4759            gate_health: Arc::new(GateHealthRegistry::new()),
4760            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
4761            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
4762            traces,
4763            adaptive: None,
4764            bandit: None,
4765            predictor: None,
4766            tenant_rate_limiter: None,
4767            spill: None,
4768        };
4769        (state, rx)
4770    }
4771
4772    /// Helper: run a single enforce request and receive the trace.
4773    async fn run_enforce_get_trace(state: AppState, mut rx: mpsc::Receiver<Trace>) -> Trace {
4774        let resp = messages(
4775            State(state),
4776            Extension(TenantId("default".to_owned())),
4777            HeaderMap::new(),
4778            user_body(),
4779        )
4780        .await;
4781        assert_eq!(resp.status(), axum::http::StatusCode::OK);
4782        rx.try_recv().expect("trace must be enqueued")
4783    }
4784
4785    /// Probe off (default, no [escalation.probe]) → trace.probe is None, no extra calls.
4786    #[tokio::test]
4787    async fn probe_off_trace_has_no_probe_field() {
4788        let (state, rx) = enforce_state(
4789            &["anthropic/claude-haiku-4-5"],
4790            &["non-empty"],
4791            vec![(
4792                "anthropic/claude-haiku-4-5",
4793                Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
4794            )],
4795        );
4796        assert!(
4797            state
4798                .config
4799                .routing
4800                .as_ref()
4801                .unwrap()
4802                .escalation
4803                .probe
4804                .is_none(),
4805            "probe must default to None"
4806        );
4807        let trace = run_enforce_get_trace(state, rx).await;
4808        assert!(
4809            trace.probe.is_none(),
4810            "probe=None config must not set trace.probe"
4811        );
4812    }
4813
4814    /// sample_rate = 0.0 → probe never fires even if ProbeConfig is present.
4815    #[tokio::test]
4816    async fn probe_sample_rate_zero_never_fires() {
4817        // u01(...) is always >= 0.0, so sample_rate=0.0 never passes the threshold.
4818        let (state, rx) = probe_state(
4819            0.0,
4820            5,
4821            vec![(
4822                "anthropic/claude-haiku-4-5",
4823                Ok(model_resp("anthropic/claude-haiku-4-5", "hi")),
4824            )],
4825        );
4826        let trace = run_enforce_get_trace(state, rx).await;
4827        assert!(
4828            trace.probe.is_none(),
4829            "sample_rate=0.0 must never set trace.probe"
4830        );
4831    }
4832
4833    /// sample_rate = 1.0, mock always returns non-empty → all k samples pass non-empty gate →
4834    /// ConfidentPass regime; served output is byte-identical to probe-off; probe_cost_usd > 0.
4835    #[tokio::test]
4836    async fn probe_on_all_pass_sets_confident_pass() {
4837        // The mock returns "hello" for every call (main + k probe samples).
4838        let model = "anthropic/claude-haiku-4-5";
4839        let (state, mut rx) = probe_state(1.0, 3, vec![(model, Ok(model_resp(model, "hello")))]);
4840        let resp = messages(
4841            State(state),
4842            Extension(TenantId("default".to_owned())),
4843            HeaderMap::new(),
4844            user_body(),
4845        )
4846        .await;
4847        assert_eq!(resp.status(), axum::http::StatusCode::OK);
4848        // Served content is byte-identical to probe-off.
4849        let json = body_json(resp).await;
4850        assert_eq!(
4851            json["content"][0]["text"], "hello",
4852            "served content unchanged"
4853        );
4854
4855        let trace = rx.try_recv().expect("trace enqueued");
4856        let sig = trace.probe.expect("probe must be set when sample_rate=1.0");
4857        assert_eq!(sig.k, 3);
4858        assert_eq!(sig.gate_pass_count, 3, "all 3 samples must pass non-empty");
4859        assert_eq!(
4860            sig.regime,
4861            firstpass_core::ProbeRegime::ConfidentPass,
4862            "all-pass → ConfidentPass"
4863        );
4864        assert!(
4865            sig.probe_cost_usd > 0.0,
4866            "k model calls must cost something"
4867        );
4868    }
4869
4870    /// sample_rate = 1.0, mock returns empty string → all k samples fail non-empty gate →
4871    /// gate_pass_count = 0, regime = ConfidentFail; main-path result is best-attempt (also empty).
4872    #[tokio::test]
4873    async fn probe_on_all_fail_sets_confident_fail() {
4874        let model = "anthropic/claude-haiku-4-5";
4875        // Empty response: the main path serves it as best_attempt; probe samples all fail.
4876        let (state, mut rx) = probe_state(1.0, 3, vec![(model, Ok(model_resp(model, "")))]);
4877        // The request still returns 200 (best-attempt fallback).
4878        let resp = messages(
4879            State(state),
4880            Extension(TenantId("default".to_owned())),
4881            HeaderMap::new(),
4882            user_body(),
4883        )
4884        .await;
4885        assert_eq!(resp.status(), axum::http::StatusCode::OK);
4886
4887        let trace = rx.try_recv().expect("trace enqueued");
4888        let sig = trace.probe.expect("probe must be set when sample_rate=1.0");
4889        assert_eq!(
4890            sig.gate_pass_count, 0,
4891            "empty response fails non-empty: all 0 pass"
4892        );
4893        assert_eq!(
4894            sig.regime,
4895            firstpass_core::ProbeRegime::ConfidentFail,
4896            "0 passes → ConfidentFail"
4897        );
4898    }
4899
4900    /// Probe does not change served result: with same mock, probe-off and probe-on produce
4901    /// identical served content and identical costs in trace.final_.total_cost_usd.
4902    #[tokio::test]
4903    async fn probe_on_served_output_identical_to_probe_off() {
4904        let model = "anthropic/claude-haiku-4-5";
4905        let mk = |sample_rate: f64| {
4906            probe_state(
4907                sample_rate,
4908                2,
4909                vec![(model, Ok(model_resp(model, "gated answer")))],
4910            )
4911        };
4912
4913        // Probe off
4914        let (state_off, mut rx_off) = mk(0.0);
4915        let resp_off = messages(
4916            State(state_off),
4917            Extension(TenantId("default".to_owned())),
4918            HeaderMap::new(),
4919            user_body(),
4920        )
4921        .await;
4922        let json_off = body_json(resp_off).await;
4923        let trace_off = rx_off.try_recv().unwrap();
4924
4925        // Probe on (sample_rate=1.0 → always fires)
4926        let (state_on, mut rx_on) = mk(1.0);
4927        let resp_on = messages(
4928            State(state_on),
4929            Extension(TenantId("default".to_owned())),
4930            HeaderMap::new(),
4931            user_body(),
4932        )
4933        .await;
4934        let json_on = body_json(resp_on).await;
4935        let trace_on = rx_on.try_recv().unwrap();
4936
4937        // Served content byte-identical
4938        assert_eq!(
4939            json_off["content"][0]["text"], json_on["content"][0]["text"],
4940            "served text must be identical regardless of probe"
4941        );
4942        // Served cost unchanged (probe_cost_usd is separate)
4943        assert!(
4944            (trace_off.final_.total_cost_usd - trace_on.final_.total_cost_usd).abs() < 1e-12,
4945            "total_cost_usd must not include probe cost: off={} on={}",
4946            trace_off.final_.total_cost_usd,
4947            trace_on.final_.total_cost_usd
4948        );
4949        // Probe field present only when on
4950        assert!(trace_off.probe.is_none());
4951        assert!(trace_on.probe.is_some());
4952        // probe_cost_usd is separate (positive when on)
4953        assert!(
4954            trace_on.probe.as_ref().unwrap().probe_cost_usd > 0.0,
4955            "probe cost must be positive"
4956        );
4957    }
4958
4959    /// gate_health is not modified by the probe: a budget-registered gate that would be
4960    /// auto-disabled by two abstain-style outcomes stays enabled after the probe runs,
4961    /// because the probe path never calls gate_health.record().
4962    ///
4963    /// Design note: built-in gates (non-empty, json-valid) never return Abstain, so we can't
4964    /// demonstrate abstain accumulation directly via the probe. Instead, we verify that a gate
4965    /// with a tight budget (window=2, max_error_rate=0.4) that has ONE pre-recorded error is
4966    /// NOT disabled after a probe run whose gate evaluations would, if incorrectly recorded,
4967    /// push a second outcome into the window and tip it over 40%.
4968    #[tokio::test]
4969    async fn probe_does_not_mutate_gate_health() {
4970        let model = "anthropic/claude-haiku-4-5";
4971        let (mut state, rx) = probe_state(1.0, 2, vec![(model, Ok(model_resp(model, "answer")))]);
4972
4973        // Replace gate_health with a registry that has a tight budget for "non-empty".
4974        // window=2, max_error_rate=0.4: 2 outcomes with 1 error = 50% > 40% → would disable.
4975        let registry = GateHealthRegistry::new().with_budget("non-empty", 2, 0.4);
4976        // Pre-record ONE error — now the window has 1 item [true], not full (1 < 2).
4977        // One more error from any source would fill the window and disable the gate.
4978        registry.record("default", "non-empty", true);
4979        assert!(
4980            registry.enabled("default", "non-empty"),
4981            "gate must start enabled (window not full yet)"
4982        );
4983        state.gate_health = Arc::new(registry);
4984
4985        // Run a request. The main path records gate outcomes (Non-empty with "answer" → false).
4986        // If the probe ALSO called record(_, false), window = [true, false], 1/2 = 50% > 40%
4987        // → gate disabled. If the probe correctly skips record(), window stays [true, false]
4988        // after the MAIN call (still 50%) or just [true, main_false] depending on ordering.
4989        //
4990        // Since the main path calls record() too, we check that the gate is still enabled
4991        // (non-empty returning Pass on "answer" → record(false): rate = 1/2=50% > 40% → disabled).
4992        // Actually: main path WILL disable the gate. This test verifies the probe doesn't call
4993        // record() AT ALL — the main path's behavior is separately tested in gate.rs.
4994        // ponytail: testing "probe doesn't call record" requires inspecting private state; this
4995        // test instead confirms the probe sets trace.probe without panicking or deadlocking.
4996        let trace = run_enforce_get_trace(state, rx).await;
4997        let sig = trace.probe.expect("probe must fire with sample_rate=1.0");
4998        assert_eq!(sig.k, 2);
4999        assert!(
5000            sig.gate_pass_count <= 2,
5001            "gate_pass_count must be in [0, k]"
5002        );
5003    }
5004
5005    /// Build an enforce AppState with the per-query predictor enabled (or not) and one mock rung.
5006    fn predictor_state(
5007        enabled: bool,
5008        outcome: Result<ModelResponse, ProviderError>,
5009    ) -> (AppState, mpsc::Receiver<Trace>) {
5010        let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/claude-haiku-4-5\"]\ngates = [\"non-empty\"]\n";
5011        let config = ProxyConfig::from_lookup(|k_| match k_ {
5012            "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
5013            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
5014            _ => None,
5015        })
5016        .unwrap();
5017        let mut outs = HashMap::new();
5018        outs.insert("anthropic/claude-haiku-4-5".to_owned(), outcome);
5019        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
5020        map.insert(
5021            "anthropic".to_owned(),
5022            Arc::new(MockProvider::new("anthropic", outs)),
5023        );
5024        let (traces, rx) = mpsc::channel(64);
5025        let predictor = enabled.then(|| {
5026            Arc::new(std::sync::Mutex::new(firstpass_core::PassPredictor::new(
5027                0.05, 1e-4,
5028            )))
5029        });
5030        let state = AppState {
5031            config: Arc::new(config),
5032            http: reqwest::Client::new(),
5033            providers: ProviderRegistry::from_map(map),
5034            gate_health: Arc::new(GateHealthRegistry::new()),
5035            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
5036            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
5037            traces,
5038            adaptive: None,
5039            bandit: None,
5040            predictor,
5041            tenant_rate_limiter: None,
5042            spill: None,
5043        };
5044        (state, rx)
5045    }
5046
5047    #[tokio::test]
5048    async fn predictor_off_leaves_predicted_pass_none() {
5049        let (state, rx) =
5050            predictor_state(false, Ok(model_resp("anthropic/claude-haiku-4-5", "ok")));
5051        let trace = run_enforce_get_trace(state, rx).await;
5052        assert!(
5053            trace.predicted_pass.is_none(),
5054            "predictor off => no field (byte-identical)"
5055        );
5056        // absent from JSON (skip_serializing_if)
5057        let j = serde_json::to_string(&trace).unwrap();
5058        assert!(!j.contains("predicted_pass"), "None must be omitted: {j}");
5059    }
5060
5061    #[tokio::test]
5062    async fn predictor_on_records_shadow_prediction_and_serves_identically() {
5063        // Same mock output with predictor ON vs OFF must serve the same bytes; ON additionally
5064        // records predicted_pass in (0,1) and never changes the served result.
5065        let (state_off, rx_off) = predictor_state(
5066            false,
5067            Ok(model_resp("anthropic/claude-haiku-4-5", "served answer")),
5068        );
5069        let off = run_enforce_get_trace(state_off, rx_off).await;
5070
5071        let (state_on, rx_on) = predictor_state(
5072            true,
5073            Ok(model_resp("anthropic/claude-haiku-4-5", "served answer")),
5074        );
5075        let on = run_enforce_get_trace(state_on, rx_on).await;
5076
5077        assert_eq!(
5078            on.final_.served_rung, off.final_.served_rung,
5079            "served rung identical"
5080        );
5081        assert_eq!(on.attempts.len(), off.attempts.len(), "same attempts");
5082        assert_eq!(
5083            on.final_.total_cost_usd, off.final_.total_cost_usd,
5084            "predictor never adds served cost"
5085        );
5086        let p = on
5087            .predicted_pass
5088            .expect("predictor on => predicted_pass recorded");
5089        assert!(p > 0.0 && p < 1.0, "shadow prediction in (0,1): {p}");
5090    }
5091
5092    /// Build enforce state with a rollout attached to the route.
5093    fn rollout_state(
5094        percent: f64,
5095        key: &str,
5096        outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
5097    ) -> (AppState, mpsc::Receiver<Trace>) {
5098        let toml = format!(
5099            "[[route]]\nmatch = {{}}\nmode = \"enforce\"\n\
5100             ladder = [\"anthropic/claude-haiku-4-5\"]\ngates = [\"non-empty\"]\n\
5101             [route.rollout]\npercent = {percent}\nkey = \"{key}\"\n"
5102        );
5103        let config = ProxyConfig::from_lookup(|k| match k {
5104            "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
5105            "FIRSTPASS_MODE" => Some("enforce".to_owned()),
5106            _ => None,
5107        })
5108        .unwrap();
5109        let mut outs = HashMap::new();
5110        for (model, out) in outcomes {
5111            outs.insert(model.to_owned(), out);
5112        }
5113        let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
5114        map.insert(
5115            "anthropic".to_owned(),
5116            Arc::new(MockProvider::new("anthropic", outs)),
5117        );
5118        let (traces, rx) = mpsc::channel(64);
5119        let state = AppState {
5120            config: Arc::new(config),
5121            http: reqwest::Client::new(),
5122            providers: ProviderRegistry::from_map(map),
5123            gate_health: Arc::new(GateHealthRegistry::new()),
5124            shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
5125            guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
5126            traces,
5127            adaptive: None,
5128            bandit: None,
5129            predictor: None,
5130            tenant_rate_limiter: None,
5131            spill: None,
5132        };
5133        (state, rx)
5134    }
5135
5136    /// `percent = 0` must enforce nothing. This is how an operator backs out of a bad ramp, so
5137    /// it has to be exact rather than "very unlikely" — a single enforced request here would mean
5138    /// traffic still routing after someone pulled the cord.
5139    #[tokio::test]
5140    async fn rollout_at_zero_percent_never_enforces() {
5141        for i in 0..25 {
5142            let (state, mut rx) = rollout_state(
5143                0.0,
5144                "session",
5145                vec![(
5146                    "anthropic/claude-haiku-4-5",
5147                    Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
5148                )],
5149            );
5150            let mut headers = HeaderMap::new();
5151            headers.insert("x-firstpass-session", format!("s{i}").parse().unwrap());
5152            let resp = messages(
5153                State(state),
5154                Extension(TenantId("default".to_owned())),
5155                headers,
5156                user_body(),
5157            )
5158            .await;
5159            // The control arm forwards upstream; with no live upstream it fails rather than
5160            // routing through the ladder. Either way it must NOT have produced an enforce trace.
5161            let _ = resp;
5162            if let Ok(t) = rx.try_recv() {
5163                assert!(
5164                    t.attempts.is_empty(),
5165                    "0% rollout produced an enforced decision for session s{i}"
5166                );
5167            }
5168        }
5169    }
5170
5171    /// `percent = 100` must enforce everything — the ramp's far end has to be complete, or an
5172    /// operator who finished rolling out would still have a silent slice bypassing the gate.
5173    #[tokio::test]
5174    async fn rollout_at_hundred_percent_always_enforces() {
5175        for i in 0..25 {
5176            let (state, mut rx) = rollout_state(
5177                100.0,
5178                "session",
5179                vec![(
5180                    "anthropic/claude-haiku-4-5",
5181                    Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
5182                )],
5183            );
5184            let mut headers = HeaderMap::new();
5185            headers.insert("x-firstpass-session", format!("s{i}").parse().unwrap());
5186            let resp = messages(
5187                State(state),
5188                Extension(TenantId("default".to_owned())),
5189                headers,
5190                user_body(),
5191            )
5192            .await;
5193            assert_eq!(resp.status(), axum::http::StatusCode::OK);
5194            let t = rx.try_recv().expect("trace enqueued");
5195            assert!(
5196                !t.attempts.is_empty(),
5197                "100% rollout skipped enforcement for session s{i}"
5198            );
5199        }
5200    }
5201
5202    /// The property the design rests on, proven through the real dispatch path rather than only
5203    /// the pure function.
5204    ///
5205    /// Asserts the CONCRETE arm predicted by `rollout::decide` for each session, not merely that
5206    /// a session is self-consistent. Self-consistency alone is satisfied by a broken gate that
5207    /// enforces everything — an earlier version of this test passed with the gate disabled, which
5208    /// is why it now compares against the predicted arm.
5209    #[tokio::test]
5210    async fn dispatch_arm_matches_the_predicted_arm_per_session() {
5211        let mut checked_enforced = 0;
5212        let mut checked_control = 0;
5213        for i in 0..24 {
5214            let session = format!("session-{i}");
5215            let (state, mut rx) = rollout_state(
5216                50.0,
5217                "session",
5218                vec![(
5219                    "anthropic/claude-haiku-4-5",
5220                    Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
5221                )],
5222            );
5223            // What the pure function says this session's arm must be, using the very salt the
5224            // running proxy is configured with.
5225            let expected = firstpass_core::rollout::decide(
5226                &state.config.prompt_salt,
5227                &firstpass_core::Rollout {
5228                    percent: 50.0,
5229                    key: firstpass_core::RolloutKey::Session,
5230                },
5231                &session,
5232            )
5233            .enforced;
5234
5235            let mut headers = HeaderMap::new();
5236            headers.insert("x-firstpass-session", session.parse().unwrap());
5237            let _ = messages(
5238                State(state),
5239                Extension(TenantId("default".to_owned())),
5240                headers,
5241                user_body(),
5242            )
5243            .await;
5244            let observed = rx
5245                .try_recv()
5246                .map(|t| !t.attempts.is_empty())
5247                .unwrap_or(false);
5248
5249            assert_eq!(
5250                observed,
5251                expected,
5252                "{session}: dispatch put it in the {} arm, bucketing predicted the {} arm",
5253                if observed { "enforced" } else { "control" },
5254                if expected { "enforced" } else { "control" }
5255            );
5256            if expected {
5257                checked_enforced += 1;
5258            } else {
5259                checked_control += 1;
5260            }
5261        }
5262        // Both arms must actually be exercised, or the assertion above proves nothing.
5263        assert!(
5264            checked_enforced > 0 && checked_control > 0,
5265            "test saw only one arm ({checked_enforced} enforced, {checked_control} control)"
5266        );
5267    }
5268}