1use crate::pipeline::AppState;
4use axum::body::{Body, Bytes};
5use axum::extract::{Path, State};
6use axum::http::{HeaderMap, HeaderValue, Request, StatusCode};
7use axum::middleware::Next;
8use axum::response::{IntoResponse, Response};
9use axum::routing::{get, post};
10use axum::{Json, Router};
11use futures::stream::BoxStream;
12use futures::{FutureExt, Stream, StreamExt};
13use pv_core::time::elapsed_us;
14use pv_events::StopReason;
15use pv_sandbox::ToolVerdict;
16use serde_json::{json, Value};
17use std::future::Future as _;
18use std::pin::Pin;
19use std::sync::Arc;
20use std::task::{Context, Poll};
21use tracing::Instrument as _;
22
23pub fn build_router(state: AppState) -> Router {
25 let max_body = state.config.max_request_bytes;
26 let dashboard_enabled = state.config.dashboard_enabled;
27 let mut router = Router::new()
28 .route("/health", get(health))
29 .route("/livez", get(livez))
30 .route("/readyz", get(readyz))
31 .route("/metrics", get(metrics))
32 .route("/v1/chat/completions", post(chat_completions).options(cors_deny))
33 .route(
34 "/v1/mcp",
35 post(mcp_call).get(mcp_get).delete(mcp_delete).options(cors_deny),
36 )
37 .route(
38 "/mcp",
39 post(mcp_call).get(mcp_get).delete(mcp_delete).options(cors_deny),
40 )
41 .route(
42 "/v1/mcp/{backend}",
43 post(mcp_backend_call)
44 .get(mcp_get)
45 .delete(mcp_delete)
46 .options(cors_deny),
47 )
48 .route(
49 "/mcp/{backend}",
50 post(mcp_backend_call)
51 .get(mcp_get)
52 .delete(mcp_delete)
53 .options(cors_deny),
54 )
55 .route("/v1/sessions/{id}/close", post(close_session).options(cors_deny))
56 .route(
57 "/v1/sessions/{id}/promote",
58 post(promote_session).options(cors_deny),
59 )
60 .route("/v1/token", post(token_exchange).options(cors_deny))
61 .route("/v1/revoke", post(crate::token_routes::revoke).options(cors_deny))
62 .route("/.well-known/jwks.json", get(jwks));
63 if !state.config.introspection_tokens.is_empty() {
64 router = router.route(
65 "/v1/introspect",
66 post(crate::token_routes::introspect).options(cors_deny),
67 );
68 }
69 if !state.config.operator_tokens.is_empty() {
70 router = router.route(
71 "/admin/v1/revocations",
72 post(crate::token_routes::admin_revoke).options(cors_deny),
73 );
74 }
75 if dashboard_enabled {
76 router = router
77 .route("/dashboard", get(crate::dashboard::index))
78 .route("/dashboard/", get(crate::dashboard::index))
79 .route("/dashboard/style.css", get(crate::dashboard::style_css))
80 .route("/dashboard/app.js", get(crate::dashboard::app_js))
81 .route("/api/v1/dashboard/stats", get(crate::dashboard::stats))
82 .route("/api/v1/dashboard/sessions", get(crate::dashboard::list_sessions))
83 .route(
84 "/api/v1/dashboard/sessions/{id}",
85 get(crate::dashboard::session_detail),
86 );
87 }
88 router
89 .layer(axum::middleware::from_fn_with_state(
90 state.clone(),
91 enforce_allowed_hosts,
92 ))
93 .layer(axum::middleware::from_fn_with_state(
94 state.clone(),
95 request_metrics,
96 ))
97 .layer(axum::middleware::from_fn(trace_request))
98 .layer(axum::middleware::from_fn_with_state(state.clone(), explain_body_limit))
106 .layer(axum::extract::DefaultBodyLimit::max(max_body))
112 .with_state(state)
113}
114
115async fn explain_body_limit(
119 State(state): State<AppState>,
120 request: Request<Body>,
121 next: axum::middleware::Next,
122) -> Response {
123 let response = next.run(request).await;
124 if response.status() != StatusCode::PAYLOAD_TOO_LARGE {
125 return response;
126 }
127 let is_json = response
130 .headers()
131 .get(axum::http::header::CONTENT_TYPE)
132 .and_then(|value| value.to_str().ok())
133 .is_some_and(|value| value.starts_with("application/json"));
134 if is_json {
135 return response;
136 }
137 let limit = state.config.max_request_bytes;
138 (
139 StatusCode::PAYLOAD_TOO_LARGE,
140 Json(serde_json::json!({
141 "error": {
142 "message": format!(
143 "request body exceeds max_request_bytes ({limit} bytes); \
144 raise `max_request_bytes` in the harness config if larger \
145 payloads are intended"
146 ),
147 "type": "invalid_request_error",
148 "param": serde_json::Value::Null,
149 "code": StatusCode::PAYLOAD_TOO_LARGE.as_u16(),
155 }
156 })),
157 )
158 .into_response()
159}
160
161async fn enforce_allowed_hosts(
170 State(state): State<AppState>,
171 request: Request<Body>,
172 next: Next,
173) -> Response {
174 if state.config.allowed_hosts.is_empty() {
175 return next.run(request).await;
176 }
177 let host = request
178 .headers()
179 .get(axum::http::header::HOST)
180 .and_then(|value| value.to_str().ok())
181 .map(|value| {
182 let stripped = value
184 .rsplit_once(':')
185 .filter(|(head, port)| !head.is_empty() && port.chars().all(|c| c.is_ascii_digit()))
186 .map_or(value, |(head, _)| head);
187 stripped.trim_start_matches('[').trim_end_matches(']')
188 });
189 let allowed = host.is_some_and(|host| {
190 state.config.allowed_hosts.iter().any(|entry| {
191 entry
192 .trim_start_matches('[')
193 .trim_end_matches(']')
194 .eq_ignore_ascii_case(host)
195 })
196 });
197 if !allowed {
198 state
199 .metrics
200 .counter(
201 "pv_host_header_refused_total",
202 "Requests refused by the allowed_hosts Host-header allowlist",
203 )
204 .inc();
205 return (
206 StatusCode::FORBIDDEN,
207 Json(json!({
208 "error": {
209 "message": "Host header is not in allowed_hosts",
210 "type": "permission_error",
211 "param": null,
212 "code": 403,
213 }
214 })),
215 )
216 .into_response();
217 }
218 next.run(request).await
219}
220
221async fn request_metrics(State(state): State<AppState>, request: Request<Body>, next: Next) -> Response {
239 let route = route(request.uri().path());
240 let started = std::time::Instant::now();
241 let response = next.run(request).await;
242 let status_class = crate::pipeline::StatusClass::from_u16(response.status().as_u16());
243 #[allow(clippy::indexing_slicing)]
250 {
251 state.hot_metrics.request_counters[route.index()][status_class.index()].inc();
252 state.hot_metrics.request_duration_histograms[route.index()].observe_us(elapsed_us(started));
253 }
254 response
255}
256
257fn route(path: &str) -> crate::pipeline::Route {
261 use crate::pipeline::Route;
262 match path {
263 "/v1/chat/completions" => Route::Chat,
264 "/v1/mcp" | "/mcp" => Route::Mcp,
265 "/health" => Route::Health,
266 "/livez" => Route::Livez,
267 "/readyz" => Route::Readyz,
268 "/metrics" => Route::Metrics,
269 "/v1/token" => Route::TokenExchange,
270 "/v1/revoke" => Route::Revoke,
271 "/v1/introspect" => Route::Introspect,
272 "/admin/v1/revocations" => Route::AdminRevocation,
273 "/.well-known/jwks.json" => Route::Jwks,
274 _ if path.starts_with("/v1/mcp/") || path.starts_with("/mcp/") => Route::Mcp,
275 _ if path.starts_with("/v1/sessions/") && path.ends_with("/close") => Route::SessionClose,
276 _ if path.starts_with("/v1/sessions/") && path.ends_with("/promote") => Route::SessionPromote,
277 _ if path.starts_with("/dashboard") || path.starts_with("/api/v1/dashboard") => Route::Dashboard,
278 _ => Route::Other,
279 }
280}
281
282async fn trace_request(request: Request<Body>, next: Next) -> Response {
283 let method = request.method().clone();
284 let path = request.uri().path().to_owned();
285 let headers = request.headers();
304 let session_id = match crate::pipeline::single_header(headers, crate::pipeline::SESSION_HEADER) {
305 Ok(Some(value)) => value
306 .to_str()
307 .ok()
308 .and_then(|v| pv_core::SessionId::parse(v).ok())
309 .map(|id| id.to_string())
310 .unwrap_or_else(|| " rejected".to_owned()),
311 Ok(None) => " unbound".to_owned(),
312 Err(_) => " duplicate-header".to_owned(),
313 };
314 let span = tracing::info_span!(
315 "proofvisor.request",
316 otel.kind = "server",
317 http.request.method = %method,
318 url.path = %path,
319 session.id = %session_id,
320 http.response.status_code = tracing::field::Empty,
321 );
322 let response = next.run(request).instrument(span.clone()).await;
323 span.record("http.response.status_code", response.status().as_u16());
324 response
325}
326
327async fn health() -> impl IntoResponse {
328 Json(json!({
336 "status": "ok",
337 "service": "proofvisor",
342 }))
343}
344
345async fn livez() -> impl IntoResponse {
355 Json(json!({
356 "status": "alive",
357 "service": "proofvisor",
358 }))
359}
360
361async fn readyz(State(state): State<AppState>) -> impl IntoResponse {
375 let draining = state.draining.load(std::sync::atomic::Ordering::SeqCst);
376 let spool_writable = {
393 let dir = state.config.atif_spool_dir.clone();
394 tokio::task::spawn_blocking(move || {
395 let probe = std::path::Path::new(&dir).join(format!(".readyz-{}.tmp", pv_core::new_event_uid()));
396 std::fs::write(&probe, b"readyz")
397 .and_then(|()| {
398 std::fs::remove_file(&probe)
402 })
403 .is_ok()
404 })
405 .await
406 .unwrap_or(false)
407 };
408 let revocation_health: Vec<_> = state
411 .identity
412 .as_ref()
413 .and_then(|validator| validator.revocation_store())
414 .and_then(|store| store.health())
415 .into_iter()
416 .chain(state.exchanged_revocations.health())
417 .collect();
418 let revocation_available = revocation_health.iter().all(|health| health.available);
419 let local_revocations: usize = revocation_health.iter().map(|health| health.local_entries).sum();
420 let ready = !draining && spool_writable;
421 let body = Json(json!({
422 "status": if ready { "ready" } else { "not_ready" },
423 "service": "proofvisor",
424 "checks": {
425 "draining": draining,
426 "spool_dir_writable": spool_writable,
427 "revocation_available": revocation_available,
428 "revocation_local_entries": local_revocations,
429 },
430 }));
431 if ready {
432 (StatusCode::OK, body).into_response()
433 } else {
434 (StatusCode::SERVICE_UNAVAILABLE, body).into_response()
435 }
436}
437
438async fn cors_deny() -> Response {
450 let mut response = Response::new(Body::empty());
451 *response.status_mut() = StatusCode::NO_CONTENT;
452 response.headers_mut().insert(
453 axum::http::header::CACHE_CONTROL,
454 HeaderValue::from_static("no-store"),
455 );
456 response
460}
461
462async fn metrics(State(state): State<AppState>) -> Response {
463 state
470 .metrics
471 .gauge(
472 "pv_open_sessions",
473 "Currently-open sessions in the in-memory registry",
474 )
475 .set(state.sessions.len() as u64);
476 state
482 .metrics
483 .gauge(
484 "pv_worker_queue_depth",
485 "Accepted-but-not-yet-completed worker jobs",
486 )
487 .set(state.worker.queue_depth());
488 let spool_dir = state.config.atif_spool_dir.clone();
502 let (spool_bytes, spool_files) =
503 tokio::task::spawn_blocking(move || spool_footprint(std::path::Path::new(&spool_dir)))
504 .await
505 .unwrap_or((0, 0));
506 state
507 .metrics
508 .gauge(
509 "pv_spool_bytes",
510 "Total bytes in the top-level ATIF spool directory",
511 )
512 .set(spool_bytes);
513 state
514 .metrics
515 .gauge(
516 "pv_spool_files",
517 "File count in the top-level ATIF spool directory",
518 )
519 .set(spool_files);
520 let mut response = Response::new(Body::from(state.metrics.render()));
521 response.headers_mut().insert(
522 axum::http::header::CONTENT_TYPE,
523 HeaderValue::from_static("text/plain; version=0.0.4; charset=utf-8"),
524 );
525 response
526}
527
528fn spool_footprint(dir: &std::path::Path) -> (u64, u64) {
533 let Ok(entries) = std::fs::read_dir(dir) else {
534 return (0, 0);
535 };
536 let mut bytes = 0u64;
537 let mut files = 0u64;
538 for entry in entries.flatten() {
539 if let Ok(metadata) = entry.metadata() {
540 if metadata.is_file() {
541 bytes = bytes.saturating_add(metadata.len());
542 files = files.saturating_add(1);
543 }
544 }
545 }
546 (bytes, files)
547}
548
549fn refuse_non_json_content_type(headers: &HeaderMap, context_label: &str) -> Option<Response> {
560 let value = headers.get(axum::http::header::CONTENT_TYPE)?;
561 let is_json = value
562 .to_str()
563 .ok()
564 .and_then(|s| s.split(';').next())
565 .map(|main| main.trim().eq_ignore_ascii_case("application/json"))
566 .unwrap_or(false);
567 if is_json {
568 return None;
569 }
570 let display = value.to_str().unwrap_or("<non-ascii>");
571 Some(pipeline_error(crate::pipeline::PipelineError::bad_request(
572 format!("{context_label} requires Content-Type: application/json, got {display:?}"),
573 )))
574}
575
576async fn chat_completions(State(state): State<AppState>, headers: HeaderMap, body: Bytes) -> Response {
577 if let Some(err) = refuse_non_json_content_type(&headers, "chat completions") {
589 return err;
590 }
591 let caller = match state.authenticate_chat_nonblocking(&headers).await {
603 Ok(caller) => caller,
604 Err(error) => {
605 if let Err(audit_error) = state.audit_authentication_error(&headers, &error) {
606 return pipeline_error(audit_error);
607 }
608 return pipeline_error(error);
609 }
610 };
611 if let Err(reason) = pv_sandbox::refuse_duplicate_json_keys(&body) {
623 return pipeline_error(crate::pipeline::PipelineError::bad_request(format!(
624 "chat request rejected: {reason}"
625 )));
626 }
627 let payload: Value = match serde_json::from_slice(&body) {
628 Ok(value) => value,
629 Err(error) => return pipeline_error(crate::pipeline::PipelineError::bad_request_source(error)),
630 };
631 if pv_sandbox::rpc::depth_of(&payload, 0) > pv_sandbox::rpc::MAX_JSON_DEPTH {
639 return pipeline_error(crate::pipeline::PipelineError::bad_request(format!(
640 "chat request nesting depth exceeds {}",
641 pv_sandbox::rpc::MAX_JSON_DEPTH
642 )));
643 }
644 let admission_started = std::time::Instant::now();
645 let mut prepared = match state
646 .prepare_authenticated_chat(&headers, payload, body.len(), caller)
647 .await
648 {
649 Ok(prepared) => prepared,
650 Err(error) => return pipeline_error(error),
651 };
652 prepared.middleware_us = elapsed_us(admission_started);
653 let session = Arc::clone(&prepared.session);
654 let identity = prepared.identity.clone();
655 let session_id = prepared.session.id.clone();
656 let middleware_us = prepared.middleware_us;
657 let forwarded = match state.forward_chat(prepared).await {
658 Ok(response) => response,
659 Err(error) => return pipeline_error_for_session(error, &session_id),
660 };
661 let crate::pipeline::ForwardedResponse {
662 status,
663 headers: upstream_headers,
664 body: upstream_body,
665 billed_prompt_tokens,
666 lease,
667 mut capture_guard,
668 } = forwarded;
669 let Some(response_permit) = capture_guard.take_permit() else {
670 let (response_marker, response_attempt_id) = capture_guard.disarm();
678 let capture_session = Arc::clone(&session);
679 if state
680 .worker
681 .try_submit(crate::pipeline::refused_response_failure_job(
682 session,
683 identity,
684 "response_capture_permit_missing".to_owned(),
685 response_marker,
686 response_attempt_id,
687 ))
688 .is_err()
689 {
690 capture_session.mark_capture_failed();
691 }
692 let mut response = lifecycle_error("durable response capture permit is missing".to_owned());
693 if let Ok(value) = HeaderValue::from_str(&session_id) {
694 response
695 .headers_mut()
696 .insert(crate::pipeline::SESSION_HEADER, value);
697 }
698 return response;
699 };
700
701 let mut refused_encoding: Option<(String, String)> = None;
718 'encodings: for value in upstream_headers.get_all(axum::http::header::CONTENT_ENCODING) {
719 let raw = value.to_str().unwrap_or_default();
720 for token in raw.split(',') {
721 let token = token.trim();
722 if token.is_empty() || token.eq_ignore_ascii_case("identity") {
723 continue;
724 }
725 refused_encoding = Some((token.to_owned(), raw.to_owned()));
726 break 'encodings;
727 }
728 }
729 if let Some((token, raw)) = refused_encoding {
730 let (response_marker, response_attempt_id) = capture_guard.disarm();
741 let capture_session = Arc::clone(&session);
742 if response_permit
743 .submit(
744 &state.worker,
745 crate::pipeline::refused_response_failure_job(
746 session,
747 identity,
748 "upstream_unsupported_content_encoding".to_owned(),
749 response_marker,
750 response_attempt_id,
751 ),
752 )
753 .is_err()
754 {
755 capture_session.mark_capture_failed();
756 }
757 return pipeline_error_for_session(
758 crate::pipeline::PipelineError::upstream(format!(
759 "upstream responded with unsupported Content-Encoding token {token:?} \
760 (full header: {raw:?}) — the proxy is built without decompression \
761 support; enable it upstream (Accept-Encoding: identity) or rebuild \
762 with the reqwest `gzip` feature"
763 )),
764 &session_id,
765 );
766 }
767 let is_sse = is_sse_content_type(&upstream_headers);
778 let (response_marker, response_attempt_id) = capture_guard.disarm();
782 let principal_budget = state
787 .config
788 .principal_budget
789 .as_ref()
790 .map(|spec| (crate::pipeline::principal_id_for_budget(&identity), spec.clone()));
791 let stream = AbortFinalizingStream {
792 inner: upstream_body,
793 session,
794 identity,
795 response_permit: Some(response_permit),
796 worker: state.worker.clone(),
797 store: Arc::clone(&state.store),
798 budget: state.effective_budget(),
799 rates: Arc::clone(&state.pushed_rates),
800 finalizer: state.finalizer.clone(),
801 _lease: lease,
802 response_marker,
803 response_attempt_id,
804 response_choices: std::collections::BTreeMap::new(),
805 response_model: None,
806 upstream_status: status,
807 response_cost_usd_micros: 0,
808 response_tool_calls: std::collections::BTreeMap::new(),
809 response_metrics: pv_events::EventMetrics::default(),
810 charged_completion_tokens: 0,
811 last_reported_completion_tokens: None,
812 last_reported_prompt_tokens: None,
813 last_reported_cached_tokens: None,
814 last_reported_cost_usd_micros: None,
815 saw_chunk: false,
816 capture_attempted: false,
817 is_sse,
818 protocol_buffer: Vec::new(),
819 frame_scratch: Vec::new(),
820 pending_output: std::collections::VecDeque::new(),
821 pending_budget: None,
822 principal_budget,
823 captured_bytes: 0,
824 billed_prompt_tokens,
825 completed: false,
826 ephemeral: !headers.contains_key(crate::pipeline::SESSION_HEADER),
829 };
830 let mut response = if is_sse {
831 let mut relay = Box::pin(stream);
844 match relay.next().await {
845 Some(Err(error)) => {
846 let mapped = if error.kind() == std::io::ErrorKind::QuotaExceeded {
847 crate::pipeline::PipelineError::blocked_source(error)
848 } else {
849 crate::pipeline::PipelineError::upstream_source(error)
850 };
851 return pipeline_error_for_session(mapped, &session_id);
854 }
855 Some(Ok(first)) => {
856 let head = futures::stream::once(async move { Ok::<_, std::io::Error>(first) });
857 Response::new(Body::from_stream(head.chain(relay)))
858 }
859 None => Response::new(Body::empty()),
862 }
863 } else {
864 let mut relay = Box::pin(stream);
872 let mut buffered: Vec<u8> = Vec::new();
873 loop {
874 match relay.next().await {
875 Some(Ok(bytes)) => buffered.extend_from_slice(&bytes),
876 Some(Err(error)) => {
877 let mapped = if error.kind() == std::io::ErrorKind::QuotaExceeded {
887 crate::pipeline::PipelineError::blocked_source(error)
888 } else {
889 crate::pipeline::PipelineError::upstream_source(error)
890 };
891 return pipeline_error_for_session(mapped, &session_id);
895 }
896 None => break,
897 }
898 }
899 Response::new(Body::from(buffered))
900 };
901 *response.status_mut() = status;
902 for (name, value) in &upstream_headers {
903 if is_forwardable_upstream_header(name, &state.config.upstream_auth_header) {
904 response.headers_mut().append(name.clone(), value.clone());
908 }
909 }
910 if let Ok(value) = HeaderValue::from_str(&session_id) {
911 response
912 .headers_mut()
913 .insert(crate::pipeline::SESSION_HEADER, value);
914 }
915 if let Ok(value) = HeaderValue::from_str(&middleware_us.to_string()) {
916 response
917 .headers_mut()
918 .insert(crate::pipeline::MIDDLEWARE_US_HEADER, value);
919 }
920 response.headers_mut().insert(
932 axum::http::HeaderName::from_static("x-content-type-options"),
933 HeaderValue::from_static("nosniff"),
934 );
935 if is_sse {
945 response.headers_mut().insert(
946 axum::http::HeaderName::from_static("x-accel-buffering"),
947 HeaderValue::from_static("no"),
948 );
949 }
950 response
951}
952
953fn is_forwardable_upstream_header(name: &axum::http::HeaderName, upstream_auth_header: &str) -> bool {
985 use axum::http::header;
986 if name.as_str().eq_ignore_ascii_case(upstream_auth_header) {
990 return false;
991 }
992 let is_denied = *name == header::CONTENT_LENGTH
993 || *name == header::TRANSFER_ENCODING
994 || *name == header::CONNECTION
995 || *name == header::UPGRADE
996 || *name == header::TE
997 || *name == header::TRAILER
998 || *name == header::PROXY_AUTHENTICATE
999 || *name == header::PROXY_AUTHORIZATION
1000 || *name == header::AUTHORIZATION
1004 || *name == header::SET_COOKIE
1005 || *name == header::SERVER
1006 || *name == header::VIA
1007 || name.as_str().eq_ignore_ascii_case("keep-alive")
1008 || name.as_str().eq_ignore_ascii_case("x-powered-by")
1009 || name.as_str().eq_ignore_ascii_case("x-request-id")
1010 || name.as_str().eq_ignore_ascii_case("api-key")
1015 || name.as_str().eq_ignore_ascii_case("x-api-key")
1016 || name.as_str().eq_ignore_ascii_case("x-goog-api-key")
1017 || name.as_str().eq_ignore_ascii_case("anthropic-api-key")
1018 || name.as_str().eq_ignore_ascii_case("openai-api-key")
1019 || name.as_str().eq_ignore_ascii_case("x-auth-token")
1020 || name.as_str().eq_ignore_ascii_case("x-amz-security-token")
1021 || name.as_str().to_ascii_lowercase().starts_with("access-control-");
1024 !is_denied
1025}
1026
1027fn is_sse_content_type(headers: &HeaderMap) -> bool {
1036 headers
1037 .get(axum::http::header::CONTENT_TYPE)
1038 .and_then(|value| value.to_str().ok())
1039 .is_some_and(|value| {
1040 let head = value.split(';').next().unwrap_or("").trim();
1041 head.eq_ignore_ascii_case("text/event-stream")
1042 })
1043}
1044
1045fn tool_audit_gate(state: &AppState, key: &str) -> Arc<tokio::sync::Mutex<()>> {
1049 let mut gates = state.tool_audit_gates.lock();
1050 gates.retain(|_, gate| Arc::strong_count(gate) > 1);
1051 Arc::clone(gates.entry(key.to_owned()).or_default())
1052}
1053
1054async fn mcp_call(State(state): State<AppState>, headers: HeaderMap, body: Bytes) -> Response {
1055 mcp_entry(state, headers, body, crate::mcp::Scope::All).await
1056}
1057
1058async fn mcp_backend_call(
1061 State(state): State<AppState>,
1062 Path(backend): Path<String>,
1063 headers: HeaderMap,
1064 body: Bytes,
1065) -> Response {
1066 if state.backend_router.get(&backend).is_none() {
1067 return (
1068 StatusCode::NOT_FOUND,
1069 Json(openai_error_body(StatusCode::NOT_FOUND, "unknown MCP backend")),
1070 )
1071 .into_response();
1072 }
1073 mcp_entry(state, headers, body, crate::mcp::Scope::Backend(backend)).await
1074}
1075
1076async fn mcp_get(State(state): State<AppState>, headers: HeaderMap) -> Response {
1079 if let Err(response) = crate::mcp::check_origin(&state, &headers) {
1080 return *response;
1081 }
1082 let mut response = StatusCode::METHOD_NOT_ALLOWED.into_response();
1083 response.headers_mut().insert(
1084 axum::http::header::ALLOW,
1085 HeaderValue::from_static("POST, DELETE"),
1086 );
1087 response
1088}
1089
1090async fn mcp_delete(State(state): State<AppState>, headers: HeaderMap) -> Response {
1094 if let Err(response) = crate::mcp::check_origin(&state, &headers) {
1095 return *response;
1096 }
1097 let caller = match lifecycle_caller(&state, &headers, &state.config.session_close_scope).await {
1098 Ok(caller) => caller,
1099 Err(error) => return pipeline_error(error),
1100 };
1101 let mut headers = headers;
1102 let session = match crate::mcp::bind_session(&state, &mut headers, &caller) {
1103 Ok(Some(session)) => session,
1104 Ok(None) => {
1105 return (
1106 StatusCode::BAD_REQUEST,
1107 Json(openai_error_body(
1108 StatusCode::BAD_REQUEST,
1109 "MCP-Session-Id is required",
1110 )),
1111 )
1112 .into_response()
1113 }
1114 Err(response) => return *response,
1115 };
1116 state.mcp_backends.forget_agent_session(&session.agent_session);
1117 let Some(agent_session) = state.sessions.get(&session.agent_session) else {
1118 return StatusCode::NO_CONTENT.into_response();
1120 };
1121 if let Err(error) = caller.authorize_session(&state, &agent_session, &state.config.session_close_scope) {
1122 return pipeline_error(error);
1123 }
1124 match state
1125 .finalizer
1126 .close_session(agent_session, StopReason::SessionClosed)
1127 .await
1128 {
1129 Ok(outcome) => Json(outcome).into_response(),
1130 Err(error) => finalize_error_response(&error),
1131 }
1132}
1133
1134async fn mcp_entry(state: AppState, headers: HeaderMap, body: Bytes, scope: crate::mcp::Scope) -> Response {
1135 if let Some(err) = refuse_non_json_content_type(&headers, "MCP") {
1144 return err;
1145 }
1146 let admission_permit = match Arc::clone(&state.mcp_admission).try_acquire_owned() {
1156 Ok(permit) => permit,
1157 Err(_) => {
1158 state.mcp_admission_refusals_counter.inc();
1163 let mut response = (
1164 StatusCode::SERVICE_UNAVAILABLE,
1165 Json(openai_error_body(
1166 StatusCode::SERVICE_UNAVAILABLE,
1167 "MCP admission at capacity; retry shortly",
1168 )),
1169 )
1170 .into_response();
1171 response
1172 .headers_mut()
1173 .insert(axum::http::header::RETRY_AFTER, HeaderValue::from_static("1"));
1174 return response;
1175 }
1176 };
1177 let inflight_guard = state.mcp_inflight.enter();
1199 match tokio::spawn(async move {
1200 let _admission_permit = admission_permit;
1201 let _inflight_guard = inflight_guard;
1202 mcp_call_inner(state, headers, body, scope).await
1203 })
1204 .await
1205 {
1206 Ok(response) => response,
1207 Err(join_error) => lifecycle_error(format!("tool call task failed: {join_error}")),
1208 }
1209}
1210
1211async fn refund_tool_admission(
1214 state: &AppState,
1215 execution: &ToolExecution,
1216 payout_micros: u64,
1217 principal_id: Option<&str>,
1218) -> Result<(), String> {
1219 let store = Arc::clone(&state.store);
1220 let session_id = execution.session_id.clone();
1221 let tool = execution.tool.clone();
1222 let budget = state.effective_budget();
1223 let principal = principal_id
1224 .zip(state.config.principal_budget.as_ref())
1225 .map(|(id, spec)| (id.to_owned(), spec.clone()));
1226 tokio::task::spawn_blocking(move || {
1227 pv_state::ActionBudget::new(store.as_ref(), &session_id, &budget)
1228 .refund_tool_call(&tool, payout_micros);
1229 if let Some((principal_id, spec)) = principal {
1230 pv_state::ActionBudget::for_principal(store.as_ref(), &principal_id, &spec)
1231 .refund_tool_call(&tool, payout_micros);
1232 }
1233 })
1234 .await
1235 .map_err(|_| "tool admission refund task failed".to_owned())
1236}
1237
1238async fn mcp_call_inner(
1239 state: AppState,
1240 headers: HeaderMap,
1241 body: Bytes,
1242 scope: crate::mcp::Scope,
1243) -> Response {
1244 if let Err(response) = crate::mcp::check_origin(&state, &headers) {
1247 return *response;
1248 }
1249 if let Err(response) = crate::mcp::check_protocol_version(&headers) {
1250 return *response;
1251 }
1252 let mut headers = headers;
1253 let forwarding_configured = state.config.tool_upstream_url.is_some() || !state.backend_router.is_empty();
1257 let auth_state = state.clone();
1258 let auth_headers = headers.clone();
1259 let authenticated =
1260 match tokio::task::spawn_blocking(move || auth_state.authenticate(&auth_headers)).await {
1261 Ok(Ok(caller)) => caller,
1262 Ok(Err(error)) => {
1263 if let Err(audit_error) = state.audit_authentication_error(&headers, &error) {
1264 return pipeline_error(audit_error);
1265 }
1266 return pipeline_error(error);
1267 }
1268 Err(error) => return pipeline_error(crate::pipeline::PipelineError::unavailable_source(error)),
1269 };
1270 let mcp_session = match crate::mcp::bind_session(&state, &mut headers, &authenticated) {
1273 Ok(session) => session,
1274 Err(response) => return *response,
1275 };
1276 match crate::mcp::answer(
1277 &state,
1278 &headers,
1279 &authenticated,
1280 crate::mcp::classify(&body),
1281 &scope,
1282 mcp_session.as_ref(),
1283 )
1284 .await
1285 {
1286 Ok(Some(response)) => return response,
1287 Ok(None) => {}
1288 Err(error) => {
1289 if let Err(audit_error) = state.audit_authentication_error(&headers, &error) {
1290 return pipeline_error(audit_error);
1291 }
1292 return pipeline_error(error);
1293 }
1294 }
1295 let (execution, unaudited_outcome) = if forwarding_configured {
1296 if let Err(error) = state.refuse_unauthenticated_tool_call(&headers) {
1302 return pipeline_error(error);
1303 }
1304 match ToolExecution::from_request(&state.config.atif_spool_dir, &headers, &body, state.journal_key) {
1305 Ok(mut execution) => {
1306 let required_scope = crate::pipeline::tool_scope(&execution.tool);
1307 let identity = match authenticated.authorize(&state, Some(&required_scope)) {
1308 Ok(identity) => identity,
1309 Err(error) => {
1310 if let Err(audit_error) = state.audit_authentication_error(&headers, &error) {
1311 return pipeline_error(audit_error);
1312 }
1313 return pipeline_error(error);
1314 }
1315 };
1316 if let Err(error) = execution.bind_principal(&identity) {
1317 return pipeline_error(error);
1318 }
1319 if let Some(principal) = authenticated.principal_digest() {
1320 execution.principal_digest = pv_core::digest::sha256_hex(
1321 format!("{}:{principal}", execution.principal_digest).as_bytes(),
1322 );
1323 }
1324 let authorized_session = |missing_context: &str| -> Result<(), Box<Response>> {
1331 let Some(session) = state.sessions.get(&execution.session_id) else {
1332 return Err(Box::new(pipeline_error(
1333 crate::pipeline::PipelineError::bad_request(format!(
1334 "unknown session for {missing_context}"
1335 )),
1336 )));
1337 };
1338 authenticated
1339 .authorize_session(&state, &session, &required_scope)
1340 .map_err(|error| Box::new(pipeline_error(error)))
1341 };
1342 match execution.load().await {
1343 Ok(ToolExecutionState::Completed(outcome)) => {
1344 if let Err(response) = authorized_session("cached tool result") {
1345 return *response;
1346 }
1347 return outcome.into_response();
1348 }
1349 Ok(ToolExecutionState::Unaudited(outcome)) => {
1350 if let Err(response) = authorized_session("pending tool audit") {
1351 return *response;
1352 }
1353 (Some(execution), Some(outcome))
1354 }
1355 Ok(ToolExecutionState::Pending) => {
1356 if let Err(response) = authorized_session("pending tool execution") {
1357 return *response;
1358 }
1359 return (
1360 StatusCode::CONFLICT,
1361 Json(json!({"error": TOOL_OUTCOME_UNCERTAIN})),
1362 )
1363 .into_response();
1364 }
1365 Ok(ToolExecutionState::Missing) => (Some(execution), None),
1366 Err(error) if error == TOOL_REQUEST_MISMATCH => {
1367 let Some(session) = state.sessions.get(&execution.session_id) else {
1376 return pipeline_error(crate::pipeline::PipelineError::bad_request(
1377 "unknown session for tool execution".to_owned(),
1378 ));
1379 };
1380 if let Err(error) = authenticated.authorize_session(
1381 &state,
1382 &session,
1383 &crate::pipeline::tool_scope(&execution.tool),
1384 ) {
1385 return pipeline_error(error);
1386 }
1387 return (
1388 StatusCode::CONFLICT,
1389 Json(json!({"error": TOOL_REQUEST_MISMATCH})),
1390 )
1391 .into_response();
1392 }
1393 Err(error) => return lifecycle_error(error),
1394 }
1395 }
1396 Err(error) => return pipeline_error(error),
1397 }
1398 } else {
1399 (None, None)
1400 };
1401 if let (Some(execution), Some(_)) = (execution.as_ref(), unaudited_outcome.as_ref()) {
1402 let _lease = match state.lease_session(&headers) {
1403 Ok(lease) => lease,
1404 Err(error) => return pipeline_error(error),
1405 };
1406 let gate = tool_audit_gate(&state, &execution.key);
1412 let _gate = gate.lock().await;
1413 let outcome = match execution.load().await {
1414 Ok(ToolExecutionState::Completed(outcome)) => return outcome.into_response(),
1415 Ok(ToolExecutionState::Unaudited(outcome)) => outcome,
1416 Ok(_) => return lifecycle_error("tool execution state changed during audit".to_owned()),
1417 Err(error) => return lifecycle_error(error),
1418 };
1419 let completion_permit = match state.worker.try_reserve(&execution.session_id) {
1420 Ok(permit) => permit,
1421 Err(error) => {
1422 return pipeline_error(crate::pipeline::PipelineError::unavailable_source(error));
1423 }
1424 };
1425 let Some(session) = state.sessions.get(&execution.session_id) else {
1426 return lifecycle_error("tool session disappeared".to_owned());
1427 };
1428 return complete_tool_audit(&state, execution, outcome, completion_permit, session).await;
1429 }
1430 let mut exchange_token = None;
1439 let mut preflight_denial = None;
1440 if let crate::mcp::Scope::Backend(name) = &scope {
1443 if let Ok(request) = pv_sandbox::parse_tool_call(&body) {
1444 let routed = state
1445 .backend_router
1446 .resolve(&request.tool)
1447 .map(|backend| backend.name.as_str());
1448 if routed != Some(name.as_str()) {
1449 let code = pv_sandbox::DenialCode::NoBackend;
1450 let reason = format!("tool {:?} is not served by backend {name:?}", request.tool);
1451 preflight_denial = Some(ToolVerdict::Blocked {
1452 tool: request.tool.clone(),
1453 stage: "policy",
1454 policy: Some("backend.route".into()),
1455 denial_code: code,
1456 response: pv_sandbox::rpc::denial_error(request.id.as_ref(), code, &reason),
1457 reason,
1458 elapsed_us: 0,
1459 });
1460 }
1461 }
1462 }
1463 if let (None, Some(authzen)) = (&preflight_denial, &state.authzen) {
1468 if let Ok(request) = pv_sandbox::parse_tool_call(&body) {
1469 let facts = authenticated.facts();
1470 let now_s = pv_core::time::now_ms() / 1000;
1471 if let crate::authz::AuthzDecision::Permit { intent } =
1472 state.pdp.decide_for(&request.tool, now_s, &facts)
1473 {
1474 let missions: Vec<String> = state
1475 .pdp
1476 .applicable_missions(&facts)
1477 .into_iter()
1478 .map(|mission| mission.id.clone())
1479 .collect();
1480 let backend = state
1481 .backend_router
1482 .resolve(&request.tool)
1483 .map(|backend| backend.name.as_str());
1484 let session = crate::pipeline::single_header(&headers, crate::pipeline::SESSION_HEADER)
1485 .ok()
1486 .flatten()
1487 .and_then(|value| value.to_str().ok());
1488 let evaluation = crate::authzen::evaluation_request(
1489 &authenticated,
1490 &request.tool,
1491 &intent,
1492 backend,
1493 session,
1494 &missions,
1495 );
1496 match authzen.evaluate(&evaluation).await {
1497 Ok(crate::authzen::AuthzenDecision::Permit) => {
1498 state
1499 .metrics
1500 .counter(
1501 "pv_authzen_decisions_total{decision=\"permit\"}",
1502 "External AuthZEN decisions by outcome",
1503 )
1504 .inc();
1505 }
1506 Ok(crate::authzen::AuthzenDecision::Deny(reason)) => {
1507 state
1508 .metrics
1509 .counter(
1510 "pv_authzen_decisions_total{decision=\"deny\"}",
1511 "External AuthZEN decisions by outcome",
1512 )
1513 .inc();
1514 let code = pv_sandbox::DenialCode::PdpDenied;
1515 preflight_denial = Some(ToolVerdict::Blocked {
1516 tool: request.tool.clone(),
1517 stage: "policy",
1518 policy: Some("pdp.authzen".into()),
1519 denial_code: code,
1520 response: pv_sandbox::rpc::denial_error(request.id.as_ref(), code, &reason),
1521 reason,
1522 elapsed_us: 0,
1523 });
1524 }
1525 Err(error) => {
1526 state
1527 .metrics
1528 .counter(
1529 "pv_authzen_decisions_total{decision=\"error\"}",
1530 "External AuthZEN decisions by outcome",
1531 )
1532 .inc();
1533 tracing::warn!(%error, "external policy decision point unavailable; refusing the call");
1534 return pipeline_error(crate::pipeline::PipelineError::unavailable(
1535 "the external policy decision point is unavailable; retry shortly".to_owned(),
1536 ));
1537 }
1538 }
1539 }
1540 }
1541 }
1542 if let (Some(execution), None) = (&execution, &preflight_denial) {
1543 if let Some(backend) = state
1544 .backend_router
1545 .resolve(&execution.tool)
1546 .filter(|backend| backend.auth_mode == crate::backend::BackendAuthMode::Exchange)
1547 {
1548 match exchange_for_backend(&state, &authenticated, &backend.name, &execution.tool) {
1549 Ok(token) => exchange_token = Some(token),
1550 Err(crate::pipeline::PipelineError::Blocked { context: reason, .. }) => {
1551 let request = pv_sandbox::parse_tool_call(&body).ok();
1552 let code = pv_sandbox::DenialCode::PolicyDenied;
1553 preflight_denial = Some(ToolVerdict::Blocked {
1554 tool: execution.tool.clone(),
1555 stage: "policy",
1556 policy: Some("backend.exchange".into()),
1557 denial_code: code,
1558 response: pv_sandbox::rpc::denial_error(
1559 request.as_ref().and_then(|r| r.id.as_ref()),
1560 code,
1561 &reason,
1562 ),
1563 reason,
1564 elapsed_us: 0,
1565 });
1566 }
1567 Err(error) => return pipeline_error(error),
1568 }
1569 }
1570 }
1571 let intent_caller = authenticated.clone();
1573 match state
1574 .intercept_authenticated_nonblocking(&headers, &body, authenticated, preflight_denial)
1575 .await
1576 {
1577 Ok(ToolVerdict::Allowed {
1578 tool,
1579 budget_remaining,
1580 elapsed_us,
1581 payout_micros,
1582 principal_id,
1583 }) => {
1584 let resolved_backend = state.backend_router.resolve(&tool);
1585 if let Some(url) = resolved_backend
1586 .map(|b| b.url.as_str())
1587 .or(state.config.tool_upstream_url.as_deref())
1588 {
1589 let execution = match execution {
1590 Some(execution) => execution,
1591 None => return lifecycle_error("tool execution state is missing".to_owned()),
1592 };
1593 let refund_admission =
1602 || refund_tool_admission(&state, &execution, payout_micros, principal_id.as_deref());
1603 let now_s = pv_core::time::now_ms() / 1000;
1604 let intent = state.pdp.intent_for(&tool);
1608 let audience =
1609 resolved_backend.map_or(crate::backend::DEFAULT_BACKEND_NAME, |b| b.name.as_str());
1610 let intent_token =
1611 state
1612 .pdp
1613 .mint_intent_token(&tool, &intent, &intent_caller.facts(), audience, now_s);
1614 let credential: Option<(axum::http::HeaderName, HeaderValue)> =
1618 if let Some(token) = &exchange_token {
1619 crate::mcp::bearer_header(token)
1620 } else if let Some(backend) = resolved_backend {
1621 if backend.auth_mode == crate::backend::BackendAuthMode::ForwardToken {
1622 let Some(header) = crate::mcp::caller_bearer(&headers) else {
1623 if let Err(error) = refund_admission().await {
1624 return lifecycle_error(error);
1625 }
1626 return pipeline_error(crate::pipeline::PipelineError::Unauthorized(
1627 "this backend requires the caller's bearer token".to_owned(),
1628 ));
1629 };
1630 Some(header)
1631 } else {
1632 backend.auth_header.clone()
1633 }
1634 } else {
1635 state
1636 .tool_auth
1637 .clone()
1638 .map(|bearer| (axum::http::header::AUTHORIZATION, bearer))
1639 };
1640 let intent_header = intent_token
1641 .as_deref()
1642 .and_then(|token| HeaderValue::from_str(token).ok())
1643 .map(|value| (axum::http::HeaderName::from_static("x-pv-intent-token"), value));
1644 let _lease = match state.lease_session(&headers) {
1645 Ok(lease) => lease,
1646 Err(error) => {
1647 if let Err(error) = refund_admission().await {
1648 return lifecycle_error(error);
1649 }
1650 return pipeline_error(error);
1651 }
1652 };
1653 let completion_permit = match state.worker.try_reserve(&execution.session_id) {
1654 Ok(permit) => permit,
1655 Err(error) => {
1656 if let Err(error) = refund_admission().await {
1657 return lifecycle_error(error);
1658 }
1659 return pipeline_error(crate::pipeline::PipelineError::unavailable(
1660 error.to_string(),
1661 ));
1662 }
1663 };
1664 match execution.claim().await {
1665 Ok(()) => {}
1666 Err(ClaimError::Race) => {
1667 tracing::warn!(
1671 session = %execution.session_id,
1672 "concurrent tool execution claim lost; refunding budget"
1673 );
1674 if let Err(error) = refund_admission().await {
1678 return lifecycle_error(error);
1679 }
1680 return (
1681 StatusCode::CONFLICT,
1682 Json(json!({"error": TOOL_OUTCOME_UNCERTAIN})),
1683 )
1684 .into_response();
1685 }
1686 Err(ClaimError::Backend(error)) => {
1687 tracing::warn!(
1693 session = %execution.session_id,
1694 %error,
1695 "tool execution claim failed; refunding budget"
1696 );
1697 if let Err(error) = refund_admission().await {
1698 return lifecycle_error(error);
1699 }
1700 return pipeline_error(crate::pipeline::PipelineError::unavailable(
1701 "tool execution intent could not be persisted".to_owned(),
1702 ));
1703 }
1704 }
1705 let mark_incomplete = || {
1710 if let Some(session) = state.sessions.get(&execution.session_id) {
1711 session.mark_capture_failed();
1712 }
1713 };
1714 let tool_timeout = std::time::Duration::from_secs(
1718 state
1719 .config
1720 .mcp_request_timeout_s
1721 .unwrap_or(crate::config::DEFAULT_MCP_REQUEST_TIMEOUT_S),
1722 );
1723 if let Some(backend) = resolved_backend
1724 .filter(|backend| backend.transport == crate::config::BackendTransport::Mcp)
1725 {
1726 let request_id = serde_json::from_slice::<Value>(&body)
1730 .ok()
1731 .and_then(|request| request.get("id").cloned())
1732 .unwrap_or(Value::Null);
1733 let endpoint = crate::mcp_client::BackendEndpoint {
1734 url: url.to_owned(),
1735 auth: credential,
1736 extra: intent_header.into_iter().collect(),
1737 deadline: tokio::time::Instant::now() + tool_timeout,
1738 };
1739 let exchanged = state
1740 .mcp_backends
1741 .request(
1742 &execution.session_id,
1743 &backend.name,
1744 &endpoint,
1745 &body,
1746 &request_id,
1747 )
1748 .await;
1749 let outcome = match exchanged {
1750 Ok((status, bytes, content_type)) => ToolOutcome {
1751 status,
1752 body_hex: hex::encode(&bytes),
1753 content_type,
1754 },
1755 Err(crate::mcp_client::McpError::NotDelivered(category)) => {
1756 match execution.release_unexecuted().await {
1760 Ok(()) => {
1761 if let Err(error) = refund_admission().await {
1762 return lifecycle_error(error);
1763 }
1764 }
1765 Err(release_error) => {
1766 mark_incomplete();
1767 tracing::warn!(
1768 session = %execution.session_id,
1769 error = %release_error,
1770 "failed to release undelivered MCP tool intent; charge kept, key stays claimed"
1771 );
1772 }
1773 }
1774 return pipeline_error(crate::pipeline::PipelineError::upstream(format!(
1775 "forward tool call: {category}"
1776 )));
1777 }
1778 Err(crate::mcp_client::McpError::AfterDelivery(error)) => ToolOutcome {
1779 status: StatusCode::BAD_GATEWAY.as_u16(),
1780 body_hex: hex::encode(
1781 serde_json::to_vec(&json!({
1782 "error": format!("tool executed but its response could not be read: {error}"),
1783 }))
1784 .unwrap_or_default(),
1785 ),
1786 content_type: Some("application/json".to_owned()),
1787 },
1788 };
1789 let gate = tool_audit_gate(&state, &execution.key);
1790 let _gate = gate.lock().await;
1791 if let Err(error) = execution.persist(&outcome).await {
1792 mark_incomplete();
1793 return lifecycle_error(error);
1794 }
1795 let Some(session) = state.sessions.get(&execution.session_id) else {
1796 return lifecycle_error("tool session disappeared".to_owned());
1797 };
1798 return complete_tool_audit(&state, &execution, outcome, completion_permit, session)
1799 .await;
1800 }
1801 let mut tool_request = state
1802 .backend_client
1803 .post(url)
1804 .timeout(tool_timeout)
1805 .header(axum::http::header::CONTENT_TYPE, "application/json")
1806 .body(body);
1807 if let Some((name, value)) = &credential {
1808 tool_request = tool_request.header(name.clone(), value.clone());
1809 }
1810 if let Some((name, value)) = &intent_header {
1811 tool_request = tool_request.header(name.clone(), value.clone());
1812 }
1813 match tool_request.send().await {
1814 Ok(upstream) => {
1815 let status = upstream.status();
1816 let upstream_headers = upstream.headers();
1827 let mut refused_encoding: Option<(String, String)> = None;
1828 'tool_encodings: for value in
1829 upstream_headers.get_all(axum::http::header::CONTENT_ENCODING)
1830 {
1831 let raw = value.to_str().unwrap_or_default();
1832 for token in raw.split(',') {
1833 let token = token.trim();
1834 if token.is_empty() || token.eq_ignore_ascii_case("identity") {
1835 continue;
1836 }
1837 refused_encoding = Some((token.to_owned(), raw.to_owned()));
1838 break 'tool_encodings;
1839 }
1840 }
1841 if let Some((token, raw)) = refused_encoding {
1842 let failure = ToolOutcome {
1843 status: StatusCode::BAD_GATEWAY.as_u16(),
1844 body_hex: hex::encode(
1845 serde_json::to_vec(&json!({
1846 "error": format!(
1847 "tool upstream responded with unsupported Content-Encoding token {token:?} (full header: {raw:?})",
1848 ),
1849 }))
1850 .unwrap_or_default(),
1851 ),
1852 content_type: Some("application/json".to_owned()),
1853 };
1854 let gate = tool_audit_gate(&state, &execution.key);
1855 let _gate = gate.lock().await;
1856 if let Err(error) = execution.persist(&failure).await {
1857 mark_incomplete();
1858 return lifecycle_error(error);
1859 }
1860 let session = match state.sessions.get(&execution.session_id) {
1861 Some(session) => session,
1862 None => return lifecycle_error("tool session disappeared".to_owned()),
1863 };
1864 return complete_tool_audit(
1865 &state,
1866 &execution,
1867 failure,
1868 completion_permit,
1869 session,
1870 )
1871 .await;
1872 }
1873 match read_limited_tool_response(upstream).await {
1874 Ok((bytes, content_type)) => {
1875 let outcome = ToolOutcome {
1876 status: status.as_u16(),
1877 body_hex: hex::encode(&bytes),
1878 content_type,
1879 };
1880 let gate = tool_audit_gate(&state, &execution.key);
1886 let _gate = gate.lock().await;
1887 if let Err(error) = execution.persist(&outcome).await {
1888 mark_incomplete();
1889 return lifecycle_error(error);
1890 }
1891 let session = match state.sessions.get(&execution.session_id) {
1892 Some(session) => session,
1893 None => return lifecycle_error("tool session disappeared".to_owned()),
1894 };
1895 complete_tool_audit(&state, &execution, outcome, completion_permit, session)
1896 .await
1897 }
1898 Err(error) => {
1899 let failure = ToolOutcome {
1911 status: StatusCode::BAD_GATEWAY.as_u16(),
1912 body_hex: hex::encode(
1913 serde_json::to_vec(&json!({
1914 "error": format!("tool executed but its response could not be read: {error}"),
1915 }))
1916 .unwrap_or_default(),
1917 ),
1918 content_type: Some("application/json".to_owned()),
1919 };
1920 let gate = tool_audit_gate(&state, &execution.key);
1921 let _gate = gate.lock().await;
1922 if execution.persist(&failure).await.is_ok() {
1923 if let Some(session) = state.sessions.get(&execution.session_id) {
1924 return complete_tool_audit(
1925 &state,
1926 &execution,
1927 failure,
1928 completion_permit,
1929 session,
1930 )
1931 .await;
1932 }
1933 }
1934 mark_incomplete();
1935 pipeline_error(crate::pipeline::PipelineError::upstream(format!(
1936 "read tool response: {error}"
1937 )))
1938 }
1939 }
1940 }
1941 Err(error) => {
1942 let category = crate::pipeline::classify_upstream_error(&error);
1957 tracing::warn!(
1958 category = category,
1959 error.status = ?error.status(),
1960 error.is_timeout = error.is_timeout(),
1961 error.is_connect = error.is_connect(),
1962 error.is_request = error.is_request(),
1963 "tool upstream forwarding failed"
1964 );
1965 if error.is_connect() {
1976 match execution.release_unexecuted().await {
1977 Ok(()) => {
1978 if let Err(error) = refund_admission().await {
1979 return lifecycle_error(error);
1980 }
1981 }
1982 Err(release_error) => {
1983 mark_incomplete();
1984 tracing::warn!(
1985 session = %execution.session_id,
1986 error = %release_error,
1987 "failed to release unexecuted tool intent; charge kept, key stays claimed"
1988 );
1989 }
1990 }
1991 } else {
1992 mark_incomplete();
1995 }
1996 pipeline_error(crate::pipeline::PipelineError::upstream(format!(
2000 "forward tool call: {category}"
2001 )))
2002 }
2003 }
2004 } else {
2005 let request_id = serde_json::from_slice::<Value>(&body)
2011 .ok()
2012 .and_then(|request| request.get("id").cloned())
2013 .unwrap_or(Value::Null);
2014 (
2015 StatusCode::OK,
2016 Json(json!({
2017 "jsonrpc": "2.0",
2018 "id": request_id,
2019 "result": {
2020 "allowed": true,
2021 "tool": tool,
2022 "budget_remaining": budget_remaining,
2023 "decision_us": elapsed_us,
2024 },
2025 })),
2026 )
2027 .into_response()
2028 }
2029 }
2030 Ok(ToolVerdict::Blocked { response, stage, .. }) => {
2031 let status = if stage == "parse" {
2035 StatusCode::BAD_REQUEST
2036 } else {
2037 StatusCode::FORBIDDEN
2038 };
2039 (status, Json(response)).into_response()
2040 }
2041 Err(error) => pipeline_error(error),
2042 Ok(_) => {
2050 tracing::error!(
2051 "unhandled ToolVerdict variant reached routes::tool_forward; failing closed. \
2052 This is a bug — extend the match to handle the new variant."
2053 );
2054 pipeline_error(crate::pipeline::PipelineError::unavailable(
2055 "tool sandbox produced an unhandled verdict; refusing to proceed".to_owned(),
2056 ))
2057 }
2058 }
2059}
2060
2061const MAX_TOOL_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
2062
2063async fn read_limited_tool_response(response: reqwest::Response) -> Result<(Bytes, Option<String>), String> {
2072 let content_type = response
2073 .headers()
2074 .get(axum::http::header::CONTENT_TYPE)
2075 .and_then(|value| value.to_str().ok())
2076 .map(str::to_owned);
2077 let mut stream = response.bytes_stream();
2078 let mut body = Vec::new();
2079 while let Some(chunk) = stream.next().await {
2080 let chunk = chunk.map_err(|error| {
2086 let category = crate::pipeline::classify_upstream_error(&error);
2087 tracing::warn!(
2088 category = category,
2089 error.status = ?error.status(),
2090 error.is_timeout = error.is_timeout(),
2091 error.is_connect = error.is_connect(),
2092 error.is_body = error.is_body(),
2093 "tool upstream stream chunk failed"
2094 );
2095 category.to_owned()
2096 })?;
2097 let next = body
2098 .len()
2099 .checked_add(chunk.len())
2100 .ok_or_else(|| "tool response size overflow".to_owned())?;
2101 if next > MAX_TOOL_RESPONSE_BYTES {
2102 return Err(format!("tool response exceeds {MAX_TOOL_RESPONSE_BYTES} bytes"));
2103 }
2104 body.extend_from_slice(&chunk);
2105 }
2106 Ok((Bytes::from(body), content_type))
2107}
2108
2109fn tool_response(status: StatusCode, bytes: Bytes, content_type: Option<&str>) -> Response {
2114 let mut response = (status, bytes).into_response();
2115 let value = content_type
2116 .and_then(|ct| HeaderValue::from_str(ct).ok())
2117 .unwrap_or_else(|| HeaderValue::from_static("application/json"));
2118 response
2119 .headers_mut()
2120 .insert(axum::http::header::CONTENT_TYPE, value);
2121 response.headers_mut().insert(
2127 axum::http::HeaderName::from_static("x-content-type-options"),
2128 HeaderValue::from_static("nosniff"),
2129 );
2130 response
2131}
2132
2133async fn complete_tool_audit(
2134 state: &AppState,
2135 execution: &ToolExecution,
2136 outcome: ToolOutcome,
2137 completion_permit: crate::worker::WorkerPermit,
2138 session: Arc<crate::session::Session>,
2139) -> Response {
2140 let status = StatusCode::from_u16(outcome.status).unwrap_or(StatusCode::BAD_GATEWAY);
2141 let bytes = match decode_cached_tool_body(&outcome.body_hex) {
2142 Ok(bytes) => bytes,
2143 Err(response) => return *response,
2144 };
2145 let success = status.is_success();
2146 let already_emitted = state.tool_audits_emitted.lock().contains(&execution.key);
2151 if !already_emitted {
2152 completion_permit.submit(crate::worker::WorkerJob {
2153 session: Arc::clone(&session),
2154 identity: session.current_identity(),
2155 class: pv_events::EventClass::Session,
2156 payload: json!({
2157 "action": "tool_completed",
2158 "execution_key": &execution.key,
2159 "status": status.as_u16(),
2160 "response_sha256": pv_core::digest::sha256_hex(&bytes),
2161 }),
2162 text: String::new(),
2163 analyze_loop: false,
2164 status: if success {
2165 pv_events::StatusId::Success
2166 } else {
2167 pv_events::StatusId::Failure
2168 },
2169 stop_reason: (!success).then_some(StopReason::Other),
2170 native_stop_reason: None,
2171 metrics: pv_events::EventMetrics::default(),
2172 cost_usd_micros: 0,
2173 prompt_token_correction: 0,
2174 atif: Some(crate::worker::AtifCapture {
2175 source: pv_atif::Source::System,
2176 message: Value::String(String::from_utf8_lossy(&bytes).into_owned()),
2177 reasoning_content: None,
2178 model_name: None,
2179 tool_calls: None,
2180 observation: None,
2181 llm_call_count: None,
2182 }),
2183 response_marker: None,
2184 response_attempt: None,
2185 });
2186 session.wait_for_worker_jobs().await;
2187 if session.capture_failed() {
2188 return lifecycle_error("tool completed but completion audit failed".to_owned());
2189 }
2190 state.tool_audits_emitted.lock().insert(execution.key.clone());
2193 }
2194 if let Err(error) = execution.mark_audited().await {
2195 return lifecycle_error(error);
2196 }
2197 state.tool_audits_emitted.lock().remove(&execution.key);
2198 tool_response(status, bytes, outcome.content_type.as_deref())
2203}
2204
2205#[derive(Clone, serde::Serialize, serde::Deserialize)]
2206struct ToolOutcome {
2207 status: u16,
2208 body_hex: String,
2209 #[serde(default)]
2216 content_type: Option<String>,
2217}
2218
2219impl ToolOutcome {
2220 fn into_response(self) -> Response {
2221 let status = StatusCode::from_u16(self.status).unwrap_or(StatusCode::BAD_GATEWAY);
2222 match decode_cached_tool_body(&self.body_hex) {
2223 Ok(body) => tool_response(status, body, self.content_type.as_deref()),
2224 Err(response) => *response,
2225 }
2226 }
2227}
2228
2229fn decode_cached_tool_body(body_hex: &str) -> Result<Bytes, Box<Response>> {
2251 match hex::decode(body_hex) {
2252 Ok(bytes) => Ok(Bytes::from(bytes)),
2253 Err(error) => {
2254 tracing::error!(
2255 %error,
2256 "cached tool outcome body_hex failed to decode; \
2257 possible control-key compromise or on-disk MAC bypass"
2258 );
2259 Err(Box::new(lifecycle_error(
2260 "cached tool outcome is corrupt".to_owned(),
2261 )))
2262 }
2263 }
2264}
2265
2266enum ToolExecutionState {
2267 Missing,
2268 Pending,
2269 Unaudited(ToolOutcome),
2270 Completed(ToolOutcome),
2271}
2272
2273#[derive(Debug)]
2282enum ClaimError {
2283 Race,
2288 Backend(String),
2292}
2293
2294const TOOL_REQUEST_MISMATCH: &str = "JSON-RPC id is already bound to a different tool request or principal";
2295const TOOL_OUTCOME_UNCERTAIN: &str = "tool execution outcome is uncertain; refusing duplicate execution";
2299
2300pub(crate) async fn ensure_tool_executions_resolved(
2303 spool: &std::path::Path,
2304 session_id: &str,
2305 control_key: &[u8; 32],
2306) -> Result<(), String> {
2307 let directory = spool.join(crate::spool::TOOL_EXECUTIONS);
2308 let session_id = session_id.to_owned();
2309 let control_key = *control_key;
2310 tokio::task::spawn_blocking(move || {
2311 ensure_tool_executions_resolved_with_limit(
2312 &directory,
2313 &session_id,
2314 &control_key,
2315 crate::reconciler::MAX_RECOVERY_DIRENTS_PER_TICK,
2316 )
2317 })
2318 .await
2319 .map_err(|_| "tool evidence lookup task failed".to_owned())?
2320}
2321
2322fn ensure_tool_executions_resolved_with_limit(
2323 directory: &std::path::Path,
2324 session_id: &str,
2325 control_key: &[u8; 32],
2326 max_dirents: usize,
2327) -> Result<(), String> {
2328 let entries = match std::fs::read_dir(directory) {
2329 Ok(entries) => entries,
2330 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()),
2331 Err(_) => return Err("tool evidence directory is unavailable".to_owned()),
2332 };
2333 for (index, entry) in entries.enumerate() {
2334 if index >= max_dirents {
2335 return Err("tool evidence directory exceeds the bounded scan limit".to_owned());
2336 }
2337 let path = entry
2338 .map_err(|_| "tool evidence entry is unavailable".to_owned())?
2339 .path();
2340 let Some(name) = path.file_name().and_then(std::ffi::OsStr::to_str) else {
2341 continue;
2342 };
2343 let Some(key) = name.strip_suffix(crate::spool::TOOL_INTENT_SUFFIX) else {
2344 continue;
2345 };
2346 let bytes = pv_core::fsutil::read_capped(&path, pv_core::fsutil::MAX_CONTROL_BYTES)
2347 .map_err(|_| "tool execution intent is unavailable".to_owned())?;
2348 let intent: ToolIntent = crate::journal::open(
2349 control_key,
2350 &format!("{}:{key}", crate::journal::TOOL_INTENT_DOMAIN),
2351 0,
2352 &bytes,
2353 )
2354 .map_err(|_| "tool execution intent could not be authenticated".to_owned())?;
2355 if intent.execution_key != key {
2356 return Err("tool intent path does not match its authenticated key".to_owned());
2357 }
2358 if intent.session_id == session_id && !tool_completion_is_authenticated(directory, key, control_key) {
2359 return Err("tool execution outcome or completion audit is unresolved".to_owned());
2360 }
2361 }
2362 Ok(())
2363}
2364
2365pub(crate) fn tool_completion_is_authenticated(
2369 directory: &std::path::Path,
2370 key: &str,
2371 control_key: &[u8; 32],
2372) -> bool {
2373 let outcome_path = directory.join(format!("{key}{}", crate::spool::TOOL_OUTCOME_SUFFIX));
2374 let Ok(outcome_bytes) = pv_core::fsutil::read_capped(&outcome_path, pv_core::fsutil::MAX_ATIF_BYTES)
2375 else {
2376 return false;
2377 };
2378 if crate::journal::open::<ToolOutcome>(
2379 control_key,
2380 &format!("{}:{key}", crate::journal::TOOL_OUTCOME_DOMAIN),
2381 0,
2382 &outcome_bytes,
2383 )
2384 .is_err()
2385 {
2386 return false;
2387 }
2388 let audited_path = directory.join(format!("{key}{}", crate::spool::TOOL_AUDITED_SUFFIX));
2389 let Ok(audited_bytes) = pv_core::fsutil::read_capped(&audited_path, pv_core::fsutil::MAX_CONTROL_BYTES)
2390 else {
2391 return false;
2392 };
2393 crate::journal::open::<serde_json::Value>(
2394 control_key,
2395 &format!("{}:{key}", crate::journal::TOOL_AUDITED_DOMAIN),
2396 0,
2397 &audited_bytes,
2398 )
2399 .is_ok()
2400}
2401
2402#[derive(PartialEq, serde::Serialize, serde::Deserialize)]
2403struct ToolIntent {
2404 pub(crate) execution_key: String,
2405 pub(crate) session_id: String,
2406 pub(crate) tool: String,
2407 pub(crate) request_digest: String,
2408 pub(crate) principal_digest: String,
2409}
2410
2411#[cfg(test)]
2412pub(crate) async fn ensure_no_unresolved_tool_executions(
2413 spool: &std::path::Path,
2414 control_key: &[u8; 32],
2415) -> Result<(), String> {
2416 let sessions = unresolved_tool_sessions(spool, control_key).await?;
2417 if let Some(session_id) = sessions.into_iter().next() {
2418 Err(format!("session {session_id} has an unresolved tool execution"))
2419 } else {
2420 Ok(())
2421 }
2422}
2423
2424pub(crate) async fn unresolved_tool_sessions(
2425 spool: &std::path::Path,
2426 control_key: &[u8; 32],
2427) -> Result<std::collections::HashSet<String>, String> {
2428 let directory = spool.join(crate::spool::TOOL_EXECUTIONS);
2429 let control_key = *control_key;
2430 tokio::task::spawn_blocking(move || {
2431 let entries = match std::fs::read_dir(&directory) {
2432 Ok(entries) => entries,
2433 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
2434 return Ok(std::collections::HashSet::new());
2435 }
2436 Err(error) => return Err(error.to_string()),
2437 };
2438 let mut intent_keys = std::collections::HashSet::new();
2439 let mut unresolved_sessions = std::collections::HashSet::new();
2440 let mut dirents_seen = 0usize;
2441 for entry in entries {
2442 dirents_seen = dirents_seen.saturating_add(1);
2455 if dirents_seen > crate::reconciler::MAX_RECOVERY_DIRENTS_PER_TICK {
2456 return Err(format!(
2457 "tool-executions scan aborted: more than {} directory entries; \
2458 the spool directory is poisoned or leaking files",
2459 crate::reconciler::MAX_RECOVERY_DIRENTS_PER_TICK
2460 ));
2461 }
2462 let path = entry.map_err(|error| error.to_string())?.path();
2463 let Some(name) = path.file_name().and_then(std::ffi::OsStr::to_str) else {
2464 continue;
2465 };
2466 let Some(key) = name.strip_suffix(crate::spool::TOOL_INTENT_SUFFIX) else {
2467 continue;
2468 };
2469 intent_keys.insert(key.to_owned());
2470 let intent_check: Result<ToolIntent, String> =
2487 match pv_core::fsutil::read_capped(&path, pv_core::fsutil::MAX_CONTROL_BYTES) {
2488 Ok(intent_bytes) => crate::journal::open(
2489 &control_key,
2490 &format!("{}:{key}", crate::journal::TOOL_INTENT_DOMAIN),
2491 0,
2492 &intent_bytes,
2493 ),
2494 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
2497 intent_keys.remove(key);
2498 continue;
2499 }
2500 Err(error)
2501 if matches!(
2502 error.kind(),
2503 std::io::ErrorKind::InvalidData | std::io::ErrorKind::InvalidInput
2504 ) =>
2505 {
2506 Err(error.to_string())
2507 }
2508 Err(error) => return Err(error.to_string()),
2509 };
2510 let intent: ToolIntent = match intent_check {
2511 Ok(intent) => intent,
2512 Err(error) => {
2513 let file_name = path
2524 .file_stem()
2525 .and_then(std::ffi::OsStr::to_str)
2526 .and_then(|stem| stem.strip_suffix(".intent"))
2527 .map(|key| format!("{key}.intent.torn"));
2528 let quarantine = match file_name {
2529 Some(name) => path.with_file_name(name),
2530 None => path.with_extension("intent.torn"),
2531 };
2532 tracing::warn!(
2533 %error,
2534 original = %pv_core::fsutil::basename(&path),
2535 quarantine = %pv_core::fsutil::basename(&quarantine),
2536 "torn tool-execution intent quarantined so recovery can proceed"
2537 );
2538 if let Err(rename_err) = std::fs::rename(&path, &quarantine) {
2539 tracing::warn!(%rename_err, "failed to quarantine torn intent — leaving in place");
2540 }
2541 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_OUTCOME_SUFFIX, "outcome");
2548 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_AUDITED_SUFFIX, "audited");
2549 intent_keys.remove(key);
2550 continue;
2551 }
2552 };
2553 if intent.execution_key != key {
2554 return Err("tool intent path does not match authenticated execution key".to_owned());
2555 }
2556 let outcome_path = directory.join(format!("{key}{}", crate::spool::TOOL_OUTCOME_SUFFIX));
2557 let outcome_check: Result<ToolOutcome, String> =
2565 match pv_core::fsutil::read_capped(&outcome_path, pv_core::fsutil::MAX_ATIF_BYTES) {
2566 Ok(outcome_bytes) => crate::journal::open(
2567 &control_key,
2568 &format!("{}:{key}", crate::journal::TOOL_OUTCOME_DOMAIN),
2569 0,
2570 &outcome_bytes,
2571 ),
2572 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
2573 unresolved_sessions.insert(intent.session_id);
2574 continue;
2575 }
2576 Err(error)
2577 if matches!(
2578 error.kind(),
2579 std::io::ErrorKind::InvalidData | std::io::ErrorKind::InvalidInput
2580 ) =>
2581 {
2582 Err(error.to_string())
2583 }
2584 Err(error) => return Err(error.to_string()),
2585 };
2586 let _: ToolOutcome = match outcome_check {
2587 Ok(outcome) => outcome,
2588 Err(error) => {
2589 tracing::warn!(
2594 %error,
2595 key = %key,
2596 "torn tool-execution outcome quarantined so recovery can proceed"
2597 );
2598 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_OUTCOME_SUFFIX, "outcome");
2599 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_AUDITED_SUFFIX, "audited");
2600 unresolved_sessions.insert(intent.session_id);
2601 continue;
2602 }
2603 };
2604 let audited_path = directory.join(format!("{key}{}", crate::spool::TOOL_AUDITED_SUFFIX));
2605 let audited_check: Result<serde_json::Value, String> =
2606 match pv_core::fsutil::read_capped(&audited_path, pv_core::fsutil::MAX_CONTROL_BYTES) {
2607 Ok(audited_bytes) => crate::journal::open(
2608 &control_key,
2609 &format!("{}:{key}", crate::journal::TOOL_AUDITED_DOMAIN),
2610 0,
2611 &audited_bytes,
2612 ),
2613 Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
2614 unresolved_sessions.insert(intent.session_id);
2615 continue;
2616 }
2617 Err(error)
2618 if matches!(
2619 error.kind(),
2620 std::io::ErrorKind::InvalidData | std::io::ErrorKind::InvalidInput
2621 ) =>
2622 {
2623 Err(error.to_string())
2624 }
2625 Err(error) => return Err(error.to_string()),
2626 };
2627 let _: serde_json::Value = match audited_check {
2628 Ok(value) => value,
2629 Err(error) => {
2630 tracing::warn!(
2631 %error,
2632 key = %key,
2633 "torn tool-execution audited record quarantined so recovery can proceed"
2634 );
2635 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_AUDITED_SUFFIX, "audited");
2636 unresolved_sessions.insert(intent.session_id);
2637 continue;
2638 }
2639 };
2640 }
2641 let mut orphan_dirents = 0usize;
2642 for entry in std::fs::read_dir(&directory).map_err(|error| error.to_string())? {
2643 orphan_dirents = orphan_dirents.saturating_add(1);
2645 if orphan_dirents > crate::reconciler::MAX_RECOVERY_DIRENTS_PER_TICK {
2646 return Err(format!(
2647 "tool-executions orphan scan aborted: more than {} directory entries; \
2648 the spool directory is poisoned or leaking files",
2649 crate::reconciler::MAX_RECOVERY_DIRENTS_PER_TICK
2650 ));
2651 }
2652 let path = entry.map_err(|error| error.to_string())?.path();
2653 let Some(name) = path.file_name().and_then(std::ffi::OsStr::to_str) else {
2654 continue;
2655 };
2656 if let Some(key) = name.strip_suffix(crate::spool::TOOL_OUTCOME_SUFFIX) {
2657 if !intent_keys.contains(key) {
2658 tracing::warn!(
2665 key = %key,
2666 "orphan tool outcome without authenticated intent quarantined"
2667 );
2668 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_OUTCOME_SUFFIX, "outcome");
2669 quarantine_orphan_sibling(&directory, key, crate::spool::TOOL_AUDITED_SUFFIX, "audited");
2670 }
2671 }
2672 }
2673 Ok(unresolved_sessions)
2674 })
2675 .await
2676 .map_err(|error| error.to_string())?
2677}
2678
2679fn quarantine_orphan_sibling(directory: &std::path::Path, key: &str, suffix: &str, role: &str) {
2683 let path = directory.join(format!("{key}{suffix}"));
2684 if !path.exists() {
2685 return;
2686 }
2687 let uid = pv_core::new_event_uid();
2688 let quarantine = directory.join(format!("{key}{suffix}.corrupt-{uid}"));
2689 if let Err(rename_err) = std::fs::rename(&path, &quarantine) {
2690 tracing::warn!(
2691 %rename_err,
2692 role = %role,
2693 original = %pv_core::fsutil::basename(&path),
2694 "failed to quarantine orphan tool-execution sibling — leaving in place"
2695 );
2696 } else {
2697 tracing::warn!(
2698 role = %role,
2699 original = %pv_core::fsutil::basename(&path),
2700 quarantine = %pv_core::fsutil::basename(&quarantine),
2701 "orphan tool-execution sibling quarantined"
2702 );
2703 }
2704}
2705
2706#[derive(Clone)]
2707struct ToolExecution {
2708 key: String,
2709 session_id: String,
2710 tool: String,
2711 request_digest: String,
2712 principal_digest: String,
2713 control_key: [u8; 32],
2714 intent_path: std::path::PathBuf,
2715 outcome_path: std::path::PathBuf,
2716 audited_path: std::path::PathBuf,
2717}
2718
2719impl ToolExecution {
2720 fn from_request(
2721 spool: &str,
2722 headers: &HeaderMap,
2723 body: &[u8],
2724 control_key: [u8; 32],
2725 ) -> Result<Self, crate::pipeline::PipelineError> {
2726 let raw_header = crate::pipeline::single_header(headers, crate::pipeline::SESSION_HEADER)?;
2738 let session_id = match raw_header {
2739 None => {
2740 return Err(crate::pipeline::PipelineError::bad_request(
2741 "missing x-pv-session".to_owned(),
2742 ));
2743 }
2744 Some(value) => value.to_str().map_err(|_| {
2745 crate::pipeline::PipelineError::bad_request("X-PV-Session is not valid text".to_owned())
2746 })?,
2747 };
2748 let session_id = pv_core::SessionId::parse(session_id)
2751 .map(|id| id.to_string())
2752 .map_err(crate::pipeline::PipelineError::bad_request_source)?;
2753 let call =
2754 pv_sandbox::parse_tool_call(body).map_err(crate::pipeline::PipelineError::bad_request_source)?;
2755 let id = call.id.filter(|id| !id.is_null()).ok_or_else(|| {
2761 crate::pipeline::PipelineError::bad_request(
2762 "forwarded tool calls require a non-null JSON-RPC id for idempotency".to_owned(),
2763 )
2764 })?;
2765 let request: Value =
2766 serde_json::from_slice(body).map_err(crate::pipeline::PipelineError::bad_request_source)?;
2767 let canonical = pv_receipts::canonicalize(&request)
2768 .map_err(crate::pipeline::PipelineError::bad_request_source)?;
2769 let request_digest = pv_core::digest::sha256_hex(canonical.as_bytes());
2770 let key_material = format!("{session_id}:{}", id);
2771 let key = pv_core::digest::sha256_hex(key_material.as_bytes());
2772 let directory = std::path::Path::new(spool).join(crate::spool::TOOL_EXECUTIONS);
2773 Ok(Self {
2774 intent_path: directory.join(format!("{key}{}", crate::spool::TOOL_INTENT_SUFFIX)),
2775 outcome_path: directory.join(format!("{key}{}", crate::spool::TOOL_OUTCOME_SUFFIX)),
2776 audited_path: directory.join(format!("{key}{}", crate::spool::TOOL_AUDITED_SUFFIX)),
2777 key,
2778 session_id,
2779 tool: call.tool,
2780 request_digest,
2781 principal_digest: String::new(),
2782 control_key,
2783 })
2784 }
2785
2786 fn bind_principal(
2787 &mut self,
2788 identity: &pv_events::AgentIdentity,
2789 ) -> Result<(), crate::pipeline::PipelineError> {
2790 let stable_identity = json!({
2791 "version": identity.version,
2792 "charter": identity.charter,
2793 "instance_uid": identity.instance_uid,
2794 });
2795 let canonical = pv_receipts::canonicalize(&stable_identity)
2796 .map_err(crate::pipeline::PipelineError::bad_request_source)?;
2797 self.principal_digest = pv_core::digest::sha256_hex(canonical.as_bytes());
2798 Ok(())
2799 }
2800
2801 async fn load(&self) -> Result<ToolExecutionState, String> {
2802 let execution = self.clone();
2803 tokio::task::spawn_blocking(move || execution.load_sync())
2804 .await
2805 .map_err(|error| error.to_string())?
2806 }
2807
2808 fn load_sync(&self) -> Result<ToolExecutionState, String> {
2809 let intent_bytes = match std::fs::read(&self.intent_path) {
2815 Ok(bytes) => Some(bytes),
2816 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
2817 Err(error) => return Err(error.to_string()),
2818 };
2819 let Some(intent_bytes) = intent_bytes else {
2820 if self.outcome_path.exists() || self.audited_path.exists() {
2821 return Err("tool outcome exists without an authenticated intent".to_owned());
2822 }
2823 return Ok(ToolExecutionState::Missing);
2824 };
2825 let intent: ToolIntent = crate::journal::open(
2826 &self.control_key,
2827 &format!("{}:{}", crate::journal::TOOL_INTENT_DOMAIN, self.key),
2828 0,
2829 &intent_bytes,
2830 )?;
2831 if intent != self.intent() {
2832 return Err(TOOL_REQUEST_MISMATCH.to_owned());
2833 }
2834 let outcome_bytes = match std::fs::read(&self.outcome_path) {
2835 Ok(bytes) => Some(bytes),
2836 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
2837 Err(error) => return Err(error.to_string()),
2838 };
2839 if let Some(outcome_bytes) = outcome_bytes {
2840 let outcome: ToolOutcome = crate::journal::open(
2841 &self.control_key,
2842 &format!("{}:{}", crate::journal::TOOL_OUTCOME_DOMAIN, self.key),
2843 0,
2844 &outcome_bytes,
2845 )?;
2846 let audited = match std::fs::read(&self.audited_path) {
2847 Ok(bytes) => crate::journal::open::<serde_json::Value>(
2848 &self.control_key,
2849 &format!("{}:{}", crate::journal::TOOL_AUDITED_DOMAIN, self.key),
2850 0,
2851 &bytes,
2852 )
2853 .is_ok(),
2854 Err(_) => false,
2855 };
2856 return if audited {
2857 Ok(ToolExecutionState::Completed(outcome))
2858 } else {
2859 Ok(ToolExecutionState::Unaudited(outcome))
2860 };
2861 }
2862 Ok(ToolExecutionState::Pending)
2863 }
2864
2865 async fn claim(&self) -> Result<(), ClaimError> {
2866 let execution = self.clone();
2867 tokio::task::spawn_blocking(move || execution.claim_sync())
2868 .await
2869 .map_err(|error| ClaimError::Backend(error.to_string()))?
2870 }
2871
2872 async fn release_unexecuted(&self) -> Result<(), String> {
2887 let execution = self.clone();
2888 tokio::task::spawn_blocking(move || {
2889 match std::fs::remove_file(&execution.intent_path) {
2890 Ok(()) => {}
2891 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()),
2892 Err(error) => return Err(error.to_string()),
2893 }
2894 if let Some(directory) = execution.intent_path.parent() {
2895 if let Err(error) = pv_core::fsutil::sync_directory(directory) {
2896 tracing::warn!(
2897 path = %pv_core::fsutil::basename(&execution.intent_path),
2898 %error,
2899 "release_unexecuted: intent removed but directory fsync failed; treating as released"
2900 );
2901 }
2902 }
2903 Ok(())
2904 })
2905 .await
2906 .map_err(|error| error.to_string())?
2907 }
2908
2909 fn claim_sync(&self) -> Result<(), ClaimError> {
2910 use std::io::Write as _;
2911 let directory = self
2912 .intent_path
2913 .parent()
2914 .ok_or_else(|| ClaimError::Backend("tool execution directory is missing".to_owned()))?;
2915 if !directory.is_dir() {
2919 pv_core::fsutil::create_dir_all_synced(directory)
2924 .map_err(|error| ClaimError::Backend(error.to_string()))?;
2925 }
2926 let mut open_options = std::fs::OpenOptions::new();
2948 open_options.write(true).create_new(true);
2949 #[cfg(unix)]
2950 {
2951 use std::os::unix::fs::OpenOptionsExt as _;
2952 open_options.mode(0o600);
2953 }
2954 let mut file = open_options
2955 .open(&self.intent_path)
2956 .map_err(|error| match error.kind() {
2957 std::io::ErrorKind::AlreadyExists => ClaimError::Race,
2958 _ => ClaimError::Backend(format!("tool execution intent unavailable: {error}")),
2959 })?;
2960 let intent = crate::journal::seal(
2961 &self.control_key,
2962 &format!("{}:{}", crate::journal::TOOL_INTENT_DOMAIN, self.key),
2963 0,
2964 &self.intent(),
2965 )
2966 .map_err(ClaimError::Backend)?;
2967 file.write_all(&intent)
2968 .map_err(|error| ClaimError::Backend(error.to_string()))?;
2969 file.sync_all()
2970 .map_err(|error| ClaimError::Backend(error.to_string()))?;
2971 pv_core::fsutil::sync_directory(directory).map_err(|error| ClaimError::Backend(error.to_string()))
2972 }
2973
2974 async fn persist(&self, outcome: &ToolOutcome) -> Result<(), String> {
2975 let execution = self.clone();
2976 let outcome = outcome.clone();
2977 tokio::task::spawn_blocking(move || {
2978 let sealed = crate::journal::seal(
2979 &execution.control_key,
2980 &format!("{}:{}", crate::journal::TOOL_OUTCOME_DOMAIN, execution.key),
2981 0,
2982 &outcome,
2983 )?;
2984 write_atomic_bytes(&execution.outcome_path, &sealed)
2985 })
2986 .await
2987 .map_err(|error| error.to_string())?
2988 }
2989
2990 async fn mark_audited(&self) -> Result<(), String> {
2991 let execution = self.clone();
2992 tokio::task::spawn_blocking(move || {
2993 let sealed = crate::journal::seal(
2994 &execution.control_key,
2995 &format!("{}:{}", crate::journal::TOOL_AUDITED_DOMAIN, execution.key),
2996 0,
2997 &json!({"audited": true}),
2998 )?;
2999 write_atomic_bytes(&execution.audited_path, &sealed)
3000 })
3001 .await
3002 .map_err(|error| error.to_string())?
3003 }
3004
3005 fn intent(&self) -> ToolIntent {
3006 ToolIntent {
3007 execution_key: self.key.clone(),
3008 session_id: self.session_id.clone(),
3009 tool: self.tool.clone(),
3010 request_digest: self.request_digest.clone(),
3011 principal_digest: self.principal_digest.clone(),
3012 }
3013 }
3014}
3015
3016fn write_atomic_bytes(path: &std::path::Path, bytes: &[u8]) -> Result<(), String> {
3017 pv_core::fsutil::write_atomic(path, bytes).map_err(|error| error.to_string())
3018}
3019
3020async fn lifecycle_caller(
3021 state: &AppState,
3022 headers: &HeaderMap,
3023 scope: &str,
3024) -> Result<crate::pipeline::AuthenticatedCaller, crate::pipeline::PipelineError> {
3025 let state = state.clone();
3026 let headers = headers.clone();
3027 let scope = scope.to_owned();
3028 tokio::task::spawn_blocking(move || {
3029 let caller = state.authenticate(&headers)?;
3030 caller.authorize(&state, Some(&scope))?;
3031 Ok(caller)
3032 })
3033 .await
3034 .map_err(crate::pipeline::PipelineError::unavailable_source)?
3035}
3036
3037async fn close_session(
3038 State(state): State<AppState>,
3039 Path(id): Path<String>,
3040 headers: HeaderMap,
3041) -> Response {
3042 let caller = match lifecycle_caller(&state, &headers, &state.config.session_close_scope).await {
3043 Ok(caller) => caller,
3044 Err(error) => return pipeline_error(error),
3045 };
3046 let Some(session) = state.sessions.get(&id) else {
3047 return (
3048 StatusCode::NOT_FOUND,
3049 Json(openai_error_body(StatusCode::NOT_FOUND, "unknown session")),
3050 )
3051 .into_response();
3052 };
3053 if let Err(error) = caller.authorize_session(&state, &session, &state.config.session_close_scope) {
3054 return pipeline_error(error);
3055 }
3056 match state
3057 .finalizer
3058 .close_session(session, StopReason::SessionClosed)
3059 .await
3060 {
3061 Ok(outcome) => Json(outcome).into_response(),
3062 Err(error) => finalize_error_response(&error),
3063 }
3064}
3065
3066async fn promote_session(
3067 State(state): State<AppState>,
3068 Path(id): Path<String>,
3069 headers: HeaderMap,
3070) -> Response {
3071 let caller = match lifecycle_caller(&state, &headers, &state.config.session_promote_scope).await {
3072 Ok(caller) => caller,
3073 Err(error) => return pipeline_error(error),
3074 };
3075 let Some(session) = state.sessions.get(&id) else {
3076 return (
3077 StatusCode::NOT_FOUND,
3078 Json(openai_error_body(StatusCode::NOT_FOUND, "unknown session")),
3079 )
3080 .into_response();
3081 };
3082 if let Err(error) = caller.authorize_session(&state, &session, &state.config.session_promote_scope) {
3083 return pipeline_error(error);
3084 }
3085 if !session.is_closed() {
3094 if let Err(error) = caller.authorize_session(&state, &session, &state.config.session_close_scope) {
3095 return pipeline_error(error);
3096 }
3097 }
3098 match state.finalizer.promote(session).await {
3099 Ok(receipt) => Json(receipt).into_response(),
3100 Err(error) => finalize_error_response(&error),
3101 }
3102}
3103
3104async fn jwks(State(state): State<AppState>) -> Response {
3107 match &state.token_signer {
3108 Some(signer) => Json(signer.jwks_json()).into_response(),
3109 None => Json(serde_json::json!({"keys": []})).into_response(),
3110 }
3111}
3112
3113async fn token_exchange(State(state): State<AppState>, body: Bytes) -> Response {
3118 crate::token_routes::control_request(state, move |state| async move {
3119 tokio::task::spawn_blocking(move || token_exchange_inner(state, body))
3120 .await
3121 .unwrap_or_else(|_| oauth_unavailable())
3122 })
3123 .await
3124}
3125
3126fn token_exchange_inner(state: AppState, body: Bytes) -> Response {
3127 let grant_type = serde_urlencoded::from_bytes::<Vec<(String, String)>>(&body)
3130 .ok()
3131 .and_then(|pairs| pairs.into_iter().find(|(key, _)| key == "grant_type"))
3132 .map(|(_, value)| value);
3133 if grant_type.as_deref() == Some(pv_identity::JWT_BEARER_GRANT_TYPE) {
3134 return crate::workload::grant(&state, &body);
3135 }
3136 if !state.config.token_exchange_enabled {
3137 return StatusCode::NOT_FOUND.into_response();
3138 }
3139 let validator = match &state.identity {
3140 Some(v) => v,
3141 None => {
3142 return (
3143 StatusCode::INTERNAL_SERVER_ERROR,
3144 Json(oauth_error(
3145 "server_error",
3146 "token exchange requires an identity validator",
3147 )),
3148 )
3149 .into_response();
3150 }
3151 };
3152 let req: pv_identity::TokenExchangeRequest = match serde_urlencoded::from_bytes(&body) {
3153 Ok(r) => r,
3154 Err(e) => {
3155 return (
3156 StatusCode::BAD_REQUEST,
3157 Json(oauth_error(
3158 "invalid_request",
3159 &format!("invalid token exchange request: {e}"),
3160 )),
3161 )
3162 .into_response();
3163 }
3164 };
3165 if req.grant_type != pv_identity::TOKEN_EXCHANGE_GRANT_TYPE {
3166 return (
3167 StatusCode::BAD_REQUEST,
3168 Json(oauth_error(
3169 "unsupported_grant_type",
3170 &format!("unsupported grant_type: {:?}", req.grant_type),
3171 )),
3172 )
3173 .into_response();
3174 }
3175 if req.subject_token_type != pv_identity::exchange::JWT_TOKEN_TYPE {
3176 return (
3177 StatusCode::BAD_REQUEST,
3178 Json(oauth_error(
3179 "invalid_request",
3180 &format!("unsupported subject_token_type: {:?}", req.subject_token_type),
3181 )),
3182 )
3183 .into_response();
3184 }
3185 state
3186 .metrics
3187 .counter(
3188 "pv_identity_validations_total",
3189 "Full inbound identity validations",
3190 )
3191 .inc();
3192 let validated = match validator.validate(&req.subject_token) {
3193 Ok(v) => v,
3194 Err(pv_identity::IdentityError::RevocationUnavailable(detail)) => {
3195 tracing::error!(%detail, "identity revocation list unavailable; refusing token exchange");
3196 return oauth_unavailable();
3197 }
3198 Err(e) => {
3199 tracing::warn!(error = %e, "token exchange subject token rejected");
3200 return (
3203 StatusCode::BAD_REQUEST,
3204 Json(oauth_error(
3205 "invalid_request",
3206 crate::pipeline::classify_identity_error(&e),
3207 )),
3208 )
3209 .into_response();
3210 }
3211 };
3212 let now_s = pv_core::time::now_ms() / 1000;
3213 let ttl = state.config.token_exchange_ttl_s.min(900);
3214 let azp = &validated.claims.instance_uid;
3215 let allowed_audiences: Vec<String> = state.config.backends.iter().map(|b| b.name.clone()).collect();
3216 if allowed_audiences.is_empty() {
3220 return (
3221 StatusCode::BAD_REQUEST,
3222 Json(oauth_error(
3223 "invalid_target",
3224 "no [[backends]] are configured, so no exchange audience is permitted",
3225 )),
3226 )
3227 .into_response();
3228 }
3229 let actor = match (req.actor_token.as_deref(), req.actor_token_type.as_deref()) {
3233 (None, None) => None,
3234 (Some(token), Some(pv_identity::exchange::JWT_TOKEN_TYPE)) => match validator.validate(token) {
3235 Ok(actor) => Some(actor),
3236 Err(pv_identity::IdentityError::RevocationUnavailable(detail)) => {
3237 tracing::error!(%detail, "identity revocation list unavailable; refusing token exchange");
3238 return oauth_unavailable();
3239 }
3240 Err(error) => {
3241 tracing::warn!(%error, "token exchange actor token rejected");
3242 return (
3243 StatusCode::BAD_REQUEST,
3244 Json(oauth_error(
3245 "invalid_request",
3246 crate::pipeline::classify_identity_error(&error),
3247 )),
3248 )
3249 .into_response();
3250 }
3251 },
3252 (Some(_), Some(other)) => {
3253 return (
3254 StatusCode::BAD_REQUEST,
3255 Json(oauth_error(
3256 "invalid_request",
3257 &format!("unsupported actor_token_type: {other:?}"),
3258 )),
3259 )
3260 .into_response();
3261 }
3262 _ => {
3263 return (
3264 StatusCode::BAD_REQUEST,
3265 Json(oauth_error(
3266 "invalid_request",
3267 "actor_token and actor_token_type must be sent together",
3268 )),
3269 )
3270 .into_response();
3271 }
3272 };
3273 let params = pv_identity::ExchangeParams {
3274 subject: &validated.claims,
3275 subject_chain: &validated.chain_tokens,
3276 subject_delegation_depth: validated.total_delegation_depth,
3277 target_audience: &req.audience,
3278 requested_scopes: req.scope.as_deref(),
3279 azp,
3280 now_s,
3281 ttl_s: ttl,
3282 issuer: &state.config.audience,
3283 max_depth: state.config.max_delegation_depth,
3284 allowed_audiences: &allowed_audiences,
3285 };
3286 let actor = actor.as_ref().map(|actor| pv_identity::ExchangeActor {
3287 claims: &actor.claims,
3288 chain: &actor.chain_tokens,
3289 });
3290 let exchanged = match pv_identity::build_exchanged_claims_with_actor(¶ms, actor.as_ref()) {
3291 Ok(c) => c,
3292 Err(e) => {
3293 let code = match &e {
3294 pv_identity::ExchangeError::EmptyIntersection
3295 | pv_identity::ExchangeError::ScopeExceedsSubject(_)
3296 | pv_identity::ExchangeError::InvalidScopes => "invalid_scope",
3297 pv_identity::ExchangeError::DelegationExceeded { .. }
3298 | pv_identity::ExchangeError::ActorNotDirect => "invalid_grant",
3299 pv_identity::ExchangeError::TargetNotAllowed(_) => "invalid_target",
3300 _ => "invalid_request",
3301 };
3302 return (StatusCode::BAD_REQUEST, Json(oauth_error(code, &e.to_string()))).into_response();
3303 }
3304 };
3305 let signer = match &state.token_signer {
3306 Some(s) => s,
3307 None => {
3308 return (
3309 StatusCode::INTERNAL_SERVER_ERROR,
3310 Json(oauth_error(
3311 "server_error",
3312 "token exchange is enabled but no signing key is configured",
3313 )),
3314 )
3315 .into_response();
3316 }
3317 };
3318 let access_token = match signer.sign(&exchanged) {
3319 Ok(jwt) => jwt,
3320 Err(e) => {
3321 return (
3322 StatusCode::INTERNAL_SERVER_ERROR,
3323 Json(oauth_error("server_error", &format!("token signing failed: {e}"))),
3324 )
3325 .into_response();
3326 }
3327 };
3328 let resp = pv_identity::TokenExchangeResponse {
3329 access_token,
3330 issued_token_type: pv_identity::exchange::JWT_TOKEN_TYPE.to_owned(),
3331 token_type: "Bearer".to_owned(),
3332 expires_in: exchanged.exp.saturating_sub(now_s),
3333 scope: if exchanged.scopes.is_empty() {
3334 None
3335 } else {
3336 Some(exchanged.scopes.join(" "))
3337 },
3338 };
3339 (
3340 [
3341 (axum::http::header::CACHE_CONTROL, "no-store"),
3342 (axum::http::header::PRAGMA, "no-cache"),
3343 ],
3344 Json(resp),
3345 )
3346 .into_response()
3347}
3348
3349fn finalize_error_response(error: &crate::reconciler::FinalizeError) -> Response {
3356 use crate::reconciler::FinalizeError;
3357 if error.is_transient_io() {
3368 let mut response = (
3369 StatusCode::SERVICE_UNAVAILABLE,
3370 Json(openai_error_body(
3371 StatusCode::SERVICE_UNAVAILABLE,
3372 &error.to_string(),
3373 )),
3374 )
3375 .into_response();
3376 response
3377 .headers_mut()
3378 .insert(axum::http::header::RETRY_AFTER, HeaderValue::from_static("5"));
3379 return response;
3380 }
3381 let status = match error {
3382 FinalizeError::CaptureIncomplete | FinalizeError::Promotion(_) => StatusCode::CONFLICT,
3387 FinalizeError::BridgeConfig { .. } => StatusCode::BAD_REQUEST,
3393 FinalizeError::Bridge { .. } | FinalizeError::ToolEvidence(_) => StatusCode::SERVICE_UNAVAILABLE,
3396 _ => StatusCode::INTERNAL_SERVER_ERROR,
3398 };
3399 let mut response = (status, Json(openai_error_body(status, &error.to_string()))).into_response();
3400 if status == StatusCode::SERVICE_UNAVAILABLE {
3401 response
3403 .headers_mut()
3404 .insert(axum::http::header::RETRY_AFTER, HeaderValue::from_static("5"));
3405 }
3406 response
3407}
3408
3409fn openai_error_body(status: StatusCode, message: &str) -> serde_json::Value {
3420 let error_type = match status.as_u16() {
3421 400 | 404 | 409 | 413 => "invalid_request_error",
3422 401 => "authentication_error",
3423 403 => "permission_error",
3424 429 => "rate_limit_error",
3425 _ => "api_error",
3426 };
3427 json!({
3428 "error": {
3429 "message": message,
3430 "type": error_type,
3431 "param": null,
3432 "code": status.as_u16(),
3433 }
3434 })
3435}
3436
3437pub(crate) const TOKENS_REVOKED_METRIC: &str = "pv_tokens_revoked_total";
3438pub(crate) const TOKENS_REVOKED_HELP: &str = "Successful holder revocations and operator revocation commands";
3439
3440fn exchange_for_backend(
3445 state: &AppState,
3446 caller: &crate::pipeline::AuthenticatedCaller,
3447 target_audience: &str,
3448 tool: &str,
3449) -> Result<String, crate::pipeline::PipelineError> {
3450 exchange_token_for(state, caller, target_audience, &crate::pipeline::tool_scope(tool))
3451}
3452
3453pub(crate) fn exchange_token_for(
3456 state: &AppState,
3457 caller: &crate::pipeline::AuthenticatedCaller,
3458 target_audience: &str,
3459 requested_scopes: &str,
3460) -> Result<String, crate::pipeline::PipelineError> {
3461 use crate::pipeline::PipelineError;
3462 let refused = |reason: String| PipelineError::blocked(format!("backend token exchange failed: {reason}"));
3463 let Some(signer) = state.token_signer.as_ref() else {
3464 return Err(PipelineError::unavailable(
3465 "backend token exchange is not configured on this gateway".to_owned(),
3466 ));
3467 };
3468 let validated = caller
3469 .validated
3470 .as_ref()
3471 .ok_or_else(|| refused("no validated caller identity".into()))?;
3472 let now_s = pv_core::time::now_ms() / 1000;
3473 let ttl = state.config.token_exchange_ttl_s.min(900);
3474 let allowed_audiences: Vec<String> = state.config.backends.iter().map(|b| b.name.clone()).collect();
3475 let params = pv_identity::ExchangeParams {
3476 subject: &validated.claims,
3477 subject_chain: &validated.chain_tokens,
3478 subject_delegation_depth: validated.total_delegation_depth,
3479 target_audience,
3480 requested_scopes: Some(requested_scopes),
3481 azp: &validated.claims.instance_uid,
3482 now_s,
3483 ttl_s: ttl,
3484 issuer: &state.config.audience,
3485 max_depth: state.config.max_delegation_depth,
3486 allowed_audiences: &allowed_audiences,
3487 };
3488 let exchanged = pv_identity::build_exchanged_claims(¶ms).map_err(|e| refused(e.to_string()))?;
3489 signer.sign_with_typ(&exchanged, "av-tool+jwt").map_err(|error| {
3490 tracing::error!(%error, "exchanged token signing failed");
3491 PipelineError::unavailable("backend token exchange could not sign a token".to_owned())
3492 })
3493}
3494
3495fn oauth_error(error: &str, description: &str) -> serde_json::Value {
3496 json!({
3497 "error": error,
3498 "error_description": description,
3499 })
3500}
3501
3502fn oauth_unavailable() -> Response {
3505 let mut response = (
3506 StatusCode::SERVICE_UNAVAILABLE,
3507 Json(oauth_error(
3508 "temporarily_unavailable",
3509 "the token revocation list is unavailable; retry shortly",
3510 )),
3511 )
3512 .into_response();
3513 response
3514 .headers_mut()
3515 .insert(axum::http::header::RETRY_AFTER, HeaderValue::from_static("5"));
3516 response
3517}
3518
3519fn pipeline_error_for_session(error: crate::pipeline::PipelineError, session_id: &str) -> Response {
3528 let mut response = pipeline_error(error);
3529 if let Ok(value) = HeaderValue::from_str(session_id) {
3530 response
3531 .headers_mut()
3532 .insert(crate::pipeline::SESSION_HEADER, value);
3533 }
3534 response
3535}
3536
3537pub(crate) fn pipeline_error(error: crate::pipeline::PipelineError) -> Response {
3538 use crate::pipeline::PipelineError;
3539 let close = matches!(error, PipelineError::Abort(_));
3540 let status = error.status();
3541 let mut response = (status, Json(openai_error_body(status, &error.to_string()))).into_response();
3542 if close {
3543 response
3544 .headers_mut()
3545 .insert(axum::http::header::CONNECTION, HeaderValue::from_static("close"));
3546 }
3547 match status.as_u16() {
3556 429 | 502 | 503 | 504 => {
3561 response
3562 .headers_mut()
3563 .insert(axum::http::header::RETRY_AFTER, HeaderValue::from_static("5"));
3564 }
3565 401 => {
3568 response.headers_mut().insert(
3569 axum::http::header::WWW_AUTHENTICATE,
3570 HeaderValue::from_static("Bearer realm=\"proofvisor\", error=\"invalid_token\""),
3571 );
3572 }
3573 _ => {}
3574 }
3575 response
3576}
3577
3578fn lifecycle_error(error: String) -> Response {
3579 (
3583 StatusCode::INTERNAL_SERVER_ERROR,
3584 Json(openai_error_body(StatusCode::INTERNAL_SERVER_ERROR, &error)),
3585 )
3586 .into_response()
3587}
3588
3589struct AbortFinalizingStream {
3590 inner: BoxStream<'static, Result<Bytes, std::io::Error>>,
3591 session: Arc<crate::session::Session>,
3592 identity: pv_events::AgentIdentity,
3593 response_permit: Option<crate::worker::ResponsePermit>,
3594 worker: crate::worker::WorkerHandle,
3599 store: Arc<dyn pv_state::StateStore>,
3600 budget: pv_state::BudgetSpec,
3601 rates: Arc<crate::policy_push::RateLimits>,
3604 finalizer: crate::reconciler::Finalizer,
3605 _lease: crate::session::SessionLease,
3606 response_marker: Option<String>,
3607 response_attempt_id: String,
3608 response_choices: std::collections::BTreeMap<u64, ChoiceCapture>,
3613 response_model: Option<String>,
3614 upstream_status: StatusCode,
3615 response_cost_usd_micros: u64,
3616 response_tool_calls: std::collections::BTreeMap<(u64, u64), PartialToolCall>,
3617 response_metrics: pv_events::EventMetrics,
3618 charged_completion_tokens: u64,
3619 last_reported_completion_tokens: Option<u64>,
3620 last_reported_prompt_tokens: Option<u64>,
3621 last_reported_cached_tokens: Option<u64>,
3622 last_reported_cost_usd_micros: Option<u64>,
3623 saw_chunk: bool,
3624 capture_attempted: bool,
3632 is_sse: bool,
3633 protocol_buffer: Vec<u8>,
3634 frame_scratch: Vec<u8>,
3644 pending_output: std::collections::VecDeque<Bytes>,
3645 pending_budget: Option<PendingBudget>,
3646 principal_budget: Option<(String, pv_state::BudgetSpec)>,
3655 captured_bytes: usize,
3656 billed_prompt_tokens: u64,
3659 completed: bool,
3660 ephemeral: bool,
3669}
3670
3671const MAX_PROVIDER_CAPTURE_BYTES: usize = 16 * 1024 * 1024;
3672const MAX_PROVIDER_FIELD_BYTES: usize = 8 * 1024 * 1024;
3673const MAX_PROVIDER_TOOL_CALLS: usize = 128;
3674
3675impl AbortFinalizingStream {
3676 fn absorb_network_chunk(&mut self, raw: &[u8]) -> Result<u64, String> {
3677 self.saw_chunk |= !raw.is_empty();
3678 self.captured_bytes = self
3679 .captured_bytes
3680 .checked_add(raw.len())
3681 .ok_or_else(|| "provider response size overflow".to_owned())?;
3682 if self.captured_bytes > MAX_PROVIDER_CAPTURE_BYTES {
3683 return Err(format!(
3684 "provider response exceeds {MAX_PROVIDER_CAPTURE_BYTES} capture bytes"
3685 ));
3686 }
3687 self.protocol_buffer.extend_from_slice(raw);
3688 if !self.is_sse {
3689 return Ok(0);
3690 }
3691 if !self.upstream_status.is_success() {
3700 self.protocol_buffer.clear();
3701 return Ok(0);
3702 }
3703 let mut budget_delta = 0u64;
3704 let mut scratch: Vec<u8> = std::mem::take(&mut self.frame_scratch);
3715 let absorb_result = loop {
3716 let Some(end) = sse_frame_end(&self.protocol_buffer) else {
3717 break Ok(());
3718 };
3719 scratch.clear();
3720 #[allow(clippy::indexing_slicing)]
3723 scratch.extend_from_slice(&self.protocol_buffer[..end]);
3724 self.protocol_buffer.drain(..end);
3725 let frame = match std::str::from_utf8(&scratch) {
3726 Ok(frame) => frame,
3727 Err(error) => break Err(format!("provider SSE frame is not UTF-8: {error}")),
3728 };
3729 match self.absorb_frame(frame) {
3730 Ok(delta) => match budget_delta.checked_add(delta) {
3731 Some(total) => budget_delta = total,
3732 None => break Err("provider completion-token delta overflow".to_owned()),
3733 },
3734 Err(error) => break Err(error),
3735 }
3736 };
3737 self.frame_scratch = scratch;
3738 if let Err(error) = absorb_result {
3739 self.charged_completion_tokens = self.charged_completion_tokens.saturating_sub(budget_delta);
3751 return Err(error);
3752 }
3753 if self.protocol_buffer.len() > MAX_PROVIDER_FIELD_BYTES {
3754 self.charged_completion_tokens = self.charged_completion_tokens.saturating_sub(budget_delta);
3759 return Err("unterminated provider SSE frame exceeds limit".to_owned());
3760 }
3761 Ok(budget_delta)
3762 }
3763
3764 fn flush_protocol_buffer(&mut self) -> Result<u64, String> {
3765 if self.protocol_buffer.is_empty() {
3766 return Ok(0);
3767 }
3768 if !self.upstream_status.is_success() {
3781 let _ = std::mem::take(&mut self.protocol_buffer);
3782 return Ok(0);
3783 }
3784 let frame = std::mem::take(&mut self.protocol_buffer);
3785 let frame = std::str::from_utf8(&frame)
3786 .map_err(|error| format!("provider response is not UTF-8: {error}"))?;
3787 self.absorb_frame(frame)
3788 }
3789
3790 fn absorb_frame(&mut self, raw: &str) -> Result<u64, String> {
3791 let Some(parsed) = parse_provider_chunk(raw)? else {
3795 return Ok(0);
3796 };
3797 if self.upstream_status.is_success() && !parsed.has_choices {
3798 return Err("successful provider response has no choices array".to_owned());
3799 }
3800 for choice in parsed.choices {
3801 if choice.index >= MAX_CHOICE_INDEX {
3802 return Err(format!("provider choice index {} is out of range", choice.index));
3803 }
3804 let capture = self.response_choices.entry(choice.index).or_default();
3805 push_bounded(&mut capture.content, &choice.content, "response message")?;
3806 push_bounded(&mut capture.reasoning, &choice.reasoning, "response reasoning")?;
3807 if choice.finish_reason.is_some() {
3808 capture.finish_reason = choice.finish_reason;
3809 }
3810 }
3811 if self.response_model.is_none() {
3812 self.response_model = parsed.model_name;
3813 }
3814 self.response_cost_usd_micros = self.response_cost_usd_micros.max(parsed.cost_usd_micros);
3829 for delta in parsed.tool_call_deltas {
3830 let key = if delta.complete {
3831 let next = self
3840 .response_tool_calls
3841 .range((delta.choice_index, 0)..=(delta.choice_index, u64::MAX))
3842 .next_back()
3843 .map_or(0, |(&(_, index), _)| index.saturating_add(1));
3844 (delta.choice_index, next)
3845 } else {
3846 if delta.index >= MAX_PROVIDER_TOOL_CALLS as u64 {
3847 return Err(format!(
3848 "provider tool-call index {} is out of range",
3849 delta.index
3850 ));
3851 }
3852 (delta.choice_index, delta.index)
3853 };
3854 if !self.response_tool_calls.contains_key(&key)
3855 && self.response_tool_calls.len() >= MAX_PROVIDER_TOOL_CALLS
3856 {
3857 return Err(format!(
3858 "provider response exceeds {MAX_PROVIDER_TOOL_CALLS} tool calls"
3859 ));
3860 }
3861 let partial = self.response_tool_calls.entry(key).or_default();
3862 if let (Some(existing), Some(new)) = (partial.id.as_deref(), delta.id.as_deref()) {
3870 if existing != new {
3871 return Err(format!(
3872 "provider reused tool-call index {} with conflicting id",
3873 delta.index
3874 ));
3875 }
3876 }
3877 if let (Some(existing), Some(new)) = (partial.name.as_deref(), delta.name.as_deref()) {
3878 if existing != new {
3879 return Err(format!(
3880 "provider reused tool-call index {} with conflicting name",
3881 delta.index
3882 ));
3883 }
3884 }
3885 if delta.id.is_some() {
3886 partial.id = delta.id;
3887 }
3888 if delta.name.is_some() {
3889 partial.name = delta.name;
3890 }
3891 push_bounded(&mut partial.arguments, &delta.arguments, "tool-call arguments")?;
3892 }
3893 if parsed.usage_reported {
3894 reject_metric_regression(
3895 "prompt tokens",
3896 &mut self.last_reported_prompt_tokens,
3897 parsed.metrics.prompt_tokens,
3898 )?;
3899 reject_metric_regression(
3900 "cached tokens",
3901 &mut self.last_reported_cached_tokens,
3902 parsed.metrics.cached_tokens,
3903 )?;
3904 }
3905 if parsed.cost_reported {
3906 reject_metric_regression(
3907 "cost",
3908 &mut self.last_reported_cost_usd_micros,
3909 Some(parsed.cost_usd_micros),
3910 )?;
3911 }
3912 if let Some(prompt) = parsed.metrics.prompt_tokens {
3913 self.response_metrics.prompt_tokens = Some(
3914 self.response_metrics
3915 .prompt_tokens
3916 .map_or(prompt, |current| current.max(prompt)),
3917 );
3918 }
3919 if let Some(cached) = parsed.metrics.cached_tokens {
3920 self.response_metrics.cached_tokens = Some(
3921 self.response_metrics
3922 .cached_tokens
3923 .map_or(cached, |current| current.max(cached)),
3924 );
3925 }
3926 let reported_completion = parsed.metrics.completion_tokens.unwrap_or(0);
3927 let delta = if parsed.usage_reported && parsed.completion_reported {
3938 let previous = self.last_reported_completion_tokens.unwrap_or(0);
3939 if reported_completion < previous {
3940 return Err("provider completion usage regressed".to_owned());
3941 }
3942 self.last_reported_completion_tokens = Some(reported_completion);
3943 let delta = reported_completion.saturating_sub(self.charged_completion_tokens);
3944 self.response_metrics.completion_tokens = Some(reported_completion);
3957 delta
3958 } else {
3959 let total = self
3960 .response_metrics
3961 .completion_tokens
3962 .unwrap_or(0)
3963 .checked_add(reported_completion)
3964 .filter(|total| *total <= pv_core::error::JCS_SAFE_MAX)
3965 .ok_or_else(|| "provider completion-token total exceeds JCS-safe bounds".to_owned())?;
3966 self.response_metrics.completion_tokens = Some(total);
3967 reported_completion
3968 };
3969 if delta == 0 {
3970 return Ok(0);
3971 }
3972 self.charged_completion_tokens = self
3973 .charged_completion_tokens
3974 .checked_add(delta)
3975 .ok_or_else(|| "charged completion-token counter overflow".to_owned())?;
3976 self.rates
3977 .record(&crate::pipeline::principal_id_for_budget(&self.identity), delta);
3978 Ok(delta)
3979 }
3980
3981 fn begin_budget_check(&mut self, delta: u64, continuation: BudgetContinuation) {
3982 let store = Arc::clone(&self.store);
3983 let session_id = self.session.id.clone();
3984 let budget = self.budget.clone();
3985 let principal = self.principal_budget.clone();
3986 let task = tokio::task::spawn_blocking(move || {
3987 if let Some((principal_id, spec)) = principal.as_ref() {
3992 match pv_state::ActionBudget::for_principal(store.as_ref(), principal_id, spec)
3993 .try_tokens(delta)
3994 .map_err(|error| format!("token budget backend failed closed (principal): {error}"))?
3995 {
3996 pv_state::BudgetDecision::Refused { limit, cap } => {
3997 return Ok(pv_state::BudgetDecision::Refused {
3998 limit: format!("principal.{limit}"),
3999 cap,
4000 });
4001 }
4002 pv_state::BudgetDecision::Allowed { .. } => {}
4003 }
4004 }
4005 let decision =
4012 match pv_state::ActionBudget::new(store.as_ref(), &session_id, &budget).try_tokens(delta) {
4013 Ok(decision) => decision,
4014 Err(error) => {
4015 if let Some((principal_id, spec)) = principal.as_ref() {
4016 pv_state::ActionBudget::for_principal(store.as_ref(), principal_id, spec)
4017 .refund_tokens(delta);
4018 }
4019 return Err(format!("token budget backend failed closed: {error}"));
4020 }
4021 };
4022 if matches!(decision, pv_state::BudgetDecision::Refused { .. }) {
4023 if let Some((principal_id, spec)) = principal.as_ref() {
4024 pv_state::ActionBudget::for_principal(store.as_ref(), principal_id, spec)
4025 .refund_tokens(delta);
4026 }
4027 }
4028 Ok(decision)
4029 });
4030 self.pending_budget = Some(PendingBudget {
4031 task,
4032 continuation,
4033 delta,
4034 });
4035 }
4036
4037 fn submit_response_capture(
4038 &mut self,
4039 failure: Option<(pv_events::StopReason, String)>,
4040 ) -> Result<(), crate::worker::SubmitError> {
4041 if self.capture_attempted {
4042 return Ok(());
4043 }
4044 self.capture_attempted = true;
4045 let choices = std::mem::take(&mut self.response_choices);
4046 let multi_choice = choices.len() > 1;
4047 let reasoning = {
4048 let joined = choices
4049 .values()
4050 .filter(|capture| !capture.reasoning.is_empty())
4051 .map(|capture| capture.reasoning.as_str())
4052 .collect::<Vec<_>>()
4053 .join("\n\n");
4054 (!joined.is_empty()).then_some(joined)
4055 };
4056 let message_text = choices
4059 .values()
4060 .map(|capture| capture.content.as_str())
4061 .collect::<Vec<_>>()
4062 .join("\n");
4063 let message = if multi_choice {
4064 Value::Array(
4065 choices
4066 .iter()
4067 .map(|(index, capture)| {
4068 json!({"type": "text", "text": capture.content, "choice_index": index})
4069 })
4070 .collect(),
4071 )
4072 } else {
4073 Value::String(message_text.clone())
4074 };
4075 let native_finish_reason = aggregate_finish_reason(&choices);
4076 let tool_calls: Vec<pv_atif::ToolCall> = std::mem::take(&mut self.response_tool_calls)
4086 .into_iter()
4087 .map(|((choice_index, _), partial)| {
4088 let arguments = normalize_tool_call_arguments(&partial.arguments);
4089 pv_atif::ToolCall {
4090 tool_call_id: partial.id.unwrap_or_else(pv_core::new_event_uid),
4091 function_name: partial.name.unwrap_or_else(|| "unknown".to_owned()),
4092 arguments,
4093 extra: multi_choice.then(|| json!({"choice_index": choice_index})),
4097 }
4098 })
4099 .collect();
4100 let analysis_text = if let Some(text) = reasoning.clone() {
4101 text
4102 } else if !message_text.is_empty() {
4103 message_text.clone()
4104 } else if !tool_calls.is_empty() {
4105 tool_calls
4106 .iter()
4107 .map(|tc| format!("{}({})", tc.function_name, tc.arguments))
4108 .collect::<Vec<_>>()
4109 .join(" ")
4110 } else {
4111 String::new()
4112 };
4113 let analyze_loop = !analysis_text.is_empty();
4117 let (class, status, stop_reason, payload) = if let Some((failure_reason, reason)) = failure {
4124 (
4125 pv_events::EventClass::StopReason,
4126 pv_events::StatusId::Failure,
4127 Some(failure_reason),
4128 json!({"reason": reason, "direction": "upstream_response"}),
4129 )
4130 } else if !self.upstream_status.is_success() {
4131 (
4132 pv_events::EventClass::StopReason,
4133 pv_events::StatusId::Failure,
4134 Some(pv_events::StopReason::Other),
4135 json!({
4136 "direction": "upstream_response",
4137 "http_status": self.upstream_status.as_u16(),
4138 }),
4139 )
4140 } else if let Some(native) = &native_finish_reason {
4141 let mapped = map_finish_reason(native);
4157 let normalized = if mapped == pv_events::StopReason::Stop && !tool_calls.is_empty() {
4158 pv_events::StopReason::ToolUse
4159 } else {
4160 mapped
4161 };
4162 (
4163 pv_events::EventClass::StopReason,
4164 pv_events::StatusId::Success,
4165 Some(normalized),
4166 json!({
4167 "direction": "upstream_response",
4168 "finish_reason": native,
4169 "http_status": self.upstream_status.as_u16(),
4170 }),
4171 )
4172 } else if self.is_sse {
4173 (
4184 pv_events::EventClass::StopReason,
4185 pv_events::StatusId::Unknown,
4186 Some(pv_events::StopReason::Other),
4187 json!({
4188 "direction": "upstream_response",
4189 "http_status": self.upstream_status.as_u16(),
4190 "truncated": true,
4191 "reason": "stream ended without a finish_reason chunk",
4192 }),
4193 )
4194 } else if message_text.is_empty() && reasoning.is_none() && tool_calls.is_empty() {
4195 (
4205 pv_events::EventClass::StopReason,
4206 pv_events::StatusId::Unknown,
4207 Some(pv_events::StopReason::Other),
4208 json!({
4209 "direction": "upstream_response",
4210 "http_status": self.upstream_status.as_u16(),
4211 "truncated": true,
4212 "reason": "response carried no choices, content, or finish_reason",
4213 }),
4214 )
4215 } else {
4216 (
4217 pv_events::EventClass::Session,
4218 pv_events::StatusId::Success,
4219 None,
4220 json!({
4221 "direction": "upstream_response",
4222 "http_status": self.upstream_status.as_u16(),
4223 }),
4224 )
4225 };
4226 if let Some(provider) = self.response_metrics.prompt_tokens {
4237 let over_charge = self.billed_prompt_tokens.saturating_sub(provider);
4238 if over_charge > 0 {
4239 pv_state::ActionBudget::new(self.store.as_ref(), &self.session.id, &self.budget)
4240 .refund_tokens(over_charge);
4241 }
4242 }
4243 if let Some(provider) = self.response_metrics.completion_tokens {
4252 let over_charge = self.charged_completion_tokens.saturating_sub(provider);
4253 if over_charge > 0 {
4254 pv_state::ActionBudget::new(self.store.as_ref(), &self.session.id, &self.budget)
4255 .refund_tokens(over_charge);
4256 }
4257 }
4258 let permit = self
4259 .response_permit
4260 .take()
4261 .ok_or(crate::worker::SubmitError::Closed)?;
4262 permit.submit(
4263 &self.worker,
4264 crate::worker::WorkerJob {
4265 session: Arc::clone(&self.session),
4266 identity: self.identity.clone(),
4267 class,
4268 payload,
4269 text: analysis_text,
4270 analyze_loop,
4271 status,
4272 stop_reason,
4273 native_stop_reason: native_finish_reason,
4274 metrics: self.response_metrics,
4275 cost_usd_micros: self.response_cost_usd_micros,
4276 prompt_token_correction: self.response_metrics.prompt_tokens.map_or(0i64, |provider| {
4284 let provider = i64::try_from(provider).unwrap_or(i64::MAX);
4285 let billed = i64::try_from(self.billed_prompt_tokens).unwrap_or(i64::MAX);
4286 provider.saturating_sub(billed)
4287 }),
4288 atif: Some(crate::worker::AtifCapture {
4289 source: pv_atif::Source::Agent,
4290 message,
4291 reasoning_content: reasoning,
4292 model_name: self.response_model.take(),
4293 tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
4294 observation: None,
4295 llm_call_count: Some(1),
4296 }),
4297 response_marker: self.response_marker.take(),
4298 response_attempt: Some(crate::worker::ResponseAttempt {
4299 id: self.response_attempt_id.clone(),
4300 terminal: true,
4301 }),
4302 },
4303 )?;
4304 Ok(())
4305 }
4306}
4307
4308enum BudgetContinuation {
4309 Emit(Bytes),
4310 ContinueNonSse,
4311 FinishSse,
4312 FinishNonSse,
4313}
4314
4315struct PendingBudget {
4316 task: tokio::task::JoinHandle<Result<pv_state::BudgetDecision, String>>,
4317 continuation: BudgetContinuation,
4318 delta: u64,
4327}
4328
4329#[derive(Default)]
4330struct PartialToolCall {
4331 id: Option<String>,
4332 name: Option<String>,
4333 arguments: String,
4334}
4335
4336const MAX_CHOICE_INDEX: u64 = 128;
4340
4341#[derive(Default, Debug)]
4343struct ChoiceCapture {
4344 content: String,
4345 reasoning: String,
4346 finish_reason: Option<String>,
4347}
4348
4349fn aggregate_finish_reason(choices: &std::collections::BTreeMap<u64, ChoiceCapture>) -> Option<String> {
4355 if choices.is_empty() || choices.values().any(|capture| capture.finish_reason.is_none()) {
4356 return None;
4357 }
4358 let rank = |reason: &str| match map_finish_reason(reason) {
4359 pv_events::StopReason::ContentFilter => 4,
4360 pv_events::StopReason::MaxTokens => 3,
4361 pv_events::StopReason::ToolUse => 2,
4362 pv_events::StopReason::Stop => 0,
4363 _ => 1,
4364 };
4365 choices
4366 .values()
4367 .filter_map(|capture| capture.finish_reason.as_deref())
4368 .max_by_key(|reason| rank(reason))
4369 .map(str::to_owned)
4370}
4371
4372#[derive(Debug, Default)]
4374pub(crate) struct ParsedChoiceDelta {
4375 pub(crate) index: u64,
4376 pub(crate) content: String,
4377 pub(crate) reasoning: String,
4378 pub(crate) finish_reason: Option<String>,
4379}
4380
4381#[derive(Debug)]
4382pub(crate) struct ProviderToolCallDelta {
4383 pub(crate) choice_index: u64,
4389 pub(crate) index: u64,
4395 pub(crate) complete: bool,
4401 pub(crate) id: Option<String>,
4402 pub(crate) name: Option<String>,
4403 pub(crate) arguments: String,
4404}
4405
4406#[derive(Debug)]
4407pub(crate) struct ParsedProviderChunk {
4408 pub(crate) model_name: Option<String>,
4409 pub(crate) metrics: pv_events::EventMetrics,
4410 pub(crate) usage_reported: bool,
4411 pub(crate) completion_reported: bool,
4418 pub(crate) cost_usd_micros: u64,
4419 pub(crate) cost_reported: bool,
4420 pub(crate) has_choices: bool,
4421 pub(crate) tool_call_deltas: Vec<ProviderToolCallDelta>,
4422 pub(crate) choices: Vec<ParsedChoiceDelta>,
4425}
4426
4427fn reject_metric_regression(
4428 field: &str,
4429 previous: &mut Option<u64>,
4430 current: Option<u64>,
4431) -> Result<(), String> {
4432 let Some(current) = current else {
4433 return Ok(());
4434 };
4435 if previous.is_some_and(|previous| current < previous) {
4436 return Err(format!("provider {field} regressed"));
4437 }
4438 *previous = Some(current);
4439 Ok(())
4440}
4441
4442fn push_bounded(target: &mut String, value: &str, field: &str) -> Result<(), String> {
4443 let size = target
4444 .len()
4445 .checked_add(value.len())
4446 .ok_or_else(|| format!("{field} size overflow"))?;
4447 if size > MAX_PROVIDER_FIELD_BYTES {
4448 return Err(format!("{field} exceeds {MAX_PROVIDER_FIELD_BYTES} bytes"));
4449 }
4450 target.push_str(value);
4451 Ok(())
4452}
4453
4454#[doc(hidden)]
4460pub fn __fuzz_parse_provider_chunk(raw: &[u8]) {
4461 if let Ok(s) = std::str::from_utf8(raw) {
4462 let _ = parse_provider_chunk(s);
4463 }
4464}
4465
4466#[doc(hidden)]
4469pub fn __fuzz_sse_frame_end(raw: &[u8]) {
4470 let _ = sse_frame_end(raw);
4471}
4472
4473pub(crate) fn sse_frame_end(buffer: &[u8]) -> Option<usize> {
4474 let mut line_start = 0usize;
4475 let mut index = 0usize;
4476 while index < buffer.len() {
4477 let byte = buffer.get(index).copied()?;
4478 let newline = match byte {
4479 b'\r' if buffer.get(index + 1) == Some(&b'\n') => 2,
4480 b'\r' if buffer.get(index + 1).is_none() => return None,
4481 b'\r' | b'\n' => 1,
4482 _ => {
4483 index += 1;
4484 continue;
4485 }
4486 };
4487 if index == line_start {
4488 return Some(index + newline);
4489 }
4490 index += newline;
4491 line_start = index;
4492 }
4493 None
4494}
4495
4496impl AbortFinalizingStream {
4497 fn abort_error(&self, reason: String) -> std::io::Error {
4503 tracing::warn!(
4504 session = %self.session.id,
4505 %reason,
4506 "aborting client response; upstream reply could not be captured"
4507 );
4508 std::io::Error::other(reason)
4509 }
4510
4511 fn record_response_failure(&mut self, reason: &str) {
4522 if let Err(error) =
4531 self.submit_response_capture(Some((pv_events::StopReason::Other, reason.to_owned())))
4532 {
4533 tracing::warn!(
4534 %error,
4535 session = %self.session.id,
4536 "response-failure capture submit failed; sealing session capture"
4537 );
4538 self.session.mark_capture_failed();
4539 }
4540 }
4541}
4542
4543impl Stream for AbortFinalizingStream {
4544 type Item = Result<Bytes, std::io::Error>;
4545
4546 fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
4547 if self.pending_budget.is_some() {
4548 let result = {
4549 let Some(pending) = self.pending_budget.as_mut() else {
4550 self.session.mark_capture_failed();
4551 return Poll::Ready(Some(Err(std::io::Error::other(
4552 "pending completion budget state disappeared",
4553 ))));
4554 };
4555 match Pin::new(&mut pending.task).poll(context) {
4556 Poll::Pending => return Poll::Pending,
4557 Poll::Ready(result) => result,
4558 }
4559 };
4560 let Some(pending) = self.pending_budget.take() else {
4561 self.session.mark_capture_failed();
4562 return Poll::Ready(Some(Err(std::io::Error::other(
4563 "completed completion budget state disappeared",
4564 ))));
4565 };
4566 let failure = match result {
4567 Ok(Ok(pv_state::BudgetDecision::Allowed { .. })) => None,
4568 Ok(Ok(pv_state::BudgetDecision::Refused { limit, cap })) => {
4569 self.charged_completion_tokens =
4577 self.charged_completion_tokens.saturating_sub(pending.delta);
4578 self.session
4586 .latch_enforcement(pv_events::StopReason::BudgetExceeded);
4587 Some((
4588 pv_events::StopReason::BudgetExceeded,
4589 format!("{limit} exceeded (cap {cap})"),
4590 ))
4591 }
4592 Ok(Err(reason)) => {
4598 self.charged_completion_tokens =
4599 self.charged_completion_tokens.saturating_sub(pending.delta);
4600 Some((pv_events::StopReason::Other, reason))
4601 }
4602 Err(error) => {
4603 self.session.mark_capture_failed();
4604 self.pending_output.clear();
4605 return Poll::Ready(Some(Err(std::io::Error::other(format!(
4606 "completion budget task failed: {error}"
4607 )))));
4608 }
4609 };
4610 if let Some((failure_reason, reason)) = failure {
4611 if let Err(error) = self.submit_response_capture(Some((failure_reason, reason.clone()))) {
4612 self.session.mark_capture_failed();
4613 self.pending_output.clear();
4614 return Poll::Ready(Some(Err(std::io::Error::other(format!(
4615 "response capture failed closed: {error}"
4616 )))));
4617 }
4618 self.pending_output.clear();
4619 return Poll::Ready(Some(Err(std::io::Error::new(
4626 std::io::ErrorKind::QuotaExceeded,
4627 format!("response blocked by token budget: {reason}"),
4628 ))));
4629 }
4630 match pending.continuation {
4631 BudgetContinuation::Emit(bytes) => return Poll::Ready(Some(Ok(bytes))),
4632 BudgetContinuation::ContinueNonSse => {}
4633 BudgetContinuation::FinishSse => {
4634 if let Err(error) = self.submit_response_capture(None) {
4635 self.session.mark_capture_failed();
4636 return Poll::Ready(Some(Err(std::io::Error::other(format!(
4637 "response capture failed closed: {error}"
4638 )))));
4639 }
4640 self.completed = true;
4641 return Poll::Ready(None);
4642 }
4643 BudgetContinuation::FinishNonSse => {
4644 if let Err(error) = self.submit_response_capture(None) {
4645 self.session.mark_capture_failed();
4646 self.pending_output.clear();
4647 return Poll::Ready(Some(Err(std::io::Error::other(format!(
4648 "response capture failed closed: {error}"
4649 )))));
4650 }
4651 self.completed = true;
4652 return Poll::Ready(self.pending_output.pop_front().map(Ok));
4653 }
4654 }
4655 }
4656
4657 if !self.is_sse {
4658 if self.completed {
4659 return Poll::Ready(self.pending_output.pop_front().map(Ok));
4660 }
4661 loop {
4662 match self.inner.as_mut().poll_next(context) {
4663 Poll::Ready(Some(Ok(bytes))) => {
4664 let delta = match self.absorb_network_chunk(&bytes) {
4665 Ok(delta) => delta,
4666 Err(error) => {
4667 self.record_response_failure(&error);
4668 self.pending_output.clear();
4669 let error = self.abort_error(error);
4670 return Poll::Ready(Some(Err(error)));
4671 }
4672 };
4673 self.pending_output.push_back(bytes);
4674 if delta > 0 {
4675 self.begin_budget_check(delta, BudgetContinuation::ContinueNonSse);
4676 context.waker().wake_by_ref();
4677 return Poll::Pending;
4678 }
4679 }
4680 Poll::Ready(Some(Err(error))) => {
4681 let reason = error.to_string();
4682 self.record_response_failure(&reason);
4683 self.pending_output.clear();
4684 let error = self.abort_error(reason);
4685 return Poll::Ready(Some(Err(error)));
4686 }
4687 Poll::Pending => return Poll::Pending,
4688 Poll::Ready(None) => {
4689 let delta = match self.flush_protocol_buffer() {
4690 Ok(delta) => delta,
4691 Err(error) => {
4692 self.record_response_failure(&error);
4693 self.pending_output.clear();
4694 let error = self.abort_error(error);
4695 return Poll::Ready(Some(Err(error)));
4696 }
4697 };
4698 if delta > 0 {
4699 self.begin_budget_check(delta, BudgetContinuation::FinishNonSse);
4700 context.waker().wake_by_ref();
4701 return Poll::Pending;
4702 }
4703 if self.upstream_status.is_success() && !self.saw_chunk {
4712 let reason = "successful provider response has an empty body".to_owned();
4720 self.record_response_failure(&reason);
4721 self.session.mark_capture_failed();
4722 self.pending_output.clear();
4723 let error = self.abort_error(reason);
4724 return Poll::Ready(Some(Err(error)));
4725 }
4726 if let Err(error) = self.submit_response_capture(None) {
4727 self.session.mark_capture_failed();
4728 self.pending_output.clear();
4729 return Poll::Ready(Some(Err(std::io::Error::other(format!(
4730 "response capture failed closed: {error}"
4731 )))));
4732 }
4733 self.completed = true;
4734 return Poll::Ready(self.pending_output.pop_front().map(Ok));
4735 }
4736 }
4737 }
4738 }
4739
4740 match self.inner.as_mut().poll_next(context) {
4741 Poll::Ready(Some(Ok(bytes))) => match self.absorb_network_chunk(&bytes) {
4742 Ok(0) => Poll::Ready(Some(Ok(bytes))),
4743 Ok(delta) => {
4744 self.begin_budget_check(delta, BudgetContinuation::Emit(bytes));
4745 context.waker().wake_by_ref();
4746 Poll::Pending
4747 }
4748 Err(error) => {
4749 self.record_response_failure(&error);
4750 let error = self.abort_error(error);
4751 Poll::Ready(Some(Err(error)))
4752 }
4753 },
4754 Poll::Ready(Some(Err(error))) => {
4755 let reason = error.to_string();
4756 self.record_response_failure(&reason);
4757 let error = self.abort_error(reason);
4758 Poll::Ready(Some(Err(error)))
4759 }
4760 Poll::Pending => Poll::Pending,
4761 Poll::Ready(None) => {
4762 let delta = match self.flush_protocol_buffer() {
4763 Ok(delta) => delta,
4764 Err(error) => {
4765 self.record_response_failure(&error);
4766 let error = self.abort_error(error);
4767 return Poll::Ready(Some(Err(error)));
4768 }
4769 };
4770 if delta > 0 {
4771 self.begin_budget_check(delta, BudgetContinuation::FinishSse);
4772 context.waker().wake_by_ref();
4773 return Poll::Pending;
4774 }
4775 if self.upstream_status.is_success() && !self.saw_chunk {
4779 let reason = "successful provider response has an empty body".to_owned();
4784 self.record_response_failure(&reason);
4785 self.session.mark_capture_failed();
4786 let error = self.abort_error(reason);
4787 return Poll::Ready(Some(Err(error)));
4788 }
4789 if let Err(error) = self.submit_response_capture(None) {
4790 self.session.mark_capture_failed();
4791 return Poll::Ready(Some(Err(std::io::Error::other(format!(
4792 "response capture failed closed: {error}"
4793 )))));
4794 }
4795 self.completed = true;
4796 Poll::Ready(None)
4797 }
4798 }
4799 }
4800}
4801
4802pub(crate) fn parse_provider_chunk(raw: &str) -> Result<Option<ParsedProviderChunk>, String> {
4803 let raw = raw.strip_prefix('\u{feff}').unwrap_or(raw);
4812 let mut message = String::new();
4813 let mut reasoning = String::new();
4814 let mut prompt_tokens = None;
4815 let mut completion_tokens = None;
4816 let mut cached_tokens = None;
4817 let mut usage_reported = false;
4818 let mut cost_usd_micros = 0u64;
4819 let mut cost_reported = false;
4820 let mut tool_call_deltas = Vec::new();
4821 let mut choice_deltas: Vec<ParsedChoiceDelta> = Vec::new();
4822 let mut is_sse = false;
4823 let mut data = Vec::new();
4824 let mut event_type = String::new();
4839 for line in raw.split(['\r', '\n']) {
4840 if let Some(value) = line.strip_prefix("data:") {
4841 is_sse = true;
4842 data.push(value.strip_prefix(' ').unwrap_or(value));
4848 } else if let Some(value) = line.strip_prefix("event:") {
4849 is_sse = true;
4850 event_type = value.strip_prefix(' ').unwrap_or(value).to_owned();
4851 } else if line == "data"
4852 || line.starts_with("id:")
4853 || line.starts_with("retry:")
4854 || line.starts_with(':')
4855 {
4856 is_sse = true;
4857 }
4858 }
4859 if is_sse && !event_type.is_empty() && event_type != "message" {
4860 if data.iter().all(|entry| {
4878 let trimmed = entry.trim();
4879 trimmed.is_empty() || trimmed == "[DONE]"
4880 }) {
4881 return Ok(None);
4882 }
4883 return Err(format!(
4884 "provider SSE frame carries unsupported event type {event_type:?}; \
4885 ProofVisor only captures the default `message` event because non-message \
4886 events (error, ping, custom) are dispatched to different client listeners \
4887 per SSE §9.2.8 and would be attributed to the wrong audit surface"
4888 ));
4889 }
4890 let candidate = if is_sse {
4891 if data.is_empty() {
4892 return Ok(None);
4893 }
4894 data.join("\n")
4895 } else {
4896 raw.trim().to_owned()
4897 };
4898 if candidate.is_empty()
4904 || candidate.trim() == "[DONE]"
4905 || data
4906 .iter()
4907 .all(|entry| entry.trim().is_empty() || entry.trim() == "[DONE]")
4908 && !data.is_empty()
4909 {
4910 return Ok(None);
4911 }
4912 let value = serde_json::from_str::<Value>(&candidate)
4913 .map_err(|error| format!("invalid provider JSON frame: {error}"))?;
4914 if let Some(err_obj) = value.get("error").filter(|v| !v.is_null()) {
4926 let message = err_obj
4927 .get("message")
4928 .and_then(Value::as_str)
4929 .unwrap_or("provider ended stream with an unspecified error");
4930 let code = err_obj.get("code").and_then(Value::as_str).unwrap_or("");
4931 return Err(if code.is_empty() {
4932 format!("provider stream error: {message}")
4933 } else {
4934 format!("provider stream error [{code}]: {message}")
4935 });
4936 }
4937 let model_name = value.get("model").and_then(Value::as_str).map(str::to_owned);
4938 if let Some(usage) = value.get("usage").filter(|usage| !usage.is_null()) {
4948 usage_reported = true;
4949 prompt_tokens = provider_u64(usage.get("prompt_tokens"), "prompt_tokens")?.or(prompt_tokens);
4950 completion_tokens =
4951 provider_u64(usage.get("completion_tokens"), "completion_tokens")?.or(completion_tokens);
4952 cached_tokens = provider_u64(
4953 usage.pointer("/prompt_tokens_details/cached_tokens"),
4954 "cached_tokens",
4955 )?
4956 .or(cached_tokens);
4957 }
4958 let cost_value = value.pointer("/usage/cost_usd").or_else(|| value.get("cost_usd"));
4959 if let Some(cost_value) = cost_value {
4960 let cost = cost_value
4961 .as_f64()
4962 .ok_or_else(|| "provider cost_usd is not a number".to_owned())?;
4963 if !cost.is_finite() || cost < 0.0 {
4964 return Err("provider cost_usd is not finite and nonnegative".to_owned());
4965 }
4966 #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
4967 {
4968 let micros = (cost * pv_core::units::USD_MICROS_PER_DOLLAR as f64).round();
4969 if micros > pv_core::error::JCS_SAFE_MAX as f64 {
4970 return Err("provider cost exceeds JCS-safe receipt bounds".to_owned());
4971 }
4972 cost_usd_micros = micros as u64;
4973 cost_reported = true;
4974 }
4975 }
4976 let has_choices = value.get("choices").is_some();
4977 let choices = match value.get("choices") {
4978 Some(Value::Array(choices)) => Some(choices),
4979 Some(_) => return Err("provider choices is not an array".to_owned()),
4980 None => None,
4981 };
4982 if let Some(choices) = choices {
4983 for (choice_position, choice) in choices.iter().enumerate() {
4984 let choice_index = choice
4990 .get("index")
4991 .and_then(Value::as_u64)
4992 .or_else(|| u64::try_from(choice_position).ok())
4993 .unwrap_or(u64::MAX);
4994 let mut choice_delta = ParsedChoiceDelta {
4995 index: choice_index,
4996 ..ParsedChoiceDelta::default()
4997 };
4998 if let Some(reason) = choice
4999 .get("finish_reason")
5000 .and_then(Value::as_str)
5001 .filter(|s| !s.is_empty())
5002 {
5003 choice_delta.finish_reason = Some(reason.to_owned());
5010 }
5011 let content = choice
5012 .pointer("/delta/content")
5013 .or_else(|| choice.pointer("/message/content"))
5014 .and_then(Value::as_str);
5015 if let Some(content) = content {
5016 message.push_str(content);
5017 choice_delta.content.push_str(content);
5018 }
5019 let reasoning_content = choice
5020 .pointer("/delta/reasoning_content")
5021 .or_else(|| choice.pointer("/message/reasoning_content"))
5022 .and_then(Value::as_str);
5023 if let Some(content) = reasoning_content {
5024 reasoning.push_str(content);
5025 choice_delta.reasoning.push_str(content);
5026 }
5027 choice_deltas.push(choice_delta);
5028 let streaming_calls = choice.pointer("/delta/tool_calls");
5029 let calls = streaming_calls
5030 .or_else(|| choice.pointer("/message/tool_calls"))
5031 .map(|calls| {
5032 calls
5033 .as_array()
5034 .ok_or_else(|| "provider tool_calls is not an array".to_owned())
5035 })
5036 .transpose()?;
5037 if let Some(calls) = calls {
5038 for (position, call) in calls.iter().enumerate() {
5039 let index = match call.get("index").and_then(Value::as_u64) {
5040 Some(index) => index,
5041 None if streaming_calls.is_some() => {
5042 return Err("streaming provider tool call has no index".to_owned());
5043 }
5044 None => u64::try_from(position)
5045 .map_err(|_| "provider tool-call index overflow".to_owned())?,
5046 };
5047 tool_call_deltas.push(ProviderToolCallDelta {
5048 choice_index,
5049 index,
5050 complete: false,
5051 id: call.get("id").and_then(Value::as_str).map(str::to_owned),
5052 name: call
5053 .pointer("/function/name")
5054 .and_then(Value::as_str)
5055 .map(str::to_owned),
5056 arguments: call
5057 .pointer("/function/arguments")
5058 .and_then(Value::as_str)
5059 .unwrap_or_default()
5060 .to_owned(),
5061 });
5062 }
5063 }
5064 }
5065 }
5066 let completion_reported = completion_tokens.is_some();
5067 let completion_tokens = completion_tokens.unwrap_or_else(|| {
5068 let mut estimated = pv_core::tokens::approx_tokens(&message)
5069 .saturating_add(pv_core::tokens::approx_tokens(&reasoning));
5070 for call in &tool_call_deltas {
5071 estimated = estimated
5072 .saturating_add(call.name.as_deref().map_or(0, pv_core::tokens::approx_tokens))
5073 .saturating_add(pv_core::tokens::approx_tokens(&call.arguments));
5074 }
5075 estimated
5076 });
5077 Ok(Some(ParsedProviderChunk {
5078 model_name,
5079 metrics: pv_events::EventMetrics {
5080 prompt_tokens,
5081 completion_tokens: Some(completion_tokens),
5082 cached_tokens,
5083 pruned_tokens: None,
5084 pruning_ratio_millis: None,
5085 },
5086 usage_reported,
5087 completion_reported,
5088 cost_usd_micros,
5089 cost_reported,
5090 has_choices,
5091 tool_call_deltas,
5092 choices: choice_deltas,
5093 }))
5094}
5095
5096pub(crate) fn provider_u64(value: Option<&Value>, field: &str) -> Result<Option<u64>, String> {
5097 let Some(value) = value else {
5098 return Ok(None);
5099 };
5100 if value.is_null() {
5104 return Ok(None);
5105 }
5106 let value = value
5107 .as_u64()
5108 .ok_or_else(|| format!("provider {field} is not a nonnegative integer"))?;
5109 if value > pv_core::error::JCS_SAFE_MAX {
5110 return Err(format!("provider {field} exceeds JCS-safe bounds"));
5111 }
5112 Ok(Some(value))
5113}
5114
5115pub(crate) fn normalize_tool_call_arguments(raw: &str) -> serde_json::Value {
5138 let candidate = if raw.is_empty() { "{}" } else { raw };
5139 serde_json::from_str(candidate).unwrap_or_else(|_| json!({"raw": raw}))
5140}
5141
5142pub(crate) fn map_finish_reason(native: &str) -> pv_events::StopReason {
5143 match native {
5144 "stop" | "stop_sequence" | "end_turn" => pv_events::StopReason::Stop,
5145 "length" | "max_tokens" | "model_context_window_exceeded" => pv_events::StopReason::MaxTokens,
5151 "tool_calls" | "function_call" | "tool_use" => pv_events::StopReason::ToolUse,
5152 "pause_turn" => pv_events::StopReason::ToolUse,
5159 "content_filter" | "refusal" => pv_events::StopReason::ContentFilter,
5167 _ => pv_events::StopReason::Other,
5168 }
5169}
5170
5171impl Drop for AbortFinalizingStream {
5172 fn drop(&mut self) {
5173 if self.completed {
5178 if self.ephemeral && !self.session.is_closed() && self.session.active_streams_count() <= 1 {
5190 let session = Arc::clone(&self.session);
5191 let finalizer = self.finalizer.clone();
5192 let session_id = session.id.clone();
5193 if let Ok(runtime) = tokio::runtime::Handle::try_current() {
5194 runtime.spawn(async move {
5195 let outcome = std::panic::AssertUnwindSafe(
5208 finalizer.close_session(session, StopReason::SessionClosed),
5209 )
5210 .catch_unwind()
5211 .await;
5212 match outcome {
5213 Ok(Ok(_)) => {}
5214 Ok(Err(error)) => {
5215 finalizer
5216 .metrics()
5217 .counter(
5218 "pv_ephemeral_close_failures_total",
5219 "Auto-close of a completed ephemeral one-shot failed; \
5220 the session is left open until the idle sweeper reaps it",
5221 )
5222 .inc();
5223 tracing::warn!(
5224 %error,
5225 session = %session_id,
5226 "auto-close of completed ephemeral session failed"
5227 );
5228 }
5229 Err(panic) => {
5230 let msg = panic
5231 .downcast_ref::<&'static str>()
5232 .copied()
5233 .or_else(|| panic.downcast_ref::<String>().map(String::as_str))
5234 .unwrap_or("panic payload was not a string");
5235 finalizer
5236 .metrics()
5237 .counter(
5238 "pv_ephemeral_close_panics_total",
5239 "Auto-close of a completed ephemeral one-shot panicked \
5240 and was caught; the session is left open until the idle \
5241 sweeper reaps it",
5242 )
5243 .inc();
5244 tracing::warn!(
5245 session = %session_id,
5246 panic = %msg,
5247 "auto-close of completed ephemeral session PANICKED (caught)"
5248 );
5249 }
5250 }
5251 });
5252 }
5253 }
5256 return;
5257 }
5258 let is_closed = self.session.is_closed();
5278 let budget_incomplete = self.pending_budget.is_some();
5279 if let Some(pending) = self.pending_budget.take() {
5291 pending.task.abort();
5292 self.charged_completion_tokens = self.charged_completion_tokens.saturating_sub(pending.delta);
5303 }
5304 let (flush_reason, budget_delta) = match self.flush_protocol_buffer() {
5305 Ok(delta) => (None, delta),
5306 Err(error) => {
5307 tracing::warn!(%error, "provider stream flush failed on drop");
5308 (Some(format!("provider stream flush failed: {error}")), 0u64)
5309 }
5310 };
5311 self.charged_completion_tokens = self.charged_completion_tokens.saturating_sub(budget_delta);
5317 let response_failure: Option<String> = if budget_incomplete {
5324 Some("client disconnected before response budget check completed".to_owned())
5325 } else if let Some(reason) = flush_reason {
5326 Some(reason)
5327 } else if budget_delta > 0 {
5328 Some(format!(
5329 "client disconnected with {budget_delta} unbilled response bytes"
5330 ))
5331 } else {
5332 None
5333 };
5334 if let Err(error) = self.submit_response_capture(
5335 response_failure.map(|reason| (pv_events::StopReason::Other, reason)),
5337 ) {
5338 tracing::warn!(%error, "response-capture submit failed on drop");
5339 if !is_closed {
5340 self.session.mark_capture_failed();
5344 }
5345 }
5346 let nothing_captured = !self.saw_chunk && self.captured_bytes == 0;
5371 let other_streams = self.session.active_streams_count() > 1;
5372 if self.ephemeral && !is_closed && nothing_captured && !other_streams {
5373 let session = Arc::clone(&self.session);
5374 let finalizer = self.finalizer.clone();
5375 let session_id = session.id.clone();
5376 let stop_reason = StopReason::Other;
5377 match tokio::runtime::Handle::try_current() {
5378 Ok(runtime) => {
5379 runtime.spawn(async move {
5392 let outcome =
5393 std::panic::AssertUnwindSafe(finalizer.close_session(session, stop_reason))
5394 .catch_unwind()
5395 .await;
5396 match outcome {
5397 Ok(Ok(_)) => {}
5398 Ok(Err(error)) => {
5399 finalizer
5400 .metrics()
5401 .counter(
5402 "pv_stream_abort_close_failures_total",
5403 "Background close after a stream abort failed; \
5404 the session is left open until the idle sweeper reaps it",
5405 )
5406 .inc();
5407 tracing::warn!(
5408 %error,
5409 session = %session_id,
5410 "background close on stream abort failed"
5411 );
5412 }
5413 Err(panic) => {
5414 let msg = panic
5415 .downcast_ref::<&'static str>()
5416 .copied()
5417 .or_else(|| panic.downcast_ref::<String>().map(String::as_str))
5418 .unwrap_or("panic payload was not a string");
5419 finalizer
5420 .metrics()
5421 .counter(
5422 "pv_stream_abort_panics_total",
5423 "Background close after a stream abort panicked and was \
5424 caught; the session is left open until the idle sweeper \
5425 reaps it",
5426 )
5427 .inc();
5428 tracing::warn!(
5429 session = %session_id,
5430 panic = %msg,
5431 "background close on stream abort PANICKED (caught)"
5432 );
5433 }
5434 }
5435 });
5436 }
5437 Err(error) => {
5438 self.finalizer
5444 .metrics()
5445 .counter(
5446 "pv_stream_abort_no_runtime_total",
5447 "Stream abort observed no tokio runtime; capture marked failed for \
5448 reconciler retry",
5449 )
5450 .inc();
5451 tracing::warn!(
5452 %error,
5453 session = %session_id,
5454 "no tokio runtime available for stream-abort close; marking capture failed for reconciler retry"
5455 );
5456 self.session.mark_capture_failed();
5457 }
5458 }
5459 }
5460 }
5461}
5462
5463#[cfg(test)]
5464mod tests {
5465 #![allow(
5466 clippy::expect_used,
5467 clippy::indexing_slicing,
5468 clippy::panic,
5469 clippy::unwrap_used
5470 )]
5471
5472 use super::*;
5473 use axum::http::Request;
5474 use pv_bridge::{BusError, EventBus, PublishAck, StoredEvent};
5475 use pv_receipts::Ed25519Signer;
5476 use pv_sandbox::{Sandbox, SandboxConfig};
5477 use pv_state::{InMemoryStore, Spend, StateError, StateStore};
5478 use std::time::Duration;
5479 use tower::ServiceExt;
5480
5481 #[test]
5495 fn openai_compatibility_documents_every_client_facing_route() {
5496 let doc = include_str!("../../../docs/reference/OPENAI-COMPATIBILITY.md");
5497 for (route, must_contain) in [
5510 ("/v1/chat/completions", "`/v1/chat/completions`"),
5511 ("/v1/mcp", "`/v1/mcp`"),
5512 ("/mcp (bare, MCP SDK convention)", "bare `/mcp`"),
5516 ("/v1/sessions/{id}/close", "`/v1/sessions/{id}/close`"),
5517 ("/v1/sessions/{id}/promote", "`/v1/sessions/{id}/promote`"),
5518 ] {
5519 assert!(
5520 doc.contains(must_contain),
5521 "route {route} is registered by build_router but the doc is missing \
5522 its distinguishing marker {must_contain:?} — an integration \
5523 engineer following the doc will not know it exists. Add a row to \
5524 the route table (or the /v1 vs /api/v1 split section) and re-run."
5525 );
5526 }
5527 }
5528
5529 #[tokio::test]
5537 async fn unsupported_openai_routes_stay_404() {
5538 let doc = include_str!("../../../docs/reference/OPENAI-COMPATIBILITY.md");
5539 let directory = tempfile::tempdir().unwrap();
5540 let (state, _reconciler) = test_state(directory.path()).await;
5541 let app = build_router(state);
5542 for (route, marker) in [
5543 ("/v1/embeddings", "`/v1/embeddings`"),
5544 ("/v1/audio/speech", "`/v1/audio/*`"),
5545 ("/v1/images/generations", "`/v1/images/*`"),
5546 ("/v1/completions", "`/v1/completions`"),
5547 ("/v1/models", "`/v1/models`"),
5548 ("/v1/files", "`/v1/files`"),
5549 ("/v1/moderations", "`/v1/moderations`"),
5550 ("/v1/assistants", "`/v1/assistants`"),
5551 ("/v1/fine_tuning/jobs", "`/v1/fine_tuning/*`"),
5552 ] {
5553 assert!(
5554 doc.contains(marker),
5555 "the Unsupported list lost {marker:?} — if {route} is now \
5556 implemented, update OPENAI-COMPATIBILITY.md; if not, restore it"
5557 );
5558 let response = app
5559 .clone()
5560 .oneshot(
5561 Request::builder()
5562 .method("POST")
5563 .uri(route)
5564 .header("content-type", "application/json")
5565 .body(Body::from("{}"))
5566 .unwrap(),
5567 )
5568 .await
5569 .unwrap();
5570 assert_eq!(
5571 response.status(),
5572 StatusCode::NOT_FOUND,
5573 "{route} no longer 404s — OPENAI-COMPATIBILITY.md's Unsupported \
5574 section promises it does; update the doc with the new behavior"
5575 );
5576 }
5577 }
5578
5579 #[test]
5593 fn product_console_pins_match_real_schema_and_taxonomy() {
5594 let console_js = include_str!("../../../docs/app/pitch/console.js");
5599 let console_html = include_str!("../../../docs/app/pitch/index.html");
5600 let combined = format!("{console_js}\n{console_html}");
5601
5602 assert_eq!(
5606 pv_receipts::RECEIPT_VERSION,
5607 2,
5608 "receipt version bumped from 2; update docs/app/pitch/console.js's \
5609 receiptHtml() to render the new version",
5610 );
5611 assert!(
5612 combined.contains(r#""receipt_version""#) && combined.contains(">2<"),
5613 "the console must render receipt_version = 2 to match \
5614 pv_receipts::RECEIPT_VERSION",
5615 );
5616
5617 for key in [
5621 "receipt_version",
5622 "receipt_id",
5623 "session_id",
5624 "issued_at",
5625 "issued_at_iso",
5626 "ai_agent",
5627 "subject",
5628 "tool_calls",
5629 "cost",
5630 "stop_reason_id",
5631 "stop_reason",
5632 "key_id",
5633 "public_key_b64",
5634 "signature_b64",
5635 ] {
5636 assert!(
5637 combined.contains(&format!(r#""{key}""#)),
5638 "receipt field `{key}` is not rendered by the product \
5639 console — either the field was renamed (update the \
5640 console) or the console was truncated (update this test)",
5641 );
5642 }
5643
5644 for cost_key in [
5649 "prompt_tokens",
5650 "completion_tokens",
5651 "cached_tokens",
5652 "cost_usd_micros",
5653 ] {
5654 assert!(
5655 combined.contains(&format!(r#""{cost_key}""#)),
5656 "cost field `{cost_key}` is not rendered by the console",
5657 );
5658 }
5659
5660 assert!(
5664 combined.contains("cost_usd_micros"),
5665 "console tampering demo relies on cost_usd_micros being \
5666 the integer field it changes; rename would break the demo",
5667 );
5668
5669 assert_eq!(
5673 pv_events::StopReason::BudgetExceeded.id(),
5674 92,
5675 "BudgetExceeded id changed from 92; update the console's \
5676 sample receipt receiptHtml()",
5677 );
5678 assert_eq!(
5679 pv_events::StopReason::BudgetExceeded.caption(),
5680 "Budget Exceeded",
5681 "BudgetExceeded caption changed; update the console",
5682 );
5683 assert!(
5684 combined.contains(r#""stop_reason_id""#) && combined.contains(">92<"),
5685 "console must render stop_reason_id = 92 to match \
5686 StopReason::BudgetExceeded",
5687 );
5688 assert!(
5689 combined.contains("Budget Exceeded"),
5690 "console must render the 'Budget Exceeded' caption for the \
5691 sample receipt",
5692 );
5693
5694 for topic in ["agent.tool_call", "agent.session"] {
5699 assert!(
5700 combined.contains(topic),
5701 "OCSF topic `{topic}` is not rendered by the console but is \
5702 shown as the file that lands in the operator's SIEM",
5703 );
5704 }
5705
5706 assert!(
5711 combined.contains("HTTP 403"),
5712 "console must render the 403 policy-refusal status the \
5713 pipeline emits (deliberately not 429, which SDKs auto-retry)",
5714 );
5715
5716 assert!(
5719 combined.contains("127.0.0.1:8484"),
5720 "console must render the default listen address 127.0.0.1:8484; \
5721 raise or lower the default and the setup/curl blocks lie",
5722 );
5723
5724 assert!(
5733 combined.contains("budget.max_tool_calls"),
5734 "console's TOML preview must render the per-tool cap as the \
5735 `[budget.max_tool_calls]` table the real BudgetSpec requires",
5736 );
5737 assert!(
5738 !console_js.contains(r#"max_tool_calls</span> = "#),
5739 "console's TOML preview renders `max_tool_calls = <scalar>`, \
5740 which the real config parser rejects (BudgetSpec.max_tool_calls \
5741 is a BTreeMap; the valid shape is a [budget.max_tool_calls] table)",
5742 );
5743 }
5744
5745 struct NullBus;
5746
5747 struct SlowStore {
5748 inner: InMemoryStore,
5749 calls: Arc<std::sync::atomic::AtomicUsize>,
5750 }
5751
5752 impl StateStore for SlowStore {
5753 fn add(&self, key: &str, delta: u64) -> Result<u64, StateError> {
5754 self.inner.add(key, delta)
5755 }
5756
5757 fn get(&self, key: &str) -> Result<u64, StateError> {
5758 self.inner.get(key)
5759 }
5760
5761 fn try_spend(&self, key: &str, amount: u64, limit: u64) -> Result<bool, StateError> {
5762 self.calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
5763 std::thread::sleep(Duration::from_millis(100));
5764 self.inner.try_spend(key, amount, limit)
5765 }
5766
5767 fn try_spend_many(&self, spends: &[Spend]) -> Result<pv_state::TrySpendOutcome, StateError> {
5768 self.calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
5777 std::thread::sleep(Duration::from_millis(100));
5778 self.inner.try_spend_many(spends)
5779 }
5780
5781 fn remove(&self, key: &str) {
5782 self.inner.remove(key);
5783 }
5784 }
5785
5786 impl EventBus for NullBus {
5787 fn publish(&self, topic: &str, _key: &str, _value: &Value) -> Result<PublishAck, BusError> {
5788 Ok(PublishAck {
5789 topic: topic.to_owned(),
5790 partition: 0,
5791 offset: 0,
5792 })
5793 }
5794
5795 fn fetch(
5796 &self,
5797 _topic: &str,
5798 _partition: u32,
5799 _offset: u64,
5800 _max: usize,
5801 ) -> Result<Vec<StoredEvent>, BusError> {
5802 Ok(Vec::new())
5803 }
5804
5805 fn partitions(&self, _topic: &str) -> Result<u32, BusError> {
5806 Ok(1)
5807 }
5808
5809 fn topics(&self) -> Vec<String> {
5810 pv_events::EventClass::all()
5811 .iter()
5812 .map(|class| class.topic().to_owned())
5813 .collect()
5814 }
5815 }
5816
5817 async fn mock_chat(Json(payload): Json<Value>) -> Response {
5818 if payload.get("model").and_then(Value::as_str) == Some("empty-error") {
5819 let mut response = Response::new(Body::empty());
5820 *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
5821 response.headers_mut().insert(
5822 axum::http::header::CONTENT_TYPE,
5823 HeaderValue::from_static("application/json"),
5824 );
5825 return response;
5826 }
5827 if payload.get("model").and_then(Value::as_str) == Some("json-error") {
5828 return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({}))).into_response();
5829 }
5830 if payload.get("model").and_then(Value::as_str) == Some("empty-success") {
5831 return Json(json!({})).into_response();
5832 }
5833 if payload.get("model").and_then(Value::as_str) == Some("split-sse") {
5834 let frame = "data: {\"choices\":[{\"delta\":{\"content\":\"héllo\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n"
5840 .as_bytes()
5841 .to_vec();
5842 let split = frame.iter().position(|byte| *byte == 0xc3).unwrap() + 1;
5843 let chunks = vec![
5844 Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(&frame[..split])),
5845 Ok(Bytes::copy_from_slice(&frame[split..])),
5846 ];
5847 let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
5848 response.headers_mut().insert(
5849 axum::http::header::CONTENT_TYPE,
5850 HeaderValue::from_static("text/event-stream"),
5851 );
5852 return response;
5853 }
5854 if payload.get("model").and_then(Value::as_str) == Some("truncated-stream") {
5855 let mut response = Response::new(Body::from(
5859 "data: {\"choices\":[{\"delta\":{\"content\":\"partial ans\"}}]}\n\n",
5860 ));
5861 response.headers_mut().insert(
5862 axum::http::header::CONTENT_TYPE,
5863 HeaderValue::from_static("text/event-stream"),
5864 );
5865 return response;
5866 }
5867 if payload.get("model").and_then(Value::as_str) == Some("severed-stream") {
5868 let chunks = futures::stream::unfold(0u8, |state| async move {
5876 match state {
5877 0 => Some((
5878 Ok::<_, std::io::Error>(Bytes::from_static(
5879 b"data: {\"choices\":[{\"delta\":{\"content\":\"partial ans\"}}]}\n\n",
5880 )),
5881 1,
5882 )),
5883 1 => {
5884 tokio::time::sleep(Duration::from_millis(150)).await;
5885 Some((Err(std::io::Error::other("connection reset by peer")), 2))
5886 }
5887 _ => None,
5888 }
5889 });
5890 let mut response = Response::new(Body::from_stream(chunks));
5891 response.headers_mut().insert(
5892 axum::http::header::CONTENT_TYPE,
5893 HeaderValue::from_static("text/event-stream"),
5894 );
5895 return response;
5896 }
5897 if payload.get("model").and_then(Value::as_str) == Some("provisional-usage") {
5898 let chunks = vec![
5899 Ok::<_, std::convert::Infallible>(Bytes::from_static(
5900 b"data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n",
5901 )),
5902 Ok(Bytes::from_static(
5903 b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2}}\n\n",
5904 )),
5905 ];
5906 let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
5907 response.headers_mut().insert(
5908 axum::http::header::CONTENT_TYPE,
5909 HeaderValue::from_static("text/event-stream"),
5910 );
5911 return response;
5912 }
5913 if payload.get("model").and_then(Value::as_str) == Some("tool-no-usage") {
5914 let arguments = "x".repeat(4_096);
5915 let event = format!(
5916 "data: {{\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":0,\"function\":{{\"name\":\"read\",\"arguments\":\"{arguments}\"}}}}]}}}}]}}\n\n"
5917 );
5918 let mut response = Response::new(Body::from(event));
5919 response.headers_mut().insert(
5920 axum::http::header::CONTENT_TYPE,
5921 HeaderValue::from_static("text/event-stream"),
5922 );
5923 return response;
5924 }
5925 if payload.get("model").and_then(Value::as_str) == Some("tool-oob-index") {
5926 let event = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":18446744073709551615,\"function\":{\"name\":\"read\",\"arguments\":\"{}\"}}]}}]}\n\n";
5927 let mut response = Response::new(Body::from(event));
5928 response.headers_mut().insert(
5929 axum::http::header::CONTENT_TYPE,
5930 HeaderValue::from_static("text/event-stream"),
5931 );
5932 return response;
5933 }
5934 if payload.get("model").and_then(Value::as_str) == Some("tool-conflicting-id") {
5935 let event = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-a\",\"function\":{\"name\":\"read\",\"arguments\":\"{\"}},{\"index\":0,\"id\":\"call-b\",\"function\":{\"arguments\":\"}\"}}]}}]}\n\n";
5939 let mut response = Response::new(Body::from(event));
5940 response.headers_mut().insert(
5941 axum::http::header::CONTENT_TYPE,
5942 HeaderValue::from_static("text/event-stream"),
5943 );
5944 return response;
5945 }
5946 if payload.get("model").and_then(Value::as_str) == Some("tool-conflicting-name") {
5947 let event = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call-a\",\"function\":{\"name\":\"read\",\"arguments\":\"{\"}},{\"index\":0,\"id\":\"call-a\",\"function\":{\"name\":\"db_write\",\"arguments\":\"}\"}}]}}]}\n\n";
5950 let mut response = Response::new(Body::from(event));
5951 response.headers_mut().insert(
5952 axum::http::header::CONTENT_TYPE,
5953 HeaderValue::from_static("text/event-stream"),
5954 );
5955 return response;
5956 }
5957 if payload.get("model").and_then(Value::as_str) == Some("overcharged-completion") {
5958 let chunks = vec![
5962 Ok::<_, std::convert::Infallible>(Bytes::from_static(
5963 b"data: {\"choices\":[{\"delta\":{\"content\":\"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\"},\"finish_reason\":\"stop\"}]}\n\n",
5964 )),
5965 Ok(Bytes::from_static(
5966 b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1}}\n\n",
5967 )),
5968 ];
5969 let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
5970 response.headers_mut().insert(
5971 axum::http::header::CONTENT_TYPE,
5972 HeaderValue::from_static("text/event-stream"),
5973 );
5974 return response;
5975 }
5976 if payload.get("model").and_then(Value::as_str) == Some("regressive-usage") {
5977 let chunks = vec![
5978 Ok::<_, std::convert::Infallible>(Bytes::from_static(
5979 b"data: {\"choices\":[{\"delta\":{\"content\":\"first\"}}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":100,\"prompt_tokens_details\":{\"cached_tokens\":5}},\"cost_usd\":0.01}\n\n",
5980 )),
5981 Ok(Bytes::from_static(
5982 b"data: {\"choices\":[{\"delta\":{\"content\":\"second\"}}],\"usage\":{\"prompt_tokens\":9,\"completion_tokens\":1,\"prompt_tokens_details\":{\"cached_tokens\":4}},\"cost_usd\":0.001}\n\n",
5983 )),
5984 ];
5985 let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
5986 response.headers_mut().insert(
5987 axum::http::header::CONTENT_TYPE,
5988 HeaderValue::from_static("text/event-stream"),
5989 );
5990 return response;
5991 }
5992 if payload.get("model").and_then(Value::as_str) == Some("split-json") {
5993 let body = serde_json::to_vec(&json!({
5994 "model": "split-json",
5995 "choices": [{"message": {"role": "assistant", "content": "héllo"}}],
5996 "usage": {"prompt_tokens": 5, "completion_tokens": 2}
5997 }))
5998 .unwrap();
5999 let split = body.iter().position(|byte| *byte == 0xc3).unwrap() + 1;
6000 let chunks = vec![
6001 Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(&body[..split])),
6002 Ok(Bytes::copy_from_slice(&body[split..])),
6003 ];
6004 let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
6005 response.headers_mut().insert(
6006 axum::http::header::CONTENT_TYPE,
6007 HeaderValue::from_static("application/json"),
6008 );
6009 return response;
6010 }
6011 if payload.get("model").and_then(Value::as_str) == Some("gzip-encoded") {
6012 let mut response = Response::new(Body::from("compressed-bytes"));
6015 response.headers_mut().insert(
6016 axum::http::header::CONTENT_TYPE,
6017 HeaderValue::from_static("application/json"),
6018 );
6019 response.headers_mut().insert(
6020 axum::http::header::CONTENT_ENCODING,
6021 HeaderValue::from_static("gzip"),
6022 );
6023 return response;
6024 }
6025 if payload.get("model").and_then(Value::as_str) == Some("malformed-json") {
6026 let mut response = Response::new(Body::from("{\"choices\":["));
6027 response.headers_mut().insert(
6028 axum::http::header::CONTENT_TYPE,
6029 HeaderValue::from_static("application/json"),
6030 );
6031 return response;
6032 }
6033 if payload.get("stream").and_then(Value::as_bool).unwrap_or(false) {
6034 let event = "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":3,\"cost_usd\":0.00125,\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\ndata: [DONE]\n\n";
6035 let bytes = event.as_bytes();
6036 let chunks = vec![
6037 Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(&bytes[..19])),
6038 Ok(Bytes::copy_from_slice(&bytes[19..97])),
6039 Ok(Bytes::copy_from_slice(&bytes[97..])),
6040 ];
6041 let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
6042 response.headers_mut().insert(
6043 axum::http::header::CONTENT_TYPE,
6044 HeaderValue::from_static("text/event-stream"),
6045 );
6046 response
6047 } else {
6048 Json(json!({"choices": [{"message": {"role": "assistant", "content": "hello"}}]})).into_response()
6049 }
6050 }
6051
6052 async fn counting_tool(
6053 State(calls): State<Arc<std::sync::atomic::AtomicUsize>>,
6054 body: Bytes,
6055 ) -> Response {
6056 calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
6057 (StatusCode::OK, body).into_response()
6058 }
6059
6060 async fn redirecting_tool(State(calls): State<Arc<std::sync::atomic::AtomicUsize>>) -> Response {
6061 calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
6062 (
6063 StatusCode::TEMPORARY_REDIRECT,
6064 [(axum::http::header::LOCATION, "/effect")],
6065 )
6066 .into_response()
6067 }
6068
6069 async fn redirected_effect(State(calls): State<Arc<std::sync::atomic::AtomicUsize>>) -> Response {
6070 calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
6071 StatusCode::OK.into_response()
6072 }
6073
6074 struct HeldToolState {
6075 arrived: std::sync::atomic::AtomicBool,
6076 release: tokio::sync::Notify,
6077 }
6078
6079 async fn held_tool(State(state): State<Arc<HeldToolState>>, body: Bytes) -> Response {
6080 state.arrived.store(true, std::sync::atomic::Ordering::Release);
6081 state.release.notified().await;
6082 (StatusCode::OK, body).into_response()
6083 }
6084
6085 async fn test_state(spool: &std::path::Path) -> (AppState, tokio::task::JoinHandle<()>) {
6086 test_state_with_token_cap(spool, None).await
6087 }
6088
6089 async fn mock_anthropic_chat() -> Response {
6092 let frames = concat!(
6093 "event: message_start\n",
6094 "data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-sonnet-4-5\",\"usage\":{\"input_tokens\":25,\"output_tokens\":1}}}\n\n",
6095 "event: ping\n",
6096 "data: {\"type\": \"ping\"}\n\n",
6097 "event: content_block_start\n",
6098 "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
6099 "event: content_block_delta\n",
6100 "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello from anthropic\"}}\n\n",
6101 "event: content_block_stop\n",
6102 "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
6103 "event: message_delta\n",
6104 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":15}}\n\n",
6105 "event: message_stop\n",
6106 "data: {\"type\":\"message_stop\"}\n\n",
6107 );
6108 let mut response = Response::new(Body::from(frames));
6109 response.headers_mut().insert(
6110 axum::http::header::CONTENT_TYPE,
6111 HeaderValue::from_static("text/event-stream"),
6112 );
6113 response
6114 }
6115
6116 async fn test_state_anthropic(spool: &std::path::Path) -> (AppState, tokio::task::JoinHandle<()>) {
6117 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6118 let address = listener.local_addr().unwrap();
6119 let provider = Router::new().route("/v1/messages", post(mock_anthropic_chat));
6120 let server = tokio::spawn(async move {
6121 axum::serve(listener, provider).await.unwrap();
6122 });
6123 let mut config = crate::config::HarnessConfig::for_tests(
6124 &format!("http://{address}"),
6125 &spool.to_string_lossy(),
6126 &spool.to_string_lossy(),
6127 );
6128 config.provider = "anthropic".to_owned();
6129 let state = AppState::new(
6130 config,
6131 Arc::new(InMemoryStore::new()),
6132 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6133 Arc::new(NullBus),
6134 None,
6135 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
6136 )
6137 .unwrap();
6138 (state, server)
6139 }
6140
6141 #[tokio::test]
6145 async fn oversize_body_413_is_json_and_names_the_knob() {
6146 let directory = tempfile::tempdir().unwrap();
6147 let (state, provider) = test_state(directory.path()).await;
6148 let limit = state.config.max_request_bytes;
6149 let app = build_router(state);
6150 let oversize = "x".repeat(limit + 1);
6151 let response = app
6152 .oneshot(
6153 Request::builder()
6154 .method("POST")
6155 .uri("/v1/chat/completions")
6156 .header("content-type", "application/json")
6157 .header("x-pv-session", "too-big")
6158 .body(Body::from(oversize))
6159 .unwrap(),
6160 )
6161 .await
6162 .unwrap();
6163 assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
6164 let content_type = response
6165 .headers()
6166 .get(axum::http::header::CONTENT_TYPE)
6167 .and_then(|value| value.to_str().ok())
6168 .unwrap_or_default()
6169 .to_owned();
6170 assert!(
6171 content_type.starts_with("application/json"),
6172 "413 must honour the JSON error contract, got {content_type}"
6173 );
6174 let body: Value = serde_json::from_slice(
6175 &axum::body::to_bytes(response.into_body(), 64 * 1024)
6176 .await
6177 .unwrap(),
6178 )
6179 .unwrap();
6180 let message = body["error"]["message"].as_str().unwrap();
6181 assert!(
6182 message.contains("max_request_bytes") && message.contains(&limit.to_string()),
6183 "413 must name the knob and the configured limit: {message}"
6184 );
6185 provider.abort();
6186 }
6187
6188 async fn mock_gemini_chat() -> Response {
6191 let frames = concat!(
6192 "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello from gemini\"}],\"role\":\"model\"},\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":25,\"candidatesTokenCount\":1},\"modelVersion\":\"gemini-2.0-flash\"}\n\n",
6193 "data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"!\"}]},\"finishReason\":\"STOP\",\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":25,\"candidatesTokenCount\":15}}\n\n",
6194 );
6195 let mut response = Response::new(Body::from(frames));
6196 response.headers_mut().insert(
6197 axum::http::header::CONTENT_TYPE,
6198 HeaderValue::from_static("text/event-stream"),
6199 );
6200 response
6201 }
6202
6203 async fn test_state_gemini(spool: &std::path::Path) -> (AppState, tokio::task::JoinHandle<()>) {
6204 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6205 let address = listener.local_addr().unwrap();
6206 let provider = Router::new().route(
6207 "/v1beta/models/mock:streamGenerateContent",
6208 post(mock_gemini_chat),
6209 );
6210 let server = tokio::spawn(async move {
6211 axum::serve(listener, provider).await.unwrap();
6212 });
6213 let mut config = crate::config::HarnessConfig::for_tests(
6214 &format!("http://{address}"),
6215 &spool.to_string_lossy(),
6216 &spool.to_string_lossy(),
6217 );
6218 config.provider = "gemini".to_owned();
6219 let state = AppState::new(
6220 config,
6221 Arc::new(InMemoryStore::new()),
6222 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6223 Arc::new(NullBus),
6224 None,
6225 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
6226 )
6227 .unwrap();
6228 (state, server)
6229 }
6230
6231 #[tokio::test]
6236 async fn gemini_session_records_the_same_audit_shape() {
6237 let directory = tempfile::tempdir().unwrap();
6238 let (state, provider) = test_state_gemini(directory.path()).await;
6239 let app = build_router(state.clone());
6240 let response = app.clone().oneshot(chat_request("gemini-flow")).await.unwrap();
6241 assert_eq!(response.status(), StatusCode::OK);
6242 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
6243 .await
6244 .unwrap();
6245 let body = String::from_utf8_lossy(&body);
6246 assert!(
6247 body.contains("hello from gemini"),
6248 "the translated stream must carry the upstream text: {body}"
6249 );
6250 assert!(
6251 body.contains("\"object\":\"chat.completion.chunk\"")
6252 && !body.contains("candidates")
6253 && body.trim_end().ends_with("data: [DONE]"),
6254 "an OpenAI client must receive OpenAI chunks, not the native stream: {body}"
6255 );
6256
6257 let session = state.sessions.get("gemini-flow").unwrap();
6258 session.wait_for_worker_jobs().await;
6259 use std::sync::atomic::Ordering;
6260 assert_eq!(session.totals.prompt_tokens.load(Ordering::Acquire), 25);
6261 assert_eq!(session.totals.completion_tokens.load(Ordering::Acquire), 15);
6262
6263 let close = app
6264 .clone()
6265 .oneshot(
6266 Request::builder()
6267 .method("POST")
6268 .uri("/v1/sessions/gemini-flow/close")
6269 .body(Body::empty())
6270 .unwrap(),
6271 )
6272 .await
6273 .unwrap();
6274 assert_eq!(close.status(), StatusCode::OK);
6275 let promoted = app
6276 .clone()
6277 .oneshot(
6278 Request::builder()
6279 .method("POST")
6280 .uri("/v1/sessions/gemini-flow/promote")
6281 .body(Body::empty())
6282 .unwrap(),
6283 )
6284 .await
6285 .unwrap();
6286 assert_eq!(promoted.status(), StatusCode::OK);
6287 let receipt: pv_receipts::Receipt = serde_json::from_slice(
6288 &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
6289 .await
6290 .unwrap(),
6291 )
6292 .unwrap();
6293 receipt.verify_embedded().unwrap();
6294 assert_eq!(receipt.body.stop_reason_id, pv_events::StopReason::Stop.id());
6295 assert_eq!(receipt.body.cost.prompt_tokens, 25);
6296 assert_eq!(receipt.body.cost.completion_tokens, 15);
6297 provider.abort();
6298 }
6299
6300 async fn mock_gemini_tool_chat() -> Response {
6306 let frames = concat!(
6307 "data: {\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"Paris\"}}}]},\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":15,\"candidatesTokenCount\":10},\"modelVersion\":\"gemini-2.0-flash\"}\n\n",
6308 "data: {\"candidates\":[{\"content\":{\"parts\":[]},\"finishReason\":\"STOP\",\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":15,\"candidatesTokenCount\":10}}\n\n",
6309 );
6310 let mut response = Response::new(Body::from(frames));
6311 response.headers_mut().insert(
6312 axum::http::header::CONTENT_TYPE,
6313 HeaderValue::from_static("text/event-stream"),
6314 );
6315 response
6316 }
6317
6318 async fn test_state_gemini_tool(spool: &std::path::Path) -> (AppState, tokio::task::JoinHandle<()>) {
6319 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6320 let address = listener.local_addr().unwrap();
6321 let provider = Router::new().route(
6322 "/v1beta/models/mock:streamGenerateContent",
6323 post(mock_gemini_tool_chat),
6324 );
6325 let server = tokio::spawn(async move {
6326 axum::serve(listener, provider).await.unwrap();
6327 });
6328 let mut config = crate::config::HarnessConfig::for_tests(
6329 &format!("http://{address}"),
6330 &spool.to_string_lossy(),
6331 &spool.to_string_lossy(),
6332 );
6333 config.provider = "gemini".to_owned();
6334 let state = AppState::new(
6335 config,
6336 Arc::new(InMemoryStore::new()),
6337 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6338 Arc::new(NullBus),
6339 None,
6340 Arc::new(Ed25519Signer::from_seed(&[13; 32])),
6341 )
6342 .unwrap();
6343 (state, server)
6344 }
6345
6346 #[tokio::test]
6360 async fn gemini_tool_call_with_finish_stop_attests_tool_use() {
6361 let directory = tempfile::tempdir().unwrap();
6362 let (state, provider) = test_state_gemini_tool(directory.path()).await;
6363 let app = build_router(state.clone());
6364 let response = app.clone().oneshot(chat_request("gemini-tool")).await.unwrap();
6365 assert_eq!(response.status(), StatusCode::OK);
6366 let _ = axum::body::to_bytes(response.into_body(), 64 * 1024).await;
6367 let session = state.sessions.get("gemini-tool").unwrap();
6368 session.wait_for_worker_jobs().await;
6369
6370 let close = app
6373 .clone()
6374 .oneshot(
6375 Request::builder()
6376 .method("POST")
6377 .uri("/v1/sessions/gemini-tool/close")
6378 .body(Body::empty())
6379 .unwrap(),
6380 )
6381 .await
6382 .unwrap();
6383 assert_eq!(close.status(), StatusCode::OK);
6384 let promoted = app
6385 .clone()
6386 .oneshot(
6387 Request::builder()
6388 .method("POST")
6389 .uri("/v1/sessions/gemini-tool/promote")
6390 .body(Body::empty())
6391 .unwrap(),
6392 )
6393 .await
6394 .unwrap();
6395 assert_eq!(promoted.status(), StatusCode::OK);
6396 let receipt: pv_receipts::Receipt = serde_json::from_slice(
6397 &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
6398 .await
6399 .unwrap(),
6400 )
6401 .unwrap();
6402 receipt.verify_embedded().unwrap();
6403 assert_eq!(
6404 receipt.body.stop_reason_id,
6405 pv_events::StopReason::ToolUse.id(),
6406 "a Gemini session that ended with a functionCall part must attest \
6407 StopReason::ToolUse (id {}), not Stop (id {}) — the adapter must \
6408 normalize provider-specific quirks so downstream stop-reason \
6409 analytics are provider-agnostic. Received: stop_reason_id={}, \
6410 stop_reason={:?}.",
6411 pv_events::StopReason::ToolUse.id(),
6412 pv_events::StopReason::Stop.id(),
6413 receipt.body.stop_reason_id,
6414 receipt.body.stop_reason,
6415 );
6416 provider.abort();
6417 }
6418
6419 async fn mock_gemini_parallel_tool_chat() -> Response {
6428 let frames = concat!(
6429 "data: {\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_weather\",\"args\":{\"city\":\"Paris\"}}}]},\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":15,\"candidatesTokenCount\":10},\"modelVersion\":\"gemini-2.0-flash\"}\n\n",
6430 "data: {\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"name\":\"get_time\",\"args\":{\"city\":\"London\"}}}]},\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":15,\"candidatesTokenCount\":20}}\n\n",
6431 "data: {\"candidates\":[{\"content\":{\"parts\":[]},\"finishReason\":\"STOP\",\"index\":0}],\"usageMetadata\":{\"promptTokenCount\":15,\"candidatesTokenCount\":20}}\n\n",
6432 );
6433 let mut response = Response::new(Body::from(frames));
6434 response.headers_mut().insert(
6435 axum::http::header::CONTENT_TYPE,
6436 HeaderValue::from_static("text/event-stream"),
6437 );
6438 response
6439 }
6440
6441 #[tokio::test]
6442 async fn gemini_parallel_function_calls_across_chunks_attest_as_two_calls() {
6443 let directory = tempfile::tempdir().unwrap();
6444 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6445 let address = listener.local_addr().unwrap();
6446 let upstream = Router::new().route(
6447 "/v1beta/models/mock:streamGenerateContent",
6448 post(mock_gemini_parallel_tool_chat),
6449 );
6450 let provider = tokio::spawn(async move {
6451 axum::serve(listener, upstream).await.unwrap();
6452 });
6453 let mut config = crate::config::HarnessConfig::for_tests(
6454 &format!("http://{address}"),
6455 &directory.path().to_string_lossy(),
6456 &directory.path().to_string_lossy(),
6457 );
6458 config.provider = "gemini".to_owned();
6459 let state = AppState::new(
6460 config,
6461 Arc::new(InMemoryStore::new()),
6462 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6463 Arc::new(NullBus),
6464 None,
6465 Arc::new(Ed25519Signer::from_seed(&[17; 32])),
6466 )
6467 .unwrap();
6468 let app = build_router(state.clone());
6469 let response = app
6470 .clone()
6471 .oneshot(chat_request("gemini-parallel-tools"))
6472 .await
6473 .unwrap();
6474 assert_eq!(response.status(), StatusCode::OK, "the stream must not hard-fail");
6475 let _ = axum::body::to_bytes(response.into_body(), 64 * 1024).await;
6476 let session = state.sessions.get("gemini-parallel-tools").unwrap();
6477 session.wait_for_worker_jobs().await;
6478 let close = app
6479 .clone()
6480 .oneshot(
6481 Request::builder()
6482 .method("POST")
6483 .uri("/v1/sessions/gemini-parallel-tools/close")
6484 .body(Body::empty())
6485 .unwrap(),
6486 )
6487 .await
6488 .unwrap();
6489 assert_eq!(close.status(), StatusCode::OK);
6490
6491 let digest = pv_core::digest::sha256_hex(b"gemini-parallel-tools");
6493 let stem = digest.get(..32).unwrap();
6494 let trajectory: serde_json::Value =
6495 serde_json::from_slice(&std::fs::read(directory.path().join(format!("{stem}.json"))).unwrap())
6496 .unwrap();
6497 let mut calls: Vec<(String, serde_json::Value)> = Vec::new();
6498 for step in trajectory
6499 .pointer("/steps")
6500 .and_then(serde_json::Value::as_array)
6501 .cloned()
6502 .unwrap_or_default()
6503 {
6504 for call in step
6505 .pointer("/tool_calls")
6506 .and_then(serde_json::Value::as_array)
6507 .cloned()
6508 .unwrap_or_default()
6509 {
6510 calls.push((
6511 call.pointer("/function_name")
6512 .and_then(serde_json::Value::as_str)
6513 .unwrap_or_default()
6514 .to_owned(),
6515 call.pointer("/arguments").cloned().unwrap_or_default(),
6516 ));
6517 }
6518 }
6519 assert_eq!(
6520 calls.len(),
6521 2,
6522 "two functionCall parts in separate chunks must attest as two \
6523 distinct tool calls, not merge at frame-local index 0: {trajectory}"
6524 );
6525 assert_eq!(calls[0].0, "get_weather");
6526 assert_eq!(calls[0].1, serde_json::json!({"city": "Paris"}));
6527 assert_eq!(calls[1].0, "get_time");
6528 assert_eq!(calls[1].1, serde_json::json!({"city": "London"}));
6529 provider.abort();
6530 }
6531
6532 async fn mock_anthropic_refusal() -> Response {
6537 let frames = concat!(
6538 "event: message_start\n",
6539 "data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-sonnet-4-5\",\"usage\":{\"input_tokens\":15,\"output_tokens\":1}}}\n\n",
6540 "event: content_block_start\n",
6541 "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
6542 "event: content_block_delta\n",
6543 "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"I can't help with that.\"}}\n\n",
6544 "event: content_block_stop\n",
6545 "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
6546 "event: message_delta\n",
6547 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"refusal\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":6}}\n\n",
6548 "event: message_stop\n",
6549 "data: {\"type\":\"message_stop\"}\n\n",
6550 );
6551 let mut response = Response::new(Body::from(frames));
6552 response.headers_mut().insert(
6553 axum::http::header::CONTENT_TYPE,
6554 HeaderValue::from_static("text/event-stream"),
6555 );
6556 response
6557 }
6558
6559 async fn test_state_anthropic_refusal(
6560 spool: &std::path::Path,
6561 ) -> (AppState, tokio::task::JoinHandle<()>) {
6562 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6563 let address = listener.local_addr().unwrap();
6564 let provider = Router::new().route("/v1/messages", post(mock_anthropic_refusal));
6565 let server = tokio::spawn(async move {
6566 axum::serve(listener, provider).await.unwrap();
6567 });
6568 let mut config = crate::config::HarnessConfig::for_tests(
6569 &format!("http://{address}"),
6570 &spool.to_string_lossy(),
6571 &spool.to_string_lossy(),
6572 );
6573 config.provider = "anthropic".to_owned();
6574 let state = AppState::new(
6575 config,
6576 Arc::new(InMemoryStore::new()),
6577 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6578 Arc::new(NullBus),
6579 None,
6580 Arc::new(Ed25519Signer::from_seed(&[15; 32])),
6581 )
6582 .unwrap();
6583 (state, server)
6584 }
6585
6586 #[tokio::test]
6597 async fn anthropic_refusal_stop_reason_attests_content_filter() {
6598 let directory = tempfile::tempdir().unwrap();
6599 let (state, provider) = test_state_anthropic_refusal(directory.path()).await;
6600 let app = build_router(state.clone());
6601 let response = app
6602 .clone()
6603 .oneshot(chat_request("anthropic-refusal"))
6604 .await
6605 .unwrap();
6606 assert_eq!(response.status(), StatusCode::OK);
6607 let _ = axum::body::to_bytes(response.into_body(), 64 * 1024).await;
6608 let session = state.sessions.get("anthropic-refusal").unwrap();
6609 session.wait_for_worker_jobs().await;
6610 let close = app
6611 .clone()
6612 .oneshot(
6613 Request::builder()
6614 .method("POST")
6615 .uri("/v1/sessions/anthropic-refusal/close")
6616 .body(Body::empty())
6617 .unwrap(),
6618 )
6619 .await
6620 .unwrap();
6621 assert_eq!(close.status(), StatusCode::OK);
6622 let promoted = app
6623 .clone()
6624 .oneshot(
6625 Request::builder()
6626 .method("POST")
6627 .uri("/v1/sessions/anthropic-refusal/promote")
6628 .body(Body::empty())
6629 .unwrap(),
6630 )
6631 .await
6632 .unwrap();
6633 assert_eq!(promoted.status(), StatusCode::OK);
6634 let receipt: pv_receipts::Receipt = serde_json::from_slice(
6635 &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
6636 .await
6637 .unwrap(),
6638 )
6639 .unwrap();
6640 receipt.verify_embedded().unwrap();
6641 assert_eq!(
6642 receipt.body.stop_reason_id,
6643 pv_events::StopReason::ContentFilter.id(),
6644 "an Anthropic session that ended with stop_reason:refusal must attest \
6645 StopReason::ContentFilter (id {}), not Other (id {}). OpenAI's \
6646 content_filter and Gemini's SAFETY cluster both map to ContentFilter; \
6647 leaving Anthropic's refusal as Other silently splits safety-refusal \
6648 analytics along provider lines. Received: stop_reason_id={}, \
6649 stop_reason={:?}.",
6650 pv_events::StopReason::ContentFilter.id(),
6651 pv_events::StopReason::Other.id(),
6652 receipt.body.stop_reason_id,
6653 receipt.body.stop_reason,
6654 );
6655 provider.abort();
6656 }
6657
6658 async fn mock_anthropic_context_window_exceeded() -> Response {
6664 let frames = concat!(
6665 "event: message_start\n",
6666 "data: {\"type\":\"message_start\",\"message\":{\"model\":\"claude-sonnet-4-5\",\"usage\":{\"input_tokens\":200000,\"output_tokens\":1}}}\n\n",
6667 "event: content_block_start\n",
6668 "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n",
6669 "event: content_block_delta\n",
6670 "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"partial\"}}\n\n",
6671 "event: content_block_stop\n",
6672 "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n",
6673 "event: message_delta\n",
6674 "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"model_context_window_exceeded\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":3}}\n\n",
6675 "event: message_stop\n",
6676 "data: {\"type\":\"message_stop\"}\n\n",
6677 );
6678 let mut response = Response::new(Body::from(frames));
6679 response.headers_mut().insert(
6680 axum::http::header::CONTENT_TYPE,
6681 HeaderValue::from_static("text/event-stream"),
6682 );
6683 response
6684 }
6685
6686 async fn test_state_anthropic_ctx_exceeded(
6687 spool: &std::path::Path,
6688 ) -> (AppState, tokio::task::JoinHandle<()>) {
6689 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6690 let address = listener.local_addr().unwrap();
6691 let provider = Router::new().route("/v1/messages", post(mock_anthropic_context_window_exceeded));
6692 let server = tokio::spawn(async move {
6693 axum::serve(listener, provider).await.unwrap();
6694 });
6695 let mut config = crate::config::HarnessConfig::for_tests(
6696 &format!("http://{address}"),
6697 &spool.to_string_lossy(),
6698 &spool.to_string_lossy(),
6699 );
6700 config.provider = "anthropic".to_owned();
6701 let state = AppState::new(
6702 config,
6703 Arc::new(InMemoryStore::new()),
6704 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6705 Arc::new(NullBus),
6706 None,
6707 Arc::new(Ed25519Signer::from_seed(&[17; 32])),
6708 )
6709 .unwrap();
6710 (state, server)
6711 }
6712
6713 #[tokio::test]
6723 async fn anthropic_context_window_exceeded_attests_max_tokens() {
6724 let directory = tempfile::tempdir().unwrap();
6725 let (state, provider) = test_state_anthropic_ctx_exceeded(directory.path()).await;
6726 let app = build_router(state.clone());
6727 let response = app.clone().oneshot(chat_request("anthropic-ctx")).await.unwrap();
6728 assert_eq!(response.status(), StatusCode::OK);
6729 let _ = axum::body::to_bytes(response.into_body(), 64 * 1024).await;
6730 let session = state.sessions.get("anthropic-ctx").unwrap();
6731 session.wait_for_worker_jobs().await;
6732 let close = app
6733 .clone()
6734 .oneshot(
6735 Request::builder()
6736 .method("POST")
6737 .uri("/v1/sessions/anthropic-ctx/close")
6738 .body(Body::empty())
6739 .unwrap(),
6740 )
6741 .await
6742 .unwrap();
6743 assert_eq!(close.status(), StatusCode::OK);
6744 let promoted = app
6745 .clone()
6746 .oneshot(
6747 Request::builder()
6748 .method("POST")
6749 .uri("/v1/sessions/anthropic-ctx/promote")
6750 .body(Body::empty())
6751 .unwrap(),
6752 )
6753 .await
6754 .unwrap();
6755 assert_eq!(promoted.status(), StatusCode::OK);
6756 let receipt: pv_receipts::Receipt = serde_json::from_slice(
6757 &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
6758 .await
6759 .unwrap(),
6760 )
6761 .unwrap();
6762 receipt.verify_embedded().unwrap();
6763 assert_eq!(
6764 receipt.body.stop_reason_id,
6765 pv_events::StopReason::MaxTokens.id(),
6766 "an Anthropic session that ended with \
6767 stop_reason:model_context_window_exceeded must attest \
6768 StopReason::MaxTokens (id {}), not Other (id {}). OpenAI's `length` \
6769 and Gemini's MAX_TOKENS both map to MaxTokens; leaving Anthropic's \
6770 overflow variant as Other hides capacity signals from operators. \
6771 Received: stop_reason_id={}, stop_reason={:?}.",
6772 pv_events::StopReason::MaxTokens.id(),
6773 pv_events::StopReason::Other.id(),
6774 receipt.body.stop_reason_id,
6775 receipt.body.stop_reason,
6776 );
6777 provider.abort();
6778 }
6779
6780 #[tokio::test]
6786 async fn anthropic_session_records_the_same_audit_shape() {
6787 let directory = tempfile::tempdir().unwrap();
6788 let (state, provider) = test_state_anthropic(directory.path()).await;
6789 let app = build_router(state.clone());
6790 let response = app.clone().oneshot(chat_request("anthropic-flow")).await.unwrap();
6791 assert_eq!(response.status(), StatusCode::OK);
6792 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
6793 .await
6794 .unwrap();
6795 let body = String::from_utf8_lossy(&body);
6796 assert!(
6797 body.contains("hello from anthropic"),
6798 "the translated stream must carry the upstream text: {body}"
6799 );
6800 assert!(
6801 body.contains("\"object\":\"chat.completion.chunk\"")
6802 && !body.contains("message_start")
6803 && body.trim_end().ends_with("data: [DONE]"),
6804 "an OpenAI client must receive OpenAI chunks, not the native stream: {body}"
6805 );
6806
6807 let session = state.sessions.get("anthropic-flow").unwrap();
6808 session.wait_for_worker_jobs().await;
6809 use std::sync::atomic::Ordering;
6813 assert_eq!(session.totals.prompt_tokens.load(Ordering::Acquire), 25);
6814 assert_eq!(session.totals.completion_tokens.load(Ordering::Acquire), 15);
6815
6816 let close = app
6817 .clone()
6818 .oneshot(
6819 Request::builder()
6820 .method("POST")
6821 .uri("/v1/sessions/anthropic-flow/close")
6822 .body(Body::empty())
6823 .unwrap(),
6824 )
6825 .await
6826 .unwrap();
6827 assert_eq!(close.status(), StatusCode::OK);
6828 let promoted = app
6829 .clone()
6830 .oneshot(
6831 Request::builder()
6832 .method("POST")
6833 .uri("/v1/sessions/anthropic-flow/promote")
6834 .body(Body::empty())
6835 .unwrap(),
6836 )
6837 .await
6838 .unwrap();
6839 assert_eq!(promoted.status(), StatusCode::OK);
6840 let receipt: pv_receipts::Receipt = serde_json::from_slice(
6841 &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
6842 .await
6843 .unwrap(),
6844 )
6845 .unwrap();
6846 receipt.verify_embedded().unwrap();
6847 assert_eq!(receipt.body.stop_reason_id, pv_events::StopReason::Stop.id());
6849 assert_eq!(receipt.body.cost.prompt_tokens, 25);
6850 assert_eq!(receipt.body.cost.completion_tokens, 15);
6851 provider.abort();
6852 }
6853
6854 async fn test_state_with_token_cap(
6855 spool: &std::path::Path,
6856 max_tokens: Option<u64>,
6857 ) -> (AppState, tokio::task::JoinHandle<()>) {
6858 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6859 let address = listener.local_addr().unwrap();
6860 let provider = Router::new().route("/v1/chat/completions", post(mock_chat));
6861 let server = tokio::spawn(async move {
6862 axum::serve(listener, provider).await.unwrap();
6863 });
6864 let mut config = crate::config::HarnessConfig::for_tests(
6865 &format!("http://{address}"),
6866 &spool.to_string_lossy(),
6867 &spool.to_string_lossy(),
6868 );
6869 config.budget.max_tokens = max_tokens;
6870 let state = AppState::new(
6871 config,
6872 Arc::new(InMemoryStore::new()),
6873 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6874 Arc::new(NullBus),
6875 None,
6876 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
6877 )
6878 .unwrap();
6879 (state, server)
6880 }
6881
6882 fn chat_request(session: &str) -> Request<Body> {
6883 chat_request_with_payload(session, chat_payload())
6884 }
6885
6886 #[tokio::test]
6893 async fn header_less_one_shot_closes_immediately_after_the_response() {
6894 let directory = tempfile::tempdir().unwrap();
6895 let (state, provider) = test_state(directory.path()).await;
6896 let request = Request::builder()
6897 .method("POST")
6898 .uri("/v1/chat/completions")
6899 .header(axum::http::header::CONTENT_TYPE, "application/json")
6900 .header("x-pv-workflow", "unsigned")
6902 .body(Body::from(serde_json::to_vec(&chat_payload()).unwrap()))
6903 .unwrap();
6904 let response = build_router(state.clone()).oneshot(request).await.unwrap();
6905 assert_eq!(response.status(), StatusCode::OK);
6906 let session_id = response
6907 .headers()
6908 .get("x-pv-session")
6909 .and_then(|v| v.to_str().ok())
6910 .expect("generated session id must be echoed in the response header")
6911 .to_owned();
6912 axum::body::to_bytes(response.into_body(), 1024 * 1024)
6913 .await
6914 .unwrap();
6915
6916 let session = state.sessions.get(&session_id).expect("session registered");
6917 let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10);
6918 while !session.is_closed() {
6919 assert!(
6920 std::time::Instant::now() < deadline,
6921 "an ephemeral one-shot must close immediately after the response, not wait for the idle sweeper"
6922 );
6923 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
6924 }
6925 provider.abort();
6926 }
6927
6928 #[tokio::test]
6940 async fn principal_budget_charges_completion_tokens_across_rotated_sessions() {
6941 let directory = tempfile::tempdir().unwrap();
6942 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
6943 let address = listener.local_addr().unwrap();
6944 let provider = Router::new().route(
6945 "/v1/chat/completions",
6946 post(|| async {
6947 Json(json!({
6948 "model": "mock",
6949 "choices": [{
6950 "message": {"role": "assistant", "content": "big completion"},
6951 "finish_reason": "stop"
6952 }],
6953 "usage": {"prompt_tokens": 5, "completion_tokens": 1000}
6954 }))
6955 }),
6956 );
6957 let server = tokio::spawn(async move {
6958 axum::serve(listener, provider).await.unwrap();
6959 });
6960 let spool = directory.path().to_string_lossy().into_owned();
6961 let mut config =
6962 crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
6963 config.budget.max_tokens = Some(1_000_000);
6965 config.principal_budget = Some(pv_state::BudgetSpec {
6968 max_tokens: Some(1500),
6969 ..pv_state::BudgetSpec::default()
6970 });
6971 let state = AppState::new(
6972 config,
6973 Arc::new(InMemoryStore::new()),
6974 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
6975 Arc::new(NullBus),
6976 None,
6977 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
6978 )
6979 .unwrap();
6980 let app = build_router(state);
6981 let non_streamed = json!({
6982 "model": "mock",
6983 "messages": [{"role": "user", "content": "hello"}]
6984 });
6985
6986 let first = app
6987 .clone()
6988 .oneshot(chat_request_with_payload("principal-rot-1", non_streamed.clone()))
6989 .await
6990 .unwrap();
6991 assert_eq!(
6992 first.status(),
6993 StatusCode::OK,
6994 "first rotated request fits the principal budget"
6995 );
6996 axum::body::to_bytes(first.into_body(), 1024 * 1024)
6997 .await
6998 .unwrap();
6999
7000 let second = app
7005 .oneshot(chat_request_with_payload("principal-rot-2", non_streamed))
7006 .await
7007 .unwrap();
7008 let status = second.status();
7009 let body = axum::body::to_bytes(second.into_body(), 1024 * 1024)
7010 .await
7011 .unwrap();
7012 let body = String::from_utf8_lossy(&body);
7013 assert_eq!(
7014 status,
7015 StatusCode::FORBIDDEN,
7016 "rotated completion spend must be refused by the principal ledger; got {status}: {body}"
7017 );
7018 assert!(
7019 body.contains("principal.max_tokens"),
7020 "refusal must name the principal-scoped limit: {body}"
7021 );
7022 server.abort();
7023 }
7024
7025 #[tokio::test]
7038 async fn upstream_authorization_contract_by_posture() {
7039 type Captured = Arc<parking_lot::Mutex<Vec<(Option<String>, Option<String>, Option<String>)>>>;
7040 for posture in ["keyless_ignore", "static_key_file", "passthrough"] {
7041 let directory = tempfile::tempdir().unwrap();
7042 let captured: Captured = Arc::default();
7043 let sink = Arc::clone(&captured);
7044 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
7045 let address = listener.local_addr().unwrap();
7046 let provider = Router::new().route(
7047 "/v1/chat/completions",
7048 post(
7049 move |headers: axum::http::HeaderMap, Json(payload): Json<Value>| {
7050 let sink = Arc::clone(&sink);
7051 async move {
7052 let get = |name: &str| {
7053 headers.get(name).and_then(|v| v.to_str().ok()).map(str::to_owned)
7054 };
7055 sink.lock().push((
7056 get("authorization"),
7057 get("openai-beta"),
7058 get("x-custom-header"),
7059 ));
7060 mock_chat(Json(payload)).await
7061 }
7062 },
7063 ),
7064 );
7065 let server = tokio::spawn(async move {
7066 axum::serve(listener, provider).await.unwrap();
7067 });
7068 let spool = directory.path().to_string_lossy().into_owned();
7069 let mut config =
7070 crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
7071 match posture {
7072 "keyless_ignore" => config.ignore_client_authorization = true,
7073 "static_key_file" => {
7074 let key_path = directory.path().join("upstream.key");
7075 std::fs::write(&key_path, "sk-SERVER-SECRET").unwrap();
7076 #[cfg(unix)]
7077 {
7078 use std::os::unix::fs::PermissionsExt as _;
7079 std::fs::set_permissions(&key_path, std::fs::Permissions::from_mode(0o600)).unwrap();
7080 }
7081 config.upstream_api_key_file = Some(key_path.to_string_lossy().into_owned());
7082 }
7083 "passthrough" => config.upstream_authorization_passthrough = true,
7084 _ => unreachable!(),
7085 }
7086 let state = AppState::new(
7087 config,
7088 Arc::new(InMemoryStore::new()),
7089 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
7090 Arc::new(NullBus),
7091 None,
7092 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
7093 )
7094 .unwrap();
7095 let request = Request::builder()
7096 .method("POST")
7097 .uri("/v1/chat/completions")
7098 .header(axum::http::header::CONTENT_TYPE, "application/json")
7099 .header("authorization", "Bearer sk-CLIENT-JUNK")
7100 .header("openai-beta", "assistants=v2")
7101 .header("x-custom-header", "evil")
7102 .header("x-pv-session", format!("auth-{posture}"))
7103 .header("x-pv-workflow", "unsigned")
7104 .body(Body::from(serde_json::to_vec(&chat_payload()).unwrap()))
7105 .unwrap();
7106 let response = build_router(state).oneshot(request).await.unwrap();
7107 assert_eq!(response.status(), StatusCode::OK, "posture {posture} must serve");
7108 axum::body::to_bytes(response.into_body(), 1024 * 1024)
7109 .await
7110 .unwrap();
7111 let seen = captured.lock().clone();
7112 let (auth, beta, custom) = seen
7113 .first()
7114 .cloned()
7115 .unwrap_or_else(|| panic!("posture {posture}: upstream never saw the request"));
7116 match posture {
7117 "keyless_ignore" => assert_eq!(
7118 auth, None,
7119 "keyless posture must DISCARD the client placeholder bearer"
7120 ),
7121 "static_key_file" => assert_eq!(
7122 auth.as_deref(),
7123 Some("Bearer sk-SERVER-SECRET"),
7124 "static-key posture must inject the server key, never the client's"
7125 ),
7126 "passthrough" => assert_eq!(
7127 auth.as_deref(),
7128 Some("Bearer sk-CLIENT-JUNK"),
7129 "passthrough posture must forward the client bearer verbatim"
7130 ),
7131 _ => unreachable!(),
7132 }
7133 assert_eq!(
7134 beta.as_deref(),
7135 Some("assistants=v2"),
7136 "posture {posture}: allowlisted provider header must relay"
7137 );
7138 assert_eq!(
7139 custom, None,
7140 "posture {posture}: non-allowlisted client headers must never reach the upstream"
7141 );
7142 server.abort();
7143 }
7144 }
7145
7146 #[tokio::test]
7152 async fn quarantined_session_close_and_promote_map_to_conflict_not_500() {
7153 let directory = tempfile::tempdir().unwrap();
7154 let (state, provider) = test_state(directory.path()).await;
7155 let app = build_router(state.clone());
7156 let opened = app.clone().oneshot(chat_request("quarantine-409")).await.unwrap();
7158 assert_eq!(opened.status(), StatusCode::OK);
7159 axum::body::to_bytes(opened.into_body(), 1024 * 1024)
7160 .await
7161 .unwrap();
7162 let session = state.sessions.get("quarantine-409").unwrap();
7163 session.wait_for_worker_jobs().await;
7164 session.mark_capture_failed();
7165
7166 let close = app
7167 .clone()
7168 .oneshot(
7169 Request::builder()
7170 .method("POST")
7171 .uri("/v1/sessions/quarantine-409/close")
7172 .body(Body::empty())
7173 .unwrap(),
7174 )
7175 .await
7176 .unwrap();
7177 assert_eq!(
7178 close.status(),
7179 StatusCode::CONFLICT,
7180 "capture-incomplete close refusal is a state conflict, not a server fault"
7181 );
7182 let promote = app
7183 .oneshot(
7184 Request::builder()
7185 .method("POST")
7186 .uri("/v1/sessions/quarantine-409/promote")
7187 .body(Body::empty())
7188 .unwrap(),
7189 )
7190 .await
7191 .unwrap();
7192 assert_eq!(
7193 promote.status(),
7194 StatusCode::CONFLICT,
7195 "unfulfillable promotion refusal is a state conflict, not a server fault"
7196 );
7197 provider.abort();
7198 }
7199
7200 fn captured_choice(
7203 text: &str,
7204 finish_reason: Option<&str>,
7205 ) -> std::collections::BTreeMap<u64, ChoiceCapture> {
7206 let mut choices = std::collections::BTreeMap::new();
7207 choices.insert(
7208 0,
7209 ChoiceCapture {
7210 content: text.to_owned(),
7211 reasoning: String::new(),
7212 finish_reason: finish_reason.map(str::to_owned),
7213 },
7214 );
7215 choices
7216 }
7217
7218 fn chat_request_with_payload(session: &str, payload: Value) -> Request<Body> {
7219 Request::builder()
7220 .method("POST")
7221 .uri("/v1/chat/completions")
7222 .header(axum::http::header::CONTENT_TYPE, "application/json")
7223 .header("x-pv-session", session)
7224 .header("x-pv-workflow", "unsigned")
7225 .body(Body::from(serde_json::to_vec(&payload).unwrap()))
7226 .unwrap()
7227 }
7228
7229 fn chat_payload() -> Value {
7230 json!({
7231 "model": "mock",
7232 "stream": true,
7233 "messages": [{"role": "user", "content": "hello"}],
7234 })
7235 }
7236
7237 #[tokio::test]
7247 async fn chat_completions_refuses_duplicate_json_keys() {
7248 let directory = tempfile::tempdir().unwrap();
7249 let (state, provider) = test_state(directory.path()).await;
7250 let app = build_router(state.clone());
7251 let dup = br#"{"model":"mock","messages":[{"role":"user","content":"safe"}],"messages":[{"role":"user","content":"hostile"}]}"#;
7255 let request = Request::builder()
7256 .method("POST")
7257 .uri("/v1/chat/completions")
7258 .header(axum::http::header::CONTENT_TYPE, "application/json")
7259 .header("x-pv-session", "dup-key")
7260 .header("x-pv-workflow", "unsigned")
7261 .body(Body::from(&dup[..]))
7262 .unwrap();
7263 let response = app.oneshot(request).await.unwrap();
7264 assert_eq!(
7265 response.status(),
7266 StatusCode::BAD_REQUEST,
7267 "duplicate `messages` key must be refused with 400, not silently taken"
7268 );
7269 let body = axum::body::to_bytes(response.into_body(), 4 * 1024)
7270 .await
7271 .unwrap();
7272 let text = std::str::from_utf8(&body).unwrap();
7273 assert!(
7274 text.contains("duplicate"),
7275 "response body must name the reason, got {text:?}"
7276 );
7277 provider.abort();
7278 }
7279
7280 fn active_records(
7281 directory: &std::path::Path,
7282 state: &AppState,
7283 session_id: &str,
7284 ) -> Vec<crate::worker::ActiveJournalRecord> {
7285 let digest = pv_core::digest::sha256_hex(session_id.as_bytes());
7286 let journal =
7287 std::fs::read_to_string(directory.join(format!("{}.events.ndjson", &digest[..32]))).unwrap();
7288 journal
7289 .lines()
7290 .enumerate()
7291 .map(|(index, line)| {
7292 crate::journal::open(
7293 &state.journal_key,
7294 &format!("{session_id}:active"),
7295 index as u64,
7296 line.as_bytes(),
7297 )
7298 .unwrap()
7299 })
7300 .collect()
7301 }
7302
7303 #[tokio::test]
7304 async fn full_chat_close_and_promotion_flow() {
7305 let directory = tempfile::tempdir().unwrap();
7306 let (state, provider) = test_state(directory.path()).await;
7307 let app = build_router(state.clone());
7308
7309 let response = app.clone().oneshot(chat_request("http-flow")).await.unwrap();
7310 assert_eq!(response.status(), StatusCode::OK);
7311 assert_eq!(response.headers().get("x-pv-session").unwrap(), "http-flow");
7312 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
7313 .await
7314 .unwrap();
7315 assert!(String::from_utf8_lossy(&body).contains("[DONE]"));
7316
7317 let close = app
7318 .clone()
7319 .oneshot(
7320 Request::builder()
7321 .method("POST")
7322 .uri("/v1/sessions/http-flow/close")
7323 .body(Body::empty())
7324 .unwrap(),
7325 )
7326 .await
7327 .unwrap();
7328 let close_status = close.status();
7329 let close_bytes = axum::body::to_bytes(close.into_body(), 64 * 1024).await.unwrap();
7330 assert_eq!(
7331 close_status,
7332 StatusCode::OK,
7333 "{}",
7334 String::from_utf8_lossy(&close_bytes)
7335 );
7336 let close_body: Value = serde_json::from_slice(&close_bytes).unwrap();
7337 assert_eq!(close_body["kind"], "atif");
7338 let artifact = state
7339 .sessions
7340 .get("http-flow")
7341 .unwrap()
7342 .atif_path
7343 .lock()
7344 .clone()
7345 .unwrap();
7346 let trajectory: pv_atif::Trajectory =
7347 serde_json::from_slice(&tokio::fs::read(&artifact).await.unwrap()).unwrap();
7348 if let Some(destination) = std::env::var_os("PV_HARBOR_INTEROP_OUT") {
7349 std::fs::copy(&artifact, destination).unwrap();
7350 }
7351 assert_eq!(
7352 trajectory.steps.len(),
7353 2,
7354 "request and response must both be captured"
7355 );
7356 assert_eq!(trajectory.steps[0].source, pv_atif::Source::User);
7357 assert!(trajectory.steps[0].metrics.is_none());
7358 assert_eq!(trajectory.steps[1].source, pv_atif::Source::Agent);
7359 assert_eq!(trajectory.steps[1].message, Value::String("hello".to_owned()));
7360 assert_eq!(
7361 trajectory.steps[1].metrics.as_ref().unwrap().cached_tokens,
7362 Some(4)
7363 );
7364
7365 let promoted = app
7366 .oneshot(
7367 Request::builder()
7368 .method("POST")
7369 .uri("/v1/sessions/http-flow/promote")
7370 .body(Body::empty())
7371 .unwrap(),
7372 )
7373 .await
7374 .unwrap();
7375 assert_eq!(promoted.status(), StatusCode::OK);
7376 let receipt: pv_receipts::Receipt = serde_json::from_slice(
7377 &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
7378 .await
7379 .unwrap(),
7380 )
7381 .unwrap();
7382 receipt.verify_embedded().unwrap();
7383 assert_eq!(receipt.body.stop_reason_id, pv_events::StopReason::Stop.id());
7384 assert_eq!(receipt.body.cost.completion_tokens, 3);
7385 assert_eq!(receipt.body.cost.cached_tokens, 4);
7386 assert_eq!(receipt.body.cost.cost_usd_micros, 1_250);
7387 assert_eq!(
7388 receipt.body.cost.prompt_tokens, 10,
7389 "prompt total must reconcile to the provider-reported usage, \
7390 not the admission-time heuristic (and never both summed)"
7391 );
7392 provider.abort();
7393 }
7394
7395 #[tokio::test]
7396 async fn mcp_metrics_and_abort_paths_are_enforced() {
7397 let directory = tempfile::tempdir().unwrap();
7398 let (state, provider) = test_state(directory.path()).await;
7399 let app = build_router(state.clone());
7400 let allowed = app
7401 .clone()
7402 .oneshot(
7403 Request::builder()
7404 .method("POST")
7405 .uri("/v1/mcp")
7406 .header("x-pv-session", "mcp-flow")
7407 .body(Body::from(
7408 r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"read","arguments":{}}}"#,
7409 ))
7410 .unwrap(),
7411 )
7412 .await
7413 .unwrap();
7414 assert_eq!(allowed.status(), StatusCode::OK);
7415 let allowed_body: Value = serde_json::from_slice(
7416 &axum::body::to_bytes(allowed.into_body(), 64 * 1024)
7417 .await
7418 .unwrap(),
7419 )
7420 .unwrap();
7421 assert_eq!(allowed_body["jsonrpc"], "2.0");
7424 assert_eq!(allowed_body["id"], 1);
7425 assert!(allowed_body["result"]["decision_us"].as_u64().unwrap() < 5_000);
7426
7427 let blocked = app
7428 .clone()
7429 .oneshot(
7430 Request::builder()
7431 .method("POST")
7432 .uri("/v1/mcp")
7433 .header("x-pv-session", "mcp-flow")
7434 .body(Body::from("not-json"))
7435 .unwrap(),
7436 )
7437 .await
7438 .unwrap();
7439 assert_eq!(blocked.status(), StatusCode::BAD_REQUEST);
7442 let blocked_body: Value = serde_json::from_slice(
7443 &axum::body::to_bytes(blocked.into_body(), 64 * 1024)
7444 .await
7445 .unwrap(),
7446 )
7447 .unwrap();
7448 assert_eq!(blocked_body["error"]["code"], -32700);
7449
7450 let metrics_response = app
7451 .clone()
7452 .oneshot(Request::builder().uri("/metrics").body(Body::empty()).unwrap())
7453 .await
7454 .unwrap();
7455 let metrics = String::from_utf8_lossy(
7456 &axum::body::to_bytes(metrics_response.into_body(), 128 * 1024)
7457 .await
7458 .unwrap(),
7459 )
7460 .into_owned();
7461 assert!(metrics.contains("pv_stage_duration_seconds") || metrics.contains("pv_sessions"));
7462
7463 let aborted = app.clone().oneshot(chat_request("abort-flow")).await.unwrap();
7464 drop(aborted);
7465 let session = tokio::time::timeout(Duration::from_secs(1), async {
7472 loop {
7473 if let Some(session) = state.sessions.get("abort-flow") {
7474 if session.active_streams_count() == 0 {
7475 break session;
7476 }
7477 }
7478 tokio::task::yield_now().await;
7479 }
7480 })
7481 .await
7482 .expect("abort stream never drained");
7483 assert!(
7484 !session.is_closed(),
7485 "an abort after captured content must leave the conversation open"
7486 );
7487 provider.abort();
7488 }
7489
7490 #[tokio::test]
7502 async fn stream_drop_during_concurrent_close_still_captures_response() {
7503 let directory = tempfile::tempdir().unwrap();
7504 let (state, provider) = test_state(directory.path()).await;
7505 let app = build_router(state.clone());
7506
7507 let request = Request::builder()
7508 .method("POST")
7509 .uri("/v1/chat/completions")
7510 .header(axum::http::header::CONTENT_TYPE, "application/json")
7511 .header("x-pv-session", "stream-close-race")
7512 .header("x-pv-workflow", "signed")
7513 .body(Body::from(serde_json::to_vec(&chat_payload()).unwrap()))
7514 .unwrap();
7515 let response = app.oneshot(request).await.unwrap();
7516 assert_eq!(response.status(), StatusCode::OK);
7517
7518 let session = state.sessions.get("stream-close-race").unwrap();
7519
7520 let close_task = {
7521 let finalizer = state.finalizer.clone();
7522 let session = Arc::clone(&session);
7523 tokio::spawn(async move { finalizer.close_session(session, StopReason::SessionClosed).await })
7524 };
7525
7526 tokio::time::timeout(Duration::from_secs(1), async {
7527 while !session.is_closed() {
7528 tokio::task::yield_now().await;
7529 }
7530 })
7531 .await
7532 .expect("concurrent close did not initiate before the response was dropped");
7533
7534 drop(response);
7535
7536 close_task
7537 .await
7538 .expect("close_session task panicked")
7539 .expect("close_session returned an error");
7540
7541 let final_count = session.chain.lock().count();
7542 assert!(
7543 final_count >= 2,
7544 "AbortFinalizingStream::drop must submit the response capture even when the session is already being closed — otherwise a concurrent close during an active stream discards the response and the receipt attests to fewer events than actually flowed; got chain.count() = {final_count}",
7545 );
7546
7547 provider.abort();
7548 }
7549
7550 #[tokio::test]
7574 async fn drop_with_pending_budget_during_concurrent_close_does_not_stick_session() {
7575 let directory = tempfile::tempdir().unwrap();
7576 let (state, provider) = test_state(directory.path()).await;
7577
7578 let identity = pv_events::AgentIdentity {
7579 version: "1".into(),
7580 charter: "c".into(),
7581 instance_uid: "i".into(),
7582 ttl_remaining_s: None,
7583 };
7584 let session = state.sessions.get_or_open(
7585 "pending-budget-close",
7586 crate::session::Workflow::Signed,
7587 &identity,
7588 &state.config.breaker,
7589 );
7590 let permits = state
7595 .worker
7596 .try_reserve_pair("pending-budget-close")
7597 .expect("test setup: fused worker/response permit");
7598 drop(permits.worker);
7602 let permit = permits.response;
7603 let lease = crate::session::SessionLease::new(Arc::clone(&session));
7606
7607 let budget_task =
7611 tokio::spawn(async { Ok::<_, String>(pv_state::BudgetDecision::Allowed { remaining: 100 }) });
7612
7613 let stream = AbortFinalizingStream {
7614 inner: futures::stream::empty::<Result<Bytes, std::io::Error>>().boxed(),
7615 session: Arc::clone(&session),
7616 identity: identity.clone(),
7617 response_permit: Some(permit),
7618 worker: state.worker.clone(),
7619 store: Arc::clone(&state.store),
7620 budget: state.config.budget.clone(),
7621 rates: Arc::clone(&state.pushed_rates),
7622 billed_prompt_tokens: 0,
7623 principal_budget: None,
7624 finalizer: state.finalizer.clone(),
7625 _lease: lease,
7626 response_marker: None,
7627 response_attempt_id: "pending-budget-attempt".into(),
7628 response_choices: captured_choice("captured content", Some("stop")),
7629 response_model: None,
7630 upstream_status: StatusCode::OK,
7631 response_cost_usd_micros: 0,
7632 response_tool_calls: std::collections::BTreeMap::new(),
7633 response_metrics: pv_events::EventMetrics::default(),
7634 charged_completion_tokens: 3,
7635 last_reported_completion_tokens: Some(3),
7636 last_reported_prompt_tokens: None,
7637 last_reported_cached_tokens: None,
7638 last_reported_cost_usd_micros: None,
7639 saw_chunk: true,
7640 capture_attempted: false,
7641 is_sse: true,
7642 protocol_buffer: Vec::new(),
7643 frame_scratch: Vec::new(),
7644 pending_output: std::collections::VecDeque::new(),
7645 pending_budget: Some(PendingBudget {
7646 task: budget_task,
7647 continuation: BudgetContinuation::FinishSse,
7648 delta: 0,
7649 }),
7650 captured_bytes: 100,
7651 completed: false,
7652 ephemeral: false,
7653 };
7654
7655 assert!(session.try_close(), "test setup: try_close should have won");
7657 assert!(session.is_closed());
7658
7659 drop(stream);
7662
7663 assert!(
7664 !session.capture_failed(),
7665 "AbortFinalizingStream::drop must not fail the capture just because a completion-token budget check was still in flight while a concurrent close is already draining this stream via wait_for_streams — otherwise close_session_locked's capture_failed guard seals the session and returns CaptureIncomplete: the receipt is never signed and the capture is lost",
7666 );
7667
7668 let outcome = tokio::time::timeout(
7674 Duration::from_secs(3),
7675 state
7676 .finalizer
7677 .close_session(Arc::clone(&session), StopReason::SessionClosed),
7678 )
7679 .await
7680 .expect("close_session hung after drop released the stream");
7681 outcome.expect(
7682 "close_session must succeed after AbortFinalizingStream::drop submits the response capture — a session sealed without its artifact is worse than a marginally-over-budget response record, because the receipt is lost entirely",
7683 );
7684
7685 provider.abort();
7686 }
7687
7688 #[tokio::test]
7695 async fn absorb_error_rolls_back_undebited_completion_charges() {
7696 let directory = tempfile::tempdir().unwrap();
7697 let (state, provider) = test_state(directory.path()).await;
7698 let identity = pv_events::AgentIdentity {
7699 version: "1".into(),
7700 charter: "c".into(),
7701 instance_uid: "i".into(),
7702 ttl_remaining_s: None,
7703 };
7704 let session = state.sessions.get_or_open(
7705 "absorb-rollback",
7706 crate::session::Workflow::Signed,
7707 &identity,
7708 &state.config.breaker,
7709 );
7710 let lease = crate::session::SessionLease::new(Arc::clone(&session));
7711 let mut stream = AbortFinalizingStream {
7712 inner: futures::stream::empty::<Result<Bytes, std::io::Error>>().boxed(),
7713 session: Arc::clone(&session),
7714 identity,
7715 response_permit: None,
7716 worker: state.worker.clone(),
7717 store: Arc::clone(&state.store),
7718 budget: state.config.budget.clone(),
7719 rates: Arc::clone(&state.pushed_rates),
7720 billed_prompt_tokens: 0,
7721 principal_budget: None,
7722 finalizer: state.finalizer.clone(),
7723 _lease: lease,
7724 response_marker: None,
7725 response_attempt_id: "absorb-rollback-attempt".into(),
7726 response_choices: std::collections::BTreeMap::new(),
7727 response_model: None,
7728 upstream_status: StatusCode::OK,
7729 response_cost_usd_micros: 0,
7730 response_tool_calls: std::collections::BTreeMap::new(),
7731 response_metrics: pv_events::EventMetrics::default(),
7732 charged_completion_tokens: 0,
7733 last_reported_completion_tokens: None,
7734 last_reported_prompt_tokens: None,
7735 last_reported_cached_tokens: None,
7736 last_reported_cost_usd_micros: None,
7737 saw_chunk: false,
7738 capture_attempted: true,
7739 is_sse: true,
7740 protocol_buffer: Vec::new(),
7741 frame_scratch: Vec::new(),
7742 pending_output: std::collections::VecDeque::new(),
7743 pending_budget: None,
7744 captured_bytes: 0,
7745 completed: false,
7746 ephemeral: false,
7747 };
7748 let chunk = concat!(
7753 "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"some streamed words here\"}}]}\n\n",
7754 "data: {not-json\n\n",
7755 );
7756 let result = stream.absorb_network_chunk(chunk.as_bytes());
7757 assert!(result.is_err(), "the malformed frame must fail the chunk");
7758 assert_eq!(
7759 stream.charged_completion_tokens, 0,
7760 "an errored chunk's never-debited frame deltas must be rolled back \
7761 from charged_completion_tokens, or the terminal refund credits \
7762 tokens that were never debited"
7763 );
7764 provider.abort();
7765 }
7766
7767 #[tokio::test]
7776 async fn refused_budget_check_rolls_back_undebited_charge() {
7777 use futures::StreamExt as _;
7778 let directory = tempfile::tempdir().unwrap();
7779 let (state, provider) = test_state(directory.path()).await;
7780 let identity = pv_events::AgentIdentity {
7781 version: "1".into(),
7782 charter: "c".into(),
7783 instance_uid: "i".into(),
7784 ttl_remaining_s: None,
7785 };
7786 let session = state.sessions.get_or_open(
7787 "refused-rollback",
7788 crate::session::Workflow::Signed,
7789 &identity,
7790 &state.config.breaker,
7791 );
7792 let lease = crate::session::SessionLease::new(Arc::clone(&session));
7793 let refused = tokio::spawn(async {
7794 Ok::<_, String>(pv_state::BudgetDecision::Refused {
7795 limit: "tokens".to_owned(),
7796 cap: 100,
7797 })
7798 });
7799 let mut stream = AbortFinalizingStream {
7800 inner: futures::stream::empty::<Result<Bytes, std::io::Error>>().boxed(),
7801 session: Arc::clone(&session),
7802 identity,
7803 response_permit: None,
7804 worker: state.worker.clone(),
7805 store: Arc::clone(&state.store),
7806 budget: state.config.budget.clone(),
7807 rates: Arc::clone(&state.pushed_rates),
7808 billed_prompt_tokens: 0,
7809 principal_budget: None,
7810 finalizer: state.finalizer.clone(),
7811 _lease: lease,
7812 response_marker: None,
7813 response_attempt_id: "refused-rollback-attempt".into(),
7814 response_choices: std::collections::BTreeMap::new(),
7815 response_model: None,
7816 upstream_status: StatusCode::OK,
7817 response_cost_usd_micros: 0,
7818 response_tool_calls: std::collections::BTreeMap::new(),
7819 response_metrics: pv_events::EventMetrics::default(),
7820 charged_completion_tokens: 10,
7823 last_reported_completion_tokens: None,
7824 last_reported_prompt_tokens: None,
7825 last_reported_cached_tokens: None,
7826 last_reported_cost_usd_micros: None,
7827 saw_chunk: true,
7828 capture_attempted: true,
7831 is_sse: true,
7832 protocol_buffer: Vec::new(),
7833 frame_scratch: Vec::new(),
7834 pending_output: std::collections::VecDeque::new(),
7835 pending_budget: Some(PendingBudget {
7836 task: refused,
7837 continuation: BudgetContinuation::FinishSse,
7838 delta: 4,
7839 }),
7840 captured_bytes: 0,
7841 completed: false,
7842 ephemeral: false,
7843 };
7844 let item = stream.next().await.expect("refusal must yield a terminal item");
7845 let error = item.expect_err("a refused budget check must surface as an error");
7846 assert_eq!(error.kind(), std::io::ErrorKind::QuotaExceeded);
7847 assert_eq!(
7848 stream.charged_completion_tokens, 6,
7849 "the refused (never-debited) delta must be rolled back from \
7850 charged_completion_tokens, or the terminal refund credits \
7851 tokens that were never spent"
7852 );
7853 assert!(
7854 stream.session.enforcement_tripped(),
7855 "a genuine mid-stream refusal still latches enforcement"
7856 );
7857 provider.abort();
7858 }
7859
7860 fn empty_abort_stream(
7866 state: &AppState,
7867 session: &Arc<crate::session::Session>,
7868 identity: pv_events::AgentIdentity,
7869 ephemeral: bool,
7870 ) -> AbortFinalizingStream {
7871 AbortFinalizingStream {
7872 inner: futures::stream::empty::<Result<Bytes, std::io::Error>>().boxed(),
7873 session: Arc::clone(session),
7874 identity,
7875 response_permit: None,
7876 worker: state.worker.clone(),
7877 store: Arc::clone(&state.store),
7878 budget: state.config.budget.clone(),
7879 rates: Arc::clone(&state.pushed_rates),
7880 billed_prompt_tokens: 0,
7881 principal_budget: None,
7882 finalizer: state.finalizer.clone(),
7883 _lease: crate::session::SessionLease::new(Arc::clone(session)),
7884 response_marker: None,
7885 response_attempt_id: "empty-abort-attempt".into(),
7886 response_choices: std::collections::BTreeMap::new(),
7887 response_model: None,
7888 upstream_status: StatusCode::OK,
7889 response_cost_usd_micros: 0,
7890 response_tool_calls: std::collections::BTreeMap::new(),
7891 response_metrics: pv_events::EventMetrics::default(),
7892 charged_completion_tokens: 0,
7893 last_reported_completion_tokens: None,
7894 last_reported_prompt_tokens: None,
7895 last_reported_cached_tokens: None,
7896 last_reported_cost_usd_micros: None,
7897 saw_chunk: false,
7898 capture_attempted: true,
7899 is_sse: true,
7900 protocol_buffer: Vec::new(),
7901 frame_scratch: Vec::new(),
7902 pending_output: std::collections::VecDeque::new(),
7903 pending_budget: None,
7904 captured_bytes: 0,
7905 completed: false,
7906 ephemeral,
7907 }
7908 }
7909
7910 #[tokio::test]
7921 async fn non_ephemeral_empty_abort_leaves_session_open() {
7922 let directory = tempfile::tempdir().unwrap();
7923 let (state, provider) = test_state(directory.path()).await;
7924 let identity = pv_events::AgentIdentity {
7925 version: "1".into(),
7926 charter: "c".into(),
7927 instance_uid: "i".into(),
7928 ttl_remaining_s: None,
7929 };
7930 let session = state.sessions.get_or_open(
7931 "client-managed-multi-turn",
7932 crate::session::Workflow::Signed,
7933 &identity,
7934 &state.config.breaker,
7935 );
7936 let stream = empty_abort_stream(&state, &session, identity, false);
7937 drop(stream);
7938 for _ in 0..20 {
7941 tokio::time::sleep(Duration::from_millis(10)).await;
7942 if session.is_closed() {
7943 break;
7944 }
7945 }
7946 assert!(
7947 !session.is_closed(),
7948 "a client-managed (non-ephemeral) session must NOT be auto-closed \
7949 when a turn aborts with nothing captured — multi-turn sessions \
7950 close only via POST /close or the idle sweeper, otherwise the \
7951 client's retry with the same id fragments the audit chain"
7952 );
7953 provider.abort();
7954 }
7955
7956 #[tokio::test]
7960 async fn ephemeral_empty_abort_closes_session_in_background() {
7961 let directory = tempfile::tempdir().unwrap();
7962 let (state, provider) = test_state(directory.path()).await;
7963 let identity = pv_events::AgentIdentity {
7964 version: "1".into(),
7965 charter: "c".into(),
7966 instance_uid: "i".into(),
7967 ttl_remaining_s: None,
7968 };
7969 let session = state.sessions.get_or_open(
7970 "ephemeral-empty-abort",
7971 crate::session::Workflow::Signed,
7972 &identity,
7973 &state.config.breaker,
7974 );
7975 let stream = empty_abort_stream(&state, &session, identity, true);
7976 drop(stream);
7977 let deadline = tokio::time::Instant::now() + Duration::from_secs(3);
7978 while !session.is_closed() && tokio::time::Instant::now() < deadline {
7979 tokio::time::sleep(Duration::from_millis(10)).await;
7980 }
7981 assert!(
7982 session.is_closed(),
7983 "an ephemeral one-shot aborted before any upstream byte must be \
7984 closed by the background abort-cleanup spawn, not linger until \
7985 the idle sweeper"
7986 );
7987 provider.abort();
7988 }
7989
7990 #[tokio::test]
7991 async fn duplicate_tool_id_replays_outcome_without_reexecution() {
7992 let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
7993 let tool_address = tool_listener.local_addr().unwrap();
7994 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
7995 let tool_router = Router::new()
7996 .route("/mcp", post(counting_tool))
7997 .with_state(Arc::clone(&calls));
7998 let tool_server = tokio::spawn(async move {
7999 axum::serve(tool_listener, tool_router).await.unwrap();
8000 });
8001 let directory = tempfile::tempdir().unwrap();
8002 let (mut state, provider) = test_state(directory.path()).await;
8003 Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
8004 Some(format!("http://{tool_address}/mcp"));
8005 let app = build_router(state.clone());
8006 let body = r#"{"jsonrpc":"2.0","id":"once","method":"tools/call","params":{"name":"read","arguments":{"id":7}}}"#;
8007 for _ in 0..2 {
8008 let response = app
8009 .clone()
8010 .oneshot(
8011 Request::builder()
8012 .method("POST")
8013 .uri("/v1/mcp")
8014 .header("x-pv-session", "tool-once")
8015 .body(Body::from(body))
8016 .unwrap(),
8017 )
8018 .await
8019 .unwrap();
8020 assert_eq!(response.status(), StatusCode::OK);
8021 assert_eq!(
8022 axum::body::to_bytes(response.into_body(), 64 * 1024)
8023 .await
8024 .unwrap(),
8025 body
8026 );
8027 }
8028 assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
8029 let session = state.sessions.get("tool-once").unwrap();
8030 assert_eq!(
8031 session
8032 .totals
8033 .tool_calls
8034 .load(std::sync::atomic::Ordering::Acquire),
8035 1
8036 );
8037 let audited_path = std::fs::read_dir(directory.path().join(crate::spool::TOOL_EXECUTIONS))
8038 .unwrap()
8039 .map(|entry| entry.unwrap().path())
8040 .find(|path| path.extension().is_some_and(|extension| extension == "audited"))
8041 .unwrap();
8042 std::fs::remove_file(audited_path).unwrap();
8043 let unaudited_replay = app
8044 .clone()
8045 .oneshot(
8046 Request::builder()
8047 .method("POST")
8048 .uri("/v1/mcp")
8049 .header("x-pv-session", "tool-once")
8050 .body(Body::from(body))
8051 .unwrap(),
8052 )
8053 .await
8054 .unwrap();
8055 assert_eq!(unaudited_replay.status(), StatusCode::OK);
8056 assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
8057 assert_eq!(
8058 session
8059 .totals
8060 .tool_calls
8061 .load(std::sync::atomic::Ordering::Acquire),
8062 1,
8063 "unaudited outcome recovery must not spend tool budget again"
8064 );
8065
8066 let changed_request = app
8067 .clone()
8068 .oneshot(
8069 Request::builder()
8070 .method("POST")
8071 .uri("/v1/mcp")
8072 .header("x-pv-session", "tool-once")
8073 .body(Body::from(
8074 r#"{"jsonrpc":"2.0","id":"once","method":"tools/call","params":{"name":"read","arguments":{"id":8}}}"#,
8075 ))
8076 .unwrap(),
8077 )
8078 .await
8079 .unwrap();
8080 assert_eq!(changed_request.status(), StatusCode::CONFLICT);
8081 assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
8082 tool_server.abort();
8083 provider.abort();
8084 }
8085
8086 #[tokio::test]
8091 async fn upstream_auth_header_reaches_provider_wire() {
8092 type Captured = Arc<parking_lot::Mutex<Vec<(Option<String>, Option<String>)>>>;
8093 async fn capture_chat(
8094 State(captured): State<Captured>,
8095 headers: HeaderMap,
8096 Json(_): Json<Value>,
8097 ) -> Response {
8098 captured.lock().push((
8099 headers
8100 .get(axum::http::header::AUTHORIZATION)
8101 .and_then(|value| value.to_str().ok())
8102 .map(str::to_owned),
8103 headers
8104 .get("api-key")
8105 .and_then(|value| value.to_str().ok())
8106 .map(str::to_owned),
8107 ));
8108 Json(json!({"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 1}}))
8109 .into_response()
8110 }
8111 let captured: Captured = Arc::new(parking_lot::Mutex::new(Vec::new()));
8112 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8113 let address = listener.local_addr().unwrap();
8114 let provider_router = Router::new()
8115 .route("/v1/chat/completions", post(capture_chat))
8116 .route("/openai/deployments/d1/chat/completions", post(capture_chat))
8117 .with_state(Arc::clone(&captured));
8118 let provider = tokio::spawn(async move {
8119 axum::serve(listener, provider_router).await.unwrap();
8120 });
8121 let directory = tempfile::tempdir().unwrap();
8122 let key_path = directory.path().join("upstream.key");
8123 std::fs::write(&key_path, "sk-wire-test\n").unwrap();
8124 #[cfg(unix)]
8125 {
8126 use std::os::unix::fs::PermissionsExt as _;
8127 std::fs::set_permissions(&key_path, std::fs::Permissions::from_mode(0o600)).unwrap();
8128 }
8129 let spool = directory.path().to_string_lossy();
8130 let build_state = |config: crate::config::HarnessConfig| {
8131 AppState::new(
8132 config,
8133 Arc::new(InMemoryStore::new()),
8134 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
8135 Arc::new(NullBus),
8136 None,
8137 Arc::new(Ed25519Signer::from_seed(&[12; 32])),
8138 )
8139 .unwrap()
8140 };
8141 let send_chat = |state: AppState, authorization: Option<&'static str>| async move {
8142 let mut request = Request::builder()
8143 .method("POST")
8144 .uri("/v1/chat/completions")
8145 .header(axum::http::header::CONTENT_TYPE, "application/json")
8146 .header("x-pv-session", "auth-wire");
8147 if let Some(value) = authorization {
8148 request = request.header(axum::http::header::AUTHORIZATION, value);
8149 }
8150 let response = build_router(state)
8151 .oneshot(request.body(Body::from(chat_payload().to_string())).unwrap())
8152 .await
8153 .unwrap();
8154 assert_eq!(response.status(), StatusCode::OK);
8155 };
8156
8157 let mut config =
8159 crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
8160 config.upstream_api_key_file = Some(key_path.to_string_lossy().into_owned());
8161 send_chat(build_state(config), None).await;
8162 assert_eq!(
8163 captured.lock().pop().unwrap(),
8164 (Some("Bearer sk-wire-test".to_owned()), None),
8165 "Bearer credential must reach the provider"
8166 );
8167
8168 let mut config =
8170 crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
8171 config.upstream_api_key_file = Some(key_path.to_string_lossy().into_owned());
8172 config.upstream_auth_header = "api-key".into();
8173 config.upstream_auth_scheme = String::new();
8174 config.upstream_chat_path = "/openai/deployments/d1/chat/completions".into();
8175 send_chat(build_state(config), None).await;
8176 assert_eq!(
8177 captured.lock().pop().unwrap(),
8178 (None, Some("sk-wire-test".to_owned())),
8179 "raw key must reach the provider on the configured path and header"
8180 );
8181
8182 let mut config =
8184 crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
8185 config.upstream_authorization_passthrough = true;
8186 send_chat(build_state(config), Some("Bearer client-owned-key")).await;
8187 assert_eq!(
8188 captured.lock().pop().unwrap(),
8189 (Some("Bearer client-owned-key".to_owned()), None),
8190 "passthrough must relay the client credential"
8191 );
8192
8193 let config = crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
8200 let request = Request::builder()
8201 .method("POST")
8202 .uri("/v1/chat/completions")
8203 .header(axum::http::header::CONTENT_TYPE, "application/json")
8204 .header("x-pv-session", "auth-wire")
8205 .header(axum::http::header::AUTHORIZATION, "Bearer client-owned-key")
8206 .body(Body::from(chat_payload().to_string()))
8207 .unwrap();
8208 let response = build_router(build_state(config)).oneshot(request).await.unwrap();
8209 assert_eq!(
8210 response.status(),
8211 StatusCode::UNAUTHORIZED,
8212 "presenting a bearer with no validator configured must be refused rather than \
8213 silently recorded as anonymous — that would let a signed request end up \
8214 attributed to `charter=anonymous` in the receipt"
8215 );
8216 assert!(
8217 captured.lock().is_empty(),
8218 "no credential must reach upstream when the request was refused"
8219 );
8220 provider.abort();
8221 }
8222
8223 #[tokio::test]
8226 async fn tool_upstream_bearer_reaches_tool_server() {
8227 type Captured = Arc<parking_lot::Mutex<Vec<Option<String>>>>;
8228 async fn capture_tool(State(captured): State<Captured>, headers: HeaderMap, body: Bytes) -> Response {
8229 captured.lock().push(
8230 headers
8231 .get(axum::http::header::AUTHORIZATION)
8232 .and_then(|value| value.to_str().ok())
8233 .map(str::to_owned),
8234 );
8235 (StatusCode::OK, body).into_response()
8236 }
8237 let captured: Captured = Arc::new(parking_lot::Mutex::new(Vec::new()));
8238 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8239 let address = listener.local_addr().unwrap();
8240 let tool_router = Router::new()
8241 .route("/mcp", post(capture_tool))
8242 .with_state(Arc::clone(&captured));
8243 let tool_server = tokio::spawn(async move {
8244 axum::serve(listener, tool_router).await.unwrap();
8245 });
8246 let directory = tempfile::tempdir().unwrap();
8247 let token_path = directory.path().join("mcp.token");
8248 std::fs::write(&token_path, "tool-secret\n").unwrap();
8249 #[cfg(unix)]
8250 {
8251 use std::os::unix::fs::PermissionsExt as _;
8252 std::fs::set_permissions(&token_path, std::fs::Permissions::from_mode(0o600)).unwrap();
8253 }
8254 let (mut state, provider) = test_state(directory.path()).await;
8255 {
8256 let config = Arc::get_mut(&mut state.config).unwrap();
8257 config.tool_upstream_url = Some(format!("http://{address}/mcp"));
8258 config.tool_upstream_bearer_file = Some(token_path.to_string_lossy().into_owned());
8259 }
8260 state.tool_auth = crate::pipeline::resolve_tool_auth(&state.config).unwrap();
8262 let response = build_router(state)
8263 .oneshot(
8264 Request::builder()
8265 .method("POST")
8266 .uri("/v1/mcp")
8267 .header("x-pv-session", "tool-auth")
8268 .body(Body::from(
8269 r#"{"jsonrpc":"2.0","id":"t1","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
8270 ))
8271 .unwrap(),
8272 )
8273 .await
8274 .unwrap();
8275 assert_eq!(response.status(), StatusCode::OK);
8276 assert_eq!(
8277 captured.lock().pop().unwrap().as_deref(),
8278 Some("Bearer tool-secret"),
8279 "tool bearer must reach the tool server"
8280 );
8281 tool_server.abort();
8282 provider.abort();
8283 }
8284
8285 #[tokio::test]
8286 async fn startup_refuses_pending_and_unaudited_tool_effects() {
8287 let directory = tempfile::tempdir().unwrap();
8288 let control_key = [31; 32];
8289 let mut headers = HeaderMap::new();
8290 headers.insert("x-pv-session", HeaderValue::from_static("tool-recovery"));
8291 let body = br#"{"jsonrpc":"2.0","id":"recover","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
8292 let mut execution =
8293 ToolExecution::from_request(&directory.path().to_string_lossy(), &headers, body, control_key)
8294 .unwrap();
8295 execution
8296 .bind_principal(&pv_events::AgentIdentity {
8297 version: "dev".to_owned(),
8298 charter: "anonymous".into(),
8299 instance_uid: "anonymous".to_owned(),
8300 ttl_remaining_s: None,
8301 })
8302 .unwrap();
8303 execution.claim().await.unwrap();
8304 assert!(
8305 ensure_no_unresolved_tool_executions(directory.path(), &control_key)
8306 .await
8307 .is_err()
8308 );
8309 execution
8310 .persist(&ToolOutcome {
8311 status: 200,
8312 body_hex: hex::encode(body),
8313 content_type: Some("application/json".to_owned()),
8314 })
8315 .await
8316 .unwrap();
8317 assert!(
8318 ensure_no_unresolved_tool_executions(directory.path(), &control_key)
8319 .await
8320 .is_err()
8321 );
8322 execution.mark_audited().await.unwrap();
8323 ensure_no_unresolved_tool_executions(directory.path(), &control_key)
8324 .await
8325 .unwrap();
8326 }
8327
8328 #[tokio::test]
8336 async fn oversize_tool_intent_is_quarantined_not_loaded() {
8337 let directory = tempfile::tempdir().unwrap();
8338 let control_key = [32; 32];
8339 let executions = directory.path().join(crate::spool::TOOL_EXECUTIONS);
8340 std::fs::create_dir_all(&executions).unwrap();
8341 let plant = executions.join("oversizekey.intent.json");
8342 std::fs::write(&plant, vec![b'x'; 1024 * 1024 + 1]).unwrap();
8344
8345 let unresolved = unresolved_tool_sessions(directory.path(), &control_key)
8346 .await
8347 .expect("an oversize plant must not abort the recovery scan");
8348 assert!(unresolved.is_empty(), "{unresolved:?}");
8349 assert!(
8350 !plant.exists(),
8351 "the oversize intent must be quarantined out of the scan glob"
8352 );
8353 let quarantined = std::fs::read_dir(&executions)
8354 .unwrap()
8355 .filter_map(Result::ok)
8356 .any(|entry| {
8357 entry
8358 .file_name()
8359 .to_str()
8360 .is_some_and(|name| name.ends_with(".intent.torn"))
8361 });
8362 assert!(quarantined, "quarantined plant must stay on disk as evidence");
8363 }
8364
8365 #[cfg(unix)]
8372 #[tokio::test]
8373 async fn unreadable_tool_intent_aborts_the_scan_not_quarantines() {
8374 use std::os::unix::fs::PermissionsExt as _;
8375 let directory = tempfile::tempdir().unwrap();
8376 let control_key = [33; 32];
8377 let executions = directory.path().join(crate::spool::TOOL_EXECUTIONS);
8378 std::fs::create_dir_all(&executions).unwrap();
8379 let unreadable = executions.join("unreadablekey.intent.json");
8380 std::fs::write(&unreadable, b"sealed-bytes").unwrap();
8381 std::fs::set_permissions(&unreadable, std::fs::Permissions::from_mode(0o000)).unwrap();
8382 if std::fs::read(&unreadable).is_ok() {
8383 return;
8385 }
8386
8387 let outcome = unresolved_tool_sessions(directory.path(), &control_key).await;
8388 assert!(
8389 outcome.is_err(),
8390 "a permission fault must abort the scan, not quarantine evidence"
8391 );
8392 assert!(
8393 unreadable.exists(),
8394 "the unreadable intent must be left in place for the operator"
8395 );
8396 }
8397
8398 struct BlockingRefundStore {
8399 inner: InMemoryStore,
8400 batches: parking_lot::Mutex<Vec<Vec<pv_state::Refund>>>,
8401 entered: tokio::sync::Notify,
8402 released: parking_lot::Mutex<bool>,
8403 wake: parking_lot::Condvar,
8404 fallback_used: std::sync::atomic::AtomicBool,
8405 }
8406
8407 impl StateStore for BlockingRefundStore {
8408 fn add(&self, key: &str, delta: u64) -> Result<u64, StateError> {
8409 self.inner.add(key, delta)
8410 }
8411
8412 fn get(&self, key: &str) -> Result<u64, StateError> {
8413 self.inner.get(key)
8414 }
8415
8416 fn try_spend(&self, key: &str, amount: u64, limit: u64) -> Result<bool, StateError> {
8417 self.inner.try_spend(key, amount, limit)
8418 }
8419
8420 fn try_spend_many(&self, spends: &[Spend]) -> Result<pv_state::TrySpendOutcome, StateError> {
8421 self.inner.try_spend_many(spends)
8422 }
8423
8424 fn remove(&self, key: &str) {
8425 self.inner.remove(key);
8426 }
8427
8428 fn refund_many(&self, refunds: &[pv_state::Refund]) {
8429 let first = {
8430 let mut batches = self.batches.lock();
8431 batches.push(refunds.to_vec());
8432 batches.len() == 1
8433 };
8434 if first {
8435 self.entered.notify_one();
8436 let deadline = std::time::Instant::now() + Duration::from_secs(3);
8437 let mut released = self.released.lock();
8438 while !*released {
8439 let remaining = deadline.saturating_duration_since(std::time::Instant::now());
8440 if remaining.is_zero() || self.wake.wait_for(&mut released, remaining).timed_out() {
8441 self.fallback_used
8444 .store(true, std::sync::atomic::Ordering::Release);
8445 break;
8446 }
8447 }
8448 }
8449 self.inner.refund_many(refunds);
8450 }
8451 }
8452
8453 #[tokio::test(flavor = "current_thread")]
8454 async fn mcp_refund_keeps_reactor_live_and_survives_caller_cancellation() {
8455 let directory = tempfile::tempdir().unwrap();
8456 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
8457 let address = listener.local_addr().unwrap();
8458 drop(listener);
8460 let mut config = crate::config::HarnessConfig::for_tests(
8461 &format!("http://{address}"),
8462 &directory.path().to_string_lossy(),
8463 &directory.path().to_string_lossy(),
8464 );
8465 config.tool_upstream_url = Some(format!("http://{address}/mcp"));
8466 config.budget.max_total_tool_calls = Some(1);
8467 config.budget.max_tool_calls.insert("read".into(), 1);
8468 config.principal_budget = Some(config.budget.clone());
8469 config.allow_anonymous_principal_budget = true;
8470 config.mcp_concurrency = 1;
8471 let store = Arc::new(BlockingRefundStore {
8472 inner: InMemoryStore::new(),
8473 batches: parking_lot::Mutex::new(Vec::new()),
8474 entered: tokio::sync::Notify::new(),
8475 released: parking_lot::Mutex::new(false),
8476 wake: parking_lot::Condvar::new(),
8477 fallback_used: std::sync::atomic::AtomicBool::new(false),
8478 });
8479 let sandbox = Arc::new(
8480 Sandbox::new(
8481 SandboxConfig {
8482 budget: config.budget.clone(),
8483 ..SandboxConfig::default()
8484 },
8485 Vec::new(),
8486 )
8487 .unwrap(),
8488 );
8489 let state = AppState::new(
8490 config,
8491 store.clone(),
8492 sandbox,
8493 Arc::new(NullBus),
8494 None,
8495 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
8496 )
8497 .unwrap();
8498 let app = build_router(state.clone());
8499 let request = || {
8500 Request::builder()
8501 .method("POST")
8502 .uri("/v1/mcp")
8503 .header("x-pv-session", "refund-cancel")
8504 .body(Body::from(
8505 r#"{"jsonrpc":"2.0","id":"refund","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
8506 ))
8507 .unwrap()
8508 };
8509 let mut expected_keys = Vec::new();
8510 for prefix in [
8511 pv_state::ActionBudget::session_prefix("refund-cancel"),
8512 pv_state::ActionBudget::principal_prefix("anonymous"),
8513 ] {
8514 expected_keys.push(format!("{prefix}total_calls"));
8515 expected_keys.push(format!("{prefix}tool:read"));
8516 }
8517 let first = tokio::spawn(app.clone().oneshot(request()));
8518 tokio::time::timeout(Duration::from_secs(5), store.entered.notified())
8519 .await
8520 .unwrap();
8521 let progressed_while_refund_waited = !store.fallback_used.load(std::sync::atomic::Ordering::Acquire);
8525 let charged_before_release: Vec<_> =
8526 expected_keys.iter().map(|key| store.get(key).unwrap()).collect();
8527 first.abort();
8528 *store.released.lock() = true;
8529 store.wake.notify_all();
8530 let caller_cancelled = first.await.is_err_and(|error| error.is_cancelled());
8531 tokio::time::timeout(Duration::from_secs(5), state.mcp_inflight.wait_drained())
8532 .await
8533 .unwrap();
8534 let first_batches = store.batches.lock().clone();
8535 let after_first: Vec<_> = expected_keys.iter().map(|key| store.get(key).unwrap()).collect();
8536 let retry = app.oneshot(request()).await.unwrap();
8537 let retry_status = retry.status();
8538 tokio::time::timeout(Duration::from_secs(5), state.mcp_inflight.wait_drained())
8539 .await
8540 .unwrap();
8541 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
8542 .await
8543 .unwrap();
8544 let batches = store.batches.lock().clone();
8545 let session = state.sessions.get("refund-cancel").unwrap();
8546
8547 assert!(
8548 progressed_while_refund_waited,
8549 "synchronous budget compensation blocked the only Tokio thread until the fallback released it"
8550 );
8551 assert!(
8552 charged_before_release.iter().all(|value| *value == 1),
8553 "both session and principal ledgers must actually be debited before compensation"
8554 );
8555 assert!(
8556 caller_cancelled,
8557 "the outer caller must be cancelled while compensation is in progress"
8558 );
8559 assert_eq!(
8560 first_batches.len(),
8561 2,
8562 "one session and one principal refund must complete despite cancellation"
8563 );
8564 assert!(
8565 after_first.iter().all(|value| *value == 0),
8566 "both ledgers must be refunded before retry"
8567 );
8568 assert_eq!(
8569 retry_status,
8570 StatusCode::BAD_GATEWAY,
8571 "retry must attempt the unexecuted call, not return a pending claim or spent-budget refusal"
8572 );
8573 assert_eq!(
8574 batches.len(),
8575 4,
8576 "each of two attempts refunds each ledger exactly once"
8577 );
8578 for key in expected_keys {
8579 let amounts: Vec<_> = batches
8580 .iter()
8581 .flatten()
8582 .filter(|refund| refund.key == key)
8583 .map(|refund| refund.amount)
8584 .collect();
8585 assert_eq!(amounts, [1, 1], "unexpected compensation for {key}");
8586 assert_eq!(store.get(&key).unwrap(), 0);
8587 }
8588 assert_eq!(state.mcp_admission.available_permits(), 1);
8589 assert_eq!(session.active_streams_count(), 0);
8590 assert_eq!(session.pending_jobs_count(), 0);
8591 }
8592
8593 async fn assert_orphan_mcp_intent_refuses_startup(unaudited: bool) {
8594 let directory = tempfile::tempdir().unwrap();
8595 let (state, provider) = test_state(directory.path()).await;
8596 let session_id = "orphan-uncertain";
8597 let mut headers = HeaderMap::new();
8598 headers.insert("x-pv-session", HeaderValue::from_static("orphan-uncertain"));
8599 let body = br#"{"jsonrpc":"2.0","id":"orphan","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
8600 let mut execution =
8601 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
8602 .unwrap();
8603 execution
8604 .bind_principal(&pv_events::AgentIdentity {
8605 version: "dev".to_owned(),
8606 charter: "anonymous".into(),
8607 instance_uid: "anonymous".to_owned(),
8608 ttl_remaining_s: None,
8609 })
8610 .unwrap();
8611 execution.claim().await.unwrap();
8615 if unaudited {
8616 execution
8617 .persist(&ToolOutcome {
8618 status: 200,
8619 body_hex: hex::encode(b"effect happened"),
8620 content_type: Some("text/plain".to_owned()),
8621 })
8622 .await
8623 .unwrap();
8624 }
8625 let intent_before = std::fs::read(&execution.intent_path).unwrap();
8626 assert!(state.sessions.get(session_id).is_none());
8627 assert!(!std::fs::read_dir(directory.path())
8628 .unwrap()
8629 .filter_map(Result::ok)
8630 .any(|entry| entry.file_name().to_string_lossy().ends_with(".session.json")));
8631 let recovered = tokio::time::timeout(
8632 Duration::from_secs(5),
8633 state
8634 .finalizer
8635 .recover_spooled_sessions(&state.sessions, &state.config.breaker),
8636 )
8637 .await
8638 .unwrap();
8639 let mut minted_after_new_traffic = false;
8640 if recovered.is_ok() {
8641 let app = build_router(state.clone());
8645 let chat = app.clone().oneshot(chat_request(session_id)).await.unwrap();
8646 axum::body::to_bytes(chat.into_body(), 64 * 1024).await.unwrap();
8647 let _ = app
8648 .clone()
8649 .oneshot(
8650 Request::builder()
8651 .method("POST")
8652 .uri("/v1/sessions/orphan-uncertain/close")
8653 .body(Body::empty())
8654 .unwrap(),
8655 )
8656 .await
8657 .unwrap();
8658 let _ = app
8659 .oneshot(
8660 Request::builder()
8661 .method("POST")
8662 .uri("/v1/sessions/orphan-uncertain/promote")
8663 .body(Body::empty())
8664 .unwrap(),
8665 )
8666 .await
8667 .unwrap();
8668 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
8669 .await
8670 .unwrap();
8671 minted_after_new_traffic = state
8672 .sessions
8673 .get(session_id)
8674 .is_some_and(|session| session.receipt.lock().is_some());
8675 }
8676 provider.abort();
8677 let _ = provider.await;
8678 assert!(!minted_after_new_traffic,
8679 "startup accepted an authenticated orphan intent, then new traffic minted a receipt omitting its uncertain effect");
8680 let error = recovered.expect_err(
8681 "an authenticated unresolved intent without recoverable session metadata must refuse startup",
8682 );
8683 assert!(error.to_string().contains("unresolved tool execution"));
8684 assert_eq!(std::fs::read(&execution.intent_path).unwrap(), intent_before);
8685 assert!(
8686 state.sessions.get(session_id).is_none(),
8687 "recovery must not invent session attribution"
8688 );
8689 }
8690
8691 #[tokio::test]
8692 async fn mcp_durable_orphan_pending_intent_refuses_startup_and_new_receipt() {
8693 assert_orphan_mcp_intent_refuses_startup(false).await;
8694 }
8695
8696 #[tokio::test]
8697 async fn mcp_durable_orphan_unaudited_intent_refuses_startup_and_new_receipt() {
8698 assert_orphan_mcp_intent_refuses_startup(true).await;
8699 }
8700
8701 async fn assert_evicted_mcp_intent_refuses_receipt(workflow: &str) {
8702 let directory = tempfile::tempdir().unwrap();
8703 let (state, provider) = test_state(directory.path()).await;
8704 let session_id = "evicted-uncertain";
8705 let identity = pv_events::AgentIdentity {
8706 version: "dev".to_owned(),
8707 charter: "anonymous".into(),
8708 instance_uid: "anonymous".to_owned(),
8709 ttl_remaining_s: None,
8710 };
8711 let original = state.sessions.get_or_open(
8712 session_id,
8713 crate::session::Workflow::parse(workflow).unwrap(),
8714 &identity,
8715 &state.config.breaker,
8716 );
8717 let mut headers = HeaderMap::new();
8718 headers.insert("x-pv-session", HeaderValue::from_static("evicted-uncertain"));
8719 let body = br#"{"jsonrpc":"2.0","id":"evicted","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
8720 let mut execution =
8721 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
8722 .unwrap();
8723 execution.bind_principal(&identity).unwrap();
8724 execution.claim().await.unwrap();
8725 let intent_before = std::fs::read(&execution.intent_path).unwrap();
8726 original.mark_capture_failed();
8727 state.sessions.remove(session_id);
8730 let app = build_router(state.clone());
8731 let request = Request::builder()
8732 .method("POST")
8733 .uri("/v1/chat/completions")
8734 .header("content-type", "application/json")
8735 .header("x-pv-session", session_id)
8736 .header("x-pv-workflow", workflow)
8737 .body(Body::from(serde_json::to_vec(&chat_payload()).unwrap()))
8738 .unwrap();
8739 let chat = app.clone().oneshot(request).await.unwrap();
8740 assert_eq!(chat.status(), StatusCode::OK);
8741 axum::body::to_bytes(chat.into_body(), 64 * 1024).await.unwrap();
8742 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
8743 .await
8744 .unwrap();
8745 let fresh = state.sessions.get(session_id).unwrap();
8746 assert!(!Arc::ptr_eq(&original, &fresh));
8747 assert!(
8748 !fresh.capture_failed(),
8749 "the durable guard cannot rely on the old RAM flag"
8750 );
8751 let close = app
8752 .clone()
8753 .oneshot(
8754 Request::builder()
8755 .method("POST")
8756 .uri("/v1/sessions/evicted-uncertain/close")
8757 .body(Body::empty())
8758 .unwrap(),
8759 )
8760 .await
8761 .unwrap()
8762 .status();
8763 let promote = app
8764 .oneshot(
8765 Request::builder()
8766 .method("POST")
8767 .uri("/v1/sessions/evicted-uncertain/promote")
8768 .body(Body::empty())
8769 .unwrap(),
8770 )
8771 .await
8772 .unwrap()
8773 .status();
8774 provider.abort();
8775 let _ = provider.await;
8776 assert!(
8777 fresh.receipt.lock().is_none(),
8778 "evicting the RAM quarantine must not permit a receipt omitting the durable uncertain effect"
8779 );
8780 assert_eq!(close, StatusCode::SERVICE_UNAVAILABLE);
8781 assert!(!promote.is_success());
8782 assert_eq!(std::fs::read(&execution.intent_path).unwrap(), intent_before);
8783 assert!(
8784 !fresh.capture_failed(),
8785 "the durable refusal must preserve retryable audit completion"
8786 );
8787 }
8788
8789 #[tokio::test]
8790 async fn mcp_durable_evicted_signed_quarantine_cannot_mint_receipt() {
8791 assert_evicted_mcp_intent_refuses_receipt("signed").await;
8792 }
8793
8794 #[tokio::test]
8795 async fn mcp_durable_evicted_unsigned_quarantine_cannot_promote() {
8796 assert_evicted_mcp_intent_refuses_receipt("unsigned").await;
8797 }
8798
8799 #[tokio::test]
8800 async fn mcp_durable_closed_unsigned_artifact_checks_durable_intent_before_promotion() {
8801 let directory = tempfile::tempdir().unwrap();
8802 let (state, provider) = test_state(directory.path()).await;
8803 let app = build_router(state.clone());
8804 let chat = app
8805 .clone()
8806 .oneshot(chat_request("closed-uncertain"))
8807 .await
8808 .unwrap();
8809 axum::body::to_bytes(chat.into_body(), 64 * 1024).await.unwrap();
8810 let close = app
8811 .clone()
8812 .oneshot(
8813 Request::builder()
8814 .method("POST")
8815 .uri("/v1/sessions/closed-uncertain/close")
8816 .body(Body::empty())
8817 .unwrap(),
8818 )
8819 .await
8820 .unwrap();
8821 assert_eq!(close.status(), StatusCode::OK);
8822 let session = state.sessions.get("closed-uncertain").unwrap();
8823 assert!(session.is_closed());
8824 assert!(session.receipt.lock().is_none());
8825 let mut headers = HeaderMap::new();
8826 headers.insert("x-pv-session", HeaderValue::from_static("closed-uncertain"));
8827 let body = br#"{"jsonrpc":"2.0","id":"closed","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
8828 let mut execution =
8829 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
8830 .unwrap();
8831 execution.bind_principal(&session.current_identity()).unwrap();
8832 execution.claim().await.unwrap();
8833 let before = std::fs::read(&execution.intent_path).unwrap();
8834 let promotion = app
8835 .oneshot(
8836 Request::builder()
8837 .method("POST")
8838 .uri("/v1/sessions/closed-uncertain/promote")
8839 .body(Body::empty())
8840 .unwrap(),
8841 )
8842 .await
8843 .unwrap()
8844 .status();
8845 provider.abort();
8846 let _ = provider.await;
8847 assert!(
8848 session.receipt.lock().is_none(),
8849 "a closed unsigned artifact must not bypass the durable uncertainty guard"
8850 );
8851 assert_eq!(promotion, StatusCode::SERVICE_UNAVAILABLE);
8852 assert_eq!(std::fs::read(&execution.intent_path).unwrap(), before);
8853 assert!(!session.capture_failed());
8854 }
8855
8856 #[tokio::test]
8857 async fn mcp_durable_missing_audit_marker_keeps_finalization_retryable() {
8858 let directory = tempfile::tempdir().unwrap();
8859 let (state, provider) = test_state(directory.path()).await;
8860 let app = build_router(state.clone());
8861 let chat = app.clone().oneshot(chat_request("audit-retry")).await.unwrap();
8862 axum::body::to_bytes(chat.into_body(), 64 * 1024).await.unwrap();
8863 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
8864 .await
8865 .unwrap();
8866 let session = state.sessions.get("audit-retry").unwrap();
8867 let mut headers = HeaderMap::new();
8868 headers.insert("x-pv-session", HeaderValue::from_static("audit-retry"));
8869 let body = br#"{"jsonrpc":"2.0","id":"audit","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
8870 let mut execution =
8871 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
8872 .unwrap();
8873 execution.bind_principal(&session.current_identity()).unwrap();
8874 execution.claim().await.unwrap();
8875 execution
8876 .persist(&ToolOutcome {
8877 status: 200,
8878 body_hex: hex::encode(b"captured outcome"),
8879 content_type: None,
8880 })
8881 .await
8882 .unwrap();
8883 let close_request = || {
8884 Request::builder()
8885 .method("POST")
8886 .uri("/v1/sessions/audit-retry/close")
8887 .body(Body::empty())
8888 .unwrap()
8889 };
8890 let first = app.clone().oneshot(close_request()).await.unwrap().status();
8891 let refused_without_poison = !session.capture_failed()
8892 && !session.is_closed()
8893 && execution.intent_path.exists()
8894 && execution.outcome_path.exists();
8895 execution.mark_audited().await.unwrap();
8897 let second = app.clone().oneshot(close_request()).await.unwrap().status();
8898 let promotion = app
8899 .oneshot(
8900 Request::builder()
8901 .method("POST")
8902 .uri("/v1/sessions/audit-retry/promote")
8903 .body(Body::empty())
8904 .unwrap(),
8905 )
8906 .await
8907 .unwrap()
8908 .status();
8909 provider.abort();
8910 let _ = provider.await;
8911 assert_eq!(first, StatusCode::SERVICE_UNAVAILABLE);
8912 assert!(
8913 refused_without_poison,
8914 "an unfinished audit must remain retryable with its evidence intact"
8915 );
8916 assert_eq!(second, StatusCode::OK);
8917 assert_eq!(promotion, StatusCode::OK);
8918 assert!(session.receipt.lock().is_some());
8919 assert!(!session.capture_failed());
8920 }
8921
8922 #[tokio::test]
8923 async fn mcp_durable_lifecycle_scan_io_failure_preserves_evidence_and_can_retry() {
8924 let directory = tempfile::tempdir().unwrap();
8925 let (state, provider) = test_state(directory.path()).await;
8926 let app = build_router(state.clone());
8927 let chat = app.clone().oneshot(chat_request("scan-error")).await.unwrap();
8928 axum::body::to_bytes(chat.into_body(), 64 * 1024).await.unwrap();
8929 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
8930 .await
8931 .unwrap();
8932 let session = state.sessions.get("scan-error").unwrap();
8933 let execution_dir = directory.path().join(crate::spool::TOOL_EXECUTIONS);
8934 let evidence = b"operator recovery evidence at an invalid directory path";
8935 std::fs::write(&execution_dir, evidence).unwrap();
8936 let close_request = || {
8937 Request::builder()
8938 .method("POST")
8939 .uri("/v1/sessions/scan-error/close")
8940 .body(Body::empty())
8941 .unwrap()
8942 };
8943 let failed = app.clone().oneshot(close_request()).await.unwrap().status();
8944 let retained = std::fs::read(&execution_dir).unwrap() == evidence;
8945 let issued_on_failure = session.receipt.lock().is_some() || session.atif_path.lock().is_some();
8946 let retryable = !session.capture_failed() && !session.is_closed();
8947 std::fs::remove_file(&execution_dir).unwrap();
8948 let repaired = app.oneshot(close_request()).await.unwrap().status();
8949 provider.abort();
8950 let _ = provider.await;
8951 assert_eq!(failed, StatusCode::SERVICE_UNAVAILABLE);
8952 assert!(retained);
8953 assert!(!issued_on_failure);
8954 assert!(retryable);
8955 assert_eq!(repaired, StatusCode::OK);
8956 }
8957
8958 #[test]
8959 fn mcp_durable_readonly_scan_preserves_evidence_on_cap_and_authentication_errors() {
8960 let directory = tempfile::tempdir().unwrap();
8961 for name in ["one.evidence", "two.evidence", "three.evidence"] {
8962 std::fs::write(directory.path().join(name), b"retained evidence").unwrap();
8963 }
8964 let capped = ensure_tool_executions_resolved_with_limit(directory.path(), "session", &[7; 32], 2);
8965 assert!(capped.unwrap_err().contains("bounded scan limit"));
8966 for name in ["one.evidence", "two.evidence", "three.evidence"] {
8967 assert_eq!(
8968 std::fs::read(directory.path().join(name)).unwrap(),
8969 b"retained evidence"
8970 );
8971 }
8972 let writing_claim = directory.path().join("partial.intent.json");
8973 std::fs::write(&writing_claim, b"unfinished concurrent claim").unwrap();
8974 let unauthenticated =
8975 ensure_tool_executions_resolved_with_limit(directory.path(), "session", &[7; 32], 10);
8976 assert!(unauthenticated.unwrap_err().contains("authenticated"));
8977 assert_eq!(
8978 std::fs::read(writing_claim).unwrap(),
8979 b"unfinished concurrent claim"
8980 );
8981 assert_eq!(
8982 std::fs::read_dir(directory.path()).unwrap().count(),
8983 4,
8984 "the lifecycle scan must never rename or delete another live request's incomplete files"
8985 );
8986 }
8987
8988 struct WithheldHeadersToolState {
8989 effects: std::sync::atomic::AtomicUsize,
8990 release: tokio::sync::Notify,
8991 returned: std::sync::atomic::AtomicBool,
8992 terminated: std::sync::atomic::AtomicBool,
8993 }
8994
8995 struct WithheldHeadersGuard(Arc<WithheldHeadersToolState>);
8996
8997 impl Drop for WithheldHeadersGuard {
8998 fn drop(&mut self) {
8999 self.0
9000 .terminated
9001 .store(true, std::sync::atomic::Ordering::Release);
9002 }
9003 }
9004
9005 async fn tool_effect_before_headers(State(state): State<Arc<WithheldHeadersToolState>>) -> Response {
9006 let _guard = WithheldHeadersGuard(Arc::clone(&state));
9007 state.effects.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9010 let _ = tokio::time::timeout(Duration::from_secs(6), state.release.notified()).await;
9011 state.returned.store(true, std::sync::atomic::Ordering::Release);
9012 Json(json!({"ok": true})).into_response()
9013 }
9014
9015 async fn assert_mcp_header_deadline(workflow: &str, cancel_caller: bool, restart: bool) {
9016 let directory = tempfile::tempdir().unwrap();
9017 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9018 let address = listener.local_addr().unwrap();
9019 let observed = Arc::new(WithheldHeadersToolState {
9020 effects: std::sync::atomic::AtomicUsize::new(0),
9021 release: tokio::sync::Notify::new(),
9022 returned: std::sync::atomic::AtomicBool::new(false),
9023 terminated: std::sync::atomic::AtomicBool::new(false),
9024 });
9025 let tool = Router::new()
9026 .route("/mcp", post(tool_effect_before_headers))
9027 .with_state(Arc::clone(&observed));
9028 let server = tokio::spawn(async move { axum::serve(listener, tool).await.unwrap() });
9029 let mut config = crate::config::HarnessConfig::for_tests(
9030 &format!("http://{address}"),
9031 &directory.path().to_string_lossy(),
9032 &directory.path().to_string_lossy(),
9033 );
9034 config.default_workflow = workflow.to_owned();
9035 config.tool_upstream_url = Some(format!("http://{address}/mcp"));
9036 config.upstream_read_timeout_s = Some(10);
9038 config.mcp_request_timeout_s = Some(1);
9039 config.mcp_concurrency = 1;
9040 config.budget.max_total_tool_calls = Some(2);
9041 config.budget.max_tool_calls.insert("read".into(), 2);
9042 config.principal_budget = Some(config.budget.clone());
9043 config.allow_anonymous_principal_budget = true;
9044 let make_state = |store: Arc<InMemoryStore>| {
9045 AppState::new(
9046 config.clone(),
9047 store,
9048 Arc::new(
9049 Sandbox::new(
9050 SandboxConfig {
9051 budget: config.budget.clone(),
9052 ..SandboxConfig::default()
9053 },
9054 Vec::new(),
9055 )
9056 .unwrap(),
9057 ),
9058 Arc::new(NullBus),
9059 None,
9060 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
9061 )
9062 .unwrap()
9063 };
9064 let store = Arc::new(InMemoryStore::new());
9065 let state = make_state(Arc::clone(&store));
9066 let app = build_router(state.clone());
9067 let body = br#"{"jsonrpc":"2.0","id":"headers","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
9068 let request = || {
9069 Request::builder()
9070 .method("POST")
9071 .uri("/v1/mcp")
9072 .header("content-type", "application/json")
9073 .header("x-pv-session", "withheld-headers")
9074 .body(Body::from(Bytes::from_static(body)))
9075 .unwrap()
9076 };
9077 let first = tokio::spawn(app.clone().oneshot(request()));
9078 tokio::time::timeout(Duration::from_secs(5), async {
9079 while observed.effects.load(std::sync::atomic::Ordering::Acquire) == 0 {
9080 tokio::time::sleep(Duration::from_millis(10)).await;
9081 }
9082 })
9083 .await
9084 .unwrap();
9085 if cancel_caller {
9086 first.abort();
9087 }
9088 tokio::time::timeout(Duration::from_secs(5), state.mcp_inflight.wait_drained())
9089 .await
9090 .unwrap();
9091 let first_status = if cancel_caller {
9092 assert!(first.await.unwrap_err().is_cancelled());
9093 None
9094 } else {
9095 Some(first.await.unwrap().unwrap().status())
9096 };
9097 assert!(!observed.returned.load(std::sync::atomic::Ordering::Acquire));
9098 let retry = app.clone().oneshot(request()).await.unwrap();
9099 let retry_status = retry.status();
9100 let retry_body = axum::body::to_bytes(retry.into_body(), 64 * 1024).await.unwrap();
9101 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
9102 .await
9103 .unwrap();
9104 let session = state.sessions.get("withheld-headers").unwrap();
9105 let mut headers = HeaderMap::new();
9106 headers.insert("x-pv-session", HeaderValue::from_static("withheld-headers"));
9107 let mut execution =
9108 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
9109 .unwrap();
9110 execution.bind_principal(&session.current_identity()).unwrap();
9111 assert!(matches!(
9112 execution.load().await.unwrap(),
9113 ToolExecutionState::Pending
9114 ));
9115 assert!(execution.intent_path.exists());
9116 assert!(!execution.outcome_path.exists());
9117 assert!(!execution.audited_path.exists());
9118 for prefix in [
9119 pv_state::ActionBudget::session_prefix("withheld-headers"),
9120 pv_state::ActionBudget::principal_prefix("anonymous"),
9121 ] {
9122 for suffix in ["total_calls", "tool:read"] {
9123 assert_eq!(
9124 store.get(&format!("{prefix}{suffix}")).unwrap(),
9125 1,
9126 "uncertain execution must retain its charge without charging the refused retry"
9127 );
9128 }
9129 }
9130 assert_eq!(state.mcp_inflight.count(), 0);
9131 assert_eq!(state.mcp_admission.available_permits(), 1);
9132 assert_eq!(session.active_streams_count(), 0);
9133 assert_eq!(session.pending_jobs_count(), 0);
9134 drop(session);
9135 drop(app);
9136 let lifecycle_state = if restart {
9139 drop(state);
9140 let recovered = make_state(Arc::new(InMemoryStore::new()));
9141 recovered
9142 .finalizer
9143 .recover_spooled_sessions(&recovered.sessions, &recovered.config.breaker)
9144 .await
9145 .unwrap();
9146 assert!(recovered
9147 .sessions
9148 .get("withheld-headers")
9149 .unwrap()
9150 .capture_failed());
9151 recovered
9152 } else {
9153 state
9154 };
9155 let lifecycle_app = build_router(lifecycle_state.clone());
9156 let replay = lifecycle_app.clone().oneshot(request()).await.unwrap();
9157 let replay_status = replay.status();
9158 let close = lifecycle_app
9159 .clone()
9160 .oneshot(
9161 Request::builder()
9162 .method("POST")
9163 .uri("/v1/sessions/withheld-headers/close")
9164 .body(Body::empty())
9165 .unwrap(),
9166 )
9167 .await
9168 .unwrap();
9169 let close_status = close.status();
9170 let close_body = axum::body::to_bytes(close.into_body(), 64 * 1024).await.unwrap();
9171 let promote = lifecycle_app
9172 .clone()
9173 .oneshot(
9174 Request::builder()
9175 .method("POST")
9176 .uri("/v1/sessions/withheld-headers/promote")
9177 .body(Body::empty())
9178 .unwrap(),
9179 )
9180 .await
9181 .unwrap();
9182 let promote_status = promote.status();
9183 tokio::time::timeout(Duration::from_secs(5), lifecycle_state.worker.wait_idle())
9184 .await
9185 .unwrap();
9186 let session = lifecycle_state.sessions.get("withheld-headers").unwrap();
9187 let receipt_minted = session.receipt.lock().is_some();
9188 let retained_intent_after_close = execution.intent_path.exists();
9189 let retry_after_close = lifecycle_app.clone().oneshot(request()).await.unwrap().status();
9190 drop(session);
9191 drop(lifecycle_app);
9192 drop(lifecycle_state);
9193 let recovered_again = make_state(Arc::new(InMemoryStore::new()));
9194 tokio::time::timeout(
9195 Duration::from_secs(5),
9196 recovered_again
9197 .finalizer
9198 .recover_spooled_sessions(&recovered_again.sessions, &recovered_again.config.breaker),
9199 )
9200 .await
9201 .unwrap()
9202 .unwrap();
9203 let session = recovered_again.sessions.get("withheld-headers").unwrap();
9204 let still_quarantined = session.capture_failed();
9205 let recovered_receipt = session.receipt.lock().is_some();
9206 let recovered_app = build_router(recovered_again.clone());
9207 let retry_after_restart = recovered_app.clone().oneshot(request()).await.unwrap().status();
9208 let close_after_restart = recovered_app
9209 .clone()
9210 .oneshot(
9211 Request::builder()
9212 .method("POST")
9213 .uri("/v1/sessions/withheld-headers/close")
9214 .body(Body::empty())
9215 .unwrap(),
9216 )
9217 .await
9218 .unwrap();
9219 let close_after_restart_status = close_after_restart.status();
9220 let close_after_restart_body = axum::body::to_bytes(close_after_restart.into_body(), 64 * 1024)
9221 .await
9222 .unwrap();
9223 let promote_after_restart = recovered_app
9224 .oneshot(
9225 Request::builder()
9226 .method("POST")
9227 .uri("/v1/sessions/withheld-headers/promote")
9228 .body(Body::empty())
9229 .unwrap(),
9230 )
9231 .await
9232 .unwrap()
9233 .status();
9234 tokio::time::timeout(Duration::from_secs(5), recovered_again.worker.wait_idle())
9235 .await
9236 .unwrap();
9237 observed.release.notify_one();
9240 tokio::time::timeout(Duration::from_secs(2), async {
9241 while !observed.terminated.load(std::sync::atomic::Ordering::Acquire) {
9242 tokio::time::sleep(Duration::from_millis(10)).await;
9243 }
9244 })
9245 .await
9246 .unwrap();
9247 server.abort();
9248 let _ = server.await;
9249
9250 if let Some(status) = first_status {
9251 assert_eq!(status, StatusCode::BAD_GATEWAY);
9252 }
9253 assert_eq!(retry_status, StatusCode::CONFLICT);
9254 assert!(String::from_utf8_lossy(&retry_body).contains(TOOL_OUTCOME_UNCERTAIN));
9255 assert_eq!(replay_status, StatusCode::CONFLICT);
9256 assert_eq!(observed.effects.load(std::sync::atomic::Ordering::Acquire), 1);
9257 assert!(
9258 !receipt_minted,
9259 "uncertain tool execution must never mint a complete receipt"
9260 );
9261 assert_eq!(promote_status, StatusCode::CONFLICT);
9262 assert!(
9263 retained_intent_after_close,
9264 "closing a quarantine must retain the authenticated uncertain intent"
9265 );
9266 assert_eq!(retry_after_close, StatusCode::CONFLICT);
9267 assert!(
9268 still_quarantined,
9269 "quarantine must survive closing and another restart"
9270 );
9271 assert!(
9272 !recovered_receipt,
9273 "restart must not turn an uncertain execution into a complete receipt"
9274 );
9275 assert_eq!(retry_after_restart, StatusCode::CONFLICT);
9276 assert_eq!(promote_after_restart, StatusCode::CONFLICT);
9277 if close_after_restart_status == StatusCode::OK {
9278 assert_eq!(
9279 serde_json::from_slice::<Value>(&close_after_restart_body).unwrap(),
9280 json!({"kind": "already_closed"})
9281 );
9282 } else {
9283 assert_eq!(close_after_restart_status, StatusCode::CONFLICT);
9284 }
9285 assert!(session.receipt.lock().is_none());
9286 assert!(session.atif_path.lock().is_none());
9287 if restart && close_status == StatusCode::OK {
9290 assert_eq!(
9291 serde_json::from_slice::<Value>(&close_body).unwrap(),
9292 json!({"kind": "already_closed"})
9293 );
9294 } else {
9295 assert_eq!(close_status, StatusCode::CONFLICT);
9296 }
9297 }
9298
9299 #[tokio::test]
9300 async fn mcp_header_deadline_refuses_live_signed_receipt() {
9301 assert_mcp_header_deadline("signed", false, false).await;
9302 }
9303
9304 #[tokio::test]
9305 async fn mcp_header_deadline_refuses_live_unsigned_promotion() {
9306 assert_mcp_header_deadline("unsigned", false, false).await;
9307 }
9308
9309 #[tokio::test]
9310 async fn mcp_header_deadline_cancellation_keeps_live_intent() {
9311 assert_mcp_header_deadline("unsigned", true, false).await;
9312 }
9313
9314 #[tokio::test]
9315 async fn mcp_header_deadline_restart_quarantines_signed_session() {
9316 assert_mcp_header_deadline("signed", false, true).await;
9317 }
9318
9319 #[tokio::test]
9320 async fn mcp_header_deadline_cancel_then_restart_quarantines_unsigned_session() {
9321 assert_mcp_header_deadline("unsigned", true, true).await;
9322 }
9323
9324 #[tokio::test]
9325 async fn mcp_outcome_persist_failure_refuses_receipt_and_keeps_charge() {
9326 let directory = tempfile::tempdir().unwrap();
9327 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9328 let address = listener.local_addr().unwrap();
9329 let observed = Arc::new(WithheldHeadersToolState {
9330 effects: std::sync::atomic::AtomicUsize::new(0),
9331 release: tokio::sync::Notify::new(),
9332 returned: std::sync::atomic::AtomicBool::new(false),
9333 terminated: std::sync::atomic::AtomicBool::new(false),
9334 });
9335 let tool = Router::new()
9336 .route("/mcp", post(tool_effect_before_headers))
9337 .with_state(Arc::clone(&observed));
9338 let server = tokio::spawn(async move { axum::serve(listener, tool).await.unwrap() });
9339 let mut config = crate::config::HarnessConfig::for_tests(
9340 &format!("http://{address}"),
9341 &directory.path().to_string_lossy(),
9342 &directory.path().to_string_lossy(),
9343 );
9344 config.default_workflow = "signed".to_owned();
9345 config.tool_upstream_url = Some(format!("http://{address}/mcp"));
9346 config.mcp_request_timeout_s = Some(5);
9347 config.mcp_concurrency = 1;
9348 config.budget.max_total_tool_calls = Some(2);
9349 config.principal_budget = Some(config.budget.clone());
9350 config.allow_anonymous_principal_budget = true;
9351 let store = Arc::new(InMemoryStore::new());
9352 let sandbox = Arc::new(
9353 Sandbox::new(
9354 SandboxConfig {
9355 budget: config.budget.clone(),
9356 ..SandboxConfig::default()
9357 },
9358 Vec::new(),
9359 )
9360 .unwrap(),
9361 );
9362 let state = AppState::new(
9363 config,
9364 store.clone(),
9365 sandbox,
9366 Arc::new(NullBus),
9367 None,
9368 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
9369 )
9370 .unwrap();
9371 let app = build_router(state.clone());
9372 let body = br#"{"jsonrpc":"2.0","id":"persist","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
9373 let request = || {
9374 Request::builder()
9375 .method("POST")
9376 .uri("/v1/mcp")
9377 .header("x-pv-session", "persist-failure")
9378 .body(Body::from(Bytes::from_static(body)))
9379 .unwrap()
9380 };
9381 let first = tokio::spawn(app.clone().oneshot(request()));
9382 tokio::time::timeout(Duration::from_secs(5), async {
9383 while observed.effects.load(std::sync::atomic::Ordering::Acquire) == 0 {
9384 tokio::time::sleep(Duration::from_millis(10)).await;
9385 }
9386 })
9387 .await
9388 .unwrap();
9389 let session = state.sessions.get("persist-failure").unwrap();
9390 let mut headers = HeaderMap::new();
9391 headers.insert("x-pv-session", HeaderValue::from_static("persist-failure"));
9392 let mut execution =
9393 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
9394 .unwrap();
9395 execution.bind_principal(&session.current_identity()).unwrap();
9396 std::fs::create_dir(&execution.outcome_path).unwrap();
9399 observed.release.notify_one();
9400 let first_status = tokio::time::timeout(Duration::from_secs(5), first)
9401 .await
9402 .unwrap()
9403 .unwrap()
9404 .unwrap()
9405 .status();
9406 tokio::time::timeout(Duration::from_secs(5), state.mcp_inflight.wait_drained())
9407 .await
9408 .unwrap();
9409 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
9410 .await
9411 .unwrap();
9412 std::fs::remove_dir(&execution.outcome_path).unwrap();
9413 assert!(matches!(
9414 execution.load().await.unwrap(),
9415 ToolExecutionState::Pending
9416 ));
9417 let capture_failed = session.capture_failed();
9418 let retry_status = app.clone().oneshot(request()).await.unwrap().status();
9419 for prefix in [
9420 pv_state::ActionBudget::session_prefix("persist-failure"),
9421 pv_state::ActionBudget::principal_prefix("anonymous"),
9422 ] {
9423 assert_eq!(store.get(&format!("{prefix}total_calls")).unwrap(), 1);
9424 }
9425 let close = app
9426 .clone()
9427 .oneshot(
9428 Request::builder()
9429 .method("POST")
9430 .uri("/v1/sessions/persist-failure/close")
9431 .body(Body::empty())
9432 .unwrap(),
9433 )
9434 .await
9435 .unwrap()
9436 .status();
9437 let promote = app
9438 .oneshot(
9439 Request::builder()
9440 .method("POST")
9441 .uri("/v1/sessions/persist-failure/promote")
9442 .body(Body::empty())
9443 .unwrap(),
9444 )
9445 .await
9446 .unwrap()
9447 .status();
9448 server.abort();
9449 let _ = server.await;
9450 assert_eq!(first_status, StatusCode::INTERNAL_SERVER_ERROR);
9451 assert_eq!(retry_status, StatusCode::CONFLICT);
9452 assert_eq!(observed.effects.load(std::sync::atomic::Ordering::Acquire), 1);
9453 assert!(
9454 capture_failed,
9455 "an executed tool whose outcome cannot persist must poison complete evidence"
9456 );
9457 assert_eq!(close, StatusCode::CONFLICT);
9458 assert_eq!(promote, StatusCode::CONFLICT);
9459 assert!(session.receipt.lock().is_none());
9460 assert!(execution.intent_path.exists());
9461 assert!(!execution.outcome_path.exists());
9462 assert!(!execution.audited_path.exists());
9463 assert_eq!(state.mcp_inflight.count(), 0);
9464 assert_eq!(state.mcp_admission.available_permits(), 1);
9465 assert_eq!(session.active_streams_count(), 0);
9466 assert_eq!(session.pending_jobs_count(), 0);
9467 }
9468
9469 #[tokio::test]
9470 async fn mcp_quarantine_retains_only_unresolved_execution_evidence() {
9471 let directory = tempfile::tempdir().unwrap();
9472 let (state, provider) = test_state(directory.path()).await;
9473 let app = build_router(state.clone());
9474 let opened = app.clone().oneshot(chat_request("retention")).await.unwrap();
9475 assert_eq!(opened.status(), StatusCode::OK);
9476 axum::body::to_bytes(opened.into_body(), 64 * 1024).await.unwrap();
9477 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
9478 .await
9479 .unwrap();
9480 let session = state.sessions.get("retention").unwrap();
9481 let mut headers = HeaderMap::new();
9482 headers.insert("x-pv-session", HeaderValue::from_static("retention"));
9483 let outcome = ToolOutcome {
9484 status: 502,
9485 body_hex: hex::encode(b"partial failure evidence"),
9486 content_type: Some("text/plain".to_owned()),
9487 };
9488 let mut executions = Vec::new();
9489 for kind in [
9490 "pending",
9491 "unaudited",
9492 "corrupt-outcome",
9493 "corrupt-audit",
9494 "completed",
9495 ] {
9496 let body = serde_json::to_vec(&json!({
9497 "jsonrpc": "2.0", "id": kind, "method": "tools/call",
9498 "params": {"name": "read", "arguments": {}}
9499 }))
9500 .unwrap();
9501 let mut execution =
9502 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, &body, state.journal_key)
9503 .unwrap();
9504 execution.bind_principal(&session.current_identity()).unwrap();
9505 execution.claim().await.unwrap();
9506 if kind != "pending" {
9507 execution.persist(&outcome).await.unwrap();
9508 }
9509 if matches!(kind, "completed" | "corrupt-outcome" | "corrupt-audit") {
9510 execution.mark_audited().await.unwrap();
9511 }
9512 if kind == "corrupt-outcome" {
9513 std::fs::write(&execution.outcome_path, b"not an authenticated outcome").unwrap();
9514 }
9515 if kind == "corrupt-audit" {
9516 std::fs::write(&execution.audited_path, b"not an authenticated audit marker").unwrap();
9517 }
9518 let intent_before = std::fs::read(&execution.intent_path).unwrap();
9519 let outcome_before = std::fs::read(&execution.outcome_path).ok();
9520 executions.push((kind, execution, intent_before, outcome_before));
9521 }
9522 session.mark_capture_failed();
9523 let refused = app
9524 .clone()
9525 .oneshot(
9526 Request::builder()
9527 .method("POST")
9528 .uri("/v1/sessions/retention/close")
9529 .body(Body::empty())
9530 .unwrap(),
9531 )
9532 .await
9533 .unwrap()
9534 .status();
9535 let other = app.clone().oneshot(chat_request("other-complete")).await.unwrap();
9538 let other_status = other.status();
9539 axum::body::to_bytes(other.into_body(), 64 * 1024).await.unwrap();
9540 let other_close = app
9541 .clone()
9542 .oneshot(
9543 Request::builder()
9544 .method("POST")
9545 .uri("/v1/sessions/other-complete/close")
9546 .body(Body::empty())
9547 .unwrap(),
9548 )
9549 .await
9550 .unwrap()
9551 .status();
9552 let other_promote = app
9553 .oneshot(
9554 Request::builder()
9555 .method("POST")
9556 .uri("/v1/sessions/other-complete/promote")
9557 .body(Body::empty())
9558 .unwrap(),
9559 )
9560 .await
9561 .unwrap()
9562 .status();
9563 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
9564 .await
9565 .unwrap();
9566 provider.abort();
9567 let _ = provider.await;
9568 assert_eq!(refused, StatusCode::CONFLICT);
9569 assert!(session.receipt.lock().is_none());
9570 assert_eq!(other_status, StatusCode::OK);
9571 assert_eq!(other_close, StatusCode::OK);
9572 assert_eq!(other_promote, StatusCode::OK);
9573 for (kind, execution, intent_before, outcome_before) in executions {
9574 if kind == "completed" {
9575 assert!(!execution.intent_path.exists());
9576 assert!(!execution.outcome_path.exists());
9577 assert!(!execution.audited_path.exists());
9578 let archived_intent = std::fs::read_dir(execution.intent_path.parent().unwrap())
9579 .unwrap()
9580 .filter_map(Result::ok)
9581 .find(|entry| {
9582 entry
9583 .file_name()
9584 .to_string_lossy()
9585 .starts_with(&format!("{}.intent.json.capturefailed-", execution.key))
9586 })
9587 .expect("completed evidence must be archived, not deleted");
9588 assert_eq!(std::fs::read(archived_intent.path()).unwrap(), intent_before);
9589 } else {
9590 assert_eq!(
9591 std::fs::read(&execution.intent_path).unwrap(),
9592 intent_before,
9593 "{kind} must retain its authenticated primary intent"
9594 );
9595 assert_eq!(
9596 std::fs::read(&execution.outcome_path).ok(),
9597 outcome_before,
9598 "{kind} must preserve all partial outcome evidence"
9599 );
9600 }
9601 }
9602 }
9603
9604 async fn slow_active_chat() -> Response {
9605 let frames = futures::stream::unfold(0u8, |index| async move {
9606 if index == 7 {
9607 return None;
9608 }
9609 if index != 0 {
9610 tokio::time::sleep(Duration::from_millis(250)).await;
9611 }
9612 let frame = if index == 6 {
9613 format!(
9614 "data: {}\n\ndata: [DONE]\n\n",
9615 json!({
9616 "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
9617 "usage": {"prompt_tokens": 2, "completion_tokens": 6}
9618 })
9619 )
9620 } else {
9621 format!(
9624 "data: {}\n\n",
9625 json!({
9626 "choices": [{"index": 0, "delta": {"content": format!("chat exceeds MCP {}", "x".repeat(280))}, "finish_reason": null}]
9627 })
9628 )
9629 };
9630 Some((Ok::<_, std::convert::Infallible>(Bytes::from(frame)), index + 1))
9631 });
9632 let mut response = Response::new(Body::from_stream(frames));
9633 response.headers_mut().insert(
9634 axum::http::header::CONTENT_TYPE,
9635 HeaderValue::from_static("text/event-stream"),
9636 );
9637 response
9638 }
9639
9640 #[tokio::test]
9641 async fn mcp_total_deadline_does_not_limit_active_chat_streams() {
9642 let directory = tempfile::tempdir().unwrap();
9643 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9644 let address = listener.local_addr().unwrap();
9645 let provider = Router::new().route("/v1/chat/completions", post(slow_active_chat));
9646 let server = tokio::spawn(async move { axum::serve(listener, provider).await.unwrap() });
9647 let mut config = crate::config::HarnessConfig::for_tests(
9648 &format!("http://{address}"),
9649 &directory.path().to_string_lossy(),
9650 &directory.path().to_string_lossy(),
9651 );
9652 config.upstream_read_timeout_s = Some(1);
9653 config.mcp_request_timeout_s = Some(1);
9654 let state = AppState::new(
9655 config,
9656 Arc::new(InMemoryStore::new()),
9657 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
9658 Arc::new(NullBus),
9659 None,
9660 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
9661 )
9662 .unwrap();
9663 let started = std::time::Instant::now();
9664 let response = build_router(state.clone())
9665 .oneshot(chat_request_with_payload("chat-idle-boundary", chat_payload()))
9666 .await
9667 .unwrap();
9668 let status = response.status();
9669 let body = tokio::time::timeout(
9670 Duration::from_secs(5),
9671 axum::body::to_bytes(response.into_body(), 64 * 1024),
9672 )
9673 .await
9674 .unwrap()
9675 .unwrap();
9676 let elapsed = started.elapsed();
9677 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
9678 .await
9679 .unwrap();
9680 server.abort();
9681 let _ = server.await;
9682 let text = String::from_utf8(body.to_vec()).unwrap();
9683 assert_eq!(status, StatusCode::OK);
9684 assert!(elapsed > Duration::from_secs(1));
9685 assert_eq!(text.matches("chat exceeds MCP").count(), 6);
9686 assert!(text.contains("\"finish_reason\":\"stop\""));
9687 assert!(text.contains("data: [DONE]"));
9688 assert!(!text.contains("\"error\""));
9689 let session = state.sessions.get("chat-idle-boundary").unwrap();
9690 assert!(!session.capture_failed());
9691 assert_eq!(session.active_streams_count(), 0);
9692 assert_eq!(session.pending_jobs_count(), 0);
9693 }
9694
9695 struct SlowDripToolState {
9696 calls: std::sync::atomic::AtomicUsize,
9697 chunks: std::sync::atomic::AtomicUsize,
9698 dropped: std::sync::atomic::AtomicBool,
9699 }
9700
9701 struct SlowDripBodyGuard(Arc<SlowDripToolState>);
9702
9703 impl Drop for SlowDripBodyGuard {
9704 fn drop(&mut self) {
9705 self.0.dropped.store(true, std::sync::atomic::Ordering::Release);
9706 }
9707 }
9708
9709 async fn slow_drip_tool(State(state): State<Arc<SlowDripToolState>>) -> Response {
9710 state.calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9711 let chunks = futures::stream::unfold((0u8, SlowDripBodyGuard(state)), |(index, guard)| async move {
9715 if index == 16 {
9716 return None;
9717 }
9718 if index != 0 {
9719 tokio::time::sleep(Duration::from_millis(200)).await;
9720 }
9721 let bytes = match index {
9722 0 => Bytes::from_static(b"{\"ok\":"),
9723 15 => Bytes::from_static(b"true}"),
9724 _ => Bytes::from_static(b" "),
9725 };
9726 guard.0.chunks.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9727 Some((Ok::<_, std::convert::Infallible>(bytes), (index + 1, guard)))
9728 });
9729 let mut response = Response::new(Body::from_stream(chunks));
9730 response.headers_mut().insert(
9731 axum::http::header::CONTENT_TYPE,
9732 HeaderValue::from_static("application/json"),
9733 );
9734 response
9735 }
9736
9737 async fn assert_mcp_slow_drip_deadline(cancel_caller: bool) {
9738 let directory = tempfile::tempdir().unwrap();
9739 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9740 let address = listener.local_addr().unwrap();
9741 let observed = Arc::new(SlowDripToolState {
9742 calls: std::sync::atomic::AtomicUsize::new(0),
9743 chunks: std::sync::atomic::AtomicUsize::new(0),
9744 dropped: std::sync::atomic::AtomicBool::new(false),
9745 });
9746 let tool = Router::new()
9747 .route("/mcp", post(slow_drip_tool))
9748 .with_state(Arc::clone(&observed));
9749 let server = tokio::spawn(async move { axum::serve(listener, tool).await.unwrap() });
9750 let mut config = crate::config::HarnessConfig::for_tests(
9751 &format!("http://{address}"),
9752 &directory.path().to_string_lossy(),
9753 &directory.path().to_string_lossy(),
9754 );
9755 config.tool_upstream_url = Some(format!("http://{address}/mcp"));
9756 config.upstream_read_timeout_s = Some(1);
9757 config.mcp_request_timeout_s = Some(1);
9758 config.mcp_concurrency = 1;
9759 let state = AppState::new(
9760 config,
9761 Arc::new(InMemoryStore::new()),
9762 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
9763 Arc::new(NullBus),
9764 None,
9765 Arc::new(Ed25519Signer::from_seed(&[11; 32])),
9766 )
9767 .unwrap();
9768 let app = build_router(state.clone());
9769 let body = br#"{"jsonrpc":"2.0","id":"slow-drip","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
9770 let request = || {
9771 Request::builder()
9772 .method("POST")
9773 .uri("/v1/mcp")
9774 .header("content-type", "application/json")
9775 .header("x-pv-session", "slow-drip")
9776 .body(Body::from(Bytes::from_static(body)))
9777 .unwrap()
9778 };
9779 let first = tokio::spawn(app.clone().oneshot(request()));
9780 tokio::time::timeout(Duration::from_secs(5), async {
9781 while observed.chunks.load(std::sync::atomic::Ordering::Acquire) < 2 {
9782 tokio::time::sleep(Duration::from_millis(10)).await;
9783 }
9784 })
9785 .await
9786 .unwrap();
9787 if cancel_caller {
9788 first.abort();
9789 }
9790 tokio::time::timeout(Duration::from_secs(5), state.mcp_inflight.wait_drained())
9793 .await
9794 .unwrap();
9795 let first_status = if cancel_caller {
9796 assert!(first.await.unwrap_err().is_cancelled());
9797 None
9798 } else {
9799 Some(first.await.unwrap().unwrap().status())
9800 };
9801 let retry = app.oneshot(request()).await.unwrap();
9802 let retry_status = retry.status();
9803 let retry_body = axum::body::to_bytes(retry.into_body(), 64 * 1024).await.unwrap();
9804 tokio::time::timeout(Duration::from_secs(5), state.worker.wait_idle())
9805 .await
9806 .unwrap();
9807 let session = state.sessions.get("slow-drip").unwrap();
9808 let mut headers = HeaderMap::new();
9809 headers.insert("x-pv-session", HeaderValue::from_static("slow-drip"));
9810 let mut execution =
9811 ToolExecution::from_request(&state.config.atif_spool_dir, &headers, body, state.journal_key)
9812 .unwrap();
9813 execution.bind_principal(&session.current_identity()).unwrap();
9814 let audited = matches!(execution.load().await.unwrap(), ToolExecutionState::Completed(_));
9815 let upstream_dropped = tokio::time::timeout(Duration::from_secs(2), async {
9816 while !observed.dropped.load(std::sync::atomic::Ordering::Acquire) {
9817 tokio::time::sleep(Duration::from_millis(10)).await;
9818 }
9819 })
9820 .await
9821 .is_ok();
9822 server.abort();
9823 let _ = server.await;
9824
9825 assert_eq!(
9826 retry_status,
9827 StatusCode::BAD_GATEWAY,
9828 "the total deadline must terminate a response that never exceeds the idle timeout"
9829 );
9830 if let Some(status) = first_status {
9831 assert_eq!(status, retry_status, "the retry must replay the original failure");
9832 }
9833 assert!(String::from_utf8_lossy(&retry_body).contains("response could not be read"));
9834 assert_eq!(
9835 observed.calls.load(std::sync::atomic::Ordering::Acquire),
9836 1,
9837 "retry must not repeat tool effects"
9838 );
9839 assert!(
9840 observed.chunks.load(std::sync::atomic::Ordering::Acquire) < 16,
9841 "the upstream response must be dropped before its finite completion"
9842 );
9843 assert!(upstream_dropped, "timing out must release the upstream body");
9844 assert!(audited, "the timeout outcome must carry its durable audit marker");
9845 assert_eq!(state.mcp_inflight.count(), 0);
9846 assert_eq!(state.mcp_admission.available_permits(), 1);
9847 assert_eq!(session.active_streams_count(), 0);
9848 assert_eq!(session.pending_jobs_count(), 0);
9849 }
9850
9851 #[tokio::test]
9852 async fn mcp_slow_drip_total_deadline_records_failure_once() {
9853 assert_mcp_slow_drip_deadline(false).await;
9854 }
9855
9856 #[tokio::test]
9857 async fn mcp_slow_drip_total_deadline_survives_caller_cancellation() {
9858 assert_mcp_slow_drip_deadline(true).await;
9859 }
9860
9861 #[tokio::test]
9862 async fn tool_redirect_is_terminal_and_never_replays_post() {
9863 let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9864 let tool_address = tool_listener.local_addr().unwrap();
9865 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
9866 let tool_router = Router::new()
9867 .route("/redirect", post(redirecting_tool))
9868 .route("/effect", post(redirected_effect))
9869 .with_state(Arc::clone(&calls));
9870 let tool_server = tokio::spawn(async move {
9871 axum::serve(tool_listener, tool_router).await.unwrap();
9872 });
9873 let directory = tempfile::tempdir().unwrap();
9874 let (mut state, provider) = test_state(directory.path()).await;
9875 Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
9876 Some(format!("http://{tool_address}/redirect"));
9877 let response = build_router(state)
9878 .oneshot(
9879 Request::builder()
9880 .method("POST")
9881 .uri("/v1/mcp")
9882 .header("x-pv-session", "tool-redirect")
9883 .body(Body::from(
9884 r#"{"jsonrpc":"2.0","id":"redirect","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
9885 ))
9886 .unwrap(),
9887 )
9888 .await
9889 .unwrap();
9890 assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT);
9891 assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
9892 tool_server.abort();
9893 provider.abort();
9894 }
9895
9896 #[tokio::test]
9897 async fn concurrent_duplicate_tool_claims_conflict_without_leaking_io_detail() {
9898 let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9899 let tool_address = tool_listener.local_addr().unwrap();
9900 let held = Arc::new(HeldToolState {
9901 arrived: std::sync::atomic::AtomicBool::new(false),
9902 release: tokio::sync::Notify::new(),
9903 });
9904 let tool_router = Router::new()
9905 .route("/mcp", post(held_tool))
9906 .with_state(Arc::clone(&held));
9907 let tool_server = tokio::spawn(async move {
9908 axum::serve(tool_listener, tool_router).await.unwrap();
9909 });
9910 let directory = tempfile::tempdir().unwrap();
9911 let (mut state, provider) = test_state(directory.path()).await;
9912 Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
9913 Some(format!("http://{tool_address}/mcp"));
9914 let app = build_router(state);
9915 let request = || {
9916 Request::builder()
9917 .method("POST")
9918 .uri("/v1/mcp")
9919 .header("x-pv-session", "claim-race")
9920 .body(Body::from(
9921 r#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"read","arguments":{}}}"#,
9922 ))
9923 .unwrap()
9924 };
9925 let winner_app = app.clone();
9926 let winner = tokio::spawn(async move { winner_app.oneshot(request()).await.unwrap() });
9927 tokio::time::timeout(Duration::from_secs(1), async {
9928 while !held.arrived.load(std::sync::atomic::Ordering::Acquire) {
9929 tokio::task::yield_now().await;
9930 }
9931 })
9932 .await
9933 .unwrap();
9934 let loser = app.oneshot(request()).await.unwrap();
9938 assert_eq!(loser.status(), StatusCode::CONFLICT);
9939 let body = axum::body::to_bytes(loser.into_body(), 64 * 1024).await.unwrap();
9940 let text = String::from_utf8(body.to_vec()).unwrap();
9941 assert!(text.contains(TOOL_OUTCOME_UNCERTAIN), "{text}");
9942 assert!(!text.contains("os error"), "leaked io detail: {text}");
9943 held.release.notify_one();
9944 assert_eq!(winner.await.unwrap().status(), StatusCode::OK);
9945 tool_server.abort();
9946 provider.abort();
9947 }
9948
9949 #[tokio::test]
9950 async fn close_waits_for_tool_execution_and_completion_capture() {
9951 let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
9952 let tool_address = tool_listener.local_addr().unwrap();
9953 let held = Arc::new(HeldToolState {
9954 arrived: std::sync::atomic::AtomicBool::new(false),
9955 release: tokio::sync::Notify::new(),
9956 });
9957 let tool_router = Router::new()
9958 .route("/mcp", post(held_tool))
9959 .with_state(Arc::clone(&held));
9960 let tool_server = tokio::spawn(async move {
9961 axum::serve(tool_listener, tool_router).await.unwrap();
9962 });
9963 let directory = tempfile::tempdir().unwrap();
9964 let (mut state, provider) = test_state(directory.path()).await;
9965 Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
9966 Some(format!("http://{tool_address}/mcp"));
9967 let app = build_router(state);
9968 let tool_app = app.clone();
9969 let tool_task = tokio::spawn(async move {
9970 tool_app
9971 .oneshot(
9972 Request::builder()
9973 .method("POST")
9974 .uri("/v1/mcp")
9975 .header("x-pv-session", "held-tool")
9976 .body(Body::from(
9977 r#"{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"read","arguments":{}}}"#,
9978 ))
9979 .unwrap(),
9980 )
9981 .await
9982 .unwrap()
9983 });
9984 tokio::time::timeout(Duration::from_secs(1), async {
9985 while !held.arrived.load(std::sync::atomic::Ordering::Acquire) {
9986 tokio::task::yield_now().await;
9987 }
9988 })
9989 .await
9990 .unwrap();
9991 let close_task = tokio::spawn(async move {
9992 app.oneshot(
9993 Request::builder()
9994 .method("POST")
9995 .uri("/v1/sessions/held-tool/close")
9996 .body(Body::empty())
9997 .unwrap(),
9998 )
9999 .await
10000 .unwrap()
10001 });
10002 tokio::task::yield_now().await;
10003 assert!(!close_task.is_finished(), "close overtook the executing tool");
10004 held.release.notify_one();
10005 assert_eq!(tool_task.await.unwrap().status(), StatusCode::OK);
10006 let close_response = close_task.await.unwrap();
10007 let close_status = close_response.status();
10008 let close_body = axum::body::to_bytes(close_response.into_body(), 64 * 1024)
10009 .await
10010 .unwrap();
10011 assert_eq!(
10012 close_status,
10013 StatusCode::OK,
10014 "close failed: {}",
10015 String::from_utf8_lossy(&close_body)
10016 );
10017 tool_server.abort();
10018 provider.abort();
10019 }
10020
10021 #[test]
10022 #[allow(clippy::cast_possible_truncation)] fn fuzz_parse_provider_chunk_is_total() {
10024 let mut state = 0x9E37_79B9_7F4A_7C15u64;
10028 let mut next = move || {
10029 state ^= state << 13;
10030 state ^= state >> 7;
10031 state ^= state << 17;
10032 state
10033 };
10034 let seeds: Vec<&str> = vec![
10035 "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}],\"usage\":null}\n\n",
10036 "data: {\"model\":\"m\",\"choices\":[{\"delta\":{\"content\":\"a\",\"reasoning_content\":\"r\"},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":3,\"cost_usd\":0.001,\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\n",
10037 "data: {\"choices\":[{\"index\":1,\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"c\",\"function\":{\"name\":\"f\",\"arguments\":\"{}\"}}]}}]}\n\n",
10038 "data: [DONE]\n\n",
10039 "event: message\ndata: {\"choices\":[]}\n\n",
10040 "{\"choices\":[{\"message\":{\"content\":\"x\",\"tool_calls\":[{\"id\":\"a\",\"function\":{\"name\":\"n\",\"arguments\":\"{}\"}}]}}],\"usage\":{\"completion_tokens\":1}}",
10041 ];
10042 for _ in 0..200_000 {
10043 let seed = seeds[(next() as usize) % seeds.len()];
10044 let mut bytes = seed.as_bytes().to_vec();
10045 for _ in 0..=(next() % 8) {
10046 match next() % 3 {
10047 0 if !bytes.is_empty() => {
10048 let idx = (next() as usize) % bytes.len();
10049 bytes[idx] = (next() & 0xFF) as u8;
10050 }
10051 1 if !bytes.is_empty() => {
10052 let idx = (next() as usize) % bytes.len();
10053 bytes.truncate(idx);
10054 }
10055 _ => {
10056 let idx = if bytes.is_empty() {
10057 0
10058 } else {
10059 (next() as usize) % bytes.len()
10060 };
10061 bytes.insert(idx, (next() & 0xFF) as u8);
10062 }
10063 }
10064 }
10065 let raw = String::from_utf8_lossy(&bytes).into_owned();
10066 let _ = parse_provider_chunk(&raw);
10067 let _ = sse_frame_end(&bytes);
10068 }
10069 }
10070
10071 #[test]
10072 fn provider_usage_and_reasoning_are_preserved() {
10073 let raw = concat!(
10074 "data: {\"model\":\"gpt-test\",\"choices\":[{\"delta\":{",
10075 "\"content\":\"answer\",\"reasoning_content\":\"thought\"},\"finish_reason\":\"stop\"}],",
10076 "\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":3,\"cost_usd\":0.00125,",
10077 "\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\n"
10078 );
10079 let parsed = parse_provider_chunk(raw).unwrap().unwrap();
10080 assert_eq!(parsed.choices[0].content, "answer");
10081 assert_eq!(parsed.choices[0].reasoning, "thought");
10082 assert_eq!(parsed.model_name.as_deref(), Some("gpt-test"));
10083 assert_eq!(parsed.metrics.prompt_tokens, Some(10));
10084 assert_eq!(parsed.metrics.completion_tokens, Some(3));
10085 assert_eq!(parsed.metrics.cached_tokens, Some(4));
10086 assert!(parsed.usage_reported);
10087 assert_eq!(parsed.choices[0].finish_reason.as_deref(), Some("stop"));
10088 assert_eq!(parsed.cost_usd_micros, 1250);
10089 assert!(parsed.tool_call_deltas.is_empty());
10090 }
10091
10092 #[test]
10093 fn null_usage_chunks_are_not_reported_usage() {
10094 let raw = r#"data: {"choices":[{"delta":{"content":"Hello"}}],"usage":null}"#;
10102 let parsed = parse_provider_chunk(raw).unwrap().unwrap();
10103 assert!(!parsed.usage_reported, "null usage must not count as reported");
10104 assert!(!parsed.completion_reported);
10105 assert_eq!(parsed.choices[0].content, "Hello");
10106 assert_eq!(
10108 parsed.metrics.completion_tokens,
10109 Some(pv_core::tokens::approx_tokens("Hello"))
10110 );
10111 let raw = r#"data: {"choices":[{"delta":{"content":"x"}}],"usage":{"prompt_tokens":7,"completion_tokens":null}}"#;
10114 let parsed = parse_provider_chunk(raw).unwrap().unwrap();
10115 assert!(parsed.usage_reported);
10116 assert!(
10117 !parsed.completion_reported,
10118 "a null completion count must not drive cumulative-usage math"
10119 );
10120 assert_eq!(parsed.metrics.prompt_tokens, Some(7));
10121 assert_eq!(
10122 parsed.metrics.completion_tokens,
10123 Some(pv_core::tokens::approx_tokens("x"))
10124 );
10125 }
10126
10127 #[test]
10128 fn hostile_provider_metrics_are_rejected() {
10129 for raw in [
10130 r#"{"usage":{"completion_tokens":9007199254740993}}"#,
10131 r#"{"usage":{"prompt_tokens":-1}}"#,
10132 r#"{"usage":{"prompt_tokens_details":{"cached_tokens":"many"}}}"#,
10133 r#"{"usage":{"cost_usd":1e30}}"#,
10134 r#"{"usage":{"cost_usd":-1}}"#,
10135 ] {
10136 assert!(
10137 parse_provider_chunk(raw).is_err(),
10138 "accepted hostile metrics: {raw}"
10139 );
10140 }
10141 }
10142
10143 #[test]
10144 fn streaming_tool_call_arguments_reassemble() {
10145 let first = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-1","function":{"name":"db_write","arguments":"{\"table\":"}}]}}]}"#;
10146 let second = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"items\",\"row\":{}}"}}]},"finish_reason":"tool_calls"}]}"#;
10147 let first_delta = parse_provider_chunk(first).unwrap().unwrap().tool_call_deltas;
10148 let second_delta = parse_provider_chunk(second).unwrap().unwrap().tool_call_deltas;
10149 assert_eq!(first_delta.len(), 1);
10150 assert_eq!(second_delta.len(), 1);
10151 let arguments = format!("{}{}", first_delta[0].arguments, second_delta[0].arguments);
10152 assert_eq!(
10153 serde_json::from_str::<Value>(&arguments).unwrap(),
10154 json!({"table": "items", "row": {}})
10155 );
10156 assert_eq!(map_finish_reason("tool_calls"), StopReason::ToolUse);
10157 }
10158
10159 #[test]
10160 fn hostile_provider_tool_call_index_is_rejected() {
10161 let raw = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":18446744073709551615,"function":{"name":"n","arguments":"{}"}}]}}]}"#;
10162 let parsed = parse_provider_chunk(raw).unwrap().unwrap();
10163 assert_eq!(parsed.tool_call_deltas.len(), 1);
10164 assert_eq!(parsed.tool_call_deltas[0].index, u64::MAX);
10165 }
10166
10167 #[test]
10168 fn non_stream_tool_calls_receive_distinct_positional_indices() {
10169 let parsed = parse_provider_chunk(
10170 r#"{"choices":[{"message":{"tool_calls":[{"id":"a","function":{"name":"read","arguments":"{}"}},{"id":"b","function":{"name":"write","arguments":"{}"}}]}}]}"#,
10171 )
10172 .unwrap()
10173 .unwrap();
10174 assert_eq!(parsed.tool_call_deltas.len(), 2);
10175 assert_eq!(parsed.tool_call_deltas[0].index, 0);
10176 assert_eq!(parsed.tool_call_deltas[1].index, 1);
10177 }
10178
10179 #[test]
10180 fn sse_frame_boundary_accepts_crlf_and_cr() {
10181 assert_eq!(sse_frame_end(b"data: {}\r\n\r"), None);
10182 assert_eq!(sse_frame_end(b"data: {}\r\n\r\nrest"), Some(12));
10183 assert_eq!(sse_frame_end(b"data: {}\r\rrest"), Some(10));
10184 }
10185
10186 #[test]
10193 fn parse_provider_chunk_strips_leading_bom() {
10194 let raw = "\u{feff}data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n";
10195 let parsed = parse_provider_chunk(raw)
10196 .expect("BOM-prefixed SSE frame must parse")
10197 .expect("must yield a chunk");
10198 assert_eq!(parsed.choices[0].content, "hi");
10199 }
10200
10201 #[test]
10212 fn parse_provider_chunk_refuses_non_message_sse_event_types() {
10213 for event in ["error", "ping", "custom_signal"] {
10214 let raw = format!(
10215 "event: {event}\ndata: {{\"choices\":[{{\"delta\":{{\"content\":\"forged\"}}}}]}}\n\n"
10216 );
10217 let err = match parse_provider_chunk(&raw) {
10218 Ok(_) => panic!("event type {event:?} must be refused"),
10219 Err(error) => error,
10220 };
10221 assert!(
10222 err.contains("unsupported event type") && err.contains(event),
10223 "expected diagnostic naming event type {event:?}, got: {err}"
10224 );
10225 }
10226 let ok_raw = "event: message\ndata: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n";
10229 let parsed = match parse_provider_chunk(ok_raw) {
10230 Ok(Some(chunk)) => chunk,
10231 Ok(None) => panic!("explicit event: message must yield a chunk"),
10232 Err(error) => panic!("explicit event: message must be accepted, got: {error}"),
10233 };
10234 assert_eq!(parsed.choices[0].content, "hi");
10235 }
10236
10237 #[test]
10246 fn parse_provider_chunk_treats_dataless_named_events_as_keepalives() {
10247 for raw in [
10248 "event: ping\n\n",
10249 "event: heartbeat\n\n",
10250 "event: keep-alive\ndata:\n\n",
10251 "event: status\ndata: \n\n",
10252 ] {
10253 match parse_provider_chunk(raw) {
10254 Ok(None) => {}
10255 Ok(Some(_)) => panic!("dataless named event must not yield a chunk: {raw:?}"),
10256 Err(error) => panic!("dataless named event {raw:?} must not fail the stream (D7): {error}"),
10257 }
10258 }
10259 let poison = "event: error\ndata: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n";
10263 assert!(
10264 parse_provider_chunk(poison).is_err(),
10265 "named event with non-empty data must still be refused"
10266 );
10267 }
10268
10269 #[test]
10270 fn parse_provider_chunk_treats_done_variants_as_stream_end() {
10271 for raw in [
10272 "data: [DONE]\ndata: [DONE]\n\n",
10273 "data: [DONE] \n\n",
10274 "data:\ndata: [DONE]\n\n",
10275 ] {
10276 let result = parse_provider_chunk(raw);
10277 match result {
10278 Ok(None) => {}
10279 Ok(Some(_)) => panic!("[DONE] variant must terminate stream (None) for {raw:?}"),
10280 Err(error) => panic!("[DONE] variant must not fail parse: {raw:?}: {error}"),
10281 }
10282 }
10283 }
10284
10285 #[test]
10294 fn parse_provider_chunk_done_under_named_event_is_keepalive() {
10295 for raw in [
10296 "event: summary\ndata: [DONE]\n\n",
10297 "event: error\ndata: [DONE]\n\n",
10298 "event: custom\ndata:\ndata: [DONE]\n\n",
10299 "event: batched\ndata: [DONE]\ndata: [DONE]\n\n",
10300 ] {
10301 match parse_provider_chunk(raw) {
10302 Ok(None) => {}
10303 Ok(Some(_)) => panic!(
10304 "[DONE] under named event {raw:?} must not yield a chunk (no attributable content)"
10305 ),
10306 Err(error) => {
10307 panic!("[DONE] under named event {raw:?} must not fail-close the stream: {error}")
10308 }
10309 }
10310 }
10311 let mixed = "event: error\ndata: [DONE]\ndata: {\"choices\":[{\"delta\":{\"content\":\"x\"}}]}\n\n";
10315 assert!(
10316 parse_provider_chunk(mixed).is_err(),
10317 "named event with mixed [DONE] + real content must still refuse: {mixed:?}"
10318 );
10319 }
10320
10321 #[test]
10352 fn parse_provider_chunk_strips_exactly_one_leading_space() {
10353 let raw = "data: \thello\n\n";
10362 let result = parse_provider_chunk(raw);
10363 let error = match result {
10364 Ok(_) => panic!("bare hello must not parse as valid provider frame"),
10365 Err(error) => error,
10366 };
10367 assert!(
10368 error.contains("column 2") || error.contains("column 3"),
10369 "expected error to name a column past the retained leading space+tab; \
10370 got: {error}"
10371 );
10372 }
10373
10374 #[tokio::test]
10375 async fn split_non_sse_utf8_is_captured_as_one_json_document() {
10376 let directory = tempfile::tempdir().unwrap();
10377 let (state, provider) = test_state(directory.path()).await;
10378 let app = build_router(state.clone());
10379 let response = app
10380 .oneshot(chat_request_with_payload(
10381 "split-json",
10382 json!({
10383 "model": "split-json",
10384 "stream": false,
10385 "messages": [{"role": "user", "content": "hello"}]
10386 }),
10387 ))
10388 .await
10389 .unwrap();
10390 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
10391 .await
10392 .unwrap();
10393 assert_eq!(
10394 serde_json::from_slice::<Value>(&body).unwrap()["choices"][0]["message"]["content"],
10395 "héllo"
10396 );
10397 let session = state.sessions.get("split-json").unwrap();
10398 session.wait_for_worker_jobs().await;
10399 let crate::reconciler::FinalizeOutcome::Atif { path } = state
10400 .finalizer
10401 .close_session(session, StopReason::SessionClosed)
10402 .await
10403 .unwrap()
10404 else {
10405 panic!("expected ATIF")
10406 };
10407 let trajectory: pv_atif::Trajectory = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
10408 assert_eq!(trajectory.steps[1].message, Value::String("héllo".to_owned()));
10409 provider.abort();
10410 }
10411
10412 #[tokio::test]
10413 async fn malformed_non_sse_json_fails_capture_before_delivery() {
10414 let directory = tempfile::tempdir().unwrap();
10415 let (state, provider) = test_state(directory.path()).await;
10416 let app = build_router(state.clone());
10417 let response = app
10418 .oneshot(chat_request_with_payload(
10419 "malformed-json",
10420 json!({
10421 "model": "malformed-json",
10422 "stream": false,
10423 "messages": [{"role": "user", "content": "hello"}]
10424 }),
10425 ))
10426 .await
10427 .unwrap();
10428 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
10432 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
10433 .await
10434 .unwrap();
10435 {
10436 let parsed = serde_json::from_slice::<Value>(&body).unwrap();
10439 assert!(parsed["error"]["message"].is_string(), "{parsed}");
10440 assert!(parsed["error"]["type"].is_string(), "{parsed}");
10441 assert!(parsed["error"]["code"].is_number(), "{parsed}");
10442 }
10443 let session = state.sessions.get("malformed-json").unwrap();
10444 assert!(
10449 !session.capture_failed(),
10450 "a single malformed provider frame must not brick the session"
10451 );
10452 provider.abort();
10453 }
10454
10455 #[tokio::test]
10461 async fn unsupported_content_encoding_refusal_retires_marker_and_attempt() {
10462 let directory = tempfile::tempdir().unwrap();
10463 let (state, provider) = test_state(directory.path()).await;
10464 let response = build_router(state.clone())
10465 .oneshot(chat_request_with_payload(
10466 "gzip-encoded",
10467 json!({
10468 "model": "gzip-encoded",
10469 "stream": false,
10470 "messages": [{"role": "user", "content": "hello"}]
10471 }),
10472 ))
10473 .await
10474 .unwrap();
10475 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
10476 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
10477 .await
10478 .unwrap();
10479 let error = serde_json::from_slice::<Value>(&body).unwrap()["error"]["message"]
10480 .as_str()
10481 .unwrap()
10482 .to_owned();
10483 assert!(error.contains("Content-Encoding"), "unexpected error: {error}");
10484 let session = state.sessions.get("gzip-encoded").unwrap();
10485 session.wait_for_worker_jobs().await;
10486 crate::worker::ensure_no_inflight_responses(directory.path(), &state.journal_key)
10487 .await
10488 .expect("refusal must retire the in-flight response marker");
10489 let records = active_records(directory.path(), &state, "gzip-encoded");
10490 assert!(
10491 records
10492 .iter()
10493 .filter_map(|record| record.response_attempt.as_ref())
10494 .any(|attempt| attempt.terminal),
10495 "refusal must journal a terminal response attempt"
10496 );
10497 let event: pv_events::OcsfEvent =
10498 serde_json::from_value(records.last().unwrap().event.clone()).unwrap();
10499 assert_eq!(event.class_name, pv_events::EventClass::StopReason);
10500 assert_eq!(event.status_id, pv_events::StatusId::Failure.id());
10501 assert_eq!(event.payload["reason"], "upstream_unsupported_content_encoding");
10502 provider.abort();
10503 }
10504
10505 #[tokio::test]
10506 async fn repeated_provider_responses_open_the_loop_breaker() {
10507 let directory = tempfile::tempdir().unwrap();
10508 let (mut state, provider) = test_state(directory.path()).await;
10509 let config = Arc::get_mut(&mut state.config).unwrap();
10510 config.breaker.min_tokens = 0;
10511 config.breaker.window = 3;
10512 let app = build_router(state.clone());
10513 for _ in 0..4 {
10514 let response = app.clone().oneshot(chat_request("response-loop")).await.unwrap();
10515 assert_eq!(response.status(), StatusCode::OK);
10516 axum::body::to_bytes(response.into_body(), 64 * 1024)
10517 .await
10518 .unwrap();
10519 state
10520 .sessions
10521 .get("response-loop")
10522 .unwrap()
10523 .wait_for_worker_jobs()
10524 .await;
10525 }
10526 let blocked = app.oneshot(chat_request("response-loop")).await.unwrap();
10527 assert_eq!(blocked.status(), StatusCode::FORBIDDEN);
10532 provider.abort();
10533 }
10534
10535 #[tokio::test]
10536 async fn empty_and_failed_upstream_responses_are_audited_as_failures() {
10537 let directory = tempfile::tempdir().unwrap();
10538 let (state, provider) = test_state(directory.path()).await;
10539 let app = build_router(state.clone());
10540 for (session_id, model) in [("empty-error", "empty-error"), ("json-error", "json-error")] {
10541 let response = app
10542 .clone()
10543 .oneshot(chat_request_with_payload(
10544 session_id,
10545 json!({
10546 "model": model,
10547 "stream": false,
10548 "messages": [{"role": "user", "content": "hello"}]
10549 }),
10550 ))
10551 .await
10552 .unwrap();
10553 assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
10554 axum::body::to_bytes(response.into_body(), 64 * 1024)
10555 .await
10556 .unwrap();
10557 let session = state.sessions.get(session_id).unwrap();
10558 session.wait_for_worker_jobs().await;
10559 assert!(!session.capture_failed());
10560 let records = active_records(directory.path(), &state, session_id);
10561 assert_eq!(records.len(), 2);
10562 let event: pv_events::OcsfEvent =
10563 serde_json::from_value(records.last().unwrap().event.clone()).unwrap();
10564 assert_eq!(event.class_name, pv_events::EventClass::StopReason);
10565 assert_eq!(event.status_id, pv_events::StatusId::Failure.id());
10566 assert_eq!(event.payload["http_status"], 500);
10567 }
10568 provider.abort();
10569 }
10570
10571 #[tokio::test]
10572 async fn regressing_provider_usage_fails_capture_before_later_bytes() {
10573 let directory = tempfile::tempdir().unwrap();
10574 let (state, provider) = test_state(directory.path()).await;
10575 let response = build_router(state.clone())
10576 .oneshot(chat_request_with_payload(
10577 "regressive-usage",
10578 json!({
10579 "model": "regressive-usage",
10580 "stream": true,
10581 "messages": [{"role": "user", "content": "hello"}]
10582 }),
10583 ))
10584 .await
10585 .unwrap();
10586 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10587 .await
10588 .is_err());
10589 let session = state.sessions.get("regressive-usage").unwrap();
10590 assert!(
10594 !session.capture_failed(),
10595 "a regressive usage frame must be scoped to the response, not the session"
10596 );
10597 session.wait_for_worker_jobs().await;
10598 let records = active_records(directory.path(), &state, "regressive-usage");
10599 assert_eq!(records.len(), 2);
10603 assert!(records.iter().any(|record| record.response_attempt.is_some()));
10604 provider.abort();
10605 }
10606
10607 #[tokio::test]
10608 async fn successful_provider_response_requires_choices() {
10609 let directory = tempfile::tempdir().unwrap();
10610 let (state, provider) = test_state(directory.path()).await;
10611 let response = build_router(state.clone())
10612 .oneshot(chat_request_with_payload(
10613 "empty-success",
10614 json!({
10615 "model": "empty-success",
10616 "stream": false,
10617 "messages": [{"role": "user", "content": "hello"}]
10618 }),
10619 ))
10620 .await
10621 .unwrap();
10622 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
10625 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
10626 .await
10627 .unwrap();
10628 {
10629 let parsed = serde_json::from_slice::<Value>(&body).unwrap();
10632 assert!(parsed["error"]["message"].is_string(), "{parsed}");
10633 assert!(parsed["error"]["type"].is_string(), "{parsed}");
10634 assert!(parsed["error"]["code"].is_number(), "{parsed}");
10635 }
10636 assert!(
10639 !state.sessions.get("empty-success").unwrap().capture_failed(),
10640 "a choices-less response must not brick the session"
10641 );
10642 provider.abort();
10643 }
10644
10645 #[tokio::test]
10646 async fn cumulative_usage_settles_against_provisional_charge() {
10647 let directory = tempfile::tempdir().unwrap();
10648 let payload = json!({
10649 "model": "provisional-usage",
10650 "stream": true,
10651 "messages": [{"role": "user", "content": "hello"}]
10652 });
10653 let prompt_tokens = pv_core::tokens::approx_tokens(&payload.to_string());
10654 let (state, provider) = test_state_with_token_cap(directory.path(), Some(prompt_tokens + 2)).await;
10655 let response = build_router(state.clone())
10656 .oneshot(chat_request_with_payload("provisional-usage", payload))
10657 .await
10658 .unwrap();
10659 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10660 .await
10661 .is_ok());
10662 let session = state.sessions.get("provisional-usage").unwrap();
10663 session.wait_for_worker_jobs().await;
10664 assert_eq!(
10665 session
10666 .totals
10667 .completion_tokens
10668 .load(std::sync::atomic::Ordering::Acquire),
10669 2
10670 );
10671 provider.abort();
10672 }
10673
10674 #[tokio::test]
10675 async fn provider_tool_call_index_out_of_range_fails_capture() {
10676 let directory = tempfile::tempdir().unwrap();
10677 let (state, provider) = test_state(directory.path()).await;
10678 let response = build_router(state.clone())
10679 .oneshot(chat_request_with_payload(
10680 "tool-oob-index",
10681 json!({
10682 "model": "tool-oob-index",
10683 "stream": true,
10684 "messages": [{"role": "user", "content": "hello"}]
10685 }),
10686 ))
10687 .await
10688 .unwrap();
10689 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
10694 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10695 .await
10696 .is_ok());
10697 assert!(
10701 !state.sessions.get("tool-oob-index").unwrap().capture_failed(),
10702 "an out-of-range tool-call index must be scoped to the response"
10703 );
10704 provider.abort();
10705 }
10706
10707 #[tokio::test]
10714 async fn provider_tool_call_id_conflict_fails_capture() {
10715 let directory = tempfile::tempdir().unwrap();
10716 let (state, provider) = test_state(directory.path()).await;
10717 let response = build_router(state.clone())
10718 .oneshot(chat_request_with_payload(
10719 "tool-conflicting-id",
10720 json!({
10721 "model": "tool-conflicting-id",
10722 "stream": true,
10723 "messages": [{"role": "user", "content": "hello"}]
10724 }),
10725 ))
10726 .await
10727 .unwrap();
10728 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
10731 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10732 .await
10733 .is_ok());
10734 assert!(
10735 !state
10736 .sessions
10737 .get("tool-conflicting-id")
10738 .unwrap()
10739 .capture_failed(),
10740 "a conflicting tool-call id must be scoped to the response"
10741 );
10742 let response = build_router(state.clone())
10746 .oneshot(chat_request_with_payload(
10747 "tool-conflicting-name",
10748 json!({
10749 "model": "tool-conflicting-name",
10750 "stream": true,
10751 "messages": [{"role": "user", "content": "hello"}]
10752 }),
10753 ))
10754 .await
10755 .unwrap();
10756 assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
10757 provider.abort();
10758 }
10759
10760 #[tokio::test]
10768 async fn completion_estimate_over_charge_is_refunded_after_authoritative_usage() {
10769 let directory = tempfile::tempdir().unwrap();
10770 let payload = json!({
10771 "model": "overcharged-completion",
10772 "stream": true,
10773 "messages": [{"role": "user", "content": "hello"}]
10774 });
10775 let estimate = pv_core::tokens::approx_tokens(&"a".repeat(32));
10778 assert!(estimate > 1, "test premise: estimate must exceed usage");
10779 let prompt_tokens = pv_core::tokens::approx_tokens(&payload.to_string());
10780 let cap = prompt_tokens + estimate + 8;
10781 let (state, provider) = test_state_with_token_cap(directory.path(), Some(cap)).await;
10782 let response = build_router(state.clone())
10783 .oneshot(chat_request_with_payload("overcharged-completion", payload))
10784 .await
10785 .unwrap();
10786 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10787 .await
10788 .is_ok());
10789 let session = state.sessions.get("overcharged-completion").unwrap();
10790 session.wait_for_worker_jobs().await;
10791 let ledger = pv_state::ActionBudget::new(
10795 state.store.as_ref(),
10796 "overcharged-completion",
10797 &state.config.budget,
10798 );
10799 assert!(
10800 ledger.try_tokens(cap - 2).unwrap().is_allowed(),
10801 "the completion-estimate over-charge must be refunded once authoritative usage arrives"
10802 );
10803 assert!(
10804 !ledger.try_tokens(1).unwrap().is_allowed(),
10805 "refunds must settle to the provider-billed spend, not more"
10806 );
10807 provider.abort();
10808 }
10809
10810 #[tokio::test]
10816 async fn sse_frame_split_across_chunks_reassembles_before_capture() {
10817 let directory = tempfile::tempdir().unwrap();
10818 let (state, provider) = test_state(directory.path()).await;
10819 let response = build_router(state.clone())
10820 .oneshot(chat_request_with_payload(
10821 "split-sse",
10822 json!({
10823 "model": "split-sse",
10824 "stream": true,
10825 "messages": [{"role": "user", "content": "hello"}]
10826 }),
10827 ))
10828 .await
10829 .unwrap();
10830 let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
10831 .await
10832 .unwrap();
10833 assert!(
10835 String::from_utf8_lossy(&body).contains("héllo"),
10836 "client bytes truncated"
10837 );
10838 let session = state.sessions.get("split-sse").unwrap();
10839 session.wait_for_worker_jobs().await;
10840 assert!(
10841 !session.capture_failed(),
10842 "a chunk-split frame must not fail capture"
10843 );
10844 let crate::reconciler::FinalizeOutcome::Atif { path } = state
10845 .finalizer
10846 .close_session(session, StopReason::SessionClosed)
10847 .await
10848 .unwrap()
10849 else {
10850 panic!("expected ATIF")
10851 };
10852 let trajectory: pv_atif::Trajectory = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
10853 assert_eq!(
10854 trajectory.steps[1].message,
10855 Value::String("héllo".to_owned()),
10856 "the audit trail must carry the reassembled multi-byte content"
10857 );
10858 provider.abort();
10859 }
10860
10861 #[tokio::test]
10869 async fn truncated_stream_is_not_attested_as_a_complete_success() {
10870 let directory = tempfile::tempdir().unwrap();
10871 let (state, provider) = test_state(directory.path()).await;
10872 let response = build_router(state.clone())
10873 .oneshot(chat_request_with_payload(
10874 "truncated-stream",
10875 json!({
10876 "model": "truncated-stream",
10877 "stream": true,
10878 "messages": [{"role": "user", "content": "hello"}]
10879 }),
10880 ))
10881 .await
10882 .unwrap();
10883 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10886 .await
10887 .is_ok());
10888 let session = state.sessions.get("truncated-stream").unwrap();
10889 session.wait_for_worker_jobs().await;
10890 let records = active_records(directory.path(), &state, "truncated-stream");
10891 let truncated = records.iter().any(|record| {
10892 record
10893 .event
10894 .get("payload")
10895 .is_some_and(|payload| payload.get("truncated").and_then(Value::as_bool) == Some(true))
10896 });
10897 assert!(
10898 truncated,
10899 "a finish_reason-less SSE end must be journalled with a truncated marker; records: {:?}",
10900 records
10901 .iter()
10902 .map(|r| r.event.get("payload").cloned())
10903 .collect::<Vec<_>>()
10904 );
10905 provider.abort();
10906 }
10907
10908 #[tokio::test]
10918 async fn severed_stream_is_attested_other_not_budget_exceeded() {
10919 let directory = tempfile::tempdir().unwrap();
10920 let (state, provider) = test_state(directory.path()).await;
10921 let response = build_router(state.clone())
10922 .oneshot(chat_request_with_payload(
10923 "severed-stream",
10924 json!({
10925 "model": "severed-stream",
10926 "stream": true,
10927 "messages": [{"role": "user", "content": "hello"}]
10928 }),
10929 ))
10930 .await
10931 .unwrap();
10932 let _ = axum::body::to_bytes(response.into_body(), 64 * 1024).await;
10935 let session = state.sessions.get("severed-stream").unwrap();
10936 session.wait_for_worker_jobs().await;
10937 let records = active_records(directory.path(), &state, "severed-stream");
10938 let stop_reasons: Vec<(u64, Option<String>)> = records
10939 .iter()
10940 .filter_map(|record| {
10941 let id = record.event.get("stop_reason_id").and_then(Value::as_u64)?;
10942 let reason = record
10943 .event
10944 .pointer("/payload/reason")
10945 .and_then(Value::as_str)
10946 .map(str::to_owned);
10947 Some((id, reason))
10948 })
10949 .collect();
10950 assert!(
10951 stop_reasons.iter().all(|(id, _)| *id != 92),
10952 "a provider cut must never be attested as Budget Exceeded (92); got {stop_reasons:?}"
10953 );
10954 assert!(
10955 stop_reasons
10956 .iter()
10957 .any(|(id, reason)| *id == 99 && reason.as_deref().is_some_and(|r| !r.is_empty())),
10958 "the cut turn must be attested Other (99) with a free-text failure reason; got {stop_reasons:?}"
10959 );
10960 provider.abort();
10961 }
10962
10963 #[tokio::test]
10964 async fn tool_arguments_without_usage_are_budgeted_before_delivery() {
10965 let directory = tempfile::tempdir().unwrap();
10966 let payload = json!({
10967 "model": "tool-no-usage",
10968 "stream": true,
10969 "messages": [{"role": "user", "content": "hello"}]
10970 });
10971 let prompt_tokens = pv_core::tokens::approx_tokens(&payload.to_string());
10972 let (state, provider) = test_state_with_token_cap(directory.path(), Some(prompt_tokens)).await;
10973 let response = build_router(state)
10974 .oneshot(chat_request_with_payload("tool-no-usage", payload))
10975 .await
10976 .unwrap();
10977 assert_eq!(response.status(), StatusCode::FORBIDDEN);
10981 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
10982 .await
10983 .is_ok());
10984 provider.abort();
10985 }
10986
10987 #[tokio::test(flavor = "current_thread")]
10988 async fn completion_budget_store_does_not_block_stream_runtime() {
10989 let directory = tempfile::tempdir().unwrap();
10990 let (mut state, provider) = test_state_with_token_cap(directory.path(), Some(10_000)).await;
10991 let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
10992 state.store = Arc::new(SlowStore {
10993 inner: InMemoryStore::new(),
10994 calls: Arc::clone(&calls),
10995 });
10996 let response = build_router(state)
10997 .oneshot(chat_request("slow-budget"))
10998 .await
10999 .unwrap();
11000 let body = tokio::spawn(async move {
11001 axum::body::to_bytes(response.into_body(), 64 * 1024)
11002 .await
11003 .unwrap()
11004 });
11005
11006 let started = std::time::Instant::now();
11007 tokio::time::sleep(Duration::from_millis(10)).await;
11008 assert!(
11015 started.elapsed() < Duration::from_millis(250),
11016 "completion budget storage blocked the Tokio reactor"
11017 );
11018 assert!(!body.await.unwrap().is_empty());
11019 assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 2);
11020 provider.abort();
11021 }
11022
11023 #[tokio::test]
11024 async fn completion_tokens_are_blocked_before_delivery() {
11025 let directory = tempfile::tempdir().unwrap();
11026 let prompt_tokens = pv_core::tokens::approx_tokens(&chat_payload().to_string());
11027 let (state, provider) = test_state_with_token_cap(directory.path(), Some(prompt_tokens)).await;
11028 let app = build_router(state.clone());
11029 let response = app.oneshot(chat_request("completion-budget")).await.unwrap();
11030 assert_eq!(response.status(), StatusCode::OK);
11031 assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
11033 .await
11034 .is_err());
11035 let session = state.sessions.get("completion-budget").unwrap();
11041 session.wait_for_streams().await;
11042 assert!(
11043 !session.capture_failed(),
11044 "budget refusal must be scoped to the response, not the session"
11045 );
11046 assert!(
11047 !session.is_closed(),
11048 "budget refusal must not force-close the conversation"
11049 );
11050 provider.abort();
11051 }
11052
11053 #[test]
11071 fn upstream_response_headers_do_not_cross_proxy_trust_boundary() {
11072 use axum::http::HeaderName;
11073 let dangerous = [
11077 "set-cookie",
11078 "access-control-allow-origin",
11079 "access-control-allow-credentials",
11080 "access-control-allow-methods",
11081 "access-control-allow-headers",
11082 "access-control-expose-headers",
11083 "access-control-max-age",
11084 "connection",
11085 "keep-alive",
11086 "transfer-encoding",
11087 "upgrade",
11088 "te",
11089 "trailer",
11090 "proxy-authenticate",
11091 "proxy-authorization",
11092 "content-length",
11093 "server",
11094 "via",
11095 "x-powered-by",
11096 "x-request-id",
11097 "authorization",
11101 "api-key",
11102 "x-api-key",
11103 "x-goog-api-key",
11104 "anthropic-api-key",
11105 "openai-api-key",
11106 "x-auth-token",
11107 "x-amz-security-token",
11108 ];
11109 for name in dangerous {
11110 let header = HeaderName::from_static(name);
11111 assert!(
11112 !is_forwardable_upstream_header(&header, "authorization"),
11113 "dangerous header {name:?} must not be forwarded to the client"
11114 );
11115 }
11116
11117 let exotic = HeaderName::from_static("x-my-secret");
11121 assert!(
11122 !is_forwardable_upstream_header(&exotic, "x-my-secret"),
11123 "operator-configured upstream_auth_header must be refused",
11124 );
11125
11126 let benign = [
11127 "content-type",
11128 "content-encoding",
11129 "cache-control",
11130 "etag",
11131 "last-modified",
11132 "vary",
11133 "expires",
11134 ];
11135 for name in benign {
11136 let header = HeaderName::from_static(name);
11137 assert!(
11138 is_forwardable_upstream_header(&header, "authorization"),
11139 "benign header {name:?} must still be forwarded"
11140 );
11141 }
11142 }
11143
11144 #[tokio::test]
11157 async fn mcp_tool_upstream_failure_does_not_leak_configured_url() {
11158 let sentinel_host = "internal-mcp-sentinel-host.corp.example";
11160 let sentinel_url = format!("http://{sentinel_host}:65002/mcp");
11161 let directory = tempfile::tempdir().unwrap();
11162 let (mut state, provider) = test_state(directory.path()).await;
11163 Arc::get_mut(&mut state.config).unwrap().tool_upstream_url = Some(sentinel_url);
11164 let response = build_router(state)
11165 .oneshot(
11166 Request::builder()
11167 .method("POST")
11168 .uri("/v1/mcp")
11169 .header("x-pv-session", "mcp-cwe-209-check")
11170 .body(Body::from(
11171 r#"{"jsonrpc":"2.0","id":"leak","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
11172 ))
11173 .unwrap(),
11174 )
11175 .await
11176 .unwrap();
11177 assert_eq!(
11178 response.status(),
11179 StatusCode::BAD_GATEWAY,
11180 "tool-upstream faults must surface as 502 like the chat relay, not 500"
11181 );
11182 let bytes = axum::body::to_bytes(response.into_body(), 64 * 1024)
11183 .await
11184 .unwrap();
11185 let body = String::from_utf8(bytes.to_vec()).unwrap();
11186 assert!(
11187 !body.contains(sentinel_host),
11188 "MCP forward-tool error body {body:?} leaks the configured tool-upstream host"
11189 );
11190 assert!(
11191 !body.contains("65002"),
11192 "MCP forward-tool error body {body:?} leaks the configured tool-upstream port"
11193 );
11194 assert!(
11195 !body.contains("corp.example"),
11196 "MCP forward-tool error body {body:?} leaks the configured tool-upstream domain"
11197 );
11198 provider.abort();
11199 }
11200
11201 #[tokio::test]
11202 async fn pipeline_error_carries_retry_after_and_www_authenticate() {
11203 use crate::pipeline::PipelineError;
11210
11211 let unavailable = pipeline_error(PipelineError::unavailable("worker queue full"));
11212 assert_eq!(unavailable.status(), StatusCode::SERVICE_UNAVAILABLE);
11213 assert_eq!(
11214 unavailable
11215 .headers()
11216 .get(axum::http::header::RETRY_AFTER)
11217 .and_then(|v| v.to_str().ok()),
11218 Some("5")
11219 );
11220
11221 let upstream = pipeline_error(PipelineError::upstream("bad gateway"));
11222 assert_eq!(upstream.status(), StatusCode::BAD_GATEWAY);
11223 assert!(upstream.headers().get(axum::http::header::RETRY_AFTER).is_some());
11224
11225 let unauthorized = pipeline_error(PipelineError::Unauthorized("bad token".into()));
11226 assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED);
11227 let auth = unauthorized
11228 .headers()
11229 .get(axum::http::header::WWW_AUTHENTICATE)
11230 .and_then(|v| v.to_str().ok())
11231 .unwrap_or_default();
11232 assert!(
11233 auth.starts_with("Bearer"),
11234 "WWW-Authenticate must be a Bearer challenge, got {auth:?}"
11235 );
11236
11237 let blocked = pipeline_error(PipelineError::blocked("loop breaker"));
11240 assert_eq!(blocked.status(), StatusCode::FORBIDDEN);
11241 assert!(blocked.headers().get(axum::http::header::RETRY_AFTER).is_none());
11242 let abort = pipeline_error(PipelineError::Abort("stop retrying".into()));
11243 assert_eq!(abort.status(), StatusCode::CONFLICT);
11244 assert!(abort.headers().get(axum::http::header::RETRY_AFTER).is_none());
11245 }
11246
11247 #[test]
11255 fn is_sse_content_type_is_case_insensitive_and_param_tolerant() {
11256 fn ct(value: &str) -> HeaderMap {
11257 let mut headers = HeaderMap::new();
11258 headers.insert(
11259 axum::http::header::CONTENT_TYPE,
11260 value.parse().expect("header value"),
11261 );
11262 headers
11263 }
11264 assert!(is_sse_content_type(&ct("text/event-stream")));
11265 assert!(is_sse_content_type(&ct("Text/Event-Stream")));
11266 assert!(is_sse_content_type(&ct("TEXT/EVENT-STREAM")));
11267 assert!(is_sse_content_type(&ct("text/event-stream; charset=utf-8")));
11268 assert!(is_sse_content_type(&ct("Text/Event-Stream ; charset=utf-8")));
11269 assert!(!is_sse_content_type(&ct("application/json")));
11271 assert!(!is_sse_content_type(&ct("text/event-stream-json")));
11272 assert!(!is_sse_content_type(&HeaderMap::new()));
11274 }
11275
11276 #[tokio::test]
11283 async fn options_returns_204_without_cors_headers() {
11284 let directory = tempfile::tempdir().unwrap();
11285 let (state, provider) = test_state(directory.path()).await;
11286 for path in [
11287 "/v1/chat/completions",
11288 "/v1/mcp",
11289 "/mcp",
11290 "/v1/sessions/some-id/close",
11291 "/v1/sessions/some-id/promote",
11292 ] {
11293 let response = build_router(state.clone())
11294 .oneshot(
11295 Request::builder()
11296 .method(axum::http::Method::OPTIONS)
11297 .uri(path)
11298 .header("origin", "http://attacker.example")
11299 .header("access-control-request-method", "POST")
11300 .body(Body::empty())
11301 .unwrap(),
11302 )
11303 .await
11304 .unwrap();
11305 assert_eq!(
11306 response.status(),
11307 StatusCode::NO_CONTENT,
11308 "OPTIONS {path} should return 204, got {}",
11309 response.status()
11310 );
11311 for header in [
11312 "access-control-allow-origin",
11313 "access-control-allow-methods",
11314 "access-control-allow-headers",
11315 "access-control-allow-credentials",
11316 ] {
11317 assert!(
11318 response.headers().get(header).is_none(),
11319 "OPTIONS {path} must not emit {header} — got {:?}",
11320 response.headers().get(header)
11321 );
11322 }
11323 }
11324 provider.abort();
11325 }
11326
11327 #[test]
11341 fn round_32_f2_tool_outcome_preserves_content_type() {
11342 let outcome = ToolOutcome {
11344 status: 200,
11345 body_hex: hex::encode(b"{\"jsonrpc\":\"2.0\"}"),
11346 content_type: Some("application/problem+json".to_owned()),
11347 };
11348 let response = outcome.into_response();
11349 assert_eq!(
11350 response
11351 .headers()
11352 .get(axum::http::header::CONTENT_TYPE)
11353 .and_then(|v| v.to_str().ok()),
11354 Some("application/problem+json")
11355 );
11356 let outcome_no_ct = ToolOutcome {
11358 status: 200,
11359 body_hex: hex::encode(b"{}"),
11360 content_type: None,
11361 };
11362 let response = outcome_no_ct.into_response();
11363 assert_eq!(
11364 response
11365 .headers()
11366 .get(axum::http::header::CONTENT_TYPE)
11367 .and_then(|v| v.to_str().ok()),
11368 Some("application/json")
11369 );
11370 let sealed = serde_json::to_string(&ToolOutcome {
11372 status: 429,
11373 body_hex: hex::encode(b"{}"),
11374 content_type: Some("text/plain".to_owned()),
11375 })
11376 .unwrap();
11377 let recovered: ToolOutcome = serde_json::from_str(&sealed).unwrap();
11378 assert_eq!(recovered.content_type.as_deref(), Some("text/plain"));
11379 let legacy = r#"{"status":200,"body_hex":"7b7d"}"#;
11381 let outcome: ToolOutcome = serde_json::from_str(legacy).unwrap();
11382 assert!(outcome.content_type.is_none());
11383 }
11384
11385 #[tokio::test]
11391 async fn livez_and_readyz_reflect_service_readiness() {
11392 let scratch = tempfile::tempdir().unwrap();
11393 let spool = scratch.path().join("spool");
11394 std::fs::create_dir_all(&spool).unwrap();
11395 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
11396 let address = listener.local_addr().unwrap();
11397 let provider = Router::new().route("/v1/chat/completions", post(mock_chat));
11398 let _server = tokio::spawn(async move {
11399 axum::serve(listener, provider).await.unwrap();
11400 });
11401 let config = crate::config::HarnessConfig::for_tests(
11402 &format!("http://{address}"),
11403 &spool.to_string_lossy(),
11404 &spool.to_string_lossy(),
11405 );
11406 let state = AppState::new(
11407 config,
11408 Arc::new(InMemoryStore::new()),
11409 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
11410 Arc::new(NullBus),
11411 None,
11412 Arc::new(Ed25519Signer::from_seed(&[19; 32])),
11413 )
11414 .unwrap();
11415 let app = crate::build_router(state.clone());
11416
11417 let live = app
11420 .clone()
11421 .oneshot(Request::builder().uri("/livez").body(Body::empty()).unwrap())
11422 .await
11423 .unwrap();
11424 assert_eq!(live.status(), StatusCode::OK);
11425
11426 let ready = app
11428 .clone()
11429 .oneshot(Request::builder().uri("/readyz").body(Body::empty()).unwrap())
11430 .await
11431 .unwrap();
11432 assert_eq!(ready.status(), StatusCode::OK);
11433
11434 state.draining.store(true, std::sync::atomic::Ordering::SeqCst);
11436 let not_ready = app
11437 .clone()
11438 .oneshot(Request::builder().uri("/readyz").body(Body::empty()).unwrap())
11439 .await
11440 .unwrap();
11441 assert_eq!(not_ready.status(), StatusCode::SERVICE_UNAVAILABLE);
11442
11443 let still_live = app
11445 .clone()
11446 .oneshot(Request::builder().uri("/livez").body(Body::empty()).unwrap())
11447 .await
11448 .unwrap();
11449 assert_eq!(still_live.status(), StatusCode::OK);
11450 }
11451
11452 #[cfg(unix)]
11461 #[tokio::test]
11462 async fn readyz_fails_when_the_spool_is_not_writable() {
11463 use std::os::unix::fs::PermissionsExt as _;
11464 let scratch = tempfile::tempdir().unwrap();
11465 let spool = scratch.path().join("spool");
11466 std::fs::create_dir_all(&spool).unwrap();
11467 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
11468 let address = listener.local_addr().unwrap();
11469 let provider = Router::new().route("/v1/chat/completions", post(mock_chat));
11470 let _server = tokio::spawn(async move {
11471 axum::serve(listener, provider).await.unwrap();
11472 });
11473 let config = crate::config::HarnessConfig::for_tests(
11474 &format!("http://{address}"),
11475 &spool.to_string_lossy(),
11476 &spool.to_string_lossy(),
11477 );
11478 let state = AppState::new(
11479 config,
11480 Arc::new(InMemoryStore::new()),
11481 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
11482 Arc::new(NullBus),
11483 None,
11484 Arc::new(Ed25519Signer::from_seed(&[23; 32])),
11485 )
11486 .unwrap();
11487 let app = crate::build_router(state);
11488
11489 std::fs::set_permissions(&spool, std::fs::Permissions::from_mode(0o555)).unwrap();
11491 if std::fs::write(spool.join("root-probe"), b"x").is_ok() {
11492 let _ = std::fs::remove_file(spool.join("root-probe"));
11494 std::fs::set_permissions(&spool, std::fs::Permissions::from_mode(0o755)).unwrap();
11495 return;
11496 }
11497 let not_ready = app
11498 .clone()
11499 .oneshot(Request::builder().uri("/readyz").body(Body::empty()).unwrap())
11500 .await
11501 .unwrap();
11502 assert_eq!(
11503 not_ready.status(),
11504 StatusCode::SERVICE_UNAVAILABLE,
11505 "an unwritable (readable) spool must flip readiness — the full-disk incident shape"
11506 );
11507
11508 std::fs::set_permissions(&spool, std::fs::Permissions::from_mode(0o755)).unwrap();
11510 let ready = app
11511 .clone()
11512 .oneshot(Request::builder().uri("/readyz").body(Body::empty()).unwrap())
11513 .await
11514 .unwrap();
11515 assert_eq!(ready.status(), StatusCode::OK);
11516 let residue = std::fs::read_dir(&spool)
11518 .unwrap()
11519 .filter_map(Result::ok)
11520 .any(|entry| {
11521 entry
11522 .file_name()
11523 .to_str()
11524 .is_some_and(|name| name.starts_with(".readyz-"))
11525 });
11526 assert!(!residue, "probe files must not accumulate in the spool");
11527 }
11528
11529 #[tokio::test]
11536 async fn allowed_hosts_refuses_rebound_host_headers() {
11537 let scratch = tempfile::tempdir().unwrap();
11538 let spool = scratch.path().join("spool");
11539 std::fs::create_dir_all(&spool).unwrap();
11540 let mut config = crate::config::HarnessConfig::for_tests(
11541 "http://127.0.0.1:9",
11542 &spool.to_string_lossy(),
11543 &spool.to_string_lossy(),
11544 );
11545 config.allowed_hosts = vec!["localhost".to_owned(), "[::1]".to_owned()];
11546 let state = AppState::new(
11547 config,
11548 Arc::new(InMemoryStore::new()),
11549 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
11550 Arc::new(NullBus),
11551 None,
11552 Arc::new(Ed25519Signer::from_seed(&[23; 32])),
11553 )
11554 .unwrap();
11555 let app = crate::build_router(state);
11556
11557 let probe = |host: Option<&'static str>| {
11558 let mut builder = Request::builder().uri("/health");
11559 if let Some(host) = host {
11560 builder = builder.header("host", host);
11561 }
11562 builder.body(Body::empty()).unwrap()
11563 };
11564
11565 for host in ["localhost", "localhost:8484", "LOCALHOST", "[::1]:8484"] {
11567 let response = app.clone().oneshot(probe(Some(host))).await.unwrap();
11568 assert_eq!(response.status(), StatusCode::OK, "host {host:?} must be allowed");
11569 }
11570 for host in [
11572 Some("evil.attacker.example"),
11573 Some("localhost.attacker.example"),
11574 None,
11575 ] {
11576 let response = app.clone().oneshot(probe(host)).await.unwrap();
11577 assert_eq!(
11578 response.status(),
11579 StatusCode::FORBIDDEN,
11580 "host {host:?} must be refused"
11581 );
11582 }
11583 }
11584
11585 #[tokio::test]
11590 async fn metrics_scrape_emits_open_sessions_gauge_and_signing_key_info() {
11591 let scratch = tempfile::tempdir().unwrap();
11592 let spool = scratch.path().join("spool");
11593 std::fs::create_dir_all(&spool).unwrap();
11594 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
11595 let address = listener.local_addr().unwrap();
11596 let provider = Router::new().route("/v1/chat/completions", post(mock_chat));
11597 let _server = tokio::spawn(async move {
11598 axum::serve(listener, provider).await.unwrap();
11599 });
11600 let config = crate::config::HarnessConfig::for_tests(
11601 &format!("http://{address}"),
11602 &spool.to_string_lossy(),
11603 &spool.to_string_lossy(),
11604 );
11605 let signer = Ed25519Signer::from_seed(&[31; 32]);
11606 use pv_receipts::Signer as _;
11607 let key_id = signer.key_id().to_owned();
11608 let public_key_hex = hex::encode(signer.public_key_bytes());
11609 let state = AppState::new(
11610 config,
11611 Arc::new(InMemoryStore::new()),
11612 Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
11613 Arc::new(NullBus),
11614 None,
11615 Arc::new(signer),
11616 )
11617 .unwrap();
11618 state
11621 .metrics
11622 .gauge(
11623 &format!("pv_signing_key_info{{key_id=\"{key_id}\",public_key_hex=\"{public_key_hex}\"}}"),
11624 "test signing key info",
11625 )
11626 .set(1);
11627 let app = crate::build_router(state.clone());
11628
11629 let response = app
11631 .clone()
11632 .oneshot(Request::builder().uri("/metrics").body(Body::empty()).unwrap())
11633 .await
11634 .unwrap();
11635 let body = String::from_utf8_lossy(
11636 &axum::body::to_bytes(response.into_body(), 256 * 1024)
11637 .await
11638 .unwrap(),
11639 )
11640 .into_owned();
11641 assert!(body.contains("# TYPE pv_open_sessions gauge"), "{body}");
11642 assert!(body.contains("pv_open_sessions 0"), "{body}");
11643 assert!(body.contains("# TYPE pv_worker_queue_depth gauge"), "{body}");
11645 assert!(body.contains("# TYPE pv_spool_bytes gauge"), "{body}");
11646 assert!(body.contains("# TYPE pv_requests_total counter"), "{body}");
11647 assert!(
11648 body.contains("pv_requests_total{route=\"chat\",status_class=\"2xx\"}"),
11649 "{body}"
11650 );
11651 assert!(
11652 body.contains("pv_upstream_errors_total{kind=\"timeout\"}"),
11653 "{body}"
11654 );
11655 assert!(body.contains("pv_upstream_latency_seconds"), "{body}");
11656 assert!(
11657 body.contains(&format!("pv_signing_key_info{{key_id=\"{key_id}\"")),
11658 "signing-key info gauge missing from scrape: {body}"
11659 );
11660
11661 state.sessions.insert_recovered(crate::session::Session::new(
11663 "gauge-test".to_owned(),
11664 crate::session::Workflow::Unsigned,
11665 pv_events::AgentIdentity {
11666 version: "1".to_owned(),
11667 charter: "test".into(),
11668 instance_uid: "instance-1".to_owned(),
11669 ttl_remaining_s: Some(600),
11670 },
11671 Default::default(),
11672 ));
11673 let response = app
11674 .clone()
11675 .oneshot(Request::builder().uri("/metrics").body(Body::empty()).unwrap())
11676 .await
11677 .unwrap();
11678 let body = String::from_utf8_lossy(
11679 &axum::body::to_bytes(response.into_body(), 256 * 1024)
11680 .await
11681 .unwrap(),
11682 )
11683 .into_owned();
11684 assert!(body.contains("pv_open_sessions 1"), "{body}");
11685 }
11686
11687 #[tokio::test]
11697 async fn decode_cached_tool_body_returns_stable_error_message() {
11698 let helper_err = decode_cached_tool_body("zz").expect_err("`zz` must not hex-decode");
11699 let body = String::from_utf8(
11700 axum::body::to_bytes((*helper_err).into_body(), 16 * 1024)
11701 .await
11702 .unwrap()
11703 .to_vec(),
11704 )
11705 .expect("lifecycle_error body is UTF-8");
11706 assert!(
11707 body.contains("cached tool outcome is corrupt"),
11708 "leak-free stable error message expected in body: {body}"
11709 );
11710 assert!(
11714 !body.contains("at 0"),
11715 "hex::FromHexError position must not leak: {body}"
11716 );
11717 let outcome_response = ToolOutcome {
11720 status: 200,
11721 body_hex: "zz".to_owned(),
11722 content_type: None,
11723 }
11724 .into_response();
11725 let outcome_body = String::from_utf8(
11726 axum::body::to_bytes(outcome_response.into_body(), 16 * 1024)
11727 .await
11728 .unwrap()
11729 .to_vec(),
11730 )
11731 .expect("into_response body is UTF-8");
11732 assert!(
11733 outcome_body.contains("cached tool outcome is corrupt"),
11734 "ToolOutcome::into_response must route through the shared helper: {outcome_body}"
11735 );
11736 }
11737}