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, FEATURE_VERSION, Features, FinalOutcome, GENESIS_HASH, Mode,
19 PolicyRef, RequestInfo, Score, ServedFrom, TaskKind, Trace, Verdict,
20};
21use serde::Deserialize;
22use serde_json::Value;
23use tokio::sync::mpsc::error::TrySendError;
24use uuid::Uuid;
25
26use crate::config::ProxyConfig;
27use crate::error::ProxyError;
28use crate::gate::{GateHealthRegistry, resolve_gates};
29use crate::provider::{Auth, ChatMessage, ModelRequest, ModelResponse, ProviderRegistry};
30use crate::router::{EnforceCtx, EngineOutcome, route_enforce};
31use crate::store;
32use crate::tenant_auth::{TenantId, auth_middleware};
33use crate::upstream::{forward_anthropic, forward_anthropic_streaming};
34use firstpass_core::Route;
35
36#[derive(Clone)]
39pub struct AppState {
40 pub config: Arc<ProxyConfig>,
42 pub http: reqwest::Client,
44 pub providers: ProviderRegistry,
46 pub gate_health: Arc<GateHealthRegistry>,
48 pub traces: store::TraceSender,
50 pub adaptive: Option<Arc<std::sync::Mutex<firstpass_core::conformal::AdaptiveConformal>>>,
54 pub tenant_rate_limiter: Option<Arc<governor::DefaultKeyedRateLimiter<String>>>,
58}
59
60#[must_use]
64pub fn build_tenant_rate_limiter(
65 config: &ProxyConfig,
66) -> Option<Arc<governor::DefaultKeyedRateLimiter<String>>> {
67 let per_sec = config.tenant_rate_per_sec?;
68 Some(Arc::new(governor::RateLimiter::keyed(
69 governor::Quota::per_second(per_sec),
70 )))
71}
72
73pub async fn tenant_rate_limit_middleware(
77 State(state): State<AppState>,
78 Extension(tenant): Extension<TenantId>,
79 req: Request,
80 next: Next,
81) -> Response {
82 if let Some(limiter) = &state.tenant_rate_limiter
83 && limiter.check_key(&tenant.0).is_err()
84 {
85 return ProxyError::RateLimited.into_response();
86 }
87 next.run(req).await
88}
89
90impl std::fmt::Debug for AppState {
91 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92 f.debug_struct("AppState")
93 .field("config", &self.config)
94 .finish_non_exhaustive()
95 }
96}
97
98fn offer_trace(traces: &store::TraceSender, trace: Trace) {
103 record_trace_metrics(&trace);
104 match traces.try_send(trace) {
105 Ok(()) => {}
106 Err(TrySendError::Full(_)) => {
107 tracing::warn!("trace channel full; dropping trace (writer behind under load)");
108 metrics::counter!("firstpass_traces_dropped_total").increment(1);
109 }
110 Err(TrySendError::Closed(_)) => {
111 tracing::warn!("trace writer is gone; dropping trace");
112 }
113 }
114}
115
116fn record_trace_metrics(trace: &Trace) {
120 if trace.mode == Mode::Enforce {
121 metrics::histogram!("firstpass_enforce_latency_ms")
122 .record(trace.final_.total_latency_ms as f64);
123 if trace.final_.escalations > 0 {
124 metrics::counter!("firstpass_escalations_total")
125 .increment(u64::from(trace.final_.escalations));
126 }
127 }
128 let served_from = match trace.final_.served_from {
129 ServedFrom::Attempt => "attempt",
130 ServedFrom::BestAttempt => "best_attempt",
131 ServedFrom::Error => "error",
132 };
133 metrics::counter!("firstpass_served_total", "served_from" => served_from).increment(1);
134 if trace.final_.served_from == ServedFrom::Error {
135 metrics::counter!("firstpass_upstream_failures_total").increment(1);
136 }
137}
138
139const MAX_BODY_BYTES: usize = 16 * 1024 * 1024;
143
144pub fn app(state: AppState) -> Result<Router, ProxyError> {
151 crate::metrics::install()?;
152 let max_concurrency = state.config.max_concurrency;
153
154 let business = Router::new()
163 .route("/v1/messages", post(messages))
164 .route("/v1/feedback", post(feedback))
165 .route("/v1/capabilities", get(capabilities))
166 .layer(axum::middleware::from_fn_with_state(
167 state.clone(),
168 tenant_rate_limit_middleware,
169 ))
170 .layer(axum::middleware::from_fn_with_state(
171 state.clone(),
172 auth_middleware,
173 ));
174
175 Ok(Router::new()
176 .merge(business)
177 .route("/healthz", get(healthz))
178 .route("/metrics", get(crate::metrics::handler))
179 .layer(axum::extract::DefaultBodyLimit::max(MAX_BODY_BYTES))
181 .layer(tower::limit::GlobalConcurrencyLimitLayer::new(
184 max_concurrency,
185 ))
186 .with_state(state))
187}
188
189async fn healthz() -> impl IntoResponse {
191 Json(serde_json::json!({ "status": "ok" }))
192}
193
194async fn capabilities(State(state): State<AppState>) -> impl IntoResponse {
197 let (ladder, gates) = state
200 .config
201 .routing
202 .as_ref()
203 .and_then(|c| c.routes.iter().find(|r| r.mode == Mode::Enforce))
204 .map(|r| (r.ladder.clone(), r.gates.clone()))
205 .unwrap_or_default();
206 Json(serde_json::json!({
207 "service": "firstpass",
208 "version": env!("CARGO_PKG_VERSION"),
209 "feature_version": FEATURE_VERSION,
210 "modes": ["observe", "enforce"],
211 "wire_apis": ["anthropic.messages"],
212 "ladder": ladder,
213 "gates": gates,
214 "feedback_api": "POST /v1/feedback",
215 "offboarding": "unset ANTHROPIC_BASE_URL",
216 }))
217}
218
219#[derive(Debug, Deserialize)]
221struct FeedbackRequest {
222 trace_id: String,
224 gate_id: String,
226 verdict: String,
228 #[serde(default)]
230 score: Option<f64>,
231 reporter: String,
233}
234
235async fn feedback(
240 State(state): State<AppState>,
241 Extension(TenantId(tenant)): Extension<TenantId>,
242 body: Bytes,
243) -> Response {
244 let req: FeedbackRequest = match serde_json::from_slice(&body) {
245 Ok(r) => r,
246 Err(e) => {
247 return ProxyError::BadRequest(format!("invalid feedback body: {e}")).into_response();
248 }
249 };
250 let verdict = match req.verdict.as_str() {
251 "pass" => Verdict::Pass,
252 "fail" => Verdict::Fail,
253 "abstain" => Verdict::Abstain,
254 other => {
255 return ProxyError::BadRequest(format!("unknown verdict {other:?}")).into_response();
256 }
257 };
258 let score = match req.score {
259 Some(s) => match Score::new(s) {
260 Ok(sc) => Some(sc),
261 Err(_) => {
262 return ProxyError::BadRequest(format!("score {s} out of range [0,1]"))
263 .into_response();
264 }
265 },
266 None => None,
267 };
268
269 let db = state.config.db_path.clone();
270
271 let (db_check, tenant_check, tid_check) = (db.clone(), tenant.clone(), req.trace_id.clone());
276 match tokio::task::spawn_blocking(move || {
277 store::trace_exists(&db_check, &tenant_check, &tid_check)
278 })
279 .await
280 {
281 Ok(Ok(true)) => {}
282 Ok(Ok(false)) => {
283 return ProxyError::NotFound(format!("unknown trace_id {:?}", req.trace_id))
284 .into_response();
285 }
286 Ok(Err(e)) => {
287 tracing::error!(%e, "feedback: trace_exists check failed");
288 return ProxyError::Internal(e.to_string()).into_response();
289 }
290 Err(e) => {
291 tracing::error!(%e, "feedback: trace_exists task panicked");
292 return ProxyError::Internal(e.to_string()).into_response();
293 }
294 }
295
296 let feedback_signal = match verdict {
298 Verdict::Pass => Some(true),
299 Verdict::Fail => Some(false),
300 Verdict::Abstain => None,
301 };
302 let dv = DeferredVerdict {
303 gate_id: req.gate_id,
304 verdict,
305 score,
306 reported_at: jiff::Timestamp::now(),
307 reporter: req.reporter,
308 };
309 let trace_id = req.trace_id.clone();
310 match tokio::task::spawn_blocking(move || store::append_deferred(&db, &req.trace_id, &dv)).await
311 {
312 Ok(Ok(())) => {
313 if let (Some(a), Some(correct)) = (state.adaptive.as_ref(), feedback_signal)
315 && let Ok(mut g) = a.lock()
316 {
317 g.observe_served(correct);
318 }
319 (
320 axum::http::StatusCode::ACCEPTED,
321 Json(serde_json::json!({ "status": "recorded", "trace_id": trace_id })),
322 )
323 .into_response()
324 }
325 Ok(Err(e)) => {
326 tracing::error!(%e, "feedback: append_deferred failed");
327 ProxyError::Internal(e.to_string()).into_response()
328 }
329 Err(e) => {
330 tracing::error!(%e, "feedback: append_deferred task panicked");
331 ProxyError::Internal(e.to_string()).into_response()
332 }
333 }
334}
335
336const SESSION_HEADER: &str = "x-firstpass-session";
339
340const AGENT_HEADER: &str = "x-firstpass-agent";
342const SUBAGENT_HEADER: &str = "x-firstpass-subagent";
344
345async fn messages(
350 State(state): State<AppState>,
351 Extension(TenantId(tenant)): Extension<TenantId>,
352 headers: HeaderMap,
353 body: Bytes,
354) -> Response {
355 let session_header = header_str(&headers, SESSION_HEADER);
356
357 if let Some(routing) = state.config.routing.as_ref() {
360 let features = extract_features(&headers, &body);
361 if let Some(route) = routing
362 .route_for(&features)
363 .filter(|r| r.mode == Mode::Enforce && !r.ladder.is_empty())
364 {
365 let route = route.clone();
368 if enforce_can_handle(&features, &body, routing.escalation.enforce_structured) {
369 return handle_enforce(
370 &state,
371 &headers,
372 &body,
373 features,
374 &route,
375 session_header,
376 tenant,
377 )
378 .await;
379 }
380 tracing::info!(
385 "enforce route matched but request carries tools/images; serving via observe passthrough"
386 );
387 }
388 }
389 observe_passthrough(state, headers, body, session_header, tenant).await
390}
391
392fn enforce_can_handle(features: &Features, body: &[u8], enforce_structured: bool) -> bool {
402 if enforce_structured {
403 return true;
404 }
405 features.tool_count == 0 && !features.has_images && !messages_have_tool_blocks(body)
406}
407
408fn messages_have_tool_blocks(body: &[u8]) -> bool {
411 serde_json::from_slice::<Value>(body)
412 .ok()
413 .and_then(|json| {
414 json.get("messages")
415 .and_then(Value::as_array)
416 .map(|messages| messages.iter().any(message_has_tool_block))
417 })
418 .unwrap_or(false)
419}
420
421fn message_has_tool_block(message: &Value) -> bool {
423 message
424 .get("content")
425 .and_then(Value::as_array)
426 .is_some_and(|blocks| {
427 blocks.iter().any(|block| {
428 matches!(
429 block.get("type").and_then(Value::as_str),
430 Some("tool_use" | "tool_result")
431 )
432 })
433 })
434}
435
436fn is_stream_request(body: &[u8]) -> bool {
438 serde_json::from_slice::<Value>(body)
439 .ok()
440 .and_then(|json| json.get("stream").and_then(Value::as_bool))
441 .unwrap_or(false)
442}
443
444fn header_str(headers: &HeaderMap, name: &str) -> Option<String> {
446 headers
447 .get(name)
448 .and_then(|v| v.to_str().ok())
449 .map(str::to_owned)
450}
451
452fn extract_features(headers: &HeaderMap, body: &[u8]) -> Features {
455 let (_model, tool_count, has_images) = request_features(body);
456 let mut f = Features::new(TaskKind::Other);
457 f.agent = header_str(headers, AGENT_HEADER);
458 f.subagent = header_str(headers, SUBAGENT_HEADER);
459 f.tool_count = tool_count;
460 f.has_images = has_images;
461 f.prompt_token_bucket = token_bucket(body.len() as u64);
464 f.hour_bucket = hour_bucket(jiff::Timestamp::now());
465 f
466}
467
468async fn handle_enforce(
471 state: &AppState,
472 headers: &HeaderMap,
473 body: &Bytes,
474 features: Features,
475 route: &Route,
476 session_header: Option<String>,
477 tenant: String,
478) -> Response {
479 let Some(base_request) = parse_model_request(body) else {
480 return ProxyError::BadRequest(
481 "request body is not a valid Anthropic Messages request".to_owned(),
482 )
483 .into_response();
484 };
485 let auth = Auth::from_headers(headers);
486 let gate_defs = state
487 .config
488 .routing
489 .as_ref()
490 .map_or(&[][..], |cfg| &cfg.gate_defs);
491 let gates = resolve_gates(&route.gates, gate_defs, &state.providers, &auth);
492 let session_id = session_header.unwrap_or_else(|| Uuid::now_v7().to_string());
493 let (budget, max_rungs, speculation, serve_threshold) = match state.config.routing.as_ref() {
494 Some(cfg) => (
495 cfg.budget.per_request_usd,
496 cfg.escalation.max_rungs_per_request,
497 cfg.escalation.speculation,
498 cfg.escalation.serve_threshold,
499 ),
500 None => (None, 3, 0, None),
501 };
502 let serve_threshold = state
505 .adaptive
506 .as_ref()
507 .and_then(|a| a.lock().ok().map(|g| g.threshold()))
508 .or(serve_threshold);
509
510 let ctx = EnforceCtx {
511 ladder: &route.ladder,
512 gates: &gates,
513 health: &state.gate_health,
514 base_request: &base_request,
515 providers: &state.providers,
516 auth: &auth,
517 prices: &state.config.prices,
518 budget_per_request_usd: budget,
519 max_rungs,
520 speculation,
521 serve_threshold,
522 features,
523 tenant_id: tenant,
526 session_id,
527 prompt_hash: prompt_hash(&state.config.prompt_salt, body),
528 api: "anthropic.messages".to_owned(),
529 policy_id: "static-ladder@v0".to_owned(),
530 };
531
532 let (outcome, trace) = route_enforce(ctx).await;
533 offer_trace(&state.traces, trace);
535
536 match outcome {
537 EngineOutcome::Served(resp) => {
538 let message = anthropic_response_json(&resp);
539 if is_stream_request(body) {
543 (
544 axum::http::StatusCode::OK,
545 [(
546 axum::http::header::CONTENT_TYPE,
547 "text/event-stream; charset=utf-8",
548 )],
549 anthropic_sse_from_message(&message),
550 )
551 .into_response()
552 } else {
553 (axum::http::StatusCode::OK, Json(message)).into_response()
554 }
555 }
556 EngineOutcome::Failed(msg) => ProxyError::Engine(msg).into_response(),
557 }
558}
559
560fn parse_model_request(body: &[u8]) -> Option<ModelRequest> {
569 let json: Value = serde_json::from_slice(body).ok()?;
570 let messages_json = json.get("messages")?.as_array()?;
571 let messages = messages_json
572 .iter()
573 .map(|m| ChatMessage {
574 role: m
575 .get("role")
576 .and_then(Value::as_str)
577 .unwrap_or("user")
578 .to_owned(),
579 content: m
580 .get("content")
581 .cloned()
582 .unwrap_or_else(|| Value::String(String::new())),
583 })
584 .collect();
585 let system = json
586 .get("system")
587 .and_then(Value::as_str)
588 .map(str::to_owned);
589 let max_tokens = json
590 .get("max_tokens")
591 .and_then(Value::as_u64)
592 .and_then(|n| u32::try_from(n).ok())
593 .unwrap_or(1024);
594 let tools = json.get("tools").cloned().unwrap_or(Value::Null);
595 Some(ModelRequest {
596 model: json
597 .get("model")
598 .and_then(Value::as_str)
599 .unwrap_or_default()
600 .to_owned(),
601 system,
602 messages,
603 max_tokens,
604 tools,
605 })
606}
607
608fn anthropic_response_json(resp: &ModelResponse) -> Value {
618 let content = resp
619 .raw
620 .get("content")
621 .filter(|c| c.is_array())
622 .cloned()
623 .unwrap_or_else(|| serde_json::json!([{ "type": "text", "text": resp.text }]));
624 serde_json::json!({
625 "id": format!("msg_{}", Uuid::now_v7()),
626 "type": "message",
627 "role": "assistant",
628 "model": resp.model,
629 "content": content,
630 "usage": { "input_tokens": resp.in_tokens, "output_tokens": resp.out_tokens },
631 })
632}
633
634fn sse_event(out: &mut String, event: &str, data: &Value) {
636 out.push_str("event: ");
637 out.push_str(event);
638 out.push_str("\ndata: ");
639 out.push_str(&data.to_string());
640 out.push_str("\n\n");
641}
642
643fn anthropic_sse_from_message(message: &Value) -> String {
650 let mut out = String::new();
651
652 let mut start_msg = message.clone();
654 start_msg["content"] = Value::Array(Vec::new());
655 sse_event(
656 &mut out,
657 "message_start",
658 &serde_json::json!({ "type": "message_start", "message": start_msg }),
659 );
660
661 let empty = Vec::new();
662 let blocks = message
663 .get("content")
664 .and_then(Value::as_array)
665 .unwrap_or(&empty);
666 for (i, block) in blocks.iter().enumerate() {
667 match block.get("type").and_then(Value::as_str) {
668 Some("tool_use") => {
669 let mut shell = block.clone();
671 shell["input"] = serde_json::json!({});
672 sse_event(
673 &mut out,
674 "content_block_start",
675 &serde_json::json!({ "type": "content_block_start", "index": i, "content_block": shell }),
676 );
677 let input_json = block
678 .get("input")
679 .map_or_else(|| "{}".to_owned(), std::string::ToString::to_string);
680 sse_event(
681 &mut out,
682 "content_block_delta",
683 &serde_json::json!({ "type": "content_block_delta", "index": i,
684 "delta": { "type": "input_json_delta", "partial_json": input_json } }),
685 );
686 }
687 _ => {
688 let text = block.get("text").and_then(Value::as_str).unwrap_or("");
690 sse_event(
691 &mut out,
692 "content_block_start",
693 &serde_json::json!({ "type": "content_block_start", "index": i,
694 "content_block": { "type": "text", "text": "" } }),
695 );
696 sse_event(
697 &mut out,
698 "content_block_delta",
699 &serde_json::json!({ "type": "content_block_delta", "index": i,
700 "delta": { "type": "text_delta", "text": text } }),
701 );
702 }
703 }
704 sse_event(
705 &mut out,
706 "content_block_stop",
707 &serde_json::json!({ "type": "content_block_stop", "index": i }),
708 );
709 }
710
711 let out_tokens = message
712 .pointer("/usage/output_tokens")
713 .cloned()
714 .unwrap_or_else(|| Value::from(0));
715 sse_event(
716 &mut out,
717 "message_delta",
718 &serde_json::json!({ "type": "message_delta", "delta": { "stop_reason": "end_turn" },
719 "usage": { "output_tokens": out_tokens } }),
720 );
721 sse_event(
722 &mut out,
723 "message_stop",
724 &serde_json::json!({ "type": "message_stop" }),
725 );
726 out
727}
728
729async fn observe_passthrough(
731 state: AppState,
732 headers: HeaderMap,
733 body: Bytes,
734 session_header: Option<String>,
735 tenant: String,
736) -> Response {
737 if is_stream_request(&body) {
739 return observe_stream(state, headers, body, session_header, tenant).await;
740 }
741 let start = Instant::now();
742 let result = forward_anthropic(
743 &state.http,
744 &state.config.upstream_anthropic,
745 &headers,
746 body.clone(),
747 )
748 .await;
749 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
750
751 match result {
752 Ok((status, resp_headers, resp_body)) => {
753 spawn_trace(
757 &state,
758 body,
759 Some(resp_body.clone()),
760 latency_ms,
761 session_header,
762 tenant,
763 );
764 (status, resp_headers, resp_body).into_response()
765 }
766 Err(err) => {
767 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
768 err.into_response()
769 }
770 }
771}
772
773async fn observe_stream(
782 state: AppState,
783 headers: HeaderMap,
784 body: Bytes,
785 session_header: Option<String>,
786 tenant: String,
787) -> Response {
788 let start = Instant::now();
789 let result = forward_anthropic_streaming(
790 &state.http,
791 &state.config.upstream_anthropic,
792 &headers,
793 body.clone(),
794 )
795 .await;
796 let latency_ms = u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX);
797
798 match result {
799 Ok((status, resp_headers, response)) => {
800 spawn_stream_trace(&state, body, latency_ms, session_header, tenant);
801 let stream_body = Body::from_stream(response.bytes_stream());
802 (status, resp_headers, stream_body).into_response()
803 }
804 Err(err) => {
805 spawn_trace(&state, body, None, latency_ms, session_header, tenant);
806 err.into_response()
807 }
808 }
809}
810
811fn spawn_stream_trace(
813 state: &AppState,
814 req_body: Bytes,
815 latency_ms: u64,
816 session_header: Option<String>,
817 tenant: String,
818) {
819 let config = state.config.clone();
820 let traces = state.traces.clone();
821 tokio::spawn(async move {
822 let mut trace =
823 build_stream_trace(&config, &req_body, latency_ms, session_header.as_deref());
824 trace.tenant_id = tenant;
826 offer_trace(&traces, trace);
827 });
828}
829
830fn spawn_trace(
835 state: &AppState,
836 req_body: Bytes,
837 resp_body: Option<Bytes>,
838 latency_ms: u64,
839 session_header: Option<String>,
840 tenant: String,
841) {
842 let config = state.config.clone();
843 let traces = state.traces.clone();
844 tokio::spawn(async move {
845 let mut trace = match resp_body {
846 Some(resp) => build_trace(
847 &config,
848 &req_body,
849 &resp,
850 latency_ms,
851 session_header.as_deref(),
852 ),
853 None => build_error_trace(&config, &req_body, latency_ms, session_header.as_deref()),
854 };
855 trace.tenant_id = tenant;
857 offer_trace(&traces, trace);
858 });
859}
860
861fn session_id(session_header: Option<&str>, trace_id: Uuid) -> String {
863 session_header
864 .map(str::to_owned)
865 .unwrap_or_else(|| trace_id.to_string())
866}
867
868fn prompt_hash(salt: &str, body: &[u8]) -> String {
871 let mut salted = Vec::with_capacity(salt.len() + body.len());
872 salted.extend_from_slice(salt.as_bytes());
873 salted.extend_from_slice(body);
874 sha256_hex(&salted)
875}
876
877fn request_features(body: &[u8]) -> (Option<String>, u32, bool) {
881 let Ok(json) = serde_json::from_slice::<Value>(body) else {
882 return (None, 0, false);
883 };
884 let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
885 let tool_count = json
886 .get("tools")
887 .and_then(Value::as_array)
888 .map_or(0, |tools| u32::try_from(tools.len()).unwrap_or(u32::MAX));
889 let has_images = json
890 .get("messages")
891 .and_then(Value::as_array)
892 .is_some_and(|messages| messages.iter().any(message_has_image));
893 (model, tool_count, has_images)
894}
895
896fn message_has_image(message: &Value) -> bool {
898 message
899 .get("content")
900 .and_then(Value::as_array)
901 .is_some_and(|blocks| {
902 blocks
903 .iter()
904 .any(|block| block.get("type").and_then(Value::as_str) == Some("image"))
905 })
906}
907
908fn response_usage(body: &[u8]) -> (Option<String>, u64, u64) {
911 let Ok(json) = serde_json::from_slice::<Value>(body) else {
912 return (None, 0, 0);
913 };
914 let model = json.get("model").and_then(Value::as_str).map(str::to_owned);
915 let in_tokens = json
916 .pointer("/usage/input_tokens")
917 .and_then(Value::as_u64)
918 .unwrap_or(0);
919 let out_tokens = json
920 .pointer("/usage/output_tokens")
921 .and_then(Value::as_u64)
922 .unwrap_or(0);
923 (model, in_tokens, out_tokens)
924}
925
926fn build_trace(
928 config: &ProxyConfig,
929 req_body: &Bytes,
930 resp_body: &Bytes,
931 latency_ms: u64,
932 session_header: Option<&str>,
933) -> Trace {
934 let (req_model, tool_count, has_images) = request_features(req_body);
935 let (resp_model, in_tokens, out_tokens) = response_usage(resp_body);
936 let model = resp_model
937 .or(req_model)
938 .unwrap_or_else(|| "unknown".to_owned());
939
940 let cost_usd = config
941 .prices
942 .cost_usd(&format!("anthropic/{model}"), in_tokens, out_tokens)
943 .unwrap_or(0.0);
944
945 let attempt = Attempt {
946 rung: 0,
947 model,
948 provider: "anthropic".to_owned(),
949 in_tokens,
950 out_tokens,
951 cost_usd,
952 latency_ms,
953 gates: Vec::new(),
954 verdict: Verdict::Pass,
955 };
956
957 let mut trace = base_trace(config, req_body, latency_ms, session_header);
958 trace.request.features.prompt_token_bucket = token_bucket(in_tokens);
959 trace.request.features.tool_count = tool_count;
960 trace.request.features.has_images = has_images;
961 trace.attempts.push(attempt);
962 trace.final_ = FinalOutcome {
963 served_rung: Some(0),
964 served_from: ServedFrom::Attempt,
965 total_cost_usd: cost_usd,
966 gate_cost_usd: 0.0,
967 total_latency_ms: latency_ms,
968 escalations: 0,
969 counterfactual_baseline_usd: cost_usd,
970 savings_usd: 0.0,
971 };
972 trace.recompute_savings();
973 trace
974}
975
976fn build_stream_trace(
980 config: &ProxyConfig,
981 req_body: &Bytes,
982 latency_ms: u64,
983 session_header: Option<&str>,
984) -> Trace {
985 let (req_model, tool_count, has_images) = request_features(req_body);
986 let model = req_model.unwrap_or_else(|| "unknown".to_owned());
987
988 let attempt = Attempt {
989 rung: 0,
990 model,
991 provider: "anthropic".to_owned(),
992 in_tokens: 0,
993 out_tokens: 0,
994 cost_usd: 0.0,
995 latency_ms,
996 gates: Vec::new(),
997 verdict: Verdict::Pass,
998 };
999
1000 let mut trace = base_trace(config, req_body, latency_ms, session_header);
1001 trace.request.features.tool_count = tool_count;
1002 trace.request.features.has_images = has_images;
1003 trace.attempts.push(attempt);
1004 trace.final_ = FinalOutcome {
1005 served_rung: Some(0),
1006 served_from: ServedFrom::Attempt,
1007 total_cost_usd: 0.0,
1008 gate_cost_usd: 0.0,
1009 total_latency_ms: latency_ms,
1010 escalations: 0,
1011 counterfactual_baseline_usd: 0.0,
1012 savings_usd: 0.0,
1013 };
1014 trace.recompute_savings();
1015 trace
1016}
1017
1018fn build_error_trace(
1022 config: &ProxyConfig,
1023 req_body: &Bytes,
1024 latency_ms: u64,
1025 session_header: Option<&str>,
1026) -> Trace {
1027 let (_, tool_count, has_images) = request_features(req_body);
1028 let mut trace = base_trace(config, req_body, latency_ms, session_header);
1029 trace.request.features.tool_count = tool_count;
1030 trace.request.features.has_images = has_images;
1031 trace.final_ = FinalOutcome {
1032 served_rung: None,
1033 served_from: ServedFrom::Error,
1034 total_cost_usd: 0.0,
1035 gate_cost_usd: 0.0,
1036 total_latency_ms: latency_ms,
1037 escalations: 0,
1038 counterfactual_baseline_usd: 0.0,
1039 savings_usd: 0.0,
1040 };
1041 trace.recompute_savings();
1042 trace
1043}
1044
1045fn base_trace(
1048 config: &ProxyConfig,
1049 req_body: &Bytes,
1050 latency_ms: u64,
1051 session_header: Option<&str>,
1052) -> Trace {
1053 let trace_id = Uuid::now_v7();
1054 let mut features = Features::new(TaskKind::Other);
1055 features.hour_bucket = hour_bucket(jiff::Timestamp::now());
1056
1057 Trace {
1058 trace_id,
1059 prev_hash: GENESIS_HASH.to_owned(),
1060 tenant_id: config.tenant_id.clone(),
1061 session_id: session_id(session_header, trace_id),
1062 ts: jiff::Timestamp::now(),
1063 mode: Mode::Observe,
1064 policy: PolicyRef {
1065 id: "observe-passthrough@v0".to_owned(),
1066 explore: false,
1067 },
1068 request: RequestInfo {
1069 api: "anthropic.messages".to_owned(),
1070 prompt_hash: prompt_hash(&config.prompt_salt, req_body),
1071 features,
1072 },
1073 attempts: Vec::new(),
1074 deferred: Vec::new(),
1075 final_: FinalOutcome {
1076 served_rung: None,
1077 served_from: ServedFrom::Error,
1078 total_cost_usd: 0.0,
1079 gate_cost_usd: 0.0,
1080 total_latency_ms: latency_ms,
1081 escalations: 0,
1082 counterfactual_baseline_usd: 0.0,
1083 savings_usd: 0.0,
1084 },
1085 }
1086}
1087
1088#[cfg(test)]
1089mod tests {
1090 use bytes::Bytes;
1091
1092 use super::*;
1093
1094 fn test_config() -> ProxyConfig {
1095 ProxyConfig::from_lookup(|_| None).unwrap()
1096 }
1097
1098 #[test]
1099 fn build_trace_maps_request_and_response_fields() {
1100 let config = test_config();
1101 let req = Bytes::from_static(
1102 br#"{"model":"claude-haiku-4-5","tools":[{"name":"a"}],"messages":[]}"#,
1103 );
1104 let resp = Bytes::from_static(
1105 br#"{"model":"claude-haiku-4-5","usage":{"input_tokens":1200,"output_tokens":300}}"#,
1106 );
1107
1108 let trace = build_trace(&config, &req, &resp, 42, Some("sess-1"));
1109
1110 assert_eq!(trace.request.api, "anthropic.messages");
1111 assert_eq!(trace.session_id, "sess-1");
1112 assert_eq!(trace.attempts.len(), 1);
1113 let attempt = &trace.attempts[0];
1114 assert_eq!(attempt.model, "claude-haiku-4-5");
1115 assert_eq!(attempt.provider, "anthropic");
1116 assert_eq!(attempt.in_tokens, 1200);
1117 assert_eq!(attempt.out_tokens, 300);
1118 assert!(attempt.cost_usd > 0.0);
1119 assert_eq!(trace.request.features.tool_count, 1);
1120 assert!(!trace.request.features.has_images);
1121 assert_eq!(trace.final_.served_rung, Some(0));
1122 }
1123
1124 #[test]
1125 fn build_trace_falls_back_to_trace_id_session_when_header_absent() {
1126 let config = test_config();
1127 let req = Bytes::from_static(b"{}");
1128 let resp = Bytes::from_static(b"{}");
1129
1130 let trace = build_trace(&config, &req, &resp, 1, None);
1131
1132 assert_eq!(trace.session_id, trace.trace_id.to_string());
1133 }
1134
1135 #[test]
1136 fn build_error_trace_has_no_attempts_and_served_from_error() {
1137 let config = test_config();
1138 let req = Bytes::from_static(br#"{"model":"claude-haiku-4-5"}"#);
1139
1140 let trace = build_error_trace(&config, &req, 7, None);
1141
1142 assert!(trace.attempts.is_empty());
1143 assert_eq!(trace.final_.served_from, ServedFrom::Error);
1144 assert_eq!(trace.final_.served_rung, None);
1145 }
1146
1147 #[test]
1148 fn message_with_image_block_sets_has_images() {
1149 let req = br#"{"messages":[{"role":"user","content":[{"type":"image"}]}]}"#;
1150 let (_, _, has_images) = request_features(req);
1151 assert!(has_images);
1152 }
1153
1154 #[test]
1155 fn prompt_hash_never_contains_raw_prompt_text() {
1156 let hash = prompt_hash("salt", b"super secret prompt");
1157 assert!(!hash.contains("secret"));
1158 assert_eq!(hash.len(), 64);
1159 }
1160
1161 #[test]
1162 fn parse_model_request_preserves_content_verbatim_and_projects_text() {
1163 let body = br#"{"model":"m","system":"sys","max_tokens":50,
1164 "messages":[{"role":"user","content":[{"type":"text","text":"a"},{"type":"text","text":"b"}]},
1165 {"role":"assistant","content":"c"}]}"#;
1166 let req = parse_model_request(body).unwrap();
1167 assert_eq!(req.system.as_deref(), Some("sys"));
1168 assert_eq!(req.max_tokens, 50);
1169 assert_eq!(req.messages.len(), 2);
1170 assert_eq!(
1172 req.messages[0].content,
1173 serde_json::json!([{"type":"text","text":"a"},{"type":"text","text":"b"}])
1174 );
1175 assert_eq!(req.messages[1].content, Value::String("c".to_owned()));
1177 assert_eq!(req.messages[0].text_view(), "a\nb");
1179 assert_eq!(req.messages[1].text_view(), "c");
1180 }
1181
1182 #[test]
1183 fn tool_and_image_blocks_survive_the_request_round_trip() {
1184 let body = br#"{"model":"m","max_tokens":50,"messages":[
1186 {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
1187 {"role":"user","content":[
1188 {"type":"tool_result","tool_use_id":"t1","content":"2"},
1189 {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
1190 ]}]}"#;
1191 let req = parse_model_request(body).unwrap();
1192 let round_tripped = serde_json::to_value(&req.messages).unwrap();
1193 assert_eq!(
1194 round_tripped,
1195 serde_json::json!([
1196 {"role":"assistant","content":[{"type":"tool_use","id":"t1","name":"calc","input":{"x":1}}]},
1197 {"role":"user","content":[
1198 {"type":"tool_result","tool_use_id":"t1","content":"2"},
1199 {"type":"image","source":{"type":"base64","media_type":"image/png","data":"AA=="}}
1200 ]}
1201 ])
1202 );
1203 }
1204
1205 #[test]
1206 fn text_message_serializes_byte_identical_to_a_plain_string() {
1207 let m = ChatMessage::text("user", "hello");
1209 assert_eq!(
1210 serde_json::to_string(&m).unwrap(),
1211 r#"{"role":"user","content":"hello"}"#
1212 );
1213 }
1214
1215 #[test]
1216 fn parse_model_request_rejects_non_message_bodies() {
1217 assert!(parse_model_request(b"not json").is_none());
1218 assert!(parse_model_request(br#"{"no":"messages"}"#).is_none());
1219 }
1220
1221 use crate::provider::{MockProvider, ModelResponse, Provider, ProviderError, ProviderRegistry};
1224 use axum::extract::State;
1225 use std::collections::HashMap;
1226 use std::sync::Arc;
1227 use tokio::sync::mpsc;
1228
1229 fn model_resp(model: &str, text: &str) -> ModelResponse {
1230 ModelResponse {
1231 model: model.to_owned(),
1232 text: text.to_owned(),
1233 in_tokens: 1000,
1234 out_tokens: 400,
1235 raw: serde_json::Value::Null,
1236 }
1237 }
1238
1239 fn enforce_state(
1242 ladder: &[&str],
1243 gates: &[&str],
1244 outcomes: Vec<(&str, Result<ModelResponse, ProviderError>)>,
1245 ) -> (AppState, mpsc::Receiver<Trace>) {
1246 let toml = format!(
1247 "[[route]]\nmatch = {{}}\nmode = \"enforce\"\nladder = [{}]\ngates = [{}]\n",
1248 ladder
1249 .iter()
1250 .map(|m| format!("\"{m}\""))
1251 .collect::<Vec<_>>()
1252 .join(", "),
1253 gates
1254 .iter()
1255 .map(|g| format!("\"{g}\""))
1256 .collect::<Vec<_>>()
1257 .join(", "),
1258 );
1259 let config = ProxyConfig::from_lookup(|k| match k {
1260 "FIRSTPASS_CONFIG_TOML" => Some(toml.clone()),
1261 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
1262 _ => None,
1263 })
1264 .unwrap();
1265
1266 let mut outs = HashMap::new();
1267 for (model, out) in outcomes {
1268 outs.insert(model.to_owned(), out);
1269 }
1270 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
1271 map.insert(
1272 "anthropic".to_owned(),
1273 Arc::new(MockProvider::new("anthropic", outs)),
1274 );
1275 let providers = ProviderRegistry::from_map(map);
1276
1277 let (traces, rx) = mpsc::channel(64);
1278 let state = AppState {
1279 config: Arc::new(config),
1280 http: reqwest::Client::new(),
1281 providers,
1282 gate_health: Arc::new(GateHealthRegistry::new()),
1283 traces,
1284 adaptive: None,
1285 tenant_rate_limiter: None,
1286 };
1287 (state, rx)
1288 }
1289
1290 fn user_body() -> Bytes {
1291 Bytes::from_static(
1292 br#"{"model":"ignored","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}"#,
1293 )
1294 }
1295
1296 async fn body_json(resp: Response) -> Value {
1297 let bytes = axum::body::to_bytes(resp.into_body(), 1 << 20)
1298 .await
1299 .unwrap();
1300 serde_json::from_slice(&bytes).unwrap()
1301 }
1302
1303 #[tokio::test]
1304 async fn enforce_serves_first_pass_and_returns_anthropic_shape() {
1305 let (state, mut rx) = enforce_state(
1306 &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
1307 &["non-empty"],
1308 vec![(
1309 "anthropic/claude-haiku-4-5",
1310 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
1311 )],
1312 );
1313 let resp = messages(
1314 State(state),
1315 Extension(TenantId("default".to_owned())),
1316 HeaderMap::new(),
1317 user_body(),
1318 )
1319 .await;
1320 assert_eq!(resp.status(), axum::http::StatusCode::OK);
1321 let json = body_json(resp).await;
1322 assert_eq!(json["type"], "message");
1323 assert_eq!(json["content"][0]["text"], "hello");
1324 assert_eq!(json["model"], "anthropic/claude-haiku-4-5");
1325
1326 let trace = rx.try_recv().expect("a trace was enqueued");
1327 assert_eq!(trace.mode, Mode::Enforce);
1328 assert_eq!(trace.final_.served_rung, Some(0));
1329 assert_eq!(trace.attempts.len(), 1);
1330 }
1331
1332 #[tokio::test]
1333 async fn enforce_escalates_then_serves_and_traces_two_attempts() {
1334 let (state, mut rx) = enforce_state(
1335 &["anthropic/claude-haiku-4-5", "anthropic/claude-sonnet-5"],
1336 &["non-empty"],
1337 vec![
1338 (
1339 "anthropic/claude-haiku-4-5",
1340 Ok(model_resp("anthropic/claude-haiku-4-5", " ")),
1341 ), (
1343 "anthropic/claude-sonnet-5",
1344 Ok(model_resp("anthropic/claude-sonnet-5", "answer")),
1345 ),
1346 ],
1347 );
1348 let resp = messages(
1349 State(state),
1350 Extension(TenantId("default".to_owned())),
1351 HeaderMap::new(),
1352 user_body(),
1353 )
1354 .await;
1355 let json = body_json(resp).await;
1356 assert_eq!(json["content"][0]["text"], "answer");
1357
1358 let trace = rx.try_recv().expect("trace enqueued");
1359 assert_eq!(trace.attempts.len(), 2);
1360 assert_eq!(trace.final_.escalations, 1);
1361 assert_eq!(trace.final_.served_rung, Some(1));
1362 }
1363
1364 #[tokio::test]
1365 async fn enforce_all_rungs_error_returns_502() {
1366 let (state, mut rx) = enforce_state(
1367 &["anthropic/claude-haiku-4-5"],
1368 &["non-empty"],
1369 vec![(
1370 "anthropic/claude-haiku-4-5",
1371 Err(ProviderError::Transport("down".into())),
1372 )],
1373 );
1374 let resp = messages(
1375 State(state),
1376 Extension(TenantId("default".to_owned())),
1377 HeaderMap::new(),
1378 user_body(),
1379 )
1380 .await;
1381 assert_eq!(resp.status(), axum::http::StatusCode::BAD_GATEWAY);
1382 assert!(rx.try_recv().is_ok());
1384 }
1385
1386 #[tokio::test]
1387 async fn no_routing_config_falls_through_to_observe_not_enforce() {
1388 let config = ProxyConfig::from_lookup(|k| match k {
1392 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
1393 _ => None,
1394 })
1395 .unwrap();
1396 let (traces, _rx) = mpsc::channel(64);
1397 let state = AppState {
1398 config: Arc::new(config),
1399 http: reqwest::Client::new(),
1400 providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
1401 gate_health: Arc::new(GateHealthRegistry::new()),
1402 traces,
1403 adaptive: None,
1404 tenant_rate_limiter: None,
1405 };
1406 let resp = messages(
1407 State(state),
1408 Extension(TenantId("default".to_owned())),
1409 HeaderMap::new(),
1410 user_body(),
1411 )
1412 .await;
1413 assert_ne!(resp.status(), axum::http::StatusCode::OK);
1415 }
1416
1417 #[test]
1418 fn detects_stream_requests() {
1419 assert!(is_stream_request(br#"{"stream": true}"#));
1420 assert!(!is_stream_request(br#"{"stream": false}"#));
1421 assert!(!is_stream_request(br#"{"model":"m"}"#));
1422 assert!(!is_stream_request(b"not json"));
1423 }
1424
1425 #[test]
1426 fn detects_tool_blocks_in_messages() {
1427 let with =
1428 br#"{"messages":[{"role":"user","content":[{"type":"tool_result","content":"42"}]}]}"#;
1429 let without = br#"{"messages":[{"role":"user","content":"hi"}]}"#;
1430 assert!(messages_have_tool_blocks(with));
1431 assert!(!messages_have_tool_blocks(without));
1432 }
1433
1434 #[test]
1435 fn enforce_only_handles_plain_text() {
1436 let plain =
1437 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
1438 let tools = Bytes::from_static(
1439 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
1440 );
1441 let f_plain = extract_features(&HeaderMap::new(), &plain);
1442 let f_tools = extract_features(&HeaderMap::new(), &tools);
1443 assert!(enforce_can_handle(&f_plain, &plain, false));
1445 assert!(!enforce_can_handle(&f_tools, &tools, false));
1446 }
1447
1448 #[test]
1449 fn structured_enforce_routes_tools_and_streaming() {
1450 let tools = Bytes::from_static(
1453 br#"{"model":"m","tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
1454 );
1455 let streaming_tools = Bytes::from_static(
1456 br#"{"model":"m","stream":true,"tools":[{"name":"t"}],"messages":[{"role":"user","content":"hi"}]}"#,
1457 );
1458 let f = extract_features(&HeaderMap::new(), &tools);
1459 assert!(enforce_can_handle(&f, &tools, true));
1460 assert!(enforce_can_handle(&f, &streaming_tools, true));
1461 }
1462
1463 #[test]
1464 fn enforce_sse_reemission_preserves_text_and_tool_use() {
1465 let resp = ModelResponse {
1468 model: "anthropic/claude-haiku-4-5".to_owned(),
1469 text: "let me check".to_owned(),
1470 in_tokens: 5,
1471 out_tokens: 7,
1472 raw: serde_json::json!({
1473 "content": [
1474 { "type": "text", "text": "let me check" },
1475 { "type": "tool_use", "id": "tu_1", "name": "get_weather", "input": { "city": "Paris" } }
1476 ]
1477 }),
1478 };
1479 let sse = anthropic_sse_from_message(&anthropic_response_json(&resp));
1480
1481 let frames: Vec<Value> = sse
1483 .lines()
1484 .filter_map(|l| l.strip_prefix("data: "))
1485 .map(|d| serde_json::from_str::<Value>(d).expect("each SSE data frame is valid JSON"))
1486 .collect();
1487
1488 assert_eq!(frames.first().unwrap()["type"], "message_start");
1490 assert_eq!(frames.last().unwrap()["type"], "message_stop");
1491 assert!(frames.iter().any(|f| f["delta"]["type"] == "text_delta"
1493 && f["delta"]["text"] == "let me check"));
1494 assert!(
1497 frames
1498 .iter()
1499 .any(|f| f["content_block"]["type"] == "tool_use"
1500 && f["content_block"]["name"] == "get_weather"
1501 && f["content_block"]["id"] == "tu_1")
1502 );
1503 assert!(
1504 frames
1505 .iter()
1506 .any(|f| f["delta"]["type"] == "input_json_delta"
1507 && f["delta"]["partial_json"] == r#"{"city":"Paris"}"#)
1508 );
1509 }
1510
1511 #[tokio::test]
1515 async fn enforce_falls_back_to_observe_for_tool_requests() {
1516 let toml = "[[route]]\nmatch = {}\nmode = \"enforce\"\nladder = [\"anthropic/m\"]\ngates = [\"non-empty\"]\n";
1517 let config = ProxyConfig::from_lookup(|k| match k {
1518 "FIRSTPASS_CONFIG_TOML" => Some(toml.to_owned()),
1519 "FIRSTPASS_MODE" => Some("enforce".to_owned()),
1520 "FIRSTPASS_UPSTREAM_ANTHROPIC" => Some("http://127.0.0.1:1".to_owned()),
1521 _ => None,
1522 })
1523 .unwrap();
1524 let mut outs = HashMap::new();
1525 outs.insert(
1526 "anthropic/m".to_owned(),
1527 Ok(model_resp("anthropic/m", "hello")),
1528 );
1529 let mut map: HashMap<String, Arc<dyn Provider>> = HashMap::new();
1530 map.insert(
1531 "anthropic".to_owned(),
1532 Arc::new(MockProvider::new("anthropic", outs)),
1533 );
1534 let (traces, _rx) = mpsc::channel(64);
1535 let state = AppState {
1536 config: Arc::new(config),
1537 http: reqwest::Client::new(),
1538 providers: ProviderRegistry::from_map(map),
1539 gate_health: Arc::new(GateHealthRegistry::new()),
1540 traces,
1541 adaptive: None,
1542 tenant_rate_limiter: None,
1543 };
1544
1545 let plain =
1547 Bytes::from_static(br#"{"model":"m","messages":[{"role":"user","content":"hi"}]}"#);
1548 let resp = messages(
1549 State(state.clone()),
1550 Extension(TenantId("default".to_owned())),
1551 HeaderMap::new(),
1552 plain,
1553 )
1554 .await;
1555 assert_eq!(
1556 resp.status(),
1557 axum::http::StatusCode::OK,
1558 "plain text should enforce"
1559 );
1560
1561 let tools = Bytes::from_static(
1563 br#"{"model":"m","tools":[{"name":"get_weather"}],"messages":[{"role":"user","content":"hi"}]}"#,
1564 );
1565 let resp = messages(
1566 State(state.clone()),
1567 Extension(TenantId("default".to_owned())),
1568 HeaderMap::new(),
1569 tools,
1570 )
1571 .await;
1572 assert_ne!(
1573 resp.status(),
1574 axum::http::StatusCode::OK,
1575 "tool request must fall back to observe, not enforce"
1576 );
1577
1578 let toolres = Bytes::from_static(
1580 br#"{"model":"m","messages":[{"role":"user","content":[{"type":"tool_result","tool_use_id":"x","content":"42"}]}]}"#,
1581 );
1582 let resp = messages(
1583 State(state),
1584 Extension(TenantId("default".to_owned())),
1585 HeaderMap::new(),
1586 toolres,
1587 )
1588 .await;
1589 assert_ne!(
1590 resp.status(),
1591 axum::http::StatusCode::OK,
1592 "tool_result request must fall back to observe, not enforce"
1593 );
1594 }
1595
1596 async fn feedback_state() -> (AppState, std::path::PathBuf, String) {
1600 let db = std::env::temp_dir().join(format!("firstpass-feedback-{}.db", Uuid::now_v7()));
1601 let (tx, handle) = crate::store::open(&db).unwrap();
1602
1603 let mut trace = build_error_trace(
1604 &ProxyConfig::from_lookup(|_| None).unwrap(),
1605 &Bytes::from_static(b"{}"),
1606 5,
1607 Some("sess-fb"),
1608 );
1609 trace.attempts.push(Attempt {
1610 rung: 0,
1611 model: "anthropic/claude-haiku-4-5".into(),
1612 provider: "anthropic".into(),
1613 in_tokens: 10,
1614 out_tokens: 5,
1615 cost_usd: 0.001,
1616 latency_ms: 5,
1617 gates: vec![],
1618 verdict: Verdict::Pass,
1619 });
1620 let trace_id = trace.trace_id.to_string();
1621 tx.try_send(trace).unwrap();
1622 drop(tx);
1623 handle.await.unwrap();
1624
1625 let db_str = db.to_string_lossy().into_owned();
1626 let config = ProxyConfig::from_lookup(move |k| match k {
1627 "FIRSTPASS_DB" => Some(db_str.clone()),
1628 _ => None,
1629 })
1630 .unwrap();
1631 let (traces, _rx) = mpsc::channel(64);
1632 let state = AppState {
1633 config: Arc::new(config),
1634 http: reqwest::Client::new(),
1635 providers: ProviderRegistry::new("http://127.0.0.1:1", "http://127.0.0.1:1"),
1636 gate_health: Arc::new(GateHealthRegistry::new()),
1637 traces,
1638 adaptive: None,
1639 tenant_rate_limiter: None,
1640 };
1641 (state, db, trace_id)
1642 }
1643
1644 #[tokio::test]
1645 async fn feedback_nudges_the_adaptive_threshold() {
1646 use firstpass_core::conformal::AdaptiveConformal;
1647 let (mut state, _db, trace_id) = feedback_state().await;
1648 let aci = Arc::new(std::sync::Mutex::new(AdaptiveConformal::new(0.1, 0.2, 0.5)));
1649 state.adaptive = Some(aci.clone());
1650 let before = aci.lock().unwrap().threshold();
1651
1652 let fail = Bytes::from(
1654 serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "fail", "reporter": "ci" })
1655 .to_string(),
1656 );
1657 assert_eq!(
1658 feedback(
1659 State(state.clone()),
1660 Extension(TenantId("default".to_owned())),
1661 fail
1662 )
1663 .await
1664 .status(),
1665 axum::http::StatusCode::ACCEPTED
1666 );
1667 let after_fail = aci.lock().unwrap().threshold();
1668 assert!(
1669 after_fail > before,
1670 "served fail should raise the live threshold: {before} -> {after_fail}"
1671 );
1672
1673 let pass = Bytes::from(
1675 serde_json::json!({ "trace_id": trace_id, "gate_id": "tests", "verdict": "pass", "reporter": "ci" })
1676 .to_string(),
1677 );
1678 let _ = feedback(
1679 State(state),
1680 Extension(TenantId("default".to_owned())),
1681 pass,
1682 )
1683 .await;
1684 assert!(aci.lock().unwrap().threshold() < after_fail);
1685 }
1686
1687 #[tokio::test]
1688 async fn feedback_records_a_deferred_verdict_without_breaking_the_chain() {
1689 let (state, db, trace_id) = feedback_state().await;
1690 let body = Bytes::from(
1691 serde_json::json!({
1692 "trace_id": trace_id,
1693 "gate_id": "tests",
1694 "verdict": "pass",
1695 "score": 1.0,
1696 "reporter": "ci",
1697 })
1698 .to_string(),
1699 );
1700 let resp = feedback(
1701 State(state),
1702 Extension(TenantId("default".to_owned())),
1703 body,
1704 )
1705 .await;
1706 assert_eq!(resp.status(), axum::http::StatusCode::ACCEPTED);
1707
1708 let view = crate::store::load_trace_view(&db, "default", &trace_id)
1710 .unwrap()
1711 .unwrap();
1712 assert_eq!(view.deferred.len(), 1);
1713 assert_eq!(view.deferred[0].gate_id, "tests");
1714 let traces = crate::store::load_all_traces(&db).unwrap();
1716 firstpass_core::verify_chain(&traces, GENESIS_HASH).unwrap();
1717
1718 let _ = std::fs::remove_file(&db);
1719 }
1720
1721 #[tokio::test]
1722 async fn feedback_for_unknown_trace_is_404() {
1723 let (state, db, _trace_id) = feedback_state().await;
1724 let body = Bytes::from(
1725 serde_json::json!({
1726 "trace_id": "does-not-exist",
1727 "gate_id": "tests",
1728 "verdict": "pass",
1729 "reporter": "ci",
1730 })
1731 .to_string(),
1732 );
1733 let resp = feedback(
1734 State(state),
1735 Extension(TenantId("default".to_owned())),
1736 body,
1737 )
1738 .await;
1739 assert_eq!(resp.status(), axum::http::StatusCode::NOT_FOUND);
1740 let _ = std::fs::remove_file(&db);
1741 }
1742
1743 #[tokio::test]
1746 async fn feedback_across_tenants_is_404_not_403() {
1747 let (state, db, trace_id) = feedback_state().await;
1748 let body = Bytes::from(
1749 serde_json::json!({
1750 "trace_id": trace_id,
1751 "gate_id": "tests",
1752 "verdict": "pass",
1753 "score": 1.0,
1754 "reporter": "attacker",
1755 })
1756 .to_string(),
1757 );
1758 let resp = feedback(
1760 State(state),
1761 Extension(TenantId("tenant-b".to_owned())),
1762 body,
1763 )
1764 .await;
1765 assert_eq!(
1766 resp.status(),
1767 axum::http::StatusCode::NOT_FOUND,
1768 "cross-tenant feedback must look exactly like a missing trace"
1769 );
1770 let _ = std::fs::remove_file(&db);
1771 }
1772
1773 #[tokio::test]
1774 async fn feedback_rejects_bad_verdict_and_score() {
1775 let (state, db, trace_id) = feedback_state().await;
1776 let bad_verdict = Bytes::from(
1777 serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "maybe", "reporter": "x" })
1778 .to_string(),
1779 );
1780 assert_eq!(
1781 feedback(
1782 State(state.clone()),
1783 Extension(TenantId("default".to_owned())),
1784 bad_verdict
1785 )
1786 .await
1787 .status(),
1788 axum::http::StatusCode::BAD_REQUEST
1789 );
1790 let bad_score = Bytes::from(
1791 serde_json::json!({ "trace_id": trace_id, "gate_id": "g", "verdict": "pass", "score": 9.0, "reporter": "x" })
1792 .to_string(),
1793 );
1794 assert_eq!(
1795 feedback(
1796 State(state),
1797 Extension(TenantId("default".to_owned())),
1798 bad_score
1799 )
1800 .await
1801 .status(),
1802 axum::http::StatusCode::BAD_REQUEST
1803 );
1804 let _ = std::fs::remove_file(&db);
1805 }
1806
1807 #[tokio::test]
1808 async fn metrics_endpoint_renders_after_a_real_request() {
1809 use tower::ServiceExt;
1810
1811 let (state, mut rx) = enforce_state(
1812 &["anthropic/claude-haiku-4-5"],
1813 &["non-empty"],
1814 vec![(
1815 "anthropic/claude-haiku-4-5",
1816 Ok(model_resp("anthropic/claude-haiku-4-5", "hello")),
1817 )],
1818 );
1819 let router = app(state).expect("prometheus recorder installs");
1820
1821 let req = axum::http::Request::builder()
1822 .method("POST")
1823 .uri("/v1/messages")
1824 .header("content-type", "application/json")
1825 .body(Body::from(user_body()))
1826 .unwrap();
1827 let resp = router.clone().oneshot(req).await.unwrap();
1828 assert_eq!(resp.status(), axum::http::StatusCode::OK);
1829 rx.try_recv().expect("a trace was enqueued");
1830
1831 let metrics_req = axum::http::Request::builder()
1832 .method("GET")
1833 .uri("/metrics")
1834 .body(Body::empty())
1835 .unwrap();
1836 let metrics_resp = router.oneshot(metrics_req).await.unwrap();
1837 assert_eq!(metrics_resp.status(), axum::http::StatusCode::OK);
1838 let bytes = axum::body::to_bytes(metrics_resp.into_body(), 1 << 20)
1839 .await
1840 .unwrap();
1841 let body = String::from_utf8(bytes.to_vec()).unwrap();
1842 assert!(
1843 body.contains("firstpass_enforce_latency_ms"),
1844 "metrics body missing enforce latency histogram: {body}"
1845 );
1846 assert!(
1847 body.contains("firstpass_served_total"),
1848 "metrics body missing served counter: {body}"
1849 );
1850 }
1851
1852 fn auth_state(require_auth: bool, keys_json: Option<String>) -> AppState {
1856 auth_state_rated(require_auth, keys_json, None)
1857 }
1858
1859 fn auth_state_rated(
1862 require_auth: bool,
1863 keys_json: Option<String>,
1864 rate_per_sec: Option<u32>,
1865 ) -> AppState {
1866 let config = ProxyConfig::from_lookup(|k| match k {
1867 "FIRSTPASS_REQUIRE_AUTH" => require_auth.then(|| "true".to_owned()),
1868 "FIRSTPASS_TENANT_KEYS_JSON" => keys_json.clone(),
1869 "FIRSTPASS_TENANT_RATE_PER_SEC" => rate_per_sec.map(|n| n.to_string()),
1870 _ => None,
1871 })
1872 .unwrap();
1873 let (traces, _rx) = mpsc::channel(64);
1874 std::mem::forget(_rx);
1877 let providers: HashMap<String, Arc<dyn Provider>> = HashMap::new();
1878 let tenant_rate_limiter = build_tenant_rate_limiter(&config);
1879 AppState {
1880 config: Arc::new(config),
1881 http: reqwest::Client::new(),
1882 providers: ProviderRegistry::from_map(providers),
1883 gate_health: Arc::new(GateHealthRegistry::new()),
1884 traces,
1885 adaptive: None,
1886 tenant_rate_limiter,
1887 }
1888 }
1889
1890 fn cap_request(auth_header: Option<&str>) -> axum::http::Request<Body> {
1891 let mut b = axum::http::Request::builder()
1892 .method("GET")
1893 .uri("/v1/capabilities");
1894 if let Some(h) = auth_header {
1895 b = b.header("authorization", h);
1896 }
1897 b.body(Body::empty()).unwrap()
1898 }
1899
1900 #[tokio::test]
1901 async fn auth_off_allows_unauthenticated_request() {
1902 use tower::ServiceExt;
1903 let router = app(auth_state(false, None)).expect("router");
1904 let resp = router.oneshot(cap_request(None)).await.unwrap();
1905 assert_eq!(resp.status(), axum::http::StatusCode::OK);
1907 }
1908
1909 #[tokio::test]
1910 async fn auth_on_missing_key_is_401_opaque() {
1911 use tower::ServiceExt;
1912 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1913 let keys = format!("{{\"tenant-a\": {hash:?}}}");
1914 let router = app(auth_state(true, Some(keys))).expect("router");
1915
1916 let resp = router.oneshot(cap_request(None)).await.unwrap();
1917 assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
1918 let json = body_json(resp).await;
1919 assert_eq!(json["error"]["type"], "unauthorized");
1920 let msg = json["error"]["message"].as_str().unwrap();
1922 assert!(!msg.contains("tenant"), "no tenant oracle in body: {msg}");
1923 }
1924
1925 #[tokio::test]
1926 async fn auth_on_invalid_key_is_401() {
1927 use tower::ServiceExt;
1928 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1929 let keys = format!("{{\"tenant-a\": {hash:?}}}");
1930 let router = app(auth_state(true, Some(keys))).expect("router");
1931
1932 let resp = router
1933 .oneshot(cap_request(Some("Bearer wrong-key")))
1934 .await
1935 .unwrap();
1936 assert_eq!(resp.status(), axum::http::StatusCode::UNAUTHORIZED);
1937 }
1938
1939 #[tokio::test]
1940 async fn auth_on_valid_key_proceeds() {
1941 use tower::ServiceExt;
1942 let hash = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1943 let keys = format!("{{\"tenant-a\": {hash:?}}}");
1944 let router = app(auth_state(true, Some(keys))).expect("router");
1945
1946 let resp = router
1948 .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
1949 .await
1950 .unwrap();
1951 assert_eq!(resp.status(), axum::http::StatusCode::OK);
1953 }
1954
1955 fn two_tenant_state(rate_per_sec: Option<u32>) -> AppState {
1957 let hash_a = crate::tenant_auth::TenantKeys::hash_key("key-a").unwrap();
1958 let hash_b = crate::tenant_auth::TenantKeys::hash_key("key-b").unwrap();
1959 let keys = format!("{{\"tenant-a\": {hash_a:?}, \"tenant-b\": {hash_b:?}}}");
1960 auth_state_rated(true, Some(keys), rate_per_sec)
1961 }
1962
1963 #[tokio::test]
1964 async fn tenant_exceeding_rate_limit_gets_429_opaque() {
1965 use tower::ServiceExt;
1966 let router = app(two_tenant_state(Some(1))).expect("router");
1972
1973 let (r1, r2, r3, r4) = tokio::join!(
1974 router
1975 .clone()
1976 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
1977 router
1978 .clone()
1979 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
1980 router
1981 .clone()
1982 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
1983 router
1984 .clone()
1985 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
1986 );
1987 let responses = [r1.unwrap(), r2.unwrap(), r3.unwrap(), r4.unwrap()];
1988 let ok = responses
1989 .iter()
1990 .filter(|r| r.status() == axum::http::StatusCode::OK)
1991 .count();
1992 assert!(ok >= 1, "the burst's first request must pass");
1993 let limited: Vec<_> = responses
1994 .into_iter()
1995 .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
1996 .collect();
1997 assert!(
1998 !limited.is_empty(),
1999 "a 4-request burst against 1 req/sec must trip the limiter"
2000 );
2001
2002 let json = body_json(limited.into_iter().next().unwrap()).await;
2003 assert_eq!(json["error"]["type"], "rate_limited");
2004 let msg = json["error"]["message"].as_str().unwrap();
2006 assert!(!msg.contains('1'), "no limit value in body: {msg}");
2007 }
2008
2009 #[tokio::test]
2010 async fn rate_limit_buckets_are_independent_per_tenant() {
2011 use tower::ServiceExt;
2012 let router = app(two_tenant_state(Some(1))).expect("router");
2013
2014 let (a1, a2, a3) = tokio::join!(
2017 router
2018 .clone()
2019 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
2020 router
2021 .clone()
2022 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
2023 router
2024 .clone()
2025 .oneshot(cap_request(Some("Bearer tenant-a.key-a"))),
2026 );
2027 let a_limited = [a1.unwrap(), a2.unwrap(), a3.unwrap()]
2028 .iter()
2029 .filter(|r| r.status() == axum::http::StatusCode::TOO_MANY_REQUESTS)
2030 .count();
2031 assert!(a_limited >= 1, "tenant A's burst must trip its limiter");
2032
2033 let b1 = router
2035 .clone()
2036 .oneshot(cap_request(Some("Bearer tenant-b.key-b")))
2037 .await
2038 .unwrap();
2039 assert_eq!(b1.status(), axum::http::StatusCode::OK);
2040 }
2041
2042 #[tokio::test]
2043 async fn rate_limit_unset_never_429s() {
2044 use tower::ServiceExt;
2045 let router = app(two_tenant_state(None)).expect("router");
2048 for _ in 0..20 {
2049 let resp = router
2050 .clone()
2051 .oneshot(cap_request(Some("Bearer tenant-a.key-a")))
2052 .await
2053 .unwrap();
2054 assert_eq!(resp.status(), axum::http::StatusCode::OK);
2055 }
2056 }
2057}