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