av_harness/
routes.rs

1//! Axum HTTP routes for proxy, MCP interception, lifecycle, and operations.
2
3use crate::pipeline::AppState;
4use av_core::time::elapsed_us;
5use av_events::StopReason;
6use av_sandbox::ToolVerdict;
7use axum::body::{Body, Bytes};
8use axum::extract::{Path, State};
9use axum::http::{HeaderMap, HeaderValue, Request, StatusCode};
10use axum::middleware::Next;
11use axum::response::{IntoResponse, Response};
12use axum::routing::{get, post};
13use axum::{Json, Router};
14use futures::stream::BoxStream;
15use futures::{Stream, StreamExt};
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
23/// Build the complete harness router.
24pub 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("/metrics", get(metrics))
30        .route("/v1/chat/completions", post(chat_completions).options(cors_deny))
31        .route("/v1/mcp", post(mcp_call).options(cors_deny))
32        .route("/mcp", post(mcp_call).options(cors_deny))
33        .route("/v1/sessions/{id}/close", post(close_session).options(cors_deny))
34        .route(
35            "/v1/sessions/{id}/promote",
36            post(promote_session).options(cors_deny),
37        );
38    if dashboard_enabled {
39        router = router
40            .route("/dashboard", get(crate::dashboard::index))
41            .route("/dashboard/", get(crate::dashboard::index))
42            .route("/dashboard/style.css", get(crate::dashboard::style_css))
43            .route("/dashboard/app.js", get(crate::dashboard::app_js))
44            .route("/api/v1/dashboard/stats", get(crate::dashboard::stats))
45            .route("/api/v1/dashboard/sessions", get(crate::dashboard::list_sessions))
46            .route(
47                "/api/v1/dashboard/sessions/{id}",
48                get(crate::dashboard::session_detail),
49            );
50    }
51    router
52        .layer(axum::middleware::from_fn(trace_request))
53        // axum's default body limit is 2 MiB, which silently rejects large
54        // chat contexts (Claude 200k, GPT-4 128k) before the sandbox even
55        // sees the payload. `max_request_bytes` (default 4 MiB, matching
56        // the sandbox `MAX_PAYLOAD_BYTES`) is the single knob operators
57        // control.
58        .layer(axum::extract::DefaultBodyLimit::max(max_body))
59        .with_state(state)
60}
61
62async fn trace_request(request: Request<Body>, next: Next) -> Response {
63    let method = request.method().clone();
64    let path = request.uri().path().to_owned();
65    // Round-13 F5: sanitize the session id BEFORE binding it into the
66    // span. Previously the raw header value went in verbatim, which
67    // (1) risked unbounded label cardinality on OTLP exporters that
68    // map span attributes to metric labels — every distinct
69    // client-supplied session id (including hostile garbage) became
70    // its own series; (2) created a split-brain when a client sent
71    // two `X-AV-Session` headers — `HeaderMap::get` returned the
72    // first while `pipeline::single_header` (round-13) refused the
73    // whole request, so traces named a "friendly" id for a hard-400;
74    // (3) accepted values that pass `HeaderValue::to_str` but fail
75    // `SessionId::parse` (too long, control chars in disguise). Use
76    // `single_header` + `SessionId::parse` to bind ONE consistent
77    // value or the sentinel `"invalid"`.
78    // Round-14 F7: use a sentinel that CANNOT pass `SessionId::parse`
79    // (0x21..=0x7e visible-ASCII only), so a client cannot legitimately
80    // send `X-AV-Session: invalid` and share a trace label with a
81    // rejected request. The leading space (0x20) is outside the
82    // allowed range → collision-free.
83    let headers = request.headers();
84    let session_id = match crate::pipeline::single_header(headers, crate::pipeline::SESSION_HEADER) {
85        Ok(Some(value)) => value
86            .to_str()
87            .ok()
88            .and_then(|v| av_core::SessionId::parse(v).ok())
89            .map(|id| id.to_string())
90            .unwrap_or_else(|| " rejected".to_owned()),
91        Ok(None) => "unbound".to_owned(),
92        Err(_) => " duplicate-header".to_owned(),
93    };
94    let span = tracing::info_span!(
95        "agentvisor.request",
96        otel.kind = "server",
97        http.request.method = %method,
98        url.path = %path,
99        session.id = %session_id,
100        http.response.status_code = tracing::field::Empty,
101    );
102    let response = next.run(request).instrument(span.clone()).await;
103    span.record("http.response.status_code", response.status().as_u16());
104    response
105}
106
107async fn health() -> impl IntoResponse {
108    // Round-29 F6: DO NOT expose the CARGO_PKG_VERSION on this
109    // unauthenticated endpoint. Version disclosure lets a LAN
110    // attacker correlate an agentvisor deployment to a specific
111    // known-vulnerable release without needing a chat probe. No
112    // harness endpoint surfaces build info (`/metrics` included);
113    // the version appears only in the *outbound* upstream
114    // `User-Agent` built in `pipeline.rs`.
115    Json(json!({
116        "status": "ok",
117        // Product identifier so callers (avctl start) can tell a real
118        // AgentVisor AI apart from an unrelated service squatting the
119        // port. This is not a version number and reveals no
120        // vulnerability-relevant information.
121        "service": "agentvisor",
122    }))
123}
124
125/// Round-31 F5: explicit deny-CORS OPTIONS handler.
126///
127/// The harness is a same-origin proxy; no cross-origin client is
128/// expected or supported. Without this route, axum's default reply to
129/// a preflight (`OPTIONS /v1/chat/completions`) is `405 Method Not
130/// Allowed` with an `Allow: POST` header — inconsistent with the
131/// round-29 F6 "no discoverable posture" hygiene, and confusing to any
132/// operator whose LAN browser client accidentally triggers a preflight.
133/// Reply with `204 No Content` and NO `Access-Control-Allow-Origin`
134/// header: browsers correctly treat this as "cross-origin denied" and
135/// refuse the actual request, making the posture explicit at the wire.
136async fn cors_deny() -> Response {
137    let mut response = Response::new(Body::empty());
138    *response.status_mut() = StatusCode::NO_CONTENT;
139    response.headers_mut().insert(
140        axum::http::header::CACHE_CONTROL,
141        HeaderValue::from_static("no-store"),
142    );
143    // Do not echo the requester's Origin. Do not include any
144    // Access-Control-Allow-* header. This deliberately fails the
145    // browser's preflight check.
146    response
147}
148
149async fn metrics(State(state): State<AppState>) -> Response {
150    let mut response = Response::new(Body::from(state.metrics.render()));
151    response.headers_mut().insert(
152        axum::http::header::CONTENT_TYPE,
153        HeaderValue::from_static("text/plain; version=0.0.4; charset=utf-8"),
154    );
155    response
156}
157
158async fn chat_completions(
159    State(state): State<AppState>,
160    headers: HeaderMap,
161    Json(payload): Json<Value>,
162) -> Response {
163    let admission_started = std::time::Instant::now();
164    let mut prepared = match state.prepare_chat_nonblocking(&headers, payload).await {
165        Ok(prepared) => prepared,
166        Err(error) => return pipeline_error(error),
167    };
168    prepared.middleware_us = elapsed_us(admission_started);
169    let session = Arc::clone(&prepared.session);
170    let identity = prepared.identity.clone();
171    let session_id = prepared.session.id.clone();
172    let middleware_us = prepared.middleware_us;
173    let forwarded = match state.forward_chat(prepared).await {
174        Ok(response) => response,
175        Err(error) => return pipeline_error(error),
176    };
177    let crate::pipeline::ForwardedResponse {
178        response: upstream,
179        lease,
180        response_permit,
181        response_marker,
182        response_attempt_id,
183    } = forwarded;
184    let Some(response_permit) = response_permit else {
185        return lifecycle_error("durable response capture permit is missing".to_owned());
186    };
187
188    let status = upstream.status();
189    let upstream_headers = upstream.headers().clone();
190    // reqwest is built without the `gzip` feature (see Cargo.toml), so
191    // the client does not decode `Content-Encoding: gzip`. If the
192    // provider or a CDN in front of it compresses the SSE response,
193    // `absorb_network_chunk` immediately fails
194    // `std::str::from_utf8(&frame)` and aborts every stream. Refuse
195    // the response up front with a clear error rather than corrupt
196    // the audit trail with UTF-8-error frames.
197    //
198    // Multi-value / multi-line semantics: an upstream that sends
199    // `Content-Encoding: identity` on one line and
200    // `Content-Encoding: gzip` on another (some layered reverse
201    // proxies do this) would slip past a `.get(...)` that only reads
202    // the first value. Iterate `get_all` and require EVERY value to
203    // be empty or `identity` case-insensitively; also split each
204    // value on `,` so `gzip, identity` (single header, two tokens)
205    // is caught.
206    for value in upstream_headers.get_all(axum::http::header::CONTENT_ENCODING) {
207        let raw = value.to_str().unwrap_or_default();
208        for token in raw.split(',') {
209            let token = token.trim();
210            if token.is_empty() || token.eq_ignore_ascii_case("identity") {
211                continue;
212            }
213            return pipeline_error(crate::pipeline::PipelineError::Upstream(format!(
214                "upstream responded with unsupported Content-Encoding token {token:?} \
215                 (full header: {raw:?}) — the proxy is built without decompression \
216                 support; enable it upstream (Accept-Encoding: identity) or rebuild \
217                 with the reqwest `gzip` feature"
218            )));
219        }
220    }
221    // Round-25 F3: RFC 7231 §3.1.1.1 says media type/subtype are
222    // case-insensitive and the header value may carry parameters
223    // (`; charset=utf-8`). Byte-exact `starts_with("text/event-stream")`
224    // misses `Text/Event-Stream` (some CDNs re-title-case) and can also
225    // mis-fire on hypothetical `text/event-stream-json`. Split on `;`
226    // then compare with `eq_ignore_ascii_case` so SSE is detected
227    // regardless of casing and parameters. A misclassified SSE stream
228    // gets buffered to the 16 MiB provider cap in the non-SSE branch
229    // and either loses streaming semantics (client sees no deltas
230    // until upstream EOF) or is refused with 502 despite being valid.
231    let is_sse = is_sse_content_type(&upstream_headers);
232    let stream = upstream
233        .bytes_stream()
234        .map(|chunk| chunk.map_err(std::io::Error::other));
235    let stream = AbortFinalizingStream {
236        inner: stream.boxed(),
237        session,
238        identity,
239        response_permit: Some(response_permit),
240        worker: state.worker.clone(),
241        store: Arc::clone(&state.store),
242        budget: state.config.budget.clone(),
243        finalizer: state.finalizer.clone(),
244        _lease: lease,
245        response_marker,
246        response_attempt_id,
247        response_message: String::new(),
248        response_reasoning: String::new(),
249        response_model: None,
250        response_finish_reason: None,
251        upstream_status: status,
252        response_cost_usd_micros: 0,
253        response_tool_calls: std::collections::BTreeMap::new(),
254        response_metrics: av_events::EventMetrics::default(),
255        charged_completion_tokens: 0,
256        last_reported_completion_tokens: None,
257        last_reported_prompt_tokens: None,
258        last_reported_cached_tokens: None,
259        last_reported_cost_usd_micros: None,
260        saw_chunk: false,
261        capture_submitted: false,
262        is_sse,
263        protocol_buffer: Vec::new(),
264        pending_output: std::collections::VecDeque::new(),
265        pending_budget: None,
266        captured_bytes: 0,
267        completed: false,
268    };
269    let mut response = if is_sse {
270        Response::new(Body::from_stream(stream))
271    } else {
272        // A non-SSE body is fully buffered inside the relay before any byte
273        // is released: worker capture and the completion-token budget gate
274        // run on the complete body. Handing axum a stream here would commit
275        // the upstream status line before those gates run, so a budget
276        // refusal degraded into a 200 head followed by an empty body with
277        // no explanation. Drain the relay first (no extra buffering beyond
278        // what the relay already does) and surface refusals as real errors.
279        let mut relay = Box::pin(stream);
280        let mut buffered: Vec<u8> = Vec::new();
281        loop {
282            match relay.next().await {
283                Some(Ok(bytes)) => buffered.extend_from_slice(&bytes),
284                Some(Err(error)) => {
285                    let refusal_status = if error.kind() == std::io::ErrorKind::QuotaExceeded {
286                        StatusCode::TOO_MANY_REQUESTS
287                    } else {
288                        StatusCode::BAD_GATEWAY
289                    };
290                    // Dropping the relay here runs its finalization Drop
291                    // (evidence capture + session seal), same as when a
292                    // client observed the severed stream.
293                    return (refusal_status, Json(json!({"error": error.to_string()}))).into_response();
294                }
295                None => break,
296            }
297        }
298        Response::new(Body::from(buffered))
299    };
300    *response.status_mut() = status;
301    for (name, value) in &upstream_headers {
302        if is_forwardable_upstream_header(name) {
303            // `append`, not `insert`: iterating a HeaderMap repeats the name
304            // for each value of a multi-valued header, and `insert` would
305            // keep only the last one.
306            response.headers_mut().append(name.clone(), value.clone());
307        }
308    }
309    if let Ok(value) = HeaderValue::from_str(&session_id) {
310        response
311            .headers_mut()
312            .insert(crate::pipeline::SESSION_HEADER, value);
313    }
314    if let Ok(value) = HeaderValue::from_str(&middleware_us.to_string()) {
315        response
316            .headers_mut()
317            .insert(crate::pipeline::MIDDLEWARE_US_HEADER, value);
318    }
319    // Round-29 F4: pin `X-Content-Type-Options: nosniff` on every
320    // upstream-relayed response. The relay forwards the upstream's
321    // Content-Type verbatim (validated by our `is_sse_content_type`
322    // for framing decisions, but not sanitised for the client).
323    // A rogue upstream, MITM at egress, or CDN mis-config could
324    // otherwise flip Content-Type to `text/html; charset=utf-8` on
325    // a body carrying prompt-echoed attacker bytes — turning what
326    // the audit trail attests as "assistant output" into HTML the
327    // browser might render. `nosniff` prevents browser-side MIME
328    // sniffing from disagreeing with the declared type and closes
329    // the reflected-content path.
330    response.headers_mut().insert(
331        axum::http::HeaderName::from_static("x-content-type-options"),
332        HeaderValue::from_static("nosniff"),
333    );
334    response
335}
336
337/// Decide whether an upstream response header may be forwarded to the client.
338///
339/// This is a *proxy* trust boundary: the upstream LLM provider is not on the
340/// same trust domain as our client, so blindly forwarding response headers
341/// would let the upstream (or an attacker who influences the upstream)
342/// - set cookies in our domain (`Set-Cookie` — cookie injection),
343/// - open CORS on our origin (`Access-Control-Allow-*`),
344/// - inject invalid or contradictory framing/hop-by-hop metadata
345///   (`Transfer-Encoding`, `Content-Length`, `Connection`, `Keep-Alive`,
346///   `Trailer`, `TE`, `Upgrade`, `Proxy-Authenticate`, `Proxy-Authorization`),
347/// - leak upstream implementation identity (`Server`, `X-Powered-By`,
348///   `Via`, `X-Request-ID`).
349///
350/// Hyper computes framing metadata (`Content-Length`, `Transfer-Encoding`)
351/// itself when we build the response; forwarding the upstream's copies risks
352/// double-encoded or contradictory headers and — with `Transfer-Encoding` —
353/// classical HTTP request smuggling. Hop-by-hop headers are forbidden by
354/// RFC 7230 §6.1 from crossing a proxy.
355fn is_forwardable_upstream_header(name: &axum::http::HeaderName) -> bool {
356    use axum::http::header;
357    let is_denied = *name == header::CONTENT_LENGTH
358        || *name == header::TRANSFER_ENCODING
359        || *name == header::CONNECTION
360        || *name == header::UPGRADE
361        || *name == header::TE
362        || *name == header::TRAILER
363        || *name == header::PROXY_AUTHENTICATE
364        || *name == header::PROXY_AUTHORIZATION
365        || *name == header::SET_COOKIE
366        || *name == header::SERVER
367        || *name == header::VIA
368        || name.as_str().eq_ignore_ascii_case("keep-alive")
369        || name.as_str().eq_ignore_ascii_case("x-powered-by")
370        || name.as_str().eq_ignore_ascii_case("x-request-id")
371        // Never let the upstream open CORS on our origin — if we want CORS
372        // we set it deliberately in our own router.
373        || name.as_str().to_ascii_lowercase().starts_with("access-control-");
374    !is_denied
375}
376
377/// Round-25 F3: detect `text/event-stream` in `Content-Type`
378/// case-insensitively, accepting parameters like `; charset=utf-8`.
379/// RFC 7231 §3.1.1.1 says media type/subtype are case-insensitive.
380/// Byte-exact matching missed `Text/Event-Stream` (some CDNs
381/// re-title-case) and mis-fired on hypothetical
382/// `text/event-stream-json`. Misclassification cost: the non-SSE
383/// branch buffers up to the 16 MiB provider cap and either loses
384/// streaming semantics or refuses valid streams with 502.
385fn is_sse_content_type(headers: &HeaderMap) -> bool {
386    headers
387        .get(axum::http::header::CONTENT_TYPE)
388        .and_then(|value| value.to_str().ok())
389        .is_some_and(|value| {
390            let head = value.split(';').next().unwrap_or("").trim();
391            head.eq_ignore_ascii_case("text/event-stream")
392        })
393}
394
395async fn mcp_call(State(state): State<AppState>, headers: HeaderMap, body: Bytes) -> Response {
396    let (execution, unaudited_outcome) = if state.config.tool_upstream_url.is_some() {
397        match ToolExecution::from_request(&state.config.atif_spool_dir, &headers, &body, state.journal_key) {
398            Ok(mut execution) => {
399                let required_scope = crate::pipeline::tool_scope(&execution.tool);
400                let identity = match state.resolve_identity(&headers, Some(&required_scope)) {
401                    Ok(identity) => identity,
402                    Err(error) => return pipeline_error(error),
403                };
404                if let Err(error) = execution.bind_principal(&identity) {
405                    return pipeline_error(error);
406                }
407                match execution.load().await {
408                    Ok(ToolExecutionState::Completed(outcome)) => {
409                        let Some(session) = state.sessions.get(&execution.session_id) else {
410                            return pipeline_error(crate::pipeline::PipelineError::BadRequest(
411                                "unknown session for cached tool result".to_owned(),
412                            ));
413                        };
414                        if let Err(error) = state.authorize_session(
415                            &headers,
416                            &session,
417                            &crate::pipeline::tool_scope(&execution.tool),
418                        ) {
419                            return pipeline_error(error);
420                        }
421                        return outcome.into_response();
422                    }
423                    Ok(ToolExecutionState::Unaudited(outcome)) => {
424                        let Some(session) = state.sessions.get(&execution.session_id) else {
425                            return pipeline_error(crate::pipeline::PipelineError::BadRequest(
426                                "unknown session for pending tool audit".to_owned(),
427                            ));
428                        };
429                        if let Err(error) = state.authorize_session(
430                            &headers,
431                            &session,
432                            &crate::pipeline::tool_scope(&execution.tool),
433                        ) {
434                            return pipeline_error(error);
435                        }
436                        (Some(execution), Some(outcome))
437                    }
438                    Ok(ToolExecutionState::Pending) => {
439                        let Some(session) = state.sessions.get(&execution.session_id) else {
440                            return pipeline_error(crate::pipeline::PipelineError::BadRequest(
441                                "unknown session for pending tool execution".to_owned(),
442                            ));
443                        };
444                        if let Err(error) = state.authorize_session(
445                            &headers,
446                            &session,
447                            &crate::pipeline::tool_scope(&execution.tool),
448                        ) {
449                            return pipeline_error(error);
450                        }
451                        return (
452                            StatusCode::CONFLICT,
453                            Json(json!({"error": TOOL_OUTCOME_UNCERTAIN})),
454                        )
455                            .into_response();
456                    }
457                    Ok(ToolExecutionState::Missing) => (Some(execution), None),
458                    Err(error) if error == TOOL_REQUEST_MISMATCH => {
459                        return (
460                            StatusCode::CONFLICT,
461                            Json(json!({"error": TOOL_REQUEST_MISMATCH})),
462                        )
463                            .into_response();
464                    }
465                    Err(error) => return lifecycle_error(error),
466                }
467            }
468            Err(error) => return pipeline_error(error),
469        }
470    } else {
471        (None, None)
472    };
473    if let (Some(execution), Some(outcome)) = (execution.as_ref(), unaudited_outcome) {
474        let _lease = match state.lease_session(&headers) {
475            Ok(lease) => lease,
476            Err(error) => return pipeline_error(error),
477        };
478        let completion_permit = match state.worker.try_reserve(&execution.session_id) {
479            Ok(permit) => permit,
480            Err(error) => {
481                return pipeline_error(crate::pipeline::PipelineError::Unavailable(error.to_string()));
482            }
483        };
484        let Some(session) = state.sessions.get(&execution.session_id) else {
485            return lifecycle_error("tool session disappeared".to_owned());
486        };
487        return complete_tool_audit(execution, outcome, completion_permit, session).await;
488    }
489    // Round-33 F1: closes the round-32 F3 concurrent-MCP budget
490    // double-spend by threading the debited `payout_micros` out of
491    // `ToolVerdict::Allowed` and calling `ActionBudget::refund_
492    // tool_call` on the lost-claim branch. `refund` is best-effort
493    // (backend errors are silently absorbed) so a Redis blip on the
494    // compensation path cannot turn the CONFLICT response into 5xx.
495    match state.intercept_tool_nonblocking(&headers, &body).await {
496        Ok(ToolVerdict::Allowed {
497            tool,
498            budget_remaining,
499            elapsed_us,
500            payout_micros,
501        }) => {
502            if let Some(url) = state.config.tool_upstream_url.as_deref() {
503                let _lease = match state.lease_session(&headers) {
504                    Ok(lease) => lease,
505                    Err(error) => return pipeline_error(error),
506                };
507                let execution = match execution {
508                    Some(execution) => execution,
509                    None => return lifecycle_error("tool execution state is missing".to_owned()),
510                };
511                let completion_permit = match state.worker.try_reserve(&execution.session_id) {
512                    Ok(permit) => permit,
513                    Err(error) => {
514                        return pipeline_error(crate::pipeline::PipelineError::Unavailable(
515                            error.to_string(),
516                        ));
517                    }
518                };
519                if let Err(error) = execution.claim().await {
520                    // A lost claim race means another in-flight request owns
521                    // this execution; answer exactly like the Pending state
522                    // and keep the underlying io detail out of the wire.
523                    tracing::warn!(
524                        session = %execution.session_id,
525                        error = %error,
526                        "concurrent tool execution claim lost; refunding budget"
527                    );
528                    // Round-33 F1: refund the exact amount debited so
529                    // `payout_remaining` and per-tool counters reflect
530                    // only admitted work, not the lost race.
531                    av_state::ActionBudget::new(
532                        state.store.as_ref(),
533                        &execution.session_id,
534                        &state.config.budget,
535                    )
536                    .refund_tool_call(&execution.tool, payout_micros);
537                    return (
538                        StatusCode::CONFLICT,
539                        Json(json!({"error": TOOL_OUTCOME_UNCERTAIN})),
540                    )
541                        .into_response();
542                }
543                let mut tool_request = state
544                    .client
545                    .post(url)
546                    .header(axum::http::header::CONTENT_TYPE, "application/json")
547                    .body(body);
548                if let Some(bearer) = &state.tool_auth {
549                    tool_request = tool_request.header(axum::http::header::AUTHORIZATION, bearer.clone());
550                }
551                match tool_request.send().await {
552                    Ok(upstream) => {
553                        let status = upstream.status();
554                        match read_limited_tool_response(upstream).await {
555                            Ok((bytes, content_type)) => {
556                                let outcome = ToolOutcome {
557                                    status: status.as_u16(),
558                                    body_hex: hex::encode(&bytes),
559                                    content_type,
560                                };
561                                if let Err(error) = execution.persist(&outcome).await {
562                                    return lifecycle_error(error);
563                                }
564                                let session = match state.sessions.get(&execution.session_id) {
565                                    Some(session) => session,
566                                    None => return lifecycle_error("tool session disappeared".to_owned()),
567                                };
568                                complete_tool_audit(&execution, outcome, completion_permit, session).await
569                            }
570                            Err(error) => pipeline_error(crate::pipeline::PipelineError::Upstream(format!(
571                                "read tool response: {error}"
572                            ))),
573                        }
574                    }
575                    Err(error) => {
576                        // CWE-209: `reqwest::Error::Display` embeds the request URL —
577                        // returning it verbatim would leak the operator-configured
578                        // tool-upstream URL (potentially an internal hostname) to
579                        // whichever client called `/mcp`. Report a stable category
580                        // and preserve the raw detail server-side for operators.
581                        //
582                        // Round-34 F4: also do not log the raw error to
583                        // tracing::warn — `reqwest::Error::Display` embeds
584                        // the same URL, and the tracing subscriber flows
585                        // to Vector -> OTLP -> SIEM per the deploy
586                        // topology. If OTLP is exported to a third party
587                        // with a lower trust boundary than the operator,
588                        // the internal `tool_upstream_url` leaks there.
589                        // Log structured fields only.
590                        let category = crate::pipeline::classify_upstream_error(&error);
591                        tracing::warn!(
592                            category = category,
593                            error.status = ?error.status(),
594                            error.is_timeout = error.is_timeout(),
595                            error.is_connect = error.is_connect(),
596                            error.is_request = error.is_request(),
597                            "tool upstream forwarding failed"
598                        );
599                        // Upstream faults must surface as 502 (as the chat
600                        // relay does), not 500: a 500 blames the harness and
601                        // misroutes operator alerting/retry policy.
602                        pipeline_error(crate::pipeline::PipelineError::Upstream(format!(
603                            "forward tool call: {category}"
604                        )))
605                    }
606                }
607            } else {
608                (
609                    StatusCode::OK,
610                    Json(json!({
611                        "allowed": true,
612                        "tool": tool,
613                        "budget_remaining": budget_remaining,
614                        "decision_us": elapsed_us,
615                    })),
616                )
617                    .into_response()
618            }
619        }
620        Ok(ToolVerdict::Blocked { response, .. }) => (StatusCode::FORBIDDEN, Json(response)).into_response(),
621        Err(error) => pipeline_error(error),
622    }
623}
624
625const MAX_TOOL_RESPONSE_BYTES: usize = 16 * 1024 * 1024;
626
627/// Round-32 F2: capture and preserve the upstream tool response's
628/// `Content-Type` so the MCP client sees exactly what the tool
629/// upstream sent. Without this, `Bytes: IntoResponse` stamps
630/// `application/octet-stream`, which strict JSON-RPC 2.0 clients
631/// (spec: MUST be `application/json`) reject and downstream SIEM
632/// filters mis-classify as binary. We store the string (not
633/// HeaderValue) so it round-trips through the on-disk `ToolOutcome`
634/// journal for cached-outcome replay.
635async fn read_limited_tool_response(response: reqwest::Response) -> Result<(Bytes, Option<String>), String> {
636    let content_type = response
637        .headers()
638        .get(axum::http::header::CONTENT_TYPE)
639        .and_then(|value| value.to_str().ok())
640        .map(str::to_owned);
641    let mut stream = response.bytes_stream();
642    let mut body = Vec::new();
643    while let Some(chunk) = stream.next().await {
644        // CWE-209: `reqwest::Error::Display` embeds the request URL — leaking
645        // the operator-configured tool-upstream URL to the client if we
646        // returned it verbatim. Use the stable classifier and log
647        // structured fields for operators (round-34 F4: never `%error`
648        // — that renders the URL into any downstream OTLP sink).
649        let chunk = chunk.map_err(|error| {
650            let category = crate::pipeline::classify_upstream_error(&error);
651            tracing::warn!(
652                category = category,
653                error.status = ?error.status(),
654                error.is_timeout = error.is_timeout(),
655                error.is_connect = error.is_connect(),
656                error.is_body = error.is_body(),
657                "tool upstream stream chunk failed"
658            );
659            category.to_owned()
660        })?;
661        let next = body
662            .len()
663            .checked_add(chunk.len())
664            .ok_or_else(|| "tool response size overflow".to_owned())?;
665        if next > MAX_TOOL_RESPONSE_BYTES {
666            return Err(format!("tool response exceeds {MAX_TOOL_RESPONSE_BYTES} bytes"));
667        }
668        body.extend_from_slice(&chunk);
669    }
670    Ok((Bytes::from(body), content_type))
671}
672
673/// Build a tool-response with the round-32 F2 Content-Type
674/// preserved. Defaults to `application/json` — MCP is JSON-RPC 2.0
675/// by convention — when the upstream did not set one or set an
676/// unrepresentable value.
677fn tool_response(status: StatusCode, bytes: Bytes, content_type: Option<&str>) -> Response {
678    let mut response = (status, bytes).into_response();
679    let value = content_type
680        .and_then(|ct| HeaderValue::from_str(ct).ok())
681        .unwrap_or_else(|| HeaderValue::from_static("application/json"));
682    response
683        .headers_mut()
684        .insert(axum::http::header::CONTENT_TYPE, value);
685    response
686}
687
688async fn complete_tool_audit(
689    execution: &ToolExecution,
690    outcome: ToolOutcome,
691    completion_permit: crate::worker::WorkerPermit,
692    session: Arc<crate::session::Session>,
693) -> Response {
694    let status = StatusCode::from_u16(outcome.status).unwrap_or(StatusCode::BAD_GATEWAY);
695    let bytes = match hex::decode(&outcome.body_hex) {
696        Ok(bytes) => Bytes::from(bytes),
697        Err(error) => return lifecycle_error(format!("decode tool response: {error}")),
698    };
699    let success = status.is_success();
700    completion_permit.submit(crate::worker::WorkerJob {
701        session: Arc::clone(&session),
702        identity: session.current_identity(),
703        class: av_events::EventClass::Session,
704        payload: json!({
705            "action": "tool_completed",
706            "execution_key": &execution.key,
707            "status": status.as_u16(),
708            "response_sha256": av_core::digest::sha256_hex(&bytes),
709        }),
710        text: String::new(),
711        analyze_loop: false,
712        status: if success {
713            av_events::StatusId::Success
714        } else {
715            av_events::StatusId::Failure
716        },
717        stop_reason: (!success).then_some(StopReason::Other),
718        native_stop_reason: None,
719        metrics: av_events::EventMetrics::default(),
720        cost_usd_micros: 0,
721        atif: Some(crate::worker::AtifCapture {
722            source: av_atif::Source::System,
723            message: Value::String(String::from_utf8_lossy(&bytes).into_owned()),
724            reasoning_content: None,
725            model_name: None,
726            tool_calls: None,
727            observation: None,
728            llm_call_count: None,
729        }),
730        response_marker: None,
731        response_attempt: None,
732    });
733    session.wait_for_worker_jobs().await;
734    if session.capture_failed() {
735        return lifecycle_error("tool completed but completion audit failed".to_owned());
736    }
737    if let Err(error) = execution.mark_audited().await {
738        return lifecycle_error(error);
739    }
740    // Round-32 F2: preserve the upstream Content-Type so a spec-
741    // conforming JSON-RPC 2.0 client sees `application/json`
742    // (default) or whatever the tool upstream declared, not axum's
743    // `application/octet-stream`.
744    tool_response(status, bytes, outcome.content_type.as_deref())
745}
746
747#[derive(Clone, serde::Serialize, serde::Deserialize)]
748struct ToolOutcome {
749    status: u16,
750    body_hex: String,
751    /// Round-32 F2: MCP client requires the upstream tool response's
752    /// `Content-Type` to round-trip on cached-outcome replay too, so
753    /// strict JSON-RPC 2.0 clients see `application/json` on replay
754    /// just as they did on the fresh forward. `#[serde(default)]` so
755    /// journals persisted before the field existed decode as `None`
756    /// (which the response builder will map to the JSON default).
757    #[serde(default)]
758    content_type: Option<String>,
759}
760
761impl ToolOutcome {
762    fn into_response(self) -> Response {
763        let status = StatusCode::from_u16(self.status).unwrap_or(StatusCode::BAD_GATEWAY);
764        match hex::decode(self.body_hex) {
765            Ok(body) => tool_response(status, Bytes::from(body), self.content_type.as_deref()),
766            Err(error) => lifecycle_error(format!("decode cached tool response: {error}")),
767        }
768    }
769}
770
771enum ToolExecutionState {
772    Missing,
773    Pending,
774    Unaudited(ToolOutcome),
775    Completed(ToolOutcome),
776}
777
778const TOOL_REQUEST_MISMATCH: &str = "JSON-RPC id is already bound to a different tool request or principal";
779/// Canonical duplicate-execution refusal, shared by the pre-flight Pending
780/// state and a lost concurrent claim race so neither reveals more than the
781/// other (a raw claim error would leak filesystem detail, CWE-209).
782const TOOL_OUTCOME_UNCERTAIN: &str = "tool execution outcome is uncertain; refusing duplicate execution";
783
784#[derive(PartialEq, serde::Serialize, serde::Deserialize)]
785struct ToolIntent {
786    pub(crate) execution_key: String,
787    pub(crate) session_id: String,
788    pub(crate) tool: String,
789    pub(crate) request_digest: String,
790    pub(crate) principal_digest: String,
791}
792
793#[cfg(test)]
794pub(crate) async fn ensure_no_unresolved_tool_executions(
795    spool: &std::path::Path,
796    control_key: &[u8; 32],
797) -> Result<(), String> {
798    let sessions = unresolved_tool_sessions(spool, control_key).await?;
799    if let Some(session_id) = sessions.into_iter().next() {
800        Err(format!("session {session_id} has an unresolved tool execution"))
801    } else {
802        Ok(())
803    }
804}
805
806pub(crate) async fn unresolved_tool_sessions(
807    spool: &std::path::Path,
808    control_key: &[u8; 32],
809) -> Result<std::collections::HashSet<String>, String> {
810    let directory = spool.join(crate::spool::TOOL_EXECUTIONS);
811    let control_key = *control_key;
812    tokio::task::spawn_blocking(move || {
813        let entries = match std::fs::read_dir(&directory) {
814            Ok(entries) => entries,
815            Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
816                return Ok(std::collections::HashSet::new());
817            }
818            Err(error) => return Err(error.to_string()),
819        };
820        let mut intent_keys = std::collections::HashSet::new();
821        let mut unresolved_sessions = std::collections::HashSet::new();
822        for entry in entries {
823            let path = entry.map_err(|error| error.to_string())?.path();
824            let Some(name) = path.file_name().and_then(std::ffi::OsStr::to_str) else {
825                continue;
826            };
827            let Some(key) = name.strip_suffix(crate::spool::TOOL_INTENT_SUFFIX) else {
828                continue;
829            };
830            intent_keys.insert(key.to_owned());
831            // A torn intent (crash between create_new and sync_all in
832            // claim_sync) leaves a 0-byte or partial file at the final
833            // intent_path. If we bubble the journal::open error via `?`
834            // the recovery loop aborts on the first torn intent and
835            // every subsequent tool session stays unrecovered
836            // indefinitely. Quarantine the file so the reconciler makes
837            // progress on the rest of the spool; the tool execution
838            // itself is still uncertain — a subsequent retry by the
839            // client will get TOOL_OUTCOME_UNCERTAIN and can decide.
840            let intent_bytes = std::fs::read(&path).map_err(|error| error.to_string())?;
841            let intent: ToolIntent = match crate::journal::open(
842                &control_key,
843                &format!("{}:{key}", crate::journal::TOOL_INTENT_DOMAIN),
844                0,
845                &intent_bytes,
846            ) {
847                Ok(intent) => intent,
848                Err(error) => {
849                    let quarantine = path.with_extension("intent.torn");
850                    tracing::warn!(
851                        %error,
852                        original = %av_core::fsutil::basename(&path),
853                        quarantine = %av_core::fsutil::basename(&quarantine),
854                        "torn tool-execution intent quarantined so recovery can proceed"
855                    );
856                    if let Err(rename_err) = std::fs::rename(&path, &quarantine) {
857                        tracing::warn!(%rename_err, "failed to quarantine torn intent — leaving in place");
858                    }
859                    intent_keys.remove(key);
860                    continue;
861                }
862            };
863            if intent.execution_key != key {
864                return Err("tool intent path does not match authenticated execution key".to_owned());
865            }
866            let outcome_path = directory.join(format!("{key}{}", crate::spool::TOOL_OUTCOME_SUFFIX));
867            if !outcome_path.exists() {
868                unresolved_sessions.insert(intent.session_id);
869                continue;
870            }
871            let _: ToolOutcome = crate::journal::open(
872                &control_key,
873                &format!("{}:{key}", crate::journal::TOOL_OUTCOME_DOMAIN),
874                0,
875                &std::fs::read(&outcome_path).map_err(|error| error.to_string())?,
876            )?;
877            let audited_path = directory.join(format!("{key}{}", crate::spool::TOOL_AUDITED_SUFFIX));
878            if !audited_path.exists() {
879                unresolved_sessions.insert(intent.session_id);
880                continue;
881            }
882            let _: serde_json::Value = crate::journal::open(
883                &control_key,
884                &format!("{}:{key}", crate::journal::TOOL_AUDITED_DOMAIN),
885                0,
886                &std::fs::read(&audited_path).map_err(|error| error.to_string())?,
887            )?;
888        }
889        for entry in std::fs::read_dir(&directory).map_err(|error| error.to_string())? {
890            let path = entry.map_err(|error| error.to_string())?.path();
891            let Some(name) = path.file_name().and_then(std::ffi::OsStr::to_str) else {
892                continue;
893            };
894            if let Some(key) = name.strip_suffix(crate::spool::TOOL_OUTCOME_SUFFIX) {
895                if !intent_keys.contains(key) {
896                    return Err("tool outcome exists without an authenticated intent".to_owned());
897                }
898            }
899        }
900        Ok(unresolved_sessions)
901    })
902    .await
903    .map_err(|error| error.to_string())?
904}
905
906#[derive(Clone)]
907struct ToolExecution {
908    key: String,
909    session_id: String,
910    tool: String,
911    request_digest: String,
912    principal_digest: String,
913    control_key: [u8; 32],
914    intent_path: std::path::PathBuf,
915    outcome_path: std::path::PathBuf,
916    audited_path: std::path::PathBuf,
917}
918
919impl ToolExecution {
920    fn from_request(
921        spool: &str,
922        headers: &HeaderMap,
923        body: &[u8],
924        control_key: [u8; 32],
925    ) -> Result<Self, crate::pipeline::PipelineError> {
926        // Refuse duplicate x-av-session headers (see
927        // `pipeline::single_header` for the rationale — proxies
928        // sometimes merge duplicates on the wire and log aggregators
929        // can then observe a comma-joined value that leaks a
930        // client-desync into audit).
931        let session_id = crate::pipeline::single_header(headers, crate::pipeline::SESSION_HEADER)?
932            .and_then(|value| value.to_str().ok())
933            .ok_or_else(|| crate::pipeline::PipelineError::BadRequest("missing x-av-session".to_owned()))?;
934        // Same validation as the pipeline's `session_id`: an id the intercept
935        // path would reject must not key a tool-execution intent either.
936        let session_id = av_core::SessionId::parse(session_id)
937            .map(|id| id.to_string())
938            .map_err(|error| crate::pipeline::PipelineError::BadRequest(error.to_string()))?;
939        let call = av_sandbox::parse_tool_call(body)
940            .map_err(|error| crate::pipeline::PipelineError::BadRequest(error.to_string()))?;
941        let id = call.id.ok_or_else(|| {
942            crate::pipeline::PipelineError::BadRequest(
943                "forwarded tool calls require a JSON-RPC id for idempotency".to_owned(),
944            )
945        })?;
946        let request: Value = serde_json::from_slice(body)
947            .map_err(|error| crate::pipeline::PipelineError::BadRequest(error.to_string()))?;
948        let canonical = av_receipts::canonicalize(&request)
949            .map_err(|error| crate::pipeline::PipelineError::BadRequest(error.to_string()))?;
950        let request_digest = av_core::digest::sha256_hex(canonical.as_bytes());
951        let key_material = format!("{session_id}:{}", id);
952        let key = av_core::digest::sha256_hex(key_material.as_bytes());
953        let directory = std::path::Path::new(spool).join(crate::spool::TOOL_EXECUTIONS);
954        Ok(Self {
955            intent_path: directory.join(format!("{key}{}", crate::spool::TOOL_INTENT_SUFFIX)),
956            outcome_path: directory.join(format!("{key}{}", crate::spool::TOOL_OUTCOME_SUFFIX)),
957            audited_path: directory.join(format!("{key}{}", crate::spool::TOOL_AUDITED_SUFFIX)),
958            key,
959            session_id,
960            tool: call.tool,
961            request_digest,
962            principal_digest: String::new(),
963            control_key,
964        })
965    }
966
967    fn bind_principal(
968        &mut self,
969        identity: &av_events::AgentIdentity,
970    ) -> Result<(), crate::pipeline::PipelineError> {
971        let stable_identity = json!({
972            "version": identity.version,
973            "charter": identity.charter,
974            "instance_uid": identity.instance_uid,
975        });
976        let canonical = av_receipts::canonicalize(&stable_identity)
977            .map_err(|error| crate::pipeline::PipelineError::BadRequest(error.to_string()))?;
978        self.principal_digest = av_core::digest::sha256_hex(canonical.as_bytes());
979        Ok(())
980    }
981
982    async fn load(&self) -> Result<ToolExecutionState, String> {
983        let execution = self.clone();
984        tokio::task::spawn_blocking(move || execution.load_sync())
985            .await
986            .map_err(|error| error.to_string())?
987    }
988
989    fn load_sync(&self) -> Result<ToolExecutionState, String> {
990        if !self.intent_path.exists() {
991            if self.outcome_path.exists() || self.audited_path.exists() {
992                return Err("tool outcome exists without an authenticated intent".to_owned());
993            }
994            return Ok(ToolExecutionState::Missing);
995        }
996        let intent: ToolIntent = crate::journal::open(
997            &self.control_key,
998            &format!("{}:{}", crate::journal::TOOL_INTENT_DOMAIN, self.key),
999            0,
1000            &std::fs::read(&self.intent_path).map_err(|error| error.to_string())?,
1001        )?;
1002        if intent != self.intent() {
1003            return Err(TOOL_REQUEST_MISMATCH.to_owned());
1004        }
1005        if self.outcome_path.exists() {
1006            let outcome: ToolOutcome = crate::journal::open(
1007                &self.control_key,
1008                &format!("{}:{}", crate::journal::TOOL_OUTCOME_DOMAIN, self.key),
1009                0,
1010                &std::fs::read(&self.outcome_path).map_err(|error| error.to_string())?,
1011            )?;
1012            return if self.audited_path.exists()
1013                && crate::journal::open::<serde_json::Value>(
1014                    &self.control_key,
1015                    &format!("{}:{}", crate::journal::TOOL_AUDITED_DOMAIN, self.key),
1016                    0,
1017                    &std::fs::read(&self.audited_path).map_err(|error| error.to_string())?,
1018                )
1019                .is_ok()
1020            {
1021                Ok(ToolExecutionState::Completed(outcome))
1022            } else {
1023                Ok(ToolExecutionState::Unaudited(outcome))
1024            };
1025        }
1026        Ok(ToolExecutionState::Pending)
1027    }
1028
1029    async fn claim(&self) -> Result<(), String> {
1030        let execution = self.clone();
1031        tokio::task::spawn_blocking(move || execution.claim_sync())
1032            .await
1033            .map_err(|error| error.to_string())?
1034    }
1035
1036    fn claim_sync(&self) -> Result<(), String> {
1037        use std::io::Write as _;
1038        let directory = self
1039            .intent_path
1040            .parent()
1041            .ok_or_else(|| "tool execution directory is missing".to_owned())?;
1042        std::fs::create_dir_all(directory).map_err(|error| error.to_string())?;
1043        let mut file = std::fs::OpenOptions::new()
1044            .write(true)
1045            .create_new(true)
1046            .open(&self.intent_path)
1047            .map_err(|error| format!("tool execution already claimed or unavailable: {error}"))?;
1048        let intent = crate::journal::seal(
1049            &self.control_key,
1050            &format!("{}:{}", crate::journal::TOOL_INTENT_DOMAIN, self.key),
1051            0,
1052            &self.intent(),
1053        )?;
1054        file.write_all(&intent).map_err(|error| error.to_string())?;
1055        file.sync_all().map_err(|error| error.to_string())?;
1056        std::fs::File::open(directory)
1057            .and_then(|directory| directory.sync_all())
1058            .map_err(|error| error.to_string())
1059    }
1060
1061    async fn persist(&self, outcome: &ToolOutcome) -> Result<(), String> {
1062        let execution = self.clone();
1063        let outcome = outcome.clone();
1064        tokio::task::spawn_blocking(move || {
1065            let sealed = crate::journal::seal(
1066                &execution.control_key,
1067                &format!("{}:{}", crate::journal::TOOL_OUTCOME_DOMAIN, execution.key),
1068                0,
1069                &outcome,
1070            )?;
1071            write_atomic_bytes(&execution.outcome_path, &sealed)
1072        })
1073        .await
1074        .map_err(|error| error.to_string())?
1075    }
1076
1077    async fn mark_audited(&self) -> Result<(), String> {
1078        let execution = self.clone();
1079        tokio::task::spawn_blocking(move || {
1080            let sealed = crate::journal::seal(
1081                &execution.control_key,
1082                &format!("{}:{}", crate::journal::TOOL_AUDITED_DOMAIN, execution.key),
1083                0,
1084                &json!({"audited": true}),
1085            )?;
1086            write_atomic_bytes(&execution.audited_path, &sealed)
1087        })
1088        .await
1089        .map_err(|error| error.to_string())?
1090    }
1091
1092    fn intent(&self) -> ToolIntent {
1093        ToolIntent {
1094            execution_key: self.key.clone(),
1095            session_id: self.session_id.clone(),
1096            tool: self.tool.clone(),
1097            request_digest: self.request_digest.clone(),
1098            principal_digest: self.principal_digest.clone(),
1099        }
1100    }
1101}
1102
1103fn write_atomic_bytes(path: &std::path::Path, bytes: &[u8]) -> Result<(), String> {
1104    av_core::fsutil::write_atomic(path, bytes).map_err(|error| error.to_string())
1105}
1106
1107async fn close_session(
1108    State(state): State<AppState>,
1109    Path(id): Path<String>,
1110    headers: HeaderMap,
1111) -> Response {
1112    let Some(session) = state.sessions.get(&id) else {
1113        return (StatusCode::NOT_FOUND, Json(json!({"error": "unknown session"}))).into_response();
1114    };
1115    if let Err(error) = state.authorize_session(&headers, &session, &state.config.session_close_scope) {
1116        return pipeline_error(error);
1117    }
1118    match state
1119        .finalizer
1120        .close_session(session, StopReason::SessionClosed)
1121        .await
1122    {
1123        Ok(outcome) => Json(outcome).into_response(),
1124        Err(error) => lifecycle_error(error.to_string()),
1125    }
1126}
1127
1128async fn promote_session(
1129    State(state): State<AppState>,
1130    Path(id): Path<String>,
1131    headers: HeaderMap,
1132) -> Response {
1133    let Some(session) = state.sessions.get(&id) else {
1134        return (StatusCode::NOT_FOUND, Json(json!({"error": "unknown session"}))).into_response();
1135    };
1136    if let Err(error) = state.authorize_session(&headers, &session, &state.config.session_promote_scope) {
1137        return pipeline_error(error);
1138    }
1139    // Round-27 F4: `promote()` silently drives `close_session_locked`
1140    // on any still-open session. If an operator split the two scopes
1141    // (compliance auditor gets `session:promote`, on-call gets
1142    // `session:close`), a `session:promote`-only bearer could
1143    // otherwise force-close any live agent session by promoting it —
1144    // bypassing `session:close` entirely. When the session is still
1145    // open, additionally require the close scope so promote ⊇ close
1146    // in the scope authority sense.
1147    if !session.is_closed() {
1148        if let Err(error) = state.authorize_session(&headers, &session, &state.config.session_close_scope) {
1149            return pipeline_error(error);
1150        }
1151    }
1152    match state.finalizer.promote(session).await {
1153        Ok(receipt) => Json(receipt).into_response(),
1154        Err(error) => lifecycle_error(error.to_string()),
1155    }
1156}
1157
1158fn pipeline_error(error: crate::pipeline::PipelineError) -> Response {
1159    use crate::pipeline::PipelineError;
1160    let close = matches!(error, PipelineError::Abort(_));
1161    let status = error.status();
1162    let mut response = (status, Json(json!({"error": error.to_string()}))).into_response();
1163    if close {
1164        response
1165            .headers_mut()
1166            .insert(axum::http::header::CONNECTION, HeaderValue::from_static("close"));
1167    }
1168    // Drive advisory-header attachment off the numeric HTTP status
1169    // rather than variant identity. `PipelineError` is
1170    // `#[non_exhaustive]`, so a future variant (e.g. `Timeout` → 504,
1171    // `RateLimited` → 429) would silently skip Retry-After under a
1172    // `matches!(error, PipelineError::Unavailable(_))` chain — exactly
1173    // the SDK-hammering regression this arm exists to prevent. Using
1174    // the status code centralises the semantic once and covers every
1175    // present and future variant that maps to the same class.
1176    match status.as_u16() {
1177        // 429 / 502 / 503 / 504 — retryable per RFC 7231 §7.1.3; the
1178        // header value is deliberately short so the audit-capture
1179        // recovery window is fast, but long enough that clients do not
1180        // hammer during transient failures.
1181        429 | 502 | 503 | 504 => {
1182            response
1183                .headers_mut()
1184                .insert(axum::http::header::RETRY_AFTER, HeaderValue::from_static("5"));
1185        }
1186        // 401 — RFC 7235 §3.1 MUST send `WWW-Authenticate`; RFC 6750 §3
1187        // defines the Bearer challenge shape used with NHI JWTs.
1188        401 => {
1189            response.headers_mut().insert(
1190                axum::http::header::WWW_AUTHENTICATE,
1191                HeaderValue::from_static("Bearer realm=\"agentvisor\", error=\"invalid_token\""),
1192            );
1193        }
1194        _ => {}
1195    }
1196    response
1197}
1198
1199fn lifecycle_error(error: String) -> Response {
1200    (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({"error": error}))).into_response()
1201}
1202
1203struct AbortFinalizingStream {
1204    inner: BoxStream<'static, Result<Bytes, std::io::Error>>,
1205    session: Arc<crate::session::Session>,
1206    identity: av_events::AgentIdentity,
1207    response_permit: Option<crate::worker::ResponsePermit>,
1208    /// Handle to submit the response-capture job at end-of-stream. The
1209    /// permit above only holds the `response_capacity` semaphore slot
1210    /// — the mpsc queue slot is re-acquired at submit time via
1211    /// `ResponsePermit::submit(&worker, job)`.
1212    worker: crate::worker::WorkerHandle,
1213    store: Arc<dyn av_state::StateStore>,
1214    budget: av_state::BudgetSpec,
1215    finalizer: crate::reconciler::Finalizer,
1216    _lease: crate::session::SessionLease,
1217    response_marker: Option<String>,
1218    response_attempt_id: String,
1219    response_message: String,
1220    response_reasoning: String,
1221    response_model: Option<String>,
1222    response_finish_reason: Option<String>,
1223    upstream_status: StatusCode,
1224    response_cost_usd_micros: u64,
1225    response_tool_calls: std::collections::BTreeMap<u64, PartialToolCall>,
1226    response_metrics: av_events::EventMetrics,
1227    charged_completion_tokens: u64,
1228    last_reported_completion_tokens: Option<u64>,
1229    last_reported_prompt_tokens: Option<u64>,
1230    last_reported_cached_tokens: Option<u64>,
1231    last_reported_cost_usd_micros: Option<u64>,
1232    saw_chunk: bool,
1233    capture_submitted: bool,
1234    is_sse: bool,
1235    protocol_buffer: Vec<u8>,
1236    pending_output: std::collections::VecDeque<Bytes>,
1237    pending_budget: Option<PendingBudget>,
1238    captured_bytes: usize,
1239    completed: bool,
1240}
1241
1242const MAX_PROVIDER_CAPTURE_BYTES: usize = 16 * 1024 * 1024;
1243const MAX_PROVIDER_FIELD_BYTES: usize = 8 * 1024 * 1024;
1244const MAX_PROVIDER_TOOL_CALLS: usize = 128;
1245
1246impl AbortFinalizingStream {
1247    fn absorb_network_chunk(&mut self, raw: &[u8]) -> Result<u64, String> {
1248        self.saw_chunk |= !raw.is_empty();
1249        self.captured_bytes = self
1250            .captured_bytes
1251            .checked_add(raw.len())
1252            .ok_or_else(|| "provider response size overflow".to_owned())?;
1253        if self.captured_bytes > MAX_PROVIDER_CAPTURE_BYTES {
1254            return Err(format!(
1255                "provider response exceeds {MAX_PROVIDER_CAPTURE_BYTES} capture bytes"
1256            ));
1257        }
1258        self.protocol_buffer.extend_from_slice(raw);
1259        if !self.is_sse {
1260            return Ok(0);
1261        }
1262        // Round-29 F1: non-success upstream bodies are relayed
1263        // verbatim; do NOT try to parse them as chat-completion SSE
1264        // frames. An `event: error` / `data: {"error":...}` frame
1265        // from a 4xx/5xx stream would otherwise fail the strict
1266        // parser and collapse the true status into a 502. Drop the
1267        // buffered bytes on the floor (the wire body has already
1268        // been relayed to the client via `pending_output`) and
1269        // let flush_protocol_buffer's own guard finalise.
1270        if !self.upstream_status.is_success() {
1271            self.protocol_buffer.clear();
1272            return Ok(0);
1273        }
1274        let mut budget_delta = 0u64;
1275        while let Some(end) = sse_frame_end(&self.protocol_buffer) {
1276            let frame: Vec<u8> = self.protocol_buffer.drain(..end).collect();
1277            let frame = std::str::from_utf8(&frame)
1278                .map_err(|error| format!("provider SSE frame is not UTF-8: {error}"))?;
1279            budget_delta = budget_delta
1280                .checked_add(self.absorb_frame(frame)?)
1281                .ok_or_else(|| "provider completion-token delta overflow".to_owned())?;
1282        }
1283        if self.protocol_buffer.len() > MAX_PROVIDER_FIELD_BYTES {
1284            return Err("unterminated provider SSE frame exceeds limit".to_owned());
1285        }
1286        Ok(budget_delta)
1287    }
1288
1289    fn flush_protocol_buffer(&mut self) -> Result<u64, String> {
1290        if self.protocol_buffer.is_empty() {
1291            return Ok(0);
1292        }
1293        // Round-29 F1: never fail-closed on a non-success upstream body.
1294        // Providers ship text/plain and HTML error pages on 4xx/5xx
1295        // (OpenAI's Cloudflare frontend returns text/html on 429;
1296        // Anthropic ships 503 HTML from AWS ALBs during backend
1297        // restarts). The strict JSON parse below would fail and be
1298        // mapped to a fresh 502, silently dropping the real status +
1299        // Retry-After header. SDKs treat "502 without Retry-After" as
1300        // an immediate retry candidate, so a rate-limited upstream
1301        // gets hammered instead of backed off. Skip the parse for
1302        // non-success responses; the buffered body still relays
1303        // through the buffered non-SSE branch and the true
1304        // upstream_status is preserved into the response later.
1305        if !self.upstream_status.is_success() {
1306            let _ = std::mem::take(&mut self.protocol_buffer);
1307            return Ok(0);
1308        }
1309        let frame = std::mem::take(&mut self.protocol_buffer);
1310        let frame = std::str::from_utf8(&frame)
1311            .map_err(|error| format!("provider response is not UTF-8: {error}"))?;
1312        self.absorb_frame(frame)
1313    }
1314
1315    fn absorb_frame(&mut self, raw: &str) -> Result<u64, String> {
1316        let Some(parsed) = parse_provider_chunk(raw)? else {
1317            return Ok(0);
1318        };
1319        if self.upstream_status.is_success() && !parsed.has_choices {
1320            return Err("successful provider response has no choices array".to_owned());
1321        }
1322        push_bounded(&mut self.response_message, &parsed.message, "response message")?;
1323        if let Some(reasoning) = parsed.reasoning {
1324            push_bounded(&mut self.response_reasoning, &reasoning, "response reasoning")?;
1325        }
1326        if self.response_model.is_none() {
1327            self.response_model = parsed.model_name;
1328        }
1329        if parsed.finish_reason.is_some() {
1330            self.response_finish_reason = parsed.finish_reason;
1331        }
1332        self.response_cost_usd_micros = self.response_cost_usd_micros.max(parsed.cost_usd_micros);
1333        for delta in parsed.tool_call_deltas {
1334            if delta.index >= MAX_PROVIDER_TOOL_CALLS as u64 {
1335                return Err(format!(
1336                    "provider tool-call index {} is out of range",
1337                    delta.index
1338                ));
1339            }
1340            if !self.response_tool_calls.contains_key(&delta.index)
1341                && self.response_tool_calls.len() >= MAX_PROVIDER_TOOL_CALLS
1342            {
1343                return Err(format!(
1344                    "provider response exceeds {MAX_PROVIDER_TOOL_CALLS} tool calls"
1345                ));
1346            }
1347            let partial = self.response_tool_calls.entry(delta.index).or_default();
1348            if delta.id.is_some() {
1349                partial.id = delta.id;
1350            }
1351            if delta.name.is_some() {
1352                partial.name = delta.name;
1353            }
1354            push_bounded(&mut partial.arguments, &delta.arguments, "tool-call arguments")?;
1355        }
1356        if parsed.usage_reported {
1357            reject_metric_regression(
1358                "prompt tokens",
1359                &mut self.last_reported_prompt_tokens,
1360                parsed.metrics.prompt_tokens,
1361            )?;
1362            reject_metric_regression(
1363                "cached tokens",
1364                &mut self.last_reported_cached_tokens,
1365                parsed.metrics.cached_tokens,
1366            )?;
1367        }
1368        if parsed.cost_reported {
1369            reject_metric_regression(
1370                "cost",
1371                &mut self.last_reported_cost_usd_micros,
1372                Some(parsed.cost_usd_micros),
1373            )?;
1374        }
1375        if let Some(prompt) = parsed.metrics.prompt_tokens {
1376            self.response_metrics.prompt_tokens = Some(
1377                self.response_metrics
1378                    .prompt_tokens
1379                    .map_or(prompt, |current| current.max(prompt)),
1380            );
1381        }
1382        if let Some(cached) = parsed.metrics.cached_tokens {
1383            self.response_metrics.cached_tokens = Some(
1384                self.response_metrics
1385                    .cached_tokens
1386                    .map_or(cached, |current| current.max(cached)),
1387            );
1388        }
1389        let reported_completion = parsed.metrics.completion_tokens.unwrap_or(0);
1390        let delta = if parsed.usage_reported {
1391            let previous = self.last_reported_completion_tokens.unwrap_or(0);
1392            if reported_completion < previous {
1393                return Err("provider completion usage regressed".to_owned());
1394            }
1395            self.last_reported_completion_tokens = Some(reported_completion);
1396            let delta = reported_completion.saturating_sub(self.charged_completion_tokens);
1397            self.response_metrics.completion_tokens = Some(
1398                self.response_metrics
1399                    .completion_tokens
1400                    .map_or(reported_completion, |current| current.max(reported_completion)),
1401            );
1402            delta
1403        } else {
1404            let total = self
1405                .response_metrics
1406                .completion_tokens
1407                .unwrap_or(0)
1408                .checked_add(reported_completion)
1409                .filter(|total| *total <= av_core::error::JCS_SAFE_MAX)
1410                .ok_or_else(|| "provider completion-token total exceeds JCS-safe bounds".to_owned())?;
1411            self.response_metrics.completion_tokens = Some(total);
1412            reported_completion
1413        };
1414        if delta == 0 {
1415            return Ok(0);
1416        }
1417        self.charged_completion_tokens = self
1418            .charged_completion_tokens
1419            .checked_add(delta)
1420            .ok_or_else(|| "charged completion-token counter overflow".to_owned())?;
1421        Ok(delta)
1422    }
1423
1424    fn begin_budget_check(&mut self, delta: u64, continuation: BudgetContinuation) {
1425        let store = Arc::clone(&self.store);
1426        let session_id = self.session.id.clone();
1427        let budget = self.budget.clone();
1428        let task = tokio::task::spawn_blocking(move || {
1429            av_state::ActionBudget::new(store.as_ref(), &session_id, &budget)
1430                .try_tokens(delta)
1431                .map_err(|error| format!("token budget backend failed closed: {error}"))
1432        });
1433        self.pending_budget = Some(PendingBudget { task, continuation });
1434    }
1435
1436    fn submit_response_capture(&mut self, failure: Option<String>) -> Result<(), crate::worker::SubmitError> {
1437        if self.capture_submitted {
1438            return Ok(());
1439        }
1440        self.capture_submitted = true;
1441        let reasoning =
1442            (!self.response_reasoning.is_empty()).then(|| std::mem::take(&mut self.response_reasoning));
1443        let analysis_text = reasoning.clone().unwrap_or_else(|| self.response_message.clone());
1444        let native_finish_reason = self.response_finish_reason.clone();
1445        let (class, status, stop_reason, payload) = if let Some(reason) = failure {
1446            (
1447                av_events::EventClass::StopReason,
1448                av_events::StatusId::Failure,
1449                Some(av_events::StopReason::BudgetExceeded),
1450                json!({"reason": reason, "direction": "upstream_response"}),
1451            )
1452        } else if !self.upstream_status.is_success() {
1453            (
1454                av_events::EventClass::StopReason,
1455                av_events::StatusId::Failure,
1456                Some(av_events::StopReason::Other),
1457                json!({
1458                    "direction": "upstream_response",
1459                    "http_status": self.upstream_status.as_u16(),
1460                }),
1461            )
1462        } else if let Some(native) = &native_finish_reason {
1463            (
1464                av_events::EventClass::StopReason,
1465                av_events::StatusId::Success,
1466                Some(map_finish_reason(native)),
1467                json!({
1468                    "direction": "upstream_response",
1469                    "finish_reason": native,
1470                    "http_status": self.upstream_status.as_u16(),
1471                }),
1472            )
1473        } else {
1474            (
1475                av_events::EventClass::Session,
1476                av_events::StatusId::Success,
1477                None,
1478                json!({
1479                    "direction": "upstream_response",
1480                    "http_status": self.upstream_status.as_u16(),
1481                }),
1482            )
1483        };
1484        let tool_calls: Vec<av_atif::ToolCall> = std::mem::take(&mut self.response_tool_calls)
1485            .into_values()
1486            .map(|partial| {
1487                let arguments = serde_json::from_str(&partial.arguments)
1488                    .unwrap_or_else(|_| json!({"raw": partial.arguments}));
1489                av_atif::ToolCall {
1490                    tool_call_id: partial.id.unwrap_or_else(av_core::new_event_uid),
1491                    function_name: partial.name.unwrap_or_else(|| "unknown".to_owned()),
1492                    arguments,
1493                    extra: None,
1494                }
1495            })
1496            .collect();
1497        let permit = self
1498            .response_permit
1499            .take()
1500            .ok_or(crate::worker::SubmitError::Closed)?;
1501        permit.submit(
1502            &self.worker,
1503            crate::worker::WorkerJob {
1504                session: Arc::clone(&self.session),
1505                identity: self.identity.clone(),
1506                class,
1507                payload,
1508                text: analysis_text,
1509                analyze_loop: true,
1510                status,
1511                stop_reason,
1512                native_stop_reason: native_finish_reason,
1513                metrics: self.response_metrics,
1514                cost_usd_micros: self.response_cost_usd_micros,
1515                atif: Some(crate::worker::AtifCapture {
1516                    source: av_atif::Source::Agent,
1517                    message: Value::String(std::mem::take(&mut self.response_message)),
1518                    reasoning_content: reasoning,
1519                    model_name: self.response_model.take(),
1520                    tool_calls: (!tool_calls.is_empty()).then_some(tool_calls),
1521                    observation: None,
1522                    llm_call_count: Some(1),
1523                }),
1524                response_marker: self.response_marker.take(),
1525                response_attempt: Some(crate::worker::ResponseAttempt {
1526                    id: self.response_attempt_id.clone(),
1527                    terminal: true,
1528                }),
1529            },
1530        )?;
1531        Ok(())
1532    }
1533}
1534
1535enum BudgetContinuation {
1536    Emit(Bytes),
1537    ContinueNonSse,
1538    FinishSse,
1539    FinishNonSse,
1540}
1541
1542struct PendingBudget {
1543    task: tokio::task::JoinHandle<Result<av_state::BudgetDecision, String>>,
1544    continuation: BudgetContinuation,
1545}
1546
1547#[derive(Default)]
1548struct PartialToolCall {
1549    id: Option<String>,
1550    name: Option<String>,
1551    arguments: String,
1552}
1553
1554struct ProviderToolCallDelta {
1555    index: u64,
1556    id: Option<String>,
1557    name: Option<String>,
1558    arguments: String,
1559}
1560
1561struct ParsedProviderChunk {
1562    message: String,
1563    reasoning: Option<String>,
1564    model_name: Option<String>,
1565    metrics: av_events::EventMetrics,
1566    usage_reported: bool,
1567    finish_reason: Option<String>,
1568    cost_usd_micros: u64,
1569    cost_reported: bool,
1570    has_choices: bool,
1571    tool_call_deltas: Vec<ProviderToolCallDelta>,
1572}
1573
1574fn reject_metric_regression(
1575    field: &str,
1576    previous: &mut Option<u64>,
1577    current: Option<u64>,
1578) -> Result<(), String> {
1579    let Some(current) = current else {
1580        return Ok(());
1581    };
1582    if previous.is_some_and(|previous| current < previous) {
1583        return Err(format!("provider {field} regressed"));
1584    }
1585    *previous = Some(current);
1586    Ok(())
1587}
1588
1589fn push_bounded(target: &mut String, value: &str, field: &str) -> Result<(), String> {
1590    let size = target
1591        .len()
1592        .checked_add(value.len())
1593        .ok_or_else(|| format!("{field} size overflow"))?;
1594    if size > MAX_PROVIDER_FIELD_BYTES {
1595        return Err(format!("{field} exceeds {MAX_PROVIDER_FIELD_BYTES} bytes"));
1596    }
1597    target.push_str(value);
1598    Ok(())
1599}
1600
1601fn sse_frame_end(buffer: &[u8]) -> Option<usize> {
1602    let mut line_start = 0usize;
1603    let mut index = 0usize;
1604    while index < buffer.len() {
1605        let byte = buffer.get(index).copied()?;
1606        let newline = match byte {
1607            b'\r' if buffer.get(index + 1) == Some(&b'\n') => 2,
1608            b'\r' if buffer.get(index + 1).is_none() => return None,
1609            b'\r' | b'\n' => 1,
1610            _ => {
1611                index += 1;
1612                continue;
1613            }
1614        };
1615        if index == line_start {
1616            return Some(index + newline);
1617        }
1618        index += newline;
1619        line_start = index;
1620    }
1621    None
1622}
1623
1624impl AbortFinalizingStream {
1625    /// Fail-closed relay abort: the client connection is severed (the
1626    /// status line already went out, so a clean error response is no
1627    /// longer possible). Without this log the reason would vanish into
1628    /// an io::Error that hyper discards — leaving "empty reply from
1629    /// server" as the only symptom of e.g. an upstream returning HTML.
1630    fn abort_error(&self, reason: String) -> std::io::Error {
1631        tracing::warn!(
1632            session = %self.session.id,
1633            %reason,
1634            "aborting client response; upstream reply could not be captured"
1635        );
1636        std::io::Error::other(reason)
1637    }
1638}
1639
1640impl Stream for AbortFinalizingStream {
1641    type Item = Result<Bytes, std::io::Error>;
1642
1643    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1644        if self.pending_budget.is_some() {
1645            let result = {
1646                let Some(pending) = self.pending_budget.as_mut() else {
1647                    self.session.mark_capture_failed();
1648                    return Poll::Ready(Some(Err(std::io::Error::other(
1649                        "pending completion budget state disappeared",
1650                    ))));
1651                };
1652                match Pin::new(&mut pending.task).poll(context) {
1653                    Poll::Pending => return Poll::Pending,
1654                    Poll::Ready(result) => result,
1655                }
1656            };
1657            let Some(pending) = self.pending_budget.take() else {
1658                self.session.mark_capture_failed();
1659                return Poll::Ready(Some(Err(std::io::Error::other(
1660                    "completed completion budget state disappeared",
1661                ))));
1662            };
1663            let failure = match result {
1664                Ok(Ok(av_state::BudgetDecision::Allowed { .. })) => None,
1665                Ok(Ok(av_state::BudgetDecision::Refused { limit, cap })) => {
1666                    Some(format!("{limit} exceeded (cap {cap})"))
1667                }
1668                Ok(Err(reason)) => Some(reason),
1669                Err(error) => {
1670                    self.session.mark_capture_failed();
1671                    self.pending_output.clear();
1672                    return Poll::Ready(Some(Err(std::io::Error::other(format!(
1673                        "completion budget task failed: {error}"
1674                    )))));
1675                }
1676            };
1677            if let Some(reason) = failure {
1678                if let Err(error) = self.submit_response_capture(Some(reason.clone())) {
1679                    self.session.mark_capture_failed();
1680                    self.pending_output.clear();
1681                    return Poll::Ready(Some(Err(std::io::Error::other(format!(
1682                        "response capture failed closed: {error}"
1683                    )))));
1684                }
1685                self.pending_output.clear();
1686                // QuotaExceeded is an in-process marker: the non-SSE drain in
1687                // `chat_completions` maps it to HTTP 429 so budget refusals
1688                // reach the client as a real error instead of a severed body.
1689                return Poll::Ready(Some(Err(std::io::Error::new(
1690                    std::io::ErrorKind::QuotaExceeded,
1691                    format!("response blocked by token budget: {reason}"),
1692                ))));
1693            }
1694            match pending.continuation {
1695                BudgetContinuation::Emit(bytes) => return Poll::Ready(Some(Ok(bytes))),
1696                BudgetContinuation::ContinueNonSse => {}
1697                BudgetContinuation::FinishSse => {
1698                    if let Err(error) = self.submit_response_capture(None) {
1699                        self.session.mark_capture_failed();
1700                        return Poll::Ready(Some(Err(std::io::Error::other(format!(
1701                            "response capture failed closed: {error}"
1702                        )))));
1703                    }
1704                    self.completed = true;
1705                    return Poll::Ready(None);
1706                }
1707                BudgetContinuation::FinishNonSse => {
1708                    if let Err(error) = self.submit_response_capture(None) {
1709                        self.session.mark_capture_failed();
1710                        self.pending_output.clear();
1711                        return Poll::Ready(Some(Err(std::io::Error::other(format!(
1712                            "response capture failed closed: {error}"
1713                        )))));
1714                    }
1715                    self.completed = true;
1716                    return Poll::Ready(self.pending_output.pop_front().map(Ok));
1717                }
1718            }
1719        }
1720
1721        if !self.is_sse {
1722            if self.completed {
1723                return Poll::Ready(self.pending_output.pop_front().map(Ok));
1724            }
1725            loop {
1726                match self.inner.as_mut().poll_next(context) {
1727                    Poll::Ready(Some(Ok(bytes))) => {
1728                        let delta = match self.absorb_network_chunk(&bytes) {
1729                            Ok(delta) => delta,
1730                            Err(error) => {
1731                                self.session.mark_capture_failed();
1732                                self.pending_output.clear();
1733                                let error = self.abort_error(error);
1734                                return Poll::Ready(Some(Err(error)));
1735                            }
1736                        };
1737                        self.pending_output.push_back(bytes);
1738                        if delta > 0 {
1739                            self.begin_budget_check(delta, BudgetContinuation::ContinueNonSse);
1740                            context.waker().wake_by_ref();
1741                            return Poll::Pending;
1742                        }
1743                    }
1744                    Poll::Ready(Some(Err(error))) => {
1745                        self.session.mark_capture_failed();
1746                        self.pending_output.clear();
1747                        let error = self.abort_error(error.to_string());
1748                        return Poll::Ready(Some(Err(error)));
1749                    }
1750                    Poll::Pending => return Poll::Pending,
1751                    Poll::Ready(None) => {
1752                        let delta = match self.flush_protocol_buffer() {
1753                            Ok(delta) => delta,
1754                            Err(error) => {
1755                                self.session.mark_capture_failed();
1756                                self.pending_output.clear();
1757                                let error = self.abort_error(error);
1758                                return Poll::Ready(Some(Err(error)));
1759                            }
1760                        };
1761                        if delta > 0 {
1762                            self.begin_budget_check(delta, BudgetContinuation::FinishNonSse);
1763                            context.waker().wake_by_ref();
1764                            return Poll::Pending;
1765                        }
1766                        if let Err(error) = self.submit_response_capture(None) {
1767                            self.session.mark_capture_failed();
1768                            self.pending_output.clear();
1769                            return Poll::Ready(Some(Err(std::io::Error::other(format!(
1770                                "response capture failed closed: {error}"
1771                            )))));
1772                        }
1773                        self.completed = true;
1774                        return Poll::Ready(self.pending_output.pop_front().map(Ok));
1775                    }
1776                }
1777            }
1778        }
1779
1780        match self.inner.as_mut().poll_next(context) {
1781            Poll::Ready(Some(Ok(bytes))) => match self.absorb_network_chunk(&bytes) {
1782                Ok(0) => Poll::Ready(Some(Ok(bytes))),
1783                Ok(delta) => {
1784                    self.begin_budget_check(delta, BudgetContinuation::Emit(bytes));
1785                    context.waker().wake_by_ref();
1786                    Poll::Pending
1787                }
1788                Err(error) => {
1789                    self.session.mark_capture_failed();
1790                    let error = self.abort_error(error);
1791                    Poll::Ready(Some(Err(error)))
1792                }
1793            },
1794            Poll::Ready(Some(Err(error))) => {
1795                self.session.mark_capture_failed();
1796                let error = self.abort_error(error.to_string());
1797                Poll::Ready(Some(Err(error)))
1798            }
1799            Poll::Pending => Poll::Pending,
1800            Poll::Ready(None) => {
1801                let delta = match self.flush_protocol_buffer() {
1802                    Ok(delta) => delta,
1803                    Err(error) => {
1804                        self.session.mark_capture_failed();
1805                        let error = self.abort_error(error);
1806                        return Poll::Ready(Some(Err(error)));
1807                    }
1808                };
1809                if delta > 0 {
1810                    self.begin_budget_check(delta, BudgetContinuation::FinishSse);
1811                    context.waker().wake_by_ref();
1812                    return Poll::Pending;
1813                }
1814                if let Err(error) = self.submit_response_capture(None) {
1815                    self.session.mark_capture_failed();
1816                    return Poll::Ready(Some(Err(std::io::Error::other(format!(
1817                        "response capture failed closed: {error}"
1818                    )))));
1819                }
1820                self.completed = true;
1821                Poll::Ready(None)
1822            }
1823        }
1824    }
1825}
1826
1827fn parse_provider_chunk(raw: &str) -> Result<Option<ParsedProviderChunk>, String> {
1828    // SSE spec (§9.2.4) requires the leading U+FEFF (BOM) to be
1829    // discarded once, at the start of the stream. Rust's `.trim()` does
1830    // not strip U+FEFF (char::is_whitespace() returns false for it), so
1831    // a provider that ships a BOM would otherwise cause every following
1832    // parse to fail on either `strip_prefix("data:")` (BOM before
1833    // "data") or `serde_json::from_str` (BOM before "{"). Do the strip
1834    // exactly once here; subsequent chunks in a stream will not carry
1835    // another BOM.
1836    let raw = raw.strip_prefix('\u{feff}').unwrap_or(raw);
1837    let mut message = String::new();
1838    let mut reasoning = String::new();
1839    let mut prompt_tokens = None;
1840    let mut completion_tokens = None;
1841    let mut cached_tokens = None;
1842    let mut usage_reported = false;
1843    let mut finish_reason = None;
1844    let mut cost_usd_micros = 0u64;
1845    let mut cost_reported = false;
1846    let mut tool_call_deltas = Vec::new();
1847    let mut is_sse = false;
1848    let mut data = Vec::new();
1849    // Track the last `event:` type seen in this frame. SSE §9.2.8 says
1850    // an SSE client dispatches by event name; "message" (empty or
1851    // omitted) is the default. AgentVisor AI attributes the `data:`
1852    // payload to the audit trail (receipt, ATIF) as if it were the
1853    // model's output — but a hostile upstream (rogue provider,
1854    // TLS-MITM at egress, misconfigured cache) can send
1855    // `event: error\ndata: {"choices":[{"delta":{"content":"…"}}]}` —
1856    // spec-compliant SSE clients (browsers, OpenAI SDK's error hook)
1857    // would dispatch that to the ERROR listener, so the user sees an
1858    // error UI while our receipt records the payload as the model's
1859    // response. That defeats the cryptographic-attestation posture.
1860    // Only accept `event:` empty or `"message"` for capture; other
1861    // types abort the frame with a diagnostic and let the caller
1862    // mark_capture_failed rather than sign attributable content.
1863    let mut event_type = String::new();
1864    for line in raw.split(['\r', '\n']) {
1865        if let Some(value) = line.strip_prefix("data:") {
1866            is_sse = true;
1867            // SSE spec (§9.2.6): strip exactly ONE leading U+0020
1868            // SPACE — the field is otherwise verbatim. `.trim_start()`
1869            // would eat runs of Unicode whitespace and silently
1870            // mis-account any provider that sends `data:  {...}` with a
1871            // meaningful second byte.
1872            data.push(value.strip_prefix(' ').unwrap_or(value));
1873        } else if let Some(value) = line.strip_prefix("event:") {
1874            is_sse = true;
1875            event_type = value.strip_prefix(' ').unwrap_or(value).to_owned();
1876        } else if line == "data"
1877            || line.starts_with("id:")
1878            || line.starts_with("retry:")
1879            || line.starts_with(':')
1880        {
1881            is_sse = true;
1882        }
1883    }
1884    if is_sse && !event_type.is_empty() && event_type != "message" {
1885        return Err(format!(
1886            "provider SSE frame carries unsupported event type {event_type:?}; \
1887             AgentVisor AI only captures the default `message` event because non-message \
1888             events (error, ping, custom) are dispatched to different client listeners \
1889             per SSE §9.2.8 and would be attributed to the wrong audit surface"
1890        ));
1891    }
1892    let candidate = if is_sse {
1893        if data.is_empty() {
1894            return Ok(None);
1895        }
1896        data.join("\n")
1897    } else {
1898        raw.trim().to_owned()
1899    };
1900    // Robust `[DONE]` sentinel handling. Some providers double-terminate
1901    // (`data: [DONE]\ndata: [DONE]`), send trailing whitespace, or emit
1902    // an empty keepalive line followed by `[DONE]`. The strict
1903    // byte-exact check used to fail every subsequent stream on such
1904    // benign variants — `trim` + per-line probe catches them.
1905    if candidate.is_empty()
1906        || candidate.trim() == "[DONE]"
1907        || data
1908            .iter()
1909            .all(|entry| entry.trim().is_empty() || entry.trim() == "[DONE]")
1910            && !data.is_empty()
1911    {
1912        return Ok(None);
1913    }
1914    let value = serde_json::from_str::<Value>(&candidate)
1915        .map_err(|error| format!("invalid provider JSON frame: {error}"))?;
1916    let model_name = value.get("model").and_then(Value::as_str).map(str::to_owned);
1917    if let Some(usage) = value.get("usage") {
1918        usage_reported = true;
1919        prompt_tokens = provider_u64(usage.get("prompt_tokens"), "prompt_tokens")?.or(prompt_tokens);
1920        completion_tokens =
1921            provider_u64(usage.get("completion_tokens"), "completion_tokens")?.or(completion_tokens);
1922        cached_tokens = provider_u64(
1923            usage.pointer("/prompt_tokens_details/cached_tokens"),
1924            "cached_tokens",
1925        )?
1926        .or(cached_tokens);
1927    }
1928    let cost_value = value.pointer("/usage/cost_usd").or_else(|| value.get("cost_usd"));
1929    if let Some(cost_value) = cost_value {
1930        let cost = cost_value
1931            .as_f64()
1932            .ok_or_else(|| "provider cost_usd is not a number".to_owned())?;
1933        if !cost.is_finite() || cost < 0.0 {
1934            return Err("provider cost_usd is not finite and nonnegative".to_owned());
1935        }
1936        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
1937        {
1938            let micros = (cost * av_core::units::USD_MICROS_PER_DOLLAR as f64).round();
1939            if micros > av_core::error::JCS_SAFE_MAX as f64 {
1940                return Err("provider cost exceeds JCS-safe receipt bounds".to_owned());
1941            }
1942            cost_usd_micros = micros as u64;
1943            cost_reported = true;
1944        }
1945    }
1946    let has_choices = value.get("choices").is_some();
1947    let choices = match value.get("choices") {
1948        Some(Value::Array(choices)) => Some(choices),
1949        Some(_) => return Err("provider choices is not an array".to_owned()),
1950        None => None,
1951    };
1952    if let Some(choices) = choices {
1953        for choice in choices {
1954            if let Some(reason) = choice.get("finish_reason").and_then(Value::as_str) {
1955                finish_reason = Some(reason.to_owned());
1956            }
1957            let content = choice
1958                .pointer("/delta/content")
1959                .or_else(|| choice.pointer("/message/content"))
1960                .and_then(Value::as_str);
1961            if let Some(content) = content {
1962                message.push_str(content);
1963            }
1964            let reasoning_content = choice
1965                .pointer("/delta/reasoning_content")
1966                .or_else(|| choice.pointer("/message/reasoning_content"))
1967                .and_then(Value::as_str);
1968            if let Some(content) = reasoning_content {
1969                reasoning.push_str(content);
1970            }
1971            let streaming_calls = choice.pointer("/delta/tool_calls");
1972            let calls = streaming_calls
1973                .or_else(|| choice.pointer("/message/tool_calls"))
1974                .map(|calls| {
1975                    calls
1976                        .as_array()
1977                        .ok_or_else(|| "provider tool_calls is not an array".to_owned())
1978                })
1979                .transpose()?;
1980            if let Some(calls) = calls {
1981                for (position, call) in calls.iter().enumerate() {
1982                    let index = match call.get("index").and_then(Value::as_u64) {
1983                        Some(index) => index,
1984                        None if streaming_calls.is_some() => {
1985                            return Err("streaming provider tool call has no index".to_owned());
1986                        }
1987                        None => u64::try_from(position)
1988                            .map_err(|_| "provider tool-call index overflow".to_owned())?,
1989                    };
1990                    tool_call_deltas.push(ProviderToolCallDelta {
1991                        index,
1992                        id: call.get("id").and_then(Value::as_str).map(str::to_owned),
1993                        name: call
1994                            .pointer("/function/name")
1995                            .and_then(Value::as_str)
1996                            .map(str::to_owned),
1997                        arguments: call
1998                            .pointer("/function/arguments")
1999                            .and_then(Value::as_str)
2000                            .unwrap_or_default()
2001                            .to_owned(),
2002                    });
2003                }
2004            }
2005        }
2006    }
2007    let completion_tokens = completion_tokens.unwrap_or_else(|| {
2008        let mut estimated = av_core::tokens::approx_tokens(&message)
2009            .saturating_add(av_core::tokens::approx_tokens(&reasoning));
2010        for call in &tool_call_deltas {
2011            estimated = estimated
2012                .saturating_add(call.name.as_deref().map_or(0, av_core::tokens::approx_tokens))
2013                .saturating_add(av_core::tokens::approx_tokens(&call.arguments));
2014        }
2015        estimated
2016    });
2017    Ok(Some(ParsedProviderChunk {
2018        message,
2019        reasoning: (!reasoning.is_empty()).then_some(reasoning),
2020        model_name,
2021        metrics: av_events::EventMetrics {
2022            prompt_tokens,
2023            completion_tokens: Some(completion_tokens),
2024            cached_tokens,
2025            pruned_tokens: None,
2026            pruning_ratio_millis: None,
2027        },
2028        usage_reported,
2029        finish_reason,
2030        cost_usd_micros,
2031        cost_reported,
2032        has_choices,
2033        tool_call_deltas,
2034    }))
2035}
2036
2037fn provider_u64(value: Option<&Value>, field: &str) -> Result<Option<u64>, String> {
2038    let Some(value) = value else {
2039        return Ok(None);
2040    };
2041    let value = value
2042        .as_u64()
2043        .ok_or_else(|| format!("provider {field} is not a nonnegative integer"))?;
2044    if value > av_core::error::JCS_SAFE_MAX {
2045        return Err(format!("provider {field} exceeds JCS-safe bounds"));
2046    }
2047    Ok(Some(value))
2048}
2049
2050fn map_finish_reason(native: &str) -> av_events::StopReason {
2051    match native {
2052        "stop" | "stop_sequence" | "end_turn" => av_events::StopReason::Stop,
2053        "length" | "max_tokens" => av_events::StopReason::MaxTokens,
2054        "tool_calls" | "function_call" | "tool_use" => av_events::StopReason::ToolUse,
2055        "content_filter" => av_events::StopReason::ContentFilter,
2056        _ => av_events::StopReason::Other,
2057    }
2058}
2059
2060impl Drop for AbortFinalizingStream {
2061    fn drop(&mut self) {
2062        // Session-being-closed intentionally still captures: `wait_for_streams`
2063        // in `close_session_locked` is blocked on this stream's `SessionLease`,
2064        // so the response_capture we submit here will land before finalize
2065        // reads the chain / atif.
2066        if self.completed {
2067            return;
2068        }
2069        // Capture is_closed once so the fail-closed guards below agree on
2070        // the same view. When a concurrent close is already draining this
2071        // stream via `wait_for_streams`, the finalize path is imminent and
2072        // we must not `mark_capture_failed` on the last budget check /
2073        // trailing frame — otherwise `close_session_locked`'s post-drain
2074        // `capture_failed` guard seals the session (`mark_artifact_committed`
2075        // + committed claim) and returns `CaptureIncomplete`: the close is
2076        // never retried and the session finalizes with no artifact at all.
2077        // On a normal abort (no concurrent close) the fail-closed marks
2078        // still fire, so a mid-flight budget verdict or garbled trailing
2079        // frame still refuses the capture on the client's next request.
2080        let is_closed = self.session.is_closed();
2081        if !is_closed && self.pending_budget.is_some() {
2082            self.session.mark_capture_failed();
2083        }
2084        // Round-29 F5: abort the pending budget task. `spawn_blocking`
2085        // returns a JoinHandle whose Drop does NOT cancel the queued
2086        // closure; the blocking pool would otherwise run
2087        // `ActionBudget::try_tokens(delta)` AFTER the session was
2088        // sealed by the preceding `mark_capture_failed`, silently
2089        // debiting the session's budget key for a request the client
2090        // never received. `abort()` is best-effort for a closure
2091        // already picked up by the blocking pool (blocking tasks
2092        // have no cancellation points) but reliably cancels a
2093        // still-queued task — closing the common case. A full fix
2094        // would thread a shutdown token into `try_tokens`; deferred
2095        // until that helper takes a cancellation argument.
2096        if let Some(pending) = self.pending_budget.take() {
2097            pending.task.abort();
2098        }
2099        let budget_delta = match self.flush_protocol_buffer() {
2100            Ok(delta) => delta,
2101            Err(error) => {
2102                tracing::warn!(%error, "provider stream flush failed on drop");
2103                if !is_closed {
2104                    self.session.mark_capture_failed();
2105                }
2106                0
2107            }
2108        };
2109        if !is_closed && budget_delta > 0 {
2110            self.session.mark_capture_failed();
2111        }
2112        if !self.session.capture_failed() {
2113            if let Err(error) = self.submit_response_capture(None) {
2114                tracing::warn!(%error, "response-capture submit failed on drop");
2115                self.session.mark_capture_failed();
2116            }
2117        }
2118        if !is_closed {
2119            let session = Arc::clone(&self.session);
2120            let finalizer = self.finalizer.clone();
2121            let session_id = session.id.clone();
2122            match tokio::runtime::Handle::try_current() {
2123                Ok(runtime) => {
2124                    // Detach: the drop is sync and cannot await. Mirror
2125                    // the outcome to tracing *and* a Prometheus counter
2126                    // so a PromQL alert can catch the class — the
2127                    // fallback path is otherwise invisible until the
2128                    // idle sweeper reaps the "still open" session.
2129                    runtime.spawn(async move {
2130                        if let Err(error) = finalizer.close_session(session, StopReason::Other).await {
2131                            finalizer
2132                                .metrics()
2133                                .counter(
2134                                    "av_stream_abort_close_failures_total",
2135                                    "Background close after a stream abort failed; \
2136                                     the session is left open until the idle sweeper reaps it",
2137                                )
2138                                .inc();
2139                            tracing::warn!(
2140                                %error,
2141                                %session_id,
2142                                "background close on stream abort failed"
2143                            );
2144                        }
2145                    });
2146                }
2147                Err(error) => {
2148                    // Runtime is gone (shutdown, drop from a blocking
2149                    // thread) — there is no place to await the close.
2150                    // Mark capture failed so the reconciler retries on
2151                    // the next tick instead of finalising a session
2152                    // that never ran to completion.
2153                    self.finalizer
2154                        .metrics()
2155                        .counter(
2156                            "av_stream_abort_no_runtime_total",
2157                            "Stream abort observed no tokio runtime; capture marked failed for \
2158                             reconciler retry",
2159                        )
2160                        .inc();
2161                    tracing::warn!(
2162                        %error,
2163                        %session_id,
2164                        "no tokio runtime available for stream-abort close; marking capture failed for reconciler retry"
2165                    );
2166                    self.session.mark_capture_failed();
2167                }
2168            }
2169        }
2170    }
2171}
2172
2173#[cfg(test)]
2174mod tests {
2175    #![allow(
2176        clippy::expect_used,
2177        clippy::indexing_slicing,
2178        clippy::panic,
2179        clippy::unwrap_used
2180    )]
2181
2182    use super::*;
2183    use av_bridge::{BusError, EventBus, PublishAck, StoredEvent};
2184    use av_receipts::Ed25519Signer;
2185    use av_sandbox::{Sandbox, SandboxConfig};
2186    use av_state::{InMemoryStore, Spend, StateError, StateStore};
2187    use axum::http::Request;
2188    use std::time::Duration;
2189    use tower::ServiceExt;
2190
2191    struct NullBus;
2192
2193    struct SlowStore {
2194        inner: InMemoryStore,
2195        calls: Arc<std::sync::atomic::AtomicUsize>,
2196    }
2197
2198    impl StateStore for SlowStore {
2199        fn add(&self, key: &str, delta: u64) -> Result<u64, StateError> {
2200            self.inner.add(key, delta)
2201        }
2202
2203        fn get(&self, key: &str) -> Result<u64, StateError> {
2204            self.inner.get(key)
2205        }
2206
2207        fn try_spend(&self, key: &str, amount: u64, limit: u64) -> Result<bool, StateError> {
2208            self.calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2209            std::thread::sleep(Duration::from_millis(100));
2210            self.inner.try_spend(key, amount, limit)
2211        }
2212
2213        fn try_spend_many(&self, spends: &[Spend]) -> Result<Option<usize>, StateError> {
2214            self.inner.try_spend_many(spends)
2215        }
2216
2217        fn remove(&self, key: &str) {
2218            self.inner.remove(key);
2219        }
2220    }
2221
2222    impl EventBus for NullBus {
2223        fn publish(&self, topic: &str, _key: &str, _value: &Value) -> Result<PublishAck, BusError> {
2224            Ok(PublishAck {
2225                topic: topic.to_owned(),
2226                partition: 0,
2227                offset: 0,
2228            })
2229        }
2230
2231        fn fetch(
2232            &self,
2233            _topic: &str,
2234            _partition: u32,
2235            _offset: u64,
2236            _max: usize,
2237        ) -> Result<Vec<StoredEvent>, BusError> {
2238            Ok(Vec::new())
2239        }
2240
2241        fn partitions(&self, _topic: &str) -> Result<u32, BusError> {
2242            Ok(1)
2243        }
2244
2245        fn topics(&self) -> Vec<String> {
2246            av_events::EventClass::all()
2247                .iter()
2248                .map(|class| class.topic().to_owned())
2249                .collect()
2250        }
2251    }
2252
2253    async fn mock_chat(Json(payload): Json<Value>) -> Response {
2254        if payload.get("model").and_then(Value::as_str) == Some("empty-error") {
2255            let mut response = Response::new(Body::empty());
2256            *response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
2257            response.headers_mut().insert(
2258                axum::http::header::CONTENT_TYPE,
2259                HeaderValue::from_static("application/json"),
2260            );
2261            return response;
2262        }
2263        if payload.get("model").and_then(Value::as_str) == Some("json-error") {
2264            return (StatusCode::INTERNAL_SERVER_ERROR, Json(json!({}))).into_response();
2265        }
2266        if payload.get("model").and_then(Value::as_str) == Some("empty-success") {
2267            return Json(json!({})).into_response();
2268        }
2269        if payload.get("model").and_then(Value::as_str) == Some("provisional-usage") {
2270            let chunks = vec![
2271                Ok::<_, std::convert::Infallible>(Bytes::from_static(
2272                    b"data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\n",
2273                )),
2274                Ok(Bytes::from_static(
2275                    b"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2}}\n\n",
2276                )),
2277            ];
2278            let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
2279            response.headers_mut().insert(
2280                axum::http::header::CONTENT_TYPE,
2281                HeaderValue::from_static("text/event-stream"),
2282            );
2283            return response;
2284        }
2285        if payload.get("model").and_then(Value::as_str) == Some("tool-no-usage") {
2286            let arguments = "x".repeat(4_096);
2287            let event = format!(
2288                "data: {{\"choices\":[{{\"delta\":{{\"tool_calls\":[{{\"index\":0,\"function\":{{\"name\":\"read\",\"arguments\":\"{arguments}\"}}}}]}}}}]}}\n\n"
2289            );
2290            let mut response = Response::new(Body::from(event));
2291            response.headers_mut().insert(
2292                axum::http::header::CONTENT_TYPE,
2293                HeaderValue::from_static("text/event-stream"),
2294            );
2295            return response;
2296        }
2297        if payload.get("model").and_then(Value::as_str) == Some("tool-oob-index") {
2298            let event = "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":18446744073709551615,\"function\":{\"name\":\"read\",\"arguments\":\"{}\"}}]}}]}\n\n";
2299            let mut response = Response::new(Body::from(event));
2300            response.headers_mut().insert(
2301                axum::http::header::CONTENT_TYPE,
2302                HeaderValue::from_static("text/event-stream"),
2303            );
2304            return response;
2305        }
2306        if payload.get("model").and_then(Value::as_str) == Some("regressive-usage") {
2307            let chunks = vec![
2308                Ok::<_, std::convert::Infallible>(Bytes::from_static(
2309                    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",
2310                )),
2311                Ok(Bytes::from_static(
2312                    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",
2313                )),
2314            ];
2315            let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
2316            response.headers_mut().insert(
2317                axum::http::header::CONTENT_TYPE,
2318                HeaderValue::from_static("text/event-stream"),
2319            );
2320            return response;
2321        }
2322        if payload.get("model").and_then(Value::as_str) == Some("split-json") {
2323            let body = serde_json::to_vec(&json!({
2324                "model": "split-json",
2325                "choices": [{"message": {"role": "assistant", "content": "héllo"}}],
2326                "usage": {"prompt_tokens": 5, "completion_tokens": 2}
2327            }))
2328            .unwrap();
2329            let split = body.iter().position(|byte| *byte == 0xc3).unwrap() + 1;
2330            let chunks = vec![
2331                Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(&body[..split])),
2332                Ok(Bytes::copy_from_slice(&body[split..])),
2333            ];
2334            let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
2335            response.headers_mut().insert(
2336                axum::http::header::CONTENT_TYPE,
2337                HeaderValue::from_static("application/json"),
2338            );
2339            return response;
2340        }
2341        if payload.get("model").and_then(Value::as_str) == Some("malformed-json") {
2342            let mut response = Response::new(Body::from("{\"choices\":["));
2343            response.headers_mut().insert(
2344                axum::http::header::CONTENT_TYPE,
2345                HeaderValue::from_static("application/json"),
2346            );
2347            return response;
2348        }
2349        if payload.get("stream").and_then(Value::as_bool).unwrap_or(false) {
2350            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";
2351            let bytes = event.as_bytes();
2352            let chunks = vec![
2353                Ok::<_, std::convert::Infallible>(Bytes::copy_from_slice(&bytes[..19])),
2354                Ok(Bytes::copy_from_slice(&bytes[19..97])),
2355                Ok(Bytes::copy_from_slice(&bytes[97..])),
2356            ];
2357            let mut response = Response::new(Body::from_stream(futures::stream::iter(chunks)));
2358            response.headers_mut().insert(
2359                axum::http::header::CONTENT_TYPE,
2360                HeaderValue::from_static("text/event-stream"),
2361            );
2362            response
2363        } else {
2364            Json(json!({"choices": [{"message": {"role": "assistant", "content": "hello"}}]})).into_response()
2365        }
2366    }
2367
2368    async fn counting_tool(
2369        State(calls): State<Arc<std::sync::atomic::AtomicUsize>>,
2370        body: Bytes,
2371    ) -> Response {
2372        calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2373        (StatusCode::OK, body).into_response()
2374    }
2375
2376    async fn redirecting_tool(State(calls): State<Arc<std::sync::atomic::AtomicUsize>>) -> Response {
2377        calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2378        (
2379            StatusCode::TEMPORARY_REDIRECT,
2380            [(axum::http::header::LOCATION, "/effect")],
2381        )
2382            .into_response()
2383    }
2384
2385    async fn redirected_effect(State(calls): State<Arc<std::sync::atomic::AtomicUsize>>) -> Response {
2386        calls.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
2387        StatusCode::OK.into_response()
2388    }
2389
2390    struct HeldToolState {
2391        arrived: std::sync::atomic::AtomicBool,
2392        release: tokio::sync::Notify,
2393    }
2394
2395    async fn held_tool(State(state): State<Arc<HeldToolState>>, body: Bytes) -> Response {
2396        state.arrived.store(true, std::sync::atomic::Ordering::Release);
2397        state.release.notified().await;
2398        (StatusCode::OK, body).into_response()
2399    }
2400
2401    async fn test_state(spool: &std::path::Path) -> (AppState, tokio::task::JoinHandle<()>) {
2402        test_state_with_token_cap(spool, None).await
2403    }
2404
2405    async fn test_state_with_token_cap(
2406        spool: &std::path::Path,
2407        max_tokens: Option<u64>,
2408    ) -> (AppState, tokio::task::JoinHandle<()>) {
2409        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2410        let address = listener.local_addr().unwrap();
2411        let provider = Router::new().route("/v1/chat/completions", post(mock_chat));
2412        let server = tokio::spawn(async move {
2413            axum::serve(listener, provider).await.unwrap();
2414        });
2415        let mut config = crate::config::HarnessConfig::for_tests(
2416            &format!("http://{address}"),
2417            &spool.to_string_lossy(),
2418            &spool.to_string_lossy(),
2419        );
2420        config.budget.max_tokens = max_tokens;
2421        let state = AppState::new(
2422            config,
2423            Arc::new(InMemoryStore::new()),
2424            Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
2425            Arc::new(NullBus),
2426            None,
2427            Arc::new(Ed25519Signer::from_seed(&[11; 32])),
2428        )
2429        .unwrap();
2430        (state, server)
2431    }
2432
2433    fn chat_request(session: &str) -> Request<Body> {
2434        chat_request_with_payload(session, chat_payload())
2435    }
2436
2437    fn chat_request_with_payload(session: &str, payload: Value) -> Request<Body> {
2438        Request::builder()
2439            .method("POST")
2440            .uri("/v1/chat/completions")
2441            .header(axum::http::header::CONTENT_TYPE, "application/json")
2442            .header("x-av-session", session)
2443            .header("x-av-workflow", "unsigned")
2444            .body(Body::from(serde_json::to_vec(&payload).unwrap()))
2445            .unwrap()
2446    }
2447
2448    fn chat_payload() -> Value {
2449        json!({
2450            "model": "mock",
2451            "stream": true,
2452            "messages": [{"role": "user", "content": "hello"}],
2453        })
2454    }
2455
2456    fn active_records(
2457        directory: &std::path::Path,
2458        state: &AppState,
2459        session_id: &str,
2460    ) -> Vec<crate::worker::ActiveJournalRecord> {
2461        let digest = av_core::digest::sha256_hex(session_id.as_bytes());
2462        let journal =
2463            std::fs::read_to_string(directory.join(format!("{}.events.ndjson", &digest[..32]))).unwrap();
2464        journal
2465            .lines()
2466            .enumerate()
2467            .map(|(index, line)| {
2468                crate::journal::open(
2469                    &state.journal_key,
2470                    &format!("{session_id}:active"),
2471                    index as u64,
2472                    line.as_bytes(),
2473                )
2474                .unwrap()
2475            })
2476            .collect()
2477    }
2478
2479    #[tokio::test]
2480    async fn full_chat_close_and_promotion_flow() {
2481        let directory = tempfile::tempdir().unwrap();
2482        let (state, provider) = test_state(directory.path()).await;
2483        let app = build_router(state.clone());
2484
2485        let response = app.clone().oneshot(chat_request("http-flow")).await.unwrap();
2486        assert_eq!(response.status(), StatusCode::OK);
2487        assert_eq!(response.headers().get("x-av-session").unwrap(), "http-flow");
2488        let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
2489            .await
2490            .unwrap();
2491        assert!(String::from_utf8_lossy(&body).contains("[DONE]"));
2492
2493        let close = app
2494            .clone()
2495            .oneshot(
2496                Request::builder()
2497                    .method("POST")
2498                    .uri("/v1/sessions/http-flow/close")
2499                    .body(Body::empty())
2500                    .unwrap(),
2501            )
2502            .await
2503            .unwrap();
2504        let close_status = close.status();
2505        let close_bytes = axum::body::to_bytes(close.into_body(), 64 * 1024).await.unwrap();
2506        assert_eq!(
2507            close_status,
2508            StatusCode::OK,
2509            "{}",
2510            String::from_utf8_lossy(&close_bytes)
2511        );
2512        let close_body: Value = serde_json::from_slice(&close_bytes).unwrap();
2513        assert_eq!(close_body["kind"], "atif");
2514        let artifact = state
2515            .sessions
2516            .get("http-flow")
2517            .unwrap()
2518            .atif_path
2519            .lock()
2520            .clone()
2521            .unwrap();
2522        let trajectory: av_atif::Trajectory =
2523            serde_json::from_slice(&tokio::fs::read(&artifact).await.unwrap()).unwrap();
2524        if let Some(destination) = std::env::var_os("AV_HARBOR_INTEROP_OUT") {
2525            std::fs::copy(&artifact, destination).unwrap();
2526        }
2527        assert_eq!(
2528            trajectory.steps.len(),
2529            2,
2530            "request and response must both be captured"
2531        );
2532        assert_eq!(trajectory.steps[0].source, av_atif::Source::User);
2533        assert!(trajectory.steps[0].metrics.is_none());
2534        assert_eq!(trajectory.steps[1].source, av_atif::Source::Agent);
2535        assert_eq!(trajectory.steps[1].message, Value::String("hello".to_owned()));
2536        assert_eq!(
2537            trajectory.steps[1].metrics.as_ref().unwrap().cached_tokens,
2538            Some(4)
2539        );
2540
2541        let promoted = app
2542            .oneshot(
2543                Request::builder()
2544                    .method("POST")
2545                    .uri("/v1/sessions/http-flow/promote")
2546                    .body(Body::empty())
2547                    .unwrap(),
2548            )
2549            .await
2550            .unwrap();
2551        assert_eq!(promoted.status(), StatusCode::OK);
2552        let receipt: av_receipts::Receipt = serde_json::from_slice(
2553            &axum::body::to_bytes(promoted.into_body(), 64 * 1024)
2554                .await
2555                .unwrap(),
2556        )
2557        .unwrap();
2558        receipt.verify_embedded().unwrap();
2559        assert_eq!(receipt.body.stop_reason_id, av_events::StopReason::Stop.id());
2560        assert_eq!(receipt.body.cost.completion_tokens, 3);
2561        assert_eq!(receipt.body.cost.cached_tokens, 4);
2562        assert_eq!(receipt.body.cost.cost_usd_micros, 1_250);
2563        assert_eq!(
2564            receipt.body.cost.prompt_tokens,
2565            av_core::tokens::approx_tokens(&chat_payload().to_string()),
2566            "provider prompt usage must not be added a second time"
2567        );
2568        provider.abort();
2569    }
2570
2571    #[tokio::test]
2572    async fn mcp_metrics_and_abort_paths_are_enforced() {
2573        let directory = tempfile::tempdir().unwrap();
2574        let (state, provider) = test_state(directory.path()).await;
2575        let app = build_router(state.clone());
2576        let allowed = app
2577            .clone()
2578            .oneshot(
2579                Request::builder()
2580                    .method("POST")
2581                    .uri("/v1/mcp")
2582                    .header("x-av-session", "mcp-flow")
2583                    .body(Body::from(
2584                        r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"read","arguments":{}}}"#,
2585                    ))
2586                    .unwrap(),
2587            )
2588            .await
2589            .unwrap();
2590        assert_eq!(allowed.status(), StatusCode::OK);
2591        let allowed_body: Value = serde_json::from_slice(
2592            &axum::body::to_bytes(allowed.into_body(), 64 * 1024)
2593                .await
2594                .unwrap(),
2595        )
2596        .unwrap();
2597        assert!(allowed_body["decision_us"].as_u64().unwrap() < 5_000);
2598
2599        let blocked = app
2600            .clone()
2601            .oneshot(
2602                Request::builder()
2603                    .method("POST")
2604                    .uri("/v1/mcp")
2605                    .header("x-av-session", "mcp-flow")
2606                    .body(Body::from("not-json"))
2607                    .unwrap(),
2608            )
2609            .await
2610            .unwrap();
2611        assert_eq!(blocked.status(), StatusCode::FORBIDDEN);
2612
2613        let metrics_response = app
2614            .clone()
2615            .oneshot(Request::builder().uri("/metrics").body(Body::empty()).unwrap())
2616            .await
2617            .unwrap();
2618        let metrics = String::from_utf8_lossy(
2619            &axum::body::to_bytes(metrics_response.into_body(), 128 * 1024)
2620                .await
2621                .unwrap(),
2622        )
2623        .into_owned();
2624        assert!(metrics.contains("av_stage_duration_seconds") || metrics.contains("av_sessions"));
2625
2626        let aborted = app.clone().oneshot(chat_request("abort-flow")).await.unwrap();
2627        drop(aborted);
2628        tokio::time::timeout(Duration::from_secs(1), async {
2629            loop {
2630                if state
2631                    .sessions
2632                    .get("abort-flow")
2633                    .is_some_and(|session| session.is_closed())
2634                {
2635                    break;
2636                }
2637                tokio::task::yield_now().await;
2638            }
2639        })
2640        .await
2641        .expect("abort finalization timed out");
2642        provider.abort();
2643    }
2644
2645    /// The response-body stream drops mid-flight while another task is
2646    /// closing the session (idle timeout or explicit close). Recorded
2647    /// scenario: `try_close` in `close_session_locked` transitions
2648    /// `closed` 0→1, then blocks in `wait_for_streams` because the
2649    /// client still holds the response body's `SessionLease`. When
2650    /// the client eventually drops the response, `AbortFinalizingStream::drop`
2651    /// used to see `session.is_closed() == true` and return early —
2652    /// silently skipping `submit_response_capture`. The reader's
2653    /// response bytes never reached the chain, but `wait_for_streams`
2654    /// then unblocked and the finalize path signed a receipt whose
2655    /// `subject.event_count` reflects only the compression event.
2656    #[tokio::test]
2657    async fn stream_drop_during_concurrent_close_still_captures_response() {
2658        let directory = tempfile::tempdir().unwrap();
2659        let (state, provider) = test_state(directory.path()).await;
2660        let app = build_router(state.clone());
2661
2662        let request = Request::builder()
2663            .method("POST")
2664            .uri("/v1/chat/completions")
2665            .header(axum::http::header::CONTENT_TYPE, "application/json")
2666            .header("x-av-session", "stream-close-race")
2667            .header("x-av-workflow", "signed")
2668            .body(Body::from(serde_json::to_vec(&chat_payload()).unwrap()))
2669            .unwrap();
2670        let response = app.oneshot(request).await.unwrap();
2671        assert_eq!(response.status(), StatusCode::OK);
2672
2673        let session = state.sessions.get("stream-close-race").unwrap();
2674
2675        let close_task = {
2676            let finalizer = state.finalizer.clone();
2677            let session = Arc::clone(&session);
2678            tokio::spawn(async move { finalizer.close_session(session, StopReason::SessionClosed).await })
2679        };
2680
2681        tokio::time::timeout(Duration::from_secs(1), async {
2682            while !session.is_closed() {
2683                tokio::task::yield_now().await;
2684            }
2685        })
2686        .await
2687        .expect("concurrent close did not initiate before the response was dropped");
2688
2689        drop(response);
2690
2691        close_task
2692            .await
2693            .expect("close_session task panicked")
2694            .expect("close_session returned an error");
2695
2696        let final_count = session.chain.lock().count();
2697        assert!(
2698            final_count >= 2,
2699            "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}",
2700        );
2701
2702        provider.abort();
2703    }
2704
2705    /// Regression for the follow-on to the concurrent-close capture bug. The
2706    /// original fix removed the `is_closed()` early-return from
2707    /// `AbortFinalizingStream::drop` so the response capture would still be
2708    /// submitted while another task drained the stream via
2709    /// `wait_for_streams`. But `drop` also carried a defensive
2710    /// `if self.pending_budget.is_some() { self.session.mark_capture_failed(); }`
2711    /// intended for genuine abort-mid-flight: an in-flight completion-token
2712    /// budget task means we don't yet know if the token spend would have been
2713    /// refused, so fail-closed is safe for a normal client disconnect. When
2714    /// composed with the removal of the `is_closed()` early-return, however,
2715    /// this fail-closed mark now fires during the concurrent-close scenario
2716    /// too — even though the concurrent close is imminent and skipping the
2717    /// last budget verdict is harmless (the response bytes are captured
2718    /// internally regardless of whether we would have emitted them to the
2719    /// client). The resulting `capture_failed = 1` then makes the concurrent
2720    /// close's post-`wait_for_worker_jobs` `if session.capture_failed()` guard
2721    /// seal the session (`mark_artifact_committed` + committed claim) and
2722    /// return `Err(FinalizeError::CaptureIncomplete)`: the close is not
2723    /// retried, and the session finalizes without any artifact — the
2724    /// capture is silently lost. The fix: gate the pending-budget mark
2725    /// (and the budget-delta > 0 mark and the flush-error mark) on
2726    /// `!is_closed`, so during a concurrent close the drop submits the
2727    /// response capture without fail-closing the session.
2728    #[tokio::test]
2729    async fn drop_with_pending_budget_during_concurrent_close_does_not_stick_session() {
2730        let directory = tempfile::tempdir().unwrap();
2731        let (state, provider) = test_state(directory.path()).await;
2732
2733        let identity = av_events::AgentIdentity {
2734            version: "1".into(),
2735            charter: "c".into(),
2736            instance_uid: "i".into(),
2737            ttl_remaining_s: None,
2738        };
2739        let session = state.sessions.get_or_open(
2740            "pending-budget-close",
2741            crate::session::Workflow::Signed,
2742            &identity,
2743            &state.config.breaker,
2744        );
2745        // Reserve a full permit pair up front (worker slot + response
2746        // slot) so `submit_response_capture` in drop has somewhere to
2747        // send the job; the response permit lives on
2748        // `AbortFinalizingStream` for the stream's lifetime.
2749        let permits = state
2750            .worker
2751            .try_reserve_pair("pending-budget-close")
2752            .expect("test setup: fused worker/response permit");
2753        // The worker permit is dropped here — the test does not
2754        // actually submit a job through it; the response permit is
2755        // what the stream drop cares about.
2756        drop(permits.worker);
2757        let permit = permits.response;
2758        // Bump active_streams so a concurrent close would block in
2759        // `wait_for_streams`, mirroring the production scenario.
2760        let lease = crate::session::SessionLease::new(Arc::clone(&session));
2761
2762        // Set up an in-flight budget task. Its result is irrelevant to drop —
2763        // drop only checks `pending_budget.is_some()` — but the JoinHandle
2764        // has to be a real one so the field type-checks.
2765        let budget_task =
2766            tokio::spawn(async { Ok::<_, String>(av_state::BudgetDecision::Allowed { remaining: 100 }) });
2767
2768        let stream = AbortFinalizingStream {
2769            inner: futures::stream::empty::<Result<Bytes, std::io::Error>>().boxed(),
2770            session: Arc::clone(&session),
2771            identity: identity.clone(),
2772            response_permit: Some(permit),
2773            worker: state.worker.clone(),
2774            store: Arc::clone(&state.store),
2775            budget: state.config.budget.clone(),
2776            finalizer: state.finalizer.clone(),
2777            _lease: lease,
2778            response_marker: None,
2779            response_attempt_id: "pending-budget-attempt".into(),
2780            response_message: "captured content".into(),
2781            response_reasoning: String::new(),
2782            response_model: None,
2783            response_finish_reason: Some("stop".into()),
2784            upstream_status: StatusCode::OK,
2785            response_cost_usd_micros: 0,
2786            response_tool_calls: std::collections::BTreeMap::new(),
2787            response_metrics: av_events::EventMetrics::default(),
2788            charged_completion_tokens: 3,
2789            last_reported_completion_tokens: Some(3),
2790            last_reported_prompt_tokens: None,
2791            last_reported_cached_tokens: None,
2792            last_reported_cost_usd_micros: None,
2793            saw_chunk: true,
2794            capture_submitted: false,
2795            is_sse: true,
2796            protocol_buffer: Vec::new(),
2797            pending_output: std::collections::VecDeque::new(),
2798            pending_budget: Some(PendingBudget {
2799                task: budget_task,
2800                continuation: BudgetContinuation::FinishSse,
2801            }),
2802            captured_bytes: 100,
2803            completed: false,
2804        };
2805
2806        // Simulate the concurrent close's `try_close` having already fired.
2807        assert!(session.try_close(), "test setup: try_close should have won");
2808        assert!(session.is_closed());
2809
2810        // Drop the stream. This is what happens when the response body is
2811        // released while another task is blocked in `wait_for_streams`.
2812        drop(stream);
2813
2814        assert!(
2815            !session.capture_failed(),
2816            "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",
2817        );
2818
2819        // Now run the concurrent close to completion. It must succeed —
2820        // wait_for_streams sees active=0 (lease dropped), wait_for_worker_jobs
2821        // drains the response job we just submitted, and the capture_failed
2822        // check must find `false` so the finalize path can proceed to
2823        // sign the receipt.
2824        let outcome = tokio::time::timeout(
2825            Duration::from_secs(3),
2826            state
2827                .finalizer
2828                .close_session(Arc::clone(&session), StopReason::SessionClosed),
2829        )
2830        .await
2831        .expect("close_session hung after drop released the stream");
2832        outcome.expect(
2833            "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",
2834        );
2835
2836        provider.abort();
2837    }
2838
2839    #[tokio::test]
2840    async fn duplicate_tool_id_replays_outcome_without_reexecution() {
2841        let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2842        let tool_address = tool_listener.local_addr().unwrap();
2843        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2844        let tool_router = Router::new()
2845            .route("/mcp", post(counting_tool))
2846            .with_state(Arc::clone(&calls));
2847        let tool_server = tokio::spawn(async move {
2848            axum::serve(tool_listener, tool_router).await.unwrap();
2849        });
2850        let directory = tempfile::tempdir().unwrap();
2851        let (mut state, provider) = test_state(directory.path()).await;
2852        Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
2853            Some(format!("http://{tool_address}/mcp"));
2854        let app = build_router(state.clone());
2855        let body = r#"{"jsonrpc":"2.0","id":"once","method":"tools/call","params":{"name":"read","arguments":{"id":7}}}"#;
2856        for _ in 0..2 {
2857            let response = app
2858                .clone()
2859                .oneshot(
2860                    Request::builder()
2861                        .method("POST")
2862                        .uri("/v1/mcp")
2863                        .header("x-av-session", "tool-once")
2864                        .body(Body::from(body))
2865                        .unwrap(),
2866                )
2867                .await
2868                .unwrap();
2869            assert_eq!(response.status(), StatusCode::OK);
2870            assert_eq!(
2871                axum::body::to_bytes(response.into_body(), 64 * 1024)
2872                    .await
2873                    .unwrap(),
2874                body
2875            );
2876        }
2877        assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
2878        let session = state.sessions.get("tool-once").unwrap();
2879        assert_eq!(
2880            session
2881                .totals
2882                .tool_calls
2883                .load(std::sync::atomic::Ordering::Acquire),
2884            1
2885        );
2886        let audited_path = std::fs::read_dir(directory.path().join(crate::spool::TOOL_EXECUTIONS))
2887            .unwrap()
2888            .map(|entry| entry.unwrap().path())
2889            .find(|path| path.extension().is_some_and(|extension| extension == "audited"))
2890            .unwrap();
2891        std::fs::remove_file(audited_path).unwrap();
2892        let unaudited_replay = app
2893            .clone()
2894            .oneshot(
2895                Request::builder()
2896                    .method("POST")
2897                    .uri("/v1/mcp")
2898                    .header("x-av-session", "tool-once")
2899                    .body(Body::from(body))
2900                    .unwrap(),
2901            )
2902            .await
2903            .unwrap();
2904        assert_eq!(unaudited_replay.status(), StatusCode::OK);
2905        assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
2906        assert_eq!(
2907            session
2908                .totals
2909                .tool_calls
2910                .load(std::sync::atomic::Ordering::Acquire),
2911            1,
2912            "unaudited outcome recovery must not spend tool budget again"
2913        );
2914
2915        let changed_request = app
2916            .clone()
2917            .oneshot(
2918                Request::builder()
2919                    .method("POST")
2920                    .uri("/v1/mcp")
2921                    .header("x-av-session", "tool-once")
2922                    .body(Body::from(
2923                        r#"{"jsonrpc":"2.0","id":"once","method":"tools/call","params":{"name":"read","arguments":{"id":8}}}"#,
2924                    ))
2925                    .unwrap(),
2926            )
2927            .await
2928            .unwrap();
2929        assert_eq!(changed_request.status(), StatusCode::CONFLICT);
2930        assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
2931        tool_server.abort();
2932        provider.abort();
2933    }
2934
2935    /// End-to-end proof that a configured upstream API key is injected on
2936    /// the upstream wire: a capturing mock provider records the header the
2937    /// harness actually sent for both Bearer and Azure raw-key styles, and
2938    /// the passthrough mode relays the client's own Authorization header.
2939    #[tokio::test]
2940    async fn upstream_auth_header_reaches_provider_wire() {
2941        type Captured = Arc<parking_lot::Mutex<Vec<(Option<String>, Option<String>)>>>;
2942        async fn capture_chat(
2943            State(captured): State<Captured>,
2944            headers: HeaderMap,
2945            Json(_): Json<Value>,
2946        ) -> Response {
2947            captured.lock().push((
2948                headers
2949                    .get(axum::http::header::AUTHORIZATION)
2950                    .and_then(|value| value.to_str().ok())
2951                    .map(str::to_owned),
2952                headers
2953                    .get("api-key")
2954                    .and_then(|value| value.to_str().ok())
2955                    .map(str::to_owned),
2956            ));
2957            Json(json!({"choices": [], "usage": {"prompt_tokens": 1, "completion_tokens": 1}}))
2958                .into_response()
2959        }
2960        let captured: Captured = Arc::new(parking_lot::Mutex::new(Vec::new()));
2961        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2962        let address = listener.local_addr().unwrap();
2963        let provider_router = Router::new()
2964            .route("/v1/chat/completions", post(capture_chat))
2965            .route("/openai/deployments/d1/chat/completions", post(capture_chat))
2966            .with_state(Arc::clone(&captured));
2967        let provider = tokio::spawn(async move {
2968            axum::serve(listener, provider_router).await.unwrap();
2969        });
2970        let directory = tempfile::tempdir().unwrap();
2971        let key_path = directory.path().join("upstream.key");
2972        std::fs::write(&key_path, "sk-wire-test\n").unwrap();
2973        #[cfg(unix)]
2974        {
2975            use std::os::unix::fs::PermissionsExt as _;
2976            std::fs::set_permissions(&key_path, std::fs::Permissions::from_mode(0o600)).unwrap();
2977        }
2978        let spool = directory.path().to_string_lossy();
2979        let build_state = |config: crate::config::HarnessConfig| {
2980            AppState::new(
2981                config,
2982                Arc::new(InMemoryStore::new()),
2983                Arc::new(Sandbox::new(SandboxConfig::default(), Vec::new()).unwrap()),
2984                Arc::new(NullBus),
2985                None,
2986                Arc::new(Ed25519Signer::from_seed(&[12; 32])),
2987            )
2988            .unwrap()
2989        };
2990        let send_chat = |state: AppState, authorization: Option<&'static str>| async move {
2991            let mut request = Request::builder()
2992                .method("POST")
2993                .uri("/v1/chat/completions")
2994                .header(axum::http::header::CONTENT_TYPE, "application/json")
2995                .header("x-av-session", "auth-wire");
2996            if let Some(value) = authorization {
2997                request = request.header(axum::http::header::AUTHORIZATION, value);
2998            }
2999            let response = build_router(state)
3000                .oneshot(request.body(Body::from(chat_payload().to_string())).unwrap())
3001                .await
3002                .unwrap();
3003            assert_eq!(response.status(), StatusCode::OK);
3004        };
3005
3006        // 1) Bearer key from an owner-only file.
3007        let mut config =
3008            crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
3009        config.upstream_api_key_file = Some(key_path.to_string_lossy().into_owned());
3010        send_chat(build_state(config), None).await;
3011        assert_eq!(
3012            captured.lock().pop().unwrap(),
3013            (Some("Bearer sk-wire-test".to_owned()), None),
3014            "Bearer credential must reach the provider"
3015        );
3016
3017        // 2) Azure style: raw key in a custom header on a custom path.
3018        let mut config =
3019            crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
3020        config.upstream_api_key_file = Some(key_path.to_string_lossy().into_owned());
3021        config.upstream_auth_header = "api-key".into();
3022        config.upstream_auth_scheme = String::new();
3023        config.upstream_chat_path = "/openai/deployments/d1/chat/completions".into();
3024        send_chat(build_state(config), None).await;
3025        assert_eq!(
3026            captured.lock().pop().unwrap(),
3027            (None, Some("sk-wire-test".to_owned())),
3028            "raw key must reach the provider on the configured path and header"
3029        );
3030
3031        // 3) Passthrough: the client's own Authorization header is relayed.
3032        let mut config =
3033            crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
3034        config.upstream_authorization_passthrough = true;
3035        send_chat(build_state(config), Some("Bearer client-owned-key")).await;
3036        assert_eq!(
3037            captured.lock().pop().unwrap(),
3038            (Some("Bearer client-owned-key".to_owned()), None),
3039            "passthrough must relay the client credential"
3040        );
3041
3042        // 4) No identity validator configured but the client sent a
3043        //    bearer: refuse with 401 rather than record the request as
3044        //    anonymous (a repudiation vector — see resolve_identity's
3045        //    "bearer presented but validator not configured" arm).
3046        //    The response is the block; the fact that no wire capture
3047        //    landed also proves the credential does not leak.
3048        let config = crate::config::HarnessConfig::for_tests(&format!("http://{address}"), &spool, &spool);
3049        let request = Request::builder()
3050            .method("POST")
3051            .uri("/v1/chat/completions")
3052            .header(axum::http::header::CONTENT_TYPE, "application/json")
3053            .header("x-av-session", "auth-wire")
3054            .header(axum::http::header::AUTHORIZATION, "Bearer client-owned-key")
3055            .body(Body::from(chat_payload().to_string()))
3056            .unwrap();
3057        let response = build_router(build_state(config)).oneshot(request).await.unwrap();
3058        assert_eq!(
3059            response.status(),
3060            StatusCode::UNAUTHORIZED,
3061            "presenting a bearer with no validator configured must be refused rather than \
3062             silently recorded as anonymous — that would let a signed request end up \
3063             attributed to `charter=anonymous` in the receipt"
3064        );
3065        assert!(
3066            captured.lock().is_empty(),
3067            "no credential must reach upstream when the request was refused"
3068        );
3069        provider.abort();
3070    }
3071
3072    /// The tool-upstream bearer token must be injected on forwarded MCP
3073    /// calls (and must not depend on the caller's own headers).
3074    #[tokio::test]
3075    async fn tool_upstream_bearer_reaches_tool_server() {
3076        type Captured = Arc<parking_lot::Mutex<Vec<Option<String>>>>;
3077        async fn capture_tool(State(captured): State<Captured>, headers: HeaderMap, body: Bytes) -> Response {
3078            captured.lock().push(
3079                headers
3080                    .get(axum::http::header::AUTHORIZATION)
3081                    .and_then(|value| value.to_str().ok())
3082                    .map(str::to_owned),
3083            );
3084            (StatusCode::OK, body).into_response()
3085        }
3086        let captured: Captured = Arc::new(parking_lot::Mutex::new(Vec::new()));
3087        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3088        let address = listener.local_addr().unwrap();
3089        let tool_router = Router::new()
3090            .route("/mcp", post(capture_tool))
3091            .with_state(Arc::clone(&captured));
3092        let tool_server = tokio::spawn(async move {
3093            axum::serve(listener, tool_router).await.unwrap();
3094        });
3095        let directory = tempfile::tempdir().unwrap();
3096        let token_path = directory.path().join("mcp.token");
3097        std::fs::write(&token_path, "tool-secret\n").unwrap();
3098        #[cfg(unix)]
3099        {
3100            use std::os::unix::fs::PermissionsExt as _;
3101            std::fs::set_permissions(&token_path, std::fs::Permissions::from_mode(0o600)).unwrap();
3102        }
3103        let (mut state, provider) = test_state(directory.path()).await;
3104        {
3105            let config = Arc::get_mut(&mut state.config).unwrap();
3106            config.tool_upstream_url = Some(format!("http://{address}/mcp"));
3107            config.tool_upstream_bearer_file = Some(token_path.to_string_lossy().into_owned());
3108        }
3109        // test_state resolved auth from the pre-mutation config; re-resolve.
3110        state.tool_auth = crate::pipeline::resolve_tool_auth(&state.config).unwrap();
3111        let response = build_router(state)
3112            .oneshot(
3113                Request::builder()
3114                    .method("POST")
3115                    .uri("/v1/mcp")
3116                    .header("x-av-session", "tool-auth")
3117                    .body(Body::from(
3118                        r#"{"jsonrpc":"2.0","id":"t1","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
3119                    ))
3120                    .unwrap(),
3121            )
3122            .await
3123            .unwrap();
3124        assert_eq!(response.status(), StatusCode::OK);
3125        assert_eq!(
3126            captured.lock().pop().unwrap().as_deref(),
3127            Some("Bearer tool-secret"),
3128            "tool bearer must reach the tool server"
3129        );
3130        tool_server.abort();
3131        provider.abort();
3132    }
3133
3134    #[tokio::test]
3135    async fn startup_refuses_pending_and_unaudited_tool_effects() {
3136        let directory = tempfile::tempdir().unwrap();
3137        let control_key = [31; 32];
3138        let mut headers = HeaderMap::new();
3139        headers.insert("x-av-session", HeaderValue::from_static("tool-recovery"));
3140        let body = br#"{"jsonrpc":"2.0","id":"recover","method":"tools/call","params":{"name":"read","arguments":{}}}"#;
3141        let mut execution =
3142            ToolExecution::from_request(&directory.path().to_string_lossy(), &headers, body, control_key)
3143                .unwrap();
3144        execution
3145            .bind_principal(&av_events::AgentIdentity {
3146                version: "dev".to_owned(),
3147                charter: "anonymous".into(),
3148                instance_uid: "anonymous".to_owned(),
3149                ttl_remaining_s: None,
3150            })
3151            .unwrap();
3152        execution.claim().await.unwrap();
3153        assert!(
3154            ensure_no_unresolved_tool_executions(directory.path(), &control_key)
3155                .await
3156                .is_err()
3157        );
3158        execution
3159            .persist(&ToolOutcome {
3160                status: 200,
3161                body_hex: hex::encode(body),
3162                content_type: Some("application/json".to_owned()),
3163            })
3164            .await
3165            .unwrap();
3166        assert!(
3167            ensure_no_unresolved_tool_executions(directory.path(), &control_key)
3168                .await
3169                .is_err()
3170        );
3171        execution.mark_audited().await.unwrap();
3172        ensure_no_unresolved_tool_executions(directory.path(), &control_key)
3173            .await
3174            .unwrap();
3175    }
3176
3177    #[tokio::test]
3178    async fn tool_redirect_is_terminal_and_never_replays_post() {
3179        let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3180        let tool_address = tool_listener.local_addr().unwrap();
3181        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
3182        let tool_router = Router::new()
3183            .route("/redirect", post(redirecting_tool))
3184            .route("/effect", post(redirected_effect))
3185            .with_state(Arc::clone(&calls));
3186        let tool_server = tokio::spawn(async move {
3187            axum::serve(tool_listener, tool_router).await.unwrap();
3188        });
3189        let directory = tempfile::tempdir().unwrap();
3190        let (mut state, provider) = test_state(directory.path()).await;
3191        Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
3192            Some(format!("http://{tool_address}/redirect"));
3193        let response = build_router(state)
3194            .oneshot(
3195                Request::builder()
3196                    .method("POST")
3197                    .uri("/v1/mcp")
3198                    .header("x-av-session", "tool-redirect")
3199                    .body(Body::from(
3200                        r#"{"jsonrpc":"2.0","id":"redirect","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
3201                    ))
3202                    .unwrap(),
3203            )
3204            .await
3205            .unwrap();
3206        assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT);
3207        assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 1);
3208        tool_server.abort();
3209        provider.abort();
3210    }
3211
3212    #[tokio::test]
3213    async fn concurrent_duplicate_tool_claims_conflict_without_leaking_io_detail() {
3214        let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3215        let tool_address = tool_listener.local_addr().unwrap();
3216        let held = Arc::new(HeldToolState {
3217            arrived: std::sync::atomic::AtomicBool::new(false),
3218            release: tokio::sync::Notify::new(),
3219        });
3220        let tool_router = Router::new()
3221            .route("/mcp", post(held_tool))
3222            .with_state(Arc::clone(&held));
3223        let tool_server = tokio::spawn(async move {
3224            axum::serve(tool_listener, tool_router).await.unwrap();
3225        });
3226        let directory = tempfile::tempdir().unwrap();
3227        let (mut state, provider) = test_state(directory.path()).await;
3228        Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
3229            Some(format!("http://{tool_address}/mcp"));
3230        let app = build_router(state);
3231        let request = || {
3232            Request::builder()
3233                .method("POST")
3234                .uri("/v1/mcp")
3235                .header("x-av-session", "claim-race")
3236                .body(Body::from(
3237                    r#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"read","arguments":{}}}"#,
3238                ))
3239                .unwrap()
3240        };
3241        let winner_app = app.clone();
3242        let winner = tokio::spawn(async move { winner_app.oneshot(request()).await.unwrap() });
3243        tokio::time::timeout(Duration::from_secs(1), async {
3244            while !held.arrived.load(std::sync::atomic::Ordering::Acquire) {
3245                tokio::task::yield_now().await;
3246            }
3247        })
3248        .await
3249        .unwrap();
3250        // The winner holds the claim while its upstream call is in flight;
3251        // an identical duplicate must lose the race with the canonical
3252        // uncertainty message and no io/filesystem detail.
3253        let loser = app.oneshot(request()).await.unwrap();
3254        assert_eq!(loser.status(), StatusCode::CONFLICT);
3255        let body = axum::body::to_bytes(loser.into_body(), 64 * 1024).await.unwrap();
3256        let text = String::from_utf8(body.to_vec()).unwrap();
3257        assert!(text.contains(TOOL_OUTCOME_UNCERTAIN), "{text}");
3258        assert!(!text.contains("os error"), "leaked io detail: {text}");
3259        held.release.notify_one();
3260        assert_eq!(winner.await.unwrap().status(), StatusCode::OK);
3261        tool_server.abort();
3262        provider.abort();
3263    }
3264
3265    #[tokio::test]
3266    async fn close_waits_for_tool_execution_and_completion_capture() {
3267        let tool_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
3268        let tool_address = tool_listener.local_addr().unwrap();
3269        let held = Arc::new(HeldToolState {
3270            arrived: std::sync::atomic::AtomicBool::new(false),
3271            release: tokio::sync::Notify::new(),
3272        });
3273        let tool_router = Router::new()
3274            .route("/mcp", post(held_tool))
3275            .with_state(Arc::clone(&held));
3276        let tool_server = tokio::spawn(async move {
3277            axum::serve(tool_listener, tool_router).await.unwrap();
3278        });
3279        let directory = tempfile::tempdir().unwrap();
3280        let (mut state, provider) = test_state(directory.path()).await;
3281        Arc::get_mut(&mut state.config).unwrap().tool_upstream_url =
3282            Some(format!("http://{tool_address}/mcp"));
3283        let app = build_router(state);
3284        let tool_app = app.clone();
3285        let tool_task = tokio::spawn(async move {
3286            tool_app
3287                .oneshot(
3288                    Request::builder()
3289                        .method("POST")
3290                        .uri("/v1/mcp")
3291                        .header("x-av-session", "held-tool")
3292                        .body(Body::from(
3293                            r#"{"jsonrpc":"2.0","id":9,"method":"tools/call","params":{"name":"read","arguments":{}}}"#,
3294                        ))
3295                        .unwrap(),
3296                )
3297                .await
3298                .unwrap()
3299        });
3300        tokio::time::timeout(Duration::from_secs(1), async {
3301            while !held.arrived.load(std::sync::atomic::Ordering::Acquire) {
3302                tokio::task::yield_now().await;
3303            }
3304        })
3305        .await
3306        .unwrap();
3307        let close_task = tokio::spawn(async move {
3308            app.oneshot(
3309                Request::builder()
3310                    .method("POST")
3311                    .uri("/v1/sessions/held-tool/close")
3312                    .body(Body::empty())
3313                    .unwrap(),
3314            )
3315            .await
3316            .unwrap()
3317        });
3318        tokio::task::yield_now().await;
3319        assert!(!close_task.is_finished(), "close overtook the executing tool");
3320        held.release.notify_one();
3321        assert_eq!(tool_task.await.unwrap().status(), StatusCode::OK);
3322        let close_response = close_task.await.unwrap();
3323        let close_status = close_response.status();
3324        let close_body = axum::body::to_bytes(close_response.into_body(), 64 * 1024)
3325            .await
3326            .unwrap();
3327        assert_eq!(
3328            close_status,
3329            StatusCode::OK,
3330            "close failed: {}",
3331            String::from_utf8_lossy(&close_body)
3332        );
3333        tool_server.abort();
3334        provider.abort();
3335    }
3336
3337    #[test]
3338    fn provider_usage_and_reasoning_are_preserved() {
3339        let raw = concat!(
3340            "data: {\"model\":\"gpt-test\",\"choices\":[{\"delta\":{",
3341            "\"content\":\"answer\",\"reasoning_content\":\"thought\"},\"finish_reason\":\"stop\"}],",
3342            "\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":3,\"cost_usd\":0.00125,",
3343            "\"prompt_tokens_details\":{\"cached_tokens\":4}}}\n\n"
3344        );
3345        let parsed = parse_provider_chunk(raw).unwrap().unwrap();
3346        assert_eq!(parsed.message, "answer");
3347        assert_eq!(parsed.reasoning.as_deref(), Some("thought"));
3348        assert_eq!(parsed.model_name.as_deref(), Some("gpt-test"));
3349        assert_eq!(parsed.metrics.prompt_tokens, Some(10));
3350        assert_eq!(parsed.metrics.completion_tokens, Some(3));
3351        assert_eq!(parsed.metrics.cached_tokens, Some(4));
3352        assert!(parsed.usage_reported);
3353        assert_eq!(parsed.finish_reason.as_deref(), Some("stop"));
3354        assert_eq!(parsed.cost_usd_micros, 1250);
3355        assert!(parsed.tool_call_deltas.is_empty());
3356    }
3357
3358    #[test]
3359    fn hostile_provider_metrics_are_rejected() {
3360        for raw in [
3361            r#"{"usage":{"completion_tokens":9007199254740993}}"#,
3362            r#"{"usage":{"prompt_tokens":-1}}"#,
3363            r#"{"usage":{"prompt_tokens_details":{"cached_tokens":"many"}}}"#,
3364            r#"{"usage":{"cost_usd":1e30}}"#,
3365            r#"{"usage":{"cost_usd":-1}}"#,
3366        ] {
3367            assert!(
3368                parse_provider_chunk(raw).is_err(),
3369                "accepted hostile metrics: {raw}"
3370            );
3371        }
3372    }
3373
3374    #[test]
3375    fn streaming_tool_call_arguments_reassemble() {
3376        let first = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call-1","function":{"name":"db_write","arguments":"{\"table\":"}}]}}]}"#;
3377        let second = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"items\",\"row\":{}}"}}]},"finish_reason":"tool_calls"}]}"#;
3378        let first_delta = parse_provider_chunk(first).unwrap().unwrap().tool_call_deltas;
3379        let second_delta = parse_provider_chunk(second).unwrap().unwrap().tool_call_deltas;
3380        assert_eq!(first_delta.len(), 1);
3381        assert_eq!(second_delta.len(), 1);
3382        let arguments = format!("{}{}", first_delta[0].arguments, second_delta[0].arguments);
3383        assert_eq!(
3384            serde_json::from_str::<Value>(&arguments).unwrap(),
3385            json!({"table": "items", "row": {}})
3386        );
3387        assert_eq!(map_finish_reason("tool_calls"), StopReason::ToolUse);
3388    }
3389
3390    #[test]
3391    fn hostile_provider_tool_call_index_is_rejected() {
3392        let raw = r#"data: {"choices":[{"delta":{"tool_calls":[{"index":18446744073709551615,"function":{"name":"n","arguments":"{}"}}]}}]}"#;
3393        let parsed = parse_provider_chunk(raw).unwrap().unwrap();
3394        assert_eq!(parsed.tool_call_deltas.len(), 1);
3395        assert_eq!(parsed.tool_call_deltas[0].index, u64::MAX);
3396    }
3397
3398    #[test]
3399    fn non_stream_tool_calls_receive_distinct_positional_indices() {
3400        let parsed = parse_provider_chunk(
3401            r#"{"choices":[{"message":{"tool_calls":[{"id":"a","function":{"name":"read","arguments":"{}"}},{"id":"b","function":{"name":"write","arguments":"{}"}}]}}]}"#,
3402        )
3403        .unwrap()
3404        .unwrap();
3405        assert_eq!(parsed.tool_call_deltas.len(), 2);
3406        assert_eq!(parsed.tool_call_deltas[0].index, 0);
3407        assert_eq!(parsed.tool_call_deltas[1].index, 1);
3408    }
3409
3410    #[test]
3411    fn sse_frame_boundary_accepts_crlf_and_cr() {
3412        assert_eq!(sse_frame_end(b"data: {}\r\n\r"), None);
3413        assert_eq!(sse_frame_end(b"data: {}\r\n\r\nrest"), Some(12));
3414        assert_eq!(sse_frame_end(b"data: {}\r\rrest"), Some(10));
3415    }
3416
3417    /// Leading BOM (U+FEFF) on an SSE stream must be stripped per
3418    /// §9.2.4. Rust's `.trim()` and `.trim_start()` do NOT remove
3419    /// U+FEFF (char::is_whitespace returns false for it), so without
3420    /// an explicit strip the very first `data:` line would parse as
3421    /// `"\u{FEFF}data:"` and the frame would take the non-SSE path
3422    /// where `serde_json::from_str` chokes on the BOM prefix.
3423    #[test]
3424    fn parse_provider_chunk_strips_leading_bom() {
3425        let raw = "\u{feff}data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n";
3426        let parsed = parse_provider_chunk(raw)
3427            .expect("BOM-prefixed SSE frame must parse")
3428            .expect("must yield a chunk");
3429        assert_eq!(parsed.message, "hi");
3430    }
3431
3432    /// SSE §9.2.8: an SSE client dispatches events by name, so
3433    /// `event: error\ndata: {model-shaped-JSON}` is delivered to the
3434    /// browser/SDK's `error` listener — NOT the default `message`
3435    /// listener. AgentVisor AI attributes captured `data:` payloads to
3436    /// the signed receipt / ATIF as if they were model output, so a
3437    /// hostile upstream (rogue provider, TLS-MITM at egress mesh, a
3438    /// misconfigured caching proxy) could forge receipt content while
3439    /// the user's UI showed nothing suspicious. Refuse the frame with
3440    /// a diagnostic so the caller marks capture-failed rather than
3441    /// signs attributable content that never was displayed.
3442    #[test]
3443    fn parse_provider_chunk_refuses_non_message_sse_event_types() {
3444        for event in ["error", "ping", "custom_signal"] {
3445            let raw = format!(
3446                "event: {event}\ndata: {{\"choices\":[{{\"delta\":{{\"content\":\"forged\"}}}}]}}\n\n"
3447            );
3448            let err = match parse_provider_chunk(&raw) {
3449                Ok(_) => panic!("event type {event:?} must be refused"),
3450                Err(error) => error,
3451            };
3452            assert!(
3453                err.contains("unsupported event type") && err.contains(event),
3454                "expected diagnostic naming event type {event:?}, got: {err}"
3455            );
3456        }
3457        // Explicit `event: message` is accepted (spec default) — must
3458        // NOT be refused by the same guard.
3459        let ok_raw = "event: message\ndata: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n";
3460        let parsed = match parse_provider_chunk(ok_raw) {
3461            Ok(Some(chunk)) => chunk,
3462            Ok(None) => panic!("explicit event: message must yield a chunk"),
3463            Err(error) => panic!("explicit event: message must be accepted, got: {error}"),
3464        };
3465        assert_eq!(parsed.message, "hi");
3466    }
3467
3468    #[test]
3469    fn parse_provider_chunk_treats_done_variants_as_stream_end() {
3470        for raw in [
3471            "data: [DONE]\ndata: [DONE]\n\n",
3472            "data: [DONE]  \n\n",
3473            "data:\ndata: [DONE]\n\n",
3474        ] {
3475            let result = parse_provider_chunk(raw);
3476            match result {
3477                Ok(None) => {}
3478                Ok(Some(_)) => panic!("[DONE] variant must terminate stream (None) for {raw:?}"),
3479                Err(error) => panic!("[DONE] variant must not fail parse: {raw:?}: {error}"),
3480            }
3481        }
3482    }
3483
3484    /// SSE §9.2.6: the `data:` field strips exactly one leading
3485    /// U+0020 SPACE — no more, and no other whitespace class. A
3486    /// regression to `.trim_start()` would eat runs of Unicode
3487    /// whitespace and silently corrupt payloads whose second byte
3488    /// after `data:` is another space or a tab.
3489    ///
3490    /// Distinguishing input: `data:  hi\n\n` — two spaces then a
3491    /// literal `hi`. Non-SSE parse via serde_json would fail on
3492    /// `"hi"` unquoted, so the frame goes through the SSE path where
3493    /// the `data:` accumulator kept one leading space (spec-compliant)
3494    /// and serde_json then fails on ` hi` — this is exactly what the
3495    /// spec's "leave everything after the first space verbatim"
3496    /// behaviour dictates. A `.trim_start()` regression would silently
3497    /// consume both spaces and produce the same (still-invalid) input.
3498    /// We instead pick a JSON payload where the number of leading
3499    /// spaces changes the parse result: `{"choices":[{"delta":{"content":" hi"}}]}`
3500    /// with wire `data:  {"choices"...}` — under `strip_prefix(' ')`
3501    /// the accumulated frame starts with a space then `{`, both parse
3502    /// fine and the extracted `content` field is `" hi"`; under
3503    /// `.trim_start()` the frame starts with `{`, also parses fine and
3504    /// yields the same `" hi"`. So parse content doesn't discriminate.
3505    ///
3506    /// Instead observe the raw data buffer that the parser builds:
3507    /// use a plain non-JSON `data:` value and assert the trimmed
3508    /// prefix. Since the parser now runs serde_json on the assembled
3509    /// frame, the only way to expose the trim behaviour is a test on
3510    /// a lower-level helper. We keep this test as a smoke check on a
3511    /// case where the frame *fails* differently: two-space prefix
3512    /// with a leading space kept produces a JSON parse failure the
3513    /// old `trim_start` would not have produced.
3514    #[test]
3515    fn parse_provider_chunk_strips_exactly_one_leading_space() {
3516        // Bare non-JSON `hello` — invalid JSON either way. With
3517        // `strip_prefix(' ')` (spec-correct) the parser accumulates
3518        // ` \thello` (space+tab+hello). With `.trim_start()` (buggy)
3519        // it accumulates `hello`. Both produce `Err(...)` from
3520        // serde_json but the reported column differs — the trimmed
3521        // form reports column 1 (immediate `h`); the correctly
3522        // preserved form reports a later column because of the
3523        // retained tab.
3524        let raw = "data:  \thello\n\n";
3525        let result = parse_provider_chunk(raw);
3526        let error = match result {
3527            Ok(_) => panic!("bare hello must not parse as valid provider frame"),
3528            Err(error) => error,
3529        };
3530        assert!(
3531            error.contains("column 2") || error.contains("column 3"),
3532            "expected error to name a column past the retained leading space+tab; \
3533             got: {error}"
3534        );
3535    }
3536
3537    #[tokio::test]
3538    async fn split_non_sse_utf8_is_captured_as_one_json_document() {
3539        let directory = tempfile::tempdir().unwrap();
3540        let (state, provider) = test_state(directory.path()).await;
3541        let app = build_router(state.clone());
3542        let response = app
3543            .oneshot(chat_request_with_payload(
3544                "split-json",
3545                json!({
3546                    "model": "split-json",
3547                    "stream": false,
3548                    "messages": [{"role": "user", "content": "hello"}]
3549                }),
3550            ))
3551            .await
3552            .unwrap();
3553        let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
3554            .await
3555            .unwrap();
3556        assert_eq!(
3557            serde_json::from_slice::<Value>(&body).unwrap()["choices"][0]["message"]["content"],
3558            "héllo"
3559        );
3560        let session = state.sessions.get("split-json").unwrap();
3561        session.wait_for_worker_jobs().await;
3562        let crate::reconciler::FinalizeOutcome::Atif { path } = state
3563            .finalizer
3564            .close_session(session, StopReason::SessionClosed)
3565            .await
3566            .unwrap()
3567        else {
3568            panic!("expected ATIF")
3569        };
3570        let trajectory: av_atif::Trajectory = serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap();
3571        assert_eq!(trajectory.steps[1].message, Value::String("héllo".to_owned()));
3572        provider.abort();
3573    }
3574
3575    #[tokio::test]
3576    async fn malformed_non_sse_json_fails_capture_before_delivery() {
3577        let directory = tempfile::tempdir().unwrap();
3578        let (state, provider) = test_state(directory.path()).await;
3579        let app = build_router(state.clone());
3580        let response = app
3581            .oneshot(chat_request_with_payload(
3582                "malformed-json",
3583                json!({
3584                    "model": "malformed-json",
3585                    "stream": false,
3586                    "messages": [{"role": "user", "content": "hello"}]
3587                }),
3588            ))
3589            .await
3590            .unwrap();
3591        // The non-SSE path drains the relay server-side, so a capture
3592        // failure surfaces as a clean 502 + JSON error instead of a
3593        // committed 200 with a severed body.
3594        assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
3595        let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
3596            .await
3597            .unwrap();
3598        assert!(serde_json::from_slice::<Value>(&body).unwrap()["error"].is_string());
3599        let session = state.sessions.get("malformed-json").unwrap();
3600        assert!(session.capture_failed());
3601        provider.abort();
3602    }
3603
3604    #[tokio::test]
3605    async fn repeated_provider_responses_open_the_loop_breaker() {
3606        let directory = tempfile::tempdir().unwrap();
3607        let (mut state, provider) = test_state(directory.path()).await;
3608        let config = Arc::get_mut(&mut state.config).unwrap();
3609        config.breaker.min_tokens = 0;
3610        config.breaker.window = 3;
3611        let app = build_router(state.clone());
3612        for _ in 0..4 {
3613            let response = app.clone().oneshot(chat_request("response-loop")).await.unwrap();
3614            assert_eq!(response.status(), StatusCode::OK);
3615            axum::body::to_bytes(response.into_body(), 64 * 1024)
3616                .await
3617                .unwrap();
3618            state
3619                .sessions
3620                .get("response-loop")
3621                .unwrap()
3622                .wait_for_worker_jobs()
3623                .await;
3624        }
3625        let blocked = app.oneshot(chat_request("response-loop")).await.unwrap();
3626        // Loop-breaker verdicts are permanent for the current session:
3627        // 403 stops mainstream LLM-SDK auto-retry loops that would
3628        // otherwise burn budget re-hitting the same breaker (see
3629        // pipeline.rs::PipelineError::status).
3630        assert_eq!(blocked.status(), StatusCode::FORBIDDEN);
3631        provider.abort();
3632    }
3633
3634    #[tokio::test]
3635    async fn empty_and_failed_upstream_responses_are_audited_as_failures() {
3636        let directory = tempfile::tempdir().unwrap();
3637        let (state, provider) = test_state(directory.path()).await;
3638        let app = build_router(state.clone());
3639        for (session_id, model) in [("empty-error", "empty-error"), ("json-error", "json-error")] {
3640            let response = app
3641                .clone()
3642                .oneshot(chat_request_with_payload(
3643                    session_id,
3644                    json!({
3645                        "model": model,
3646                        "stream": false,
3647                        "messages": [{"role": "user", "content": "hello"}]
3648                    }),
3649                ))
3650                .await
3651                .unwrap();
3652            assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
3653            axum::body::to_bytes(response.into_body(), 64 * 1024)
3654                .await
3655                .unwrap();
3656            let session = state.sessions.get(session_id).unwrap();
3657            session.wait_for_worker_jobs().await;
3658            assert!(!session.capture_failed());
3659            let records = active_records(directory.path(), &state, session_id);
3660            assert_eq!(records.len(), 2);
3661            let event: av_events::OcsfEvent =
3662                serde_json::from_value(records.last().unwrap().event.clone()).unwrap();
3663            assert_eq!(event.class_name, av_events::EventClass::StopReason);
3664            assert_eq!(event.status_id, av_events::StatusId::Failure.id());
3665            assert_eq!(event.payload["http_status"], 500);
3666        }
3667        provider.abort();
3668    }
3669
3670    #[tokio::test]
3671    async fn regressing_provider_usage_fails_capture_before_later_bytes() {
3672        let directory = tempfile::tempdir().unwrap();
3673        let (state, provider) = test_state(directory.path()).await;
3674        let response = build_router(state.clone())
3675            .oneshot(chat_request_with_payload(
3676                "regressive-usage",
3677                json!({
3678                    "model": "regressive-usage",
3679                    "stream": true,
3680                    "messages": [{"role": "user", "content": "hello"}]
3681                }),
3682            ))
3683            .await
3684            .unwrap();
3685        assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
3686            .await
3687            .is_err());
3688        let session = state.sessions.get("regressive-usage").unwrap();
3689        assert!(session.capture_failed());
3690        session.wait_for_worker_jobs().await;
3691        let records = active_records(directory.path(), &state, "regressive-usage");
3692        assert_eq!(records.len(), 1);
3693        assert!(records[0]
3694            .response_attempt
3695            .as_ref()
3696            .is_some_and(|attempt| !attempt.terminal));
3697        provider.abort();
3698    }
3699
3700    #[tokio::test]
3701    async fn successful_provider_response_requires_choices() {
3702        let directory = tempfile::tempdir().unwrap();
3703        let (state, provider) = test_state(directory.path()).await;
3704        let response = build_router(state.clone())
3705            .oneshot(chat_request_with_payload(
3706                "empty-success",
3707                json!({
3708                    "model": "empty-success",
3709                    "stream": false,
3710                    "messages": [{"role": "user", "content": "hello"}]
3711                }),
3712            ))
3713            .await
3714            .unwrap();
3715        // Missing choices fails capture; the server-side drain converts
3716        // the refusal into a 502 + JSON error rather than a severed body.
3717        assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
3718        let body = axum::body::to_bytes(response.into_body(), 64 * 1024)
3719            .await
3720            .unwrap();
3721        assert!(serde_json::from_slice::<Value>(&body).unwrap()["error"].is_string());
3722        assert!(state.sessions.get("empty-success").unwrap().capture_failed());
3723        provider.abort();
3724    }
3725
3726    #[tokio::test]
3727    async fn cumulative_usage_settles_against_provisional_charge() {
3728        let directory = tempfile::tempdir().unwrap();
3729        let payload = json!({
3730            "model": "provisional-usage",
3731            "stream": true,
3732            "messages": [{"role": "user", "content": "hello"}]
3733        });
3734        let prompt_tokens = av_core::tokens::approx_tokens(&payload.to_string());
3735        let (state, provider) = test_state_with_token_cap(directory.path(), Some(prompt_tokens + 2)).await;
3736        let response = build_router(state.clone())
3737            .oneshot(chat_request_with_payload("provisional-usage", payload))
3738            .await
3739            .unwrap();
3740        assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
3741            .await
3742            .is_ok());
3743        let session = state.sessions.get("provisional-usage").unwrap();
3744        session.wait_for_worker_jobs().await;
3745        assert_eq!(
3746            session
3747                .totals
3748                .completion_tokens
3749                .load(std::sync::atomic::Ordering::Acquire),
3750            2
3751        );
3752        provider.abort();
3753    }
3754
3755    #[tokio::test]
3756    async fn provider_tool_call_index_out_of_range_fails_capture() {
3757        let directory = tempfile::tempdir().unwrap();
3758        let (state, provider) = test_state(directory.path()).await;
3759        let response = build_router(state.clone())
3760            .oneshot(chat_request_with_payload(
3761                "tool-oob-index",
3762                json!({
3763                    "model": "tool-oob-index",
3764                    "stream": true,
3765                    "messages": [{"role": "user", "content": "hello"}]
3766                }),
3767            ))
3768            .await
3769            .unwrap();
3770        assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
3771            .await
3772            .is_err());
3773        assert!(state.sessions.get("tool-oob-index").unwrap().capture_failed());
3774        provider.abort();
3775    }
3776
3777    #[tokio::test]
3778    async fn tool_arguments_without_usage_are_budgeted_before_delivery() {
3779        let directory = tempfile::tempdir().unwrap();
3780        let payload = json!({
3781            "model": "tool-no-usage",
3782            "stream": true,
3783            "messages": [{"role": "user", "content": "hello"}]
3784        });
3785        let prompt_tokens = av_core::tokens::approx_tokens(&payload.to_string());
3786        let (state, provider) = test_state_with_token_cap(directory.path(), Some(prompt_tokens)).await;
3787        let response = build_router(state)
3788            .oneshot(chat_request_with_payload("tool-no-usage", payload))
3789            .await
3790            .unwrap();
3791        assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
3792            .await
3793            .is_err());
3794        provider.abort();
3795    }
3796
3797    #[tokio::test(flavor = "current_thread")]
3798    async fn completion_budget_store_does_not_block_stream_runtime() {
3799        let directory = tempfile::tempdir().unwrap();
3800        let (mut state, provider) = test_state_with_token_cap(directory.path(), Some(10_000)).await;
3801        let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
3802        state.store = Arc::new(SlowStore {
3803            inner: InMemoryStore::new(),
3804            calls: Arc::clone(&calls),
3805        });
3806        let response = build_router(state)
3807            .oneshot(chat_request("slow-budget"))
3808            .await
3809            .unwrap();
3810        let body = tokio::spawn(async move {
3811            axum::body::to_bytes(response.into_body(), 64 * 1024)
3812                .await
3813                .unwrap()
3814        });
3815
3816        let started = std::time::Instant::now();
3817        tokio::time::sleep(Duration::from_millis(10)).await;
3818        assert!(
3819            started.elapsed() < Duration::from_millis(75),
3820            "completion budget storage blocked the Tokio reactor"
3821        );
3822        assert!(!body.await.unwrap().is_empty());
3823        assert_eq!(calls.load(std::sync::atomic::Ordering::Acquire), 2);
3824        provider.abort();
3825    }
3826
3827    #[tokio::test]
3828    async fn completion_tokens_are_blocked_before_delivery() {
3829        let directory = tempfile::tempdir().unwrap();
3830        let prompt_tokens = av_core::tokens::approx_tokens(&chat_payload().to_string());
3831        let (state, provider) = test_state_with_token_cap(directory.path(), Some(prompt_tokens)).await;
3832        let app = build_router(state.clone());
3833        let response = app.oneshot(chat_request("completion-budget")).await.unwrap();
3834        assert_eq!(response.status(), StatusCode::OK);
3835        assert!(axum::body::to_bytes(response.into_body(), 64 * 1024)
3836            .await
3837            .is_err());
3838        tokio::time::timeout(Duration::from_secs(1), async {
3839            loop {
3840                if state
3841                    .sessions
3842                    .get("completion-budget")
3843                    .is_some_and(|session| session.is_closed())
3844                {
3845                    break;
3846                }
3847                tokio::task::yield_now().await;
3848            }
3849        })
3850        .await
3851        .unwrap();
3852        provider.abort();
3853    }
3854
3855    /// Regression lock for the upstream-response header forwarding
3856    /// trust-boundary bug (CWE-346 / CWE-16). Before the fix, every
3857    /// upstream response header except `Content-Length` was blindly
3858    /// forwarded to the client — letting a hostile or MITM'd upstream
3859    ///
3860    /// - set cookies on our domain (`Set-Cookie`),
3861    /// - open CORS on our origin (`Access-Control-Allow-*`),
3862    /// - leak upstream implementation identity (`Server`, `Via`,
3863    ///   `X-Powered-By`, `X-Request-ID`),
3864    /// - inject conflicting framing/hop-by-hop metadata (`Connection`,
3865    ///   `Transfer-Encoding`, `Keep-Alive`, `Upgrade`, `TE`, `Trailer`,
3866    ///   `Proxy-Authenticate`, `Proxy-Authorization`).
3867    ///
3868    /// `is_forwardable_upstream_header` is now the sole gate on which
3869    /// upstream headers cross the proxy trust boundary — this test locks
3870    /// each dangerous class out and confirms benign headers (`Content-Type`,
3871    /// `Cache-Control`, `ETag`) still pass through.
3872    #[test]
3873    fn upstream_response_headers_do_not_cross_proxy_trust_boundary() {
3874        use axum::http::HeaderName;
3875        let dangerous = [
3876            // Cookie injection on our domain from a hostile upstream.
3877            "set-cookie",
3878            // CORS bypass — we never let the upstream open our origin.
3879            "access-control-allow-origin",
3880            "access-control-allow-credentials",
3881            "access-control-allow-methods",
3882            "access-control-allow-headers",
3883            "access-control-expose-headers",
3884            "access-control-max-age",
3885            // RFC 7230 §6.1 hop-by-hop headers — a proxy must not
3886            // forward these; forwarding `Transfer-Encoding` enables
3887            // classical HTTP request smuggling.
3888            "connection",
3889            "keep-alive",
3890            "transfer-encoding",
3891            "upgrade",
3892            "te",
3893            "trailer",
3894            "proxy-authenticate",
3895            "proxy-authorization",
3896            // Framing metadata that hyper computes from the response
3897            // body; the upstream's value would be wrong.
3898            "content-length",
3899            // Implementation-identity leaks.
3900            "server",
3901            "via",
3902            "x-powered-by",
3903            "x-request-id",
3904        ];
3905        for name in dangerous {
3906            let header = HeaderName::from_static(name);
3907            assert!(
3908                !is_forwardable_upstream_header(&header),
3909                "dangerous header {name:?} must not be forwarded to the client"
3910            );
3911        }
3912
3913        let benign = [
3914            "content-type",
3915            "content-encoding",
3916            "cache-control",
3917            "etag",
3918            "last-modified",
3919            "vary",
3920            "expires",
3921        ];
3922        for name in benign {
3923            let header = HeaderName::from_static(name);
3924            assert!(
3925                is_forwardable_upstream_header(&header),
3926                "benign header {name:?} must still be forwarded"
3927            );
3928        }
3929    }
3930
3931    /// Regression lock for CWE-209 information exposure on the MCP
3932    /// tool-call path — the second occurrence of pass 17's chat/completion
3933    /// leak. Before this fix, when the tool-upstream request failed
3934    /// (unroutable host, timeout, TLS error, …) `mcp_call` returned
3935    /// `lifecycle_error(format!("forward tool call: {reqwest_err}"))`
3936    /// which — because `reqwest::Error::Display` embeds the request URL —
3937    /// leaked the operator-configured tool-upstream URL (potentially an
3938    /// internal hostname) to any client that could hit `/mcp`. The
3939    /// `read_limited_tool_response` mid-stream error path had the same
3940    /// bug. The client-facing message must now be a stable, non-identifying
3941    /// category (e.g. `"upstream unreachable"`), matching pass 17's
3942    /// classifier categories.
3943    #[tokio::test]
3944    async fn mcp_tool_upstream_failure_does_not_leak_configured_url() {
3945        // Distinctive sentinel host so any regression is unmissable.
3946        let sentinel_host = "internal-mcp-sentinel-host.corp.example";
3947        let sentinel_url = format!("http://{sentinel_host}:65002/mcp");
3948        let directory = tempfile::tempdir().unwrap();
3949        let (mut state, provider) = test_state(directory.path()).await;
3950        Arc::get_mut(&mut state.config).unwrap().tool_upstream_url = Some(sentinel_url);
3951        let response = build_router(state)
3952            .oneshot(
3953                Request::builder()
3954                    .method("POST")
3955                    .uri("/v1/mcp")
3956                    .header("x-av-session", "mcp-cwe-209-check")
3957                    .body(Body::from(
3958                        r#"{"jsonrpc":"2.0","id":"leak","method":"tools/call","params":{"name":"read","arguments":{}}}"#,
3959                    ))
3960                    .unwrap(),
3961            )
3962            .await
3963            .unwrap();
3964        assert_eq!(
3965            response.status(),
3966            StatusCode::BAD_GATEWAY,
3967            "tool-upstream faults must surface as 502 like the chat relay, not 500"
3968        );
3969        let bytes = axum::body::to_bytes(response.into_body(), 64 * 1024)
3970            .await
3971            .unwrap();
3972        let body = String::from_utf8(bytes.to_vec()).unwrap();
3973        assert!(
3974            !body.contains(sentinel_host),
3975            "MCP forward-tool error body {body:?} leaks the configured tool-upstream host"
3976        );
3977        assert!(
3978            !body.contains("65002"),
3979            "MCP forward-tool error body {body:?} leaks the configured tool-upstream port"
3980        );
3981        assert!(
3982            !body.contains("corp.example"),
3983            "MCP forward-tool error body {body:?} leaks the configured tool-upstream domain"
3984        );
3985        provider.abort();
3986    }
3987
3988    #[tokio::test]
3989    async fn pipeline_error_carries_retry_after_and_www_authenticate() {
3990        // Contract: 503 SHOULD carry Retry-After (RFC 7231 §7.1.3);
3991        // 401 MUST carry WWW-Authenticate (RFC 7235 §3.1). These
3992        // headers are the mechanism by which intermediaries and SDKs
3993        // decide to back off / prompt for credentials — without them,
3994        // mainstream LLM SDKs interpret 503 as "retry immediately"
3995        // and 401 as "broken proxy".
3996        use crate::pipeline::PipelineError;
3997
3998        let unavailable = pipeline_error(PipelineError::Unavailable("worker queue full".into()));
3999        assert_eq!(unavailable.status(), StatusCode::SERVICE_UNAVAILABLE);
4000        assert_eq!(
4001            unavailable
4002                .headers()
4003                .get(axum::http::header::RETRY_AFTER)
4004                .and_then(|v| v.to_str().ok()),
4005            Some("5")
4006        );
4007
4008        let upstream = pipeline_error(PipelineError::Upstream("bad gateway".into()));
4009        assert_eq!(upstream.status(), StatusCode::BAD_GATEWAY);
4010        assert!(upstream.headers().get(axum::http::header::RETRY_AFTER).is_some());
4011
4012        let unauthorized = pipeline_error(PipelineError::Unauthorized("bad token".into()));
4013        assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED);
4014        let auth = unauthorized
4015            .headers()
4016            .get(axum::http::header::WWW_AUTHENTICATE)
4017            .and_then(|v| v.to_str().ok())
4018            .unwrap_or_default();
4019        assert!(
4020            auth.starts_with("Bearer"),
4021            "WWW-Authenticate must be a Bearer challenge, got {auth:?}"
4022        );
4023
4024        // Permanent verdicts must NOT carry Retry-After — that would
4025        // encourage the very retry loop we're trying to stop.
4026        let blocked = pipeline_error(PipelineError::Blocked("loop breaker".into()));
4027        assert_eq!(blocked.status(), StatusCode::FORBIDDEN);
4028        assert!(blocked.headers().get(axum::http::header::RETRY_AFTER).is_none());
4029        let abort = pipeline_error(PipelineError::Abort("stop retrying".into()));
4030        assert_eq!(abort.status(), StatusCode::CONFLICT);
4031        assert!(abort.headers().get(axum::http::header::RETRY_AFTER).is_none());
4032    }
4033
4034    /// Round-25 F3: SSE detection is case-insensitive and tolerates
4035    /// media-type parameters. Byte-exact `starts_with` previously
4036    /// misclassified `Text/Event-Stream` (some CDNs re-title-case)
4037    /// and any `text/event-stream; charset=utf-8` with a leading
4038    /// title-cased subtype. Misclassified SSE streams get buffered
4039    /// to the 16 MiB provider cap and either lose streaming
4040    /// semantics or 502 despite being valid.
4041    #[test]
4042    fn is_sse_content_type_is_case_insensitive_and_param_tolerant() {
4043        fn ct(value: &str) -> HeaderMap {
4044            let mut headers = HeaderMap::new();
4045            headers.insert(
4046                axum::http::header::CONTENT_TYPE,
4047                value.parse().expect("header value"),
4048            );
4049            headers
4050        }
4051        assert!(is_sse_content_type(&ct("text/event-stream")));
4052        assert!(is_sse_content_type(&ct("Text/Event-Stream")));
4053        assert!(is_sse_content_type(&ct("TEXT/EVENT-STREAM")));
4054        assert!(is_sse_content_type(&ct("text/event-stream; charset=utf-8")));
4055        assert!(is_sse_content_type(&ct("Text/Event-Stream ; charset=utf-8")));
4056        // Non-SSE and superstring both must not match.
4057        assert!(!is_sse_content_type(&ct("application/json")));
4058        assert!(!is_sse_content_type(&ct("text/event-stream-json")));
4059        // Absent header returns false.
4060        assert!(!is_sse_content_type(&HeaderMap::new()));
4061    }
4062
4063    /// Round-31 F5: OPTIONS to every mutating route replies with
4064    /// `204 No Content` and NO Access-Control-Allow-* headers.
4065    /// Browsers must interpret this as "cross-origin denied" and
4066    /// refuse the actual request — making the same-origin-only
4067    /// posture explicit at the wire. Guards against a future PR
4068    /// accidentally adding `CorsLayer::permissive()`.
4069    #[tokio::test]
4070    async fn options_returns_204_without_cors_headers() {
4071        let directory = tempfile::tempdir().unwrap();
4072        let (state, provider) = test_state(directory.path()).await;
4073        for path in [
4074            "/v1/chat/completions",
4075            "/v1/mcp",
4076            "/mcp",
4077            "/v1/sessions/some-id/close",
4078            "/v1/sessions/some-id/promote",
4079        ] {
4080            let response = build_router(state.clone())
4081                .oneshot(
4082                    Request::builder()
4083                        .method(axum::http::Method::OPTIONS)
4084                        .uri(path)
4085                        .header("origin", "http://attacker.example")
4086                        .header("access-control-request-method", "POST")
4087                        .body(Body::empty())
4088                        .unwrap(),
4089                )
4090                .await
4091                .unwrap();
4092            assert_eq!(
4093                response.status(),
4094                StatusCode::NO_CONTENT,
4095                "OPTIONS {path} should return 204, got {}",
4096                response.status()
4097            );
4098            for header in [
4099                "access-control-allow-origin",
4100                "access-control-allow-methods",
4101                "access-control-allow-headers",
4102                "access-control-allow-credentials",
4103            ] {
4104                assert!(
4105                    response.headers().get(header).is_none(),
4106                    "OPTIONS {path} must not emit {header} — got {:?}",
4107                    response.headers().get(header)
4108                );
4109            }
4110        }
4111        provider.abort();
4112    }
4113
4114    /// Round-32 F2: cached MCP tool-outcome replays preserve the
4115    /// upstream Content-Type end-to-end (including through the
4116    /// on-disk journal roundtrip). Strict JSON-RPC 2.0 clients
4117    /// (spec: MUST be `application/json`) would previously receive
4118    /// `application/octet-stream` on both the fresh forward and the
4119    /// cached replay. This test locks in:
4120    ///   1. `ToolOutcome::into_response` honours a set content_type.
4121    ///   2. A `None` content_type defaults to `application/json`
4122    ///      (MCP is JSON-RPC 2.0 by convention).
4123    ///   3. Serialising and re-deserialising a `ToolOutcome` (i.e.
4124    ///      the on-disk journal round-trip) preserves the field.
4125    ///   4. Legacy journals without the field decode as `None` and
4126    ///      pick up the default via #1.
4127    #[test]
4128    fn round_32_f2_tool_outcome_preserves_content_type() {
4129        // (1) Custom content_type survives.
4130        let outcome = ToolOutcome {
4131            status: 200,
4132            body_hex: hex::encode(b"{\"jsonrpc\":\"2.0\"}"),
4133            content_type: Some("application/problem+json".to_owned()),
4134        };
4135        let response = outcome.into_response();
4136        assert_eq!(
4137            response
4138                .headers()
4139                .get(axum::http::header::CONTENT_TYPE)
4140                .and_then(|v| v.to_str().ok()),
4141            Some("application/problem+json")
4142        );
4143        // (2) None -> application/json default.
4144        let outcome_no_ct = ToolOutcome {
4145            status: 200,
4146            body_hex: hex::encode(b"{}"),
4147            content_type: None,
4148        };
4149        let response = outcome_no_ct.into_response();
4150        assert_eq!(
4151            response
4152                .headers()
4153                .get(axum::http::header::CONTENT_TYPE)
4154                .and_then(|v| v.to_str().ok()),
4155            Some("application/json")
4156        );
4157        // (3) Journal round-trip preserves the field.
4158        let sealed = serde_json::to_string(&ToolOutcome {
4159            status: 429,
4160            body_hex: hex::encode(b"{}"),
4161            content_type: Some("text/plain".to_owned()),
4162        })
4163        .unwrap();
4164        let recovered: ToolOutcome = serde_json::from_str(&sealed).unwrap();
4165        assert_eq!(recovered.content_type.as_deref(), Some("text/plain"));
4166        // (4) Legacy journal without the field decodes as None.
4167        let legacy = r#"{"status":200,"body_hex":"7b7d"}"#;
4168        let outcome: ToolOutcome = serde_json::from_str(legacy).unwrap();
4169        assert!(outcome.content_type.is_none());
4170    }
4171}
av_harness/routes.rs

av_harness/
routes.rs
