1use 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#[derive(Clone)]
45pub struct AppState {
46 pub config: Arc<ProxyConfig>,
48 pub http: reqwest::Client,
50 pub providers: ProviderRegistry,
52 pub gate_health: Arc<GateHealthRegistry>,
54 pub shadow_ledger: Arc<crate::shadow::ShadowLedger>,
57 pub guardrails: Arc<crate::guard::GuardrailRegistry>,
60 pub traces: store::TraceSender,
62 pub adaptive: Option<Arc<std::sync::Mutex<firstpass_core::conformal::AdaptiveConformal>>>,
66 pub bandit: Option<Arc<std::sync::Mutex<crate::bandit::StartRungBandit>>>,
71 pub predictor: Option<Arc<std::sync::Mutex<firstpass_core::PassPredictor>>>,
77 pub tenant_rate_limiter: Option<Arc<governor::DefaultKeyedRateLimiter<String>>>,
81 pub spill: Option<store::SpillHandle>,
85}
86
87#[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
100pub 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
125fn 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
167fn 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 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 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
213const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
217
218pub(crate) fn u01(seed: u128) -> f64 {
229 let lo = splitmix64_finalise(seed as u64);
230 let hi = splitmix64_finalise((seed >> 64) as u64);
231 ((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#[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
254pub fn app(state: AppState) -> Result<Router, ProxyError> {
261 crate::metrics::install()?;
262 let max_concurrency = state.config.max_concurrency;
263
264 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 .layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES))
292 .layer(tower::limit::GlobalConcurrencyLimitLayer::new(
295 max_concurrency,
296 ))
297 .with_state(state))
298}
299
300async fn healthz() -> impl IntoResponse {
302 Json(serde_json::json!({ "status": "ok" }))
303}
304
305async fn capabilities(State(state): State<AppState>) -> impl IntoResponse {
308 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#[derive(Debug, Deserialize)]
344struct FeedbackRequest {
345 trace_id: String,
347 gate_id: String,
349 verdict: String,
351 #[serde(default)]
353 score: Option<f64>,
354 reporter: String,
356}
357
358async 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 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 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 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 let attributed_route = tokio::task::spawn_blocking({
451 let (db, ten, tid) = (
452 state.config.db_path.clone(),
453 tenant.clone(),
454 trace_id.clone(),
455 );
456 move || store::trace_route(&db, &ten, &tid)
457 })
458 .await
459 .ok()
460 .and_then(Result::ok)
461 .flatten()
462 .flatten()
463 .unwrap_or(0) as usize;
464
465 if let (Some(cfg), Some(correct)) = (
470 state.config.routing.as_ref().and_then(|r| r.guardrail),
471 feedback_signal,
472 ) {
473 let cooldown = state
474 .config
475 .routing
476 .as_ref()
477 .map_or(3_600, |r| r.guardrail_cooldown_secs);
478 let reaction = state.guardrails.record(
479 &tenant,
480 attributed_route,
481 &cfg,
482 correct,
483 jiff::Timestamp::now().as_second(),
484 cooldown,
485 );
486 match &reaction {
487 crate::guard::Reaction::Demoted(v) => {
488 metrics::counter!("firstpass_guardrail_demotions_total").increment(1);
489 tracing::error!(
490 tenant = %tenant,
491 n = v.n,
492 rate = v.rate,
493 bound = v.bound,
494 alpha = cfg.alpha,
495 "GUARDRAIL: served-failure bound exceeded target — route demoted to observe; traffic now serves as it would without Firstpass"
496 );
497 }
498 crate::guard::Reaction::Alarmed(v) => {
499 metrics::counter!("firstpass_guardrail_breaches_total").increment(1);
500 tracing::error!(
501 tenant = %tenant,
502 n = v.n,
503 rate = v.rate,
504 bound = v.bound,
505 alpha = cfg.alpha,
506 "GUARDRAIL: served-failure bound exceeded target (alarm only — routing unchanged)"
507 );
508 }
509 crate::guard::Reaction::None => {}
510 }
511 }
512 (
513 axum::http::StatusCode::ACCEPTED,
514 Json(serde_json::json!({ "status": "recorded", "trace_id": trace_id })),
515 )
516 .into_response()
517 }
518 Ok(Err(e)) => {
519 tracing::error!(%e, "feedback: append_deferred failed");
520 ProxyError::Internal(e.to_string()).into_response()
521 }
522 Err(e) => {
523 tracing::error!(%e, "feedback: append_deferred task panicked");
524 ProxyError::Internal(e.to_string()).into_response()
525 }
526 }
527}
528
529const SESSION_HEADER: &str = "x-firstpass-session";
532
533const AGENT_HEADER: &str = "x-firstpass-agent";
535const SUBAGENT_HEADER: &str = "x-firstpass-subagent";
537#[expect(
543 clippy::too_many_arguments,
544 reason = "mirrors the enforce context it evaluates; a wrapper struct would only move the list"
545)]
546fn spawn_shadow(
547 state: &AppState,
548 route: &firstpass_core::Route,
549 route_ix: usize,
550 body: Bytes,
551 auth: Auth,
552 features: Features,
553 tenant: String,
554 session_id: String,
555 api: &str,
556) {
557 let Some(shadow) = route.shadow else {
558 return;
559 };
560 if !shadow.sampled(&state.config.prompt_salt, &session_id) {
563 return;
564 }
565 let (state, route, api) = (state.clone(), route.clone(), api.to_owned());
566 tokio::spawn(async move {
567 let signal = evaluate_shadow(
568 &state, &route, shadow, route_ix, &body, auth, &features, tenant, session_id, &api, 0.0,
569 )
570 .await;
571 tracing::debug!(
573 would_pass = signal.would_pass,
574 projected_usd = signal.projected_cost_usd,
575 skipped = ?signal.skipped,
576 "shadow evaluation complete"
577 );
578 });
579}
580
581#[expect(
588 clippy::too_many_arguments,
589 reason = "mirrors the enforce context; grouping them into a struct would only move the list"
590)]
591async fn evaluate_shadow(
592 state: &AppState,
593 route: &firstpass_core::Route,
594 shadow: firstpass_core::rollout::Shadow,
595 route_ix: usize,
596 body: &Bytes,
597 auth: Auth,
598 features: &Features,
599 tenant: String,
600 session_id: String,
601 api: &str,
602 actual_cost_usd: f64,
603) -> ShadowSignal {
604 let skipped = |why: &str| ShadowSignal {
605 would_serve_rung: None,
606 would_pass: false,
607 projected_cost_usd: 0.0,
608 actual_cost_usd,
609 skipped: Some(why.to_owned()),
610 };
611
612 let now = jiff::Timestamp::now();
613 if !state
614 .shadow_ledger
615 .may_spend(&tenant, route_ix, shadow.max_usd_per_day, now)
616 {
617 return skipped("budget_exhausted");
620 }
621
622 let Some(base_request) = parse_model_request(body) else {
623 return skipped("unparseable_request");
624 };
625 let gate_defs = state
626 .config
627 .routing
628 .as_ref()
629 .map_or(&[][..], |cfg| &cfg.gate_defs);
630 let gates = resolve_gates(
631 &route.gates,
632 gate_defs,
633 &state.providers,
634 &auth,
635 &state.config.prices,
636 );
637 let (budget, max_rungs) = state
638 .config
639 .routing
640 .as_ref()
641 .map_or((None, u32::MAX), |cfg| {
642 (
643 cfg.budget.per_request_usd,
644 cfg.escalation.max_rungs_per_request,
645 )
646 });
647
648 let ctx = EnforceCtx {
649 ladder: &route.ladder,
650 gates: &gates,
651 health: &state.gate_health,
652 base_request: &base_request,
653 providers: &state.providers,
654 auth: &auth,
655 prices: &state.config.prices,
656 budget_per_request_usd: budget,
657 max_rungs,
658 speculation: 0,
661 serve_threshold: None,
662 elastic: None,
663 features: features.clone(),
664 start_rung: 0,
665 tenant_id: tenant.clone(),
666 session_id,
667 prompt_hash: prompt_hash(&state.config.prompt_salt, body),
668 api: api.to_owned(),
669 policy_id: "shadow".to_owned(),
670 };
671
672 let (outcome, trace) = route_enforce(ctx).await;
673 let spent = trace.final_.total_cost_usd;
674 state.shadow_ledger.debit(&tenant, route_ix, spent, now);
675
676 let would_pass = matches!(outcome, EngineOutcome::Served(_));
679 let would_serve_rung = if would_pass {
680 trace.final_.served_rung
681 } else {
682 None
683 };
684 ShadowSignal {
685 would_serve_rung,
686 would_pass,
687 projected_cost_usd: spent,
688 actual_cost_usd,
689 skipped: None,
690 }
691}
692
693const MODE_PROFILE_HEADER: &str = "x-firstpass-mode";
696
697fn resolve_mode(headers: &HeaderMap, route: &Route, config: &ProxyConfig) -> RoutingMode {
706 if let Some(val) = header_str(headers, MODE_PROFILE_HEADER) {
708 match val.trim().to_ascii_lowercase().as_str() {
709 "observe" => return RoutingMode::Observe,
710 "cost" => return RoutingMode::Cost,
711 "balanced" => return RoutingMode::Balanced,
712 "quality" => return RoutingMode::Quality,
713 "latency" => return RoutingMode::Latency,
714 "max" => return RoutingMode::Max,
715 other => {
716 tracing::warn!(
717 value = other,
718 "unknown x-firstpass-mode value; ignoring \
719 (valid: observe|cost|balanced|quality|latency|max)"
720 );
721 }
722 }
723 }
724 if let Some(m) = route.routing_mode {
726 return m;
727 }
728 config.default_routing_mode
730}
731
732async fn messages(
737 State(state): State<AppState>,
738 Extension(TenantId(tenant)): Extension<TenantId>,
739 headers: HeaderMap,
740 body: Bytes,
741) -> Response {
742 let session_header = header_str(&headers, SESSION_HEADER);
743
744 if let Some(routing) = state.config.routing.as_ref() {
747 let features = extract_features(&headers, &body);
748 if let Some(route) = routing
749 .route_for(&features)
750 .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
751 {
752 let route_ix = routing
755 .routes
756 .iter()
757 .position(|r| std::ptr::eq(r, route))
758 .unwrap_or(0);
759 let route = route.clone();
762 let routing_mode = resolve_mode(&headers, &route, &state.config);
764 if routing_mode == RoutingMode::Observe {
766 return observe_passthrough(state, headers, body, session_header, tenant).await;
767 }
768 if state
772 .guardrails
773 .is_demoted(&tenant, route_ix, jiff::Timestamp::now().as_second())
774 {
775 return observe_passthrough(state, headers, body, session_header, tenant).await;
776 }
777 if let Some(rollout) = route.rollout.as_ref() {
782 let key_value = match rollout.key {
783 firstpass_core::RolloutKey::Session => session_header
784 .clone()
785 .unwrap_or_else(|| firstpass_core::rollout::request_identity(&body)),
786 firstpass_core::RolloutKey::Request => {
792 firstpass_core::rollout::request_identity(&body)
793 }
794 firstpass_core::RolloutKey::Tenant => tenant.to_string(),
795 };
796 let decision =
797 firstpass_core::rollout::decide(&state.config.prompt_salt, rollout, &key_value);
798 if !decision.enforced {
799 let shadow_ctx = route.shadow.map(|_| {
801 (
802 body.clone(),
803 Auth::from_headers(&headers),
804 features.clone(),
805 tenant.clone(),
806 session_header
807 .clone()
808 .unwrap_or_else(|| Uuid::now_v7().to_string()),
809 )
810 });
811 let resp =
812 observe_passthrough(state.clone(), headers, body, session_header, tenant)
813 .await;
814 if let Some((b, auth, feats, ten, sess)) = shadow_ctx {
817 spawn_shadow(
818 &state,
819 &route,
820 route_ix,
821 b,
822 auth,
823 feats,
824 ten,
825 sess,
826 "anthropic",
827 );
828 }
829 return resp;
830 }
831 }
832 if enforce_can_handle(
833 &features,
834 &body,
835 routing.escalation.enforce_structured,
836 &route.ladder,
837 &state.providers,
838 Dialect::Anthropic,
839 ) {
840 return handle_enforce(
841 &state,
842 &headers,
843 &body,
844 features,
845 &route,
846 route_ix,
847 session_header,
848 tenant,
849 routing_mode,
850 )
851 .await;
852 }
853 tracing::info!(
857 "enforce route matched but structured request can't be routed faithfully (flag/ladder); serving via observe passthrough"
858 );
859 }
860 }
861 observe_passthrough(state, headers, body, session_header, tenant).await
862}
863
864fn enforce_can_handle(
878 features: &Features,
879 body: &[u8],
880 enforce_structured: bool,
881 ladder: &[String],
882 providers: &crate::provider::ProviderRegistry,
883 inbound: Dialect,
884) -> bool {
885 let structured = features.tool_count > 0
886 || features.has_images
887 || match inbound {
888 Dialect::Anthropic => messages_have_tool_blocks(body),
889 Dialect::Openai => openai_messages_have_tool_calls(body),
890 Dialect::Gemini => false,
891 };
892 if !structured {
893 return true;
894 }
895 if !enforce_structured {
896 return false;
897 }
898 let all_verbatim = ladder.iter().all(|rung| {
900 let provider_id = rung.split('/').next().unwrap_or_default();
901 providers
902 .get(provider_id)
903 .is_some_and(|p| p.carries_structured_verbatim(inbound))
904 });
905 if all_verbatim {
906 return true;
907 }
908 if inbound == Dialect::Openai && !openai_has_http_images(body) {
912 let all_anthropic = ladder.iter().all(|rung| {
913 let pid = rung.split('/').next().unwrap_or_default();
914 providers
915 .get(pid)
916 .is_some_and(|p| p.carries_structured_verbatim(Dialect::Anthropic))
917 });
918 if all_anthropic {
919 return true;
920 }
921 }
922 false
923}
924
925fn messages_have_tool_blocks(body: &[u8]) -> bool {
928 serde_json::from_slice::<Value>(body)
929 .ok()
930 .and_then(|json| {
931 json.get("messages")
932 .and_then(Value::as_array)
933 .map(|messages| messages.iter().any(message_has_tool_block))
934 })
935 .unwrap_or(false)
936}
937
938fn message_has_tool_block(message: &Value) -> bool {
940 message
941 .get("content")
942 .and_then(Value::as_array)
943 .is_some_and(|blocks| {
944 blocks.iter().any(|block| {
945 matches!(
946 block.get("type").and_then(Value::as_str),
947 Some("tool_use" | "tool_result")
948 )
949 })
950 })
951}
952
953fn is_stream_request(body: &[u8]) -> bool {
955 serde_json::from_slice::<Value>(body)
956 .ok()
957 .and_then(|json| json.get("stream").and_then(Value::as_bool))
958 .unwrap_or(false)
959}
960
961fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
963 headers
964 .get(name)
965 .and_then(|v| v.to_str().ok())
966 .map(str::to_owned)
967}
968
969fn extract_features(headers: &HeaderMap, body: &[u8]) -> Features {
972 let (_model, tool_count, has_images) = request_features(body);
973 let mut f = Features::new(TaskKind::Other);
974 f.agent = header_str(headers, AGENT_HEADER);
975 f.subagent = header_str(headers, SUBAGENT_HEADER);
976 f.tool_count = tool_count;
977 f.has_images = has_images;
978 f.prompt_token_bucket = token_bucket(body.len() as u64);
981 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
982 f
983}
984
985#[allow(clippy::too_many_arguments)]
988async fn handle_enforce(
989 state: &AppState,
990 headers: &HeaderMap,
991 body: &Bytes,
992 features: Features,
993 route: &Route,
994 route_ix: usize,
997 session_header: Option<String>,
998 tenant: String,
999 routing_mode: RoutingMode,
1000) -> Response {
1001 if is_stream_request(body) {
1009 let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
1010 let (state_c, headers_c, body_c, route_c) =
1011 (state.clone(), headers.clone(), body.clone(), route.clone());
1012 tokio::spawn(async move {
1013 let out = enforce_pipeline(
1014 &state_c,
1015 &headers_c,
1016 &body_c,
1017 features,
1018 &route_c,
1019 route_ix,
1020 session_header,
1021 tenant,
1022 routing_mode,
1023 )
1024 .await;
1025 let _ = tx.send(out);
1026 });
1027 return sse_keepalive_response(rx, anthropic_sse_from_message);
1028 }
1029 match enforce_pipeline(
1030 state,
1031 headers,
1032 body,
1033 features,
1034 route,
1035 route_ix,
1036 session_header,
1037 tenant,
1038 routing_mode,
1039 )
1040 .await
1041 {
1042 Ok(message) => (axum::http::StatusCode::OK, Json(message)).into_response(),
1043 Err(e) => e.into_response(),
1044 }
1045}
1046
1047#[allow(clippy::too_many_arguments)] async fn enforce_pipeline_inner(
1052 state: &AppState,
1053 body: &Bytes,
1054 base_request: ModelRequest,
1055 auth: Auth,
1056 features: Features,
1057 route: &Route,
1058 session_header: Option<String>,
1059 tenant: String,
1060 api: &str,
1061 routing_mode: RoutingMode,
1062 route_ix: usize,
1063) -> Result<ModelResponse, ProxyError> {
1064 let gate_defs = state
1065 .config
1066 .routing
1067 .as_ref()
1068 .map_or(&[][..], |cfg| &cfg.gate_defs);
1069 let gates = resolve_gates(
1070 &route.gates,
1071 gate_defs,
1072 &state.providers,
1073 &auth,
1074 &state.config.prices,
1075 );
1076 let session_id = session_header.unwrap_or_else(|| Uuid::now_v7().to_string());
1077 let (budget, max_rungs, speculation, serve_threshold, elastic) =
1078 match state.config.routing.as_ref() {
1079 Some(cfg) => (
1080 cfg.budget.per_request_usd,
1081 cfg.escalation.max_rungs_per_request,
1082 cfg.escalation.speculation,
1083 cfg.escalation.serve_threshold,
1084 cfg.escalation.elastic.as_ref(),
1085 ),
1086 None => (None, 3, 0, None, None),
1087 };
1088 let serve_threshold = state
1091 .adaptive
1092 .as_ref()
1093 .and_then(|a| a.lock().ok().map(|g| g.threshold()))
1094 .or(serve_threshold);
1095
1096 let preset = routing_mode.preset();
1099 let max_rungs = if let Some(delta) = preset.max_rungs_delta {
1100 (max_rungs as i32 + delta).max(1) as u32
1101 } else {
1102 max_rungs
1103 };
1104 let speculation = preset.speculation.unwrap_or(speculation);
1105
1106 let bandit_ctx = crate::bandit::ContextBucket::from_features(&features);
1110
1111 let (greedy_rung, base_policy_id, ts_propensity) = {
1115 let (chosen, ts_p) = state
1116 .bandit
1117 .as_ref()
1118 .and_then(|b| b.lock().ok())
1119 .map(|mut b| {
1120 b.choose_start_with_propensity(&bandit_ctx, &route.ladder, &state.config.prices)
1121 })
1122 .unwrap_or((0, None));
1123 let policy = if ts_p.is_some() {
1124 "bandit@v2-ts".to_owned()
1125 } else if chosen > 0 {
1126 "bandit@v1".to_owned()
1127 } else {
1128 "static-ladder@v0".to_owned()
1129 };
1130 (chosen, policy, ts_p)
1131 };
1132
1133 let exploration_epsilon = state
1138 .config
1139 .routing
1140 .as_ref()
1141 .and_then(|cfg| cfg.escalation.exploration.as_ref())
1142 .map(|e| e.epsilon);
1143
1144 let (start_rung, policy_id, explore_flag, propensity) = if ts_propensity.is_some() {
1145 if exploration_epsilon.is_some() {
1148 tracing::warn!(
1149 "bandit.algorithm = thompson already logs propensities; \
1150 [escalation.exploration] epsilon is ignored"
1151 );
1152 }
1153 (greedy_rung, base_policy_id, false, ts_propensity)
1154 } else if let Some(epsilon) = exploration_epsilon {
1155 let k = route.ladder.len().max(1);
1156 let u = u01(Uuid::now_v7().as_u128());
1158 let (chosen, eps_branch) = if u < epsilon {
1159 let idx = ((u / epsilon) * k as f64) as u32;
1161 (idx.min(k as u32 - 1), true)
1162 } else {
1163 (greedy_rung, false)
1164 };
1165 let p = epsilon_propensity(chosen, greedy_rung, epsilon, k);
1166 (chosen, format!("{base_policy_id}+eps"), eps_branch, Some(p))
1167 } else {
1168 (greedy_rung, base_policy_id, false, None)
1169 };
1170
1171 let start_rung = if preset.start_at_top {
1174 route.ladder.len().saturating_sub(1) as u32
1175 } else {
1176 start_rung
1177 };
1178
1179 let speculation = match state
1185 .config
1186 .routing
1187 .as_ref()
1188 .and_then(|cfg| cfg.escalation.speculation_band)
1189 {
1190 Some([lo, hi]) if speculation > 0 => {
1191 let estimate = state
1192 .bandit
1193 .as_ref()
1194 .and_then(|b| b.lock().ok())
1195 .and_then(|b| b.pass_estimate(&bandit_ctx, start_rung));
1196 match estimate {
1197 Some(p) if p < lo || p > hi => {
1198 metrics::counter!("firstpass_speculation_skipped_total").increment(1);
1199 0
1200 }
1201 _ => speculation,
1202 }
1203 }
1204 _ => speculation,
1205 };
1206
1207 if state.bandit.is_some() {
1209 metrics::counter!(
1210 "firstpass_bandit_start_rung",
1211 "rung" => start_rung.to_string()
1212 )
1213 .increment(1);
1214 }
1215
1216 let ctx = EnforceCtx {
1217 ladder: &route.ladder,
1218 gates: &gates,
1219 health: &state.gate_health,
1220 base_request: &base_request,
1221 providers: &state.providers,
1222 auth: &auth,
1223 prices: &state.config.prices,
1224 budget_per_request_usd: budget,
1225 max_rungs,
1226 speculation,
1227 serve_threshold,
1228 elastic,
1229 features,
1230 start_rung,
1231 tenant_id: tenant,
1234 session_id,
1235 prompt_hash: prompt_hash(&state.config.prompt_salt, body),
1236 api: api.to_owned(),
1237 policy_id,
1238 };
1239
1240 let (outcome, mut trace) = route_enforce(ctx).await;
1241
1242 trace.route_ix = Some(u32::try_from(route_ix).unwrap_or(0));
1246 trace.policy.explore = explore_flag;
1249 trace.policy.propensity = propensity;
1250 if routing_mode != RoutingMode::Balanced {
1253 trace.policy.mode_profile = Some(routing_mode.as_str().to_owned());
1254 }
1255
1256 if let Some(bandit) = state.bandit.as_ref()
1260 && let Ok(mut b) = bandit.lock()
1261 {
1262 for attempt in &trace.attempts {
1263 b.observe(&bandit_ctx, attempt.rung, attempt.verdict);
1264 }
1265 }
1266
1267 if let Some(predictor) = state.predictor.as_ref()
1273 && let Ok(mut p) = predictor.lock()
1274 {
1275 let predicted = p.predict(&trace.request.features, start_rung);
1277 for attempt in &trace.attempts {
1278 match attempt.verdict {
1279 Verdict::Pass => p.update(&trace.request.features, attempt.rung, true),
1280 Verdict::Fail => p.update(&trace.request.features, attempt.rung, false),
1281 Verdict::Abstain => {} }
1283 }
1284 trace.predicted_pass = Some(predicted);
1285 metrics::histogram!("firstpass_predictor_pass_prob").record(predicted);
1286 }
1287
1288 if let Some(probe_cfg) = state
1295 .config
1296 .routing
1297 .as_ref()
1298 .and_then(|c| c.escalation.probe)
1299 && u01(Uuid::now_v7().as_u128()) < probe_cfg.sample_rate
1300 {
1301 let probe_rung = (start_rung as usize).min(route.ladder.len().saturating_sub(1));
1303 if let Some(probe_model_str) = route.ladder.get(probe_rung).cloned() {
1304 let probe_provider = ModelRef::parse(&probe_model_str)
1305 .ok()
1306 .and_then(|m| state.providers.get(&m.provider));
1307
1308 if let Some(probe_provider) = probe_provider {
1309 let mut join_set = tokio::task::JoinSet::new();
1312 for _ in 0..probe_cfg.k {
1313 let mut probe_req = base_request.clone();
1314 probe_req.model = probe_model_str.clone();
1315 let probe_auth = auth.clone();
1316 let prov = probe_provider.clone();
1317 join_set.spawn(async move { prov.complete(&probe_req, &probe_auth).await });
1318 }
1319
1320 let fail_closed_owned: std::collections::HashSet<String> = gates
1323 .iter()
1324 .filter(|g| g.abstain_fails_closed())
1325 .map(|g| g.id().to_owned())
1326 .collect();
1327
1328 let mut gate_pass_count = 0u32;
1329 let mut probe_cost_usd = 0.0f64;
1330
1331 while let Some(task_result) = join_set.join_next().await {
1332 let Ok(Ok(probe_resp)) = task_result else {
1333 continue; };
1335 probe_cost_usd += state
1336 .config
1337 .prices
1338 .cost_usd(
1339 &probe_model_str,
1340 probe_resp.in_tokens,
1341 probe_resp.out_tokens,
1342 )
1343 .unwrap_or(0.0);
1344
1345 let mut probe_gate_req = base_request.clone();
1346 probe_gate_req.model = probe_model_str.clone();
1347
1348 let mut probe_gate_results = Vec::with_capacity(gates.len());
1351 for g in &gates {
1352 if !state.gate_health.enabled(&trace.tenant_id, g.id()) {
1354 continue;
1355 }
1356 let r = g.evaluate(&probe_gate_req, &probe_resp).await;
1357 probe_gate_results.push(r);
1359 }
1360
1361 let fail_closed_refs: std::collections::HashSet<&str> =
1362 fail_closed_owned.iter().map(|s| s.as_str()).collect();
1363 let verdict = aggregate_with_policy(&probe_gate_results, &fail_closed_refs);
1364 let passes = match serve_threshold {
1367 None => verdict == Verdict::Pass,
1368 Some(t) => crate::calibrate::gate_score(&probe_gate_results, verdict) >= t,
1369 };
1370 if passes {
1371 gate_pass_count += 1;
1372 }
1373 }
1374
1375 let regime = ProbeRegime::classify(gate_pass_count, probe_cfg.k);
1376 let regime_label = match regime {
1377 ProbeRegime::ConfidentPass => "confident_pass",
1378 ProbeRegime::ConfidentFail => "confident_fail",
1379 ProbeRegime::Ambiguous => "ambiguous",
1380 };
1381 metrics::counter!(
1382 "firstpass_probe_regime_total",
1383 "regime" => regime_label
1384 )
1385 .increment(1);
1386 metrics::gauge!("firstpass_probe_cost_usd_total").increment(probe_cost_usd);
1387 trace.probe = Some(ProbeSignal {
1388 k: probe_cfg.k,
1389 gate_pass_count,
1390 regime,
1391 probe_cost_usd,
1392 });
1393 }
1394 }
1395 }
1396 offer_trace(&state.traces, state.spill.as_ref(), trace);
1400
1401 match outcome {
1402 EngineOutcome::Served(resp) => Ok(resp),
1403 EngineOutcome::Failed(msg) => Err(ProxyError::Engine(msg)),
1404 }
1405}
1406
1407#[allow(clippy::too_many_arguments)]
1410async fn enforce_pipeline(
1411 state: &AppState,
1412 headers: &HeaderMap,
1413 body: &Bytes,
1414 features: Features,
1415 route: &Route,
1416 route_ix: usize,
1418 session_header: Option<String>,
1419 tenant: String,
1420 routing_mode: RoutingMode,
1421) -> Result<Value, ProxyError> {
1422 let Some(base_request) = parse_model_request(body) else {
1423 return Err(ProxyError::BadRequest(
1424 "request body is not a valid Anthropic Messages request".to_owned(),
1425 ));
1426 };
1427 let auth = Auth::from_headers(headers);
1428 let resp = enforce_pipeline_inner(
1429 state,
1430 body,
1431 base_request,
1432 auth,
1433 features,
1434 route,
1435 session_header,
1436 tenant,
1437 "anthropic.messages",
1438 routing_mode,
1439 route_ix,
1440 )
1441 .await?;
1442 Ok(anthropic_response_json(&resp))
1443}
1444
1445#[allow(clippy::too_many_arguments)]
1448async fn enforce_pipeline_openai(
1449 state: &AppState,
1450 headers: &HeaderMap,
1451 body: &Bytes,
1452 features: Features,
1453 route: &Route,
1454 route_ix: usize,
1456 session_header: Option<String>,
1457 tenant: String,
1458 routing_mode: RoutingMode,
1459) -> Result<Value, ProxyError> {
1460 let providers = &state.providers;
1462 let all_openai = route.ladder.iter().all(|rung| {
1463 let pid = rung.split('/').next().unwrap_or_default();
1464 providers
1465 .get(pid)
1466 .is_some_and(|p| p.carries_structured_verbatim(Dialect::Openai))
1467 });
1468 let Some(base_request) = parse_openai_request(body, all_openai) else {
1469 return Err(ProxyError::BadRequest(
1470 "request body is not a valid OpenAI Chat Completions request".to_owned(),
1471 ));
1472 };
1473 let auth = Auth::from_headers(headers);
1474 let resp = enforce_pipeline_inner(
1475 state,
1476 body,
1477 base_request,
1478 auth,
1479 features,
1480 route,
1481 session_header,
1482 tenant,
1483 "openai.chat_completions",
1484 routing_mode,
1485 route_ix,
1486 )
1487 .await?;
1488 Ok(openai_response_json(&resp))
1489}
1490
1491const SSE_KEEPALIVE_EVERY: Duration = Duration::from_secs(5);
1493
1494fn sse_keepalive_response(
1503 rx: tokio::sync::oneshot::Receiver<Result<Value, ProxyError>>,
1504 format_message: fn(&Value) -> String,
1505) -> Response {
1506 let mut ticks = tokio::time::interval(SSE_KEEPALIVE_EVERY);
1507 ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
1508 ticks.reset(); let stream = KeepaliveStream {
1510 rx: Some(rx),
1511 ticks,
1512 format_message,
1513 };
1514 (
1515 axum::http::StatusCode::OK,
1516 [(
1517 axum::http::header::CONTENT_TYPE,
1518 "text/event-stream; charset=utf-8",
1519 )],
1520 axum::body::Body::from_stream(stream),
1521 )
1522 .into_response()
1523}
1524
1525struct KeepaliveStream {
1528 rx: Option<tokio::sync::oneshot::Receiver<Result<Value, ProxyError>>>,
1530 ticks: tokio::time::Interval,
1531 format_message: fn(&Value) -> String,
1533}
1534
1535impl futures_core::Stream for KeepaliveStream {
1536 type Item = Result<Bytes, std::convert::Infallible>;
1537
1538 fn poll_next(
1539 mut self: std::pin::Pin<&mut Self>,
1540 cx: &mut std::task::Context<'_>,
1541 ) -> std::task::Poll<Option<Self::Item>> {
1542 use std::task::Poll;
1543 let Some(rx) = self.rx.as_mut() else {
1544 return Poll::Ready(None); };
1546 if let Poll::Ready(out) = std::pin::Pin::new(rx).poll(cx) {
1547 let fmt = self.format_message;
1548 let frame = match out {
1549 Ok(Ok(message)) => fmt(&message),
1550 Ok(Err(e)) => sse_error_event(&e),
1551 Err(_) => sse_error_event(&ProxyError::Internal(
1552 "enforce pipeline task dropped".to_owned(),
1553 )),
1554 };
1555 self.rx = None;
1556 return Poll::Ready(Some(Ok(Bytes::from(frame))));
1557 }
1558 if self.ticks.poll_tick(cx).is_ready() {
1559 return Poll::Ready(Some(Ok(Bytes::from_static(b": firstpass routing\n\n"))));
1560 }
1561 Poll::Pending
1562 }
1563}
1564
1565fn sse_error_event(e: &ProxyError) -> String {
1568 let mut out = String::new();
1569 sse_event(
1570 &mut out,
1571 "error",
1572 &serde_json::json!({
1573 "type": "error",
1574 "error": { "type": "api_error", "message": e.client_message() }
1575 }),
1576 );
1577 out
1578}
1579
1580fn parse_model_request(body: &[u8]) -> Option<ModelRequest> {
1589 let json: Value = serde_json::from_slice(body).ok()?;
1590 let raw = json.clone();
1591 let messages_json = json.get("messages")?.as_array()?;
1592 let messages = messages_json
1593 .iter()
1594 .map(|m| ChatMessage {
1595 role: m
1596 .get("role")
1597 .and_then(Value::as_str)
1598 .unwrap_or("user")
1599 .to_owned(),
1600 content: m
1601 .get("content")
1602 .cloned()
1603 .unwrap_or_else(|| Value::String(String::new())),
1604 })
1605 .collect();
1606 let system = json
1607 .get("system")
1608 .and_then(Value::as_str)
1609 .map(str::to_owned);
1610 let max_tokens = json
1611 .get("max_tokens")
1612 .and_then(Value::as_u64)
1613 .and_then(|n| u32::try_from(n).ok())
1614 .unwrap_or(1024);
1615 let tools = json.get("tools").cloned().unwrap_or(Value::Null);
1616 Some(ModelRequest {
1617 model: json
1618 .get("model")
1619 .and_then(Value::as_str)
1620 .unwrap_or_default()
1621 .to_owned(),
1622 system,
1623 messages,
1624 max_tokens,
1625 tools,
1626 raw,
1627 })
1628}
1629
1630fn anthropic_response_json(resp: &ModelResponse) -> Value {
1640 let content = resp
1641 .raw
1642 .get("content")
1643 .filter(|c| c.is_array())
1644 .cloned()
1645 .unwrap_or_else(|| serde_json::json!([{ "type": "text", "text": resp.text }]));
1646 serde_json::json!({
1647 "id": format!("msg_{}", Uuid::now_v7()),
1648 "type": "message",
1649 "role": "assistant",
1650 "model": resp.model,
1651 "content": content,
1652 "usage": { "input_tokens": resp.in_tokens, "output_tokens": resp.out_tokens },
1653 })
1654}
1655
1656fn sse_event(out: &mut String, event: &str, data: &Value) {
1658 out.push_str("event: ");
1659 out.push_str(event);
1660 out.push_str("\ndata: ");
1661 out.push_str(&data.to_string());
1662 out.push_str("\n\n");
1663}
1664
1665fn anthropic_sse_from_message(message: &Value) -> String {
1672 let mut out = String::new();
1673
1674 let mut start_msg = message.clone();
1676 start_msg["content"] = Value::Array(Vec::new());
1677 sse_event(
1678 &mut out,
1679 "message_start",
1680 &serde_json::json!({ "type": "message_start", "message": start_msg }),
1681 );
1682
1683 let empty = Vec::new();
1684 let blocks = message
1685 .get("content")
1686 .and_then(Value::as_array)
1687 .unwrap_or(&empty);
1688 for (i, block) in blocks.iter().enumerate() {
1689 match block.get("type").and_then(Value::as_str) {
1690 Some("tool_use") => {
1691 let mut shell = block.clone();
1693 shell["input"] = serde_json::json!({});
1694 sse_event(
1695 &mut out,
1696 "content_block_start",
1697 &serde_json::json!({ "type": "content_block_start", "index": i, "content_block": shell }),
1698 );
1699 let input_json = block
1700 .get("input")
1701 .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
1702 sse_event(
1703 &mut out,
1704 "content_block_delta",
1705 &serde_json::json!({ "type": "content_block_delta", "index": i,
1706 "delta": { "type": "input_json_delta", "partial_json": input_json } }),
1707 );
1708 }
1709 _ => {
1710 let text = block.get("text").and_then(Value::as_str).unwrap_or("");
1712 sse_event(
1713 &mut out,
1714 "content_block_start",
1715 &serde_json::json!({ "type": "content_block_start", "index": i,
1716 "content_block": { "type": "text", "text": "" } }),
1717 );
1718 sse_event(
1719 &mut out,
1720 "content_block_delta",
1721 &serde_json::json!({ "type": "content_block_delta", "index": i,
1722 "delta": { "type": "text_delta", "text": text } }),
1723 );
1724 }
1725 }
1726 sse_event(
1727 &mut out,
1728 "content_block_stop",
1729 &serde_json::json!({ "type": "content_block_stop", "index": i }),
1730 );
1731 }
1732
1733 let out_tokens = message
1734 .pointer("/usage/output_tokens")
1735 .cloned()
1736 .unwrap_or_else(|| Value::from(0));
1737 sse_event(
1738 &mut out,
1739 "message_delta",
1740 &serde_json::json!({ "type": "message_delta", "delta": { "stop_reason": "end_turn" },
1741 "usage": { "output_tokens": out_tokens } }),
1742 );
1743 sse_event(
1744 &mut out,
1745 "message_stop",
1746 &serde_json::json!({ "type": "message_stop" }),
1747 );
1748 out
1749}
1750
1751async fn observe_passthrough(
1753 state: AppState,
1754 headers: HeaderMap,
1755 body: Bytes,
1756 session_header: Option<String>,
1757 tenant: String,
1758) -> Response {
1759 if is_stream_request(&body) {
1761 return observe_stream(state, headers, body, session_header, tenant).await;
1762 }
1763 let start = Instant::now();
1764 let result = forward_anthropic(
1765 &state.http,
1766 &state.config.upstream_anthropic,
1767 &headers,
1768 body.clone(),
1769 )
1770 .await;
1771 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
1772
1773 match result {
1774 Ok((status, resp_headers, resp_body)) => {
1775 spawn_trace(
1779 &state,
1780 body,
1781 Some(resp_body.clone()),
1782 latency_ms,
1783 session_header,
1784 tenant,
1785 );
1786 (status, resp_headers, resp_body).into_response()
1787 }
1788 Err(err) => {
1789 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
1790 err.into_response()
1791 }
1792 }
1793}
1794
1795async fn observe_stream(
1804 state: AppState,
1805 headers: HeaderMap,
1806 body: Bytes,
1807 session_header: Option<String>,
1808 tenant: String,
1809) -> Response {
1810 let start = Instant::now();
1811 let result = forward_anthropic_streaming(
1812 &state.http,
1813 &state.config.upstream_anthropic,
1814 &headers,
1815 body.clone(),
1816 )
1817 .await;
1818 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
1819
1820 match result {
1821 Ok((status, resp_headers, response)) => {
1822 spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
1823 let stream_body = Body::from_stream(response.bytes_stream());
1824 (status, resp_headers, stream_body).into_response()
1825 }
1826 Err(err) => {
1827 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
1828 err.into_response()
1829 }
1830 }
1831}
1832
1833fn spawn_stream_trace(
1835 state: &AppState,
1836 req_body: Bytes,
1837 latency_ms: u64,
1838 session_header: Option<String>,
1839 tenant: String,
1840) {
1841 let config = state.config.clone();
1842 let traces = state.traces.clone();
1843 let spill = state.spill.clone();
1844 tokio::spawn(async move {
1845 let mut trace =
1846 build_stream_trace(&config, &req_body, latency_ms, session_header.as_deref());
1847 trace.tenant_id = tenant;
1849 offer_trace(&traces, spill.as_ref(), trace);
1850 });
1851}
1852
1853fn spawn_trace(
1858 state: &AppState,
1859 req_body: Bytes,
1860 resp_body: Option<Bytes>,
1861 latency_ms: u64,
1862 session_header: Option<String>,
1863 tenant: String,
1864) {
1865 let config = state.config.clone();
1866 let traces = state.traces.clone();
1867 let spill = state.spill.clone();
1868 tokio::spawn(async move {
1869 let mut trace = match resp_body {
1870 Some(resp) => build_trace(
1871 &config,
1872 &req_body,
1873 &resp,
1874 latency_ms,
1875 session_header.as_deref(),
1876 ),
1877 None => build_error_trace(&config, &req_body, latency_ms, session_header.as_deref()),
1878 };
1879 trace.tenant_id = tenant;
1881 offer_trace(&traces, spill.as_ref(), trace);
1882 });
1883}
1884
1885fn session_id(session_header: Option<&str>, trace_id: Uuid) -> String {
1887 session_header
1888 .map(str::to_owned)
1889 .unwrap_or_else(|| trace_id.to_string())
1890}
1891
1892fn prompt_hash(salt: &str, body: &[u8]) -> String {
1895 let mut salted = Vec::with_capacity(salt.len() + body.len());
1896 salted.extend_from_slice(salt.as_bytes());
1897 salted.extend_from_slice(body);
1898 sha256_hex(&salted)
1899}
1900
1901fn request_features(body: &[u8]) -> (Option<String>, u32, bool) {
1905 let Ok(json) = serde_json::from_slice::<Value>(body) else {
1906 return (None, 0, false);
1907 };
1908 let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
1909 let tool_count = json
1910 .get("tools")
1911 .and_then(Value::as_array)
1912 .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
1913 let has_images = json
1914 .get("messages")
1915 .and_then(Value::as_array)
1916 .is_some_and(|messages| messages.iter().any(message_has_image));
1917 (model, tool_count, has_images)
1918}
1919
1920fn message_has_image(message: &Value) -> bool {
1922 message
1923 .get("content")
1924 .and_then(Value::as_array)
1925 .is_some_and(|blocks| {
1926 blocks
1927 .iter()
1928 .any(|block| block.get("type").and_then(Value::as_str) == Some("image"))
1929 })
1930}
1931
1932fn response_usage(body: &[u8]) -> (Option<String>, u64, u64) {
1935 let Ok(json) = serde_json::from_slice::<Value>(body) else {
1936 return (None, 0, 0);
1937 };
1938 let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
1939 let in_tokens = json
1940 .pointer("/usage/input_tokens")
1941 .and_then(Value::as_u64)
1942 .unwrap_or(0);
1943 let out_tokens = json
1944 .pointer("/usage/output_tokens")
1945 .and_then(Value::as_u64)
1946 .unwrap_or(0);
1947 (model, in_tokens, out_tokens)
1948}
1949
1950fn build_trace(
1952 config: &ProxyConfig,
1953 req_body: &Bytes,
1954 resp_body: &Bytes,
1955 latency_ms: u64,
1956 session_header: Option<&str>,
1957) -> Trace {
1958 let (req_model, tool_count, has_images) = request_features(req_body);
1959 let (resp_model, in_tokens, out_tokens) = response_usage(resp_body);
1960 let model = resp_model
1961 .or(req_model)
1962 .unwrap_or_else(|| "unknown".to_owned());
1963
1964 let cost_usd = config
1965 .prices
1966 .cost_usd(&format!("anthropic/{model}"), in_tokens, out_tokens)
1967 .unwrap_or(0.0);
1968
1969 let attempt = Attempt {
1970 rung: 0,
1971 model,
1972 provider: "anthropic".to_owned(),
1973 in_tokens,
1974 out_tokens,
1975 cost_usd,
1976 latency_ms,
1977 gates: Vec::new(),
1978 verdict: Verdict::Pass,
1979 };
1980
1981 let mut trace = base_trace(config, req_body, latency_ms, session_header);
1982 trace.request.features.prompt_token_bucket = token_bucket(in_tokens);
1983 trace.request.features.tool_count = tool_count;
1984 trace.request.features.has_images = has_images;
1985 trace.attempts.push(attempt);
1986 trace.final_ = FinalOutcome {
1987 served_rung: Some(0),
1988 served_from: ServedFrom::Attempt,
1989 total_cost_usd: cost_usd,
1990 gate_cost_usd: 0.0,
1991 total_latency_ms: latency_ms,
1992 escalations: 0,
1993 counterfactual_baseline_usd: cost_usd,
1994 savings_usd: 0.0,
1995 };
1996 trace.recompute_savings();
1997 trace
1998}
1999
2000fn build_stream_trace(
2004 config: &ProxyConfig,
2005 req_body: &Bytes,
2006 latency_ms: u64,
2007 session_header: Option<&str>,
2008) -> Trace {
2009 let (req_model, tool_count, has_images) = request_features(req_body);
2010 let model = req_model.unwrap_or_else(|| "unknown".to_owned());
2011
2012 let attempt = Attempt {
2013 rung: 0,
2014 model,
2015 provider: "anthropic".to_owned(),
2016 in_tokens: 0,
2017 out_tokens: 0,
2018 cost_usd: 0.0,
2019 latency_ms,
2020 gates: Vec::new(),
2021 verdict: Verdict::Pass,
2022 };
2023
2024 let mut trace = base_trace(config, req_body, latency_ms, session_header);
2025 trace.request.features.tool_count = tool_count;
2026 trace.request.features.has_images = has_images;
2027 trace.attempts.push(attempt);
2028 trace.final_ = FinalOutcome {
2029 served_rung: Some(0),
2030 served_from: ServedFrom::Attempt,
2031 total_cost_usd: 0.0,
2032 gate_cost_usd: 0.0,
2033 total_latency_ms: latency_ms,
2034 escalations: 0,
2035 counterfactual_baseline_usd: 0.0,
2036 savings_usd: 0.0,
2037 };
2038 trace.recompute_savings();
2039 trace
2040}
2041
2042fn build_error_trace(
2046 config: &ProxyConfig,
2047 req_body: &Bytes,
2048 latency_ms: u64,
2049 session_header: Option<&str>,
2050) -> Trace {
2051 let (_, tool_count, has_images) = request_features(req_body);
2052 let mut trace = base_trace(config, req_body, latency_ms, session_header);
2053 trace.request.features.tool_count = tool_count;
2054 trace.request.features.has_images = has_images;
2055 trace.final_ = FinalOutcome {
2056 served_rung: None,
2057 served_from: ServedFrom::Error,
2058 total_cost_usd: 0.0,
2059 gate_cost_usd: 0.0,
2060 total_latency_ms: latency_ms,
2061 escalations: 0,
2062 counterfactual_baseline_usd: 0.0,
2063 savings_usd: 0.0,
2064 };
2065 trace.recompute_savings();
2066 trace
2067}
2068
2069fn base_trace(
2072 config: &ProxyConfig,
2073 req_body: &Bytes,
2074 latency_ms: u64,
2075 session_header: Option<&str>,
2076) -> Trace {
2077 let trace_id = Uuid::now_v7();
2078 let mut features = Features::new(TaskKind::Other);
2079 features.hour_bucket = hour_bucket(jiff::Timestamp::now());
2080
2081 Trace {
2082 trace_id,
2083 prev_hash: GENESIS_HASH.to_owned(),
2084 tenant_id: config.tenant_id.clone(),
2085 session_id: session_id(session_header, trace_id),
2086 ts: jiff::Timestamp::now(),
2087 mode: Mode::Observe,
2088 policy: PolicyRef {
2089 id: "observe-passthrough@v0".to_owned(),
2090 explore: false,
2091 propensity: None,
2092 mode_profile: None,
2093 },
2094 request: RequestInfo {
2095 api: "anthropic.messages".to_owned(),
2096 prompt_hash: prompt_hash(&config.prompt_salt, req_body),
2097 features,
2098 },
2099 attempts: Vec::new(),
2100 deferred: Vec::new(),
2101 final_: FinalOutcome {
2102 served_rung: None,
2103 served_from: ServedFrom::Error,
2104 total_cost_usd: 0.0,
2105 gate_cost_usd: 0.0,
2106 total_latency_ms: latency_ms,
2107 escalations: 0,
2108 counterfactual_baseline_usd: 0.0,
2109 savings_usd: 0.0,
2110 },
2111 probe: None,
2112 rollout: None,
2113 shadow: None,
2114 route_ix: None,
2115 predicted_pass: None,
2116 elastic: None,
2117 }
2118}
2119
2120fn openai_messages_have_tool_calls(body: &[u8]) -> bool {
2125 serde_json::from_slice::<Value>(body)
2126 .ok()
2127 .and_then(|json| {
2128 json.get("messages").and_then(Value::as_array).map(|msgs| {
2129 msgs.iter().any(|m| {
2130 m.get("tool_calls").is_some()
2131 || m.get("role").and_then(Value::as_str) == Some("tool")
2132 })
2133 })
2134 })
2135 .unwrap_or(false)
2136}
2137
2138fn openai_has_http_images(body: &[u8]) -> bool {
2142 serde_json::from_slice::<Value>(body)
2143 .ok()
2144 .and_then(|json| {
2145 json.get("messages").and_then(Value::as_array).map(|msgs| {
2146 msgs.iter().any(|m| {
2147 m.get("content")
2148 .and_then(Value::as_array)
2149 .is_some_and(|parts| {
2150 parts.iter().any(|p| {
2151 p.get("type").and_then(Value::as_str) == Some("image_url")
2152 && p.pointer("/image_url/url")
2153 .and_then(Value::as_str)
2154 .is_some_and(|u| {
2155 u.starts_with("http://") || u.starts_with("https://")
2156 })
2157 })
2158 })
2159 })
2160 })
2161 })
2162 .unwrap_or(false)
2163}
2164
2165fn openai_messages_have_images(body: &[u8]) -> bool {
2168 serde_json::from_slice::<Value>(body)
2169 .ok()
2170 .and_then(|json| {
2171 json.get("messages").and_then(Value::as_array).map(|msgs| {
2172 msgs.iter().any(|m| {
2173 m.get("content")
2174 .and_then(Value::as_array)
2175 .is_some_and(|parts| {
2176 parts
2177 .iter()
2178 .any(|p| p.get("type").and_then(Value::as_str) == Some("image_url"))
2179 })
2180 })
2181 })
2182 })
2183 .unwrap_or(false)
2184}
2185
2186fn extract_openai_features(headers: &HeaderMap, body: &[u8]) -> Features {
2189 let Ok(json) = serde_json::from_slice::<Value>(body) else {
2190 let mut f = Features::new(TaskKind::Other);
2191 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
2192 return f;
2193 };
2194 let tool_count = json
2195 .get("tools")
2196 .and_then(Value::as_array)
2197 .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
2198 let has_images = openai_messages_have_images(body);
2199 let mut f = Features::new(TaskKind::Other);
2200 f.agent = header_str(headers, AGENT_HEADER);
2201 f.subagent = header_str(headers, SUBAGENT_HEADER);
2202 f.tool_count = tool_count;
2203 f.has_images = has_images;
2204 f.prompt_token_bucket = token_bucket(body.len() as u64);
2205 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
2206 f
2207}
2208
2209fn parse_data_url(url: &str) -> Option<(&str, &str)> {
2213 let rest = url.strip_prefix("data:")?;
2214 let (meta, data) = rest.split_once(',')?;
2215 let media_type = meta.strip_suffix(";base64")?;
2216 Some((media_type, data))
2217}
2218
2219fn translate_openai_user_content(content: &Value) -> Option<Value> {
2222 match content {
2223 Value::String(_) => Some(content.clone()),
2225 Value::Array(parts) => {
2226 let mut blocks: Vec<Value> = Vec::with_capacity(parts.len());
2227 for part in parts {
2228 match part.get("type").and_then(Value::as_str) {
2229 Some("text") => {
2230 let text = part.get("text").and_then(Value::as_str).unwrap_or("");
2231 blocks.push(serde_json::json!({ "type": "text", "text": text }));
2232 }
2233 Some("image_url") => {
2234 let url = part.pointer("/image_url/url").and_then(Value::as_str)?;
2235 if url.starts_with("http://") || url.starts_with("https://") {
2236 return None; }
2238 let (media_type, data) = parse_data_url(url)?;
2240 blocks.push(serde_json::json!({
2241 "type": "image",
2242 "source": { "type": "base64", "media_type": media_type, "data": data }
2243 }));
2244 }
2245 _ => {} }
2247 }
2248 Some(Value::Array(blocks))
2249 }
2250 _ => Some(Value::String(String::new())),
2251 }
2252}
2253
2254fn translate_openai_tools(tools: &Value) -> Value {
2258 let Some(arr) = tools.as_array() else {
2259 return Value::Null;
2260 };
2261 let converted: Vec<Value> = arr
2262 .iter()
2263 .map(|tool| {
2264 let func = tool.get("function").unwrap_or(&Value::Null);
2265 let mut out = serde_json::json!({
2266 "name": func.get("name").cloned().unwrap_or(Value::String(String::new())),
2267 "input_schema": func.get("parameters").cloned()
2268 .unwrap_or_else(|| serde_json::json!({ "type": "object" })),
2269 });
2270 if let Some(desc) = func.get("description") {
2271 out["description"] = desc.clone();
2272 }
2273 out
2274 })
2275 .collect();
2276 Value::Array(converted)
2277}
2278
2279fn translate_openai_tool_choice(tc: &Value) -> Value {
2281 match tc {
2282 Value::String(s) => match s.as_str() {
2283 "auto" => serde_json::json!({ "type": "auto" }),
2284 "required" => serde_json::json!({ "type": "any" }),
2285 _ => serde_json::json!({ "type": "auto" }),
2287 },
2288 Value::Object(_) => {
2289 if tc.get("type").and_then(Value::as_str) == Some("function") {
2291 let name = tc.pointer("/function/name").cloned().unwrap_or(Value::Null);
2292 serde_json::json!({ "type": "tool", "name": name })
2293 } else {
2294 serde_json::json!({ "type": "auto" })
2295 }
2296 }
2297 _ => serde_json::json!({ "type": "auto" }),
2298 }
2299}
2300
2301pub fn parse_openai_request(body: &[u8], carry_raw: bool) -> Option<ModelRequest> {
2314 let json: Value = serde_json::from_slice(body).ok()?;
2315 let raw = if carry_raw { json.clone() } else { Value::Null };
2316
2317 let messages_json = json.get("messages")?.as_array()?;
2318
2319 let mut system: Option<String> = None;
2320 let mut messages: Vec<ChatMessage> = Vec::with_capacity(messages_json.len());
2321 let mut tools = Value::Null;
2322 let mut tool_choice_override: Option<Value> = None;
2323
2324 for msg in messages_json {
2325 let role = msg.get("role").and_then(Value::as_str).unwrap_or("user");
2326 match role {
2327 "system" => {
2328 if let Some(s) = msg.get("content").and_then(Value::as_str) {
2333 system = Some(s.to_owned());
2334 }
2335 }
2336 "user" => {
2337 let content_val = msg.get("content").unwrap_or(&Value::Null);
2338 let translated = translate_openai_user_content(content_val)?;
2339 messages.push(ChatMessage {
2340 role: "user".to_owned(),
2341 content: translated,
2342 });
2343 }
2344 "assistant" => {
2345 if let Some(tc_arr) = msg.get("tool_calls").and_then(Value::as_array) {
2346 let mut blocks: Vec<Value> = Vec::new();
2348 if let Some(text) = msg.get("content").and_then(Value::as_str)
2350 && !text.is_empty()
2351 {
2352 blocks.push(serde_json::json!({ "type": "text", "text": text }));
2353 }
2354 for tc in tc_arr {
2355 let id = tc.get("id").and_then(Value::as_str).unwrap_or("");
2356 let func = tc.get("function").unwrap_or(&Value::Null);
2357 let name = func.get("name").and_then(Value::as_str).unwrap_or("");
2358 let args_str = func
2359 .get("arguments")
2360 .and_then(Value::as_str)
2361 .unwrap_or("{}");
2362 let input: Value = serde_json::from_str(args_str)
2363 .unwrap_or_else(|_| serde_json::json!({}));
2364 blocks.push(serde_json::json!({
2365 "type": "tool_use",
2366 "id": id,
2367 "name": name,
2368 "input": input,
2369 }));
2370 }
2371 messages.push(ChatMessage {
2372 role: "assistant".to_owned(),
2373 content: Value::Array(blocks),
2374 });
2375 } else {
2376 let content = msg
2378 .get("content")
2379 .cloned()
2380 .unwrap_or_else(|| Value::String(String::new()));
2381 messages.push(ChatMessage {
2382 role: "assistant".to_owned(),
2383 content,
2384 });
2385 }
2386 }
2387 "tool" => {
2388 let tool_call_id = msg
2390 .get("tool_call_id")
2391 .and_then(Value::as_str)
2392 .unwrap_or("");
2393 let content = msg
2394 .get("content")
2395 .cloned()
2396 .unwrap_or_else(|| Value::String(String::new()));
2397 let result_block = serde_json::json!({
2398 "type": "tool_result",
2399 "tool_use_id": tool_call_id,
2400 "content": content,
2401 });
2402 messages.push(ChatMessage {
2403 role: "user".to_owned(),
2404 content: Value::Array(vec![result_block]),
2405 });
2406 }
2407 _ => {} }
2409 }
2410
2411 if !carry_raw {
2413 if let Some(t) = json.get("tools") {
2414 tools = translate_openai_tools(t);
2415 }
2416 if let Some(tc) = json.get("tool_choice") {
2417 tool_choice_override = Some(translate_openai_tool_choice(tc));
2418 }
2419 } else {
2420 tools = json.get("tools").cloned().unwrap_or(Value::Null);
2421 }
2422
2423 let max_tokens = json
2424 .get("max_tokens")
2425 .and_then(Value::as_u64)
2426 .and_then(|n| u32::try_from(n).ok())
2427 .unwrap_or(1024);
2428
2429 let model = json
2430 .get("model")
2431 .and_then(Value::as_str)
2432 .unwrap_or_default()
2433 .to_owned();
2434
2435 let _ = tool_choice_override; Some(ModelRequest {
2446 model,
2447 system,
2448 messages,
2449 max_tokens,
2450 tools,
2451 raw,
2452 })
2453}
2454
2455fn extract_openai_content_and_tools(raw: &Value, text: &str) -> (Value, Option<Value>) {
2462 if let Some(blocks) = raw.get("content").and_then(Value::as_array) {
2464 let mut text_parts: Vec<&str> = Vec::new();
2465 let mut tool_calls: Vec<Value> = Vec::new();
2466 for block in blocks {
2467 match block.get("type").and_then(Value::as_str) {
2468 Some("text") => {
2469 if let Some(t) = block.get("text").and_then(Value::as_str) {
2470 text_parts.push(t);
2471 }
2472 }
2473 Some("tool_use") => {
2474 let id = block.get("id").and_then(Value::as_str).unwrap_or("");
2475 let name = block.get("name").and_then(Value::as_str).unwrap_or("");
2476 let input_str = block
2477 .get("input")
2478 .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
2479 tool_calls.push(serde_json::json!({
2480 "id": id,
2481 "type": "function",
2482 "function": { "name": name, "arguments": input_str },
2483 }));
2484 }
2485 _ => {}
2486 }
2487 }
2488 let content_text = if tool_calls.is_empty() || !text_parts.is_empty() {
2489 Value::String(text_parts.join(""))
2490 } else {
2491 Value::Null };
2493 let tc = if tool_calls.is_empty() {
2494 None
2495 } else {
2496 Some(Value::Array(tool_calls))
2497 };
2498 return (content_text, tc);
2499 }
2500
2501 if let Some(msg) = raw.pointer("/choices/0/message") {
2503 let content = msg
2504 .get("content")
2505 .cloned()
2506 .unwrap_or(Value::String(text.to_owned()));
2507 let tc = msg.get("tool_calls").cloned();
2508 return (content, tc);
2509 }
2510
2511 (Value::String(text.to_owned()), None)
2513}
2514
2515fn openai_response_json(resp: &ModelResponse) -> Value {
2518 let (content_text, tool_calls) = extract_openai_content_and_tools(&resp.raw, &resp.text);
2519 let finish_reason = if tool_calls.is_some() {
2520 "tool_calls"
2521 } else {
2522 "stop"
2523 };
2524 let mut message = serde_json::json!({
2525 "role": "assistant",
2526 "content": content_text,
2527 });
2528 if let Some(tc) = tool_calls {
2529 message["tool_calls"] = tc;
2530 }
2531 serde_json::json!({
2532 "id": format!("chatcmpl-{}", Uuid::now_v7()),
2533 "object": "chat.completion",
2534 "created": jiff::Timestamp::now().as_second(),
2535 "model": resp.model,
2536 "choices": [{
2537 "index": 0,
2538 "message": message,
2539 "finish_reason": finish_reason,
2540 }],
2541 "usage": {
2542 "prompt_tokens": resp.in_tokens,
2543 "completion_tokens": resp.out_tokens,
2544 "total_tokens": resp.in_tokens + resp.out_tokens,
2545 }
2546 })
2547}
2548
2549fn openai_sse_from_message(message: &Value) -> String {
2553 let id = message
2554 .get("id")
2555 .and_then(Value::as_str)
2556 .unwrap_or("chatcmpl-unknown")
2557 .to_owned();
2558 let created = message.get("created").cloned().unwrap_or(Value::from(0));
2559 let model = message
2560 .get("model")
2561 .and_then(Value::as_str)
2562 .unwrap_or("unknown");
2563
2564 let choices = message.get("choices").and_then(Value::as_array);
2565 let msg = choices
2566 .and_then(|c| c.first())
2567 .and_then(|c| c.get("message"));
2568 let content = msg
2569 .and_then(|m| m.get("content"))
2570 .cloned()
2571 .unwrap_or(Value::Null);
2572 let tool_calls = msg.and_then(|m| m.get("tool_calls")).cloned();
2573 let finish_reason = choices
2574 .and_then(|c| c.first())
2575 .and_then(|c| c.get("finish_reason"))
2576 .cloned()
2577 .unwrap_or_else(|| Value::String("stop".to_owned()));
2578
2579 let mut out = String::new();
2580 let chunk = |delta: Value| {
2581 serde_json::json!({
2582 "id": id,
2583 "object": "chat.completion.chunk",
2584 "created": created,
2585 "model": model,
2586 "choices": [{ "index": 0, "delta": delta, "finish_reason": Value::Null }]
2587 })
2588 };
2589
2590 let role_chunk = chunk(serde_json::json!({ "role": "assistant", "content": "" }));
2592 out.push_str("data: ");
2593 out.push_str(&role_chunk.to_string());
2594 out.push_str("\n\n");
2595
2596 if let Value::String(text) = &content
2598 && !text.is_empty()
2599 {
2600 let content_chunk = chunk(serde_json::json!({ "content": text }));
2601 out.push_str("data: ");
2602 out.push_str(&content_chunk.to_string());
2603 out.push_str("\n\n");
2604 }
2605
2606 if let Some(tc) = tool_calls {
2608 let tc_chunk = chunk(serde_json::json!({ "tool_calls": tc }));
2609 out.push_str("data: ");
2610 out.push_str(&tc_chunk.to_string());
2611 out.push_str("\n\n");
2612 }
2613
2614 let finish_chunk = serde_json::json!({
2616 "id": id,
2617 "object": "chat.completion.chunk",
2618 "created": created,
2619 "model": model,
2620 "choices": [{ "index": 0, "delta": {}, "finish_reason": finish_reason }]
2621 });
2622 out.push_str("data: ");
2623 out.push_str(&finish_chunk.to_string());
2624 out.push_str("\n\n");
2625
2626 out.push_str("data: [DONE]\n\n");
2627 out
2628}
2629
2630#[allow(clippy::too_many_arguments)]
2635async fn handle_enforce_openai(
2636 state: &AppState,
2637 headers: &HeaderMap,
2638 body: &Bytes,
2639 features: Features,
2640 route: &Route,
2641 route_ix: usize,
2644 session_header: Option<String>,
2645 tenant: String,
2646 routing_mode: RoutingMode,
2647) -> Response {
2648 if is_stream_request(body) {
2649 let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
2650 let (state_c, headers_c, body_c, route_c) =
2651 (state.clone(), headers.clone(), body.clone(), route.clone());
2652 tokio::spawn(async move {
2653 let out = enforce_pipeline_openai(
2654 &state_c,
2655 &headers_c,
2656 &body_c,
2657 features,
2658 &route_c,
2659 route_ix,
2660 session_header,
2661 tenant,
2662 routing_mode,
2663 )
2664 .await;
2665 let _ = tx.send(out);
2666 });
2667 return sse_keepalive_response(rx, openai_sse_from_message);
2668 }
2669 match enforce_pipeline_openai(
2670 state,
2671 headers,
2672 body,
2673 features,
2674 route,
2675 route_ix,
2676 session_header,
2677 tenant,
2678 routing_mode,
2679 )
2680 .await
2681 {
2682 Ok(message) => (axum::http::StatusCode::OK, Json(message)).into_response(),
2683 Err(e) => e.into_response(),
2684 }
2685}
2686
2687async fn observe_passthrough_openai(
2693 state: AppState,
2694 headers: HeaderMap,
2695 body: Bytes,
2696 session_header: Option<String>,
2697 tenant: String,
2698) -> Response {
2699 if is_stream_request(&body) {
2700 return observe_stream_openai(state, headers, body, session_header, tenant).await;
2701 }
2702 let start = Instant::now();
2703 let result = forward_openai(
2704 &state.http,
2705 &state.config.upstream_openai,
2706 &headers,
2707 body.clone(),
2708 )
2709 .await;
2710 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
2711 match result {
2712 Ok((status, resp_headers, resp_body)) => {
2713 spawn_trace(
2714 &state,
2715 body,
2716 Some(resp_body.clone()),
2717 latency_ms,
2718 session_header,
2719 tenant,
2720 );
2721 (status, resp_headers, resp_body).into_response()
2722 }
2723 Err(err) => {
2724 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
2725 err.into_response()
2726 }
2727 }
2728}
2729
2730async fn observe_stream_openai(
2732 state: AppState,
2733 headers: HeaderMap,
2734 body: Bytes,
2735 session_header: Option<String>,
2736 tenant: String,
2737) -> Response {
2738 let start = Instant::now();
2739 let result = forward_openai_streaming(
2740 &state.http,
2741 &state.config.upstream_openai,
2742 &headers,
2743 body.clone(),
2744 )
2745 .await;
2746 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
2747 match result {
2748 Ok((status, resp_headers, response)) => {
2749 spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
2750 let stream_body = Body::from_stream(response.bytes_stream());
2751 (status, resp_headers, stream_body).into_response()
2752 }
2753 Err(err) => {
2754 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
2755 err.into_response()
2756 }
2757 }
2758}
2759
2760async fn chat_completions(
2765 State(state): State<AppState>,
2766 Extension(TenantId(tenant)): Extension<TenantId>,
2767 headers: HeaderMap,
2768 body: Bytes,
2769) -> Response {
2770 let session_header = header_str(&headers, SESSION_HEADER);
2771
2772 if let Some(routing) = state.config.routing.as_ref() {
2773 let features = extract_openai_features(&headers, &body);
2774 if let Some(route) = routing
2775 .route_for(&features)
2776 .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
2777 {
2778 let route_ix = routing
2780 .routes
2781 .iter()
2782 .position(|r| std::ptr::eq(r, route))
2783 .unwrap_or(0);
2784 let route = route.clone();
2785 let routing_mode = resolve_mode(&headers, &route, &state.config);
2787 if routing_mode == RoutingMode::Observe {
2789 return observe_passthrough_openai(state, headers, body, session_header, tenant)
2790 .await;
2791 }
2792 if enforce_can_handle(
2793 &features,
2794 &body,
2795 routing.escalation.enforce_structured,
2796 &route.ladder,
2797 &state.providers,
2798 Dialect::Openai,
2799 ) {
2800 return handle_enforce_openai(
2801 &state,
2802 &headers,
2803 &body,
2804 features,
2805 &route,
2806 route_ix,
2807 session_header,
2808 tenant,
2809 routing_mode,
2810 )
2811 .await;
2812 }
2813 tracing::info!(
2814 "enforce route matched but OpenAI structured request can't be routed faithfully (flag/ladder); serving via observe passthrough"
2815 );
2816 }
2817 }
2818 observe_passthrough_openai(state, headers, body, session_header, tenant).await
2819}
2820
2821#[cfg(test)]
2822mod tests {
2823 use bytes::Bytes;
2824
2825 use super::*;
2826
2827 fn test_config() -> ProxyConfig {
2828 ProxyConfig::from_lookup(|_| None).unwrap()
2829 }
2830
2831 #[test]
2832 fn build_trace_maps_request_and_response_fields() {
2833 let config = test_config();
2834 let req = Bytes::from_static(
2835 br#"{"model":"claude-haiku-4-5","tools":[{"name":"a"}],"messages":[]}"#,
2836 );
2837 let resp = Bytes::from_static(
2838 br#"{"model":"claude-haiku-4-5","usage":{"input_tokens":1200,"output_tokens":300}}"#,
2839 );
2840
2841 let trace = build_trace(&config, &req, &resp, 42, Some("sess-1"));
2842
2843 assert_eq!(trace.request.api, "anthropic.messages");
2844 assert_eq!(trace.session_id, "sess-1");
2845 assert_eq!(trace.attempts.len(), 1);
2846 let attempt = &trace.attempts[0];
2847 assert_eq!(attempt.model, "claude-haiku-4-5");
2848 assert_eq!(attempt.provider, "anthropic");
2849 assert_eq!(attempt.in_tokens, 1200);
2850 assert_eq!(attempt.out_tokens, 300);
2851 assert!(attempt.cost_usd > 0.0);
2852 assert_eq!(trace.request.features.tool_count, 1);
2853 assert!(!trace.request.features.has_images);
2854 assert_eq!(trace.final_.served_rung, Some(0));
2855 }
2856
2857 #[test]
2858 fn build_trace_falls_back_to_trace_id_session_when_header_absent() {
2859 let config = test_config();
2860 let req = Bytes::from_static(b"{}");
2861 let resp = Bytes::from_static(b"{}");
2862
2863 let trace = build_trace(&config, &req, &resp, 1, None);
2864
2865 assert_eq!(trace.session_id, trace.trace_id.to_string());
2866 }
2867
2868 #[test]
2869 fn build_error_trace_has_no_attempts_and_served_from_error() {
2870 let config = test_config();
2871 let req = Bytes::from_static(br#"{"model":"claude-haiku-4-5"}"#);
2872
2873 let trace = build_error_trace(&config, &req, 7, None);
2874
2875 assert!(trace.attempts.is_empty());
2876 assert_eq!(trace.final_.served_from, ServedFrom::Error);
2877 assert_eq!(trace.final_.served_rung, None);
2878 }
2879
2880 #[test]
2881 fn message_with_image_block_sets_has_images() {
2882 let req = br#"{"messages":[{"role":"user","content":[{"type":"image"}]}]}"#;
2883 let (_, _, has_images) = request_features(req);
2884 assert!(has_images);
2885 }
2886
2887 #[test]
2888 fn prompt_hash_never_contains_raw_prompt_text() {
2889 let hash = prompt_hash("salt", b"super secret prompt");
2890 assert!(!hash.contains("secret"));
2891 assert_eq!(hash.len(), 64);
2892 }
2893
2894 #[test]
2895 fn parse_model_request_preserves_content_verbatim_and_projects_text() {
2896 let body = br#"{"model":"m","system":"sys","max_tokens":50,
2897 "messages":[{"role":"user","content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]},
2898 {"role":"assistant","content":"c"}]}"#;
2899 let req = parse_model_request(body).unwrap();
2900 assert_eq!(req.system.as_deref(), Some("sys"));
2901 assert_eq!(req.max_tokens, 50);
2902 assert_eq!(req.messages.len(), 2);
2903 assert_eq!(
2905 req.messages[0].content,
2906 serde_json::json!([{"type":"text","text":"a"},{"type":"text","text":"b"}])
2907 );
2908 assert_eq!(req.messages[1].content, Value::String("c".to_owned()));
2910 assert_eq!(req.messages[0].text_view(), "a\nb");
2912 assert_eq!(req.messages[1].text_view(), "c");
2913 }
2914
2915 #[test]
2916 fn tool_and_image_blocks_survive_the_request_round_trip() {
2917 let body = br#"{"model":"m","max_tokens":50,"messages":[
2919 {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
2920 {"role":"user","content":[
2921 {"type":"tool_result","tool_use_id":"t1","content":"2"},
2922 {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
2923 ]}]}"#;
2924 let req = parse_model_request(body).unwrap();
2925 let round_tripped = serde_json::to_value(&req.messages).unwrap();
2926 assert_eq!(
2927 round_tripped,
2928 serde_json::json!([
2929 {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
2930 {"role":"user","content":[
2931 {"type":"tool_result","tool_use_id":"t1","content":"2"},
2932 {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
2933 ]}
2934 ])
2935 );
2936 }
2937
2938 #[test]
2939 fn text_message_serializes_byte_identical_to_a_plain_string() {
2940 let m = ChatMessage::text("user", "hello");
2942 assert_eq!(
2943 serde_json::to_string(&m).unwrap(),
2944 r#"{"role":"user","content":"hello"}"#
2945 );
2946 }
2947
2948 #[test]
2949 fn parse_model_request_rejects_non_message_bodies() {
2950 assert!(parse_model_request(b"not json").is_none());
2951 assert!(parse_model_request(br#"{"no":"messages"}"#).is_none());
2952 }
2953
2954 use crate::provider::{MockProvider, ModelResponse, Provider, ProviderError, ProviderRegistry};
2957 use axum::extract::State;
2958 use std::collections::HashMap;
2959 use std::sync::Arc;
2960 use tokio::sync::mpsc;
2961
2962 fn model_resp(model: &str, text: &str) -> ModelResponse {
2963 ModelResponse {
2964 model: model.to_owned(),
2965 text: text.to_owned(),
2966 in_tokens: 1000,
2967 out_tokens: 400,
2968 raw: serde_json::Value::Null,
2969 }
2970 }
2971
2972 fn enforce_state(
2975 ladder: &[&str],
2976 gates: &[&str],
2977 outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
2978 ) -> (AppState, mpsc::Receiver<Trace>) {
2979 let toml = format!(
2980 "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [{}]\ngates = [{}]\n",
2981 ladder
2982 .iter()
2983 .map(|m| format!("\"{m}\""))
2984 .collect::<Vec<_>>()
2985 .join(", "),
2986 gates
2987 .iter()
2988 .map(|g| format!("\"{g}\""))
2989 .collect::<Vec<_>>()
2990 .join(", "),
2991 );
2992 let config = ProxyConfig::from_lookup(|k| match k {
2993 "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
2994 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
2995 _ => None,
2996 })
2997 .unwrap();
2998
2999 let mut outs = HashMap::new();
3000 for (model, out) in outcomes {
3001 outs.insert(model.to_owned(), out);
3002 }
3003 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3004 map.insert(
3005 "anthropic".to_owned(),
3006 Arc::new(MockProvider::new("anthropic", outs)),
3007 );
3008 let providers = ProviderRegistry::from_map(map);
3009
3010 let (traces, rx) = mpsc::channel(64);
3011 let state = AppState {
3012 config: Arc::new(config),
3013 http: reqwest::Client::new(),
3014 providers,
3015 gate_health: Arc::new(GateHealthRegistry::new()),
3016 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3017 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3018 traces,
3019 adaptive: None,
3020 bandit: None,
3021 predictor: None,
3022 tenant_rate_limiter: None,
3023 spill: None,
3024 };
3025 (state, rx)
3026 }
3027
3028 fn user_body() -> Bytes {
3029 Bytes::from_static(
3030 br#"{"model":"ignored","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}"#,
3031 )
3032 }
3033
3034 async fn body_json(resp: Response) -> Value {
3035 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
3036 .await
3037 .unwrap();
3038 serde_json::from_slice(&bytes).unwrap()
3039 }
3040
3041 #[tokio::test]
3042 async fn mode_profile_stamped_on_trace_when_non_balanced() {
3043 let (state, mut rx) = enforce_state(
3044 &["anthropic/claude-haiku-4-5"],
3045 &["non-empty"],
3046 vec![(
3047 "anthropic/claude-haiku-4-5",
3048 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3049 )],
3050 );
3051 let mut headers = HeaderMap::new();
3052 headers.insert("x-firstpass-mode", "quality".parse().unwrap());
3053 let resp = messages(
3054 State(state),
3055 Extension(TenantId("default".to_owned())),
3056 headers,
3057 user_body(),
3058 )
3059 .await;
3060 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3061 let trace = rx.try_recv().expect("trace enqueued");
3062 assert_eq!(
3063 trace.policy.mode_profile.as_deref(),
3064 Some("quality"),
3065 "mode_profile must be stamped when quality mode is active"
3066 );
3067 }
3068
3069 #[tokio::test]
3070 async fn mode_profile_absent_from_trace_when_balanced() {
3071 let (state, mut rx) = enforce_state(
3072 &["anthropic/claude-haiku-4-5"],
3073 &["non-empty"],
3074 vec![(
3075 "anthropic/claude-haiku-4-5",
3076 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3077 )],
3078 );
3079 let resp = messages(
3081 State(state),
3082 Extension(TenantId("default".to_owned())),
3083 HeaderMap::new(),
3084 user_body(),
3085 )
3086 .await;
3087 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3088 let trace = rx.try_recv().expect("trace enqueued");
3089 assert!(
3090 trace.policy.mode_profile.is_none(),
3091 "mode_profile must be absent when Balanced (byte-identical invariant)"
3092 );
3093 }
3094
3095 #[tokio::test]
3096 async fn enforce_serves_first_pass_and_returns_anthropic_shape() {
3097 let (state, mut rx) = enforce_state(
3098 &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
3099 &["non-empty"],
3100 vec![(
3101 "anthropic/claude-haiku-4-5",
3102 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3103 )],
3104 );
3105 let resp = messages(
3106 State(state),
3107 Extension(TenantId("default".to_owned())),
3108 HeaderMap::new(),
3109 user_body(),
3110 )
3111 .await;
3112 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3113 let json = body_json(resp).await;
3114 assert_eq!(json["type"], "message");
3115 assert_eq!(json["content"][0]["text"], "hello");
3116 assert_eq!(json["model"], "anthropic/claude-haiku-4-5");
3117
3118 let trace = rx.try_recv().expect("a trace was enqueued");
3119 assert_eq!(trace.mode, Mode::Enforce);
3120 assert_eq!(trace.final_.served_rung, Some(0));
3121 assert_eq!(trace.attempts.len(), 1);
3122 }
3123
3124 #[tokio::test]
3125 async fn enforce_escalates_then_serves_and_traces_two_attempts() {
3126 let (state, mut rx) = enforce_state(
3127 &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
3128 &["non-empty"],
3129 vec![
3130 (
3131 "anthropic/claude-haiku-4-5",
3132 Ok(model_resp("anthropic/claude-haiku-4-5", " ")),
3133 ), (
3135 "anthropic/claude-sonnet-5",
3136 Ok(model_resp("anthropic/claude-sonnet-5", "answer")),
3137 ),
3138 ],
3139 );
3140 let resp = messages(
3141 State(state),
3142 Extension(TenantId("default".to_owned())),
3143 HeaderMap::new(),
3144 user_body(),
3145 )
3146 .await;
3147 let json = body_json(resp).await;
3148 assert_eq!(json["content"][0]["text"], "answer");
3149
3150 let trace = rx.try_recv().expect("trace enqueued");
3151 assert_eq!(trace.attempts.len(), 2);
3152 assert_eq!(trace.final_.escalations, 1);
3153 assert_eq!(trace.final_.served_rung, Some(1));
3154 }
3155
3156 #[tokio::test]
3157 async fn enforce_all_rungs_error_returns_502() {
3158 let (state, mut rx) = enforce_state(
3159 &["anthropic/claude-haiku-4-5"],
3160 &["non-empty"],
3161 vec![(
3162 "anthropic/claude-haiku-4-5",
3163 Err(ProviderError::Transport("down".into())),
3164 )],
3165 );
3166 let resp = messages(
3167 State(state),
3168 Extension(TenantId("default".to_owned())),
3169 HeaderMap::new(),
3170 user_body(),
3171 )
3172 .await;
3173 assert_eq!(resp.status(), axum::http::StatusCode::BAD_GATEWAY);
3174 assert!(rx.try_recv().is_ok());
3176 }
3177
3178 #[tokio::test]
3179 async fn no_routing_config_falls_through_to_observe_not_enforce() {
3180 let config = ProxyConfig::from_lookup(|k| match k {
3184 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3185 _ => None,
3186 })
3187 .unwrap();
3188 let (traces, _rx) = mpsc::channel(64);
3189 let state = AppState {
3190 config: Arc::new(config),
3191 http: reqwest::Client::new(),
3192 providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
3193 gate_health: Arc::new(GateHealthRegistry::new()),
3194 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3195 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3196 traces,
3197 adaptive: None,
3198 bandit: None,
3199 predictor: None,
3200 tenant_rate_limiter: None,
3201 spill: None,
3202 };
3203 let resp = messages(
3204 State(state),
3205 Extension(TenantId("default".to_owned())),
3206 HeaderMap::new(),
3207 user_body(),
3208 )
3209 .await;
3210 assert_ne!(resp.status(), axum::http::StatusCode::OK);
3212 }
3213
3214 fn bare_enforce_route() -> Route {
3217 use firstpass_core::config::{Match, Mode};
3218 Route {
3219 match_: Match::default(),
3220 mode: Mode::Enforce,
3221 ladder: vec!["anthropic/claude-haiku-4-5".to_owned()],
3222 gates: vec![],
3223 deferred_gates: vec![],
3224 routing_mode: None,
3225 rollout: None,
3226 shadow: None,
3227 }
3228 }
3229
3230 #[test]
3231 fn balanced_preset_application_is_noop() {
3232 let base_max_rungs = 3u32;
3234 let base_speculation = 2u32;
3235 let preset = RoutingMode::Balanced.preset();
3236 let max_rungs = if let Some(d) = preset.max_rungs_delta {
3237 (base_max_rungs as i32 + d).max(1) as u32
3238 } else {
3239 base_max_rungs
3240 };
3241 let speculation = preset.speculation.unwrap_or(base_speculation);
3242 let start_at_top = preset.start_at_top;
3243 assert_eq!(max_rungs, 3, "Balanced must not change max_rungs");
3244 assert_eq!(speculation, 2, "Balanced must not change speculation");
3245 assert!(!start_at_top, "Balanced must not set start_at_top");
3246 }
3247
3248 #[test]
3249 fn resolve_mode_header_wins_over_route_and_global() {
3250 let mut headers = HeaderMap::new();
3251 headers.insert("x-firstpass-mode", "cost".parse().unwrap());
3252 let mut route = bare_enforce_route();
3253 route.routing_mode = Some(RoutingMode::Quality); let mut config = test_config();
3255 config.default_routing_mode = RoutingMode::Max; assert_eq!(
3257 resolve_mode(&headers, &route, &config),
3258 RoutingMode::Cost,
3259 "header must win"
3260 );
3261 }
3262
3263 #[test]
3264 fn resolve_mode_route_wins_over_global_when_no_header() {
3265 let mut route = bare_enforce_route();
3266 route.routing_mode = Some(RoutingMode::Latency);
3267 let mut config = test_config();
3268 config.default_routing_mode = RoutingMode::Max;
3269 assert_eq!(
3270 resolve_mode(&HeaderMap::new(), &route, &config),
3271 RoutingMode::Latency,
3272 "route must beat global default"
3273 );
3274 }
3275
3276 #[test]
3277 fn resolve_mode_global_when_header_and_route_absent() {
3278 let mut config = test_config();
3279 config.default_routing_mode = RoutingMode::Quality;
3280 assert_eq!(
3281 resolve_mode(&HeaderMap::new(), &bare_enforce_route(), &config),
3282 RoutingMode::Quality
3283 );
3284 }
3285
3286 #[test]
3287 fn resolve_mode_unknown_header_falls_through_to_route() {
3288 let mut headers = HeaderMap::new();
3289 headers.insert("x-firstpass-mode", "turbo-mode".parse().unwrap());
3291 let mut route = bare_enforce_route();
3292 route.routing_mode = Some(RoutingMode::Cost);
3293 let config = test_config();
3294 assert_eq!(
3295 resolve_mode(&headers, &route, &config),
3296 RoutingMode::Cost,
3297 "unknown header value must fall through to route"
3298 );
3299 }
3300
3301 #[test]
3302 fn resolve_mode_header_case_insensitive() {
3303 let mut headers = HeaderMap::new();
3304 headers.insert("x-firstpass-mode", "QUALITY".parse().unwrap());
3305 assert_eq!(
3306 resolve_mode(&headers, &bare_enforce_route(), &test_config()),
3307 RoutingMode::Quality
3308 );
3309 }
3310
3311 #[test]
3312 fn resolve_mode_no_mode_set_returns_balanced() {
3313 assert_eq!(
3315 resolve_mode(&HeaderMap::new(), &bare_enforce_route(), &test_config()),
3316 RoutingMode::Balanced
3317 );
3318 }
3319
3320 #[test]
3321 fn capabilities_json_includes_routing_modes() {
3322 let modes: Vec<&'static str> = RoutingMode::ALL.iter().map(|m| m.as_str()).collect();
3323 assert!(modes.contains(&"balanced"));
3324 assert!(modes.contains(&"cost"));
3325 assert!(modes.contains(&"quality"));
3326 assert!(modes.contains(&"latency"));
3327 assert!(modes.contains(&"max"));
3328 assert!(modes.contains(&"observe"));
3329 }
3330
3331 #[test]
3332 fn detects_stream_requests() {
3333 assert!(is_stream_request(br#"{"stream": true}"#));
3334 assert!(!is_stream_request(br#"{"stream": false}"#));
3335 assert!(!is_stream_request(br#"{"model":"m"}"#));
3336 assert!(!is_stream_request(b"not json"));
3337 }
3338
3339 #[test]
3340 fn detects_tool_blocks_in_messages() {
3341 let with =
3342 br#"{"messages":[{"role":"user","content":[{"type":"tool_result","content":"42"}]}]}"#;
3343 let without = br#"{"messages":[{"role":"user","content":"hi"}]}"#;
3344 assert!(messages_have_tool_blocks(with));
3345 assert!(!messages_have_tool_blocks(without));
3346 }
3347
3348 #[test]
3349 fn enforce_only_handles_plain_text() {
3350 let plain =
3351 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
3352 let tools = Bytes::from_static(
3353 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3354 );
3355 let f_plain = extract_features(&HeaderMap::new(), &plain);
3356 let f_tools = extract_features(&HeaderMap::new(), &tools);
3357 let anthropic_ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
3358 let providers = test_registry();
3359 assert!(enforce_can_handle(
3361 &f_plain,
3362 &plain,
3363 false,
3364 &anthropic_ladder,
3365 &providers,
3366 Dialect::Anthropic,
3367 ));
3368 assert!(!enforce_can_handle(
3369 &f_tools,
3370 &tools,
3371 false,
3372 &anthropic_ladder,
3373 &providers,
3374 Dialect::Anthropic,
3375 ));
3376 }
3377
3378 #[test]
3379 fn structured_enforce_routes_tools_and_streaming() {
3380 let tools = Bytes::from_static(
3383 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3384 );
3385 let streaming_tools = Bytes::from_static(
3386 br#"{"model":"m","stream":true,"tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3387 );
3388 let f = extract_features(&HeaderMap::new(), &tools);
3389 let anthropic_ladder = vec![
3390 "anthropic/claude-haiku-4-5".to_owned(),
3391 "anthropic/claude-sonnet-5".to_owned(),
3392 ];
3393 let providers = test_registry();
3394 assert!(enforce_can_handle(
3395 &f,
3396 &tools,
3397 true,
3398 &anthropic_ladder,
3399 &providers,
3400 Dialect::Anthropic,
3401 ));
3402 assert!(enforce_can_handle(
3403 &f,
3404 &streaming_tools,
3405 true,
3406 &anthropic_ladder,
3407 &providers,
3408 Dialect::Anthropic,
3409 ));
3410 }
3411
3412 fn test_registry() -> crate::provider::ProviderRegistry {
3414 crate::provider::ProviderRegistry::new("http://localhost", "http://localhost")
3415 }
3416
3417 #[test]
3418 fn fidelity_guard_blocks_structured_on_non_verbatim_ladder() {
3419 let tools = Bytes::from_static(
3423 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
3424 );
3425 let plain =
3426 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
3427 let f_tools = extract_features(&HeaderMap::new(), &tools);
3428 let f_plain = extract_features(&HeaderMap::new(), &plain);
3429 let providers = test_registry();
3430 let mixed_ladder = vec![
3431 "openai/gpt-4.1-mini".to_owned(),
3432 "anthropic/claude-sonnet-5".to_owned(),
3433 ];
3434 assert!(!enforce_can_handle(
3435 &f_tools,
3436 &tools,
3437 true,
3438 &mixed_ladder,
3439 &providers,
3440 Dialect::Anthropic,
3441 ));
3442 assert!(enforce_can_handle(
3443 &f_plain,
3444 &plain,
3445 true,
3446 &mixed_ladder,
3447 &providers,
3448 Dialect::Anthropic,
3449 ));
3450 }
3451
3452 #[test]
3453 fn enforce_sse_reemission_preserves_text_and_tool_use() {
3454 let resp = ModelResponse {
3457 model: "anthropic/claude-haiku-4-5".to_owned(),
3458 text: "let me check".to_owned(),
3459 in_tokens: 5,
3460 out_tokens: 7,
3461 raw: serde_json::json!({
3462 "content": [
3463 { "type": "text", "text": "let me check" },
3464 { "type": "tool_use", "id": "tu_1", "name": "get_weather", "input": { "city": "Paris" } }
3465 ]
3466 }),
3467 };
3468 let sse = anthropic_sse_from_message(&anthropic_response_json(&resp));
3469
3470 let frames: Vec<Value> = sse
3472 .lines()
3473 .filter_map(|l| l.strip_prefix("data: "))
3474 .map(|d| serde_json::from_str::<Value>(d).expect("each SSE data frame is valid JSON"))
3475 .collect();
3476
3477 assert_eq!(frames.first().unwrap()["type"], "message_start");
3479 assert_eq!(frames.last().unwrap()["type"], "message_stop");
3480 assert!(frames.iter().any(|f| f["delta"]["type"] == "text_delta"
3482 && f["delta"]["text"] == "let me check"));
3483 assert!(
3486 frames
3487 .iter()
3488 .any(|f| f["content_block"]["type"] == "tool_use"
3489 && f["content_block"]["name"] == "get_weather"
3490 && f["content_block"]["id"] == "tu_1")
3491 );
3492 assert!(
3493 frames
3494 .iter()
3495 .any(|f| f["delta"]["type"] == "input_json_delta"
3496 && f["delta"]["partial_json"] == r#"{"city":"Paris"}"#)
3497 );
3498 }
3499
3500 #[tokio::test]
3505 async fn enforce_falls_back_to_observe_for_tool_requests() {
3506 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
3507 let config = ProxyConfig::from_lookup(|k| match k {
3508 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
3509 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
3510 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3511 _ => None,
3512 })
3513 .unwrap();
3514 let mut outs = HashMap::new();
3515 outs.insert(
3516 "anthropic/m".to_owned(),
3517 Ok(model_resp("anthropic/m", "hello")),
3518 );
3519 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3520 map.insert(
3521 "anthropic".to_owned(),
3522 Arc::new(MockProvider::new("anthropic", outs)),
3523 );
3524 let (traces, _rx) = mpsc::channel(64);
3525 let state = AppState {
3526 config: Arc::new(config),
3527 http: reqwest::Client::new(),
3528 providers: ProviderRegistry::from_map(map),
3529 gate_health: Arc::new(GateHealthRegistry::new()),
3530 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3531 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3532 traces,
3533 adaptive: None,
3534 bandit: None,
3535 predictor: None,
3536 tenant_rate_limiter: None,
3537 spill: None,
3538 };
3539
3540 let plain =
3542 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
3543 let resp = messages(
3544 State(state.clone()),
3545 Extension(TenantId("default".to_owned())),
3546 HeaderMap::new(),
3547 plain,
3548 )
3549 .await;
3550 assert_eq!(
3551 resp.status(),
3552 axum::http::StatusCode::OK,
3553 "plain text should enforce"
3554 );
3555
3556 let tools = Bytes::from_static(
3559 br#"{"model":"m","tools":[{"name":"get_weather"}],"messages":[{"role":"user","content":"hi"}]}"#,
3560 );
3561 let resp = messages(
3562 State(state.clone()),
3563 Extension(TenantId("default".to_owned())),
3564 HeaderMap::new(),
3565 tools.clone(),
3566 )
3567 .await;
3568 assert_eq!(
3569 resp.status(),
3570 axum::http::StatusCode::OK,
3571 "tool request must route through enforce by default (ADR 0005 default-on)"
3572 );
3573
3574 let toml_off = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n[escalation]\nenforce_structured = false\n";
3577 let config_off = ProxyConfig::from_lookup(|k| match k {
3578 "FIRSTPASS_CONFIG_TOML" => Some(toml_off.to_owned()),
3579 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
3580 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3581 _ => None,
3582 })
3583 .unwrap();
3584 let state_off = AppState {
3585 config: Arc::new(config_off),
3586 ..state.clone()
3587 };
3588 let resp = messages(
3589 State(state_off),
3590 Extension(TenantId("default".to_owned())),
3591 HeaderMap::new(),
3592 tools,
3593 )
3594 .await;
3595 assert_ne!(
3596 resp.status(),
3597 axum::http::StatusCode::OK,
3598 "with enforce_structured = false a tool request must fall back to observe"
3599 );
3600
3601 let toolres = Bytes::from_static(
3603 br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"x","content":"42"}]}]}"#,
3604 );
3605 let resp = messages(
3606 State(state),
3607 Extension(TenantId("default".to_owned())),
3608 HeaderMap::new(),
3609 toolres,
3610 )
3611 .await;
3612 assert_eq!(
3613 resp.status(),
3614 axum::http::StatusCode::OK,
3615 "tool_result blocks route through enforce by default too (verbatim carry)"
3616 );
3617 }
3618
3619 async fn feedback_state() -> (AppState, std::path::PathBuf, String) {
3623 let db = std::env::temp_dir().join(format!("firstpass-feedback-{}.db", Uuid::now_v7()));
3624 let (tx, handle) = crate::store::open(&db).unwrap();
3625
3626 let mut trace = build_error_trace(
3627 &ProxyConfig::from_lookup(|_| None).unwrap(),
3628 &Bytes::from_static(b"{}"),
3629 5,
3630 Some("sess-fb"),
3631 );
3632 trace.attempts.push(Attempt {
3633 rung: 0,
3634 model: "anthropic/claude-haiku-4-5".into(),
3635 provider: "anthropic".into(),
3636 in_tokens: 10,
3637 out_tokens: 5,
3638 cost_usd: 0.001,
3639 latency_ms: 5,
3640 gates: vec![],
3641 verdict: Verdict::Pass,
3642 });
3643 let trace_id = trace.trace_id.to_string();
3644 tx.try_send(trace).unwrap();
3645 drop(tx);
3646 handle.await.unwrap();
3647
3648 let db_str = db.to_string_lossy().into_owned();
3649 let config = ProxyConfig::from_lookup(move |k| match k {
3650 "FIRSTPASS_DB" => Some(db_str.clone()),
3651 _ => None,
3652 })
3653 .unwrap();
3654 let (traces, _rx) = mpsc::channel(64);
3655 let state = AppState {
3656 config: Arc::new(config),
3657 http: reqwest::Client::new(),
3658 providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
3659 gate_health: Arc::new(GateHealthRegistry::new()),
3660 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3661 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3662 traces,
3663 adaptive: None,
3664 bandit: None,
3665 predictor: None,
3666 tenant_rate_limiter: None,
3667 spill: None,
3668 };
3669 (state, db, trace_id)
3670 }
3671
3672 #[tokio::test]
3673 async fn feedback_nudges_the_adaptive_threshold() {
3674 use firstpass_core::conformal::AdaptiveConformal;
3675 let (mut state, _db, trace_id) = feedback_state().await;
3676 let aci = Arc::new(std::sync::Mutex::new(AdaptiveConformal::new(0.1, 0.2, 0.5)));
3677 state.adaptive = Some(aci.clone());
3678 let before = aci.lock().unwrap().threshold();
3679
3680 let fail = Bytes::from(
3682 serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "fail", "reporter": "ci" })
3683 .to_string(),
3684 );
3685 assert_eq!(
3686 feedback(
3687 State(state.clone()),
3688 Extension(TenantId("default".to_owned())),
3689 fail
3690 )
3691 .await
3692 .status(),
3693 axum::http::StatusCode::ACCEPTED
3694 );
3695 let after_fail = aci.lock().unwrap().threshold();
3696 assert!(
3697 after_fail > before,
3698 "served fail should raise the live threshold: {before} -> {after_fail}"
3699 );
3700
3701 let pass = Bytes::from(
3703 serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "pass", "reporter": "ci" })
3704 .to_string(),
3705 );
3706 let _ = feedback(
3707 State(state),
3708 Extension(TenantId("default".to_owned())),
3709 pass,
3710 )
3711 .await;
3712 assert!(aci.lock().unwrap().threshold() < after_fail);
3713 }
3714
3715 #[tokio::test]
3716 async fn feedback_records_a_deferred_verdict_without_breaking_the_chain() {
3717 let (state, db, trace_id) = feedback_state().await;
3718 let body = Bytes::from(
3719 serde_json::json!({
3720 "trace_id": trace_id,
3721 "gate_id": "tests",
3722 "verdict": "pass",
3723 "score": 1.0,
3724 "reporter": "ci",
3725 })
3726 .to_string(),
3727 );
3728 let resp = feedback(
3729 State(state),
3730 Extension(TenantId("default".to_owned())),
3731 body,
3732 )
3733 .await;
3734 assert_eq!(resp.status(), axum::http::StatusCode::ACCEPTED);
3735
3736 let view = crate::store::load_trace_view(&db, "default", &trace_id)
3738 .unwrap()
3739 .unwrap();
3740 assert_eq!(view.deferred.len(), 1);
3741 assert_eq!(view.deferred[0].gate_id, "tests");
3742 let traces = crate::store::load_all_traces(&db).unwrap();
3744 firstpass_core::verify_chain(&traces, GENESIS_HASH).unwrap();
3745
3746 let _ = std::fs::remove_file(&db);
3747 }
3748
3749 #[tokio::test]
3750 async fn feedback_for_unknown_trace_is_404() {
3751 let (state, db, _trace_id) = feedback_state().await;
3752 let body = Bytes::from(
3753 serde_json::json!({
3754 "trace_id": "does-not-exist",
3755 "gate_id": "tests",
3756 "verdict": "pass",
3757 "reporter": "ci",
3758 })
3759 .to_string(),
3760 );
3761 let resp = feedback(
3762 State(state),
3763 Extension(TenantId("default".to_owned())),
3764 body,
3765 )
3766 .await;
3767 assert_eq!(resp.status(), axum::http::StatusCode::NOT_FOUND);
3768 let _ = std::fs::remove_file(&db);
3769 }
3770
3771 #[tokio::test]
3774 async fn feedback_across_tenants_is_404_not_403() {
3775 let (state, db, trace_id) = feedback_state().await;
3776 let body = Bytes::from(
3777 serde_json::json!({
3778 "trace_id": trace_id,
3779 "gate_id": "tests",
3780 "verdict": "pass",
3781 "score": 1.0,
3782 "reporter": "attacker",
3783 })
3784 .to_string(),
3785 );
3786 let resp = feedback(
3788 State(state),
3789 Extension(TenantId("tenant-b".to_owned())),
3790 body,
3791 )
3792 .await;
3793 assert_eq!(
3794 resp.status(),
3795 axum::http::StatusCode::NOT_FOUND,
3796 "cross-tenant feedback must look exactly like a missing trace"
3797 );
3798 let _ = std::fs::remove_file(&db);
3799 }
3800
3801 #[tokio::test]
3802 async fn feedback_rejects_bad_verdict_and_score() {
3803 let (state, db, trace_id) = feedback_state().await;
3804 let bad_verdict = Bytes::from(
3805 serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "maybe", "reporter": "x" })
3806 .to_string(),
3807 );
3808 assert_eq!(
3809 feedback(
3810 State(state.clone()),
3811 Extension(TenantId("default".to_owned())),
3812 bad_verdict
3813 )
3814 .await
3815 .status(),
3816 axum::http::StatusCode::BAD_REQUEST
3817 );
3818 let bad_score = Bytes::from(
3819 serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "pass", "score": 9.0, "reporter": "x" })
3820 .to_string(),
3821 );
3822 assert_eq!(
3823 feedback(
3824 State(state),
3825 Extension(TenantId("default".to_owned())),
3826 bad_score
3827 )
3828 .await
3829 .status(),
3830 axum::http::StatusCode::BAD_REQUEST
3831 );
3832 let _ = std::fs::remove_file(&db);
3833 }
3834
3835 #[tokio::test]
3836 async fn metrics_endpoint_renders_after_a_real_request() {
3837 use tower::ServiceExt;
3838
3839 let (state, mut rx) = enforce_state(
3840 &["anthropic/claude-haiku-4-5"],
3841 &["non-empty"],
3842 vec![(
3843 "anthropic/claude-haiku-4-5",
3844 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3845 )],
3846 );
3847 let router = app(state).expect("prometheus recorder installs");
3848
3849 let req = axum::http::Request::builder()
3850 .method("POST")
3851 .uri("/v1/messages")
3852 .header("content-type", "application/json")
3853 .body(Body::from(user_body()))
3854 .unwrap();
3855 let resp = router.clone().oneshot(req).await.unwrap();
3856 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3857 rx.try_recv().expect("a trace was enqueued");
3858
3859 let metrics_req = axum::http::Request::builder()
3860 .method("GET")
3861 .uri("/metrics")
3862 .body(Body::empty())
3863 .unwrap();
3864 let metrics_resp = router.oneshot(metrics_req).await.unwrap();
3865 assert_eq!(metrics_resp.status(), axum::http::StatusCode::OK);
3866 let bytes = axum::body::to_bytes(metrics_resp.into_body(), 1 << 20)
3867 .await
3868 .unwrap();
3869 let body = String::from_utf8(bytes.to_vec()).unwrap();
3870 assert!(
3871 body.contains("firstpass_enforce_latency_ms"),
3872 "metrics body missing enforce latency histogram: {body}"
3873 );
3874 assert!(
3875 body.contains("firstpass_served_total"),
3876 "metrics body missing served counter: {body}"
3877 );
3878 }
3879
3880 fn auth_state(require_auth: bool, keys_json: Option<String>) -> AppState {
3884 auth_state_rated(require_auth, keys_json, None)
3885 }
3886
3887 fn auth_state_rated(
3890 require_auth: bool,
3891 keys_json: Option<String>,
3892 rate_per_sec: Option<u32>,
3893 ) -> AppState {
3894 let config = ProxyConfig::from_lookup(|k| match k {
3895 "FIRSTPASS_REQUIRE_AUTH" => require_auth.then(|| "true".to_owned()),
3896 "FIRSTPASS_TENANT_KEYS_JSON" => keys_json.clone(),
3897 "FIRSTPASS_TENANT_RATE_PER_SEC" => rate_per_sec.map(|n| n.to_string()),
3898 _ => None,
3899 })
3900 .unwrap();
3901 let (traces, _rx) = mpsc::channel(64);
3902 std::mem::forget(_rx);
3905 let providers: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3906 let tenant_rate_limiter = build_tenant_rate_limiter(&config);
3907 AppState {
3908 config: Arc::new(config),
3909 http: reqwest::Client::new(),
3910 providers: ProviderRegistry::from_map(providers),
3911 gate_health: Arc::new(GateHealthRegistry::new()),
3912 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
3913 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
3914 traces,
3915 adaptive: None,
3916 bandit: None,
3917 predictor: None,
3918 tenant_rate_limiter,
3919 spill: None,
3920 }
3921 }
3922
3923 fn cap_request(auth_header: Option<&str>) -> axum::http::Request<Body> {
3924 let mut b = axum::http::Request::builder()
3925 .method("GET")
3926 .uri("/v1/capabilities");
3927 if let Some(h) = auth_header {
3928 b = b.header("authorization", h);
3929 }
3930 b.body(Body::empty()).unwrap()
3931 }
3932
3933 #[tokio::test]
3934 async fn auth_off_allows_unauthenticated_request() {
3935 use tower::ServiceExt;
3936 let router = app(auth_state(false, None)).expect("router");
3937 let resp = router.oneshot(cap_request(None)).await.unwrap();
3938 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3940 }
3941
3942 #[tokio::test]
3943 async fn auth_on_missing_key_is_401_opaque() {
3944 use tower::ServiceExt;
3945 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3946 let keys = format!("{{\"tenant-a\": {hash:?}}}");
3947 let router = app(auth_state(true, Some(keys))).expect("router");
3948
3949 let resp = router.oneshot(cap_request(None)).await.unwrap();
3950 assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
3951 let json = body_json(resp).await;
3952 assert_eq!(json["error"]["type"], "unauthorized");
3953 let msg = json["error"]["message"].as_str().unwrap();
3955 assert!(!msg.contains("tenant"), "no tenant oracle in body: {msg}");
3956 }
3957
3958 #[tokio::test]
3959 async fn auth_on_invalid_key_is_401() {
3960 use tower::ServiceExt;
3961 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3962 let keys = format!("{{\"tenant-a\": {hash:?}}}");
3963 let router = app(auth_state(true, Some(keys))).expect("router");
3964
3965 let resp = router
3966 .oneshot(cap_request(Some("Bearer wrong-key")))
3967 .await
3968 .unwrap();
3969 assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
3970 }
3971
3972 #[tokio::test]
3973 async fn auth_on_valid_key_proceeds() {
3974 use tower::ServiceExt;
3975 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3976 let keys = format!("{{\"tenant-a\": {hash:?}}}");
3977 let router = app(auth_state(true, Some(keys))).expect("router");
3978
3979 let resp = router
3981 .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
3982 .await
3983 .unwrap();
3984 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3986 }
3987
3988 fn two_tenant_state(rate_per_sec: Option<u32>) -> AppState {
3990 let hash_a = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3991 let hash_b = crate::tenant_auth::TenantKeys::hash_key("key-b").unwrap();
3992 let keys = format!("{{\"tenant-a\": {hash_a:?}, \"tenant-b\": {hash_b:?}}}");
3993 auth_state_rated(true, Some(keys), rate_per_sec)
3994 }
3995
3996 #[tokio::test]
3997 async fn tenant_exceeding_rate_limit_gets_429_opaque() {
3998 use tower::ServiceExt;
3999 let router = app(two_tenant_state(Some(1))).expect("router");
4005
4006 let (r1, r2, r3, r4) = tokio::join!(
4007 router
4008 .clone()
4009 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4010 router
4011 .clone()
4012 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4013 router
4014 .clone()
4015 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4016 router
4017 .clone()
4018 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4019 );
4020 let responses = [r1.unwrap(), r2.unwrap(), r3.unwrap(), r4.unwrap()];
4021 let ok = responses
4022 .iter()
4023 .filter(|r| r.status() == axum::http::StatusCode::OK)
4024 .count();
4025 assert!(ok >= 1, "the burst's first request must pass");
4026 let limited: Vec<_> = responses
4027 .into_iter()
4028 .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
4029 .collect();
4030 assert!(
4031 !limited.is_empty(),
4032 "a 4-request burst against 1 req/sec must trip the limiter"
4033 );
4034
4035 let json = body_json(limited.into_iter().next().unwrap()).await;
4036 assert_eq!(json["error"]["type"], "rate_limited");
4037 let msg = json["error"]["message"].as_str().unwrap();
4039 assert!(!msg.contains('1'), "no limit value in body: {msg}");
4040 }
4041
4042 #[tokio::test]
4043 async fn rate_limit_buckets_are_independent_per_tenant() {
4044 use tower::ServiceExt;
4045 let router = app(two_tenant_state(Some(1))).expect("router");
4046
4047 let (a1, a2, a3) = tokio::join!(
4050 router
4051 .clone()
4052 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4053 router
4054 .clone()
4055 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4056 router
4057 .clone()
4058 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
4059 );
4060 let a_limited = [a1.unwrap(), a2.unwrap(), a3.unwrap()]
4061 .iter()
4062 .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
4063 .count();
4064 assert!(a_limited >= 1, "tenant A's burst must trip its limiter");
4065
4066 let b1 = router
4068 .clone()
4069 .oneshot(cap_request(Some("Bearer tenant-b.key-b")))
4070 .await
4071 .unwrap();
4072 assert_eq!(b1.status(), axum::http::StatusCode::OK);
4073 }
4074
4075 #[tokio::test]
4076 async fn rate_limit_unset_never_429s() {
4077 use tower::ServiceExt;
4078 let router = app(two_tenant_state(None)).expect("router");
4081 for _ in 0..20 {
4082 let resp = router
4083 .clone()
4084 .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
4085 .await
4086 .unwrap();
4087 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4088 }
4089 }
4090
4091 #[test]
4094 fn u01_is_deterministic_and_in_range() {
4095 let s1 = u01(0xDEAD_BEEF_CAFE_1234_u128);
4096 let s2 = u01(0xDEAD_BEEF_CAFE_1234_u128);
4097 assert_eq!(s1, s2, "u01 must be deterministic for the same seed");
4098 assert!((0.0..1.0).contains(&s1), "u01 must return [0, 1), got {s1}");
4099
4100 let s3 = u01(0x1234_5678_9ABC_DEF0_u128);
4102 assert_ne!(s1, s3, "different seeds should give different values");
4103
4104 for i in 0u64..256 {
4106 let v = u01(i as u128);
4107 assert!((0.0..1.0).contains(&v), "seed {i}: u01={v} out of [0,1)");
4108 }
4109 }
4110
4111 #[test]
4112 fn epsilon_propensity_formula() {
4113 let epsilon = 0.2_f64;
4114 let k = 3_usize;
4115 let greedy = 1_u32;
4116
4117 let p_greedy = epsilon_propensity(greedy, greedy, epsilon, k);
4119 let expected_greedy = (1.0 - epsilon) + epsilon / k as f64;
4120 assert!(
4121 (p_greedy - expected_greedy).abs() < 1e-12,
4122 "{p_greedy} != {expected_greedy}"
4123 );
4124
4125 let p_other = epsilon_propensity(0, greedy, epsilon, k);
4127 let expected_other = epsilon / k as f64;
4128 assert!(
4129 (p_other - expected_other).abs() < 1e-12,
4130 "{p_other} != {expected_other}"
4131 );
4132
4133 for chosen in 0..k as u32 {
4135 let p = epsilon_propensity(chosen, greedy, epsilon, k);
4136 assert!(
4137 p > 0.0 && p <= 1.0,
4138 "propensity {p} out of (0,1] for chosen={chosen}"
4139 );
4140 }
4141 }
4142
4143 #[test]
4144 fn epsilon_branch_and_greedy_branch_both_occur_over_many_seeds() {
4145 let epsilon = 0.3_f64;
4147 let mut saw_explore = false;
4148 let mut saw_greedy = false;
4149 for i in 0u64..200 {
4150 let u = u01(i as u128);
4151 if u < epsilon {
4152 saw_explore = true;
4153 } else {
4154 saw_greedy = true;
4155 }
4156 if saw_explore && saw_greedy {
4157 break;
4158 }
4159 }
4160 assert!(
4161 saw_explore,
4162 "epsilon branch must fire with epsilon=0.3 over 200 seeds"
4163 );
4164 assert!(
4165 saw_greedy,
4166 "greedy branch must occur with epsilon=0.3 over 200 seeds"
4167 );
4168 }
4169
4170 #[test]
4171 fn epsilon_propensity_sums_to_one_over_all_rungs() {
4172 let epsilon = 0.15_f64;
4174 let k = 4_usize;
4175 let greedy = 2_u32;
4176 let total: f64 = (0..k as u32)
4177 .map(|r| epsilon_propensity(r, greedy, epsilon, k))
4178 .sum();
4179 assert!(
4180 (total - 1.0).abs() < 1e-12,
4181 "propensities must sum to 1, got {total}"
4182 );
4183 }
4184
4185 #[tokio::test(start_paused = true)]
4188 async fn keepalive_stream_ticks_then_emits_final_frame() {
4189 let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
4190 let mut ticks = tokio::time::interval(SSE_KEEPALIVE_EVERY);
4191 ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
4192 ticks.reset();
4193 let mut stream = KeepaliveStream {
4194 rx: Some(rx),
4195 ticks,
4196 format_message: anthropic_sse_from_message,
4197 };
4198 async fn next(
4200 stream: &mut KeepaliveStream,
4201 ) -> Option<Option<Result<Bytes, std::convert::Infallible>>> {
4202 std::future::poll_fn(|cx| {
4203 std::task::Poll::Ready(
4204 match futures_core::Stream::poll_next(std::pin::Pin::new(&mut *stream), cx) {
4205 std::task::Poll::Ready(item) => Some(item),
4206 std::task::Poll::Pending => None,
4207 },
4208 )
4209 })
4210 .await
4211 }
4212
4213 assert!(
4215 next(&mut stream).await.is_none(),
4216 "no frame before an interval"
4217 );
4218
4219 tokio::time::advance(SSE_KEEPALIVE_EVERY + Duration::from_millis(1)).await;
4221 let frame = next(&mut stream)
4222 .await
4223 .expect("keepalive due")
4224 .unwrap()
4225 .unwrap();
4226 assert!(
4227 frame.starts_with(b": "),
4228 "keepalive must be an SSE comment (ignored by every conforming parser)"
4229 );
4230
4231 let message = serde_json::json!({
4233 "id": "msg_1", "type": "message", "role": "assistant", "model": "m",
4234 "content": [{ "type": "text", "text": "done" }],
4235 "usage": { "input_tokens": 1, "output_tokens": 1 }
4236 });
4237 tx.send(Ok(message)).unwrap();
4238 let frame = next(&mut stream)
4239 .await
4240 .expect("final frame")
4241 .unwrap()
4242 .unwrap();
4243 let text = String::from_utf8(frame.to_vec()).unwrap();
4244 assert!(text.contains("event: message_start"));
4245 assert!(text.contains("event: message_stop"));
4246 let eos = next(&mut stream)
4247 .await
4248 .expect("stream must end after the final frame");
4249 assert!(eos.is_none(), "end-of-stream after the final frame");
4250 }
4251
4252 #[tokio::test]
4255 async fn streaming_enforce_serves_full_sse_sequence() {
4256 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
4257 let config = ProxyConfig::from_lookup(|k| match k {
4258 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
4259 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
4260 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
4261 _ => None,
4262 })
4263 .unwrap();
4264 let mut outs = HashMap::new();
4265 outs.insert(
4266 "anthropic/m".to_owned(),
4267 Ok(model_resp("anthropic/m", "gated answer")),
4268 );
4269 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
4270 map.insert(
4271 "anthropic".to_owned(),
4272 Arc::new(MockProvider::new("anthropic", outs)),
4273 );
4274 let (traces, _rx) = mpsc::channel(64);
4275 let state = AppState {
4276 config: Arc::new(config),
4277 http: reqwest::Client::new(),
4278 providers: ProviderRegistry::from_map(map),
4279 gate_health: Arc::new(GateHealthRegistry::new()),
4280 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
4281 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
4282 traces,
4283 adaptive: None,
4284 bandit: None,
4285 predictor: None,
4286 tenant_rate_limiter: None,
4287 spill: None,
4288 };
4289 let body = Bytes::from_static(
4290 br#"{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}"#,
4291 );
4292 let resp = messages(
4293 State(state),
4294 Extension(TenantId("default".to_owned())),
4295 HeaderMap::new(),
4296 body,
4297 )
4298 .await;
4299 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4300 assert!(
4301 resp.headers()
4302 .get(axum::http::header::CONTENT_TYPE)
4303 .and_then(|v| v.to_str().ok())
4304 .is_some_and(|ct| ct.starts_with("text/event-stream")),
4305 "streaming client must get SSE"
4306 );
4307 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
4308 .await
4309 .unwrap();
4310 let text = String::from_utf8(bytes.to_vec()).unwrap();
4311 assert!(text.contains("event: message_start"));
4312 assert!(text.contains("gated answer"));
4313 assert!(text.contains("event: message_stop"));
4314 }
4315
4316 #[test]
4321 fn parse_openai_request_plain_text() {
4322 let body = br#"{"model":"gpt-4o","max_tokens":256,"messages":[{"role":"user","content":"hello"}]}"#;
4324 let req = parse_openai_request(body, false).expect("must parse");
4325 assert_eq!(req.model, "gpt-4o");
4326 assert_eq!(req.max_tokens, 256);
4327 assert_eq!(req.messages.len(), 1);
4328 assert_eq!(req.messages[0].role, "user");
4329 assert_eq!(req.messages[0].content, Value::String("hello".to_owned()));
4330 assert!(req.system.is_none());
4331 assert_eq!(req.raw, Value::Null);
4333 }
4334
4335 #[test]
4336 fn parse_openai_request_system_message() {
4337 let body = br#"{"model":"gpt-4o","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hi"}]}"#;
4338 let req = parse_openai_request(body, false).expect("must parse");
4339 assert_eq!(req.system.as_deref(), Some("be concise"));
4340 assert_eq!(req.messages.len(), 1);
4341 assert_eq!(req.messages[0].role, "user");
4342 }
4343
4344 #[test]
4345 fn parse_openai_request_tool_calls_translate_to_tool_use() {
4346 let body = br#"{
4347 "model":"gpt-4o",
4348 "messages":[
4349 {"role":"user","content":"what's the weather?"},
4350 {"role":"assistant","content":null,"tool_calls":[
4351 {"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Paris\"}"}}
4352 ]},
4353 {"role":"tool","tool_call_id":"call_1","content":"15C, cloudy"}
4354 ]
4355 }"#;
4356 let req = parse_openai_request(body, false).expect("must parse");
4357 assert_eq!(req.messages.len(), 3);
4359 let asst = &req.messages[1];
4361 assert_eq!(asst.role, "assistant");
4362 let blocks = asst.content.as_array().expect("content array");
4363 assert_eq!(blocks[0]["type"], "tool_use");
4364 assert_eq!(blocks[0]["name"], "get_weather");
4365 assert_eq!(blocks[0]["id"], "call_1");
4366 assert_eq!(blocks[0]["input"]["city"], "Paris");
4367 let tool_msg = &req.messages[2];
4369 assert_eq!(tool_msg.role, "user");
4370 let result_blocks = tool_msg.content.as_array().expect("result blocks");
4371 assert_eq!(result_blocks[0]["type"], "tool_result");
4372 assert_eq!(result_blocks[0]["tool_use_id"], "call_1");
4373 }
4374
4375 #[test]
4376 fn parse_openai_request_tools_translate_to_anthropic_format() {
4377 let body = br#"{
4378 "model":"gpt-4o",
4379 "messages":[{"role":"user","content":"use a tool"}],
4380 "tools":[{"type":"function","function":{"name":"search","description":"web search","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}}]
4381 }"#;
4382 let req = parse_openai_request(body, false).expect("must parse");
4383 let tools = req.tools.as_array().expect("tools array");
4384 assert_eq!(tools.len(), 1);
4385 assert_eq!(tools[0]["name"], "search");
4386 assert_eq!(tools[0]["description"], "web search");
4387 assert_eq!(tools[0]["input_schema"]["type"], "object");
4388 }
4389
4390 #[test]
4391 fn parse_openai_request_raw_carry_preserves_full_body() {
4392 let body =
4394 br#"{"model":"gpt-4o","max_tokens":100,"messages":[{"role":"user","content":"hi"}]}"#;
4395 let req = parse_openai_request(body, true).expect("must parse");
4396 assert!(req.raw.is_object(), "raw must be the full JSON object");
4397 assert_eq!(req.raw["model"], "gpt-4o");
4398 assert_eq!(req.raw["max_tokens"], 100);
4399 assert!(req.tools.is_null(), "no tools in this request");
4401 }
4402
4403 #[test]
4404 fn parse_openai_request_http_image_returns_none() {
4405 let body = br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
4407 {"type":"text","text":"describe this"},
4408 {"type":"image_url","image_url":{"url":"https://example.com/cat.png"}}
4409 ]}]}"#;
4410 let result = parse_openai_request(body, false);
4411 assert!(result.is_none(), "http image URL must fail translation");
4412 }
4413
4414 #[test]
4415 fn parse_openai_request_data_url_image_translates_to_anthropic_base64() {
4416 let body = br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
4417 {"type":"text","text":"describe"},
4418 {"type":"image_url","image_url":{"url":"data:image/png;base64,iVBORw0KGgo="}}
4419 ]}]}"#;
4420 let req = parse_openai_request(body, false).expect("data URL must parse");
4421 let blocks = req.messages[0].content.as_array().expect("blocks");
4422 let img = blocks
4423 .iter()
4424 .find(|b| b["type"] == "image")
4425 .expect("image block");
4426 assert_eq!(img["source"]["type"], "base64");
4427 assert_eq!(img["source"]["media_type"], "image/png");
4428 assert_eq!(img["source"]["data"], "iVBORw0KGgo=");
4429 }
4430
4431 #[test]
4434 fn openai_response_json_renders_text_response() {
4435 let resp = ModelResponse {
4436 model: "gpt-4o".to_owned(),
4437 text: "Hello!".to_owned(),
4438 in_tokens: 10,
4439 out_tokens: 5,
4440 raw: serde_json::json!({
4441 "content": [{ "type": "text", "text": "Hello!" }]
4442 }),
4443 };
4444 let json = openai_response_json(&resp);
4445 assert_eq!(json["object"], "chat.completion");
4446 assert_eq!(json["model"], "gpt-4o");
4447 assert_eq!(json["choices"][0]["message"]["role"], "assistant");
4448 assert_eq!(json["choices"][0]["message"]["content"], "Hello!");
4449 assert_eq!(json["choices"][0]["finish_reason"], "stop");
4450 assert_eq!(json["usage"]["prompt_tokens"], 10);
4451 assert_eq!(json["usage"]["completion_tokens"], 5);
4452 }
4453
4454 #[test]
4455 fn openai_response_json_renders_tool_call() {
4456 let resp = ModelResponse {
4457 model: "gpt-4o".to_owned(),
4458 text: String::new(),
4459 in_tokens: 20,
4460 out_tokens: 15,
4461 raw: serde_json::json!({
4462 "content": [{
4463 "type": "tool_use",
4464 "id": "call_abc",
4465 "name": "search",
4466 "input": {"q": "Rust async"}
4467 }]
4468 }),
4469 };
4470 let json = openai_response_json(&resp);
4471 assert_eq!(json["choices"][0]["finish_reason"], "tool_calls");
4472 let tc = &json["choices"][0]["message"]["tool_calls"][0];
4473 assert_eq!(tc["id"], "call_abc");
4474 assert_eq!(tc["type"], "function");
4475 assert_eq!(tc["function"]["name"], "search");
4476 assert_eq!(json["choices"][0]["message"]["content"], Value::Null);
4478 }
4479
4480 #[test]
4481 fn openai_sse_from_message_plain_text_has_role_then_content_then_stop() {
4482 let resp = ModelResponse {
4483 model: "gpt-4o".to_owned(),
4484 text: "Hi there!".to_owned(),
4485 in_tokens: 5,
4486 out_tokens: 3,
4487 raw: serde_json::json!({
4488 "content": [{ "type": "text", "text": "Hi there!" }]
4489 }),
4490 };
4491 let sse = openai_sse_from_message(&openai_response_json(&resp));
4492 for line in sse.lines() {
4494 assert!(
4495 line.is_empty() || line.starts_with("data: "),
4496 "bad SSE line: {line:?}"
4497 );
4498 }
4499 let frames: Vec<&str> = sse
4500 .lines()
4501 .filter_map(|l| l.strip_prefix("data: "))
4502 .collect();
4503 assert_eq!(*frames.last().unwrap(), "[DONE]");
4505 let role_frame: Value = serde_json::from_str(frames[0]).unwrap();
4507 assert_eq!(role_frame["choices"][0]["delta"]["role"], "assistant");
4508 assert!(frames.iter().any(|f| {
4510 if *f == "[DONE]" {
4511 return false;
4512 }
4513 serde_json::from_str::<Value>(f)
4514 .ok()
4515 .is_some_and(|v| v["choices"][0]["delta"]["content"] == "Hi there!")
4516 }));
4517 assert!(frames.iter().any(|f| {
4519 if *f == "[DONE]" {
4520 return false;
4521 }
4522 serde_json::from_str::<Value>(f)
4523 .ok()
4524 .is_some_and(|v| v["choices"][0]["finish_reason"] == "stop")
4525 }));
4526 }
4527
4528 #[test]
4531 fn detects_openai_tool_calls() {
4532 let with_tool_calls = Bytes::from_static(br#"{"messages":[
4533 {"role":"user","content":"hi"},
4534 {"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"f","arguments":"{}"}}]}
4535 ]}"#);
4536 let without = Bytes::from_static(br#"{"messages":[{"role":"user","content":"hi"}]}"#);
4537 let with_tool_msg = Bytes::from_static(
4538 br#"{"messages":[
4539 {"role":"tool","tool_call_id":"c1","content":"result"}
4540 ]}"#,
4541 );
4542 assert!(openai_messages_have_tool_calls(&with_tool_calls));
4543 assert!(!openai_messages_have_tool_calls(&without));
4544 assert!(openai_messages_have_tool_calls(&with_tool_msg));
4545 }
4546
4547 #[test]
4548 fn detects_openai_http_images() {
4549 let http_img = Bytes::from_static(
4550 br#"{"messages":[{"role":"user","content":[
4551 {"type":"image_url","image_url":{"url":"https://example.com/img.png"}}
4552 ]}]}"#,
4553 );
4554 let data_img = Bytes::from_static(
4555 br#"{"messages":[{"role":"user","content":[
4556 {"type":"image_url","image_url":{"url":"data:image/png;base64,abc"}}
4557 ]}]}"#,
4558 );
4559 let no_img = Bytes::from_static(br#"{"messages":[{"role":"user","content":"hi"}]}"#);
4560 assert!(openai_has_http_images(&http_img));
4561 assert!(!openai_has_http_images(&data_img));
4562 assert!(!openai_has_http_images(&no_img));
4563 }
4564
4565 #[test]
4566 fn enforce_can_handle_openai_inbound_all_openai_ladder() {
4567 let tools_body = Bytes::from_static(br#"{"model":"gpt-4o","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"c","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#);
4569 let f = extract_openai_features(&HeaderMap::new(), &tools_body);
4570 let ladder = vec!["openai/gpt-4o-mini".to_owned(), "openai/gpt-4o".to_owned()];
4571 let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
4572 assert!(enforce_can_handle(
4573 &f,
4574 &tools_body,
4575 true,
4576 &ladder,
4577 &providers,
4578 Dialect::Openai
4579 ));
4580 }
4581
4582 #[test]
4583 fn enforce_can_handle_openai_inbound_all_anthropic_ladder_no_http_image() {
4584 let tools_body = Bytes::from_static(br#"{"model":"gpt-4o","messages":[{"role":"assistant","content":null,"tool_calls":[{"id":"c","type":"function","function":{"name":"f","arguments":"{}"}}]}]}"#);
4586 let f = extract_openai_features(&HeaderMap::new(), &tools_body);
4587 let ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
4588 let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
4589 assert!(enforce_can_handle(
4590 &f,
4591 &tools_body,
4592 true,
4593 &ladder,
4594 &providers,
4595 Dialect::Openai,
4596 ));
4597 }
4598
4599 #[test]
4600 fn enforce_can_handle_openai_inbound_http_image_falls_back() {
4601 let img_body = Bytes::from_static(
4603 br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
4604 {"type":"image_url","image_url":{"url":"https://example.com/img.png"}}
4605 ]}]}"#,
4606 );
4607 let f = extract_openai_features(&HeaderMap::new(), &img_body);
4608 let ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
4609 let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
4610 assert!(!enforce_can_handle(
4611 &f,
4612 &img_body,
4613 true,
4614 &ladder,
4615 &providers,
4616 Dialect::Openai,
4617 ));
4618 }
4619
4620 fn openai_enforce_state(mock_resp: ModelResponse) -> AppState {
4624 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"mock/m\"]\ngates = [\"non-empty\"]\n";
4625 let config = ProxyConfig::from_lookup(|k| match k {
4626 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
4627 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
4628 _ => None,
4629 })
4630 .unwrap();
4631 let mut outs = HashMap::new();
4632 outs.insert("mock/m".to_owned(), Ok(mock_resp));
4633 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
4634 map.insert("mock".to_owned(), Arc::new(MockProvider::new("mock", outs)));
4635 let (traces, _rx) = mpsc::channel(64);
4636 std::mem::forget(_rx);
4637 let tenant_rate_limiter = build_tenant_rate_limiter(&config);
4638 AppState {
4639 config: Arc::new(config),
4640 http: reqwest::Client::new(),
4641 providers: ProviderRegistry::from_map(map),
4642 gate_health: Arc::new(GateHealthRegistry::new()),
4643 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
4644 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
4645 traces,
4646 adaptive: None,
4647 bandit: None,
4648 predictor: None,
4649 tenant_rate_limiter,
4650 spill: None,
4651 }
4652 }
4653
4654 #[tokio::test]
4655 async fn chat_completions_plain_text_enforce_returns_openai_shape() {
4656 let mock = model_resp("mock/m", "gated answer");
4657 let state = openai_enforce_state(mock);
4658 let body = Bytes::from_static(
4659 br#"{"model":"gpt-4o","messages":[{"role":"user","content":"hello"}]}"#,
4660 );
4661 let resp = chat_completions(
4662 State(state),
4663 Extension(TenantId("default".to_owned())),
4664 HeaderMap::new(),
4665 body,
4666 )
4667 .await;
4668 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4669 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
4670 .await
4671 .unwrap();
4672 let json: Value = serde_json::from_slice(&bytes).expect("must be JSON");
4673 assert_eq!(json["object"], "chat.completion");
4674 assert_eq!(json["choices"][0]["message"]["role"], "assistant");
4675 assert_eq!(json["choices"][0]["message"]["content"], "gated answer");
4676 assert_eq!(json["choices"][0]["finish_reason"], "stop");
4677 assert!(
4678 json["id"]
4679 .as_str()
4680 .is_some_and(|id| id.starts_with("chatcmpl-"))
4681 );
4682 }
4683
4684 #[tokio::test]
4685 async fn chat_completions_stream_true_returns_sse_with_openai_chunks() {
4686 let mock = model_resp("mock/m", "gated answer");
4687 let state = openai_enforce_state(mock);
4688 let body = Bytes::from_static(
4689 br#"{"model":"gpt-4o","stream":true,"messages":[{"role":"user","content":"hello"}]}"#,
4690 );
4691 let resp = chat_completions(
4692 State(state),
4693 Extension(TenantId("default".to_owned())),
4694 HeaderMap::new(),
4695 body,
4696 )
4697 .await;
4698 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4699 assert!(
4700 resp.headers()
4701 .get(axum::http::header::CONTENT_TYPE)
4702 .and_then(|v| v.to_str().ok())
4703 .is_some_and(|ct| ct.starts_with("text/event-stream")),
4704 "stream:true must return SSE"
4705 );
4706 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
4707 .await
4708 .unwrap();
4709 let text = String::from_utf8(bytes.to_vec()).unwrap();
4710 assert!(
4712 text.contains("chat.completion.chunk"),
4713 "must have OpenAI chunk frames"
4714 );
4715 assert!(text.contains("[DONE]"), "must end with [DONE]");
4716 assert!(
4717 text.contains("gated answer"),
4718 "content must be in the stream"
4719 );
4720 assert!(
4722 !text.contains("message_start"),
4723 "must not have Anthropic event types"
4724 );
4725 }
4726
4727 fn probe_state(
4731 sample_rate: f64,
4732 k: u32,
4733 outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
4734 ) -> (AppState, mpsc::Receiver<Trace>) {
4735 let toml = format!(
4736 "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [\"anthropic/claude-haiku-4-5\"]\ngates = [\"non-empty\"]\n\
4737 [escalation.probe]\nk = {k}\nsample_rate = {sample_rate}\n"
4738 );
4739 let config = ProxyConfig::from_lookup(|k_| match k_ {
4740 "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
4741 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
4742 _ => None,
4743 })
4744 .unwrap();
4745 let mut outs = HashMap::new();
4746 for (model, out) in outcomes {
4747 outs.insert(model.to_owned(), out);
4748 }
4749 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
4750 map.insert(
4751 "anthropic".to_owned(),
4752 Arc::new(MockProvider::new("anthropic", outs)),
4753 );
4754 let (traces, rx) = mpsc::channel(64);
4755 let state = AppState {
4756 config: Arc::new(config),
4757 http: reqwest::Client::new(),
4758 providers: ProviderRegistry::from_map(map),
4759 gate_health: Arc::new(GateHealthRegistry::new()),
4760 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
4761 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
4762 traces,
4763 adaptive: None,
4764 bandit: None,
4765 predictor: None,
4766 tenant_rate_limiter: None,
4767 spill: None,
4768 };
4769 (state, rx)
4770 }
4771
4772 async fn run_enforce_get_trace(state: AppState, mut rx: mpsc::Receiver<Trace>) -> Trace {
4774 let resp = messages(
4775 State(state),
4776 Extension(TenantId("default".to_owned())),
4777 HeaderMap::new(),
4778 user_body(),
4779 )
4780 .await;
4781 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4782 rx.try_recv().expect("trace must be enqueued")
4783 }
4784
4785 #[tokio::test]
4787 async fn probe_off_trace_has_no_probe_field() {
4788 let (state, rx) = enforce_state(
4789 &["anthropic/claude-haiku-4-5"],
4790 &["non-empty"],
4791 vec![(
4792 "anthropic/claude-haiku-4-5",
4793 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
4794 )],
4795 );
4796 assert!(
4797 state
4798 .config
4799 .routing
4800 .as_ref()
4801 .unwrap()
4802 .escalation
4803 .probe
4804 .is_none(),
4805 "probe must default to None"
4806 );
4807 let trace = run_enforce_get_trace(state, rx).await;
4808 assert!(
4809 trace.probe.is_none(),
4810 "probe=None config must not set trace.probe"
4811 );
4812 }
4813
4814 #[tokio::test]
4816 async fn probe_sample_rate_zero_never_fires() {
4817 let (state, rx) = probe_state(
4819 0.0,
4820 5,
4821 vec![(
4822 "anthropic/claude-haiku-4-5",
4823 Ok(model_resp("anthropic/claude-haiku-4-5", "hi")),
4824 )],
4825 );
4826 let trace = run_enforce_get_trace(state, rx).await;
4827 assert!(
4828 trace.probe.is_none(),
4829 "sample_rate=0.0 must never set trace.probe"
4830 );
4831 }
4832
4833 #[tokio::test]
4836 async fn probe_on_all_pass_sets_confident_pass() {
4837 let model = "anthropic/claude-haiku-4-5";
4839 let (state, mut rx) = probe_state(1.0, 3, vec![(model, Ok(model_resp(model, "hello")))]);
4840 let resp = messages(
4841 State(state),
4842 Extension(TenantId("default".to_owned())),
4843 HeaderMap::new(),
4844 user_body(),
4845 )
4846 .await;
4847 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4848 let json = body_json(resp).await;
4850 assert_eq!(
4851 json["content"][0]["text"], "hello",
4852 "served content unchanged"
4853 );
4854
4855 let trace = rx.try_recv().expect("trace enqueued");
4856 let sig = trace.probe.expect("probe must be set when sample_rate=1.0");
4857 assert_eq!(sig.k, 3);
4858 assert_eq!(sig.gate_pass_count, 3, "all 3 samples must pass non-empty");
4859 assert_eq!(
4860 sig.regime,
4861 firstpass_core::ProbeRegime::ConfidentPass,
4862 "all-pass → ConfidentPass"
4863 );
4864 assert!(
4865 sig.probe_cost_usd > 0.0,
4866 "k model calls must cost something"
4867 );
4868 }
4869
4870 #[tokio::test]
4873 async fn probe_on_all_fail_sets_confident_fail() {
4874 let model = "anthropic/claude-haiku-4-5";
4875 let (state, mut rx) = probe_state(1.0, 3, vec![(model, Ok(model_resp(model, "")))]);
4877 let resp = messages(
4879 State(state),
4880 Extension(TenantId("default".to_owned())),
4881 HeaderMap::new(),
4882 user_body(),
4883 )
4884 .await;
4885 assert_eq!(resp.status(), axum::http::StatusCode::OK);
4886
4887 let trace = rx.try_recv().expect("trace enqueued");
4888 let sig = trace.probe.expect("probe must be set when sample_rate=1.0");
4889 assert_eq!(
4890 sig.gate_pass_count, 0,
4891 "empty response fails non-empty: all 0 pass"
4892 );
4893 assert_eq!(
4894 sig.regime,
4895 firstpass_core::ProbeRegime::ConfidentFail,
4896 "0 passes → ConfidentFail"
4897 );
4898 }
4899
4900 #[tokio::test]
4903 async fn probe_on_served_output_identical_to_probe_off() {
4904 let model = "anthropic/claude-haiku-4-5";
4905 let mk = |sample_rate: f64| {
4906 probe_state(
4907 sample_rate,
4908 2,
4909 vec![(model, Ok(model_resp(model, "gated answer")))],
4910 )
4911 };
4912
4913 let (state_off, mut rx_off) = mk(0.0);
4915 let resp_off = messages(
4916 State(state_off),
4917 Extension(TenantId("default".to_owned())),
4918 HeaderMap::new(),
4919 user_body(),
4920 )
4921 .await;
4922 let json_off = body_json(resp_off).await;
4923 let trace_off = rx_off.try_recv().unwrap();
4924
4925 let (state_on, mut rx_on) = mk(1.0);
4927 let resp_on = messages(
4928 State(state_on),
4929 Extension(TenantId("default".to_owned())),
4930 HeaderMap::new(),
4931 user_body(),
4932 )
4933 .await;
4934 let json_on = body_json(resp_on).await;
4935 let trace_on = rx_on.try_recv().unwrap();
4936
4937 assert_eq!(
4939 json_off["content"][0]["text"], json_on["content"][0]["text"],
4940 "served text must be identical regardless of probe"
4941 );
4942 assert!(
4944 (trace_off.final_.total_cost_usd - trace_on.final_.total_cost_usd).abs() < 1e-12,
4945 "total_cost_usd must not include probe cost: off={} on={}",
4946 trace_off.final_.total_cost_usd,
4947 trace_on.final_.total_cost_usd
4948 );
4949 assert!(trace_off.probe.is_none());
4951 assert!(trace_on.probe.is_some());
4952 assert!(
4954 trace_on.probe.as_ref().unwrap().probe_cost_usd > 0.0,
4955 "probe cost must be positive"
4956 );
4957 }
4958
4959 #[tokio::test]
4969 async fn probe_does_not_mutate_gate_health() {
4970 let model = "anthropic/claude-haiku-4-5";
4971 let (mut state, rx) = probe_state(1.0, 2, vec![(model, Ok(model_resp(model, "answer")))]);
4972
4973 let registry = GateHealthRegistry::new().with_budget("non-empty", 2, 0.4);
4976 registry.record("default", "non-empty", true);
4979 assert!(
4980 registry.enabled("default", "non-empty"),
4981 "gate must start enabled (window not full yet)"
4982 );
4983 state.gate_health = Arc::new(registry);
4984
4985 let trace = run_enforce_get_trace(state, rx).await;
4997 let sig = trace.probe.expect("probe must fire with sample_rate=1.0");
4998 assert_eq!(sig.k, 2);
4999 assert!(
5000 sig.gate_pass_count <= 2,
5001 "gate_pass_count must be in [0, k]"
5002 );
5003 }
5004
5005 fn predictor_state(
5007 enabled: bool,
5008 outcome: Result<ModelResponse, ProviderError>,
5009 ) -> (AppState, mpsc::Receiver<Trace>) {
5010 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/claude-haiku-4-5\"]\ngates = [\"non-empty\"]\n";
5011 let config = ProxyConfig::from_lookup(|k_| match k_ {
5012 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
5013 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
5014 _ => None,
5015 })
5016 .unwrap();
5017 let mut outs = HashMap::new();
5018 outs.insert("anthropic/claude-haiku-4-5".to_owned(), outcome);
5019 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
5020 map.insert(
5021 "anthropic".to_owned(),
5022 Arc::new(MockProvider::new("anthropic", outs)),
5023 );
5024 let (traces, rx) = mpsc::channel(64);
5025 let predictor = enabled.then(|| {
5026 Arc::new(std::sync::Mutex::new(firstpass_core::PassPredictor::new(
5027 0.05, 1e-4,
5028 )))
5029 });
5030 let state = AppState {
5031 config: Arc::new(config),
5032 http: reqwest::Client::new(),
5033 providers: ProviderRegistry::from_map(map),
5034 gate_health: Arc::new(GateHealthRegistry::new()),
5035 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
5036 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
5037 traces,
5038 adaptive: None,
5039 bandit: None,
5040 predictor,
5041 tenant_rate_limiter: None,
5042 spill: None,
5043 };
5044 (state, rx)
5045 }
5046
5047 #[tokio::test]
5048 async fn predictor_off_leaves_predicted_pass_none() {
5049 let (state, rx) =
5050 predictor_state(false, Ok(model_resp("anthropic/claude-haiku-4-5", "ok")));
5051 let trace = run_enforce_get_trace(state, rx).await;
5052 assert!(
5053 trace.predicted_pass.is_none(),
5054 "predictor off => no field (byte-identical)"
5055 );
5056 let j = serde_json::to_string(&trace).unwrap();
5058 assert!(!j.contains("predicted_pass"), "None must be omitted: {j}");
5059 }
5060
5061 #[tokio::test]
5062 async fn predictor_on_records_shadow_prediction_and_serves_identically() {
5063 let (state_off, rx_off) = predictor_state(
5066 false,
5067 Ok(model_resp("anthropic/claude-haiku-4-5", "served answer")),
5068 );
5069 let off = run_enforce_get_trace(state_off, rx_off).await;
5070
5071 let (state_on, rx_on) = predictor_state(
5072 true,
5073 Ok(model_resp("anthropic/claude-haiku-4-5", "served answer")),
5074 );
5075 let on = run_enforce_get_trace(state_on, rx_on).await;
5076
5077 assert_eq!(
5078 on.final_.served_rung, off.final_.served_rung,
5079 "served rung identical"
5080 );
5081 assert_eq!(on.attempts.len(), off.attempts.len(), "same attempts");
5082 assert_eq!(
5083 on.final_.total_cost_usd, off.final_.total_cost_usd,
5084 "predictor never adds served cost"
5085 );
5086 let p = on
5087 .predicted_pass
5088 .expect("predictor on => predicted_pass recorded");
5089 assert!(p > 0.0 && p < 1.0, "shadow prediction in (0,1): {p}");
5090 }
5091
5092 fn rollout_state(
5094 percent: f64,
5095 key: &str,
5096 outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
5097 ) -> (AppState, mpsc::Receiver<Trace>) {
5098 let toml = format!(
5099 "[[route]]\nmatch = {{}}\nmode = \"enforce\"\n\
5100 ladder = [\"anthropic/claude-haiku-4-5\"]\ngates = [\"non-empty\"]\n\
5101 [route.rollout]\npercent = {percent}\nkey = \"{key}\"\n"
5102 );
5103 let config = ProxyConfig::from_lookup(|k| match k {
5104 "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
5105 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
5106 _ => None,
5107 })
5108 .unwrap();
5109 let mut outs = HashMap::new();
5110 for (model, out) in outcomes {
5111 outs.insert(model.to_owned(), out);
5112 }
5113 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
5114 map.insert(
5115 "anthropic".to_owned(),
5116 Arc::new(MockProvider::new("anthropic", outs)),
5117 );
5118 let (traces, rx) = mpsc::channel(64);
5119 let state = AppState {
5120 config: Arc::new(config),
5121 http: reqwest::Client::new(),
5122 providers: ProviderRegistry::from_map(map),
5123 gate_health: Arc::new(GateHealthRegistry::new()),
5124 shadow_ledger: Arc::new(crate::shadow::ShadowLedger::new()),
5125 guardrails: Arc::new(crate::guard::GuardrailRegistry::new()),
5126 traces,
5127 adaptive: None,
5128 bandit: None,
5129 predictor: None,
5130 tenant_rate_limiter: None,
5131 spill: None,
5132 };
5133 (state, rx)
5134 }
5135
5136 #[tokio::test]
5140 async fn rollout_at_zero_percent_never_enforces() {
5141 for i in 0..25 {
5142 let (state, mut rx) = rollout_state(
5143 0.0,
5144 "session",
5145 vec![(
5146 "anthropic/claude-haiku-4-5",
5147 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
5148 )],
5149 );
5150 let mut headers = HeaderMap::new();
5151 headers.insert("x-firstpass-session", format!("s{i}").parse().unwrap());
5152 let resp = messages(
5153 State(state),
5154 Extension(TenantId("default".to_owned())),
5155 headers,
5156 user_body(),
5157 )
5158 .await;
5159 let _ = resp;
5162 if let Ok(t) = rx.try_recv() {
5163 assert!(
5164 t.attempts.is_empty(),
5165 "0% rollout produced an enforced decision for session s{i}"
5166 );
5167 }
5168 }
5169 }
5170
5171 #[tokio::test]
5174 async fn rollout_at_hundred_percent_always_enforces() {
5175 for i in 0..25 {
5176 let (state, mut rx) = rollout_state(
5177 100.0,
5178 "session",
5179 vec![(
5180 "anthropic/claude-haiku-4-5",
5181 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
5182 )],
5183 );
5184 let mut headers = HeaderMap::new();
5185 headers.insert("x-firstpass-session", format!("s{i}").parse().unwrap());
5186 let resp = messages(
5187 State(state),
5188 Extension(TenantId("default".to_owned())),
5189 headers,
5190 user_body(),
5191 )
5192 .await;
5193 assert_eq!(resp.status(), axum::http::StatusCode::OK);
5194 let t = rx.try_recv().expect("trace enqueued");
5195 assert!(
5196 !t.attempts.is_empty(),
5197 "100% rollout skipped enforcement for session s{i}"
5198 );
5199 }
5200 }
5201
5202 #[tokio::test]
5210 async fn dispatch_arm_matches_the_predicted_arm_per_session() {
5211 let mut checked_enforced = 0;
5212 let mut checked_control = 0;
5213 for i in 0..24 {
5214 let session = format!("session-{i}");
5215 let (state, mut rx) = rollout_state(
5216 50.0,
5217 "session",
5218 vec![(
5219 "anthropic/claude-haiku-4-5",
5220 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
5221 )],
5222 );
5223 let expected = firstpass_core::rollout::decide(
5226 &state.config.prompt_salt,
5227 &firstpass_core::Rollout {
5228 percent: 50.0,
5229 key: firstpass_core::RolloutKey::Session,
5230 },
5231 &session,
5232 )
5233 .enforced;
5234
5235 let mut headers = HeaderMap::new();
5236 headers.insert("x-firstpass-session", session.parse().unwrap());
5237 let _ = messages(
5238 State(state),
5239 Extension(TenantId("default".to_owned())),
5240 headers,
5241 user_body(),
5242 )
5243 .await;
5244 let observed = rx
5245 .try_recv()
5246 .map(|t| !t.attempts.is_empty())
5247 .unwrap_or(false);
5248
5249 assert_eq!(
5250 observed,
5251 expected,
5252 "{session}: dispatch put it in the {} arm, bucketing predicted the {} arm",
5253 if observed { "enforced" } else { "control" },
5254 if expected { "enforced" } else { "control" }
5255 );
5256 if expected {
5257 checked_enforced += 1;
5258 } else {
5259 checked_control += 1;
5260 }
5261 }
5262 assert!(
5264 checked_enforced > 0 && checked_control > 0,
5265 "test saw only one arm ({checked_enforced} enforced, {checked_control} control)"
5266 );
5267 }
5268}