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