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 PolicyRef, RequestInfo, Score, ServedFrom, TaskKind, Trace, Verdict,
20};
21use serde::Deserialize;
22use serde_json::Value;
23use std::future::Future;
24use std::time::Duration;
25use tokio::sync::mpsc::error::TrySendError;
26use uuid::Uuid;
27
28use crate::config::ProxyConfig;
29use crate::error::ProxyError;
30use crate::gate::{GateHealthRegistry, resolve_gates};
31use crate::provider::{Auth, ChatMessage, ModelRequest, ModelResponse, ProviderRegistry};
32use crate::router::{EnforceCtx, EngineOutcome, route_enforce};
33use crate::store;
34use crate::tenant_auth::{TenantId, auth_middleware};
35use crate::upstream::{
36 forward_anthropic, forward_anthropic_streaming, forward_openai, forward_openai_streaming,
37};
38use firstpass_core::Route;
39
40#[derive(Clone)]
43pub struct AppState {
44 pub config: Arc<ProxyConfig>,
46 pub http: reqwest::Client,
48 pub providers: ProviderRegistry,
50 pub gate_health: Arc<GateHealthRegistry>,
52 pub traces: store::TraceSender,
54 pub adaptive: Option<Arc<std::sync::Mutex<firstpass_core::conformal::AdaptiveConformal>>>,
58 pub bandit: Option<Arc<std::sync::Mutex<crate::bandit::StartRungBandit>>>,
63 pub tenant_rate_limiter: Option<Arc<governor::DefaultKeyedRateLimiter<String>>>,
67 pub spill: Option<store::SpillHandle>,
71}
72
73#[must_use]
77pub fn build_tenant_rate_limiter(
78 config: &ProxyConfig,
79) -> Option<Arc<governor::DefaultKeyedRateLimiter<String>>> {
80 let per_sec = config.tenant_rate_per_sec?;
81 Some(Arc::new(governor::RateLimiter::keyed(
82 governor::Quota::per_second(per_sec),
83 )))
84}
85
86pub async fn tenant_rate_limit_middleware(
90 State(state): State<AppState>,
91 Extension(tenant): Extension<TenantId>,
92 req: Request,
93 next: Next,
94) -> Response {
95 if let Some(limiter) = &state.tenant_rate_limiter
96 && limiter.check_key(&tenant.0).is_err()
97 {
98 return ProxyError::RateLimited.into_response();
99 }
100 next.run(req).await
101}
102
103impl std::fmt::Debug for AppState {
104 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105 f.debug_struct("AppState")
106 .field("config", &self.config)
107 .finish_non_exhaustive()
108 }
109}
110
111fn offer_trace(traces: &store::TraceSender, spill: Option<&store::SpillHandle>, trace: Trace) {
128 record_trace_metrics(&trace);
129 match traces.try_send(trace) {
130 Ok(()) => {}
131 Err(TrySendError::Full(t)) => {
132 if let Some(handle) = spill {
133 match store::append_to_spill(handle, &t) {
134 Ok(()) => {
135 metrics::counter!("firstpass_receipts_spilled_total").increment(1);
136 }
137 Err(e) => {
138 tracing::error!(%e, "durable mode: spill write failed; trace lost");
139 metrics::counter!("firstpass_traces_dropped_total").increment(1);
140 }
141 }
142 } else {
143 tracing::warn!("trace channel full; dropping trace (writer behind under load)");
144 metrics::counter!("firstpass_traces_dropped_total").increment(1);
145 }
146 }
147 Err(TrySendError::Closed(_)) => {
148 tracing::warn!("trace writer is gone; dropping trace");
149 }
150 }
151}
152
153fn record_trace_metrics(trace: &Trace) {
157 if trace.mode == Mode::Enforce {
158 metrics::histogram!("firstpass_enforce_latency_ms")
159 .record(trace.final_.total_latency_ms as f64);
160 if trace.final_.escalations > 0 {
161 metrics::counter!("firstpass_escalations_total")
162 .increment(u64::from(trace.final_.escalations));
163 }
164 }
165 let served_from = match trace.final_.served_from {
166 ServedFrom::Attempt => "attempt",
167 ServedFrom::BestAttempt => "best_attempt",
168 ServedFrom::Error => "error",
169 };
170 metrics::counter!("firstpass_served_total", "served_from" => served_from).increment(1);
171 if trace.final_.served_from == ServedFrom::Error {
172 metrics::counter!("firstpass_upstream_failures_total").increment(1);
173 }
174 metrics::gauge!("firstpass_cost_usd_total").increment(trace.final_.total_cost_usd);
178 metrics::gauge!("firstpass_gate_cost_usd_total").increment(trace.final_.gate_cost_usd);
179 metrics::gauge!("firstpass_baseline_usd_total")
180 .increment(trace.final_.counterfactual_baseline_usd);
181 metrics::gauge!("firstpass_savings_usd_total").increment(trace.final_.savings_usd);
182 if let Some(rung) = trace.final_.served_rung {
184 let model = trace
185 .attempts
186 .iter()
187 .find(|a| a.rung == rung)
188 .map(|a| a.model.clone())
189 .unwrap_or_else(|| "unknown".to_owned());
190 metrics::counter!(
191 "firstpass_served_rung_total",
192 "rung" => rung.to_string(),
193 "model" => model
194 )
195 .increment(1);
196 }
197}
198
199const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
203
204pub(crate) fn u01(seed: u128) -> f64 {
215 let lo = splitmix64_finalise(seed as u64);
216 let hi = splitmix64_finalise((seed >> 64) as u64);
217 ((lo ^ hi) >> 11) as f64 * (1.0_f64 / (1u64 << 53) as f64)
219}
220
221fn splitmix64_finalise(mut z: u64) -> u64 {
222 z = z.wrapping_add(0x9E37_79B9_7F4A_7C15);
223 z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
224 z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
225 z ^ (z >> 31)
226}
227
228#[must_use]
235pub(crate) fn epsilon_propensity(chosen: u32, greedy: u32, epsilon: f64, k: usize) -> f64 {
236 let greedy_term = if chosen == greedy { 1.0 - epsilon } else { 0.0 };
237 greedy_term + epsilon / k as f64
238}
239
240pub fn app(state: AppState) -> Result<Router, ProxyError> {
247 crate::metrics::install()?;
248 let max_concurrency = state.config.max_concurrency;
249
250 let business = Router::new()
259 .route("/v1/messages", post(messages))
260 .route("/v1/chat/completions", post(chat_completions))
261 .route("/v1/feedback", post(feedback))
262 .route("/v1/capabilities", get(capabilities))
263 .layer(axum::middleware::from_fn_with_state(
264 state.clone(),
265 tenant_rate_limit_middleware,
266 ))
267 .layer(axum::middleware::from_fn_with_state(
268 state.clone(),
269 auth_middleware,
270 ));
271
272 Ok(Router::new()
273 .merge(business)
274 .route("/healthz", get(healthz))
275 .route("/metrics", get(crate::metrics::handler))
276 .layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES))
278 .layer(tower::limit::GlobalConcurrencyLimitLayer::new(
281 max_concurrency,
282 ))
283 .with_state(state))
284}
285
286async fn healthz() -> impl IntoResponse {
288 Json(serde_json::json!({ "status": "ok" }))
289}
290
291async fn capabilities(State(state): State<AppState>) -> impl IntoResponse {
294 let (ladder, gates) = state
297 .config
298 .routing
299 .as_ref()
300 .and_then(|c| c.routes.iter().find(|r| r.mode == Mode::Enforce))
301 .map(|r| (r.ladder.clone(), r.gates.clone()))
302 .unwrap_or_default();
303 Json(serde_json::json!({
304 "service": "firstpass",
305 "version": env!("CARGO_PKG_VERSION"),
306 "feature_version": FEATURE_VERSION,
307 "modes": ["observe", "enforce"],
308 "wire_apis": ["anthropic.messages", "openai.chat_completions"],
309 "ladder": ladder,
310 "gates": gates,
311 "feedback_api": "POST /v1/feedback",
312 "offboarding": "unset ANTHROPIC_BASE_URL (or OPENAI_BASE_URL for OpenAI clients)",
313 }))
314}
315
316#[derive(Debug, Deserialize)]
318struct FeedbackRequest {
319 trace_id: String,
321 gate_id: String,
323 verdict: String,
325 #[serde(default)]
327 score: Option<f64>,
328 reporter: String,
330}
331
332async fn feedback(
337 State(state): State<AppState>,
338 Extension(TenantId(tenant)): Extension<TenantId>,
339 body: Bytes,
340) -> Response {
341 let req: FeedbackRequest = match serde_json::from_slice(&body) {
342 Ok(r) => r,
343 Err(e) => {
344 return ProxyError::BadRequest(format!("invalid feedback body: {e}")).into_response();
345 }
346 };
347 let verdict = match req.verdict.as_str() {
348 "pass" => Verdict::Pass,
349 "fail" => Verdict::Fail,
350 "abstain" => Verdict::Abstain,
351 other => {
352 return ProxyError::BadRequest(format!("unknown verdict {other:?}")).into_response();
353 }
354 };
355 let score = match req.score {
356 Some(s) => match Score::new(s) {
357 Ok(sc) => Some(sc),
358 Err(_) => {
359 return ProxyError::BadRequest(format!("score {s} out of range [0,1]"))
360 .into_response();
361 }
362 },
363 None => None,
364 };
365
366 let db = state.config.db_path.clone();
367
368 let (db_check, tenant_check, tid_check) = (db.clone(), tenant.clone(), req.trace_id.clone());
373 match tokio::task::spawn_blocking(move || {
374 store::trace_exists(&db_check, &tenant_check, &tid_check)
375 })
376 .await
377 {
378 Ok(Ok(true)) => {}
379 Ok(Ok(false)) => {
380 return ProxyError::NotFound(format!("unknown trace_id {:?}", req.trace_id))
381 .into_response();
382 }
383 Ok(Err(e)) => {
384 tracing::error!(%e, "feedback: trace_exists check failed");
385 return ProxyError::Internal(e.to_string()).into_response();
386 }
387 Err(e) => {
388 tracing::error!(%e, "feedback: trace_exists task panicked");
389 return ProxyError::Internal(e.to_string()).into_response();
390 }
391 }
392
393 let feedback_signal = match verdict {
395 Verdict::Pass => Some(true),
396 Verdict::Fail => Some(false),
397 Verdict::Abstain => None,
398 };
399 let dv = DeferredVerdict {
400 gate_id: req.gate_id,
401 verdict,
402 score,
403 reported_at: jiff::Timestamp::now(),
404 reporter: req.reporter,
405 };
406 let trace_id = req.trace_id.clone();
407 match tokio::task::spawn_blocking(move || store::append_deferred(&db, &req.trace_id, &dv)).await
408 {
409 Ok(Ok(())) => {
410 if let (Some(a), Some(correct)) = (state.adaptive.as_ref(), feedback_signal)
412 && let Ok(mut g) = a.lock()
413 {
414 g.observe_served(correct);
415 metrics::gauge!("firstpass_serve_threshold").set(g.threshold());
416 metrics::gauge!("firstpass_realized_served_failure")
417 .set(g.realized_served_failure());
418 }
419 (
420 axum::http::StatusCode::ACCEPTED,
421 Json(serde_json::json!({ "status": "recorded", "trace_id": trace_id })),
422 )
423 .into_response()
424 }
425 Ok(Err(e)) => {
426 tracing::error!(%e, "feedback: append_deferred failed");
427 ProxyError::Internal(e.to_string()).into_response()
428 }
429 Err(e) => {
430 tracing::error!(%e, "feedback: append_deferred task panicked");
431 ProxyError::Internal(e.to_string()).into_response()
432 }
433 }
434}
435
436const SESSION_HEADER: &str = "x-firstpass-session";
439
440const AGENT_HEADER: &str = "x-firstpass-agent";
442const SUBAGENT_HEADER: &str = "x-firstpass-subagent";
444
445async fn messages(
450 State(state): State<AppState>,
451 Extension(TenantId(tenant)): Extension<TenantId>,
452 headers: HeaderMap,
453 body: Bytes,
454) -> Response {
455 let session_header = header_str(&headers, SESSION_HEADER);
456
457 if let Some(routing) = state.config.routing.as_ref() {
460 let features = extract_features(&headers, &body);
461 if let Some(route) = routing
462 .route_for(&features)
463 .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
464 {
465 let route = route.clone();
468 if enforce_can_handle(
469 &features,
470 &body,
471 routing.escalation.enforce_structured,
472 &route.ladder,
473 &state.providers,
474 Dialect::Anthropic,
475 ) {
476 return handle_enforce(
477 &state,
478 &headers,
479 &body,
480 features,
481 &route,
482 session_header,
483 tenant,
484 )
485 .await;
486 }
487 tracing::info!(
491 "enforce route matched but structured request can't be routed faithfully (flag/ladder); serving via observe passthrough"
492 );
493 }
494 }
495 observe_passthrough(state, headers, body, session_header, tenant).await
496}
497
498fn enforce_can_handle(
512 features: &Features,
513 body: &[u8],
514 enforce_structured: bool,
515 ladder: &[String],
516 providers: &crate::provider::ProviderRegistry,
517 inbound: Dialect,
518) -> bool {
519 let structured = features.tool_count > 0
520 || features.has_images
521 || match inbound {
522 Dialect::Anthropic => messages_have_tool_blocks(body),
523 Dialect::Openai => openai_messages_have_tool_calls(body),
524 Dialect::Gemini => false,
525 };
526 if !structured {
527 return true;
528 }
529 if !enforce_structured {
530 return false;
531 }
532 let all_verbatim = ladder.iter().all(|rung| {
534 let provider_id = rung.split('/').next().unwrap_or_default();
535 providers
536 .get(provider_id)
537 .is_some_and(|p| p.carries_structured_verbatim(inbound))
538 });
539 if all_verbatim {
540 return true;
541 }
542 if inbound == Dialect::Openai && !openai_has_http_images(body) {
546 let all_anthropic = ladder.iter().all(|rung| {
547 let pid = rung.split('/').next().unwrap_or_default();
548 providers
549 .get(pid)
550 .is_some_and(|p| p.carries_structured_verbatim(Dialect::Anthropic))
551 });
552 if all_anthropic {
553 return true;
554 }
555 }
556 false
557}
558
559fn messages_have_tool_blocks(body: &[u8]) -> bool {
562 serde_json::from_slice::<Value>(body)
563 .ok()
564 .and_then(|json| {
565 json.get("messages")
566 .and_then(Value::as_array)
567 .map(|messages| messages.iter().any(message_has_tool_block))
568 })
569 .unwrap_or(false)
570}
571
572fn message_has_tool_block(message: &Value) -> bool {
574 message
575 .get("content")
576 .and_then(Value::as_array)
577 .is_some_and(|blocks| {
578 blocks.iter().any(|block| {
579 matches!(
580 block.get("type").and_then(Value::as_str),
581 Some("tool_use" | "tool_result")
582 )
583 })
584 })
585}
586
587fn is_stream_request(body: &[u8]) -> bool {
589 serde_json::from_slice::<Value>(body)
590 .ok()
591 .and_then(|json| json.get("stream").and_then(Value::as_bool))
592 .unwrap_or(false)
593}
594
595fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
597 headers
598 .get(name)
599 .and_then(|v| v.to_str().ok())
600 .map(str::to_owned)
601}
602
603fn extract_features(headers: &HeaderMap, body: &[u8]) -> Features {
606 let (_model, tool_count, has_images) = request_features(body);
607 let mut f = Features::new(TaskKind::Other);
608 f.agent = header_str(headers, AGENT_HEADER);
609 f.subagent = header_str(headers, SUBAGENT_HEADER);
610 f.tool_count = tool_count;
611 f.has_images = has_images;
612 f.prompt_token_bucket = token_bucket(body.len() as u64);
615 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
616 f
617}
618
619async fn handle_enforce(
622 state: &AppState,
623 headers: &HeaderMap,
624 body: &Bytes,
625 features: Features,
626 route: &Route,
627 session_header: Option<String>,
628 tenant: String,
629) -> Response {
630 if is_stream_request(body) {
638 let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
639 let (state_c, headers_c, body_c, route_c) =
640 (state.clone(), headers.clone(), body.clone(), route.clone());
641 tokio::spawn(async move {
642 let out = enforce_pipeline(
643 &state_c,
644 &headers_c,
645 &body_c,
646 features,
647 &route_c,
648 session_header,
649 tenant,
650 )
651 .await;
652 let _ = tx.send(out);
653 });
654 return sse_keepalive_response(rx, anthropic_sse_from_message);
655 }
656 match enforce_pipeline(
657 state,
658 headers,
659 body,
660 features,
661 route,
662 session_header,
663 tenant,
664 )
665 .await
666 {
667 Ok(message) => (axum::http::StatusCode::OK, Json(message)).into_response(),
668 Err(e) => e.into_response(),
669 }
670}
671
672#[allow(clippy::too_many_arguments)] async fn enforce_pipeline_inner(
677 state: &AppState,
678 body: &Bytes,
679 base_request: ModelRequest,
680 auth: Auth,
681 features: Features,
682 route: &Route,
683 session_header: Option<String>,
684 tenant: String,
685 api: &str,
686) -> Result<ModelResponse, ProxyError> {
687 let gate_defs = state
688 .config
689 .routing
690 .as_ref()
691 .map_or(&[][..], |cfg| &cfg.gate_defs);
692 let gates = resolve_gates(
693 &route.gates,
694 gate_defs,
695 &state.providers,
696 &auth,
697 &state.config.prices,
698 );
699 let session_id = session_header.unwrap_or_else(|| Uuid::now_v7().to_string());
700 let (budget, max_rungs, speculation, serve_threshold) = match state.config.routing.as_ref() {
701 Some(cfg) => (
702 cfg.budget.per_request_usd,
703 cfg.escalation.max_rungs_per_request,
704 cfg.escalation.speculation,
705 cfg.escalation.serve_threshold,
706 ),
707 None => (None, 3, 0, None),
708 };
709 let serve_threshold = state
712 .adaptive
713 .as_ref()
714 .and_then(|a| a.lock().ok().map(|g| g.threshold()))
715 .or(serve_threshold);
716
717 let bandit_ctx = crate::bandit::ContextBucket::from_features(&features);
721
722 let (greedy_rung, base_policy_id, ts_propensity) = {
726 let (chosen, ts_p) = state
727 .bandit
728 .as_ref()
729 .and_then(|b| b.lock().ok())
730 .map(|mut b| {
731 b.choose_start_with_propensity(&bandit_ctx, &route.ladder, &state.config.prices)
732 })
733 .unwrap_or((0, None));
734 let policy = if ts_p.is_some() {
735 "bandit@v2-ts".to_owned()
736 } else if chosen > 0 {
737 "bandit@v1".to_owned()
738 } else {
739 "static-ladder@v0".to_owned()
740 };
741 (chosen, policy, ts_p)
742 };
743
744 let exploration_epsilon = state
749 .config
750 .routing
751 .as_ref()
752 .and_then(|cfg| cfg.escalation.exploration.as_ref())
753 .map(|e| e.epsilon);
754
755 let (start_rung, policy_id, explore_flag, propensity) = if ts_propensity.is_some() {
756 if exploration_epsilon.is_some() {
759 tracing::warn!(
760 "bandit.algorithm = thompson already logs propensities; \
761 [escalation.exploration] epsilon is ignored"
762 );
763 }
764 (greedy_rung, base_policy_id, false, ts_propensity)
765 } else if let Some(epsilon) = exploration_epsilon {
766 let k = route.ladder.len().max(1);
767 let u = u01(Uuid::now_v7().as_u128());
769 let (chosen, eps_branch) = if u < epsilon {
770 let idx = ((u / epsilon) * k as f64) as u32;
772 (idx.min(k as u32 - 1), true)
773 } else {
774 (greedy_rung, false)
775 };
776 let p = epsilon_propensity(chosen, greedy_rung, epsilon, k);
777 (chosen, format!("{base_policy_id}+eps"), eps_branch, Some(p))
778 } else {
779 (greedy_rung, base_policy_id, false, None)
780 };
781
782 let speculation = match state
788 .config
789 .routing
790 .as_ref()
791 .and_then(|cfg| cfg.escalation.speculation_band)
792 {
793 Some([lo, hi]) if speculation > 0 => {
794 let estimate = state
795 .bandit
796 .as_ref()
797 .and_then(|b| b.lock().ok())
798 .and_then(|b| b.pass_estimate(&bandit_ctx, start_rung));
799 match estimate {
800 Some(p) if p < lo || p > hi => {
801 metrics::counter!("firstpass_speculation_skipped_total").increment(1);
802 0
803 }
804 _ => speculation,
805 }
806 }
807 _ => speculation,
808 };
809
810 if state.bandit.is_some() {
812 metrics::counter!(
813 "firstpass_bandit_start_rung",
814 "rung" => start_rung.to_string()
815 )
816 .increment(1);
817 }
818
819 let ctx = EnforceCtx {
820 ladder: &route.ladder,
821 gates: &gates,
822 health: &state.gate_health,
823 base_request: &base_request,
824 providers: &state.providers,
825 auth: &auth,
826 prices: &state.config.prices,
827 budget_per_request_usd: budget,
828 max_rungs,
829 speculation,
830 serve_threshold,
831 features,
832 start_rung,
833 tenant_id: tenant,
836 session_id,
837 prompt_hash: prompt_hash(&state.config.prompt_salt, body),
838 api: api.to_owned(),
839 policy_id,
840 };
841
842 let (outcome, mut trace) = route_enforce(ctx).await;
843
844 trace.policy.explore = explore_flag;
847 trace.policy.propensity = propensity;
848
849 if let Some(bandit) = state.bandit.as_ref()
853 && let Ok(mut b) = bandit.lock()
854 {
855 for attempt in &trace.attempts {
856 b.observe(&bandit_ctx, attempt.rung, attempt.verdict);
857 }
858 }
859
860 offer_trace(&state.traces, state.spill.as_ref(), trace);
862
863 match outcome {
864 EngineOutcome::Served(resp) => Ok(resp),
865 EngineOutcome::Failed(msg) => Err(ProxyError::Engine(msg)),
866 }
867}
868
869async fn enforce_pipeline(
872 state: &AppState,
873 headers: &HeaderMap,
874 body: &Bytes,
875 features: Features,
876 route: &Route,
877 session_header: Option<String>,
878 tenant: String,
879) -> Result<Value, ProxyError> {
880 let Some(base_request) = parse_model_request(body) else {
881 return Err(ProxyError::BadRequest(
882 "request body is not a valid Anthropic Messages request".to_owned(),
883 ));
884 };
885 let auth = Auth::from_headers(headers);
886 let resp = enforce_pipeline_inner(
887 state,
888 body,
889 base_request,
890 auth,
891 features,
892 route,
893 session_header,
894 tenant,
895 "anthropic.messages",
896 )
897 .await?;
898 Ok(anthropic_response_json(&resp))
899}
900
901async fn enforce_pipeline_openai(
904 state: &AppState,
905 headers: &HeaderMap,
906 body: &Bytes,
907 features: Features,
908 route: &Route,
909 session_header: Option<String>,
910 tenant: String,
911) -> Result<Value, ProxyError> {
912 let providers = &state.providers;
914 let all_openai = route.ladder.iter().all(|rung| {
915 let pid = rung.split('/').next().unwrap_or_default();
916 providers
917 .get(pid)
918 .is_some_and(|p| p.carries_structured_verbatim(Dialect::Openai))
919 });
920 let Some(base_request) = parse_openai_request(body, all_openai) else {
921 return Err(ProxyError::BadRequest(
922 "request body is not a valid OpenAI Chat Completions request".to_owned(),
923 ));
924 };
925 let auth = Auth::from_headers(headers);
926 let resp = enforce_pipeline_inner(
927 state,
928 body,
929 base_request,
930 auth,
931 features,
932 route,
933 session_header,
934 tenant,
935 "openai.chat_completions",
936 )
937 .await?;
938 Ok(openai_response_json(&resp))
939}
940
941const SSE_KEEPALIVE_EVERY: Duration = Duration::from_secs(5);
943
944fn sse_keepalive_response(
953 rx: tokio::sync::oneshot::Receiver<Result<Value, ProxyError>>,
954 format_message: fn(&Value) -> String,
955) -> Response {
956 let mut ticks = tokio::time::interval(SSE_KEEPALIVE_EVERY);
957 ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
958 ticks.reset(); let stream = KeepaliveStream {
960 rx: Some(rx),
961 ticks,
962 format_message,
963 };
964 (
965 axum::http::StatusCode::OK,
966 [(
967 axum::http::header::CONTENT_TYPE,
968 "text/event-stream; charset=utf-8",
969 )],
970 axum::body::Body::from_stream(stream),
971 )
972 .into_response()
973}
974
975struct KeepaliveStream {
978 rx: Option<tokio::sync::oneshot::Receiver<Result<Value, ProxyError>>>,
980 ticks: tokio::time::Interval,
981 format_message: fn(&Value) -> String,
983}
984
985impl futures_core::Stream for KeepaliveStream {
986 type Item = Result<Bytes, std::convert::Infallible>;
987
988 fn poll_next(
989 mut self: std::pin::Pin<&mut Self>,
990 cx: &mut std::task::Context<'_>,
991 ) -> std::task::Poll<Option<Self::Item>> {
992 use std::task::Poll;
993 let Some(rx) = self.rx.as_mut() else {
994 return Poll::Ready(None); };
996 if let Poll::Ready(out) = std::pin::Pin::new(rx).poll(cx) {
997 let fmt = self.format_message;
998 let frame = match out {
999 Ok(Ok(message)) => fmt(&message),
1000 Ok(Err(e)) => sse_error_event(&e),
1001 Err(_) => sse_error_event(&ProxyError::Internal(
1002 "enforce pipeline task dropped".to_owned(),
1003 )),
1004 };
1005 self.rx = None;
1006 return Poll::Ready(Some(Ok(Bytes::from(frame))));
1007 }
1008 if self.ticks.poll_tick(cx).is_ready() {
1009 return Poll::Ready(Some(Ok(Bytes::from_static(b": firstpass routing\n\n"))));
1010 }
1011 Poll::Pending
1012 }
1013}
1014
1015fn sse_error_event(e: &ProxyError) -> String {
1018 let mut out = String::new();
1019 sse_event(
1020 &mut out,
1021 "error",
1022 &serde_json::json!({
1023 "type": "error",
1024 "error": { "type": "api_error", "message": e.client_message() }
1025 }),
1026 );
1027 out
1028}
1029
1030fn parse_model_request(body: &[u8]) -> Option<ModelRequest> {
1039 let json: Value = serde_json::from_slice(body).ok()?;
1040 let raw = json.clone();
1041 let messages_json = json.get("messages")?.as_array()?;
1042 let messages = messages_json
1043 .iter()
1044 .map(|m| ChatMessage {
1045 role: m
1046 .get("role")
1047 .and_then(Value::as_str)
1048 .unwrap_or("user")
1049 .to_owned(),
1050 content: m
1051 .get("content")
1052 .cloned()
1053 .unwrap_or_else(|| Value::String(String::new())),
1054 })
1055 .collect();
1056 let system = json
1057 .get("system")
1058 .and_then(Value::as_str)
1059 .map(str::to_owned);
1060 let max_tokens = json
1061 .get("max_tokens")
1062 .and_then(Value::as_u64)
1063 .and_then(|n| u32::try_from(n).ok())
1064 .unwrap_or(1024);
1065 let tools = json.get("tools").cloned().unwrap_or(Value::Null);
1066 Some(ModelRequest {
1067 model: json
1068 .get("model")
1069 .and_then(Value::as_str)
1070 .unwrap_or_default()
1071 .to_owned(),
1072 system,
1073 messages,
1074 max_tokens,
1075 tools,
1076 raw,
1077 })
1078}
1079
1080fn anthropic_response_json(resp: &ModelResponse) -> Value {
1090 let content = resp
1091 .raw
1092 .get("content")
1093 .filter(|c| c.is_array())
1094 .cloned()
1095 .unwrap_or_else(|| serde_json::json!([{ "type": "text", "text": resp.text }]));
1096 serde_json::json!({
1097 "id": format!("msg_{}", Uuid::now_v7()),
1098 "type": "message",
1099 "role": "assistant",
1100 "model": resp.model,
1101 "content": content,
1102 "usage": { "input_tokens": resp.in_tokens, "output_tokens": resp.out_tokens },
1103 })
1104}
1105
1106fn sse_event(out: &mut String, event: &str, data: &Value) {
1108 out.push_str("event: ");
1109 out.push_str(event);
1110 out.push_str("\ndata: ");
1111 out.push_str(&data.to_string());
1112 out.push_str("\n\n");
1113}
1114
1115fn anthropic_sse_from_message(message: &Value) -> String {
1122 let mut out = String::new();
1123
1124 let mut start_msg = message.clone();
1126 start_msg["content"] = Value::Array(Vec::new());
1127 sse_event(
1128 &mut out,
1129 "message_start",
1130 &serde_json::json!({ "type": "message_start", "message": start_msg }),
1131 );
1132
1133 let empty = Vec::new();
1134 let blocks = message
1135 .get("content")
1136 .and_then(Value::as_array)
1137 .unwrap_or(&empty);
1138 for (i, block) in blocks.iter().enumerate() {
1139 match block.get("type").and_then(Value::as_str) {
1140 Some("tool_use") => {
1141 let mut shell = block.clone();
1143 shell["input"] = serde_json::json!({});
1144 sse_event(
1145 &mut out,
1146 "content_block_start",
1147 &serde_json::json!({ "type": "content_block_start", "index": i, "content_block": shell }),
1148 );
1149 let input_json = block
1150 .get("input")
1151 .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
1152 sse_event(
1153 &mut out,
1154 "content_block_delta",
1155 &serde_json::json!({ "type": "content_block_delta", "index": i,
1156 "delta": { "type": "input_json_delta", "partial_json": input_json } }),
1157 );
1158 }
1159 _ => {
1160 let text = block.get("text").and_then(Value::as_str).unwrap_or("");
1162 sse_event(
1163 &mut out,
1164 "content_block_start",
1165 &serde_json::json!({ "type": "content_block_start", "index": i,
1166 "content_block": { "type": "text", "text": "" } }),
1167 );
1168 sse_event(
1169 &mut out,
1170 "content_block_delta",
1171 &serde_json::json!({ "type": "content_block_delta", "index": i,
1172 "delta": { "type": "text_delta", "text": text } }),
1173 );
1174 }
1175 }
1176 sse_event(
1177 &mut out,
1178 "content_block_stop",
1179 &serde_json::json!({ "type": "content_block_stop", "index": i }),
1180 );
1181 }
1182
1183 let out_tokens = message
1184 .pointer("/usage/output_tokens")
1185 .cloned()
1186 .unwrap_or_else(|| Value::from(0));
1187 sse_event(
1188 &mut out,
1189 "message_delta",
1190 &serde_json::json!({ "type": "message_delta", "delta": { "stop_reason": "end_turn" },
1191 "usage": { "output_tokens": out_tokens } }),
1192 );
1193 sse_event(
1194 &mut out,
1195 "message_stop",
1196 &serde_json::json!({ "type": "message_stop" }),
1197 );
1198 out
1199}
1200
1201async fn observe_passthrough(
1203 state: AppState,
1204 headers: HeaderMap,
1205 body: Bytes,
1206 session_header: Option<String>,
1207 tenant: String,
1208) -> Response {
1209 if is_stream_request(&body) {
1211 return observe_stream(state, headers, body, session_header, tenant).await;
1212 }
1213 let start = Instant::now();
1214 let result = forward_anthropic(
1215 &state.http,
1216 &state.config.upstream_anthropic,
1217 &headers,
1218 body.clone(),
1219 )
1220 .await;
1221 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
1222
1223 match result {
1224 Ok((status, resp_headers, resp_body)) => {
1225 spawn_trace(
1229 &state,
1230 body,
1231 Some(resp_body.clone()),
1232 latency_ms,
1233 session_header,
1234 tenant,
1235 );
1236 (status, resp_headers, resp_body).into_response()
1237 }
1238 Err(err) => {
1239 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
1240 err.into_response()
1241 }
1242 }
1243}
1244
1245async fn observe_stream(
1254 state: AppState,
1255 headers: HeaderMap,
1256 body: Bytes,
1257 session_header: Option<String>,
1258 tenant: String,
1259) -> Response {
1260 let start = Instant::now();
1261 let result = forward_anthropic_streaming(
1262 &state.http,
1263 &state.config.upstream_anthropic,
1264 &headers,
1265 body.clone(),
1266 )
1267 .await;
1268 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
1269
1270 match result {
1271 Ok((status, resp_headers, response)) => {
1272 spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
1273 let stream_body = Body::from_stream(response.bytes_stream());
1274 (status, resp_headers, stream_body).into_response()
1275 }
1276 Err(err) => {
1277 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
1278 err.into_response()
1279 }
1280 }
1281}
1282
1283fn spawn_stream_trace(
1285 state: &AppState,
1286 req_body: Bytes,
1287 latency_ms: u64,
1288 session_header: Option<String>,
1289 tenant: String,
1290) {
1291 let config = state.config.clone();
1292 let traces = state.traces.clone();
1293 let spill = state.spill.clone();
1294 tokio::spawn(async move {
1295 let mut trace =
1296 build_stream_trace(&config, &req_body, latency_ms, session_header.as_deref());
1297 trace.tenant_id = tenant;
1299 offer_trace(&traces, spill.as_ref(), trace);
1300 });
1301}
1302
1303fn spawn_trace(
1308 state: &AppState,
1309 req_body: Bytes,
1310 resp_body: Option<Bytes>,
1311 latency_ms: u64,
1312 session_header: Option<String>,
1313 tenant: String,
1314) {
1315 let config = state.config.clone();
1316 let traces = state.traces.clone();
1317 let spill = state.spill.clone();
1318 tokio::spawn(async move {
1319 let mut trace = match resp_body {
1320 Some(resp) => build_trace(
1321 &config,
1322 &req_body,
1323 &resp,
1324 latency_ms,
1325 session_header.as_deref(),
1326 ),
1327 None => build_error_trace(&config, &req_body, latency_ms, session_header.as_deref()),
1328 };
1329 trace.tenant_id = tenant;
1331 offer_trace(&traces, spill.as_ref(), trace);
1332 });
1333}
1334
1335fn session_id(session_header: Option<&str>, trace_id: Uuid) -> String {
1337 session_header
1338 .map(str::to_owned)
1339 .unwrap_or_else(|| trace_id.to_string())
1340}
1341
1342fn prompt_hash(salt: &str, body: &[u8]) -> String {
1345 let mut salted = Vec::with_capacity(salt.len() + body.len());
1346 salted.extend_from_slice(salt.as_bytes());
1347 salted.extend_from_slice(body);
1348 sha256_hex(&salted)
1349}
1350
1351fn request_features(body: &[u8]) -> (Option<String>, u32, bool) {
1355 let Ok(json) = serde_json::from_slice::<Value>(body) else {
1356 return (None, 0, false);
1357 };
1358 let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
1359 let tool_count = json
1360 .get("tools")
1361 .and_then(Value::as_array)
1362 .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
1363 let has_images = json
1364 .get("messages")
1365 .and_then(Value::as_array)
1366 .is_some_and(|messages| messages.iter().any(message_has_image));
1367 (model, tool_count, has_images)
1368}
1369
1370fn message_has_image(message: &Value) -> bool {
1372 message
1373 .get("content")
1374 .and_then(Value::as_array)
1375 .is_some_and(|blocks| {
1376 blocks
1377 .iter()
1378 .any(|block| block.get("type").and_then(Value::as_str) == Some("image"))
1379 })
1380}
1381
1382fn response_usage(body: &[u8]) -> (Option<String>, u64, u64) {
1385 let Ok(json) = serde_json::from_slice::<Value>(body) else {
1386 return (None, 0, 0);
1387 };
1388 let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
1389 let in_tokens = json
1390 .pointer("/usage/input_tokens")
1391 .and_then(Value::as_u64)
1392 .unwrap_or(0);
1393 let out_tokens = json
1394 .pointer("/usage/output_tokens")
1395 .and_then(Value::as_u64)
1396 .unwrap_or(0);
1397 (model, in_tokens, out_tokens)
1398}
1399
1400fn build_trace(
1402 config: &ProxyConfig,
1403 req_body: &Bytes,
1404 resp_body: &Bytes,
1405 latency_ms: u64,
1406 session_header: Option<&str>,
1407) -> Trace {
1408 let (req_model, tool_count, has_images) = request_features(req_body);
1409 let (resp_model, in_tokens, out_tokens) = response_usage(resp_body);
1410 let model = resp_model
1411 .or(req_model)
1412 .unwrap_or_else(|| "unknown".to_owned());
1413
1414 let cost_usd = config
1415 .prices
1416 .cost_usd(&format!("anthropic/{model}"), in_tokens, out_tokens)
1417 .unwrap_or(0.0);
1418
1419 let attempt = Attempt {
1420 rung: 0,
1421 model,
1422 provider: "anthropic".to_owned(),
1423 in_tokens,
1424 out_tokens,
1425 cost_usd,
1426 latency_ms,
1427 gates: Vec::new(),
1428 verdict: Verdict::Pass,
1429 };
1430
1431 let mut trace = base_trace(config, req_body, latency_ms, session_header);
1432 trace.request.features.prompt_token_bucket = token_bucket(in_tokens);
1433 trace.request.features.tool_count = tool_count;
1434 trace.request.features.has_images = has_images;
1435 trace.attempts.push(attempt);
1436 trace.final_ = FinalOutcome {
1437 served_rung: Some(0),
1438 served_from: ServedFrom::Attempt,
1439 total_cost_usd: cost_usd,
1440 gate_cost_usd: 0.0,
1441 total_latency_ms: latency_ms,
1442 escalations: 0,
1443 counterfactual_baseline_usd: cost_usd,
1444 savings_usd: 0.0,
1445 };
1446 trace.recompute_savings();
1447 trace
1448}
1449
1450fn build_stream_trace(
1454 config: &ProxyConfig,
1455 req_body: &Bytes,
1456 latency_ms: u64,
1457 session_header: Option<&str>,
1458) -> Trace {
1459 let (req_model, tool_count, has_images) = request_features(req_body);
1460 let model = req_model.unwrap_or_else(|| "unknown".to_owned());
1461
1462 let attempt = Attempt {
1463 rung: 0,
1464 model,
1465 provider: "anthropic".to_owned(),
1466 in_tokens: 0,
1467 out_tokens: 0,
1468 cost_usd: 0.0,
1469 latency_ms,
1470 gates: Vec::new(),
1471 verdict: Verdict::Pass,
1472 };
1473
1474 let mut trace = base_trace(config, req_body, latency_ms, session_header);
1475 trace.request.features.tool_count = tool_count;
1476 trace.request.features.has_images = has_images;
1477 trace.attempts.push(attempt);
1478 trace.final_ = FinalOutcome {
1479 served_rung: Some(0),
1480 served_from: ServedFrom::Attempt,
1481 total_cost_usd: 0.0,
1482 gate_cost_usd: 0.0,
1483 total_latency_ms: latency_ms,
1484 escalations: 0,
1485 counterfactual_baseline_usd: 0.0,
1486 savings_usd: 0.0,
1487 };
1488 trace.recompute_savings();
1489 trace
1490}
1491
1492fn build_error_trace(
1496 config: &ProxyConfig,
1497 req_body: &Bytes,
1498 latency_ms: u64,
1499 session_header: Option<&str>,
1500) -> Trace {
1501 let (_, tool_count, has_images) = request_features(req_body);
1502 let mut trace = base_trace(config, req_body, latency_ms, session_header);
1503 trace.request.features.tool_count = tool_count;
1504 trace.request.features.has_images = has_images;
1505 trace.final_ = FinalOutcome {
1506 served_rung: None,
1507 served_from: ServedFrom::Error,
1508 total_cost_usd: 0.0,
1509 gate_cost_usd: 0.0,
1510 total_latency_ms: latency_ms,
1511 escalations: 0,
1512 counterfactual_baseline_usd: 0.0,
1513 savings_usd: 0.0,
1514 };
1515 trace.recompute_savings();
1516 trace
1517}
1518
1519fn base_trace(
1522 config: &ProxyConfig,
1523 req_body: &Bytes,
1524 latency_ms: u64,
1525 session_header: Option<&str>,
1526) -> Trace {
1527 let trace_id = Uuid::now_v7();
1528 let mut features = Features::new(TaskKind::Other);
1529 features.hour_bucket = hour_bucket(jiff::Timestamp::now());
1530
1531 Trace {
1532 trace_id,
1533 prev_hash: GENESIS_HASH.to_owned(),
1534 tenant_id: config.tenant_id.clone(),
1535 session_id: session_id(session_header, trace_id),
1536 ts: jiff::Timestamp::now(),
1537 mode: Mode::Observe,
1538 policy: PolicyRef {
1539 id: "observe-passthrough@v0".to_owned(),
1540 explore: false,
1541 propensity: None,
1542 },
1543 request: RequestInfo {
1544 api: "anthropic.messages".to_owned(),
1545 prompt_hash: prompt_hash(&config.prompt_salt, req_body),
1546 features,
1547 },
1548 attempts: Vec::new(),
1549 deferred: Vec::new(),
1550 final_: FinalOutcome {
1551 served_rung: None,
1552 served_from: ServedFrom::Error,
1553 total_cost_usd: 0.0,
1554 gate_cost_usd: 0.0,
1555 total_latency_ms: latency_ms,
1556 escalations: 0,
1557 counterfactual_baseline_usd: 0.0,
1558 savings_usd: 0.0,
1559 },
1560 }
1561}
1562
1563fn openai_messages_have_tool_calls(body: &[u8]) -> bool {
1568 serde_json::from_slice::<Value>(body)
1569 .ok()
1570 .and_then(|json| {
1571 json.get("messages").and_then(Value::as_array).map(|msgs| {
1572 msgs.iter().any(|m| {
1573 m.get("tool_calls").is_some()
1574 || m.get("role").and_then(Value::as_str) == Some("tool")
1575 })
1576 })
1577 })
1578 .unwrap_or(false)
1579}
1580
1581fn openai_has_http_images(body: &[u8]) -> bool {
1585 serde_json::from_slice::<Value>(body)
1586 .ok()
1587 .and_then(|json| {
1588 json.get("messages").and_then(Value::as_array).map(|msgs| {
1589 msgs.iter().any(|m| {
1590 m.get("content")
1591 .and_then(Value::as_array)
1592 .is_some_and(|parts| {
1593 parts.iter().any(|p| {
1594 p.get("type").and_then(Value::as_str) == Some("image_url")
1595 && p.pointer("/image_url/url")
1596 .and_then(Value::as_str)
1597 .is_some_and(|u| {
1598 u.starts_with("http://") || u.starts_with("https://")
1599 })
1600 })
1601 })
1602 })
1603 })
1604 })
1605 .unwrap_or(false)
1606}
1607
1608fn openai_messages_have_images(body: &[u8]) -> bool {
1611 serde_json::from_slice::<Value>(body)
1612 .ok()
1613 .and_then(|json| {
1614 json.get("messages").and_then(Value::as_array).map(|msgs| {
1615 msgs.iter().any(|m| {
1616 m.get("content")
1617 .and_then(Value::as_array)
1618 .is_some_and(|parts| {
1619 parts
1620 .iter()
1621 .any(|p| p.get("type").and_then(Value::as_str) == Some("image_url"))
1622 })
1623 })
1624 })
1625 })
1626 .unwrap_or(false)
1627}
1628
1629fn extract_openai_features(headers: &HeaderMap, body: &[u8]) -> Features {
1632 let Ok(json) = serde_json::from_slice::<Value>(body) else {
1633 let mut f = Features::new(TaskKind::Other);
1634 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
1635 return f;
1636 };
1637 let tool_count = json
1638 .get("tools")
1639 .and_then(Value::as_array)
1640 .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
1641 let has_images = openai_messages_have_images(body);
1642 let mut f = Features::new(TaskKind::Other);
1643 f.agent = header_str(headers, AGENT_HEADER);
1644 f.subagent = header_str(headers, SUBAGENT_HEADER);
1645 f.tool_count = tool_count;
1646 f.has_images = has_images;
1647 f.prompt_token_bucket = token_bucket(body.len() as u64);
1648 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
1649 f
1650}
1651
1652fn parse_data_url(url: &str) -> Option<(&str, &str)> {
1656 let rest = url.strip_prefix("data:")?;
1657 let (meta, data) = rest.split_once(',')?;
1658 let media_type = meta.strip_suffix(";base64")?;
1659 Some((media_type, data))
1660}
1661
1662fn translate_openai_user_content(content: &Value) -> Option<Value> {
1665 match content {
1666 Value::String(_) => Some(content.clone()),
1668 Value::Array(parts) => {
1669 let mut blocks: Vec<Value> = Vec::with_capacity(parts.len());
1670 for part in parts {
1671 match part.get("type").and_then(Value::as_str) {
1672 Some("text") => {
1673 let text = part.get("text").and_then(Value::as_str).unwrap_or("");
1674 blocks.push(serde_json::json!({ "type": "text", "text": text }));
1675 }
1676 Some("image_url") => {
1677 let url = part.pointer("/image_url/url").and_then(Value::as_str)?;
1678 if url.starts_with("http://") || url.starts_with("https://") {
1679 return None; }
1681 let (media_type, data) = parse_data_url(url)?;
1683 blocks.push(serde_json::json!({
1684 "type": "image",
1685 "source": { "type": "base64", "media_type": media_type, "data": data }
1686 }));
1687 }
1688 _ => {} }
1690 }
1691 Some(Value::Array(blocks))
1692 }
1693 _ => Some(Value::String(String::new())),
1694 }
1695}
1696
1697fn translate_openai_tools(tools: &Value) -> Value {
1701 let Some(arr) = tools.as_array() else {
1702 return Value::Null;
1703 };
1704 let converted: Vec<Value> = arr
1705 .iter()
1706 .map(|tool| {
1707 let func = tool.get("function").unwrap_or(&Value::Null);
1708 let mut out = serde_json::json!({
1709 "name": func.get("name").cloned().unwrap_or(Value::String(String::new())),
1710 "input_schema": func.get("parameters").cloned()
1711 .unwrap_or_else(|| serde_json::json!({ "type": "object" })),
1712 });
1713 if let Some(desc) = func.get("description") {
1714 out["description"] = desc.clone();
1715 }
1716 out
1717 })
1718 .collect();
1719 Value::Array(converted)
1720}
1721
1722fn translate_openai_tool_choice(tc: &Value) -> Value {
1724 match tc {
1725 Value::String(s) => match s.as_str() {
1726 "auto" => serde_json::json!({ "type": "auto" }),
1727 "required" => serde_json::json!({ "type": "any" }),
1728 _ => serde_json::json!({ "type": "auto" }),
1730 },
1731 Value::Object(_) => {
1732 if tc.get("type").and_then(Value::as_str) == Some("function") {
1734 let name = tc.pointer("/function/name").cloned().unwrap_or(Value::Null);
1735 serde_json::json!({ "type": "tool", "name": name })
1736 } else {
1737 serde_json::json!({ "type": "auto" })
1738 }
1739 }
1740 _ => serde_json::json!({ "type": "auto" }),
1741 }
1742}
1743
1744pub fn parse_openai_request(body: &[u8], carry_raw: bool) -> Option<ModelRequest> {
1757 let json: Value = serde_json::from_slice(body).ok()?;
1758 let raw = if carry_raw { json.clone() } else { Value::Null };
1759
1760 let messages_json = json.get("messages")?.as_array()?;
1761
1762 let mut system: Option<String> = None;
1763 let mut messages: Vec<ChatMessage> = Vec::with_capacity(messages_json.len());
1764 let mut tools = Value::Null;
1765 let mut tool_choice_override: Option<Value> = None;
1766
1767 for msg in messages_json {
1768 let role = msg.get("role").and_then(Value::as_str).unwrap_or("user");
1769 match role {
1770 "system" => {
1771 if let Some(s) = msg.get("content").and_then(Value::as_str) {
1776 system = Some(s.to_owned());
1777 }
1778 }
1779 "user" => {
1780 let content_val = msg.get("content").unwrap_or(&Value::Null);
1781 let translated = translate_openai_user_content(content_val)?;
1782 messages.push(ChatMessage {
1783 role: "user".to_owned(),
1784 content: translated,
1785 });
1786 }
1787 "assistant" => {
1788 if let Some(tc_arr) = msg.get("tool_calls").and_then(Value::as_array) {
1789 let mut blocks: Vec<Value> = Vec::new();
1791 if let Some(text) = msg.get("content").and_then(Value::as_str)
1793 && !text.is_empty()
1794 {
1795 blocks.push(serde_json::json!({ "type": "text", "text": text }));
1796 }
1797 for tc in tc_arr {
1798 let id = tc.get("id").and_then(Value::as_str).unwrap_or("");
1799 let func = tc.get("function").unwrap_or(&Value::Null);
1800 let name = func.get("name").and_then(Value::as_str).unwrap_or("");
1801 let args_str = func
1802 .get("arguments")
1803 .and_then(Value::as_str)
1804 .unwrap_or("{}");
1805 let input: Value = serde_json::from_str(args_str)
1806 .unwrap_or_else(|_| serde_json::json!({}));
1807 blocks.push(serde_json::json!({
1808 "type": "tool_use",
1809 "id": id,
1810 "name": name,
1811 "input": input,
1812 }));
1813 }
1814 messages.push(ChatMessage {
1815 role: "assistant".to_owned(),
1816 content: Value::Array(blocks),
1817 });
1818 } else {
1819 let content = msg
1821 .get("content")
1822 .cloned()
1823 .unwrap_or_else(|| Value::String(String::new()));
1824 messages.push(ChatMessage {
1825 role: "assistant".to_owned(),
1826 content,
1827 });
1828 }
1829 }
1830 "tool" => {
1831 let tool_call_id = msg
1833 .get("tool_call_id")
1834 .and_then(Value::as_str)
1835 .unwrap_or("");
1836 let content = msg
1837 .get("content")
1838 .cloned()
1839 .unwrap_or_else(|| Value::String(String::new()));
1840 let result_block = serde_json::json!({
1841 "type": "tool_result",
1842 "tool_use_id": tool_call_id,
1843 "content": content,
1844 });
1845 messages.push(ChatMessage {
1846 role: "user".to_owned(),
1847 content: Value::Array(vec![result_block]),
1848 });
1849 }
1850 _ => {} }
1852 }
1853
1854 if !carry_raw {
1856 if let Some(t) = json.get("tools") {
1857 tools = translate_openai_tools(t);
1858 }
1859 if let Some(tc) = json.get("tool_choice") {
1860 tool_choice_override = Some(translate_openai_tool_choice(tc));
1861 }
1862 } else {
1863 tools = json.get("tools").cloned().unwrap_or(Value::Null);
1864 }
1865
1866 let max_tokens = json
1867 .get("max_tokens")
1868 .and_then(Value::as_u64)
1869 .and_then(|n| u32::try_from(n).ok())
1870 .unwrap_or(1024);
1871
1872 let model = json
1873 .get("model")
1874 .and_then(Value::as_str)
1875 .unwrap_or_default()
1876 .to_owned();
1877
1878 let _ = tool_choice_override; Some(ModelRequest {
1889 model,
1890 system,
1891 messages,
1892 max_tokens,
1893 tools,
1894 raw,
1895 })
1896}
1897
1898fn extract_openai_content_and_tools(raw: &Value, text: &str) -> (Value, Option<Value>) {
1905 if let Some(blocks) = raw.get("content").and_then(Value::as_array) {
1907 let mut text_parts: Vec<&str> = Vec::new();
1908 let mut tool_calls: Vec<Value> = Vec::new();
1909 for block in blocks {
1910 match block.get("type").and_then(Value::as_str) {
1911 Some("text") => {
1912 if let Some(t) = block.get("text").and_then(Value::as_str) {
1913 text_parts.push(t);
1914 }
1915 }
1916 Some("tool_use") => {
1917 let id = block.get("id").and_then(Value::as_str).unwrap_or("");
1918 let name = block.get("name").and_then(Value::as_str).unwrap_or("");
1919 let input_str = block
1920 .get("input")
1921 .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
1922 tool_calls.push(serde_json::json!({
1923 "id": id,
1924 "type": "function",
1925 "function": { "name": name, "arguments": input_str },
1926 }));
1927 }
1928 _ => {}
1929 }
1930 }
1931 let content_text = if tool_calls.is_empty() || !text_parts.is_empty() {
1932 Value::String(text_parts.join(""))
1933 } else {
1934 Value::Null };
1936 let tc = if tool_calls.is_empty() {
1937 None
1938 } else {
1939 Some(Value::Array(tool_calls))
1940 };
1941 return (content_text, tc);
1942 }
1943
1944 if let Some(msg) = raw.pointer("/choices/0/message") {
1946 let content = msg
1947 .get("content")
1948 .cloned()
1949 .unwrap_or(Value::String(text.to_owned()));
1950 let tc = msg.get("tool_calls").cloned();
1951 return (content, tc);
1952 }
1953
1954 (Value::String(text.to_owned()), None)
1956}
1957
1958fn openai_response_json(resp: &ModelResponse) -> Value {
1961 let (content_text, tool_calls) = extract_openai_content_and_tools(&resp.raw, &resp.text);
1962 let finish_reason = if tool_calls.is_some() {
1963 "tool_calls"
1964 } else {
1965 "stop"
1966 };
1967 let mut message = serde_json::json!({
1968 "role": "assistant",
1969 "content": content_text,
1970 });
1971 if let Some(tc) = tool_calls {
1972 message["tool_calls"] = tc;
1973 }
1974 serde_json::json!({
1975 "id": format!("chatcmpl-{}", Uuid::now_v7()),
1976 "object": "chat.completion",
1977 "created": jiff::Timestamp::now().as_second(),
1978 "model": resp.model,
1979 "choices": [{
1980 "index": 0,
1981 "message": message,
1982 "finish_reason": finish_reason,
1983 }],
1984 "usage": {
1985 "prompt_tokens": resp.in_tokens,
1986 "completion_tokens": resp.out_tokens,
1987 "total_tokens": resp.in_tokens + resp.out_tokens,
1988 }
1989 })
1990}
1991
1992fn openai_sse_from_message(message: &Value) -> String {
1996 let id = message
1997 .get("id")
1998 .and_then(Value::as_str)
1999 .unwrap_or("chatcmpl-unknown")
2000 .to_owned();
2001 let created = message.get("created").cloned().unwrap_or(Value::from(0));
2002 let model = message
2003 .get("model")
2004 .and_then(Value::as_str)
2005 .unwrap_or("unknown");
2006
2007 let choices = message.get("choices").and_then(Value::as_array);
2008 let msg = choices
2009 .and_then(|c| c.first())
2010 .and_then(|c| c.get("message"));
2011 let content = msg
2012 .and_then(|m| m.get("content"))
2013 .cloned()
2014 .unwrap_or(Value::Null);
2015 let tool_calls = msg.and_then(|m| m.get("tool_calls")).cloned();
2016 let finish_reason = choices
2017 .and_then(|c| c.first())
2018 .and_then(|c| c.get("finish_reason"))
2019 .cloned()
2020 .unwrap_or_else(|| Value::String("stop".to_owned()));
2021
2022 let mut out = String::new();
2023 let chunk = |delta: Value| {
2024 serde_json::json!({
2025 "id": id,
2026 "object": "chat.completion.chunk",
2027 "created": created,
2028 "model": model,
2029 "choices": [{ "index": 0, "delta": delta, "finish_reason": Value::Null }]
2030 })
2031 };
2032
2033 let role_chunk = chunk(serde_json::json!({ "role": "assistant", "content": "" }));
2035 out.push_str("data: ");
2036 out.push_str(&role_chunk.to_string());
2037 out.push_str("\n\n");
2038
2039 if let Value::String(text) = &content
2041 && !text.is_empty()
2042 {
2043 let content_chunk = chunk(serde_json::json!({ "content": text }));
2044 out.push_str("data: ");
2045 out.push_str(&content_chunk.to_string());
2046 out.push_str("\n\n");
2047 }
2048
2049 if let Some(tc) = tool_calls {
2051 let tc_chunk = chunk(serde_json::json!({ "tool_calls": tc }));
2052 out.push_str("data: ");
2053 out.push_str(&tc_chunk.to_string());
2054 out.push_str("\n\n");
2055 }
2056
2057 let finish_chunk = serde_json::json!({
2059 "id": id,
2060 "object": "chat.completion.chunk",
2061 "created": created,
2062 "model": model,
2063 "choices": [{ "index": 0, "delta": {}, "finish_reason": finish_reason }]
2064 });
2065 out.push_str("data: ");
2066 out.push_str(&finish_chunk.to_string());
2067 out.push_str("\n\n");
2068
2069 out.push_str("data: [DONE]\n\n");
2070 out
2071}
2072
2073async fn handle_enforce_openai(
2078 state: &AppState,
2079 headers: &HeaderMap,
2080 body: &Bytes,
2081 features: Features,
2082 route: &Route,
2083 session_header: Option<String>,
2084 tenant: String,
2085) -> Response {
2086 if is_stream_request(body) {
2087 let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
2088 let (state_c, headers_c, body_c, route_c) =
2089 (state.clone(), headers.clone(), body.clone(), route.clone());
2090 tokio::spawn(async move {
2091 let out = enforce_pipeline_openai(
2092 &state_c,
2093 &headers_c,
2094 &body_c,
2095 features,
2096 &route_c,
2097 session_header,
2098 tenant,
2099 )
2100 .await;
2101 let _ = tx.send(out);
2102 });
2103 return sse_keepalive_response(rx, openai_sse_from_message);
2104 }
2105 match enforce_pipeline_openai(
2106 state,
2107 headers,
2108 body,
2109 features,
2110 route,
2111 session_header,
2112 tenant,
2113 )
2114 .await
2115 {
2116 Ok(message) => (axum::http::StatusCode::OK, Json(message)).into_response(),
2117 Err(e) => e.into_response(),
2118 }
2119}
2120
2121async fn observe_passthrough_openai(
2127 state: AppState,
2128 headers: HeaderMap,
2129 body: Bytes,
2130 session_header: Option<String>,
2131 tenant: String,
2132) -> Response {
2133 if is_stream_request(&body) {
2134 return observe_stream_openai(state, headers, body, session_header, tenant).await;
2135 }
2136 let start = Instant::now();
2137 let result = forward_openai(
2138 &state.http,
2139 &state.config.upstream_openai,
2140 &headers,
2141 body.clone(),
2142 )
2143 .await;
2144 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
2145 match result {
2146 Ok((status, resp_headers, resp_body)) => {
2147 spawn_trace(
2148 &state,
2149 body,
2150 Some(resp_body.clone()),
2151 latency_ms,
2152 session_header,
2153 tenant,
2154 );
2155 (status, resp_headers, resp_body).into_response()
2156 }
2157 Err(err) => {
2158 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
2159 err.into_response()
2160 }
2161 }
2162}
2163
2164async fn observe_stream_openai(
2166 state: AppState,
2167 headers: HeaderMap,
2168 body: Bytes,
2169 session_header: Option<String>,
2170 tenant: String,
2171) -> Response {
2172 let start = Instant::now();
2173 let result = forward_openai_streaming(
2174 &state.http,
2175 &state.config.upstream_openai,
2176 &headers,
2177 body.clone(),
2178 )
2179 .await;
2180 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
2181 match result {
2182 Ok((status, resp_headers, response)) => {
2183 spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
2184 let stream_body = Body::from_stream(response.bytes_stream());
2185 (status, resp_headers, stream_body).into_response()
2186 }
2187 Err(err) => {
2188 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
2189 err.into_response()
2190 }
2191 }
2192}
2193
2194async fn chat_completions(
2199 State(state): State<AppState>,
2200 Extension(TenantId(tenant)): Extension<TenantId>,
2201 headers: HeaderMap,
2202 body: Bytes,
2203) -> Response {
2204 let session_header = header_str(&headers, SESSION_HEADER);
2205
2206 if let Some(routing) = state.config.routing.as_ref() {
2207 let features = extract_openai_features(&headers, &body);
2208 if let Some(route) = routing
2209 .route_for(&features)
2210 .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
2211 {
2212 let route = route.clone();
2213 if enforce_can_handle(
2214 &features,
2215 &body,
2216 routing.escalation.enforce_structured,
2217 &route.ladder,
2218 &state.providers,
2219 Dialect::Openai,
2220 ) {
2221 return handle_enforce_openai(
2222 &state,
2223 &headers,
2224 &body,
2225 features,
2226 &route,
2227 session_header,
2228 tenant,
2229 )
2230 .await;
2231 }
2232 tracing::info!(
2233 "enforce route matched but OpenAI structured request can't be routed faithfully (flag/ladder); serving via observe passthrough"
2234 );
2235 }
2236 }
2237 observe_passthrough_openai(state, headers, body, session_header, tenant).await
2238}
2239
2240#[cfg(test)]
2241mod tests {
2242 use bytes::Bytes;
2243
2244 use super::*;
2245
2246 fn test_config() -> ProxyConfig {
2247 ProxyConfig::from_lookup(|_| None).unwrap()
2248 }
2249
2250 #[test]
2251 fn build_trace_maps_request_and_response_fields() {
2252 let config = test_config();
2253 let req = Bytes::from_static(
2254 br#"{"model":"claude-haiku-4-5","tools":[{"name":"a"}],"messages":[]}"#,
2255 );
2256 let resp = Bytes::from_static(
2257 br#"{"model":"claude-haiku-4-5","usage":{"input_tokens":1200,"output_tokens":300}}"#,
2258 );
2259
2260 let trace = build_trace(&config, &req, &resp, 42, Some("sess-1"));
2261
2262 assert_eq!(trace.request.api, "anthropic.messages");
2263 assert_eq!(trace.session_id, "sess-1");
2264 assert_eq!(trace.attempts.len(), 1);
2265 let attempt = &trace.attempts[0];
2266 assert_eq!(attempt.model, "claude-haiku-4-5");
2267 assert_eq!(attempt.provider, "anthropic");
2268 assert_eq!(attempt.in_tokens, 1200);
2269 assert_eq!(attempt.out_tokens, 300);
2270 assert!(attempt.cost_usd > 0.0);
2271 assert_eq!(trace.request.features.tool_count, 1);
2272 assert!(!trace.request.features.has_images);
2273 assert_eq!(trace.final_.served_rung, Some(0));
2274 }
2275
2276 #[test]
2277 fn build_trace_falls_back_to_trace_id_session_when_header_absent() {
2278 let config = test_config();
2279 let req = Bytes::from_static(b"{}");
2280 let resp = Bytes::from_static(b"{}");
2281
2282 let trace = build_trace(&config, &req, &resp, 1, None);
2283
2284 assert_eq!(trace.session_id, trace.trace_id.to_string());
2285 }
2286
2287 #[test]
2288 fn build_error_trace_has_no_attempts_and_served_from_error() {
2289 let config = test_config();
2290 let req = Bytes::from_static(br#"{"model":"claude-haiku-4-5"}"#);
2291
2292 let trace = build_error_trace(&config, &req, 7, None);
2293
2294 assert!(trace.attempts.is_empty());
2295 assert_eq!(trace.final_.served_from, ServedFrom::Error);
2296 assert_eq!(trace.final_.served_rung, None);
2297 }
2298
2299 #[test]
2300 fn message_with_image_block_sets_has_images() {
2301 let req = br#"{"messages":[{"role":"user","content":[{"type":"image"}]}]}"#;
2302 let (_, _, has_images) = request_features(req);
2303 assert!(has_images);
2304 }
2305
2306 #[test]
2307 fn prompt_hash_never_contains_raw_prompt_text() {
2308 let hash = prompt_hash("salt", b"super secret prompt");
2309 assert!(!hash.contains("secret"));
2310 assert_eq!(hash.len(), 64);
2311 }
2312
2313 #[test]
2314 fn parse_model_request_preserves_content_verbatim_and_projects_text() {
2315 let body = br#"{"model":"m","system":"sys","max_tokens":50,
2316 "messages":[{"role":"user","content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]},
2317 {"role":"assistant","content":"c"}]}"#;
2318 let req = parse_model_request(body).unwrap();
2319 assert_eq!(req.system.as_deref(), Some("sys"));
2320 assert_eq!(req.max_tokens, 50);
2321 assert_eq!(req.messages.len(), 2);
2322 assert_eq!(
2324 req.messages[0].content,
2325 serde_json::json!([{"type":"text","text":"a"},{"type":"text","text":"b"}])
2326 );
2327 assert_eq!(req.messages[1].content, Value::String("c".to_owned()));
2329 assert_eq!(req.messages[0].text_view(), "a\nb");
2331 assert_eq!(req.messages[1].text_view(), "c");
2332 }
2333
2334 #[test]
2335 fn tool_and_image_blocks_survive_the_request_round_trip() {
2336 let body = br#"{"model":"m","max_tokens":50,"messages":[
2338 {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
2339 {"role":"user","content":[
2340 {"type":"tool_result","tool_use_id":"t1","content":"2"},
2341 {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
2342 ]}]}"#;
2343 let req = parse_model_request(body).unwrap();
2344 let round_tripped = serde_json::to_value(&req.messages).unwrap();
2345 assert_eq!(
2346 round_tripped,
2347 serde_json::json!([
2348 {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
2349 {"role":"user","content":[
2350 {"type":"tool_result","tool_use_id":"t1","content":"2"},
2351 {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
2352 ]}
2353 ])
2354 );
2355 }
2356
2357 #[test]
2358 fn text_message_serializes_byte_identical_to_a_plain_string() {
2359 let m = ChatMessage::text("user", "hello");
2361 assert_eq!(
2362 serde_json::to_string(&m).unwrap(),
2363 r#"{"role":"user","content":"hello"}"#
2364 );
2365 }
2366
2367 #[test]
2368 fn parse_model_request_rejects_non_message_bodies() {
2369 assert!(parse_model_request(b"not json").is_none());
2370 assert!(parse_model_request(br#"{"no":"messages"}"#).is_none());
2371 }
2372
2373 use crate::provider::{MockProvider, ModelResponse, Provider, ProviderError, ProviderRegistry};
2376 use axum::extract::State;
2377 use std::collections::HashMap;
2378 use std::sync::Arc;
2379 use tokio::sync::mpsc;
2380
2381 fn model_resp(model: &str, text: &str) -> ModelResponse {
2382 ModelResponse {
2383 model: model.to_owned(),
2384 text: text.to_owned(),
2385 in_tokens: 1000,
2386 out_tokens: 400,
2387 raw: serde_json::Value::Null,
2388 }
2389 }
2390
2391 fn enforce_state(
2394 ladder: &[&str],
2395 gates: &[&str],
2396 outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
2397 ) -> (AppState, mpsc::Receiver<Trace>) {
2398 let toml = format!(
2399 "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [{}]\ngates = [{}]\n",
2400 ladder
2401 .iter()
2402 .map(|m| format!("\"{m}\""))
2403 .collect::<Vec<_>>()
2404 .join(", "),
2405 gates
2406 .iter()
2407 .map(|g| format!("\"{g}\""))
2408 .collect::<Vec<_>>()
2409 .join(", "),
2410 );
2411 let config = ProxyConfig::from_lookup(|k| match k {
2412 "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
2413 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
2414 _ => None,
2415 })
2416 .unwrap();
2417
2418 let mut outs = HashMap::new();
2419 for (model, out) in outcomes {
2420 outs.insert(model.to_owned(), out);
2421 }
2422 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
2423 map.insert(
2424 "anthropic".to_owned(),
2425 Arc::new(MockProvider::new("anthropic", outs)),
2426 );
2427 let providers = ProviderRegistry::from_map(map);
2428
2429 let (traces, rx) = mpsc::channel(64);
2430 let state = AppState {
2431 config: Arc::new(config),
2432 http: reqwest::Client::new(),
2433 providers,
2434 gate_health: Arc::new(GateHealthRegistry::new()),
2435 traces,
2436 adaptive: None,
2437 bandit: None,
2438 tenant_rate_limiter: None,
2439 spill: None,
2440 };
2441 (state, rx)
2442 }
2443
2444 fn user_body() -> Bytes {
2445 Bytes::from_static(
2446 br#"{"model":"ignored","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}"#,
2447 )
2448 }
2449
2450 async fn body_json(resp: Response) -> Value {
2451 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
2452 .await
2453 .unwrap();
2454 serde_json::from_slice(&bytes).unwrap()
2455 }
2456
2457 #[tokio::test]
2458 async fn enforce_serves_first_pass_and_returns_anthropic_shape() {
2459 let (state, mut rx) = enforce_state(
2460 &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
2461 &["non-empty"],
2462 vec![(
2463 "anthropic/claude-haiku-4-5",
2464 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
2465 )],
2466 );
2467 let resp = messages(
2468 State(state),
2469 Extension(TenantId("default".to_owned())),
2470 HeaderMap::new(),
2471 user_body(),
2472 )
2473 .await;
2474 assert_eq!(resp.status(), axum::http::StatusCode::OK);
2475 let json = body_json(resp).await;
2476 assert_eq!(json["type"], "message");
2477 assert_eq!(json["content"][0]["text"], "hello");
2478 assert_eq!(json["model"], "anthropic/claude-haiku-4-5");
2479
2480 let trace = rx.try_recv().expect("a trace was enqueued");
2481 assert_eq!(trace.mode, Mode::Enforce);
2482 assert_eq!(trace.final_.served_rung, Some(0));
2483 assert_eq!(trace.attempts.len(), 1);
2484 }
2485
2486 #[tokio::test]
2487 async fn enforce_escalates_then_serves_and_traces_two_attempts() {
2488 let (state, mut rx) = enforce_state(
2489 &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
2490 &["non-empty"],
2491 vec![
2492 (
2493 "anthropic/claude-haiku-4-5",
2494 Ok(model_resp("anthropic/claude-haiku-4-5", " ")),
2495 ), (
2497 "anthropic/claude-sonnet-5",
2498 Ok(model_resp("anthropic/claude-sonnet-5", "answer")),
2499 ),
2500 ],
2501 );
2502 let resp = messages(
2503 State(state),
2504 Extension(TenantId("default".to_owned())),
2505 HeaderMap::new(),
2506 user_body(),
2507 )
2508 .await;
2509 let json = body_json(resp).await;
2510 assert_eq!(json["content"][0]["text"], "answer");
2511
2512 let trace = rx.try_recv().expect("trace enqueued");
2513 assert_eq!(trace.attempts.len(), 2);
2514 assert_eq!(trace.final_.escalations, 1);
2515 assert_eq!(trace.final_.served_rung, Some(1));
2516 }
2517
2518 #[tokio::test]
2519 async fn enforce_all_rungs_error_returns_502() {
2520 let (state, mut rx) = enforce_state(
2521 &["anthropic/claude-haiku-4-5"],
2522 &["non-empty"],
2523 vec![(
2524 "anthropic/claude-haiku-4-5",
2525 Err(ProviderError::Transport("down".into())),
2526 )],
2527 );
2528 let resp = messages(
2529 State(state),
2530 Extension(TenantId("default".to_owned())),
2531 HeaderMap::new(),
2532 user_body(),
2533 )
2534 .await;
2535 assert_eq!(resp.status(), axum::http::StatusCode::BAD_GATEWAY);
2536 assert!(rx.try_recv().is_ok());
2538 }
2539
2540 #[tokio::test]
2541 async fn no_routing_config_falls_through_to_observe_not_enforce() {
2542 let config = ProxyConfig::from_lookup(|k| match k {
2546 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
2547 _ => None,
2548 })
2549 .unwrap();
2550 let (traces, _rx) = mpsc::channel(64);
2551 let state = AppState {
2552 config: Arc::new(config),
2553 http: reqwest::Client::new(),
2554 providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
2555 gate_health: Arc::new(GateHealthRegistry::new()),
2556 traces,
2557 adaptive: None,
2558 bandit: None,
2559 tenant_rate_limiter: None,
2560 spill: None,
2561 };
2562 let resp = messages(
2563 State(state),
2564 Extension(TenantId("default".to_owned())),
2565 HeaderMap::new(),
2566 user_body(),
2567 )
2568 .await;
2569 assert_ne!(resp.status(), axum::http::StatusCode::OK);
2571 }
2572
2573 #[test]
2574 fn detects_stream_requests() {
2575 assert!(is_stream_request(br#"{"stream": true}"#));
2576 assert!(!is_stream_request(br#"{"stream": false}"#));
2577 assert!(!is_stream_request(br#"{"model":"m"}"#));
2578 assert!(!is_stream_request(b"not json"));
2579 }
2580
2581 #[test]
2582 fn detects_tool_blocks_in_messages() {
2583 let with =
2584 br#"{"messages":[{"role":"user","content":[{"type":"tool_result","content":"42"}]}]}"#;
2585 let without = br#"{"messages":[{"role":"user","content":"hi"}]}"#;
2586 assert!(messages_have_tool_blocks(with));
2587 assert!(!messages_have_tool_blocks(without));
2588 }
2589
2590 #[test]
2591 fn enforce_only_handles_plain_text() {
2592 let plain =
2593 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
2594 let tools = Bytes::from_static(
2595 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
2596 );
2597 let f_plain = extract_features(&HeaderMap::new(), &plain);
2598 let f_tools = extract_features(&HeaderMap::new(), &tools);
2599 let anthropic_ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
2600 let providers = test_registry();
2601 assert!(enforce_can_handle(
2603 &f_plain,
2604 &plain,
2605 false,
2606 &anthropic_ladder,
2607 &providers,
2608 Dialect::Anthropic,
2609 ));
2610 assert!(!enforce_can_handle(
2611 &f_tools,
2612 &tools,
2613 false,
2614 &anthropic_ladder,
2615 &providers,
2616 Dialect::Anthropic,
2617 ));
2618 }
2619
2620 #[test]
2621 fn structured_enforce_routes_tools_and_streaming() {
2622 let tools = Bytes::from_static(
2625 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
2626 );
2627 let streaming_tools = Bytes::from_static(
2628 br#"{"model":"m","stream":true,"tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
2629 );
2630 let f = extract_features(&HeaderMap::new(), &tools);
2631 let anthropic_ladder = vec![
2632 "anthropic/claude-haiku-4-5".to_owned(),
2633 "anthropic/claude-sonnet-5".to_owned(),
2634 ];
2635 let providers = test_registry();
2636 assert!(enforce_can_handle(
2637 &f,
2638 &tools,
2639 true,
2640 &anthropic_ladder,
2641 &providers,
2642 Dialect::Anthropic,
2643 ));
2644 assert!(enforce_can_handle(
2645 &f,
2646 &streaming_tools,
2647 true,
2648 &anthropic_ladder,
2649 &providers,
2650 Dialect::Anthropic,
2651 ));
2652 }
2653
2654 fn test_registry() -> crate::provider::ProviderRegistry {
2656 crate::provider::ProviderRegistry::new("http://localhost", "http://localhost")
2657 }
2658
2659 #[test]
2660 fn fidelity_guard_blocks_structured_on_non_verbatim_ladder() {
2661 let tools = Bytes::from_static(
2665 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
2666 );
2667 let plain =
2668 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
2669 let f_tools = extract_features(&HeaderMap::new(), &tools);
2670 let f_plain = extract_features(&HeaderMap::new(), &plain);
2671 let providers = test_registry();
2672 let mixed_ladder = vec![
2673 "openai/gpt-4.1-mini".to_owned(),
2674 "anthropic/claude-sonnet-5".to_owned(),
2675 ];
2676 assert!(!enforce_can_handle(
2677 &f_tools,
2678 &tools,
2679 true,
2680 &mixed_ladder,
2681 &providers,
2682 Dialect::Anthropic,
2683 ));
2684 assert!(enforce_can_handle(
2685 &f_plain,
2686 &plain,
2687 true,
2688 &mixed_ladder,
2689 &providers,
2690 Dialect::Anthropic,
2691 ));
2692 }
2693
2694 #[test]
2695 fn enforce_sse_reemission_preserves_text_and_tool_use() {
2696 let resp = ModelResponse {
2699 model: "anthropic/claude-haiku-4-5".to_owned(),
2700 text: "let me check".to_owned(),
2701 in_tokens: 5,
2702 out_tokens: 7,
2703 raw: serde_json::json!({
2704 "content": [
2705 { "type": "text", "text": "let me check" },
2706 { "type": "tool_use", "id": "tu_1", "name": "get_weather", "input": { "city": "Paris" } }
2707 ]
2708 }),
2709 };
2710 let sse = anthropic_sse_from_message(&anthropic_response_json(&resp));
2711
2712 let frames: Vec<Value> = sse
2714 .lines()
2715 .filter_map(|l| l.strip_prefix("data: "))
2716 .map(|d| serde_json::from_str::<Value>(d).expect("each SSE data frame is valid JSON"))
2717 .collect();
2718
2719 assert_eq!(frames.first().unwrap()["type"], "message_start");
2721 assert_eq!(frames.last().unwrap()["type"], "message_stop");
2722 assert!(frames.iter().any(|f| f["delta"]["type"] == "text_delta"
2724 && f["delta"]["text"] == "let me check"));
2725 assert!(
2728 frames
2729 .iter()
2730 .any(|f| f["content_block"]["type"] == "tool_use"
2731 && f["content_block"]["name"] == "get_weather"
2732 && f["content_block"]["id"] == "tu_1")
2733 );
2734 assert!(
2735 frames
2736 .iter()
2737 .any(|f| f["delta"]["type"] == "input_json_delta"
2738 && f["delta"]["partial_json"] == r#"{"city":"Paris"}"#)
2739 );
2740 }
2741
2742 #[tokio::test]
2747 async fn enforce_falls_back_to_observe_for_tool_requests() {
2748 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
2749 let config = ProxyConfig::from_lookup(|k| match k {
2750 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
2751 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
2752 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
2753 _ => None,
2754 })
2755 .unwrap();
2756 let mut outs = HashMap::new();
2757 outs.insert(
2758 "anthropic/m".to_owned(),
2759 Ok(model_resp("anthropic/m", "hello")),
2760 );
2761 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
2762 map.insert(
2763 "anthropic".to_owned(),
2764 Arc::new(MockProvider::new("anthropic", outs)),
2765 );
2766 let (traces, _rx) = mpsc::channel(64);
2767 let state = AppState {
2768 config: Arc::new(config),
2769 http: reqwest::Client::new(),
2770 providers: ProviderRegistry::from_map(map),
2771 gate_health: Arc::new(GateHealthRegistry::new()),
2772 traces,
2773 adaptive: None,
2774 bandit: None,
2775 tenant_rate_limiter: None,
2776 spill: None,
2777 };
2778
2779 let plain =
2781 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
2782 let resp = messages(
2783 State(state.clone()),
2784 Extension(TenantId("default".to_owned())),
2785 HeaderMap::new(),
2786 plain,
2787 )
2788 .await;
2789 assert_eq!(
2790 resp.status(),
2791 axum::http::StatusCode::OK,
2792 "plain text should enforce"
2793 );
2794
2795 let tools = Bytes::from_static(
2798 br#"{"model":"m","tools":[{"name":"get_weather"}],"messages":[{"role":"user","content":"hi"}]}"#,
2799 );
2800 let resp = messages(
2801 State(state.clone()),
2802 Extension(TenantId("default".to_owned())),
2803 HeaderMap::new(),
2804 tools.clone(),
2805 )
2806 .await;
2807 assert_eq!(
2808 resp.status(),
2809 axum::http::StatusCode::OK,
2810 "tool request must route through enforce by default (ADR 0005 default-on)"
2811 );
2812
2813 let toml_off = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n[escalation]\nenforce_structured = false\n";
2816 let config_off = ProxyConfig::from_lookup(|k| match k {
2817 "FIRSTPASS_CONFIG_TOML" => Some(toml_off.to_owned()),
2818 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
2819 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
2820 _ => None,
2821 })
2822 .unwrap();
2823 let state_off = AppState {
2824 config: Arc::new(config_off),
2825 ..state.clone()
2826 };
2827 let resp = messages(
2828 State(state_off),
2829 Extension(TenantId("default".to_owned())),
2830 HeaderMap::new(),
2831 tools,
2832 )
2833 .await;
2834 assert_ne!(
2835 resp.status(),
2836 axum::http::StatusCode::OK,
2837 "with enforce_structured = false a tool request must fall back to observe"
2838 );
2839
2840 let toolres = Bytes::from_static(
2842 br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"x","content":"42"}]}]}"#,
2843 );
2844 let resp = messages(
2845 State(state),
2846 Extension(TenantId("default".to_owned())),
2847 HeaderMap::new(),
2848 toolres,
2849 )
2850 .await;
2851 assert_eq!(
2852 resp.status(),
2853 axum::http::StatusCode::OK,
2854 "tool_result blocks route through enforce by default too (verbatim carry)"
2855 );
2856 }
2857
2858 async fn feedback_state() -> (AppState, std::path::PathBuf, String) {
2862 let db = std::env::temp_dir().join(format!("firstpass-feedback-{}.db", Uuid::now_v7()));
2863 let (tx, handle) = crate::store::open(&db).unwrap();
2864
2865 let mut trace = build_error_trace(
2866 &ProxyConfig::from_lookup(|_| None).unwrap(),
2867 &Bytes::from_static(b"{}"),
2868 5,
2869 Some("sess-fb"),
2870 );
2871 trace.attempts.push(Attempt {
2872 rung: 0,
2873 model: "anthropic/claude-haiku-4-5".into(),
2874 provider: "anthropic".into(),
2875 in_tokens: 10,
2876 out_tokens: 5,
2877 cost_usd: 0.001,
2878 latency_ms: 5,
2879 gates: vec![],
2880 verdict: Verdict::Pass,
2881 });
2882 let trace_id = trace.trace_id.to_string();
2883 tx.try_send(trace).unwrap();
2884 drop(tx);
2885 handle.await.unwrap();
2886
2887 let db_str = db.to_string_lossy().into_owned();
2888 let config = ProxyConfig::from_lookup(move |k| match k {
2889 "FIRSTPASS_DB" => Some(db_str.clone()),
2890 _ => None,
2891 })
2892 .unwrap();
2893 let (traces, _rx) = mpsc::channel(64);
2894 let state = AppState {
2895 config: Arc::new(config),
2896 http: reqwest::Client::new(),
2897 providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
2898 gate_health: Arc::new(GateHealthRegistry::new()),
2899 traces,
2900 adaptive: None,
2901 bandit: None,
2902 tenant_rate_limiter: None,
2903 spill: None,
2904 };
2905 (state, db, trace_id)
2906 }
2907
2908 #[tokio::test]
2909 async fn feedback_nudges_the_adaptive_threshold() {
2910 use firstpass_core::conformal::AdaptiveConformal;
2911 let (mut state, _db, trace_id) = feedback_state().await;
2912 let aci = Arc::new(std::sync::Mutex::new(AdaptiveConformal::new(0.1, 0.2, 0.5)));
2913 state.adaptive = Some(aci.clone());
2914 let before = aci.lock().unwrap().threshold();
2915
2916 let fail = Bytes::from(
2918 serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "fail", "reporter": "ci" })
2919 .to_string(),
2920 );
2921 assert_eq!(
2922 feedback(
2923 State(state.clone()),
2924 Extension(TenantId("default".to_owned())),
2925 fail
2926 )
2927 .await
2928 .status(),
2929 axum::http::StatusCode::ACCEPTED
2930 );
2931 let after_fail = aci.lock().unwrap().threshold();
2932 assert!(
2933 after_fail > before,
2934 "served fail should raise the live threshold: {before} -> {after_fail}"
2935 );
2936
2937 let pass = Bytes::from(
2939 serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "pass", "reporter": "ci" })
2940 .to_string(),
2941 );
2942 let _ = feedback(
2943 State(state),
2944 Extension(TenantId("default".to_owned())),
2945 pass,
2946 )
2947 .await;
2948 assert!(aci.lock().unwrap().threshold() < after_fail);
2949 }
2950
2951 #[tokio::test]
2952 async fn feedback_records_a_deferred_verdict_without_breaking_the_chain() {
2953 let (state, db, trace_id) = feedback_state().await;
2954 let body = Bytes::from(
2955 serde_json::json!({
2956 "trace_id": trace_id,
2957 "gate_id": "tests",
2958 "verdict": "pass",
2959 "score": 1.0,
2960 "reporter": "ci",
2961 })
2962 .to_string(),
2963 );
2964 let resp = feedback(
2965 State(state),
2966 Extension(TenantId("default".to_owned())),
2967 body,
2968 )
2969 .await;
2970 assert_eq!(resp.status(), axum::http::StatusCode::ACCEPTED);
2971
2972 let view = crate::store::load_trace_view(&db, "default", &trace_id)
2974 .unwrap()
2975 .unwrap();
2976 assert_eq!(view.deferred.len(), 1);
2977 assert_eq!(view.deferred[0].gate_id, "tests");
2978 let traces = crate::store::load_all_traces(&db).unwrap();
2980 firstpass_core::verify_chain(&traces, GENESIS_HASH).unwrap();
2981
2982 let _ = std::fs::remove_file(&db);
2983 }
2984
2985 #[tokio::test]
2986 async fn feedback_for_unknown_trace_is_404() {
2987 let (state, db, _trace_id) = feedback_state().await;
2988 let body = Bytes::from(
2989 serde_json::json!({
2990 "trace_id": "does-not-exist",
2991 "gate_id": "tests",
2992 "verdict": "pass",
2993 "reporter": "ci",
2994 })
2995 .to_string(),
2996 );
2997 let resp = feedback(
2998 State(state),
2999 Extension(TenantId("default".to_owned())),
3000 body,
3001 )
3002 .await;
3003 assert_eq!(resp.status(), axum::http::StatusCode::NOT_FOUND);
3004 let _ = std::fs::remove_file(&db);
3005 }
3006
3007 #[tokio::test]
3010 async fn feedback_across_tenants_is_404_not_403() {
3011 let (state, db, trace_id) = feedback_state().await;
3012 let body = Bytes::from(
3013 serde_json::json!({
3014 "trace_id": trace_id,
3015 "gate_id": "tests",
3016 "verdict": "pass",
3017 "score": 1.0,
3018 "reporter": "attacker",
3019 })
3020 .to_string(),
3021 );
3022 let resp = feedback(
3024 State(state),
3025 Extension(TenantId("tenant-b".to_owned())),
3026 body,
3027 )
3028 .await;
3029 assert_eq!(
3030 resp.status(),
3031 axum::http::StatusCode::NOT_FOUND,
3032 "cross-tenant feedback must look exactly like a missing trace"
3033 );
3034 let _ = std::fs::remove_file(&db);
3035 }
3036
3037 #[tokio::test]
3038 async fn feedback_rejects_bad_verdict_and_score() {
3039 let (state, db, trace_id) = feedback_state().await;
3040 let bad_verdict = Bytes::from(
3041 serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "maybe", "reporter": "x" })
3042 .to_string(),
3043 );
3044 assert_eq!(
3045 feedback(
3046 State(state.clone()),
3047 Extension(TenantId("default".to_owned())),
3048 bad_verdict
3049 )
3050 .await
3051 .status(),
3052 axum::http::StatusCode::BAD_REQUEST
3053 );
3054 let bad_score = Bytes::from(
3055 serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "pass", "score": 9.0, "reporter": "x" })
3056 .to_string(),
3057 );
3058 assert_eq!(
3059 feedback(
3060 State(state),
3061 Extension(TenantId("default".to_owned())),
3062 bad_score
3063 )
3064 .await
3065 .status(),
3066 axum::http::StatusCode::BAD_REQUEST
3067 );
3068 let _ = std::fs::remove_file(&db);
3069 }
3070
3071 #[tokio::test]
3072 async fn metrics_endpoint_renders_after_a_real_request() {
3073 use tower::ServiceExt;
3074
3075 let (state, mut rx) = enforce_state(
3076 &["anthropic/claude-haiku-4-5"],
3077 &["non-empty"],
3078 vec![(
3079 "anthropic/claude-haiku-4-5",
3080 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
3081 )],
3082 );
3083 let router = app(state).expect("prometheus recorder installs");
3084
3085 let req = axum::http::Request::builder()
3086 .method("POST")
3087 .uri("/v1/messages")
3088 .header("content-type", "application/json")
3089 .body(Body::from(user_body()))
3090 .unwrap();
3091 let resp = router.clone().oneshot(req).await.unwrap();
3092 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3093 rx.try_recv().expect("a trace was enqueued");
3094
3095 let metrics_req = axum::http::Request::builder()
3096 .method("GET")
3097 .uri("/metrics")
3098 .body(Body::empty())
3099 .unwrap();
3100 let metrics_resp = router.oneshot(metrics_req).await.unwrap();
3101 assert_eq!(metrics_resp.status(), axum::http::StatusCode::OK);
3102 let bytes = axum::body::to_bytes(metrics_resp.into_body(), 1 << 20)
3103 .await
3104 .unwrap();
3105 let body = String::from_utf8(bytes.to_vec()).unwrap();
3106 assert!(
3107 body.contains("firstpass_enforce_latency_ms"),
3108 "metrics body missing enforce latency histogram: {body}"
3109 );
3110 assert!(
3111 body.contains("firstpass_served_total"),
3112 "metrics body missing served counter: {body}"
3113 );
3114 }
3115
3116 fn auth_state(require_auth: bool, keys_json: Option<String>) -> AppState {
3120 auth_state_rated(require_auth, keys_json, None)
3121 }
3122
3123 fn auth_state_rated(
3126 require_auth: bool,
3127 keys_json: Option<String>,
3128 rate_per_sec: Option<u32>,
3129 ) -> AppState {
3130 let config = ProxyConfig::from_lookup(|k| match k {
3131 "FIRSTPASS_REQUIRE_AUTH" => require_auth.then(|| "true".to_owned()),
3132 "FIRSTPASS_TENANT_KEYS_JSON" => keys_json.clone(),
3133 "FIRSTPASS_TENANT_RATE_PER_SEC" => rate_per_sec.map(|n| n.to_string()),
3134 _ => None,
3135 })
3136 .unwrap();
3137 let (traces, _rx) = mpsc::channel(64);
3138 std::mem::forget(_rx);
3141 let providers: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3142 let tenant_rate_limiter = build_tenant_rate_limiter(&config);
3143 AppState {
3144 config: Arc::new(config),
3145 http: reqwest::Client::new(),
3146 providers: ProviderRegistry::from_map(providers),
3147 gate_health: Arc::new(GateHealthRegistry::new()),
3148 traces,
3149 adaptive: None,
3150 bandit: None,
3151 tenant_rate_limiter,
3152 spill: None,
3153 }
3154 }
3155
3156 fn cap_request(auth_header: Option<&str>) -> axum::http::Request<Body> {
3157 let mut b = axum::http::Request::builder()
3158 .method("GET")
3159 .uri("/v1/capabilities");
3160 if let Some(h) = auth_header {
3161 b = b.header("authorization", h);
3162 }
3163 b.body(Body::empty()).unwrap()
3164 }
3165
3166 #[tokio::test]
3167 async fn auth_off_allows_unauthenticated_request() {
3168 use tower::ServiceExt;
3169 let router = app(auth_state(false, None)).expect("router");
3170 let resp = router.oneshot(cap_request(None)).await.unwrap();
3171 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3173 }
3174
3175 #[tokio::test]
3176 async fn auth_on_missing_key_is_401_opaque() {
3177 use tower::ServiceExt;
3178 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3179 let keys = format!("{{\"tenant-a\": {hash:?}}}");
3180 let router = app(auth_state(true, Some(keys))).expect("router");
3181
3182 let resp = router.oneshot(cap_request(None)).await.unwrap();
3183 assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
3184 let json = body_json(resp).await;
3185 assert_eq!(json["error"]["type"], "unauthorized");
3186 let msg = json["error"]["message"].as_str().unwrap();
3188 assert!(!msg.contains("tenant"), "no tenant oracle in body: {msg}");
3189 }
3190
3191 #[tokio::test]
3192 async fn auth_on_invalid_key_is_401() {
3193 use tower::ServiceExt;
3194 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3195 let keys = format!("{{\"tenant-a\": {hash:?}}}");
3196 let router = app(auth_state(true, Some(keys))).expect("router");
3197
3198 let resp = router
3199 .oneshot(cap_request(Some("Bearer wrong-key")))
3200 .await
3201 .unwrap();
3202 assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
3203 }
3204
3205 #[tokio::test]
3206 async fn auth_on_valid_key_proceeds() {
3207 use tower::ServiceExt;
3208 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3209 let keys = format!("{{\"tenant-a\": {hash:?}}}");
3210 let router = app(auth_state(true, Some(keys))).expect("router");
3211
3212 let resp = router
3214 .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
3215 .await
3216 .unwrap();
3217 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3219 }
3220
3221 fn two_tenant_state(rate_per_sec: Option<u32>) -> AppState {
3223 let hash_a = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
3224 let hash_b = crate::tenant_auth::TenantKeys::hash_key("key-b").unwrap();
3225 let keys = format!("{{\"tenant-a\": {hash_a:?}, \"tenant-b\": {hash_b:?}}}");
3226 auth_state_rated(true, Some(keys), rate_per_sec)
3227 }
3228
3229 #[tokio::test]
3230 async fn tenant_exceeding_rate_limit_gets_429_opaque() {
3231 use tower::ServiceExt;
3232 let router = app(two_tenant_state(Some(1))).expect("router");
3238
3239 let (r1, r2, r3, r4) = tokio::join!(
3240 router
3241 .clone()
3242 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3243 router
3244 .clone()
3245 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3246 router
3247 .clone()
3248 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3249 router
3250 .clone()
3251 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3252 );
3253 let responses = [r1.unwrap(), r2.unwrap(), r3.unwrap(), r4.unwrap()];
3254 let ok = responses
3255 .iter()
3256 .filter(|r| r.status() == axum::http::StatusCode::OK)
3257 .count();
3258 assert!(ok >= 1, "the burst's first request must pass");
3259 let limited: Vec<_> = responses
3260 .into_iter()
3261 .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
3262 .collect();
3263 assert!(
3264 !limited.is_empty(),
3265 "a 4-request burst against 1 req/sec must trip the limiter"
3266 );
3267
3268 let json = body_json(limited.into_iter().next().unwrap()).await;
3269 assert_eq!(json["error"]["type"], "rate_limited");
3270 let msg = json["error"]["message"].as_str().unwrap();
3272 assert!(!msg.contains('1'), "no limit value in body: {msg}");
3273 }
3274
3275 #[tokio::test]
3276 async fn rate_limit_buckets_are_independent_per_tenant() {
3277 use tower::ServiceExt;
3278 let router = app(two_tenant_state(Some(1))).expect("router");
3279
3280 let (a1, a2, a3) = tokio::join!(
3283 router
3284 .clone()
3285 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3286 router
3287 .clone()
3288 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3289 router
3290 .clone()
3291 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
3292 );
3293 let a_limited = [a1.unwrap(), a2.unwrap(), a3.unwrap()]
3294 .iter()
3295 .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
3296 .count();
3297 assert!(a_limited >= 1, "tenant A's burst must trip its limiter");
3298
3299 let b1 = router
3301 .clone()
3302 .oneshot(cap_request(Some("Bearer tenant-b.key-b")))
3303 .await
3304 .unwrap();
3305 assert_eq!(b1.status(), axum::http::StatusCode::OK);
3306 }
3307
3308 #[tokio::test]
3309 async fn rate_limit_unset_never_429s() {
3310 use tower::ServiceExt;
3311 let router = app(two_tenant_state(None)).expect("router");
3314 for _ in 0..20 {
3315 let resp = router
3316 .clone()
3317 .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
3318 .await
3319 .unwrap();
3320 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3321 }
3322 }
3323
3324 #[test]
3327 fn u01_is_deterministic_and_in_range() {
3328 let s1 = u01(0xDEAD_BEEF_CAFE_1234_u128);
3329 let s2 = u01(0xDEAD_BEEF_CAFE_1234_u128);
3330 assert_eq!(s1, s2, "u01 must be deterministic for the same seed");
3331 assert!((0.0..1.0).contains(&s1), "u01 must return [0, 1), got {s1}");
3332
3333 let s3 = u01(0x1234_5678_9ABC_DEF0_u128);
3335 assert_ne!(s1, s3, "different seeds should give different values");
3336
3337 for i in 0u64..256 {
3339 let v = u01(i as u128);
3340 assert!((0.0..1.0).contains(&v), "seed {i}: u01={v} out of [0,1)");
3341 }
3342 }
3343
3344 #[test]
3345 fn epsilon_propensity_formula() {
3346 let epsilon = 0.2_f64;
3347 let k = 3_usize;
3348 let greedy = 1_u32;
3349
3350 let p_greedy = epsilon_propensity(greedy, greedy, epsilon, k);
3352 let expected_greedy = (1.0 - epsilon) + epsilon / k as f64;
3353 assert!(
3354 (p_greedy - expected_greedy).abs() < 1e-12,
3355 "{p_greedy} != {expected_greedy}"
3356 );
3357
3358 let p_other = epsilon_propensity(0, greedy, epsilon, k);
3360 let expected_other = epsilon / k as f64;
3361 assert!(
3362 (p_other - expected_other).abs() < 1e-12,
3363 "{p_other} != {expected_other}"
3364 );
3365
3366 for chosen in 0..k as u32 {
3368 let p = epsilon_propensity(chosen, greedy, epsilon, k);
3369 assert!(
3370 p > 0.0 && p <= 1.0,
3371 "propensity {p} out of (0,1] for chosen={chosen}"
3372 );
3373 }
3374 }
3375
3376 #[test]
3377 fn epsilon_branch_and_greedy_branch_both_occur_over_many_seeds() {
3378 let epsilon = 0.3_f64;
3380 let mut saw_explore = false;
3381 let mut saw_greedy = false;
3382 for i in 0u64..200 {
3383 let u = u01(i as u128);
3384 if u < epsilon {
3385 saw_explore = true;
3386 } else {
3387 saw_greedy = true;
3388 }
3389 if saw_explore && saw_greedy {
3390 break;
3391 }
3392 }
3393 assert!(
3394 saw_explore,
3395 "epsilon branch must fire with epsilon=0.3 over 200 seeds"
3396 );
3397 assert!(
3398 saw_greedy,
3399 "greedy branch must occur with epsilon=0.3 over 200 seeds"
3400 );
3401 }
3402
3403 #[test]
3404 fn epsilon_propensity_sums_to_one_over_all_rungs() {
3405 let epsilon = 0.15_f64;
3407 let k = 4_usize;
3408 let greedy = 2_u32;
3409 let total: f64 = (0..k as u32)
3410 .map(|r| epsilon_propensity(r, greedy, epsilon, k))
3411 .sum();
3412 assert!(
3413 (total - 1.0).abs() < 1e-12,
3414 "propensities must sum to 1, got {total}"
3415 );
3416 }
3417
3418 #[tokio::test(start_paused = true)]
3421 async fn keepalive_stream_ticks_then_emits_final_frame() {
3422 let (tx, rx) = tokio::sync::oneshot::channel::<Result<Value, ProxyError>>();
3423 let mut ticks = tokio::time::interval(SSE_KEEPALIVE_EVERY);
3424 ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
3425 ticks.reset();
3426 let mut stream = KeepaliveStream {
3427 rx: Some(rx),
3428 ticks,
3429 format_message: anthropic_sse_from_message,
3430 };
3431 async fn next(
3433 stream: &mut KeepaliveStream,
3434 ) -> Option<Option<Result<Bytes, std::convert::Infallible>>> {
3435 std::future::poll_fn(|cx| {
3436 std::task::Poll::Ready(
3437 match futures_core::Stream::poll_next(std::pin::Pin::new(&mut *stream), cx) {
3438 std::task::Poll::Ready(item) => Some(item),
3439 std::task::Poll::Pending => None,
3440 },
3441 )
3442 })
3443 .await
3444 }
3445
3446 assert!(
3448 next(&mut stream).await.is_none(),
3449 "no frame before an interval"
3450 );
3451
3452 tokio::time::advance(SSE_KEEPALIVE_EVERY + Duration::from_millis(1)).await;
3454 let frame = next(&mut stream)
3455 .await
3456 .expect("keepalive due")
3457 .unwrap()
3458 .unwrap();
3459 assert!(
3460 frame.starts_with(b": "),
3461 "keepalive must be an SSE comment (ignored by every conforming parser)"
3462 );
3463
3464 let message = serde_json::json!({
3466 "id": "msg_1", "type": "message", "role": "assistant", "model": "m",
3467 "content": [{ "type": "text", "text": "done" }],
3468 "usage": { "input_tokens": 1, "output_tokens": 1 }
3469 });
3470 tx.send(Ok(message)).unwrap();
3471 let frame = next(&mut stream)
3472 .await
3473 .expect("final frame")
3474 .unwrap()
3475 .unwrap();
3476 let text = String::from_utf8(frame.to_vec()).unwrap();
3477 assert!(text.contains("event: message_start"));
3478 assert!(text.contains("event: message_stop"));
3479 let eos = next(&mut stream)
3480 .await
3481 .expect("stream must end after the final frame");
3482 assert!(eos.is_none(), "end-of-stream after the final frame");
3483 }
3484
3485 #[tokio::test]
3488 async fn streaming_enforce_serves_full_sse_sequence() {
3489 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
3490 let config = ProxyConfig::from_lookup(|k| match k {
3491 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
3492 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
3493 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
3494 _ => None,
3495 })
3496 .unwrap();
3497 let mut outs = HashMap::new();
3498 outs.insert(
3499 "anthropic/m".to_owned(),
3500 Ok(model_resp("anthropic/m", "gated answer")),
3501 );
3502 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3503 map.insert(
3504 "anthropic".to_owned(),
3505 Arc::new(MockProvider::new("anthropic", outs)),
3506 );
3507 let (traces, _rx) = mpsc::channel(64);
3508 let state = AppState {
3509 config: Arc::new(config),
3510 http: reqwest::Client::new(),
3511 providers: ProviderRegistry::from_map(map),
3512 gate_health: Arc::new(GateHealthRegistry::new()),
3513 traces,
3514 adaptive: None,
3515 bandit: None,
3516 tenant_rate_limiter: None,
3517 spill: None,
3518 };
3519 let body = Bytes::from_static(
3520 br#"{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}"#,
3521 );
3522 let resp = messages(
3523 State(state),
3524 Extension(TenantId("default".to_owned())),
3525 HeaderMap::new(),
3526 body,
3527 )
3528 .await;
3529 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3530 assert!(
3531 resp.headers()
3532 .get(axum::http::header::CONTENT_TYPE)
3533 .and_then(|v| v.to_str().ok())
3534 .is_some_and(|ct| ct.starts_with("text/event-stream")),
3535 "streaming client must get SSE"
3536 );
3537 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
3538 .await
3539 .unwrap();
3540 let text = String::from_utf8(bytes.to_vec()).unwrap();
3541 assert!(text.contains("event: message_start"));
3542 assert!(text.contains("gated answer"));
3543 assert!(text.contains("event: message_stop"));
3544 }
3545
3546 #[test]
3551 fn parse_openai_request_plain_text() {
3552 let body = br#"{"model":"gpt-4o","max_tokens":256,"messages":[{"role":"user","content":"hello"}]}"#;
3554 let req = parse_openai_request(body, false).expect("must parse");
3555 assert_eq!(req.model, "gpt-4o");
3556 assert_eq!(req.max_tokens, 256);
3557 assert_eq!(req.messages.len(), 1);
3558 assert_eq!(req.messages[0].role, "user");
3559 assert_eq!(req.messages[0].content, Value::String("hello".to_owned()));
3560 assert!(req.system.is_none());
3561 assert_eq!(req.raw, Value::Null);
3563 }
3564
3565 #[test]
3566 fn parse_openai_request_system_message() {
3567 let body = br#"{"model":"gpt-4o","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hi"}]}"#;
3568 let req = parse_openai_request(body, false).expect("must parse");
3569 assert_eq!(req.system.as_deref(), Some("be concise"));
3570 assert_eq!(req.messages.len(), 1);
3571 assert_eq!(req.messages[0].role, "user");
3572 }
3573
3574 #[test]
3575 fn parse_openai_request_tool_calls_translate_to_tool_use() {
3576 let body = br#"{
3577 "model":"gpt-4o",
3578 "messages":[
3579 {"role":"user","content":"what's the weather?"},
3580 {"role":"assistant","content":null,"tool_calls":[
3581 {"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":\"Paris\"}"}}
3582 ]},
3583 {"role":"tool","tool_call_id":"call_1","content":"15C, cloudy"}
3584 ]
3585 }"#;
3586 let req = parse_openai_request(body, false).expect("must parse");
3587 assert_eq!(req.messages.len(), 3);
3589 let asst = &req.messages[1];
3591 assert_eq!(asst.role, "assistant");
3592 let blocks = asst.content.as_array().expect("content array");
3593 assert_eq!(blocks[0]["type"], "tool_use");
3594 assert_eq!(blocks[0]["name"], "get_weather");
3595 assert_eq!(blocks[0]["id"], "call_1");
3596 assert_eq!(blocks[0]["input"]["city"], "Paris");
3597 let tool_msg = &req.messages[2];
3599 assert_eq!(tool_msg.role, "user");
3600 let result_blocks = tool_msg.content.as_array().expect("result blocks");
3601 assert_eq!(result_blocks[0]["type"], "tool_result");
3602 assert_eq!(result_blocks[0]["tool_use_id"], "call_1");
3603 }
3604
3605 #[test]
3606 fn parse_openai_request_tools_translate_to_anthropic_format() {
3607 let body = br#"{
3608 "model":"gpt-4o",
3609 "messages":[{"role":"user","content":"use a tool"}],
3610 "tools":[{"type":"function","function":{"name":"search","description":"web search","parameters":{"type":"object","properties":{"q":{"type":"string"}}}}}]
3611 }"#;
3612 let req = parse_openai_request(body, false).expect("must parse");
3613 let tools = req.tools.as_array().expect("tools array");
3614 assert_eq!(tools.len(), 1);
3615 assert_eq!(tools[0]["name"], "search");
3616 assert_eq!(tools[0]["description"], "web search");
3617 assert_eq!(tools[0]["input_schema"]["type"], "object");
3618 }
3619
3620 #[test]
3621 fn parse_openai_request_raw_carry_preserves_full_body() {
3622 let body =
3624 br#"{"model":"gpt-4o","max_tokens":100,"messages":[{"role":"user","content":"hi"}]}"#;
3625 let req = parse_openai_request(body, true).expect("must parse");
3626 assert!(req.raw.is_object(), "raw must be the full JSON object");
3627 assert_eq!(req.raw["model"], "gpt-4o");
3628 assert_eq!(req.raw["max_tokens"], 100);
3629 assert!(req.tools.is_null(), "no tools in this request");
3631 }
3632
3633 #[test]
3634 fn parse_openai_request_http_image_returns_none() {
3635 let body = br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
3637 {"type":"text","text":"describe this"},
3638 {"type":"image_url","image_url":{"url":"https://example.com/cat.png"}}
3639 ]}]}"#;
3640 let result = parse_openai_request(body, false);
3641 assert!(result.is_none(), "http image URL must fail translation");
3642 }
3643
3644 #[test]
3645 fn parse_openai_request_data_url_image_translates_to_anthropic_base64() {
3646 let body = br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
3647 {"type":"text","text":"describe"},
3648 {"type":"image_url","image_url":{"url":"data:image/png;base64,iVBORw0KGgo="}}
3649 ]}]}"#;
3650 let req = parse_openai_request(body, false).expect("data URL must parse");
3651 let blocks = req.messages[0].content.as_array().expect("blocks");
3652 let img = blocks
3653 .iter()
3654 .find(|b| b["type"] == "image")
3655 .expect("image block");
3656 assert_eq!(img["source"]["type"], "base64");
3657 assert_eq!(img["source"]["media_type"], "image/png");
3658 assert_eq!(img["source"]["data"], "iVBORw0KGgo=");
3659 }
3660
3661 #[test]
3664 fn openai_response_json_renders_text_response() {
3665 let resp = ModelResponse {
3666 model: "gpt-4o".to_owned(),
3667 text: "Hello!".to_owned(),
3668 in_tokens: 10,
3669 out_tokens: 5,
3670 raw: serde_json::json!({
3671 "content": [{ "type": "text", "text": "Hello!" }]
3672 }),
3673 };
3674 let json = openai_response_json(&resp);
3675 assert_eq!(json["object"], "chat.completion");
3676 assert_eq!(json["model"], "gpt-4o");
3677 assert_eq!(json["choices"][0]["message"]["role"], "assistant");
3678 assert_eq!(json["choices"][0]["message"]["content"], "Hello!");
3679 assert_eq!(json["choices"][0]["finish_reason"], "stop");
3680 assert_eq!(json["usage"]["prompt_tokens"], 10);
3681 assert_eq!(json["usage"]["completion_tokens"], 5);
3682 }
3683
3684 #[test]
3685 fn openai_response_json_renders_tool_call() {
3686 let resp = ModelResponse {
3687 model: "gpt-4o".to_owned(),
3688 text: String::new(),
3689 in_tokens: 20,
3690 out_tokens: 15,
3691 raw: serde_json::json!({
3692 "content": [{
3693 "type": "tool_use",
3694 "id": "call_abc",
3695 "name": "search",
3696 "input": {"q": "Rust async"}
3697 }]
3698 }),
3699 };
3700 let json = openai_response_json(&resp);
3701 assert_eq!(json["choices"][0]["finish_reason"], "tool_calls");
3702 let tc = &json["choices"][0]["message"]["tool_calls"][0];
3703 assert_eq!(tc["id"], "call_abc");
3704 assert_eq!(tc["type"], "function");
3705 assert_eq!(tc["function"]["name"], "search");
3706 assert_eq!(json["choices"][0]["message"]["content"], Value::Null);
3708 }
3709
3710 #[test]
3711 fn openai_sse_from_message_plain_text_has_role_then_content_then_stop() {
3712 let resp = ModelResponse {
3713 model: "gpt-4o".to_owned(),
3714 text: "Hi there!".to_owned(),
3715 in_tokens: 5,
3716 out_tokens: 3,
3717 raw: serde_json::json!({
3718 "content": [{ "type": "text", "text": "Hi there!" }]
3719 }),
3720 };
3721 let sse = openai_sse_from_message(&openai_response_json(&resp));
3722 for line in sse.lines() {
3724 assert!(
3725 line.is_empty() || line.starts_with("data: "),
3726 "bad SSE line: {line:?}"
3727 );
3728 }
3729 let frames: Vec<&str> = sse
3730 .lines()
3731 .filter_map(|l| l.strip_prefix("data: "))
3732 .collect();
3733 assert_eq!(*frames.last().unwrap(), "[DONE]");
3735 let role_frame: Value = serde_json::from_str(frames[0]).unwrap();
3737 assert_eq!(role_frame["choices"][0]["delta"]["role"], "assistant");
3738 assert!(frames.iter().any(|f| {
3740 if *f == "[DONE]" {
3741 return false;
3742 }
3743 serde_json::from_str::<Value>(f)
3744 .ok()
3745 .is_some_and(|v| v["choices"][0]["delta"]["content"] == "Hi there!")
3746 }));
3747 assert!(frames.iter().any(|f| {
3749 if *f == "[DONE]" {
3750 return false;
3751 }
3752 serde_json::from_str::<Value>(f)
3753 .ok()
3754 .is_some_and(|v| v["choices"][0]["finish_reason"] == "stop")
3755 }));
3756 }
3757
3758 #[test]
3761 fn detects_openai_tool_calls() {
3762 let with_tool_calls = Bytes::from_static(br#"{"messages":[
3763 {"role":"user","content":"hi"},
3764 {"role":"assistant","content":null,"tool_calls":[{"id":"c1","type":"function","function":{"name":"f","arguments":"{}"}}]}
3765 ]}"#);
3766 let without = Bytes::from_static(br#"{"messages":[{"role":"user","content":"hi"}]}"#);
3767 let with_tool_msg = Bytes::from_static(
3768 br#"{"messages":[
3769 {"role":"tool","tool_call_id":"c1","content":"result"}
3770 ]}"#,
3771 );
3772 assert!(openai_messages_have_tool_calls(&with_tool_calls));
3773 assert!(!openai_messages_have_tool_calls(&without));
3774 assert!(openai_messages_have_tool_calls(&with_tool_msg));
3775 }
3776
3777 #[test]
3778 fn detects_openai_http_images() {
3779 let http_img = Bytes::from_static(
3780 br#"{"messages":[{"role":"user","content":[
3781 {"type":"image_url","image_url":{"url":"https://example.com/img.png"}}
3782 ]}]}"#,
3783 );
3784 let data_img = Bytes::from_static(
3785 br#"{"messages":[{"role":"user","content":[
3786 {"type":"image_url","image_url":{"url":"data:image/png;base64,abc"}}
3787 ]}]}"#,
3788 );
3789 let no_img = Bytes::from_static(br#"{"messages":[{"role":"user","content":"hi"}]}"#);
3790 assert!(openai_has_http_images(&http_img));
3791 assert!(!openai_has_http_images(&data_img));
3792 assert!(!openai_has_http_images(&no_img));
3793 }
3794
3795 #[test]
3796 fn enforce_can_handle_openai_inbound_all_openai_ladder() {
3797 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":"{}"}}]}]}"#);
3799 let f = extract_openai_features(&HeaderMap::new(), &tools_body);
3800 let ladder = vec!["openai/gpt-4o-mini".to_owned(), "openai/gpt-4o".to_owned()];
3801 let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
3802 assert!(enforce_can_handle(
3803 &f,
3804 &tools_body,
3805 true,
3806 &ladder,
3807 &providers,
3808 Dialect::Openai
3809 ));
3810 }
3811
3812 #[test]
3813 fn enforce_can_handle_openai_inbound_all_anthropic_ladder_no_http_image() {
3814 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":"{}"}}]}]}"#);
3816 let f = extract_openai_features(&HeaderMap::new(), &tools_body);
3817 let ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
3818 let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
3819 assert!(enforce_can_handle(
3820 &f,
3821 &tools_body,
3822 true,
3823 &ladder,
3824 &providers,
3825 Dialect::Openai,
3826 ));
3827 }
3828
3829 #[test]
3830 fn enforce_can_handle_openai_inbound_http_image_falls_back() {
3831 let img_body = Bytes::from_static(
3833 br#"{"model":"gpt-4o","messages":[{"role":"user","content":[
3834 {"type":"image_url","image_url":{"url":"https://example.com/img.png"}}
3835 ]}]}"#,
3836 );
3837 let f = extract_openai_features(&HeaderMap::new(), &img_body);
3838 let ladder = vec!["anthropic/claude-haiku-4-5".to_owned()];
3839 let providers = crate::provider::ProviderRegistry::new("http://x", "http://x");
3840 assert!(!enforce_can_handle(
3841 &f,
3842 &img_body,
3843 true,
3844 &ladder,
3845 &providers,
3846 Dialect::Openai,
3847 ));
3848 }
3849
3850 fn openai_enforce_state(mock_resp: ModelResponse) -> AppState {
3854 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"mock/m\"]\ngates = [\"non-empty\"]\n";
3855 let config = ProxyConfig::from_lookup(|k| match k {
3856 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
3857 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
3858 _ => None,
3859 })
3860 .unwrap();
3861 let mut outs = HashMap::new();
3862 outs.insert("mock/m".to_owned(), Ok(mock_resp));
3863 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
3864 map.insert("mock".to_owned(), Arc::new(MockProvider::new("mock", outs)));
3865 let (traces, _rx) = mpsc::channel(64);
3866 std::mem::forget(_rx);
3867 let tenant_rate_limiter = build_tenant_rate_limiter(&config);
3868 AppState {
3869 config: Arc::new(config),
3870 http: reqwest::Client::new(),
3871 providers: ProviderRegistry::from_map(map),
3872 gate_health: Arc::new(GateHealthRegistry::new()),
3873 traces,
3874 adaptive: None,
3875 bandit: None,
3876 tenant_rate_limiter,
3877 spill: None,
3878 }
3879 }
3880
3881 #[tokio::test]
3882 async fn chat_completions_plain_text_enforce_returns_openai_shape() {
3883 let mock = model_resp("mock/m", "gated answer");
3884 let state = openai_enforce_state(mock);
3885 let body = Bytes::from_static(
3886 br#"{"model":"gpt-4o","messages":[{"role":"user","content":"hello"}]}"#,
3887 );
3888 let resp = chat_completions(
3889 State(state),
3890 Extension(TenantId("default".to_owned())),
3891 HeaderMap::new(),
3892 body,
3893 )
3894 .await;
3895 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3896 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
3897 .await
3898 .unwrap();
3899 let json: Value = serde_json::from_slice(&bytes).expect("must be JSON");
3900 assert_eq!(json["object"], "chat.completion");
3901 assert_eq!(json["choices"][0]["message"]["role"], "assistant");
3902 assert_eq!(json["choices"][0]["message"]["content"], "gated answer");
3903 assert_eq!(json["choices"][0]["finish_reason"], "stop");
3904 assert!(
3905 json["id"]
3906 .as_str()
3907 .is_some_and(|id| id.starts_with("chatcmpl-"))
3908 );
3909 }
3910
3911 #[tokio::test]
3912 async fn chat_completions_stream_true_returns_sse_with_openai_chunks() {
3913 let mock = model_resp("mock/m", "gated answer");
3914 let state = openai_enforce_state(mock);
3915 let body = Bytes::from_static(
3916 br#"{"model":"gpt-4o","stream":true,"messages":[{"role":"user","content":"hello"}]}"#,
3917 );
3918 let resp = chat_completions(
3919 State(state),
3920 Extension(TenantId("default".to_owned())),
3921 HeaderMap::new(),
3922 body,
3923 )
3924 .await;
3925 assert_eq!(resp.status(), axum::http::StatusCode::OK);
3926 assert!(
3927 resp.headers()
3928 .get(axum::http::header::CONTENT_TYPE)
3929 .and_then(|v| v.to_str().ok())
3930 .is_some_and(|ct| ct.starts_with("text/event-stream")),
3931 "stream:true must return SSE"
3932 );
3933 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
3934 .await
3935 .unwrap();
3936 let text = String::from_utf8(bytes.to_vec()).unwrap();
3937 assert!(
3939 text.contains("chat.completion.chunk"),
3940 "must have OpenAI chunk frames"
3941 );
3942 assert!(text.contains("[DONE]"), "must end with [DONE]");
3943 assert!(
3944 text.contains("gated answer"),
3945 "content must be in the stream"
3946 );
3947 assert!(
3949 !text.contains("message_start"),
3950 "must not have Anthropic event types"
3951 );
3952 }
3953}